#!/usr/bin/env python3
"""
Extract Whisper Encoder hidden states from an audio file.
Loads openai/whisper-large-v3-turbo from HuggingFace transformers
and produces (seq_len, 1280) hidden states.
"""

import argparse
import logging
import sys

import librosa
import numpy as np
import torch

logging.basicConfig(level=logging.INFO, format="%(levelname)s:%(name)s:%(message)s")
logger = logging.getLogger(__name__)


class WhisperEncoder:
    """Wrapper around Whisper-large-v3-turbo encoder for feature extraction."""

    MODEL_ID = "openai/whisper-large-v3-turbo"

    def __init__(self, device: str = "cpu"):
        self.device = device
        logger.info(f"Loading Whisper encoder from {self.MODEL_ID}...")
        from transformers import WhisperModel

        self.model = WhisperModel.from_pretrained(
            self.MODEL_ID,
            torch_dtype=torch.float32,
            low_cpu_mem_usage=True,
        ).to(device)
        self.model.eval()
        # We only need the encoder
        self.encoder = self.model.encoder
        # Freeze encoder
        for param in self.encoder.parameters():
            param.requires_grad = False
        logger.info("Whisper encoder loaded successfully.")

    @torch.no_grad()
    def encode(self, audio: np.ndarray, sampling_rate: int = 16000) -> np.ndarray:
        """
        Encode audio waveform to hidden states.

        Args:
            audio: 1D numpy array of audio samples
            sampling_rate: Sample rate of input audio (will resample to 16000 if needed)

        Returns:
            numpy array of shape (seq_len, 1280) with encoder hidden states
        """
        # Resample to 16000 if needed
        if sampling_rate != 16000:
            audio = librosa.resample(audio, orig_sr=sampling_rate, target_sr=16000)
            sampling_rate = 16000

        # Whisper expects 30-second chunks (80 mel bins, 3000 time steps)
        # We'll use the feature extractor to convert audio to log-mel spectrogram
        from transformers import WhisperFeatureExtractor

        feature_extractor = WhisperFeatureExtractor.from_pretrained(self.MODEL_ID)
        # Convert audio to input features
        input_features = feature_extractor(
            audio, sampling_rate=sampling_rate, return_tensors="pt"
        ).input_features  # shape: (1, 80, 3000)
        input_features = input_features.to(self.device)

        # Run encoder
        encoder_outputs = self.encoder(input_features)
        # encoder_outputs.last_hidden_state shape: (1, seq_len, 1280)
        hidden_states = encoder_outputs.last_hidden_state.squeeze(0).cpu().numpy()
        # hidden_states shape: (seq_len, 1280)
        return hidden_states


def main():
    parser = argparse.ArgumentParser(
        description="Extract Whisper encoder hidden states from audio"
    )
    parser.add_argument("audio_path", help="Path to audio file")
    parser.add_argument(
        "--device",
        default="cpu",
        choices=["cpu", "cuda"],
        help="Device to run encoder on",
    )
    args = parser.parse_args()

    # Load audio
    logger.info(f"Loading audio: {args.audio_path}")
    audio, sr = librosa.load(args.audio_path, sr=16000, mono=True)
    logger.info(f"Audio loaded: {len(audio)} samples, {sr}Hz, {len(audio)/sr:.2f}s")

    # Initialize encoder
    encoder = WhisperEncoder(device=args.device)

    # Encode
    hidden_states = encoder.encode(audio, sampling_rate=sr)
    logger.info(f"Hidden states shape: {hidden_states.shape}")
    logger.info(f"Hidden states dtype: {hidden_states.dtype}")
    logger.info(f"Hidden states range: [{hidden_states.min():.4f}, {hidden_states.max():.4f}]")

    print(f"\nHidden states shape: {hidden_states.shape}")
    print(f"Sequence length: {hidden_states.shape[0]}")
    print(f"Feature dimension: {hidden_states.shape[1]}")


if __name__ == "__main__":
    main()