# -*- coding: iso-8859-1 -*-
"""
Purpose of this code is to form mixing ratio files and T/P profiles for generic prebiotic Earths. The idea is to closely follow Rugheimer et al. 2015, i.e. to start with N2-CO2 atmospheres, 1 bar total pressure. 

Each function call produces two files as output. These are a 1D T/P profile and a 1D molar concentration profile.

We output this model in a form usable by our RT code. (high resolution, preferably ~0.1 km or better, required)
"""

########################
###Import useful libraries
########################
import numpy as np
import matplotlib.pyplot as plt
import pdb
import scipy.integrate
import scipy.optimize
from matplotlib.pyplot import cm
import thermodynamic_functions as thermo

########################
###Define useful constants, all in CGS (via http://www.astro.wisc.edu/~dolan/constants.html)
########################

#Unit conversions
km2m=1.e3 #1 km in m
km2cm=1.e5 #1 km in cm
cm2km=1.e-5 #1 cm in km
amu2g=1.66054e-24 #1 amu in g
bar2atm=0.9869 #1 bar in atm
Pascal2bar=1.e-5 #1 Pascal in bar
Pa2bar=1.e-5 #1 Pascal in bar
bar2Pa=1.e5 #1 bar in Pascal
deg2rad=np.pi/180.
bar2barye=1.e6 #1 Bar in Barye (the cgs unit of pressure)
barye2bar=1.e-6 #1 Barye in Bar
micron2m=1.e-6 #1 micron in m
micron2cm=1.e-4 #1 micron in cm
metricton2kg=1000. #1 metric ton in kg

#Fundamental constants
c=2.997924e10 #speed of light, cm/s
h=6.6260755e-27 #planck constant, erg/s
k=1.380658e-16 #boltzmann constant, erg/K
sigma=5.67051e-5 #Stefan-Boltzmann constant, erg/(cm^2 K^4 s)
R_earth=6371.*km2m#radius of earth in m
R_sun=69.63e9 #radius of sun in cm
AU=1.496e13#1AU in cm

#Mean molecular masses
m_co2=44.01*amu2g #co2, in g
m_h2o=18.02*amu2g #h2o, in g
m_n2=28.01*amu2g #n2, in g

#Mars parameters
g=981. #surface gravity of Earth, cm/s**2


