"""Robust and Unbiased Model-BAsed Spherical Deconvolution (RUMBA-SD)"""

import logging
import warnings

import numpy as np

from dipy.core.geometry import vec2vec_rotmat
from dipy.core.gradients import get_bval_indices, gradient_table, unique_bvals_tolerance
from dipy.core.onetime import auto_attr
from dipy.core.sphere import Sphere
from dipy.data import get_sphere
from dipy.reconst.csdeconv import AxSymShResponse
from dipy.reconst.odf import OdfFit, OdfModel
from dipy.reconst.shm import lazy_index, normalize_data
from dipy.segment.mask import bounding_box, crop
from dipy.sims.voxel import all_tensor_evecs, single_tensor
from dipy.testing.decorators import warning_for_keywords

# Machine precision for numerical stability in division
_EPS = np.finfo(float).eps
logger = logging.getLogger(__name__)


class RumbaSDModel(OdfModel):
    @warning_for_keywords()
    def __init__(
        self,
        gtab,
        *,
        wm_response=(1.7e-3, 0.2e-3, 0.2e-3),
        gm_response=0.8e-3,
        csf_response=3.0e-3,
        n_iter=600,
        recon_type="smf",
        n_coils=1,
        R=1,
        voxelwise=True,
        use_tv=False,
        sphere=None,
        verbose=False,
    ):
        """
        Robust and Unbiased Model-BAsed Spherical Deconvolution (RUMBA-SD).

        RUMBA-SD :footcite:p:`CanalesRodriguez2015` is a modification of the
        Richardson-Lucy algorithm accounting for Rician and Noncentral Chi noise
        distributions, which more accurately represent MRI noise. Computes a
        maximum likelihood estimation of the fiber orientation density function
        (fODF) at each voxel. Includes white matter compartments alongside
        optional GM and CSF compartments to account for partial volume effects.
        This fit can be performed voxelwise or globally. The global fit will
        proceed more quickly than the voxelwise fit provided that the computer
        has adequate RAM (>= 16 GB should be sufficient for most datasets).

        Kernel for deconvolution constructed using a priori knowledge of white
        matter response function, as well as the mean diffusivity of GM and/or
        CSF. RUMBA-SD is robust against impulse response imprecision, and thus
        the default diffusivity values are often adequate
        :footcite:p:`DellAcqua2007`.


        Parameters
        ----------
        gtab : GradientTable
            Gradient table.
        wm_response : 1d ndarray or 2d ndarray or AxSymShResponse, optional
            Tensor eigenvalues as a (3,) ndarray, multishell eigenvalues as
            a (len(unique_bvals_tolerance(gtab.bvals))-1, 3) ndarray in
            order of smallest to largest b-value, or an AxSymShResponse.
            Default: np.array([1.7e-3, 0.2e-3, 0.2e-3])
        gm_response : float, optional
            Mean diffusivity for GM compartment. If `None`, then grey
            matter volume fraction is not computed. Default: 0.8e-3
        csf_response : float, optional
            Mean diffusivity for CSF compartment. If `None`, then CSF
            volume fraction is not computed. Default: 3.0e-3
        n_iter : int, optional
            Number of iterations for fODF estimation. Must be a positive int.
            Default: 600
        recon_type : {'smf', 'sos'}, optional
            MRI reconstruction method: spatial matched filter (SMF) or
            sum-of-squares (SoS). SMF reconstruction generates Rician noise
            while SoS reconstruction generates Noncentral Chi noise.
            Default: 'smf'
        n_coils : int, optional
            Number of coils in MRI scanner -- only relevant in SoS
            reconstruction. Must be a positive int. Default: 1
        R : int, optional
            Acceleration factor of the acquisition. For SIEMENS,
            R = iPAT factor. For GE, R = ASSET factor. For PHILIPS,
            R = SENSE factor. Typical values are 1 or 2. Must be a positive
            int. Default: 1
        voxelwise : bool, optional
            If true, performs a voxelwise fit. If false, performs a global fit
            on the entire brain at once. The global fit requires a 4D brain
            volume in `fit`.
        use_tv : bool, optional
            If true, applies total variation regularization. This only takes
            effect in a global fit (`voxelwise` is set to `False`). TV can only
            be applied to 4D brain volumes with no singleton dimensions.
        sphere : Sphere, optional
            Sphere on which to construct fODF. If None, uses `repulsion724`.
        verbose : bool, optional
            If true, logs updates on estimated signal-to-noise ratio after each
            iteration. This only takes effect in a global fit (`voxelwise` is
            set to `False`).

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

        """

        if not np.any(gtab.b0s_mask):
            raise ValueError("Gradient table has no b0 measurements")

        self.gtab_orig = gtab  # save for prediction

        # Masks to extract b0/non-b0 measurements
        self.where_b0s = lazy_index(gtab.b0s_mask)
        self.where_dwi = lazy_index(~gtab.b0s_mask)

        # Correct gradient table to contain b0 data at the beginning
        bvals_cor = np.concatenate(([0], gtab.bvals[self.where_dwi]))
        bvecs_cor = np.concatenate(([[0, 0, 0]], gtab.bvecs[self.where_dwi]))
        gtab_cor = gradient_table(bvals_cor, bvecs=bvecs_cor)

        # Initialize self.gtab
        OdfModel.__init__(self, gtab_cor)

        if isinstance(wm_response, tuple):
            wm_response = np.asarray(wm_response)

        # Store responses
        self.wm_response = wm_response
        self.gm_response = gm_response
        self.csf_response = csf_response

        # Initializing remaining parameters
        if R < 1 or n_iter < 1 or n_coils < 1:
            raise ValueError(
                f"R, n_iter, and n_coils must be >= 1, but R={R},"
                + f"n_iter={n_iter}, and n_coils={n_coils} "
            )

        self.R = R
        self.n_iter = n_iter
        self.recon_type = recon_type
        self.n_coils = n_coils

        if voxelwise and use_tv:
            raise ValueError("Total variation has no effect in voxelwise fit")
        if voxelwise and verbose:
            warnings.warn(
                "Verbosity has no effect in voxelwise fit", UserWarning, stacklevel=2
            )

        self.voxelwise = voxelwise
        self.use_tv = use_tv

        self.verbose = verbose

        if sphere is None:
            self.sphere = get_sphere(name="repulsion724")
        else:
            self.sphere = sphere

        if voxelwise:
            self.fit = self._voxelwise_fit
        else:
            self.fit = self._global_fit

        # Fitting parameters
        self.kernel = None

    @warning_for_keywords()
    def _global_fit(self, data, *, mask=None):
        """
        Fit fODF and GM/CSF volume fractions globally.

        Parameters
        ----------
        data : ndarray (x, y, z, N)
            Signal values for each voxel. Must be 4D.
        mask : ndarray (x, y, z), optional
            Binary mask specifying voxels of interest with 1; results will only
            be fit at these voxels (0 elsewhere). If `None`, fits all voxels.
            Default: None.

        Returns
        -------
        model_fit : RumbaFit
            Fit object storing model parameters.

        """

        # Checking data and mask shapes
        if len(data.shape) != 4:
            raise ValueError(f"Data should be 4D, received shape f{data.shape}")

        if mask is None:  # default mask includes all voxels
            mask = np.ones(data.shape[:3])

        if data.shape[:3] != mask.shape:
            raise ValueError(
                "Mask shape should match first 3 dimensions of "
                + f"data, but data dimensions are f{data.shape} "
                + f"while mask dimensions are f{mask.shape}"
            )

        # Signal repair, normalization

        # Normalize data to mean b0 image
        data = normalize_data(data, self.where_b0s, min_signal=_EPS)
        # Rearrange data to match corrected gradient table
        data = np.concatenate(
            (np.ones([*data.shape[:3], 1]), data[..., self.where_dwi]), axis=3
        )
        data[data > 1] = 1  # clip values between 0 and 1

        # All arrays are converted to float32 to reduce memory load
        data = data.astype(np.float32)

        # Generate kernel
        self.kernel = generate_kernel(
            self.gtab,
            self.sphere,
            self.wm_response,
            self.gm_response,
            self.csf_response,
        ).astype(np.float32)

        # Fit fODF
        model_params = rumba_deconv_global(
            data,
            self.kernel,
            mask,
            n_iter=self.n_iter,
            recon_type=self.recon_type,
            n_coils=self.n_coils,
            R=self.R,
            use_tv=self.use_tv,
            verbose=self.verbose,
        )

        model_fit = RumbaFit(self, model_params)
        return model_fit

    @warning_for_keywords()
    def _voxelwise_fit(self, data, *, mask=None):
        """
        Fit fODF and GM/CSF volume fractions voxelwise.

        Parameters
        ----------
        data : ndarray ([x, y, z], N)
            Signal values for each voxel.
        mask : ndarray ([x, y, z]), optional
            Binary mask specifying voxels of interest with 1; results will only
            be fit at these voxels (0 elsewhere). If `None`, fits all voxels.
            Default: None.

        Returns
        -------
        model_fit : RumbaFit
            Fit object storing model parameters.

        """

        if mask is None:  # default mask includes all voxels
            mask = np.ones(data.shape[:-1])

        if data.shape[:-1] != mask.shape:
            raise ValueError(
                "Mask shape should match first dimensions of "
                + f"data, but data dimensions are f{data.shape} "
                + f"while mask dimensions are f{mask.shape}"
            )

        self.kernel = generate_kernel(
            self.gtab,
            self.sphere,
            self.wm_response,
            self.gm_response,
            self.csf_response,
        )

        model_params = np.zeros(data.shape[:-1] + (len(self.sphere.vertices) + 2,))

        for ijk in np.ndindex(data.shape[:-1]):
            if mask[ijk]:
                vox_data = data[ijk]
                # Normalize data to mean b0 image
                vox_data = normalize_data(vox_data, self.where_b0s, min_signal=_EPS)
                # Rearrange data to match corrected gradient table
                vox_data = np.concatenate(([1], vox_data[self.where_dwi]))
                vox_data[vox_data > 1] = 1  # clip values between 0 and 1

                # Fitting
                model_param = rumba_deconv(
                    vox_data,
                    self.kernel,
                    n_iter=self.n_iter,
                    recon_type=self.recon_type,
                    n_coils=self.n_coils,
                )

                model_params[ijk] = model_param

        model_fit = RumbaFit(self, model_params)
        return model_fit


