"""Cross-fit calibration — leakage-safe calibrated OOF + C_final.

Design invariants (from SKILL calibration):
- Calibrator NEVER sees X; it only receives (oof_scores, y).
- C_final is trained on ALL (oof_scores, y) and is only used for inference.
- The calibrated OOF used for evaluation is produced via cross-fit
  (each fold's calibrator is trained on the *other* folds — no same-row leak).
"""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from typing import Any

import numpy as np
import numpy.typing as npt

from lizyml.calibration.base import BaseCalibratorAdapter


@dataclass
class CalibrationResult:
    """Outcome of cross-fit calibration.

    Attributes:
        c_final: Calibrator trained on ALL (oof_scores, y).
            Used for inference only.
        calibrated_oof: Cross-fit calibrated OOF probabilities.
            Length equals ``n_samples``; same row ordering as ``oof_scores``.
            Used for evaluation — never produced by ``c_final``.
        method: Calibration method name.
        split_indices: Per-fold ``(train_idx, valid_idx)`` used in cross-fit.
            Stored for reproducibility / audit.  Same ordering as the folds
            that produced ``calibrated_oof``.
    """

    c_final: BaseCalibratorAdapter
    calibrated_oof: npt.NDArray[np.float64]
    method: str
    split_indices: list[tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]]


def cross_fit_calibrate(
    oof_scores: npt.NDArray[np.float64],
    y: npt.NDArray[Any],
    calibrator_factory: Callable[[], BaseCalibratorAdapter],
    *,
    split_indices: list[tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]],
) -> CalibrationResult:
    """Produce leakage-safe calibrated OOF scores and a C_final calibrator.

    For each fold:
    - Calibrator is fitted on (oof_scores[train_idx], y[train_idx]).
    - Calibrated probabilities are predicted for valid_idx rows.

    C_final is trained on ALL (oof_scores, y) and is only used for inference.

    Split generation is the caller's responsibility (H-0044).  This function
    only performs the cross-fit loop over the provided indices.

    Args:
        oof_scores: 1-D base-model OOF probabilities (n_samples,).
        y: 1-D binary ground-truth labels (n_samples,).
        calibrator_factory: Callable returning a fresh calibrator instance.
        split_indices: Pre-computed ``(train_idx, valid_idx)`` pairs.
            Generated by the appropriate splitter for the configured
            ``split.method``.

    Returns:
        :class:`CalibrationResult` with ``c_final`` and ``calibrated_oof``.
    """
    calibrated_oof = np.empty_like(oof_scores, dtype=float)

    for train_idx, val_idx in split_indices:
        cal = calibrator_factory()
        cal.fit(oof_scores[train_idx], y[train_idx])
        calibrated_oof[val_idx] = cal.predict(oof_scores[val_idx])

    # C_final: trained on ALL data — for inference only, never for evaluation
    c_final = calibrator_factory()
    c_final.fit(oof_scores, y)

    return CalibrationResult(
        c_final=c_final,
        calibrated_oof=calibrated_oof,
        method=calibrator_factory().name,
        split_indices=split_indices,
    )
