"""Synthesize code-switched ASR training audio with Dynamic Accent Flash-DPO (public-safe, licensed-data TTS).
~/venv/bin/python synth_cs.py <worker> <n_workers>  -> /scratch/asr/synth/wav/*.wav + parts/part_<k>.jsonl
Voices: 1 train clip per WAXAL speaker (ibo/yor/hau 1-8, pcm F1-4/M1-4; CC-BY-SA-4.0 Media Trust) as zero-shot prompt + LM speaker id.
Pace: random LM rate token (slow/normal/fast) + random speed perturbation 0.9-1.1. Channel/noise/reverb augmentation is applied on the fly in training."""
import os, sys, json, random, re, unicodedata, time
os.environ.setdefault("CV3_DIR", "/home/dev/models/Fun-CosyVoice3-0.5B-2512")
REL = "/home/dev/rel_flash"; sys.path[:0] = [REL]
import numpy as np, soundfile as sf, librosa, torch
K, NW = int(sys.argv[1]), int(sys.argv[2])
S_DIR = "/scratch/asr/synth"; os.makedirs(f"{S_DIR}/wav", exist_ok=True); os.makedirs(f"{S_DIR}/voices", exist_ok=True); os.makedirs(f"{S_DIR}/parts", exist_ok=True)
import flash_lm
TONE = {"\u0300", "\u0301", "\u0304", "\u0302", "\u030c"}
def norm(s, lang):
    s = unicodedata.normalize("NFC", s.lower())
    if lang == "ibo": s = unicodedata.normalize("NFC", "".join(c for c in unicodedata.normalize("NFD", s) if c not in TONE))
    s = s.replace("\u2019", "'"); s = re.sub(r"[^\w\s']", " ", s); s = re.sub(r"(^|\s)'+|'+(\s|$)", " ", s)
    return re.sub(r"\s+", " ", s).strip()
TAG = re.compile(r"<(eng|ibo|yor|hau|pcm)>(.*?)</\1>", re.S)

# voice pool (deterministic, built by every worker identically; worker 0 writes the wavs)
W = "/home/dev/data3/waxal"; meta = [json.loads(l) for l in open(f"{W}/meta.jsonl")]
pcm = np.memmap(f"{W}/pcm24k.i16", dtype=np.int16, mode="r")
rng0 = random.Random(3); by = {}
for r in meta:
    if r["split"] == "train" and 4.0 <= r["dur"] <= 9.0 and r.get("ascore", 0) > 0.4: by.setdefault(r["spk"], []).append(r)
voices = {}
for spk in sorted(by):
    r = rng0.choice(by[spk]); p = f"{S_DIR}/voices/{spk}.wav"
    if not os.path.exists(p):
        tmp = p + f".{K}.tmp.wav"; sf.write(tmp, pcm[r["off24"]: r["off24"] + r["n24"]].astype(np.float32) / 32768, 24000); os.replace(tmp, p)
    voices[spk] = dict(wav=p, spk=spk, lang=r["lang"], text=f"<{r['lang']}>{r['text']}</{r['lang']}>")
for v, d in voices.items(): flash_lm.VOICES[v] = dict(wav=d["wav"], spk=d["spk"])
S = flash_lm.Synth(f"{REL}/checkpoints/flash/flash_dpo/best.pt")
for v, d in voices.items(): S.vtext[v] = d["text"]
VL = sorted(voices); print("voices", len(VL), VL, flush=True)
lines = [json.loads(l) for l in open(f"{S_DIR}/cs_text.jsonl")]
mine = [l for i, l in enumerate(lines) if i % NW == K]
out = open(f"{S_DIR}/parts/part_{K}.jsonl", "a"); done = {json.loads(l)["id"] for l in open(f"{S_DIR}/parts/part_{K}.jsonl")}
rng = random.Random(100 + K); t0 = time.time(); n = 0
for L in mine:
    if L["id"] in done: continue
    langs = L["pair"].split("+")
    # prefer a voice whose native language is in the pair (natural accent), sometimes any voice (accent variety)
    pool = [v for v in VL if voices[v]["lang"] in langs] if rng.random() < 0.75 else VL
    v = rng.choice(pool or VL); rate = rng.choice(["slow", "normal", "normal", "normal", "fast"]); seed = rng.randrange(10 ** 6)
    try:
        toks = S.tokens(L["tagged"], v, rate, seed)[0]
        if len(toks) < 10 or len(toks) > 25 * 30: continue
        w = S.wav(toks, v); w = w.squeeze().detach().cpu().numpy() if torch.is_tensor(w) else np.asarray(w).squeeze()
    except Exception as e:
        print("ERR", L["id"], repr(e)[:200], flush=True); continue
    y = librosa.resample(w.astype(np.float32), orig_sr=S.cosy.sample_rate, target_sr=16000)
    sp = rng.uniform(0.9, 1.1)
    if abs(sp - 1) > 0.02: y = librosa.resample(y, orig_sr=16000, target_sr=int(16000 / sp))   # speed perturb (played at 16 kHz)
    y = y / (np.abs(y).max() + 1e-6) * rng.uniform(0.5, 0.95)
    p = f"{S_DIR}/wav/{L['id']}.wav"; sf.write(p, y, 16000, subtype="PCM_16")
    text = " ".join(norm(t, l) for l, t in TAG.findall(L["tagged"]))
    out.write(json.dumps(dict(id=L["id"], wav=p, dur=round(len(y) / 16000, 2), lang="cs", src="synth_flash_dpo", text=text, tagged=L["tagged"],
                              pair=L["pair"], voice=v, rate=rate, speed=round(sp, 3), licence="synthetic (Flash-DPO, licensed-data TTS)"), ensure_ascii=False) + "\n")
    out.flush(); n += 1
    if n % 20 == 0: print(f"w{K} {n}/{len(mine)} {(time.time() - t0) / n:.2f}s/clip", flush=True)
print("SYNTH_DONE", K, n, flush=True)