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

import os
import shutil
from copy import deepcopy

import numpy as np
import pytest
from numpy.testing import (
    assert_allclose,
    assert_array_almost_equal,
    assert_array_equal,
    assert_equal,
)
from scipy import io

import mne
from mne import read_epochs_eeglab, write_events
from mne.annotations import events_from_annotations, read_annotations
from mne.channels import read_custom_montage
from mne.datasets import testing
from mne.io import read_raw_eeglab
from mne.io.eeglab._eeglab import _readmat
from mne.io.eeglab.eeglab import _dol_to_lod, _get_montage_information
from mne.io.tests.test_raw import _test_raw_reader
from mne.utils import Bunch, _check_pymatreader_installed, _record_warnings

base_dir = testing.data_path(download=False) / "EEGLAB"
raw_fname_mat = base_dir / "test_raw.set"
raw_fname_onefile_mat = base_dir / "test_raw_onefile.set"
raw_fname_event_duration = base_dir / "test_raw_event_duration.set"
epochs_fname_mat = base_dir / "test_epochs.set"
epochs_fname_onefile_mat = base_dir / "test_epochs_onefile.set"
epochs_mat_fnames = [epochs_fname_mat, epochs_fname_onefile_mat]
raw_fname_chanloc = base_dir / "test_raw_chanloc.set"
raw_fname_chanloc_fids = base_dir / "test_raw_chanloc_fids.set"
raw_fname_2021 = base_dir / "test_raw_2021.set"
raw_fname_h5 = base_dir / "test_raw_h5.set"
epochs_fname_h5 = base_dir / "test_epochs_h5.set"
epochs_fname_onefile_h5 = base_dir / "test_epochs_onefile_h5.set"
epochs_h5_fnames = [epochs_fname_h5, epochs_fname_onefile_h5]
montage_path = base_dir / "test_chans.locs"


@testing.requires_testing_data
@pytest.mark.parametrize(
    "fname",
    [
        raw_fname_mat,
        pytest.param(
            raw_fname_h5,
            marks=[
                pytest.mark.skipif(
                    not _check_pymatreader_installed(strict=False),
                    reason="pymatreader not installed",
                )
            ],
        ),
        raw_fname_chanloc,
    ],
    ids=os.path.basename,
)
def test_io_set_raw(fname):
    """Test importing EEGLAB .set files."""
    montage = read_custom_montage(montage_path)
    montage.ch_names = [f"EEG {ii:03d}" for ii in range(len(montage.ch_names))]

    kws = dict(reader=read_raw_eeglab, input_fname=fname)
    if fname.name == "test_raw_chanloc.set":
        with pytest.warns(RuntimeWarning, match="The data contains 'boundary' events"):
            raw0 = _test_raw_reader(**kws)
    elif "_h5" in fname.name:  # should be safe enough, and much faster
        raw0 = read_raw_eeglab(fname, preload=True)
    else:
        raw0 = _test_raw_reader(**kws)

    # test that preloading works
    if fname.name == "test_raw_chanloc.set":
        raw0.set_montage(montage, on_missing="ignore")
        # crop to check if the data has been properly preloaded; we cannot
        # filter as the snippet of raw data is very short
        raw0.crop(0, 1)
    else:
        raw0.set_montage(montage)
        raw0.filter(
            1, None, l_trans_bandwidth="auto", filter_length="auto", phase="zero"
        )

    # test that using uint16_codec does not break stuff
    read_raw_kws = dict(input_fname=fname, preload=False, uint16_codec="ascii")
    if fname.name == "test_raw_chanloc.set":
        with pytest.warns(RuntimeWarning, match="The data contains 'boundary' events"):
            raw0 = read_raw_eeglab(**read_raw_kws)
            raw0.set_montage(montage, on_missing="ignore")
    else:
        raw0 = read_raw_eeglab(**read_raw_kws)
        raw0.set_montage(montage)

    # Annotations
    if fname != raw_fname_chanloc:
        assert len(raw0.annotations) == 154
        assert set(raw0.annotations.description) == {"rt", "square"}
        assert_array_equal(raw0.annotations.duration, 0.0)


