from __future__ import annotations

import torch
from torch import nn

VOCAB_SIZE = 16
CONTROL_STATES = 8
EMBEDDING_DIM = 26


class HistoryCompressor(nn.Module):
    """A slow recurrent level updated only by informative control events."""

    def __init__(self, hidden_size: int = 20) -> None:
        super().__init__()
        self.embedding = nn.Embedding(VOCAB_SIZE, EMBEDDING_DIM)
        self.cell = nn.GRUCell(EMBEDDING_DIM, hidden_size)
        self.output = nn.Linear(hidden_size, CONTROL_STATES)

    def forward(
        self,
        tokens: torch.Tensor,
        *,
        return_states: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        batch, steps = tokens.shape
        hidden = torch.zeros(batch, self.cell.hidden_size, device=tokens.device)
        logits = []
        states = []
        embedded = self.embedding(tokens)
        for step in range(steps):
            candidate = self.cell(embedded[:, step], hidden)
            informative = tokens[:, step].lt(CONTROL_STATES)[:, None]
            hidden = torch.where(informative, candidate, hidden)
            logits.append(self.output(hidden))
            states.append(hidden)
        output = torch.stack(logits, dim=1)
        if return_states:
            return output, torch.stack(states, dim=1)
        return output


class PlainRNN(nn.Module):
    """A wider tanh RNN that updates on every control and distractor token."""

    def __init__(self, hidden_size: int = 40) -> None:
        super().__init__()
        self.embedding = nn.Embedding(VOCAB_SIZE, EMBEDDING_DIM)
        self.cell = nn.RNNCell(EMBEDDING_DIM, hidden_size)
        self.output = nn.Linear(hidden_size, CONTROL_STATES)

    def forward(
        self,
        tokens: torch.Tensor,
        *,
        return_states: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        batch, steps = tokens.shape
        hidden = torch.zeros(batch, self.cell.hidden_size, device=tokens.device)
        logits = []
        states = []
        embedded = self.embedding(tokens)
        for step in range(steps):
            hidden = self.cell(embedded[:, step], hidden)
            logits.append(self.output(hidden))
            states.append(hidden)
        output = torch.stack(logits, dim=1)
        if return_states:
            return output, torch.stack(states, dim=1)
        return output


def parameter_count(model: nn.Module) -> int:
    return sum(parameter.numel() for parameter in model.parameters())