#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Nov  6 15:27:56 2018
Elman in PyTorch
@author: owais
"""

import torch
from torch.autograd import Variable
import numpy as np
import pylab as pl
import torch.nn.init as init
import pandas as pd
from matplotlib import pyplot as plt

dtype = torch.DoubleTensor
input_size, hidden_size, output_size = 7, 5, 2
epochs = 300
seq_length = 20
lr = 0.1

#Read and Parse data into x, the input
xa = pd.read_csv('a.csv')
xb = np.zeros((len(xa),2))
for i in range(len(xa)):
    xb[i] = xa.values[i]     #Input layer
x = torch.from_numpy(xb)
    
#Build output matrix for prediction    
ya = np.roll(x,1,axis=0)
ya[0] = [0.65,0.7]        #Output layer
y = torch.from_numpy(ya)

mypred = torch.DoubleTensor(len(ya), output_size).type(dtype)

w1 = torch.DoubleTensor(input_size, hidden_size).type(dtype)
init.normal(w1, 0.0, 0.4)
w1 =  Variable(w1, requires_grad=True)
w2 = torch.DoubleTensor(hidden_size, output_size).type(dtype)
init.normal(w2, 0.0, 0.3)
w2 = Variable(w2, requires_grad=True)

def forward(input, context_state, w1, w2):
  xh = torch.cat((input, context_state), 1)
  context_state = torch.tanh(xh.mm(w1))
  out = context_state.mm(w2)
  return  (out, context_state)

for i in range(epochs):
  total_loss = 0
  context_state = Variable(torch.zeros((1, hidden_size)).type(dtype), requires_grad=True)
  for j in range(x.size(0)):
    input = x[j:(j+1)]
    target = y[j:(j+1)]
    (pred, context_state) = forward(input, context_state, w1, w2)
    mypred[j] = pred
    loss = (pred - target).pow(2).sum()/2
    total_loss += loss
    loss.backward()
    w1.data -= lr * w1.grad.data
    w2.data -= lr * w2.grad.data
    w1.grad.data.zero_()
    w2.grad.data.zero_()
    context_state = Variable(context_state.data)
  if i % 10 == 0:
     print("Epoch: {} loss {}".format(i, total_loss.data[0]))
     
context_state = Variable(torch.zeros((1, hidden_size)).type(dtype), requires_grad=False)
prediction = []

prediction = mypred.detach().numpy()
plt.plot(prediction[:,0], prediction[:,1])