#!/bin/python
# Filename egi_subs.py

# additional sub-functions to help with running egi
# v - Aug-17-2016  - 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 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
    
    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 = repsingle(yi,mx_len)
    mi = repsingle(mi,mx_len)
    di = repsingle(di,mx_len)
    hi = repsingle(hi,mx_len)
    Mni = repsingle(Mni,mx_len)
    si = 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


## 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."
    
    
    
