import torch
import torch.nn as nn
from torch.autograd import Variable
import MetaRNNs
torch.manual_seed(100)

rnn = nn.RNN(4, 8, 1, batch_first=True)
input = Variable(torch.randn(2, 5, 4))
h0 = Variable(torch.randn(1, 2, 8))
output, hn = rnn(input, h0)
print("------------input------------")
print(input)
print("-------------h0--------------")
print(h0)
print("-----------output------------")
print(output)
print("-------------hn--------------")
print(hn)


test_rnn = MetaRNNs.MetaRNN(4, 8, 4, 4, 1)
test_input = Variable(torch.randn(2, 5, 4))
# test_h0 = Variable(torch.randn(1, 8))
test_output, test_hn = test_rnn(input)
print("------------test_input------------")
print(test_input)
print("------------test_output------------")
print(test_output)
print("--------------test_hn--------------")
print(test_hn)

lstm = nn.LSTM(4, 8, 1, batch_first=True)
input = Variable(torch.randn(1, 5, 4))
h0 = Variable(torch.randn(1, 1, 8))
c0 = Variable(torch.randn(1, 1, 8))
output, hn = lstm(input, (h0, c0))
print("------------input------------")
print(input)
print("-------------h0--------------")
print(h0)
print("-------------c0--------------")
print(c0)
print("-----------output------------")
print(output)
print("-------------hn--------------")
print(hn)

test_lstm = MetaRNNs.MetaLSTM(4, 8, 4, 4, 1)
test_input = Variable(torch.randn(1, 5, 4))
# test_h0 = Variable(torch.randn(1, 8))
test_output, test_hn = test_lstm(test_input)
print("------------test_input------------")
print(test_input)
print("------------test_output------------")
print(test_output)
print("--------------test_hn------------")
print(test_hn)