"""
This module implements the GASMOTE method.
"""

import numpy as np

from sklearn.model_selection import KFold
from sklearn.tree import DecisionTreeClassifier

from ..base import coalesce, coalesce_dict
from ..base import NearestNeighborsWithMetricTensor
from ..base import OverSamplingSimplex
from .._logger import logger

_logger = logger

__all__ = ["GASMOTE"]


class GASMOTE(OverSamplingSimplex):
    """
    References:
        * BibTex::

            @Article{gasmote,
                        author="Jiang, Kun
                        and Lu, Jing
                        and Xia, Kuiliang",
                        title="A Novel Algorithm for Imbalance Data
                                Classification Based on Genetic
                                Algorithm Improved SMOTE",
                        journal="Arabian Journal for Science and
                                    Engineering",
                        year="2016",
                        month="Aug",
                        day="01",
                        volume="41",
                        number="8",
                        pages="3255--3266",
                        issn="2191-4281",
                        doi="10.1007/s13369-016-2179-2",
                        url="https://doi.org/10.1007/s13369-016-2179-2"
                        }
    """

    categories = [
        OverSamplingSimplex.cat_extensive,
        OverSamplingSimplex.cat_memetic,
        OverSamplingSimplex.cat_sample_ordinary,
        OverSamplingSimplex.cat_metric_learning,
    ]

    def __init__(
        self,
        n_neighbors=5,
        *,
        nn_params=None,
        ss_params=None,
        maxn=7,
        n_pop=10,
        popl3=5,
        pm=0.3,
        pr=0.2,
        Ge=10,
        n_jobs=1,
        random_state=None,
        **_kwargs
    ):
        """
        Constructor of the sampling object

        Args:
            n_neighbors (int): number of neighbors
            nn_params (dict): additional parameters for nearest neighbor calculations, any
                                parameter NearestNeighbors accepts, and additionally use
                                {'metric': 'precomputed', 'metric_learning': '<method>', ...}
                                with <method> in 'ITML', 'LSML' to enable the learning of
                                the metric to be used for neighborhood calculations
            ss_params (dict): simplex sampling parameters
            maxn (int): maximum number of samples to generate per minority
                        instances
            n_pop (int): size of population
            popl3 (int): number of crossovers
            pm (float): mutation probability
            pr (float): selection probability
            Ge (int): number of generations
            n_jobs (int): number of parallel jobs
            random_state (int/RandomState/None): initializer of random_state,
                                                    like in sklearn
        """
        nn_params = coalesce(nn_params, {})

        ss_params_default = {
            "n_dim": 2,
            "simplex_sampling": "random",
            "within_simplex_sampling": "random",
            "gaussian_component": None,
        }
        ss_params = coalesce_dict(ss_params, ss_params_default)

        super().__init__(**ss_params, random_state=random_state)
        self.check_greater_or_equal(n_neighbors, "n_neighbors", 1)
        self.check_greater_or_equal(maxn, "maxn", 1)
        self.check_greater_or_equal(n_pop, "n_pop", 1)
        self.check_in_range(pm, "pm", [0, 1])
        self.check_in_range(pr, "pr", [0, 1])
        self.check_greater_or_equal(Ge, "Ge", 1)
        self.check_n_jobs(n_jobs, "n_jobs")

        self.n_neighbors = n_neighbors
        self.nn_params = nn_params
        self.ga_params = {
            "maxn": maxn,
            "n_pop": n_pop,
            "popl3": popl3,
            "pm": pm,
            "pr": pr,
            "Ge": Ge,
        }
        self.n_jobs = n_jobs

    @classmethod
    def parameter_combinations(cls, raw=False):
        """
        Generates reasonable parameter combinations.

        Returns:
            list(dict): a list of meaningful parameter combinations
        """
        return cls.generate_parameter_combinations(
            {
                "n_neighbors": [7],
                "maxn": [2, 3, 4],
                "n_pop": [10],
                "popl3": [4],
                "pm": [0.3],
                "pr": [0.2],
                "Ge": [10],
            },
            raw,
        )

    def samples_from_conf(self, X_min, ind, conf):
        """
        Generate samples from a configuration.

        Args:
            X_min (np.array): the minority samples
            ind (np.array): the neighborhoods
            conf (np.array): a configuration

        Returns:
            np.array: the generated samples
        """
        samples = [np.zeros(shape=(0, X_min.shape[1]), dtype=float)]
        for idx, conf_i in enumerate(conf):
            samples.append(
                self.sample_simplex(
                    X=X_min[[idx]],
                    indices=ind[[idx]],
                    n_to_sample=conf_i,
                    X_vertices=X_min,
                )
            )
        return np.vstack(samples)

    def calculate_score(self, tests, preds):
        """
        Calculate the objective function score

        Args:
            tests (np.array): the test labels
            preds (np.array): the predicted labels

        Returns:
            float: the score
        """
        # compute fitness measure
        t_p = np.sum(np.logical_and(tests == self.min_label, tests == preds))
        t_n = np.sum(np.logical_and(tests == self.maj_label, tests == preds))
        f_p = np.sum(np.logical_and(tests == self.maj_label, tests != preds))
        f_n = np.sum(np.logical_and(tests == self.min_label, tests != preds))
        sens = t_p / (t_p + f_n)
        spec = t_n / (f_p + t_n)

        return np.sqrt(sens * spec)

    def fitness(self, *, conf, X_min, ind, X, y, kfold):
        """
        Evluate fitness of configuration

        Args:
            conf (list(list)): configuration
            X_min (np.array): the minority samples
            ind (np.array): the neighborhood relation
            X (np.array): all vectors
            y (np.array): all target labels
            kfold (obj): cross validation object

        Returns:
            float: the fitness score
        """
        # generate new samples
        samples = self.samples_from_conf(X_min, ind, conf)

        # samples = []
        # for idx, conf_i in enumerate(conf):
        #    for _ in range(conf_i):
        #        X_b = X_min[self.random_state.choice(ind[idx][1:])]
        #        samples.append(self.sample_between_points(X_min[idx], X_b))

        X_new = np.vstack([X, samples])
        y_new = np.hstack([y, np.repeat(self.min_label, len(samples))])

        # execute kfold cross validation
        preds, tests = [], []
        for train, test in kfold.split(X_new):
            clas = DecisionTreeClassifier(random_state=self._random_state_init)
            clas.fit(X_new[train], y_new[train])
            preds.append(clas.predict(X_new[test]))
            tests.append(y_new[test])
        preds = np.hstack(preds)
        tests = np.hstack(tests)

        return self.calculate_score(tests, preds)

    def crossover(self, conf_a, conf_b):
        """
        Crossover

        Args:
            conf_a (list): configuration to crossover
            conf_b (list): configuration to crossover

        Returns:
            list, list: the configurations after crossover
        """
        for _ in range(self.ga_params["popl3"]):
            kdx = self.random_state.randint(len(conf_a))
            conf_a = np.hstack([conf_a[:kdx], conf_b[kdx:]])
            conf_b = np.hstack([conf_b[:kdx], conf_a[kdx:]])
        return conf_a, conf_b

    def mutation(self, conf, ge_):
        """
        Mutation

        Args:
            conf (list): configuration to mutate
            ge (int): iteration number

        Returns:
            list: the mutated configuration
        """
        conf = conf.copy()
        if self.random_state.random_sample() >= self.ga_params["pm"]:
            for idx, conf_i in enumerate(conf):
                rand = self.random_state.random_sample()
                rand = rand ** ((1 - ge_ / self.ga_params["Ge"]) ** 3)
                if self.random_state.randint(2) == 0:
                    tmp = (self.ga_params["maxn"] - conf_i) * rand
                    conf[idx] = int(conf_i + tmp)
                else:
                    conf[idx] = int(conf_i - (conf_i - 0) * rand)
        return conf

    # generate initial population
    def init_pop(self, X_min):
        """
        Initialize the population

        Args:
            X_min (np.array): the minority vectors

        Returns:
            np.array: an initial population
        """
        return self.random_state.randint(self.ga_params["maxn"], size=X_min.shape[0])

    def genetic_algorithm(self, *, X_min, ind, X, y, kfold):
        """
        The genetic algorithm

        Args:
            X_min (np.array): the minority vectors
            ind (np.array): the neighborhood relations
            X (np.array): all feature vectors
            y (np.array): all target labels
            kfold (obj): a k-fold cross-validation object
        """
        population = [[self.init_pop(X_min), 0] for _ in range(self.ga_params["n_pop"])]

        # calculate fitness values
        for pop in population:
            pop[1] = self.fitness(
                conf=pop[0], X_min=X_min, ind=ind, X=X, y=y, kfold=kfold
            )

        # start iteration
        ge_ = 0
        while ge_ < self.ga_params["Ge"]:
            # sorting population in descending order by fitness scores
            population = sorted(population, key=lambda x: -x[1])

            # selection operation (Step 2)
            n_keep = int(self.ga_params["n_pop"] * self.ga_params["pr"])
            population_new = []
            population_new.extend(population[:n_keep])
            population_new.extend(population[: (self.ga_params["n_pop"] - n_keep)])
            population = population_new

            # crossover
            for _ in range(int(self.ga_params["n_pop"] / 2)):
                pop_0 = population[self.random_state.randint(self.ga_params["n_pop"])][
                    0
                ]
                pop_1 = population[self.random_state.randint(self.ga_params["n_pop"])][
                    0
                ]
                conf_a, conf_b = self.crossover(pop_0, pop_1)
                population.append(
                    [
                        conf_a,
                        self.fitness(
                            conf=conf_a, X_min=X_min, ind=ind, X=X, y=y, kfold=kfold
                        ),
                    ]
                )
                population.append(
                    [
                        conf_b,
                        self.fitness(
                            conf=conf_b, X_min=X_min, ind=ind, X=X, y=y, kfold=kfold
                        ),
                    ]
                )

            # mutation
            for _ in range(int(self.ga_params["n_pop"] / 2)):
                pop_0 = population[self.random_state.randint(self.ga_params["n_pop"])][
                    0
                ]
                conf_a = self.mutation(pop_0, ge_)
                population.append(
                    [
                        conf_a,
                        self.fitness(
                            conf=conf_a, X_min=X_min, ind=ind, X=X, y=y, kfold=kfold
                        ),
                    ]
                )

            ge_ = ge_ + 1

        # sorting final population
        population = sorted(population, key=lambda x: -x[1])

        # return best configuration
        return population[0][0]

    def sampling_algorithm(self, X, y):
        """
        Does the sample generation according to the class parameters.

        Args:
            X (np.ndarray): training set
            y (np.array): target labels

        Returns:
            (np.ndarray, np.array): the extended training set and target labels
        """
        X_min = X[y == self.min_label]

        # fitting nearest neighbors model to find minority neighbors of
        #  minority samples
        n_neighbors = min([self.n_neighbors + 1, len(X_min)])

        nn_params = {**self.nn_params}
        nn_params["metric_tensor"] = self.metric_tensor_from_nn_params(nn_params, X, y)

        nnmt = NearestNeighborsWithMetricTensor(
            n_neighbors=n_neighbors, n_jobs=self.n_jobs, **(nn_params)
        )
        nnmt.fit(X_min)
        ind = nnmt.kneighbors(X_min, return_distance=False)
        kfold = KFold(min([len(X), 5]))

        conf = self.genetic_algorithm(X_min=X_min, ind=ind, X=X, y=y, kfold=kfold)

        # generate final samples
        samples = self.samples_from_conf(X_min, ind, conf)

        return (
            np.vstack([X, samples]),
            np.hstack([y, np.repeat(self.min_label, len(samples))]),
        )

    def get_params(self, deep=False):
        """
        Returns:
            dict: the parameters of the current sampling object
        """
        return {
            "n_neighbors": self.n_neighbors,
            "nn_params": self.nn_params,
            "maxn": self.ga_params["maxn"],
            "n_pop": self.ga_params["n_pop"],
            "popl3": self.ga_params["popl3"],
            "pm": self.ga_params["pm"],
            "pr": self.ga_params["pr"],
            "Ge": self.ga_params["Ge"],
            "n_jobs": self.n_jobs,
            **OverSamplingSimplex.get_params(self),
        }
