"""
Fair decision tree classifier using AUC-based splitting criterion.

This module provides FairDecisionTreeClassifier, which always uses
criterion='auc'. No other criteria (gini, entropy, regression) are supported.
"""

# Authors: The scikit-learn developers
# SPDX-License-Identifier: BSD-3-Clause

import copy
import numbers
from abc import ABCMeta, abstractmethod
from math import ceil
from numbers import Integral, Real

import numpy as np
from scipy.sparse import issparse

from fair_trees.utils.metadata_routing import UNUSED as _METADATA_ROUTING_UNUSED

from ..base import (
    BaseEstimator,
    ClassifierMixin,
    MultiOutputMixin,
    _fit_context,
    clone,
    is_classifier,
)
from ..utils import Bunch, check_random_state, compute_sample_weight
from ..utils._param_validation import Hidden, Interval, RealNotInt, StrOptions
from ..utils.multiclass import check_classification_targets
from ..utils.validation import (
    _assert_all_finite_element_wise,
    _check_n_features,
    _check_sample_weight,
    assert_all_finite,
    check_is_fitted,
    validate_data,
)
from . import _criterion, _splitter, _tree
from ._criterion import Criterion
from ._splitter import Splitter
from ._tree import (
    BestFirstTreeBuilder,
    DepthFirstTreeBuilder,
    Tree,
    _build_pruned_tree_ccp,
    _build_pruned_tree_ccp_auc,
    ccp_pruning_path,
    ccp_pruning_path_auc,
)
from ._utils import _any_isnan_axis0

__all__ = [
    "FairDecisionTreeClassifier",
]


# =============================================================================
# Types and constants
# =============================================================================

DTYPE = _tree.DTYPE
DOUBLE = _tree.DOUBLE

CRITERIA_CLF = {
    "auc": _criterion.AUC,
}

DENSE_SPLITTERS = {"best": _splitter.BestSplitter, "random": _splitter.RandomSplitter}

SPARSE_SPLITTERS = {
    "best": _splitter.BestSparseSplitter,
    "random": _splitter.RandomSparseSplitter,
}

# =============================================================================
# Base decision tree
# =============================================================================


