
def create_2tracefig(ylims):
    from pylab import figure
    fig1 = figure(num=1,figsize = (5,1.8), facecolor = 'w', dpi = 150, edgecolor = 'w')
    fig1.clf()
    ax1=axes([0.12,0.3,0.75,0.65])
    sns.despine()
    ax2 = ax1.twinx()
    ax2.spines["right"].set_edgecolor('r')
    ax1.set_ylabel('Rate [Hz]')
    ax1.set_xlabel('Time [s]')
    ax2.set_ylabel('Burst. Prob.')
    ax1.set_xticks([-1.5,-1,-.5,0,.5,1,1.5])
    ax2.tick_params(axis='y', colors='r')
    ax1.spines['top'].set_visible(False)
    ax2.spines['top'].set_visible(False)
    ax2.set_ylim((-0.05,.9))
    ax1.set_xlim((-1.4,1.4))
    ax1.set_ylim(ylims)
    return fig1,ax1,ax2
def train2binned(train,dt,smboxsize):
    if train.shape[0]==0:
        smrate = zeros(train.shape[1])
    else:
        rate = mean(train,axis=0)/dt
        nbins = floor(float(int(len(rate)))/smboxsize)
        smrate = rate[:int(nbins*smboxsize)].reshape((int(nbins),smboxsize))
        smrate = array([mean(smrate,axis=1)])
        smrate = smrate.repeat(smboxsize,axis=0)
        smrate = transpose(smrate).reshape(int(nbins*smboxsize),1)
        smrate = append(smrate,ones(len(rate)-int(nbins*smboxsize))*mean(rate[int(nbins*smboxsize):]))
    return smrate
def create_1tracefig(ylims):
    fig1 = figure(num=1,figsize = (5,1.8), facecolor = 'w', dpi = 150, edgecolor = 'w')
    fig1.clf()
    ax1=axes([0.12,0.3,0.75,0.65])
    sns.despine()
    ax1.set_ylabel('Rate [Hz]')
    ax1.set_xlabel('Time [s]')
    ax1.spines['top'].set_visible(False)
    ax1.set_ylim((-1,14))
    ax1.set_xticks([-1,-.5,0,.5,1])
    ax1.set_ylim(ylims)
    return fig1,ax1
def S2Btr(spktre,t,dt,npyr,BstTh):
    """
    transform set of spike trains into set of burst trains
    """
    if spktre.shape[0]==0:
        bstr = zeros((1,spktre.shape[1]))
        sstr = zeros((1,spktre.shape[1]))
    else:
        bstr = spktre*1.0
        sstr = spktre*1.0
        for ind in range(spktre.shape[0]):
            isise = array([1e3])
            spkse = where(spktre[ind,:])[0]
            isise = append(isise,diff(spkse)*dt)
            isise = append(isise,1e3)
            if len(spkse)>1:
                WhichAPinBst = where(logical_or(isise[1:]<BstTh,isise[:-1]<BstTh))[0] # spk is in bst if either isi after or isi before is smalle than bst isi-threshold
                Is1stinBst = logical_and(isise[1:]<BstTh,isise[:-1]>BstTh)
                sstr[ind,spkse[WhichAPinBst]] = sstr[ind,spkse[WhichAPinBst]]*0
                WhichAPnot1stBsty = where(Is1stinBst!=1)[0] # not first spike in burst isi after bursty but not before
                bstr[ind,spkse[WhichAPnot1stBsty]] = bstr[ind,spkse[WhichAPnot1stBsty]]*0 # 
            else:
                bstr[ind,spkse] = bstr[ind,spkse]*0

    return bstr,sstr


legroom1 = -3000 # half window around stim for deciding ON or OFF, T/2-legroom1:T/2+legroom2. So -3000:4000 means +0.3:0.4s
legroom2 = 4000
BstTh = 15e-3
smboxsize = 1001 # 100 ms binning
lim1=[0.12,0.3,0.75,0.65]
dt = 1e-4 # discretize in 0.1 ms time bins, dt is in seconds now
T = 4.0
t = arange(0,T,dt)
allhitreps=[]
L=int(T/dt)

# Loading data
cells = os.listdir(datapath)
ncells = len(cells)
firingcells = array([])
for i in range(ncells):
    if cells[i][-4:]=='.mat':
        WS = io.loadmat(datapath+cells[i])
        StimTimes = WS.get('MicrostimTimes')[:,0]
        if logical_and(len(StimTimes)>0,cells[i][-4:]=='.mat'): # only cells containing at least one stim
            firingcells=append(firingcells,i)
ncells = len(firingcells)

# prepare data arrays and counters
FratesNR_Hits = zeros((ncells,L))
FratesON_Hits = zeros((ncells,L))
FratesOFF_Hits = zeros((ncells,L))
BratesNR_Hits = zeros((ncells,L))
BratesON_Hits = zeros((ncells,L))
BratesOFF_Hits = zeros((ncells,L))
nHits = zeros((ncells,1))
nMissed = zeros((ncells,1))
nrep = zeros((ncells,1))
counton = 0
countoff = 0
countnr = 0
nON_Hits=0
nON_Missed=0
nNR_Hits=0
nNR_Missed=0
nOFF_Hits =0
nOFF_Missed=0
oninds =array([])
offinds = array([])

