"""Phase D: splice trained mtp.* tensors back into the merged BF16 checkpoint.

Creates a new checkpoint dir where safetensors shards containing replaced
tensors are rewritten and every other file is HARDLINKED (no 230GB copy).
Then quantize with quant_oq5e_mtp.py (preserve_mtp) as usual.

Usage:
  python3 splice_mtp.py [--trained mtp_trained.safetensors]
                        [--src .../NVIDIA-Nemotron-3-Super-120B-merged3]
                        [--dst .../NVIDIA-Nemotron-3-Super-120B-merged3-mtpft]
"""

import argparse
import json
import os
from pathlib import Path

import mlx.core as mx

ROOT = Path(__file__).parent


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--trained", default=str(ROOT / "mtp_trained.safetensors"))
    ap.add_argument(
        "--src", default="/Users/david/AI/NVIDIA-Nemotron-3-Super-120B-merged3")
    ap.add_argument(
        "--dst",
        default="/Users/david/AI/NVIDIA-Nemotron-3-Super-120B-merged3-mtpft")
    args = ap.parse_args()
    src, dst = Path(args.src), Path(args.dst)

    trained = mx.load(args.trained)
    print(f"{len(trained)} trained tensors")
    index = json.load(open(src / "model.safetensors.index.json"))["weight_map"]
    missing = [k for k in trained if k not in index]
    if missing:
        raise SystemExit(f"trained tensors not in checkpoint index: {missing[:5]}")
    dirty_files = {index[k] for k in trained}
    print(f"rewriting {len(dirty_files)} shard file(s): {sorted(dirty_files)}")

    dst.mkdir(exist_ok=True)
    for f in sorted(src.iterdir()):
        if f.is_dir() or f.name.startswith("."):
            continue
        target = dst / f.name
        if target.exists():
            target.unlink()
        if f.name in dirty_files:
            shard = dict(mx.load(str(f)))
            replaced = 0
            for k in list(shard.keys()):
                if k in trained:
                    if shard[k].shape != trained[k].shape:
                        raise SystemExit(
                            f"shape mismatch {k}: {shard[k].shape} vs "
                            f"{trained[k].shape}")
                    shard[k] = trained[k].astype(shard[k].dtype)
                    replaced += 1
            mx.save_safetensors(str(target), shard,
                                metadata={"format": "pt"})
            print(f"  {f.name}: replaced {replaced} tensors")
        else:
            os.link(f, target)
    print(f"DONE -> {dst}")


if __name__ == "__main__":
    main()