"""
.. _ex-mixed-source-space-inverse:

=====================================================================
Compute MNE inverse solution on evoked data with a mixed source space
=====================================================================

Create a mixed source space and compute an MNE inverse solution on an
evoked dataset.
"""
# Author: Annalisa Pascarella <a.pascarella@iac.cnr.it>
#
# License: BSD-3-Clause
# Copyright the MNE-Python contributors.

# %%

import matplotlib.pyplot as plt
from nilearn import plotting

import mne
from mne.minimum_norm import apply_inverse, make_inverse_operator

# Set dir
data_path = mne.datasets.sample.data_path()
subject = "sample"
data_dir = data_path / "MEG" / subject
subjects_dir = data_path / "subjects"
bem_dir = subjects_dir / subject / "bem"

# Set file names
fname_mixed_src = bem_dir / f"{subject}-oct-6-mixed-src.fif"
fname_aseg = subjects_dir / subject / "mri" / "aseg.mgz"

fname_model = bem_dir / f"{subject}-5120-bem.fif"
fname_bem = bem_dir / f"{subject}-5120-bem-sol.fif"

fname_evoked = data_dir / f"{subject}_audvis-ave.fif"
fname_trans = data_dir / f"{subject}_audvis_raw-trans.fif"
fname_fwd = data_dir / f"{subject}_audvis-meg-oct-6-mixed-fwd.fif"
fname_cov = data_dir / f"{subject}_audvis-shrunk-cov.fif"

# %%
# Set up our source space
# -----------------------
# List substructures we are interested in. We select only the
# sub structures we want to include in the source space:

labels_vol = [
    "Left-Amygdala",
    "Left-Thalamus-Proper",
    "Left-Cerebellum-Cortex",
    "Brain-Stem",
    "Right-Amygdala",
    "Right-Thalamus-Proper",
    "Right-Cerebellum-Cortex",
]

# %%
# Get a surface-based source space, here with few source points for speed
# in this demonstration, in general you should use oct6 spacing!
src = mne.setup_source_space(
    subject, spacing="oct5", add_dist=False, subjects_dir=subjects_dir
)

# %%
# Now we create a mixed src space by adding the volume regions specified in the
# list labels_vol. First, read the aseg file and the source space bounds
# using the inner skull surface (here using 10mm spacing to save time,
# we recommend something smaller like 5.0 in actual analyses):

vol_src = mne.setup_volume_source_space(
    subject,
    mri=fname_aseg,
    pos=10.0,
    bem=fname_model,
    volume_label=labels_vol,
    subjects_dir=subjects_dir,
    add_interpolator=False,  # just for speed, usually this should be True
    verbose=True,
)

# Generate the mixed source space
src += vol_src
print(
    f"The source space contains {len(src)} spaces and "
    f"{sum(s['nuse'] for s in src)} vertices"
)

# %%
# View the source space
# ---------------------

src.plot(subjects_dir=subjects_dir)

# %%
# We could write the mixed source space with::
#
#    >>> write_source_spaces(fname_mixed_src, src, overwrite=True)
#
# We can also export source positions to NIfTI file and visualize it again:

nii_fname = bem_dir / f"{subject}-mixed-src.nii"
src.export_volume(nii_fname, mri_resolution=True, overwrite=True)
plotting.plot_img(str(nii_fname), cmap="nipy_spectral")

# %%
# Compute the fwd matrix
# ----------------------
fwd = mne.make_forward_solution(
    fname_evoked,
    fname_trans,
    src,
    fname_bem,
    mindist=5.0,  # ignore sources<=5mm from innerskull
    meg=True,
    eeg=False,
    n_jobs=None,
)
del src  # save memory

leadfield = fwd["sol"]["data"]
ns, nd = leadfield.shape
print(f"Leadfield size : {ns} sensors x {nd} dipoles")
print(
    f"The fwd source space contains {len(fwd['src'])} spaces and "
    f"{sum(s['nuse'] for s in fwd['src'])} vertices"
)

# Load data
condition = "Left Auditory"
evoked = mne.read_evokeds(fname_evoked, condition=condition, baseline=(None, 0))
noise_cov = mne.read_cov(fname_cov)

# %%
# Compute inverse solution
# ------------------------
snr = 3.0  # use smaller SNR for raw data
inv_method = "dSPM"  # sLORETA, MNE, dSPM
parc = "aparc"  # the parcellation to use, e.g., 'aparc' 'aparc.a2009s'
loose = dict(surface=0.2, volume=1.0)

lambda2 = 1.0 / snr**2

inverse_operator = make_inverse_operator(
    evoked.info, fwd, noise_cov, depth=None, loose=loose, verbose=True
)
del fwd

stc = apply_inverse(evoked, inverse_operator, lambda2, inv_method, pick_ori=None)
src = inverse_operator["src"]

# %%
# Plot the mixed source estimate
# ------------------------------

# sphinx_gallery_thumbnail_number = 3
initial_time = 0.1
stc_vec = apply_inverse(
    evoked, inverse_operator, lambda2, inv_method, pick_ori="vector"
)
brain = stc_vec.plot(
    hemi="both",
    src=inverse_operator["src"],
    views="coronal",
    initial_time=initial_time,
    subjects_dir=subjects_dir,
    brain_kwargs=dict(silhouette=True),
    smoothing_steps=7,
)

# %%
# Plot the surface
# ----------------
brain = stc.surface().plot(
    initial_time=initial_time, subjects_dir=subjects_dir, smoothing_steps=7
)
# %%
# Plot the volume
# ---------------

fig = stc.volume().plot(initial_time=initial_time, src=src, subjects_dir=subjects_dir)

# %%
# Process labels
# --------------
# Average the source estimates within each label of the cortical parcellation
# and each sub structure contained in the src space

# Get labels for FreeSurfer 'aparc' cortical parcellation with 34 labels/hemi
labels_parc = mne.read_labels_from_annot(subject, parc=parc, subjects_dir=subjects_dir)

label_ts = mne.extract_label_time_course(
    [stc], labels_parc, src, mode="mean", allow_empty=True
)

# plot the times series of 2 labels
fig, axes = plt.subplots(1, layout="constrained")
axes.plot(1e3 * stc.times, label_ts[0][0, :], "k", label="bankssts-lh")
axes.plot(1e3 * stc.times, label_ts[0][-1, :].T, "r", label="Brain-stem")
axes.set(xlabel="Time (ms)", ylabel="MNE current (nAm)")
axes.legend()