from pylab import *
import scipy.io
import mytools
from matplotlib.collections import PatchCollection
from os.path import exists
from scipy.ndimage import gaussian_filter1d

jitters = [0,2,5,10,20]
filenames = ['MMNs_2pm_sep_noISDIDD_limtau_Nperpop40_paramSD0.3_stimA125_105_gAMPA17.5_40.0_100.0_gNMDA5.827500000000001_13.32_33.300000000000004_gGABA10.0_10.0_jitter'+str(jitter)+'_dep1000_0.0_0.9_tau10.0_10.0_10.0_200.0_seed2.mat' for jitter in jitters]

Nperpop = 40
Nsamp = 8
def boxoff(ax,whichxoff='top'):
    ax.spines[whichxoff].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.get_xaxis().tick_bottom()
    ax.get_yaxis().tick_left()
def mybar(ax,x,y,facecolor=[],linewidth=0.1,w=0.4):
  qs = quantile(y, [0,0.25,0.5,0.75,1])
  polygon = Polygon(array([[x-w,x+w,x+w,x-w],[qs[1],qs[1],qs[3],qs[3]]]).T)
  p = PatchCollection([polygon], cmap=matplotlib.cm.jet)
  if type(facecolor) is not list or len(facecolor) > 0:
    p.set_facecolor(facecolor)
  p.set_edgecolor('#000000')
  p.set_linewidth(0.3)
  ax.add_collection(p)
  a2 = ax.plot([x-w,x+w,x,x,x-w,x+w,x,x,x-w,x+w],[qs[0],qs[0],qs[0],qs[2],qs[2],qs[2],qs[2],qs[4],qs[4],qs[4]],'k-',lw=linewidth)
  return [p,a2]

fig1, axs = subplots(2,4)
axarr = axs.reshape(prod(axs.shape),).tolist()
axnew = []
for iax in range(0,4):
  #axs[iax].set_position([0.08+0.24*iax, 0.78-0*0.2-0.12,0.18,0.2])
  axnew.append(fig1.add_axes([0.08+0.24*iax+0.15,0.78-1*0.2+0.055,0.03,0.05]))
  for iay in range(0,2):
    axs[iay,iax].set_position([0.08+0.24*iax, 0.78-iay*0.2,0.18,0.16])
for iax in range(0,len(axarr)):
  axarr[iax].tick_params(axis='both', which='major', labelsize=4)
  boxoff(axarr[iax])
  for axis in ['top','bottom','left','right']:
    axarr[iax].spines[axis].set_linewidth(0.2)
  axarr[iax].set_xlim([0,3800])
  axarr[iax].set_ylim([0,250])

for iax in range(0,len(axnew)):
  axnew[iax].tick_params(axis='both', which='major', labelsize=4)
  boxoff(axnew[iax])
  for axis in ['top','bottom','left','right']:
    axnew[iax].spines[axis].set_linewidth(0.2)

#axs[0,0].text(0,Nperpop+7,'Excitatory deviant detecting output (EO)',fontsize=4,ha='left',va='top',fontweight='bold')
#axs[1,0].text(0,Nperpop+7,'Excitatory population for standards (ES)',fontsize=4,ha='left',va='top')
#axs[2,0].text(0,Nperpop+7,'Inhibitory population for standards (IS)',fontsize=4,ha='left',va='top')
#axs[3,0].text(0,Nperpop+7,'Excitatory population for standards, delayed (ESD)',fontsize=4,ha='left',va='top')
#axs[4,0].text(0,Nperpop+7,'Excitatory population for deviants (ED)',fontsize=4,ha='left',va='top')
#axs[5,0].text(0,Nperpop+7,'Inhibitory population for deviants (ID)',fontsize=4,ha='left',va='top')
#axs[6,0].text(0,Nperpop+7,'Excitatory pop. for deviants, delayed (ESD)',fontsize=4,ha='left',va='top')
#axs[7,0].text(0,Nperpop+7,'Exc. timer pop. receiving phase-locked input (EP)',fontsize=4,ha='left',va='top')
#axs[8,0].text(0,Nperpop+7,'Exc. timer pop. receiving phase-locked input, alt. phase (EP2)',fontsize=4,ha='left',va='top')


