import os
import json
from safetensors.torch import load_file, save_file

def fix_sharded_mtp(base_path, merged_path):
    # Load indices
    with open(os.path.join(base_path, "model.safetensors.index.json"), "r") as f:
        base_idx = json.load(f)
    with open(os.path.join(merged_path, "model.safetensors.index.json"), "r") as f:
        merged_idx = json.load(f)

    # Find missing MTP keys
    mtp_keys = {k: v for k, v in base_idx["weight_map"].items() if "mtp" in k.lower()}
    missing_keys = [k for k in mtp_keys if k not in merged_idx["weight_map"]]

    if not missing_keys:
        return print("No missing MTP keys found.")

    # Group by source shard to minimize loads
    shard_to_keys = {}
    for k in missing_keys:
        shard = mtp_keys[k]
        shard_to_keys.setdefault(shard, []).append(k)

    # Collect tensors
    restored_tensors = {}
    for shard, keys in shard_to_keys.items():
        weights = load_file(os.path.join(base_path, shard))
        for k in keys:
            restored_tensors[k] = weights[k]

    # Save to a new dedicated MTP shard
    new_shard_name = "model-mtp-restored.safetensors"
    save_file(restored_tensors, os.path.join(merged_path, new_shard_name))

    # Update merged index
    for k in missing_keys:
        merged_idx["weight_map"][k] = new_shard_name
    
    with open(os.path.join(merged_path, "model.safetensors.index.json"), "w") as f:
        json.dump(merged_idx, f, indent=2)
    
    print(f"Fixed! Added {len(missing_keys)} keys to {new_shard_name}")

#
# Usage: fix_sharded_mtp("path/to/base", "path/to/merged")

fix_sharded_mtp("MODEL_W_MTP_KEYS", "MODEL_TO_ADD_MTP_KEYS")