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:

rota_generator/workers.py

@@ -135,6 +135,32 @@ class AvoidShiftOnDates(BaseModel):
         raise ValueError(f"Cannot parse date: {v}")
 
 
+class MaxDaysPerWeekBlockConstraint(BaseModel):
+    max_days: int = Field(
+        ...,
+        ge=1,
+        description="Maximum number of worked days allowed inside each sliding week block.",
+    )
+    week_block: int = Field(
+        ...,
+        ge=1,
+        description="Length of the sliding week block (for example, 2 means every consecutive 2-week window).",
+    )
+
+
+class MaxUniqueShiftsPerWeekBlockConstraint(BaseModel):
+    max_unique_shifts: int = Field(
+        ...,
+        ge=1,
+        description="Maximum number of distinct shift names allowed inside each sliding week block.",
+    )
+    week_block: int = Field(
+        ...,
+        ge=1,
+        description="Length of the sliding week block (for example, 2 means every consecutive 2-week window).",
+    )
+
+
 class Worker(BaseModel):
     name: str
     site: str
@@ -171,6 +197,8 @@ class Worker(BaseModel):
     forced_assignments: List[tuple[int, str, str]] = []  # (week, day, shift_name)
     forced_assignments_by_date: List[tuple[datetime.date, str]] = []  # (date, shift_name)
     max_days_worked_per_week: int | None= None # Default: no limit, can be set to a specific number
+    max_days_per_week_block: list[MaxDaysPerWeekBlockConstraint] = []
+    max_unique_shifts_per_week_block: list[MaxUniqueShiftsPerWeekBlockConstraint] = []
 
 
     assign_as_block_preferences: dict[str, float] = {}  # shift_name -> weight (positive = prefer block)