#!/bin/python
# Filename egi_subs.py

# additional sub-functions to help with running egi
# v - Feb-11-2018  - JKH


import numpy as np
import re
#import datetime as dt


def frange(start, end=None, inc=None):
    "A range function, that does accept float increments..."

    if end == None:
        end = start + 0.0
        start = 0.0
    else: start += 0.0 # force it to be a float

    if inc == None:
        inc = 1.0

    count = int((end - start) / inc)
    if start + count * inc != end:
        # need to adjust the count.
        # AFAIKT, it always comes up one short.
        count += 1

    L = [None,] * count #pre-allocating for speed
    for i in range(count):
        L[i] = start + i * inc

    return L


def is_number(s):
    "Checks if a string is a number"
    
    try:
        float(s)
        return True
    except ValueError:
        return False

def is_vec(v):
    "Checks if an input is in a vector shape"
    
    import numpy as np
    if v.ndim==1 or np.min(v.shape)==1:
        return True
    else:
        return False
    

def isleapyr(year):
    import numpy as np
    if not(hasattr(year, '__iter__')): #for single row cases
        yeari = np.array([year])
    else:
        yeari = np.array(year)
    
    leap4_ind = np.mod(yeari,4)==0
    leap100_ind = np.mod(yeari,100)!=0
    leap400_ind = np.mod(yeari,400)==0
    leap_all = ((leap4_ind) & (leap100_ind)) | (leap400_ind)
    return leap_all

def vec_1D(x):                                                                                                                          
     #makes value a 1D vector (adds a dimension if needed)
     
     if np.shape(np.shape(x))[0]==1:
         out = np.transpose(x*np.ones([1,1]))
     else:
         out = x
      
     return out


def month_days(year,month,day):
#    year=np.array([2016])
#    month=np.array([3])
#    day=np.array([1])
        
    month[month==0] = 1 #month 0 is also January
    
    mdays=np.array([0, 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30])
    mdays_cum=np.cumsum(mdays) #cumulative sum of days
    mth_ind=month-1
    if np.size(mth_ind) == 1:
        mth_ind=int(np.squeeze(mth_ind))
    else:
        mth_ind=mth_ind.astype(np.int64)
    
    beg_month=mdays_cum[mth_ind]
    dayofmonth=day
    
    leap_ind = isleapyr(year)
    dayofyear_1 = beg_month+ dayofmonth #day of year minus 1
    dayofyear_1 = vec_1D(dayofyear_1) #adds a dimension, 1D vector
    
    aft_feb29=(dayofyear_1>=60) & (month>=3) #add additional day leap year (not needed when month input is)
    add_day_ind=((aft_feb29) & (leap_ind)) | ((month==3) & (day==1) & leap_ind) #after februrary 29th & a leap year
    zero_list = np.zeros([len(year),1])
    
    zero_list[add_day_ind]=zero_list[add_day_ind]+1
    dayofyear_1_lp=dayofyear_1+zero_list #day of year including leap years
    return dayofyear_1_lp

 
def mkarylst(x): 
    #make a value into an array list
    if not(hasattr(x, '__iter__')): #for single row cases
        xi = np.array([x])
    else:
        xi = np.array(x)
    return xi 

def repsingle(x,n):
    # repeat a single value if it is length 1 n times
    if len(x) == 1 and n!=1:
        out = np.repeat(x,n)
    else:
        out = x
    
    return out

     

