from typing import Union
import numpy as np
import torch
import librosa
from transformers import Wav2Vec2Processor, Wav2Vec2Model
from .base import BaseEncoder


class AudioEncoder(BaseEncoder):
    """Audio encoder using Wav2Vec2."""

    def __init__(self, model_name: str = "facebook/wav2vec2-base-960h", revision: str = "main"):
        self.model_name = model_name
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        # The revision parameter is configurable by callers to allow pinning to
        # a specific, trusted commit.  Bandit is informed that the revision is
        # intentionally forwarded.
        self.processor = Wav2Vec2Processor.from_pretrained(
            model_name, revision=revision)  # nosec B615
        self.model = Wav2Vec2Model.from_pretrained(
            model_name, revision=revision).to(self.device)  # nosec B615
        self.target_sr = 16000

    def encode(self, data: Union[str, np.ndarray]) -> np.ndarray:
        """
        Encode audio into a vector.

        Args:
            data: File path to audio or numpy array of audio samples

        Returns:
            Numpy array representing the embedding vector (mean pooled)
        """
        if isinstance(data, str):
            audio, sr = librosa.load(data, sr=self.target_sr)
        else:
            audio = data

        # Process audio
        inputs = self.processor(
            audio,
            sampling_rate=self.target_sr,
            return_tensors="pt",
            padding=True).to(
            self.device)

        with torch.no_grad():
            outputs = self.model(**inputs)

        # Mean pooling of last hidden state
        # outputs.last_hidden_state shape: (batch, sequence_length,
        # hidden_size)
        hidden_states = outputs.last_hidden_state
        embedding = torch.mean(hidden_states, dim=1)

        return embedding.cpu().numpy().flatten()

    @property
    def dimension(self) -> int:
        """Return the dimension of the embedding vector."""
        return self.model.config.hidden_size
