from __future__ import annotations

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

import gpytorch
import numpy as np
import torch
from gpytorch.distributions import MultivariateNormal
from gpytorch.kernels import AdditiveStructureKernel, Kernel, MaternKernel, ScaleKernel
from gpytorch.likelihoods import GaussianLikelihood, Likelihood
from gpytorch.means import ConstantMean
from gpytorch.models import ExactGP
from hyperopt import STATUS_OK, Trials, fmin, hp, tpe
from hyperopt.pyll.base import Literal
from torch import Tensor

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

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 ExactGPModel(ExactGP):
    """
    Class for defining the GP model for regression (inherits from ExactGP in GPyTorch)
    """

    def __init__(self, trainX: Tensor, trainY: Tensor, likelihood: Likelihood, kernel: Kernel):
        super(ExactGPModel, self).__init__(trainX, trainY, likelihood)
        self.mean_module = ConstantMean()
        self.covar_module = kernel

    def forward(self, x: Tensor) -> MultivariateNormal:
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return MultivariateNormal(mean_x, covar_x)


class GPModel(Base_Model):

    """
    Class to instantiate for fitting GP 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 = "normalisation"
        self.default_Y_transform_type = "standardisation"
        self.include_validation = False
        self.name = "gp"

    def fit(
        self, X: np.ndarray, Y: np.ndarray, test_ratio: Optional[float] = None, sort_before_split: bool = None, **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 in this function are limited to the type of kernel and the number of training iterations

        """
        del test_ratio
        del sort_before_split
        if self.variables.num_cat_var > 0:
            self.X_train = self.variables.categorical_transform(X).astype("float")

        self.X_train, self.Y_train = self.transform_by_predictor_type(self.X_train, Y=Y)

        self.fit_model(params)

        self.Y_pred, self.Y_pred_stddev = self.predict(X)
        self.performance_metrics = pm.all_performance_metrics(Y, self.Y_pred)

        self.Y_pred_error = 1.96 * self.Y_pred_stddev

    def fit_model(self, params: Dict[str, Any]) -> None:

        X_tensor = torch.tensor(
            self.X_train.astype(np.float32),
            device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu"),
        )
        Y_tensor = torch.tensor(
            self.Y_train.astype(np.float32),
            device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu"),
        ).flatten()

        self.n_iterations = int(params.get("n_iterations", 50))
        self.kernel = params.get("kernel", ScaleKernel(MaternKernel(nu=2.5)))

        self.likelihood = GaussianLikelihood()
        self.model = ExactGPModel(X_tensor, Y_tensor, self.likelihood, copy.deepcopy(self.kernel))
        if torch.cuda.is_available():
            self.likelihood = self.likelihood.to("cuda")
            self.model = self.model.to("cuda")

        self.model.train()
        self.likelihood.train()

        self.optimizer = torch.optim.Adam(self.model.parameters())

        mll = gpytorch.mlls.ExactMarginalLogLikelihood(self.likelihood, self.model)

        for i in range(self.n_iterations):
            self.optimizer.zero_grad()
            output = self.model(X_tensor)
            mll_loss = -mll(output, Y_tensor).sum()
            mll_loss.backward()
            self.optimizer.step()

        self.model.eval()
        self.likelihood.eval()

    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)

        # Split to reduce memory requirements
        split_size = 50000
        X_split = np.split(X, np.array((range(1, int(X.shape[0] / split_size) + 1))) * split_size)
        observed_Y_pred_list = []
        for X_split_each in X_split:
            if X_split_each.shape[0] == 0:
                continue
            with torch.no_grad():
                observed_Y_pred = self.likelihood(
                    self.model(
                        torch.tensor(
                            X_split_each.astype(np.float32),
                            device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu"),
                        )
                    )
                )
            observed_Y_pred_list.append(observed_Y_pred)

        Y_pred = np.hstack([observed_Y_pred.mean.cpu().detach().numpy() for observed_Y_pred in observed_Y_pred_list])
        Y_pred_upper = np.hstack(
            [observed_Y_pred.confidence_region()[1].cpu().detach().numpy() for observed_Y_pred in observed_Y_pred_list]
        )

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

        return Y_pred.flatten(), Y_pred_stddev.flatten()

    @staticmethod
    def default_params(n_var, gp_kernel_choices: Optional[List[Kernel]] = None) -> Dict[str, Any]:
        if gp_kernel_choices is None:
            gp_kernel_choices = [
                AdditiveStructureKernel(
                    base_kernel=ScaleKernel(MaternKernel(ard_num_dims=n_var, nu=0.5)),
                    num_dims=n_var,
                ),
                AdditiveStructureKernel(
                    base_kernel=ScaleKernel(MaternKernel(ard_num_dims=n_var, nu=1.5)),
                    num_dims=n_var,
                ),
                AdditiveStructureKernel(
                    base_kernel=ScaleKernel(MaternKernel(ard_num_dims=n_var, nu=2.5)),
                    num_dims=n_var,
                ),
                ScaleKernel(MaternKernel(nu=0.5, ard_num_dims=n_var)),
                ScaleKernel(MaternKernel(nu=1.5, ard_num_dims=n_var)),
                ScaleKernel(MaternKernel(nu=2.5, ard_num_dims=n_var)),
            ]
        return {
            "n_iterations": hp.quniform("n_iterations", 5, 500, 1),
            "kernel": hp.choice("kernel", gp_kernel_choices),
        }

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

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

        if predictor_params_dict is None:
            predictor_params_dict = self.default_params(self.variables.n_var, gp_kernel_choices)
        else:
            predictor_params_dict = predictor_params_dict[self.name]

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

        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, **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,
        )

        # HyperOpt returns an integer for the best GP kernel. This links back to the actual kernel.
        gp_kernel_choices = [i.obj for i in predictor_params_dict["kernel"].pos_args if isinstance(i, Literal)]
        model_params["kernel"] = gp_kernel_choices[model_params["kernel"]]

        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
