import torch
import torch.nn as nn
import torch.nn.functional as F


class LSTMCell(nn.Module):
    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.hidden_size = hidden_size
        # combine all 4 gates into one matmul for efficiency
        self.gates = nn.Linear(
            input_size + hidden_size, 4 * hidden_size
        )

    def forward(self, x, state):
        h_prev, c_prev = state
        combined = torch.cat([x, h_prev], dim=1)
        gates = self.gates(combined)
        f, i, g, o = gates.chunk(4, dim=1)
        f = torch.sigmoid(f)
        i = torch.sigmoid(i)
        g = torch.tanh(g)
        o = torch.sigmoid(o)
        c = f * c_prev + i * g
        h = o * torch.tanh(c)
        return h, c


class CharLSTM(nn.Module):
    def __init__(self, vocab_size, hidden_size=128):
        super().__init__()
        self.hidden_size = hidden_size
        self.embed = nn.Embedding(vocab_size, hidden_size)
        self.cell = LSTMCell(hidden_size, hidden_size)
        self.output_proj = nn.Linear(hidden_size, vocab_size)

    def forward(self, x, state=None):
        # x: (batch, seq_len) integer token ids
        batch, seq_len = x.shape
        if state is None:
            h = torch.zeros(batch, self.hidden_size, device=x.device)
            c = torch.zeros(batch, self.hidden_size, device=x.device)
        else:
            h, c = state

        embedded = self.embed(x)
        logits = []
        for t in range(seq_len):
            h, c = self.cell(embedded[:, t, :], (h, c))
            logits.append(self.output_proj(h))
        return torch.stack(logits, dim=1), (h, c)
