aaredb/backend/app/schemas.py

from typing import List, Optional, Union, Literal
from datetime import datetime
from pydantic import BaseModel, EmailStr, constr, Field, field_validator
from datetime import date
import logging
import re

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)


class loginToken(BaseModel):
    access_token: str
    token_type: str


class loginData(BaseModel):
    username: str
    pgroups: List[str]
    # role: Optional[str] = "user"


class DewarTypeBase(BaseModel):
    dewar_type: str


class DewarTypeCreate(DewarTypeBase):
    pass


class DewarType(DewarTypeBase):
    id: int

    class Config:
        from_attributes = True


class DewarSerialNumberBase(BaseModel):
    serial_number: str
    dewar_type_id: int


class DewarSerialNumberCreate(DewarSerialNumberBase):
    pass


class DewarSerialNumber(DewarSerialNumberBase):
    id: int
    dewar_type: DewarType

    class Config:
        from_attributes = True


class DataCollectionParameters(BaseModel):
    directory: Optional[str] = None
    oscillation: Optional[float] = None  # Only accept positive float
    exposure: Optional[float] = None  # Only accept positive floats between 0 and 1
    totalrange: Optional[int] = None  # 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
    aperture: Optional[str] = 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
    cellparameters: Optional[
        str
    ] = None  # 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
    userresolution: Optional[float] = None
    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

    def to_dict(self):
        """Convert the model instance to a dictionary."""
        return self.dict(
            exclude_unset=True
        )  # Use this built-in method for serialization

    class Config:
        from_attributes = True

    @field_validator("directory", mode="after")
    @classmethod
    def directory_characters(cls, v):
        logger.debug(f"Validating 'directory' field with initial value: {repr(v)}")

        # Default directory value if empty
        if not v:  # Handles None or empty cases
            default_value = "{sgPuck}/{sgPosition}"
            logger.warning(
                f"'directory' field is empty or None. Assigning default value: "
                f"{default_value}"
            )
            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(" ", "_")
        logger.debug(f"Corrected 'directory', spaces replaced: {repr(v)}")

        # Validate directory pattern with macros and allowed characters
        valid_macros = [
            "{date}",
            "{prefix}",
            "{sgPuck}",
            "{sgPosition}",
            "{beamline}",
            "{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."
            )

        # Log and return corrected value
        if v != original_value:
            logger.info(f"Directory was corrected from '{original_value}' to '{v}'")
        return v

    @field_validator("aperture", mode="before")
    @classmethod
    def aperture_selection(cls, v):
        if v is not None:
            try:
                v = int(float(v))
                if v not in {1, 2, 3}:
                    raise ValueError(f" '{v}' is not valid. Value must be 1, 2, or 3.")
            except (ValueError, TypeError) as e:
                raise ValueError(
                    f" '{v}' is not valid. Value must be 1, 2, or 3."
                ) from e
        return v

    @field_validator("oscillation", mode="before")
    @classmethod
    def positive_float_validator(cls, v):
        if v is None:
            return None
        try:
            v = float(v)
            if v <= 0:
                raise ValueError(f"'{v}' is not valid. Value must be a positive float.")
        except (ValueError, TypeError) as e:
            raise ValueError(
                f"'{v}' is not valid. Value must be a positive float."
            ) from e
        return v

    @field_validator("exposure", mode="before")
    @classmethod
    def exposure_in_range(cls, v):
        if v is not None:
            try:
                v = float(v)
                if not (0 <= v <= 1):
                    raise ValueError(
                        f" '{v}' is not valid. Value must be a float between 0 and 1."
                    )
            except (ValueError, TypeError) as e:
                raise ValueError(
                    f" '{v}' is not valid. Value must be a float between 0 and 1."
                ) from e
        return v

    @field_validator("totalrange", mode="before")
    @classmethod
    def totalrange_in_range(cls, v):
        if v is not None:
            try:
                v = int(v)
                if not (0 <= v <= 360):
                    raise ValueError(
                        f" '{v}' is not valid."
                        f"Value must be an integer between 0 and 360."
                    )
            except (ValueError, TypeError) as e:
                raise ValueError(
                    f" '{v}' is not valid."
                    f"Value must be an integer between 0 and 360."
                ) from e
        return v

    @field_validator("transmission", mode="before")
    @classmethod
    def transmission_fraction(cls, v):
        if v is not None:
            try:
                v = int(v)
                if not (0 <= v <= 100):
                    raise ValueError(
                        f" '{v}' is not valid."
                        f"Value must be an integer between 0 and 100."
                    )
            except (ValueError, TypeError) as e:
                raise ValueError(
                    f" '{v}' is not valid."
                    f"Value must be an integer between 0 and 100."
                ) from e
        return v

    @field_validator("datacollectiontype", mode="before")
    @classmethod
    def datacollectiontype_allowed(cls, v):
        allowed = {"standard"}  # Other types of data collection might be added later
        if v and v.lower() not in allowed:
            raise ValueError(f" '{v}' is not valid." f"Value must be one of {allowed}.")
        return v

    @field_validator("processingpipeline", mode="before")
    @classmethod
    def processingpipeline_allowed(cls, v):
        allowed = {"gopy", "autoproc", "xia2dials"}
        if v and v.lower() not in allowed:
            raise ValueError(f" '{v}' is not valid." f"Value must be one of {allowed}.")
        return v

    @field_validator("spacegroupnumber", mode="before")
    @classmethod
    def spacegroupnumber_allowed(cls, v):
        if v is not None:
            try:
                v = int(v)
                if not (1 <= v <= 230):
                    raise ValueError(
                        f" '{v}' is not valid."
                        f"Value must be an integer between 1 and 230."
                    )
            except (ValueError, TypeError) as e:
                raise ValueError(
                    f" '{v}' is not valid."
                    f"Value must be an integer between 1 and 230."
                ) from e
        return v

    @field_validator("cellparameters", mode="before")
    @classmethod
    def cellparameters_format(cls, v):
        if v:
            values = [float(i) for i in v.split(",")]
            if len(values) != 6 or any(val <= 0 for val in values):
                raise ValueError(
                    f" '{v}' is not valid."
                    f"Value must be a set of six positive floats or integers."
                )
        return v

    # @field_validator("rescutkey", "rescutvalue", mode="before")
    # @classmethod
    # def rescutkey_value_pair(cls, values):
    #    rescutkey = values.get("rescutkey")
    #    rescutvalue = values.get("rescutvalue")
    #    if rescutkey and rescutvalue:
    #        if rescutkey not in {"is", "cchalf"}:
    #            raise ValueError("Rescutkey must be either 'is' or 'cchalf'")
    #        if not isinstance(rescutvalue, float) or rescutvalue <= 0:
    #            raise ValueError(
    #                "Rescutvalue must be a positive float if rescutkey is provided"
    #            )
    #    return values

    @field_validator("trustedhigh", mode="before")
    @classmethod
    def trustedhigh_allowed(cls, v):
        if v is not None:
            try:
                v = float(v)
                if not (0 <= v <= 2.0):
                    raise ValueError(
                        f" '{v}' is not valid."
                        f"Value must be a float between 0 and 2.0."
                    )
            except (ValueError, TypeError) as e:
                raise ValueError(
                    f" '{v}' is not valid." f"Value must be a float between 0 and 2.0."
                ) from e
        return v

    @field_validator("chiphiangles", mode="before")
    @classmethod
    def chiphiangles_allowed(cls, v):
        if v is not None:
            try:
                v = float(v)
                if not (0 <= v <= 30):
                    raise ValueError(
                        f" '{v}' is not valid."
                        f"Value must be a float between 0 and 30."
                    )
            except (ValueError, TypeError) as e:
                raise ValueError(
                    f" '{v}' is not valid. Value must be a float between 0 and 30."
                ) from e
        return v

    @field_validator("dose", mode="before")
    @classmethod
    def dose_positive(cls, v):
        if v is not None:
            try:
                v = float(v)
                if v <= 0:
                    raise ValueError(
                        f" '{v}' is not valid. Value must be a positive float."
                    )
            except (ValueError, TypeError) as e:
                raise ValueError(
                    f" '{v}' is not valid. Value must be a positive float."
                ) from e
        return v