MMNorder = [1,0,2,3]
axarr[MMNorder[0]].set_title('   Omission',fontsize=6,pad=12)
axarr[MMNorder[1]].set_title('Frequency deviant',fontsize=6,pad=12)
axarr[MMNorder[2]].set_title('Duration deviant',fontsize=6,pad=12)
axarr[MMNorder[3]].set_title('Inv. dur. deviant',fontsize=6,pad=12)
cols = ['#000000','#333300','#666600', '#999900', '#BBBB00']
#dimcols = ['#FFFF22','#CCCC55','#999999','#77CCCC','#55FFFF']
dimcols = ['#EEEEEE','#FFFFDD','#FFFFBB','#FFFF90','#FFFF66']
for ifile in range(0,len(filenames)):
  filename = filenames[ifile]
  yfiles = [200,150,100,50,0]

  print('Loading '+filename)
  A = scipy.io.loadmat(filename)
  for q in ['standard', 'deviant', 'pacemaker', 'pacemaker2', 'output', 'standardBoost', 'deviantBoost']:
    try:
      shp = A[q].shape
      for iy in range(0,shp[0]):
        for ix in range(0,shp[1]):
          if A[q][iy,ix].shape[0] == 1 and A[q][iy,ix].shape[1] > 1:
            A[q][iy,ix] = A[q][iy,ix][0]
    except:
      pass
  # Plotting the spikes for standardPopulationSpikeMonitor
  #stimvec = [stimulusStandard,stimulusPaceMaker,stimulusDeviant]
  #for stimind in [0,1,2]:
  #  thisstim = stimvec[stimind]
  #  lastval = 0
  #  dt = thisstim.dt*1000
  #  vals = thisstim.values
  #  for itime in range(0,len(vals)):
  #      axarr[0].plot([itime*dt,itime*dt,(itime+1)*dt],[lastval-150*stimind,vals[itime]-150*stimind,vals[itime]-150*stimind],lw=0.5)
  #      lastval = vals[itime]
  plotteds = []
  for iiMMN in range(0,4):
    polygon = Polygon(array([[0,3900,3900,0],[yfiles[ifile],yfiles[ifile],yfiles[ifile]+Nperpop,yfiles[ifile]+Nperpop]]).T)
    p = PatchCollection([polygon], cmap=matplotlib.cm.jet)
    p.set_facecolor(dimcols[ifile])
    p.set_edgecolor(None)
    axarr[iiMMN].add_collection(p)
    iMMN = MMNorder[iiMMN]
    plotteds_this = []
    axarr[iiMMN].plot(A['output'][iMMN,0], A['output'][iMMN,1]+yfiles[ifile], 'r.', lw=0.35, ms=0.35, mew=0.35, color=cols[ifile])
    plotteds_this.append(len(A['output'][iMMN,0]))

  axarr[0].text(-500,20+yfiles[ifile],'$\\theta$='+str(jitters[ifile])+' ms',rotation=0,ha='right',va='center',fontsize=5,fontweight='bold' if ifile==0 else 'normal')

standard_xs = [[0+x,400+x,400+x,450+x,450+x,500+x,500+x] for x in [0,500,1000,1500,2000,2500,3000]]
for iMMN in range(0,4):
    st_on = [1,1,1,1,0,1,1] if iMMN < 2 else ([1,1,1,1,2,1,1] if iMMN == 2 else [2,2,2,2,1,2,2])
    dev_on = [0,0,0,0,1,0,0] if iMMN == 0 else [0,0,0,0,0,0,0]

    standard_ys = [[0,0,1*(x>1),1*(x>1),1*(x>0),1*(x>0),0] for x in st_on]
    deviant_ys = [[0,0,1*(x>1),1*(x>1),1*(x>0),1*(x>0),0] for x in dev_on]
    standard_xs_this = [x for y in standard_xs for x in y]+[3800]
    standard_ys_this = [x for y in standard_ys for x in y]+[0]
    deviant_ys_this = [x for y in deviant_ys for x in y]+[0]

    axarr[iMMN].plot(standard_xs_this,[4*y+208 for y in standard_ys_this],'b-',lw=0.3,clip_on=False)
    axarr[iMMN].plot(standard_xs_this,[4*y+200 for y in deviant_ys_this],'r-',lw=0.3,clip_on=False)

pos = axarr[0].get_position()
fig1.text(pos.x0 - 0.05, pos.y1 + 0.02, 'D', fontsize=11)
pos = axarr[4].get_position()
fig1.text(pos.x0 - 0.07, pos.y1 - 0.01, 'E', fontsize=11)


