import numpy as np

from patch_view.patching import apply_patching, compute_patch_statistics


def test_single_constraint_backward_compatibility():
    metric = np.array([2.0, 3.0, 6.0, 2.0, 3.0, 4.0])

    patches = apply_patching(metric, constraint="global", threshold=5.0)

    assert len(patches) == 2
    assert patches[0] == [0, 1]
    assert patches[1] == [2, 3, 4, 5]


def test_multi_constraint_combination():
    metric = np.array([1.0, 2.0, 3.0, 8.0, 2.0, 3.0])
    patches = apply_patching(
        metric,
        constraints={
            "global": 5.0,
            "additive": 3.0,
        },
    )

    assert len(patches) == 2
    assert patches[0] == [0, 1, 2]
    assert patches[1] == [3, 4, 5]


def test_cumulative_constraint_behavior():
    metric = np.array([2.0, 3.0, 4.0, 2.0, 3.0, 4.0])
    patches = apply_patching(metric, constraints={"cumulative": 8.0})

    assert len(patches) == 3
    assert patches[0] == [0, 1]
    assert patches[1] == [2, 3, 4]
    assert patches[2] == [5]


def test_cumulative_max_constraint_keeps_sums_small():
    metric = np.array([2.0, 3.0, 4.0, 2.0, 3.0])
    patches = apply_patching(metric, constraints={"cumulative_max": 8.0})

    for patch in patches:
        patch_sum = np.sum(metric[patch])
        assert patch_sum <= 11.0  # allow threshold plus one token overflow


def test_compute_patch_statistics_matches_expected_values():
    patches = [[0, 1], [2, 3, 4], [5]]
    stats = compute_patch_statistics(patches)

    assert stats["num_patches"] == 3
    assert stats["avg_size"] == np.mean([2, 3, 1])
    assert stats["min_size"] == 1
    assert stats["max_size"] == 3