"""Functionality to estimate CCA/PLS models and analyze outcomes."""

import numbers
import warnings

import numpy as np
import xarray as xr
import pandas as pd
from scipy.stats import pearsonr
from scipy.spatial.distance import cdist
from sklearn.utils import check_random_state

from tqdm.autonotebook import tqdm

from joblib import Parallel, delayed

from ..generative_model import setup_model, generate_data
from ..util import align_weights
from ..estimators import SVDPLS, SVDCCA
from ..estimators.helpers import _calc_cov, _calc_corr

try:
    from ..estimators import SparseCCA
except ImportError:
    pass


__all__ = [
    'analyze_dataset', 'analyze_resampled', 'analyze_subsampled',
    'analyze_model', 'analyze_model_parameters',
]


class JointCovarianceWarning(Warning):
    pass


warnings.simplefilter('always', JointCovarianceWarning)


def analyze_dataset(estr, X, Y, Xorig=None, Yorig=None,
                    x_align_ref=None, y_align_ref=None,
                    addons=tuple(), **kwargs
                    ):
    """Analyze a given dataset with a given estimator

    Fits an estimator to a given dataset, stores association strengths, weights
    and loadings and optionally runs additional analyses with fitted estimator.

    Parameters
    ----------
    estr : **sklearn**-style estimator
        for performing CCA or PLS. Must have method ``fit`` and (after fitting)
        attributes ``assocs_``, ``x_rotations_``, ``y_rotations_``,
        ``x_scores_``, ``y_scores_``
    X : np.ndarray (n_samples, n_features)
        dataset `X`
    Y : np.ndarray (n_samples, n_features)
        dataset `Y`
    Xorig : ``None`` or np.ndarray (n_samples, n_orig_features)
        if ``None`` set to ``X``. Allows to provide an alternative set of `X`
        features for calculating loadings. I.e. an implicit assumption is that
        the rows in ``X`` and ``Xorig`` correspond to the same samples
        (subjects).
    Yorig : ``None`` or np.ndarray (n_samples, n_orig_features)
        if ``None`` set to ``Y``. Allows to provide an alternative set of `Y`
        features for calculating loadings. I.e. an implicit assumption is that
        the rows in ``Y`` and ``Yorig`` correspond to the same samples
        (subjects).
    x_align_ref : (n_features,)
        after fitting, the sign of `X` weights is chosen such that the
        cosine-distance between fitted `X` weights and ``x_align_ref`` is
        positive
    y_align_ref : (n_features,)
        after fitting, the sign of `Y` weights is chosen such that the
        cosine-distance between fitted `Y` weights and ``y_align_ref`` is
        positive
    addons : list-like of add-on functions
        After fitting the estimator and saving association strengths, weights
        and loadings in ``results`` additional analyses can be performed with
        these functions. They are called in the given order, and must have the
        signature

        .. code-block:: python

            addana_fun(estr, X, Y, Xorig, Yorig, x_align_ref, y_align_ref,
                results, **kwargs)

        and are expected to save their respective outcome features ``results``.
        Various such functions are provided in module
        :mod:`sample_analysis_addons`

    kwargs : dict
        forwarded to additional analysis functions

    Returns
    -------
    results : xr.Dataset
        containing data variables for outcome features generated by analyses
    """

    if Xorig is None:
        Xorig = X
    if Yorig is None:
        Yorig = Y

    px = X.shape[1]
    py = Y.shape[1]

    # x_covs = np.cov(X.T)
    # y_covs = np.cov(Y.T)
    # assert(x_covs.shape == (px, px))

    try:
        estr.fit(X, Y)

    except ValueError as e:
        warnings.warn('Fitting error: {}'.format(e))
        between_assocs = np.nan
        between_covs_sample = np.nan
        between_corrs_sample = np.nan
        x_weights = x_loadings = np.nan * np.empty((X.shape[1], 1))
        y_weights = y_loadings = np.nan * np.empty((Y.shape[1], 1))

    else:  # fit successful

        between_assocs = estr.assocs_[0]  # estr.corrs_[0]#estr.s_[0]

        if np.isfinite(estr.x_scores_[:, 0]).all() and \
                np.isfinite(estr.y_scores_[:, 0]).all():
            between_covs_sample = _calc_cov(
                estr.x_scores_[:, 0], estr.y_scores_[:, 0])
            between_corrs_sample = _calc_corr(
                estr.x_scores_[:, 0], estr.y_scores_[:, 0])
        else:
            between_covs_sample = np.nan
            between_corrs_sample = np.nan

        x_weights = estr.x_rotations_.copy()
        y_weights = estr.y_rotations_.copy()

        # x_loadings = np.dot(Xorig_z_.T, zscore(estr.x_scores_[:, 0])) / (n)
        # y_loadings = np.dot(Yorig_z_.T,
        #   zscore(estr.y_scores_[Ymask, 0])) / (n_Ymask)
        x_loadings = _calc_loadings(Xorig, estr.x_scores_[:, [0]])
        y_loadings = _calc_loadings(Yorig, estr.y_scores_[:, [0]])

        # align
        for mode in range(x_weights.shape[1]):
            if x_align_ref is not None:
                _x_weights, signs = align_weights(
                    x_weights[:, mode], x_align_ref[:, mode], return_sign=True)
                assert signs.shape == (1, 1)
                x_weights[:, mode] = _x_weights[0]
                x_loadings[:, mode] *= signs[0, 0]
                # CAREFUL: x_weights needs to be a COPY of x_rotations_,
                # otherwise "*= signs[0, 0]" will reverse the sign
                # modification!!!
                estr.x_rotations_[:, mode] *= signs[0, 0]

        for mode in range(y_weights.shape[1]):
            if y_align_ref is not None:
                _y_weights, signs = align_weights(
                    y_weights[:, mode], y_align_ref[:, mode], return_sign=True)
                assert signs.shape == (1, 1)
                y_weights[:, mode] = _y_weights[0]
                y_loadings[:, mode] *= signs[0, 0]
                # CAREFUL: y_weights needs to be a COPY of y_rotations_,
                # otherwise "*= signs[0, 0]" will reverse the sign
                # modification!!!
                estr.y_rotations_[:, mode] *= signs[0, 0]

    x_coords = dict(x_feature=np.arange(px))
    y_coords = dict(y_feature=np.arange(py))
    results = xr.Dataset(dict(
        between_assocs=between_assocs,
        between_covs_sample=between_covs_sample,
        between_corrs_sample=between_corrs_sample,
        x_weights=xr.DataArray(x_weights,
                               dims=('x_feature', 'mode'), coords=x_coords),
        y_weights=xr.DataArray(y_weights,
                               dims=('y_feature', 'mode'), coords=y_coords),
        x_loadings=xr.DataArray(x_loadings, dims=('x_orig_feature', 'mode'),
                                coords=dict(
                                    x_orig_feature=np.arange(Xorig.shape[1]))),
        y_loadings=xr.DataArray(y_loadings, dims=('y_orig_feature', 'mode'),
                                coords=dict(
                                    y_orig_feature=np.arange(Yorig.shape[1]))),
    ))

    for ana_fun in addons:
        ana_fun(estr, X, Y, Xorig, Yorig, x_align_ref, y_align_ref, results,
                **kwargs)

    return results


def _calc_loadings(X, scores):
    """Calculate loadings, i.e. Pearson correlations between features and
    scores.

    Parameters
    ----------
    X : np.ndarray (n_samples, n_features)
        a dataset containing features in columns for which to calculate
        loadings
    scores : np.ndarray (n_samples, 1)
        scores used to calculate loadings. Samples in rows must correspond to
        samples in rows of ``X``

    Returns
    -------
    loadings : np.ndarray (n_features,)
        calculated loadings
    """
    if not np.ma.isMaskedArray(X):
        loadings = 1 - cdist(X.T, scores.T, metric='correlation')
    else:
        loadings = np.nan * np.empty((X.shape[1], scores.shape[1]))
        for i in range(X.shape[1]):
            for j in range(scores.shape[1]):
                # scores are assumed to be always valid
                is_finite = ~(X.mask[:, i])
                loadings[i, j] = pearsonr(
                    X[is_finite, i], scores[is_finite, j])[0]
    return loadings


def analyze_resampled(estr, X, Y, Xorig=None, Yorig=None,
                      x_align_ref=None, y_align_ref=None,
                      addons=tuple(),
                      n_bs=0, perm=None, loo=False,
                      random_state=None,
                      saved_perm_features='all',
                      show_progress=True,
                      n_jobs=1,
                      **kwargs):
    """Analyze a given dataset and resampled versions of it with a given
    estimator.

    In addition to the dataset itself resampled versions of it are analyzed in
    the same way. Provided resampling strategies include:

    * bootstrap
    * leave-one-out
    * permutation

    Parameters
    ----------
    estr : **sklearn**-style estimator
        for performing CCA or PLS. Must have method ``fit`` and (after fitting)
        attributes ``assocs_``, ``x_rotations_``, ``y_rotations_``,
        ``x_scores_``, ``y_scores_``
    X : np.ndarray (n_samples, n_features)
        dataset `X`
    Y : np.ndarray (n_samples, n_features)
        dataset `Y`
    Xorig : ``None`` or np.ndarray (n_samples, n_orig_features)
        if ``None`` set to ``X``. Allows to provide an alternative set of `X`
        features for calculating loadings. I.e. an implicit assumption is that
        the rows in ``X`` and ``Xorig`` correspond to the same samples
        (subjects).
    Yorig : ``None`` or np.ndarray (n_samples, n_orig_features)
        if ``None`` set to ``Y``. Allows to provide an alternative set of `Y`
        features for calculating loadings. I.e. an implicit assumption is that
        the rows in ``Y`` and ``Yorig`` correspond to the same samples
        (subjects).
    x_align_ref : (n_features,)
        after fitting, the sign of `X` weights is chosen such that the
        cosine-distance between fitted `X` weights and ``x_align_ref`` is
        positive
    y_align_ref : (n_features,)
        after fitting, the sign of `Y` weights is chosen such that the
        cosine-distance between fitted `Y` weights and ``y_align_ref`` is
        positive
    addons : list-like of add-on functions
        After fitting the estimator and saving association strengths, weights
        and loadings in ``results`` additional analyses can be performed with
        these functions. They are called in the given order, and must have the
        signature

        .. code-block:: python

            addana_fun(estr, X, Y, Xorig, Yorig, x_align_ref, y_align_ref,
                results, **kwargs)

        and are expected to save their respective outcome features ``results``.
        Various such functions are provided in module
        :mod:`sample_analysis_addons`
    n_bs : int
        number of bootstrap iterations to perform on the data
    perm : None, int, or iterable
        if ``None`` no permutations are performed, if ``int`` gives the number
        of permutations to perform on the data, if iterable each element is an
        array giving the permuted indices to use
    loo : bool
        if ``True`` leave-one-out analysis is performed on the data
    random_state : ``None``, int or random number generator instance
        used to generate random numbers
    saved_perm_features : 'all' or list-like of 'str'
        if 'all' all outcome features resulting from the permutations are
        returned, otherwise only data variables indicated by this list are
    show_progress : bool
        whether to show progress bar
    n_jobs : int or None
        number of parallel jobs (see :class:`joblib.Parallel`)
    kwargs : dict
        forwarded to additional analysis functions

    Returns
    -------
    results : xr.Dataset
        containing data variables for outcome features generated by analyses.
        Data variables obtained from bootstrapping are suffixed '_bs', from
        permutations '_perm' and from leave-one-out '_loo'
    """

    _tqdm = _prep_progressbar(show_progress)

    parallel = Parallel(n_jobs=n_jobs)

    if Xorig is None:
        Xorig = X
    if Yorig is None:
        Yorig = Y

    results = analyze_dataset(estr, X, Y, Xorig, Yorig, x_align_ref,
                              y_align_ref, addons, **kwargs)

    rng = check_random_state(random_state)

    if n_bs > 0:

        if x_align_ref is None:
            x_align_ref_ = results.x_weights.values
        else:
            x_align_ref_ = x_align_ref

        if y_align_ref is None:
            y_align_ref_ = results.y_weights.values
        else:
            y_align_ref_ = y_align_ref

        results_bs = []
        min_unique_inds = min(X.shape[1], Y.shape[1]) + 1
        for bsi in _tqdm(range(n_bs), total=n_bs, leave=False,
                         desc='bootstrap'):
            bs_inds = np.array([])
            # make sure that there are more UNIQUE bs-samples than features so
            # that cov-matrix has full rank
            while len(np.unique(bs_inds)) < min_unique_inds:
                bs_inds = rng.choice(np.arange(len(X)), size=len(X),
                                     replace=True)
            results_bs.append(
                analyze_dataset(
                    estr,
                    X[bs_inds], Y[bs_inds], Xorig[bs_inds], Yorig[bs_inds],
                    x_align_ref=x_align_ref_, y_align_ref=y_align_ref_,
                    addons=addons, **kwargs
                )
            )

        results_bs = xr.concat(results_bs, 'bs')

        for v in results_bs.data_vars:
            results[v + '_bs'] = results_bs[v]

    if loo:

        if x_align_ref is None:
            x_align_ref_ = results.x_weights.values
        else:
            x_align_ref_ = x_align_ref

        if y_align_ref is None:
            y_align_ref_ = results.y_weights.values
        else:
            y_align_ref_ = y_align_ref

        results_loo = []
        for looi in _tqdm(range(len(X)), total=len(X), leave=False,
                          desc='leave-one-out'):
            loo_inds = np.ones(len(X), dtype=bool)
            loo_inds[looi] = False
            results_loo.append(
                analyze_dataset(
                    estr, X[loo_inds], Y[loo_inds],
                    Xorig[loo_inds], Yorig[loo_inds],
                    x_align_ref=x_align_ref_, y_align_ref=y_align_ref_,
                    addons=addons, **kwargs
                )
            )

        results_loo = xr.concat(results_loo, 'loo')

        for v in results_loo.data_vars:
            results[v + '_loo'] = results_loo[v]

    if (perm is not None) and \
            ((not isinstance(perm, numbers.Integral)) or perm > 0):

        if isinstance(perm, numbers.Integral):
            n_permutations = perm
            perm_iter = _naive_permutations(Y, rng, n_permutations)
        else:
            n_permutations = len(perm)
            perm_iter = perm

        def _ana_ds_perm(perm_inds):
            return analyze_dataset(
                estr, X, Y[perm_inds], Xorig, Yorig[perm_inds],
                # for permuted data do not align
                # x_align_ref=None, y_align_ref=None,
                x_align_ref=x_align_ref,
                y_align_ref=y_align_ref,
                addons=addons, **kwargs
            )
        results_perm = parallel(
            delayed(_ana_ds_perm)(perm_inds)
            for perm_inds in perm_iter
        )

        results_perm = xr.concat(results_perm, 'perm')

        if saved_perm_features == 'all':
            saved_perm_features = results_perm.data_vars
        # else: assume saved_perm_features is a list of names
        for v in saved_perm_features:
            results[v + '_perm'] = results_perm[v]

    return results


