import sys
import math
import numpy as np

from .utils.utils import check_version
from .utils.distance import distance, pairwise_distance
from .utils.mean import mean_covariance
from .classification import MDM, _check_metric

if check_version('sklearn', '0.18'):
    from sklearn.model_selection import cross_val_score
else:
    from sklearn.cross_validation import cross_val_score


def multiset_perm_number(y):
    """return the number of unique permutation in a multiset."""
    pr = 1
    for i in np.unique(y):
        pr *= math.factorial(np.sum(y == i))
    return math.factorial(len(y)) / pr


def unique_permutations(elements):
    """Return the list of unique permutations."""
    if len(elements) == 1:
        yield (elements[0], )
    else:
        unique_elements = set(elements)
        for first_element in unique_elements:
            remaining_elements = list(elements)
            remaining_elements.remove(first_element)
            for sub_permutation in unique_permutations(remaining_elements):
                yield (first_element, ) + sub_permutation


class BasePermutation():
    """Base object for permutations test"""

    def test(self, X, y, groups=None, verbose=True):
        """Performs the permutation test

        Parameters
        ----------
        X : array-like
            The data to fit. Can be, for example a list, or an array at
            least 2d.

        y : array-like
            The target variable to try to predict in the case of
            supervised learning.

        verbose : bool, default=True
            If true, print progress.
        """
        Npe = multiset_perm_number(y)
        self.scores_ = np.zeros(np.min([self.n_perms, int(Npe)]))

        # initial fit. This is usefull for transform data or for estimating
        # parameter that does not change across permutation, like the mean of
        # all data, the pairwise distance matrix, etc.
        X = self._initial_transform(X)

        # get the non permuted score
        self.scores_[0] = self.score(X, y, groups=groups)

        if Npe <= self.n_perms:
            print("Warning, number of unique permutations : %d" % Npe)
            perms = unique_permutations(y)
            ii = 0
            for perm in perms:
                if not np.array_equal(perm, y):
                    self.scores_[ii + 1] = self.score(X, perm, groups=groups)
                    ii += 1
                    if verbose:
                        self._print_progress(ii)

        else:
            rs = np.random.RandomState(self.random_state)
            for ii in range(self.n_perms - 1):
                perm = self._shuffle(y, groups, rs)
                self.scores_[ii + 1] = self.score(X, perm, groups=groups)
                if verbose:
                    self._print_progress(ii)
        if verbose:
            print("")
        self.p_value_ = (self.scores_[0] <= self.scores_).mean()

        return self.p_value_, self.scores_

    def _print_progress(self, ii):
        """Print permutation progress"""
        sys.stdout.write("Performing permutations : [%.1f%%]\r" %
                         ((100. * (ii + 1)) / self.n_perms))
        sys.stdout.flush()

    def _initial_transform(self, X):
        """Initial transformation. By default return X."""
        return X

    def _shuffle(self, y, groups, rs):
        """Return a shuffled copy of y eventually shuffle among same groups."""
        if groups is None:
            indices = rs.permutation(len(y))
        else:
            indices = np.arange(len(groups))
            for group in np.unique(groups):
                this_mask = (groups == group)
                indices[this_mask] = rs.permutation(indices[this_mask])
        return y[indices]

    def plot(self, nbins=10, range=None, axes=None):
        """Plot results of the permutation test.

        Parameters
        ----------
        nbins : integer or array_like or 'auto', default=10
            If an integer is given, bins + 1 bin edges are returned,
            consistently with np.histogram() for numpy version >= 1.3.
            Unequally spaced bins are supported if bins is a sequence.

        range : tuple or None, default=None
            The lower and upper range of the bins. Lower and upper outliers are
            ignored. If not provided, range is (x.min(), x.max()).
            Range has no effect if bins is a sequence.
            If bins is a sequence or range is specified, autoscaling is based
            on the specified bin range instead of the range of x.

        axes : axes handle, default=None
            Axes handle for matplotlib. if None a new figure will be created.
        """
        try:
            import matplotlib.pyplot as plt
        except ImportError:
            raise ImportError("Install matplotlib to plot permutation")

        if axes is None:
            fig, axes = plt.subplots(1, 1)
        axes.hist(self.scores_[1:], nbins, range)
        x_val = self.scores_[0]
        y_max = axes.get_ylim()[1]
        axes.plot([x_val, x_val], [0, y_max], '--r', lw=2)
        x_max = axes.get_xlim()[1]
        x_min = axes.get_xlim()[0]
        x_pos = x_min + ((x_max - x_min) * 0.25)
        axes.text(x_pos, y_max * 0.8, 'p-value: %.3f' % self.p_value_)
        axes.set_xlabel('Score')
        axes.set_ylabel('Count')
        return axes


