"""Run the deflated-ASP NVFP4 checkpoint. Real 4-bit execution on the FP4 tensor cores.

Per layer, exactly the contract the export wrote and the search scored:

    x~ = (x / s) H
    y  = (x~ V)(W~V)^T  +  NVFP4GEMM( (I - VV^T) x~ ,  W~(I-VV^T) )   + bias

The weight never leaves E2M1: it is packed nibbles with pre-swizzled E4M3 block scales, handed
straight to `torch._scaled_mm_v2` with recipe BlockWise1x16. Nothing is dequantised to bf16 and no
fp16 GEMM runs on the bulk path -- that is what makes this real quantisation rather than a
simulation of it.

The ACTIVATION is quantised per forward: block 16, E4M3 scales, one level (SVDQuant's nvfp4 recipe
for inputs), then swizzled. The swizzle is unavoidable -- `_scaled_mm_v2` rejects row-major scales
-- and it is done in torch here rather than fused, so this path is correct first and fast later.
The rank-r branch stays in bf16 by design: it carries the directions the action depends on, which
is the whole point of protecting them.

  FW_ASP_CKPT=/path/to/ur3_step7000_asp_nvfp4.pt
"""
from __future__ import annotations

import math
import os

import torch
import torch.nn as nn

from nvfp4 import nvfp4_mm, quantize_nvfp4, recip_scale, swizzle_scales


