"""
Self-contained loader for AudioMosaic-vit-b16-finetune-esc-split2.

Usage:
    from huggingface_hub import snapshot_download
    import sys
    local = snapshot_download("hanxunh/AudioMosaic-vit-b16-finetune-esc-split2")
    sys.path.insert(0, local)
    from load_model import load_classifier
    model = load_classifier()
"""
import json
import os
import torch
from safetensors.torch import load_file


def load_classifier(repo_dir: str = None, device: str = "cpu"):
    """Load AudioMosaic ViT-B/16 classifier finetuned on ESC-50 (fold 2) (Acc 98.75)."""
    if repo_dir is None:
        repo_dir = os.path.dirname(os.path.abspath(__file__))
    cfg_path = os.path.join(repo_dir, "config.json")
    weights_path = os.path.join(repo_dir, "model.safetensors")

    with open(cfg_path) as f:
        cfg = json.load(f)
    cfg.pop("model_class", None)

    from modeling import AudioMosaicClassifier
    model = AudioMosaicClassifier(**cfg)
    state = load_file(weights_path)
    msg = model.load_state_dict(state, strict=False)
    print(msg)
    return model.to(device).eval()


def preprocess(audio, repo_dir: str = None, **overrides):
    """Convenience wrapper: read audio_preprocessing defaults from config.json and call modeling.preprocess.

    Pass `repo_dir` if calling from outside the snapshot, otherwise auto-detected.
    Use `**overrides` to override any default (e.g., target_length=2048).
    """
    if repo_dir is None:
        repo_dir = os.path.dirname(os.path.abspath(__file__))
    import json as _json
    with open(os.path.join(repo_dir, "config.json")) as f:
        cfg = _json.load(f)
    defaults = cfg.get("audio_preprocessing", {})
    defaults.update(overrides)
    from modeling import preprocess as _preprocess
    return _preprocess(audio, **defaults)