# process each cell type
for i in firingcells.astype('int'):
    WS = io.loadmat(datapath+cells[i])
    spks = WS.get('SpikeTimes')[:,0]
    StimTimes = WS.get('MicrostimTimes')[:,0]
    Hit = WS.get('Hit')[:,0]
    ind = where(firingcells==i)
    hitreps = where(Hit==1)[0]
    nHits[ind] = sum(Hit==1)
    nMissed[ind] = sum(Hit!=1)
    nrep[ind] = nHits[ind]+nMissed[ind]
    allhitreps = [allhitreps,hitreps]
    spktr = zeros((int(len(Hit)),int(T/dt)),dtype=dtype('bool') )
    for j in range(len(spks)):
    	dvec = abs(spks[j]-StimTimes)
    	mdvec = min(dvec)
    	rep = where(dvec==mdvec)[0][0]
    	if logical_and(mdvec<T/2,Hit[rep]==1):	########### Hit or Miss selection	
    		spktr[rep,int((spks[j]-StimTimes[rep]+T/2)/dt)]=True		
    	elif mdvec<T/2:
    		spktr[rep,int((spks[j]-StimTimes[rep]+T/2)/dt)]=True

    btr,sstr = S2Btr(spktr,t,dt,nrep[ind],BstTh)
    frate = train2binned(spktr,dt,smboxsize)
    brate = train2binned(btr,dt,smboxsize)
    spktrHits = spktr[hitreps,:]
    btrHits,sstrHits = S2Btr(spktrHits,t,dt,nHits[ind],BstTh) # separates bursts from single spikes based in ISI interval
    frateHits = train2binned(spktrHits,dt,smboxsize)
    brateHits = train2binned(btrHits,dt,smboxsize)

    e0 = mean(frateHits[:int((T/2-.02)/dt)])
    sde = std(frateHits[:int((T/2-.02)/dt)])
    close('all')
    if any(frateHits[int(T/2/dt)-legroom1:int(T/2/dt)+legroom2]>e0+nsigmaON*sde):	
        BratesON_Hits[counton,:]=brateHits*1.0
        FratesON_Hits[counton,:]=frateHits*1.0
        nON_Hits += nHits[ind][0][0]
        nON_Missed += nHits[ind][0][0]

        cellno=counton
        oninds = append(oninds,i)
        counton +=1
    elif any(frateHits[int(T/2/dt)-legroom1:int(T/2/dt)+legroom2]<e0-nsigmaOFF*sde): 
        BratesOFF_Hits[countoff,:]=brateHits*1.0
        FratesOFF_Hits[countoff,:]=frateHits*1.0
        nOFF_Hits += nHits[ind][0][0]
        nOFF_Missed += nHits[ind][0][0]
        cellno=countoff
        offinds= append(offinds,i)
        countoff +=1
    else:
        BratesNR_Hits[countnr,:]=brateHits*1.0
        FratesNR_Hits[countnr,:]=frateHits*1.0
        nNR_Hits += nHits[ind][0][0]
        nNR_Missed += nHits[ind][0][0]
        countnr +=1
        

# Plotting
close()
fig1,ax1=create_1tracefig((-.1,6))
mu = mean(BratesON_Hits[:counton,:],axis=0)
ax1.plot(t-T/2.0,mu/mean(mu[:20000]),'orange')
mu = mean(FratesON_Hits[:counton,:],axis=0)
ax1.plot(t-T/2.0,mu/mean(mu[:20000]),'green')
mu = mean(FratesOFF_Hits[:countoff,:],axis=0)
ax1.plot(t-T/2.0,mu/mean(mu[:19000]),'cyan')
mu = mean(FratesNR_Hits[:countnr,:],axis=0)
ax1.plot(t-T/2.0,mu/mean(mu[:19000]),'gray')
ax1.set_xlim((-1,2))
ax1.set_xticks([-1,0,1,2])
ax1.set_ylim((0,6 ))
ax1.set_ylabel('Norm. Rate')
fig1.savefig(task+'StimTrigTrialAvgRates_ONOFF_Normalized.pdf',format='pdf')

close()
fig1,ax1,ax2=create_2tracefig((-.1,6))
mu = mean(BratesON_Hits[:counton,:],axis=0)
ax2.plot(t-T/2.0,mu,'red')
mu = mean(FratesON_Hits[:counton,:],axis=0)
ax1.plot(t-T/2.0,mu,'green')
mu = mean(FratesOFF_Hits[:countoff,:],axis=0)
ax1.plot(t-T/2.0,mu,'cyan')
mu = mean(FratesNR_Hits[:countnr,:],axis=0)
ax1.plot(t-T/2.0,mu,'gray')
ax1.set_xlim((-1,2))
ax1.set_xticks([-1,0,1,2])
ax1.set_ylim((0,34 ))
ax2.set_ylim((0,4))
ax2.set_ylabel('Burst Rate [Hz]')
ax1.set_ylabel('Firing Rate [Hz]')
fig1.savefig(task+'StimTrigTrialAvgRates_ONOFF_Dual.pdf',format='pdf')

# Saving
downsampletime = arange(0,FratesON_Hits.shape[1],smboxsize-1)
time = downsampletime*dt*1e3 # time in ms
io.savemat(task+'PSTHs_3Classes_3sigOn03to04s',{'time':time,'BratesON_Hits':BratesON_Hits[:counton,downsampletime],'BratesNR_Hits':BratesNR_Hits[:countnr,downsampletime],'BratesOFF_Hits':BratesOFF_Hits[:countoff,downsampletime],'FratesON_Hits':FratesON_Hits[:counton,downsampletime],'FratesNR_Hits':FratesNR_Hits[:countnr,downsampletime],'FratesOFF_Hits':FratesOFF_Hits[:countoff,downsampletime]})



