import torch.nn as nn


class CharLSTM(nn.Module):
    """Stacked LSTM for next-character prediction on the Shakespeare dataset.

    Input:  LongTensor of shape (batch, seq_len=80) — character indices in [0, 79]
    Output: Tensor of shape (batch, 80) — logits over the 80-character alphabet
    """

    def __init__(self):
        super(CharLSTM, self).__init__()
        self.embed = nn.Embedding(80, 8)
        self.lstm = nn.LSTM(8, 256, 2, batch_first=True)
        self.drop = nn.Dropout()
        self.out = nn.Linear(256, 80)
        self.apply(self.weights_init)

    def forward(self, x):
        x = self.embed(x)
        x, _ = self.lstm(x)
        x = self.drop(x)
        return self.out(x[:, -1, :])

    def weights_init(self, module):
        if isinstance(module, nn.Linear):
            nn.init.xavier_uniform_(module.weight)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
