#!/usr/bin/env python3
"""Merge the AngelSpec MTP head into a base HYV3 model.

The base model (e.g. ``Hy3_opensource``) already has the full transformer body
(layers 0-79 + embeddings + lm_head + final norm) *and* its own MTP head at
``model.layers.80.*``. This script produces a NEW model directory that is a copy
of the base, but whose MTP head (layer 80) is replaced by the AngelSpec MTP
weights exported in ``mtp.safetensors``.

Result = base body  +  AngelSpec MTP head.

Only the 5 shards that contain layer-80 tensors are rewritten. The other 94
shards, together with the config / tokenizer / index files, are linked (or
copied) verbatim from the base model, so the merge is fast and cheap.

Usage:
    python merge_mtp.py \
        --base ORIGINAL_MODEL_PATH \
        --mtp  MTP_WEIGHT_PATH \
        --out  OUTPUT_WEIGHT_PATH \
        --link copy        # symlink | hardlink | copy  (how to bring over the 94 unchanged shards)
"""
import argparse
import json
import os
import shutil

from safetensors import safe_open
from safetensors.torch import save_file

MTP_PREFIX = "model.layers.80."


def bring_over(src, dst, mode):
    """Materialise an unchanged shard in the output dir according to *mode*."""
    if os.path.lexists(dst):
        os.remove(dst)
    if mode == "symlink":
        os.symlink(os.path.abspath(src), dst)
    elif mode == "hardlink":
        os.link(src, dst)
    elif mode == "copy":
        shutil.copy2(src, dst)
    else:
        raise ValueError(f"unknown link mode: {mode}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", required=True, help="Base full model directory (body comes from here)")
    ap.add_argument("--mtp", required=True, help="MTP export dir (contains mtp.safetensors)")
    ap.add_argument("--out", required=True, help="Output directory for the merged model")
    ap.add_argument("--link", default="symlink", choices=["symlink", "hardlink", "copy"],
                    help="How to bring over the 94 unchanged shards (default: symlink)")
    ap.add_argument("--prefix", default=MTP_PREFIX, help="Tensor prefix of the MTP layer")
    args = ap.parse_args()

    os.makedirs(args.out, exist_ok=True)

    # --- 1. Load the base index and locate the MTP (layer-80) tensors ----------
    base_index_path = os.path.join(args.base, "model.safetensors.index.json")
    base_index = json.load(open(base_index_path))
    weight_map = base_index["weight_map"]

    mtp_keys = sorted(k for k in weight_map if k.startswith(args.prefix))
    if not mtp_keys:
        raise SystemExit(f"No {args.prefix!r} tensors found in {base_index_path}")

    mtp_shards = sorted({weight_map[k] for k in mtp_keys})
    all_shards = sorted(set(weight_map.values()))
    print(f"MTP layer has {len(mtp_keys)} tensors across {len(mtp_shards)} shards.")
    print(f"Base model has {len(all_shards)} shards total; "
          f"{len(all_shards) - len(mtp_shards)} will be linked verbatim.")

    # --- 2. Load replacement MTP tensors from the export -----------------------
    mtp_file = os.path.join(args.mtp, "mtp.safetensors")
    new_mtp = {}
    with safe_open(mtp_file, framework="pt", device="cpu") as f:
        exported = set(f.keys())
        missing = set(mtp_keys) - exported
        if missing:
            raise SystemExit(f"{len(missing)} MTP tensors missing from {mtp_file}, "
                             f"e.g. {sorted(missing)[:3]}")
        for k in mtp_keys:
            new_mtp[k] = f.get_tensor(k)

    # Sanity check: shapes/dtypes must match what the base expects.
    for shard in mtp_shards:
        with safe_open(os.path.join(args.base, shard), framework="pt", device="cpu") as f:
            for k in [k for k in mtp_keys if weight_map[k] == shard]:
                bt = f.get_slice(k)
                bshape = tuple(bt.get_shape())
                if bshape != tuple(new_mtp[k].shape):
                    raise SystemExit(f"Shape mismatch for {k}: base {bshape} vs mtp {tuple(new_mtp[k].shape)}")

    # --- 3. Bring over the unchanged shards ------------------------------------
    for shard in all_shards:
        if shard in mtp_shards:
            continue
        bring_over(os.path.join(args.base, shard), os.path.join(args.out, shard), args.link)
    print(f"Linked {len(all_shards) - len(mtp_shards)} unchanged shards ({args.link}).")

    # --- 4. Rewrite the MTP-containing shards with the new head ----------------
    for shard in mtp_shards:
        src = os.path.join(args.base, shard)
        tensors = {}
        with safe_open(src, framework="pt", device="cpu") as f:
            for k in f.keys():
                tensors[k] = f.get_tensor(k)
        replaced = 0
        for k in tensors:
            if k in new_mtp:
                tensors[k] = new_mtp[k]
                replaced += 1
        save_file(tensors, os.path.join(args.out, shard), metadata={"format": "pt"})
        print(f"  rewrote {shard}: replaced {replaced} MTP tensors "
              f"(kept {len(tensors) - replaced} base tensors)")

    # --- 5. Copy every non-shard file (config, tokenizer, index, ...) ----------
    for name in sorted(os.listdir(args.base)):
        src = os.path.join(args.base, name)
        if not os.path.isfile(src):
            continue
        if name.endswith(".safetensors"):   # shards handled above
            continue
        shutil.copy2(src, os.path.join(args.out, name))
    print("Copied config / tokenizer / index and other metadata files.")

    print(f"\nDone. Merged model written to: {args.out}")
    if args.link == "symlink":
        print("NOTE: unchanged shards are symlinks into the base model. Keep the base "
              "model in place, or re-run with --link copy to make a self-contained copy.")


if __name__ == "__main__":
    main()