AareCommon/src/aarecommon/models/aerotech.py

from enum import Enum

from pydantic import BaseModel, ConfigDict, Field

from aarecommon.models.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: AerotechAxisStatus | None = None
    y: AerotechAxisStatus | None = None
    z: AerotechAxisStatus | None = None
    u: AerotechAxisStatus | None = 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: float | None = None
    y: float | None = None
    z: float | None = None
    u: float | None = 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)