#!/usr/bin/env python3
"""
predict.py — live-vs-studio music classifier inference.

Pipeline (per audio file):
  1. Resample to 16 kHz mono
  2. Extract three 30-second windows at 25/50/75% of the usable span
  3. Compute AST log-mel features for each window
  4. AST forward pass -> 768-dim pooler_output
  5. Linear probe -> (p_studio, p_live)
  6. Majority vote across windows

Usage as CLI:
  python predict.py path/to/song.mp3
  python predict.py *.mp3
  python predict.py --json song.mp3       # machine-readable output
  python predict.py --device cpu song.mp3 # force CPU instead of MPS/CUDA

Usage as library (e.g. from a Gradio Space):
  import predict
  classifier = predict.Classifier()              # loads AST + probe once
  result = classifier.classify_file("song.mp3")  # dict with prediction, confidence, per-window probs
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Optional

import librosa
import numpy as np
import torch
import torch.nn as nn
from transformers import ASTFeatureExtractor, ASTModel

# The AST backbone we use. Same one the browser demo loads (its int8 ONNX
# variant lives at Xenova/<same id>; here we use HuggingFace's reference
# PyTorch weights, which is the easier path for a Python user).
MODEL_ID = "MIT/ast-finetuned-audioset-10-10-0.4593"

# Clip-windowing constants — must match the JS reference implementation so
# the CLI and any browser-side path produce identical input tensors.
SR = 16000
CLIP_SEC = 30
CLIP_LEN = SR * CLIP_SEC
HEAD_SKIP = 10
TAIL_SKIP = 10
POSITIONS = [0.25, 0.5, 0.75]

# The probe weights live next to this script when installed as a HF model repo.
# Don't .resolve() __file__ — when this module is imported from a HuggingFace
# snapshot cache, __file__ is a symlink into ../../blobs/<sha>, and resolving
# it puts us in the blobs dir where there's no probe-weights.json by name.
# Path(__file__).parent stays in the snapshots/<rev>/ dir where the probe-
# weights.json symlink (alongside this predict.py symlink) lives.
DEFAULT_PROBE = Path(__file__).parent / "probe-weights.json"


def pick_device(requested: Optional[str] = None) -> torch.device:
    if requested:
        return torch.device(requested)
    if torch.backends.mps.is_available():
        return torch.device("mps")
    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


def load_audio(path) -> np.ndarray:
    """Decode any common audio file as 16 kHz mono float32."""
    y, _ = librosa.load(str(path), sr=SR, mono=True)
    return y.astype(np.float32)


def extract_windows(y: np.ndarray) -> tuple[list[np.ndarray], str]:
    """Return up to three 30s windows from 16 kHz mono audio.

    Same logic as the reference web/audio.js so both paths produce the same
    inputs for the same audio source.
    """
    total_sec = len(y) / SR
    if len(y) < CLIP_LEN:
        if len(y) < SR * 5:
            return [], "audio is under 5 seconds — too short to classify"
        padded = np.zeros(CLIP_LEN, dtype=np.float32)
        padded[: len(y)] = y
        return [padded], f"padded {total_sec:.1f}s to 30s"

    if total_sec < CLIP_SEC * 2:
        start = max(0, (len(y) - CLIP_LEN) // 2)
        return [y[start:start + CLIP_LEN]], (
            f"single middle window from {total_sec:.1f}s clip"
        )

    usable = total_sec - HEAD_SKIP - TAIL_SKIP
    if usable < CLIP_SEC:
        start = max(0, (len(y) - CLIP_LEN) // 2)
        return [y[start:start + CLIP_LEN]], (
            f"single window from {total_sec:.1f}s clip"
        )

    windows = []
    for p in POSITIONS:
        center_sec = HEAD_SKIP + usable * p
        start = int((center_sec - CLIP_SEC / 2) * SR)
        start = max(0, min(start, len(y) - CLIP_LEN))
        windows.append(y[start:start + CLIP_LEN])
    return windows, f"three windows at 25/50/75% of usable span ({total_sec:.1f}s clip)"


class Probe(nn.Module):
    """Linear-probe head, loaded from `probe-weights.json`.

    Keeping weights as JSON means the Python and browser (`transformers.js`)
    inference paths share a single artifact — one source of truth.
    """

    def __init__(self, weights_path: Path):
        super().__init__()
        cfg = json.loads(weights_path.read_text())
        self.fc = nn.Linear(cfg["in_dim"], cfg["num_classes"])
        with torch.no_grad():
            self.fc.weight.copy_(torch.tensor(cfg["weight"], dtype=torch.float32))
            self.fc.bias.copy_(torch.tensor(cfg["bias"], dtype=torch.float32))
        # id2label keys may be strings or ints depending on the JSON writer; normalize.
        id2label = cfg["id2label"]
        self.id2label = {int(k): v for k, v in id2label.items()}
        self.version = cfg.get("version", "?")
        self.metrics = {
            "trained_val_acc": cfg.get("trained_val_acc"),
            "trained_test_acc": cfg.get("trained_test_acc"),
            "source_quality_test_acc": cfg.get("source_quality_test_acc"),
        }

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fc(x)


class Classifier:
    """Stateful classifier — loads AST + probe once, predicts many files.

    Use this when running multiple predictions (e.g. from a Gradio Space).
    For a single prediction the module-level `classify_file()` helper is
    fine and saves you a line.
    """

    def __init__(self, probe_path: Optional[Path] = None,
                 device: Optional[str] = None):
        self.device = pick_device(device)
        self.feature_extractor = ASTFeatureExtractor.from_pretrained(MODEL_ID)
        self.model = ASTModel.from_pretrained(MODEL_ID).to(self.device).eval()
        self.probe = Probe(probe_path or DEFAULT_PROBE).to(self.device).eval()

    @torch.no_grad()
    def classify(self, audio_path) -> dict:
        """Classify one file. Returns a dict; never raises on per-file errors
        (they're returned in result["error"])."""
        try:
            y = load_audio(audio_path)
        except Exception as e:
            return {"file": str(audio_path), "error": f"failed to decode: {e}"}

        windows, info = extract_windows(y)
        if not windows:
            return {"file": str(audio_path), "error": info, "windows": []}

        per_window = []
        for w in windows:
            features = self.feature_extractor(w, sampling_rate=SR, return_tensors="pt")
            x = features["input_values"].to(self.device)
            pooled = self.model(x).pooler_output  # (1, 768)
            logits = self.probe(pooled)           # (1, 2)
            probs = torch.softmax(logits, dim=1).cpu().numpy()[0]
            per_window.append({
                "p_studio": float(probs[0]),
                "p_live": float(probs[1]),
                "label": "live" if probs[1] >= probs[0] else "studio",
            })

        n_live = sum(1 for w in per_window if w["label"] == "live")
        final_label = "live" if n_live >= (len(per_window) + 1) // 2 else "studio"
        avg_p_live = sum(w["p_live"] for w in per_window) / len(per_window)
        final_conf = avg_p_live if final_label == "live" else 1 - avg_p_live

        return {
            "file": str(audio_path),
            "info": info,
            "label": final_label,
            "confidence": float(final_conf),
            "p_live": float(avg_p_live),
            "p_studio": float(1 - avg_p_live),
            "windows": per_window,
            "agreement": f"{n_live}/{len(per_window)} live",
        }


# Module-level singleton for the one-shot helper. Lazy so importing this module
# is cheap; first call to classify_file() pays the AST load.
_default_classifier: Optional[Classifier] = None


def classify_file(audio_path, device: Optional[str] = None) -> dict:
    """One-shot prediction. For multiple files, instantiate Classifier directly
    to avoid re-loading the model every call."""
    global _default_classifier
    if _default_classifier is None:
        _default_classifier = Classifier(device=device)
    return _default_classifier.classify(audio_path)


def render_text(result: dict) -> str:
    if "error" in result:
        return f"{result['file']}\n  ERROR: {result['error']}\n"
    lines = [
        result["file"],
        f"  prediction: {result['label']}",
        f"  confidence: {result['confidence']:.3f}  ({result['agreement']})",
        f"  per-window p(live): "
        + " ".join(f"{w['p_live']:.3f}" for w in result["windows"]),
        f"  ({result['info']})",
    ]
    return "\n".join(lines) + "\n"


def main() -> None:
    p = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    p.add_argument("files", nargs="+", type=Path,
                   help="audio file(s) to classify (mp3, wav, flac, m4a, ogg …)")
    p.add_argument("--probe-weights", type=Path, default=DEFAULT_PROBE,
                   help=f"path to probe weights JSON (default: probe-weights.json "
                        "next to this script)")
    p.add_argument("--device", choices=["cpu", "cuda", "mps"],
                   help="force a specific torch device (default: auto-detect)")
    p.add_argument("--json", action="store_true",
                   help="print one JSON object per file on stdout instead of text")
    args = p.parse_args()

    if not args.probe_weights.exists():
        sys.exit(f"probe weights not found: {args.probe_weights}")

    classifier = Classifier(probe_path=args.probe_weights, device=args.device)
    if not args.json:
        print(f"Loaded {MODEL_ID} on {classifier.device}", file=sys.stderr)
        print(f"Probe {classifier.probe.version} "
              f"(test acc: {classifier.probe.metrics['trained_test_acc']:.3f})",
              file=sys.stderr)

    for path in args.files:
        if not path.exists():
            err = {"file": str(path), "error": "file not found"}
            print(json.dumps(err) if args.json else f"{path}\n  ERROR: file not found\n")
            continue
        try:
            result = classifier.classify(path)
        except Exception as e:
            result = {"file": str(path), "error": str(e)}
        print(json.dumps(result) if args.json else render_text(result))


if __name__ == "__main__":
    main()