From 6fc6347f81b10fc225b2a59d11fa3d72394a66ba Mon Sep 17 00:00:00 2001
From: fedem-p <federico.puppo@tum.de>
Date: Tue, 1 Mar 2022 14:52:44 +0100
Subject: [PATCH 01/13] calibrator initialized from the protocol runner; new
 init_external get's also a calibrator object

---
 stytra/stimulation/__init__.py                |  2 +-
 stytra/stimulation/stimuli/conditional.py     | 28 ++++++++---------
 stytra/stimulation/stimuli/external.py        |  2 +-
 stytra/stimulation/stimuli/generic_stimuli.py | 16 +++++-----
 stytra/stimulation/stimuli/kinematograms.py   |  4 +--
 stytra/stimulation/stimuli/visual.py          | 30 +++++++++----------
 6 files changed, 42 insertions(+), 40 deletions(-)

diff --git a/stytra/stimulation/__init__.py b/stytra/stimulation/__init__.py
index 4749c334..b9139f39 100644
--- a/stytra/stimulation/__init__.py
+++ b/stytra/stimulation/__init__.py
@@ -112,7 +112,7 @@ def update_protocol(self):
 
         # pass experiment to stimuli for calibrator and asset folders:
         for stimulus in self.stimuli:
-            stimulus.initialise_external(self.experiment)
+            stimulus.initialise_external(self.experiment, self.experiment.calibrator)
 
         if self.dynamic_log is None:
             self.dynamic_log = DynamicLog(self.stimuli, experiment=self.experiment)
diff --git a/stytra/stimulation/stimuli/conditional.py b/stytra/stimulation/stimuli/conditional.py
index 1dd6f69f..c5d98a25 100644
--- a/stytra/stimulation/stimuli/conditional.py
+++ b/stytra/stimulation/stimuli/conditional.py
@@ -36,9 +36,9 @@ def get_dynamic_state(self):
             state.update(self.active.get_dynamic_state())
         return state
 
-    def initialise_external(self, experiment):
-        super().initialise_external(experiment)
-        self.active.initialise_external(experiment)
+    def initialise_external(self, experiment, calibrator):
+        super().initialise_external(experiment, calibrator)
+        self.active.initialise_external(experiment, calibrator)
 
     def get_state(self):
         state = super().get_state()
@@ -50,7 +50,7 @@ def start(self):
         self.active.start()
 
     def check_condition(self):
-        y, x, theta = self._experiment.estimator.get_position()
+        y, x, theta = self._experiment.estimator.get_position()#! TOFIX: Remove
         return not np.isnan(y)
 
     def update(self):
@@ -157,10 +157,10 @@ def get_dynamic_state(self):
             state.update(self._stim_on.get_dynamic_state())
         return state
 
-    def initialise_external(self, experiment):
-        super().initialise_external(experiment)
-        self._stim_on.initialise_external(experiment)
-        self._stim_off.initialise_external(experiment)
+    def initialise_external(self, experiment, calibrator):
+        super().initialise_external(experiment, calibrator)
+        self._stim_on.initialise_external(experiment, calibrator)
+        self._stim_off.initialise_external(experiment, calibrator)
 
     def get_state(self):
         state = super().get_state()
