# -*- coding: utf-8 -*-
"""
Created on Fri Oct 21 11:32:51 2016

v - 2016-10-22 

@author: Jacob Hedelius

Contains a collection of functions to calculate running values
"""

import os
import numpy as np
import warnings as warn
import findclosestdaily_n as fcdn
curr_dir = os.getcwd()
os.chdir('../../python_subs')
import egi_subs as egs
os.chdir(curr_dir)



def nancorrcoef(x_set, y_set ):
    # Calculates the correlation coefficient of data in matrix form with NaNs
        
    #example inputs
    #x_set = [[1, 2, 3], [4, 5, 6]]
    #y_set = [[2, 3.5, 6], [np.NaN, 10.5, 11]]
    # y_set = [[np.NaN, np.NaN, np.NaN], [np.NaN, np.NaN, np.NaN]]
    
    
    if np.size(x_set) != np.size(y_set):
        raise ValueError('Inputs are different sizes')
    
    x_ary = np.array(x_set) #convert to numpy array (in case input is a list)
    y_ary = np.array(y_set)
    
    x_linear = x_ary.flatten() #connverts matricies to vectors
    y_linear = y_ary.flatten()
    
    xy_combo = np.vstack((x_linear, y_linear)) #concatenates both vectors to 1 matrix
    xy_nonan_combo = xy_combo[:, ~np.any(np.isnan(xy_combo),0)] #removes nan values
    
    if xy_nonan_combo.size == 0: #too many NaNs
        out = []
    else:
        cc_xy = np.corrcoef(xy_nonan_combo[0,:], xy_nonan_combo[1,:])
        out = cc_xy[0][1]
    
    return out




#%%
def nanpolyfit(x,y,n):
    # This function runs polyfit, but excludes NaN values if they exist
    ##example inputs
        
    #x=[[1, 2, 3], [4, 5, 6]]
    #y = [[2, 3.5, 6], [np.NaN, 10.5, 11]]
    #n=1
    #y = [[np.NaN, np.NaN, np.NaN], [np.NaN, np.NaN, np.NaN]]
    
    
    if np.size(x) != np.size(y):
        raise ValueError('Sizes of x & y do not match')
    
    x_ary = np.array(x) #converts x to an array if needed
    y_ary = np.array(y) #converts y to an array if needed
    
    x_vec = x_ary.flatten() #connverts matricies to vectors
    y_vec = y_ary.flatten()
    
    xy_combo = np.vstack((x_vec, y_vec)) #concatenates both vectors to 1 matrix
    xy_nonan_combo = xy_combo[:, ~np.any(np.isnan(xy_combo),0)] #removes nan values
    
    if xy_nonan_combo.size == 0: #too many NaNs
        p_ary = []
    else:
        p_ary = np.polyfit(xy_nonan_combo[0,:], xy_nonan_combo[1,:],n) #polynomial fit
    
    return p_ary


#%%
def running_corr_vals(UTC_time, field_val1, field_val2, calc_types, bin_width_min):
    # This function was written to calculate the running value(s) of various
    # pairs of data (e.g. lines & slopes)
    # Inputs: 1) UTC_time should be a vector, 2) field_val1 and field_val2 
    # should be a vector that are the same size as UTC_time
    # 3) calc_types is a list of % the different calculations that can be run 
    # for a correlated running value. This input can be a single string, or a 
    # cell. Calculation types include:
    # linem - calculate the running slope of a linear fit line
    # lineb - calculate the running intercept of a linear fit line
    # r - linear correlation coefficient
    # 4) bin_width_min is the width of the running bin in minutes. Note: it is
    # assumed that there are no more than 11 values per minute.
        
    
    ## example inputs
#    UTC_time = egs.frange(egs.dt_from6vec(2015,1,1,10,0,0), egs.dt_from6vec(2015,1,1,15,0,0), (1/(60*24)))
    #field_val1 = np.random.random(len(UTC_time))
    #field_val2 = np.random.random(len(UTC_time))
