import numpy
import pylab
import os
from pyrnn.elmannet import ElmanNet
from pyrnn.learn import Learn, phase_plot_tixo

num_o = 2
num_c = 5
num_i = 2
ec_scale = 1
steps = 30
tch_T1 = 8
tch_T2 = 12
tch_dT = 0
iter_max = 10000
cb_step =  1
learn_rate = 1.0
momentum = 0.9

def make_lissajous(T1, T2, dT, steps):
    t = numpy.arange(steps)
    a = numpy.zeros((steps, 2), dtype=numpy.float)
    pi = numpy.pi
    a[:,0] = numpy.cos(2*pi*t/T1)
    a[:,1] = numpy.sin(2*pi*(t+dT)/T2)
    return a

def get_cb_phase_plot(ln):
    def gene_cb_phase_plot():
        i = 0
        while 1:
            issave = False
            if   i < 100               : issave = True
            elif i < 1000 and i%10 == 0: issave = True
            elif             i%100 == 0: issave = True
            if issave:
                phase_plot_tixo(ln)
                pylab.savefig('anime/nn%010d.png'%i)
            i += 1
            yield
    cb_phase_plot = gene_cb_phase_plot()
    return cb_phase_plot.next

nn = ElmanNet(num_o,num_c,num_i,steps)
ln = Learn(nn,iter_max)
cb_func = get_cb_phase_plot(ln)
ln.set_callback(cb_func,cb_step)

nn.set_learn_rate(learn_rate)
nn.momentum = momentum
nn.ns.to = make_lissajous(tch_T1,tch_T2,tch_dT,steps) * 0.8
nn.ns.ti[1:] = nn.ns.to[:-1]
nn.set_input()
nn.randomize_param()
nn.ns.ec *= ec_scale

os.mkdir('anime')
ln.learn()