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

import datetime
import re
from functools import partial

import numpy as np
import pytest
from matplotlib.colors import same_color
from numpy.testing import assert_allclose, assert_array_equal

from mne import (
    Annotations,
    BaseEpochs,
    create_info,
    grand_average,
    make_fixed_length_epochs,
)
from mne.channels import equalize_channels
from mne.io import RawArray
from mne.time_frequency import read_spectrum
from mne.time_frequency.multitaper import _psd_from_mt
from mne.time_frequency.spectrum import (
    EpochsSpectrumArray,
    SpectrumArray,
    combine_spectrum,
)
from mne.utils import _import_h5io_funcs, _record_warnings


def test_compute_psd_errors(raw):
    """Test for expected errors in the .compute_psd() method."""
    with pytest.raises(ValueError, match="must not exceed ½ the sampling"):
        raw.compute_psd(fmax=raw.info["sfreq"] * 0.51)
    with pytest.raises(TypeError, match="unexpected keyword argument foo for"):
        raw.compute_psd(foo=None)
    with pytest.raises(TypeError, match="keyword arguments foo, bar for"):
        raw.compute_psd(foo=None, bar=None)
    raw.set_annotations(Annotations(onset=0.01, duration=0.01, description="bad_foo"))
    with pytest.raises(NotImplementedError, match='Cannot use method="multitaper"'):
        raw.compute_psd(method="multitaper", reject_by_annotation=True)


@pytest.mark.parametrize("method", ("welch", "multitaper"))
@pytest.mark.parametrize(
    (
        "fmin, fmax, tmin, tmax, picks, proj, n_fft, n_overlap, n_per_seg, "
        "average, window, bandwidth, adaptive, low_bias, normalization, output"
    ),
    [
        [
            0,
            np.inf,
            None,
            None,
            None,
            False,
            256,
            0,
            None,
            "mean",
            "hamming",
            None,
            False,
            True,
            "length",
            "power",
        ],  # defaults
        [
            5,
            50,
            1,
            6,
            "grad",
            True,
            128,
            8,
            32,
            "median",
            "triang",
            10,
            True,
            False,
            "full",
            "power",
        ],  # non-defaults (excluding output)
        [
            0,
            np.inf,
            None,
            None,
            None,
            False,
            256,
            0,
            None,
            "mean",
            "hamming",
            None,
            False,
            True,
            "length",
            "complex",
        ],  # complex (testing with non-defaults doesn't increase coverage)
    ],
)
def test_spectrum_params(
    method,
    fmin,
    fmax,
    tmin,
    tmax,
    picks,
    proj,
    n_fft,
    n_overlap,
    n_per_seg,
    average,
    window,
    bandwidth,
    adaptive,
    low_bias,
    normalization,
    output,
    raw,
):
    """Test valid parameter combinations in the .compute_psd() method."""
    kwargs = dict(
        method=method,
        fmin=fmin,
        fmax=fmax,
        tmin=tmin,
        tmax=tmax,
        picks=picks,
        proj=proj,
        output=output,
    )
    if method == "welch":
        kwargs.update(
            n_fft=n_fft,
            n_overlap=n_overlap,
            n_per_seg=n_per_seg,
            average=average,
            window=window,
        )
    else:
        kwargs.update(
            bandwidth=bandwidth,
            adaptive=adaptive,
            low_bias=low_bias,
            normalization=normalization,
        )
    raw.compute_psd(**kwargs)


def test_n_welch_windows(raw):
    """Test computation of welch windows https://mne.discourse.group/t/5734."""
    with raw.info._unlock():
        raw.info["sfreq"] = 999.412109375
    raw.compute_psd(
        method="welch", n_fft=999, n_per_seg=999, n_overlap=250, average=None
    )


def _get_inst(inst, request, *, evoked=None, average_tfr=None):
    # ↓ XXX workaround:
    # ↓ parametrized fixtures are not accessible via request.getfixturevalue
    # ↓ https://github.com/pytest-dev/pytest/issues/4666#issuecomment-456593913
    if inst == "evoked":
        return evoked
    elif inst == "average_tfr":
        return average_tfr
    return request.getfixturevalue(inst)


