"""Wildcat Mini served at Bonsai's own ternary size: student-base's language Linears (every PQ2_0 matmul weight of
Ternary Bonsai 2 27B) held on the GPU exactly as the GGUF stores them -- 2-bit codes (value (q-1)*d) plus one fp16
scale d per 128 weights, 2.125 bits/weight -- in the Hadamard-rotated input basis prism quantized them in.

The math (prism converter, undone by tools/gguf2hf.py, which made student-base):
    W_hf = W_t . H . D        W_t = ternary codes * group scales (the GGUF), H = block-1024 normalized Walsh-Hadamard
                              on the input axis, D = diag(explicit +-1 signs per input width)
so  y = W_hf x = W_t (H (D x)): flip the activation's signs, rotate each 1024 block with H (one small bf16 matmul,
shared by every Linear that reads the same tensor), then a plain matmul with the unpacked ternary weight. The 2-bit
codes are unpacked by one Triton kernel into a shared bf16 scratch buffer right before F.linear (as d3_q3.py).
token_embd is PQ2_0 too, with the inverse fold (plain row = (e_t H) D): rows are gathered packed and decoded per call.
The LoRA stays separate on the unrotated input: y += f * B(A x). in_proj_a / in_proj_b (BF16 in the GGUF), norms,
conv, A_log, dt_bias, the vision tower and the readout stay as student-base stores them.

At load (from student-base) every packed Linear is decoded back to W_t H D in fp32 and compared with student-base's
bf16 weight; a relative error above 1e-2 (a mapping bug, not rounding) stops the load.

    TERN_GGUF=<PQ2_0 gguf>  source of the codes (default: flamma's Ternary-Bonsai-2-27B-PQ2_0.gguf)
    TERN_SAVE=<file>        after packing, save everything the engine needs (packed + the non-ternary tensors)
    TERN_LOAD=<file>        build the backbone from that file alone (no 49 GB bf16 checkpoint read); .pt or
                            .safetensors. Default: <model dir>/mini-base.safetensors when it exists.
    PYTHONPATH=<model dir> python -m decision_index run --engine d3_tern:D3TernEngine --option model=<model dir> ...
        (the LoRA in <model dir>/adapter attaches on load; --option adapter=<dir> picks another)
"""
import json
import os
import re
import sys

import numpy as np
import torch
import triton
import triton.language as tl

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

GGUF = os.environ.get("TERN_GGUF",
                      "/mnt/tosh/llm-bench/bonsai2/Ternary-Bonsai-2-27B-gguf/Ternary-Bonsai-2-27B-PQ2_0.gguf")
GGUF_PY = os.environ.get("TERN_GGUF_PY", "/home/josh/llm-bench/llama.cpp-prism/gguf-py")
QK, GB, BS = 128, 34, 1024  # PQ2_0: 128 weights = fp16 d + 32 bytes of 2-bit codes; Hadamard block 1024
_SCRATCH = {}
_HAD = {}
_ROT = [None]  # (x, version, rotated): the last rotated activation, held so its address cannot be reused


