mxbiflow 0.1.1__py3-none-any.whl

mxbiflow/tasks/two_alternative_choice/stages/size_reduction_stage/size_reduction_stage.py

@@ -0,0 +1,122 @@
+from typing import TYPE_CHECKING, Final
+
+from mxbi.data_logger import DataLogger, DataLoggerType
+from mxbi.tasks.two_alternative_choice.models import PersistentData, Result
+from mxbi.tasks.two_alternative_choice.stages.size_reduction_stage.size_reduction_models import (
+    config,
+)
+from mxbi.tasks.two_alternative_choice.tasks.touch.touch_scene import TwoACTouchScene
+from mxbi.utils.logger import logger
+
+if TYPE_CHECKING:
+    from mxbi.models.animal import AnimalState, ScheduleCondition
+    from mxbi.models.session import SessionState
+    from mxbi.models.task import Feedback
+    from mxbi.tasks.two_alternative_choice.stages.size_reduction_stage.size_reduction_models import (
+        SizeReductionStageConfig,
+    )
+    from mxbi.theater import Theater
+
+_presistent_data: dict[str, PersistentData] = {}
+
+
+class TWOACSizeReductionStage:
+    STAGE_NAME: Final[str] = "twoac_size_reduction_stage"
+
+    def __init__(
+        self,
+        theater: "Theater",
+        session_state: "SessionState",
+        animal_state: "AnimalState",
+    ) -> None:
+        self._theater = theater
+        self._session_state = session_state
+        self._animal_state = animal_state
+
+        self._stage_config = self._load_stage_config(animal_state.name)
+
+        _levels_config = self._stage_config.levels_table[animal_state.level]
+        _config = self._stage_config.trial_config
+
+        _config.level = _levels_config.level
+        _config.stimulation_size = _levels_config.stimulation_size
+
+        self._data_logger = DataLogger(
+            self._session_state,
+            self._animal_state.name,
+            self.STAGE_NAME,
+            DataLoggerType.JSONL,
+        )
+
+        self._presistent_data = _presistent_data.get(self._animal_state.name)
+
+        if self._presistent_data is None:
+            self._presistent_data = PersistentData(
+                rewards=0,
+                correct=0,
+                incorrect=0,
+                timeout=0,
+            )
+            _presistent_data[self._animal_state.name] = self._presistent_data
+
+        self._task = TwoACTouchScene(
+            theater,
+            session_state.session_config,
+            animal_state,
+            session_state.session_config.screen_type,
+            _config,
+            self._presistent_data,
+        )
+
+    def start(self) -> "Feedback":
+        trial_data = self._task.start()
+        self._data_logger.save(trial_data.model_dump())
+
+        feedback = self._handle_result(trial_data.result)
+        logger.debug(
+            f"{self.STAGE_NAME}: "
+            f"session_id={self._session_state.session_id}, "
+            f"animal_name={self._animal_state.name}, "
+            f"animal_level={self._animal_state.level}, "
+            f"state_name={self.STAGE_NAME}, "
+            f"result={trial_data}, "
+            f"feedback={feedback}"
+        )
+
+        return feedback
+
+    def _load_stage_config(self, monkey: str) -> "SizeReductionStageConfig":
+        stage_config = config.root.get(monkey) or config.root.get("default")
+        if stage_config is None:
+            raise ValueError("No default stage config found")
+        return stage_config
+
+    def _handle_result(self, result: "Result") -> "Feedback":
+        feedback = False
+        match result:
+            case Result.CORRECT:
+                _presistent_data[self._animal_state.name].correct += 1
+                feedback = True
+            case Result.INCORRECT:
+                _presistent_data[self._animal_state.name].incorrect += 1
+                feedback = False
+            case Result.TIMEOUT:
+                _presistent_data[self._animal_state.name].timeout += 1
+                feedback = False
+            case Result.CANCEL:
+                feedback = False
+
+        return feedback
+
+    def quit(self) -> None:
+        self._task.cancle()
+
+    def on_idle(self) -> None:
+        self._task.cancle()
+
+    def on_return(self) -> None:
+        self._task.cancle()
+
+    @property
+    def condition(self) -> "ScheduleCondition | None":
+        return self._stage_config.condition

