#!/usr/bin/env python3
"""PARI STEP-CHANGE : fine-tune Whisper-large-v3 (seq2seq 1.5B) sur le LINGALA WAXAL (70 h).
Motivation : notre champion est un CTC w2v-BERT 2.0 (580M) ; sur benchmark africain comparable
le seq2seq Whisper fait -16% relatif de WER vs w2v-BERT CTC. Le lingala (combine 0.346) est LE
goulot ; le shona (0.125) est déjà saturé. Whisper connaît nativement 'ln' et produit
casse+ponctuation (que le scoreur WAXAL compte).
Conforme : entraînement sur le SEUL texte/audio du train WAXAL, fine-tuning autorisé (fil 34204).

Gate : devhard-lin (439 clips, locuteurs held-out) — fiable pour les GROS effets.
"""
import json, os, sys
import numpy as np, soundfile as sf, torch
from dataclasses import dataclass
from torch.utils.data import Dataset
from transformers import (WhisperForConditionalGeneration, WhisperProcessor,
                          Seq2SeqTrainer, Seq2SeqTrainingArguments)

BASE = os.environ.get("BASE_MODEL", "openai/whisper-large-v3")
LANG = os.environ.get("LANG_ASR", "lin")
WLANG = {"lin": "ln", "sna": "sn"}[LANG]
OUT = os.environ.get("OUT", "/scratch/runs/whisper_%s" % LANG)
EPOCHS = float(os.environ.get("EPOCHS", "3"))
BS = int(os.environ.get("BS", "4"))
ACC = int(os.environ.get("ACC", "4"))
LR = float(os.environ.get("LR", "1e-5"))
MAXDUR = float(os.environ.get("MAXDUR", "30"))
SR = 16000


def load_manifest(path):
    rows = []
    for l in open(path, encoding="utf-8"):
        r = json.loads(l)
        if not r.get("text", "").strip():
            continue
        a = r["audio"]
        if not os.path.exists(a):                     # remap si le chemin a bougé
            for cand in ("/scratch/prep/audio/", "/scratch/restore/devhard_audio/",
                         "/root/devhard_audio/"):
                p = os.path.join(cand, os.path.basename(a))
                if os.path.exists(p):
                    r["audio"] = p; break
        if os.path.exists(r["audio"]) and r.get("duration", 0) <= MAXDUR:
            rows.append(r)
    return rows


class DS(Dataset):
    def __init__(self, rows, proc):
        self.rows = rows; self.proc = proc

    def __len__(self):
        return len(self.rows)

    def __getitem__(self, i):
        r = self.rows[i]
        au = sf.read(r["audio"], dtype="float32")[0]
        if au.ndim > 1:
            au = au.mean(1)
        feats = self.proc.feature_extractor(au, sampling_rate=SR).input_features[0]
        ids = self.proc.tokenizer(r["text"], max_length=200, truncation=True).input_ids
        return {"input_features": feats, "labels": ids}


@dataclass
class Collate:
    proc: object

    def __call__(self, batch):
        feats = torch.tensor(np.stack([b["input_features"] for b in batch]))
        lab = self.proc.tokenizer.pad([{"input_ids": b["labels"]} for b in batch],
                                      return_tensors="pt")
        labels = lab["input_ids"].masked_fill(lab.attention_mask.ne(1), -100)
        if (labels[:, 0] == self.proc.tokenizer.convert_tokens_to_ids("<|startoftranscript|>")).all():
            labels = labels[:, 1:]
        return {"input_features": feats, "labels": labels}


def main():
    proc = WhisperProcessor.from_pretrained(BASE, language=WLANG, task="transcribe")
    train = load_manifest("/root/devhard/train_%s_min.jsonl" % LANG)
    dev = load_manifest("/root/devhard/devhard_%s.jsonl" % LANG)[:120]   # eval rapide
    print("train %d clips (%.1f h) | eval %d" % (
        len(train), sum(r.get("duration", 0) for r in train) / 3600, len(dev)), flush=True)
    if not train:
        print("AUCUNE DONNEE — prep_data non termine ?", flush=True); sys.exit(1)

    model = WhisperForConditionalGeneration.from_pretrained(BASE, dtype=torch.bfloat16)
    model.config.forced_decoder_ids = None
    model.config.suppress_tokens = []
    model.generation_config.language = WLANG
    model.generation_config.task = "transcribe"
    model.generation_config.forced_decoder_ids = None
    model.config.use_cache = False

    args = Seq2SeqTrainingArguments(
        output_dir=OUT, per_device_train_batch_size=BS, gradient_accumulation_steps=ACC,
        learning_rate=LR, warmup_steps=300, num_train_epochs=EPOCHS,
        gradient_checkpointing=True, bf16=True, optim="adamw_torch_fused",
        logging_steps=50, save_strategy="epoch", save_total_limit=2,
        eval_strategy="no", report_to=[], dataloader_num_workers=4,
        remove_unused_columns=False, lr_scheduler_type="linear")
    tr = Seq2SeqTrainer(model=model, args=args, train_dataset=DS(train, proc),
                        data_collator=Collate(proc))
    tr.train()
    tr.save_model(OUT + "/final")
    proc.save_pretrained(OUT + "/final")
    print("WHISPER_TRAIN_DONE %s/final" % OUT, flush=True)


if __name__ == "__main__":
    main()