"""IsotonicCalibrator — isotonic regression via LGBM monotone constraints.

Uses a single-feature LightGBM Booster with ``monotone_constraints=[1]``
to learn a monotone mapping from raw scores to calibrated probabilities
(BLUEPRINT §12.2, H-0047).
"""

from __future__ import annotations

import math
from pathlib import Path
from typing import Any

import lightgbm as lgbm
import numpy as np
import numpy.typing as npt

from lizyml.calibration.base import BaseCalibratorAdapter
from lizyml.core.registries import CalibratorRegistry

_ISOTONIC_DEFAULTS: dict[str, Any] = {
    "objective": "binary",
    "metric": "binary_logloss",
    "monotone_constraints": [1],
    "monotone_constraints_method": "advanced",
    "num_leaves": 7,
    "max_depth": 3,
    "min_data_in_leaf_ratio": 0.01,
    "learning_rate": 0.03,
    "lambda_l2": 5.0,
    "min_gain_to_split": 0.0,
    "feature_fraction": 1.0,
    "bagging_fraction": 1.0,
    "bagging_freq": 0,
}

_NUM_BOOST_ROUND = 1000
_EARLY_STOPPING_ROUNDS = 100
_VALIDATION_RATIO = 0.1
_DEFAULT_SEED = 42
_MIN_SAMPLES_FOR_EARLY_STOPPING = 20


@CalibratorRegistry.register("isotonic")
class IsotonicCalibrator(BaseCalibratorAdapter):
    """Isotonic calibration using LGBM Booster with monotone constraints.

    Learns a monotone non-decreasing mapping from raw scores (logits)
    to calibrated probabilities using a single-feature LightGBM Booster.

    Args:
        params: Optional parameter overrides.  ``monotone_constraints``
            is always forced to ``[1]`` and cannot be overridden.
            Special keys (not passed to LightGBM):
            - ``num_boost_round``: number of boosting rounds (default 1000)
            - ``validation_ratio``: fraction for early stopping (default 0.1)
            - ``seed``: random seed for validation split (default 42)
    """

    def __init__(self, params: dict[str, Any] | None = None) -> None:
        # Copy to avoid mutating the caller's dict
        user = dict(params) if params else {}

        # Extract non-LightGBM params
        self._num_boost_round = int(user.pop("num_boost_round", _NUM_BOOST_ROUND))
        self._validation_ratio = float(user.pop("validation_ratio", _VALIDATION_RATIO))
        self._seed = int(user.pop("seed", _DEFAULT_SEED))

        if not (0.0 < self._validation_ratio < 1.0):
            from lizyml.core.exceptions import ErrorCode, LizyMLError

            raise LizyMLError(
                code=ErrorCode.CONFIG_INVALID,
                user_message="validation_ratio must be in (0.0, 1.0)",
                context={"validation_ratio": self._validation_ratio},
            )

        merged = {**_ISOTONIC_DEFAULTS, **user}
        # Always enforce monotone constraint
        merged["monotone_constraints"] = [1]
        merged["verbose"] = -1
        merged["seed"] = self._seed
        self._lgbm_params = merged
        self._model: lgbm.Booster | None = None

    @property
    def name(self) -> str:
        return "isotonic"

    def fit(
        self, oof_scores: npt.NDArray[np.float64], y: npt.NDArray[Any]
    ) -> IsotonicCalibrator:
        X_cal = oof_scores.reshape(-1, 1)
        y_float = y.astype(float)
        n_samples = len(X_cal)

        # Resolve min_data_in_leaf_ratio to absolute value
        params = dict(self._lgbm_params)
        ratio = params.pop("min_data_in_leaf_ratio", None)
        if ratio is not None:
            params["min_data_in_leaf"] = max(1, math.ceil(n_samples * ratio))

        train_ds, valid_sets, callbacks = self._prepare_training(
            X_cal, y_float, n_samples, params
        )
        self._model = lgbm.train(
            params,
            train_ds,
            num_boost_round=self._num_boost_round,
            valid_sets=valid_sets,
            valid_names=["valid"] if valid_sets else None,
            callbacks=callbacks,
        )
        return self

    def _prepare_training(
        self,
        X_cal: npt.NDArray[np.float64],
        y_float: npt.NDArray[np.float64],
        n_samples: int,
        params: dict[str, Any],
    ) -> tuple[lgbm.Dataset, list[lgbm.Dataset], list[Any]]:
        """Build train dataset, optional valid sets, and callbacks."""
        callbacks: list[Any] = [lgbm.log_evaluation(period=0)]

        if n_samples < _MIN_SAMPLES_FOR_EARLY_STOPPING:
            return lgbm.Dataset(X_cal, label=y_float), [], callbacks

        # Split for early stopping validation
        rng = np.random.default_rng(self._seed)
        n_valid = max(1, int(n_samples * self._validation_ratio))
        indices = rng.permutation(n_samples)
        valid_idx, train_idx = indices[:n_valid], indices[n_valid:]

        train_ds = lgbm.Dataset(X_cal[train_idx], label=y_float[train_idx])
        valid_ds = lgbm.Dataset(
            X_cal[valid_idx], label=y_float[valid_idx], reference=train_ds
        )
        callbacks.insert(
            0,
            lgbm.early_stopping(stopping_rounds=_EARLY_STOPPING_ROUNDS, verbose=False),
        )
        return train_ds, [valid_ds], callbacks

    def predict(self, scores: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        if self._model is None:
            raise RuntimeError("IsotonicCalibrator has not been fitted.")
        X = scores.reshape(-1, 1)
        # objective="binary" Booster.predict() returns probabilities (sigmoid
        # is applied internally by LightGBM), so no manual sigmoid needed.
        raw_pred = self._model.predict(X)
        proba: npt.NDArray[np.float64] = np.asarray(raw_pred, dtype=np.float64)
        clipped: npt.NDArray[np.float64] = np.clip(proba, 0.0, 1.0)
        return clipped

    def export_params(self) -> dict[str, Any]:
        """Export isotonic calibrator metadata."""
        if self._model is None:
            raise RuntimeError("IsotonicCalibrator has not been fitted.")
        return {"method": "isotonic", "model_file": "calibrator_model.txt"}

    def save_model_text(self, path: str | Path) -> Path:
        """Save the internal Booster to a human-readable text file."""
        if self._model is None:
            raise RuntimeError("IsotonicCalibrator has not been fitted.")
        p = Path(path)
        self._model.save_model(str(p))
        return p
