eeg.py

""" Abstraction for the various supported EEG devices.

    1. Determine which backend to use for the board.
    2.

"""

import sys
import time
import logging
from time import sleep
from multiprocessing import Process
import threading 

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import keyboard

# from PyQt5.QtWidgets import QApplication, QMainWindow, QSizePolicy

from brainflow.board_shim import BoardShim, BoardIds, BrainFlowInputParams
from muselsl import stream, list_muses, record, constants as mlsl_cnsts
from pylsl import StreamInfo, StreamOutlet, StreamInlet, resolve_byprop

from eegnb.devices.rolling_buffer import RollingBuffer
from eegnb.devices.eeg_rt_plot_mpl import EEGRealTimePlotMPL
from eegnb.devices.EMA_Filters import EMA_Filters

from eegnb.devices.utils import (
    get_openbci_usb,
    create_stim_array,
    create_filt_array,
    SAMPLE_FREQS,
    EEG_INDICES,
    EEG_CHANNELS,
)

logger = logging.getLogger(__name__)

# list of brainflow devices
brainflow_devices = [
    "ganglion",
    "ganglion_wifi",
    "cyton",
    "cyton_wifi",
    "cyton_daisy",
    "cyton_daisy_wifi",
    "brainbit",
    "unicorn",
    "synthetic",
    "brainbit",
    "notion1",
    "notion2",
    "freeeeg32",
    "crown",
    "museS_bfn",  # bfn = brainflow with native bluetooth;
    "museS_bfb",  # bfb = brainflow with BLED dongle bluetooth
    "muse2_bfn",
    "muse2_bfb",
    "muse2016_bfn",
    "muse2016_bfb",
]