class SampleEventCreate(BaseModel):
    event_type: Literal[
        "Mounting", "Centering", "Failed", "Lost", "Collecting", "Unmounting"
    ]


class SampleEventResponse(SampleEventCreate):
    id: int
    sample_id: int
    timestamp: datetime

    class Config:
        from_attributes = True


class Results(BaseModel):
    pipeline: str
    resolution: float
    unit_cell: str
    spacegroup: str
    rmerge: float
    rmeas: float
    isig: float
    cc: float
    cchalf: float
    completeness: float
    multiplicity: float
    nobs: int
    total_refl: int
    unique_refl: int
    comments: Optional[constr(max_length=200)] = None


class ContactCreate(BaseModel):
    pgroups: str
    firstname: str
    lastname: str
    phone_number: str
    email: EmailStr


class Contact(ContactCreate):
    id: int
    status: str = "active"

    class Config:
        from_attributes = True


class ContactUpdate(BaseModel):
    pgroups: str
    firstname: Optional[str] = None
    lastname: Optional[str] = None
    phone_number: Optional[str] = None
    email: Optional[EmailStr] = None


class ContactMinimal(BaseModel):
    firstname: str
    lastname: str
    email: EmailStr
    id: int


class Proposal(BaseModel):
    id: int
    number: str

    class Config:
        from_attributes = True


