#!/usr/bin/env Python3
import numpy as np
import matplotlib.pyplot as plt
from AstroAtmosphere import *
from ADC import *
import sys
import scipy.optimize as opt



'''
@ Author:		Joost van den Born 
@ Contact:		born@astron.nl
@ Description:	This script finds the filter constant using
				- Direct application of equation 18 and 19
				- The minimization algorithm implemented in the ADC class
				and compares these two.
'''







class filter:
	def __init__(self, name, l1, l2):
		self.name 				= name
		self.lmin 				= np.min([l1,l2])
		self.lmax 				= np.max([l1,l2])
		self.bandwidth 			= self.lmax - self.lmin
		self.optimumAngle 		= None
		self.residualDispersion = None
		self.optimumAngleForAllZ = None
		self.cf 				= None

	def analytic_optimum_position_ADC(self, z, l1, l2, dispersion_object, adc_object):
		disp = dispersion_object
		adc = adc_object

		_R = disp.rc / (disp.rc + disp.H)
		n1  = at.n_tph(l=self.lmin, T=T, p=P, RH=H, xc=xc)
		n2  = at.n_tph(l=self.lmax, T=T, p=P, RH=H, xc=xc)
		dn_atm = n2 - n1
		_dispersion = np.deg2rad(disp.cassini(n1,n2,z))

		_A1 = np.deg2rad(adc.p1.tilt2x_OG - adc.p1.tilt1x_OG)
		_A2 = np.deg2rad(adc.p2.tilt2x_OG - adc.p2.tilt1x_OG)

		_A3 = -1 * np.deg2rad(adc.p3.tilt2x_OG - adc.p3.tilt1x_OG) # -1 because doublet gets flipped after setting rotation
		_A4 = -1 * np.deg2rad(adc.p4.tilt2x_OG - adc.p4.tilt1x_OG) # -1 because doublet gets flipped after setting rotation

		_dn_p1 = adc.p1.indexOfRefraction(l1) - adc.p1.indexOfRefraction(l2)
		_dn_p2 = adc.p2.indexOfRefraction(l1) - adc.p2.indexOfRefraction(l2)

		_pupilRatio = 0.07943/38.542
		_Psi = (_dn_p1 * (_A1 + _A4) + _dn_p2 * (_A2 + _A3)) * _pupilRatio

		_cf = -1/_pupilRatio * (dn_atm)/((_dn_p1 * (_A1 + _A4) + _dn_p2 * (_A2 + _A3)))
		_thetaOpt = np.arccos(1/_Psi * _dispersion)
		return np.degrees(_thetaOpt), _cf

def ADCposAnalytical(z, Cf):
	return np.arccos(Cf*np.tan(np.deg2rad(z)))*180/np.pi


if __name__ == "__main__":
	Iband = filter('I-band',0.8,0.934)
	Jband = filter('J-band',1.15,1.345)
	Hband = filter('H-band',1.49,1.78)
	Kband = filter('K-band',1.97,2.32)


	##########################################################
	# Atmospheric model
	# Parameters at Cerro Armazones
	T   = 279.65    # k
	P   = 71200     # Pa
	H   = 0.22
	xc  = 450       # ppm
	lat = -24.5983  # degrees
	h   = 3064      # m

	# Initializing dispersion model
	at  = Observatory()

	# Density of the atmosphere (following CIPM-81/91 equations)
	rho = at.rho(p=P, T=T, RH=H, xc=xc)

	# Initializing dispersion object
	disp = dispersion(lat, h)
	disp.setReducedHeight(P, rho)
	
	##########################################################

	adc = ADC()


	# Unlock this when the above works
	# zRange = range(2,62)
	zRange = np.arange(5,65.1,0.1)
	optimumAngleFilterList = []
	fig = plt.figure()
	ax = fig.add_subplot(111)


	for f in [Hband, Iband, Jband, Kband]:
		# For new filters or if the zenith values are different, the new values are calculated		
		print ('Starting %s...' %f.name	)
		# Calculating indices of refraction for l1 and l2
		n1  = at.n_tph(l=f.lmin, T=T, p=P, RH=H, xc=xc)
		n2  = at.n_tph(l=f.lmax, T=T, p=P, RH=H, xc=xc)
		dispersion_obs = disp.cassini(n1, n2, zRange)*3.6e6

		adc.setOrientation(0)
		theta_opt = [adc.optimize(disp, f.lmin, f.lmax) for disp in dispersion_obs]
		f.optimumAngleForAllZ = theta_opt
		# Plotting the data
		ax.plot(zRange, f.optimumAngleForAllZ, label=f.name)

	for f in [Hband, Iband, Jband, Kband]:
		popt, pcov = opt.curve_fit(ADCposAnalytical, zRange[:-50], f.optimumAngleForAllZ[:-50], p0=0.5)
		f.cf = popt[0]
		ax.plot(zRange, ADCposAnalytical(zRange, f.cf), 'k--', label='best fit %s' %f.name)
		theta_opt_analytical, cf_analytical = f.analytic_optimum_position_ADC(zRange, f.lmin, f.lmax, disp, adc)
		ax.plot(zRange, ADCposAnalytical(zRange, cf_analytical), 'b--', label='analytical cf %s' %f.name)
		print('{} - Best fit: {} - Analytical: {}'.format(f.name, f.cf, cf_analytical))

	# Plot properties
	ax.set_title('Optimum rotation angle of the first ADC prism', fontsize=9)
	ax.set_ylabel('Rotation angle (deg)', fontsize=9)
	ax.set_xlabel('Zenith angle (deg)', fontsize=9)	
	ax.grid(True, alpha=0.5)
	ax.legend(fontsize=9)

	plt.show()