"""Shared machinery for workspace-band looping (WORKSPACE_LOOPING.md).

Codifies probe 2d: loop L14-30 through a merge layer at the L13->L14 boundary,
    L14_in = (1-alpha)*e + alpha*(s renormalized to |e|) + MLP([e; s_hat])
with e = L13 output (fixed anchor) and s = looped-back band output.
The MLP is zero-initialized, so the untrained adapter reproduces the
hand-built alpha-merge exactly.

Band forward is done by re-calling the decoder layers with the per-layer
(args, kwargs) captured from a normal forward pass -- position embeddings and
attention masks do not depend on hidden states, so they are reusable across
loop iterations.
"""

import re
import sys
from pathlib import Path

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from jlens.core import _text_model, load_model  # noqa: E402

import os as _os
BAND = tuple(int(x) for x in
             _os.environ.get("JLENS_BAND", "14,30").split(","))
# default: workspace band on gemma-4-E2B (inclusive); 12B: JLENS_BAND=36,45
ALPHA = 0.3      # anchor-dominant merge weight from probe 2d


class MergeAdapter(nn.Module):
    """(1-a)*e + a*s_hat + MLP([e; s_hat]); MLP zero-init => starts at probe 2d."""

    def __init__(self, d=1536, hidden=512, alpha=ALPHA):
        super().__init__()
        self.alpha = alpha
        self.mlp = nn.Sequential(
            nn.Linear(2 * d, hidden), nn.GELU(), nn.Linear(hidden, d)
        )
        nn.init.zeros_(self.mlp[2].weight)
        nn.init.zeros_(self.mlp[2].bias)

    def forward(self, e, s):
        dt = e.dtype
        e32, s32 = e.float(), s.float()
        s_hat = s32 * (
            e32.norm(dim=-1, keepdim=True) / (s32.norm(dim=-1, keepdim=True) + 1e-6)
        )
        base = (1 - self.alpha) * e32 + self.alpha * s_hat
        out = base + self.mlp(torch.cat([e32, s_hat], dim=-1))
        return out.to(dt)


class AdaptiveMergeAdapter(nn.Module):
    """Merge with state-dependent anchor coefficient (Lys-inspired, trained).

    alpha(e, s) = sigmoid(w·[e;ŝ] + b), per position; w zero-init and
    b = logit(0.3), so at init this is exactly the fixed alpha=0.3 merge."""

    def __init__(self, d=1536, hidden=512, alpha0=0.3):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(2 * d, hidden), nn.GELU(), nn.Linear(hidden, d)
        )
        nn.init.zeros_(self.mlp[2].weight)
        nn.init.zeros_(self.mlp[2].bias)
        self.alpha_head = nn.Linear(2 * d, 1)
        nn.init.zeros_(self.alpha_head.weight)
        import math as _m
        nn.init.constant_(self.alpha_head.bias,
                          _m.log(alpha0 / (1 - alpha0)))

    def forward(self, e, s):
        dt = e.dtype
        e32, s32 = e.float(), s.float()
        s_hat = s32 * (
            e32.norm(dim=-1, keepdim=True) / (s32.norm(dim=-1, keepdim=True) + 1e-6)
        )
        cat = torch.cat([e32, s_hat], dim=-1)
        a = torch.sigmoid(self.alpha_head(cat))
        out = (1 - a) * e32 + a * s_hat + self.mlp(cat)
        return out.to(dt)


