From 05edb93427a95d031b85da20d65f9a51d50d6a99 Mon Sep 17 00:00:00 2001
From: Ross <ross@xkjq.uk>
Date: Tue, 9 Dec 2025 13:41:26 +0000
Subject: [PATCH] start improving constraint options

---
 rota_generator/shifts.py | 44 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/rota_generator/shifts.py b/rota_generator/shifts.py
index 6e9aab1..74be937 100644
--- a/rota_generator/shifts.py
+++ b/rota_generator/shifts.py
@@ -321,6 +321,50 @@ class SingleShift(BaseModel):
         return self.name[0]
 
 
+class RotaConstraintOptions:
+    """Class to hold rota-level constraint options"""
+
+    def __init__(self):
+        # self.options: { name : (default_value, type, description) }
+        # For each option we keep a tuple: (default_value, python_type, description)
+        # The types are informational and can be used when generating forms/docs.
+        self.constraint_options = {
+            "balance_nights": (True, bool, "Balance night shifts across workers"),
+            "constrain_time_off_after_nights": (False, bool, "Apply time-off constraints after night shifts"),
+            "balance_nights_across_sites": (True, bool, "Balance night shifts across sites"),
+            "balance_bank_holidays": (True, bool, "Balance bank-holiday shifts across staff"),
+            "balance_blocks": (True, bool, "Balance shifts across week-blocks"),
+            "balance_shifts": (True, bool, "Balance shifts across workers"),
+            "balance_shifts_true_quadratic": (False, bool, "Use true quadratic balancing for shifts"),
+            "balance_shifts_quadratic": (False, bool, "Use quadratic penalty to discourage spread across shifts"),
+            "balance_shifts_over_workers": (True, bool, "Balance shift distribution over workers"),
+            "minimise_shift_diffs": (False, bool, "Use simpler minimisation to reduce shift differences"),
+            "maximum_allowed_shift_diff": (None, Optional[int], "Maximum allowed per-worker shift difference (None = no limit)"),
+            "balance_weekends": (True, bool, "Balance weekend allocations"),
+            "max_weekends": (100, int, "Maximum number of weekends considered"),
+            "max_weekend_frequency": (1, int, "Don't assign multiple shifts every N weeks (requires balance_weekends)"),
+            "max_night_frequency": (1, int, "Don't assign multiple night shifts every N weeks"),
+            "max_night_frequency_week_exclusions": ([], list, "Weeks to exclude from night frequency checks"),
+            "max_shifts_per_week": (7, int, "Max shifts per worker per week"),
+            "max_shifts_per_month": (40, int, "Max shifts per worker per month"),
+            "max_days_per_week_block": ([], list, "Optional list of (max_days, week_block) tuples to limit days in blocks"),
+            "prevent_monday_after_full_weekends": ([], list, "Prevent assignment on Monday after full weekends"),
+            "prevent_monday_and_tuesday_after_full_weekends": ([], list, "Prevent Monday/Tuesday after full weekends"),
+            "prevent_fridays_before_full_weekends": ([], list, "Prevent Fridays before full weekends"),
+            "prevent_thursdays_before_full_weekends": ([], list, "Prevent Thursdays before full weekends"),
+            "hard_constrain_pair_separation": (False, bool, "Enforce hard pair separation constraint"),
+            "avoid_shifts_by_grades": ([], list, "List of (grade, shift_name) pairs or rules to avoid assigning a grade to a shift"),
+            "avoid_shifts_by_worker_names": ([], list, "List of worker names (or patterns) to avoid assigning to particular shifts"),
+            "allocate_locum_shifts": (True, bool, "Allow allocation of locum shifts"),
+            "distribute_locum_shifts": (True, bool, "Distribute locum shifts among available workers"),
+            "balance_locum_shifts": (False, bool, "Balance locum shifts similarly to normal shifts"),
+            "maximum_allowed_locum_shifts_per_worker": (None, Optional[int], "Maximum locum shifts allowed per worker (None = no limit)"),
+            "minimum_allowed_locum_shifts_per_worker": (None, Optional[int], "Minimum locum shifts required per worker (None = no minimum)"),
+            # This is a typed list of Pydantic constraint models; items should follow
+            # the schema of MinSummedGradeByShiftsPerDayConstraint
+            "min_summed_grade_by_shifts_per_day": ([], List[MinSummedGradeByShiftsPerDayConstraint], "List of minimum-summed-grade constraints per day"),
+        }
+
 class RotaBuilder(object):
     """Class to hold and manipulate shifts"""