class PermutationModel(BasePermutation):
    """Permutation test using any scikit-learn model for scoring.

    Perform a permutation test using the cross-validation score of any
    scikit-learn compatible model. Score is obtained with `cross_val_score`
    from scikit-learn. The score should be a "higer is better" metric.

    Parameters
    ----------
    n_perms : int, default=100
        The number of permutation. The minimum should be 20 for a resolution of
        0.05 p-value.

    model : sklearn compatible model, default=MDM()
        The model for scoring.

    cv : int or cross-validation generator or an iterable, default=3
        Determines the cross-validation splitting strategy.
        Possible inputs for cv are:

        - None, to use the default 3-fold cross validation;
        - integer, to specify the number of folds in a `(Stratified)KFold`;
        - an object to be used as a cross-validation generator;
        - an iterable yielding train, test splits.

        For integer/None inputs, if the estimator is a classifier and `y` is
        either binary or multiclass, :class:`StratifiedKFold` is used. In all
        other cases, `KFold` is used.

    scoring : string or callable or None, default=None
        A string (see model evaluation documentation) or
        a scorer callable object / function with signature
        `scorer(estimator, X, y)`.

    n_jobs : integer, default=1
        The number of CPUs to use to do the computation. -1 means
        'all CPUs'.

    random_state : int, default=42
        random state for the permutation test.

    Attributes
    ----------
    p_value_ : float
        the p-value of the test
    scores_ : list
        contain all scores for all permutations. The fist element is the
        non-permuted score.

    See Also
    --------
    PermutationDistance
    """

    def __init__(self,
                 n_perms=100,
                 model=MDM(),
                 cv=3,
                 scoring=None,
                 n_jobs=1,
                 random_state=42):
        """Init."""
        self.n_perms = n_perms
        self.model = model
        self.cv = cv
        self.scoring = scoring
        self.n_jobs = n_jobs
        self.random_state = random_state

    def score(self, X, y, groups=None):
        """Score one permutation.

        Parameters
        ----------
        X : array-like
            The data to fit. Can be, for example a list, or an array at
            least 2d.

        y : array-like
            The target variable to try to predict in the case of
            supervised learning.
        """
        score = cross_val_score(
            self.model,
            X,
            y,
            cv=self.cv,
            n_jobs=self.n_jobs,
            scoring=self.scoring,
            groups=groups)
        return score.mean()


