"""Unit tests for sentiment macro-F1, micro-F1, direction accuracy, and calibration metrics.

Validates: Requirements 16.3, 16.4
"""
from __future__ import annotations

from services.intelligence_pipeline_v3.evaluation.sentiment_metrics import (
    CalibrationResult,
    DirectionAccuracyResult,
    SentimentEvaluationReport,
    SentimentF1Result,
    SentimentLabel,
    SentimentPrediction,
    compute_calibration,
    compute_direction_accuracy,
    compute_sentiment_f1,
    evaluate_sentiment,
)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _pred(
    company_entity_id: str,
    label: str,
    pos: float = 0.0,
    neg: float = 0.0,
    neu: float = 0.0,
    mix: float = 0.0,
) -> SentimentPrediction:
    return SentimentPrediction(
        company_entity_id=company_entity_id,
        label=SentimentLabel(label),
        positive_prob=pos,
        negative_prob=neg,
        neutral_prob=neu,
        mixed_prob=mix,
    )


# ---------------------------------------------------------------------------
# Sentiment F1 Metrics
# ---------------------------------------------------------------------------


class TestSentimentF1:
    def test_perfect_match(self) -> None:
        gold = [
            _pred("c1", "positive", pos=0.9, neg=0.05, neu=0.05),
            _pred("c2", "negative", pos=0.1, neg=0.8, neu=0.1),
            _pred("c3", "neutral", pos=0.1, neg=0.1, neu=0.8),
        ]
        pred = [
            _pred("c1", "positive", pos=0.85, neg=0.1, neu=0.05),
            _pred("c2", "negative", pos=0.05, neg=0.9, neu=0.05),
            _pred("c3", "neutral", pos=0.05, neg=0.05, neu=0.9),
        ]
        result = compute_sentiment_f1(pred, gold)
        assert result.macro_f1 == 1.0
        assert result.micro_f1 == 1.0
        assert result.support == 3

    def test_all_wrong(self) -> None:
        gold = [
            _pred("c1", "positive"),
            _pred("c2", "negative"),
            _pred("c3", "neutral"),
        ]
        pred = [
            _pred("c1", "negative"),
            _pred("c2", "neutral"),
            _pred("c3", "positive"),
        ]
        result = compute_sentiment_f1(pred, gold)
        assert result.macro_f1 == 0.0
        assert result.micro_f1 == 0.0

    def test_partial_match(self) -> None:
        gold = [
            _pred("c1", "positive"),
            _pred("c2", "positive"),
            _pred("c3", "negative"),
            _pred("c4", "neutral"),
        ]
        pred = [
            _pred("c1", "positive"),  # correct
            _pred("c2", "negative"),  # wrong
            _pred("c3", "negative"),  # correct
            _pred("c4", "neutral"),   # correct
        ]
        result = compute_sentiment_f1(pred, gold)
        # micro: overall accuracy across all label comparisons
        # TP: c1=pos correct, c3=neg correct, c4=neu correct = 3
        # Total predictions that match across all labels = 3
        assert result.micro_f1 == 0.75
        assert result.support == 4

    def test_unmatched_predictions_ignored(self) -> None:
        gold = [_pred("c1", "positive")]
        pred = [
            _pred("c1", "positive"),
            _pred("c_unknown", "negative"),  # no match in gold
        ]
        result = compute_sentiment_f1(pred, gold)
        assert result.macro_f1 == 1.0
        assert result.support == 1

    def test_empty_inputs(self) -> None:
        result = compute_sentiment_f1([], [])
        assert result.macro_f1 == 1.0
        assert result.micro_f1 == 1.0
        assert result.support == 0

    def test_per_label_breakdown(self) -> None:
        gold = [
            _pred("c1", "positive"),
            _pred("c2", "positive"),
            _pred("c3", "negative"),
        ]
        pred = [
            _pred("c1", "positive"),  # TP for positive
            _pred("c2", "neutral"),   # FN for positive, FP for neutral
            _pred("c3", "negative"),  # TP for negative
        ]
        result = compute_sentiment_f1(pred, gold)

        # Positive: TP=1, FP=0, FN=1 -> P=1.0, R=0.5, F1=2/3
        assert result.per_label["positive"].precision == 1.0
        assert result.per_label["positive"].recall == 0.5
        assert abs(result.per_label["positive"].f1 - 2 / 3) < 1e-9

        # Negative: TP=1, FP=0, FN=0 -> P=1.0, R=1.0, F1=1.0
        assert result.per_label["negative"].f1 == 1.0

        # Neutral: TP=0, FP=1, FN=0 -> P=0.0, R=1.0, F1=0.0
        assert result.per_label["neutral"].precision == 0.0
        assert result.per_label["neutral"].recall == 1.0
        assert result.per_label["neutral"].f1 == 0.0

    def test_mixed_label_support(self) -> None:
        gold = [_pred("c1", "mixed")]
        pred = [_pred("c1", "mixed", mix=0.7)]
        result = compute_sentiment_f1(pred, gold)
        assert result.per_label["mixed"].f1 == 1.0
        assert result.per_label["mixed"].support_gold == 1