def _naive_permutations(Y, rng, n):
    for _ in range(n):
        yield rng.permutation(len(Y))


def analyze_subsampled(estr, X, Y, Xorig=None, Yorig=None, x_align_ref=None,
                       y_align_ref=None, addons=tuple(), ns=tuple(), n_rep=10,
                       n_perm=100, n_test=0, postprocessors=tuple(), n_jobs=1,
                       show_progress=True, random_state=None, **kwargs):
    """Analyze subsampled versions of a dataset with a given estimator.

    Parameters
    ----------
    estr : **sklearn**-style estimator
        for performing CCA or PLS. Must have method ``fit`` and (after fitting)
        attributes ``assocs_``, ``x_rotations_``, ``y_rotations_``,
        ``x_scores_``, ``y_scores_``
    X : np.ndarray (n_samples, n_features)
        dataset `X`
    Y : np.ndarray (n_samples, n_features)
        dataset `Y`
    Xorig : ``None`` or np.ndarray (n_samples, n_orig_features)
        if ``None`` set to ``X``. Allows to provide an alternative set of `X`
        features for calculating loadings. I.e. an implicit assumption is that
        the rows in ``X`` and ``Xorig`` correspond to the same samples
        (subjects).
    Yorig : ``None`` or np.ndarray (n_samples, n_orig_features)
        if ``None`` set to ``Y``. Allows to provide an alternative set of `Y`
        features for calculating loadings. I.e. an implicit assumption is that
        the rows in ``Y`` and ``Yorig`` correspond to the same samples
        (subjects).
    x_align_ref : (n_features,)
        after fitting, the sign of `X` weights is chosen such that the
        cosine-distance between fitted `X` weights and ``x_align_ref`` is
        positive
    y_align_ref : (n_features,)
        after fitting, the sign of `Y` weights is chosen such that the
        cosine-distance between fitted `Y` weights and ``y_align_ref`` is
        positive
    addons : list-like of add-on functions
        After fitting the estimator and saving association strengths, weights
        and loadings in ``results`` additional analyses can be performed with
        these functions. They are called in the given order, and must have the
        signature

        .. code-block:: python

            addana_fun(estr, X, Y, Xorig, Yorig, x_align_ref, y_align_ref,
                results, **kwargs)

        and are expected to save their respective outcome features ``results``.
        Various such functions are provided in module
        :mod:`sample_analysis_addons`
    ns : list-like of int
        subsamples of these sizes are used
    n_rep : int
        number of times a subsample of a given size is drawn
    n_perm : int
        each subsample is permuted ``n_perm`` times to generate a
        null-distribution of outcome quantities
    n_test : int
        number of subjects to use as test set. ``max(ns) + n_test`` must be <=
        ``n_samples``
    postprocessors : list-like of functions
        functions are called after the final dataset has been concatenated and
        take that xr.Dataset as only argument
    n_jobs : int or None
        number of parallel jobs (see :class:`joblib.Parallel`)
    show_progress : bool
        whether to show progress bar
    random_state : ``None``, int or random number generator instance
        used to generate random numbers
    kwargs : dict
        forwarded to additional analysis functions

    Returns
    -------
    results : xr.Dataset
        containing data variables for outcome features generated by analyses
    """

    _tqdm = _prep_progressbar(show_progress)

    rng = check_random_state(random_state)
    n_subjects = len(X)

    if (len(ns) == 0) or (n_rep <= 0):
        raise ValueError('Nothing to do, need "ns" and nrep > 0')
    if not (np.asarray(ns) < n_subjects).all():
        raise ValueError('All subsample sizes must be smaller than total '
                         'number of available samples ({})'.format(n_subjects))
    if max(ns) + n_test > n_subjects:
        raise ValueError('max(ns) + n_test must be <= n_samples')

    if Xorig is None:
        Xorig = X
    if Yorig is None:
        Yorig = Y

    permuted_subjects = rng.permutation(n_subjects)
    train_subjects = permuted_subjects[n_test:]
    test_subjects = permuted_subjects[:n_test]
    Xtest, Ytest = X[test_subjects], Y[test_subjects]

    results_n = []
    for n in _tqdm(ns, total=len(ns), desc='n'):

        results_rep = []
        for repi in _tqdm(range(n_rep), total=n_rep, leave=False, desc='rep'):

            inds = rng.permutation(train_subjects)[:n]
            X_, Y_ = X[inds], Y[inds]
            Xorig_, Yorig_ = Xorig[inds], Yorig[inds]

            results_rep.append(
                analyze_resampled(
                    estr, X_, Y_, Xorig_, Yorig_,
                    x_align_ref=x_align_ref, y_align_ref=y_align_ref,
                    perm=n_perm, random_state=rng,
                    addons=addons, n_jobs=n_jobs,
                    Xtest=Xtest, Ytest=Ytest, **kwargs
                )
            )

        results_n.append(
            xr.concat(results_rep, 'rep')
        )

    results_n = xr.concat(results_n, pd.Index(ns, name='n'))

    for postproc in postprocessors:
        postproc(results_n)

    return results_n


