#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""Prediction calibration per domain — Item 169.

  - Item 169: Prediction calibration per domain (Metacognition)

Track prediction accuracy separately per domain.
Some domains may be well-calibrated while others are overconfident.
"""
import logging
from typing import Dict, Any, List

log = logging.getLogger("pulse.brain.prediction_calibration")

# {domain: {"predicted_sum": float, "actual_sum": float, "count": int}}
_DOMAIN_CALIBRATION: Dict[str, Dict[str, float]] = {}
_MIN_SAMPLES = 5    # need this many samples before computing calibration

# N54: Raw observations for resolution (point-biserial) computation
_RAW_OBSERVATIONS: List[tuple] = []  # [(predicted_confidence, binary_outcome)]
_MAX_RAW_OBSERVATIONS = 200


def record_domain_prediction(
    domain: str,
    predicted_confidence: float,
    task_succeeded: bool,
) -> None:
    """Record a prediction outcome for a domain. Fails open."""
    try:
        if not domain:
            domain = "general"
        entry = _DOMAIN_CALIBRATION.setdefault(domain, {
            "predicted_sum": 0.0,
            "actual_sum": 0.0,
            "count": 0,
        })
        entry["predicted_sum"] += float(predicted_confidence)
        entry["actual_sum"] += 1.0 if task_succeeded else 0.0
        entry["count"] += 1
        # N54: track raw pairs for resolution computation
        _RAW_OBSERVATIONS.append((float(predicted_confidence), 1.0 if task_succeeded else 0.0))
        if len(_RAW_OBSERVATIONS) > _MAX_RAW_OBSERVATIONS:
            _RAW_OBSERVATIONS.pop(0)
    except Exception as exc:
        log.debug("record_domain_prediction failed: %s", exc)


def get_domain_calibration() -> Dict[str, Any]:
    """Return calibration stats per domain.

    Returns:
        {
            domains: {domain: {avg_predicted, avg_actual, count,
                                calibration_error, status}},
            overconfident_domains: [domain],
            underconfident_domains: [domain],
        }
    calibration_error = avg_predicted - avg_actual (positive = overconfident).
    status: "well_calibrated" | "overconfident" | "underconfident" | "insufficient_data"
    Fails open (returns empty) on any error.
    """
    result: Dict[str, Any] = {
        "domains": {},
        "overconfident_domains": [],
        "underconfident_domains": [],
    }
    try:
        overconfident = []
        underconfident = []

        for domain, entry in _DOMAIN_CALIBRATION.items():
            n = entry.get("count", 0)
            domain_result: Dict[str, Any] = {
                "count": n,
                "avg_predicted": 0.0,
                "avg_actual": 0.0,
                "calibration_error": 0.0,
                "status": "insufficient_data",
            }
            if n >= _MIN_SAMPLES:
                avg_p = entry["predicted_sum"] / n
                avg_a = entry["actual_sum"] / n
                error = round(avg_p - avg_a, 3)
                if error > 0.15:
                    status = "overconfident"
                    overconfident.append(domain)
                elif error < -0.15:
                    status = "underconfident"
                    underconfident.append(domain)
                else:
                    status = "well_calibrated"
                domain_result.update({
                    "avg_predicted": round(avg_p, 3),
                    "avg_actual": round(avg_a, 3),
                    "calibration_error": error,
                    "status": status,
                })

            result["domains"][domain] = domain_result

        result["overconfident_domains"] = overconfident
        result["underconfident_domains"] = underconfident

    except Exception as exc:
        log.debug("get_domain_calibration failed: %s", exc)

    return result


def compute_resolution() -> float:
    """Compute point-biserial correlation between predicted confidence and binary outcome.

    This is the resolution (gamma/d-prime) metric: measures discrimination ability.
    Returns float in [-1, 1]. Higher = better discrimination. Returns 0.0 on insufficient data.
    """
    import math

    # Collect all (predicted, actual) pairs across all domains
    pairs: list[tuple[float, float]] = []
    for entry in _DOMAIN_CALIBRATION.values():
        n = entry.get("count", 0)
        if n < 1:
            continue
        # We only store aggregates, so reconstruct a single representative pair per domain
        # when raw pairs are not stored.
        pass

    # Use per-observation storage if available (populated by record_domain_prediction)
    pairs = list(_RAW_OBSERVATIONS)

    if len(pairs) < 4:
        return 0.0

    n = len(pairs)
    preds = [p for p, _ in pairs]
    actuals = [a for _, a in pairs]

    mean_p = sum(preds) / n
    mean_a = sum(actuals) / n

    # Point-biserial correlation: r = (M1 - M0) / S_total * sqrt(n1*n0 / n^2)
    n1 = sum(1 for a in actuals if a >= 0.5)
    n0 = n - n1
    if n1 == 0 or n0 == 0:
        return 0.0

    mean1 = sum(p for p, a in pairs if a >= 0.5) / n1
    mean0 = sum(p for p, a in pairs if a < 0.5) / n0

    var_total = sum((p - mean_p) ** 2 for p in preds) / n
    if var_total <= 0:
        return 0.0

    s_total = math.sqrt(var_total)
    r = ((mean1 - mean0) / s_total) * math.sqrt((n1 * n0) / (n ** 2))
    return round(max(-1.0, min(1.0, r)), 4)
