app.py

"""
Pre-Visit Eye Intake Agent — Hugging Face Space (Gradio)

A hackathon prototype. NOT a medical device. All outputs are suggestions for
clinician review only.

Flow:
  1. Patient uploads photo AND answers initial questions (upfront).
  2. On submit: quality gate -> (retake loop) -> perception/findings.
  3. Dynamic follow-up questions appear, chosen from findings + answers.
  4. On answering: red-flag fast path -> triage -> chart.
"""

import os
import re
import json

import numpy as np
import gradio as gr

try:
    import cv2
    _HAVE_CV2 = True
except Exception:
    _HAVE_CV2 = False

# ----------------------------------------------------------------------------
# Config (Space -> Settings -> Variables and secrets)
# ----------------------------------------------------------------------------
# Triage LLM via any OpenAI-compatible endpoint (Modal vLLM or NVIDIA API).
LLM_API_KEY = os.environ.get("LLM_API_KEY")
LLM_BASE_URL = os.environ.get("LLM_BASE_URL", "https://integrate.api.nvidia.com/v1")
LLM_MODEL = os.environ.get("LLM_MODEL", "nvidia/NVIDIA-Nemotron-Nano-9B-v2")
FINDINGS_MODEL = os.environ.get("FINDINGS_MODEL")  # "hf-hub:your-user/eye-findings"

# Quality-gate thresholds
BLUR_MIN = 100.0
BRIGHT_MIN = 40.0
BRIGHT_MAX = 220.0

MAX_FOLLOWUPS = 3  # number of dynamic question slots in the UI

# Emergency phrases in the free text that force URGENT
RED_FLAGS = [
    "sudden vision loss", "sudden loss of vision", "lost vision", "cant see",
    "can't see", "curtain", "flashes", "floaters", "chemical", "bleach",
    "severe pain", "trauma", "hit in the eye", "double vision",
]


# ----------------------------------------------------------------------------
# Step: Quality gate (local, no model required)
# ----------------------------------------------------------------------------
def check_quality(img):
    if img is None:
        return False, "No image received. Please upload an eye photo."
    if not _HAVE_CV2:
        return True, "Quality check skipped (OpenCV unavailable)."
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    blur = cv2.Laplacian(gray, cv2.CV_64F).var()
    brightness = float(gray.mean())
    if blur < BLUR_MIN:
        return False, "Image looks **blurry**. Hold steady, tap to focus, and retake."
    if brightness < BRIGHT_MIN:
        return False, "Image looks **too dark**. Move to better light and retake."
    if brightness > BRIGHT_MAX:
        return False, "Image looks **overexposed / glary**. Reduce direct light and retake."
    return True, f"Quality OK (sharpness {blur:.0f}, brightness {brightness:.0f})."


# ----------------------------------------------------------------------------
# Step: Perception / findings (placeholder heuristic; swap in your timm model)
# ----------------------------------------------------------------------------
_findings_model = None


def _load_findings_model():
    global _findings_model
    if _findings_model is not None or not FINDINGS_MODEL:
        return _findings_model
    try:
        import timm  # noqa
        _findings_model = timm.create_model(FINDINGS_MODEL, pretrained=True)
        _findings_model.eval()
    except Exception as e:
        print(f"[findings] could not load {FINDINGS_MODEL}: {e}")
        _findings_model = None
    return _findings_model


def detect_findings(img):
    model = _load_findings_model()
    if model is not None:
        # TODO: real preprocessing + your label map.
        pass
    r, g, b = img[..., 0].mean(), img[..., 1].mean(), img[..., 2].mean()
    redness = max(0.0, (r - (g + b) / 2) / 255.0)
    redness_conf = min(1.0, redness * 4)
    return {
        "redness": round(float(redness_conf), 2),
        "note": "PLACEHOLDER heuristic — swap in your fine-tuned timm model.",
    }


# ----------------------------------------------------------------------------
# Step: Dynamic follow-up questions (ask-then-act)
# Driven by findings + the patient's initial answers. Red-flag screens always
# come first. Returns up to MAX_FOLLOWUPS question dicts.
# ----------------------------------------------------------------------------
def generate_followups(findings, pain, free_text):
    redness = findings.get("redness", 0)
    candidates = [
        {"key": "sudden_vision_loss",
         "question": "Sudden loss of vision or a 'curtain' over part of your sight?",
         "options": ["No", "Yes"], "red_flag": True, "prio": 1},
        {"key": "flashes_floaters",
         "question": "New flashes of light or a sudden shower of floaters?",
         "options": ["No", "Yes"], "red_flag": True, "prio": 2},
    ]
    if pain in ("Moderate", "Severe"):
        candidates.append(
            {"key": "photophobia",
             "question": "Is light painful to look at (light sensitivity)?",
             "options": ["No", "Yes"], "red_flag": False, "prio": 3})
    if redness > 0.35:
        candidates.append(
            {"key": "discharge",
             "question": "Any discharge from the eye?",
             "options": ["None", "Clear / watery", "Thick / colored"],
             "red_flag": False, "prio": 4})
    if redness > 0.35 or pain != "None":
        candidates.append(
            {"key": "contacts",
             "question": "Do you wear contact lenses?",
             "options": ["No", "Yes"], "red_flag": False, "prio": 5})
    candidates.append(
        {"key": "vision_change",
         "question": "Any change in your vision?",
         "options": ["No", "Slightly blurry", "Much worse"],
         "red_flag": False, "prio": 6})
    candidates.sort(key=lambda c: c["prio"])
    return candidates[:MAX_FOLLOWUPS]