class RecurrentAdapter(nn.Module):
    """Huginn-style recurrent-state update on the frozen band
    (arXiv 2502.05171 regime: h_{t+1} = A·h_t + B·e + Transformer(h_t, e)).

    The band's residual stream supplies the "+Transformer" term, so the
    adapter computes the band input x_t = A·ĥ_t + B·e + MLP([e;ĥ_t]) with
    LEARNED d×d maps A, B (init A=α·I, B=(1−α)·I: starts exactly at the
    fixed merge). h_0 is norm-scaled noise via init_state — combined with
    randomized-depth training this targets depth-monotone iteration rather
    than our anchor-dominant fixed point.
    """

    def __init__(self, d=1536, hidden=512, alpha=ALPHA, sigma=1.0):
        super().__init__()
        self.A = nn.Linear(d, d, bias=False)
        self.B = nn.Linear(d, d, bias=False)
        with torch.no_grad():
            self.A.weight.copy_(alpha * torch.eye(d))
            self.B.weight.copy_((1 - alpha) * torch.eye(d))
        self.mlp = nn.Sequential(
            nn.Linear(2 * d, hidden), nn.GELU(), nn.Linear(hidden, d)
        )
        nn.init.zeros_(self.mlp[2].weight)
        nn.init.zeros_(self.mlp[2].bias)
        self.sigma = sigma

    def init_state(self, e):
        """h_0: per-position Gaussian noise scaled to the anchor's norm."""
        n = torch.randn_like(e.float())
        n = n * (e.float().norm(dim=-1, keepdim=True)
                 / (n.norm(dim=-1, keepdim=True) + 1e-6)) * self.sigma
        return n.to(e.dtype)

    def forward(self, e, s):
        dt = e.dtype
        e32, s32 = e.float(), s.float()
        s_hat = s32 * (
            e32.norm(dim=-1, keepdim=True) / (s32.norm(dim=-1, keepdim=True) + 1e-6)
        )
        out = (self.A(s_hat) + self.B(e32)
               + self.mlp(torch.cat([e32, s_hat], dim=-1)))
        return out.to(dt)

    def rho(self):
        """Spectral radius of A (power iteration) — the stability dial."""
        with torch.no_grad():
            v = torch.randn(self.A.weight.shape[0],
                            device=self.A.weight.device)
            for _ in range(30):
                v = self.A.weight.T @ (self.A.weight @ v)
                v = v / (v.norm() + 1e-12)
            return (self.A.weight @ v).norm().item()


class ParcaeAdapter(RecurrentAdapter):
    """Recurrent adapter with ρ(A) < 1 GUARANTEED by construction
    (Parcae parameterization, Prairie et al. 2026): A is the ZOH
    discretization of a continuous negative-diagonal system,
    A = exp(−Δt·exp(a)) with learned per-dim a and scalar log-Δt, so every
    entry lies in (0,1) regardless of optimizer noise. Init matches the
    fixed merge exactly: a = log(−log α), Δt = 1 → A = α·I."""

    def __init__(self, d=1536, hidden=512, alpha=ALPHA, sigma=1.0):
        nn.Module.__init__(self)
        import math as _m
        self.log_a = nn.Parameter(
            torch.full((d,), _m.log(-_m.log(alpha))))
        self.log_dt = nn.Parameter(torch.zeros(()))
        self.B = nn.Linear(d, d, bias=False)
        with torch.no_grad():
            self.B.weight.copy_((1 - alpha) * torch.eye(d))
        self.mlp = nn.Sequential(
            nn.Linear(2 * d, hidden), nn.GELU(), nn.Linear(hidden, d)
        )
        nn.init.zeros_(self.mlp[2].weight)
        nn.init.zeros_(self.mlp[2].bias)
        self.sigma = sigma

    def A_diag(self):
        return torch.exp(-torch.exp(self.log_dt) * torch.exp(self.log_a))

    def forward(self, e, s):
        dt = e.dtype
        e32, s32 = e.float(), s.float()
        s_hat = s32 * (
            e32.norm(dim=-1, keepdim=True) / (s32.norm(dim=-1, keepdim=True) + 1e-6)
        )
        out = (self.A_diag().to(e32.dtype) * s_hat + self.B(e32)
               + self.mlp(torch.cat([e32, s_hat], dim=-1)))
        return out.to(dt)

    def rho(self):
        with torch.no_grad():
            return self.A_diag().max().item()


class TiedAlphaAdapter(nn.Module):
    """Merge with LEARNED per-dim constant alpha, B tied to (1-a):
        x = (1-a)*e + a*s_hat + MLP([e;s_hat]),  a = sigmoid(a_hat)
    Convex combination => the LTI fixed point cannot leave the e-s_hat
    segment (substrate-anchored by construction) and rho = max(a) < 1.
    Init a = alpha everywhere: bit-identical to MergeAdapter at step 0."""

    def __init__(self, d=1536, hidden=512, alpha=ALPHA):
        super().__init__()
        import math as _m
        self.a_hat = nn.Parameter(
            torch.full((d,), _m.log(alpha / (1 - alpha))))
        self.mlp = nn.Sequential(
            nn.Linear(2 * d, hidden), nn.GELU(), nn.Linear(hidden, d)
        )
        nn.init.zeros_(self.mlp[2].weight)
        nn.init.zeros_(self.mlp[2].bias)

    def forward(self, e, s):
        dt = e.dtype
        e32, s32 = e.float(), s.float()
        s_hat = s32 * (
            e32.norm(dim=-1, keepdim=True) / (s32.norm(dim=-1, keepdim=True) + 1e-6)
        )
        a = torch.sigmoid(self.a_hat)
        out = ((1 - a) * e32 + a * s_hat
               + self.mlp(torch.cat([e32, s_hat], dim=-1)))
        return out.to(dt)

    def rho(self):
        with torch.no_grad():
            return torch.sigmoid(self.a_hat).max().item()


