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,
    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 (même poste, même jour, créneau juste avant) : sert au
    # malus de chaînage (voir `engine._score`). Calculé sur tous les créneaux
    # du poste, pas seulement ceux à compléter.
    predecessor_id: dict = {}
    by_position: dict = defaultdict(list)
    for s in shifts:
        by_position[s.position_id].append(s)
    for position_shifts in by_position.values():
        by_day: dict = defaultdict(list)
        for s in position_shifts:
            by_day[s.starts_at.date()].append(s)
        for day_shifts in by_day.values():
            day_shifts.sort(key=lambda s: s.starts_at)
            for previous, current in zip(day_shifts, day_shifts[1:], strict=False):
                predecessor_id[current.id] = str(previous.id)

    assigned_count: dict = defaultdict(int)
    shifts_by_volunteer: dict = defaultdict(list)
    for a in assignments:
        assigned_count[a.shift_id] += 1
        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_id=predecessor_id.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))

    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_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)