"""Judge evaluation metrics for EB1.

Completeness and soundness are fundamental to understanding judge quality:
- Completeness (TPR): Does the judge accept correct answers?
- Soundness (TNR): Does the judge reject incorrect answers?

Example:
    Basic usage with labels:

    >>> y_true = [True, True, False, False]  # Correct, Correct, Wrong, Wrong
    >>> y_pred = [True, False, False, True]  # Accept, Reject, Reject, Accept
    >>> cm = ConfusionMatrix.from_labels(y_true, y_pred)
    >>> cm.completeness  # TPR: 1 correct accepted / 2 total correct
    0.5
    >>> cm.soundness  # TNR: 1 wrong rejected / 2 total wrong
    0.5

    Judge-specific usage:

    >>> class Verdict:
    ...     def __init__(self, accepts: bool, is_correct: bool):
    ...         self._accepts = accepts
    ...         self._is_correct = is_correct
    ...     def get_predicted_label(self) -> bool:
    ...         return self._accepts
    ...     def get_true_label(self) -> bool:
    ...         return self._is_correct
    >>>
    >>> verdicts = [
    ...     Verdict(accepts=True, is_correct=True),   # TP
    ...     Verdict(accepts=False, is_correct=False), # TN
    ...     Verdict(accepts=True, is_correct=False),  # FP
    ...     Verdict(accepts=False, is_correct=True),  # FN
    ... ]
    >>> metrics = JudgeMetrics.from_verdicts(verdicts, judge_name="gpt-5.1")
    >>> metrics.completeness
    0.5
    >>> metrics.soundness
    0.5
    >>> metrics.is_well_calibrated(min_completeness=0.85, min_soundness=0.90)
    False
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Protocol


class HasVerdict(Protocol):
    """Protocol for objects with judge verdict and ground truth.

    Any object implementing these two methods can be used with JudgeMetrics.
    """

    def get_predicted_label(self) -> bool:
        """Return True if judge accepts, False if judge rejects."""
        ...

    def get_true_label(self) -> bool:
        """Return True if answer is correct, False if incorrect."""
        ...


@dataclass(frozen=True)
class ConfusionMatrix:
    """Confusion matrix for binary classification.

    Attributes:
        tp: True positives (correct answers accepted)
        tn: True negatives (incorrect answers rejected)
        fp: False positives (incorrect answers accepted)
        fn: False negatives (correct answers rejected)
    """

    tp: int
    tn: int
    fp: int
    fn: int

    @property
    def completeness(self) -> float:
        """True positive rate (sensitivity, recall).

        P(judge accepts | answer is correct)

        For judges, this measures: Do we catch good answers?
        Low completeness means the judge is too strict.
        """
        total_positive = self.tp + self.fn
        return self.tp / total_positive if total_positive > 0 else 0.0

    @property
    def soundness(self) -> float:
        """True negative rate (specificity).

        P(judge rejects | answer is incorrect)

        For judges, this measures: Do we filter out bad answers?
        Low soundness means the judge accepts wrong answers (dangerous!).
        """
        total_negative = self.tn + self.fp
        return self.tn / total_negative if total_negative > 0 else 0.0

    @property
    def precision(self) -> float:
        """Positive predictive value.

        P(answer is correct | judge accepts)

        Measures: When the judge says "accept", how often is it right?
        """
        total_predicted_positive = self.tp + self.fp
        return self.tp / total_predicted_positive if total_predicted_positive > 0 else 0.0

    @property
    def recall(self) -> float:
        """Recall (same as completeness/TPR)."""
        return self.completeness

    @property
    def f1_score(self) -> float:
        """Harmonic mean of precision and recall."""
        p = self.precision
        r = self.completeness
        return 2 * p * r / (p + r) if (p + r) > 0 else 0.0

    @property
    def accuracy(self) -> float:
        """Overall accuracy: (TP + TN) / total."""
        total = self.tp + self.tn + self.fp + self.fn
        return (self.tp + self.tn) / total if total > 0 else 0.0

    @classmethod
    def from_labels(
        cls,
        y_true: list[bool],
        y_pred: list[bool],
    ) -> ConfusionMatrix:
        """Compute confusion matrix from true and predicted labels.

        Args:
            y_true: Ground truth labels (True = correct, False = incorrect)
            y_pred: Predicted labels (True = accept, False = reject)

        Returns:
            ConfusionMatrix with counts.

        Raises:
            ValueError: If lists have different lengths.
        """
        if len(y_true) != len(y_pred):
            raise ValueError(
                f"Label length mismatch: {len(y_true)} true vs {len(y_pred)} predicted"
            )

        tp = sum(1 for t, p in zip(y_true, y_pred, strict=True) if t and p)
        tn = sum(1 for t, p in zip(y_true, y_pred, strict=True) if not t and not p)
        fp = sum(1 for t, p in zip(y_true, y_pred, strict=True) if not t and p)
        fn = sum(1 for t, p in zip(y_true, y_pred, strict=True) if t and not p)

        return cls(tp=tp, tn=tn, fp=fp, fn=fn)

    @classmethod
    def from_verdicts(cls, verdicts: list[HasVerdict]) -> ConfusionMatrix:
        """Compute confusion matrix from verdict objects.

        Args:
            verdicts: List of objects implementing the HasVerdict protocol.

        Returns:
            ConfusionMatrix with counts.
        """
        y_true = [v.get_true_label() for v in verdicts]
        y_pred = [v.get_predicted_label() for v in verdicts]
        return cls.from_labels(y_true, y_pred)

    def to_dict(self) -> dict[str, int | float]:
        """Export all metrics as a dictionary.

        Returns:
            Dictionary with counts (tp, tn, fp, fn) and all computed metrics.
        """
        return {
            "tp": self.tp,
            "tn": self.tn,
            "fp": self.fp,
            "fn": self.fn,
            "completeness": self.completeness,
            "soundness": self.soundness,
            "precision": self.precision,
            "recall": self.recall,
            "f1": self.f1_score,
            "accuracy": self.accuracy,
        }


@dataclass
class JudgeMetrics:
    """Metrics for evaluating judge quality.

    Wraps a ConfusionMatrix with judge-specific semantics and terminology.
    Provides methods for calibration checking and exporting results.

    Attributes:
        confusion_matrix: Underlying confusion matrix with counts.
        judge_name: Optional identifier for the judge being evaluated.
    """

    confusion_matrix: ConfusionMatrix
    judge_name: str | None = None

    @property
    def completeness(self) -> float:
        """Fraction of correct answers the judge accepts (TPR).

        High completeness = judge catches good answers.
        Low completeness = judge is too strict, rejects valid answers.

        For a production judge, you typically want completeness >= 0.85.
        """
        return self.confusion_matrix.completeness

    @property
    def soundness(self) -> float:
        """Fraction of incorrect answers the judge rejects (TNR).

        High soundness = judge filters out bad answers.
        Low soundness = judge accepts wrong answers (dangerous!).

        For a production judge, you typically want soundness >= 0.90.
        Soundness is usually more critical than completeness.
        """
        return self.confusion_matrix.soundness

    @property
    def precision(self) -> float:
        """Precision (positive predictive value)."""
        return self.confusion_matrix.precision

    @property
    def recall(self) -> float:
        """Recall (same as completeness)."""
        return self.confusion_matrix.recall

    @property
    def f1_score(self) -> float:
        """F1 score (harmonic mean of precision and recall)."""
        return self.confusion_matrix.f1_score

    @property
    def accuracy(self) -> float:
        """Overall accuracy."""
        return self.confusion_matrix.accuracy

    def is_well_calibrated(
        self,
        min_completeness: float = 0.85,
        min_soundness: float = 0.90,
    ) -> bool:
        """Check if judge meets minimum quality thresholds.

        Args:
            min_completeness: Minimum acceptable TPR (default: 0.85).
            min_soundness: Minimum acceptable TNR (default: 0.90).

        Returns:
            True if both thresholds are met, False otherwise.

        Note:
            Default thresholds are conservative. Soundness is set higher
            because accepting wrong answers is more dangerous than rejecting
            borderline-correct ones.
        """
        return self.completeness >= min_completeness and self.soundness >= min_soundness

    @classmethod
    def from_verdicts(
        cls,
        verdicts: list[HasVerdict],
        judge_name: str | None = None,
    ) -> JudgeMetrics:
        """Compute judge metrics from verdict objects.

        Args:
            verdicts: List of objects implementing the HasVerdict protocol.
            judge_name: Optional identifier for the judge.

        Returns:
            JudgeMetrics instance with computed metrics.
        """
        cm = ConfusionMatrix.from_verdicts(verdicts)
        return cls(confusion_matrix=cm, judge_name=judge_name)

    def to_dict(self) -> dict[str, int | float | str | None]:
        """Export metrics as dictionary.

        Returns:
            Dictionary with all metrics and judge_name if set.
        """
        result: dict[str, int | float | str | None] = {
            **self.confusion_matrix.to_dict(),
            "judge_name": self.judge_name,
        }
        return result
