"""Sanity check for the AnimeScore HuBERT release.

What it verifies:
  1. modeling_animescore.AnimeScoreRankNet can be built from config.json.
  2. model.safetensors loads with zero missing/unexpected non-SSL keys.
  3. A forward pass on a 3-second random waveform runs and returns a finite scalar.
  4. (Optional) If --wav is given, prints the AnimeScore for that file.

Usage:
  python sanity_test.py
  python sanity_test.py --wav path/to/clip.wav
"""

import argparse
import math
import os
import sys

import torch
from safetensors.torch import load_file


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--wav", help="Optional: score this wav file as a real-data check.")
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    args = ap.parse_args()

    here = os.path.dirname(os.path.abspath(__file__))
    sys.path.insert(0, here)

    from modeling_animescore import AnimeScoreConfig, AnimeScoreRankNet

    print("[1/4] Building model from config.json...")
    cfg_path = os.path.join(here, "config.json")
    if not os.path.exists(cfg_path):
        raise FileNotFoundError(cfg_path)
    cfg = AnimeScoreConfig.from_json_file(cfg_path)
    model = AnimeScoreRankNet(cfg).to(args.device).eval()
    print(f"      backbone = {cfg.ssl_backbone}")
    n_head = sum(p.numel() for n, p in model.named_parameters() if not n.startswith("ssl."))
    n_ssl = sum(p.numel() for n, p in model.named_parameters() if n.startswith("ssl."))
    print(f"      ssl = {n_ssl/1e6:.2f} M, head = {n_head/1e6:.2f} M")

    print("[2/4] Loading head weights from model.safetensors...")
    sd = load_file(os.path.join(here, "model.safetensors"))
    missing, unexpected = model.load_state_dict(sd, strict=False)
    head_missing = [m for m in missing if not m.startswith("ssl.")]
    if head_missing:
        raise RuntimeError(f"head keys missing after load: {head_missing}")
    if unexpected:
        raise RuntimeError(f"unexpected keys in safetensors: {unexpected}")
    print(f"      loaded {len(sd)} head tensors, 0 missing, 0 unexpected.")

    print("[3/4] Forward pass on 3 s of random audio...")
    wav = torch.randn(1, 16000 * 3).to(args.device)
    with torch.no_grad():
        s = model.score(wav).item()
    if not math.isfinite(s):
        raise RuntimeError(f"non-finite score: {s}")
    print(f"      score = {s:+.4f}  (random audio; value is uninformative, just non-NaN check)")

    if args.wav:
        print(f"[4/4] Scoring real audio: {args.wav}")
        import torchaudio
        try:
            import soundfile as sf
            data, sr = sf.read(args.wav, dtype="float32", always_2d=True)
            wav = torch.from_numpy(data.T).contiguous()
        except Exception:
            wav, sr = torchaudio.load(args.wav)
        if wav.size(0) > 1:
            wav = wav.mean(0, keepdim=True)
        if sr != cfg.target_sr:
            wav = torchaudio.functional.resample(wav, sr, cfg.target_sr)
        with torch.no_grad():
            s = model.score(wav.to(args.device)).item()
        print(f"      AnimeScore({os.path.basename(args.wav)}) = {s:+.4f}")
    else:
        print("[4/4] (skipped) Pass --wav to score a real audio file.")

    print("\nAll sanity checks passed.")


if __name__ == "__main__":
    main()