.

gen.py

@@ -1,6 +1,7 @@
 import datetime
 import os
 import sys
+import time
 from rota.shifts import NoWorkers, RotaBuilder, SingleShift, days
 
 
@@ -32,12 +33,13 @@ time_to_run = 258500
 ratio = 0.5
 
 
-start_date = datetime.date(2022, 9, 5)
+start_date = datetime.date(2022, 11, 7)
+#start_date = datetime.date(2022, 10, 10)
 
 Rota = RotaBuilder(start_date, weeks_to_rota=12)
 Rota.constraint_options["balance_weekends"] = True
-Rota.constraint_options["max_night_frequency"] = 1
+Rota.constraint_options["max_night_frequency"] = 4
-Rota.constraint_options["max_weekend_frequency"] = 1
+Rota.constraint_options["max_weekend_frequency"] = 4
 # Rota.constraint_options["avoid_st2_first_month"] = True
 
 Rota.add_shifts(
@@ -47,7 +49,7 @@ Rota.add_shifts(
         length=12.5,
         days=days[:5],
         balance_offset=3,
-        constraint=[{"name": "max_shifts_per_week", "options": 2}],
+        constraint=[{"name": "max_shifts_per_week", "options": 2, }],
     ),
     SingleShift(
         sites=("truro", "truro twilights"),
@@ -81,8 +83,8 @@ Rota.add_shifts(
         balance_offset=4,
         rota_on_nwds=True,
         force_as_block=True,
-        constraint=[{"name": "preclear2"}, {"name": "postclear2"}, {"name": "single_shift_block_per_week"}],
+        constraint=[{"name": "post", "options": 2}, {"name": "pre", "options": 5}],
-        #constraint=[{"name": "preclear2"}, {"name": "postclear2"},],
+        #constraint=[{"name": "pre", "options": 2}, {"name": "post", "options": 2},],
     ),
     SingleShift(
         sites=("truro",),
@@ -92,7 +94,7 @@ Rota.add_shifts(
         balance_offset=4,
         rota_on_nwds=True,
         force_as_block=True,
-        constraint=[{"name": "preclear2"}, {"name": "postclear2"}],
+        constraint=[{"name": "pre", "options": 2}, {"name": "post", "options": 2}],
         # force_as_block_unless_nwd=True
     ),
     SingleShift(
@@ -103,7 +105,7 @@ Rota.add_shifts(
         balance_offset=4,
         rota_on_nwds=True,
         force_as_block=True,
-        constraint=[{"name": "preclear2"}, {"name": "postclear2"}],
+        constraint=[{"name": "pre", "options": 2}, {"name": "post", "options": 2}],
         # force_as_block_unless_nwd=True
     ),
     SingleShift(
@@ -115,7 +117,7 @@ Rota.add_shifts(
         workers_required=1,
         rota_on_nwds=True,
         force_as_block=True,
-        constraint=[{"name": "preclear2"}, {"name": "postclear2"}],
+        constraint=[{"name": "pre", "options": 2}, {"name": "post", "options": 2}],
     ),
     SingleShift(
         sites=("plymouth",),
@@ -126,7 +128,7 @@ Rota.add_shifts(
         workers_required=1,
         rota_on_nwds=True,
         force_as_block=True,
-        constraint=[{"name": "preclear2"}, {"name": "postclear2"}],
+        constraint=[{"name": "pre", "options": 2}, {"name": "post", "options": 2}],
     ),
     SingleShift(
         sites=("plymouth", "plymouth twilights"),
@@ -152,8 +154,8 @@ Rota.add_shifts(
         rota_on_nwds=True,
         constraint=[
             {"name": "night"},
-            {"name": "preclear2"},
+            {"name": "pre", "options": 5},
-            {"name": "postclear2"},
+            {"name": "post", "options": 5},
             {"name": "require_remote_site_presence_week", "options": ("plymouth", 1)},
         ],
     ),
@@ -170,8 +172,8 @@ Rota.add_shifts(
         rota_on_nwds=True,
         constraint=[
             {"name": "night"},
-            {"name": "preclear2"},
+            {"name": "pre", "options": 5},
-            {"name": "postclear2"},
+            {"name": "post", "options": 5},
             {"name": "require_remote_site_presence_week", "options": ("plymouth", 1)},
         ],
     ),
@@ -284,9 +286,13 @@ if load_leave:
 Rota.build_workers()
 Rota.build_model()
 
-solver_options = {"ratio": ratio, "seconds": time_to_run, "threads": 10}
+solver_options = {"ratio": ratio, "seconds": time_to_run, "threads": 6}
 
+start_time = time.time()
 Rota.solve_model(options=solver_options)
+end_time = time.time()
+
+print(f"Time taken {end_time-start_time}")
 
 # print(Rota.get_worker_details())
 

rota/shifts.py

@@ -162,7 +162,6 @@ class RotaBuilder(object):
     ):
 
         print(f"Start time: {datetime.datetime.now()}")
-        print(f"Weeks to rota: {weeks_to_rota}")
         print(f"Use previous shifts {use_previous_shifts}")
         print(f"Use previous shifts (balance_extra) {use_shift_balance_extra}")
 
@@ -188,6 +187,10 @@ class RotaBuilder(object):
         self.rota_days_length = len(self.weeks) * 7
         self.rota_end_date = self.start_date + datetime.timedelta(self.rota_days_length)
 
+        print(
+            f"Weeks to rota: {weeks_to_rota} ({self.start_date} - {self.rota_end_date})"
+        )
+
         self.unavailable_to_work = set()
         self.unavailable_to_work_reason = {}
         self.pref_not_to_work = {}
@@ -205,7 +208,6 @@ class RotaBuilder(object):
         self.balance_offset_modifier = balance_offset_modifier
         self.ltft_balance_offset = ltft_balance_offset
 
-
         self.constraint_options = {
             "balance_nights": True,
             "constrain_time_off_after_nights": False,
@@ -221,7 +223,7 @@ class RotaBuilder(object):
             # Don't assign multiple shifts every (n) weeks
             "max_weekend_frequency": 1,  # Requires balance weekends
             # Don't assign multiple shifts every (n) weeks
-            "max_night_frequency": 4,
+            "max_night_frequency": 1,
             "max_shifts_per_week": 7,
             "max_shifts_per_month": 40,
             # The following will cause an unsolvable problem if a shift
@@ -635,9 +637,8 @@ class RotaBuilder(object):
             else:
                 work_request_sets.add(work_request[3])
 
-
         # Validate shifts are valid
-        #work_request_sets = set([i[3] for i in self.work_requests])
+        # work_request_sets = set([i[3] for i in self.work_requests])
         invalid_shifts = work_request_sets - set(self.get_shift_names())
         if invalid_shifts:
             raise InvalidShift(
@@ -1433,32 +1434,36 @@ class RotaBuilder(object):
                                 )
 
             if self.constraint_options["balance_blocks"]:
-                for week_blocks in self.get_week_block_iterator(
+                if self.constraint_options["max_night_frequency"]:
-                    self.constraint_options["max_night_frequency"]
+                    for week_blocks in self.get_week_block_iterator(
-                ):
+                        self.constraint_options["max_night_frequency"]
-                    if self.get_shifts_with_constraint("night"):
+                    ):
-                        # Prevent nights more than once every n weeks
+                        if self.get_shifts_with_constraint("night"):
-                        self.model.constraints.add(
+                            # Prevent nights more than once every n weeks
-                            1
+                            self.model.constraints.add(
-                            >= sum(
+                                1
-                                self.model.shift_week_worker_assigned[
+                                >= sum(
-                                    shift.name, week, worker.id
+                                    self.model.shift_week_worker_assigned[
-                                ]
+                                        shift.name, week, worker.id
-                                for week in week_blocks
+                                    ]
-                                for shift in self.get_shifts_with_constraint("night")
+                                    for week in week_blocks
+                                    for shift in self.get_shifts_with_constraint("night")
+                                )
                             )
-                        )
 
             # if self.constraint_options["balance_weekends"]:
-            for week_blocks in self.get_week_block_iterator(self.constraint_options["max_weekend_frequency"]):
+            if self.constraint_options["max_weekend_frequency"]:
-                # Prevent weekend shifts more than once every n weeks
+                for week_blocks in self.get_week_block_iterator(
-                self.model.constraints.add(
+                    self.constraint_options["max_weekend_frequency"]
-                    1
+                ):
-                    >= sum(
+                    # Prevent weekend shifts more than once every n weeks
-                        self.model.works_weekend[worker.id, week]
+                    self.model.constraints.add(
-                        for week in week_blocks
+                        1
+                        >= sum(
+                            self.model.works_weekend[worker.id, week]
+                            for week in week_blocks
+                        )
                     )
-                )
 
             for constraint_shift in self.get_shifts_with_constraints(
                 "max_shifts_per_week_block"
@@ -1504,7 +1509,6 @@ class RotaBuilder(object):
                         )
                     )
 
-                
                 self.model.constraints.add(
                     self.constraint_options["max_shifts_per_week"]
                     >= sum(
@@ -1667,7 +1671,7 @@ class RotaBuilder(object):
                 #     self.model.shift_week_worker_assigned[s.name, week, worker.id]
                 #     for s in self.get_shifts_with_constraint("night")))
 
-                #for shift in self.get_shifts_with_constraint("single_block_per_week"):
+                # for shift in self.get_shifts_with_constraint("single_block_per_week"):
                 #    if shift.force_as_block:
                 #        # Force nights to be assigned in blocks
                 #        # self.model.constraints.add(8* self.model.shift_week_worker_assigned[shift.name, week, worker.id] >= sum(self.model.works[worker.id, week, day, shift.name] for day in self.days)
@@ -1709,64 +1713,171 @@ class RotaBuilder(object):
 
                 #
 
+            #def get_pre_may(week_days, max_pre)
+
             for n in range(len(weeks_days)):
 
                 week, day = weeks_days[n]
 
-                p1 = 1
+                pre_map = []
-                p2 = 1
+                for pre_n in range(1, self.max_pre + 1):
-                p3 = 1
+                    try:
-                p4 = 1
+                        pweek, pday = weeks_days[n - pre_n]
-                p5 = 1
+                        p = 1
-                if n > 0:
+                    except:
-                    pweek, pday = weeks_days[n - 1]
+                        pweek, pday = weeks_days[n]
-                else:
+                        p = 0
-                    p1 = 0
-                    pweek, pday = weeks_days[n]
 
-                try:
+                    pre_map.append((p, pweek, pday))
-                    p2week, p2day = weeks_days[n - 2]
-                except IndexError:
-                    p2week, p2day = weeks_days[n]
-                    p2 = 0
 
-                try:
+                post_map = []
-                    p3week, p3day = weeks_days[n - 3]
+                for post_n in range(1, self.max_post + 1):
-                except IndexError:
+                    try:
-                    p3week, p3day = weeks_days[n]
+                        pweek, pday = weeks_days[n + post_n]
-                    p3 = 0
+                        p = 1
+                    except:
+                        pweek, pday = weeks_days[n]
+                        p = 0
 
-                try:
+                    post_map.append((p, pweek, pday))
-                    p4week, p4day = weeks_days[n - 4]
-                except IndexError:
-                    p4week, p4day = weeks_days[n]
-                    p4 = 0
 
-                try:
+                # p1 = 1
-                    p5week, p5day = weeks_days[n - 5]
+                # p2 = 1
-                except IndexError:
+                # p3 = 1
-                    p5week, p5day = weeks_days[n]
+                # p4 = 1
-                    p5 = 0
+                # p5 = 1
+                # p6 = 1
+                # p7 = 1
+                # if n > 0:
+                #    pweek, pday = weeks_days[n - 1]
+                # else:
+                #    p1 = 0
+                #    pweek, pday = weeks_days[n]
 
-                n1 = 1
+                # try:
-                n2 = 1
+                #    p2week, p2day = weeks_days[n - 2]
-                try:
+                # except IndexError:
-                    nweek, nday = weeks_days[n + 1]
+                #    p2week, p2day = weeks_days[n]
-                except IndexError:
+                #    p2 = 0
-                    nweek, nday = weeks_days[n]
-                    n1 = 0
 
-                try:
+                # try:
-                    n2week, n2day = weeks_days[n + 2]
+                #    p3week, p3day = weeks_days[n - 3]
-                except IndexError:
+                # except IndexError:
-                    n2week, n2day = weeks_days[n]
+                #    p3week, p3day = weeks_days[n]
-                    n2 = 0
+                #    p3 = 0
 
-                try:
+                # try:
-                    n3week, n3day = weeks_days[n + 3]
+                #    p4week, p4day = weeks_days[n - 4]
-                except IndexError:
+                # except IndexError:
-                    n3week, n3day = weeks_days[n]
+                #    p4week, p4day = weeks_days[n]
-                    n3 = 0
+                #    p4 = 0
+
+                # try:
+                #    p5week, p5day = weeks_days[n - 5]
+                # except IndexError:
+                #    p5week, p5day = weeks_days[n]
+                #    p5 = 0
+
+                # try:
+                #    p6week, p6day = weeks_days[n - 6]
+                # except IndexError:
+                #    p6week, p6day = weeks_days[n]
+                #    p6 = 0
+
+                # try:
+                #    p7week, p7day = weeks_days[n - 7]
+                # except IndexError:
+                #    p7week, p7day = weeks_days[n]
+                #    p7 = 0
+
+                # pre_map = [
+                #    (p1, pweek, pday),
+                #    (p2, p2week, p2day),
+                #    (p3, p3week, p3day),
+                #    (p4, p4week, p4day),
+                #    (p5, p5week, p5day),
+                #    (p6, p6week, p6day),
+                #    (p7, p7week, p7day),
+                # ]
+
+                # n1 = 1
+                # n2 = 1
+                # n3 = 1
+                # n4 = 1
+                # n5 = 1
+                # n6 = 1
+                # n7 = 1
+                # if n > 0:
+                #    nweek, nday = weeks_days[n - 1]
+                # else:
+                #    n1 = 0
+                #    nweek, nday = weeks_days[n]
+
+                # try:
+                #    n2week, n2day = weeks_days[n - 2]
+                # except IndexError:
+                #    n2week, n2day = weeks_days[n]
+                #    n2 = 0
+
+                # try:
+                #    n3week, n3day = weeks_days[n - 3]
+                # except IndexError:
+                #    n3week, n3day = weeks_days[n]
+                #    n3 = 0
+
+                # try:
+                #    n4week, n4day = weeks_days[n - 4]
+                # except IndexError:
+                #    n4week, n4day = weeks_days[n]
+                #    n4 = 0
+
+                # try:
+                #    n5week, n5day = weeks_days[n - 5]
+                # except IndexError:
+                #    n5week, n5day = weeks_days[n]
+                #    n5 = 0
+
+                # try:
+                #    n6week, n6day = weeks_days[n - 6]
+                # except IndexError:
+                #    n6week, n6day = weeks_days[n]
+                #    n6 = 0
+
+                # try:
+                #    n7week, n7day = weeks_days[n - 7]
+                # except IndexError:
+                #    n7week, n7day = weeks_days[n]
+                #    n7 = 0
+
+                # post_map = [
+                #    (n1, nweek, nday),
+                #    (n2, n2week, n2day),
+                #    (n3, n3week, n3day),
+                #    (n4, n4week, n4day),
+                #    (n5, n5week, n5day),
+                #    (n6, n6week, n6day),
+                #    (n7, n7week, n7day),
+                # ]
+
+                # n1 = 1
+                # n2 = 1
+                # try:
+                #    nweek, nday = weeks_days[n + 1]
+                # except IndexError:
+                #    nweek, nday = weeks_days[n]
+                #    n1 = 0
+
+                # try:
+                #    n2week, n2day = weeks_days[n + 2]
+                # except IndexError:
+                #    n2week, n2day = weeks_days[n]
+                #    n2 = 0
+
+                # try:
+                #    n3week, n3day = weeks_days[n + 3]
+                # except IndexError:
+                #    n3week, n3day = weeks_days[n]
+                #    n3 = 0
 
                 # IF paired we check the following against both workers
                 workers = [worker]
@@ -1810,96 +1921,173 @@ class RotaBuilder(object):
                                 )
 
                 for constraint_shift in self.get_shifts_with_constraints(
-                    "preclear", "preclear2"
+                   "pre",
                 ):
-                    if day in constraint_shift.days:
+                    for n in range(0, constraint_shift.constraint_options["pre"]):
-                        self.model.constraints.add(
+                        if day in constraint_shift.days:
-                            1
+                            self.model.constraints.add(
-                            >= self.model.works[
+                                1
-                                worker.id, week, day, constraint_shift.name
+                                >= self.model.works[
-                            ]
+                                    worker.id, week, day, constraint_shift.name
-                            + sum(
+                                ]
-                                p1 * self.model.works[w.id, pweek, pday, shiftname]
+                                + sum(
-                                for shiftname in self.get_shift_names_by_week_day(
+                                    pre_map[n][0] * self.model.works[w.id, pre_map[n][1], pre_map[n][2], shiftname]
-                                    pweek, pday
+                                    for shiftname in self.get_shift_names_by_week_day(
+                                        pre_map[n][1], pre_map[n][2]
+                                    )
+                                    if shiftname != constraint_shift.name
+                                    for w in workers
                                 )
-                                if shiftname != constraint_shift.name
-                                for w in workers
                             )
-                        )
+                #for constraint_shift in self.get_shifts_with_constraint( "pre"):
+                #    to_sum = []
+                #    for n in range(
+                #        0, constraint_shift.constraint_options["pre"]
+                #    ):
+                #        for w in workers:
+                #            for wk, dy, shiftname in self.get_week_day_shift_names_by_week_day(
+                #                pre_map[n][1], pre_map[n][2]
+                #            ):
+                #                if shiftname != constraint_shift.name:
+                #                    to_sum.append(pre_map[n][0]
+                #                    * self.model.works[
+                #                        w.id, wk, dy, shiftname
+                #                    ])
 
-                for constraint_shift in self.get_shifts_with_constraint("preclear2"):
+
-                    if day in constraint_shift.days:
+                #    if day in constraint_shift.days:
-                        self.model.constraints.add(
+                #        self.model.constraints.add(
-                            1
+                #            1
-                            >= self.model.works[
+                #            >= self.model.works[
-                                worker.id, week, day, constraint_shift.name
+                #                worker.id, week, day, constraint_shift.name
-                            ]
+                #            ]
-                            + sum(
+                #            + sum(
-                                p2 * self.model.works[w.id, p2week, p2day, shiftname]
+                #                to_sum
-                                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"
+                    "post",
                 ):
-                    if day in constraint_shift.days:
+                    for n in range(0, constraint_shift.constraint_options["post"]):
-                        self.model.constraints.add(
+                        if day in constraint_shift.days:
-                            1
+                            self.model.constraints.add(
-                            >= self.model.works[
+                                1
-                                worker.id, week, day, constraint_shift.name
+                                >= self.model.works[
-                            ]
+                                    worker.id, week, day, constraint_shift.name
-                            + sum(
+                                ]
-                                n1 * self.model.works[w.id, nweek, nday, shiftname]
+                                + sum(
-                                for shiftname in self.get_shift_names_by_week_day(
+                                    post_map[n][0]
-                                    nweek, nday
+                                    * self.model.works[
+                                        w.id, post_map[n][1], post_map[n][2], shiftname
+                                    ]
+                                    for shiftname in self.get_shift_names_by_week_day(
+                                        post_map[n][1], post_map[n][2]
+                                    )
+                                    if shiftname != constraint_shift.name
+                                    for w in workers
                                 )
-                                if shiftname != constraint_shift.name
-                                for w in workers
                             )
-                        )
 
-                for constraint_shift in self.get_shifts_with_constraints("postclear2"):
+                #                for constraint_shift in self.get_shifts_with_constraints(
-                    if day in constraint_shift.days:
+                #                    "preclear", "preclear2"
-                        self.model.constraints.add(
+                #                ):
-                            1
+                #                    if day in constraint_shift.days:
-                            >= self.model.works[
+                #                        self.model.constraints.add(
-                                worker.id, week, day, constraint_shift.name
+                #                            1
-                            ]
+                #                            >= self.model.works[
-                            + sum(
+                #                                worker.id, week, day, constraint_shift.name
-                                n2 * self.model.works[w.id, n2week, n2day, shiftname]
+                #                            ]
-                                for shiftname in self.get_shift_names_by_week_day(
+                #                            + sum(
-                                    n2week, n2day
+                #                                p1 * self.model.works[w.id, pweek, pday, shiftname]
-                                )
+                #                                for shiftname in self.get_shift_names_by_week_day(
-                                if shiftname != constraint_shift.name
+                #                                    pweek, pday
-                                for w in workers
+                #                                )
-                            )
+                #                                if shiftname != constraint_shift.name
-                        )
+                #                                for w in workers
+                #                            )
+                #                        )
+                #
+                #                for constraint_shift in self.get_shifts_with_constraint("preclear2"):
+                #                    if day in constraint_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.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.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
+                #            )
+                #        )
 
                 # Night constraint means we won't assign a shift the day before
                 # an unavailability
                 for constraint_shift in self.get_shifts_with_constraint("night"):
                     if (
                         worker.id,
-                        pweek,
+                        # pweek,
-                        pday,
+                        pre_map[0][1],
+                        # pday,
+                        pre_map[0][2],
                         constraint_shift.name,
                     ) in self.model.works:
                         # Ensure night prior to unavalibity is not assigned
                         self.model.constraints.add(
                             self.model.available[worker.id, week, day]
                             >= self.model.works[
-                                worker.id, pweek, pday, constraint_shift.name
+                                # worker.id, pweek, pday, constraint_shift.name
+                                worker.id,
+                                pre_map[0][1],
+                                pre_map[0][2],
+                                constraint_shift.name,
                             ]
                         )
 
         self.define_objectives()
 
+        print("Building model completed")
+
     def define_objectives(self):
         # Define an objective function with model as input, to pass later
         def obj_rule(m):
@@ -2171,6 +2359,15 @@ class RotaBuilder(object):
         """
         return self.week_day_shifts_dict[(week, day)]
 
+    def get_week_day_shift_names_by_week_day(self, week, day: DayStr) -> set:
+        """Returns the shifts required for a specific day
+
+        Returns:
+            set: Set of ShiftName
+
+        """
+        return set((week, day, shift_name) for  shift_name in self.week_day_shifts_dict[(week, day)])
+
     def get_shift_by_name(self, name: str) -> SingleShift:
         """
 
@@ -2240,6 +2437,14 @@ class RotaBuilder(object):
                         self.shift_counts[s.name] = self.shift_counts[s.name] + 1
         # print(self.week_day_shift_product)
 
+        self.max_pre = 1
+        self.max_post = 1
+        for shift in self.get_shifts_with_constraint("pre"):
+            self.max_pre = max(shift.constraint_options["pre"], self.max_pre)
+
+        for shift in self.get_shifts_with_constraint("post"):
+            self.max_post = max(shift.constraint_options["post"], self.max_post)
+
     #        # Todo replace with week_day.....
     #        self.day_shift_product = []
     #        self.day_shiftclass_product = []

test/test_nights2.py

@@ -0,0 +1,160 @@
+import datetime
+import pytest
+from rota.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, days
+
+from rota.workers import Worker
+
+def generate_basic_rota(weeks_to_rota=10):
+    start_date = datetime.date(2022, 3, 7)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+
+    # Add a few workers
+    Rota.add_workers(
+        [
+            Worker(name="worker1", site="group1", grade=1),
+            Worker(name="worker2", site="group1", grade=1),
+        ]
+    )
+
+    return Rota
+
+
+def date_generator(from_date, days):
+    n = 0
+    while True:
+        yield from_date
+
+        n = n + 1
+        if n >= days:
+            break
+        from_date = from_date + datetime.timedelta(days=1)
+
+
+class TestNightShifts:
+    def test_nights(self):
+        Rota = generate_basic_rota()
+
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="n", length= 12.5, days=days[5:],
+                force_as_block=True,
+                constraint=[{"name": "night"}, {"name": "pre", "options": 2}],
+                workers_required=2,
+            ),
+            SingleShift(
+                sites=("group1", "group2"), name="b", length= 12.5, days=days[:3],
+                force_as_block=False,
+                #constraint=[{"name": "night"}],
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("night")
+
+        assert Rota.results.solver.status == "ok"
+        assert Rota.results.solver.termination_condition == "optimal"
+
+    def test_nights_pre3(self):
+        Rota = generate_basic_rota()
+
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="n", length= 12.5, days=days[5:],
+                force_as_block=True,
+                constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
+                workers_required=2,
+            ),
+            SingleShift(
+                sites=("group1", "group2"), name="b", length= 12.5, days=days[:2],
+                force_as_block=False,
+                #constraint=[{"name": "night"}],
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("night_pre3")
+
+        assert Rota.results.solver.status == "ok"
+
+    def test_nights_pre3_2(self):
+        Rota = generate_basic_rota()
+
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="n", length= 12.5, days=days[5:],
+                force_as_block=True,
+                constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
+                workers_required=1,
+            ),
+            SingleShift(
+                sites=("group1", "group2"), name="b", length= 12.5, days=days[:3],
+                force_as_block=False,
+                #constraint=[{"name": "night"}],
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("night_pre3")
+
+        assert Rota.results.solver.status == "ok"
+
+    def test_nights_fail(self):
+        Rota = generate_basic_rota()
+
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="n", length= 12.5, days=days[5:],
+                force_as_block=True,
+                constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
+                workers_required=2,
+            ),
+            SingleShift(
+                sites=("group1", "group2"), name="b", length= 12.5, days=days[:3],
+                force_as_block=False,
+                #constraint=[{"name": "night"}],
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+
+        assert Rota.results.solver.status == "warning"
+
+    def test_nights2(self):
+        Rota = generate_basic_rota()
+
+        Rota.constraint_options["max_night_frequency"] = 0
+
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="n", length= 12.5, days=days[5:],
+                force_as_block=True,
+                constraint=[{"name": "night"}],
+                workers_required=2,
+            ),
+            SingleShift(
+                sites=("group1", "group2"), name="m", length= 12.5, days=days[:5],
+                force_as_block=True,
+                #constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
+                #constraint=[{"name": "night"}],
+                workers_required=2,
+            ),
+            #SingleShift(
+            #    sites=("group1", "group2"), name="b", length= 12.5, days=days[:3],
+            #    force_as_block=False,
+            #    #constraint=[{"name": "night"}],
+            #),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("night")
+
+        assert Rota.results.solver.status == "ok"

test/test_rota.py

@@ -203,6 +203,7 @@ class TestDemoRotaNights:
     Rota.constraint_options["balance_nights"] = False
     Rota.constraint_options["minimise_shift_diffs"] = False
     Rota.constraint_options["balance_blocks"] = False
+    Rota.constraint_options["max_night_frequency"] = 0
     # Rota.constraint_options["balance_weekends"] = False
 
     # Add a few workers
@@ -227,7 +228,7 @@ class TestDemoRotaNights:
             length=12.5,
             days=days[:4],
             workers_required=1,
-            constraint=[{"name": "preclear2"}, {"name": "postclear2"}],
+            constraint=[{"name": "pre", "options": 1}, {"name": "post", "options": 2}],
             # constraint=["night"],
             force_as_block=True
             # We use a different balance weighting for each shift as otherwise
@@ -240,7 +241,7 @@ class TestDemoRotaNights:
             length=12.5,
             days=days[4:],
             # balance_weighting=0.5
-            constraint=[{"name": "preclear2"}, {"name": "postclear2"}],
+            constraint=[{"name": "pre", "options": 1}, {"name": "post", "options": 2}],
             # constraint=["night"],
             force_as_block=True,
         ),
@@ -274,6 +275,8 @@ class TestDemoRotaNights:
         assert self.Rota.start_date == self.start_date
 
     def test_night_assignment(self):
+        assert self.Rota.results.solver.status == "ok"
+
         shift_summary = self.Rota.get_shift_summary_dict()
 
         for worker_name in shift_summary:
@@ -316,7 +319,7 @@ class TestDemoRotaClear:
         (worker1, worker2, worker3, worker4, worker5, worker6, worker7, worker8)
     )
 
-    def test_preclear(self):
+    def test_pre(self):
         self.Rota.add_shifts(
             SingleShift(
                 sites=("group1", "group2"),
@@ -325,7 +328,7 @@ class TestDemoRotaClear:
                 days=days[:5],
                 balance_offset=40,
                 workers_required=2,
-                constraint=[{"name": "preclear"}],
+                constraint=[{"name": "pre", "options": 1}],
             ),
             SingleShift(
                 sites=("group1", "group2"),
@@ -342,7 +345,7 @@ class TestDemoRotaClear:
                 days=days[3],
                 balance_offset=40,
                 workers_required=1,
-                constraint=[{"name": "preclear2"}, {"name": "postclear2"}],
+                constraint=[{"name": "pre", "options": 1}, {"name": "post", "options": 1}],
             ),
         )
 
@@ -369,12 +372,72 @@ class TestDemoRotaClear:
             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 "a-c" not in shifts_string
-            assert "c-a" 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 "b-c" not in shifts_string
-            assert "c-b" not in shifts_string
+            #assert "c-b" not in shifts_string
+
+    #def test_preclear(self):
+    #    self.Rota.add_shifts(
+    #        SingleShift(
+    #            sites=("group1", "group2"),
+    #            name="a",
+    #            length=12.5,
+    #            days=days[:5],
+    #            balance_offset=40,
+    #            workers_required=2,
+    #            constraint=[{"name": "preclear"}],
+    #        ),
+    #        SingleShift(
+    #            sites=("group1", "group2"),
+    #            name="b",
+    #            length=12.5,
+    #            days=days[:5],
+    #            balance_offset=40,
+    #            workers_required=2,
+    #        ),
+    #        SingleShift(
+    #            sites=("group1", "group2"),
+    #            name="c",
+    #            length=12.5,
+    #            days=days[3],
+    #            balance_offset=40,
+    #            workers_required=1,
+    #            constraint=[{"name": "pre", "options": 2}, {"name": "post", "options": 2}],
+    #        ),
+    #    )
+
+    #    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:
@@ -1045,7 +1108,7 @@ class TestNightUnavailable:
 
     Rota.add_workers((worker1, worker2))
 
-    def test_basic_assignment(self, capsys):
+    def test_basic_assignment(self):
         self.Rota.shifts = []
 
         # Add a weekday and weekend shift
@@ -1074,7 +1137,6 @@ class TestNightUnavailable:
     def test_assign_night_prior_to_unavailablity(self):
         self.Rota.shifts = []
 
-        # Add a weekday and weekend shift
         self.Rota.add_shifts(
             SingleShift(
                 sites =("group1",),
@@ -1087,10 +1149,30 @@ class TestNightUnavailable:
         )
 
         self.Rota.build_and_solve(options={"ratio": 0.00})
+        self.Rota.export_rota_to_html("testnight_unavail")
 
         assert self.Rota.results.solver.status == "warning"
         assert self.Rota.results.solver.termination_condition == "infeasible"
 
+    def test_assign_non_night_prior_to_unavailablity(self):
+        self.Rota.shifts = []
+
+        self.Rota.add_shifts(
+            SingleShift(
+                sites =("group1",),
+                name = "night_weekend",
+                length = 12.5,
+                days = days[5:],
+                workers_required=2,
+                #constraint=[{"name": "night"}],
+            ),
+        )
+
+        self.Rota.build_and_solve(options={"ratio": 0.00})
+        self.Rota.export_rota_to_html("testnight_unavail")
+
+        assert self.Rota.results.solver.status == "ok"
+
     def test_assign_prior_to_unavailablity_non_night(self):
         self.Rota.shifts = []
 

test/test_shifts.py

@@ -173,3 +173,142 @@ class TestShiftConstraints:
         Rota.export_rota_to_html("max_shifts_per_week_block")
 
         assert Rota.results.solver.status == "warning"
+
+    def test_pre_shift_constraint(self):
+        Rota = generate_basic_rota()
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="a", length= 12.5, days=days[0],
+                constraint=[{"name": "pre", "options": 1}],
+            ),
+            SingleShift(
+                sites=("group1", "group2"), name="b", length= 12.5, days=days[1],
+                constraint=[{"name": "pre", "options": 1}],
+            ),
+            SingleShift(
+                sites=("group1", "group2"), name="c", length= 12.5, days=days[2],
+                constraint=[{"name": "pre", "options": 1}],
+            ),
+            SingleShift(
+                sites=("group1", "group2"), name="d", length= 12.5, days=days[3],
+                constraint=[{"name": "pre", "options": 1}],
+            ),
+            SingleShift(
+                sites=("group1", "group2"), name="e", length= 12.5, days=days[4],
+                constraint=[{"name": "pre", "options": 1}],
+            ),
+            SingleShift(
+                sites=("group1", "group2"), name="f", length= 12.5, days=days[5],
+                constraint=[{"name": "pre", "options": 1}],
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("pre")
+
+        assert Rota.results.solver.status == "ok"
+
+        for worker in Rota.get_workers():
+            shift_string = Rota.get_worker_shift_list_string(worker)
+
+            print(shift_string)
+            for s in ["ab", "bc", "cd", "de", "ef", "fg"]:
+                assert s not in shift_string
+
+            
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="g", length= 12.5, days=days[6],
+                constraint=[{"name": "pre", "options": 2}],
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+
+        assert Rota.results.solver.status == "warning"
+
+    def test_pre_shift_constraint2(self):
+        Rota = generate_basic_rota()
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="a", length= 12.5, days=days[0],
+                constraint=[{"name": "pre", "options": 2}],
+                workers_required=2
+            ),
+            SingleShift(
+                sites=("group1", "group2"), name="b", length= 12.5, days=days[4],
+                constraint=[{"name": "pre", "options": 3}],
+                workers_required=2
+            ),
+
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("pre")
+
+        assert Rota.results.solver.status == "ok"
+
+    def test_pre_shift_constraint3(self):
+        Rota = generate_basic_rota()
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="a", length= 12.5, days=days[0],
+                constraint=[{"name": "pre", "options": 3}],
+                workers_required=2
+            ),
+            SingleShift(
+                sites=("group1", "group2"), name="b", length= 12.5, days=days[4],
+                constraint=[{"name": "pre", "options": 3}],
+                workers_required=2
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("pre")
+
+        assert Rota.results.solver.status == "warning"
+
+    def test_post_shift_constraint(self):
+        Rota = generate_basic_rota()
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="a", length= 12.5, days=days[0],
+                constraint=[{"name": "post", "options": 3}],
+                workers_required=2
+            ),
+            SingleShift(
+                sites=("group1", "group2"), name="b", length= 12.5, days=days[4],
+                constraint=[{"name": "post", "options": 2}],
+                workers_required=2
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("post")
+
+        assert Rota.results.solver.status == "ok"
+
+    def test_post_shift_constraint2(self):
+        Rota = generate_basic_rota()
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="a", length= 12.5, days=days[0],
+                constraint=[{"name": "post", "options": 3}],
+                workers_required=2
+            ),
+            SingleShift(
+                sites=("group1", "group2"), name="b", length= 12.5, days=days[4],
+                constraint=[{"name": "post", "options": 3}],
+                workers_required=2
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("post")
+
+        assert Rota.results.solver.status == "warning"