pages/YAML_editor_testing.py

import streamlit as st
import pandas as pd
import os
from modeci_mdf.mdf import Model, Graph, Node, Parameter, OutputPort
from modeci_mdf.utils import load_mdf_json, load_mdf, load_mdf_yaml
from modeci_mdf.execution_engine import EvaluableGraph
from code_editor import code_editor

st.set_page_config(layout="wide")
st.title("🔋 MDF Simulator")

height = [19, 22]
theme = "default"
shortcuts = "vscode"
focus = False
wrap = True
editor_btns = [{
    "name": "Run",
    "feather": "Play",
    "primary": True,
    "hasText": True,
    "showWithIcon": True,
    "commands": ["submit"],
    "style": {"bottom": "0.44rem", "right": "0.4rem"}
}]

def mdf_model_to_paramters_dictionary(param_inputs, mdf_model):
    mod_graph = mdf_model.graphs[0]
    nodes = mod_graph.nodes[0]
    parameters = nodes.parameters
    for param in parameters:
        if isinstance(param.value, str) or param.value is None:
            continue
        st.session_state.param_inputs[param.id] = param.value
    return st.session_state.param_inputs

def run_simulation(param_inputs, mdf_model):
    mod_graph = mdf_model.graphs[0]
    nodes = mod_graph.nodes[0]
    parameters = nodes.parameters
    outputs = nodes.output_ports
    eg = EvaluableGraph(mod_graph, verbose=False)
    duration = param_inputs["Simulation Duration (s)"]
    dt = param_inputs["Time Step (s)"]
    t = 0
    times = []
    output_values = {op.value: [] for op in outputs}
    while t <= duration:
        times.append(t)
        if t == 0:
            eg.evaluate()
        else:
            eg.evaluate(time_increment=dt)

        for param in output_values:
            eval_param = eg.enodes[nodes.id].evaluable_parameters[param]
            output_values[param].append(eval_param.curr_value)
        t += dt
    chart_data = pd.DataFrame(output_values)
    chart_data['Time'] = times
    chart_data.set_index('Time', inplace=True)
    return chart_data

def show_simulation_results():
    if 'chart_data' in st.session_state:
        st.line_chart(st.session_state.chart_data, use_container_width=True, height=400)
        st.write("Input Values")
        st.write(st.session_state.param_inputs)
        st.write("Output Values")
        st.write(st.session_state.chart_data)

def show_mdf_graph():
    if 'mdf_model_graph' in st.session_state:
        st.subheader("MDF Graph")
        image_path = st.session_state.mdf_model.id + ".png"
        st.image(image_path, caption="Model Graph Visualization")

def show_json_output():
    if 'mdf_model_json' in st.session_state:
        st.subheader("JSON Output")
        st.json(st.session_state.mdf_model_json)

def update_view_tabs(mdf_model, param_inputs):
    tab1, tab2, tab3, tab4 = st.tabs(["Simulation Results", "MDF Graph", "Json Output", "Yaml Editor"])
    with tab1:
        show_simulation_results()
    with tab2:
        show_mdf_graph()
    with tab3:
        show_json_output()
    with tab4:
        changes_in_yaml_to_update_model()

