"""Arm B, Phase 2. CosyVoice2-0.5B (Apache-2.0) + language embeddings + speaker token + GRL speaker classifier.

New in Phase 2
  * data from ~/data2 (24 kHz, forced-aligned, pause-normalised; long WAXAL clips segmented; BibleTTS Hausa; CS pairs)
  * in-context prompts during LM training (p=0.5): [sos, prompt_text+text, spk, task, prompt_speech, speech] with
    loss on the target speech only -> at inference the anchor clip steers prosody and pause style
  * keep-best on validation CE (every --val_every steps), cosine LR, early stop
  * flow-matching decoder fine-tune on 24 kHz studio sources (tones / pitch / timing detail)
  * synthesis: anchor prompt at LM and flow, sentence-level generation for multi-sentence text with anchor-matched
    inter-sentence pauses, N candidates per text (seed + prompt clip vary)

  python train_b2.py lm    --out ~/runs/p2_lm  --minutes 90
  python train_b2.py flow  --out ~/runs/p2_flow --minutes 60
  python train_b2.py synth --ckpt ~/runs/p2_lm/best.pt [--flow ~/runs/p2_flow/best.pt] --out DIR --n 4
"""
import argparse, inspect, json, math, os, random, re, sys, time
import numpy as np, torch, torch.nn as nn, torch.nn.functional as F, soundfile as sf, librosa

HOME = os.path.expanduser("~")
sys.path[:0] = [f"{HOME}/CosyVoice", f"{HOME}/CosyVoice/third_party/Matcha-TTS", f"{HOME}/da"]
from train_b import load_cosy, text_tokens, DALM, LANGS, L2I, MODEL_DIR, token2wav  # noqa: E402
from da.model import grad_reverse  # noqa: E402
from da.text import segments  # noqa: E402

D2 = os.path.expanduser(os.environ.get("P2OUT", f"{HOME}/data2"))
ANCHOR = "waxal_yor_4"


# ------------------------------------------------------------------ model
class DALM2(DALM):
    def build2(self, ids, langs, spk, sp=None, pids=None, plangs=None, psp=None):
        """[sos, (prompt_text) text, spk, task, (prompt_speech) speech] -> embeddings, n_text_total, n_prompt_speech"""
        sos = self.q.llm_embedding.weight[self.q.sos][None]
        task = self.q.llm_embedding.weight[self.q.task_id][None]
        if pids is not None:
            ids = torch.cat([pids, ids]); langs = torch.cat([plangs, langs])
        parts = [sos, self.embed_text(ids, langs), self.spk_emb(spk)[None], task]
        npsp = 0
        if psp is not None:
            parts.append(self.q.speech_embedding(psp)); npsp = len(psp)
        if sp is not None:
            parts.append(self.q.speech_embedding(sp))
        return torch.cat(parts, 0), len(ids), npsp

    def forward2(self, batch, adv_lambda=1.0):
        seqs, tgts, tpos, spks = [], [], [], []
        for b in batch:
            ids, langs, spk, sp = b[:4]
            pr = b[4] if len(b) > 4 else None
            x, nt, npsp = self.build2(ids, langs, spk, sp, *(pr if pr else (None, None, None)))
            tgt = torch.full((x.shape[0] + 1,), -100, device=x.device, dtype=torch.long)
            start = 1 + nt + 2 + npsp - 1          # position whose output predicts speech[0]
            tgt[start: start + len(sp)] = sp
            tgt[start + len(sp)] = self.V
            seqs.append(x); tgts.append(tgt[:x.shape[0]]); tpos.append((1, 1 + nt)); spks.append(int(spk))
        T = max(s.shape[0] for s in seqs)
        D = seqs[0].shape[1]
        X = seqs[0].new_zeros(len(seqs), T, D)
        Y = torch.full((len(seqs), T), -100, device=X.device, dtype=torch.long)
        M = torch.zeros(len(seqs), T, device=X.device, dtype=torch.long)
        for i, (s, t) in enumerate(zip(seqs, tgts)):
            X[i, :len(s)] = s; Y[i, :len(t)] = t; M[i, :len(s)] = 1
        H = self.hidden(X, M)
        logits = self.q.llm_decoder(H)
        ce = F.cross_entropy(logits.float().reshape(-1, logits.shape[-1]), Y.reshape(-1), ignore_index=-100)
        valid = Y != -100
        acc = (logits.argmax(-1)[valid] == Y[valid]).float().mean()
        hs, ys = [], []
        for i, (a, b) in enumerate(tpos):
            hs.append(H[i, a:b]); ys.append(torch.full((b - a,), spks[i], device=X.device))
        hs = torch.cat(hs); ys = torch.cat(ys)
        al = self.adv(grad_reverse(hs.float(), adv_lambda))
        adv = F.cross_entropy(al, ys)
        return ce, acc, adv, (al.argmax(-1) == ys).float().mean()

    @torch.no_grad()
    def generate2(self, ids, langs, spk, prompt=None, top_p=0.8, top_k=25, win=10, tau_r=0.1, seed=0,
                  min_ratio=2.0, max_ratio=20.0):
        g = torch.Generator(device=ids.device).manual_seed(seed)
        x, _, _ = self.build2(ids, langs, spk, None, *(prompt if prompt else (None, None, None)))
        x = x[None]
        n = len(ids)
        min_len, max_len = int(n * min_ratio), int(n * max_ratio)
        past, out, inp = None, [], x
        for i in range(max_len):
            o = self.q.llm.model.model(inputs_embeds=inp, past_key_values=past, use_cache=True, return_dict=True)
            past = o.past_key_values
            logp = self.q.llm_decoder(o.last_hidden_state[:, -1]).float().log_softmax(-1)[0]
            if i < min_len:
                logp[self.V] = -float("inf")
            logp[self.V + 1:] = -float("inf")
            p = logp.exp()
            sp, si = p.sort(descending=True)
            keep = (sp.cumsum(0) - sp) < top_p
            keep[top_k:] = False
            cp, ci = sp[keep], si[keep]
            tokid = int(ci[torch.multinomial(cp / cp.sum(), 1, generator=g)])
            if out and sum(1 for t in out[-win:] if t == tokid) >= win * tau_r and tokid != self.V:
                tokid = int(torch.multinomial(p, 1, generator=g))
            if tokid >= self.V:
                break
            out.append(tokid)
            inp = self.q.speech_embedding.weight[tokid][None, None]
        return out


