mxbiflow 0.1.1__py3-none-any.whl

mxbiflow/__init__.py

@@ -0,0 +1,3 @@
+from .mxbiflow import MXBIFlow, get_mxbiflow
+
+__all__ = ["MXBIFlow", "get_mxbiflow"]

mxbiflow/assets/__init__.py

@@ -0,0 +1,5 @@
+from pathlib import Path
+
+ROOT = Path(__file__).parent
+
+ASSET_CLICKER_PATH = ROOT / "clicker.wav"

mxbiflow/config_store.py

@@ -0,0 +1,68 @@
+import os
+from pathlib import Path
+from typing import Generic, TypeVar
+
+from pydantic import BaseModel
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class ConfigStore(Generic[T]):
+    def __init__(
+        self,
+        config_path: Path,
+        config_class: type[T],
+        *,
+        create_default: bool = True,
+    ) -> None:
+        self._config_path = config_path
+        self._config_class = config_class
+        self._create_default = create_default
+        self._config = self._load_config()
+
+    @property
+    def value(self) -> T:
+        return self._config
+
+    def _ensure_config_readable(self) -> None:
+        if not self._config_path.exists():
+            raise FileNotFoundError(f"Config file {self._config_path} not found")
+
+        if not os.access(self._config_path, os.R_OK):
+            raise PermissionError(f"Config file {self._config_path} is not readable")
+
+    def _create_default_config(self) -> T:
+        config = self._config_class()
+        self._config_path.parent.mkdir(parents=True, exist_ok=True)
+        self._config_path.write_text(
+            config.model_dump_json(indent=4),
+            encoding="utf-8",
+        )
+        return config
+
+    def _load_config(self) -> T:
+        try:
+            self._ensure_config_readable()
+            text = self._config_path.read_text(encoding="utf-8")
+            return self._config_class.model_validate_json(text)
+        except Exception as e:
+            if self._create_default:
+                return self._create_default_config()
+            raise RuntimeError(
+                f"Failed to load config file {self._config_path}: {e}"
+            ) from e
+
+    def save(self, data: T | None = None) -> None:
+        if data is not None:
+            self._config = data
+
+        self._config_path.parent.mkdir(parents=True, exist_ok=True)
+        try:
+            self._config_path.write_text(
+                self._config.model_dump_json(indent=4),
+                encoding="utf-8",
+            )
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to save config file {self._config_path}: {e}"
+            ) from e

mxbiflow/data_logger.py

@@ -0,0 +1,114 @@
+import csv
+import json
+import sys
+from datetime import datetime
+from enum import StrEnum
+from pathlib import Path
+
+from .utils.logger import logger
+
+now = datetime.now()
+
+
+class DataLoggerType(StrEnum):
+    JSONL = "jsonl"
+    JSON = "json"
+    CSV = "csv"
+
+
+class DataLogger:
+    def __init__(
+        self,
+        path: Path,
+        session_id: int,
+        monkey: str,
+        filename: str,
+        type: DataLoggerType = DataLoggerType.JSONL,
+    ) -> None:
+        self._path = path
+        self._session_id = session_id
+        self._monkey = monkey
+        self._filename = filename
+        self._type = type
+
+        self._data_dir = self._ensure_data_dir()
+        self._data_path = self._get_path(f".{self._type.value}")
+
+    @property
+    def path(self) -> Path:
+        return self._data_path
+
+    def _ensure_data_dir(self) -> Path:
+        date_path = Path(f"{now.year}{now.month:02d}{now.day:02d}")
+        session_path = Path(f"{self._session_id}")
+        monkey_path = Path(f"{self._monkey}")
+
+        base_dir = self._path / date_path / session_path / monkey_path
+
+        try:
+            base_dir.mkdir(parents=True, exist_ok=True)
+            return base_dir
+        except Exception as e:
+            logger.error(f"failed to create {base_dir}: {e}")
+            sys.exit(1)
+
+    def _get_path(self, suffix: str) -> Path:
+        return self._data_dir / f"{self._filename}{suffix}"
+
+    def save(self, data: dict) -> None:
+        match self._type:
+            case DataLoggerType.JSONL:
+                self._save_jsonl(data)
+            case DataLoggerType.JSON:
+                self._save_json(data)
+            case DataLoggerType.CSV:
+                self.save_csv_row(data)
+
+    def _save_jsonl(self, data: dict) -> None:
+        try:
+            json_line = json.dumps(data, ensure_ascii=False)
+
+            with open(self._data_path, "a", encoding="utf-8") as f:
+                f.write(json_line + "\n")
+
+        except TypeError as e:
+            logger.error(f"Data is not JSON serializable: {e}")
+            raise
+        except IOError as e:
+            logger.error(f"Failed to write to file {self._data_path}: {e}")
+            raise
+        except Exception as e:
+            logger.error(f"Unexpected error while writing data: {e}")
+            raise
+
+    def _save_json(self, data: dict) -> None:
+        try:
+            self._data_path.parent.mkdir(parents=True, exist_ok=True)
+            with open(self._data_path, "w", encoding="utf-8") as f:
+                json.dump(data, f, ensure_ascii=False, indent=2)
+        except TypeError as e:
+            logger.error(f"Data is not JSON serializable: {e}")
+            raise
+        except IOError as e:
+            logger.error(f"Failed to write to file {self._data_path}: {e}")
+            raise
+        except Exception as e:
+            logger.error(f"Unexpected error while writing JSON data: {e}")
+            raise
+
+    def save_csv_row(self, data: dict) -> None:
+        csv_path = self._get_path(".csv")
+        try:
+            csv_path.parent.mkdir(parents=True, exist_ok=True)
+            file_exists = csv_path.exists() and csv_path.stat().st_size > 0
+
+            fieldnames = sorted(data.keys())
+
+            with csv_path.open("a", newline="", encoding="utf-8") as f:
+                writer = csv.DictWriter(f, fieldnames=fieldnames)
+                if not file_exists:
+                    writer.writeheader()
+                writer.writerow({k: data.get(k, "") for k in fieldnames})
+        except Exception as e:
+            logger.error(f"Failed to write CSV row to {csv_path}: {e}")
+            raise