axarr[0].plot([3000,3500],[20,20],'k-',lw=0.5)
axarr[0].text(3250,22,'500 ms',ha='center',va='bottom',fontsize=5)

for iax in range(0,4): 
  axarr[iax].set_yticks([])
  axarr[iax].set_xticks([])

for ax in axarr:
  for line in ax.yaxis.get_ticklines():
    line.set_markeredgewidth(0.3)
    line.set_markersize(2)
fig1.savefig('fig_rob_jitter.pdf')


#filenames = ['MMNs_2pm_sep_noISDIDD_limtau_Nperpop40_paramSD0.3_stimA125_105_gAMPA17.5_40.0_100.0_gNMDA5.827500000000001_13.32_33.300000000000004_gGABA10.0_10.0_jitter'+str(jitter)+'_dep1000_0.0_0.9_tau10.0_10.0_10.0_200.0_seed2.mat' for jitter in [0,10,20,40]]
goodOnes = ['MMNs_2pm_sep_noISDIDD_limtau_Nperpop40_paramSD0.3_stimA150_130_gAMPA17.5_30.0_80.0_gNMDA5.827500000000001_9.99_26.64_gGABA25.0_25.0_dep1000_0.0_0.95_tau10.0_10.0_10.0_250.0.mat',                           
            'MMNs_2pm_sep_noISDIDD_limtau_Nperpop40_paramSD0.3_stimA150_130_gAMPA17.5_30.0_80.0_gNMDA5.827500000000001_9.99_26.64_gGABA20.0_20.0_dep1000_0.0_0.9_tau10.0_10.0_10.0_250.0.mat',                              
            'MMNs_2pm_sep_noISDIDD_limtau_Nperpop40_paramSD0.3_stimA140_120_gAMPA17.5_30.0_80.0_gNMDA5.827500000000001_9.99_26.64_gGABA10.0_10.0_dep1000_0.0_0.9_tau10.0_10.0_10.0_250.0.mat',                             
            'MMNs_2pm_sep_noISDIDD_limtau_Nperpop40_paramSD0.3_stimA150_130_gAMPA17.5_30.0_80.0_gNMDA5.827500000000001_9.99_26.64_gGABA35.0_35.0_dep1000_0.0_0.9_tau10.0_10.0_10.0_250.0.mat',                             
            'MMNs_2pm_sep_noISDIDD_limtau_Nperpop40_paramSD0.3_stimA150_130_gAMPA17.5_30.0_80.0_gNMDA5.827500000000001_9.99_26.64_gGABA30.0_30.0_dep1000_0.0_0.9_tau10.0_10.0_10.0_250.0.mat',                             
            'MMNs_2pm_sep_noISDIDD_limtau_Nperpop40_paramSD0.3_stimA140_120_gAMPA17.5_30.0_80.0_gNMDA5.827500000000001_9.99_26.64_gGABA35.0_35.0_dep1000_0.0_0.9_tau10.0_10.0_10.0_250.0.mat',                              
            'MMNs_2pm_sep_noISDIDD_limtau_Nperpop40_paramSD0.3_stimA150_130_gAMPA17.5_30.0_80.0_gNMDA5.827500000000001_9.99_26.64_gGABA35.0_35.0_dep1000_0.0_0.95_tau10.0_10.0_10.0_250.0.mat',                             
            'MMNs_2pm_sep_noISDIDD_limtau_Nperpop40_paramSD0.3_stimA140_120_gAMPA17.5_30.0_80.0_gNMDA5.827500000000001_9.99_26.64_gGABA30.0_30.0_dep1000_0.0_0.9_tau10.0_10.0_10.0_250.0.mat',                              
            'MMNs_2pm_sep_noISDIDD_limtau_Nperpop40_paramSD0.3_stimA125_105_gAMPA17.5_40.0_140.0_gNMDA5.827500000000001_13.32_46.620000000000005_gGABA10.0_10.0_dep1000_0.0_0.9_tau10.0_10.0_10.0_250.0.mat',               
            'MMNs_2pm_sep_noISDIDD_limtau_Nperpop40_paramSD0.3_stimA125_105_gAMPA17.5_40.0_140.0_gNMDA5.827500000000001_13.32_46.620000000000005_gGABA25.0_25.0_dep1000_0.0_0.9_tau10.0_10.0_10.0_250.0.mat',              
            'MMNs_2pm_sep_noISDIDD_limtau_Nperpop40_paramSD0.3_stimA150_130_gAMPA17.5_30.0_80.0_gNMDA5.827500000000001_9.99_26.64_gGABA15.0_15.0_dep1000_0.0_0.9_tau10.0_10.0_10.0_250.0.mat',                             
            'MMNs_2pm_sep_noISDIDD_limtau_Nperpop40_paramSD0.3_stimA150_130_gAMPA17.5_30.0_80.0_gNMDA5.827500000000001_9.99_26.64_gGABA30.0_30.0_dep1000_0.0_0.95_tau10.0_10.0_10.0_250.0.mat',                            
            'MMNs_2pm_sep_noISDIDD_limtau_Nperpop40_paramSD0.3_stimA150_130_gAMPA17.5_30.0_80.0_gNMDA5.827500000000001_9.99_26.64_gGABA25.0_25.0_dep1000_0.0_0.9_tau10.0_10.0_10.0_250.0.mat',                            
            'MMNs_2pm_sep_noISDIDD_limtau_Nperpop40_paramSD0.3_stimA125_105_gAMPA17.5_40.0_100.0_gNMDA5.827500000000001_13.32_33.300000000000004_gGABA25.0_25.0_dep1000_0.0_0.9_tau10.0_10.0_10.0_200.0.mat',              
            'MMNs_2pm_sep_noISDIDD_limtau_Nperpop40_paramSD0.3_stimA150_130_gAMPA17.5_30.0_80.0_gNMDA5.827500000000001_9.99_26.64_gGABA15.0_15.0_dep1000_0.0_0.95_tau10.0_10.0_10.0_250.0.mat',                            
            'MMNs_2pm_sep_noISDIDD_limtau_Nperpop40_paramSD0.3_stimA125_105_gAMPA17.5_40.0_100.0_gNMDA5.827500000000001_13.32_33.300000000000004_gGABA10.0_10.0_dep1000_0.0_0.9_tau10.0_10.0_10.0_200.0.mat']              

it_target = 2300
its_nontarget = [800,1300,1800,2800,3300]

for iMMN in [0,1,2,3]:
  for itarget in [0,1]:
    x = 600*itarget
    polygon = Polygon(array([[x,x+500,x+500,x],[-1e5,-1e5,1e5,1e5]]).T)
    print('iMMN='+str(iMMN)+', x = '+str(x))
    p = PatchCollection([polygon], cmap=matplotlib.cm.jet)
    p.set_facecolor('#EEEEEE' if itarget == 1 else '#DDDDFF')
    p.set_edgecolor(None)
    axarr[4+iMMN].add_collection(p)


for ijitter in range(0,len(jitters)):
  myjitter = jitters[ijitter]
  Nspikes_target = []
  Nspikes_nontarget = []
  curves = []
  for imodel in range(0,len(goodOnes)):

   curves_thismodel = []
   Nspikes_target_thismodel = []
   Nspikes_nontarget_thismodel = []
   for myseed in range(1,1+Nsamp):
    filename = goodOnes[imodel]
    filename = filename.replace('_dep','_jitter'+str(myjitter)+'_dep')
    if myseed > 1:
      filename = filename.replace('.mat','_seed'+str(myseed)+'.mat')
    if not exists(filename):
      print(filename+' does not exist')
      continue
    A = scipy.io.loadmat(filename)

    curves_thisseed = []
    Nspikes_target_thisseed = []
    Nspikes_nontarget_thisseed = []
    for iMMNtype in range(0,4):
      A_iMMNtype_order = [1,0,2,3]
      spikes = A['output'][A_iMMNtype_order[iMMNtype]][0][0] 
      spikers = A['output'][A_iMMNtype_order[iMMNtype]][1][0]

      Nspikes_target_thisseed.append(len([1 for x in spikes if x >= it_target and x < it_target+500]))

      #print('iMMNtype = '+str(iMMNtype)+', standards, '+str(len([spikes[i]-it_target+iMMNtype*1400+0 for i in range(0,len(spikes)) if spikes[i] >= it_target and spikes[i] < it_target+500]))+' plotted')
      mysigma = 25 #25 ms std
      # Parameters
      n_samples = 4000
      spike_train = np.zeros(n_samples)
      rounded_spikes = np.round(spikes).astype(int)
      valid_spikes = rounded_spikes[(rounded_spikes >= 0) & (rounded_spikes < n_samples)]
      spike_train = np.bincount(valid_spikes, minlength=n_samples).astype(float)

      # Apply Gaussian smoothing (convolution)
      thiscurve = gaussian_filter1d(spike_train, sigma=mysigma)

      #Non-target (standards):
      Nspikes_nontarget_this = 0
      nontarget_curve = zeros([500])
      for it in its_nontarget:
        nontarget_curve = nontarget_curve +  thiscurve[it:it+500]
        Nspikes_nontarget_this = Nspikes_nontarget_this + len([1 for x in spikes if x >= it and x < it+500])

      nontarget_curve = nontarget_curve/len(its_nontarget) #normalize by the number of stimuli (in [900,1400,1900,2900,3400])
      Nspikes_nontarget_thisseed.append(Nspikes_nontarget_this/len(its_nontarget))
      curves_thisseed.append(thiscurve[:])
    curves_thismodel.append(curves_thisseed[:])
    Nspikes_target_thismodel.append(Nspikes_target_thisseed[:])
    Nspikes_nontarget_thismodel.append(Nspikes_nontarget_thisseed[:])
   Nspikes_target.append(Nspikes_target_thismodel[:])
   Nspikes_nontarget.append(Nspikes_nontarget_thismodel[:])
   mean_curves = mean(array(curves_thismodel),axis=0)
   #mean_curves = curves_thismodel
   curves.append(mean_curves[:])

  for iMMNtype in range(0,4):
    mean_curve = mean(array([curves[imodel][iMMNtype] for imodel in range(0,len(curves))]),axis=0)
    target_curve = mean_curve[it_target:it_target+500]
  
    nontarget_curve = zeros([500])
    for it in its_nontarget:
      nontarget_curve = nontarget_curve +  mean_curve[it:it+500]
    nontarget_curve = nontarget_curve/len(its_nontarget) #normalize by the number of stimuli (in [900,1400,1900,2900,3400])

    #axarr[4+iMMNtype].plot(range(0,500),target_curve,'-',lw=0.45,color=cols[ijitter],label='$\Delta t$='+str(myjitter)+' ms')
    axarr[4+iMMNtype].plot(range(0,500),target_curve,'-',lw=0.45,color=cols[ijitter],label='$\\theta$='+str(int(myjitter))+' ms')
    axarr[4+iMMNtype].plot(range(600,1100),nontarget_curve,'-',lw=0.45,color=cols[ijitter])

    Q=mybar(axnew[iMMNtype],ijitter,[mean([Nspikes_target[imodel][i][iMMNtype]/40/0.5-Nspikes_nontarget[imodel][i][iMMNtype]/40/0.5 for imodel in range(0,len(Nspikes_target))]) for i in range(0,len(Nspikes_target[0]))],facecolor=dimcols[ijitter],linewidth=0.3)
    Q[1][0].set_color(cols[ijitter])
    Q[0].set_edgecolor(cols[ijitter])
    Q[0].set_linewidth(0.3)
    
