carrramba-encore-rate/backend/backend/models/stop.py

from __future__ import annotations

from typing import Iterable, Self, Sequence, TYPE_CHECKING

from sqlalchemy import (
    BigInteger,
    Column,
    Enum,
    Float,
    ForeignKey,
    select,
    String,
    Table,
)
from sqlalchemy.orm import (
    mapped_column,
    Mapped,
    relationship,
    selectinload,
    with_polymorphic,
)
from sqlalchemy.schema import Index

from ..db import Base, db
from ..idfm_interface.idfm_types import TransportMode, IdfmState, StopAreaType

if TYPE_CHECKING:
    from .line import Line


stop_area_stop_association_table = Table(
    "stop_area_stop_association_table",
    Base.metadata,
    Column("stop_id", ForeignKey("_stops.id")),
    Column("stop_area_id", ForeignKey("stop_areas.id")),
)


class _Stop(Base):

    db = db

    id = mapped_column(BigInteger, primary_key=True)
    kind = mapped_column(String)

    name = mapped_column(String, nullable=False, index=True)
    town_name = mapped_column(String, nullable=False)
    postal_region = mapped_column(String, nullable=False)
    xepsg2154 = mapped_column(BigInteger, nullable=False)
    yepsg2154 = mapped_column(BigInteger, nullable=False)
    version = mapped_column(String, nullable=False)
    created_ts = mapped_column(BigInteger)
    changed_ts = mapped_column(BigInteger, nullable=False)
    lines: Mapped[list[Line]] = relationship(
        "Line",
        secondary="line_stop_association_table",
        back_populates="stops",
        # lazy="joined",
        lazy="selectin",
    )
    areas: Mapped[list["StopArea"]] = relationship(
        "StopArea", secondary=stop_area_stop_association_table, back_populates="stops"
    )

    __tablename__ = "_stops"
    __mapper_args__ = {"polymorphic_identity": "_stops", "polymorphic_on": kind}
    __table_args__ = (
        # To optimize the ilike requests
        Index(
            "name_idx_gin",
            name,
            postgresql_ops={"name": "gin_trgm_ops"},
            postgresql_using="gin",
        ),
    )

    # TODO: Test https://www.cybertec-postgresql.com/en/postgresql-more-performance-for-like-and-ilike-statements/
    # TODO: Should be able to remove with_polymorphic ?
    @classmethod
    async def get_by_name(cls, name: str) -> Sequence[type[_Stop]] | None:
        session = cls.db.session
        if session is None:
            return None

        stop_stop_area = with_polymorphic(_Stop, [Stop, StopArea])
        stmt = (
            select(stop_stop_area)
            .where(stop_stop_area.name.ilike(f"%{name}%"))
            .options(
                selectinload(stop_stop_area.areas),
                selectinload(stop_stop_area.lines),
            )
        )

        res = await session.execute(stmt)
        stops = res.scalars().all()

        return stops


class Stop(_Stop):

    id = mapped_column(BigInteger, ForeignKey("_stops.id"), primary_key=True)

    latitude = mapped_column(Float, nullable=False)
    longitude = mapped_column(Float, nullable=False)
    transport_mode = mapped_column(Enum(TransportMode), nullable=False)
    accessibility = mapped_column(Enum(IdfmState), nullable=False)
    visual_signs_available = mapped_column(Enum(IdfmState), nullable=False)
    audible_signs_available = mapped_column(Enum(IdfmState), nullable=False)
    record_id = mapped_column(String, nullable=False)
    record_ts = mapped_column(BigInteger, nullable=False)

    __tablename__ = "stops"
    __mapper_args__ = {"polymorphic_identity": "stops", "polymorphic_load": "inline"}


class StopArea(_Stop):

    id = mapped_column(BigInteger, ForeignKey("_stops.id"), primary_key=True)

    type = mapped_column(Enum(StopAreaType), nullable=False)
    stops: Mapped[list["_Stop"]] = relationship(
        "_Stop",
        secondary=stop_area_stop_association_table,
        back_populates="areas",
        lazy="selectin",
        # lazy="joined",
    )

    __tablename__ = "stop_areas"
    __mapper_args__ = {
        "polymorphic_identity": "stop_areas",
        "polymorphic_load": "inline",
    }

    @classmethod
    async def add_stops(
        cls, stop_area_to_stop_ids: Iterable[tuple[int, int]]
    ) -> int | None:
        session = cls.db.session
        if session is None:
            return None

        stop_area_ids, stop_ids = set(), set()
        for stop_area_id, stop_id in stop_area_to_stop_ids:
            stop_area_ids.add(stop_area_id)
            stop_ids.add(stop_id)

        stop_areas_res = await session.execute(
            select(StopArea)
            .where(StopArea.id.in_(stop_area_ids))
            .options(selectinload(StopArea.stops))
        )
        stop_areas: dict[int, StopArea] = {
            stop_area.id: stop_area for stop_area in stop_areas_res.scalars()
        }

        stop_res = await session.execute(select(_Stop).where(_Stop.id.in_(stop_ids)))
        stops: dict[int, _Stop] = {stop.id: stop for stop in stop_res.scalars()}

        found = 0
        for stop_area_id, stop_id in stop_area_to_stop_ids:
            if (stop_area := stop_areas.get(stop_area_id)) is not None:
                if (stop := stops.get(stop_id)) is not None:
                    stop_area.stops.append(stop)
                    found += 1
                else:
                    print(f"No stop found for {stop_id} id")
            else:
                print(f"No stop area found for {stop_area_id}")

        await session.commit()

        return found