"""Arm B: CosyVoice2-0.5B (Apache-2.0) fine-tune with the Zhang et al. 2019 recipe added to its text->speech-token LM.

Additions (new, ours):
  * per-token language embedding added to the Qwen text-token embeddings (zero-init)
  * learned speaker embedding inserted as one token AFTER the text, before the task token (zero-init), so text
    positions cannot see it; speech positions attend to it (anchor-voice conditioning)
  * domain-adversarial speaker classifier with gradient reversal on the LM hidden states at text positions
Flow-matching decoder + HiFT vocoder stay frozen; at inference they are prompted with an anchor-voice clip.

  python train_b.py train --out ~/runs/b1 --minutes 30
  python train_b.py synth --ckpt ~/runs/b1/last.pt --out ~/eval/sys/armB_ft
  python train_b.py zeroshot --out ~/eval/sys/cosyvoice2_zeroshot      (pretrained model, cross-lingual, no fine-tune)
"""
import argparse, inspect, json, math, os, random, 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"]
MODEL_DIR = f"{HOME}/models/CosyVoice2-0.5B"
LANGS = ["eng", "ibo", "hau", "yor", "pcm"]
L2I = {l: i for i, l in enumerate(LANGS)}
sys.path.insert(0, f"{HOME}/da")
from da.text import segments  # noqa: E402

from da.model import grad_reverse  # noqa: E402


def _sf_load_wav(wav, target_sr, min_sr=16000):
    """soundfile replacement for cosyvoice load_wav (torchaudio>=2.9 load needs torchcodec)."""
    y, sr = sf.read(wav, dtype="float32", always_2d=True)
    y = y.mean(1)
    if sr != target_sr:
        y = librosa.resample(y, orig_sr=sr, target_sr=target_sr)
    return torch.from_numpy(y)[None]


def load_cosy():
    import cosyvoice.utils.file_utils as fu
    fu.load_wav = _sf_load_wav
    import cosyvoice.cli.frontend as cf
    cf.load_wav = _sf_load_wav
    from cosyvoice.cli.cosyvoice import CosyVoice2
    kw = dict(load_jit=False, load_trt=False, fp16=False)
    sig = inspect.signature(CosyVoice2.__init__).parameters
    kw = {k: v for k, v in kw.items() if k in sig}
    return CosyVoice2(MODEL_DIR, **kw)


def text_tokens(tok, text, default="eng"):
    """tagged text -> (qwen ids, per-token lang ids)."""
    ids, langs = [], []
    for l, s in segments(text, default):
        t = tok(s, add_special_tokens=False)["input_ids"]
        ids += t; langs += [L2I[l]] * len(t)
    return ids, langs


