#!/usr/bin/env python2
# -*- coding: utf-8 -*-

import multiprocessing as mp, os, struct, sys, time
from StringIO import StringIO

from DisplayCAL import ICCProfile as ICCP, config, colormath, madvr, worker_base
from DisplayCAL.log import safe_print


H3D_HEADER = ("3DLT\x01\x00\x00\x00DisplayCAL\x00\x00\x00\x00\x00\x00\x00\x00"
			  "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00 \x00"
			  "\x00\x00\x00\x00\x00\x08\x00\x00\x00\x08\x00\x00\x00\x08\x00\x00"
			  "\x00\x00\x00\x00\x00\x10\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00"
			  "\x00\x00\x00\x00\x00\x00@\x00\x00\x00\x00\x00\x00\x00\x00\x00"
			  "\x06\x00\x00\x00\x06")

H3D_PARAMETERS = ("Input_Primaries %.6f %.6f %.6f %.6f %.6f %.6f %.6f %.6f\r\n"
				  "Input_Range 16 235\r\nOutput_Range 16 235\r\n\x00")


def eecolor_to_madvr(eecolor_3dlut_filename, unity=False):
	"""
	Convert eeColor 65^3 3D LUT to madVR 256^3 3D LUT using interpolation
	
	madvr 3D LUT will be written to:
	<eeColor 3D LUT filename without extension> + '.madVR256.3dlut'
	
	"""

	filename, ext = os.path.splitext(eecolor_3dlut_filename)

	if not unity:
		# Read in eeColor 3D LUT data
		safe_print("Reading in eeColor 3D LUT", '"%s"' % eecolor_3dlut_filename)
		eecolor_3dlut_data = []
		with open(eecolor_3dlut_filename, "rb") as eecolor_3dlut:
			for line in eecolor_3dlut:
				line = line.strip()
				if line.startswith("#"):
					continue
				values = [float(v) for v in line.split()]
				eecolor_3dlut_data.append(values)

		# Sort eeColor 3D LUT data
		# (fastest increasing column from right to left: B -> G -> R)
		eecolor_3dlut_data.sort()

		# Create ICC device link profile from eeColor 3D LUT data
		safe_print("Creating device link from eeColor 3D LUT data...")
		link = ICCP.ICCProfile()
		link.connectionColorSpace = "RGB"
		link.profileClass = "link"
		link.tags.desc = ICCP.TextDescriptionType()
		link.tags.desc.ASCII = os.path.basename(filename)
		link.tags.cprt = ICCP.TextType("text\0\0\0\0No copyright", "cprt")
		link.tags.A2B0 = ICCP.LUT16Type()
		link.tags.A2B0.input = link.tags.A2B0.output = [[0, 65535]] * 3
		link.tags.A2B0.matrix = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
		link.tags.A2B0.clut = []

		# Fill cLUT
		clutres = 65
		row = 0
		for a in xrange(clutres):
			for b in xrange(clutres):
				link.tags.A2B0.clut.append([])
				for c in xrange(clutres):
					values = [v * 65535 for v in eecolor_3dlut_data[row][3:]]
					link.tags.A2B0.clut[-1].append(values)
					row += 1

		# Write device link
		link.write(filename + ".icc")

	icc_device_link_to_madvr(filename + ".icc", unity=False)


