"""
============================================================================
Tracking with Robust Unbiased Model-BAsed Spherical Deconvolution (RUMBA-SD)
============================================================================

Here, we demonstrate fiber tracking using a probabilistic direction getter
and RUMBA-SD, a model introduced in [CanalesRodriguez2015]_. This model adapts
Richardson-Lucy deconvolution by assuming Rician or Noncentral Chi noise
instead of Gaussian, which more accurately reflects the noise from MRI
scanners (see also :ref:`example_reconst_rumba`). This tracking tutorial is an
extension on :ref:`example_tracking_probabilistic`.

We start by loading sample data and identifying a fiber response function.
"""

from numpy.linalg import inv

from dipy.core.gradients import gradient_table
from dipy.data import get_fnames, small_sphere
from dipy.io.gradients import read_bvals_bvecs
from dipy.io.image import load_nifti, load_nifti_data
from dipy.reconst.csdeconv import auto_response_ssst
from dipy.tracking import utils
from dipy.tracking.local_tracking import LocalTracking
from dipy.tracking.streamline import Streamlines, transform_streamlines
from dipy.tracking.stopping_criterion import ThresholdStoppingCriterion
from dipy.viz import window, actor, colormap
from dipy.reconst.rumba import RumbaSDModel

# Enables/disables interactive visualization
interactive = False

hardi_fname, hardi_bval_fname, hardi_bvec_fname = get_fnames('stanford_hardi')
label_fname = get_fnames('stanford_labels')
t1_fname = get_fnames('stanford_t1')

data, affine, hardi_img = load_nifti(hardi_fname, return_img=True)
labels = load_nifti_data(label_fname)
t1_data, t1_aff = load_nifti(t1_fname)
bvals, bvecs = read_bvals_bvecs(hardi_bval_fname, hardi_bvec_fname)
gtab = gradient_table(bvals, bvecs)

seed_mask = (labels == 2)
white_matter = (labels == 1) | (labels == 2)
seeds = utils.seeds_from_mask(seed_mask, affine, density=2)

response, ratio = auto_response_ssst(gtab, data, roi_radii=10, fa_thr=0.7)

sphere = small_sphere

"""
We can now initialize a `RumbaSdModel` model and fit it globally by setting
`voxelwise` to `False`. For this example, TV regularization (`use_tv`) will be
turned off for efficiency, although its usage can provide more coherent results
in fiber tracking. The fit will take about 5 minutes to complete.
"""

rumba = RumbaSDModel(gtab, wm_response=response[0], n_iter=200,
                     voxelwise=False, use_tv=False, sphere=sphere)
rumba_fit = rumba.fit(data, mask=white_matter)
odf = rumba_fit.odf()  # fODF
f_wm = rumba_fit.f_wm  # white matter volume fractions


"""
To establish stopping criterion, a common technique is to use the Generalized
Fractional Anisotropy (GFA). One point of caution is that RUMBA-SD by default
separates the fODF from an isotropic compartment. This can bias GFA results
computed on the fODF, although it will still generally be an effective
criterion.

However, an alternative stopping criterion that takes advantage of this
feature is to use RUMBA-SD's white matter volume fraction map.
"""

stopping_criterion = ThresholdStoppingCriterion(f_wm, .25)

"""
We can visualize a slice of this mask.
"""

import matplotlib.pyplot as plt

sli = f_wm.shape[2] // 2
plt.figure()

plt.subplot(1, 2, 1).set_axis_off()
plt.imshow(f_wm[:, :, sli].T, cmap='gray', origin='lower')

plt.subplot(1, 2, 2).set_axis_off()
plt.imshow((f_wm[:, :, sli] > 0.25).T, cmap='gray', origin='lower')

plt.savefig('f_wm_tracking_mask.png')

"""
.. figure:: f_wm_tracking_mask.png
   :align: center

   White matter volume fraction slice
"""

"""
These discrete fODFs can be used as a PMF in the `ProbabilisticDirectionGetter`
for sampling tracking directions. The PMF must be strictly non-negative;
RUMBA-SD already adheres to this constraint so no further manipulation of the
fODFs is necessary.
"""

from dipy.direction import ProbabilisticDirectionGetter
from dipy.io.stateful_tractogram import Space, StatefulTractogram
from dipy.io.streamline import save_trk

prob_dg = ProbabilisticDirectionGetter.from_pmf(odf, max_angle=30.,
                                                sphere=sphere)
streamline_generator = LocalTracking(prob_dg, stopping_criterion, seeds,
                                     affine, step_size=.5)
streamlines = Streamlines(streamline_generator)

color = colormap.line_colors(streamlines)
streamlines_actor = actor.streamtube(
    list(transform_streamlines(streamlines, inv(t1_aff))),
    color, linewidth=0.1)

vol_actor = actor.slicer(t1_data)
vol_actor.display(x=40)
vol_actor2 = vol_actor.copy()
vol_actor2.display(z=35)

scene = window.Scene()
scene.add(vol_actor)
scene.add(vol_actor2)
scene.add(streamlines_actor)
if interactive:
    window.show(scene)

window.record(scene, out_path='tractogram_probabilistic_rumba.png',
              size=(800, 800))

sft = StatefulTractogram(streamlines, hardi_img, Space.RASMM)
save_trk(sft, "tractogram_probabilistic_rumba.trk")

"""
.. figure:: tractogram_probabilistic_rumba.png
   :align: center

   RUMBA-SD tractogram
"""

"""
References
----------

.. [CanalesRodriguez2015] Canales-Rodríguez, E. J., Daducci, A., Sotiropoulos,
   S. N., Caruyer, E., Aja-Fernández, S., Radua, J., Mendizabal, J. M. Y.,
   Iturria-Medina, Y., Melie-García, L., Alemán-Gómez, Y., Thiran, J.-P.,
   Sarró, S., Pomarol-Clotet, E., & Salvador, R. (2015). Spherical
   Deconvolution of Multichannel Diffusion MRI Data with Non-Gaussian Noise
   Models and Spatial Regularization. PLOS ONE, 10(10), e0138910.
   https://doi.org/10.1371/journal.pone.0138910

.. include:: ../links_names.inc

"""