# Author: Tommy Clausner <tommy.clausner@gmail.com>
#
# License: BSD-3-Clause

import os

import nibabel as nib
import mne
from mne.datasets import sample, fetch_fsaverage
from mne.minimum_norm import apply_inverse, read_inverse_operator
from nilearn.plotting import plot_glass_brain

print(__doc__)

"""
Explanation: Morph volumetric source estimate
This example demonstrates how to morph an individual subject's
:class:mne.VolSourceEstimate to a common reference space. We achieve this
using :class:mne.SourceMorph. Data will be morphed based on
an affine transformation and a nonlinear registration method
known as Symmetric Diffeomorphic Registration (SDR) by
:footcite:AvantsEtAl2008.
Transformation is estimated from the subject's anatomical T1 weighted MRI
(brain) to FreeSurfer's 'fsaverage' T1 weighted MRI (brain)_.
Afterwards the transformation will be applied to the volumetric source
estimate. The result will be plotted, showing the fsaverage T1 weighted
anatomical MRI, overlaid with the morphed volumetric source estimate.
End of explanation
"""


sample_dir_raw = sample.data_path()
sample_dir = os.path.join(sample_dir_raw, 'MEG', 'sample')
subjects_dir = os.path.join(sample_dir_raw, 'subjects')

fname_evoked = os.path.join(sample_dir, 'sample_audvis-ave.fif')
fname_inv = os.path.join(sample_dir, 'sample_audvis-meg-vol-7-meg-inv.fif')

fname_t1_fsaverage = os.path.join(subjects_dir, 'fsaverage', 'mri',
                                  'brain.mgz')
fetch_fsaverage(subjects_dir)  # ensure fsaverage src exists
fname_src_fsaverage = subjects_dir + '/fsaverage/bem/fsaverage-vol-5-src.fif'

"""
Explanation: Setup paths
End of explanation
"""


evoked = mne.read_evokeds(fname_evoked, condition=0, baseline=(None, 0))
inverse_operator = read_inverse_operator(fname_inv)

# Apply inverse operator
stc = apply_inverse(evoked, inverse_operator, 1.0 / 3.0 ** 2, "dSPM")

# To save time
stc.crop(0.09, 0.09)

"""
Explanation: Compute example data. For reference see ex-inverse-volume.
Load data:
End of explanation
"""


src_fs = mne.read_source_spaces(fname_src_fsaverage)
morph = mne.compute_source_morph(
    inverse_operator['src'], subject_from='sample', subjects_dir=subjects_dir,
    niter_affine=[10, 10, 5], niter_sdr=[10, 10, 5],  # just for speed
    src_to=src_fs, verbose=True)

"""
Explanation: Get a SourceMorph object for VolSourceEstimate
subject_from can typically be inferred from
:class:src &lt;mne.SourceSpaces&gt;,
and subject_to is  set to 'fsaverage' by default. subjects_dir can be
None when set in the environment. In that case SourceMorph can be initialized
taking src as only argument. See :class:mne.SourceMorph for more
details.
The default parameter setting for zooms will cause the reference volumes
to be resliced before computing the transform. A value of '5' would cause
the function to reslice to an isotropic voxel size of 5 mm. The higher this
value the less accurate but faster the computation will be.
The recommended way to use this is to morph to a specific destination source
space so that different subject_from morphs will go to the same space.`
A standard usage for volumetric data reads:
End of explanation
"""


stc_fsaverage = morph.apply(stc)

"""
Explanation: Apply morph to VolSourceEstimate
The morph can be applied to the source estimate data, by giving it as the
first argument to the :meth:morph.apply() &lt;mne.SourceMorph.apply&gt; method.
<div class="alert alert-info"><h4>Note</h4><p>Volumetric morphing is much slower than surface morphing because the
     volume for each time point is individually resampled and SDR morphed.
     The :meth:`mne.SourceMorph.compute_vol_morph_mat` method can be used
     to compute an equivalent sparse matrix representation by computing the
     transformation for each source point individually. This generally takes
     a few minutes to compute, but can be
     :meth:`saved <mne.SourceMorph.save>` to disk and be reused. The
     resulting sparse matrix operation is very fast (about 400× faster) to
     :meth:`apply <mne.SourceMorph.apply>`. This approach is more efficient
     when the number of time points to be morphed exceeds the number of
     source space points, which is generally in the thousands. This can
     easily occur when morphing many time points and multiple conditions.</p></div>
End of explanation
"""


# Create mri-resolution volume of results
img_fsaverage = morph.apply(stc, mri_resolution=2, output='nifti1')

"""
Explanation: Convert morphed VolSourceEstimate into NIfTI
We can convert our morphed source estimate into a NIfTI volume using
:meth:morph.apply(..., output='nifti1') &lt;mne.SourceMorph.apply&gt;.
End of explanation
"""


# Load fsaverage anatomical image
t1_fsaverage = nib.load(fname_t1_fsaverage)

# Plot glass brain (change to plot_anat to display an overlaid anatomical T1)
display = plot_glass_brain(t1_fsaverage,
                           title='subject results to fsaverage',
                           draw_cross=False,
                           annotate=True)

# Add functional data as overlay
display.add_overlay(img_fsaverage, alpha=0.75)

"""
Explanation: Plot results
End of explanation
"""