class PerDepthAdapter(nn.Module):
    """Depth-wise relaxation (Bae et al. 2024, entrance-level): iteration t
    gets its OWN merge adapter — breaks time-invariance, so each loop step
    can perform a different computation phase instead of converging to a
    fixed point. Depths beyond n_depth reuse the last adapter."""

    def __init__(self, d=1536, hidden=512, alpha=ALPHA, n_depth=4):
        super().__init__()
        self.steps = nn.ModuleList(
            MergeAdapter(d=d, hidden=hidden, alpha=alpha)
            for _ in range(n_depth))

    def at(self, t):
        return self.steps[min(t, len(self.steps) - 1)]

    def forward(self, e, s):  # fallback: first-depth adapter
        return self.steps[0](e, s)


class BandLooper:
    """Capture layer-call kwargs once per forward, then re-run L14-30 manually."""

    def __init__(self, model, band=BAND):
        self.model = model
        self.tm = _text_model(model)
        self.l0, self.l1 = band
        self.n_layers = len(self.tm.layers)

    def capture(self, input_ids, attention_mask=None, logits_to_keep=0,
                position_ids=None):
        """Plain forward; returns (calls dict, logits). calls[i] = (args, kwargs)."""
        calls, hin, handles = {}, {}, []
        for i in range(self.l0, self.n_layers):
            def pre(mod, args, kwargs, i=i):
                # normalize: strip hidden_states, keep the rest for re-calls
                drop = {"hidden_states", "past_key_value", "past_key_values",
                        "use_cache"}
                kw = {k: v for k, v in kwargs.items() if k not in drop}
                if "hidden_states" in kwargs:
                    hin[i] = kwargs["hidden_states"]
                    calls[i] = (args, kw)
                else:
                    hin[i] = args[0]
                    calls[i] = (args[1:], kw)
            handles.append(
                self.tm.layers[i].register_forward_pre_hook(pre, with_kwargs=True)
            )
        try:
            with torch.no_grad():
                out = self.model(input_ids=input_ids, attention_mask=attention_mask,
                                 use_cache=False, logits_to_keep=logits_to_keep,
                                 position_ids=position_ids)
        finally:
            for h in handles:
                h.remove()
        self._hin = hin
        return calls, out.logits

    def _run(self, h, calls, lo, hi):
        for i in range(lo, hi + 1):
            args, kwargs = calls[i]
            h = self.tm.layers[i](h, *args, **kwargs)
            if isinstance(h, tuple):
                h = h[0]
        return h

    def band(self, h, calls):
        try:
            from lora_band import loop_active
            loop_active(True)
            out = self._run(h, calls, self.l0, self.l1)
            loop_active(False)
            return out
        except ImportError:
            return self._run(h, calls, self.l0, self.l1)

    def suffix_logits(self, h, calls, last_only=False):
        h = self._run(h, calls, self.l1 + 1, self.n_layers - 1)
        if last_only:
            h = h[:, -1:]
        h = self.tm.norm(h.to(self.tm.norm.weight.dtype))
        head = self.model.lm_head if hasattr(self.model, "lm_head") else self.model.get_output_embeddings()
        logits = head(h)
        cap = getattr(self.model.config.get_text_config(), "final_logit_softcapping", None)
        if cap:
            logits = cap * torch.tanh(logits / cap)
        return logits

    def loop_logits(self, adapter, input_ids, k, attention_mask=None,
                    use_checkpoint=False, return_states=False, last_only=False,
                    loop_mask=None, feedforward=False, bptt=None):
        """bptt: backprop only through the last `bptt` iterations (McLeish-
        style truncated BPTT); earlier iterations run under no_grad."""
        """Teacher-forced logits after k merge->band loops. k=0 = plain forward.

        loop_mask (B, T) bool: positions where the merge applies; elsewhere the
        band input stays the anchor e ("latent planning" over the prompt span —
        unmasked positions still re-attend to the looped states each iteration).
        """
        calls, base_logits = self.capture(input_ids, attention_mask,
                                          logits_to_keep=1 if last_only else 0)
        if k == 0:
            return (base_logits, None) if return_states else base_logits
        del base_logits  # full-vocab logits — do not hold across the loop
        e = self._hin[self.l0].detach()
        if feedforward:
            # no-recurrence control: adapter sees (e, e), applied exactly once
            x = adapter(e, e)
            if loop_mask is not None:
                x = torch.where(loop_mask[..., None], x, e)
            s = (checkpoint(lambda x_: self.band(x_, calls), x,
                            use_reentrant=False) if use_checkpoint
                 else self.band(x, calls))
            logits = self.suffix_logits(s, calls, last_only=last_only)
            return (logits, [s]) if return_states else logits
        with torch.no_grad():
            # s_0: no trainable params upstream (noise state for recurrent
            # adapters — the band(e) warm start would hide the B·e path)
            s = (adapter.init_state(e) if hasattr(adapter, "init_state")
                 else self.band(e, calls))
        states = [s]
        n_nograd = max(0, k - bptt) if bptt else 0
        for i in range(k):
            if i < n_nograd:
                with torch.no_grad():
                    x = (adapter.at(i) if hasattr(adapter, "at")
                         else adapter)(e, s)
                    if loop_mask is not None:
                        x = torch.where(loop_mask[..., None], x, e)
                    s = self.band(x, calls)
                s = s.detach()
                states.append(s)
                continue
            x = (adapter.at(i) if hasattr(adapter, "at") else adapter)(e, s)
            if loop_mask is not None:
                x = torch.where(loop_mask[..., None], x, e)
            if use_checkpoint:
                s = checkpoint(lambda x_: self.band(x_, calls), x, use_reentrant=False)
            else:
                s = self.band(x, calls)
            states.append(s)
        logits = self.suffix_logits(s, calls, last_only=last_only)
        return (logits, states) if return_states else logits

    @torch.no_grad()
    def loop_generate(self, adapter, tok, input_ids, k, max_new_tokens=12,
                      attention_mask=None, loop_prompt_only=False,
                      stop_strs=()):
        """Batched greedy decode with the looped forward (no KV cache).

        loop_prompt_only: merge applies only to the initial prompt span;
        generated tokens go through the plain band (latent planning)."""
        ids = input_ids
        mask = attention_mask
        lmask = None
        if loop_prompt_only:
            lmask = (mask if mask is not None
                     else torch.ones_like(ids)).bool().clone()
        n_prompt = ids.shape[1]
        texts = [""] * ids.shape[0]
        eos = {tok.eos_token_id}
        eot = tok.convert_tokens_to_ids("<end_of_turn>")
        if eot is not None and eot >= 0:
            eos.add(eot)
        done = torch.zeros(ids.shape[0], dtype=torch.bool, device=ids.device)
        for _ in range(max_new_tokens):
            logits = self.loop_logits(adapter, ids, k, attention_mask=mask,
                                      last_only=True, loop_mask=lmask)
            nxt = logits[:, -1].argmax(-1)
            nxt = torch.where(done, torch.full_like(nxt, list(eos)[0]), nxt)
            ids = torch.cat([ids, nxt[:, None]], 1)
            if mask is not None:
                mask = torch.cat([mask, (~done)[:, None].long()], 1)
            if lmask is not None:
                lmask = torch.cat(
                    [lmask, torch.zeros_like(lmask[:, :1])], 1)
            for b, t in enumerate(nxt.tolist()):
                if not done[b] and stop_strs:
                    texts[b] += tok.decode([t])
            done |= torch.tensor([t.item() in eos for t in nxt], device=ids.device)
            if stop_strs:
                done |= torch.tensor(
                    [any(ss in tx for ss in stop_strs) for tx in texts],
                    device=ids.device)
            if done.all():
                break
        return ids

    @torch.no_grad()
    def generate_frozen_prompt(self, adapter, tok, input_ids, k,
                               max_new_tokens=220, attention_mask=None,
                               stop_strs=(), feedforward=False,
                               conv_out=None):
        """Fast equivalent of loop_generate(loop_prompt_only=True).

        The looped prompt states are constant across token steps (causality),
        so: loop the prompt once to the merged input x* = adapter(e, s*), then
        run ONE cached prefill with a hook swapping the band input to x*, and
        generate normally with the KV cache (native speed)."""
        if k == 0:
            x_star = None
        elif feedforward:
            calls, _ = self.capture(input_ids, attention_mask,
                                    logits_to_keep=1)
            x_star = adapter(self._hin[self.l0], self._hin[self.l0])
            del calls
        else:
            calls, _ = self.capture(input_ids, attention_mask,
                                    logits_to_keep=1)
            e = self._hin[self.l0]
            s = (adapter.init_state(e) if hasattr(adapter, "init_state")
                 else self.band(e, calls))
            x_star = e
            conv = (torch.full((e.shape[0],), -1, dtype=torch.long)
                    if conv_out is not None else None)
            for _i in range(k):
                x_star = (adapter.at(_i) if hasattr(adapter, "at")
                          else adapter)(e, s)
                s_new = self.band(x_star, calls)
                if conv is not None:
                    c = nn.functional.cosine_similarity(
                        s_new.float().flatten(1), s.float().flatten(1), dim=1)
                    conv[((c > 0.9995).cpu()) & (conv < 0)] = _i + 1
                s = s_new
            if conv is not None:
                conv[conv < 0] = k
                conv_out["k_conv"] = conv.tolist()
            del calls  # prefill re-runs band(x_star) -> same final s as slow path

        hook = None
        if x_star is not None:
            def swap(mod, args, kwargs):
                h = kwargs.get("hidden_states", args[0] if args else None)
                if h is not None and h.shape[1] == x_star.shape[1]:  # prefill
                    if "hidden_states" in kwargs:
                        kwargs["hidden_states"] = x_star.to(h.dtype)
                        return args, kwargs
                    return (x_star.to(h.dtype),) + args[1:], kwargs
                return None
            hook = self.tm.layers[self.l0].register_forward_pre_hook(
                swap, with_kwargs=True)
        try:
            unk = tok.unk_token_id
            eos = [t for t in (tok.eos_token_id,
                               tok.convert_tokens_to_ids("<end_of_turn>"),
                               tok.convert_tokens_to_ids("<turn|>"))
                   if t is not None and t >= 0 and t != unk]
            out = self.model.generate(
                input_ids=input_ids, attention_mask=attention_mask,
                max_new_tokens=max_new_tokens, do_sample=False,
                eos_token_id=eos, pad_token_id=tok.pad_token_id or 0,
                stop_strings=list(stop_strs) or None, tokenizer=tok)
        finally:
            if hook:
                hook.remove()
        return out


