app.py

import gradio as gr
import json
import os
import re
import requests
from datetime import datetime, date
from typing import Optional

# ── Synthetic patient database ─────────────────────────────────────────────────

PATIENT_RECORDS = {
    "P001": {
        "name": "Jana Nováková",
        "age": 67,
        "gender": "F",
        "history": [
            {"date": "2024-05-10", "event": "CABG bypass operace srdce"},
            {"date": "2024-11-02", "event": "Kontrola kardiologie – stabilní stav"},
        ],
        "medications": ["Warfarin 5mg", "Metoprolol 50mg", "Atorvastatin 20mg"],
        "allergies": ["Penicilin"],
        "blood_type": "A+",
    },
    "P002": {
        "name": "Tomáš Kratochvíl",
        "age": 34,
        "gender": "M",
        "history": [
            {"date": "2025-01-15", "event": "Zlomenina levého zápěstí – ošetřena"},
        ],
        "medications": [],
        "allergies": [],
        "blood_type": "0-",
    },
    "P003": {
        "name": "Marta Horáčková",
        "age": 52,
        "gender": "F",
        "history": [
            {"date": "2025-04-20", "event": "Appendektomie"},
            {"date": "2025-04-28", "event": "Kontrola po operaci – hojení v pořádku"},
        ],
        "medications": ["Metformin 1000mg", "Ramipril 5mg"],
        "allergies": ["Ibuprofen"],
        "blood_type": "B+",
    },
    "P004": {
        "name": "Karel Dvořák",
        "age": 71,
        "gender": "M",
        "history": [
            {"date": "2025-12-01", "event": "Diabetes mellitus 2. typu – dispenzarizace"},
        ],
        "medications": ["Inzulin Lantus", "Aspirin 100mg"],
        "allergies": [],
        "blood_type": "A-",
    },
}

EPIDEMIOLOGY = {
    "current_date": str(date.today()),
    "alerts": [
        "Zvýšený výskyt chřipky A (H1N1) – týden 24/2026, kraj Plzeňský: +34 % oproti průměru",
        "Lokální gastroenteritida – hlášeno 12 případů z Plzně-sever, pravděpodobně norovirus",
    ],
    "season": "léto",
}

# ── Patient queue (in-memory) ──────────────────────────────────────────────────

patient_queue = []

# ── OpenRouter ─────────────────────────────────────────────────────────────────
# Free models (no credit card required on openrouter.ai):
#   meta-llama/llama-3.3-70b-instruct:free   ← doporučeno, nejlepší kvalita zdarma
#   google/gemma-3-27b-it:free
#   mistralai/mistral-7b-instruct:free        ← záloha pokud 70B má frontu

OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"
FREE_MODEL = "meta-llama/llama-3.3-70b-instruct:free"

TRIAGE_SYSTEM_PROMPT = """Jsi specializovaný triážní agent pro urgentní příjem v nemocnici.
Pracuješ podle platného českého triážního standardu Manchester Triage System (MTS).

PRIORITY (Manchester Triage System):
- P1 RESUSCITACE (červená): okamžité ohrožení života, čas do ošetření: ihned
- P2 URGENTNÍ (oranžová): potenciální ohrožení života, čas do ošetření: max 10 minut
- P3 NALÉHAVÉ (žlutá): závažný stav, čas do ošetření: max 60 minut
- P4 MÉNĚ NALÉHAVÉ (zelená): stabilní stav, čas do ošetření: max 120 minut
- P5 NEURGENTNÍ (modrá): plánovaná péče, čas do ošetření: max 240 minut

KLÍČOVÁ PRAVIDLA:
1. Vždy upozorni na "red flags" – příznaky skrytě nebezpečných stavů (atypický infarkt, disekce aorty, SAK, apod.)
2. U žen nad 40 let, diabetiků a starších pacientů VŽDY zvažuj atypické prezentace infarktu
3. Pokud jsi nejistý, raději eskaluj na vyšší prioritu (undertriage je nebezpečnější než overtriage)
4. Vysvětli PROČ jsi přiřadil danou prioritu – lékař musí rozumět tvému rozhodnutí
5. Nikdy nerozhoduješ sám – tvůj výstup je PODKLAD pro lékaře, ne konečné rozhodnutí

DŮLEŽITÉ: Odpovídej POUZE čistým JSON objektem. Žádný text před ani za. Žádné ```json bloky.

Přesný formát výstupu:
{
  "priority": "P2",
  "priority_label": "URGENTNÍ",
  "priority_color": "orange",
  "confidence": "vysoká",
  "main_concern": "Jedna věta s hlavním důvodem priority",
  "red_flags": ["Red flag 1", "Red flag 2"],
  "differential_diagnoses": [
    {"diagnosis": "Název diagnózy", "probability": "vysoká", "reasoning": "Stručné zdůvodnění"},
    {"diagnosis": "Název diagnózy", "probability": "střední", "reasoning": "Stručné zdůvodnění"}
  ],
  "recommended_immediate_actions": ["Akce 1", "Akce 2", "Akce 3"],
  "escalate_to_doctor": true,
  "escalation_reason": "Důvod proč přivolat lékaře, nebo null",
  "data_sources_used": ["příznaky od recepce", "kartotéka pacienta", "epidemiologický kontext"]
}"""


