"""Tests for research utilities."""

import json

import numpy as np
import pytest

from sltk.data.core import Segment, SegmentList
from sltk.exceptions import ConfigurationError
from sltk.utils.research import (
    DatasetStats,
    compare_vocabularies,
    compute_pose_coverage,
    create_vocabulary_mapping,
    export_segments_for_training,
    find_gloss_examples,
)


class TestDatasetStats:
    """Tests for DatasetStats dataclass."""

    def test_creation(self):
        """Test basic DatasetStats creation."""
        stats = DatasetStats(
            name="test",
            num_samples=10,
            vocabulary_size=100,
        )
        assert stats.name == "test"
        assert stats.num_samples == 10
        assert stats.vocabulary_size == 100

    def test_to_dict(self):
        """Test conversion to dictionary."""
        stats = DatasetStats(
            name="test",
            num_samples=10,
            gloss_frequencies={"HELLO": 50},
        )
        d = stats.to_dict()
        assert d["name"] == "test"
        assert d["num_samples"] == 10
        assert d["gloss_frequencies"]["HELLO"] == 50

    def test_json_roundtrip(self, tmp_path):
        """Test JSON save and load."""
        stats = DatasetStats(
            name="test",
            num_samples=10,
            vocabulary_size=100,
            gloss_frequencies={"HELLO": 50, "WORLD": 30},
        )

        json_path = tmp_path / "stats.json"
        stats.to_json(json_path)

        loaded = DatasetStats.from_json(json_path)
        assert loaded.name == stats.name
        assert loaded.num_samples == stats.num_samples
        assert loaded.vocabulary_size == stats.vocabulary_size
        assert loaded.gloss_frequencies == stats.gloss_frequencies


class TestCompareVocabularies:
    """Tests for vocabulary comparison."""

    def test_identical_vocabularies(self):
        """Test comparison of identical vocabularies."""
        stats1 = DatasetStats(
            name="A",
            gloss_frequencies={"HELLO": 50, "WORLD": 30},
        )
        stats2 = DatasetStats(
            name="B",
            gloss_frequencies={"HELLO": 40, "WORLD": 35},
        )

        result = compare_vocabularies(stats1, stats2)

        assert result["intersection_size"] == 2
        assert result["jaccard_index"] == 1.0
        assert result["overlap_ratio_1"] == 1.0
        assert result["overlap_ratio_2"] == 1.0

    def test_disjoint_vocabularies(self):
        """Test comparison of disjoint vocabularies."""
        stats1 = DatasetStats(
            name="A",
            gloss_frequencies={"HELLO": 50, "WORLD": 30},
        )
        stats2 = DatasetStats(
            name="B",
            gloss_frequencies={"GOODBYE": 40, "THANKS": 35},
        )

        result = compare_vocabularies(stats1, stats2)

        assert result["intersection_size"] == 0
        assert result["jaccard_index"] == 0.0
        assert result["only_in_1_count"] == 2
        assert result["only_in_2_count"] == 2

    def test_partial_overlap(self):
        """Test comparison with partial overlap."""
        stats1 = DatasetStats(
            name="A",
            gloss_frequencies={"HELLO": 50, "WORLD": 30, "FOO": 10},
        )
        stats2 = DatasetStats(
            name="B",
            gloss_frequencies={"HELLO": 40, "BAR": 35},
        )

        result = compare_vocabularies(stats1, stats2)

        assert result["intersection_size"] == 1
        assert result["vocab_size_1"] == 3
        assert result["vocab_size_2"] == 2
        assert "HELLO" in result["common_glosses"]


