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

import numpy as np
import pytest
from numpy.testing import assert_allclose, assert_array_almost_equal, assert_array_equal
from scipy.signal import welch

from mne.time_frequency import psd_array_multitaper, psd_array_welch
from mne.time_frequency.multitaper import _psd_from_mt
from mne.time_frequency.psd import _median_biases
from mne.utils import catch_logging


def test_psd_nan():
    """Test handling of NaN in psd_array_welch."""
    n_samples, n_fft, n_overlap = 2048, 1024, 512
    x = np.random.RandomState(0).randn(1, n_samples)
    psds, freqs = psd_array_welch(
        x[:, : n_fft + n_overlap], float(n_fft), n_fft=n_fft, n_overlap=n_overlap
    )
    x[:, n_fft + n_overlap :] = np.nan  # what Raw.get_data() will give us
    psds_2, freqs_2 = psd_array_welch(x, float(n_fft), n_fft=n_fft, n_overlap=n_overlap)
    assert_allclose(freqs, freqs_2)
    assert_allclose(psds, psds_2)
    # 1-d
    psds_2, freqs_2 = psd_array_welch(
        x[0], float(n_fft), n_fft=n_fft, n_overlap=n_overlap
    )
    assert_allclose(freqs, freqs_2)
    assert_allclose(psds[0], psds_2)
    # defaults
    with catch_logging() as log:
        psd_array_welch(x, float(n_fft), verbose="debug")
    log = log.getvalue()
    assert "using 256-point FFT on 256 samples with 0 overlap" in log
    assert "hamming window" in log


def test_bad_annot_handling():
    """Make sure results equivalent with/without Annotations."""
    n_per_seg = 256
    n_chan = 3
    n_times = 5 * n_per_seg
    x = np.random.default_rng(seed=42).standard_normal(size=(n_chan, n_times))
    want = psd_array_welch(x, sfreq=100)
    # now simulate an annotation that breaks up the array into unequal spans. Using
    # `n_per_seg` as the cut point is unrealistic/idealized, but it allows us to test
    # whether we get results ~identical to `want` (which we should in this edge case)
    x2 = np.concatenate(
        (x[..., :n_per_seg], np.full((n_chan, 1), np.nan), x[..., n_per_seg:]), axis=-1
    )
    got = psd_array_welch(x2, sfreq=100)
    # freqs should be identical
    np.testing.assert_array_equal(got[1], want[1])
    # powers should be very very close
    np.testing.assert_allclose(got[0], want[0], rtol=1e-15, atol=0)


def _make_psd_data():
    """Make noise data with sinusoids in 2 out of 7 channels."""
    rng = np.random.default_rng(0)
    n_chan, n_times, sfreq = 7, 8000, 1000
    data = 0.1 * rng.random((n_chan, n_times))
    times = np.arange(n_times) / sfreq
    sinusoid_freqs = [8.0, 50.0]
    chs_with_sinusoids = [0, 1]
    for ix, freq in zip(chs_with_sinusoids, sinusoid_freqs):
        data[ix, :] += 2 * np.sin(np.pi * 2.0 * freq * times)
    return data, sfreq, sinusoid_freqs


@pytest.mark.parametrize(
    "psd_func, psd_kwargs",
    [
        (psd_array_welch, dict(n_fft=128, window="hann")),
        (psd_array_multitaper, dict(low_bias=True)),
    ],
)
def test_psd(psd_func, psd_kwargs):
    """Tests the welch and multitaper PSD."""
    data, sfreq, sinusoid_freqs = _make_psd_data()
    # prepare kwargs
    psd_kwargs.update(dict(fmin=2, fmax=70, verbose="debug"))
    # compute PSD and test basic conformity
    with catch_logging() as log:
        psds, freqs = psd_func(data, sfreq, **psd_kwargs)
    if psd_func is psd_array_welch:
        log = log.getvalue()
        n_fft = psd_kwargs["n_fft"]
        assert f"{n_fft}-point FFT on {n_fft} samples with 0 overl" in log
        assert "hann window" in log
    assert psds.shape == (data.shape[0], len(freqs))
    assert np.sum(freqs < 0) == 0
    assert np.sum(psds < 0) == 0
    # Is power found where it should be?
    ixs_max = np.argmax(psds, axis=1)
    for ixmax, ifreq in zip(ixs_max, sinusoid_freqs):
        # Find nearest frequency to the "true" freq
        ixtrue = np.argmin(np.abs(ifreq - freqs))
        assert np.abs(ixmax - ixtrue) < 2


def test_psd_array_welch_nperseg_kwarg():
    """Test n_per_seg and padding in psd_array_welch()."""
    data, sfreq, _ = _make_psd_data()
    # prepare kwargs
    kwargs = dict(fmin=2, fmax=70, n_per_seg=128)
    # test n_per_seg in psd_array_welch (and padding)
    psds1, freqs1 = psd_array_welch(data, sfreq, n_fft=128, **kwargs)
    psds2, freqs2 = psd_array_welch(data, sfreq, n_fft=256, **kwargs)
    assert len(freqs1) == np.floor(len(freqs2) / 2.0)
    assert psds1.shape[-1] == np.floor(psds2.shape[-1] / 2.0)
    # test bad n_fft
    with pytest.raises(ValueError, match="n_fft is not allowed to be > n_tim"):
        kwargs.update(n_per_seg=None)
        bad_n_fft = int(data.shape[-1] * 1.1)
        psd_array_welch(data, sfreq, n_fft=bad_n_fft, **kwargs)
    # test bad n_overlap
    with pytest.raises(ValueError, match="n_overlap cannot be greater"):
        kwargs.update(n_per_seg=64)
        psd_array_welch(data, sfreq, n_fft=128, n_overlap=90, **kwargs)
    # test bad fmin/fmax
    with pytest.raises(ValueError, match="No frequencies found"):
        psd_array_welch(data, sfreq, fmin=10, fmax=1)


