"""PANNs (AudioSet) sample tagging → fleet-family probability vector.

A second, independent opinion to ensemble with CLAP. PANNs Cnn14 is trained on
AudioSet's 527 sound classes — many are exactly our percussion/instrument families
("Bass drum", "Snare drum", "Hi-hat", "Bass guitar", "Synthesizer", "Singing"…).
We sum the model's clip-wise probabilities over the AudioSet labels mapped to each
fleet family and normalize → a comparable 12-way vector.

Audio is decoded to 32 kHz mono (Cnn14's rate) via ffmpeg — no librosa needed here.
`family_vector(path)` returns a dict {family: prob} (None if nothing relevant fired).
The map is conservative: only confident, unambiguous AudioSet labels are wired; generic
classes ("Drum", "Music") are deliberately left out so they can't pollute the vote."""
import subprocess

import numpy as np

import sample_ontology as ONT

SR = 32000
FAMILIES = ONT.FAMILIES

# AudioSet display-name → fleet family. Exact substrings matched against panns labels.
AUDIOSET_FAMILY = {
    "Bass drum": "kick",
    "Snare drum": "snare", "Rimshot": "snare", "Clapping": "snare",
    "Hi-hat": "hat", "Cymbal": "hat", "Crash cymbal": "hat",
    "Tabla": "perc", "Tambourine": "perc", "Maraca": "perc", "Cowbell": "perc",
    "Wood block": "perc", "Tom-tom": "perc", "Conga": "perc", "Marimba, xylophone": "perc",
    "Bass guitar": "bass", "Double bass": "bass",
    "Piano": "keys", "Electric piano": "keys", "Organ": "keys",
    "Electronic organ": "keys", "Hammond organ": "keys", "Harpsichord": "keys",
    "Keyboard (musical)": "keys", "Mellotron": "keys", "Guitar": "keys",
    "Electric guitar": "keys", "Acoustic guitar": "keys", "Clavinet": "keys",
    "Synthesizer": "synth", "Sampler": "synth", "Sine wave": "synth",
    "Brass instrument": "lead", "Trumpet": "lead", "Saxophone": "lead", "Trombone": "lead",
    "String section": "pad", "Orchestra": "pad", "Violin, fiddle": "pad",
    "Sound effect": "fx", "Effects unit": "fx", "Whoosh, swoosh, swish": "fx",
    "Noise": "fx", "White noise": "fx", "Wind": "fx",
    "Scratching (performance technique)": "fx",
    "Singing": "vox", "Choir": "vox", "Male singing": "vox", "Female singing": "vox",
    "Rapping": "vox", "Speech": "vox", "Vocal music": "vox", "Chant": "vox",
    "Mantra": "vox", "Humming": "vox", "Yodeling": "vox",
}

_M = {}


def _model():
    if not _M:
        from panns_inference import AudioTagging
        from panns_inference.config import labels
        at = AudioTagging(checkpoint_path=None, device="cpu")
        # precompute label-index → family (sum many labels per family)
        idx2fam = {}
        for i, lab in enumerate(labels):
            if lab in AUDIOSET_FAMILY:
                idx2fam[i] = AUDIOSET_FAMILY[lab]
        _M.update(at=at, idx2fam=idx2fam)
    return _M


def load_audio(path):
    try:
        p = subprocess.run(
            ["ffmpeg", "-v", "quiet", "-i", str(path), "-ac", "1", "-ar", str(SR),
             "-t", "10", "-f", "f32le", "-"], capture_output=True, timeout=30)
        a = np.frombuffer(p.stdout, dtype=np.float32).copy()   # writable for torch
        if a.size < SR // 100:
            return None
        if a.size < SR:           # Cnn14 needs ≥~1s of mel frames; zero-pad short one-shots
            a = np.pad(a, (0, SR - a.size))
        return a
    except Exception:
        return None


def family_vector(path):
    """{family: prob} from AudioSet tags mapped to fleet families, or None."""
    a = load_audio(path)
    if a is None:
        return None
    M = _model()
    clip, _ = M["at"].inference(a[None, :])     # (1, 527)
    clip = clip[0]
    fam = {f: 0.0 for f in FAMILIES}
    for i, f in M["idx2fam"].items():
        fam[f] += float(clip[i])
    tot = sum(fam.values())
    if tot <= 1e-6:
        return None
    return {f: v / tot for f, v in fam.items()}


def classify_file(path):
    """(family, confidence) or None — the dominant mapped AudioSet family."""
    v = family_vector(path)
    if not v:
        return None
    f = max(v, key=v.get)
    return f, round(v[f], 3)
