"""
Copyright (c) 2021, Zolisa Bleki
SPDX-License-Identifier: BSD-3-Clause
"""
from collections import defaultdict
import math
from numbers import Number
from typing import NamedTuple

from imblearn.over_sampling.base import BaseOverSampler
from imblearn.utils import Substitution
from imblearn.utils._docstring import (
    _random_state_docstring,
    _n_jobs_docstring,
)
from imblearn.utils._validation import check_neighbors_object
import numpy as np

from ._common import check_random_state


class MinorityData(NamedTuple):
    # Minority class data points
    X: np.ndarray
    # Non-overlapping partitions of the minority class data points
    clusters: list
    # Number of prowras synthetic samples to generate per cluster
    n_new_samples: list
    # Number of convex combinations of each new synthetic sample per cluster
    n_affine: list


@Substitution(
    sampling_strategy=BaseOverSampler._sampling_strategy_docstring,
    random_state=_random_state_docstring,
    n_jobs=_n_jobs_docstring
)
class ProWRAS(BaseOverSampler):
    """Proximity Weighted Random Affine Shadowsampling (ProWRAS).

    Parameters
    ----------
    {sampling_strategy}
    max_clusters : int, default=5
        The maximum number of clusters to partition the data points of each
        class included in ``sampling_strategy``. This parameter is equivalent
        to ``max_levels`` in [1]_.
    n_neighbors_max : int or estimator object, default=None
        If ``int``, the maximum number of minority class nearest neighbours of
        all the majority class data points. If object, an estimator that
        inherits from :class:`~sklearn.neighbors.base.KNeighborsMixin`
        that will be used to find the neighbors. If None, a maximum of 5
        neighbors is used.
    decay_rate : int, default=1
        The rate of decay of the weights assigned to the clusters formed by the
        data points of the minority class(es). Defaults to 1.
    n_cluster_neighbors : int or estimator object, default=5
        If ``int``, the maximum size of the minority class neighborhood per
        cluster that is used to generate the synthetic samples. If object, an
        estimator that inherits from :class:`~sklearn.neighbors.base.KNeighborsMixin`
        that will be used to find the neighbors. If larger than the size of the
        cluster itself, the whole cluster will be used as a neighbourhood.
        Defaults to 5. This parameter is equivalent to ``neb_conv`` in [1]_.
    max_affine : int, default=10
        The maximum number of shadow samples to use when generating the synthetic
        samples through random affine combinations. If ``2``, SMOTE-like
        sythetic samples will be generated that have high variance. If greater
        than 2, then LoRAS-like synthetic samples will be generated by taking
        a convex combination of them. In this case, the actual number of convex
        combinations for each cluster is decided by the normalized proximity
        weight of the respective cluster; which are then multiplied with the
        parameter to obtain the number of appropriate convex combinations of
        shadowsamples for each cluster.
    n_shadow : int, default=100
        The number of shadow samples to generate per minority class data point
        in the case where `max_affine` is greater than 2.
    std : float or sequence, default=0.0001
        The standard deviation of the Normal distribution to add to each data
        point feature when generating shadow samples. If the input is a
        sequence, its size must be equal to the number of features of ``X``
        when calling the ``fit_resample`` method. If ``float``, then same
        standard deviation will be used for all shadow samples generated.
    {random_state}
    {n_jobs}

    References
    ----------
    .. [1] Bej, S., Schultz, K., Srivastava, P., Wolfien, M., & Wolkenhauer,
       O. (2021). A multi-schematic classifier-independent oversampling
       approach for imbalanced datasets. ArXiv, abs/2107.07349.

    Examples
    --------
    >>> from pyloras import ProWRAS
    >>> from sklearn.datasets import make_classification
    >>> from collections import Counter
    >>> l = ProWRAS()
    >>> X, y = make_classification(n_classes=3, class_sep=3,
    ... weights=[0.7, 0.2, 0.1], n_informative=3, n_redundant=1, flip_y=0,
    ... n_features=20, n_clusters_per_class=1, n_samples=2000, random_state=10)
    >>> print('Original dataset shape %s' % Counter(y))
    Original dataset shape Counter({{1: 400, 2: 200, 0: 1400}})
    >>> X_res, y_res = l.fit_resample(X, y)
    >>> print(f"Resampled dataset shape % Counter(y_res))
    Resampled dataset shape Counter({{1: 1400, 2: 1400, 0: 1400}})

    """
    def __init__(
        self,
        *,
        sampling_strategy="auto",
        max_clusters=15,
        n_neighbors_max=None,
        decay_rate=1.,
        n_cluster_neighbors=5,
        max_affine=10,
        n_shadow=100,
        std=0.001,
        random_state=None,
        n_jobs=None,
    ):
        super().__init__(sampling_strategy=sampling_strategy)
        self.max_clusters = max_clusters
        self.n_neighbors_max = n_neighbors_max
        self.decay_rate = decay_rate
        self.n_cluster_neighbors = n_cluster_neighbors
        self.max_affine = max_affine
        self.n_shadow = n_shadow
        self.std = std
        self.random_state = random_state
        self.n_jobs = n_jobs

    def _make_cluster_data(self, minority_data, X_not_minority, samples_to_make):
        nn = check_neighbors_object(
            "n_neighbors_max", self.n_neighbors_max if self.n_neighbors_max else 5
        ).set_params(n_jobs=self.n_jobs)

        X = minority_data.X
        max_weight = 0
        weights = []
        for i in range(1, self.max_clusters):
            if X.shape[0] == 0:
                break
            k = min(X.shape[0], nn.n_neighbors)
            nn.set_params(n_neighbors=k).fit(X)
            cluster_indices = np.unique(
                nn.kneighbors(X_not_minority, return_distance=False)
            )
            minority_data.clusters.append(X[cluster_indices])
            weights.append(math.exp(-self.decay_rate * (i - 1)))
            max_weight = weights[-1] if weights[-1] > max_weight else max_weight
            X = X[~np.in1d(np.arange(X.shape[0]), cluster_indices)]

        # append the remaining partition (if it exists) as the last cluster
        # if the last (max_cluster - 1) clusters have been used up.
        if X.shape[0]:  # pragma: no cover
            weights.append(
                math.exp(-self.decay_rate * (self.max_clusters - 1))
            )
            max_weight = weights[-1] if weights[-1] > max_weight else max_weight
            minority_data.clusters.append(X)

        weight_sum = sum(weights)
        for i in range(len(weights)):
            # normalize weights
            weights[i] /= weight_sum
            minority_data.n_new_samples.append(
                math.ceil(samples_to_make * weights[i])
            )
            minority_data.n_affine.append(
                math.ceil(self.max_affine * weights[i] / max_weight)
            )
        return minority_data

    def _fit_resample(self, X, y):
        random_state = check_random_state(self.random_state)
        X_res = [X.copy()]
        y_res = [y.copy()]

        n_features = X.shape[1]
        prowras_samples = defaultdict(lambda: [])

        if not isinstance(self.std, Number) and len(self.std) != n_features:
            raise ValueError(
                "``std`` and number of features of `X` must be equal."
            )

        self.cdata_ = []
        for minority_class, samples_to_make in self.sampling_strategy_.items():
            if samples_to_make == 0:
                continue

            mask = (y == minority_class)
            cdata = self._make_cluster_data(
                MinorityData(X[mask], [], [], []),
                X[~mask],
                samples_to_make
            )
            self.cdata_.append(cdata)

            cnn = check_neighbors_object(
                "n_cluster_neighbors", self.n_cluster_neighbors
            )

            for i, c in enumerate(cdata.clusters):
                n_samples = cdata.n_new_samples[i]
                if cnn.n_neighbors < c.shape[0]:
                    cnn.set_params(n_jobs=self.n_jobs).fit(c)
                    neighborhoods = c[cnn.kneighbors(c, return_distance=False)]
                    rand_idx = random_state.integers(
                        0, neighborhoods.shape[0], size=n_samples
                    )
                else:
                    neighborhoods = c[None, :, :]
                    rand_idx = [0] * n_samples
                n_affine = (
                    2 if cdata.n_affine[i] < n_features else cdata.n_affine[i]
                )

                if cdata.n_affine[i] < n_features:
                    shadows = neighborhoods[rand_idx].reshape(-1, n_features)
                else:
                    groups = neighborhoods[rand_idx]
                    group_size = groups.shape[1]
                    size = (n_samples, self.n_shadow, group_size, n_features)
                    norms = random_state.normal(scale=self.std, size=size)
                    shadows = groups[:, None, :, :] + norms
                    shadows = shadows.reshape(-1, n_features)

                shadow_idx = random_state.integers(
                    0, shadows.shape[0], size=(n_samples, n_affine)
                )

                affine_weights = random_state.dirichlet(
                    [1] * n_affine, size=n_samples
                )[:, None]
                synthetic_samples = (
                    affine_weights @ shadows[shadow_idx]
                ).reshape(-1, n_features)

                prowras_samples[minority_class].append(synthetic_samples)

            samples_to_drop = sum(cdata.n_new_samples) - samples_to_make
            random_state.shuffle(prowras_samples[minority_class])
            X_res.append(
                np.vstack(prowras_samples[minority_class])[samples_to_drop:]
            )
            y_res.append([minority_class] * X_res[-1].shape[0])

        return np.concatenate(X_res), np.concatenate(y_res)
