app_logic

  1"""
  2'app_logic' is the primary module in the Photologic-Experiment-Rig codebase. It contains the
  3StateMachine class which holds the logic for each 'state' an experiment can be in.
  4
  5Before performing any actions, the StateMachine class initializes the instances of `models.experiment_process_data`,
  6`controllers.arduino_control`, and `views.main_gui` classes.
  7Launching main_gui initializes a tkinter root and allows for a GUI to be created for the program.
  8
  9This module is launched from main to make restarting the program easier, which is done by destroying the
 10instance of state machine and launcing a new one.
 11"""
 12
 13# external imports
 14import time
 15import threading
 16import logging
 17import toml
 18import queue
 19from typing import Callable
 20import datetime
 21from logging import FileHandler
 22
 23# imports for locally used modules and classes
 24from models.experiment_process_data import ExperimentProcessData
 25from views.gui_common import GUIUtils
 26import system_config
 27from views.main_gui import MainGUI
 28
 29# these are just use for type hinting here
 30from controllers.arduino_control import ArduinoManager
 31
 32logger = logging.getLogger()
 33"""Logger used to log program runtime details for debugging"""
 34logger.setLevel(logging.INFO)
 35
 36console_handler = logging.StreamHandler()
 37"""The console handler is used to print errors to the console"""
 38console_handler.setLevel(logging.ERROR)
 39console_handler.setFormatter(
 40    logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
 41)
 42logger.addHandler(console_handler)
 43
 44RIG_CONFIG = system_config.get_rig_config()
 45"""
 46This utilizes `system_config` module to obtain a constant path for this machine regardless of OS to the Rig Files foler 
 47at the current users documents folder.
 48"""
 49
 50with open(RIG_CONFIG, "r") as f:
 51    DOOR_CONFIG = toml.load(f)["door_motor_config"]
 52
 53DOOR_MOVE_TIME = DOOR_CONFIG["DOOR_MOVE_TIME"]
 54"""Constant value stored in the door_motor_config section of `RIG CONFIG` config"""
 55
 56
 57now = datetime.datetime.now()
 58# configure logfile details such as name and the level at which information should be added to the log (INFO here)
 59logfile_name = f"{now.hour}_{now.minute}_{now.second} - {now.date()} experiment log"
 60logfile_path = system_config.get_log_path(logfile_name)
 61
 62file_handler = FileHandler(logfile_path)
 63file_handler.setLevel(logging.INFO)
 64file_handler.setFormatter(
 65    logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
 66)
 67logger.addHandler(file_handler)
 68
 69
 70class StateMachine:
 71    """
 72    This class is the heart of the program. Defines and handles program state transitions.
 73    It coordinates co-operation between view (gui) and models (data). We inherit a TkinterApp class to include
 74    attributes of a tkinter app such as gui, arduino controller, etc.
 75
 76    Attributes
 77    ----------
 78    - **exp_data** (*ExperimentProcessData*): An instance of `models.experiment_process_data` ExperimentProcessData, this attribute allows for access and
 79    modification of experiment variables, and access for to more specific models like `models.event_data` and `models.arduino_data`.
 80    - **arduino_controller** (*ArduinoManager*): An instance of `controllers.arduino_control` ArduinoManager, this allows for communication between this program and the
 81    Arduino board.
 82    - **main_gui** (*MainGUI*): An instance of `views.main_gui` MainGUI, this allows for the creation and modification of all GUI attributes in the program. All GUI windows
 83    are created and managed here.
 84    - **state** (*str*): Contains the current state the program is in.
 85    - **prev_state** (*str*): Contains the state the program was previously in. Useful to restore state in case of erroneous transitions.
 86    - **app_result** (*list*): Mutable list with one element. Is a reference to list defined in `main`.
 87    - **transitions**  (*dict*): Program state transition table.
 88
 89    Methods
 90    -------
 91    - `trigger`(event)
 92        Handles state transition events. Decides if a transition is valid and warns user of destructive transitions. If valid,
 93        passes state to `execute_state` method.
 94    - `execute_state`(new_state: str)
 95        Takes the state passed from trigger event and decides appropriate action.
 96    - `process_queue`(data_queue)
 97        Processes incoming data from the Arduino board. Reads from queue that is added to by `controllers.arduino_control` module
 98        `listen_for_serial` method.
 99    - `reject_actions`(event)
100        A static method that handles the rejection of actions that cannot be performed given a certain state transition.
101    """
102
103    def __init__(self, result_container):
104        """init method for `StateMachine`. Takes `main` module `result_container`.
105        the 'super' parent class which is the gui"""
106
107        self.exp_data = ExperimentProcessData()
108
109        self.arduino_controller = ArduinoManager(self.exp_data)
110
111        self.main_gui = MainGUI(self.exp_data, self.trigger, self.arduino_controller)
112        logging.info("GUI started successfully.")
113
114        self.state = "IDLE"
115        """Default state for program is set at IDLE"""
116
117        self.prev_state = None
118        """Default previous state for program is set to None. Utilized in `trigger`."""
119
120        self.app_result = result_container
121        """
122        Here we store a reference to `main` module `result_container` to store decision of whether to restart the 
123        program or just terminate the current instance.
124        """
125
126        self.transitions = {
127            ("IDLE", "GENERATE SCHEDULE"): "GENERATE SCHEDULE",
128            (
129                "GENERATE SCHEDULE",
130                "IDLE",
131            ): "IDLE",
132            ("IDLE", "START"): "START PROGRAM",
133            ("IDLE", "RESET"): "RESET PROGRAM",
134            ("STOP PROGRAM", "RESET"): "RESET PROGRAM",
135            ("START PROGRAM", "ITI"): "ITI",
136            ("ITI", "DOOR OPEN"): "OPENING DOOR",
137            ("OPENING DOOR", "TTC"): "TTC",
138            ("TTC", "SAMPLE"): "SAMPLE",  # -> if rat engages in trial
139            ("TTC", "TRIAL END"): "TRIAL END",  # -> if rat does NOT engage
140            ("SAMPLE", "TRIAL END"): "TRIAL END",
141            ("TRIAL END", "ITI"): "ITI",  # -> save trial data, update gui, door up
142            ("ITI", "STOP"): "STOP PROGRAM",  # can move to 'stop' state in all states
143            ("OPENING DOOR", "STOP"): "STOP PROGRAM",
144            ("TTC", "STOP"): "STOP PROGRAM",
145            ("SAMPLE", "STOP"): "STOP PROGRAM",
146            ("TRIAL END", "STOP"): "STOP PROGRAM",
147        }
148        """State transition table defines all transitions that the program can possibly take"""
149
150        # don't start the mainloop until AFTER the gui is setup so that the app_result property
151        # is available if reset is desired
152        self.main_gui.mainloop()
153
154    def trigger(self, event):
155        """
156        This function takes a requested state and decides if the transition from this state is allowed.
157
158        Parameters
159        ----------
160        - **event** (*str*): Contains the desired state to move to.
161        """
162
163        transition = (self.state, event)
164        self.prev_state = self.state
165        new_state = None
166
167        logger.info(f"state transition -> {transition}")
168
169        # checking if the attemped transition is valid according to the table
170        if transition in self.transitions:
171            new_state = self.transitions[transition]
172        else:
173            # if the key is not in the transition table, handle rejection based on desired state
174            self.reject_actions(event)
175
176        # if the key was in the transition table, and the new state wants to reset the program, that means this is a valid action at this time. MAKE SURE the
177        # user REALLY wants to do that
178        if new_state == "RESET PROGRAM":
179            response = GUIUtils.askyesno(
180                "============WARNING============",
181                "THIS ACTION WILL ERASE ALL DATA CURRENTLY STORED FOR THIS EXPERIMENT... ARE YOU SURE YOU WANT TO CONTINUE?",
182            )
183            # if true, go ahead with the reset, if false return to prev_state
184            if not response:
185                new_state = self.prev_state
186
187        # if attempt to start program with no experiment schedule, tell user to do that
188        if new_state == "START PROGRAM":
189            if self.exp_data.program_schedule_df.empty:
190                GUIUtils.display_error(
191                    "Experiement Schedule Not Yet Generated!!!",
192                    "Please generate the program schedule before attempting to start the experiment. You can do this by clicking the 'Valve / Stimuli' button in the main screen.",
193                )
194                new_state = self.prev_state
195
196        # if we have a new state, perform the associated action for that state
197        if new_state:
198            logger.info(f"new state -> {new_state}")
199            self.execute_state(new_state)
200
201    def execute_state(self, new_state: str) -> None:
202        """
203        Executes the action corresponding to new_state. if the action is defined under state_target,
204        that action will be taken in a new thread to avoid overloading the main tkinter thread. otherwise
205        it is executed immediately in the main thread.
206
207        Parameters
208        ----------
209        - **new_state** (*str*): Contains the desired state to move to. Used to locate the desired action.
210        """
211
212        thread_target = None
213
214        match new_state:
215            case "START PROGRAM":
216
217                def thread_target():
218                    StartProgram(
219                        self.exp_data,
220                        self.main_gui,
221                        self.arduino_controller,
222                        self.trigger,
223                    )
224            case "RESET PROGRAM":
225                ResetProgram(self.main_gui, self.app_result, self.arduino_controller)
226            case "STOP PROGRAM":
227                StopProgram(self.main_gui, self.arduino_controller, self.trigger)
228            case "GENERATE SCHEDULE":
229                # because this will run in main thread, we need to return early to avoid self.state
230                # assignment confusion
231                self.state = new_state
232                GenerateSchedule(
233                    self.main_gui,
234                    self.arduino_controller,
235                    self.process_queue,
236                    self.trigger,
237                )
238                return
239            case "ITI":
240
241                def thread_target():
242                    InitialTimeInterval(
243                        self.exp_data,
244                        self.main_gui,
245                        new_state,
246                        self.trigger,
247                    )
248            case "OPENING DOOR":
249
250                def thread_target():
251                    OpeningDoor(
252                        self.exp_data,
253                        self.main_gui,
254                        self.arduino_controller,
255                        new_state,
256                        self.trigger,
257                    )
258            case "TTC":
259
260                def thread_target():
261                    TimeToContact(
262                        self.exp_data,
263                        self.arduino_controller,
264                        self.main_gui,
265                        new_state,
266                        self.trigger,
267                    )
268            case "SAMPLE":
269
270                def thread_target():
271                    SampleTime(
272                        self.exp_data,
273                        self.main_gui,
274                        self.arduino_controller,
275                        new_state,
276                        self.trigger,
277                    )
278            case "TRIAL END":
279
280                def thread_target():
281                    TrialEnd(
282                        self.exp_data,
283                        self.main_gui,
284                        self.arduino_controller,
285                        self.prev_state,
286                        self.trigger,
287                    )
288
289        self.prev_state = self.state
290        self.state = new_state
291
292        if thread_target:
293            threading.Thread(target=thread_target).start()
294
295    def process_queue(self, data_queue: queue.Queue[tuple[str, str]]) -> None:
296        """
297        Process the data queue. Pulls in data from the thread that is reading data from the arduinos constantly. if there is anything
298        in the queue at time of function call, we will stay here until its all been dealt with.
299
300        Parameters
301        ----------
302        - **data_queue** (*tuple[str, str]*): This is a shared queue between a thread initiated in the `controllers.arduino_control` module's
303        `listen_for_serial` method. That thread constantly reads info from Arduino so it is not missed, and the main thread calls this process method
304        to process accumulated data to avoid overloading the thread.
305        """
306
307        arduino_data = self.exp_data.arduino_data
308        try:
309            while not data_queue.empty():
310                source, data = data_queue.get()
311                arduino_data.process_data(source, data, self.state, self.trigger)
312
313            # run this command every 250ms
314            queue_ps_id = self.main_gui.after(
315                250, lambda: self.process_queue(data_queue)
316            )
317            self.main_gui.scheduled_tasks["PROCESS QUEUE"] = queue_ps_id
318        except Exception as e:
319            logging.error(f"Error processing data queue: {e}")
320            raise
321
322    @staticmethod
323    def reject_actions(event):
324        """
325        This method handles the rejection of the 'START' and 'RESET' actions. These are rejected when there is no state transition for their current state.
326        This is usually during program runtime for reset or before schedule generation for start.
327
328        Parameters
329        ----------
330        - **event** (*str*): This is the desired state being rejected here.
331        """
332        match event:
333            case "RESET":
334                GUIUtils.display_error(
335                    "CANNOT PERFORM THIS ACTION",
336                    "RESET cannot be performed during runtime. Stop the program to reset.",
337                )
338            # if key is not in table but we try to start (for example, STOP PROGRAM -> START) we do NOT want to allow this until a full reset has been completed
339            case "START":
340                GUIUtils.display_error(
341                    "CANNOT PERFORM THIS ACTION",
342                    "Experiement has already been executed. Before running another, you need to RESET the application to ensure all data is reset to defaults",
343                )
344            case _:
345                return
346
347
348class GenerateSchedule:
349    """
350    This class handles the schedule generation state. It is called once the user has input the amount of simuli for this experiment, the
351    number of trial blocks, changed the names of the stimuli in each cylinder/valve in the 'Valve / Stimuli' window, and pressed generate schedule.
352
353    The main things handled here are the updating of gui objects that needed the previously discussed info to be created (e.g raster plots need
354    total trials, program schedule window needed the experiment data, etc.)
355    """
356
357    def __init__(
358        self,
359        main_gui: MainGUI,
360        arduino_controller: ArduinoManager,
361        process_queue: Callable[[queue.Queue[tuple[str, str]]], None],
362        trigger: Callable[[str], None],
363    ):
364        """
365        Initialize and handle the GenerateSchedule class state.
366
367        Parameters
368        ----------
369        - **main_gui** (*MainGUI*): A reference to the `views.main_gui` MainGUI instance, this is used to update elements inside the GUI window(s). Here we use it to
370        update show the program schedule and create raster plots.
371        - **arduino_controller** (*ArduinoManager*): A reference to the `controllers.arduino_control` Arduino controller instance, this is the method by which the Arduino is communicated
372        with in the program. Used to send schedules, variables, and valve durations here.
373        - **process_queue** (*Callback method*): This callback method is passed here so that the Arduino listener process thread can be started once the listener thread
374        is running. It does NOT run prior to this so that sending individual bytes (send schedule, etc.) is performed without having data stolen away by the constantly running
375        listener thread.
376        - **trigger** (*Callback method*): This callback is passed in so that this state can trigger a transition back to `IDLE` when it is finished with its work.
377        """
378
379        # show the program sched window via the callback passed into the class at initialization
380        main_gui.show_secondary_window("Program Schedule")
381
382        # create plots using generated number of trials as max Y values
383        for window in main_gui.windows["Raster Plot"]:
384            window.create_plot()
385
386        # send exp variables, schedule, and valve open durations stored in arduino_data.toml to the arduino
387        arduino_controller.send_experiment_variables()
388        arduino_controller.send_schedule_data()
389        arduino_controller.send_valve_durations()
390
391        # start Arduino process queue and listener thread so we know when Arduino tries to tell us something
392        process_queue(arduino_controller.data_queue)
393
394        arduino_controller.listener_thread = threading.Thread(
395            target=arduino_controller.listen_for_serial
396        )
397
398        arduino_controller.listener_thread.start()
399
400        logger.info("Started listening thread for Arduino serial input.")
401
402        # transition back to idle
403        trigger("IDLE")
404
405
406class StartProgram:
407    """
408    This class handles the remaining set-up steps to prepare for experiment runtime. It then triggers the experiment.
409    """
410
411    def __init__(
412        self,
413        exp_data: ExperimentProcessData,
414        main_gui: MainGUI,
415        arduino_controller: ArduinoManager,
416        trigger: Callable[[str], None],
417    ):
418        """
419        Parameters
420        ----------
421        - **exp_data** (*ExperimentProcessData*): Reference to the `models.experiment_process_data`. Here we use it to mark experiment and state start times.
422        - **main_gui** (*MainGUI*): A reference to `views.main_gui` MainGUI instance, this is used to update elements inside the GUI window(s). Here we use it to
423        update clock labels and max program runtime.
424        - **arduino_controller** (*ArduinoManager*): A reference to the `controllers.arduino_control` instance, this is the method by which the Arduino is communicated
425        with in the program. Used to tell Arduino that program begins now.
426        - **trigger** (*Callback method*): This callback is passed in so that this state can trigger a transition to `ITI` when it is finished with its work.
427        """
428        try:
429            # update experiment data model program start time and state start time variables with current time
430            exp_data.start_time = time.time()
431            exp_data.state_start_time = time.time()
432
433            # tell Arduino that experiment starts now so it knows how to calculate timestamps
434            start_command = "T=0\n".encode("utf-8")
435            arduino_controller.send_command(command=start_command)
436
437            main_gui.update_clock_label()
438            main_gui.update_max_time()
439
440            # change the green start button into a red stop button, update the associated command
441            main_gui.start_button.configure(
442                text="Stop", bg="red", command=lambda: trigger("STOP")
443            )
444
445            logging.info("==========EXPERIMENT BEGINS NOW==========")
446
447            trigger("ITI")
448        except Exception as e:
449            logging.error(f"Error starting program: {e}")
450            raise
451
452
453class StopProgram:
454    """
455    Stops and updates program to reflect `IDLE` state.
456
457    Methods
458    -------
459    - `finalize_program`(main_gui, arduino_controller)
460        This method waits until the door closes for the last time, then cancels the last scheduled task, stops the listener thread, closes Arduino connections, and
461        instructs the user to save their data.
462    """
463
464    def __init__(
465        self,
466        main_gui: MainGUI,
467        arduino_controller: ArduinoManager,
468        trigger: Callable[[str], None],
469    ) -> None:
470        """
471        Parameters
472        ----------
473        - **main_gui** (*MainGUI*): A reference to `views.main_gui` MainGUI instance, this is used to update elements inside the GUI window(s). Here we use it to
474        call update_on_stop, configure the start/stop button back to start, and cancel tkinter scheduled tasks.
475        - **arduino_controller** (*ArduinoManager*): A reference to the `controllers.arduino_control` instance, this is the method by which the Arduino is communicated
476        with in the program. Used to reset Arduino and clear all left-over experiment data.
477        - **trigger** (*Callback method*): This callback is passed in so the start button can be configured to command a `START` state.
478        """
479        try:
480            main_gui.update_on_stop()
481            # change the command back to start for the start button. (STOP PROGRAM, START) is not a defined transition, so the
482            # trigger function will call reject_actions to let the user know the program has already ran and they need to reset the app.
483            main_gui.start_button.configure(
484                text="Start", bg="green", command=lambda: trigger("START")
485            )
486
487            # stop all scheduled tasks EXCEPT for the process queue, we are still waiting for the last door close timestamp
488            for desc, sched_task in main_gui.scheduled_tasks.items():
489                if desc != "PROCESS QUEUE":
490                    main_gui.after_cancel(sched_task)
491
492            # schedule finalization after door will be down
493            main_gui.scheduled_tasks["FINALIZE"] = main_gui.after(
494                5000, lambda: self.finalize_program(main_gui, arduino_controller)
495            )
496
497            logging.info("Program stopped... waiting to finalize...")
498        except Exception as e:
499            logging.error(f"Error stopping program: {e}")
500            raise
501
502    def finalize_program(
503        self, main_gui: MainGUI, arduino_controller: ArduinoManager
504    ) -> None:
505        """
506        Finalize the program by resetting the Arduino board, offer to save the data frames into xlsx files.
507
508        Parameters
509        ----------
510        - **main_gui** (*MainGUI*): A reference to `views.main_gui` MainGUI instance, this is used to update elements inside the GUI window(s). Here we use it to
511        cancel the last tkinter.after call.
512        - **arduino_controller** (*ArduinoManager*): A reference to the `controllers.arduino_control` instance, this is the method by which the Arduino is communicated
513        with in the program. Used to reset Arduino board and close the connection to it.
514        """
515        try:
516            # stop the arduino listener so that the program can shut down
517            # and not be blocked
518            arduino_controller.stop_listener_thread()
519
520            queue_id = main_gui.scheduled_tasks["PROCESS QUEUE"]
521            main_gui.after_cancel(queue_id)
522
523            arduino_controller.close_connection()
524
525            main_gui.save_button_handler()
526
527            logging.info("Program finalized, arduino boards reset.")
528        except Exception as e:
529            logging.error(f"Error completing program: {e}")
530            raise
531
532
533class ResetProgram:
534    """
535    Handle resetting the program on click of the reset button. This class has no instance variable or methods.
536    """
537
538    def __init__(
539        self, main_gui: MainGUI, app_result: list, arduino_controller: ArduinoManager
540    ) -> None:
541        try:
542            # these two calls will stop the gui, halting the programs mainloop.
543            main_gui.quit()
544            main_gui.destroy()
545
546            # cancel all scheduled tasks
547            for sched_task in main_gui.scheduled_tasks.values():
548                main_gui.after_cancel(sched_task)
549            # if we have an listener thread and arduino connected, stop the thread and close the connection.
550            if arduino_controller.listener_thread is not None:
551                arduino_controller.stop_listener_thread()
552            if arduino_controller.arduino is not None:
553                arduino_controller.close_connection()
554
555            # set first bit in app_result list to 1 to instruct main to restart.
556            app_result[0] = 1
557
558        except Exception as e:
559            logging.error(f"Error resetting the program: {e}")
560
561
562class InitialTimeInterval:
563    """
564    State class for initial time interval experiment state.
565    """
566
567    def __init__(
568        self,
569        exp_data: ExperimentProcessData,
570        main_gui: MainGUI,
571        state: str,
572        trigger: Callable[[str], None],
573    ) -> None:
574        """
575        This function transitions the program into the `ITI` state. It does this by resetting lick counts, setting new state time,
576        updating the models, and setting the .after call to transition to TTC.
577
578        Parameters
579        ----------
580        - **exp_data** (*ExperimentProcessData*): Reference to the `models.experiment_process_data`. Here we use it to get a reference to event_data to set trial
581        licks to 0, get logical trial number, and get state duration time.
582        - **main_gui** (*MainGUI*): A reference to `views.main_gui` MainGUI instance, this is used to update elements inside the GUI window(s). Here we use it to
583        update GUI with new trial information and configure the state timer.
584        - **state** (*str*): State is used to update GUI with current state and grab the state duration time from the program schedule df.
585        - **trigger** (*Callback method*): This callback is passed in so the `OPENING DOOR` state can be triggered after the `ITI` time has passed.
586        """
587        try:
588            # get a reference to the event_data from exp_data.
589            event_data = exp_data.event_data
590
591            # set trial licks to 0 for both sides
592            event_data.side_one_licks = 0
593            event_data.side_two_licks = 0
594
595            # logical df rows will always be 1 behind the current trial because of zero based indexing in dataframes vs
596            # 1 based for trials
597            logical_trial = exp_data.current_trial_number - 1
598
599            # Call gui updates needed every trial (e.g program schedule window highlighting, progress bar, trial stimuli, etc)
600            # make sure we convert received vals to str to avoid type errors
601            main_gui.update_on_new_trial(
602                str(exp_data.program_schedule_df.loc[logical_trial, "Port 1"]),
603                str(exp_data.program_schedule_df.loc[logical_trial, "Port 2"]),
604            )
605
606            # clear state timer and reset the state start time
607            main_gui.state_timer_text.configure(text=(state + "Time:"))
608            exp_data.state_start_time = time.time()
609
610            initial_time_interval = exp_data.program_schedule_df.loc[
611                logical_trial, state
612            ]
613
614            main_gui.update_on_state_change(initial_time_interval, state)
615
616            # tell tkinter main loop that we want to trigger DOOR OPEN state after initial_time_interval milliseconds
617            iti_ttc_transition = main_gui.after(
618                int(initial_time_interval),
619                lambda: trigger("DOOR OPEN"),
620            )
621
622            # add the transition to scheduled tasks dict so that
623            main_gui.scheduled_tasks["ITI TO DOOR OPEN"] = iti_ttc_transition
624
625            logging.info(
626                f"STATE CHANGE: ITI BEGINS NOW for trial -> {exp_data.current_trial_number}, completes in {initial_time_interval}."
627            )
628        except Exception as e:
629            logging.error(f"Error in initial time interval: {e}")
630            raise
631
632
633class OpeningDoor:
634    """
635    Intermediate state between `ITI` and `TTC`.
636    """
637
638    # we are going to need to pass in the arduino controller as well
639    def __init__(
640        self,
641        exp_data: ExperimentProcessData,
642        main_gui: MainGUI,
643        arduino_controller: ArduinoManager,
644        state: str,
645        trigger: Callable[[str], None],
646    ):
647        """
648        In the initialization steps for this state we command the door down, then wait `DOOR_MOVE_TIME` to move to the `TTC` state.
649        """
650        # send comment to arduino to move the door down
651        down_command = "DOWN\n".encode("utf-8")
652        arduino_controller.send_command(command=down_command)
653
654        # after the door is down, then we will begin the ttc state logic, found in run_ttc
655        main_gui.after(DOOR_MOVE_TIME, lambda: trigger("TTC"))
656
657        # state start time begins
658        exp_data.state_start_time = time.time()
659
660        # show current_time / door close time to avoid confusion
661        main_gui.update_on_state_change(DOOR_MOVE_TIME, state)
662
663
664class TimeToContact:
665    """
666    State class for time to contact experiment state
667    """
668
669    def __init__(
670        self,
671        exp_data: ExperimentProcessData,
672        arduino_controller: ArduinoManager,
673        main_gui: MainGUI,
674        state: str,
675        trigger: Callable[[str], None],
676    ):
677        """
678        This function transitions the program into the `TTC` state. We can only reach this state from `OPENING DOOR`.
679
680        We transition into `TTC` by resetting the state timer, instructing the Arduino that the trial has started, and scheduling a transition into
681        `TRIAL END` if the rat does not engage in this trial. This is cancelled in sample time if 3 licks or more from `TTC` are detected in the `models.arduino_data`
682        module's `handle_licks` method.
683
684        Parameters
685        ----------
686        - **exp_data** (*ExperimentProcessData*): Reference to the `models.experiment_process_data`. Here we use it to get a reference to event_data to reset state time
687        and find `TTC` state time.
688        - **arduino_controller** (*ArduinoManager*): A reference to the `controllers.arduino_control` instance, this is the method by which the Arduino is communicated
689        with in the program. Used here to tell Arduino board trial has begun.
690        - **main_gui** (*MainGUI*): A reference to `views.main_gui` MainGUI instance, this is used to update elements inside the GUI window(s). Here we use it to
691        update GUI with new state time.
692        - **state** (*str*): State is used to update GUI with current state and grab the state duration time from the program schedule df.
693        - **trigger** (*Callback method*): This callback is passed in so the `TRIAL END` state can be triggered if trial is not engaged.
694        """
695        try:
696            logging.info(
697                f"STATE CHANGE: DOOR OPEN -> TTC NOW for trial -> {exp_data.current_trial_number}."
698            )
699
700            logical_trial = exp_data.current_trial_number - 1
701
702            # tell the arduino that the trial begins now, becuase the rats are able to licks beginning now
703            command = "TRIAL START\n".encode("utf-8")
704            arduino_controller.send_command(command)
705
706            main_gui.state_timer_text.configure(text=(state + " Time:"))
707
708            # state time starts now, as does trial because lick availabilty starts now
709            exp_data.state_start_time = time.time()
710            exp_data.trial_start_time = time.time()
711
712            # find the amount of time available for TTC for this trial
713            time_to_contact = exp_data.program_schedule_df.loc[logical_trial, state]
714
715            main_gui.update_on_state_change(time_to_contact, state)
716
717            # set a state change to occur after the time_to_contact time, this will be cancelled if the laser arduino
718            # sends 3 licks befote the TTC_time
719            ttc_iti_transition = main_gui.after(
720                int(time_to_contact), lambda: trigger("TRIAL END")
721            )
722
723            # store this task id so that it can be cancelled if the sample time transition is taken.
724            main_gui.scheduled_tasks["TTC TO TRIAL END"] = ttc_iti_transition
725
726            logging.info(
727                f"STATE CHANGE: TTC BEGINS NOW for trial -> {exp_data.current_trial_number}, completes in {time_to_contact}."
728            )
729        except Exception as e:
730            logging.error(f"Error in TTC state start: {e}")
731            raise
732
733
734class SampleTime:
735    """
736    Sample state class to handle setup of `Sample` state.
737    """
738
739    def __init__(
740        self,
741        exp_data: ExperimentProcessData,
742        main_gui: MainGUI,
743        arduino_controller: ArduinoManager,
744        state: str,
745        trigger: Callable[[str], None],
746    ):
747        """
748        This function transitions the program into the `SAMPLE` state. We can only reach this state from `TTC` if 3 licks or more are detected in `TTC`.
749
750        We transition into `SAMPLE` by resetting licks to zero to count only sample time licks for trial, updating the actual `TTC` time taken before engaging in the
751        trial, cancelling the scheduled `TTC` to `TRIAL END` transition, instructing Arduino to begin opening valves when licks occur, and updating state start time.
752
753        Parameters
754        ----------
755        - **exp_data** (*ExperimentProcessData*): Reference to the `models.experiment_process_data`. Here we use it to get a reference to event_data to reset state time
756        and find `SAMPLE` state time, reset trial lick data.
757        - **main_gui** (*MainGUI*): A reference to `views.main_gui` MainGUI instance, this is used to update elements inside the GUI window(s). Here we use it to
758        update GUI with new state info and cancel `TTC` to `TRIAL END` transition.
759        - **arduino_controller** (*ArduinoManager*): A reference to the `controllers.arduino_control` instance, this is the method by which the Arduino is communicated
760        with in the program. Used here to tell Arduino to begin opening valves when licks are detected.
761        - **state** (*str*): State is used to update GUI with current state and grab the state duration time from the program schedule df.
762        - **trigger** (*Callback method*): This callback is passed in so the `TRIAL END` state can be triggered after `SAMPLE` time.
763
764        Methods
765        -------
766        - `update_ttc_time`(exp_data, logical_trial)
767        Updates program schedule dataframe with actual time used in the `TTC` state.
768        """
769        try:
770            logical_trial = exp_data.current_trial_number - 1
771
772            event_data = exp_data.event_data
773
774            # reset trial licks to count only sample licks
775            event_data.side_one_licks = 0
776            event_data.side_two_licks = 0
777
778            self.update_ttc_time(exp_data, logical_trial)
779
780            # since the trial was engaged, cancel the planned transition from ttc to trial end. program has branched to new state
781            main_gui.after_cancel(main_gui.scheduled_tasks["TTC TO TRIAL END"])
782
783            # Tell the laser arduino to begin accepting licks and opening valves
784            # resetting the lick counters for both spouts
785            # tell the laser to begin opening valves on licks
786            open_command = "BEGIN OPEN VALVES\n".encode("utf-8")
787            arduino_controller.send_command(command=open_command)
788
789            exp_data.state_start_time = time.time()
790
791            sample_interval_value = exp_data.program_schedule_df.loc[
792                logical_trial, state
793            ]
794
795            main_gui.update_on_state_change(sample_interval_value, state)
796
797            main_gui.after(
798                int(sample_interval_value),
799                lambda: trigger("TRIAL END"),
800            )
801
802            logging.info(
803                f"STATE CHANGE: SAMPLE BEGINS NOW for trial-> {exp_data.current_trial_number}, completes in {sample_interval_value}."
804            )
805        except Exception as e:
806            logging.error(
807                f"Error in sample time trial no. {exp_data.current_trial_number}: {e}"
808            )
809            raise
810
811    def update_ttc_time(
812        self, exp_data: ExperimentProcessData, logical_trial: int
813    ) -> None:
814        """
815        If we reach this code that means the rat licked at least 3 times in `TTC` state
816        so we didn't take all the allocated time. update program schedule window with actual time taken
817        """
818        ttc_time = (time.time() - exp_data.state_start_time) * 1000
819
820        exp_data.program_schedule_df.loc[logical_trial, "TTC Actual"] = round(
821            ttc_time, 3
822        )
823
824
825class TrialEnd:
826    """
827    Handle end of trial operations such as data storage and GUI updates with gathered trial data.
828
829    Methods
830    -------
831    - `arduino_trial_end`(arduino_controller): Handle Arduino end of trial, send door up, stop accepting licks.
832    - `end_trial`(exp_data: ExperimentProcessData): Determine if this trial is the last, if not increment trial number in `models.experiment_process_data`.
833    - `handle_from_ttc`(logical_trial, trigger), exp_data): Call `update_ttc_actual` to update TTC time. Trigger transition to `ITI`.
834    - `update_schedule_licks`(logical_trial, exp_data) -> None: Update program schedule df with licks for this trial on each port.
835    - `update_raster_plots`(exp_data, logical_trial, main_gui) -> None: Update raster plot with licks from this trial.
836    """
837
838    def __init__(
839        self,
840        exp_data: ExperimentProcessData,
841        main_gui: MainGUI,
842        arduino_controller: ArduinoManager,
843        prev_state: str,
844        trigger: Callable[[str], None],
845    ) -> None:
846        """
847        This function transitions the program into the `SAMPLE` state. We can only reach this state from `TTC` if 3 licks or more are detected in `TTC`.
848
849        We transition into `SAMPLE` by resetting licks to zero to count only sample time licks for trial, updating the actual `TTC` time taken before engaging in the
850        trial, cancelling the scheduled `TTC` to `TRIAL END` transition, instructing Arduino to begin opening valves when licks occur, and updating state start time.
851
852        Parameters
853        ----------
854        - **exp_data** (*ExperimentProcessData*): Reference to the `models.experiment_process_data`. Here we use it to update program df with trial licks, increment trial,
855        check if current trial is the final trial, and other data operations.
856        - **main_gui** (*MainGUI*): A reference to `views.main_gui` MainGUI instance, this is used to update elements inside the GUI window(s). Here we use it to
857        update program schedule and populate raster plots with trial licks.
858        - **arduino_controller** (*ArduinoManager*): A reference to the `controllers.arduino_control` instance, this is the method by which the Arduino is communicated
859        with in the program. Used here to tell Arduino to stop opening valves when licks are detected.
860        - **prev_state** (*str*): prev_state is used to decide which end of trial operations should be executed.
861        - **trigger** (*Callback method*): This callback is passed in so the `ITI` state can be triggered after `TRIAL END` logic is complete.
862
863        Methods
864        -------
865        - `update_ttc_time`(exp_data, logical_trial)
866        Updates program schedule dataframe with actual time used in the `TTC` state.
867        """
868
869        # use the same logical trial for all functions current_trial_number - 1
870        # to account for 1 indexing
871        logical_trial = exp_data.current_trial_number - 1
872
873        self.arduino_trial_end(arduino_controller)
874
875        #######TRIAL NUMBER IS INCREMENTED INSIDE OF END_TRIAL#######
876        if self.end_trial(exp_data):
877            program_schedule = main_gui.windows["Program Schedule"]
878            # if the experiment is over update the licks for the final trial
879            self.update_schedule_licks(logical_trial, exp_data)
880            program_schedule.refresh_end_trial(logical_trial)
881
882            trigger("STOP")
883            # return from the call / kill the working thread
884            return
885
886        # update licks for this trial
887        self.update_schedule_licks(logical_trial, exp_data)
888
889        match prev_state:
890            case "TTC":
891                self.update_ttc_actual(logical_trial, exp_data)
892                trigger("ITI")
893            case "SAMPLE":
894                self.update_raster_plots(exp_data, logical_trial, main_gui)
895                trigger("ITI")
896            case _:
897                # cases not explicitly defined go here
898                logger.error("UNDEFINED PREVIOUS TRANSITION IN TRIAL END STATE")
899
900        main_gui.windows["Program Schedule"].refresh_end_trial(logical_trial)
901
902    def arduino_trial_end(self, arduino_controller: ArduinoManager) -> None:
903        """
904        Send rig door up and tell Arduino to stop accepting licks.
905        """
906        up_command = "UP\n".encode("utf-8")
907        arduino_controller.send_command(command=up_command)
908
909        stop_command = "STOP OPEN VALVES\n".encode("utf-8")
910        arduino_controller.send_command(command=stop_command)
911
912    def end_trial(self, exp_data: ExperimentProcessData) -> bool:
913        """
914        This method checks if current trial is the last trial. If it is, we return true to the calling code.
915        If false, we increment current_trial_number, return false, and proceed to next trial.
916        """
917        current_trial = exp_data.current_trial_number
918        max_trials = exp_data.exp_var_entries["Num Trials"]
919
920        if current_trial >= max_trials:
921            return True
922
923        exp_data.current_trial_number += 1
924        return False
925
926    def update_ttc_actual(
927        self, logical_trial: int, exp_data: ExperimentProcessData
928    ) -> None:
929        """
930        Update the actual ttc time take in the program schedule with the max time value,
931        since we reached trial end we know we took the max time allowed
932        """
933        program_df = exp_data.program_schedule_df
934
935        ttc_column_value = program_df.loc[logical_trial, "TTC"]
936        program_df.loc[logical_trial, "TTC Actual"] = ttc_column_value
937
938    def update_schedule_licks(
939        self, logical_trial: int, exp_data: ExperimentProcessData
940    ) -> None:
941        """Update the licks for each side in the program schedule df for this trial"""
942        program_df = exp_data.program_schedule_df
943        event_data = exp_data.event_data
944
945        licks_sd_one = event_data.side_one_licks
946        licks_sd_two = event_data.side_two_licks
947
948        program_df.loc[logical_trial, "Port 1 Licks"] = licks_sd_one
949
950        program_df.loc[logical_trial, "Port 2 Licks"] = licks_sd_two
951
952    def update_raster_plots(
953        self, exp_data: ExperimentProcessData, logical_trial: int, main_gui: MainGUI
954    ) -> None:
955        """Instruct raster windows to update with new lick timestamps"""
956        # gather lick timestamp data from the event data model
957        lick_stamps = exp_data.event_data.get_lick_timestamps(logical_trial)
958
959        # send it to raster windows
960        for i, window in enumerate(main_gui.windows["Raster Plot"]):
961            window.update_plot(lick_stamps[i], logical_trial)