class RumbaFit(OdfFit):
    def __init__(self, model, model_params):
        """
        Constructs fODF, GM/CSF volume fractions, and other derived results.

        fODF and GM/CSF fractions are normalized to collectively sum to 1 for
        each voxel.

        Parameters
        ----------
        model : RumbaSDModel
            RumbaSDModel-SD model.
        model_params : ndarray ([x, y, z], M)
            fODF and GM/CSF volume fractions for each voxel.

        """

        self.model = model
        self.model_params = model_params

    @warning_for_keywords()
    def odf(self, *, sphere=None):
        """
        Constructs fODF at discrete vertices on model sphere for each voxel.

        Parameters
        ----------
        sphere : Sphere, optional
            Sphere on which to construct fODF. If specified, must be the same
            sphere used by the `RumbaSDModel` model. Default: None.

        Returns
        -------
        odf : ndarray ([x, y, z], M-2)
            fODF computed at each vertex on model sphere.

        """
        if sphere is not None and sphere != self.model.sphere:
            raise ValueError(
                "Reconstruction sphere must be the same as used"
                + " in the RUMBA-SD model."
            )

        odf = self.model_params[..., :-2]
        return odf

    @auto_attr
    def f_gm(self):
        """
        Constructs GM volume fraction for each voxel.

        Returns
        -------
        f_gm : ndarray ([x, y, z])
            GM volume fraction.
        """

        f_gm = self.model_params[..., -2]
        return f_gm

    @auto_attr
    def f_csf(self):
        """
        Constructs CSF volume fraction for each voxel.

        Returns
        -------
        f_csf : ndarray ([x, y, z])
            CSF volume fraction.
        """

        f_csf = self.model_params[..., -1]
        return f_csf

    @auto_attr
    def f_wm(self):
        """
        Constructs white matter volume fraction for each voxel.

        Equivalent to sum of fODF.

        Returns
        -------
        f_wm : ndarray ([x, y, z])
            White matter volume fraction.
        """

        f_wm = np.sum(self.odf(), axis=-1)
        return f_wm

    @auto_attr
    def f_iso(self):
        """
        Constructs isotropic volume fraction for each voxel.

        Equivalent to sum of GM and CSF volume fractions.

        Returns
        -------
        f_iso : ndarray ([x, y, z])
            Isotropic volume fraction.
        """

        f_iso = self.f_gm + self.f_csf
        return f_iso

    @auto_attr
    def combined_odf_iso(self):
        """
        Constructs fODF combined with isotropic volume fraction at discrete
        vertices on model sphere.

        Distributes isotropic compartments evenly along each fODF direction.
        Sums to 1.

        Returns
        -------
        combined : ndarray ([x, y, z], M-2)
            fODF combined with isotropic volume fraction.
        """

        odf = self.odf()
        combined = odf + self.f_iso[..., None] / odf.shape[-1]
        return combined

    @warning_for_keywords()
    def predict(self, *, gtab=None, S0=None):
        """
        Compute signal prediction on model gradient table given given fODF
        and GM/CSF volume fractions for each voxel.

        Parameters
        ----------
        gtab : GradientTable, optional
            The gradients for which the signal will be predicted. Use the
            model's gradient table if `None`. Default: None
        S0 : ndarray ([x, y, z]) or float, optional
            The non diffusion-weighted signal value for each voxel. If a float,
            the same value is used for each voxel. If `None`, 1 is used for
            each voxel. Default: None

        Returns
        -------
        pred_sig : ndarray ([x, y, z], N)
            The predicted signal.

        """
        model_params = self.model_params

        if gtab is None:
            gtab = self.model.gtab_orig

        pred_kernel = generate_kernel(
            gtab,
            self.model.sphere,
            self.model.wm_response,
            self.model.gm_response,
            self.model.csf_response,
        )

        if S0 is None:
            S0 = np.ones(model_params.shape[:-1] + (1,))
        elif isinstance(S0, np.ndarray):
            S0 = S0[..., None]
        else:
            S0 = S0 * np.ones(model_params.shape[:-1] + (1,))

        pred_sig = np.empty(model_params.shape[:-1] + (len(gtab.bvals),))
        for ijk in np.ndindex(model_params.shape[:-1]):
            pred_sig[ijk] = S0[ijk] * np.dot(pred_kernel, model_params[ijk])

        return pred_sig


