# Authors: The MNE-Python contributors.
# License: BSD-3-Clause
# Copyright the MNE-Python contributors.

import sys

import numpy as np
import pytest
from numpy.testing import assert_array_almost_equal, assert_array_equal

pytest.importorskip("sklearn")

from sklearn.pipeline import Pipeline

from mne import create_info, io
from mne.decoding import CSP
from mne.decoding.ssd import SSD
from mne.filter import filter_data
from mne.time_frequency import psd_array_welch

freqs_sig = 9, 12
freqs_noise = 8, 13


def simulate_data(
    freqs_sig=(9, 12),
    n_trials=100,
    n_channels=20,
    n_samples=500,
    samples_per_second=250,
    n_components=5,
    SNR=0.05,
    random_state=42,
):
    """Simulate data according to an instantaneous mixin model.

    Data are simulated in the statistical source space, where n=n_components
    sources contain the peak of interest.
    """
    rng = np.random.RandomState(random_state)

    filt_params_signal = dict(
        l_freq=freqs_sig[0],
        h_freq=freqs_sig[1],
        l_trans_bandwidth=1,
        h_trans_bandwidth=1,
        fir_design="firwin",
    )

    # generate an orthogonal mixin matrix
    mixing_mat = np.linalg.svd(rng.randn(n_channels, n_channels))[0]
    # define sources
    S_s = rng.randn(n_trials * n_samples, n_components)
    # filter source in the specific freq. band of interest
    S_s = filter_data(S_s.T, samples_per_second, **filt_params_signal).T
    S_n = rng.randn(n_trials * n_samples, n_channels - n_components)
    S = np.hstack((S_s, S_n))
    # mix data
    X_s = np.dot(mixing_mat[:, :n_components], S_s.T).T
    X_n = np.dot(mixing_mat[:, n_components:], S_n.T).T
    # add noise
    X_s = X_s / np.linalg.norm(X_s, "fro")
    X_n = X_n / np.linalg.norm(X_n, "fro")
    X = SNR * X_s + (1 - SNR) * X_n
    X = X.T
    S = S.T
    return X, mixing_mat, S


