From 48d6b748d1767b48bf860d255d7235c9302fb948 Mon Sep 17 00:00:00 2001
From: Ross <xkjq@ymail.com>
Date: Sun, 22 Jun 2025 22:34:05 +0100
Subject: [PATCH] fix the weekend count for multishits

---
 rota/shifts.py          | 56 +++++++++++++++++++++++++----------------
 test/test_rota_clear.py |  4 +--
 2 files changed, 37 insertions(+), 23 deletions(-)

diff --git a/rota/shifts.py b/rota/shifts.py
index de1cfd2..ceae404 100644
--- a/rota/shifts.py
+++ b/rota/shifts.py
@@ -2145,6 +2145,20 @@ class RotaBuilder(object):
                     )
 
             for week in self.weeks:
+
+                weekend_days = self.days[5:]  # ["Sat", "Sun"]
+                total_assigned = sum(
+                    self.model.works[worker.id, week, day, shiftname]
+                    for w, day, shiftname in self.get_all_shiftname_combinations(week=week)
+                    if day in weekend_days
+                )
+                # If any shift is assigned on Sat or Sun, works_weekend is 1
+                self.model.constraints.add(total_assigned >= self.model.works_weekend[worker.id, week])
+                # If works_weekend is 1, at least one shift must be assigned
+                self.model.constraints.add(total_assigned <= 1000 * self.model.works_weekend[worker.id, week])
+
+
+
                 for constraint_shift in self.get_shifts_with_constraints(
                     "max_shifts_per_week"
                 ):
@@ -2301,27 +2315,27 @@ class RotaBuilder(object):
 
 
                 # Try disabling tihs to allow more than 2 shifts ot be assigned on a weekend
-                self.model.constraints.add(
-                    self.model.works_weekend[worker.id, week]
-                    >= sum(
-                        self.model.works[worker.id, week, day, shiftname]
-                        for w, day, shiftname in self.get_all_shiftname_combinations(
-                            week=week
-                        )
-                        if day in self.days[5:]
-                    )
-                    / 2
-                )
-                self.model.constraints.add(
-                    self.model.works_weekend[worker.id, week]
-                    <= sum(
-                        self.model.works[worker.id, week, day, shiftname]
-                        for w, day, shiftname in self.get_all_shiftname_combinations(
-                            week=week
-                        )
-                        if day in self.days[5:]
-                    )
-                )
+                #self.model.constraints.add(
+                #    self.model.works_weekend[worker.id, week]
+                #    >= sum(
+                #        self.model.works[worker.id, week, day, shiftname]
+                #        for w, day, shiftname in self.get_all_shiftname_combinations(
+                #            week=week
+                #        )
+                #        if day in self.days[5:]
+                #    )
+                #    / 2
+                #)
+                #self.model.constraints.add(
+                #    self.model.works_weekend[worker.id, week]
+                #    <= sum(
+                #        self.model.works[worker.id, week, day, shiftname]
+                #        for w, day, shiftname in self.get_all_shiftname_combinations(
+                #            week=week
+                #        )
+                #        if day in self.days[5:]
+                #    )
+                #)
                 # self.model.constraints.add(
                 #     self.model.works_saturday[worker.id, week] == sum(
                 #         self.model.works[worker.id, week, "Sat", shift.name]
diff --git a/test/test_rota_clear.py b/test/test_rota_clear.py
index edc0d2a..4950e5f 100644
--- a/test/test_rota_clear.py
+++ b/test/test_rota_clear.py
@@ -365,7 +365,7 @@ def test_pre_post_clear_force_assign(demo_rota_clear):
             #],
         ),
     )
-    Rota.build_and_solve(options={"ratio": 0.0})
+    Rota.build_and_solve(options={"ratio": 0.01})
     Rota.export_rota_to_html("test_pre_post_clear", folder="tests")
 
     assert Rota.results.solver.status == "ok"
@@ -441,7 +441,7 @@ def test_pre_post_clear_force_assign2(demo_rota_clear):
             ],
         ),
     )
-    Rota.build_and_solve(options={"ratio": 0.0})
+    Rota.build_and_solve(options={"ratio": 0.01})
     Rota.export_rota_to_html("test_pre_post_clear", folder="tests")
 
     assert Rota.results.solver.status == "ok"