"""

Over-sampling module for MGS-GRF strategy.

"""

import warnings

import numpy as np
from imblearn.over_sampling.base import BaseOverSampler
from imblearn.utils import check_target_type
from sklearn.covariance import empirical_covariance, ledoit_wolf, oas
from sklearn.neighbors import NearestNeighbors

from .generalized_forest import DrfSk
from .knn import KNNTies

CLASSIFIERS = {
    "grf": DrfSk(random_state=0),
    "1-nn": KNNTies(n_neighbors=1),
    "5-nn": KNNTies(n_neighbors=5),
}


class MGSGRFOverSampler(BaseOverSampler):
    """
    MGS-GRF oversampling strategy.
    MGSGRFOverSampler generates the continuous featrues with MGS,
      while the catgorical ones are generated by a Classifer, a GRF by default.
    """

    def __init__(
        self,
        K=None,
        categorical_features=None,
        classifier="grf",
        kind_sampling="cholesky",
        kind_cov="EmpCov",
        mucentered=True,
        llambda=1.0,
        sampling_strategy="auto",
        random_state=None,
    ):
        """
        Initializes the MGSGRFOverSampler.

        Parameters
        ----------
        K : int
            The number of nearest neighbors considered for the parameters estimation of the gaussians.
        categorical_features : np.array like
            The indices of the categorical features
        classifier : str
            String referring to the classifer used for handling categorical features.
            Available values : 'grf','1-nn','5-nn'.
        kind_sampling : str, optional
            The implementation of MGS, by default "cholesky"
        kind_cov : str, optional
            The desired covariance methodology estimation, by default "EmpCov" which the usual emprical coavriance estimator
        mucentered : bool, optional
            If True the gaussian are centered on the empirical mean mu. Otherwise they are centered on the central point, by default True
        llambda : float, optional
            The dilatation factors of the covariances., by default 1.0
        sampling_strategy : str or float, optional
            Sampling information to resample the data set., by default "auto"
        random_state : int or str, optional
            Control the randomization of the algorithm., by default None
        """
        super().__init__(sampling_strategy=sampling_strategy)
        self.K = K
        self.llambda = llambda
        self.categorical_features = categorical_features
        self.classifier = classifier
        self.kind_cov = kind_cov
        self.kind_sampling = kind_sampling
        self.mucentered = mucentered
        self.random_state = random_state

    def _check_X_y(self, X, y):
        """
        Input validation for standard estimators from scikit learn.
        Intially chehck for consistencey length, enforces X to be 2D and y 1D.

        -> In particualr, the checking is changed to let some string for categorical
        features pass.

        Parameters
        ----------
        X : numpy array
            Input data
        y : numpy array
            Labels

        Returns
        -------
        _type_
            _description_
        """
        y, binarize_y = check_target_type(y, indicate_one_vs_all=True)
        # X = _check_X(X)
        # self._check_n_features(X, reset=True)
        # self._check_feature_names(X, reset=True)
        return X, y, binarize_y

    def _validate_estimator(self):
        """Check if the estimator instance is valid and has categorical features.

        Raises
        ------
        ValueError
            If no categorical features are provided.
        """
        super()._validate_estimator()
        if self.classifier not in CLASSIFIERS:
            raise ValueError(
                "MGSGRFOverSampler called with unsupported classsifier."
                "Available values : 'grf','1-nn','5-nn'"
            )

    def array_of_lists_to_array(self, arr):
        """Convert an array of lists to a single NumPy array.

        Parameters
        ----------
        arr : array-like
            Input array of lists.

        Returns
        -------
        np.ndarray
            The converted NumPy array.
        """
        return np.apply_along_axis(lambda a: np.array(a[0]), -1, arr[..., None])

    def _fit_resample_continuous(self, n_synthetic_sample, X_positifs):
        """Generate the synthetic continuous features.
        The categorical features are only used for the distance derivation if needed.

        Parameters
        ----------
        n_synthetic_sample : int
            The number of synthetic samples to generate.
        X_positifs : np.ndarray
            The positive continuous original samples.
        X_positifs_categorical : np.ndarray, optional
            The positive categorical original samples, by default None

        Returns
        -------
        np.ndarray
            New synthetic samples, continuous features only.
        """
        X_positifs = X_positifs.astype(float)
        n_minoritaire = X_positifs.shape[0]
        dimension_continuous = X_positifs.shape[1]  ## features continues seulement

        np.random.seed(self.random_state)

        # We fit the nn estimator only on the continuous features
        neigh = NearestNeighbors(n_neighbors=self.K, algorithm="ball_tree")
        neigh.fit(X_positifs)
        neighbor_by_index = neigh.kneighbors(
            X=X_positifs, n_neighbors=self.K + 1, return_distance=False
        )

        if self.mucentered:
            # We sample from mean of neighbors
            all_neighbors = X_positifs[neighbor_by_index.flatten()]
            mus = (1 / (self.K + 1)) * all_neighbors.reshape(
                len(X_positifs), self.K + 1, dimension_continuous
            ).sum(axis=1)
        else:
            # We sample from central point
            mus = X_positifs

        if self.kind_cov == "EmpCov" or self.kind_cov == "InvWeightCov":
            centered_X = X_positifs[neighbor_by_index.flatten()] - np.repeat(
                mus, self.K + 1, axis=0
            )
            centered_X = centered_X.reshape(len(X_positifs), self.K + 1, dimension_continuous)
            if self.kind_cov == "InvWeightCov":
                distances = (centered_X**2).sum(axis=-1)
                distances[distances > 1e-10] = distances[distances > 1e-10] ** -0.25

                # inv sqrt for positives only and half of power for multiplication below
                distances /= distances.sum(axis=-1)[:, np.newaxis]
                centered_X = (
                    np.repeat(distances[:, :, np.newaxis] ** 0.5, dimension_continuous, axis=2)
                    * centered_X
                )

            covs = (
                self.llambda * np.matmul(np.swapaxes(centered_X, 1, 2), centered_X) / (self.K + 1)
            )
            if self.kind_sampling == "svd":
                # spectral decomposition of all covariances
                eigen_values, eigen_vectors = np.linalg.eigh(covs)  ## long
                eigen_values[eigen_values > 1e-10] = eigen_values[eigen_values > 1e-10] ** 0.5
                As = [eigen_vectors[i].dot(eigen_values[i]) for i in range(len(eigen_values))]
            elif self.kind_sampling == "cholesky":
                As = np.linalg.cholesky(covs + 1e-8 * np.identity(dimension_continuous))
            else:
                raise ValueError(
                    "kind_sampling of MGS not supportedAvailable values : 'cholescky','svd' "
                )

        elif self.kind_cov in ["LWCov","OASCov","TraceCov","IdCov""ExpCov"]:
            As = []
            p = X_positifs.shape[1]
            for i in range(n_minoritaire):
                if self.kind_cov=="LWCov":
                    covariance, shrinkage = ledoit_wolf(
                    X_positifs[neighbor_by_index[i, 1:], :] - mus[neighbor_by_index[i, 0]],
                    assume_centered=True,
                    )
                    As.append(self.llambda * covariance)

                elif self.kind_cov == "OASCov":
                    covariance, shrinkage = oas(
                        X_positifs[neighbor_by_index[i, 1:], :] - mus[neighbor_by_index[i, 0]],
                        assume_centered=True,
                    )
                    As.append(self.llambda * covariance)
                elif self.kind_cov == "TraceCov":
                    covariance = empirical_covariance(
                        X_positifs[neighbor_by_index[i, 1:], :] - mus[neighbor_by_index[i, 0]],
                        assume_centered=True,
                    )
                    final_covariance = (np.trace(covariance) / p) * np.eye(p)
                    As.append(self.llambda * final_covariance)
                elif self.kind_cov == "IdCov":
                    final_covariance = (1 / p) * np.eye(p)
                    As.append(self.llambda * final_covariance)
                elif self.kind_cov == "ExpCov":
                    diffs = X_positifs[neighbor_by_index[i, 1:], :] - mus[neighbor_by_index[i, 0]]
                    exp_dist = np.exp(-np.linalg.norm(diffs, axis=1))
                    weights = exp_dist / (np.sum(exp_dist))
                    final_covariance = (diffs.T.dot(np.diag(weights)).dot(diffs)) + np.eye(
                        dimension_continuous
                    ) * 1e-10
                    As.append(self.llambda * final_covariance)
            As = np.array(As)
        else:
            raise ValueError(
                "kind_cov of MGS not supported"
                "Available values : 'EmpCov','InvWeightCov','LWCov',\
                'OASCov','TraceCov','IdCov','ExpCov' "
            )
        # sampling all new points
        indices = np.random.randint(n_minoritaire, size=n_synthetic_sample)
        u = np.random.normal(loc=0, scale=1, size=(len(indices), dimension_continuous))
        new_samples = np.zeros((n_synthetic_sample, dimension_continuous))
        for i, central_point in enumerate(indices):
            # u = np.random.normal(loc=0, scale=1, size=dimension_continuous)
            new_observation = mus[central_point, :] + As[central_point].dot(u[i])
            new_samples[i, :] = new_observation

        return new_samples

    def _fit_resample_categorical(self, new_samples, X_positifs, X_positifs_categorical):
        """ "
        Generate the categorical features based on the continuous ones already generated.
        The categorical features are generated by a classifier, a GRF by default.

        Parameters
        ----------
        new_samples : np.ndarray
            The continuous features of the new samples.
        X_positifs : np.ndarray
            The positive continuous original samples.
        X_positifs_categorical : np.ndarray
            The positive categorical original samples.

        Returns
        -------
        np.ndarray
            New synthetic samples, categorical features only.

        """
        self.clf = CLASSIFIERS[self.classifier]
        if len(self.categorical_features) == 1:  # ravel in case of one categorical freatures
            self.clf.fit(
                X_positifs, X_positifs_categorical.ravel().astype(str)
            )  # learn on continuous features in order to predict categorical features
        else:
            self.clf.fit(
                X_positifs, X_positifs_categorical.astype(str)
            )  # learn on continuous features in order to predict categorical features

        if len(self.categorical_features) == 1:  # Ravel in case of one categorical freatures
            new_samples_cat = self.clf.predict(new_samples).reshape(-1, 1)

        else:
            new_samples_cat = self.clf.predict(new_samples)

        return new_samples_cat

    def _fit_resample(self, X, y=None, to_return_classifier=False):
        """
        Resample the dataset.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Matrix containing the data which have to be sampled.
        y : array-like, shape (n_samples,)
            Corresponding label for each sample in X.
            if y=None, all points are considered positive, and oversampling on all X
        to_return_classifier : bool, default=False
            If True, return the classifier used for categorical features.

        Returns
        -------
        X_resampled : array-like, shape (n_samples, n_features)
            The resampled data.
        y_resampled : array-like, shape (n_samples,)
            The corresponding labels for the resampled data.
        classifier : object, optional
            The classifier used for categorical features, if to_return_classifier is True.
        enc : object, optional
            The encoder used for categorical features,
            if to_return_classifier is True and to_encode is True.
        """
        if self.categorical_features is not None:
            if len(self.categorical_features) == X.shape[1]:
                raise ValueError(
                    "MGSGRFOverSampler is not designed to work only with categorical "
                    "features. It requires some numerical features."
                )
        else:
            warnings.warn("MGSGRFOverSampler called with continuous features only.")

        if self.K is None:
            if self.categorical_features is None:
                self.K = int(len(X[0]))
                warnings.warn(
                    "MGSGRFOverSampler called with K=None. We set K to the number of continuous features: K="
                    + str(self.K)
                )
            else:
                self.K = int(len(X[0]) - len(self.categorical_features))
                warnings.warn(
                    "MGSGRFOverSampler called with K=None. We set K to the number of continuous features: K="
                    + str(self.K)
                )

        np.random.seed(self.random_state)

        for class_sample, n_samples in self.sampling_strategy_.items():
            if n_samples == 0:
                continue
            X_positifs = X[y == class_sample]  ## current class

            continuous = np.ones((X_positifs.shape[1]), dtype=bool)
            if self.categorical_features is not None:
                continuous[self.categorical_features] = False

            oversampled_X = np.zeros((len(X) + n_samples, X_positifs.shape[1]), dtype=object)
            oversampled_X[:len(X)] = X
            new_samples = self._fit_resample_continuous(
                n_samples, X_positifs[:, continuous]
            )  # Generate continuous features

            if self.categorical_features is not None:
                new_samples_cat = self._fit_resample_categorical(
                    new_samples, X_positifs[:, continuous], X_positifs[:, ~continuous]
                )  # Generate categorical features

            oversampled_X[len(X):, continuous] = new_samples
            del new_samples
            if self.categorical_features is not None:
                oversampled_X[len(X):, ~continuous] = new_samples_cat
                del new_samples_cat

            oversampled_X = np.array(oversampled_X)
            oversampled_y = np.hstack((y, np.full(n_samples, class_sample)))

        if to_return_classifier:
            return oversampled_X, oversampled_y, self.clf

        return oversampled_X, oversampled_y
