#!/usr/bin/env python3
"""Render a SWMM model's spatial layout (network map) as a PNG.

Companion to ``plot_rain_runoff_si.py`` (which is the hydrograph view).
``aiswmm plot`` answers "what does the simulated flow look like over
time?". ``aiswmm map`` answers "what does the network look like in
space?" — the question every reviewer asks before they trust a SWMM
model. Both scripts share the same skill (``swmm-plot``) and the same
stylesheet (``plot_style.apply_style()``: the vendored Nature spec, 89 mm
single column, 7 pt sans-serif, ticks out, no grid, Wong palette, no
title) so a paper with both figures reads as one consistent diagnostic
pair. Output is the ``--out-png`` preview plus its vector twin
``<stem>.pdf`` (the submission file).

The data path is deliberately two-tier:

* **Preferred — SWMManywhere geoparquet artefacts.** When the run
  directory was produced by the swmm-anywhere chain (PRD
  swmmanywhere_integration), the runner copies
  ``nodes.geoparquet`` / ``edges.geoparquet`` /
  ``subcatchments.geoparquet`` under the canonical upstream box
  ``<run-dir>/10_upstream/swmmanywhere/`` (ADR-0004; older runs may carry
  it under the legacy flat ``<run-dir>/10_swmmanywhere/``).
  These carry the real WGS84 polygons SWMManywhere downloaded from OSM
  plus extra columns (``node_type``, ``outfall_id``) that drive the
  per-outfall colouring. We load them via ``geopandas`` *lazily* so the
  default aiswmm install (no ``[anywhere]`` extra) still has a working
  ``aiswmm map``.
* **Fallback — INP text parsing.** Every SWMM ``.inp`` file ships with
  ``[COORDINATES]`` (nodes), ``[VERTICES]`` (conduit shape points),
  ``[Polygons]`` (subcatchment boundaries), and ``[SUBCATCHMENTS]``
  (subcatchment→outlet linkage). Pure-text parsing of these four
  sections gives us everything the renderer needs without pulling in
  geopandas/pyarrow. This is the path that runs on a bare aiswmm
  install or against a hand-built INP that never touched
  SWMManywhere.

Colouring strategy: every conduit is traced upstream→downstream to its
terminal outfall via the ``[CONDUITS]`` adjacency. All conduits in the
same drainage area share that outfall's colour (the Wong palette's
categorical cycle), so a sub-network jumps out visually. Subcatchments
are tinted by the colour of the outfall their outlet drains to
(transitively). Junctions are small grey dots; outfalls are black ``★``
markers (the palette's first colour, kept out of the sub-network cycle)
so reviewers can find the model discharge points in one glance.

CLI surface (the driver — ``aiswmm map`` — forwards these):

    --inp <path>      explicit INP (overrides discovery)
    --run-dir <path>  run directory (used only to find the geoparquet trio)
    --out-png <path>  output PNG path (required); <stem>.pdf is written beside it
    --dpi <int>       resolution of the PNG preview (default 450; the PDF is vector)
    --no-subcatchments  skip the polygon layer
    --no-vertices       draw conduits as straight lines (ignore [VERTICES])
"""
from __future__ import annotations

import argparse
import sys
from collections import defaultdict, deque
from pathlib import Path
from typing import Any

# Headless backend — same convention as plot_rain_runoff_si.py.
import matplotlib

matplotlib.use("Agg")
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt

# Shared style module next to this script (tests load the script through
# ``importlib.util.spec_from_file_location``, which does not put this
# directory on sys.path).
_HERE = str(Path(__file__).resolve().parent)
if _HERE not in sys.path:
    sys.path.insert(0, _HERE)

from plot_style import (  # noqa: E402
    CATEGORY_CYCLE,
    WONG,
    apply_style,
    legend_outside,
    map_figsize,
    save_figure,
)


# --------------------------------------------------------------------- #
# INP text parsing (the fallback data source).
#
# Each helper consumes the full INP text once and returns the section
# it owns. We do NOT use a single multi-pass tokeniser — the script's
# value is that it works without any extra deps, so adding even a
# minimal class hierarchy would be over-engineered.
# --------------------------------------------------------------------- #


def _read_section_lines(inp_text: str, section: str) -> list[list[str]]:
    """Return a list of ``parts`` lists for every non-comment row in
    ``[section]``. ``parts`` is the result of ``.split()`` on the row,
    so callers index by position. Empty lists and comment lines are
    dropped.
    """
    rows: list[list[str]] = []
    in_section = False
    upper_section = f"[{section.upper()}]"
    for raw in inp_text.splitlines():
        line = raw.strip()
        if not line:
            continue
        if line.upper() == upper_section:
            in_section = True
            continue
        if in_section and line.startswith("[") and line.endswith("]"):
            break
        if not in_section:
            continue
        if line.startswith(";"):
            continue
        rows.append(line.split())
    return rows


def parse_inp_coordinates(inp_text: str) -> dict[str, tuple[float, float]]:
    """``[COORDINATES]`` -> {node_name: (x, y)}.

    SWMM stores ``;;Name X Y`` rows. We tolerate ``;`` and ``;;``
    comments and silently skip unparseable rows (some hand-edited INPs
    have ragged columns).
    """
    coords: dict[str, tuple[float, float]] = {}
    for parts in _read_section_lines(inp_text, "COORDINATES"):
        if len(parts) < 3:
            continue
        try:
            coords[parts[0]] = (float(parts[1]), float(parts[2]))
        except ValueError:
            continue
    return coords


def parse_inp_vertices(inp_text: str) -> dict[str, list[tuple[float, float]]]:
    """``[VERTICES]`` -> {conduit_name: [(x, y), ...]}.

    ``[VERTICES]`` is optional in SWMM; many INPs draw conduits as
    straight lines and skip the section entirely. Returns an empty
    dict in that case.
    """
    verts: dict[str, list[tuple[float, float]]] = defaultdict(list)
    for parts in _read_section_lines(inp_text, "VERTICES"):
        if len(parts) < 3:
            continue
        try:
            verts[parts[0]].append((float(parts[1]), float(parts[2])))
        except ValueError:
            continue
    return dict(verts)


def parse_inp_polygons(inp_text: str) -> dict[str, list[tuple[float, float]]]:
    """``[Polygons]`` (or ``[POLYGONS]``) -> {subcatchment_name: [(x, y), ...]}.

    SWMM rings the polygon implicitly (last vertex may or may not
    repeat the first). The renderer closes the ring itself, so we
    just collect rows in order.
    """
    polys: dict[str, list[tuple[float, float]]] = defaultdict(list)
    for parts in _read_section_lines(inp_text, "POLYGONS"):
        if len(parts) < 3:
            continue
        try:
            polys[parts[0]].append((float(parts[1]), float(parts[2])))
        except ValueError:
            continue
    return dict(polys)


def parse_inp_subcatchments(inp_text: str) -> dict[str, str]:
    """``[SUBCATCHMENTS]`` -> {subcatchment_name: outlet_node_name}.

    Row layout is ``Name Raingage Outlet Area %Imperv Width %Slope ...``,
    so the outlet is column index 2.
    """
    out: dict[str, str] = {}
    for parts in _read_section_lines(inp_text, "SUBCATCHMENTS"):
        if len(parts) < 3:
            continue
        out[parts[0]] = parts[2]
    return out


def parse_inp_conduits(inp_text: str) -> list[tuple[str, str, str]]:
    """``[CONDUITS]`` -> [(name, from_node, to_node), ...]."""
    out: list[tuple[str, str, str]] = []
    for parts in _read_section_lines(inp_text, "CONDUITS"):
        if len(parts) < 3:
            continue
        out.append((parts[0], parts[1], parts[2]))
    return out


def parse_inp_node_kinds(inp_text: str) -> dict[str, str]:
    """Tag every node in ``[JUNCTIONS]`` / ``[OUTFALLS]`` / ``[STORAGE]``.

    Returns a {node_name: kind} mapping where ``kind`` is one of
    ``junction``, ``outfall``, ``storage``. Nodes only referenced
    from ``[COORDINATES]`` (no kind row) default to ``junction``
    in the renderer.
    """
    kinds: dict[str, str] = {}
    for parts in _read_section_lines(inp_text, "JUNCTIONS"):
        if parts:
            kinds[parts[0]] = "junction"
    for parts in _read_section_lines(inp_text, "OUTFALLS"):
        if parts:
            kinds[parts[0]] = "outfall"
    for parts in _read_section_lines(inp_text, "STORAGE"):
        if parts:
            kinds[parts[0]] = "storage"
    return kinds


