"""Tests for metrics module."""

import numpy as np
import pytest

from sltk.data.core import Segment, SegmentList
from sltk.metrics.segmentation import (
    SegmentationResult,
    boundary_f1,
    boundary_f1_multi_tolerance,
    comprehensive_segmentation_metrics,
    confusion_matrix,
    frame_accuracy,
    label_accuracy,
    over_segmentation_rate,
    per_class_metrics,
    per_segment_iou,
    segment_iou,
    segmentation_metrics,
    segmentation_summary,
    under_segmentation_rate,
)


class TestSegmentationMetrics:
    """Tests for segmentation metrics."""

    def test_perfect_segmentation(self):
        """Test metrics with perfect predictions."""
        segments = SegmentList(
            [
                Segment(start=0.0, end=1.0),
                Segment(start=2.0, end=3.0),
            ]
        )

        metrics = segmentation_metrics(segments, segments)

        assert metrics["boundary_precision"] == 1.0
        assert metrics["boundary_recall"] == 1.0
        assert metrics["boundary_f1"] == 1.0
        assert metrics["mean_iou"] == 1.0

    def test_no_overlap(self):
        """Test metrics with no overlap."""
        predictions = SegmentList([Segment(start=0.0, end=1.0)])
        references = SegmentList([Segment(start=2.0, end=3.0)])

        metrics = segmentation_metrics(predictions, references)

        assert metrics["mean_iou"] == 0.0

    def test_boundary_tolerance(self):
        """Test boundary metrics with tolerance."""
        predictions = SegmentList([Segment(start=0.0, end=1.05)])
        references = SegmentList([Segment(start=0.0, end=1.0)])

        # With default tolerance of 0.1, boundaries should match
        metrics = segmentation_metrics(predictions, references, tolerance_sec=0.1)

        assert metrics["boundary_precision"] == 1.0
        assert metrics["boundary_recall"] == 1.0

        # With stricter tolerance, boundaries should not match
        metrics_strict = segmentation_metrics(predictions, references, tolerance_sec=0.01)

        assert metrics_strict["boundary_precision"] < 1.0


class TestBoundaryF1:
    """Tests for boundary F1 score."""

    def test_perfect_boundaries(self):
        """Test F1 with perfect boundaries."""
        segments = SegmentList(
            [
                Segment(start=0.0, end=1.0),
                Segment(start=1.0, end=2.0),
            ]
        )

        f1 = boundary_f1(segments, segments)

        assert f1 == 1.0

    def test_partial_match(self):
        """Test F1 with partial boundary matches."""
        predictions = SegmentList(
            [
                Segment(start=0.0, end=1.0),
                Segment(start=1.0, end=2.0),
            ]
        )
        references = SegmentList(
            [
                Segment(start=0.0, end=1.5),
                Segment(start=1.5, end=2.0),
            ]
        )

        f1 = boundary_f1(predictions, references, tolerance_sec=0.1)

        # Some boundaries match (0.0, 2.0), some don't (1.0 vs 1.5)
        assert 0.0 < f1 < 1.0


class TestSegmentIoU:
    """Tests for segment IoU computation."""

    def test_full_overlap(self):
        """Test IoU with full overlap."""
        predictions = SegmentList([Segment(start=0.0, end=1.0)])
        references = SegmentList([Segment(start=0.0, end=1.0)])

        iou = segment_iou(predictions, references)

        assert iou == 1.0

    def test_partial_overlap(self):
        """Test IoU with partial overlap."""
        predictions = SegmentList([Segment(start=0.0, end=2.0)])
        references = SegmentList([Segment(start=1.0, end=3.0)])

        iou = segment_iou(predictions, references)

        # Intersection: 1.0-2.0 = 1.0
        # Union: 0.0-3.0 = 3.0
        # IoU = 1/3
        assert abs(iou - 1 / 3) < 0.01

    def test_multiple_segments(self):
        """Test IoU with multiple segments."""
        predictions = SegmentList(
            [
                Segment(start=0.0, end=1.0),
                Segment(start=2.0, end=3.0),
            ]
        )
        references = SegmentList(
            [
                Segment(start=0.0, end=1.0),  # Perfect match
                Segment(start=2.5, end=3.5),  # 50% overlap
            ]
        )

        iou = segment_iou(predictions, references)

        # Average of 1.0 and partial overlap
        assert 0.5 < iou < 1.0

    def test_empty_segments(self):
        """Test IoU with empty input."""
        empty = SegmentList([])
        non_empty = SegmentList([Segment(start=0.0, end=1.0)])

        assert segment_iou(empty, non_empty) == 0.0
        assert segment_iou(non_empty, empty) == 0.0
        assert segment_iou(empty, empty) == 0.0


