proc-rota/rota/workers.py

import datetime
from collections import defaultdict
from typing import Iterable, List, Literal, Optional
from pydantic import BaseModel, Extra, validator

from rich.pretty import pprint

# from .shifts import RotaBuilder, days, sites
import uuid

# from rota.shifts import RotaBuilder
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from rota.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 NonWorkingDays(BaseModel):
    day: days
    start_date: datetime.date | None = None
    end_date: datetime.date | None = None

    @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):
    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.shift_target_number = 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:
                raise ValueError("End date must be after start date")

        if self.name == "Steph bailey":
            print(self.start_date, self.end_date)

        # 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:
            # ? test if start date is valid
            self.calculated_start_date = self.start_date

        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}"

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

        # 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

        # 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 item in self.work_requests:
            days_from_start = (item.date - Rota.start_date).days
            week = days_from_start // 7 + 1
            day = Rota.days[(days_from_start % 7)]
            if item.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, item.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

        if self.fte_adj > 100:
            # Shouldn't happen ? bug if it does
            raise ValueError("{} : fte_ajd = {}".format(self.name, self.fte_adj))

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

        if days_to_work < 1:
            self.fte_adj = 0

    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:

        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