"""Entity and ticker precision, recall, F1, and ambiguity accuracy metrics.

Implements evaluation metrics for entity extraction quality against a gold
standard corpus. Supports both strict matching (exact span) and relaxed
matching (overlapping span with same type), with per-type breakdowns.

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

from enum import Enum
from typing import Literal

from pydantic import BaseModel, Field

# ---------------------------------------------------------------------------
# Domain Models
# ---------------------------------------------------------------------------


class MatchMode(str, Enum):
    """Entity matching strategy."""

    strict = "strict"
    relaxed = "relaxed"


class EntitySpan(BaseModel):
    """A single entity mention with character offsets and type."""

    text: str
    entity_type: str
    start_char: int
    end_char: int
    document_id: str = ""
    canonical_id: str | None = None
    is_ambiguous: bool = False

    @property
    def span(self) -> tuple[int, int]:
        return (self.start_char, self.end_char)


class TickerMention(BaseModel):
    """A resolved ticker/company mention."""

    text: str
    ticker: str
    start_char: int
    end_char: int
    document_id: str = ""
    canonical_company_id: str | None = None
    is_ambiguous: bool = False

    @property
    def span(self) -> tuple[int, int]:
        return (self.start_char, self.end_char)


# ---------------------------------------------------------------------------
# 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 EntityMetricsResult(BaseModel):
    """Full entity evaluation result with per-type breakdowns."""

    match_mode: Literal["strict", "relaxed"]
    overall: PRF1
    per_type: dict[str, PRF1]


class TickerMetricsResult(BaseModel):
    """Ticker/company resolution evaluation result."""

    match_mode: Literal["strict", "relaxed"]
    overall: PRF1
    per_type: dict[str, PRF1] = Field(
        default_factory=dict,
        description="Breakdown by canonical company or sector if available",
    )


class AmbiguityResult(BaseModel):
    """Ambiguity detection accuracy."""

    accuracy: float = Field(ge=0.0, le=1.0)
    true_positives: int = Field(ge=0)
    true_negatives: int = Field(ge=0)
    false_positives: int = Field(ge=0)
    false_negatives: int = Field(ge=0)
    support: int = Field(ge=0)


class EntityEvaluationReport(BaseModel):
    """Complete entity evaluation report."""

    entity_metrics: EntityMetricsResult
    ticker_metrics: TickerMetricsResult
    ambiguity_accuracy: AmbiguityResult
    document_count: int = Field(ge=0)


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


def _spans_overlap(a: tuple[int, int], b: tuple[int, int]) -> bool:
    """Return True if two character spans overlap."""
    return a[0] < b[1] and b[0] < a[1]


def _entity_matches_strict(pred: EntitySpan, gold: EntitySpan) -> bool:
    """Strict match: exact span boundaries and same entity type."""
    return (
        pred.entity_type == gold.entity_type
        and pred.start_char == gold.start_char
        and pred.end_char == gold.end_char
    )


def _entity_matches_relaxed(pred: EntitySpan, gold: EntitySpan) -> bool:
    """Relaxed match: overlapping span with same entity type."""
    return pred.entity_type == gold.entity_type and _spans_overlap(
        pred.span, gold.span
    )


def _ticker_matches_strict(pred: TickerMention, gold: TickerMention) -> bool:
    """Strict match: exact span and same resolved ticker."""
    return (
        pred.ticker == gold.ticker
        and pred.start_char == gold.start_char
        and pred.end_char == gold.end_char
    )


def _ticker_matches_relaxed(pred: TickerMention, gold: TickerMention) -> bool:
    """Relaxed match: overlapping span with same resolved ticker."""
    return pred.ticker == gold.ticker and _spans_overlap(pred.span, gold.span)


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


def _compute_prf1(
    predicted: list[EntitySpan] | list[TickerMention],
    gold: list[EntitySpan] | list[TickerMention],
    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,
        )

    # Greedy matching: for each predicted, find first unmatched gold
    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,
    )


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def compute_entity_metrics(
    predicted: list[EntitySpan],
    gold: list[EntitySpan],
    mode: MatchMode = MatchMode.strict,
) -> EntityMetricsResult:
    """Compute entity precision, recall, F1 with per-type breakdowns.

    Args:
        predicted: Predicted entity spans.
        gold: Gold standard entity spans.
        mode: Matching strategy (strict or relaxed).

    Returns:
        EntityMetricsResult with overall and per-type PRF1.
    """
    match_fn = _entity_matches_strict if mode == MatchMode.strict else _entity_matches_relaxed

    # Overall
    overall = _compute_prf1(predicted, gold, match_fn)

    # Per-type breakdown
    all_types = {e.entity_type for e in predicted} | {e.entity_type for e in gold}
    per_type: dict[str, PRF1] = {}

    for entity_type in sorted(all_types):
        type_predicted = [e for e in predicted if e.entity_type == entity_type]
        type_gold = [e for e in gold if e.entity_type == entity_type]
        per_type[entity_type] = _compute_prf1(type_predicted, type_gold, match_fn)

    return EntityMetricsResult(
        match_mode=mode.value,
        overall=overall,
        per_type=per_type,
    )


def compute_ticker_metrics(
    predicted: list[TickerMention],
    gold: list[TickerMention],
    mode: MatchMode = MatchMode.strict,
) -> TickerMetricsResult:
    """Compute ticker/company resolution precision, recall, F1.

    Args:
        predicted: Predicted ticker mentions with resolved tickers.
        gold: Gold standard ticker mentions.
        mode: Matching strategy (strict or relaxed).

    Returns:
        TickerMetricsResult with overall and optional per-ticker PRF1.
    """
    match_fn = _ticker_matches_strict if mode == MatchMode.strict else _ticker_matches_relaxed

    overall = _compute_prf1(predicted, gold, match_fn)

    # Per-ticker breakdown
    all_tickers = {t.ticker for t in predicted} | {t.ticker for t in gold}
    per_type: dict[str, PRF1] = {}

    for ticker in sorted(all_tickers):
        ticker_predicted = [t for t in predicted if t.ticker == ticker]
        ticker_gold = [t for t in gold if t.ticker == ticker]
        per_type[ticker] = _compute_prf1(ticker_predicted, ticker_gold, match_fn)

    return TickerMetricsResult(
        match_mode=mode.value,
        overall=overall,
        per_type=per_type,
    )


def compute_ambiguity_accuracy(
    predicted: list[EntitySpan] | list[TickerMention],
    gold: list[EntitySpan] | list[TickerMention],
) -> AmbiguityResult:
    """Compute ambiguity detection accuracy.

    Measures how well the system identifies entities that require
    adjudication (ambiguous entities). Uses the `is_ambiguous` flag
    on each span/mention.

    Entities are aligned by position (exact start_char, end_char match)
    to compare ambiguity labels.

    Args:
        predicted: Predicted entities/tickers with ambiguity flags.
        gold: Gold standard entities/tickers with ambiguity flags.

    Returns:
        AmbiguityResult with accuracy and confusion counts.
    """
    # Build a lookup from gold spans to ambiguity flag
    gold_lookup: dict[tuple[int, int], bool] = {}
    for g in gold:
        gold_lookup[(g.start_char, g.end_char)] = g.is_ambiguous

    tp = 0  # predicted ambiguous, gold ambiguous
    tn = 0  # predicted not ambiguous, gold not ambiguous
    fp = 0  # predicted ambiguous, gold not ambiguous
    fn = 0  # predicted not ambiguous, gold ambiguous

    matched_count = 0

    for p in predicted:
        key = (p.start_char, p.end_char)
        if key in gold_lookup:
            matched_count += 1
            gold_ambiguous = gold_lookup[key]
            pred_ambiguous = p.is_ambiguous

            if pred_ambiguous and gold_ambiguous:
                tp += 1
            elif not pred_ambiguous and not gold_ambiguous:
                tn += 1
            elif pred_ambiguous and not gold_ambiguous:
                fp += 1
            else:
                fn += 1

    support = tp + tn + fp + fn
    accuracy = (tp + tn) / support if support > 0 else 1.0

    return AmbiguityResult(
        accuracy=accuracy,
        true_positives=tp,
        true_negatives=tn,
        false_positives=fp,
        false_negatives=fn,
        support=support,
    )


def evaluate_entities(
    predicted_entities: list[EntitySpan],
    gold_entities: list[EntitySpan],
    predicted_tickers: list[TickerMention],
    gold_tickers: list[TickerMention],
    mode: MatchMode = MatchMode.strict,
    document_count: int = 1,
) -> EntityEvaluationReport:
    """Run full entity evaluation producing a complete report.

    Args:
        predicted_entities: All predicted entity spans.
        gold_entities: All gold standard entity spans.
        predicted_tickers: All predicted ticker mentions.
        gold_tickers: All gold standard ticker mentions.
        mode: Matching strategy.
        document_count: Number of documents evaluated.

    Returns:
        EntityEvaluationReport with entity metrics, ticker metrics,
        and ambiguity accuracy.
    """
    entity_metrics = compute_entity_metrics(predicted_entities, gold_entities, mode)
    ticker_metrics = compute_ticker_metrics(predicted_tickers, gold_tickers, mode)
    ambiguity_accuracy = compute_ambiguity_accuracy(predicted_entities, gold_entities)

    return EntityEvaluationReport(
        entity_metrics=entity_metrics,
        ticker_metrics=ticker_metrics,
        ambiguity_accuracy=ambiguity_accuracy,
        document_count=document_count,
    )