#!/usr/bin/env Python3
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from matplotlib import cm
from matplotlib.colors import LogNorm
from AstroAtmosphere import *
from scipy import interpolate
import json


'''
@ Author:		Joost van den Born 
@ Contact:		born@astron.nl
@ Description:	This script uses the NCEP NARR data to model a real structured atmosphere 
				from data with a resolution of 2.5 km per data point. We interpolate this
				data to a higher resolution.
				Then a full integration over this atmosphere is done and compared to the 
				values for the refraction and dispersion if we had only taken the atmospheric
				conditions directly at the observatory.
'''


# Load a variable from a json text file.
def loadVariableFromTXT(filename):
	variableFile = open(filename, 'r')
	variable = json.load(variableFile)
	variableFile.close()
	return np.asarray(variable)


plt.rc('font', family='serif')
plt.rc('axes', labelsize='large')

# h, _, _, _, Temp, Press, Dens, vSound, _,_ = np.genfromtxt('si1py_modified.txt', unpack=True, skip_header=6)

dist 	= loadVariableFromTXT('export_dist.txt')
heights = loadVariableFromTXT('export_heights.txt')
press 	= loadVariableFromTXT('export_press_atmosphere.txt')
temp 	= loadVariableFromTXT('export_temp_atmosphere.txt')
print("Loading complete!")

# Conditions at Cerro Armazones
T   = temp[0,0]	# k
P   = press[0,0]# Pa
H   = 0.0
xc  = 314   	# ppm
lat = -24.5983 	# degrees
h0   = 0  		# m
l1  = 1.49 		# micron
l2  = 1.78 		# micron
zenith = 60 	# 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)

RHlist = np.asarray([H]*len(press[:,0]))
xclist = np.asarray([xc]*len(press[:,0]))
Dens = at.rho(p=press[:,0], T=temp[:,0], RH=RHlist, xc=xclist)

refr = refraction(lat, h0) 
refr.setReducedHeight(P, rho)
disp = dispersion(lat, h0)
disp.setReducedHeight(P, rho)


# We must create a cube of weather data:
# We assume a flat atmosphere for now, so that h does not need to be adjusted
# x_max = 200 km, gives zenith_max = ~66.7 degrees
# Every 20 km in horizontal direction a new perturbation and interpolation
# for every 500 m.
x_max = dist[-1] # km

Temp_cube_n = []
Press_cube_n = []

for x in range(len(dist)):
	Temp_cube_n.append(temp[:,0])
	Press_cube_n.append(press[:,0])

Temp_cube_n = np.asarray(Temp_cube_n).T
Press_cube_n = np.asarray(Press_cube_n).T



fig = plt.figure()
ax1 = fig.add_subplot(211)
im = ax1.pcolormesh(dist,heights, temp, cmap=cm.coolwarm)
fig.colorbar(im, ax=ax1, label='Temperature (K)')
ax1.set_aspect(1)
ax1.set_ylabel('Height (km)')

ax2 = fig.add_subplot(212)
im2 = ax2.pcolormesh(dist, heights, press, cmap=cm.coolwarm, norm=LogNorm())
ax2.set_aspect(1)
ax2.set_xlabel('Horizontal distance (km)')
ax2.set_ylabel('Height (km)')
fig.colorbar(im2, ax=ax2, label='Pressure (Pa)')
plt.tight_layout()
plt.show(block=False)

# Now for a given zenith angle we need to find the temp and press along the line of sight
f_Temp_2d  		= interpolate.interp2d(dist, heights, temp, kind='linear')
f_Press_2d 		= interpolate.interp2d(dist, heights, press, kind='linear')
f_Temp_2d_n  	= interpolate.interp2d(dist, heights, Temp_cube_n, kind='linear')
f_Press_2d_n 	= interpolate.interp2d(dist, heights, Press_cube_n, kind='linear')
print('Interpolation complete!')



zenithRange = np.linspace(30,65,10)
refrDiff, dispDiff = [], []

for zenith in zenithRange:
	z = np.deg2rad(zenith)
	y_los = heights
	x_los = np.linspace(0, heights.max()*np.tan(z), len(heights))

	Temp_los, Press_los = [], []
	Temp_los_n, Press_los_n = [], []
	for i in range(len(heights)):
		 Temp_los.append(f_Temp_2d(x_los[i], y_los[i])[0])
		 Press_los.append(f_Press_2d(x_los[i], y_los[i])[0])
		 Temp_los_n.append(f_Temp_2d_n(x_los[i], y_los[i])[0])
		 Press_los_n.append(f_Press_2d_n(x_los[i], y_los[i])[0])

	Temp_los  	= np.asarray(Temp_los)
	Press_los 	= np.asarray(Press_los)
	Temp_los_n 	= np.asarray(Temp_los_n)
	Press_los_n = np.asarray(Press_los_n)
	rho_los   	= at.rho(p=Press_los, T=Temp_los, RH=H, xc=xc)
	rho_los_n 	= at.rho(p=Press_los_n, T=Temp_los_n, RH=H, xc=xc)

	refrInt = refr.refractionIntegral(n1, z, useStandardAtmosphere=False, heightData=heights, rhoData=rho_los_n)*3600*1000
	refrInt_pert = refr.refractionIntegral(n1, z, useStandardAtmosphere=False, heightData=heights, rhoData=rho_los)*3600*1000
	refrDiff.append(refrInt_pert - refrInt)

	dispInt = disp.refractionIntegral(n1, n2, z, useStandardAtmosphere=False, heightData=heights, rhoData=rho_los_n)*3.6e6
	dispInt_pert = disp.refractionIntegral(n1, n2, z, useStandardAtmosphere=False, heightData=heights, rhoData=rho_los)*3.6e6
	dispDiff.append(dispInt_pert - dispInt)





fig3 = plt.figure(figsize=(9,7))
ax1 = fig3.add_subplot(211)
ax2 = fig3.add_subplot(212)
ax1.tick_params(axis ='both', direction='in', which='both', labelsize=16) 
ax2.tick_params(axis ='both', direction='in', which='both', labelsize=16) 
ax1.grid(True, alpha=0.3)#, which='major')
ax2.grid(True, alpha=0.3)#, which='major')


ax1.plot(zenithRange, np.log10(np.abs(refrDiff)), 'k-', linewidth=2)
ax2.plot(zenithRange, np.log10(np.abs(dispDiff)), 'k-', linewidth=2)
ax1.set_xlabel('Zenith (degrees)', fontsize=18)
ax2.set_xlabel('Zenith (degrees)', fontsize=18)
ax1.set_ylabel('$\log_{10}($|$R - R_{o}$|$)$ (mas)', fontsize=18)
ax2.set_ylabel('$\log_{10}($|$\Delta R - \Delta R_{o}$|$)$ (mas)', fontsize=18)
plt.tight_layout()
plt.show()