"""arms.py — the detachable arm runtime for mini-beatrix-2.5s.

An *arm* is a small aleph-addressed patch head placed after every decoder
block of one exact frozen core. It is not a low-rank matrix: each site
projects the block output into `n_slots` slots of dimension D, reads them
through a K-atom aleph address (a closed-form signed read, never a
selector), and adds a gated patch back into the residual stream:

    slots   = proj(x).view(B, n, n_slots, D)
    m_hat   = sum_k sinh(u_k) A_k / sum_k cosh(u_k),  u = (x_hat . A) / tau
    x       = x + sigmoid(gate) * consume(m_hat)

`mount()` makes the live core carry exactly one recipe and returns a
handle; `handle.detach(verify=True)` restores the untouched blocks and
asserts the restored logits are BIT-EXACT with a pre-attach fingerprint.
That exactness is the point of the design: an arm is a guest, never a
migration of the core.

Recipes carried by this package:
  anchor   one arm, always on
  stack    several anchor arms applied always-on, nested inner to outer,
           with no mixer between them (the ungated composition)
  routed   a dense signed dispatch over member arms, optionally with a
           detached bounded per-member gain read from a causal usage EMA

THE DAMPING LAW (do not "fix" it): the dispatch denominator sums cosh
over ALL members, including masked ones, and masking a member removes
its delta WITHOUT renormalizing. Masked is not solo.

This file is self-contained (torch + safetensors). It is a vendored,
verified re-expression of the amoe-lora runtime the arms were trained
with (github.com/AbstractEyes/amoe-lora): the packaged arms were checked
to produce bit-identical logits through both paths before release.
"""
from __future__ import annotations

import contextlib
import json
import os
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F

ARM_TYPES = ("anchor", "stack", "routed")


# --------------------------------------------------------------- geometry
@dataclass
class AdapterSpec:
    n_slots: int = 16
    K: int = 64
    D: int = 4
    tau: float = 0.1
    hidden: int = 178
    gate_init: float = -3.0


class AlephAddress(nn.Module):
    """m_hat(x) = sum_k sinh(u_k) A_k / sum_k cosh(u_k), max-|u| stabilized.

    `home` snapshots the codebook at init and is required in every arm
    state (the drift gauge)."""

    def __init__(self, K: int = 64, D: int = 4, tau: float = 0.1):
        super().__init__()
        self.K, self.D, self.tau = K, D, tau
        self.codebook = nn.Parameter(F.normalize(torch.randn(K, D), dim=-1))
        self.register_buffer("home", self.codebook.detach().clone())

    def m_hat(self, x: torch.Tensor) -> torch.Tensor:
        A = F.normalize(self.codebook, dim=-1)
        u = (F.normalize(x, dim=-1) @ A.transpose(-1, -2)) / self.tau
        m = u.abs().amax(dim=-1, keepdim=True)
        ep, en = torch.exp(u - m), torch.exp(-u - m)
        return ((ep - en) @ A) / (ep + en).sum(dim=-1, keepdim=True)

    @torch.no_grad()
    def drift(self) -> float:
        """Mean angular drift (radians) of the codebook from `home`."""
        A = F.normalize(self.codebook, dim=-1)
        H = F.normalize(self.home, dim=-1)
        return float(torch.arccos((A * H).sum(-1).clamp(-1.0, 1.0)).mean())


class SquaredReLU(nn.Module):
    def forward(self, x):
        return F.relu(x) ** 2


