import math
import os
import random
from attrs import define, field
from typing import List, Optional, Union

import numpy as np
import torch
from botorch.utils.sampling import get_polytope_samples
from scipy.stats import qmc

import nemo_bo.utils.logger as logging_nemo
from nemo_bo.opt.constraints import ConstraintsList
from nemo_bo.opt.utils_samplers.sampling import *
from nemo_bo.opt.variables import ContinuousVariable, 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__))


@define
class SampleGenerator:
    """

    Base class for the sample generators

    Parameters
    ----------
    num_new_points: int = None
        The number of samples to generate

    """

    num_new_points: Optional[int] = None

    def generate_samples(self, *args, **kwargs) -> None:
        """

        Function that returns an X array of self.num_new_points number of samples

        """
        pass


@define
class PolytopeSampling(SampleGenerator):
    """

    Class that is used to generate samples from a polytope that is limited by the variables lower/upper bounds and
    constraints.  The sampler uses the implementation from the BoTorch package:
    M. Balandat, B. Karrer, D. R. Jiang, S. Daulton, B. Letham, A. G. Wilson, and E. Bakshy. BoTorch: A Framework for
    Efficient Monte-Carlo Bayesian Optimization. Advances in Neural Information Processing Systems 33, 2020.
    http://arxiv.org/abs/1910.06403
    https://github.com/pytorch/botorch
    https://botorch.org/

    """

    def generate_samples(
        self,
        variables: VariablesList,
        constraints: Optional[ConstraintsList] = None,
        seed: int = None,
        n_burnin: int = 10000,
        thinning: int = 32,
    ) -> np.ndarray:
        """

        Parameters
        ----------
        variables: VariablesList
            VariablesList object that contains information about all variables
        constraints: ConstraintsList, Default = None,
            ConstraintsList object that contains information about all constraints
        seed: int = None,
            Seed for a random number generator for the manual_seed function in botorch.utils.sampling
        n_burnin: int = 10000,
            The number of burn-in samples for the Markov chain sampler, where it drops the first n_burnin number of
            sampled points
        thinning: int = 32,
            The amount of thinning from the total number of generated by the hit-and-run Markov chain sampler

        Returns
        -------
        X: np.ndarray
            Array of generated samples that are within the defined bounds and passes all constraints

        """
        kwargs = {"seed": seed, "n_burnin": n_burnin, "thinning": thinning}
        if constraints is not None:
            for constr in constraints:
                if constr.constraint_type == "ineq":
                    kwargs["inequality_constraints"] = [constr.build_polytope_sampler_constraint()]
                else:
                    kwargs["equality_constraints"] = [constr.build_polytope_sampler_constraint()]

        X = (
            get_polytope_samples(
                n=self.num_new_points,
                bounds=torch.tensor([variables.lower_bounds, variables.upper_bounds]),
                **kwargs,
            )
            .cpu()
            .numpy()
        )

        if variables.num_cat_var > 0:
            return variables.descriptor_to_name(X)
        else:
            return X