def analyze_model(estr, Sigma, px, ns, x_align_ref=None, y_align_ref=None,
                  addons=tuple(), n_rep=10, n_perm=100, random_state=None,
                  **kwargs):
    """Synthetic datasets drawn from a model are analyzed with a given
    stimator.

    The model is specified by the covariance matrix ``Sigma``. Synthetic
    datasets are drawn from the corresponding normal distribution and analyzed.

    Parameters
    ----------
    estr : **sklearn**-style estimator
        for performing CCA or PLS. Must have method ``fit`` and (after fitting)
        attributes ``assocs_``, ``x_rotations_``, ``y_rotations_``,
        ``x_scores_``, ``y_scores_``
    Sigma : (total number of features in `X` and `Y`, total number of features
        in `X` and `Y`) model covariance matrix
    px : int
        number of `X` features (number of `Y` features is inferred from size of
        ``Sigma``)
    ns : list-like of int
        datasets of these sizes are generated from the model and analyzed
    x_align_ref : (n_features,)
        after fitting, the sign of `X` weights is chosen such that the
        cosine-distance between fitted `X` weights and ``x_align_ref`` is
        positive
    y_align_ref : (n_features,)
        after fitting, the sign of `Y` weights is chosen such that the
        cosine-distance between fitted `Y` weights and ``y_align_ref`` is
        positive
    addons : list-like of add-on functions
        After fitting the estimator and saving association strengths, weights
        and loadings in ``results`` additional analyses can be performed with
        these functions. They are called in the given order, and must have the
        signature

        .. code-block:: python

            addana_fun(estr, X, Y, Xorig, Yorig, x_align_ref, y_align_ref,
                results, **kwargs)

        and are expected to save their respective outcome features ``results``.
        Various such functions are provided in module
        :mod:`sample_analysis_addons`
    n_rep : int
        number of times a dataset of a given size is generated
    n_perm : int
        each generated dataset is permuted ``n_perm`` times to generate a
        null-distribution of outcome quantities
    random_state : ``None``, int or random number generator instance
        used to generate random numbers
    kwargs : dict
        forwarded to additional analysis functions

    Returns
    -------
    results : xr.Dataset
        containing data variables for outcome features generated by analyses
    """

    rng = check_random_state(random_state)

    if (len(ns) == 0) or (n_rep <= 0):
        raise ValueError('Nothing to do, need "ns" and nrep > 0')

    results_n = []
    for n in ns:

        results_rep = []
        for repi in range(n_rep):

            X, Y = generate_data(Sigma, px, n, rng)

            results_rep.append(
                analyze_resampled(
                    estr, X, Y,
                    x_align_ref=x_align_ref, y_align_ref=y_align_ref,
                    perm=n_perm, random_state=rng,
                    addons=addons, **kwargs
                )
            )

        results_n.append(
            xr.concat(results_rep, 'rep')
        )

    results_n = xr.concat(results_n, pd.Index(ns, name='n'))
    return results_n


