AareDAQ/common/src/aaredaqlib/models.py

from enum import Enum
from typing import Annotated, Literal, Tuple, List, Optional

from pydantic import BaseModel, Field

from aaredaqlib.coordinate import Coordinate
from aaredaqlib.diffraction_geometry import DiffractionGeometry
from aaredaqlib.sample_geometry import SampleGeometryModel

class TokenData(BaseModel):
    sub: str # Username
    pgroups: List[str]
    session: int
    staff: bool = False

class DewarAddress(BaseModel):
    segment: Literal["A", "B", "C", "D", "E", "F", "X", "R"]
    pos: Annotated[int, Field(ge=1, le=5)]


class SampleDewarAddress(BaseModel):
    puck: DewarAddress
    pin: Annotated[int, Field(ge=1, le=16)]


# From the database for puck loading
class PuckInfo(BaseModel):
    db_id: int
    puck_name: str
    dewar_name: str
    user: str = ""
    location: Optional[DewarAddress] = None


# From TELL to database after loading
class PuckLoadedInfo(BaseModel):
    puck_name: str
    location: DewarAddress


# From database to TELL after loading
class SampleShortInfo(BaseModel):
    db_id: int
    puck_name: str
    dewar_name: str
    sample_name: str
    user: str = ""
    pin: Annotated[int, Field(ge=1, le=16)]
    location: DewarAddress | None = None
    priority: Optional[float] = 1.0
    comment: Optional[str] = None
    mount_count: int = 0
    def tell_address(self) -> SampleDewarAddress:
        return SampleDewarAddress(puck=self.location, pin=self.pin)

    def loc_str(self) -> str:
        if self.location is None:
            return "-"
        else:
            return f"{self.location.segment}{self.location.pos}-{self.pin}"

    def loc_str_sort(self) -> str:
        if self.location is None:
            return ""
        else:
            return f"{self.location.segment}{self.location.pos}-{self.pin:02d}"


class SampleShortInfoList(BaseModel):
    s: List[SampleShortInfo]


class BeamMarkCoeffModel(BaseModel):
    """
    Model to calculate beam center for a given zoom level based on a quadratic approximation
    For both x and y it contains tuple of coefficients a, b, c (a*zoom^2 + b*zoom + c)
    Default is beam center in 1000,1000 at any zoom level
    """

    coeff_x: Tuple[float, float, float] = (0, 0, 1000.0)
    coeff_y: Tuple[float, float, float] = (0, 0, 1000.0)

    def apply(self, zoom: float):
        return Coordinate(
            x=self.coeff_x[0] * zoom**2 + self.coeff_x[1] * zoom + self.coeff_x[2],
            y=self.coeff_y[0] * zoom**2 + self.coeff_y[1] * zoom + self.coeff_y[2],
        )


class LoopCenteringZoomModelElem(BaseModel):
    zoom_value: float
    sam_cam_gain: float
    sam_cam_exp: float


class LoopCenteringZoomModel(BaseModel):
    z: List[LoopCenteringZoomModelElem]


# zoom gain exp (X06DA)
#    1  50 0.05
#  200  50 0.05
#  500  70 0.05
#  700 150 0.05
#  800 100 0.10
# 1000 200 0.10


def def_loop_centering_zoom() -> LoopCenteringZoomModel:
    return LoopCenteringZoomModel(
        z=[
            LoopCenteringZoomModelElem(
                zoom_value=200, sam_cam_gain=50.0, sam_cam_exp=0.05
            )
        ]
    )


class BeamlineStateEnum(Enum):
    Maintenance = 1
    SampleExchange = 2
    SampleAlignment = 3
    DataCollection = 4
    DewarTransfer = 5
    XrayFluorescence = 6
    BeamLocation = 7
    Moving = 8
    RobotSampleExchange = 9
    # Busy state is used whenever transition between two states happen or mounting/data collection procedure happens


class SessionsStateEnum(Enum):
    OwnedByYou = 1
    OwnedByElse = 2
    Vacant = 3

class SampleCameraSettings(BaseModel):
    gain: float
    exposure: float

class AutofocusSettings(BaseModel):
    center_x_pxl: float
    center_y_pxl: float
    radius_pxl: float
    z_range_um: float
    z_steps: int

class BeamlineStatus(BaseModel):
    ring_current_mA: float
    light: Annotated[float, Field(ge=0.0, le=100.0)]
    cryojet_K: float
    shutter_open: bool
    flux_ph_s: float
    sample_camera: SampleCameraSettings

class SessionStatus(BaseModel):
    session: SessionsStateEnum = SessionsStateEnum.Vacant
    current_pgroup: str | None = None
    staff: bool = False

class DAQStatusModel(BaseModel):
    geom: SampleGeometryModel
    diffraction: DiffractionGeometry
    bl: BeamlineStatus
    state: BeamlineStateEnum
    busy: bool
    sample: SampleShortInfo | None = None
    session: SessionStatus

class BeamlineSettingsModel(BaseModel):
    # Float properties
    cryojet_park_position: float | None = 12.0
    cryojet_measurement_position: float | None = 5.0

    dtz_max: float | None = 1600.0
    dtz_min: float | None = 120.0
    dtz_collection: float | None = 130.0
    dtz_park: float | None = 150.0
    dtz_wash_sample_distance: float | None = 150.0
    dtz_bsz_safety_margin: float | None = 50.0
    bsz: float | None = 25.0

    camera_max_magnification: float | None = 1.0
    camera_min_magnification: float | None = 500.0
    camera_translation_factor_a: float | None = 0.00253
    camera_translation_factor_b: float | None = 255.0

    cryojet_in_use: bool | None = True