"""Cross-check SPipe's server analyses against USEPA WNTR (an independent EPANET
front end). Skipped when wntr is not installed (it is a dev-only dependency).

WNTR works in SI (m, m³/s); Net1/Net3 are US units (psi, GPM), so values are
converted before comparing. WNTR bundles its own EPANET build, so small
numerical differences are allowed.
"""
import os
import tempfile
from pathlib import Path

import numpy as np
import pytest

from backend.batch import Network
from backend.epanet import EN
from backend.jobs import job_criticality, job_fireflow

wntr = pytest.importorskip("wntr")

PUBLIC = Path(__file__).resolve().parents[2] / "public"
PSI_TO_M = 0.3048 / 0.4333  # EPANET's own psi <-> ft of water factor
GPM_TO_M3S = 6.30902e-5


def wntr_pressure(wn) -> "object":
    with tempfile.TemporaryDirectory() as tmp:
        res = wntr.sim.EpanetSimulator(wn).run_sim(file_prefix=os.path.join(tmp, "w"))
    return res.node["pressure"]


def spipe_pressure(inp: bytes) -> dict:
    with Network(inp) as n:
        js = n.junction_indices()
        r = n.simulate(node_props=(EN.PRESSURE,), nodes=js)
        return {n.node_ids[j - 1]: r.nodes[EN.PRESSURE][:, k] * PSI_TO_M for k, j in enumerate(js)}, r.times


@pytest.mark.parametrize("name", ["Net1.inp", "Net3.inp"])
def test_baseline_pressures_match(name):
    path = PUBLIC / name
    ours, times = spipe_pressure(path.read_bytes())
    p = wntr_pressure(wntr.network.WaterNetworkModel(str(path)))
    diffs = []
    for nid, series in ours.items():
        w = p[nid].reindex(times).to_numpy()
        diffs.append(np.nanmax(np.abs(series - w)))
    assert max(diffs) < 0.01, f"max pressure difference {max(diffs):.3f} m"


def test_pipe_criticality_matches_closing_pipes_in_wntr():
    path = PUBLIC / "Net1.inp"
    ours = {e["id"]: e for e in job_criticality(path.read_bytes(), {"threshold": 40})["links"]}
    checked = 0
    for pid, e in ours.items():
        if e.get("disconnected") or e.get("minPressure") is None:
            continue  # isolated nodes: both engines report meaningless negative heads
        wn = wntr.network.WaterNetworkModel(str(path))
        wn.get_link(pid).initial_status = wntr.network.LinkStatus.Closed
        p = wntr_pressure(wn)
        junctions = [n for n, _ in wn.junctions()]
        wmin = float(p[junctions].min().min())
        assert abs(e["minPressure"] * PSI_TO_M - wmin) < 0.01, f"pipe {pid}: {e['minPressure'] * PSI_TO_M:.2f} vs {wmin:.2f} m"
        checked += 1
    assert checked >= 10


def test_fire_flow_pressure_drop_matches_extra_demand_in_wntr():
    path = PUBLIC / "Net1.inp"
    flow_gpm = 500.0
    r = job_fireflow(path.read_bytes(), {"nodes": ["22"], "flow": flow_gpm, "start": 0, "end": 24})
    ours = r["hydrants"][0]["firePressure"] * PSI_TO_M
    wn = wntr.network.WaterNetworkModel(str(path))
    # A demand without a pattern would get Net1's default pattern (up to 1.6x); fire flow is constant.
    wn.add_pattern("fire_const", [1.0])
    wn.get_node("22").add_demand(flow_gpm * GPM_TO_M3S, "fire_const", "fire")
    p = wntr_pressure(wn)
    assert abs(ours - float(p["22"].min())) < 0.01