# -*- coding: utf-8 -*-
"""
Created on Fri Oct 26 11:44:34 2018

@author: Karolina Szabone Andre
"""
import numpy as np
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt

def CalculateLinefitGradients(data, npoints):
    """
    Calculates gradient of the data.
    
    Gradient data series is calculated using the following method:
        - grad[k] is the slope of the line fitted to points around a chosen datapoint (k)
        - the number of points used for line-fitting is:
            - given by the 'npoints' variable (in general)
            - at least the half of 'npoints' plus one (at the beginning 
              and the end of the data series)
    
    The length of the data series must be greater than 'npoints'.
    
    IN:
        - data: data series (numpy array or a list)
        - npoints: number of points to be used for fitting (integer)
                   it has to be an odd number and >= 3
    OUT:
        - grad: gradient of the data array described above (numpy array or
                a list depending on the input)


    See and run the function TestLinefitGradient() at the end of this module
    for examples.
    matplotlib is required only for running TestLinefitGradient()

    The code was developed under WinPython-32bit-2.7.6.4
    """
    
    # Auxiliary function
    def Line(x,a,b):
        return a*x+b
    
    # Checking whether the npoints is odd or not
    if (npoints%2==0):
        print "ERROR: 'npoints' is an even number, odd is needed"
        return []
    
    if (len(data)<npoints):
        print "ERROR: Length of the data serie is less than 'npoints'"
        return []
    
    # This variable is needed to define the start and end index 
    #  of the subset used for line fitting
    sub_add = int(npoints//2)
    
    # Initializing the gradient variable
    grad = []

    # Looping trough the data serie
    for itim in range(0,len(data)):
        # Determining the indexes of the data points to be used to calculate the actual gradient
        start = itim-sub_add
        end = itim+sub_add+1
        # Correcting the start index (at the beginning of the data serie)
        if (start <0):
            start = 0
        # Correcting the end index (at the end of the data serie)
        if (end > len(data)):
            end = len(data)
        
        # Cutting data
        vals = data[start:end]
        # If there are no missing value in the subset, fitting line 
        # otherwise appending nan to the grad list
        if (not np.isnan(np.min(vals))):
            x = np.array(range(0,len(vals)))
            popt, pcov = curve_fit(Line, x, vals)
            grad.append(popt[0])
        else:
            grad.append(np.nan)
    
    # Returning result using the type of the input data serie
    if (type(data) == type(np.array([]))):
        return np.array(grad)
    elif (type(data) == type([])):
        return grad
    
###############################################################################
def TestLinefitGradient(data_type, plot):
    """
    Tester subroutine
      - data_type: 'numpy array' or 'list
    """
    if (data_type == 'numpy array'):
        data = np.array([3,0,0,1,1,1,2,3,4,5,5,5,4,3,2,1])
    elif (data_type == 'list'):
        data = [3,0,0,1,1,1,2,3,4,5,5,5,4,3,2,1]
    npoints = 3
    gradient = CalculateLinefitGradients(data, npoints)
    print "Data type:", data_type
    print "Data:", data
    print "npoints:", npoints
    print "Grad:", gradient
    
    if plot:
        plt.figure()
        plt.plot(data,'b.-', label='data')
        plt.plot(gradient,'r.-', label='gradient')
        plt.axhline(y=0, color='grey')
        plt.legend()
    
    

###############################################################################

#TestLinefitGradient('numpy array', plot=True)
#TestLinefitGradient('list', plot=False)