"""Tests for core data structures."""

import numpy as np
import pytest

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


class TestPoseSequence:
    """Tests for PoseSequence."""

    def test_creation(self):
        """Test basic PoseSequence creation."""
        data = np.random.randn(100, 21, 3).astype(np.float32)
        poses = PoseSequence(data=data, fps=25.0, format="test")

        assert poses.num_frames == 100
        assert poses.num_keypoints == 21
        assert poses.num_coords == 3
        assert poses.duration == 4.0  # 100 frames / 25 fps

    def test_slicing_frames(self):
        """Test slicing by frame index."""
        data = np.random.randn(100, 21, 3).astype(np.float32)
        poses = PoseSequence(data=data, fps=25.0, format="test")

        subset = poses[10:50]
        assert subset.num_frames == 40
        assert subset.fps == 25.0
        assert subset.format == "test"

    def test_slicing_time(self):
        """Test slicing by time."""
        data = np.random.randn(100, 21, 3).astype(np.float32)
        poses = PoseSequence(data=data, fps=25.0, format="test")

        subset = poses.slice_time(1.0, 2.0)
        assert subset.num_frames == 25  # 1 second at 25 fps

    def test_with_confidence(self):
        """Test PoseSequence with confidence scores."""
        data = np.random.randn(100, 21, 3).astype(np.float32)
        confidence = np.random.rand(100, 21).astype(np.float32)
        poses = PoseSequence(data=data, fps=25.0, format="test", confidence=confidence)

        assert poses.confidence is not None
        assert poses.confidence.shape == (100, 21)

        # Confidence should be sliced too
        subset = poses[10:50]
        assert subset.confidence.shape == (40, 21)


class TestSegment:
    """Tests for Segment."""

    def test_creation(self):
        """Test basic Segment creation."""
        seg = Segment(start=1.0, end=2.5, label="HELLO", tier="Gloss")

        assert seg.start == 1.0
        assert seg.end == 2.5
        assert seg.duration == 1.5
        assert seg.midpoint == 1.75
        assert seg.label == "HELLO"
        assert seg.tier == "Gloss"

    def test_invalid_segment(self):
        """Test that end < start raises error."""
        with pytest.raises(ValueError):
            Segment(start=2.0, end=1.0)

    def test_overlaps(self):
        """Test overlap detection."""
        seg1 = Segment(start=1.0, end=2.0)
        seg2 = Segment(start=1.5, end=2.5)
        seg3 = Segment(start=2.5, end=3.0)

        assert seg1.overlaps(seg2)
        assert seg2.overlaps(seg1)
        assert not seg1.overlaps(seg3)

    def test_iou(self):
        """Test IoU calculation."""
        seg1 = Segment(start=0.0, end=2.0)
        seg2 = Segment(start=1.0, end=3.0)

        iou = seg1.iou(seg2)
        # Intersection: 1.0-2.0 = 1.0
        # Union: 0.0-3.0 = 3.0
        # IoU = 1/3
        assert abs(iou - 1 / 3) < 1e-6

    def test_contains(self):
        """Test time point containment."""
        seg = Segment(start=1.0, end=2.0)

        assert seg.contains(1.5)
        assert seg.contains(1.0)
        assert seg.contains(2.0)
        assert not seg.contains(0.5)
        assert not seg.contains(2.5)


class TestSegmentList:
    """Tests for SegmentList."""

    def test_creation(self):
        """Test SegmentList creation."""
        segments = SegmentList(
            [
                Segment(0.0, 1.0, "A", tier="Gloss"),
                Segment(1.5, 2.0, "B", tier="Gloss"),
                Segment(0.5, 1.5, "C", tier="Other"),
            ]
        )

        assert len(segments) == 3
        assert segments.tiers == ["Gloss", "Other"]

    def test_filter_by_tier(self):
        """Test filtering by tier."""
        segments = SegmentList(
            [
                Segment(0.0, 1.0, "A", tier="Gloss"),
                Segment(1.5, 2.0, "B", tier="Gloss"),
                Segment(0.5, 1.5, "C", tier="Other"),
            ]
        )

        gloss_segs = segments.filter_by_tier("Gloss")
        assert len(gloss_segs) == 2

    def test_filter_by_label(self):
        """Test filtering by label."""
        segments = SegmentList(
            [
                Segment(0.0, 1.0, "A"),
                Segment(1.0, 2.0, "B"),
                Segment(2.0, 3.0, "A"),
            ]
        )

        a_segs = segments.filter_by_label("A")
        assert len(a_segs) == 2

    def test_sorted(self):
        """Test sorting."""
        segments = SegmentList(
            [
                Segment(2.0, 3.0, "C"),
                Segment(0.0, 1.0, "A"),
                Segment(1.0, 2.0, "B"),
            ]
        )

        sorted_segs = segments.sorted()
        assert sorted_segs[0].label == "A"
        assert sorted_segs[1].label == "B"
        assert sorted_segs[2].label == "C"


class TestSample:
    """Tests for Sample."""

    def test_creation(self):
        """Test basic Sample creation."""
        sample = Sample(
            id="test_001",
            text="Hello world",
            glosses=["HELLO", "WORLD"],
            signer_id="signer_01",
            dataset="test",
            split="train",
        )

        assert sample.id == "test_001"
        assert sample.has_text()
        assert not sample.has_poses()
        assert not sample.has_video()

    def test_with_poses(self):
        """Test Sample with pose data."""
        data = np.random.randn(100, 21, 3).astype(np.float32)
        poses = PoseSequence(data=data, fps=25.0, format="test")

        sample = Sample(id="test_001", poses=poses)

        assert sample.has_poses()
        assert sample.duration == 4.0
