"""P(dialect) for an Arabic speech clip: spoken dialect vs Modern Standard Arabic (MSA).

  pip install torch transformers librosa safetensors huggingface_hub
  python inference.py clip.wav            # prints P(dialect), 0..1

  from inference import DialectScorer
  DialectScorer().p_dialect("clip.wav")   # downloads nsa01n/najdi-classifier + openai/whisper-large-v3-turbo

Features: mean over the real audio frames of Whisper large-v3-turbo encoder layers 20, 24, 28, 32
(float32), concatenated; then standardize and logistic regression (weights in model.safetensors).
"""

import json
import sys
from pathlib import Path

import numpy as np

REPO = "nsa01n/najdi-classifier"
SR = 16000


def probability(x, w):
    """x: n x features -> P(dialect) per row."""
    z = ((np.asarray(x, np.float32) - w["scaler_mean"]) / w["scaler_scale"]) @ w["weight"] + w["bias"][0]
    return 1.0 / (1.0 + np.exp(-z))


class DialectScorer:
    def __init__(self, repo=REPO, device=None):
        import torch
        from huggingface_hub import snapshot_download
        from safetensors.numpy import load_file
        from transformers import WhisperFeatureExtractor, WhisperModel

        path = Path(repo) if Path(repo).is_dir() else Path(snapshot_download(repo))
        self.cfg = json.loads((path / "config.json").read_text())
        self.w = load_file(path / "model.safetensors")
        self.device = device or ("cuda" if torch.cuda.is_available() else
                                 "mps" if torch.backends.mps.is_available() else "cpu")
        self.fe = WhisperFeatureExtractor.from_pretrained(self.cfg["encoder"])
        self.encoder = WhisperModel.from_pretrained(self.cfg["encoder"], dtype=torch.float32).get_encoder()
        self.encoder = self.encoder.to(self.device).eval()
        self.torch = torch

    def features(self, audio):
        audio = np.asarray(audio, np.float32)[: self.cfg["max_seconds"] * SR]
        x = self.fe(audio, sampling_rate=SR, return_tensors="pt").input_features.to(self.device)
        with self.torch.no_grad():
            hidden = self.encoder(x, output_hidden_states=True).hidden_states
        n = max(1, int(np.ceil(len(audio) / SR * 50)))  # 50 encoder frames per second; padding left out
        return np.concatenate([hidden[l][0, :n].mean(0).float().cpu().numpy() for l in self.cfg["layers"]])

    def p_dialect(self, audio):
        if isinstance(audio, (str, Path)):
            import librosa

            audio, _ = librosa.load(str(audio), sr=SR, mono=True)
        return float(probability(self.features(audio)[None], self.w)[0])


if __name__ == "__main__":
    scorer = DialectScorer()
    for f in sys.argv[1:]:
        print(f"{f}\tP(dialect) = {scorer.p_dialect(f):.3f}")