start generic group balance

shifts.py

@@ -360,6 +360,39 @@ class RotaBuilder(object):
                 initialize=0,
             )
 
+        # Replacement for balance_nights_across_sites
+        self.model.shift_per_site = Var(
+            (
+                (week, shift.name, site)
+                for site in self.sites
+                for week in self.weeks
+                for shift in self.get_shifts_with_constraint("balance_across_groups")
+            ),
+            # within=Binary,
+            within=NonNegativeIntegers,
+            initialize=0,
+        )
+
+        self.model.shift_per_site_t1 = Var(
+            (
+                (week, shift.name, site)
+                for site in self.sites
+                for week in self.weeks
+                for shift in self.get_shifts_with_constraint("balance_across_groups")
+            ),
+            domain=NonNegativeIntegers,
+            initialize=0,
+        )
+        self.model.shift_per_site_t2 = Var(
+            (
+                (week, shift.name, site)
+                for site in self.sites
+                for week in self.weeks
+                for shift in self.get_shifts_with_constraint("balance_across_groups")
+            ),
+            domain=NonNegativeIntegers,
+            initialize=0,
+        )
 
         self.model.shift_count = Var(
             (
@@ -748,6 +781,26 @@ class RotaBuilder(object):
                             - self.model.night_per_site_t2[week, block, site]
                             == self.model.night_per_site[week, block, site] - 1
                         )
+        for site in self.sites:
+            for shift in self.get_shifts_with_constraint("balance_across_groups"):
+                block = shift.name
+                for week in self.weeks:
+                    self.model.constraints.add(
+                        self.model.shift_per_site[week, block, site]
+                        == sum(
+                            self.model.shift_week_worker_assigned[
+                                block, week, worker.id
+                            ]
+                            for worker in self.workers
+                            if worker.site == site
+                        )
+                        # 1 >= sum(self.model.shift_week_worker_assigned[shift.name, week, worker.id] for worker in self.workers if worker.site == site)
+                    )
+                    self.model.constraints.add(
+                        self.model.shift_per_site_t1[week, block, site]
+                        - self.model.shift_per_site_t2[week, block, site]
+                        == self.model.shift_per_site[week, block, site] - 1
+                    )
 
         weeks_days = self.get_week_day_combinations()