# -*- coding: utf-8 -*-
"""
Created on Sat Oct 22 18:52:41 2016

@author: Jacob Hedelius

Main script to calculate I0 and AOD from EM27 measurements

v - 23-Oct-2016
"""

import numpy as np
#import scipy as sc
from scipy import interpolate
import os
curr_dir = os.getcwd()
os.chdir('../python_subs')
import egi_subs as egs
import imp
imp.reload(egs)
os.chdir(curr_dir + '/python_subs')
import sub_readaod as sra
import run_calc as rc
import rayleigh_OT as rOT
os.chdir(curr_dir)



#def calc_1win1day_aod(rl_fname, daymap_fname, h2oam_fname, settings_fname, sat_fname, ic_fname):

# This function calculates I0, and all the tau parameters for 1 day, and 1
# window

# example inputs
rl_fname = 'dn20160825.grl' #runlog file for flags
daymap_fname = 'dn20160825.map' #a priori, for pressure profile
h2oam_fname = 'dn20160825.h2o_.sf' #contains H2O scaling factors
sat_fname = 'dn20160825_synrl.sat' #contains a (sparse) list of SZAs, AMs, and T's
ic_fname = ['870_nm.dn20160825.ic', '1020_nm.dn20160825.ic', '1640_nm.dn20160825.ic'] #window raw intensities, cell containing a list
settings_fname = 'i0_aod_settings.dat' 
# all_in = {'20160826/dn20160826.grl', '20160826/dn20160826.map', '20160826/dn20160826.h2o_.sf', 'i0_aod_settings.dat', '20160826/dn20160826_synrl.sat', {'20160826/870_nm.dn20160826.ic', '20160826/1640_nm.dn20160826.ic'}};
# rl_fname = all_in{1}; %runlog file
# daymap_fname = all_in{2}; %a priori 
# h2oam_fname = all_in{3}; %contains H2O scaling factors and AM
# settings_fname = all_in{4};
# sat_fname = all_in{5};
# ic_fname = all_in{6}; %window raw intensities


## version
v_num = 'calc_1win1day_aod   V: 2016-10-23   - JKH'  

#%% read in data
# read runlog
rl_raw = sra.read_runlog(rl_fname)

# read gamma H2O file
h2oam_raw = sra.read_h2ogam(h2oam_fname)

# read .ic files
ics_raw = [None] * len(ic_fname) # pre-allocate list
i1 = -1
for ic_filei1 in ic_fname: # read in each .ic file
    i1 += 1
    ics_raw[i1] = sra.read_ic(ic_filei1) #read in given ic file

# read .sat file
sat_raw = sra.read_sat(sat_fname)

# read the aod_settings file
set_raw = sra.read_aod_settings(settings_fname) 

#%% unique runlog times - get rid of duplicated values (e.g. for 2 detector instruments)
used_rl_flds = ['Spectrum_File_Name', 'time', 'fvsi', 'SNR', 'sia', 'Year', 'Day', 'Hour', 'pout', 'ASZA'] # fields needed in the runlog
rl_unq = rl_raw #default is regular runlog
if len(rl_raw['time']) > 1:
   spec1 = rl_raw['Spectrum_File_Name'][0].replace('b', 'a') #assuming spectra from detectors detectors are "a" and "b"
   spec2 = rl_raw['Spectrum_File_Name'][1].replace('b', 'a') #assuming spectra from detectors detectors are "a" and "b"
   if spec1 == spec2: #just comparing first 2 spectra for a test, will have same number if 2 detectors
       for fld_i1 in used_rl_flds:
           rl_unq[fld_i1] = rl_raw[fld_i1][::2] # get every other row


