"""Build held-out ASR/LID eval set -> /scratch/duplex/asr_eval/{wav/*.wav, list.jsonl}
Mono: WAXAL test (ibo/yor/hau/pcm, 40 each, spread over speakers), FLEURS val eng (30).
CS: 40 synthetic code-switched clips = two held-out clips of different languages concatenated (0.25 s gap), ref = joined text + per-word lang labels.
All splits are test/val/dev splits never used in TTS training; whether the ASR models saw them is unknown (noted in report)."""
import json, os, random, numpy as np, soundfile as sf, librosa
H = os.path.expanduser("~"); OUT = "/scratch/duplex/asr_eval"; os.makedirs(OUT + "/wav", exist_ok=True)
random.seed(7)
def rows(d, pred):
    R = [json.loads(l) for l in open(f"{d}/meta.jsonl")]
    return [r for r in R if pred(r)]
def audio(d, r):
    a = np.memmap(f"{d}/pcm24k.i16", dtype=np.int16, mode="r")
    y = a[r["off24"]: r["off24"] + r["n24"]].astype(np.float32) / 32768
    return librosa.resample(y, orig_sr=24000, target_sr=16000)
def pick(R, n):
    by = {}
    for r in R: by.setdefault(r.get("spk"), []).append(r)
    for v in by.values(): random.shuffle(v)
    out = []
    while len(out) < n and any(by.values()):
        for k in list(by):
            if by[k] and len(out) < n: out.append(by[k].pop())
    return out
items = []
W = f"{H}/data3/waxal"
for L in ["ibo", "yor", "hau", "pcm"]:
    for r in pick(rows(W, lambda r: r["lang"] == L and r["split"] == "test" and 1.5 < r["dur"] < 14), 40):
        items.append(("waxal_test", W, r, L))
F = f"{H}/data3/fleurs"
for r in pick(rows(F, lambda r: r["lang"] == "eng" and r["split"] == "val" and r["dur"] < 16), 30): items.append(("fleurs_val", F, r, "eng"))
recs = []; pool = {}
for src, d, r, L in items:
    y = audio(d, r); i = f"{L}_{len(recs):03d}"
    sf.write(f"{OUT}/wav/{i}.wav", y, 16000)
    rec = dict(id=i, wav=f"{OUT}/wav/{i}.wav", lang=L, text=r["text"], src=f"{src}:{r['id']}", dur=round(len(y) / 16000, 2),
               wlang=[L] * len(r["text"].split()))
    recs.append(rec); pool.setdefault(L, []).append((rec, y))
pairs = [("eng", "ibo"), ("ibo", "eng"), ("eng", "yor"), ("yor", "eng"), ("eng", "hau"), ("hau", "eng"), ("pcm", "eng"), ("eng", "pcm"), ("pcm", "ibo"), ("pcm", "yor")]
for k in range(40):
    a, b = pairs[k % len(pairs)]
    (ra, ya), (rb, yb) = random.choice(pool[a]), random.choice(pool[b])
    y = np.concatenate([ya, np.zeros(4000, np.float32), yb]); i = f"cs_{k:03d}"
    sf.write(f"{OUT}/wav/{i}.wav", y, 16000)
    recs.append(dict(id=i, wav=f"{OUT}/wav/{i}.wav", lang="cs", langs=[a, b], text=ra["text"] + " " + rb["text"], src=f"concat:{ra['id']}+{rb['id']}",
                     dur=round(len(y) / 16000, 2), wlang=ra["wlang"] + rb["wlang"]))
with open(f"{OUT}/list.jsonl", "w") as f:
    for r in recs: f.write(json.dumps(r, ensure_ascii=False) + "\n")
from collections import Counter; print(Counter(r["lang"] for r in recs)); print("BUILD_DONE", len(recs))