# --------------------------------------------------------------------- #
# Outfall partitioning — turn (conduits, node_kinds) into a colouring.
#
# We BFS each outfall in reverse through the conduit graph: any node
# that can reach this outfall by following ``from_node -> to_node``
# downstream belongs to that outfall's sub-network. Conduits between
# two such nodes inherit the colour. Cycles (rare but legal in SWMM
# with pumps/orifices) are broken by the visited set.
# --------------------------------------------------------------------- #


def assign_outfall_colours(
    conduits: list[tuple[str, str, str]],
    node_kinds: dict[str, str],
) -> tuple[dict[str, str], dict[str, str]]:
    """Return (node_outfall, conduit_outfall) keyed by name.

    ``node_outfall[n]`` = the outfall id ``n`` ultimately drains to
    (None when ``n`` is disconnected). ``conduit_outfall[c]`` =
    same, applied to the conduit's downstream node.
    """
    # Build reverse adjacency: who drains INTO this node?
    incoming: dict[str, list[str]] = defaultdict(list)
    for _, src, dst in conduits:
        incoming[dst].append(src)

    outfalls = sorted(n for n, k in node_kinds.items() if k == "outfall")
    node_outfall: dict[str, str] = {}
    # BFS upstream from each outfall.
    for o in outfalls:
        queue = deque([o])
        node_outfall.setdefault(o, o)
        while queue:
            cur = queue.popleft()
            for upstream in incoming.get(cur, []):
                if upstream in node_outfall:
                    continue  # first claim wins -> deterministic
                node_outfall[upstream] = o
                queue.append(upstream)

    conduit_outfall: dict[str, str] = {}
    for name, _src, dst in conduits:
        if dst in node_outfall:
            conduit_outfall[name] = node_outfall[dst]
    return node_outfall, conduit_outfall


def palette_for(outfalls: list[str]) -> dict[str, tuple[float, float, float]]:
    """Deterministic colour wheel for outfall sub-networks.

    Uses the Wong colour-blind-safe categorical cycle from ``plot_style``
    (black excluded: it marks the outfalls themselves), cycling when
    there are more outfalls than colours. The sorted outfall list drives
    the order so the same network always yields the same picture.
    """
    from matplotlib.colors import to_rgb

    out: dict[str, tuple[float, float, float]] = {}
    for i, o in enumerate(sorted(outfalls)):
        out[o] = to_rgb(CATEGORY_CYCLE[i % len(CATEGORY_CYCLE)])
    return out


# --------------------------------------------------------------------- #
# Geoparquet path (preferred when SWMManywhere chain ran).
#
# Lazy import: importing this script must not require geopandas.
# When geopandas is missing or the files are missing, the caller
# falls back to the INP path.
# --------------------------------------------------------------------- #


def try_load_geoparquet(synth_dir: Path) -> dict[str, Any] | None:
    """Return a dict of GeoDataFrames or ``None`` if unavailable.

    Returns ``None`` on any of: missing files, missing geopandas,
    missing pyarrow, or any read error. Callers fall back to the INP
    parsing path on ``None``.
    """
    needed = {
        "nodes": synth_dir / "nodes.geoparquet",
        "edges": synth_dir / "edges.geoparquet",
        "subcatchments": synth_dir / "subcatchments.geoparquet",
    }
    if not all(p.exists() for p in needed.values()):
        return None
    try:
        import geopandas as gpd  # type: ignore  # noqa: F401
    except ImportError:
        return None
    try:
        return {key: gpd.read_parquet(path) for key, path in needed.items()}
    except Exception:
        return None


# --------------------------------------------------------------------- #
# Rendering.
# --------------------------------------------------------------------- #


def _finish_map(fig, ax, out_png: Path, dpi: int) -> None:
    """Common tail: equal aspect, plain coordinates, legend above, PDF + PNG."""
    from matplotlib.ticker import MaxNLocator, ScalarFormatter

    ax.set_aspect("equal", adjustable="datalim")
    # Coordinates read as coordinates: no "x 10^6" offset pulled out of
    # the tick labels, and few enough ticks that 6-7 digit labels do not
    # touch at 89 mm (same choices as plot_study_area.py).
    for axis in (ax.xaxis, ax.yaxis):
        fmt = ScalarFormatter(useOffset=False)
        fmt.set_scientific(False)
        axis.set_major_formatter(fmt)
        axis.set_major_locator(MaxNLocator(nbins=5))
    handles, labels = ax.get_legend_handles_labels()
    if labels:
        legend_outside(fig, handles, labels)
    save_figure(fig, out_png, dpi=dpi)
    plt.close(fig)


_JUNCTION_GREY = "#555555"


def render_from_inp(
    *,
    inp_path: Path,
    out_png: Path,
    dpi: int,
    draw_subcatchments: bool,
    draw_vertices: bool,
) -> None:
    """Render the layout from a SWMM INP file alone.

    This is the dependency-light path — pure matplotlib, no
    geopandas. Always works as long as the INP carries at least
    ``[COORDINATES]``.
    """
    inp_text = inp_path.read_text(encoding="utf-8", errors="ignore")

    coords = parse_inp_coordinates(inp_text)
    if not coords:
        raise SystemExit(
            f"--inp has no [COORDINATES] section: {inp_path}\n"
            "Cannot render a layout from an INP without node coordinates."
        )
    polygons = parse_inp_polygons(inp_text) if draw_subcatchments else {}
    vertices = parse_inp_vertices(inp_text) if draw_vertices else {}
    sub_outlet = parse_inp_subcatchments(inp_text)
    conduits = parse_inp_conduits(inp_text)
    node_kinds = parse_inp_node_kinds(inp_text)

    node_outfall, conduit_outfall = assign_outfall_colours(conduits, node_kinds)
    outfalls = sorted(n for n, k in node_kinds.items() if k == "outfall")
    colours = palette_for(outfalls) if outfalls else {}

    apply_style()
    xs_all = [x for x, _ in coords.values()]
    ys_all = [y for _, y in coords.values()]
    fig, ax = plt.subplots(
        figsize=map_figsize((min(xs_all), min(ys_all), max(xs_all), max(ys_all))),
        layout="constrained",
    )

    # Layer 1: subcatchment polygons (light blue, optional). Tint by
    # the outlet's outfall colour when available so sub-networks are
    # visible even when the polygons hide most conduits.
    if draw_subcatchments and polygons:
        for sub_name, ring in polygons.items():
            if len(ring) < 3:
                continue
            outlet = sub_outlet.get(sub_name)
            tint = colours.get(node_outfall.get(outlet, ""), (0.6, 0.8, 0.95))
            patch = mpatches.Polygon(
                ring,
                closed=True,
                facecolor=(*tint, 0.18),
                edgecolor=(*tint, 0.6),
                linewidth=0.3,
                zorder=1,
            )
            ax.add_patch(patch)

    # Layer 2: conduits. Use [VERTICES] when present for the polyline
    # geometry; otherwise draw a straight line between end-node coords.
    fallback_colour = (0.4, 0.4, 0.4)
    for name, src, dst in conduits:
        if src not in coords or dst not in coords:
            continue
        xs = [coords[src][0]]
        ys = [coords[src][1]]
        if draw_vertices and name in vertices:
            for vx, vy in vertices[name]:
                xs.append(vx)
                ys.append(vy)
        xs.append(coords[dst][0])
        ys.append(coords[dst][1])
        colour = colours.get(conduit_outfall.get(name, ""), fallback_colour)
        ax.plot(xs, ys, color=colour, linewidth=0.6, zorder=3)

    # Layer 3: junctions (small grey dot) + storage (blue square).
    jx, jy, sx, sy = [], [], [], []
    for name, (x, y) in coords.items():
        kind = node_kinds.get(name, "junction")
        if kind == "outfall":
            continue  # drawn last
        if kind == "storage":
            sx.append(x)
            sy.append(y)
        else:
            jx.append(x)
            jy.append(y)
    if jx:
        ax.scatter(jx, jy, s=4, c=_JUNCTION_GREY, marker="o", linewidths=0, zorder=4, label="Junction")
    if sx:
        ax.scatter(sx, sy, s=9, c=WONG["blue"], marker="s", linewidths=0, zorder=4, label="Storage")

    # Layer 4: outfalls — drawn last so they stack on top.
    ox, oy = [], []
    for o in outfalls:
        if o not in coords:
            continue
        ox.append(coords[o][0])
        oy.append(coords[o][1])
    if ox:
        ax.scatter(
            ox,
            oy,
            s=40,
            c=WONG["black"],
            marker="*",
            edgecolor="white",
            linewidth=0.3,
            zorder=5,
            label="Outfall",
        )

    # Axis tidy-up. The INP might be in metres, feet, or projected XY;
    # we don't know, so the axis says so instead of guessing a unit.
    ax.set_xlabel("X (INP units)")
    ax.set_ylabel("Y (INP units)")
    _finish_map(fig, ax, out_png, dpi)


