#!/usr/bin/env python3
"""Génère le CSV de soumission Phase 1 avec le set gagnant par langue (modèles mixtes
MMS + w2v-BERT via AutoModelForCTC) + fill-empty (short -> top-4gram du train)."""
import csv
import json
import os
from collections import Counter

import numpy as np
import soundfile as sf
import torch
from transformers import AutoModelForCTC, AutoProcessor

SR = 16000
MANIF = "/scratch/prep/manifests"
# set gagnant : lin=MMS-1B (GO), sna=w2v-BERT champion, lug=w2v-BERT champion
MODELS = {
    "lin": "/root/models/mms1b_lin_best",
    "sna": "/root/models/sna_ps_best",
    "lug": "/root/models/lug_ps_best",
}
OUT = "/root/sub_mms1b_lin.csv"


def norm(s):
    return " ".join(str(s).replace("|", " ").split())


def top_gram(lang):
    g = Counter()
    for l in open(f"{MANIF}/waxal_{lang}_train.jsonl", encoding="utf-8"):
        w = norm(json.loads(l).get("text", "")).lower().split()
        for i in range(len(w) - 3):
            g[" ".join(w[i:i + 4])] += 1
    return g.most_common(1)[0][0] if g else "a"


def main():
    test_ids = [r["ID"] for r in csv.DictReader(open("/root/Test.csv", encoding="utf-8"))]
    priors = {l: top_gram(l) for l in MODELS}
    preds = {}
    for lang, path in MODELS.items():
        tman = {}
        for l in open(f"{MANIF}/waxal_{lang}_test.jsonl", encoding="utf-8"):
            r = json.loads(l)
            tman[r["id"]] = r
        tids = [t for t in test_ids if t.startswith(lang + "_")]
        rows = [tman[t] for t in tids if t in tman]
        proc = AutoProcessor.from_pretrained(path)
        m = AutoModelForCTC.from_pretrained(path, torch_dtype=torch.bfloat16).cuda().eval()
        rows.sort(key=lambda r: -r["duration"])
        bs, cur, bud = [], [], 0.0
        for r in rows:
            if cur and bud + r["duration"] > 140:
                bs.append(cur); cur, bud = [], 0.0
            cur.append(r); bud += r["duration"]
        if cur:
            bs.append(cur)
        out = {}
        with torch.inference_mode():
            for b in bs:
                au = [sf.read(r["audio"], dtype="float32")[0] for r in b]
                f = proc(au, sampling_rate=SR, return_tensors="pt", padding=True)
                f = {k: v.to("cuda", dtype=torch.bfloat16 if v.dtype == torch.float32 else v.dtype)
                     for k, v in f.items()}
                ids = m(**f).logits.float().argmax(-1).cpu().numpy()
                for r, s in zip(b, proc.batch_decode(ids)):
                    out[r["id"]] = norm(s)
        missing = 0
        for t in tids:
            h = out.get(t, "")
            if len(h.split()) > 2:
                preds[t] = h
            else:
                preds[t] = priors[lang]; missing += 1
        del m
        torch.cuda.empty_cache()
        print(f"{lang}: {len(tids)} clips ({path.split('/')[-1]}) | {missing} fill-empty", flush=True)

    with open(OUT, "w", newline="", encoding="utf-8") as fo:
        w = csv.writer(fo)
        w.writerow(["ID", "Target"])
        for t in test_ids:
            w.writerow([t, preds.get(t) or "a"])
    print(f"SUB_DONE {OUT} ({len(test_ids)} lignes)", flush=True)


if __name__ == "__main__":
    main()