"""
This module implements the SMOTE_PSOBAT 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 OverSampling
from ._smote import SMOTE

from .._logger import logger

_logger = logger

__all__ = ["SMOTE_PSOBAT"]


class SMOTE_PSOBAT(OverSampling):
    """
    References:
        * BibTex::

            @INPROCEEDINGS{smote_psobat,
                            author={Li, J. and Fong, S. and Zhuang, Y.},
                            booktitle={2015 3rd International Symposium on
                                        Computational and Business
                                        Intelligence (ISCBI)},
                            title={Optimizing SMOTE by Metaheuristics with
                                    Neural Network and Decision Tree},
                            year={2015},
                            volume={},
                            number={},
                            pages={26-32},
                            keywords={data mining;particle swarm
                                        optimisation;pattern classification;
                                        data mining;classifier;metaherustics;
                                        SMOTE parameters;performance
                                        indicators;selection optimization;
                                        PSO;particle swarm optimization
                                        algorithm;BAT;bat-inspired algorithm;
                                        metaheuristic optimization algorithms;
                                        nearest neighbors;imbalanced dataset
                                        problem;synthetic minority
                                        over-sampling technique;decision tree;
                                        neural network;Classification
                                        algorithms;Neural networks;Decision
                                        trees;Training;Optimization;Particle
                                        swarm optimization;Data mining;SMOTE;
                                        Swarm Intelligence;parameter
                                        selection optimization},
                            doi={10.1109/ISCBI.2015.12},
                            ISSN={},
                            month={Dec}}

    Notes:
        * The parameters of the memetic algorithms are not specified.
        * I have checked multiple paper describing the BAT algorithm, but the
            meaning of "Generate a new solution by flying randomly" is still
            unclear.
        * It is also unclear if best solutions are recorded for each bat, or
            the entire population.
    """

    categories = [
        OverSampling.cat_extensive,
        OverSampling.cat_uses_clustering,
        OverSampling.cat_sample_ordinary,
        OverSampling.cat_memetic,
        OverSampling.cat_metric_learning,
    ]

    def __init__(
        self,
        *,
        maxit=50,
        c1=0.3,
        c2=0.1,
        c3=0.1,
        alpha=0.9,
        gamma=0.9,
        method="bat",
        nn_params=None,
        ss_params=None,
        n_jobs=1,
        random_state=None,
        **_kwargs
    ):
        """
        Constructor of the sampling object

        Args:
            maxit (int): maximum number of iterations
            c1 (float): intertia weight of PSO
            c2 (float): attraction of local maximums in PSO
            c3 (float): attraction of global maximum in PSO
            alpha (float): alpha parameter of the method
            gamma (float): gamma parameter of the method
            method (str): optimization technique to be used
            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
            n_jobs (int): number of parallel jobs
            random_state (int/RandomState/None): initializer of random_state,
                                                    like in sklearn
        """
        ss_params_default = {
            "n_dim": 2,
            "simplex_sampling": "random",
            "within_simplex_sampling": "random",
            "gaussian_component": None,
        }

        super().__init__(random_state=random_state)
        self.check_greater_or_equal(maxit, "maxit", 1)
        self.check_greater_or_equal(c1, "c1", 0)
        self.check_greater_or_equal(c2, "c2", 0)
        self.check_greater_or_equal(c3, "c3", 0)
        self.check_greater_or_equal(alpha, "alpha", 0)
        self.check_greater_or_equal(gamma, "gamma", 0)
        self.check_isin(method, "method", ["pso", "bat"])
        self.check_n_jobs(n_jobs, "n_jobs")

        self.params = {
            "maxit": maxit,
            "c1": c1,
            "c2": c2,
            "c3": c3,
            "alpha": alpha,
            "gamma": gamma,
            "n_pop": 10,
            "n_keep": 5,
            "f_max": 10,
        }
        self.method = method
        self.nn_params = coalesce(nn_params, {})
        self.ss_params = coalesce_dict(ss_params, ss_params_default)
        self.n_jobs = n_jobs

        # a reasonable range of proportions
        self.proportion_range = np.array([0.1, 2.0])
        self.k_range = np.array([1, 10])

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

        Returns:
            list(dict): a list of meaningful parameter combinations
        """
        bat_pc = cls.generate_parameter_combinations(
            {
                "maxit": [50],
                "alpha": [0.7, 0.9],
                "gamma": [0.7, 0.9],
                "method": ["bat"],
            },
            raw,
        )
        pso_pc = cls.generate_parameter_combinations(
            {
                "maxit": [50],
                "c1": [0.2, 0.5],
                "c2": [0.1, 0.2],
                "c3": [0.1, 0.2],
                "method": ["pso"],
            },
            raw,
        )
        if not raw:
            bat_pc.extend(pso_pc)
            return bat_pc

        return {
            "maxit": [50],
            "alpha": [0.7, 0.9],
            "gamma": [0.7, 0.9],
            "method": ["bat", "pso"],
            "c1": [0.2, 0.5],
            "c2": [0.1, 0.2],
            "c3": [0.1, 0.2],
        }

    def pred_test_vectors(self, *, proportion, K, nn_params, X, y):
        """
        Determine the prediction and test vectors.

        Args:
            proportion (float): the proportion to sample
            K (int): the number of neighbors
            nn_params (dict): the nearest neighbor parameters
            X (np.array): all training vectors
            y (np.array): all target labels

        Returns:
            np.array, np.array: the predictions and the test labels
        """
        ss_params = self.ss_params.copy()
        ss_params["n_dim"] = np.min([ss_params["n_dim"], int(K)])

        X_samp, y_samp = SMOTE(
            proportion=proportion,
            n_neighbors=K,
            nn_params=nn_params,
            ss_params=ss_params,
            n_jobs=self.n_jobs,
            random_state=self._random_state_init,
        ).sample(X, y)

        # doing k-fold cross validation
        kfold = KFold(5)
        preds = []
        tests = []
        for train, test in kfold.split(X_samp):
            dtree = DecisionTreeClassifier(random_state=self._random_state_init)
            dtree.fit(X_samp[train], y_samp[train])
            preds.append(dtree.predict(X_samp[test]))
            tests.append(y_samp[test])
        preds = np.hstack(preds)
        tests = np.hstack(tests)

        return preds, tests

    def calculate_scores(self, truep, truen, falsep, falsen):
        """
        Calculate the target scores.

        Args:
            truep (int): the number of true positives
            truen (int): the number of true negatives
            falsep (int): the number of false positives
            falsen (int): the number of false negatives

        Returns:
            float, float: kappa, accuracy
        """
        p_o = (truep + truen) / (truep + falsen + truen + falsep)
        term_0 = (
            (truep + falsen) * (truep + falsep) / (truep + falsen + truen + falsep) ** 2
        )
        term_1 = (
            (falsep + truen) * (falsen + truen) / (truep + falsen + truen + falsep) ** 2
        )
        p_e = term_0 + term_1

        kappa = (p_o - p_e) / (1.0 - p_e)

        return kappa, p_o

    def evaluate(self, *, K, proportion, nn_params, X, y):
        """
        Evaluate given configuration

        Args:
            K (int): number of neighbors in nearest neighbors component
            proportion (float): proportion of missing data to generate
            nn_params (dict): nearest neighbor parameters
            X (np.array): the training vectors
            y (np.array): the target labels

        Returns:
            float, float: kappa and accuracy scores
        """
        preds, tests = self.pred_test_vectors(
            proportion=proportion, K=K, nn_params=nn_params, X=X, y=y
        )

        # computing the kappa score
        truep = np.sum(np.logical_and(preds == tests, tests == self.min_label))
        falsen = np.sum(np.logical_and(preds != tests, tests == self.min_label))
        truen = np.sum(np.logical_and(preds == tests, tests == self.maj_label))
        falsep = np.sum(np.logical_and(preds != tests, tests == self.maj_label))

        return self.calculate_scores(truep, truen, falsep, falsen)

    def init_particle_or_bat(self):
        """
        Initialize a PSO particle.

        Args:
            k_range (tuple): the range of k values
            proportion_range (tuple): the range of proportion values

        Returns:
            np.array: an initial particle
        """
        k_rand = self.random_state.randint(self.k_range[0], self.k_range[1])
        rand = self.random_state.random_sample()
        diff = self.proportion_range[1] - self.proportion_range[0]
        vect = rand * diff + self.proportion_range[0]
        return np.array([k_rand, vect])

    def init_population(self):
        """
        Initialize the population.

        Returns:
            np.array: the initial population
        """
        return np.array(
            [self.init_particle_or_bat() for _ in range(self.params["n_pop"])]
        )

    def pso_update_local_scores(self, scores, particles, local_scores, local_best):
        """
        Update the local scores according to the rule in the paper.

        Args:
            scores (np.array): the new scores
            particles (np.array): the particles
            local_scores (np.array): the local best scores
            local_best (np.array): the local best particles

        Returns:
            np.array, np.array, bool: local scores, local best particles,
                                        not changed flag
        """
        not_changed = True

        mask = (
            np.min(np.vstack([local_scores[:, 0], scores[:, 0]]).T, axis=1) > 0.4
        ) & (local_scores[:, 1] > scores[:, 1])

        local_scores[mask] = scores[mask].copy()
        local_best[mask] = particles[mask].copy()
        if np.any(mask):
            not_changed = False

        mask = (~mask) & (scores[:, 0] > 0.4) & (local_scores[:, 0] <= 0.4)
        local_scores[mask] = scores[mask].copy()
        local_best[mask] = particles[mask].copy()

        if np.any(mask):
            not_changed = False

        return local_scores, local_best, not_changed

    def pso_update_global_scores(self, scores, particles, global_scores, global_best):
        """
        Update global scores for PSO

        Args:
            scores (np.array): the actual scores
            particles (np.array): the particles
            global_scores (np.array): the global scores
            global_best (np.array): the globally best particle

        Returns:
            np.array, np.array, bool: the global best score, best particle,
                                        the not changed flag
        """
        not_changed = True
        mask = (
            np.min(
                np.vstack([np.repeat(global_scores[0], len(scores)), scores[:, 0]]).T,
                axis=1,
            )
            > 0.4
        ) & (global_scores[1] > scores[:, 1])
        if np.any(mask):
            max_idx = np.argmax(scores[mask, 1])
            global_scores = scores[mask][max_idx].copy()
            global_best = particles[mask][max_idx].copy()
            not_changed = False

        mask = (~mask) & (scores[:, 0] > 0.4) & (global_scores[0] <= 0.4)
        if np.any(mask):
            max_idx = np.argmax(scores[mask, 1])
            global_scores = scores[mask][max_idx].copy()
            global_best = particles[mask][max_idx].copy()
            not_changed = False

        return (global_scores, global_best, not_changed)

    def pso_update_velocities(self, particles, velocities, local_best, global_best):
        """
        Update the PSO velocities

        Args:
            particles (np.array): the particles
            velocities (np.array): the actual velocities
            local_best (np.array): the locally best particles
            global_best (np.array): the globally best particle

        Returns:
            np.array: the updated velocities
        """
        # update velocities
        velocities = (
            self.params["c1"] * velocities
            + (local_best - particles) * self.params["c2"]
            + (global_best - particles) * self.params["c3"]
        )

        ratios = np.abs(velocities[:, 0]) / (self.k_range[1] - self.k_range[0])
        velocities[ratios > 1.0, 0] = velocities[ratios > 1.0, 0] / np.ceil(
            ratios[ratios > 1.0]
        )

        proportion_diff = self.proportion_range[1] - self.proportion_range[0]
        ratios = np.abs(velocities[:, 1]) / proportion_diff
        velocities[ratios > 1.0, 1] = velocities[ratios > 1.0, 1] / np.ceil(
            ratios[ratios > 1.0]
        )

        return velocities

    def pso_vectorized(self, X, y, nn_params):
        """
        PSO optimization

        Args:
            X (np.array): all training vectors
            y (np.array): all target labels
            nn_params (dict): the nearest neighbor parameters

        Returns:
            np.array: the best K and proportion values
        """

        # initial particles

        particles = self.init_population()
        # initial velocities
        velocities = np.array([np.array([0, 0]) for _ in range(self.params["n_pop"])])

        local_best = np.array(
            [particles[i].copy() for i in range(self.params["n_pop"])]
        )
        local_scores = np.zeros((self.params["n_pop"], 2))
        global_best = particles[0].copy()
        global_scores = np.array([0.0, 0.0])

        # executing the particle swarm optimization
        not_changed = 0
        for _ in range(self.params["maxit"]):
            # if the configurations didn't change for 10 iterations, stop
            if not_changed > len(particles):
                break

            not_changed = not_changed + 1

            scores = np.array(
                [
                    self.evaluate(
                        K=int(particle[0]),
                        proportion=particle[1],
                        nn_params=nn_params,
                        X=X,
                        y=y,
                    )
                    for particle in particles
                ]
            )

            (local_scores, local_best, changed) = self.pso_update_local_scores(
                scores, particles, local_scores, local_best
            )
            (global_scores, global_best, changed) = self.pso_update_global_scores(
                scores, particles, global_scores, global_best
            )

            not_changed = not_changed * changed

            velocities = self.pso_update_velocities(
                particles, velocities, local_best, global_best
            )

            # update positions
            particles = particles + velocities
            particles[:, 0] = np.clip(particles[:, 0], self.k_range[0], self.k_range[1])
            particles[:, 1] = np.clip(
                particles[:, 1], self.proportion_range[0], self.proportion_range[1]
            )

        return global_best

    def bat_check_global(self, scores, global_scores):
        """
        Check the global performances for the bat optimization

        Args:
            scores (np.array): the actual scores
            global_scores (np.array): the globally best scores

        Returns:
            int: the index maximizing the global score or -1 if no improvement
                 happened
        """
        max_idx = -1
        tmp0 = np.repeat(global_scores[0], len(scores))
        tmp1 = scores[:, 0]

        mask = (np.min(np.vstack([tmp0, tmp1]).T, axis=1) > 0.4) & (
            global_scores[1] > scores[:, 1]
        )
        if np.any(mask):
            max_idx = np.argmax(scores[mask, 1])

        mask = (~mask) & (scores[:, 0] > 0.4) & (global_scores[0] <= 0.4)
        if np.any(mask):
            max_idx = np.argmax(scores[mask, 1])

        return max_idx

    def bat_check_local(self, scores, local_scores):
        """
        Check the performance locally.

        Args:
            scores (np.array): the actual scores
            local_scores (np.array): the local scores

        Returns:
            np.array: the mask of particles that improved
        """
        local_scores = np.array([local[0] for local in local_scores])
        mask = (
            np.min(np.vstack([local_scores[:, 0], scores[:, 0]]).T, axis=1) > 0.4
        ) & (local_scores[:, 1] > scores[:, 1])

        mask = (~mask) & (scores[:, 0] > 0.4) & (local_scores[:, 0] <= 0.4)

        return mask

    def bat_local_search(self, bats, local_best, local_scores, pulse_loudness):
        """
        The local search algorithm.

        Args:
            bats (np.array): the bats
            local_best (np.array): the locally best bats
            local_scores (np.array): the local scores
            pulse_loudness (np.array): the pulse and loudness parameters

        Returns:
            np.array: the updated bats
        """
        n_pop = self.params["n_pop"]

        local_lens = np.sum(local_scores[:, :, 0] >= 0.0, axis=1)
        mask = self.random_state.random_sample(n_pop) > pulse_loudness[:, 1]

        random_best_idx = np.floor(
            local_lens * self.random_state.random_sample(n_pop)
        ).astype(int)
        random_best = local_best[np.arange(n_pop), random_best_idx][mask]
        rand = (self.random_state.random_sample(size=random_best.shape) - 0.5) * 2.0
        rand = rand / 20.0
        rand[:, 0] = rand[:, 0] * (self.k_range[1] - self.k_range[0])
        rand[:, 1] = rand[:, 1] * (self.proportion_range[1] - self.proportion_range[0])
        bats[mask] = random_best + rand * np.mean(pulse_loudness[:, 2])

        bats[:, 0] = np.clip(bats[:, 0], self.k_range[0], self.k_range[1])
        bats[:, 1] = np.clip(
            bats[:, 1], self.proportion_range[0], self.proportion_range[1]
        )

        return bats

    def bat_update_local(
        self, *, bats, scores, local_best, local_scores, improved_local
    ):
        """
        Update the local scores for the bat optimization

        Args:
            bats (np.array): the bats
            scores (np.array): the actual scores
            local_best (np.array): the locally best bats
            local_scores (np.array): the local scores
            improved_local (np.array): the mask of particles that improved locally

        Returns:
            np.array, np.array: the updated local scores, and the locally best bats
        """
        local_lens = np.sum(local_scores[:, :, 0] >= 0.0, axis=1)
        local_lens[improved_local] = np.minimum(
            local_lens[improved_local] + 1, self.params["n_keep"]
        )
        local_best[improved_local, local_lens[improved_local]] = bats[improved_local]
        local_scores[improved_local, local_lens[improved_local]] = scores[
            improved_local
        ]

        sorting = local_scores[:, :, 0].argsort(axis=1)[:, ::-1]
        sorting = (
            sorting + (np.arange(self.params["n_pop"]) * self.params["n_keep"])[:, None]
        )
        n_total = self.params["n_pop"] * self.params["n_keep"]

        local_scores = local_scores.reshape(n_total, 2)[sorting]
        local_scores = local_scores.reshape(
            self.params["n_pop"], self.params["n_keep"], 2
        )

        local_best = local_best.reshape(n_total, 2)[sorting]
        local_best = local_best.reshape(self.params["n_pop"], self.params["n_keep"], 2)

        return local_scores, local_best

    def bat_update_pulse_loudness(self, pulse_loudness, improved_local, iteration):
        """
        Update the pulse and loudness values.

        Args:
            pulse_loudness (np.array): the pulse and loudness values
            improved_local (np.array): the mask of bats that improved locally
            iteration (int): the iteration number

        Returns:
            np.array: the updated pulse-loudness parameters
        """
        random = self.random_state.random_sample(self.params["n_pop"])
        mask = random[improved_local] < pulse_loudness[improved_local, 1]
        pulse_loudness[improved_local, 2][mask] = (
            pulse_loudness[improved_local, 2][mask] * self.params["alpha"]
        )
        pulse_loudness[improved_local, 1][mask] = pulse_loudness[improved_local, 1][
            mask
        ] * (1 - np.exp(-self.params["gamma"] * iteration))

        return pulse_loudness

    def bat_update_local_scores(
        self, *, bats, scores, local_best, local_scores, pulse_loudness, iteration
    ):
        """
        Update the local scores in the bat algorithm.

        Args:
            bats (np.array): all bats
            scores (np.array): the actual scores
            local_best (np.array): the locally best bats
            local_scores (np.array): the local scores
            pulse_loudness (np.array): the pulse-loudness parameters
            iteration (int): the iteration number

        Returns:
            np.array, np.array, np.array, bool: the updated locally best bats,
                                                local scores, pulse-loudness
                                                values and the flag indicating
                                                any update
        """
        improved_local = self.bat_check_local(scores, local_scores)
        local_scores, local_best = self.bat_update_local(
            bats=bats,
            scores=scores,
            local_best=local_best,
            local_scores=local_scores,
            improved_local=improved_local,
        )
        pulse_loudness = self.bat_update_pulse_loudness(
            pulse_loudness, improved_local, iteration
        )

        return local_best, local_scores, pulse_loudness, np.any(improved_local)

    def bat_update_global_scores(
        self, *, bats, scores, global_best, global_scores, pulse_loudness
    ):
        """
        Update the global scores in the bat algorithm.

        Args:
            bats (np.array): the bats
            scores (np.array): the actual scores
            global_best (np.array): the globally best bat
            global_scores (np.array): the globally best scores
            pulse_loudness (np.array): the pulse-loudness parameters

        Returns:
            np.array, np.array, bool: the updated globally best bat, globally
                                        best scores and the improvement flag
        """
        improved_global = self.bat_check_global(scores, global_scores)

        rand = self.random_state.random_sample()

        if improved_global >= 0 and rand < pulse_loudness[improved_global, 2]:
            global_best = bats[improved_global]
            global_scores = scores[improved_global]

        return global_best, global_scores, improved_global >= 0

    def bat_vectorized(self, X, y, nn_params):
        """
        BAT optimization

        Args:
            X (np.array): all training vectors
            y (np.array): all target labels
            nn_params (dict): the nearest neighbor parameters

        Returns:
            np.array: the best K and proportion values
        """

        # initial bat positions
        bats = self.init_population()
        # initial velocities
        velocities = np.zeros(shape=(self.params["n_pop"], 2))

        # best configurations of particles
        local_best = np.array(
            [[bats[i].copy()] * self.params["n_keep"] for i in range(len(bats))]
        )

        local_scores = np.zeros(shape=(self.params["n_pop"], self.params["n_keep"], 2))
        local_scores[:, :, :] = -1

        # scores of best configurations of particles
        global_best = bats[0].copy()
        global_scores = np.array([0.0, 0.0])

        # columns: pulse frequency, pulse rate, loudness
        pulse_loudness = self.random_state.random_sample(size=(self.params["n_pop"], 3))

        not_changed = 0
        for iteration in range(self.params["maxit"]):
            not_changed = not_changed + 1

            if not_changed > 10:
                break

            # update frequencies
            pulse_loudness[:, 0] = (
                self.random_state.random_sample(size=self.params["n_pop"])
                * self.params["f_max"]
            )

            velocities = velocities + (bats - global_best) * pulse_loudness[:, 0, None]

            bats = bats + velocities

            bats = self.bat_local_search(bats, local_best, local_scores, pulse_loudness)

            scores = np.array(
                [
                    self.evaluate(
                        K=int(bat[0]), proportion=bat[1], nn_params=nn_params, X=X, y=y
                    )
                    for bat in bats
                ]
            )

            (
                local_best,
                local_scores,
                pulse_loudness,
                changed,
            ) = self.bat_update_local_scores(
                bats=bats,
                scores=scores,
                local_best=local_best,
                local_scores=local_scores,
                pulse_loudness=pulse_loudness,
                iteration=iteration,
            )

            global_best, global_scores, changed = self.bat_update_global_scores(
                bats=bats,
                scores=scores,
                global_best=global_best,
                global_scores=global_scores,
                pulse_loudness=pulse_loudness,
            )

            not_changed = not_changed * changed

        return global_best

    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
        """

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

        # a reasonable range of nearest neighbors to use with SMOTE
        self.k_range = np.array(
            [1, np.min([np.sum(y == self.min_label), self.k_range[1]])]
        )

        if self.method == "pso":
            best_combination = self.pso_vectorized(X, y, nn_params)
        elif self.method == "bat":
            best_combination = self.bat_vectorized(X, y, nn_params)

        ss_params = self.ss_params.copy()
        ss_params["n_dim"] = np.min([ss_params["n_dim"], int(best_combination[0])])

        return SMOTE(
            proportion=best_combination[1],
            n_neighbors=int(best_combination[0]),
            nn_params=nn_params,
            ss_params=ss_params,
            n_jobs=self.n_jobs,
            random_state=self._random_state_init,
        ).sample(X, y)

    def get_params(self, deep=False):
        """
        Returns:
            dict: the parameters of the current sampling object
        """
        return {
            "maxit": self.params["maxit"],
            "c1": self.params["c1"],
            "c2": self.params["c2"],
            "c3": self.params["c3"],
            "alpha": self.params["alpha"],
            "gamma": self.params["gamma"],
            "method": self.method,
            "nn_params": self.nn_params,
            "ss_params": self.ss_params,
            "n_jobs": self.n_jobs,
            **OverSampling.get_params(self),
        }
