breakpilot-lehrer/timetable-solver-service/app/constraints.py

"""Timefold constraint provider for the school timetable.

Three categories:
  * universal hard — no double-booking class/teacher/room. These can't be
    turned off; the school can't physically run lessons that overlap.
  * DB-driven hard — soft-fallback if is_hard=False. Each constraint joins
    Lesson against a rule-fact collection from the corresponding tt_
    constraint_* table.
  * Quality soft — preferred periods, etc.

Scoring uses HardSoftScore. Hard violations are weighted by 1; soft
violations use the rule's stored `weight` (0-100). The UI rejects any
solution where hard_score < 0.
"""

from timefold.solver.score import (
    constraint_provider,
    HardSoftScore,
    ConstraintFactory,
    Constraint,
    Joiners,
)

from .domain import Lesson
from .rules import (
    TeacherUnavailableDayRule, TeacherUnavailableWindowRule, TeacherExcludedRoomRule,
    RoomUnavailableRule, SubjectPreferredPeriodRule, RoomRequiresTypeRule,
)


@constraint_provider
def define_constraints(factory: ConstraintFactory) -> list[Constraint]:
    return [
        # ---------- Universal hard ----------
        _class_conflict(factory),
        _teacher_conflict(factory),
        _room_conflict(factory),

        # ---------- DB-driven hard or soft ----------
        _teacher_unavailable_day(factory),
        _teacher_unavailable_window(factory),
        _teacher_excluded_room(factory),
        _room_unavailable(factory),
        _room_requires_type(factory),

        # ---------- Quality soft ----------
        _subject_preferred_period(factory),
    ]


# ==========================================================================
# Universal hard constraints
# ==========================================================================

def _class_conflict(factory: ConstraintFactory) -> Constraint:
    """A class can't sit in two lessons at once."""
    return (
        factory.for_each_unique_pair(
            Lesson,
            Joiners.equal(lambda l: l.school_class.id),
            Joiners.equal(lambda l: l.timeslot.id if l.timeslot else None),
        )
        .filter(lambda l1, l2: l1.timeslot is not None and l2.timeslot is not None)
        .penalize(HardSoftScore.ONE_HARD)
        .as_constraint("class_conflict")
    )


def _teacher_conflict(factory: ConstraintFactory) -> Constraint:
    """A teacher can't run two lessons at once."""
    return (
        factory.for_each_unique_pair(
            Lesson,
            Joiners.equal(lambda l: l.teacher.id),
            Joiners.equal(lambda l: l.timeslot.id if l.timeslot else None),
        )
        .filter(lambda l1, l2: l1.timeslot is not None and l2.timeslot is not None)
        .penalize(HardSoftScore.ONE_HARD)
        .as_constraint("teacher_conflict")
    )


def _room_conflict(factory: ConstraintFactory) -> Constraint:
    """A room can't host two lessons at once."""
    return (
        factory.for_each_unique_pair(
            Lesson,
            Joiners.equal(lambda l: l.room.id if l.room else None),
            Joiners.equal(lambda l: l.timeslot.id if l.timeslot else None),
        )
        .filter(lambda l1, l2: l1.room is not None and l2.room is not None
                and l1.timeslot is not None and l2.timeslot is not None)
        .penalize(HardSoftScore.ONE_HARD)
        .as_constraint("room_conflict")
    )


# ==========================================================================
# DB-driven constraints
# ==========================================================================

def _score_for(rule, *, hard_per_violation: int = 1) -> HardSoftScore:
    """Pick HardSoftScore from a rule's is_hard + weight."""
    if rule.is_hard:
        return HardSoftScore.of(hard_per_violation, 0)
    return HardSoftScore.of(0, max(rule.weight, 1))


def _teacher_unavailable_day(factory: ConstraintFactory) -> Constraint:
    return (
        factory.for_each(Lesson)
        .filter(lambda l: l.timeslot is not None)
        .join(
            TeacherUnavailableDayRule,
            Joiners.equal(lambda l: l.teacher.id, lambda r: r.teacher_id),
            Joiners.equal(lambda l: l.timeslot.day_of_week, lambda r: r.day_of_week),
        )
        .penalize(HardSoftScore.ONE_HARD, lambda l, r: 1 if r.is_hard else 0)
        .penalize(HardSoftScore.ONE_SOFT, lambda l, r: r.weight if not r.is_hard else 0)
        .as_constraint("teacher_unavailable_day")
    )


