import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from math import sin

#def prepare_sequence(seq, to_ix):
#    idxs = [to_ix[w] for w in seq]
#    return torch.tensor(idxs, dtype=torch.long)

def geni(i):
    if i % 2 == 0:
        return sin(i/3) + 0.05#, sin(i) - 0.08]
    else:
        return sin(i/3) - 0.1#, sin(i) - 0.1, sin(i)]

input_seq = [geni(i) for i in range(1024)]

EMBEDDING_DIM = 7
HIDDEN_DIM = 7

class LSTMTagger(nn.Module):

    def __init__(self, input_dim, hidden_dim):
        super(LSTMTagger, self).__init__()
        self.hidden_dim = hidden_dim

        # The LSTM takes word embeddings as inputs, and outputs hidden states
        # with dimensionality hidden_dim.
        self.lstm = nn.LSTM(1, hidden_dim)

        # The linear layer that maps from hidden state space to tag space
        self.out = nn.Linear(hidden_dim, 1)

    def forward(self, sample_seq, lstm_state):
        lstm_out, lstm_state = self.lstm(sample_seq.view(len(sample_seq), 1, 1))
        score = self.out(torch.tanh(lstm_out[-1:]))
        return score, lstm_state


model = LSTMTagger(128, 128)
loss_function = nn.NLLLoss()
# optimizer = optim.SGD(model.parameters(), lr=0.1)
optimizer = optim.RMSprop(model.parameters(), lr=0.05, alpha=0.99, eps=1e-08, weight_decay=0, momentum=0.75, centered=False)

# See what the scores are before training
# Note that element i,j of the output is the score for tag j for word i.
# Here we don't need to train, so the code is wrapped in torch.no_grad()
#with torch.no_grad():
#    inputs = prepare_sequence(training_data[0][0], word_to_ix)
#    tag_scores = model(inputs)
#    print(tag_scores)

real_xy, guess_xy = [], []

lstm_state = torch.zeros([1,1,128]), torch.zeros([1,1,128])
for i in range(len(input_seq)):
    if i+130 >= len(input_seq):
        continue
    inputseq = input_seq[i:i+128]
    # Step 1. Remember that Pytorch accumulates gradients.
    # We need to clear them out before each instance
    model.zero_grad()

    # Step 2. Get our inputs ready for the network, that is, turn them into
    # Tensors of word indices.
    sentence_in = torch.tensor(inputseq, dtype=torch.float)

    # Step 3. Run our forward pass.
    #print("MODEL INPUT: ", sentence_in.shape)
    scout, lstm_state = model(sentence_in, lstm_state)

    # Step 4. Compute the loss, gradients, and update the parameters by
    #  calling optimizer.step()
    # loss = loss_function(tag_scores, targets)
    #loss = loss_function(scout, torch.tensor(input_seq[i+129]))
    loss = torch.square(scout - input_seq[i+129])
    print("DEBUG: LOSS=", loss)
    loss.backward(retain_graph=True)
    optimizer.step()

    if i > len(input_seq)-1024:
        real_xy.append(((i+129)/3, input_seq[i+129]))
        guess_xy.append(((i+129)/3, scout[0,0,0].tolist()))

#print("REAL=", real_xy)
#print("GUESS=", guess_xy)

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm

x = np.arange(10)
ys = [i+x+(i*x)**2 for i in range(10)]

colors = cm.rainbow(np.linspace(0, 1, len(ys)))
for x, y in real_xy:
    plt.scatter(x, y, color=colors[1])
for x, y in guess_xy:
    plt.scatter(x, y, color=colors[6])
plt.show()

#for epoch in range(3):  # again, normally you would NOT do 300 epochs, it is toy data
#    for sentence, tags in training_data:
#        # Step 1. Remember that Pytorch accumulates gradients.
#        # We need to clear them out before each instance
#        model.zero_grad()
#
#        # Step 2. Get our inputs ready for the network, that is, turn them into
#        # Tensors of word indices.
#        sentence_in = prepare_sequence(sentence, word_to_ix)
#        targets = prepare_sequence(tags, tag_to_ix)
#
#        # Step 3. Run our forward pass.
#        print("MODEL INPUT: ", sentence_in.shape)
#        tag_scores = model(sentence_in)
#
#        # Step 4. Compute the loss, gradients, and update the parameters by
#        #  calling optimizer.step()
#        loss = loss_function(tag_scores, targets)
#        loss.backward()
#        optimizer.step()

# See what the scores are after training
#with torch.no_grad():
#    inputs = prepare_sequence(training_data[0][0], word_to_ix)
#    tag_scores = model(inputs)
#
#    # The sentence is "the dog ate the apple".  i,j corresponds to score for tag j
#    # for word i. The predicted tag is the maximum scoring tag.
#    # Here, we can see the predicted sequence below is 0 1 2 0 1
#    # since 0 is index of the maximum value of row 1,
#    # 1 is the index of maximum value of row 2, etc.
#    # Which is DET NOUN VERB DET NOUN, the correct sequence!
#    print(tag_scores)

