'''Written by Corinne Teeter 3-31-14
'''

import matplotlib
import matplotlib.pyplot as plt
import numpy as np


def checkPreprocess(originalStim_list, processedStim_list, originalVoltage_list, processedVoltage_list, config, blockME=False):
        
    timeOriginal=np.arange(len(np.concatenate(originalStim_list)))*config.neuron['dt']
    if 'subSample' in config.dictOfPreprocessMethods.keys():
        timeProcessed=np.arange(len(np.concatenate(processedStim_list)))*config.dictOfPreprocessMethods['subSample']['desired_time_step']
    else:
        timeProcessed=timeOriginal
        
    plt.figure(figsize=(20,10))
    plt.subplot(4,1,1)
    plt.plot(timeOriginal, np.concatenate(originalStim_list), 'b')
    plt.title('original stimulation')
    
    plt.subplot(4,1,2)
    plt.plot(timeProcessed, np.concatenate(processedStim_list), 'r')
    plt.title('processed stimulation')        

    plt.subplot(4,1,3)
    plt.plot(timeOriginal, np.concatenate(originalVoltage_list), 'b')
    plt.title('original voltage')
    
    plt.subplot(4,1,4)
    plt.plot(timeProcessed, np.concatenate(processedVoltage_list), 'r')
    plt.title('processed voltage') 
    
    plt.annotate(config.cellName+': View result of preprocessing', xy=(.4, .975),
         xycoords='figure fraction',
         horizontalalignment='left', verticalalignment='top',
         fontsize=20)

#    plt.show(block=blockME)

def plotSpikes(voltage_list, spike_ind_list, dt, blockME=False, method=False):
    
    converted_spike_ind_list=[]
    time=np.arange(len(np.concatenate(voltage_list)))*dt
    #--find the length of each vector
    thelength=0
    for ii, voltage in enumerate(voltage_list):
        converted_spike_ind_list.append(spike_ind_list[ii]+thelength)
        thelength=thelength+len(voltage)
    
    subsampled_time=[time[ii] for ii in np.concatenate(converted_spike_ind_list)]

    plt.figure(figsize=(20, 5))
    plt.plot(time, np.concatenate(voltage_list), 'b')
    plt.plot(subsampled_time, [np.concatenate(voltage_list)[ii] for ii in np.concatenate(converted_spike_ind_list)], 'r.', ms=16)
    if method==False:
        plt.title('Spikes')
    else:
        plt.title('Spikes. Method used: '+method)
    
    plt.ylabel('voltage (V)')
    plt.xlabel('time (s)')
#    plt.show(block=blockME)
    
def checkSpikeCutting(originalStim_list, cutStim_list, originalVoltage_list, cutVoltage_list, allindOfNonSpiking_list, config, blockME=False):
    
    if len(originalStim_list)!=len(cutStim_list) or  \
        len(originalStim_list)!=len(originalVoltage_list) or \
        len(originalStim_list)!=len(cutVoltage_list) or \
        len(originalStim_list)!=len(allindOfNonSpiking_list):
        raise Exception('lists are not the same length')
    
        
    lengthGoingToAdd=0
    whole_ind=np.array([])
    whole_v=np.array([])
    for trace, ind_array in zip(originalVoltage_list, allindOfNonSpiking_list):
        ind=np.arange(0, len(trace))+lengthGoingToAdd
        whole_ind=np.append(whole_ind, [ind[ii] for ii in ind_array])
        whole_v=np.append(whole_v, [trace[ii] for ii in ind_array])
        lengthGoingToAdd=lengthGoingToAdd+len(trace)
             
             
    time=np.arange(len(np.concatenate(originalStim_list)))*config.neuron['dt']        
    plt.figure(figsize=(20,10))
    
    plt.subplot(1,1,1)
    plt.plot(time, np.concatenate(originalVoltage_list))
    plt.title('voltage')
    
    plt.plot(whole_ind*config.neuron['dt'], whole_v, '--r', lw=2)
    
    plt.annotate(config.cellName+': check spike cutting', xy=(.4, .975),
                 xycoords='figure fraction',
                 horizontalalignment='left', verticalalignment='top',
                 fontsize=20)
#    plt.show(block=blockME)

def plotLineRegress1(slope, intercept, r,xlim):
    y=slope*xlim+intercept
    print 'slope=', slope, 'intercept=', intercept, 'xlim', xlim
    plt.plot(xlim, y, '-k', lw=4, label='slope='+"%.2f"%slope+', intercept='+"%.3f"%intercept+', r='+"%.2f"%r)

def plotLineRegressRed(slope, intercept, r,xlim):
    y=slope*xlim+intercept
    print 'slope=', slope, 'intercept=', intercept, 'xlim', xlim
    plt.plot(xlim, y, '-r', lw=4, label='slope='+"%.2f"%slope+', intercept='+"%.3f"%intercept+', r='+"%.2f"%r)
