import numpy as np
from explorer_helpers import compute_party_centroids


def test_full_coverage():
    windows = ["w1", "w2"]
    positions_by_window = {
        "w1": {"mp1": (0.0, 0.0), "mp2": (2.0, 0.0)},
        "w2": {"mp1": (1.0, 1.0), "mp2": (3.0, 1.0)},
    }
    party_map = {"mp1": "P1", "mp2": "P2"}

    centroids, meta = compute_party_centroids(positions_by_window, party_map, windows)

    # both parties present in both windows -> no nans and correct lengths
    assert set(centroids.keys()) == {"P1", "P2"}
    for vals in centroids.values():
        assert len(vals) == len(windows)
        for x, y in vals:
            assert not (np.isnan(x) or np.isnan(y))


def test_partial_coverage():
    windows = ["w1", "w2", "w3"]
    positions_by_window = {
        "w1": {"mp1": (0.0, 0.0), "mp2": (2.0, 0.0)},
        "w2": {"mp1": (1.0, 1.0)},
        "w3": {"mp2": (3.0, 1.0)},
    }
    party_map = {"mp1": "P1", "mp2": "P2"}

    centroids, meta = compute_party_centroids(positions_by_window, party_map, windows)

    # Expect P1 present in w1,w2 but missing in w3
    assert centroids["P1"][0] == (0.0, 0.0)
    assert centroids["P1"][1] == (1.0, 1.0)
    assert np.isnan(centroids["P1"][2][0]) and np.isnan(centroids["P1"][2][1])

    # Expect P2 present in w1,w3 but missing in w2
    assert centroids["P2"][0] == (2.0, 0.0)
    assert np.isnan(centroids["P2"][1][0]) and np.isnan(centroids["P2"][1][1])
    assert centroids["P2"][2] == (3.0, 1.0)

    # metadata counts should reflect non-nan entries
    assert meta["per_party_counts"]["P1"] == 2
    assert meta["per_party_counts"]["P2"] == 2
    assert meta["total_windows"] == len(windows)


def test_no_parties():
    windows = ["w1", "w2"]
    positions_by_window = {}
    party_map = {}

    centroids, meta = compute_party_centroids(positions_by_window, party_map, windows)

    assert centroids == {}
    assert meta["total_windows"] == len(windows)