#!/usr/bin/env python3
"""Diagnostic routage Phase 2 : LID vs confiance-CTC.
Sur un échantillon, compare la langue prédite par le LID et par la confiance CTC
(3 modèles w2v-BERT comparables : lin_s4, sna_ps, lug_ps). La confiance = moyenne
du max-softmax par trame (peakiness) — le bon modèle de langue est plus peaky/confiant."""
import glob
import json

import numpy as np
import soundfile as sf
import torch
from transformers import (AutoModelForCTC, AutoProcessor,
                          Wav2Vec2BertForSequenceClassification, SeamlessM4TFeatureExtractor)

SR = 16000
N = 120
CTC = {"lin": "/root/models/lin_s4_best", "sna": "/root/models/sna_ps_best", "lug": "/root/models/lug_ps_best"}
files = sorted(glob.glob("/root/phase2_audio/audio/*.wav"))[:N]

# --- LID ---
fe = SeamlessM4TFeatureExtractor.from_pretrained("/root/models/lid_best")
lid = Wav2Vec2BertForSequenceClassification.from_pretrained("/root/models/lid_best", torch_dtype=torch.bfloat16).cuda().eval()
id2label = json.load(open("/root/models/lid_best/config.json"))["id2label"]
lid_pred = {}
with torch.inference_mode():
    for f in files:
        w = sf.read(f, dtype="float32")[0][:20 * SR]
        x = fe([w], sampling_rate=SR, return_tensors="pt", padding=True)
        x = {k: v.to("cuda", dtype=torch.bfloat16 if v.dtype == torch.float32 else v.dtype) for k, v in x.items()}
        lid_pred[f] = id2label[str(int(lid(**x).logits.float().argmax(-1)))]
del lid; torch.cuda.empty_cache()

# --- confiance CTC par modèle ---
conf = {f: {} for f in files}
for lang, path in CTC.items():
    proc = AutoProcessor.from_pretrained(path)
    m = AutoModelForCTC.from_pretrained(path, torch_dtype=torch.bfloat16).cuda().eval()
    with torch.inference_mode():
        for f in files:
            w = sf.read(f, dtype="float32")[0]
            x = proc(w, sampling_rate=SR, return_tensors="pt", padding=True)
            x = {k: v.to("cuda", dtype=torch.bfloat16 if v.dtype == torch.float32 else v.dtype) for k, v in x.items()}
            probs = m(**x).logits.float().softmax(-1)[0]
            conf[f][lang] = float(probs.max(-1).values.mean())  # peakiness moyenne
    del m; torch.cuda.empty_cache()

ctc_pred = {f: max(conf[f], key=conf[f].get) for f in files}

from collections import Counter
lid_dist = Counter(lid_pred.values())
ctc_dist = Counter(ctc_pred.values())
agree = sum(1 for f in files if lid_pred[f] == ctc_pred[f])
print(f"N={len(files)}")
print(f"LID  distribution: {dict(lid_dist)}")
print(f"CTC  distribution: {dict(ctc_dist)}")
print(f"accord LID/CTC: {agree}/{len(files)} ({100*agree//len(files)}%)")
# focus : clips que le LID dit 'lug'
lugs = [f for f in files if lid_pred[f] == "lug"]
print(f"\nParmi les {len(lugs)} clips LID='lug', ce que dit la confiance CTC:")
print(" ", dict(Counter(ctc_pred[f] for f in lugs)))
print("\n8 exemples (LID vs CTC + confiances):")
for f in files[:8]:
    print(f"  {f.split('/')[-1]}: LID={lid_pred[f]} CTC={ctc_pred[f]} | " +
          " ".join(f"{l}={conf[f][l]:.2f}" for l in CTC))
print("DIAG_DONE")