#!/usr/bin/env python3
# ------------------------------------------------------------
# Whisper Small LS-100 Demo (local or Hugging Face Spaces)
# Author: Murali Chandran (codeninja3d)
# ------------------------------------------------------------

import torch
import torchaudio
import gradio as gr
from transformers import WhisperProcessor, WhisperForConditionalGeneration

# -------------------------------
# Config
# -------------------------------
MODEL_ID = "codeninja3d/whisper-small-ls100"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# print(f"✅ Loading model from: {MODEL_ID}")
processor = WhisperProcessor.from_pretrained(MODEL_ID)
model = WhisperForConditionalGeneration.from_pretrained(MODEL_ID).to(DEVICE).eval()
# print(f"✅ #Params: {sum(p.numel() for p in model.parameters()) / 1e6:.2f}M")
# print("✅ Sample rate:", processor.feature_extractor.sampling_rate)


# -------------------------------
# Transcribe function
# -------------------------------
def transcribe(audio, language, task):
    if audio is None:
        return "⚠️ Please record or upload an audio clip."

    # Handle tuple (mic) or file path (upload)
    if isinstance(audio, tuple):
        sr, waveform_np = audio
        waveform = torch.tensor(waveform_np, dtype=torch.float32)
        if waveform.ndim > 1:                     # [samples, channels] → [1, samples]
            waveform = waveform.mean(dim=-1, keepdim=True).T
    else:
        waveform, sr = torchaudio.load(audio)

    # Ensure mono and correct shape [1, samples]
    if waveform.ndim > 1:
        waveform = waveform.mean(dim=0, keepdim=True)

    # Resample to 16 kHz
    if sr != 16000:
        waveform = torchaudio.functional.resample(waveform, sr, 16000)
        sr = 16000

     # ✅ Normalize amplitude to [-1, 1]
    waveform = waveform / waveform.abs().max()

    # ✅ Force 1D numpy array (no 2D padding issues)
    waveform_np = waveform.squeeze().numpy().astype("float32")

    # Preprocess
    inputs = processor.feature_extractor(
        waveform_np,
        sampling_rate=sr,
        return_tensors="pt"
    ).to(DEVICE)


    # print(f"[DEBUG] waveform shape: {waveform.shape}, dtype: {waveform.dtype}, sr={sr}")
    # print(f"[DEBUG] min={waveform.min().item():.3f}, max={waveform.max().item():.3f}, mean={waveform.mean().item():.3f}")

    # Decoder setup
    forced_decoder_ids = processor.get_decoder_prompt_ids(language="en", task="transcribe")

    # print("[DEBUG] forced_decoder_ids:", forced_decoder_ids)


    with torch.no_grad():
        pred_ids = model.generate(
            **inputs,
            forced_decoder_ids=forced_decoder_ids,
            max_new_tokens=256,
            num_beams=5,
            no_repeat_ngram_size=3,
            repetition_penalty=1.2,
            early_stopping=True
        )

    text = processor.tokenizer.batch_decode(pred_ids, skip_special_tokens=True)[0]
    return text.strip()




# -------------------------------
# Gradio Interface
# -------------------------------


iface = gr.Interface(
    fn=transcribe,
    inputs=[
        gr.Audio(type="numpy", label="🎙️ Record or Upload Audio"),
    ],
    outputs=gr.Textbox(label="📄 Transcription"),
    title="Whisper Small LS-100 (Fine-tuned)",
    description="Upload or record audio to test Murali's fine-tuned Whisper model on LibriSpeech-100.",
    theme="soft",
    allow_flagging="auto"
)

iface.launch(server_name="0.0.0.0", server_port=7860)