#!/usr/bin/env Python
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import host_subplot
from AstroAtmosphere import *
from pyslalib import slalib


'''
@ Author:		Joost van den Born 
@ Contact:		born@astron.nl
@ Description:	This script was used to compare the impact of using the old 
				B&S equations, and the particular implementation of SLALIB. 
				The rest of the script is almost a one-to-one copy of other 
				script that compared the atmospheric models.

				Modeling the atmospheric refraction using multiple models of 
				the atmosphere and then comparing them.

				Models included:
				- Full integration method
				- Plane-parallel atmosphere
				- Cassini's model
				- Tan5 model (A*tan(z) + B*tan^3(z) + C*tan^5(z))
				- Oriani's theorem (A*tan(z) + B*tan^3(z))
				- Erf model (from Corbard et al., 2019)
				- SLALIB
'''
class atmosphereSLALIB:
	def __init__(self, T=279.65, P=712, H=0.22):
		self.name 	= 'SLALIB atmosphere'
		self.T 		= T
		self.P 		= P
		self.H 		= H
		self.h 		= 3064
		self.lat 	= -24.5983
		self.TLR 	= 0.0065
		self.prec 	= 1e-8

	def atmModel(self, A, B, zenith):
		zenith = np.radians(zenith)
		return A * np.tan(zenith) + B * (np.tan(zenith)**3)

	def dispersion(self, l1, l2, z):
		A1, B1 = slalib.sla_refco(self.h, self.T, self.P, self.H, l1, np.deg2rad(self.lat), self.TLR, self.prec)
		A2, B2 = slalib.sla_atmdsp(self.T, self.P, self.H, l1, A1, B1, l2)
		D1 = np.degrees(self.atmModel(A1, B1, z)) * 3600
		D2 = np.degrees(self.atmModel(A2, B2, z)) * 3600
		dispersionValue = ((D1 - D2) * 1000)
		return dispersionValue

	def refractivity(self, l=None):
		'''
			SLALIB uses the same atmospheric model as ZEMAX
			based on the works of Hohenkerk & Sinclair and Seidelmann.
			They take the B&S refraction formula as a basis.
		'''
		zObs = np.deg2rad(45.06563)
		h = 0
		T = self.T
		P = self.P
		H = self.H
		lat = 24
		tlr = 0.0065
		eps = 1e-10
		if not l:
			l = self.lrange
		R = []
		if type(l) == float:
			l = [l]
		for wl in l:
			R.append(slalib.sla_refro(zObs, h, T, P, H, wl, lat, tlr, eps)+1)
		return np.asarray(R)


