from layer_interactions.result import Order2Result


def test_interactions_and_selection():
    # D0 = 1.0. First-order deletion costs: [0.1, 0.2, 0.05]
    # Pair (0,2) has a positive interaction of +0.3, so after deleting 2,
    # deleting 1 is cheaper than deleting 0.
    r = Order2Result(
        depth=3,
        baseline_nll=1.0,
        single_nll={0: 1.1, 1: 1.2, 2: 1.05},
        pair_nll={
            (0, 1): 1.3,
            (0, 2): 1.45,
            (1, 2): 1.25,
        },
    )
    r.build_interactions().build_greedy_path(max_delete=2)
    assert r.delete_order == [2, 1]
    assert r.select(1)["retained_layers"] == [0]


def test_roundtrip_dict():
    r = Order2Result(
        depth=2,
        baseline_nll=1.0,
        single_nll={0: 1.1, 1: 1.2},
        pair_nll={(0, 1): 1.35},
    ).build_interactions().build_greedy_path(max_delete=1)
    rr = Order2Result.from_dict(r.to_dict())
    assert rr.depth == 2
    assert rr.single_nll == r.single_nll
    assert rr.pair_nll == r.pair_nll
    assert rr.delete_order == r.delete_order