mxbiflow/tasks/two_alternative_choice/tasks/touch/touch_models.py

@@ -0,0 +1,19 @@
+from mxbi.tasks.two_alternative_choice.models import (
+    BaseDataToShow,
+    BaseTrialConfig,
+    BaseTrialData,
+)
+
+
+class TrialConfig(BaseTrialConfig):
+    stimulus_freq: int
+    stimulus_freq_duration: int
+
+    stimulus_interval: int
+
+
+class TrialData(BaseTrialData):
+    trial_config: TrialConfig
+
+
+class DataToShow(BaseDataToShow): ...

mxbiflow/tasks/two_alternative_choice/tasks/touch/touch_scene.py

@@ -0,0 +1,249 @@
+from concurrent.futures import Future
+from datetime import datetime
+from math import ceil
+from tkinter import CENTER, Canvas, Event
+from typing import TYPE_CHECKING, Final
+
+from mxbi.tasks.two_alternative_choice.assets.starter import Starter
+from mxbi.tasks.two_alternative_choice.models import Result, TouchEvent
+from mxbi.tasks.two_alternative_choice.tasks.touch.touch_models import (
+    DataToShow,
+    TrialData,
+)
+from mxbi.utils.aplayer import ToneConfig
+from mxbi.utils.tkinter.components.canvas_with_border import CanvasWithInnerBorder
+from mxbi.utils.tkinter.components.showdata_widget import ShowDataWidget
+from numpy import int16
+from numpy.typing import NDArray
+
+if TYPE_CHECKING:
+    from mxbi.models.animal import AnimalState
+    from mxbi.models.session import ScreenConfig, SessionConfig
+    from mxbi.tasks.two_alternative_choice.models import PersistentData
+    from mxbi.tasks.two_alternative_choice.tasks.touch.touch_models import TrialConfig
+    from mxbi.theater import Theater
+
+
+class TwoACTouchScene:
+    def __init__(
+        self,
+        theater: "Theater",
+        session: "SessionConfig",
+        animal_state: "AnimalState",
+        screen_type: "ScreenConfig",
+        trial_config: "TrialConfig",
+        persistent_data: "PersistentData",
+    ) -> None:
+        self._theater: "Final[Theater]" = theater
+        self._animal_state: "Final[AnimalState]" = animal_state
+        self._screen_type: "Final[ScreenConfig]" = screen_type
+        self._trial_config: "Final[TrialConfig]" = trial_config
+        self._persistent_data: "Final[PersistentData]" = persistent_data
+
+        self._tone = self._prepare_stimulus()
+        self._standard_reward_stimulus = self._theater.new_standard_reward_stimulus(
+            self._trial_config.stimulus_duration
+        )
+
+        self._on_trial_start()
+
+    # region public api
+    def start(self) -> "TrialData":
+        self._theater.mainloop()
+
+        return self._data
+
+    def cancle(self) -> None:
+        self._data.result = Result.CANCEL
+        self._theater.aplayer.stop()
+        self._on_trial_end()
+
+    # endregion
+
+    # region lifecycle
+    def _on_trial_start(self) -> None:
+        self._create_view()
+        self._init_data()
+        self._bind_events()
+
+    def _on_inter_trial(self) -> None:
+        self._background.after(
+            self._trial_config.inter_trial_interval, self._on_trial_end
+        )
+
+    def _on_trial_end(self) -> None:
+        self._background.destroy()
+        self._theater.root.quit()
+
+    # endregion
+
+    # region views
+    def _create_view(self) -> None:
+        self._create_background()
+        self._create_show_data_widget()
+        self._create_target()
+
+    def _create_background(self) -> None:
+        self._background = CanvasWithInnerBorder(
+            master=self._theater.root,
+            bg="black",
+            width=self._screen_type.width,
+            height=self._screen_type.height,
+            border_width=40,
+        )
+
+        self._background.place(relx=0.5, rely=0.5, anchor="center")
+
+    def _create_show_data_widget(self) -> None:
+        self._show_data_widget = ShowDataWidget(self._background)
+        self._show_data_widget.place(relx=0, rely=1, anchor="sw")
+        data = DataToShow(
+            name=self._animal_state.name,
+            id=self._animal_state.trial_id,
+            level_id=self._animal_state.current_level_trial_id,
+            level=self._trial_config.level,
+            rewards=self._persistent_data.rewards,
+            correct=self._animal_state.correct_trial,
+            incorrect=self._persistent_data.incorrect,
+            timeout=self._persistent_data.timeout,
+        )
+        self._show_data_widget.show_data(data.model_dump())
+
+    def _create_target(self) -> None:
+        xshift = 240
+        xcenter = self._screen_type.width * 0.5 + xshift
+        ycenter = self._screen_type.height * 0.5
+
+        self._trigger_canvas = Starter(
+            self._background, self._trial_config.stimulation_size
+        )
+        self._trigger_canvas.place(x=xcenter, y=ycenter, anchor="center")
+
+    def _create_wrong_view(self) -> None:
+        self._trigger_canvas = Canvas(
+            self._background,
+            bg="grey",
+            width=self._screen_type.width,
+            height=self._screen_type.height,
+        )
+        self._trigger_canvas.place(relx=0.5, rely=0.5, anchor=CENTER)
+
+    # endregion
+
+    # region event binding
+    def _bind_events(self) -> None:
+        # Manual reward
+        self._background.focus_set()
+        self._background.bind("<r>", lambda e: self._give_standard_stimulus())
+        self._background.bind("<s>", lambda e: self._theater.caputre(self._background))
+
+        # Trigger event
+        self._background.bind("<ButtonPress>", self._on_background_touched)
+        self._trigger_canvas.bind("<ButtonPress>", self._on_touched)
+
+        # Timeout event
+        self._trigger_canvas.after(self._trial_config.time_out, self._on_timeout)
+
+    # endregion
+
+    # region event handlers
+    def _on_touched(self, event: Event) -> None:
+        self._data.touch_events.append(
+            TouchEvent(time=datetime.now().timestamp(), x=event.x, y=event.y)
+        )
+        self._background.unbind("<ButtonPress>")
+        self._trigger_canvas.destroy()
+
+        future = self._give_stimulus()
+        future.add_done_callback(self._on_stimulus_complete)
+
+    def _on_background_touched(self, event: Event) -> None:
+        self._data.touch_events.append(
+            TouchEvent(time=datetime.now().timestamp(), x=event.x, y=event.y)
+        )
+        self._background.unbind("<ButtonPress>")
+        self._trigger_canvas.destroy()
+
+        self._on_incorrect()
+
+    def _on_correct(self) -> None:
+        self._give_reward()
+        self._data.result = Result.CORRECT
+        self._data.correct_rate = (self._animal_state.correct_trial + 1) / (
+            self._animal_state.current_level_trial_id + 1
+        )
+
+        self._on_inter_trial()
+
+    def _on_incorrect(self) -> None:
+        self._data.result = Result.INCORRECT
+        self._data.correct_rate = self._animal_state.correct_trial / (
+            self._animal_state.current_level_trial_id + 1
+        )
+
+        self._create_wrong_view()
+        self._on_inter_trial()
+
+    def _on_timeout(self) -> None:
+        self._data.result = Result.TIMEOUT
+        self._data.correct_rate = self._animal_state.correct_trial / (
+            self._animal_state.current_level_trial_id + 1
+        )
+
+        self._create_wrong_view()
+        self._on_inter_trial()
+
+    # endregion
+
+    # region sitimulus and reward
+    def _prepare_stimulus(self) -> "NDArray[int16]":
+        unit_duration = (
+            self._trial_config.stimulus_freq_duration
+            + self._trial_config.stimulus_freq_duration
+        )
+
+        times = ceil(self._trial_config.stimulus_duration / unit_duration)
+        times = max(times, 1)
+
+        freq_1 = ToneConfig(
+            frequency=self._trial_config.stimulus_freq,
+            duration=self._trial_config.stimulus_freq_duration,
+        )
+        freq_2 = ToneConfig(frequency=0, duration=self._trial_config.stimulus_interval)
+
+        return self._theater.aplayer.generate_stimulus([freq_1, freq_2], times)
+
+    def _give_stimulus(self) -> "Future[bool]":
+        return self._theater.aplayer.play_stimulus(self._tone)
+
+    def _on_stimulus_complete(self, future: "Future[bool]") -> None:
+        if future.result():
+            self._trigger_canvas.after(
+                self._trial_config.reward_delay,
+                lambda: self._on_correct(),
+            )
+
+    def _give_reward(self) -> None:
+        self._persistent_data.rewards += 1
+        self._theater.reward.give_reward(self._trial_config.reward_duration)
+
+    def _give_standard_stimulus(self) -> None:
+        self._standard_reward_stimulus.play(self._trial_config.reward_duration)
+
+    # endregion
+
+    # region data
+    def _init_data(self) -> None:
+        self._data = TrialData(
+            animal=self._animal_state.name,
+            trial_id=self._animal_state.trial_id,
+            current_level_trial_id=self._animal_state.current_level_trial_id,
+            trial_config=self._trial_config,
+            trial_start_time=datetime.now().timestamp(),
+            trial_end_time=0,
+            result=Result.TIMEOUT,
+            correct_rate=0,
+            touch_events=[],
+        )
+
+    # endregion

