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

def format4(x):
    return format(x, ".4f").rstrip("0").rstrip(".")

f,axarr = subplots(1,1)
axarr.tick_params(axis='both', which='major', labelsize=4, direction='out', width=0.4, length=2)

axarr.set_position([0.09,0.8,0.86,0.1])

polygon = Polygon(array([[0,0,0+0.1,0+0.1],[0,200,200,0]]).T)
p = PatchCollection([polygon], cmap=matplotlib.cm.jet)
p.set_facecolor('#DDDDDD')
p.set_edgecolor(None)
axarr.add_collection(p)

filename = 'synfirefiles/stdpsynfire_synstim_N50_L70_T12.0_16.0_2.0Hz_gEE12.0_gEI20.0_gIE20.0_A0.0_0.0_0.03.mat'
A = scipy.io.loadmat(filename)

iN = filename.find('_N')
N = int(filename[iN+2:iN+2+filename[iN+2:].find('_')])
iL = filename.find('_L')
Lmax = int(filename[iL+2:iL+2+filename[iL+2:].find('_')])
iT = filename.find('_T')
stim_duration = float(filename[iT+2:iT+2+filename[iT+2:].find('_')])

Nskip = 10
for i in range(0, Lmax):  # plot every pool
  ts = A['spikes'][0][1+i][0][::Nskip]
  inds = i*N+0.95*A['spikes'][1][1+i][0][::Nskip]
  axarr.plot(A['spikes'][0][1+i][0][::Nskip], i*N+0.95*A['spikes'][1][1+i][0][::Nskip], 'ks', markersize=0.3,mew=0.3,lw=0.3)
axarr.set_xlabel('Time (s)',fontsize=6)
axarr.set_ylabel('Neuron ID',fontsize=6)
axarr.set_xlim([0,16])
axarr.set_ylim([0,3500])

axnew = []
pos = axarr.get_position()
axnew.append(f.add_axes([pos.x0+0.72,pos.y0+0.025,0.07,0.065]))

Nskip = 1
for i in range(0, Lmax):  # plot every pool
  ts = A['spikes'][0][1+i][0][::Nskip]
  inds = i*N+0.95*A['spikes'][1][1+i][0][::Nskip]
  axnew[0].plot([ts[i]*1000 for i in range(0,len(ts)) if 0<=ts[i]<=0.6], [inds[i] for i in range(0,len(ts)) if 0<=ts[i]<=0.6], 'ks', markersize=0.3,mew=0.3,lw=0.3)

for iax in range(0,len(axnew)):
  axnew[iax].tick_params(axis='both', which='major', labelsize=4, direction='out', width=0.4, length=2.0, pad=1.5)
  axnew[iax].set_xticks([0,25,50],['0','25 ms','50 ms'])
  
axnew[0].set_xlim([0.0,50])
axnew[0].set_ylim([0,200])
axnew[0].set_ylabel('Neuron ID',fontsize=4.5)

pos = axarr.get_position()
f.text(pos.x0 - 0.05, pos.y1 + 0.01, chr(ord('A')+0), fontsize=11)

f.savefig("figstdp0.pdf")