"""
Shared AST (Audio Spectrogram Transformer) encoder — used IDENTICALLY in training
and in the Hugging Face Space so an uploaded clip is embedded exactly as the
training clips were.

We take a pretrained AST (MIT/ast-finetuned-audioset-10-10-0.4593), a state-of-the-art
audio transformer, and use it as a frozen feature extractor: each ~2-minute hive
recording is summarised by the MEAN of the AST pooled embeddings over N evenly
spaced 10.24 s windows -> a single 768-d audio vector.
"""
import numpy as np

AST_MODEL = "MIT/ast-finetuned-audioset-10-10-0.4593"
SR = 16000
WIN_SAMPLES = int(10.24 * SR)   # 163840 -> AST's native 1024 frames
N_WINDOWS = 4
EMB_DIM = 768


def load_ast(device=None):
    import torch
    from transformers import ASTFeatureExtractor, ASTModel
    if device is None:
        device = "mps" if torch.backends.mps.is_available() else (
            "cuda" if torch.cuda.is_available() else "cpu")
    fe = ASTFeatureExtractor.from_pretrained(AST_MODEL)
    model = ASTModel.from_pretrained(AST_MODEL).to(device).eval()
    return fe, model, device


def _windows(y):
    """Return a list of fixed-length windows covering the clip."""
    n = len(y)
    if n <= WIN_SAMPLES:
        return [np.pad(y, (0, WIN_SAMPLES - n))]
    centers = np.linspace(WIN_SAMPLES // 2, n - WIN_SAMPLES // 2, N_WINDOWS)
    segs = []
    for c in centers:
        s = int(c - WIN_SAMPLES // 2)
        segs.append(y[s:s + WIN_SAMPLES])
    return segs


def encode_waveform(y, sr, fe, model, device):
    """waveform -> 768-d AST embedding (mean over windows)."""
    import torch
    import librosa
    if sr != SR:
        y = librosa.resample(np.asarray(y, dtype=np.float32), orig_sr=sr, target_sr=SR)
    y = np.asarray(y, dtype=np.float32)
    if y.ndim > 1:
        y = y.mean(axis=0)
    segs = _windows(y)
    with torch.no_grad():
        inp = fe(segs, sampling_rate=SR, return_tensors="pt")
        out = model(inp["input_values"].to(device))
        emb = out.pooler_output.mean(dim=0).cpu().numpy()   # (768,)
    return emb.astype(np.float32)


def encode_file(path, fe, model, device):
    import librosa
    y, sr = librosa.load(path, sr=SR, mono=True)
    return encode_waveform(y, sr, fe, model, device)