Update app.py

app.py

@@ -1,19 +1,67 @@
 import streamlit as st
-from tsnet_agents import run_agentic_query
-import base64
-from PIL import Image
-import io
-
-st.set_page_config(page_title="TSNet + LLM Final App", layout="wide")
-st.title("🌊 TSNet Transient Simulation with LLM + Plot + RAG")
-
-query = st.text_area("Ask a question or describe a valve closure:")
-if st.button("Run"):
-    response = run_agentic_query(query)
-    if response["type"] == "plot":
-        img_data = base64.b64decode(response["image"])
-        image = Image.open(io.BytesIO(img_data))
-        st.image(image, caption="Simulation Result")
-    else:
-        st.markdown("**LLM Answer:**")
-        st.write(response["content"])
+import tempfile
+import matplotlib.pyplot as plt
+import pickle
+
+st.set_page_config(page_title="TSNet Transient Simulator", layout="wide")
+st.title("🌊 TSNet Transient Simulation App")
+
+# Step 1: Upload EPANET .inp File
+st.header("1. Upload EPANET .inp File")
+uploaded_file = st.file_uploader("Upload your .inp file", type=["inp"])
+
+if uploaded_file is not None:
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".inp") as temp_inp:
+        temp_inp.write(uploaded_file.read())
+        inp_path = temp_inp.name
+
+    st.success("File uploaded and saved.")
+    st.header("2. Set Simulation Parameters")
+
+    wave_speed = st.number_input("Wave Speed (m/s)", value=1000.0)
+    dt = st.number_input("Time Step (s)", value=0.01)
+    duration = st.number_input("Simulation Duration (s)", value=10.0)
+
+    if st.button("Run TSNet Simulation"):
+        import tsnet
+        from tsnet.network.model import TransientModel
+        from tsnet.simulation import MOCSimulator
+
+        net = TransientModel(inp_path)
+        net.set_wavespeed(wave_speed)
+        net.set_time(duration, dt)
+
+        for pipe in net.pipes():
+            if not hasattr(pipe, 'initial_head'):
+                pipe.initial_head = (pipe.start_node.head + pipe.end_node.head) / 2
+
+        sim = MOCSimulator(net, 'results')
+        with open("results.obj", "wb") as f:
+            pickle.dump(sim, f)
+
+        st.session_state["sim_results"] = "results.obj"
+        st.success("Simulation complete.")
+
+if "sim_results" in st.session_state:
+    st.header("3. Plot Results")
+    sim = pickle.load(open(st.session_state["sim_results"], "rb"))
+
+    node_id = st.selectbox("Select Node to Plot Head", sim.node_name_list)
+    pipe_id = st.selectbox("Select Pipe to Plot Flow", sim.pipe_name_list)
+
+    head = sim.get_node_head(node_id)
+    flow = sim.get_link_flow(pipe_id)
+
+    fig1, ax1 = plt.subplots()
+    ax1.plot(sim.Time, head)
+    ax1.set_title(f"Head at Node {node_id}")
+    ax1.set_xlabel("Time (s)")
+    ax1.set_ylabel("Head (m)")
+    st.pyplot(fig1)
+
+    fig2, ax2 = plt.subplots()
+    ax2.plot(sim.Time, flow)
+    ax2.set_title(f"Flow in Pipe {pipe_id}")
+    ax2.set_xlabel("Time (s)")
+    ax2.set_ylabel("Flow (L/s)")
+    st.pyplot(fig2)