"""Gender detection using a fine-tuned wav2vec2 classifier.

Model: prithivMLmods/Common-Voice-Gender-Detection
  https://huggingface.co/prithivMLmods/Common-Voice-Gender-Detection
A `facebook/wav2vec2-base-960h` model fine-tuned for binary (female/male)
speaker-gender classification. We feed it a 16 kHz mono waveform and read the
softmax probabilities. If the top probability is below a confidence floor (or
there isn't enough audio) we report "uncertain" instead of guessing.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from functools import lru_cache

import numpy as np
import torch
from transformers import Wav2Vec2FeatureExtractor, Wav2Vec2ForSequenceClassification

MODEL_NAME = "prithivMLmods/Common-Voice-Gender-Detection"
TARGET_SR = 16000
MIN_SAMPLES = int(0.3 * TARGET_SR)  # need ~0.3s of audio to bother
CONFIDENCE_FLOOR = 0.6  # below this we say "uncertain"


@dataclass
class GenderResult:
    label: str  # "male" | "female" | "uncertain"
    confidence: float  # top-class softmax probability (0..1)
    probs: dict = field(default_factory=dict)  # {"male": p, "female": p}
    audio_seconds: float = 0.0


@lru_cache(maxsize=1)
def _load():
    """Load (and cache) the model + feature extractor once per process."""
    model = Wav2Vec2ForSequenceClassification.from_pretrained(MODEL_NAME)
    extractor = Wav2Vec2FeatureExtractor.from_pretrained(MODEL_NAME)
    model.eval()
    id2label = {int(k): str(v).lower() for k, v in model.config.id2label.items()}
    return model, extractor, id2label


def warmup() -> None:
    """Eagerly load the model (e.g. at server startup) to avoid a cold first request."""
    _load()


def classify(samples: np.ndarray, sr: int) -> GenderResult:
    """Classify gender from a mono waveform (float32)."""
    audio_seconds = float(len(samples) / sr) if sr else 0.0

    if samples is None or samples.size < MIN_SAMPLES:
        return GenderResult("uncertain", 0.0, {}, round(audio_seconds, 2))

    samples = np.asarray(samples, dtype=np.float32)
    if sr != TARGET_SR:
        import librosa

        samples = librosa.resample(samples, orig_sr=sr, target_sr=TARGET_SR)
        sr = TARGET_SR

    model, extractor, id2label = _load()

    inputs = extractor(samples, sampling_rate=sr, return_tensors="pt", padding=True)
    with torch.no_grad():
        logits = model(**inputs).logits
        probs = torch.softmax(logits, dim=1).squeeze(0).tolist()

    prob_map = {id2label[i]: float(probs[i]) for i in range(len(probs))}
    label = max(prob_map, key=prob_map.get)
    confidence = prob_map[label]
    if confidence < CONFIDENCE_FLOOR:
        label = "uncertain"

    return GenderResult(
        label=label,
        confidence=round(confidence, 3),
        probs={k: round(v, 3) for k, v in prob_map.items()},
        audio_seconds=round(audio_seconds, 2),
    )