"""Unit tests for event and relation macro/micro F1 metrics.

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

from services.intelligence_pipeline_v3.evaluation.event_metrics import (
    EventMetricsResult,
    EventRelationEvaluationReport,
    GoldEvent,
    GoldRelation,
    PredictedEvent,
    PredictedRelation,
    RelationMetricsResult,
    compute_event_metrics,
    compute_relation_metrics,
    evaluate_events_and_relations,
)
from services.intelligence_pipeline_v3.schemas.annotations import (
    EventClass,
    RelationType,
)

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


def _pred_event(
    event_class: EventClass,
    evidence_ids: list[str] | None = None,
    primary_company_ids: list[str] | None = None,
    confidence: float = 1.0,
) -> PredictedEvent:
    return PredictedEvent(
        event_class=event_class,
        evidence_ids=evidence_ids or [],
        primary_company_ids=primary_company_ids or [],
        confidence=confidence,
    )


def _gold_event(
    event_class: EventClass,
    evidence_ids: list[str] | None = None,
    primary_company_ids: list[str] | None = None,
) -> GoldEvent:
    return GoldEvent(
        event_class=event_class,
        evidence_ids=evidence_ids or [],
        primary_company_ids=primary_company_ids or [],
    )


def _pred_relation(
    relation_type: RelationType,
    source_id: str,
    target_id: str,
    confidence: float = 1.0,
) -> PredictedRelation:
    return PredictedRelation(
        relation_type=relation_type,
        source_id=source_id,
        target_id=target_id,
        confidence=confidence,
    )


def _gold_relation(
    relation_type: RelationType,
    source_id: str,
    target_id: str,
) -> GoldRelation:
    return GoldRelation(
        relation_type=relation_type,
        source_id=source_id,
        target_id=target_id,
    )


# ---------------------------------------------------------------------------
# Event Metrics — Basic
# ---------------------------------------------------------------------------


class TestEventMetricsBasic:
    def test_both_empty(self) -> None:
        result = compute_event_metrics([], [])
        # All per-class are vacuously 1.0 (no predictions, no gold)
        assert result.micro.precision == 1.0
        assert result.micro.recall == 1.0
        assert result.micro.f1 == 1.0
        assert result.macro_f1 == 1.0

    def test_perfect_match_evidence(self) -> None:
        """Events with same class and overlapping evidence match."""
        pred = [_pred_event(EventClass.EARNINGS_BEAT, evidence_ids=["e1", "e2"])]
        gold = [_gold_event(EventClass.EARNINGS_BEAT, evidence_ids=["e2", "e3"])]
        result = compute_event_metrics(pred, gold)
        assert result.micro.precision == 1.0
        assert result.micro.recall == 1.0
        assert result.micro.f1 == 1.0

    def test_perfect_match_company(self) -> None:
        """Events with same class and overlapping primary company match."""
        pred = [_pred_event(EventClass.MA_ANNOUNCEMENT, primary_company_ids=["c1"])]
        gold = [_gold_event(EventClass.MA_ANNOUNCEMENT, primary_company_ids=["c1", "c2"])]
        result = compute_event_metrics(pred, gold)
        assert result.micro.precision == 1.0
        assert result.micro.recall == 1.0

    def test_no_predictions(self) -> None:
        gold = [_gold_event(EventClass.EARNINGS_BEAT, evidence_ids=["e1"])]
        result = compute_event_metrics([], gold)
        assert result.micro.recall == 0.0

    def test_no_gold(self) -> None:
        pred = [_pred_event(EventClass.EARNINGS_BEAT, evidence_ids=["e1"])]
        result = compute_event_metrics(pred, [])
        assert result.micro.precision == 0.0

    def test_wrong_class_no_match(self) -> None:
        """Different event_class means no match regardless of evidence."""
        pred = [_pred_event(EventClass.EARNINGS_BEAT, evidence_ids=["e1"])]
        gold = [_gold_event(EventClass.EARNINGS_MISS, evidence_ids=["e1"])]
        result = compute_event_metrics(pred, gold)
        assert result.micro.precision == 0.0
        assert result.micro.recall == 0.0

    def test_no_overlap_no_match(self) -> None:
        """Same class but no overlapping evidence or companies means no match."""
        pred = [_pred_event(EventClass.EARNINGS_BEAT, evidence_ids=["e1"], primary_company_ids=["c1"])]
        gold = [_gold_event(EventClass.EARNINGS_BEAT, evidence_ids=["e2"], primary_company_ids=["c2"])]
        result = compute_event_metrics(pred, gold)
        assert result.micro.precision == 0.0
        assert result.micro.recall == 0.0


# ---------------------------------------------------------------------------
# Event Metrics — Per-Class
# ---------------------------------------------------------------------------


class TestEventMetricsPerClass:
    def test_per_class_breakdown_all_13_classes(self) -> None:
        """Result always contains all 13 event classes."""
        result = compute_event_metrics([], [])
        assert len(result.per_class) == 13
        for ec in EventClass:
            assert ec.value in result.per_class

    def test_per_class_single_class(self) -> None:
        pred = [
            _pred_event(EventClass.PRODUCT_LAUNCH, evidence_ids=["e1"]),
            _pred_event(EventClass.PRODUCT_LAUNCH, evidence_ids=["e2"]),
        ]
        gold = [
            _gold_event(EventClass.PRODUCT_LAUNCH, evidence_ids=["e1"]),
            _gold_event(EventClass.PRODUCT_LAUNCH, evidence_ids=["e3"]),
        ]
        result = compute_event_metrics(pred, gold)
        pl = result.per_class["product_launch"]
        # 1 TP (e1 match), 1 FP, 1 FN
        assert pl.precision == 0.5
        assert pl.recall == 0.5
        assert abs(pl.f1 - 0.5) < 1e-9

    def test_per_class_mixed(self) -> None:
        pred = [
            _pred_event(EventClass.EARNINGS_BEAT, evidence_ids=["e1"]),
            _pred_event(EventClass.LEGAL_REGULATORY, primary_company_ids=["c1"]),
        ]
        gold = [
            _gold_event(EventClass.EARNINGS_BEAT, evidence_ids=["e1"]),
            _gold_event(EventClass.LEGAL_REGULATORY, primary_company_ids=["c1"]),
            _gold_event(EventClass.MACRO_EVENT, evidence_ids=["e5"]),
        ]
        result = compute_event_metrics(pred, gold)
        assert result.per_class["earnings_beat"].f1 == 1.0
        assert result.per_class["legal_regulatory"].f1 == 1.0
        assert result.per_class["macro_event"].recall == 0.0


# ---------------------------------------------------------------------------
# Event Metrics — Macro vs Micro
# ---------------------------------------------------------------------------


class TestEventMetricsMacroMicro:
    def test_macro_averages_across_classes(self) -> None:
        """Macro-F1 averages per-class F1, including classes with no data."""
        pred = [_pred_event(EventClass.EARNINGS_BEAT, evidence_ids=["e1"])]
        gold = [_gold_event(EventClass.EARNINGS_BEAT, evidence_ids=["e1"])]
        result = compute_event_metrics(pred, gold)
        # earnings_beat has F1=1.0, all other 12 classes have F1=1.0 (empty/empty)
        assert result.macro_f1 == 1.0

    def test_macro_penalizes_missing_class(self) -> None:
        """A class with only gold items drags macro-F1 down."""
        pred = [_pred_event(EventClass.EARNINGS_BEAT, evidence_ids=["e1"])]
        gold = [
            _gold_event(EventClass.EARNINGS_BEAT, evidence_ids=["e1"]),
            _gold_event(EventClass.EARNINGS_MISS, evidence_ids=["e2"]),
        ]
        result = compute_event_metrics(pred, gold)
        # earnings_beat: F1=1.0, earnings_miss: recall=0 -> F1=0, rest: F1=1.0
        # macro = (1.0 + 0.0 + 11*1.0) / 13 = 12/13
        assert abs(result.macro_f1 - 12 / 13) < 1e-9

    def test_micro_aggregates_tp_fp_fn(self) -> None:
        """Micro-F1 sums TP/FP/FN across all classes."""
        pred = [
            _pred_event(EventClass.EARNINGS_BEAT, evidence_ids=["e1"]),  # TP
            _pred_event(EventClass.RATING_CHANGE, evidence_ids=["e99"]),  # FP
        ]
        gold = [
            _gold_event(EventClass.EARNINGS_BEAT, evidence_ids=["e1"]),
            _gold_event(EventClass.MACRO_EVENT, evidence_ids=["e5"]),  # FN
        ]
        result = compute_event_metrics(pred, gold)
        # TP=1, FP=1, FN=1 -> P=1/2, R=1/2, F1=1/2
        assert result.micro.support_predicted == 2
        assert result.micro.support_gold == 2
        assert result.micro.precision == 0.5
        assert result.micro.recall == 0.5
        assert abs(result.micro.f1 - 0.5) < 1e-9


# ---------------------------------------------------------------------------
# Relation Metrics — Basic
# ---------------------------------------------------------------------------


class TestRelationMetricsBasic:
    def test_both_empty(self) -> None:
        result = compute_relation_metrics([], [])
        assert result.micro.f1 == 1.0
        assert result.macro_f1 == 1.0

    def test_perfect_match(self) -> None:
        pred = [_pred_relation(RelationType.DIRECTLY_AFFECTS, "ev1", "comp1")]
        gold = [_gold_relation(RelationType.DIRECTLY_AFFECTS, "ev1", "comp1")]
        result = compute_relation_metrics(pred, gold)
        assert result.micro.f1 == 1.0

    def test_wrong_type_no_match(self) -> None:
        pred = [_pred_relation(RelationType.DIRECTLY_AFFECTS, "ev1", "comp1")]
        gold = [_gold_relation(RelationType.INFERRED_EXPOSURE, "ev1", "comp1")]
        result = compute_relation_metrics(pred, gold)
        assert result.micro.precision == 0.0
        assert result.micro.recall == 0.0

    def test_wrong_source_no_match(self) -> None:
        pred = [_pred_relation(RelationType.COMPETES_WITH, "c1", "c2")]
        gold = [_gold_relation(RelationType.COMPETES_WITH, "c3", "c2")]
        result = compute_relation_metrics(pred, gold)
        assert result.micro.precision == 0.0

    def test_wrong_target_no_match(self) -> None:
        pred = [_pred_relation(RelationType.SUPPLIES, "c1", "c2")]
        gold = [_gold_relation(RelationType.SUPPLIES, "c1", "c3")]
        result = compute_relation_metrics(pred, gold)
        assert result.micro.precision == 0.0

    def test_no_predictions(self) -> None:
        gold = [_gold_relation(RelationType.COMPETES_WITH, "c1", "c2")]
        result = compute_relation_metrics([], gold)
        assert result.micro.recall == 0.0

    def test_no_gold(self) -> None:
        pred = [_pred_relation(RelationType.COMPETES_WITH, "c1", "c2")]
        result = compute_relation_metrics(pred, [])
        assert result.micro.precision == 0.0


# ---------------------------------------------------------------------------
# Relation Metrics — Per-Type
# ---------------------------------------------------------------------------


class TestRelationMetricsPerType:
    def test_per_type_breakdown_all_4_types(self) -> None:
        """Result always contains all 4 relation types."""
        result = compute_relation_metrics([], [])
        assert len(result.per_type) == 4
        for rt in RelationType:
            assert rt.value in result.per_type

    def test_per_type_mixed(self) -> None:
        pred = [
            _pred_relation(RelationType.DIRECTLY_AFFECTS, "ev1", "c1"),
            _pred_relation(RelationType.COMPETES_WITH, "c1", "c2"),
            _pred_relation(RelationType.COMPETES_WITH, "c3", "c4"),  # FP
        ]
        gold = [
            _gold_relation(RelationType.DIRECTLY_AFFECTS, "ev1", "c1"),
            _gold_relation(RelationType.COMPETES_WITH, "c1", "c2"),
            _gold_relation(RelationType.SUPPLIES, "c5", "c6"),  # FN
        ]
        result = compute_relation_metrics(pred, gold)
        assert result.per_type["directly_affects"].f1 == 1.0
        assert result.per_type["competes_with"].precision == 0.5
        assert result.per_type["competes_with"].recall == 1.0
        assert result.per_type["supplies"].recall == 0.0


# ---------------------------------------------------------------------------
# Relation Metrics — Macro vs Micro
# ---------------------------------------------------------------------------


class TestRelationMetricsMacroMicro:
    def test_macro_averages_across_types(self) -> None:
        pred = [_pred_relation(RelationType.DIRECTLY_AFFECTS, "ev1", "c1")]
        gold = [_gold_relation(RelationType.DIRECTLY_AFFECTS, "ev1", "c1")]
        result = compute_relation_metrics(pred, gold)
        # directly_affects: F1=1.0, other 3: F1=1.0 (empty)
        assert result.macro_f1 == 1.0

    def test_macro_penalizes_missing_type(self) -> None:
        pred = [_pred_relation(RelationType.DIRECTLY_AFFECTS, "ev1", "c1")]
        gold = [
            _gold_relation(RelationType.DIRECTLY_AFFECTS, "ev1", "c1"),
            _gold_relation(RelationType.SUPPLIES, "c5", "c6"),
        ]
        result = compute_relation_metrics(pred, gold)
        # directly_affects: F1=1.0, supplies: recall=0 -> F1=0, other 2: F1=1.0
        # macro = (1.0 + 0.0 + 1.0 + 1.0) / 4 = 3/4
        assert abs(result.macro_f1 - 0.75) < 1e-9

    def test_micro_aggregates(self) -> None:
        pred = [
            _pred_relation(RelationType.DIRECTLY_AFFECTS, "ev1", "c1"),  # TP
            _pred_relation(RelationType.COMPETES_WITH, "c1", "c99"),  # FP
        ]
        gold = [
            _gold_relation(RelationType.DIRECTLY_AFFECTS, "ev1", "c1"),
            _gold_relation(RelationType.INFERRED_EXPOSURE, "ev2", "c3"),  # FN
        ]
        result = compute_relation_metrics(pred, gold)
        # TP=1, FP=1, FN=1 -> P=1/2, R=1/2, F1=1/2
        assert result.micro.precision == 0.5
        assert result.micro.recall == 0.5
        assert abs(result.micro.f1 - 0.5) < 1e-9


# ---------------------------------------------------------------------------
# Combined Evaluation Report
# ---------------------------------------------------------------------------


class TestEvaluateEventsAndRelations:
    def test_full_report(self) -> None:
        pred_events = [_pred_event(EventClass.EARNINGS_BEAT, evidence_ids=["e1"])]
        gold_events = [_gold_event(EventClass.EARNINGS_BEAT, evidence_ids=["e1"])]
        pred_relations = [_pred_relation(RelationType.DIRECTLY_AFFECTS, "ev1", "c1")]
        gold_relations = [_gold_relation(RelationType.DIRECTLY_AFFECTS, "ev1", "c1")]

        report = evaluate_events_and_relations(
            pred_events, gold_events, pred_relations, gold_relations,
            document_count=5,
        )

        assert isinstance(report, EventRelationEvaluationReport)
        assert isinstance(report.event_metrics, EventMetricsResult)
        assert isinstance(report.relation_metrics, RelationMetricsResult)
        assert report.document_count == 5
        assert report.event_metrics.micro.f1 == 1.0
        assert report.relation_metrics.micro.f1 == 1.0

    def test_report_with_failures(self) -> None:
        pred_events = [
            _pred_event(EventClass.EARNINGS_BEAT, evidence_ids=["e1"]),
            _pred_event(EventClass.SUPPLY_CHAIN, evidence_ids=["e99"]),
        ]
        gold_events = [
            _gold_event(EventClass.EARNINGS_BEAT, evidence_ids=["e1"]),
            _gold_event(EventClass.MACRO_EVENT, primary_company_ids=["c7"]),
        ]
        pred_relations = []
        gold_relations = [_gold_relation(RelationType.SUPPLIES, "c1", "c2")]

        report = evaluate_events_and_relations(
            pred_events, gold_events, pred_relations, gold_relations,
            document_count=2,
        )

        assert report.event_metrics.micro.precision == 0.5
        assert report.event_metrics.micro.recall == 0.5
        assert report.relation_metrics.micro.recall == 0.0
        assert report.document_count == 2


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


class TestEdgeCases:
    def test_event_match_requires_both_class_and_overlap(self) -> None:
        """Same class but completely empty evidence and companies — no match."""
        pred = [_pred_event(EventClass.BUYBACK)]
        gold = [_gold_event(EventClass.BUYBACK)]
        result = compute_event_metrics(pred, gold)
        # No evidence or companies to overlap -> no match
        assert result.per_class["buyback"].precision == 0.0

    def test_multiple_events_greedy_matching(self) -> None:
        """Greedy matching: first match consumes the gold item."""
        pred = [
            _pred_event(EventClass.EARNINGS_BEAT, evidence_ids=["e1"]),
            _pred_event(EventClass.EARNINGS_BEAT, evidence_ids=["e1"]),
        ]
        gold = [_gold_event(EventClass.EARNINGS_BEAT, evidence_ids=["e1"])]
        result = compute_event_metrics(pred, gold)
        # 1 TP, 1 FP -> precision = 0.5, recall = 1.0
        assert result.per_class["earnings_beat"].precision == 0.5
        assert result.per_class["earnings_beat"].recall == 1.0

    def test_relation_duplicates(self) -> None:
        """Duplicate predictions can only match once."""
        pred = [
            _pred_relation(RelationType.COMPETES_WITH, "c1", "c2"),
            _pred_relation(RelationType.COMPETES_WITH, "c1", "c2"),
        ]
        gold = [_gold_relation(RelationType.COMPETES_WITH, "c1", "c2")]
        result = compute_relation_metrics(pred, gold)
        assert result.per_type["competes_with"].precision == 0.5
        assert result.per_type["competes_with"].recall == 1.0

    def test_event_confidence_does_not_affect_matching(self) -> None:
        """Confidence is stored but doesn't affect match logic."""
        pred = [_pred_event(EventClass.DIVIDEND_CHANGE, evidence_ids=["e1"], confidence=0.1)]
        gold = [_gold_event(EventClass.DIVIDEND_CHANGE, evidence_ids=["e1"])]
        result = compute_event_metrics(pred, gold)
        assert result.per_class["dividend_change"].f1 == 1.0