#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon May the 4 be with you 2020

@author: JKH


The purpose of this function is to ingest *.oof and *.eof files and to average
them as specified (e.g., by time, SZA, or number) and write the results to a new
file

Possible future options: boxcar (moving average, maybe with some convolution
where points further away in time get deweighted)



Options: (how to specify "and" or "or"?)
all_flags = ['xco2_ppm', 'lt', 450; 'xco2_ppm_error', 'lt', 5]


"""

import numpy as np
import os
import copy
import pandas as pd
from itertools import compress
import time
import re


curr_dir = os.getcwd()
os.chdir('../python_subs')
import read_ggg_misc as rgggm
import egi_subs as egs
os.chdir(curr_dir)





# Example input:
#oof_ffname = '/home/osboxes/research/tccon/egi2020/results/cn_spec/daily/20140609/cn20140609.vav.ada.aia.oof'
#oof_ffname = '/home/osboxes/research/tccon/egi2020/results/xa/cumulative/xa20140603_20140721/xa20140603_20140721.vav.ada.aia.oof'


#  flag_list=['xco2(ppm)', 'lt', 450, 'and', 'flag', 'eq', 0, 'or', 'flag', 'ge', 70, 'or', 'flag', 'eq', 21] #testing!!!
#  avg_style=[['n', 2], ['time', 200], ['sza', 10]]




def avg_egi(oof_ffname, flag_list=['xco2(ppm)', 'lt', 450, 'and', 'flag', 'eq', 0, 'or', 'flag', 'ge', 70], avg_style=[['n', 2], ['time', 200]]):
  """ This is a helper function for EGI, it includes some defaults on the flags 
  and helps find the corresponding *.eof file. 
  
  On the default flag_list - a large flag may represent wrong *.mod files used
  """
  
  out_dir = os.path.dirname(oof_ffname)
  xxYYYYMMDD = os.path.basename(oof_ffname).split('.')[0]
  if oof_ffname[-4:] == '.oof':
    eof_ffname = oof_ffname[:-4] + '.eof'
  elif oof_ffname[-8:] == '.oof.csv':
    eof_ffname = oof_ffname[:-8] + '.eof.csv'
  else:
    raise ValueError("I don't understand extension on oof file: " + oof_ffname)
  
  if not os.path.exists(eof_ffname) or os.path.getsize(eof_ffname) < 100:
    # eof file exists, but is less than 100 bytes, likely was an error on writing.
    eof_ffname = None 
  
  
  if eof_ffname:
      out_oof, out_eof = avg_oof(oof_ffname, eof_ffname=eof_ffname, flag_list=flag_list, avg_style=avg_style)
      oof_fmts, eof_fmts = avgout_fmts(oof_ffname, eof_ffname=eof_ffname)
  else:
    out_oof = avg_oof(oof_ffname, eof_ffname=eof_ffname, flag_list=flag_list, avg_style=avg_style)
    oof_fmts = avgout_fmts(oof_ffname, eof_ffname=eof_ffname)
  
  ## Write averaged results to file
  for i1, i_avg in enumerate(avg_style):
    if i_avg[0] == 'n':
      a_str = 'n' + '%04d' % i_avg[1] 
    elif i_avg[0] == 'time': #input in HMMSS format
      a_str = 't' + '%05d' % i_avg[1]
    elif i_avg[0] == 'sza': #input in number of tenths of degrees
      a_str = 's' + '%03d' % i_avg[1]
      
    oof_out_name = out_dir + '/' + xxYYYYMMDD + '_avg.oof.' + a_str + '.csv'
    eof_out_name = out_dir + '/' + xxYYYYMMDD + '_avg.eof.' + a_str + '.csv'
    write_avg_csv(out_oof[i1], oof_out_name, oof_fmts)
    if eof_ffname:
      write_avg_csv(out_eof[i1], eof_out_name, eof_fmts)
    

  

#%%
def avg_oof(oof_ffname, eof_ffname=None, avg_style=[['n', 2]], flag_list=['flag', 'eq', 0]):
  """ 
  This function takes an *.oof file (and possibly an *.eof file) and averages 
  in the ways specified.
  avg_style options (pass as a list of list, with each output option as a 2 element sub-list):
    * n - number of sequential values to average. Skips/starts a new bin when the 
          time gap between missing files is larger than 10x the 2nd most frequent
          occurence
    * time - time in HMMSS format of bin sizes to average. E.g., 200=2 minutes
          Bins start/end at UTC midnight  
    * sza - solar zenith angle in tenths of degrees (e.g., 50 for bins of size 5 
          degrees). Bins start/end at 0 degrees in both directions (morning & afternoon)
  flag_list: pass as a long list with 'and' or 'or' dividing up groups of 3. Syntax 
    for the groups of 3 goes: 1) oof (or eof) field name as it is read into Python, 
    2) operator ('eq', 'ne', 'lt', 'gt', 'le', 'ge') and 3) numeric value
    When combining the 'or' works between all sequential statements. E.g.,
    ['xco2(ppm)', 'lt', 450, 'and', 'flag', 'eq', 0, 'or', 'flag', 'eq', 70, 'or', 
       'flag', 'eq', 71, 'and', 'xco2(ppm)_error', 'le', 5]
    would be interpreted
    'xco2(ppm)'<450 & ('flag'==0 | 'flag'==70 | 'flag'==71) & 'xco2(ppm)_error'<=5
    (If additional boolean options are needed later, could add in a list of lists 
    or just passing in a boolean flag directly...)
  
  A couple examples:
  flag_list=['xco2(ppm)', 'lt', 450, 'and', 'flag', 'eq', 0, 'or', 'flag', 'ge', 70, 'or', 'flag', 'eq', 21]
  avg_style=[['n', 2], ['time', 200], ['sza', 10]]
  """
  
  ######## Read in input files  ########
    
  oof_raw = rgggm.read_oof(oof_ffname)
  #oof_hdr = rgggm.parse_oof_comments(oof_raw['comments']) #parse the header comments
    
  
  if eof_ffname is not None:
    eof_raw = rgggm.read_eof(eof_ffname)
  else:
    eof_raw = None
  
  n_lns = len(oof_raw['flag']) #number of data lines
  
  ############# Start by making the boolean flag #################
  # using the same for *.oof and *.eof. The files *should* have the same number of
  # entries.
  comb_flags = egs.oof_bool_flag(oof_raw, flag_list, eof_raw)
  
  
  ############# Set bins indicies for each of the values ###############
  pseudo_hr = (oof_raw['year']-2000)*366 + oof_raw['day'] + oof_raw['hour']/24 #this is just to make a unique time field and may not have real meaning
  pseudo_min = (pseudo_hr - np.floor(np.min(pseudo_hr)))*24*60
  bin_idxs = np.zeros((len(oof_raw['flag']), len(avg_style) ), dtype=int) #pre-allocate
  for i_Avg in range(len(avg_style)):
    
    if avg_style[i_Avg][0] == 'time':
      t_space_min = (avg_style[i_Avg][1] % 100)/60 + (np.floor(avg_style[i_Avg][1]/100) % 100) + np.floor(avg_style[i_Avg][1]/10000)*60
      bin_idxs[:,i_Avg] = np.digitize(pseudo_min,  np.arange(0, np.max(pseudo_min)+120, t_space_min))
    
    elif avg_style[i_Avg][0] == 'n':
      t_diff = pseudo_min[1:]-pseudo_min[:-1]
      typ_spc, n_freq = np.unique(np.round(t_diff*60), return_counts=True) #typical spacing in seconds
      if len(typ_spc)>1:
        srt_idx = np.argsort(n_freq) #
        t_new = typ_spc[srt_idx[-2]]*60*10 #second to last value 
      else: #just one time spacing. (This may not be needed...)
        t_new = typ_spc*60*10
      
      i_bin = 1
      t_diff = np.append(0, t_diff)
      n_in_bin = 0
      for iPt in range(n_lns):
        if n_in_bin>=avg_style[i_Avg][1] or t_diff[iPt]>t_new:
          i_bin += 1 #advance bin index
          n_in_bin = 0 #reset to 0
        bin_idxs[iPt, i_Avg] = i_bin
        n_in_bin +=1 #add one
      
    elif avg_style[i_Avg][0] == 'sza':
      pseudo_day = np.floor(pseudo_hr/24 + oof_raw['long(deg)']/360) #split the local days
      is_morn = oof_raw['azim(deg)']<180 #index for morning values
      sza_bin = np.digitize(oof_raw['solzen(deg)'],  np.arange(0, 100, avg_style[i_Avg][1]/10)  )
      sza_bin[is_morn] = -1*sza_bin[is_morn]
      
  #    pseudo_day = np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2])
  #    is_morn = np.array(   [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 1, 1, 1])==1
  #    sza_bin = np.array(   [5, 5, 5, 4, 4, 3, 3, 4, 4, 8, 5, 6, 6, 6, 7, 7, 5, 5, 4, 4]) #for testing
      
      sA = pd.Series(np.arange(n_lns))
      bin_idxs[:, i_Avg] = np.array(sA.groupby([pseudo_day, sza_bin]).ngroup() )
     
      
    else:
      raise ValueError('Avg style ' + avg_style[i_Avg][0] + ' is not currently supported.')
  
  
  
  
  ########## Figure out field types and if there is an existing error field with it #########
  oof_fields = list(oof_raw.keys())
  if 'comments' in oof_fields:
    oof_fields.remove('comments')
  oof_out = [None]*len(avg_style)
  #oof_dtypes = [list if type(oof_raw[x]) is list else np.float64 for x in oof_fields ]  
  oof_dtypes = ['s' if type(oof_raw[x]) is list else 'n' for x in oof_fields ] 
  
  oof_f_werr = [None]*len(oof_fields) #over-allocate
  oof_df_cnames = ['n']
  for iFld, fld_name in enumerate(oof_fields): 
    if oof_dtypes[iFld] == 's':
      oof_f_werr[iFld] = fld_name
      oof_df_cnames.append(fld_name)
    else: #numeric
      if (len(fld_name)>6 and fld_name[-6:] == '_error') or (len(fld_name)>4 and fld_name[-4:] == '_err'):
        continue #if it's an error field, skip it (should be included with main field)
      
      if fld_name + '_error' in oof_fields:
        oof_f_werr[iFld] = [fld_name, fld_name + '_error']
        oof_df_cnames.extend(oof_f_werr[iFld])
      elif fld_name + '_err' in oof_fields:
        oof_f_werr[iFld] = [fld_name, fld_name + '_err']
        oof_df_cnames.extend(oof_f_werr[iFld])
      else:
        oof_f_werr[iFld] = [fld_name]
        oof_df_cnames.append(fld_name)
        oof_df_cnames.append(fld_name + '_std')
    
  oof_f_werr = [x for x in oof_f_werr if x is not None ] 
  
  
  
  eof_fields = []
  eof_out = [] 
  if eof_ffname is not None:
    eof_fields = list(eof_raw.keys())
    eof_out = [None]*len(avg_style)
  #  eof_dtypes = [list if type(eof_raw[x]) is list else np.float64 for x in eof_fields ] 
    eof_dtypes = ['s' if type(eof_raw[x]) is list else 'n' for x in eof_fields ] 
    
    eof_f_werr = [None]*len(eof_fields) #over-allocate
    eof_df_cnames = ['n']
    for iFld, fld_name in enumerate(eof_fields): 
      if eof_dtypes[iFld] == 's':
        eof_f_werr[iFld] = fld_name
        eof_df_cnames.append(fld_name)
      else: #numeric
        if (len(fld_name)>6 and fld_name[-6:] == '_error') or (len(fld_name)>4 and fld_name[-4:] == '_err'):
          continue #if it's an error field, skip it (should be included with main field)
        
        if fld_name + '_error' in eof_fields:
          eof_f_werr[iFld] = [fld_name, fld_name + '_error']
          eof_df_cnames.extend(eof_f_werr[iFld])
        elif fld_name + '_err' in eof_fields:
          eof_f_werr[iFld] = [fld_name, fld_name + '_err']
          eof_df_cnames.extend(eof_f_werr[iFld])
        else:
          eof_f_werr[iFld] = [fld_name]
          eof_df_cnames.append(fld_name)
          eof_df_cnames.append(fld_name + '_std')
    eof_f_werr = [x for x in eof_f_werr if x is not None ]
    
  
  
  ########## Average within the bin for each of the averaging types ##########
  for i_Avg in range(len(avg_style)):
    unq_idxs, n_cnts = np.unique(bin_idxs[comb_flags, i_Avg], return_counts=True) #unique bin indicies that are not completely flagged out
    
    # pre-allocate dataframes
    oof_out[i_Avg] = pd.DataFrame(index=np.arange(len(unq_idxs)), columns=oof_df_cnames)
    oof_out[i_Avg].loc[:,'n'] = n_cnts
    if eof_fields:
      eof_out[i_Avg] = pd.DataFrame(index=np.arange(len(unq_idxs)), columns=eof_df_cnames)
      eof_out[i_Avg].loc[:,'n'] = n_cnts
    
    for i_bin, i_unq in enumerate(unq_idxs):
      i_bool = np.all((comb_flags, bin_idxs[:, i_Avg]==i_unq), 0)
      
      for iFld, fld_name in enumerate(oof_f_werr):
        if type(fld_name) is not list: #for strings, just take the entry from the first part of the bin
          fld_subset = list(compress(oof_raw[fld_name], i_bool))
          oof_out[i_Avg].loc[i_bin, fld_name] = fld_subset[0]
        else: #numbers
          fld_subset = oof_raw[fld_name[0]][i_bool] #subset for current field
          if len(fld_name)==1: #no error field
            if np.all(np.isnan(fld_subset)): #all NaNs
              oof_out[i_Avg].loc[i_bin, fld_name[0]] = np.NaN
              oof_out[i_Avg].loc[i_bin, fld_name[0] + '_std'] = np.NaN
            else: #not all NaNs
              # TODO: might consider an exception for windspeed and direction...
              oof_out[i_Avg].loc[i_bin, fld_name[0]] = np.mean(fld_subset)
              oof_out[i_Avg].loc[i_bin, fld_name[0] + '_std'] = np.std(fld_subset)
            
          elif len(fld_name)==2: #error field included
            err_subset = oof_raw[fld_name[1]][i_bool]
            err_subset[err_subset==0] = np.NaN
            wgt_subset = 1/err_subset**2 #subset for error field
            wgt_mn, wgt_sem, n_pts = egs.nanwgt_sem(fld_subset, wgt_subset) #natively handles NaNs
            oof_out[i_Avg].loc[i_bin, fld_name[0]] = wgt_mn
            oof_out[i_Avg].loc[i_bin, fld_name[1]] = wgt_sem*np.sqrt(n_pts)
            
      ###### Do the same for *.eof file ######
      if eof_fields:
        for iFld, fld_name in enumerate(eof_f_werr):
          if type(fld_name) is not list: #for strings, just take the entry from the first part of the bin
            fld_subset = list(compress(eof_raw[fld_name], i_bool))
            eof_out[i_Avg].loc[i_bin, fld_name] = fld_subset[0]
          
          else: #numbers
            fld_subset = eof_raw[fld_name[0]][i_bool] #subset for current field
            
            if len(fld_name)==1: #no error field
              if np.all(np.isnan(fld_subset)): #all NaNs
                eof_out[i_Avg].loc[i_bin, fld_name[0]] = np.NaN
                eof_out[i_Avg].loc[i_bin, fld_name[0] + '_std'] = np.NaN
              else: #not all NaNs
                # TODO: might consider an exception for windspeed and direction...
                eof_out[i_Avg].loc[i_bin, fld_name[0]] = np.mean(fld_subset)
                eof_out[i_Avg].loc[i_bin, fld_name[0] + '_std'] = np.std(fld_subset)
            elif len(fld_name)==2: #error field included
              err_subset = eof_raw[fld_name[1]][i_bool]
              err_subset[err_subset==0] = np.NaN
              wgt_subset = 1/err_subset**2 #subset for error field
              wgt_mn, wgt_sem, n_pts = egs.nanwgt_sem(fld_subset, wgt_subset) #natively handles NaNs
              eof_out[i_Avg].loc[i_bin, fld_name[0]] = wgt_mn
              eof_out[i_Avg].loc[i_bin, fld_name[1]] = wgt_sem*np.sqrt(n_pts) #out error approximates std
  
  if not eof_out:
    return oof_out           
  else:
    return oof_out, eof_out



#%% Get formats of inputs and guess output formats
def avgout_fmts(oof_ffname, eof_ffname=None):
  """This function gets the input format of files and uses them to set the output
  formats. Input formats are based on number of decimal points (if it's a number)
  This will currently struggle if there are NaNs in the first line but not subsequent
  lines.
  """
    
  oof_raw = rgggm.read_oof(oof_ffname)
  with open(oof_ffname) as f:
    oof_nhc1 = f.readline()
    oof_nhc1 = np.int32(re.split('[ |,]+', oof_nhc1.strip() ))
    for i1 in range(oof_nhc1[0]):
      oof_line1 = f.readline()
  oof_line1 = re.split('[ |,]+', oof_line1.strip() )
    
  oof_fields = list(oof_raw.keys())
  if 'comments' in oof_fields:
    oof_fields.remove('comments')
  oof_dtypes = ['s' if type(oof_raw[x]) is list else 'n' for x in oof_fields ] 
  oof_infmt = [None]*len(oof_dtypes)
  oof_outfmt = ['%d']
  oof_outfields = ['n']
  for i1, fld_name in enumerate(oof_fields):
    if oof_dtypes[i1] == 's':
      oof_infmt[i1] = '%s'
      oof_outfmt.append(oof_infmt[i1])
      oof_outfields.append(oof_fields[i1])
      
    else: #number of some sort
      if 'E' in oof_line1[i1].upper():
        fld_i1_num = re.split('E', oof_line1[i1].upper() )[0]
        oof_infmt[i1] = 'E'
      else:
        fld_i1_num = copy.copy(oof_line1[i1])
        oof_infmt[i1] = 'f'
        
      fld_i1_dec = re.split('\.', fld_i1_num ) #split on decimal
      if len(fld_i1_dec)>1:
        oof_infmt[i1] = '%.' + '%d' % len(fld_i1_dec[1]) + oof_infmt[i1]
      else:
        oof_infmt[i1] = '%.0' + oof_infmt[i1]
      
  #    if fld_i1_num.upper() == 'NAN': #should maybe do something different if first line has NaN...
  #      raise ValueError('p')
      
      
      ##### Set output formats
      if (len(fld_name)>6 and fld_name[-6:] == '_error') or (len(fld_name)>4 and fld_name[-4:] == '_err'):
        continue #if it's an error field, skip it (should be included with main field)
      
      err_fldname = fld_name + '_std'
      if fld_name + '_error' in oof_fields:
        err_fldname = fld_name + '_error'
      elif fld_name + '_err' in oof_fields:
        err_fldname = fld_name + '_err'
      
      oof_outfmt.append(oof_infmt[i1])
      oof_outfields.append(fld_name)
      oof_outfmt.append(oof_infmt[i1])
      oof_outfields.append(err_fldname)
      
  oof_fmt_df = pd.Series(oof_outfmt, index=oof_outfields)
      
      
      
  
  if eof_ffname is not None:
    eof_raw = rgggm.read_eof(eof_ffname)
    with open(eof_ffname) as f:
      eof_nhc1 = f.readline()
      eof_nhc1 = np.int32(re.split('[ |,]+', eof_nhc1))
      for i1 in range(eof_nhc1[0]):
        eof_line1 = f.readline()
    eof_line1 = re.split('[ |,]+', eof_line1.strip() )
    
    eof_fields = list(eof_raw.keys())
    eof_dtypes = ['s' if type(eof_raw[x]) is list else 'n' for x in eof_fields ] 
    eof_infmt = [None]*len(eof_dtypes)
    eof_outfmt = ['%d']
    eof_outfields = ['n']
    for i1, fld_name in enumerate(eof_fields):
      if eof_dtypes[i1] == 's':
        eof_infmt[i1] = '%s'
        eof_outfmt.append(eof_infmt[i1])
        eof_outfields.append(eof_fields[i1])
        
      else:
        if 'E' in eof_line1[i1].upper():
          fld_i1_num = re.split('E', eof_line1[i1].upper() )[0]
          eof_infmt[i1] = 'E'
        else:
          fld_i1_num = copy.copy(eof_line1[i1])
          eof_infmt[i1] = 'f'
          
        fld_i1_dec = re.split('\.', fld_i1_num ) #split on decimal
        if len(fld_i1_dec)>1:
          eof_infmt[i1] = '%.' + '%d' % len(fld_i1_dec[1]) + eof_infmt[i1]
        else:
          eof_infmt[i1] = '%.0' + eof_infmt[i1]
        
  #    if fld_i1_num.upper() == 'NAN': #should maybe do something different if first line has NaN...
  #      raise ValueError('p')
        
        
        ##### Set output formats
        if (len(fld_name)>6 and fld_name[-6:] == '_error') or (len(fld_name)>4 and fld_name[-4:] == '_err'):
          continue #if it's an error field, skip it (should be included with main field)
        
        err_fldname = fld_name + '_std'
        if fld_name + '_error' in eof_fields:
          err_fldname = fld_name + '_error'
        elif fld_name + '_err' in eof_fields:
          err_fldname = fld_name + '_err'
        
        eof_outfmt.append(eof_infmt[i1])
        eof_outfields.append(fld_name)
        eof_outfmt.append(eof_infmt[i1])
        eof_outfields.append(err_fldname)
    
    
    eof_fmt_df = pd.Series(eof_outfmt, index=eof_outfields)
  
  
  if eof_ffname is None:
    return oof_fmt_df           
  else:
    return oof_fmt_df, eof_fmt_df


#%% Write averaged outputs
def write_avg_csv(in_df, out_fname, out_fmts):
  """
  This function takes an input dataframe and writes to output file out_fname
  using the formats in out_fmts (which is a series with fieldnames matching 
  those of in_df)
  
  in_df = out_eof[1]
  out_fmts = eof_fmts 
  """
  
  with open(out_fname, 'w') as fid:
    fid.write('2,' + '%d' % in_df.shape[1] + '\n') #number of header rows, and number of columns
    
    fld_names = list(in_df.keys())
    jn_fnames = ','.join(fld_names)
    fid.write('%s' % jn_fnames + '\n') #write headers
    
    # Make sure that out formats are in correct order (they likely already are)
    ordrd_fmts = [None]*len(fld_names)
    for i1, i_fld in enumerate(fld_names):
      ordrd_fmts[i1] = out_fmts.loc[i_fld]
    
    
    for iRow in range(in_df.shape[0]):
      str_list = [ordrd_fmts[k] % in_df.iloc[iRow,k] for k in range(len(fld_names))  ] #format strings for all fields
      fid.write('%s' % ','.join(str_list) + '\n') #write string to file
    

    






