import gradio as gr
import torch
import librosa
import re
import math
from transformers import AutoProcessor, SpeechT5ForSpeechToText
from jiwer import wer, cer

MODEL_ID = "nadaeltanekhy/artst-mgb3-finetuned"
DEVICE   = "cuda" if torch.cuda.is_available() else "cpu"

print(f"Loading model on {DEVICE}...")
processor = AutoProcessor.from_pretrained("MBZUAI/artst_asr_v3")   # ← load processor from original model
model     = SpeechT5ForSpeechToText.from_pretrained(MODEL_ID).to(DEVICE)
model.eval()

def normalize_arabic(text):
    text = re.sub(r'[\u064B-\u065F]', '', text)
    text = re.sub(r'\u0640', '', text)
    text = re.sub(r'[\u0622\u0623\u0625]', '\u0627', text)
    text = re.sub(r'\u0629', '\u0647', text)
    text = re.sub(r'[^\u0600-\u06FF\s]', '', text)
    return ' '.join(text.split())

def transcribe(audio_path, ground_truth=""):
    if audio_path is None:
        return "Please upload or record audio.", "N/A", "N/A"

    speech, _ = librosa.load(audio_path, sr=16000)
    inputs = processor(audio=speech, sampling_rate=16000, return_tensors="pt").input_values.to(DEVICE)

    with torch.no_grad():
        outputs = model.generate(
            inputs,
            max_length=200,
            num_beams=5,
            return_dict_in_generate=True,
            output_scores=True
        )

    text = processor.batch_decode(outputs.sequences, skip_special_tokens=True)[0]

    confidence = "N/A"
    if hasattr(outputs, "sequences_scores") and outputs.sequences_scores is not None:
        confidence = f"{min(1.0, math.exp(outputs.sequences_scores[0].item()))*100:.1f}%"

    metrics = "No ground truth provided."
    if ground_truth.strip():
        ref = normalize_arabic(ground_truth)
        hyp = normalize_arabic(text)
        metrics = f"WER: {wer(ref, hyp):.4f}  |  CER: {cer(ref, hyp):.4f}"

    return text, confidence, metrics

with gr.Blocks(title="Arabic ASR — ArTST v3") as demo:
    gr.Markdown("""
    # 🎙️ Arabic ASR — ArTST v3
    Fine-tuned on **MGB-3 Egyptian Arabic** broadcast speech.  
    Upload a `.wav` file or record directly from your microphone.
    """)
    with gr.Row():
        with gr.Column():
            audio  = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Arabic Speech Input")
            truth  = gr.Textbox(label="Ground Truth (optional)", placeholder="أدخل النص الصحيح هنا للحصول على WER/CER")
            btn    = gr.Button("Transcribe 🎯", variant="primary")
        with gr.Column():
            out_text = gr.Textbox(label="Transcription", lines=4, rtl=True)
            out_conf = gr.Textbox(label="Confidence Score")
            out_wer  = gr.Textbox(label="WER / CER")
    btn.click(transcribe, [audio, truth], [out_text, out_conf, out_wer])

demo.launch()