#!/usr/bin/env python3
"""Evaluate every model folder with Whisper (CER/WER) and TitaNet (speaker similarity).

Folder layout:
    <root>/<model>/sample_00.wav ...

Usage:
    python evaluate_compare.py --root /path/eval_in \
        --models bayan,magpie,xtts,chatterbox,mms \
        --text-file ../texts.txt --out ../metrics.json --csv ../results.csv \
        --whisper /path/whisper-large-v3 --titanet /path/titanet.nemo \
        --speaker-ref-dir /path/wav22k
"""
from __future__ import annotations

import argparse
import csv
import json
import re
import unicodedata
from pathlib import Path

import numpy as np
import soundfile as sf
import torch
import torch.nn.functional as F
from faster_whisper import WhisperModel
from jiwer import cer, wer
from nemo.collections.asr.models import EncDecSpeakerLabelModel

DIACRITICS = re.compile(r"[\u0610-\u061a\u064b-\u065f\u0670\u06d6-\u06ed\u0640]")
NUMBER_WORDS = (
    "واحد|واحده|اثنان|اثنين|ثلاثه|اربعه|خمسه|سته|سبعه|ثمانيه|تسعه|عشره|عشر|"
    "عشرون|عشرين|ثلاثون|ثلاثين|اربعون|اربعين|خمسون|خمسين|ستون|ستين|"
    "سبعون|سبعين|ثمانون|ثمانين|تسعون|تسعين|مايه|مائه|ميه|ميتين|مائتين|مئتين|مائتان|"
    "الف|الاف|الفين|الفان|مليون|مليار|مليارات|"
    "ثلاثمائه|ثلاثمايه|اربعمائه|اربعمايه|خمسمائه|خمسمايه|ستمائه|ستمايه|"
    "سبعمائه|سبعمايه|ثمانمائه|ثمانمايه|تسعمائه|تسعمايه"
)
UNITS = "كم|كيلومتر|كيلومترات|كيلومترا|متر|مترا|امتار|عام|عاما|اعوام|سنوات|سنه|دوله|دول|شخص|اشخاص"
NUM_RE = re.compile(rf"\bو?(?:{NUMBER_WORDS}|{UNITS}|[0-9]+)\b")


def normalize_ar(text: str) -> str:
    text = unicodedata.normalize("NFKC", text)
    text = DIACRITICS.sub("", text)
    text = (
        text.replace("أ", "ا").replace("إ", "ا").replace("آ", "ا")
        .replace("ى", "ي").replace("ة", "ه").replace("ؤ", "و").replace("ئ", "ي")
    )
    text = re.sub(r"[^\w\s]", " ", text)
    return re.sub(r"\s+", " ", text).strip()


def mask_numbers(text: str) -> str:
    out = NUM_RE.sub("NUM", normalize_ar(text))
    out = re.sub(r"(?:NUM\s*)+", "NUM ", out)
    return re.sub(r"\s+", " ", out).strip()


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", required=True)
    ap.add_argument("--models", required=True)
    ap.add_argument("--text-file", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--csv", required=True)
    ap.add_argument("--whisper", required=True)
    ap.add_argument("--titanet", required=True)
    ap.add_argument("--speaker-ref-dir", required=True)
    ap.add_argument("--ref-stride", type=int, default=50)
    ap.add_argument("--ref-count", type=int, default=12)
    args = ap.parse_args()

    root = Path(args.root)
    texts = [ln.strip() for ln in Path(args.text_file).read_text(encoding="utf-8").splitlines() if ln.strip()]
    whisper = WhisperModel(args.whisper, device="cpu", compute_type="int8")
    sv = EncDecSpeakerLabelModel.restore_from(args.titanet, map_location="cpu").eval()

    refs = sorted(Path(args.speaker_ref_dir).glob("*.wav"))[:: args.ref_stride][: args.ref_count]
    with torch.no_grad():
        embs = torch.stack([F.normalize(sv.get_embedding(str(p)).reshape(-1), dim=0) for p in refs])
    voice = F.normalize(embs.mean(0), dim=0)
    print(f"voice centroid from {len(refs)} clips")

    report: dict[str, dict] = {}
    for model in [m.strip() for m in args.models.split(",") if m.strip()]:
        files = sorted((root / model).glob("sample_*.wav"))
        rows = []
        for f, ref in zip(files, texts):
            segs, _ = whisper.transcribe(str(f), language="ar", beam_size=5, vad_filter=False, condition_on_previous_text=False)
            hyp = " ".join(s.text.strip() for s in segs).strip()
            c = float(cer(normalize_ar(ref), normalize_ar(hyp)))
            c_norm = float(cer(mask_numbers(ref), mask_numbers(hyp)))
            w = float(wer(normalize_ar(ref), normalize_ar(hyp)))
            with torch.no_grad():
                e = F.normalize(sv.get_embedding(str(f)).reshape(-1), dim=0)
            sim = float(e @ voice)
            dur = float(sf.info(str(f)).duration)
            rows.append({"file": f.name, "cer": c, "cer_norm": c_norm, "wer": w,
                         "sim": sim, "duration": dur, "hyp": hyp, "ref": ref})
            print(f"{model}/{f.name} CER={c:.3f} CER*={c_norm:.3f} WER={w:.3f} SIM={sim:.3f}")
        report[model] = {
            "n": len(rows),
            "mean_cer": float(np.mean([r["cer"] for r in rows])),
            "mean_cer_norm": float(np.mean([r["cer_norm"] for r in rows])),
            "mean_wer": float(np.mean([r["wer"] for r in rows])),
            "mean_sim": float(np.mean([r["sim"] for r in rows])),
            "total_duration": float(np.sum([r["duration"] for r in rows])),
            "samples": rows,
        }

    Path(args.out).write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
    with Path(args.csv).open("w", newline="", encoding="utf-8") as fh:
        wr = csv.writer(fh)
        wr.writerow(["model", "file", "cer", "cer_numeric_normalized", "wer", "sim", "duration", "hypothesis"])
        for model, rep in report.items():
            for r in rep["samples"]:
                wr.writerow([model, r["file"], f"{r['cer']:.4f}", f"{r['cer_norm']:.4f}",
                             f"{r['wer']:.4f}", f"{r['sim']:.4f}", f"{r['duration']:.2f}", r["hyp"]])

    for model, rep in report.items():
        print(f"\n{model}: CER={rep['mean_cer']:.4f} CER*={rep['mean_cer_norm']:.4f} "
              f"WER={rep['mean_wer']:.4f} SIM={rep['mean_sim']:.3f} dur={rep['total_duration']:.1f}s")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())