#    calc_types = ['line', 'lineb', 'r'] #list of calculations to be run
#    bin_width_min = 10 #bin widths in minutes
    
    
    ## data check
    if not(egs.is_vec(field_val1)) or not(egs.is_vec(field_val2)) or \
          len(field_val1) != len(UTC_time) or len(field_val2) != len(UTC_time):
        raise ValueError('Variables field_vals1 and field_vals2 should be vectors the same length as UTC_time')
    
    field_val1 = field_val1.reshape(len(field_val1),1) #reshape to give consistent size input
    field_val2 = field_val2.reshape(len(field_val2),1) #reshape to give consistent size input
    
      
    n_tpoints = len(UTC_time)
    UTC_time_ary = np.array(UTC_time) #convert UTC_time to array
    UTC_time_ary = UTC_time_ary.flatten() #flatten array
    
    if isinstance(calc_types, (list)):
        # already a list, do nothing
        pass
    elif isinstance(calc_types, (str)):
        calc_types = [calc_types] #convert to a list
    else:
        raise ValueError('Calc_types needs to be a list or string of the possible calculation types')
    
    
    internal_ct_list = ['linem', 'lineb', 'r'] # internal list of all possible calculation types
    i1=-1 #initialize
    for type_i1 in calc_types:
        i1 += 1 #increment by 1
        if not(type_i1 in internal_ct_list):
            warn.warn(type_i1 + ' is not a recognized calculation type...skipping.')
            calc_types[i1] = '' #remove non-recognized calc types
    
    calc_types = [x for x in calc_types if x != ''] #filter out the unrecognized entries
    
    
    ## sometimes more outputs are calculated than what is called for. This subsection adds these in
    calc_types2 = calc_types #values calculated internally (includes ancillary types)
    if ('linem' in calc_types) | ('lineb' in calc_types):
        calc_types2 = calc_types2 + ['linem', 'lineb']
    
    calc_types2 = list(set(calc_types2)) #removes duplicates
    
    ## pre-allocate some values
    t_spacing = np.abs(np.diff(UTC_time_ary))*24*60 #time spacing in minutes
    t_spacing = t_spacing[t_spacing>0.0001] 
    min_t_space = np.min(t_spacing) #minimum time spacing in minutes - this is likely a slight underestimate in I2S data
    
    self_index = fcdn.findclosestdaily_n(UTC_time_ary, UTC_time_ary, int(np.floor(bin_width_min/min_t_space))) #index closest neighboring data points
    running_npoints = egs.nan_ary([n_tpoints, 1]) #pre-allocate, number of points in calculation
    
    pre_out={} #pre-allocate dictionary
    for type_i1 in calc_types2:
        pre_out[type_i1] = egs.nan_ary([len(field_val1), 1]) 
    
    
    ##
    for i1 in np.arange(0,n_tpoints): #iterative loop
        i1_indx_indx = np.abs(self_index['closest_ind_diff'][i1,:])<(bin_width_min/2)/(60*24) # within specified number of minutes of self
        running_npoints[i1] = np.sum(i1_indx_indx) #number of points in range
        
        i1_indx = np.int32(self_index['closest_ind'][i1, i1_indx_indx]) #index from main dataset
        i1_vals1 = field_val1[i1_indx, :] #pulls values close to current time point
        i1_vals2 = field_val2[i1_indx, :] #pulls values close to current time point
        
        for type_i2 in calc_types2:
            if type_i2 == 'r':
                pre_out[type_i2][i1] = nancorrcoef(i1_vals1, i1_vals2) # calculates correlation coefficient
            elif type_i2 == 'linem':
                linefit = nanpolyfit(i1_vals1, i1_vals2, 1) #calculate coefficients for best fit line
                pre_out['linem'][i1] = linefit[0] #saves the slope value
                pre_out['lineb'][i1] = linefit[1] #saves the intercept value
            elif type_i2 == 'lineb':
                pass #do nothing for now
    
    out = {} #pre-form an output dictionary
    out['n'] = running_npoints
    out['UTC_time'] = UTC_time
    for type_i1 in calc_types:
        out[type_i1] = pre_out[type_i1]
    
    return out


