import streamlit as st
import wntr
import tempfile
import os
import pandas as pd
import plotly.graph_objects as go

st.set_page_config(page_title="EPANET Simulation with WNTR", layout="wide")
st.title("💧 EPANET Simulation with WNTR")

uploaded_file = st.file_uploader("Upload your EPANET .inp file", type=["inp"])

if uploaded_file:
    with tempfile.NamedTemporaryFile(delete=False, suffix=".inp") as temp:
        temp.write(uploaded_file.read())
        inp_path = temp.name

    st.success("File uploaded. Running simulation...")

    wn = wntr.network.WaterNetworkModel(inp_path)
    sim = wntr.sim.EpanetSimulator(wn)
    results = sim.run_sim()

    st.subheader("✅ Simulation Complete")

    # Custom groups
    pressure = results.node["pressure"]
    head = results.node["head"]
    tank_levels = results.node["head"][[name for name, _ in wn.tanks()]]
    demand = results.node["demand"]

    simulation_groups = {
        "🟦 Node Pressures": {
            "data": pressure,
            "ylabel": "Pressure (m)",
            "color": "#dbeafe"
        },
        "🟩 Tank Levels": {
            "data": tank_levels,
            "ylabel": "Level (m)",
            "color": "#d1fae5"
        },
        "🟧 Demands": {
            "data": demand,
            "ylabel": "Demand (m³/s)",
            "color": "#fef3c7"
        }
    }

    for title, group in simulation_groups.items():
        df = group["data"]
        color = group["color"]

        st.markdown(
            f"<div style='background-color:{color}; padding: 10px; border-radius: 8px;'>"
            f"<h4 style='margin: 0px;'>{title}</h4></div>",
            unsafe_allow_html=True
        )

        csv = df.to_csv().encode("utf-8")
        st.download_button(
            label=f"⬇️ Download {title} as CSV",
            data=csv,
            file_name=f"{title.replace(' ', '_')}.csv",
            mime="text/csv"
        )

        with st.expander(f"Show plots for {title}"):
            for node_id in df.columns:
                try:
                    node = wn.get_node(node_id)
                    node_type = type(node).__name__
                except:
                    node_type = "Unknown"

                time_minutes = df.index / 60
                fig = go.Figure()
                fig.add_trace(go.Scatter(x=time_minutes, y=df[node_id],
                                         mode='lines',
                                         name=node_id))
                fig.update_layout(
                    title=f"{title} for {node_type} '{node_id}'",
                    xaxis_title="Time (min)",
                    yaxis_title=group["ylabel"],
                    height=400
                )
                st.plotly_chart(fig, use_container_width=True)

    # ───────────────────────────────────────────────────────
    # 🔍 Full Auto Node + Link Results (All Types)
    # ───────────────────────────────────────────────────────
    st.header("📍 Node Results (All Types)")
    for result_type in results.node.keys():
        df = results.node[result_type]
        st.markdown(f"### 🟦 {result_type.title()} at Nodes")
        csv = df.to_csv().encode("utf-8")
        st.download_button(
            label=f"⬇️ Download {result_type} (nodes)",
            data=csv,
            file_name=f"node_{result_type}.csv",
            mime="text/csv"
        )
        with st.expander(f"Show plots for node result: {result_type}"):
            for node_id in df.columns:
                try:
                    node = wn.get_node(node_id)
                    node_type = type(node).__name__
                except:
                    node_type = "Unknown"
                time_minutes = df.index / 60
                fig = go.Figure()
                fig.add_trace(go.Scatter(x=time_minutes, y=df[node_id],
                                         mode='lines',
                                         name=node_id))
                fig.update_layout(
                    title=f"{result_type.title()} at {node_type} '{node_id}'",
                    xaxis_title="Time (min)",
                    yaxis_title=result_type.title(),
                    height=400
                )
                st.plotly_chart(fig, use_container_width=True)

    st.header("🔗 Link Results (All Types)")
    for result_type in results.link.keys():
        df = results.link[result_type]
        st.markdown(f"### 🟧 {result_type.title()} at Links")
        csv = df.to_csv().encode("utf-8")
        st.download_button(
            label=f"⬇️ Download {result_type} (links)",
            data=csv,
            file_name=f"link_{result_type}.csv",
            mime="text/csv"
        )
        with st.expander(f"Show plots for link result: {result_type}"):
            for link_id in df.columns:
                try:
                    link = wn.get_link(link_id)
                    link_type = type(link).__name__
                except:
                    link_type = "Unknown"
                time_minutes = df.index / 60
                fig = go.Figure()
                fig.add_trace(go.Scatter(x=time_minutes, y=df[link_id],
                                         mode='lines',
                                         name=link_id))
                fig.update_layout(
                    title=f"{result_type.title()} at {link_type} '{link_id}'",
                    xaxis_title="Time (min)",
                    yaxis_title=result_type.title(),
                    height=400
                )
                st.plotly_chart(fig, use_container_width=True)