def dt_from6vec(y,m,d,h,Mn,s):
    
    "creates a datetime array in days since Jan 1, 0 (no leap seconds)"
    yi = mkarylst(y)
    mi = mkarylst(m)
    di = mkarylst(d)
    hi = mkarylst(h)
    Mni = mkarylst(Mn)
    si = mkarylst(s)
    
    mx_len = np.max([len(yi), len(mi), len(di), len(hi), len(Mni), len(si)]) #get max list length
    
    yi = vec_1D(repsingle(yi,mx_len))
    mi = vec_1D(repsingle(mi,mx_len))
    di = vec_1D(repsingle(di,mx_len))
    hi = vec_1D(repsingle(hi,mx_len))
    Mni = vec_1D(repsingle(Mni,mx_len))
    si = vec_1D(repsingle(si,mx_len))
    
        
    if (len(yi)!=mx_len or len(mi)!=mx_len or len(di)!=mx_len or \
        len(hi)!=mx_len or len(Mni)!=mx_len or len(si)!=mx_len):
        raise ValueError('Improper array length!!  Non-matching')
    
    base_leapy=yi//4+1 #plus 1 for year zero
    base_nleap100=yi//100+1 #plus 1 for year zero
    base_leap400=yi//400+1 #plus 1 for year zero
    add_days=base_leapy-base_nleap100+base_leap400 #plus 1 because days start on 1 not 0
    
    year_days=yi*365+add_days-isleapyr(yi) #need to subtract 1 if _current_ year is still leap year
    
    
    days_md = month_days(yi,mi,di)
    days_ymd = days_md + year_days
    days_hms = hi/24 + Mni/(24*60) + si/(24*3600)
    days_ymdhms = days_ymd+days_hms
    days_srl = days_ymdhms #-719529
     
    return days_srl


