# 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 mne import create_info
from mne._fiff.constants import FIFF
from mne._fiff.pick import _picks_to_idx
from mne.datasets import testing
from mne.datasets.testing import data_path
from mne.io import RawArray, read_raw_nirx, read_raw_snirf
from mne.preprocessing.nirs import (
    _channel_chromophore,
    _channel_frequencies,
    _check_channels_ordered,
    _fnirs_optode_names,
    _fnirs_spread_bads,
    _optode_position,
    _validate_nirs_info,
    beer_lambert_law,
    optical_density,
    scalp_coupling_index,
    tddr,
)

fname_nirx_15_0 = (
    data_path(download=False) / "NIRx" / "nirscout" / "nirx_15_0_recording"
)
fname_nirx_15_2 = (
    data_path(download=False) / "NIRx" / "nirscout" / "nirx_15_2_recording"
)
fname_nirx_15_2_short = (
    data_path(download=False) / "NIRx" / "nirscout" / "nirx_15_2_recording_w_short"
)
fname_labnirs_multi_wavelength = (
    data_path(download=False) / "SNIRF" / "Labnirs" / "labnirs_3wl_raw_recording.snirf"
)


def read_raw_snirf_safe(fname):
    """Wrap to read_raw_snirf, skipping if h5py is not installed."""
    pytest.importorskip("h5py")
    return read_raw_snirf(fname)


@testing.requires_testing_data
@pytest.mark.parametrize(
    "fname, readerfn",
    [
        (fname_nirx_15_0, read_raw_nirx),
        (fname_labnirs_multi_wavelength, read_raw_snirf_safe),
    ],
)
def test_fnirs_picks(fname, readerfn):
    """Test picking of fnirs types after different conversions."""
    raw = readerfn(fname)
    picks = _picks_to_idx(raw.info, "fnirs_cw_amplitude")
    assert len(picks) == len(raw.ch_names)
    raw_subset = raw.copy().pick(picks="fnirs_cw_amplitude")
    for ch in raw_subset.info["chs"]:
        assert ch["coil_type"] == FIFF.FIFFV_COIL_FNIRS_CW_AMPLITUDE

    picks = _picks_to_idx(raw.info, ["fnirs_cw_amplitude", "fnirs_od"])
    assert len(picks) == len(raw.ch_names)
    picks = _picks_to_idx(raw.info, ["fnirs_cw_amplitude", "fnirs_od", "hbr"])
    assert len(picks) == len(raw.ch_names)
    pytest.raises(ValueError, _picks_to_idx, raw.info, "fnirs_od")
    pytest.raises(ValueError, _picks_to_idx, raw.info, "hbo")
    pytest.raises(ValueError, _picks_to_idx, raw.info, ["hbr"])
    pytest.raises(ValueError, _picks_to_idx, raw.info, "fnirs_fd_phase")
    pytest.raises(ValueError, _picks_to_idx, raw.info, "junk")

    raw = optical_density(raw)
    picks = _picks_to_idx(raw.info, "fnirs_od")
    assert len(picks) == len(raw.ch_names)
    raw_subset = raw.copy().pick(picks="fnirs_od")
    for ch in raw_subset.info["chs"]:
        assert ch["coil_type"] == FIFF.FIFFV_COIL_FNIRS_OD

    picks = _picks_to_idx(raw.info, ["fnirs_cw_amplitude", "fnirs_od"])
    assert len(picks) == len(raw.ch_names)
    picks = _picks_to_idx(raw.info, ["fnirs_cw_amplitude", "fnirs_od", "hbr"])
    assert len(picks) == len(raw.ch_names)
    pytest.raises(ValueError, _picks_to_idx, raw.info, "fnirs_cw_amplitude")
    pytest.raises(ValueError, _picks_to_idx, raw.info, "hbo")
    pytest.raises(ValueError, _picks_to_idx, raw.info, "hbr")
    pytest.raises(ValueError, _picks_to_idx, raw.info, "fnirs_fd_phase")
    pytest.raises(ValueError, _picks_to_idx, raw.info, "junk")

    raw = beer_lambert_law(raw)
    picks = _picks_to_idx(raw.info, "hbo")
    assert len(picks) == len(raw.ch_names) / 2
    raw_subset = raw.copy().pick(picks="hbo")
    for ch in raw_subset.info["chs"]:
        assert ch["coil_type"] == FIFF.FIFFV_COIL_FNIRS_HBO

    picks = _picks_to_idx(raw.info, ["hbr"])
    assert len(picks) == len(raw.ch_names) / 2
    raw_subset = raw.copy().pick(picks=["hbr"])
    for ch in raw_subset.info["chs"]:
        assert ch["coil_type"] == FIFF.FIFFV_COIL_FNIRS_HBR

    picks = _picks_to_idx(raw.info, ["hbo", "hbr"])
    assert len(picks) == len(raw.ch_names)
    picks = _picks_to_idx(raw.info, ["hbo", "fnirs_od", "hbr"])
    assert len(picks) == len(raw.ch_names)
    picks = _picks_to_idx(raw.info, ["hbo", "fnirs_od"])
    assert len(picks) == len(raw.ch_names) / 2
    pytest.raises(ValueError, _picks_to_idx, raw.info, "fnirs_cw_amplitude")
    pytest.raises(ValueError, _picks_to_idx, raw.info, ["fnirs_od"])
    pytest.raises(ValueError, _picks_to_idx, raw.info, "junk")
    pytest.raises(ValueError, _picks_to_idx, raw.info, "fnirs_fd_phase")