@warning_for_keywords()
def rumba_deconv(data, kernel, *, n_iter=600, recon_type="smf", n_coils=1):
    r"""
    Fit fODF and GM/CSF volume fractions for a voxel using RUMBA-SD.

    Deconvolves the kernel from the diffusion-weighted signal by computing a
    maximum likelihood estimation of the fODF
    :footcite:p:`CanalesRodriguez2015`. Minimizes the negative log-likelihood of
    the data under Rician or Noncentral Chi noise distributions by adapting the
    iterative technique developed in Richardson-Lucy deconvolution.

    Parameters
    ----------
    data : 1d ndarray (N,)
        Signal values for a single voxel.
    kernel : 2d ndarray (N, M)
        Deconvolution kernel mapping volume fractions of the M compartments to
        N-length signal. Last two columns should be for GM and CSF.
    n_iter : int, optional
        Number of iterations for fODF estimation. Must be a positive int.
        Default: 600
    recon_type : {'smf', 'sos'}, optional
        MRI reconstruction method: spatial matched filter (SMF) or
        sum-of-squares (SoS). SMF reconstruction generates Rician noise while
        SoS reconstruction generates Noncentral Chi noise. Default: 'smf'
    n_coils : int, optional
        Number of coils in MRI scanner -- only relevant in SoS reconstruction.
        Must be a positive int. Default: 1

    Returns
    -------
    fit_vec : 1d ndarray (M,)
        Vector containing fODF and GM/CSF volume fractions. First M-2
        components are fODF while last two are GM and CSF respectively.

    Notes
    -----
    The diffusion MRI signal measured at a given voxel is a sum of
    contributions from each intra-voxel compartment, including parallel white
    matter (WM) fiber populations in a given orientation as well as effects
    from GM and CSF. The equation governing these  contributions is:

    $S_i = S_0\left(\sum_{j=1}^{M}f_j\exp(-b_i\textbf{v}_i^T\textbf{D}_j
    \textbf{v}_i) + f_{GM}\exp(-b_iD_{GM})+f_{CSF}\exp(-b_iD_{CSF})\right)$

    Where $S_i$ is the resultant signal along the diffusion-sensitizing
    gradient unit vector $\textbf{v_i}; i = 1, ..., N$ with a b-value of $b_i$.
    $f_j; j = 1, ..., M$ is the volume fraction of the $j^{th}$ fiber
    population with an anisotropic diffusion tensor $\textbf{D_j}$.

    $f_{GM}$ and $f_{CSF}$ are the volume fractions and $D_{GM}$ and $D_{CSF}$
    are the mean diffusivities of GM and CSF respectively.

    This equation is linear in $f_j, f_{GM}, f_{CSF}$ and can be simplified to
    a single matrix multiplication:

    $\textbf{S} = \textbf{Hf}$

    Where $\textbf{S}$ is the signal vector at a certain voxel, $\textbf{H}$ is
    the deconvolution kernel, and $\textbf{f}$ is the vector of volume
    fractions for each compartment.

    Modern MRI scanners produce noise following a Rician or Noncentral Chi
    distribution, depending on their signal reconstruction technique
    `footcite:p:`Constantinides1997`. Using this linear model, it can be shown
    that the likelihood of a signal under a Noncentral Chi noise model is:

    $P(\textbf{S}|\textbf{H}, \textbf{f}, \sigma^2, n) = \prod_{i=1}^{N}\left(
    \frac{S_i}{\bar{S_i}}\right)^n\exp\left\{-\frac{1}{2\sigma^2}\left[
    S_i^2 + \bar{S}_i^2\right]\right\}I_{n-1}\left(\frac{S_i\bar{S}_i}
    {\sigma^2}\right)u(S_i)$

    Where $S_i$ and $\bar{S}_i = \textbf{Hf}$ are the measured and expected
    signals respectively, and $n$ is the number of coils in the scanner, and
    $I_{n-1}$ is the modified Bessel function of first kind of order $n-1$.
    This gives the likelihood under a Rician distribution when $n$ is set to 1.

    By taking the negative log of this with respect to $\textbf{f}$ and setting
    the derivative to 0, the $\textbf{f}$ maximizing likelihood is found to be:

    $\textbf{f} = \textbf{f} \circ \frac{\textbf{H}^T\left[\textbf{S}\circ
    \frac{I_n(\textbf{S}\circ \textbf{Hf}/\sigma^2)} {I_{n-1}(\textbf{S}
    \circ\textbf{Hf}\sigma^2)} \right ]} {\textbf{H}^T\textbf{Hf}}$

    The solution can be found using an iterative scheme, just as in the
    Richardson-Lucy algorithm:

    $\textbf{f}^{k+1} = \textbf{f}^k \circ \frac{\textbf{H}^T\left[\textbf{S}
    \circ\frac{I_n(\textbf{S}\circ\textbf{Hf}^k/\sigma^2)} {I_{n-1}(\textbf{S}
    \circ\textbf{Hf}^k/\sigma^2)} \right ]} {\textbf{H}^T\textbf{Hf}^k}$

    In order to apply this, a reasonable estimate of $\sigma^2$ is required.
    To find this, a separate iterative scheme is found using the derivative
    of the negative log with respect to $\sigma^2$, and is run in parallel.
    This is shown here:

    $\alpha^{k+1} = \frac{1}{nN}\left\{ \frac{\textbf{S}^T\textbf{S} +
    \textbf{f}^T\textbf{H}^T\textbf{Hf}}{2} - \textbf{1}^T_N\left[(\textbf{S}
    \circ\textbf{Hf})\circ\frac{I_n(\textbf{S}\circ\textbf{Hf}/\alpha^k)}
    {I_{n-1}(\textbf{S}\circ\textbf{Hf}/\alpha^k)} \right ]\right \}$

    For more details, see :footcite:p:`CanalesRodriguez2015`.

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

    n_comp = kernel.shape[1]  # number of compartments
    n_grad = len(data)  # gradient directions

    fodf = np.ones((n_comp, 1))  # initial guess is iso-probable
    fodf = fodf / np.sum(fodf, axis=0)  # normalize initial guess

    if recon_type == "smf":
        n_order = 1  # Rician noise (same as Noncentral Chi with order 1)
    elif recon_type == "sos":
        n_order = n_coils  # Noncentral Chi noise (order = # of coils)
    else:
        raise ValueError(
            f"Invalid recon_type. Should be 'smf' or 'sos', received {recon_type}"
        )

    data = data.reshape(-1, 1)
    reblurred = np.matmul(kernel, fodf)

    # For use later
    kernel_t = kernel.T
    f_zero = 0  # minimum value allowed in fODF

    # Initialize variance
    sigma0 = 1 / 15
    sigma2 = sigma0**2 * np.ones(data.shape)  # Expand into vector

    reblurred_s = data * reblurred / sigma2

    for _ in range(n_iter):
        fodf_i = fodf
        ratio = mbessel_ratio(n_order, reblurred_s)
        rl_factor = np.matmul(kernel_t, data * ratio) / (
            np.matmul(kernel_t, reblurred) + _EPS
        )

        fodf = fodf_i * rl_factor  # result of iteration
        fodf = np.maximum(f_zero, fodf)  # apply positivity constraint

        # Update other variables
        reblurred = np.matmul(kernel, fodf)
        reblurred_s = data * reblurred / sigma2

        # Iterate variance
        sigma2_i = (1 / (n_grad * n_order)) * np.sum(
            (data**2 + reblurred**2) / 2 - (sigma2 * reblurred_s) * ratio, axis=0
        )
        sigma2_i = np.minimum((1 / 8) ** 2, np.maximum(sigma2_i, (1 / 80) ** 2))
        sigma2 = sigma2_i * np.ones(data.shape)  # Expand into vector

    # Normalize final result
    fit_vec = np.squeeze(fodf / (np.sum(fodf, axis=0) + _EPS))

    return fit_vec


def mbessel_ratio(n, x):
    r"""
    Fast computation of modified Bessel function ratio (first kind).

    Computes:

    $I_{n}(x) / I_{n-1}(x)$

    using Perron's continued fraction equation where $I_n$ is the modified
    Bessel function of first kind, order $n$ :footcite:p:`Gautschi1978`.

    Parameters
    ----------
    n : int
        Order of Bessel function in numerator (denominator is of order n-1).
        Must be a positive int.
    x : float or ndarray
        Value or array of values with which to compute ratio.

    Returns
    -------
    y : float or ndarray
        Result of ratio computation.

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

    y = x / (
        (2 * n + x)
        - (
            2
            * x
            * (n + 1 / 2)
            / (
                2 * n
                + 1
                + 2 * x
                - (
                    2
                    * x
                    * (n + 3 / 2)
                    / (2 * n + 2 + 2 * x - (2 * x * (n + 5 / 2) / (2 * n + 3 + 2 * x)))
                )
            )
        )
    )

    return y


