"""Classify any audio file into speech_noise / speech_music / singing_music / none.

Command line:
    python predict_crnn.py song.mp3

Library:
    from predict_crnn import SmadClassifier
    clf = SmadClassifier()
    print(clf.predict("song.mp3")["overall"])

Audio of any length, sample rate or channel count is accepted: it is converted
to 16 kHz mono, cut into 4-second windows, and each window classified. That
window length is not a free choice -- it is what the model was trained on.
"""

import argparse
import json
from pathlib import Path

import numpy as np
import torch

from crnn_model import TinyAudioCRNN

CLASS_NAMES = ("speech_noise", "speech_music", "singing_music", "none")

DESCRIPTIONS = {
    "speech_noise": "spoken voice over non-music noise",
    "speech_music": "spoken voice over a music bed",
    "singing_music": "sung voice over music",
    "none": "no human voice: instrumental music, noise, or silence",
}

# Must match mel_features.py exactly -- these numbers define the input the
# model was trained on, and changing any of them silently degrades accuracy
# rather than raising an error. Duplicated here so this file works standalone,
# and checked against the training module below when it is importable.
SAMPLE_RATE = 16000
SEGMENT_SECONDS = 4.0
N_MELS = 80
N_FFT = 400          # 25 ms
HOP_LENGTH = 160     # 10 ms

DEFAULT_CHECKPOINT = Path(__file__).with_name("models") / "crnn_scratch_v1.pt"


def _check_constants_match_training():
    """Fail loudly if the training module and this file ever disagree."""
    try:
        import mel_features
    except Exception:  # noqa: BLE001 - standalone use, nothing to compare against
        return
    mismatches = [
        name for name, here, there in [
            ("N_MELS", N_MELS, mel_features.N_MELS),
            ("N_FFT", N_FFT, mel_features.N_FFT),
            ("HOP_LENGTH", HOP_LENGTH, mel_features.HOP_LENGTH),
        ] if here != there
    ]
    if mismatches:
        raise RuntimeError(
            f"feature settings drifted from mel_features.py: {mismatches}. "
            "Inference must use the same settings as training.")


def load_audio(path, sample_rate=SAMPLE_RATE):
    """Any audio file -> mono float32 at `sample_rate`."""
    import librosa

    audio, _ = librosa.load(str(path), sr=sample_rate, mono=True)
    return audio.astype(np.float32)


def waveform_to_mel(waveform, sample_rate=SAMPLE_RATE):
    """(N,) float32 -> (T, N_MELS) log-mel in dB."""
    import librosa

    mel = librosa.feature.melspectrogram(
        y=waveform, sr=sample_rate, n_fft=N_FFT,
        hop_length=HOP_LENGTH, n_mels=N_MELS, power=2.0)
    return librosa.power_to_db(mel).T.astype(np.float32)