mxbiflow/default/__init__.py

@@ -0,0 +1,4 @@
+from .habituation.habituarion import Habituarion
+from .idle.idle import IDLE
+
+__all__ = ["Habituarion", "IDLE"]

mxbiflow/default/idle/idle.py

@@ -0,0 +1,57 @@
+from dataclasses import dataclass
+from pathlib import Path
+from random import choice
+
+from pygame import Event, Rect, Surface, image, transform
+
+from mxbiflow import get_mxbiflow
+from mxbiflow.scene.scene_protocol import SceneProtocol
+
+ASSETS_PATH = Path(__file__).parent / "assets"
+
+
+@dataclass
+class Asset:
+    image: Surface
+    rect: Rect
+
+
+class IDLE:
+    _running: bool
+
+    def __init__(self) -> None:
+        self._mxbiflow = get_mxbiflow()
+
+        self._screen_size = self._mxbiflow.mxbi.screen_size
+        self._pos = ((self._screen_size.width // 4) * 3, self._screen_size.height // 2)
+        self._vstimulus_size = self._screen_size.width // 2 * 0.75
+
+        self._assets = [
+            Asset(
+                asset_image := transform.scale(
+                    image.load(path).convert_alpha(),
+                    (self._vstimulus_size, self._vstimulus_size),
+                ),
+                asset_image.get_rect(center=self._pos),
+            )
+            for path in ASSETS_PATH.glob("*.png")
+        ]
+
+        self._asset = choice(self._assets)
+
+    def start(self) -> None:
+        self._running = True
+
+    def quit(self) -> None:
+        self._running = False
+
+    @property
+    def running(self) -> bool:
+        return self._running
+
+    def handle_event(self, event: Event) -> None: ...
+
+    def update(self, dt_s: float) -> None: ...
+
+    def draw(self, screen: Surface) -> None:
+        screen.blit(self._asset.image, self._asset.rect)

mxbiflow/detector_bridge.py

@@ -0,0 +1,87 @@
+from dataclasses import dataclass
+from queue import Empty, SimpleQueue
+
+import pygame
+from pygame import Event, event
+from pymxbi.detector import MockDetector
+from pymxbi.detector.detector import DetectionResult, Detector, DetectorEvent
+
+EVT_DETECTOR = pygame.USEREVENT + 1
+
+
+@dataclass(frozen=True)
+class DetectorMsg:
+    kind: DetectorEvent
+    animal: str | None
+
+
+class DetectorBridge:
+    def __init__(self, detector: Detector) -> None:
+        self._detector = detector
+        self._q = SimpleQueue()
+        self._started = False
+
+    def start(self) -> None:
+        if self._started:
+            return
+        self._started = True
+
+        self._detector.begin()
+        self._detector.register_event(DetectorEvent.ANIMAL_ENTERED, self._emit_entered)
+        self._detector.register_event(DetectorEvent.ANIMAL_LEFT, self._emit_left)
+        self._detector.register_event(DetectorEvent.FAULT_DETECTED, self._emit_fault)
+
+    def _emit(self, kind: DetectorEvent, animal: str | None) -> None:
+        self._q.put(DetectorMsg(kind=kind, animal=animal))
+
+    def _emit_entered(self, detection_result: DetectionResult) -> None:
+        self._emit(
+            DetectorEvent.ANIMAL_ENTERED,
+            detection_result.animal_id or detection_result.animal_name,
+        )
+
+    def _emit_left(self, detection_result: DetectionResult) -> None:
+        self._emit(
+            DetectorEvent.ANIMAL_LEFT,
+            detection_result.animal_id or detection_result.animal_name,
+        )
+
+    def _emit_fault(self, detection_result: DetectionResult) -> None:
+        self._emit(
+            DetectorEvent.FAULT_DETECTED,
+            detection_result.animal_id or detection_result.animal_name,
+        )
+
+    def emit_pygame_event(self) -> None:
+        while True:
+            try:
+                msg = self._q.get_nowait()
+            except Empty:
+                break
+
+            event.post(event.Event(EVT_DETECTOR, msg=msg))
+
+    def manaul_emit(self, animal_idx: int | None = None) -> None:
+        if isinstance(self._detector, MockDetector):
+            if animal_idx is None:
+                self._detector.animal_left()
+            else:
+                self._detector.animal_present(animal_idx)
+
+    def handle_event(self, event: Event) -> None:
+        if event.type == pygame.KEYDOWN:
+            match event.key:
+                case pygame.K_0:
+                    self.manaul_emit(0)
+                case pygame.K_1:
+                    self.manaul_emit(1)
+                case pygame.K_2:
+                    self.manaul_emit(2)
+                case pygame.K_3:
+                    self.manaul_emit(3)
+                case pygame.K_4:
+                    self.manaul_emit(4)
+                case pygame.K_5:
+                    self.manaul_emit(5)
+                case pygame.K_l:
+                    self.manaul_emit()

mxbiflow/game.py

@@ -0,0 +1,84 @@
+import pygame
+from pygame import Event
+from pymxbi import MXBI
+
+from .detector_bridge import DetectorBridge
+from .models.session import Session
+from .mxbiflow import MXBIFlow, set_mxbiflow
+from .scene import SceneManager
+from .scheduler import Scheduler
+
+
+class Game:
+    def __init__(
+        self,
+        session: Session,
+        scene_manager: SceneManager,
+        detector_bridge: DetectorBridge,
+        mxbi: MXBI,
+    ) -> None:
+        pygame.init()
+
+        self._scene_manager = scene_manager
+        self._session = session
+        self._mxbi = mxbi
+        self._aplayer = self._mxbi.aplayer
+
+        self._detector_binder = detector_bridge
+        self._detector_binder.start()
+
+        self._scheduler = Scheduler(self._session, self._scene_manager)
+
+        self._mxbiflow = MXBIFlow(self._session, self._mxbi)
+        set_mxbiflow(self._mxbiflow)
+
+        self._screen = pygame.display.set_mode(
+            (self._mxbi.screen_size.width, self._mxbi.screen_size.height)
+        )
+        self._clock = pygame.time.Clock()
+        self._running = True
+
+    def play(self) -> None:
+        while self._running:
+            dt = self._clock.tick(60) / 1000.0
+
+            self._detector_binder.emit_pygame_event()
+
+            for event in pygame.event.get():
+                self._handle_event(event)
+                self._scheduler.handle_event(event)
+                self._scene_manager.handle_event(event)
+                self._detector_binder.handle_event(event)
+
+            self._scene_manager.update(dt)
+            self._aplayer.update()
+            self._scheduler.update()
+            self._mxbiflow.update()
+
+            self._screen.fill((0, 0, 0))
+
+            self._scene_manager.draw(self._screen)
+            self._scene_manager.apply_pending()
+
+            pygame.display.flip()
+
+        self.quit()
+
+    def _handle_event(self, event: Event) -> None:
+        match event.type:
+            case pygame.QUIT:
+                self._running = False
+            case pygame.KEYDOWN:
+                self._handle_keyboard_event(event)
+
+    def _handle_keyboard_event(self, event: Event) -> None:
+        match event.key:
+            case pygame.K_ESCAPE:
+                self._running = False
+            case pygame.K_q:
+                self._running = False
+
+    def quit(self) -> None:
+        if self._scene_manager.current is not None:
+            self._scene_manager.current.quit()
+        pygame.quit()

mxbiflow/infra/eventbus.py

@@ -0,0 +1,31 @@
+from typing import Callable
+
+
+class EventBus:
+    def __init__(self) -> None:
+        self._events_dict: dict[str, list[Callable]] = {}
+
+    def subscribe(self, event: str, handler: Callable) -> None:
+        if event in self._events_dict:
+            return
+
+        self._events_dict[event] = [handler]
+
+    def unsubscribe(self, event: str, handler: Callable) -> None:
+        if event not in self._events_dict:
+            return
+
+        self._events_dict[event].remove(handler)
+
+    def publish(self, event: str) -> None:
+        if event not in self._events_dict:
+            return
+
+        for handler in self._events_dict[event]:
+            handler()
+
+    def clear(self) -> None:
+        self._events_dict.clear()
+
+
+event_bus = EventBus()

mxbiflow/main.py

@@ -0,0 +1,106 @@
+from pymxbi import MXBI
+
+from .game import Game
+from .models.session import Session
+
+
+def main():
+    run_config()
+
+    mxbi = build_mxbi()
+
+    session = init_session()
+
+    game = init_mxbiflow(mxbi, session)
+    game.play()
+
+
+def run_config():
+    import sys
+
+    from PySide6.QtWidgets import QApplication
+
+    from .ui.experiment_panel import ExperimentPanel
+    from .ui.mxbi_panel import MXBIPanel
+
+    app = QApplication(sys.argv)
+
+    mxbi_panel = MXBIPanel()
+    experiment_panel = ExperimentPanel()
+    mxbi_panel.accepted.connect(experiment_panel.show)
+    experiment_panel.accepted.connect(app.quit)
+
+    mxbi_panel.show()
+
+    app.exec()
+
+
+def build_mxbi() -> MXBI:
+    from loguru import logger
+    from pymxbi import MXBIModel, build_mxbi
+
+    from .config_store import ConfigStore
+    from .models.session import SessionConfig
+    from .path import MXBI_CONFIG_PATH, SESSION_CONFIG_PATH
+
+    mxbi_config = ConfigStore(MXBI_CONFIG_PATH, MXBIModel).value
+    session_config = ConfigStore(SESSION_CONFIG_PATH, SessionConfig).value
+
+    mxbi = build_mxbi(mxbi_config, logger)
+    mxbi.register_animal(
+        {animal.rfid_id: animal.name for animal in session_config.animals}
+    )
+
+    return mxbi
+
+
+def init_session() -> Session:
+    from .config_store import ConfigStore
+    from .models.animal import Animal, StageState
+    from .models.session import DailySessionIdStore, Session, SessionConfig
+    from .path import SESSION_CONFIG_PATH, SESSION_COUNTER_PATH
+
+    session_config = ConfigStore(SESSION_CONFIG_PATH, SessionConfig).value
+    store = DailySessionIdStore(SESSION_COUNTER_PATH)
+
+    animal_dict: dict[str, Animal] = {}
+    for animal_config in session_config.animals:
+        train_state = StageState(
+            stage_name=animal_config.stage, level=animal_config.level
+        )
+        animal_state = Animal(
+            rfid_id=animal_config.rfid_id,
+            name=animal_config.name,
+        )
+        animal_state.set_current_stage(train_state)
+        animal_dict[animal_config.name] = animal_state
+
+    session = Session(
+        session_id=store.session_id,
+        experimenter=session_config.experimenter,
+        reward_type=session_config.reward_type,
+        send_email=False,
+        sync_data=False,
+        note=session_config.note,
+        animals=animal_dict,
+    )
+    session.start()
+    print(session.session_id)
+
+    return session
+
+
+def init_mxbiflow(mxbi, session) -> Game:
+    from .default import IDLE, Habituarion
+    from .detector_bridge import DetectorBridge
+    from .GNGSiD import SizeReduction
+    from .path import STAGE_PATH
+    from .scene import SceneManager
+
+    scene_manager = SceneManager()
+    scene_manager.register(Habituarion)
+    scene_manager.register(IDLE)
+    scene_manager.register(SizeReduction)
+    scene_manager.persist(STAGE_PATH)
+    detector_bridge = DetectorBridge(mxbi.detector)
+    return Game(session, scene_manager, detector_bridge, mxbi)

mxbiflow/models/animal.py

@@ -0,0 +1,130 @@
+from time import time
+
+from pydantic import BaseModel, Field, PrivateAttr
+
+
+class AnimalConfig(BaseModel):
+    rfid_id: str = Field(default="", frozen=True)
+    name: str = Field(default="mock", frozen=True)
+    stage: str = Field(default="idle", frozen=True)
+    level: int = Field(default=0, ge=0, frozen=True)
+
+
+class AnimalBaseInfo(BaseModel):
+    animal: str
+    trial_id: int
+    level: int
+    level_trial_id: int
+    animal_session_id: int
+    animal_session_trial_id: int
+
+
+class StageState(BaseModel):
+    stage_name: str
+    stage_trial_id: int = Field(default=0, ge=0)
+
+    level: int = Field(default=0, ge=0)
+    level_trial_id: int = Field(default=0, ge=0)
+
+    def level_up(self):
+        self.level_trial_id = 0
+        self.level += 1
+
+    def level_down(self):
+        self.level_trial_id = 0
+        self.level -= 1
+
+
+class AnimalSessionState(BaseModel):
+    session_id: int = Field(ge=0)
+    start_at: float = Field(default_factory=lambda: time())
+    end_at: float | None = None
+    trial_id: int = Field(default=0, ge=0)
+
+
+class Animal(BaseModel):
+    rfid_id: str = Field(frozen=True)
+    name: str = Field(frozen=True)
+
+    trial_id: int = Field(default=0, ge=0)
+
+    _current_stage: str = PrivateAttr(default="idle")
+    _stages: dict[str, StageState] = PrivateAttr(default_factory=dict)
+    _current_animal_session: AnimalSessionState | None = PrivateAttr(default=None)
+    _sessions: list[AnimalSessionState] = PrivateAttr(default_factory=list)
+
+    def add_trial(self) -> None:
+        self.trial_id += 1
+        self.current_stage.stage_trial_id += 1
+        self.current_stage.level_trial_id += 1
+        assert self.current_animal_session is not None
+        self.current_animal_session.trial_id += 1
+
+    @property
+    def current_animal_session(self) -> AnimalSessionState | None:
+        return self._current_animal_session
+
+    def start_animal_session(self):
+        session_id = 1
+
+        if self._current_animal_session is not None:
+            raise ValueError("Animal session is already started")
+
+        if self._sessions:
+            session_id = len(self._sessions) + 1
+
+        session = AnimalSessionState(session_id=session_id)
+        self._current_animal_session = session
+        self._sessions.append(session)
+
+    def end_animal_session(self):
+        if self._current_animal_session is None:
+            raise ValueError("Animal session is not started")
+
+        self._current_animal_session.end_at = time()
+        self._current_animal_session = None
+
+    @property
+    def current_stage(self) -> StageState:
+        key = self._current_stage
+
+        state = self._stages.get(key)
+        if state is None:
+            raise ValueError(f"Unknown stage: {key}")
+
+        return state
+
+    def set_current_stage(self, stage: StageState | str) -> None:
+        if isinstance(stage, str):
+            stage = StageState(stage_name=stage)
+
+        self._current_stage = stage.stage_name
+        if stage.stage_name in self._stages:
+            return
+
+        self._stages[stage.stage_name] = stage
+
+    def clear_current_stage(self) -> None:
+        self._current_stage = "idle"
+
+    def level_up(self) -> int:
+        self.current_stage.level_up()
+        return self.current_stage.level
+
+    def level_down(self) -> int:
+        self.current_stage.level_down()
+        return self.current_stage.level
+
+    @property
+    def base_info(self) -> AnimalBaseInfo:
+        if self.current_animal_session is None:
+            raise ValueError("Animal session is not started")
+
+        return AnimalBaseInfo(
+            animal=self.name,
+            trial_id=self.trial_id,
+            level=self.current_stage.level,
+            level_trial_id=self.current_stage.level_trial_id,
+            animal_session_id=self.current_animal_session.session_id,
+            animal_session_trial_id=self.current_animal_session.trial_id,
+        )