foodforthought-cli 0.2.7__py3-none-any.whl → 0.3.0__py3-none-any.whl

ate/robot/calibration_state.py

@@ -0,0 +1,388 @@
+"""
+Calibration State Machine with Enforced Prerequisites.
+
+This module ensures calibration steps are completed in the correct order,
+making it IMPOSSIBLE to skip critical steps like direction calibration.
+
+The Problem This Solves:
+-----------------------
+Previously, a user could run:
+    ate robot calibrate start    # Get servo ranges
+    ate robot behavior pickup    # CRASH! Wrong direction!
+
+The behavior would fail catastrophically because we knew WHICH servos
+existed but not WHETHER positive commands moved TOWARD or AWAY from targets.
+
+The Solution:
+------------
+Each calibration stage must be completed before the next can begin.
+Commands check prerequisites and refuse to run if not met.
+
+Stages:
+    1. DEVICE_DISCOVERED  - Robot found on USB/network
+    2. SERVO_RANGES_KNOWN - Min/max/neutral for each servo
+    3. STRUCTURE_IDENTIFIED - Which servo controls which joint
+    4. DIRECTION_CALIBRATED - For each servo: does + move toward target?
+    5. VALIDATED - Pickup test passed
+
+Usage:
+    from ate.robot.calibration_state import CalibrationState, check_prerequisite
+
+    # Check if we can run a behavior
+    state = CalibrationState.load("my_robot")
+    if not state.direction_calibrated:
+        print("Run 'ate robot calibrate direction' first")
+        sys.exit(1)
+
+    # Or use the decorator
+    @check_prerequisite("direction_calibrated")
+    def robot_behavior_command(...):
+        ...
+"""
+
+import json
+import os
+from dataclasses import dataclass, field, asdict
+from datetime import datetime
+from pathlib import Path
+from typing import Optional, Dict, List, Any
+from functools import wraps
+import sys
+
+
+# State storage directory
+STATE_DIR = Path.home() / ".ate" / "calibration_state"
+
+
+@dataclass
+class DirectionMapping:
+    """Records whether positive servo values move toward or away from target."""
+    servo_id: int
+    toward_direction: str  # "positive" or "negative"
+    confidence: float  # 0.0 to 1.0
+    delta_pixels: float  # measured movement in pixels
+    calibrated_at: str  # ISO timestamp
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "DirectionMapping":
+        return cls(**d)
+
+
+@dataclass
+class CalibrationState:
+    """
+    Tracks calibration progress for a robot.
+
+    Enforces that each stage is completed before the next can begin.
+    Persists to disk so progress is not lost between sessions.
+    """
+
+    robot_name: str
+
+    # Stage 1: Device discovered
+    device_discovered: bool = False
+    serial_port: Optional[str] = None
+    camera_ip: Optional[str] = None
+    discovered_at: Optional[str] = None
+
+    # Stage 2: Servo ranges known
+    servo_ranges_known: bool = False
+    servo_count: int = 0
+    servo_ranges: Dict[int, Dict[str, int]] = field(default_factory=dict)
+    ranges_calibrated_at: Optional[str] = None
+
+    # Stage 3: Kinematic structure identified
+    structure_identified: bool = False
+    joint_mappings: Dict[int, str] = field(default_factory=dict)  # servo_id → joint_name
+    structure_identified_at: Optional[str] = None
+
+    # Stage 4: Direction calibrated (THE CRITICAL MISSING STEP)
+    direction_calibrated: bool = False
+    direction_mappings: Dict[int, DirectionMapping] = field(default_factory=dict)
+    direction_calibrated_at: Optional[str] = None
+
+    # Stage 5: Validated via task execution
+    validated: bool = False
+    validation_task: Optional[str] = None
+    validated_at: Optional[str] = None
+
+    # Metadata
+    created_at: str = field(default_factory=lambda: datetime.now().isoformat())
+    updated_at: str = field(default_factory=lambda: datetime.now().isoformat())
+
+    def save(self) -> bool:
+        """Persist state to disk."""
+        STATE_DIR.mkdir(parents=True, exist_ok=True)
+        path = STATE_DIR / f"{self.robot_name}.json"
+
+        self.updated_at = datetime.now().isoformat()
+
+        # Convert to JSON-serializable dict
+        data = asdict(self)
+
+        # Convert DirectionMapping objects to dicts
+        if self.direction_mappings:
+            data["direction_mappings"] = {
+                str(k): asdict(v) if isinstance(v, DirectionMapping) else v
+                for k, v in self.direction_mappings.items()
+            }
+
+        try:
+            with open(path, "w") as f:
+                json.dump(data, f, indent=2)
+            return True
+        except Exception as e:
+            print(f"Error saving calibration state: {e}", file=sys.stderr)
+            return False
+
+    @classmethod
+    def load(cls, robot_name: str) -> "CalibrationState":
+        """Load state from disk, or create new if not exists."""
+        path = STATE_DIR / f"{robot_name}.json"
+
+        if not path.exists():
+            return cls(robot_name=robot_name)
+
+        try:
+            with open(path) as f:
+                data = json.load(f)
+
+            # Convert direction_mappings back to DirectionMapping objects
+            if "direction_mappings" in data and data["direction_mappings"]:
+                data["direction_mappings"] = {
+                    int(k): DirectionMapping.from_dict(v)
+                    for k, v in data["direction_mappings"].items()
+                }
+
+            return cls(**data)
+        except Exception as e:
+            print(f"Error loading calibration state: {e}", file=sys.stderr)
+            return cls(robot_name=robot_name)
+
+    @classmethod
+    def list_all(cls) -> List[str]:
+        """List all saved calibration states."""
+        if not STATE_DIR.exists():
+            return []
+        return [p.stem for p in STATE_DIR.glob("*.json")]
+
+    def get_current_stage(self) -> str:
+        """Get the name of the current (incomplete) stage."""
+        if not self.device_discovered:
+            return "DEVICE_DISCOVERY"
+        if not self.servo_ranges_known:
+            return "SERVO_RANGE_DISCOVERY"
+        if not self.structure_identified:
+            return "STRUCTURE_IDENTIFICATION"
+        if not self.direction_calibrated:
+            return "DIRECTION_CALIBRATION"
+        if not self.validated:
+            return "VALIDATION"
+        return "COMPLETE"
+
+    def get_next_command(self) -> str:
+        """Get the command to run for the next stage."""
+        stage = self.get_current_stage()
+        commands = {
+            "DEVICE_DISCOVERY": "ate robot discover",
+            "SERVO_RANGE_DISCOVERY": "ate robot calibrate start",
+            "STRUCTURE_IDENTIFICATION": "ate robot calibrate structure",
+            "DIRECTION_CALIBRATION": "ate robot calibrate direction",
+            "VALIDATION": "ate robot calibrate validate",
+            "COMPLETE": "(calibration complete)",
+        }
+        return commands.get(stage, "unknown")
+
+    def check_prerequisite(self, required_stage: str) -> tuple[bool, str]:
+        """
+        Check if a prerequisite stage is complete.
+
+        Returns:
+            (passed, message) - passed is True if OK, message explains what's needed
+        """
+        stages = {
+            "device_discovered": (self.device_discovered, "ate robot discover"),
+            "servo_ranges_known": (self.servo_ranges_known, "ate robot calibrate start"),
+            "structure_identified": (self.structure_identified, "ate robot calibrate structure"),
+            "direction_calibrated": (self.direction_calibrated, "ate robot calibrate direction"),
+            "validated": (self.validated, "ate robot calibrate validate"),
+        }
+
+        if required_stage not in stages:
+            return True, ""
+
+        passed, command = stages[required_stage]
+
+        if passed:
+            return True, ""
+
+        # Build error message showing what's needed
+        msg = f"\n{'='*60}\n"
+        msg += "PREREQUISITE NOT MET\n"
+        msg += f"{'='*60}\n\n"
+        msg += f"Required: {required_stage.replace('_', ' ').title()}\n\n"
+        msg += "You must complete calibration steps in order:\n\n"
+
+        checklist = [
+            ("device_discovered", "Device Discovery", "ate robot discover"),
+            ("servo_ranges_known", "Servo Ranges", "ate robot calibrate start"),
+            ("structure_identified", "Kinematic Structure", "ate robot calibrate structure"),
+            ("direction_calibrated", "Direction Mapping", "ate robot calibrate direction"),
+            ("validated", "Task Validation", "ate robot calibrate validate"),
+        ]
+
+        for stage_key, name, cmd in checklist:
+            status = "✓" if getattr(self, stage_key) else "○"
+            arrow = "  ← NEXT" if stage_key == required_stage else ""
+            msg += f"  {status} {name}{arrow}\n"
+            if stage_key == required_stage:
+                break
+
+        msg += f"\nRun: {command}\n"
+
+        return False, msg
+
+    # Stage completion methods
+
+    def complete_device_discovery(
+        self,
+        serial_port: Optional[str] = None,
+        camera_ip: Optional[str] = None,
+    ) -> None:
+        """Mark device discovery as complete."""
+        self.device_discovered = True
+        self.serial_port = serial_port
+        self.camera_ip = camera_ip
+        self.discovered_at = datetime.now().isoformat()
+        self.save()
+
+    def complete_servo_ranges(
+        self,
+        servo_count: int,
+        servo_ranges: Dict[int, Dict[str, int]],
+    ) -> None:
+        """Mark servo range discovery as complete."""
+        self.servo_ranges_known = True
+        self.servo_count = servo_count
+        self.servo_ranges = servo_ranges
+        self.ranges_calibrated_at = datetime.now().isoformat()
+        self.save()
+
+    def complete_structure_identification(
+        self,
+        joint_mappings: Dict[int, str],
+    ) -> None:
+        """Mark structure identification as complete."""
+        self.structure_identified = True
+        self.joint_mappings = joint_mappings
+        self.structure_identified_at = datetime.now().isoformat()
+        self.save()
+
+    def complete_direction_calibration(
+        self,
+        direction_mappings: Dict[int, DirectionMapping],
+    ) -> None:
+        """Mark direction calibration as complete."""
+        self.direction_calibrated = True
+        self.direction_mappings = direction_mappings
+        self.direction_calibrated_at = datetime.now().isoformat()
+        self.save()
+
+    def complete_validation(self, task_name: str) -> None:
+        """Mark validation as complete."""
+        self.validated = True
+        self.validation_task = task_name
+        self.validated_at = datetime.now().isoformat()
+        self.save()
+
+    def reset(self) -> None:
+        """Reset all calibration state."""
+        self.__init__(robot_name=self.robot_name)
+        self.save()
+
+    def print_status(self) -> None:
+        """Print current calibration status."""
+        print(f"\nCalibration Status: {self.robot_name}")
+        print("=" * 50)
+
+        stages = [
+            ("device_discovered", "Device Discovery", self.discovered_at),
+            ("servo_ranges_known", "Servo Ranges", self.ranges_calibrated_at),
+            ("structure_identified", "Kinematic Structure", self.structure_identified_at),
+            ("direction_calibrated", "Direction Mapping", self.direction_calibrated_at),
+            ("validated", "Task Validation", self.validated_at),
+        ]
+
+        for key, name, timestamp in stages:
+            status = "✓" if getattr(self, key) else "○"
+            time_str = f" ({timestamp[:10]})" if timestamp else ""
+            print(f"  {status} {name}{time_str}")
+
+        print()
+
+        current = self.get_current_stage()
+        if current != "COMPLETE":
+            print(f"Next step: {self.get_next_command()}")
+        else:
+            print("Robot is fully calibrated and validated!")
+
+
+def check_prerequisite(required_stage: str):
+    """
+    Decorator that checks calibration prerequisites before running a command.
+
+    Usage:
+        @check_prerequisite("direction_calibrated")
+        def robot_behavior_command(profile_name: str, ...):
+            # Will only run if direction_calibrated is True
+            ...
+    """
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            # Try to get robot name from various sources
+            robot_name = None
+
+            # Check kwargs
+            for key in ["profile_name", "name", "robot_name"]:
+                if key in kwargs and kwargs[key]:
+                    robot_name = kwargs[key]
+                    break
+
+            # Check positional args (common patterns)
+            if not robot_name and args:
+                robot_name = args[0] if isinstance(args[0], str) else None
+
+            if not robot_name:
+                # Can't check without robot name - let the function handle it
+                return func(*args, **kwargs)
+
+            # Load state and check prerequisite
+            state = CalibrationState.load(robot_name)
+            passed, message = state.check_prerequisite(required_stage)
+
+            if not passed:
+                print(message, file=sys.stderr)
+                sys.exit(1)
+
+            return func(*args, **kwargs)
+
+        return wrapper
+    return decorator
+
+
+def require_direction_calibration(robot_name: str) -> bool:
+    """
+    Check if direction calibration is complete. Exit if not.
+
+    Use this in commands that need direction calibration.
+    Returns True if calibrated, exits with error if not.
+    """
+    state = CalibrationState.load(robot_name)
+    passed, message = state.check_prerequisite("direction_calibrated")
+
+    if not passed:
+        print(message, file=sys.stderr)
+        return False
+
+    return True