"""Sentiment macro-F1, micro-F1, direction accuracy, and probability calibration metrics.

Implements evaluation metrics for company-specific sentiment extraction quality
and probability calibration against a gold standard corpus. Includes Expected
Calibration Error (ECE), Brier score, and reliability diagram data.

Sentiments are matched by company_entity_id between predicted and gold sets.

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

from enum import Enum

from pydantic import BaseModel, Field

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

SENTIMENT_LABELS = ("positive", "negative", "neutral", "mixed")


class SentimentLabel(str, Enum):
    """Supported sentiment labels."""

    positive = "positive"
    negative = "negative"
    neutral = "neutral"
    mixed = "mixed"


class SentimentPrediction(BaseModel):
    """A predicted or gold sentiment for a specific company entity."""

    company_entity_id: str
    label: SentimentLabel
    positive_prob: float = Field(ge=0.0, le=1.0, default=0.0)
    negative_prob: float = Field(ge=0.0, le=1.0, default=0.0)
    neutral_prob: float = Field(ge=0.0, le=1.0, default=0.0)
    mixed_prob: float = Field(ge=0.0, le=1.0, default=0.0)
    document_id: str = ""


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


class LabelF1(BaseModel):
    """Per-label precision, recall, F1."""

    label: str
    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 SentimentF1Result(BaseModel):
    """Sentiment classification F1 metrics."""

    macro_f1: float = Field(ge=0.0, le=1.0)
    micro_f1: float = Field(ge=0.0, le=1.0)
    per_label: dict[str, LabelF1]
    support: int = Field(ge=0)


class DirectionAccuracyResult(BaseModel):
    """Binary direction accuracy (positive vs negative, ignoring neutral/mixed)."""

    accuracy: float = Field(ge=0.0, le=1.0)
    correct: int = Field(ge=0)
    total: int = Field(ge=0)


class CalibrationBin(BaseModel):
    """A single bin in the reliability diagram."""

    bin_lower: float = Field(ge=0.0, le=1.0)
    bin_upper: float = Field(ge=0.0, le=1.0)
    mean_predicted_prob: float = Field(ge=0.0, le=1.0)
    fraction_positive: float = Field(ge=0.0, le=1.0)
    count: int = Field(ge=0)


class CalibrationResult(BaseModel):
    """Probability calibration metrics."""

    ece: float = Field(ge=0.0, le=1.0, description="Expected Calibration Error")
    brier_score: float = Field(ge=0.0, description="Brier score (mean squared error)")
    reliability_bins: list[CalibrationBin]
    n_samples: int = Field(ge=0)


class SentimentEvaluationReport(BaseModel):
    """Complete sentiment evaluation report."""

    f1_metrics: SentimentF1Result
    direction_accuracy: DirectionAccuracyResult
    calibration: CalibrationResult
    document_count: int = Field(ge=0)


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


def _compute_label_f1(
    predicted_labels: list[str],
    gold_labels: list[str],
    label: str,
) -> LabelF1:
    """Compute precision, recall, F1 for a single label (one-vs-rest)."""
    tp = 0
    fp = 0
    fn = 0

    for pred, gold in zip(predicted_labels, gold_labels):
        if pred == label and gold == label:
            tp += 1
        elif pred == label and gold != label:
            fp += 1
        elif pred != label and gold == label:
            fn += 1

    support_predicted = tp + fp
    support_gold = tp + fn

    precision = tp / (tp + fp) if (tp + fp) > 0 else 1.0
    recall = tp / (tp + fn) if (tp + fn) > 0 else 1.0

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

    return LabelF1(
        label=label,
        precision=precision,
        recall=recall,
        f1=f1,
        support_predicted=support_predicted,
        support_gold=support_gold,
    )


def compute_sentiment_f1(
    predicted: list[SentimentPrediction],
    gold: list[SentimentPrediction],
) -> SentimentF1Result:
    """Compute macro-F1, micro-F1, and per-label F1 for sentiment classification.

    Matches predictions to gold by company_entity_id. Only matched pairs are
    evaluated (unmatched predictions/gold are ignored).

    Args:
        predicted: Predicted sentiment labels with probabilities.
        gold: Gold standard sentiment labels.

    Returns:
        SentimentF1Result with macro-F1, micro-F1, and per-label breakdown.
    """
    # Match by company_entity_id
    gold_by_id = {g.company_entity_id: g for g in gold}
    matched_pred_labels: list[str] = []
    matched_gold_labels: list[str] = []

    for p in predicted:
        if p.company_entity_id in gold_by_id:
            matched_pred_labels.append(p.label.value)
            matched_gold_labels.append(gold_by_id[p.company_entity_id].label.value)

    support = len(matched_pred_labels)

    if support == 0:
        empty_per_label = {
            label: LabelF1(
                label=label, precision=1.0, recall=1.0, f1=1.0,
                support_predicted=0, support_gold=0,
            )
            for label in SENTIMENT_LABELS
        }
        return SentimentF1Result(
            macro_f1=1.0,
            micro_f1=1.0,
            per_label=empty_per_label,
            support=0,
        )

    # Per-label F1
    per_label: dict[str, LabelF1] = {}
    for label in SENTIMENT_LABELS:
        per_label[label] = _compute_label_f1(matched_pred_labels, matched_gold_labels, label)

    # Macro-F1: average of per-label F1 scores (only labels with support)
    active_labels = [
        label for label in SENTIMENT_LABELS
        if per_label[label].support_predicted > 0 or per_label[label].support_gold > 0
    ]
    if active_labels:
        label_f1_values = [per_label[label].f1 for label in active_labels]
        macro_f1 = sum(label_f1_values) / len(label_f1_values)
    else:
        macro_f1 = 1.0

    # Micro-F1: global TP, FP, FN across all labels
    total_tp = 0
    total_fp = 0
    total_fn = 0

    for label in SENTIMENT_LABELS:
        for pred, gold_label in zip(matched_pred_labels, matched_gold_labels):
            if pred == label and gold_label == label:
                total_tp += 1
            elif pred == label and gold_label != label:
                total_fp += 1
            elif pred != label and gold_label == label:
                total_fn += 1

    micro_precision = total_tp / (total_tp + total_fp) if (total_tp + total_fp) > 0 else 1.0
    micro_recall = total_tp / (total_tp + total_fn) if (total_tp + total_fn) > 0 else 1.0

    if micro_precision + micro_recall > 0:
        micro_f1 = 2 * micro_precision * micro_recall / (micro_precision + micro_recall)
    else:
        micro_f1 = 0.0

    return SentimentF1Result(
        macro_f1=macro_f1,
        micro_f1=micro_f1,
        per_label=per_label,
        support=support,
    )


def compute_direction_accuracy(
    predicted: list[SentimentPrediction],
    gold: list[SentimentPrediction],
) -> DirectionAccuracyResult:
    """Compute binary direction accuracy (positive vs negative).

    Only considers matched pairs where BOTH predicted and gold labels are
    either 'positive' or 'negative'. Neutral and mixed are ignored.

    Args:
        predicted: Predicted sentiment labels.
        gold: Gold standard sentiment labels.

    Returns:
        DirectionAccuracyResult with accuracy and counts.
    """
    gold_by_id = {g.company_entity_id: g for g in gold}
    correct = 0
    total = 0

    directional_labels = {SentimentLabel.positive, SentimentLabel.negative}

    for p in predicted:
        if p.company_entity_id not in gold_by_id:
            continue
        g = gold_by_id[p.company_entity_id]
        # Both must be directional (positive or negative)
        if p.label in directional_labels and g.label in directional_labels:
            total += 1
            if p.label == g.label:
                correct += 1

    accuracy = correct / total if total > 0 else 1.0

    return DirectionAccuracyResult(
        accuracy=accuracy,
        correct=correct,
        total=total,
    )


def compute_calibration(
    predicted: list[SentimentPrediction],
    gold: list[SentimentPrediction],
    n_bins: int = 10,
) -> CalibrationResult:
    """Compute Expected Calibration Error (ECE), Brier score, and reliability diagram.

    For each matched pair, we evaluate how well the predicted probability for
    the true label reflects observed frequency. Uses the maximum predicted
    probability (confidence) and checks if the predicted label matches gold.

    Args:
        predicted: Predicted sentiments with probability distributions.
        gold: Gold standard sentiments.
        n_bins: Number of bins for ECE and reliability diagram.

    Returns:
        CalibrationResult with ECE, Brier score, and per-bin data.
    """
    gold_by_id = {g.company_entity_id: g for g in gold}

    # Collect (confidence, correct) pairs
    confidences: list[float] = []
    corrects: list[int] = []
    brier_terms: list[float] = []

    for p in predicted:
        if p.company_entity_id not in gold_by_id:
            continue
        g = gold_by_id[p.company_entity_id]

        # Confidence = probability assigned to the predicted label
        confidence = _get_label_prob(p, p.label)
        is_correct = 1 if p.label == g.label else 0

        confidences.append(confidence)
        corrects.append(is_correct)

        # Brier score: sum of squared errors across all label probabilities
        # For each label, the "true" probability is 1 if it matches gold, else 0
        brier_term = 0.0
        for label in SENTIMENT_LABELS:
            pred_prob = _get_label_prob(p, SentimentLabel(label))
            true_indicator = 1.0 if label == g.label.value else 0.0
            brier_term += (pred_prob - true_indicator) ** 2
        brier_terms.append(brier_term)

    n_samples = len(confidences)

    if n_samples == 0:
        return CalibrationResult(
            ece=0.0,
            brier_score=0.0,
            reliability_bins=[],
            n_samples=0,
        )

    # Brier score: mean of per-sample squared error sums
    brier_score = sum(brier_terms) / n_samples

    # ECE and reliability diagram
    bin_width = 1.0 / n_bins
    reliability_bins: list[CalibrationBin] = []
    weighted_abs_diff_sum = 0.0

    for i in range(n_bins):
        bin_lower = i * bin_width
        bin_upper = (i + 1) * bin_width

        # Collect samples in this bin
        bin_confidences: list[float] = []
        bin_corrects: list[int] = []

        for conf, correct in zip(confidences, corrects):
            # Include in bin if conf is in [bin_lower, bin_upper)
            # Last bin includes the upper boundary
            if i == n_bins - 1:
                in_bin = bin_lower <= conf <= bin_upper
            else:
                in_bin = bin_lower <= conf < bin_upper
            if in_bin:
                bin_confidences.append(conf)
                bin_corrects.append(correct)

        bin_count = len(bin_confidences)

        if bin_count > 0:
            mean_predicted = sum(bin_confidences) / bin_count
            fraction_positive = sum(bin_corrects) / bin_count
            weighted_abs_diff_sum += bin_count * abs(mean_predicted - fraction_positive)
        else:
            mean_predicted = (bin_lower + bin_upper) / 2
            fraction_positive = 0.0

        reliability_bins.append(
            CalibrationBin(
                bin_lower=bin_lower,
                bin_upper=bin_upper,
                mean_predicted_prob=mean_predicted,
                fraction_positive=fraction_positive,
                count=bin_count,
            )
        )

    ece = weighted_abs_diff_sum / n_samples

    return CalibrationResult(
        ece=ece,
        brier_score=brier_score,
        reliability_bins=reliability_bins,
        n_samples=n_samples,
    )


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


def evaluate_sentiment(
    predicted: list[SentimentPrediction],
    gold: list[SentimentPrediction],
    n_bins: int = 10,
    document_count: int = 1,
) -> SentimentEvaluationReport:
    """Run full sentiment evaluation producing a complete report.

    Args:
        predicted: All predicted sentiment records.
        gold: All gold standard sentiment records.
        n_bins: Number of bins for calibration metrics.
        document_count: Number of documents evaluated.

    Returns:
        SentimentEvaluationReport with F1, direction accuracy, and calibration.
    """
    f1_metrics = compute_sentiment_f1(predicted, gold)
    direction_accuracy = compute_direction_accuracy(predicted, gold)
    calibration = compute_calibration(predicted, gold, n_bins=n_bins)

    return SentimentEvaluationReport(
        f1_metrics=f1_metrics,
        direction_accuracy=direction_accuracy,
        calibration=calibration,
        document_count=document_count,
    )


# ---------------------------------------------------------------------------
# Internal Helpers
# ---------------------------------------------------------------------------


def _get_label_prob(prediction: SentimentPrediction, label: SentimentLabel) -> float:
    """Get the predicted probability for a specific label."""
    if label == SentimentLabel.positive:
        return prediction.positive_prob
    elif label == SentimentLabel.negative:
        return prediction.negative_prob
    elif label == SentimentLabel.neutral:
        return prediction.neutral_prob
    elif label == SentimentLabel.mixed:
        return prediction.mixed_prob
    return 0.0
