Add constraints for maximum days and unique shifts per week block in worker model

rota_generator/shifts.py

@@ -1055,6 +1055,35 @@ class RotaBuilder(object):
             initialize=0,
         )
 
+        # Helper binary variables for per-worker unique-shift block limits.
+        # Index is (worker_id, block_length, block_start_week, shift_name).
+        unique_shift_block_indices = []
+        for worker in self.workers:
+            for constraint in worker.max_unique_shifts_per_week_block:
+                for week_blocks in self.get_week_block_iterator(constraint.week_block):
+                    if not week_blocks:
+                        continue
+                    block_start_week = week_blocks[0]
+                    for shift_name in self.get_shift_names():
+                        if any(
+                            shift_name in self.get_shift_names_by_week_day(
+                                week,
+                                day,
+                                return_empty_if_week_day_not_found=True,
+                            )
+                            for week in week_blocks
+                            for day in days
+                        ):
+                            unique_shift_block_indices.append(
+                                (worker.id, constraint.week_block, block_start_week, shift_name)
+                            )
+
+        self.model.worker_shift_used_in_week_block = Var(
+            unique_shift_block_indices,
+            within=Binary,
+            initialize=0,
+        )
+
         if self.constraint_options_model.balance_bank_holidays:
             # Bank holidays
             self.model.bank_holiday_count = Var(
@@ -2374,6 +2403,55 @@ class RotaBuilder(object):
                         <= max_days
                     )
 
+            for constraint in worker.max_days_per_week_block:
+                for week_blocks in self.get_week_block_iterator(constraint.week_block):
+                    self.model.constraints.add(
+                        sum(
+                            self.model.works_day[worker.id, week, day]
+                            for week, day in self.get_week_day_combinations()
+                            if week in week_blocks
+                        )
+                        <= constraint.max_days
+                    )
+
+            for constraint in worker.max_unique_shifts_per_week_block:
+                for week_blocks in self.get_week_block_iterator(constraint.week_block):
+                    if not week_blocks:
+                        continue
+
+                    block_start_week = week_blocks[0]
+                    used_shift_vars = []
+                    for shift_name in self.get_shift_names():
+                        shift_assignments = [
+                            self.model.works[worker.id, week, day, shift_name]
+                            for week, day in self.get_week_day_combinations()
+                            if week in week_blocks
+                            and shift_name in self.get_shift_names_by_week_day(
+                                week,
+                                day,
+                                return_empty_if_week_day_not_found=True,
+                            )
+                        ]
+                        if not shift_assignments:
+                            continue
+
+                        used_shift_var = self.model.worker_shift_used_in_week_block[
+                            worker.id,
+                            constraint.week_block,
+                            block_start_week,
+                            shift_name,
+                        ]
+                        used_shift_vars.append(used_shift_var)
+                        self.model.constraints.add(
+                            sum(shift_assignments)
+                            <= len(shift_assignments) * used_shift_var
+                        )
+
+                    if used_shift_vars:
+                        self.model.constraints.add(
+                            sum(used_shift_vars) <= constraint.max_unique_shifts
+                        )
+
             for week_blocks in self.get_week_block_iterator(4):
                 # Prevent more than n number shifts per 4 weeks
                 try: