aaredb/backend/app/models.py

from sqlalchemy import (
    Column,
    Integer,
    String,
    Date,
    ForeignKey,
    JSON,
    DateTime,
    Boolean,
    func,
    Enum,
)
from sqlalchemy.orm import relationship
from .database import Base
from datetime import datetime
import enum


class Shipment(Base):
    __tablename__ = "shipments"

    id = Column(Integer, primary_key=True, index=True, autoincrement=True)
    pgroups = Column(String(255), nullable=False)
    shipment_name = Column(String(255), index=True)
    shipment_date = Column(Date, nullable=True)
    shipment_status = Column(String(255), nullable=True)
    comments = Column(String(200), nullable=True)
    contact_id = Column(Integer, ForeignKey("contacts.id"), nullable=False)
    return_address_id = Column(Integer, ForeignKey("addresses.id"), nullable=False)
    proposal_id = Column(Integer, ForeignKey("proposals.id"), nullable=True)

    contact = relationship("Contact", back_populates="shipments")
    return_address = relationship("Address", back_populates="shipments")
    proposal = relationship("Proposal", back_populates="shipments")
    dewars = relationship("Dewar", back_populates="shipment")


class Contact(Base):
    __tablename__ = "contacts"

    id = Column(Integer, primary_key=True, index=True, autoincrement=True)
    status = Column(String(255), default="active")
    pgroups = Column(String(255), nullable=False)
    firstname = Column(String(255), nullable=False)
    lastname = Column(String(255))
    phone_number = Column(String(255))
    email = Column(String(255))
    shipments = relationship("Shipment", back_populates="contact")


class LocalContact(Base):
    __tablename__ = "local_contacts"

    id = Column(Integer, primary_key=True, index=True, autoincrement=True)
    status = Column(String(255), default="active")
    firstname = Column(String(255), nullable=False)
    lastname = Column(String(255), nullable=False)
    phone_number = Column(String(255), nullable=False)
    email = Column(String(255), nullable=False)


class Address(Base):
    __tablename__ = "addresses"

    id = Column(Integer, primary_key=True, index=True, autoincrement=True)
    status = Column(String(255), default="active")
    pgroups = Column(String(255), nullable=False)
    street = Column(String(255), nullable=False)
    house_number = Column(String(255), nullable=True)
    city = Column(String(255), nullable=False)
    state = Column(String(255), nullable=True)
    zipcode = Column(String(255), nullable=False)
    country = Column(String(255), nullable=False)

    shipments = relationship("Shipment", back_populates="return_address")


class DewarType(Base):
    __tablename__ = "dewar_types"
    id = Column(Integer, primary_key=True, index=True, autoincrement=True)
    dewar_type = Column(String(255), unique=True, index=True)
    serial_numbers = relationship("DewarSerialNumber", back_populates="dewar_type")


class DewarSerialNumber(Base):
    __tablename__ = "dewar_serial_numbers"
    id = Column(Integer, primary_key=True, index=True, autoincrement=True)
    serial_number = Column(String(255), index=True)
    dewar_type_id = Column(Integer, ForeignKey("dewar_types.id"))
    dewar_type = relationship("DewarType", back_populates="serial_numbers")


class Dewar(Base):
    __tablename__ = "dewars"

    id = Column(Integer, primary_key=True, index=True, autoincrement=True)
    pgroups = Column(String(255), nullable=False)
    dewar_name = Column(String(255), nullable=False)
    dewar_type_id = Column(Integer, ForeignKey("dewar_types.id"), nullable=True)
    dewar_serial_number_id = Column(
        Integer, ForeignKey("dewar_serial_numbers.id"), nullable=True
    )
    status = Column(String(255), nullable=True)
    unique_id = Column(String(255), unique=True, index=True, nullable=True)
    tracking_number = Column(String(255), nullable=True)
    shipment_id = Column(Integer, ForeignKey("shipments.id"))
    return_address_id = Column(Integer, ForeignKey("addresses.id"))
    contact_id = Column(Integer, ForeignKey("contacts.id"))

    shipment = relationship("Shipment", back_populates="dewars")
    return_address = relationship("Address")
    contact = relationship("Contact")
    pucks = relationship("Puck", back_populates="dewar")

    dewar_type = relationship("DewarType")
    dewar_serial_number = relationship("DewarSerialNumber")
    slot = relationship("Slot", back_populates="dewar")
    events = relationship("LogisticsEvent", back_populates="dewar")
    beamline_location = None
    local_contact_id = Column(Integer, ForeignKey("local_contacts.id"), nullable=True)
    local_contact = relationship("LocalContact")
    beamtime = relationship("Beamtime", back_populates="dewars")
    beamtime_id = Column(Integer, ForeignKey("beamtimes.id"), nullable=True)

    @property
    def number_of_pucks(self) -> int:
        return len(self.pucks) if self.pucks else 0

    @property
    def number_of_samples(self) -> int:
        if not self.pucks:
            return 0
        return sum(len(puck.samples) for puck in self.pucks)


class Proposal(Base):
    __tablename__ = "proposals"

    id = Column(Integer, primary_key=True, index=True, autoincrement=True)
    number = Column(String(255))
    shipments = relationship("Shipment", back_populates="proposal")


class Puck(Base):
    __tablename__ = "pucks"

    id = Column(Integer, primary_key=True, index=True, autoincrement=True)
    puck_name = Column(String(255), index=True)
    puck_type = Column(String(255))
    puck_location_in_dewar = Column(Integer)

    # Foreign keys and relationships
    dewar_id = Column(Integer, ForeignKey("dewars.id"))
    dewar = relationship("Dewar", back_populates="pucks")
    samples = relationship("Sample", back_populates="puck")
    events = relationship("PuckEvent", back_populates="puck")
    beamtime_id = Column(Integer, ForeignKey("beamtimes.id"), nullable=True)
    beamtime = relationship(
        "Beamtime", back_populates="pucks", foreign_keys=[beamtime_id]
    )


class Sample(Base):
    __tablename__ = "samples"

    id = Column(Integer, primary_key=True, index=True, autoincrement=True)
    sample_name = Column(String(255), index=True)
    proteinname = Column(String(255), index=True)
    position = Column(Integer)
    priority = Column(Integer)
    comments = Column(String(255))
    data_collection_parameters = Column(JSON, nullable=True)

    # Foreign keys and relationships
    dewar_id = Column(Integer, ForeignKey("dewars.id"))
    puck_id = Column(Integer, ForeignKey("pucks.id"))
    puck = relationship("Puck", back_populates="samples")
    events = relationship("SampleEvent", back_populates="sample", lazy="joined")
    images = relationship("Image", back_populates="sample", lazy="joined")
    beamtime_id = Column(Integer, ForeignKey("beamtimes.id"), nullable=True)
    beamtime = relationship("Beamtime", back_populates="samples")

    @property
    def mount_count(self) -> int:
        # Dynamically calculate mount_count
        return len([event for event in self.events if event.event_type == "Mounted"])

    @property
    def unmount_count(self) -> int:
        # Dynamically calculate unmount_count
        return len([event for event in self.events if event.event_type == "Unmounted"])


class Slot(Base):
    __tablename__ = "slots"

    id = Column(Integer, primary_key=True, index=True)
    qr_code = Column(String(255), unique=True, index=True)
    label = Column(String(255))
    qr_base = Column(String(255), nullable=True)
    occupied = Column(Boolean, default=False)
    needs_refill = Column(Boolean, default=False)
    dewar_unique_id = Column(String(255), ForeignKey("dewars.unique_id"), nullable=True)
    dewar = relationship("Dewar", back_populates="slot")
    events = relationship("LogisticsEvent", back_populates="slot")


class LogisticsEvent(Base):
    __tablename__ = "logistics_events"

    id = Column(Integer, primary_key=True, index=True)
    dewar_id = Column(Integer, ForeignKey("dewars.id"))
    slot_id = Column(Integer, ForeignKey("slots.id"))
    event_type = Column(String(255), index=True)
    timestamp = Column(DateTime, default=datetime.now)
    dewar = relationship("Dewar", back_populates="events")
    slot = relationship("Slot", back_populates="events")


class SampleEvent(Base):
    __tablename__ = "sample_events"

    id = Column(Integer, primary_key=True, index=True)
    sample_id = Column(Integer, ForeignKey("samples.id"), nullable=False)
    event_type = Column(String(255), nullable=False)
    timestamp = Column(DateTime, default=datetime.now)

    images = relationship("Image", back_populates="sample_event")
    sample = relationship("Sample", back_populates="events")


class PuckEvent(Base):
    __tablename__ = "puck_events"

    id = Column(Integer, primary_key=True, index=True)
    puck_id = Column(Integer, ForeignKey("pucks.id"))
    tell = Column(String(255), nullable=True)
    tell_position = Column(String(255), nullable=True)
    event_type = Column(String(255), index=True)
    timestamp = Column(DateTime, default=datetime.now)

    puck = relationship("Puck", back_populates="events")


SHIFT_CHOICES = ("morning", "afternoon", "night")


class Beamtime(Base):
    __tablename__ = "beamtimes"

    id = Column(Integer, primary_key=True, index=True, autoincrement=True)
    pgroups = Column(String(255), nullable=False)
    shift = Column(Enum(*SHIFT_CHOICES, name="shift_enum"), nullable=False, index=True)
    beamtime_name = Column(String(255), index=True)
    beamline = Column(String(255), nullable=True)
    start_date = Column(Date, nullable=True)
    end_date = Column(Date, nullable=True)
    status = Column(String(255), nullable=True)
    comments = Column(String(200), nullable=True)
    proposal_id = Column(Integer, ForeignKey("proposals.id"), nullable=True)
    local_contact_id = Column(Integer, ForeignKey("local_contacts.id"), nullable=False)

    local_contact = relationship("LocalContact")
    dewars = relationship("Dewar", back_populates="beamtime")
    pucks = relationship("Puck", back_populates="beamtime")
    samples = relationship("Sample", back_populates="beamtime")


class Image(Base):
    __tablename__ = "images"

    id = Column(Integer, primary_key=True, index=True, autoincrement=True)
    pgroup = Column(String(255), nullable=False)
    comment = Column(String(200), nullable=True)
    filepath = Column(String(255), nullable=False)
    status = Column(String(255), nullable=True)
    sample_id = Column(Integer, ForeignKey("samples.id"), nullable=False)
    sample_event_id = Column(Integer, ForeignKey("sample_events.id"), nullable=False)

    sample = relationship("Sample", back_populates="images")
    sample_event = relationship("SampleEvent", back_populates="images")


class ExperimentParameters(Base):
    __tablename__ = "experiment_parameters"
    id = Column(Integer, primary_key=True, index=True, autoincrement=True)
    run_number = Column(Integer, nullable=False)
    type = Column(String(255), nullable=False)
    beamline_parameters = Column(JSON, nullable=True)
    dataset = Column(JSON, nullable=True)
    sample_id = Column(Integer, ForeignKey("samples.id"), nullable=False)


class Results(Base):
    __tablename__ = "results"

    id = Column(Integer, primary_key=True, index=True, autoincrement=True)
    status = Column(String(255), nullable=False)
    result = Column(JSON, nullable=False)  # store the full result object as JSON
    sample_id = Column(Integer, ForeignKey("samples.id"), nullable=False)
    run_id = Column(Integer, ForeignKey("experiment_parameters.id"), nullable=False)

    # optional relationships if you wish to query easily
    # sample = relationship("SampleModel", backref="results")
    # experiment_parameters = relationship("ExperimentParametersModel",
    # backref="results")


#    method = Column(String(255), nullable=False)
#    #resolution: Column(Float(255), nullable=False)
#    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 JobStatus(str, enum.Enum):
    TO_DO = "to_do"
    SUBMITTED = "submitted"
    DONE = "done"
    TO_CANCEL = "to_cancel"
    CANCELLED = "cancelled"
    FAILED = "failed"


class Jobs(Base):
    __tablename__ = "jobs"

    id = Column(Integer, primary_key=True, index=True)
    sample_id = Column(Integer, ForeignKey("samples.id"), nullable=False)
    run_id = Column(Integer, ForeignKey("experiment_parameters.id"), nullable=False)
    status = Column(String, nullable=False)
    experiment_parameters = relationship(ExperimentParameters)
    created_at = Column(DateTime, server_default=func.now())
    updated_at = Column(DateTime, onupdate=func.now())
    slurm_id = Column(Integer, nullable=True)