#!/usr/bin/env python
# Based on the "Clockwork RNN" paper at 
# http://jmlr.org/proceedings/papers/v32/koutnik14.pdf
#'For LSTM it was also crucial to initialize the bias of the 
# forget gates to a high value (5.0 in this case) to encourage 
# the long-term memory'

import numpy as np
import sys
sys.path.append('../../python')
import caffe
from pudb import set_trace 

def gen_signal(seq_len):
    a = np.arange(0, seq_len/10, 0.01)
    s = 0.5 * np.sin(2*a) - \
        0.05 * np.cos( 17*a + 0.8  ) + \
        0.05 * np.sin( 25 * a + 10 ) - \
        0.02 * np.cos( 45 * a + 0.3)
    s = s / max(np.max(s), -np.min(s))
    s = s - np.mean(s)
    print s
    print "length of d {}".format(len(s))
    return s

caffe.set_mode_gpu()
caffe.set_device(0)
# set_trace()
solver = caffe.SGDSolver('toy_lstm_solver.prototxt')

print "blob names {}".format(solver.net.blobs.keys()) 
print "blob shapes {}".format([(k,v.data.shape) for k, v in solver.net.blobs.items()])

print "param names {}".format(solver.net.params.keys()) 
print "param shape 0 {}".format(solver.net.params['lstm'][0].data.shape)
print "param shape 1 {}".format(solver.net.params['lstm'][1].data.shape)
print "param shape 2 {}".format(solver.net.params['lstm'][2].data.shape)

# set_trace()

# Set the bias of the forget gate to 5.0 as explained in the clockwork RNN paper
#solver.net.params['lstm'][1].data[15:30] = 5.0
solver.net.params['lstm'][1].data[15:30] = 5.0

seq_len = np.shape(solver.net.blobs['data'].data)[0]
streams_len = np.shape(solver.net.blobs['data'].data)[1]
print "sequence length: {}".format (seq_len)
print "streams length: {}".format (streams_len)

solver.net.blobs['clip'].data.reshape(seq_len, 1)
solver.net.blobs['label'].data.reshape(seq_len, 1)

solver.net.blobs['clip'].data[...] = 1

sig = gen_signal(seq_len)
print "signal length: {}".format (len(sig))

# Train
niter = 50000
iter = 0
train_loss = np.zeros(niter)
while iter < niter :
    seq_idx = iter % (len(sig) / seq_len)
    # print "seq_idx val: {} clip value: {}".format(seq_idx, (seq_idx > 0))
    solver.net.blobs['clip'].data[0] = seq_idx > 0
    # solver.net.blobs['label'].data[:,0] = d[ seq_idx * seq_len : (seq_idx+1) * seq_len ]
    solver.net.blobs['label'].data[:,0] = sig[seq_idx * seq_len : (seq_idx + 1) * seq_len]
    # print solver.net.blobs['label'].data
    solver.step(1)
    train_loss[iter] = solver.net.blobs['loss'].data
    iter += 1

# Save 
solver.net.save("lstm.caffemodel")

# test
# set_trace()
test_net = caffe.Net('deploy_toy_lstm.prototxt', 'lstm.caffemodel', caffe.TEST)
test_net.blobs['data'].data.reshape(2, 1)
test_net.blobs['clip'].data.reshape(2, 1)
#test_net.blobs['label'].data.reshape(2, 1)
test_net.blobs['clip'].data[...] = 1
preds = np.zeros(len(sig))
with open('wave_tracking.txt', 'w') as f:
    for i in range(len(sig)):
        test_net.blobs['clip'].data[0] = i > 0
        preds[i] =  test_net.forward()['ip'][0][0]
        print "{} {}".format(sig[i], preds[i])
        f.write("%f %f\n" %(sig[i], preds[i]))