from joblib import Parallel, delayed

import numpy as np

from pyriemann.classification import MDM
from pyriemann.utils.distance import distance
from pyriemann.utils.geodesic import geodesic
from pyriemann.utils.mean import mean_covariance


class MDWM(MDM):
    def __init__(self, L=0, **kwargs):
        """Init."""
        if L < 0. or L > 1.0:
            raise ValueError("L should be chosen between 0.0 and 1.0")
        self.L = L
        super().__init__(**kwargs)

    def fit(self, X, y, X_domain, y_domain, sample_weight=None):
        """Fit (estimates) the centroids.

        Parameters
        ----------
        X : ndarray, shape (n_trials, n_channels, n_channels)
            ndarray of SPD matrices.
        y : ndarray shape (n_trials, 1)
            labels corresponding to each trial.
        sample_weight : None | ndarray shape (n_trials, 1)
            the weights of each sample from the domain. if None, each sample
            is treated with equal weights.

        Returns
        -------
        self : MDM instance
            The MDM instance.
        """
        self.classes_ = np.unique(y)
        # TODO: ajouter un test pour verifier que y et y_domain
        #       ont les meme classes

        if sample_weight is None:
            sample_weight = np.ones(X_domain.shape[0])

        if self.n_jobs == 1:
            self.target_means_ = [
                mean_covariance(X[y == label], metric=self.metric_mean)
                # sample_weight=sample_weight_target[y == l])
                for label in self.classes_
            ]

            self.domain_means_ = [
                mean_covariance(
                    X_domain[y_domain == label],
                    metric=self.metric_mean,
                    sample_weight=sample_weight[y_domain == label],
                )
                for label in self.classes_
            ]
        else:
            self.target_means_ = Parallel(n_jobs=self.n_jobs)(
                delayed(mean_covariance)(X[y == label],
                                         metric=self.metric_mean)
                for label in self.classes_
            )  # sample_weight=sample_weight_target[y == l])
            self.domain_means_ = Parallel(n_jobs=self.n_jobs)(
                delayed(mean_covariance)(
                    X_domain[y_domain == label],
                    metric=self.metric_mean,
                    sample_weight=sample_weight[y_domain == label],
                )
                for label in self.classes_
            )

        self.class_center_ = [
            geodesic(self.target_means_[i], self.domain_means_[i],
                     self.L, self.metric)
            for i, _ in enumerate(self.classes_)
        ]

        return self

    def _predict_distances(self, covtest):
        """Helper to predict the distance. equivalent to transform."""
        Nc = len(self.class_center_)

        if self.n_jobs == 1:
            dist = [
                distance(covtest, self.class_center_[m], self.metric_dist)
                for m in range(Nc)
            ]
        else:
            dist = Parallel(n_jobs=self.n_jobs)(
                delayed(distance)(covtest, self.class_center_[m],
                                  self.metric_dist)
                for m in range(Nc)
            )

        dist = np.concatenate(dist, axis=1)
        return dist

    def predict(self, covtest):
        """get the predictions.

        Parameters
        ----------
        X : ndarray, shape (n_trials, n_channels, n_channels)
            ndarray of SPD matrices.

        Returns
        -------
        pred : ndarray of int, shape (n_trials, 1)
            the prediction for each trials according to the closest centroid.
        """
        dist = self._predict_distances(covtest)
        return self.classes_[dist.argmin(axis=1)]

    def transform(self, X):
        """get the distance to each centroid.

        Parameters
        ----------
        X : ndarray, shape (n_trials, n_channels, n_channels)
            ndarray of SPD matrices.

        Returns
        -------
        dist : ndarray, shape (n_trials, n_classes)
            the distance to each centroid according to the metric.
        """
        return self._predict_distances(X)

from joblib import Parallel, delayed

import numpy as np

from pyriemann.classification import MDM
from pyriemann.utils.distance import distance
from pyriemann.utils.geodesic import geodesic
from pyriemann.utils.mean import mean_covariance


