TJWaterServerBinary/app/native/wndb/gis/network_views.py

"""Read-only GIS projections and authoritative network topology snapshots."""

from typing import Any

from ..core.database import read, read_all


def _node_coord_rows(
    rows: list[Any],
) -> dict[str, dict[str, Any]]:
    return {
        str(row["id"]): {
            "x": float(row["x"]),
            "y": float(row["y"]),
            "type": str(row["node_type"]),
        }
        for row in rows
    }


def get_network_node_coords(name: str) -> dict[str, dict[str, Any]]:
    """Return every publishable node with one view-backed query."""
    rows = read_all(
        name,
        """
        SELECT id, x, y, node_type
        FROM gis.network_nodes
        ORDER BY id
        """,
    )
    return _node_coord_rows(rows)


def get_major_node_coords(
    name: str, diameter: int
) -> dict[str, dict[str, Any]]:
    """Return endpoints of pipes above the requested diameter."""
    rows = read_all(
        name,
        """
        SELECT n.id, n.x, n.y, n.node_type
        FROM gis.pipes AS p
        CROSS JOIN LATERAL (
            VALUES (p.start_node_id), (p.end_node_id)
        ) AS endpoint(id)
        JOIN gis.network_nodes AS n ON n.id = endpoint.id
        WHERE p.diameter > %s
        GROUP BY n.id, n.x, n.y, n.node_type
        ORDER BY n.id
        """,
        (diameter,),
    )
    return _node_coord_rows(rows)


def get_network_link_nodes(name: str) -> list[str]:
    """Return every publishable link in the established API wire format."""
    rows = read_all(
        name,
        """
        SELECT id, link_type, start_node_id, end_node_id
        FROM gis.network_links
        ORDER BY id
        """,
    )
    return [
        f"{row['id']}:{row['link_type']}:{row['start_node_id']}:{row['end_node_id']}"
        for row in rows
    ]


def get_valve_isolation_topology(name: str) -> tuple[list[str], list[str], list[str]]:
    """Read links, closed statuses and sources from one consistent SQL snapshot.

    Isolation must include model links without published geometries, and use
    initial statuses rather than treating every link as open. UNION ALL keeps
    source nodes in the same statement snapshot without repeating them for
    every link or relying on materialized view refresh timing.
    """
    rows = read_all(
        name,
        """
        SELECT l.id, l.link_type::text AS element_type,
               l.start_node_id, l.end_node_id,
               COALESCE(s.status::text, p.status::text) AS status
        FROM network.links AS l
        LEFT JOIN network.pipes AS p ON p.link_id = l.id
        LEFT JOIN network.link_initial_settings AS s ON s.link_id = l.id
        UNION ALL
        SELECT n.id, n.node_type::text AS element_type,
               NULL, NULL, NULL
        FROM network.nodes AS n
        WHERE n.node_type::text IN ('reservoir', 'tank')
        """,
    )
    link_entries: list[str] = []
    source_nodes: list[str] = []
    closed_links: list[str] = []
    for row in rows:
        element_id = str(row["id"])
        if row["element_type"] in {"reservoir", "tank"}:
            source_nodes.append(element_id)
        else:
            link_entries.append(
                f"{element_id}:{row['element_type']}:{row['start_node_id']}:{row['end_node_id']}"
            )
            if row["status"] == "CLOSED":
                closed_links.append(element_id)
    return link_entries, source_nodes, closed_links


def get_major_pipe_nodes(name: str, diameter: int) -> list[str]:
    """Return large pipes in the established API wire format."""
    rows = read_all(
        name,
        """
        SELECT id, start_node_id, end_node_id
        FROM gis.pipes
        WHERE diameter > %s
        ORDER BY id
        """,
        (diameter,),
    )
    return [
        f"{row['id']}:pipe:{row['start_node_id']}:{row['end_node_id']}"
        for row in rows
    ]


def get_topology_rows(
    name: str, node_ids: list[str]
) -> tuple[list[Any], list[Any]]:
    """Load selected nodes and their internal links with two batch queries."""
    if not node_ids:
        return [], []

    nodes = read_all(
        name,
        """
        SELECT n.id,
               ST_X(g.geom) AS x,
               ST_Y(g.geom) AS y,
               n.node_type::text AS node_type
        FROM network.nodes AS n
        JOIN gis.node_geometries AS g ON g.node_id = n.id
        WHERE n.id = ANY(%s)
        ORDER BY n.id
        """,
        (node_ids,),
    )
    links = read_all(
        name,
        """
        SELECT l.id,
               l.start_node_id,
               l.end_node_id,
               COALESCE(p.length, 0.0) AS length
        FROM network.links AS l
        LEFT JOIN network.pipes AS p ON p.link_id = l.id
        WHERE l.start_node_id = ANY(%s)
          AND l.end_node_id = ANY(%s)
        ORDER BY l.id
        """,
        (node_ids, node_ids),
    )
    return nodes, links


def get_boundary_link_ids(name: str, node_ids: list[str]) -> list[str]:
    """Return links with exactly one endpoint inside the supplied node set."""
    if not node_ids:
        return []
    rows = read_all(
        name,
        """
        SELECT id
        FROM network.links
        WHERE (start_node_id = ANY(%s)) <> (end_node_id = ANY(%s))
        ORDER BY id
        """,
        (node_ids, node_ids),
    )
    return [str(row["id"]) for row in rows]


def get_junction_demands(
    name: str, node_ids: list[str]
) -> dict[str, list[dict[str, Any]]]:
    """Read authoritative demands for selected junctions from model tables."""
    if not node_ids:
        return {}
    rows = read_all(
        name,
        """
        SELECT junction_id,
               sequence_no,
               base_demand,
               pattern_id,
               category
        FROM network.demands
        WHERE junction_id = ANY(%s)
        ORDER BY junction_id, sequence_no
        """,
        (node_ids,),
    )
    result: dict[str, list[dict[str, Any]]] = {}
    for row in rows:
        result.setdefault(str(row["junction_id"]), []).append(
            {
                "demand": float(row["base_demand"]),
                "pattern": row["pattern_id"],
                "category": row["category"],
            }
        )
    return result


def sum_junction_base_demand(name: str, node_ids: list[str]) -> float:
    """Sum selected junction demand from the authoritative model table."""
    if not node_ids:
        return 0.0
    row = read(
        name,
        """
        SELECT COALESCE(SUM(base_demand), 0.0) AS total_base_demand
        FROM network.demands
        WHERE junction_id = ANY(%s)
        """,
        (node_ids,),
    )
    return float(row["total_base_demand"])