#!/usr/bin/env Python
from __future__ import division
import numpy as np
import matplotlib.pyplot as plt
import AstroAtmosphere as atm
import sys
from ADC import *

'''
@author: 		Joost van den Born
@Description: 	Calculates differential refraction (dispersion) between stars with a different 
				effective temperature. Similar to Gubler & Tytler 1998
'''

def blackbody(l, T):
	'''
		Calculates the intensity output of a star with effective temperature T at wavelength l
	'''
	h = 6.626070150e-34	# J*s 
	k = 1.38064852e-23	# J/K
	c = 299792458		# m/s
	return 2 * np.pi * h * c**2 / l**5 * 1 / (np.exp( h * c / (l * k * T)) - 1)

def effectiveWavelength(T, lmin, lmax):
	'''
		Calculates effective wavelength of a star from the effective temperature.
	'''
	lmin = lmin/1e6
	lmax = lmax/1e6
	_l = np.linspace(lmin, lmax, 1000)
	_I = blackbody(_l, T)
	return np.sum(_I*_l)/np.sum(_I)*1e6

if __name__ == '__main__':
	# Parameters at Cerro Armazones
	T   = 279.65 	# k
	P   = 71200  	# Pa
	H   = 0.22
	xc  = 450   	# ppm
	lat = -24.5983 	# degrees
	h   = 3064  	# m
	l1  = 1.49 		# micron
	l2  = 1.78 		# micron
	zenith = 60 	# degrees
	offset = 0#.285/2		# degrees, ADC offset to simulate finite step size

	S2temps = [2500, 4000, 5500, 7000, 8500]
	S1temps = np.linspace(2500, 1e4, 100)


	########################################################################################
	########################################################################################
	# 1. We look at the differential refraction through the atmosphere, without an ADC
	# at = atm.dispersion(lat, h)
	at = atm.refraction(lat, h)
	at.setReducedHeight(P, at.rho(P, T, H, xc) )


	# Calculating the differential atmospheric refraction for stars of different temperatures
	data = []
	for T2 in S2temps:
		# Properties of star 2
		l_eff_2 = effectiveWavelength(T2, l1, l2)
		n2 = at.n_tph(l_eff_2, T, P, H, xc)
		R2 = at.cassini(n2, zenith)

		# Calculating the differences with star 1
		dataEntry = []
		for T1 in S1temps:
			l_eff_1 = effectiveWavelength(T1, l1, l2)
			n1 		= at.n_tph(l_eff_1, T, P, H, xc)
			R1 		= at.cassini(n1, zenith)
			diffRef = (R1 - R2) * 3.6e9
			dataEntry.append(diffRef)
		data.append(dataEntry)

	# Plotting the results
	plt.rc('font', family='serif')
	plt.rc('axes', labelsize='large')
	fig = plt.figure(figsize=(9,6))
	ax = fig.add_subplot(111)
	ax2 = ax.twiny()
	ax.grid(True, alpha=0.3)
	linestylesList = [':', '-.', '--', '-', (0, (3, 1, 1, 1, 1, 1))]
	for i in range(len(data)):
		ax.plot(S1temps, data[i], label=str(S2temps[i]) + ' K', linestyle=linestylesList[i], color='k', linewidth=2)
	
	S1temp_ticks = ax.get_xticks()
	S1temp_tick_loc = ax.get_xticks()
	
	l_eff_ticks = []
	for T1 in S1temp_ticks:
		l_eff_ticks.append(round( (effectiveWavelength(T1, l1, l2)*1000-1620),1))
	
	ax2.set_xticklabels(l_eff_ticks)
	ax2.set_xlim(np.min(S1temps), np.max(S1temps))
	ax2.set_xlabel('$\lambda_{eff} - 1620$ (nm) ', fontsize=18)
	ax2.tick_params(axis ='both', direction='in', labelsize=16) 

	ax.set_xlabel('Temperature of Star 1 (K)', fontsize=18)
	ax.set_ylabel('$\Delta$R ($\mu$as)', fontsize=18)
	ax.legend(fontsize=16, labelspacing=0.3)
	ax.set_xlim(np.min(S1temps), np.max(S1temps))
	ax.tick_params(axis ='both', direction='in', labelsize=16) 
	# plt.show()

	########################################################################################
	########################################################################################
	


	########################################################################################
	########################################################################################
	# 2. We look now at the differential refraction through the atmosphere, including an ADC
	at = atm.refraction(lat, h)
	at.setReducedHeight(P, at.rho(P, T, H, xc) )

	adc = ADC()

	# Conditions at the band cut-on and cut-off wavelength
	n_l1 = at.n_tph(l1, T, P, H, xc)
	n_l2 = at.n_tph(l2, T, P, H, xc)


	atd = atm.dispersion(lat, h)
	atd.setReducedHeight(P, at.rho(P, T, H, xc) )
	# Optimize ADC for broad band edges at zenith distance.
	atm_disp_full_band = atd.cassini(n_l1, n_l2, zenith)
	print(atm_disp_full_band*3.6e6)
	adcOptimumAngle = adc.optimize(atm_disp_full_band*3.6e6, l1, l2) 
	adc.setOrientation(adcOptimumAngle+offset)


	# Calculating the differential atmospheric refraction for stars of different temperatures
	data = []
	for T2 in S2temps:
		# Properties of star 2
		l_eff_2 = effectiveWavelength(T2, l1, l2)
		n2 = at.n_tph(l_eff_2, T, P, H, xc)
		R2_atm = at.cassini(n2, zenith)
		R2_adc = adc.calculateDispersion(l_eff_2, l1, refraction=True)

		# Calculating the differences with star 1
		dataEntry = []
		for T1 in S1temps:
			l_eff_1 = effectiveWavelength(T1, l1, l2)
			n1 		= at.n_tph(l_eff_1, T, P, H, xc)
			R1_atm	= at.cassini(n1, zenith)
			R1_adc 	= adc.calculateDispersion(l_eff_1, l1, refraction=True)
			diffRef = ( (R1_adc[1] - R2_adc[1]) + ((R1_atm - R2_atm) * 3.6e6 ) ) * 1e3
			dataEntry.append(diffRef)
		data.append(dataEntry)

	# Plotting the results
	plt.rc('font', family='serif')
	plt.rc('axes', labelsize='large')
	fig = plt.figure(figsize=(9,6))
	ax = fig.add_subplot(111)
	ax2 = ax.twiny()
	ax.grid(True, alpha=0.3)
	linestylesList = [':', '-.', '--', '-', (0, (3, 1, 1, 1, 1, 1))]
	for i in range(len(data)):
		ax.plot(S1temps, data[i], label=str(S2temps[i]) + ' K', linestyle=linestylesList[i], color='k', linewidth=2)
	
	S1temp_ticks = ax.get_xticks()
	S1temp_tick_loc = ax.get_xticks()
	
	l_eff_ticks = []
	for T1 in S1temp_ticks:
		l_eff_ticks.append(round( (effectiveWavelength(T1, l1, l2)*1000-1620),1))
	
	ax2.set_xticklabels(l_eff_ticks)
	ax2.set_xlim(np.min(S1temps), np.max(S1temps))
	ax2.set_xlabel('$\lambda_{eff} - 1620$ (nm) ', fontsize=18)
	ax2.tick_params(axis ='both', direction='in', labelsize=16) 

	ax.set_xlabel('Temperature of Star 1 (K)', fontsize=18)
	ax.set_ylabel('$\Delta$R ($\mu$as)', fontsize=18)
	ax.legend(fontsize=16, labelspacing=0.3, loc='lower right')
	ax.set_xlim(np.min(S1temps), np.max(S1temps))
	ax.tick_params(axis ='both', direction='in', labelsize=16) 
	plt.show()

	########################################################################################
	########################################################################################



	########################################################################################
	########################################################################################
	# 3. Now we look at the differential dispersion, without an ADC
	at = atm.dispersion(lat, h)
	at.setReducedHeight(P, at.rho(P, T, H, xc) )

	# Conditions at the band cut-on and cut-off wavelength
	n_l1 = at.n_tph(l1, T, P, H, xc)
	n_l2 = at.n_tph(l2, T, P, H, xc)

	# Optimize ADC for broad band edges at zenith distance.
	atm_disp_full_band = at.cassini(n_l1, n_l2, zenith)
	print(atm_disp_full_band*3.6e6)


	# Check that we can use dispersion? If not, how to calculate refraction...?
	# We plot below the difference in atmospheric dispersion of the two stars, incl. an ADC.
	data = []
	for T2 in S2temps:
		l_eff_2 = effectiveWavelength(T2, l1, l2)
		n2 = at.n_tph(l_eff_2, T, P, H, xc)
		dataEntry = []
		for T1 in S1temps:
			l_eff_1 = effectiveWavelength(T1, l1, l2)
			n1 		= at.n_tph(l_eff_1, T, P, H, xc)
			atm_disp = at.cassini(n1, n2, zenith) * 3.6e9
			dataEntry.append(atm_disp)
		data.append(dataEntry)

	plt.rc('font', family='serif')
	plt.rc('axes', labelsize='large')
	fig2 = plt.figure(figsize=(9,6))
	ax = fig2.add_subplot(111)
	ax2 = ax.twiny()
	ax.grid(True, alpha=0.3)
	linestylesList = [':', '-.', '--', '-', (0, (3, 1, 1, 1, 1, 1))]
	for i in range(len(data)):
		ax.plot(S1temps, data[i], label=str(S2temps[i]) + ' K', linestyle=linestylesList[i], color='k', linewidth=2)
	
	S1temp_ticks = ax.get_xticks()
	S1temp_tick_loc = ax.get_xticks()
	
	l_eff_ticks = []
	for T1 in S1temp_ticks:
		l_eff_ticks.append(round( (effectiveWavelength(T1, l1, l2)*1000-1620),1))
	
	# ax2.plot(S1temps, data[0], alpha=0)
	ax2.set_xticklabels(l_eff_ticks)
	ax2.set_xlim(np.min(S1temps), np.max(S1temps))
	ax2.set_xlabel('$\lambda_{eff} - 1620$ (nm) ', fontsize=18)
	ax2.tick_params(axis ='both', direction='in', labelsize=16) 

	ax.set_xlabel('Temperature of Star 1 (K)', fontsize=18)
	ax.set_ylabel('$\Delta(\Delta$R)  ($\mu$as), excl. ADC', fontsize=18)
	ax.legend(fontsize=16, labelspacing=0.3, loc='lower right')
	ax.set_xlim(np.min(S1temps), np.max(S1temps))
	ax.tick_params(axis ='both', direction='in', labelsize=16) 
	# plt.show()

	########################################################################################
	########################################################################################
	########################################################################################




	########################################################################################
	########################################################################################
	# 4. Now we add the ADC and look again at the differential refraction, but on the detector
	at = atm.dispersion(lat, h)
	at.setReducedHeight(P, at.rho(P, T, H, xc) )
	adc = ADC()

	# Conditions at the band cut-on and cut-off wavelength
	n_l1 = at.n_tph(l1, T, P, H, xc)
	n_l2 = at.n_tph(l2, T, P, H, xc)

	# Optimize ADC for broad band edges at zenith distance.
	atm_disp_full_band = at.cassini(n_l1, n_l2, zenith)
	print(atm_disp_full_band*3.6e6)
	adcOptimumAngle = adc.optimize(atm_disp_full_band*3.6e6, l1, l2) 
	adc.setOrientation(adcOptimumAngle+offset)


	# Check that we can use dispersion? If not, how to calculate refraction...?
	# We plot below the difference in atmospheric dispersion of the two stars, incl. an ADC.
	data = []
	for T2 in S2temps:
		l_eff_2 	= effectiveWavelength(T2, l1, l2)
		n2 			= at.n_tph(l_eff_2, T, P, H, xc)
		atm_disp2 	= at.cassini(n_l1, n2, zenith) * 3.6e6
		adc_disp2 	= adc.calculateDispersion(l1, l_eff_2)[1]

		dataEntry = []
		for T1 in S1temps:
			l_eff_1 	= effectiveWavelength(T1, l1, l2)
			n1 			= at.n_tph(l_eff_1, T, P, H, xc)
			atm_disp1 	= at.cassini(n_l1, n1, zenith) * 3.6e6
			adc_disp1 	= adc.calculateDispersion(l1,l_eff_1)[1]
			diffDisp 	= ((adc_disp1 - adc_disp2) - (atm_disp1 - atm_disp2)) * 1e3
			dataEntry.append(diffDisp)
		data.append(dataEntry)

	plt.rc('font', family='serif')
	plt.rc('axes', labelsize='large')
	fig2 = plt.figure(figsize=(9,6))
	ax = fig2.add_subplot(111)
	ax2 = ax.twiny()
	ax.grid(True, alpha=0.3)
	linestylesList = [':', '-.', '--', '-', (0, (3, 1, 1, 1, 1, 1))]
	for i in range(len(data)):
		ax.plot(S1temps, data[i], label=str(S2temps[i]) + ' K', linestyle=linestylesList[i], color='k', linewidth=2)
	
	S1temp_ticks = ax.get_xticks()
	S1temp_tick_loc = ax.get_xticks()
	
	l_eff_ticks = []
	for T in S1temp_ticks:
		l_eff_ticks.append(round( (effectiveWavelength(T, l1, l2)*1000-1620),1))
	
	# ax2.plot(S1temps, data[0], alpha=0)
	ax2.set_xticklabels(l_eff_ticks)
	ax2.set_xlim(np.min(S1temps), np.max(S1temps))
	ax2.set_xlabel('$\lambda_{eff} - 1620$ (nm) ', fontsize=18)
	ax2.tick_params(axis ='both', direction='in', labelsize=16) 

	ax.set_xlabel('Temperature of Star 1 (K)', fontsize=18)
	ax.set_ylabel('$\Delta(\Delta$R)  ($\mu$as), incl. ADC', fontsize=18)
	ax.legend(fontsize=16, labelspacing=0.3, loc='lower right')
	ax.set_xlim(np.min(S1temps), np.max(S1temps))
	ax.tick_params(axis ='both', direction='in', labelsize=16) 
	plt.show()

	########################################################################################
	########################################################################################
	########################################################################################