@pytest.mark.parametrize("inst", ("raw", "epochs_full", "evoked"))
def test_spectrum_io(inst, tmp_path, request, evoked):
    """Test save/load of spectrum objects."""
    pytest.importorskip("h5io")
    h5py = pytest.importorskip("h5py")
    fname = tmp_path / f"{inst}-spectrum.h5"
    inst = _get_inst(inst, request, evoked=evoked)
    if isinstance(inst, BaseEpochs):
        # fake HED-like tags (https://mne.discourse.group/t/10634)
        inst.events[-2:, -1] = 2
        inst.event_id = {"foo/bar": 1, "foo/qux": 2}
    orig = inst.compute_psd()
    if isinstance(inst, BaseEpochs):
        orig = orig["foo"]
    orig.save(fname)
    loaded = read_spectrum(fname)
    assert orig == loaded
    # Only check following for one type
    if not isinstance(inst, BaseEpochs):
        return
    # Test loading with old-style birthday format
    fname_subject_info = tmp_path / "subject-info.h5"
    _, write_hdf5 = _import_h5io_funcs()
    write_hdf5(fname_subject_info, dict(birthday=(2000, 1, 1)), title="subject_info")
    with h5py.File(fname, "r+") as f:
        del f["mnepython/key_info/key_subject_info"]
        f["mnepython/key_info/key_subject_info"] = h5py.ExternalLink(
            fname_subject_info, "subject_info"
        )
    loaded = read_spectrum(fname)
    assert isinstance(loaded.info["subject_info"]["birthday"], datetime.date)
    # Test Spectrum from EpochsSpectrum.average() can be read (gh-13521)
    origavg = orig.average()
    origavg.save(fname, overwrite=True)
    loadedavg = read_spectrum(fname)
    assert origavg == loadedavg


def test_spectrum_copy(raw_spectrum):
    """Test copying Spectrum objects."""
    spect_copy = raw_spectrum.copy()
    assert raw_spectrum == spect_copy
    assert id(raw_spectrum) != id(spect_copy)
    spect_copy._freqs = None
    assert raw_spectrum.freqs is not None


@pytest.mark.parametrize("weights", ["nave", "equal", [1, -1]])
def test_combine_spectrum(raw_spectrum, weights):
    """Test `combine_spectrum()` works."""
    spectrum1 = raw_spectrum.copy()
    spectrum2 = raw_spectrum.copy()
    if weights == "nave":
        spectrum1.nave = 1
        spectrum2.nave = 2
        spectrum2._data *= 2
        new_spectrum = combine_spectrum([spectrum1, spectrum2], weights=weights)
        assert_allclose(new_spectrum.data, spectrum1.data * (5 / 3))
    elif weights == "equal":
        spectrum2._data *= 2
        new_spectrum = combine_spectrum([spectrum1, spectrum2], weights=weights)
        assert_allclose(new_spectrum.data, spectrum1.data * 1.5)
    else:
        new_spectrum = combine_spectrum([spectrum1, spectrum2], weights=weights)
        assert_allclose(new_spectrum.data, 0)


def test_combine_spectrum_error_catch(raw_spectrum):
    """Test `combine_spectrum()` catches errors."""
    # Test bad weights
    with pytest.raises(
        ValueError, match='Weights must be a list of float, or "nave" or "equal"'
    ):
        combine_spectrum([raw_spectrum, raw_spectrum], weights="foo")
    with pytest.raises(
        ValueError, match="Weights must be the same size as all_spectrum"
    ):
        combine_spectrum([raw_spectrum, raw_spectrum], weights=[1, 1, 1])

    # Test bad nave
    with pytest.raises(ValueError, match="The 'nave' attribute is not specified"):
        combine_spectrum([raw_spectrum, raw_spectrum], weights="nave")

    # Test inconsistent channels
    raw_spectrum2 = raw_spectrum.copy()
    raw_spectrum2.drop_channels(raw_spectrum2.ch_names[0])
    with pytest.raises(AssertionError, match=".* do not contain the same channels"):
        combine_spectrum([raw_spectrum, raw_spectrum2], weights="equal")

    # Test inconsistent frequencies
    raw_spectrum2 = raw_spectrum.copy()
    raw_spectrum2._freqs = raw_spectrum2._freqs + 1
    with pytest.raises(AssertionError, match=".* do not contain the same frequencies"):
        combine_spectrum([raw_spectrum, raw_spectrum2], weights="equal")


def test_grand_average(raw_spectrum):
    """Test `grand_average()` works for instances of `BaseSpectrum`."""
    spectrum1 = raw_spectrum.copy()
    spectrum2 = raw_spectrum.copy()
    spectrum2._data *= 2
    new_spectrum = grand_average([spectrum1, spectrum2])
    assert_allclose(new_spectrum.data, spectrum1.data * 1.5)


def test_equalize_channels(raw_spectrum):
    """Test equalization of channels for instances of `BaseSpectrum`."""
    spect1 = raw_spectrum.copy()
    spect2 = spect1.copy().pick(["MEG 0122", "MEG 0111"])
    spect1, spect2 = equalize_channels([spect1, spect2])

    assert spect1.ch_names == ["MEG 0111", "MEG 0122"]
    assert spect2.ch_names == ["MEG 0111", "MEG 0122"]


def test_spectrum_reject_by_annot(raw):
    """Test rejecting by annotation.

    Cannot use raw_spectrum fixture here because we're testing reject_by_annotation in
    .compute_psd() method.
    """
    kw = dict(n_per_seg=512)  # smaller than shortest good span, to avoid warning
    spect_no_annot = raw.compute_psd(**kw)
    raw.set_annotations(Annotations([1, 5], [3, 3], ["test", "test"]))
    spect_benign_annot = raw.compute_psd(**kw)
    raw.annotations.description = np.array(["bad_test", "bad_test"])
    spect_reject_annot = raw.compute_psd(**kw)
    spect_ignored_annot = raw.compute_psd(**kw, reject_by_annotation=False)
    # the only one that should be different is `spect_reject_annot`
    assert spect_no_annot == spect_benign_annot
    assert spect_no_annot == spect_ignored_annot
    assert spect_no_annot != spect_reject_annot


def test_spectrum_bads_exclude(raw):
    """Test bads are not removed unless exclude="bads"."""
    raw.pick("mag")  # get rid of IAS channel
    spect_no_excld = raw.compute_psd()
    spect_with_excld = raw.compute_psd(exclude="bads")
    assert raw.info["bads"] == spect_no_excld.info["bads"]
    assert spect_with_excld.info["bads"] == []
    assert set(raw.ch_names) - set(spect_with_excld.ch_names) == set(raw.info["bads"])


def test_spectrum_getitem_raw(raw_spectrum):
    """Test Spectrum.__getitem__ for Raw-derived spectra."""
    want = raw_spectrum.get_data(slice(1, 3), fmax=7)
    freq_idx = np.searchsorted(raw_spectrum.freqs, 7)
    got = raw_spectrum[1:3, :freq_idx]
    assert_array_equal(want, got)


def test_spectrum_getitem_epochs(epochs_spectrum):
    """Test Spectrum.__getitem__ for Epochs-derived spectra."""
    # testing data has just one epoch, its event_id label is "1"
    want = epochs_spectrum.get_data()
    got = epochs_spectrum["1"].get_data()
    assert_array_equal(want, got)


@pytest.mark.parametrize("method", ("mean", partial(np.std, axis=0)))
def test_epochs_spectrum_average(epochs_spectrum, method):
    """Test EpochsSpectrum.average()."""
    avg_spect = epochs_spectrum.average(method=method)
    assert avg_spect.shape == epochs_spectrum.shape[1:]
    assert avg_spect._dims == ("channel", "freq")  # no 'epoch'
    assert repr(avg_spect).startswith("<Averaged Power Spectrum (from Epochs")


