From 500c9e23a04a2b88773456f3117901783ac88846 Mon Sep 17 00:00:00 2001
From: Ross <xkjq@ymail.com>
Date: Fri, 28 Jan 2022 17:38:07 +0000
Subject: [PATCH] start generic group balance

---
 shifts.py | 53 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 53 insertions(+)

diff --git a/shifts.py b/shifts.py
index 5d72cb8..3d6b3c4 100644
--- a/shifts.py
+++ b/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()