def _teacher_unavailable_window(factory: ConstraintFactory) -> Constraint:
    def overlaps(l: Lesson, r: TeacherUnavailableWindowRule) -> bool:
        if l.timeslot is None:
            return False
        # Compare HH:MM strings — they sort correctly when zero-padded.
        return l.timeslot.start_time < r.end_time and l.timeslot.end_time > r.start_time

    return (
        factory.for_each(Lesson)
        .filter(lambda l: l.timeslot is not None)
        .join(
            TeacherUnavailableWindowRule,
            Joiners.equal(lambda l: l.teacher.id, lambda r: r.teacher_id),
            Joiners.equal(lambda l: l.timeslot.day_of_week, lambda r: r.day_of_week),
        )
        .filter(overlaps)
        .penalize(HardSoftScore.ONE_HARD, lambda l, r: 1 if r.is_hard else 0)
        .penalize(HardSoftScore.ONE_SOFT, lambda l, r: r.weight if not r.is_hard else 0)
        .as_constraint("teacher_unavailable_window")
    )


def _teacher_excluded_room(factory: ConstraintFactory) -> Constraint:
    return (
        factory.for_each(Lesson)
        .filter(lambda l: l.room is not None)
        .join(
            TeacherExcludedRoomRule,
            Joiners.equal(lambda l: l.teacher.id, lambda r: r.teacher_id),
            Joiners.equal(lambda l: l.room.id, lambda r: r.room_id),
        )
        .penalize(HardSoftScore.ONE_HARD, lambda l, r: 1 if r.is_hard else 0)
        .penalize(HardSoftScore.ONE_SOFT, lambda l, r: r.weight if not r.is_hard else 0)
        .as_constraint("teacher_excluded_room")
    )


def _room_unavailable(factory: ConstraintFactory) -> Constraint:
    return (
        factory.for_each(Lesson)
        .filter(lambda l: l.room is not None and l.timeslot is not None)
        .join(
            RoomUnavailableRule,
            Joiners.equal(lambda l: l.room.id, lambda r: r.room_id),
            Joiners.equal(lambda l: l.timeslot.day_of_week, lambda r: r.day_of_week),
            Joiners.equal(lambda l: l.timeslot.period_index, lambda r: r.period_index),
        )
        .penalize(HardSoftScore.ONE_HARD, lambda l, r: 1 if r.is_hard else 0)
        .penalize(HardSoftScore.ONE_SOFT, lambda l, r: r.weight if not r.is_hard else 0)
        .as_constraint("room_unavailable")
    )


def _room_requires_type(factory: ConstraintFactory) -> Constraint:
    """If a subject requires a specific room type, the assigned room must match."""
    return (
        factory.for_each(Lesson)
        .filter(lambda l: l.room is not None)
        .join(
            RoomRequiresTypeRule,
            Joiners.equal(lambda l: l.subject.id, lambda r: r.subject_id),
        )
        .filter(lambda l, r: l.room.room_type != r.room_type)
        .penalize(HardSoftScore.ONE_HARD, lambda l, r: 1 if r.is_hard else 0)
        .penalize(HardSoftScore.ONE_SOFT, lambda l, r: r.weight if not r.is_hard else 0)
        .as_constraint("room_requires_type")
    )


# ==========================================================================
# Quality soft
# ==========================================================================

def _subject_preferred_period(factory: ConstraintFactory) -> Constraint:
    """Soft penalty when a lesson lands outside the subject's preferred period range."""
    return (
        factory.for_each(Lesson)
        .filter(lambda l: l.timeslot is not None)
        .join(
            SubjectPreferredPeriodRule,
            Joiners.equal(lambda l: l.subject.id, lambda r: r.subject_id),
        )
        .filter(lambda l, r: not (r.period_from <= l.timeslot.period_index <= r.period_to))
        .penalize(HardSoftScore.ONE_SOFT, lambda l, r: r.weight)
        .as_constraint("subject_preferred_period")
    )