"""Fine-tune Meta omniASR CTC (default omniASR_CTC_1B_v2, Apache-2.0) on Nigerian speech (+ optional synthetic code-switch).
Full fine-tune (CNN feature extractor frozen), bf16 autocast, AdamW, warmup + linear decay, duration-bucketed batches.
The model stays a plain CTC wav2vec2 (same architecture/tokenizer), so the duplex asr_server streams it unchanged.
  ~/venv_omni/bin/python train_asr.py --run /scratch/asr/runs/A --steps 4000 [--synth /scratch/asr/synth/manifest.jsonl --synth_frac 0.2]
"""
import argparse, json, math, os, random, time, re, unicodedata
import numpy as np, torch, torch.nn.functional as F
import soundfile  # noqa: F401
from scipy.signal import resample_poly, fftconvolve
from fairseq2.models.hub import load_model
from fairseq2.data.tokenizers.hub import load_tokenizer
from fairseq2.nn.batch_layout import BatchLayout

ap = argparse.ArgumentParser()
ap.add_argument("--card", default="omniASR_CTC_1B_v2"); ap.add_argument("--run", required=True)
ap.add_argument("--data", default="/scratch/asr/data"); ap.add_argument("--synth", default="")
ap.add_argument("--synth_frac", type=float, default=0.2); ap.add_argument("--steps", type=int, default=4000)
ap.add_argument("--lr", type=float, default=2e-5); ap.add_argument("--warmup", type=int, default=300)
ap.add_argument("--batch_sec", type=float, default=280.0); ap.add_argument("--max_sec", type=float, default=20.0)
ap.add_argument("--real_aug_p", type=float, default=0.15); ap.add_argument("--synth_aug_p", type=float, default=0.8)
ap.add_argument("--val_every", type=int, default=500); ap.add_argument("--seed", type=int, default=1)
ap.add_argument("--workers", type=int, default=10); ap.add_argument("--smoke", action="store_true")
a = ap.parse_args()
os.makedirs(a.run, exist_ok=True); json.dump(vars(a), open(f"{a.run}/args.json", "w"), indent=1)
random.seed(a.seed); np.random.seed(a.seed); torch.manual_seed(a.seed)
dev = torch.device("cuda")
if os.environ.get("MEM_FRAC"): torch.cuda.set_per_process_memory_fraction(float(os.environ["MEM_FRAC"]))
model = load_model(a.card, device=dev, dtype=torch.float32)
tok = load_tokenizer(a.card); enc = tok.create_encoder(); dec = tok.create_decoder(skip_special_tokens=True)
VI = tok.vocab_info; BLANK = 0; PAD = VI.pad_idx if VI.pad_idx is not None else 1
print("vocab", VI.size, "pad", VI.pad_idx, "unk", VI.unk_idx, "bos", VI.bos_idx, "eos", VI.eos_idx, flush=True)
SPECIAL = {i for i in (VI.pad_idx, VI.bos_idx, VI.eos_idx) if i is not None}

def encode(t):
    ids = [int(x) for x in enc(t).tolist()]
    return [i for i in ids if i not in SPECIAL]

def rows(p):
    return [json.loads(l) for l in open(p)]
train = rows(f"{a.data}/train.jsonl"); devset = rows(f"{a.data}/dev.jsonl")
synth = rows(a.synth) if a.synth else []
for r in train: r["kind"] = "real"
for r in synth: r["kind"] = "synth"; r.setdefault("w", 1.0)
unk = VI.unk_idx
def ok(r):
    ids = encode(r["text"]); r["n_ids"] = len(ids)
    return len(ids) > 0 and (unk is None or unk not in ids) and r["dur"] <= a.max_sec
n0 = len(train); train = [r for r in train if ok(r)]; synth = [r for r in synth if ok(r)]
print(f"train {len(train)}/{n0} after tokenizer filter, synth {len(synth)}", flush=True)

_mm = {}
def load_audio(r):
    if r.get("wav"):
        y, sr = soundfile.read(r["wav"], dtype="float32")
        if y.ndim > 1: y = y.mean(1)
        if sr != 16000: y = resample_poly(y, 16000, sr).astype(np.float32)
        return y
    if r["pcm"] not in _mm: _mm[r["pcm"]] = np.memmap(r["pcm"], dtype=np.int16, mode="r")
    y = _mm[r["pcm"]][r["off24"]: r["off24"] + r["n24"]].astype(np.float32) / 32768
    return resample_poly(y, 2, 3).astype(np.float32)

