"""TextPro-SLM (rev 2): speech understanding through transcript tokens and compact prosody.

Demo of "Minimizing Modality Gap from the Input Side: Your Speech LLM Can Be a
Prosody-Aware Text LLM" (NeurIPS 2026). Speech is encoded by WhisperPro into a
transcript plus a compact prosody vector; a 2-layer MLP projector maps the
prosody into the Qwen2-7B embedding space, where it replaces the embedding of a
single trained marker token prepended to the transcript.

Two input modes are supported:
  1. Spoken QA       -- the whole question is spoken in the audio (VoiceBench
                        style spoken multiple-choice / open questions).
  2. Audio + text    -- the audio carries the speech signal, and a text
                        instruction/question is appended after the transcript
                        (AIR-Bench style emotion classification, paralinguistic QA).
"""

import time

import spaces  # must precede any torch/CUDA-touching import
import gradio as gr
import torch

from textpro_slm.model import SYSTEM_PROMPTS, TextProSLM

MODEL_ID = "qqjz/TextPro-SLM-7B-Global"

# AIR-Bench emotion instruction shipped with the TextPro-SLM repository.
EMOTION_INSTRUCTION = (
    "Choose the most suitable answer from options A, B, C, and D to respond the "
    "question in next line, you may only choose A or B or C or D.\n"
    "What emotion is at the forefront of the speaker's words?\n"
    "A. angry\nB. sad\nC. happy\nD. neutral"
)

MAX_AUDIO_SECONDS = 30  # original WhisperPro models accept at most 30 s

# Load once at module scope; ZeroGPU streams these weights to the worker.
model = TextProSLM(model=MODEL_ID, device="cuda")


@spaces.GPU(duration=25)
def respond(
    audio,
    instruction: str = "",
    task: str = "semantic",
    max_new_tokens: int = 512,
    temperature: float = 0.7,
    top_p: float = 0.9,
    greedy: bool = False,
):
    """Answer a spoken question, or an audio + text-instruction question.

    Args:
        audio: recorded or uploaded speech (16 kHz mono is resampled internally).
        instruction: optional text appended after the ASR transcript. Empty for
            spoken QA where the whole question is in the audio.
        task: system prompt preset -- "semantic" for spoken QA, "emotion" for
            AIR-Bench emotion classification, "audio" for other paralinguistic QA.
        max_new_tokens: maximum reply tokens (paper uses 512).
        temperature: sampling temperature (paper uses 0.7).
        top_p: nucleus sampling threshold (paper uses 0.9).
        greedy: decode greedily instead of the paper's sampling policy.

    Returns:
        (transcript, answer, stats) as plain strings.
    """
    if audio is None:
        raise gr.Error("Please record or upload some speech audio first.")
    path = audio if isinstance(audio, str) else audio[0] if isinstance(audio, (tuple, list)) else audio
    torch.manual_seed(42)
    task = task if task in SYSTEM_PROMPTS else "semantic"
    started = time.perf_counter()
    result = model.respond(
        audio=path,
        instruction=instruction.strip(),
        task=task,
        max_new_tokens=int(max_new_tokens),
        do_sample=not greedy,
        temperature=float(temperature),
        top_p=float(top_p),
    )
    elapsed = time.perf_counter() - started
    stats = (
        f"**Audio duration:** {result['duration']:.1f} s  |  "
        f"**Segments:** {result['num_segments']}  |  "
        f"**Inference:** {elapsed:.1f} s"
    )
    return result["transcript"], result["response"], stats


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

with gr.Blocks(title="TextPro-SLM") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            f"""
            # 🗣️ TextPro-SLM: Prosody-Aware Text LLM

            Speech understanding from the input side: WhisperPro turns speech into a
            **transcript** plus a compact **prosody vector**, a 2-layer MLP projector maps that
            vector into the Qwen2-7B embedding space, and the LLM answers with the prosody
            embedding injected at a single trained marker token. This Space runs
            [`{MODEL_ID}`](https://huggingface.co/{MODEL_ID}) (global prepending fusion).

            **Two modes** — leave *Text instruction* empty and speak the whole question
            (VoiceBench spoken QA), or provide audio + a text question (AIR-Bench style).

            *Paper: "Minimizing Modality Gap from the Input Side: Your Speech LLM Can Be a
            Prosody-Aware Text LLM" (NeurIPS 2026) · [Code](https://github.com/dreamtheater123/TextPro-SLM)*
            """
        )
        with gr.Row():
            with gr.Column(scale=1):
                audio_in = gr.Audio(
                    label="Speech audio (≤ 30 s)",
                    type="filepath",
                    sources=["microphone", "upload"],
                    format="wav",
                )
                instruction = gr.Textbox(
                    label="Text instruction (optional)",
                    placeholder="Empty = the question is spoken in the audio. "
                    'Otherwise, e.g. paste the emotion question.',
                    lines=4,
                )
                task = gr.Radio(
                    label="Task / system prompt",
                    choices=["semantic", "emotion", "audio"],
                    value="semantic",
                    info="semantic: spoken QA · emotion: AIR-Bench emotion · audio: paralinguistic QA",
                )
                run = gr.Button("Answer", variant="primary")
                with gr.Accordion("Advanced settings", open=False):
                    max_new_tokens = gr.Slider(
                        label="Max new tokens", minimum=64, maximum=1024, value=512, step=32
                    )
                    temperature = gr.Slider(
                        label="Temperature", minimum=0.0, maximum=1.5, value=0.7, step=0.05
                    )
                    top_p = gr.Slider(
                        label="Top-p", minimum=0.1, maximum=1.0, value=0.9, step=0.05
                    )
                    greedy = gr.Checkbox(label="Greedy decoding", value=False)
            with gr.Column(scale=1):
                transcript_out = gr.Textbox(label="WhisperPro transcript", lines=4)
                answer_out = gr.Markdown(label="Answer")
                stats_out = gr.Markdown(label="Stats")

        run.click(
            respond,
            inputs=[audio_in, instruction, task, max_new_tokens, temperature, top_p, greedy],
            outputs=[transcript_out, answer_out, stats_out],
            api_name="respond",
        )

        gr.Examples(
            examples=[
                ["examples/predators.wav", "", "semantic"],
                ["examples/fog.wav", "", "semantic"],
                ["examples/cactus.wav", "", "semantic"],
            ],
            inputs=[audio_in, instruction, task],
            outputs=[transcript_out, answer_out, stats_out],
            fn=respond,
            cache_examples=True,
            cache_mode="lazy",
            label="VoiceBench spoken multiple-choice questions (openbookqa)",
        )

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