def dt_to6vec(days_srl):
    "creates a datetime array into a 6 column vector given days since Jan 1, 0 (no leap seconds)"
    import numpy as np
    
    
    if not(hasattr(days_srl, '__iter__')): #for single row cases
        days_srli = np.array([days_srl])
    else:
        days_srli = np.array(days_srl)
    
    # make a vector of total running days for all years
    mx_day = np.max(days_srli)
    yrs_days = np.zeros(int(mx_day//365)+1) #pre-allocate, total number of days each year
    
    days_this_year = 367 #starting with year 0 (year 0 is 366+1 for proper indexing)
    run_days = 0 #running total number of days
    for year_i2 in range(len(yrs_days)):
        run_days = run_days + days_this_year
        if isleapyr(year_i2+1):
            days_this_year = 366
        else:
            days_this_year = 365
        
        yrs_days[year_i2] = run_days #running total number of days
        
    
    day_vec = np.zeros([len(days_srli),6]) #pre-allocating
    
    mdaysstd = [32, 60, 91, 121, 152, 182, 213, 244, 274, 305, 335, 366]
    mdaysstd_lr = [32, 61, 92, 122, 153, 183, 214, 245, 275, 306, 336, 367] #leap year
#    mdaysstd = [31, 59, 90, 120, 151, 181, 212, 243, 273, 304, 334, 365]
#    mdaysstd_lr = [31, 60, 91, 121, 152, 182, 213, 244, 274, 305, 335, 366] #leap year
    
    i1 = 0
    for time_i1 in days_srli:
        day_hms_i1 = time_i1 - time_i1//1 #finds floor value
        day_ymd_i1 = time_i1 - day_hms_i1
        day_ms_i1 = day_hms_i1*24 - day_hms_i1*24//1 #in hour FRACtions
        day_s_i1 = (day_ms_i1*60 - day_ms_i1*60//1)*60 #in seconds  
        day_h_i1 = day_hms_i1*24 - day_ms_i1
        day_m_i1 = day_ms_i1*60 - day_s_i1/60
        
        yr_diffs = day_ymd_i1 - yrs_days
        yr_i1 = np.argmax(yr_diffs<0)
        
        if yr_i1 > 0:
            days_remain = yr_diffs[yr_i1 - 1] + 1
        elif yr_i1 == 0:
            days_remain = day_ymd_i1
        else:
            raise ValueError('negative values not currently defined')
        
        if isleapyr(yr_i1):
            mth_diffs = days_remain - mdaysstd_lr
        else:
            mth_diffs = days_remain - mdaysstd
        
        mth_i1 = np.argmax(mth_diffs<0)
        
        if mth_i1 > 0:
            days_remain = mth_diffs[mth_i1 - 1] + 1
        
         
        day_vec[i1,0] = yr_i1
        day_vec[i1,1] = mth_i1 + 1 #add 1 to month index
        day_vec[i1,2] = days_remain
        day_vec[i1,3] = day_h_i1
        day_vec[i1,4] = day_m_i1
        day_vec[i1,5] = day_s_i1           
        
        i1 += 1
    
    return day_vec


def file_lines(fname):
    i=-1
    with open(fname) as f:
        for i, l in enumerate(f):
            pass
    return i + 1


def nanmean(in_ary):
    #finds the mean of an array, ignoring NaN values
    import numpy as np
    not_nan_ind = ~np.isnan(in_ary)   
    out = np.mean(in_ary[not_nan_ind])

    return out  

def nan_ary(shape, dtype=float): #returns a NaN array
    a = np.empty(shape, dtype)
    a.fill(np.nan)
    return a

def avg_wdir(wdirs_deg, wspds):
    wind_y = nanmean(wspds*np.sin(wdirs_deg*np.pi/180)) 
    wind_x = nanmean(wspds*np.cos(wdirs_deg*np.pi/180))
    J = 1j #complex number        
    wind_c = wind_x + wind_y*J
    wind_ang = np.mod(np.angle(wind_c)*180/np.pi, 360) #average angle of wind
    wspd_avg = np.sqrt(wind_y**2 + wind_x**2)
#    wind_a = 80
    return [wind_ang, wspd_avg] 


def read_sunrun(sr_filename):  #reads sunrun, but may be generic enough to read other files too
    fobj = open(sr_filename, 'r') # gets file object for reading
    file_lines = fobj.readlines() #reads in entire file all at once, line by line
    fobj.close()
    i0linep = re.split('[ |\n]+', file_lines[0]) #splits line
    i0line = list(filter(None, i0linep)) #remove empties   
    hd_lines = int(i0line[0]) #number of header lines
    hdrsp = re.split('[ |\n]+', file_lines[hd_lines-1]) 
    hdrs = list(filter(None, hdrsp)) #list of the headers
    flength = len(file_lines)-hd_lines #filelength exculding headers
    i1 = 0
    data_array = {'pre_al': 0} #pre-allocate file as dictionary
    for line_i1 in file_lines:
        if i1 >= hd_lines:
            line_i1_spt_p_2 = file_lines[hd_lines]
            line_i1_spt_p_2 = line_i1_spt_p_2.replace('-', ' -') #add in an extra space before "-" (helps with splitting)
            line_i1_spt_p = re.split('[ |\n]+', line_i1_spt_p_2)
            line_i1_spt = list(filter(None, line_i1_spt_p))
            if i1==hd_lines: #pre-allocate first time
                i2 = 0
                for var_hdr in hdrs: #loop over each header
                    data_array[var_hdr] = flength*[None]
                
            i2 = 0
            for var_hdr in hdrs:
                data_array[var_hdr][i1-hd_lines] = line_i1_spt[i2]
                i2 += 1
            
        i1 +=1
    del data_array['pre_al']
    return data_array


def reshape_lcols(mylist, ncols, nrows, fillvalue=None):
    # reshape a list, into a list of columns
    if len(mylist) != ncols*nrows:
        raise ValueError('reshape length of list not equal to ncols*nrows')
    newlist = [mylist[i::ncols] for i in range(ncols)]
    return newlist
    
def reshape_lrows(mylist, ncols, nrows, fillvalue=None):
    # reshape a list, into a list of rows
    if len(mylist) != ncols*nrows:
        raise ValueError('reshape length of list not equal to ncols*nrows')
    newlist = [mylist[i*ncols : (i+1)*ncols] for i in range(nrows)]
    return newlist

def listfind(mylist, matchlist):
    # find value in list, allowing "or"
    if type(matchlist) != list:
        matchlist = [matchlist]
        
    err_i1 = 0    
    for i1_search in matchlist:
        try:
            indx = mylist.index(i1_search)
        except ValueError:
            err_i1 = err_i1 + 1
    
    if err_i1 == len(matchlist): #if there was no match found
        indx = []
    return indx

def find(bool_vec):
    # acts like the matlab "find" option, creates a list from a boolean input array
    # of a numerical rather than logical index
    out = [x for xs in np.where(bool_vec) for x in xs]
    return out

def flt_or_str(item):
    "converts numbers to float, but leaves strings alone"
    try:
        return float(item)
    except ValueError:
        return item



def sunrise_noon_sunset(day_in_year, lon, lat):
    """Get approximate solar noon, and sunrise/sunset times. Probably only good 
    to 5-10 minutes.
    The ephem toolbox would be more accurate, but is not
    natively included in Anaconda Python.
    See: https://en.wikipedia.org/wiki/Position_of_the_Sun.
    day_in_year is the day of year. If using matlab time ('MLdate') you can get using: 
    day_in_year = (MLdate - 730486) % 365.24 #approximate date in year (730486 is for Jan 1, 2000)
    """
    
    utc_noon = -1*lon/15 + 12 #approximate hour for UTC solar noon
    solar_decl = -1*np.arcsin(0.39779 * np.cos(np.pi/180*(0.98565*(day_in_year+10) + 1.914*np.sin(0.98565*np.pi/180*(day_in_year-2))  ))  )
    hour_angle_rad = np.arccos(-1*np.tan(np.pi/180*lat)*np.squeeze(np.tan(solar_decl)) )
    hour_angle_hr = np.array(hour_angle_rad/(2*np.pi)*24)
    utc_srise = utc_noon - hour_angle_hr 
    utc_sset = utc_noon + hour_angle_hr 
    
    return [utc_srise, utc_noon, utc_sset]


def nanwgt_sem(in_dat, weights):
  """Finds the weighted average of the input data as well as the weighted SEM.
  This will work with NaNs."""
  # Note: functions like sum and mean are ~30% faster than their NaN couterparts like nansum and nanmean
  
  dw_mat = np.vstack((in_dat, weights)) #data & weights matrix
  dw_mat = dw_mat[:, ~np.any(np.isnan(dw_mat),0)] #remove rows with NaNs
  
  wgt_sum = np.sum(dw_mat[1,:]  ) #sum of the weights (for normalzation)
  n_pts = len(dw_mat[0,:]) #number of input datapoints
  
  
  if n_pts>1:
    wgt_mn = np.sum(dw_mat[0,:]*dw_mat[1,:])/wgt_sum #straightforward weighted average
    wgt_sem = np.sqrt(n_pts/(n_pts-1) * (1/wgt_sum**2) * np.sum((dw_mat[1,:]**2) * (dw_mat[0,:]-wgt_mn)**2) ) #weighted standard error of the mean. Empirically seems to agree with the bootstrap
  elif n_pts==1: #only one data point
    wgt_mn = dw_mat[0,0]
    wgt_sem = np.sqrt(1/dw_mat[1,0]) #assuming weights are 1/err^2
  else: #no points (all were NaNs in either in_dat or weights column)
    if np.any(~np.isnan(in_dat)): #if there are data in the main column
      wgt_mn = np.nanmean(in_dat)
      wgt_sem = np.nanstd(in_dat)/np.sqrt(np.sum(~np.isnan(in_dat)))
    else:
      wgt_mn = np.NaN
      wgt_sem = np.NaN
  
  return wgt_mn, wgt_sem, n_pts


def oof_bool_flag(oof_in, flag_list=['flag', 'eq', 0], eof_in=None):
  """
  oof_in: dictionary of read in oof file
  eof_in: dictionary of read in eof file (if using both eof and oof, they *should* 
    have the same number of entries)
  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]
  """
  ############# Start by making the boolean flag #################
  # using the same for *.oof and *.eof. The files *should* have the same number of
  # entries.
  if 'year' in oof_in:
    n_lns = len(oof_in['year']) #number of data lines
  elif 'Year' in oof_in:
    n_lns = len(oof_in['Year']) #number of data lines
  # else - add in other field that works for you
  
  avg_flags = np.ones((n_lns, (len(flag_list)-3)//4 + 1), dtype=bool) #pre-allocate
  
  n_ors = 0
  for i1 in range(np.shape(avg_flags)[1]):
    if i1>0 and flag_list[4*i1-1] == 'or':
      n_ors = n_ors+1
    
    set_of_3 = flag_list[4*i1:4*i1+3]
    if set_of_3[0] in oof_in:
      fld_vec = oof_in[set_of_3[0]]
    elif eof_in is not None and set_of_3[0] in eof_in:
      fld_vec = eof_in[set_of_3[0]]
    else: #field not found...
      continue #skip this loop iteration
    
    if set_of_3[1] == 'eq':
      avg_flags[:, i1] = fld_vec==set_of_3[2]
    elif set_of_3[1] == 'ne':
      avg_flags[:, i1] = fld_vec!=set_of_3[2]
    elif set_of_3[1] == 'lt':
      avg_flags[:, i1] = fld_vec<set_of_3[2]
    elif set_of_3[1] == 'gt':
      avg_flags[:, i1] = fld_vec>set_of_3[2]
    elif set_of_3[1] == 'le':
      avg_flags[:, i1] = fld_vec<=set_of_3[2]
    elif set_of_3[1] == 'ge':
      avg_flags[:, i1] = fld_vec>=set_of_3[2]
    else: #field not found
      if i1>0 and flag_list[4*i1-1] == 'or':
        avg_flags[:, i1] = False
      
  
  ## combine the flags
  comb_flags = np.zeros((n_lns, (len(flag_list)-3)//4 + 1 - n_ors), dtype=bool) #pre-allocate
  i_col = -1
  for i1 in range(np.shape(avg_flags)[1]):
    i_col += 1 #iterate column number by 1
    if i1>0 and flag_list[4*i1-1] == 'or':
      i_col -= 1 #backtrack for "or" statements
      comb_flags[:, i_col] = np.any((comb_flags[:, i_col], avg_flags[:, i1]), 0)
    else:
      comb_flags[:, i_col] = avg_flags[:, i1]
  
  out_flags = np.all(comb_flags, 1) #overall boolean flag
  
  return out_flags





## from this line, the next class and 3 functions are part of the itertools
## avaliable from: https://docs.python.org/3/library/itertools.html
#class ZipExhausted(Exception):
#    pass
#
#def izip_longest(*args, **kwds):
#    # zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-
#    fillvalue = kwds.get('fillvalue')
#    counter = len(args) - 1
#    def sentinel():
#        nonlocal counter
#        if not counter:
#            raise ZipExhausted
#        counter -= 1
#        yield fillvalue
#    fillers = repeat(fillvalue)
#    iterators = [chain(it, sentinel(), fillers) for it in args]
#    try:
#        while iterators:
#            yield tuple(map(next, iterators))
#    except ZipExhausted:
#        pass
#
#def repeat(object, times=None):
#    # repeat(10, 3) --> 10 10 10
#    if times is None:
#        while True:
#            yield object
#    else:
#        for i in range(times):
#            yield object
#
#def chain(*iterables):
#    # chain('ABC', 'DEF') --> A B C D E F
#    for it in iterables:
#        for element in it:
#            yield element

    
#def int_leadz2(in_num):
#    "takes an input array of integers and makes strings in list 2 digits long"
#    "pads with leading zeros as necessary.  Will make numbers integers as well."
    
    
    
