"""
QuantileGBM: CatBoost-backed quantile and expectile regression for insurance pricing.

Design rationale:
- Quantile mode uses CatBoost's MultiQuantile loss (single model, all quantiles at once).
  This is far more efficient than fitting one model per quantile and ensures the learned
  feature representations are shared across all probability levels.
- Expectile mode fits separate CatBoost models per alpha. CatBoost has no MultiExpectile
  loss (as of 1.2.x), so we trade off efficiency for the expectile's theoretical advantages:
  it is both elicitable and satisfies subadditivity under elliptical distributions.
- Isotonic regression post-processing fixes quantile crossing (CatBoost issue #2317).
  The underlying MultiQuantile loss does not guarantee monotonicity across quantile levels
  at the prediction stage. Isotonic regression is O(n * k) and negligible in practice.
- Polars in, Polars out. Internal CatBoost conversion goes through numpy; we never expose
  pandas or CatBoost Pool objects in the public API.
- Exposure is passed as sample_weight to CatBoost. This weights the loss contribution of
  each row by its exposure. It is NOT offset modelling — there is no log-link adjustment.
  Document this clearly and recommend modelling severity (claims / exposure) as the target
  when exposure varies significantly across risks.
"""

from __future__ import annotations

from typing import Any

import numpy as np
import polars as pl
from sklearn.isotonic import IsotonicRegression

from ._types import QuantileSpec, TailModel

__all__ = ["QuantileGBM"]

# numpy<2.0 compat: trapezoid was added in 2.0, trapz deprecated in 2.0
_trapezoid = getattr(np, "trapezoid", None) or np.trapz


def _to_numpy(df) -> np.ndarray:
    """Convert a Polars DataFrame or numpy array to a float64 numpy array."""
    if isinstance(df, np.ndarray):
        return df.astype(np.float64)
    return df.to_numpy().astype(np.float64)


def _series_to_numpy(s) -> np.ndarray:
    """Convert a Polars Series or numpy array to a float64 numpy array."""
    if isinstance(s, np.ndarray):
        return s.astype(np.float64)
    return s.to_numpy().astype(np.float64)


def _apply_isotonic(predictions: np.ndarray) -> np.ndarray:
    """
    Apply per-row isotonic regression to enforce monotone quantile predictions.

    Parameters
    ----------
    predictions:
        Array of shape (n_samples, n_quantiles). Each row is a set of quantile
        predictions for one risk; we enforce q_{alpha1} <= q_{alpha2} when
        alpha1 < alpha2.

    Returns
    -------
    Array of the same shape with crossing corrected.
    """
    n_rows, n_cols = predictions.shape
    if n_cols == 1:
        return predictions
    iso = IsotonicRegression(increasing=True, out_of_bounds="clip")
    # Use equally spaced x-coordinates (the exact values don't matter,
    # only monotonicity of the fitted values matters)
    x = np.arange(n_cols, dtype=np.float64)
    corrected = np.empty_like(predictions)
    for i in range(n_rows):
        corrected[i] = iso.fit_transform(x, predictions[i])
    return corrected