@testing.requires_testing_data
def test_io_set_raw_more(tmp_path):
    """Test importing EEGLAB .set files."""
    eeg = io.loadmat(raw_fname_mat, struct_as_record=False, squeeze_me=True)["EEG"]

    # test reading file with one event (read old version)
    negative_latency_fname = tmp_path / "test_negative_latency.set"
    events = deepcopy(eeg.event[0])
    events.latency = 0
    io.savemat(
        negative_latency_fname,
        {
            "EEG": {
                "trials": eeg.trials,
                "srate": eeg.srate,
                "nbchan": eeg.nbchan,
                "data": "test_negative_latency.fdt",
                "epoch": eeg.epoch,
                "event": events,
                "chanlocs": eeg.chanlocs,
                "pnts": eeg.pnts,
            }
        },
        appendmat=False,
        oned_as="row",
    )
    shutil.copyfile(
        base_dir / "test_raw.fdt", negative_latency_fname.with_suffix(".fdt")
    )
    with (
        _record_warnings(),
        pytest.warns(RuntimeWarning, match="has a sample index of -1."),
    ):
        read_raw_eeglab(input_fname=negative_latency_fname, preload=True)

    # test negative event latencies
    events.latency = -1
    io.savemat(
        negative_latency_fname,
        {
            "EEG": {
                "trials": eeg.trials,
                "srate": eeg.srate,
                "nbchan": eeg.nbchan,
                "data": "test_negative_latency.fdt",
                "epoch": eeg.epoch,
                "event": events,
                "chanlocs": eeg.chanlocs,
                "pnts": eeg.pnts,
            }
        },
        appendmat=False,
        oned_as="row",
    )
    with pytest.raises(ValueError, match="event sample index is negative"):
        with _record_warnings():
            read_raw_eeglab(input_fname=negative_latency_fname, preload=True)

    # test overlapping events
    overlap_fname = tmp_path / "test_overlap_event.set"
    io.savemat(
        overlap_fname,
        {
            "EEG": {
                "trials": eeg.trials,
                "srate": eeg.srate,
                "nbchan": eeg.nbchan,
                "data": "test_overlap_event.fdt",
                "epoch": eeg.epoch,
                "event": [eeg.event[0], eeg.event[0]],
                "chanlocs": eeg.chanlocs,
                "pnts": eeg.pnts,
            }
        },
        appendmat=False,
        oned_as="row",
    )
    shutil.copyfile(base_dir / "test_raw.fdt", overlap_fname.with_suffix(".fdt"))
    read_raw_eeglab(input_fname=overlap_fname, preload=True)

    # test reading file with empty event durations
    empty_dur_fname = tmp_path / "test_empty_durations.set"
    events = deepcopy(eeg.event)
    for ev in events:
        ev.duration = np.array([], dtype="float")

    io.savemat(
        empty_dur_fname,
        {
            "EEG": {
                "trials": eeg.trials,
                "srate": eeg.srate,
                "nbchan": eeg.nbchan,
                "data": "test_negative_latency.fdt",
                "epoch": eeg.epoch,
                "event": events,
                "chanlocs": eeg.chanlocs,
                "pnts": eeg.pnts,
            }
        },
        appendmat=False,
        oned_as="row",
    )
    shutil.copyfile(base_dir / "test_raw.fdt", empty_dur_fname.with_suffix(".fdt"))
    raw = read_raw_eeglab(input_fname=empty_dur_fname, preload=True)
    assert (raw.annotations.duration == 0).all()

    # test reading file when the EEG.data name is wrong
    io.savemat(
        overlap_fname,
        {
            "EEG": {
                "trials": eeg.trials,
                "srate": eeg.srate,
                "nbchan": eeg.nbchan,
                "data": "test_overla_event.fdt",
                "epoch": eeg.epoch,
                "event": [eeg.event[0], eeg.event[0]],
                "chanlocs": eeg.chanlocs,
                "pnts": eeg.pnts,
            }
        },
        appendmat=False,
        oned_as="row",
    )
    with pytest.warns(RuntimeWarning, match="must have changed on disk"):
        read_raw_eeglab(input_fname=overlap_fname, preload=True)

    # raise error when both EEG.data and fdt name from set are wrong
    overlap_fname = tmp_path / "test_ovrlap_event.set"
    io.savemat(
        overlap_fname,
        {
            "EEG": {
                "trials": eeg.trials,
                "srate": eeg.srate,
                "nbchan": eeg.nbchan,
                "data": "test_overla_event.fdt",
                "epoch": eeg.epoch,
                "event": [eeg.event[0], eeg.event[0]],
                "chanlocs": eeg.chanlocs,
                "pnts": eeg.pnts,
            }
        },
        appendmat=False,
        oned_as="row",
    )
    with pytest.raises(FileNotFoundError, match="not find the .fdt data file"):
        read_raw_eeglab(input_fname=overlap_fname, preload=True)

    # test reading file with one channel
    one_chan_fname = tmp_path / "test_one_channel.set"
    io.savemat(
        one_chan_fname,
        {
            "EEG": {
                "trials": eeg.trials,
                "srate": eeg.srate,
                "nbchan": 1,
                "data": np.random.random((1, 3)),
                "epoch": eeg.epoch,
                "event": eeg.epoch,
                "chanlocs": {"labels": "E1", "Y": -6.6069, "X": 6.3023, "Z": -2.9423},
                "times": eeg.times[:3],
                "pnts": 3,
            }
        },
        appendmat=False,
        oned_as="row",
    )
    read_raw_eeglab(input_fname=one_chan_fname, preload=True, montage_units="cm")

    # test reading file with 3 channels - one without position information
    # first, create chanlocs structured array
    ch_names = ["F3", "unknown", "FPz"]
    x, y, z = [1.0, 2.0, np.nan], [4.0, 5.0, np.nan], [7.0, 8.0, np.nan]
    dt = [("labels", "S10"), ("X", "f8"), ("Y", "f8"), ("Z", "f8")]
    nopos_dt = [("labels", "S10"), ("Z", "f8")]
    chanlocs = np.zeros((3,), dtype=dt)
    nopos_chanlocs = np.zeros((3,), dtype=nopos_dt)
    for ind, vals in enumerate(zip(ch_names, x, y, z)):
        for fld in range(4):
            chanlocs[ind][dt[fld][0]] = vals[fld]
            if fld in (0, 3):
                nopos_chanlocs[ind][dt[fld][0]] = vals[fld]
    # In theory this should work and be simpler, but there is an obscure
    # SciPy writing bug that pops up sometimes:
    # nopos_chanlocs = np.array(chanlocs[['labels', 'Z']])

    # test reading channel names but not positions when there is no X (only Z)
    # field in the EEG.chanlocs structure
    nopos_fname = tmp_path / "test_no_chanpos.set"
    io.savemat(
        nopos_fname,
        {
            "EEG": {
                "trials": eeg.trials,
                "srate": eeg.srate,
                "nbchan": 3,
                "data": np.random.random((3, 2)),
                "epoch": eeg.epoch,
                "event": eeg.epoch,
                "chanlocs": nopos_chanlocs,
                "times": eeg.times[:2],
                "pnts": 2,
            }
        },
        appendmat=False,
        oned_as="row",
    )
    # load the file
    raw = read_raw_eeglab(input_fname=nopos_fname, preload=True, montage_units="cm")

    # test that channel names have been loaded but not channel positions
    for i in range(3):
        assert_equal(raw.info["chs"][i]["ch_name"], ch_names[i])
        assert_array_equal(
            raw.info["chs"][i]["loc"][:3], np.array([np.nan, np.nan, np.nan])
        )


