simulator.py

"""
CDSS Simulator Component
"""

import random
import time
from dataclasses import asdict
from typing import Dict, Any, Tuple, List
from datetime import datetime

import gradio as gr
import pandas as pd
import plotly.express as px

from models import Vitals, PatientState
from rules import rule_based_cdss
from google.generativeai.types import HarmCategory, HarmBlockThreshold

# --- Gemini setup (simplified) ---
try:
    import google.generativeai as genai
    import os

    genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
    GEMINI_MODEL = genai.GenerativeModel("gemini-2.5-flash")
    GEMINI_ERR = None
except Exception as e:
    GEMINI_MODEL, GEMINI_ERR = None, f"Gemini import/config error: {e}"


# --- Data structures & Scenarios (Full list included) ---


def scenario_A0_Normal() -> PatientState:
    return PatientState(
        "A0 Normal Case",
        "Mother",
        "Pemeriksaan rutin.",
        {"Hb": 12.5},
        Vitals(110, 70, 80, 16, 36.7, 99),
    )


def scenario_A1_PPH() -> PatientState:
    return PatientState(
        "A1 PPH",
        "Mother",
        "30 menit postpartum; kehilangan darah ~900 ml.",
        {"Hb": 9},
        Vitals(90, 60, 120, 24, 36.8, 96),
    )


def scenario_A2_Preeclampsia() -> PatientState:
    return PatientState(
        "A2 Preeklampsia",
        "Mother",
        "36 minggu; sakit kepala, pandangan kabur.",
        {"Proteinuria": "3+"},
        Vitals(165, 105, 98, 20, 36.9, 98),
    )


def scenario_A3_MaternalSepsis() -> PatientState:
    return PatientState(
        "A3 Sepsis Maternal",
        "Mother",
        "POD2 pasca SC; luka purulen.",
        {"Leukosit": 17000},
        Vitals(95, 60, 110, 24, 39.0, 96),
    )


def scenario_B1_Prematurity() -> PatientState:
    return PatientState(
        "B1 Prematuritas/BBLR",
        "Neonate",
        "34 minggu; berat 1900 g; hipotermia ringan; SpO2 borderline",
        {"BB": 1900, "UsiaGestasi_mgg": 34},
        Vitals(60, 35, 150, 50, 35.0, 90),
    )


def scenario_B2_Asphyxia() -> PatientState:
    return PatientState(
        "B2 Asfiksia Perinatal",
        "Neonate",
        "APGAR 3 menit 1; tidak menangis >1 menit",
        {"APGAR_1m": 3},
        Vitals(55, 30, 80, 10, 36.5, 82),
    )


def scenario_B3_NeonatalSepsis() -> PatientState:
    return PatientState(
        "B3 Sepsis Neonatal",
        "Neonate",
        "Hari ke-4; lemas, malas minum",
        {"CRP": 25, "Leukosit": 19000},
        Vitals(60, 35, 170, 60, 38.5, 93),
    )


def scenario_C1_GynSurgComp() -> PatientState:
    return PatientState(
        "C1 Komplikasi Bedah Ginekologis",
        "Gyn",
        "Pasca histerektomi; nyeri perut bawah; urine output turun",
        {"UrineOutput_ml_hr": 10},
        Vitals(100, 65, 105, 20, 37.8, 98),
    )


def scenario_C2_PostOpInfection() -> PatientState:
    return PatientState(
        "C2 Infeksi Pasca-Bedah",
        "Gyn",
        "Pasca kistektomi; luka bengkak & kemerahan; demam",
        {"Luka": "bengkak+kemerahan"},
        Vitals(105, 70, 108, 22, 38.0, 98),
    )


def scenario_C3_DelayedGynCancer() -> PatientState:
    return PatientState(
        "C3 Keterlambatan Diagnostik Kanker Ginekologi",
        "Gyn",
        "45 th; perdarahan pascamenopause; Pap abnormal 6 bulan lalu tanpa tindak lanjut",
        {"PapSmear": "abnormal 6 bln lalu"},
        Vitals(120, 78, 86, 18, 36.8, 99),
    )


SCENARIOS = {
    "A0": scenario_A0_Normal,
    "A1": scenario_A1_PPH,
    "A2": scenario_A2_Preeclampsia,
    "A3": scenario_A3_MaternalSepsis,
    "B1": scenario_B1_Prematurity,
    "B2": scenario_B2_Asphyxia,
    "B3": scenario_B3_NeonatalSepsis,
    "C1": scenario_C1_GynSurgComp,
    "C2": scenario_C2_PostOpInfection,
    "C3": scenario_C3_DelayedGynCancer,
}