def render_from_geoparquet(
    *,
    gdfs: dict[str, Any],
    out_png: Path,
    dpi: int,
    draw_subcatchments: bool,
) -> None:
    """Render the layout from SWMManywhere geoparquet trio.

    Reached only when ``geopandas`` and the three files were
    available. The schema convention follows SWMManywhere v0.2.x:
    nodes carry a ``node_type`` column ("junction"/"outfall") and
    edges have ``LINESTRING`` geometries with their downstream node
    encoded in ``v`` (NetworkX-style ``(u, v)`` columns).
    """
    nodes = gdfs["nodes"]
    edges = gdfs["edges"]
    subs = gdfs["subcatchments"]

    # Build the per-outfall colouring on the geopandas frames. We
    # mirror the INP path's logic (BFS upstream from each outfall)
    # rather than relying on SWMManywhere-specific columns so a hand-
    # built geoparquet trio still renders.
    node_kinds: dict[str, str] = {}
    if "node_type" in nodes.columns:
        for _, row in nodes.iterrows():
            node_kinds[str(row["id"])] = str(row["node_type"]).lower()
    else:
        # No type column — treat every node as a junction; outfalls
        # surface from edges' end-of-line nodes that have no outgoing
        # edges. (Best effort; the INP path is the canonical surface.)
        downstream = set(str(v) for v in edges["v"]) if "v" in edges.columns else set()
        upstream = set(str(u) for u in edges["u"]) if "u" in edges.columns else set()
        for _, row in nodes.iterrows():
            nid = str(row["id"])
            if nid in downstream and nid not in upstream:
                node_kinds[nid] = "outfall"
            else:
                node_kinds[nid] = "junction"

    conduits = [
        (str(row.get("id", i)), str(row["u"]), str(row["v"]))
        for i, row in edges.iterrows()
        if "u" in edges.columns and "v" in edges.columns
    ]
    node_outfall, conduit_outfall = assign_outfall_colours(conduits, node_kinds)
    outfalls = sorted(n for n, k in node_kinds.items() if k == "outfall")
    colours = palette_for(outfalls) if outfalls else {}

    apply_style()
    frame = subs if (draw_subcatchments and not subs.empty) else edges
    fig, ax = plt.subplots(
        figsize=map_figsize(tuple(frame.total_bounds)),
        layout="constrained",
    )

    if draw_subcatchments and not subs.empty:
        if "outlet" in subs.columns:
            tint_for = lambda r: colours.get(  # noqa: E731 — tiny local helper
                node_outfall.get(str(r["outlet"]), ""), (0.6, 0.8, 0.95)
            )
        else:
            tint_for = lambda r: (0.6, 0.8, 0.95)  # noqa: E731
        for _, row in subs.iterrows():
            geom = row.geometry
            if geom is None or geom.is_empty:
                continue
            tint = tint_for(row)
            polys = [geom] if geom.geom_type == "Polygon" else list(geom.geoms)
            for poly in polys:
                xs, ys = poly.exterior.xy
                ax.fill(xs, ys, color=(*tint, 0.18), edgecolor=(*tint, 0.6), linewidth=0.3, zorder=1)

    # Conduits via LINESTRING geometry; falls back to (u, v) end nodes
    # when geometry is missing.
    coord_lookup: dict[str, tuple[float, float]] = {}
    if "geometry" in nodes.columns:
        for _, row in nodes.iterrows():
            pt = row.geometry
            if pt is None or pt.is_empty:
                continue
            coord_lookup[str(row["id"])] = (pt.x, pt.y)

    fallback_colour = (0.4, 0.4, 0.4)
    for _, row in edges.iterrows():
        u, v = str(row.get("u", "")), str(row.get("v", ""))
        colour = colours.get(node_outfall.get(v, ""), fallback_colour)
        geom = row.geometry if "geometry" in edges.columns else None
        if geom is not None and not geom.is_empty:
            xs, ys = geom.xy
            ax.plot(xs, ys, color=colour, linewidth=0.6, zorder=3)
        elif u in coord_lookup and v in coord_lookup:
            ax.plot(
                [coord_lookup[u][0], coord_lookup[v][0]],
                [coord_lookup[u][1], coord_lookup[v][1]],
                color=colour,
                linewidth=0.6,
                zorder=3,
            )

    # Nodes
    jx, jy, ox, oy = [], [], [], []
    for nid, (x, y) in coord_lookup.items():
        if node_kinds.get(nid) == "outfall":
            ox.append(x)
            oy.append(y)
        else:
            jx.append(x)
            jy.append(y)
    if jx:
        ax.scatter(jx, jy, s=4, c=_JUNCTION_GREY, marker="o", linewidths=0, zorder=4, label="Junction")
    if ox:
        ax.scatter(
            ox,
            oy,
            s=40,
            c=WONG["black"],
            marker="*",
            edgecolor="white",
            linewidth=0.3,
            zorder=5,
            label="Outfall",
        )

    xlabel, ylabel = _crs_axis_labels(getattr(nodes, "crs", None))
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    _finish_map(fig, ax, out_png, dpi)


def _crs_axis_labels(crs: Any) -> tuple[str, str]:
    """Axis labels with units, read off the geoparquet CRS when it has one."""
    try:
        if crs is None:
            return ("X (CRS units)", "Y (CRS units)")
        if crs.is_geographic:
            return ("Longitude (°)", "Latitude (°)")
        unit = crs.axis_info[0].unit_name if crs.axis_info else ""
        unit = {"metre": "m", "meter": "m", "foot": "ft", "US survey foot": "ft"}.get(unit, unit)
        return (f"Easting ({unit or 'CRS units'})", f"Northing ({unit or 'CRS units'})")
    except Exception:  # pragma: no cover - defensive against exotic CRS objects
        return ("X (CRS units)", "Y (CRS units)")


# --------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------- #


def _build_argparser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description=(
            "Render a SWMM model's spatial layout (subcatchments, conduits, "
            "nodes, outfalls) as a PNG."
        )
    )
    p.add_argument(
        "--inp",
        type=Path,
        default=None,
        help=(
            "Path to a SWMM .inp file. Required when no SWMManywhere "
            "geoparquet trio is found under --synth-dir."
        ),
    )
    p.add_argument(
        "--synth-dir",
        type=Path,
        default=None,
        help=(
            "Optional path to a directory containing nodes.geoparquet, "
            "edges.geoparquet, subcatchments.geoparquet (the SWMManywhere "
            "chain output under <run-dir>/10_upstream/swmmanywhere/, or the "
            "legacy <run-dir>/10_swmmanywhere/ for older runs)."
        ),
    )
    p.add_argument(
        "--out-png",
        type=Path,
        required=True,
        help="Preview PNG path; the vector twin <stem>.pdf is written beside it.",
    )
    p.add_argument(
        "--dpi",
        type=int,
        default=450,
        help="Resolution of the preview PNG (the PDF is vector). Spec minimum for images is 450.",
    )
    p.add_argument(
        "--no-subcatchments",
        action="store_true",
        help="Skip the polygon (subcatchment) layer; render conduits + nodes only.",
    )
    p.add_argument(
        "--no-vertices",
        action="store_true",
        help=(
            "Draw conduits as straight lines between their end nodes, "
            "ignoring [VERTICES]. Useful when [VERTICES] is noisy or absent."
        ),
    )
    return p


