"""Post-hoc confidence calibration for answer/path/verifier scores."""

from __future__ import annotations

import math
from dataclasses import dataclass

import torch
from torch import Tensor

#: Beyond this the fit is degenerate: the calibrator becomes a constant predictor.
MAX_TEMPERATURE = 100.0


@dataclass(frozen=True)
class TemperatureCalibrator:
    temperature: float = 1.0

    def transform(self, logits: Tensor) -> Tensor:
        if self.temperature <= 0:
            raise ValueError("temperature must be positive")
        return torch.sigmoid(logits / self.temperature)

    def fit(
        self, logits: Tensor, labels: Tensor, iterations: int = 100
    ) -> TemperatureCalibrator:
        value = torch.tensor(
            float(self.temperature),
            dtype=logits.dtype,
            device=logits.device,
            requires_grad=True,
        )
        optimizer = torch.optim.LBFGS(
            [value], max_iter=iterations, line_search_fn="strong_wolfe"
        )

        def closure() -> float:
            optimizer.zero_grad()
            loss = torch.nn.functional.binary_cross_entropy_with_logits(
                logits / value.clamp_min(1e-3), labels.float()
            )
            loss.backward()
            return float(loss.detach())

        optimizer.step(closure)
        fitted = float(value.detach().clamp_min(1e-3))
        if not math.isfinite(fitted) or fitted > MAX_TEMPERATURE:
            # A near-flat logit distribution drives temperature scaling to
            # divergence: on the 100M run LBFGS returned T=1954, which maps every
            # input to ~0.5 and is a degenerate fit rather than a calibration.
            # Report the failure by returning the identity so the promotion rule
            # compares against something meaningful instead of silently shipping a
            # constant predictor.
            return TemperatureCalibrator(1.0)
        return TemperatureCalibrator(fitted)


@dataclass(frozen=True)
class IsotonicCalibrator:
    """Dependency-free monotonic probability calibrator using PAVA."""

    thresholds: tuple[float, ...]
    values: tuple[float, ...]

    @classmethod
    def fit(cls, probabilities: Tensor, labels: Tensor) -> IsotonicCalibrator:
        rows = sorted(
            (float(probability), float(label))
            for probability, label in zip(
                probabilities.detach().cpu().tolist(),
                labels.detach().cpu().tolist(),
                strict=True,
            )
        )
        if not rows:
            raise ValueError("isotonic calibration requires at least one example")
        blocks: list[list[float]] = []
        for probability, label in rows:
            blocks.append([probability, probability, label, 1.0])
            while (
                len(blocks) >= 2
                and blocks[-2][2] / blocks[-2][3] > blocks[-1][2] / blocks[-1][3]
            ):
                right = blocks.pop()
                left = blocks.pop()
                blocks.append(
                    [left[0], right[1], left[2] + right[2], left[3] + right[3]]
                )
        return cls(
            tuple(block[1] for block in blocks),
            tuple(block[2] / block[3] for block in blocks),
        )

    def transform(self, probabilities: Tensor) -> Tensor:
        thresholds = torch.tensor(
            self.thresholds, device=probabilities.device, dtype=probabilities.dtype
        )
        values = torch.tensor(
            self.values, device=probabilities.device, dtype=probabilities.dtype
        )
        indices = torch.bucketize(probabilities, thresholds, right=True).clamp_max(
            len(values) - 1
        )
        return values[indices]


def calibration_metrics(
    probabilities: Tensor, labels: Tensor, bins: int = 10
) -> dict[str, float]:
    probabilities = probabilities.detach().float().clamp(0, 1)
    labels = labels.detach().float()
    brier = float(((probabilities - labels) ** 2).mean())
    ece = 0.0
    for index in range(bins):
        low, high = index / bins, (index + 1) / bins
        mask = (probabilities >= low) & (
            probabilities < high if index < bins - 1 else probabilities <= high
        )
        if mask.any():
            ece += float(mask.float().mean()) * abs(
                float(probabilities[mask].mean()) - float(labels[mask].mean())
            )
    return {"ece": ece, "brier": brier}