import argparse, numpy as np, sys
import scipy.interpolate
import PIL.Image
import crl

parser = argparse.ArgumentParser()
parser.add_argument("colorspace", type=str, nargs=1)
parser.add_argument("output", type=str, nargs=1)
parser.add_argument("--resolution", type=int, default=500) 
args = parser.parse_args()

try:
	CS, polar_x0, polar_y0 = {
		"xy": (crl.xyY, 0.25, 0.40),
		"uv": (crl.uvY, 0.25, 0.20),
		"uv76": (crl.uvY76, 0.25, 0.20),
		"UVstar": (crl.UVstarNormed, 0, 0),
		"Labstar": (crl.LabstarD65, 0, 0)
	}[args.colorspace[0]]
except KeyError:
	print("error: '%s' is not a supported color space" % args.colorspace[0])
	parser.print_help()
	exit(1)



def diagram_xy():
	def polar(u, v):
		du = u - polar_x0
		dv = v - polar_y0
		r = np.sqrt(du ** 2 + dv ** 2)
		phi = np.arctan2(du, dv)
		return r, phi

	# First, the spectral color locus
	locus = []
	wvls = list(np.linspace(380, 780, 400))

	for i, wvl in enumerate(wvls):
		x, y, _ = CS(crl.PointSpectrum(wvl))

		if i == 0:
			x0, y0 = x, y
		elif i == len(wvls) - 1:
			x1, y1 = x, y

		r, phi = polar(x, y)
		locus.append([phi, r])

	# Then, the line of purples (interpolate from x0,y0 to x1,y1) 
	for t in np.linspace(0, 1, 40):
		s = 1 - t
		x = s * x0 + t * x1
		y = s * y0 + t * y1
		r, phi = polar(x, y)
		locus.append([phi, r])


	locus = sorted(np.array(locus), key=lambda v: v[0])
	locus = np.array(locus)

	cull_func = scipy.interpolate.interp1d(
		locus[:, 0], locus[:, 1], "linear", fill_value="extrapolate")


	bounds = crl.plot_data[CS]["diagram_bounds"]

	im = np.zeros([args.resolution, args.resolution, 4])
	for i in range(im.shape[0]):
		for j in range(im.shape[1]):
			x = i / (im.shape[0] - 1) * (bounds[1] - bounds[0]) + bounds[0]
			y = (1 - j / (im.shape[1] - 1)) * (bounds[3] - bounds[2]) + bounds[2]

			r, phi = polar(x, y)
			if r > cull_func(phi):
				continue
			im[j, i] = np.array([*crl.sRGB(CS(x, y, 1)), 1])

	return im

def diagram_ab():
	bounds = crl.plot_data[CS]["diagram_bounds"]

	im = np.zeros([args.resolution, args.resolution, 4])
	for i in range(im.shape[0]):
		for j in range(im.shape[1]):
			x = i / (im.shape[0] - 1) * (bounds[1] - bounds[0]) + bounds[0]
			y = (1 - j / (im.shape[1] - 1)) * (bounds[3] - bounds[2]) + bounds[2]
			im[j, i] = np.array([*crl.sRGB100(CS(75, x, y)), 1])

	return im


if CS == crl.LabstarD65:
	im = diagram_ab()
else:
	im = diagram_xy()

np.clip(im, 0, 1, out=im)
image = PIL.Image.fromarray((im * 255).astype("uint8"))
image.save(args.output[0])