class MDWM(MDM):
    def __init__(self, L=0, **kwargs):
        """Init."""
        if L < 0. or L > 1.0:
            raise ValueError("L should be chosen between 0.0 and 1.0")
        self.L = L
        super().__init__(**kwargs)

    def fit(self, X, y, X_domain, y_domain, sample_weight=None):
        """Fit (estimates) the centroids.

        Parameters
        ----------
        X : ndarray, shape (n_trials, n_channels, n_channels)
            ndarray of SPD matrices.
        y : ndarray shape (n_trials, 1)
            labels corresponding to each trial.
        sample_weight : None | ndarray shape (n_trials, 1)
            the weights of each sample from the domain. if None, each sample
            is treated with equal weights.

        Returns
        -------
        self : MDM instance
            The MDM instance.
        """
        self.classes_ = np.unique(y)
        # TODO: ajouter un test pour verifier que y et y_domain
        #       ont les meme classes

        if sample_weight is None:
            sample_weight = np.ones(X_domain.shape[0])

        if self.n_jobs == 1:
            self.target_means_ = [
                mean_covariance(X[y == label], metric=self.metric_mean)
                # sample_weight=sample_weight_target[y == l])
                for label in self.classes_
            ]

            self.domain_means_ = [
                mean_covariance(
                    X_domain[y_domain == label],
                    metric=self.metric_mean,
                    sample_weight=sample_weight[y_domain == label],
                )
                for label in self.classes_
            ]
        else:
            self.target_means_ = Parallel(n_jobs=self.n_jobs)(
                delayed(mean_covariance)(X[y == label],
                                         metric=self.metric_mean)
                for label in self.classes_
            )  # sample_weight=sample_weight_target[y == l])
            self.domain_means_ = Parallel(n_jobs=self.n_jobs)(
                delayed(mean_covariance)(
                    X_domain[y_domain == label],
                    metric=self.metric_mean,
                    sample_weight=sample_weight[y_domain == label],
                )
                for label in self.classes_
            )

        self.class_center_ = [
            geodesic(self.target_means_[i], self.domain_means_[i],
                     self.L, self.metric)
            for i, _ in enumerate(self.classes_)
        ]

        return self

    def _predict_distances(self, covtest):
        """Helper to predict the distance. equivalent to transform."""
        Nc = len(self.class_center_)

        if self.n_jobs == 1:
            dist = [
                distance(covtest, self.class_center_[m], self.metric_dist)
                for m in range(Nc)
            ]
        else:
            dist = Parallel(n_jobs=self.n_jobs)(
                delayed(distance)(covtest, self.class_center_[m],
                                  self.metric_dist)
                for m in range(Nc)
            )

        dist = np.concatenate(dist, axis=1)
        return dist

    def predict(self, covtest):
        """get the predictions.

        Parameters
        ----------
        X : ndarray, shape (n_trials, n_channels, n_channels)
            ndarray of SPD matrices.

        Returns
        -------
        pred : ndarray of int, shape (n_trials, 1)
            the prediction for each trials according to the closest centroid.
        """
        dist = self._predict_distances(covtest)
        return self.classes_[dist.argmin(axis=1)]

    def transform(self, X):
        """get the distance to each centroid.

        Parameters
        ----------
        X : ndarray, shape (n_trials, n_channels, n_channels)
            ndarray of SPD matrices.

        Returns
        -------
        dist : ndarray, shape (n_trials, n_classes)
            the distance to each centroid according to the metric.
        """
        return self._predict_distances(X)

Init. Fit (estimates) the centroids.          Parameters         ----------         X : ndarray, shape (n_trials, n_channels, n_channels)             ndarray of SPD matrices.         y : ndarray shape (n_trials, 1)             labels corresponding to each trial.         sample_weight : None | ndarray shape (n_trials, 1)             the weights of each sample from the domain. if None, each sample             is treated with equal weights.          Returns         -------         self : MDM instance             The MDM instance. # TODO: ajouter un test pour verifier que y et y_domain #       ont les meme classes # sample_weight=sample_weight_target[y == l]) # sample_weight=sample_weight_target[y == l]) Helper to predict the distance. equivalent to transform. get the predictions.          Parameters         ----------         X : ndarray, shape (n_trials, n_channels, n_channels)             ndarray of SPD matrices.          Returns         -------         pred : ndarray of int, shape (n_trials, 1)             the prediction for each trials according to the closest centroid. get the distance to each centroid.          Parameters         ----------         X : ndarray, shape (n_trials, n_channels, n_channels)             ndarray of SPD matrices.          Returns         -------         dist : ndarray, shape (n_trials, n_classes)             the distance to each centroid according to the metric.