########################
###Exponential atmosphere, 0.9 bar N2, 0.1 bar CO2, T_0=288K
########################
def generate_profiles_prebioticearth_exponential(P_0_bar, T_0, r_co2,r_so2, r_h2s, savefilename):
	"""
	Generates input files for an exponential atmosphere
	P_0_bar: Dry total surface pressure in Bar
	T_0: surface temperature in K
	r_co2: mole fraction of CO2 relative to bulk dry air (CO2+N2)
	r_so2: mole fraction of SO2 relative to bulk dry air (CO2+N2)
	r_h2s: mole fraction of H2S relative to bulk dry air (CO2+N2)
	savefilename: name the output files generated by the function will be saved with
	"""
	P_0=P_0_bar*bar2barye #Dry CO2 surface pressure, from bar-->barye. Specified when you assume an atmospheric thickness

	#gas composition of atmosphere
	r_n2=1.-r_co2 #dry air is just N2 and CO2
	mmm=r_co2*m_co2+r_n2*m_n2 #mean molecular mass in g

	H=(k*T_0)/(mmm*g) #scale height in cm

	#########################
	###Compute exponential atmosphere
	#########################
	z_edges=np.arange(0., 100., step=0.1)*km2cm #edges of altitude bins
	P_edges=P_0*np.exp(-z_edges/H) #pressure at each edge; simple isothermal atmosphere
	T_edges=np.zeros(len(z_edges))+T_0 #temperature at each edge; taken to be isothermal in this approx
	n_edges=P_edges/(k*T_edges) #number densities at each edge

	#########################
	###Print Results in Usable Format
	#########################
	####Form T/P Profile
	towrite_tp=np.zeros([len(z_edges),4])
	towrite_tp[:,0]=z_edges #altitude in cm
	towrite_tp[:,1]=T_edges #temperature in K
	towrite_tp[:,2]=n_edges #number density in cm**-3
	towrite_tp[:,3]=P_edges*barye2bar #Pressure in bar
	np.savetxt('./TPProfiles/'+savefilename+'_tpprofile.dat', towrite_tp, delimiter='	', fmt='%1.7e', newline='\n', header='\n'.join(['Formed from form_atmospheric_profiles_paper_model.py','Z (cm)	T (K)	DEN (cm**-3)	P (bar)']))

	####Form Molar Concentration File
	towrite_mc=np.zeros([len(z_edges),9])
	towrite_mc[:,0]=towrite_tp[:,0] #altitude in cm
	towrite_mc[:,1]+=r_n2 #N2
	towrite_mc[:,2]+=r_co2#CO2
	towrite_mc[:,3]=0#H2O
	towrite_mc[:,4]=0. #CH4
	towrite_mc[:,5]+=r_so2 #SO2
	towrite_mc[:,6]=0. #O2
	towrite_mc[:,7]=0. #O3
	towrite_mc[:,8]+=r_h2s #H2S
	np.savetxt('./MolarConcentrations/'+savefilename+'_molarconcentrations.dat', towrite_mc, delimiter='	', fmt='%1.7e', newline='\n', header='\n'.join(['Formed from prebioticearth_atmosphere_models.generate_profiles_prebioticearth_exponential','Z (cm)		N2	CO2	H2O	CH4	SO2	O2	O3	H2S']))
	

	#########################
	###Plot Results
	#########################
	###First, plot dry adiabat
	fig, (ax1, ax2, ax3)=plt.subplots(3, figsize=(8,10), sharey=True)

	ax1.plot(P_edges*barye2bar,z_edges*cm2km,  color='black',linestyle='-')
	ax1.set_ylabel('Altitude (km)')
	ax1.set_xlabel('Pressure (bar)')
	ax1.set_xscale('log')

	ax2.plot(T_edges, z_edges*cm2km, color='black',linestyle='-')
	ax2.set_ylabel('Altitude (km)')
	ax2.set_xlabel('Temperature (K)')

	ax3.plot(n_edges, z_edges*cm2km, color='black',linestyle='-')
	ax3.set_ylabel('Altitude (km)')
	ax3.set_xlabel('Number Density (cm**-3)')
	ax3.set_xscale('log')

	plt.show()
	
#generate_profiles_prebioticearth_exponential(1., 288., 'test_exponential') #Test run of function



