from brian import *
import time
import pickle
import joblib
import gc

def brianreset():
    clear(True,True)
    gc.collect()
    defaultclock.t = 0*second

%%timeit
for i in range(1000):
    for j in range(1000):
        i*j

%%timeit
for i in xrange(1000):
    for j in xrange(1000):
        i*j

A = rand(100000)
B = rand(100000)

%%timeit
for a, b in zip(A, B):
    a*b

import itertools

%%timeit
for a, b in itertools.izip(A, B):
    a*b

import pickle
obj = ['a', 'list', 'of', 'words']
pickle.dump(obj, open('alistofwords.pkl', 'wb'), -1)
newobj = pickle.load(open('alistofwords.pkl', 'rb'))
print ' '.join(newobj)

import shelve
shelf = shelve.open('test.shelf')
shelf['meaning'] = 42
shelf.close()
newshelf = shelve.open('test.shelf')
print newshelf['meaning']

from collections import namedtuple
paramA = linspace(0, 1, 10)
paramB = linspace(0, 1, 10)
trials = []
Trial = namedtuple("Trial", ['a', 'b', 'spikes'])
for i in xrange(100):
    a = paramA[randint(len(paramA))]
    b = paramB[randint(len(paramB))]
    r = (a-0.5)**2+(b-0.5)**2
    spikes = rand(int(r*100))
    trials.append(Trial(a=a, b=b, spikes=spikes))

from collections import defaultdict

counts = defaultdict(int)
numtrials = defaultdict(int)

for trial in trials:
    counts[trial.a, trial.b] += len(trial.spikes)
    numtrials[trial.a, trial.b] += 1
    
rates = dict()

for a, b in counts.keys():
    rate = counts[a,b]*1.0/numtrials[a,b]
    plot([a], [b], 'o', ms=rate, c='k')
    rates[a, b] = rates

pickle.dump(rates, open('rates.pkl', 'wb'), -1)

a = linspace(0, 1, 100)
b = linspace(0, 1, 100)

def getimg():
    img = zeros((100, 100))
    for i in xrange(100):
        for j in xrange(100):
            img[i, j] = (a[i]-0.5)**2+(b[j]-0.5)**2
    return img

def getimg_vectorised():
    return (a[:, newaxis]-0.5)**2+(b[newaxis, :]-0.5)**2

subplot(121)
imshow(getimg())
subplot(122)
imshow(getimg_vectorised())

%timeit getimg()
%timeit getimg_vectorised()

delta = 0.005
train1 = sort(rand(50))
train2 = sort(rand(55))

def cf_loop():
    num_coinc = 0
    for t1 in train1:
        for t2 in train2:
            if abs(t1-t2)<delta:
                num_coinc += 1
                break
    return num_coinc

def cf_vectorised():
    ends = 0.5*(train2[1:]+train2[:-1])
    i = digitize(train1, ends)
    return sum(abs(train1-train2[i])<delta)

print cf_loop(), cf_vectorised()

%timeit cf_loop()
%timeit cf_vectorised()

from scipy import weave
N = 100
x = zeros(N)
code = '''
for(int i=0; i<N; i++)
{
    x[i] = i;
}
'''
namespace = {'x': x, 'N': N}
weave.inline(code, namespace.keys(), namespace, compiler='gcc')
print x[3]

import joblib

cache_mem = joblib.Memory(cachedir='.', verbose=0)
cache_mem.clear() # don't do this normally, it's just for the example

@cache_mem.cache
def f(x):
    print "--- I am currently computing x*x for x =", str(x)
    return x*x

print "Computing f(2.3)"
print f(2.3)
print "Computing f(3.5)"
print f(3.5)
print "Computing f(2.3) again"
print f(2.3)

from brian.globalprefs import *
set_global_preferences(
    useweave=True,
    weavecompiler='gcc',
    gcc_options=['-ffast-math', '-march=native'],
    usecodegen=True,
    usecodegenweave=True,
    usenewpropagate=True,
    usecstdp=True,
    )

from brian import *
import time
def runsim(N, duration=1*second, codegen=False):
    set_global_preferences(usecodegen=codegen)
    clk = Clock(dt=0.1*ms)
    G = NeuronGroup(N, 'dv/dt=(2-v**2)/(10*ms):1', reset=0, threshold=1, clock=clk)
    G.v = rand(len(G))
    net = Network(G)
    start = time.time()
    net.run(duration)
    end = time.time()
    return end-start

