import os.path as op
import mne
from mne.datasets import sample
data_path = sample.data_path()

# the raw file containing the channel location + types
sample_dir = op.join(data_path, 'MEG', 'sample',)
raw_fname = op.join(sample_dir, 'sample_audvis_raw.fif')
# The paths to Freesurfer reconstructions
subjects_dir = op.join(data_path, 'subjects')
subject = 'sample'

"""
Explanation: Head model and forward computation
The aim of this tutorial is to be a getting started for forward computation.
For more extensive details and presentation of the general concepts for forward
modeling, see ch_forward.
End of explanation
"""


plot_bem_kwargs = dict(
    subject=subject, subjects_dir=subjects_dir,
    brain_surfaces='white', orientation='coronal',
    slices=[50, 100, 150, 200])

mne.viz.plot_bem(**plot_bem_kwargs)

"""
Explanation: Computing the forward operator
To compute a forward operator we need:

a -trans.fif file that contains the coregistration info.
a source space
the :term:BEM surfaces

Compute and visualize BEM surfaces
The :term:BEM surfaces are the triangulations of the interfaces between
different tissues needed for forward computation. These surfaces are for
example the inner skull surface, the outer skull surface and the outer skin
surface, a.k.a. scalp surface.
Computing the BEM surfaces requires FreeSurfer and makes use of
the command-line tools mne watershed_bem or mne flash_bem, or
the related functions :func:mne.bem.make_watershed_bem or
:func:mne.bem.make_flash_bem.
Here we'll assume it's already computed. It takes a few minutes per subject.
For EEG we use 3 layers (inner skull, outer skull, and skin) while for
MEG 1 layer (inner skull) is enough.
Let's look at these surfaces. The function :func:mne.viz.plot_bem
assumes that you have the bem folder of your subject's FreeSurfer
reconstruction, containing the necessary surface files. Here we use a smaller
than default subset of slices for speed.
End of explanation
"""


# The transformation file obtained by coregistration
trans = op.join(sample_dir, 'sample_audvis_raw-trans.fif')

info = mne.io.read_info(raw_fname)
# Here we look at the dense head, which isn't used for BEM computations but
# is useful for coregistration.
mne.viz.plot_alignment(info, trans, subject=subject, dig=True,
                       meg=['helmet', 'sensors'], subjects_dir=subjects_dir,
                       surfaces='head-dense')

"""
Explanation: Visualizing the coregistration
The coregistration is the operation that allows to position the head and the
sensors in a common coordinate system. In the MNE software the transformation
to align the head and the sensors in stored in a so-called trans file.
It is a FIF file that ends with -trans.fif. It can be obtained with
:func:mne.gui.coregistration (or its convenient command line
equivalent mne coreg), or mrilab if you're using a Neuromag
system.
Here we assume the coregistration is done, so we just visually check the
alignment with the following code.
End of explanation
"""


src = mne.setup_source_space(subject, spacing='oct4', add_dist='patch',
                             subjects_dir=subjects_dir)
print(src)

"""
Explanation: Compute Source Space
The source space defines the position and orientation of the candidate source
locations. There are two types of source spaces:


surface-based source space when the candidates are confined to a
  surface.


volumetric or discrete source space when the candidates are discrete,
  arbitrarily located source points bounded by the surface.


Surface-based source space is computed using
:func:mne.setup_source_space, while volumetric source space is computed
using :func:mne.setup_volume_source_space.
We will now compute a surface-based source space with an 'oct4'
resolution. See setting_up_source_space for details on source space
definition and spacing parameter.
<div class="alert alert-danger"><h4>Warning</h4><p>``'oct4'`` is used here just for speed, for real analyses the recommended
    spacing is ``'oct6'``.</p></div>
End of explanation
"""


mne.viz.plot_bem(src=src, **plot_bem_kwargs)

"""
Explanation: The surface based source space src contains two parts, one for the left
hemisphere (258 locations) and one for the right hemisphere (258
locations). Sources can be visualized on top of the BEM surfaces in purple.
End of explanation
"""


sphere = (0.0, 0.0, 0.04, 0.09)
vol_src = mne.setup_volume_source_space(
    subject, subjects_dir=subjects_dir, sphere=sphere, sphere_units='m',
    add_interpolator=False)  # just for speed!