def _fwht(x):
    """Orthonormal fast Walsh-Hadamard transform along the last dim (size a power of 2)."""
    orig = x.shape; m = orig[-1]; x = x.reshape(-1, m).clone(); h = 1
    while h < m:
        x = x.view(-1, m // (2 * h), 2, h); a = x[:, :, 0, :]; b = x[:, :, 1, :]
        x = torch.stack([a + b, a - b], dim=2).reshape(-1, m); h *= 2
    return (x / math.sqrt(m)).reshape(orig)


_HMAT_CACHE = {}


def _hadamard_matrix(B, dtype, device):
    """H with x @ H == _fwht(x). Built by running _fwht on the identity, so the butterfly
    ORDERING and the 1/sqrt(B) normalisation match the ones the weights were rotated with at
    export. A Sylvester construction would give a different ordering, and the resulting error
    would look like a quantisation regression rather than a transform mismatch."""
    key = (B, dtype, str(device))
    if key not in _HMAT_CACHE:
        eye = torch.eye(B, device=device, dtype=torch.float32)
        _HMAT_CACHE[key] = _fwht(eye).to(dtype).contiguous()
    return _HMAT_CACHE[key]


_H: dict = {}


def _hmat(B, device, dtype):
    k = (B, str(device), dtype)
    if k not in _H:
        _H[k] = _hadamard_matrix(B, dtype, device)
    return _H[k]


class ASPNVFP4Linear(nn.Module):
    """One exported Linear: NVFP4 deflated weight + bf16 rank-r action subspace."""

    def __init__(self, e: dict, dtype, device):
        super().__init__()
        self.in_features = int(e["in_features"])
        self.out_features = int(e["out_features"])
        self.block = int(e["fwht_block"])
        self.a_block = int(e["a_block"])
        self.rank = int(e["asp_rank"])
        self.register_buffer("wq", e["wq"].to(device))

        self.register_buffer("wsc", e["wscale_swizzled"].to(device))
        # The kernel wants the RECIPROCAL of the per-tensor scale, and it is a constant: built
        # once here, not per forward. The transposed weight view is hoisted for the same reason.
        self.register_buffer("w_rglob", recip_scale(e["wglobal"].to(device), device))
        self.register_buffer("a_rglob", recip_scale(None, device).clone())
        # KEPT IN FP32, NOT THE MODEL DTYPE. These are stored fp16; casting them to bf16 throws
        # away three mantissa bits before a 4-bit grid ever sees them, and the grid's codes are
        # ~33% apart, so those bits decide which code an element lands on. Measured: bf16 buffers
        # put the layer output 2e-2 to 5e-2 from the fp32 reference; fp16 -> fp32 is exact and
        # costs a few hundred MB across the model for tensors that are r=32 columns wide.
        self.register_buffer("inv_smooth", e["inv_smooth"].to(device=device, dtype=torch.float32))
        self.register_buffer("bias", e["bias"].to(device=device, dtype=torch.float32)
                             if "bias" in e else None)
        # V as [in, r] for the projection, W~V as [out, r] for the epilogue
        self.register_buffer("V", e["lr_a"].t().contiguous().to(device=device,
                                                                dtype=torch.float32)
                             if self.rank else None)
        self.register_buffer("WV", e["lr_b"].to(device=device, dtype=torch.float32)
                             if self.rank else None)
        self._wqt = self.wq.t()          # transposed view, hoisted out of the forward

    def forward(self, x):
        # THE PROLOGUE RUNS IN FP32, deliberately. What follows it is a 4-bit grid whose codes are
        # ~33% apart, so a bf16 rounding of the smoothed/rotated activation moves elements across
        # code boundaries: measured on the exported layers, a bf16 prologue put the layer output
        # 3.8e-2 to 9.0e-2 from the fp32 one. The search that CHOSE this layer's (alpha, beta) and
        # its subspace scored an fp32 prologue, so a bf16 one would deploy a different arm than the
        # one that was selected. It is also cheap: the rotation is O(d*sqrt(d)) beside an FP4 GEMM.
        shp = x.shape
        x2 = x.reshape(-1, shp[-1]).float()
        xh = x2 * self.inv_smooth
        D = xh.shape[-1]
        xt = (xh.reshape(-1, D // self.block, self.block)
              @ _hmat(self.block, xh.device, torch.float32)).reshape(-1, D)
        if self.rank:
            Vf = self.V
            proj = xt @ Vf                              # [tok, r]
            xp = xt - proj @ Vf.t()                     # (I - VV^T) x~
        else:
            proj, xp = None, xt
        ak, asc, _ = quantize_nvfp4(xp, block=self.a_block, two_level=False)
        y = nvfp4_mm(ak, swizzle_scales(asc), self.wq, self.wsc, self.a_rglob, self.w_rglob,
                     out_dtype=torch.float32, b_packed_t=self._wqt)[: xp.shape[0]]
        # EVERYTHING IN FP32, THEN ONE CAST AT THE END. The epilogue and the bias are fp32
        # buffers, so casting before adding them promotes the result straight back to fp32 and the
        # next layer_norm fails on a dtype it did not expect. One cast, last.
        if self.rank:
            y = y + proj @ self.WV.t()
        if self.bias is not None:
            y = y + self.bias
        y = y.to(x.dtype)
        return y.reshape(*shp[:-1], self.out_features)

    def extra_repr(self):
        return (f"in={self.in_features}, out={self.out_features}, NVFP4 w{self.a_block}, "
                f"rank={self.rank}, H={self.block}")


def install_asp_nvfp4(model, ckpt_path=None, verbose=True) -> int:
    """Swap every exported Linear. A partial swap is not a defined arm, so it raises."""
    ckpt_path = ckpt_path or os.environ.get("FW_ASP_CKPT")
    blob = torch.load(str(ckpt_path), map_location="cpu", weights_only=False)
    meta, layers = blob["meta"], blob["layers"]
    if meta.get("lowrank_mode") != "asp_deflated":
        raise RuntimeError(f"{ckpt_path}: lowrank_mode={meta.get('lowrank_mode')!r}; this runtime "
                           f"implements the DEFLATED ASP contract and applying it to another "
                           f"would give a well-formed GEMM of the wrong bilinear form.")
    if verbose:
        f = meta["format"]
        print(f"[asp-nvfp4] {meta['scheme']}: {meta['n_layers']} Linears "
              f"({meta['asp_layers']} with ASP r{meta['rank']}), {meta['bpw']} BPW", flush=True)
        print(f"[asp-nvfp4]   {f['element']} w{f['w_block']}/a{f['a_block']}, {f['scale']} scales, "
              f"{f['weight_scale_levels']}-level weights, {f['scale_layout']}", flush=True)
    named = dict(model.named_modules())
    done, missing = 0, []
    for name, e in layers.items():
        mod = named.get(name)
        if not isinstance(mod, nn.Linear):
            missing.append(name)
            continue
        new = ASPNVFP4Linear(e, mod.weight.dtype, mod.weight.device)
        parent = model.get_submodule(name.rsplit(".", 1)[0])
        setattr(parent, name.rsplit(".", 1)[-1], new)
        named[name] = None
        del mod, new
        done += 1
    if missing:
        raise RuntimeError(f"{ckpt_path}: {len(missing)} layers did not resolve, e.g. {missing[:3]}")
    del blob, layers, named
    import gc
    gc.collect()
    torch.cuda.empty_cache()
    if verbose:
        free, tot = torch.cuda.mem_get_info()
        print(f"[asp-nvfp4] replaced {done} Linears | GPU {(tot-free)/2**30:.1f}/"
              f"{tot/2**30:.0f} GiB", flush=True)
    return done


def load_quantized_asp_model(ckpt_path, build_model, verbose=True):
    """Build the model straight from the self-contained NVFP4 checkpoint.

    `build_model` is your own bf16 FastWAM constructor and must return `(model, cfg)`; the
    checkpoint carries every tensor the quantised model needs, so it is only used for the
    module graph and the non-Linear submodules. Nothing is read from the bf16 weights.
    """
    blob = torch.load(str(ckpt_path), map_location="cpu", weights_only=False)
    if "mot_rest" not in blob:
        raise RuntimeError(f"{ckpt_path} is not self-contained")
    model, cfg = build_model()
    missing, unexpected = model.mot.load_state_dict(blob["mot_rest"], strict=False)
    if unexpected:
        raise RuntimeError(f"{len(unexpected)} unexpected mot tensors, e.g. {list(unexpected)[:3]}")
    if model.proprio_encoder is not None:
        model.proprio_encoder.load_state_dict(blob["proprio_encoder"], strict=True)
    del blob
    install_asp_nvfp4(model, ckpt_path, verbose=verbose)
    return model, cfg