import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("IMAGE_MAX_TOKEN_NUM", "1024")
os.environ.setdefault("VIDEO_MAX_TOKEN_NUM", "128")
os.environ.setdefault("FPS_MAX_FRAMES", "16")

import spaces  # MUST come before torch / any CUDA-touching import
import torch
import gradio as gr
from transformers import Qwen3_5ForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info

MODEL_ID = "Jiaha0Hu4ng/OneEmo"

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

processor = AutoProcessor.from_pretrained(MODEL_ID)

# ── Task prompts (ported 1:1 from the OneEmo repo prompts/ directory) ───────

MER_PROMPT = "[MER]<video>Please identify the emotions of the person in the video. Answer with a single emotion."
OVMER_PROMPT = "[OVMER]<video>What feelings does the character show? List them briefly."
MSA_PROMPT = "[MSA]<video>Please analyze the sentiment of the characters in the video and label them as positive, neutral or negative."
MHD_PROMPT = "[MHD]<video>Based on the context of the speaker and visual cues, determine whether there is a humorous expression, answer with yes or no."
MSD_PROMPT = "[MSD]<video>Please judge based on the context whether the speaker in this round is being sarcastic, answer with yes or no."
ERG_PROMPT = (
    "[ERG]<video>You are an empathetic listener, your goal is to understand the user's emotions "
    "and intentions, and respond or comfort them with appropriate language that helps them feel "
    "understood and cared for.\n Avoid rushing into your response; instead, carefully engage in "
    "a step-by-step, in-depth analysis before providing an answer.\n Please analyze using Chain "
    "of Empathy (Firstly, Event scenario:Reflect on the event scenarios that arise from the "
    "ongoing dialogue. Secondly, User's emotion:Analyze both the implicit and explicit emotions "
    "conveyed by the user. Thirdly, the emotion cause:Infer the underlying reasons for the "
    "user's emotions. Fourthly, determine the goal of your response in this particular instance, "
    "such as alleviating anxiety, offering reassurance, or expressing understanding.) in imd "
    "monospace tags and with Line break, then provide your empathetic response."
)

TASK_PROMPTS = {
    "MER — Basic Emotion Recognition": MER_PROMPT,
    "OVMER — Open-Vocabulary Emotion Recognition": OVMER_PROMPT,
    "MSA — Multimodal Sentiment Analysis": MSA_PROMPT,
    "MHD — Humor Detection": MHD_PROMPT,
    "MSD — Sarcasm Detection": MSD_PROMPT,
    "ERG — Empathetic Response Generation": ERG_PROMPT,
}


@spaces.GPU(duration=60)
def analyze_emotion(
    video,
    task: str,
    transcript: str = "",
    enable_thinking: bool = True,
    max_tokens: int = 2048,
    temperature: float = 0.7,
    progress=gr.Progress(track_tqdm=True),
):
    """Analyze emotions in a video using OneEmo, a unified multimodal reasoning model.

    Args:
        video: Input video file path.
        task: Emotion analysis task type.
        transcript: Optional transcript of speech in the video.
        enable_thinking: Whether to enable chain-of-thought reasoning.
        max_tokens: Maximum number of new tokens to generate.
        temperature: Sampling temperature.
    """
    if video is None:
        return "Please upload a video file."

    base_prompt = TASK_PROMPTS[task]
    if transcript.strip():
        if task.startswith("ERG"):
            prompt = f"{base_prompt}\n{transcript}"
        else:
            prompt = f"{base_prompt}\nHere is what the character says: {transcript}"
    else:
        prompt = base_prompt

    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "video",
                    "video": video,
                    "max_pixels": 360 * 420,
                    "fps": 2.0,
                },
                {
                    "type": "text",
                    "text": prompt,
                },
            ],
        }
    ]

    text = processor.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=enable_thinking,
    )
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    ).to("cuda")

    with torch.no_grad():
        generated_ids = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            temperature=temperature,
            top_p=0.9,
            top_k=50,
            do_sample=temperature > 0,
        )

    generated_ids_trimmed = [
        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    output_text = processor.batch_decode(
        generated_ids_trimmed,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
    )
    return output_text[0]


# ── UI ───────────────────────────────────────────────────────────────────

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

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # OneEmo: Unified Multimodal Reasoning for Emotion Perception

        Upload a video and select an emotion analysis task. OneEmo is a unified
        multimodal reasoning model that supports emotion recognition, sentiment
        analysis, humor/sarcasm detection, and empathetic response generation.

        [Paper](https://arxiv.org/abs/2608.06013) · [Model](https://huggingface.co/Jiaha0Hu4ng/OneEmo) · [Code](https://github.com/waHAHJIAHAO/OneEmo)
        """
    )

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            with gr.Column(scale=1):
                video_input = gr.Video(label="Input Video")
                task_dropdown = gr.Dropdown(
                    choices=list(TASK_PROMPTS.keys()),
                    value="MER — Basic Emotion Recognition",
                    label="Task",
                    info="Choose the emotion analysis task",
                )
                transcript_input = gr.Textbox(
                    label="Transcript (optional)",
                    placeholder="Optional transcript of speech in the video…",
                    lines=2,
                )
                run_btn = gr.Button("Analyze", variant="primary")

            with gr.Column(scale=1):
                output = gr.Markdown(
                    label="Result",
                    value="Upload a video and click **Analyze** to see the emotion analysis.",
                )

        with gr.Accordion("Advanced settings", open=False):
            enable_thinking = gr.Checkbox(
                label="Enable thinking (chain-of-thought)",
                value=True,
            )
            max_tokens = gr.Slider(
                label="Max new tokens",
                minimum=256,
                maximum=4096,
                value=2048,
                step=256,
            )
            temperature = gr.Slider(
                label="Temperature",
                minimum=0.0,
                maximum=1.5,
                value=0.7,
                step=0.1,
            )

        gr.Examples(
            examples=[
                ["examples/man_laughing.mp4", "MER — Basic Emotion Recognition", ""],
                ["examples/man_sad.mp4", "OVMER — Open-Vocabulary Emotion Recognition", ""],
                ["examples/2_women_arguing.mp4", "MSA — Multimodal Sentiment Analysis", ""],
                ["examples/man_smiling_studio.mp4", "ERG — Empathetic Response Generation", ""],
            ],
            inputs=[video_input, task_dropdown, transcript_input],
            outputs=output,
            fn=analyze_emotion,
            cache_examples=True,
            cache_mode="lazy",
        )

    run_btn.click(
        fn=analyze_emotion,
        inputs=[video_input, task_dropdown, transcript_input, enable_thinking, max_tokens, temperature],
        outputs=output,
        api_name="analyze",
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)