"""Demo of a simple LSTM language model.

Code is adapted from the PyTorch examples at
https://github.com/pytorch/examples/blob/main/word_language_model

"""

from collections.abc import Iterator, Sized
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torch.optim import Optimizer
from torch.utils.data import DataLoader, Sampler

from pytorch_lightning.core import LightningModule
from pytorch_lightning.demos.transformer import WikiText2


class SimpleLSTM(nn.Module):
    def __init__(
        self, vocab_size: int = 33278, ninp: int = 512, nhid: int = 512, nlayers: int = 4, dropout: float = 0.2
    ):
        super().__init__()
        self.vocab_size = vocab_size
        self.drop = nn.Dropout(dropout)
        self.encoder = nn.Embedding(vocab_size, ninp)
        self.rnn = nn.LSTM(ninp, nhid, nlayers, dropout=dropout, batch_first=True)
        self.decoder = nn.Linear(nhid, vocab_size)
        self.nlayers = nlayers
        self.nhid = nhid
        self.init_weights()

    def init_weights(self) -> None:
        nn.init.uniform_(self.encoder.weight, -0.1, 0.1)
        nn.init.zeros_(self.decoder.bias)
        nn.init.uniform_(self.decoder.weight, -0.1, 0.1)

    def forward(self, input: Tensor, hidden: tuple[Tensor, Tensor]) -> tuple[Tensor, Tensor]:
        emb = self.drop(self.encoder(input))
        output, hidden = self.rnn(emb, hidden)
        output = self.drop(output)
        decoded = self.decoder(output).view(-1, self.vocab_size)
        return F.log_softmax(decoded, dim=1), hidden

    def init_hidden(self, batch_size: int) -> tuple[Tensor, Tensor]:
        weight = next(self.parameters())
        return (
            weight.new_zeros(self.nlayers, batch_size, self.nhid),
            weight.new_zeros(self.nlayers, batch_size, self.nhid),
        )


class SequenceSampler(Sampler[list[int]]):
    def __init__(self, dataset: Sized, batch_size: int) -> None:
        super().__init__()
        self.dataset = dataset
        self.batch_size = batch_size
        self.chunk_size = len(self.dataset) // self.batch_size

    def __iter__(self) -> Iterator[list[int]]:
        n = len(self.dataset)
        for i in range(self.chunk_size):
            yield list(range(i, n - (n % self.batch_size), self.chunk_size))

    def __len__(self) -> int:
        return self.chunk_size


class LightningLSTM(LightningModule):
    def __init__(self, vocab_size: int = 33278):
        super().__init__()
        self.model = SimpleLSTM(vocab_size=vocab_size)
        self.hidden: Optional[tuple[Tensor, Tensor]] = None

    def on_train_epoch_end(self) -> None:
        self.hidden = None

    def training_step(self, batch: tuple[Tensor, Tensor], batch_idx: int) -> Tensor:
        input, target = batch
        if self.hidden is None:
            self.hidden = self.model.init_hidden(input.size(0))
        self.hidden = (self.hidden[0].detach(), self.hidden[1].detach())
        output, self.hidden = self.model(input, self.hidden)
        loss = F.nll_loss(output, target.view(-1))
        self.log("train_loss", loss, prog_bar=True)
        return loss

    def prepare_data(self) -> None:
        WikiText2(download=True)

    def train_dataloader(self) -> DataLoader:
        dataset = WikiText2()
        return DataLoader(dataset, batch_sampler=SequenceSampler(dataset, batch_size=20))

    def configure_optimizers(self) -> Optimizer:
        return torch.optim.SGD(self.parameters(), lr=20.0)