@pytest.mark.timeout(60)  # ~60 s on Travis OSX
@testing.requires_testing_data
@pytest.mark.parametrize(
    "fnames",
    [
        epochs_mat_fnames,
        pytest.param(
            epochs_h5_fnames,
            marks=[
                pytest.mark.slowtest,
                pytest.mark.skipif(
                    not _check_pymatreader_installed(strict=False),
                    reason="pymatreader not installed",
                ),
            ],
        ),
    ],
)
def test_io_set_epochs(fnames):
    """Test importing EEGLAB .set epochs files."""
    epochs_fname, epochs_fname_onefile = fnames
    with _record_warnings(), pytest.warns(RuntimeWarning, match="multiple events"):
        epochs = read_epochs_eeglab(epochs_fname)
    with _record_warnings(), pytest.warns(RuntimeWarning, match="multiple events"):
        epochs2 = read_epochs_eeglab(epochs_fname_onefile)
    # one warning for each read_epochs_eeglab because both files have epochs
    # associated with multiple events
    assert_array_equal(epochs.get_data(copy=False), epochs2.get_data(copy=False))


@testing.requires_testing_data
def test_io_set_epochs_events(tmp_path):
    """Test different combinations of events and event_ids."""
    out_fname = tmp_path / "test-eve.fif"
    events = np.array([[4, 0, 1], [12, 0, 2], [20, 0, 3], [26, 0, 3]])
    write_events(out_fname, events)
    event_id = {"S255/S8": 1, "S8": 2, "S255/S9": 3}
    epochs = read_epochs_eeglab(epochs_fname_mat, events, event_id)
    assert_equal(len(epochs.events), 4)
    assert epochs.preload
    assert epochs._bad_dropped
    epochs = read_epochs_eeglab(epochs_fname_mat, out_fname, event_id)
    pytest.raises(ValueError, read_epochs_eeglab, epochs_fname_mat, None, event_id)
    pytest.raises(ValueError, read_epochs_eeglab, epochs_fname_mat, epochs.events, None)