class TestCreateVocabularyMapping:
    """Tests for vocabulary mapping creation."""

    def test_with_elan_files(self, tmp_path):
        """Test vocabulary mapping from ELAN files."""
        from sltk.io.elan import write_eaf

        # Create test ELAN files
        segments1 = SegmentList(
            [
                Segment(0, 1, "HELLO"),
                Segment(1, 2, "WORLD"),
                Segment(2, 3, "HELLO"),  # Duplicate
            ]
        )
        segments2 = SegmentList(
            [
                Segment(0, 1, "HELLO"),
                Segment(1, 2, "FOO"),
            ]
        )

        eaf1 = tmp_path / "test1.eaf"
        eaf2 = tmp_path / "test2.eaf"
        write_eaf(segments1, eaf1)
        write_eaf(segments2, eaf2)

        mapping = create_vocabulary_mapping(
            elan_files=[eaf1, eaf2],
            min_frequency=1,
            special_tokens=["<PAD>", "<UNK>"],
        )

        assert mapping["vocab_size"] >= 5  # 2 special + 3 glosses
        assert "<PAD>" in mapping["gloss_to_id"]
        assert "<UNK>" in mapping["gloss_to_id"]
        assert "HELLO" in mapping["gloss_to_id"]
        assert mapping["frequencies"]["HELLO"] == 3

    def test_min_frequency_filter(self, tmp_path):
        """Test minimum frequency filtering."""
        from sltk.io.elan import write_eaf

        # Create test file with varying frequencies
        segments = SegmentList(
            [
                Segment(0, 1, "COMMON"),
                Segment(1, 2, "COMMON"),
                Segment(2, 3, "COMMON"),
                Segment(3, 4, "RARE"),  # Only once
            ]
        )

        eaf = tmp_path / "test.eaf"
        write_eaf(segments, eaf)

        mapping = create_vocabulary_mapping(
            elan_files=[eaf],
            min_frequency=2,  # Filter out RARE
            special_tokens=None,
        )

        assert "COMMON" in mapping["gloss_to_id"]
        assert "RARE" not in mapping["gloss_to_id"]


class TestExportSegmentsForTraining:
    """Tests for training data export."""

    def test_csv_export(self, tmp_path):
        """Test CSV export."""
        import csv

        from sltk.io.elan import write_eaf

        segments = SegmentList(
            [
                Segment(0, 1, "HELLO"),
                Segment(1, 2, "WORLD"),
            ]
        )

        eaf = tmp_path / "test.eaf"
        write_eaf(segments, eaf)

        output = tmp_path / "output.csv"
        export_segments_for_training(
            elan_files=[eaf],
            output_path=output,
            format="csv",
            include_context=False,
        )

        assert output.exists()
        with open(output) as f:
            reader = csv.DictReader(f)
            rows = list(reader)

        assert len(rows) == 2
        assert rows[0]["label"] == "HELLO"
        assert rows[1]["label"] == "WORLD"

    def test_json_export(self, tmp_path):
        """Test JSON export."""
        from sltk.io.elan import write_eaf

        segments = SegmentList(
            [
                Segment(0, 1, "HELLO"),
            ]
        )

        eaf = tmp_path / "test.eaf"
        write_eaf(segments, eaf)

        output = tmp_path / "output.json"
        export_segments_for_training(
            elan_files=[eaf],
            output_path=output,
            format="json",
        )

        with open(output) as f:
            data = json.load(f)

        assert len(data) == 1
        assert data[0]["label"] == "HELLO"

    def test_context_inclusion(self, tmp_path):
        """Test context window inclusion."""
        import csv

        from sltk.io.elan import write_eaf

        segments = SegmentList(
            [
                Segment(0, 1, "A"),
                Segment(1, 2, "B"),
                Segment(2, 3, "C"),
            ]
        )

        eaf = tmp_path / "test.eaf"
        write_eaf(segments, eaf)

        output = tmp_path / "output.csv"
        export_segments_for_training(
            elan_files=[eaf],
            output_path=output,
            format="csv",
            include_context=True,
            context_window=1,
        )

        with open(output) as f:
            reader = csv.DictReader(f)
            rows = list(reader)

        # Middle segment should have context
        middle = rows[1]
        assert middle["label"] == "B"
        assert "A" in middle["context_before"]
        assert "C" in middle["context_after"]