# ------------------------------------------------------------------ data
def load_items(datasets, split, max_tok=500, flow=False):
    items, mm = [], {}
    for ds in datasets:
        od = f"{D2}/{ds}"
        if not os.path.exists(f"{od}/s3tok.jsonl"):
            print("missing", ds, flush=True); continue
        mm[ds] = np.memmap(f"{od}/s3tok.i32", dtype=np.int32, mode="r")
        tk = [json.loads(l) for l in open(f"{od}/s3tok.jsonl")]
        for k, l in enumerate(open(f"{od}/meta.jsonl", encoding="utf-8")):
            m = json.loads(l)
            t = tk[k]
            assert t["id"] == m["id"]
            if m["split"] != split:
                continue
            if flow and not m.get("flow_ok"):
                continue
            m["toff"], m["ntok"], m["row"] = t["toff"], t["ntok"], k
            if 5 < m["ntok"] <= max_tok:
                items.append(m)
    return items, mm


def tagged(m):
    return m["text"] if m["lang"] == "cs" else f"<{m['lang']}>{m['text']}</{m['lang']}>"


DS_W = {"waxal": 2.0, "bible_yor": 0.5, "bible_hau": 0.6, "fleurs": 1.0, "slr86": 1.0, "slr70": 1.0, "cs_pairs": 1.0}