class BaseDecisionTree(MultiOutputMixin, BaseEstimator, metaclass=ABCMeta):
    """Base class for decision trees.

    Warning: This class should not be used directly.
    Use derived classes instead.
    """

    # "check_input" is used for optimisation and isn't something to be passed
    # around in a pipeline.
    __metadata_request__predict = {"check_input": _METADATA_ROUTING_UNUSED}

    _parameter_constraints: dict = {
        "splitter": [StrOptions({"best", "random"})],
        "max_depth": [Interval(Integral, 1, None, closed="left"), None],
        "min_samples_split": [
            Interval(Integral, 2, None, closed="left"),
            Interval(RealNotInt, 0.0, 1.0, closed="right"),
        ],
        "min_samples_leaf": [
            Interval(Integral, 1, None, closed="left"),
            Interval(RealNotInt, 0.0, 1.0, closed="neither"),
        ],
        "min_weight_fraction_leaf": [Interval(Real, 0.0, 0.5, closed="both")],
        "max_features": [
            Interval(Integral, 1, None, closed="left"),
            Interval(RealNotInt, 0.0, 1.0, closed="right"),
            StrOptions({"sqrt", "log2"}),
            None,
        ],
        "random_state": ["random_state"],
        "max_leaf_nodes": [Interval(Integral, 2, None, closed="left"), None],
        "min_impurity_decrease": [Interval(Real, 0.0, None, closed="left")],
        "ccp_alpha": [Interval(Real, 0.0, None, closed="left")],
        "monotonic_cst": ["array-like", None],
    }

    @abstractmethod
    def __init__(
        self,
        *,
        splitter,
        max_depth,
        min_samples_split,
        min_samples_leaf,
        min_weight_fraction_leaf,
        max_features,
        max_leaf_nodes,
        random_state,
        min_impurity_decrease,
        class_weight=None,
        ccp_alpha=0.0,
        monotonic_cst=None,
    ):
        self.splitter = splitter
        self.max_depth = max_depth
        self.min_samples_split = min_samples_split
        self.min_samples_leaf = min_samples_leaf
        self.min_weight_fraction_leaf = min_weight_fraction_leaf
        self.max_features = max_features
        self.max_leaf_nodes = max_leaf_nodes
        self.random_state = random_state
        self.min_impurity_decrease = min_impurity_decrease
        self.class_weight = class_weight
        self.ccp_alpha = ccp_alpha
        self.monotonic_cst = monotonic_cst

    def get_depth(self):
        """Return the depth of the decision tree.

        The depth of a tree is the maximum distance between the root
        and any leaf.

        Returns
        -------
        self.tree_.max_depth : int
            The maximum depth of the tree.
        """
        check_is_fitted(self)
        return self.tree_.max_depth

    def get_n_leaves(self):
        """Return the number of leaves of the decision tree.

        Returns
        -------
        self.tree_.n_leaves : int
            Number of leaves.
        """
        check_is_fitted(self)
        return self.tree_.n_leaves

    def _support_missing_values(self, X):
        return (
            not issparse(X)
            and self.__sklearn_tags__().input_tags.allow_nan
            and self.monotonic_cst is None
        )

    def _compute_missing_values_in_feature_mask(self, X, estimator_name=None):
        """Return boolean mask denoting if there are missing values for each feature.

        This method also ensures that X is finite.

        Parameter
        ---------
        X : array-like of shape (n_samples, n_features), dtype=DOUBLE
            Input data.

        estimator_name : str or None, default=None
            Name to use when raising an error. Defaults to the class name.

        Returns
        -------
        missing_values_in_feature_mask : ndarray of shape (n_features,), or None
            Missing value mask. If missing values are not supported or there
            are no missing values, return None.
        """
        estimator_name = estimator_name or self.__class__.__name__
        common_kwargs = dict(estimator_name=estimator_name, input_name="X")

        if not self._support_missing_values(X):
            assert_all_finite(X, **common_kwargs)
            return None

        with np.errstate(over="ignore"):
            overall_sum = np.sum(X)

        if not np.isfinite(overall_sum):
            # Raise a ValueError in case of the presence of an infinite element.
            _assert_all_finite_element_wise(X, xp=np, allow_nan=True, **common_kwargs)

        # If the sum is not nan, then there are no missing values
        if not np.isnan(overall_sum):
            return None

        missing_values_in_feature_mask = _any_isnan_axis0(X)
        return missing_values_in_feature_mask

    def _fit(
        self,
        X,
        y,
        sample_weight=None,
        check_input=True,
        missing_values_in_feature_mask=None,
        Z=None,
    ):
        random_state = check_random_state(self.random_state)

        if check_input:
            check_X_params = dict(
                dtype=DTYPE, accept_sparse="csc", ensure_all_finite=False
            )
            check_y_params = dict(ensure_2d=False, dtype=None)
            X, y = validate_data(
                self, X, y, validate_separately=(check_X_params, check_y_params)
            )

            missing_values_in_feature_mask = (
                self._compute_missing_values_in_feature_mask(X)
            )
            if issparse(X):
                X.sort_indices()

                if X.indices.dtype != np.intc or X.indptr.dtype != np.intc:
                    raise ValueError(
                        "No support for np.int64 index based sparse matrices"
                    )

        # Determine output settings
        n_samples, self.n_features_in_ = X.shape
        is_classification = is_classifier(self)

        y = np.atleast_1d(y)
        expanded_class_weight = None

        if y.ndim == 1:
            y = np.reshape(y, (-1, 1))

        self.n_outputs_ = y.shape[1]

        if is_classification:
            check_classification_targets(y)
            y = np.copy(y)

            self.classes_ = []
            self.n_classes_ = []

            if self.class_weight is not None:
                y_original = np.copy(y)

            y_encoded = np.zeros(y.shape, dtype=int)
            for k in range(self.n_outputs_):
                classes_k, y_encoded[:, k] = np.unique(y[:, k], return_inverse=True)
                self.classes_.append(classes_k)
                self.n_classes_.append(classes_k.shape[0])
            y = y_encoded

            if self.class_weight is not None:
                expanded_class_weight = compute_sample_weight(
                    self.class_weight, y_original
                )

            self.n_classes_ = np.array(self.n_classes_, dtype=np.intp)

        if getattr(y, "dtype", None) != DOUBLE or not y.flags.contiguous:
            y = np.ascontiguousarray(y, dtype=DOUBLE)

        max_depth = np.iinfo(np.int32).max if self.max_depth is None else self.max_depth

        if isinstance(self.min_samples_leaf, numbers.Integral):
            min_samples_leaf = self.min_samples_leaf
        else:  # float
            min_samples_leaf = ceil(self.min_samples_leaf * n_samples)

        if isinstance(self.min_samples_split, numbers.Integral):
            min_samples_split = self.min_samples_split
        else:  # float
            min_samples_split = ceil(self.min_samples_split * n_samples)
            min_samples_split = max(2, min_samples_split)

        min_samples_split = max(min_samples_split, 2 * min_samples_leaf)

        if isinstance(self.max_features, str):
            if self.max_features == "sqrt":
                max_features = max(1, int(np.sqrt(self.n_features_in_)))
            elif self.max_features == "log2":
                max_features = max(1, int(np.log2(self.n_features_in_)))
        elif self.max_features is None:
            max_features = self.n_features_in_
        elif isinstance(self.max_features, numbers.Integral):
            max_features = self.max_features
        else:  # float
            if self.max_features > 0.0:
                max_features = max(1, int(self.max_features * self.n_features_in_))
            else:
                max_features = 0

        self.max_features_ = max_features

        max_leaf_nodes = -1 if self.max_leaf_nodes is None else self.max_leaf_nodes

        if len(y) != n_samples:
            raise ValueError(
                "Number of labels=%d does not match number of samples=%d"
                % (len(y), n_samples)
            )

        if sample_weight is not None:
            sample_weight = _check_sample_weight(sample_weight, X, dtype=DOUBLE)

        if expanded_class_weight is not None:
            if sample_weight is not None:
                sample_weight = sample_weight * expanded_class_weight
            else:
                sample_weight = expanded_class_weight

        # Set min_weight_leaf from min_weight_fraction_leaf
        if sample_weight is None:
            min_weight_leaf = self.min_weight_fraction_leaf * n_samples
        else:
            min_weight_leaf = self.min_weight_fraction_leaf * np.sum(sample_weight)

        # Build tree — criterion is always AUC
        # Check if we need biased criterion (Z-based regularization)
        use_biased = (Z is not None and
                     hasattr(self, 'theta') and self.theta > 0)

        if use_biased:
            # Validate and prepare Z
            from ..utils.validation import check_array
            Z_array = check_array(
                Z, dtype=DOUBLE,
                ensure_2d=False, ensure_all_finite=True,
                copy=True
            )

            # Ensure it's a numpy array
            if not isinstance(Z_array, np.ndarray):
                Z_array = np.asarray(Z_array, dtype=DOUBLE)

            # Make a copy to avoid modifying the original
            Z_array = np.array(Z_array, dtype=DOUBLE, copy=True)

            if Z_array.shape[0] != n_samples:
                raise ValueError(
                    f"Z must have the same number of samples as X. "
                    f"Got Z.shape[0]={Z_array.shape[0]}, expected {n_samples}"
                )

            # Ensure Z is 2D
            if Z_array.ndim == 1:
                Z_array = Z_array.reshape(-1, 1)

            # Encode Z categories and get class counts
            n_z_outputs = Z_array.shape[1]
            Z_encoded = np.zeros((n_samples, n_z_outputs), dtype=DOUBLE)
            n_z_classes = np.zeros(n_z_outputs, dtype=np.intp)

            for k in range(n_z_outputs):
                z_classes_k, Z_encoded[:, k] = np.unique(
                    Z_array[:, k], return_inverse=True
                )
                n_z_classes[k] = z_classes_k.shape[0]

            # Ensure Z_encoded is contiguous
            if not Z_encoded.flags.c_contiguous:
                Z_encoded = np.ascontiguousarray(Z_encoded, dtype=DOUBLE)

            # Convert Z_agg to mode
            Z_agg_mode = 0 if self.Z_agg == "mean" else 1

            # AUC-based objective: (1-theta)*AUC_y - theta*AUC_z
            criterion = _criterion.create_biased_auc(
                self.n_outputs_, self.n_classes_,
                Z_encoded, n_z_classes, self.theta, Z_agg_mode
            )
        else:
            # Standard AUC criterion without bias
            criterion = CRITERIA_CLF["auc"](
                self.n_outputs_, self.n_classes_
            )

        SPLITTERS = SPARSE_SPLITTERS if issparse(X) else DENSE_SPLITTERS

        splitter = self.splitter
        if self.monotonic_cst is None:
            monotonic_cst = None
        else:
            if self.n_outputs_ > 1:
                raise ValueError(
                    "Monotonicity constraints are not supported with multiple outputs."
                )
            monotonic_cst = np.asarray(self.monotonic_cst)
            if monotonic_cst.shape[0] != X.shape[1]:
                raise ValueError(
                    "monotonic_cst has shape {} but the input data "
                    "X has {} features.".format(monotonic_cst.shape[0], X.shape[1])
                )
            valid_constraints = np.isin(monotonic_cst, (-1, 0, 1))
            if not np.all(valid_constraints):
                unique_constaints_value = np.unique(monotonic_cst)
                raise ValueError(
                    "monotonic_cst must be None or an array-like of -1, 0 or 1, but"
                    f" got {unique_constaints_value}"
                )
            monotonic_cst = np.asarray(monotonic_cst, dtype=np.int8)
            if is_classifier(self):
                if self.n_classes_[0] > 2:
                    raise ValueError(
                        "Monotonicity constraints are not supported with multiclass "
                        "classification"
                    )
                monotonic_cst *= -1

        if not isinstance(self.splitter, Splitter):
            splitter = SPLITTERS[self.splitter](
                criterion,
                self.max_features_,
                min_samples_leaf,
                min_weight_leaf,
                random_state,
                monotonic_cst,
            )

        self.tree_ = Tree(self.n_features_in_, self.n_classes_, self.n_outputs_)

        # Use BestFirst if max_leaf_nodes given; use DepthFirst otherwise
        if max_leaf_nodes < 0:
            builder = DepthFirstTreeBuilder(
                splitter,
                min_samples_split,
                min_samples_leaf,
                min_weight_leaf,
                max_depth,
                self.min_impurity_decrease,
            )
        else:
            builder = BestFirstTreeBuilder(
                splitter,
                min_samples_split,
                min_samples_leaf,
                min_weight_leaf,
                max_depth,
                max_leaf_nodes,
                self.min_impurity_decrease,
            )

        builder.build(self.tree_, X, y, sample_weight, missing_values_in_feature_mask)

        if self.n_outputs_ == 1 and is_classifier(self):
            self.n_classes_ = self.n_classes_[0]
            self.classes_ = self.classes_[0]

        self._prune_tree()

        return self

    def _validate_X_predict(self, X, check_input):
        """Validate the training data on predict (probabilities)."""
        if check_input:
            if self._support_missing_values(X):
                ensure_all_finite = "allow-nan"
            else:
                ensure_all_finite = True
            X = validate_data(
                self,
                X,
                dtype=DTYPE,
                accept_sparse="csr",
                reset=False,
                ensure_all_finite=ensure_all_finite,
            )
            if issparse(X) and (
                X.indices.dtype != np.intc or X.indptr.dtype != np.intc
            ):
                raise ValueError("No support for np.int64 index based sparse matrices")
        else:
            # The number of features is checked regardless of `check_input`
            _check_n_features(self, X, reset=False)
        return X

    def predict(self, X, check_input=True):
        """Predict class for X.

        Parameters
        ----------
        X : {array-like, sparse matrix} of shape (n_samples, n_features)
            The input samples.

        check_input : bool, default=True
            Allow to bypass several input checking.

        Returns
        -------
        y : array-like of shape (n_samples,) or (n_samples, n_outputs)
            The predicted classes.
        """
        check_is_fitted(self)
        X = self._validate_X_predict(X, check_input)
        proba = self.tree_.predict(X)
        n_samples = X.shape[0]

        # Classification
        if self.n_outputs_ == 1:
            return self.classes_.take(np.argmax(proba, axis=1), axis=0)
        else:
            class_type = self.classes_[0].dtype
            predictions = np.zeros((n_samples, self.n_outputs_), dtype=class_type)
            for k in range(self.n_outputs_):
                predictions[:, k] = self.classes_[k].take(
                    np.argmax(proba[:, k], axis=1), axis=0
                )
            return predictions

    def apply(self, X, check_input=True):
        """Return the index of the leaf that each sample is predicted as.

        Parameters
        ----------
        X : {array-like, sparse matrix} of shape (n_samples, n_features)
            The input samples.

        check_input : bool, default=True
            Allow to bypass several input checking.

        Returns
        -------
        X_leaves : array-like of shape (n_samples,)
            For each datapoint x in X, return the index of the leaf x
            ends up in.
        """
        check_is_fitted(self)
        X = self._validate_X_predict(X, check_input)
        return self.tree_.apply(X)

    def decision_path(self, X, check_input=True):
        """Return the decision path in the tree.

        Parameters
        ----------
        X : {array-like, sparse matrix} of shape (n_samples, n_features)
            The input samples.

        check_input : bool, default=True
            Allow to bypass several input checking.

        Returns
        -------
        indicator : sparse matrix of shape (n_samples, n_nodes)
            Return a node indicator CSR matrix where non zero elements
            indicates that the samples goes through the nodes.
        """
        X = self._validate_X_predict(X, check_input)
        return self.tree_.decision_path(X)

    def _prune_tree(self):
        """Prune tree using Minimal Cost-Complexity Pruning."""
        check_is_fitted(self)

        if self.ccp_alpha == 0.0:
            return

        # build pruned tree
        n_classes = np.atleast_1d(self.n_classes_)
        pruned_tree = Tree(self.n_features_in_, n_classes, self.n_outputs_)

        # Always use AUC-specific pruning
        _build_pruned_tree_ccp_auc(pruned_tree, self.tree_, self.ccp_alpha)

        self.tree_ = pruned_tree

    def cost_complexity_pruning_path(self, X, y, Z=None, sample_weight=None):
        """Compute the pruning path during Minimal Cost-Complexity Pruning.

        Parameters
        ----------
        X : {array-like, sparse matrix} of shape (n_samples, n_features)
            The training input samples.

        y : array-like of shape (n_samples,) or (n_samples, n_outputs)
            The target values (class labels) as integers or strings.

        Z : array-like of shape (n_samples,) or (n_samples, n_z_features), default=None
            Protected/sensitive attributes for fairness-aware splitting.

        sample_weight : array-like of shape (n_samples,), default=None
            Sample weights.

        Returns
        -------
        ccp_path : :class:`~sklearn.utils.Bunch`
            Dictionary-like object, with the following attributes.

            ccp_alphas : ndarray
                Effective alphas of subtree during pruning.

            gains : ndarray
                Sum of the gains of the subtree leaves for the
                corresponding alpha value in ``ccp_alphas``.
        """
        est = clone(self).set_params(ccp_alpha=0.0)
        est.fit(X, y, Z, sample_weight=sample_weight)
        out = ccp_pruning_path_auc(est.tree_)
        return Bunch(ccp_alphas=out["ccp_alphas"], gains=out["gains"])

    @property
    def feature_importances_(self):
        """Return the feature importances.

        The importance of a feature is computed as the sum of weighted
        gains (weighted_n_samples * gain) for each split on that feature.

        Returns
        -------
        feature_importances_ : ndarray of shape (n_features,)
            Normalized total reduction of criteria by feature.
        """
        check_is_fitted(self)
        # Always use AUC-based importances
        return self.tree_.compute_feature_importances(use_auc_importances=True)

    def __sklearn_tags__(self):
        tags = super().__sklearn_tags__()
        tags.input_tags.sparse = True
        return tags