def icc_device_link_to_madvr(icc_device_link_filename, unity=False):
	"""
	Convert ICC device link profile to madVR 256^3 3D LUT using interpolation
	
	madvr 3D LUT will be written to:
	<device link filename without extension> + '.madVR256.3dlut'
	
	"""
	t = time.time()

	filename, ext = os.path.splitext(icc_device_link_filename)

	h3d_params = H3D_PARAMETERS

	if filename.endswith(".HDR"):
		name = os.path.splitext(filename)[0]
		h3d_params = "Input_Transfer_Function PQ\r\nOutput_Transfer_Function PQ\r\n" + h3d_params
	elif filename.endswith(".HDR2SDR"):
		name = os.path.splitext(filename)[0]
		h3d_params = "Input_Transfer_Function PQ\r\n" + h3d_params
	else:
		name = filename

	colorspace = os.path.splitext(name)[1]
	colorspace = colorspace[1:]
	key = {"BT709": "Rec. 709",
		   "SMPTE_C": "SMPTE-C",
		   "EBU_PAL": "PAL/SECAM",
		   "BT2020": "Rec. 2020",
		   "DCI_P3": "DCI P3 D65"}.get(colorspace)
	if not key:
		if not colorspace:
			safe_print("ERROR - no target color space suffix in filename",
					   colorspace)
		else:
			safe_print("ERROR - invalid target color space suffix in filename:",
					   colorspace)
		safe_print("Possible target color spaces:",
				   "BT709, SMPTE_C, EBU_PAL, BT2020, DCI_P3")
		sys.exit(1)

	rgb_space = colormath.get_rgb_space(key)

	# Create madVR 3D LUT
	h3d_stream = StringIO(H3D_HEADER)
	h3dlut = madvr.H3DLUT(h3d_stream, check_lut_size=False)
	(rx, ry, rY), (gx, gy, gY), (bx, by, bY) = rgb_space[2:5]
	wx, wy = colormath.XYZ2xyY(*rgb_space[1])[:2]
	h3dlut.parametersData = h3d_params % (rx, ry, gx, gy, bx, by, wx, wy)
	h3dlut.write(filename + ".madVR256.3dlut")
	raw = open(filename + ".madVR256.3dlut", "r+b")
	raw.seek(h3dlut.lutFileOffset)
	# Make sure no longer needed h3DLUT instance can be garbage collected
	del h3dlut

	# Lookup 256^3 values through device link and fill madVR cLUT
	safe_print("Generating 256^3 input values...")
	clutres = 256
	clutmax = clutres - 1.0
	prevperc = -1
	if not unity:
		link = ICCP.ICCProfile(icc_device_link_filename)
		xicclu = worker_base.MP_Xicclu(link, scale=clutmax, logfile=sys.stdout,
									   output_format=("<H", 65535 / clutmax),
									   reverse=True, output_stream=raw)
	for a in xrange(clutres):
		for b in xrange(clutres):
			for c in xrange(clutres):
				if not unity:
					xicclu((a, b, c))
				else:
					# Optimize for speed
					B, G, R = chr(c), chr(b), chr(a)
					raw.write(B + B + G + G + R + R)
		perc = round(a / clutmax * 100)
		if perc > prevperc:
			sys.stdout.write("\r%i%%" % perc)
			prevperc = perc
	if not unity:
		safe_print("")
		safe_print("Looking up 256^3 input values through device link and writing madVR 3D LUT",
				   '"%s"' % (filename + ".madVR256.3dlut"))
		xicclu.exit()
		xicclu.get()
	safe_print("")
	safe_print("Finished in", time.time() - t, "seconds")
	if filename.endswith(".HDR"):
		safe_print("Gamut (rx ry gx gy bx by wx wy):",
				   "%.4f %.4f %.4f %.4f %.4f %.4f %.4f %.4f" %
				   (rx, ry, gx, gy, bx, by, wx, wy))


if __name__ == "__main__":
	mp.freeze_support()
	if len(sys.argv[1:]) and not "--help" in sys.argv[1:]:
		config.initcfg()
		for arg in sys.argv[1:]:
			if arg.lower().endswith(".txt"):
				fn = eecolor_to_madvr
			elif arg.lower().endswith(".icc") or arg.lower().endswith(".icm"):
				fn = icc_device_link_to_madvr
			else:
				continue
			fn(arg, unity="--unity" in sys.argv)
			safe_print("")
	else:
		safe_print("Convert eeColor 65^3 to madVR 256^3 3D LUT (video levels in, video levels out)")
		safe_print("Author: Florian Hoech, licensed under the GPL version 3")
		safe_print("Usage: %s [--unity] <name>.<target color space>.txt" % os.path.basename(sys.argv[0]))
		safe_print("Possible target color spaces: BT709, SMPTE_C, EBU_PAL, BT2020, DCI_P3")
		safe_print("Write output to <name>.<target color space>.madVR256.3dlut")
		safe_print("")
		safe_print("Examples:")
		safe_print("  Convert SDR BT709 eeColor 3D LUT")
		safe_print("  %s eeColor.BT709.txt"  % os.path.basename(sys.argv[0]))
		safe_print("")
		safe_print("  Convert HDR BT2020 eeColor 3D LUT")
		safe_print("  %s eeColor.BT2020.HDR.txt"  % os.path.basename(sys.argv[0]))
		safe_print("")
		safe_print("  Convert HDR to SDR BT2020 eeColor 3D LUT")
		safe_print("  %s eeColor.BT2020.HDR2SDR.txt"  % os.path.basename(sys.argv[0]))
	raw_input("Press RETURN to exit")