def main(argv: list[str] | None = None) -> int:
    args = _build_argparser().parse_args(argv)

    draw_subcatchments = not args.no_subcatchments
    draw_vertices = not args.no_vertices

    # Try geoparquet first — only when the user pointed us at a synth_dir.
    if args.synth_dir is not None and args.synth_dir.exists():
        gdfs = try_load_geoparquet(args.synth_dir)
        if gdfs is not None:
            render_from_geoparquet(
                gdfs=gdfs,
                out_png=args.out_png,
                dpi=args.dpi,
                draw_subcatchments=draw_subcatchments,
            )
            print(f"wrote {args.out_png}")
            return 0

    # Fallback: INP parsing.
    if args.inp is None:
        print(
            "error: must provide --inp (or a --synth-dir containing the "
            "SWMManywhere geoparquet trio).",
            file=sys.stderr,
        )
        return 2
    if not args.inp.exists():
        print(f"error: --inp not found: {args.inp}", file=sys.stderr)
        return 1

    render_from_inp(
        inp_path=args.inp,
        out_png=args.out_png,
        dpi=args.dpi,
        draw_subcatchments=draw_subcatchments,
        draw_vertices=draw_vertices,
    )
    print(f"wrote {args.out_png}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

#!/usr/bin/env python3
"""Plot rainfall (inverted) vs runoff/outfall hydrograph to the Nature figure spec.

Figure rules (the swmm-plot skill's standing standard, see SKILL.md):
- SI units, every axis labelled with units in parentheses
- Inverted rainfall axis (hyetograph convention), flow rising from the bottom
- Hydrograph shape preserved (assumes output time step is sufficiently fine, e.g., 5-min)
- Style comes from ``plot_style.apply_style()`` (vendored nature.mplstyle):
  89 mm single / 183 mm double column, 7 pt sans-serif, 0.25-1 pt lines,
  ticks outward, no gridlines, Wong colour-blind-safe palette
- No title (titles belong to the surrounding document)

Inputs:
- INP path (to read TIMESERIES for rainfall)
- OUT path (to read node total inflow from SWMM output)

Output:
- ``--out-png`` (450 dpi preview) plus its vector twin ``<stem>.pdf`` (the
  submission file), written by ``plot_style.save_figure``

Agent-flow invariant (issue #125):
    The ``--rain-ts`` / ``--node`` / ``--node-attr`` defaults below are
    self-documenting placeholders (``<rainfall-series-name>`` /
    ``<outfall-or-junction>``) that are only reachable from a *manual
    CLI invocation*. The agent-driven path always passes explicit
    values resolved against the run's actual INP/OUT:

        agent goal
          -> planner._extract_plot_choice (agentic_swmm/agent/planner.py)
             which reads inspect_plot_options output and picks real names
          -> tool_registry._plot_run_args (agentic_swmm/agent/tool_registry.py)
             which forwards the explicit values to the MCP server
          -> mcp/swmm-plot/server.js
             whose Zod defaults are also placeholders and are likewise
             never reached on an agent call.

    If a manual CLI invocation hits one of these placeholder defaults,
    the script errors informatively (``rainfall series
    '<rainfall-series-name>' not found in INP``) instead of failing
    with a Tecnopolo-shaped error that misleads users on a different
    watershed. ``Total_inflow`` is a SWMM-universal attribute name (not
    watershed-specific) so its default remains literal.

    Regression test: ``tests/test_plot_run_args_overrides_defaults.py``.
"""

from __future__ import annotations

import argparse
import sys
from datetime import datetime, timedelta
from pathlib import Path


def _warn_if_cold_start() -> None:
    """Emit a one-line stderr hint if matplotlib's font cache is missing.

    The MCP server preheats matplotlib + swmmtoolbox at boot (see
    issue #109) so the user normally never sees this. If the preheat
    failed (no Python, no deps) or hasn't finished yet, this warning
    tells the user why the first plot call is taking a while instead
    of leaving them staring at a silent ``you>`` prompt.
    """
    try:
        import matplotlib  # cheap; just reads metadata
        cachedir = Path(matplotlib.get_cachedir())
    except Exception:
        return
    # matplotlib names the cache ``fontlist-vNNN.json``; if any file
    # matching that glob exists we treat the cache as warm.
    if not any(cachedir.glob('fontlist-v*.json')):
        sys.stderr.write(
            '[swmm-plot] First plot warms up matplotlib + swmmtoolbox '
            '(~60-90s). Subsequent plots are fast.\n'
        )
        sys.stderr.flush()


_warn_if_cold_start()

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np

from swmmtoolbox import extract

# The shared style module lives next to this script. Tests load the script
# through ``importlib.util.spec_from_file_location`` (no sys.path[0] entry
# for this directory), so resolve the sibling explicitly.
_HERE = str(Path(__file__).resolve().parent)
if _HERE not in sys.path:
    sys.path.insert(0, _HERE)

from plot_style import WONG, apply_style, figsize, legend_outside, save_figure  # noqa: E402


def parse_timeseries_file(path: Path) -> tuple[list[datetime], list[float]]:
    times: list[datetime] = []
    vals: list[float] = []
    for raw in path.read_text(errors='ignore').splitlines():
        s = raw.strip()
        if not s or s.startswith(';'):
            continue
        parts = s.split()
        if len(parts) < 3:
            continue
        dt = _parse_inline_ts_datetime(parts[0], parts[1])
        times.append(dt)
        vals.append(float(parts[2]))
    if not times:
        raise SystemExit(f'No timeseries values found in {path}')
    return times, vals


def _find_raingages_file(inp_path: Path, gage_id: str) -> Path | None:
    """Return the .dat path referenced by ``[RAINGAGES] FILE`` for ``gage_id``.

    SWMManywhere emits INPs whose rainfall input lives in an external
    .dat file referenced from ``[RAINGAGES]`` instead of an inline
    ``[TIMESERIES]`` block, e.g.::

        [RAINGAGES]
        rg1   INTENSITY  0:05   1.0   FILE   "storm.dat"

    Returns ``None`` (not raise) so the caller can produce a context-rich
    error pointing at both INP and the missing series.
    """
    in_raingages = False
    for raw in inp_path.read_text(errors='ignore').splitlines():
        s = raw.strip()
        if s.upper() == '[RAINGAGES]':
            in_raingages = True
            continue
        if in_raingages:
            if s.startswith('[') and s.endswith(']'):
                break
            if not s or s.startswith(';'):
                continue
            parts = s.split()
            upper_parts = [p.upper() for p in parts]
            if parts[0].strip('"') != gage_id:
                continue
            if 'FILE' in upper_parts:
                idx = upper_parts.index('FILE')
                if idx + 1 < len(parts):
                    return inp_path.parent / parts[idx + 1].strip('"')
    return None


def parse_raingages_file(path: Path, gage_id: str | None = None) -> tuple[list[datetime], list[float]]:
    """Parse the SWMM5 ``[RAINGAGES] FILE`` format::

        <gage_id> <YYYY> <MM> <DD> <HH> <mm> <value>

    Lines whose first token does not match ``gage_id`` (when supplied)
    are skipped — SWMM5 allows a single .dat to hold multiple gages.
    """
    times: list[datetime] = []
    vals: list[float] = []
    for raw in path.read_text(errors='ignore').splitlines():
        s = raw.strip()
        if not s or s.startswith(';'):
            continue
        parts = s.split()
        if len(parts) < 7:
            continue
        if gage_id is not None and parts[0] != gage_id:
            continue
        try:
            dt = datetime(
                int(parts[1]), int(parts[2]), int(parts[3]),
                int(parts[4]), int(parts[5]),
            )
            v = float(parts[6])
        except (ValueError, IndexError):
            continue
        times.append(dt)
        vals.append(v)
    if not times:
        raise SystemExit(
            f'No RAINGAGES FILE rows found in {path}'
            + (f' for gage {gage_id!r}' if gage_id else '')
        )
    return times, vals


_INLINE_TS_DATETIME_FORMATS = ('%m/%d/%Y %H:%M', '%m/%d/%Y %H:%M:%S')


def _parse_inline_ts_datetime(date_str: str, time_str: str) -> datetime:
    """Parse a [TIMESERIES] date+time pair, tolerating seconds.

    SWMM writes both HH:MM and HH:MM:SS in inline timeseries rows; the
    SWMMCanada upstream emits second-precision stamps, and the previous
    single hardcoded format crashed on every fetched Canadian model
    (found live 2026-08-09: "unconverted data remains: :00").
    """
    raw = f"{date_str} {time_str}"
    for fmt in _INLINE_TS_DATETIME_FORMATS:
        try:
            return datetime.strptime(raw, fmt)
        except ValueError:
            continue
    raise SystemExit(
        f"unsupported [TIMESERIES] datetime {raw!r}; expected one of "
        f"{', '.join(_INLINE_TS_DATETIME_FORMATS)}"
    )


def parse_timeseries_from_inp(inp_path: Path, ts_name: str) -> tuple[list[datetime], list[float]]:
    """Return (times, values) from [TIMESERIES]. Values are whatever units the INP encodes.

    Fallback (strict additive): if no ``[TIMESERIES]`` row matches
    ``ts_name``, look for a ``[RAINGAGES] <ts_name> ... FILE <storm.dat>``
    entry and parse that .dat. This supports SWMManywhere-generated INPs
    which omit ``[TIMESERIES]`` entirely.
    """
    times: list[datetime] = []
    vals: list[float] = []
    reading = False
    for line in inp_path.read_text(errors='ignore').splitlines():
        s = line.strip()
        if s.upper() == '[TIMESERIES]':
            reading = True
            continue
        if reading:
            if s.startswith('[') and s.endswith(']'):
                break
            if (not s) or s.startswith(';;'):
                continue
            parts = s.split()
            if parts[0] != ts_name:
                continue
            if len(parts) >= 3 and parts[1].upper() == 'FILE':
                return parse_timeseries_file(inp_path.parent / parts[2].strip('"'))
            dt = _parse_inline_ts_datetime(parts[1], parts[2])
            times.append(dt)
            vals.append(float(parts[3]))
    if not times:
        # Strict additive fallback: try RAINGAGES FILE (SWMManywhere case).
        raingages_path = _find_raingages_file(inp_path, ts_name)
        if raingages_path is not None:
            if not raingages_path.exists():
                raise SystemExit(
                    f'RAINGAGES FILE referenced by gage {ts_name!r} not found: '
                    f'{raingages_path} (referenced from {inp_path})'
                )
            return parse_raingages_file(raingages_path, gage_id=ts_name)
        raise SystemExit(f'No TIMESERIES values found for {ts_name} in {inp_path}')
    return times, vals


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--inp', required=True, type=Path)
    ap.add_argument('--out', dest='out_file', required=True, type=Path)
    # Issue #125: ``--rain-ts`` / ``--node`` defaults are SELF-DOCUMENTING
    # PLACEHOLDERS, not portability rot. They are unreachable in the
    # agent-driven path (planner always overrides via
    # ``agentic_swmm/agent/tool_registry.py::_plot_run_args``). A manual
    # CLI invocation that hits them errors with a clear message that
    # names the placeholder string instead of misleading the user with
    # ``TS_RAIN``/``O1``. ``--node-attr`` keeps its literal default
    # because ``Total_inflow`` is a SWMM-universal attribute name.
    ap.add_argument('--rain-ts', default='<rainfall-series-name>',
                    help='Name of the SWMM [TIMESERIES] block holding rainfall. The agent flow passes the resolved value; manual CLI users must supply this.')
    ap.add_argument('--rain-kind', choices=['intensity_mm_per_hr', 'depth_mm_per_dt', 'cumulative_depth_mm'], default='depth_mm_per_dt',
                    help='How to interpret TIMESERIES values for plotting. Use depth_mm_per_dt for (mm/Δt) hyetograph (inverted).')
    ap.add_argument('--dt-min', type=float, default=5.0, help='Used only when rain-kind=depth_mm_per_dt or to convert intensity to depth.')
    # ``--node`` and ``--link`` are alternate entity selectors. The
    # argparse group below enforces mutual exclusion: a single render
    # plots either a node-level series or a link-level (Flow_rate)
    # series. ``--node`` keeps its placeholder default so existing
    # callers / tests that omit the flag still get the same fail-fast
    # message that pre-dates --link.
    entity_group = ap.add_mutually_exclusive_group()
    entity_group.add_argument('--node', default='<outfall-or-junction>',
                              help='SWMM node id (outfall or junction) whose attribute is plotted. The agent flow passes the resolved value; manual CLI users must supply this.')
    entity_group.add_argument('--link', default=None,
                              help='SWMM link/conduit id; when set, the lower panel plots Flow_rate for the link instead of a node attribute. Mutually exclusive with --node.')
    ap.add_argument('--node-attr', default='Total_inflow')
    ap.add_argument('--out-png', required=True, type=Path,
                    help='Preview PNG path; the vector twin <stem>.pdf is written beside it.')
    ap.add_argument('--dpi', type=int, default=450,
                    help='Resolution of the preview PNG (the PDF is vector). Spec minimum for images is 450.')
    ap.add_argument('--width', choices=['single', 'double'], default='single',
                    help='Figure width: single column (89 mm, default) or double column (183 mm).')
    ap.add_argument('--focus-day', type=str, default=None,
                    help='If set (YYYY-MM-DD), base day for x-axis formatting.')
    ap.add_argument('--window-start', type=str, default=None,
                    help='Optional HH:MM. If provided with --focus-day, x-axis will be limited to this time window within the day.')
    ap.add_argument('--window-end', type=str, default=None,
                    help='Optional HH:MM. If provided with --focus-day, x-axis will be limited to this time window within the day.')
    ap.add_argument('--pad-hours', type=float, default=2.0,
                    help='When focus-day is not set, auto-window uses nonzero rainfall extent ± pad-hours.')
    ap.add_argument('--rain-ymax-factor', type=float, default=3.0,
                    help='Multiplier applied to the plotted rainfall maximum so inverted bars stay in the upper part of the panel.')
    ap.add_argument('--flow-ymax-factor', type=float, default=2.5,
                    help='Multiplier applied to the plotted flow maximum so the hydrograph does not visually collide with rainfall bars.')
    args = ap.parse_args()

    # Issue #125: catch manual CLI users who relied on the old
    # ``TS_RAIN``/``O1`` defaults. The placeholder strings cannot resolve
    # against any real INP, so we fail fast with a message that names
    # the missing flag instead of letting ``parse_timeseries_from_inp``
    # surface a confusing "No TIMESERIES values found for
    # '<rainfall-series-name>'" error deeper in the stack.
    if args.rain_ts == '<rainfall-series-name>':
        raise SystemExit(
            "--rain-ts is a placeholder ('<rainfall-series-name>'); pass an "
            "actual TIMESERIES name from your INP, e.g. --rain-ts MyRainSeries. "
            "(The agent-driven path resolves this automatically via "
            "inspect_plot_options; this error only appears in manual CLI use.)"
        )
    # The ``--node`` placeholder check only fires when ``--link`` was
    # not supplied. ``--link`` is the alternate selector and provides
    # its own (non-placeholder) id, so the user must not be told to
    # pass --node in that case.
    if not args.link and args.node == '<outfall-or-junction>':
        raise SystemExit(
            "--node is a placeholder ('<outfall-or-junction>'); pass an "
            "actual SWMM node id, e.g. --node OUT_0, or use --link for "
            "a conduit-level hydrograph. "
            "(The agent-driven path resolves this automatically via "
            "inspect_plot_options; this error only appears in manual CLI use.)"
        )

    # Bug #236: windowStart/windowEnd are only honoured together with
    # focusDay (they are HH:MM offsets within that day). Supplying them
    # without focusDay was previously a silent no-op. Raise a clear error
    # so the user knows exactly what is missing.
    if (args.window_start or args.window_end) and not args.focus_day:
        raise SystemExit(
            "--window-start/--window-end require --focus-day (YYYY-MM-DD) "
            "and use HH:MM format. They cannot be used without --focus-day."
        )

    # Matplotlib styling: the vendored Nature stylesheet (7 pt sans-serif,
    # ticks out, no grid, Wong palette, TrueType-embedded text).
    apply_style()

    rain_t, rain_v = parse_timeseries_from_inp(args.inp, args.rain_ts)
    rain_v = np.asarray(rain_v, dtype=float)

    # Render every gage format as depth (mm) per recording interval so bar
    # heights are comparable. Only INTENSITY (mm/hr) needs a unit conversion;
    # VOLUME is already depth-per-interval and CUMULATIVE is a running total.
    if args.rain_kind == 'intensity_mm_per_hr':
        # mm/hr -> depth over the interval. This is the ONLY branch that may
        # multiply by dt/60 (review P1-9).
        rain_plot = rain_v * (args.dt_min / 60.0)
        rain_ylabel = f'Rainfall depth (mm/{int(args.dt_min)} min)'
    elif args.rain_kind == 'cumulative_depth_mm':
        rain_plot = np.diff(rain_v, prepend=rain_v[0])
        rain_plot = np.where(rain_plot < 0, 0.0, rain_plot)
        rain_ylabel = f'Rainfall depth (mm/{int(args.dt_min)} min)'
    else:
        # depth_mm_per_dt: values are already depth (mm) per interval (e.g. a
        # VOLUME-format gage). Plot as-is; a second dt/60 multiply here would
        # double-convert and shrink the hyetograph (review P1-9).
        rain_plot = rain_v
        rain_ylabel = f'Rainfall depth (mm/{int(args.dt_min)} min)'

    # Flow series (SI): CMS = m^3/s.
    # Node-level path reads ``node,<id>,<attr>`` (Total_inflow by
    # default); link-level path reads ``link,<id>,Flow_rate`` so the
    # lower panel renders a conduit hydrograph. Both share the same
    # paired-axis layout (rain top, flow bottom).
    if args.link:
        key = f'link,{args.link},Flow_rate'
        flow_label = 'Flow'
        flow_ylabel = 'Flow (m³/s)'
    else:
        key = f'node,{args.node},{args.node_attr}'
        flow_label = 'Flow'
        flow_ylabel = 'Flow (m³/s)'
    try:
        flow_df = extract(str(args.out_file), key)
    except Exception:
        # swmmtoolbox surfaces a missing id as a raw pandas
        # "No objects to concatenate" ValueError. Turn that into an
        # actionable message listing what the .out actually contains.
        kind = 'link' if args.link else 'node'
        ident = args.link or args.node
        try:
            from swmmtoolbox import catalog

            available = sorted({str(row[1]) for row in catalog(str(args.out_file), kind) if len(row) > 1})
        except Exception:
            available = []
        if not available:
            # Not a wrong id: the .out carries NO per-element series at
            # all. SWMM only stores node/link time series for elements
            # named in the INP's [REPORT] section; a model without one
            # (e.g. some upstream-generated INPs) yields a system-only
            # .out no id can ever be found in (found live 2026-08-09).
            print(
                f"error: {args.out_file} contains no per-{kind} time series "
                "at all, so no id can be plotted from it. The model's "
                "[REPORT] section likely omits NODES/LINKS lines. Add "
                "'[REPORT]' with 'NODES ALL' and 'LINKS ALL' to the INP, "
                "re-run the simulation, then plot again.",
                file=sys.stderr,
            )
            raise SystemExit(2)
        listing = ', '.join(available[:20])
        print(
            f"error: {kind} '{ident}' not found in {args.out_file}. "
            f"Available {kind} ids: {listing}. "
            "Pick one with --node/--link, or run `aiswmm plot` without "
            "--node to use the model's first outfall.",
            file=sys.stderr,
        )
        raise SystemExit(2)
    flow_t = flow_df.index.to_pydatetime()
    flow_v = flow_df.iloc[:, 0].to_numpy(dtype=float)

    # Keep only the rainfall inside the reported period (plus the pad).
    # Continuous-simulation INPs often point at multi-year climate files
    # (a 30-year 5-min file is 3M rows): drawing every bar stalls
    # matplotlib for minutes and would bloat the vector PDF, and rain
    # outside the run says nothing about the run.
    if flow_t.size and rain_t:
        pad = timedelta(hours=float(args.pad_hours))
        rt = np.asarray(rain_t, dtype='datetime64[s]')
        keep = (rt >= np.datetime64(flow_t[0] - pad, 's')) & (rt <= np.datetime64(flow_t[-1] + pad, 's'))
        rain_t = [t for t, k in zip(rain_t, keep) if k]
        rain_plot = rain_plot[keep]

    # Figure: spec-legal column width; a hydrograph wants a wide panel, so
    # the double-column variant is flatter than the golden-ratio single.
    fig, ax_rain = plt.subplots(
        figsize=figsize(args.width, aspect=0.618 if args.width == 'single' else 0.40),
        layout='constrained',
    )

    # Rain bars: sky blue from the Wong palette, solid (no alpha, no hatching).
    bar_width_days = (args.dt_min / 60.0) / 24.0
    ax_rain.bar(
        rain_t,
        rain_plot,
        width=bar_width_days,
        color=WONG['sky_blue'],
        edgecolor='none',
        linewidth=0,
        label='Rain',
        zorder=1,
    )
    ax_rain.set_ylabel(rain_ylabel)
    ax_rain.set_xlabel('Time of day (hh:mm)' if args.focus_day else 'Time')

    # invert rain axis (hyetograph convention)
    ax_rain.invert_yaxis()

    # Flow line (draw above rain): black, the palette's first colour, for
    # the primary series. The twin axis carries its own right-hand spine
    # (the stylesheet hides top/right spines by default, and its left one
    # would double-draw the rain axis).
    ax_flow = ax_rain.twinx()
    ax_flow.plot(flow_t, flow_v, color=WONG['black'], linewidth=0.75, label=flow_label, zorder=3)
    ax_flow.set_ylabel(flow_ylabel)
    ax_flow.spines['right'].set_visible(True)
    ax_flow.spines['left'].set_visible(False)

    rain_max = float(np.nanmax(rain_plot)) if rain_plot.size else 0.0
    if rain_max > 0:
        ax_rain.set_ylim(rain_max * max(args.rain_ymax_factor, 1.0), 0.0)
    flow_max = float(np.nanmax(flow_v)) if flow_v.size else 0.0
    if flow_max > 0:
        ax_flow.set_ylim(0.0, flow_max * max(args.flow_ymax_factor, 1.0))

    # Focus x-axis: one day or auto-window
    import matplotlib.dates as mdates
    if args.focus_day:
        d0 = datetime.strptime(args.focus_day, '%Y-%m-%d')
        if args.window_start and args.window_end:
            ws = datetime.strptime(args.window_start, '%H:%M').time()
            we = datetime.strptime(args.window_end, '%H:%M').time()
            t0 = d0.replace(hour=ws.hour, minute=ws.minute)
            t1 = d0.replace(hour=we.hour, minute=we.minute)
            ax_rain.set_xlim(t0, t1)
            ax_rain.xaxis.set_major_locator(mdates.HourLocator(interval=1))
            ax_rain.xaxis.set_major_formatter(mdates.DateFormatter('%H:%M'))
        else:
            ax_rain.set_xlim(d0, d0 + timedelta(hours=24))
            ax_rain.xaxis.set_major_locator(mdates.HourLocator(interval=3))
            ax_rain.xaxis.set_major_formatter(mdates.DateFormatter('%H:%M'))
    else:
        nz = np.where(np.asarray(rain_plot) > 0)[0]
        if nz.size:
            tmin = rain_t[int(nz.min())]
            tmax = rain_t[int(nz.max())]
            pad = timedelta(hours=float(args.pad_hours))
            ax_rain.set_xlim(tmin - pad, tmax + pad)
        # Auto-pick a readable tick density regardless of duration. The
        # legacy ``HourLocator(interval=2)`` here exploded into hundreds
        # of overlapping labels on multi-week / multi-month runs (#112
        # black-blur). ``ConciseDateFormatter`` picks the smallest
        # readable format (``HH:MM`` for sub-day, ``MM-DD`` for sub-year,
        # ``YYYY`` otherwise) and stores the calendar context in the
        # offset string above the axis.
        locator = mdates.AutoDateLocator(maxticks=12)
        ax_rain.xaxis.set_major_locator(locator)
        ax_rain.xaxis.set_major_formatter(mdates.ConciseDateFormatter(locator))

    # Ticks: outward on every axis, from the stylesheet. No title (per spec).

    # Legend: one combined key above the panel so it never sits on the
    # bars or the hydrograph (text over data is a spec violation).
    h1, l1 = ax_rain.get_legend_handles_labels()
    h2, l2 = ax_flow.get_legend_handles_labels()
    legend_outside(fig, h1 + h2, l1 + l2)

    # Vector PDF (submission) + PNG preview, default bbox (never 'tight').
    save_figure(fig, args.out_png, dpi=args.dpi)


if __name__ == '__main__':
    main()

#!/usr/bin/env python3
"""Render a study-area map from a SWMMCanada upstream bundle.

The classic report figure: WHERE the modeled area is and what it holds.
Composition (all layers ship inside the run's own upstream bundle, so
the figure is fully self-contained, no web tiles, no network access):

* DEM hillshade background (``dem_dtm.tif``),
* subcatchment polygons (``preview/network.geojson`` kind=subcatchment),
* conduit lines (kind=conduit),
* outfall markers (kind=outfall),
* scale bar, north arrow, and an annotation cartouche with the AOI's
  WGS84 extent, element counts, and CRS.

Inputs
------
``--run-dir`` locates ``10_upstream/swmmcanada/swmm_model.zip`` per the
canonical layout; ``--bundle`` points at a zip explicitly. Runs that
were not fetched from SWMMCanada have no bundle: the script exits 2
with a plain explanation instead of guessing.

Dependencies: geopandas + rasterio (the aiswmm ``gis`` extra). Absent
dependencies exit 2 with the install command.

Style follows the skill's standing figure standard (``plot_style``: the
vendored Nature spec, 89 mm single column, 7 pt sans-serif, ticks out,
no grid, Wong palette, no title); the cartouche carries the identifying
text instead. Output is the ``--out-png`` preview plus its vector twin
``<stem>.pdf``.
"""

from __future__ import annotations

import argparse
import json
import sys
import tempfile
import zipfile
from pathlib import Path

# Shared style module next to this script (see plot_network_layout.py).
_HERE = str(Path(__file__).resolve().parent)
if _HERE not in sys.path:
    sys.path.insert(0, _HERE)

from plot_style import WONG, apply_style, legend_outside, map_figsize, save_figure  # noqa: E402


_NETWORK_MEMBER = "preview/network.geojson"
_DEM_MEMBER = "dem_dtm.tif"
_BUNDLE_RELPATH = Path("10_upstream/swmmcanada/swmm_model.zip")


def _fail(message: str) -> "SystemExit":
    print(f"error: {message}", file=sys.stderr)
    return SystemExit(2)


def _load_geo_stack():
    try:
        import geopandas as gpd
        import rasterio
        from rasterio.warp import (
            Resampling,
            calculate_default_transform,
            reproject,
        )
    except ImportError as exc:
        raise _fail(
            f"study-area map needs the gis extra ({exc.name} missing); "
            "install with: pip install 'aiswmm[gis]'"
        )
    return gpd, rasterio, Resampling, calculate_default_transform, reproject


def _resolve_bundle(args: argparse.Namespace) -> Path:
    if args.bundle:
        bundle = Path(args.bundle).expanduser()
        if not bundle.is_file():
            raise _fail(f"bundle not found: {bundle}")
        return bundle
    run_dir = Path(args.run_dir).expanduser()
    bundle = run_dir / _BUNDLE_RELPATH
    if not bundle.is_file():
        raise _fail(
            f"no SWMMCanada bundle at {bundle}. The study-area map is "
            "built from the upstream bundle's GIS layers, so it is "
            "available for runs fetched via the Canada service; for "
            "other runs pass --bundle pointing at a swmm_model.zip."
        )
    return bundle


def _extract_members(bundle: Path, dest: Path) -> tuple[Path, Path | None]:
    with zipfile.ZipFile(bundle) as zf:
        names = set(zf.namelist())
        if _NETWORK_MEMBER not in names:
            raise _fail(
                f"{bundle.name} has no {_NETWORK_MEMBER}; cannot draw the "
                "study area from this bundle."
            )
        zf.extract(_NETWORK_MEMBER, dest)
        dem_path: Path | None = None
        if _DEM_MEMBER in names:
            zf.extract(_DEM_MEMBER, dest)
            dem_path = dest / _DEM_MEMBER
    return dest / _NETWORK_MEMBER, dem_path


def _hillshade(dem_path: Path, dst_crs: str, rasterio, Resampling, calc, reproject):
    """Return (shade_array, extent) reprojected to ``dst_crs``."""
    import numpy as np
    from matplotlib.colors import LightSource

    with rasterio.open(dem_path) as src:
        transform, w, h = calc(src.crs, dst_crs, src.width, src.height, *src.bounds)
        dem = np.empty((h, w), dtype=np.float32)
        reproject(
            rasterio.band(src, 1),
            dem,
            dst_transform=transform,
            dst_crs=dst_crs,
            resampling=Resampling.bilinear,
        )
        left, top = transform.c, transform.f
        right = left + transform.a * w
        bottom = top + transform.e * h
    dem = np.where(dem < -1000, np.nan, dem)
    fill = float(np.nanmean(dem)) if np.isfinite(np.nanmean(dem)) else 0.0
    shade = LightSource(azdeg=315, altdeg=45).hillshade(
        np.nan_to_num(dem, nan=fill), vert_exag=2
    )
    return shade, (left, right, bottom, top)


def _scale_bar_length(width_m: float) -> int:
    """A round bar length near a fifth of the map width."""
    target = width_m / 5
    for candidate in (5000, 2000, 1000, 500, 200, 100, 50, 20, 10):
        if candidate <= target:
            return candidate
    return 10


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Render a study-area map from a SWMMCanada bundle."
    )
    parser.add_argument("--run-dir", type=Path, help="Run directory (canonical layout).")
    parser.add_argument("--bundle", type=Path, help="Explicit swmm_model.zip path.")
    parser.add_argument(
        "--out-png",
        type=Path,
        default=None,
        help="Preview PNG path; the vector twin <stem>.pdf is written beside it.",
    )
    parser.add_argument(
        "--dpi",
        type=int,
        default=450,
        help="Resolution of the preview PNG (the PDF is vector). Spec minimum for images is 450.",
    )
    parser.add_argument(
        "--place",
        default=None,
        help="Optional place label for the cartouche (e.g. 'James Bay, Victoria BC').",
    )
    args = parser.parse_args(argv)
    if not args.run_dir and not args.bundle:
        parser.error("pass --run-dir or --bundle")

    gpd, rasterio, Resampling, calc, reproject = _load_geo_stack()

    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from matplotlib.ticker import MaxNLocator, ScalarFormatter

    bundle = _resolve_bundle(args)
    with tempfile.TemporaryDirectory() as tmp:
        geojson_path, dem_path = _extract_members(bundle, Path(tmp))
        net = gpd.read_file(geojson_path)
        # Work in the local UTM zone so meters are meters.
        utm = net.estimate_utm_crs()
        net = net.to_crs(utm)
        subs = net[net["kind"] == "subcatchment"]
        links = net[net["kind"] == "conduit"]
        outfalls = net[net["kind"] == "outfall"]
        if subs.empty and links.empty:
            raise _fail("bundle network.geojson holds no subcatchments or conduits")

        shade = extent = None
        if dem_path is not None:
            try:
                shade, extent = _hillshade(
                    dem_path, str(utm), rasterio, Resampling, calc, reproject
                )
            except Exception as exc:  # DEM is decoration; never fatal.
                print(f"note: DEM hillshade skipped ({exc})", file=sys.stderr)

        frame = subs if not subs.empty else links
        minx, miny, maxx, maxy = frame.total_bounds
        pad = max((maxx - minx), (maxy - miny)) * 0.04

        apply_style()
        fig, ax = plt.subplots(
            figsize=map_figsize((minx - pad, miny - pad, maxx + pad, maxy + pad)),
            layout="constrained",
        )
        if shade is not None:
            ax.imshow(shade, cmap="gray", extent=extent, alpha=0.55, zorder=1)
        # Wong palette throughout: sky-blue subcatchments over the grey
        # hillshade, black conduits, black outfall stars (same marker as
        # the network map so the two figures of a run read as one set).
        if not subs.empty:
            subs.plot(
                ax=ax,
                facecolor=WONG["sky_blue"],
                edgecolor=WONG["blue"],
                linewidth=0.3,
                alpha=0.35,
                zorder=2,
            )
        if not links.empty:
            links.plot(ax=ax, color=WONG["black"], linewidth=0.6, zorder=3)
        if not outfalls.empty:
            outfalls.plot(
                ax=ax, color=WONG["black"], marker="*", markersize=40,
                edgecolor="white", linewidth=0.3, zorder=4,
            )

        ax.set_xlim(minx - pad, maxx + pad)
        ax.set_ylim(miny - pad, maxy + pad)
        # 'datalim': the panel keeps the space constrained layout gives it
        # and shows a little more context instead of a blank band.
        ax.set_aspect("equal", adjustable="datalim")
        # Plain 6-7 digit coordinates: no offset, and few enough ticks
        # that the labels do not touch at 89 mm.
        fmt = ScalarFormatter(useOffset=False)
        fmt.set_scientific(False)
        for axis in (ax.xaxis, ax.yaxis):
            axis.set_major_formatter(fmt)
            axis.set_major_locator(MaxNLocator(nbins=5))
        zone = str(utm).split(":")[-1]
        ax.set_xlabel(f"Easting (m, EPSG:{zone})")
        ax.set_ylabel(f"Northing (m, EPSG:{zone})")

        # Scale bar and north arrow inside the 0.25-1 pt line-weight band.
        bar = _scale_bar_length(maxx - minx)
        bx, by = minx + pad * 0.5, maxy - pad * 1.2
        ax.plot([bx, bx + bar], [by, by], color="k", lw=1.0, zorder=5)
        bar_label = f"{bar} m" if bar < 1000 else f"{bar // 1000} km"
        ax.text(bx + bar / 2, by + pad * 0.25, bar_label, ha="center")
        ax.annotate(
            "N",
            xy=(0.955, 0.94),
            xytext=(0.955, 0.865),
            xycoords="axes fraction",
            ha="center",
            fontweight="bold",
            arrowprops=dict(arrowstyle="-|>", color="k", lw=0.75),
        )

        # Cartouche: dense identifying text sits at the bottom of the
        # 5-7 pt band, on a white box so it never reads over the hillshade.
        # Short lines: at 89 mm a long line would spill past the panel and
        # constrained layout would shrink the map to make room for it.
        wgs = frame.to_crs("EPSG:4326").total_bounds
        cartouche = [
            f"{len(subs)} subcatchments, {len(links)} conduits, {len(outfalls)} outfalls",
            f"AOI {wgs[0]:.3f} to {wgs[2]:.3f} E, {wgs[1]:.3f} to {wgs[3]:.3f} N (WGS84)",
            "Data: SWMMCanada bundle",
        ]
        if args.place:
            cartouche.insert(0, args.place)
        ax.text(
            0.012,
            0.012,
            "\n".join(cartouche),
            transform=ax.transAxes,
            fontsize=6,
            va="bottom",
            bbox=dict(facecolor="white", alpha=0.85, edgecolor="none"),
            zorder=6,
        )

        # Key above the panel (proxy handles: geopandas layers do not
        # register legend entries themselves).
        from matplotlib.lines import Line2D
        from matplotlib.patches import Patch

        handles = [
            Patch(facecolor=WONG["sky_blue"], edgecolor=WONG["blue"], alpha=0.35, linewidth=0.3),
            Line2D([], [], color=WONG["black"], linewidth=0.6),
            Line2D([], [], linestyle="none", marker="*", markersize=6,
                   markerfacecolor=WONG["black"], markeredgecolor="white", markeredgewidth=0.3),
        ]
        legend_outside(fig, handles, ["Subcatchment", "Conduit", "Outfall"])

        out_png = args.out_png
        if out_png is None:
            if not args.run_dir:
                raise _fail("--out-png is required when only --bundle is given")
            out_png = Path(args.run_dir) / "08_plot" / "study_area.png"
        out_png = Path(out_png)
        written = save_figure(fig, out_png, dpi=args.dpi)
        print(json.dumps({
            "ok": True,
            "out_png": str(out_png),
            "out_pdf": str(written["pdf"]),
            "subcatchments": int(len(subs)),
            "conduits": int(len(links)),
            "outfalls": int(len(outfalls)),
            "crs": str(utm),
            "dem_hillshade": shade is not None,
        }, indent=2))
    return 0


