diff --git a/shiftregister/fallback/distribution.py b/shiftregister/fallback/distribution.py new file mode 100644 index 0000000..60f76d5 --- /dev/null +++ b/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 diff --git a/shiftregister/fallback/models.py b/shiftregister/fallback/models.py index c7b46b1..81f3d76 100644 --- a/shiftregister/fallback/models.py +++ b/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,10 +42,38 @@ 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 # after 20:00 with the time 20:00 in the current timezone @@ -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) + restricted_buckets = [ + Bucket(Q(calendar=calendar)) + for calendar in Calendar.objects.filter( + needs_fallback=True, restricted=True + ) + ] - 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"), + 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 )