cdss.py

"""
Gradio CDSS Simulator for RSIA (Obstetri, Neonatologi, Ginekologi)

FINAL INTEGRATED VERSION - Now with:
- **Custom Plotly graphs** using gr.Plot for full control.
- **X-Axis shows only the first and last timestamps.**
- Tabbed plots for each vital sign, guaranteeing graph rendering.
- Historic Clinical Text, 30s countdown, CSV loader, and all scenarios.

Requirements:
pip install gradio google-generativeai pandas plotly

Set environment variable for Gemini:
export GOOGLE_API_KEY=your_key_here

Run:
python app.py
"""

from __future__ import annotations
import time

import gradio as gr
import pandas as pd

from models import Vitals, PatientState
from simulator import (
    simulator_ui,
    inject_scenario,
    manual_edit,
    tick_timer,
    load_csv,
    countdown_tick,
    SCENARIOS,
)
from editor import editor_ui, save_rules
from validator import validator_ui, test_condition, add_rule_to_set
from diagnosis import diagnosis_ui, generate_diagnosis, check_medication_interaction, generate_medical_record_from_image


# --- Build UI ---
with gr.Blocks(
    css=".gradio-container { max-width: 1200px !important; margin: auto !important; }"
) as demo:
    gr.Markdown("# RSIA CDSS Simulator with Custom Plotly Graphs")
    state = gr.State({})
    history_df = gr.State(pd.DataFrame())
    historic_text = gr.State("")
    last_tick_ts = gr.State(time.time())
    interpretation = gr.Textbox(label="CDSS Interpretation", lines=2, interactive=False)

    with gr.Tabs():
        bp_plot, hr_plot, rr_plot, temp_plot, spo2_plot = simulator_ui()
        (df_mother, df_neonate, df_gyn, save_button, status_textbox) = editor_ui()
        (
            patient_type_validate,
            sbp_validate,
            dbp_validate,
            hr_validate,
            rr_validate,
            temp_c_validate,
            spo2_validate,
            labs_validate,
            condition_validate,
            alert_validate,
            test_button,
            validation_result,
            add_rule_button,
            add_rule_status,
        ) = validator_ui()
        (
            generate_button,
            diagnosis_output,
            medication_output,
            medication_input,
            check_button,
            interaction_output,
            image_input,
            generate_from_image_button,
            medical_record_output,
        ) = diagnosis_ui()

    with gr.Row():
        with gr.Column(scale=2):
            with gr.Accordion("Inject New Scenario (resets chart)", open=True):
                with gr.Row():
                    btn_A0 = gr.Button("A0: Normal", elem_classes="small-btn")
                    btn_A1 = gr.Button("A1: PPH", elem_classes="small-btn")
                    btn_A2 = gr.Button("A2: Preeklampsia", elem_classes="small-btn")
                    btn_A3 = gr.Button("A3: Sepsis", elem_classes="small-btn")
                with gr.Row():
                    btn_B1 = gr.Button("B1: Prematuritas", elem_classes="small-btn")
                    btn_B2 = gr.Button("B2: Asfiksia", elem_classes="small-btn")
                    btn_B3 = gr.Button("B3: Sepsis", elem_classes="small-btn")
                with gr.Row():
                    btn_C1 = gr.Button("C1: Bedah Komplikasi", elem_classes="small-btn")
                    btn_C2 = gr.Button(
                        "C2: Infeksi Pasca-Bedah", elem_classes="small-btn"
                    )
                    btn_C3 = gr.Button("C3: Kanker", elem_classes="small-btn")
            notes = gr.Textbox(label="Catatan Klinis", lines=2)
            labs_text = gr.Textbox(label="Lab (dict or text)", value="{}")
            labs_show = gr.Textbox(label="Labs (Parsed)", interactive=False)
        with gr.Column(scale=1):
            with gr.Group():
                patient_type_radio = gr.Radio(
                    ["Mother", "Neonate", "Gyn"], label="Patient Type", value="Mother"
                )
                sbp = gr.Number(label="SBP")
                dbp = gr.Number(label="DBP")
                hr = gr.Number(label="HR")
                rr = gr.Number(label="RR")
                temp_c = gr.Number(label="Temp (°C)")
                spo2 = gr.Number(label="SpO₂ (%)")
            apply_manual = gr.Button("Apply Manual Edits", variant="secondary")

    with gr.Row():
        csv_file = gr.File(label="Load CSV to Graph")
        df_view = gr.Dataframe(label="History Table", wrap=True, interactive=False)
        cdss_toggle = gr.Checkbox(value=False, label="With CDSS (Gemini)")
        scenario_lbl = gr.Textbox(label="Active Scenario", interactive=False)
        countdown_lbl = gr.Label()
        historic_box = gr.Textbox(label="Historic Text", lines=12, interactive=False)

    # --- Event Handlers ---
    def update_medication_input(patient_type):
        if patient_type == "Mother":
            return gr.update(value="Aspirin, Ibuprofen")
        elif patient_type == "Gyn":
            return gr.update(value="Clopidogrel, Omeprazole")
        elif patient_type == "Neonate":
            return gr.update(value="Ceftriaxone, Calcium")

    patient_type_radio.change(update_medication_input, inputs=patient_type_radio, outputs=medication_input)

    ui_outputs = [
        state,
        scenario_lbl,
        patient_type_radio,
        notes,
        sbp,
        dbp,
        hr,
        rr,
        temp_c,
        spo2,
        labs_text,
        labs_show,
        interpretation,
        history_df,
        df_view,
        historic_box,
        last_tick_ts,
        bp_plot,
        hr_plot,
        rr_plot,
        temp_plot,
        spo2_plot,
    ]

    manual_inputs = [
        sbp,
        dbp,
        hr,
        rr,
        temp_c,
        spo2,
        notes,
        labs_text,
        cdss_toggle,
        patient_type_radio,
        state,
        history_df,
        historic_text,
    ]
    apply_manual.click(manual_edit, manual_inputs, ui_outputs)

    for tag, btn in zip(
        SCENARIOS.keys(),
        [
            btn_A0,
            btn_A1,
            btn_A2,
            btn_A3,
            btn_B1,
            btn_B2,
            btn_B3,
            btn_C1,
            btn_C2,
            btn_C3,
        ],
    ):
        btn.click(
            inject_scenario,
            [gr.State(tag), cdss_toggle, history_df, historic_text],
            ui_outputs,
        )

    csv_outputs = [history_df, df_view, bp_plot, hr_plot, rr_plot, temp_plot, spo2_plot]
    csv_file.change(load_csv, [csv_file, history_df], csv_outputs)

    timer_inputs = [cdss_toggle, state, history_df, historic_text]
    gr.Timer(30.0).tick(tick_timer, timer_inputs, ui_outputs)
    gr.Timer(1.0).tick(countdown_tick, [last_tick_ts], [countdown_lbl])

    generate_button.click(generate_diagnosis, inputs=state, outputs=[diagnosis_output, medication_output])
    check_button.click(check_medication_interaction, inputs=[patient_type_radio, medication_input], outputs=interaction_output)

    generate_from_image_button.click(
        generate_medical_record_from_image,
        inputs=[state, image_input],
        outputs=medical_record_output
    )

    demo.load(inject_scenario, [gr.State("A0"), cdss_toggle, history_df, historic_text], ui_outputs)

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