"""Wildcat served for real at 3 bits: d3-ultra's Linears (outside SKIP) held packed on the GPU as 3-bit HQQ codes
(group 64, a 2-bit plane + a 1-bit plane, bf16 scale/zero per group), unpacked by one Triton kernel into a shared bf16
scratch buffer right before each matmul. The LoRA adapters stay separate (y += f * B(A x)); several adapters are a
seed soup (their A/B stacked, each at lora_scale / count), the same sum of deltas d3_fakeq.py merges into bf16.

The unpacked weight is bit-identical to HQQLinear.dequantize() (checked per layer at quantize time), so this computes
what d3_fakeq.py computes, with ~3.5 bits per weight resident instead of 16.

    Q3_SAVE=<file>  after quantizing, save the packed tensors (the release checkpoint)
    Q3_LOAD=<file>  load packed tensors instead of quantizing (no bf16 copy of the quantized Linears is kept)
    PYTHONPATH=<d3ultra dir>:<this dir> python -m decision_index run --engine d3_q3:D3Q3Engine \
        --option model=<d3ultra dir> --option adapter=<out>[,<out2>] [--option lora_scale=1.0] ...
"""
import json
import os
import sys

import torch
import triton
import triton.language as tl

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

SKIP = ["visual", "in_proj_a", "in_proj_b", "lm_head"]
BITS, GROUP = 3, 64
_SCRATCH = {}


