editor.py

"""
CDSS Rule Editor Component
"""

import gradio as gr
import pandas as pd
import json
import ast


def parse_rules():
    with open("rules.py", "r") as f:
        tree = ast.parse(f.read())

    rules = {"Mother": [], "Neonate": [], "Gyn": []}

    for node in ast.walk(tree):
        if isinstance(node, ast.FunctionDef) and node.name == "rule_based_cdss":
            for body_item in node.body:
                if (
                    isinstance(body_item, ast.If)
                    and isinstance(body_item.test, ast.Compare)
                    and isinstance(body_item.test.left, ast.Attribute)
                    and isinstance(body_item.test.left.value, ast.Name)
                    and body_item.test.left.value.id == "state"
                    and body_item.test.left.attr == "patient_type"
                    and isinstance(body_item.test.ops[0], ast.Eq)
                    and isinstance(body_item.test.comparators[0], ast.Constant)
                ):

                    patient_type = body_item.test.comparators[0].value
                    if patient_type in rules:
                        for rule_node in body_item.body:
                            if isinstance(rule_node, ast.If):
                                conditions = ast.unparse(rule_node.test)
                                alert = ""
                                for item in rule_node.body:
                                    if (
                                        isinstance(item, ast.Expr)
                                        and isinstance(item.value, ast.Call)
                                        and hasattr(item.value.func, "value")
                                        and hasattr(item.value.func.value, "id")
                                        and item.value.func.value.id == "alerts"
                                        and item.value.func.attr == "append"
                                        and isinstance(item.value.args[0], ast.Constant)
                                    ):
                                        alert = item.value.args[0].value
                                rules[patient_type].append(
                                    {"conditions": conditions, "alert": alert}
                                )
    return rules


def rules_to_dataframes(rules):
    dataframes = {}
    for patient_type, rules_list in rules.items():
        data = {"Conditions": [], "Alert": []}
        for rule in rules_list:
            data["Conditions"].append(rule["conditions"])
            data["Alert"].append(rule["alert"])
        df = pd.DataFrame(data)
        dataframes[patient_type] = df
    return dataframes


def dataframes_to_rules(dfs):
    rules = {"Mother": [], "Neonate": [], "Gyn": []}
    for patient_type, df in dfs.items():
        if df is not None:
            for index, row in df.iterrows():
                if row["Conditions"] and row["Alert"]:
                    rules[patient_type].append(
                        {"conditions": row["Conditions"], "alert": row["Alert"]}
                    )
    return rules


def add_row(df):
    if df is None:
        df = pd.DataFrame(columns=["Conditions", "Alert"])
    df.loc[len(df)] = ["", ""]
    return df


def save_rules(df_mother, df_neonate, df_gyn):
    dfs = {"Mother": df_mother, "Neonate": df_neonate, "Gyn": df_gyn}
    for patient_type, df in dfs.items():
        if not isinstance(df, pd.DataFrame):
            dfs[patient_type] = pd.DataFrame(df, columns=["Conditions", "Alert"])

    rules = dataframes_to_rules(dfs)

    with open("rules.py", "r") as f:
        tree = ast.parse(f.read())

    for node in ast.walk(tree):
        if isinstance(node, ast.FunctionDef) and node.name == "rule_based_cdss":
            node.body = []
            node.body.append(ast.parse("v = state.vitals").body[0])
            node.body.append(ast.parse("labs = state.labs").body[0])
            node.body.append(ast.parse("alerts = []").body[0])

            for patient_type, rule_list in rules.items():
                if_patient_type_body = []
                for rule in rule_list:
                    conditions = (
                        rule["conditions"].replace("\r", " ").replace("\n", " ")
                    )
                    if_rule_str = f"if {conditions}:\n    alerts.append({json.dumps(rule['alert'])})"
                    if_rule = ast.parse(if_rule_str).body[0]
                    if_patient_type_body.append(if_rule)

                if if_patient_type_body:
                    if_patient_type = ast.If(
                        test=ast.Compare(
                            left=ast.Attribute(
                                value=ast.Name(id="state", ctx=ast.Load()),
                                attr="patient_type",
                                ctx=ast.Load(),
                            ),
                            ops=[ast.Eq()],
                            comparators=[ast.Constant(value=patient_type)],
                        ),
                        body=if_patient_type_body,
                        orelse=[],
                    )
                    node.body.append(if_patient_type)

            node.body.append(
                ast.parse(
                    'if not alerts:\\n    return "Tidak ada alert prioritas tinggi. Lanjutkan pemantauan dan dokumentasi."'
                ).body[0]
            )
            node.body.append(
                ast.parse(
                    'return "\\n- ".join(["ALERT:"] + alerts)', mode="single"
                ).body[0]
            )

    new_code = ast.unparse(tree)
    with open("rules.py", "w") as f:
        f.write(new_code)

    return "Rules saved successfully."


def editor_ui():
    with gr.TabItem("Rule Editor"):
        with gr.Tabs():
            with gr.TabItem("Edit Rules"):
                gr.Markdown("## CDSS Rule Editor")

                initial_rules = parse_rules()
                initial_dfs = rules_to_dataframes(initial_rules)

                with gr.Tabs():
                    with gr.Tab("Mother"):
                        df_mother = gr.DataFrame(
                            value=initial_dfs["Mother"],
                            headers=["Conditions", "Alert"],
                            interactive=True,
                            row_count=(len(initial_dfs["Mother"]) + 1, "dynamic"),
                            type="pandas",
                        )
                        add_mother_btn = gr.Button("➕ Add Mother Rule")
                        add_mother_btn.click(
                            add_row, inputs=df_mother, outputs=df_mother
                        )

                    with gr.Tab("Neonate"):
                        df_neonate = gr.DataFrame(
                            value=initial_dfs["Neonate"],
                            headers=["Conditions", "Alert"],
                            interactive=True,
                            row_count=(len(initial_dfs["Neonate"]) + 1, "dynamic"),
                            type="pandas",
                        )
                        add_neonate_btn = gr.Button("➕ Add Neonate Rule")
                        add_neonate_btn.click(
                            add_row, inputs=df_neonate, outputs=df_neonate
                        )

                    with gr.Tab("Gyn"):
                        df_gyn = gr.DataFrame(
                            value=initial_dfs["Gyn"],
                            headers=["Conditions", "Alert"],
                            interactive=True,
                            row_count=(len(initial_dfs["Gyn"]) + 1, "dynamic"),
                            type="pandas",
                        )
                        add_gyn_btn = gr.Button("➕ Add Gyn Rule")
                        add_gyn_btn.click(add_row, inputs=df_gyn, outputs=df_gyn)

                save_button = gr.Button("💾 Save Rules")
                status_textbox = gr.Textbox(label="Status", interactive=False)

                save_button.click(
                    save_rules,
                    inputs=[df_mother, df_neonate, df_gyn],
                    outputs=status_textbox,
                )
    return df_mother, df_neonate, df_gyn, save_button, status_textbox