if __name__ == "__main__":
    sys.exit(main())

"""Shared figure style for the swmm-plot skill: the Nature figure specification.

Every renderer in this directory (hydrograph, network map, study area) goes
through the same three calls, so the figures of one run read as one set and
none of them hand-sets rcParams:

    from plot_style import apply_style, figsize, save_figure

    apply_style()                                   # vendored nature.mplstyle
    fig, ax = plt.subplots(figsize=figsize("single"), layout="constrained")
    ...
    save_figure(fig, out_png)                       # out.pdf (vector) + out.png (preview)

Spec: https://research-figure-guide.nature.com/figures/preparing-figures-our-specifications/
Mirrors the nature-figures skill's ``nature_export.py`` (same numbers, same
guards) so a figure produced here passes its ``check_figure.py`` unchanged.

Why two files per figure: journals take vector PDF/EPS/AI for main figures and
reject PNG outright, so the PDF is the submission file and the PNG is the
preview that reports, chat and slides embed. They share a stem, which is how
the checker knows the PNG is a preview and not an orphan raster.

Never save with ``bbox_inches='tight'``: it silently changes the physical
width and breaks the 89/183 mm rule. Use ``layout='constrained'`` on the
figure and let ``save_figure`` write with the default bbox.
"""

from __future__ import annotations

