import copy
from warnings import warn

import numpy as np
from scipy.linalg import eigh
from scipy.linalg.lapack import dgesvd as svd

from dipy.denoise.pca_noise_estimate import pca_noise_estimate
from dipy.testing.decorators import warning_for_keywords


def dimensionality_problem_message(arr, num_samples, spr):
    """Message about the number of samples being smaller than one less the
    dimensionality of the data to be denoised.

    Parameters
    ----------
    arr : ndarray
        Data to be denoised.
    num_samples : int
        Number of samples.
    spr : int
        Suggested patch radius.

    Returns
    -------
    str
        Message.
    """

    return (
        f"Number of samples {arr.shape[-1]} - 1 < Dimensionality "
        f"{num_samples}. This might have a performance impact. Increase p"
        f"atch_radius to {spr} to avoid this."
    )


def _pca_classifier(L, nvoxels):
    """Classifies which PCA eigenvalues are related to noise and estimates the
    noise variance

    Parameters
    ----------
    L : array (n,)
        Array containing the PCA eigenvalues in ascending order.
    nvoxels : int
        Number of voxels used to compute L

    Returns
    -------
    var : float
        Estimation of the noise variance
    ncomps : int
        Number of eigenvalues related to noise

    Notes
    -----
    This is based on the algorithm described in :footcite:p:`Veraart2016c`.

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

    # if num_samples - 1 (to correct for mean subtraction) is less than number
    # of features, discard the zero eigenvalues
    if L.size > nvoxels - 1:
        L = L[-(nvoxels - 1) :]

    # Note that the condition expressed in the while-loop is expressed in terms
    # of the variance of equation (12), not equation (11) as in
    # :footcite:p:`Veraart2016c`. Also, this code implements ascending
    # eigenvalues, unlike :footcite:p:`Veraart2016c`.
    var = np.mean(L)
    c = L.size - 1
    r = L[c] - L[0] - 4 * np.sqrt((c + 1.0) / nvoxels) * var
    while r > 0:
        var = np.mean(L[:c])
        c = c - 1
        r = L[c] - L[0] - 4 * np.sqrt((c + 1.0) / nvoxels) * var
    ncomps = c + 1

    return var, ncomps


def create_patch_radius_arr(arr, pr):
    """Create the patch radius array from the data to be denoised and the patch
    radius.

    Parameters
    ----------
    arr : ndarray
        Data to be denoised.
    pr : int or ndarray
        Patch radius.

    Returns
    -------
    patch_radius : ndarray
        Number of samples.
    """

    patch_radius = copy.deepcopy(pr)

    if isinstance(patch_radius, int):
        patch_radius = np.ones(3, dtype=int) * patch_radius
    if len(patch_radius) != 3:
        raise ValueError("patch_radius should have length 3")
    else:
        patch_radius = np.asarray(patch_radius).astype(int)
    patch_radius[arr.shape[0:3] == np.ones(3)] = 0  # account for dim of size 1

    return patch_radius


def compute_patch_size(patch_radius):
    """Compute patch size from the patch radius.

    it is twice the radius plus one.

    Parameters
    ----------
    patch_radius : ndarray
        Patch radius.

    Returns
    -------
    int
        Number of samples.
    """

    return 2 * patch_radius + 1


def compute_num_samples(patch_size):
    """Compute the number of samples as the dot product of the elements.

    Parameters
    ----------
    patch_size : ndarray
        Patch size.

    Returns
    -------
    int
        Number of samples.
    """

    return np.prod(patch_size)


def compute_suggested_patch_radius(arr, patch_size):
    """Compute the suggested patch radius.

    Parameters
    ----------
    arr : ndarray
        Data to be denoised.
    patch_size : ndarray
        Patch size.

    Returns
    -------
    int
        Suggested patch radius.
    """

    tmp = np.sum(patch_size == 1)  # count spatial dimensions with size 1
    if tmp == 0:
        root = np.ceil(arr.shape[-1] ** (1.0 / 3))  # 3D
    if tmp == 1:
        root = np.ceil(arr.shape[-1] ** (1.0 / 2))  # 2D
    if tmp == 2:
        root = arr.shape[-1]  # 1D
    root = root + 1 if (root % 2) == 0 else root  # make odd

    return int((root - 1) / 2)


@warning_for_keywords()
def genpca(
    arr,
    *,
    sigma=None,
    mask=None,
    patch_radius=2,
    pca_method="eig",
    tau_factor=None,
    return_sigma=False,
    out_dtype=None,
    suppress_warning=False,
):
    r"""General function to perform PCA-based denoising of diffusion datasets.

    Parameters
    ----------
    arr : 4D array
        Array of data to be denoised. The dimensions are (X, Y, Z, N), where N
        are the diffusion gradient directions. The first 3 dimension must have
        size >= 2 * patch_radius + 1 or size = 1.
    sigma : float or 3D array, optional
        Standard deviation of the noise estimated from the data. If no sigma
        is given, this will be estimated based on random matrix theory
        :footcite:p:`Veraart2016b`, :footcite:p:`Veraart2016c`.
    mask : 3D boolean array, optional
        A mask with voxels that are true inside the brain and false outside of
        it. The function denoises within the true part and returns zeros
        outside of those voxels.
    patch_radius : int or 1D array, optional
        The radius of the local patch to be taken around each voxel (in
        voxels). E.g. patch_radius=2 gives 5x5x5 patches.
    pca_method : 'eig' or 'svd', optional
        Use either eigenvalue decomposition (eig) or singular value
        decomposition (svd) for principal component analysis. The default
        method is 'eig' which is faster. However, occasionally 'svd' might be
        more accurate.
    tau_factor : float, optional
        Thresholding of PCA eigenvalues is done by nulling out eigenvalues that
        are smaller than:

        .. math::

                \tau = (\tau_{factor} \sigma)^2

        $\tau_{factor}$ can be set to a predefined values (e.g. $\tau_{factor} =
        2.3$ :footcite:p:`Manjon2013`), or automatically calculated using random
        matrix theory (in case that $\tau_{factor}$ is set to None).
    return_sigma : bool, optional
        If true, the Standard deviation of the noise will be returned.
    out_dtype : str or dtype, optional
        The dtype for the output array. Default: output has the same dtype as
        the input.
    suppress_warning : bool, optional
        If true, suppress warning caused by patch_size < arr.shape[-1].

    Returns
    -------
    denoised_arr : 4D array
        This is the denoised array of the same size as that of the input data,
        clipped to non-negative values.

    References
    ----------
    .. footbibliography::
    """
    if mask is None:
        # If mask is not specified, use the whole volume
        mask = np.ones_like(arr, dtype=bool)[..., 0]

    if out_dtype is None:
        out_dtype = arr.dtype

    # We retain float64 precision, iff the input is in this precision:
    if arr.dtype == np.float64:
        calc_dtype = np.float64
    # Otherwise, we'll calculate things in float32 (saving memory)
    else:
        calc_dtype = np.float32

    if not arr.ndim == 4:
        raise ValueError("PCA denoising can only be performed on 4D arrays.", arr.shape)

    if pca_method.lower() == "svd":
        is_svd = True
    elif pca_method.lower() == "eig":
        is_svd = False
    else:
        raise ValueError("pca_method should be either 'eig' or 'svd'")

    patch_radius_arr = create_patch_radius_arr(arr, patch_radius)
    patch_size = compute_patch_size(patch_radius_arr)

    ash = arr.shape[0:3]
    if np.any((ash != np.ones(3)) * (ash < patch_size)):
        raise ValueError("Array 'arr' is incorrect shape")

    num_samples = compute_num_samples(patch_size)
    if num_samples == 1:
        raise ValueError(
            "Cannot have only 1 sample,\
                          please increase patch_radius."
        )
    # account for mean subtraction by testing #samples - 1
    if (num_samples - 1) < arr.shape[-1] and not suppress_warning:
        spr = compute_suggested_patch_radius(arr, patch_size)
        warn(
            dimensionality_problem_message(arr, num_samples, spr),
            UserWarning,
            stacklevel=2,
        )

    if isinstance(sigma, np.ndarray):
        var = sigma**2
        if not sigma.shape == arr.shape[:-1]:
            e_s = "You provided a sigma array with a shape"
            e_s += f"{sigma.shape} for data with"
            e_s += f"shape {arr.shape}. Please provide a sigma array"
            e_s += " that matches the spatial dimensions of the data."
            raise ValueError(e_s)
    elif isinstance(sigma, (int, float)):
        var = sigma**2 * np.ones(arr.shape[:-1])

    dim = arr.shape[-1]
    if tau_factor is None:
        tau_factor = 1 + np.sqrt(dim / num_samples)

    theta = np.zeros(arr.shape, dtype=calc_dtype)
    thetax = np.zeros(arr.shape, dtype=calc_dtype)

    if return_sigma is True and sigma is None:
        var = np.zeros(arr.shape[:-1], dtype=calc_dtype)
        thetavar = np.zeros(arr.shape[:-1], dtype=calc_dtype)

    # loop around and find the 3D patch for each direction at each pixel
    for k in range(patch_radius_arr[2], arr.shape[2] - patch_radius_arr[2]):
        for j in range(patch_radius_arr[1], arr.shape[1] - patch_radius_arr[1]):
            for i in range(patch_radius_arr[0], arr.shape[0] - patch_radius_arr[0]):
                # Shorthand for indexing variables:
                if not mask[i, j, k]:
                    continue
                ix1 = i - patch_radius_arr[0]
                ix2 = i + patch_radius_arr[0] + 1
                jx1 = j - patch_radius_arr[1]
                jx2 = j + patch_radius_arr[1] + 1
                kx1 = k - patch_radius_arr[2]
                kx2 = k + patch_radius_arr[2] + 1

                X = arr[ix1:ix2, jx1:jx2, kx1:kx2].reshape(num_samples, dim)
                # compute the mean
                M = np.mean(X, axis=0)
                # Upcast the dtype for precision in the SVD
                X = X - M

                if is_svd:
                    # PCA using an SVD
                    svd_args = [1, 0]
                    U, S, Vt = svd(X, *svd_args)[:3]
                    # Items in S are the eigenvalues, but in ascending order
                    # We invert the order (=> descending), square and normalize
                    # \lambda_i = s_i^2 / n
                    d = S[::-1] ** 2 / X.shape[0]
                    # Rows of Vt are eigenvectors, but also in ascending
                    # eigenvalue order:
                    W = Vt[::-1].T

                else:
                    # PCA using an Eigenvalue decomposition
                    C = np.transpose(X).dot(X)
                    C = C / X.shape[0]
                    [d, W] = eigh(C)

                if sigma is None:
                    # Random matrix theory
                    this_var, _ = _pca_classifier(d, num_samples)
                else:
                    # Predefined variance
                    this_var = var[i, j, k]

                # Threshold by tau:
                tau = tau_factor**2 * this_var

                # Update ncomps according to tau_factor
                ncomps = np.sum(d < tau)
                W[:, :ncomps] = 0

                # This is equations 1 and 2 in Manjon 2013:
                Xest = X.dot(W).dot(W.T) + M
                Xest = Xest.reshape(patch_size[0], patch_size[1], patch_size[2], dim)
                # This is equation 3 in Manjon 2013:
                this_theta = 1.0 / (1.0 + dim - ncomps)
                theta[ix1:ix2, jx1:jx2, kx1:kx2] += this_theta
                thetax[ix1:ix2, jx1:jx2, kx1:kx2] += Xest * this_theta
                if return_sigma is True and sigma is None:
                    var[ix1:ix2, jx1:jx2, kx1:kx2] += this_var * this_theta
                    thetavar[ix1:ix2, jx1:jx2, kx1:kx2] += this_theta

    denoised_arr = thetax / theta
    denoised_arr.clip(min=0, out=denoised_arr)
    denoised_arr[mask == 0] = 0
    if return_sigma is True:
        if sigma is None:
            var = var / thetavar
            var[mask == 0] = 0
            return denoised_arr.astype(out_dtype), np.sqrt(var)
        else:
            return denoised_arr.astype(out_dtype), sigma
    else:
        return denoised_arr.astype(out_dtype)


@warning_for_keywords()
def localpca(
    arr,
    *,
    sigma=None,
    mask=None,
    patch_radius=2,
    gtab=None,
    patch_radius_sigma=1,
    pca_method="eig",
    tau_factor=2.3,
    return_sigma=False,
    correct_bias=True,
    out_dtype=None,
    suppress_warning=False,
):
    r"""Performs local PCA denoising.

    The method follows the algorithm in :footcite:t:`Manjon2013`.

    Parameters
    ----------
    arr : 4D array
        Array of data to be denoised. The dimensions are (X, Y, Z, N), where N
        are the diffusion gradient directions.
    sigma : float or 3D array, optional
        Standard deviation of the noise estimated from the data. If not given,
        calculate using method in :footcite:t:`Manjon2013`.
    mask : 3D boolean array, optional
        A mask with voxels that are true inside the brain and false outside of
        it. The function denoises within the true part and returns zeros
        outside of those voxels.
    patch_radius : int or 1D array, optional
        The radius of the local patch to be taken around each voxel (in
        voxels). E.g. patch_radius=2 gives 5x5x5 patches.
    gtab: gradient table object (optional if sigma is provided)
        gradient information for the data gives us the bvals and bvecs of
        diffusion data, which is needed to calculate noise level if sigma is
        not provided.
    patch_radius_sigma : int, optional
        The radius of the local patch to be taken around each voxel (in
        voxels) for estimating sigma. E.g. patch_radius_sigma=2 gives
        5x5x5 patches.
    pca_method : 'eig' or 'svd', optional
        Use either eigenvalue decomposition (eig) or singular value
        decomposition (svd) for principal component analysis. The default
        method is 'eig' which is faster. However, occasionally 'svd' might be
        more accurate.
    tau_factor : float, optional
        Thresholding of PCA eigenvalues is done by nulling out eigenvalues that
        are smaller than:

        .. math::

                \tau = (\tau_{factor} \sigma)^2

        \tau_{factor} can be change to adjust the relationship between the
        noise standard deviation and the threshold \tau. If \tau_{factor} is
        set to None, it will be automatically calculated using the
        Marcenko-Pastur distribution :footcite:p:`Veraart2016c`. Default: 2.3
        according to :footcite:t:`Manjon2013`.
    return_sigma : bool, optional
        If true, a noise standard deviation estimate based on the
        Marcenko-Pastur distribution is returned :footcite:p:`Veraart2016c`.
    correct_bias : bool, optional
        Whether to correct for bias due to Rician noise. This is an
        implementation of equation 8 in :footcite:p:`Manjon2013`.
    out_dtype : str or dtype, optional
        The dtype for the output array. Default: output has the same dtype as
        the input.
    suppress_warning : bool, optional
        If true, suppress warning caused by patch_size < arr.shape[-1].

    Returns
    -------
    denoised_arr : 4D array
        This is the denoised array of the same size as that of the input data,
        clipped to non-negative values

    References
    ----------
    .. footbibliography::
    """
    # check gtab is given, if sigma is not given
    if sigma is None and gtab is None:
        raise ValueError("gtab must be provided if sigma is not given")

    # calculate sigma
    if sigma is None:
        sigma = pca_noise_estimate(
            arr,
            gtab,
            correct_bias=correct_bias,
            patch_radius=patch_radius_sigma,
            images_as_samples=True,
        )

    return genpca(
        arr,
        sigma=sigma,
        mask=mask,
        patch_radius=patch_radius,
        pca_method=pca_method,
        tau_factor=tau_factor,
        return_sigma=return_sigma,
        out_dtype=out_dtype,
        suppress_warning=suppress_warning,
    )


@warning_for_keywords()
def mppca(
    arr,
    *,
    mask=None,
    patch_radius=2,
    pca_method="eig",
    return_sigma=False,
    out_dtype=None,
    suppress_warning=False,
):
    r"""Performs PCA-based denoising using the Marcenko-Pastur
    distribution.

    See :footcite:p:`Veraart2016c` for further details about the method.

    Parameters
    ----------
    arr : 4D array
        Array of data to be denoised. The dimensions are (X, Y, Z, N), where N
        are the diffusion gradient directions.
    mask : 3D boolean array, optional
        A mask with voxels that are true inside the brain and false outside of
        it. The function denoises within the true part and returns zeros
        outside of those voxels.
    patch_radius : int or 1D array, optional
        The radius of the local patch to be taken around each voxel (in
        voxels). E.g. patch_radius=2 gives 5x5x5 patches.
    pca_method : 'eig' or 'svd', optional
        Use either eigenvalue decomposition (eig) or singular value
        decomposition (svd) for principal component analysis. The default
        method is 'eig' which is faster. However, occasionally 'svd' might be
        more accurate.
    return_sigma : bool, optional
        If true, a noise standard deviation estimate based on the
        Marcenko-Pastur distribution is returned :footcite:p:`Veraart2016b`.
    out_dtype : str or dtype, optional
        The dtype for the output array. Default: output has the same dtype as
        the input.
    suppress_warning : bool, optional
        If true, suppress warning caused by patch_size < arr.shape[-1].

    Returns
    -------
    denoised_arr : 4D array
        This is the denoised array of the same size as that of the input data,
        clipped to non-negative values
    sigma : 3D array (when return_sigma=True)
        Estimate of the spatial varying standard deviation of the noise

    References
    ----------
    .. footbibliography::
    """
    return genpca(
        arr,
        sigma=None,
        mask=mask,
        patch_radius=patch_radius,
        pca_method=pca_method,
        tau_factor=None,
        return_sigma=return_sigma,
        out_dtype=out_dtype,
        suppress_warning=suppress_warning,
    )
