start pydanitification

rota/runfixnights.py

@@ -227,7 +227,7 @@ if use_test_workers:
         ]
     )
     Rota.add_workers(
-        [Worker(Rota, i, "Plym {}".format(i), "plymouth", 4, 50) for i in range(9, 17)]
+        [Worker(i, "Plym {}".format(i), "plymouth", 4, 50) for i in range(9, 17)]
     )
     Rota.add_workers(
         [

rota/shifts.py

@@ -40,7 +40,7 @@ for bank_holiday in bank_holidays.get_holidays(division="england-and-wales"):
 SHIFT_BOUNDS = {
     "bank_holiday": (0, 9),
     "shift_count": (0, 400),
-    "night_shift_count": (0, 19),
+    "night_shift_count": (0, 100),
     "weekend_count": (0, 60),
 }
 
@@ -55,7 +55,7 @@ class SingleShift(object):
 
     postclear(2) / preclear(2)
 
-    require_remote_site_presence:
+    require_remote_site_presence_week:
         options: (site, required_number)
 
     night
@@ -185,6 +185,7 @@ class RotaBuilder(object):
         self.unavailable_to_work = set()
         self.unavailable_to_work_reason = {}
         self.pref_not_to_work = {}
+        self.pref_not_to_work_reason = {}
         self.work_requests = set()
         self.work_requests_map = {}
 
@@ -203,14 +204,13 @@ class RotaBuilder(object):
         self.max_weekend_frequency = max_weekend_frequency
 
         self.constraint_options = {
-            "ensure_1_st4_plus_on_nights": True,
             "balance_nights": True,
             "constrain_time_off_after_nights": False,
             "balance_nights_across_sites": True,
             "balance_bank_holidays": True,
             "balance_blocks": True,
             "balance_shifts": True,  # Does not use a quadratic function
-            "balance_shifts_quadratic": False,
+            "balance_shifts_quadratic": False, # Will prevent spreading of spreading across different shifts
             "balance_shifts_over_workers": True,
             "minimise_shift_diffs": False,  # less sophisticated version of balance_shifts_over_workers
             "balance_weekends": True,
@@ -223,8 +223,6 @@ class RotaBuilder(object):
             "prevent_monday_and_tuesday_after_full_weekends": [],
             "prevent_fridays_before_full_weekends": [], 
             "prevent_thursdays_before_full_weekends": [],
-            # This may be better defined on the shift?
-            "require_presence_at_site_overnight": [],
             "avoid_st2_first_month": False,
             "hard_constrain_pair_separation": False,
         }
@@ -804,7 +802,7 @@ class RotaBuilder(object):
                 ),
             )
 
-        def presenceAtRemoteSite(model, week, shift, required_site, required_number):
+        def presenceAtRemoteSiteWeek(model, week, shift, required_site, required_number):
             required_site_workers = [
                 w for w in self.workers if required_site == w.remote_site
             ]
@@ -817,6 +815,40 @@ class RotaBuilder(object):
                 >= required_number
             )
 
