from pyNN.nest import *

### Cell parameters ###
tau_m   = 20. # ms Membrane time constant
c_m     = 0.2 # nF Capacitance
tau_exc =  3. # ms Synaptic time constant (excitatory)
tau_inh =  7. # ms Synaptic time constant (inhibitory)
tau_ref =  5. # ms Refractory period
El      = -49 # mV Leak potential
Vt      = -50 # mV Spike Threhold 
Vr      = -60 # mV Spike Reset

cell_params = {
    'tau_m'      : tau_m, 'tau_syn_E'  : tau_exc, 'tau_syn_I'  : tau_inh,
    'v_rest'     : El,    'v_reset'    : Vr,      'v_thresh'   : Vt,
    'cm'         : c_m,   'tau_refrac' : tau_ref}

celltype      = IF_curr_exp             # The pyNN celltype we are using       
n_cells       = 1000                    # Total number of cells
n_exc         = int(0.8 * n_cells)      # 4:1 ratio for exc/inh
n_inh         = n_cells - n_exc         
size          = 1.                      # [arbitrary units] Size of the network
simtime       = 10000                   # [ms] Simulation time
epsilon       = 0.2                     # Probability density for the gaussian connections
s_lat         = 0.15                    # [unit] Spread of the lateral connections
g_exc         = 0.015                   # [nS] Excitatory conductance
g_inh         = -0.15                   # [nS] Inhibitory conductance
velocity      = 0.3                     # [mm/ms]  velocity of axonal delays
dt            = 0.1                     # [ms] timestep of the simulation
max_distance  = size/numpy.sqrt(2)      # [unit] Since this is a torus
max_delay     = dt + max_distance/velocity # [ms] Needed for the connectors
min_delay     = dt                      # [ms]
rngseed       = 92847459                # Random seed
parallel_safe = False                   # To have results independant on number of nodes

node_id = setup(timestep=dt, min_delay=min_delay, max_delay=max_delay+dt)

all_cells              = Population(n_exc+n_inh, celltype, cell_params, label="All Cells")
numpy.random.seed(12)
all_cells.positions    = size*numpy.random.rand(3, n_cells)
all_cells.positions[2] = 0
exc_cells              = all_cells[0:n_exc]
inh_cells              = all_cells[n_exc:n_cells]

rng = NumpyRNG(seed=rngseed, parallel_safe=parallel_safe)
uniformDistr = RandomDistribution('uniform', [Vr, 1.1*Vt], rng=rng)
all_cells.initialize('v', uniformDistr)

probas   = "%g * exp(-d**2/(2*%g**2))" %(epsilon, s_lat)
delays   = "%g + d/%g" %(dt, velocity)
space    = Space(periodic_boundaries=((0, 1), (0, 1), (0, 1))) #This is a torus, so we are wrapping at edges

exc_conn = DistanceDependentProbabilityConnector(probas, weights=g_exc, delays=delays, space=space)
inh_conn = DistanceDependentProbabilityConnector(probas, weights=g_inh, delays=delays, space=space)

connections        = {}
connections['exc'] = Projection(exc_cells, all_cells, exc_conn, target='excitatory', rng=rng)
connections['inh'] = Projection(inh_cells, all_cells, inh_conn, target='inhibitory', rng=rng)

def draw_rf(cell, positions, connections, color='k'):
    idx     = numpy.where(connections[:,1] == cell)[0]
    sources = connections[idx, 0]   
    for src in sources:        
        plot([positions[cell, 1], positions[src, 1]], [positions[cell, 2], positions[src, 2]], c=color)

def distances(pos_1, pos_2, N):
    dx = abs(pos_1[:,0]-pos_2[:,0])
    dy = abs(pos_1[:,1]-pos_2[:,1])
    dx = numpy.minimum(dx, N-dx)
    dy = numpy.minimum(dy, N-dy)
    return sqrt(dx*dx + dy*dy)

connections['exc'].saveConnections('connections.dat')
all_cells.save_positions('positions.dat')
positions    = numpy.loadtxt('positions.dat') ## Reading the positions from file
conn         = numpy.loadtxt('connections.dat')  ## Reading connections from file

idx_pre          = conn[:,0].astype(int)
idx_post         = conn[:,1].astype(int)
d                = distances(positions[idx_pre,1:3], positions[idx_post,1:3], 1)

plot(positions[:,1], positions[:,2], '.')
title('Cells positions')

d=hist(conn[:,3], 50)

ids   = numpy.random.permutation(positions[:,0])[0:6]
colors = ['k', 'r', 'b', 'g', 'c', 'y'] 
for count, cell in enumerate(ids):
    draw_rf(cell, positions, conn, colors[count])

exc_cells.record(to_file=False)

run(simtime)
exc_cells.printSpikes("spikes.dat")

from NeuroTools.signals import *

spikes = load_spikelist('spikes.dat', id_list=xrange(800))

spikes.raster_plot(t_stop=simtime)

fig=subplot(131)
spikes.firing_rate(100, display=fig)
fig=subplot(132)
spikes.rate_distribution(display=fig)
fig=subplot(133)
spikes.cv_isi_hist(display=fig)

spikes.activity_map(float_positions=positions[:800, 1:3].T, display=True)