#%% calculate seasonal correction for sun-earth distance
lst_ones = np.ones([len(rl_unq['Year']),1])
lst_zeros = np.zeros([len(rl_unq['Year']),1])
j_days = rl_unq['time'] - egs.dt_from6vec(np.reshape(rl_unq['Year'], np.shape(lst_ones)),lst_ones,lst_ones,lst_zeros,lst_zeros,lst_zeros) + 1
s_corr = 1 + 0.033*np.cos(2*np.pi*(j_days-4)/365.25) # From Marenco and Brewer 2002

#%% interpolate the sparse .sat file to SZAs
sat_interp = {} #sat file interpolation

for hdr_i1 in sat_raw['headers']: #interpolating all values
    tck = interpolate.splrep(sat_raw['SZA'], sat_raw[hdr_i1], s=0) #initial curve fit to data
    sat_interp[hdr_i1] = interpolate.splev(rl_unq['ASZA'], tck, der=0) #spline interpolation

airmass = sat_interp['airmass'] #list of airmasses

#%% filter for clouds
cloud_uflg = {} #pre-initialiaze as dictionary

cld_param = ['fvsi', 'SNR', 'sia'] #list of parameters for cloud filtering
for cld_p_i1 in cld_param:
    cloud_uflg[cld_p_i1] = np.logical_and(rl_unq[cld_p_i1]>=set_raw[cld_p_i1][1], rl_unq[cld_p_i1]<=set_raw[cld_p_i1][2]) #filter for clouds based on specified settings

cloud_uflg['negI'] = np.ones([len(cloud_uflg['sia']), len(ics_raw)], dtype=bool) # flag for positive I values only


for i1 in np.arange(0,len(ics_raw)): #iterative loop, for each .ic file
    cloud_uflg['negI'][:,i1] = np.logical_and(cloud_uflg['negI'][:,i1], ics_raw[i1]['iavg']>0) # positive values required across all windows

cloud_uflg['all'] = cloud_uflg['fvsi'] & cloud_uflg['SNR'] & cloud_uflg['sia'] & np.all(cloud_uflg['negI'],1) # filter where 1="no clouds"
#cloud_uflg['all'][1] = False #test

filtered = {} #pre-initialize as dictionary
filtered['s_corr'] = s_corr[cloud_uflg['all']]
filtered['airmass'] = airmass[cloud_uflg['all']]
filtered['utc_time'] = rl_unq['time'][cloud_uflg['all']] #utc time for a running slope 
filtered['ln_avg_intens'] = egs.nan_ary([len(filtered['s_corr']), len(ics_raw)]) # pre-allocating, columns=wavelengths
for i1 in np.arange(0,len(ics_raw)): #iterative loop, for each .ic file
    filtered['ln_avg_intens'][:,i1] = np.log(ics_raw[i1]['iavg'][cloud_uflg['all']]/  \
      filtered['s_corr'].flatten()) # natural log of intensity, corrected for annual solar variance

fltr_sz = [len(filtered['utc_time']),1]
ufltr_sz = [len(cloud_uflg['all']),1]

#%% calculate running values
run_range = rc.running_vals(filtered['utc_time'], filtered['airmass'], ['range'], set_raw['deltime'][2])

filter_int = {} #pre-initialize as dictionary
filter_int['AM'] = (run_range['range']>=set_raw['delAM'][1]) & (run_range['range']<=set_raw['delAM'][2]) # filtering for minimum airmass change
filter_int['n_run'] = (run_range['n']>=set_raw['n_uf_run'][1]) & (run_range['n']<=set_raw['n_uf_run'][2]) # filtering for number of points


#%% loop over each ic file
run_mbr = [None]*len(ics_raw) #pre-allocating shape
tau_gas = [None]*len(ics_raw) # pafs, cell of structures of tau values for the gases
filenames = []