@pytest.mark.slowtest
def test_ssd():
    """Test Common Spatial Patterns algorithm on raw data."""
    X, A, S = simulate_data()
    sf = 250
    n_channels = X.shape[0]
    info = create_info(ch_names=n_channels, sfreq=sf, ch_types="eeg")
    n_components_true = 5

    # Init
    filt_params_signal = dict(
        l_freq=freqs_sig[0],
        h_freq=freqs_sig[1],
        l_trans_bandwidth=1,
        h_trans_bandwidth=1,
    )
    filt_params_noise = dict(
        l_freq=freqs_noise[0],
        h_freq=freqs_noise[1],
        l_trans_bandwidth=1,
        h_trans_bandwidth=1,
    )
    ssd = SSD(info, filt_params_signal, filt_params_noise)
    # freq no int
    freq = "foo"
    filt_params_signal = dict(
        l_freq=freq, h_freq=freqs_sig[1], l_trans_bandwidth=1, h_trans_bandwidth=1
    )
    filt_params_noise = dict(
        l_freq=freqs_noise[0],
        h_freq=freqs_noise[1],
        l_trans_bandwidth=1,
        h_trans_bandwidth=1,
    )
    with pytest.raises(TypeError, match="must be an instance "):
        ssd = SSD(info, filt_params_signal, filt_params_noise)

    # Wrongly specified noise band
    freq = 2
    filt_params_signal = dict(
        l_freq=freq, h_freq=freqs_sig[1], l_trans_bandwidth=1, h_trans_bandwidth=1
    )
    filt_params_noise = dict(
        l_freq=freqs_noise[0],
        h_freq=freqs_noise[1],
        l_trans_bandwidth=1,
        h_trans_bandwidth=1,
    )
    with pytest.raises(ValueError, match="Wrongly specified "):
        ssd = SSD(info, filt_params_signal, filt_params_noise)

    # filt param no dict
    filt_params_signal = freqs_sig
    filt_params_noise = freqs_noise
    with pytest.raises(ValueError, match="must be defined"):
        ssd = SSD(info, filt_params_signal, filt_params_noise)

    # Data type
    filt_params_signal = dict(
        l_freq=freqs_sig[0],
        h_freq=freqs_sig[1],
        l_trans_bandwidth=1,
        h_trans_bandwidth=1,
    )
    filt_params_noise = dict(
        l_freq=freqs_noise[0],
        h_freq=freqs_noise[1],
        l_trans_bandwidth=1,
        h_trans_bandwidth=1,
    )
    ssd = SSD(info, filt_params_signal, filt_params_noise)
    raw = io.RawArray(X, info)

    pytest.raises(TypeError, ssd.fit, raw)

    # check non-boolean return_filtered
    with pytest.raises(ValueError, match="return_filtered"):
        ssd = SSD(info, filt_params_signal, filt_params_noise, return_filtered=0)

    # check non-boolean sort_by_spectral_ratio
    with pytest.raises(ValueError, match="sort_by_spectral_ratio"):
        ssd = SSD(info, filt_params_signal, filt_params_noise, sort_by_spectral_ratio=0)

    # More than 1 channel type
    ch_types = np.reshape([["mag"] * 10, ["eeg"] * 10], n_channels)
    info_2 = create_info(ch_names=n_channels, sfreq=sf, ch_types=ch_types)

    ssd = SSD(info_2, filt_params_signal, filt_params_noise)
    with pytest.raises(ValueError, match="At this point SSD"):
        ssd.fit(X)

    # Number of channels
    info_3 = create_info(ch_names=n_channels + 1, sfreq=sf, ch_types="eeg")
    ssd = SSD(info_3, filt_params_signal, filt_params_noise)
    pytest.raises(ValueError, ssd.fit, X)

    # Fit
    n_components = 10
    ssd = SSD(info, filt_params_signal, filt_params_noise, n_components=n_components)

    # Call transform before fit
    pytest.raises(AttributeError, ssd.transform, X)

    # Check outputs
    ssd.fit(X)

    assert ssd.filters_.shape == (n_channels, n_channels)
    assert ssd.patterns_.shape == (n_channels, n_channels)

    # Transform
    X_ssd = ssd.fit_transform(X)
    assert X_ssd.shape[0] == n_components
    # back and forward
    ssd = SSD(
        info,
        filt_params_signal,
        filt_params_noise,
        n_components=None,
        sort_by_spectral_ratio=False,
    )
    ssd.fit(X)
    X_denoised = ssd.apply(X)
    assert_array_almost_equal(X_denoised, X)
    # denoised by low-rank-factorization
    ssd = SSD(
        info,
        filt_params_signal,
        filt_params_noise,
        n_components=n_components,
        sort_by_spectral_ratio=True,
    )
    ssd.fit(X)
    X_denoised = ssd.apply(X)
    assert np.linalg.matrix_rank(X_denoised) == n_components

    # Power ratio ordering
    ssd = SSD(
        info,
        filt_params_signal,
        filt_params_noise,
        n_components=None,
        sort_by_spectral_ratio=False,
    )
    ssd.fit(X)
    spec_ratio, sorter_spec = ssd.get_spectral_ratio(ssd.transform(X))
    # since we now that the number of true components is 5, the relative
    # difference should be low for the first 5 components and then increases
    index_diff = np.argmax(-np.diff(spec_ratio))
    assert index_diff == n_components_true - 1
    # Check detected peaks
    # fit ssd
    n_components = n_components_true
    filt_params_signal = dict(
        l_freq=freqs_sig[0],
        h_freq=freqs_sig[1],
        l_trans_bandwidth=1,
        h_trans_bandwidth=1,
    )
    filt_params_noise = dict(
        l_freq=freqs_noise[0],
        h_freq=freqs_noise[1],
        l_trans_bandwidth=1,
        h_trans_bandwidth=1,
    )
    ssd = SSD(
        info,
        filt_params_signal,
        filt_params_noise,
        n_components=n_components,
        sort_by_spectral_ratio=False,
    )
    ssd.fit(X)

    out = ssd.transform(X)
    psd_out, _ = psd_array_welch(out[0], sfreq=250, n_fft=250)
    psd_S, _ = psd_array_welch(S[0], sfreq=250, n_fft=250)
    corr = np.abs(np.corrcoef((psd_out, psd_S))[0, 1])
    assert np.abs(corr) > 0.95
    # Check pattern estimation
    # Since there is no exact ordering of the recovered patterns
    # a pair-wise greedy search will be done
    error = list()
    for ii in range(n_channels):
        corr = np.abs(np.corrcoef(ssd.patterns_[ii, :].T, A[:, 0])[0, 1])
        error.append(1 - corr)
        min_err = np.min(error)
    assert min_err < 0.3  # threshold taken from SSD original paper


