# Source code for dipy.reconst.qti

    """Classes and functions for fitting the covariance tensor model of q-space
    trajectory imaging (QTI) by :footcite:t:`Westin2016`.

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

    from warnings import warn

    import numpy as np

    from dipy.reconst.base import ReconstModel
    from dipy.reconst.dti import auto_attr
    from dipy.testing.decorators import warning_for_keywords
    from dipy.utils.optpkg import optional_package

    cp, have_cvxpy, _ = optional_package("cvxpy", min_version="1.4.1")

    # XXX Eventually to be replaced with `reconst.dti.lower_triangular`

    def from_3x3_to_6x1(T):
        """Convert symmetric 3 x 3 matrices into 6 x 1 vectors.

        Parameters
        ----------
        T : numpy.ndarray
            An array of size (..., 3, 3).

        Returns
        -------
        V : numpy.ndarray
            Converted vectors of size (..., 6, 1).

        Notes
        -----
        The conversion of a matrix into a vector is defined as

            .. math::

                \\mathbf{V} = \\begin{bmatrix}
                T_{11} & T_{22} & T_{33} &
                \\sqrt{2} T_{23} & \\sqrt{2} T_{13} & \\sqrt{2} T_{12}
                \\end{bmatrix}^T
        """
        if T.shape[-2::] != (3, 3):
            raise ValueError("The shape of the input array must be (..., 3, 3).")
        if not np.all(np.isclose(T, np.swapaxes(T, -1, -2))):
            warn(
                "All matrices converted to Voigt notation are not symmetric.", stacklevel=2
            )
        C = np.sqrt(2)
        V = np.stack(
            (
                T[..., 0, 0],
                T[..., 1, 1],
                T[..., 2, 2],
                C * T[..., 1, 2],
                C * T[..., 0, 2],
                C * T[..., 0, 1],
            ),
            axis=-1,
        )[..., np.newaxis]
        return V

    def from_6x1_to_3x3(V):
        """Convert 6 x 1 vectors into symmetric 3 x 3 matrices.

        Parameters
        ----------
        V : numpy.ndarray
            An array of size (..., 6, 1).

        Returns
        -------
        T : numpy.ndarray
            Converted matrices of size (..., 3, 3).

        Notes
        -----
        The conversion of a matrix into a vector is defined as

            .. math::

                \\mathbf{V} = \\begin{bmatrix}
                T_{11} & T_{22} & T_{33} &
                \\sqrt{2} T_{23} & \\sqrt{2} T_{13} & \\sqrt{2} T_{12}
                \\end{bmatrix}^T
        """
        if V.shape[-2::] != (6, 1):
            raise ValueError("The shape of the input array must be (..., 6, 1).")
        C = 1 / np.sqrt(2)
        T = np.array(
            (
                [V[..., 0, 0], C * V[..., 5, 0], C * V[..., 4, 0]],
                [C * V[..., 5, 0], V[..., 1, 0], C * V[..., 3, 0]],
                [C * V[..., 4, 0], C * V[..., 3, 0], V[..., 2, 0]],
            )
        )
        T = np.moveaxis(T, (0, 1), (-2, -1))
        return T

    def from_6x6_to_21x1(T):
        """Convert symmetric 6 x 6 matrices into 21 x 1 vectors.

        Parameters
        ----------
        T : numpy.ndarray
            An array of size (..., 6, 6).

        Returns
        -------
        V : numpy.ndarray
            Converted vectors of size (..., 21, 1).

        Notes
        -----
        The conversion of a matrix into a vector is defined as

            .. math::

                \\begin{matrix}
                \\mathbf{V} = & \\big[
                T_{11} & T_{22} & T_{33} \\
                & \\sqrt{2} T_{23} & \\sqrt{2} T_{13} & \\sqrt{2} T_{12} \\
                & \\sqrt{2} T_{14} & \\sqrt{2} T_{15} & \\sqrt{2} T_{16} \\
                & \\sqrt{2} T_{24} & \\sqrt{2} T_{25} & \\sqrt{2} T_{26} \\
                & \\sqrt{2} T_{34} & \\sqrt{2} T_{35} & \\sqrt{2} T_{36} \\
                & T_{44} & T_{55} & T_{66} \\
                & \\sqrt{2} T_{45} & \\sqrt{2} T_{56} & \\sqrt{2} T_{46} \\big]^T
                \\end{matrix}
        """
        if T.shape[-2::] != (6, 6):
            raise ValueError("The shape of the input array must be (..., 6, 6).")
        if not np.all(np.isclose(T, np.swapaxes(T, -1, -2), equal_nan=True)):
            warn(
                "All matrices converted to Voigt notation are not symmetric.", stacklevel=2
            )
        C = np.sqrt(2)
        V = np.stack(
            (
                [
                    T[..., 0, 0],
                    T[..., 1, 1],
                    T[..., 2, 2],
                    C * T[..., 1, 2],
                    C * T[..., 0, 2],
                    C * T[..., 0, 1],
                    C * T[..., 0, 3],
                    C * T[..., 0, 4],
                    C * T[..., 0, 5],
                    C * T[..., 1, 3],
                    C * T[..., 1, 4],
                    C * T[..., 1, 5],
                    C * T[..., 2, 3],
                    C * T[..., 2, 4],
                    C * T[..., 2, 5],
                    T[..., 3, 3],
                    T[..., 4, 4],
                    T[..., 5, 5],
                    C * T[..., 3, 4],
                    C * T[..., 4, 5],
                    C * T[..., 3, 5],
                ]
            ),
            axis=-1,
        )[..., np.newaxis]
        return V

    def from_21x1_to_6x6(V):
        """Convert 21 x 1 vectors into symmetric 6 x 6 matrices.

        Parameters
        ----------
        V : numpy.ndarray
            An array of size (..., 21, 1).

        Returns
        -------
        T : numpy.ndarray
            Converted matrices of size (..., 6, 6).

        Notes
        -----
        The conversion of a matrix into a vector is defined as

            .. math::

                \\begin{matrix}
                \\mathbf{V} = & \\big[
                T_{11} & T_{22} & T_{33} \\
                & \\sqrt{2} T_{23} & \\sqrt{2} T_{13} & \\sqrt{2} T_{12} \\
                & \\sqrt{2} T_{14} & \\sqrt{2} T_{15} & \\sqrt{2} T_{16} \\
                & \\sqrt{2} T_{24} & \\sqrt{2} T_{25} & \\sqrt{2} T_{26} \\
                & \\sqrt{2} T_{34} & \\sqrt{2} T_{35} & \\sqrt{2} T_{36} \\
                & T_{44} & T_{55} & T_{66} \\
                & \\sqrt{2} T_{45} & \\sqrt{2} T_{56} & \\sqrt{2} T_{46} \\big]^T
                \\end{matrix}
        """
        if V.shape[-2::] != (21, 1):
            raise ValueError("The shape of the input array must be (..., 21, 1).")
        C = 1 / np.sqrt(2)
        T = np.array(
            (
                [
                    V[..., 0, 0],
                    C * V[..., 5, 0],
                    C * V[..., 4, 0],
                    C * V[..., 6, 0],
                    C * V[..., 7, 0],
                    C * V[..., 8, 0],
                ],
                [
                    C * V[..., 5, 0],
                    V[..., 1, 0],
                    C * V[..., 3, 0],
                    C * V[..., 9, 0],
                    C * V[..., 10, 0],
                    C * V[..., 11, 0],
                ],
                [
                    C * V[..., 4, 0],
                    C * V[..., 3, 0],
                    V[..., 2, 0],
                    C * V[..., 12, 0],
                    C * V[..., 13, 0],
                    C * V[..., 14, 0],
                ],
                [
                    C * V[..., 6, 0],
                    C * V[..., 9, 0],
                    C * V[..., 12, 0],
                    V[..., 15, 0],
                    C * V[..., 18, 0],
                    C * V[..., 20, 0],
                ],
                [
                    C * V[..., 7, 0],
                    C * V[..., 10, 0],
                    C * V[..., 13, 0],
                    C * V[..., 18, 0],
                    V[..., 16, 0],
                    C * V[..., 19, 0],
                ],
                [
                    C * V[..., 8, 0],
                    C * V[..., 11, 0],
                    C * V[..., 14, 0],
                    C * V[..., 20, 0],
                    C * V[..., 19, 0],
                    V[..., 17, 0],
                ],
            )
        )
        T = np.moveaxis(T, (0, 1), (-2, -1))
        return T

    def cvxpy_1x6_to_3x3(V):
        """Convert a 1 x 6 vector into a symmetric 3 x 3 matrix.

        Parameters
        ----------
        V : numpy.ndarray
            An array of size (1, 6).

        Returns
        -------
        T : cvxpy.bmat
            Converted matrix of size (3, 3).

        Notes
        -----
        The conversion of a matrix into a vector is defined as

            .. math::

                \\mathbf{V} = \\begin{bmatrix}
                T_{11} & T_{22} & T_{33} &
                \\sqrt{2} T_{23} & \\sqrt{2} T_{13} & \\sqrt{2} T_{12}
                \\end{bmatrix}^T
        """
        if V.shape[0] == 6:
            V = V.T

        f = 1 / np.sqrt(2)

        T = cp.bmat(
            [
                [V[0, 0], f * V[0, 5], f * V[0, 4]],
                [f * V[0, 5], V[0, 1], f * V[0, 3]],
                [f * V[0, 4], f * V[0, 3], V[0, 2]],
            ]
        )
        return T

    def cvxpy_1x21_to_6x6(V):
        """Convert 1 x 21 vector into a symmetric 6 x 6 matrix.

        Parameters
        ----------
        V : numpy.ndarray
            An array of size (1, 21).

        Returns
        -------
        T : cvxpy.bmat
            Converted matrices of size (6, 6).

        Notes
        -----
        The conversion of a matrix into a vector is defined as

            .. math::

                \\begin{matrix}
                \\mathbf{V} = & \\big[
                T_{11} & T_{22} & T_{33} \\
                & \\sqrt{2} T_{23} & \\sqrt{2} T_{13} & \\sqrt{2} T_{12} \\
                & \\sqrt{2} T_{14} & \\sqrt{2} T_{15} & \\sqrt{2} T_{16} \\
                & \\sqrt{2} T_{24} & \\sqrt{2} T_{25} & \\sqrt{2} T_{26} \\
                & \\sqrt{2} T_{34} & \\sqrt{2} T_{35} & \\sqrt{2} T_{36} \\
                & T_{44} & T_{55} & T_{66} \\
                & \\sqrt{2} T_{45} & \\sqrt{2} T_{56} & \\sqrt{2} T_{46} \\big]^T
                \\end{matrix}
        """
        if V.shape[0] == 21:
            V = V.T

        f = 1 / np.sqrt(2)

        T = cp.bmat(
            [
                [V[0, 0], f * V[0, 5], f * V[0, 4], f * V[0, 6], f * V[0, 7], f * V[0, 8]],
                [
                    f * V[0, 5],
                    V[0, 1],
                    f * V[0, 3],
                    f * V[0, 9],
                    f * V[0, 10],
                    f * V[0, 11],
                ],
                [
                    f * V[0, 4],
                    f * V[0, 3],
                    V[0, 2],
                    f * V[0, 12],
                    f * V[0, 13],
                    f * V[0, 14],
                ],
                [
                    f * V[0, 6],
                    f * V[0, 9],
                    f * V[0, 12],
                    V[0, 15],
                    f * V[0, 18],
                    f * V[0, 20],
                ],
                [
                    f * V[0, 7],
                    f * V[0, 10],
                    f * V[0, 13],
                    f * V[0, 18],
                    V[0, 16],
                    f * V[0, 19],
                ],
                [
                    f * V[0, 8],
                    f * V[0, 11],
                    f * V[0, 14],
                    f * V[0, 20],
                    f * V[0, 19],
                    V[0, 17],
                ],
            ]
        )
        return T

    # These tensors are used in the calculation of the QTI parameters
    e_iso = np.eye(3) / 3
    E_iso = np.eye(6) / 3
    E_bulk = from_3x3_to_6x1(e_iso) @ from_3x3_to_6x1(e_iso).T
    E_shear = E_iso - E_bulk
    E_tsym = E_bulk + 0.4 * E_shear

    def dtd_covariance(DTD):
        """Calculate covariance of a diffusion tensor distribution (DTD).

        Parameters
        ----------
        DTD : numpy.ndarray
            Diffusion tensor distribution of shape (number of tensors, 3, 3) or
            (number of tensors, 6, 1).

        Returns
        -------
        C : numpy.ndarray
            Covariance tensor of shape (6, 6).

        Notes
        -----
        The covariance tensor is calculated according to the following equation and
        converted into a rank-2 tensor :footcite:p:`Westin2016`:

            .. math::

                \\mathbb{C} = \\langle \\mathbf{D} \\otimes \\mathbf{D} \\rangle -
                \\langle \\mathbf{D} \\rangle \\otimes \\langle \\mathbf{D}
                \\rangle

        References
        ----------
        .. footbibliography::
        """
        dims = DTD.shape
        if len(dims) != 3 or (dims[1:3] != (3, 3) and dims[1:3] != (6, 1)):
            raise ValueError(
                "The shape of DTD must be (number of tensors, 3, 3) or (number of "
                + "tensors, 6, 1)."
            )
        if dims[1:3] == (3, 3):
            DTD = from_3x3_to_6x1(DTD)
        D = np.mean(DTD, axis=0)
        C = np.mean(DTD @ np.swapaxes(DTD, -2, -1), axis=0) - D @ np.swapaxes(D, -2, -1)
        return C

    @warning_for_keywords()
    def qti_signal(gtab, D, C, *, S0=1):
        """Generate signals using the covariance tensor signal representation.

        Parameters
        ----------
        gtab : dipy.core.gradients.GradientTable
            Gradient table with b-tensors.
        D : numpy.ndarray
            Diffusion tensors of shape (..., 3, 3), (..., 6, 1), or (..., 6).
        C : numpy.ndarray
            Covariance tensors of shape (..., 6, 6), (..., 21, 1), or (..., 21).
        S0 : numpy.ndarray, optional
            Signal magnitudes without diffusion-weighting. Must be a single number
            or an array of same shape as D and C without the last two dimensions.

        Returns
        -------
        S : numpy.ndarray
            Simulated signals.

        Notes
        -----
        The signal is generated according to

            .. math::

                S = S_0 \\exp \\left(- \\mathbf{b} : \\langle \\mathbf{D} \\rangle
                + \\frac{1}{2}(\\mathbf{b} \\otimes \\mathbf{b}) : \\mathbb{C}
                \\right)
        """

        # Validate input and convert to Voigt notation if necessary
        if gtab.btens is None:
            raise ValueError("QTI requires b-tensors to be defined in the gradient table.")
        if D.shape[-2::] != (6, 1):
            if D.shape[-2::] == (3, 3):
                D = from_3x3_to_6x1(D)
            elif D.shape[-1] == 6:
                D = D[..., np.newaxis]
            else:
                raise ValueError(
                    "The shape of D must be (..., 3, 3), (..., 6, 1) or (..., 6)."
                )
        if C.shape[-2::] != (21, 1):
            if C.shape[-2::] == (6, 6):
                C = from_6x6_to_21x1(C)
            elif C.shape[-1] == 21:
                C = C[..., np.newaxis]
            else:
                raise ValueError(
                    "The shape of C must be (..., 6, 6), (..., 21, 1), or (..., 21)."
                )
        if D.shape[0:-2] != C.shape[0:-2]:
            raise ValueError("The shapes of C and D are not compatible")
        if not isinstance(S0, (int, float)):
            if S0.shape != (1,) and S0.shape != D.shape[0:-2]:
                raise ValueError(
                    "S0 must be a single number or an array of the same shape "
                    + " compatible with D and C."
                )

        # Generate signals
        S = np.zeros(D.shape[0:-2] + (gtab.btens.shape[0],))
        for i, bten in enumerate(gtab.btens):
            b = from_3x3_to_6x1(bten)
            b_sq = from_6x6_to_21x1(b @ np.swapaxes(b, -2, -1))
            S[..., i] = (
                S0
                * np.exp(-np.swapaxes(b, -2, -1) @ D + 0.5 * np.swapaxes(b_sq, -2, -1) @ C)[
                    ..., 0, 0
                ]
            )
        return S

    def design_matrix(btens):
        """Calculate the design matrix from the b-tensors.

        Parameters
        ----------
        btens : numpy.ndarray
            An array of b-tensors of shape (number of acquisitions, 3, 3).

        Returns
        -------
        X : numpy.ndarray
            Design matrix.

        Notes
        -----
        The design matrix is generated according to

            .. math::

                X = \\begin{pmatrix} 1 & -\\mathbf{b}_1^T & \\frac{1}{2}(
                \\mathbf{b}_1
                \\otimes\\mathbf{b}_1)^T \\ \\vdots & \\vdots & \\vdots \\ 1 &
                -\\mathbf{b}_n^T & \\frac{1}{2}(\\mathbf{b}_n\\otimes
                \\mathbf{b}_n)^T
                \\end{pmatrix}
        """
        X = np.zeros((btens.shape[0], 28))
        for i, bten in enumerate(btens):
            b = from_3x3_to_6x1(bten)
            b_sq = from_6x6_to_21x1(b @ b.T)
            X[i] = np.concatenate(([1], (-b.T)[0, :], (0.5 * b_sq.T)[0, :]))
        return X

    @warning_for_keywords()
    def _ols_fit(data, mask, X, *, step=int(1e4)):
        """Estimate the model parameters using ordinary least squares.

        Parameters
        ----------
        data : numpy.ndarray
            Array of shape (..., number of acquisitions).
        mask : numpy.ndarray
            Boolean array with the same shape as the data array of a single
            acquisition.
        X : numpy.ndarray
            Design matrix of shape (number of acquisitions, 28).
        step : int, optional
            The number of voxels over which the fit is calculated simultaneously.

        Returns
        -------
        params : numpy.ndarray
            Array of shape (..., 28) containing the estimated model parameters.
            Element 0 is the natural logarithm of the estimated signal without
            diffusion-weighting, elements 1-6 are the estimated diffusion tensor
            elements in Voigt notation, and elements 7-27 are the estimated
            covariance tensor elements in Voigt notation.
        """
        params = np.zeros((np.prod(mask.shape), 28)) * np.nan
        data_masked = data[mask]
        size = len(data_masked)
        X_inv = np.linalg.pinv(X.T @ X)  # Independent of data
        if step >= size:  # Fit over all data simultaneously
            S = np.log(data_masked)[..., np.newaxis]
            params_masked = (X_inv @ X.T @ S)[..., 0]
        else:  # Iterate over data
            params_masked = np.zeros((size, 28))
            for i in range(0, size, step):
                S = np.log(data_masked[i : i + step])[..., np.newaxis]
                params_masked[i : i + step] = (X_inv @ X.T @ S)[..., 0]
        params[np.where(mask.ravel())] = params_masked
        params = params.reshape((mask.shape + (28,)))
        return params

    @warning_for_keywords()
    def _wls_fit(data, mask, X, *, step=int(1e4)):
        """Estimate the model parameters using weighted least squares with the
        signal magnitudes as weights.

        Parameters
        ----------
        data : numpy.ndarray
            Array of shape (..., number of acquisitions).
        mask : numpy.ndarray
            Array with the same shape as the data array of a single acquisition.
        X : numpy.ndarray
            Design matrix of shape (number of acquisitions, 28).
        step : int, optional
            The number of voxels over which the fit is calculated simultaneously.

        Returns
        -------
        params : numpy.ndarray
            Array of shape (..., 28) containing the estimated model parameters.
            Element 0 is the natural logarithm of the estimated signal without
            diffusion-weighting, elements 1-6 are the estimated diffusion tensor
            elements in Voigt notation, and elements 7-27 are the estimated
            covariance tensor elements in Voigt notation.
        """
        params = np.zeros((np.prod(mask.shape), 28)) * np.nan
        data_masked = data[mask]
        size = len(data_masked)
        if step >= size:  # Fit over all data simultaneously
            S = np.log(data_masked)[..., np.newaxis]
            C = data_masked[:, np.newaxis, :]
            B = X.T * C
            A = np.linalg.pinv(B @ X)
            params_masked = (A @ B @ S)[..., 0]
        else:  # Iterate over data
            params_masked = np.zeros((size, 28))
            for i in range(0, size, step):
                S = np.log(data_masked[i : i + step])[..., np.newaxis]
                C = data_masked[i : i + step][:, np.newaxis, :]
                B = X.T * C
                A = np.linalg.pinv(B @ X)
                params_masked[i : i + step] = (A @ B @ S)[..., 0]
        params[np.where(mask.ravel())] = params_masked
        params = params.reshape((mask.shape + (28,)))
        return params

    def _sdpdc_fit(data, mask, X, cvxpy_solver):
        """Estimate the model parameters using Semidefinite Programming (SDP),
        while enforcing positivity constraints on the D and C tensors (SDPdc).

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

        Parameters
        ----------
        data : numpy.ndarray
            Array of shape (..., number of acquisitions).
        mask : numpy.ndarray
            Array with the same shape as the data array of a single acquisition.
        X : numpy.ndarray
            Design matrix of shape (number of acquisitions, 28).
        cvxpy_solver: string, required
            The name of the SDP solver to be used. Default: 'SCS'

        Returns
        -------
        params : numpy.ndarray
            Array of shape (..., 28) containing the estimated model parameters.
            Element 0 is the natural logarithm of the estimated signal without
            diffusion-weighting, elements 1-6 are the estimated diffusion tensor
            elements in Voigt notation, and elements 7-27 are the estimated
            covariance tensor elements in Voigt notation.

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

        if not have_cvxpy:
            raise ImportError("CVXPY package needed to enforce constraints")

        if cvxpy_solver not in cp.installed_solvers():
            raise ValueError("The selected solver is not available")

        params = np.zeros((np.prod(mask.shape), 28)) * np.nan
        data_masked = data[mask]
        size, nvols = data_masked.shape
        scale = np.maximum(np.max(data_masked, axis=1, keepdims=True), 1)
        data_masked = data_masked / scale
        data_masked[data_masked < 0] = 0
        log_data = np.log(data_masked)
        params_masked = np.zeros((size, 28))

        x = cp.Variable((28, 1))
        y = cp.Parameter((nvols, 1))
        A = cp.Parameter((nvols, 28))
        dc = cvxpy_1x6_to_3x3(x[1:7])
        cc = cvxpy_1x21_to_6x6(x[7:])
        constraints = [dc >> 0, cc >> 0]
        objective = cp.Minimize(cp.norm(A @ x - y))
        prob = cp.Problem(objective, constraints)
        unconstrained = cp.Problem(objective)

        for i in range(0, size, 1):
            vox_data = data_masked[i : i + 1, :].T
            vox_log_data = log_data[i : i + 1, :].T
            vox_log_data[np.isinf(vox_log_data)] = 0
            y.value = vox_data * vox_log_data
            A.value = vox_data * X

            try:
                prob.solve(solver=cvxpy_solver, verbose=False)
                m = x.value
            except Exception:
                msg = "Constrained optimization failed, attempting unconstrained"
                msg += " optimization."
                warn(msg, stacklevel=2)
                try:
                    unconstrained.solve(solver=cvxpy_solver)
                    m = x.value
                except Exception:
                    msg = "Unconstrained optimization failed,"
                    msg += " returning zero array."
                    warn(msg, stacklevel=2)
                    m = np.zeros(x.shape)

            params_masked[i : i + 1, :] = m.T

        params_masked[:, 0] += np.log(scale[:, 0])
        params[np.where(mask.ravel())] = params_masked
        params = params.reshape((mask.shape + (28,)))
        return params

    class QtiModel(ReconstModel):
        @warning_for_keywords()
        def __init__(self, gtab, *, fit_method="WLS", cvxpy_solver="SCS"):
            """Covariance tensor model of q-space trajectory imaging.

            See :footcite:t:`Westin2016` for further details about the model.

            Parameters
            ----------
            gtab : dipy.core.gradients.GradientTable
                Gradient table with b-tensors.
            fit_method : str, optional
                Must be one of the following:

                - 'OLS' for ordinary least squares :func:`qti._ols_fit`
                - 'WLS' for weighted least squares :func:`qti._wls_fit`
                - 'SDPDc' for semidefinite programming with positivity constraints
                  applied :footcite:p:`Herberthson2021` :func:`qti._sdpdc_fit`

            cvxpy_solver: str, optionals
                solver for the SDP formulation. default: 'SCS'

            References
            ----------
            .. footbibliography::
            """
            ReconstModel.__init__(self, gtab)

            if self.gtab.btens is None:
                raise ValueError(
                    "QTI requires b-tensors to be defined in the gradient table."
                )
            self.X = design_matrix(self.gtab.btens)
            rank = np.linalg.matrix_rank(self.X.T @ self.X)
            if rank < 28:
                warn(
                    "The combination of the b-tensor shapes, sizes, and "
                    + "orientations does not enable all elements of the covariance "
                    + f"tensor to be estimated (rank(X.T @ X) = {rank} < 28).",
                    stacklevel=2,
                )

            try:
                self.fit_method = common_fit_methods[fit_method]
            except KeyError as e:
                raise ValueError(
                    f"Invalid value ({fit_method}) for 'fit_method'."
                    + " Options: 'OLS', 'WLS', 'SDPdc'."
                ) from e

            self.cvxpy_solver = cvxpy_solver
            self.fit_method_name = fit_method

        @warning_for_keywords()
        def fit(self, data, *, mask=None):
            """Fit QTI to data.

            Parameters
            ----------
            data : numpy.ndarray
                Array of shape (..., number of acquisitions).
            mask : numpy.ndarray, optional
                Array with the same shape as the data array of a single acquisition.

            Returns
            -------
            qtifit : dipy.reconst.qti.QtiFit
                The fitted model.
            """
            if mask is None:
                mask = np.ones(data.shape[:-1], dtype=bool)
            else:
                if mask.shape != data.shape[:-1]:
                    raise ValueError("Mask is not the same shape as data.")
                mask = np.asarray(mask, dtype=bool)
            if self.fit_method_name == "SDPdc":
                params = self.fit_method(data, mask, self.X, self.cvxpy_solver)
            else:
                params = self.fit_method(data, mask, self.X)
            return QtiFit(params)

        def predict(self, params):
            """Generate signals from this model class instance and given parameters.

            Parameters
            ----------
            params : numpy.ndarray
                Array of shape (..., 28) containing the model parameters. Element 0
                is the natural logarithm of the signal without diffusion-weighting,
                elements 1-6 are the diffusion tensor elements in Voigt notation,
                and elements 7-27 are the covariance tensor elements in Voigt
                notation.

            Returns
            -------
            S : numpy.ndarray
                Signals.
            """
            S0 = np.exp(params[..., 0])
            D = params[..., 1:7, np.newaxis]
            C = params[..., 7::, np.newaxis]
            S = qti_signal(self.gtab, D, C, S0=S0)
            return S

    class QtiFit:
        def __init__(self, params):
            """Fitted QTI model.

            Parameters
            ----------
            params : numpy.ndarray
                Array of shape (..., 28) containing the model parameters. Element 0
                is the natural logarithm of the signal without diffusion-weighting,
                elements 1-6 are the diffusion tensor elements in Voigt notation,
                and elements 7-27 are the covariance tensor elements in Voigt
                notation.
            """
            self.params = params

        def predict(self, gtab):
            """Generate signals from this model fit and a given gradient table.

            Parameters
            ----------
            gtab : dipy.core.gradients.GradientTable
                Gradient table with b-tensors.

            Returns
            -------
            S : numpy.ndarray
                Signals.
            """
            if gtab.btens is None:
                raise ValueError(
                    "QTI requires b-tensors to be defined in the gradient table."
                )
            S0 = self.S0_hat
            D = self.params[..., 1:7, np.newaxis]
            C = self.params[..., 7::, np.newaxis]
            S = qti_signal(gtab, D, C, S0=S0)
            return S

        @auto_attr
        def S0_hat(self):
            """Estimated signal without diffusion-weighting.

            Returns
            -------
            S0 : numpy.ndarray
            """
            S0 = np.exp(self.params[..., 0])
            return S0

        @auto_attr
        def md(self):
            """Mean diffusivity.

            Returns
            -------
            md : numpy.ndarray

            Notes
            -----
            Mean diffusivity is calculated as

                .. math::

                    \\text{MD} = \\langle \\mathbf{D} \\rangle :
                    \\mathbf{E}_\\text{iso}
            """
            md = np.matmul(self.params[..., np.newaxis, 1:7], from_3x3_to_6x1(e_iso))[
                ..., 0, 0
            ]
            return md

        @auto_attr
        def v_md(self):
            """Variance of microscopic mean diffusivities.

            Returns
            -------
            v_md : numpy.ndarray

            Notes
            -----
            Variance of microscopic mean diffusivities is calculated as

                .. math::

                    V_\\text{MD} = \\mathbb{C} : \\mathbb{E}_\\text{bulk}
            """
            v_md = np.matmul(self.params[..., np.newaxis, 7::], from_6x6_to_21x1(E_bulk))[
                ..., 0, 0
            ]
            return v_md

        @auto_attr
        def v_shear(self):
            """Shear variance.

            Returns
            -------
            v_shear : numpy.ndarray

            Notes
            -----
            Shear variance is calculated as

                .. math::

                    V_\\text{shear} = \\mathbb{C} : \\mathbb{E}_\\text{shear}
            """
            v_shear = np.matmul(
                self.params[..., np.newaxis, 7::], from_6x6_to_21x1(E_shear)
            )[..., 0, 0]
            return v_shear

        @auto_attr
        def v_iso(self):
            """Total isotropic variance.

            Returns
            -------
            v_iso : numpy.ndarray

            Notes
            -----
            Total isotropic variance is calculated as

                .. math::

                    V_\\text{iso} = \\mathbb{C} : \\mathbb{E}_\\text{iso}
            """
            v_iso = np.matmul(self.params[..., np.newaxis, 7::], from_6x6_to_21x1(E_iso))[
                ..., 0, 0
            ]
            return v_iso

        @auto_attr
        def d_sq(self):
            """Diffusion tensor's outer product with itself.

            Returns
            -------
            d_sq : numpy.ndarray
            """
            d_sq = np.matmul(
                self.params[..., 1:7, np.newaxis], self.params[..., np.newaxis, 1:7]
            )
            return d_sq

        @auto_attr
        def mean_d_sq(self):
            """Average of microscopic diffusion tensors' outer products with
            themselves.

            Returns
            -------
            mean_d_sq : numpy.ndarray

            Notes
            -----
            Average of microscopic diffusion tensors' outer products with
            themselves is calculated as

                .. math::

                    \\langle \\mathbf{D} \\otimes \\mathbf{D} \\rangle =
                    \\mathbb{C} +
                    \\langle \\mathbf{D} \\rangle \\otimes \\langle \\mathbf{D}
                    \\rangle
            """
            mean_d_sq = from_21x1_to_6x6(self.params[..., 7::, np.newaxis]) + self.d_sq
            return mean_d_sq

        @auto_attr
        def c_md(self):
            """Normalized variance of mean diffusivities.

            Returns
            -------
            c_md : numpy.ndarray

            Notes
            -----
            Normalized variance of microscopic mean diffusivities is calculated as

                .. math::

                    C_\\text{MD} = \\frac{\\mathbb{C} : \\mathbb{E}_\\text{bulk}}
                    {\\langle \\mathbf{D} \\otimes \\mathbf{D} \\rangle :
                    \\mathbb{E}_\\text{bulk}}
            """
            c_md = (
                self.v_md
                / np.matmul(
                    np.swapaxes(from_6x6_to_21x1(self.mean_d_sq), -1, -2),
                    from_6x6_to_21x1(E_bulk),
                )[..., 0, 0]
            )
            return c_md

        @auto_attr
        def c_mu(self):
            """Normalized microscopic anisotropy.

            Returns
            -------
            c_mu : numpy.ndarray

            Notes
            -----
            Normalized microscopic anisotropy is calculated as

                .. math::

                    C_\\mu = \\frac{3}{2} \\frac{\\langle \\mathbf{D} \\otimes
                    \\mathbf{D}
                    \\rangle : \\mathbb{E}_\\text{shear}}{\\langle \\mathbf{D}
                    \\otimes
                    \\mathbf{D} \\rangle : \\mathbb{E}_\\text{iso}}
            """
            c_mu = (
                1.5
                * np.matmul(
                    np.swapaxes(from_6x6_to_21x1(self.mean_d_sq), -1, -2),
                    from_6x6_to_21x1(E_shear),
                )
                / np.matmul(
                    np.swapaxes(from_6x6_to_21x1(self.mean_d_sq), -1, -2),
                    from_6x6_to_21x1(E_iso),
                )
            )[..., 0, 0]
            return c_mu

        @auto_attr
        def ufa(self):
            """Microscopic fractional anisotropy.

            Returns
            -------
            ufa : numpy.ndarray

            Notes
            -----
            Microscopic fractional anisotropy is calculated as

                .. math::

                    \\mu\\text{FA} = \\sqrt{C_\\mu}
            """
            ufa = np.sqrt(self.c_mu)
            return ufa

        @auto_attr
        def c_m(self):
            """Normalized macroscopic anisotropy.

            Returns
            -------
            c_m : numpy.ndarray

            Notes
            -----
            Normalized macroscopic anisotropy is calculated as

                .. math::

                    C_\\text{M} = \\frac{3}{2} \\frac{\\langle \\mathbf{D} \\rangle
                    \\otimes \\langle \\mathbf{D} \\rangle :
                    \\mathbb{E}_\\text{shear}}
                    {\\langle \\mathbf{D} \\rangle \\otimes \\langle \\mathbf{D}
                    \\rangle :
                    \\mathbb{E}_\\text{iso}}
            """
            c_m = (
                1.5
                * np.matmul(
                    np.swapaxes(from_6x6_to_21x1(self.d_sq), -1, -2),
                    from_6x6_to_21x1(E_shear),
                )
                / np.matmul(
                    np.swapaxes(from_6x6_to_21x1(self.d_sq), -1, -2),
                    from_6x6_to_21x1(E_iso),
                )
            )[..., 0, 0]
            return c_m

        @auto_attr
        def fa(self):
            """Fractional anisotropy.

            Returns
            -------
            fa : numpy.ndarray

            Notes
            -----
            Fractional anisotropy is calculated as

                .. math::

                    \\text{FA} = \\sqrt{C_\\text{M}}
            """
            fa = np.sqrt(self.c_m)
            return fa

        @auto_attr
        def c_c(self):
            """Microscopic orientation coherence.

            Returns
            -------
            c_c : numpy.ndarray

            Notes
            -----
            Microscopic orientation coherence is calculated as

                .. math::

                    C_c = \\frac{C_\\text{M}}{C_\\mu}
            """
            c_c = self.c_m / self.c_mu
            return c_c

        @auto_attr
        def mk(self):
            """Mean kurtosis.

            Returns
            -------
            mk : numpy.ndarray

            Notes
            -----
            Mean kurtosis is calculated as

                .. math::

                    \\text{MK} = K_\\text{bulk} + K_\\text{shear}
            """
            mk = self.k_bulk + self.k_shear
            return mk

        @auto_attr
        def k_bulk(self):
            """Bulk kurtosis.

            Returns
            -------
            k_bulk : numpy.ndarray

            Notes
            -----
            Bulk kurtosis is calculated as

                .. math::

                    K_\\text{bulk} = 3 \\frac{\\mathbb{C} :
                    \\mathbb{E}_\\text{bulk}}
                    {\\langle \\mathbf{D} \\rangle \\otimes \\langle \\mathbf{D}
                    \\rangle : \\mathbb{E}_\\text{bulk}}
            """
            k_bulk = (
                3
                * np.matmul(self.params[..., np.newaxis, 7::], from_6x6_to_21x1(E_bulk))
                / np.matmul(
                    np.swapaxes(from_6x6_to_21x1(self.d_sq), -1, -2),
                    from_6x6_to_21x1(E_bulk),
                )
            )[..., 0, 0]
            return k_bulk

        @auto_attr
        def k_shear(self):
            """Shear kurtosis.

            Returns
            -------
            k_shear : numpy.ndarray

            Notes
            -----
            Shear kurtosis is calculated as

                .. math::

                    K_\\text{shear} = \\frac{6}{5} \\frac{\\mathbb{C} :
                    \\mathbb{E}_\\text{shear}}{\\langle \\mathbf{D} \\rangle
                    \\otimes
                    \\langle \\mathbf{D} \\rangle : \\mathbb{E}_\\text{bulk}}
            """
            k_shear = (
                6
                / 5
                * np.matmul(self.params[..., np.newaxis, 7::], from_6x6_to_21x1(E_shear))
                / np.matmul(
                    np.swapaxes(from_6x6_to_21x1(self.d_sq), -1, -2),
                    from_6x6_to_21x1(E_bulk),
                )
            )[..., 0, 0]
            return k_shear

        @auto_attr
        def k_mu(self):
            """Microscopic kurtosis.

            Returns
            -------
            k_mu : numpy.ndarray

            Notes
            -----
            Microscopic kurtosis is calculated as

                .. math::

                    K_\\mu = \\frac{6}{5} \\frac{\\langle \\mathbf{D} \\otimes
                    \\mathbf{D}
                    \\rangle : \\mathbb{E}_\\text{shear}}{\\langle \\mathbf{D}
                    \\rangle
                    \\otimes \\langle \\mathbf{D} \\rangle :
                    \\mathbb{E}_\\text{bulk}}
            """
            k_mu = (
                6
                / 5
                * np.matmul(
                    np.swapaxes(from_6x6_to_21x1(self.mean_d_sq), -1, -2),
                    from_6x6_to_21x1(E_shear),
                )
                / np.matmul(
                    np.swapaxes(from_6x6_to_21x1(self.d_sq), -1, -2),
                    from_6x6_to_21x1(E_bulk),
                )
            )[..., 0, 0]
            return k_mu

    common_fit_methods = {"OLS": _ols_fit, "WLS": _wls_fit, "SDPdc": _sdpdc_fit}