#!/usr/bin/env python3
"""Applique le restaurateur de ponctuation a un CSV de soumission (ID,Target).
Garde anti-regression : rejette toute sortie qui modifie les mots hors ponct/casse.
Langue deduite du prefixe de l'ID (lin_/lug_/sna_)."""
import argparse
import csv
import unicodedata

import torch
from transformers import AutoTokenizer, T5ForConditionalGeneration

MODEL = "/root/models/punct_best"


def strip_form(t):
    t = unicodedata.normalize("NFC", t).lower()
    t = "".join(c if c.isalnum() or c.isspace() or c in "'-" else " " for c in t)
    return " ".join(t.split())


@torch.inference_mode()
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--inp", required=True)
    ap.add_argument("--out", required=True)
    a = ap.parse_args()
    tok = AutoTokenizer.from_pretrained(MODEL)
    model = T5ForConditionalGeneration.from_pretrained(MODEL, torch_dtype=torch.bfloat16).cuda().eval()
    rows = list(csv.DictReader(open(a.inp, encoding="utf-8")))
    kept = rej = 0
    B = 48
    out_map = {}
    for k in range(0, len(rows), B):
        chunk = rows[k:k+B]
        srcs, langs = [], []
        for r in chunk:
            lang = r["ID"].split("_")[0]
            langs.append(lang)
            srcs.append(f"{lang}: {strip_form(r['Target'])}")
        enc = tok(srcs, return_tensors="pt", padding=True, truncation=True, max_length=512).to("cuda")
        gen = model.generate(**enc, max_length=512, num_beams=1)
        outs = tok.batch_decode(gen, skip_special_tokens=True)
        for r, o in zip(chunk, outs):
            o = " ".join(o.split())
            if o and strip_form(o) == strip_form(r["Target"]):
                out_map[r["ID"]] = o
                kept += 1
            else:
                out_map[r["ID"]] = r["Target"] or "a"
                rej += 1
    with open(a.out, "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["ID", "Target"])
        for r in rows:
            w.writerow([r["ID"], out_map[r["ID"]] or "a"])
    print(f"PUNCT_CSV_DONE {a.out}: {kept} restaures, {rej} gardes-original")


if __name__ == "__main__":
    main()