"""
==============================================
Reconstruct with Constant Solid Angle (Q-Ball)
==============================================

We show how to apply a Constant Solid Angle ODF (Q-Ball) model from Aganj et
al. [Aganj2010]_ to your datasets.

First import the necessary modules:
"""

import numpy as np
from dipy.core.gradients import gradient_table
from dipy.data import get_fnames, default_sphere
from dipy.io.gradients import read_bvals_bvecs
from dipy.io.image import load_nifti
from dipy.reconst.shm import CsaOdfModel
from dipy.direction import peaks_from_model
from dipy.segment.mask import median_otsu

"""
Download and read the data for this tutorial and load the raw diffusion data
and the affine.
"""

hardi_fname, hardi_bval_fname, hardi_bvec_fname = get_fnames('stanford_hardi')

data, affine = load_nifti(hardi_fname)

bvals, bvecs = read_bvals_bvecs(hardi_bval_fname, hardi_bvec_fname)
gtab = gradient_table(bvals, bvecs)

"""
img contains a nibabel Nifti1Image object (data) and gtab contains a
GradientTable object (gradient information e.g. b-values). For example to
read the b-values it is possible to write print(gtab.bvals).
"""

print('data.shape (%d, %d, %d, %d)' % data.shape)

"""
data.shape ``(81, 106, 76, 160)``

Remove most of the background using DIPY's mask module.
"""

maskdata, mask = median_otsu(data, vol_idx=range(10, 50), median_radius=3,
                             numpass=1, autocrop=True, dilate=2)

"""
We instantiate our CSA model with spherical harmonic order of 4
"""

csamodel = CsaOdfModel(gtab, 4)

"""
`Peaks_from_model` is used to calculate properties of the ODFs (Orientation
Distribution Function) and return for
example the peaks and their indices, or GFA which is similar to FA but for ODF
based models. This function mainly needs a reconstruction model, the data and a
sphere as input. The sphere is an object that represents the spherical discrete
grid where the ODF values will be evaluated.
"""

csapeaks = peaks_from_model(model=csamodel,
                            data=maskdata,
                            sphere=default_sphere,
                            relative_peak_threshold=.5,
                            min_separation_angle=25,
                            mask=mask,
                            return_odf=False,
                            normalize_peaks=True)

GFA = csapeaks.gfa

print('GFA.shape (%d, %d, %d)' % GFA.shape)

"""
GFA.shape ``(81, 106, 76)``

Apart from GFA, csapeaks also has the attributes peak_values, peak_indices and
ODF. peak_values shows the maxima values of the ODF and peak_indices gives us
their position on the discrete sphere that was used to do the reconstruction of
the ODF. In order to obtain the full ODF, return_odf should be True. Before
enabling this option, make sure that you have enough memory.

Let's visualize the ODFs of a small rectangular area in an axial slice of the
splenium of the corpus callosum (CC).
"""

data_small = maskdata[13:43, 44:74, 28:29]

from dipy.viz import window, actor

# Enables/disables interactive visualization
interactive = False

scene = window.Scene()

csaodfs = csamodel.fit(data_small).odf(default_sphere)

"""
It is common with CSA ODFs to produce negative values, we can remove those
using ``np.clip``
"""

csaodfs = np.clip(csaodfs, 0, np.max(csaodfs, -1)[..., None])
csa_odfs_actor = actor.odf_slicer(csaodfs, sphere=default_sphere,
                                  colormap='plasma', scale=0.4)
csa_odfs_actor.display(z=0)

scene.add(csa_odfs_actor)
print('Saving illustration as csa_odfs.png')
window.record(scene, n_frames=1, out_path='csa_odfs.png', size=(600, 600))
if interactive:
    window.show(scene)

"""
.. figure:: csa_odfs.png
   :align: center

   Constant Solid Angle ODFs.

.. include:: ../links_names.inc

References
----------

.. [Aganj2010] Aganj I, Lenglet C, Sapiro G, Yacoub E, Ugurbil K, Harel N.
   "Reconstruction of the orientation distribution function in single- and
   multiple-shell q-ball imaging within constant solid angle", Magnetic
   Resonance in Medicine. 2010 Aug;64(2):554-66. doi: 10.1002/mrm.22365


"""