class TestFindGlossExamples:
    """Tests for gloss example search."""

    def test_find_examples(self, tmp_path):
        """Test finding gloss examples."""
        from sltk.io.elan import write_eaf

        segments = SegmentList(
            [
                Segment(0, 1, "CONTEXT1"),
                Segment(1, 2, "TARGET"),
                Segment(2, 3, "CONTEXT2"),
                Segment(5, 6, "TARGET"),  # Second occurrence
            ]
        )

        eaf = tmp_path / "test.eaf"
        write_eaf(segments, eaf)

        examples = find_gloss_examples(
            gloss="TARGET",
            elan_files=[eaf],
            max_examples=10,
        )

        assert len(examples) == 2
        assert all(ex["gloss"].upper() == "TARGET" for ex in examples)

    def test_case_insensitive(self, tmp_path):
        """Test case-insensitive search."""
        from sltk.io.elan import write_eaf

        segments = SegmentList(
            [
                Segment(0, 1, "Hello"),
            ]
        )

        eaf = tmp_path / "test.eaf"
        write_eaf(segments, eaf)

        examples = find_gloss_examples(
            gloss="hello",  # lowercase
            elan_files=[eaf],
        )

        assert len(examples) == 1


class TestComputePoseCoverage:
    """Tests for pose coverage computation."""

    def test_full_coverage(self, tmp_path):
        """Test with full coverage (all frames valid)."""
        from sltk.data.core import PoseSequence
        from sltk.io.h5 import save_poses_h5

        poses = PoseSequence(
            data=np.ones((100, 21, 3), dtype=np.float32),  # All non-zero
            fps=25.0,
            format="test",
        )

        h5_path = tmp_path / "poses.h5"
        save_poses_h5(poses, h5_path)

        coverage = compute_pose_coverage(h5_path)

        assert coverage["total_frames"] == 100
        assert coverage["valid_frames"] == 100
        assert coverage["coverage_ratio"] == 1.0
        assert coverage["num_gaps"] == 0

    def test_partial_coverage(self, tmp_path):
        """Test with partial coverage (some zero frames)."""
        from sltk.data.core import PoseSequence
        from sltk.io.h5 import save_poses_h5

        data = np.ones((100, 21, 3), dtype=np.float32)
        data[20:30] = 0  # Gap of 10 frames

        poses = PoseSequence(
            data=data,
            fps=25.0,
            format="test",
        )

        h5_path = tmp_path / "poses.h5"
        save_poses_h5(poses, h5_path)

        coverage = compute_pose_coverage(h5_path)

        assert coverage["total_frames"] == 100
        assert coverage["valid_frames"] == 90
        assert coverage["coverage_ratio"] == 0.9
        assert coverage["num_gaps"] >= 1


class TestComputeDatasetStatistics:
    """Tests for compute_dataset_statistics function."""

    def test_basic_statistics(self, tmp_path):
        """Test basic statistics computation from ELAN files."""
        from sltk.io.elan import write_eaf
        from sltk.utils.research import compute_dataset_statistics

        # Create test ELAN files
        segments1 = SegmentList(
            [
                Segment(0, 1, "HELLO", tier="Gloss"),
                Segment(1, 2, "WORLD", tier="Gloss"),
            ]
        )
        segments2 = SegmentList(
            [
                Segment(0, 0.5, "HELLO", tier="Gloss"),
                Segment(1, 1.5, "FOO", tier="Gloss"),
            ]
        )

        eaf1 = tmp_path / "test1.eaf"
        eaf2 = tmp_path / "test2.eaf"
        write_eaf(segments1, eaf1)
        write_eaf(segments2, eaf2)

        stats = compute_dataset_statistics([eaf1, eaf2], name="test")

        assert stats.name == "test"
        assert stats.num_samples == 2
        assert stats.num_segments == 4
        assert stats.vocabulary_size == 3  # HELLO, WORLD, FOO
        assert "HELLO" in stats.gloss_frequencies

    def test_empty_files_list(self):
        """Test with empty files list."""
        from sltk.utils.research import compute_dataset_statistics

        stats = compute_dataset_statistics([], name="empty")

        assert stats.name == "empty"
        assert stats.num_samples == 0
        assert stats.num_segments == 0

    def test_tier_counts(self, tmp_path):
        """Test tier counting."""
        from sltk.io.elan import write_eaf
        from sltk.utils.research import compute_dataset_statistics

        segments = SegmentList(
            [
                Segment(0, 1, "A", tier="Tier1"),
                Segment(1, 2, "B", tier="Tier1"),
                Segment(2, 3, "C", tier="Tier2"),
            ]
        )

        eaf = tmp_path / "test.eaf"
        write_eaf(segments, eaf)

        stats = compute_dataset_statistics([eaf], name="test")

        assert "Tier1" in stats.tier_counts
        assert "Tier2" in stats.tier_counts
        assert stats.tier_counts["Tier1"] == 2
        assert stats.tier_counts["Tier2"] == 1

    def test_gloss_frequency_limit(self, tmp_path):
        """Test top_k_glosses parameter."""
        from sltk.io.elan import write_eaf
        from sltk.utils.research import compute_dataset_statistics

        # Create many different glosses
        segments = SegmentList(
            [Segment(i, i + 0.5, f"GLOSS_{i % 100}", tier="Gloss") for i in range(200)]
        )

        eaf = tmp_path / "test.eaf"
        write_eaf(segments, eaf)

        stats = compute_dataset_statistics([eaf], name="test", top_k_glosses=10)

        # Should only have top 10 glosses
        assert len(stats.gloss_frequencies) <= 10


