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

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

# This is the model of spectral reflectivity described in the article.
def model(wvl, cc):
	# One thing they didn't mention in the text is that their polynomial
	# is nondimensionalized (mapped from 360--800 nm to 0--1).
	def remap(x):
		return (x - 360) / (830 - 360)

	x = cc[0] * remap(wvl) ** 2 + cc[1] * remap(wvl) + cc[2]
	return 1 / 2 + x / (2 * np.sqrt(1 + x ** 2))

# The goal is to minimize the color difference between a given distrbution
# and the one computed from the model above.
def error_function(cc, target, illuminant):
	ev = model(wvl, cc)
	sd = SpectralDistribution(ev, wvl)
	Lab = XYZ_to_Lab(sd_to_XYZ(sd), illuminant)
	return delta_E_CIE1976(target, Lab)



# This callback to scipy.optimize.basinhopping tells the solver to stop once
# the error is small enough. The threshold was chosen arbitrarily, as a small
# fraction of the JND (about 2.3 with this metric).
def cb_basinhopping(x, f, accept):
	return f < 0.1

# Finds coefficients for Jakob and Hanika's function
def jakob_hanika(target, illuminant):
	# First, a conventional solver is called. For 'yellow green' this
	# actually fails: gets stuck at a local minimum that's far away
	# from the global one.
	# FIXME: better parameters, a better x0, a better method?
	# FIXME: stop iterating as soon as delta E is negligible 
	opt = optimize.minimize(
		error_function, (0, 0, 0), (target, illuminant),
		method="L-BFGS-B", options={"disp": True, "ftol": 1e-5}
	)
	print(opt)

	error = error_function(opt.x, target, illuminant)
	print("Delta E is %g" % error)

	# Basin hopping is far more likely to find the actual minimum we're
	# looking for, but it's extremely slow in comparison.
	if error > 0.1:
		print("Error too large, trying global optimization")
		opt = optimize.basinhopping(
			lambda cc: error_function(cc, target, illuminant),
			(0, 0, 0), disp=True, callback=cb_basinhopping
		)
		print(opt)
		error = error_function(opt.x, target)
		print("Global delta E is %g" % error)

	matched_sd = SpectralDistribution(model(wvl, opt.x), wvl, name="Model")
	return opt.x, matched_sd, error