From d1df11dcf2727ca455589ea11a6150218771c8c8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olivier=20H=C3=A9riveaux?= <olivier.heriveaux@ledger.fr>
Date: Fri, 16 Aug 2024 15:18:25 +0200
Subject: [PATCH] Moved autofocus helper to instruments

---
 laserstudio/instruments/instruments.py       | 19 +++++++++++++++++--
 laserstudio/laserstudio.py                   | 12 +++++++++---
 laserstudio/restserver/server.py             |  7 ++-----
 laserstudio/widgets/toolbars/focustoolbar.py |  6 +++---
 4 files changed, 31 insertions(+), 13 deletions(-)

diff --git a/laserstudio/instruments/instruments.py b/laserstudio/instruments/instruments.py
index 28544dc..138376a 100644
--- a/laserstudio/instruments/instruments.py
+++ b/laserstudio/instruments/instruments.py
@@ -1,3 +1,6 @@
+from pystages import Autofocus, Vector
+from typing import Optional, cast, Any
+import logging
 from .stage import StageInstrument
 from .list_serials import DeviceSearchError
 from .camera import CameraInstrument
@@ -8,8 +11,6 @@
 from .laserdriver import LaserDriverInstrument, LaserDriver  # type: ignore
 from .pdm import PDMInstrument
 from .probe import ProbeInstrument
-from typing import Optional, cast, Any
-import logging
 
 
 class Instruments:
@@ -91,8 +92,22 @@ def __init__(self, config: dict):
                     continue
                 self.probes.append(ProbeInstrument(config=probe_config))
 
+        # Autofocus helper: stores registered position in order to do automatic camera
+        # focusing. This can be considered as an abstract instrument.
+        self.autofocus_helper = Autofocus()
+
     def go_next(self) -> dict[str, Any]:
         results = []
         for laser in self.lasers:
             results.append(laser.go_next())
         return {"lasers": results}
+
+    def autofocus(self):
+        if self.stage is None:
+            return
+        if len(self.autofocus_helper.registered_points) < 3:
+            return
+        pos = self.stage.position
+        z = self.autofocus_helper.focus(pos.x, pos.y)
+        assert abs(z - pos.z) < 500
+        self.stage.move_to(Vector(pos.x, pos.y, z), wait=True)
diff --git a/laserstudio/laserstudio.py b/laserstudio/laserstudio.py
index 848a108..7ab2083 100644
--- a/laserstudio/laserstudio.py
+++ b/laserstudio/laserstudio.py
@@ -26,7 +26,7 @@
     PDMToolbar,
     LaserDriverToolbar,
     CameraNITToolBar,
-    FocusToolbar
+    FocusToolbar,
 )
 import yaml
 from .restserver.server import RestProxy
@@ -93,8 +93,14 @@ def __init__(self, config: Optional[dict]):
             self.addToolBar(Qt.ToolBarArea.BottomToolBarArea, toolbar)
 
         # Toolbar: Focusing
-        if (self.instruments.stage is not None) and (self.instruments.camera is not None):
-            toolbar = FocusToolbar(self.instruments.stage, self.instruments.camera)
+        if (self.instruments.stage is not None) and (
+            self.instruments.camera is not None
+        ):
+            toolbar = FocusToolbar(
+                self.instruments.stage,
+                self.instruments.camera,
+                self.instruments.autofocus_helper,
+            )
             self.addToolBar(toolbar)
 
         # Toolbar: Scanning zone definition and usage
diff --git a/laserstudio/restserver/server.py b/laserstudio/restserver/server.py
index f7a3016..f3c8000 100644
--- a/laserstudio/restserver/server.py
+++ b/laserstudio/restserver/server.py
@@ -41,11 +41,8 @@ def handle_go_next(self):
 
     @pyqtSlot(result="QVariant")
     def handle_autofocus(self):
-        return QVariant({"error": "Not implemented"})
-        try:
-            return QVariant(self.laser_studio.autofocus(no_warning=True))
-        except RuntimeError as e:
-            return {"error": str(e)}
+        self.laser_studio.instruments.autofocus()
+        return QVariant({})
 
     @pyqtSlot(result="QVariant")
     def handle_go_to_memory_point(self, index: int):
diff --git a/laserstudio/widgets/toolbars/focustoolbar.py b/laserstudio/widgets/toolbars/focustoolbar.py
index ecae01b..7821a60 100644
--- a/laserstudio/widgets/toolbars/focustoolbar.py
+++ b/laserstudio/widgets/toolbars/focustoolbar.py
@@ -13,9 +13,9 @@
 class FocusToolbar(QToolBar):
     """Toolbar for focus registration and autofocus."""
 
-    def __init__(self, stage, camera: CameraNITInstrument):
+    def __init__(self, stage, camera: CameraNITInstrument, autofocus_helper: Autofocus):
         """
-        :param autofocus: Stores the registered points and calculates focus on demand.
+        :param autofocus_helper: Stores the registered points and calculates focus on demand.
         """
         super().__init__("Focus")
         self.setObjectName("toolbar-focus")  # For settings save and restore
@@ -26,7 +26,7 @@ def __init__(self, stage, camera: CameraNITInstrument):
         # This is used to prevent launching two search threads at the same time.
         self.focus_thread: FocusThread | None = None
 
-        self.autofocus_helper = Autofocus()
+        self.autofocus_helper = autofocus_helper
         self.stage = stage
         self.camera = camera