def test_complex_multitaper():
    """Test complex-valued multitaper output."""
    data, sfreq, _ = _make_psd_data()
    psd_complex, freq_complex, weights = psd_array_multitaper(
        data[:4, :500], sfreq, output="complex"
    )
    psd, freq = psd_array_multitaper(data[:4, :500], sfreq, output="power")
    assert_array_equal(freq_complex, freq)
    assert psd_complex.ndim == 3  # channels x tapers x freqs
    psd_from_complex = _psd_from_mt(psd_complex, weights)
    assert_allclose(psd_from_complex, psd)


# Copied from SciPy
def _median_bias(n):
    ii_2 = 2 * np.arange(1.0, (n - 1) // 2 + 1)
    return 1 + np.sum(1.0 / (ii_2 + 1) - 1.0 / ii_2)


@pytest.mark.parametrize("crop", (False, True))
def test_psd_array_welch_average_kwarg(crop):
    """Test `average` kwarg of psd_array_welch()."""
    data, sfreq, _ = _make_psd_data()
    # prepare kwargs
    n_per_seg = 32
    kwargs = dict(fmin=0, fmax=np.inf, n_fft=64, n_per_seg=n_per_seg, n_overlap=0)
    # optionally crop data by n_per_seg so that we are sure to test both an
    # odd number and an even number of estimates (for median bias)
    if crop:
        data = data[..., :-n_per_seg]
    # run with average=mean/median/None
    psds_mean, freqs_mean = psd_array_welch(data, sfreq, average="mean", **kwargs)
    psds_median, freqs_median = psd_array_welch(data, sfreq, average="median", **kwargs)
    psds_unagg, freqs_unagg = psd_array_welch(data, sfreq, average=None, **kwargs)
    # Frequencies should be equal across all "average" types, as we feed in
    # the exact same data.
    assert_array_equal(freqs_mean, freqs_median)
    assert_array_equal(freqs_mean, freqs_unagg)
    # For `average=None`, the last dimension contains the un-aggregated
    # segments.
    assert psds_mean.shape == psds_median.shape
    assert psds_mean.shape == psds_unagg.shape[:-1]
    assert_array_equal(psds_mean, psds_unagg.mean(axis=-1))
    # Compare with manual median calculation (_median_bias copied from SciPy)
    bias = _median_bias(psds_unagg.shape[-1])
    assert_allclose(psds_median, np.median(psds_unagg, axis=-1) / bias)
    # check shape of unagg
    n_chan, n_times = data.shape
    n_freq = len(freqs_unagg)
    n_segs = np.ceil(n_times / n_per_seg).astype(int)
    assert n_segs % 2 == (1 if crop else 0)
    assert psds_unagg.shape == (n_chan, n_freq, n_segs)


@pytest.mark.parametrize("n", (2, 3, 5, 8, 12, 13, 14, 15))
def test_median_biases(n):
    """Test vectorization of median_biases."""
    want_biases = np.concatenate(
        ([1.0, 1.0], [_median_bias(ii) for ii in range(2, n + 1)])
    )
    got_biases = _median_biases(n)
    assert_allclose(want_biases, got_biases)
    assert_allclose(got_biases[n], _median_bias(n))
    assert_allclose(got_biases[:3], 1.0)


@pytest.mark.slowtest
def test_compares_psd():
    """Test PSD estimation on raw for plt.psd and scipy.signal.welch."""
    data, sfreq, _ = _make_psd_data()
    n_fft = 2048
    fmin, fmax = 2, 70
    # Compute PSD with psd_array_welch
    psds_mne, freqs_mne = psd_array_welch(
        data, sfreq, fmin=fmin, fmax=fmax, n_fft=n_fft
    )
    # Compute psds with scipy.signal.welch
    freqs_scipy, psds_scipy = welch(
        data, fs=sfreq, nperseg=n_fft, noverlap=0, window="hamming"
    )
    # restrict to the relevant frequencies
    mask = (freqs_scipy >= fmin) & (freqs_scipy <= fmax)
    freqs_scipy = freqs_scipy[mask]
    psds_scipy = psds_scipy[:, mask]
    # make sure they match
    assert_array_almost_equal(psds_mne, psds_scipy)
    assert_array_equal(freqs_mne, freqs_scipy)
    assert psds_mne.shape == (data.shape[0], len(freqs_mne))
    assert psds_scipy.shape == (data.shape[0], len(freqs_scipy))
    assert np.sum(freqs_mne < 0) == 0
    assert np.sum(freqs_scipy < 0) == 0
    assert np.sum(psds_mne < 0) == 0
    assert np.sum(psds_scipy < 0) == 0


def test_psd_array_welch_n_jobs():
    """Test that n_jobs works even with more jobs than channels."""
    data = np.zeros((1, 2048))
    psd_array_welch(data, 1024, n_jobs=1)
    psd_array_welch(data, 1024, n_jobs=2)