From 8d39c1fdb3d9368913bb693ee695d85d7b982e91 Mon Sep 17 00:00:00 2001
From: Clemens Brunner <clemens.brunner@gmail.com>
Date: Mon, 19 Feb 2024 18:15:52 +0100
Subject: [PATCH] Include export functions from pybv

---
 mne/export/_brainvision.py | 146 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 144 insertions(+), 2 deletions(-)

diff --git a/mne/export/_brainvision.py b/mne/export/_brainvision.py
index 0da7647ecb7..12212c91343 100644
--- a/mne/export/_brainvision.py
+++ b/mne/export/_brainvision.py
@@ -4,11 +4,153 @@
 # Copyright the MNE-Python contributors.
 
 import os
+from pathlib import Path
 
-from ..utils import _check_pybv_installed
+import numpy as np
+
+from mne.channels.channels import _unit2human
+from mne.io.constants import FIFF
+from mne.utils import _check_pybv_installed, warn
 
 _check_pybv_installed()
-from pybv._export import _export_mne_raw  # noqa: E402
+from pybv import write_brainvision  # noqa: E402
+
+
+def _export_mne_raw(*, raw, fname, events=None, overwrite=False):
+    """Export raw data from MNE-Python.
+
+    Parameters
+    ----------
+    raw : mne.io.Raw
+        The raw data to export.
+    fname : str | pathlib.Path
+        The name of the file where raw data will be exported to. Must end with
+        ``".vhdr"``, and accompanying *.vmrk* and *.eeg* files will be written inside
+        the same directory.
+    events : np.ndarray | None
+        Events to be written to the marker file (*.vmrk*). If array, must be in
+        `MNE-Python format <https://mne.tools/stable/glossary.html#term-events>`_. If
+        ``None`` (default), events will be written based on ``raw.annotations``.
+    overwrite : bool
+        Whether or not to overwrite existing data. Defaults to ``False``.
+
+    """
+    # prepare file location
+    if not str(fname).endswith(".vhdr"):
+        raise ValueError("`fname` must have the '.vhdr' extension for BrainVision.")
+    fname = Path(fname)
+    folder_out = fname.parents[0]
+    fname_base = fname.stem
+
+    # prepare data from raw
+    data = raw.get_data()  # gets data starting from raw.first_samp
+    sfreq = raw.info["sfreq"]  # in Hz
+    meas_date = raw.info["meas_date"]  # datetime.datetime
+    ch_names = raw.ch_names
+
+    # write voltage units as micro-volts and all other units without scaling
+    # write units that we don't know as n/a
+    unit = []
+    for ch in raw.info["chs"]:
+        if ch["unit"] == FIFF.FIFF_UNIT_V:
+            unit.append("µV")
+        elif ch["unit"] == FIFF.FIFF_UNIT_CEL:
+            unit.append("°C")
+        else:
+            unit.append(_unit2human.get(ch["unit"], "n/a"))
+    unit = [u if u != "NA" else "n/a" for u in unit]
+
+    # enforce conversion to float32 format
+    # XXX: Could add a feature that checks data and optimizes `unit`, `resolution`, and
+    #      `format` so that raw.orig_format could be retained if reasonable.
+    if raw.orig_format != "single":
+        warn(
+            f"Encountered data in '{raw.orig_format}' format. "
+            "Converting to float32.",
+            RuntimeWarning,
+        )
+
+    fmt = "binary_float32"
+    resolution = 0.1
+
+    # handle events
+    # if we got an ndarray, this is in MNE-Python format
+    msg = "`events` must be None or array in MNE-Python format."
+    if events is not None:
+        # Subtract raw.first_samp because brainvision marks events starting from the
+        # first available data point and ignores the raw.first_samp
+        assert isinstance(events, np.ndarray), msg
+        assert events.ndim == 2, msg
+        assert events.shape[-1] == 3, msg
+        events[:, 0] -= raw.first_samp
+        events = events[:, [0, 2]]  # reorder for pybv required order
+
+    # else, prepare pybv style events from raw.annotations
+    else:
+        events = _mne_annots2pybv_events(raw)
+
+    # no information about reference channels in mne currently
+    ref_ch_names = None
+
+    # write to BrainVision
+    write_brainvision(
+        data=data,
+        sfreq=sfreq,
+        ch_names=ch_names,
+        ref_ch_names=ref_ch_names,
+        fname_base=fname_base,
+        folder_out=folder_out,
+        overwrite=overwrite,
+        events=events,
+        resolution=resolution,
+        unit=unit,
+        fmt=fmt,
+        meas_date=meas_date,
+    )
+
+
+def _mne_annots2pybv_events(raw):
+    """Convert mne Annotations to pybv events."""
+    events = []
+    for annot in raw.annotations:
+        # handle onset and duration: seconds to sample, relative to
+        # raw.first_samp / raw.first_time
+        onset = annot["onset"] - raw.first_time
+        onset = raw.time_as_index(onset).astype(int)[0]
+        duration = int(annot["duration"] * raw.info["sfreq"])
+
+        # triage type and description
+        # defaults to type="Comment" and the full description
+        etype = "Comment"
+        description = annot["description"]
+        for start in ["Stimulus/S", "Response/R", "Comment/"]:
+            if description.startswith(start):
+                etype = start.split("/")[0]
+                description = description.replace(start, "")
+                break
+
+        if etype in ["Stimulus", "Response"] and description.strip().isdigit():
+            description = int(description.strip())
+        else:
+            # if cannot convert to int, we must use this as "Comment"
+            etype = "Comment"
+
+        event_dict = dict(
+            onset=onset,  # in samples
+            duration=duration,  # in samples
+            description=description,
+            type=etype,
+        )
+
+        if "ch_names" in annot:
+            # handle channels
+            channels = list(annot["ch_names"])
+            event_dict["channels"] = channels
+
+        # add a "pybv" event
+        events += [event_dict]
+
+    return events
 
 
 def _export_raw(fname, raw, overwrite):