def batches(items, budget, rng, extra_w=None, cost=None):
    idx = []
    for i, m in enumerate(items):
        w = DS_W.get(m["ds"], 1.0) * (extra_w(m) if extra_w else 1.0)
        idx += [i] * int(w) + ([i] if rng.random() < w - int(w) else [])
    cost = cost or (lambda m: m["ntok"] + len(m["text"]) // 2 + 4)
    idx.sort(key=lambda i: cost(items[i]) + rng.random() * 10)
    out, cur, mx = [], [], 0
    for i in idx:
        L = cost(items[i])
        if cur and max(mx, L) * (len(cur) + 1) > budget:
            out.append(cur); cur, mx = [], 0
        cur.append(i); mx = max(mx, L)
    if cur:
        out.append(cur)
    rng.shuffle(out)
    return out


# ------------------------------------------------------------------ LM training
def train_lm(a):
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained(f"{MODEL_DIR}/CosyVoice-BlankEN")
    cosy = load_cosy()
    dss = a.datasets.split(",")
    items, mm = load_items(dss, "train")
    val_all, _ = load_items(dss, "val")
    if a.anchor_only_w > 0:
        print("anchor items", sum(m["spk"] == ANCHOR for m in items), flush=True)
    spks = sorted({m["spk"] for m in items})
    if a.init:
        ck = torch.load(a.init, map_location="cpu")
        spk2id = json.load(open(os.path.join(os.path.dirname(a.init), "spk2id.json")))
        for s in spks:
            if s not in spk2id:
                spk2id[s] = len(spk2id)
    else:
        spk2id = {s: i for i, s in enumerate(spks)}
    n_spk = max(len(spk2id), ck["n_spk"]) if a.init else len(spk2id)
    # balanced val: up to 60 per language (+ anchor), never more than 15 per speaker
    rng = random.Random(1); rng.shuffle(val_all)
    val, per_l, per_s = [], {}, {}
    for m in val_all:
        if m["spk"] not in spk2id or m["lang"] == "cs":
            continue
        L = "anchor" if m["spk"] == ANCHOR else m["lang"]
        if per_l.get(L, 0) >= 60 or per_s.get(m["spk"], 0) >= 15:
            continue
        per_l[L] = per_l.get(L, 0) + 1; per_s[m["spk"]] = per_s.get(m["spk"], 0) + 1; val.append(m)
    by_spk = {}
    for m in items:
        if m["ntok"] <= 160 and m["lang"] != "cs":
            by_spk.setdefault(m["spk"], []).append(m)
    os.makedirs(a.out, exist_ok=True)
    json.dump(spk2id, open(f"{a.out}/spk2id.json", "w"))
    print("lm items", len(items), "hours", round(sum(m["dur"] for m in items) / 3600, 1), "spk", len(spk2id),
          "val", len(val), per_l, flush=True)
    model = DALM2(cosy.model.llm, n_spk).cuda().train()
    if a.init:
        sd = ck["model"]
        if sd["spk_emb.weight"].shape[0] < n_spk:
            for k in ["spk_emb.weight"]:
                w = model.state_dict()[k].clone(); w[:sd[k].shape[0]] = sd[k]; sd[k] = w
            for k in ["adv.2.weight", "adv.2.bias"]:
                w = model.state_dict()[k].clone(); w[:sd[k].shape[0]] = sd[k]; sd[k] = w
        print(model.load_state_dict(sd, strict=False), flush=True)
    for p in cosy.model.flow.parameters():
        p.requires_grad = False
    for p in cosy.model.hift.parameters():
        p.requires_grad = False
    new = list(model.lang_emb.parameters()) + list(model.spk_emb.parameters()) + list(model.adv.parameters())
    newids = {id(p) for p in new}
    old = [p for p in model.parameters() if id(p) not in newids and p.requires_grad]
    opt = torch.optim.AdamW([{"params": old, "lr": a.lr}, {"params": new, "lr": a.lr_new}], weight_decay=0.01)
    base = [a.lr, a.lr_new]
    cache = {}

    def tt(text, key):
        if key not in cache:
            cache[key] = text_tokens(tok, text)
        return cache[key]

    def enc(m, prng=None, force_prompt=None):
        ids, langs = tt(tagged(m), m["id"])
        sp = np.asarray(mm[m["ds"]][m["toff"]: m["toff"] + m["ntok"]], dtype=np.int64)
        r = [torch.tensor(ids).cuda(), torch.tensor(langs).cuda(), torch.tensor(spk2id[m["spk"]]).cuda(),
             torch.from_numpy(sp).cuda()]
        use = force_prompt if force_prompt is not None else (prng is not None and prng.random() < a.p_prompt)
        if use:
            c = by_spk.get(m["spk"], [])
            c = [x for x in (prng or random).sample(c, min(6, len(c))) if x["id"] != m["id"]] if c else []
            if c and c[0]["ntok"] + m["ntok"] <= 600:
                p = c[0]
                pids, plangs = tt(tagged(p), p["id"])
                psp = np.asarray(mm[p["ds"]][p["toff"]: p["toff"] + p["ntok"]], dtype=np.int64)
                r.append((torch.tensor(pids).cuda(), torch.tensor(plangs).cuda(), torch.from_numpy(psp).cuda()))
        return r

    def validate():
        model.eval(); out = {}
        with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
            for mode in (False, True):
                vr = random.Random(7); L, W = [], []
                for i in range(0, len(val), 16):
                    bt = val[i:i + 16]
                    ce, _, _, _ = model.forward2([enc(m, vr, force_prompt=mode) for m in bt])
                    L.append(float(ce) * len(bt)); W.append(len(bt))
                out["p" if mode else "np"] = sum(L) / sum(W)
            # anchor-only val
            av = [m for m in val if m["spk"] == ANCHOR]
            if av:
                vr = random.Random(7)
                ce, _, _, _ = model.forward2([enc(m, vr, force_prompt=True) for m in av[:32]])
                out["anchor_p"] = float(ce)
        model.train()
        out["mean"] = (out["np"] + out["p"]) / 2
        return out

    extra = (lambda m: a.anchor_only_w if m["spk"] == ANCHOR else 1.0) if a.anchor_only_w > 0 else None
    rng = random.Random(0); prng = random.Random(3)
    step, nskip, t0 = 0, 0, time.time()
    logf = open(f"{a.out}/log.jsonl", "a")
    v = validate(); best = v["mean"] if not a.select_anchor else 0.5 * (v["mean"] + v.get("anchor_p", v["mean"])); bad = 0
    logf.write(json.dumps({"step": 0, "val": v}) + "\n"); print(json.dumps({"step": 0, "val": v}), flush=True)
    torch.save(dict(model=model.state_dict(), step=0, n_spk=n_spk, val=v), f"{a.out}/best.pt")
    total = a.max_steps
    done = False
    while not done:
        for bi in batches(items, a.tokens, rng, extra, cost=lambda m: m["ntok"] + len(m["text"]) // 2 + 4 + int(a.p_prompt * 200)):
            lrm = min(1.0, (step + 1) / a.warmup) * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * min(1.0, step / total))))
            for g, b in zip(opt.param_groups, base):
                g["lr"] = b * lrm
            batch = [enc(items[i], prng) for i in bi]
            # keep the padded length inside the budget when prompts were added
            with torch.autocast("cuda", dtype=torch.bfloat16):
                ce, acc, adv, adv_acc = model.forward2(batch, adv_lambda=min(1.0, step / 1000))
            loss = ce + a.adv_w * adv
            opt.zero_grad(set_to_none=True)
            if not torch.isfinite(loss):
                nskip += 1; print(json.dumps({"skip_nonfinite_fwd": step, "ce": float(ce), "adv": float(adv)}), flush=True)
                continue
            loss.backward()
            gn = torch.nn.utils.clip_grad_norm_([p for p in model.parameters() if p.requires_grad], 1.0)
            if not torch.isfinite(gn):
                nskip += 1; opt.zero_grad(set_to_none=True); print(json.dumps({"skip_nonfinite_grad": step}), flush=True)
                continue
            opt.step(); step += 1
            if step % 25 == 0:
                r = dict(step=step, ce=round(float(ce), 4), acc=round(float(acc), 4), adv=round(float(adv), 3),
                         adv_acc=round(float(adv_acc), 3), gn=round(float(gn), 2), B=len(bi), lr=round(lrm * a.lr, 9),
                         skips=nskip, min=round((time.time() - t0) / 60, 2))
                print(json.dumps(r), flush=True); logf.write(json.dumps(r) + "\n"); logf.flush()
            if step % a.val_every == 0:
                v = validate()
                score = v["mean"] if not a.select_anchor else 0.5 * (v["mean"] + v.get("anchor_p", v["mean"]))
                imp = score < best - 1e-4
                rec = {"step": step, "val": v, "best": imp}
                print(json.dumps(rec), flush=True); logf.write(json.dumps(rec) + "\n"); logf.flush()
                if imp:
                    best, bad = score, 0
                    torch.save(dict(model=model.state_dict(), step=step, n_spk=n_spk, val=v), f"{a.out}/best.pt")
                else:
                    bad += 1
                    if bad >= a.patience:
                        print("EARLY_STOP", step, flush=True); done = True; break
            if (time.time() - t0) / 60 > a.minutes or step >= total:
                done = True; break
    print("TRAIN_LM_DONE", step, "best", best, "skips", nskip, flush=True)