if __name__ == '__main__':
	plt.rc('font', family='serif')
	# plt.rc('xtick', labelsize='x-small')
	# plt.rc('ytick', labelsize='x-small')
	plt.rc('axes', labelsize='large')

	# Standard atmospheric conditions
	# T   = 279.65 	# k
	# P   = 71200  	# Pa
	# H   = 0.22
	# xc  = 450   	# ppm
	T   = 288.15 	# k
	P   = 101325  	# Pa
	H   = 0.0
	xc  = 314   	# ppm
	lat = -24.5983 	# degrees
	h   = 3064  	# m
	l1  = 1.49 		# micron
	l2  = 1.78 		# micron
	zenith = 30 	# degrees

	# Calculate the index of refraction
	at  = Observatory()
	n1  = at.n_tph(l=l1, T=T, p=P, RH=H, xc=xc)
	n2  = at.n_tph(l=l2, T=T, p=P, RH=H, xc=xc)
	rho = at.rho(p=P, T=T, RH=H, xc=xc)
	
	#######################################################
	# DIRECT COMPARISON
	ref = dispersion(lat, h)
	ref.setReducedHeight(P, rho)
	print('Ciddor refractivity: ', n1, n2)
	print ('Cassini dispersion: %s mas' %(ref.cassini(n1, n2, zenith) * 3.6e6))
	n1c = n1
	n2c = n2

	refraction = refraction(lat, h)
	refraction.setReducedHeight(P, rho)
	refractionValueCiddor = refraction.cassini(n1c, zenith) * 3.6e6 # in mas

	# SLALIB
	atSLALIB = atmosphereSLALIB(T, P/100, H)
	n1 = atSLALIB.refractivity(l1)
	n2 = atSLALIB.refractivity(l2)

	print('B&S refractivity: ', n1, n2)
	print ('SLALIB model: %s mas' %(atSLALIB.dispersion(l1, l2, zenith) ) )

	# SLALIB refraction:
	A1, B1 = slalib.sla_refco(atSLALIB.h, atSLALIB.T, atSLALIB.P, atSLALIB.H, l1, np.deg2rad(atSLALIB.lat), atSLALIB.TLR, atSLALIB.prec)
	refractionValueSLALIB = np.degrees(atSLALIB.atmModel(A1, B1, zenith)) * 3.6e6 # in mas

	print('Differential dispersion = {} mas'.format(atSLALIB.dispersion(l1, l2, zenith) - ref.cassini(n1c, n2c, zenith) * 3.6e6) )
	print('Differential refraction = {} mas'.format( refractionValueSLALIB - refractionValueCiddor ) )

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

	refractionAnalysis  = True
	dispersionAnalysis  = True
	tableCompare 		= True
	if dispersionAnalysis:
		# Initiate the refraction class
		ref = dispersion(lat, h)
		ref.setReducedHeight(P, rho)

		# Compare the refraction models to check if the output is okay
		print ('Plane - parallel atmosphere dispersion: %s mas' %(ref.planeParallel(n1, n2, zenith) * 3.6e6))
		print ('Cassini dispersion: %s mas' %(ref.cassini(n1, n2, zenith) * 3.6e6))
		print ('Tan5 dispersion: %s mas' %(ref.tan5(n1, n2, zenith) * 3.6e6))
		print ('Oriani dispersion: %s mas' %(ref.oriani(n1, n2, zenith) * 3.6e6))
		print ('Corbard dispersion: %s mas' %(ref.corbard(n1, n2, zenith)*3.6e6))
		print ('Full integration: %s mas' %(ref.refractionIntegral(n1, n2, zenith)*3.6e6))
		print ('SLALIB model: %s mas' %(atSLALIB.dispersion(l1, l2, zenith) ) )

		# Now we calculate the dispersion over a wide range of zenith angles
		# for ideal gas atmosphere
		zenithRange = np.linspace(30, 85, 1000)
		ppList 		= ref.planeParallel(n1, n2, zenithRange) * 3.6e6
		cassiniList = ref.cassini(n1, n2, zenithRange) * 3.6e6
		tan5List 	= ref.tan5(n1, n2, zenithRange) * 3.6e6
		orianiList  = ref.oriani(n1, n2, zenithRange) * 3.6e6
		corbardList = [ref.corbard(n1, n2, z)[0] * 3.6e6 for z in zenithRange]
		fullIntegrationList = np.asarray([ref.refractionIntegral(n1, n2, z) * 3.6e6 for z in zenithRange])
		slalibList	= atSLALIB.dispersion(l1, l2, zenithRange)

		# for isothermal atmosphere
		cassiniList_isotherm = ref.cassini(n1, n2, zenithRange, reduced_height=ref.H_isotherm(T)) * 3.6e6
		tan5List_isotherm 	 = ref.tan5(n1, n2, zenithRange, reduced_height=ref.H_isotherm(T)) * 3.6e6
		orianiList_isotherm  = ref.oriani(n1, n2, zenithRange, reduced_height=ref.H_isotherm(T)) * 3.6e6
		corbardList_isotherm = [ref.corbard(n1, n2, z, reduced_height=ref.H_isotherm(T))[0] * 3.6e6 for z in zenithRange]

		# check when the variation between the different models pass different values
		dispersionWidth = []
		widthList = []
		for z in range(len(zenithRange)):
			compList = [cassiniList[z], tan5List[z], orianiList[z], corbardList[z], cassiniList_isotherm[z], tan5List_isotherm[z], orianiList_isotherm[z], corbardList_isotherm[z], slalibList[z]]#, fullIntegrationList[z]]
			width    = np.max(compList) - np.min(compList)
			widthList.append(width)
			dispersionWidth.append((zenithRange[z], width))

		masBorder = np.argmax(np.asarray(widthList)>1)
		p1masBorder = np.argmax(np.asarray(widthList)>0.1)
		p01masBorder = np.argmax(np.asarray(widthList)>0.01)
		p001masBorder = np.argmax(np.asarray(widthList)>0.001)

		plotAbsolute = 1
		if plotAbsolute:
			# Plotting it all, with vertical lines where the variation between the models exceed 1, 0.1, 0.01 and 0.001 mas.
			fig = plt.figure(figsize=(9,5))
			ax  = fig.add_subplot(111)
			ax.set_xlim([30,85])
			ax.set_xlabel('Zenith (degrees)')
			ax.set_ylabel('Dispersion (mas)')
			ax.grid(True, alpha=0.3)
			ax.axvline(dispersionWidth[masBorder][0], linestyle='--', color='k', linewidth=0.5)
			ax.text(dispersionWidth[masBorder][0]-0.5, 402, '1 mas', ha='center', va='bottom', rotation='vertical')
			ax.axvline(dispersionWidth[p1masBorder][0], linestyle='--', color='k', linewidth=0.5)
			ax.text(dispersionWidth[p1masBorder][0]-0.5, 402, '0.1 mas', ha='center', va='bottom', rotation='vertical')
			ax.axvline(dispersionWidth[p01masBorder][0], linestyle='--', color='k', linewidth=0.5)
			ax.text(dispersionWidth[p01masBorder][0]-0.5, 402, '10 $\mu$as', ha='center', va='bottom', rotation='vertical')
			ax.axvline(dispersionWidth[p001masBorder][0], linestyle='--', color='k', linewidth=0.5)
			ax.text(dispersionWidth[p001masBorder][0]-0.5, 402, '1 $\mu$as', ha='center', va='bottom', rotation='vertical')

			ax.plot(zenithRange, ppList, 'k-', label='Plane-parallel', alpha=0.2)
			ax.plot(zenithRange, cassiniList, 'g-', label='Cassini')
			ax.plot(zenithRange, tan5List, 'b-', label='$Atan(z) - Btan^3(z) + Ctan^5(z)$')
			ax.plot(zenithRange, orianiList, 'y-', label='$Atan(z) - Btan^3(z)$')
			ax.plot(zenithRange, corbardList, 'r-', label='Error function')
			ax.plot(zenithRange, fullIntegrationList, 'c-', label='Full integration')

			ax.plot(zenithRange, cassiniList_isotherm, 'g--', label='Cassini (isotherm)')
			ax.plot(zenithRange, tan5List_isotherm, 'b--', label='$Atan(z) - Btan^3(z) + Ctan^5(z)$ (isotherm)')
			ax.plot(zenithRange, orianiList_isotherm, 'y--', label='$Atan(z) - Btan^3(z)$ (isotherm)')
			ax.plot(zenithRange, corbardList_isotherm, 'r--', label='Error function (isotherm)')


			ax.tick_params(axis ='both', direction='in') 
			ax.legend()
			plt.tight_layout()
			plt.show()

		plotRelative = 1
		if plotRelative:
			# RELATIVE TO FULL INTEGRATION
			fig = plt.figure(figsize=(9,6))
			ax  = fig.add_subplot(111)
			ax.set_xlim([30,85])
			
			ax.set_xlabel('Zenith (degrees)', fontsize=18)
			ax.set_ylabel(r'$\log_{10}(|(\Delta R - \Delta R_{int})|$ mas$^{-1})$', fontsize=18)
			
			ax.axhline(np.log10(0.4), linestyle='-', linewidth=1, color='red')
			ax.text(30.5, np.log10(0.4)+0.05, 'Residual dispersion requirement', fontsize=12)

			ax.grid(True, which='major', alpha=0.3)
			ax.grid(True, which='minor', alpha=0.2)

			ax.plot(zenithRange, np.log10(np.abs(ppList-fullIntegrationList)), 					color='k', linestyle='-',  linewidth=2, label='Plane-parallel')
			ax.plot(zenithRange, np.log10(np.abs(cassiniList-fullIntegrationList)), 			color='g', linestyle=(0, (1, 1)),  linewidth=2, label='Cassini')
			ax.plot(zenithRange, np.log10(np.abs(orianiList-fullIntegrationList)), 				color='y', linestyle='-.', linewidth=2, label='Oriani')
			ax.plot(zenithRange, np.log10(np.abs(tan5List-fullIntegrationList)), 				color='b', linestyle=(0, (5, 1)), linewidth=2, label='Oriani incl. fifth order')
			ax.plot(zenithRange, np.log10(np.abs(np.asarray(corbardList)-fullIntegrationList)), color='r', linestyle=(0, (3, 1, 1, 1, 1, 1)), linewidth=2, label='Error function')
			ax.plot(zenithRange, np.log10(np.abs(np.asarray(slalibList)-fullIntegrationList)), 	color='m', linewidth=2, label='SLALIB')
			ax.plot(zenithRange, np.log10(np.abs(cassiniList-fullIntegrationList)), 			color='g', linestyle=(0, (1, 1)),  linewidth=2)

			ax.set_ylim(top=2.6, bottom=-7.2)
			ax.tick_params(axis ='both', direction='in', labelsize=16) 
			ax.legend(fontsize=14, labelspacing=0.3, columnspacing=0.5, handletextpad=0.3,  ncol=1, loc='lower right')
			plt.tight_layout()
			plt.show()