app.py

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
from typing import List, Optional
import uvicorn
import tempfile
import os
import wntr
import matplotlib.pyplot as plt
import pandas as pd

app = FastAPI(title="EPANET Simulation API")

@app.post("/simulate")
async def simulate_epanet(
    inp_file: UploadFile = File(...),
    simulation_duration: Optional[int] = 24,
    output_results: List[str] = ["pressure", "flow"],
    nodes_to_report: Optional[List[str]] = None,
    links_to_report: Optional[List[str]] = None,
    generate_plot: Optional[bool] = True
):
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".inp") as temp_file:
            temp_file.write(await inp_file.read())
            inp_path = temp_file.name

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

        response_data = {}

        if "pressure" in output_results:
            pressures = results.node['pressure']
            if nodes_to_report:
                pressures = pressures[nodes_to_report]
            response_data["pressure"] = pressures.describe().to_dict()

        if "flow" in output_results:
            flows = results.link['flowrate']
            if links_to_report:
                flows = flows[links_to_report]
            response_data["flow"] = flows.describe().to_dict()

        if generate_plot:
            plot_dir = tempfile.mkdtemp()
            plot_paths = []
            if "pressure" in output_results and nodes_to_report:
                for node in nodes_to_report:
                    plt.figure()
                    results.node['pressure'][node].plot(title=f"Pressure at {node}")
                    img_path = os.path.join(plot_dir, f"pressure_{node}.png")
                    plt.savefig(img_path)
                    plot_paths.append(img_path)
                    plt.close()
            if "flow" in output_results and links_to_report:
                for link in links_to_report:
                    plt.figure()
                    results.link['flowrate'][link].plot(title=f"Flow in {link}")
                    img_path = os.path.join(plot_dir, f"flow_{link}.png")
                    plt.savefig(img_path)
                    plot_paths.append(img_path)
                    plt.close()
            response_data["plots"] = plot_paths

        return JSONResponse(content=response_data)

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        if os.path.exists(inp_path):
            os.remove(inp_path)