from pathlib import Path

MM_PER_IN = 25.4
SINGLE_MM, DOUBLE_MM, MAX_H_MM = 89.0, 183.0, 170.0
SINGLE_IN = SINGLE_MM / MM_PER_IN   # 3.5039
DOUBLE_IN = DOUBLE_MM / MM_PER_IN   # 7.2047
MAX_H_IN = MAX_H_MM / MM_PER_IN     # 6.6929

# "For images, minimum 450 dpi" -- the preview PNG is exported at this by default.
PREVIEW_DPI = 450

STYLE_PATH = Path(__file__).resolve().parent.parent / "assets" / "nature.mplstyle"

# Wong colour-blind-safe palette, in the order the spec says to use it
# (black first). Reach for these by name instead of ad-hoc hex codes.
WONG = {
    "black": "#000000",
    "orange": "#E69F00",
    "sky_blue": "#56B4E9",
    "bluish_green": "#009E73",
    "yellow": "#F0E442",
    "blue": "#0072B2",
    "vermillion": "#D55E00",
    "reddish_purple": "#CC79A7",
}

# Categorical cycle for "one colour per group" layers (map sub-networks).
# Black is kept out because it is the colour of the outfall markers, and
# yellow goes last because thin yellow lines vanish on white paper.
CATEGORY_CYCLE = (
    WONG["orange"],
    WONG["sky_blue"],
    WONG["bluish_green"],
    WONG["blue"],
    WONG["vermillion"],
    WONG["reddish_purple"],
    WONG["yellow"],
)


