"""Arm mounting for the Beatrix Space.

`mount(model, recipe, path)` makes the live model carry exactly one arm
recipe and returns a handle whose `detach()` restores the bare blocks.
Every recipe is served from the files it was trained into; nothing is
re-derived at serve time.

Recipe types (the `mount.type` field of an arm-index row):
  anchor      one patchwork arm (an amoe anchor file)
  stack       several anchor arms applied always-on, nested inner to
              outer, with no mixer between them
  routed      a dense signed dispatch over member arms; the checkpoint
              carries the members and the router for every block
  routed + gain
              the same dispatch with a detached, bounded per-member gain
              on the numerator, read per site or once at the first site
              and shared by every site
  relay_ema   the patchwork organ plus two fixed-decay causal averages
              of its own read feeding its head
  recurrent   a small recurrent cell per site fed by the attention read
  program     a co-trained stage-arm program (mini-beatrix-3): one anchor
              arm per curriculum stage, trained together with the core,
              nested in stage order, each member switchable on its own;
              the arm sets fitted on the finished core mount the same way
  lookup      a zero-born causal softmax-attention arm after chosen
              blocks (amoe.arms.lookup), read from the lookup_* tensors
              of a trainer checkpoint or of an arm-only file

The relay_ema, recurrent, gain and lookup recipes carry state across
cached decode steps; a fresh prefill starts that state over, and the
uncached path treats every forward as a fresh context.

Trained arms and their measurements:
huggingface.co/AbstractPhil/alephllm-mini-beatrix-training
(mini-beatrix-2s/arms/index.json and arms/dispatch_index.json; the
mini-beatrix-3 stage arms ship beside every checkpoint as
mini-beatrix-3/arms/<arm>_step<N>.safetensors).
"""
from __future__ import annotations

import contextlib
import re

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

import amoe
from amoe.binding.resolver import resolve
from amoe.core.adapter import AdapterSpec, RelayPatchwork
from amoe.core.dispatch import AnchorDispatch, BlockWithDispatch
from amoe.io.checkpoint import load_anchor, load_dispatch
from amoe.runtime.attach import AttachHandle, _probe

MOUNT_TYPES = ("anchor", "stack", "routed", "relay_ema", "recurrent",
               "program", "lookup")


# ------------------------------------------------------------- geometry
def spec_from_state(adapters: dict, prefix: str | None = None) -> AdapterSpec:
    """Adapter geometry read from the arm's own tensors (the arm library
    holds more than one shape, and modules are built before state loads).
    `prefix` defaults to the first block the arm actually carries."""
    if prefix is None:
        blocks = sorted(int(k.split(".")[0]) for k in adapters
                        if k.endswith(".addr.codebook") and k.split(".")[0].isdigit())
        prefix = f"{blocks[0]}." if blocks else "0."
    cb = adapters.get(f"{prefix}addr.codebook")
    proj = adapters.get(f"{prefix}proj.weight")
    cons = adapters.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]))


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


def _load(path: str, *, weights_only: bool):
    """torch.load by CONTENT, never by name. torch >= 2.13 reads any path
    ending in .safetensors as a safetensors file, and the mini-beatrix-3
    arm files are torch archives under that extension (amoe's
    AnchorCheckpoint.save is torch.save): loaded by path they fail with
    'header too large'. An open file handle skips the name check."""
    with open(path, "rb") as fh:
        return torch.load(fh, map_location="cpu", weights_only=weights_only)


def _anchor(path: str):
    """amoe's load_anchor, fed an open handle for the same reason."""
    with open(path, "rb") as fh:
        return load_anchor(fh)


try:        # amoe >= 0.2.6: block anchors may be real safetensors files
    from amoe.io.safetensors_io import is_safetensors
except ImportError:       # pragma: no cover - older pins: torch files only
    def is_safetensors(_src):
        return False


