"""
UTMOS naturalness evaluator.

Matches the F5-TTS evaluation approach: SpeechMOS UTMOS22 strong predictor
loaded through torch.hub. Kept independent and lazy.
"""

import gc
import threading
from pathlib import Path

import numpy as np
import soundfile as sf
import torch

_LOCK = threading.RLock()
_PREDICTOR = None
_DEVICE = "cpu"


def _load():
    global _PREDICTOR
    with _LOCK:
        if _PREDICTOR is not None:
            return _PREDICTOR
        _PREDICTOR = torch.hub.load(
            "tarepan/SpeechMOS:v1.2.0",
            "utmos22_strong",
            trust_repo=True,
        )
        _PREDICTOR = _PREDICTOR.to(_DEVICE).eval()
        return _PREDICTOR


def evaluate_utmos(audio_path: str) -> dict:
    if not audio_path or not Path(audio_path).exists():
        raise ValueError("الصوت الناتج غير موجود.")

    wav, sr = sf.read(audio_path, dtype="float32", always_2d=False)
    wav = np.asarray(wav, dtype=np.float32)
    if wav.ndim > 1:
        wav = wav.mean(axis=1)
    if wav.size == 0:
        raise ValueError("الصوت الناتج فارغ.")

    predictor = _load()
    tensor = torch.from_numpy(wav).unsqueeze(0).to(_DEVICE)
    with torch.inference_mode():
        score = float(predictor(tensor, int(sr))[0].detach().cpu().item())

    return {
        "metric": "UTMOS Naturalness",
        "score": score,
        "scale": "1-5",
        "higher_is_better": True,
        "model": "tarepan/SpeechMOS:v1.2.0 / utmos22_strong",
    }


def release_utmos():
    global _PREDICTOR
    with _LOCK:
        _PREDICTOR = None
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()