# --- Simulation & CDSS Logic (simplified) ---
def drift_vitals(state: PatientState) -> PatientState:
    v = state.vitals
    clamp = lambda val, lo, hi: max(lo, min(hi, val))
    drift_factor = 0 if state.scenario.startswith("A0") else 1
    v.hr = clamp(v.hr + random.randint(-2, 2) * drift_factor, 40, 200)
    v.sbp = clamp(v.sbp + random.randint(-2, 2) * drift_factor, 50, 220)
    v.rr = clamp(v.rr + random.randint(-1, 1) * drift_factor, 8, 80)
    state.vitals = v
    return state


# --- Rule-based fallback (no AI or AI disabled) ---


def gemini_cdss(state: PatientState) -> str:
    if not GEMINI_MODEL:
        return f"[CDSS AI ERROR] {GEMINI_ERR}"
    try:
        v = state.vitals
        prompt = f"CDSS for {state.scenario}. Vitals: SBP {v.sbp}/{v.dbp}, HR {v.hr}. Analyze risks, give concise steps in Indonesian."
        response = GEMINI_MODEL.generate_content(
            prompt,
            safety_settings={
                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
            },
        )
        print(response)
        if response.parts:
            return response.text or "[CDSS AI] No response."
        else:
            return "[CDSS AI] No response due to safety settings."
    except Exception as e:
        return f"[CDSS AI error] {e}"


# --- Plotting & Data Helpers ---
def create_vital_plot(
    df: pd.DataFrame, y_cols: List[str] | str, title: str, y_lim: List[int]
):
    """Creates a customized Plotly figure for a specific vital sign."""
    # Create an empty plot if there is no data to prevent errors
    if df.empty:
        fig = px.line(title=title)
    else:
        fig = px.line(df, x="timestamp", y=y_cols, title=title, markers=True)
        # Customize x-axis to show only first and last tick
        if len(df) > 1:
            fig.update_xaxes(
                tickvals=[df["timestamp"].iloc[0], df["timestamp"].iloc[-1]]
            )

    # Apply standard layout settings
    fig.update_layout(
        height=250,
        yaxis_range=y_lim,
        margin=dict(t=40, b=10, l=10, r=10),  # Tighten margins
    )
    return fig


def _row_from_state(ps: PatientState) -> Dict[str, Any]:
    return {"timestamp": datetime.now(), "scenario": ps.scenario, **asdict(ps.vitals)}


def prepare_df_for_display(df: pd.DataFrame) -> pd.DataFrame:
    if df is None or df.empty:
        return pd.DataFrame(
            columns=[
                "timestamp",
                "scenario",
                "sbp",
                "dbp",
                "hr",
                "rr",
                "temp_c",
                "spo2",
            ]
        )
    df_display = df.copy()
    df_display["timestamp"] = pd.to_datetime(df_display["timestamp"])
    df_display = df_display.sort_values("timestamp")
    df_display["timestamp"] = df_display["timestamp"].dt.strftime("%Y-%m-%d %H:%M:%S")
    return df_display


def generate_all_plots(df: pd.DataFrame):
    """Helper to generate all 5 plot figures from a dataframe."""
    df_display = prepare_df_for_display(df)
    bp_fig = create_vital_plot(
        df_display,
        y_cols=["sbp", "dbp"],
        title="Blood Pressure (mmHg)",
        y_lim=[40, 200],
    )
    hr_fig = create_vital_plot(
        df_display, y_cols="hr", title="Heart Rate (bpm)", y_lim=[40, 200]
    )
    rr_fig = create_vital_plot(
        df_display, y_cols="rr", title="Respiratory Rate (/min)", y_lim=[0, 70]
    )
    temp_fig = create_vital_plot(
        df_display, y_cols="temp_c", title="Temperature (°C)", y_lim=[34, 42]
    )
    spo2_fig = create_vital_plot(
        df_display, y_cols="spo2", title="Oxygen Saturation (%)", y_lim=[70, 101]
    )
    return df_display, bp_fig, hr_fig, rr_fig, temp_fig, spo2_fig