class EEG:
    device_name: str
    stream_started: bool = False

    def __init__(
        self,
        device=None,
        serial_port=None,
        serial_num=None,
        mac_addr=None,
        other=None,
        ip_addr=None,
    ):
        """The initialization function takes the name of the EEG device and determines whether or not
        the device belongs to the Muse or Brainflow families and initializes the appropriate backend.

        Parameters:
            device (str): name of eeg device used for reading data.
        """
        # determine if board uses brainflow or muselsl backend
        self.device_name = device
        self.serial_num = serial_num
        self.serial_port = serial_port
        self.mac_address = mac_addr
        self.ip_addr = ip_addr
        self.other = other
        self.backend = self._get_backend(self.device_name)
        self.initialize_backend()
        self.n_channels = len(EEG_INDICES[self.device_name])
        self.sfreq = SAMPLE_FREQS[self.device_name]
        self.channels = EEG_CHANNELS[self.device_name]

        self._stop_event = threading.Event() # used to stop threads
        self.filt_data = [] # filtered data

    def initialize_backend(self):
        if self.backend == "brainflow":
            self._init_brainflow()
            self.timestamp_channel = BoardShim.get_timestamp_channel(self.brainflow_id)
        elif self.backend == "muselsl":
            self._init_muselsl()
            self._muse_get_recent()  # run this at initialization to get some
            # stream metadata into the eeg class

    def _get_backend(self, device_name):
        if device_name in brainflow_devices:
            return "brainflow"
        elif device_name in ["muse2016", "muse2", "museS"]:
            return "muselsl"

    #####################
    #   MUSE functions  #
    #####################
    def _init_muselsl(self):
        # Currently there's nothing we need to do here. However keeping the
        # option open to add things with this init method.
        self._muse_recent_inlet = None

    def _start_muse(self, duration):
        if sys.platform in ["linux", "linux2", "darwin"]:
            # Look for muses
            self.muses = list_muses()
            # self.muse = muses[0]

            # Start streaming process
            self.stream_process = Process(
                target=stream, args=(self.muses[0]["address"],)
            )
            self.stream_process.start()

        # Create markers stream outlet
        self.muse_StreamInfo = StreamInfo(
            "Markers", "Markers", 1, 0, "int32", "myuidw43536"
        )
        self.muse_StreamOutlet = StreamOutlet(self.muse_StreamInfo)

        # Start a background process that will stream data from the first available Muse
        print("starting background recording process")
        if self.save_fn:
            print("will save to file: %s" % self.save_fn)
        self.recording = Process(target=record, args=(duration, self.save_fn))
        self.recording.start()

        time.sleep(5)
        self.stream_started = True
        self.push_sample([99], timestamp=time.time())

    def _stop_muse(self):
        pass

    def _muse_push_sample(self, marker, timestamp):
        self.muse_StreamOutlet.push_sample(marker, timestamp)

    def _muse_get_recent(self, n_samples: int = 256, restart_inlet: bool = False):
        if self._muse_recent_inlet and not restart_inlet:
            inlet = self._muse_recent_inlet
        else:
            # Initiate a new lsl stream
            streams = resolve_byprop("type", "EEG", timeout=mlsl_cnsts.LSL_SCAN_TIMEOUT)
            if not streams:
                raise Exception("Couldn't find any stream, is your device connected?")
            inlet = StreamInlet(streams[0], max_chunklen=mlsl_cnsts.LSL_EEG_CHUNK)
            self._muse_recent_inlet = inlet

        info = inlet.info()
        sfreq = info.nominal_srate()
        description = info.desc()
        n_chans = info.channel_count()

        self.sfreq = sfreq
        self.info = info
        self.n_chans = n_chans

        timeout = (n_samples / sfreq) + 0.5
        samples, timestamps = inlet.pull_chunk(timeout=timeout, max_samples=n_samples)

        samples = np.array(samples)
        timestamps = np.array(timestamps)

        ch = description.child("channels").first_child()
        ch_names = [ch.child_value("label")]
        for i in range(n_chans):
            ch = ch.next_sibling()
            lab = ch.child_value("label")
            if lab != "":
                ch_names.append(lab)

        df = pd.DataFrame(samples, index=timestamps, columns=ch_names)
        return df

    ##########################
    #   BrainFlow functions  #
    ##########################
    def _init_brainflow(self):
        """This function initializes the brainflow backend based on the input device name. It calls
        a utility function to determine the appropriate USB port to use based on the current operating system.
        Additionally, the system allows for passing a serial number in the case that they want to use either
        the BraintBit or the Unicorn EEG devices from the brainflow family.

        Parameters:
             serial_num (str or int): serial number for either the BrainBit or Unicorn devices.
        """
        # Initialize brainflow parameters
        self.brainflow_params = BrainFlowInputParams()

        if self.device_name == "ganglion":
            self.brainflow_id = BoardIds.GANGLION_BOARD.value
            if self.serial_port is None:
                self.brainflow_params.serial_port = get_openbci_usb()
            # set mac address parameter in case
            if self.mac_address is None:
                print("No MAC address provided, attempting to connect without one")
            else:
                self.brainflow_params.mac_address = self.mac_address

        elif self.device_name == "ganglion_wifi":
            self.brainflow_id = BoardIds.GANGLION_WIFI_BOARD.value
            if self.ip_addr is not None:
                self.brainflow_params.ip_address = self.ip_addr
                self.brainflow_params.ip_port = 6677

        elif self.device_name == "cyton":
            self.brainflow_id = BoardIds.CYTON_BOARD.value
            if self.serial_port is None:
                self.brainflow_params.serial_port = get_openbci_usb()

        elif self.device_name == "cyton_wifi":
            self.brainflow_id = BoardIds.CYTON_WIFI_BOARD.value
            if self.ip_addr is not None:
                self.brainflow_params.ip_address = self.ip_addr
                self.brainflow_params.ip_port = 6677

        elif self.device_name == "cyton_daisy":
            self.brainflow_id = BoardIds.CYTON_DAISY_BOARD.value
            if self.serial_port is None:
                self.brainflow_params.serial_port = get_openbci_usb()

        elif self.device_name == "cyton_daisy_wifi":
            self.brainflow_id = BoardIds.CYTON_DAISY_WIFI_BOARD.value
            if self.ip_addr is not None:
                self.brainflow_params.ip_address = self.ip_addr

        elif self.device_name == "brainbit":
            self.brainflow_id = BoardIds.BRAINBIT_BOARD.value

        elif self.device_name == "unicorn":
            self.brainflow_id = BoardIds.UNICORN_BOARD.value

        elif self.device_name == "callibri_eeg":
            self.brainflow_id = BoardIds.CALLIBRI_EEG_BOARD.value
            if self.other:
                self.brainflow_params.other_info = str(self.other)

        elif self.device_name == "notion1":
            self.brainflow_id = BoardIds.NOTION_1_BOARD.value

        elif self.device_name == "notion2":
            self.brainflow_id = BoardIds.NOTION_2_BOARD.value

        elif self.device_name == "crown":
            self.brainflow_id = BoardIds.CROWN_BOARD.value

        elif self.device_name == "freeeeg32":
            self.brainflow_id = BoardIds.FREEEEG32_BOARD.value
            if self.serial_port is None:
                self.brainflow_params.serial_port = get_openbci_usb()

        elif self.device_name == "museS_bfn":
            self.brainflow_id = BoardIds.MUSE_S_BOARD.value

        elif self.device_name == "museS_bfb":
            self.brainflow_id = BoardIds.MUSE_S_BLED_BOARD.value

        elif self.device_name == "muse2_bfn":
            self.brainflow_id = BoardIds.MUSE_2_BOARD.value

        elif self.device_name == "muse2_bfb":
            self.brainflow_id = BoardIds.MUSE_2_BLED_BOARD.value

        elif self.device_name == "muse2016_bfn":
            self.brainflow_id = BoardIds.MUSE_2016_BOARD.value

        elif self.device_name == "muse2016_bfb":
            self.brainflow_id = BoardIds.MUSE_2016_BLED_BOARD.value

        elif self.device_name == "synthetic":
            self.brainflow_id = BoardIds.SYNTHETIC_BOARD.value

        # some devices allow for an optional serial number parameter for better connection
        if self.serial_num:
            serial_num = str(self.serial_num)
            self.brainflow_params.serial_number = serial_num

        if self.serial_port:
            serial_port = str(self.serial_port)
            self.brainflow_params.serial_port = serial_port

        # Initialize board_shim
        self.sfreq = BoardShim.get_sampling_rate(self.brainflow_id)
        self.board = BoardShim(self.brainflow_id, self.brainflow_params)
        self.board.prepare_session()

    def _start_brainflow(self):
        # only start stream if non exists
        if not self.stream_started:
            self.board.start_stream()

        self.stream_started = True

        # wait for signal to settle
        if (self.device_name.find("cyton") != -1) or (
            self.device_name.find("ganglion") != -1
        ):
            # wait longer for openbci cyton / ganglion
            sleep(10)
        else:
            sleep(10) # also wait longer for unicorn

    def _stop_brainflow(self):
        """This functions kills the brainflow backend and saves the data to a CSV file."""

        # Collect session data and kill session
        data = self.board.get_board_data()  # will clear board buffer
        self.board.stop_stream()
        self.board.release_session()

        # Extract relevant metadata from board
        ch_names, eeg_data, timestamps = self._brainflow_extract(data)

        # Create a column for the stimuli to append to the EEG data
        stim_array = create_stim_array(timestamps, self.markers)
        timestamps = timestamps[..., None]

        # Add an additional dimension so that shapes match
        total_data = np.append(timestamps, eeg_data, 1)

        # Append the stim array to data.
        total_data = np.append(total_data, stim_array, 1)

        # Subtract five seconds of settling time from beginning
        total_data = total_data[5 * self.sfreq :]
        data_df = pd.DataFrame(total_data, columns=["timestamps"] + ch_names + ["stim"])
        data_df.to_csv(self.save_fn, index=False)

    def _brainflow_extract(self, data): 
        """
        Formats the data returned from brainflow to get
        ch_names; list of channel names
        eeg_data: NDArray of eeg samples
        timestamps: NDArray of timestamps
        """

        # transform data for saving
        data = data.T  # transpose data

        # get the channel names for EEG data
        if (
            self.brainflow_id == BoardIds.GANGLION_BOARD.value
            or self.brainflow_id == BoardIds.GANGLION_WIFI_BOARD.value
        ):
            # if a ganglion is used, use recommended default EEG channel names
            ch_names = ["fp1", "fp2", "tp7", "tp8"]
        elif self.brainflow_id == BoardIds.FREEEEG32_BOARD.value:
            ch_names = [f"eeg_{i}" for i in range(0, 32)]
        else:
            # otherwise select eeg channel names via brainflow API
            ch_names = BoardShim.get_eeg_names(self.brainflow_id)

        # pull EEG channel data via brainflow API
        eeg_data = data[:, BoardShim.get_eeg_channels(self.brainflow_id)]
        timestamps = data[:, BoardShim.get_timestamp_channel(self.brainflow_id)]

        return ch_names, eeg_data, timestamps

    def _brainflow_push_sample(self, marker):
        last_timestamp = self.board.get_current_board_data(1)[self.timestamp_channel][0]
        self.markers.append([marker, last_timestamp])

    def _brainflow_get_recent(self, n_samples=256):

        # initialize brainflow if not set
        if self.board == None:
            self._init_brainflow()

        # start branflow stream
        self._start_brainflow()

        # get the latest data
        data = self.board.get_current_board_data(n_samples)

        ch_names, eeg_data, timestamps = self._brainflow_extract(data)

        eeg_data = np.array(eeg_data)
        timestamps = np.array(timestamps)

        df = pd.DataFrame(eeg_data, index=timestamps, columns=ch_names)

        return df

    #################################
    #   Highlevel device functions  #
    #################################

    def start(self, fn, duration=None):
        """Starts the EEG device based on the defined backend.

        Parameters:
            fn (str): name of the file to save the sessions data to.
        """
        if fn:
            self.save_fn = fn

        if self.backend == "brainflow":
            self._start_brainflow()
            self.markers = []
        elif self.backend == "muselsl":
            self._start_muse(duration)
    

        # muse backend, not applicable until we have muse
        elif self.backend == "muselsl":
            self._start_muse(duration)
    
    def push_sample(self, marker, timestamp): # add code here to log in keyboard input
        """
        Universal method for pushing a marker and its timestamp to store alongside the EEG data.

        Parameters:
            marker (int): marker number for the stimuli being presented.
            timestamp (float): timestamp of stimulus onset from time.time() function.
        """
        if self.backend == "brainflow":
            self._brainflow_push_sample(marker=marker)
        elif self.backend == "muselsl":
            self._muse_push_sample(marker=marker, timestamp=timestamp)

    def stop(self):
        if self.backend == "brainflow":
            self._stop_brainflow()
        elif self.backend == "muselsl":
            pass

    def get_recent(self, n_samples: int = 256):
        """
        Usage:
        -------
        from eegnb.devices.eeg import EEG
        this_eeg = EEG(device='museS')
        df_rec = this_eeg.get_recent()
        """

        if self.backend == "brainflow":
            df = self._brainflow_get_recent(n_samples)
        elif self.backend == "muselsl":
            df = self._muse_get_recent(n_samples)
        else:
            raise ValueError(f"Unknown backend {self.backend}")

        # Sort out the sensor coils
        sorted_cols = sorted(df.columns)
        df = df[sorted_cols]

        return df
    
    ##############################################################
    #   Unicorn functions by JHUBCIS for streaming and analysis  #
    ##############################################################


    # def filt_eeg(self, eeg_data, bp_fc_high = 60, bp_fc_low = 1, n_fc = 60):
    #     '''bandpass and notch filter eeg data, incomplete'''
    #     sfreq = self.sfreq
    #     emaFilt = EMA_Filters()
    #     if emaFilt:
    #         print("ema filter initiated: bandpass filter", bp_fc_low, "to", bp_fc_high, "Hz, notch filter at", n_fc, "Hz")
        
    #     if self.stream_started:
    #         eeg_data_filt = emaFilt.BPF(eeg_data, bp_fc_low, bp_fc_high, sfreq) #bandpass filter
    #         eeg_data_filt = emaFilt.Notch(eeg_data_filt, n_fc, sfreq) #notch filter


    def _stop_brainflow_save_filt(self):
        """This functions kills the brainflow backend and saves the raw and filtered data to a CSV file."""

        # Collect session data and kill session
        data = self.board.get_board_data()  # will clear board buffer
        self.board.stop_stream()
        self.board.release_session()

        # Extract relevant metadata from board
        ch_names, eeg_data, timestamps = self._brainflow_extract(data)

        # Create a column for the stimuli to append to the EEG data
        stim_array = create_stim_array(timestamps, self.markers)
        filt_array = create_filt_array(timestamps, self.filt_data, self.n_channels)
        timestamps = timestamps[..., None]

        # Add an additional dimension so that shapes match
        total_data = np.append(timestamps, eeg_data, 1)

        total_data = np.append(total_data, filt_array, 1)

        # Append the stim array to data.
        total_data = np.append(total_data, stim_array, 1)

        # Subtract five seconds of settling time from beginning
        total_data = total_data[5 * self.sfreq :]
        data_df = pd.DataFrame(total_data, columns=["timestamps"] + ch_names + [s + "_filt" for s in ch_names] + ["stim"])
        data_df.to_csv(self.save_fn, index=False)

    # Stream raw unicorn input and prints in terminal. press `q` to quit and save data
    def stream(self, fn, duration=None):
        print("press 'q' to stop stream and save data")    
        self.save_fn = fn

        # brainflow backend, works with unicorn
        if self.backend == "brainflow":
            self._start_brainflow()
            self.markers = []
            running = True

            while running:
                # get the latest data
                data = self.board.get_current_board_data(1)

                ch_names, eeg_data, timestamps = self._brainflow_extract(data)

                eeg_data = np.array(eeg_data)
                timestamps = np.array(timestamps)

                df = pd.DataFrame(eeg_data, index=timestamps, columns=ch_names)
                print (df) # JHUBCIS: modify output accordingly

                if keyboard.is_pressed('q'):
                    print("stopping stream and saving data")
                    self._stop_brainflow()
                    print(self.save_fn) 
                    running = not running

    # streams band-passed & notch filtered unicorn output and plots in real-time. recorded data saved to SCV st the end.
    def stream_plot(self, fn, buffer_time=5, bp_fc_high = 60, bp_fc_low = 1, n_fc = 60): 
        print("first press 'q' to stop stream and save data, then close plot window to end program")    
        self.save_fn = fn

        sfreq = self.sfreq
        channel_names = self.channels
        num_channels = self.n_channels
        print("channel names = ", channel_names, "\nsampling frequency = ", sfreq)

        emaFilt = EMA_Filters()
        if emaFilt:
            print("ema filter initiated: bandpass filter", bp_fc_low, "to", bp_fc_high, "Hz, notch filter at", n_fc, "Hz")

        rolling_buffer = RollingBuffer(buffer_time, sfreq, num_channels)
        if rolling_buffer:
            print("rolling buffer initiated with buffer time of", buffer_time, "seconds")
        
        
        plotter = EEGRealTimePlotMPL(rolling_buffer, channel_names)
        
        # brainflow backend, works with unicorn
        if self.backend == "brainflow":
            self._start_brainflow()
            self.markers = []
            print("brainflow from unicorn started")
        
        # thread to stream filtered EEG data
        def eeg_stream_thread(rolling_buffer):
            while self.stream_started and not self._stop_event.is_set():
                data = self.board.get_current_board_data(1)
                _, eeg_data, timestamps = self._brainflow_extract(data)
                eeg_data_filt = emaFilt.BPF(eeg_data, bp_fc_low, bp_fc_high, sfreq) #bandpass filter
                eeg_data_filt = emaFilt.Notch(eeg_data_filt, n_fc, sfreq) #notch filter
                if len(eeg_data) > 0 and len(timestamps) > 0: # only update buffer if neither is empty
                    rolling_buffer.update(eeg_data_filt, timestamps)
                    last_timestamp = data[self.timestamp_channel][0]
                    # print([eeg_data_filt[0].tolist(), last_timestamp])
                    self.filt_data.append([eeg_data_filt[0].tolist(), last_timestamp])
                else:
                    time.sleep(1)
                    continue
                if self._stop_event.is_set():
                    # print("thread terminated")
                    break
        # thread to end streaming
        def listen_for_q():
            keyboard.wait('q')
            self._stop_event.set()  # Signal the eeg_data_thread to stop
            print("Stop brainflow, saving raw and filtered data")
            self._stop_brainflow_save_filt() #edit this line to save filtered data
            # print(self.filt_data[-50:])
            print("Data saved at:")
            print(self.save_fn)
        

        q_thread = threading.Thread(target=listen_for_q)
        q_thread.daemon = True
        q_thread.start()

        eeg_data_thread = threading.Thread(target=eeg_stream_thread, args=(rolling_buffer,))
        eeg_data_thread.daemon = True
        eeg_data_thread.start()
        print("eeg_data_thread initiated")
        
        plotter.animate()

        # make sure threads have finished before exiting stream_plot
        q_thread.join()
        eeg_data_thread.join()