# Author: Arjun Roy (arjun.roy@unibw.de, arjunroyihrpa@gmail.com) https://orcid.org/0000-0002-4279-9442 
# Apache License Version 2.0
import numpy as np
from time import time
from functools import partial
from numbers import Integral, Real
import itertools
from fairbench import v2 as fb

##sklearn ensemble
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.ensemble._hist_gradient_boosting.grower import TreeGrower
from sklearn.ensemble._hist_gradient_boosting._gradient_boosting import _update_raw_predictions
from sklearn.ensemble._hist_gradient_boosting.binning import _BinMapper
from sklearn.ensemble._hist_gradient_boosting.common import G_H_DTYPE, X_DTYPE, Y_DTYPE
from sklearn._loss.loss import BaseLoss

#sklearn data processing and metrics
from sklearn.metrics import accuracy_score, confusion_matrix, balanced_accuracy_score
from sklearn.compose import ColumnTransformer
from sklearn.metrics import check_scoring
from sklearn.metrics._scorer import _SCORERS
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import FunctionTransformer, LabelEncoder, OrdinalEncoder

#sklearn utils
from sklearn.utils import check_random_state, compute_sample_weight, resample
from sklearn.utils._missing import is_scalar_nan
from sklearn.utils._openmp_helpers import _openmp_effective_n_threads
from sklearn.utils._param_validation import Interval, RealNotInt, StrOptions
from sklearn.utils.multiclass import check_classification_targets
from sklearn.utils import check_random_state, compute_sample_weight, resample, compute_class_weight
from sklearn.utils.validation import (
    _check_monotonic_cst,
    _check_sample_weight,
    _check_y,
    _is_pandas_df,
    check_array,
    check_consistent_length,
    check_is_fitted,
    validate_data,
)

from copy import deepcopy

from sklearn._loss import HalfBinomialLoss
from pymoo.mcdm.pseudo_weights import PseudoWeights
from typing import Literal

__all__ = ["MMM_Fair_GradientBoostedClassifier"]


def logistic_chain_rule(prob):
    """
    d(prob)/d(raw_pred) = prob * (1 - prob).
    """
    return prob * (1.0 - prob)