N = 10
runsim(N, codegen=True, duration=1*ms)
t_without = runsim(N, codegen=False)
t_with = runsim(N, codegen=True)
print 'N =', N
print 'Without code generation:', t_without 
print 'With code generation:   ', t_with
print 'Speedup:                ', t_without/t_with

import gc
for params in []:
    clear(True, True)
    gc.collect()
    defaultclock.t = 0*second
    # your code here

def brianreset():
    clear(True,True)
    gc.collect()
    defaultclock.t = 0*second

def IF(v0):
    brianreset()
    
    eqs = '''
    dv/dt=(v0-v)/(10*ms) : volt
    '''
    group = NeuronGroup(1, eqs, threshold='v>10*mV', reset='v = 0*mV')   
    monitor = SpikeMonitor(group)
    duration = 1*second
    run(duration)
    return monitor.nspikes/duration

V0 = [v0*mV for v0 in range(20)]

log_level_error() # suppress a few warning messages

%time plot(V0, [IF(v0) for v0 in V0])

if 0:
    import multiprocessing
    pool = multiprocessing.Pool()
    plot(V0, pool.map(IF, V0))

from IPython import parallel
rc = parallel.Client()
view = rc[:]

%%px
from brian import *
import gc

%time plot(V0, view.map_sync(IF, V0))

def IF_vectorised(V0):
    brianreset()
    
    eqs = '''
    dv/dt=(v0-v)/(10*ms) : volt
    v0 : volt
    '''
    group = NeuronGroup(len(V0), eqs, threshold='v>10*mV', reset='v = 0*mV')   
    group.v0 = V0
    monitor = SpikeCounter(group)
    duration = 1*second
    run(duration)
    return monitor.count/duration

%time plot(V0, IF_vectorised(V0))


if 0: # doesn't work in IPython
    from brian import *
    from brian.tools.datamanager import *
    from brian.tools.taskfarm import *
    
    def find_rate(k, report):
        eqs = '''
        dV/dt = (k-V)/(10*ms) : 1
        '''
        G = NeuronGroup(1000, eqs, reset=0, threshold=1)
        M = SpikeCounter(G)
        run(30*second, report=report)
        return (k, mean(M.count)/30)
    
    if __name__=='__main__':
        N = 20
        dataman = DataManager('taskfarmexample')
        if dataman.itemcount()<N:
            M = N-dataman.itemcount()
            run_tasks(dataman, find_rate, rand(M)*19+1)
        X, Y = zip(*dataman.values())
        plot(X, Y, '.')
        xlabel('k')
        ylabel('Firing rate (Hz)')
        show()

from IPython.core.display import Image
Image(filename='taskcontrol.jpg')

brianreset()

N = 5
G = NeuronGroup(N, 'dV/dt = -V/(10*ms) : 1')
@network_operation(clock=EventClock(dt=5*ms))
def random_input():
    i = randint(N)
    G.V[i] += 1
M = StateMonitor(G, 'V', record=True)
run(100*ms)
plot(M.times, M.values.T+3*arange(N)[newaxis, :]);

brianreset()

N = 5
duration = 100 * ms
Vr = -60 * mV
Vt = -50 * mV
tau = 10 * ms
Rmin = 1 * Mohm
Rmax = 10 * Mohm
freq = 50 * Hz
k = 10 * nA

eqs = '''
dV/dt = (-(V-Vr)+R*I)/tau : volt
R : ohm
I : amp
'''

G = NeuronGroup(N, eqs, reset='V=Vr', threshold='V>Vt')
G.R = linspace(Rmin, Rmax, N)

t = linspace(0 * second, duration, int(duration / defaultclock.dt))
I = clip(k * sin(2 * pi * freq * t), 0, Inf)
G.I = TimedArray(I)

M = MultiStateMonitor(G, record=True)

run(duration)

subplot(211)
M['I'].plot()
ylabel('I (amp)')
subplot(212)
M['V'].plot()
ylabel('V (volt)')

brianreset()

N = 5
f = 50 * Hz
a_min = 1.0
a_max = 100.0
tau_haircell = 50 * ms
tau = 10 * ms
duration = 100 * ms

