from __future__ import annotations

from functools import lru_cache

import numpy as np
import soundfile as sf
import torch
from scipy.signal import resample_poly
from transformers import AutoFeatureExtractor, AutoModelForAudioClassification

MODEL_ID = "mispeech/ced-tiny"
MODEL_REVISION = "ace276d29dd0bb3f3517b0fa8cf300738c409019"
TARGET_SAMPLE_RATE = 16000
WINDOW_SAMPLES = 10 * TARGET_SAMPLE_RATE


@lru_cache(maxsize=1)
def _load_model():
    feature_extractor = AutoFeatureExtractor.from_pretrained(
        MODEL_ID,
        revision=MODEL_REVISION,
        trust_remote_code=True,
    )
    model = AutoModelForAudioClassification.from_pretrained(
        MODEL_ID,
        revision=MODEL_REVISION,
        trust_remote_code=True,
    )
    model.to("cpu")
    model.eval()
    return feature_extractor, model


def _load_audio(path: str) -> np.ndarray:
    audio, sample_rate = sf.read(path, dtype="float32", always_2d=True)
    if sample_rate <= 0 or audio.size == 0:
        raise ValueError("The audio file is empty or has an invalid sample rate.")
    if not np.isfinite(audio).all():
        raise ValueError("The audio file contains non-finite samples.")
    audio = audio.mean(axis=1)
    if sample_rate != TARGET_SAMPLE_RATE:
        divisor = int(np.gcd(sample_rate, TARGET_SAMPLE_RATE))
        audio = resample_poly(
            audio,
            TARGET_SAMPLE_RATE // divisor,
            sample_rate // divisor,
        ).astype(np.float32, copy=False)
    return np.ascontiguousarray(audio, dtype=np.float32)


def classify_audio(path: str, top_k: int = 10) -> list[dict[str, float | str]]:
    audio = _load_audio(path)
    feature_extractor, model = _load_model()
    window_probabilities = []
    for start in range(0, len(audio), WINDOW_SAMPLES):
        window = audio[start : start + WINDOW_SAMPLES]
        inputs = feature_extractor(
            window,
            sampling_rate=TARGET_SAMPLE_RATE,
            return_tensors="pt",
        )
        with torch.inference_mode():
            logits = model(**inputs).logits[0]
            window_probabilities.append(torch.sigmoid(logits))
    probabilities = torch.stack(window_probabilities).mean(dim=0)
    if not torch.isfinite(probabilities).all():
        raise ValueError("The model returned non-finite scores.")

    count = min(max(1, int(top_k)), probabilities.numel())
    scores, indices = torch.topk(probabilities, k=count)
    labels = model.config.id2label
    return [
        {
            "rank": rank,
            "label": str(labels.get(int(index), int(index))),
            "score": round(float(score), 6),
        }
        for rank, (score, index) in enumerate(zip(scores, indices), start=1)
    ]