@pytest.mark.parametrize("inst", ("raw_spectrum", "epochs_spectrum", "evoked"))
def test_spectrum_to_data_frame(inst, request, evoked):
    """Test the to_data_frame method for Spectrum."""
    pytest.importorskip("pandas")
    from pandas.testing import assert_frame_equal

    # setup
    is_already_psd = inst in ("raw_spectrum", "epochs_spectrum")
    is_epochs = inst == "epochs_spectrum"
    inst = _get_inst(inst, request, evoked=evoked)
    extra_dim = () if is_epochs else (1,)
    extra_cols = ["freq", "condition", "epoch"] if is_epochs else ["freq"]
    # compute PSD
    spectrum = inst if is_already_psd else inst.compute_psd(exclude="bads")
    n_epo, n_chan, n_freq = extra_dim + spectrum.get_data().shape
    # test wide format
    df_wide = spectrum.to_data_frame()
    n_row, n_col = df_wide.shape
    assert n_row == n_freq
    assert n_col == n_chan + len(extra_cols)
    assert set(spectrum.ch_names + extra_cols) == set(df_wide.columns)
    # test long format
    df_long = spectrum.to_data_frame(long_format=True)
    n_row, n_col = df_long.shape
    assert n_row == n_epo * n_freq * n_chan
    base_cols = ["channel", "ch_type", "value"]
    assert n_col == len(base_cols + extra_cols)
    assert set(base_cols + extra_cols) == set(df_long.columns)
    # test index
    index = extra_cols[-2:]  # ['freq'] or ['condition', 'epoch']
    df = spectrum.to_data_frame(index=index)
    if is_epochs:
        index_tuple = (
            list(spectrum.event_id)[0],  # condition
            spectrum.selection[0],
        )  # epoch number
        subset = df.loc[index_tuple]
        assert subset.shape == (n_freq, n_chan + 1)  # + 1 is the freq column
    with pytest.raises(ValueError, match='"time" is not a valid option'):
        spectrum.to_data_frame(index="time")
    # test picks
    picks = [0, 1]
    _pick_first = spectrum.pick(picks).to_data_frame()
    _pick_last = spectrum.to_data_frame(picks=picks)
    assert_frame_equal(_pick_first, _pick_last)


def _agg_helper(df, weights, group_cols):
    """Aggregate complex multitaper spectrum after conversion to DataFrame."""
    from pandas import Series

    unagged_columns = df[group_cols].iloc[0].values.tolist()
    x_mt = df.drop(columns=group_cols).values[np.newaxis].T
    psd = _psd_from_mt(x_mt, weights)
    psd = np.atleast_1d(np.squeeze(psd)).tolist()
    _df = dict(zip(df.columns, unagged_columns + psd))
    return Series(_df)