+        for shift in self.get_shifts_with_constraint("require_remote_site_presence_week"):
+            site, required_number = shift.constraint_options[
+                "require_remote_site_presence_week"
+            ]
+            # self.model.require_presence_at_site_overnight_rule = Constraint(
+            setattr(
+                self.model,
+                f"require_remote_site_presence_week_{shift.name}",
+                Constraint(
+                    [week for week in self.weeks],
+                    [shift.name],
+                    [site],
+                    [required_number],
+                    rule=presenceAtRemoteSiteWeek,
+                ),
+            )
+
+        def presenceAtRemoteSite(model, day, week, shift, required_site, required_number):
+            required_site_workers = [
+                w for w in self.workers if required_site == w.remote_site
+            ]
+
+            for w in required_site_workers:
+                if (w.id, week, day, shift) not in model.works:
+                    return Constraint.Skip
+
+            return (
+                sum(
+                    model.works[w.id, week, day, shift]
+                    for w in required_site_workers
+                )
+                >= required_number
+            )
+
         for shift in self.get_shifts_with_constraint("require_remote_site_presence"):
             site, required_number = shift.constraint_options[
                 "require_remote_site_presence"
@@ -826,6 +858,7 @@ class RotaBuilder(object):
                 self.model,
                 f"require_remote_site_presence_{shift.name}",
                 Constraint(
+                    [day for day in self.days],
                     [week for week in self.weeks],
                     [shift.name],
                     [site],
@@ -912,10 +945,9 @@ class RotaBuilder(object):
                     full_workers = []
 
                     for w in workers:
-                        if w.nwd:
+                        if w.non_working_day_list:
                             l = []
-                            # This should take into account dates!
-                            for nwd, start_nwd_date, end_nwd_date in w.nwd:
+                            for nwd, start_nwd_date, end_nwd_date in w.non_working_day_list:
                                 if nwd in shift.shift_days:
                                     if start_nwd_date > self.get_week_start_date(week):
                                         continue
@@ -1121,7 +1153,7 @@ class RotaBuilder(object):
 
                 if self.constraint_options["balance_shifts_quadratic"]:
                     # This may need to be updated
-                    xU = 25
+                    xU = 10
                     xL = 1
                     self.model.constraints.add(
                         inequality(
@@ -1333,9 +1365,9 @@ class RotaBuilder(object):
                     )
 
             # Ensure worker is not allocated shifts on non working days
-            if worker.nwd:
+            if worker.non_working_day_list:
                 for week, day, shift in self.get_all_shiftclass_combinations():
-                    for n, start, end in worker.nwd:
+                    for n, start, end in worker.non_working_day_list:
                         if not shift.rota_on_nwds and day == n:
                             # print(start, self.week_day_date_map[(week, day)], end)
                             if start <= self.week_day_date_map[(week, day)] < end:
@@ -1932,7 +1964,9 @@ class RotaBuilder(object):
         Args:
             workers (List(Worker)):
         """
-        self.workers.extend(workers)
+        for worker in workers:
+            self.add_worker(worker)
+        #self.workers.extend(workers)
 
     def build_workers(self) -> None:
         """Process loaded shifts and workers
@@ -1947,6 +1981,7 @@ class RotaBuilder(object):
         self.workers_name_map = {}
 
         for worker in self.workers:
+            worker.load_rota(self)
             wid = worker.id
             if wid in self.workers_id_map:
                 raise ValueError(f"Worker with id '{wid}' has been added twice")
@@ -2251,16 +2286,23 @@ class RotaBuilder(object):
         return self.workers_name_map[name]
 
     def get_workers_by_group(self) -> Dict[str, Worker]:
-        group_workers = defaultdict(set)
+        group_workers = defaultdict(list)
         for worker in self.workers:
-            group_workers[worker.site].add(worker)
+            group_workers[worker.site].append(worker)
+
+        return group_workers
+
+    def get_workers_by_remote_group(self) -> Dict[str, Worker]:
+        group_workers = defaultdict(list)
+        for worker in self.workers:
+            group_workers[worker.remote_site].append(worker)
 
         return group_workers
 
     def get_workers_by_grade(self) -> Dict[str, Worker]:
-        group_workers = defaultdict(set)
+        group_workers = defaultdict(list)
         for worker in self.workers:
-            group_workers[worker.grade].add(worker)
+            group_workers[worker.grade].append(worker)
 
         return group_workers
 
@@ -2462,7 +2504,7 @@ class RotaBuilder(object):
                 current_site = worker.site
 
             shifts = []
-            nwds = json.dumps(worker.nwd, default=str)
+            nwds = json.dumps(worker.non_working_day_list, default=str)
             #if worker.nwd:
             #    # TODO: limit to dates
             #    nwds = ", ".join([i[0] for i in worker.nwd])
@@ -2525,8 +2567,8 @@ class RotaBuilder(object):
                 remote_site = ""
                 if shift_name:
                     shift = self.get_shift_by_name(shift_name)
-                    if "require_remote_site_presence" in shift.constraints:
-                        remote_site = f" data-shift-remote-site='{shift.constraint_options['require_remote_site_presence'][0]}'"
+                    if "require_remote_site_presence_week" in shift.constraints:
+                        remote_site = f" data-shift-remote-site='{shift.constraint_options['require_remote_site_presence_week'][0]}'"
 
                 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}'{remote_site}{requests}{bank_holiday}>{a}</td>"
@@ -2564,8 +2606,8 @@ class RotaBuilder(object):
                 name=worker.name,
                 fte=worker.fte,
                 fte_adj=worker.fte_adj,
-                start_date=worker.start_date,
-                end_date=worker.end_date,
+                start_date=worker.calculated_start_date,
+                end_date=worker.calculated_end_date,
                 targets=worker_targets,
                 worker_shift_counts=json.dumps(shift_count_dict),
                 weekend_target=worker.weekend_shift_target_number,
@@ -2698,6 +2740,15 @@ class RotaBuilder(object):
 
         return shifts
 
+    def get_workers_total_shifts(self) -> Dict[str, int]:
+        shifts = {}
+
+        for worker in self.workers:
+            shifts[worker.name] = len([i for i in self.get_worker_shift_list(worker) if i != ""])
+
+        return shifts
+
+
 
     def get_worker_shift_list(self, worker: Worker) -> List:
         shifts = []

rota/workers.py

@@ -1,6 +1,9 @@
 import datetime
 from collections import defaultdict
-from typing import Iterable, List
+from typing import Iterable, List, Literal, Optional
+from pydantic import BaseModel, Extra
+
+from rich.pretty import pprint
 
 # from .shifts import RotaBuilder, days, sites
 import uuid
@@ -11,175 +14,204 @@ from typing import TYPE_CHECKING
 if TYPE_CHECKING:
     from rota.shifts import RotaBuilder
 
+days = Literal["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]
 
-class Worker:
-    def __init__(
-        self,
-        Rota: "RotaBuilder",  # Aim to remove (workers should be independent of rotas)
-        name: str,
-        site: str,
-        grade: int,
-        id=None,
-        fte: int = 100,
-        nwd: list[str] | list[str, datetime.datetime, datetime.datetime] = [],
-        start_date: datetime.datetime | None = None,
-        end_date: datetime.datetime | None = None,
-        oop: list[(datetime.datetime, datetime.datetime, str)] | None = None,
-        not_available_to_work: list[(datetime.datetime, str)] | None = None,
-        pref_not_to_work: list[datetime.datetime] | None = None,
-        work_requests: list[(datetime.datetime, str)] | None = None,
-        remote_site: str = "plymouth",  # We set a default proc_site
-        previous_shifts: dict = {},
-        shift_balance_extra: dict = {},
-        bank_holiday_extra: int = 0,
-        pair: int = 0,
-    ):
+class NotAvailableToWork(BaseModel):
+    date: datetime.date
+    reason: str = "unknown"
 
-        # We can either have a user generated ID
-        if id is not None:
-            self.id = id
-        else:
-            self.id = uuid.uuid4()
-        self.name = name
-        self.site = site
-        # Grade are equivalent to roles, by keeping them integer defining model
-        # rules is easier.
-        self.grade = grade
-        self.fte = fte
-        self.nwd = []
+class NonWorkingDays(BaseModel):
+    day: days
+    start_date: datetime.date | None = None
+    end_date: datetime.date | None = None
+
+class WorkRequests(BaseModel):
+    date: datetime.date
+    shift: str = ""
+
+class PreferenceNotToWork(BaseModel):
+    date: datetime.date
+    reason: str = ""
+
+class OutOfProgramme(BaseModel):
+    start_date: datetime.date
+    end_date: datetime.date
+    reason: str = ""
+
+class Worker(BaseModel):
+    name: str
+    site: str
+    # Grade are equivalent to roles, by keeping them integer defining model
+    # rules is easier.
+    grade: int
+    # We can either have a user generated ID
+    id: Optional[int]=None
+    fte: int = 100
+    nwds: list[NonWorkingDays] = []
+    start_date: datetime.date | None = None
+    end_date: datetime.date | None = None
+    oop: list[OutOfProgramme] = []
+    not_available_to_work: list[NotAvailableToWork] = []
+    pref_not_to_work: list[PreferenceNotToWork] = []
+    work_requests: list[WorkRequests] = []
+    remote_site: str = "plymouth",  # We set a default proc_site
+    previous_shifts: dict = {}
+    shift_balance_extra: dict = {}
+    bank_holiday_extra: int = 0
+    pair: int = 0
+
+    class Config:
+        extra = Extra.allow
+
+#    def __init__(
+#        self,
+#        #Rota: "RotaBuilder",  # Aim to remove (workers should be independent of rotas)
+#    ):
+#
+#        self.name = name
+#        self.site = site
+#        self.grade = grade
+#        self.fte = fte
+#        self.nwd = nwd
+#        self.pair = pair
+#        self.shift_balance_extra = shift_balance_extra
+#        self.bank_holiday_extra = bank_holiday_extra
+#        self.start_date
+#
+#        self.remote_site = remote_site
+#
+#        self.previous_shifts = previous_shifts
+#
+#
+#        # days_to_work = Rota.rota_days_length
+
+    def load_rota(self, Rota):
         self.proportion_rota_to_work = 1
-        self.pair = pair
-        self.shift_balance_extra = shift_balance_extra
-        self.bank_holiday_extra = bank_holiday_extra
-
-        self.remote_site = remote_site
-
-        self.previous_shifts = previous_shifts
-
         self.shift_target_number = defaultdict(int)
 
-        # days_to_work = Rota.rota_days_length
+        if self.id is None:
+            self.id = uuid.uuid4()
 
-        if start_date is not None and end_date is not None:
-            if start_date >= end_date:
+        if self.start_date is not None and self.end_date is not None:
+            if self.start_date >= self.end_date:
                 raise ValueError("End date must be after start date")
 
         # if no start date default to the start of the rota
-        if start_date is None:
-            self.start_date = Rota.start_date
+        if self.start_date is None:
+            self.calculated_start_date = Rota.start_date
         else:
             # ? test if start date is valid
-            self.start_date = start_date
+            self.calculated_start_date = self.start_date
 
-        if end_date is None:
-            self.end_date = Rota.rota_end_date
+        if self.end_date is None:
+            self.calculated_end_date = Rota.rota_end_date
         else:
-            self.end_date = end_date
-            # self.end_date = datetime.datetime.strptime(end_date, "%d/%m/%y").date()
+            self.calculated_end_date = self.end_date
+            # self.calculated_end_date = datetime.datetime.strptime(end_date, "%d/%m/%y").date()
 
-            if self.end_date > Rota.rota_end_date:
-                self.end_date = Rota.rota_end_date
+            if self.calculated_end_date > Rota.rota_end_date:
+                self.calculated_end_date = Rota.rota_end_date
 
-        days_to_work = (self.end_date - self.start_date).days
+        days_to_work = (self.calculated_end_date - self.calculated_start_date).days
 
         for week, day in Rota.weeks_days_product:
             date = Rota.week_day_date_map[(week, day)]
 
-            if date < self.start_date:
+            if date < self.calculated_start_date:
                 Rota.unavailable_to_work.add((self.id, week, day))
                 Rota.unavailable_to_work_reason[
                     (self.id, week, day)
-                ] = f"START DATE: {self.start_date}"
+                ] = f"START DATE: {self.calculated_start_date}"
 
-            if date >= self.end_date:
+            if date >= self.calculated_end_date:
                 Rota.unavailable_to_work.add((self.id, week, day))
                 Rota.unavailable_to_work_reason[
                     (self.id, week, day)
-                ] = f"END DATE: {self.end_date}"
+                ] = f"END DATE: {self.calculated_end_date}"
 
-        if oop is not None:
-            for start_oop, end_oop, oop_name in oop:
-                # start_oop, end_oop = oop
-                if isinstance(start_oop, datetime.date):
-                    start_oop_date = start_oop
+        for item in self.oop:
+            start_oop = item.start_date
+            end_oop = item.end_date
+            oop_name = item.reason
+            # start_oop, end_oop = oop
+            if isinstance(start_oop, datetime.date):
+                start_oop_date = start_oop
+            else:
+                start_oop_date = datetime.datetime.strptime(
+                    start_oop, "%d/%m/%y"
+                ).date()
+
+            if isinstance(end_oop, datetime.date):
+                end_oop_date = end_oop
+            else:
+                end_oop_date = datetime.datetime.strptime(
+                    end_oop, "%d/%m/%y"
+                ).date()
+
+            if start_oop_date >= end_oop_date:
+                raise ValueError("End OOP date must be after start date")
+
+            # ignore oops if they finish before the rota (or worker) start date
+            if end_oop_date > self.calculated_start_date:
+                if start_oop_date > self.calculated_end_date:
+                    pass
                 else:
-                    start_oop_date = datetime.datetime.strptime(
-                        start_oop, "%d/%m/%y"
-                    ).date()
+                    if end_oop_date > Rota.rota_end_date:
+                        end_oop_date = Rota.rota_end_date
 
-                if isinstance(end_oop, datetime.date):
-                    end_oop_date = end_oop
-                else:
-                    end_oop_date = datetime.datetime.strptime(
-                        end_oop, "%d/%m/%y"
-                    ).date()
+                    if start_oop_date < Rota.start_date:
+                        start_oop_date = Rota.start_date
 
-                if start_oop_date >= end_oop_date:
-                    raise ValueError("End OOP date must be after start date")
+                    oop_length = (end_oop_date - start_oop_date).days
+                    days_to_work = days_to_work - oop_length
 
-                # ignore oops if they finish before the rota (or worker) start date
-                if end_oop_date > self.start_date:
-                    if start_oop_date > self.end_date:
-                        pass
-                    else:
-                        if end_oop_date > Rota.rota_end_date:
-                            end_oop_date = Rota.rota_end_date
+                    days_until_oop = (start_oop_date - Rota.start_date).days
 
-                        if start_oop_date < Rota.start_date:
-                            start_oop_date = Rota.start_date
+                    for weeks_days in Rota.weeks_days_product[
+                        days_until_oop : days_until_oop + oop_length
+                    ]:
+                        week, day = weeks_days
+                        Rota.unavailable_to_work.add((self.id, week, day))
+                        Rota.unavailable_to_work_reason[
+                            (self.id, week, day)
+                        ] = f"OOP ({oop_name})".format(self.oop)
 
-                        oop_length = (end_oop_date - start_oop_date).days
-                        days_to_work = days_to_work - oop_length
+        # loop throught dates converting to week / day combination
+        for item in self.pref_not_to_work:
+            days_from_start = (
+                item.date
+                - Rota.start_date
+            ).days
+            week = days_from_start // 7 + 1
+            day = Rota.days[(days_from_start % 7)]
+            # Weight the value to take into account the number of preferences
+            # 1 is added to the total number of requests to ensure they do not outweight
+            # a single (or fewer) request(s)
+            Rota.pref_not_to_work[(self.id, week, day)] = 1 / (
+                len(self.pref_not_to_work) + 1
+            )
+            Rota.pref_not_to_work_reason[(self.id, week, day)] = item.reason
 
-                        days_until_oop = (start_oop_date - Rota.start_date).days
+        # print(not_available_to_work)
+        # loop throught dates converting to week / day combination
+        for item in self.not_available_to_work:
+            days_from_start = (
+                item.date
+                - Rota.start_date
+            ).days
+            week = days_from_start // 7 + 1
+            day = Rota.days[(days_from_start % 7)]
+            Rota.unavailable_to_work.add((self.id, week, day))
+            Rota.unavailable_to_work_reason[(self.id, week, day)] = item.reason
 
-                        for weeks_days in Rota.weeks_days_product[
-                            days_until_oop : days_until_oop + oop_length
-                        ]:
-                            week, day = weeks_days
-                            Rota.unavailable_to_work.add((self.id, week, day))
-                            Rota.unavailable_to_work_reason[
-                                (self.id, week, day)
-                            ] = f"OOP ({oop_name})".format(oop)
-
-        if pref_not_to_work is not None:
-            # loop throught dates converting to week / day combination
-            for date in pref_not_to_work:
+        for item in self.work_requests:
                 days_from_start = (
-                    date
+                    item.date
                     - Rota.start_date
                 ).days
                 week = days_from_start // 7 + 1
                 day = Rota.days[(days_from_start % 7)]
-                # Weight the value to take into account the number of preferences
-                # 1 is added to the total number of requests to ensure they do not outweight
-                # a single (or fewer) request(s)
-                Rota.pref_not_to_work[(self.id, week, day)] = 1 / (
-                    len(pref_not_to_work) + 1
-                )
-
-        if not_available_to_work is not None:
-            # print(not_available_to_work)
-            # loop throught dates converting to week / day combination
-            for date, reason in not_available_to_work:
-                days_from_start = (
-                    date
-                    - Rota.start_date
-                ).days
-                week = days_from_start // 7 + 1
-                day = Rota.days[(days_from_start % 7)]
-                Rota.unavailable_to_work.add((self.id, week, day))
-                Rota.unavailable_to_work_reason[(self.id, week, day)] = reason
-
-        if work_requests is not None:
-            for date, shift in work_requests:
-                days_from_start = (
-                    date
-                    - Rota.start_date
-                ).days
-                week = days_from_start // 7 + 1
-                day = Rota.days[(days_from_start % 7)]
-                Rota.work_requests.add((self.id, week, day, shift))
+                Rota.work_requests.add((self.id, week, day, item.shift))
 
         # Calculate the proportion of the rota that is being worked
         self.proportion_rota_to_work = days_to_work / Rota.rota_days_length
@@ -189,15 +221,20 @@ class Worker:
         self.fte_adj = self.fte * self.proportion_rota_to_work
 
         if self.fte_adj > 100:
+            # Shouldn't happen ? bug if it does
             raise ValueError("{} : fte_ajd = {}".format(self.name, self.fte_adj))
 
-        assert type(nwd) == list, f"nwd must be a list: {nwd=}"
-        for item in nwd:
-            match item:
-                case (nwd_day, start_date, end_date):
-                    self.nwd.append((nwd_day, start_date, end_date))
-                case nwd_day:
-                    self.nwd.append((nwd_day, Rota.start_date, Rota.rota_end_date))
+        # TODO: this has already been validated, consider moving
+        self.non_working_day_list = []
+        for item in self.nwds:
+            start_date = Rota.start_date
+            end_date = Rota.rota_end_date
+            if item.start_date is not None:
+                start_date = item.start_date
+            if item.end_date is not None:
+                end_date = item.end_date
+
+            self.non_working_day_list.append((item.day, start_date, end_date))
 
     def __lt__(self, other) -> bool:
         return (self.site, self.grade, self.fte_adj, self.name) < (
@@ -209,11 +246,11 @@ class Worker:
 
     def __str__(self) -> str:
 
-        nwd = ", ".join([str(i) for i in self.nwd]) if self.nwd is not None else ""
+        nwds = ", ".join([str(i) for i, start, end in self.non_working_day_list]) if self.non_working_day_list is not None else ""
         return "{} {} [{}] REMOTE SITE:{}".format(
             self.name,
             (self.site, self.grade, self.fte_adj),
-            nwd,
+            nwds,
             self.remote_site,
         )