"""Connectivity matrices."""

import warnings
from math import floor, sqrt

import numpy as np
from scipy import linalg
from sklearn.base import TransformerMixin, clone
from sklearn.covariance import LedoitWolf
from sklearn.utils import check_array
from sklearn.utils.estimator_checks import check_is_fitted

from nilearn import signal
from nilearn._base import NilearnBaseEstimator
from nilearn._utils.docs import fill_doc
from nilearn._utils.extmath import is_spd
from nilearn._utils.logger import find_stack_level, log
from nilearn._utils.param_validation import check_parameter_in_allowed
from nilearn._utils.versions import SKLEARN_LT_1_6


def _check_square(matrix: np.ndarray) -> None:
    """Raise a ValueError if the input matrix is square.

    Parameters
    ----------
    matrix : numpy.ndarray
        Input array.

    """
    if matrix.ndim != 2 or (matrix.shape[0] != matrix.shape[-1]):
        raise ValueError(
            f"Expected a square matrix, got array of shape {matrix.shape}."
        )


def _check_spd(matrix: np.ndarray) -> None:
    """Raise a ValueError if the input matrix is not symmetric positive \
    definite.

    Parameters
    ----------
    matrix : numpy.ndarray
        Input array.

    """
    if not is_spd(matrix, decimal=7):
        raise ValueError("Expected a symmetric positive definite matrix.")


def _form_symmetric(
    function, eigenvalues: np.ndarray, eigenvectors: np.ndarray
) -> np.ndarray:
    """Return the symmetric matrix with the given eigenvectors and \
    eigenvalues transformed by function.

    Parameters
    ----------
    function : function numpy.ndarray -> numpy.ndarray
        The transform to apply to the eigenvalues.

    eigenvalues : numpy.ndarray, shape (n_features, )
        Input argument of the function.

    eigenvectors : numpy.ndarray, shape (n_features, n_features)
        Unitary matrix.

    Returns
    -------
    output : numpy.ndarray, shape (n_features, n_features)
        The symmetric matrix obtained after transforming the eigenvalues, while
        keeping the same eigenvectors.

    """
    return np.dot(eigenvectors * function(eigenvalues), eigenvectors.T)


def _map_eigenvalues(function, symmetric):
    """Matrix function, for real symmetric matrices.

    The function is applied to the eigenvalues of symmetric.

    Parameters
    ----------
    function : function numpy.ndarray -> numpy.ndarray
        The transform to apply to the eigenvalues.

    symmetric : numpy.ndarray, shape (n_features, n_features)
        The input symmetric matrix.

    Returns
    -------
    output : numpy.ndarray, shape (n_features, n_features)
        The new symmetric matrix obtained after transforming the eigenvalues,
        while keeping the same eigenvectors.

    Notes
    -----
    If input matrix is not real symmetric, no error is reported but result will
    be wrong.

    """
    eigenvalues, eigenvectors = linalg.eigh(symmetric)
    return _form_symmetric(function, eigenvalues, eigenvectors)