mxbiflow/timer/__init__.py

@@ -0,0 +1,3 @@
+from .frame_timer import FrameTimer
+
+__all__ = ["FrameTimer"]

mxbiflow/timer/frame_timer.py

@@ -0,0 +1,47 @@
+from dataclasses import dataclass
+from typing import Callable
+
+from pygame import time
+
+
+@dataclass
+class _Task:
+    interval_ms: int
+    callback: Callable[[], None]
+    repeat: bool
+    next_fire_ms: int
+    cancelled: bool = False
+
+
+class FrameTimer:
+    def __init__(self) -> None:
+        self._tasks: list[_Task] = []
+
+    def after(self, delay_ms: int, callback: Callable[[], None]) -> _Task:
+        now = time.get_ticks()
+        task = _Task(delay_ms, callback, False, now + delay_ms)
+        self._tasks.append(task)
+        return task
+
+    def every(self, interval_ms: int, callback: Callable[[], None]) -> _Task:
+        now = time.get_ticks()
+        task = _Task(interval_ms, callback, True, now + interval_ms)
+        self._tasks.append(task)
+        return task
+
+    def cancel(self, task: _Task) -> None:
+        task.cancelled = True
+
+    def update(self) -> None:
+        now = time.get_ticks()
+        for task in self._tasks:
+            if task.cancelled:
+                continue
+            if now >= task.next_fire_ms:
+                task.callback()
+                if task.repeat and not task.cancelled:
+                    task.next_fire_ms += task.interval_ms
+                else:
+                    task.cancelled = True
+
+        self._tasks = [task for task in self._tasks if not task.cancelled]

