"""Reusable gender-ID pipeline: sherpa-onnx speaker embedding + this repo's ONNX head.

No custom encoder here - sherpa-onnx's pretrained 3D-Speaker CAM++ speaker-embedding
extractor (VoxCeleb, 16 kHz) does the acoustic modeling; this repo only ships the
tiny MLP head trained on top of its embeddings.
"""
from __future__ import annotations

import numpy as np
import onnxruntime as ort
import sherpa_onnx
import soundfile as sf
from huggingface_hub import hf_hub_download

EMBED_MODEL_REPO = "csukuangfj/speaker-embedding-models"
EMBED_MODEL_FILE = "3dspeaker_speech_campplus_sv_en_voxceleb_16k.onnx"
HEAD_MODEL_REPO = "AfriSpeech/afrispeech-gender-id"
TARGET_SR = 16000


class GenderClassifier:
    def __init__(self, num_threads: int = 2):
        embed_path = hf_hub_download(EMBED_MODEL_REPO, EMBED_MODEL_FILE)
        config = sherpa_onnx.SpeakerEmbeddingExtractorConfig(
            model=embed_path, num_threads=num_threads, debug=False, provider="cpu"
        )
        self.extractor = sherpa_onnx.SpeakerEmbeddingExtractor(config)

        # config.json holds the output-index -> gender label mapping (as
        # well as being one of HF Hub's default download-count query files,
        # so loading it here also makes real use of this class register as a
        # tracked download on the model page).
        import json

        head_path = hf_hub_download(HEAD_MODEL_REPO, "onnx/model.onnx")
        config_path = hf_hub_download(HEAD_MODEL_REPO, "config.json")

        with open(config_path, encoding="utf-8") as f:
            self.label_map = json.load(f)["label_map"]
        self.session = ort.InferenceSession(head_path, providers=["CPUExecutionProvider"])

    def embed(self, samples: np.ndarray, sr: int) -> np.ndarray:
        stream = self.extractor.create_stream()
        stream.accept_waveform(sample_rate=sr, waveform=samples)
        stream.input_finished()
        return np.asarray(self.extractor.compute(stream), dtype=np.float32)

    def predict(self, samples: np.ndarray, sr: int) -> tuple[str, float]:
        embedding = self.embed(samples, sr).reshape(1, -1)
        logits = self.session.run(["logits"], {"embedding": embedding})[0][0]
        probs = np.exp(logits - logits.max())
        probs /= probs.sum()
        pred_idx = int(probs.argmax())
        return self.label_map[str(pred_idx)], float(probs[pred_idx])

    def predict_file(self, path: str) -> tuple[str, float]:
        audio, sr = sf.read(path, dtype="float32", always_2d=False)
        if audio.ndim > 1:
            audio = audio.mean(axis=1)
        return self.predict(audio, sr)