# https://github.com/CDonnerer/xgboost-distribution
from __future__ import annotations

import os
from functools import partial
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional

from hyperopt import STATUS_OK, Trials, fmin, hp, tpe
from hyperopt.pyll.base import scope
from sklearn.model_selection import train_test_split
from xgboost_distribution import XGBDistribution

import nemo_bo.utils.logger as logging_nemo
import nemo_bo.utils.perf_metrics as pm
from nemo_bo.models.base.base_model import Base_Model
from nemo_bo.utils.data_proc import *

if TYPE_CHECKING:
    from nemo_bo.opt.objectives import RegressionObjective
    from nemo_bo.opt.variables import VariablesList

try:
    logging_nemo.logging_path
    logger = logging_nemo.logging_nemo_child(os.path.basename(__file__))
except AttributeError:
    logger = logging_nemo.logging_nemo_master(os.path.basename(__file__))


class XGBoostModel(Base_Model):
    """
    Class to instantiate for fitting NGBoost models
    """

    def __init__(
        self, variables: VariablesList, objective: RegressionObjective, always_hyperparam_opt: bool = True, **kwargs
    ):
        super().__init__(variables, objective, always_hyperparam_opt)
        self.default_X_transform_type = "none"
        self.default_Y_transform_type = "none"
        self.include_validation = True
        self.name = "xgb"

    def fit(
        self, X: np.ndarray, Y: np.ndarray, test_ratio: Optional[float] = 0.2, sort_before_split: bool = True, **params
    ) -> None:
        """
        Function for fitting the model

        Parameters
        ----------
        plot_parity : bool type for generating the parity plot for the model fitting performance
        params : dict type containing keyword arguments that can be generated by the HyperOpt package for the model being fitted or chosen manually passing a dictionary
            The keyword arguments that can be used for XGBoost Distribution models in this function relate to the XGBDistribution class and its intrinsically inherited XGBModel class within.

        """
        if sort_before_split:
            (
                X_train,
                X_val,
                Y_train,
                Y_val,
            ) = sort_train_test_split_shuffle(X, Y, test_ratio=test_ratio, seed=1)
        else:
            X_train, X_val, Y_train, Y_val = train_test_split(X, Y, test_size=test_ratio, random_state=1)

        Y_train = Y_train.astype("float")
        Y_val = Y_val.astype("float")

        self.X_train, self.X_val, self.Y_train, self.Y_val = (
            X_train,
            X_val,
            Y_train,
            Y_val,
        )

        if self.variables.num_cat_var > 0:
            self.X_train = self.variables.categorical_transform(X_train).astype("float")
            self.X_val = self.variables.categorical_transform(X_val).astype("float")
        self.X_train, self.Y_train = self.transform_by_predictor_type(self.X_train, Y=self.Y_train)
        self.X_val, self.Y_val = self.transform_only_by_predictor_type(self.X_val, self.Y_val)

        for key in params:
            if key == "max_depth":
                params["max_depth"] = int(params["max_depth"])
            if key == "n_estimators":
                params["n_estimators"] = int(params["n_estimators"])

        self.model = XGBDistribution(early_stopping_rounds=10, **params)

        self.fit_model(params)

        self.Y_train_pred, self.Y_train_pred_stddev = self.predict(X_train)
        self.Y_val_pred, self.Y_val_pred_stddev = self.predict(X_val)
        self.performance_metrics = pm.all_performance_metrics_train_val(
            Y_train, self.Y_train_pred, Y_val, self.Y_val_pred
        )

        self.Y_train_error = 1.96 * self.Y_train_pred_stddev
        self.Y_val_error = 1.96 * self.Y_val_pred_stddev

    def fit_model(self, params: Dict[str, Any]) -> None:
        for key in params:
            if key == "max_depth":
                params["max_depth"] = int(params["max_depth"])
            if key == "n_estimators":
                params["n_estimators"] = int(params["n_estimators"])

        self.model = XGBDistribution(early_stopping_rounds=10, **params)

        # Re-attempts fitting
        for x in range(10):
            try:
                self.model.fit(
                    self.X_train,
                    self.Y_train,
                    eval_set=[(self.X_val, self.Y_val)],
                    verbose=False,
                )
                break
            except np.linalg.LinAlgError as e:
                print(e)
            except TypeError as e:
                print(e)
            except sklearn.exceptions.NotFittedError as e:
                print(e)

    def predict(self, X: np.ndarray, X_transform: bool = True) -> Tuple[np.ndarray, np.ndarray]:
        # The functions below require the data to be a 2D array
        if X.ndim == 1:
            X = X.reshape(1, -1)

        if X_transform:
            if self.variables.num_cat_var > 0:
                X = self.variables.categorical_transform(X).astype("float")
            X = self.transform_only_by_predictor_type(X)

        Y_pred = self.model.predict(X.astype(float))
        Y_mean = Y_pred.loc
        Y_pred_upper = Y_mean + (1.96 * Y_pred.scale)

        Y_pred = self.objective.inverse_transform(Y_mean)
        Y_pred_upper = self.objective.inverse_transform(Y_pred_upper)
        Y_pred_stddev = (np.subtract(Y_pred_upper, Y_pred)) / 1.96

        return Y_mean.flatten(), Y_pred_stddev.flatten()

    @staticmethod
    def default_params():
        return {
            "eta": hp.uniform("eta", 0.01, 0.5),
            "max_depth": scope.int(hp.quniform("max_depth", 3, 18, 1)),
            "min_child_weight": hp.uniform("min_child_weight", 0.01, 0.75),
            "colsample_bytree": hp.uniform("colsample_bytree", 0.5, 1),
        }

    def cv(
        self,
        X: np.ndarray,
        Y: np.ndarray,
        model_params: Dict[str, Any],
        test_ratio: float = 0.2,
        sort_before_split: Optional[bool] = True,
        **kwargs,
    ) -> Dict[str, Any]:
        return self.cv_train_val_test(
            X, Y, model_params, test_ratio=test_ratio, sort_before_split=sort_before_split, **kwargs
        )

    def hyperparam_opt(
        self,
        X: np.ndarray,
        Y: np.ndarray,
        test_ratio: float,
        sort_before_split: Optional[bool] = True,
        predictor_params_dict: Optional[Dict[str, Dict[str, Callable]]] = None,
        **kwargs,
    ) -> Tuple[float, Dict[str, Any]]:

        if predictor_params_dict is None:
            predictor_params_dict = self.default_params()
        else:
            predictor_params_dict = predictor_params_dict[self.name]

        max_evals = kwargs.get("max_evals", 1)
        # max_evals = kwargs.get("max_evals", 100)

        def func(X: np.ndarray, Y: np.ndarray, model_params: Dict[str, Any]) -> Dict[str, Any]:
            cv_results = self.cv(
                X, Y, model_params, test_ratio=test_ratio, sort_before_split=sort_before_split, **kwargs
            )

            return {
                "loss": cv_results["Mean Test RMSE"],
                "status": STATUS_OK,
            }

        trials = Trials()
        model_params = fmin(
            fn=partial(func, X, Y),
            space=predictor_params_dict,
            algo=tpe.suggest,
            max_evals=max_evals,
            trials=trials,
        )

        model = self.new_instance(self.__class__, **kwargs)
        model.fit(X, Y, test_ratio=test_ratio, **self.model_params, **kwargs)

        logger.info(f"Completed hyperparameter opt")
        return model, model_params