# --- Gradio App Logic ---
def process_and_update(
    ps: PatientState, history_df: pd.DataFrame, historic_text: str, cdss_on: bool
):
    """Centralized function to process state, update history, and generate all UI component outputs."""
    interpretation = gemini_cdss(ps) if cdss_on else rule_based_cdss(ps)
    new_row = _row_from_state(ps)
    history_df = pd.concat([history_df, pd.DataFrame([new_row])], ignore_index=True)

    df_for_table, bp_fig, hr_fig, rr_fig, temp_fig, spo2_fig = generate_all_plots(
        history_df
    )

    return (
        asdict(ps),
        *state_to_panels(ps),
        str(ps.labs),  # For labs_text
        str(ps.labs),  # For labs_show
        interpretation,
        history_df,
        df_for_table,
        historic_text.strip(),
        time.time(),
        bp_fig,
        hr_fig,
        rr_fig,
        temp_fig,
        spo2_fig,
    )


def state_to_panels(state: PatientState) -> Tuple:
    v = state.vitals
    return (
        state.scenario,
        state.patient_type,
        state.notes,
        v.sbp,
        v.dbp,
        v.hr,
        v.rr,
        v.temp_c,
        v.spo2,
    )


def inject_scenario(
    tag: str, cdss_on: bool, history_df: pd.DataFrame, historic_text: str
):
    ps = SCENARIOS[tag]()
    if historic_text:  # Add a newline if text already exists
        historic_text += f"\n[{datetime.now().strftime('%H:%M:%S')}] Scenario Injected: {ps.scenario}"
    else:
        historic_text = (
            f"[{datetime.now().strftime('%H:%M:%S')}] Scenario Injected: {ps.scenario}"
        )
    return process_and_update(ps, history_df, historic_text, cdss_on)


def manual_edit(
    sbp,
    dbp,
    hr,
    rr,
    temp_c,
    spo2,
    notes,
    labs_text,
    cdss_on,
    patient_type,
    current_state,
    history_df,
    historic_text,
):
    try:
        labs = eval(labs_text)
    except:
        labs = {"raw": labs_text}
    ps = PatientState(
        current_state.get("scenario", "Manual"),
        patient_type,
        notes,
        labs,
        Vitals(int(sbp), int(dbp), int(hr), int(rr), float(temp_c), int(spo2)),
    )
    if ps.notes and ps.notes.strip():
        historic_text += f"\n[{datetime.now().strftime('%H:%M:%S')}] {ps.notes}"
    return process_and_update(ps, history_df, historic_text, cdss_on)


def tick_timer(cdss_on, current_state, history_df, historic_text):
    if not current_state:
        return [gr.update()] * 22
    ps = PatientState(**current_state)
    ps.vitals = Vitals(**ps.vitals)
    ps = drift_vitals(ps)
    return process_and_update(ps, history_df, historic_text, cdss_on)


def load_csv(file, history_df: pd.DataFrame):
    try:
        if file is not None:
            df_new = pd.read_csv(file.name)
            df_new["timestamp"] = pd.to_datetime(df_new["timestamp"])
            history_df = (
                pd.concat([history_df, df_new], ignore_index=True)
                if not history_df.empty
                else df_new
            )
    except Exception as e:
        print(f"Error loading CSV: {e}")
    df_for_table, bp_fig, hr_fig, rr_fig, temp_fig, spo2_fig = generate_all_plots(
        history_df
    )
    return history_df, df_for_table, bp_fig, hr_fig, rr_fig, temp_fig, spo2_fig


def countdown_tick(last_tick_ts: float):
    if not last_tick_ts:
        return "Next update in —"
    return f"Next update in {max(0, 30 - int(time.time() - last_tick_ts))}s"


def simulator_ui():
    with gr.TabItem("CDSS Simulator"):
        with gr.Accordion("History, Trends, and Data Loading", open=True):
            with gr.Row():
                with gr.Tabs():
                    with gr.Tab("Blood Pressure"):
                        bp_plot = gr.Plot()
                    with gr.Tab("Heart Rate"):
                        hr_plot = gr.Plot()
                    with gr.Tab("Respiration"):
                        rr_plot = gr.Plot()
                    with gr.Tab("Temperature"):
                        temp_plot = gr.Plot()
                    with gr.Tab("SpO₂"):
                        spo2_plot = gr.Plot()
    return bp_plot, hr_plot, rr_plot, temp_plot, spo2_plot