class DALM(nn.Module):
    def __init__(self, qlm, n_spk, d=896):
        super().__init__()
        self.q = qlm                                     # cosyvoice Qwen2LM
        self.lang_emb = nn.Embedding(len(LANGS), d); nn.init.zeros_(self.lang_emb.weight)
        self.spk_emb = nn.Embedding(n_spk, d); nn.init.normal_(self.spk_emb.weight, std=0.02)
        self.adv = nn.Sequential(nn.Linear(d, 256), nn.ReLU(), nn.Linear(256, n_spk))
        self.V = qlm.speech_token_size

    def embed_text(self, ids, langs):
        return self.q.llm.model.model.embed_tokens(ids) + self.lang_emb(langs)

    def build(self, ids, langs, spk, speech=None):
        """one sequence: [sos, text+lang, spk, task, speech...]"""
        sos = self.q.llm_embedding.weight[self.q.sos][None]
        task = self.q.llm_embedding.weight[self.q.task_id][None]
        parts = [sos, self.embed_text(ids, langs), self.spk_emb(spk)[None], task]
        if speech is not None:
            parts.append(self.q.speech_embedding(speech))
        return torch.cat(parts, 0)

    def hidden(self, x, mask):
        out = self.q.llm.model.model(inputs_embeds=x, attention_mask=mask, return_dict=True)
        return out.last_hidden_state

    def forward(self, batch, adv_lambda=1.0):
        seqs, tgts, tpos = [], [], []
        for ids, langs, spk, sp in batch:
            x = self.build(ids, langs, spk, sp)
            n_text = len(ids)
            tgt = torch.full((x.shape[0],), -100, device=x.device, dtype=torch.long)
            # position of task token = 1 + n_text + 1; it predicts speech[0]; last speech pos predicts EOS
            start = 1 + n_text + 1
            tgt[start: start + len(sp)] = sp
            tgt[start + len(sp)] = self.V                         # EOS
            seqs.append(x); tgts.append(tgt); tpos.append((1, 1 + n_text))
        T = max(s.shape[0] for s in seqs) + 1
        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[:len(s)] if len(t) > len(s) else t; M[i, :len(s)] = 1
        # targets were laid out on positions; shift is implicit (position p predicts tgt[p])
        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()
        # adversarial speaker classifier on text positions
        hs, ys = [], []
        for i, (a, b) in enumerate(tpos):
            hs.append(H[i, a:b]); ys.append(torch.full((b - a,), int(batch[i][2]), 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)
        adv_acc = (al.argmax(-1) == ys).float().mean()
        return ce, acc, adv, adv_acc

    @torch.no_grad()
    def generate(self, ids, langs, spk, top_p=0.8, top_k=25, win=10, tau_r=0.1, seed=0):
        g = torch.Generator(device=ids.device).manual_seed(seed)
        x = self.build(ids, langs, spk)[None]
        n = len(ids)
        min_len, max_len = int(n * 2), int(n * 20)
        past, out = None, []
        inp = 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
            cand_p, cand_i = sp[keep], si[keep]
            tokid = int(cand_i[torch.multinomial(cand_p / cand_p.sum(), 1, generator=g)])
            if out and sum(1 for t in out[-win:] if t == tokid) >= win * tau_r and tokid != self.V:  # RAS
                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(root, datasets, split, max_tok=500):
    items, mm = [], {}
    for ds in datasets:
        if not os.path.exists(f"{root}/{ds}/s3tok.jsonl"):
            continue
        mm[ds] = np.memmap(f"{root}/{ds}/s3tok.i32", dtype=np.int32, mode="r")
        tk = {}
        for l in open(f"{root}/{ds}/s3tok.jsonl"):
            r = json.loads(l); tk[r["id"]] = (r["toff"], r["ntok"])
        for l in open(f"{root}/{ds}/meta.jsonl", encoding="utf-8"):
            m = json.loads(l)
            if m["split"] != split or m["id"] not in tk:
                continue
            m["toff"], m["ntok"] = tk[m["id"]]
            if 5 < m["ntok"] <= max_tok:
                items.append(m)
    return items, mm


def batches(items, budget, weights, rng):
    idx = []
    for i, m in enumerate(items):
        w = weights.get(m["ds"], 1.0)
        idx += [i] * int(w) + ([i] if rng.random() < w - int(w) else [])
    idx.sort(key=lambda i: items[i]["ntok"] + rng.random() * 10)
    out, cur, mx = [], [], 0
    for i in idx:
        L = items[i]["ntok"] + len(items[i]["text"]) // 2 + 4
        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


def train(a):
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained(f"{MODEL_DIR}/CosyVoice-BlankEN")
    cosy = load_cosy()
    items, mm = load_items(a.data, a.datasets.split(","), "train")
    val, _ = load_items(a.data, a.datasets.split(","), "val")
    spks = sorted({m["spk"] for m in items}); spk2id = {s: i for i, s in enumerate(spks)}
    val = [m for m in val if m["spk"] in spk2id][:200]
    os.makedirs(a.out, exist_ok=True)
    json.dump(spk2id, open(f"{a.out}/spk2id.json", "w"))
    print("armB items", len(items), "hours", round(sum(m["dur"] for m in items) / 3600, 1), "spk", len(spks), flush=True)
    model = DALM(cosy.model.llm, len(spks)).cuda().train()
    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 enc(m):
        if m["id"] not in cache:
            cache[m["id"]] = text_tokens(tok, f"<{m['lang']}>{m['text']}</{m['lang']}>")
        ids, langs = cache[m["id"]]
        sp = np.asarray(mm[m["ds"]][m["toff"]: m["toff"] + m["ntok"]], dtype=np.int64)
        return (torch.tensor(ids).cuda(), torch.tensor(langs).cuda(), torch.tensor(spk2id[m["spk"]]).cuda(),
                torch.from_numpy(sp).cuda())

    rng = random.Random(0); step = 0; nskip = 0; t0 = time.time(); logf = open(f"{a.out}/log.jsonl", "a")

    def validate():
        model.eval(); L = []
        with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
            for i in range(0, len(val), 16):
                ce, acc, _, _ = model([enc(m) for m in val[i:i + 16]]); L.append(float(ce))
        model.train(); return float(np.mean(L))

    v0 = validate(); print(json.dumps({"step": 0, "val_ce": v0}), flush=True); logf.write(json.dumps({"step": 0, "val_ce": v0}) + "\n")
    done = False
    while not done:
        for bi in batches(items, a.tokens, {"waxal": a.waxal_w}, rng):
            lrm = min(1.0, (step + 1) / 200)
            for g, b in zip(opt.param_groups, base):
                g["lr"] = b * lrm
            with torch.autocast("cuda", dtype=torch.bfloat16):
                ce, acc, adv, adv_acc = model([enc(items[i]) for i in bi])
            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), "B": len(bi),
                                  "maxtok": max(items[i]["ntok"] for i in bi)}), 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), min=round((time.time() - t0) / 60, 2))
                print(json.dumps(r), flush=True); logf.write(json.dumps(r) + "\n"); logf.flush()
            if step % 500 == 0:
                v = validate(); print(json.dumps({"step": step, "val_ce": v}), flush=True)
                logf.write(json.dumps({"step": step, "val_ce": v}) + "\n"); logf.flush()
            if (time.time() - t0) / 60 > a.minutes:
                done = True; break
    v = validate(); logf.write(json.dumps({"step": step, "val_ce": v, "final": True}) + "\n")
    sd = {k: v for k, v in model.state_dict().items()}
    torch.save(dict(model=sd, step=step, n_spk=len(spks)), f"{a.out}/last.pt")
    print("TRAIN_B_DONE", step, v, flush=True)


