import os
import numpy as np


def test_select_trajectory_plot_data_with_party_centroids():
    # Synthetic positions_by_window: two windows with MPs mapping to parties
    positions_by_window = {
        "2024-Q1": {
            "A": (0.1, 0.2),
            "B": (0.2, 0.25),
        },
        "2024-Q2": {
            "A": (0.15, 0.22),
            "B": (0.21, 0.27),
        },
    }

    party_map = {"A": "P1", "B": "P2"}
    windows = sorted(list(positions_by_window.keys()))
    selected_parties = ["P1", "P2"]

    from explorer import select_trajectory_plot_data

    fig, trace_count, banner = select_trajectory_plot_data(
        positions_by_window, party_map, windows, selected_parties, smooth_alpha=0.35
    )

    assert hasattr(fig, "data")
    assert trace_count > 0
    # traces should include party names
    names = [getattr(t, "name", None) for t in fig.data]
    assert "P1" in names or "P2" in names
    assert banner is None or banner == ""


def test_select_trajectory_plot_data_fallback_to_mps():
    # No parties known in party_map -> centroids will be all NaN
    positions_by_window = {
        "2024-Q1": {"mp1": (0.1, 0.2)},
        "2024-Q2": {"mp2": (0.2, 0.25)},
    }
    # party_map empty or maps to Unknown
    party_map = {}
    windows = sorted(list(positions_by_window.keys()))
    selected_parties = []

    # make fallback threshold small for test
    os.environ.pop("EXPLORER_MP_FALLBACK_COUNT", None)

    from explorer import select_trajectory_plot_data

    fig, trace_count, banner = select_trajectory_plot_data(
        positions_by_window, party_map, windows, selected_parties, smooth_alpha=0.35
    )

    assert hasattr(fig, "data")
    assert trace_count > 0
    assert (
        banner
        == "Partijcentroiden niet beschikbaar — tonen individuele MP-trajecten als fallback."
    )