AareDAQ/src/aare/common/models.py

from pathlib import Path
import re
from enum import Enum
from typing import Annotated, Literal, Tuple, List, Optional
from dataclasses import dataclass
from pydantic import BaseModel, Field, field_validator, AfterValidator, ConfigDict, AliasChoices

from aare.common.coordinate import Coordinate, positive_coords
from aare.common.diffraction_geometry import DiffractionGeometry
from aare.common.sample_geometry import SampleGeometryModel
from jfjoch_client.models.scan_result import ScanResult

from aare.common.beamline import MXBeamline
from aare.common.tell_models import TellStateModel


class StagePositionEnum(Enum):
    MEASURE = 0
    PARK = 1
    DOWN = 2
    UNKNOWN = 3


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


class DataCollectionParameters(BaseModel):
    model_config = ConfigDict(from_attributes=True, populate_by_name=True)

    directory: Optional[str] = None
    oscillation: Optional[float] = None  # Only accept positive float
    exposure: Optional[float] = None  # Only accept positive floats between 0 and 1
    totalangle: Optional[int] = Field(  # was totalrange
        default=None,
        validation_alias=AliasChoices('totalangle', 'totalrange')
    )  # Only accept positive integers between 0 and 360
    transmission: Optional[int] = None  # Only accept positive integers between 0 and 100
    targetresolution: Optional[float] = None  # Only accept positive float
    beamsize: Optional[str] = None
    aperture: Optional[int] = None  # Optional string field
    datacollectiontype: Optional[str] = None  # Only accept "standard", other types might be added later
    processingpipeline: Optional[str] = ""  # Only accept "gopy", "autoproc", "xia2dials"
    spacegroupnumber: Optional[int] = None  # Only accept positive integers between 1 and 230
    unitcell: Optional[str] = Field(  # was cellparameters
        default=None,
        validation_alias=AliasChoices('unitcell', 'cellparameters')
    )  # Must be a set of six positive floats or integers
    rescutkey: Optional[str] = None  # Only accept "is" or "cchalf"
    rescutvalue: Optional[float] = None  # Must be a positive float if rescutkey is provided
    processingresolution: Optional[float] = Field(  # was userresolution
        default=None,
        validation_alias=AliasChoices('processingresolution', 'userresolution')
    )
    pdbid: Optional[str] = ""  # Accepts either the format of the protein data bank code or {provided}
    autoprocfull: Optional[bool] = None
    procfull: Optional[bool] = None
    adpenabled: Optional[bool] = None
    noano: Optional[bool] = None
    ffcscampaign: Optional[bool] = None
    trustedhigh: Optional[float] = None  # Should be a float between 0 and 2.0
    autoprocextraparams: Optional[str] = None  # Optional string field
    chiphiangles: Optional[float] = None  # Optional float field between 0 and 30
    dose: Optional[float] = None  # Optional float field
    cloud: bool = True
    pdbmodel: Optional[str] = None

    @field_validator("directory", mode="after")
    @classmethod
    def directory_characters(cls, v):

        # Default directory value if empty
        if not v:  # Handles None or empty cases
            default_value = "{date}/{prefix}"
            return default_value

        # Strip trailing slashes and store original value for comparison
        v = str(v).strip("/")  # Ensure it's a string and no trailing slashes
        original_value = v

        # Replace spaces with underscores
        v = v.replace(" ", "_")

        # Validate directory pattern with macros and allowed characters
        valid_macros = [
            # Current macros
            "{puck}",
            "{position}",
            "{prefix}",
            "{date}",
            "{run}",
            "{beamline}",
            # Legacy macros — accepted for back-compat, no longer documented
            "{sgPuck}",
            "{sgPosition}",
            "{sgPrefix}",
            "{sgPriority}",
            "{protein}",
            "{method}",
        ]
        valid_macro_pattern = re.compile(
            "|".join(re.escape(macro) for macro in valid_macros)
        )

        # Check if the value contains valid macros
        allowed_chars_pattern = "[a-z0-9_.+-/]"
        v_without_macros = valid_macro_pattern.sub("macro", v)

        allowed_path_pattern = re.compile(
            f"^(({allowed_chars_pattern}+|macro)*/*)*$", re.IGNORECASE
        )
        if not allowed_path_pattern.match(v_without_macros):
            raise ValueError(
                f"'{v}' is not valid. Value must be a valid path or macro."
            )
        return v

    @field_validator("unitcell", mode="before")
    @classmethod
    def unitcell_format(cls, v):
        if v:
            tokens = v.replace(",", " ").split()
            try:
                values = [float(i) for i in tokens]
            except ValueError:
                raise ValueError(
                    f"'{v}' is not valid. "
                    "Value must be six positive floats or integers (space or comma separated)."
                )
            if len(values) != 6 or any(val <= 0 for val in values):
                raise ValueError(
                    f"'{v}' is not valid. "
                    "Value must be six positive floats or integers (space or comma separated)."
                )
        return v

    @field_validator("cloud", mode="before")
    @classmethod
    def coerce_cloud_default(cls, v):
        if v in ("", None):
            return True
        if isinstance(v, bool):
            return v
        v_str = str(v).strip().lower()
        if v_str in {"true", "yes", "1"}:
            return True
        if v_str in {"false", "no", "0"}:
            return False
        raise ValueError("cloud must be blank for default, or True/False")