#%%
def running_vals(UTC_time, field_vals, calc_types, bin_width_min):
    
    # This function was written to calculate the running value(s) of various
    # fields.
    # Inputs: 1) UTC_time should be a vector, 2) field_vals should be a vector
    # or a matrix with one side equal to the UTC_time length, these are values
    # that the running product is calculated over, 3) calc_types is a list of
    # the different calculations that can be run for a running value. This 
    # input can be a single string, or a cell. Calculation types include:
    # min - calculate the running minimum
    # max - calculate the running maximum
    # range - calculate the running range
    # mean - calculate the running mean
    # median - calculate the running median
    # std - calculate the running std
    # 4) bin_width_min is the width of the running bin in minutes. Note: it is
    # assumed that there are no more than 11 values per minute.
    
    
    
    ## example inputs
    #UTC_time = egs.frange(egs.dt_from6vec(2015,1,1,10,0,0), egs.dt_from6vec(2015,1,1,15,0,0), (1/(60*24)))
    #field_vals = np.random.random(len(UTC_time)*2)
    #field_vals = field_vals.reshape(len(UTC_time),2)
    ##field_vals = egs.nan_ary([len(UTC_time),2])
    #calc_types = ['min', 'max', 'range', 'mean', 'median', 'std']
    #bin_width_min = 10
    #% UTC_time = filtered.utc_time;
    #% field_vals = filtered.airmass;
    #% calc_types = {'min', 'max', 'range', 'mean', 'median', 'std'};
    #% bin_width_min = 10;
    
    
    ## data check
    sfv = field_vals.shape #size of input
    if sfv[0] == len(UTC_time): #input correct format
        if len(sfv) == 1:
            field_vals = field_vals.reshape(len(UTC_time),1) #format to have 2 dimensions, even though the second is just 1
    elif sfv[1] == len(UTC_time):
        field_vals = np.transpose(field_vals)
    else:
        raise ValueError('UTC_time and spec_fields appear to have different lengths')
    
    sfv = np.int32(field_vals.shape) #new size
    sfv = sfv.reshape([len(sfv),1])
    
    n_tpoints = len(UTC_time)
    UTC_time_ary = np.array(UTC_time) #convert UTC_time to array
    UTC_time_ary = UTC_time_ary.flatten() #flatten array
    
    if isinstance(calc_types, (list)):
        # already a list, do nothing
        pass
    elif isinstance(calc_types, (str)):
        calc_types = [calc_types] #convert to a list
    else:
        raise ValueError('Calc_types needs to be a list or string of the possible calculation types')
    
    
    internal_ct_list = ['min', 'max', 'range', 'mean', 'median', 'std'] #internal list of calculation types
    i1=-1 #initialize
    for type_i1 in calc_types:
        i1 += 1 #increment by 1
        if not(type_i1 in internal_ct_list):
            warn.warn(type_i1 + ' is not a recognized calculation type...skipping.')
            calc_types[i1] = '' #remove non-recognized calc types
    
    calc_types = [x for x in calc_types if x != ''] #filter out the unrecognized entries
        
    calc_types2 = calc_types #values calculated internally (includes ancillary types)
    if ('range' in calc_types) :
        calc_types2 = calc_types2 + ['min', 'max']
    
    calc_types2 = list(set(calc_types2)) #removes duplicates
    
    ## t-spacing and pre-allocating outputs
    t_spacing = np.abs(np.diff(UTC_time_ary))*24*60 #time spacing in minutes
    t_spacing = t_spacing[t_spacing>0.0001] 
    min_t_space = np.min(t_spacing) #minimum time spacing in minutes - this is likely a slight underestimate in I2S data
    
    self_index = fcdn.findclosestdaily_n(UTC_time_ary, UTC_time_ary, int(np.floor(bin_width_min/min_t_space))) #index closest neighboring data points
    running_npoints = egs.nan_ary([n_tpoints, 1]) #pre-allocate, number of points in calculation
    
    pre_out={} #pre-allocate dictionary
    for type_i1 in calc_types2:
        pre_out[type_i1] = egs.nan_ary(sfv) 
    
    
    ## loop over all points
    for i1 in np.arange(0,n_tpoints): #iterative loop
        i1_indx_indx = np.abs(self_index['closest_ind_diff'][i1,:])<(bin_width_min/2)/(60*24) # within specified number of minutes of self
        running_npoints[i1] = np.sum(i1_indx_indx) #number of points in range
        
        i1_indx = np.int32(self_index['closest_ind'][i1, i1_indx_indx]) #index from main dataset
        i1_vals = field_vals[i1_indx, :] #pulls values close to current time point
        
        if ~np.all(np.isnan(i1_vals)): #requires that some values are not NaNs
            for type_i2 in calc_types2:
                if type_i2 == 'min':
                    pre_out[type_i2][i1] = np.nanmin(i1_vals,0) #calculates the running min
                elif type_i2 == 'max':
                    pre_out[type_i2][i1] = np.nanmax(i1_vals,0) #calculates the running max
                elif type_i2 == 'mean':
                    pre_out[type_i2][i1] = np.nanmean(i1_vals,0)
                elif type_i2 == 'median':
                    pre_out[type_i2][i1] = np.nanmedian(i1_vals,0)
                elif type_i2 == 'std':
                    pre_out[type_i2][i1] = np.nanstd(i1_vals,0)
                elif type_i2 == 'range':
                    pass #do nothing (for now)
    
    if 'range' in calc_types:
        pre_out['range'] = pre_out['max'] - pre_out['min']
    
    ## format the outputs
    out = {} #pre-form an output dictionary
    out['n'] = running_npoints
    out['UTC_time'] = UTC_time
    for type_i1 in calc_types:
        out[type_i1] = pre_out[type_i1]
    
    return out




