"""Prepare the existing BF16 MTP captures for the owner's new Trellis MTP arm."""
import hashlib
import json
import os
from pathlib import Path

import numpy as np
import torch
from huggingface_hub import snapshot_download
from safetensors.torch import save_file

root = Path(os.environ["CV2_ROOT"])
out = Path(os.environ["CV2_OUT"])
dest = out / "mtp-inputs"
dest.mkdir(exist_ok=True)
repo = "brandonmusic/GLM-5.3-Flash-BF16-Teacher-Logits"
rev = "95f4fdd94bf29989db2e0d1054e4931f55edb6aa"
snapshot_download(repo, repo_type="dataset", revision=rev, local_dir=dest / "source",
                  allow_patterns=["calibration/mtp45-ep4-full/*"], max_workers=4)
source = dest / "source/calibration/mtp45-ep4-full"
manifest = json.loads((source / "capture-manifest.json").read_text())
receipt = json.loads((source / "capture-receipt.json").read_text())
assert receipt["complete"] is True
assert hashlib.sha256((source / "capture-manifest.json").read_bytes()).hexdigest() == receipt["capture_manifest_file_sha256"]
assert manifest["layer"] == 45 and manifest["geometry"]["experts"] == 288
assert manifest["model_revision"] == "a6c167b62691b2bac901344b65cb651a70f53e43"
for info in manifest["files"].values():
    path = source / info["path"]
    assert path.stat().st_size == info["bytes"]
    with path.open("rb") as stream:
        assert hashlib.file_digest(stream, "sha256").hexdigest() == info["sha256"]
roles = json.loads((root / "data/roles-v5.json").read_text())["roles"]["fit"]
assert len(roles) == 64
role_map = {r["id"]: r for r in roles}
windows = [w for w in manifest["windows"] if w["window_id"] in role_map]
assert len(windows) == 64
offset = 0
offsets = {}
for window in manifest["windows"]:
    offsets[window["window_id"]] = offset
    offset += window["rows"]
assert offset == manifest["rows"]
for window in windows:
    assert window["role"] == "fit" and window["rows"] == 2047
    assert window["token_ids_sha256"] == role_map[window["window_id"]]["input_sha256"]
streams = [("hidden_bf16", "<u2", 4096, "hidden.bf16.bin"),
           ("topk_ids_u16le", "<u2", 8, "topk_ids.u16le.bin"),
           ("topk_weights_f32le", "<f4", 8, "topk_weights.f32le.bin")]
for key, dtype, width, name in streams:
    array = np.memmap(source / manifest["files"][key]["path"], mode="r", dtype=dtype,
                      shape=(manifest["rows"], width))
    with (dest / name).open("wb") as stream:
        for window in windows:
            start = offsets[window["window_id"]]
            stream.write(np.asarray(array[start:start + window["rows"]]).tobytes())
# No prior Trellis/EXL3 MTP scale vectors exist in the supplied r27 overlay.
# New signed-unit vectors provide an orthogonal coupled H512/H128 transform;
# the unchanged codec-v2 block-scale refit handles weight magnitudes.
generator = torch.Generator(device="cpu").manual_seed(530045)
def signs(shape):
    return torch.randint(0, 2, shape, generator=generator).mul_(2).sub_(1).to(torch.float16).contiguous()
scales = {"gate_up_suh": signs((4096,)), "down_svh": signs((4096,)),
          "gate_svh": signs((288, 2048)), "up_svh": signs((288, 2048)),
          "down_suh": signs((288, 2048))}
save_file(scales, str(dest / "coupled-scales.safetensors"))
info = {"schema": "trellismx-codecv2-mtp-inputs.v1", "layer": 45,
        "capture_repo": repo, "capture_revision": rev, "source_capture_sha256": manifest["capture_sha256"],
        "windows": windows, "rows_per_window": 2047, "rows": 64 * 2047,
        "scale_policy": "signed-unit coupled vectors; CPU torch seed 530045; no prior r27 MTP scales",
        "scale_sha256": hashlib.sha256((dest / "coupled-scales.safetensors").read_bytes()).hexdigest()}
(dest / "inputs.json").write_text(json.dumps(info, indent=2) + "\n")
(out / "status/mtp-inputs.ready").write_text(info["scale_sha256"] + "\n")
print(json.dumps({k: v for k, v in info.items() if k != "windows"}), flush=True)