shiftregister/shiftregister/fallback/models.py

from shiftregister.importer.models import *
from django.db.models import Max, Sum
from django.db.models import Count, Exists, OuterRef, ExpressionWrapper
from django.db.models.lookups import LessThan
from django.db.models.fields import DateTimeField
import math

night_shift_query = Q(start_at__hour__gte=21) | Q(start_at__hour__lte=10)


class TeamMember(models.Model):
    name = models.CharField(max_length=100)
    fallback_shifts = models.ManyToManyField(Shift, through="FallbackAssignment")

    def assign_random_shifts(self):
        if self.fallback_shifts.count() != 0:
            return
        selector1 = (~night_shift_query) & Q(
            deleted=False, calendar__needs_fallback=True
        )
        selector2 = night_shift_query & Q(deleted=False, calendar__needs_fallback=True)
        self._assign_from_bucket(selector1)
        self._assign_from_bucket(selector2)

    def _assign_from_bucket(self, bucket_selector):
        free_bucket = (
            Event.objects.filter(bucket_selector)
            .annotate(
                fallback_count=Count("fallbackassignment"),
            )
            .filter(fallback_count__lt=F("required_helpers"))
        )

        total_slot_count = Event.objects.filter(bucket_selector).aggregate(
            sum=Sum("required_helpers")
        )["sum"]
        free_slot_count = free_bucket.annotate(
            needed_helpers=F("required_helpers") - F("fallback_count")
        ).aggregate(sum=Sum("needed_helpers"))["sum"]
        active_team_members = (
            TeamMember.objects.filter(~Q(fallback_shifts=None)).count() + 1
        )

        quota = global_preferences["helper__fallback_quota"]
        number_of_team_members = TeamMember.objects.count()
        max_shifts_per_member = math.ceil(
            total_slot_count / max((number_of_team_members * quota), 1)
        )
        shifts_per_member = math.ceil(total_slot_count / (active_team_members))
        shift_count = min(max_shifts_per_member, shifts_per_member)

        blocked_times = []
        for shift in self.fallback_shifts.all():
            blocked_times.append(
                Q(start_at__gte=(shift.start_at + shift.duration))
                | Q(end_at__lte=shift.start_at)
            )

        # easy part: enough free shifts for everyone:
        assigned_shift_count = 0
        for _ in range(shift_count):
            shift = (
                free_bucket.annotate(
                    end_at=ExpressionWrapper(
                        F("start_at") + F("duration"),
                        output_field=models.DateTimeField(),
                    )
                )
                .filter(*blocked_times)
                .order_by("?")
                .first()
            )
            if not shift:
                break
            self.fallback_shifts.add(shift)
            assigned_shift_count += 1
            blocked_times.append(
                Q(start_at__gte=(shift.start_at + shift.duration))
                | Q(end_at__lte=shift.start_at)
            )

        # there is a chance that even if quota*teammembers team members are activated, there are still unassigned shifts
        # this happens if there are shifts with multiple people left, as we cannot assign multiple slots for
        # the same shift to one member.
        # for now we will just reduce the quota a bit to calculate for these cases.

        shifts_needed = shift_count - assigned_shift_count
        # this is not done very often so we can do this kinda inefficient but readable and maintainable:
        # for each missing shift, get the team member who has the most shifts in our bucket and take one of them at random
        # but also take care to not take any slots for shifts we already have...

        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"
            )
            relevant_assignments = FallbackAssignment.objects.filter(
                shift_id__in=candidate_shift_ids
            )

            # get teammembers sorted by the most shifts in the relevant bucket
            sorted_members = (
                TeamMember.objects.annotate(
                    relevant_fallback_count=Count(
                        "fallback_shifts",
                        distinct=True,
                        filter=Q(fallback_shifts__id__in=candidate_shift_ids),
                    ),
                    overall_fallback_count=Count("fallback_shifts"),
                )
                .exclude(pk=self.pk)
                .order_by("-relevant_fallback_count", "-overall_fallback_count", "?")
            )
            assignment = False
            for member in sorted_members:
                # now get all their assignments in the relevant bucket but exclude the ones where we already have a slot in the same shift...
                assignment = (
                    FallbackAssignment.objects.annotate(
                        start_at=F("shift__start_at"),
                        end_at=ExpressionWrapper(
                            F("shift__start_at") + F("shift__duration"),
                            output_field=models.DateTimeField(),
                        ),
                    )
                    .filter(team_member_id=member.id, shift_id__in=candidate_shift_ids)
                    .exclude(
                        shift_id__in=FallbackAssignment.objects.filter(
                            team_member_id=self.pk
                        ).values("shift_id")
                    )
                    .filter(*blocked_times)
                    .order_by("?")
                    .first()
                )
                if assignment:
                    break
            if not assignment:
                print("could not find any matching assignments to take away")
                return
            shifts_needed -= 1
            blocked_times.append(
                Q(start_at__gte=(assignment.shift.start_at + assignment.shift.duration))
                | Q(end_at__lte=assignment.shift.start_at)
            )
            self.fallback_shifts.add(assignment.shift)
            assignment.delete()

    def __str__(self):
        return f"{self.name}"


class FallbackAssignment(models.Model):
    shift = models.ForeignKey(Shift, on_delete=models.CASCADE)
    team_member = models.ForeignKey(TeamMember, on_delete=models.CASCADE)

    def __str__(self):
        return f"{self.shift} {self.team_member.name}"