many improvements

run.py

@@ -47,7 +47,7 @@ Rota.constraint_options["constrain_time_off_after_nights"] = True
 Rota.constraint_options["balance_nights_across_sites"] = True
 Rota.constraint_options["balance_blocks"] = True
 Rota.constraint_options["balance_weekends"] = True
-Rota.constraint_options["balance_shifts"] = True
+Rota.constraint_options["balance_shifts_over_workers"] = True
 Rota.constraint_options["balance_bank_holidays"] = True
 Rota.constraint_options["avoid_st2_first_month"] = False
 Rota.constraint_options["hard_constrain_pair_separation"] = True

runfixnights.py

@@ -47,7 +47,7 @@ Rota.constraint_options["constrain_time_off_after_nights"] = True
 Rota.constraint_options["balance_nights_across_sites"] = True
 Rota.constraint_options["balance_blocks"] = True
 Rota.constraint_options["balance_weekends"] = True
-Rota.constraint_options["balance_shifts"] = True
+Rota.constraint_options["balance_shifts_over_workers"] = True
 Rota.constraint_options["balance_bank_holidays"] = True
 Rota.constraint_options["avoid_st2_first_month"] = False
 Rota.constraint_options["hard_constrain_pair_separation"] = True

shifts.py

@@ -1,6 +1,6 @@
 import datetime
 import itertools
-from typing import Iterable, List, Sequence, Tuple
+from typing import Iterable, List, Sequence, Tuple, Set
 
 import datetime
 