@define
class LatinHyperCubeSampling(SampleGenerator):
    """

    Class that is used to generate samples via Latin Hypercube Sampling for continuous variables or a mixture of
    continuous and categorical variables. Samples generated will be limited by the variables lower/upper bounds,
    categorical variable levels, and constraints.  This sampler uses the implementation from the smt package by M. A.
    Bouhlel and J. T. Hwang and N. Bartoli and R. Lafage and J. Morlier and J. R. R. A. Martins, A Python surrogate
    modeling framework with derivatives. Advances in Engineering Software, 2019, 102662
    https://doi.org/10.1016/j.advengsoft.2019.03.005
    https://smt.readthedocs.io/en/latest/

    """

    def generate_samples(
        self, variables: VariablesList, constraints: Optional[ConstraintsList] = None, **kwargs
    ) -> np.ndarray:
        """

        Parameters
        ----------
        variables: VariablesList
            VariablesList object that contains information about all variables
        constraints: ConstraintsList, Default = None,
            ConstraintsList object that contains information about all constraints

        Returns
        -------
        X: np.ndarray
            Array of generated samples that are within the defined bounds and passes all constraints

        """
        xlimits = np.array(
            [
                [var.lower_bound, var.upper_bound] if isinstance(var, ContinuousVariable) else var.categorical_levels
                for var in variables.variables
            ],
            dtype=object,
        )

        if variables.num_cat_var > 0:
            xtypes = [
                FLOAT if isinstance(var, ContinuousVariable) else (ENUM, var.num_categorical_levels)
                for var in variables.variables
            ]
            sampler = MixedIntegerSamplingMethod(xtypes, xlimits, LHS, random_state=random.randint(0, 1000000))

            X = sampler(self.num_new_points)

            X = X.astype(np.object)
            for var_index, var in enumerate(variables.variables):
                if not isinstance(var, ContinuousVariable):
                    X[:, var_index] = var.enum_to_cat_level(X[:, var_index])

            # To do: Improve the implementation for constraints for LHS.
            # Currently this does not sample the available space in a guaranteed LHS manner
            if constraints is not None:
                num_points_temp = self.num_new_points
                X = X[constraints.bool_evaluate_constraints(X)]
                while X.shape[0] < self.num_new_points:
                    num_points_temp = math.ceil(num_points_temp * 1.2)
                    X = sampler(num_points_temp)

                    X = X.astype(np.object)
                    for var_index, var in enumerate(variables.variables):
                        if not isinstance(var, ContinuousVariable):
                            X[:, var_index] = var.enum_to_cat_level(X[:, var_index])

                    X = X[constraints.bool_evaluate_constraints(X)]

            logger.info(f"Generated {X.shape[0]} mixed integer Latin hypercube sampling points in the variable space")

        else:
            sampler = LHS(xlimits=xlimits, **kwargs)

            X = sampler(self.num_new_points)

            # To do: Improve the implementation for constraints for LHS.
            # Currently this does not sample the available space in a guaranteed LHS manner after constraints are
            # applied
            if constraints is not None:
                num_points_temp = self.num_new_points
                X = X[constraints.bool_evaluate_constraints(X)]
                while X.shape[0] < self.num_new_points:
                    num_points_temp = math.ceil(num_points_temp * 1.2)
                    X = sampler(num_points_temp)
                    X = X[constraints.bool_evaluate_constraints(X)]

            logger.info(f"Generated {X.shape[0]} continuous Latin hypercube sampling points in the variable space")

        return X


@define
class RandomSampling(SampleGenerator):
    """

    Class that is used to generate samples via Random Sampling for continuous variables or a mixture of
    continuous and categorical variables. Samples generated will be limited by the variables lower/upper bounds,
    categorical variable levels, and constraints.  This sampler uses the implementation from the smt package by M. A.
    Bouhlel and J. T. Hwang and N. Bartoli and R. Lafage and J. Morlier and J. R. R. A. Martins, A Python surrogate
    modeling framework with derivatives. Advances in Engineering Software, 2019, 102662
    https://doi.org/10.1016/j.advengsoft.2019.03.005
    https://smt.readthedocs.io/en/latest/

    """

    def generate_samples(
        self, variables: VariablesList, constraints: Optional[ConstraintsList] = None, **kwargs
    ) -> np.ndarray:
        """

        Parameters
        ----------
        variables: VariablesList
            VariablesList object that contains information about all variables
        constraints: ConstraintsList, Default = None,
            ConstraintsList object that contains information about all constraints

        Returns
        -------
        X: np.ndarray
            Array of generated samples that are within the defined bounds and passes all constraints

        """
        xlimits = np.array(
            [
                [var.lower_bound, var.upper_bound] if isinstance(var, ContinuousVariable) else var.categorical_levels
                for var in variables.variables
            ],
            dtype=object,
        )

        if variables.num_cat_var > 0:
            logger.info(f"Generating {self.num_new_points} mixed integer random sampling points in the variable space")
            xtypes = [
                FLOAT if isinstance(var, ContinuousVariable) else (ENUM, var.num_categorical_levels)
                for var in variables.variables
            ]
            sampler = MixedIntegerSamplingMethod(xtypes, xlimits, Random)
            X = sampler(self.num_new_points)

            X = X.astype(np.object)
            for var_index, var in enumerate(variables.variables):
                if not isinstance(var, ContinuousVariable):
                    X[:, var_index] = var.enum_to_cat_level(X[:, var_index])

            if constraints is not None:
                num_points_temp = self.num_new_points
                X = X[constraints.bool_evaluate_constraints(X)]
                while X.shape[0] < self.num_new_points:
                    num_points_temp = math.ceil(num_points_temp * 1.2)
                    X = sampler(num_points_temp)

                    X = X.astype(np.object)
                    for var_index, var in enumerate(variables.variables):
                        if not isinstance(var, ContinuousVariable):
                            X[:, var_index] = var.enum_to_cat_level(X[:, var_index])

                    X = X[constraints.bool_evaluate_constraints(X)]

        else:
            logger.info(f"Generating {self.num_new_points} continuous random sampling points in the variable space")
            sampler = Random(xlimits=xlimits, **kwargs)
            X = sampler(self.num_new_points)

            if constraints is not None:
                num_points_temp = self.num_new_points
                X = X[constraints.bool_evaluate_constraints(X)]
                while X.shape[0] < self.num_new_points:
                    num_points_temp = math.ceil(num_points_temp * 1.2)
                    X = sampler(num_points_temp)
                    X = X[constraints.bool_evaluate_constraints(X)]

        return X