# From database to TELL after loading
class SampleShortInfo(BaseModel):
    db_id: int
    puck_name: str
    dewar_name: str
    sample_name: str
    run_number: int
    aaredb_params: Optional[DataCollectionParameters] = None
    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
    rotation_count: int = 0
    raster_count: int = 0
    screening_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}"

    @classmethod
    def from_dict(cls, data: dict):
        return cls(**data)


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 FluorescenceSpectrumParameterModel(BaseModel):
    erase: bool = True
    acq_time_s: float
    transmission: Annotated[float, Field(ge=0,le=1)] | None = None


class FluorescenceSpectrumOutputModel(BaseModel):
    bkg: list[float] | None = None
    spectrum: list[float]
    energy_eV: list[float]
    average_dead_time: Annotated[float, Field(ge=0,le=1)]


class MLBoxType(Enum):
    LOOP_ALL  = 0
    PIN       = 1
    CRYSTAL   = 2
    LOOP_FACE = 3
    ICE       = 4
    NEEDLE    = 5


class BoundingBoxModel(BaseModel):
    top_x: float
    top_y: float
    bottom_x: float
    bottom_y: float


class MLBoxModel(BaseModel):
    cls: MLBoxType
    box: BoundingBoxModel
    conf: float

    @staticmethod
    def from_tuple(cls:MLBoxType, box_tuple: tuple[float, float, float, float], conf: float) -> "MLBoxModel":
        x1, y1, x2, y2 = box_tuple
        return MLBoxModel(
            cls=cls,
            box=BoundingBoxModel(top_x=float(x1), top_y=float(y1), bottom_x=float(x2), bottom_y=float(y2)),
            conf=float(conf)
        )


