"""Full-model B200 reference KLD; explicitly distinct from native SM120 true-decode KLD."""
from __future__ import annotations

import argparse
import dataclasses
import hashlib
import json
import os
import re
import sys
import time
import types
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from accelerate import dispatch_model, init_empty_weights
from accelerate.utils import set_module_tensor_to_device
from safetensors import safe_open
from transformers import AutoConfig, Glm5NextForConditionalGeneration

import encode_layer_v2 as encoder
from glm53_nvfp4 import p8_coupled_scale as coupled
from glm53_nvfp4.p8_layer_rate_damage import RankLayerStreams, decode_expert, state_codebook
from kld_core import score_window


def sha(path):
    h = hashlib.sha256()
    with Path(path).open("rb") as f:
        for part in iter(lambda: f.read(64 << 20), b""):
            h.update(part)
    return h.hexdigest()


def qdq_rp2(values, global_scale, policy):
    if global_scale != 1.0 or policy != "amax":
        raise ValueError("Only the fixed RP2 amax carrier is supported")
    blocks = values.float().reshape(*values.shape[:-1], values.shape[-1] // 32, 32)

    def term(x):
        maximum = x.abs().amax(-1, keepdim=True)
        exponent = torch.ceil(torch.log2((maximum / 448).clamp_min(2.0 ** -127))).clamp(-127, 128)
        scale = torch.exp2(exponent)
        return (x / scale).clamp(-448, 448).to(torch.float8_e4m3fn).float() * scale

    hi = term(blocks)
    return (hi + term(blocks - hi)).reshape_as(values)


def experts_forward(self, hidden_states, top_k_index, top_k_weights):
    final = torch.zeros_like(hidden_states, dtype=torch.float32)
    # Every expert shares this row-local input transform. Evaluate it once per layer;
    # selecting rows afterward preserves the reference's casts and arithmetic.
    fc1_input = coupled.coupled_input_carrier(
        hidden_states, self._cv2_scales.gate_up_suh, quantize=True)
    # These are exact decoded physical weights; scales and rotations stay in the forward path.
    for expert in torch.unique(top_k_index).tolist():
        token_idx, route_idx = torch.where(top_k_index == expert)
        gate, up = self.gate_up_proj[expert].chunk(2, dim=0)
        down_input = coupled.coupled_middle_carrier(
            fc1_input[token_idx], gate, up, self._cv2_scales,
            expert=expert, intermediate_draw=coupled.COUPLED_SIGN_DRAW, quantize=True)
        down = F.linear(down_input.float(), self.down_proj[expert].float()).to(torch.float16).float()
        route = coupled.block_hadamard(down, block_size=coupled.HADAMARD_INPUT)
        route = route * self._cv2_scales.down_svh.float()
        result = coupled.block_hadamard(route, block_size=coupled.HADAMARD_RESIDUAL)
        final.index_add_(0, token_idx, result * top_k_weights[token_idx, route_idx, None].float())
    return final.to(hidden_states.dtype)


def tr3_experts_forward(self, hidden_states, top_k_index, top_k_weights):
    final = torch.zeros_like(hidden_states, dtype=torch.float32)
    for expert in torch.unique(top_k_index).tolist():
        token_idx, route_idx = torch.where(top_k_index == expert)
        gate, up = self.gate_up_proj[expert].chunk(2, dim=0)
        result = coupled.source_expert_reference(hidden_states[token_idx], gate, up, self.down_proj[expert])
        final.index_add_(0, token_idx, result * top_k_weights[token_idx, route_idx, None].float())
    return final.to(hidden_states.dtype)


def layer_device(layer):
    # Balance the 42 routed layers 11/11/10/10, leaving decoding/forward headroom.
    return 0 if layer < 14 else 1 if layer < 25 else 2 if layer < 35 else 3


@torch.inference_mode()
def load_model(source, checkpoint, mode):
    config = AutoConfig.from_pretrained(source, local_files_only=True)
    config._attn_implementation = "sdpa"
    config.text_config._attn_implementation = "sdpa"
    old_dtype = torch.get_default_dtype()
    torch.set_default_dtype(torch.bfloat16)
    try:
        with init_empty_weights():
            model = Glm5NextForConditionalGeneration(config)
    finally:
        torch.set_default_dtype(old_dtype)
    wm = json.loads((source / "model.safetensors.index.json").read_text())["weight_map"]
    # A root entry makes Accelerate move every descendant to GPU0 before the
    # per-layer hooks attach. Map complete sibling blocks explicitly instead.
    device_map = {"model.language_model.embed_tokens": 0,
                  "model.language_model.norm": 3,
                  "model.language_model.hc_head": 3, "lm_head": 3}
    for name, _ in model.model.named_children():
        if name != "language_model":
            device_map[f"model.{name}"] = 0
    for layer in range(45):
        device_map[f"model.language_model.layers.{layer}"] = layer_device(layer)

    def target_device(name):
        match = re.search(r"language_model\.layers\.(\d+)\.", name)
        if match:
            return layer_device(int(match.group(1)))
        return 3 if name.startswith(("lm_head.", "model.language_model.norm.", "model.language_model.hc_head.")) else 0

    def source_names(name):
        # Exact GLM5Next renames/concatenation from Transformers 5.17 conversion_mapping.py.
        if ".self_attn.conv1d.weight" in name:
            return [name.replace("self_attn.conv1d.weight", f"self_attn.{axis}_conv1d.weight")
                    for axis in ("q", "k", "v")]
        name = name.replace("self_attn.forget_gate.", "self_attn.")
        for hop in ("attn", "ffn"):
            for part in ("fn", "base", "scale"):
                name = name.replace(f"{hop}_hc.{part}", f"hc_{hop}_{part}")
        return [name]

    # Only non-routed tensors are read from BF16. Routed storage is filled directly from the checkpoint.
    missing = []
    for name, tensor in list(model.named_parameters()) + list(model.named_buffers()):
        if ".mlp.experts." in name:
            continue
        names = source_names(name)
        if any(key not in wm for key in names):
            if tensor.is_meta:
                missing.append(name)
            continue
        values = []
        for key in names:
            with safe_open(str(source / wm[key]), framework="pt", device="cpu") as f:
                values.append(f.get_tensor(key))
        value = torch.cat(values, dim=0) if len(values) > 1 else values[0]
        if tuple(value.shape) != tuple(tensor.shape):
            raise ValueError(f"Non-expert shape mismatch: {name}: {value.shape} vs {tensor.shape}")
        set_module_tensor_to_device(model, name, target_device(name), value=value, dtype=value.dtype)
    if missing:
        raise ValueError(f"Missing source parameters: {missing}")

    tr3_map = None
    if mode == "tr3":
        tr3_map = json.loads((checkpoint / "model.safetensors.index.json").read_text())["weight_map"]
        sys.path.insert(0, str(Path(os.environ["GLM53_KQUANT_ROOT"])))
        from kquant.exl3_reference import decode_exl3_weight
        from glm53_nvfp4 import trellis_nvfp4 as tn

    @torch.inference_mode()
    def load_layers(gpu):
        torch.cuda.set_device(gpu)
        device = torch.device(f"cuda:{gpu}")
        for layer in range(3, 45):
            if layer_device(layer) != gpu:
                continue
            start = time.time()
            prefix = f"model.language_model.layers.{layer}.mlp.experts"
            module = model.get_submodule(prefix)
            for name in ("gate_up_proj", "down_proj"):
                shape = getattr(module, name).shape
                value = torch.empty(shape, dtype=torch.float16 if mode == "tr3" else torch.bfloat16, device=device)
                set_module_tensor_to_device(model, f"{prefix}.{name}", device, value=value, dtype=value.dtype)
            if mode == "tr3":
                mcg = tn.procedural_state_values("mcg").float().to(device)
                def get(name):
                    with safe_open(str(checkpoint / tr3_map[name]), framework="pt", device="cpu") as f:
                        return f.get_tensor(name).to(device)
                for expert in range(288):
                    for projection in ("gate", "up", "down"):
                        base = f"model.language_model.layers.{layer}.mlp.experts.{expert}.{projection}_proj"
                        trellis = get(base + ".trellis")
                        bits = trellis.shape[-1] // 16
                        assert bits == 4
                        states = tn.reconstruct_trellis_states(tn.unpack_trellis_edges(trellis, bits), bits)
                        weight = decode_exl3_weight(states, get(base + ".suh").float(), get(base + ".svh").float(),
                                                    codebook_values=mcg, bits=bits).t()
                        target = (module.down_proj[expert] if projection == "down" else
                                  module.gate_up_proj[expert, :2048] if projection == "gate" else
                                  module.gate_up_proj[expert, 2048:])
                        target.copy_(weight)
                module.forward = types.MethodType(tr3_experts_forward, module)
                torch.cuda.synchronize(device)
                print(json.dumps({"decoded_tr3_layer": layer, "gpu": gpu, "seconds": round(time.time() - start, 2)}), flush=True)
                continue
            streams = RankLayerStreams([checkpoint / "sidecars" / f"p8-layer-{layer:03d}-tp4-rank-{r}.safetensors"
                                        for r in range(4)])
            lut = state_codebook(streams.bits, device)
            for expert in range(288):
                gate, up, down = decode_expert(streams.expert(expert), streams.bits, lut, device)
                module.gate_up_proj[expert, :2048].copy_(gate)
                module.gate_up_proj[expert, 2048:].copy_(up)
                module.down_proj[expert].copy_(down)
            scales = encoder.r27_scales(layer)
            module._cv2_scales = dataclasses.replace(scales, **{
                field.name: getattr(scales, field.name).to(device)
                for field in dataclasses.fields(scales)
                if isinstance(getattr(scales, field.name), torch.Tensor)})
            module.forward = types.MethodType(experts_forward, module)
            torch.cuda.synchronize(device)
            print(json.dumps({"decoded_layer": layer, "gpu": gpu, "seconds": round(time.time() - start, 2)}), flush=True)
    with ThreadPoolExecutor(4) as pool:
        list(pool.map(load_layers, range(4)))
    if any(p.is_meta for p in model.parameters()):
        raise RuntimeError("Unloaded model parameter")
    if mode == "p8":
        coupled._qdq_e4m3_k32 = qdq_rp2
    model = dispatch_model(model, device_map=device_map, force_hooks=True)
    model.eval()
    return model


@torch.inference_mode()
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--checkpoint", type=Path, required=True)
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--mode", choices=("p8", "tr3"), default="p8")
    args = ap.parse_args()
    source = Path(os.environ["CV2_BF16"])
    root = Path(os.environ["CV2_ROOT"])
    teacher = Path(os.environ["CV2_DATA"]) / "kld-teacher"
    manifest = json.loads((root / "ops/verified-inputs.json").read_text())
    args.out.mkdir(parents=True, exist_ok=True)
    torch.set_num_threads(16)
    torch.backends.cuda.matmul.allow_tf32 = False
    model = load_model(source, args.checkpoint, args.mode)
    rows = []
    for window in manifest["windows"]:
        wid = window["id"]
        summary_path = args.out / f"{wid}.json"
        if summary_path.exists():
            rows.append(json.loads(summary_path.read_text()))
            continue
        started = time.time()
        token_path = teacher / f"calibration/panel-v1/arrays/{wid}.tokens.npy"
        teacher_path = teacher / window["teacher_path"]
        assert sha(token_path) == window["input_sha256"], wid
        assert sha(teacher_path) == window["teacher_sha256"], wid
        tokens = np.load(token_path, allow_pickle=False).reshape(-1).astype(np.int64)
        assert tokens.shape == (2048,)
        assert hashlib.sha256(tokens.astype("<i8").tobytes()).hexdigest() == window["token_values_sha256"]
        inputs = torch.from_numpy(tokens[None]).to("cuda:0")
        logits = model(input_ids=inputs, use_cache=False).logits[0, :2047, :154880].float().cpu().numpy()
        with safe_open(str(teacher_path), framework="pt", device="cpu") as f:
            tensors = [f.get_tensor(key) for key in f.keys() if f.get_slice(key).get_shape() == [2047, 154880]]
        assert len(tensors) == 1, wid
        scores = score_window(tensors[0].numpy(), logits, tokens, 154880, chunk_rows=128)
        np.savez(args.out / f"{wid}.npz", kld=scores.kld)
        row = {"window": wid, "domain": window["domain"], "mean_kld_rows1_2046": float(scores.kld[1:].mean()),
               "teacher_sha256": window["teacher_sha256"], "token_values_sha256": window["token_values_sha256"],
               "seconds": round(time.time() - started, 2)}
        summary_path.write_text(json.dumps(row, indent=2) + "\n")
        rows.append(row)
        print(json.dumps(row), flush=True)
    report = {"checkpoint": str(args.checkpoint), "windows": len(rows),
              "mean_kld": float(np.mean([r["mean_kld_rows1_2046"] for r in rows])),
              "backend": ("Transformers 5.17 GLM5Next BF16 backbone; " +
                          ("decoded coupled P8 with two-term E4M3 reference carriers" if args.mode == "p8" else
                           "TR3 FP16 MCG decoded with H128 and suh/svh, FP16 weight storage, FP32 reference expert math")),
              "mode": args.mode,
              "execution": "B200 full-window causal forward, no KV quantization; rows 1..2046 scored",
              "native_sm120_true_decode_equivalence": False,
              "historical_tr3_fp8_true_decode_kld": 0.030099944949203105,
              "historical_r27_rp2_fp8_true_decode_kld": 0.03260408965478527,
              "comparison_limit": "Historical scores use a different backbone, KV and execution path; this result alone does not establish a production TR3 win or runtime speed.",
              "checkpoint_manifest_sha256": sha(args.checkpoint / ("trellismx-manifest.json" if args.mode == "p8"
                                                                    else "model.safetensors.index.json"))}
    (args.out / "REPORT.json").write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(report), flush=True)


if __name__ == "__main__":
    main()