# =============================================================================
# Public estimator
# =============================================================================


class FairDecisionTreeClassifier(ClassifierMixin, BaseDecisionTree):
    """A fair decision tree classifier using AUC-based splitting.

    This classifier always uses the AUC criterion for splitting. It supports
    fairness-aware splitting through the theta and Z_agg parameters.

    Parameters
    ----------
    theta : float, default=0.0
        Fairness penalty weight. When theta > 0, the split criterion becomes:
        (1-theta)*performance - theta*fairness_violation.
        Range: [0, 1]. Set to 0 to disable fairness constraints.

    Z_agg : {"mean", "max"}, default="max"
        How to aggregate fairness violations across sensitive attributes.
        Only used when Z has multiple columns and theta > 0.
        - "mean": Average disparity across attributes
        - "max": Worst-case disparity across attributes

    splitter : {"best", "random"}, default="best"
        The strategy used to choose the split at each node.

    max_depth : int, default=None
        The maximum depth of the tree. If None, then nodes are expanded until
        all leaves are pure or until all leaves contain less than
        min_samples_split samples.

    min_samples_split : int or float, default=2
        The minimum number of samples required to split an internal node.

    min_samples_leaf : int or float, default=1
        The minimum number of samples required to be at a leaf node.

    min_weight_fraction_leaf : float, default=0.0
        The minimum weighted fraction of the sum total of weights required
        to be at a leaf node.

    max_features : int, float or {"sqrt", "log2"}, default=None
        The number of features to consider when looking for the best split.

    random_state : int, RandomState instance or None, default=None
        Controls the randomness of the estimator.

    max_leaf_nodes : int, default=None
        Grow a tree with ``max_leaf_nodes`` in best-first fashion.

    min_impurity_decrease : float, default=0.0
        A node will be split if this split induces a decrease of the impurity
        greater than or equal to this value.

    class_weight : dict, list of dict or "balanced", default=None
        Weights associated with classes.

    ccp_alpha : non-negative float, default=0.0
        Complexity parameter used for Minimal Cost-Complexity Pruning.

    monotonic_cst : array-like of int of shape (n_features), default=None
        Indicates the monotonicity constraint to enforce on each feature.

    Attributes
    ----------
    classes_ : ndarray of shape (n_classes,) or list of ndarray
        The classes labels.

    feature_importances_ : ndarray of shape (n_features,)
        The AUC-based feature importances.

    max_features_ : int
        The inferred value of max_features.

    n_classes_ : int or list of int
        The number of classes.

    n_features_in_ : int
        Number of features seen during :term:`fit`.

    feature_names_in_ : ndarray of shape (`n_features_in_`,)
        Names of features seen during :term:`fit`.

    n_outputs_ : int
        The number of outputs when ``fit`` is performed.

    tree_ : Tree instance
        The underlying Tree object.

    Examples
    --------
    >>> from sklearn.datasets import load_iris
    >>> from sklearn.model_selection import cross_val_score
    >>> from sklearn.tree import FairDecisionTreeClassifier
    >>> clf = FairDecisionTreeClassifier(random_state=0)
    >>> iris = load_iris()
    >>> cross_val_score(clf, iris.data, iris.target, cv=10)
    ...                             # doctest: +SKIP
    """

    # "check_input" is used for optimisation and isn't something to be passed
    # around in a pipeline.
    __metadata_request__predict_proba = {"check_input": _METADATA_ROUTING_UNUSED}
    __metadata_request__fit = {"check_input": _METADATA_ROUTING_UNUSED}

    _parameter_constraints: dict = {
        **BaseDecisionTree._parameter_constraints,
        "class_weight": [dict, list, StrOptions({"balanced"}), None],
        "theta": [Interval(Real, 0, 1, closed="both")],
        "Z_agg": [StrOptions({"mean", "max"})],
    }

    def __init__(
        self,
        *,
        splitter="best",
        max_depth=None,
        min_samples_split=2,
        min_samples_leaf=1,
        min_weight_fraction_leaf=0.0,
        max_features=None,
        random_state=None,
        max_leaf_nodes=None,
        min_impurity_decrease=0.0,
        class_weight=None,
        ccp_alpha=0.0,
        monotonic_cst=None,
        theta=0.0,
        Z_agg="max",
    ):
        super().__init__(
            splitter=splitter,
            max_depth=max_depth,
            min_samples_split=min_samples_split,
            min_samples_leaf=min_samples_leaf,
            min_weight_fraction_leaf=min_weight_fraction_leaf,
            max_features=max_features,
            max_leaf_nodes=max_leaf_nodes,
            class_weight=class_weight,
            random_state=random_state,
            min_impurity_decrease=min_impurity_decrease,
            monotonic_cst=monotonic_cst,
            ccp_alpha=ccp_alpha,
        )

        # Store Z-related hyperparameters for bias regularization
        self.theta = theta
        self.Z_agg = Z_agg

    @_fit_context(prefer_skip_nested_validation=True)
    def fit(self, X, y, Z=None, sample_weight=None, check_input=True):
        """Build a fair decision tree classifier from the training set (X, y).

        Parameters
        ----------
        X : {array-like, sparse matrix} of shape (n_samples, n_features)
            The training input samples.

        y : array-like of shape (n_samples,) or (n_samples, n_outputs)
            The target values (class labels) as integers or strings.

        Z : array-like of shape (n_samples,) or (n_samples, n_z_features), default=None
            Protected/sensitive attributes for fairness-aware splitting.
            When provided with theta > 0, the tree will use a biased splitting
            criterion that balances accuracy with fairness.

        sample_weight : array-like of shape (n_samples,), default=None
            Sample weights.

        check_input : bool, default=True
            Allow to bypass several input checking.

        Returns
        -------
        self : FairDecisionTreeClassifier
            Fitted estimator.
        """

        super()._fit(
            X,
            y,
            sample_weight=sample_weight,
            check_input=check_input,
            Z=Z,
        )
        return self

    def predict_proba(self, X, check_input=True):
        """Predict class probabilities of the input samples X.

        Parameters
        ----------
        X : {array-like, sparse matrix} of shape (n_samples, n_features)
            The input samples.

        check_input : bool, default=True
            Allow to bypass several input checking.

        Returns
        -------
        proba : ndarray of shape (n_samples, n_classes) or list of n_outputs
            such arrays if n_outputs > 1
            The class probabilities of the input samples.
        """
        check_is_fitted(self)
        X = self._validate_X_predict(X, check_input)
        proba = self.tree_.predict(X)

        if self.n_outputs_ == 1:
            return proba[:, : self.n_classes_]
        else:
            all_proba = []
            for k in range(self.n_outputs_):
                proba_k = proba[:, k, : self.n_classes_[k]]
                all_proba.append(proba_k)
            return all_proba

    def predict_log_proba(self, X):
        """Predict class log-probabilities of the input samples X.

        Parameters
        ----------
        X : {array-like, sparse matrix} of shape (n_samples, n_features)
            The input samples.

        Returns
        -------
        proba : ndarray of shape (n_samples, n_classes) or list of n_outputs
            such arrays if n_outputs > 1
            The class log-probabilities of the input samples.
        """
        proba = self.predict_proba(X)

        if self.n_outputs_ == 1:
            return np.log(proba)

        else:
            for k in range(self.n_outputs_):
                proba[k] = np.log(proba[k])

            return proba

    def __sklearn_tags__(self):
        tags = super().__sklearn_tags__()
        # AUC criterion does not support NaN
        tags.classifier_tags.multi_label = True
        tags.input_tags.allow_nan = False
        return tags
