# -*- coding: utf-8 -*-
"""
Plot point-spread functions (PSFs) and cross-talk functions (CTFs)
==================================================================

Visualise PSF and CTF at one vertex for sLORETA.
"""
# Authors: Olaf Hauk <olaf.hauk@mrc-cbu.cam.ac.uk>
#          Alexandre Gramfort <alexandre.gramfort@inria.fr>
#
# License: BSD (3-clause)

import mne
mne.viz.set_3d_backend('pyvista')
from mne.datasets import sample
from mne.minimum_norm import (make_inverse_resolution_matrix, get_cross_talk,
                              get_point_spread)
import numpy as np
print(__doc__)

data_path = sample.data_path()
subjects_dir = data_path + '/subjects/'
fname_fwd = data_path + '/MEG/sample/sample_audvis-meg-eeg-oct-6-fwd.fif'
fname_cov = data_path + '/MEG/sample/sample_audvis-cov.fif'
fname_evo = data_path + '/MEG/sample/sample_audvis-ave.fif'

# read forward solution
forward = mne.read_forward_solution(fname_fwd)
# forward operator with fixed source orientations
forward = mne.convert_forward_solution(forward, surf_ori=True,
                                       force_fixed=True)

# noise covariance matrix
noise_cov = mne.read_cov(fname_cov)

# evoked data for info
evoked = mne.read_evokeds(fname_evo, 0)

# make inverse operator from forward solution
# free source orientation
inverse_operator = mne.minimum_norm.make_inverse_operator(
    info=evoked.info, forward=forward, noise_cov=noise_cov, loose=0.,
    depth=None)

# regularisation parameter
snr = 3.0
lambda2 = 1.0 / snr ** 2
method = 'dSPM'  # can be 'MNE' or 'sLORETA'

# compute resolution matrix for sLORETA
rm_lor = make_inverse_resolution_matrix(forward, inverse_operator,
                                        method=method, lambda2=lambda2)


# read the labels from sample data 
labels = mne.read_labels_from_annot(subject='sample', parc='aparc.a2009s', subjects_dir=subjects_dir)
# label to be used
label_pick = 'S_temporal_transverse-rh'  # auditory sources
label_names = [l.name for l in labels]
# lets get the label idx 
idx = label_names.index(label_pick)
label = labels[idx]  
hemi = 'rh'   # define hemisphere
if hemi == 'lh':
    hemi_value = 0
else:
    hemi_value = 1

# compute the centre dipole location of the label
hemi_to_ind = {'lh': 0, 'rh': 1}
# Restrict the eligible vertices to be those on the surface under
# consideration and within the label.
surf_vertices = forward['src'][hemi_to_ind[label.hemi]]['vertno']
restrict_verts = np.intersect1d(surf_vertices, label.vertices)
vertno_hemi = forward['src'][hemi_value]['vertno']
centre_vertex = label.center_of_mass(subject='sample', subjects_dir=subjects_dir, restrict_vertices=restrict_verts, surf='white')

# get the source idx
sources = [np.where(vertno_hemi == centre_vertex)[0][0]]

# get PSF for dSPM at one vertex

stc_psf = get_point_spread(rm_lor, forward['src'], sources, norm=True)

##############################################################################
# Visualise
# ---------

# Which vertex corresponds to selected source
vertno = forward['src'][hemi_value]['vertno']
verttrue = [vertno[sources[0]]]  # just one vertex

# find vertices with maxima in PSF and CTF
vert_max_psf = vertno[stc_psf.data.argmax()]

brain_psf = stc_psf.plot('sample', 'inflated', hemi='split', subjects_dir=subjects_dir,
                         figure=1)
brain_psf.show_view('ventral')
brain_psf.add_text(0.1, 0.9, '%s PSF'%method, 'title', font_size=16)

# True source location for PSF
brain_psf.add_foci(verttrue, coords_as_verts=True, scale_factor=1., hemi=hemi,
                   color='green')

# Maximum of PSF
brain_psf.add_foci(vert_max_psf, coords_as_verts=True, scale_factor=1.,
                   hemi=hemi, color='black')

print('The green spheres indicate the true source location, and the black '
      'spheres the maximum of the distribution.')