print(vol_src)

mne.viz.plot_bem(src=vol_src, **plot_bem_kwargs)

"""
Explanation: To compute a volume based source space defined with a grid of candidate
dipoles inside a sphere of radius 90mm centered at (0.0, 0.0, 40.0) mm
you can use the following code.
Obviously here, the sphere is not perfect. It is not restricted to the
brain and it can miss some parts of the cortex.
End of explanation
"""


surface = op.join(subjects_dir, subject, 'bem', 'inner_skull.surf')
vol_src = mne.setup_volume_source_space(
    subject, subjects_dir=subjects_dir, surface=surface,
    add_interpolator=False)  # Just for speed!
print(vol_src)

mne.viz.plot_bem(src=vol_src, **plot_bem_kwargs)

"""
Explanation: To compute a volume based source space defined with a grid of candidate
dipoles inside the brain (requires the :term:BEM surfaces) you can use the
following.
End of explanation
"""


fig = mne.viz.plot_alignment(subject=subject, subjects_dir=subjects_dir,
                             surfaces='white', coord_frame='mri',
                             src=src)
mne.viz.set_3d_view(fig, azimuth=173.78, elevation=101.75,
                    distance=0.30, focalpoint=(-0.03, -0.01, 0.03))

"""
Explanation: <div class="alert alert-info"><h4>Note</h4><p>Some sources may appear to be outside the BEM inner skull contour.
          This is because the ``slices`` are decimated for plotting here.
          Each slice in the figure actually represents several MRI slices,
          but only the MRI voxels and BEM boundaries for a single (midpoint
          of the given slice range) slice are shown, whereas the source space
          points plotted on that midpoint slice consist of all points
          for which that slice (out of all slices shown) was the closest.</p></div>

Now let's see how to view all sources in 3D.
End of explanation
"""


conductivity = (0.3,)  # for single layer
# conductivity = (0.3, 0.006, 0.3)  # for three layers
model = mne.make_bem_model(subject='sample', ico=4,
                           conductivity=conductivity,
                           subjects_dir=subjects_dir)
bem = mne.make_bem_solution(model)

"""
Explanation: Compute forward solution
We can now compute the forward solution.
To reduce computation we'll just compute a single layer BEM (just inner
skull) that can then be used for MEG (not EEG).
We specify if we want a one-layer or a three-layer BEM using the
conductivity parameter.
The BEM solution requires a BEM model which describes the geometry
of the head the conductivities of the different tissues.
End of explanation
"""


fwd = mne.make_forward_solution(raw_fname, trans=trans, src=src, bem=bem,
                                meg=True, eeg=False, mindist=5.0, n_jobs=None,
                                verbose=True)
print(fwd)

"""
Explanation: Note that the :term:BEM does not involve any use of the trans file. The BEM
only depends on the head geometry and conductivities.
It is therefore independent from the MEG data and the head position.
Let's now compute the forward operator, commonly referred to as the
gain or leadfield matrix.
See :func:mne.make_forward_solution for details on the meaning of each
parameter.
End of explanation
"""


print(f'Before: {src}')
print(f'After:  {fwd["src"]}')

"""
Explanation: <div class="alert alert-danger"><h4>Warning</h4><p>Forward computation can remove vertices that are too close to (or outside)
   the inner skull surface. For example, here we have gone from 516 to 474
   vertices in use. For many functions, such as
   :func:`mne.compute_source_morph`, it is important to pass ``fwd['src']``
   or ``inv['src']`` so that this removal is adequately accounted for.</p></div>
End of explanation
"""


leadfield = fwd['sol']['data']
print("Leadfield size : %d sensors x %d dipoles" % leadfield.shape)

"""
Explanation: We can explore the content of fwd to access the numpy array that contains
the gain matrix.
End of explanation
"""


fwd_fixed = mne.convert_forward_solution(fwd, surf_ori=True, force_fixed=True,
                                         use_cps=True)
leadfield = fwd_fixed['sol']['data']
print("Leadfield size : %d sensors x %d dipoles" % leadfield.shape)

"""
Explanation: To extract the numpy array containing the forward operator corresponding to
the source space fwd['src'] with cortical orientation constraint
we can use the following:
End of explanation
"""


