knapsack scheduling algorithm

backend/test_window.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from datetime import date, timedelta
+
+
+def is_off_day(current_date: date, holiday_dates: set[str]) -> bool:
+    return current_date.weekday() >= 5 or current_date.isoformat() in holiday_dates
+
+def get_shift_sequence(
+    current_date: date,
+    holiday_dates: set[str],
+    daytime_shift2_only: bool = False,
+) -> list[tuple[date, int]]:
+    is_holiday = current_date.isoformat() in holiday_dates
+    weekday = current_date.weekday()
+    shifts: list[tuple[date, int]] = []
+
+    if daytime_shift2_only and weekday < 5 and not is_holiday:
+        return [(current_date, 2)]
+
+    if is_off_day(current_date, holiday_dates):
+        return [(current_date, 1), (current_date, 2), (current_date, 3)]
+
+    shifts.append((current_date, 3))
+    next_day = current_date + timedelta(days=1)
+    if is_off_day(next_day, holiday_dates):
+        cursor = next_day
+        while is_off_day(cursor, holiday_dates):
+            shifts.extend([(cursor, 1), (cursor, 2), (cursor, 3)])
+            cursor += timedelta(days=1)
+        shifts.append((cursor, 1))
+        return shifts
+
+    shifts.append((next_day, 1))
+    return shifts
+
+
+def get_shift_capacity_for_date(
+    current_date: date,
+    holiday_dates: set[str],
+    daytime_shift2_only: bool = False,
+) -> dict[int, int]:
+    if daytime_shift2_only:
+        return {1: 0, 2: 480, 3: 0}
+    if is_off_day(current_date, holiday_dates):
+        return {1: 480, 2: 480, 3: 480}
+    return {1: 480, 2: 0, 3: 480}
+
+
+def get_shift_sequence_with_capacity(
+    current_date: date,
+    holiday_dates: set[str],
+    daytime_shift2_only: bool = False,
+) -> tuple[list[tuple[date, int]], int]:
+    shifts = get_shift_sequence(current_date, holiday_dates, daytime_shift2_only)
+    capacity = sum(
+        get_shift_capacity_for_date(day, holiday_dates, daytime_shift2_only).get(shift_index, 0)
+        for day, shift_index in shifts
+    )
+    return shifts, capacity
+
+
+def next_window_start_date(shifts: list[tuple[date, int]]) -> date:
+    last_date, last_shift = shifts[-1]
+    if last_shift in (1, 2):
+        return last_date
+    return last_date + timedelta(days=1)