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)