@define
class SobolSampling(SampleGenerator):
    """

    Class that is used to generate scrambled Sobol sequences for continuous variables. Generator descriptor values will
    be converted to their corresponding categorical variables. Samples generated will be limited by the variables
    lower/upper bounds and constraints

    """

    def generate_samples(self, variables: VariablesList, constraints: Optional[ConstraintsList] = None) -> np.ndarray:
        """

        Parameters
        ----------
        variables: VariablesList
            VariablesList object that contains information about all variables
        constraints: ConstraintsList, Default = None,
            ConstraintsList object that contains information about all constraints

        Returns
        -------
        X: np.ndarray
            Array of generated samples that are within the defined bounds and passes all constraints

        """
        sampler = qmc.Sobol(d=variables.n_var)

        num_samples_power = math.ceil(math.log2(self.num_new_points))
        if 2**num_samples_power != self.num_new_points:
            logger.info(
                f"Generating {2 ** num_samples_power} sobol sampling points in the variable "
                f"space (different to the desired {self.num_new_points} to guarantee the balance properties "
                f"of the sobol sequence)."
            )

        X = sampler.random_base2(m=num_samples_power)
        X = qmc.scale(X, variables.lower_bounds, variables.upper_bounds)

        # To do: Improve the implementation for constraints for Sobol sampling.
        # Currently this does not guarantee the balance properties of the sobol sequence after constraints are applied
        if constraints is not None:
            X = X[constraints.bool_evaluate_constraints(X)]
            while X.shape[0] < self.num_new_points:
                num_samples_power += 1
                sampler.reset()
                X = sampler.random_base2(m=num_samples_power)
                X = qmc.scale(X, variables.lower_bounds, variables.upper_bounds)
                X = X[constraints.bool_evaluate_constraints(X)]

            logger.info(f"Retained {X.shape[0]} sobol sampling points in the variable space after applying constraints")

        if variables.num_cat_var > 0:
            return variables.descriptor_to_name(X)
        else:
            return X


@define
class PoolBased:
    """

    Class that is used to store and take samples from designated X and Y arrays (a pool of samples). For simulated
    closed-loop optimisation problems, this is typically used to obtain true results whereas the alternative benchmark
    function would return predicted results instead

    Parameters
    ----------
    X_pool: np.ndarray
        The X array to be used in the sample pool
    Y_pool: np.ndarray
        The Y array to be used in the sample pool

    """

    X_pool: np.ndarray
    Y_pool: np.ndarray
    index: str = field(init=False)

    def __attrs_post_init__(self):
        self.index = None

    def generate_samples(self) -> np.ndarray:
        """

        Returns
        -------
        np.ndarray
            Array of X-values that are in the user-supplied pool

        """

        return self.X_pool

    def Y_and_update_pool(self, index: Union[int, List[int]]) -> np.ndarray:
        """

        Uses the index to select Y-values, delete them from the Y sample pool and return them

        Parameters
        ----------
        index: int, List[int]
            Index/indexes used to reference the Y sample pool

        Returns
        -------
        Y: np.ndarray
            The selected Y-values based on the index/indexes

        """
        if isinstance(index, int):
            index = [index]

        # Identifies the Y values to be returned
        Y = np.array([self.Y_pool[i] for i in index])

        # Removes the samples in the sample pool based on the supplied index/indexes
        self.X_pool = np.delete(self.X_pool, index, 0)
        self.Y_pool = np.delete(self.Y_pool, index, 0)

        return Y