def augment(y, rng):
    """noise (white/pink/babble-ish), simple synthetic reverb, phone band-limit, gain; ~ mic/channel variety"""
    if rng.random() < 0.5:   # reverb: exponentially decaying noise RIR, RT60 0.15-0.7 s
        rt = rng.uniform(0.15, 0.7); n = int(16000 * rt); t = np.arange(n) / 16000
        rir = rng.standard_normal(n) * np.exp(-6.9 * t / rt); rir[0] = 1.0; rir /= np.sqrt((rir ** 2).sum())
        y = fftconvolve(y, rir)[: len(y)].astype(np.float32)
    if rng.random() < 0.35:  # telephone / cheap mic channel
        y = resample_poly(resample_poly(y, 1, 2), 2, 1)[: len(y)].astype(np.float32) if rng.random() < 0.5 else \
            resample_poly(resample_poly(y, 8, 16), 16, 8)[: len(y)].astype(np.float32)
        if rng.random() < 0.5: y = np.tanh(y * rng.uniform(1.5, 4.0)).astype(np.float32)
    if rng.random() < 0.7:   # additive noise at SNR 5-30 dB
        n = rng.standard_normal(len(y)).astype(np.float32)
        if rng.random() < 0.5: n = np.cumsum(n); n -= np.convolve(n, np.ones(400) / 400, "same")   # brownish
        snr = rng.uniform(5, 30); ps = (y ** 2).mean() + 1e-8; pn = (n ** 2).mean() + 1e-8
        y = y + n * np.sqrt(ps / (pn * 10 ** (snr / 10)))
    y = y * rng.uniform(0.3, 1.5)
    return np.clip(y, -1, 1).astype(np.float32)

class DS(torch.utils.data.Dataset):
    def __init__(self, items, train_mode): self.items, self.tm = items, train_mode
    def __len__(self): return len(self.items)
    def __getitem__(self, i):
        r = self.items[i]; y = load_audio(r)
        if self.tm:
            rng = np.random.default_rng((i * 7919 + int(time.time() * 1000)) % 2 ** 32)
            p = a.synth_aug_p if r["kind"] == "synth" else a.real_aug_p
            if rng.random() < p: y = augment(y, rng)
        return torch.from_numpy(y), torch.tensor(encode(r["text"])), i

def epoch_order(rng):
    """weighted sampling (source weights) + synth mixed at synth_frac of audio hours; bucketed by duration"""
    pick = []
    for r in train:
        k = r["w"]; n = int(k) + (1 if rng.random() < k - int(k) else 0); pick += [r] * n
    if synth:
        real_h = sum(r["dur"] for r in pick); want = real_h * a.synth_frac / (1 - a.synth_frac); have = sum(r["dur"] for r in synth)
        rep = want / have; s = []
        for r in synth:
            n = int(rep) + (1 if rng.random() < rep - int(rep) else 0); s += [r] * n
        pick += s
    rng.shuffle(pick)
    batches = []
    for c in range(0, len(pick), 4000):   # sort within chunks of 4000 to bucket by duration
        ch = sorted(pick[c: c + 4000], key=lambda r: r["dur"]); cur, cs = [], 0.0
        for r in ch:
            if cs + r["dur"] > a.batch_sec and cur: batches.append(cur); cur, cs = [], 0.0
            cur.append(r); cs += r["dur"]
        if cur: batches.append(cur)
    rng.shuffle(batches); return batches

def collate(b):
    ys = [F.layer_norm(x, x.shape) for x, _, _ in b]; L = torch.tensor([len(x) for x in ys])
    X = torch.zeros(len(ys), int(L.max())); [X[i, : len(x)].copy_(x) for i, x in enumerate(ys)]
    T = [t for _, t, _ in b]; TL = torch.tensor([len(t) for t in T]); Y = torch.cat(T)
    return X, L, Y, TL

def forward(X, L):
    bl = BatchLayout(X.shape, seq_lens=L.tolist(), device=dev)
    logits, blo = model(X, bl)
    return logits, torch.tensor(blo.seq_lens, device=dev)

