import numpy as np, os, multiprocessing
from colour import *
from colour.models import eotf_inverse_sRGB
from jakob_hanika import jakob_hanika, model_sd



illuminant = "D65"
ill_xy = ILLUMINANTS["CIE 1931 2 Degree Standard Observer"][illuminant]
ill_sd = SpectralDistribution(ILLUMINANT_SDS[illuminant]) # Note: changed in dev (to ILLUMINANT_SD)

# Resolution of the discretization cubes
# Both should be 64, but this code is a bit too slow at the moment
CHROMA_STEPS = 8
LIGHTNESS_STEPS = 8



# Solve for a specific RGB color
def optimize_RGB(linear_RGB, ccp0):
	RGB = eotf_inverse_sRGB(linear_RGB)
	target = sRGB_to_XYZ(RGB, ill_xy)
	ccp, error = jakob_hanika(target, ill_sd, ill_xy, ccp0, verbose=False)
	if error > 0.1:
		print("WARNING: convergence failed with L=%s, starting from %s" \
		      % (linear_RGB, ccp0))

	return ccp, error

# Solve for all lightness values of a fully saturated RGB color
def optimize_chromaticity(argv):
	linear_RGB, dirname = argv

	# alpha's aren't spaced equally (see the article)
	def smoothstep(x):
		return x**2 * (3 - 2 * x)

	ii = np.arange(0, LIGHTNESS_STEPS)
	aa = smoothstep(smoothstep(ii / LIGHTNESS_STEPS))
	ccps = np.empty((LIGHTNESS_STEPS, 3)) # FIXME: dtype?

	def f(alpha, ccp0):
		path = "%s/%.3f" % (dirname, alpha)
		print("%s..." % path)

		ccp, error = optimize_RGB(linear_RGB * alpha, ccp0)

		try:
			os.makedirs(os.path.dirname(path))
		except:
			pass
		with open(path, "w") as fd:
			fd.write("%s, %s" % (ccp, error))
			if error > 0.1:
				fd.write(" FAILED")
			fd.write("\n")

		return ccp, error


	i_mid = LIGHTNESS_STEPS // 2
	ccps[i_mid, :], _ = f(aa[i_mid], (0, 0, 0))

	ccp0 = ccps[i_mid, :]
	for i in range(i_mid + 1, LIGHTNESS_STEPS):
		ccps[i, :], _ = f(aa[i], ccp0)

	ccp0 = ccps[i_mid, :]
	for i in reversed(range(0, i_mid)):
		ccps[i, :], _ = f(aa[i], ccp0)



# This demo discretizes the sRGB space to three 8x8x8 cubes and tries to find
# the model parameters for all the colors. Coefficients are saved to a file
# somewhere in the 'out' directory (each color gets its own file).
# FIXME: This dumpster fire of an output format is to be replaced ASAP.
# FIXME: This takes a good while (on the order of an hour).
if __name__ == "__main__":
	args = []

	done = 0
	total = 0
	for A in np.linspace(0, 1, CHROMA_STEPS):
		for B in np.linspace(0, 1, CHROMA_STEPS):
			for RGB in [np.array([1, A, B]),
			            np.array([A, 1, B]),
			            np.array([A, B, 1])]:

				dirname = "out/R%.3fG%.3fB%.3f" % (*RGB,)
				total += 1

				if os.path.exists(dirname) and len(os.listdir(dirname)) == LIGHTNESS_STEPS:
					done += 1
					continue

				args.append((RGB, dirname))

	print("%d/%d already done" % (done, total))

	pool = multiprocessing.Pool()
	pool.map(optimize_chromaticity, args)