class MultiFairBinomialLoss(HalfBinomialLoss):
    """
    A multi-attribute fairness-aware binomial loss for binary classification.
    We combine multiple fairness penalties via a smooth max operator,
    approximating max across attributes in a differentiable manner.

    Each attribute is a binary array (0=protected,1=non-protected).
    We do e.g. 'DP' for each attribute => difference in predicted positives,
    then combine them via smoothmax with parameter 'beta'.

    Parameters
    ----------
    saIndex : array of protected values shape (n_samples,n_protected)
        For each protected attribute i, a value of 0/1 indicating group membership.
    alpha_fair : float, default=0.1
        Weight of the fairness penalty relative to logistic loss.
    beta : float, default=10.0
        Smoothmax parameter. Larger means closer to actual max but less smooth.
    fairness_type : str in {"DP","EP","EO"}, default="DP"
        Which measure to compute for each attribute.
        For multiple attributes, we do the same measure but you can adapt.
    sample_weight : array-like, optional
        For base logistic. 
    """

    def __init__(
        self,
        saIndex,
        saValue=None,
        alpha_fair=0.1,
        beta=10.0,
        fairness_type="DP",
        sample_weight=None,
    ):
        super().__init__(sample_weight=sample_weight)
        # store protected attributes as a list of arrays
        # each array shape (n_samples,)
        self.saIndex = saIndex
        if saValue is not None:
            self.saValue = saValue
        else:
            self.saValue = {}
            for j in range(self.saIndex.shape[-1]):
                self.saValue[str(j)]=0
        self.alpha_fair = alpha_fair
        self.beta = beta
        self.fairness_type = fairness_type

    def _compute_single_fairness(self, y_true, p, s):
        """
        Compute a single-attribute fairness measure:
        DP, EP, or EO (toy version with squared difference).
        
        s is shape (n_samples,), 0=protected,1=non-protected
        p is predicted prob of pos class, shape (n_samples,).
        y_true is in {0,1}.
        Return:
          scalar fairness measure f_j
          array partial derivative df_j/d p[i] in shape (n_samples,).
        """
        idx0 = (s == 0)
        idx1 = (s == 1)
        n0 = max(1, idx0.sum())
        n1 = max(1, idx1.sum())

        if self.fairness_type.lower() == "dp":
            # difference in mean predicted positives
            mean0 = p[idx0].mean() if np.any(idx0) else 0.0
            mean1 = p[idx1].mean() if np.any(idx1) else 0.0
            diff = (mean1 - mean0)
            # penalty = diff^2
            penalty = diff*diff

            # derivative wrt each sample's p[i]
            grad = np.zeros_like(p)
            factor = 2.0 * diff
            for i in range(len(p)):
                if s[i] == 1:
                    chain = factor*(1./n1)
                else:
                    chain = factor*(-1./n0)
                grad[i] = chain * logistic_chain_rule(p[i])
            return penalty, grad

        elif self.fairness_type.lower() == "ep":
            # difference in TPR => subset y=1
            idx_pos = (y_true == 1)
            idx0p = idx0 & idx_pos
            idx1p = idx1 & idx_pos
            n0p = max(1, idx0p.sum())
            n1p = max(1, idx1p.sum())
            mean0p = p[idx0p].mean() if np.any(idx0p) else 0.
            mean1p = p[idx1p].mean() if np.any(idx1p) else 0.
            diff = mean1p - mean0p
            penalty = diff*diff

            grad = np.zeros_like(p)
            factor = 2.*diff
            for i in range(len(p)):
                if idx_pos[i]:
                    if s[i] == 1:
                        chain = factor*(1./n1p)
                    else:
                        chain = factor*(-1./n0p)
                    grad[i] = chain * logistic_chain_rule(p[i])
            return penalty, grad

        elif self.fairness_type.lower() == "eo":
            # TPR difference + TNR difference => for y=1, y=0
            # define penalty = (diff_tpr^2 + diff_tnr^2).
            idx_pos = (y_true == 1)
            idx_neg = (y_true == 0)
            idx1p = idx1 & idx_pos
            idx0p = idx0 & idx_pos
            idx1n = idx1 & idx_neg
            idx0n = idx0 & idx_neg

            n0p = max(1, idx0p.sum())
            n1p = max(1, idx1p.sum())
            n0n = max(1, idx0n.sum())
            n1n = max(1, idx1n.sum())

            mean0p = p[idx0p].mean() if np.any(idx0p) else 0.
            mean1p = p[idx1p].mean() if np.any(idx1p) else 0.
            diff_tpr = (mean1p - mean0p)

            # TNR => predicted negative => 1-p
            mean0n = (1.-p[idx0n]).mean() if np.any(idx0n) else 0.
            mean1n = (1.-p[idx1n]).mean() if np.any(idx1n) else 0.
            diff_tnr = (mean1n - mean0n)

            penalty = diff_tpr*diff_tpr + diff_tnr*diff_tnr

            grad = np.zeros_like(p)
            # partial derivative wrt TPR part
            factor_tpr = 2.*diff_tpr
            # partial derivative wrt TNR part
            factor_tnr = 2.*diff_tnr

            for i in range(len(p)):
                chain = 0.
                if y_true[i] == 1:
                    # TPR
                    if s[i] == 1:
                        chain += factor_tpr*(1./n1p)
                    else:
                        chain += factor_tpr*(-1./n0p)
                    # chain => derivative wrt p => multiply logistic_chain_rule
                    chain *= logistic_chain_rule(p[i])
                else:
                    # TNR => partial wrt p => negative sign
                    if s[i] == 1:
                        chain += factor_tnr*(1./n1n)*(-1.)
                    else:
                        chain += factor_tnr*(-1./n0n)*(-1.)
                    chain *= logistic_chain_rule(p[i])
                grad[i] = chain
            return penalty, grad

        else:
            # no fairness
            return 0., np.zeros_like(p)

    def _fairness_smoothmax(self, y_true, raw_pred):
        """
        For each attribute j:
          f_j, grad_j
        Then combine via softmax with param beta => smoothmax of f_j.

        Return total_penalty (scalar), grad_penalty (array shape=(n_samples,)) 
        """
        p = 1.0 / (1.0 + np.exp(-raw_pred))

        # For each attribute j => compute penalty_j, grad_j
        penalty_list = []
        grad_list = []
        for j in range(self.saIndex.shape[-1]):
            s_j=self.saIndex[:,j]
            pen_j, grad_j = self._compute_single_fairness(y_true, p, s_j)
            penalty_list.append(pen_j)
            grad_list.append(grad_j)
        
        # penalty_list => shape (m,) for m attributes
        # we define smoothmax( f_1, ..., f_m ) with param self.beta
        # => c = (1/beta)*log( sum( exp(beta * f_j) ) )
        # partial derivative wrt f_j => softmax weighting
        # partial wrt raw_pred => sum_j( w_j * d f_j/d raw_pred )
        f_arr = np.array(penalty_list, dtype=float)  # shape (m,)
        # avoid overflow
        max_f = np.max(f_arr)
        exps = np.exp(self.beta * (f_arr - max_f))
        sum_exps = np.sum(exps)
        c = (1.0/self.beta)*( np.log(sum_exps) ) + max_f  # shift back

        # derivative of c wrt f_j => w_j = exp( beta*(f_j - max_f) ) / sum_exps
        w = exps / sum_exps  # shape (m,)

        # partial derivative wrt each sample i => sum_j( w_j * grad_j[i] )
        # Combine all attributes' partial derivatives
        grad_c = np.zeros_like(grad_list[0])
        for j in range(self.saIndex.shape[-1]):
            grad_c += w[j] * grad_list[j]
        
        return c, grad_c

    def _fairness_penalty_grad(self, y_true, raw_pred):
        """
        The single-sample partial derivative of the multi-attribute
        fairness penalty. We do a smoothmax across attributes.

        Return grad_fair shape (n_samples,).
        """
        total_penalty, grad_penalty = self._fairness_smoothmax(y_true, raw_pred)
        return total_penalty, grad_penalty

    def gradient(self, y_true, raw_prediction, sample_weight, gradient_out, n_threads=1):
        """
        Overriding gradient: logistic + alpha_fair * fairness derivative
        ignoring hessian. scikit-learn calls gradient() if self.constant_hessian = True.
        """
        # call base logistic gradient
        super().gradient(y_true, raw_prediction, sample_weight, gradient_out, n_threads)
        # compute multi-attribute fairness
        _, grad_fair = self._fairness_penalty_grad(y_true, raw_prediction)
        gradient_out += self.alpha_fair * grad_fair

        return gradient_out

    def gradient_hessian(
        self,
        y_true,
        raw_prediction,
        sample_weight,
        gradient_out,
        hessian_out,
        n_threads=1
    ):
        """
        Overriding gradient & hessian. We'll do logistic for hessian, and
        add alpha_fair * fairness to gradient.
        """

        super().gradient_hessian(
            y_true, raw_prediction, sample_weight, gradient_out, hessian_out, n_threads
        )

        _, grad_fair = self._fairness_penalty_grad(y_true, raw_prediction)

        if len(gradient_out.shape)<len(grad_fair.shape):
            grad_fair = grad_fair.ravel()

        mean_log = np.mean(np.abs(gradient_out))
        mean_fair = np.mean(np.abs(grad_fair))
        ratio = mean_log / (mean_fair + 1e-7) ##slack to avoid division by zero
        gradient_out += self.alpha_fair * grad_fair*ratio

        return gradient_out,hessian_out