# ------------------------------------------------------------------ flow training
def train_flow(a):
    cosy = load_cosy()
    flow = cosy.model.flow
    if a.flow_init:
        flow.load_state_dict(torch.load(a.flow_init, map_location="cuda")["flow"])
    flow.train()
    for p in cosy.model.llm.parameters():
        p.requires_grad = False
    fe = cosy.frontend.feat_extractor
    dss = a.datasets.split(",")
    items, mm = load_items(dss, "train", max_tok=400, flow=True)
    val, _ = load_items(dss, "val", max_tok=400, flow=True)
    rng0 = random.Random(1); rng0.shuffle(val); val = val[:256]
    pcm = {ds: np.memmap(f"{D2}/{ds}/pcm24k.i16", dtype=np.int16, mode="r") for ds in dss if os.path.exists(f"{D2}/{ds}/pcm24k.i16")}
    emb = {ds: np.load(f"{D2}/{ds}/campplus.npy", mmap_mode="r") for ds in dss if os.path.exists(f"{D2}/{ds}/campplus.npy")}
    os.makedirs(a.out, exist_ok=True)
    print("flow items", len(items), "hours", round(sum(m["dur"] for m in items) / 3600, 1), "val", len(val), flush=True)
    extra = (lambda m: a.anchor_only_w if m["spk"] == ANCHOR else 1.0) if a.anchor_only_w > 0 else None

    def make(bt):
        toks, feats, es = [], [], []
        for m in bt:
            y = torch.from_numpy(np.asarray(pcm[m["ds"]][m["off24"]: m["off24"] + m["n24"]], dtype=np.float32) / 32768.0)
            n = int(math.ceil(len(y) / 960)) * 960
            y = F.pad(y, (0, n - len(y)))[None].cuda()
            ft = fe(y).squeeze(0).transpose(0, 1)         # [T,80]
            tk = torch.from_numpy(np.asarray(mm[m["ds"]][m["toff"]: m["toff"] + m["ntok"]], dtype=np.int64))
            k = min(ft.shape[0] // 2, len(tk))
            toks.append(tk[:k]); feats.append(ft[:2 * k]); es.append(torch.from_numpy(np.array(emb[m["ds"]][m["row"]])))
        tl = torch.tensor([len(t) for t in toks], dtype=torch.int32)
        fl = torch.tensor([len(f) for f in feats], dtype=torch.int32)
        T = int(tl.max()); Fm = int(fl.max())
        tok = torch.zeros(len(bt), T, dtype=torch.long); feat = torch.zeros(len(bt), Fm, 80)
        for i, (t, f) in enumerate(zip(toks, feats)):
            tok[i, :len(t)] = t; feat[i, :len(f)] = f.cpu()
        return dict(speech_token=tok, speech_token_len=tl, speech_feat=feat, speech_feat_len=fl,
                    embedding=torch.stack(es).float())

    params = [p for p in flow.parameters() if p.requires_grad]
    opt = torch.optim.AdamW(params, lr=a.lr, weight_decay=0.0)
    logf = open(f"{a.out}/log.jsonl", "a")

    def validate():
        flow.eval(); L = []
        st = torch.random.get_rng_state(); cst = torch.cuda.get_rng_state(); pst = random.getstate()
        torch.manual_seed(0); random.seed(0)
        with torch.no_grad():
            for i in range(0, len(val), 16):
                L.append(float(flow(make(val[i:i + 16]), torch.device("cuda"))["loss"]))
        torch.random.set_rng_state(st); torch.cuda.set_rng_state(cst); random.setstate(pst)
        flow.train(); return float(np.mean(L))

    v = validate(); best, bad = v, 0
    print(json.dumps({"step": 0, "val": v}), flush=True); logf.write(json.dumps({"step": 0, "val": v}) + "\n")
    torch.save(dict(flow=flow.state_dict(), step=0, val=v), f"{a.out}/best.pt")
    rng = random.Random(0); step, t0, done = 0, time.time(), False
    while not done:
        for bi in batches(items, a.tokens, rng, extra, cost=lambda m: m["ntok"]):
            lrm = min(1.0, (step + 1) / a.warmup)
            for g in opt.param_groups:
                g["lr"] = a.lr * lrm
            with torch.autocast("cuda", dtype=torch.bfloat16):
                loss = flow(make([items[i] for i in bi]), torch.device("cuda"))["loss"]
            opt.zero_grad(set_to_none=True)
            if not torch.isfinite(loss):
                continue
            loss.backward()
            gn = torch.nn.utils.clip_grad_norm_(params, 5.0)
            if not torch.isfinite(gn):
                continue
            opt.step(); step += 1
            if step % 50 == 0:
                r = dict(step=step, loss=round(float(loss), 4), gn=round(float(gn), 2), B=len(bi), min=round((time.time() - t0) / 60, 2))
                print(json.dumps(r), flush=True); logf.write(json.dumps(r) + "\n"); logf.flush()
            if step % a.val_every == 0:
                v = validate(); imp = v < best - 1e-4
                print(json.dumps({"step": step, "val": v, "best": imp}), flush=True)
                logf.write(json.dumps({"step": step, "val": v, "best": imp}) + "\n"); logf.flush()
                if imp:
                    best, bad = v, 0
                    torch.save(dict(flow=flow.state_dict(), step=step, val=v), f"{a.out}/best.pt")
                else:
                    bad += 1
                    if bad >= a.patience:
                        done = True; break
            if (time.time() - t0) / 60 > a.minutes or step >= a.max_steps:
                done = True; break
    print("TRAIN_FLOW_DONE", step, "best", best, flush=True)


# ------------------------------------------------------------------ synthesis
_SENT = re.compile(r"(?<=[.?!])\s+")


def split_sentences(text):
    """tagged text -> list of tagged sentences. A sentence ends at . ? ! inside a span or at the end of a span;
    language spans are cut there and re-tagged, so a sentence keeps its in-sentence switches."""
    out, cur = [], []
    for l, s in segments(text):
        parts = _SENT.split(s.strip())
        for k, p in enumerate(parts):
            p = p.strip()
            if p:
                cur.append(f"<{l}>{p}</{l}>")
            if (k < len(parts) - 1 or p[-1:] in ".?!") and cur:
                out.append(" ".join(cur)); cur = []
    if cur:
        out.append(" ".join(cur))
    return [o for o in out if o.strip()]


def trim_edges(w, sr, keep_lead=0.12, keep_tail=0.18):
    from da.p2common import silence_mask
    sil = silence_mask(w, sr)
    sp = np.where(~sil)[0]
    if not len(sp):
        return w
    a, b = sp[0] * 0.01, (sp[-1] + 1) * 0.01
    out = w[int(max(0, a - keep_lead) * sr): int(min(len(w) / sr, b + keep_tail) * sr)]
    lead = a - max(0, a - keep_lead)
    if lead < keep_lead:
        out = np.concatenate([np.zeros(int((keep_lead - lead) * sr), out.dtype), out])
    return out


def cap_internal(w, sr, cap=0.6):
    """shorten every internal silence longer than cap seconds to cap (the centre of the silent run is removed)."""
    from da.p2common import silence_mask, runs
    sil = silence_mask(w, sr)
    hop = int(0.01 * sr)
    R = [(a, b) for a, b in runs(sil) if a > 0 and b < len(sil) and (b - a) * 0.01 > cap]
    if not R:
        return w
    keep, cur = [], 0
    for a, b in R:
        cut0 = a * hop + int(cap / 2 * sr)
        cut1 = b * hop - int(cap / 2 * sr)
        keep.append(w[cur:cut0]); cur = cut1
    keep.append(w[cur:])
    return np.concatenate(keep)


class Synth:
    def __init__(self, ckpt, flow_ckpt=None, sent_pause=0.5):
        from transformers import AutoTokenizer
        self.tok = AutoTokenizer.from_pretrained(f"{MODEL_DIR}/CosyVoice-BlankEN")
        self.cosy = load_cosy()
        if flow_ckpt:
            self.cosy.model.flow.load_state_dict(torch.load(flow_ckpt, map_location="cuda")["flow"])
        ck = torch.load(ckpt, map_location="cuda")
        self.spk2id = json.load(open(os.path.join(os.path.dirname(ckpt), "spk2id.json")))
        self.m = DALM2(self.cosy.model.llm, ck["n_spk"]).cuda().eval()
        self.m.load_state_dict(ck["model"])
        self.sent_pause = sent_pause
        self.cap = 0.6
        # anchor prompts: normalised anchor val clips from data2 (never in train), 3-7 s, highest SNR first
        meta = [json.loads(l) for l in open(f"{D2}/waxal/meta.jsonl", encoding="utf-8")]
        tk = [json.loads(l) for l in open(f"{D2}/waxal/s3tok.jsonl")]
        pcm = np.memmap(f"{D2}/waxal/pcm24k.i16", dtype=np.int16, mode="r")
        s3 = np.memmap(f"{D2}/waxal/s3tok.i32", dtype=np.int32, mode="r")
        cand = [(m, t) for m, t in zip(meta, tk) if m["spk"] == ANCHOR and m["split"] == "val" and 3.0 <= m["dur"] <= 7.0
                and not m.get("longform")]
        cand.sort(key=lambda x: -x[0]["ascore"])
        self.prompts = []
        os.makedirs(f"{HOME}/eval/prompts", exist_ok=True)
        for m, t in cand[:4]:
            y = np.asarray(pcm[m["off24"]: m["off24"] + m["n24"]], dtype=np.float32) / 32768.0
            p = f"{HOME}/eval/prompts/{m['id']}.wav"; sf.write(p, y, 24000)
            from train_b import prompt_feats
            ptok, pfeat, pemb = prompt_feats(self.cosy, p)
            ids, langs = text_tokens(self.tok, tagged(m))
            lmp = (torch.tensor(ids).cuda(), torch.tensor(langs).cuda(),
                   torch.from_numpy(np.asarray(s3[t["toff"]: t["toff"] + t["ntok"]], dtype=np.int64)).cuda())
            self.prompts.append(dict(id=m["id"], text=m["text"], lm=lmp, flow=(ptok, pfeat, pemb)))
        print("anchor prompts", [p["id"] for p in self.prompts], flush=True)

    def one(self, tagged_text, seed=0, prompt_k=0, use_lm_prompt=True):
        ids, langs = text_tokens(self.tok, tagged_text)
        P = self.prompts[prompt_k % len(self.prompts)]
        with torch.autocast("cuda", dtype=torch.bfloat16):
            sp = self.m.generate2(torch.tensor(ids).cuda(), torch.tensor(langs).cuda(),
                                  torch.tensor(self.spk2id[ANCHOR]).cuda(), prompt=P["lm"] if use_lm_prompt else None,
                                  seed=seed)
        if len(sp) < 3:
            sp = sp + [0] * 3
        self.last_tokens.append(sp)
        return token2wav(self.cosy, sp, *P["flow"])

    def join(self, ws):
        sr = self.cosy.sample_rate
        if len(ws) == 1:
            return ws[0]
        out = []
        for i, w in enumerate(ws):
            if i:
                out.append(np.zeros(max(0, int((self.sent_pause - 0.30) * sr)), np.float32))
            out.append(w)
        return np.concatenate(out)

    def redecode(self, toks, prompt_k=0):
        """same LM tokens, (possibly different) flow -> identical post-processing"""
        sr = self.cosy.sample_rate
        P = self.prompts[prompt_k % len(self.prompts)]
        return self.join([cap_internal(trim_edges(token2wav(self.cosy, t, *P["flow"]), sr), sr, self.cap) for t in toks])

    def say(self, tagged_text, seed=0, prompt_k=0, stitch=True, use_lm_prompt=True):
        sr = self.cosy.sample_rate
        self.last_tokens = []
        sents = split_sentences(tagged_text) if stitch else [tagged_text]
        ws = [cap_internal(trim_edges(self.one(s, seed=seed + 1000 * i, prompt_k=prompt_k, use_lm_prompt=use_lm_prompt), sr), sr, self.cap)
              for i, s in enumerate(sents)]
        # trailing 0.18 + leading 0.12 already present; join() tops up to the anchor's sentence pause
        return self.join(ws)


def redecode(a):
    """decode the LM tokens saved by synth (--src DIR) with another flow checkpoint -> --out"""
    S = Synth(a.ckpt, a.flow or None, a.sent_pause)
    os.makedirs(a.out, exist_ok=True)
    import glob
    for f in sorted(glob.glob(f"{a.src}/*.tok.json")):
        d = json.load(open(f))
        i = os.path.basename(f)[:-9]
        sf.write(f"{a.out}/{i}.wav", S.redecode(d["tokens"], d["prompt_k"]), S.cosy.sample_rate)
    print("REDECODE_DONE", a.out, flush=True)


def synth(a):
    S = Synth(a.ckpt, a.flow or None, a.sent_pause)
    rows = [json.loads(l) for l in open(a.eval_set, encoding="utf-8")]
    if a.only:
        keep = set(a.only.split(",")); rows = [r for r in rows if r["id"] in keep or r["lang"] in keep]
    if a.nper:
        cnt, sel = {}, []
        for r in rows:
            if r["lang"] in ("cs", "para") or cnt.get(r["lang"], 0) < a.nper:
                sel.append(r); cnt[r["lang"]] = cnt.get(r["lang"], 0) + 1
        rows = sel
    ks = [a.k] if a.k >= 0 else list(range(a.n))
    for k in ks:
        od = a.out if len(ks) == 1 else f"{a.out}/cand{k}"
        os.makedirs(od, exist_ok=True)
        for r in rows:
            text = r["text"] if r["lang"] in ("cs", "para") else f"<{r['lang']}>{r['text']}</{r['lang']}>"
            w = S.say(text, seed=k, prompt_k=k, stitch=not a.no_stitch, use_lm_prompt=not a.no_lm_prompt)
            sf.write(f"{od}/{r['id']}.wav", w, S.cosy.sample_rate)
            json.dump(dict(tokens=S.last_tokens, prompt_k=k), open(f"{od}/{r['id']}.tok.json", "w"))
        print("synth cand", k, "done", od, flush=True)
    print("SYNTH2_DONE", a.out, flush=True)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("cmd")
    ap.add_argument("--datasets", default="waxal,fleurs,slr86,slr70,bible_yor,bible_hau,cs_pairs")
    ap.add_argument("--out", default=f"{HOME}/runs/p2_lm")
    ap.add_argument("--init", default="")
    ap.add_argument("--flow_init", default="")
    ap.add_argument("--ckpt", default="")
    ap.add_argument("--flow", default="")
    ap.add_argument("--minutes", type=float, default=90)
    ap.add_argument("--max_steps", type=int, default=4000)
    ap.add_argument("--lr", type=float, default=1e-5)
    ap.add_argument("--lr_new", type=float, default=5e-4)
    ap.add_argument("--warmup", type=int, default=300)
    ap.add_argument("--tokens", type=int, default=16000)
    ap.add_argument("--adv_w", type=float, default=0.02)
    ap.add_argument("--p_prompt", type=float, default=0.5)
    ap.add_argument("--val_every", type=int, default=200)
    ap.add_argument("--patience", type=int, default=5)
    ap.add_argument("--anchor_only_w", type=float, default=0.0)
    ap.add_argument("--select_anchor", action="store_true")
    ap.add_argument("--eval_set", default=f"{HOME}/eval/eval_set2.jsonl")
    ap.add_argument("--n", type=int, default=1)
    ap.add_argument("--nper", type=int, default=0)
    ap.add_argument("--k", type=int, default=-1)
    ap.add_argument("--src", default="")
    ap.add_argument("--only", default="")
    ap.add_argument("--sent_pause", type=float, default=0.6)
    ap.add_argument("--no_stitch", action="store_true")
    ap.add_argument("--no_lm_prompt", action="store_true")
    a = ap.parse_args()
    {"lm": train_lm, "flow": train_flow, "synth": synth, "redecode": redecode}[a.cmd](a)