#!/bin/env python
# -*- coding: utf-8 -*-

from brian2 import *
import random
import sys

if(len(sys.argv) != 2):
    print("Usage : ", sys.argv[0] , " shuffle(0,1) ")
    sys.exit(-1)

shuffle = int(sys.argv[1])

tau = 20 * msecond        # constant de temps
Vt = -55 * mvolt          # seuil pour spiker
Vr = -75 * mvolt          # potentiel de reset
Vs = -65 * mvolt          # potentiel de repos
dt = defaultclock.dt


# Pas de couplage
g_weight = 0 * mvolt

# Couplage excitateur
g_weight = 0.7 * mvolt


print("Poids du couplage" ,g_weight)

I0 = 40 * mV
t_init = [20*ms, 70* ms,  125 * ms]

# On fixe la décroissance du premier stimulus à -0.05 mvolt/msecond
alpha = [-0.05*mvolt/msecond] +  ([0] * (len(t_init)-1))
t_cross = 200 * msecond
for i,(t1,a1) in enumerate(zip (t_init[1:],alpha[1:])):
    alpha[i+1] = alpha[0] * (t_cross - t_init[0])/(t_cross - t1)

# On shuffle la liste des temps
if(shuffle):
    random.shuffle(t_init)
print(t_init)

tmax = 800*msecond


# Notre groupe contient 3 neurones intègre et tire
eqs = '''
dV/dt = (-(V-Vs) + I)/tau : volt
I : volt
'''

N = len(t_init)
Ga = NeuronGroup(N=N, model=eqs,
                 threshold='V > Vt', reset='V = Vr', method='linear')

S = Synapses(Ga, Ga, 'w : volt', on_pre='V_post += w')
S.connect(condition='i!=j')
S.w = g_weight


# On se définit les entrées
inputs = zeros((N, (int)(tmax/dt)+1))
for i in range(N):
    inputs[i,:] = 0
    ti = t_init[i]
    for j in range((int)(ti/dt), (int)(tmax/dt)):
        t = j * dt
        inputs[i,j] = I0 + (t-ti) * alpha[i]
        if(inputs[i,j] <= 0):
            inputs[i,j] = 0.0

# On monitore ses spikes et son potentiel
spikemon = SpikeMonitor(Ga)
statemon_input = StateMonitor(Ga, 'I', record=True)
statemon_V = StateMonitor(Ga, 'V', record=2)

# On initialise les variables des neurones
Ga.V = Vs
Ga.I = 0

# Pour injecter le courant, on se définit
# une fonction appelée à chaque itération
@network_operation(when='start')
def set_current():
    global Ga
    # On modifie les entrées du neurones
    Ga.I = inputs[:,(int)(defaultclock.t/defaultclock.dt)] * volt

# On lance la simulation
run(tmax)

figure(figsize=(10,10))
subplot(211)
for i, ts in zip(spikemon.i, spikemon.t):
    plot([ts/msecond, ts/msecond], [i, i+1], 'k')
xlim([0, tmax/msecond])
ylim([-1, N+2])
title("Trains de spikes")

subplot(212)
plot(statemon_input.t, statemon_input.I.T)
xlim([0, tmax/second])
title(u'Courants d\'entrée')

#savefig('synchronisation_transitoire_nocoupling.png')
if(shuffle):
    savefig('synchronisation_transitoire_shuffle.png')
else:
    savefig('synchronisation_transitoire.png')

show()