class LocalContactCreate(BaseModel):
    firstname: str
    lastname: str
    phone_number: str
    email: EmailStr
    status: str = "active"

    class Config:
        from_attributes = True


class LocalContact(LocalContactCreate):
    id: int


class AddressCreate(BaseModel):
    pgroups: str
    house_number: Optional[str] = None
    street: str
    city: str
    state: Optional[str] = None
    zipcode: str
    country: str


class Address(AddressCreate):
    id: int
    status: str = "active"

    class Config:
        from_attributes = True


class AddressUpdate(BaseModel):
    pgroups: str
    house_number: Optional[str] = None
    street: Optional[str] = None
    city: Optional[str] = None
    state: Optional[str] = None
    zipcode: Optional[str] = None
    country: Optional[str] = None


class AddressMinimal(BaseModel):
    house_number: str
    street: str
    city: str
    state: Optional[str] = None
    zipcode: str
    country: str
    id: int


class Sample(BaseModel):
    id: int
    sample_name: str
    position: int  # Position within the puck
    puck_id: int
    crystalname: Optional[str] = Field(None)
    proteinname: Optional[str] = None
    positioninpuck: Optional[int] = Field(None)
    priority: Optional[int] = None
    comments: Optional[str] = None
    data_collection_parameters: Optional[DataCollectionParameters]
    events: List[SampleEventResponse] = []
    mount_count: Optional[int] = None
    unmount_count: Optional[int] = None
    # results: Optional[Results] = None

    class Config:
        from_attributes = True


class SampleCreate(BaseModel):
    sample_name: str = Field(..., alias="crystalname")
    proteinname: Optional[str] = None
    position: int = Field(..., alias="positioninpuck")
    data_collection_parameters: Optional[DataCollectionParameters] = Field(default=None)
    priority: Optional[int] = None
    comments: Optional[str] = None
    results: Optional[Results] = None
    events: Optional[List[str]] = None

    class Config:
        populate_by_name = True


class PuckEvent(BaseModel):
    id: int
    puck_id: int
    tell_position: Optional[str] = None
    event_type: str
    timestamp: datetime

    class Config:
        from_attributes = True


class PuckBase(BaseModel):
    puck_name: str
    puck_type: str
    puck_location_in_dewar: int


class PuckCreate(BaseModel):
    puck_name: str
    puck_type: str
    puck_location_in_dewar: int
    samples: List[SampleCreate] = []


class PuckUpdate(BaseModel):
    puck_name: Optional[str] = None
    puck_type: Optional[str] = None
    puck_location_in_dewar: Optional[int] = None
    dewar_id: Optional[int] = None


class Puck(BaseModel):
    id: int
    puck_name: str
    puck_type: str
    puck_location_in_dewar: int
    dewar_id: int
    events: List[PuckEvent] = []
    samples: List[Sample] = []

    class Config:
        from_attributes = True


class DewarBase(BaseModel):
    pgroups: str
    dewar_name: str
    dewar_type_id: Optional[int] = None
    dewar_serial_number_id: Optional[int] = None
    unique_id: Optional[str] = None
    tracking_number: str
    number_of_pucks: Optional[int] = None
    number_of_samples: Optional[int] = None
    status: str
    contact_id: Optional[int]
    return_address_id: Optional[int]
    pucks: List[PuckCreate] = []

    class Config:
        from_attributes = True


class DewarCreate(DewarBase):
    pass