def call_openrouter(user_message: str) -> dict:
    api_key = os.environ.get("OPENROUTER_API_KEY")
    if not api_key:
        raise ValueError(
            "OPENROUTER_API_KEY není nastaven. Jdi na Settings → Secrets a přidej ho."
        )

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://huggingface.co",
        "X-Title": "TriazAI",
    }

    payload = {
        "model": FREE_MODEL,
        "messages": [
            {"role": "system", "content": TRIAGE_SYSTEM_PROMPT},
            {"role": "user", "content": user_message},
        ],
        "max_tokens": 1200,
        "temperature": 0.1,
    }

    resp = requests.post(OPENROUTER_URL, headers=headers, json=payload, timeout=60)
    resp.raise_for_status()
    raw = resp.json()["choices"][0]["message"]["content"].strip()

    # Strip markdown fences if model ignores instructions
    raw = re.sub(r"^```(?:json)?\s*", "", raw)
    raw = re.sub(r"\s*```$", "", raw)
    raw = raw.strip()

    # Find JSON object in response (fallback if model adds prose)
    match = re.search(r"\{[\s\S]*\}", raw)
    if match:
        raw = match.group(0)

    return json.loads(raw)


def run_triage_agent(
    patient_data: dict,
    medical_record: Optional[dict],
    symptoms: str,
    vitals: dict,
) -> dict:
    record_text = (
        json.dumps(medical_record, ensure_ascii=False, indent=2)
        if medical_record
        else "Kartotéka nedostupná – nový pacient bez záznamu"
    )

    user_message = f"""NOVÝ PACIENT K TRIÁŽI:

Základní údaje:
- Jméno: {patient_data.get("name", "Neznámý")}
- Věk: {patient_data.get("age", "Neznámý")} let
- Pohlaví: {patient_data.get("gender", "Neznámé")}
- Čas příchodu: {datetime.now().strftime("%H:%M")}

Příznaky (popsala recepce):
{symptoms}

Vitální funkce:
{json.dumps(vitals, ensure_ascii=False, indent=2)}

Zdravotní kartotéka:
{record_text}

Epidemiologická situace:
{json.dumps(EPIDEMIOLOGY, ensure_ascii=False, indent=2)}

Proveď triáž. Vrať POUZE JSON objekt."""

    return call_openrouter(user_message)


# ── API functions ──────────────────────────────────────────────────────────────

def lookup_patient(patient_id: str):
    if not patient_id or not patient_id.strip():
        return "{}", "⚠️ Zadejte ID pacienta"
    record = PATIENT_RECORDS.get(patient_id.strip().upper())
    if record:
        return json.dumps(record, ensure_ascii=False), f"✅ Nalezeno: {record['name']}"
    return "{}", "⚠️ Pacient nenalezen – nový pacient nebo chybné ID"


def add_patient_to_queue(
    patient_id, name, age, gender, symptoms,
    bp_sys, bp_dia, hr, spo2, temp, gcs,
):
    if not symptoms or not symptoms.strip():
        return "❌ Příznaky jsou povinné.", json.dumps(patient_queue, ensure_ascii=False)

    pid_clean = (patient_id or "").strip().upper()
    medical_record = PATIENT_RECORDS.get(pid_clean) if pid_clean else None

    patient_data = {
        "id": pid_clean if pid_clean else f"TEMP_{len(patient_queue)+1:03d}",
        "name": (name or "Neznámý").strip(),
        "age": int(age) if age else 0,
        "gender": gender or "X",
        "arrival_time": datetime.now().strftime("%H:%M"),
    }

    vitals = {
        "tlak": f"{int(bp_sys or 0)}/{int(bp_dia or 0)} mmHg",
        "puls": f"{int(hr or 0)} /min",
        "SpO2": f"{float(spo2 or 98)} %",
        "teplota": f"{float(temp or 36.8)} °C",
        "GCS": int(gcs or 15),
    }

    try:
        triage_result = run_triage_agent(patient_data, medical_record, symptoms, vitals)
    except Exception as e:
        return f"❌ Chyba triáže: {str(e)}", json.dumps(patient_queue, ensure_ascii=False)

    entry = {
        **patient_data,
        "symptoms": symptoms,
        "vitals": vitals,
        "medical_record": medical_record,
        "triage": triage_result,
    }
    patient_queue.append(entry)

    order = {"P1": 0, "P2": 1, "P3": 2, "P4": 3, "P5": 4}
    patient_queue.sort(key=lambda x: order.get(x["triage"].get("priority", "P5"), 4))

    t = triage_result
    msg = f"✅ {patient_data['name']} přidán – {t.get('priority','?')} {t.get('priority_label','')}"
    return msg, json.dumps(patient_queue, ensure_ascii=False)


