"""Merge synth parts, split held-out synthetic test (every 25th text id), and (after eval_asr.py --hyp_only) filter TTS failures.
python3 build_synth.py merge   -> /scratch/asr/synth/all.jsonl
python3 build_synth.py filter  -> /scratch/asr/synth/{manifest.jsonl (train), test.jsonl}  (keep clips the stock v2 model reads at CER <= 0.35)"""
import json, glob, sys, collections
D = "/scratch/asr/synth"
if sys.argv[1] == "merge":
    R = [json.loads(l) for f in sorted(glob.glob(f"{D}/parts/part_*.jsonl")) for l in open(f)]
    R = [r for r in R if 1.0 <= r["dur"] <= 20 and r["text"]]
    with open(f"{D}/all.jsonl", "w") as f:
        for r in R: f.write(json.dumps(r, ensure_ascii=False) + "\n")
    print("merged", len(R), "hours", round(sum(r["dur"] for r in R) / 3600, 2))
else:
    R = [json.loads(l) for l in open(f"{D}/all.jsonl")]; H = {json.loads(l)["id"]: json.loads(l) for l in open(f"{D}/all.jsonl.hyp.jsonl")}
    tr, te, drop = [], [], collections.Counter()
    for r in R:
        c = H[r["id"]]["cer"]; r["base_cer"] = c
        if c > 0.35: drop[r["pair"]] += 1; continue
        (te if int(r["id"][2:]) % 25 == 0 else tr).append(r)
    for name, X in (("manifest", tr), ("test", te)):
        with open(f"{D}/{name}.jsonl", "w") as f:
            for r in X: f.write(json.dumps(r, ensure_ascii=False) + "\n")
    print("train", len(tr), round(sum(r["dur"] for r in tr) / 3600, 2), "h; test", len(te), "; dropped (base CER>0.35)", sum(drop.values()), dict(drop))
    print("pairs train", collections.Counter(r["pair"] for r in tr))