from __future__ import annotations

import math

from blocksys.boundary import BoundaryLink, apply_blocks_to_model
from blocksys.validate import hydraulic_validation


def test_baseline_simulation_has_all_nodes(bundle, baseline):
    assert set(baseline.node_pressure_mean.keys()) == set(bundle.wn.node_name_list)
    assert set(baseline.link_flow_mean_abs.keys()) == set(bundle.wn.link_name_list)


def test_baseline_pressures_are_finite(baseline):
    assert all(math.isfinite(v) for v in baseline.node_pressure_mean.values())


def test_hydraulic_validation_runs_on_unmodified_network(bundle, baseline, cfg):
    result = hydraulic_validation(bundle.wn, baseline, cfg)
    assert result.status in ("ok", "violated")
    assert not result.isolated_nodes


def test_isolation_is_detected_and_rejected(bundle, baseline, cfg):
    graph = bundle.graph
    source_names = set(bundle.wn.reservoir_name_list) | set(bundle.wn.tank_name_list)

    leaf = next(n for n in graph.nodes if graph.degree(n) == 1 and n not in source_names)
    _, neighbor, edge_data = next(iter(graph.edges(leaf, data=True)))
    pipe_id = edge_data["pipe_ids"][0]

    boundary = [
        BoundaryLink(
            pipe_id=pipe_id,
            from_node=leaf,
            to_node=neighbor,
            block_a=0,
            block_b=1,
            action="close",
            baseline_flow=0.0,
            diameter=0.0,
        )
    ]
    wn_blocked = apply_blocks_to_model(bundle.wn, boundary)

    result = hydraulic_validation(wn_blocked, baseline, cfg)
    assert leaf in result.isolated_nodes
    assert result.status == "rejected"
    assert result.rejection_reason


def test_allow_isolated_nodes_still_simulates(bundle, baseline, cfg):
    graph = bundle.graph
    source_names = set(bundle.wn.reservoir_name_list) | set(bundle.wn.tank_name_list)
    leaf = next(n for n in graph.nodes if graph.degree(n) == 1 and n not in source_names)
    _, neighbor, edge_data = next(iter(graph.edges(leaf, data=True)))
    pipe_id = edge_data["pipe_ids"][0]

    boundary = [
        BoundaryLink(
            pipe_id=pipe_id, from_node=leaf, to_node=neighbor, block_a=0, block_b=1,
            action="close", baseline_flow=0.0, diameter=0.0,
        )
    ]
    wn_blocked = apply_blocks_to_model(bundle.wn, boundary)

    cfg.hydraulics.allow_isolated_nodes = True
    try:
        result = hydraulic_validation(wn_blocked, baseline, cfg)
        assert leaf in result.isolated_nodes
        assert result.status == "violated"
    finally:
        cfg.hydraulics.allow_isolated_nodes = False