"""Encode MTP45 with the supplied codec-v2 math and its 2047-row fit captures."""
import argparse
import hashlib
import json
import os
import time
from pathlib import Path

import numpy as np
import torch
from safetensors.torch import load_file, save_file

import encode_layer_v2 as cv2
from glm53_nvfp4.p8_coupled_scale import CoupledScaleSet, COUPLED_SIGN_DRAW, encode_coupled_scale_weights, coupled_input_carrier, coupled_middle_carrier
from glm53_nvfp4.shard_index import IndexedCheckpoint

ap = argparse.ArgumentParser()
ap.add_argument("--expert-start", type=int, required=True)
ap.add_argument("--expert-end", type=int, required=True)
ap.add_argument("--out", type=Path, required=True)
args = ap.parse_args()
assert 0 <= args.expert_start < args.expert_end <= 288
assert not args.out.exists()
inputs = Path(os.environ["CV2_OUT"]) / "mtp-inputs"
info = json.loads((inputs / "inputs.json").read_text())
assert info["rows"] == 64 * 2047
hidden = np.memmap(inputs / "hidden.bf16.bin", mode="r", dtype="<u2", shape=(info["rows"], 4096))
ids = np.memmap(inputs / "topk_ids.u16le.bin", mode="r", dtype="<u2", shape=(info["rows"], 8))
routes = np.memmap(inputs / "topk_weights.f32le.bin", mode="r", dtype="<f4", shape=(info["rows"], 8))
scales = load_file(str(inputs / "coupled-scales.safetensors"))
scales_all = CoupledScaleSet(**scales, source_path=inputs, source_sha256=info["scale_sha256"], source_metadata={}, tensor_hashes={})
scales_all.validate()
device = torch.device("cuda:0")
torch.cuda.set_device(device)
every = np.asarray([w * 2047 + t for w in range(64) for t in range(0, 2047, 4)])
hid_every = torch.from_numpy(np.array(hidden[every], copy=True)).view(torch.bfloat16)
ones_every = torch.ones(len(every))
source_path = Path(os.environ["CV2_BF16"])
source = IndexedCheckpoint(source_path, source_path / "model.safetensors.index.json")
payload, metrics, h_all_in = {}, [], None
started = time.time()
for expert in range(args.expert_start, args.expert_end):
    t0 = time.time()
    base = source.expert_prefix(45, expert)
    orig = {p: source.get(f"{base}{p}.weight").to(device).float() for p in cv2.PROJ}
    sc = CoupledScaleSet(gate_up_suh=scales_all.gate_up_suh.to(device),
        down_svh=scales_all.down_svh.to(device),
        gate_svh=scales_all.gate_svh[expert:expert + 1].to(device),
        up_svh=scales_all.up_svh[expert:expert + 1].to(device),
        down_suh=scales_all.down_suh[expert:expert + 1].to(device),
        source_path=inputs, source_sha256=info["scale_sha256"], source_metadata={}, tensor_hashes={})
    tw = encode_coupled_scale_weights(orig["gate_proj"], orig["up_proj"], orig["down_proj"], sc,
                                       expert=0, intermediate_draw=COUPLED_SIGN_DRAW)
    carrier = lambda h: coupled_input_carrier(h, sc.gate_up_suh, quantize=True)
    row_ids, slots = np.nonzero(ids == expert)
    assert row_ids.size, f"No MTP fit routes for expert {expert}"
    hid_r = torch.from_numpy(np.array(hidden[row_ids], copy=True)).view(torch.bfloat16)
    rt_r = torch.from_numpy(np.array(routes[row_ids, slots], copy=True))
    if h_all_in is None:
        h_all_in = cv2.hessian_chunks(hid_every, ones_every, carrier, device)
    h_in = cv2.mix(cv2.hessian_chunks(hid_r, rt_r, carrier, device), h_all_in, 0.5)
    g_rec, g_tr, g_sc = cv2.encode_v2(tw[0], h_in, 4, 0.3)
    u_rec, u_tr, u_sc = cv2.encode_v2(tw[1], h_in, 4, 0.3)
    mid = lambda h: coupled_middle_carrier(carrier(h), g_rec, u_rec, sc, expert=0,
                                           intermediate_draw=COUPLED_SIGN_DRAW, quantize=True)
    h_dn = cv2.mix(cv2.hessian_chunks(hid_r, rt_r, mid, device),
                   cv2.hessian_chunks(hid_every, ones_every, mid, device), 0.5)
    d_rec, d_tr, d_sc = cv2.encode_v2(tw[2], h_dn, 4, 0.3)
    row = {"expert": expert, "routed_tokens": len(row_ids), "seconds": round(time.time() - t0, 2)}
    for name, rec, tr, codes, target in (("gate", g_rec, g_tr, g_sc, tw[0]),
            ("up", u_rec, u_tr, u_sc, tw[1]), ("down", d_rec, d_tr, d_sc, tw[2])):
        payload[f"{expert}.{name}.trellis"] = tr
        payload[f"{expert}.{name}.scale_ue8m0"] = codes
        row[f"{name}_tw_nmse"] = float((rec - target).double().square().sum() / target.double().square().sum())
    assert all(np.isfinite(v) for v in row.values())
    metrics.append(row)
    print(json.dumps(row), flush=True)
    del orig, tw, h_in, h_dn, g_rec, u_rec, d_rec, hid_r, rt_r
    torch.cuda.empty_cache()
meta = {"schema": "trellismx-codec-v2-mtp-chunk.v1", "layer": "45", "bits": "4",
        "expert_range": f"{args.expert_start}:{args.expert_end}", "beta": "0.5", "percdamp": "0.3",
        "encoder_sha256": hashlib.sha256(Path(cv2.__file__).read_bytes()).hexdigest(),
        "mtp_wrapper_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
        "mtp_inputs_sha256": hashlib.sha256((inputs / "inputs.json").read_bytes()).hexdigest(),
        "scale_policy": info["scale_policy"]}
args.out.parent.mkdir(parents=True, exist_ok=True)
tmp = args.out.with_suffix(".partial")
save_file(payload, str(tmp), metadata=meta)
tmp.rename(args.out)
args.out.with_suffix(".metrics.json").write_text(json.dumps({"meta": meta, "rows": metrics, "seconds": time.time() - started}, indent=2) + "\n")
print(json.dumps({"done": str(args.out), "seconds": time.time() - started}), flush=True)