mxbiflow/tmp_email.py

@@ -0,0 +1,13 @@
+from mxbi.utils.logger import logger
+from pymotego.send_email import EmailAttachment, EmailClient
+
+
+def send_email(subject: str, body: str, attachments: list[EmailAttachment]) -> None:
+    with EmailClient() as client:
+        response = client.send(
+            subject=subject,
+            html_body=body,
+            attachments=attachments,
+        )
+    logger.info(f"Status: {response.status_code}")
+    logger.info(f"Body: {response.text}")

mxbiflow/ui/components/animal.py

@@ -0,0 +1,87 @@
+from PySide6.QtCore import Qt, Signal
+from PySide6.QtWidgets import QComboBox, QFormLayout, QLabel, QLineEdit, QMenu
+
+from ...config_store import ConfigStore
+from ...models.animal import AnimalConfig
+from ...path import STAGE_PATH
+from ...scene import Scenes
+from .card import CardFrame
+
+
+class AnimalCard(CardFrame):
+    remove_requested = Signal()
+
+    def __init__(self, parent, animals: dict[str, str]):
+        super().__init__(parent=parent, object_name="card")
+        self.setContextMenuPolicy(Qt.ContextMenuPolicy.CustomContextMenu)
+        self.customContextMenuRequested.connect(self._on_context_menu)
+
+        items = list(animals.items())
+        self._animal_ids = [animal_id for animal_id, _name in items]
+        self._animal_names = [name for _animal_id, name in items]
+
+        layout = QFormLayout(self)
+        layout.setContentsMargins(8, 8, 8, 8)
+
+        label_animal_name = QLabel("animal", self)
+        self.combo_animal_name = QComboBox(self)
+        self.combo_animal_name.addItems(self._animal_names)
+        self.combo_animal_name.currentIndexChanged.connect(self._on_animal_changed)
+        layout.addRow(label_animal_name, self.combo_animal_name)
+
+        lable_animal_id = QLabel("id", self)
+        self.line_animal_id = QLineEdit(self)
+        self.line_animal_id.setReadOnly(True)
+        self.line_animal_id.setFocusPolicy(Qt.FocusPolicy.NoFocus)
+        layout.addRow(lable_animal_id, self.line_animal_id)
+
+        label_stage = QLabel("stage", self)
+        items = ConfigStore(STAGE_PATH, Scenes).value
+        self.combo_stage = QComboBox(self)
+        self.combo_stage.addItems([i for i in items.root])
+        self.combo_stage.setCurrentText("idle")
+        layout.addRow(label_stage, self.combo_stage)
+
+        label_level = QLabel("level", self)
+        self.combo_level = QComboBox(self)
+        self.combo_level.addItems(["0", "1"])
+        self.combo_level.setCurrentText("0")
+        layout.addRow(label_level, self.combo_level)
+
+        self._sync_animal_id()
+
+    def _on_animal_changed(self, _index: int) -> None:
+        self._sync_animal_id()
+
+    def _sync_animal_id(self) -> None:
+        index = self.combo_animal_name.currentIndex()
+        if index < 0 or index >= len(self._animal_ids):
+            self.line_animal_id.setText("")
+            return
+        self.line_animal_id.setText(self._animal_ids[index])
+
+    def animal_id(self) -> str:
+        return self.line_animal_id.text()
+
+    def animal_name(self) -> str:
+        return self.combo_animal_name.currentText()
+
+    def _on_context_menu(self, pos):
+        menu = QMenu(self)
+        action = menu.addAction("Remove")
+        action.triggered.connect(lambda _checked=False: self.remove_requested.emit())
+        menu.exec(self.mapToGlobal(pos))
+
+    def load_config(self, config: AnimalConfig) -> None:
+        self.combo_animal_name.setCurrentText(config.name)
+        self.combo_stage.setCurrentText(config.stage)
+        self.combo_level.setCurrentText(str(config.level))
+
+    @property
+    def result(self) -> AnimalConfig:
+        return AnimalConfig(
+            rfid_id=self.animal_id(),
+            name=self.animal_name(),
+            stage=self.combo_stage.currentText(),
+            level=int(self.combo_level.currentText()),
+        )

