import gradio as gr
import wntr
import tempfile
import pandas as pd
import plotly.graph_objects as go
import os
import shutil

def run_simulation(inp_file):
    try:
        # Create temporary working directory
        work_dir = tempfile.mkdtemp()
        inp_path = os.path.join(work_dir, "network.inp")

        # Save uploaded file to temp path
        shutil.copy(inp_file.name, inp_path)

        # Load and simulate the EPANET model
        wn = wntr.network.WaterNetworkModel(inp_path)
        sim = wntr.sim.EpanetSimulator(wn)
        results = sim.run_sim()

        output_list = []

        # Simulation result groups
        simulation_groups = {
            "Node Pressures": {
                "data": results.node["pressure"],
                "ylabel": "Pressure (m)"
            },
            "Tank Levels": {
                "data": results.node["head"][[name for name, _ in wn.tanks()]],
                "ylabel": "Tank Head (m)"
            },
            "Demands": {
                "data": results.node["demand"],
                "ylabel": "Demand (m³/s)"
            }
        }

        # Append simulation group results
        for title, group in simulation_groups.items():
            df = group["data"]
            csv_path = os.path.join(work_dir, f"{title.replace(' ', '_')}.csv")
            df.to_csv(csv_path)

            fig = go.Figure()
            for col in df.columns:
                fig.add_trace(go.Scatter(
                    x=df.index / 60,  # seconds → minutes
                    y=df[col],
                    mode="lines",
                    name=str(col)
                ))
            fig.update_layout(
                title=title,
                xaxis_title="Time (min)",
                yaxis_title=group["ylabel"],
                height=400
            )

            output_list.extend([csv_path, df.reset_index(), fig])

        # Tank metadata
        tank_data = [(name, tank.elevation, tank.min_level, tank.max_level, tank.diameter)
                     for name, tank in wn.tanks()]
        df_tanks = pd.DataFrame(tank_data, columns=["Tank ID", "Elevation", "Min Level", "Max Level", "Diameter"])
        output_list.append(df_tanks)

        return tuple(output_list)

    except Exception as e:
        return [f"❌ ERROR: {str(e)}"] + [None] * 9

# Gradio interface
input_file = gr.File(label="📁 Upload .inp File", file_types=[".inp"])

outputs = [
    # Node Pressures
    gr.File(label="📥 Pressure CSV"), gr.Dataframe(label="Pressure Table"), gr.Plot(label="Pressure Plot"),
    # Tank Levels
    gr.File(label="📥 Tank Level CSV"), gr.Dataframe(label="Tank Level Table"), gr.Plot(label="Tank Level Plot"),
    # Demands
    gr.File(label="📥 Demand CSV"), gr.Dataframe(label="Demand Table"), gr.Plot(label="Demand Plot"),
    # Tank metadata table
    gr.Dataframe(label="📊 Tank Metadata")
]

demo = gr.Interface(
    fn=run_simulation,
    inputs=input_file,
    outputs=outputs,
    title="💧 EPANET Simulation Viewer (Gradio)",
    description="Upload a valid EPANET `.inp` file. This app runs hydraulic simulation using WNTR and returns pressure, tank levels, demands, and tank properties.",
    allow_flagging="never"
)

if __name__ == "__main__":
    demo.launch()