"""Event and relation macro/micro F1 evaluation metrics.

Implements evaluation metrics for event classification and relation extraction
quality against a gold standard corpus. Supports both macro-F1 (average across
classes) and micro-F1 (global TP/FP/FN) with per-class breakdowns.

Matching logic:
- Events match if they share the same event_class AND have overlapping evidence
  spans OR the same primary company.
- Relations match if they share the same relation_type, source_id, and target_id.

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

from pydantic import BaseModel, Field

from services.intelligence_pipeline_v3.schemas.annotations import (
    EventClass,
    RelationType,
)

# ---------------------------------------------------------------------------
# Input Models
# ---------------------------------------------------------------------------


class PredictedEvent(BaseModel):
    """A predicted event for evaluation."""

    event_class: EventClass
    evidence_ids: list[str] = Field(default_factory=list)
    primary_company_ids: list[str] = Field(default_factory=list)
    confidence: float = Field(ge=0.0, le=1.0, default=1.0)


class GoldEvent(BaseModel):
    """A gold standard event for evaluation."""

    event_class: EventClass
    evidence_ids: list[str] = Field(default_factory=list)
    primary_company_ids: list[str] = Field(default_factory=list)


class PredictedRelation(BaseModel):
    """A predicted relation for evaluation."""

    relation_type: RelationType
    source_id: str
    target_id: str
    confidence: float = Field(ge=0.0, le=1.0, default=1.0)


class GoldRelation(BaseModel):
    """A gold standard relation for evaluation."""

    relation_type: RelationType
    source_id: str
    target_id: str


# ---------------------------------------------------------------------------
# Result Models
# ---------------------------------------------------------------------------


class PRF1(BaseModel):
    """Precision, recall, F1 triple."""

    precision: float = Field(ge=0.0, le=1.0)
    recall: float = Field(ge=0.0, le=1.0)
    f1: float = Field(ge=0.0, le=1.0)
    support_predicted: int = Field(ge=0)
    support_gold: int = Field(ge=0)


class EventMetricsResult(BaseModel):
    """Full event evaluation result with macro/micro F1 and per-class breakdown."""

    macro_f1: float = Field(ge=0.0, le=1.0)
    micro: PRF1
    per_class: dict[str, PRF1]


class RelationMetricsResult(BaseModel):
    """Full relation evaluation result with macro/micro F1 and per-type breakdown."""

    macro_f1: float = Field(ge=0.0, le=1.0)
    micro: PRF1
    per_type: dict[str, PRF1]


class EventRelationEvaluationReport(BaseModel):
    """Complete event and relation evaluation report."""

    event_metrics: EventMetricsResult
    relation_metrics: RelationMetricsResult
    document_count: int = Field(ge=0)


# ---------------------------------------------------------------------------
# Matching Logic
# ---------------------------------------------------------------------------


def _events_match(pred: PredictedEvent, gold: GoldEvent) -> bool:
    """Events match if same event_class AND overlapping evidence OR same primary company.

    Overlap means at least one evidence_id in common, OR at least one
    primary_company_id in common.
    """
    if pred.event_class != gold.event_class:
        return False

    # Check overlapping evidence spans
    if pred.evidence_ids and gold.evidence_ids:
        if set(pred.evidence_ids) & set(gold.evidence_ids):
            return True

    # Check same primary company
    if pred.primary_company_ids and gold.primary_company_ids:
        if set(pred.primary_company_ids) & set(gold.primary_company_ids):
            return True

    return False


def _relations_match(pred: PredictedRelation, gold: GoldRelation) -> bool:
    """Relations match if same type, source, and target."""
    return (
        pred.relation_type == gold.relation_type
        and pred.source_id == gold.source_id
        and pred.target_id == gold.target_id
    )


# ---------------------------------------------------------------------------
# Core Metric Computation
# ---------------------------------------------------------------------------


def _compute_prf1_greedy(
    predicted: list,
    gold: list,
    match_fn: object,
) -> PRF1:
    """Compute precision, recall, F1 using greedy bipartite matching.

    Each predicted item can match at most one gold item and vice versa.
    """
    n_pred = len(predicted)
    n_gold = len(gold)

    if n_pred == 0 and n_gold == 0:
        return PRF1(
            precision=1.0, recall=1.0, f1=1.0,
            support_predicted=0, support_gold=0,
        )

    if n_pred == 0:
        return PRF1(
            precision=1.0, recall=0.0, f1=0.0,
            support_predicted=0, support_gold=n_gold,
        )

    if n_gold == 0:
        return PRF1(
            precision=0.0, recall=1.0, f1=0.0,
            support_predicted=n_pred, support_gold=0,
        )

    matched_gold: set[int] = set()
    true_positives = 0

    for p in predicted:
        for g_idx, g in enumerate(gold):
            if g_idx in matched_gold:
                continue
            if match_fn(p, g):  # type: ignore[operator]
                true_positives += 1
                matched_gold.add(g_idx)
                break

    precision = true_positives / n_pred if n_pred > 0 else 0.0
    recall = true_positives / n_gold if n_gold > 0 else 0.0

    if precision + recall > 0:
        f1 = 2 * precision * recall / (precision + recall)
    else:
        f1 = 0.0

    return PRF1(
        precision=precision,
        recall=recall,
        f1=f1,
        support_predicted=n_pred,
        support_gold=n_gold,
    )


def _compute_micro_prf1(
    predicted: list,
    gold: list,
    match_fn: object,
    class_key_pred: object,
    class_key_gold: object,
    all_classes: set[str],
) -> PRF1:
    """Compute micro-averaged PRF1 by summing TP/FP/FN across all classes."""
    total_tp = 0
    total_pred = 0
    total_gold = 0

    for cls in all_classes:
        cls_predicted = [p for p in predicted if class_key_pred(p) == cls]
        cls_gold = [g for g in gold if class_key_gold(g) == cls]

        total_pred += len(cls_predicted)
        total_gold += len(cls_gold)

        # Greedy match within this class
        matched_gold: set[int] = set()
        for p in cls_predicted:
            for g_idx, g in enumerate(cls_gold):
                if g_idx in matched_gold:
                    continue
                if match_fn(p, g):  # type: ignore[operator]
                    total_tp += 1
                    matched_gold.add(g_idx)
                    break

    if total_pred == 0 and total_gold == 0:
        return PRF1(
            precision=1.0, recall=1.0, f1=1.0,
            support_predicted=0, support_gold=0,
        )

    precision = total_tp / total_pred if total_pred > 0 else 0.0
    recall = total_tp / total_gold if total_gold > 0 else 0.0

    if precision + recall > 0:
        f1 = 2 * precision * recall / (precision + recall)
    else:
        f1 = 0.0

    return PRF1(
        precision=precision,
        recall=recall,
        f1=f1,
        support_predicted=total_pred,
        support_gold=total_gold,
    )


# ---------------------------------------------------------------------------
# Public API — Events
# ---------------------------------------------------------------------------


def compute_event_metrics(
    predicted: list[PredictedEvent],
    gold: list[GoldEvent],
) -> EventMetricsResult:
    """Compute event macro-F1, micro-F1, and per-class F1.

    Args:
        predicted: Predicted events.
        gold: Gold standard events.

    Returns:
        EventMetricsResult with macro, micro, and per-class breakdowns.
    """
    all_classes = {e.value for e in EventClass}

    # Per-class breakdown
    per_class: dict[str, PRF1] = {}
    f1_scores: list[float] = []

    for cls in sorted(all_classes):
        cls_predicted = [p for p in predicted if p.event_class.value == cls]
        cls_gold = [g for g in gold if g.event_class.value == cls]
        prf1 = _compute_prf1_greedy(cls_predicted, cls_gold, _events_match)
        per_class[cls] = prf1
        f1_scores.append(prf1.f1)

    # Macro-F1: average F1 across all event classes
    macro_f1 = sum(f1_scores) / len(f1_scores) if f1_scores else 0.0

    # Micro-F1: global TP/FP/FN
    micro = _compute_micro_prf1(
        predicted, gold, _events_match,
        lambda p: p.event_class.value,
        lambda g: g.event_class.value,
        all_classes,
    )

    return EventMetricsResult(
        macro_f1=macro_f1,
        micro=micro,
        per_class=per_class,
    )


# ---------------------------------------------------------------------------
# Public API — Relations
# ---------------------------------------------------------------------------


def compute_relation_metrics(
    predicted: list[PredictedRelation],
    gold: list[GoldRelation],
) -> RelationMetricsResult:
    """Compute relation macro-F1, micro-F1, and per-type F1.

    Args:
        predicted: Predicted relations.
        gold: Gold standard relations.

    Returns:
        RelationMetricsResult with macro, micro, and per-type breakdowns.
    """
    all_types = {r.value for r in RelationType}

    # Per-type breakdown
    per_type: dict[str, PRF1] = {}
    f1_scores: list[float] = []

    for rtype in sorted(all_types):
        type_predicted = [p for p in predicted if p.relation_type.value == rtype]
        type_gold = [g for g in gold if g.relation_type.value == rtype]
        prf1 = _compute_prf1_greedy(type_predicted, type_gold, _relations_match)
        per_type[rtype] = prf1
        f1_scores.append(prf1.f1)

    # Macro-F1: average F1 across all relation types
    macro_f1 = sum(f1_scores) / len(f1_scores) if f1_scores else 0.0

    # Micro-F1: global TP/FP/FN
    micro = _compute_micro_prf1(
        predicted, gold, _relations_match,
        lambda p: p.relation_type.value,
        lambda g: g.relation_type.value,
        all_types,
    )

    return RelationMetricsResult(
        macro_f1=macro_f1,
        micro=micro,
        per_type=per_type,
    )


# ---------------------------------------------------------------------------
# Public API — Combined Report
# ---------------------------------------------------------------------------


def evaluate_events_and_relations(
    predicted_events: list[PredictedEvent],
    gold_events: list[GoldEvent],
    predicted_relations: list[PredictedRelation],
    gold_relations: list[GoldRelation],
    document_count: int = 1,
) -> EventRelationEvaluationReport:
    """Run full event and relation evaluation producing a complete report.

    Args:
        predicted_events: All predicted events.
        gold_events: All gold standard events.
        predicted_relations: All predicted relations.
        gold_relations: All gold standard relations.
        document_count: Number of documents evaluated.

    Returns:
        EventRelationEvaluationReport with event metrics, relation metrics.
    """
    event_metrics = compute_event_metrics(predicted_events, gold_events)
    relation_metrics = compute_relation_metrics(predicted_relations, gold_relations)

    return EventRelationEvaluationReport(
        event_metrics=event_metrics,
        relation_metrics=relation_metrics,
        document_count=document_count,
    )