@testing.requires_testing_data
@pytest.mark.filterwarnings("ignore:At least one epoch has multiple events")
@pytest.mark.filterwarnings("ignore:The data contains 'boundary' events")
def test_degenerate(tmp_path):
    """Test some degenerate conditions."""
    # test if .dat file raises an error
    eeg = io.loadmat(epochs_fname_mat, struct_as_record=False, squeeze_me=True)["EEG"]
    eeg.data = "epochs_fname.dat"
    bad_epochs_fname = tmp_path / "test_epochs.set"
    io.savemat(
        bad_epochs_fname,
        {
            "EEG": {
                "trials": eeg.trials,
                "srate": eeg.srate,
                "nbchan": eeg.nbchan,
                "data": eeg.data,
                "epoch": eeg.epoch,
                "event": eeg.event,
                "chanlocs": eeg.chanlocs,
                "pnts": eeg.pnts,
            }
        },
        appendmat=False,
        oned_as="row",
    )
    shutil.copyfile(base_dir / "test_epochs.fdt", tmp_path / "test_epochs.dat")
    pytest.raises(NotImplementedError, read_epochs_eeglab, bad_epochs_fname)

    # error when montage units incorrect
    with pytest.raises(ValueError, match=r"Invalid value"):
        read_epochs_eeglab(epochs_fname_mat, montage_units="mV")

    # warning when head radius too large
    with pytest.warns(RuntimeWarning, match="is above"):
        read_raw_eeglab(raw_fname_chanloc, montage_units="m")

    # warning when head radius too small
    m_fname = tmp_path / "test_montage_m.set"
    _create_eeg_with_scaled_montage_units(raw_fname_chanloc, m_fname, 1e-3)
    with pytest.warns(RuntimeWarning, match="is below"):
        read_raw_eeglab(m_fname, montage_units="mm")