def _geometric_mean(matrices, init=None, max_iter=10, tol=1e-7):
    """Compute the geometric mean of symmetric positive definite matrices.

    The geometric mean of n positive definite matrices
    M_1, ..., M_n is the minimizer of the sum of squared distances from an
    arbitrary matrix to each input matrix M_k

    gmean(M_1, ..., M_n) = argmin_X sum_{k=1}^N dist(X, M_k)^2

    where the used distance is related to matrices logarithm

    dist(X, M_k) = ||log(X^{-1/2} M_k X^{-1/2)}||

    In case of positive numbers, this mean is the usual geometric mean.

    See Algorithm 3 of :footcite:t:`Fletcher2007`.

    References
    ----------
    .. footbibliography::

    Parameters
    ----------
    matrices : list of numpy.ndarray, all of shape (n_features, n_features)
        List of matrices whose geometric mean to compute. Raise an error if the
        matrices are not all symmetric positive definite of the same shape.

    init : numpy.ndarray of shape (n_features, n_features) or None, \
            default=None
        Initialization matrix, default to the arithmetic mean of matrices.
        Raise an error if the matrix is not symmetric positive definite of the
        same shape as the elements of matrices.

    %(max_iter10)s

    tol : positive float or None, default=1e-7
        The tolerance to declare convergence: if the gradient norm goes below
        this value, the gradient descent is stopped. If None, no  check is
        performed.

    Returns
    -------
    gmean : numpy.ndarray, shape (n_features, n_features)
        Geometric mean of the matrices.

    """
    # Shape and symmetry positive definiteness checks
    n_features = matrices[0].shape[0]
    for matrix in matrices:
        _check_square(matrix)
        if matrix.shape[0] != n_features:
            raise ValueError("Matrices are not of the same shape.")
        _check_spd(matrix)

    # Initialization
    matrices = np.array(matrices)
    if init is None:
        gmean = np.mean(matrices, axis=0)
    else:
        _check_square(init)
        if init.shape[0] != n_features:
            raise ValueError("Initialization has incorrect shape.")
        _check_spd(init)
        gmean = init

    norm_old = np.inf
    step = 1.0

    # Gradient descent
    for _ in range(max_iter):
        # Computation of the gradient
        vals_gmean, vecs_gmean = linalg.eigh(gmean)
        gmean_inv_sqrt = _form_symmetric(np.sqrt, 1.0 / vals_gmean, vecs_gmean)
        whitened_matrices = [
            gmean_inv_sqrt.dot(matrix).dot(gmean_inv_sqrt)
            for matrix in matrices
        ]
        logs = [_map_eigenvalues(np.log, w_mat) for w_mat in whitened_matrices]
        # Covariant derivative is - gmean.dot(logms_mean)
        logs_mean = np.mean(logs, axis=0)
        if np.any(np.isnan(logs_mean)):
            raise FloatingPointError("Nan value after logarithm operation.")

        # Norm of the covariant derivative on the tangent space at point gmean
        norm = np.linalg.norm(logs_mean)

        # Update of the minimizer
        vals_log, vecs_log = linalg.eigh(logs_mean)
        gmean_sqrt = _form_symmetric(np.sqrt, vals_gmean, vecs_gmean)
        # Move along the geodesic
        gmean = gmean_sqrt.dot(
            _form_symmetric(np.exp, vals_log * step, vecs_log)
        ).dot(gmean_sqrt)

        # Update the norm and the step size
        if norm < norm_old:
            norm_old = norm
        elif norm > norm_old:
            step = step / 2.0
            norm = norm_old
        if tol is not None and norm / gmean.size < tol:
            break
    if tol is not None and norm / gmean.size >= tol:
        warnings.warn(
            f"Maximum number of iterations {max_iter} reached without "
            f"getting to the requested tolerance level {tol}.",
            stacklevel=find_stack_level(),
        )

    return gmean


def sym_matrix_to_vec(symmetric, discard_diagonal: bool = False) -> np.ndarray:
    """Return the flattened lower triangular part of an array.

    If diagonal is kept, diagonal elements are divided by sqrt(2) to conserve
    the norm.

    Acts on the last two dimensions of the array if not 2-dimensional.

    .. nilearn_versionadded:: 0.3

    Parameters
    ----------
    symmetric : numpy.ndarray or :obj:`list` of numpy arrays, shape\
        (..., n_features, n_features)
        Input array.

    discard_diagonal : :obj:`bool`, default=False
        If True, the values of the diagonal are not returned.

    Returns
    -------
    output : numpy.ndarray
        The output flattened lower triangular part of symmetric. Shape is
        (..., n_features * (n_features + 1) / 2) if discard_diagonal is False
        and (..., (n_features - 1) * n_features / 2) otherwise.

    """
    if discard_diagonal:
        # No scaling, we directly return the values
        tril_mask = np.tril(np.ones(symmetric.shape[-2:]), k=-1).astype(bool)
        return symmetric[..., tril_mask]
    scaling = np.ones(symmetric.shape[-2:])
    np.fill_diagonal(scaling, sqrt(2.0))
    tril_mask = np.tril(np.ones(symmetric.shape[-2:])).astype(bool)
    return symmetric[..., tril_mask] / scaling[tril_mask]


