#!/usr/bin/python3

# https://pytorch.org/tutorials/beginner/nlp/word_embeddings_tutorial.html

import torch
from torch import nn, optim

history_length = 32
history_encoded = [ord('\n')] * history_length
nb_of_char_codes = 128
embedding_size = 30
step = 1000

device = torch.device('cpu')

f = open('shakespeare.txt')
def char_source():
    for line in f:
        for c in line:
            c_code = ord(c)
            if c_code < nb_of_char_codes:
                yield(c_code)


class EncoderRNN(nn.Module):
    def __init__(self, input_size, hidden_size, embedding_size):
        super(EncoderRNN, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size

        self.embedding = nn.Embedding(input_size, embedding_size)
        self.gru = nn.GRU(embedding_size, nb_of_char_codes)
        self.softmax = nn.LogSoftmax(dim=1)

    def forward(self, input, hidden):
        embedded = self.embedding(input)
        output = embedded
        output, hidden = self.gru(output, hidden)
        output = self.softmax(output)
        return output, hidden

    def initHidden(self):
        return torch.zeros(1, self.hidden_size, self.input_size, device=device)

    def generate(self, n, encoder_hidden):
        t = (" " * 200 + "To be or not to be")[-history_length:]
        history = [ord(c) for c in t]

        with torch.no_grad():
            for _ in range(n):
                x = torch.tensor(history, dtype=torch.long, device=device)
                x = x.unsqueeze(0)
                y = model(x,encoder_hidden)[0][:,-1,:][0]
                y = torch.exp(y)
                best = (sorted(range(nb_of_char_codes), key=lambda i: -y[i]))[0:2]
                yb = torch.tensor([(y[ix] if ix in best else 0.0) for ix in range(nb_of_char_codes)])
                c = torch.multinomial(yb, 1)[0].item()

                t += chr(c)

                history.pop(0)
                history.append(c)

            print(t)


model = EncoderRNN(nb_of_char_codes, history_length, embedding_size).to(device)
criterion = nn.NLLLoss().to(device)
optimizer = optim.Adam(model.parameters())


counter = 0
losses = []

for c in char_source():
    x = torch.tensor(history_encoded, dtype=torch.long, device=device)
    model.zero_grad()
    x = x.unsqueeze(0)
    encoder_hidden = model.initHidden()
    y = model(x,encoder_hidden)[0][:,-1,:]

    loss = criterion(y, torch.tensor([c]).to(device))
    losses += [loss.item()]
    if len(losses) > step:
        losses.pop(0)

    counter += 1

    if counter % step == 0:
        avg_loss = sum(losses)/len(losses)
        print(f"{counter}: {loss} {avg_loss}")
        model.generate(200, encoder_hidden)

    loss.backward()
    optimizer.step()

    history_encoded.pop(0)
    history_encoded.append(c)