def get_queue_json():
    return json.dumps(patient_queue, ensure_ascii=False)


def clear_patient(patient_id: str):
    global patient_queue
    pid = (patient_id or "").strip().upper()
    before = len(patient_queue)
    patient_queue = [p for p in patient_queue if p["id"] != pid]
    if len(patient_queue) < before:
        return f"✅ Pacient {pid} odebrán z fronty.", json.dumps(patient_queue, ensure_ascii=False)
    return f"⚠️ ID {pid} nenalezeno.", json.dumps(patient_queue, ensure_ascii=False)


# ── Gradio UI ──────────────────────────────────────────────────────────────────

with gr.Blocks(title="TriážAI API", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
# TriážAI – Backend API
Backend pro HTML frontend. Volá **OpenRouter** (free tier) – model: `meta-llama/llama-3.3-70b-instruct:free`

> **Setup:** Přidej `OPENROUTER_API_KEY` do **Settings → Variables and Secrets** tohoto Space.
> Klíč získáš zdarma na [openrouter.ai](https://openrouter.ai) po registraci.

Přímé API endpointy pro HTML frontend:
- `POST /call/lookup_patient` 
- `POST /call/add_patient_to_queue`
- `POST /call/get_queue_json`
- `POST /call/clear_patient`
""")

    with gr.Tab("🔍 Vyhledat pacienta"):
        pid_in = gr.Textbox(label="ID pacienta", placeholder="P001, P002, P003 nebo P004")
        lookup_btn = gr.Button("Vyhledat", variant="secondary")
        with gr.Row():
            stat_out = gr.Textbox(label="Status")
            rec_out = gr.Textbox(label="Záznam (JSON)", lines=12)
        lookup_btn.click(lookup_patient, inputs=[pid_in], outputs=[rec_out, stat_out])

    with gr.Tab("🚨 Přidat pacienta (test)"):
        with gr.Row():
            pid2 = gr.Textbox(label="ID pacienta", placeholder="P001 nebo nechej prázdné")
            pname = gr.Textbox(label="Jméno", placeholder="Jan Novák")
            page = gr.Number(label="Věk", value=50)
            pgender = gr.Radio(["M", "F", "X"], label="Pohlaví", value="M")
        psymptoms = gr.Textbox(
            label="Příznaky",
            lines=4,
            placeholder="Pacient udává bolest na hrudi vyzařující do levé paže, trvá 20 minut, nausea..."
        )
        with gr.Row():
            bps = gr.Number(label="TK systola", value=140)
            bpd = gr.Number(label="TK diastola", value=90)
            phr = gr.Number(label="Puls /min", value=95)
            pspo2 = gr.Number(label="SpO2 %", value=97)
            ptemp = gr.Number(label="Teplota °C", value=36.8)
            pgcs = gr.Number(label="GCS", value=15)
        add_btn = gr.Button("🚨 Odeslat k triáži", variant="primary")
        add_status = gr.Textbox(label="Výsledek")
        queue_out1 = gr.Textbox(label="Fronta po přidání (JSON)", lines=10)
        add_btn.click(
            add_patient_to_queue,
            inputs=[pid2, pname, page, pgender, psymptoms, bps, bpd, phr, pspo2, ptemp, pgcs],
            outputs=[add_status, queue_out1],
        )

    with gr.Tab("👨‍⚕️ Fronta (doktor)"):
        refresh_btn = gr.Button("🔄 Obnovit frontu", variant="secondary")
        queue_out2 = gr.Textbox(label="Aktuální fronta (JSON)", lines=20)
        refresh_btn.click(get_queue_json, outputs=[queue_out2])
        with gr.Row():
            clr_id = gr.Textbox(label="ID pacienta k odebrání")
            clr_btn = gr.Button("✅ Označit jako ošetřeného", variant="secondary")
        clr_status = gr.Textbox(label="Status")
        clr_btn.click(clear_patient, inputs=[clr_id], outputs=[clr_status, queue_out2])

demo.launch()