def base_model_id(path: str) -> str | None:
    """The core an arm file was trained on, from its own metadata, read in
    whichever format the file actually is: a real safetensors block anchor
    through amoe's load_anchor (which reads either anchor format by
    content), every other arm file (anchor, dispatch, relay-EMA, recurrent
    torch archives) as the torch archive it is."""
    if is_safetensors(path):
        return _anchor(path).meta.get("base_model_id")
    blob = _load(path, weights_only=False)
    return (blob.get("meta") or {}).get("base_model_id")


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

    def __init__(self, handles):
        self.handles = list(handles)

    def reset_state(self):
        for h in self.handles:
            getattr(h, "reset_state", lambda: None)()

    def detach(self, *, verify: bool = False):
        for h in reversed(self.handles):
            h.detach(verify=verify)


def stage_index(phase) -> int | None:
    """curriculum_s3 -> 3; anything else -> None."""
    m = re.fullmatch(r"curriculum_s(\d+)", str(phase or ""))
    return int(m.group(1)) if m else None


class ProgramHandle(StackHandle):
    """A co-trained stage-arm program: every member is an anchor arm that
    trained in the same optimizer steps as the core, and each is nested
    around the blocks in the order the trainer attached it (stage order).

    Switching a member off sets its block wrappers to pass their input
    through unchanged, so the switch is an exact removal of that member's
    write, and the others read exactly what they would read without it.
    Nothing renormalizes. `detach` restores the bare blocks."""

    def __init__(self, handles, names, meta):
        super().__init__(handles)
        self.names = list(names)
        self.meta = list(meta)
        self.by_name = dict(zip(self.names, self.handles))
        self.live = {n: True for n in self.names}

    def set_live(self, live):
        want = set(live)
        for n, h in self.by_name.items():
            h.set_mask({n: n in want})
            self.live[n] = n in want

    @contextlib.contextmanager
    def masked(self, name):
        """One member off for the duration (the others as they are)."""
        with self.by_name[name].all_off():
            yield self

    def detach(self, *, verify: bool = False):
        # every member back on first: each member's detach check compares
        # against the model as it stood when that member attached, and
        # every member inside it was on then
        self.set_live(self.names)
        super().detach(verify=verify)


# ------------------------------------------------------------- gain dispatch
class GainDispatch(AnchorDispatch):
    """AnchorDispatch with a detached, bounded per-member gain alpha on the
    numerator (the denominator still sums every member, masked or not):

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

    mode 'growth' lets the mean gain rise toward the cap; 'conserved'
    renormalizes alpha to mean 1 over members at every byte.
    coherent None reads alpha at every site; 'alpha' reads it once at the
    first site and shares it; 'full' also shares the first site's signed
    routing weight. Never a selector: no top-k, no softmax over members."""

    def configure(self, *, mode, rho, gamma, a_min, a_max, coherent=None,
                  role="ref", holder=None):
        self.mode, self.rho, self.gamma = mode, rho, gamma
        self.a_min, self.a_max = a_min, a_max
        self.coherent, self.role = coherent, role
        self.holder = holder if holder is not None else {}
        self._state, self._call = None, "forward"
        return self

    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 GainBlock(BlockWithDispatch):
    """Marks cached-decode calls so the gain state is carried across steps."""

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

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


# ------------------------------------------------------------- mounts
def _mount_anchor(model, path):
    ck = _anchor(path)
    return amoe.attach(model, ck, spec=spec_from_state(ck.adapters), strict=False)


def _mount_stack(model, paths):
    handles = []
    try:
        for p in paths:
            handles.append(_mount_anchor(model, p))
    except Exception:
        StackHandle(handles).detach()      # never leave half a stack attached
        raise
    return StackHandle(handles)


