proc-rota/rota_generator/workers.py

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).",
    )


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).",
    )


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

    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))