"""Whisper Small ASR — inference example using HuggingFace Transformers.

The optimized.pte and original.pte artifacts in the sibling pte_optimized/ and
pte_original/ directories are ExecuTorch-optimized models (8da8w INT8 dynamic
quantization) intended for on-device ARM inference (Android / Graviton).
Running those artifacts directly requires the ExecuTorch C++ runtime and a
specialized seq2seq runner — they are not suitable for a simple Python
``method.execute()`` call.

This script demonstrates equivalent inference using the HuggingFace Transformers
library, which is the recommended path for Python-based evaluation and prototyping.

Requirements:
    pip install transformers torch soundfile numpy
"""

import json
from pathlib import Path

import torch
import torch.nn.functional as F
from transformers import WhisperForConditionalGeneration, WhisperProcessor

# ── Configuration ──────────────────────────────────────────────────────────────
AUDIO_PATH = "sample_input.flac"
MODEL_NAME = "openai/whisper-small"
LANGUAGE = "en"
TASK = "transcribe"
MAX_NEW_TOKENS = 128
SAMPLE_RATE = 16000


# ── Audio Loading ──────────────────────────────────────────────────────────────
def load_audio(audio_path: str) -> tuple[torch.Tensor, int]:
    """Load audio file and return (waveform_1d_float32, sample_rate)."""
    try:
        import soundfile as sf

        audio_np, sr = sf.read(audio_path, dtype="float32")
        if audio_np.ndim == 2:
            audio_np = audio_np.mean(axis=1)
        import numpy as np

        return torch.from_numpy(audio_np.astype(np.float32)), int(sr)
    except ImportError:
        import torchaudio

        waveform, sr = torchaudio.load(audio_path)
        if waveform.shape[0] > 1:
            waveform = waveform.mean(dim=0)
        else:
            waveform = waveform.squeeze(0)
        return waveform.float(), int(sr)


# ── Resampling ─────────────────────────────────────────────────────────────────
def resample_to_16k(waveform: torch.Tensor, sample_rate: int) -> torch.Tensor:
    """Resample waveform to 16 kHz using linear interpolation."""
    if sample_rate == SAMPLE_RATE:
        return waveform
    new_len = int(round(len(waveform) * SAMPLE_RATE / sample_rate))
    return F.interpolate(
        waveform.view(1, 1, -1), size=new_len, mode="linear", align_corners=False
    ).view(-1)


# ── Preprocessing ──────────────────────────────────────────────────────────────
def preprocess(audio_path: str, processor: WhisperProcessor) -> torch.Tensor:
    """Load audio and extract 80-bin log-mel spectrogram features [1, 80, 3000]."""
    waveform, sr = load_audio(audio_path)
    waveform_16k = resample_to_16k(waveform, sr)
    features = processor.feature_extractor(
        waveform_16k.numpy(), sampling_rate=SAMPLE_RATE, return_tensors="pt"
    ).input_features
    return features  # [1, 80, 3000]


# ── Inference ──────────────────────────────────────────────────────────────────
def transcribe(
    audio_path: str,
    model: WhisperForConditionalGeneration,
    processor: WhisperProcessor,
) -> str:
    """Run Whisper inference and return transcribed text."""
    features = preprocess(audio_path, processor)
    with torch.no_grad():
        output_ids = model.generate(
            features,
            language=LANGUAGE,
            task=TASK,
            max_new_tokens=MAX_NEW_TOKENS,
        )
    return processor.tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0].strip()


# ── Save Results ───────────────────────────────────────────────────────────────
def save_results(audio_path: str, text: str) -> None:
    """Save transcription result to JSON in the same directory as this script."""
    out_dir = Path(__file__).parent
    result = {
        "audio_file": str(Path(audio_path).name),
        "transcription": text,
        "model": MODEL_NAME,
        "language": LANGUAGE,
        "task": TASK,
    }
    out_path = out_dir / "transcription.json"
    with open(out_path, "w") as f:
        json.dump(result, f, indent=2)
    print(f"Transcription saved to: {out_path}")


# ── Main ───────────────────────────────────────────────────────────────────────
def main() -> None:
    print(f"Loading model: {MODEL_NAME}")
    processor = WhisperProcessor.from_pretrained(MODEL_NAME)
    model = WhisperForConditionalGeneration.from_pretrained(MODEL_NAME)
    model.eval()

    audio_path = str(Path(__file__).parent / AUDIO_PATH)
    print(f"Transcribing: {audio_path}")
    text = transcribe(audio_path, model, processor)

    print(f"\nTranscription: {text}")
    save_results(audio_path, text)


if __name__ == "__main__":
    main()