r"""Fairness-aware gradient boosting classifier.

Provides ``FairClassifier``, a wrapper around ``PerpetualBooster`` that adds
an in-processing fairness regularization penalty to the log-loss gradient.
Two fairness criteria are supported:

* **Demographic Parity** — penalizes correlation between predictions and a
  sensitive attribute :math:`S`.
* **Equalized Odds** — penalizes correlation between predictions and
  :math:`S` *conditionally* within each class of the true label :math:`Y`.

The objective is:

.. math::

    L = \\text{LogLoss} + \\lambda \\cdot \\text{Penalty}(\\hat{Y}, S)
"""

from typing import Optional

import numpy as np
from typing_extensions import Self

from perpetual.booster import PerpetualBooster

from .perpetual import FairnessObjective

_VALID_FAIRNESS_TYPES = ("demographic_parity", "equalized_odds")


def stable_sigmoid(x):
    """Numerically stable sigmoid function."""
    x = np.asarray(x)
    out = np.empty_like(x, dtype=float)
    pos_mask = x >= 0
    neg_mask = ~pos_mask
    out[pos_mask] = 1.0 / (1.0 + np.exp(-x[pos_mask]))
    exp_x = np.exp(x[neg_mask])
    out[neg_mask] = exp_x / (1.0 + exp_x)
    return out


# ---------------------------------------------------------------------------
# FairClassifier
# ---------------------------------------------------------------------------


class FairClassifier:
    r"""Fairness-aware gradient boosting classifier.

    Wraps a ``PerpetualBooster`` with an in-processing fairness penalty that
    regularizes the log-loss gradient to reduce dependence of predictions on
    a sensitive attribute.

    Parameters
    ----------
    sensitive_feature : int
        Column index of the sensitive attribute in *X*.  The column must be
        binary (0 or 1).
    fairness_type : str, default="demographic_parity"
        Fairness criterion.  One of:

        - ``"demographic_parity"`` — penalize overall disparity.
        - ``"equalized_odds"`` — penalize disparity within each label class.
    lam : float, default=1.0
        Strength of the fairness penalty (:math:`\\lambda`).
    budget : float, default=0.5
        Fitting budget forwarded to ``PerpetualBooster``.
    **kwargs
        Additional keyword arguments forwarded to ``PerpetualBooster``.

    Attributes
    ----------
    feature_importances_ : ndarray of shape (n_features,)
        Feature importances from the fitted model.

    Examples
    --------
    >>> from perpetual.fairness import FairClassifier
    >>> import numpy as np
    >>> X = np.column_stack([np.random.randn(200, 3),
    ...                      np.random.binomial(1, 0.5, 200)])
    >>> y = (X[:, 0] > 0).astype(float)
    >>> clf = FairClassifier(sensitive_feature=3, lam=2.0)
    >>> clf.fit(X, y)  # doctest: +SKIP
    >>> probs = clf.predict_proba(X)  # doctest: +SKIP

    Notes
    -----
    The fairness penalty is applied only through the gradient; the reported
    loss is standard log-loss.  This mirrors the Rust ``FairnessObjective``
    implementation.
    """

    def __init__(
        self,
        sensitive_feature: int,
        fairness_type: str = "demographic_parity",
        lam: float = 1.0,
        budget: float = 0.5,
        **kwargs,
    ):
        if fairness_type not in _VALID_FAIRNESS_TYPES:
            raise ValueError(
                f"fairness_type must be one of {_VALID_FAIRNESS_TYPES}, "
                f"got {fairness_type!r}."
            )
        self.sensitive_feature = sensitive_feature
        self.fairness_type = fairness_type
        self.lam = lam
        self.budget = budget
        self._kwargs = kwargs
        self.feature_importances_: Optional[np.ndarray] = None

    def fit(self, X, y) -> Self:
        """Fit the fair classifier.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Feature matrix.  The column at index ``sensitive_feature`` must
            contain binary (0/1) values.
        y : array-like of shape (n_samples,)
            Binary target variable (0 or 1).

        Returns
        -------
        self
            Fitted estimator.
        """
        X_arr = np.asarray(X, dtype=float)
        y_arr = np.asarray(y, dtype=float)
        # Ensure int32 for Rust compatibility
        sensitive_attr = X_arr[:, self.sensitive_feature].astype(np.int32)

        objective = FairnessObjective(sensitive_attr, self.lam, self.fairness_type)

        self._model = PerpetualBooster(
            budget=self.budget, objective=objective, **self._kwargs
        )
        self._model.fit(X_arr, y_arr)
        self.feature_importances_ = self._model.feature_importances_
        return self

    def predict(self, X) -> np.ndarray:
        """Predict class labels.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Feature matrix.

        Returns
        -------
        ndarray of shape (n_samples,)
            Predicted class labels (0 or 1).
        """
        scores = self._model.predict(X)
        return (scores > 0).astype(int)

    def predict_proba(self, X) -> np.ndarray:
        """Predict class probabilities.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Feature matrix.

        Returns
        -------
        ndarray of shape (n_samples, 2)
            Predicted probabilities for class 0 and class 1.
        """
        scores = self._model.predict(X)
        p1 = stable_sigmoid(scores)
        return np.column_stack([1.0 - p1, p1])

    def predict_contributions(
        self, X, method: str = "Average", parallel: Optional[bool] = None
    ) -> np.ndarray:
        """Predict feature contributions (SHAP-like values) for new data.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Input features.
        method : str, default="Average"
            Method to calculate contributions.
        parallel : bool, optional
            Whether to run prediction in parallel.

        Returns
        -------
        ndarray of shape (n_samples, n_features + 1)
            Feature contributions.
        """
        return self._model.predict_contributions(X, method=method, parallel=parallel)

    def decision_function(self, X) -> np.ndarray:
        """Return raw log-odds scores.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Feature matrix.

        Returns
        -------
        ndarray of shape (n_samples,)
            Raw boosted scores (log-odds).
        """
        return self._model.predict(X)
