api.py

"""
TSNet Transient Simulator — FastAPI Backend
===========================================
Designed for use as a Custom GPT Action endpoint.
Deployed on Hugging Face Spaces (port 7860).

Endpoints:
  GET  /              — API info
  GET  /health        — Health check (for GPT Action verification)
  POST /simulate      — Run MOC transient simulation
  POST /network-info  — Inspect .inp file without simulating
"""

import sys
import os
import tempfile
import traceback

# Ensure bundled tsnet package is importable
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from typing import Optional, Literal
import numpy as np

# ── TSNet imports ─────────────────────────────────────────────────────────────
try:
    from tsnet.network.model import TransientModel
    from tsnet.simulation import MOCSimulator, Initializer
    TSNET_OK = True
except Exception as e:
    TSNET_OK = False
    TSNET_ERR = str(e)

# ── App ───────────────────────────────────────────────────────────────────────
app = FastAPI(
    title="TSNet Transient Simulator API",
    description="Water-hammer MOC simulation API for Custom GPT Actions",
    version="1.0.0",
)

# Allow GPT Action calls (OpenAI origin)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


# ── Request / Response models ─────────────────────────────────────────────────
class SimRequest(BaseModel):
    inp_content: str = Field(
        ...,
        description="Full EPANET .inp file content as a plain string"
    )
    wave_speed: float = Field(
        1200.0, ge=100.0, le=2000.0,
        description="Acoustic wave speed in m/s (100–2000)"
    )
    duration: float = Field(
        10.0, ge=0.5, le=300.0,
        description="Simulation duration in seconds"
    )
    dt: float = Field(
        0.01, ge=0.001, le=1.0,
        description="Time step in seconds (must satisfy CFL: dt ≤ L_min/2a)"
    )
    engine: Literal["DD", "PDD"] = Field(
        "DD",
        description="Hydraulic engine: DD = demand-driven, PDD = pressure-dependent"
    )
    friction: Literal["steady", "quasi-steady", "unsteady"] = Field(
        "steady",
        description="Friction model"
    )
    event_type: Optional[Literal[
        "valve_closure", "valve_opening", "pump_shutoff", "none"
    ]] = Field(
        "none",
        description="Transient event type"
    )
    element_id: Optional[str] = Field(
        None,
        description="Valve or pump element name from the .inp file"
    )
    tc: float = Field(0.0, ge=0.0, description="Event duration in seconds")
    ts: float = Field(2.0, ge=0.0, description="Event start time in seconds")
    se: float = Field(0.0, ge=0.0, le=1.0,
                      description="Final open fraction (0=closed, 1=open)")
    m:  float = Field(1.0, ge=0.1, description="Shape constant (1=linear)")


class NetworkInfoRequest(BaseModel):
    inp_content: str = Field(..., description="EPANET .inp file content")


# ── Helpers ───────────────────────────────────────────────────────────────────
def write_temp_inp(content: str) -> str:
    tf = tempfile.NamedTemporaryFile(
        delete=False, suffix=".inp", mode="w", encoding="utf-8"
    )
    tf.write(content)
    tf.close()
    return tf.name


def safe_float(v) -> float:
    f = float(v)
    return f if np.isfinite(f) else 0.0


# ── Routes ────────────────────────────────────────────────────────────────────
@app.get("/")
def root():
    return {
        "name": "TSNet Transient Simulator API",
        "version": "1.0.0",
        "tsnet_available": TSNET_OK,
        "endpoints": {
            "GET  /health":       "Health check",
            "POST /simulate":     "Run MOC transient simulation",
            "POST /network-info": "Inspect .inp file without simulating",
        },
        "docs": "/docs",
    }


@app.get("/health")
def health():
    if not TSNET_OK:
        raise HTTPException(status_code=503,
                            detail=f"TSNet unavailable: {TSNET_ERR}")
    return {"status": "healthy", "tsnet": "ok"}


@app.post("/network-info")
def network_info(req: NetworkInfoRequest):
    """
    Parse an EPANET .inp file and return network metadata.
    Use this before /simulate to discover element names.
    """
    if not TSNET_OK:
        raise HTTPException(503, f"TSNet unavailable: {TSNET_ERR}")

    inp_path = write_temp_inp(req.inp_content)
    try:
        tm = TransientModel(inp_path)
        return {
            "num_junctions": tm.num_junctions,
            "num_pipes":     tm.num_pipes,
            "num_valves":    tm.num_valves,
            "num_pumps":     tm.num_pumps,
            "num_reservoirs": tm.num_reservoirs,
            "pipe_ids":   [n for n, _ in tm.pipes()],
            "valve_ids":  [n for n, _ in tm.valves()],
            "pump_ids":   [n for n, _ in tm.pumps()],
            "node_ids":   [n for n, _ in tm.nodes()],
        }
    except Exception as e:
        raise HTTPException(400, f"Failed to parse .inp file: {e}")
    finally:
        os.unlink(inp_path)