def vec_to_sym_matrix(vec, diagonal=None):
    """Return the symmetric matrix given its flattened lower triangular part.

    Acts on the last dimension of the array if not 1-dimensional.
    Diagonal can be encompassed in vec or given separately. In both cases, note
    that diagonal elements are multiplied by sqrt(2).

    .. nilearn_versionadded:: 0.3

    Parameters
    ----------
    vec : numpy.ndarray or :obj:`list` of numpy arrays, shape \
        (..., n_columns * (n_columns + 1) /2) or
        (..., (n_columns - 1) * n_columns / 2) if diagonal is given separately.
        The input array.

    diagonal : numpy.ndarray, shape (..., n_columns), default=None
        The diagonal array to be stacked to vec. If None, the diagonal is
        assumed to be included in vec.

    Returns
    -------
    sym : numpy.ndarray, shape (..., n_columns, n_columns).
        The output symmetric matrix.

    Notes
    -----
    This function is meant to be the inverse of sym_matrix_to_vec. If you have
    discarded the diagonal in sym_matrix_to_vec, you need to provide it
    separately to reconstruct the symmetric matrix. For instance this can be
    useful for correlation matrices for which we know the diagonal is 1.

    See Also
    --------
    nilearn.connectome.sym_matrix_to_vec

    """
    n = vec.shape[-1]
    # Compute the number of the symmetric matrix columns
    # solve n_columns * (n_columns + 1) / 2 = n subject to n_columns > 0
    n_columns = (sqrt(8 * n + 1) - 1.0) / 2
    if diagonal is not None:
        n_columns += 1

    if n_columns > floor(n_columns):
        raise ValueError(
            f"Vector of unsuitable shape {vec.shape} cannot be transformed to "
            "a symmetric matrix."
        )

    n_columns = int(n_columns)
    first_shape = vec.shape[:-1]
    if diagonal is not None and (
        diagonal.shape[:-1] != first_shape or diagonal.shape[-1] != n_columns
    ):
        raise ValueError(
            f"diagonal of shape {diagonal.shape} incompatible "
            f"with vector of shape {vec.shape}"
        )

    sym = np.zeros((*first_shape, n_columns, n_columns))

    # Fill lower triangular part
    skip_diagonal = diagonal is not None
    mask = np.tril(np.ones((n_columns, n_columns)), k=-skip_diagonal).astype(
        bool
    )
    sym[..., mask] = vec

    # Fill upper triangular part
    sym.swapaxes(-1, -2)[..., mask] = vec

    # (Fill and) rescale diagonal terms
    mask.fill(False)
    np.fill_diagonal(mask, True)
    if diagonal is not None:
        sym[..., mask] = diagonal

    sym[..., mask] *= sqrt(2)

    return sym


def cov_to_corr(covariance: np.ndarray) -> np.ndarray:
    """Return correlation matrix for a given covariance matrix.

    Parameters
    ----------
    covariance : 2D numpy.ndarray
        The input covariance matrix.

    Returns
    -------
    correlation : 2D numpy.ndarray
        The output correlation matrix.

    """
    diagonal = np.atleast_2d(1.0 / np.sqrt(np.diag(covariance)))
    correlation = covariance * diagonal * diagonal.T

    # Force exact 1. on diagonal
    np.fill_diagonal(correlation, 1.0)
    return correlation


def prec_to_partial(precision: np.ndarray) -> np.ndarray:
    """Return partial correlation matrix for a given precision matrix.

    Parameters
    ----------
    precision : 2D numpy.ndarray
        The input precision matrix.

    Returns
    -------
    partial_correlation : 2D numpy.ndarray
        The 2D output partial correlation matrix.

    """
    partial_correlation = -cov_to_corr(precision)
    np.fill_diagonal(partial_correlation, 1.0)
    return partial_correlation


