#!/usr/bin/env python3
"""mini-next: our mini implementation of the Flash-Next recipe on Whittle-16B.

Composition (each piece independently verifiable):
  1. MoE FFN     - carved 192-shared + 67x192 experts (see carve.py, CARVE_GATE)
  2. Hyper-connections - n residual streams with learned mixing, wrapped AROUND
     unmodified HF decoder layers via hooks. The layer computes out = h0 + T(h0)
     internally, so T(h0) = out - h0, and the HC update is
         H_i <- sum_j Ar[i,j] H_j + B_i * T(h0),   h0 = sum_i Am[k,i] H_i.
     Eq-14 identity init (HC paper, ICLR 2025): Am = e_{k mod n}, Ar = I, B = 1
     -> all streams stay equal to the standard residual, and the final row-sum's
     factor n cancels in the scale-invariant RMSNorm => logits IDENTICAL (HC_GATE).
  3. mHC constraint - Ar is parameterised through Sinkhorn-Knopp so residual
     mixing is doubly stochastic (can average, never amplify; the 3000x
     divergence fix from mHC).
  4. PLE - n-gram table with per-layer gated injection (train_ple.py checkpoint).
"""
import json, math
import torch
import torch.nn as nn
import torch.nn.functional as F


def sinkhorn(logits, iters=8):
    """Project exp(logits) onto (approx) doubly-stochastic via Sinkhorn-Knopp."""
    M = torch.exp(logits - logits.max())
    for _ in range(iters):
        M = M / (M.sum(-1, keepdim=True) + 1e-9)
        M = M / (M.sum(-2, keepdim=True) + 1e-9)
    return M


class HCState:
    def __init__(self): self.H = None
    def reset(self): self.H = None


class HyperConnections(nn.Module):
    """Static-matrix hyper-connections for L layers, expansion n (mHC-constrained)."""
    def __init__(self, n_layers, n=2, sinkhorn_iters=8):
        super().__init__()
        self.n, self.L, self.si = n, n_layers, sinkhorn_iters
        # Eq-14 identity init
        am = torch.zeros(n_layers, n)
        for k in range(n_layers): am[k, k % n] = 1.0
        self.Am_logit = nn.Parameter(torch.log(am + 1e-4))          # softmax -> ~e_{k mod n}
        eye = torch.eye(n).unsqueeze(0).repeat(n_layers, 1, 1)
        self.Ar_logit = nn.Parameter(torch.log(eye * 8.0 + 1.0))    # sinkhorn(exp) ~= I
        self.B = nn.Parameter(torch.ones(n_layers, n))               # write weights
        # Eq-14 init taken LITERALLY: while in identity mode, read/write use
        # exact index/add paths (zero arithmetic). The soft softmax/Sinkhorn
        # parameterisation is only engaged when training starts - in bf16 the
        # soft mix injects ~2^-8 error per layer and compounds to ~0.6 rel over
        # 44 layers, which is noise, not signal.
        self.identity_mode = True

    def release_identity(self):
        self.identity_mode = False

    def Am(self, k): return F.softmax(self.Am_logit[k], -1)          # non-neg, sums to 1
    def Ar(self, k): return sinkhorn(self.Ar_logit[k], self.si)      # doubly stochastic
    def read(self, k, H):
        if self.identity_mode:
            return H[k % self.n]
        # mixing coefficients are tiny (n, n^2); follow the streams' device -
        # layers span GPU boundaries under device_map.
        a = self.Am(k).to(dtype=H[0].dtype, device=H[0].device)
        return sum(a[i] * H[i] for i in range(self.n))
    def write(self, k, H, T_out):
        if self.identity_mode:
            return [H[i].to(T_out.device) + T_out for i in range(self.n)]
        R = self.Ar(k).to(dtype=H[0].dtype, device=H[0].device)
        b = self.B[k].to(dtype=H[0].dtype, device=H[0].device)
        return [sum(R[i, j] * H[j] for j in range(self.n)) + b[i] * T_out
                for i in range(self.n)]


def attach_hc(model, n=2):
    """Wrap every decoder layer of a HF qwen3_5(_moe) model in hyper-connections."""
    layers = model.model.layers
    hc = HyperConnections(len(layers), n=n)
    dev = next(layers[0].parameters()).device
    hc.to(dev).to(next(model.parameters()).dtype)
    st = HCState()
    inbuf = {}

    def mk_pre(k):
        def pre(mod, args, kwargs):
            h = kwargs.get("hidden_states", args[0] if args else None)
            if k == 0 or st.H is None:
                st.H = [h.clone() for _ in range(hc.n)]
            h0 = hc.read(k, [x.to(h.device) for x in st.H])
            inbuf[k] = h0
            if "hidden_states" in kwargs:
                kwargs["hidden_states"] = h0; return (args, kwargs)
            return ((h0,) + tuple(args[1:]), kwargs)
        return pre

    def mk_post(k, last):
        def post(mod, args, kwargs, out):
            o = out[0] if isinstance(out, tuple) else out
            if hc.identity_mode:
                # Eq-14 identity, taken to its bit-exact conclusion: with B=1,
                # Ar=I and equal streams, H_i <- H_i + (o - h0) == o. Assign
                # directly - zero extra arithmetic, so the wrapped model IS the
                # base model, bitwise. (T = o - h0 re-add costs one extra bf16
                # rounding per layer and flipped 5% of top-1s by layer 44.)
                inbuf.pop(k, None)
                st.H = [o for _ in range(hc.n)]
                new = o
                if last: st.reset()
                if isinstance(out, tuple): return (new,) + tuple(out[1:])
                return new
            T = o - inbuf.pop(k).to(o.device)   # layer may span a GPU boundary
            st.H = hc.write(k, [x.to(o.device) for x in st.H], T)
            new = sum(st.H) if last else st.H[0]
            # note: what we return only matters for the LAST layer (final norm
            # consumes it); intermediate layers are re-mixed by the next pre-hook.
            if last: st.reset()
            if isinstance(out, tuple): return (new,) + tuple(out[1:])
            return new
        return post

    hs = []
    for k, layer in enumerate(layers):
        hs.append(layer.register_forward_pre_hook(mk_pre(k), with_kwargs=True))
        hs.append(layer.register_forward_hook(mk_post(k, k == len(layers) - 1), with_kwargs=True))
    model._hc = hc
    model._hc_hooks = hs
    return hc


def detach_hc(model):
    for h in getattr(model, "_hc_hooks", []): h.remove()
    model._hc_hooks = []