# ---------------------------------------------------------------------------
# Direction Accuracy
# ---------------------------------------------------------------------------


class TestDirectionAccuracy:
    def test_all_correct(self) -> None:
        gold = [
            _pred("c1", "positive"),
            _pred("c2", "negative"),
        ]
        pred = [
            _pred("c1", "positive"),
            _pred("c2", "negative"),
        ]
        result = compute_direction_accuracy(pred, gold)
        assert result.accuracy == 1.0
        assert result.correct == 2
        assert result.total == 2

    def test_all_wrong(self) -> None:
        gold = [
            _pred("c1", "positive"),
            _pred("c2", "negative"),
        ]
        pred = [
            _pred("c1", "negative"),
            _pred("c2", "positive"),
        ]
        result = compute_direction_accuracy(pred, gold)
        assert result.accuracy == 0.0
        assert result.correct == 0
        assert result.total == 2

    def test_neutral_ignored(self) -> None:
        gold = [
            _pred("c1", "positive"),
            _pred("c2", "neutral"),
            _pred("c3", "negative"),
        ]
        pred = [
            _pred("c1", "positive"),
            _pred("c2", "positive"),  # gold is neutral, ignored
            _pred("c3", "negative"),
        ]
        result = compute_direction_accuracy(pred, gold)
        assert result.accuracy == 1.0
        assert result.total == 2  # c2 excluded

    def test_mixed_ignored(self) -> None:
        gold = [
            _pred("c1", "positive"),
            _pred("c2", "mixed"),
        ]
        pred = [
            _pred("c1", "positive"),
            _pred("c2", "negative"),  # gold is mixed, ignored
        ]
        result = compute_direction_accuracy(pred, gold)
        assert result.accuracy == 1.0
        assert result.total == 1

    def test_pred_neutral_ignored(self) -> None:
        """If predicted is neutral but gold is positive, pair is excluded."""
        gold = [_pred("c1", "positive")]
        pred = [_pred("c1", "neutral")]
        result = compute_direction_accuracy(pred, gold)
        assert result.total == 0
        assert result.accuracy == 1.0  # vacuously true

    def test_empty_inputs(self) -> None:
        result = compute_direction_accuracy([], [])
        assert result.accuracy == 1.0
        assert result.total == 0

    def test_unmatched_ignored(self) -> None:
        gold = [_pred("c1", "positive")]
        pred = [_pred("c_other", "negative")]
        result = compute_direction_accuracy(pred, gold)
        assert result.total == 0


# ---------------------------------------------------------------------------
# Calibration Metrics
# ---------------------------------------------------------------------------