@@ -104,12 +104,12 @@ class RotaBuilder(object):
 
     def __init__(
         self,
-        start_date: datetime.datetime,
+        start_date: datetime.datetime = datetime.datetime.now(),
         weeks_to_rota: int = 26,
         balance_offset_modifier: int = 1,
         ltft_balance_offset: int = 1,
         max_night_frequency: int = 2,
-        max_weekend_frequency: int = 2,  # Requires balance weekends
+        max_weekend_frequency: int = 1,  # Requires balance weekends
         use_previous_shifts: bool = False,
         use_shift_balance_extra: bool = False,
         use_bank_holiday_extra: bool = False,
@@ -171,11 +171,12 @@ class RotaBuilder(object):
             "balance_bank_holidays": True,
             "balance_blocks": True,
             "balance_shifts": True,
-            "minimise_shift_diffs": False,  # less sophisticated version of balance_shifts
+            "balance_shifts_over_workers": True,
+            "minimise_shift_diffs": False,  # less sophisticated version of balance_shifts_over_workers
             "balance_weekends": True,
             "max_weekends": 100,
-            "max_shifts_per_week": 4,
-            "max_shifts_per_month": 12,
+            "max_shifts_per_week": 7,
+            "max_shifts_per_month": 40,
             "prevent_monday_after_full_weekends": [],
             "prevent_monday_and_tuesday_after_full_weekends": [],
             "prevent_fridays_before_full_weekends": [],
@@ -226,6 +227,14 @@ class RotaBuilder(object):
         if not results.solver.status:
             sys.exit(0)
 
+    def build_and_solve(self, options={"ratio": 0.1, "seconds": 1000, "threads": 10}):
+        self.build_shifts()
+        self.build_workers()
+        self.build_model()
+
+        self.solve_model(options=options)
+
+
     def build_model(self):
         # Initialize model
         self.model = ConcreteModel()
@@ -234,10 +243,9 @@ class RotaBuilder(object):
         # NOTE: this will assign to all possible shift combinations (which we probably don't want)
         self.model.works = Var(
             (
-                (worker.id, week, day, shift)
+                (worker.id, week, day, shiftname)
                 for worker in self.workers
-                for week, day in self.get_week_day_combinations()
-                for shift in self.get_shift_names()
+                for week, day, shiftname in self.get_all_shiftname_combinations()
             ),
             within=Binary,
             initialize=0,
@@ -352,6 +360,7 @@ class RotaBuilder(object):
                 initialize=0,
             )
 
+
         self.model.shift_count = Var(
             (
                 (worker.id, shift)
@@ -379,20 +388,40 @@ class RotaBuilder(object):
             initialize=0,
         )
 
-        if self.constraint_options["balance_shifts"]:
-
-            self.model.shift_count_t1 = Var(
+        if self.constraint_options["balance_shifts_over_workers"]:
+            self.model.worker_shift_count_t1 = Var(
                 ((worker.id) for worker in self.workers),
                 domain=NonNegativeReals,
                 initialize=0,
             )
 
-            self.model.shift_count_t2 = Var(
+            self.model.worker_shift_count_t2 = Var(
                 ((worker.id) for worker in self.workers),
                 domain=NonNegativeReals,
                 initialize=0,
             )
 
+        if self.constraint_options["balance_shifts"]:
+            self.model.shift_count_t1 = Var(
+                (
+                    (worker.id, shift.name)
+                    for worker in self.workers
+                    for shift in self.get_shifts()
+                ),
+                domain=NonNegativeReals,
+                initialize=0,
+            )
+
+            self.model.shift_count_t2 = Var(
+                (
+                    (worker.id, shift.name)
+                    for worker in self.workers
+                    for shift in self.get_shifts()
+                ),
+                domain=NonNegativeReals,
+                initialize=0,
+            )
+
         if self.constraint_options["balance_nights"]:
             # We also try to even out the night shifts seperately
             self.model.night_shift_count = Var(
@@ -453,7 +482,6 @@ class RotaBuilder(object):
         #     initialize=0,
         # )
 
-        if self.constraint_options["balance_weekends"]:
 
         self.model.works_weekend = Var(
             ((worker.id, week) for worker in self.workers for week in self.weeks),
@@ -468,6 +496,7 @@ class RotaBuilder(object):
             bounds=SHIFT_BOUNDS["weekend_count"],
         )
 
+        if self.constraint_options["balance_weekends"]:
             self.model.weekend_shift_count_t1 = Var(
                 ((worker.id) for worker in self.workers),
                 domain=NonNegativeReals,
@@ -591,15 +620,16 @@ class RotaBuilder(object):
             #     self.model.works[worker.id, week, day, shift]
             #     for worker in self.workers if worker.site not in site_required))
 
+        # THIS SHOULD NO LONGER BE REQUIRED
         # Ensure no workers are assigned to shifts that are not required
-        for week, day, shift in self.get_not_required_shifts():
-            self.model.constraints.add(
-                0
-                == sum(
-                    self.model.works[worker.id, week, day, shift.name]
-                    for worker in self.workers
-                )
-            )
+        #for week, day, shift in self.get_not_required_shifts():
+        #    self.model.constraints.add(
+        #        0
+        #        == sum(
+        #            self.model.works[worker.id, week, day, shift.name]
+        #            for worker in self.workers
+        #        )
+        #    )
 
         # Constraint: total hours worked hours worked
         # for worker in self.workers:
@@ -820,7 +850,7 @@ class RotaBuilder(object):
                                 ]
                                 for worker in self.workers
                             )
-                            <= 1
+                            <= shift.workers_required
                         )
 
         # Most of our constraints apply per worker
@@ -840,10 +870,9 @@ class RotaBuilder(object):
                 self.model.constraints.add(
                     self.constraint_options["max_shifts_per_month"]
                     >= sum(
-                        self.model.works[worker.id, week, day, shift.name]
-                        for week in week_blocks
-                        for day in self.days
-                        for shift in self.get_shifts()
+                        self.model.works[worker.id, week, day, shiftname]
+                        for week, day, shiftname in self.get_all_shiftname_combinations()
+                        if week in week_blocks
                     )
                 )
 
@@ -862,8 +891,9 @@ class RotaBuilder(object):
                 except ValueError as e:
                     print("Failure setting constraint", "avoid_st2_first_month")
                     print(e)
+
             # Count number of weekends an worker works
-            if self.constraint_options["balance_weekends"]:
+            #if self.constraint_options["balance_weekends"]:
             self.model.constraints.add(
                 self.model.worker_weekend_count[worker.id]
                 == sum(
@@ -933,7 +963,9 @@ class RotaBuilder(object):
                                 min_shifts,
                                 sum(
                                     self.model.works[worker.id, week, day, shift.name]
-                                    for week, day in self.get_week_day_combinations()
+                                    for week, day in self.get_week_day_combinations_for_shift(shift)
+                                    #for week, day in self.get_week_day_combinations()
+                                    #if shift.name in self.get_shift_names_by_week_day(week, day)
                                 ),
                                 max_shifts,
                             )
@@ -944,7 +976,8 @@ class RotaBuilder(object):
                         0
                         == sum(
                             self.model.works[worker.id, week, day, shift.name]
-                            for week, day in self.get_week_day_combinations()
+                            for week, day in self.get_week_day_combinations_for_shift(shift)
+                            #for week, day in self.get_week_day_combinations()
                         )
                     )
 
@@ -954,7 +987,7 @@ class RotaBuilder(object):
                     self.model.shift_count[worker.id, shift.name]
                     == sum(
                         self.model.works[worker.id, week, day, shift.name]
-                        for week, day in self.get_week_day_combinations()
+                        for week, day in self.get_week_day_combinations_for_shift(shift)
                     )
                 )
 
@@ -964,16 +997,23 @@ class RotaBuilder(object):
                     - worker.shift_target_number[shift.name]
                 )
 
-            # Define shift_count_t1 and shift_count_t2 constraints for the object
+                if self.constraint_options["balance_shifts"]:
+                    self.model.constraints.add(
+                        self.model.shift_count_t1[worker.id, shift.name]
+                        - self.model.shift_count_t2[worker.id, shift.name]
+                        == self.model.shift_count_diff[worker.id, shift.name]
+                    )
+
+            # Define worker_shift_count_t1 and worker_shift_count_t2 constraints for the object
             # Thus bypassing the need for a quadratic solver
             # t1-t2 is the target
             # As the objective is to minimise t1+t2 and t1 and t2 are positive reals
             # t1+t2 approximates the absolute target (which otherwise requires a quadratic solver)
             # TODO: quadratic implementation so perfect solutions will be chosen
-            if self.constraint_options["balance_shifts"]:
+            if self.constraint_options["balance_shifts_over_workers"]:
                 self.model.constraints.add(
-                    self.model.shift_count_t1[worker.id]
-                    - self.model.shift_count_t2[worker.id]
+                    self.model.worker_shift_count_t1[worker.id]
+                    - self.model.worker_shift_count_t2[worker.id]
                     == sum(
                         (self.model.shift_count_diff[worker.id, shift.name])
                         * shift.balance_weighting
@@ -1081,7 +1121,6 @@ class RotaBuilder(object):
 
             # We use a similar method to balance the number of weekends worked
             # This works as long as weekend shifts are assigned as blocks!
-            if self.constraint_options["balance_weekends"]:
             weekend_shift_target_number = sum(
                 worker.shift_target_number[shift.name] / len(shift.shift_days)
                 for shift in self.get_shifts()
@@ -1090,6 +1129,8 @@ class RotaBuilder(object):
 
             worker.weekend_shift_target_number = weekend_shift_target_number
 
+            if self.constraint_options["balance_weekends"]:
+
                 if weekend_shift_target_number > 0:
                     self.model.constraints.add(
                         self.model.weekend_shift_count_t1[worker.id]
@@ -1166,7 +1207,7 @@ class RotaBuilder(object):
                             )
                         )
 
-            if self.constraint_options["balance_weekends"]:
+            #if self.constraint_options["balance_weekends"]:
             for week_blocks in self.get_week_block_iterator(
                 self.max_weekend_frequency
             ):
@@ -1184,20 +1225,22 @@ class RotaBuilder(object):
                 self.model.constraints.add(
                     self.constraint_options["max_shifts_per_week"]
                     >= sum(
-                        self.model.works[worker.id, week, day, shift.name]
-                        for day in self.days
-                        for shift in self.get_shifts()
+                        self.model.works[worker.id, week, day, shiftname]
+                        #for shiftname in self.get_shift_names_by_week_day(week, day)
+                        #for day in self.days
+                        for day, shiftname in self.get_day_shiftname_combinations()
                     )
                 )
 
-                for shift in self.get_shifts_with_constraint("max_2_shifts_per_week"):
-                    self.model.constraints.add(
-                        2
-                        >= sum(
-                            self.model.works[worker.id, week, day, shift.name]
-                            for day in self.days
-                        )
-                    )
+                # TODO: generic implementation
+                #for shift in self.get_shifts_with_constraint("max_2_shifts_per_week"):
+                #    self.model.constraints.add(
+                #        2
+                #        >= sum(
+                #            self.model.works[worker.id, week, day, shift.name]
+                #            for day in self.days
+                #        )
+                #    )
 
                 if (
                     worker.site
@@ -1305,22 +1348,24 @@ class RotaBuilder(object):
                 #         self.model.works[worker.id, week, day, shift.name]
                 #         for day in self.days[5:]
                 #         for shift in self.get_shifts()) / 2)
-                if self.constraint_options["balance_weekends"]:
+                #if self.constraint_options["balance_weekends"]:
                 self.model.constraints.add(
                     self.model.works_weekend[worker.id, week]
                     >= sum(
-                            self.model.works[worker.id, week, day, shift.name]
-                            for day in self.days[5:]
-                            for shift in self.get_shifts()
+                        self.model.works[worker.id, week, day, shiftname]
+                        for day, shiftname in self.get_day_shiftname_combinations() if day in self.days[5:]
+                        #for day in self.days[5:]
+                        #for shift in self.get_shifts()
                     )
                     / 2
                 )
                 self.model.constraints.add(
                     self.model.works_weekend[worker.id, week]
                     <= sum(
-                            self.model.works[worker.id, week, day, shift.name]
-                            for day in self.days[5:]
-                            for shift in self.get_shifts()
+                        self.model.works[worker.id, week, day, shiftname]
+                        for day, shiftname in self.get_day_shiftname_combinations() if day in self.days[5:]
+                        #for day in self.days[5:]
+                        #for shift in self.get_shifts()
                     )
                 )
                 # self.model.constraints.add(
@@ -1393,6 +1438,13 @@ class RotaBuilder(object):
                     n3week, n3day = weeks_days[n]
                     n3 = 0
 
+                # IF paired we check the following against both workers
+                workers = [worker]
+                if self.constraint_options["hard_constrain_pair_separation"]:
+                    for worker_pairs in self.worker_pairs:
+                        if worker_pairs[0] == worker:
+                            workers = worker_pairs
+
                 if self.get_shift_names_by_week_day(week, day):
                     # Unable to work (hard constraint not preference)
                     self.model.constraints.add(
@@ -1426,9 +1478,68 @@ class RotaBuilder(object):
                                     )
                                 )
 
+                for constraint_shift in self.get_shifts_with_constraints(
+                    "preclear", "preclear2"
+                ):
+                    if day in constraint_shift.shift_days:
+                        self.model.constraints.add(
+                            1
+                            >= self.model.works[worker.id, week, day, constraint_shift.name]
+                            + sum(
+                                p1 * self.model.works[w.id, pweek, pday, shiftname]
+                                for shiftname in self.get_shift_names_by_week_day(pweek, pday)
+                                #for day, shiftname in self.get_day_shiftname_combinations() if day == pday
+                                if shiftname != constraint_shift.name
+                                for w in workers
+                            )
+                        )
+
+                for constraint_shift in self.get_shifts_with_constraint("preclear2"):
+                    if day in constraint_shift.shift_days:
+                        self.model.constraints.add(
+                            1
+                            >= self.model.works[worker.id, week, day, constraint_shift.name]
+                            + sum(
+                                p2 * self.model.works[w.id, p2week, p2day, shiftname]
+                                for shiftname in self.get_shift_names_by_week_day(p2week, p2day)
+                                if shiftname != constraint_shift.name
+                                for w in workers
+                            )
+                        )
+
+                for constraint_shift in self.get_shifts_with_constraints(
+                    "postclear", "postclear2"
+                ):
+                    if day in constraint_shift.shift_days:
+                        self.model.constraints.add(
+                            1
+                            >= self.model.works[worker.id, week, day, constraint_shift.name]
+                            + sum(
+                                n1 * self.model.works[w.id, nweek, nday, shiftname]
+                                for shiftname in self.get_shift_names_by_week_day(nweek, nday)
+                                if shiftname != constraint_shift.name
+                                for w in workers
+                            )
+                        )
+
+                for constraint_shift in self.get_shifts_with_constraints("postclear2"):
+                    if day in constraint_shift.shift_days:
+                        self.model.constraints.add(
+                            1
+                            >= self.model.works[worker.id, week, day, constraint_shift.name]
+                            + sum(
+                                n2 * self.model.works[w.id, n2week, n2day, shiftname]
+                                for shiftname in self.get_shift_names_by_week_day(n2week, n2day)
+                                if shiftname != constraint_shift.name
+                                for w in workers
+                            )
+                        )
+
+                # NOTE: use pre / post clear instead
                 # if working a night ensure preceeding (1) or subsequent (2) shifts can only be nights
                 if self.constraint_options["constrain_time_off_after_nights"]:
                     for constraint_shift in self.get_shifts_with_constraint("night"):
+                        if day in constraint_shift.shift_days:
                             # Ensure night prior to unavalibity is not assigned
                             self.model.constraints.add(
                                 self.model.available[worker.id, week, day]
@@ -1437,13 +1548,6 @@ class RotaBuilder(object):
                                 ]
                             )
 
-                        # IF paired we check the following against both workers
-                        workers = [worker]
-                        if self.constraint_options["hard_constrain_pair_separation"]:
-                            for worker_pairs in self.worker_pairs:
-                                if worker_pairs[0] == worker:
-                                    workers = worker_pairs
-
                             # print("Workers", workers)
 
                             self.model.constraints.add(
@@ -1521,14 +1625,27 @@ class RotaBuilder(object):
 
             balance_modifier_constant = 1
 
+            if self.constraint_options["balance_shifts_over_workers"]:
+                worker_shift_balancing = sum(
+                    balance_modifier_constant
+                    * (
+                        self.model.worker_shift_count_t1[(worker.id)]
+                        + self.model.worker_shift_count_t2[(worker.id)]
+                    )
+                    for worker in self.workers
+                )
+            else:
+                worker_shift_balancing = 0
+
             if self.constraint_options["balance_shifts"]:
                 shift_balancing = sum(
                     balance_modifier_constant
                     * (
-                        self.model.shift_count_t1[(worker.id)]
-                        + self.model.shift_count_t2[(worker.id)]
+                        self.model.shift_count_t1[worker.id, shift.name]
+                        + self.model.shift_count_t2[worker.id, shift.name]
                     )
                     for worker in self.workers
+                    for shift in self.get_shifts()
                 )
             else:
                 shift_balancing = 0
@@ -1598,7 +1715,6 @@ class RotaBuilder(object):
             )
 
             # # Spread nights
-
             if self.constraint_options["balance_nights_across_sites"]:
                 nights_site_balancing = sum(
                     (
@@ -1631,14 +1747,10 @@ class RotaBuilder(object):
             #        print(worker_a, worker_b)
             #        for week, day in self.get_week_day_combinations():
             #            worker_pairs_balancing = worker_pairs_balancing + sum(self.model.works[worker_a.id, week, day, shift] for shift in self.get_shift_names_by_week_day(week, day)) * sum(self.model.works[worker_b.id, week, day, shift] for shift in self.get_shift_names_by_week_day(week, day)) * worker_pairs_constant
-
-            # return weekend_shift_balancing + blocks_balancing
-            # return bank_holiday_balancing
-            # return shift_balancing + preferences + blocks_balancing
-            # return shift_balancing + preferences + nights_site_balancing + blocks_balancing
             return (
                 weekend_shift_balancing
                 + shift_balancing
+                + worker_shift_balancing
                 + night_shift_balancing
                 + shift_diff_balancing
                 + preferences
@@ -1712,7 +1824,7 @@ class RotaBuilder(object):
         """
         self.shifts.extend(shifts)
 
-    def get_shift_names_by_week_day(self, week, day: DayStr) -> Set():
+    def get_shift_names_by_week_day(self, week, day: DayStr) -> set():
         """Returns the shifts required for a specific day
 
         Returns:
@@ -1833,6 +1945,9 @@ class RotaBuilder(object):
         """
         return self.weeks_days_product
 
+    def get_week_day_combinations_for_shift(self, shift) -> list:
+        return [(week, day) for week, day in self.get_week_day_combinations() if shift.name in self.get_shift_names_by_week_day(week, day)]
+
     def get_shifts(self) -> List[SingleShift]:
         """Returns a list of all the registered shifts
 
@@ -1844,6 +1959,15 @@ class RotaBuilder(object):
     def get_shifts_with_constraint(self, constraint) -> List[SingleShift]:
         return [shift for shift in self.shifts if constraint in shift.constraints]
 
+    def get_shifts_with_constraints(self, *constraints) -> set[SingleShift]:
+        shifts = set()
+
+        for constraint in constraints:
+            shifts.update(
+                [shift for shift in self.shifts if constraint in shift.constraints]
+            )
+        return shifts
+
     def get_shift_names(self) -> List[ShiftName]:
         """Returns a list of all the registered shift names
         
@@ -1944,42 +2068,37 @@ class RotaBuilder(object):
             if self.week_day_date_map[(week, day)] in bank_holiday_map
         ]
 
-
-class RotaResults(object):
-    def __init__(self, rota):
-        self.rota = rota
-        # self.results = results
-
+    # RESULTS
     def export_rota_to_html(self, filename: str = "rota"):
         with open(f"{filename}.html", "w") as f:
             f.write(self.get_worker_timetable_html(True))
 
     def export_rota_to_csv(self, filename: str = "rota"):
-        works = self.rota.model.works
+        works = self.model.works
         with open(f"{filename}.csv", "w", newline="") as f:
             wr = csv.writer(f, quoting=csv.QUOTE_ALL)
             l = ["Name"]
-            l.extend([worker.name for worker in self.rota.workers])
+            l.extend([worker.name for worker in self.workers])
             wr.writerow(l)
 
             l2 = ["Site"]
-            l2.extend([worker.site for worker in self.rota.workers])
+            l2.extend([worker.site for worker in self.workers])
             wr.writerow(l2)
 
             l3 = ["Grade"]
-            l3.extend([worker.grade for worker in self.rota.workers])
+            l3.extend([worker.grade for worker in self.workers])
             wr.writerow(l3)
 
             l4 = ["FTE"]
-            l4.extend([worker.fte for worker in self.rota.workers])
+            l4.extend([worker.fte for worker in self.workers])
             wr.writerow(l4)
 
-            for week, day in self.rota.get_week_day_combinations():
+            for week, day in self.get_week_day_combinations():
                 d = [f"Week {week} Day {day}"]
 
-                for worker in self.rota.workers:
+                for worker in self.workers:
                     i = ""
-                    for shift in self.rota.get_shift_names_by_week_day(week, day):
+                    for shift in self.get_shift_names_by_week_day(week, day):
                         if works[worker.id, week, day, shift].value == 1:
                             i = shift
 
@@ -1988,30 +2107,27 @@ class RotaResults(object):
 
     def get_work_table(self):
         """Build a timetable of the week as a dictionary from the model's optimal solution."""
-        works = self.rota.model.works
+        works = self.model.works
         week_table = {
-            week: {
-                day: {shift: [] for shift in self.rota.get_shift_names()}
-                for day in days
+            week: {day: {shift: [] for shift in self.get_shift_names()} for day in days}
+            for week in self.weeks
         }
-            for week in self.rota.weeks
-        }
-        for week in self.rota.weeks:
-            for worker in self.rota.workers:
-                for day, shift in self.rota.get_day_shiftname_combinations():
+        for week in self.weeks:
+            for worker in self.workers:
+                for day, shift in self.get_day_shiftname_combinations():
                     if works[worker.id, week, day, shift].value == 1:
                         week_table[week][day][shift].append(worker.get_details())
         return week_table
 
     def get_worker_timetable(self):
-        works = self.rota.model.works
+        works = self.model.works
         timetable = {
-            worker.name: {week: {day: "" for day in days} for week in self.rota.weeks}
-            for worker in self.rota.workers
+            worker.name: {week: {day: "" for day in days} for week in self.weeks}
+            for worker in self.workers
         }
-        for worker in self.rota.workers:
-            for week in self.rota.weeks:
-                for day, shift in self.rota.get_day_shiftname_combinations():
+        for worker in self.workers:
+            for week in self.weeks:
+                for day, shift in self.get_day_shiftname_combinations():
                     if works[worker.id, week, day, shift].value == 1:
                         timetable[worker.name][week][day] = shift
         return timetable
@@ -2019,18 +2135,18 @@ class RotaResults(object):
     def get_worker_timetable_brief(
         self, show_prefs=False, marker_every=30, show_unavailable=False
     ):
-        model = self.rota.model
+        model = self.model
         week_string = f"{'-Week-':20}" + "".join(
-            [7 * str(f"{w}")[-1:] for w in self.rota.weeks]
+            [7 * str(f"{w}")[-1:] for w in self.weeks]
         )
-        days_string = f"{'-Day-':20}" + "".join("MTWTFSS" * len(self.rota.weeks))
+        days_string = f"{'-Day-':20}" + "".join("MTWTFSS" * len(self.weeks))
         timetable = []
-        for worker in self.rota.workers:
+        for worker in self.workers:
             shifts = []
             w = [f"{worker.name:20}"]
-            for week, day in self.rota.get_week_day_combinations():
+            for week, day in self.get_week_day_combinations():
                 a = "-"
-                for shift in self.rota.get_shift_names_by_week_day(week, day):
+                for shift in self.get_shift_names_by_week_day(week, day):
                     if model.works[worker.id, week, day, shift].value > 0:
                         shifts.append(shift)
                         a = shift[0]
@@ -2049,7 +2165,7 @@ class RotaResults(object):
             if show_prefs:
                 # prefs
                 w = [f"{'Preferences':20}"]
-                for week, day in self.rota.get_week_day_combinations():
+                for week, day in self.get_week_day_combinations():
                     if model.pref_not_to_work[worker.id, week, day] > 0:
                         w.append("Y")
                     else:
@@ -2059,7 +2175,7 @@ class RotaResults(object):
             if show_unavailable:
                 # prefs
                 w = [f"{'Unavailable':20}"]
-                for week, day in self.rota.get_week_day_combinations():
+                for week, day in self.get_week_day_combinations():
                     if model.available[worker.id, week, day] > 0:
                         w.append("A")
                     else:
@@ -2080,30 +2196,30 @@ class RotaResults(object):
     def get_worker_timetable_html(
         self, include_html_tag=False, table_name="rota-table"
     ):
-        model = self.rota.model
+        model = self.model
 
         timetable = []
 
         timetable.append(
-            f"<h2>Rota start date: {self.rota.start_date.isoformat()} ({self.rota.weeks[-1]} weeks)</h2>"
+            f"<h2>Rota start date: {self.start_date.isoformat()} ({self.weeks[-1]} weeks)</h2>"
         )
 
         date_row = ["<th class='worker'></th>"]
 
         n = 0
-        for week, day in self.rota.get_week_day_combinations():
-            d = self.rota.start_date + datetime.timedelta(n)
+        for week, day in self.get_week_day_combinations():
+            d = self.start_date + datetime.timedelta(n)
             date_row.append(f"<th title='{d}'>Week {week}: {day}</th>")
             n = n + 1
         timetable.append(f"<tr class='data-row'>{''.join(date_row)}</tr>")
 
         current_site = ""
 
-        for worker in self.rota.workers:
+        for worker in self.workers:
             if worker.site != current_site:
                 try:
                     timetable.append(
-                        f"<tr><th class='site-title'>{worker.site} n={len(self.rota.workers_at_sites[worker.site])} fte={self.rota.full_time_equivalent_sites[worker.site]}</th><tr>"
+                        f"<tr><th class='site-title'>{worker.site} n={len(self.workers_at_sites[worker.site])} fte={self.full_time_equivalent_sites[worker.site]}</th><tr>"
                     )
                 except KeyError as e:
                     print(e)
@@ -2123,14 +2239,14 @@ class RotaResults(object):
             n = 0
 
             shift_tds = []
-            for week, day in self.rota.get_week_day_combinations():
-                d = self.rota.start_date + datetime.timedelta(n)
+            for week, day in self.get_week_day_combinations():
+                d = self.start_date + datetime.timedelta(n)
 
                 n = n + 1
                 a = "-"
                 shift_name = ""
 
-                for shift in self.rota.get_shift_names_by_week_day(week, day):
+                for shift in self.get_shift_names_by_week_day(week, day):
                     if model.works[worker.id, week, day, shift].value > 0:
                         shifts.append(shift)
                         a = shift[0]
@@ -2146,7 +2262,7 @@ class RotaResults(object):
                     available = False
                     css_class = "unavailable"
                     try:
-                        unavailable_reason = self.rota.unavailable_to_work_reason[
+                        unavailable_reason = self.unavailable_to_work_reason[
                             (worker.id, week, day)
                         ]
                     except KeyError:
@@ -2162,10 +2278,10 @@ class RotaResults(object):
                     bank_holiday = f" data-bank-holiday='{bank_holiday_map[d]}'"
 
                 requests = ""
-                if (worker.id, week, day) in self.rota.work_requests_map:
+                if (worker.id, week, day) in self.work_requests_map:
                     css_class = " ".join((css_class, "shift-requested"))
                     title = " ".join((title, "[REQUESTED]"))
-                    requests = f" data-shift-request='{self.rota.work_requests_map[worker.id, week, day]}'"
+                    requests = f" data-shift-request='{self.work_requests_map[worker.id, week, day]}'"
 
                 shift_tds.append(
                     f"<td title='{title}' class='rota-day {css_class}' data-shift='{shift_name}' data-available='{available}' data-unavailable_reason='{unavailable_reason}' data-date='{d}' data-week='{week}' data-day='{day}'{requests}{bank_holiday}>{a}</td>"
@@ -2179,7 +2295,7 @@ class RotaResults(object):
                 shift_count = shift_count + f"{s}: {c}, "
 
             shift_diff_dict = {}
-            for shift in self.rota.get_shifts():
+            for shift in self.get_shifts():
                 diff = model.shift_count_diff[worker.id, shift.name].value
                 shift_diff_dict[shift.name] = diff
 
@@ -2219,8 +2335,15 @@ class RotaResults(object):
                 + f"#weekends_worked: {model.worker_weekend_count[worker.id].value}\\#"
             )
 
+            if self.constraint_options["balance_bank_holidays"]:
                 bank_holiday_count = model.bank_holiday_count[worker.id].value
+
                 bank_holiday_count_w = model.bank_holiday_count_w[worker.id].value - 1
+            else:
+                bank_holiday_count = -1
+                bank_holiday_count_w = -1
+
+
             # print(worker.name, bank_holiday_count, bank_holiday_count_w)
             timetable.append(
                 f"<tr class='worker-row'>{worker_td}{''.join(shift_tds)}</tr>"
@@ -2230,7 +2353,7 @@ class RotaResults(object):
             # if show_prefs:
             #     # prefs
             #     w = ["{:20}".format("Preferences")]
-            #     for week, day in self.rota.get_week_day_combinations():
+            #     for week, day in self.get_week_day_combinations():
             #         if model.pref_not_to_work[worker.id, week, day] > 0:
             #             w.append("Y")
             #         else:
@@ -2240,7 +2363,7 @@ class RotaResults(object):
             # if show_unavailable:
             #     # prefs
             #     w = ["{:20}".format("Unavailable")]
-            #     for week, day in self.rota.get_week_day_combinations():
+            #     for week, day in self.get_week_day_combinations():
             #         if model.available[worker.id, week, day] > 0:
             #             w.append("A")
             #         else:
@@ -2249,9 +2372,9 @@ class RotaResults(object):
 
         result_stream = StringIO()
 
-        if self.rota.results is not None:
-            self.rota.results.write(ostream=result_stream)
-            self.rota.results.write_json(ostream=result_stream)
+        if self.results is not None:
+            self.results.write(ostream=result_stream)
+            self.results.write_json(ostream=result_stream)
             result_string = result_stream.getvalue()
         else:
             result_string = "Rota not run"
@@ -2267,14 +2390,14 @@ class RotaResults(object):
         <summary><h2>Rota settings</h2></summary>
         <div>
         <pre>
-        {json.dumps(self.rota.constraint_options, indent=4)}
+        {json.dumps(self.constraint_options, indent=4)}
         </pre>
         </details>
         </div>
         <details>
         <summary><h2>Shifts settings</h2></summary>
-        <div id="shifts-container" data-shifts='{json.dumps([i.name for i in self.rota.shifts])}'>
-        {"<br/>".join(str(i) for i in self.rota.shifts)}
+        <div id="shifts-container" data-shifts='{json.dumps([i.name for i in self.shifts])}'>
+        {"<br/>".join(str(i) for i in self.shifts)}
         </div>
         </details>
         <details>
@@ -2299,54 +2422,79 @@ class RotaResults(object):
 
         return html
 
-    def get_shift_summary(self):
-        works = self.rota.model.works
+    def get_shift_summary_dict(self):
+        """"""
         timetable = {
-            worker.get_details(): {
-                shift[0]: "" for shift in self.rota.get_shift_names()
+            worker.name: {shift: "" for shift in self.get_shift_names()}
+            for worker in self.workers
         }
-            for worker in self.rota.workers
+        for worker in self.workers:
+            for shift in self.get_shift_names():
+                timetable[worker.name][shift] = self.model.shift_count[
+                    worker.id, shift
+                ].value
+
+        return timetable
+
+    def get_worker_shift_list(self, worker: Worker) -> List:
+        shifts = []
+
+        for week, day in self.get_week_day_combinations():
+            # d = self.start_date + datetime.timedelta(n)
+            # n = n + 1
+
+            temp = ""
+            for shift in self.get_shift_names_by_week_day(week, day):
+                if self.model.works[worker.id, week, day, shift].value > 0:
+                    temp = shift
+            shifts.append(temp)
+
+        return shifts
+
+    def get_shift_summary(self):
+        works = self.model.works
+        timetable = {
+            worker.get_details(): {shift: "" for shift in self.get_shift_names()}
+            for worker in self.workers
         }
         # timetable = { worker.get_details() : { shift : "" for shift in ["truro_twilight"] } for worker in workers }
         t = []
-        for worker in self.rota.workers:
+        for worker in self.workers:
             l = []
 
             total_shifts = 0
-            for shift in self.rota.get_shift_names():
+            for shift in self.get_shift_names():
                 # for shift in ["truro_twilight"]:
                 c = [
-                    works[worker.id, week, day, shift]
-                    for week in self.rota.weeks
+                    works[worker.id, week, day, shift].value
+                    for week in self.weeks
                     for day in days
                 ].count(1)
                 if c > 0:
                     l.append(f"{shift} ({c})")
                     total_shifts = total_shifts + c
                 # print(worker.id, shift)
-                timetable[worker.get_details()][shift[0]] = c
+                timetable[worker.get_details()][shift] = c
             t.append(f"{worker.get_full_details()} [{total_shifts}]: {', '.join(l)}")
         return "\n".join(t)
 
     def get_shift_summary_html(self):
-        works = self.rota.model.works
+        works = self.model.works
         timetable = {
-            worker.get_details(): {
-                shift[0]: "" for shift in self.rota.get_shift_names()
-            }
-            for worker in self.rota.workers
+            worker.get_details(): {shift[0]: "" for shift in self.get_shift_names()}
+            for worker in self.workers
         }
         # timetable = { worker.get_details() : { shift : "" for shift in ["truro_twilight"] } for worker in workers }
         t = []
-        for worker in self.rota.workers:
+        for worker in self.workers:
             l = []
 
             total_shifts = 0
-            for shift in self.rota.get_shift_names():
+            for shift in self.get_shift_names():
                 # for shift in ["truro_twilight"]:
                 c = [
                     works[worker.id, week, day, shift]
-                    for week in self.rota.weeks
+                    for week in self.weeks
                     for day in days
                 ].count(1)
                 if c > 0:
@@ -2357,80 +2505,6 @@ class RotaResults(object):
             t.append(f"{worker.get_full_details()} [{total_shifts}]: {', '.join(l)}")
         return "\n".join(t)
 
-    # def get_no_preference(no_pref):
-    #    """Extract to a list the workers not satisfied with their weekend preference."""
-    #    return [worker.id for worker in workers if no_pref[worker.id].value == 1]
-
-    # def get_night_blocks(self):
-    #     nights = self.rota.model.nights
-    #     timetable = {
-    #         worker.name: {
-    #             week: {block: ""
-    #                    for block in self.rota.night_blocks}
-    #             for week in self.rota.weeks
-    #         }
-    #         for worker in self.rota.workers
-    #     }
-    #     for worker in self.rota.workers:
-    #         for week in self.rota.weeks:
-    #             for block in self.rota.night_blocks:
-    #                 if nights[worker.id, week, block].value == 1:
-    #                     timetable[worker.name][week][block] = "true"
-    #     return timetable
-
-    def get_multinight(self):
-        for site in self.rota.sites:
-            for shift in self.get_shifts_with_constraint("night"):
-                block = shift.name
-                for week in self.rota.weeks:
-                    print(
-                        site,
-                        week,
-                        block,
-                        self.rota.model.night_per_site[week, block, site].value,
-                        self.rota.model.night_per_site_t1[week, block, site].value,
-                        self.rota.model.night_per_site_t2[week, block, site].value,
-                    )
-
-    """
-    def get_night_details(self):
-        for worker in self.rota.workers:
-            print(
-                "{:20}".format(worker.name),
-                "worked: {},".format(
-                    self.rota.model.night_shift_count[(worker.id)].value),
-                    worker.shift_target_number["night_weekday"] +
-                    worker.shift_target_number["night_weekend"]),
-                "target_diff: {},".format(
-                    self.rota.model.night_shift_count_t1[(worker.id)].value +
-                    self.rota.model.night_shift_count_t2[(worker.id)].value +
-                    1),
-                "balance: {},".format(
-                    self.rota.model.night_shift_count_w[worker.id].value),
-            )
-
-    def get_weekend_details(self):
-        for worker in self.rota.workers:
-            print(
-                "{:20}".format(worker.name),
-                "worked: {},".format(
-                    self.rota.model.worker_weekend_count[(worker.id)].value),
-                "target: {},".format(
-                    sum([
-                        worker.shift_target_number[shift.name]
-                        for shift in self.rota.get_shifts()
-                        if not set(("Sat", "Sun")).isdisjoint(shift.shift_days)
-                    ])),
-                "target_diff: {},".format(
-                    self.rota.model.weekend_shift_count_t1[(worker.id)].value +
-                    self.rota.model.weekend_shift_count_t2[(worker.id)].value +
-                    1),
-                "balance: {},".format(
-                    self.rota.model.weekend_shift_count_w[worker.id].value),
-            )
-            
-    """
-
 
 class NoActiveSites(Exception):
     """Raised when there are no active sites"""

test_rota.py

@@ -1,6 +1,6 @@
 from black import main
 import pytest
-from shifts import NoWorkers, RotaBuilder, RotaResults, SingleShift, days
+from shifts import NoWorkers, RotaBuilder, SingleShift, days
 
 import datetime
 
@@ -27,8 +27,10 @@ class TestDemoRota:
     Rota.constraint_options["constrain_time_off_after_nights"] = False
     Rota.constraint_options["balance_nights_across_site"] = False
     Rota.constraint_options["balance_shifts"] = True
+    Rota.constraint_options["balance_shifts_over_workers"] = True
     Rota.constraint_options["balance_nights"] = False
-    Rota.constraint_options["minimise_shift_diffs"] = True
+    Rota.constraint_options["minimise_shift_diffs"] = False
+    Rota.constraint_options["balance_blocks"] = False
     #Rota.constraint_options["balance_weekends"] = False
 
     # Add a few workers
@@ -52,7 +54,7 @@ class TestDemoRota:
             workers_required=1,
             # We use a different balance weighting for each shift as otherwise
             # the shifts can be treated equally (and no target diff)
-            balance_weighting=0.5, 
+            #balance_weighting=0.6, 
         ),
         SingleShift(("group1",), "weekend", 12.5, days[5:], balance_offset=40,
         #balance_weighting=0.5
@@ -80,15 +82,17 @@ class TestDemoRota:
     # result.Solver.Status = SolverStatus.warning
     # prob.solutions.load_from(result)
 
-    ResultsHolder = RotaResults(Rota)
+    #worker_timetable_brief = Rota.get_worker_timetable_brief(
+    #    show_prefs=False, show_unavailable=False
+    #)
 
+    #print(worker_timetable_brief)
+    #print(Rota.get_shift_summary_dict())
+    Rota.export_rota_to_html("test2")
 
-    worker_timetable_brief = ResultsHolder.get_worker_timetable_brief(
-        show_prefs=False, show_unavailable=False
-    )
-
-    print(worker_timetable_brief)
-    ResultsHolder.export_rota_to_html("test")
+    def test_optimal_solution(self):
+        assert self.Rota.results.solver.status == "ok"
+        assert self.Rota.results.solver.termination_condition == "optimal"
 
     def test_start_date(self):
         assert self.Rota.start_date == self.start_date
@@ -145,23 +149,349 @@ class TestDemoRota:
                 case _:
                     assert 1 == 0
 
+    def test_worker_shift_targets(self):
+        shift_number = self.Rota.get_shift_summary_dict()
+
+        for worker in self.Rota.workers:
+            match worker.site:
+                case "group1":
+                    assert shift_number[worker.name]["weekday"] == 10
+                    assert shift_number[worker.name]["weekend"] == 5
+                case "group2":
+                    assert shift_number[worker.name]["weekend"] == 0
+                    match worker.fte:
+                        case 60:
+                            assert shift_number[worker.name]["weekday"] == 6
+                        case 40:
+                            assert shift_number[worker.name]["weekday"] == 4
+                        case _:
+                            assert 1 == 0
+                case _:
+                    assert 1 == 0
+
+
     def test_rota_sites(self):
         assert self.Rota.sites == {"group1", "group2"}
 
+class TestDemoRotaNights:
+    # Set up rota
+    weeks_to_rota = 12
+    start_date = datetime.date(2022, 3, 14)
 
-class TestNoWorkerRota:
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+        balance_offset_modifier=1,
+        max_weekend_frequency=1,
+    )
+
+    Rota.constraint_options["max_shifts_per_week"] = 5
+    Rota.constraint_options["max_shifts_per_month"] = 31
+    Rota.constraint_options["ensure_derriford_reg_for_nights"] = False
+    Rota.constraint_options["limit_to_1_st2_on_nights"] = False
+    Rota.constraint_options["ensure_1_st4_plus_on_nights"] = False
+    Rota.constraint_options["constrain_time_off_after_nights"] = False
+    Rota.constraint_options["balance_nights_across_site"] = False
+    Rota.constraint_options["balance_shifts"] = True
+    Rota.constraint_options["balance_shifts_over_workers"] = True
+    Rota.constraint_options["balance_nights"] = False
+    Rota.constraint_options["minimise_shift_diffs"] = False
+    Rota.constraint_options["balance_blocks"] = False
+    #Rota.constraint_options["balance_weekends"] = False
+
+    # Add a few workers
+    worker1 = Worker(Rota, "worker1", "group1", 1)
+    worker2 = Worker(Rota, "worker2", "group1", 1)
+    worker3 = Worker(Rota, "worker3", "group1", 1)
+    worker4 = Worker(Rota, "worker4", "group1", 1)
+    worker5 = Worker(Rota, "worker5", "group2", 1, fte=60)
+    worker6 = Worker(Rota, "worker6", "group2", 1, fte=40)
+    worker7 = Worker(Rota, "worker7", "group2", 1)
+    worker8 = Worker(Rota, "worker8", "group2", 1)
+
+    Rota.add_workers((worker1, worker2, worker3, worker4, worker5, worker6, worker7, worker8))
+
+    # Add a weekday and weekend shift
+    Rota.add_shifts(
+        SingleShift(
+            ("group1", "group2"),
+            "night_weekday",
+            12.5,
+            days[:4],
+            balance_offset=10,
+            workers_required=1,
+            constraints=["preclear2", "postclear2"],
+            #constraints=["night"],
+            force_as_block=True
+        
+            # We use a different balance weighting for each shift as otherwise
+            # the shifts can be treated equally (and no target diff)
+            #balance_weighting=0.6, 
+        ),
+        SingleShift(("group1","group2"), "night_weekend", 12.5, days[4:], balance_offset=40,
+        #balance_weighting=0.5
+            constraints=["preclear2", "postclear2"],
+            #constraints=["night"],
+            force_as_block=True
+        ),
+        SingleShift(("group1","group2"), "twilight", 12.5, days[:5], balance_offset=40,
+        #balance_weighting=0.5
+        workers_required=2,
+            #constraints=["preclear2", "postclear2"],
+        ),
+        #SingleShift(("group1","group2"), "atwilight", 12.5, days[:5], balance_offset=40,
+        ##balance_weighting=0.5
+        #workers_required=2,
+        #    #constraints=["postclear2"],
+        #),
+
+    )
+
+    Rota.build_shifts()
+    Rota.build_workers()
+    Rota.build_model()
+
+    print(Rota.get_worker_details())
+
+    for w in Rota.get_workers():
+        print(w)
+        print(w.get_shift_targets())
+
+    solver_options = {"ratio": 0.1, "seconds": 1000, "threads": 10}
+
+    Rota.solve_model(options=solver_options)
+
+    # print(Rota.get_worker_details())
+
+    # optimizer = SolverFactory('cbc')
+    # result = optimizer.solve(prob,tee=True)
+    # result.Solver.Status = SolverStatus.warning
+    # prob.solutions.load_from(result)
+
+    worker_timetable_brief = Rota.get_worker_timetable_brief(
+        show_prefs=False, show_unavailable=False
+    )
+
+    print(worker_timetable_brief)
+    print(Rota.get_shift_summary_dict())
+    Rota.export_rota_to_html("test3")
+
+    print(Rota.get_worker_shift_list(worker1))
+
+    def test_start_date(self):
+        assert self.Rota.start_date == self.start_date
+
+class TestDemoRotaClear:
     weeks_to_rota = 10
     start_date = datetime.date(2022, 3, 7)
 
     Rota = RotaBuilder(
         start_date,
         weeks_to_rota=weeks_to_rota,
+        balance_offset_modifier=1,
+        max_weekend_frequency=1,
+    )
+    #Rota.constraint_options["max_shifts_per_week"] = 5
+    Rota.constraint_options["max_shifts_per_month"] = 20
+    #Rota.constraint_options["ensure_derriford_reg_for_nights"] = False
+    #Rota.constraint_options["limit_to_1_st2_on_nights"] = False
+    #Rota.constraint_options["ensure_1_st4_plus_on_nights"] = False
+    #Rota.constraint_options["constrain_time_off_after_nights"] = False
+    #Rota.constraint_options["balance_nights_across_site"] = False
+    #Rota.constraint_options["balance_shifts"] = True
+    #Rota.constraint_options["balance_shifts_over_workers"] = True
+    #Rota.constraint_options["balance_nights"] = False
+    #Rota.constraint_options["minimise_shift_diffs"] = False
+    #Rota.constraint_options["balance_blocks"] = False
+
+    worker1 = Worker(Rota, "worker1", "group1", 1)
+    worker2 = Worker(Rota, "worker2", "group1", 1)
+    worker3 = Worker(Rota, "worker3", "group1", 1)
+    worker4 = Worker(Rota, "worker4", "group1", 1)
+    worker5 = Worker(Rota, "worker5", "group2", 1, fte=60)
+    worker6 = Worker(Rota, "worker6", "group2", 1, fte=40)
+    worker7 = Worker(Rota, "worker7", "group2", 1)
+    worker8 = Worker(Rota, "worker8", "group2", 1)
+
+    Rota.add_workers((worker1, worker2, worker3, worker4, worker5, worker6, worker7, worker8))
+
+    def test_preclear(self):
+        self.Rota.add_shifts(
+            SingleShift(("group1","group2"), "a", 12.5, days[:5], balance_offset=40,
+            workers_required=2,
+                constraints=["preclear"],
+            ),
+            SingleShift(("group1","group2"), "b", 12.5, days[:5], balance_offset=40,
+            workers_required=2,
+            ),
+            SingleShift(("group1","group2"), "c", 12.5, days[3], balance_offset=40,
+            workers_required=1,
+                constraints=["preclear2", "postclear2"],
+            ),
         )
 
-    def test_no_workers(self):
-        with pytest.raises(NoWorkers):
+        self.Rota.build_shifts()
         self.Rota.build_workers()
+        self.Rota.build_model()
+
+
+        solver_options = {"ratio": 0.1, "seconds": 1000, "threads": 10}
+
+        self.Rota.solve_model(options=solver_options)
+        self.Rota.export_rota_to_html("test4")
+
+        assert self.Rota.results.solver.status == "ok"
+        assert self.Rota.results.solver.termination_condition == "optimal"
+
+        for worker in self.Rota.workers:
+            shifts = self.Rota.get_worker_shift_list(worker)
+
+            # Convert shift to a string representation
+            shifts_string = "".join([i if i != "" else "-" for i in shifts])
+
+            # check for patterns that should not occur
+            assert "ba" not in shifts_string
+            assert "ca" not in shifts_string
+            assert "ac" not in shifts_string
+            assert "ca" not in shifts_string
+            assert "a-c" not in shifts_string
+            assert "c-a" not in shifts_string
+            assert "bc" not in shifts_string
+            assert "cb" not in shifts_string
+            assert "b-c" not in shifts_string
+            assert "c-b" not in shifts_string
+
+class TestDemoRotaShiftConstraints:
+    weeks_to_rota = 10
+    start_date = datetime.date(2022, 3, 7)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+        balance_offset_modifier=1,
+        max_weekend_frequency=1,
+    )
+    #Rota.constraint_options["max_shifts_per_week"] = 5
+    #Rota.constraint_options["max_shifts_per_month"] = 20
+    worker1 = Worker(Rota, "worker1", "group1", 1)
+    worker2 = Worker(Rota, "worker2", "group1", 1)
+
+    Rota.add_workers((worker1, worker2))
+    Rota.add_shifts(
+        SingleShift(("group1",), "a", 12.5, days, balance_offset=40,
+        workers_required=1,
+        ),
+    )
+
+    def test_max_shifts(self):
+        self.Rota.constraint_options["max_shifts_per_month"] = 14
+        self.Rota.build_and_solve()
+
+        assert self.Rota.results.solver.status == "ok"
+        assert self.Rota.results.solver.termination_condition == "optimal"
+
+    def test_max_shifts_fail(self):
+        self.Rota.constraint_options["max_shifts_per_month"] = 13
+        self.Rota.build_and_solve()
+        assert self.Rota.results.solver.status == "warning"
+        assert self.Rota.results.solver.termination_condition == "infeasible"
+
+    def test_max_shifts_extra_worker(self):
+        self.Rota.constraint_options["max_shifts_per_month"] = 12
+        worker3 = Worker(self.Rota, "worker3", "group1", 1)
+
+        self.Rota.add_worker(worker3)
+        self.Rota.build_and_solve()
+
+        assert self.Rota.results.solver.status == "ok"
+        assert self.Rota.results.solver.termination_condition == "optimal"
+
+    def test_max_shifts_extra_worker_wrong_group(self):
+        self.Rota.constraint_options["max_shifts_per_month"] = 12
+        worker3 = Worker(self.Rota, "worker3", "group2", 1)
+        self.Rota.add_worker(worker3)
+        self.Rota.build_and_solve()
+
+        assert self.Rota.results.solver.status == "warning"
+        assert self.Rota.results.solver.termination_condition == "infeasible"
+
+    def test_max_shifts_per_week_fail(self):
+        self.Rota.constraint_options["max_shifts_per_week"] = 3
+        self.Rota.build_and_solve()
+
+        assert self.Rota.results.solver.status == "warning"
+        assert self.Rota.results.solver.termination_condition == "infeasible"
+
+    def test_max_shifts_per_week_pass(self):
+        self.Rota.constraint_options["max_shifts_per_week"] = 4
+        self.Rota.build_and_solve()
+
+        assert self.Rota.results.solver.status == "ok"
+        
+    def test_max_shifts_per_week_extra_worker_pass(self):
+        self.Rota.constraint_options["max_shifts_per_week"] = 3
+        worker3 = Worker(self.Rota, "worker3", "group1", 1)
+        self.Rota.add_worker(worker3)
+        self.Rota.build_and_solve()
+
+        assert self.Rota.results.solver.status == "ok"
+
+class TestDemoRotaBalanceShiftSites:
+    weeks_to_rota = 8
+    start_date = datetime.date(2022, 3, 7)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+    #Rota.constraint_options["max_shifts_per_week"] = 5
+    #Rota.constraint_options["max_shifts_per_month"] = 20
+    worker1 = Worker(Rota, "worker1", "group1", 1)
+    worker2 = Worker(Rota, "worker2", "group1", 1)
+    worker3 = Worker(Rota, "worker3", "group2", 1)
+    worker4 = Worker(Rota, "worker4", "group2", 1)
+    worker5 = Worker(Rota, "worker3", "group2", 1)
+    worker6 = Worker(Rota, "worker4", "group2", 1)
+
+    Rota.add_workers((worker1, worker2, worker3, worker4))
+    #Rota.add_workers((worker5, worker6))
+    Rota.add_shifts(
+        SingleShift(("group1", "group2"), "weekday_night", 12.5, days[:4], balance_offset=40,
+        workers_required=2,
+        force_as_block=True,
+        constraints=["balance_across_groups"]
+        ),
+        #SingleShift(("group1", "group2"), "weekend_night", 12.5, days[4:], balance_offset=40,
+        #workers_required=2,
+        #force_as_block=True,
+        #),
+    )
+
+    def test_balance_blocks_across_groups(self):
+        self.Rota.constraint_options["balance_nights_across_sites"] = True
+        #self.Rota.constraint_options["balance_bank_holidays"] = False
+        #self.Rota.constraint_options["balance_weekends"] = False
+        #self.Rota.constraint_options["balance_blocks"] = False
+        self.Rota.build_and_solve(options={"ratio": 0.0, "seconds": 1000, "threads": 10})
+        self.Rota.export_rota_to_html("test5")
+
+        assert self.Rota.results.solver.status == "ok"
+        assert self.Rota.results.solver.termination_condition == "optimal"
+
+#class TestNoWorkerRota:
+#    weeks_to_rota = 10
+#    start_date = datetime.date(2022, 3, 7)
+#
+#    Rota = RotaBuilder(
+#        start_date,
+#        weeks_to_rota=weeks_to_rota,
+#    )
+#
+#    def test_no_workers(self):
+#        with pytest.raises(NoWorkers):
+#            self.Rota.build_workers()
 
 
 if __name__ == "__main__":
-    t = TestDemoRota()
+    t = TestDemoRotaShiftConstraints()
+    t.test_balance_blocks_across_groups()