"""Autoregressive MLP generative model (Stage C, for the N5 architecture-generality axis).

A causal feed-forward next-token model: token ``i`` is predicted from the concatenated
(causally-masked) embeddings of all earlier tokens. Deliberately a *different* inductive
bias from the GRU — if collapse appears here too, it is a property of the transmission
operator, not of any one architecture.
"""

from __future__ import annotations

import numpy as np

from .config import ModelCfg, SyntheticCfg
from .oracle import Oracle
from .torch_models import _BaseTorchGenerator
from .train import device_generator, seed_everything


def _make_mlp_net(V: int, L: int, embed: int, hidden: int):
    import torch
    import torch.nn as nn

    class MLPNet(nn.Module):
        """Predict every position from a causally-masked flatten of prior embeddings."""

        def __init__(self) -> None:
            super().__init__()
            self.V, self.L, self.E = V, L, embed
            self.bos = V
            self.embed = nn.Embedding(V + 1, embed)
            self.net = nn.Sequential(
                nn.Linear(L * embed, hidden), nn.ReLU(),
                nn.Linear(hidden, hidden), nn.ReLU(),
                nn.Linear(hidden, V),
            )
            # lower-triangular INCLUSIVE mask over input positions: position i sees inputs
            # 0..i (the input is already shifted by one, so this is strictly causal on x).
            mask = torch.tril(torch.ones(L, L))
            self.register_buffer("mask", mask)

        def _context(self, inp):                       # inp: (B, L) input tokens
            B = inp.shape[0]
            emb = self.embed(inp)                       # (B, L, E)
            m = self.mask.to(emb.dtype)                 # (L, L)
            # ctx[b, i] = concat_j ( emb[b, j] * mask[i, j] )  -> (B, L, L*E)
            ctx = emb.unsqueeze(1) * m.unsqueeze(0).unsqueeze(-1)   # (B, L, L, E)
            return ctx.reshape(B, self.L, self.L * self.E)

        def forward(self, x):                           # x: (B, L) targets
            B = x.shape[0]
            bos = torch.full((B, 1), self.bos, dtype=torch.long, device=x.device)
            inp = torch.cat([bos, x[:, :-1]], dim=1)
            ctx = self._context(inp)
            return self.net(ctx)                        # (B, L, V)

        def step_logits(self, prefix):                  # prefix: (B, pos) tokens so far
            """Logits for the next token given the tokens generated so far."""
            B, pos = prefix.shape
            bos = torch.full((B, 1), self.bos, dtype=torch.long, device=prefix.device)
            inp = torch.cat([bos, prefix], dim=1)[:, : self.L]      # (B, <=L)
            if inp.shape[1] < self.L:
                pad = torch.zeros((B, self.L - inp.shape[1]), dtype=torch.long,
                                  device=prefix.device)
                inp = torch.cat([inp, pad], dim=1)
            ctx = self._context(inp)                    # (B, L, L*E)
            return self.net(ctx[:, pos, :])             # logits at position `pos`

    return MLPNet()


class MLPGenerator(_BaseTorchGenerator):
    """Autoregressive feed-forward generative model over token sequences."""

    def fit(self, X: np.ndarray, rng: np.random.Generator) -> None:
        import torch

        g = seed_everything(int(rng.integers(2 ** 31)))
        net = _make_mlp_net(self.V, self.L, self.mcfg.embed, self.mcfg.hidden).to(self.device)
        net.train()
        opt = torch.optim.Adam(net.parameters(), lr=self.mcfg.lr)
        loss_fn = torch.nn.CrossEntropyLoss()
        data = torch.as_tensor(np.asarray(X), dtype=torch.long, device=self.device)
        n, bs = data.shape[0], self.mcfg.batch_size
        for _ in range(self.mcfg.epochs):
            perm = torch.randperm(n, generator=g).to(self.device)
            for i in range(0, n, bs):
                batch = data[perm[i:i + bs]]
                logits = net(batch)
                loss = loss_fn(logits.reshape(-1, self.V), batch.reshape(-1))
                opt.zero_grad()
                loss.backward()
                opt.step()
        net.eval()
        self.net = net

    def sample(self, n: int, rng: np.random.Generator) -> np.ndarray:
        import torch

        if self.net is None:
            raise RuntimeError("MLPGenerator.sample called before fit")
        g = device_generator(int(rng.integers(2 ** 31)), self.device)
        prefix = torch.empty((n, 0), dtype=torch.long, device=self.device)
        with torch.no_grad():
            for pos in range(self.L):
                logits = self.net.step_logits(prefix)
                probs = torch.softmax(logits, dim=-1)
                tok = torch.multinomial(probs, 1, generator=g)
                prefix = torch.cat([prefix, tok], dim=1)
        return prefix.cpu().numpy()