class TestFrameAccuracy:
    """Tests for frame-level accuracy."""

    def test_perfect_accuracy(self):
        """Test accuracy with perfect predictions."""
        segments = SegmentList([Segment(start=0.0, end=1.0)])

        acc = frame_accuracy(segments, segments, fps=25.0, total_frames=50)

        assert acc == 1.0

    def test_zero_accuracy(self):
        """Test accuracy with inverted predictions."""
        predictions = SegmentList([Segment(start=0.0, end=1.0)])
        references = SegmentList([Segment(start=1.0, end=2.0)])

        acc = frame_accuracy(predictions, references, fps=25.0, total_frames=50)

        # No overlap: all frames are wrong
        assert acc == 0.0

    def test_half_accuracy(self):
        """Test accuracy with 50% overlap."""
        predictions = SegmentList([Segment(start=0.0, end=1.0)])  # Frames 0-24
        references = SegmentList([Segment(start=0.0, end=2.0)])  # Frames 0-49

        acc = frame_accuracy(predictions, references, fps=25.0, total_frames=50)

        # First 25 frames match (both in segment)
        # Last 25 frames: pred says no, ref says yes
        assert acc == 0.5


class TestTranslationMetrics:
    """Tests for translation metrics (basic tests without sacrebleu)."""

    def test_wer_perfect(self):
        """Test WER with perfect match."""
        from sltk.metrics.translation import wer

        predictions = ["hello world", "how are you"]
        references = ["hello world", "how are you"]

        error_rate = wer(predictions, references)

        assert error_rate == 0.0

    def test_wer_all_wrong(self):
        """Test WER with completely wrong predictions."""
        from sltk.metrics.translation import wer

        predictions = ["foo bar baz"]
        references = ["hello world"]

        error_rate = wer(predictions, references)

        # All words wrong
        assert error_rate > 0.0


# =============================================================================
# Comprehensive Segmentation Metrics Tests
# =============================================================================


class TestComprehensiveSegmentationMetrics:
    """Tests for comprehensive_segmentation_metrics function."""

    def test_perfect_match(self):
        """Test with perfect predictions."""
        segments = SegmentList(
            [
                Segment(start=0.0, end=1.0, label="A"),
                Segment(start=2.0, end=3.0, label="B"),
            ]
        )

        result = comprehensive_segmentation_metrics(segments, segments)

        assert isinstance(result, SegmentationResult)
        assert result.boundary_precision == 1.0
        assert result.boundary_recall == 1.0
        assert result.boundary_f1 == 1.0
        assert result.mean_iou == 1.0
        assert result.label_accuracy == 1.0
        assert result.num_matched == 2
        assert len(result.unmatched_predictions) == 0
        assert len(result.unmatched_references) == 0

    def test_no_overlap(self):
        """Test with no overlap between predictions and references."""
        predictions = SegmentList([Segment(start=0.0, end=1.0, label="A")])
        references = SegmentList([Segment(start=5.0, end=6.0, label="B")])

        result = comprehensive_segmentation_metrics(predictions, references)

        assert result.mean_iou == 0.0
        assert result.num_matched == 0
        assert len(result.unmatched_predictions) == 1
        assert len(result.unmatched_references) == 1

    def test_partial_label_match(self):
        """Test label accuracy with partial matches."""
        predictions = SegmentList(
            [
                Segment(start=0.0, end=1.0, label="A"),
                Segment(start=2.0, end=3.0, label="WRONG"),
            ]
        )
        references = SegmentList(
            [
                Segment(start=0.0, end=1.0, label="A"),
                Segment(start=2.0, end=3.0, label="B"),
            ]
        )

        result = comprehensive_segmentation_metrics(predictions, references, match_labels=False)

        # Both segments match by IoU, but only one has correct label
        assert result.num_matched == 2
        assert result.label_accuracy == 0.5

    def test_empty_segments(self):
        """Test with empty segment lists."""
        empty = SegmentList([])

        result = comprehensive_segmentation_metrics(empty, empty)

        # Perfect score for empty
        assert result.boundary_f1 == 1.0
        assert result.mean_iou == 1.0
        assert result.label_accuracy == 1.0

    def test_multi_tolerance_boundary_f1(self):
        """Test boundary F1 at multiple tolerances."""
        predictions = SegmentList([Segment(start=0.0, end=1.08)])
        references = SegmentList([Segment(start=0.0, end=1.0)])

        result = comprehensive_segmentation_metrics(predictions, references)

        # Boundary at 1.08 vs 1.0 - 80ms difference
        # Should match at 100ms and 200ms tolerance, not at 50ms
        assert result.boundary_f1_50ms < 1.0
        assert result.boundary_f1_100ms == 1.0
        assert result.boundary_f1_200ms == 1.0

    def test_to_dict(self):
        """Test conversion to dictionary."""
        segments = SegmentList([Segment(start=0.0, end=1.0, label="A")])

        result = comprehensive_segmentation_metrics(segments, segments)
        result_dict = result.to_dict()

        assert isinstance(result_dict, dict)
        assert "boundary_f1" in result_dict
        assert "mean_iou" in result_dict
        assert "per_class_f1" in result_dict

    def test_to_summary_dict(self):
        """Test conversion to summary dictionary."""
        segments = SegmentList([Segment(start=0.0, end=1.0, label="A")])

        result = comprehensive_segmentation_metrics(segments, segments)
        summary = result.to_summary_dict()

        assert isinstance(summary, dict)
        assert "boundary_f1" in summary
        assert "mean_iou" in summary
        # Summary should not have detailed metrics
        assert "per_class_f1" not in summary


