from pathlib import Path
import tempfile

from eeglabio.epochs import export_set
import mne

from numpy.testing import assert_allclose

from pymatreader import read_mat

test_data_dir = Path(mne.__file__).parent / "io" / "tests" / "data"
test_data_fpath = test_data_dir / "test_raw.fif"
events_fpath = test_data_dir / "test-eve.fif"
raw = mne.io.read_raw_fif(test_data_fpath, preload=True).pick("eeg")
events = mne.read_events(events_fpath)

raw.resample(600, events=events)
# let's drop one epoch to prove the point
bad_annot = mne.Annotations(onset=[4.4], duration=[0.1], description=["BAD_manual"])
raw.set_annotations(raw.annotations + bad_annot)
epochs = mne.Epochs(raw, events, preload=True)
data = epochs.get_data()

print(f"Epoch Indices: {epochs.selection}")

want_latencies = epochs.events[:, 0]
with tempfile.TemporaryDirectory() as tmpdir:
    export_set(
        fname=f"{tmpdir}/foo.set",
        data=data,
        sfreq=epochs.info["sfreq"],
        events=epochs.events,
        event_id=epochs.event_id,
        tmin=epochs.tmin,
        tmax=epochs.tmax,
        ch_names=epochs.ch_names,
        epoch_indices=epochs.selection,  # <--- Remove this arg if calling from Main
    )
    epochs_in = mne.read_epochs_eeglab(
        f"{tmpdir}/foo.set", events=epochs.events, event_id=epochs.event_id
        )
    mat = read_mat(f"{tmpdir}/foo.set")  # If you want it
    epochs_in = mne.read_epochs_eeglab(f"{tmpdir}/foo.set")
    got_latencies = epochs_in.events[:, 0]
    assert_allclose(want_latencies, got_latencies)