# Backward compat wrapper for simplicity below
def _fnirs_check_bads(info):
    _validate_nirs_info(info)


@testing.requires_testing_data
@pytest.mark.parametrize(
    "fname, readerfn",
    [
        (fname_nirx_15_0, read_raw_nirx),
        (fname_nirx_15_2_short, read_raw_nirx),
        (fname_nirx_15_2, read_raw_nirx),
        (fname_labnirs_multi_wavelength, read_raw_snirf_safe),
    ],
)
def test_fnirs_check_bads(fname, readerfn):
    """Test checking of bad markings."""
    # No bad channels, so these should all pass
    raw = readerfn(fname)
    _fnirs_check_bads(raw.info)
    raw = optical_density(raw)
    _fnirs_check_bads(raw.info)
    raw = beer_lambert_law(raw)
    _fnirs_check_bads(raw.info)

    # Mark pairs of bad channels, so these should all pass
    raw = readerfn(fname)
    nfreqs = len(set(_channel_frequencies(raw.info)))
    raw.info["bads"] = raw.ch_names[0:nfreqs]
    _fnirs_check_bads(raw.info)
    raw = optical_density(raw)
    _fnirs_check_bads(raw.info)
    raw = beer_lambert_law(raw)
    _fnirs_check_bads(raw.info)

    # Mark single channel as bad, so these should all fail
    raw = readerfn(fname)
    raw.info["bads"] = raw.ch_names[0:1]
    pytest.raises(RuntimeError, _fnirs_check_bads, raw.info)
    with pytest.raises(RuntimeError, match="bad labelling"):
        raw = optical_density(raw)
    raw.info["bads"] = []
    raw = optical_density(raw)
    raw.info["bads"] = raw.ch_names[0:1]
    pytest.raises(RuntimeError, _fnirs_check_bads, raw.info)
    with pytest.raises(RuntimeError, match="bad labelling"):
        raw = beer_lambert_law(raw)
    pytest.raises(RuntimeError, _fnirs_check_bads, raw.info)