class TestCalibration:
    def test_perfect_calibration(self) -> None:
        """When confidence exactly matches accuracy, ECE should be 0."""
        # All predictions are correct with confidence 1.0
        gold = [
            _pred("c1", "positive"),
            _pred("c2", "negative"),
        ]
        pred = [
            _pred("c1", "positive", pos=1.0, neg=0.0, neu=0.0),
            _pred("c2", "negative", pos=0.0, neg=1.0, neu=0.0),
        ]
        result = compute_calibration(pred, gold, n_bins=10)
        assert result.ece == 0.0
        assert result.n_samples == 2

    def test_brier_score_perfect(self) -> None:
        """Perfect predictions should have Brier score of 0."""
        gold = [_pred("c1", "positive")]
        pred = [_pred("c1", "positive", pos=1.0, neg=0.0, neu=0.0, mix=0.0)]
        result = compute_calibration(pred, gold, n_bins=10)
        assert result.brier_score == 0.0

    def test_brier_score_worst_case(self) -> None:
        """Completely wrong confidence should have high Brier score."""
        gold = [_pred("c1", "positive")]
        # Predicted negative with full confidence, gold is positive
        pred = [_pred("c1", "negative", pos=0.0, neg=1.0, neu=0.0, mix=0.0)]
        result = compute_calibration(pred, gold, n_bins=10)
        # Brier: (0-1)^2 + (1-0)^2 + (0-0)^2 + (0-0)^2 = 2.0
        assert abs(result.brier_score - 2.0) < 1e-9

    def test_brier_score_uniform_probs(self) -> None:
        """Uniform probabilities across 4 labels."""
        gold = [_pred("c1", "positive")]
        pred = [_pred("c1", "positive", pos=0.25, neg=0.25, neu=0.25, mix=0.25)]
        result = compute_calibration(pred, gold, n_bins=10)
        # Brier: (0.25-1)^2 + (0.25-0)^2 + (0.25-0)^2 + (0.25-0)^2
        # = 0.5625 + 0.0625 + 0.0625 + 0.0625 = 0.75
        assert abs(result.brier_score - 0.75) < 1e-9

    def test_ece_with_overconfidence(self) -> None:
        """High confidence but wrong predictions -> high ECE."""
        gold = [
            _pred("c1", "positive"),
            _pred("c2", "positive"),
        ]
        pred = [
            # Predicts negative with 0.9 confidence, wrong
            _pred("c1", "negative", pos=0.05, neg=0.9, neu=0.05, mix=0.0),
            # Predicts negative with 0.9 confidence, wrong
            _pred("c2", "negative", pos=0.05, neg=0.9, neu=0.05, mix=0.0),
        ]
        result = compute_calibration(pred, gold, n_bins=10)
        # Both have confidence 0.9, both wrong -> fraction_positive=0.0
        # ECE = |0.9 - 0.0| = 0.9
        assert abs(result.ece - 0.9) < 1e-9

    def test_reliability_bins_structure(self) -> None:
        """Reliability bins should cover [0, 1] range."""
        gold = [_pred("c1", "positive")]
        pred = [_pred("c1", "positive", pos=0.7, neg=0.2, neu=0.1)]
        result = compute_calibration(pred, gold, n_bins=5)
        assert len(result.reliability_bins) == 5
        assert result.reliability_bins[0].bin_lower == 0.0
        assert result.reliability_bins[-1].bin_upper == 1.0

    def test_empty_inputs(self) -> None:
        result = compute_calibration([], [], n_bins=10)
        assert result.ece == 0.0
        assert result.brier_score == 0.0
        assert result.n_samples == 0
        assert result.reliability_bins == []

    def test_unmatched_predictions_ignored(self) -> None:
        gold = [_pred("c1", "positive")]
        pred = [_pred("c_other", "positive", pos=0.9)]
        result = compute_calibration(pred, gold, n_bins=10)
        assert result.n_samples == 0

    def test_single_bin(self) -> None:
        """Single bin should contain all samples."""
        gold = [
            _pred("c1", "positive"),
            _pred("c2", "negative"),
        ]
        pred = [
            _pred("c1", "positive", pos=0.8, neg=0.1, neu=0.1),
            _pred("c2", "negative", pos=0.1, neg=0.7, neu=0.2),
        ]
        result = compute_calibration(pred, gold, n_bins=1)
        assert len(result.reliability_bins) == 1
        assert result.reliability_bins[0].count == 2

    def test_calibration_bins_count_sum(self) -> None:
        """Total count across bins should equal n_samples."""
        gold = [
            _pred("c1", "positive"),
            _pred("c2", "negative"),
            _pred("c3", "neutral"),
        ]
        pred = [
            _pred("c1", "positive", pos=0.7, neg=0.2, neu=0.1),
            _pred("c2", "negative", pos=0.1, neg=0.6, neu=0.3),
            _pred("c3", "neutral", pos=0.1, neg=0.1, neu=0.8),
        ]
        result = compute_calibration(pred, gold, n_bins=10)
        total_count = sum(b.count for b in result.reliability_bins)
        assert total_count == result.n_samples