class TestBoundaryF1MultiTolerance:
    """Tests for boundary_f1_multi_tolerance function."""

    def test_default_tolerances(self):
        """Test with default tolerance values."""
        segments = SegmentList([Segment(start=0.0, end=1.0)])

        scores = boundary_f1_multi_tolerance(segments, segments)

        assert "f1_50ms" in scores
        assert "f1_100ms" in scores
        assert "f1_200ms" in scores
        assert scores["f1_50ms"] == 1.0
        assert scores["f1_100ms"] == 1.0
        assert scores["f1_200ms"] == 1.0

    def test_custom_tolerances(self):
        """Test with custom tolerance values."""
        segments = SegmentList([Segment(start=0.0, end=1.0)])

        scores = boundary_f1_multi_tolerance(segments, segments, tolerances_sec=[0.01, 0.025, 0.5])

        assert "f1_10ms" in scores
        assert "f1_25ms" in scores
        assert "f1_500ms" in scores


class TestLabelAccuracy:
    """Tests for label_accuracy function."""

    def test_perfect_labels(self):
        """Test with all correct labels."""
        segments = SegmentList(
            [
                Segment(start=0.0, end=1.0, label="A"),
                Segment(start=2.0, end=3.0, label="B"),
            ]
        )

        acc = label_accuracy(segments, segments)

        assert acc == 1.0

    def test_wrong_labels(self):
        """Test with wrong labels."""
        predictions = SegmentList(
            [
                Segment(start=0.0, end=1.0, label="X"),
                Segment(start=2.0, end=3.0, label="Y"),
            ]
        )
        references = SegmentList(
            [
                Segment(start=0.0, end=1.0, label="A"),
                Segment(start=2.0, end=3.0, label="B"),
            ]
        )

        acc = label_accuracy(predictions, references)

        assert acc == 0.0

    def test_empty_input(self):
        """Test with empty input."""
        empty = SegmentList([])
        non_empty = SegmentList([Segment(start=0.0, end=1.0, label="A")])

        assert label_accuracy(empty, non_empty) == 0.0
        assert label_accuracy(non_empty, empty) == 0.0


class TestOverSegmentationRate:
    """Tests for over_segmentation_rate function."""

    def test_no_over_segmentation(self):
        """Test with no over-segmentation."""
        predictions = SegmentList([Segment(start=0.0, end=1.0)])
        references = SegmentList([Segment(start=0.0, end=1.0)])

        rate = over_segmentation_rate(predictions, references)

        assert rate == 0.0

    def test_over_segmentation(self):
        """Test with over-segmentation (multiple preds per ref)."""
        predictions = SegmentList(
            [
                Segment(start=0.0, end=0.5),
                Segment(start=0.5, end=1.0),
            ]
        )
        references = SegmentList([Segment(start=0.0, end=1.0)])

        rate = over_segmentation_rate(predictions, references, iou_threshold=0.1)

        # 2 predictions for 1 reference -> rate = 2 - 1 = 1.0
        assert rate == 1.0

    def test_empty_references(self):
        """Test with empty references."""
        predictions = SegmentList([Segment(start=0.0, end=1.0)])
        references = SegmentList([])

        rate = over_segmentation_rate(predictions, references)

        assert rate == 0.0