def parameter_form_to_update_model_and_view(mdf_model, param_inputs):
    mod_graph = mdf_model.graphs[0]
    nodes = mod_graph.nodes[0]
    parameters = nodes.parameters
    form = st.form(key="parameter_form")
    for i, param in enumerate(parameters):
        if isinstance(param.value, str) or param.value is None:
            continue  
        key = f"{param.id}_{i}"
        if mdf_model.metadata:
            param_inputs[param.id] = float(form.text_input(f"{param.metadata.get('description', param.id)} ({param.id})", value=param.value, key=key))
        else:
            param_inputs[param.id] = float(form.text_input(f"{param.id}", value=param.value, key=key))
    param_inputs["Simulation Duration (s)"] = float(form.text_input("Simulation Duration (s)", value=st.session_state.param_inputs.get("Simulation Duration (s)", 10), key="sim_duration"))
    param_inputs["Time Step (s)"] = float(form.text_input("Time Step (s)", value=st.session_state.param_inputs.get("Time Step (s)", 0.1), key="time_step"))
    
    run_button = form.form_submit_button("Run Simulation")
    if run_button:
        st.session_state.run_button_clicked = True
        for i in param_inputs:
            st.session_state.param_inputs[i] = float(param_inputs[i])
        st.session_state.chart_data = run_simulation(param_inputs, mdf_model)
        st.session_state.mdf_model_graph = mdf_model.to_graph_image(engine="dot", output_format="png", view_on_render=False, level=3, filename_root=mdf_model.id, only_warn_on_fail=(os.name == "nt"))
        st.session_state.mdf_model_json = mdf_model.to_json()
        st.session_state.update_button_clicked = False

def upload_file_and_load_to_model():
    uploaded_file = st.file_uploader("Choose a JSON/YAML/BSON file", type=["json", "yaml", "bson"])
    if uploaded_file is not None:
        filename = uploaded_file.name
        with open(filename, "wb") as f:
            f.write(uploaded_file.getbuffer())
        mdf_model = load_mdf(filename)
        st.session_state.original_mdf_model = mdf_model  # Save the original model
        st.session_state.mdf_model_yaml = mdf_model  # Save the current model state
        return mdf_model


def changes_in_yaml_to_update_model():
    response = code_editor(st.session_state.mdf_model_yaml.to_yaml(),  height=height, lang="yaml", theme=theme, shortcuts=shortcuts, focus=focus, buttons=editor_btns, options={"wrap": wrap})
    st.write("First, make changes in the YAML code above, press Ctrl + Enter (or Run button inside code editor) and then click on the 'Update Model' button below to update the model. Sometimes you have to follow these steps twice. Currently, the plot, mdf_graph and json output get updtes, but parameter input form needs to be worked on")
    if len(response['id']) != 0 and (response['type'] == "selection" or response['type'] == "submit"):
        code_text = response['text']
        update_button = st.button("Update Model")
        if update_button:
            mdf_model = Model.from_yaml(code_text)
            mdf_model_to_paramters_dictionary(st.session_state.param_inputs, mdf_model)
            st.session_state.mdf_model_yaml = mdf_model
            st.session_state.chart_data = run_simulation(st.session_state.param_inputs, mdf_model)
            st.session_state.mdf_model_graph = mdf_model.to_graph_image(engine="dot", output_format="png", view_on_render=False, level=3, filename_root=mdf_model.id, only_warn_on_fail=(os.name == "nt"))
            st.session_state.mdf_model_json = mdf_model.to_json()
            st.session_state.update_button_clicked = True
            
        


def main():
    mdf_model = upload_file_and_load_to_model()
    st.write("Working on making YAML editor.")
    if mdf_model:
        st.session_state.mdf_model = mdf_model
        if 'run_button_clicked' not in st.session_state:
            st.session_state.run_button_clicked = False
        if 'update_button_clicked' not in st.session_state:
            st.session_state.update_button_clicked = False
        param_inputs = {}
        if mdf_model.metadata:
            preferred_duration = float(mdf_model.metadata.get("preferred_duration", 10))
            preferred_dt = float(mdf_model.metadata.get("preferred_dt", 0.1))
        else:
            preferred_duration = 100
            preferred_dt = 0.1
        param_inputs["Simulation Duration (s)"] = preferred_duration
        param_inputs["Time Step (s)"] = preferred_dt
        st.session_state.param_inputs = param_inputs
        parameter_form_to_update_model_and_view(st.session_state.mdf_model, st.session_state.param_inputs)
        update_view_tabs(mdf_model, param_inputs)
        
if __name__ == "__main__":
    main()