"""
DeepMed-R1: Medical Reasoning AI Demo
"""
import os
import gradio as gr
from huggingface_hub import InferenceClient

HF_TOKEN = os.environ.get("HF_TOKEN")
client = InferenceClient(token=HF_TOKEN)

MODEL = "Qwen/Qwen2.5-72B-Instruct"

SYSTEM_PROMPT = """You are DeepMed-R1, a medical reasoning AI trained with GRPO on AMD MI300X.

For EVERY question you MUST follow this format exactly:

**Step 1 - Key Findings:**
[List the important clinical features from the question]

**Step 2 - Differential Diagnosis:**
[List 3-4 possible diagnoses ranked by likelihood]

**Step 3 - Pathophysiology:**
[Explain WHY the most likely diagnosis causes these symptoms]

**Step 4 - Option Analysis:**
- A: [evaluate this option]
- B: [evaluate this option]
- C: [evaluate this option]
- D: [evaluate this option]

**Step 5 - Final Answer:**
The answer is **X** because [one sentence justification].

IMPORTANT: Always show ALL 5 steps. Never skip steps. Be detailed and specific."""


EXAMPLES = [
    ["A 65-year-old male with hypertension presents with sudden 'worst headache of my life,' neck stiffness, photophobia. BP 180/100. Most likely diagnosis?\nA. Migraine\nB. Subarachnoid hemorrhage\nC. Meningitis\nD. Tension headache"],
    ["A 28-year-old woman: fatigue, weight gain, cold intolerance. TSH 12 mIU/L, Free T4 0.5 ng/dL. Initial treatment?\nA. Levothyroxine\nB. Liothyronine\nC. Methimazole\nD. Radioactive iodine"],
    ["3-month-old with projectile non-bilious vomiting, olive-shaped RUQ mass, metabolic alkalosis. Diagnosis?\nA. Pyloric stenosis\nB. Intussusception\nC. Malrotation\nD. Hirschsprung disease"],
    ["55-year-old diabetic: RUQ pain, fever 39.2C, jaundice (Charcot triad). WBC 18K, bilirubin 5.2, CBD stone. Next step?\nA. Cholecystectomy\nB. ERCP with sphincterotomy\nC. MRCP\nD. PTC"],
    ["22-year-old post-MVC: left chest pain, absent breath sounds left, trachea deviated right, JVD, BP 80/50. Immediate management?\nA. Chest X-ray\nB. CT chest\nC. Needle decompression left chest\nD. Intubation"],
]


def respond(message, history):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for h in history:
        if h[0]: messages.append({"role": "user", "content": h[0]})
        if h[1]: messages.append({"role": "assistant", "content": h[1]})
    messages.append({"role": "user", "content": message})

    response = ""
    try:
        stream = client.chat_completion(
            model=MODEL,
            messages=messages,
            max_tokens=4000,
            temperature=0.4,
            top_p=0.95,
            stream=True,
        )
        for chunk in stream:
            if chunk.choices and len(chunk.choices) > 0:
                delta = chunk.choices[0].delta
                if hasattr(delta, "content") and delta.content:
                    response += delta.content
                    yield response
        if not response:
            yield "No response received. Please try again."
    except Exception as e:
        yield f"⚠️ Error: {str(e)}"


with gr.Blocks(title="DeepMed-R1", theme=gr.themes.Soft(primary_hue="blue")) as demo:
    gr.Markdown("""
# 🏥 DeepMed-R1: Medical Reasoning AI

**Systematic clinical reasoning powered by GRPO + Multi-Objective Clinical Rewards**

Built for **AMD Developer Hackathon 2026** — Track 2: Fine-Tuning on AMD GPUs

> ⚠️ Research demo only. Not for clinical use.
---
    """)

    chatbot = gr.ChatInterface(
        fn=respond,
        examples=EXAMPLES,
        title="",
    )

    gr.Markdown("""
---
### 🔬 Architecture

| Component | Detail |
|-----------|--------|
| Base Model | Qwen3 (RL-trained reasoning) |
| Training | GRPO + DAPO loss + CRPO rewards |
| Innovations | iGRPO (Feb 2026) + AERO (Feb 2026) + Curriculum (Mar 2026) |
| Hardware | AMD MI300X (192GB HBM3) |
| Reward | Accuracy (2.0) + Reasoning (1.0) + Consistency (0.5) + Length (0.3) |

📄 [Code & Training](https://huggingface.co/asusf15/DeepMed-R1) |
📚 [Gazal-R1](https://arxiv.org/abs/2506.21594) |
📚 [Clinical-R1](https://arxiv.org/abs/2512.00601) |
📚 [iGRPO](https://arxiv.org/abs/2602.09000)
    """)

if __name__ == "__main__":
    demo.launch()