TJWaterServerBinary/app/algorithms/valve_isolation/topology_search.py

"""Topology-only valve isolation search."""

from collections import defaultdict, deque
from typing import Any, Iterable, Literal


VALVE_LINK_TYPE = "valve"
AccidentElementType = Literal["auto", "link", "node"]


def _parse_link_entry(link_entry: str) -> tuple[str, str, str, str]:
    parts = link_entry.split(":", 3)
    if len(parts) != 4:
        raise ValueError(f"Invalid link entry format: {link_entry}")
    return parts[0], parts[1], parts[2], parts[3]


def _reachable_nodes(
    adjacency: dict[str, list[tuple[str, str]]],
    start_nodes: Iterable[str],
    blocked_links: set[str],
) -> set[str]:
    visited = set(start_nodes)
    queue = deque(visited)
    while queue:
        for neighbor, link_id in adjacency.get(queue.popleft(), ()):
            if link_id not in blocked_links and neighbor not in visited:
                visited.add(neighbor)
                queue.append(neighbor)
    return visited


def valve_isolation_analysis(
    link_entries: Iterable[str],
    accident_elements: str | list[str],
    disabled_valves: list[str] | None = None,
    *,
    accident_element_type: AccidentElementType = "auto",
    water_source_nodes: Iterable[str] | None = None,
    closed_links: Iterable[str] = (),
) -> dict[str, Any]:
    """Find nearest valve boundaries, then check source connectivity.

    Legacy untyped IDs prefer links; explicit node mode handles cross-type ID
    collisions. Supply impact is a before/after connectivity difference, not
    a hydraulic pressure or demand calculation. Without source metadata only
    the provisional boundary search is available.
    """

    disabled_valves_set = set(disabled_valves or [])
    target_elements = (
        [accident_elements]
        if isinstance(accident_elements, str)
        else accident_elements
    )
    if not target_elements:
        raise ValueError("At least one accident element is required")
    if accident_element_type not in {"auto", "link", "node"}:
        raise ValueError(f"Invalid accident element type: {accident_element_type}")
    closed_links_set = set(closed_links)

    adjacency: dict[str, list[tuple[str, str]]] = defaultdict(list)
    all_valves: dict[str, tuple[str, str]] = {}
    link_lookup: dict[str, tuple[str, str, str]] = {}
    node_set: set[str] = set()
    for link_entry in link_entries:
        link_id, link_type, node1, node2 = _parse_link_entry(link_entry)
        link_type_name = str(link_type).lower()
        link_lookup[link_id] = (node1, node2, link_type_name)
        node_set.update((node1, node2))
        if link_type_name == VALVE_LINK_TYPE:
            all_valves[link_id] = (node1, node2)
        if link_id not in closed_links_set:
            adjacency[node1].append((node2, link_id))
            adjacency[node2].append((node1, link_id))

    unknown_disabled = disabled_valves_set - all_valves.keys()
    if unknown_disabled:
        raise ValueError(f"Disabled valve was not found in topology: {sorted(unknown_disabled)}")

    start_nodes: set[str] = set()
    for element in target_elements:
        if accident_element_type != "node" and element in link_lookup:
            node1, node2, _ = link_lookup[element]
            start_nodes.update((node1, node2))
        elif accident_element_type != "link" and element in node_set:
            start_nodes.add(element)
        else:
            raise ValueError(f"Accident element {element} was not found in topology")

    closable_valves = set(all_valves) - disabled_valves_set - closed_links_set
    affected_nodes = _reachable_nodes(adjacency, start_nodes, closable_valves)

    must_close_valves: list[str] = []
    optional_valves: list[str] = []
    for valve_id in closable_valves:
        node1, node2 = all_valves[valve_id]
        node1_affected = node1 in affected_nodes
        node2_affected = node2 in affected_nodes
        if node1_affected and node2_affected:
            optional_valves.append(valve_id)
        elif node1_affected or node2_affected:
            must_close_valves.append(valve_id)

    must_close_valves.sort()
    optional_valves.sort()
    isolatable = bool(must_close_valves)
    failure_reason = None if isolatable else "no_boundary_valves"
    sources = set(water_source_nodes or ())
    source_check_performed = bool(sources)
    supply_affected_nodes = None
    if water_source_nodes is not None:
        if not sources:
            isolatable = False
            failure_reason = "water_sources_unavailable"
        else:
            after_closure = _reachable_nodes(adjacency, sources, set(must_close_valves))
            if affected_nodes & after_closure:
                isolatable = False
                failure_reason = "connected_to_water_source"
            if isolatable:
                before_closure = _reachable_nodes(adjacency, sources, set())
                supply_affected_nodes = sorted(before_closure - after_closure)
    result: dict[str, Any] = {
        "accident_elements": target_elements,
        "accident_element_type": accident_element_type,
        "disabled_valves": disabled_valves,
        "affected_nodes": sorted(affected_nodes) if isolatable else [],
        "affected_node_count": len(affected_nodes),
        "must_close_valves": must_close_valves,
        "optional_valves": optional_valves,
        "isolatable": isolatable,
        "source_check_performed": source_check_performed,
        "isolation_failure_reason": failure_reason,
        "supply_affected_nodes": supply_affected_nodes,
        "supply_affected_node_count": (
            len(supply_affected_nodes) if supply_affected_nodes is not None else None
        ),
    }
    if len(target_elements) == 1:
        result["accident_element"] = target_elements[0]
    return result