import numpy as np, scipy.interpolate, scipy.optimize
import crl.tables as tables
import crl.tables2 as tables2

class Spectrum:
	def __init__(self, X, Y):
		self.X = X
		self.Y = Y
		self.Yi = scipy.interpolate.interp1d(self.X, self.Y, \
		          "linear", fill_value=0, bounds_error=False)

	def dot(A, B, Xmin=380, Xmax=780, Xstep=0.01):
		if not isinstance(B, Spectrum):
			raise TypeError("B is an instance of %s, not Spectrum" \
			      % type(B))

		if isinstance(A, PointSpectrum):
			if isinstance(B, PointSpectrum):
				raise TypeError("dotting two instances of PointSpectum isn't supported")
				# ^ could be, but who cares?
			return B.dot(A)

		if isinstance(B, PointSpectrum):
			return A.Yi(B.x_0) * B.S

		X = np.arange(Xmin, Xmax, Xstep)
		dX = X[1] - X[0]
		return dX * np.sum(A.Yi(X) * B.Yi(X))

	def normed(self, norm=100):
		return NormedSpectrum(self, norm)

	def __str__(self):
		return "a tabulated Spectrum"

class NormedSpectrum(Spectrum):
	def __init__(self, spec, norm=100):
		self.spec = spec

		_, Y, _ = XYZ(self.spec)
		self.factor = norm / Y

	def Yi(self, wvl):
		return self.factor * self.spec.Yi(wvl)

	def __repr__(self):
		return "NormedSpectrum(" + repr(self.spec) + ")"

	def __str__(self):
		return "normed  " + str(self.spec)

class PointSpectrum(Spectrum):
	# the "S" stands for "sum" (being the integral of the Dirac delta)
	def __init__(self, x_0, S=1):
		self.x_0 = x_0
		self.S = S

		# a crappy hack for easy plotting
		self.X = np.array([380, x_0 - 0.001, x_0, x_0 + 0.001, 780])
		self.Y = np.array([0, 0, S, 0, 0])

	def __repr__(self):
		return "PointSpectrum(x_0=%g, S=%g)" % (self.x_0, self.S)

	def __str__(self):
		return self.__repr__()

class BlackBodySpectrum(Spectrum):
	def __init__(self, T):
		self.T = T

	def Yi(self, wvl):
		h = 6.62607015e-34 # kg m^2 / s
		c = 299792458 # m / s
		k_B = 1.3806485e-23 # kg m^2 / K s^2
		return 2 * h * c ** 2 / ((wvl / 1e9) ** 5) \
		       / (np.exp(h * c / (wvl / 1e9 * k_B * self.T)) - 1) \
		       * 1e-9 # kg / (nm s^3)

	def __repr__(self):
		return "BlackBodySpectrum(T=%g)" % self.T

	def __str__(self):
		return repr(self)

# Schanda, page 41
class CIEDaylightSpectrum(Spectrum):
	S0 = Spectrum(tables.D_eigenvectors[:, 0], tables.D_eigenvectors[:, 1])
	S1 = Spectrum(tables.D_eigenvectors[:, 0], tables.D_eigenvectors[:, 2])
	S2 = Spectrum(tables.D_eigenvectors[:, 0], tables.D_eigenvectors[:, 3])

	def __init__(self, T):
		if T < 4000 or T > 25000:
			raise ValueError("D illuminants are not defined for CCTs outside the 4000 K-25000 K range")
		elif T <= 7000:
			x = -4.6070e9 / T**3 \
			    + 2.9678e6 / T**2 \
			    + 0.09911e3 / T \
			    + 0.244063
		else:
			x = -2.0064e9 / T**3 \
			    + 1.9018e6 / T**2 \
			    + 0.0247483 / T \
			    + 0.237040

		y = -3 * x**2 + 2.870 * x - 0.275

		self.M1 = (-1.3515 - 1.7703 * x + 5.9114 * y) \
		          / (0.0241 + 0.2562 * x - 0.7341 * y)
		self.M2 = (0.0300 - 31.4424 * x + 30.0717 * y) \
		          / (0.0241 + 0.2562 * x - 0.7341 * y)

		self.M1 = round(self.M1, 3)
		self.M2 = round(self.M2, 3)

	def Yi(self, wvl):
		return self.S0.Yi(wvl) \
		       + self.M1 * self.S1.Yi(wvl) \
		       + self.M2 * self.S2.Yi(wvl)