def _sample_within_variance_powerlaw(n_Sigmas, rng, axPlusay_range=(-3, 0)):
    if len(axPlusay_range) != 2:
        raise ValueError('Invalid axPlusay_range')
    elif axPlusay_range[0] > axPlusay_range[1]:
        raise ValueError('axPlusay_range must be tuple with increasing values')
    for _ in range(n_Sigmas):
        axPlusay = rng.uniform(*axPlusay_range)
        ax = rng.uniform(0, 1) * axPlusay
        ay = axPlusay - ax
        yield ax, ay


def analyze_model_parameters(model, estr=None, n_rep=100, n_bs=0, n_perm=0,
                             n_per_ftrs=(2, 10, 50), pxs=(4, 8, 16, 32, 64),
                             pys='px', rs=(.1, .3, .5, .7, .9),
                             n_between_modes=1, n_Sigmas=1,
                             axPlusay_range=(-2, 0), rotate_XY=False,
                             qx=.9, qy=.9, expl_var_ratio_thr=1./2,
                             max_n_sigma_trials=10000, addons=tuple(),
                             n_test=0, mk_test_statistics=None,
                             postprocessors=tuple(), verbose=False,
                             random_state=0, show_progress=True, **kwargs):
    """Parameter-dependent models are set up and resulting synthetic datasets
    are analyzed.

    For each model, differing by the number of features, ground-truth
    correlations, within-set principal component spectra and direction of
    weight vectors relative to the principal component axes, synthetic datasets
    are generated and analyzed.

    Parameters
    ----------
    model : 'cca' or 'pls'
        whether synthetic data is generated and analyzed for CCA or PLS
    estr : ``None`` or **sklearn**-style estimator
        if ``None`` either :class:`.estimators.SVDCCA` or
        :class:`.estimators.SVDPLS` is used depending on the value of
        ``model``. Otherwise, the given estimator should correspond to
        ``model`` and must have a method ``fit`` and (after fitting) attributes
        ``assocs_``, ``x_rotations_``, ``y_rotations_``, ``x_scores_``,
        ``y_scores_``
    n_rep : int
        for each investigated model (i.e. for each joint covariance matrix)
        specified by the number of features (argument ``pxs``), ground-truth
        correlations (argument ``rs``) and the principal component spectra
        (argument ``n_Sigmas``) ``n_rep`` datasets are drawn from this
        particular model and analyzed
    n_bs : int
        number of bootstrap iterations to perform on each synthetic dataset
    n_perm : int
        number of permutations to perform on each synthetic dataset
    n_per_ftrs : 'auto' or list-like of int
        multiplied by ``px+py`` specifies the size of samples generated from
        the model. If 'auto' values
        are chosen heuristically.
    pxs : list-like of int
        number of `X`-features to use
    pys : 'px', function or int
        if 'px' uses ``px`` `Y` features, if function uses ``function(px)``
        `Y` features, if int uses ``int(pys)`` `Y` features
    rs : list-like of float between 0 and 1
        assumed ground-truth correlations
    n_between_modes : int
        number of between-set association modes
    n_Sigmas : int
        number of covariance matrices generated for each `px` and `r`. Given
        `px` and `r` covariance matrices differ by their within-set principal
        component spectra (specified by parameter ``axPlusay_range`` and the
        directions of the between-set mode (i.e. CCA  / PLS weight) vectors
        relative to the principal component axes
    axPlusay_range : tuple of floats <= 0
        separately for `X` and `Y` the within-set principal component spectrum
        is assumed to follow a power-law with exponent drawn from a uniform
        distribution with bounds given by ``axPlusay_range``
    rotate_XY : bool
        if ``True`` a random rotation is applied to each generated dataset, the
        same rotation is applied to datasets drawn from the same model (i.e.
        with the same joint covariance matrix)
    qx : int or float between 0 and 1
        specifies the number of dominant basis vectors from which to choose the
        dominant component of the latent mode vectors for `X`. See
        :func:`.generative_model._mk_Sigmaxy` for details
    qy : int or float between 0 and 1
        specifies the number of dominant basis vectors from which to choose the
        dominant component of the latent mode vectors for `Y`. See
        :func:`.generative_model._mk_Sigmaxy` for details
    expl_var_ratio_thr : float
        threshold for required within-modality variance along latent mode
        vectors
    max_n_sigma_trials : int
        number of times an attempt is made to find suitable latent mode
        vectors. See `_mk_Sigmaxy` for details.
    addons : list-like of functions
        After fitting the estimator and saving association strengths, weights
        and loadings in ``results`` additional analyses can be performed with
        these functions. They are called in the given order, and must have the
        signature

        .. code-block:: python

            addana_fun(estr, X, Y, Xorig, Yorig, x_align_ref, y_align_ref,
                results, **kwargs)

        and are expected to save their respective outcome features ``results``.
        Various such functions are provided in module
        :mod:`sample_analysis_addons`
    n_test : 'auto' or int >= 0
        to analyze some consistency properties across repeated draws from the
        same model, a test set of size ``n_test`` is generated for each joint
        covariance matrix and provided to down-stream analyses via
        keyword-arguments ``Xtest`` and ``Ytest``. If set to ``'auto'`` a test
        set of size max(n_per_ftrs) * (px + py) is used.
    mk_test_statistics : ``None`` or function
        if not ``None`` the function must have the signature

        .. code-block:: python

            fun(Xtest, Ytest, x_weights_true, y_weights_true)

        where ``Xtest`` and ``Ytest`` are, respectively, ``np.ndarray`` of
        dimension ``(n_test, n_x_features)`` and ``(n_test, n_y_features)``,
        and ``x_weights_true`` and ``y_weights_true`` are ``np.ndarray`` of
        dimension ``(n_x_features, n_components)`` and
        ``(n_y_features, n_components)`` containing the true weight vectors
    postprocessors : list-like of functions
        functions are called after the final dataset has been concatenated and
        take that xr.Dataset as only argument
    verbose : bool
        whether some status messages are printed
    random_state : ``None``, int or random number generator instance
        used to generate random numbers    additional_analyses
    show_progress : bool
        if ``True`` progress bars are shown, if ``False`` not
    kwargs : dict
        forwarded to additional analysis functions

    Returns
    -------
    results : xr.Dataset
        containing data variables for outcome features generated by analyses
    """

    _tqdm = _prep_progressbar(show_progress)

    estr = _check_model_and_estr(model, estr)

    rng = check_random_state(random_state)

    px_results = []
    for pxi, px in _tqdm(enumerate(pxs), total=len(pxs), desc='px'):

        # px = px
        py = _get_py(pys, px)

        r_results = []
        for ri, r_between in _tqdm(enumerate(rs), total=len(rs), leave=False,
                                   desc='r'):

            sigma_results = []
            for sigmai, (ax, ay) in _tqdm(enumerate(
                    _sample_within_variance_powerlaw(n_Sigmas, rng,
                                                     axPlusay_range)
            ), total=n_Sigmas, leave=False, desc='Sigma'):
                try:
                    Sigma, Sigmaxy_svals, true_corrs, U, V, \
                        latent_expl_var_ratios_x, latent_expl_var_ratios_y, \
                        U_latent, V_latent, \
                        cosine_sim_pc1_latentMode_x, \
                        cosine_sim_pc1_latentMode_y, \
                        latent_mode_vector_algo \
                        = setup_model(model, rng, m=n_between_modes,
                                      px=px, py=py, qx=qx, qy=qy,
                                      ax=ax, ay=ay, r_between=r_between,
                                      expl_var_ratio_thr=expl_var_ratio_thr,
                                      max_n_sigma_trials=max_n_sigma_trials,
                                      verbose=verbose, return_full=True)
                except ValueError as e:
                    warnings.warn("Couldn't find Sigma: {}".format(e),
                                  category=JointCovarianceWarning)
                    continue

                if rotate_XY:
                    warnings.warn('rotate_XY functionality not tested')
                    Qx = np.linalg.qr(rng.normal(size=(px, px)))[0]
                    Qy = np.linalg.qr(rng.normal(size=(py, py)))[0]
                    assert np.allclose(np.dot(Qx.T, Qx), np.eye(len(Qx)))
                    assert np.allclose(np.dot(Qy.T, Qy), np.eye(len(Qy)))
                    Qxy = np.vstack([
                        np.hstack([Qx, np.zeros((Qx.shape[0], Qy.shape[0]))]),
                        np.hstack([np.zeros((Qy.shape[0], Qx.shape[0])), Qy])
                    ])
                    Sigma = (Qxy.T).dot(Sigma).dot(Qxy)
                    U_latent = np.dot(Qx.T, U_latent)  # Qx.T = Qx^{-1}
                    V_latent = np.dot(Qy.T, V_latent)  # Qy.T = Qy^{-1}

                true_loadings = _calc_true_loadings(Sigma, px,
                                                    U_latent, V_latent)

                n_per_ftrs = _select_n_per_ftrs(n_per_ftrs, model, ax, ay,
                                                r_between)
                if (n_test == 'auto') or (n_test > 0):
                    n_max = np.max(n_per_ftrs) * (px + py)
                    if n_test == 'auto':
                        n_test = n_max
                    elif n_max > n_test:
                        warnings.warn(
                            "It's recommended to have n_test > max n, but got "
                            "n_test ({}) < max n ({}).".format(n_test, n_max),
                            category=UserWarning)

                Xtest, Ytest = generate_data(Sigma, px, n_test, rng)

                if mk_test_statistics is None:
                    test_statistics = dict()
                else:
                    test_statistics = mk_test_statistics(Xtest, Ytest,
                                                         U_latent, V_latent)

                n_results = []
                for ni, n_per_ftr in _tqdm(
                        enumerate(n_per_ftrs),
                        total=len(n_per_ftrs), leave=False, desc='n_per_ftr'
                ):

                    # px + py = total number of variables in X and Y
                    n = int((px + py) * n_per_ftr)

                    rep_results = []
                    for repi in _tqdm(range(n_rep), total=n_rep, leave=False,
                                      desc='repetition'):

                        X, Y = generate_data(Sigma, px, n, rng)

                        rep_results.append(
                            analyze_resampled(estr, X, Y, X, Y, n_bs=n_bs,
                                              perm=n_perm, random_state=rng,
                                              addons=addons,
                                              x_align_ref=U_latent,
                                              y_align_ref=V_latent,
                                              true_loadings=true_loadings,
                                              Xtest=Xtest, Ytest=Ytest,
                                              test_statistics=test_statistics,
                                              show_progress=show_progress,
                                              **kwargs
                                              )
                        )

                    n_results.append(
                        xr.concat(rep_results, 'rep')
                    )

                coords_x_ftr = dict(x_feature=np.arange(len(U_latent)))
                coords_y_ftr = dict(y_feature=np.arange(len(V_latent)))

                _sigma_results = xr.concat(
                    n_results, pd.Index(n_per_ftrs, name='n_per_ftr'))
                _sigma_results['between_assocs_true'] = \
                    xr.DataArray(Sigmaxy_svals, dims=('mode',))
                _sigma_results['x_weights_true'] = \
                    xr.DataArray(U_latent, dims=('x_feature', 'mode'),
                                 coords=coords_x_ftr)
                _sigma_results['y_weights_true'] = \
                    xr.DataArray(V_latent, dims=('y_feature', 'mode'),
                                 coords=coords_y_ftr)
                _sigma_results['ax'] = ax
                _sigma_results['ay'] = ay
                _sigma_results['latent_expl_var_ratios_x'] = \
                    xr.DataArray(latent_expl_var_ratios_x, dims=('mode',))
                _sigma_results['latent_expl_var_ratios_y'] = \
                    xr.DataArray(latent_expl_var_ratios_y, dims=('mode',))
                _sigma_results['weight_selection_algorithm'] = \
                    latent_mode_vector_algo

                _sigma_results['x_loadings_true'] = \
                    xr.DataArray(true_loadings['x_loadings_true'],
                                 dims=('x_feature', 'mode'),
                                 coords=coords_x_ftr)
                _sigma_results['x_crossloadings_true'] = \
                    xr.DataArray(true_loadings['x_crossloadings_true'],
                                 dims=('x_feature', 'mode'),
                                 coords=coords_x_ftr)
                _sigma_results['y_loadings_true'] = \
                    xr.DataArray(true_loadings['y_loadings_true'],
                                 dims=('y_feature', 'mode'),
                                 coords=coords_y_ftr)
                _sigma_results['y_crossloadings_true'] = \
                    xr.DataArray(true_loadings['y_crossloadings_true'],
                                 dims=('y_feature', 'mode'),
                                 coords=coords_y_ftr)

                for k, v in test_statistics.items():
                    _sigma_results[k] = v

                sigma_results.append(
                    _sigma_results
                )

            r_results.append(xr.concat(
                sigma_results,
                pd.Index(np.arange(len(sigma_results)), name='Sigma_id')
            ))

        _r_results = xr.concat(r_results, pd.Index(rs, name='r'))
        _r_results['py'] = py
        px_results.append(
            _r_results
        )

    result = xr.concat(px_results, pd.Index(pxs, name='px'))

    for postproc in postprocessors:
        postproc(result)

    result = result.sortby(['n_per_ftr', 'px', 'r'])

    return result