mxbiflow/ui/components/baseconfig.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+from pymxbi import MXBIModel
+from pymxbi.platform import PlatformEnum
+from pymxbi.screen import get_screen_size
+from PySide6.QtCore import Signal
+from PySide6.QtWidgets import QComboBox, QFormLayout, QGroupBox, QLabel
+
+
+class BaseConfig(QGroupBox):
+    changed = Signal(str)
+
+    def __init__(self, parent, mxbi_options: list[str]):
+        super().__init__(parent)
+
+        self.setTitle("Base config")
+
+        self._layout = QFormLayout()
+        self.setLayout(self._layout)
+
+        self._label_mxbi = QLabel("mxbi id:")
+        self._combo_mxbi = QComboBox()
+        self._combo_mxbi.addItems(mxbi_options)
+        self._layout.addRow(self._label_mxbi, self._combo_mxbi)
+
+        self._label_platform = QLabel("platform:")
+        self._combo_platform = QComboBox()
+        self._combo_platform.addItems([platform.value for platform in PlatformEnum])
+        self._layout.addRow(self._label_platform, self._combo_platform)
+
+        self._label_screen = QLabel("screen:")
+        self._combo_screen = QComboBox()
+
+        for screen in get_screen_size():
+            self._combo_screen.addItem(
+                f"{screen.width} * {screen.height}", (screen.width, screen.height)
+            )
+        self._layout.addRow(self._label_screen, self._combo_screen)
+
+        self._bind_events()
+
+    def _emit_changed(self, msg: str) -> None:
+        self.changed.emit(msg)
+
+    def _bind_events(self) -> None:
+        self._combo_mxbi.currentTextChanged.connect(self._emit_changed)
+        self._combo_platform.currentTextChanged.connect(self._emit_changed)
+
+    @property
+    def mxbi_id(self) -> str:
+        return self._combo_mxbi.currentText()
+
+    @property
+    def platform(self) -> str:
+        return self._combo_platform.currentText()
+
+    def load_from_model(self, model: MXBIModel) -> None:
+        self._combo_mxbi.setCurrentText(str(model.mxbi_id))
+        self._combo_platform.setCurrentText(str(model.platform))
+
+    def apply_to_model(self, model: MXBIModel) -> None:
+        mxbi_id_text = self._combo_mxbi.currentText().strip()
+        try:
+            model.mxbi_id = int(mxbi_id_text) if mxbi_id_text else 0
+        except ValueError:
+            model.mxbi_id = 0
+
+        model.platform = PlatformEnum(self._combo_platform.currentText())

