#!/usr/bin/env python3
"""RANK-1 acoustic probe: synth N training clips with FORCED (ground-truth) durations through an
ONNX dir, isolating the phones->mel mapping from the duration predictor + g2p frontend.
Reads phone/tone/lang ids + GT durations directly from m2_align.jsonl (no frontend).
Run in moss-train-venv. Then ASR the wavs (probe_forced_asr via xasr_offline) -> CER.
Pairs the 0.80(forced)-vs-0.145(GT-mel) acoustic gap to a single number per config."""
import argparse, json, sys
from pathlib import Path
import numpy as np, soundfile as sf, onnxruntime as ort
ZT = "/home/luigi/jetson-tts/mossnano/zhtw8k"
sys.path.insert(0, ZT)
from synth_from_text import host_regulate


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--onnx-dir", required=True)
    ap.add_argument("--out-dir", required=True)
    ap.add_argument("--n", type=int, default=30)
    ap.add_argument("--align-jsonl", default=f"{ZT}/m2_align.jsonl")
    args = ap.parse_args()
    meta = json.load(open(f"{args.onnx_dir}/meta.json"))
    so = ort.SessionOptions(); so.intra_op_num_threads = 4
    sA = ort.InferenceSession(f"{args.onnx_dir}/acoustic_encoder.onnx", so, providers=["CPUExecutionProvider"])
    sB = ort.InferenceSession(f"{args.onnx_dir}/acoustic_decoder.onnx", so, providers=["CPUExecutionProvider"])
    sV = ort.InferenceSession(f"{args.onnx_dir}/vocoder.onnx", so, providers=["CPUExecutionProvider"])
    bn = ["frames", "frame_meta", "local_ctx_raw", "abs_pos", "pitch_frame", "frame_mask"]
    Path(args.out_dir).mkdir(parents=True, exist_ok=True)
    rows = [json.loads(l) for l in open(args.align_jsonl) if l.strip()][:args.n]
    out = open(f"{args.out_dir}/synth.jsonl", "w")
    for i, r in enumerate(rows):
        phone = np.array([r["phone_ids"]], np.int64); tone = np.array([r["tone_ids"]], np.int64)
        lang = np.array([r["lang_ids"]], np.int64); spk = np.zeros(1, np.int64)
        cond, _dur_pred, pitch = sA.run(None, {"phone": phone, "tone": tone, "lang": lang, "speaker": spk})
        # substitute GT (forced) durations, rescaled so total ~ predicted length (stable regulator)
        df = np.array([r["hifigan_durations"]], np.float32)
        df = df * (_dur_pred.sum() / max(1.0, df.sum()))
        reg = host_regulate(cond, df, pitch, meta["abs_frame_bins"], meta["max_frames"])
        feeds = {n: (reg[n].astype(np.float32) if reg[n].dtype != bool else reg[n]) for n in bn}
        feeds["abs_pos"] = reg["abs_pos"].astype(np.int64)
        mel = sB.run(None, feeds)[0]
        wav = sV.run(None, {"mel": mel.astype(np.float32)})[0].reshape(-1)
        wp = f"{args.out_dir}/p{i:03d}.wav"; sf.write(wp, wav, meta["sample_rate"])
        out.write(json.dumps({"id": f"p{i:03d}", "text": r["text"], "wav": wp}, ensure_ascii=False) + "\n")
    out.close()
    print(f"PROBE SYNTH DONE {len(rows)} clips -> {args.out_dir}/synth.jsonl")


if __name__ == "__main__":
    main()