class SpecError(ValueError):
    """The figure cannot be exported without violating the specification."""


def apply_style() -> None:
    """Load the vendored Nature stylesheet into matplotlib's rcParams."""
    import matplotlib.pyplot as plt

    plt.style.use(str(STYLE_PATH))


def figsize(width: str | float = "single", height_mm: float | None = None,
            aspect: float = 0.618) -> tuple[float, float]:
    """Return a spec-legal ``(w, h)`` in inches.

    width      -- ``"single"`` (89 mm), ``"double"`` (183 mm), or a number in mm.
    height_mm  -- explicit height; otherwise ``width * aspect``, capped at 170 mm.
    """
    if width == "single":
        w_in = SINGLE_IN
    elif width == "double":
        w_in = DOUBLE_IN
    else:
        w_in = float(width) / MM_PER_IN
        if w_in > DOUBLE_IN + 1e-6:
            raise SpecError(f"{width} mm exceeds the 183 mm maximum width.")

    h_in = (height_mm / MM_PER_IN) if height_mm is not None else w_in * aspect
    if h_in > MAX_H_IN + 1e-6:
        raise SpecError(
            f"height {h_in * MM_PER_IN:.1f} mm exceeds the 170 mm maximum height."
        )
    return (w_in, h_in)


def map_figsize(bounds: tuple[float, float, float, float],
                width: str | float = "single") -> tuple[float, float]:
    """Figure size for an equal-aspect map: spec width, height from the data.

    ``bounds`` is ``(minx, miny, maxx, maxy)`` in data units. A fixed square
    canvas would leave bands of white space around a wide or tall network,
    so the height follows the extent's own aspect ratio, clamped to
    [half the width, the 170 mm maximum].
    """
    w_in, _ = figsize(width, height_mm=1.0)
    w_mm = w_in * MM_PER_IN
    minx, miny, maxx, maxy = bounds
    dx, dy = max(maxx - minx, 1e-9), max(maxy - miny, 1e-9)
    height_mm = min(max(w_mm * dy / dx, w_mm * 0.5), MAX_H_MM)
    return figsize(width, height_mm=height_mm)


