# Tony Duan, Anand Avati, Daisy Yi Ding, Khanh K. Thai, Sanjay Basu, Andrew Y. Ng, Alejandro Schuler. 2019. NGBoost: Natural Gradient Boosting for Probabilistic Prediction.
# https://arxiv.org/abs/1910.03225
# https://github.com/stanfordmlgroup/ngboost

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 ngboost import NGBRegressor
from ngboost.distns import Normal
from ngboost.scores import LogScore
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeRegressor

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 NGBoostModel(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 = "ngb"

    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
        verbose : bool type for if outputs are printed during fitting. Default: True
        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 NGBoost models relate to the NGBRegressor class and its intrinsic DecisionTreeRegressor that is used as the default base learner in NGBoost.
            This was found to give the most reliable performance for regression. Nonetheless, the base learner can be changed to any sklearn regressor using the Base keyword

        """
        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(self.X_train).astype("float")
            self.X_val = self.variables.categorical_transform(self.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)

        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:
        decision_tree_params = {
            "criterion": params.get("criterion", "friedman_mse"),
            "splitter": params.get("splitter", "best"),
            "max_depth": int(params.get("max_depth")),
            "min_samples_split": params.get("min_samples_split", 2),
            "min_samples_leaf": params.get("min_samples_leaf", 1),
            "min_weight_fraction_leaf": params.get("min_weight_fraction_leaf", 0.0),
            "max_features": params.get("max_features"),
            "max_leaf_nodes": params.get("max_leaf_nodes"),
            "min_impurity_decrease": params.get("min_impurity_decrease", 0.0),
            "ccp_alpha": params.get("ccp_alpha", 0.0),
        }
        if decision_tree_params["max_leaf_nodes"] is not None:
            decision_tree_params["max_leaf_nodes"] = int(decision_tree_params["max_leaf_nodes"])

        regressor_params = {
            "Dist": params.get("Dist", Normal),
            "Score": params.get("Score", LogScore),
            "Base": params.get("base", DecisionTreeRegressor(**decision_tree_params)),
            "natural_gradient": params.get("natural_gradient", True),
            "n_estimators": int(params.get("n_estimators", 2000)),
            "learning_rate": params.get("learning_rate", 0.01),
            "minibatch_frac": params.get("minibatch_frac", 1.0),
            "col_sample": params.get("col_sample", 1.0),
            "tol": params.get("tol", 1e-4),
        }

        self.model = NGBRegressor(verbose=True, verbose_eval=20, random_state=1, **regressor_params)

        # Re-attempts fitting
        for x in range(10):
            try:
                self.model.fit(
                    self.X_train,
                    self.Y_train.flatten(),
                    X_val=self.X_val,
                    Y_val=self.Y_val.flatten(),
                    early_stopping_rounds=10,
                )
                break
            except np.linalg.LinAlgError as e:
                print(e)
            except TypeError 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)

        X = self.transform_only_by_predictor_type(X)

        Y_pred = self.model.pred_dist(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 {
            "learning_rate": hp.uniform("learning_rate", 0.005, 0.02),
            "max_depth": scope.int(hp.quniform("max_depth", 3, 10, 1)),
            "min_samples_split": hp.uniform("min_samples_split", 0.01, 0.75),
            "col_sample": hp.uniform("col_sample", 0.5, 1.0),
        }

    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", 30)

        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
