import datetime from collections import defaultdict from typing import Iterable, List, Literal, Optional, Set from pydantic import BaseModel, ConfigDict, field_validator, Field from rich.pretty import pprint from rota_generator.console import console # from .shifts import RotaBuilder, days, sites import uuid # from rota.shifts import RotaBuilder from typing import TYPE_CHECKING if TYPE_CHECKING: from rota_generator.shifts import RotaBuilder days = Literal["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"] whole_days = [ "sunday", "monday", "tuesday", "wednesday", "thursday", "friday", "saturday", ] class NotAvailableToWork(BaseModel): date: datetime.date reason: str = "unknown" class HardDayDependency(BaseModel): days: Set[str] = Field(..., description="List of days for the dependency group") weeks: Optional[List[int]] = None # If None, applies to all weeks start_date: Optional[datetime.date] = None # Optional start date for the dependency end_date: Optional[datetime.date] = None # Optional end date for the dependency ignore_dates: list[datetime.date] = Field(default_factory=list, description="List of specific dates to ignore for this dependency") @field_validator("weeks") def weeks_none_if_dates_set(cls, weeks, values): data = values.data if (data.get("start_date") is not None or data.get("end_date") is not None) and weeks is not None: raise ValueError("If start_date or end_date is set, weeks must not be set") return weeks class HardDayExclusion(BaseModel): days: set[str] = Field(..., description="List of days to exclude together") weeks: Optional[List[int]] = None # If None, applies to all weeks start_date: Optional[datetime.date] = None # Optional start date for the dependency end_date: Optional[datetime.date] = None # Optional end date for the dependency @field_validator("weeks") def weeks_none_if_dates_set(cls, weeks, values): data = values.data if (data.get("start_date") is not None or data.get("end_date") is not None) and weeks is not None: raise ValueError("If start_date or end_date is set, weeks must not be set") return weeks def generate_not_available_to_works( start_date: datetime.date, end_date: datetime.date | None = None, days: int | None = None, reason: str = "unknown", ) -> List[NotAvailableToWork]: if end_date is None and days is None: raise ValueError("Either end_date or days must be provided") if end_date is not None and days is not None: raise ValueError("Only one of end_date or days must be provided") if end_date is not None: days = (end_date - start_date).days not_available_to_works = [] for day in range(days): not_available_to_works.append(NotAvailableToWork(date=start_date + datetime.timedelta(days=day), reason=reason)) return not_available_to_works class NonWorkingDays(BaseModel): day: days start_date: datetime.date | None = None end_date: datetime.date | None = None @field_validator("day") def day_in(cls, day): for whole_day in whole_days: if day.lower() in whole_day: return whole_day[:3].capitalize() raise ValueError("Invalid day") class WorkRequests(BaseModel): """ Shift is the name of the shift requested, * is a wildcard for any shift """ 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 AvoidShiftOnDates(BaseModel): """Worker-level constraint: avoid listed shifts on given dates or date range.""" shifts: list[str] start_date: Optional[datetime.date] = None end_date: Optional[datetime.date] = None dates: Optional[list[datetime.date]] = None reason: Optional[str] = None @field_validator("start_date", "end_date", mode="before") def coerce_dates(cls, v): # allow strings in common formats if v is None: return None if isinstance(v, datetime.date): return v if isinstance(v, str): for fmt in ("%d/%m/%Y", "%d/%m/%y", "%Y-%m-%d"): try: return datetime.datetime.strptime(v, fmt).date() except Exception: continue raise ValueError(f"Cannot parse date: {v}") class MaxDaysPerWeekBlockConstraint(BaseModel): max_days: int = Field( ..., ge=1, description="Maximum number of worked days allowed inside each sliding week block.", ) week_block: int = Field( ..., ge=1, description="Length of the sliding week block (for example, 2 means every consecutive 2-week window).", ) ignore_dates: list[datetime.date] = Field( default_factory=list, description="List of specific dates to exclude from this constraint.", ) class MaxUniqueShiftsPerWeekBlockConstraint(BaseModel): max_unique_shifts: int = Field( ..., ge=1, description="Maximum number of distinct shift names allowed inside each sliding week block.", ) week_block: int = Field( ..., ge=1, description="Length of the sliding week block (for example, 2 means every consecutive 2-week window).", ) ignore_dates: list[datetime.date] = Field( default_factory=list, description="List of specific dates to exclude from this constraint.", ) class Worker(BaseModel): name: str site: str # Grade are equivalent to roles, by keeping them integer defining model # rules is easier. grade: int = 1 # We can either have a user generated ID id: Optional[int | uuid.UUID | str] = 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] = [] locum_availability: list[WorkRequests] = [] locum_max_shifts: int = 0 locum_max_shifts_per_week: int = 0 locum_on_nwds: bool = False remote_site: str = "plymouth" # We set a default proc_site previous_shifts: dict = {} shift_balance_extra: dict = {} bank_holiday_extra: int = 0 pair: int | str | None = None locum: bool = False allowed_multi_shift_sets: list[set] = [] # or list/set of frozensets prefer_multi_shift_together: int = 0 # Default: no preference # Set to a positive integer to prefer, negative to discourage, 0 for neutral hard_day_dependencies: list[HardDayDependency] = [] hard_day_exclusions: list[HardDayExclusion] = [] force_assign_with: dict[str, list[str]] = {} # e.g. {"oncall": ["weekend"]} max_shifts_per_week_by_shift_name: dict[str, int] = {} # e.g. {"night": 2, "a": 3} forced_assignments: List[tuple[int, str, str]] = [] # (week, day, shift_name) forced_assignments_by_date: List[tuple[datetime.date, str]] = [] # (date, shift_name) max_days_worked_per_week: int | None= None # Default: no limit, can be set to a specific number max_days_per_week_block: list[MaxDaysPerWeekBlockConstraint] = [] max_unique_shifts_per_week_block: list[MaxUniqueShiftsPerWeekBlockConstraint] = [] assign_as_block_preferences: dict[str, float] = {} # shift_name -> weight (positive = prefer block) shift_fte_overrides: dict[str, int] = {} # Need checks to ensure shifts exist exact_shifts: dict[str, int] = {} # Map shift_name to exact number of shifts groups: set[str] = set() # Groups that the worker belongs to weekend_shift_target_number: float = 0.0 avoid_shifts_on_dates: list[AvoidShiftOnDates] = [] model_config = ConfigDict( 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: "RotaBuilder"): self.proportion_rota_to_work: float = 1 self.shift_target_number: dict[str, int] = defaultdict(int) if self.id is None: self.id = uuid.uuid4() if self.start_date is not None and self.end_date is not None: if self.start_date >= self.end_date: Rota.add_warning( "Worker/Early end date", f"{self.name} [{self.id}] end date is before rota start date: {self.end_date}", ) raise ValueError(f"End date must be after start date: {self.name}") # if no start date default to the start of the rota if self.start_date is None: self.calculated_start_date = Rota.start_date else: self.calculated_start_date = self.start_date # ? test if start date is valid if self.start_date < Rota.start_date: # If worker start date is before rota start date use the rota start date self.calculated_start_date = Rota.start_date elif self.start_date >= Rota.rota_end_date: # If it is after add a warning Rota.add_warning( "Worker/Late start date", f"{self.name} [{self.id}] start date is after rota end date: {self.start_date}", ) else: pass if self.end_date is None: self.calculated_end_date = Rota.rota_end_date else: self.calculated_end_date = self.end_date # self.calculated_end_date = datetime.datetime.strptime(end_date, "%d/%m/%y").date() if self.calculated_end_date > Rota.rota_end_date: self.calculated_end_date = Rota.rota_end_date 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.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.calculated_start_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.calculated_end_date}" ) if not self.not_available_to_work: Rota.add_warning( "Worker/No unavailability", f"{self.name} [{self.id}] has no unavailabilities (leave)", ) 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(f"End OOP date must be after start date [{self.name} - Start date: {self.start_date} / End date: {self.end_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: if end_oop_date > Rota.rota_end_date: end_oop_date = Rota.rota_end_date if start_oop_date < Rota.start_date: start_oop_date = Rota.start_date oop_length = (end_oop_date - start_oop_date).days print(f"oop length {oop_length}") days_to_work = days_to_work - oop_length days_until_oop = (start_oop_date - Rota.start_date).days 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) ) days_to_end = (self.calculated_end_date - Rota.start_date).days # loop throught dates converting to week / day combination for preference in self.pref_not_to_work: days_from_start = (preference.date - Rota.start_date).days # Ignore dates past the end of the rota (or end date) if days_from_start < days_to_end: 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)] = preference.reason # print(not_available_to_work) # loop throught dates converting to week / day combination unavailable_set = set() for unavalability in self.not_available_to_work: days_from_start = (unavalability.date - Rota.start_date).days # Ignore dates past the end of the rota (or end date) if days_from_start < days_to_end: 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)] = ( unavalability.reason ) unavailable_set.add((week, day)) for request in self.work_requests: days_from_start = (request.date - Rota.start_date).days week = days_from_start // 7 + 1 day = Rota.days[(days_from_start % 7)] if request.shift == "*": for shift in Rota.get_shifts_for_worker_site(self.site): Rota.work_requests.add((self.id, week, day, shift.name)) else: Rota.work_requests.add((self.id, week, day, request.shift)) for request in self.locum_availability: days_from_start = (request.date - Rota.start_date).days week = days_from_start // 7 + 1 day = Rota.days[(days_from_start % 7)] if request.shift == "*": # We may want to add a limit to shifts that can be filled by locums? for shift in Rota.get_shifts(): Rota.locum_availability.add((self.id, week, day, shift.name)) else: Rota.locum_availability.add((self.id, week, day, request.shift)) # Calculate the proportion of the rota that is being worked self.proportion_rota_to_work = days_to_work / Rota.rota_days_length self.days_to_work = days_to_work # We have to adjust the full time equivalent for people who CCT / leave the rota early self.fte_adj = self.fte * self.proportion_rota_to_work self.fte_adj_shifts = {} if self.shift_fte_overrides: for shift, fte in self.shift_fte_overrides.items(): self.fte_adj_shifts[shift] = fte * self.proportion_rota_to_work if self.fte_adj > 100: # Shouldn't happen ? bug if it does Rota.add_warning( "FTE > 100%", f"{self.name} [{self.id}] FTE = {self.fte_adj}" ) raise ValueError("{} : fte_ajd = {}".format(self.name, self.fte_adj)) # TODO: this has already been validated, consider moving self.non_working_day_list = [] for non_working_day in self.nwds: start_date = Rota.start_date end_date = Rota.rota_end_date if non_working_day.start_date is not None: start_date = non_working_day.start_date if non_working_day.end_date is not None: end_date = non_working_day.end_date self.non_working_day_list.append( (non_working_day.day, start_date, end_date) ) if days_to_work < 1: self.fte_adj = 0 if ( days_to_work > 0 and (days_to_work - len(unavailable_set)) / days_to_work < 0.3 ): console.print(f"{ self.name }: {days_to_work}, {Rota.rota_days_length}") console.print( f"Warning {self.name} has excessive unavailability [{len(unavailable_set)}/{days_to_work}, adjusted fte={self.fte_adj}] (this may be infeasible)", style="red", ) # Register any per-worker avoid-shift-on-dates rules with the rota builder try: for av in getattr(self, "avoid_shifts_on_dates", []): entry = { "names": [self.name], "shifts": av.shifts, "start_date": av.start_date, "end_date": av.end_date, "dates": av.dates, "reason": getattr(av, "reason", None), } # append to rota-level list; the model builder will enforce these if hasattr(Rota.constraint_options_model, "avoid_shifts_by_worker_dates"): Rota.constraint_options_model.avoid_shifts_by_worker_dates.append(entry) except Exception: Rota.add_warning("Worker/avoid_shifts_on_dates", f"Failed to register avoid_shifts_on_dates for {self.name}") def __lt__(self, other) -> bool: return (self.site, self.grade, self.fte_adj, self.name) < ( other.site, other.grade, other.fte_adj, other.name, ) def __str__(self) -> str: return f"{self.name}, {self.site}, {self.oop}" # 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), # nwds, # self.remote_site, # ) def get_details(self) -> str: return "{} {} {}".format(self.id, self.site[0], self.grade) def get_full_details(self) -> str: return "{}/{}/{}".format(self.site, self.grade, self.name) def get_shift_targets(self) -> dict: return self.shift_target_number def get_fte(self, shift: str | None=None) -> int: if shift is not None: if shift in self.fte_adj_shifts: return self.fte_adj_shifts[shift] return self.fte_adj def allow_shifts_together(self, *shifts): if not hasattr(self, "allowed_multi_shift_sets"): self.allowed_multi_shift_sets = [] self.allowed_multi_shift_sets.append(frozenset(shifts)) def add_hard_day_dependency(self, days: set[str], weeks: Optional[List[int]] = None, start_date: Optional[datetime.date] = None, end_date: Optional[datetime.date] = None, ignore_dates: Optional[List[datetime.date]] = None): """ Add a hard day dependency for this worker. days: List of days that must be grouped together. weeks: Optional list of weeks this dependency applies to. If None, applies to all weeks. """ self.hard_day_dependencies.append(HardDayDependency(days=days, weeks=weeks, start_date=start_date, end_date=end_date, ignore_dates=ignore_dates or [])) def add_hard_day_exclusion(self, days: set[str], weeks: Optional[List[int]] = None, start_date: Optional[datetime.date] = None, end_date: Optional[datetime.date] = None): """ Add a hard day exclusion for this worker. days: List of days that must be excluded together. weeks: Optional list of weeks this exclusion applies to. If None, applies to all weeks. """ self.hard_day_exclusions.append(HardDayExclusion(days=days, weeks=weeks, start_date=start_date, end_date=end_date)) def add_force_assign_with(self, shift: str, with_shifts: list[str] | str): if not hasattr(self, "force_assign_with"): self.force_assign_with = {} if isinstance(with_shifts, str): with_shifts = [with_shifts] self.force_assign_with[shift] = with_shifts def set_max_shifts_per_week(self, shift_name: str, max_per_week: int): """Helper to set the maximum number of shifts of a given type per week for this worker.""" if not hasattr(self, "max_shifts_per_week_by_shift_name"): self.max_shifts_per_week_by_shift_name = {} self.max_shifts_per_week_by_shift_name[shift_name] = max_per_week def force_assign_shift(self, week: int, day: str, shift_name: str): """Force this worker to be assigned to a shift on a specific week/day.""" self.forced_assignments.append((week, day, shift_name)) def force_assign_shift_by_date(self, date: datetime.date, shift_name: str): """Force this worker to be assigned to a shift on a specific date.""" self.forced_assignments_by_date.append((date, shift_name)) def add_max_days_per_week_block_constraint(self, max_days: int, week_block: int, ignore_dates: Optional[List[datetime.date]] = None): """Add a constraint to limit the number of days worked in any sliding block of weeks.""" self.max_days_per_week_block.append(MaxDaysPerWeekBlockConstraint(max_days=max_days, week_block=week_block, ignore_dates=ignore_dates or []))