import numpy as np
import theano
import theano.tensor as T
%pylab inline
import time

rng = np.random
N = 4
dt= theano.shared(0.01, name="dt")

rn=lambda a:rng.randn(N)*0.0001+a
x = theano.shared(rn(0.1), name="x")
y = theano.shared(rn(0.1), name="y")
z = theano.shared(rn(0.), name="z")

#rossler                                                                                                                                                                                
a=T.dscalar("a")
b=T.dscalar("b")
c=T.dscalar("c")

rossler_param=[a,b,c]

xn=x+dt*(-y-z)
yn=y+dt*(x+a*y)
zn=z+dt*(b+z*(x-c))

rossler = theano.function(
                inputs=rossler_param,
                outputs=[xn],
                updates={x:xn,y:yn,z:zn},
                on_unused_input='warn'
                )

tn=20000

tp=np.array([0]*N)
t0=time.clock()
for i in xrange(tn):
    tp=np.vstack((tp,rossler(0.1,0.1,14.)))
print time.clock()-t0

x=T.dvector("X")
y=T.dvector("Y")
z=T.dvector("Z")

n=T.iscalar("n")

def rossler(x,y,z,a,b,c):
    xn=x+dt*(-y-z)
    yn=y+dt*(x+a*y)
    zn=z+dt*(b+z*(x-c))
    return [xn,yn,zn]

ros_result, ros_updates = theano.scan(fn=lambda x,y,z,aa,bb,cc:rossler(x,y,z,a,b,c),
                              sequences=None,
                              outputs_info=[x,y,z],
                              non_sequences=[a,b,c],
                              n_steps=n)

rosslers= theano.function(inputs=[x,y,z,a,b,c,n], outputs=ros_result, updates=ros_updates)

an=rn(0.1)
t0=time.clock()
res=rosslers(an,an,an,0.1,0.1,14.,tn)
print time.clock()-t0

#CPU
an=rn(0.1)
t0=time.clock()
res=rosslers(an,an,an,0.1,0.1,14.,tn)
print time.clock()-t0

#len(res)
an
#len(res[0])

values, updates = theano.scan(lambda a,b,c:{x:xn,y:yn,z:zn},
                              non_sequences=[a,b,c],
                              n_steps=n)
rosslers = theano.function([a,b,c,n], values, updates=updates, allow_input_downcast=True)

tp=np.array([0]*N)
t0=time.clock()
res=rosslers(0.1,0.1,14.,tn)
print time.clock()-t0

cn=tp.T
for cc in cn:
    plt.plot(cc[1:])

for cc in cn:
    plt.plot(cc[:10000])

rossler_param=T.dvector("param")
dt= theano.shared(0.01, name="dt")

x=T.dvector("X")
y=T.dvector("Y")
z=T.dvector("Z")

n=T.iscalar("n")

def dx(xx):
    return -xx[1]-xx[2]

def rossler(x,y,z,p):
    a=p[0]
    b=p[1]
    c=p[2]
    xx=[x,y,z]
#    xn=x+dt*(-y-z)
    xn=x+dt*dx(xx)
    yn=y+dt*(x+a*y)
    zn=z+dt*(b+z*(x-c))
    return [xn,yn,zn]

ros_result, ros_updates = theano.scan(rossler,
                              sequences=None,
                              outputs_info=[x,y,z],
                              non_sequences=[rossler_param],
                              n_steps=n)

rosslers= theano.function(inputs=[x,y,z,rossler_param,n],
                          outputs=ros_result, 
                          updates=ros_updates)


an=rn(0.1)
t0=time.clock()
res=rosslers(an,an,an,[0.1,0.1,14.],tn)
print time.clock()-t0

cn=res.T
for cc in cn:
    plt.plot(cc[1:])

res[0][0][0]

res[0][0]