# Authors: Chris Holdgraf <choldgraf@gmail.com>
#
# License: BSD (3-clause)
import os.path as op

import pytest
import numpy as np

from numpy.testing import assert_array_equal, assert_allclose, assert_equal

from mne import io, pick_types
from mne.fixes import einsum
from mne.utils import requires_version, requires_sklearn, run_tests_if_main
from mne.decoding import ReceptiveField, TimeDelayingRidge
from mne.decoding.receptive_field import (_delay_time_series, _SCORERS,
                                          _times_to_delays, _delays_to_slice)
from mne.decoding.time_delaying_ridge import (_compute_reg_neighbors,
                                              _compute_corrs)


data_dir = op.join(op.dirname(__file__), '..', '..', 'io', 'tests', 'data')
raw_fname = op.join(data_dir, 'test_raw.fif')
event_name = op.join(data_dir, 'test-eve.fif')

tmin, tmax = -0.1, 0.5
event_id = dict(aud_l=1, vis_l=3)

# Loading raw data
raw = io.read_raw_fif(raw_fname, preload=True)
picks = pick_types(raw.info, meg=True, stim=False, ecg=False,
                   eog=False, exclude='bads')
picks = picks[:2]

n_jobs_test = (1, 'cuda')


def test_compute_reg_neighbors():
    """Test fast calculation of laplacian regularizer."""
    for reg_type in (
            ('ridge', 'ridge'),
            ('ridge', 'laplacian'),
            ('laplacian', 'ridge'),
            ('laplacian', 'laplacian')):
        for n_ch_x, n_delays in (
                (1, 1), (1, 2), (2, 1), (1, 3), (3, 1), (1, 4), (4, 1),
                (2, 2), (2, 3), (3, 2), (3, 3),
                (2, 4), (4, 2), (3, 4), (4, 3), (4, 4),
                (5, 4), (4, 5), (5, 5),
                (20, 9), (9, 20)):
            for normed in (True, False):
                reg_direct = _compute_reg_neighbors(
                    n_ch_x, n_delays, reg_type, 'direct', normed=normed)
                reg_csgraph = _compute_reg_neighbors(