def generate_kernel(gtab, sphere, wm_response, gm_response, csf_response):
    """
    Generate deconvolution kernel

    Compute kernel mapping orientation densities of white matter fiber
    populations (along each vertex of the sphere) and isotropic volume
    fractions to a diffusion weighted signal.

    Parameters
    ----------
    gtab : GradientTable
        Gradient table.
    sphere : Sphere
        Sphere with which to sample discrete fiber orientations in order to
        construct kernel
    wm_response : 1d ndarray or 2d ndarray or AxSymShResponse, optional
        Tensor eigenvalues as a (3,) ndarray, multishell eigenvalues as
        a (len(unique_bvals_tolerance(gtab.bvals))-1, 3) ndarray in
        order of smallest to largest b-value, or an AxSymShResponse.
    gm_response : float, optional
        Mean diffusivity for GM compartment. If `None`, then grey
        matter compartment set to all zeros.
    csf_response : float, optional
        Mean diffusivity for CSF compartment. If `None`, then CSF
        compartment set to all zeros.

    Returns
    -------
    kernel : 2d ndarray (N, M)
        Computed kernel; can be multiplied with a vector consisting of volume
        fractions for each of M-2 fiber populations as well as GM and CSF
        fractions to produce a diffusion weighted signal.
    """

    # Coordinates of sphere vertices
    sticks = sphere.vertices

    n_grad = len(gtab.gradients)  # number of gradient directions
    n_wm_comp = sticks.shape[0]  # number of fiber populations
    n_comp = n_wm_comp + 2  # plus isotropic compartments

    kernel = np.zeros((n_grad, n_comp))

    # White matter compartments
    list_bvals = unique_bvals_tolerance(gtab.bvals)
    n_bvals = len(list_bvals) - 1  # number of unique b-values

    if isinstance(wm_response, AxSymShResponse):
        # Data-driven response
        where_dwi = lazy_index(~gtab.b0s_mask)
        gradients = gtab.gradients[where_dwi]
        gradients = gradients / np.linalg.norm(gradients, axis=1)[..., None]
        S0 = wm_response.S0
        for i in range(n_wm_comp):
            # Response oriented along [0, 0, 1], so must rotate sticks[i]
            rot_mat = vec2vec_rotmat(sticks[i], np.array([0, 0, 1]))
            rot_gradients = np.dot(rot_mat, gradients.T).T
            rot_sphere = Sphere(xyz=rot_gradients)
            # Project onto rotated sphere and scale
            rot_response = wm_response.on_sphere(rot_sphere) / S0
            kernel[where_dwi, i] = rot_response

        # Set b0 components
        kernel[gtab.b0s_mask, :] = 1

    elif wm_response.shape == (n_bvals, 3):
        # Multi-shell response
        bvals = gtab.bvals
        bvecs = gtab.bvecs
        for n, bval in enumerate(list_bvals[1:]):
            indices = get_bval_indices(bvals, bval)
            with warnings.catch_warnings():  # extract relevant b-value
                warnings.simplefilter("ignore")
                gtab_sub = gradient_table(bvals[indices], bvecs=bvecs[indices])

            for i in range(n_wm_comp):
                # Signal generated by WM-fiber for each gradient direction
                S = single_tensor(
                    gtab_sub, evals=wm_response[n], evecs=all_tensor_evecs(sticks[i])
                )
                kernel[indices, i] = S

        # Set b0 components
        b0_indices = get_bval_indices(bvals, list_bvals[0])
        kernel[b0_indices, :] = 1

    else:
        # Single-shell response
        for i in range(n_wm_comp):
            # Signal generated by WM-fiber for each gradient direction
            S = single_tensor(
                gtab, evals=wm_response, evecs=all_tensor_evecs(sticks[i])
            )
            kernel[:, i] = S

        # Set b0 components
        kernel[gtab.b0s_mask, :] = 1

    # GM compartment
    if gm_response is None:
        S_gm = np.zeros(n_grad)
    else:
        S_gm = single_tensor(
            gtab, evals=np.array([gm_response, gm_response, gm_response])
        )

    if csf_response is None:
        S_csf = np.zeros(n_grad)
    else:
        S_csf = single_tensor(
            gtab, evals=np.array([csf_response, csf_response, csf_response])
        )

    kernel[:, n_comp - 2] = S_gm
    kernel[:, n_comp - 1] = S_csf

    return kernel


