benevoles/backend/app/solver/loader.py

"""Construit un `SolverInput` depuis la base pour une édition. Couche de
liaison : c'est le seul endroit du module `solver` qui connaît la base."""

import datetime as dt
from collections import defaultdict

from sqlmodel import Session, select

from app.models import (
    Artist,
    Assignment,
    Availability,
    Edition,
    Shift,
    Volunteer,
    VolunteerArtist,
    VolunteerCompanionWish,
    VolunteerPositionWish,
)
from app.solver.models import SolverInput, SolverShift, SolverVolunteer
from app.validation.models import AvailabilityWindow, ShiftInfo


def _artist_window(artist: Artist) -> ShiftInfo:
    """Fenêtre datée du passage de l'artiste (même convention que les
    créneaux : `end_time <= start_time` ⇒ le passage finit le lendemain)."""
    starts_at = dt.datetime.combine(artist.performance_date, artist.start_time)
    end_date = (
        artist.performance_date
        if artist.end_time > artist.start_time
        else artist.performance_date + dt.timedelta(days=1)
    )
    ends_at = dt.datetime.combine(end_date, artist.end_time)
    return ShiftInfo(id=str(artist.id), starts_at=starts_at, ends_at=ends_at)


def build_solver_input(*, session: Session, edition: Edition) -> SolverInput:
    shifts = list(session.exec(select(Shift).where(Shift.edition_id == edition.id)))
    assignments = list(
        session.exec(select(Assignment).where(Assignment.edition_id == edition.id))
    )
    shifts_by_id = {s.id: s for s in shifts}

    # Prédécesseur(s) : créneaux (n'importe quel poste) du même jour dont la
    # fin coïncide avec le début de celui-ci — sert au malus de chaînage
    # (voir `engine._score`), pour limiter les enchaînements de shifts plutôt
    # que cibler la rotation d'un poste en particulier. Calculé sur tous les
    # créneaux de l'édition, pas seulement ceux à compléter.
    predecessor_ids: dict = {}
    by_day: dict = defaultdict(list)
    for s in shifts:
        by_day[s.starts_at.date()].append(s)
    for day_shifts in by_day.values():
        ids_by_end: dict = defaultdict(list)
        for s in day_shifts:
            ids_by_end[s.ends_at].append(str(s.id))
        distinct_ends = sorted(ids_by_end)
        for s in day_shifts:
            earlier_ends = [e for e in distinct_ends if e <= s.starts_at]
            if earlier_ends:
                predecessor_ids[s.id] = frozenset(ids_by_end[earlier_ends[-1]])

    assigned_count: dict = defaultdict(int)
    shifts_by_volunteer: dict = defaultdict(list)
    assigned_volunteers_by_shift: dict = defaultdict(set)
    for a in assignments:
        assigned_count[a.shift_id] += 1
        assigned_volunteers_by_shift[a.shift_id].add(str(a.volunteer_id))
        shift = shifts_by_id.get(a.shift_id)
        if shift is not None:
            shifts_by_volunteer[a.volunteer_id].append(shift)

    solver_shifts = tuple(
        SolverShift(
            id=str(s.id),
            position_id=str(s.position_id),
            starts_at=s.starts_at,
            ends_at=s.ends_at,
            min_staff=s.min_staff,
            ideal_staff=s.ideal_staff,
            already_assigned=assigned_count.get(s.id, 0),
            predecessor_shift_ids=predecessor_ids.get(s.id, frozenset()),
            assigned_volunteer_ids=frozenset(assigned_volunteers_by_shift.get(s.id, ())),
        )
        for s in shifts
        if assigned_count.get(s.id, 0) < s.ideal_staff
    )

    availabilities: dict = defaultdict(list)
    for a in session.exec(
        select(Availability).where(Availability.edition_id == edition.id)
    ):
        availabilities[a.volunteer_id].append(
            AvailabilityWindow(date=a.date, start_time=a.start_time, end_time=a.end_time)
        )

    wishes: dict = defaultdict(set)
    for w in session.exec(
        select(VolunteerPositionWish).join(
            Volunteer, Volunteer.id == VolunteerPositionWish.volunteer_id
        ).where(Volunteer.edition_id == edition.id)
    ):
        wishes[w.volunteer_id].add(str(w.position_id))

    companion_wishes: dict = defaultdict(set)
    for w in session.exec(
        select(VolunteerCompanionWish).join(
            Volunteer, Volunteer.id == VolunteerCompanionWish.volunteer_id
        ).where(Volunteer.edition_id == edition.id)
    ):
        companion_wishes[w.volunteer_id].add(str(w.companion_id))

    artists_by_id = {
        a.id: a
        for a in session.exec(select(Artist).where(Artist.edition_id == edition.id))
    }
    artist_wishes: dict = defaultdict(list)
    for w in session.exec(
        select(VolunteerArtist).join(
            Volunteer, Volunteer.id == VolunteerArtist.volunteer_id
        ).where(Volunteer.edition_id == edition.id)
    ):
        artist = artists_by_id.get(w.artist_id)
        if artist is not None:
            artist_wishes[w.volunteer_id].append(_artist_window(artist))

    volunteers = list(
        session.exec(select(Volunteer).where(Volunteer.edition_id == edition.id))
    )
    solver_volunteers = tuple(
        SolverVolunteer(
            id=str(v.id),
            availability=tuple(availabilities.get(v.id, ())),
            wished_position_ids=frozenset(wishes.get(v.id, ())),
            wished_companion_ids=frozenset(companion_wishes.get(v.id, ())),
            wished_artist_windows=tuple(artist_wishes.get(v.id, ())),
            existing_shifts=tuple(
                ShiftInfo(id=str(s.id), starts_at=s.starts_at, ends_at=s.ends_at)
                for s in shifts_by_volunteer.get(v.id, ())
            ),
            existing_minutes=sum(
                int((s.ends_at - s.starts_at).total_seconds() // 60)
                for s in shifts_by_volunteer.get(v.id, ())
            ),
        )
        for v in volunteers
    )

    return SolverInput(shifts=solver_shifts, volunteers=solver_volunteers)