class MLOutputModel(BaseModel):
    boxes: dict[str, MLBoxModel] = {}

    @staticmethod
    def get_class_str(cls: MLBoxType) -> str:
        if cls == MLBoxType.LOOP_ALL:
            return "Loop_all"
        if cls == MLBoxType.PIN:
            return "Pin"
        if cls == MLBoxType.CRYSTAL:
            return "Crystal"
        if cls == MLBoxType.LOOP_FACE:
            return "Loop_face"
        if cls == MLBoxType.ICE:
            return "Ice"
        if cls == MLBoxType.NEEDLE:
            return "Needle"
        return "Unknown"

    def _next_unique_key(self, base: str) -> str:
        if base not in self.boxes:
            return base
        i = 2
        while f"{base}_{i}" in self.boxes:
            i += 1
        return f"{base}_{i}"

    def add_box(self, cls: MLBoxType, box_tuple: tuple[float, float, float, float],
                conf: float | None = None) -> str:
        key_base = self.get_class_str(cls)
        key = self._next_unique_key(key_base)
        self.boxes[key] = MLBoxModel.from_tuple(cls, box_tuple, conf)
        return key

    def get_box_model(self, key: str) -> MLBoxModel | None:
        return self.boxes.get(key)

    def get_box_tuple(self, key: str) -> tuple[float, float, float, float] | None:
        m = self.get_box_model(key)
        if not m or not m.box:
            return None
        return (m.box.top_x, m.box.top_y, m.box.bottom_x, m.box.bottom_y)

    def get_box_tuple_with_conf(self, key: str) -> tuple[float, float, float, float, float] | None:
        m = self.get_box_model(key)
        if not m or not m.box:
            return None
        conf = float(m.conf) if m.conf is not None else 0.0
        return (m.box.top_x, m.box.top_y, m.box.bottom_x, m.box.bottom_y, conf)

    def get_keys_for_class(self, cls: MLBoxType) -> list[str]:
        base = self.get_class_str(cls)
        return [k for k, v in self.boxes.items() if v.cls == cls and (k == base or k.startswith(base + "_"))]

    def get_models_for_class(self, cls: MLBoxType) -> list[MLBoxModel]:
        keys = self.get_keys_for_class(cls)
        return [self.boxes[k] for k in keys]

    def get_tuples_for_class(self, cls: MLBoxType) -> list[tuple[float, float, float, float]]:
        out: list[tuple[float, float, float, float]] = []
        for m in self.get_models_for_class(cls):
            if m.box:
                out.append((m.box.top_x, m.box.top_y, m.box.bottom_x, m.box.bottom_y))
        return out

    def get_tuples_with_conf_for_class(self, cls: MLBoxType) -> list[tuple[float, float, float, float, float]]:
        out: list[tuple[float, float, float, float, float]] = []
        for m in self.get_models_for_class(cls):
            if m.box:
                out.append((m.box.top_x, m.box.top_y, m.box.bottom_x, m.box.bottom_y, float(m.conf)))
        return out

    def get_best_for_class(self, cls: MLBoxType) -> MLBoxModel | None:
        models = self.get_models_for_class(cls)
        if not models:
            return None
        return max(models, key=lambda m: (m.conf if m.conf is not None else 0.0))


class BeamlineStateEnum(Enum):
    Maintenance = 1
    SampleExchange = 2
    SampleAlignment = 3
    DataCollection = 4
    DewarTransfer = 5
    XrayFluorescence = 6
    BeamLocation = 7
    Moving = 8
    RobotSampleExchange = 9
    XtalSnapshot = 10
    BeamstopAlignment = 11
    FluxMeasurement = 12

    def display_name(self) -> str:
        return {
            BeamlineStateEnum.Maintenance: "Maintenance",
            BeamlineStateEnum.SampleExchange: "Sample exchange",
            BeamlineStateEnum.SampleAlignment: "Sample alignment",
            BeamlineStateEnum.DataCollection: "Data collection",
            BeamlineStateEnum.DewarTransfer: "Dewar transfer",
            BeamlineStateEnum.XrayFluorescence: "X-ray fluorescence",
            BeamlineStateEnum.BeamLocation: "Beam location",
            BeamlineStateEnum.Moving: "Moving",
            BeamlineStateEnum.RobotSampleExchange: "Robot sample exchange",
            BeamlineStateEnum.XtalSnapshot: "Xtal snapshot",
            BeamlineStateEnum.BeamstopAlignment: "Beamstop alignment",
            BeamlineStateEnum.FluxMeasurement: "Flux measurement",
        }.get(self, "-")


class DAQOperation(str, Enum):
    AUTOMATION = "automation"
    MOUNT = "mount"
    UNMOUNT = "unmount"
    LOOP_CENTERING = "loop_centering"
    FACE_CENTERING = "face_centering"
    RASTER = "raster"
    ROTATION = "rotation"
    MEASURE = "measure"


class SessionsStateEnum(Enum):
    Vacant = 0
    OwnedByYou = 1
    OwnedByElse = 2
    PendingYouToElse = 3
    PendingElseToYou = 4


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


class ZoomModeEnum(Enum):
    User = 1
    LoopCenter = 2
    BeamLocation = 3