class QuantileGBM:
    """
    Gradient boosted machine quantile/expectile regressor for actuarial use.

    Wraps CatBoost's MultiQuantile loss (quantile mode) or fits per-alpha
    CatBoost models (expectile mode) and returns predictions in Polars with
    actuarial column naming (q_0.5, q_0.9, etc.).

    Parameters
    ----------
    quantiles:
        Probability levels to model, e.g. [0.5, 0.75, 0.9, 0.95, 0.99].
    use_expectile:
        If True, fit expectile regression instead of quantile regression.
        Expectile regression minimises an asymmetric squared loss — unlike
        quantiles, expectiles are elicitable and satisfy subadditivity under
        elliptical distributions. They are better suited to heavy-tailed lines
        (motor BI, liability) where the tail shape matters. Default is False.
    fix_crossing:
        Apply isotonic regression post-processing to prevent quantile crossing
        at prediction time. Strongly recommended; defaults to True. Crossing
        is a known CatBoost MultiQuantile limitation (GitHub issue #2317).
    **catboost_kwargs:
        Passed directly to CatBoostRegressor. Common overrides: iterations,
        learning_rate, depth, l2_leaf_reg, random_seed. The loss_function
        and quantile parameters are set internally and should not be overridden.

    Examples
    --------
    >>> import polars as pl
    >>> from insurance_quantile import QuantileGBM
    >>> model = QuantileGBM(quantiles=[0.5, 0.75, 0.9, 0.95])
    >>> model.fit(X_train, y_train)
    >>> predictions = model.predict(X_val)
    >>> predictions.columns  # ['q_0.5', 'q_0.75', 'q_0.9', 'q_0.95']

    Notes
    -----
    Zero-inflated targets (common in insurance where most policies have zero
    claims) are handled by modelling on the full dataset. However, if the
    zero mass is large (>40% of observations), consider modelling severity
    separately on non-zero claims and using this class only for the severity
    component. The quantiles of the full (zero-included) distribution will
    bunch at zero for lower alpha levels, which is mathematically correct but
    practically less useful for large loss loading.
    """

    def __init__(
        self,
        quantiles: list[float],
        use_expectile: bool = False,
        fix_crossing: bool = True,
        **catboost_kwargs: Any,
    ) -> None:
        self._spec = QuantileSpec(
            quantiles=quantiles,
            mode="expectile" if use_expectile else "quantile",
        )
        self._fix_crossing = fix_crossing
        self._catboost_kwargs = catboost_kwargs
        self._models: list[Any] = []  # CatBoostRegressor instances
        self._metadata: TailModel | None = None
        self._is_fitted: bool = False

    # ------------------------------------------------------------------
    # Properties
    # ------------------------------------------------------------------

    @property
    def spec(self) -> QuantileSpec:
        """The QuantileSpec used to build this model."""
        return self._spec

    @property
    def metadata(self) -> TailModel:
        """Fit metadata. Raises RuntimeError if model not yet fitted."""
        if self._metadata is None:
            raise RuntimeError("Model is not fitted. Call fit() first.")
        return self._metadata

    @property
    def is_fitted(self) -> bool:
        """True if fit() has been called successfully."""
        return self._is_fitted

    # ------------------------------------------------------------------
    # Fit
    # ------------------------------------------------------------------

    def fit(
        self,
        X: pl.DataFrame,
        y: pl.Series,
        exposure: pl.Series | None = None,
    ) -> "QuantileGBM":
        """
        Fit quantile or expectile models on training data.

        Parameters
        ----------
        X:
            Feature matrix as a Polars DataFrame. All columns must be numeric.
            Categorical features should be encoded before calling fit().
        y:
            Target variable (e.g. claim cost or loss ratio) as a Polars Series.
            Must be non-negative for insurance applications.
        exposure:
            Optional exposure per row (e.g. earned car years). Passed as
            sample_weight to CatBoost. Note: this weights the loss, it is not
            an offset. If your target is aggregate cost (not severity), pass
            exposure so that high-exposure risks count more in training.

        Returns
        -------
        self, for method chaining.
        """
        from catboost import CatBoostRegressor

        X_np = _to_numpy(X)
        y_np = _series_to_numpy(y)
        w_np = _series_to_numpy(exposure) if exposure is not None else None

        base_params: dict[str, Any] = {
            "iterations": 500,
            "learning_rate": 0.05,
            "depth": 6,
            "random_seed": 42,
            "verbose": 0,
        }
        base_params.update(self._catboost_kwargs)

        self._models = []

        if self._spec.mode == "quantile":
            # Use Quantile loss for a single quantile level; MultiQuantile requires >= 2.
            # MultiQuantile: single model, all quantiles at once (efficient, shared features).
            if len(self._spec.quantiles) == 1:
                alpha = self._spec.quantiles[0]
                params = {**base_params, "loss_function": f"Quantile:alpha={alpha}"}
            else:
                quantile_str = ",".join(str(q) for q in self._spec.quantiles)
                params = {**base_params, "loss_function": f"MultiQuantile:alpha={quantile_str}"}
            model = CatBoostRegressor(**params)
            model.fit(X_np, y_np, sample_weight=w_np)
            self._models = [model]
        else:
            # Separate model per expectile level
            for alpha in self._spec.quantiles:
                params = {**base_params, "loss_function": f"Expectile:alpha={alpha}"}
                model = CatBoostRegressor(**params)
                model.fit(X_np, y_np, sample_weight=w_np)
                self._models.append(model)

        self._metadata = TailModel(
            spec=self._spec,
            n_features=X_np.shape[1],
            feature_names=list(X.columns) if hasattr(X, 'columns') else [f'x{i}' for i in range(X_np.shape[1])],
            n_training_rows=X_np.shape[0],
            catboost_params=base_params,
            fix_crossing=self._fix_crossing,
        )
        self._is_fitted = True
        return self

    # ------------------------------------------------------------------
    # Predict
    # ------------------------------------------------------------------

    def predict(self, X: pl.DataFrame) -> pl.DataFrame:
        """
        Predict quantile (or expectile) values for each row.

        Parameters
        ----------
        X:
            Feature matrix, same columns as used in fit().

        Returns
        -------
        Polars DataFrame with one column per quantile level,
        named q_0.5, q_0.9, etc. Columns are in the same order
        as the quantiles passed to __init__.
        """
        if not self._is_fitted:
            raise RuntimeError("Model is not fitted. Call fit() first.")

        X_np = _to_numpy(X)

        if self._spec.mode == "quantile":
            raw = self._models[0].predict(X_np)
            # CatBoost MultiQuantile returns shape (n_samples, n_quantiles)
            if raw.ndim == 1:
                raw = raw.reshape(-1, 1)
        else:
            # Stack per-alpha model predictions
            parts = [m.predict(X_np).reshape(-1, 1) for m in self._models]
            raw = np.hstack(parts)

        if self._fix_crossing:
            raw = _apply_isotonic(raw)

        return pl.DataFrame(
            {name: raw[:, i] for i, name in enumerate(self._spec.column_names)}
        )

    # ------------------------------------------------------------------
    # TVaR
    # ------------------------------------------------------------------

    def predict_tvar(
        self,
        X: pl.DataFrame,
        alpha: float,
    ) -> pl.Series:
        """
        Estimate Tail Value at Risk (TVaR) at level alpha for each risk.

        TVaR_alpha = E[Y | Y > VaR_alpha(Y)].

        Uses trapezoidal integration over the stored quantile levels above
        alpha. For best accuracy, include several levels above your alpha
        (e.g. 0.9, 0.95, 0.99 for TVaR at alpha=0.9).

        The module-level function ``per_risk_tvar`` is the recommended
        interface — it is equivalent to this method and is exported at the
        package level.

        Parameters
        ----------
        X:
            Feature matrix.
        alpha:
            Probability threshold, e.g. 0.95 for TVaR_95.

        Returns
        -------
        Polars Series of TVaR estimates, one per row of X.

        Raises
        ------
        ValueError
            If the model has no fitted quantile levels above alpha.
        """
        if not self._is_fitted:
            raise RuntimeError("Model is not fitted. Call fit() first.")
        if not (0.0 < alpha < 1.0):
            raise ValueError(f"alpha must be in (0, 1), got {alpha}")

        above_any = [q for q in self._spec.quantiles if q > alpha]
        if not above_any:
            raise ValueError(
                f"Model has no quantile levels above alpha={alpha}. "
                "Add higher quantiles (e.g. 0.99) or call predict_tvar "
                "with a lower alpha."
            )

        max_quantile = max(self._spec.quantiles)
        if alpha >= max_quantile:
            import warnings
            warnings.warn(
                f"alpha={alpha} is at or above the highest modelled quantile level "
                f"({max_quantile}). The TVaR estimate will be unreliable because "
                "trapezoidal integration requires quantile levels well above alpha. "
                f"Add quantile levels above {alpha} when fitting the model.",
                UserWarning,
                stacklevel=2,
            )

        col_names = [f"q_{q}" for q in above_any]
        preds_df = self.predict(X)
        tail_vals = np.stack([preds_df[c].to_numpy() for c in col_names], axis=1)
        above_arr = np.array(above_any)

        # Trapezoidal integration: weight each quantile prediction by its interval width.
        # Use Q(alpha) as the left boundary when available (most accurate and guarantees
        # TVaR >= VaR). Fall back to flat extrapolation with Q(first_above) otherwise.
        at_alpha = [q for q in self._spec.quantiles if q == alpha]
        if at_alpha:
            boundary_vals = preds_df[f"q_{at_alpha[0]}"].to_numpy().reshape(-1, 1)
        else:
            boundary_vals = tail_vals[:, 0:1]

        levels = np.concatenate([[alpha], above_arr])
        values_with_boundary = np.concatenate([boundary_vals, tail_vals], axis=1)
        integral = _trapezoid(values_with_boundary, levels, axis=1)
        tvar = integral / (1.0 - alpha)

        # VaR at alpha — clip TVaR to be >= VaR for numerical safety
        at_or_below = [q for q in self._spec.quantiles if q <= alpha]
        if at_or_below:
            var_vals = preds_df[f"q_{at_or_below[-1]}"].to_numpy()
        else:
            var_vals = preds_df[f"q_{self._spec.quantiles[0]}"].to_numpy()
        tvar = np.maximum(tvar, var_vals)

        return pl.Series("tvar", tvar)

    # ------------------------------------------------------------------
    # Calibration report
    # ------------------------------------------------------------------

    def calibration_report(
        self,
        X: pl.DataFrame,
        y: pl.Series,
    ) -> dict[str, Any]:
        """
        Compute calibration statistics on a held-out validation set.

        Returns a dict containing:
        - 'coverage': dict mapping each quantile to observed coverage fraction.
          A well-calibrated q_0.9 model should have coverage close to 0.9.
        - 'pinball_loss': dict mapping each quantile to its pinball loss.
        - 'mean_pinball_loss': average pinball loss across all quantile levels.

        Parameters
        ----------
        X:
            Validation feature matrix.
        y:
            Validation targets.

        Returns
        -------
        dict with keys 'coverage', 'pinball_loss', 'mean_pinball_loss'.
        """
        from ._calibration import coverage_check, pinball_loss

        preds = self.predict(X)
        y_np = _series_to_numpy(y)

        coverage: dict[str, float] = {}
        pinball: dict[str, float] = {}

        for q, col in zip(self._spec.quantiles, self._spec.column_names):
            q_pred = preds[col].to_numpy()
            coverage[col] = float(np.mean(y_np <= q_pred))
            pinball[col] = pinball_loss(y, preds[col], alpha=q)

        return {
            "coverage": coverage,
            "pinball_loss": pinball,
            "mean_pinball_loss": float(np.mean(list(pinball.values()))),
        }