eqs_haircells = '''
input = a*sin(2*pi*f*t) : 1
x = clip(input, 0, Inf)**(1.0/3.0) : 1
a : 1
dy/dt = (x-y)/tau_haircell : 1 
'''

haircells = NeuronGroup(N, eqs_haircells)
haircells.a = linspace(a_min, a_max, N)
M_haircells = MultiStateMonitor(haircells, vars=('input', 'y'), record=True)

eqs_nervefibres = '''
dV/dt = (I-V)/tau : 1
I : 1
'''
nervefibres = NeuronGroup(N, eqs_nervefibres, reset=0, threshold=1)
nervefibres.I = linked_var(haircells, 'y')
M_nervefibres = MultiStateMonitor(nervefibres, record=True)

run(duration)

subplot(221)
M_haircells['input'].plot()
ylabel('haircell.input')
subplot(222)
M_haircells['y'].plot()
ylabel('haircell.y')
subplot(223)
M_nervefibres['I'].plot()
ylabel('nervefibres.I')
subplot(224)
M_nervefibres['V'].plot()
ylabel('nervefibres.V')

brianreset()

eqs = '''
dV/dt = (I-V)/(10*ms)+0.1*xi*(2/(10*ms))**.5 : 1
I : 1
'''

G = NeuronGroup(3, eqs, reset=0, threshold=1)
M = RecentStateMonitor(G, 'V', duration=50*ms)

server = RemoteControlServer()

run(1e10*second)

if 0:
    from brian import *
    client = RemoteControlClient()
    clf(); plot(*client.evaluate('(M.times, M.values)'))
    client.execute('G.I = 1.1')
    clf(); plot(*client.evaluate('(M.times, M.values)'))
    client.stop()

def getsol_exact():
    brianreset()
    
    tau = 10*ms
    eqs = '''
    dv/dt = -v/tau : 1
    '''
    G = NeuronGroup(1, eqs)
    G.v = 1
    M = StateMonitor(G, 'v', record=True)
    run(100*ms)
    return M[0]

def getsol_numerical():
    brianreset()
    
    eqs = '''
    dv/dt = -v/tau : 1
    tau : second
    '''
    G = NeuronGroup(1, eqs)
    G.v = 1
    G.tau = 10*ms
    M = StateMonitor(G, 'v', record=True)
    run(100*ms)
    return M[0]

exact = getsol_exact()
numerical = getsol_numerical()
figure(figsize=(10,4))
subplot(121)
title('Solutions')
plot(exact)
plot(numerical)
subplot(122)
title('Difference')
plot(exact-numerical)

brianreset()
eqs = '''
dv/dt = -v/(10*ms) : 1
dvt/dt = (1-vt)/(100*ms) : 1
'''
def reset(G, spikes):
    G.vt[spikes] += 1
    G.v[spikes] = 0
inp = PoissonGroup(1, 100*Hz)
refrac = SimpleCustomRefractoriness(reset, 10*ms, 'v')
G = NeuronGroup(1, eqs, threshold='v>vt', reset=refrac)
G.vt = 1
C = Connection(inp, G, 'v', weight=0.5)
M = MultiStateMonitor(G, record=True)
run(500*ms)
figure(figsize=(10, 8))
M.plot()
ylim(-0.1, 3.0)
legend()

brianreset()
eqs = '''
dv/dt = -v/(10*ms) : 1
dvt/dt = (1-vt)/(100*ms) : 1
savevt : 1
'''
def resetfunc(G, spikes):
    G.vt[spikes] += 1
    G.savevt[spikes] = G.vt[spikes]
    G.v[spikes] = 0
def refracfunc(G, indices):
    G.vt[indices] = G.savevt[indices]
    G.v[indices] = 0
inp = PoissonGroup(1, 100*Hz)
refrac = CustomRefractoriness(resetfunc, 10*ms, refracfunc)
G = NeuronGroup(1, eqs, threshold='v>vt', reset=refrac)
G.vt = 1
C = Connection(inp, G, 'v', weight=0.5)
M = MultiStateMonitor(G, record=True)
run(500*ms)
figure(figsize=(10, 8))
M.plot()
ylim(-0.1, 3.0)
legend()