#!/usr/bin/env Python3
import numpy as np
from scipy.optimize import fsolve


class ADC:
	def __init__(self, wedge=(1.6554, 6.28927, -1.48306)):
		self.p1 = prism('S-LAH71', 15, wedge[0], wedge[1])
		self.p2 = prism('S-FPL51', 15, wedge[1], wedge[2])
		self.p3 = prism('S-FPL51', 15, wedge[2], wedge[1])
		self.p4 = prism('S-LAH71', 15, wedge[1], wedge[0])
		self.theta = 90

	def setOrientation(self, theta):
		self.p1.setPrismOrientation(theta)
		self.p2.setPrismOrientation(theta)
		self.p3.setPrismOrientation(180-1*theta)
		self.p4.setPrismOrientation(180-1*theta)
		self.theta = theta

	def setFieldRotation(self, FR):
		self.p1.setPrismOrientation(self.theta-FR)
		self.p2.setPrismOrientation(self.theta-FR)
		self.p3.setPrismOrientation(180-1*self.theta-FR)
		self.p4.setPrismOrientation(180-1*self.theta-FR)

	def calculateDispersion(self, l1, l2, hx=0, hy=0, refraction=False):
		'''
			Calculates the dispersion for the ADC using the current position, 
			for a normalized on sky field angle (hx, hy). Default field position
			is in the center of the FoV.
		'''
		_prismList = [self.p1, self.p2, self.p3, self.p4]
		_d1x_list, _d1y_list = [], []
		_pupilRatio = 0.07943/38.542

		hx = hx * 7.5e-3
		hy = hy * 7.5e-3
		ix = hx * 1/_pupilRatio
		iy = hy * 1/_pupilRatio

		for l in [l1,l2]:
			_deltaListX = np.zeros(len(_prismList))
			_deltaListY = np.zeros(len(_prismList))
			for p in range(4):
				_n_g = _prismList[p].indexOfRefraction(l)
				_b1x = _prismList[p].tilt1x
				_b1y = _prismList[p].tilt1y

				_b2x = _prismList[p].tilt2x
				_b2y = _prismList[p].tilt2y

				if p==0:
				    _t_iz_n_x = ix # Zero before first prism
				    _t_iz_n_y = iy # Zero before first prism

				# X direction    
				_t_i1_n_x = _t_iz_n_x
				_d1x = deviationCalculation3(_t_i1_n_x, _b1x, _b2x, _n_g)
				_t_tz_n_x = _t_iz_n_x - _d1x
				_t_iz_n_x = _t_tz_n_x         # Set output to new input (for the next prism)

				# Y direction
				_t_i1_n_y = _t_iz_n_y
				_d1y = deviationCalculation3(_t_i1_n_y, _b1y, _b2y, _n_g)
				_t_tz_n_y = _t_iz_n_y - _d1y
				_t_iz_n_y = _t_tz_n_y         # Set output to new input (for the next prism)

				# Adding the deviations to the list with delta terms
				_deltaListX[p] = _d1x 
				_deltaListY[p] = _d1y
			_d1x_list.append(_deltaListX.sum())
			_d1y_list.append(_deltaListY.sum())

		self.dispX = ( _d1x_list[1] - _d1x_list[0] ) * 3600 * 1000 * _pupilRatio
		self.dispY = ( _d1y_list[1] - _d1y_list[0] ) * 3600 * 1000 * _pupilRatio
		
		if refraction:
			rx = ( _d1x_list[0] ) * 3600 * 1000 * _pupilRatio
			ry = ( _d1y_list[0] ) * 3600 * 1000 * _pupilRatio
			return rx, ry
		return self.dispX, self.dispY

	def optimize(self, inputDispersion, l1, l2, x0=50):
		_solver = fsolve(optimizeFunction, x0=x0, args=(self, inputDispersion, l1, l2))
		return _solver[0]




