import matplotlib.pyplot as plt

from matplotlib import animation

from util.Generators import *

import os

current_folder = os.path.dirname(os.path.realpath(__file__))

save_folder = current_folder+"\\video\\"

os.chdir(current_folder)

print "Video will be save in: \n", save_folder

RECORD_VIDEO = False

simulation_time = 1000 #frames

## Note that we use two stimuli which are not equal in strength

## equal competition doesn't work well when sigmoid gain function + method of Euler for integration.

## --- Old Comments ---

## for repulsion only: A = 1.0 // GC1 = 1.0

## for repulsion and selection: A = 1.0 // GC1 = 100.0

## for selection only: A = 0.2 // GC1 = 100.0

## --------------------

start_bar = 450; end_bar = 550 # one bump

#start_bar = 350; end_bar = 650 # one bump

#start_bar = 325; end_bar = 675 # one bump

#start_bar = 315; end_bar = 685 # one bump ## size 370

#start_bar = 310; end_bar = 690 # two bumps ## size 380

#start_bar = 300; end_bar = 700 # two bumps

#start_bar = 250; end_bar = 750 # two bumps

# here we modified the Trappenberg set up of parameters:

GC1 = GeneratorConstante(250.0*1.0)

GC1.color = [1,0,0]

GC2 = GeneratorConstante(250.0*0)

width = 1000

NF = NeuralField1D(width, 100) ## time constant of 100 just to slow the process down and see it on the animation

NF.u0 = 0 ##Trappenberg biases the equilibirum state of build-up neurons

NF.getProbe(0).position = start_bar + 50

NF.getProbe(0).color = [1.0, 0.7, 0.2]

NF.getProbe(0).name = "P1"

NF.addProbe("P2", end_bar-50).color = [0.2, 0.7, 1.0]

NF.addProbe("P3", 500).color = [0.7, 0.2, 1.0]

## uncomment this line to get noise:

#NF.noise = lambda x: np.random.uniform(0, 1, x)*20

## --- Definition of the Gain function:

## Here we use the default gain function, which is a sigmoid defined by beta slope and an x_offset:

NF.beta = 0.02 # 0.07 in Trappenberg

NF.firing_offset = 350 ## this allows the model to be stable when no input (quite important, isn't it?'), it can't fire for negative value of u(t), which seems quite reasonnable. However, this parameters was set to 0 in Trappenberg 2001

# But we can also change the gain function and use an inverse exponential:

def exponentialRise(u):

return fr

# you just need to replace the function in the neural field NF:

#NF.GainFunction = exponentialRise

## -- Definition of the Connections:

## It seems that the parameters of Trappenberg was way to strong (but note that we have more neurons in our model)

Sa = 60.0; A = 144.0/200/5; ## Sa = 0.6mm while 1 neurons is 0.01

Sb = 180.0; B = 44.0/100/5;

c = 16.0/100/5;

## I create the lateral connections and attached them to the neural field

def DoF(x, Sa, Aa, Sb, Ab, c):

return DoG

connections = lambda x: DoF(x, Sa, A, Sb, B, c) ## that my connection function

NF.createConnectionCurve(connections) ## I give my connection function to my Neural Field

## I connect the Contanst Generator to my NeuralField as external input

#w1 = np.zeros(width); w1[390:410] = 1.0

w1 = np.zeros(width); w1[start_bar:end_bar] = 1.0

W1 = NF.set_external(w1, GC1)

w2 = np.zeros(width); w2[590:610] = 1.0

W2 = NF.set_external(w2, GC2)

list_obj = [GC1, GC2, NF] ## they need to be updated at each time step

## -- Plotting and Loop section:

# I plot the connection function and the gain function:

plt.figure()

plt.subplot(211)

plt.title("Close the window to start the simulation... \n Connections Pattern")

s = NF.myConnection.shape[0]/2

plt.plot(NF.myConnection[s-width/2:s+width/2], color = "pink", linewidth=2.0)

plt.xlabel("neurons")

plt.hlines(0, 0, width, color="gray", linestyle="--")

plt.subplot(212)

plt.plot(NF.GainFunction(np.arange(0,2000)), color = "purple")

plt.xlabel("Current/Membrane potential")

plt.ylabel("Firing Rate (Hz)")

plt.title("Gain Function")

plt.tight_layout()

plt.show()

## I define which probes we want to plot:

displayed_probes = ["Default", "P1", "P2", "P3"]

## First set up the figure, the axis, and the plot element we want to animate

fig = plt.figure(figsize = (7,11))

print "Hello"

ax1 = fig.add_subplot(411); ax1.set_xlim(0, 500); ax1.set_ylim(-1000, 1000)

ax1.set_title("Recorded Probes")

ax1.set_xlabel("time"); ax1.set_ylabel("Membrane Potential")

list_lines = []

for o in list_obj:

for probe in o.getProbes():

list_lines.append(*ax1.plot([], [], color = probe.color, linestyle = o.lns, lw=2))

ax2 = fig.add_subplot(412); ax2.set_xlim(0, width); ax2.set_ylim(-1000, 1000)

ax2.set_title("Membrane Potential over the Map (1D)")

ax2.set_xlabel("neuron positions"); ax2.set_ylabel("membrane potential")

NF_mp, = ax2.plot([], [], color = "red", linestyle = "-", lw=2)

ax2.hlines(0, 0, 2000, color = "gray", linestyle="--")

for probe in NF.getProbes():

ax2.scatter(probe.position, 0, s=10, color=probe.color)

ax2.hlines(0, start_bar, end_bar, color= "black", linewidth = 2.5)

list_lines.append(NF_mp)

ax3 = fig.add_subplot(413); ax3.set_xlim(0, width); ax3.set_ylim(0, 600)

ax3.set_title("Firing Rate over the Map (1D)")

ax3.set_xlabel("neuron positions"); ax3.set_ylabel("Firing Rate")

NF_firing, = ax3.plot([], [], color = "green", linestyle = "-", lw=2)

list_lines.append(NF_firing)

plt.tight_layout()

## Now the graph is ready, I run the sim and plot the results in real time:

def init():

return list_lines

## animation function. This is called sequentially

def animate(i):

return list_lines

from time import time

t0 = time()

print list_obj

## call the animator. blit=True means only re-draw the parts that have changed.

if not RECORD_VIDEO:

anim = animation.FuncAnimation(fig, animate, frames=simulation_time, interval=5 , blit=True, init_func= init, repeat=False)

else:

from moviepy.editor import VideoClip

from moviepy.video.io.bindings import mplfig_to_npimage

anim = VideoClip(animate, duration = simulation_time/60)

anim.write_videofile("bartest_"+str(end_bar-start_bar)+".mp4", fps = 60)

t1 = time()

print((t1 - t0))

plt.show()