# coding: utf-8
""" module for training and creating the model """
import torch
import click
from torch.utils.data import DataLoader
from tqdm import tqdm  # type: ignore
import numpy as np
from dotd.data import ShakespeareDataset, preprocess, get_shakespeare


class GRU(torch.nn.Module):
    """simple GRU model"""

    def __init__(self, vocab_size, hidden_size, embedding_size, batch=32, layers=2):
        super().__init__()
        self.hidden_size = hidden_size  # size of the GRU layers
        self.batch = batch
        self.layers = layers  # how many GRU layers
        self.word_embeds = torch.nn.Embedding(
            vocab_size, embedding_size
        )  # Embedding layer
        self.gru = torch.nn.GRU(
            embedding_size, hidden_size, layers, batch_first=True
        )  # GRU layer(s)
        self.output_layer = torch.nn.Linear(hidden_size, vocab_size)

    def forward(self, inputs, hidden):
        x = self.word_embeds(
            inputs
        )  # transform the input integer into a high dimensional embedding
        output, hidden = self.gru(x, hidden)  # Compute the output of the GRU layer(s)
        output = self.output_layer(output)  # compute the logits
        return output, hidden

    def init_hidden(self):
        return torch.zeros(self.layers, self.batch, self.hidden_size)


def train_epoch(model, loader, optim, loss, device):
    current_loss = []  # record running loss
    model.to(device)  # put the model on the specified device
    hidden = model.init_hidden().to(device)  # create the hidden state
    model.train()  # tell the model its training time
    for X, y in loader:
        X, y = X.to(device), y.to(
            device
        )  # collect the data and labels from the dataloader and put them on the device
        optim.zero_grad()  # empty the gradients
        output, hidden = model(X, hidden)  # compute the output
        hidden = hidden.detach()  # take the hidden state out of the graph
        batch_loss = loss(output.transpose(1, 2), y)  # compute loss
        batch_loss.backward()  # compute gradients
        optim.step()  # update weights
        current_loss.append(batch_loss.item())  # record loss
    epoch_loss = np.mean(current_loss)
    return epoch_loss


def train(epochs: int, save_path: str):
    _, _, alphabet_size, _ = preprocess(
        get_shakespeare()
    )  # get amount of characters for one-hot to embedding
    gru = GRU(alphabet_size, 1024, 256, layers=2)  # instantiate model
    loss = torch.nn.CrossEntropyLoss()
    optim = torch.optim.Adam(gru.parameters())
    device = "cuda" if torch.cuda.is_available() else "cpu"
    loader = DataLoader(
        ShakespeareDataset(seq_len=100),
        batch_size=32,
        shuffle=True,
        num_workers=2,
        drop_last=True,
    )
    for e in tqdm(range(epochs)):
        l = train_epoch(gru, loader, optim, loss, device)
        print(f"Epoch: {e}, Loss: {l}")
    torch.save(gru.state_dict(), save_path + f"gru_{epochs+1}epochs.pth")


@click.command()
@click.option(
    "--save_path",
    type=str,
    default="./saved_models",
    help="Folder to store the weights in.",
)
@click.option("--epochs", type=int, default=20, help="Number of epochs to train")
def main(epochs, save_path):
    train(epochs, save_path)


if __name__ == "__main__":
    main()
