#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Apr 13  2020

@author: JKH

Reads misc GGG files


"""
import os
import numpy as np
import re
import egi_subs as egs
import pandas as pd
import shlex
import netCDF4 as nc4
import datetime as dt


#%%
def read_spt(in_fname):
  """ 
  Read the spectral fits files.
  Based on read_spt_fast.m by Debra
  """
  
  with open(in_fname, 'r') as fid:
    raw_lines = fid.read().splitlines() #quickly reads in full file
  
  n_hd_col = np.int32(raw_lines[0].split())
  
  stic_val = np.float64(raw_lines[1].split())
  headers = raw_lines[2].split()
  
  if len(headers) != n_hd_col[1]:
    raise ValueError('The number of headers, ' + '%d' % len(headers) + 
                     ' does not match the first line of the file, ' + '%d' % n_hd_col[1])
  
  # append 'mol_' to beginning of names to make them valid
#  headers = ['mol_' + x if x[0] in ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'] else x for x in headers ]

  dat_ary = np.NaN*np.zeros((len(raw_lines)-n_hd_col[0], len(headers))) #pre-allocate
  for i1, i_row in enumerate(raw_lines[n_hd_col[0]:]):
    dat_ary[i1, :] = np.float64(i_row.split())
  
  spt_fit={}
  for i1, i_hdr in enumerate(headers):
    spt_fit[i_hdr] = dat_ary[:, i1]
  
  spt_fit['Freq_min'] = stic_val[0]
  spt_fit['Freq_max'] = stic_val[1]
  spt_fit['npts'] = stic_val[2]
  #3 ??
  spt_fit['ASZA'] = stic_val[4]
  spt_fit['obalt'] = stic_val[5]
  spt_fit['p_alt'] = stic_val[6]
  spt_fit['rms'] = stic_val[7] #slightly different from my calc using Tm & Tc, I probably missed something...
#  spt_fit['intdz'] = stic_val[8] #???
#  spt_fit[??] = stic_val[9] # ???
  spt_fit['tot_cols'] = stic_val[10:]
  
  spt_fit['filename'] = os.path.basename(in_fname)
  
  return spt_fit


#%% read oof
def read_oof(oof_ffname):
  """ Reads in a GGG *.oof file into a dictionary. Converts number columns to 
  numpy arrays
  """
  
  with open(oof_ffname, 'r') as fid:
    raw_lines = fid.read().splitlines() #quickly reads in full file
      
  n_hd_col = np.int32(re.split('[ |,]+', raw_lines[0].strip() ))
  
  out = {'comments': raw_lines[1:n_hd_col[0]-1]}
  
  if  not raw_lines[n_hd_col[0]-1][0:23] in [' flag  spectrum  year  ', ' flag,spectrum,year,day']:
    raise ValueError('oof header is not where expected for file ' + oof_ffname)
  
  
  
  col_hdr = re.split('[ |,]+', raw_lines[n_hd_col[0]-1].strip())
  num_drows = len(raw_lines)-n_hd_col[0] #number of data rows
  
  dig_or_str = np.zeros([num_drows, len(col_hdr)], dtype=bool) #pre-allocating logic vector if values are strings or digits
  lst_by_elem = [[None]*num_drows for i1 in col_hdr]
  for i1, line_i1 in enumerate(raw_lines[n_hd_col[0]:]):
    line_i1_spt = re.split('[ |,]+', line_i1.strip()) #splits line based on delimiters
    
    for i2, field_i1 in enumerate(line_i1_spt): #checking if fields are numbers
      dig_or_str[i1, i2] = egs.is_number(field_i1) #0 if number, 1 otherwise
      
      lst_by_elem[i2][i1] = field_i1
  
  tot_col_ds = np.sum(dig_or_str, 0) #
  
  
  for i1, iField in enumerate(col_hdr):
    if tot_col_ds[i1] == num_drows: #all digits
      out[iField] = np.float64(lst_by_elem[i1]) #convert to numpy array
    else:
      out[iField] = lst_by_elem[i1]
    
  return out


#%% further parse the *.oof header, pull out other info that may be of interest
# this will be finiky to how the header is styled and may break then
def parse_oof_comments(hdr_lst):
  #hdr_lst = out['comments']
  
  
  hdr_prsd = {} #parsed header
  
  ###### Get version numbers
  list_version_strs = [' write_official_output_file   Version', ' apply_insitu_correction      Version', 
     ' apply_airmass_correction     Version', ' average_results          Version', 
     ' collate_results          Version', ' GFIT                     Version', 
     ' GSETUP                   Version']
  for iVer in list_version_strs:
    matchd_lines = [s for s in hdr_lst if iVer in s]
    if len(matchd_lines)==1:
      if not 'versions' in hdr_prsd:
        hdr_prsd['versions'] = {}
      
      line_i1_spt = re.split('[ |,]+', matchd_lines[0].strip())
      hdr_prsd['versions'][line_i1_spt[0]] = {'version': line_i1_spt[2], 
              'date': line_i1_spt[3], 'initials': line_i1_spt[4]}
  
  
  ####### ADCF and AICF
  idx_adcf = np.int64([k for k,s in enumerate(hdr_lst) if ' Airmass-Dependent Correction Factors:' in s])
  idx_aicf = np.int64([k for k,s in enumerate(hdr_lst) if 'Airmass-Independent/In-Situ Correction Factors:' in s])
#  end_aicf = np.int64([k for k,s in enumerate(hdr_lst) if 'format:' in s or 'format=' in s])
  adcf_rc = re.split('[ |,]+', hdr_lst[idx_adcf[0]].strip()) #number of rows and columns for adcf
  adcf_rc = [np.int32(adcf_rc[-2]), np.int32(adcf_rc[-1])]
  aicf_rc = re.split('[ |,]+', hdr_lst[idx_aicf[0]].strip()) #number of rows and columns for aicf
  aicf_rc = [np.int32(aicf_rc[-2]), np.int32(aicf_rc[-1])]
  
  
  hdr_prsd['ADCFs'] = {}
  for i1 in range(idx_adcf[0]+1, idx_adcf[0]+adcf_rc[0]+1):
    line_i1_spt = re.split('[ |,]+', hdr_lst[i1].strip())
    hdr_prsd['ADCFs'][line_i1_spt[0]] = np.float64(line_i1_spt[1:])
  
  
  hdr_prsd['AICFs'] = {}
  for i1 in range(idx_aicf[0]+1, idx_aicf[0]+aicf_rc[0]+1): #, end_aicf[0]-1):
    line_i1_spt = re.split('[ |,]+', hdr_lst[i1].strip())
    hdr_prsd['AICFs'][line_i1_spt[0]] = np.float64(line_i1_spt[1:])
  
  
  ####### EGI notes
  # This function (parse_oof_comments) is only used by EGI make plots. So no exception needs to be made now to skip this
  # EGI notes section if no EGI notes were found in the *.oof 
  idx_eginotes = np.int64([k for k,s in enumerate(hdr_lst) if '##  EGI Notes  ##' in s])
  idxs_pounds = np.int64([k for k,s in enumerate(hdr_lst) if '######' == s[0:6] ])
  egi_comments = hdr_lst[idx_eginotes[0]+1:idxs_pounds[np.where(idxs_pounds>idx_eginotes)[0][0]]] #subset for EGI notes section
  
  hdr_prsd['EGI'] = {}
  d_keys = ['system', 'name', 'date']
  mtch_lst = ['Retrievals were run on system: ', 'and were run by: ', '*.oof file annotations written on: ']
  for i1, i_str in enumerate(mtch_lst):
    mtch_ln = [s for s in egi_comments if i_str in s]
    trm_ln = mtch_ln[0][len(i_str):]
    hdr_prsd['EGI'][d_keys[i1]] = trm_ln
  
  
  mod_notes = [s for s in egi_comments if ' *.mod and *.vmr ' in s]
  hdr_prsd['EGI']['mod_notes'] = mod_notes[0]
  
  
  hdr_prsd['EGI']['overrides'] = []
  idx_ovrride = np.int64([k for k,s in enumerate(egi_comments) if 'Override file parameters:' in s])
  if len(idx_ovrride)>0: #only if override parameters were found
    i1 = 1
    i_str = egi_comments[idx_ovrride[0] + i1].strip()
    while i_str: #while string is not empty
      hdr_prsd['EGI']['overrides'].append([i_str])
      i1 += 1
      i_str = egi_comments[idx_ovrride[0] + i1].strip() #get the next string
  
  
  ####### List out flag limits
  idx_flag = np.int64([k for k,s in enumerate(hdr_lst) if ' # Variable' in s])
  df_flags = pd.DataFrame(columns=['Number', 'Variable', 'Output', 'Scale', 'Format', 'Unit', 'Vmin', 'Vmax', 'Description'])
  
  i1 = 1
  i_str = hdr_lst[idx_flag[-1] + i1]
  while i_str[0:4] != '----': #not the ends of the flags section...
    line_i1_spt = shlex.split(i_str.strip()) #split current line
  #  line_i1_spt = re.findall(r'(\w+|".*?")', i_str.strip())
  #  line_i1_spt = re.split('[ ]+', i_str.strip())
    df_flags = df_flags.append(pd.Series(), ignore_index=True) #append empty row
    df_flags.iloc[i1-1, 0:8] = line_i1_spt[0:8]
    df_flags.iloc[i1-1, -1] = " ".join(line_i1_spt[8:])
    
    i1 += 1
    i_str = hdr_lst[idx_flag[-1] + i1].strip() #get the next string
  
  df_flags['Number'] = df_flags['Number'].astype(int)
  df_flags['Output'] = df_flags['Output'].astype(int)
  df_flags['Scale'] = df_flags['Scale'].astype(float)
  df_flags['Vmin'] = df_flags['Vmin'].astype(float)
  df_flags['Vmax'] = df_flags['Vmax'].astype(float)
  
  hdr_prsd['flags'] = df_flags
  
  return hdr_prsd



#%% read eof files
def read_eof(eof_ffname):
  """ Reads in a GGG *.eof file into a dictionary. Converts number columns to 
  numpy arrays
  """
#eof_ffname  = '/home/osboxes/research/tccon/egi2020/results/cn_spec/daily/20140609/cn20140609.vav.ada.aia.eof.csv'

  hdr_comments = []
  with open(eof_ffname, 'r') as fid:
    raw_line1 = fid.readline() # reads in first line
    n_hd_col = np.int32(re.split('[ |,]+', raw_line1))
    for i1 in range(1, n_hd_col[0]-1):
      hdr_comments.append(fid.readline())
  
  
  # for reading in with Pandas, setting the option: low_memory=True may be useful
  # if files are large
  if ',' in raw_line1: #comma delimited
    raw_df = pd.read_csv(eof_ffname, header=n_hd_col[0]-1) #, low_memory=True)
  else:
    raw_df = pd.read_csv(eof_ffname, header=n_hd_col[0]-1, delim_whitespace=True) #, low_memory=True)
  
  
  ## convert to dictionary
  eof_dict = raw_df.to_dict('list')
  df_dtypes = raw_df.dtypes
  for i1, iField in enumerate(eof_dict.keys()):
    if df_dtypes[i1] == np.dtype('O'):
      pass
    elif df_dtypes[i1] == np.dtype('int64'):
      eof_dict[iField] = np.int64(eof_dict[iField])
    elif df_dtypes[i1] == np.dtype('float64'):
      eof_dict[iField] = np.float64(eof_dict[iField])
    else:
      # if you get this error, probably just need to add in some more elif statements
      raise ValueError('Unrecognized datatype for conversion to dict: ' + '%s' % df_dtypes[i1])
  
  return eof_dict


#%% read private *.nc files
def read_prv_nc(nc_ffname):
  """ Reads in GGG2020 private *.nc files into a dictionary.
  
  """

# nc_ffname = '/home/osboxes/research/tccon/egi2020/results/xa/daily/20140709/xa20140709_20140709.private.nc'

  ## initial calling of dataset
  tc_raw = nc4.Dataset(nc_ffname)
  
  ## process some variables
  tc_proc = {'null': 0} #initialize
  
  ## read in the info (site ID, software version, etc.)
  for attr in tc_raw.ncattrs():
     tc_proc[attr] = getattr(tc_raw, attr)
  
  tc_proc.pop('null', None) #delete the null key in dictionary
  
  ## save the variable attributes
  def var_atts(vname):
      vatts_i1 = {'null': None} #pre-initialize
      for attr in vname.ncattrs():
          vatts_i1[attr] = getattr(vname, attr)
      vatts_i1.pop('null', None) #delete the null key in dictionary
      return vatts_i1
  
  
  ## convert to np.array if possible
  def conv_ary(var_in):
      try:
          var_out = np.array(var_in) #attempt to convert to array
      except:
          var_out = var_in #if array conversion fails, just return the initial value
      return var_out
  
  
  ##
  tc_proc['vatts'] = {'null': 0} #initialize dictionary
  for var_nm in tc_raw.variables.keys():
      tc_proc['vatts'][var_nm] = var_atts(tc_raw[var_nm]) #pull out attribute names
      tc_proc[var_nm] = conv_ary(tc_raw[var_nm])
  
  diff_1970_1 = dt.datetime(1970,1,1) - dt.datetime(1,1,1) #raw format comes in fractional days since Jan 1, 1970
  
  if 'time' in tc_proc:
    tc_proc['utc_time'] =  diff_1970_1.days + tc_proc['time'] + 1 #NOTE: about 1 year different from Matlab's datetime. Add 366
  
  # NOTE: this conversion line takes ~9x as long to run as the rest of the function combined!
  #tc_proc['utc_dt'] = mpl.dates.num2date(tc_proc['utc_time']+1) #use matplotlib to convert back to datetime. 
  
  return tc_proc





