mxbiflow/ui/components/card.py

@@ -0,0 +1,18 @@
+from PySide6.QtWidgets import QFrame
+
+
+class CardFrame(QFrame):
+    def __init__(self, parent=None, object_name: str = "card"):
+        super().__init__(parent=parent)
+
+        self.setObjectName(object_name)
+        self.setFrameShape(QFrame.Shape.Box)
+        self.setLineWidth(1)
+        self.setStyleSheet(
+            f"""
+            QFrame#{object_name} {{
+                border: 1px solid #C8C8C8;
+                border-radius: 6px;
+            }}
+            """
+        )

mxbiflow/ui/components/device_card/__init__.py

@@ -0,0 +1,17 @@
+from .detector.beambreak_detector_card import BeambreakDetectorCard
+from .detector.fusion_detector import FusionDetectorCard
+from .detector.mock_detector_card import MockDetectorCard
+from .detector.rfid_detector import RFIDDetectorCard
+from .device_card import DeviceCard
+from .rewarder.mock_rewarder_card import MockRewarderCard
+from .rewarder.rpi_gpio_rewarder import RPIGpioPumpCard
+
+__all__ = [
+    "DeviceCard",
+    "MockRewarderCard",
+    "RPIGpioPumpCard",
+    "MockDetectorCard",
+    "RFIDDetectorCard",
+    "BeambreakDetectorCard",
+    "FusionDetectorCard",
+]

mxbiflow/ui/components/device_card/detector/beambreak_detector_card.py

@@ -0,0 +1,29 @@
+from pymxbi.detector import FusionContinuousDetectorModel
+from PySide6.QtWidgets import QLabel, QLineEdit
+
+from ..device_card import DeviceCard
+
+
+class BeambreakDetectorCard(DeviceCard[FusionContinuousDetectorModel]):
+    def __init__(self):
+        super().__init__()
+
+        self.set_title("Beambreak Detector")
+
+        label_pin = QLabel("Pin")
+        self._line_pin = QLineEdit()
+        self._line_pin.setPlaceholderText("Pin")
+        self.layout_config.addRow(label_pin, self._line_pin)
+
+    def load_config(self, model: FusionContinuousDetectorModel) -> None:
+        self.checkbox_enabled.setChecked(model.enabled)
+        self.line_device_id.setText(str(model.id))
+        self._line_pin.setText(str(model.pin))
+
+    @property
+    def result(self) -> FusionContinuousDetectorModel:
+        return FusionContinuousDetectorModel(
+            enabled=self.checkbox_enabled.isChecked(),
+            id=int(self.line_device_id.text()),
+            pin=int(self._line_pin.text()),
+        )