class TestUnderSegmentationRate:
    """Tests for under_segmentation_rate function."""

    def test_no_under_segmentation(self):
        """Test with no under-segmentation."""
        predictions = SegmentList([Segment(start=0.0, end=1.0)])
        references = SegmentList([Segment(start=0.0, end=1.0)])

        rate = under_segmentation_rate(predictions, references)

        assert rate == 0.0

    def test_under_segmentation(self):
        """Test with under-segmentation (one pred covers multiple refs)."""
        predictions = SegmentList([Segment(start=0.0, end=2.0)])
        references = SegmentList(
            [
                Segment(start=0.0, end=1.0),
                Segment(start=1.0, end=2.0),
            ]
        )

        rate = under_segmentation_rate(predictions, references, iou_threshold=0.1)

        # 1 prediction covers 2 references -> rate = 2 - 1 = 1.0
        assert rate == 1.0

    def test_empty_predictions(self):
        """Test with empty predictions."""
        predictions = SegmentList([])
        references = SegmentList([Segment(start=0.0, end=1.0)])

        rate = under_segmentation_rate(predictions, references)

        assert rate == 0.0


class TestPerSegmentIoU:
    """Tests for per_segment_iou function."""

    def test_all_matched(self):
        """Test with all segments matched."""
        segments = SegmentList(
            [
                Segment(start=0.0, end=1.0, label="A"),
                Segment(start=2.0, end=3.0, label="B"),
            ]
        )

        details = per_segment_iou(segments, segments)

        assert len(details) == 2
        assert all(d["matched"] for d in details)
        assert all(d["iou"] == 1.0 for d in details)

    def test_unmatched_segments(self):
        """Test with unmatched segments."""
        predictions = SegmentList([Segment(start=0.0, end=1.0, label="A")])
        references = SegmentList([Segment(start=5.0, end=6.0, label="B")])

        details = per_segment_iou(predictions, references)

        # Should have 2 entries: 1 unmatched pred, 1 unmatched ref
        assert len(details) == 2
        unmatched = [d for d in details if not d["matched"]]
        assert len(unmatched) == 2

    def test_boundary_errors(self):
        """Test boundary error calculation."""
        predictions = SegmentList([Segment(start=0.1, end=1.1)])
        references = SegmentList([Segment(start=0.0, end=1.0)])

        details = per_segment_iou(predictions, references)

        assert len(details) == 1
        assert details[0]["matched"]
        assert abs(details[0]["boundary_error_start"] - 0.1) < 0.001
        assert abs(details[0]["boundary_error_end"] - 0.1) < 0.001


class TestPerClassMetrics:
    """Tests for per_class_metrics function."""

    def test_single_class(self):
        """Test with single class."""
        segments = SegmentList(
            [
                Segment(start=0.0, end=1.0, label="A"),
                Segment(start=2.0, end=3.0, label="A"),
            ]
        )

        metrics = per_class_metrics(segments, segments)

        assert "A" in metrics
        assert metrics["A"]["precision"] == 1.0
        assert metrics["A"]["recall"] == 1.0
        assert metrics["A"]["f1"] == 1.0
        assert metrics["A"]["count"] == 2

    def test_multiple_classes(self):
        """Test with multiple classes."""
        segments = SegmentList(
            [
                Segment(start=0.0, end=1.0, label="A"),
                Segment(start=2.0, end=3.0, label="B"),
                Segment(start=4.0, end=5.0, label="C"),
            ]
        )

        metrics = per_class_metrics(segments, segments)

        assert len(metrics) == 3
        assert all(m["f1"] == 1.0 for m in metrics.values())

    def test_class_confusion(self):
        """Test metrics when classes are confused."""
        predictions = SegmentList(
            [
                Segment(start=0.0, end=1.0, label="A"),
                Segment(start=2.0, end=3.0, label="A"),  # Wrong, should be B
            ]
        )
        references = SegmentList(
            [
                Segment(start=0.0, end=1.0, label="A"),
                Segment(start=2.0, end=3.0, label="B"),
            ]
        )

        metrics = per_class_metrics(predictions, references)

        # Class A: 1 TP (correct pred), 1 FP (extra pred) -> precision = 0.5
        assert metrics["A"]["precision"] == 0.5
        # Class B: 0 TP, 1 FN (missed) -> recall = 0
        assert metrics["B"]["recall"] == 0.0


