"""Hierarchical audio tagger: frozen Dasheng-0.6B features + an MLP head over a 4,416-node sound taxonomy.

    pip install torch dasheng soundfile scipy safetensors
    python inference.py clip1.flac clip2.wav ...

Each clip is read from its first 10 s, resampled to 16 kHz mono and zero-padded to 10 s. Dasheng-0.6B frame
features are mean-pooled to 1280-d and fed to the head, which has three outputs: 11 top-level groups, 822 second-level
nodes and 3,718 classes (nodes at depth <= 3). Class scores use the hierarchical readout
logP_L1 + logP_L2 + logP_L3 (each class scored with its own top-level and second-level ancestor).
The Dasheng weights are downloaded by the `dasheng` package on first use.
"""
import json
import sys
from math import gcd

import numpy as np
import soundfile as sf
import torch
import torch.nn as nn
from dasheng import dasheng_06B
from safetensors import safe_open
from safetensors.torch import load_file
from scipy.signal import resample_poly

SR, SECS = 16000, 10


def load_audio(path):
    """First 10 s of `path` as a (160000,) float32 array at 16 kHz mono, zero-padded."""
    with sf.SoundFile(path) as f:
        sr = f.samplerate
        a = f.read(SECS * sr, dtype="float32", always_2d=True).mean(1)
    if sr != SR:
        g = gcd(sr, SR)
        a = resample_poly(a, SR // g, sr // g).astype(np.float32)
    x = np.zeros(SECS * SR, np.float32)
    x[: min(len(a), SECS * SR)] = a[: SECS * SR]
    return x


class TaxonomyClassifier(nn.Module):
    def __init__(self, head_path="head.safetensors"):
        super().__init__()
        with safe_open(head_path, "pt") as f:
            meta = f.metadata()
        cfg, self.classes = json.loads(meta["config"]), json.loads(meta["classes"])
        l1 = [c.split("/")[0] for c in self.classes]
        l2 = ["/".join(c.split("/")[:2]) for c in self.classes]
        self.top_names, l2_names = sorted(set(l1)), sorted(set(l2))
        self.register_buffer("t1", torch.tensor([self.top_names.index(x) for x in l1]), persistent=False)
        self.register_buffer("t2", torch.tensor([l2_names.index(x) for x in l2]), persistent=False)
        h = cfg["hidden"]
        self.trunk = nn.Sequential(nn.Linear(cfg["in_dim"], h), nn.BatchNorm1d(h), nn.GELU(), nn.Dropout(cfg["dropout"]))
        self.l3, self.h1, self.h2 = nn.Linear(h, len(self.classes)), nn.Linear(h, len(self.top_names)), nn.Linear(h, len(l2_names))
        self.load_state_dict(load_file(head_path))
        self.encoder = dasheng_06B()
        self.eval()

    @torch.inference_mode()
    def forward(self, x):
        """x: (B, 160000) float32 at 16 kHz -> (top-level probs (B, 11), class probs (B, 3718))."""
        xl = torch.full((x.shape[0],), x.shape[1], dtype=torch.long, device=x.device)
        with torch.autocast(x.device.type, dtype=torch.bfloat16, enabled=x.device.type == "cuda"):
            feats = self.encoder(x, xl)  # (B, T, 1280)
        z = self.trunk(feats.float().mean(1))
        o3, o1, o2 = self.l3(z), self.h1(z), self.h2(z)
        score = o1.log_softmax(1)[:, self.t1] + o2.log_softmax(1)[:, self.t2] + o3.log_softmax(1)
        return o1.softmax(1), score.softmax(1)


if __name__ == "__main__":
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    model = TaxonomyClassifier().to(dev)
    paths = sys.argv[1:]
    top, cls = model(torch.from_numpy(np.stack([load_audio(p) for p in paths])).to(dev))
    for path, pt, pc in zip(paths, top.cpu(), cls.cpu()):
        print(f"{path}\n  top-level: {model.top_names[int(pt.argmax())]} ({float(pt.max()):.3f})")
        for p, i in zip(*pc.topk(5)):
            print(f"  {float(p):.3f}  {model.classes[int(i)]}")