def _prep_progressbar(show_progress):
    if show_progress:
        _tqdm = tqdm
    else:
        def _tqdm(x, *args, **kwargs):
            return x
    return _tqdm


def _check_model_and_estr(model, estr):

    if model not in ['pls', 'cca']:
        raise ValueError(
            'Invalid model: {}. Must be "pls" or "cca"'.format(model))

    if (estr is None) or (estr == 'auto'):
        if (model == 'pls'):
            estr = 'PLS'  # see below
        elif (model == 'cca'):
            estr = 'CCA'  # see below

    if isinstance(estr, str):
        if estr.lower() == 'cca':
            estr = SVDCCA(n_components=1, scale=False, std_ddof=1,
                          normalize_weights=True, cov_out_of_bounds='raise')
        elif estr.lower() == 'pls':
            estr = SVDPLS(n_components=1, scale=False, std_ddof=1)
        elif estr.lower() == 'sparsecca':
            estr = SparseCCA(n_components=1, scale=False, std_ddof=1)
        else:
            raise ValueError("Unknown estimator type: '{}'. Must be 'auto', "
                             "'CCA', 'PLS' or 'SparseCCA'".format(estr))
    else:
        # heuristic check if estr class is compatible with model

        if not (hasattr(estr, 'n_components') and hasattr(estr, 'fit')):
            raise ValueError("Estimator ({}) doesn't have required attributes "
                             "'n_components' or 'fit'".format(estr))

        estr_name = estr.__class__.__name__.lower()
        if 'pls' in estr_name:
            if model != 'pls':
                warnings.warn("requested model is 'pls', but estimator doesn't"
                              " seem to be a PLS estimator",
                              category=UserWarning)
        elif 'cca' in estr_name:
            if model != 'cca':
                warnings.warn("requested model is 'cca', but estimator doesn't"
                              " seem to be a CCA estimator",
                              category=UserWarning)
        else:
            warnings.warn("requested model doesn't seem to be a PLS or CCA"
                          " estimator", category=UserWarning)

    return estr


