"""Shared helpers for the M1 notebooks (notes 0101-0110).

Everything is fp64 CPU torch.  The "language model" is a deliberately tiny
autoregressive policy whose *entire* sequence space can be enumerated, so every
expectation in the notes can be computed EXACTLY (no Monte Carlo) and compared
against the hand-derived formula to ~1e-15.  Monte Carlo checks are kept for the
few identities that are about sampling (and are reported against 4 standard errors).

Policy:  pi_theta(a | prev, t) = softmax(theta_tok[prev] + theta_pos[t])_a
  - vocabulary {0,..,V-1} where token V-1 is EOS
  - `prev` is the previous token (BOS = index V for t=0)
  - a sequence stops at EOS or after L tokens (truncation, no EOS)
This is a valid autoregressive factorisation  pi(o) = prod_t pi(o_t | o_<t)
(with a Markov-1 context) so every identity about the general case holds verbatim.
"""
import os, sys, math, itertools
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")      # never touch the GPUs
import torch
torch.set_default_dtype(torch.float64)
torch.set_num_threads(2)

CHECKS = []

def check(name, err, tol, kind="exact"):
    """Record + assert a numerical identity.  kind='exact' (fp64 arithmetic) or 'mc' (4-sigma Monte Carlo)."""
    err = float(err)
    ok = err <= tol
    CHECKS.append((name, err, tol, kind, ok))
    print(f"CHECK|{kind}|{name}|{err:.3e}|tol={tol:.1e}|{'PASS' if ok else 'FAIL'}")
    assert ok, f"{name}: err {err:.3e} > tol {tol:.1e}"

def maxabs(a, b):
    return float((torch.as_tensor(a) - torch.as_tensor(b)).abs().max())

class TinyLM:
    def __init__(self, V=4, L=4, seed=0, scale=1.0):
        self.V, self.L = V, L
        self.EOS = V - 1
        g = torch.Generator().manual_seed(seed)
        self.n_tok = (V + 1) * V          # rows: prev in {0..V-1, BOS}
        self.n_pos = L * V
        self.theta0 = scale * torch.randn(self.n_tok + self.n_pos, generator=g)
        self.seqs, self.tokens, self.mask = self._enumerate()
        self.N = len(self.seqs)
        self.lengths = self.mask.sum(1)

    # ---- enumeration of the whole sequence space --------------------------------
    def _enumerate(self):
        V, L, EOS = self.V, self.L, self.EOS
        seqs = []
        def rec(prefix):
            if len(prefix) == L or (prefix and prefix[-1] == EOS):
                seqs.append(tuple(prefix)); return
            for a in range(V):
                rec(prefix + [a])
        rec([])
        tokens = torch.full((len(seqs), L), EOS, dtype=torch.long)
        mask = torch.zeros(len(seqs), L)
        for i, s in enumerate(seqs):
            tokens[i, :len(s)] = torch.tensor(s)
            mask[i, :len(s)] = 1.0
        return seqs, tokens, mask

    def _logits_tables(self, theta):
        V, L = self.V, self.L
        tok = theta[:self.n_tok].reshape(V + 1, V)
        pos = theta[self.n_tok:].reshape(L, V)
        return tok, pos

    def token_logprobs(self, theta):
        """[N, L] log pi(o_t | o_<t) for every enumerated sequence (0 where masked)."""
        V, L = self.V, self.L
        tok, pos = self._logits_tables(theta)
        prev = torch.cat([torch.full((self.N, 1), V, dtype=torch.long), self.tokens[:, :-1]], 1)  # BOS = V
        logits = tok[prev] + pos[None, :, :]                       # [N, L, V]
        lp = torch.log_softmax(logits, -1)
        lp = lp.gather(-1, self.tokens[..., None]).squeeze(-1)
        return lp * self.mask

    def seq_logprob(self, theta):
        return self.token_logprobs(theta).sum(1)

    def probs(self, theta):
        return self.seq_logprob(theta).exp()

    # ---- naive, loop-based reference (used to cross-check the vectorised code) --
    def seq_logprob_naive(self, theta, seq):
        V = self.V
        tok, pos = self._logits_tables(theta)
        lp, prev = 0.0, V
        for t, a in enumerate(seq):
            logits = tok[prev] + pos[t]
            z = logits - logits.max()
            lp = lp + (z[a] - torch.log(torch.exp(z).sum()))
            prev = a
        return lp

    # ---- exact state machinery for value functions ------------------------------
    def prefixes(self):
        """All non-terminal prefixes (the MDP states), shortest first."""
        out = [()]
        for k in range(1, self.L):
            out += [p for p in itertools.product(range(self.V - 1), repeat=k)]  # non-EOS tokens only
        return out

    def cond_probs(self, theta, prefix):
        V = self.V
        tok, pos = self._logits_tables(theta)
        prev = V if len(prefix) == 0 else prefix[-1]
        return torch.softmax(tok[prev] + pos[len(prefix)], -1)

def default_reward(lm):
    """Binary verifier-style reward: success iff the sequence ends in EOS and its token sum is odd and length>=2."""
    R = torch.zeros(lm.N)
    for i, s in enumerate(lm.seqs):
        if s[-1] == lm.EOS and len(s) >= 2 and (sum(s[:-1]) % 2 == 1):
            R[i] = 1.0
    return R

def flat_grad(scalar, theta):
    g, = torch.autograd.grad(scalar, theta, retain_graph=True, allow_unused=False)
    return g

def score_matrix(lm, theta):
    """S[o] = grad_theta log pi(o)  as an [N, P] matrix (autograd, one backward per sequence)."""
    th = theta.clone().requires_grad_(True)
    lp = lm.seq_logprob(th)
    rows = [torch.autograd.grad(lp[i], th, retain_graph=True)[0] for i in range(lm.N)]
    return torch.stack(rows)

def token_score_tensor(lm, theta):
    """T[o, t] = grad_theta log pi(o_t | o_<t) as [N, L, P]."""
    th = theta.clone().requires_grad_(True)
    lp = lm.token_logprobs(th)
    P = theta.numel()
    out = torch.zeros(lm.N, lm.L, P)
    for i in range(lm.N):
        for t in range(int(lm.lengths[i])):
            out[i, t] = torch.autograd.grad(lp[i, t], th, retain_graph=True)[0]
    return out

def perturb(theta, delta, seed=1):
    g = torch.Generator().manual_seed(seed)
    d = torch.randn(theta.numel(), generator=g)
    return theta + delta * d / d.norm()