class MMM_Fair_GradientBoostedClassifier(HistGradientBoostingClassifier):
    def __init__(
        self,
        saIndex=None,
        saValue=None,
        constraint='DP',
        alpha=0.1,
        store_iter_metrics=True,
        preference=None,
        # plus any standard HistGradientBoostingClassifier params
        early_stopping="auto",
        validation_fraction=0.1,
        max_iter=100,
        random_state=None,
        # ...
        **kwargs
    ):

        if constraint is None:
            # your default or fallback
            super().__init__(
                loss="log_loss",
                early_stopping=early_stopping,
                validation_fraction=validation_fraction,
                max_iter=max_iter,
                random_state=random_state,
                **kwargs
            )
            self.saValue = saValue
            self.saIndex = saIndex
        else:
            if saIndex is None:
                raise ValueError("Protected features cannot be None for fairness intervention")
            else:
                self.saIndex = saIndex
                if saValue is not None:
                    self.saValue = saValue
                else:
                    self.saValue = {}
                    for j in range(self.saIndex.shape[-1]):
                        self.saValue[str(j)]=0
                self.sensitives=list(self.saValue.keys())
            super().__init__(
                loss=MultiFairBinomialLoss(
                            saIndex=self.saIndex,
                            saValue=self.saValue,
                            alpha_fair=alpha,
                            beta=10.0,
                            fairness_type=constraint,  # or "EP","EO"
                        ),
                early_stopping=early_stopping,
                validation_fraction=validation_fraction,
                max_iter=max_iter,
                random_state=random_state,
                **kwargs
            )
        self.alpha=alpha
        self.constraint=constraint

        if self.saValue is None or len(self.saValue)!=self.saIndex.shape[-1]:
            if saIndex is not None:
                self.saValue = {}
                for j in range(self.saIndex.shape[-1]):
                        self.saValue[str(j)]=0
                self.sensitives=list(self.saValue.keys())
            else:
                self.sensitives=[]
        self.store_iter_metrics = store_iter_metrics
        self.acc_loss_ = []
        self.bal_loss_ = []
        self.mmm_loss_ = []
        self.all_estimators=[]
        self.ob=[]
        self.fairobs=[]
        self.theta=-1
        self.pseudo=None
        self.pareto=False
        self.PF=None
        self.preference=preference
        

    def fit(self, X, y, sample_weight=None):
        """
        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            The input samples.
    
        y : array-like of shape (n_samples,)
            Target values.
    
        sample_weight : array-like of shape (n_samples,) default=None
            Weights of training data.
    
            
    
        Returns
        -------
        self : object
            Fitted estimator.
        """
        fit_start_time = time()
        acc_find_split_time = 0.0  # time spent finding the best splits
        acc_apply_split_time = 0.0  # time spent splitting nodes
        acc_compute_hist_time = 0.0  # time spent computing histograms
        # time spent predicting X for gradient and hessians update
        acc_prediction_time = 0.0
        X, known_categories = self._preprocess_X(X, reset=True)
        y = _check_y(y, estimator=self)
        y = self._encode_y(y)
        check_consistent_length(X, y)
        if self.saIndex is None:
            self.saIndex = np.random.randint(2, size=(len(X),1))
            print(f"Caution!! saIndex not provided. Model initiated with random protected values, sample: {self.saIndex[:50,0]}")
            
        check_consistent_length(X, self.saIndex)
        # Do not create unit sample weights by default to later skip some
        # computation
        
        if sample_weight is not None:
            sample_weight = _check_sample_weight(sample_weight, X, dtype=np.float64)
            # TODO: remove when PDP supports sample weights
            self._fitted_with_sw = True
    
        sample_weight = self._finalize_sample_weight(sample_weight, y)
    
        rng = check_random_state(self.random_state)
    
        # When warm starting, we want to reuse the same seed that was used
        # the first time fit was called (e.g. train/val split).
        # For feature subsampling, we want to continue with the rng we started with.
        if not self.warm_start or not self._is_fitted():
            self._random_seed = rng.randint(np.iinfo(np.uint32).max, dtype="u8")
            feature_subsample_seed = rng.randint(np.iinfo(np.uint32).max, dtype="u8")
            self._feature_subsample_rng = np.random.default_rng(feature_subsample_seed)
    
        self._validate_parameters()
        monotonic_cst = _check_monotonic_cst(self, self.monotonic_cst)
        # _preprocess_X places the categorical features at the beginning,
        # change the order of monotonic_cst accordingly
        if self.is_categorical_ is not None:
            monotonic_cst_remapped = np.concatenate(
                (
                    monotonic_cst[self.is_categorical_],
                    monotonic_cst[~self.is_categorical_],
                )
            )
        else:
            monotonic_cst_remapped = monotonic_cst
    
        # used for validation in predict
        n_samples, self._n_features = X.shape
    
        # Encode constraints into a list of sets of features indices (integers).
        interaction_cst = self._check_interaction_cst(self._n_features)
        self._in_fit = True
    
        # `_openmp_effective_n_threads` is used to take cgroups CPU quotes
        # into account when determine the maximum number of threads to use.
        n_threads = _openmp_effective_n_threads()
    
    
        if isinstance(self.loss, str):
            self._loss = self._get_loss(sample_weight=sample_weight)
        elif isinstance(self.loss, BaseLoss):
            self._loss = self.loss
    
        if self.early_stopping == "auto":
            self.do_early_stopping_ = n_samples > 10000
        else:
            self.do_early_stopping_ = self.early_stopping
    
        # create validation data if needed
        self._use_validation_data = self.validation_fraction is not None
        if self.do_early_stopping_ and self._use_validation_data:
            # stratify for classification
            # instead of checking predict_proba, loss.n_classes >= 2 would also work
            stratify = y if hasattr(self._loss, "predict_proba") else None
    
            # Save the state of the RNG for the training and validation split.
            # This is needed in order to have the same split when using
            # warm starting.
    
            if sample_weight is None:
                X_train, X_val, y_train, y_val, sa_train, sa_val = train_test_split(
                    X,
                    y,
                    self.saIndex,
                    test_size=self.validation_fraction,
                    stratify=stratify,
                    random_state=self._random_seed,
                )
                sample_weight_train = sample_weight_val = None
            else:
                # TODO: incorporate sample_weight in sampling here, as well as
                # stratify
                (
                    X_train,
                    X_val,
                    y_train,
                    y_val,
                    sa_train,
                    sa_val,
                    sample_weight_train,
                    sample_weight_val,
                ) = train_test_split(
                    X,
                    y,
                    self.saIndex,
                    sample_weight,
                    test_size=self.validation_fraction,
                    stratify=stratify,
                    random_state=self._random_seed,
                )
        else:
            X_train, y_train, sa_train, sample_weight_train = X, y, self.saIndex, sample_weight
            X_val = y_val = sa_val = sample_weight_val = None
    
    
        n_bins = self.max_bins + 1  # + 1 for missing values
        self._bin_mapper = _BinMapper(
            n_bins=n_bins,
            is_categorical=self._is_categorical_remapped,
            known_categories=known_categories,
            random_state=self._random_seed,
            n_threads=n_threads,
        )
        X_binned_train = self._bin_data(X_train, is_training_data=True)
        if X_val is not None:
            X_binned_val = self._bin_data(X_val, is_training_data=False)
        else:
            X_binned_val = None
    
        # Uses binned data to check for missing values
        has_missing_values = (
            (X_binned_train == self._bin_mapper.missing_values_bin_idx_)
            .any(axis=0)
            .astype(np.uint8)
        )
    
        if self.verbose:
            print("Fitting gradient boosted rounds:")
    
        n_samples = X_binned_train.shape[0]
        scoring_is_predefined_string = self.scoring in _SCORERS
        need_raw_predictions_val = X_binned_val is not None and (
            scoring_is_predefined_string or self.scoring == "loss"
        )
        
        # First time calling fit, or no warm start
        if not (self._is_fitted() and self.warm_start):
            # Clear random state and score attributes
            self._clear_state()
    
            # initialize raw_predictions: those are the accumulated values
            # predicted by the trees for the training data. raw_predictions has
            # shape (n_samples, n_trees_per_iteration) where
            # n_trees_per_iterations is n_classes in multiclass classification,
            # else 1.
            # self._baseline_prediction has shape (1, n_trees_per_iteration)
            self._baseline_prediction = self._loss.fit_intercept_only(
                y_true=y_train, sample_weight=sample_weight_train
            ).reshape((1, -1))
            raw_predictions = np.zeros(
                shape=(n_samples, self.n_trees_per_iteration_),
                dtype=self._baseline_prediction.dtype,
                order="F",
            )
            raw_predictions += self._baseline_prediction
    
            # predictors is a matrix (list of lists) of TreePredictor objects
            # with shape (n_iter_, n_trees_per_iteration)
            self._predictors = predictors = []
    
            # Initialize structures and attributes related to early stopping
            self._scorer = None  # set if scoring != loss
            raw_predictions_val = None  # set if use val and scoring is a string
            self.train_score_ = []
            self.validation_score_ = []
            
            
        
            if self.do_early_stopping_:
                # populate train_score and validation_score with the
                # predictions of the initial model (before the first tree)
    
                # Create raw_predictions_val for storing the raw predictions of
                # the validation data.
                if need_raw_predictions_val:
                    raw_predictions_val = np.zeros(
                        shape=(X_binned_val.shape[0], self.n_trees_per_iteration_),
                        dtype=self._baseline_prediction.dtype,
                        order="F",
                    )
    
                    raw_predictions_val += self._baseline_prediction
    
    
                ##self.scoring can be accuracy, loss, etc.
                """
                from sklearn.metrics import make_scorer, accuracy_score, mean_squared_log_error
                scoring = {
                        "accuracy": make_scorer(accuracy_score),
                        "mean_squared_log_error": make_scorer(mean_squared_log_error),
                    }
    
                """
    
                if self.scoring == "loss":
                    # we're going to compute scoring w.r.t the loss. As losses
                    # take raw predictions as input (unlike the scorers), we
                    # can optimize a bit and avoid repeating computing the
                    # predictions of the previous trees. We'll reuse
                    # raw_predictions (as it's needed for training anyway) for
                    # evaluating the training loss.
    
                    self._check_early_stopping_loss(
                        raw_predictions=raw_predictions,
                        y_train=y_train,
                        sa_train=sa_train,
                        sample_weight_train=sample_weight_train,
                        raw_predictions_val=raw_predictions_val,
                        y_val=y_val,
                        sa_val=sa_val,
                        sample_weight_val=sample_weight_val,
                        n_threads=n_threads,
                    )
                else:
                    self._scorer = check_scoring(self, self.scoring)
                    # _scorer is a callable with signature (est, X, y) and
                    # calls est.predict() or est.predict_proba() depending on
                    # its nature.
                    # Unfortunately, each call to _scorer() will compute
                    # the predictions of all the trees. So we use a subset of
                    # the training set to compute train scores.
    
                    # Compute the subsample set
                    (
                        X_binned_small_train,
                        y_small_train,
                        sample_weight_small_train,
                        indices_small_train,
                    ) = self._get_small_trainset(
                        X_binned_train,
                        y_train,
                        sample_weight_train,
                        self._random_seed,
                    )
    
                    # If the scorer is a predefined string, then we optimize
                    # the evaluation by re-using the incrementally updated raw
                    # predictions.
                    if scoring_is_predefined_string:
                        raw_predictions_small_train = raw_predictions[
                            indices_small_train
                        ]
                    else:
                        raw_predictions_small_train = None
    
                    self._check_early_stopping_scorer(
                        X_binned_small_train,
                        y_small_train,
                        sample_weight_small_train,
                        X_binned_val,
                        y_val,
                        sample_weight_val,
                        raw_predictions_small_train=raw_predictions_small_train,
                        raw_predictions_val=raw_predictions_val,
                    )
            begin_at_stage = 0
    
    
          
    
        # warm start: this is not the first time fit was called
        else:
            # Check that the maximum number of iterations is not smaller
            # than the number of iterations from the previous fit
            if self.max_iter < self.n_iter_:
                raise ValueError(
                    "max_iter=%d must be larger than or equal to "
                    "n_iter_=%d when warm_start==True" % (self.max_iter, self.n_iter_)
                )
    
            # Convert array attributes to lists
            self.train_score_ = self.train_score_.tolist()
            self.validation_score_ = self.validation_score_.tolist()
    
            # Compute raw predictions
            raw_predictions = self._raw_predict(X_binned_train, n_threads=n_threads)
            if self.do_early_stopping_ and need_raw_predictions_val:
                raw_predictions_val = self._raw_predict(
                    X_binned_val, n_threads=n_threads
                )
            else:
                raw_predictions_val = None
    
            if self.do_early_stopping_ and self.scoring != "loss":
                # Compute the subsample set
                (
                    X_binned_small_train,
                    y_small_train,
                    sample_weight_small_train,
                    indices_small_train,
                ) = self._get_small_trainset(
                    X_binned_train, y_train, sample_weight_train, self._random_seed
                )
                sa_train=self.saIndex[indices_small_train]
    
            # Get the predictors from the previous fit
            predictors = self._predictors
    
            begin_at_stage = self.n_iter_
    
        # initialize gradients and hessians (empty arrays).
        # shape = (n_samples, n_trees_per_iteration).
        gradient, hessian = self._loss.init_gradient_and_hessian(
            n_samples=n_samples, dtype=G_H_DTYPE, order="F"
        )
    
    
        for iteration in range(begin_at_stage, self.max_iter):
            if self.verbose >= 2:
                iteration_start_time = time()
                print(
                    "[{}/{}] ".format(iteration + 1, self.max_iter), end="", flush=True
                )
    
            # Update gradients and hessians, inplace
            # Note that self._loss expects shape (n_samples,) for
            # n_trees_per_iteration = 1 else shape (n_samples, n_trees_per_iteration).
            if self._loss.constant_hessian:
                self._loss.saIndex=sa_train
                self._loss.gradient(
                    y_true=y_train,
                    raw_prediction=raw_predictions,
                    sample_weight=sample_weight_train,
                    gradient_out=gradient,
                    n_threads=n_threads,
                )
            else:
                self._loss.saIndex=sa_train
                self._loss.gradient_hessian(
                    y_true=y_train,
                    raw_prediction=raw_predictions,
                    sample_weight=sample_weight_train,
                    gradient_out=gradient,
                    hessian_out=hessian,
                    n_threads=n_threads,
                )
    
            # Append a list since there may be more than 1 predictor per iter
            predictors.append([])
    
            # 2-d views of shape (n_samples, n_trees_per_iteration_) or (n_samples, 1)
            # on gradient and hessian to simplify the loop over n_trees_per_iteration_.
            if gradient.ndim == 1:
                g_view = gradient.reshape((-1, 1))
                h_view = hessian.reshape((-1, 1))
            else:
                g_view = gradient
                h_view = hessian
    
            # Build `n_trees_per_iteration` trees.
            for k in range(self.n_trees_per_iteration_):
                grower = TreeGrower(
                    X_binned=X_binned_train,
                    gradients=g_view[:, k],
                    hessians=h_view[:, k],
                    n_bins=n_bins,
                    n_bins_non_missing=self._bin_mapper.n_bins_non_missing_,
                    has_missing_values=has_missing_values,
                    is_categorical=self._is_categorical_remapped,
                    monotonic_cst=monotonic_cst_remapped,
                    interaction_cst=interaction_cst,
                    max_leaf_nodes=self.max_leaf_nodes,
                    max_depth=self.max_depth,
                    min_samples_leaf=self.min_samples_leaf,
                    l2_regularization=self.l2_regularization,
                    feature_fraction_per_split=self.max_features,
                    rng=self._feature_subsample_rng,
                    shrinkage=self.learning_rate,
                    n_threads=n_threads,
                )
                grower.grow()
    
                acc_apply_split_time += grower.total_apply_split_time
                acc_find_split_time += grower.total_find_split_time
                acc_compute_hist_time += grower.total_compute_hist_time
    
                ### _update_leaves_values only uses loss.fit_intercept_only, so saIndex not needed
                if not self._loss.differentiable:
                    _update_leaves_values(
                        loss=self._loss,
                        grower=grower,
                        y_true=y_train,
                        raw_prediction=raw_predictions[:, k],
                        sample_weight=sample_weight_train,
                    )
    
                predictor = grower.make_predictor(
                    binning_thresholds=self._bin_mapper.bin_thresholds_
                )
                predictors[-1].append(predictor)
    
                # Update raw_predictions with the predictions of the newly
                # created tree.
                tic_pred = time()
                _update_raw_predictions(raw_predictions[:, k], grower, n_threads)
                toc_pred = time()
                acc_prediction_time += toc_pred - tic_pred
    
    
    
            should_early_stop = False
            if self.do_early_stopping_:
                # Update raw_predictions_val with the newest tree(s)
                if need_raw_predictions_val:
                    for k, pred in enumerate(self._predictors[-1]):
                        raw_predictions_val[:, k] += pred.predict_binned(
                            X_binned_val,
                            self._bin_mapper.missing_values_bin_idx_,
                            n_threads,
                        )
    
                if self.scoring == "loss":
                    should_early_stop = self._check_early_stopping_loss(
                        raw_predictions=raw_predictions,
                        y_train=y_train,
                        sa_train=sa_train,
                        sample_weight_train=sample_weight_train,
                        raw_predictions_val=raw_predictions_val,
                        y_val=y_val,
                        sa_val=sa_val,
                        sample_weight_val=sample_weight_val,
                        n_threads=n_threads,
                    )
    
                else:
                    # If the scorer is a predefined string, then we optimize the
                    # evaluation by re-using the incrementally computed raw predictions.
                    if scoring_is_predefined_string:
                        raw_predictions_small_train = raw_predictions[
                            indices_small_train
                        ]
                    else:
                        raw_predictions_small_train = None
    
                    should_early_stop = self._check_early_stopping_scorer(
                        X_binned_small_train,
                        y_small_train,
                        sample_weight_small_train,
                        X_binned_val,
                        y_val,
                        sample_weight_val,
                        raw_predictions_small_train=raw_predictions_small_train,
                        raw_predictions_val=raw_predictions_val,
                    )
                    ####hola
    
    
            if self.verbose >= 2:
                self._print_iteration_stats(iteration_start_time)
    
            # maybe we could also early stop if all the trees are stumps?
            if should_early_stop:
                break
    
        if self.verbose:
            duration = time() - fit_start_time
            n_total_leaves = sum(
                predictor.get_n_leaf_nodes()
                for predictors_at_ith_iteration in self._predictors
                for predictor in predictors_at_ith_iteration
            )
            n_predictors = sum(
                len(predictors_at_ith_iteration)
                for predictors_at_ith_iteration in self._predictors
            )
            print(
                "Fit {} trees in {:.3f} s, ({} total leaves)".format(
                    n_predictors, duration, n_total_leaves
                )
            )
            print(
                "{:<32} {:.3f}s".format(
                    "Time spent computing histograms:", acc_compute_hist_time
                )
            )
            print(
                "{:<32} {:.3f}s".format(
                    "Time spent finding best splits:", acc_find_split_time
                )
            )
            print(
                "{:<32} {:.3f}s".format(
                    "Time spent applying splits:", acc_apply_split_time
                )
            )
            print(
                "{:<32} {:.3f}s".format("Time spent predicting:", acc_prediction_time)
            )
    
        self.train_score_ = np.asarray(self.train_score_)
        self.validation_score_ = np.asarray(self.validation_score_)
        self.all_estimators=[predictor for predictor in self._predictors]
        
    
        if self.store_iter_metrics:
            X_ = check_array(X, dtype=None, accept_sparse=False)
            check_is_fitted(self)
            for ypred in self.staged_predict(X_):
                o1=1-accuracy_score(y_pred=ypred,y_true=y)
                o2=1-balanced_accuracy_score(y_pred=ypred,y_true=y).item()
                o3=[]
                for i in range(self.saIndex.shape[-1]):
                    sens = fb.categories(self.saIndex[:, i])
                    report = fb.reports.pairwise(
                        predictions=ypred,
                        labels=y,
                        sensitive=sens
                    )
                    report_dict=report.show(env=fb.export.ToDict)
                    if self.constraint.lower()=='dp':
                        o3.append(report_dict['depends'][6]['depends'][1]['value']['value'])
                    elif self.constraint.lower()=='ep':
                        o3.append(report_dict['depends'][6]['depends'][2]['value']['value'])
                    elif self.constraint.lower()=='eo':
                        o3.append(max(report_dict['depends'][6]['depends'][2]['value']['value'],
                                      report_dict['depends'][6]['depends'][3]['value']['value']))
                    else:
                        o3.append(0.99)
                self.ob.append([o1, o2, max(o3)])
                self.fairobs.append(o3)    
                
                
        self.ob = np.array(self.ob)
        self.fairobs = np.array(self.fairobs)
        del self._in_fit  # hard delete so we're sure it can't be used anymore
        return self
    
    
    def update_theta(
        self,
        criteria: Literal["all", "fairness"] = "all",
        preference=[0.33, 0.34, 0.33],
        theta=None ##for bruteforce theta update
    ):
        def is_pareto(costs, maximise=False):
            """
            :param costs: An (n_points, n_costs) array
            :maximise: boolean. True for maximising, False for minimising
            :return: A (n_points, ) boolean array, indicating whether each point is Pareto efficient
            """
            eps = 0.001
            is_efficient = np.ones(costs.shape[0], dtype=bool)
            for i, c in enumerate(costs):
                if is_efficient[i]:
                    if maximise:
                        is_efficient[is_efficient] = np.any(
                            costs[is_efficient] >= c, axis=1
                        )  # Remove dominated points
                    else:
                        is_efficient[is_efficient] = np.any(
                            costs[is_efficient] <= c, axis=1
                        )  # Remove dominated points
            return is_efficient

        if theta is None:
            #self.preference = preference
            best_theta = 0
            if criteria.lower() == "fairness":
                objective = deepcopy(self.fairobs)
                if len(preference)!= self.fairobs.shape[-1]:
                    preference = [1/self.fairobs.shape[-1] for i in range(self.fairobs.shape[-1])]
            else:
                objective = deepcopy(self.ob)
            # objective=np.round(objective,2)
            if self.pareto == False:
                PF = {i: objective[i] for i in range(len(objective)) if np.any(objective[i]==0)==False}
                F = np.array([objective[o] for o in range(len(objective))])
                self.PF = PF
            else:
                pf = is_pareto(objective)
                PF = {i: objective[i] for i in range(len(pf)) if pf[i] == True and np.any(objective[i]==0)==False}
                F = np.array(list(PF.values()))
                self.PF = PF
    
            if self.preference is None or len(self.preference)!= objective.shape[-1]:
                weights = preference  ##Preference Weights
            else:
                weights = self.preference
    
            best_theta, pseudo_weights = PseudoWeights(weights).do(
                F, return_pseudo_weights=True
            )
    
            if self.preference == None:
                sum_W = [sum((1 - pseudo_weights[w]) * F[w]) for w in range(len(PF))]
                best_theta = sum_W.index(min(sum_W))
    
            self.theta = list(PF.keys())[best_theta] + 1
            self.pseudo = pseudo_weights
            self._predictors=self.all_estimators[:self.theta]
        else:
            self.theta =theta
            self._predictors=self.all_estimators[:self.theta]
    
    def _check_early_stopping_loss(
        self,
        raw_predictions,
        y_train,
        sa_train,
        sample_weight_train,
        raw_predictions_val,
        y_val,
        sa_val,
        sample_weight_val,
        n_threads=1,
    ):
        """Check if fitting should be early-stopped based on loss.
    
        Scores are computed on validation data or on training data.
        """
        self._loss.saIndex=sa_train
        self.train_score_.append(
            -self._loss(
                y_true=y_train,
                raw_prediction=raw_predictions,
                sample_weight=sample_weight_train,
                n_threads=n_threads,
            )
        )
    
        if self._use_validation_data:
            self._loss.saIndex=sa_val
            self.validation_score_.append(
                -self._loss(
                    y_true=y_val,
                    raw_prediction=raw_predictions_val,
                    sample_weight=sample_weight_val,
                    n_threads=n_threads,
                )
            )
            return self._should_stop(self.validation_score_)
        else:
            return self._should_stop(self.train_score_)
    
    def see_pareto(self):
        import plotly.graph_objs as go
        from plotly.subplots import make_subplots
        def plot2d(x,y,axis_names):
            pass
        def plot3d(x = [1, 2, 3, 4, 5],y = [2, 3, 1, 4, 5],z = [3, 1, 2, 5, 4],weights = [1,2,3,4,5],criteria='all',axis_names=['X','Y','Z']):
            # Create the figure
            hidden_data=[str(np.around(weights[i],2)) for i in range(len(weights))]
            fig = make_subplots(rows=1, cols=1, specs=[[{'type': 'scatter3d'}]])
            
            # Add 3D scatter plot
            fig.add_trace(go.Scatter3d(x=x, y=y, z=z, mode='markers', marker=dict(color='blue'), 
                                       name='Points',text=hidden_data, customdata=hidden_data ), row=1, col=1)
            
            # Define callback function for selection event
            def update_point(trace, points, selector):
                selected_points = [(trace.x[i], trace.y[i], trace.z[i]) for i in points.point_inds]
                selected_hidden_data = [trace.customdata[i] for i in points.point_inds]
                print("Selected Points:", selected_points,selected_hidden_data)
            
            # Assign callback function to plot
            fig.data[0].on_selection(update_point)
            
            # Update layout
            fig.update_layout(title=criteria+'-objectives Pareto Plot - Select Points to Return Vector',
                             scene=dict(
                                        xaxis_title='X:'+axis_names[0],
                                        yaxis_title='Y:'+ axis_names[1],
                                        zaxis_title='Z:'+ axis_names[2]
                                    ),
                                    width=1000,  # Set the width of the plot
                                    height=700 )
            
            # Show plot
            fig.show()
        
        #self.update_theta(criteria='all')
        PF=np.array([self.ob[i] for i in range(len(self.ob))])
        plot3d(x=PF[:,0],y=PF[:,1],z=PF[:,2], weights=self.pseudo,
               axis_names=['Acc.','Balanc. Acc', 'MMM-fair'])
        
        if len(self.sensitives)>0:
            #self.theta=None
            #self.update_theta(criteria='fairness')
            PF=np.array([self.fairobs[i] for i in range(len(self.fairobs))])
            if PF.shape[-1]>2:
                plot3d(x=PF[:,0],y=PF[:,1],z=PF[:,2], weights=self.pseudo,
                       criteria='Fairness',axis_names=self.sensitives)
            elif PF.shape[-1]==2:
                    plot3d(x=PF[:,0],y=PF[:,1],z=np.zeros_like(PF[:,1]), weights=self.pseudo,
                       criteria='Fairness',axis_names=self.sensitives+[''])
        
            else:
                print('Not a MMM-fair')   
    
    





