import numpy as np
from explorer_helpers import compute_party_coords, compute_party_centroids


def test_compute_party_coords_basic():
    # synthetic positions: two windows
    positions_by_window = {
        "2024": {
            "Alice": (0.1, 0.2),
            "Bob": (0.3, 0.4),
            "Carol": (0.5, -0.1),
        }
    }
    party_map = {"Alice": "P1", "Bob": "P1", "Carol": "P2"}

    coords, fallback = compute_party_coords(positions_by_window, party_map, "2024")
    assert "P1" in coords and "P2" in coords
    # P1 mean of (0.1,0.2) and (0.3,0.4) => (0.2,0.3)
    assert abs(coords["P1"][0] - 0.2) < 1e-9
    assert abs(coords["P1"][1] - 0.3) < 1e-9
    assert abs(coords["P2"][0] - 0.5) < 1e-9
    assert abs(coords["P2"][1] - -0.1) < 1e-9
    assert fallback == set()


def test_compute_party_coords_with_fallback():
    positions_by_window = {"2024": {"Alice": (0.1, 0.1)}}
    party_map = {"Alice": "P1"}
    fallback_party_scores = {"P2": [1.234, -0.987, 0.0]}

    coords, fallback = compute_party_coords(
        positions_by_window, party_map, "2024", fallback_party_scores
    )
    assert coords["P1"][0] == 0.1
    assert coords["P2"][0] == 1.234
    assert "P2" in fallback


def test_compute_party_centroids_nan_handling():
    """Ensure compute_party_centroids fills missing windows with (np.nan, np.nan).

    Build synthetic positions where P1 has a centroid in window 'w1' but not in 'w2'.
    The resulting party_centroids for P1 should be [(x,y), (nan,nan)].
    """
    positions_by_window = {
        "w1": {"Alice": (0.1, 0.2)},
        "w2": {},
    }
    party_map = {"Alice": "P1"}
    windows = ["w1", "w2"]

    party_centroids, metadata = compute_party_centroids(
        positions_by_window, party_map, windows
    )

    assert "P1" in party_centroids
    vals = party_centroids["P1"]
    assert len(vals) == 2
    # first window has numeric coords
    assert not (np.isnan(vals[0][0]) or np.isnan(vals[0][1]))
    # second window should be nan-filled
    assert np.isnan(vals[1][0]) and np.isnan(vals[1][1])