class SmadClassifier:
    """The trained CRNN, ready to classify audio.

    The checkpoint carries its own input normalisation and confidence
    temperature, so nothing here needs configuring -- feeding the model a
    differently-scaled input than it trained on is not possible by accident.
    """

    def __init__(self, checkpoint=None, device=None):
        _check_constants_match_training()
        path = Path(checkpoint) if checkpoint else DEFAULT_CHECKPOINT
        if not path.exists():
            raise FileNotFoundError(f"checkpoint not found: {path}")

        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        ckpt = torch.load(path, map_location=self.device, weights_only=False)

        self.model = TinyAudioCRNN()
        self.model.load_state_dict(ckpt["state_dict"])   # includes feat_mean/feat_std
        self.model.to(self.device).eval()

        # Divide logits by this before softmax. Fitted on validation; it is a
        # monotonic rescale, so it never changes which class wins -- it only
        # stops the model overstating how sure it is.
        self.temperature = float(ckpt.get("temperature", 1.0))
        self.trained_val_acc = ckpt.get("val_acc")

    @torch.no_grad()
    def predict_windows(self, audio, sample_rate=SAMPLE_RATE, hop_seconds=None, batch_size=64):
        """One prediction per window. Returns a list of dicts.

        Windows are run through the model in chunks of `batch_size` rather
        than all at once -- a merged/long input can have thousands of 4s
        windows, and stacking them into a single GPU batch is what causes an
        out-of-memory error on a shared GPU, not the model itself.
        """
        window = int(SEGMENT_SECONDS * sample_rate)
        hop = int((hop_seconds or SEGMENT_SECONDS) * sample_rate)

        if audio.shape[0] < window:                 # pad a too-short clip
            audio = np.pad(audio, (0, window - audio.shape[0]))

        starts = list(range(0, audio.shape[0] - window + 1, hop))
        mels = [waveform_to_mel(audio[s:s + window], sample_rate) for s in starts]

        probs_chunks = []
        for i in range(0, len(mels), batch_size):
            batch = torch.from_numpy(np.stack(mels[i:i + batch_size])).to(self.device)
            probs_chunks.append(
                torch.softmax(self.model(batch) / self.temperature, dim=-1).cpu().numpy())
        probs = np.concatenate(probs_chunks, axis=0)

        out = []
        for start, p in zip(starts, probs):
            best = int(p.argmax())
            out.append({
                "start_s": round(start / sample_rate, 2),
                "end_s": round((start + window) / sample_rate, 2),
                "label": CLASS_NAMES[best],
                "confidence": round(float(p[best]), 4),
                "probs": {name: round(float(v), 4) for name, v in zip(CLASS_NAMES, p)},
            })
        return out

    def predict(self, path, hop_seconds=None):
        """Classify a file. Returns per-window results plus an overall verdict.

        `overall` is the label covering the most time, not an average of
        probabilities -- a 3-minute song with a spoken intro is a song, and
        averaging would let a handful of confident intro windows outvote the
        rest.
        """
        audio = load_audio(path)
        windows = self.predict_windows(audio, hop_seconds=hop_seconds)

        share = {name: 0 for name in CLASS_NAMES}
        for w in windows:
            share[w["label"]] += 1
        overall = max(share, key=share.get)

        return {
            "file": str(path),
            "duration_s": round(len(audio) / SAMPLE_RATE, 2),
            "overall": overall,
            "overall_share": round(share[overall] / max(len(windows), 1), 4),
            "share": {k: round(v / max(len(windows), 1), 4) for k, v in share.items()},
            "windows": windows,
        }


def main():
    parser = argparse.ArgumentParser(
        description="Classify audio into speech_noise / speech_music / singing_music / none")
    parser.add_argument("audio", nargs="+", help="audio file(s); any format ffmpeg reads")
    parser.add_argument("--checkpoint", default=None)
    parser.add_argument("--hop-seconds", type=float, default=None,
                        help="window step; defaults to 4.0 (no overlap)")
    parser.add_argument("--json", action="store_true", help="print JSON instead of a table")
    parser.add_argument("--quiet", action="store_true", help="only the overall label")
    args = parser.parse_args()

    clf = SmadClassifier(args.checkpoint)

    for path in args.audio:
        result = clf.predict(path, hop_seconds=args.hop_seconds)

        if args.json:
            print(json.dumps(result, indent=2))
            continue
        if args.quiet:
            print(f"{result['overall']}\t{path}")
            continue

        print(f"\n{path}  ({result['duration_s']:.1f}s)")
        for w in result["windows"]:
            print(f"  {w['start_s']:7.1f} - {w['end_s']:6.1f}s   "
                  f"{w['label']:<14s} {w['confidence'] * 100:5.1f}%")
        print(f"  {'':>7s}   {'':>6s}    {'-' * 26}")
        print(f"  overall: {result['overall']} "
              f"({result['overall_share'] * 100:.0f}% of the audio) "
              f"-- {DESCRIPTIONS[result['overall']]}")


if __name__ == "__main__":
    main()