@testing.requires_testing_data
@pytest.mark.parametrize(
    "fname, readerfn",
    [
        (fname_nirx_15_0, read_raw_nirx),
        (fname_nirx_15_2_short, read_raw_nirx),
        (fname_nirx_15_2, read_raw_nirx),
        (fname_labnirs_multi_wavelength, read_raw_snirf_safe),
    ],
)
def test_fnirs_spread_bads(fname, readerfn):
    """Test checking of bad markings."""
    # Test spreading upwards in frequency and on raw data
    raw = readerfn(fname)
    nfreqs = len(set(_channel_frequencies(raw.info)))
    raw.info["bads"] = [raw.ch_names[0]]
    info = _fnirs_spread_bads(raw.info)
    assert info["bads"] == raw.ch_names[:nfreqs]

    # Test multiple spreading directions on od data
    # For each wavelength, mark the nth item in the nth group as bad
    # e.g. 3 wavelengths:
    #   group 0 (channels 0, 1, 2) - channel 0 is marked bad,
    #   group 1 (channels 3, 4, 5) - channel 4 is bad,
    #   group 2 (channels 6, 7, 8) - channel 8 is bad.
    # This way we can test spreading from each group member in a
    # way that's agnostic of the number of wavelengths, and avoid
    # hard-coded values. Needs nfreqs**2 number of channels.
    raw_od = optical_density(raw)
    bads = [raw_od.ch_names[nfreqs * ii + ii] for ii in range(nfreqs)]
    expected_bads = raw_od.ch_names[: nfreqs**2]
    raw_od.info["bads"] = bads
    info = _fnirs_spread_bads(raw_od.info)
    # channels might not be sorted but the spreading result is
    assert info["bads"] == sorted(expected_bads)

    # Test spreading multiple bads and on chroma data
    # Hb data always has 2 channels per S-D pair, this works for any nfreqs
    raw_hb = beer_lambert_law(raw_od)
    raw_hb.info["bads"] = [raw_hb.ch_names[x] for x in [1, 8]]
    info = _fnirs_spread_bads(raw_hb.info)
    assert info["bads"] == [info.ch_names[x] for x in [0, 1, 8, 9]]


@testing.requires_testing_data
@pytest.mark.parametrize(
    "fname, readerfn",
    [
        (fname_nirx_15_0, read_raw_nirx),
        (fname_nirx_15_2_short, read_raw_nirx),
        (fname_nirx_15_2, read_raw_nirx),
        (fname_labnirs_multi_wavelength, read_raw_snirf_safe),
    ],
)
def test_fnirs_channel_frequency_ordering(fname, readerfn):
    """Test fNIRS channel frequencies ordering and related errors."""
    # fNIRS data requires specific channel naming and ordering.

    # Ensure that freqs are well-ordered after reading in from file,
    # and that there are no chroma channels
    raw = readerfn(fname)
    freqs = np.unique(_channel_frequencies(raw.info)).tolist()
    assert sorted(freqs) == freqs
    chroma = np.unique(_channel_chromophore(raw.info))
    assert len(chroma) == 0
    picks = _check_channels_ordered(raw.info, freqs)
    assert len(picks) == len(raw.ch_names)  # as all fNIRS only data

    # The ordering must be increasing for the pairs, if provided
    raw_reversed = raw.pick(list(reversed(raw.ch_names)))

    with pytest.raises(ValueError, match="The frequencies.*sorted.*"):
        _check_channels_ordered(raw_reversed.info, list(reversed(freqs)))
    # So if we flip the second argument it should pass again
    picks = _check_channels_ordered(raw_reversed.info, freqs)
    nfreqs = len(freqs)
    for ii in range(nfreqs):
        suffixes = {
            int(raw_reversed.ch_names[pick].split(" ")[1]) for pick in picks[ii::nfreqs]
        }
        assert suffixes == {freqs[ii]}

    # Check on OD data
    raw = optical_density(raw)
    freqs = np.unique(_channel_frequencies(raw.info)).tolist()
    assert sorted(freqs) == freqs
    chroma = np.unique(_channel_chromophore(raw.info))
    assert len(chroma) == 0
    picks = _check_channels_ordered(raw.info, freqs)
    assert len(picks) == len(raw.ch_names)  # as all fNIRS only data

    # Check on haemoglobin data
    raw = beer_lambert_law(raw)
    freqs = np.unique(_channel_frequencies(raw.info))
    assert len(freqs) == 0
    assert len(_channel_chromophore(raw.info)) == len(raw.ch_names)
    chroma = np.unique(_channel_chromophore(raw.info))
    assert_array_equal(chroma, ["hbo", "hbr"])
    picks = _check_channels_ordered(raw.info, chroma)
    assert len(picks) == len(raw.ch_names)
    with pytest.raises(ValueError, match="chromophore in info"):
        _check_channels_ordered(raw.info, ["hbr", "hbo"])


