"""
Test that trajectory plot renders even with edge cases.
"""

import pytest
import numpy as np
from unittest.mock import MagicMock, patch

# Import the functions to test
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from explorer_helpers import compute_party_centroids


class TestTrajectoryPlotRendering:
    """Tests to ensure trajectory plot renders in various scenarios."""

    def test_compute_party_centroids_returns_diagnostics(self):
        """Test that compute_party_centroids returns diagnostics tuple."""
        positions_by_window = {
            "2024-Q1": {"MP1": (1.0, 2.0), "MP2": (3.0, 4.0)},
            "2024-Q2": {"MP1": (1.5, 2.5), "MP2": (3.5, 4.5)},
        }
        party_map = {"MP1": "PartyA", "MP2": "PartyA"}
        windows = ["2024-Q1", "2024-Q2"]

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

        assert isinstance(centroids, dict)
        assert isinstance(diagnostics, dict)
        assert "windows_with_data_count" in diagnostics
        assert diagnostics["windows_with_data_count"] == 2

    def test_compute_party_centroids_detects_all_nan_parties(self):
        """Test that diagnostics identify parties with all NaN centroids."""
        positions_by_window = {
            "2024-Q1": {"MP1": (np.nan, np.nan)},
            "2024-Q2": {"MP1": (np.nan, np.nan)},
        }
        party_map = {"MP1": "PartyA"}
        windows = ["2024-Q1", "2024-Q2"]

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

        assert "PartyA" in diagnostics.get("parties_all_nan", [])


if __name__ == "__main__":
    pytest.main([__file__, "-v"])