@@ -270,8 +270,8 @@ def __init__(self, stimulus, *args, centering_stimulus=None, margin=45, **kwargs
         self.yc = 240
 
     def check_condition_on(self):
-        y, x, theta = self._experiment.estimator.get_position()
-        scale = self._experiment.calibrator.mm_px ** 2
+        y, x, theta = self._experiment.estimator.get_position()#! TOFIX: Remove
+        scale = self._calibrator.mm_px ** 2 
         return (
             x > 0 and ((x - self.xc) ** 2 + (y - self.yc) ** 2) <= self.margin / scale
         )
@@ -323,15 +323,15 @@ def __init__(
         self.yc = 240
 
     def check_condition_on(self):
-        y, x, theta = self._experiment.estimator.get_position()
-        scale = self._experiment.calibrator.mm_px ** 2
+        y, x, theta = self._experiment.estimator.get_position() #! TOFIX: Remove
+        scale = self._calibrator.mm_px ** 2
         return (not np.isnan(x)) and (
             (x - self.xc) ** 2 + (y - self.yc) ** 2 <= self.margin_in / scale
         )
 
     def check_condition_off(self):
-        y, x, theta = self._experiment.estimator.get_position()
-        scale = self._experiment.calibrator.mm_px ** 2
+        y, x, theta = self._experiment.estimator.get_position() #! TOFIX: Remove
+        scale = self._calibrator.mm_px ** 2
         return np.isnan(x) or (
             (x - self.xc) ** 2 + (y - self.yc) ** 2 > self.margin_out / scale
         )
diff --git a/stytra/stimulation/stimuli/external.py b/stytra/stimulation/stimuli/external.py
index 42319cb5..9df28c6f 100644
--- a/stytra/stimulation/stimuli/external.py
+++ b/stytra/stimulation/stimuli/external.py
@@ -42,7 +42,7 @@ def __init__(
         pulse_dur_str = str(pulse_dur_ms).zfill(3)
         self.mex = str("shock" + amp_dac + pulse_dur_str)
 
-    def initialise_external(self, experiment):
+    def initialise_external(self, experiment, calibrator):
         """
 
         Parameters
diff --git a/stytra/stimulation/stimuli/generic_stimuli.py b/stytra/stimulation/stimuli/generic_stimuli.py
index 4b28d218..e927906d 100644
--- a/stytra/stimulation/stimuli/generic_stimuli.py
+++ b/stytra/stimulation/stimuli/generic_stimuli.py
@@ -65,7 +65,8 @@ def __init__(self, duration=0.0):
         self._started = None
         self._elapsed = 0.0  # time from the beginning of the stimulus
         self.name = "undefined"
-        self._experiment = None
+        self._experiment = None  #! TOFIX: Remove
+        self._calibrator = None
         self.real_time_start = None
         self.real_time_stop = None
 
@@ -111,7 +112,7 @@ def stop(self):
         """
         pass
 
-    def initialise_external(self, experiment):
+    def initialise_external(self, experiment, calibrator = None):
         """ Make a reference to the Experiment class inside the Stimulus.
         This is required to access from inside the Stimulus class to the
         Calibrator, the Pyboard, the asset directories with movies or the motor
@@ -130,7 +131,8 @@ def initialise_external(self, experiment):
             None
 
         """
-        self._experiment = experiment
+        self._experiment = experiment #! TOFIX: Remove
+        self._calibrator = calibrator
 
 
 class DynamicStimulus(Stimulus):
@@ -251,7 +253,7 @@ def start(self):
 
     def update(self):
         # If trigger is set, make it end:
-        if self._experiment.trigger.start_event.is_set():
+        if self._experiment.trigger.start_event.is_set(): #! TOFIX: Remove ?
             self.duration = self._elapsed
 
 
@@ -289,10 +291,10 @@ def update(self):
             s.update()
             s._elapsed = self._elapsed
 
-    def initialise_external(self, experiment):
-        super().initialise_external(experiment)
+    def initialise_external(self, experiment, calibrator):
+        super().initialise_external(experiment, calibrator)
         for s in self._stim_list:
-            s.initialise_external(experiment)
+            s.initialise_external(experiment, calibrator)
 
     @property
     def dynamic_parameter_names(self):
diff --git a/stytra/stimulation/stimuli/kinematograms.py b/stytra/stimulation/stimuli/kinematograms.py
index 9f69355b..86556761 100644
--- a/stytra/stimulation/stimuli/kinematograms.py
+++ b/stytra/stimulation/stimuli/kinematograms.py
@@ -75,8 +75,8 @@ def get_dimensions(self):
         -------
         number of dots to display and the displacement amount in pixel coordinates
         """
-        if self._experiment.calibrator is not None:
-            mm_px = self._experiment.calibrator.mm_px
+        if self._calibrator is not None:
+            mm_px = self._calibrator.mm_px  
         else:
             mm_px = 1
 
diff --git a/stytra/stimulation/stimuli/visual.py b/stytra/stimulation/stimuli/visual.py
index 17bb8429..7b266cdb 100644
--- a/stytra/stimulation/stimuli/visual.py
+++ b/stytra/stimulation/stimuli/visual.py
@@ -206,7 +206,7 @@ def __init__(self, *args, video_path, framerate=None, duration=None, **kwargs):
 
     def initialise_external(self, *args, **kwargs):
         super().initialise_external(*args, **kwargs)
-        self._video_seq = pims.Video(self._experiment.asset_dir + "/" + self.video_path)
+        self._video_seq = pims.Video(self._experiment.asset_dir + "/" + self.video_path) #! TOFIX: Remove
 
         self._current_frame = self._video_seq.get_frame(self.i_frame)
         try:
@@ -313,8 +313,8 @@ def get_tile_ranges(self, imw, imh, w, h, tr: QTransform):
         return range(x_start, x_end + 1), range(y_start, y_end + 1)
 
     def paint(self, p, w, h):
-        if self._experiment.calibrator is not None:
-            mm_px = self._experiment.calibrator.mm_px
+        if self._calibrator is not None: 
+            mm_px = self._calibrator.mm_px
         else:
             mm_px = 1
 
@@ -396,8 +396,8 @@ def __init__(self, *args, background, background_name=None, **kwargs):
                 self.background_name = "array {}x{}".format(*self._background.shape)
         self._qbackground = None
 
-    def initialise_external(self, experiment):
-        super().initialise_external(experiment)
+    def initialise_external(self, experiment, calibrator):
+        super().initialise_external(experiment, calibrator)
 
         # Get background image from folder:
         if isinstance(self._background, str):
@@ -468,7 +468,7 @@ def __init__(
     def create_pattern(self):
         l = max(
             2,
-            int(self.grating_period / (max(self._experiment.calibrator.mm_px, 0.0001))),
+            int(self.grating_period / (max(self._calibrator.mm_px, 0.0001))),
         )
         if self.wave_shape == "square":
             self._pattern = np.ones((l, 3), np.uint8) * self.color_1
@@ -483,8 +483,8 @@ def create_pattern(self):
                 + (1 - w[:, None]) * np.array(self.color_2)[None, :]
             ).astype(np.uint8)
 
-    def initialise_external(self, experiment):
-        super().initialise_external(experiment)
+    def initialise_external(self, experiment, calibrator):
+        super().initialise_external(experiment, calibrator)
         self.create_pattern()
         # Get background image from folder:
         self._qbackground = qimage2ndarray.array2qimage(self._pattern[None, :, :])
@@ -532,7 +532,7 @@ def get_unit_dims(self, w, h):
         #TODO what does this thing define?
         """
         return (
-            int(self.grating_period / (max(self._experiment.calibrator.mm_px, 0.0001))),
+            int(self.grating_period / (max(self._calibrator.mm_px, 0.0001))),
             self.barheight,
         )
 
@@ -547,7 +547,7 @@ def draw_block(self, p, point, w, h):
             point.y(),
             int(
                 self.grating_period
-                / (2 * max(self._experiment.calibrator.mm_px, 0.0001))
+                / (2 * max(self._calibrator.mm_px, 0.0001))
             ),
             self.barheight,
         )
@@ -654,7 +654,7 @@ def update(self):
 
     def paint(self, p, w, h):
         x, y = (
-            (np.arange(d) - d / 2) * self._experiment.calibrator.mm_px for d in (w, h)
+            (np.arange(d) - d / 2) * self._calibrator.mm_px for d in (w, h)
         )
         self.image = np.round(
             np.sin(
@@ -751,8 +751,8 @@ def create_pattern(self, side_len=500):
         self._pattern = W * self.color_1 + (1 - W) * self.color_2
         self._qbackground = qimage2ndarray.array2qimage(self._pattern)
 
-    def initialise_external(self, experiment):
-        super().initialise_external(experiment)
+    def initialise_external(self, experiment, calibrator):
+        super().initialise_external(experiment, calibrator)
         self.create_pattern()
 
     def draw_block(self, p, point, w, h):
@@ -933,8 +933,8 @@ def __init__(
     def paint(self, p, w, h):
         super().paint(p, w, h)
 
-        if self._experiment.calibrator is not None:
-            mm_px = self._experiment.calibrator.mm_px
+        if self._calibrator is not None:
+            mm_px = self._calibrator.mm_px
         else:
             mm_px = 1
 

From 4e96d167e81bbf2db157adbefad39093fdefb6f0 Mon Sep 17 00:00:00 2001
From: fedem-p <federico.puppo@tum.de>
Date: Tue, 1 Mar 2022 16:09:44 +0100
Subject: [PATCH 02/13] Added warning for deprecated function inputs

---
 stytra/stimulation/__init__.py                |  8 +++++++-
 stytra/stimulation/stimuli/generic_stimuli.py | 15 +++++++++++++--
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/stytra/stimulation/__init__.py b/stytra/stimulation/__init__.py
index b9139f39..6fa415e3 100644
--- a/stytra/stimulation/__init__.py
+++ b/stytra/stimulation/__init__.py
@@ -1,5 +1,6 @@
 import datetime
 from copy import deepcopy
+import warnings
 
 from PyQt5.QtCore import pyqtSignal, QTimer, QObject
 from stytra.stimulation.stimuli import Pause, DynamicStimulus
@@ -112,7 +113,12 @@ def update_protocol(self):
 
         # pass experiment to stimuli for calibrator and asset folders:
         for stimulus in self.stimuli:
-            stimulus.initialise_external(self.experiment, self.experiment.calibrator)
+            try:
+                stimulus.initialise_external(self.experiment, self.experiment.calibrator)
+            except TypeError:
+                stimulus.initialise_external(self.experiment)
+                warnings.warn("Warning: 'initialise_external' will require a calibrator input from the new update!", FutureWarning)
+                warnings.warn("Warning: 'initialise_external' will require a calibrator input from the new update!", DeprecationWarning)
 
         if self.dynamic_log is None:
             self.dynamic_log = DynamicLog(self.stimuli, experiment=self.experiment)
diff --git a/stytra/stimulation/stimuli/generic_stimuli.py b/stytra/stimulation/stimuli/generic_stimuli.py
index e927906d..ca8f39f4 100644
--- a/stytra/stimulation/stimuli/generic_stimuli.py
+++ b/stytra/stimulation/stimuli/generic_stimuli.py
@@ -1,5 +1,6 @@
 import numpy as np
 import datetime
+import warnings
 
 
 class Stimulus:
@@ -112,7 +113,7 @@ def stop(self):
         """
         pass
 
-    def initialise_external(self, experiment, calibrator = None):
+    def initialise_external(self, experiment, calibrator = -999):
         """ Make a reference to the Experiment class inside the Stimulus.
         This is required to access from inside the Stimulus class to the
         Calibrator, the Pyboard, the asset directories with movies or the motor
@@ -131,8 +132,18 @@ def initialise_external(self, experiment, calibrator = None):
             None
 
         """
+        
+        if calibrator == -999:
+            self._calibrator = self._experiment.calibrator
+            warnings.warn("Warning: 'initialise_external' will require a calibrator input from the new update!", FutureWarning)
+            warnings.warn("Warning: 'initialise_external' will require a calibrator input from the new update!", DeprecationWarning)
+        else:
+            self._calibrator = calibrator
+            
+            
         self._experiment = experiment #! TOFIX: Remove
-        self._calibrator = calibrator
+        
+    
 
 
 class DynamicStimulus(Stimulus):

From 8e79c627b7bd200f2949b9a9142f61997df38da8 Mon Sep 17 00:00:00 2001
From: fedem-p <federico.puppo@tum.de>
Date: Tue, 1 Mar 2022 16:38:45 +0100
Subject: [PATCH 03/13] Cleaned useless comments

---
 stytra/stimulation/stimuli/conditional.py     | 8 ++++----
 stytra/stimulation/stimuli/generic_stimuli.py | 6 +++---
 stytra/stimulation/stimuli/visual.py          | 2 +-
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/stytra/stimulation/stimuli/conditional.py b/stytra/stimulation/stimuli/conditional.py
index c5d98a25..8300da63 100644
--- a/stytra/stimulation/stimuli/conditional.py
+++ b/stytra/stimulation/stimuli/conditional.py
@@ -50,7 +50,7 @@ def start(self):
         self.active.start()
 
     def check_condition(self):
-        y, x, theta = self._experiment.estimator.get_position()#! TOFIX: Remove
+        y, x, theta = self._experiment.estimator.get_position()
         return not np.isnan(y)
 
     def update(self):
@@ -270,7 +270,7 @@ def __init__(self, stimulus, *args, centering_stimulus=None, margin=45, **kwargs
         self.yc = 240
 
     def check_condition_on(self):
-        y, x, theta = self._experiment.estimator.get_position()#! TOFIX: Remove
+        y, x, theta = self._experiment.estimator.get_position()
         scale = self._calibrator.mm_px ** 2 
         return (
             x > 0 and ((x - self.xc) ** 2 + (y - self.yc) ** 2) <= self.margin / scale
@@ -323,14 +323,14 @@ def __init__(
         self.yc = 240
 
     def check_condition_on(self):
-        y, x, theta = self._experiment.estimator.get_position() #! TOFIX: Remove
+        y, x, theta = self._experiment.estimator.get_position()
         scale = self._calibrator.mm_px ** 2
         return (not np.isnan(x)) and (
             (x - self.xc) ** 2 + (y - self.yc) ** 2 <= self.margin_in / scale
         )
 
     def check_condition_off(self):
-        y, x, theta = self._experiment.estimator.get_position() #! TOFIX: Remove
+        y, x, theta = self._experiment.estimator.get_position()
         scale = self._calibrator.mm_px ** 2
         return np.isnan(x) or (
             (x - self.xc) ** 2 + (y - self.yc) ** 2 > self.margin_out / scale
diff --git a/stytra/stimulation/stimuli/generic_stimuli.py b/stytra/stimulation/stimuli/generic_stimuli.py
index ca8f39f4..9b4431cd 100644
--- a/stytra/stimulation/stimuli/generic_stimuli.py
+++ b/stytra/stimulation/stimuli/generic_stimuli.py
@@ -66,7 +66,7 @@ def __init__(self, duration=0.0):
         self._started = None
         self._elapsed = 0.0  # time from the beginning of the stimulus
         self.name = "undefined"
-        self._experiment = None  #! TOFIX: Remove
+        self._experiment = None
         self._calibrator = None
         self.real_time_start = None
         self.real_time_stop = None
@@ -141,7 +141,7 @@ def initialise_external(self, experiment, calibrator = -999):
             self._calibrator = calibrator
             
             
-        self._experiment = experiment #! TOFIX: Remove
+        self._experiment = experiment
         
     
 
@@ -264,7 +264,7 @@ def start(self):
 
     def update(self):
         # If trigger is set, make it end:
-        if self._experiment.trigger.start_event.is_set(): #! TOFIX: Remove ?
+        if self._experiment.trigger.start_event.is_set():
             self.duration = self._elapsed
 
 
diff --git a/stytra/stimulation/stimuli/visual.py b/stytra/stimulation/stimuli/visual.py
index 7b266cdb..d4899d31 100644
--- a/stytra/stimulation/stimuli/visual.py
+++ b/stytra/stimulation/stimuli/visual.py
@@ -206,7 +206,7 @@ def __init__(self, *args, video_path, framerate=None, duration=None, **kwargs):
 
     def initialise_external(self, *args, **kwargs):
         super().initialise_external(*args, **kwargs)
-        self._video_seq = pims.Video(self._experiment.asset_dir + "/" + self.video_path) #! TOFIX: Remove
+        self._video_seq = pims.Video(self._experiment.asset_dir + "/" + self.video_path)
 
         self._current_frame = self._video_seq.get_frame(self.i_frame)
         try:

From 516d8935a95f5146f7ad299a8715f92162b7efb8 Mon Sep 17 00:00:00 2001
From: fedem-p <federico.puppo@tum.de>
Date: Mon, 7 Mar 2022 17:46:05 +0100
Subject: [PATCH 04/13] added EnvironmentState dataclass; and changed
 calibrator init for stimulus

---
 stytra/stimulation/__init__.py                | 16 +++++++++++++---
 stytra/stimulation/stimuli/generic_stimuli.py | 13 +++++++------
 2 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/stytra/stimulation/__init__.py b/stytra/stimulation/__init__.py
index 6fa415e3..2392891e 100644
--- a/stytra/stimulation/__init__.py
+++ b/stytra/stimulation/__init__.py
@@ -1,3 +1,4 @@
+from dataclasses import dataclass
 import datetime
 from copy import deepcopy
 import warnings
@@ -10,6 +11,12 @@
 
 import logging
 
+@dataclass
+class EnvironmentState:
+    _calibrator = None
+    height = 0
+    width = 0
+        
 
 class ProtocolRunner(QObject):
     """Class for managing and running stimulation Protocols.
@@ -100,6 +107,9 @@ def __init__(self, experiment=None, target_dt=0, log_print=True):
         self.log = []
         self.log_print = log_print
         self.running = False
+        
+        self.environment_state = EnvironmentState(_calibrator = self.experiment.calibrator,
+                                                  )
 
         self.framerate_rec = FramerateRecorder()
         self.framerate_acc = FramerateAccumulator(experiment=self.experiment)
@@ -114,11 +124,11 @@ def update_protocol(self):
         # pass experiment to stimuli for calibrator and asset folders:
         for stimulus in self.stimuli:
             try:
-                stimulus.initialise_external(self.experiment, self.experiment.calibrator)
+                stimulus.initialise_external(self.experiment, self.environment_state,)
             except TypeError:
                 stimulus.initialise_external(self.experiment)
-                warnings.warn("Warning: 'initialise_external' will require a calibrator input from the new update!", FutureWarning)
-                warnings.warn("Warning: 'initialise_external' will require a calibrator input from the new update!", DeprecationWarning)
+                warnings.warn("Warning: 'initialise_external' will use the environment_state variable which holds the calibrator object!", FutureWarning)
+                warnings.warn("Warning: 'initialise_external' will use the environment_state variable which holds the calibrator object!", DeprecationWarning)
 
         if self.dynamic_log is None:
             self.dynamic_log = DynamicLog(self.stimuli, experiment=self.experiment)
diff --git a/stytra/stimulation/stimuli/generic_stimuli.py b/stytra/stimulation/stimuli/generic_stimuli.py
index 9b4431cd..113a3450 100644
--- a/stytra/stimulation/stimuli/generic_stimuli.py
+++ b/stytra/stimulation/stimuli/generic_stimuli.py
@@ -1,6 +1,7 @@
 import numpy as np
 import datetime
 import warnings
+from stytra.stimulation import EnvironmentState
 
 
 class Stimulus:
@@ -113,7 +114,7 @@ def stop(self):
         """
         pass
 
-    def initialise_external(self, experiment, calibrator = -999):
+    def initialise_external(self, experiment, environment_state: EnvironmentState = None):
         """ Make a reference to the Experiment class inside the Stimulus.
         This is required to access from inside the Stimulus class to the
         Calibrator, the Pyboard, the asset directories with movies or the motor
@@ -133,12 +134,12 @@ def initialise_external(self, experiment, calibrator = -999):
 
         """
         
-        if calibrator == -999:
-            self._calibrator = self._experiment.calibrator
-            warnings.warn("Warning: 'initialise_external' will require a calibrator input from the new update!", FutureWarning)
-            warnings.warn("Warning: 'initialise_external' will require a calibrator input from the new update!", DeprecationWarning)
+        if isinstance(environment_state, EnvironmentState):
+            self._calibrator = environment_state._calibrator
         else:
-            self._calibrator = calibrator
+            self._calibrator = self._experiment.calibrator
+            warnings.warn("Warning: 'initialise_external' will use the environment_state variable which holds the calibrator object!", FutureWarning)
+            warnings.warn("Warning: 'initialise_external' will use the environment_state variable which holds the calibrator object!", DeprecationWarning)
             
             
         self._experiment = experiment

From 884016a6ccd170a9c64976a0ab000defa8e7988c Mon Sep 17 00:00:00 2001
From: fedem-p <federico.puppo@tum.de>
Date: Tue, 8 Mar 2022 10:35:18 +0100
Subject: [PATCH 05/13] Move data class; printout error for feedback; moved
 initialization of environmentstate var

---
 stytra/stimulation/__init__.py                | 16 +++++-----------
 stytra/stimulation/stimuli/generic_stimuli.py |  7 ++++++-
 2 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/stytra/stimulation/__init__.py b/stytra/stimulation/__init__.py
index 2392891e..60acbf5f 100644
--- a/stytra/stimulation/__init__.py
+++ b/stytra/stimulation/__init__.py
@@ -4,18 +4,12 @@
 import warnings
 
 from PyQt5.QtCore import pyqtSignal, QTimer, QObject
-from stytra.stimulation.stimuli import Pause, DynamicStimulus
+from stytra.stimulation.stimuli import Pause, DynamicStimulus, EnvironmentState
 from stytra.collectors.accumulators import DynamicLog, FramerateAccumulator
 from stytra.utilities import FramerateRecorder
 from lightparam.param_qt import ParametrizedQt, Param
 
 import logging
-
-@dataclass
-class EnvironmentState:
-    _calibrator = None
-    height = 0
-    width = 0
         
 
 class ProtocolRunner(QObject):
@@ -99,6 +93,8 @@ def __init__(self, experiment=None, target_dt=0, log_print=True):
         self.current_stimulus = None  # current stimulus object
         self.past_stimuli_elapsed = None  # time elapsed in previous stimuli
         self.dynamic_log = None  # dynamic log for stimuli
+        self.environment_state = EnvironmentState(self.experiment.calibrator,
+                                                  )
 
         self.update_protocol()
         self.protocol.sig_param_changed.connect(self.update_protocol)
@@ -107,9 +103,6 @@ def __init__(self, experiment=None, target_dt=0, log_print=True):
         self.log = []
         self.log_print = log_print
         self.running = False
-        
-        self.environment_state = EnvironmentState(_calibrator = self.experiment.calibrator,
-                                                  )
 
         self.framerate_rec = FramerateRecorder()
         self.framerate_acc = FramerateAccumulator(experiment=self.experiment)
@@ -125,7 +118,8 @@ def update_protocol(self):
         for stimulus in self.stimuli:
             try:
                 stimulus.initialise_external(self.experiment, self.environment_state,)
-            except TypeError:
+            except TypeError as e:
+                print("Error: {}".format(e))
                 stimulus.initialise_external(self.experiment)
                 warnings.warn("Warning: 'initialise_external' will use the environment_state variable which holds the calibrator object!", FutureWarning)
                 warnings.warn("Warning: 'initialise_external' will use the environment_state variable which holds the calibrator object!", DeprecationWarning)
diff --git a/stytra/stimulation/stimuli/generic_stimuli.py b/stytra/stimulation/stimuli/generic_stimuli.py
index 113a3450..3c67e82f 100644
--- a/stytra/stimulation/stimuli/generic_stimuli.py
+++ b/stytra/stimulation/stimuli/generic_stimuli.py
@@ -1,8 +1,13 @@
 import numpy as np
 import datetime
 import warnings
-from stytra.stimulation import EnvironmentState
+from dataclasses import dataclass
 
+@dataclass
+class EnvironmentState:
+    _calibrator = None
+    height:int = 0
+    width:int = 0
 
 class Stimulus:
     """ Abstract class for a Stimulus.

From e4ffdf085c17d35d27125d583519cac72ce44289 Mon Sep 17 00:00:00 2001
From: fedem-p <federico.puppo@tum.de>
Date: Tue, 8 Mar 2022 12:16:42 +0100
Subject: [PATCH 06/13] updated variable name; added class init

---
 stytra/stimulation/__init__.py                |  7 +++---
 stytra/stimulation/stimuli/conditional.py     | 14 +++++------
 stytra/stimulation/stimuli/external.py        |  2 +-
 stytra/stimulation/stimuli/generic_stimuli.py | 23 +++++++++++--------
 stytra/stimulation/stimuli/visual.py          | 12 +++++-----
 5 files changed, 32 insertions(+), 26 deletions(-)

diff --git a/stytra/stimulation/__init__.py b/stytra/stimulation/__init__.py
index 60acbf5f..23987695 100644
--- a/stytra/stimulation/__init__.py
+++ b/stytra/stimulation/__init__.py
@@ -93,7 +93,8 @@ def __init__(self, experiment=None, target_dt=0, log_print=True):
         self.current_stimulus = None  # current stimulus object
         self.past_stimuli_elapsed = None  # time elapsed in previous stimuli
         self.dynamic_log = None  # dynamic log for stimuli
-        self.environment_state = EnvironmentState(self.experiment.calibrator,
+        self.environment_state = EnvironmentState(calibrator = self.experiment.calibrator,
+                                                  estimator = self.experiment.estimator,
                                                   )
 
         self.update_protocol()
@@ -121,8 +122,8 @@ def update_protocol(self):
             except TypeError as e:
                 print("Error: {}".format(e))
                 stimulus.initialise_external(self.experiment)
-                warnings.warn("Warning: 'initialise_external' will use the environment_state variable which holds the calibrator object!", FutureWarning)
-                warnings.warn("Warning: 'initialise_external' will use the environment_state variable which holds the calibrator object!", DeprecationWarning)
+                warnings.warn("Warning: 'initialise_external' will use the environment_state variable which holds the calibrator and the estimator object!", FutureWarning)
+                warnings.warn("Warning: 'initialise_external' will use the environment_state variable which holds the calibrator and the estimator object!", DeprecationWarning)
 
         if self.dynamic_log is None:
             self.dynamic_log = DynamicLog(self.stimuli, experiment=self.experiment)
diff --git a/stytra/stimulation/stimuli/conditional.py b/stytra/stimulation/stimuli/conditional.py
index 8300da63..ab7ebc4a 100644
--- a/stytra/stimulation/stimuli/conditional.py
+++ b/stytra/stimulation/stimuli/conditional.py
@@ -36,9 +36,9 @@ def get_dynamic_state(self):
             state.update(self.active.get_dynamic_state())
         return state
 
-    def initialise_external(self, experiment, calibrator):
-        super().initialise_external(experiment, calibrator)
-        self.active.initialise_external(experiment, calibrator)
+    def initialise_external(self, experiment, environment_state):
+        super().initialise_external(experiment, environment_state)
+        self.active.initialise_external(experiment, environment_state)
 
     def get_state(self):
         state = super().get_state()
@@ -157,10 +157,10 @@ def get_dynamic_state(self):
             state.update(self._stim_on.get_dynamic_state())
         return state
 
-    def initialise_external(self, experiment, calibrator):
-        super().initialise_external(experiment, calibrator)
-        self._stim_on.initialise_external(experiment, calibrator)
-        self._stim_off.initialise_external(experiment, calibrator)
+    def initialise_external(self, experiment, environment_state):
+        super().initialise_external(experiment, environment_state)
+        self._stim_on.initialise_external(experiment, environment_state)
+        self._stim_off.initialise_external(experiment, environment_state)
 
     def get_state(self):
         state = super().get_state()
diff --git a/stytra/stimulation/stimuli/external.py b/stytra/stimulation/stimuli/external.py
index 9df28c6f..c926c468 100644
--- a/stytra/stimulation/stimuli/external.py
+++ b/stytra/stimulation/stimuli/external.py
@@ -42,7 +42,7 @@ def __init__(
         pulse_dur_str = str(pulse_dur_ms).zfill(3)
         self.mex = str("shock" + amp_dac + pulse_dur_str)
 
-    def initialise_external(self, experiment, calibrator):
+    def initialise_external(self, experiment, environment_state):
         """
 
         Parameters
diff --git a/stytra/stimulation/stimuli/generic_stimuli.py b/stytra/stimulation/stimuli/generic_stimuli.py
index 3c67e82f..c7006283 100644
--- a/stytra/stimulation/stimuli/generic_stimuli.py
+++ b/stytra/stimulation/stimuli/generic_stimuli.py
@@ -5,9 +5,14 @@
 
 @dataclass
 class EnvironmentState:
-    _calibrator = None
-    height:int = 0
-    width:int = 0
+    def __init__(self, calibrator = None, estimator = None, height:int = 600, width:int = 800):
+        """
+        Holds Environment variables to pass from the protocol runner to the stimulus
+        """
+        self.calibrator = calibrator
+        self.estimator = estimator
+        self.height = height
+        self.width = width
 
 class Stimulus:
     """ Abstract class for a Stimulus.
@@ -73,7 +78,7 @@ def __init__(self, duration=0.0):
         self._elapsed = 0.0  # time from the beginning of the stimulus
         self.name = "undefined"
         self._experiment = None
-        self._calibrator = None
+        self._environment_state = None
         self.real_time_start = None
         self.real_time_stop = None
 
@@ -140,9 +145,9 @@ def initialise_external(self, experiment, environment_state: EnvironmentState =
         """
         
         if isinstance(environment_state, EnvironmentState):
-            self._calibrator = environment_state._calibrator
+            self._environment_state = environment_state
         else:
-            self._calibrator = self._experiment.calibrator
+            self._environment_state = experiment
             warnings.warn("Warning: 'initialise_external' will use the environment_state variable which holds the calibrator object!", FutureWarning)
             warnings.warn("Warning: 'initialise_external' will use the environment_state variable which holds the calibrator object!", DeprecationWarning)
             
@@ -308,10 +313,10 @@ def update(self):
             s.update()
             s._elapsed = self._elapsed
 
-    def initialise_external(self, experiment, calibrator):
-        super().initialise_external(experiment, calibrator)
+    def initialise_external(self, experiment, environment_state):
+        super().initialise_external(experiment, environment_state)
         for s in self._stim_list:
-            s.initialise_external(experiment, calibrator)
+            s.initialise_external(experiment, environment_state)
 
     @property
     def dynamic_parameter_names(self):
diff --git a/stytra/stimulation/stimuli/visual.py b/stytra/stimulation/stimuli/visual.py
index d4899d31..54538944 100644
--- a/stytra/stimulation/stimuli/visual.py
+++ b/stytra/stimulation/stimuli/visual.py
@@ -396,8 +396,8 @@ def __init__(self, *args, background, background_name=None, **kwargs):
                 self.background_name = "array {}x{}".format(*self._background.shape)
         self._qbackground = None
 
-    def initialise_external(self, experiment, calibrator):
-        super().initialise_external(experiment, calibrator)
+    def initialise_external(self, experiment, environment_state):
+        super().initialise_external(experiment, environment_state)
 
         # Get background image from folder:
         if isinstance(self._background, str):
@@ -483,8 +483,8 @@ def create_pattern(self):
                 + (1 - w[:, None]) * np.array(self.color_2)[None, :]
             ).astype(np.uint8)
 
-    def initialise_external(self, experiment, calibrator):
-        super().initialise_external(experiment, calibrator)
+    def initialise_external(self, experiment, environment_state):
+        super().initialise_external(experiment, environment_state)
         self.create_pattern()
         # Get background image from folder:
         self._qbackground = qimage2ndarray.array2qimage(self._pattern[None, :, :])
@@ -751,8 +751,8 @@ def create_pattern(self, side_len=500):
         self._pattern = W * self.color_1 + (1 - W) * self.color_2
         self._qbackground = qimage2ndarray.array2qimage(self._pattern)
 
-    def initialise_external(self, experiment, calibrator):
-        super().initialise_external(experiment, calibrator)
+    def initialise_external(self, experiment, environment_state):
+        super().initialise_external(experiment, environment_state)
         self.create_pattern()
 
     def draw_block(self, p, point, w, h):

From 03a99bcd752ba27a8371c164a5a81e3334a00bcf Mon Sep 17 00:00:00 2001
From: fedem-p <federico.puppo@tum.de>
Date: Tue, 8 Mar 2022 13:44:03 +0100
Subject: [PATCH 07/13] Updated calibrator variable

---
 stytra/stimulation/stimuli/conditional.py   |  6 +++---
 stytra/stimulation/stimuli/kinematograms.py |  4 ++--
 stytra/stimulation/stimuli/visual.py        | 16 ++++++++--------
 3 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/stytra/stimulation/stimuli/conditional.py b/stytra/stimulation/stimuli/conditional.py
index ab7ebc4a..227e6f7c 100644
--- a/stytra/stimulation/stimuli/conditional.py
+++ b/stytra/stimulation/stimuli/conditional.py
@@ -271,7 +271,7 @@ def __init__(self, stimulus, *args, centering_stimulus=None, margin=45, **kwargs
 
     def check_condition_on(self):
         y, x, theta = self._experiment.estimator.get_position()
-        scale = self._calibrator.mm_px ** 2 
+        scale = self._environment_state.calibrator.mm_px ** 2 
         return (
             x > 0 and ((x - self.xc) ** 2 + (y - self.yc) ** 2) <= self.margin / scale
         )
@@ -324,14 +324,14 @@ def __init__(
 
     def check_condition_on(self):
         y, x, theta = self._experiment.estimator.get_position()
-        scale = self._calibrator.mm_px ** 2
+        scale = self._environment_state.calibrator.mm_px ** 2
         return (not np.isnan(x)) and (
             (x - self.xc) ** 2 + (y - self.yc) ** 2 <= self.margin_in / scale
         )
 
     def check_condition_off(self):
         y, x, theta = self._experiment.estimator.get_position()
-        scale = self._calibrator.mm_px ** 2
+        scale = self._environment_state.calibrator.mm_px ** 2
         return np.isnan(x) or (
             (x - self.xc) ** 2 + (y - self.yc) ** 2 > self.margin_out / scale
         )
diff --git a/stytra/stimulation/stimuli/kinematograms.py b/stytra/stimulation/stimuli/kinematograms.py
index 86556761..fe19abfc 100644
--- a/stytra/stimulation/stimuli/kinematograms.py
+++ b/stytra/stimulation/stimuli/kinematograms.py
@@ -75,8 +75,8 @@ def get_dimensions(self):
         -------
         number of dots to display and the displacement amount in pixel coordinates
         """
-        if self._calibrator is not None:
-            mm_px = self._calibrator.mm_px  
+        if self._environment_state.calibrator is not None:
+            mm_px = self._environment_state.calibrator.mm_px  
         else:
             mm_px = 1
 
diff --git a/stytra/stimulation/stimuli/visual.py b/stytra/stimulation/stimuli/visual.py
index 54538944..035d42fc 100644
--- a/stytra/stimulation/stimuli/visual.py
+++ b/stytra/stimulation/stimuli/visual.py
@@ -313,8 +313,8 @@ def get_tile_ranges(self, imw, imh, w, h, tr: QTransform):
         return range(x_start, x_end + 1), range(y_start, y_end + 1)
 
     def paint(self, p, w, h):
-        if self._calibrator is not None: 
-            mm_px = self._calibrator.mm_px
+        if self._environment_state.calibrator is not None: 
+            mm_px = self._environment_state.calibrator.mm_px
         else:
             mm_px = 1
 
@@ -468,7 +468,7 @@ def __init__(
     def create_pattern(self):
         l = max(
             2,
-            int(self.grating_period / (max(self._calibrator.mm_px, 0.0001))),
+            int(self.grating_period / (max(self._environment_state.calibrator.mm_px, 0.0001))),
         )
         if self.wave_shape == "square":
             self._pattern = np.ones((l, 3), np.uint8) * self.color_1
@@ -532,7 +532,7 @@ def get_unit_dims(self, w, h):
         #TODO what does this thing define?
         """
         return (
-            int(self.grating_period / (max(self._calibrator.mm_px, 0.0001))),
+            int(self.grating_period / (max(self._environment_state.calibrator.mm_px, 0.0001))),
             self.barheight,
         )
 
@@ -547,7 +547,7 @@ def draw_block(self, p, point, w, h):
             point.y(),
             int(
                 self.grating_period
-                / (2 * max(self._calibrator.mm_px, 0.0001))
+                / (2 * max(self._environment_state.calibrator.mm_px, 0.0001))
             ),
             self.barheight,
         )
@@ -654,7 +654,7 @@ def update(self):
 
     def paint(self, p, w, h):
         x, y = (
-            (np.arange(d) - d / 2) * self._calibrator.mm_px for d in (w, h)
+            (np.arange(d) - d / 2) * self._environment_state.calibrator.mm_px for d in (w, h)
         )
         self.image = np.round(
             np.sin(
@@ -933,8 +933,8 @@ def __init__(
     def paint(self, p, w, h):
         super().paint(p, w, h)
 
-        if self._calibrator is not None:
-            mm_px = self._calibrator.mm_px
+        if self._environment_state.calibrator is not None:
+            mm_px = self._environment_state.calibrator.mm_px
         else:
             mm_px = 1
 

From bad80377f99fa91af7391e53c672cd8ba09567a4 Mon Sep 17 00:00:00 2001
From: fedem-p <federico.puppo@tum.de>
Date: Tue, 8 Mar 2022 14:07:35 +0100
Subject: [PATCH 08/13] added estimator log variable;

---
 stytra/stimulation/__init__.py            | 10 +++++----
 stytra/stimulation/stimuli/closed_loop.py | 27 +++++++++++------------
 stytra/stimulation/stimuli/conditional.py |  8 +++----
 3 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/stytra/stimulation/__init__.py b/stytra/stimulation/__init__.py
index 23987695..c572e4ad 100644
--- a/stytra/stimulation/__init__.py
+++ b/stytra/stimulation/__init__.py
@@ -93,10 +93,9 @@ def __init__(self, experiment=None, target_dt=0, log_print=True):
         self.current_stimulus = None  # current stimulus object
         self.past_stimuli_elapsed = None  # time elapsed in previous stimuli
         self.dynamic_log = None  # dynamic log for stimuli
-        self.environment_state = EnvironmentState(calibrator = self.experiment.calibrator,
-                                                  estimator = self.experiment.estimator,
-                                                  )
-
+        
+        self.environment_state = EnvironmentState(calibrator = self.experiment.calibrator,)
+            
         self.update_protocol()
         self.protocol.sig_param_changed.connect(self.update_protocol)
 
@@ -114,6 +113,9 @@ def update_protocol(self):
         self.stimuli = self.protocol._get_stimulus_list()
 
         self.current_stimulus = self.stimuli[0]
+        
+        if hasattr(self.experiment, 'estimator'):
+            self.environment_state.estimator = self.experiment.estimator
 
         # pass experiment to stimuli for calibrator and asset folders:
         for stimulus in self.stimuli:
diff --git a/stytra/stimulation/stimuli/closed_loop.py b/stytra/stimulation/stimuli/closed_loop.py
index 73b5e25e..a431f5f4 100644
--- a/stytra/stimulation/stimuli/closed_loop.py
+++ b/stytra/stimulation/stimuli/closed_loop.py
@@ -69,7 +69,7 @@ def get_fish_vel(self):
         """ Function that update estimated fish velocty. Change to add lag or
         shunting.
         """
-        self.fish_vel = self._experiment.estimator.get_velocity()
+        self.fish_vel = self._environment_state.estimator.get_velocity()
 
     def bout_started(self):
         """ Function called on bout start.
@@ -90,7 +90,7 @@ def update(self):
                 self._experiment.logger.info(
                     "Experiment aborted! {} seconds without bouts".format(
                         self._elapsed - self.prev_bout_t
-                    )
+                    ) #! TOFIX: Remove
                 )
                 self.abort_experiment()
 
@@ -130,7 +130,7 @@ def calculate_final_vel(self):
         self.vel = self.base_vel - self.fish_vel * int(self.fish_swimming)
 
     def abort_experiment(self):
-        self._experiment.protocol_runner.stop()
+        self._experiment.protocol_runner.stop() #! TOFIX: Remove?
 
 
 class CalibratingClosedLoop1D(Basic_CL_1D):
@@ -177,7 +177,7 @@ def __init__(
 
     def bout_started(self):
         super().bout_started()
-        self.est_gain = self._experiment.estimator.base_gain
+        self.est_gain = self._environment_state.estimator.base_gain
 
     def bout_occurring(self):
         self.bout_vig.append(self.fish_vel / self.est_gain)
@@ -196,7 +196,7 @@ def bout_ended(self):
                 self.median_calib = self.median_vig * self.est_gain
                 self.est_gain = self.target_avg_fish_vel / self.median_vig
 
-                self._experiment.estimator.base_gain = self.est_gain
+                self._environment_state.estimator.base_gain = self.est_gain
 
         self.bout_vel = []
 
@@ -212,14 +212,14 @@ def stop(self):
                 "Experiment aborted! N bouts: {}; gain: {}".format(
                     len(self.bouts_vig_list), self.est_gain
                 )
-            )
+            ) #! TOFIX: Remove
 
         if len(self.bouts_vig_list) > self.calibrate_after:
             self._experiment.logger.info(
                 "Calibrated! Calculated gain  {} with {} bouts".format(
                     self.est_gain, len(self.bouts_vig_list)
                 )
-            )
+            ) #! TOFIX: Remove
 
 
 class GainChangerStimulus(Stimulus):
@@ -246,7 +246,7 @@ def __init__(self, newgain=1):
         self.newgain = newgain
 
     def start(self):
-        self._experiment.estimator.base_gain = self.newgain
+        self._environment_state.estimator.base_gain = self.newgain
 
 
 class GainLagClosedLoop1D(Basic_CL_1D):
@@ -277,8 +277,7 @@ def get_fish_vel(self):
         shunting.
         """
         super(GainLagClosedLoop1D, self).get_fish_vel()
-        self.lag_vel = self._experiment.estimator.get_velocity(self.lag)
-
+        self.lag_vel = self._environment_state.estimator.get_velocity(self.lag)
     def calculate_final_vel(self):
         subtract_to_base = self.gain * self.lag_vel
 
@@ -329,7 +328,7 @@ def bout_started(self):
             # print("set: {} gain and {} lag".format(self.gain, self.lag))
 
         # refresh lag if it was changed:
-        self.lag_vel = self._experiment.estimator.get_velocity(self.lag)
+        self.lag_vel = self._environment_state.estimator.get_velocity(self.lag)
 
 
 class PerpendicularMotion(BackgroundStimulus, InterpolatedStimulus):
@@ -338,7 +337,7 @@ class PerpendicularMotion(BackgroundStimulus, InterpolatedStimulus):
     """
 
     def update(self):
-        y, x, theta = self._experiment.estimator.get_position()
+        y, x, theta = self._environment_state.estimator.get_position()
         if np.isfinite(theta):
             self.theta = theta
         super().update()
@@ -352,7 +351,7 @@ def __init__(self, *args, **kwargs):
 
     def update(self):
         if self.is_tracking:
-            y, x, theta = self._experiment.estimator.get_position()
+            y, x, theta = self._environment_state.estimator.get_position()
             if np.isfinite(theta):
                 self.x = x
                 self.y = y
@@ -362,7 +361,7 @@ def update(self):
 
 class FishRelativeStimulus(BackgroundStimulus):
     def get_transform(self, w, h, x, y):
-        y_fish, x_fish, theta_fish = self._experiment.estimator.get_position()
+        y_fish, x_fish, theta_fish = self._environment_state.estimator.get_position()
         if np.isnan(y_fish):
             return super().get_transform(w, h, x, y)
         rot_fish = (theta_fish - np.pi / 2) * 180 / np.pi
diff --git a/stytra/stimulation/stimuli/conditional.py b/stytra/stimulation/stimuli/conditional.py
index 227e6f7c..831c6c13 100644
--- a/stytra/stimulation/stimuli/conditional.py
+++ b/stytra/stimulation/stimuli/conditional.py
@@ -50,7 +50,7 @@ def start(self):
         self.active.start()
 
     def check_condition(self):
-        y, x, theta = self._experiment.estimator.get_position()
+        y, x, theta = self._environment_state.estimator.get_position()
         return not np.isnan(y)
 
     def update(self):
@@ -270,7 +270,7 @@ def __init__(self, stimulus, *args, centering_stimulus=None, margin=45, **kwargs
         self.yc = 240
 
     def check_condition_on(self):
-        y, x, theta = self._experiment.estimator.get_position()
+        y, x, theta = self._environment_state.estimator.get_position()
         scale = self._environment_state.calibrator.mm_px ** 2 
         return (
             x > 0 and ((x - self.xc) ** 2 + (y - self.yc) ** 2) <= self.margin / scale
@@ -323,14 +323,14 @@ def __init__(
         self.yc = 240
 
     def check_condition_on(self):
-        y, x, theta = self._experiment.estimator.get_position()
+        y, x, theta = self._environment_state.estimator.get_position()
         scale = self._environment_state.calibrator.mm_px ** 2
         return (not np.isnan(x)) and (
             (x - self.xc) ** 2 + (y - self.yc) ** 2 <= self.margin_in / scale
         )
 
     def check_condition_off(self):
-        y, x, theta = self._experiment.estimator.get_position()
+        y, x, theta = self._environment_state.estimator.get_position()
         scale = self._environment_state.calibrator.mm_px ** 2
         return np.isnan(x) or (
             (x - self.xc) ** 2 + (y - self.yc) ** 2 > self.margin_out / scale

From 6db29ecf17fff133a7c59e99816d6df2c2dbce1f Mon Sep 17 00:00:00 2001
From: fedem-p <federico.puppo@tum.de>
Date: Tue, 8 Mar 2022 15:03:58 +0100
Subject: [PATCH 09/13] removed useless comments

---
 stytra/stimulation/stimuli/closed_loop.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/stytra/stimulation/stimuli/closed_loop.py b/stytra/stimulation/stimuli/closed_loop.py
index a431f5f4..9626957a 100644
--- a/stytra/stimulation/stimuli/closed_loop.py
+++ b/stytra/stimulation/stimuli/closed_loop.py
@@ -90,7 +90,7 @@ def update(self):
                 self._experiment.logger.info(
                     "Experiment aborted! {} seconds without bouts".format(
                         self._elapsed - self.prev_bout_t
-                    ) #! TOFIX: Remove
+                    )
                 )
                 self.abort_experiment()
 
@@ -130,7 +130,7 @@ def calculate_final_vel(self):
         self.vel = self.base_vel - self.fish_vel * int(self.fish_swimming)
 
     def abort_experiment(self):
-        self._experiment.protocol_runner.stop() #! TOFIX: Remove?
+        self._experiment.protocol_runner.stop()
 
 
 class CalibratingClosedLoop1D(Basic_CL_1D):
@@ -212,14 +212,14 @@ def stop(self):
                 "Experiment aborted! N bouts: {}; gain: {}".format(
                     len(self.bouts_vig_list), self.est_gain
                 )
-            ) #! TOFIX: Remove
+            )
 
         if len(self.bouts_vig_list) > self.calibrate_after:
             self._experiment.logger.info(
                 "Calibrated! Calculated gain  {} with {} bouts".format(
                     self.est_gain, len(self.bouts_vig_list)
                 )
-            ) #! TOFIX: Remove
+            )
 
 
 class GainChangerStimulus(Stimulus):

From 7d9cf144eaef05915024dc603ec2e4f818281b81 Mon Sep 17 00:00:00 2001
From: fedem-p <federico.puppo@tum.de>
Date: Wed, 9 Mar 2022 10:36:34 +0100
Subject: [PATCH 10/13] change msg for deprecation

---
 stytra/stimulation/__init__.py                | 8 +++++---
 stytra/stimulation/stimuli/generic_stimuli.py | 5 +++--
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/stytra/stimulation/__init__.py b/stytra/stimulation/__init__.py
index c572e4ad..08892bee 100644
--- a/stytra/stimulation/__init__.py
+++ b/stytra/stimulation/__init__.py
@@ -124,9 +124,11 @@ def update_protocol(self):
             except TypeError as e:
                 print("Error: {}".format(e))
                 stimulus.initialise_external(self.experiment)
-                warnings.warn("Warning: 'initialise_external' will use the environment_state variable which holds the calibrator and the estimator object!", FutureWarning)
-                warnings.warn("Warning: 'initialise_external' will use the environment_state variable which holds the calibrator and the estimator object!", DeprecationWarning)
-
+                msg = "Warning: self._experiment is deprecated use self._environment_state instead, self._experiment will be unavailable from version 1.0!"
+                warnings.warn(msg, FutureWarning)
+                warnings.warn(msg, DeprecationWarning)
+                
+                
         if self.dynamic_log is None:
             self.dynamic_log = DynamicLog(self.stimuli, experiment=self.experiment)
         else:
diff --git a/stytra/stimulation/stimuli/generic_stimuli.py b/stytra/stimulation/stimuli/generic_stimuli.py
index c7006283..f52b6f1b 100644
--- a/stytra/stimulation/stimuli/generic_stimuli.py
+++ b/stytra/stimulation/stimuli/generic_stimuli.py
@@ -148,8 +148,9 @@ def initialise_external(self, experiment, environment_state: EnvironmentState =
             self._environment_state = environment_state
         else:
             self._environment_state = experiment
-            warnings.warn("Warning: 'initialise_external' will use the environment_state variable which holds the calibrator object!", FutureWarning)
-            warnings.warn("Warning: 'initialise_external' will use the environment_state variable which holds the calibrator object!", DeprecationWarning)
+            msg = "Warning: self._experiment is deprecated use self._environment_state instead, self._experiment will be unavailable from version 1.0!"
+            warnings.warn(msg, FutureWarning)
+            warnings.warn(msg, DeprecationWarning)
             
             
         self._experiment = experiment

From f1b545bc5e7c1d3b75291e90fbac9bf694532eef Mon Sep 17 00:00:00 2001
From: fedem-p <federico.puppo@tum.de>
Date: Thu, 10 Mar 2022 15:52:19 +0100
Subject: [PATCH 11/13] added params to env_state var; asset_dir, arduino board
 and logger

---
 stytra/stimulation/__init__.py                |  9 ++++++++-
 stytra/stimulation/stimuli/generic_stimuli.py | 11 ++++++++++-
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/stytra/stimulation/__init__.py b/stytra/stimulation/__init__.py
index 08892bee..0e5ec93c 100644
--- a/stytra/stimulation/__init__.py
+++ b/stytra/stimulation/__init__.py
@@ -114,9 +114,16 @@ def update_protocol(self):
 
         self.current_stimulus = self.stimuli[0]
         
+        #populate environment_state class
         if hasattr(self.experiment, 'estimator'):
             self.environment_state.estimator = self.experiment.estimator
-
+        if hasattr(self.experiment, 'arduino_board'):
+            self.environment_state.arduino_board = self.experiment.arduino_board
+        if hasattr(self.experiment, 'asset_dir'):
+            self.environment_state.asset_dir = self.experiment.asset_dir
+        if hasattr(self.experiment, 'logger'):
+            self.environment_state.logger = self.experiment.logger
+            
         # pass experiment to stimuli for calibrator and asset folders:
         for stimulus in self.stimuli:
             try:
diff --git a/stytra/stimulation/stimuli/generic_stimuli.py b/stytra/stimulation/stimuli/generic_stimuli.py
index f52b6f1b..376e5df7 100644
--- a/stytra/stimulation/stimuli/generic_stimuli.py
+++ b/stytra/stimulation/stimuli/generic_stimuli.py
@@ -5,12 +5,21 @@
 
 @dataclass
 class EnvironmentState:
-    def __init__(self, calibrator = None, estimator = None, height:int = 600, width:int = 800):
+    def __init__(self, calibrator = None, 
+                 estimator = None, 
+                 arduino_board = None, 
+                 asset_dir = None,
+                 logger = None,
+                 height:int = 600, 
+                 width:int = 800):
         """
         Holds Environment variables to pass from the protocol runner to the stimulus
         """
         self.calibrator = calibrator
         self.estimator = estimator
+        self.arduino_board = arduino_board
+        self.asset_dir = asset_dir
+        self.logger = logger
         self.height = height
         self.width = width
 

From 1f3c8cb6ba3aff874f1b2bfaf5686ecaf5386489 Mon Sep 17 00:00:00 2001
From: fedem-p <federico.puppo@tum.de>
Date: Thu, 10 Mar 2022 15:53:19 +0100
Subject: [PATCH 12/13] updated access to new params; access env state var

---
 stytra/stimulation/stimuli/arduino.py     | 6 +++---
 stytra/stimulation/stimuli/closed_loop.py | 6 +++---
 stytra/stimulation/stimuli/visual.py      | 4 ++--
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/stytra/stimulation/stimuli/arduino.py b/stytra/stimulation/stimuli/arduino.py
index 7700383c..161054ca 100644
--- a/stytra/stimulation/stimuli/arduino.py
+++ b/stytra/stimulation/stimuli/arduino.py
@@ -22,7 +22,7 @@ def __init__(self, pin_values_dict, *args, **kwargs):
 
     def start(self):
         super().start()
-        self._experiment.arduino_board.write_multiple(self.pin_values)
+        self._environment_state.arduino_board.write_multiple(self.pin_values)
 
 
 class ContinuousWriteArduinoPin(InterpolatedStimulus):
@@ -45,10 +45,10 @@ def __init__(self, pin, *args, **kwargs):
 
     def update(self):
         super().update()
-        self._experiment.arduino_board.write(self.pin, self.pin_value)
+        self._environment_state.arduino_board.write(self.pin, self.pin_value)
 
     def stop(self):
         super().update()
-        self._experiment.arduino_board.write(self.pin, 0)
+        self._environment_state.arduino_board.write(self.pin, 0)
 
 
diff --git a/stytra/stimulation/stimuli/closed_loop.py b/stytra/stimulation/stimuli/closed_loop.py
index 9626957a..9dff11b4 100644
--- a/stytra/stimulation/stimuli/closed_loop.py
+++ b/stytra/stimulation/stimuli/closed_loop.py
@@ -87,7 +87,7 @@ def bout_ended(self):
     def update(self):
         if self.max_interbout_time is not None:
             if self._elapsed - self.prev_bout_t > self.max_interbout_time:
-                self._experiment.logger.info(
+                self._environment_state.logger.info(
                     "Experiment aborted! {} seconds without bouts".format(
                         self._elapsed - self.prev_bout_t
                     )
@@ -208,14 +208,14 @@ def stop(self):
         ):
             self.abort_experiment()
 
-            self._experiment.logger.info(
+            self._environment_state.logger.info(
                 "Experiment aborted! N bouts: {}; gain: {}".format(
                     len(self.bouts_vig_list), self.est_gain
                 )
             )
 
         if len(self.bouts_vig_list) > self.calibrate_after:
-            self._experiment.logger.info(
+            self._environment_state.logger.info(
                 "Calibrated! Calculated gain  {} with {} bouts".format(
                     self.est_gain, len(self.bouts_vig_list)
                 )
diff --git a/stytra/stimulation/stimuli/visual.py b/stytra/stimulation/stimuli/visual.py
index 035d42fc..311ee29d 100644
--- a/stytra/stimulation/stimuli/visual.py
+++ b/stytra/stimulation/stimuli/visual.py
@@ -206,7 +206,7 @@ def __init__(self, *args, video_path, framerate=None, duration=None, **kwargs):
 
     def initialise_external(self, *args, **kwargs):
         super().initialise_external(*args, **kwargs)
-        self._video_seq = pims.Video(self._experiment.asset_dir + "/" + self.video_path)
+        self._video_seq = pims.Video(self._environment_state.asset_dir + "/" + self.video_path)
 
         self._current_frame = self._video_seq.get_frame(self.i_frame)
         try:
@@ -403,7 +403,7 @@ def initialise_external(self, experiment, environment_state):
         if isinstance(self._background, str):
             self._qbackground = qimage2ndarray.array2qimage(
                 existing_file_background(
-                    self._experiment.asset_dir + "/" + self._background
+                    self._environment_state.asset_dir + "/" + self._background
                 )
             )
         elif isinstance(self._background, Path):

From 27e377a68c10e26922c2384ec375900f8f374fd4 Mon Sep 17 00:00:00 2001
From: fedem-p <federico.puppo@tum.de>
Date: Wed, 16 Mar 2022 13:48:43 +0100
Subject: [PATCH 13/13] Included trigger inside environment state variable

---
 stytra/stimulation/__init__.py                | 2 ++
 stytra/stimulation/stimuli/generic_stimuli.py | 4 +++-
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/stytra/stimulation/__init__.py b/stytra/stimulation/__init__.py
index 0e5ec93c..2a7318d9 100644
--- a/stytra/stimulation/__init__.py
+++ b/stytra/stimulation/__init__.py
@@ -123,6 +123,8 @@ def update_protocol(self):
             self.environment_state.asset_dir = self.experiment.asset_dir
         if hasattr(self.experiment, 'logger'):
             self.environment_state.logger = self.experiment.logger
+        if hasattr(self.experiment, 'trigger'):
+            self.environment_state.trigger = self.experiment.trigger
             
         # pass experiment to stimuli for calibrator and asset folders:
         for stimulus in self.stimuli:
diff --git a/stytra/stimulation/stimuli/generic_stimuli.py b/stytra/stimulation/stimuli/generic_stimuli.py
index 376e5df7..9857bcea 100644
--- a/stytra/stimulation/stimuli/generic_stimuli.py
+++ b/stytra/stimulation/stimuli/generic_stimuli.py
@@ -10,6 +10,7 @@ def __init__(self, calibrator = None,
                  arduino_board = None, 
                  asset_dir = None,
                  logger = None,
+                 trigger = None,
                  height:int = 600, 
                  width:int = 800):
         """
@@ -18,6 +19,7 @@ def __init__(self, calibrator = None,
         self.calibrator = calibrator
         self.estimator = estimator
         self.arduino_board = arduino_board
+        self.trigger = trigger
         self.asset_dir = asset_dir
         self.logger = logger
         self.height = height
@@ -285,7 +287,7 @@ def start(self):
 
     def update(self):
         # If trigger is set, make it end:
-        if self._experiment.trigger.start_event.is_set():
+        if self._environment_state.trigger.start_event.is_set():
             self.duration = self._elapsed