"""Phase 2: S3 speech tokens (speech_tokenizer_v2.onnx, CUDA EP) + CAM++ speaker embeddings for ~/data2/<ds>,
plus code-switched training pairs (ds "cs_pairs").

  python -m da.s3_2 tok waxal fleurs ...        -> s3tok.i32, s3tok.jsonl, campplus.npy
  python -m da.s3_2 cs 4000                     -> ~/data2/cs_pairs/{pcm24k.i16, meta.jsonl} (+ run tok on it)
"""
import os, sys, json, random, hashlib
import numpy as np, librosa, torch
from multiprocessing import Pool

D2 = os.path.expanduser(os.environ.get("P2OUT", "~/data2"))
MD = os.path.expanduser("~/models/CosyVoice2-0.5B")
_S = {}
_FB = torch.from_numpy(librosa.filters.mel(sr=16000, n_fft=400, n_mels=128)).float()
_WIN = torch.hann_window(400)


def log_mel_128(a):
    st = torch.stft(a, 400, 160, window=_WIN, return_complex=True)
    mag = st[..., :-1].abs() ** 2
    lm = torch.clamp(_FB @ mag, min=1e-10).log10()
    lm = torch.maximum(lm, lm.max() - 8.0)
    return (lm + 4.0) / 4.0


def sess(name, gpu):
    if name not in _S:
        import onnxruntime as ort
        torch.set_num_threads(1)
        o = ort.SessionOptions(); o.intra_op_num_threads = 2; o.inter_op_num_threads = 1
        prov = ["CUDAExecutionProvider", "CPUExecutionProvider"] if gpu else ["CPUExecutionProvider"]
        _S[name] = ort.InferenceSession(f"{MD}/{name}", sess_options=o, providers=prov)
    return _S[name]


def work(args):
    path, off, n = args
    import torchaudio.compliance.kaldi as kaldi
    pcm = np.memmap(path, dtype=np.int16, mode="r")
    y24 = pcm[off:off + n].astype(np.float32) / 32768.0
    y16 = librosa.resample(y24, orig_sr=24000, target_sr=16000, res_type="soxr_hq")
    a = torch.from_numpy(y16)
    f = log_mel_128(a)[None]
    s = sess("speech_tokenizer_v2.onnx", True)
    ins = s.get_inputs()
    tok = s.run(None, {ins[0].name: f.numpy(), ins[1].name: np.array([f.shape[2]], dtype=np.int32)})[0]
    fb = kaldi.fbank(a[None], num_mel_bins=80, dither=0, sample_frequency=16000)
    fb = fb - fb.mean(0, keepdim=True)
    c = sess("campplus.onnx", False)
    emb = c.run(None, {c.get_inputs()[0].name: fb[None].numpy()})[0].reshape(-1).astype(np.float32)
    return np.asarray(tok).reshape(-1).astype(np.int32), emb


def tok(ds):
    od = f"{D2}/{ds}"
    meta = [json.loads(l) for l in open(f"{od}/meta.jsonl", encoding="utf-8")]
    jobs = [(f"{od}/pcm24k.i16", m["off24"], m["n24"]) for m in meta]
    toff, embs = 0, []
    with open(f"{od}/s3tok.i32", "wb") as tf, open(f"{od}/s3tok.jsonl", "w") as jf, Pool(12) as pool:
        for i, (m, (t, e)) in enumerate(zip(meta, pool.imap(work, jobs, chunksize=4))):
            tf.write(t.tobytes()); embs.append(e)
            jf.write(json.dumps({"id": m["id"], "toff": toff, "ntok": int(len(t))}) + "\n")
            toff += len(t)
            if i % 5000 == 0:
                print(ds, i, len(meta), "tok/s", round(len(t) / max(m["dur"], 1e-3), 1), flush=True)
    np.save(f"{od}/campplus.npy", np.stack(embs))
    print(ds, "S3_DONE", toff, flush=True)


def build_cs(n_pairs=4000, seed=0):
    """Join two train utterances in different languages (prefer same gender) into one code-switched item.
    Junction silence follows the first clip's final punctuation (sentence 0.45 s, comma 0.25 s, else 0.20 s)."""
    rng = random.Random(seed)
    pools = {}
    for ds in ["waxal", "slr70", "slr86", "fleurs"]:
        p = f"{D2}/{ds}/meta.jsonl"
        if not os.path.exists(p):
            continue
        for l in open(p, encoding="utf-8"):
            m = json.loads(l)
            if m["split"] != "train" or not (1.0 <= m["dur"] <= 6.5):
                continue
            if ds == "fleurs" and m["lang"] != "eng":
                continue
            pools.setdefault(m["lang"], []).append(m)
    langs = sorted(pools)
    print("cs pools", {l: len(v) for l, v in pools.items()})
    od = f"{D2}/cs_pairs"; os.makedirs(od, exist_ok=True)
    mm = {}
    off = 0
    with open(f"{od}/pcm24k.i16", "wb") as pf, open(f"{od}/meta.jsonl", "w", encoding="utf-8") as jf:
        for k in range(n_pairs):
            l1, l2 = rng.sample(langs, 2)
            if rng.random() < 0.5 and "eng" in langs and "eng" not in (l1, l2):
                l2 = "eng"          # English is the most common partner in real Nigerian code-switching
            a = rng.choice(pools[l1])
            g = (a.get("gender") or "").lower()[:1]
            cand = [b for b in rng.sample(pools[l2], min(50, len(pools[l2]))) if (b.get("gender") or "").lower()[:1] == g]
            b = cand[0] if cand else rng.choice(pools[l2])
            ys = []
            for m in (a, b):
                if m["ds"] not in mm:
                    mm[m["ds"]] = np.memmap(f"{D2}/{m['ds']}/pcm24k.i16", dtype=np.int16, mode="r")
                ys.append(np.asarray(mm[m["ds"]][m["off24"]:m["off24"] + m["n24"]]))
            last = a["text"].rstrip()[-1:]
            junc = 0.45 if last in ".?!" else (0.25 if last in ",;:" else 0.20)
            have = 0.18 + 0.12
            y1 = ys[0]
            if junc < have:
                y1 = y1[: len(y1) - int((have - junc) * 24000)]
                gap = np.zeros(0, np.int16)
            else:
                gap = np.zeros(int((junc - have) * 24000), np.int16)
            y = np.concatenate([y1, gap, ys[1]])
            text = f"<{l1}>{a['text']}</{l1}> <{l2}>{b['text']}</{l2}>"
            m = dict(id=f"cs_{k:05d}", ds="cs_pairs", lang="cs", langs=[l1, l2], spk=a["spk"], gender=a.get("gender", ""),
                     split="train", text=text, dur=round(len(y) / 24000, 3), off24=off, n24=int(len(y)),
                     src=[a["id"], b["id"]], flow_ok=False)
            pf.write(y.tobytes()); off += len(y)
            jf.write(json.dumps(m, ensure_ascii=False) + "\n")
    print("cs pairs", n_pairs, "hours", round(off / 24000 / 3600, 2), flush=True)


if __name__ == "__main__":
    if sys.argv[1] == "tok":
        for ds in sys.argv[2:]:
            tok(ds)
    elif sys.argv[1] == "cs":
        build_cs(int(sys.argv[2]))
        tok("cs_pairs")