for i1 in np.arange(0,len(ics_raw)): #iterative loop, for each .ic file
    try:
        run_mbr[i1] = rc.running_corr_vals(filtered['utc_time'], np.reshape(filtered['airmass'], fltr_sz), np.reshape(filtered['ln_avg_intens'][:,i1], fltr_sz), ['linem', 'lineb', 'r'], set_raw['deltime'][2])
    except NameError: #except should be the correct error to actually skip
        pass
    
    ## filter the intercepts to get estimates of the values for I0 (and tau?)
    I0s_all = np.exp(run_mbr[i1]['lineb']) # listing of all I0 estimations
    filter_int['I0'] = (I0s_all>=set_raw['I0'][1]) & (I0s_all<=set_raw['I0'][2]) # filter for I0 values
    filter_int['rcorr'] = (run_mbr[i1]['r']>=set_raw['rcorr'][1]) & (run_mbr[i1]['r']<=set_raw['rcorr'][2])
    filter_int['all'] = filter_int['AM'] & filter_int['rcorr'] & filter_int['n_run'] & filter_int['I0']
    
    
    # Estimate a single I0 for the day 
    if sum(filter_int['all'])>set_raw['n_uf_day'][1]:
        I0s_int = I0s_all[filter_int['all']] #list of I0s to test statistics over
        langley_day = {} #pre-initialize as a dictionary
        langley_day['I0_avg'] = np.nanmean(I0s_int) #estimate of I0 for day - seems to bias tau high (I0 biased low)
        langley_day['I0_std'] = np.nanstd((I0s_int)) #standard deviation on I0
        langley_day['I0_med'] = np.nanmedian(I0s_int) #median estimate of I0 -
        langley_day['I0'] = langley_day['I0_med'] #median estimate of I0 -
        
        langley_day['slopes'] = -1*np.nanmean(run_mbr[i1]['linem'][filter_int['all']]) # estimate of tau_tot for day (likely a poor estimation)
        langley_day['slopes_std'] = np.nanstd(run_mbr[i1]['linem'][filter_int['all']])
    else: # not enough points in day for I0 calcu
        langley_day['I0_avg'] = -9
        langley_day['I0_std'] = -9
        langley_day['I0'] = -9

  
    #%% list of full-sized filters
    lst_ffltr = ['AM', 'rcorr', 'n_run', 'I0']
    for fltr_i1 in lst_ffltr:
        filter_int[fltr_i1 + '_full'] = np.ones([len(cloud_uflg['all'])], dtype=bool)
        filter_int[fltr_i1 + '_full'].flat[cloud_uflg['all']] = filter_int[fltr_i1]
    
    other_full = {} #pre-shape as dictionary of other filtered values
    other_full['r'] = np.zeros([len(cloud_uflg['all']), 1])
    other_full['I0'] = np.zeros([len(cloud_uflg['all']), 1])
    other_full['n'] = np.zeros([len(cloud_uflg['all']), 1])
    other_full['range'] = np.zeros([len(cloud_uflg['all']), 1])
    other_full['r'][cloud_uflg['all']] = run_mbr[i1]['r']
    other_full['I0'][cloud_uflg['all']] = I0s_all
    other_full['n'][cloud_uflg['all']] = run_range['n']
    other_full['range'][cloud_uflg['all']] = run_range['range']
    other_full['I0_999'] = other_full['I0'][:]
    other_full['I0_999'][other_full['I0_999']>999] = 0 # too large of values get scaled down
  
   
    #%% total tau calculations:
    full_tau = {} #pre-structure as dictionary
    full_tau['I0s'] = np.repeat(langley_day['I0'], len(rl_unq['time']))
    if langley_day['I0']>=0:    
        langley_day['tau_tot'] = np.log(np.transpose(s_corr)*ics_raw[i1]['iavg']/full_tau['I0s'])*-1/airmass #estimate of total optical depth for all time points
    else: # not enough points in day for I0
        langley_day['tau_tot'] = -9*np.ones(np.shape(full_tau['I0s']))
    

    #%% rayleigh scattering
    rs_raw = rOT.rayleigh_OT(ics_raw[i1]['nm_ctr_wdt'][0], rl_unq['pout'], daymap_fname) #takes a few seconds for a day. Doesn't account for width.


    #%% calculate tau due to various gases
    nm_strg = str(int(round(ics_raw[i1]['nm_ctr_wdt'][0]))) # string for the center wavelength
    
    mtch_sat_wn = ics_raw[i1]['nm_ctr_wdt'][0] == sat_raw['win_wn'] #logical index of where there are gases in the .sat file with T's
    if any(mtch_sat_wn):  # if some gases were involved
        tau_gas[i1] = {} #pre-structure as a dictionary
        fmtch_sat_wn = np.where(mtch_sat_wn) # convert logical to numerical
        for row_i2 in fmtch_sat_wn[0]: # transmittance of gases
            fld_hdr = sat_raw['T_name_gas'][row_i2, 0] # pull out the field name
            fld_spec = sat_raw['T_name_gas'][row_i2, 1] #pull out the field species
            if (fld_spec=='h2o') or (fld_spec=='hdo'): #for water species
                sf_i1 = h2oam_raw['gam_h2o'] #scaling factor for water
            else:
                sf_i1 = 1 #scaling for everything else (using the prior)
        
            # estimation of tau for gas (using the same airmass as for aerosols may introduce a small error
            i1_tau = -1*np.log(sat_interp[fld_hdr])*sf_i1/airmass;
            gas_name = 'tau' + fld_hdr[1:]
            
            tau_gas[i1][gas_name] = i1_tau
    
    
    #%% calculate OD due to just aerosols
    aod = langley_day['tau_tot'] - rs_raw # first subtract for Rayleigh scattering
    if any(mtch_sat_wn): # if there were gases that contributed to OD
        for key in tau_gas[i1]: # 
            aod = aod-tau_gas[i1][key] #subtract the gas OD
    
    
    #%% specify flags for I0
    flag_list = np.zeros([len(rl_unq['time']),1]) # pre-allocating with zeros for flags
    flag_list[~filter_int['AM_full'],0] = set_raw['delAM'][0] # flag for delAM
    flag_list[(flag_list==0) & np.reshape(~filter_int['rcorr_full'], ufltr_sz)] = set_raw['rcorr'][0] # flag for rcorr
    flag_list[(flag_list==0) & np.reshape(~filter_int['I0_full'], ufltr_sz)] = set_raw['I0'][0] # flag for I0
    flag_list[(flag_list==0) & np.reshape(~cloud_uflg['fvsi'], ufltr_sz)] = set_raw['fvsi'][0] # flag for fvsi
    flag_list[(flag_list==0) & np.reshape(~cloud_uflg['sia'], ufltr_sz)] = set_raw['sia'][0] # flag for sia
    flag_list[(flag_list==0) & np.reshape(~cloud_uflg['SNR'], ufltr_sz)] = set_raw['SNR'][0] # flag for SNR
    flag_list[(flag_list==0) & np.reshape(~filter_int['n_run_full'], ufltr_sz)] = set_raw['n_uf_run'][0] # flag for n_uf_run
    if sum(filter_int['all'])<set_raw['n_uf_day'][1]: # not enough points for day
        flag_list[flag_list==0] = set_raw['n_uf_day'][0]
    
    
    #%% write qc labels
    inqctext_out = set_raw['inqctext'][:]
    for i2 in np.arange(2,len(inqctext_out)):
       inqctext_out[i2] = '%(flg_num)3d%(nflgd)6d  %(origtxt)s' % {"flg_num": i2, \
         "nflgd": sum(flag_list==i2), "origtxt": inqctext_out[i2]}
    
    ## write total day n_used, n_flagged
    totalstext = ['Day Totals']
    totalstext.append('meanI0    ' + '%(a)12.9f' % {"a": langley_day['I0_avg']})
    totalstext.append('medianI0  ' + '%(a)12.9f' % {"a": langley_day['I0_med']})
    totalstext.append('stdI0     ' + '%(a)12.9f' % {"a": langley_day['I0_std']})
    totalstext.append('nused     ' + '%(a)6d' % {"a": np.sum(flag_list==0)})
    totalstext.append('nflagd    ' + '%(a)6d' % {"a": np.sum(flag_list!=0)})
    
    #%% outgoing columns
    out_header_pre = ['   Spectrum_File_Name                   ', ' Year', ' Day ', \
      ' Hour ', ' I0_flag ', ' rcorr ', '   delAM ', '   I0run_'+nm_strg, \
      ' n_run', '      I_'+nm_strg, '      I0_'+nm_strg, \
      ' tauTot_'+nm_strg, '   tauR_'+nm_strg, '   AOD_'+nm_strg] 
    out_header = out_header_pre[:] 
    if any(mtch_sat_wn): #if there were gases that contributed to OD
        for key in tau_gas[i1]: #list of gases names
            out_header = out_header + [' ' + key]
         
    
    
    outcols = [None]*len(rl_unq['time']) # pre-allocate outgoing list of text
    for i2 in np.arange(0,len(rl_unq['time'])): #for all rows
        # write the following: # spectrum file name, Year, Day, Hour, flag_number,
        # running r, delAM, running I0, running n, I, I0 (used), tau_total,
        # tau_Rayleigh, AOD
        outcols[i2] = '%(a)-40s %(b)4d %(c)3d %(d)7.4f %(e)6d %(f)6.3f %(g)8.5f \
          %(h)13.9f %(i)5d %(j)12.9f %(k)12.9f %(l)8.5f %(m)10.5f %(n)10.5f' % { \
          "a": rl_unq['Spectrum_File_Name'][i2], "b": rl_unq['Year'][i2], \
          "c": rl_unq['Day'][i2], "d": rl_unq['Hour'][i2], \
          "e": flag_list[i2], "f": other_full['r'][i2], \
          "g": other_full['range'][i2], "h": other_full['I0_999'][i2], \
          "i": other_full['n'][i2], "j": ics_raw[i1]['iavg'][i2], \
          "k": full_tau['I0s'][i2], "l": langley_day['tau_tot'][0,i2], \
          "m": rs_raw[0,i2], "n": aod[0,i2]}
        
        if any(mtch_sat_wn): # if there were gases that contributed to OD
            for key in tau_gas[i1]: #for each gas
                outcols[i2] = outcols[i2] + '%(z)9.5f' % {"z": tau_gas[i1][key][i2]} #tau other gases 
     
    #%% concatenate all output
    outtextall = ['%(a)4d %(b)4d\n' % {"a": 3+len(totalstext)+len(inqctext_out), "b": len(out_header)}] # length of headers. 
    outtextall.append(v_num + '\n') #version number
    for i2_row in inqctext_out:
        outtextall.append(i2_row)
    
    for i2_row in totalstext:
        outtextall.append(i2_row + '\n')
    
    outtextall.append(''.join(out_header) + '\n') #add in header list
    
    #%% write file including:
    # 1) header of flags, day mean I0, median I0, std I0, n_used, n_flagged
    # 2) spectrum, time, I0running, nrunning, I, I0, flag_number, r, delAM, tau_tot, tau_R,
    # tau_gases
    
    out_dir = os.path.dirname(ic_fname[i1]) #pull out directory name (same directory as .ic file)
    fname_woext = os.path.splitext(ic_fname[i1])[0] #filename without extension
    out_ffname = out_dir + fname_woext + '.raod' #full outgoing filename
    #%%
    fileID = open(out_ffname, 'w') #open file for writing
    
    for outline_i1 in outtextall: #over each header line
        fileID.write(outline_i1)
    
    for outline_i1 in outcols: #over each body line
        fileID.write(outline_i1 + '\n')
    
    fileID.close()
    
    filenames = filenames + [out_ffname]
#    import time
#    print("something")
#    time.sleep(15) 

#    return filenames
    