for iax in range(0,4): 
  axarr[iax].set_yticks([])
  axarr[4+iax].set_ylim([0,1.5])
  axarr[4+iax].set_xlim([0,1100])
  axarr[iax].set_xticks([])
  axarr[4+iax].set_xticks([])
  axnew[iax].set_facecolor('#EEEEEE')
  axnew[iax].set_xticks([])
  axnew[iax].set_ylim([-1,5.5])
  axnew[iax].set_ylabel('$f_{\mathrm{dd}}$',fontsize=5,rotation=0)
  axarr[4+iax].text(250,1.51,'Deviant',fontsize=5,ha='center',va='bottom',color='#8888FF')
  axarr[4+iax].text(850,1.51,'Standard',fontsize=5,ha='center',va='bottom',color='#909090')
axarr[0].plot([3000,3500],[20,20],'k-',lw=0.5)
axarr[0].text(3250,22,'500 ms',ha='center',va='bottom',fontsize=5)
axarr[5].plot([800,1050],[1.2,1.2],'k-',lw=0.5)
axarr[5].text(925,1.27,'250 ms',ha='center',va='bottom',fontsize=5)

axarr[4].legend(fontsize=5,ncol=2,handlelength=1,labelspacing=0.25,columnspacing=1.25)
axarr[4].set_ylabel('Firing rate (spikes/sec)',fontsize=4.5)
for ax in axnew+axarr:
  for line in ax.yaxis.get_ticklines():
    line.set_markeredgewidth(0.3)
    line.set_markersize(2)

fig1.savefig('fig_rob_jitter.pdf')