def _mount_program(model, paths):
    """Nest every member in the order the trainer attached it: curriculum
    stage order, read from each anchor's own `phase` (the listing order is
    the fallback). The order matters because each member reads the block
    output the members inside it have already written to."""
    cks = [_anchor(p) for p in paths]
    order = sorted(range(len(cks)), key=lambda i: (
        stage_index(cks[i].meta.get("phase")) is None,
        stage_index(cks[i].meta.get("phase")) or 0, i))
    handles, names, meta = [], [], []
    try:
        for i in order:
            ck = cks[i]
            name = ck.meta.get("name") or f"arm{i}"
            handles.append(amoe.attach(model, ck,
                                       spec=spec_from_state(ck.adapters),
                                       strict=False))
            names.append(name)
            meta.append({"name": name, "phase": ck.meta.get("phase"),
                         "lambda": ck.meta.get("lambda"),
                         "n_params": ck.meta.get("n_params")})
    except Exception:
        StackHandle(handles).detach()      # never leave half a program on
        raise
    return ProgramHandle(handles, names, meta)


def _mount_routed(model, path):
    with open(path, "rb") as fh:
        ck = load_dispatch(fh)
    b = resolve(model, None)
    layers, d = list(b.layers(model)), b.hidden_size(model)
    fingerprint = _probe(model, b)
    names = list(ck.meta.get("anchors") or [])
    A = len(names)
    member0 = {k[len("anchors.0."):]: v for k, v in ck.dispatch[0].items()
               if k.startswith("anchors.0.")}
    spec = spec_from_state(member0, prefix="")
    gain = ck.meta.get("alpha")
    coherent = ck.meta.get("coherent_form")
    holder, blocks = {}, []
    for i, layer in enumerate(layers):
        dev = next(layer.parameters()).device
        members = nn.ModuleList([RelayPatchwork(d, spec) for _ in range(A)])
        dp = AnchorDispatch(members, d)
        dp.load_state_dict(ck.dispatch[i], strict=True)
        dp.to(dev)
        if gain:
            dp.__class__ = GainDispatch
            dp.configure(mode=ck.meta.get("alpha_mode") or "growth",
                         rho=float(gain["rho"]), gamma=float(gain["gamma"]),
                         a_min=float(gain["a_min"]), a_max=float(gain["a_max"]),
                         coherent=coherent, role="ref" if i == 0 else "follow",
                         holder=holder)
            blocks.append(GainBlock(layer, dp))
        else:
            blocks.append(BlockWithDispatch(layer, dp))
    b.set_layers(model, blocks)
    return AttachHandle(model, b, layers, names, blocks, fingerprint)


def _mount_relay_ema(model, path):
    from amoe.arms.relay_ema import RelayEMASpec, attach_relay_ema
    blob = _load(path, weights_only=False)
    st, sp = blob["state"], blob.get("spec") or {}
    spec = RelayEMASpec(base=spec_from_state(st), rho1=float(sp.get("rho1", 1 / 16)),
                        rho2=float(sp.get("rho2", 1 / 64)))
    name = (blob.get("meta") or {}).get("name", "relay_ema")
    h = attach_relay_ema(model, name, spec, zero_bias=False)
    with torch.no_grad():
        for i, w in enumerate(h.wraps()):
            w.arm.load_state_dict(_block_state(st, i), strict=True)
    return h


def _mount_recurrent(model, path):
    from amoe.arms.recurrent_wide import RecurrentSpec, attach_recurrent
    blob = _load(path, weights_only=False)
    name = (blob.get("meta") or {}).get("name", "recurrent")
    return attach_recurrent(model, name, RecurrentSpec(**blob["spec"]), state=blob["state"])


# ------------------------------------------------------------- lookup arm
LOOKUP_PREFIXES = ("lookup_arms.", "lookup_norms.")
_ST_DTYPES = {"BF16": torch.bfloat16, "F16": torch.float16,
              "F32": torch.float32, "F64": torch.float64}


def _st_tensor(meta, raw):
    t = torch.frombuffer(bytearray(raw), dtype=_ST_DTYPES[meta["dtype"]])
    return t.reshape(meta["shape"]).clone()