class CombinedSpectrum(Spectrum):
	def __init__(self, A, B):
		self.A = A
		self.B = B

	def Yi(self, X):
		return self.A.Yi(X) * self.B.Yi(X)

	def __repr__(self):
		return "CombinedSpectrum(%s, %s)" % (repr(self.A), repr(self.B))

def combine_spectra(A, B):
	if isinstance(B, PointSpectrum):
		A, B = B, A

	if isinstance(A, PointSpectrum):
		if isinstance(B, PointSpectrum):
			# technically this would make sense but would be pretty useless
			raise TypeError("can't combine two instances of PointSpectrum")

		return PointSpectrum(A.x_0, A.S * B.Ef(A.x_0))

	return CombinedSpectrum(A, B)

#
# Tables, standard observers, basic colors
#

class Illuminants:
	# 0 l,nm
	# 1 Standard Illuminant A
	# 2 Standard Illuminant D65
	# 3 Illuminant C
	# 4 IlluminantD50
	# 5 IlluminantD55
	# 6 IlluminantD75
	A = Spectrum(tables.illuminants[:, 0], tables.illuminants[:, 1])
	D50 = Spectrum(tables.illuminants[:, 0], tables.illuminants[:, 4])
	D55 = Spectrum(tables.illuminants[:, 0], tables.illuminants[:, 5])
	D65 = Spectrum(tables.illuminants[:, 0], tables.illuminants[:, 2])
	D75 = Spectrum(tables.illuminants[:, 0], tables.illuminants[:, 6])

	_wvl = np.linspace(380, 780, 801)
	_E = np.ones(_wvl.size)
	E = Spectrum(_wvl, _E)

	_E = np.zeros(_wvl.size)
	_E[120] = 3 # 440 nm
	_E[376] = 5 # 568 nm
	RedmanTest1 = Spectrum(_wvl, _E) # 6500K, CRI -12

class Observer:
	def __init__(self, name, wvl, cmf_x, cmf_y, cmf_z):
		self.name = name
		self.cmf_x = Spectrum(wvl, cmf_x)
		self.cmf_y = Spectrum(wvl, cmf_y)
		self.cmf_z = Spectrum(wvl, cmf_z)

	def __str__(self):
		return self.name

Observer_2deg = Observer("2-deg Standard Observer",
                         *(tables.cmf[:, i] for i in range(4)))
Observer_10deg = Observer("10-deg Standard Observer",
                          *(tables.cmf_1964[:, i] for i in range(4)))

class Color:
	full_name = "Color"
	observer = Observer_2deg

	# Call either:
	#   __init__(self, x, y, z)
	#   __init__(self, color)
	#   __init__(self, spectrum)
	def __init__(self, *args, **kwargs):
		observer = kwargs.get("observer", Observer_2deg)

		if len(args) == 3:
			self.array = np.array(args)
		elif len(args) == 1:
			if isinstance(args[0], Color):
				self.array = args[0].to(self).array
			elif isinstance(args[0], Spectrum):
				self.array = self.from_spectrum(args[0], observer).array
			else:
				raise TypeError("argument is not a Color or a Spectrum")
		else:
			raise TypeError("bad number of arguments")

	def __repr__(self):
		return "%s(%g, %g, %g)" % (type(self).__name__, *self)

	def __str__(self):
		return "%s (%g, %g, %g), %s" % (self.full_name, *self,
		                                self.observer)

	@classmethod
	def from_spectrum(cls, S, observer):
		return cls.from_XYZ(XYZ.from_spectrum(S, observer))

	def to(self, CS):
		return CS.from_XYZ(self.to_XYZ())

	@staticmethod
	def from_XYZ(C):
		raise NotImplementedError

	def to_XYZ(C):
		raise NotImplementedError

	def __getitem__(self, index):
		return self.array[index]

	def __setitem__(self, index, value):
		self.array[index] = value

	def __iter__(self):
		return self.array.__iter__()

	def __add__(A, B):
		return type(A)(*(A.array + B.array))

class sRGB(Color):
	full_name = "sRGB"
	matrix = np.array([
		[ 3.240479, -1.537150, -0.498353],
		[-0.969256,  1.875992,  0.041556],
		[ 0.055648, -0.204043,  1.057331]])
	matrix_inverse = np.linalg.inv(matrix)

	@classmethod
	def from_XYZ(cls, C):
		def gamma(x):
			if x <= 0.0031308:
				return 12.92 * x
			return 1.055 * x ** (1 / 2.4) - 0.055

		RGB = np.dot(cls.matrix, C.array)
		return sRGB(*[gamma(x) for x in RGB])

	def to_XYZ(self):
		def gamma_inverse(x):
			if x <= 0.040449936:
				return x / 12.92;
			return ((x + 0.055) / 1.055) ** 2.4

		RGB = [gamma_inverse(x) for x in self]
		return XYZ(*np.dot(self.matrix_inverse, RGB))

