# llmindex.io — calibration metrics: Brier score, expected calibration error
# Author:  Simon-Pierre Boucher
# Contact: contact@spboucher.ai
# License: Proprietary — © Simon-Pierre Boucher, all rights reserved

from __future__ import annotations

import numpy as np


def brier_score(confidence: np.ndarray, correct: np.ndarray) -> float:
    """Mean squared error between reported confidence and outcome."""
    c = np.asarray(confidence, dtype=float)
    y = np.asarray(correct, dtype=float)
    if c.shape != y.shape or c.size == 0:
        raise ValueError("confidence and correct must be same-shape, non-empty")
    return float(np.mean((c - y) ** 2))


def expected_calibration_error(
    confidence: np.ndarray, correct: np.ndarray, bins: int = 10
) -> float:
    """Standard ECE with equal-width confidence bins."""
    c = np.asarray(confidence, dtype=float)
    y = np.asarray(correct, dtype=float)
    if c.shape != y.shape or c.size == 0:
        raise ValueError("confidence and correct must be same-shape, non-empty")
    edges = np.linspace(0, 1, bins + 1)
    idx = np.clip(np.digitize(c, edges[1:-1]), 0, bins - 1)
    ece = 0.0
    for b in range(bins):
        sel = idx == b
        n = sel.sum()
        if n == 0:
            continue
        ece += (n / c.size) * abs(y[sel].mean() - c[sel].mean())
    return float(ece)