def test_fnirs_channel_naming_and_order_custom_raw():
    """Ensure fNIRS channel checking on manually created data."""
    data = np.random.normal(size=(6, 10))

    # Start with a correctly named raw intensity dataset
    # These are the steps required to build an fNIRS Raw object from scratch
    ch_names = [
        "S1_D1 760",
        "S1_D1 850",
        "S2_D1 760",
        "S2_D1 850",
        "S3_D1 760",
        "S3_D1 850",
    ]
    ch_types = np.repeat("fnirs_cw_amplitude", 6)
    info = create_info(ch_names=ch_names, ch_types=ch_types, sfreq=1.0)
    raw = RawArray(data, info, verbose=True)
    freqs = np.tile([760, 850], 3)
    for idx, f in enumerate(freqs):
        raw.info["chs"][idx]["loc"][9] = f

    freqs = np.unique(_channel_frequencies(raw.info))
    picks = _check_channels_ordered(raw.info, freqs)
    assert len(picks) == len(raw.ch_names)
    assert len(picks) == 6

    # Different systems use different frequencies, so ensure that works
    ch_names = [
        "S1_D1 920",
        "S1_D1 850",
        "S2_D1 920",
        "S2_D1 850",
        "S3_D1 920",
        "S3_D1 850",
    ]
    ch_types = np.repeat("fnirs_cw_amplitude", 6)
    info = create_info(ch_names=ch_names, ch_types=ch_types, sfreq=1.0)
    raw = RawArray(data, info, verbose=True)
    freqs = np.tile([920, 850], 3)
    for idx, f in enumerate(freqs):
        raw.info["chs"][idx]["loc"][9] = f

    picks = _check_channels_ordered(raw.info, [850, 920])
    assert len(picks) == len(raw.ch_names)
    assert len(picks) == 6

    # Catch expected errors

    # The frequencies named in the channel names must match the info loc field
    ch_names = [
        "S1_D1 760",
        "S1_D1 850",
        "S2_D1 760",
        "S2_D1 850",
        "S3_D1 760",
        "S3_D1 850",
    ]
    ch_types = np.repeat("fnirs_cw_amplitude", 6)
    info = create_info(ch_names=ch_names, ch_types=ch_types, sfreq=1.0)
    raw = RawArray(data, info, verbose=True)
    freqs = np.tile([920, 850], 3)
    for idx, f in enumerate(freqs):
        raw.info["chs"][idx]["loc"][9] = f
    with pytest.raises(ValueError, match="not ordered"):
        _check_channels_ordered(raw.info, [850, 920])

    # Catch if someone doesn't set the info field
    ch_names = [
        "S1_D1 760",
        "S1_D1 850",
        "S2_D1 760",
        "S2_D1 850",
        "S3_D1 760",
        "S3_D1 850",
    ]
    ch_types = np.repeat("fnirs_cw_amplitude", 6)
    info = create_info(ch_names=ch_names, ch_types=ch_types, sfreq=1.0)
    raw = RawArray(data, info, verbose=True)
    with pytest.raises(ValueError, match="missing wavelength information"):
        _check_channels_ordered(raw.info, [850, 920])

    # I have seen data encoded not in alternating frequency, but blocked.
    ch_names = [
        "S1_D1 760",
        "S2_D1 760",
        "S3_D1 760",
        "S1_D1 850",
        "S2_D1 850",
        "S3_D1 850",
    ]
    ch_types = np.repeat("fnirs_cw_amplitude", 6)
    info = create_info(ch_names=ch_names, ch_types=ch_types, sfreq=1.0)
    raw = RawArray(data, info, verbose=True)
    freqs = np.repeat([760, 850], 3)
    for idx, f in enumerate(freqs):
        raw.info["chs"][idx]["loc"][9] = f
    _check_channels_ordered(raw.info, [760, 850])