@triton.jit
def _unpack3(lo_ptr, hi_ptr, s_ptr, z_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, :]
    lo = tl.load(lo_ptr + n * (K // 4) + k // 4).to(tl.int32)
    hi = tl.load(hi_ptr + n * (K // 8) + k // 8).to(tl.int32)
    code = ((lo >> ((k % 4) * 2)) & 3) | (((hi >> (k % 8)) & 1) << 2)
    g = n * (K // 64) + k // 64
    z = tl.load(z_ptr + g).to(tl.float32)
    s = tl.load(s_ptr + g).to(tl.float32)
    t = (code.to(tl.float32) - z).to(tl.bfloat16)  # rounded per op, as torch's bf16 (W_q - zero) * scale
    tl.store(out_ptr + n * K + k, (t.to(tl.float32) * s).to(tl.bfloat16))


def pack(codes):
    """uint8 codes [N, K] in 0..7 -> (2-bit plane [N, K/4], 1-bit plane [N, K/8])."""
    c = codes.to(torch.int32)
    lo = (c & 3).view(c.shape[0], -1, 4)
    hi = (c >> 2).view(c.shape[0], -1, 8)
    lo = lo[..., 0] | lo[..., 1] << 2 | lo[..., 2] << 4 | lo[..., 3] << 6
    hi = sum(hi[..., i] << i for i in range(8))
    return lo.to(torch.uint8).contiguous(), hi.to(torch.uint8).contiguous()


class Q3Linear(torch.nn.Module):
    def __init__(self, lo, hi, scale, zero, bias=None):
        super().__init__()
        self.out_features, self.in_features = lo.shape[0], lo.shape[1] * 4
        assert self.out_features % 16 == 0 and self.in_features % 64 == 0
        for k, v in dict(lo=lo, hi=hi, scale=scale, zero=zero).items():
            self.register_buffer(k, v)
        self.bias = bias
        self.lora = None  # (A [r, K], B [N, r] already times the soup factor)

    def unpacked(self):
        n, k = self.out_features, self.in_features
        buf = _SCRATCH.get(self.lo.device)
        if buf is None or buf.numel() < n * k:
            _SCRATCH[self.lo.device] = buf = torch.empty(n * k, dtype=torch.bfloat16, device=self.lo.device)
        w = buf[: n * k].view(n, k)
        _unpack3[(n // 16, k // 64)](self.lo, self.hi, self.scale, self.zero, w, k, BN=16, BK=64)
        return w

    def forward(self, x):
        y = torch.nn.functional.linear(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


def _targets(model):
    return [(n, m) for n, m in model.named_modules()
            if isinstance(m, torch.nn.Linear) and not any(s in n.split(".") or n.startswith(s) for s in SKIP)]


def quantize(model, device):
    """Each target Linear -> Q3Linear on `device` via HQQLinear, checked bit-exact against dequantize()."""
    from hqq.core.quantize import BaseQuantizeConfig, HQQLinear, Quantizer
    saved = {}
    for name, lin in _targets(model):
        n, k = lin.out_features, lin.in_features
        q = HQQLinear(lin, BaseQuantizeConfig(nbits=BITS, group_size=GROUP), compute_dtype=torch.bfloat16,
                      device=device, del_orig=False)
        codes = Quantizer.unpack[q.meta["packing"]](q.W_q, dtype=torch.uint8).flatten()[: n * k].view(n, k)
        scale = q.meta["scale"].to(torch.bfloat16).reshape(n, k // GROUP).contiguous()
        zero = q.meta["zero"].to(torch.bfloat16).reshape(n, k // GROUP).contiguous()
        lo, hi = pack(codes)
        bias = None if lin.bias is None else torch.nn.Parameter(lin.bias.data.to(device), requires_grad=False)
        m = Q3Linear(lo, hi, scale, zero, bias)
        ref = q.dequantize().view(n, k).to(torch.bfloat16)
        diff = (m.unpacked().float() - ref.float()).abs().max().item()
        if diff != 0.0:
            raise RuntimeError(f"{name}: unpacked 3-bit weight differs from HQQ dequantize() by {diff}")
        if os.environ.get("Q3_SAVE"):
            saved[name] = {k_: v.cpu() for k_, v in dict(lo=lo, hi=hi, scale=scale, zero=zero).items()}
        parent, _, child = name.rpartition(".")
        setattr(model.get_submodule(parent) if parent else model, child, m)
        del q, ref, lin
    if saved:
        torch.save({"bits": BITS, "group": GROUP, "skip": SKIP, "layers": saved}, os.environ["Q3_SAVE"])
    return len(saved) or sum(isinstance(m, Q3Linear) for m in model.modules())


def load_packed(model, path, device):
    pk = torch.load(path, map_location="cpu")
    for name, lin in _targets(model):
        t = {k: v.to(device) for k, v in pk["layers"][name].items()}
        bias = None if lin.bias is None else torch.nn.Parameter(lin.bias.data.to(device), requires_grad=False)
        parent, _, child = name.rpartition(".")
        setattr(model.get_submodule(parent) if parent else model, child, Q3Linear(**t, bias=bias))
    return len(pk["layers"])


def _rest_to(model, device):
    for m in model.modules():
        if isinstance(m, Q3Linear):
            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 load_package(cls, path, device, kwargs):
    """A Wildcat package (q3.safetensors + model-rest.safetensors): build the backbone with empty weights, put the packed
    Linears and the remaining tensors straight on `device`. No bf16 copy of the model is ever made."""
    from accelerate import init_empty_weights
    from accelerate.utils import set_module_tensor_to_device
    from safetensors import safe_open
    config = cls.config_class.from_pretrained(path)
    with init_empty_weights(include_buffers=False):
        model = cls._from_config(config, attn_implementation=kwargs.get("attn_implementation", "sdpa"),
                                 dtype=kwargs.get("dtype", torch.bfloat16))
    targets = _targets(model)
    with safe_open(os.path.join(path, "q3.safetensors"), "pt", device=str(device)) as f:
        for name, lin in targets:
            if lin.bias is not None:
                raise RuntimeError(f"{name}: packed Linears with a bias are not supported")
            parent, _, child = name.rpartition(".")
            setattr(model.get_submodule(parent) if parent else model, child,
                    Q3Linear(**{k: f.get_tensor(f"{name}.{k}") for k in ("lo", "hi", "scale", "zero")}))
    with safe_open(os.path.join(path, "model-rest.safetensors"), "pt", device=str(device)) as f:
        for key in f.keys():
            set_module_tensor_to_device(model, key, device, value=f.get_tensor(key))
    _rest_to(model, device)
    empty = [n for n, t in list(model.named_parameters()) + list(model.named_buffers()) if t.device.type == "meta"]
    if empty:
        raise RuntimeError(f"{len(empty)} tensors were not in the package, e.g. {empty[:3]}")
    return model


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

    def from_pretrained(self, path, *args, **kwargs):
        device = kwargs["device_map"][""]
        if os.path.exists(os.path.join(path, "q3.safetensors")):
            model = load_package(self.cls, path, device, kwargs)
            n = sum(isinstance(m, Q3Linear) for m in model.modules())
        else:
            kwargs["device_map"] = {"": "cpu"}
            model = self.cls.from_pretrained(path, *args, **kwargs)
            n = load_packed(model, os.environ["Q3_LOAD"], device) if os.environ.get("Q3_LOAD") else quantize(model, device)
            _rest_to(model, device)
        torch.cuda.empty_cache()
        print(f"q3: {n} packed 3-bit Linears, {torch.cuda.memory_allocated(device) / 2**30:.2f} GiB on {device}",
              file=sys.stderr, flush=True)
        return model


def stack_adapters(paths, scale):
    """{module: (A [sum r, K], B [N, sum r])} with y += B A x = sum_i f_i * B_i A_i x, f_i = scale / len(paths) *
    alpha_i / r_i: every adapter's delta at an equal share (a seed soup)."""
    from safetensors.torch import load_file
    parts = {}
    for path in paths:
        meta = json.load(open(os.path.join(path, "adapter.json")))
        f = scale / len(paths) * meta["lora_alpha"] / meta["r"]
        state = load_file(os.path.join(path, "adapter.safetensors"))
        for key, a in state.items():
            if key.endswith(".lora_A.weight"):
                mod = key[: -len(".lora_A.weight")]
                parts.setdefault(mod, []).append((a.float(), f * state[mod + ".lora_B.weight"].float()))
    return {mod: (torch.cat([a for a, _ in ab]).to(torch.bfloat16),
                  torch.cat([b for _, b in ab], dim=1).to(torch.bfloat16)) for mod, ab in parts.items()}


def attach(lm, stacked):
    for mod, (a, b) in stacked.items():
        m = lm.get_submodule(mod)
        if not isinstance(m, Q3Linear):
            raise RuntimeError(f"adapter targets {mod}, which is not a packed 3-bit Linear")
        m.lora = (a.to(m.lo.device), b.to(m.lo.device))
    return len(stacked)


def load_lora(path):
    """lora.safetensors of a package: <module>.A / <module>.B, already stacked and scaled."""
    from safetensors.torch import load_file
    state = load_file(path)
    return {k[:-2]: (state[k], state[k[:-2] + ".B"]) for k in state if k.endswith(".A")}


FALLBACKS = [0]


def _small_card(so_div):
    """On a card under 40 GiB (a 24 GB consumer card) the state-once question pass of wide requests does not fit next
    to the weights: shrink its batch bounds by so_div (grouping only; every question block still sees the same state
    cache). Big cards keep the runtime's bounds."""
    import d3_runtime
    for k in ("STATE_ONCE_ROWS", "STATE_ONCE_KEYS", "STATE_ONCE_SCORES"):
        setattr(d3_runtime, k, max(1, getattr(d3_runtime, k) // so_div))


def _run_state_once_or_plain(self, prepared):
    """State-once first, then state-once in smaller question batches (bounds / 4, / 16: grouping only, the same state
    cache); a request that still does not fit runs on the reference path (D3_STATE_ONCE=0's one-row-per-question path,
    which needs less memory) instead of stopping the run. A ToolRet request of 32 long candidates did not fit at full
    bounds on a 96 GB card, and the reference path takes minutes for it."""
    import d3_runtime
    names = ("STATE_ONCE_ROWS", "STATE_ONCE_KEYS", "STATE_ONCE_SCORES")
    bounds = [getattr(d3_runtime, k) for k in names]
    for div in (1, 4, 16):
        for k, v in zip(names, bounds):
            setattr(d3_runtime, k, max(1, v // div))
        try:
            return _RUN_STATE_ONCE[0](self, prepared)
        except torch.OutOfMemoryError:
            pass  # retry below: inside `except` the traceback still pins the failed pass's tensors
        finally:
            for k, v in zip(names, bounds):
                setattr(d3_runtime, k, v)
        torch.cuda.empty_cache()
        print(f"q3 state-once OOM at bounds / {div}", file=sys.stderr, flush=True)
    torch.cuda.empty_cache()
    FALLBACKS[0] += 1
    print(f"q3 fallback {FALLBACKS[0]}: state-once OOM, reference path "
          f"({torch.cuda.memory_allocated() / 2**30:.2f} GiB allocated)", file=sys.stderr, flush=True)
    keep, self.state_once = self.state_once, None
    try:
        try:
            return self.run(prepared)
        except torch.OutOfMemoryError:
            if self.batch_size == 1:
                raise
        # still too big: the same reference path one question at a time (a ToolRet request of 32 long candidates
        # needed 11 GiB more than a 96 GB card had left at batch 8)
        torch.cuda.empty_cache()
        print(f"q3 fallback {FALLBACKS[0]}: reference path OOM at batch {self.batch_size}, batch 1",
              file=sys.stderr, flush=True)
        width, self.batch_size = self.batch_size, 1
        try:
            return self.run(prepared)
        finally:
            self.batch_size = width
    finally:
        self.state_once = keep
        torch.cuda.empty_cache()


_RUN_STATE_ONCE = []


def install():
    import d3_runtime
    if not getattr(d3_runtime, "_q3", False):
        base = d3_runtime.backbone_class
        d3_runtime.backbone_class = lambda kind: _Q3Loader(base(kind))
        _RUN_STATE_ONCE.append(d3_runtime.D3._run_state_once)
        d3_runtime.D3._run_state_once = _run_state_once_or_plain
        d3_runtime._q3 = True


def make_resident(model):
    """The fast path for big cards: every packed Linear becomes a plain bf16 Linear holding its unpacked weight (bit for
    bit HQQ's dequantize) with the LoRA merged in, W + B A, so a forward pass costs what d3-ultra's does."""
    swaps = [(n, m) for n, m in model.named_modules() if isinstance(m, Q3Linear)]
    for name, m in swaps:
        w = m.unpacked()
        if m.lora is not None:
            a, b = m.lora
            w = (w.float() + b.float() @ a.float()).to(torch.bfloat16)
        lin = torch.nn.Linear(m.in_features, m.out_features, bias=m.bias is not None, dtype=torch.bfloat16,
                              device=m.lo.device)
        with torch.no_grad():
            lin.weight.copy_(w)
            if m.bias is not None:
                lin.bias.copy_(m.bias)
        lin.requires_grad_(False)
        parent, _, child = name.rpartition(".")
        setattr(model.get_submodule(parent) if parent else model, child, lin)
        del m, w
    _SCRATCH.clear()
    torch.cuda.empty_cache()
    return len(swaps)


from d3_engine import D3Engine  # noqa: E402


class D3Q3Engine(D3Engine):
    name = "wildcat-q3"

    def __init__(self, adapter=None, lora_scale=1.0, **options):
        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"))
        if so_div > 1:
            _small_card(so_div)
        if total < 40 * 2**30:
            # 32 retrieval questions with long options do not fit 8 at a time next to the weights on a 24 GB card
            options.setdefault("batch_size", 1)
        super().__init__(**options)
        self.provenance = dict(self.provenance)
        info = {"quant": f"hqq-{BITS}bit-g{GROUP}, packed, Triton unpack per matmul", "skip": SKIP,
                "card_gib": round(total / 2**30, 1), "state_once_div": so_div}
        lm = self.decision.backbone.language_model
        packaged = os.path.join(str(self.decision.root), "lora.safetensors")
        if adapter:
            paths = adapter.split(",")
            info.update(adapters=[os.path.abspath(p) for p in paths], lora_scale=float(lora_scale),
                        attached=attach(lm, stack_adapters(paths, float(lora_scale))))
        elif os.path.exists(packaged):
            info.update(adapters=[packaged], attached=attach(lm, load_lora(packaged)))
        # Q3_RESIDENT=1/0 forces it; by default cards with room for the bf16 weights (~54 GB) take the fast path
        if os.environ.get("Q3_RESIDENT", "1" if total > 72 * 2**30 else "0") == "1":
            info.update(quant=f"hqq-{BITS}bit-g{GROUP}, unpacked to bf16 at load, LoRA merged (fast path)",
                        resident=make_resident(self.decision.backbone))
            print(f"q3 resident: {torch.cuda.memory_allocated(device) / 2**30:.2f} GiB on {device}",
                  file=sys.stderr, flush=True)
        self.provenance["wildcat"] = info
        print(f"q3 engine {json.dumps(info)}", file=sys.stderr, flush=True)