# Source code for dipy.reconst.forecast

    from warnings import warn

    import numpy as np
    from scipy.optimize import leastsq
    from scipy.special import gamma, hyp1f1

    from dipy.core.geometry import cart2sphere
    from dipy.data import default_sphere
    from dipy.reconst.cache import Cache
    from dipy.reconst.csdeconv import csdeconv
    from dipy.reconst.multi_voxel import multi_voxel_fit
    from dipy.reconst.odf import OdfFit, OdfModel
    from dipy.reconst.shm import real_sh_descoteaux_from_index
    from dipy.testing.decorators import warning_for_keywords
    from dipy.utils.deprecator import deprecated_params
    from dipy.utils.optpkg import optional_package

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

    class ForecastModel(OdfModel, Cache):
        r"""Fiber ORientation Estimated using Continuous Axially Symmetric Tensors
        (FORECAST).

        FORECAST :footcite:p:`Anderson2005`, :footcite:p:`Kaden2016a`,
        :footcite:p:`Zucchelli2017` is a Spherical Deconvolution reconstruction
        model for multi-shell diffusion data which enables the calculation of a
        voxel adaptive response function using the Spherical Mean Technique (SMT)
        :footcite:p:`Kaden2016a`, :footcite:p:`Zucchelli2017`.

        With FORECAST it is possible to calculate crossing invariant parallel
        diffusivity, perpendicular diffusivity, mean diffusivity, and fractional
        anisotropy :footcite:p:`Kaden2016a`.

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

        Notes
        -----
        The implementation of FORECAST may require CVXPY (https://www.cvxpy.org/).
        """

        @deprecated_params("sh_order", new_name="sh_order_max", since="1.9", until="2.0")
        @warning_for_keywords()
        def __init__(
            self,
            gtab,
            *,
            sh_order_max=8,
            lambda_lb=1e-3,
            dec_alg="CSD",
            sphere=None,
            lambda_csd=1.0,
        ):
            r"""Analytical and continuous modeling of the diffusion signal with
            respect to the FORECAST basis.

            This implementation is a modification of the original FORECAST
            model presented in :footcite:p:`Anderson2005` adapted for multi-shell
            data as in :footcite:p:`Kaden2016a`, :footcite:p:`Zucchelli2017`.

            The main idea is to model the diffusion signal as the combination of a
            single fiber response function $F(\mathbf{b})$ times the fODF
            $\rho(\mathbf{v})$

            .. math::

                E(\mathbf{b}) = \int_{\mathbf{v} \in \mathcal{S}^2} \rho(\mathbf{v}) F({\mathbf{b}} | \mathbf{v}) d \mathbf{v}

            where $\mathbf{b}$ is the b-vector (b-value times gradient direction)
            and $\mathbf{v}$ is a unit vector representing a fiber direction.

            In FORECAST $\rho$ is modeled using real symmetric Spherical Harmonics
            (SH) and $F(\mathbf(b))$ is an axially symmetric tensor.

            Parameters
            ----------
            gtab : GradientTable,
                gradient directions and bvalues container class.
            sh_order_max : unsigned int, optional
                an even integer that represent the maximal SH order ($l$) of the
                basis (max 12)
            lambda_lb: float, optional
                Laplace-Beltrami regularization weight.
            dec_alg : str, optional
                Spherical deconvolution algorithm. The possible values are Weighted Least Squares ('WLS'),
                Positivity Constraints using CVXPY ('POS') and the Constraint
                Spherical Deconvolution algorithm ('CSD'). Default is 'CSD'.
            sphere : array, shape (N,3), optional
                sphere points where to enforce positivity when 'POS' or 'CSD'
                dec_alg are selected.
            lambda_csd : float, optional
                CSD regularization weight.

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

            Examples
            --------
            In this example, where the data, gradient table and sphere tessellation
            used for reconstruction are provided, we model the diffusion signal
            with respect to the FORECAST and compute the fODF, parallel and
            perpendicular diffusivity.

            >>> import warnings
            >>> from dipy.data import default_sphere, get_3shell_gtab
            >>> gtab = get_3shell_gtab()
            >>> from dipy.sims.voxel import multi_tensor
            >>> mevals = np.array(([0.0017, 0.0003, 0.0003],
            ...                    [0.0017, 0.0003, 0.0003]))
            >>> angl = [(0, 0), (60, 0)]
            >>> data, sticks = multi_tensor(gtab,
            ...                             mevals,
            ...                             S0=100.0,
            ...                             angles=angl,
            ...                             fractions=[50, 50],
            ...                             snr=None)
            >>> from dipy.reconst.forecast import ForecastModel
            >>> from dipy.reconst.shm import descoteaux07_legacy_msg
            >>> with warnings.catch_warnings():
            ...     warnings.filterwarnings(
            ...         "ignore", message=descoteaux07_legacy_msg,
            ...         category=PendingDeprecationWarning)
            ...     fm = ForecastModel(gtab, sh_order_max=6)
            >>> f_fit = fm.fit(data)
            >>> d_par = f_fit.dpar
            >>> d_perp = f_fit.dperp
            >>> with warnings.catch_warnings():
            ...     warnings.filterwarnings(
            ...         "ignore", message=descoteaux07_legacy_msg,
            ...         category=PendingDeprecationWarning)
            ...     fodf = f_fit.odf(default_sphere)
            """  # noqa: E501
            OdfModel.__init__(self, gtab)

            # round the bvals in order to avoid numerical errors
            self.bvals = np.round(gtab.bvals / 100) * 100
            self.bvecs = gtab.bvecs

            if 0 <= sh_order_max <= 12 and not bool(sh_order_max % 2):
                self.sh_order_max = sh_order_max
            else:
                msg = "sh_order_max must be a non-zero even positive number "
                msg += "between 2 and 12"
                raise ValueError(msg)

            if sphere is None:
                sphere = default_sphere
                self.vertices = sphere.vertices[0 : int(sphere.vertices.shape[0] / 2), :]

            else:
                self.vertices = sphere

            self.b0s_mask = self.bvals == 0
            self.one_0_bvals = np.r_[0, self.bvals[~self.b0s_mask]]
            self.one_0_bvecs = np.r_[
                np.array([0, 0, 0]).reshape(1, 3), self.bvecs[~self.b0s_mask, :]
            ]

            self.rho = rho_matrix(self.sh_order_max, self.one_0_bvecs)

            # signal regularization matrix
            self.srm = rho_matrix(4, self.one_0_bvecs)
            self.lb_matrix_signal = lb_forecast(4)

            self.b_unique = np.sort(np.unique(self.bvals[self.bvals > 0]))
            self.wls = True
            self.csd = False
            self.pos = False

            if dec_alg.upper() == "POS":
                if not have_cvxpy:
                    cvxpy.import_error()

                self.wls = False
                self.pos = True

            if dec_alg.upper() == "CSD":
                self.csd = True

            self.lb_matrix = lb_forecast(self.sh_order_max)
            self.lambda_lb = lambda_lb
            self.lambda_csd = lambda_csd
            self.fod = rho_matrix(sh_order_max, self.vertices)

        @multi_voxel_fit
        def fit(self, data, **kwargs):
            data_b0 = data[self.b0s_mask].mean()
            data_single_b0 = np.r_[data_b0, data[~self.b0s_mask]]
            data_single_b0 = np.divide(
                data_single_b0,
                data_b0,
                out=np.zeros_like(data_single_b0),
                where=data_b0 != 0,
            )

            # calculates the mean signal at each b_values
            means = find_signal_means(
                self.b_unique,
                data_single_b0,
                self.one_0_bvals,
                self.srm,
                self.lb_matrix_signal,
            )

            # average diffusivity initialization
            x = np.array([np.pi / 4, np.pi / 4])

            x, status = leastsq(forecast_error_func, x, args=(self.b_unique, means))

            # transform to bound the diffusivities from 0 to 3e-03
            d_par = np.cos(x[0]) ** 2 * 3e-03
            d_perp = np.cos(x[1]) ** 2 * 3e-03

            if d_perp >= d_par:
                d_par, d_perp = d_perp, d_par

            # round to avoid memory explosion
            diff_key = str(int(np.round(d_par * 1e05))) + str(int(np.round(d_perp * 1e05)))

            M_diff = self.cache_get("forecast_matrix", key=diff_key)
            if M_diff is None:
                M_diff = forecast_matrix(self.sh_order_max, d_par, d_perp, self.one_0_bvals)
                self.cache_set("forecast_matrix", key=diff_key, value=M_diff)

            M = M_diff * self.rho
            M0 = M[:, 0]
            c0 = np.sqrt(1.0 / (4 * np.pi))

            # coefficients vector initialization
            n_c = int((self.sh_order_max + 1) * (self.sh_order_max + 2) / 2)
            coef = np.zeros(n_c)
            coef[0] = c0
            if int(np.round(d_par * 1e05)) > int(np.round(d_perp * 1e05)):
                if self.wls:
                    data_r = data_single_b0 - M0 * c0

                    Mr = M[:, 1:]
                    Lr = self.lb_matrix[1:, 1:]

                    pseudo_inv = np.dot(
                        np.linalg.inv(np.dot(Mr.T, Mr) + self.lambda_lb * Lr), Mr.T
                    )
                    coef = np.dot(pseudo_inv, data_r)
                    coef = np.r_[c0, coef]

                if self.csd:
                    coef, _ = csdeconv(data_single_b0, M, self.fod, tau=0.1, convergence=50)
                    coef = coef / coef[0] * c0

                if self.pos:
                    c = cvxpy.Variable(M.shape[1])
                    design_matrix = cvxpy.Constant(M) @ c
                    objective = cvxpy.Minimize(
                        cvxpy.sum_squares(design_matrix - data_single_b0)
                        + self.lambda_lb * cvxpy.quad_form(c, self.lb_matrix)
                    )

                    constraints = [c[0] == c0, self.fod @ c >= 0]
                    prob = cvxpy.Problem(objective, constraints)
                    try:
                        prob.solve(solver=cvxpy.OSQP, eps_abs=1e-05, eps_rel=1e-05)
                        coef = np.asarray(c.value).squeeze()
                    except Exception:
                        warn("Optimization did not find a solution", stacklevel=2)
                        coef = np.zeros(M.shape[1])
                        coef[0] = c0

            return ForecastFit(self, data, coef, d_par, d_perp)

    class ForecastFit(OdfFit):
        def __init__(self, model, data, sh_coef, d_par, d_perp):
            """Calculates diffusion properties for a single voxel

            Parameters
            ----------
            model : object,
                AnalyticalModel
            data : 1d ndarray,
                fitted data
            sh_coef : 1d ndarray,
                forecast sh coefficients
            d_par : float,
                parallel diffusivity
            d_perp : float,
                perpendicular diffusivity
            """
            OdfFit.__init__(self, model, data)
            self.model = model
            self._sh_coef = sh_coef
            self.gtab = model.gtab
            self.sh_order_max = model.sh_order_max
            self.d_par = d_par
            self.d_perp = d_perp

            self.rho = None

        @warning_for_keywords()
        def odf(self, sphere, *, clip_negative=True):
            r"""Calculates the fODF for a given discrete sphere.

            Parameters
            ----------
            sphere : Sphere,
                the odf sphere
            clip_negative : boolean, optional
                if True clip the negative odf values to 0, default True
            """
            if self.rho is None:
                self.rho = rho_matrix(self.sh_order_max, sphere.vertices)

            odf = np.dot(self.rho, self._sh_coef)

            if clip_negative:
                odf = np.clip(odf, 0, odf.max())

            return odf

        def fractional_anisotropy(self):
            r"""Calculates the fractional anisotropy."""
            fa = np.sqrt(
                0.5
                * (2 * (self.d_par - self.d_perp) ** 2)
                / (self.d_par**2 + 2 * self.d_perp**2)
            )
            return fa

        def mean_diffusivity(self):
            r"""Calculates the mean diffusivity."""
            md = (self.d_par + 2 * self.d_perp) / 3.0
            return md

        @warning_for_keywords()
        def predict(self, *, gtab=None, S0=1.0):
            r"""Calculates the fODF for a given discrete sphere.

            Parameters
            ----------
            gtab : GradientTable, optional
                gradient directions and bvalues container class.
            S0 : float, optional
                the signal at b-value=0

            """
            if gtab is None:
                gtab = self.gtab

            M_diff = forecast_matrix(self.sh_order_max, self.d_par, self.d_perp, gtab.bvals)

            rho = rho_matrix(self.sh_order_max, gtab.bvecs)
            M = M_diff * rho
            S = S0 * np.dot(M, self._sh_coef)

            return S

        @property
        def sh_coeff(self):
            """The FORECAST SH coefficients"""
            return self._sh_coef

        @property
        def dpar(self):
            """The parallel diffusivity"""
            return self.d_par

        @property
        def dperp(self):
            """The perpendicular diffusivity"""
            return self.d_perp

    @warning_for_keywords()
    def find_signal_means(b_unique, data_norm, bvals, rho, lb_matrix, *, w=1e-03):
        r"""Calculate the mean signal for each shell.

        Parameters
        ----------
        b_unique : 1d ndarray,
            unique b-values in a vector excluding zero
        data_norm : 1d ndarray,
            normalized diffusion signal
        bvals : 1d ndarray,
            the b-values
        rho : 2d ndarray,
            SH basis matrix for fitting the signal on each shell
        lb_matrix : 2d ndarray,
            Laplace-Beltrami regularization matrix
        w : float,
            weight for the Laplace-Beltrami regularization

        Returns
        -------
        means : 1d ndarray
            the average of the signal for each b-values

        """
        lb = len(b_unique)
        means = np.zeros(lb)
        for u in range(lb):
            ind = bvals == b_unique[u]
            shell = data_norm[ind]
            if np.sum(ind) > 20:
                M = rho[ind, :]
                coef = np.linalg.multi_dot(
                    [np.linalg.inv(np.dot(M.T, M) + w * lb_matrix), M.T, shell]
                )

                means[u] = coef[0] / np.sqrt(4 * np.pi)
            else:
                means[u] = shell.mean()

        return means

    def forecast_error_func(x, b_unique, E):
        r"""Calculates the difference between the mean signal calculated using
        the parameter vector x and the average signal E using FORECAST and SMT
        """
        d_par = np.cos(x[0]) ** 2 * 3e-03
        d_perp = np.cos(x[1]) ** 2 * 3e-03

        if d_perp >= d_par:
            d_par, d_perp = d_perp, d_par

        E_reconst = (
            0.5 * np.exp(-b_unique * d_perp) * psi_l(0, (b_unique * (d_par - d_perp)))
        )

        v = E - E_reconst
        return v

    def psi_l(ell, b):
        n = ell // 2
        v = (-b) ** n
        v *= gamma(n + 1.0 / 2) / gamma(2 * n + 3.0 / 2)
        v *= hyp1f1(n + 1.0 / 2, 2 * n + 3.0 / 2, -b)
        return v

    @deprecated_params("sh_order", new_name="sh_order_max", since="1.9", until="2.0")
    def forecast_matrix(sh_order_max, d_par, d_perp, bvals):
        r"""Compute the FORECAST radial matrix"""
        n_c = int((sh_order_max + 1) * (sh_order_max + 2) / 2)
        M = np.zeros((bvals.shape[0], n_c))
        counter = 0
        for ell in range(0, sh_order_max + 1, 2):
            for _ in range(-ell, ell + 1):
                M[:, counter] = (
                    2
                    * np.pi
                    * np.exp(-bvals * d_perp)
                    * psi_l(ell, bvals * (d_par - d_perp))
                )
                counter += 1
        return M

    @deprecated_params("sh_order", new_name="sh_order_max", since="1.9", until="2.0")
    def rho_matrix(sh_order_max, vecs):
        r"""Compute the SH matrix $\rho$"""

        r, theta, phi = cart2sphere(vecs[:, 0], vecs[:, 1], vecs[:, 2])
        theta[np.isnan(theta)] = 0

        n_c = int((sh_order_max + 1) * (sh_order_max + 2) / 2)
        rho = np.zeros((vecs.shape[0], n_c))
        counter = 0
        for l_values in range(0, sh_order_max + 1, 2):
            for m_values in range(-l_values, l_values + 1):
                rho[:, counter] = real_sh_descoteaux_from_index(
                    m_values, l_values, theta, phi
                )
                counter += 1
        return rho

    @deprecated_params("sh_order", new_name="sh_order_max", since="1.9", until="2.0")
    def lb_forecast(sh_order_max):
        r"""Returns the Laplace-Beltrami regularization matrix for FORECAST"""
        n_c = int((sh_order_max + 1) * (sh_order_max + 2) / 2)
        diag_lb = np.zeros(n_c)
        counter = 0
        for j in range(0, sh_order_max + 1, 2):
            stop = 2 * j + 1 + counter
            diag_lb[counter:stop] = (j * (j + 1)) ** 2
            counter = stop

        return np.diag(diag_lb)