"""Inter-annotator agreement metrics for the Gold Corpus.

Implements Cohen's kappa and weighted kappa for evaluating annotation
consistency on the double-reviewed hard-case subset.

Target thresholds:
- κ ≥ 0.80 for entities and events
- κ ≥ 0.70 for relations and sentiment
"""

from __future__ import annotations

from collections import Counter

from pydantic import BaseModel, Field

# ---------------------------------------------------------------------------
# Agreement thresholds
# ---------------------------------------------------------------------------


class AgreementThresholds(BaseModel):
    """Target inter-annotator agreement thresholds per field type."""

    entities: float = Field(default=0.80, ge=0.0, le=1.0)
    events: float = Field(default=0.80, ge=0.0, le=1.0)
    relations: float = Field(default=0.70, ge=0.0, le=1.0)
    sentiment: float = Field(default=0.70, ge=0.0, le=1.0)


# ---------------------------------------------------------------------------
# Inter-annotator report
# ---------------------------------------------------------------------------


class FieldAgreement(BaseModel):
    """Agreement score for a single annotation field."""

    field_name: str
    kappa: float = Field(description="Cohen's kappa or weighted kappa value.")
    threshold: float = Field(description="Required minimum kappa.")
    meets_threshold: bool
    n_items: int = Field(description="Number of items compared.")
    agreement_rate: float = Field(
        ge=0.0, le=1.0, description="Raw proportion of agreement."
    )


class InterAnnotatorReport(BaseModel):
    """Complete inter-annotator agreement report across all annotation fields."""

    annotator_a: str
    annotator_b: str
    n_documents: int
    field_agreements: list[FieldAgreement] = Field(default_factory=list)
    overall_kappa: float = Field(description="Mean kappa across all fields.")
    all_thresholds_met: bool


# ---------------------------------------------------------------------------
# Cohen's Kappa (categorical)
# ---------------------------------------------------------------------------


def compute_cohens_kappa(
    annotations_a: list[str],
    annotations_b: list[str],
) -> float:
    """Compute Cohen's kappa for categorical labels between two annotators.

    Cohen's kappa measures agreement between two raters while accounting
    for agreement by chance.

    κ = (p_o - p_e) / (1 - p_e)

    where:
    - p_o = observed agreement proportion
    - p_e = expected agreement by chance

    Args:
        annotations_a: Labels from annotator A.
        annotations_b: Labels from annotator B.

    Returns:
        Cohen's kappa value in [-1, 1]. Values:
        - 1.0 = perfect agreement
        - 0.0 = agreement equivalent to chance
        - <0 = less agreement than expected by chance

    Raises:
        ValueError: If annotation lists have different lengths or are empty.
    """
    if len(annotations_a) != len(annotations_b):
        raise ValueError(
            f"Annotation lists must have equal length, "
            f"got {len(annotations_a)} and {len(annotations_b)}"
        )

    if not annotations_a:
        raise ValueError("Cannot compute kappa on empty annotations.")

    n = len(annotations_a)

    # Observed agreement
    agreements = sum(1 for a, b in zip(annotations_a, annotations_b) if a == b)
    p_o = agreements / n

    # Expected agreement by chance
    categories = set(annotations_a) | set(annotations_b)
    counts_a = Counter(annotations_a)
    counts_b = Counter(annotations_b)

    p_e = sum((counts_a[cat] / n) * (counts_b[cat] / n) for cat in categories)

    # Handle edge case where p_e = 1 (both annotators always pick same category)
    if abs(1.0 - p_e) < 1e-10:
        return 1.0 if p_o == 1.0 else 0.0

    kappa = (p_o - p_e) / (1.0 - p_e)
    return kappa


# ---------------------------------------------------------------------------
# Weighted Kappa (ordinal)
# ---------------------------------------------------------------------------


def compute_weighted_kappa(
    annotations_a: list[str],
    annotations_b: list[str],
    ordered_categories: list[str] | None = None,
    weight_type: str = "linear",
) -> float:
    """Compute weighted Cohen's kappa for ordinal ratings.

    Weighted kappa accounts for the degree of disagreement between ordinal
    categories. Linear weights penalize disagreements proportional to their
    distance; quadratic weights penalize proportional to squared distance.

    Args:
        annotations_a: Labels from annotator A.
        annotations_b: Labels from annotator B.
        ordered_categories: Ordered list of categories (low to high).
            If None, categories are sorted alphabetically.
        weight_type: "linear" or "quadratic" weighting.

    Returns:
        Weighted kappa value.

    Raises:
        ValueError: If inputs are invalid.
    """
    if len(annotations_a) != len(annotations_b):
        raise ValueError(
            f"Annotation lists must have equal length, "
            f"got {len(annotations_a)} and {len(annotations_b)}"
        )

    if not annotations_a:
        raise ValueError("Cannot compute kappa on empty annotations.")

    if weight_type not in ("linear", "quadratic"):
        raise ValueError(f"weight_type must be 'linear' or 'quadratic', got '{weight_type}'")

    # Determine category ordering
    if ordered_categories is None:
        ordered_categories = sorted(set(annotations_a) | set(annotations_b))

    n_categories = len(ordered_categories)
    if n_categories < 2:
        # With only one category, kappa is undefined (perfect agreement trivially)
        return 1.0

    cat_index = {cat: i for i, cat in enumerate(ordered_categories)}
    n = len(annotations_a)

    # Build weight matrix
    def _weight(i: int, j: int) -> float:
        max_dist = n_categories - 1
        if max_dist == 0:
            return 0.0
        dist = abs(i - j) / max_dist
        if weight_type == "linear":
            return dist
        else:  # quadratic
            return dist ** 2

    # Observed disagreement
    observed_disagreement = 0.0
    for a, b in zip(annotations_a, annotations_b):
        i = cat_index.get(a)
        j = cat_index.get(b)
        if i is None or j is None:
            raise ValueError(
                f"Annotation value not in ordered_categories: a='{a}', b='{b}'"
            )
        observed_disagreement += _weight(i, j)
    observed_disagreement /= n

    # Expected disagreement by chance
    counts_a = Counter(annotations_a)
    counts_b = Counter(annotations_b)

    expected_disagreement = 0.0
    for cat_i, idx_i in cat_index.items():
        for cat_j, idx_j in cat_index.items():
            expected_disagreement += (
                (counts_a[cat_i] / n) * (counts_b[cat_j] / n) * _weight(idx_i, idx_j)
            )

    # Handle edge case
    if abs(expected_disagreement) < 1e-10:
        return 1.0 if abs(observed_disagreement) < 1e-10 else 0.0

    kappa = 1.0 - (observed_disagreement / expected_disagreement)
    return kappa