def test_ssd_epoched_data():
    """Test Common Spatial Patterns algorithm on epoched data.

    Compare the outputs when raw data is used.
    """
    X, A, S = simulate_data(n_trials=100, n_channels=20, n_samples=500)
    sf = 250
    n_channels = X.shape[0]
    info = create_info(ch_names=n_channels, sfreq=sf, ch_types="eeg")
    n_components_true = 5

    # Build epochs as sliding windows over the continuous raw file

    # Epoch length is 1 second
    X_e = np.reshape(X, (100, 20, 500))

    # Fit
    filt_params_signal = dict(
        l_freq=freqs_sig[0],
        h_freq=freqs_sig[1],
        l_trans_bandwidth=4,
        h_trans_bandwidth=4,
    )
    filt_params_noise = dict(
        l_freq=freqs_noise[0],
        h_freq=freqs_noise[1],
        l_trans_bandwidth=4,
        h_trans_bandwidth=4,
    )

    # ssd on epochs
    ssd_e = SSD(info, filt_params_signal, filt_params_noise)
    ssd_e.fit(X_e)
    # ssd on raw
    ssd = SSD(info, filt_params_signal, filt_params_noise)
    ssd.fit(X)

    # Check if the 5 first 5 components are the same for both
    _, sorter_spec_e = ssd_e.get_spectral_ratio(ssd_e.transform(X_e))
    _, sorter_spec = ssd.get_spectral_ratio(ssd.transform(X))
    assert_array_equal(
        sorter_spec_e[:n_components_true], sorter_spec[:n_components_true]
    )


def test_ssd_pipeline():
    """Test if SSD works in a pipeline."""
    sf = 250
    X, A, S = simulate_data(n_trials=100, n_channels=20, n_samples=500)
    X_e = np.reshape(X, (100, 20, 500))
    # define bynary random output
    y = np.random.randint(2, size=100)

    info = create_info(ch_names=20, sfreq=sf, ch_types="eeg")

    filt_params_signal = dict(
        l_freq=freqs_sig[0],
        h_freq=freqs_sig[1],
        l_trans_bandwidth=4,
        h_trans_bandwidth=4,
    )
    filt_params_noise = dict(
        l_freq=freqs_noise[0],
        h_freq=freqs_noise[1],
        l_trans_bandwidth=4,
        h_trans_bandwidth=4,
    )
    ssd = SSD(info, filt_params_signal, filt_params_noise)
    csp = CSP()
    pipe = Pipeline([("SSD", ssd), ("CSP", csp)])
    pipe.set_params(SSD__n_components=5)
    pipe.set_params(CSP__n_components=2)
    out = pipe.fit_transform(X_e, y)
    assert out.shape == (100, 2)
    assert pipe.get_params()["SSD__n_components"] == 5


