import numpy as np
import subprocess as sp
import threading
import gc
import sys
from functools import reduce
import struct

# Constants
presets = ("ultrafast", "superfast", "veryfast", "faster", "fast", "medium", "slow", "slower", "veryslow")

def main():
	
	arg_index = 1

	# Parse first argument
	if len(sys.argv) < 2:
		print_usage()
	COMPRESS = 0
	DECOMPRESS = 1
	if sys.argv[arg_index] == "-c":
		mode = COMPRESS
	elif sys.argv[arg_index] == "-d":
		mode = DECOMPRESS
	else:
		raise Exception("First argument must be -c or -d!")
	arg_index += 1
	
	# Parse general arguments
	if mode == COMPRESS:
		if len(sys.argv) < 9:
			print_usage()
	else:
		if len(sys.argv) < 4 or len(sys.argv) > 6:
			print_usage()
	input_file = sys.argv[arg_index]
	arg_index += 1
	output_file = sys.argv[arg_index]
	arg_index += 1
	
	# Parse compression arguments
	if mode == COMPRESS:
		crf = int(sys.argv[arg_index])
		arg_index += 1
		dtype = sys.argv[arg_index]
		arg_index += 1
		time_mode_index = int(sys.argv[arg_index])
		arg_index += 1
		sizes = []
		while sys.argv[arg_index].isnumeric():
			sizes.append(int(sys.argv[arg_index]))
			arg_index += 1
		
		preset = "medium"
		if arg_index < len(sys.argv) and sys.argv[arg_index] == "-p":
				arg_index += 1
				if arg_index >= len(sys.argv):
					print_usage()
				if sys.argv[arg_index] in presets:
					preset = sys.argv[arg_index]
				arg_index += 1
	
	threads = 1
	if arg_index < len(sys.argv) and sys.argv[arg_index] == "-t":
		arg_index += 1
		if arg_index >= len(sys.argv):
			print_usage()
		threads = int(sys.argv[arg_index])
		arg_index += 1
	
	# Perform compression/decompression
	if mode == COMPRESS:
		compress_video(input_file, output_file, crf, dtype, time_mode_index, sizes, preset=preset, threads=threads)
	else:
		decompress_video(input_file, output_file, threads=threads)

def print_usage():
	print("Usage:")
	print("- python3 x264_compressor.py -c <input-file> <output-file> <crf> <dtype> <time-mode-index> <sizes (MSI first)> [-p <preset>] [-t <threads>]")
	print("- python3 x264_compressor.py -d <input-file> <output-file> [-t <threads>]")
	exit()

def compress_video(input_file, output_file, crf, dtype, time_mode_index, sizes, preset="medium", threads=1):
	
	# Compresses using video compression (only for 3D data)
	# crf is the quality parameter. The range of the CRF scale is 0-51, where 0 is lossless, 23 is the default, and 51 is worst quality possible.
	
	# Load data from file
	size = reduce(lambda x, y: x*y, sizes)
	modes = len(sizes)
	data = np.fromfile(input_file, dtype=dtype, count=size).reshape(sizes)
	
	# Transpose and reshape data such that the data is 3D and the first mode corresponds to the requested time mode
	data = np.transpose(data, axes=[time_mode_index,] + list(range(time_mode_index)) + list(range(time_mode_index + 1, modes)))
	data = data.reshape((data.shape[0], data.shape[1], -1))
	reordered_shape = data.shape
	
	# Quantize data and convert to bytes
	min_value = np.amin(data)
	scale = (np.amax(data) - min_value)/255
	output_data = np.rint(data.__iadd__(-min_value).__imul__(1/scale)).astype("uint8").tobytes()
	del data
	gc.collect()
	
	# Set up process
	command = [
		"ffmpeg",
		"-loglevel", "error",
		"-y",
		"-pix_fmt", "gray",
		"-s", "%dx%d" % (reordered_shape[1], reordered_shape[2]),
		"-r", "30",
		"-c:v", "rawvideo",
		"-f", "rawvideo",
		"-i", "-", # Input from stdin
		"-crf", str(crf),
		"-c:v", "libx264",
		"-preset", preset,
		"-f", "matroska",
		"-threads", str(threads),
		"-" # Output to stdout
	]
	buffer_size = 2**24
	p = sp.Popen(command, stdin=sp.PIPE, stdout=sp.PIPE, bufsize=buffer_size)
		
	# Set up output handler
	video_data = bytearray()
	def read_output():
		video_data.extend(p.stdout.read())
	
	# Wait for output while sending input
	thread = threading.Thread(target=read_output)
	thread.start()
	p.stdin.write(output_data)
	p.stdin.close()
	thread.join()
	
	# Join metadata and video data into single file for decompression
	with open(output_file, "wb") as f:
		f.write(struct.pack("B", len(dtype.encode()))) # dtype string length (8-bit unsigned)
		f.write(dtype.encode())
		f.write(struct.pack("B", modes)) # Modes (8-bit unsigned)
		f.write(struct.pack("B", time_mode_index)) # Time mode index (8-bit unsigned)
		for size in sizes:
			f.write(struct.pack("I", size)) # Sizes (32-bit unsigned)
		f.write(struct.pack("d", min_value)) # Offset (64-bit float)
		f.write(struct.pack("d", scale)) # Scale (64-bit float)
		f.write(video_data) # Write video data itself

def decompress_video(input_file, output_file, threads=1):
	
	# Decompresses using video compression
	
	# Interpret file
	with open(input_file, "rb") as f:
		dtype_str_length = struct.unpack("B", f.read(1))[0]
		dtype = f.read(dtype_str_length).decode()
		modes = struct.unpack("B", f.read(1))[0]
		time_mode_index = struct.unpack("B", f.read(1))[0]
		sizes = [0,]*modes
		for i in range(modes):
			sizes[i] = struct.unpack("I", f.read(4))[0]
		size = reduce(lambda x, y: x*y, sizes)
		min_value = struct.unpack("d", f.read(8))[0]
		scale = struct.unpack("d", f.read(8))[0]
		video_data = f.read()
	
	# Derive other metadata
	reordered_shape = [sizes[time_mode_index], sizes[1 if time_mode_index == 0 else 0], 0]
	reordered_shape[2] = size/reordered_shape[0]/reordered_shape[1]
	
	# Set up process
	command = [
		"ffmpeg",
		"-loglevel", "quiet",
		"-i", "-", # Input from stdin
		"-pix_fmt", "gray",
		"-s", "%dx%d" % (reordered_shape[1], reordered_shape[2]),
		"-r", "30",
		"-c:v", "rawvideo",
		"-f", "rawvideo",
		"-threads", str(threads),
		"-" # Output to stdout
	]
	buffer_size = 2**24
	p = sp.Popen(command, stdin=sp.PIPE, stdout=sp.PIPE, bufsize=buffer_size)
	
	# Set up output handler
	data = bytearray()
	def read_output():
		data.extend(p.stdout.read())
	
	# Wait for output while sending input
	gc.collect()
	thread = threading.Thread(target=read_output)
	thread.start()
	p.stdin.write(video_data)
	p.stdin.close()
	thread.join()
	
	# Transform data to original scale, offset and mode order
	data = np.frombuffer(bytes(data), dtype="uint8").astype("float32").__imul__(scale).__iadd__(min_value).astype(dtype)
	data = data.reshape([sizes[time_mode_index],] + sizes[:time_mode_index] + sizes[time_mode_index + 1:]).transpose(list(range(1, time_mode_index + 1)) + [0,] + list(range(time_mode_index + 1, modes)))
	
	# Write reconstructed data to output file
	data.tofile(output_file)

main()
