import numpy as np
from allensdk.ephys.extract_cell_features import get_square_stim_characteristics
from scipy import stats
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt

def MLIN(voltage, current, res, cap, dt, MAKE_PLOT=False, SHOW_PLOT=False, BLOCK=False, PUBLICATION_PLOT=False):
    '''voltage, current
    input:
        voltage: numpy array of voltage with test pulse cut out 
        current: numpy array of stimulus with test pulse cut out '''
    t = np.arange(0, len(current)) * dt
    (_, _, _, start_idx, end_idx) = get_square_stim_characteristics(current, t, no_test_pulse=True)
    stim_len = end_idx - start_idx

    distribution_start_ind=start_idx + int(.5/dt)
    distribution_end_ind=start_idx + stim_len
    
    v_section=voltage[distribution_start_ind:distribution_end_ind]
    if MAKE_PLOT:
        times=np.arange(0, len(voltage))*dt
        plt.figure(figsize=(15, 11))    
        plt.subplot2grid((7,2), (0,0), colspan=2)
        plt.plot(times[distribution_start_ind:distribution_end_ind], v_section)
        plt.title('voltage for histogram')

    print v_section
    v_section=v_section-np.mean(v_section)  
    var_of_section=np.var(v_section)
    sv_for_expsymm=np.std(v_section)/np.sqrt(2)
    subthreshold_long_square_voltage_distribution=stats.norm(loc=0, scale=np.sqrt(var_of_section))
    
    #--autocorrelation
    tau_4AC=res*cap
    AC=autocorr(v_section-np.mean(v_section)) 
    ACtime=np.arange(0,len(AC))*dt

    #--fit autocorrelation with decaying exponential    
    (popt, pcov)= curve_fit(exp_decay, ACtime, AC, p0=[AC[0],tau_4AC])
    tau_from_AC=popt[1]
    
    if MAKE_PLOT:        
        plt.subplot2grid((7,2), (1,0), rowspan=3)    
        plt.hist(v_section, bins=50, normed=True, label='data')
        data_grid=np.arange(min(v_section), max(v_section), abs(min(v_section)-max(v_section))/100.)
        plt.plot(data_grid, subthreshold_long_square_voltage_distribution.pdf(data_grid), 'r', label='gauss with\nmeasured var')
        plt.plot(data_grid, expsymm_pdf(data_grid, sv_for_expsymm), 'm', lw=3, label='expsymm function')
        plt.xlabel('voltage (mV)')
        plt.title('Mean subtracted voltage hist')
        plt.legend()
        
        #--cumulative density function
        (h, edges)=np.histogram(v_section, bins=50)        
        centers=find_bin_center(edges)

        CDFx=centers
        CDFy=np.cumsum(h)/float(len(v_section))

        plt.subplot2grid((7,2), (4,0), rowspan=3)
        plt.plot(CDFx, CDFy, label='data')
#        plt.plot(CDFx, sig(CDFx, popt[0], popt[1]), label='fit')
        plt.plot(data_grid, subthreshold_long_square_voltage_distribution.cdf(data_grid), 'r', label='gauss with\nmeasured var')
        plt.plot(data_grid, expsymm_cdf(data_grid, sv_for_expsymm), 'm', lw=3, label='expsymm func')
        plt.title('Normalized cumulative sum')
        plt.xlabel('v-mean(v)')
        plt.legend()
        
        plt.subplot2grid((7,2), (1,1), rowspan=3)
        plt.plot(ACtime, AC, label='data')
        plt.xlabel('shift (s)')
        plt.title('Auto correlation')
        plt.plot(ACtime, exp_decay(ACtime, AC[0], tau_4AC), label='RC')
        plt.plot(ACtime, exp_decay(ACtime, popt[0], tau_from_AC), label='fit')
        plt.legend()
    
        plt.tight_layout()
        if SHOW_PLOT:
            plt.show(block=BLOCK)  

        if PUBLICATION_PLOT:     
            times=np.arange(0, len(voltage))*dt
            plt.figure(figsize=(14, 7))    
            plt.subplot2grid((3,3), (0,0), colspan=3)
            plt.xlabel('time (s)', fontsize=14)
            plt.ylabel('(mV)', fontsize=14)
            plt.plot(times[distribution_start_ind:distribution_end_ind], v_section*1.e3)
            plt.title('Voltage for histogram', fontsize=16)
               
            plt.subplot2grid((3,3), (1,0), rowspan=2)    
            plt.hist(v_section*1.e3, bins=50, normed=True, label='data')
            data_grid=np.arange(min(v_section), max(v_section), abs(min(v_section)-max(v_section))/100.)
#                    plt.plot(data_grid, subthreshold_long_square_voltage_distribution.pdf(data_grid), 'r', label='gauss with\nmeasured var')
            plt.plot(data_grid*1.e3, 1.e-3*expsymm_pdf(data_grid, sv_for_expsymm), 'm', lw=3, label='expsymm ')
            plt.xlabel('voltage (mV)', fontsize=14)
            plt.title('Mean subtracted voltage hist', fontsize=16)
            plt.legend(loc=1)
            
            #--cumulative density function
            (h, edges)=np.histogram(v_section, bins=50)
            centers=find_bin_center(edges)

            CDFx=centers
            CDFy=np.cumsum(h)/float(len(v_section))

            plt.subplot2grid((3,3), (1,1), rowspan=2)
            plt.plot(CDFx*1e3, CDFy, label='data')
    #        plt.plot(CDFx, sig(CDFx, popt[0], popt[1]), label='fit')
#                    plt.plot(data_grid, subthreshold_long_square_voltage_distribution.cdf(data_grid), 'r', label='gauss with\nmeasured var')
            plt.plot(data_grid*1.e3, expsymm_cdf(data_grid, sv_for_expsymm), 'm', lw=3, label='expsymm')
            plt.title('Normalized cumulative sum',  fontsize=16)
            plt.xlabel('V-mean(V) (mV)', fontsize=16)
            plt.legend(loc=2, fontsize=14)
            
            plt.subplot2grid((3,3), (1,2), rowspan=2)
            plt.plot(ACtime, AC*1.e3, label='data')
            plt.xlabel('shift (s)', fontsize=14)
            plt.title('Auto correlation',  fontsize=16)
#                    plt.plot(ACtime, exp_decay(ACtime, AC[0], tau_4AC), label='RC')
            plt.plot(ACtime, exp_decay(ACtime, popt[0]*1.e3, tau_from_AC), lw=3, label='fit')
            plt.legend(loc=1)
            plt.tight_layout()
            
    return var_of_section, sv_for_expsymm, tau_from_AC

def expsymm_pdf(v, dv):
    return 1./(2.*dv)*np.exp(-np.absolute(v)/dv)

def expsymm_cdf(v, dv):
    return 1./2.+(v*(1-np.exp(-np.absolute(v)/dv)))/(2.*np.absolute(v))

def exp_decay(time, amp, tau):
    return amp*np.exp(-time/tau) 

def find_bin_center(edges):
    centers=np.zeros(len(edges)-1)
    for ii in range(0, len(edges)-1):
        centers[ii]=np.mean([edges[ii], edges[ii+1]])
    return centers

def autocorr(x):
    result = np.correlate(x, x, mode='full')
#    return result
    return result[result.size/2:]
