# -*- coding: utf-8 -*-
"""
Created on Sat Oct 22 15:35:03 2016

@author: Jacob Hedelius

A function to calculate Rayleigh scattering for a given wavelength
v - 22-Oct-2016
"""

import numpy as np
import os
curr_dir = os.getcwd()
os.chdir('../../python_subs')
#import egi_subs as egs
os.chdir(curr_dir)
import sub_readaod as sra

def rayleigh_OT(wavelengths_nm, surf_p_hPa, map_file):
    
    ## function to calculate the Rayleigh scattering using ggg files for inputs
    # This function is based on the work of Bodhaine et al. 1999.
    # (http://dx.doi.org/10.1175/1520-0426(1999)016<1854:ORODC>2.0.CO;2)
    # If finding a .map file each time is too difficult - may consider
    # hard-coding values derived from them as constants...
    
    # example inputs
    #wavelengths_nm  = egs.frange(250, 1000, 10)
    ##wavelengths_nm = np.array(wavelengths_nm)
    ##wavelengths_nm = np.transpose(np.reshape(wavelengths_nm, [len(wavelengths_nm),1]))
    #map_file = '../dn20160825.map'
    #surf_p_hPa = egs.frange(950, 1000, 10) #list of surface pressures
    # surf_p_hPa = 1013.25;
    
    ## check inputs & reshape to a uniform format
    if not(os.path.isfile(map_file)):
        raise ValueError('Missing .map file: '+ map_file)
    
    shp_nm = np.shape(wavelengths_nm)
    if shp_nm==(): #only a single wavelength given
        shp_nm=[1,1]
    wavelengths_nm = np.reshape(wavelengths_nm, [np.product(shp_nm),1])
    
    shp_p = np.shape(surf_p_hPa)
    if shp_p==(): #only a single pressure given
        shp_p=[1,1]
    surf_p_hPa = np.reshape(surf_p_hPa, [1, np.product(shp_p)])
    
    
    ## read in map file
    map_raw = sra.read_map(map_file)
    
    # calculate approximate pressure weighting function (Connor et al 2008).  Approximate because non-exact surface p used.
    p_surf_wgt = map_raw['Pressure_hPa'][0] #surface pressure for weighting
    h_vals_aprx = np.zeros([len(map_raw['Pressure_hPa']),1])
    h_vals_aprx[0] = np.abs((-1*map_raw['Pressure_hPa'][0] + (map_raw['Pressure_hPa'][1] - \
      map_raw['Pressure_hPa'][0])/np.log(map_raw['Pressure_hPa'][1]/map_raw['Pressure_hPa'][0])))/p_surf_wgt #first value
    
    h_vals_aprx[len(map_raw['Pressure_hPa'])-1] = np.abs((map_raw['Pressure_hPa'][-1] - (map_raw['Pressure_hPa'][-1] - \
      map_raw['Pressure_hPa'][-2])/np.log(map_raw['Pressure_hPa'][-1]/map_raw['Pressure_hPa'][-2])))/p_surf_wgt #end value
    
    plsts = np.concatenate((map_raw['Pressure_hPa'][0:-2], map_raw['Pressure_hPa'][1:-1], map_raw['Pressure_hPa'][2:])) #abbreviated variable name of pressure level listings
    plsts = np.transpose(plsts.reshape([3, int(len(plsts)/3)]))
    h_vals_aprx[1:-1] = np.reshape(np.abs((-1*plsts[:,1]+(plsts[:,2]-plsts[:,1])/np.log(plsts[:,2]/plsts[:,1])) + \
      (plsts[:,1]-(plsts[:,1]-plsts[:,0])/np.log(plsts[:,1]/plsts[:,0])))/p_surf_wgt, [len(plsts),1]) #calculate remaing h-values
    
    aprx_xgas = {} #pre-form dictionary
    aprx_xgas['co2_ppm'] = np.sum(map_raw['co2_ppm']*np.transpose(h_vals_aprx)) #approximate amount of CO2, from prior
    g_wgt = np.sum(map_raw['gravity_m_s2']*np.transpose(h_vals_aprx))*100 #pressure weighted g in cm/s^2
    
    ## estimate the refractive index of dry air
    lambda_um = wavelengths_nm/1000
    bodhaine_eq18 = 8060.51 + 2480990/(132.274- lambda_um**(-2)) + 17455.7/(39.32957 - lambda_um**(-2))
    bodhaine_eq19 = 1 + 0.54*(aprx_xgas['co2_ppm']*1E-6 - 0.0003)
    
    nair = bodhaine_eq18*bodhaine_eq19*1E-8+1 #refractive index of DRY air, including current CO2 value. FOR: 288.15 K, 1013.25 hPa
    
    ## calculate depolarization factor
    F_N2 = 1.034 + 3.17*1E-4*lambda_um**(-2) #depolarization of N2
    F_O2 = 1.096 + 1.385*1E-3*lambda_um**(-2) + 1.448*1E-4*lambda_um**(-4) #depolarization of O2
    F_air_co2 = (78.084*F_N2 + 20.946*F_O2 + 0.934*1.00 + aprx_xgas['co2_ppm']*1E-6*1.15)/ \
      (78.084 + 20.946 + 0.934 + aprx_xgas['co2_ppm']*1E-6)
    # rho_air = (6*F_air_co2-6)/(7*F_air_co2+3) # depolarization ratio
    
    
    ## scattering cross section of air
    lambda_cm = wavelengths_nm/1E7 # wavelengths in cm
    N_s = 2.546899E19 # molecules/cm^-3 at 288.15 K and 1013.25 hPa
    xsect = 24*np.pi**3*(nair**2-1)**2/(lambda_cm**4*N_s**2*(nair**2+2)**2)*F_air_co2 # scattering cross section of air
    xsect_rep = np.repeat(xsect, shp_p, axis=1)  #*np.ones([len(surf_p_hPa),1]) # make the right size
    
    
    ## Rayleigh optical depths
    surf_p_dyncm2 = surf_p_hPa*1000 # pressure in dyn/cm^2
    surf_p_dyncm2_rep = np.repeat(surf_p_dyncm2, len(xsect), axis=0)   
    m_a = 15.0556*aprx_xgas['co2_ppm']*1E-6 + 28.9595 # molecular weight dry air
    
    tau_R = xsect_rep*surf_p_dyncm2_rep*map_raw['Avogadro']/(m_a*g_wgt) 
    
    return tau_R