# ----------------------------------------------------------------------------
# Step: Triage (Nemotron via OpenAI-compatible endpoint; rule-based fallback)
# ----------------------------------------------------------------------------
SYSTEM_PROMPT = (
    "You are a triage assistant for an eye clinic. You DO NOT diagnose. "
    "Given image findings and patient answers, assign a priority: "
    "ROUTINE, SAME-DAY, or URGENT. Be conservative: when unsure, escalate. "
    "Respond ONLY with compact JSON: "
    '{"category": "...", "rationale": "...", "follow_up": ["..."]}'
)


def _triage_rule_based(findings, symptoms):
    pain = symptoms.get("pain")
    redness = findings.get("redness", 0)
    dyn = symptoms.get("follow_up_answers", {})
    if pain == "Severe" or dyn.get("photophobia") == "Yes":
        return {"category": "SAME-DAY",
                "rationale": "Severe pain or light sensitivity reported.",
                "follow_up": ["Any discharge?"]}
    if dyn.get("discharge") == "Thick / colored":
        return {"category": "SAME-DAY",
                "rationale": "Thick/colored discharge suggests infection.",
                "follow_up": ["Contact lens wearer?"]}
    if redness > 0.4 and pain in ("Moderate", "Severe"):
        return {"category": "SAME-DAY",
                "rationale": "Notable redness with pain.",
                "follow_up": ["Contact lens wearer?"]}
    return {"category": "ROUTINE",
            "rationale": "No high-priority features detected.",
            "follow_up": ["Any change in vision?"]}


def triage(findings, symptoms):
    if not LLM_API_KEY:
        out = _triage_rule_based(findings, symptoms)
        out["source"] = "rule-based (set LLM_API_KEY to enable Nemotron)"
        return out
    try:
        from openai import OpenAI
        client = OpenAI(base_url=LLM_BASE_URL, api_key=LLM_API_KEY)
        user_msg = (
            f"Image findings: {json.dumps(findings)}\n"
            f"Patient answers: {json.dumps(symptoms)}\n"
            "Return the JSON now."
        )
        resp = client.chat.completions.create(
            model=LLM_MODEL,
            messages=[{"role": "system", "content": SYSTEM_PROMPT},
                      {"role": "user", "content": user_msg}],
            max_tokens=400,
            temperature=0.2,
            extra_body={"chat_template_kwargs": {"enable_thinking": False}},
        )
        raw = resp.choices[0].message.content
        match = re.search(r"\{.*\}", raw, re.DOTALL)
        data = json.loads(match.group(0)) if match else {}
        data.setdefault("category", "ROUTINE")
        data.setdefault("rationale", raw[:200])
        data.setdefault("follow_up", [])
        data["source"] = f"Nemotron ({LLM_MODEL})"
        return data
    except Exception as e:
        out = _triage_rule_based(findings, symptoms)
        out["source"] = f"rule-based fallback (LLM error: {e})"
        return out


# ----------------------------------------------------------------------------
# Orchestration
# ----------------------------------------------------------------------------
def run_intake_step(img, free_text, pain, duration, laterality):
    """Step 1 -> quality gate -> findings -> reveal dynamic follow-ups."""
    passed, msg = check_quality(img)
    hide = gr.update(visible=False)
    if not passed:
        # Quality-gate loop: refuse to assess, keep step 2 hidden.
        return (gr.update(value=f"⛔ **Retake needed.** {msg}"),
                gr.update(visible=False),  # dyn group
                hide, hide, hide,          # dyn questions
                None, None, None)          # states cleared

    findings = detect_findings(img)
    symptoms = {"free_text": (free_text or "").strip(), "pain": pain,
                "duration": duration, "laterality": laterality}
    followups = generate_followups(findings, pain, free_text)

    fmt = "\n".join(f"- **{k}**: {v}" for k, v in findings.items())
    status = (f"✅ {msg}\n\n**Findings**\n{fmt}\n\n"
              "Please answer the follow-up questions below, then run triage.")

    q_updates = []
    for i in range(MAX_FOLLOWUPS):
        if i < len(followups):
            f = followups[i]
            q_updates.append(gr.update(label=f["question"], choices=f["options"],
                                       value=f["options"][0], visible=True))
        else:
            q_updates.append(gr.update(visible=False))

    return (gr.update(value=status),
            gr.update(visible=True),
            q_updates[0], q_updates[1], q_updates[2],
            findings, symptoms, followups)