@pytest.mark.parametrize(
    "fname",
    [
        raw_fname_mat,
        raw_fname_onefile_mat,
        # We don't test the h5 variants here because they are implicitly tested
        # in test_io_set_raw
    ],
)
@pytest.mark.filterwarnings("ignore: Complex objects")
@testing.requires_testing_data
def test_eeglab_annotations(fname):
    """Test reading annotations in EEGLAB files."""
    annotations = read_annotations(fname)
    assert len(annotations) == 154
    assert set(annotations.description) == {"rt", "square"}
    assert np.all(annotations.duration == 0.0)


@testing.requires_testing_data
def test_eeglab_read_annotations():
    """Test annotations onsets are timestamps (+ validate some)."""
    annotations = read_annotations(raw_fname_mat)
    validation_samples = [0, 1, 2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31]
    expected_onset = np.array(
        [
            1.00,
            1.69,
            2.08,
            4.70,
            7.71,
            11.30,
            17.18,
            20.20,
            26.12,
            29.14,
            35.25,
            44.30,
            47.15,
        ]
    )
    assert annotations.orig_time is None
    assert_array_almost_equal(
        annotations.onset[validation_samples], expected_onset, decimal=2
    )

    # test if event durations are imported correctly
    raw = read_raw_eeglab(raw_fname_event_duration, preload=True, montage_units="dm")
    # file contains 3 annotations with 0.5 s (64 samples) duration each
    assert_allclose(raw.annotations.duration, np.ones(3) * 0.5)


@testing.requires_testing_data
def test_eeglab_event_from_annot():
    """Test all forms of obtaining annotations."""
    raw_fname_mat = base_dir / "test_raw.set"
    raw_fname = raw_fname_mat
    event_id = {"rt": 1, "square": 2}
    raw1 = read_raw_eeglab(input_fname=raw_fname, preload=False)

    annotations = read_annotations(raw_fname)
    assert len(raw1.annotations) == 154
    raw1.set_annotations(annotations)
    events_b, _ = events_from_annotations(raw1, event_id=event_id)
    assert len(events_b) == 154


def _assert_array_allclose_nan(left, right):
    assert_array_equal(np.isnan(left), np.isnan(right))
    assert_allclose(left[~np.isnan(left)], right[~np.isnan(left)], atol=1e-8)


@pytest.fixture(scope="session")
def three_chanpos_fname(tmp_path_factory):
    """Test file with 3 channels to exercise EEGLAB reader.

    File characteristics
       - ch_names: 'F3', 'unknown', 'FPz'
       - 'FPz' has no position information.
       - the rest is aleatory

    Notes from when this code was factorized:
    # test reading file with one event (read old version)
    """
    fname = str(tmp_path_factory.mktemp("data") / "test_chanpos.set")
    file_conent = dict(
        EEG={
            "trials": 1,
            "nbchan": 3,
            "pnts": 3,
            "epoch": [],
            "event": [],
            "srate": 128,
            "times": np.array([0.0, 0.1, 0.2]),
            "data": np.empty([3, 3]),
            "chanlocs": np.array(
                [
                    (b"F3", 1.0, 4.0, 7.0),
                    (b"unknown", np.nan, np.nan, np.nan),
                    (b"FPz", 2.0, 5.0, 8.0),
                ],
                dtype=[("labels", "S10"), ("X", "f8"), ("Y", "f8"), ("Z", "f8")],
            ),
        }
    )

    io.savemat(file_name=fname, mdict=file_conent, appendmat=False, oned_as="row")

    return fname


