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", "placebo")

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 = sys.argv[arg_index] # Stored as string, is equal to a number in between 0 and 51 or equal to "lossless"
		arg_index += 1
		dtype = sys.argv[arg_index]
		arg_index += 1
		time_mode_index = int(sys.argv[arg_index])
		arg_index += 1
		x_mode_index = int(sys.argv[arg_index])
		arg_index += 1
		y_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 = max(1, min(16, int(sys.argv[arg_index]))) # x265 seems to have problems with more than 16 threads
		arg_index += 1
	
	# Perform compression/decompression
	if mode == COMPRESS:
		compress_video(input_file, output_file, crf, dtype, time_mode_index, x_mode_index, y_mode_index, sizes, preset=preset, threads=threads)
	else:
		decompress_video(input_file, output_file, threads=threads)

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

quantization_dtype = np.dtype("uint8")
pixel_format = "gray" # VSC only supports gray12le on x265 if you DON'T load a new FFmpeg version after logging in

def compress_video(input_file, output_file, crf, dtype, time_mode_index, x_mode_index, y_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
	perm = list(range(modes))
	perm.remove(time_mode_index)
	perm.remove(x_mode_index)
	perm.remove(y_mode_index)
	perm += [time_mode_index, x_mode_index, y_mode_index]
	data = np.transpose(data, axes=perm)
	data = data.reshape((-1, data.shape[-2], data.shape[-1]))
	
	# Check if width and height are large enough
	min_allowed_side = 64
	if min(data.shape[-2:]) < min_allowed_side:
		"""# Pad data such that width and height are at least 64
		# Required for x265 encoding
		# However, this step shouldn't be necessary for all datasets considered in the paper
		data = np.pad(data, ((0, 0), (0, max(0, min_allowed_side - data.shape[1])), (0, max(0, min_allowed_side - data.shape[2]))), mode="constant")"""
		raise Exception("Dataset has problematic sizes, min(width, height) is %s, minimum is %s"%(min(data.shape[-2:]), min_allowed_side))
	
	# Quantize data and convert to bytes
	min_value = np.amin(data)
	scale = (np.amax(data) - min_value)/(2**(8*quantization_dtype.itemsize) - 1)
	data = np.rint(data.__iadd__(-min_value).__imul__(1/scale)).astype(quantization_dtype)
	video_shape = data.shape
	data = data.tobytes()
	gc.collect()
	
	# Set up process
	command = [
		"ffmpeg",
		"-loglevel", "error",
		"-y",
		"-pix_fmt", pixel_format,
		"-s", "%dx%d" % (video_shape[1], video_shape[2]),
		"-r", "30",
		"-c:v", "rawvideo",
		"-f", "rawvideo",
		"-i", "-", # Input from stdin
		"-c:v", "libx265",
		"-x265-params", ("lossless=1" if crf == "lossless" else "crf=%s"%crf),
		"-preset", preset,
		"-tune", "psnr",
		"-f", "matroska",
		"-threads", str(threads),
		"-" # Output to stdout
	]
	print(" ".join(command))
	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(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)
		for mode in perm:
			f.write(struct.pack("B", mode)) # Mode permutation (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]
		perm = [1,]*modes
		for i in range(modes):
			perm[i] = 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()
	
	transposed_sizes = [sizes[perm[i]] for i in range(modes)]
	inv_perm = [perm.index(i) for i in range(modes)]
	
	# Set up process
	command = [
		"ffmpeg",
		"-loglevel", "quiet",
		"-i", "-", # Input from stdin
		"-pix_fmt", pixel_format,
		"-s", "%dx%d" % (transposed_sizes[-2], transposed_sizes[-1]),
		"-r", "30",
		"-c:v", "rawvideo",
		"-f", "rawvideo",
		"-threads", str(threads),
		"-" # Output to stdout
	]
	print(" ".join(command))
	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=quantization_dtype).reshape(transposed_sizes) # Reshape to padded video tensor
	data = data.transpose(inv_perm) # Transpose to original mode order
	data = data.astype("float32").__imul__(scale).__iadd__(min_value).astype(dtype)
	
	# Write reconstructed data to output file
	data.tofile(output_file)

main()