class TestBatchExtractPoses:
    """Tests for batch_extract_poses function."""

    def test_skip_existing(self, tmp_path):
        """Test that existing files are skipped."""
        from sltk.utils.research import batch_extract_poses

        # Create some video files
        videos = [tmp_path / f"video_{i}.mp4" for i in range(3)]
        for v in videos:
            v.touch()

        # Create existing output file
        output_dir = tmp_path / "output"
        output_dir.mkdir()
        existing = output_dir / "video_0_mediapipe.h5"
        existing.touch()

        # batch_extract_poses requires actual extractor setup
        # We can at least test the interface
        try:
            result = batch_extract_poses(
                video_paths=videos,
                output_dir=output_dir,
                extractor="mediapipe",
                skip_existing=True,
            )
            # If successful, check skipped count
            assert "skipped" in result
        except (ImportError, RuntimeError):
            # Expected if mediapipe not installed
            pass

    def test_unknown_extractor_raises(self, tmp_path):
        """Test that unknown extractor raises ConfigurationError."""
        from sltk.utils.research import batch_extract_poses

        videos = [tmp_path / "video.mp4"]
        videos[0].touch()

        with pytest.raises(ConfigurationError, match="Unknown extractor"):
            batch_extract_poses(
                video_paths=videos,
                output_dir=tmp_path / "output",
                extractor="nonexistent_extractor",
            )


class TestAlignPosesToSegments:
    """Tests for align_poses_to_segments function."""

    def test_basic_alignment(self):
        """Test basic pose-segment alignment."""
        from sltk.data.core import PoseSequence
        from sltk.utils.research import align_poses_to_segments

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

        segments = SegmentList(
            [
                Segment(0.0, 1.0, "A"),  # frames 0-25
                Segment(2.0, 3.0, "B"),  # frames 50-75
            ]
        )

        clips = list(align_poses_to_segments(poses, segments))

        assert len(clips) == 2
        seg1, clip1 = clips[0]
        seg2, clip2 = clips[1]

        assert seg1.label == "A"
        assert clip1.shape[0] == 25  # 1 second at 25fps

        assert seg2.label == "B"
        assert clip2.shape[0] == 25

    def test_with_padding(self):
        """Test alignment with padding frames."""
        from sltk.data.core import PoseSequence
        from sltk.utils.research import align_poses_to_segments

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

        segments = SegmentList(
            [
                Segment(1.0, 2.0, "A"),  # frames 25-50
            ]
        )

        clips = list(align_poses_to_segments(poses, segments, padding_frames=5))

        assert len(clips) == 1
        seg, clip = clips[0]

        # Should have 25 frames + 5 padding on each side = 35 frames
        assert clip.shape[0] == 35

    def test_segment_at_boundary(self):
        """Test alignment when segment is at start boundary."""
        from sltk.data.core import PoseSequence
        from sltk.utils.research import align_poses_to_segments

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

        segments = SegmentList(
            [
                Segment(0.0, 0.5, "START"),  # frames 0-12
            ]
        )

        clips = list(align_poses_to_segments(poses, segments, padding_frames=5))

        assert len(clips) == 1
        seg, clip = clips[0]

        # Padding should be clipped to available frames
        assert clip.shape[0] <= 17  # 12 + 5 padding

    def test_empty_segments(self):
        """Test alignment with empty segment list."""
        from sltk.data.core import PoseSequence
        from sltk.utils.research import align_poses_to_segments

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

        segments = SegmentList([])

        clips = list(align_poses_to_segments(poses, segments))

        assert len(clips) == 0


