"""
.. _ex-psd-vol:

===============================================
Plot point-spread functions (PSFs) for a volume
===============================================

Visualise PSF at one volume vertex for sLORETA.
"""
# Authors: Olaf Hauk <olaf.hauk@mrc-cbu.cam.ac.uk>
#          Alexandre Gramfort <alexandre.gramfort@inria.fr>
#          Eric Larson <larson.eric.d@gmail.com>
#
# License: BSD-3-Clause
# Copyright the MNE-Python contributors.

# %%

import numpy as np

import mne
from mne.datasets import sample
from mne.minimum_norm import get_point_spread, make_inverse_resolution_matrix

print(__doc__)

data_path = sample.data_path()
subjects_dir = data_path / "subjects"
meg_path = data_path / "MEG" / "sample"
fname_cov = meg_path / "sample_audvis-cov.fif"
fname_evo = meg_path / "sample_audvis-ave.fif"
fname_trans = meg_path / "sample_audvis_raw-trans.fif"
fname_bem = subjects_dir / "sample" / "bem" / "sample-5120-bem-sol.fif"

# %%
# For the volume, create a coarse source space for speed (don't do this in
# real code!), then compute the forward using this source space.

# read noise cov and and evoked
noise_cov = mne.read_cov(fname_cov)
evoked = mne.read_evokeds(fname_evo, 0)

# create a coarse source space
src_vol = mne.setup_volume_source_space(  # this is a very course resolution!
    "sample", pos=15.0, subjects_dir=subjects_dir, add_interpolator=False
)  # usually you want True, this is just for speed

# compute the forward
forward_vol = mne.make_forward_solution(  # MEG-only for speed
    evoked.info, fname_trans, src_vol, fname_bem, eeg=False
)
del src_vol

# %%
# Now make an inverse operator and compute the PSF at a source.
inverse_operator_vol = mne.minimum_norm.make_inverse_operator(
    info=evoked.info, forward=forward_vol, noise_cov=noise_cov
)

# compute resolution matrix for sLORETA
rm_lor_vol = make_inverse_resolution_matrix(
    forward_vol, inverse_operator_vol, method="sLORETA", lambda2=1.0 / 9.0
)

# get PSF and CTF for sLORETA at one vertex
sources_vol = [100]
stc_psf_vol = get_point_spread(rm_lor_vol, forward_vol["src"], sources_vol, norm=True)
del rm_lor_vol

##############################################################################
# Visualize
# ---------
# PSF:

# Which vertex corresponds to selected source
src_vol = forward_vol["src"]
verttrue_vol = src_vol[0]["vertno"][sources_vol]

# find vertex with maximum in PSF
max_vert_idx, _ = np.unravel_index(stc_psf_vol.data.argmax(), stc_psf_vol.data.shape)
vert_max_ctf_vol = src_vol[0]["vertno"][[max_vert_idx]]

# plot them
brain_psf_vol = stc_psf_vol.plot_3d(
    "sample",
    src=forward_vol["src"],
    views="ven",
    subjects_dir=subjects_dir,
    volume_options=dict(alpha=0.5),
)
brain_psf_vol.add_text(0.1, 0.9, "Volumetric sLORETA PSF", "title", font_size=16)
brain_psf_vol.add_foci(
    verttrue_vol, coords_as_verts=True, scale_factor=1, hemi="vol", color="green"
)
brain_psf_vol.add_foci(
    vert_max_ctf_vol,
    coords_as_verts=True,
    scale_factor=1.25,
    hemi="vol",
    color="black",
    alpha=0.3,
)