@app.post("/simulate")
def simulate(req: SimRequest):
    """
    Run a full MOC water-hammer simulation.
    Returns pressure heads, velocities, and flow-rates at all nodes and pipes,
    plus a surge summary suitable for engineering interpretation.
    """
    if not TSNET_OK:
        raise HTTPException(503, f"TSNet unavailable: {TSNET_ERR}")

    inp_path = write_temp_inp(req.inp_content)

    try:
        # ── Load & configure ─────────────────────────────────────────────────
        tm = TransientModel(inp_path)
        tm.set_wavespeed(req.wave_speed)
        tm.set_time(req.duration, req.dt)

        # ── Apply transient event ────────────────────────────────────────────
        rule = [req.tc, req.ts, req.se, req.m]

        if req.event_type == "valve_closure":
            if not req.element_id:
                raise HTTPException(400, "element_id required for valve_closure")
            tm.valve_closure(req.element_id, rule)

        elif req.event_type == "valve_opening":
            if not req.element_id:
                raise HTTPException(400, "element_id required for valve_opening")
            tm.valve_opening(req.element_id, rule)

        elif req.event_type == "pump_shutoff":
            if not req.element_id:
                raise HTTPException(400, "element_id required for pump_shutoff")
            if req.tc <= 0:
                raise HTTPException(400, "tc must be > 0 for pump_shutoff")
            tm.pump_shut_off(req.element_id, rule)

        # ── Initialise & simulate ────────────────────────────────────────────
        tm = Initializer(tm, 0.0, engine=req.engine)
        rpath = os.path.join(tempfile.gettempdir(), "tsnet_api_result")
        tm_r  = MOCSimulator(tm, rpath, req.friction)

        tt = tm_r.simulation_timestamps
        dt_actual = safe_float(tt[1] - tt[0]) if len(tt) > 1 else req.dt

        # ── Build results ────────────────────────────────────────────────────
        node_results = {}
        surge_summary = []

        for name, node in tm_r.nodes():
            h = node._head
            if h is None:
                continue
            h0   = safe_float(h[0])
            hmax = safe_float(h.max())
            hmin = safe_float(h.min())
            surge = round(hmax - h0, 3)
            node_results[name] = {
                "initial_head_m": round(h0,   3),
                "max_head_m":     round(hmax, 3),
                "min_head_m":     round(hmin, 3),
                "surge_m":        surge,
                "head_timeseries": [round(safe_float(v), 3) for v in h],
            }
            if abs(surge) > 0.01:
                surge_summary.append({
                    "node":     name,
                    "surge_m":  surge,
                    "max_head_m": round(hmax, 3),
                })

        pipe_results = {}
        for name, pipe in tm_r.pipes():
            vs = pipe.start_node_velocity
            ve = pipe.end_node_velocity
            qs = pipe.start_node_flowrate
            pipe_results[name] = {
                "start_velocity_initial_ms": round(safe_float(vs[0]), 4),
                "start_velocity_max_ms":     round(safe_float(vs.max()), 4),
                "start_velocity_min_ms":     round(safe_float(vs.min()), 4),
                "end_velocity_initial_ms":   round(safe_float(ve[0]), 4),
                "start_flowrate_initial_m3s": round(safe_float(qs[0]), 6),
                "start_flowrate_max_m3s":     round(safe_float(qs.max()), 6),
            }

        # Sort surge summary by magnitude descending
        surge_summary.sort(key=lambda x: abs(x["surge_m"]), reverse=True)

        return {
            "status": "ok",
            "simulation": {
                "duration_s":   req.duration,
                "dt_actual_s":  round(dt_actual, 5),
                "num_steps":    len(tt),
                "wave_speed_ms": req.wave_speed,
                "event_type":   req.event_type,
                "element_id":   req.element_id,
                "friction":     req.friction,
            },
            "surge_summary": surge_summary,
            "nodes":  node_results,
            "pipes":  pipe_results,
        }

    except HTTPException:
        raise
    except ValueError as e:
        raise HTTPException(400, f"Simulation parameter error: {e}")
    except Exception as e:
        raise HTTPException(500, f"Simulation failed: {e}\n{traceback.format_exc()}")
    finally:
        os.unlink(inp_path)