# ------------------------------------------------------------------ inference
def prompt_feats(cosy, wav_path):
    """flow prompt (speech tokens, 24k mel feat, campplus embedding) from an anchor clip."""
    fe = cosy.frontend
    tokn, tokl = fe._extract_speech_token(wav_path)
    feat, featl = fe._extract_speech_feat(wav_path)
    # CosyVoice2 aligns: feat frames = 2 * tokens
    n = min(feat.shape[1] // 2, tokn.shape[1])
    feat, tokn = feat[:, :2 * n], tokn[:, :n]
    emb = fe._extract_spk_embedding(wav_path)
    return tokn.cuda(), feat.cuda(), emb.cuda()


@torch.no_grad()
def token2wav(cosy, tokens, ptok, pfeat, emb):
    flow = cosy.model.flow
    t = torch.tensor([tokens], dtype=torch.int32).cuda()
    kw = dict(token=t, token_len=torch.tensor([t.shape[1]], dtype=torch.int32).cuda(),
              prompt_token=ptok, prompt_token_len=torch.tensor([ptok.shape[1]], dtype=torch.int32).cuda(),
              prompt_feat=pfeat, prompt_feat_len=torch.tensor([pfeat.shape[1]], dtype=torch.int32).cuda(),
              embedding=emb)
    sig = inspect.signature(flow.inference).parameters
    if "streaming" in sig:
        kw["streaming"] = False
    if "finalize" in sig:
        kw["finalize"] = True
    if "flow_cache" in sig:
        kw["flow_cache"] = torch.zeros(1, 80, 0, 2).cuda()
    mel = flow.inference(**kw)
    mel = mel[0] if isinstance(mel, tuple) else mel
    hs = inspect.signature(cosy.model.hift.inference).parameters
    w = cosy.model.hift.inference(speech_feat=mel, **({"cache_source": torch.zeros(1, 1, 0).cuda()} if "cache_source" in hs else {}))
    w = w[0] if isinstance(w, tuple) else w
    return w.squeeze().float().cpu().numpy()


def eval_rows():
    return [json.loads(l) for l in open(f"{HOME}/eval/eval_set.jsonl", encoding="utf-8")]


def synth(a):
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained(f"{MODEL_DIR}/CosyVoice-BlankEN")
    cosy = load_cosy()
    ck = torch.load(a.ckpt, map_location="cuda")
    spk2id = json.load(open(os.path.join(os.path.dirname(a.ckpt), "spk2id.json")))
    model = DALM(cosy.model.llm, ck["n_spk"]).cuda().eval()
    model.load_state_dict(ck["model"])
    anc = json.load(open(f"{HOME}/eval/anchor.json"))
    ptok, pfeat, emb = prompt_feats(cosy, anc["ref_wavs"][0])
    os.makedirs(a.out, exist_ok=True)
    for r in eval_rows():
        text = r["text"] if r["lang"] == "cs" else f"<{r['lang']}>{r['text']}</{r['lang']}>"
        ids, langs = text_tokens(tok, text)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            sp = model.generate(torch.tensor(ids).cuda(), torch.tensor(langs).cuda(), torch.tensor(spk2id[anc["spk"]]).cuda())
        sf.write(f"{a.out}/{r['id']}.wav", token2wav(cosy, sp, ptok, pfeat, emb), cosy.sample_rate)
    print("synth_b done", a.out, flush=True)


def zeroshot(a):
    """Unmodified pretrained CosyVoice2, cross-lingual mode with the anchor prompt (reference point)."""
    cosy = load_cosy()
    anc = json.load(open(f"{HOME}/eval/anchor.json"))
    from cosyvoice.utils.file_utils import load_wav
    pw = anc["ref_wavs"][0]
    sig = inspect.signature(cosy.inference_cross_lingual).parameters
    os.makedirs(a.out, exist_ok=True)
    for r in eval_rows():
        text = " ".join(s.strip() for _, s in segments(r["text"])) if r["lang"] == "cs" else r["text"]
        prompt = pw
        outs = [o["tts_speech"] for o in cosy.inference_cross_lingual(text, prompt, stream=False)]
        w = torch.cat(outs, 1).squeeze().cpu().numpy()
        sf.write(f"{a.out}/{r['id']}.wav", w, cosy.sample_rate)
    print("zeroshot done", flush=True)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("cmd")
    ap.add_argument("--data", default=f"{HOME}/data")
    ap.add_argument("--datasets", default="waxal,fleurs,slr86,slr70,bible_yor")
    ap.add_argument("--out", default=f"{HOME}/runs/b1")
    ap.add_argument("--ckpt", default="")
    ap.add_argument("--minutes", type=float, default=30)
    ap.add_argument("--lr", type=float, default=1e-5)
    ap.add_argument("--lr_new", type=float, default=1e-3)
    ap.add_argument("--tokens", type=int, default=12000)
    ap.add_argument("--waxal_w", type=float, default=2.0)
    ap.add_argument("--adv_w", type=float, default=0.02)
    a = ap.parse_args()
    {"train": train, "synth": synth, "zeroshot": zeroshot}[a.cmd](a)