class ZoomModel(BaseModel):
    z:dict[float, SampleCameraSettings]

    def get_camera_settings(self, zoom_value: float) -> SampleCameraSettings | tuple[float, float]:
        if not self.z:
            raise ValueError("No zoom data available")

        if zoom_value in self.z:
            elem = self.z[zoom_value]
            return SampleCameraSettings(gain=elem.gain, exposure=elem.exposure)

        sorted_zooms = sorted(self.z.keys())

        if zoom_value <= sorted_zooms[0]:
            elem = self.z[sorted_zooms[0]]
            return SampleCameraSettings(gain=elem.gain, exposure=elem.exposure)

        if zoom_value >= sorted_zooms[-1]:
            elem = self.z[sorted_zooms[-1]]
            return SampleCameraSettings(gain=elem.gain, exposure=elem.exposure)
        print('interpolating zoom')
        for i in range(len(sorted_zooms) - 1):
            if sorted_zooms[i] <= zoom_value <= sorted_zooms[i + 1]:
                lower_zoom = sorted_zooms[i]
                upper_zoom = sorted_zooms[i + 1]

                lower_elem = self.z[lower_zoom]
                upper_elem = self.z[upper_zoom]

                t = (zoom_value - lower_zoom) / (upper_zoom - lower_zoom)

                interpolated_gain = lower_elem.gain + t * (upper_elem.gain - lower_elem.gain)
                interpolated_exp = lower_elem.exposure + t * (upper_elem.exposure - lower_elem.exposure)

                return SampleCameraSettings(gain=interpolated_gain, exposure=interpolated_exp)

        closest = min(self.z.keys(), key=lambda x: abs(x - zoom_value))
        elem = self.z[closest]
        return SampleCameraSettings(gain=elem.gain, exposure=elem.exposure)

def def_zoom(beamline) -> ZoomModel:
    print(f'user zoom for {beamline}')
    if beamline == MXBeamline.X06DA:
        return ZoomModel(
            z={
                1: SampleCameraSettings(gain=0, exposure=0.05),
                280: SampleCameraSettings(gain=0, exposure=0.05),
                500: SampleCameraSettings(gain=0, exposure=0.1),
                700: SampleCameraSettings(gain=0, exposure=0.15),
                800: SampleCameraSettings(gain=0, exposure=0.2),
                1000: SampleCameraSettings(gain=0, exposure=0.25)
            }
        )
    elif beamline == MXBeamline.X10SA:
        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)})
    elif beamline == MXBeamline.X06SA:
        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)})
    elif beamline == MXBeamline.SIMULATED:
        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)})
    else:
        raise ValueError(f"Invalid beamline: {beamline}")

def def_bl_zoom(beamline) -> ZoomModel:
    if beamline == MXBeamline.X06DA:
        return ZoomModel(
            z={
                1: SampleCameraSettings(gain=0, exposure=0.002),
                280: SampleCameraSettings(gain=0, exposure=0.002),
                500: SampleCameraSettings(gain=0, exposure=0.002),
                700: SampleCameraSettings(gain=0, exposure=0.002),
                800: SampleCameraSettings(gain=0, exposure=0.002),
                1000: SampleCameraSettings(gain=0, exposure=0.005)
            }
        )
    elif beamline == MXBeamline.X10SA:
        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)})
    elif beamline == MXBeamline.X06SA:
        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)})
    elif beamline == MXBeamline.SIMULATED:
        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)})
    else:
        raise ValueError(f"Invalid beamline: {beamline}")

def def_loop_centering_zoom(beamline) -> ZoomModel:
    if beamline == MXBeamline.X06DA:
        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.05)})
    elif beamline == MXBeamline.X10SA:
        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002),
                            280: SampleCameraSettings(gain=0, exposure=0.002)})
    elif beamline == MXBeamline.X06SA:
        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.05)})
    elif beamline == MXBeamline.SIMULATED:
        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.05)})
    else:
        raise ValueError(f"Invalid beamline: {beamline}")

def zoom_manager(mode: ZoomModeEnum = ZoomModeEnum.User, beamline: MXBeamline = None) -> ZoomModel | None:
    if beamline is None or beamline == MXBeamline.SIMULATED:
        print('SIMULATED')
        return def_zoom(beamline)
    if mode == ZoomModeEnum.BeamLocation:
        return def_bl_zoom(beamline)
    elif mode == ZoomModeEnum.LoopCenter:
        return def_loop_centering_zoom(beamline)
    elif mode == ZoomModeEnum.User:
        return def_zoom(beamline)
    else:
        raise ValueError(f"Invalid zoom mode: {mode}")


class AutofocusSettings(BaseModel):
    center_x_pxl: float | None # Use beam center
    center_y_pxl: float | None # Use beam center
    radius_pxl: float
    z_range_um: float
    z_steps: int


