import numpy as np, PIL, os
from matplotlib import pyplot as plt, rcParams
from crl.color import *

def assets(path):
	return os.path.join(os.path.dirname(__file__), "assets", path)

plot_data = {
	xyY: {
		"diagram_path": assets("diagram_xy.png"),
		"diagram_bounds": [0, 0.75, 0, 0.85],
		"xlabel": "$x$",
		"ylabel": "$y$"
	},
	uvY: {
		"diagram_path": assets("diagram_uv.png"),
		"diagram_bounds": [0, 0.7, 0, 0.4],
		"xlabel": "$u$",
		"ylabel": "$v$"
	},
	uvY76: {
		"diagram_path": assets("diagram_uv76.png"),
		"diagram_bounds": [0, 0.65, 0, 0.63],
		"xlabel": "$u'$",
		"ylabel": "$v'$"
	},
	UVstarNormedD65: {
		"diagram_path": assets("diagram_UVWstar.png"),
		"diagram_bounds": [-2.6, 5.6, -4, 1.1],
		"xlabel": "$\\frac{U^*}{W^*}$",
		"ylabel": "$\\frac{V^*}{W^*}$"
	},
	LabstarD65: {
		"diagram_path": assets("diagram_Labstar.png"),
		"diagram_bounds": [-50, 50, -50, 50],
		"xlabel": "$a^*$",
		"ylabel": "$b^*$"
	}
}

def draw_isotherm(ax, CS, x, y, T, plot_space_scale, labels, label_scale):
	u0, v0, _ = uvY(BlackBodySpectrum(T))
	ut, vt, _ = uvY(BlackBodySpectrum(T + 100)) # t stands for 'tangent'
	duv = np.array([ut - u0, vt - v0])
	du, dv = np.matmul([[0, -1], [1, 0]], duv) # rotate 90 degrees

	annotate = False
	X, Y = [], []
	for sign in [-1, 1]:
		xp, yp, _ =  CS.from_XYZ(uvY(
			u0 + sign * du,
			v0 + sign * dv, 1).to_XYZ())

		dxy = np.array([xp - x, yp - y])

		if T % 1000 == 500:
			size = 0.006
		elif T <= 5000 or T == 10000:
			size = 0.02
			annotate = True
		else:
			size = 0.012

		dxy *= plot_space_scale * size / np.linalg.norm(dxy)

		X.append(x + dxy[0])
		Y.append(y + dxy[1])

	ax.plot(X, Y, "k-", linewidth=1)

	if labels and annotate:
		save = rcParams['font.size']
		rcParams['font.size'] *= label_scale
		ax.annotate("%d" % T, xy=(X[1], Y[1]),
			    horizontalalignment="left",
			    verticalalignment="top")
		rcParams['font.size'] = save


def chromaticity_diagram(CS, label_scale=1/1.5, ax=plt,
                         pures=True, pures_labels=True,
                         locus=True, isotherms=False,
                         isotherm_labels = False, isotherm_label_scale=1/1.7):
	if CS not in plot_data:
		raise ValueError("color triangle plots for %s are not supported" % CS.full_name)

	if CS == LabstarD65:
		# These won't work in L*a*b*
		pures = False
		locus = False

	if not hasattr(CS, "diagram_image"):
		img = PIL.Image.open(plot_data[CS]["diagram_path"])
		img.load()
		CS.diagram_image = np.asarray(img, dtype="float") / 255

	ax.imshow(CS.diagram_image, extent=plot_data[CS]["diagram_bounds"], \
	          interpolation="bilinear")

	save = rcParams['font.size']
	rcParams['font.size'] *= label_scale

	plot_space_scale = 10 if CS == UVstarNormed else 1

	if pures:
		wvls = np.linspace(380, 780, 300)
		outline = np.array([[*CS(PointSpectrum(wvl))] for wvl in wvls])
		ax.plot(outline[:, 0], outline[:, 1], 'k')

		def tick(wvl, l, label=False):
			x, y, _ = CS(PointSpectrum(wvl))
			xe, ye, _ = CS(PointSpectrum(wvl + 5))
			l *= plot_space_scale

			dx = np.array([xe - x, ye - y])
			dx *= l / np.linalg.norm(dx)
			dx = np.matmul([[0, -1], [1, 0]], dx)
		
			ax.plot((x, x + dx[0]), (y, y + dx[1]), 'k-')
		
			if pures_labels and label:
				ax.annotate("%d" % wvl, (x, y),
					     xytext=np.array([x, y]) + dx * 2,
					     horizontalalignment="center",
					     verticalalignment="center")

		L = [380, 460, 470, 480, 490, 500, 520, 540, 560, 580, 600, 620]

		for wvl in range(400, 680, 5):
			if wvl not in L:
				if wvl % 10 == 0:
					tick(wvl, 0.012)
				else:
					tick(wvl, 0.007)
		for wvl in L:
			tick(wvl, 0.02, True)

	if locus:
		x = np.zeros(50)
		y = np.zeros(np.size(x))
		for i in range(np.size(x)):
			T = 1000 * 20 ** (i / (np.size(x) - 1))
			x[i], y[i], _ = CS(BlackBodySpectrum(T))
		ax.plot(x, y, "k-", linewidth=1)

		if type(isotherms) == list:
			Ts = isotherms
		else:
			Ts = range(2000, 10001, 500)
		for T in Ts:
			x, y, _ = CS(BlackBodySpectrum(T))

			if isotherms:
				draw_isotherm(ax, CS, x, y, T, plot_space_scale,
				              isotherm_labels, isotherm_label_scale)
			else:
				ax.plot(x, y, "ko", markersize=2)

	rcParams['font.size'] = save

	ax.xlabel(plot_data[CS]["xlabel"])
	ax.ylabel(plot_data[CS]["ylabel"])