@testing.requires_testing_data
def test_position_information(three_chanpos_fname):
    """Test reading file with 3 channels - one without position information."""
    nan = np.nan
    EXPECTED_LOCATIONS_FROM_FILE = (
        np.array(
            [
                [-4.0, 1.0, 7.0, 0.0, 0.0, 0.0, nan, nan, nan, nan, nan, nan],
                [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
                [-5.0, 2.0, 8.0, 0.0, 0.0, 0.0, nan, nan, nan, nan, nan, nan],
            ]
        )
        * 0.01
    )  # 0.01 is to scale cm to meters

    EXPECTED_LOCATIONS_FROM_MONTAGE = np.array(
        [
            [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
            [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
            [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
        ]
    )

    raw = read_raw_eeglab(
        input_fname=three_chanpos_fname, preload=True, montage_units="cm"
    )
    assert_array_equal(
        np.array([ch["loc"] for ch in raw.info["chs"]]), EXPECTED_LOCATIONS_FROM_FILE
    )

    # To accommodate the new behavior so that:
    # read_raw_eeglab(.. montage=montage) and raw.set_montage(montage)
    # behaves the same we need to flush the montage. otherwise we get
    # a mix of what is in montage and in the file
    raw = read_raw_eeglab(
        input_fname=three_chanpos_fname,
        preload=True,
        montage_units="cm",
    ).set_montage(None)  # Flush the montage builtin within input_fname

    _assert_array_allclose_nan(
        np.array([ch["loc"] for ch in raw.info["chs"]]), EXPECTED_LOCATIONS_FROM_MONTAGE
    )


def _create_eeg_with_scaled_montage_units(in_fname, out_fname, scale):
    eeg = io.loadmat(in_fname, struct_as_record=False, squeeze_me=True)["EEG"]

    # test reading file with one event (read old version)
    # chanlocs = deepcopy(eeg.chanlocs)
    chanlocs = eeg.chanlocs
    xyz = np.empty((len(chanlocs), 3))
    labels = []
    for ch_i, loc in enumerate(chanlocs):
        xyz[ch_i] = [loc.X, loc.Y, loc.Z]
        labels.append(loc.labels)
    xyz *= scale
    chanlocs = np.rec.fromarrays(
        [labels, *xyz.T],
        names=["labels", "X", "Y", "Z"],
    )

    fdt = isinstance(eeg.data, str)
    if fdt:
        shutil.copyfile(in_fname.with_suffix(".fdt"), out_fname.with_suffix(".fdt"))
    io.savemat(
        out_fname,
        {
            "EEG": {
                "trials": eeg.trials,
                "srate": eeg.srate,
                "nbchan": eeg.nbchan,
                "data": out_fname.with_suffix(".fdt").name if fdt else eeg.data,
                "epoch": eeg.epoch,
                "event": eeg.event,
                "chanlocs": chanlocs,
                "pnts": eeg.pnts,
            }
        },
        appendmat=False,
        oned_as="row",
    )


@testing.requires_testing_data
def test_estimate_montage_units(tmp_path):
    """Test automatic estimation of montage units."""
    m_fname = tmp_path / "test_montage_m.set"
    _create_eeg_with_scaled_montage_units(raw_fname_chanloc, m_fname, 1e-3)
    cm_fname = tmp_path / "test_montage_cm.set"
    _create_eeg_with_scaled_montage_units(raw_fname_chanloc, cm_fname, 1e-1)
    with pytest.warns(RuntimeWarning, match="The data contains 'boundary' events"):
        # read 3 versions of the same file, with different montage units
        raw_mm = read_raw_eeglab(raw_fname_chanloc, montage_units="auto")
        raw_m = read_raw_eeglab(m_fname, montage_units="auto")
        raw_cm = read_raw_eeglab(cm_fname, montage_units="auto")
    # All locations should be the same if the units are correctly estimated
    assert_allclose(
        np.array([ch["loc"] for ch in raw_mm.info["chs"]]),
        np.array([ch["loc"] for ch in raw_m.info["chs"]]),
    )
    assert_allclose(
        np.array([ch["loc"] for ch in raw_mm.info["chs"]]),
        np.array([ch["loc"] for ch in raw_cm.info["chs"]]),
    )


@testing.requires_testing_data
def test_io_set_raw_2021():
    """Test reading new default file format (no EEG struct)."""
    assert "EEG" not in io.loadmat(raw_fname_2021)
    _test_raw_reader(
        reader=read_raw_eeglab,
        input_fname=raw_fname_2021,
        test_preloading=False,
        preload=True,
    )


@testing.requires_testing_data
def test_read_single_epoch():
    """Test reading raw set file as an Epochs instance."""
    with pytest.raises(ValueError, match="trials less than 2"):
        read_epochs_eeglab(raw_fname_mat)


@testing.requires_testing_data
def test_get_montage_info_with_ch_type():
    """Test that the channel types are properly returned."""
    mat = _readmat(raw_fname_onefile_mat)
    n = len(mat["EEG"]["chanlocs"]["labels"])
    mat["EEG"]["chanlocs"]["type"] = ["eeg"] * (n - 2) + ["eog"] + ["stim"]
    mat["EEG"]["chanlocs"] = _dol_to_lod(mat["EEG"]["chanlocs"])
    mat["EEG"] = Bunch(**mat["EEG"])
    ch_names, ch_types, montage = _get_montage_information(
        mat["EEG"],
        get_pos=False,
        montage_units="mm",
    )
    assert len(ch_names) == len(ch_types) == n
    assert ch_types == ["eeg"] * (n - 2) + ["eog"] + ["stim"]
    assert montage is None

    # test unknown type warning
    mat = _readmat(raw_fname_onefile_mat)
    n = len(mat["EEG"]["chanlocs"]["labels"])
    mat["EEG"]["chanlocs"]["type"] = ["eeg"] * (n - 2) + ["eog"] + ["unknown"]
    mat["EEG"]["chanlocs"] = _dol_to_lod(mat["EEG"]["chanlocs"])
    mat["EEG"] = Bunch(**mat["EEG"])
    with pytest.warns(RuntimeWarning, match="Unknown types found"):
        ch_names, ch_types, montage = _get_montage_information(
            mat["EEG"],
            get_pos=False,
            montage_units="mm",
        )


@testing.requires_testing_data
@pytest.mark.parametrize("has_type", (True, False))
def test_fidsposition_information(monkeypatch, has_type):
    """Test reading file with 3 fiducial locations."""
    if not has_type:

        def get_bad_information(eeg, get_pos, *, montage_units):
            del eeg.chaninfo["nodatchans"]["type"]
            return _get_montage_information(eeg, get_pos, montage_units=montage_units)

        monkeypatch.setattr(
            mne.io.eeglab.eeglab, "_get_montage_information", get_bad_information
        )
    raw = read_raw_eeglab(raw_fname_chanloc_fids, montage_units="cm")
    montage = raw.get_montage()
    pos = montage.get_positions()
    n_eeg = 129
    if not has_type:
        # These should now be estimated from the data
        assert_allclose(pos["nasion"], [0, 0.0997, 0], atol=1e-4)
        assert_allclose(pos["lpa"], -pos["nasion"][[1, 0, 0]])
        assert_allclose(pos["rpa"], pos["nasion"][[1, 0, 0]])
    assert pos["nasion"] is not None
    assert pos["lpa"] is not None
    assert pos["rpa"] is not None
    assert len(pos["nasion"]) == 3
    assert len(pos["lpa"]) == 3
    assert len(pos["rpa"]) == 3
    assert len(raw.info["dig"]) == n_eeg + 3


@testing.requires_testing_data
def test_eeglab_drop_nan_annotations(tmp_path):
    """Test reading file with NaN annotations."""
    pytest.importorskip("eeglabio")
    from eeglabio.raw import export_set

    file_path = tmp_path / "test_nan_anno.set"
    raw = read_raw_eeglab(raw_fname_mat, preload=True)
    data = raw.get_data()
    sfreq = raw.info["sfreq"]
    ch_names = raw.ch_names
    anno = [
        raw.annotations.description,
        raw.annotations.onset,
        raw.annotations.duration,
    ]
    anno[1][0] = np.nan

    export_set(
        str(file_path),
        data,
        sfreq,
        ch_names,
        ch_locs=None,
        annotations=anno,
        ref_channels="common",
        ch_types=np.repeat("EEG", len(ch_names)),
    )

    with pytest.warns(RuntimeWarning, match="1 .* have an onset that is NaN.*"):
        raw = read_raw_eeglab(file_path, preload=True)
