"""Convert g2p_en 2.1.0 checkpoint20.npz to naptrap's safetensors layout.

Usage: python convert.py path/to/g2p_en/checkpoint20.npz OUTPUT_DIR
Arrays are copied unchanged (float32); only their names change.
"""

import hashlib
import json
import sys
from pathlib import Path

import numpy as np
import torch
from safetensors.torch import save_file

SOURCE_SHA256 = "b8af35e4596d8dd5836dfd3fe9b2ba4f97b9c311efe8879544cbcfcbd566d8c6"
NAMES = {
    "enc_emb": "enc_emb", "dec_emb": "dec_emb", "fc_w": "fc_w", "fc_b": "fc_b",
    **{f"{side}_{p}": f"{'encoder' if side == 'enc' else 'decoder'}.{p}"
       for side in ("enc", "dec") for p in ("w_ih", "w_hh", "b_ih", "b_hh")},
}


def main(source, output):
    source, output = Path(source), Path(output)
    if hashlib.sha256(source.read_bytes()).hexdigest() != SOURCE_SHA256:
        raise SystemExit("Unexpected checkpoint20.npz; expected the g2p_en 2.1.0 release")
    arrays = np.load(source, allow_pickle=False)
    if set(arrays.files) != set(NAMES):
        raise SystemExit(f"Unexpected arrays: {sorted(arrays.files)}")
    state = {NAMES[k]: torch.from_numpy(np.ascontiguousarray(arrays[k], dtype=np.float32)) for k in arrays.files}
    output.mkdir(parents=True, exist_ok=True)
    weights = output / "model.safetensors"
    save_file(state, str(weights), metadata={"format": "naptrap-g2p-en-decoder-v1"})
    config = {
        "format": "naptrap-g2p-en-decoder-v1",
        "hidden_size": int(arrays["dec_emb"].shape[1]),
        "graphemes": int(arrays["enc_emb"].shape[0]),
        "phonemes": int(arrays["dec_emb"].shape[0]),
        "max_decoder_steps": 20,
        "output": "greedy decoder hidden states per word, separated by the <pad> phoneme embedding",
        "source": {"package": "g2p_en", "version": "2.1.0", "file": "g2p_en/checkpoint20.npz",
                   "sha256": SOURCE_SHA256, "license": "Apache-2.0"},
        "sha256": {"model.safetensors": hashlib.sha256(weights.read_bytes()).hexdigest()},
    }
    (output / "config.json").write_text(json.dumps(config, indent=2) + "\n")
    for k in arrays.files:
        if not np.array_equal(arrays[k], state[NAMES[k]].numpy()):
            raise SystemExit("Conversion changed " + k)


if __name__ == "__main__":
    main(*sys.argv[1:])