# -*- coding: utf-8 -*-
'''
This is a Python code provided for calculating gamma_air, n-air and gamma_self, 
of CS2 transitions in the HITRAN database as described in:
Karlovets  et al,  2021
Note that the program can be easily modified for any other file formats.
'''

from astropy.io import ascii
from astropy.table import Table, Column
import numpy as np
import sys

#--------------read CS2 HITRAN data-------------------------------

readpath = input('input HITRAN 160 .par file to do the calculation for gamma_air, n-air and gamma_self,-broadening and temperature dependence of CS2:')

savepath = input('output file name:')

file1 = open(readpath,'r') 
table = file1.read()
file1.close


total = ascii.read(table, Reader = ascii.FixedWidthNoHeader,
                         names = ('parline', 'branch', 'J'),
                         col_starts = (0, 117, 118),
                         col_ends = (159, 117, 120),
                         ) # This work assumes the HITRAN .par format is the input data

# unused columns in next steps
Parline = np.array(total['parline'])

# column used in next steps
Branch = np.array(total['branch'])
J = np.array(total['J'])

#-----------------calcuating |m| for CS2 lines----------------------------------
# *Note that m stands for |m| which is related to the lower J rotational quantum number as follows:
# P branch: m = -J" (However in this work we are using |m| so for P branches this is just J")
# Q branch: m = J"
# R branch: m = J" + 1

m = []
for i in range(len(Branch)):
    if Branch[i]=='R':
        m.append(J[i]+1)
    if Branch[i]=='P':
        m.append(J[i])
    if Branch[i]=='Q':
        m.append(J[i])
    else:
        pass

#-----------------define function for gSelf---------------------------------------
def gself(x):
    a0 = 0.16024
    a1 = 0.01105
    a2 = -7.745E-4
    a3 = 1.6848E-5
    b1 = 0.0601
    b2 = -2.970E-3
    b3 = 4.4245E-5
    b4 = 9.0837E-7

    ggself= ((a0+a1*x+a2*x**2+a3*x**3)/(1+b1*x+b2*x**2+b3*x**3+b4*x**4))
    return ggself# x in this calculation stands for |m|

def err_gself(x):
    if x<40:
        err = 5
    elif x>=40:
        err = 4
    return err    
    
#-----------------define function for gair---------------------------------------
def gair(x):
    a0 = 0.14875
    a1 = 0.66097
    a2 = -0.03369
    a3 = 0.00108
    b1 = 5.79533
    b2 = -0.29481
    b3 = 0.00885
    b4 = 4.46479E-5
    
    ggair = ((a0+a1*x+a2*x**2+a3*x**3)/(1+b1*x+b2*x**2+b3*x**3+b4*x**4))
    return ggair 

def err_gair(x):
    if x<40:
        err = 5
    elif x>=40:
        err = 4
    return err  

#-----------------define function for nair---------------------------------------
def nair(x):
    a0 = -97.80455 
    a1 = 1382.91708
    a2 = -63.59603
    a3 = 0.85226
    b1 = 1753.25152
    b2 = -72.83604
    b3 = 0.87376
    b4 = 0.00148
    
    nnair = ((a0+a1*x+a2*x**2+a3*x**3)/(1+b1*x+b2*x**2+b3*x**3+b4*x**4))
    return nnair 

def err_nair(x):
    if x<70:
        err = 5
    elif x>=70:
        err = 4
    return err  
  
#--------------Fill empty lists with calculated broadening-------------------------------

gamma_self= []
n_self= []
err_n_self= []
ref_n_self= []
err_self= []
ref_self= []

gamma_air = []
n_air = []
err_n_air = []
ref_n_air = []
err_air = []
ref_air = []

#--The reference numbers below correspond to "global reference IDs" in the HITRAN database. The mapping is also provided here in the code."
for i in range(len(m)):
    xx = m[i]
    gamma_He.append(float(gHe(xx)))   # selfbroadening
    n_He.append('0.3000')             # selftemperature dependence
    err_n_He.append('4')              # selftemperature dependence uncertainty code
    ref_n_He.append('1501')# selftemperature dependence Data Reference: Nakamichi et al. 2006 https://doi.org/10.1039/B511772K
    err_He.append(str(err_gHe(xx)))   # selfuncertainty code
    ref_He.append("1501")  # selfBroadening Data Reference: Nakamichi et al. 2006 https://doi.org/10.1039/B511772K
    gamma_air.append(float(gair(xx))) # air broadening
    n_air.append(float(nair(xx)))     # air temperature dependence
    err_n_air.append(str(err_nair(xx)))# air temperature dependence uncertainty code
    ref_n_air.append('1273')# air temperature dependence Data Reference: Hashemi et al. 2020 https://doi.org/10.1016/j.jqsrt.2020.107283
    err_air.append(str(err_gair(xx))) # air uncertainty code
    ref_air.append("1359")# air Broadening Data Reference: Hashemi et al. 2020 https://doi.org/10.1016/j.jqsrt.2020.107283

#------------create new HITRAN data file with self and air broadening and temperature dependence for CS2--------

with open(savepath,'w') as out:
    for i in range(len(m)):
        out.write("%160s, %6.3f, %3s, %3s, %6.4f, %3s, %3s, %4.2f, %3s, %3s \n" %( Parline[i],
        gamma_self[i], err_self[i], ref_self[i], 
        gamma_air[i], err_air[i], ref_air[i], n_air[i], err_n_air[i], ref_n_air[i]))
    else:
        print('end for calculation: output "160.par + gamma_self + gamma_air + n_air" ')