#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
=================================================================
test restrict_forward_to_label vs stc.in_label method
on sensitivity profile using labels
=================================================================
"""

print(__doc__)

# --------------------------------------
# import modules
# --------------------------------------
from os.path import join
import mne
from mne.datasets import sample
from mne import make_forward_solution

# load data
data_path = sample.data_path()
subject = 'sample'
subjects_dir = join(data_path, 'subjects')
meg_path = join(data_path, 'MEG', 'sample')

# read evoked
evoked_fname = join(meg_path, 'sample_audvis-ave.fif')
evoked = mne.read_evokeds(evoked_fname)[0]
evoked.pick_types(meg=True, eeg=True).apply_baseline((None, 0))

# create a forward sol
bem = join(data_path,  'subjects', 'sample', 'bem', 'sample-5120-5120-5120-bem-sol.fif')
src = join(data_path,  'subjects', 'sample', 'bem', 'sample-oct-6-src.fif')
trans = join(meg_path, 'sample_audvis_raw-trans.fif')

# make forward solution
fwd = make_forward_solution(evoked.info, trans, src, bem, meg=True, eeg=True, mindist=3.0, n_jobs=3)

# convert the forward solution with surface orientation /normal orientation
fwd = mne.convert_forward_solution(fwd, surf_ori=True, copy=False, force_fixed=True)

# method 1: restrict_forward_to_label
# read labels
labels = mne.read_labels_from_annot("fsaverage", "HCPMMP1", subjects_dir=subjects_dir)
aud_label_lh = [label for label in labels if label.name == "L_A1_ROI-lh"][0]
aud_label_rh = [label for label in labels if label.name == "R_A1_ROI-rh"][0]

# morphed the labels for individual subject anatomy
subject_labels = mne.morph_labels([aud_label_lh, aud_label_rh], subject_to=subject, subject_from='fsaverage',
                                  subjects_dir=subjects_dir, verbose=None)
# restrict fwd to auditory labels
fwd_aud_labels = mne.forward.restrict_forward_to_label(fwd.copy(), subject_labels)

# test on different sensors
ch_types = ['grad', 'eeg', 'mag']

for ch_type in ch_types:
    # compute sensitivity maps with forward labels
    grad_map_m1 = mne.sensitivity_map(fwd_aud_labels, ch_type=ch_type, mode='fixed')
    m1_mean = grad_map_m1.data.mean()  # mean sensitivity of the labels

    # method 2: 'in_label' method

    # compute sensitivity maps with full forward
    grad_map_full = mne.sensitivity_map(fwd, ch_type=ch_type, mode='fixed')

    # now compute the sensitivity for auditory labels using stc: 'in_label' method
    grad_map_m2 = grad_map_full.in_label(subject_labels[0]+subject_labels[1])

    m2_mean = grad_map_m2.data.mean()  # mean sensitivity of the labels

    print('%s method1 output:' % ch_type, round(m1_mean, 2))
    print('%s method2 output:' % ch_type, round(m2_mean, 2))