"""
.. _reconst_sfm:

==============================================
Reconstruction with the Sparse Fascicle Model
==============================================

In this example, we will use the Sparse Fascicle Model (SFM) [Rokem2015]_, to
reconstruct the fiber Orientation Distribution Function (fODF) in every voxel.

First, we import the modules we will use in this example:
"""

import dipy.reconst.sfm as sfm
import dipy.data as dpd
import dipy.direction.peaks as dpp
from dipy.io.image import load_nifti
from dipy.io.gradients import read_bvals_bvecs
from dipy.core.gradients import gradient_table
from dipy.viz import window, actor

###############################################################################
# For the purpose of this example, we will use the Stanford HARDI dataset (150
# directions, single b-value of 2000 $s/mm^2$) that can be automatically
# downloaded. If you have not yet downloaded this data-set in one of the other
# examples, you will need to be connected to the internet the first time you run
# this example. The data will be stored for subsequent runs, and for use with
# other examples.
# 


hardi_fname, hardi_bval_fname, hardi_bvec_fname = dpd.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)

# Enables/disables interactive visualization
interactive = False

###############################################################################
# Reconstruction of the fiber ODF in each voxel guides subsequent tracking
# steps. Here, the model is the Sparse Fascicle Model, described in
# [Rokem2014]_. This model reconstructs the diffusion signal as a combination of
# the signals from different fascicles. This model can be written as:
# 
# .. math::
# 
#     y = X\beta
# 
# Where $y$ is the signal and $\beta$ are weights on different points in the
# sphere. The columns of the design matrix, $X$ are the signals in each point in
# the measurement that would be predicted if there was a fascicle oriented in the
# direction represented by that column. Typically, the signal used for this
# kernel will be a prolate tensor with axial diffusivity 3-5 times higher than
# its radial diffusivity. The exact numbers can also be estimated from examining
# parts of the brain in which there is known to be only one fascicle (e.g. in
# corpus callosum).
# 
# Sparsity constraints on the fiber ODF ($\beta$) are set through the Elastic Net
# algorithm [Zou2005]_.
# 
# Elastic Net optimizes the following cost function:
# 
# .. math::
# 
#     \sum_{i=1}^{n}{(y_i - \hat{y}_i)^2} + \alpha (\lambda \sum_{j=1}^{m}{w_j}+(1-\lambda) \sum_{j=1}^{m}{w^2_j}
# 
# where $\hat{y}$ is the signal predicted for a particular setting of $\beta$,
# such that the left part of this expression is the squared loss function;
# $\alpha$ is a parameter that sets the balance between the squared loss on
# the data, and the regularization constraints. The regularization parameter
# $\lambda$ sets the `l1_ratio`, which controls the balance between L1-sparsity
# (low sum of weights), and low L2-sparsity (low sum-of-squares of the weights).
# 
# Just like Constrained Spherical Deconvolution (see :ref:`reconst-csd`), the SFM
# requires the definition of a response function. We'll take advantage of the
# automated algorithm in the :mod:`csdeconv` module to find this response
# function:
# 


from dipy.reconst.csdeconv import auto_response_ssst
response, ratio = auto_response_ssst(gtab, data, roi_radii=10, fa_thr=0.7)

###############################################################################
# The ``response`` return value contains two entries. The first is an array with
# the eigenvalues of the response function and the second is the average S0 for
# this response.
# 
# It is a very good practice to always validate the result of
# ``auto_response_ssst``. For, this purpose we can print it and have a look
# at its values.


print(response)

###############################################################################
# (array([ 0.0014,  0.00029,  0.00029]), 416.206)
# 
# We initialize an SFM model object, using these values. We will use the default
# sphere (362  vertices, symmetrically distributed on the surface of the sphere),
# as a set of putative fascicle directions that are considered in the model


sphere = dpd.get_sphere()
sf_model = sfm.SparseFascicleModel(gtab, sphere=sphere,
                                   l1_ratio=0.5, alpha=0.001,
                                   response=response[0])

###############################################################################
# For the purpose of the example, we will consider a small volume of data
# containing parts of the corpus callosum and of the centrum semiovale


data_small = data[20:50, 55:85, 38:39]

###############################################################################
# Fitting the model to this small volume of data, we calculate the ODF of this
# model on the sphere, and plot it.


sf_fit = sf_model.fit(data_small)
sf_odf = sf_fit.odf(sphere)

fodf_spheres = actor.odf_slicer(sf_odf, sphere=sphere, scale=0.8,
                                colormap='plasma')

scene = window.Scene()
scene.add(fodf_spheres)

print('Saving illustration as sf_odfs.png')
window.record(scene, out_path='sf_odfs.png', size=(1000, 1000))
if interactive:
    window.show(scene)

###############################################################################
# We can extract the peaks from the ODF, and plot these as well


sf_peaks = dpp.peaks_from_model(sf_model,
                                data_small,
                                sphere,
                                relative_peak_threshold=.5,
                                min_separation_angle=25,
                                return_sh=False)


scene.clear()
fodf_peaks = actor.peak_slicer(sf_peaks.peak_dirs, sf_peaks.peak_values)
scene.add(fodf_peaks)

print('Saving illustration as sf_peaks.png')
window.record(scene, out_path='sf_peaks.png', size=(1000, 1000))
if interactive:
    window.show(scene)

###############################################################################
# Finally, we plot both the peaks and the ODFs, overlaid:


fodf_spheres.GetProperty().SetOpacity(0.4)
scene.add(fodf_spheres)

print('Saving illustration as sf_both.png')
window.record(scene, out_path='sf_both.png', size=(1000, 1000))
if interactive:
    window.show(scene)

###############################################################################
# .. figure:: sf_both.png
#    :align: center
# 
#    SFM Peaks and ODFs.
# 
# To see how to use this information in tracking, proceed to :ref:`sfm-track`.
# 
# References
# ----------
# 
# .. [Rokem2015] Ariel Rokem, Jason D. Yeatman, Franco Pestilli, Kendrick
#    N. Kay, Aviv Mezer, Stefan van der Walt, Brian A. Wandell
#    (2015). Evaluating the accuracy of diffusion MRI models in white
#    matter. PLoS ONE 10(4): e0123272. doi:10.1371/journal.pone.0123272
# 
# .. [Zou2005] Zou H, Hastie T (2005). Regularization and variable
#    selection via the elastic net. J R Stat Soc B:301-320
#