# ---------- GSM8K helpers ----------

NUM_RE = re.compile(r"-?\$?[\d,]*\.?\d+")


def gold_answer(ans_field):
    return ans_field.split("####")[-1].strip().replace(",", "").replace("$", "")


def last_number(text):
    hits = NUM_RE.findall(text)
    if not hits:
        return None
    x = hits[-1].replace(",", "").replace("$", "").rstrip(".")
    return x


def num_eq(a, b):
    try:
        return a is not None and b is not None and abs(float(a) - float(b)) < 1e-4
    except ValueError:
        return False


DIRECT_SUFFIX = "\n\nGive only the final numeric answer, nothing else."
COT_SUFFIX = ("\n\nThink step by step, then give the final numeric answer "
              "on the last line as: #### <number>")


def chat_prompt(tok, question, suffix=DIRECT_SUFFIX):
    return tok.apply_chat_template(
        [{"role": "user", "content": question + suffix}],
        tokenize=False, add_generation_prompt=True,
    )


def build_train_batch(tok, items, device="cuda"):
    """Right-padded (input_ids, attention_mask, labels); labels only on answer tokens."""
    seqs, labs = [], []
    for it in items:
        p = tok(chat_prompt(tok, it["question"]), add_special_tokens=False)["input_ids"]
        a = tok(it["gold"] + "<end_of_turn>", add_special_tokens=False)["input_ids"]
        seqs.append(p + a)
        labs.append([-100] * len(p) + a)
    T = max(len(s) for s in seqs)
    pad = tok.pad_token_id or 0
    ids = torch.full((len(seqs), T), pad, dtype=torch.long)
    lab = torch.full((len(seqs), T), -100, dtype=torch.long)
    msk = torch.zeros((len(seqs), T), dtype=torch.long)
    for i, (s, l) in enumerate(zip(seqs, labs)):
        ids[i, : len(s)] = torch.tensor(s)
        lab[i, : len(s)] = torch.tensor(l)
        msk[i, : len(s)] = 1
    return ids.to(device), msk.to(device), lab.to(device)