class RelayPatchwork(nn.Module):
    """One arm site: projection -> aleph address -> squared-ReLU patch head,
    added through a sigmoid gate."""

    def __init__(self, d: int, spec: AdapterSpec | None = None):
        super().__init__()
        s = spec or AdapterSpec()
        self.spec = s
        self.n_slots = s.n_slots
        self.proj = nn.Linear(d, s.n_slots * s.D, bias=False)
        self.addr = AlephAddress(s.K, s.D, s.tau)
        self.consume = nn.Sequential(
            nn.Linear(s.n_slots * s.D, s.hidden), SquaredReLU(),
            nn.LayerNorm(s.hidden), nn.Linear(s.hidden, d))
        self.gate = nn.Parameter(torch.tensor(float(s.gate_init)))

    def forward(self, x):
        B, n, _ = x.shape
        slots = self.proj(x).view(B, n, self.n_slots, self.spec.D)
        feats = self.addr.m_hat(slots).reshape(B, n, -1)
        return x + torch.sigmoid(self.gate) * self.consume(feats)


def spec_from_state(state: dict, prefix: str | None = None) -> AdapterSpec:
    """Arm geometry read from the arm's own tensors (the library holds more
    than one shape, and modules are built before the state loads)."""
    if prefix is None:
        blocks = sorted(int(k.split(".")[0]) for k in state
                        if k.endswith(".addr.codebook") and k.split(".")[0].isdigit())
        prefix = f"{blocks[0]}." if blocks else "0."
    cb = state.get(f"{prefix}addr.codebook")
    proj = state.get(f"{prefix}proj.weight")
    cons = state.get(f"{prefix}consume.0.weight")
    if cb is None or proj is None or cons is None:
        return AdapterSpec()
    D = int(cb.shape[1])
    return AdapterSpec(n_slots=int(proj.shape[0] // D), K=int(cb.shape[0]),
                       D=D, hidden=int(cons.shape[0]))


# --------------------------------------------------------------- wrappers
class BlockWithAdapter(nn.Module):
    """Wraps one decoder block. `enabled=False` returns the block output
    untouched — the single-arm half of the toggle law."""

    def __init__(self, block: nn.Module, adapter: RelayPatchwork):
        super().__init__()
        self.block = block
        self.adapter = adapter
        self.enabled = True

    def forward(self, *args, **kwargs):
        out = self.block(*args, **kwargs)
        if not self.enabled:
            return out
        if isinstance(out, tuple):
            return (self.adapter(out[0]),) + out[1:]
        return self.adapter(out)

    def prefill(self, *args, **kwargs):
        out, cache = self.block.prefill(*args, **kwargs)
        return (self.adapter(out) if self.enabled else out), cache

    def step(self, *args, **kwargs):
        out = self.block.step(*args, **kwargs)
        if not self.enabled:
            return out
        if isinstance(out, tuple):
            return (self.adapter(out[0]),) + out[1:]
        return self.adapter(out)


class AnchorDispatch(nn.Module):
    """Dense signed sinh/cosh dispatch over frozen member arms.

        out = x + sum_a (sinh u_a / sum_all cosh u) (arm_a(x) - x)

    No top-k, no softmax over members, no load balancing. The denominator
    sums over every member, masked or not (the damping law)."""

    def __init__(self, anchors: nn.ModuleList, d: int, emb: int = 64,
                 tau: float = 0.1):
        super().__init__()
        self.anchors = anchors
        self.tau = tau
        A = len(anchors)
        self.register_buffer("key_proj", torch.zeros(emb, d))
        self.dispatch = nn.Parameter(torch.zeros(A, emb))
        self.enabled = [True] * A

    def forward(self, x):
        emb = F.normalize(x @ self.key_proj.T, dim=-1)
        C = F.normalize(self.dispatch, dim=-1)
        u = (emb @ C.T) / self.tau
        m = u.abs().amax(dim=-1, keepdim=True)
        ep, en = torch.exp(u - m), torch.exp(-u - m)
        w = ep - en
        z = (ep + en).sum(dim=-1, keepdim=True)      # ALL members (law)
        out = x
        for k, a in enumerate(self.anchors):
            if self.enabled[k]:
                out = out + (w[..., k:k + 1] / z) * (a(x) - x)
        return out


class GainDispatch(AnchorDispatch):
    """AnchorDispatch with a detached bounded per-member gain on the
    numerator (the denominator is untouched, the signs are untouched):

        s_a(t)  = cosh u_a / sum_all cosh u          (the member's share)
        ubar_a  = causal EMA of s_a over bytes, rate rho, start 1/A
        alpha_a = clamp(1 + gamma (A ubar_a - 1), a_min, a_max)
        out     = x + sum_a alpha_a (sinh u_a / sum_all cosh u)(arm_a(x) - x)

    mode 'growth' lets the mean gain rise toward the cap; 'conserved'
    renormalizes alpha to mean 1 over members at every byte. `coherent`
    reads alpha once at the first site and shares it ('alpha'), and with
    'full' shares the first site's signed routing weight too."""

    def configure(self, *, mode="growth", rho=1 / 16, gamma=1.5, a_min=0.5,
                  a_max=2.0, coherent=None, role="ref", holder=None):
        self.mode, self.rho, self.gamma = mode, float(rho), float(gamma)
        self.a_min, self.a_max = float(a_min), float(a_max)
        self.coherent, self.role = coherent, role
        self.holder = {} if holder is None else holder
        self._state, self._call = None, "forward"
        return self

    def reset_state(self):
        self._state = None

    def _alpha(self, ep, en, z):
        with torch.no_grad():
            s = ((ep + en) / z).float()
            B, T, A = s.shape
            st = (self._state if (self._call == "step" and self._state is not None
                                  and self._state.shape[0] == B) else None)
            if st is None:
                st = torch.full((B, A), 1.0 / A, device=s.device, dtype=torch.float32)
            if T == 1:
                st = st + self.rho * (s[:, 0] - st)
                ubar = st.unsqueeze(1)
            elif T <= 4200:
                q = 1.0 - self.rho
                t = torch.arange(1, T + 1, device=s.device, dtype=torch.float64)
                acc = torch.cumsum((q ** -t).view(1, T, 1) * s.double(), dim=1)
                ubar = ((q ** t).view(1, T, 1)
                        * (st.double().unsqueeze(1) + self.rho * acc)).float()
                st = ubar[:, -1]
            else:
                ubar = torch.empty_like(s)
                for i in range(T):
                    st = st + self.rho * (s[:, i] - st)
                    ubar[:, i] = st
            self._state = st.clone() if self._call in ("prefill", "step") else None
            alpha = (1.0 + self.gamma * (A * ubar - 1.0)).clamp(self.a_min, self.a_max)
            if self.mode == "conserved":
                alpha = alpha * (A / alpha.sum(dim=-1, keepdim=True))
            return alpha.to(z.dtype)

    def forward(self, x):
        shared_route = self.coherent == "full" and self.role != "ref"
        if not shared_route:
            emb = F.normalize(x @ self.key_proj.T, dim=-1)
            C = F.normalize(self.dispatch, dim=-1)
            u = (emb @ C.T) / self.tau
            m = u.abs().amax(dim=-1, keepdim=True)
            ep, en = torch.exp(u - m), torch.exp(-u - m)
            z = (ep + en).sum(dim=-1, keepdim=True)
            g = (ep - en) / z
        if self.coherent and self.role != "ref":
            alpha = self.holder["alpha"]
            if shared_route:
                g = self.holder["g"]
        else:
            alpha = self._alpha(ep, en, z)
            if self.coherent:
                self.holder["alpha"], self.holder["g"] = alpha, g
        out = x
        for k, a in enumerate(self.anchors):
            if self.enabled[k]:
                out = out + (g[..., k:k + 1] * alpha[..., k:k + 1]) * (a(x) - x)
        return out


class BlockWithDispatch(nn.Module):
    def __init__(self, block: nn.Module, disp: AnchorDispatch):
        super().__init__()
        self.block = block
        self.disp = disp

    def forward(self, *args, **kwargs):
        out = self.block(*args, **kwargs)
        if isinstance(out, tuple):
            return (self.disp(out[0]),) + out[1:]
        return self.disp(out)

    def prefill(self, *args, **kwargs):
        out, cache = self.block.prefill(*args, **kwargs)
        self.disp._call = getattr(self.disp, "_call", "forward")
        if isinstance(self.disp, GainDispatch):
            self.disp._call = "prefill"
            try:
                return self.disp(out), cache
            finally:
                self.disp._call = "forward"
        return self.disp(out), cache

    def step(self, *args, **kwargs):
        out = self.block.step(*args, **kwargs)
        if isinstance(self.disp, GainDispatch):
            self.disp._call = "step"
            try:
                return self.disp(out[0]) if isinstance(out, tuple) else self.disp(out)
            finally:
                self.disp._call = "forward"
        if isinstance(out, tuple):
            return (self.disp(out[0]),) + out[1:]
        return self.disp(out)


# --------------------------------------------------------------- handles
def _probe(aleph) -> torch.Tensor:
    """Deterministic logits fingerprint on a tiny fixed input."""
    dev = next(aleph.parameters()).device
    ids = torch.arange(8, dtype=torch.long, device=dev).unsqueeze(0) % 7 + 1
    was_training = aleph.training
    aleph.eval()
    with torch.no_grad():
        out = aleph(ids)
    if was_training:
        aleph.train()
    logits = out[0] if isinstance(out, (tuple, list)) else out.logits
    return logits.detach().float().cpu()


class ArmHandle:
    """One mounted recipe. `detach(verify=True)` restores the core and
    proves the restoration is bit-exact."""

    def __init__(self, aleph, original, names, wrapped, fingerprint, card=None):
        self.aleph = aleph
        self._original = original
        self.names = list(names)
        self._blocks = wrapped
        self._fingerprint = fingerprint
        self.card = card or {}

    # -- masking (the toggle law) --------------------------------------
    def set_mask(self, mask: dict) -> None:
        enabled = [bool(mask.get(n, True)) for n in self.names]
        for b in self._blocks:
            if isinstance(b, BlockWithDispatch):
                b.disp.enabled = list(enabled)
            else:
                b.enabled = enabled[0]

    def enable(self, *names):
        self.set_mask({n: (n in names or not names) for n in self.names})

    def disable(self, *names):
        self.set_mask({n: n not in names for n in self.names})

    def _snapshot(self):
        return [list(b.disp.enabled) if isinstance(b, BlockWithDispatch) else b.enabled
                for b in self._blocks]

    def _restore(self, snap):
        for b, s in zip(self._blocks, snap):
            if isinstance(b, BlockWithDispatch):
                b.disp.enabled = list(s)
            else:
                b.enabled = s

    @contextlib.contextmanager
    def only(self, *names):
        prev = self._snapshot()
        self.set_mask({n: n in names for n in self.names})
        try:
            yield self
        finally:
            self._restore(prev)

    @contextlib.contextmanager
    def all_off(self):
        prev = self._snapshot()
        self.set_mask({n: False for n in self.names})
        try:
            yield self
        finally:
            self._restore(prev)

    # -- decode state ---------------------------------------------------
    def reset_state(self):
        """Forget the carried gain state (a fresh context starts over)."""
        for b in self._blocks:
            if isinstance(b, BlockWithDispatch):
                getattr(b.disp, "reset_state", lambda: None)()

    # -- detach ---------------------------------------------------------
    def detach(self, *, verify: bool = True):
        self.aleph.blocks = self._original
        if verify:
            post = _probe(self.aleph)
            if not torch.equal(post, self._fingerprint):
                raise RuntimeError(
                    "detach verification FAILED: the restored core's logits "
                    "are not bit-exact with the pre-mount fingerprint")
        return self.aleph


class StackHandle:
    """Nested always-on handles, detached outermost first."""

    def __init__(self, handles, names=None, card=None):
        self.handles = list(handles)
        self.names = list(names or [h.names[0] for h in self.handles])
        self.card = card or {}

    def reset_state(self):
        for h in self.handles:
            h.reset_state()

    def set_mask(self, mask: dict) -> None:
        for h, n in zip(self.handles, self.names):
            h.set_mask({h.names[0]: bool(mask.get(n, True))})

    @contextlib.contextmanager
    def all_off(self):
        prev = [h._snapshot() for h in self.handles]
        for h in self.handles:
            h.set_mask({h.names[0]: False})
        try:
            yield self
        finally:
            for h, p in zip(self.handles, prev):
                h._restore(p)

    def detach(self, *, verify: bool = True):
        for h in reversed(self.handles):
            h.detach(verify=verify)
        return self.handles[0].aleph if self.handles else None


# --------------------------------------------------------------- loading
def load_arm_file(path: str):
    """-> (tensors, meta). Reads the packaged safetensors arm; also accepts
    an amoe .pt anchor/dispatch checkpoint for local experiments."""
    if str(path).endswith(".safetensors"):
        from safetensors.torch import load_file
        from safetensors import safe_open
        with safe_open(path, framework="pt") as f:
            meta = dict(f.metadata() or {})
        card = json.loads(meta.get("arm", "{}"))
        return load_file(path), card
    blob = torch.load(path, map_location="cpu", weights_only=True)
    if "adapters" in blob:
        return blob["adapters"], dict(blob.get("meta") or {})
    if "dispatch" in blob:                      # flatten to "<block>.<key>"
        flat = {f"{i}.{k}": v for i, st in enumerate(blob["dispatch"])
                for k, v in st.items()}
        return flat, dict(blob.get("meta") or {})
    raise ValueError(f"unrecognized arm file: {path}")


def _per_block(state: dict, i: int) -> dict:
    pre = f"{i}."
    return {k[len(pre):]: v for k, v in state.items() if k.startswith(pre)}


def _target_blocks(aleph):
    return list(aleph.blocks)


# --------------------------------------------------------------- mounting
def mount_anchor(aleph, state: dict, name: str = "arm", card=None) -> ArmHandle:
    layers = _target_blocks(aleph)
    d = int(aleph.cfg.d_model)
    fingerprint = _probe(aleph)
    spec = spec_from_state(state)
    blocks, wrapped = [], []
    for i, layer in enumerate(layers):
        st = _per_block(state, i)
        if not st:
            blocks.append(layer)
            continue
        p = next(layer.parameters())
        a = RelayPatchwork(d, spec)
        a.load_state_dict(st, strict=True)
        a.to(device=p.device, dtype=p.dtype)
        w = BlockWithAdapter(layer, a)
        blocks.append(w)
        wrapped.append(w)
    original = aleph.blocks
    aleph.blocks = nn.ModuleList(blocks)
    return ArmHandle(aleph, original, [name], wrapped, fingerprint, card)


def mount_routed(aleph, state: dict, names, gain=None, mode="growth",
                 coherent=None, card=None) -> ArmHandle:
    layers = _target_blocks(aleph)
    d = int(aleph.cfg.d_model)
    fingerprint = _probe(aleph)
    A = len(names)
    member0 = {k[len("anchors.0."):]: v for k, v in _per_block(state, 0).items()
               if k.startswith("anchors.0.")}
    spec = spec_from_state(member0, prefix="")
    holder, blocks, wrapped = {}, [], []
    for i, layer in enumerate(layers):
        st = _per_block(state, i)
        p = next(layer.parameters())
        members = nn.ModuleList([RelayPatchwork(d, spec) for _ in range(A)])
        emb = int(st["key_proj"].shape[0])
        dp = AnchorDispatch(members, d, emb=emb)
        dp.load_state_dict(st, strict=True)
        dp.to(device=p.device, dtype=p.dtype)
        if gain:
            dp.__class__ = GainDispatch
            dp.configure(mode=mode, rho=gain["rho"], gamma=gain["gamma"],
                         a_min=gain["a_min"], a_max=gain["a_max"],
                         coherent=coherent, role="ref" if i == 0 else "follow",
                         holder=holder)
        w = BlockWithDispatch(layer, dp)
        blocks.append(w)
        wrapped.append(w)
    original = aleph.blocks
    aleph.blocks = nn.ModuleList(blocks)
    return ArmHandle(aleph, original, list(names), wrapped, fingerprint, card)


def mount(aleph, card: dict, files) -> "ArmHandle | StackHandle":
    """Attach one recipe. `card` is an arms/index.json row; `files` maps the
    row's weight keys to local paths (a stack passes one file per member,
    inner first)."""
    kind = (card.get("mount") or {}).get("type", "anchor")
    if kind == "anchor":
        state, _ = load_arm_file(files[0])
        return mount_anchor(aleph, state, card.get("name", "arm"), card)
    if kind == "stack":
        members = (card.get("mount") or {}).get("members") or []
        handles = []
        try:
            for f, n in zip(files, members or [f"member{i}" for i in range(len(files))]):
                state, _ = load_arm_file(f)
                handles.append(mount_anchor(aleph, state, n, card))
        except Exception:
            StackHandle(handles).detach(verify=False)   # never leave half a stack on
            raise
        return StackHandle(handles, members, card)
    if kind == "routed":
        state, _ = load_arm_file(files[0])
        m = card.get("mount") or {}
        return mount_routed(aleph, state, m.get("members") or [],
                            gain=m.get("gain"), mode=m.get("gain_mode", "growth"),
                            coherent=m.get("coherent"), card=card)
    raise ValueError(f"unknown mount type {kind!r} (packaged types: {ARM_TYPES})")


# --------------------------------------------------------------- frames
CHAT_HEADER = ("A conversation with Beatrix, a small byte-level "
               "language model still in pretraining.\n")
USER_TAG = "User: "
ASSISTANT_TAG = "Beatrix:"
TEXT_STOPS = ("\nUser:", "\nBeatrix:", "User:", "Beatrix:", "\n\n")
SPECIALS = {"DOC": 0xFF, "USER": 0xFE, "MODEL": 0xFD, "END": 0xFC, "SYS": 0xFB}


def render_text_frame(prompt: str, history=None) -> bytes:
    """The plain text frame every packaged arm was trained in:
    header + 'User: ...\\n' + 'Beatrix:' — the reply follows the tag."""
    out = CHAT_HEADER
    for turn in (history or []):
        tag = USER_TAG if turn.get("role") == "user" else ASSISTANT_TAG + " "
        out += f"{tag}{turn.get('content', '')}\n"
    out += f"{USER_TAG}{prompt}\n{ASSISTANT_TAG}"
    return out.encode("utf-8")


def render_specials_frame(prompt: str, system: str | None = None,
                          history=None) -> bytes:
    """The core's own annealed frame: [SYS] .. [END] [USER] .. [END] [MODEL].
    Encoded text can never contain a special, so the frame is unforgeable."""
    S = SPECIALS
    ids = []
    if system:
        ids += [S["SYS"]] + list(system.encode("utf-8")) + [S["END"]]
    for turn in (history or []):
        tag = S["USER"] if turn.get("role") == "user" else S["MODEL"]
        ids += [tag] + list(str(turn.get("content", "")).encode("utf-8")) + [S["END"]]
    ids += [S["USER"]] + list(prompt.encode("utf-8")) + [S["END"], S["MODEL"]]
    return bytes(ids)


def cut_text_reply(raw: bytes) -> bytes:
    """Cut a text-frame continuation at the first stop string."""
    best = len(raw)
    for s in TEXT_STOPS:
        j = raw.find(s.encode("utf-8"))
        if j != -1:
            best = min(best, j)
    return raw[:best]


def cut_specials_reply(raw: bytes) -> bytes:
    j = raw.find(bytes([SPECIALS["END"]]))
    return raw if j == -1 else raw[:j]