def _get_py(pys, px):
    """Algorithm to determine number of y-features to use in parameter sweep.

    Parameters
    ----------
    pys : 'px', function or int
        if 'px' return ``px``, if function return ``function(px)``, if int
        return ``int(pys)``
    px : int
        number of `X` features

    Returns
    -------
    py : int
        number of y-features to use
    """
    if pys == 'px':
        py = px
    else:
        try:  # allow `pys` to be a function
            py = pys(px)
        except TypeError:
            # assume pys is an integer
            try:
                py = int(pys)
            except ValueError:
                raise ValueError('Invalid py: {}'.format(pys))
    return py


def _select_n_per_ftrs(n_per_ftrs, model, ax, ay, r_between):

    if n_per_ftrs != 'auto':
        return n_per_ftrs

    else:

        for _r, _n_per_ftrs in [
            # tuples (r, n_per_ftrs), values were chosen heuristically
            (.1, (3, 4, 8, 16, 32, 64, 128, 256, 512, 1024)),
            (.3, (3, 4, 8, 16, 32, 64, 128,)),
            (.5, (3, 4, 8, 16, 32, 64)),
            (.7, (3, 4, 8, 16, 32)),
            (.8, (3, 4, 8, 16)),
            (1., (3, 4, 8)),
        ]:
            if r_between <= _r:
                n_per_ftrs = _n_per_ftrs
                break

        if model == 'pls':
            npf_mx = n_per_ftrs[-1]
            n_per_ftrs = np.r_[n_per_ftrs, [2*npf_mx, 4*npf_mx, 8*npf_mx]]

        # return in reverse order to get upper limit for time estimate from
        # tqdm
        return n_per_ftrs[::-1]


