import numpy as np
import torch

MERT_MODEL_NAME = "m-a-p/MERT-v1-95M"


class MERTEncoder:
    """Encodes audio chunks to 768-dim embeddings via HuggingFace MERT."""

    def __init__(self, model=None, processor=None):
        if model is None:
            from transformers import AutoModel
            self.model = AutoModel.from_pretrained(MERT_MODEL_NAME, trust_remote_code=True)
            self.model.eval()
        else:
            self.model = model

        if processor is None:
            from transformers import Wav2Vec2FeatureExtractor
            self.processor = Wav2Vec2FeatureExtractor.from_pretrained(
                MERT_MODEL_NAME, trust_remote_code=True
            )
        else:
            self.processor = processor

    def encode(self, y: np.ndarray, sr: int) -> np.ndarray:
        """Single chunk → 768-dim mean-pooled embedding."""
        target_sr = self.processor.sampling_rate
        if sr != target_sr:
            import librosa
            y = librosa.resample(y, orig_sr=sr, target_sr=target_sr)
            sr = target_sr
        inputs = self.processor(y.tolist(), sampling_rate=sr, return_tensors="pt", padding=True)
        with torch.no_grad():
            outputs = self.model(**inputs, return_dict=True)
        return outputs.last_hidden_state.mean(dim=1).squeeze(0).detach().numpy()

    def encode_batch(self, chunks: list[dict], sr: int) -> np.ndarray:
        """List of chunks → (n, 768) array."""
        return np.stack([self.encode(c["y"], sr) for c in chunks])