def test_fnirs_channel_naming_and_order_custom_optical_density():
    """Ensure fNIRS channel checking on manually created data."""
    data = np.random.normal(size=(6, 10))

    # Start with a correctly named raw intensity dataset
    # These are the steps required to build an fNIRS Raw object from scratch
    ch_names = [
        "S1_D1 760",
        "S1_D1 850",
        "S2_D1 760",
        "S2_D1 850",
        "S3_D1 760",
        "S3_D1 850",
    ]
    ch_types = np.repeat("fnirs_od", 6)
    info = create_info(ch_names=ch_names, ch_types=ch_types, sfreq=1.0)
    raw = RawArray(data, info, verbose=True)
    freqs = np.tile([760, 850], 3)
    for idx, f in enumerate(freqs):
        raw.info["chs"][idx]["loc"][9] = f

    freqs = np.unique(_channel_frequencies(raw.info))
    picks = _check_channels_ordered(raw.info, freqs)
    assert len(picks) == len(raw.ch_names)
    assert len(picks) == 6

    # Check block naming for optical density
    ch_names = [
        "S1_D1 760",
        "S2_D1 760",
        "S3_D1 760",
        "S1_D1 850",
        "S2_D1 850",
        "S3_D1 850",
    ]
    ch_types = np.repeat("fnirs_od", 6)
    info = create_info(ch_names=ch_names, ch_types=ch_types, sfreq=1.0)
    raw = RawArray(data, info, verbose=True)
    freqs = np.repeat([760, 850], 3)
    for idx, f in enumerate(freqs):
        raw.info["chs"][idx]["loc"][9] = f
    # no problems here
    _check_channels_ordered(raw.info, [760, 850])
    # or with this (nirx) reordering
    raw.pick(picks=[0, 3, 1, 4, 2, 5])
    _check_channels_ordered(raw.info, [760, 850])

    # Check that if you mix types you get an error
    ch_names = [
        "S1_D1 hbo",
        "S1_D1 hbr",
        "S2_D1 hbo",
        "S2_D1 hbr",
        "S3_D1 hbo",
        "S3_D1 hbr",
    ]
    ch_types = np.tile(["hbo", "hbr"], 3)
    info = create_info(ch_names=ch_names, ch_types=ch_types, sfreq=1.0)
    raw2 = RawArray(data, info, verbose=True)
    raw.add_channels([raw2])
    with pytest.raises(ValueError, match="does not support a combination"):
        _check_channels_ordered(raw.info, [760, 850])


def test_fnirs_channel_naming_and_order_custom_chroma():
    """Ensure fNIRS channel checking on manually created data."""
    data = np.random.RandomState(0).randn(6, 10)

    # Start with a correctly named raw intensity dataset
    # These are the steps required to build an fNIRS Raw object from scratch
    ch_names = [
        "S1_D1 hbo",
        "S1_D1 hbr",
        "S2_D1 hbo",
        "S2_D1 hbr",
        "S3_D1 hbo",
        "S3_D1 hbr",
    ]
    ch_types = np.tile(["hbo", "hbr"], 3)
    info = create_info(ch_names=ch_names, ch_types=ch_types, sfreq=1.0)
    raw = RawArray(data, info, verbose=True)

    chroma = np.unique(_channel_chromophore(raw.info))
    picks = _check_channels_ordered(raw.info, chroma)
    assert len(picks) == len(raw.ch_names)
    assert len(picks) == 6

    # Test block creation fails
    ch_names = [
        "S1_D1 hbo",
        "S2_D1 hbo",
        "S3_D1 hbo",
        "S1_D1 hbr",
        "S2_D1 hbr",
        "S3_D1 hbr",
    ]
    ch_types = np.repeat(["hbo", "hbr"], 3)
    info = create_info(ch_names=ch_names, ch_types=ch_types, sfreq=1.0)
    raw = RawArray(data, info, verbose=True)
    # no issue here
    _check_channels_ordered(raw.info, ["hbo", "hbr"])
    # reordering okay, too
    raw.pick(picks=[0, 3, 1, 4, 2, 5])
    _check_channels_ordered(raw.info, ["hbo", "hbr"])
    # Wrong names should fail
    with pytest.raises(ValueError, match="chromophore in info"):
        _check_channels_ordered(raw.info, ["hbb", "hbr"])

    # Test weird naming
    ch_names = [
        "S1_D1 hbb",
        "S1_D1 hbr",
        "S2_D1 hbb",
        "S2_D1 hbr",
        "S3_D1 hbb",
        "S3_D1 hbr",
    ]
    ch_types = np.tile(["hbo", "hbr"], 3)
    info = create_info(ch_names=ch_names, ch_types=ch_types, sfreq=1.0)
    raw = RawArray(data, info, verbose=True)
    with pytest.raises(ValueError, match="naming conventions"):
        _check_channels_ordered(raw.info, ["hbo", "hbr"])

    # Check more weird naming
    ch_names = [
        "S1_DX hbo",
        "S1_DX hbr",
        "S2_D1 hbo",
        "S2_D1 hbr",
        "S3_D1 hbo",
        "S3_D1 hbr",
    ]
    ch_types = np.tile(["hbo", "hbr"], 3)
    info = create_info(ch_names=ch_names, ch_types=ch_types, sfreq=1.0)
    raw = RawArray(data, info, verbose=True)
    with pytest.raises(ValueError, match="can not be parsed"):
        _check_channels_ordered(raw.info, ["hbo", "hbr"])