def _calc_true_loadings(Sigma, px, wX_true, wY_true):

    if (wX_true.shape[1] > 1) or (wY_true.shape[1] > 1):
        warnings.warn('More than 1 mode not tested')

    invsqrt_std_X = np.diag(1. / np.sqrt(np.diag(Sigma[:px, :px])))
    invsqrt_std_Y = np.diag(1. / np.sqrt(np.diag(Sigma[px:, px:])))

    SigmaXX = Sigma[:px, :px]
    SigmaXY = Sigma[:px, px:]
    SigmaYX = Sigma[px:, :px]
    SigmaYY = Sigma[px:, px:]

    x_scores_std = np.array([
        np.sqrt((wX_true[:, i].T).dot(SigmaXX).dot(wX_true[:, i]))
        for i in range(wX_true.shape[1])]).reshape(1, -1)
    y_scores_std = np.array([
        np.sqrt((wY_true[:, i].T).dot(SigmaYY).dot(wY_true[:, i]))
        for i in range(wY_true.shape[1])]).reshape(1, -1)

    x_loadings_true = invsqrt_std_X.dot(SigmaXX).dot(wX_true / x_scores_std)
    x_crossloadings_true = \
        invsqrt_std_X.dot(SigmaXY).dot(wY_true / y_scores_std)
    y_loadings_true = invsqrt_std_Y.dot(SigmaYY).dot(wY_true / y_scores_std)
    y_crossloadings_true = \
        invsqrt_std_Y.dot(SigmaYX).dot(wX_true / x_scores_std)

    return dict(
        x_loadings_true=x_loadings_true,
        x_crossloadings_true=x_crossloadings_true,
        y_loadings_true=y_loadings_true,
        y_crossloadings_true=y_crossloadings_true,
    )
