# Source code for dipy.sims.phantom

    import numpy as np

    from dipy.core.geometry import vec2vec_rotmat
    from dipy.core.gradients import gradient_table
    from dipy.data import get_fnames
    import dipy.sims.voxel as vox

    # import scipy.stats as stats
    from dipy.sims.voxel import diffusion_evals, single_tensor
    from dipy.testing.decorators import warning_for_keywords

    @warning_for_keywords()
    def add_noise(vol, *, snr=1.0, S0=None, noise_type="rician", rng=None):
        """Add noise of specified distribution to a 4D array.

        Parameters
        ----------
        vol : array, shape (X,Y,Z,W)
            Diffusion measurements in `W` directions at each ``(X, Y, Z)`` voxel
            position.
        snr : float, optional
            The desired signal-to-noise ratio.  (See notes below.)
        S0 : float, optional
            Reference signal for specifying `snr`.
        noise_type : string, optional
            The distribution of noise added. Can be either 'gaussian' for Gaussian
            distributed noise, 'rician' for Rice-distributed noise (default) or
            'rayleigh' for a Rayleigh distribution.
        rng : numpy.random.Generator class, optional
            Numpy's random generator for setting seed values when needed.

        Returns
        -------
        vol : array, same shape as vol
            Volume with added noise.

        Notes
        -----
        SNR is defined here, following :footcite:p:`Descoteaux2007`, as
        ``S0 / sigma``, where ``sigma`` is the standard deviation of the two
        Gaussian distributions forming the real and imaginary components of the
        Rician noise distribution (see :footcite:p:`Gudbjartsson1995`).

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

        Examples
        --------
        >>> signal = np.arange(800).reshape(2, 2, 2, 100)
        >>> signal_w_noise = add_noise(signal, snr=10, noise_type='rician')

        """
        orig_shape = vol.shape
        vol_flat = np.reshape(vol.copy(), (-1, vol.shape[-1]))

        if S0 is None:
            S0 = np.max(vol)

        for vox_idx, signal in enumerate(vol_flat):
            vol_flat[vox_idx] = vox.add_noise(
                signal, snr=snr, S0=S0, noise_type=noise_type, rng=rng
            )

        return np.reshape(vol_flat, orig_shape)

    def diff2eigenvectors(dx, dy, dz):
        """numerical derivatives 2 eigenvectors"""
        basis = np.eye(3)
        u = np.array([dx, dy, dz])
        u = u / np.linalg.norm(u)
        R = vec2vec_rotmat(basis[:, 0], u)
        eig0 = u
        eig1 = np.dot(R, basis[:, 1])
        eig2 = np.dot(R, basis[:, 2])
        eigs = np.zeros((3, 3))
        eigs[:, 0] = eig0
        eigs[:, 1] = eig1
        eigs[:, 2] = eig2
        return eigs, R

    @warning_for_keywords()
    def orbital_phantom(
        *,
        gtab=None,
        evals=diffusion_evals,
        func=None,
        t=None,
        datashape=(64, 64, 64, 65),
        origin=(32, 32, 32),
        scale=(25, 25, 25),
        angles=None,
        radii=None,
        S0=100.0,
        snr=None,
        rng=None,
    ):
        """Create a phantom based on a 3-D orbit ``f(t) -> (x,y,z)``.

        Parameters
        ----------
        gtab : GradientTable
            Gradient table of measurement directions.
        evals : array, shape (3,)
            Tensor eigenvalues.
        func : user defined function f(t)->(x,y,z)
            It could be desirable for ``-1=<x,y,z <=1``.
            If None creates a circular orbit.
        t : array, shape (K,)
            Represents time for the orbit. Default is
            ``np.linspace(0, 2 * np.pi, 1000)``.
        datashape : array, shape (X,Y,Z,W)
            Size of the output simulated data
        origin : tuple, shape (3,)
            Define the center for the volume
        scale : tuple, shape (3,)
            Scale the function before applying to the grid
        angles : array, shape (L,)
            Density angle points, always perpendicular to the first eigen vector
            Default np.linspace(0, 2 * np.pi, 32).
        radii : array, shape (M,)
            Thickness radii.  Default ``np.linspace(0.2, 2, 6)``.
            angles and radii define the total thickness options
        S0 : double, optional
            Maximum simulated signal. Default 100.
        snr : float, optional
            The signal to noise ratio set to apply Rician noise to the data.
            Default is to not add noise at all.
        rng : numpy.random.Generator class, optional
            Numpy's random generator for setting seed values when needed.
            Default is None.

        Returns
        -------
        data : array, shape (datashape)

        See Also
        --------
        add_noise

        Examples
        --------

        >>> def f(t):
        ...    x = np.sin(t)
        ...    y = np.cos(t)
        ...    z = np.linspace(-1, 1, len(x))
        ...    return x, y, z

        >>> data = orbital_phantom(func=f)

        """

        if t is None:
            t = np.linspace(0, 2 * np.pi, 1000)

        if angles is None:
            angles = np.linspace(0, 2 * np.pi, 32)

        if radii is None:
            radii = np.linspace(0.2, 2, 6)

        if gtab is None:
            fimg, fbvals, fbvecs = get_fnames(name="small_64D")
            gtab = gradient_table(fbvals, bvecs=fbvecs)

        if func is None:
            x = np.sin(t)
            y = np.cos(t)
            z = np.zeros(t.shape)
        else:
            x, y, z = func(t)

        dx = np.diff(x)
        dy = np.diff(y)
        dz = np.diff(z)

        x = scale[0] * x + origin[0]
        y = scale[1] * y + origin[1]
        z = scale[2] * z + origin[2]

        bx = np.zeros(len(angles))
        by = np.sin(angles)
        bz = np.cos(angles)

        # The entire volume is considered to be inside the brain.
        # Voxels without a fiber crossing through them are taken
        # to be isotropic with signal = S0.
        vol = np.zeros(datashape) + S0

        for i in range(len(dx)):
            evecs, R = diff2eigenvectors(dx[i], dy[i], dz[i])
            S = single_tensor(gtab, S0, evals=evals, evecs=evecs, snr=None)

            vol[int(x[i]), int(y[i]), int(z[i]), :] += S

            for r in radii:
                for j in range(len(angles)):
                    rb = np.dot(R, np.array([bx[j], by[j], bz[j]]))

                    ix = int(x[i] + r * rb[0])
                    iy = int(y[i] + r * rb[1])
                    iz = int(z[i] + r * rb[2])
                    vol[ix, iy, iz] = vol[ix, iy, iz] + S

        vol = vol / np.max(vol, axis=-1)[..., np.newaxis]
        vol *= S0

        if snr is not None:
            vol = add_noise(vol, snr=snr, S0=S0, noise_type="rician", rng=rng)

        return vol

    if __name__ == "__main__":
        # TODO: this can become a nice tutorial for generating phantoms

        def f(t):
            x = np.sin(t)
            y = np.cos(t)
            # z=np.zeros(t.shape)
            z = np.linspace(-1, 1, len(x))
            return x, y, z

        # helix
        vol = orbital_phantom(func=f)

        def f2(t):
            x = np.linspace(-1, 1, len(t))
            y = np.linspace(-1, 1, len(t))
            z = np.zeros(x.shape)
            return x, y, z

        # first direction
        vol2 = orbital_phantom(func=f2)

        def f3(t):
            x = np.linspace(-1, 1, len(t))
            y = -np.linspace(-1, 1, len(t))
            z = np.zeros(x.shape)
            return x, y, z

        # second direction
        vol3 = orbital_phantom(func=f3)
        # double crossing
        vol23 = vol2 + vol3

        # """
        def f4(t):
            x = np.zeros(t.shape)
            y = np.zeros(t.shape)
            z = np.linspace(-1, 1, len(t))
            return x, y, z

        # triple crossing
        vol4 = orbital_phantom(func=f4)
        vol234 = vol23 + vol4

        # unknown function
        # voln = add_rician_noise(vol234)

        # """

        # from dipy.viz import window, actor
        # scene = window.Scene()
        # scene.add(actor.volume(vol234[...,0]))
        # window.show(scene)
        # vol234n=add_rician_noise(vol234,20)