class sRGB100(sRGB):
	full_name = "sRGB (with Y=100 as white)"
	matrix = 0.01 * sRGB.matrix
	matrix_inverse = np.linalg.inv(matrix)

#
# CIE 1931
#

class XYZ(Color):
	full_name = "CIE 1931 XYZ"

	@staticmethod
	def from_XYZ(C):
		return C

	def to_XYZ(self):
		return self

	@staticmethod
	def from_spectrum(S, observer):
		X = S.dot(observer.cmf_x)
		Y = S.dot(observer.cmf_y)
		Z = S.dot(observer.cmf_z)
		return XYZ(X, Y, Z)

	def normed(self, norm=100):
		X, Y, Z = self
		return XYZ(X * norm / Y, norm, Z * norm / Y)

class xyY(Color):
	full_name = "CIE 1931 xyY"

	@staticmethod
	def from_XYZ(C):
		x = C[0] / sum(C)
		y = C[1] / sum(C)
		return xyY(x, y, C[1])

	def to_XYZ(self):
		x, y, Y = self
		try:
			X = Y / y * x
			Z = Y / y * (1 - x - y)
		except ZeroDivisionError:
			X = 0
			Z = 0
		return XYZ(X, Y, Z)

#
# CIE 1960
#

class UVW(Color):
	full_name = "CIE 1960 UCS (McAdam) UVW"

	@staticmethod
	def from_XYZ(C):
		X, Y, Z = C
		U = 2 * X / 3
		V = Y
		W = 1/2 * (-X + 3 * Y + Z)
		return UVW(U, V, W)

	def to_XYZ(self):
		U, V, W = self
		X = 3/2 * U
		Y = V
		Z = 3/2 * U - 3 * V + 2 * W
		return XYZ(X, Y, Z)

	def to_uvY(self):
		U, V, W = self
		u = U / sum(UVW)
		v = V / sum(UVW)
		return uvY(u, v, V)

class uvY(Color):
	full_name = "CIE 1960 uvY"

	@staticmethod
	def from_XYZ(C):
		X, Y, Z = C
		u = 4 * X / (X + 15 * Y + 3 * Z)
		v = 6 * Y / (X + 15 * Y + 3 * Z)
		return uvY(u, v, Y)

	def to_UVW(self):
		u, v, V = self
		try:
			U = V * u / v
			W = -(V * v + V * u - V) / v
		except ZeroDivisionError:
			U = 0
			W = 0
		return UVW(U, V, W)

	def to_XYZ(self):
		return self.to_UVW().to_XYZ()

#
# CIE 1964
#

class UVWstar(Color):
	full_name = "CIE 1964 U*V*W*"

	@classmethod
	def from_uvY(cls, C, CAT=True):
		u, v, Y = C
		un, vn, _ = cls.white_uvY if CAT else (0, 0, 0)

		Wstar = 25 * Y ** (1/3) - 17
		Ustar = 13 * Wstar * (u - un)
		Vstar = 13 * Wstar * (v - vn)

		return cls(Ustar, Vstar, Wstar)

	@classmethod
	def from_XYZ(cls, C):
		return cls.from_uvY(C.to(uvY))

	def to_uvY(self, CAT=True):
		Ustar, Vstar, Wstar = self
		un, vn, _ = self.white_uvY if CAT else (0, 0, 0)

		u = Ustar / (13 * Wstar) + un
		v = Vstar / (13 * Wstar) + vn
		Y = (Wstar**3 + 51 * Wstar**2 + 867 * Wstar + 4913) / 15625

		return uvY(u, v, Y)

	def to_XYZ(self):
		return self.to_uvY().to(XYZ)



#
# CIE 1976
#

class uvY76(Color):
	full_name = "CIE 1976 u'v'Y"

	@staticmethod
	def from_XYZ(C):
		X, Y, Z = C
		u = 4 * X / (X + 15 * Y + 3 * Z)
		v = 9 * Y / (X + 15 * Y + 3 * Z)
		return uvY76(u, v, Y)

	def to_XYZ(self):
		u, v, Y = self
		X = 9 * u / (4 * v) * Y
		Z = (12 - 3 * u - 20 * v) / (4 * v) * Y
		return XYZ(X, Y, Z)