def test_sorting():
    """Test sorting learning during training."""
    X, _, _ = simulate_data(n_trials=100, n_channels=20, n_samples=500)
    # Epoch length is 1 second
    X = np.reshape(X, (100, 20, 500))
    # split data
    Xtr, Xte = X[:80], X[80:]
    sf = 250
    n_channels = Xtr.shape[1]
    info = create_info(ch_names=n_channels, sfreq=sf, ch_types="eeg")

    filt_params_signal = dict(
        l_freq=freqs_sig[0],
        h_freq=freqs_sig[1],
        l_trans_bandwidth=4,
        h_trans_bandwidth=4,
    )
    filt_params_noise = dict(
        l_freq=freqs_noise[0],
        h_freq=freqs_noise[1],
        l_trans_bandwidth=4,
        h_trans_bandwidth=4,
    )

    # check sort_by_spectral_ratio set to False
    ssd = SSD(
        info,
        filt_params_signal,
        filt_params_noise,
        n_components=None,
        sort_by_spectral_ratio=False,
    )
    ssd.fit(Xtr)
    _, sorter_tr = ssd.get_spectral_ratio(ssd.transform(Xtr))
    _, sorter_te = ssd.get_spectral_ratio(ssd.transform(Xte))
    assert any(sorter_tr != sorter_te)

    # check sort_by_spectral_ratio set to True
    ssd = SSD(
        info,
        filt_params_signal,
        filt_params_noise,
        n_components=None,
        sort_by_spectral_ratio=True,
    )
    ssd.fit(Xtr)

    # check sorters
    sorter_in = ssd.sorter_spec
    ssd = SSD(
        info,
        filt_params_signal,
        filt_params_noise,
        n_components=None,
        sort_by_spectral_ratio=False,
    )
    ssd.fit(Xtr)
    _, sorter_out = ssd.get_spectral_ratio(ssd.transform(Xtr))

    assert all(sorter_in == sorter_out)


def test_return_filtered():
    """Test return filtered option."""
    # Check return_filtered
    # Simulated more noise data and with broader frequency than the desired
    X, _, _ = simulate_data(SNR=0.9, freqs_sig=[4, 13])
    sf = 250
    n_channels = X.shape[0]
    info = create_info(ch_names=n_channels, sfreq=sf, ch_types="eeg")

    filt_params_signal = dict(
        l_freq=freqs_sig[0],
        h_freq=freqs_sig[1],
        l_trans_bandwidth=1,
        h_trans_bandwidth=1,
    )
    filt_params_noise = dict(
        l_freq=freqs_noise[0],
        h_freq=freqs_noise[1],
        l_trans_bandwidth=1,
        h_trans_bandwidth=1,
    )

    # return filtered to true
    ssd = SSD(
        info,
        filt_params_signal,
        filt_params_noise,
        sort_by_spectral_ratio=False,
        return_filtered=True,
    )
    ssd.fit(X)

    out = ssd.transform(X)
    psd_out, freqs = psd_array_welch(out[0], sfreq=250, n_fft=250)
    freqs_up = int(freqs[psd_out > 0.5][0]), int(freqs[psd_out > 0.5][-1])
    assert freqs_up == freqs_sig

    # return filtered to false
    ssd = SSD(
        info,
        filt_params_signal,
        filt_params_noise,
        sort_by_spectral_ratio=False,
        return_filtered=False,
    )
    ssd.fit(X)

    out = ssd.transform(X)
    psd_out, freqs = psd_array_welch(out[0], sfreq=250, n_fft=250)
    freqs_up = int(freqs[psd_out > 0.5][0]), int(freqs[psd_out > 0.5][-1])
    assert freqs_up != freqs_sig


def test_non_full_rank_data():
    """Test that the method works with non-full rank data."""
    n_channels = 10
    X, _, _ = simulate_data(SNR=0.9, freqs_sig=[4, 13], n_channels=n_channels)
    info = create_info(ch_names=n_channels, sfreq=250, ch_types="eeg")

    filt_params_signal = dict(
        l_freq=freqs_sig[0],
        h_freq=freqs_sig[1],
        l_trans_bandwidth=1,
        h_trans_bandwidth=1,
    )
    filt_params_noise = dict(
        l_freq=freqs_noise[0],
        h_freq=freqs_noise[1],
        l_trans_bandwidth=1,
        h_trans_bandwidth=1,
    )

    # Make data non-full rank
    rank = 5
    X[rank:] = X[:rank]  # an extreme example, but a valid one
    assert np.linalg.matrix_rank(X) == rank

    ssd = SSD(info, filt_params_signal, filt_params_noise)
    if sys.platform == "darwin":
        pytest.skip("Unknown linalg bug (Accelerate?)")
    ssd.fit(X)