class TestConfusionMatrix:
    """Tests for confusion_matrix function."""

    def test_perfect_matrix(self):
        """Test confusion matrix with perfect predictions."""
        segments = SegmentList(
            [
                Segment(start=0.0, end=1.0, label="A"),
                Segment(start=2.0, end=3.0, label="B"),
            ]
        )

        cm = confusion_matrix(segments, segments)

        assert cm["A"]["A"] == 1
        assert cm["B"]["B"] == 1
        assert cm["A"]["B"] == 0
        assert cm["B"]["A"] == 0

    def test_confused_matrix(self):
        """Test confusion matrix with confused predictions."""
        predictions = SegmentList(
            [
                Segment(start=0.0, end=1.0, label="B"),  # Wrong
                Segment(start=2.0, end=3.0, label="A"),  # Wrong
            ]
        )
        references = SegmentList(
            [
                Segment(start=0.0, end=1.0, label="A"),
                Segment(start=2.0, end=3.0, label="B"),
            ]
        )

        cm = confusion_matrix(predictions, references)

        # Pred B matched with Ref A
        assert cm["B"]["A"] == 1
        # Pred A matched with Ref B
        assert cm["A"]["B"] == 1


class TestSegmentationSummary:
    """Tests for segmentation_summary function."""

    def test_summary_structure(self):
        """Test summary dictionary structure."""
        segments = SegmentList([Segment(start=0.0, end=1.0, label="A")])

        summary = segmentation_summary(segments, segments)

        assert "overview" in summary
        assert "boundary_metrics" in summary
        assert "iou_metrics" in summary
        assert "label_metrics" in summary
        assert "segmentation_quality" in summary
        assert "per_class" in summary

    def test_summary_with_fps(self):
        """Test summary with fps and duration."""
        segments = SegmentList([Segment(start=0.0, end=1.0, label="A")])

        summary = segmentation_summary(segments, segments, fps=25.0, total_duration=2.0)

        assert "frame_accuracy" in summary
        assert summary["frame_accuracy"] == 1.0

    def test_summary_overview_counts(self):
        """Test summary overview counts."""
        predictions = SegmentList(
            [
                Segment(start=0.0, end=1.0, label="A"),
                Segment(start=2.0, end=3.0, label="B"),
            ]
        )
        references = SegmentList(
            [
                Segment(start=0.0, end=1.0, label="A"),
            ]
        )

        summary = segmentation_summary(predictions, references)

        assert summary["overview"]["num_predictions"] == 2
        assert summary["overview"]["num_references"] == 1


class TestSegmentMatchingEdgeCases:
    """Tests for edge cases in segment matching."""

    def test_none_labels(self):
        """Test segments with None labels."""
        predictions = SegmentList([Segment(start=0.0, end=1.0, label=None)])
        references = SegmentList([Segment(start=0.0, end=1.0, label=None)])

        result = comprehensive_segmentation_metrics(predictions, references)

        assert result.num_matched == 1
        assert result.label_accuracy == 1.0

    def test_iou_threshold(self):
        """Test IoU threshold for matching."""
        predictions = SegmentList([Segment(start=0.0, end=1.0)])
        references = SegmentList([Segment(start=0.9, end=1.9)])

        # IoU = 0.1 / 1.9 ~= 0.053
        result_high = comprehensive_segmentation_metrics(predictions, references, iou_threshold=0.1)
        result_low = comprehensive_segmentation_metrics(predictions, references, iou_threshold=0.01)

        assert result_high.num_matched == 0
        assert result_low.num_matched == 1

    def test_greedy_matching(self):
        """Test greedy matching prefers highest IoU."""
        predictions = SegmentList(
            [
                Segment(start=0.0, end=1.0),
            ]
        )
        references = SegmentList(
            [
                Segment(start=0.0, end=1.0),  # Perfect match
                Segment(start=0.5, end=1.5),  # Partial match
            ]
        )

        result = comprehensive_segmentation_metrics(predictions, references, iou_threshold=0.3)

        # Should match with the perfect overlap
        assert result.num_matched == 1
        assert len(result.matches) == 1
        assert result.matches[0].iou == 1.0