@triton.jit
def _unpack2(c_ptr, s_ptr, out_ptr, K, BN: tl.constexpr, BK: tl.constexpr):
    n = tl.program_id(0) * BN + tl.arange(0, BN)[:, None]
    k = tl.program_id(1) * BK + tl.arange(0, BK)[None, :]
    c = tl.load(c_ptr + n * (K // 4) + k // 4).to(tl.int32)
    q = ((c >> ((k % 4) * 2)) & 3) - 1
    s = tl.load(s_ptr + n * (K // 128) + k // 128).to(tl.float32)
    tl.store(out_ptr + n * K + k, (q.to(tl.float32) * s).to(out_ptr.dtype.element_ty))


def unpack(codes, scales, out):
    n, k = out.shape
    _unpack2[(n // 16, k // 128)](codes, scales, out, k, BN=16, BK=128)
    return out


def hadamard(device, dtype):
    key = (str(device), dtype)
    if key not in _HAD:
        h = torch.ones(1, 1, dtype=torch.float64)
        while h.shape[0] < BS:
            h = torch.cat([torch.cat([h, h], 1), torch.cat([h, -h], 1)], 0)
        _HAD[key] = (h / BS ** 0.5).to(device, dtype)  # entries +-1/32: exact in bf16
    return _HAD[key]


def unrotate(w, signs):
    """fp32 W_t [N, K] -> W_t H D (what gguf2hf.py stored, before its bf16 cast)."""
    n, k = w.shape
    return (w.view(n, k // BS, BS) @ hadamard(w.device, torch.float32)).view(n, k) * signs.to(w.device, torch.float32)


class Rotation:
    """Per-width explicit signs; rotate(x) = H(D x) in x's dtype (fp32 accumulation), cached for the last x."""
    signs = {}

    @classmethod
    def rotate(cls, x):
        last = _ROT[0]
        try:
            ver = x._version
        except RuntimeError:  # inference-mode tensors keep no version counter: identity alone
            ver = None
        if last is not None and last[0] is x and last[1] == ver:
            return last[2]
        k = x.shape[-1]
        s = cls.signs[(k, x.device, x.dtype)] if (k, x.device, x.dtype) in cls.signs else cls._signs(k, x)
        r = ((x * s).reshape(-1, BS) @ hadamard(x.device, x.dtype)).reshape(x.shape)
        _ROT[0] = (x, ver, r)
        return r

    @classmethod
    def _signs(cls, k, x):
        s = cls.signs[(k, x.device, x.dtype)] = cls.signs[k].to(x.device, x.dtype)
        return s


class TernLinear(torch.nn.Module):
    def __init__(self, codes, scales, bias=None):
        super().__init__()
        self.out_features, self.in_features = codes.shape[0], codes.shape[1] * 4
        assert self.out_features % 16 == 0 and self.in_features % BS == 0
        self.register_buffer("codes", codes)
        self.register_buffer("scales", scales)
        self.bias = bias
        self.lora = None  # (A [r, K], B [N, r] already times alpha / r)

    def unpacked(self, dtype=torch.bfloat16):
        n, k = self.out_features, self.in_features
        if dtype != torch.bfloat16:
            return unpack(self.codes, self.scales, torch.empty(n, k, dtype=dtype, device=self.codes.device))
        buf = _SCRATCH.get(self.codes.device)
        if buf is None or buf.numel() < n * k:
            _SCRATCH[self.codes.device] = buf = torch.empty(n * k, dtype=torch.bfloat16, device=self.codes.device)
        return unpack(self.codes, self.scales, buf[: n * k].view(n, k))

    def forward(self, x):
        y = torch.nn.functional.linear(Rotation.rotate(x), self.unpacked(), self.bias)
        if self.lora is not None:
            a, b = self.lora
            y = y + torch.nn.functional.linear(torch.nn.functional.linear(x.to(a.dtype), a), b).to(y.dtype)
        return y


class TernEmbedding(torch.nn.Module):
    """PQ2_0 token_embd with the inverse fold: row = ((codes * d) H) D, decoded per call in fp32, returned in bf16."""

    def __init__(self, codes, scales, dtype=torch.bfloat16):
        super().__init__()
        self.num_embeddings, self.embedding_dim = codes.shape[0], codes.shape[1] * 4
        self.padding_idx = None
        self.register_buffer("codes", codes)
        self.register_buffer("scales", scales)
        self.out_dtype = dtype

    def plain(self, ids):
        """fp32 (codes * d) rows, before the fold; the unpack kernel works on 16-row tiles, so pad."""
        n = ids.shape[0]
        pad = (-n) % 16
        if pad:
            ids = torch.cat([ids, ids.new_zeros(pad)])
        w = unpack(self.codes[ids], self.scales[ids],
                   torch.empty(ids.shape[0], self.embedding_dim, dtype=torch.float32, device=self.codes.device))
        return w[:n]

    def rows(self, ids):
        return unrotate(self.plain(ids).contiguous(), Rotation.signs[self.embedding_dim])

    def forward(self, ids):
        flat = ids.reshape(-1)
        out = self.rows(flat).to(self.out_dtype)
        return out.view(*ids.shape, self.embedding_dim)


# ---------------------------------------------------------------- GGUF -> packed, in student-base's row order

def _gguf():
    sys.path.insert(0, GGUF_PY)
    from gguf import GGUFReader
    return GGUFReader(GGUF)


def _field(r, k):
    x = r.fields[k]
    vals = [x.parts[i] for i in x.data]
    if x.types[-1].name == "STRING":
        return [bytes(v).decode() for v in vals]
    return [v.tolist()[0] for v in vals]


HF_OF = {"attn_q": "self_attn.q_proj", "attn_k": "self_attn.k_proj", "attn_v": "self_attn.v_proj",
         "attn_output": "self_attn.o_proj", "ffn_gate": "mlp.gate_proj", "ffn_up": "mlp.up_proj",
         "ffn_down": "mlp.down_proj", "attn_qkv": "linear_attn.in_proj_qkv", "attn_gate": "linear_attn.in_proj_z",
         "ssm_out": "linear_attn.out_proj"}


def _row_order(kind, n, tc):
    """Row permutation gguf2hf.py applies (linear-attention V heads tiled [rep, nk, hd] -> grouped [nk, rep, hd])."""
    nk, nv, hd = tc["linear_num_key_heads"], tc["linear_num_value_heads"], tc["linear_value_head_dim"]
    rep = nv // nk

    def grouped(m):
        return torch.arange(m).view(rep, nk, hd).transpose(0, 1).reshape(-1)
    if kind == "attn_qkv":
        qk = 2 * nk * tc["linear_key_head_dim"]
        return torch.cat([torch.arange(qk), qk + grouped(n - qk)])
    if kind == "attn_gate":
        return grouped(n)
    return None


def read_gguf(tc):
    """{hf name (relative to language_model): (codes uint8 [N, K/4], scales fp16 [N, K/128])} + per-width signs."""
    r = _gguf()
    assert _field(r, "prism.hadamard.block_size") == [BS] and _field(r, "prism.hadamard.sign_mode") == ["explicit"]
    assert _field(r, "prism.hadamard.transform") == ["normalized-sylvester-walsh-hadamard"]
    assert _field(r, "prism.hadamard.gdn_v_grouped") == [True]  # ssm_out input axis already in HF order
    rotated = set(_field(r, "prism.hadamard.weight_names"))
    inverse = set(_field(r, "prism.hadamard.inverse_weight_names"))
    vals, o = np.array(_field(r, "prism.hadamard.sign_values"), np.float32), 0
    for w in _field(r, "prism.hadamard.sign_widths"):
        Rotation.signs[w] = torch.from_numpy(vals[o:o + w].copy()); o += w
    out = {}
    for t in r.tensors:
        if t.tensor_type.name != "PQ2_0" or t.name == "output.weight":
            continue
        k, n = int(t.shape[0]), int(np.prod([int(s) for s in t.shape[1:]]))
        raw = np.asarray(t.data).reshape(n, k // QK, GB)
        scales = torch.from_numpy(raw[:, :, :2].copy().view(np.float16).reshape(n, k // QK))
        codes = torch.from_numpy(raw[:, :, 2:].copy().reshape(n, k // 4))
        if t.name == "token_embd.weight":
            assert t.name in inverse
            out["embed_tokens"] = (codes, scales)
            continue
        m = re.fullmatch(r"blk\.(\d+)\.(\w+)\.weight", t.name)
        if int(m[1]) >= tc["num_hidden_layers"]:
            continue  # MTP head
        assert t.name in rotated, t.name
        order = _row_order(m[2], n, tc)
        if order is not None:
            codes, scales = codes[order].contiguous(), scales[order].contiguous()
        out[f"layers.{m[1]}.{HF_OF[m[2]]}"] = (codes, scales)
    return out


def _set(model, name, mod):
    parent, _, child = name.rpartition(".")
    setattr(model.get_submodule(parent) if parent else model, child, mod)


def check_layer(name, wt, ref, stats):
    """W_t (fp32, decoded from the packed tensors) -> W_t H D, compared with student-base's bf16 weight."""
    w = unrotate(wt, Rotation.signs[wt.shape[1]])
    ref = ref.to(w.device, torch.float32)
    d = (w.to(torch.bfloat16).float() - ref).abs()
    rel = float((w - ref).abs().max() / ref.abs().max())
    stats["tensors"] += 1
    stats["elements"] += ref.numel()
    stats["bf16_mismatch"] += int((d > 0).sum())
    stats["max_abs_bf16"] = max(stats["max_abs_bf16"], float(d.max()))
    stats["sum_abs_bf16"] += float(d.double().sum())
    stats["max_rel_fp32"] = max(stats["max_rel_fp32"], rel)
    if rel > 1e-2:  # bf16 rounding alone is up to 2^-8 relative; a mapping bug is O(1)
        raise RuntimeError(f"{name}: packed ternary weight does not reproduce student-base (rel max err {rel:.3g})")


def pack_model(model, device, check=True):
    lm = model.language_model
    tc = model.config.text_config.to_dict()
    packed = read_gguf(tc)
    stats = dict(tensors=0, elements=0, bf16_mismatch=0, max_abs_bf16=0.0, sum_abs_bf16=0.0, max_rel_fp32=0.0)
    for name, (codes, scales) in packed.items():
        old = lm.get_submodule(name)
        if name == "embed_tokens":
            m = TernEmbedding(codes.to(device), scales.to(device))
            ref_w = old.weight.data
            if check:
                ids = torch.arange(m.num_embeddings, device=device)
                for i in range(0, m.num_embeddings, 16384):
                    check_layer(name, m.plain(ids[i:i + 16384]), ref_w[i:i + 16384], stats)
        else:
            assert old.weight.shape == (codes.shape[0], codes.shape[1] * 4), (name, old.weight.shape)
            bias = None if old.bias is None else torch.nn.Parameter(old.bias.data.to(device), requires_grad=False)
            m = TernLinear(codes.to(device), scales.to(device), bias)
            if check:
                check_layer(name, m.unpacked(torch.float32), old.weight.data, stats)
        _set(lm, name, m)
        del old
    left = [n for n, mod in lm.named_modules() if isinstance(mod, torch.nn.Linear)]
    assert all(n.endswith(("in_proj_a", "in_proj_b")) for n in left), left[:5]
    if check:
        stats["mean_abs_bf16"] = stats["sum_abs_bf16"] / stats["elements"]
    return len(packed), stats


def _rest_to(model, device):
    for m in model.modules():
        if isinstance(m, (TernLinear, TernEmbedding)):
            continue
        for k, p in list(m.named_parameters(recurse=False)):
            m._parameters[k] = torch.nn.Parameter(p.data.to(device), requires_grad=False)
        for k, b in list(m.named_buffers(recurse=False)):
            m._buffers[k] = b.to(device)


def save_packed(model, path):
    packed, rest = {}, {}
    for n, m in model.named_modules():
        if isinstance(m, (TernLinear, TernEmbedding)):
            packed[n] = {"codes": m.codes.cpu(), "scales": m.scales.cpu(),
                         **({"bias": m.bias.cpu()} if getattr(m, "bias", None) is not None else {})}
    pk = tuple(packed)
    for k, v in model.state_dict().items():
        if not any(k.startswith(p + ".") for p in pk):
            rest[k] = v.cpu()
    signs = {k: v for k, v in Rotation.signs.items() if isinstance(k, int)}
    torch.save({"format": "d3-tern-pq2-v1", "gguf": GGUF, "signs": signs, "packed": packed, "rest": rest}, path)
    return len(packed), len(rest)


def _read_packed(path):
    """A TERN_SAVE dict, from torch.save (.pt) or flattened into safetensors (packed/<module>/<field>, rest/<key>,
    signs/<width>)."""
    if not path.endswith(".safetensors"):
        return torch.load(path, map_location="cpu")
    from safetensors.torch import load_file
    pk = {"packed": {}, "rest": {}, "signs": {}}
    for k, v in load_file(path).items():
        kind, _, name = k.partition("/")
        if kind == "packed":
            mod, _, field = name.rpartition("/")
            pk["packed"].setdefault(mod, {})[field] = v
        elif kind == "rest":
            pk["rest"][name] = v
        else:
            pk["signs"][int(name)] = v
    return pk


def load_packed(cls, path, root, device):
    """Backbone from a TERN_SAVE file: weights-free skeleton (buffers real), packed modules, then the rest."""
    from accelerate import init_empty_weights
    pk = _read_packed(path)
    Rotation.signs.update(pk["signs"])
    config = cls.config_class.from_pretrained(root)
    config._attn_implementation = "sdpa"
    with init_empty_weights(include_buffers=False):
        model = cls._from_config(config, dtype=torch.bfloat16)
    for n, t in pk["packed"].items():
        if n.endswith("embed_tokens"):
            mod = TernEmbedding(t["codes"].to(device), t["scales"].to(device))
        else:
            b = t.get("bias")
            mod = TernLinear(t["codes"].to(device), t["scales"].to(device),
                             None if b is None else torch.nn.Parameter(b.to(device), requires_grad=False))
        _set(model, n, mod)
    res = model.load_state_dict(pk["rest"], strict=False, assign=True)
    real_missing = [k for k in res.missing_keys if not any(k.startswith(p + ".") for p in pk["packed"])]
    if real_missing or res.unexpected_keys:
        raise RuntimeError(f"TERN_LOAD: {len(real_missing)} missing ({real_missing[:3]}), "
                           f"{len(res.unexpected_keys)} unexpected ({res.unexpected_keys[:3]})")
    model.eval()
    _rest_to(model, device)
    meta = [n for n, p in list(model.named_parameters()) + list(model.named_buffers()) if p.device.type == "meta"]
    if meta:
        raise RuntimeError(f"TERN_LOAD left {len(meta)} tensors on meta: {meta[:3]}")
    return model, len(pk["packed"])


class _TernLoader:
    def __init__(self, cls):
        self.cls = cls

    def from_pretrained(self, path, *args, **kwargs):
        device = torch.device(kwargs["device_map"][""])
        packed = os.environ.get("TERN_LOAD") or os.path.join(path, "mini-base.safetensors")
        if os.path.exists(packed):
            model, n = load_packed(self.cls, packed, path, device)
            info = {"packed": n, "source": os.path.basename(packed)}
        else:
            kwargs["device_map"] = {"": "cpu"}
            model = self.cls.from_pretrained(path, *args, **kwargs)
            n, stats = pack_model(model, device, check=os.environ.get("TERN_CHECK", "1") == "1")
            _rest_to(model, device)
            info = {"packed": n, "source": GGUF, "check_vs_student_base": stats}
            if os.environ.get("TERN_SAVE"):
                info["saved"] = dict(zip(("packed", "rest"), save_packed(model, os.environ["TERN_SAVE"])))
        torch.cuda.empty_cache()
        info["gib_allocated"] = round(torch.cuda.memory_allocated(device) / 2**30, 3)
        LOAD_INFO.update(info)
        print(f"tern: {json.dumps(info)}", file=sys.stderr, flush=True)
        return model


LOAD_INFO = {}


def attach_adapter(decision, path, scale=1.0):
    from safetensors.torch import load_file
    lm = decision.backbone.language_model
    meta = json.load(open(os.path.join(path, "adapter.json")))
    f = scale * meta["lora_alpha"] / meta["r"]
    state = load_file(os.path.join(path, "adapter.safetensors"))
    n = 0
    for key, a in state.items():
        if not key.endswith(".lora_A.weight"):
            continue
        mod = key[: -len(".lora_A.weight")]
        m = lm.get_submodule(mod)
        if not isinstance(m, TernLinear):
            raise RuntimeError(f"adapter targets {mod}, which is not a packed ternary Linear")
        b = state[mod + ".lora_B.weight"]
        m.lora = (a.to(m.codes.device, torch.bfloat16), (f * b.float()).to(m.codes.device, torch.bfloat16))
        n += 1
    readout = os.path.join(path, "readout.safetensors")
    if os.path.exists(readout):
        decision.readout = load_file(readout)["weight"].to(decision.device, getattr(torch, decision.readout_dtype))
    return {"adapter": os.path.abspath(path), "attached": n, "readout": os.path.exists(readout), **meta}


def install():
    import d3_q3
    import d3_runtime
    if not getattr(d3_runtime, "_tern", False):
        base = d3_runtime.backbone_class
        d3_runtime.backbone_class = lambda kind: _TernLoader(base(kind))
        if hasattr(d3_runtime.D3, "_run_state_once"):  # newer runtimes (d3-ultra); student-base's has no state-once
            d3_q3._RUN_STATE_ONCE.append(d3_runtime.D3._run_state_once)
            d3_runtime.D3._run_state_once = d3_q3._run_state_once_or_plain
        d3_runtime._tern = True


from d3_engine import D3Engine  # noqa: E402


class D3TernEngine(D3Engine):
    name = "wildcat-mini-tern"

    def __init__(self, adapter=None, lora_scale=1.0, **options):
        import d3_q3
        install()
        device = options.get("device", "cuda:0")
        total = torch.cuda.get_device_properties(device).total_memory
        so_div = int(os.environ.get("Q3_SO_DIV", "4" if total < 40 * 2**30 else "1"))
        import d3_runtime
        if not hasattr(d3_runtime, "STATE_ONCE_ROWS"):
            so_div = 0  # no state-once pass in this runtime: nothing to shrink
        if so_div > 1:
            d3_q3._small_card(so_div)
        options.setdefault("verify", "none")  # d3's MODEL_MANIFEST describes d3's weights, not these
        if adapter is None and os.path.exists(os.path.join(options.get("model", ""), "adapter", "adapter.json")):
            adapter = os.path.join(options["model"], "adapter")
        super().__init__(**options)
        self.provenance = dict(self.provenance)
        info = {"quant": "Bonsai PQ2_0 ternary (2-bit codes + fp16 scale / 128), Hadamard-rotated activations, "
                         "Triton unpack per matmul", "load": dict(LOAD_INFO),
                "card_gib": round(total / 2**30, 1), "state_once_div": so_div}
        if adapter:
            info["adapter"] = attach_adapter(self.decision, adapter, float(lora_scale))
        info["gib_allocated"] = round(torch.cuda.memory_allocated(device) / 2**30, 3)
        self.provenance["wildcat"] = info
        print(f"tern engine {json.dumps(info)}", file=sys.stderr, flush=True)

    def __call__(self, state, questions, images=None, videos=None):
        """student-base's runtime has no state-once fallback: a request that runs out of memory is retried with the
        questions per forward pass halved (down to 1) instead of stopping the run. Counted in OOM_RETRIES."""
        bs = self.decision.batch_size
        try:
            while True:
                try:
                    return super().__call__(state, questions, images, videos)
                except torch.OutOfMemoryError:
                    if self.decision.batch_size == 1:
                        raise
                # retry outside `except`: there the traceback still pins the failed pass's tensors
                torch.cuda.empty_cache()
                self.decision.batch_size = max(1, self.decision.batch_size // 2)
                OOM_RETRIES[0] += 1
                print(f"tern oom retry {OOM_RETRIES[0]}: batch_size -> {self.decision.batch_size}",
                      file=sys.stderr, flush=True)
        finally:
            self.decision.batch_size = bs


OOM_RETRIES = [0]