import numpy as np, struct
from scipy.interpolate import RegularGridInterpolator
from colour import *
from colour.plotting import *
from colour.models import eotf_inverse_sRGB
from colour.difference import delta_E_CIE1976
from jakob_hanika import jakob_hanika, model_sd
from test_colorchecker import plot_comparison



illuminant = "D65"
ill_xy = ILLUMINANTS["CIE 1931 2 Degree Standard Observer"][illuminant]
ill_sd = SpectralDistribution(ILLUMINANT_SDS[illuminant])



class JakobHanikaInterpolator:
	def __init__(self):
		pass

	def read_file(self, path):
		with open(path, "rb") as fd:
			if fd.read(4).decode("ISO-8859-1") != "SPEC":
				raise ValueError("bad magic number")

			self.res = struct.unpack("i", fd.read(4))[0]
			self.scale = np.fromfile(fd, count=self.res, dtype=np.float32) # FIXME: default dtype
			coeffs = np.fromfile(fd, count=3*self.res**3*3, dtype=np.float32) # FIXME: default dtype
			coeffs = coeffs.reshape(3, self.res, self.res, self.res, 3)

		t = np.linspace(0, 1, self.res)
		axes = [[0, 1, 2], self.scale, t, t]
		self.cubes = RegularGridInterpolator(axes, coeffs[:, :, :, :, :], bounds_error=False)

	def __call__(self, RGB):
		RGB = np.asarray(RGB, dtype=np.float) # FIXME: default dtype
		vmax = np.max(RGB, axis=-1)
		imax = np.argmax(RGB, axis=-1)
		chroma = RGB / (np.expand_dims(vmax, -1) + 1e-10) # Avoid division by zero
		vmax = np.max(RGB, axis=-1)
		v2 = np.take_along_axis(chroma, np.expand_dims((imax + 2) % 3, axis=-1), axis=-1).squeeze(axis=-1)
		v3 = np.take_along_axis(chroma, np.expand_dims((imax + 1) % 3, axis=-1), axis=-1).squeeze(axis=-1)
		coords = np.stack([imax, vmax, v2, v3], axis=-1)
		return self.cubes(coords).squeeze()



if __name__ == "__main__":
	jh = JakobHanikaInterpolator()
	jh.read_file("data/srgb.coeff")

	RGBs = np.random.random((7, 6, 5, 4, 3))
	ccs = jh(RGBs)

	RGB = eotf_inverse_sRGB(RGBs[0, 0, 0, 0, :])
	XYZ = sRGB_to_XYZ(RGB)
	Lab = XYZ_to_Lab(XYZ, ill_xy)
	cc = ccs[0, 0, 0, 0, :]

	matched_sd = model_sd(cc, primed=False)
	matched_XYZ = sd_to_XYZ(matched_sd, illuminant=ill_sd)
	matched_Lab = XYZ_to_Lab(matched_XYZ, ill_xy)
	error = delta_E_CIE1976(Lab, matched_Lab)

	print(matched_sd)

	plot_comparison(XYZ, matched_sd, "Model", error, ill_sd, ill_xy)