#!/usr/bin/env python3
"""Convert merged internal-format model.pth back to HF fish_qwen3_omni layout.

Builds the exact inverse of _remap_fish_qwen3_omni_keys by replaying the
forward mapping over the ORIGINAL s2-pro index, then writes safetensors
shards with the original shard assignment + index.json. Fidelity by
construction: same keys, same shards, byte-identical layout semantics.
"""

import json
from collections import OrderedDict, defaultdict
from pathlib import Path

import torch
from safetensors.torch import save_file

SRC = Path("/opt/work/checkpoints/s2-pro")
DST = Path("/opt/work/checkpoints/s2pro-egy-merged")


def forward_map(hf_key: str) -> str:
    if hf_key.startswith("text_model.model."):
        return hf_key[len("text_model.model."):]
    if hf_key.startswith("audio_decoder."):
        suffix = hf_key[len("audio_decoder."):]
        return suffix if suffix.startswith("codebook_embeddings.") else "fast_" + suffix
    return hf_key


index = json.loads((SRC / "model.safetensors.index.json").read_text())
weight_map = index["weight_map"]  # hf_key -> shard filename

inv = {}
for hf_key in weight_map:
    ik = forward_map(hf_key)
    assert ik not in inv, f"collision: {ik}"
    inv[ik] = hf_key

print(f"index has {len(weight_map)} keys")

sd = torch.load(DST / "model.pth", map_location="cpu", mmap=True, weights_only=True)
print(f"merged model.pth has {len(sd)} keys")

missing_in_index = [k for k in sd if k not in inv]
missing_in_model = [k for k in inv if k not in sd]
print("merged keys not in index:", missing_in_index[:5], f"({len(missing_in_index)})")
print("index keys not in merged:", missing_in_model[:5], f"({len(missing_in_model)})")
assert not missing_in_index, "unexpected keys in merged model"
assert not missing_in_model, "merged model is missing weights"

shards = defaultdict(OrderedDict)
for ik, tensor in sd.items():
    hf_key = inv[ik]
    shards[weight_map[hf_key]][hf_key] = tensor.contiguous()

total = 0
for shard_name, tensors in sorted(shards.items()):
    total += sum(t.numel() * t.element_size() for t in tensors.values())
    save_file(tensors, str(DST / shard_name), metadata={"format": "pt"})
    print(f"wrote {shard_name} ({len(tensors)} tensors)")

out_index = {"metadata": {"total_size": total}, "weight_map": weight_map}
(DST / "model.safetensors.index.json").write_text(json.dumps(out_index, indent=2))
print(f"index written, total_size={total/1e9:.2f}GB")
print("CONVERT_DONE")