class BeamlineStatus(BaseModel):
    name: str
    ring_current_mA: float
    front_light: Annotated[float, Field(ge=0.0, le=100.0)]
    back_light: Annotated[float, Field(ge=0.0, le=100.0)]
    cryojet_K: float
    shutter_open: bool
    exp_shutter_open: bool | None
    flux_ph_s: float
    sample_camera: SampleCameraSettings
    transmission: Annotated[float, Field(ge=0.0, le=1.0)] | None
    zoom: float
    commissioning_mode: bool
    dtz_min: float
    dtz_max: float
    # Hutch personnel-safety system state. ``pss_prohibited`` is True when the
    # hutch is interlocked so the robot may move (PROHIBITED-STATE); the GUI
    # blocks a mount when it is False. ``pss_alarm`` is True when ALARM-STATE
    # != 0 (warning). Defaults keep older payloads/constructors valid and avoid
    # the GUI false-blocking when an old server omits the field.
    pss_prohibited: bool = True
    pss_alarm: bool = False


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


class OpenGuiSessionInfo(BaseModel):
    session: int
    username: str
    staff: bool = False
    last_seen_ts: float
    last_interaction_ts: float | None = None
    close_requested: bool = False
    close_requested_by: str | None = None
    close_requested_at: float | None = None
    close_grace_seconds: int | None = None
    holds_baton: bool = False


class CrystalSize(BaseModel):
    x: float = 0.0
    y: float = 0.0
    z: float = 0.0


class DAQStatusModel(BaseModel):
    geom: SampleGeometryModel
    diffraction: DiffractionGeometry
    bl: BeamlineStatus
    state: BeamlineStateEnum
    busy: bool
    sample: SampleShortInfo | None = None
    session: SessionStatus
    open_guis: list[OpenGuiSessionInfo] = []
    box: BoundingBoxModel | None = None
    last_best_res: float | None = None
    last_best_b_factor: float | None = None
    crystal_size: CrystalSize = CrystalSize(x=0,y=0,z=0)

    tell_connected: bool = True
    tell_error: str | None = None
    tell_state: TellStateModel | None = None

    smargon_connected: bool = True
    smargon_error: str | None = None

    aerotech_connected: bool = True
    aerotech_error: str | None = None


class BeamlineSettingsModel(BaseModel):
    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 = 512.0


class CryojetSettingsModel(BaseModel):
    cryojet_park_position: float | None = 12.0
    cryojet_measurement_position: float | None = 5.0
    cryojet_in_use: bool | None = True


class SimpleStrategyInputModel(BaseModel):
    last_best_res: float | None = None
    last_best_b_factor: float | None = None
    crystal_size: CrystalSize = CrystalSize(x=0, y=0, z=0)
    angular_range: int = 360
    start_angle: float = 0
    incr_omega_deg:float = 0.2
    d_vis: float = 1.5
    current_temp_k: float = 100
    filename: Optional[str] = None


class SimpleScanParameters(BaseModel):
    filename: Optional[str] = None
    dtz: float = 150
    exp_time_s: float = 0.02
    start_omega_deg:float = 0
    incr_omega_deg: float = 0.2
    steps: int = 1800
    transmission: Annotated[float, Field(ge=0.0, le=1.0)] = 1.0
    last_best_res: float | None = None
    last_best_b_factor: float | None = None
    crystal_size: CrystalSize = CrystalSize(x=0, y=0, z=0)
    flux_ph_s: float | None = None
    calculated_dose_Mgy: float | None = None
    calculated_dose_rate_MGy_s: float | None = None
    xtal_size_dose_rate_MGy_s: float | None = None
    target_dose_MGy: float | None = None
    beam_size_x_um: float | None = None
    beam_size_y_um: float | None = None
    dose_rate_MGy_s: float | None = None
    d_vis: float | None = None
    d_tar: float | None = None


class ScanResultPayloadModel(BaseModel):
    result: ScanResult
    sample_id: int
    attach_image: bool = True
    beam_mark_pxl: tuple[float, float]
    beam_size_mm: Annotated[Coordinate, AfterValidator(positive_coords)]


class RecoveryActionRequest(BaseModel):
    confirmation_code: str


@dataclass
class LoopCenteringResult:
    success: bool
    comment: str | None = None
    error: Exception | None = None