class PermutationDistance(BasePermutation):
    """Permutation test based on distance.

    Perform a permutation test based on distance. You have the choice of 3
    different statistic:

    - 'pairwise' :
        the statistic is based on paiwire distance as
        descibed in [1]_. This is the fastest option for low sample size since
        the pairwise distance matrix does not need to be estimated for each
        permutation.

    - 'ttest' :
        t-test based statistic obtained by the ration of the
        distance between each riemannian centroid and the group dispersion.
        The means have to be estimated for each permutation, leading to a
        slower procedure. However, this can be used for high sample size.

    - 'ftest':
        f-test based statistic estimated using the between and
        within group variability. As for the 'ttest' stats, group centroid
        are estimated for each permutation.

    Parameters
    ----------
    n_perms : int, default=100
        The number of permutation. The minimum should be 20 for a resolution of
        0.05 p-value.

    metric : string | dict, default='riemann'
        The type of metric used for centroid and distance estimation.
        see `distance` anb `mean_covariance` for the list of supported metric.
        the metric could be a dict with two keys, `mean` and `distance` in
        order to pass different metric for the centroid estimation and the
        distance estimation. Typical usecase is to pass 'logeuclid' metric for
        the mean in order to boost the computional speed and 'riemann' for the
        distance in order to keep the good sensitivity for the classification.

    mode : string, default='pairwise'
        Type of statistic to use. could be 'pairwise', 'ttest' of 'ftest'

    n_jobs : integer, default=1
        The number of CPUs to use to do the computation. -1 means
        'all CPUs'.

    random_state : int, default=42
        random state for the permutation test.

    estimator : None or sklearn compatible estimator, default=None
        If provided, data are transformed before every permutation. should
        not be used unless a supervised opperation must be applied on the data.
        This would be the case for ERP covariance.

    Attributes
    ----------
    p_value_ : float
        the p-value of the test
    scores_ : list
        contain all scores for all permutations. The fist element is the
        non-permuted score.

    See Also
    --------
    PermutationModel

    References
    ----------
    .. [1] Anderson, J. "A new method for non-parametric multivariate analysis
        of variance." Austral ecology. 2001.
    """

    def __init__(self,
                 n_perms=100,
                 metric='riemann',
                 mode='pairwise',
                 n_jobs=1,
                 random_state=42,
                 estimator=None):
        """Init."""
        self.n_perms = n_perms
        if mode not in ['pairwise', 'ttest', 'ftest']:
            raise (ValueError("mode must be 'pairwise', 'ttest' or 'ftest'"))
        self.mode = mode
        self.metric = metric
        self.n_jobs = n_jobs
        self.random_state = random_state
        self.estimator = estimator

    def score(self, X, y, groups=None):
        """Score of a permutation.

        Parameters
        ----------
        X : array-like
            The data to fit. Can be, for example a list, or an array at
            least 2d.

        y : array-like
            The target variable to try to predict in the case of
            supervised learning.
        """
        if self.estimator:
            X = self.estimator.fit_transform(X, y)
            X = self.__init_transform(X)

        if self.mode == 'ttest':
            return self._score_ttest(X, y)
        elif self.mode == 'ftest':
            return self._score_ftest(X, y)
        elif self.mode == 'pairwise':
            return self._score_pairwise(X, y)

    def _initial_transform(self, X):
        """Initial transform"""
        # if an estimator provided, then transform w
        if self.estimator:
            return X

        return self.__init_transform(X)

    def __init_transform(self, X):
        """Init tr"""
        self.mdm = MDM(metric=self.metric, n_jobs=self.n_jobs)
        self.mdm.metric_mean, self.mdm.metric_dist = _check_metric(
            self.metric
        )
        if self.mode == 'ftest':
            self.global_mean = mean_covariance(X, metric=self.mdm.metric_mean)
        elif self.mode == 'pairwise':
            X = pairwise_distance(X, metric=self.mdm.metric_dist)**2
        return X

    def _score_ftest(self, X, y):
        """Get the score"""
        mdm = self.mdm.fit(X, y)
        covmeans = np.array(mdm.covmeans_)

        # estimates between classes variability
        n_classes = len(covmeans)
        between = 0
        for ix, classe in enumerate(mdm.classes_):
            di = distance(
                covmeans[ix], self.global_mean, metric=mdm.metric_dist)**2
            between += np.sum(y == classe) * di
        between /= (n_classes - 1)

        # estimates within class variability
        within = 0
        for ix, classe in enumerate(mdm.classes_):
            within += (distance(
                X[y == classe], covmeans[ix], metric=mdm.metric_dist)
                ** 2).sum()
        within /= (len(y) - n_classes)

        score = between / within
        return score

    def _score_ttest(self, X, y):
        """Get the score"""
        mdm = self.mdm.fit(X, y)
        covmeans = np.array(mdm.covmeans_)

        # estimates distances between means
        n_classes = len(covmeans)
        pairs = pairwise_distance(covmeans, metric=mdm.metric_dist)
        mean_dist = np.triu(pairs).sum()
        mean_dist /= (n_classes * (n_classes - 1)) / 2.0

        dist = 0
        for ix, classe in enumerate(mdm.classes_):
            di = (distance(
                X[y == classe], covmeans[ix], metric=mdm.metric_dist)
                ** 2).mean()
            dist += (di / np.sum(y == classe))
        score = mean_dist / np.sqrt(dist)
        return score

    def _score_pairwise(self, X, y):
        """Score for the pairwise distance test."""
        classes = np.unique(y)
        n_classes = len(classes)
        n_samples = len(y)
        total_ss = X.sum() / (2 * n_samples)
        pattern = np.zeros((n_samples, n_samples))
        for classe in classes:
            ix = (y == classe)
            pattern += (np.outer(ix, ix) / ix.sum())

        within_ss = (X * pattern).sum() / 2

        between_ss = total_ss - within_ss

        score = ((between_ss / (n_classes - 1)) / (within_ss /
                                                   (n_samples - n_classes)))

        return score