# ---------------------------------------------------------------------------
# Full Evaluation Report
# ---------------------------------------------------------------------------


class TestEvaluateSentiment:
    def test_full_evaluation(self) -> None:
        gold = [
            _pred("c1", "positive", pos=0.9, neg=0.05, neu=0.05),
            _pred("c2", "negative", pos=0.1, neg=0.8, neu=0.1),
            _pred("c3", "neutral", pos=0.1, neg=0.1, neu=0.8),
        ]
        pred = [
            _pred("c1", "positive", pos=0.85, neg=0.1, neu=0.05),
            _pred("c2", "negative", pos=0.05, neg=0.9, neu=0.05),
            _pred("c3", "neutral", pos=0.05, neg=0.05, neu=0.9),
        ]
        report = evaluate_sentiment(pred, gold, n_bins=10, document_count=3)

        assert isinstance(report, SentimentEvaluationReport)
        assert isinstance(report.f1_metrics, SentimentF1Result)
        assert isinstance(report.direction_accuracy, DirectionAccuracyResult)
        assert isinstance(report.calibration, CalibrationResult)
        assert report.document_count == 3
        assert report.f1_metrics.macro_f1 == 1.0
        assert report.direction_accuracy.accuracy == 1.0

    def test_report_with_errors(self) -> None:
        gold = [
            _pred("c1", "positive"),
            _pred("c2", "negative"),
        ]
        pred = [
            _pred("c1", "negative", pos=0.1, neg=0.8, neu=0.1),
            _pred("c2", "negative", pos=0.1, neg=0.8, neu=0.1),
        ]
        report = evaluate_sentiment(pred, gold, document_count=2)

        # c1 wrong direction, c2 correct
        assert report.direction_accuracy.accuracy == 0.5
        assert report.direction_accuracy.total == 2
        assert report.f1_metrics.support == 2


# ---------------------------------------------------------------------------
# Edge Cases
# ---------------------------------------------------------------------------


class TestEdgeCases:
    def test_duplicate_company_ids_uses_last_gold(self) -> None:
        """When gold has duplicate IDs, dict lookup uses last occurrence."""
        gold = [
            _pred("c1", "positive"),
            _pred("c1", "negative"),  # overwrites first
        ]
        pred = [_pred("c1", "negative")]
        result = compute_sentiment_f1(pred, gold)
        # Gold dict will have c1 -> negative (last wins)
        assert result.per_label["negative"].f1 == 1.0

    def test_all_same_label(self) -> None:
        """All predictions and gold are the same label."""
        gold = [_pred(f"c{i}", "positive") for i in range(5)]
        pred = [_pred(f"c{i}", "positive", pos=0.9) for i in range(5)]
        result = compute_sentiment_f1(pred, gold)
        assert result.per_label["positive"].f1 == 1.0
        assert result.macro_f1 == 1.0

    def test_calibration_boundary_confidence(self) -> None:
        """Confidence of exactly 1.0 should be in the last bin."""
        gold = [_pred("c1", "positive")]
        pred = [_pred("c1", "positive", pos=1.0, neg=0.0, neu=0.0, mix=0.0)]
        result = compute_calibration(pred, gold, n_bins=10)
        # Last bin [0.9, 1.0] should have count 1
        assert result.reliability_bins[-1].count == 1

    def test_calibration_zero_confidence(self) -> None:
        """Confidence of 0.0 should be in the first bin."""
        gold = [_pred("c1", "positive")]
        # Label is positive but prob is 0.0 (inconsistent but valid input)
        pred = [_pred("c1", "positive", pos=0.0, neg=0.0, neu=0.0, mix=0.0)]
        result = compute_calibration(pred, gold, n_bins=10)
        # First bin [0.0, 0.1) should have count 1
        assert result.reliability_bins[0].count == 1