class Dewar(DewarBase):
    id: int
    pgroups: str
    shipment_id: Optional[int]
    contact: Optional[Contact]
    return_address: Optional[Address]
    pucks: List[Puck] = []  # List of pucks within this dewar

    class Config:
        from_attributes = True


class DewarUpdate(BaseModel):
    pgroups: str
    dewar_name: Optional[str] = None
    dewar_type_id: Optional[int] = None
    dewar_serial_number_id: Optional[int] = None
    unique_id: Optional[str] = None
    tracking_number: Optional[str] = None
    status: Optional[str] = None
    contact_id: Optional[int] = None
    address_id: Optional[int] = None


class DewarSchema(BaseModel):
    id: int
    pgroups: str
    dewar_name: str
    tracking_number: str
    status: str
    contact_id: int
    return_address_id: int

    class Config:
        from_attributes = True

    # shipping status etc will become a logistics event.
    # Tracking will also become an event


class DewarTable(BaseModel):
    id: int
    pgroups: Optional[str] = None  # Make "pgroups" optional
    shipment_id: Optional[int] = None  # Make "shipment_id" optional
    shipment_name: str
    dewar_name: str
    tracking_number: Optional[str] = None
    dewar_type_id: Optional[int] = None
    dewar_serial_number_id: Optional[int] = None
    unique_id: Optional[str] = None
    status: Optional[str] = None
    contact: Optional[List[ContactMinimal]] = None
    address: Optional[List[AddressMinimal]] = None
    event_id: Optional[int] = None
    slot_id: Optional[int] = None
    events: Optional[Union[str, int]] = None
    last_updated: Optional[datetime] = None

    class Config:
        from_attributes = True


class Shipment(BaseModel):
    id: int
    pgroups: str
    shipment_name: str
    shipment_date: date
    shipment_status: str
    comments: Optional[str]
    contact: Optional[Contact]
    return_address: Optional[Address]
    proposal: Optional[Proposal]
    dewars: List[Dewar] = []

    class Config:
        from_attributes = True


class ShipmentCreate(BaseModel):
    pgroups: str
    shipment_name: str
    shipment_date: date
    shipment_status: str
    comments: Optional[constr(max_length=200)]
    contact_id: int
    return_address_id: int
    proposal_id: int
    dewars: List[DewarCreate] = []

    class Config:
        from_attributes = True


class UpdateShipmentComments(BaseModel):
    pgroups: str
    comments: str


class LogisticsEventCreate(BaseModel):
    dewar_qr_code: str
    location_qr_code: str
    transaction_type: str


class SlotSchema(BaseModel):
    id: int
    qr_code: str
    label: str
    qr_base: Optional[str]
    occupied: bool
    needs_refill: bool
    dewar_unique_id: Optional[str]
    dewar_name: Optional[str]
    time_until_refill: Optional[int]
    at_beamline: Optional[bool]
    retrievedTimestamp: Optional[str]
    beamlineLocation: Optional[str]
    shipment_name: Optional[str]
    contact: Optional[str]
    local_contact: Optional[str]

    class Config:
        from_attributes = True


class SampleUpdate(BaseModel):
    sample_name: Optional[str] = None
    proteinname: Optional[str] = None
    priority: Optional[int] = None
    position: Optional[int] = None
    comments: Optional[str] = None
    data_collection_parameters: Optional[DataCollectionParameters] = None

    class Config:
        from_attributes = True


class SetTellPosition(BaseModel):
    puck_name: str
    segment: Optional[str] = Field(
        None,
        pattern="^[A-F]$",  # Valid segments are A, B, C, D, E, F
        description="Segment must be one of A, B, C, D, E, or F."
        "Can be null for no tell_position.",
    )
    puck_in_segment: Optional[int] = Field(
        None,
        ge=1,
        le=5,
        description="Puck in segment must be between 1 and 5."
        "Can be null for no tell_position.",
    )

    @property
    def tell_position(self) -> Optional[str]:
        """
        Combines `segment` and `puck_in_segment` to generate the `tell_position`.
        If either value is `None`, returns `None` to indicate no `tell_position`.
        """
        if self.segment and self.puck_in_segment:
            return f"{self.segment}{self.puck_in_segment}"
        return None


class SetTellPositionRequest(BaseModel):
    tell: str
    pucks: List[SetTellPosition]

    class Config:
        from_attributes = True


