"""
=====================================================================
Compute Phase Slope Index (PSI) in source space for a visual stimulus
=====================================================================

This example demonstrates how the phase slope index (PSI) :footcite:`NolteEtAl2008` can
be computed in source space based on single trial dSPM source estimates.

The :func:`~mne_connectivity.phase_slope_index` function for computing PSI shown here
gives connectivity estimates over epochs. The
:func:`~mne_connectivity.phase_slope_index_time` function is also available to compute
PSI over time, allowing for per-epoch connectivity estimation.

In addition, the example shows advanced usage of the connectivity estimation routines
by first extracting a label time course for each epoch and then combining the label time
course with the single trial source estimates to compute the connectivity.

The result clearly shows how the activity in the visual label precedes more widespread
activity (as a postivive PSI means the label time course is leading).
"""
# Author: Martin Luessi <mluessi@nmr.mgh.harvard.edu>
#
# License: BSD (3-clause)

import mne
import numpy as np
from mne.datasets import sample
from mne.minimum_norm import apply_inverse_epochs, read_inverse_operator

from mne_connectivity import phase_slope_index, seed_target_indices

print(__doc__)

data_path = sample.data_path()
subjects_dir = data_path / "subjects"
fname_inv = data_path / "MEG/sample/sample_audvis-meg-oct-6-meg-inv.fif"
fname_raw = data_path / "MEG/sample/sample_audvis_filt-0-40_raw.fif"
fname_event = data_path / "MEG/sample/sample_audvis_filt-0-40_raw-eve.fif"
fname_label = data_path / "MEG/sample/labels/Vis-lh.label"

event_id, tmin, tmax = 4, -0.2, 0.5
method = "dSPM"  # use dSPM method (could also be MNE or sLORETA)

# Load data
inverse_operator = read_inverse_operator(fname_inv)
raw = mne.io.read_raw_fif(fname_raw, preload=True)
events = mne.read_events(fname_event)

# Pick MEG channels
picks = mne.pick_types(
    raw.info, meg=True, eeg=False, stim=False, eog=True, exclude="bads"
)

# Read epochs
epochs = mne.Epochs(
    raw,
    events,
    event_id,
    tmin,
    tmax,
    picks=picks,
    baseline=(None, 0),
    reject=dict(mag=4e-12, grad=4000e-13, eog=150e-6),
)

# %%
# Compute inverse solution for each epoch. Note that since we are passing the output to
# both :func:`~mne.extract_label_time_course` and the
# :func:`~mne_connectivity.phase_slope_index` functions, we have to use
# ``return_generator=False``, since it is only possible to iterate over generators once.

snr = 1.0  # use lower SNR for single epochs
lambda2 = 1.0 / snr**2
stcs = apply_inverse_epochs(
    epochs, inverse_operator, lambda2, method, pick_ori="normal", return_generator=True
)

# %%
# Now, we generate seed time series by averaging the activity in the left visual cortex.

label = mne.read_label(fname_label)
src = inverse_operator["src"]  # the source space used
seed_ts = mne.extract_label_time_course(
    stcs, label, src, mode="mean_flip", verbose="error"
)

# Combine the seed time course with the source estimates. There will be a total
# of 7500 signals:
# index 0: time course extracted from label
# index 1..7499: dSPM source space time courses
stcs = apply_inverse_epochs(
    epochs, inverse_operator, lambda2, method, pick_ori="normal", return_generator=True
)
comb_ts = list(zip(seed_ts, stcs))

# %%
# Construct indices to estimate connectivity between the label time course and all
# source space time courses.

vertices = [src[i]["vertno"] for i in range(2)]
n_signals_tot = 1 + len(vertices[0]) + len(vertices[1])

indices = seed_target_indices([0], np.arange(1, n_signals_tot))

# Compute the PSI in the frequency range 10Hz-20Hz. We exclude the baseline
# period from the connectivity estimation.
fmin = 10.0
fmax = 20.0
tmin_con = 0.0
sfreq = epochs.info["sfreq"]  # the sampling frequency

psi = phase_slope_index(
    comb_ts,
    mode="multitaper",
    indices=indices,
    sfreq=sfreq,
    fmin=fmin,
    fmax=fmax,
    tmin=tmin_con,
)

# %%
# Generate a :class:`~mne.SourceEstimate` with the PSI. This is simple since we used a
# single seed (inspect the ``indices`` variable to see how the PSI scores are arranged
# in the output).

psi_stc = mne.SourceEstimate(
    psi.get_data(), vertices=vertices, tmin=0, tstep=1, subject="sample"
)

# %%
# Now we can visualize the PSI using the :meth:`~mne.SourceEstimate.plot` method. We use
# a custom colormap to show signed values.

v_max = np.max(np.abs(psi.get_data()))
brain = psi_stc.plot(
    surface="inflated",
    hemi="lh",
    time_label="Phase Slope Index (PSI)",
    subjects_dir=subjects_dir,
    clim=dict(kind="percent", pos_lims=(95, 97.5, 100)),
)
brain.show_view("medial")
brain.add_label(str(fname_label), color="green", alpha=0.7)

###############################################################################
# References
# ----------
# .. footbibliography::
