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()