class PuckWithTellPosition(BaseModel):
    id: int
    puck_name: str
    puck_type: str
    puck_location_in_dewar: Optional[int]
    dewar_id: Optional[
        int
    ]  # was changed to optional but probably needs to be not optional
    dewar_name: Optional[
        str
    ]  # was changed to optional but probably needs to be not optional
    pgroup: str
    samples: Optional[List[Sample]] = None
    tell_position: Optional[str]

    class Config:
        from_attributes = True


class Beamtime(BaseModel):
    id: int
    pgroups: str
    beamtime_name: str
    beamline: str
    start_date: date
    end_date: date
    status: str
    comments: Optional[constr(max_length=200)] = None
    proposal_id: Optional[int]
    proposal: Optional[Proposal]
    local_contact_id: Optional[int]
    local_contact: Optional[LocalContact]

    class Config:
        from_attributes = True


class ImageCreate(BaseModel):
    pgroup: str
    sample_id: int
    sample_event_id: int
    filepath: str
    status: str = "active"
    comment: Optional[str] = None

    class Config:
        from_attributes = True


class Image(ImageCreate):
    id: int

    class Config:
        from_attributes = True


class ImageInfo(BaseModel):
    id: int
    filepath: str
    comment: Optional[str] = None
    event_type: str
    # run_number: Optional[int]

    class Config:
        from_attributes = True


class characterizationParameters(BaseModel):
    omegaStart_deg: float
    oscillation_deg: float
    omegaStep: float
    chi: float
    phi: float
    numberOfImages: int
    exposureTime_s: float


class RotationParameters(BaseModel):
    omegaStart_deg: float
    omegaStep: float
    chi: float
    phi: float
    numberOfImages: int
    exposureTime_s: float


class gridScanParamers(BaseModel):
    xStart: float
    xStep: float
    yStart: float
    yStep: float
    zStart: float
    zStep: float
    numberOfImages: int
    exposureTime_s: float


class jetParameters(BaseModel):
    hplc_pump_ml_min: float
    pressure_bar: float
    jetDiameter_um: int
    jetSpeed_mm_s: float
    exposureTime_s: float


class detector(BaseModel):
    manufacturer: str
    model: str
    type: str
    serialNumber: str
    detectorDistance_mm: float
    beamCenterX_px: float
    beamCenterY_px: float
    pixelSizeX_um: float
    pixelSizeY_um: float


class BeamlineParameters(BaseModel):
    synchrotron: str
    beamline: str
    detector: detector
    wavelength: float
    # energy: float
    ringCurrent_A: float
    ringMode: str
    undulator: Optional[str] = None
    undulatorgap_mm: Optional[float] = None
    monochromator: str
    # bandwidth_percent: float
    transmission: float
    focusingOptic: str
    beamlineFluxAtSample_ph_s: Optional[float] = None
    beamSizeWidth: Optional[float] = None
    beamSizeHeight: Optional[float] = None
    # dose_MGy: float
    characterization: Optional[characterizationParameters] = None
    rotation: Optional[RotationParameters] = None
    gridScan: Optional[gridScanParamers] = None
    jet: Optional[jetParameters] = None
    cryojetTemperature_K: Optional[float] = None
    humidifierTemperature_K: Optional[float] = None
    humidifierHumidity: Optional[float] = None
    # experimentalHutchTemerature_K: Optional[float] = None
    # experimentalHutchHumidity_percent: Optional[float] = None
    # beamstopDistance_mm: Optional[float] = None
    # beamstopDiameter_mm: Optional[float] = None


class ExperimentParametersBase(BaseModel):
    run_number: int
    beamline_parameters: Optional[BeamlineParameters] = None
    sample_id: int


class ExperimentParametersCreate(ExperimentParametersBase):
    run_number: Optional[int] = None


class ExperimentParametersRead(ExperimentParametersBase):
    id: int

    class Config:
        from_attributes = True


class SampleResult(BaseModel):
    sample_id: int
    sample_name: str
    puck_name: Optional[str]
    dewar_name: Optional[str]
    images: List[ImageInfo]
    experiment_runs: Optional[List[ExperimentParametersRead]] = []


class ResultCreate(BaseModel):
    sample_id: int
    run_id: int
    result: Results

    class Config:
        from_attributes = True


class ResultResponse(BaseModel):
    id: int
    sample_id: int
    run_id: int
    result: Results

    class Config:
        from_attributes = True