kraken/backend/core/mi/utils.py

"""想象运动工具类
"""
from enum import Enum

from core.sig_chain.device.connector_interface import Device


class Mark(Enum):
    REST = "rest"
    MI = "mi"


class SelectedChannel():

    def get_channel_names(self, position:str):
        if position == "左手":
            return self.left_hand_channel
        elif position == "右手":
            return self.right_hand_channel
        else:
            return self.foot_channel

    def get_channel_ids(self, position:str, channel_labels: list):
        selected = []
        if position == "左手":
            selected = self.left_hand_channel
        elif position == "右手":
            selected = self.right_hand_channel
        else:
            selected = self.foot_channel
        try:
            selected_ids = [channel_labels.index(item) for item in selected]
            return selected_ids
        except ValueError as exc:
            # pylint: disable=line-too-long
            raise Exception(
                f"Some selected channel({selected}) missing in input channel({channel_labels})"
            ) from exc
            # pylint: enable=line-too-long


class SelectedCspChannel(SelectedChannel):

    def __init__(self, device: Device):
        if device == Device.NEO:
            self.left_hand_channel = ["C4", "FC4", "CP2", "CP6"]
            self.right_hand_channel = ["C3", "FC3", "CP5", "CP1"]
        else:
            self.left_hand_channel = [
                "Fz", "F4", "F8", "Cz", "C4", "T4", "Pz", "P4", "T6"
            ]
            self.right_hand_channel = [
                "F7", "F3", "Fz", "T3", "C3", "Cz", "T5", "P3", "Pz"
            ]
            self.foot_channel = [
                "F3", "Fz", "F4", "C3", "Cz", "C4", "P3", "Pz", "P4"
            ]


class SelectedCspPlotChannel(SelectedChannel):

    def __init__(self, device: Device):
        if device == Device.NEO:
            self.left_hand_channel = [
                "C3", "FC3", "CP5", "CP1", "C4", "FC4", "CP2", "CP6"
            ]
            self.right_hand_channel = self.left_hand_channel
        else:
            self.left_hand_channel = [
                "T6", "P4", "Pz", "F8", "F4", "Fp1", "Cz", "F7", "F3", "C3",
                "T3", "Oz", "O1", "O2", "Fz", "C4", "T4", "Fp2", "T5", "P3"
            ]
            self.right_hand_channel = self.left_hand_channel
            self.foot_channel = self.left_hand_channel


class SelectedPsdChannel(SelectedChannel):

    def __init__(self, device: Device):
        self.left_hand_channel = ["C4"]
        self.right_hand_channel = ["C3"]
        if device != Device.NEO:
            self.foot_channel = ["Cz"]


class SelectedPsdPlotChannel(SelectedChannel):

    def __init__(self, device: Device):
        channels = ["C3","C4"]
        self.left_hand_channel = channels
        self.right_hand_channel = channels
        if device != Device.NEO:
            self.foot_channel = ["C3", "Cz", "C4"]


class SelectedErdsChannel(SelectedChannel):
    def __init__(self, device: Device):
        self.left_hand_channel = ["C3", "C4"]
        self.right_hand_channel = ["C3", "C4"]
        if device != Device.NEO:
            self.foot_channel = ["C3", "Cz", "C4"]


class SelectedWpliChannel(SelectedChannel):

    def __init__(self, device: Device):
        if device == Device.NEO:
            self.left_hand_channel = [
                "C3", "FC3", "CP5", "CP1", "C4", "FC4", "CP2", "CP6"
            ]
        else:
            self.left_hand_channel = [
                "T6", "P4", "Pz", "F8", "F4", "Fp1", "Cz", "F7", "F3", "C3",
                "T3", "Oz", "O1", "O2", "Fz", "C4", "T4", "Fp2", "T5", "P3"
            ]
        self.right_hand_channel = self.left_hand_channel
        self.foot_channel = self.left_hand_channel