import numpy as np, scipy.optimize as optimize
from colour import *
from colour.difference import delta_E_CIE1976
from colour.colorimetry import *



# The same wavelength grid is used throughout
wvl = np.arange(360, 830, 5)
wvlp = (wvl - 360) / (830 - 360)

# This is the model of spectral reflectivity described in the article.
def model(wvlp, ccp):
	yy = ccp[0] * wvlp**2 + ccp[1] * wvlp + ccp[2]
	return 1 / 2 + yy / (2 * np.sqrt(1 + yy ** 2))

# Create a SpectralDistribution using given coefficients
def model_sd(ccp, primed=True):
	# FIXME: don't hardcode the wavelength grid; there should be a way
	#        of creating a SpectralDistribution from the function alone
	grid = wvlp if primed else wvl
	return SpectralDistribution(model(grid, ccp), wvl, name="Model")

# The goal is to minimize the color difference between a given distrbution
# and the one computed from the model above.
def error_function(ccp, target, ill_sd, ill_xy):
	# One thing they didn't mention in the text is that their polynomial
	# is nondimensionalized (mapped from 360--800 nm to 0--1).
	sd = model_sd(ccp)
	Lab = XYZ_to_Lab(sd_to_XYZ(sd, illuminant=ill_sd) / 100, ill_xy)
	return delta_E_CIE1976(target, Lab)



# Finds the parameters for Jakob and Hanika's model
def jakob_hanika(target_XYZ, ill_sd, ill_xy, ccp0=(0, 0, 0), verbose=True, try_hard=True):
	def do_optimize(XYZ, ccp0):
		Lab = XYZ_to_Lab(XYZ, ill_xy)
		opt = optimize.minimize(
			error_function, ccp0, (Lab, ill_sd, ill_xy),
			method="Nelder-Mead", options={"disp": verbose}
		)
		if verbose:
			print(opt)
		return opt

	# A special case that's hard to solve numerically
	if np.allclose(target_XYZ, [0, 0, 0]):
		return np.array([0, 0, -1e+9]), 0 # FIXME: dtype?

	if verbose:
		print("Trying the target directly, XYZ=%s" % target_XYZ)
	opt = do_optimize(target_XYZ, ccp0)
	if opt.fun < 0.1 or not try_hard:
		return opt.x, opt.fun

	good_XYZ = (1/3, 1/3, 1/3)
	good_ccp = (2.1276356, -1.07293026, -0.29583292) # FIXME: valid only for D65

	divisions = 3
	while divisions < 10:
		if verbose:
			print("Trying with %d divisions" % divisions)

		keep_divisions = False
		ref_XYZ = good_XYZ
		ref_ccp = good_ccp

		ccp0 = ref_ccp
		for i in range(1, divisions):
			intermediate_XYZ = ref_XYZ + (target_XYZ - ref_XYZ) * i / (divisions - 1)
			if verbose:
				print("Intermediate step %d/%d, XYZ=%s with ccp0=%s" %
				      (i + 1, divisions, intermediate_XYZ, ccp0))
			opt = do_optimize(intermediate_XYZ, ccp0)
			if opt.fun > 0.1:
				if verbose:
					print("WARNING: intermediate optimization failed")
				break
			else:
				good_XYZ = intermediate_XYZ
				good_ccp = opt.x
				keep_divisions = True

			ccp0 = opt.x
		else:
			return opt.x, opt.fun

		if not keep_divisions:
			divisions += 2

	raise Exception("optimization failed for target_XYZ=%s, ccp0=%s" \
	                % (target_XYZ, ccp0))