########################
###N2-CO2-H2O Atmosphere (unsaturated, dry adiabat evolution)
########################
def generate_profiles_prebioticearth_n2_co2_h2o_unsaturated(P_0_bar, T_0, r_co2,r_so2, r_h2s, relH, savefilename):
	"""
	Generates input files for a realistic N2-CO2-H2O atmosphere, but one with relative humidity<1 (i.e., dry adiabat evolution)
	P_0_bar: Dry total surface pressure in Bar
	T_0: surface temperature in K
	r_co2: mole fraction of CO2 relative to bulk dry air (CO2+N2)
	r_so2: mole fraction of SO2 relative to bulk dry air (CO2+N2)
	r_h2s: mole fraction of H2S relative to bulk dry air (CO2+N2)
	savefilename: name the output files generated by the function will be saved with
	
	NOTE: Assumes atmospheric thermodynamics set by N2+CO2, H2O tacked on afterwards
	"""
	P_0=P_0_bar*bar2barye #Dry total surface pressure, from bar-->barye. Specified when you assume an atmospheric thickness

	P_strat_0=0.1*bar2barye #pressure at which the stratosphere starts, based on the insight by Robinson & Catling 2013 that terrestrial triatomic atmospheres tend to have stratospheres starting at p~0.1 because that's where radiative windows open up and radiative cooling/heating start to dominate

	#gas composition of atmosphere
	r_n2=1.-r_co2 #dry air is just N2 and CO2
	m_r_n2=(r_n2*m_n2)/(r_n2*m_n2+r_co2*m_co2) #N2 mass mixing ratio
	m_r_co2=(r_co2*m_co2)/(r_n2*m_n2+r_co2*m_co2) #CO2 mass mixing ratio
	
	mmm=m_co2*r_co2+m_n2*r_n2 #mean molecular mass in g

	R_atm=k/mmm #specific gas constant for (dry) atmosphere
	R_h2o=h/m_h2o #specific gas constant for H2O gas.
	
	L_h2o=3.34e5*1.e4 #Heat of fusion of CO2. Value taken from PPC Table 2.1 p 92, converted from SI (Joule/Kg) to erg/gram. Taken as constant. IS FUSION RIGHT???

	#Specific heat at constant pressures for gases
	c_p_co2=thermo.c_p_shomate('co2', T_0) #converted from kJ/(Kg*K) to erg/(g*K)
	c_p_n2=thermo.c_p_shomate('n2', T_0) #converted from kJ/(Kg*K) to erg/(g*K)
	c_p_h2o=1.97*1.e7 #converted from kJ/(Kg*K) to erg/(g*K), measured at 1 atm and 338-589 F ()
	
	c_p_atm=c_p_co2*m_r_co2+c_p_n2*m_r_n2 #mean specific heat at constant pressure for atmosphere


	########################
	###Form functions to compute the dry and moist adiabats, the stratospheric temperature, and the overall temperature
	########################
	def T_dry_func(P): #constant c_p_co2
		return T_0*(P/P_0)**(R_atm/c_p_atm)
	
	def p_sat_h2o(T):
		return thermo.p_sat_h2o_wordsworth(T) #T in K
		

	#Already know P_strat_0, the pressure where the stratosphere starts. Now, get T_strat, the temperature of the stratosphere. Get this by assuming isothermal stratosphere following discussion with Robin on 4/18/16, coupled with assumption of continuity from the dry adiabat
	T_strat_0=T_dry_func(np.array([P_strat_0])) #Stratospheric temperature in K.

	def T_strat_func(P):
		return 0.0*P+T_strat_0


	###Now, establish overall atmospheric temperature profile function. You determine which temperature to use in each bin by choosing the largest temperature, i.e. the slowest decline in temperature with altitude. This mimics the effect of letting radiative or dry lapse set your temperature.
	def T_atm_func(P):
		if P_0<=P_strat_0: #this is the pure stratosphere case
			T_atm=T_strat_func(P)
		else: #This is the case where you have a dry and stratospheric regime
			if type(P)==float: #If you get just a single float as input, you can't treat in the same way
				if (P>P_strat_0):
					return T_dry_func(P)
				if (P<=P_strat_0):
					return T_strat_func(P)
			else:
				inds0=np.where(P>P_strat_0)
				inds2=np.where(P<=P_strat_0)
				
				T_atm=np.zeros(np.shape(P))
				T_atm[inds0]=T_dry_func(P[inds0])
				T_atm[inds2]=T_strat_func(P[inds2])	
		return T_atm

	########################
	###Partition the atmosphere into layers, initialize variables to hold relevant parameters
	########################
	#Dry air pressures
	P_edges=P_0*np.exp(-np.arange(0,19.01, step=0.01)) #These are the edges of the pressure bins
	P_centers=np.sqrt(P_edges[0:-1]*P_edges[1:]) #centers of pressure bins, determined by geometric mean
	

	########################
	###Compute P-T profiles
	########################
	#Assumes thermodynamics totally dry gas
	T_dry_edges=T_dry_func(P_edges)
	T_strat_edges=T_strat_func(P_edges)

	T_atm_centers=T_atm_func(P_centers)
	T_atm_edges=T_atm_func(P_edges)

	
	#########################
	####Compute altitudes corresponding to each pressure level
	#########################
	####Way 1: Do by assuming isothermal atmosphere in each layer, and using hydrostatic equilibrium equation
	z_edges=np.zeros(np.shape(P_edges))
	z_centers=np.zeros(np.shape(P_centers))
	z_widths=np.zeros(np.shape(P_centers))

	for ind in range(0, len(P_centers)):
		H=k*T_atm_centers[ind]/(mmm*g) #scale height within the layer
		z_widths[ind]=-H*np.log(P_edges[ind+1]/P_edges[ind]) #delta-z spanned by the layer
		z_edges[ind+1]=z_widths[ind]+z_edges[ind] #add to the zero point to get absolute altitude
		z_centers[ind]=z_edges[ind]+-H*np.log(P_centers[ind]/P_edges[ind])

	###Way 2: integrate equation for hydrostatic equilibrium
	###define function to get z(P)
	def integrand(P):
		return -(k*T_atm_func(P))/(mmm*g*P)

	def altitude_pressure(P):
		z=np.zeros(np.shape(P))
		for ind in range(0, len(z)):
			z[ind]=scipy.integrate.quad(integrand, P_0, P[ind], epsrel=1.e-5, epsabs=0, limit=200)[0] #1e-5 accuracy
		return z

	#######Test difference between both methods
	####z_centers_alt=altitude_pressure(P_centers)
	####z_edges_alt=altitude_pressure(P_edges)
	####z_widths_alt=z_edges_alt[1:]-z_edges_alt[:-1]
	####fracdiff_z=(z_centers_alt-z_centers)/z_centers
	####print np.max(fracdiff_z)
	####print np.min(fracdiff_z)
	#####agree to <1%

	########################
	###Compute water vapor levels. Assuming an atmosphere that is fully saturated in H2O until the tropopause, then maintains constant mixing ratio of water vapor,
	########################	
	P_h2o_edges=relH*p_sat_h2o(T_atm_edges)
	
	conc_h2o=P_h2o_edges/(P_h2o_edges+P_edges)
	conc_h2o_strat=(conc_h2o[P_edges<=P_strat_0])[0]
	conc_h2o[P_edges<=P_strat_0]=conc_h2o_strat #...until hits stratosphere.
	P_h2o_edges[P_edges<=P_strat_0]=conc_h2o_strat*P_edges[P_edges<=P_strat_0]

	P_tot_edges=P_h2o_edges+P_edges

	conc_n2=r_n2*P_edges/P_tot_edges
	conc_co2=r_co2*P_edges/P_tot_edges
	conc_so2=r_so2*P_edges/P_tot_edges
	conc_h2s=r_h2s*P_edges/P_tot_edges
	
	print np.max(np.abs(1.-(conc_h2o+conc_co2+conc_n2+conc_so2+conc_h2s)))# CHECK: If not close to 0, is a problem. 
	
	n_tot=P_tot_edges/(k*T_atm_edges) #number density of total atmosphere (H2O+CO2+N2)
	#########################
	###Print Results in Usable Format
	#########################
	####Form T/P Profile
	towrite_tp=np.zeros([len(P_edges),4])
	towrite_tp[:,0]=z_edges #altitude in cm
	towrite_tp[:,1]=T_atm_edges#temperature in K
	towrite_tp[:,2]=n_tot #number density in cm**-3
	towrite_tp[:,3]=P_tot_edges*barye2bar #Pressure in bar
	np.savetxt('./TPProfiles/'+savefilename+'_tpprofile.dat', towrite_tp, delimiter='	', fmt='%1.7e', newline='\n', header='\n'.join(['Formed from prebioticearth_atmosphere_models.py','Z (cm)	T (K)	DEN (cm**-3)	P (bar)']))

	####Form Molar Concentration File
	towrite_mc=np.zeros([len(P_edges),9])
	towrite_mc[:,0]=towrite_tp[:,0] #altitude in cm
	towrite_mc[:,1]=conc_n2 #N2
	towrite_mc[:,2]=conc_co2#CO2
	towrite_mc[:,3]=conc_h2o#H2O
	towrite_mc[:,4]=0. #CH4
	towrite_mc[:,5]=conc_so2 #SO2
	towrite_mc[:,6]=0. #O2
	towrite_mc[:,7]=0. #O3
	towrite_mc[:,8]=conc_h2s #H2S
	np.savetxt('./MolarConcentrations/'+savefilename+'_molarconcentrations.dat', towrite_mc, delimiter='	', fmt='%1.7e', newline='\n', header='\n'.join(['Formed from prebioticearth_atmosphere_models.py','Z (cm)		N2	CO2	H2O	CH4	SO2	O2	O3	H2S'])) #REALLY MOLAR CONCENTRATIONS...HOW DOES INTERPRET?

	
	#########################
	###Plot Results
	#########################
	###First, plot dry adiabat
	fig, (ax1, ax2)=plt.subplots(2, figsize=(8,10), sharey=True)

	ax1.set_title('T-P Profile')
	ax1.plot(T_strat_edges, P_edges*barye2bar, color='green',linestyle='--', label='Stratosphere')
	ax1.plot(T_dry_edges, P_edges*barye2bar, color='red',linestyle='--', label='Dry Adiabat')
	ax1.plot(T_atm_edges, P_edges*barye2bar, color='blue',linestyle='-', label='Dry Atmosphere (CO2+N2)')
	ax1.plot(T_atm_edges, P_tot_edges*barye2bar, color='black',linestyle='--', label='Total Atmosphere (CO2+N2+H2O)')

	ax1.legend(loc='lower left', fontsize=10)
	ax1.set_ylabel('Dry Pressure (bar, CO2+N2)')
	ax1.set_xlabel('Temperature (K)')
	ax1.set_yscale('log')
	#ax1.set_ylim([np.min(P_tot_edges*barye2bar), np.max(P_tot_edges*barye2bar)])
	ax1.invert_yaxis()

	tick_locations=np.array([1., 1.e-1, 1.e-2, 1.e-3, 1.e-4])*P_0*barye2bar
	ax1b=ax1.twinx()
	ax1b.set_ylim(ax1.get_ylim())
	ax1b.set_yscale('log')
	ax1b.set_yticks(tick_locations)
	ax1b.set_yticklabels(np.round(altitude_pressure(tick_locations*bar2barye)*cm2km, 1))#(altitude_pressure(tick_locations*bar2barye)*cm2km)
	ax1b.set_ylabel('Altitude (km)')
	ax1b.minorticks_off()

	ax2.set_title('Concentrations')
	ax2.plot(conc_n2,P_edges*barye2bar, color='green', label='N2')
	ax2.plot(conc_co2,P_edges*barye2bar, color='red', label='CO2')
	ax2.plot(conc_h2o,P_edges*barye2bar, color='blue', label='H2O')

	ax2.legend(loc=0, fontsize=10)
	ax2.set_ylabel('Dry Pressure (bar, CO2+N2)')
	ax2.set_xlabel('Molar Concentration')
	ax2.set_yscale('log')
	ax2.set_xscale('log')
	ax2.set_ylim([np.min(P_edges*barye2bar), np.max(P_edges*barye2bar)])
	ax2.invert_yaxis()

	tick_locations=np.array([1., 1.e-1, 1.e-2, 1.e-3, 1.e-4])*P_0*barye2bar
	ax2b=ax2.twinx()
	ax2b.set_ylim(ax1.get_ylim())
	ax2b.set_yscale('log')
	ax2b.set_yticks(tick_locations)
	ax2b.set_yticklabels(np.round(altitude_pressure(tick_locations*bar2barye)*cm2km, 1))#(altitude_pressure(tick_locations*bar2barye)*cm2km)
	ax2b.set_ylabel('Altitude (km)')
	ax2b.minorticks_off()

	#plt.savefig('./PrelimPlots/prelim_mars_v2.pdf', orientation='portrait',papertype='letter', format='pdf')
	plt.show()

#generate_profiles_prebioticearth_n2_co2_h2o_unsaturated(1.0, 320., 0.1,1., 'test-drythermodynamics')