@pytest.mark.parametrize("long_format", (False, True))
@pytest.mark.parametrize(
    "method, output",
    [("welch", "complex"), ("welch", "power"), ("multitaper", "complex")],
)
def test_unaggregated_spectrum_to_data_frame(raw, long_format, method, output):
    """Test converting unaggregated spectra (multiple segments/tapers) to data frame."""
    pytest.importorskip("pandas")
    from pandas.testing import assert_frame_equal

    from mne.utils.dataframe import _inplace

    # aggregated spectrum → dataframe
    orig_df = raw.compute_psd(method=method).to_data_frame(long_format=long_format)
    # unaggregated welch or complex multitaper →
    #   aggregate w/ pandas (to make sure we did reshaping right)
    kwargs = dict()
    if method == "welch":
        kwargs.update(average=False)
    spectrum = raw.compute_psd(method=method, output=output, **kwargs)
    df = spectrum.to_data_frame(long_format=long_format)
    grouping_cols = ["freq"]
    drop_cols = ["segment"] if method == "welch" else ["taper"]
    if long_format:
        grouping_cols.append("channel")
        drop_cols.append("ch_type")
        orig_df.drop(columns="ch_type", inplace=True)
    # only do a couple freq bins, otherwise test takes forever for multitaper
    subset = partial(np.isin, test_elements=spectrum.freqs[:2])
    df = df.loc[subset(df["freq"])]
    orig_df = orig_df.loc[subset(orig_df["freq"])]
    # sort orig_df, because at present we can't actually prevent pandas from
    # sorting at the agg step *sigh*
    _inplace(orig_df, "sort_values", by=grouping_cols, ignore_index=True)
    # aggregate
    df = df.drop(columns=drop_cols)
    gb = df.groupby(grouping_cols, as_index=False, observed=False)
    if method == "welch":
        if output == "complex":

            def _fun(x):
                return np.mean(np.real(x * np.conj(x)))  # use mean to aggregate

            agg_df = gb.agg(_fun)
        else:
            agg_df = gb.mean()  # excludes missing values itself
    else:
        gb = gb[df.columns]  # XXX: try removing when minimum pandas >= 2.1 is required
        agg_df = gb.apply(_agg_helper, spectrum.weights, grouping_cols)
    # even with check_categorical=False, we know that the *data* matches;
    # what may differ is the order of the "levels" in the *metadata* for the
    # channel name column
    agg_df.sort_values(by=grouping_cols, ignore_index=True, inplace=True)
    orig_df.sort_values(by=grouping_cols, ignore_index=True, inplace=True)
    # One can have categorical dtype and the other plain object, so don't check that
    assert_frame_equal(agg_df, orig_df, check_categorical=False, check_dtype=False)


# not testing with Evoked because it already has projs applied
@pytest.mark.parametrize("inst", ("raw", "epochs"))
def test_spectrum_proj(inst, request):
    """Test that proj is applied correctly (gh 11177)."""
    inst = request.getfixturevalue(inst)
    has_proj = inst.compute_psd(proj=True)
    no_proj = inst.compute_psd(proj=False)
    assert not np.array_equal(has_proj.get_data(), no_proj.get_data())
    # make sure only the data (and the projs) were different
    has_proj._data = no_proj._data
    with has_proj.info._unlock():
        has_proj.info["projs"] = no_proj.info["projs"]
    assert has_proj == no_proj


