import mne
from mne.datasets import sample
data_path = sample.data_path()

# the raw file containing the channel location + types
raw_fname = data_path + '/MEG/sample/sample_audvis_raw.fif'
# The paths to Freesurfer reconstructions
subjects_dir = 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
"""


mne.viz.plot_bem(subject=subject, subjects_dir=subjects_dir,
                 brain_surfaces='white', orientation='coronal')

"""
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 BEM surfaces

Compute and visualize BEM surfaces
The 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 skill
surface.
Computing the BEM surfaces requires FreeSurfer and makes use of either of
the two following command line tools:

gen_mne_watershed_bem
gen_mne_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 the bem folder of your subject FreeSurfer
reconstruction the necessary files.
End of explanation
"""


# The transformation file obtained by coregistration
trans = data_path + '/MEG/sample/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: Visualization the coregistration
The coregistration is 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
mne_analyze (Unix tools), mne.gui.coregistration (in Python) or mrilab
if you're using a Neuromag system.
For the Python version see :func:mne.gui.coregistration
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='oct6',
                             subjects_dir=subjects_dir, add_dist=False)
print(src)

"""
Explanation: Compute Source Space
The source space defines the position of the candidate source locations.
The following code compute such a cortical source space with
an OCT-6 resolution.
See setting_up_source_space for details on source space definition
and spacing parameter.
End of explanation
"""


mne.viz.plot_bem(subject=subject, subjects_dir=subjects_dir,
                 brain_surfaces='white', src=src, orientation='coronal')

"""
Explanation: src contains two parts, one for the left hemisphere (4098 locations) and
one for the right hemisphere (4098 locations). Sources can be visualized on
top of the BEM surfaces.
End of explanation
"""


import numpy as np  # noqa
from mayavi import mlab  # noqa
from surfer import Brain  # noqa

brain = Brain('sample', 'lh', 'inflated', subjects_dir=subjects_dir)
surf = brain.geo['lh']

vertidx = np.where(src[0]['inuse'])[0]

mlab.points3d(surf.x[vertidx], surf.y[vertidx],
              surf.z[vertidx], color=(1, 1, 0), scale_factor=1.5)

"""
Explanation: However, only sources that lie in the plotted MRI slices are shown.
Let's write a few lines of mayavi to see all sources.
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=2)
print(fwd)

"""
Explanation: Note that the 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 parameters meaning.
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
"""


n_dipoles = leadfield.shape[1]
vertices = [src_hemi['vertno'] for src_hemi in fwd_fixed['src']]
stc = mne.SourceEstimate(1e-9 * np.eye(n_dipoles), vertices, tmin=0., tstep=1)
leadfield = mne.apply_forward(fwd_fixed, stc, info).data / 1e-9

"""
Explanation: This is equivalent to the following code that explicitly applies the
forward operator to a source estimate composed of the identity operator:
End of explanation
"""