@warning_for_keywords()
def rumba_deconv_global(
    data,
    kernel,
    mask,
    *,
    n_iter=600,
    recon_type="smf",
    n_coils=1,
    R=1,
    use_tv=True,
    verbose=False,
):
    r"""
    Fit fODF for all voxels simultaneously using RUMBA-SD.

    Deconvolves the kernel from the diffusion-weighted signal at each voxel by
    computing a maximum likelihood estimation of the fODF
    :footcite:p:`CanalesRodriguez2015`. Global fitting also permits the use of
    total variation regularization (RUMBA-SD + TV). The spatial dependence
    introduced by TV promotes smoother solutions (i.e. prevents oscillations),
    while still allowing for sharp discontinuities :footcite:p:`Rudin1992`. This
    promotes smoothness and continuity along individual tracts while preventing
    smoothing of adjacent tracts.

    Generally, global_fit will proceed more quickly than the voxelwise fit
    provided that the computer has adequate RAM (>= 16 GB should be more than
    sufficient).

    Parameters
    ----------
    data : 4d ndarray (x, y, z, N)
        Signal values for entire brain. None of the volume dimensions x, y, z
        can be 1 if TV regularization is required.
    kernel : 2d ndarray (N, M)
        Deconvolution kernel mapping volume fractions of the M compartments to
        N-length signal. Last two columns should be for GM and CSF.
    mask : 3d ndarray(x, y, z)
        Binary mask specifying voxels of interest with 1; fODF will only be
        fit at these voxels (0 elsewhere).
    n_iter : int, optional
        Number of iterations for fODF estimation. Must be a positive int.
    recon_type : {'smf', 'sos'}, optional
        MRI reconstruction method: spatial matched filter (SMF) or
        sum-of-squares (SoS). SMF reconstruction generates Rician noise while
        SoS reconstruction generates Noncentral Chi noise.
    n_coils : int, optional
        Number of coils in MRI scanner -- only relevant in SoS reconstruction.
        Must be a positive int.
    use_tv : bool, optional
        If true, applies total variation regularization. This requires a brain
        volume with no singleton dimensions.
    verbose : bool, optional
        If true, logs updates on estimated signal-to-noise ratio after each
        iteration.

    Returns
    -------
    fit_array : 4d ndarray (x, y, z, M)
        fODF and GM/CSF volume fractions computed for each voxel. First M-2
        components are fODF, while last two are GM and CSf respectively.

    Notes
    -----
    TV modifies our cost function as follows:

    .. math::

    J(\textbf{f}) = -\log{P(\textbf{S}|\textbf{H}, \textbf{f}, \sigma^2, n)})+
    \alpha_{TV}TV(\textbf{f})

    where the first term is the negative log likelihood described in the notes
    of `rumba_deconv`, and the second term is the TV energy, or the sum of
    gradient absolute values for the fODF across the entire brain. This results
    in a new multiplicative factor in the iterative scheme, now becoming:

    .. math::

    \textbf{f}^{k+1} = \textbf{f}^k \circ \frac{\textbf{H}^T\left[\textbf{S}
    \circ\frac{I_n(\textbf{S}\circ\textbf{Hf}^k/\sigma^2)} {I_{n-1}(\textbf{S}
    \circ\textbf{Hf}^k/\sigma^2)} \right ]} {\textbf{H}^T\textbf{Hf}^k}\circ
    \textbf{R}^k

    where $\textbf{R}^k$ is computed voxelwise by:

    .. math::

    (\textbf{R}^k)_j = \frac{1}{1 - \alpha_{TV}div\left(\frac{\triangledown[
    \textbf{f}^k_{3D}]_j}{\lvert\triangledown[\textbf{f}^k_{3D}]_j \rvert}
    \right)\biggr\rvert_{x, y, z}}

    Here, $\triangledown$ is the symbol for the 3D gradient at any voxel.

    The regularization strength, $\alpha_{TV}$ is updated after each iteration
    by the discrepancy principle -- specifically, it is selected to match the
    estimated variance after each iteration :footcite:p:`Chambolle2004`.

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

    # Crop data to reduce memory consumption
    dim_orig = data.shape
    ixmin, ixmax = bounding_box(mask)
    data = crop(data, ixmin, ixmax)
    mask = crop(mask, ixmin, ixmax)

    if np.any(np.array(data.shape[:3]) == 1) and use_tv:
        raise ValueError(
            "Cannot use TV regularization if any spatial"
            + "dimensions are 1; "
            + f"provided dimensions were {data.shape[:3]}"
        )

    epsilon = 1e-7

    n_grad = kernel.shape[0]  # gradient directions
    n_comp = kernel.shape[1]  # number of compartments
    dim = data.shape
    n_v_tot = np.prod(dim[:3])  # total number of voxels

    # Initial guess is iso-probable
    fodf0 = np.ones((n_comp, 1), dtype=np.float32)
    fodf0 = fodf0 / np.sum(fodf0, axis=0)

    if recon_type == "smf":
        n_order = 1  # Rician noise (same as Noncentral Chi with order 1)
    elif recon_type == "sos":
        n_order = n_coils  # Noncentral Chi noise (order = # of coils)
    else:
        raise ValueError(
            f"Invalid recon_type. Should be 'smf' or 'sos', received f{recon_type}"
        )

    mask_vec = np.ravel(mask)
    # Indices of target voxels
    index_mask = np.atleast_1d(np.squeeze(np.argwhere(mask_vec)))
    n_v_true = len(index_mask)  # number of target voxels

    data_2d = np.zeros((n_v_true, n_grad), dtype=np.float32)
    for i in range(n_grad):
        data_2d[:, i] = np.ravel(data[:, :, :, i])[
            index_mask
        ]  # only keep voxels of interest

    data_2d = data_2d.T
    fodf = np.tile(fodf0, (1, n_v_true))
    reblurred = np.matmul(kernel, fodf)

    # For use later
    kernel_t = kernel.T
    f_zero = 0

    # Initialize algorithm parameters
    sigma0 = 1 / 15
    sigma2 = sigma0**2
    tv_lambda = sigma2  # initial guess for TV regularization strength

    # Expand into matrix form for iterations
    sigma2 = sigma2 * np.ones(data_2d.shape, dtype=np.float32)
    tv_lambda_aux = np.zeros(n_v_tot, dtype=np.float32)

    reblurred_s = data_2d * reblurred / sigma2

    for i in range(n_iter):
        fodf_i = fodf
        ratio = mbessel_ratio(n_order, reblurred_s).astype(np.float32)
        rl_factor = np.matmul(kernel_t, data_2d * ratio) / (
            np.matmul(kernel_t, reblurred) + _EPS
        )

        if use_tv:  # apply TV regularization
            tv_factor = np.ones(fodf_i.shape, dtype=np.float32)
            fodf_4d = _reshape_2d_4d(fodf_i.T, mask)
            # Compute gradient, divergence
            gr = _grad(fodf_4d)
            d_inv = 1 / np.sqrt(epsilon**2 + np.sum(gr**2, axis=3))
            gr_norm = gr * d_inv[:, :, :, None, :]
            div_f = _divergence(gr_norm)
            g0 = np.abs(1 - tv_lambda * div_f)
            tv_factor_4d = 1 / (g0 + _EPS)

            for j in range(n_comp):
                tv_factor_1d = np.ravel(tv_factor_4d[:, :, :, j])[index_mask]
                tv_factor[j, :] = tv_factor_1d

            # Apply TV regularization to iteration factor
            rl_factor = rl_factor * tv_factor

        fodf = fodf_i * rl_factor  # result of iteration
        fodf = np.maximum(f_zero, fodf)  # positivity constraint

        # Update other variables
        reblurred = np.matmul(kernel, fodf)
        reblurred_s = data_2d * reblurred / sigma2

        # Iterate variance
        sigma2_i = (1 / (n_grad * n_order)) * np.sum(
            (data_2d**2 + reblurred**2) / 2 - (sigma2 * reblurred_s) * ratio, axis=0
        )
        sigma2_i = np.minimum((1 / 8) ** 2, np.maximum(sigma2_i, (1 / 80) ** 2))

        if verbose:
            logger.info("Iteration %d of %d", i + 1, n_iter)

            snr_mean = np.mean(1 / np.sqrt(sigma2_i))
            snr_std = np.std(1 / np.sqrt(sigma2_i))
            logger.info(
                "Mean SNR (S0/sigma) estimated to be %.3f +/- %.3f", snr_mean, snr_std
            )
        # Expand into matrix
        sigma2 = np.tile(sigma2_i[None, :], (data_2d.shape[0], 1))

        # Update TV regularization strength using the discrepancy principle
        if use_tv:
            if R == 1:
                tv_lambda = np.mean(sigma2_i)

                if tv_lambda < (1 / 30) ** 2:
                    tv_lambda = (1 / 30) ** 2
            else:  # different factor for each voxel
                tv_lambda_aux[index_mask] = sigma2_i
                tv_lambda = np.reshape(tv_lambda_aux, (*dim[:3], 1))

    fodf = fodf.astype(np.float64)
    fodf = fodf / (np.sum(fodf, axis=0)[None, ...] + _EPS)  # normalize fODF

    # Extract compartments
    fit_array = np.zeros((*dim_orig[:3], n_comp))
    _reshape_2d_4d(
        fodf.T,
        mask,
        out=fit_array[ixmin[0] : ixmax[0], ixmin[1] : ixmax[1], ixmin[2] : ixmax[2]],
    )

    return fit_array


def _grad(M):
    """
    Computes one way first difference
    """
    x_ind = list(range(1, M.shape[0])) + [M.shape[0] - 1]
    y_ind = list(range(1, M.shape[1])) + [M.shape[1] - 1]
    z_ind = list(range(1, M.shape[2])) + [M.shape[2] - 1]

    grad = np.zeros((*M.shape[:3], 3, M.shape[-1]), dtype=np.float32)
    grad[:, :, :, 0, :] = M[x_ind, :, :, :] - M
    grad[:, :, :, 1, :] = M[:, y_ind, :, :] - M
    grad[:, :, :, 2, :] = M[:, :, z_ind, :] - M

    return grad


def _divergence(F):
    """
    Computes divergence of a 3-dimensional vector field (with one way
    first difference)
    """
    Fx = F[:, :, :, 0, :]
    Fy = F[:, :, :, 1, :]
    Fz = F[:, :, :, 2, :]

    x_ind = [0] + list(range(F.shape[0] - 1))
    y_ind = [0] + list(range(F.shape[1] - 1))
    z_ind = [0] + list(range(F.shape[2] - 1))

    fx = Fx - Fx[x_ind, :, :, :]
    fx[0, :, :, :] = Fx[0, :, :, :]  # edge conditions
    fx[-1, :, :, :] = -Fx[-2, :, :, :]

    fy = Fy - Fy[:, y_ind, :, :]
    fy[:, 0, :, :] = Fy[:, 0, :, :]
    fy[:, -1, :, :] = -Fy[:, -2, :, :]

    fz = Fz - Fz[:, :, z_ind, :]
    fz[:, :, 0, :] = Fz[:, :, 0, :]
    fz[:, :, -1, :] = -Fz[:, :, -2, :]

    return fx + fy + fz


@warning_for_keywords()
def _reshape_2d_4d(M, mask, *, out=None):
    """
    Faster reshape from 2D to 4D.
    """
    if out is None:
        out = np.zeros((*mask.shape, M.shape[-1]), dtype=M.dtype)
    n = 0
    for i, j, k in np.ndindex(mask.shape):
        if mask[i, j, k]:
            out[i, j, k, :] = M[n, :]
            n += 1
    return out
