start adding locum availability

gen.py

@@ -33,7 +33,7 @@ from rota.workers import (
 
 @app.command()
 def main(
-    suspend_on_finish: bool = False,
+    suspend: bool = False,
     solve: bool = True,
     time_to_run: int = 60 * 60,
     ratio: float = 0.001,
@@ -42,6 +42,7 @@ def main(
     bom: int = 1,
 ):
     rota_start_date = start_date.date()
+    suspend_on_finish = suspend
 
     Rota = RotaBuilder(
         rota_start_date,
@@ -50,6 +51,8 @@ def main(
         name="proc_rota",
     )
     # Rota = RotaBuilder(start_date, weeks_to_rota=20, balance_offset_modifier=1)
+    Rota.constraint_options["balance_shifts"] = False
+    Rota.constraint_options["balance_shifts_quadratic"] = True
     Rota.constraint_options["balance_weekends"] = True
     Rota.constraint_options["max_night_frequency"] = 3  # 3
     Rota.constraint_options["max_night_frequency_week_exclusions"] = []
@@ -461,6 +464,20 @@ def main(
             work_requests = w["requests"]
             site_pref = w["site_pref"]
 
+            shift_fte_overrides = {}
+
+            if worker_name == "Ben Kemp":
+                shift_fte_overrides = {
+                    "plymouth_twilight": 100,
+                    "weekend_plymouth1": 50,
+                    "weekend_plymouth2": 50,
+                }
+            elif worker_name == "Joel Lim":
+                shift_fte_overrides = {
+                    "plymouth_twilight": 100,
+                    "weekend_exeter": 50,
+                }
+
             w = Worker(
                 name=worker_name,
                 site=site.lower(),
@@ -478,6 +495,7 @@ def main(
                 pair=pair,
                 shift_balance_extra=w["shift_balance_extra"],
                 bank_holiday_extra=w["bank_holiday_extra"],
+                shift_fte_overrides=shift_fte_overrides,
             )
 
             print(w)

output/rota.csv

@@ -1,95 +0,0 @@
-"Name","Madalina Drumea","Layla Bell","Alex wood","Max Ireland","Seren Peters","Charles Finan","Paul Ward","Delilah Trimmer","Tom Welsh","Kyaw Tint","Maththew O'Brien","Max Finzel","Zainab Sharaf","Jean Skumar","Michael Tomek","Dhuvresh Patel","Nicolas Dziadulewicz","Vilim Kalamar","Georgina Edwards","Jenn Haw fong","Paul Jenkins","Amoolya Mannava","Rebecca Murphy","Ross Kruger","Fiona Lyall","Wijhdan Abusrewil","Andrew MacCormick","Joel Lim","Mark Haley","Alexander Sanchez - cabello","Chong Yew Ng","Matthew Thorley","Sayed Hashim Alqarooni","Harriet Conley","Jon Skinner","Alexander Wijnberg","Cameron Bullock","Alan Eccles","Christian Greer","Hannah Lewis","Amina Odeh","Sarath Vennam","Mo Babsail","Richard Chaytor","Sherafghan PROC"
-"Site","barnstaple","exeter","exeter","exeter","exeter","exeter","exeter","exeter","exeter","plymouth","plymouth","plymouth","plymouth","plymouth","plymouth","plymouth","plymouth","plymouth","plymouth","plymouth","plymouth","plymouth","plymouth","plymouth","plymouth_twilights","torbay","torbay","torbay","torbay","torbay","torbay","torbay","torbay","torbay","torbay twilights","truro","truro","truro","truro","truro","truro","truro","truro","truro","truro no proc"
-"Grade","3","2","3","3","3","4","4","5","5","2","2","2","2","3","3","3","3","3","4","4","4","5","5","5","6","2","2","2","2","3","4","4","4","5","6","2","2","3","3","3","4","4","5","5","5"
-"FTE","80","100","100","100","100","100","100","80","100","100","100","100","100","60","80","100","100","100","80","100","100","100","100","100","60","100","100","100","100","100","100","100","100","60","100","100","100","80","100","100","100","100","100","100","100"
-"Week 1 Day Mon","","","","","","exeter_twilight","","","","","night_weekday","","","","","","night_weekday","","","","","plymouth_twilight","","","","","","","torbay_twilight","","","","","","","","","","","","","night_weekday","","","truro_twilight"
-"Week 1 Day Tue","","","","","","","","exeter_twilight","","","night_weekday","","plymouth_twilight","","","","night_weekday","","","","","","","","","","","","","","","","","torbay_twilight","","","","","","","","night_weekday","","","truro_twilight"
-"Week 1 Day Wed","","","","","","exeter_twilight","","","","","night_weekday","","plymouth_twilight","","","","night_weekday","","","","","","","","","","","","torbay_twilight","","","","","","","","","","","","","night_weekday","truro_twilight","",""
-"Week 1 Day Thu","","","","","exeter_twilight","","","","","","night_weekday","","","","","","night_weekday","","","","plymouth_twilight","","","","","","","torbay_twilight","","","","","","","","","","","","","","night_weekday","truro_twilight","",""
-"Week 1 Day Fri","","","night_weekend","weekend_exeter","","","","","","","","","","plymouth_twilight","","night_weekend","","","","","","","","","","","","","","","","","","weekend_torbay","","","","night_weekend","","","weekend_truro","","","",""
-"Week 1 Day Sat","","","night_weekend","weekend_exeter","","","","","","","","","weekend_plymouth1","","weekend_plymouth2","night_weekend","","","","","","","","","","","","","","","","","","weekend_torbay","","","","night_weekend","","","weekend_truro","","","",""
-"Week 1 Day Sun","","","night_weekend","weekend_exeter","","","","","","","","","weekend_plymouth1","","weekend_plymouth2","night_weekend","","","","","","","","","","","","","","","","","","weekend_torbay","","","","night_weekend","","","weekend_truro","","","",""
-"Week 2 Day Mon","","","","","","","exeter_twilight","","","","","","","","","","","plymouth_twilight","night_weekday","","","","","","","","night_weekday","night_weekday","torbay_twilight","","","","","","","","","","","","","truro_twilight","","",""
-"Week 2 Day Tue","","","","","","","exeter_twilight","","","","","","","","","","plymouth_twilight","","night_weekday","","","","","","","","night_weekday","night_weekday","torbay_twilight","","","","","","","","","","","","","","truro_twilight","",""
-"Week 2 Day Wed","","","exeter_twilight","","","","","","","","","","","","","","","","night_weekday","","plymouth_twilight","","","","","","night_weekday","night_weekday","torbay_twilight","","","","","","","","","","","","","","","truro_twilight",""
-"Week 2 Day Thu","","","exeter_twilight","","","","","","","","","","","","","","","","night_weekday","","","","","","plymouth_twilight","","night_weekday","night_weekday","torbay_twilight","","","","","","","","truro_twilight","","","","","","","",""
-"Week 2 Day Fri","","","","","","night_weekend","weekend_exeter","","","","","","","","","plymouth_twilight","","night_weekend","","","","","","","","","","","","","weekend_torbay","","night_weekend","","","","","","","","","weekend_truro","","",""
-"Week 2 Day Sat","","","","","","night_weekend","weekend_exeter","","","","weekend_plymouth2","","","","","","","night_weekend","","","","","weekend_plymouth1","","","","","","","","weekend_torbay","","night_weekend","","","","","","","","","weekend_truro","","",""
-"Week 2 Day Sun","","","","","","night_weekend","weekend_exeter","","","","weekend_plymouth2","","","","","","","night_weekend","","","","","weekend_plymouth1","","","","","","","","weekend_torbay","","night_weekend","","","","","","","","","weekend_truro","","",""
-"Week 3 Day Mon","","","exeter_twilight","","","","","","","night_weekday","","","","","plymouth_twilight","night_weekday","","","","","","","","","","","","torbay_twilight","","","","","","","","night_weekday","","","","truro_twilight","","","","",""
-"Week 3 Day Tue","","","exeter_twilight","","","","","","","night_weekday","","plymouth_twilight","","","","night_weekday","","","","","","","","","","","","torbay_twilight","","","","","","","","night_weekday","","","truro_twilight","","","","","",""
-"Week 3 Day Wed","","","","","","exeter_twilight","","","","night_weekday","","","plymouth_twilight","","","night_weekday","","","","","","","","","","","","torbay_twilight","","","","","","","","night_weekday","","","","","","","truro_twilight","",""
-"Week 3 Day Thu","","","","","exeter_twilight","","","","","night_weekday","","","","","","night_weekday","","","plymouth_twilight","","","","","","","","","torbay_twilight","","","","","","","","night_weekday","","","","","","","truro_twilight","",""
-"Week 3 Day Fri","","weekend_exeter","","","","","","","","","","","","","","","","plymouth_twilight","","night_weekend","","night_weekend","","","","","","","night_weekend","weekend_torbay","","","","","","","","","","","","","","","weekend_truro"
-"Week 3 Day Sat","","weekend_exeter","","","","","","","","","","","","weekend_plymouth2","weekend_plymouth1","","","","","night_weekend","","night_weekend","","","","","","","night_weekend","weekend_torbay","","","","","","","","","","","","","","","weekend_truro"
-"Week 3 Day Sun","","weekend_exeter","","","","","","","","","","","","weekend_plymouth2","weekend_plymouth1","","","","","night_weekend","","night_weekend","","","","","","","night_weekend","weekend_torbay","","","","","","","","","","","","","","","weekend_truro"
-"Week 4 Day Mon","","","","","","exeter_twilight","","night_weekday","","","","","","","","","","","","","","","plymouth_twilight","","","","","","","","","","night_weekday","","torbay_twilight","truro_twilight","","","","night_weekday","","","","",""
-"Week 4 Day Tue","","","exeter_twilight","","","","","night_weekday","","","","","","","","plymouth_twilight","","","","","","","","","","","","","","","","torbay_twilight","night_weekday","","","","","","","night_weekday","","truro_twilight","","",""
-"Week 4 Day Wed","","","","exeter_twilight","","","","night_weekday","","","","","","","","plymouth_twilight","","","","","","","","","","","","","torbay_twilight","","","","night_weekday","","","","","","","night_weekday","","","","","truro_twilight"
-"Week 4 Day Thu","","","","","","","","night_weekday","exeter_twilight","plymouth_twilight","","","","","","","","","","","","","","","","","","","","","","","night_weekday","","torbay_twilight","","","","","night_weekday","","","","","truro_twilight"
-"Week 4 Day Fri","","","","","weekend_exeter","","","","","","","","plymouth_twilight","","","","","","","","","","","night_weekend","","","weekend_torbay","","","","","","","","","weekend_truro","","","night_weekend","","","night_weekend","","",""
-"Week 4 Day Sat","","","","","weekend_exeter","","","","","","weekend_plymouth1","","","","","","","","","","","","weekend_plymouth2","night_weekend","","","weekend_torbay","","","","","","","","","weekend_truro","","","night_weekend","","","night_weekend","","",""
-"Week 4 Day Sun","","","","","weekend_exeter","","","","","","weekend_plymouth1","","","","","","","","","","","","weekend_plymouth2","night_weekend","","","weekend_torbay","","","","","","","","","weekend_truro","","","night_weekend","","","night_weekend","","",""
-"Week 5 Day Mon","","exeter_twilight","","","","","night_weekday","","","plymouth_twilight","","","","","","","","","","","","","","","","","","","night_weekday","","","","torbay_twilight","night_weekday","","","","truro_twilight","","","","","","",""
-"Week 5 Day Tue","","","","","","","night_weekday","","exeter_twilight","","","","","","","","","","","","plymouth_twilight","","","","","","","","night_weekday","","","","torbay_twilight","night_weekday","","","","","","","","","truro_twilight","",""
-"Week 5 Day Wed","","","","","","","night_weekday","","exeter_twilight","","","","","","","","","","","","plymouth_twilight","","","","","","","","night_weekday","","","","torbay_twilight","night_weekday","","","","","","","","","","truro_twilight",""
-"Week 5 Day Thu","","","","","exeter_twilight","","night_weekday","","","","","","","","","","","","","","","plymouth_twilight","","","","","","","night_weekday","","","","torbay_twilight","night_weekday","","","","","truro_twilight","","","","","",""
-"Week 5 Day Fri","","","","night_weekend","","","","weekend_exeter","","","","","","","","","","","","","","","","plymouth_twilight","","","","night_weekend","","weekend_torbay","night_weekend","","","","","","","","","","weekend_truro","","","",""
-"Week 5 Day Sat","","","","night_weekend","","","","weekend_exeter","","","","","","","weekend_plymouth2","","","","weekend_plymouth1","","","","","","","","","night_weekend","","weekend_torbay","night_weekend","","","","","","","","","","weekend_truro","","","",""
-"Week 5 Day Sun","","","","night_weekend","","","","weekend_exeter","","","","","","","weekend_plymouth2","","","","weekend_plymouth1","","","","","","","","","night_weekend","","weekend_torbay","night_weekend","","","","","","","","","","weekend_truro","","","",""
-"Week 6 Day Mon","","night_weekday","","","exeter_twilight","","","","night_weekday","","","","","","","","","","","","","plymouth_twilight","night_weekday","","","","torbay_twilight","","","","","","","","","","","truro_twilight","","","","","","",""
-"Week 6 Day Tue","","night_weekday","","","","","exeter_twilight","","night_weekday","","","","","","","","","","","plymouth_twilight","","","night_weekday","","","","torbay_twilight","","","","","","","","","","truro_twilight","","","","","","","",""
-"Week 6 Day Wed","","night_weekday","","","","","exeter_twilight","","night_weekday","","","","","","","","","","","","","","night_weekday","plymouth_twilight","","","torbay_twilight","","","","","","","","","truro_twilight","","","","","","","","",""
-"Week 6 Day Thu","","night_weekday","","","","","","exeter_twilight","night_weekday","","","","","","","","","","","","","plymouth_twilight","night_weekday","","","","torbay_twilight","","","","","","","","","","","","","","","","","truro_twilight",""
-"Week 6 Day Fri","night_weekend","","","","","weekend_exeter","","","","","night_weekend","plymouth_twilight","","","","","plymouth_bank_holidays","","","","night_weekend","","","","","","","","","","","","","weekend_torbay","","","weekend_truro","","","","","","","",""
-"Week 6 Day Sat","night_weekend","","","","","weekend_exeter","","","","weekend_plymouth2","night_weekend","","","","","","","","","","night_weekend","","","weekend_plymouth1","","","","","","","","","","weekend_torbay","","","weekend_truro","","","","","","","",""
-"Week 6 Day Sun","night_weekend","","","","","weekend_exeter","","","","weekend_plymouth2","night_weekend","","","","","","","","","","night_weekend","","","weekend_plymouth1","","","","","","","","","","weekend_torbay","","","weekend_truro","","","","","","","",""
-"Week 7 Day Mon","","","exeter_twilight","","night_weekday","","","","","","","","","","","","","plymouth_bank_holidays","","plymouth_twilight","","","","","","","torbay_twilight","","","night_weekday","night_weekday","","","","","","","truro_twilight","","","","","","",""
-"Week 7 Day Tue","","","","","night_weekday","","exeter_twilight","","","","","","","","","plymouth_twilight","","","","","","","","","","","","","","night_weekday","night_weekday","","","","torbay_twilight","","","truro_twilight","","","","","","",""
-"Week 7 Day Wed","","","","","night_weekday","exeter_twilight","","","","","","","","","","","","plymouth_twilight","","","","","","","","","torbay_twilight","","","night_weekday","night_weekday","","","","","","","","","","truro_twilight","","","",""
-"Week 7 Day Thu","","exeter_twilight","","","night_weekday","","","","","","plymouth_twilight","","","","","","","","","","","","","","","","torbay_twilight","","","night_weekday","night_weekday","","","","","","","","truro_twilight","","","","","",""
-"Week 7 Day Fri","","","","","","","","","weekend_exeter","","","night_weekend","","","","","","","","","","night_weekend","","plymouth_twilight","","","","","weekend_torbay","","","","","","","","","","","weekend_truro","","","","night_weekend",""
-"Week 7 Day Sat","","","","","","","","","weekend_exeter","","","night_weekend","","","","","weekend_plymouth1","","","weekend_plymouth2","","night_weekend","","","","","","","weekend_torbay","","","","","","","","","","","weekend_truro","","","","night_weekend",""
-"Week 7 Day Sun","","","","","","","","","weekend_exeter","","","night_weekend","","","","","weekend_plymouth1","","","weekend_plymouth2","","night_weekend","","","","","","","weekend_torbay","","","","","","","","","","","weekend_truro","","","","night_weekend",""
-"Week 8 Day Mon","night_weekday","","","exeter_twilight","","night_weekday","","","","","","","","night_weekday","plymouth_twilight","","","","","","","","","","","","","","","","torbay_twilight","","","","","","","","truro_twilight","","","","","",""
-"Week 8 Day Tue","night_weekday","","","exeter_twilight","","night_weekday","","","","","","","","night_weekday","","","","","","","","","","plymouth_twilight","","","","","","","torbay_twilight","","","","","","","truro_twilight","","","","","","",""
-"Week 8 Day Wed","night_weekday","","","","","night_weekday","","","exeter_twilight","","","","","night_weekday","","","","","","","","","","plymouth_twilight","","","","torbay_twilight","","","","","","","","","","","","","truro_twilight","","","",""
-"Week 8 Day Thu","night_weekday","","","","exeter_twilight","night_weekday","","","","","","plymouth_twilight","","night_weekday","","","","","","","","","","","","","","torbay_twilight","","","","","","","","","","","","","truro_twilight","","","",""
-"Week 8 Day Fri","","weekend_exeter","","","","","","night_weekend","","","night_weekend","plymouth_twilight","night_weekend","","","","","","","","","","","","","","","","","","","weekend_torbay","","","","","weekend_truro","","","","","","","",""
-"Week 8 Day Sat","","weekend_exeter","","","","","","night_weekend","","","night_weekend","","night_weekend","","","","","","","","weekend_plymouth1","","","weekend_plymouth2","","","","","","","","weekend_torbay","","","","","weekend_truro","","","","","","","",""
-"Week 8 Day Sun","","weekend_exeter","","","","","","night_weekend","","","night_weekend","","night_weekend","","","","","","","","weekend_plymouth1","","","weekend_plymouth2","","","","","","","","weekend_torbay","","","","","weekend_truro","","","","","","","",""
-"Week 9 Day Mon","","","night_weekday","night_weekday","exeter_twilight","","","","","","","plymouth_twilight","","","night_weekday","","","","","","","","","","plymouth_bank_holidays","torbay_twilight","","","","","","","","","","","","","","truro_twilight","","","","",""
-"Week 9 Day Tue","","","night_weekday","night_weekday","","","","","exeter_twilight","plymouth_twilight","","","","","night_weekday","","","","","","","","","","","","","","","","torbay_twilight","","","","","","","","","","truro_twilight","","","",""
-"Week 9 Day Wed","","","night_weekday","night_weekday","","","","","exeter_twilight","","","","","","night_weekday","","","","","","","","","plymouth_twilight","","","","","","","torbay_twilight","","","","","","truro_twilight","","","","","","","",""
-"Week 9 Day Thu","","","night_weekday","night_weekday","exeter_twilight","","","","","","","","","","night_weekday","","","","","","plymouth_twilight","","","","","","","","","","torbay_twilight","","","","","","","","","","","","","truro_twilight",""
-"Week 9 Day Fri","","","","","","","weekend_exeter","","","night_weekend","","","","","","","","plymouth_twilight","","","","","","","","weekend_torbay","night_weekend","","","","","","","","","","","","","weekend_truro","night_weekend","","","",""
-"Week 9 Day Sat","","","","","","","weekend_exeter","","","night_weekend","","","","","","","weekend_plymouth1","","weekend_plymouth2","","","","","","","weekend_torbay","night_weekend","","","","","","","","","","","","","weekend_truro","night_weekend","","","",""
-"Week 9 Day Sun","","","","","","","weekend_exeter","","","night_weekend","","","","","","","weekend_plymouth1","","weekend_plymouth2","","","","","","","weekend_torbay","night_weekend","","","","","","","","","","","","","weekend_truro","night_weekend","","","",""
-"Week 10 Day Mon","","","","","","","","exeter_twilight","","","","","night_weekday","","","","","plymouth_twilight","","","","","night_weekday","","","","","","","","torbay_twilight","night_weekday","","","","","","","","","","","truro_twilight","",""
-"Week 10 Day Tue","","","exeter_twilight","","","","","","","","","","night_weekday","","","","","","","plymouth_twilight","","","night_weekday","","","","","","","","","night_weekday","","","torbay_twilight","truro_twilight","","","","","","","","",""
-"Week 10 Day Wed","","exeter_twilight","","","","","","","","plymouth_twilight","","","night_weekday","","","","","","","","","","night_weekday","","","","","","","","","night_weekday","","","torbay_twilight","","","","","","","truro_twilight","","",""
-"Week 10 Day Thu","","exeter_twilight","","","","","","","","","","","night_weekday","","","","","","","","","plymouth_twilight","night_weekday","","","","","","","","","night_weekday","","","torbay_twilight","","","","","","","","truro_twilight","",""
-"Week 10 Day Fri","","","weekend_exeter","","","","","","","plymouth_twilight","","","","","","night_weekend","","","","","","","","","","","","","","","weekend_torbay","","","","","night_weekend","night_weekend","weekend_truro","","","","","","",""
-"Week 10 Day Sat","","","weekend_exeter","","","","","","","","","weekend_plymouth1","","","","night_weekend","","weekend_plymouth2","","","","","","","","","","","","","weekend_torbay","","","","","night_weekend","night_weekend","weekend_truro","","","","","","",""
-"Week 10 Day Sun","","","weekend_exeter","","","","","","","","","weekend_plymouth1","","","","night_weekend","","weekend_plymouth2","","","","","","","","","","","","","weekend_torbay","","","","","night_weekend","night_weekend","weekend_truro","","","","","","",""
-"Week 11 Day Mon","","","","","","","exeter_twilight","","","","","","","","","","","","","","","plymouth_twilight","","night_weekday","","","","","","night_weekday","","torbay_twilight","","","","","","","night_weekday","","truro_twilight","","","",""
-"Week 11 Day Tue","","exeter_twilight","","","","","","","","","","","","","plymouth_twilight","","","","","","","","","night_weekday","","","","","","night_weekday","","torbay_twilight","","","","","","","night_weekday","","","","","","truro_twilight"
-"Week 11 Day Wed","","","exeter_twilight","","","","","","","","","","","","","","","","","plymouth_twilight","","","","night_weekday","","","","","","night_weekday","","torbay_twilight","","","","","","","night_weekday","","","truro_twilight","","",""
-"Week 11 Day Thu","","exeter_twilight","","","","","","","","","","plymouth_twilight","","","","","","","","","","","","night_weekday","","","","","","night_weekday","","torbay_twilight","","","","","","","night_weekday","","","truro_twilight","","",""
-"Week 11 Day Fri","","","","","weekend_exeter","","","","","","","","","","","","night_weekend","","","","night_weekend","plymouth_twilight","","","","","","weekend_torbay","","","","","","","","","","","","night_weekend","","","","weekend_truro",""
-"Week 11 Day Sat","","","","","weekend_exeter","","","","","","weekend_plymouth2","","","","","","night_weekend","","weekend_plymouth1","","night_weekend","","","","","","","weekend_torbay","","","","","","","","","","","","night_weekend","","","","weekend_truro",""
-"Week 11 Day Sun","","","","","weekend_exeter","","","","","","weekend_plymouth2","","","","","","night_weekend","","weekend_plymouth1","","night_weekend","","","","","","","weekend_torbay","","","","","","","","","","","","night_weekend","","","","weekend_truro",""
-"Week 12 Day Mon","","","","","","","","exeter_twilight","","plymouth_twilight","","","","","","","","","","","","","","","","night_weekday","","","","torbay_twilight","","","","","","","","","truro_twilight","","night_weekday","","night_weekday","",""
-"Week 12 Day Tue","","","","exeter_twilight","","","","","","","","","","","","","","plymouth_twilight","","","","","","","","night_weekday","","","","torbay_twilight","","","","","","","","truro_twilight","","","night_weekday","","night_weekday","",""
-"Week 12 Day Wed","","","","","","","","","exeter_twilight","","","","","","","","","","plymouth_twilight","","","","","","","night_weekday","","","","torbay_twilight","","","","","","","","","","truro_twilight","night_weekday","","night_weekday","",""
-"Week 12 Day Thu","","","","","","","","","exeter_twilight","","","","","","","","","","","","","","","","plymouth_twilight","night_weekday","","","","torbay_twilight","","","","","","","","","","","night_weekday","","night_weekday","","truro_twilight"
-"Week 12 Day Fri","","night_weekend","","weekend_exeter","","","","","","","","night_weekend","","","","","","","","plymouth_twilight","","","","","","","","","","","","night_weekend","weekend_torbay","","","","","","weekend_truro","","","","","",""
-"Week 12 Day Sat","","night_weekend","","weekend_exeter","","","","","","","","night_weekend","","weekend_plymouth1","","","","","","","","","weekend_plymouth2","","","","","","","","","night_weekend","weekend_torbay","","","","","","weekend_truro","","","","","",""
-"Week 12 Day Sun","","night_weekend","","weekend_exeter","","","","","","","","night_weekend","","weekend_plymouth1","","","","","","","","","weekend_plymouth2","","","","","","","","","night_weekend","weekend_torbay","","","","","","weekend_truro","","","","","",""
-"Week 13 Day Mon","","","","","","","exeter_twilight","","","","","","","","","","","","","plymouth_twilight","","","","","","torbay_twilight","night_weekday","","","","","","","","","","night_weekday","","","","","","truro_twilight","night_weekday",""
-"Week 13 Day Tue","","","","","","","exeter_twilight","","","","","","","","","","","","","plymouth_twilight","","","","","","torbay_twilight","night_weekday","","","","","","","","","","night_weekday","","","","truro_twilight","","","night_weekday",""
-"Week 13 Day Wed","","","","exeter_twilight","","","","","","","","plymouth_twilight","","","","","","","","","","","","","","","night_weekday","","","","","","torbay_twilight","","","","night_weekday","","","","","","","night_weekday","truro_twilight"
-"Week 13 Day Thu","","exeter_twilight","","","","","","","","","","","","","","","","","","","plymouth_bank_holidays","","plymouth_twilight","","","","night_weekday","","","","","","torbay_twilight","","","","night_weekday","","","","","truro_twilight","","night_weekday",""
-"Week 13 Day Fri","","","","","","weekend_exeter","","","night_weekend","","","","","","","","","night_weekend","plymouth_bank_holidays","night_weekend","","","plymouth_twilight","","","weekend_torbay","","","","","","","","","","","","","","","","","weekend_truro","",""
-"Week 13 Day Sat","","","","","","weekend_exeter","","","night_weekend","","","","","","","weekend_plymouth2","","night_weekend","","night_weekend","","weekend_plymouth1","","","","weekend_torbay","","","","","","","","","","","","","","","","","weekend_truro","",""
-"Week 13 Day Sun","","","","","","weekend_exeter","","","night_weekend","","","","","","","weekend_plymouth2","","night_weekend","","night_weekend","","weekend_plymouth1","","","","weekend_torbay","","","","","","","","","","","","","","","","","weekend_truro","",""

rota/shifts.py

@@ -209,6 +209,7 @@ class SingleShift(BaseModel):
             return self.workers_required
 
         for wr in self.workers_required:
+            #print(date, wr)
             if wr.start_date <= date < wr.end_date:
                 return wr.number
 
@@ -315,6 +316,12 @@ class RotaBuilder(object):
 
         self.exported_rota_file = None
 
+    def set_rota_constraint(self, constraint: str, value: Any):
+        if constraint not in self.constraint_options:
+            raise ValueError(f"Constraint {constraint} not valid")
+        
+        self.constraint_options[constraint] = value
+
     def set_rota_dates(self, start_date: datetime.date, weeks_to_rota: int):
         self.weeks_to_rota = weeks_to_rota
 
@@ -354,6 +361,10 @@ class RotaBuilder(object):
             Tuple[str | uuid.UUID | int, WeekInt, DayStr, ShiftName]
         ] = set()
         self.work_requests_map = {}
+        self.locum_availability: set[
+            Tuple[str | uuid.UUID | int, WeekInt, DayStr, ShiftName]
+        ] = set()
+        self.locum_availability_map = {}
 
     def solve_model(
         self,
@@ -836,6 +847,37 @@ class RotaBuilder(object):
             initialize=availability_init,
         )
 
+        locum_request_sets = set()
+        for locum_request in self.locum_availability:
+            if locum_request[3] == "*":
+                pass
+
+            else:
+                locum_request_sets.add(locum_request[3])
+
+        # Validate locum shifts are valid
+        invalid_shifts = locum_request_sets - set(self.get_shift_names())
+        if invalid_shifts and not self.ignore_valid_shifts:
+            raise InvalidShift(
+                f"Invalid locum shift request [shift(s): ${invalid_shifts}]"
+            )
+
+
+        def locum_request_init(model, wid, week, day, shift):
+            if (wid, week, day, shift) in self.locum_availability:
+                self.locum_availability_map[(wid, week, day)] = shift
+                return 1
+            return 0
+
+        self.model.locum_availability = Param(
+            (
+                (worker.id, week, day, shift)
+                for worker in self.workers
+                for week, day, shift in self.get_all_shiftname_combinations()
+            ),
+            initialize=locum_request_init,
+        )
+
         work_request_sets = set()
         for work_request in self.work_requests:
             if work_request[3] == "*":
@@ -845,7 +887,6 @@ class RotaBuilder(object):
                 work_request_sets.add(work_request[3])
 
         # Validate shifts are valid
-        # work_request_sets = set([i[3] for i in self.work_requests])
         invalid_shifts = work_request_sets - set(self.get_shift_names())
         if invalid_shifts and not self.ignore_valid_shifts:
             raise InvalidShift(
@@ -1190,6 +1231,7 @@ class RotaBuilder(object):
                 workers_required = shift.get_worker_requirement_by_date(
                     self.get_week_start_date(week)
                 )
+                print(worker.name, week, workers_required, self.get_week_start_date(week))
 
                 if shift.force_as_block_unless_nwd:
                     workers = self.get_workers_for_shift(shift)
@@ -1376,11 +1418,11 @@ class RotaBuilder(object):
                     )
 
                     full_time_equivalent_joined = sum(
-                        self.full_time_equivalent_sites[i] for i in shift.sites
+                        self.full_time_equivalent_sites[i][shift.name] for i in shift.sites
                     )
 
                     target_shifts = (
-                        total_shifts / full_time_equivalent_joined * worker.fte_adj
+                        total_shifts / full_time_equivalent_joined * worker.get_fte(shift=shift.name)
                     )
 
                     if self.use_previous_shifts:
@@ -1408,7 +1450,7 @@ class RotaBuilder(object):
 
                     if shift.hard_constrain_shift:
                         adjusted_balance_offset = shift.balance_offset
-                        if worker.fte_adj < 100:
+                        if worker.get_fte(shift=shift.name) < 100:
                             adjusted_balance_offset = (
                                 adjusted_balance_offset * self.ltft_balance_offset
                             )
@@ -2412,6 +2454,8 @@ class RotaBuilder(object):
             # c = len(workers)
 
             balance_modifier_constant = 1
+            balance_quadratic_shift_modifier_constant = 3
+            block_shift_balancing_constant = 1
 
             if self.constraint_options["balance_shifts_over_workers"]:
                 worker_shift_balancing = sum(
@@ -2440,7 +2484,7 @@ class RotaBuilder(object):
 
             if self.constraint_options["balance_shifts_quadratic"]:
                 quadratic_shift_balancing = sum(
-                    balance_modifier_constant
+                    balance_quadratic_shift_modifier_constant
                     * (
                         self.model.shift_count_t1[worker.id, shift.name]
                         + self.model.shift_count_t2[worker.id, shift.name]
@@ -2550,7 +2594,7 @@ class RotaBuilder(object):
 
             if self.constraint_options["balance_blocks"]:
                 blocks_balancing = sum(
-                    1 * self.model.blocks_assigned[week, shift]
+                    block_shift_balancing_constant * self.model.blocks_assigned[week, shift]
                     for week in self.weeks
                     for shift in self.shifts_to_assign_as_blocks()
                 )
@@ -2646,14 +2690,18 @@ class RotaBuilder(object):
 
         self.workers = sorted(self.workers)
 
-        self.full_time_equivalent = sum(w.fte_adj for w in self.workers)
+        self.full_time_equivalent = sum(w.fte_adj for w in self.workers) # TODO: rewrite as per shift
         self.full_time_equivalent_sites = {}
         self.workers_at_sites = {}
 
         for site in track(self.sites, description="Generating site workers"):
             self.workers_at_sites[site] = [w for w in self.workers if w.site == site]
-            self.full_time_equivalent_sites[site] = sum(
+
-                w.fte_adj for w in self.workers if w.site == site
+            self.full_time_equivalent_sites[site] = {}
+
+            for shift in self.get_shifts():
+                self.full_time_equivalent_sites[site][shift.name] = sum(
+                    w.get_fte(shift=shift.name) for w in self.workers if w.site == site
                 )
 
         self.worker_pairs = []
@@ -3131,16 +3179,18 @@ class RotaBuilder(object):
         return [worker for worker in self.workers if worker.site in shift.sites]
 
     def get_workers_total_fte(self) -> float:
+        """Does not take into account shift adjusted ftes"""
         return self.full_time_equivalent
 
     def get_worker_details(self):
-        w = defaultdict(list)
+        w: dict[str, Worker] = defaultdict(list)
+        worker: Worker
         for worker in self.workers:
             w[worker.site].append(worker)
 
         l = []
         for site in w:
-            l.append(f"{site}: {sum(worker.fte_adj for worker in w[site]) / 100}")
+            l.append(f"{site}: {sum(worker.get_fte() for worker in w[site]) / 100}")
         t = "\n".join(l)
         return f"Full time equivalent trainees by site:\n{t}"
 
@@ -3158,6 +3208,8 @@ class RotaBuilder(object):
         ]
 
     def get_week_start_date(self, week: int):
+        """Week start date is 1 indexed"""
+        week = week - 1
         return self.start_date + datetime.timedelta(weeks=week)
 
     # RESULTS
@@ -3454,7 +3506,7 @@ class RotaBuilder(object):
                 worker_id=worker.id,
                 name=worker.name,
                 fte=worker.fte,
-                fte_adj=worker.fte_adj,
+                fte_adj=worker.get_fte(),
                 start_date=worker.calculated_start_date,
                 end_date=worker.calculated_end_date,
                 oops=json.dumps(json.dumps(worker.oop, default=str)),

rota/workers.py

@@ -87,7 +87,7 @@ class Worker(BaseModel):
     site: str
     # Grade are equivalent to roles, by keeping them integer defining model
     # rules is easier.
-    grade: int
+    grade: int = 1
     # We can either have a user generated ID
     id: Optional[int | uuid.UUID | str] = None
     fte: int = 100
@@ -98,11 +98,15 @@ class Worker(BaseModel):
     not_available_to_work: list[NotAvailableToWork] = []
     pref_not_to_work: list[PreferenceNotToWork] = []
     work_requests: list[WorkRequests] = []
+    locum_availability: list[WorkRequests] = []
     remote_site: str = "plymouth"  # We set a default proc_site
     previous_shifts: dict = {}
     shift_balance_extra: dict = {}
     bank_holiday_extra: int = 0
     pair: int | str | None = None
+    locum: bool = False
+
+    shift_fte_overrides: dict[str, int] = {} # Need checks to ensure shifts exist
 
     model_config = ConfigDict(
         extra="allow",
@@ -284,6 +288,16 @@ class Worker(BaseModel):
             else:
                 Rota.work_requests.add((self.id, week, day, request.shift))
 
+        for request in self.locum_availability:
+            days_from_start = (request.date - Rota.start_date).days
+            week = days_from_start // 7 + 1
+            day = Rota.days[(days_from_start % 7)]
+            if request.shift == "*":
+                for shift in Rota.get_shifts_for_worker_site(self.site):
+                    Rota.locum_availability.add((self.id, week, day, shift.name))
+            else:
+                Rota.locum_availability.add((self.id, week, day, request.shift))
+
         # Calculate the proportion of the rota that is being worked
         self.proportion_rota_to_work = days_to_work / Rota.rota_days_length
         self.days_to_work = days_to_work
@@ -291,6 +305,12 @@ class Worker(BaseModel):
         # We have to adjust the full time equivalent for people who CCT / leave the rota early
         self.fte_adj = self.fte * self.proportion_rota_to_work
 
+        self.fte_adj_shifts = {}
+        if self.shift_fte_overrides:
+            for shift, fte in self.shift_fte_overrides.items():
+                self.fte_adj_shifts[shift] = fte * self.proportion_rota_to_work
+
+
         if self.fte_adj > 100:
             # Shouldn't happen ? bug if it does
             Rota.add_warning(
@@ -356,3 +376,9 @@ class Worker(BaseModel):
 
     def get_shift_targets(self) -> dict:
         return self.shift_target_number
+
+    def get_fte(self, shift: str | None=None) -> int:
+        if shift is not None:
+            if shift in self.fte_adj_shifts:
+                return self.fte_adj_shifts[shift]
+        return self.fte_adj

test/test_workers.py

@@ -1,11 +1,13 @@
 import pytest
-from rota.shifts import NoWorkers, RotaBuilder, SingleShift, WarningTermination, days
+from rota.shifts import NoWorkers, RotaBuilder, SingleShift, WarningTermination, WorkerRequirement, days
 
 import datetime
 from rota.workers import NotAvailableToWork, Worker, generate_not_available_to_works
 
-def generate_basic_rota(weeks_to_rota=10, workers=0, start_date=datetime.date(2022, 3, 7)):
 
+def generate_basic_rota(
+    weeks_to_rota=10, workers=0, start_date=datetime.date(2022, 3, 7)
+):
     Rota = RotaBuilder(
         start_date,
         weeks_to_rota=weeks_to_rota,
@@ -15,12 +17,14 @@ def generate_basic_rota(weeks_to_rota=10, workers=0, start_date=datetime.date(20
     if workers:
         Rota.add_workers(
             [
-                Worker(name=f"worker{i}", site="group1", grade=1) for i in range(1, workers+1)
+                Worker(name=f"worker{i:02}", site="group1", grade=1)
+                for i in range(1, workers + 1)
             ]
         )
 
     return Rota
 
+
 class TestWorkers:
     def test_no_workers(self):
         """Check error raised if attempt to solve with no workers"""
@@ -48,7 +52,6 @@ class TestWorkers:
             Rota.build_and_solve(solve=False)
 
     def test_workers_duplicate_name(self):
-
         Rota = generate_basic_rota()
 
         workers = []
@@ -74,7 +77,6 @@ class TestWorkers:
             Rota.build_and_solve(solve=False)
 
     def test_worker_ftes(self):
-
         weeks_to_rota = 10
         start_date = datetime.date(2022, 3, 7)
 
@@ -157,7 +159,6 @@ class TestWorkers:
 
         # All workers should have an adjusted fte of 50 %
         for worker in Rota.get_workers():
-
             # Except those in group3
             if worker.site == "group3":
                 assert worker.fte_adj == 0
@@ -404,7 +405,6 @@ class TestWorkers:
         assert Rota.get_workers_total_shifts()["worker3"] == 50
 
     def test_not_available_to_work(self):
-        
         Rota = generate_basic_rota()
 
         Rota.add_shifts(
@@ -419,14 +419,16 @@ class TestWorkers:
 
         natws = generate_not_available_to_works(Rota.start_date, days=21)
 
-        natws.append(NotAvailableToWork(date=Rota.rota_end_date-datetime.timedelta(days=1)))
+        natws.append(
-
+            NotAvailableToWork(date=Rota.rota_end_date - datetime.timedelta(days=1))
-
+        )
 
         Rota.add_workers(
             [
                 Worker(name="worker1", site="group1", grade=1),
-                Worker(name="worker2", site="group1", grade=1, not_available_to_work=natws),
+                Worker(
+                    name="worker2", site="group1", grade=1, not_available_to_work=natws
+                ),
             ]
         )
 
@@ -438,5 +440,307 @@ class TestWorkers:
         assert Rota.get_workers_total_shifts()["worker1"] == 35
         assert Rota.get_workers_total_shifts()["worker2"] == 35
 
-        assert Rota.get_worker_shift_list_string(Rota.get_worker_by_name("worker2")).startswith("-" * 21)
+        assert Rota.get_worker_shift_list_string(
-        assert Rota.get_worker_shift_list_string(Rota.get_worker_by_name("worker2")).endswith("-")
+            Rota.get_worker_by_name("worker2")
+        ).startswith("-" * 21)
+        assert Rota.get_worker_shift_list_string(
+            Rota.get_worker_by_name("worker2")
+        ).endswith("-")
+
+
+class TestWorkerFtes:
+    def test_worker_ftes(self):
+        Rota = generate_basic_rota(workers=2)
+
+        Rota.add_workers(
+            [
+                Worker(
+                    name="worker3",
+                    site="group1",
+                    grade=1,
+                    fte=60,
+                ),
+                Worker(
+                    name="worker4",
+                    site="group1",
+                    grade=1,
+                    fte=40,
+                ),
+            ]
+        )
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"),
+                name="a",
+                length=12.5,
+                days=days[:6],
+                workers_required=1,
+            ),
+        )
+
+        Rota.build_and_solve()
+
+        Rota.export_rota_to_html("test_worker_ftes")
+
+        total_shifts = Rota.get_workers_total_shifts()
+        for worker in Rota.get_workers():
+            match worker.name:
+                case "worker1":
+                    assert worker.get_fte() == 100
+                    assert total_shifts[worker.name] == 20
+                case "worker2":
+                    assert worker.get_fte() == 100
+                    assert total_shifts[worker.name] == 20
+                case "worker3":
+                    assert worker.get_fte() == 60
+                    assert total_shifts[worker.name] == 12
+                case "worker4":
+                    assert worker.get_fte() == 40
+                    assert total_shifts[worker.name] == 8
+
+    def test_worker_fte_shift_override(self):
+        Rota = generate_basic_rota(workers=1)
+
+        Rota.add_workers(
+            [
+                Worker(
+                    name="worker2",
+                    site="group1",
+                    grade=1,
+                    fte=100,
+                    shift_fte_overrides={"b": 50}
+                ),
+            ]
+        )
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"),
+                name="a",
+                length=12.5,
+                days=days[:2],
+                workers_required=1,
+            ),
+            SingleShift(
+                sites=("group1", "group2"),
+                name="b",
+                length=12.5,
+                days=days[2:5],
+                workers_required=1,
+            ),
+        )
+
+        Rota.build_and_solve()
+
+        Rota.export_rota_to_html("test_worker_fte_shift_override")
+
+        total_shifts = Rota.get_workers_total_shifts()
+        for worker in Rota.get_workers():
+            match worker.name:
+                case "worker1":
+                    assert worker.get_fte() == 100
+                    assert total_shifts[worker.name] == 30
+                    assert worker.get_fte("b") == 100
+                case "worker2":
+                    assert worker.get_fte() == 100
+                    assert worker.get_fte("b") == 50
+                    assert total_shifts[worker.name] == 20
+
+class TestComplexWorkers:
+    def test_worker_requirement_and_force_block(self):
+        Rota = generate_basic_rota(workers=10, weeks_to_rota=10)
+
+        wr = [
+            WorkerRequirement(
+                end_date=Rota.start_date + datetime.timedelta(weeks=5),
+                number=2,
+            ),
+            WorkerRequirement(
+                start_date=Rota.start_date + datetime.timedelta(weeks=5),
+                number=3,
+            ),
+        ]
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"),
+                name="a",
+                length=12.5,
+                days=days[:2],
+                workers_required=1,
+            ),
+            SingleShift(
+                sites=("group1", "group2"),
+                name="b",
+                length=12.5,
+                days=days[2:5],
+                workers_required=1,
+            ),
+            SingleShift(
+                sites=("group1", "group2"),
+                name="c",
+                length=12.5,
+                days=days[:4],
+                workers_required=wr,
+                force_as_block=True,
+                constraint=[
+                {"name": "pre", "options": 2},
+                {"name": "post", "options": 3},
+                ]
+            ),
+            SingleShift(
+                sites=("group1", "group2"),
+                name="d",
+                length=12.5,
+                days=days[4:],
+                workers_required=wr,
+                force_as_block=True,
+                constraint=[
+                {"name": "pre", "options": 2},
+                {"name": "post", "options": 3},
+                ]
+            ),
+        )
+
+        Rota.build_and_solve()
+
+        Rota.export_rota_to_html("test_worker_complex")
+
+        #total_shifts = Rota.get_workers_total_shifts()
+        #for worker in Rota.get_workers():
+        #    match worker.name:
+        #        case "worker1":
+        #            assert worker.get_fte() == 100
+        #            assert total_shifts[worker.name] == 30
+        #            assert worker.get_fte("b") == 100
+        #        case "worker2":
+        #            assert worker.get_fte() == 100
+        #            assert worker.get_fte("b") == 50
+        #            assert total_shifts[worker.name] == 20
+        #total_shifts = Rota.get_workers_total_shifts()
+        for worker in Rota.get_workers():
+            shift_string = Rota.get_worker_shift_list_string(worker)
+            assert shift_string.count("cccc") == shift_string.count("c") / 4
+            assert shift_string.count("ddd") == shift_string.count("d") / 3
+
+    def test_split_shift(self):
+        Rota = generate_basic_rota(workers=11, weeks_to_rota=22)
+
+        Rota.set_rota_constraint("max_shifts_per_week", 2)
+        Rota.set_rota_constraint("max_weekend_frequency", 3)
+        Rota.set_rota_constraint("max_shifts_per_month", 4)
+
+
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"),
+                name="a",
+                length=12.5,
+                days=days[:4],
+                workers_required=1,
+                constraint=[
+                {"name": "pre", "options": 2},
+                {"name": "post", "options": 3},
+                {"name": "max_shifts_per_week", "options": 1},
+                ]
+            ),
+            SingleShift(
+                sites=("group1", "group2"),
+                name="b",
+                length=12.5,
+                days=(days[4], days[6]),
+                workers_required=1,
+                assign_as_block=True,
+                constraint=[
+                {"name": "pre", "options": 2},
+                {"name": "post", "options": 3},
+                ]
+            ),
+            SingleShift(
+                sites=("group1", "group2"),
+                name="c",
+                length=12.5,
+                days=days[5],
+                workers_required=1,
+                constraint=[
+                {"name": "pre", "options": 2},
+                {"name": "post", "options": 3},
+                ]
+            ),
+        )
+
+        Rota.build_and_solve()
+
+        Rota.export_rota_to_html("test_split")
+
+
+class TestLocums:
+    def test_locums(self):
+        Rota = generate_basic_rota(workers=9, weeks_to_rota=22)
+
+        Rota.set_rota_constraint("max_shifts_per_week", 2)
+        Rota.set_rota_constraint("max_weekend_frequency", 3)
+        Rota.set_rota_constraint("max_shifts_per_month", 4)
+
+        Rota.terminate_on_warning.remove("Worker/no valid shifts")
+
+        Rota.add_workers(
+            [
+                Worker(name="locum1", site="group1", locum=True),
+                Worker(
+                    name="locum2", site="group1", locum=True
+                ),
+                Worker(
+                    name="extra1", site="locum"
+                ),
+                Worker(
+                    name="extra2", site="locum"
+                ),
+
+            ]
+        )
+
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"),
+                name="a",
+                length=12.5,
+                days=days[:4],
+                workers_required=1,
+                constraint=[
+                {"name": "pre", "options": 2},
+                {"name": "post", "options": 3},
+                {"name": "max_shifts_per_week", "options": 1},
+                ]
+            ),
+            SingleShift(
+                sites=("group1", "group2"),
+                name="b",
+                length=12.5,
+                days=(days[4], days[6]),
+                workers_required=1,
+                assign_as_block=True,
+                constraint=[
+                {"name": "pre", "options": 2},
+                {"name": "post", "options": 3},
+                ]
+            ),
+            SingleShift(
+                sites=("group1", "group2"),
+                name="c",
+                length=12.5,
+                days=days[5],
+                workers_required=1,
+                constraint=[
+                {"name": "pre", "options": 2},
+                {"name": "post", "options": 3},
+                ]
+            ),
+        )
+
+        Rota.build_and_solve()
+
+        Rota.export_rota_to_html("test_locums", timestamp_filename=False)