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 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(ILLUMINANTS_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 = [self.scale, t, t]
		self.cubes = [RegularGridInterpolator(axes, coeffs[j, :, :, :, :], bounds_error=False) for j in range(3)]

	def __call__(self, RGB):
		RGB = np.asarray(RGB, dtype=np.float) # FIXME: default dtype
		vmax = RGB.max()
		imax = RGB.argmax()
		cube = self.cubes[imax]
		chroma = RGB / (vmax + 1e-10)
		chroma = np.array([
			vmax,
			chroma[..., (imax + 2) % 3],
			chroma[..., (imax + 1) % 3]
		])
		return cube(chroma).squeeze()



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

	linear_RGB = np.array([1, 0, 0]) # A random test color
	RGB = eotf_inverse_sRGB(linear_RGB)
	XYZ = sRGB_to_XYZ(RGB)
	cc = jh(linear_RGB)

	plot_comparison(XYZ, model_sd(cc, primed=False), "Model", 123, ill_sd)