def _check_size(fig) -> tuple[float, float]:
    w_in, h_in = fig.get_size_inches()
    w_mm, h_mm = w_in * MM_PER_IN, h_in * MM_PER_IN
    if w_mm > DOUBLE_MM + 0.5:
        raise SpecError(
            f"figure is {w_mm:.1f} mm wide; the maximum is {DOUBLE_MM:.0f} mm. "
            f"Build the size with figsize('single'|'double')."
        )
    if h_mm > MAX_H_MM + 0.5:
        raise SpecError(
            f"figure is {h_mm:.1f} mm tall; the maximum is {MAX_H_MM:.0f} mm "
            "(leaves room for the legend)."
        )
    return w_mm, h_mm


def save_figure(fig, out_png: str | Path, *, dpi: int = PREVIEW_DPI,
                pdf: bool = True) -> dict[str, Path]:
    """Write ``out_png`` and, by default, its vector twin ``<stem>.pdf``.

    Returns ``{"png": Path, "pdf": Path}`` (``"pdf"`` absent when ``pdf=False``).
    Creates parent directories. Never uses a tight bbox. Raises ``SpecError``
    when the figure is wider than 183 mm or taller than 170 mm.
    """
    out_png = Path(out_png)
    out_png.parent.mkdir(parents=True, exist_ok=True)
    _check_size(fig)

    written: dict[str, Path] = {}
    if pdf:
        out_pdf = out_png.with_suffix(".pdf")
        fig.savefig(out_pdf, format="pdf", facecolor="white")
        written["pdf"] = out_pdf
    fig.savefig(out_png, format="png", dpi=dpi, facecolor="white")
    written["png"] = out_png
    return written


def legend_outside(fig, handles, labels, *, ncol: int | None = None):
    """Figure-level key above the axes, right-aligned, so it never sits on data.

    Text over data is a spec violation (contrast, overlap); an unframed key
    inside the panel cannot promise that on an arbitrary hydrograph. Uses the
    ``outside`` locations of matplotlib >= 3.7 and falls back to an anchored
    corner on older releases.
    """
    kwargs = {"ncol": ncol or max(len(labels), 1), "frameon": False}
    try:
        return fig.legend(handles, labels, loc="outside upper right", **kwargs)
    except ValueError:  # matplotlib < 3.7: no "outside ..." locations
        return fig.legend(handles, labels, loc="upper right",
                          bbox_to_anchor=(1.0, 1.0), **kwargs)


__all__ = [
    "CATEGORY_CYCLE",
    "DOUBLE_IN",
    "DOUBLE_MM",
    "MAX_H_IN",
    "MAX_H_MM",
    "PREVIEW_DPI",
    "SINGLE_IN",
    "SINGLE_MM",
    "STYLE_PATH",
    "WONG",
    "SpecError",
    "apply_style",
    "figsize",
    "legend_outside",
    "map_figsize",
    "save_figure",
]