feat: improve fallback shift assignment algorithm

shiftregister/fallback/distribution.py

@@ -0,0 +1,56 @@
+import math
+from random import random
+
+from django.db.models import Case, Count, F, Sum, When
+
+from shiftregister.importer.models import Event
+
+
+class Bucket:
+    def __init__(self, selector, sub_buckets=[]):
+        if not isinstance(sub_buckets, list) or not all(
+            map(lambda b: isinstance(b, Bucket), sub_buckets)
+        ):
+            raise ValueError("sub_buckets must be a list of instances of class Bucket")
+
+        self.selector = selector
+        self.sub_buckets = sub_buckets
+
+    def __iter__(self):
+        for bucket in self.sub_buckets:
+            yield bucket.selector, bucket
+
+    def has_sub_buckets(self):
+        return len(self.sub_buckets) > 0
+
+
+def distribute(total_slots, total_assignments, pool, qs=None):
+    if not isinstance(pool, Bucket):
+        raise ValueError("pool must be an instance of class Bucket")
+
+    qs = (qs or Event.objects).filter(pool.selector)
+
+    for selector, bucket in pool:
+        bucket_qs = qs.filter(selector).annotate(
+            real_required_helpers=Case(
+                When(required_helpers=0, then=F("room__required_helpers")),
+                default=F("required_helpers"),
+            )
+        )
+
+        bucket_slots = bucket_qs.aggregate(sum=Sum("real_required_helpers"))["sum"]
+
+        extra_chance, bucket_assignments = math.modf(
+            bucket_slots / total_slots * total_assignments
+        )
+        bucket_assignments = int(bucket_assignments)
+        if extra_chance > random():
+            bucket_assignments += 1
+
+        if bucket.has_sub_buckets():
+            yield from distribute(bucket_slots, bucket_assignments, bucket, qs)
+        else:
+            yield bucket_assignments, bucket_qs
+
+        total_slots -= bucket_slots
+        total_assignments -= bucket_assignments

shiftregister/fallback/models.py

@@ -12,6 +12,8 @@ from django.utils import timezone
 
 from shiftregister.importer.models import *
 
+from .distribution import Bucket, distribute
+
 
 def generate_id():
     return int.from_bytes(secrets.token_bytes(3), byteorder="big")
@@ -40,9 +42,37 @@ class TeamMember(models.Model):
         return urlsafe_b64encode(self.id.to_bytes(3, byteorder="big")).decode("utf-8")
 
     def assign_random_shifts(self):
-        needs_fallback = Q(
-            deleted=False, calendar__needs_fallback=True, calendar__restricted=False
+        if self.fallback_shifts.count() != 0:
+            return
+
+        needs_fallback = Q(deleted=False, calendar__needs_fallback=True)
+
+        total_slots = (
+            Event.objects.filter(needs_fallback)
+            .annotate(
+                real_required_helpers=Case(
+                    When(required_helpers=0, then=F("room__required_helpers")),
+                    default=F("required_helpers"),
+                ),
             )
+            .aggregate(sum=Sum("real_required_helpers"))["sum"]
+        )
+
+        number_of_team_members = TeamMember.objects.count()
+        quota = global_preferences["helper__fallback_quota"]
+        max_shifts_per_member = total_slots / max(number_of_team_members * quota, 1)
+
+        active_team_members = (
+            TeamMember.objects.filter(~Q(fallback_shifts=None)).count() + 1
+        )
+        shifts_per_member = total_slots / active_team_members
+
+        extra_chance, total_assignments = math.modf(
+            min(max_shifts_per_member, shifts_per_member)
+        )
+        total_assignments = int(total_assignments)
+        if extra_chance > random():
+            total_assignments += 1
 
         current_tz = timezone.get_current_timezone()
         # create a datetime combining the last date having fallback shifts
@@ -58,50 +88,38 @@ class TeamMember(models.Model):
         is_last_night = Q(start_at__gte=last_night)
         is_night_shift = Q(start_at__hour__gte=21) | Q(start_at__hour__lte=10)
 
-        if self.fallback_shifts.count() != 0:
-            return
+        is_restricted = Q(calendar__restricted=True)
 
-        day_shifts = ~is_night_shift & ~is_last_night & needs_fallback
-        night_shifts = is_night_shift & ~is_last_night & needs_fallback
-        shit_shifts = is_last_night & needs_fallback
+        day_shifts = Bucket(~is_night_shift & ~is_last_night)
+        night_shifts = Bucket(is_night_shift & ~is_last_night)
+        shit_shifts = Bucket(is_last_night)
 
-        self._assign_from_bucket(day_shifts)
-        self._assign_from_bucket(night_shifts)
-        self._assign_from_bucket(shit_shifts)
-
-        for calendar in Calendar.objects.filter(needs_fallback=True, restricted=True):
-            self._assign_from_bucket(Q(deleted=False, calendar=calendar))
-
-    def _assign_from_bucket(self, bucket_selector):
-        bucket = Event.objects.filter(bucket_selector).annotate(
-            fallback_count=Count("fallbackassignment"),
-            real_required_helpers=Case(
-                When(required_helpers=0, then=F("room__required_helpers")),
-                default=F("required_helpers"),
-            ),
+        restricted_buckets = [
+            Bucket(Q(calendar=calendar))
+            for calendar in Calendar.objects.filter(
+                needs_fallback=True, restricted=True
             )
+        ]
+
+        for bucket, assignments in distribute(
+            total_slots,
+            total_assignments,
+            Bucket(
+                needs_fallback,
+                [
+                    Bucket(~is_restricted, [day_shifts, night_shifts, shit_shifts]),
+                    Bucket(is_restricted, restricted_buckets),
+                ],
+            ),
+        ):
+            self._assign_from_bucket(bucket, assignments)
+
+    def _assign_from_bucket(self, bucket, shift_count):
+        bucket = bucket.annotate(fallback_count=Count("fallbackassignment"))
         free_bucket = bucket.filter(fallback_count__lt=F("real_required_helpers"))
 
         total_slot_count = bucket.aggregate(sum=Sum("real_required_helpers"))["sum"]
 
-        quota = global_preferences["helper__fallback_quota"]
-        number_of_team_members = TeamMember.objects.count()
-        max_shifts_per_member = total_slot_count / max(
-            number_of_team_members * quota, 1
-        )
-
-        active_team_members = (
-            TeamMember.objects.filter(~Q(fallback_shifts=None)).count() + 1
-        )
-        shifts_per_member = total_slot_count / active_team_members
-
-        extra_chance, shift_count = math.modf(
-            min(max_shifts_per_member, shifts_per_member)
-        )
-        shift_count = int(shift_count)
-        if extra_chance > random():
-            shift_count += 1
-
         blocked_times = []
         for shift in self.fallback_shifts.all():
             blocked_times.append(
@@ -145,9 +163,7 @@ class TeamMember(models.Model):
         while shifts_needed > 0:
             # this is a bit more complex and uses subqueries and id lists because we want to reuse the query selector for events.
             # maybe there is a good way to transform q-expressions to add prefixes to fields but for now this has to work
-            candidate_shift_ids = Event.objects.filter(bucket_selector).values(
-                "shift_ptr_id"
-            )
+            candidate_shift_ids = bucket.values("shift_ptr_id")
             relevant_assignments = FallbackAssignment.objects.filter(
                 shift_id__in=candidate_shift_ids
             )