for p in model.encoder_frontend.feature_extractor.parameters(): p.requires_grad = False
params = [p for p in model.parameters() if p.requires_grad]
print("trainable params %.1fM" % (sum(p.numel() for p in params) / 1e6), flush=True)
opt = torch.optim.AdamW(params, lr=a.lr, betas=(0.9, 0.98), eps=1e-8, weight_decay=0.01)
sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / a.warmup) * max(0.02, 1 - max(0, s - a.warmup) / max(1, a.steps - a.warmup)))

def cer_eval(items, bs=16):
    model.eval(); hyps = []
    with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
        for i in range(0, len(items), bs):
            b = [(torch.from_numpy(load_audio(r)), torch.tensor([1]), 0) for r in items[i: i + bs]]
            X, L, _, _ = collate(b); lg, ol = forward(X.to(dev), L)
            pr = lg.argmax(-1)
            for j in range(pr.shape[0]):
                s = pr[j, : ol[j]]; m = torch.ones_like(s, dtype=torch.bool); m[1:] = s[1:] != s[:-1]; s = s[m]; s = s[s != BLANK]
                hyps.append(str(dec(s)))
    model.train()
    import jiwer
    def n(x): x = unicodedata.normalize("NFC", x.lower()); x = re.sub(r"[^\w\s']", " ", x); return re.sub(r"\s+", " ", x).strip() or "x"
    out = {}
    for lg in sorted({r["lang"] for r in items}):
        idx = [k for k, r in enumerate(items) if r["lang"] == lg]
        out[lg] = round(jiwer.cer([n(items[k]["text"]) for k in idx], [n(hyps[k]) for k in idx]), 4)
    out["mean"] = round(float(np.mean([v for k, v in out.items()])), 4)
    return out, hyps

log = open(f"{a.run}/log.jsonl", "a")
v0, _ = cer_eval(devset); print("dev step0", v0, flush=True); log.write(json.dumps(dict(step=0, dev=v0)) + "\n"); log.flush()
best = v0["mean"]; step = 0; rng = random.Random(a.seed); model.train(); t0 = time.time(); ema = None
steps_total = 60 if a.smoke else a.steps
while step < steps_total:
    batches = epoch_order(rng); print(f"epoch with {len(batches)} batches", flush=True)
    lut = {}
    flat = []
    for b in batches:
        flat.append([len(lut) + k for k in range(len(b))]); [lut.setdefault(len(lut), r) for r in b]
    items = [lut[k] for k in range(len(lut))]
    dl = torch.utils.data.DataLoader(DS(items, True), batch_sampler=flat, num_workers=a.workers, collate_fn=collate,
                                     prefetch_factor=4, persistent_workers=False)
    for X, L, Y, TL in dl:
        X, Y = X.to(dev, non_blocking=True), Y.to(dev)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            lg, ol = forward(X, L)
        lp = F.log_softmax(lg.float(), -1).transpose(0, 1)
        loss = F.ctc_loss(lp, Y, ol, TL.to(dev), blank=BLANK, reduction="sum", zero_infinity=True) / TL.sum()
        opt.zero_grad(set_to_none=True); loss.backward()
        gn = torch.nn.utils.clip_grad_norm_(params, 5.0)
        if torch.isfinite(gn): opt.step()
        sched.step(); step += 1
        ema = loss.item() if ema is None else 0.98 * ema + 0.02 * loss.item()
        if step % 25 == 0:
            el = time.time() - t0
            print(f"step {step} loss {ema:.4f} gn {gn:.2f} lr {sched.get_last_lr()[0]:.2e} {el / step:.2f}s/step mem {torch.cuda.max_memory_allocated() / 2**30:.1f}G", flush=True)
            log.write(json.dumps(dict(step=step, loss=round(ema, 4), sps=round(el / step, 3))) + "\n"); log.flush()
        if step % a.val_every == 0 or step == steps_total:
            v, _ = cer_eval(devset); print("dev", step, v, flush=True); log.write(json.dumps(dict(step=step, dev=v)) + "\n"); log.flush()
            sd = {k: t.to(torch.bfloat16) for k, t in model.state_dict().items()}
            if v["mean"] < best and not a.smoke:
                best = v["mean"]; torch.save(dict(model=sd, step=step, dev=v, card=a.card), f"{a.run}/best.pt"); print("saved best", flush=True)
        if step >= steps_total: break
if not a.smoke:
    torch.save(dict(model={k: t.to(torch.bfloat16) for k, t in model.state_dict().items()}, step=step, card=a.card), f"{a.run}/last.pt")
print("TRAIN_DONE best", best, flush=True)