# AUTOGENERATED! DO NOT EDIT! File to edit: 04_optuna_objective.ipynb (unless otherwise specified).

__all__ = ['Optuna_Objective']

# Cell
from .dataframeloader import *
from .preprocessor import *
from .logger import *

# Cell
# hide
import numpy as np

from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_absolute_error, roc_auc_score, accuracy_score
from sklearn.model_selection import cross_val_score, GridSearchCV, StratifiedKFold
# for Optuna
import optuna
#for XGB
from xgboost import XGBRegressor, XGBClassifier


# for displaying diagram of pipelines
from sklearn import set_config
set_config(display="diagram")

# for finding n_jobs in all sklearn estimators
from sklearn.utils import all_estimators
import inspect

# Just to compare fit times
import time

# for os specific settings
import os

# Cell

class Optuna_Objective:
    """
    Represent Optuna Objective class

    Attributes:\n
    pl: A tmlt instance \n
    dfl: A DataFrameLoader instance \n
    pp: A PreProcessor Instance \n
    model: The given Model
    """

    def __init__(self, dfl, tmlt, preds_metrics, xgb_model, xgb_eval_metric, use_predict_proba, use_gpu):
        self.X = dfl.X
        self.y = dfl.y
        self.use_gpu = use_gpu
        self.tmlt = tmlt
        self.xgb_model = xgb_model
        self.trial = None
        self.xgb_eval_metric = xgb_eval_metric
        self.use_predict_proba = use_predict_proba
        self.preds_metrics = preds_metrics

    def __call__(self, trial):
        self.trial = trial

        #get_params here
        xgb_params = self.get_xgb_params(self.trial, use_gpu=self.use_gpu)
        early_stopping_rounds = xgb_params["early_stopping_rounds"]
        del xgb_params["early_stopping_rounds"]

        # create xgb ml model
        model = self.xgb_model(
            random_state=42,
            eval_metric=self.xgb_eval_metric,
            use_label_encoder=False,
            **xgb_params,
        )
        # update the model on pipeline
        self.tmlt.update_model(model)

        #create train valid datasets for simple training
        self.tmlt.dfl.create_train_valid(valid_size=0.2)
        # Now fit
        logger.info("Training Started!")
        self.tmlt.spl.fit(self.tmlt.dfl.X_train, self.tmlt.dfl.y_train)
        logger.info("Training Ended!")

        # choose predict type based upon metrics type
        #TO-DO instead of single metrics use list of metrics and calculate mean using dict
        metric_result = {}
        for metric in self.preds_metrics:
            if ("log_loss" in str(metric.__name__)) or ("roc_auc_score" in str(metric.__name__)):
                #logger.info("Predicting Probablities!")
                preds_probs = self.tmlt.spl.predict_proba(self.tmlt.dfl.X_valid)[:, 1]
                metric_result[str(metric.__name__)] = metric(self.tmlt.dfl.y_valid, preds_probs)

            else:
                #logger.info("Predicting Score!")
                preds = self.tmlt.spl.predict(self.tmlt.dfl.X_valid)
                metric_result[str(metric.__name__)] = metric(self.tmlt.dfl.y_valid, preds)

        #now show value of all the given metrics
        for metric_name, metric_value in metric_result.items():
            logger.info(f"{metric_name}: {metric_value}")

        # choose return result for Optuna Optimization based upon direction
        if self.xgb_eval_metric == "logloss":
            return metric_result["log_loss"]
        elif self.xgb_eval_metric == "mlogloss":
            return metric_result["log_loss"]
        elif self.xgb_eval_metric == "rmse":
            return metric_result["mean_squared_error"]
        else:
            return None


    def __str__(self):
        """Returns human readable string reprsentation"""
        attr_str = ("X, y, task, use_gpu, tmlt, xgb_model")
        return ("Optuna Objective object with attributes:"+attr_str)

    def __repr__(self):
        return self.__str__()

    #helper methods
    #Method to get xgb_params
    def get_xgb_params(self, trial, use_gpu):
        params = {
            "learning_rate": trial.suggest_float("learning_rate", 1e-2, 0.25, log=True),
            "reg_lambda": trial.suggest_float("reg_lambda", 1e-8, 100.0, log=True),
            "reg_alpha": trial.suggest_float("reg_alpha", 1e-8, 100.0, log=True),
            "subsample": trial.suggest_float("subsample", 0.1, 1.0),
            "colsample_bytree": trial.suggest_float("colsample_bytree", 0.1, 1.0),
            "max_depth": trial.suggest_int("max_depth", 1, 9),
            "early_stopping_rounds": trial.suggest_int("early_stopping_rounds", 100, 500),
            "n_estimators": trial.suggest_categorical("n_estimators", [7000, 15000, 20000]),
        }
        if use_gpu:
            params["tree_method"] = "gpu_hist"
            params["gpu_id"] = 0
            params["predictor"] = "gpu_predictor"
        else:
            params["tree_method"] = trial.suggest_categorical("tree_method", ["exact", "approx", "hist"])
            params["booster"] = trial.suggest_categorical("booster", ["gbtree", "gblinear"])
            if params["booster"] == "gbtree":
                params["gamma"] = trial.suggest_float("gamma", 1e-8, 1.0, log=True)
                params["grow_policy"] = trial.suggest_categorical("grow_policy", ["depthwise", "lossguide"])

        return params

    # core methods