fix a few more things

rota/shifts.py

@@ -1,10 +1,10 @@
 from calendar import week
 import datetime
 import itertools
-from typing import Dict, Iterable, List, Sequence, Tuple, Set, Any, Type, Literal
+from typing import Dict, Iterable, List, Sequence, Tuple, Set, Any, Type, Literal, Optional
 import time
 
-from pydantic import BaseModel, ConfigDict, Extra, constr, model_validator, StrictInt
+from pydantic import BaseModel, ConfigDict, Extra, constr, model_validator, StrictInt, Field
 
 from pathlib import Path
 
@@ -72,11 +72,24 @@ VALID_SHIFT_CONSTRAINTS = Literal[
     "require_remote_site_presence_week",
 ]
 
+def _pydantic_to_dict(obj):
+    if isinstance(obj, list):
+        return [_pydantic_to_dict(x) for x in obj]
+    if hasattr(obj, "model_dump"):
+        return obj.model_dump()
+    return obj
+
 
 class ShiftConstraint(BaseModel):
     name: VALID_SHIFT_CONSTRAINTS
     options: bool | int | Dict | tuple = False
 
+class MinSummedGradeByShiftsPerDayConstraint(BaseModel):
+    shifts: List[str] = Field(..., description="List of shift names to sum grades over")
+    min_grade_sum: int = Field(..., description="Minimum total grade required")
+    weeks: Optional[List[int]] = None  # If None, applies to all weeks
+    days: Optional[List[str]] = None   # If None, applies to all days
+
 
 class WorkerRequirement(BaseModel):
     """
@@ -310,6 +323,7 @@ class RotaBuilder(object):
             "maximum_allowed_locum_shifts_per_worker": None,
             "minimum_allowed_locum_shifts_per_worker": None,
         }
+        self.constraint_options["min_summed_grade_by_shifts_per_day"]: List[MinSummedGradeByShiftsPerDayConstraint] = []
 
         self.terminate_on_warning = [
             "Worker/duplicate id",
@@ -1007,6 +1021,24 @@ class RotaBuilder(object):
     def build_model_constraints(self):
         self.model.constraints = ConstraintList()  # Create a set of constraints
 
+
+        for constraint in self.constraint_options.get("min_summed_grade_by_shifts_per_day", []):
+            weeks = constraint.weeks if constraint.weeks is not None else self.weeks
+            days_list = constraint.days if constraint.days is not None else self.days
+            for week in weeks:
+                for day in days_list:
+                    # Only consider if all shifts are present that day
+                    if all(shift in self.get_shift_names_by_week_day(week, day) for shift in constraint.shifts):
+                        self.model.constraints.add(
+                            sum(
+                                worker.grade * self.model.works[worker.id, week, day, shift]
+                                for shift in constraint.shifts
+                                for worker in self.get_workers_for_shift(self.get_shift_by_name(shift))
+                            ) >= constraint.min_grade_sum
+                        )
+
+
+
         # Ensure each shift has the requisit number of workers assigned
         for (
             week,
@@ -2269,17 +2301,17 @@ class RotaBuilder(object):
 
 
                 # Try disabling tihs to allow more than 2 shifts ot be assigned on a weekend
-                #self.model.constraints.add(
-                #    self.model.works_weekend[worker.id, week]
-                #    >= sum(
-                #        self.model.works[worker.id, week, day, shiftname]
-                #        for w, day, shiftname in self.get_all_shiftname_combinations(
-                #            week=week
-                #        )
-                #        if day in self.days[5:]
-                #    )
-                #    / 2
-                #)
+                self.model.constraints.add(
+                    self.model.works_weekend[worker.id, week]
+                    >= sum(
+                        self.model.works[worker.id, week, day, shiftname]
+                        for w, day, shiftname in self.get_all_shiftname_combinations(
+                            week=week
+                        )
+                        if day in self.days[5:]
+                    )
+                    / 2
+                )
                 self.model.constraints.add(
                     self.model.works_weekend[worker.id, week]
                     <= sum(
@@ -3090,6 +3122,24 @@ class RotaBuilder(object):
             {"grades": grades, "weeks": weeks, "shifts": shifts}
         )
 
+    def add_min_summed_grade_by_shifts_per_day_constraint(
+        self,
+        shifts: list[str],
+        min_grade_sum: int,
+        weeks: list[int] = None,
+        days: list[str] = None,
+    ):
+        """
+        Helper to add a minimum summed grade constraint for a set of shifts per day.
+        """
+        constraint = MinSummedGradeByShiftsPerDayConstraint(
+            shifts=shifts,
+            min_grade_sum=min_grade_sum,
+            weeks=weeks,
+            days=days,
+        )
+        self.constraint_options["min_summed_grade_by_shifts_per_day"].append(constraint)
+
     def add_shift(self, shift: SingleShift) -> None:
         """Add a shift to the collection
 
@@ -3766,6 +3816,23 @@ class RotaBuilder(object):
                         week_table[week][day][shift].append(worker.get_details())
         return week_table
 
+    def get_shift_timetable_by_week(self):
+        """Build a timetable of the week as a dictionary from the model's optimal solution."""
+        works = self.model.works
+        shift_timetable = {
+            shift.name: {week: {day: [] for day in days} for week in self.weeks}
+            for shift in self.shifts
+        }
+        for shift in self.shifts:
+            for week in self.weeks:
+                for day in days:
+                    assigned_workers = []
+                    for worker in self.workers:
+                        if works[worker.id, week, day, shift.name].value > 0.5:
+                            assigned_workers.append(worker)
+                    shift_timetable[shift.name][week][day] = assigned_workers
+        return shift_timetable
+
     def get_worker_timetable(self):
         works = self.model.works
         timetable = {
@@ -4118,7 +4185,7 @@ class RotaBuilder(object):
         <summary><h2>Rota settings</h2></summary>
         <div>
         <pre>
-        {json.dumps(self.constraint_options, indent=4)}
+        {json.dumps({k: _pydantic_to_dict(v) for k, v in self.constraint_options.items()}, indent=4)}
         </pre>
         </details>
         </div>
@@ -4222,6 +4289,28 @@ class RotaBuilder(object):
 
         return shifts
 
+    def get_shift_worker_list(self, shift_name: str) -> List[set[Worker]]:
+        """
+        Returns a list of workers assigned to a specific shift.
+        
+        Args:
+            shift_name (str): The name of the shift.
+
+        Returns:
+            List[Worker]: List of workers assigned to the specified shift.
+        """
+        shift_list = []
+        for week, day in self.get_week_day_combinations():
+            workers: set = set()
+            for worker in self.workers:
+                if self.model.works[worker.id, week, day, shift_name].value > 0.5:
+                    workers.add(worker)
+            
+            shift_list.append(workers)
+
+        return shift_list
+
+
     def get_worker_shift_list(self, worker: Worker, include_locums=False, search_multiple_assignments=False) -> List:
         """
         This function returns a list of shifts assigned to a worker for each week and day.