def fetch_lookup_tensors(repo: str, path: str, cache_dir: str) -> str:
    """The lookup arm's tensors out of a trainer checkpoint on the hub, by
    HTTP range reads: the safetensors header, then only the lookup_*
    tensors' bytes (about 34 MB of a 790 MB file per attention layer set).
    Written once to a local safetensors file under cache_dir (atomically,
    so two workers never read a half-written file) and reused after."""
    import json
    import os
    import struct
    import tempfile

    from huggingface_hub import HfFileSystem
    from safetensors.torch import save_file

    out = os.path.join(cache_dir, path.replace("/", "__") + ".lookup.safetensors")
    if os.path.exists(out):
        return out
    os.makedirs(cache_dir, exist_ok=True)
    with HfFileSystem().open(f"{repo}/{path}", "rb") as fh:
        n = struct.unpack("<Q", fh.read(8))[0]
        header = json.loads(fh.read(n))
        base = 8 + n
        spans = sorted((header[k]["data_offsets"][0], header[k]["data_offsets"][1], k)
                       for k in header if k.startswith(LOOKUP_PREFIXES))
        if not spans:
            raise KeyError(f"{path}: the checkpoint carries no lookup arm")
        lo, hi = spans[0][0], spans[-1][1]
        tensors = {}
        if hi - lo <= 2 * sum(e - b for b, e, _ in spans):    # one read
            fh.seek(base + lo)
            blob = fh.read(hi - lo)
            for b, e, k in spans:
                tensors[k] = _st_tensor(header[k], blob[b - lo:e - lo])
        else:
            for b, e, k in spans:
                fh.seek(base + b)
                tensors[k] = _st_tensor(header[k], fh.read(e - b))
    fd, tmp = tempfile.mkstemp(dir=cache_dir, suffix=".part")
    os.close(fd)
    save_file(tensors, tmp, metadata={"source": f"{repo}/{path}",
                                      "keys": str(len(tensors))})
    os.replace(tmp, out)
    return out


def _mount_lookup(model, path, name="lookup"):
    """A lookup arm: zero-born causal softmax attention after chosen blocks,
    each layer behind its own LayerNorm (amoe.arms.lookup). The geometry is
    read from the keys (sites, depth, heads from the q/k gains) and checked
    against the bound trunk; the attention runs in fp32 under autocast, as
    it was trained. Cached decode keeps the arm's own key/value cache."""
    from safetensors.torch import load_file
    from amoe.arms.lookup import attach_from_alephllm, from_alephllm_state
    sd = {k: v for k, v in load_file(path).items()
          if k.startswith(LOOKUP_PREFIXES)}
    _, sites, depth = from_alephllm_state(sd)
    b = resolve(model, None)
    n_blocks, d = len(list(b.layers(model))), b.hidden_size(model)
    if max(sites) >= n_blocks:
        raise ValueError(f"lookup arm sites {sites} outside a {n_blocks}-block trunk")
    w = sd[f"lookup_arms.{sites[0]}.qkv.weight"]
    if tuple(w.shape) != (3 * d, d):
        raise ValueError(f"lookup arm width {tuple(w.shape)} does not fit d={d}")
    return attach_from_alephllm(model, sd, name=name, attn_fp32=True)


def mount(model, recipe: dict, path: str, member_paths=None):
    """Attach one arm recipe to the bare model. `path` is the local file of
    the row's own checkpoint; `member_paths` are the local files of a
    stack's members, inner first, or of a stage program's members."""
    kind = (recipe or {}).get("type", "anchor")
    if kind == "lookup":
        return _mount_lookup(model, path, (recipe or {}).get("name", "lookup"))
    if kind == "anchor":
        return _mount_anchor(model, path)
    if kind == "stack":
        return _mount_stack(model, member_paths)
    if kind == "program":
        return _mount_program(model, member_paths)
    if kind == "routed":
        return _mount_routed(model, path)
    if kind == "relay_ema":
        return _mount_relay_ema(model, path)
    if kind == "recurrent":
        return _mount_recurrent(model, path)
    raise ValueError(f"unknown mount type {kind!r}")