# The magic CIE 1976 function for mapping Y to L*
def CIE1976_f(t):
	if t > (6/29) ** 3:
		return t ** (1/3)

	return 841 * t / 108 + 4/29

# The inverse of the above
def CIE1976_f_inverse(t):
	if t > 6/29:
		return t ** 3

	return (3132 * t - 432) / 24389

class Luvstar(Color):
	full_name = "CIE 1976 L*u*v*"

	@classmethod
	def from_uvY76(cls, C):
		u, v, Y = C
		un, vn, Yn = cls.white_uvY76

		Lstar = 116 * CIE1976_f(Y / Yn) - 16
		ustar = 13 * Lstar * (u - un)
		vstar = 13 * Lstar * (v - vn)

		return cls(Lstar, ustar, vstar)

	@classmethod
	def from_XYZ(cls, C):
		return cls.from_uvY76(C.to(uvY76))

	def to_uvY76(self):
		Lstar, ustar, vstar = self
		un, vn, Yn = self.white_uvY76

		u = ustar / (13 * Lstar) + un
		v = vstar / (13 * Lstar) + vn
		Y = Yn * CIE1976_f_inverse((Lstar + 16) / 116)

		return uvY76(u, v, Y)

	def to_XYZ(self):
		return self.to_uvY76().to_XYZ()

class Labstar(Color):
	full_name = "CIE 1976 L*a*b*"

	@classmethod
	def from_XYZ(cls, C):
		X, Y, Z = C
		Xn, Yn, Zn = cls.white

		Lstar = 116 * CIE1976_f(Y / Yn) - 16
		astar = 500 * (CIE1976_f(X / Xn) - CIE1976_f(Y / Yn))
		bstar = 200 * (CIE1976_f(Y / Yn) - CIE1976_f(Z / Zn))

		return cls(Lstar, astar, bstar)

	def to_XYZ(self):
		Lstar, astar, bstar = self
		Xn, Yn, Zn = self.white

		X = Xn * CIE1976_f_inverse(astar / 500 + (Lstar + 16) / 116)
		Y = Yn * CIE1976_f_inverse((Lstar + 16) / 116)
		Z = Zn * CIE1976_f_inverse(-bstar / 200 + (Lstar + 16) / 116)

		return XYZ(X, Y, Z)

#
# Some useful predefined color spaces
# 

class UVWstarD65(UVWstar):
	white_uvY = uvY(Illuminants.D65)

class LuvstarD65(Luvstar):
	white_uvY76 = uvY76(Illuminants.D65)

class LabstarD65(Labstar):
	white = XYZ(Illuminants.D65).normed()

class UVstarNormed(Color):
	full_name = "CIE 1964 U*V*W*, normed (plots only!)"

	@staticmethod
	def from_UVWstar(C):
		if not issubclass(type(C), UVWstar):
			raise TypeError(f"{type(C)} is not a subclass of UVWstar")

		Ustar, Vstar, Wstar = C

		Unorm = Ustar / Wstar
		Vnorm = Vstar / Wstar

		return UVstarNormed(Unorm, Vnorm, Wstar)

	@staticmethod
	def from_XYZ(C):
		raise NotImplementedError("use the from_UVWstar method directly")

	def to_UVWstar(self):
		Unorm, Vnorm, Wstar = self

		Ustar = Unorm * Wstar
		Vstar = Vnorm * Wstar

		return UVWstarD65(Ustar, Vstar, Wstar)

	def to_XYZ(self):
		return self.to_UVWstar().to_XYZ()

class UVstarNormedD65(UVstarNormed):
	@staticmethod
	def from_XYZ(C):
		return UVstarNormedD65.from_UVWstar(UVWstarD65(C))

#
# CCT
#

def cct(C):
	u, v, _ = uvY(C)
	uv = np.array([u, v])

	uv_planck = scipy.interpolate.interp1d(
		tables2.locus_uv[:, 0], tables2.locus_uv[:, 1:3], axis=0,
		bounds_error=False, fill_value="extrapolate")

	def delta(T):
		return np.linalg.norm(uv_planck(T) - uv)

	res = scipy.optimize.minimize_scalar(delta, bounds=[1000, 20000])
	return res.x, res.fun

def cct_McCamy(C):
	x, y, _ = xyY(C)

	n = (x - 0.3320) / (y - 0.1858)
	return -449 * n**3 + 3525 * n**2 - 6823.3 * n + 5520.33