import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before torch / any CUDA-touching import
import torch
import gradio as gr
import numpy as np
import librosa
from transformers import (
    Qwen2_5OmniForConditionalGeneration,
    Qwen2_5OmniProcessor,
)

MODEL_ID = "umd-zhou-lab/AudioRubrics"

processor = Qwen2_5OmniProcessor.from_pretrained(MODEL_ID)
model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda").eval()

THINK_OPEN = "<think>"
THINK_CLOSE = "</think>"
ANSWER_OPEN = "<answer>"
ANSWER_CLOSE = "</answer>"


@spaces.GPU(duration=60)
def answer_audio_question(
    audio_path: str,
    question: str,
    max_new_tokens: int = 768,
    temperature: float = 0.0,
    enable_thinking: bool = True,
    progress=gr.Progress(track_tqdm=True),
):
    """Answer an audio-grounded question using AudioRubrics (Qwen2.5-Omni-7B fine-tuned with evolving rubric rewards).

    Args:
        audio_path: Path to the input audio file (WAV/MP3/FLAC etc.).
        question: A text question about the audio content.
        max_new_tokens: Maximum number of tokens to generate.
        temperature: Sampling temperature (0.0 = greedy).
        enable_thinking: If True, the model reasons step-by-step before answering.
    """
    import re

    if audio_path is None:
        return "Please upload an audio file.", ""
    if not question.strip():
        return "Please enter a question about the audio.", ""

    # Build conversation messages matching the Qwen2.5-Omni chat template
    if enable_thinking:
        system_content = (
            "You are an expert audio understanding assistant. "
            "Listen carefully and answer questions about the audio. "
            "Always think step by step inside "
            + THINK_OPEN
            + " tags, "
            "then give the final answer inside "
            + ANSWER_OPEN
            + " tags."
        )
        user_text = (
            f"Listen to the audio carefully and answer the following question.\n\n"
            f"Question: {question}\n\n"
            "First, reason step by step inside "
            + THINK_OPEN
            + " ... "
            + THINK_CLOSE
            + " tags.\n"
            "Then output your final answer inside "
            + ANSWER_OPEN
            + " ... "
            + ANSWER_CLOSE
            + " tags."
        )
    else:
        system_content = (
            "You are an expert audio understanding assistant. "
            "Listen carefully and answer questions about the audio. "
            "Give the final answer directly."
        )
        user_text = (
            f"Listen to the audio carefully and answer the following question.\n\n"
            f"Question: {question}"
        )

    messages = [
        {"role": "system", "content": [{"type": "text", "text": system_content}]},
        {
            "role": "user",
            "content": [
                {"type": "audio", "audio": audio_path},
                {"type": "text", "text": user_text},
            ],
        },
    ]

    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )

    # Load audio as numpy array (resampled to 16kHz for Whisper feature extractor)
    audio_data, sr = librosa.load(audio_path, sr=16000, mono=True)

    inputs = processor(
        text=text,
        audio=audio_data,
        return_tensors="pt",
        padding=True,
    ).to("cuda").to(model.dtype)

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            generation_mode="text",
            thinker_max_new_tokens=max_new_tokens,
            thinker_temperature=temperature if temperature > 0 else 1.0,
            thinker_do_sample=temperature > 0,
        )

    # Strip the input tokens from the output
    input_len = inputs["input_ids"].shape[1]
    generated_ids = output_ids[0][input_len:]
    response = processor.decode(generated_ids, skip_special_tokens=True)

    # The think/answer tags may be decoded as special tokens (stripped by skip_special_tokens=True)
    # or as literal text. Try both approaches.
    # First try parsing with the known tag strings.
    think_pattern = re.escape(THINK_OPEN) + r"\s*(.*?)\s*" + re.escape(THINK_CLOSE)
    answer_pattern = re.escape(ANSWER_OPEN) + r"\s*(.*?)\s*" + re.escape(ANSWER_CLOSE)

    think_match = re.search(think_pattern, response, flags=re.DOTALL | re.IGNORECASE)
    answer_match = re.search(answer_pattern, response, flags=re.DOTALL | re.IGNORECASE)

    thinking_text = think_match.group(1).strip() if think_match else ""
    answer_text = answer_match.group(1).strip() if answer_match else ""

    # If tags not found with skip_special_tokens=True, try with False
    if not think_match and not answer_match:
        response_raw = processor.decode(generated_ids, skip_special_tokens=False)
        think_match = re.search(think_pattern, response_raw, flags=re.DOTALL | re.IGNORECASE)
        answer_match = re.search(answer_pattern, response_raw, flags=re.DOTALL | re.IGNORECASE)
        thinking_text = think_match.group(1).strip() if think_match else ""
        answer_text = answer_match.group(1).strip() if answer_match else ""
        if answer_text or thinking_text:
            response = response_raw

    # If still no tags found, return the full response as the answer
    if not answer_text and not thinking_text:
        answer_text = response.strip()
        thinking_text = ""
    elif not answer_text:
        answer_text = response.strip()

    # Format nicely
    if thinking_text:
        formatted_thinking = f"**Reasoning:**\n{thinking_text}"
    else:
        formatted_thinking = ""

    return answer_text, formatted_thinking


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# AudioRubrics: Audio Reasoning with Evolving Rubric Rewards\n"
            "Upload an audio clip and ask a question about it. The model reasons step-by-step "
            "about what it hears.\n\n"
            "Based on [Reinforcement Learning with Evolving Rubrics as Rewards for Audio Reasoning](https://huggingface.co/papers/2608.02831) | "
            "a Qwen2.5-Omni-7B model fine-tuned with GRPO using self-evolving, audio-grounded rubric rewards."
        )

        with gr.Row():
            with gr.Column(scale=1):
                audio_input = gr.Audio(
                    label="Audio Input",
                    type="filepath",
                    sources=["upload", "microphone"],
                )
                question_input = gr.Textbox(
                    label="Question",
                    placeholder="e.g., What sound do you hear in the audio?",
                    lines=3,
                )
                run_btn = gr.Button("Run", variant="primary")

                with gr.Accordion("Advanced settings", open=False):
                    max_tokens = gr.Slider(
                        label="Max new tokens",
                        minimum=64,
                        maximum=1024,
                        value=768,
                        step=64,
                    )
                    temp = gr.Slider(
                        label="Temperature",
                        minimum=0.0,
                        maximum=1.5,
                        value=0.0,
                        step=0.1,
                    )
                    think_checkbox = gr.Checkbox(
                        label="Enable step-by-step thinking",
                        value=True,
                    )

            with gr.Column(scale=1):
                answer_output = gr.Textbox(
                    label="Answer",
                    lines=4,
                    interactive=False,
                )
                thinking_output = gr.Markdown(
                    label="Reasoning",
                )

        with gr.Row():
            gr.Examples(
                examples=[
                    [
                        "examples/bird_chirp.wav",
                        "What animal is making the sound in the audio?\nChoices:\nA. dog\nB. bird\nC. cat\nD. frog",
                        512,
                        0.0,
                        True,
                    ],
                    [
                        "examples/metro_sound.wav",
                        "Where did the audio take place?\nChoices:\nA. train\nB. aquatic\nC. bus station\nD. Metro Station",
                        512,
                        0.0,
                        True,
                    ],
                    [
                        "examples/alarm_sound.wav",
                        "What's that noise?\nChoices:\nA. firecrackers\nB. Car sound\nC. tornado\nD. siren",
                        512,
                        0.0,
                        True,
                    ],
                ],
                inputs=[
                    audio_input,
                    question_input,
                    max_tokens,
                    temp,
                    think_checkbox,
                ],
                outputs=[answer_output, thinking_output],
                fn=answer_audio_question,
                cache_examples=True,
                cache_mode="lazy",
            )

        gr.Markdown(
            "\n---\n"
            "**Model:** [umd-zhou-lab/AudioRubrics](https://huggingface.co/umd-zhou-lab/AudioRubrics) | "
            "**Paper:** [arXiv:2608.02831](https://arxiv.org/abs/2608.02831) | "
            "**Code:** [GitHub](https://github.com/tianyi-lab/AudioRubrics)"
        )

    run_btn.click(
        fn=answer_audio_question,
        inputs=[audio_input, question_input, max_tokens, temp, think_checkbox],
        outputs=[answer_output, thinking_output],
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)