"""How weights/*.fp16.pt were made from the training checkpoints (kept for provenance; not needed at runtime).
v0.8-en (run on the training machine in the tts-sv container, paths relative to /w):
    python tools/export_weights.py --out weights \
        --hugin hugin-v0.8-b1=elev/runs/hugin-v0.8-en/ckpt_176000.pt:w:1:vocos-b1-40k:en_v2
(vocos-b1-40k.fp16.pt and g2p.fp16.pt are unchanged since v0.5: --vocos elev/runs/vob1/vocos_40000.pt --g2p novel/g2p/g2p.pt;
hift-s3gopen.fp16.pt (the vocoder for own voices from hugin-finetune) and g2p_en.fp16.pt since v0.7: --hift out/runs/s3gopen/s3gen_8000.pt --g2p-en novel/g2p/g2p_en.pt)
The optional 5th field FRONT is stored as info["front"]; "en_v2" switches on the v0.8-en English text path in hugin/synth.py.
Only inference weights are kept (no optimiser state, no GAN discriminators). Floating tensors are stored as fp16 and loaded
back to fp32 at runtime; tiny vectors (< 4096 values: norms, biases, mel statistics, the iSTFT window) stay fp32."""
import argparse, os, torch


def half(sd, small=4096):
    return {k: (v.half() if v.is_floating_point() and v.numel() >= small else v.clone()) for k, v in sd.items()}


def g2p(src, dst):
    g = torch.load(src, map_location="cpu", weights_only=False)       # keys model, vocab
    torch.save({"model": half(g["model"]), "vocab": {"in": list(g["vocab"]["in"]), "out": list(g["vocab"]["out"])}}, dst)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--hugin", action="append", default=[], help="NAME=CKPT:VOICE:TAG:VOCODER[:FRONT] (repeatable)")
    ap.add_argument("--vocos"); ap.add_argument("--hift"); ap.add_argument("--g2p"); ap.add_argument("--g2p-en")
    ap.add_argument("--out", default="weights")
    a = ap.parse_args(); os.makedirs(a.out, exist_ok=True)
    for spec in a.hugin:
        name, rest = spec.split("=", 1); f = rest.split(":"); ck, voice, tag, voc = f[:4]; front = f[4] if len(f) > 4 else None
        s = torch.load(ck, map_location="cpu", weights_only=False)   # keys model, vocab, arch, (opt, D, optD dropped)
        torch.save({"model": half(s["model"]), "vocab": list(s["vocab"]), "arch": dict(s.get("arch", {})),
                    "info": {"name": "futhark-hugin", "version": name.replace("hugin-", ""), "step": int(s.get("step", 0)),
                             "prominence": True, "voice": voice, "tag": int(tag), "vocoder": voc, **({"front": front} if front else {})}},
                   os.path.join(a.out, f"{name}.fp16.pt"))
    if a.vocos:
        v = torch.load(a.vocos, map_location="cpu", weights_only=False)   # keys V (generator), D, oG, oD (dropped)
        torch.save({"V": half(v["V"]), "info": {"name": "vocos-b1", "step": int(v.get("step", 0))}}, os.path.join(a.out, "vocos-b1-40k.fp16.pt"))
    if a.hift:
        g = torch.load(a.hift, map_location="cpu", weights_only=True)    # full S3Gen state dict; keep only the HiFT vocoder
        h = {k[len("mel2wav."):]: v for k, v in g.items() if k.startswith("mel2wav.")}
        torch.save({"hift": half(h), "info": {"name": "hift-s3gopen", "source": "Chatterbox S3Gen HiFT (MIT), fine-tuned in s3gopen (licence-clean Swedish data), step 8000"}},
                   os.path.join(a.out, "hift-s3gopen.fp16.pt"))
    if a.g2p: g2p(a.g2p, os.path.join(a.out, "g2p.fp16.pt"))
    if a.g2p_en: g2p(a.g2p_en, os.path.join(a.out, "g2p_en.fp16.pt"))
    for f in sorted(os.listdir(a.out)):
        if f.endswith(".pt"): print(f, f"{os.path.getsize(os.path.join(a.out, f)) / 1e6:.1f} MB")


if __name__ == "__main__":
    main()