from __future__ import annotations

from blocksys.boundary import apply_blocks_to_model, assign_valve_or_meter, find_boundary_links
from blocksys.partition import partition


def test_boundary_matches_differing_labels(weighted_graph, cfg):
    labels = partition(weighted_graph, 6, cfg)
    boundary = find_boundary_links(weighted_graph, labels)

    expected_pairs = {
        (u, v)
        for u, v in weighted_graph.edges
        if u in labels and v in labels and labels[u] != labels[v]
    }
    actual_pairs = {(b.from_node, b.to_node) for b in boundary}
    assert actual_pairs == expected_pairs


def test_each_boundary_link_has_a_real_pipe_id(weighted_graph, bundle, cfg):
    labels = partition(weighted_graph, 6, cfg)
    boundary = find_boundary_links(weighted_graph, labels)
    for b in boundary:
        assert b.pipe_id in bundle.wn.link_name_list


def test_assign_valve_or_meter_gives_every_block_an_inlet(weighted_graph, baseline, cfg):
    labels = partition(weighted_graph, 6, cfg)
    boundary = find_boundary_links(weighted_graph, labels)
    boundary = assign_valve_or_meter(boundary, baseline, inlets_per_block=cfg.boundary.inlets_per_block)

    assert {b.action for b in boundary} <= {"close", "meter"}

    blocks_touched = {b.block_a for b in boundary} | {b.block_b for b in boundary}
    for block in blocks_touched:
        links = [b for b in boundary if b.block_a == block or b.block_b == block]
        assert any(b.action == "meter" for b in links)


def test_apply_blocks_to_model_only_closes_marked_links(weighted_graph, bundle, baseline, cfg):
    labels = partition(weighted_graph, 6, cfg)
    boundary = find_boundary_links(weighted_graph, labels)
    boundary = assign_valve_or_meter(boundary, baseline, inlets_per_block=cfg.boundary.inlets_per_block)

    already_closed = {
        name for name, link in bundle.wn.links() if str(link.initial_status) == "Closed"
    }
    wn_blocked = apply_blocks_to_model(bundle.wn, boundary)

    newly_closed_expected = {b.pipe_id for b in boundary if b.action == "close"}
    closed_actual = {
        name for name, link in wn_blocked.links() if str(link.initial_status) == "Closed"
    }
    assert closed_actual == already_closed | newly_closed_expected
    # 원본 모델은 변경되지 않아야 한다
    assert all(
        str(bundle.wn.get_link(pid).initial_status) != "Closed"
        for pid in newly_closed_expected - already_closed
    )