"""Tests for visualisation module."""

import numpy as np
import pytest

from sltk.data.core import PoseSequence, Segment, SegmentList


class TestPoseVisualization:
    """Tests for pose visualization utilities."""

    def test_skeleton_connections_defined(self):
        """Test that skeleton connections are defined for formats."""
        from sltk.visualisation.poses import SKELETON_CONNECTIONS

        assert "smplx" in SKELETON_CONNECTIONS
        assert "mediapipe" in SKELETON_CONNECTIONS
        assert "wilor" in SKELETON_CONNECTIONS

        # Each should have a list of (start, end) tuples
        for format_name, connections in SKELETON_CONNECTIONS.items():
            assert isinstance(connections, list)
            for conn in connections:
                assert len(conn) == 2

    def test_color_schemes_defined(self):
        """Test that color schemes are available."""
        from sltk.visualisation.poses import COLOR_SCHEMES

        assert "default" in COLOR_SCHEMES
        assert isinstance(COLOR_SCHEMES["default"], dict)

    def test_plot_pose_2d_runs(self):
        """Test that plot_pose_2d runs without error."""
        pytest.importorskip("matplotlib")
        import matplotlib

        matplotlib.use("Agg")  # Non-interactive backend
        import matplotlib.pyplot as plt

        from sltk.visualisation.poses import plot_pose_2d

        data = np.random.randn(1, 21, 3).astype(np.float32)
        poses = PoseSequence(data=data, fps=25.0, format="wilor")

        # Should not raise
        ax = plot_pose_2d(poses, frame=0)
        assert ax is not None
        plt.close("all")

    def test_plot_pose_sequence_runs(self):
        """Test that plot_pose_sequence runs without error."""
        pytest.importorskip("matplotlib")
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

        from sltk.visualisation.poses import plot_pose_sequence

        data = np.random.randn(50, 21, 3).astype(np.float32)
        poses = PoseSequence(data=data, fps=25.0, format="wilor")

        # Should not raise
        fig = plot_pose_sequence(poses, frames=[0, 10, 20])
        assert fig is not None
        plt.close("all")


class TestSegmentVisualization:
    """Tests for segment visualization."""

    def test_plot_timeline_runs(self):
        """Test that plot_timeline runs without error."""
        pytest.importorskip("matplotlib")
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

        from sltk.visualisation.segments import plot_timeline

        segments = SegmentList(
            [
                Segment(start=0.0, end=1.0, label="A", tier="Gloss"),
                Segment(start=1.0, end=2.0, label="B", tier="Gloss"),
                Segment(start=0.0, end=2.0, label="Translation", tier="Translation"),
            ]
        )

        # Should not raise
        ax = plot_timeline(segments)
        assert ax is not None
        plt.close("all")

    def test_plot_segments_with_tier_filter(self):
        """Test timeline plot with tier filtering."""
        pytest.importorskip("matplotlib")
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

        from sltk.visualisation.segments import plot_timeline

        segments = SegmentList(
            [
                Segment(start=0.0, end=1.0, label="A", tier="Gloss"),
                Segment(start=1.0, end=2.0, label="B", tier="Translation"),
            ]
        )

        # Filter to only Gloss tier
        filtered = segments.filter_by_tier("Gloss")
        ax = plot_timeline(filtered)
        assert ax is not None
        plt.close("all")

    def test_segment_to_color(self):
        """Test segment color assignment."""
        from sltk.visualisation.segments import get_tier_colors

        tiers = ["Gloss", "Translation", "Mouthing"]
        colors = get_tier_colors(tiers)

        assert len(colors) == 3
        assert all(c is not None for c in colors)


class TestAnalysisVisualization:
    """Tests for analysis visualization."""

    def test_plot_frequency_distribution_runs(self):
        """Test frequency distribution plot."""
        pytest.importorskip("matplotlib")
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

        from sltk.visualisation.analysis import plot_frequency_distribution

        freq = {"A": 100, "B": 50, "C": 25, "D": 10}

        # Should not raise
        ax = plot_frequency_distribution(freq)
        assert ax is not None
        plt.close("all")

    def test_plot_vocabulary_growth_runs(self):
        """Test vocabulary growth plot."""
        pytest.importorskip("matplotlib")
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

        from sltk.visualisation.analysis import plot_vocabulary_growth

        growth_data = {
            "sample_counts": [10, 20, 30, 40, 50],
            "vocabulary_sizes": [5, 8, 10, 11, 12],
        }

        # Should not raise
        ax = plot_vocabulary_growth(growth_data)
        assert ax is not None
        plt.close("all")

    def test_plot_confusion_matrix_runs(self):
        """Test confusion matrix plot."""
        pytest.importorskip("matplotlib")
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

        from sltk.visualisation.analysis import plot_confusion_matrix

        matrix = np.array(
            [
                [10, 2, 1],
                [1, 15, 3],
                [2, 1, 12],
            ]
        )
        labels = ["A", "B", "C"]

        # Should not raise
        ax = plot_confusion_matrix(matrix, labels)
        assert ax is not None
        plt.close("all")


class TestAnimationUtilities:
    """Tests for pose animation utilities."""

    def test_create_animation_frames(self):
        """Test animation frame data preparation."""
        from sltk.visualisation.poses import prepare_animation_data

        data = np.random.randn(100, 21, 3).astype(np.float32)
        poses = PoseSequence(data=data, fps=25.0, format="wilor")

        frames = prepare_animation_data(poses, step=10)

        # Should have 10 frames (0, 10, 20, ..., 90)
        assert len(frames) == 10
        for frame_data in frames:
            assert frame_data.shape == (21, 3)


class TestSegmentComparison:
    """Tests for segment comparison visualization."""

    def test_plot_segment_comparison_runs(self):
        """Test segment comparison plot."""
        pytest.importorskip("matplotlib")
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

        from sltk.visualisation.segments import plot_segment_comparison

        reference = SegmentList(
            [
                Segment(start=0.0, end=1.0, label="A"),
                Segment(start=1.5, end=2.5, label="B"),
            ]
        )
        prediction = SegmentList(
            [
                Segment(start=0.1, end=1.1, label="A"),
                Segment(start=1.4, end=2.4, label="B"),
            ]
        )

        ax = plot_segment_comparison(reference, prediction)
        assert ax is not None
        plt.close("all")