class prism:
	def __init__(self, glass, thickness, tilt1, tilt2=None):
		self.glass 		= glass
		self.thickness 	= thickness
		self.tilt1x		= tilt1
		self.tilt1y 	= 0
		self.tilt2x		= tilt2
		self.tilt2y 	= 0

		# Remember original tilts for reset:
		self.tilt1x_OG	= tilt1
		self.tilt1y_OG 	= 0
		self.tilt2x_OG	= tilt2
		self.tilt2y_OG 	= 0


	def indexOfRefraction(self, l):
		if self.glass == 'S-FPL51':
			a = [1.170105050E+00, 4.757107830E-02, 7.638324450E-01]
			b = [6.162039240E-03, 2.633728760E-02, 1.418826420E+02]
		elif self.glass == 'S-LAH71':
			a = [1.982800310E+00, 3.167584500E-01, 2.444726460E+00]
			b = [1.189874560E-02, 5.271560010E-02, 2.132206970E+02]
		return dispersionFormula(l, a, b) 

	def setPrismOrientation(self, theta):
		self.tilt1x, self.tilt1y = tiltRotation(self.tilt1x_OG, self.tilt1y_OG, theta)
		self.tilt2x, self.tilt2y = tiltRotation(self.tilt2x_OG, self.tilt2y_OG, theta)
	


def tiltRotation(tx,ty, theta):
	theta = np.deg2rad(theta-90)
	tx_new = tx * np.cos(theta) - ty * np.sin(theta)
	ty_new = tx * np.sin(theta) + ty * np.cos(theta)
	return tx_new, ty_new


def dispersionFormula(l, a, b):
	n = np.sqrt( (a[0]* l**2)/(l**2 - b[0]) + (a[1] * l**2) / (l**2 - b[1]) + (a[2] * l**2) / (l**2 - b[2]) + 1)
	return n


# def deviationCalculation(t_i1, b1, b2, n):
    # Returns deviation of the ray after the prism with apex angle alpha
    # Units: degrees
    # DO NOT USE!
    # alpha = b2 - b1
    # delta = t_i1 + np.arcsin( np.sin(np.deg2rad(alpha)) * np.sqrt(n**2-(np.sin(np.deg2rad(t_i1)))**2) - np.sin(np.deg2rad(t_i1)) * np.cos(np.deg2rad(alpha)) )*180/np.pi - alpha
    # return delta


def deviationCalculation2(t_i1, b1, b2, n):
	# Returns deviation of the ray after the prism with apex angle alpha
	# only now using Snell's law only
    # Units: degrees
    t_i1  = np.deg2rad(t_i1)
    b1    = np.deg2rad(b1)
    b2    = np.deg2rad(b2)

    t_i   = b1 - t_i1
    t_f   = np.arcsin( 1 / n * np.sin( t_i ) )
    t_f_z = b1 - t_f

    t_i   = b2 - t_f_z
    t_f   = np.arcsin( n / 1 * np.sin( t_i ) )
    t_f_z = b2 - t_f

    delta = t_i1 - t_f_z
    delta = np.degrees(delta)
    return delta

def deviationCalculation3(t_i1, b1, b2, n):
	# Using the formula from Hagen & Tkazcyk (2011)
	# Gives exactly the same results as deviationCaclulation2(), above
	t_i1 	= np.deg2rad(t_i1)
	b1 		= np.deg2rad(b1)
	b2 		= np.deg2rad(b2)
	alpha 	= b2 - b1
	delta 	= t_i1 - b2 - np.arcsin( n * np.sin( np.arcsin(1/n * np.sin(t_i1 - b1) ) - alpha ) )
	return np.degrees(delta)

def optimizeFunction(theta, adcObject, inputDispersion, l1, l2, FR=None):
	if FR:
		FRdeg = np.deg2rad(FR) # 30 degrees field rotation
		inputDispersionX = np.cos(FRdeg) * inputDispersion
		inputDispersionY = np.sin(FRdeg) * inputDispersion
		adcObject.setOrientation(theta)
		ADCdisp = adcObject.calculateDispersion(l1,l2)
		return np.sqrt((ADCdisp[1] - inputDispersionY)**2 + (ADCdisp[0] - inputDispersionX)**2)
	else:	
		adcObject.setOrientation(theta)
		ADCdisp = adcObject.calculateDispersion(l1,l2)
		ADCdisp = ADCdisp[1]
		return ADCdisp - inputDispersion

def angleFunc(x, a):
	'''
		Function to fit the ideal angle on
	'''
	return np.arccos(a * np.tan(np.deg2rad(x))) * 180/np.pi