import streamlit as st
import tempfile
import os
import pickle
import matplotlib.pyplot as plt
from tsnet import network
from tsnet import simulation
from tsnet import postprocessing 
from tsnet import utils

from tsnet.network.model import TransientModel
from tsnet.network.elements import Control
from tsnet.simulation import MOCSimulator

st.set_page_config(page_title="TSNet Transient Simulator", layout="wide")
st.title("🌊 TSNet Transient Analysis App with Natural Language Assist")

# Step 1: Upload EPANET .inp File
st.header("1️⃣ Upload EPANET .inp File")
uploaded_file = st.file_uploader("Upload your EPANET .inp file", type=["inp"])

if uploaded_file:
    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 processed.")

    # Step 2: Initialize Model
    try:
        net = TransientModel(inp_path)
        st.success("✅ Transient model created.")
    except Exception as e:
        st.error(f"🚨 Error loading model: {e}")
        st.stop()

    # Step 3: Configure Simulation
    st.header("2️⃣ Configure 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)

    net.set_wavespeed(wave_speed)
    net.set_time(duration, dt)

    # Step 4: Define Control (Valve Closure)
    st.header("3️⃣ Define Transient Event (Valve Closure)")
    valve_node = st.text_input("Enter valve node ID:", value="2")
    start_time = st.number_input("Closure start time (s):", value=2.0)
    closure_duration = st.number_input("Closure duration (s):", value=1.0)

    if st.button("Add Valve Closure"):
        control = Control()
        control.add_demand_control(
            nodeid=valve_node,
            time_series=[(start_time, 1.0), (start_time + closure_duration, 0.0)]
        )
        net.add_control(control)
        st.success(f"✅ Demand closure added at node {valve_node}")

    # Step 5: Assign Initial Head
    st.header("4️⃣ Initial Condition Calculation")
    if st.button("Compute Initial Head"):
        try:
            for pipe in net.pipes():
                if not hasattr(pipe, 'initial_head'):
                    start_node = pipe.start_node
                    end_node = pipe.end_node
                    pipe.initial_head = (start_node.head + end_node.head) / 2
            st.success("✅ Initial head assigned to all pipes.")
        except Exception as e:
            st.error(f"🚨 Error in assigning initial head: {e}")

    # Step 6: Run Simulation
    if st.button("🚀 Run Simulation"):
        try:
            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, output_dir="results")
            with open("results.obj", "wb") as f:
                pickle.dump(sim, f)
            st.session_state["result_path"] = "results.obj"
            st.success("✅ Simulation complete.")
        except Exception as e:
            st.error(f"❌ Simulation failed: {e}")

    # Step 7: Plot Results
    if "result_path" in st.session_state:
        st.header("5️⃣ View Results")
        sim = pickle.load(open(st.session_state["result_path"], "rb"))
        node_to_plot = st.selectbox("Select node to plot head:", sim.node_name_list)
        link_to_plot = st.selectbox("Select pipe to plot flow:", sim.pipe_name_list)

        head = sim.get_node_head(node_to_plot)
        flow = sim.get_link_flow(link_to_plot)

        fig1, ax1 = plt.subplots()
        ax1.plot(sim.Time, head)
        ax1.set_title(f"Head at Node {node_to_plot}")
        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 {link_to_plot}")
        ax2.set_xlabel("Time (s)")
        ax2.set_ylabel("Flow (L/s)")
        st.pyplot(fig2)

else:
    st.info("📂 Please upload a valid EPANET .inp file to begin.")