From ecf4192d53932fb5cd751767e8c617d40ae92513 Mon Sep 17 00:00:00 2001
From: Anik Stark <anik.stark@psi.ch>
Date: Wed, 10 Dec 2025 17:08:29 +0100
Subject: [PATCH] frappy_psi.magfield: add trainmode

---
 frappy_psi/magfield.py | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/frappy_psi/magfield.py b/frappy_psi/magfield.py
index c57099fe..b93506e7 100644
--- a/frappy_psi/magfield.py
+++ b/frappy_psi/magfield.py
@@ -61,6 +61,7 @@ class SimpleMagfield(HasStates, Drivable):
         'trained field (positive)',
         TupleOf(FloatRange(-99, 0, unit='$'), FloatRange(0, unit='$')),
         readonly=False, default=(0, 0))
+    trainmode = Parameter('train mode flag', EnumType(off=0, on=1, undef=2), default=2)
     wait_stable_field = Parameter(
         'wait time to ensure field is stable', FloatRange(0, unit='s'), readonly=False, default=31)
     ramp_tmo = Parameter(
@@ -149,10 +150,24 @@ class SimpleMagfield(HasStates, Drivable):
         """
         raise NotImplementedError
 
+    def handle_train_mode(self):
+        self.log.info('handle %r %r', self.trained, self.value)
+        if self.trained[0] < self.value < self.trained[1]:
+            trainmode = 'off'
+        else:
+            trainmode = 'on'
+            if self.value > 0:
+                self.trained = (self.trained[0], max(self.trained[1], self.value))
+            else:
+                self.trained = (min(self.trained[0], self.value), self.trained[1])
+        if self.trainmode != trainmode:
+            self.write_trainmode(trainmode)
+
     @status_code(BUSY, 'ramping field')
     def ramp_to_target(self, sm):
         if sm.init:
             self.init_progress(sm, self.value)
+        self.handle_train_mode()
         # Remarks: assume there is a ramp limiting feature
         if abs(self.value - sm.target) > self.tolerance:
             if self.get_progress(sm, self.value) > self.ramp_tmo:
@@ -166,11 +181,15 @@ class SimpleMagfield(HasStates, Drivable):
     def stabilize_field(self, sm):
         if sm.now - sm.stabilize_start < self.wait_stable_field:
             return Retry
+        self.handle_train_mode()
         return self.final_status()
 
     def read_workingramp(self):
         return self.ramp
 
+    def write_trainmode(self, value):
+        """overwrite when needed"""
+
 
 class Magfield(SimpleMagfield):
     status = Parameter(datatype=StatusType(Status))
@@ -335,6 +354,7 @@ class Magfield(SimpleMagfield):
 
     @status_code(Status.RAMPING)
     def ramp_to_target(self, sm):
+        self.handle_train_mode()
         dif = abs(self.value - sm.target)
         if sm.init:
             sm.stabilize_start = 0  # in case current is already at target
@@ -353,6 +373,7 @@ class Magfield(SimpleMagfield):
 
     @status_code(Status.STABILIZING)
     def stabilize_field(self, sm):
+        self.handle_train_mode()
         if sm.now < sm.stabilize_start + self.wait_stable_field:
             return Retry
         return self.check_switch_off