@fill_doc
class ConnectivityMeasure(TransformerMixin, NilearnBaseEstimator):
    """A class that computes different kinds of \
       :term:`functional connectivity` matrices.

    .. nilearn_versionadded:: 0.2

    Parameters
    ----------
    cov_estimator : estimator object, \
                    default=LedoitWolf(store_precision=False)
        The covariance estimator.
        This implies that correlations are slightly shrunk
        towards zero compared to a maximum-likelihood estimate

    kind : {"covariance", "correlation", "partial correlation",\
            "tangent", "precision"}, default='covariance'
        The matrix kind.
        For the use of "tangent" see :footcite:t:`Varoquaux2010b`.

    vectorize : :obj:`bool`, default=False
        If True, connectivity matrices are reshaped into 1D arrays and only
        their flattened lower triangular parts are returned.

    discard_diagonal : :obj:`bool`, default=False
        If True, vectorized connectivity coefficients do not include the
        matrices diagonal elements. Used only when vectorize is set to True.

    %(standardize_true)s

        .. note::

            Added to control passing value to `standardize` of ``signal.clean``
            to call new behavior since passing "zscore" or True (default) is
            deprecated.
            This parameter will be changed to "zscore_sample"
            in version 0.14 and removed in version 0.15.

    %(verbose0)s

    Attributes
    ----------
    cov_estimator_ : estimator object, default=None
        A new covariance estimator with the same parameters as cov_estimator.
        If ``None`` is passed,
        defaults to ``LedoitWolf(store_precision=False)``.

    mean_ : numpy.ndarray
        The mean connectivity matrix across subjects. For 'tangent' kind,
        it is the geometric mean of covariances (a group covariance
        matrix that captures information from both correlation and partial
        correlation matrices). For other values for "kind", it is the
        mean of the corresponding matrices

    n_features_in_ : :obj:`int`
        Number of features seen during fit.

    whitening_ : numpy.ndarray or None
        The inverted square-rooted geometric mean of the covariance matrices.
        Only set when for ``kind=="tangent"``

    References
    ----------
    .. footbibliography::

    """

    def __init__(
        self,
        cov_estimator=None,
        kind="covariance",
        vectorize=False,
        discard_diagonal=False,
        standardize=True,
        verbose=0,
    ):
        self.cov_estimator = cov_estimator
        self.kind = kind
        self.vectorize = vectorize
        self.discard_diagonal = discard_diagonal
        self.standardize = standardize
        self.verbose = verbose

    def _check_input(self, X, confounds=None):
        subjects_types = [type(s) for s in X]
        if set(subjects_types) != {np.ndarray}:
            raise ValueError(
                "Each subject must be 2D numpy.ndarray.\n "
                f"You provided {subjects_types}"
            )

        subjects_dims = [s.ndim for s in X]
        if set(subjects_dims) != {2}:
            raise ValueError(
                "Each subject must be 2D numpy.ndarray.\n "
                f"You provided arrays of dimensions {subjects_dims}"
            )

        features_dims = [s.shape[1] for s in X]
        if len(set(features_dims)) > 1:
            raise ValueError(
                "All subjects must have the same number of features.\n"
                f"You provided: {features_dims}"
            )

        for s in X:
            check_array(s, accept_sparse=False)

        if confounds is not None and not hasattr(confounds, "__iter__"):
            raise TypeError(
                "'confounds' input argument must be an iterable. "
                f"You provided {confounds.__class__}"
            )

    @fill_doc
    def fit(self, X, y=None):
        """Fit the covariance estimator to the given time series for each \
        subject.

        Parameters
        ----------
        X : iterable of :class:`numpy.ndarray` \
            each of shape (n_samples, n_features)
            Each :class:`numpy.ndarray` represents a subject's time series.
            The number of samples may differ from one subject to another.

        %(y_dummy)s

        Returns
        -------
        self : ConnectivityMatrix instance
            The object itself. Useful for chaining operations.

        """
        del y
        self._fit_transform(X, do_fit=True)
        return self

    def _fit_transform(
        self, X, do_transform=False, do_fit=False, confounds=None
    ):
        """Avoid duplication of computation."""
        if self.cov_estimator is None:
            self.cov_estimator = LedoitWolf(store_precision=False)

        if not hasattr(X, "__iter__"):
            raise TypeError(
                "Input must be an iterable of numpy arrays. "
                f"Got {X.__class__.__name__}"
            )

        # casting to a list
        # to make it easier to check with sklearn estimator compliance
        if isinstance(X, np.ndarray) and X.ndim == 2:
            X = [X]
        self._check_input(X, confounds=confounds)

        if do_fit:
            self.n_features_in_ = next(iter(s.shape[1] for s in X))
            self.cov_estimator_ = clone(self.cov_estimator)

        # Compute all the matrices, stored in "connectivities"
        if self.kind == "correlation":
            # TODO (nilearn 0.14: change to "zscore_sample")
            standardize = "zscore" if self.standardize is True else None
            covariances_std = [
                self.cov_estimator_.fit(
                    signal.standardize_signal(
                        x,
                        detrend=False,
                        standardize=standardize,
                    )
                ).covariance_
                for x in X
            ]
            connectivities = [cov_to_corr(cov) for cov in covariances_std]
        else:
            covariances = [self.cov_estimator_.fit(x).covariance_ for x in X]

            allowed_kinds = (
                "correlation",
                "partial correlation",
                "tangent",
                "covariance",
                "precision",
            )
            check_parameter_in_allowed(self.kind, allowed_kinds, "kind")

            if self.kind in ("covariance", "tangent"):
                connectivities = covariances
            elif self.kind == "precision":
                connectivities = [linalg.inv(cov) for cov in covariances]
            elif self.kind == "partial correlation":
                connectivities = [
                    prec_to_partial(linalg.inv(cov)) for cov in covariances
                ]

        # Store the mean
        if do_fit:
            if self.kind == "tangent":
                self.mean_ = _geometric_mean(
                    covariances, max_iter=30, tol=1e-7
                )
                self.whitening_ = _map_eigenvalues(
                    lambda x: 1.0 / np.sqrt(x), self.mean_
                )
            else:
                self.mean_ = np.mean(connectivities, axis=0)
                # Fight numerical instabilities: make symmetric
                self.mean_ = self.mean_ + self.mean_.T
                self.mean_ *= 0.5
                self.whitening_ = None

                log("Finished fit", verbose=self.verbose)

        # Compute the vector we return on transform
        if do_transform:
            # TODO (sklearn >= 1.6.0) simplify
            if SKLEARN_LT_1_6:
                for x in X:
                    check_array(
                        x,
                        estimator=self,
                        ensure_min_features=self.n_features_in_,
                    )
            else:
                from sklearn.utils.validation import validate_data

                for x in X:
                    validate_data(self, x, reset=False)

            if self.kind == "tangent":
                connectivities = [
                    _map_eigenvalues(
                        np.log, self.whitening_.dot(cov).dot(self.whitening_)
                    )
                    for cov in connectivities
                ]

            connectivities = np.array(connectivities)

            if confounds is not None and not self.vectorize:
                error_message = (
                    "'confounds' are provided but vectorize=False. "
                    "Confounds are only cleaned on vectorized matrices "
                    "as second level connectome regression "
                    "but not on symmetric matrices."
                )
                raise ValueError(error_message)

            if self.vectorize:
                connectivities = sym_matrix_to_vec(
                    connectivities, discard_diagonal=self.discard_diagonal
                )
                if confounds is not None:
                    connectivities = signal.clean(
                        connectivities, confounds=confounds
                    )

        return connectivities

    @fill_doc
    def fit_transform(self, X, y=None, confounds=None):
        """Fit the covariance estimator to the given time series \
        for each subject. \
        Then apply transform to covariance matrices for the chosen kind.

        Parameters
        ----------
        X : iterable of :class:`numpy.ndarray` \
            each of shape (n_samples, n_features)
            Each :class:`numpy.ndarray` represents a subject's time series.
            The number of samples may differ from one subject to another.

        %(y_dummy)s

        confounds : np.ndarray with shape (n_samples) or \
                    (n_samples, n_confounds), or pandas DataFrame, default=None
            Confounds to be cleaned on the vectorized matrices. Only takes
            into effect when vetorize=True.
            This parameter is passed to signal.clean. Please see the related
            documentation for details.

        Returns
        -------
        output : numpy.ndarray, shape (n_subjects, n_features, n_features) or \
            (n_subjects, n_features * (n_features + 1) / 2) if vectorize \
            is set to True.
            The transformed individual connectivities, as matrices or vectors.
            Vectors are cleaned when vectorize=True and confounds are provided.

        """
        del y
        # casting to a list
        # to make it easier to check with sklearn estimator compliance
        if not hasattr(X, "__iter__"):
            raise TypeError(
                "Input must be an iterable of numpy arrays. "
                f"Got {X.__class__.__name__}"
            )
        if isinstance(X, np.ndarray) and X.ndim == 2:
            X = [X]
        if self.kind == "tangent" and len(X) <= 1:
            # Check that people are applying fit_transform to a group of
            # subject
            # We can only impose this in fit_transform, as it is legit to
            # fit only on a single given reference point
            raise ValueError(
                "Tangent space parametrization can only "
                "be applied to a group of subjects, as it returns "
                f"deviations to the mean. You provided {X!r}"
            )
        return self._fit_transform(
            X, do_fit=True, do_transform=True, confounds=confounds
        )

    def transform(self, X, confounds=None):
        """Apply transform to covariances matrices to get the connectivity \
        matrices for the chosen kind.

        Parameters
        ----------
        X : iterable of :class:`numpy.ndarray` \
            each of shape (n_samples, n_features)
            Each :class:`numpy.ndarray` represents a subject's time series.
            The number of samples may differ from one subject to another.

        confounds : numpy.ndarray with shape (n_samples) or \
                    (n_samples, n_confounds), default=None
            Confounds to be cleaned on the vectorized matrices. Only takes
            into effect when vetorize=True.
            This parameter is passed to signal.clean. Please see the related
            documentation for details.

        Returns
        -------
        output : numpy.ndarray, shape (n_subjects, n_features, n_features) or \
            (n_subjects, n_features * (n_features + 1) / 2) if vectorize \
            is set to True.
            The transformed individual connectivities, as matrices or vectors.
            Vectors are cleaned when vectorize=True and confounds are provided.

        """
        check_is_fitted(self)
        return self._fit_transform(X, do_transform=True, confounds=confounds)

    def __sklearn_is_fitted__(self) -> bool:
        return hasattr(self, "cov_estimator_")

    def inverse_transform(self, connectivities, diagonal=None):
        """Return connectivity matrices from connectivities, \
        vectorized or not.

        If kind is 'tangent', the covariance matrices are reconstructed.

        Parameters
        ----------
        connectivities : :obj:`list` of n_subjects numpy.ndarray with shapes \
            (n_features, n_features) or (n_features * (n_features + 1) / 2,)
            or ((n_features - 1) * n_features / 2,)
            Connectivities of each subject, vectorized or not.

        diagonal : numpy.ndarray, shape (n_subjects, n_features), default=None
            The diagonals of the connectivity matrices.

        Returns
        -------
        output : numpy.ndarray, shape (n_subjects, n_features, n_features)
            The corresponding connectivity matrices. If kind is 'correlation'/
            'partial correlation', the correlation/partial correlation
            matrices are returned.
            If kind is 'tangent', the covariance matrices are reconstructed.

        """
        check_is_fitted(self)

        connectivities = np.array(connectivities)
        if self.vectorize:
            if self.discard_diagonal and diagonal is None:
                if self.kind in ["correlation", "partial correlation"]:
                    diagonal = np.ones(
                        (connectivities.shape[0], self.mean_.shape[0])
                    ) / sqrt(2.0)
                else:
                    raise ValueError(
                        "diagonal values has been discarded and are unknown "
                        f"for {self.kind} kind, "
                        "cannot reconstruct connectivity matrices."
                    )

            connectivities = vec_to_sym_matrix(
                connectivities, diagonal=diagonal
            )

        if self.kind == "tangent":
            mean_sqrt = _map_eigenvalues(np.sqrt, self.mean_)
            connectivities = [
                mean_sqrt.dot(_map_eigenvalues(np.exp, displacement)).dot(
                    mean_sqrt
                )
                for displacement in connectivities
            ]
            connectivities = np.array(connectivities)

        return connectivities

    def set_output(self, *, transform=None):
        """Set the output container when ``"transform"`` is called.

        .. warning::

            This has not been implemented yet.
        """
        raise NotImplementedError()
