fix a few more things

gen.py

@@ -37,7 +37,7 @@ def main(
     solve: bool = True,
     time_to_run: int = 60 * 60,
     ratio: float = 0.001,
-    start_date: datetime.datetime = "2025-03-03",
+    start_date: datetime.datetime = "2025-09-01",
     weeks: int = 26,
     bom: int = 1,
 ):
@@ -60,16 +60,16 @@ def main(
     Rota.constraint_options["hard_constrain_pair_separation"] = True
     # Rota.constraint_options["avoid_st2_first_month"] = True
 
-    wr = [
+    #wr = [
-        WorkerRequirement(
+    #    WorkerRequirement(
-            end_date=datetime.datetime.strptime("2025-06-02", "%Y-%m-%d").date(),
+    #        end_date=datetime.datetime.strptime("2025-06-02", "%Y-%m-%d").date(),
-            number=3,
+    #        number=3,
-        ),
+    #    ),
-        WorkerRequirement(
+    #    WorkerRequirement(
-            start_date=datetime.datetime.strptime("2025-06-02", "%Y-%m-%d").date(),
+    #        start_date=datetime.datetime.strptime("2025-06-02", "%Y-%m-%d").date(),
-            number=4,
+    #        number=4,
-        ),
+    #    ),
-    ]
+    #]
 
     Rota.add_shifts(
         SingleShift(
@@ -226,7 +226,7 @@ def main(
             balance_offset=3.9,
             balance_weighting=1,
             # hard_constrain_shift=False,
-            workers_required=wr,
+            workers_required=4,
             force_as_block=True,
             rota_on_nwds=True,
             constraint=[
@@ -256,7 +256,7 @@ def main(
             balance_offset=2.9,
             balance_weighting=1,
             # hard_constrain_shift=False,
-            workers_required=wr,
+            workers_required=4,
             force_as_block=True,
             rota_on_nwds=True,
             constraint=[
@@ -330,6 +330,8 @@ def main(
 
     # Rota.add_grade_constraint_by_week([2], [1, 2], ["night_weekday", "night_weekend"])
 
+    Rota.add_min_summed_grade_by_shifts_per_day_constraint(["weekend_plymouth1", "weekend_plymouth2"], 6)
+
     load_leave = True
     Rota.build_shifts()
     # Rota.add_grade_constraint_by_week([2], [1], Rota.get_shift_names())
@@ -341,10 +343,17 @@ def main(
 
         n = 0
         for worker in workers:
-            n = n + 1
             worker_name = worker
+
+            if worker_name == "Name":
+                continue
+            n = n + 1
             w = workers[worker]
-            site = w["site"]
+            try:
+                site = w["site"]
+            except KeyError:
+                print(f"Worker {worker} has no site")
+                raise KeyError
             grade = w["grade"]
             try:
                 fte = float(w["fte"]) * 100
@@ -481,7 +490,7 @@ def main(
             w = Worker(
                 name=worker_name,
                 site=site.lower(),
-                grade=int(grade[2]),
+                grade=int(grade),
                 id=n,
                 fte=int(fte),
                 nwds=nwds,

leave.py

@@ -27,7 +27,8 @@ def load_leave(Rota):
         }
 
         if live_rota:
-            download = s.get("https://docs.google.com/spreadsheets/d/e/2PACX-1vSRx9VWXSlRubPyA0RhiI-Oqf5eHNYYEc6rFzlraDbR5_8qqr5g13-4uV-gn4u-TjZxiSMv1fBUaESq/pub?gid=814517272&single=true&output=csv", headers=headers)
+            download = s.get("https://docs.google.com/spreadsheets/d/e/2PACX-1vSfcQi5Qs__A8hE2CVvUA_6ULNJALrtZOBHPUZ3xyrxPif9obtQF2IqhioT_4nebjHV1Ac5iqhNtuq4/pub?gid=2024511103&single=true&output=csv", headers=headers)
+            #download = s.get("https://docs.google.com/spreadsheets/d/e/2PACX-1vSRx9VWXSlRubPyA0RhiI-Oqf5eHNYYEc6rFzlraDbR5_8qqr5g13-4uV-gn4u-TjZxiSMv1fBUaESq/pub?gid=814517272&single=true&output=csv", headers=headers)
             decoded_content = download.content.decode('utf-8')
 
             reader = csv.reader(decoded_content.splitlines(), delimiter=',')
@@ -56,6 +57,7 @@ def load_leave(Rota):
                 names = r
                 order = names
                 for name in names:
+                    name = name.strip()
 
                     workers[name] = {}
                     workers[name]["leave"] = []
@@ -67,13 +69,16 @@ def load_leave(Rota):
                     workers[name]["start_date"] = None
                     workers[name]["pair"] = None
 
+                n = n + 1
+                continue
+
             for i in range(len(r)):
                 try:
-                    a = order[i]
+                    a = order[i].strip()
                 except IndexError:
                     continue
 
-                worker = workers[order[i]]
+                worker = workers[a]
 
                 lower_item = r[i].lower()
                 if lower_item == "derriford":
@@ -82,16 +87,18 @@ def load_leave(Rota):
                 if lower_item == "derriford twilights":
                     lower_item = "plymouth_twilights"
 
-                if "PROC site" in row_title or row_title == "PROC nights from":
+                if "PROC site" in row_title or row_title == "PROC nights from" or "proc nights" in row_title.lower():
                     worker["site_pref"] = lower_item
 
                 elif row_title in ("Rotation", "Site", "Placement", "Placement location"):
                     worker["site"] = lower_item
 
-                elif row_title in ("Grade", "Grade (ST)"):
+                elif row_title in ("Grade", "Grade (ST)") or "Year of training" in row_title:
-                    worker["grade"] = lower_item
+                    worker["grade"] = lower_item.lstrip("ST")
 
-                elif row_title in ("%FTE", "%FTE on-call", "FTE", "FTE on-call"):
+                elif row_title in ("%FTE", "%FTE on-call", "FTE", "FTE on-call") or "on-call commitment" in row_title:
+                    if lower_item.endswith("%"):
+                        lower_item = float(lower_item.rstrip("%")) / 100
                     worker["fte"] = lower_item
 
                 elif "NWD" in row_title:
@@ -102,7 +109,10 @@ def load_leave(Rota):
 
                 elif row_title in ("End Date", "CCT date"):
                     if lower_item:
-                        date = datetime.strptime(lower_item, "%d/%m/%y").date()
+                        try:
+                            date = datetime.strptime(lower_item, "%d/%m/%y").date()
+                        except ValueError:
+                            date = datetime.strptime(lower_item, "%d/%m/%Y").date()
                         worker["end_date"] = date
                     else:
                         worker["end_date"] = None
@@ -121,7 +131,7 @@ def load_leave(Rota):
                     else:
                         worker["start_date"] = None
 
-                elif "OOP" in row_title:
+                elif "OOP" in row_title or "out of programme" in row_title.lower():
                     worker["oop"] = lower_item
 
                 elif row_title == "Group":
@@ -178,6 +188,8 @@ def load_leave(Rota):
 
             n = n + 1
 
+        
+
         return workers
 
 

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.constraints.add(
-                #    self.model.works_weekend[worker.id, week]
+                    self.model.works_weekend[worker.id, week]
-                #    >= sum(
+                    >= sum(
-                #        self.model.works[worker.id, week, day, shiftname]
+                        self.model.works[worker.id, week, day, shiftname]
-                #        for w, day, shiftname in self.get_all_shiftname_combinations(
+                        for w, day, shiftname in self.get_all_shiftname_combinations(
-                #            week=week
+                            week=week
-                #        )
+                        )
-                #        if day in self.days[5:]
+                        if day in self.days[5:]
-                #    )
+                    )
-                #    / 2
+                    / 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.

rota/workers.py

@@ -237,7 +237,7 @@ class Worker(BaseModel):
                 end_oop_date = datetime.datetime.strptime(end_oop, "%d/%m/%y").date()
 
             if start_oop_date >= end_oop_date:
-                raise ValueError("End OOP date must be after start date")
+                raise ValueError(f"End OOP date must be after start date [{self.name} - Start date: {self.start_date} / End date: {self.end_date}]")
 
             # ignore oops if they finish before the rota (or worker) start date
             if end_oop_date > self.calculated_start_date:

test/rota_constraints/test_summed_shift_grade_constraints.py

@@ -0,0 +1,94 @@
+import datetime
+import pytest
+from rota.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, days, MinSummedGradeByShiftsPerDayConstraint
+
+from rota.workers import Worker
+
+import itertools
+
+
+def generate_basic_rota(weeks_to_rota=9):
+    start_date = datetime.date(2022, 3, 7)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+
+    # Add a few workers
+    Rota.add_workers(
+        [
+            Worker(name="worker1", site="group1", grade=1),
+            Worker(name="worker2", site="group1", grade=2),
+            Worker(name="worker3", site="group1", grade=3),
+            Worker(name="worker4", site="group1", grade=1),
+            Worker(name="worker5", site="group1", grade=2),
+            Worker(name="worker6", site="group1", grade=3),
+        ]
+    )
+
+    return Rota
+
+
+def test_summed_shift_grade_constraints():
+    Rota = generate_basic_rota(5)
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",), name="a", length=12.5, days=days,
+            workers_required=1,
+            force_as_block=False,
+        ),
+        SingleShift(
+            sites=("group1",), name="b", length=12.5, days=days,
+            workers_required=1,
+            force_as_block=False,
+        ),
+    )
+
+    # Add summed shift grade constraints
+    Rota.constraint_options["min_summed_grade_by_shifts_per_day"].append(
+        MinSummedGradeByShiftsPerDayConstraint(
+            shifts=["a", "b"],
+            min_grade_sum=4,
+        ),
+    )
+
+    Rota.build_and_solve(options={"ratio": 0.000})
+
+    assert Rota.results.solver.status == "ok"
+
+    Rota.export_rota_to_html("summed_shift_grade_constraints", folder="tests")
+
+    a_grades = []
+    b_grades = []
+
+    for shift, shift_timetable in Rota.get_shift_timetable_by_week().items():
+        for week, week_timetable in shift_timetable.items():
+            for day, workers in week_timetable.items():
+                if shift == "a":
+                    for worker in workers:
+                        a_grades.append(worker.grade)
+                elif shift == "b":
+                    for worker in workers:
+                        b_grades.append(worker.grade)
+
+    # Check that the minimum summed grade is respected
+    summed_grades = [a + b for a, b in zip(a_grades, b_grades)]
+    for summed_grade in summed_grades:
+        assert summed_grade >= 4, f"Summed grade {summed_grade} is less than 4"
+
+    Rota.constraint_options["min_summed_grade_by_shifts_per_day"] = [
+        MinSummedGradeByShiftsPerDayConstraint(
+            shifts=["a", "b"],
+            min_grade_sum=5,
+        ),
+    ]
+
+    Rota.build_and_solve(options={"ratio": 0.000})
+    assert Rota.results.solver.status == "ok"
+
+    for worker in Rota.get_workers():
+        if worker.grade == 1:
+            assert Rota.get_worker_shift_list(worker).count("a") == 0, f"Worker {worker.name} with grade 1 should not have shift 'a'"
+            assert Rota.get_worker_shift_list(worker).count("b") == 0, f"Worker {worker.name} with grade 1 should not have shift 'b'"

test/test_limit_constraints.py

@@ -60,17 +60,17 @@ def test_constraint_limit_grades2(limit_constraint_rota):
             balance_offset=40,
             workers_required=4,
             force_as_block=True,
-            constraint=[{"name": "limit_grade_number", "options": {3: 2}}],
+            constraint=[{"name": "limit_grade_number", "options": {3: 1}}],
         ),
     )
     Rota.constraint_options["balance_shifts_quadratic"] = True
     Rota.build_and_solve(options={"ratio": 0.1})
-    Rota.export_rota_to_html("test9")
+    Rota.export_rota_to_html("test9", folder="tests")
     grade_workers = Rota.get_workers_by_grade()
     for grade in grade_workers:
         shift_patterns = [Rota.get_worker_shift_list(w) for w in grade_workers[grade]]
         zipped_lists = list(zip(*shift_patterns))
-        limit = 2 if grade == 3 else 2
+        limit = 1 if grade == 3 else 3
         for day_shifts in zipped_lists:
             assert day_shifts.count("weekday_night") == limit
     assert Rota.results.solver.status == "ok"

test/test_weekend_padding.py

@@ -46,7 +46,7 @@ class TestDemoRota:
 
         self.Rota.build_and_solve(options={"ratio": 0.00})
 
-        self.Rota.export_rota_to_html("balance_weekend")
+        self.Rota.export_rota_to_html("balance_weekend", folder="tests")
 
         assert self.Rota.results.solver.status == "ok"
         assert self.Rota.results.solver.termination_condition == "optimal"