def test_optode_names():
    """Ensure optode name extraction is correct."""
    ch_names = [
        "S11_D2 760",
        "S11_D2 850",
        "S3_D1 760",
        "S3_D1 850",
        "S2_D13 760",
        "S2_D13 850",
    ]
    ch_types = np.repeat("fnirs_od", 6)
    info = create_info(ch_names=ch_names, ch_types=ch_types, sfreq=1.0)
    src_names, det_names = _fnirs_optode_names(info)
    assert_array_equal(src_names, [f"S{n}" for n in ["2", "3", "11"]])
    assert_array_equal(det_names, [f"D{n}" for n in ["1", "2", "13"]])

    ch_names = [
        "S1_D11 hbo",
        "S1_D11 hbr",
        "S2_D17 hbo",
        "S2_D17 hbr",
        "S3_D1 hbo",
        "S3_D1 hbr",
    ]
    ch_types = np.tile(["hbo", "hbr"], 3)
    info = create_info(ch_names=ch_names, ch_types=ch_types, sfreq=1.0)
    src_names, det_names = _fnirs_optode_names(info)
    assert_array_equal(src_names, [f"S{n}" for n in range(1, 4)])
    assert_array_equal(det_names, [f"D{n}" for n in ["1", "11", "17"]])


@testing.requires_testing_data
def test_optode_loc():
    """Ensure optode location extraction is correct."""
    raw = read_raw_nirx(fname_nirx_15_2_short)
    loc = _optode_position(raw.info, "D3")
    assert_array_almost_equal(loc, [0.082804, 0.01573, 0.024852])


def test_order_agnostic(nirx_snirf):
    """Test that order does not matter to (pre)processing results."""
    raw_nirx, raw_snirf = nirx_snirf
    raw_random = raw_nirx.copy().pick(
        np.random.RandomState(0).permutation(len(raw_nirx.ch_names))
    )
    raws = dict(nirx=raw_nirx, snirf=raw_snirf, random=raw_random)
    del raw_nirx, raw_snirf, raw_random
    orders = dict()
    # continuous wave
    for key, r in raws.items():
        assert set(r.get_channel_types()) == {"fnirs_cw_amplitude"}
        orders[key] = [r.ch_names.index(name) for name in raws["nirx"].ch_names]
        assert_array_equal(raws["nirx"].ch_names, np.array(r.ch_names)[orders[key]])
        assert_allclose(raws["nirx"].get_data(), r.get_data(orders[key]), err_msg=key)
    assert_array_equal(orders["nirx"], np.arange(len(raws["nirx"].ch_names)))
    # optical density
    for key, r in raws.items():
        raws[key] = r = optical_density(r)
        assert_allclose(raws["nirx"].get_data(), r.get_data(orders[key]), err_msg=key)
        assert set(r.get_channel_types()) == {"fnirs_od"}
    # scalp-coupling index
    sci = dict()
    for key, r in raws.items():
        sci[key] = r = scalp_coupling_index(r)
        assert_allclose(sci["nirx"], r[orders[key]], err_msg=key, rtol=0.01)
    # TDDR (on optical)
    tddrs = dict()
    for key, r in raws.items():
        tddrs[key] = r = tddr(r)
        assert_allclose(
            tddrs["nirx"].get_data(), r.get_data(orders[key]), err_msg=key, atol=1e-4
        )
        assert set(r.get_channel_types()) == {"fnirs_od"}
    # beer-lambert
    for key, r in raws.items():
        raws[key] = r = beer_lambert_law(r)
        assert_allclose(
            raws["nirx"].get_data(), r.get_data(orders[key]), err_msg=key, rtol=2e-7
        )
        assert set(r.get_channel_types()) == {"hbo", "hbr"}
    # TDDR (on haemo)
    tddrs = dict()
    for key, r in raws.items():
        tddrs[key] = r = tddr(r)
        assert_allclose(
            tddrs["nirx"].get_data(), r.get_data(orders[key]), err_msg=key, atol=1e-9
        )
        assert set(r.get_channel_types()) == {"hbo", "hbr"}


@testing.requires_testing_data
@pytest.mark.parametrize(
    "fname, readerfn",
    [
        (fname_nirx_15_0, read_raw_nirx),
        (fname_nirx_15_2_short, read_raw_nirx),
        (fname_nirx_15_2, read_raw_nirx),
        (fname_labnirs_multi_wavelength, read_raw_snirf_safe),
    ],
)
def test_nirs_channel_grouping(fname, readerfn):
    """Test channel grouping related errors."""
    raw = readerfn(fname)
    freqs = np.unique(_channel_frequencies(raw.info)).tolist()
    nfreqs = len(freqs)

    # Each source-detector (S-D) optode pair may have data measured at
    # >=2 wavelengths. The channels that belong to an optode pair form
    # a group, with a size equal to the number of wavelengths (nfreqs).
    # An error is raised if these groups are incomplete.

    picks = _check_channels_ordered(raw.info, freqs)
    assert len(picks) == len(raw.ch_names)

    # Removing one channel breaks the grouping
    raw_dropped = raw.copy().drop_channels(raw.ch_names[4])
    with pytest.raises(ValueError, match="NIRS channels not ordered correctly."):
        _check_channels_ordered(raw_dropped.info, freqs)

    # Selecting incomplete groups results in an error
    raw_incomplete = raw.copy().pick(list(range(nfreqs + 1)))
    with pytest.raises(ValueError, match="NIRS channels not ordered correctly."):
        _check_channels_ordered(raw_incomplete.info, freqs)

    # Changing the frequency in one channel's name also breaks groups
    raw_extrafreq = raw.copy()
    new_ch10_name = f"{raw_extrafreq.ch_names[10].split(' ')[0]} 100"
    raw_extrafreq.rename_channels({raw_extrafreq.ch_names[10]: new_ch10_name})
    print(raw_extrafreq.ch_names)
    # checks result in error for both old and new set of frequencies
    with pytest.raises(ValueError, match="NIRS channels not ordered correctly."):
        _check_channels_ordered(raw_extrafreq.info, [100] + freqs)
    with pytest.raises(ValueError, match="NIRS channels not ordered correctly."):
        _check_channels_ordered(raw_extrafreq.info, freqs)

    # Frequency values are also stored in info['chs'][ii]['loc'][9].
    # Even though the actual values can be arbitrary, they cannot be None.
    raw_locnone = raw.copy()
    raw_locnone.info["chs"][10]["loc"][9] = None
    with pytest.raises(
        ValueError, match="NIRS channels is missing wavelength information"
    ):
        _check_channels_ordered(raw_locnone.info, freqs)
