# Author: Alexandre Gramfort <alexandre.gramfort@inria.fr>
#
# License: BSD (3-clause)

import mne
from mne.datasets import sample
from mne.beamformer import make_lcmv, apply_lcmv

print(__doc__)

"""
Explanation: Compute LCMV inverse solution in volume source space
Compute LCMV beamformer on an auditory evoked dataset in a volume source space,
and show activation on fsaverage.
End of explanation
"""


data_path = sample.data_path()
subjects_dir = data_path + '/subjects'
raw_fname = data_path + '/MEG/sample/sample_audvis_filt-0-40_raw.fif'
fname_fwd = data_path + '/MEG/sample/sample_audvis-meg-vol-7-fwd.fif'

# Get epochs
event_id, tmin, tmax = [1, 2], -0.2, 0.5

# Read forward model
forward = mne.read_forward_solution(fname_fwd)

# Setup for reading the raw data
raw = mne.io.read_raw_fif(raw_fname, preload=True)
raw.info['bads'] = ['MEG 2443', 'EEG 053']  # 2 bads channels
events = mne.find_events(raw)

# Pick the channels of interest
raw.pick(['meg', 'eog'])

# Read epochs
proj = False  # already applied
epochs = mne.Epochs(raw, events, event_id, tmin, tmax,
                    baseline=(None, 0), preload=True, proj=proj,
                    reject=dict(grad=4000e-13, mag=4e-12, eog=150e-6))
evoked = epochs.average()

# Visualize sensor space data
evoked.plot_joint()

"""
Explanation: Data preprocessing:
End of explanation
"""


# Read regularized noise covariance and compute regularized data covariance
noise_cov = mne.compute_covariance(epochs, tmin=tmin, tmax=0, method='shrunk',
                                   rank=None)
data_cov = mne.compute_covariance(epochs, tmin=0.04, tmax=0.15,
                                  method='shrunk', rank=None)

# Compute weights of free orientation (vector) beamformer with weight
# normalization (neural activity index, NAI). Providing a noise covariance
# matrix enables whitening of the data and forward solution. Source orientation
# is optimized by setting pick_ori to 'max-power'.
# weight_norm can also be set to 'unit-noise-gain'. Source orientation can also
# be 'normal' (but only when using a surface-based source space) or None,
# which computes a vector beamfomer. Note, however, that not all combinations
# of orientation selection and weight normalization are implemented yet.
filters = make_lcmv(evoked.info, forward, data_cov, reg=0.05,
                    noise_cov=noise_cov, pick_ori='max-power',
                    weight_norm='nai', rank=None)
print(filters)

# You can save these with:
# filters.save('filters-lcmv.h5')

# Apply this spatial filter to the evoked data.
stc = apply_lcmv(evoked, filters, max_ori_out='signed')

"""
Explanation: Compute covariance matrices, fit and apply  spatial filter.
End of explanation
"""


# You can save result in stc files with:
# stc.save('lcmv-vol')
lims = [0.3, 0.6, 0.9]
stc.plot(
    src=forward['src'], subject='sample', subjects_dir=subjects_dir,
    clim=dict(kind='value', pos_lims=lims), mode='stat_map',
    initial_time=0.1, verbose=True)

"""
Explanation: Plot source space activity:
End of explanation
"""


stc.plot(
    src=forward['src'], subject='sample', subjects_dir=subjects_dir,
    mode='glass_brain', clim=dict(kind='value', lims=lims),
    initial_time=0.1, verbose=True)

"""
Explanation: Now let's plot this on a glass brain, which will automatically transform the
data to MNI Talairach space:
End of explanation
"""


morph = mne.compute_source_morph(
    forward['src'], 'sample', 'fsaverage', subjects_dir=subjects_dir,
    zooms=7, verbose=True)
stc.copy().crop(0.05, 0.18).plot(
    src=morph, subject='fsaverage', subjects_dir=subjects_dir,
    mode='stat_map', clim=dict(kind='value', pos_lims=lims),
    initial_time=0.1, verbose=True)

"""
Explanation: Finally let's get another view, this time plotting again a 'stat_map'
style but using volumetric morphing to get data to fsaverage space,
which we can get by passing a :class:mne.SourceMorph as the src
argument to mne.VolSourceEstimate.plot. To save a bit of speed when
applying the morph, we will crop the STC:
End of explanation
"""