@pytest.mark.parametrize(
    "method, average", [("welch", False), ("welch", "mean"), ("multitaper", None)]
)
def test_spectrum_complex(method, average):
    """Test output='complex' support."""
    sfreq = 100
    n = 10 * sfreq
    freq = 3.0
    phase = np.pi / 4  # should be recoverable
    data = np.cos(2 * np.pi * freq * np.arange(n) / sfreq + phase)[np.newaxis]
    raw = RawArray(data, create_info(1, sfreq, "eeg"))
    epochs = make_fixed_length_epochs(raw, duration=2.0, preload=True)
    assert len(epochs) == 5
    assert len(epochs.times) == 2 * sfreq
    kwargs = dict(output="complex", method=method)
    if method == "welch":
        kwargs["n_fft"] = sfreq
        want_dims = ("epoch", "channel", "freq")
        want_shape = (5, 1, sfreq // 2 + 1)
        if not average:
            want_dims = want_dims + ("segment",)
            want_shape = want_shape + (2,)
            kwargs["average"] = average
    else:
        assert method == "multitaper"
        assert not average
        want_dims = ("epoch", "channel", "taper", "freq")
        want_shape = (5, 1, 7, sfreq + 1)
    spectrum = epochs.compute_psd(**kwargs)
    idx = np.argmin(np.abs(spectrum.freqs - freq))
    assert spectrum.freqs[idx] == freq
    assert spectrum._dims == want_dims
    assert spectrum.shape == want_shape
    data = spectrum.get_data()
    assert data.dtype == np.complex128
    coef = spectrum.get_data(fmin=freq, fmax=freq).mean(0)
    if method == "multitaper":
        coef = coef[..., 0, :]  # first taper
    elif not average:
        coef = coef.mean(-1)  # over segments
    coef = coef.item()
    # Test phase matches what was simulated
    assert_allclose(np.angle(coef), phase, rtol=1e-4)
    # Now test that it warns appropriately
    epochs._data[0, 0, :] = 0  # actually zero for one epoch and ch
    with pytest.warns(UserWarning, match="Zero value.*channel 0"):
        epochs.compute_psd(**kwargs)
    # But not if we mark that channel as bad
    epochs.info["bads"] = epochs.ch_names[:1]
    epochs.compute_psd(**kwargs)


def test_spectrum_kwarg_triaging(raw):
    """Test kwarg triaging in legacy plot_psd() method."""
    import matplotlib.pyplot as plt

    regex = r"legacy plot_psd\(\) method.*unexpected keyword.*'axes'.*Try rewriting"
    _, axes = plt.subplots(1, 2)
    # `axes` is the new param name: technically only valid for Spectrum.plot()
    with _record_warnings(), pytest.warns(RuntimeWarning, match=regex):
        raw.plot_psd(axes=axes)
    # `ax` is the correct legacy param name
    raw.plot_psd(ax=axes)


def _check_spectrum_equivalent(spect1, spect2, tmp_path):
    data1 = spect1.get_data()
    data2 = spect2.get_data()
    assert_array_equal(data1, data2)
    assert_array_equal(spect1.freqs, spect2.freqs)


def test_spectrum_array_errors():
    """Test (Epochs)SpectrumArray constructor errors."""
    n_epochs = 10
    n_chans = 5
    n_freqs = 50
    freqs = np.arange(n_freqs)
    sfreq = 100
    rng = np.random.default_rng(44)
    data = rng.random((n_epochs, n_chans, n_freqs))
    dim_names = ("epoch", "channel", "freq")
    info = create_info(n_chans, sfreq, "eeg")
    # test incorrect ndims (for SpectrumArray; allows 2-3D data)
    with pytest.raises(ValueError, match="Invalid value for the 'data.ndim' parameter"):
        SpectrumArray(data[0, 0, :], info, freqs, dim_names=dim_names)
    with pytest.raises(ValueError, match="Invalid value for the 'data.ndim' parameter"):
        SpectrumArray(np.expand_dims(data, axis=3), info, freqs, dim_names=dim_names)
    # test incorrect ndims (for EpochsSpectrumArray; allows 3-4D data)
    with pytest.raises(ValueError, match="Invalid value for the 'data.ndim' parameter"):
        EpochsSpectrumArray(data[0, :, :], info, freqs, dim_names=dim_names)
    with pytest.raises(ValueError, match="Invalid value for the 'data.ndim' parameter"):
        EpochsSpectrumArray(
            np.expand_dims(data, axis=(3, 4)), info, freqs, dim_names=dim_names
        )
    # test incorrect epochs location
    with pytest.raises(ValueError, match="'epoch' must be the first dimension"):
        EpochsSpectrumArray(data, info, freqs, dim_names=("channel", "epoch", "freq"))
    # test mismatching events shape
    events = np.vstack(
        (
            np.arange(n_epochs + 1),
            np.zeros(n_epochs + 1, dtype=int),
            np.ones(n_epochs + 1, dtype=int),
        )
    ).T
    with pytest.raises(ValueError, match=r"first dimension.*dimension of `events`"):
        EpochsSpectrumArray(data, info, freqs, events, dim_names=dim_names)
    # test data-dimname mismatch
    with pytest.raises(ValueError, match=r"Expected data to have.*dimensions, got.*"):
        EpochsSpectrumArray(data, info, freqs, dim_names=dim_names[:-1])
    # test unrecognised dim_names (for SpectrumArray; epoch not allowed)
    with pytest.raises(ValueError, match="Invalid value for the 'dim_names' parameter"):
        SpectrumArray(data[0, :, :], info, freqs, dim_names=("epoch", "channel"))
    # test unrecognised dim_names (for EpochsSpectrumArray)
    with pytest.raises(ValueError, match="Invalid value for the 'dim_names' parameter"):
        EpochsSpectrumArray(
            data, info, freqs, dim_names=("epoch", "channel", "notfreq")
        )
    # test missing dim_names
    with pytest.raises(ValueError, match="Both 'channel' and 'freq' must be present"):
        EpochsSpectrumArray(
            data, info, freqs, dim_names=("epoch", "channel", "channel")
        )
    with pytest.raises(ValueError, match="Both 'channel' and 'freq' must be present"):
        EpochsSpectrumArray(data, info, freqs, dim_names=("epoch", "freq", "freq"))
    with pytest.raises(ValueError, match="Both 'channel' and 'freq' must be present"):
        EpochsSpectrumArray(data, info, freqs, dim_names=("epoch", "epoch", "epoch"))
    # test incorrect channel location (for SpectrumArray; must be 1st dim)
    with pytest.raises(ValueError, match="'channel' must be the first dimension"):
        SpectrumArray(data[0, :, :], info, freqs, dim_names=("freq", "channel"))
    # test incorrect channel location (for EpochsSpectrumArray; must be 2nd dim)
    with pytest.raises(ValueError, match="'channel' must be the second dimension"):
        EpochsSpectrumArray(data, info, freqs, dim_names=("epoch", "freq", "channel"))
    # test mismatching number of channels
    with pytest.raises(
        ValueError, match=re.escape("number of good + bad data channels")
    ):
        EpochsSpectrumArray(data[:, :-1, :], info, freqs, dim_names=dim_names)
    # test incorrect taper position
    with pytest.raises(ValueError, match="'taper' must be the second to last dim"):
        EpochsSpectrumArray(
            np.expand_dims(data, axis=3), info, freqs, dim_names=dim_names + ("taper",)
        )
    # test incorrect weight size
    with pytest.raises(ValueError, match=r"Expected size of `weights` to be.*, got.*"):
        EpochsSpectrumArray(
            np.expand_dims(data, axis=2),
            info,
            freqs,
            dim_names=("epoch", "channel", "taper", "freq"),
            weights=None,
        )
    with pytest.raises(ValueError, match=r"Expected size of `weights` to be.*, got.*"):
        EpochsSpectrumArray(
            np.expand_dims(data, axis=2),
            info,
            freqs,
            dim_names=("epoch", "channel", "taper", "freq"),
            weights=np.ones((1, 2, 1)),
        )
    # test incorrect segment position
    with pytest.raises(ValueError, match="'segment' must be the last dim"):
        EpochsSpectrumArray(
            np.expand_dims(data, axis=2),
            info,
            freqs,
            dim_names=("epoch", "channel", "segment", "freq"),
        )
    # test mismatching number of frequencies
    with pytest.raises(ValueError, match=r"number of frequencies.*number of elements"):
        EpochsSpectrumArray(data[:, :, :-1], info, freqs, dim_names=dim_names)


@pytest.mark.parametrize(
    "kind, method, output, average",
    [
        ("raw", "welch", "power", "mean"),  # test with precomputed spectrum
        ("epochs", "welch", "power", False),  # test with segments
        ("epochs", "multitaper", "complex", None),  # test with tapers
    ],  # additional variants don't improve coverage
)
def test_spectrum_array(kind, method, output, average, tmp_path, request):
    """Test EpochsSpectrumArray and SpectrumArray constructors."""
    dim_names = ("epoch", "channel") if kind == "epochs" else ("channel",)
    if method == "welch":
        dim_names += ("freq",) if average else ("freq", "segment")
    else:  # i.e. multitaper
        dim_names += ("freq",) if output == "power" else ("taper", "freq")
    if method == "welch" and output == "power" and average:
        spectrum = request.getfixturevalue(f"{kind}_spectrum")
    else:
        data = request.getfixturevalue(kind)
        kwargs = dict()
        if method == "welch":
            kwargs.update(average=average)
        spectrum = data.compute_psd(method=method, output=output, **kwargs)
    data, freqs = spectrum.get_data(return_freqs=True)
    Klass = SpectrumArray if kind == "raw" else EpochsSpectrumArray
    spect_arr = Klass(
        data=data,
        info=spectrum.info,
        freqs=freqs,
        dim_names=dim_names,
        weights=spectrum.weights,
    )
    _check_spectrum_equivalent(spectrum, spect_arr, tmp_path)


@pytest.mark.parametrize(
    "method, output, average",
    [
        ("welch", "power", "mean"),  # test with precomputed spectrum
        ("welch", "complex", False),  # test aggr over segments & conversion to power
        ("multitaper", "complex", None),  # test aggr over tapers & conversion to power
    ],  # additional variants don't improve coverage
)
def test_plot_spectrum(method, output, average, request):
    """Test plotting EpochsSpectrum(Array).

    Testing Spectrum(Array) with raw data doesn't improve coverage.
    """
    if method == "welch" and output == "power" and average:
        spectrum = request.getfixturevalue("epochs_spectrum")
    else:
        data = request.getfixturevalue("epochs")
        kwargs = dict()
        if method == "welch":
            kwargs.update(average=average)
        spectrum = data.compute_psd(method=method, output=output, **kwargs)
    spectrum.info["bads"] = spectrum.ch_names[:1]  # one grad channel
    spectrum.plot(average=True, amplitude=True, spatial_colors=True)
    spectrum.plot(average=True, amplitude=False, spatial_colors=False)
    n_grad = sum(ch_type == "grad" for ch_type in spectrum.get_channel_types())
    for amp, sc in ((True, True), (False, False)):
        fig = spectrum.plot(average=False, amplitude=amp, spatial_colors=sc, exclude=())
        lines = fig.axes[0].lines[2:]  # grads, ignore two vlines
        assert len(lines) == n_grad
        bad_color = "0.5" if sc else "r"
        n_bad = sum(same_color(line.get_color(), bad_color) for line in lines)
        assert n_bad == 1
    spectrum.plot_topo()
    spectrum.plot_topomap()


def test_plot_spectrum_array_with_bads():
    """Test plotting a spectrum array with bads."""
    raw = RawArray(np.random.randn(3, 1000), create_info(3, 1000, "eeg"))
    raw.info["bads"] = [raw.ch_names[1]]
    spectrum = raw.compute_psd()
    with pytest.raises(
        ValueError, match=re.escape("number of good + bad data channels")
    ):
        SpectrumArray(spectrum.get_data(), spectrum.info, spectrum.freqs)
    spectrum2 = SpectrumArray(
        spectrum.get_data(exclude=()), spectrum.info, spectrum.freqs
    )
    spectrum2.plot(spatial_colors=False)


@pytest.mark.parametrize("dB", (False, True))
@pytest.mark.parametrize("amplitude", (False, True))
def test_plot_spectrum_dB(raw_spectrum, dB, amplitude):
    """Test that we properly handle amplitude/power and dB."""
    idx = 7
    power = 3
    freqs = np.linspace(1, 100, 100)
    data = np.full((1, freqs.size), np.finfo(float).tiny)
    data[0, idx] = power
    info = create_info(ch_names=["delta"], sfreq=1000, ch_types="eeg")
    psd = SpectrumArray(data=data, info=info, freqs=freqs)
    with pytest.warns(RuntimeWarning, match="Channel locations not available"):
        fig = psd.plot(dB=dB, amplitude=amplitude)
    trace = list(
        filter(lambda x: len(x.get_data()[0]) == len(freqs), fig.axes[0].lines)
    )[0]
    got = trace.get_data()[1][idx]
    want = power * 1e12  # scaling for EEG (V → μV), squared
    if amplitude:
        want = np.sqrt(want)
    if dB:
        want = (20 if amplitude else 10) * np.log10(want)

    assert want == got, f"expected {want}, got {got}"