def run_triage_step(findings, symptoms, followups, a1, a2, a3):
    """Step 2 -> collect dynamic answers -> red-flag fast path -> triage -> chart."""
    if findings is None or symptoms is None:
        return "Please submit a good-quality image and details first."

    answers = [a1, a2, a3]
    dyn, red = {}, []
    for f, a in zip(followups or [], answers):
        dyn[f["key"]] = a
        if f.get("red_flag") and a == "Yes":
            red.append(f["question"])
    red += [kw for kw in RED_FLAGS if kw in symptoms.get("free_text", "").lower()]

    symptoms_full = dict(symptoms)
    symptoms_full["follow_up_answers"] = dyn

    if red:
        category = "URGENT"
        rationale = "Red flag(s): " + "; ".join(red)
        follow_up, source = [], "red-flag fast path"
    else:
        result = triage(findings, symptoms_full)
        category = result["category"]
        rationale = result["rationale"]
        follow_up = result.get("follow_up", [])
        source = result["source"]

    badge = {"URGENT": "🔴", "SAME-DAY": "🟠", "ROUTINE": "🟢"}.get(category, "⚪")
    chart = [
        "## Pre-visit chart (for clinician review)",
        f"### {badge} Priority: **{category}**",
        f"**Rationale:** {rationale}",
        "",
        "**Image findings**",
        *[f"- {k}: {v}" for k, v in findings.items()],
        "",
        "**Initial intake**",
        f"- What's bothering you: {symptoms['free_text'] or '(none)'}",
        f"- Pain: {symptoms['pain']} · Duration: {symptoms['duration']} "
        f"· Side: {symptoms['laterality']}",
    ]
    if dyn:
        chart += ["", "**Follow-up answers**",
                  *[f"- {k.replace('_', ' ')}: {v}" for k, v in dyn.items()]]
    if follow_up:
        chart += ["", "**Suggested next questions**", *[f"- {q}" for q in follow_up]]
    chart += ["", f"_Triage source: {source}_",
              "_Not a medical device. Suggestion only._"]
    return "\n".join(chart)


# ----------------------------------------------------------------------------
# UI
# ----------------------------------------------------------------------------
with gr.Blocks(title="Pre-Visit Eye Intake Agent", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        "# 👁️ Pre-Visit Eye Intake Agent\n"
        "Upload an eye photo and tell us what's going on. The agent checks image "
        "quality, runs perception, asks a few targeted follow-ups, then triages.\n\n"
        "> Prototype for a hackathon. **Not a medical device.**"
    )

    findings_state = gr.State()
    symptoms_state = gr.State()
    followups_state = gr.State()

    with gr.Row():
        with gr.Column():
            gr.Markdown("### 1. Eye photo")
            image_in = gr.Image(type="numpy", label="Eye photo", height=240)
            gr.Markdown("### 2. About the problem")
            free_text = gr.Textbox(
                label="What's bothering you?",
                placeholder="e.g. red, watery left eye for 2 days", lines=3)
            pain = gr.Radio(["None", "Mild", "Moderate", "Severe"],
                            value="None", label="Pain")
            duration = gr.Radio(["< 1 day", "1-3 days", "> 3 days"],
                                value="1-3 days", label="Duration")
            laterality = gr.Radio(["Left", "Right", "Both"],
                                  value="Left", label="Affected eye")
            submit_btn = gr.Button("Submit & analyze", variant="primary")
            intake_status = gr.Markdown()

        with gr.Column():
            dyn_group = gr.Group(visible=False)
            with dyn_group:
                gr.Markdown("### 3. A few follow-up questions")
                dyn_q1 = gr.Radio(choices=["No", "Yes"], label="", visible=False)
                dyn_q2 = gr.Radio(choices=["No", "Yes"], label="", visible=False)
                dyn_q3 = gr.Radio(choices=["No", "Yes"], label="", visible=False)
                triage_btn = gr.Button("Run triage", variant="primary")
            chart_out = gr.Markdown()

    submit_btn.click(
        run_intake_step,
        inputs=[image_in, free_text, pain, duration, laterality],
        outputs=[intake_status, dyn_group, dyn_q1, dyn_q2, dyn_q3,
                 findings_state, symptoms_state, followups_state],
    )
    triage_btn.click(
        run_triage_step,
        inputs=[findings_state, symptoms_state, followups_state,
                dyn_q1, dyn_q2, dyn_q3],
        outputs=chart_out,
    )

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