From 3b440cd1fa18d696517345b9e3a957f0a6e078d0 Mon Sep 17 00:00:00 2001
From: appleb_m <martin.appleby@psi.ch>
Date: Tue, 24 Mar 2026 11:13:09 +0100
Subject: [PATCH] aerotech models

---
 src/aare/common/aerotech_models.py | 155 +++++++++++++++++++++++++++++
 1 file changed, 155 insertions(+)
 create mode 100644 src/aare/common/aerotech_models.py

diff --git a/src/aare/common/aerotech_models.py b/src/aare/common/aerotech_models.py
new file mode 100644
index 00000000..786ea1c3
--- /dev/null
+++ b/src/aare/common/aerotech_models.py
@@ -0,0 +1,155 @@
+from enum import Enum
+from typing import Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from aare.common.rotation_scan import RotationScanRequest
+
+
+class TaskEnum(Enum):
+    TASK_0 = 0
+    TASK_1 = 1
+    TASK_2 = 2
+    TASK_3 = 3
+    TASK_4 = 4
+    TASK_5 = 5
+    TASK_6 = 6
+    TASK_7 = 7
+    TASK_8 = 8
+    TASK_9 = 9
+
+
+class AxisEnum(Enum):
+    X = "x"
+    Y = "y"
+    Z = "z"
+    OMEGA = "u"
+
+
+class AerotechRunEnum(Enum):
+    STOP = 0
+    START = 1
+    RUN = 2
+    LOAD = 3
+    PAUSE = 4
+    RESET = 5
+
+
+class VariableTypeEnum(Enum):
+    INT = 0
+    REAL = 1
+    STRING = 2
+
+
+class AerotechAxisStatus(BaseModel):
+    enabled: bool
+    fault: int
+    homed: bool
+    is_fault: bool
+    moving: bool
+    position: float
+    status: int
+    velocity: float
+
+
+class AerotechStatus(BaseModel):
+    state: str
+    x: Optional[AerotechAxisStatus] = None
+    y: Optional[AerotechAxisStatus] = None
+    z: Optional[AerotechAxisStatus] = None
+    u: Optional[AerotechAxisStatus] = None
+
+    model_config = ConfigDict(extra="allow")
+
+    def __str__(self) -> str:
+        return self.to_pretty_string()
+
+    def to_pretty_string(self) -> str:
+        def axis_line(name: str, axis: AerotechAxisStatus | None) -> str:
+            if axis is None:
+                return f"{name.upper()}: unavailable"
+            return (
+                f"{name.upper():>2} | pos={axis.position:>12.6f} | "
+                f"homed={axis.homed!s:<5} | moving={axis.moving!s:<5} | "
+                f"enabled={axis.enabled!s:<5} | fault={axis.fault} | "
+                f"faulted={axis.is_fault!s:<5} | vel={axis.velocity:>10.6f}"
+            )
+
+        return "\n".join(
+            [
+                f"STATE: {self.state}",
+                axis_line("x", self.x),
+                axis_line("y", self.y),
+                axis_line("z", self.z),
+                axis_line("u", self.u),
+            ]
+        )
+
+    def to_compact_string(self) -> str:
+        axes = []
+        for name in ("x", "y", "z", "u"):
+            axis = getattr(self, name)
+            if axis is not None:
+                axes.append(
+                    f"{name}={axis.position:.4f} "
+                    f"({'H' if axis.homed else 'NH'}, {'M' if axis.moving else '-'})"
+                )
+        return f"state={self.state} | " + " | ".join(axes)
+
+    def to_colored_string(self) -> str:
+        # ANSI colors for terminal use
+        RESET = "\033[0m"
+        BOLD = "\033[1m"
+        CYAN = "\033[36m"
+        GREEN = "\033[32m"
+        YELLOW = "\033[33m"
+        RED = "\033[31m"
+
+        def color_bool(value: bool) -> str:
+            return f"{GREEN}True{RESET}" if value else f"{RED}False{RESET}"
+
+        def axis_line(name: str, axis: AerotechAxisStatus | None) -> str:
+            if axis is None:
+                return f"{YELLOW}{name.upper()}: unavailable{RESET}"
+            return (
+                f"{BOLD}{name.upper()}{RESET} | "
+                f"pos={CYAN}{axis.position:>12.6f}{RESET} | "
+                f"homed={color_bool(axis.homed)} | "
+                f"moving={color_bool(axis.moving)} | "
+                f"enabled={color_bool(axis.enabled)} | "
+                f"fault={axis.fault} | "
+                f"faulted={color_bool(axis.is_fault)} | "
+                f"vel={axis.velocity:>10.6f}"
+            )
+
+        return "\n".join(
+            [
+                f"{BOLD}STATE:{RESET} {CYAN}{self.state}{RESET}",
+                axis_line("x", self.x),
+                axis_line("y", self.y),
+                axis_line("z", self.z),
+                axis_line("u", self.u),
+            ]
+        )
+
+
+class AerotechTarget(BaseModel):
+    x: Optional[float] = None
+    y: Optional[float] = None
+    z: Optional[float] = None
+    u: Optional[float] = None
+
+    def to_payload(self) -> dict:
+        return self.model_dump(exclude_none=True)
+
+
+class AerotechRotationScanRequest(RotationScanRequest):
+    rotation_deg: float
+    time_sec: float
+    start_pos_deg: float
+    async_move: bool = Field(default=False, alias="async")
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    def to_payload(self) -> dict:
+        return self.model_dump(by_alias=True, exclude_none=True)
\ No newline at end of file