class TestExportFormats:
    """Tests for different export formats."""

    def test_jsonl_export(self, tmp_path):
        """Test JSONL export format."""
        from sltk.io.elan import write_eaf

        segments = SegmentList(
            [
                Segment(0, 1, "A"),
                Segment(1, 2, "B"),
            ]
        )

        eaf = tmp_path / "test.eaf"
        write_eaf(segments, eaf)

        output = tmp_path / "output.jsonl"
        export_segments_for_training(
            elan_files=[eaf],
            output_path=output,
            format="jsonl",
        )

        assert output.exists()

        # Check JSONL format (one JSON object per line)
        with open(output) as f:
            lines = f.readlines()

        assert len(lines) == 2
        for line in lines:
            data = json.loads(line)
            assert "label" in data

    def test_invalid_format_raises(self, tmp_path):
        """Test that invalid format raises error."""
        from sltk.io.elan import write_eaf

        segments = SegmentList([Segment(0, 1, "A")])
        eaf = tmp_path / "test.eaf"
        write_eaf(segments, eaf)

        from sltk.exceptions import UnsupportedFormatError

        with pytest.raises(UnsupportedFormatError, match="invalid_format"):
            export_segments_for_training(
                elan_files=[eaf],
                output_path=tmp_path / "output.xyz",
                format="invalid_format",
            )


class TestDatasetStatsEdgeCases:
    """Edge case tests for DatasetStats."""

    def test_empty_gloss_frequencies(self):
        """Test DatasetStats with no gloss frequencies."""
        stats = DatasetStats(name="test")

        assert stats.gloss_frequencies == {}

    def test_all_fields(self):
        """Test DatasetStats with all fields populated."""
        stats = DatasetStats(
            name="complete",
            num_samples=100,
            num_segments=500,
            total_duration_sec=3600.0,
            vocabulary_size=200,
            avg_segment_duration=1.5,
            segment_duration_std=0.5,
            glosses_per_sample_avg=5.0,
            tier_counts={"Gloss": 400, "Translation": 100},
            gloss_frequencies={"HELLO": 50, "WORLD": 30},
            duration_histogram=[10, 20, 30],
            signers=["signer1", "signer2"],
            num_signers=2,
        )

        d = stats.to_dict()
        assert d["name"] == "complete"
        assert d["num_samples"] == 100
        assert len(d["signers"]) == 2


class TestVocabularyEdgeCases:
    """Edge case tests for vocabulary functions."""

    def test_empty_vocabulary_comparison(self):
        """Test comparing empty vocabularies."""
        stats1 = DatasetStats(name="A", gloss_frequencies={})
        stats2 = DatasetStats(name="B", gloss_frequencies={})

        result = compare_vocabularies(stats1, stats2)

        assert result["intersection_size"] == 0
        assert result["jaccard_index"] == 0  # 0/0 case

    def test_one_empty_vocabulary(self):
        """Test comparing when one vocabulary is empty."""
        stats1 = DatasetStats(name="A", gloss_frequencies={"HELLO": 1})
        stats2 = DatasetStats(name="B", gloss_frequencies={})

        result = compare_vocabularies(stats1, stats2)

        assert result["intersection_size"] == 0
        assert result["overlap_ratio_2"] == 0  # Empty vocab
