proc-rota/shifts.py

import datetime
import itertools
from typing import List, Tuple

from workers import Worker

from pyomo.environ import *
from pyomo.opt import SolverFactory

ShiftName = str
DayStr = str
WeekInt = int

days = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]

# weeks = [i for i in range(1,weeks_to_rota+1)]

# weeks_days_product = list(itertools.product(weeks, days))

sites = ("truro", "exeter", "torquay", "barnstaple", "plymouth")


class SingleShift(object):
    """Class to hold all details for a shift"""

    def __init__(
        self,
        sites,
        name,
        length,
        shift_days,
        # balance_by_site=True,
        balance_offset=1,  # this could be generated dynamically
        balance_weighting=1,  # this could be generated dynamically
        workers_required=1,
        rota_on_nwds=False,
    ):
        self.site = sites
        self.name = name
        self.length = length
        self.shift_days = shift_days
        # self.balance_by_site = balance_by_site
        # balance_offset defines the max difference in allocated shifts
        # versus target shifts (hard constraint)
        # if 0 exactly equal shifts must be assigend (unlikely to be possible)
        self.balance_offset = balance_offset

        # weight the shift for balancing, default is 1
        self.balance_weighting = balance_weighting

        self.workers_required = workers_required
        self.rota_on_nwds = rota_on_nwds

        self.total_shifts = length * len(days) * len(shift_days)


class RotaBuilder(object):
    """Class to hold and manipulate shifts"""

    def __init__(self, weeks_to_rota=26):
        self.shifts = []  # type List[SingleShift]

        self.weeks = [i for i in range(1, weeks_to_rota + 1)]

        self.weeks_days_product = list(itertools.product(self.weeks, days))

        self.start_date = datetime.date(2020, 9, 7)
        self.rota_days_length = len(self.weeks) * 7
        self.rota_end_date = self.start_date + datetime.timedelta(self.rota_days_length)

        self.unavailable_to_work = set()

        self.workers = []

        self.night_blocks = ["weekday", "weekend", "none"]

    def build_model(self):
        # Initialize model
        self.model = ConcreteModel()

        # binary variables representing if a worker is scheduled somewhere
        # NOTE: this will assign to all possible shift combinations (which we probably don't want)
        self.model.works = Var(
            (
                (worker.id, week, day, shift)
                for worker in self.workers
                for week, day in self.get_week_day_combinations()
                for shift in self.get_shift_names()
            ),
            within=Binary,
            initialize=0,
        )

        # The nights self.model is used to ensure nights are assigned as a block (and limit the number of workers required)
        self.model.nights = Var(
            (
                (worker.id, week, block)
                for worker in self.workers
                for week in self.weeks
                for block in self.night_blocks
            ),
            within=Binary,
            initialize=0,
        )

        self.model.shift_count = Var(
            (
                (worker.id, shift)
                for worker in self.workers
                for shift in self.get_shift_names()
            ),
            domain=NonNegativeReals,
            initialize=0,
        )

        # self.model.shift_count_t1 = Var(((worker.id, shift) for worker in self.workers for shift in self.get_shift_names()), domain=NonNegativeReals, initialize=0)
        # self.model.shift_count_t2 = Var(((worker.id, shift) for worker in self.workers for shift in self.get_shift_names()), domain=NonNegativeReals, initialize=0)

        self.model.shift_count_t1 = Var(
            ((worker.id) for worker in self.workers),
            domain=NonNegativeReals,
            initialize=0,
        )
        self.model.shift_count_t2 = Var(
            ((worker.id) for worker in self.workers),
            domain=NonNegativeReals,
            initialize=0,
        )

        self.model.shift_count_t3 = Var(
            (worker.id for worker in self.workers), initialize=0
        )
        # self.model.unavailable = Var(((worker, week, day) for worker in self.workers for week in weeks for day in days),
        #                  within=Binary, initialize=0)

        def availability_init(model, wid, week, day):
            if (wid, week, day) in self.unavailable_to_work:
                # print((wid, week, day))
                return 0
            return 1

        self.model.available = Param(
            (
                (worker.id, week, day)
                for worker in self.workers
                for week, day in self.get_week_day_combinations()
            ),
            initialize=availability_init,
        )

        # binary variables representing if a worker is necessary
        # self.model.needed = Var([worker.id for worker in self.workers], within=Binary, initialize=0)

        # binary variables representing if a worker worked on sunday but not on saturday (avoid if possible)
        # self.model.no_pref = Var([worker.id for worker in self.workers], within=Binary, initialize=0)

        self.build_model_constraints()

    def build_model_constraints(self):
        self.model.constraints = ConstraintList()  # Create a set of constraints

        for (
            week,
            day,
            shift,
            self.workers_required,
            site_required,
        ) in self.get_required_workers_and_site_combinations():
            # print(week, day, shift, self.workers_required, site_required)
            self.model.constraints.add(
                self.workers_required
                == sum(
                    self.model.works[worker.id, week, day, shift]
                    for worker in self.workers
                    if worker.site in site_required
                )
            )

        for week, day, shift in self.get_not_required_shifts():
            self.model.constraints.add(
                0
                == sum(
                    self.model.works[worker.id, week, day, shift.name]
                    for worker in self.workers
                )
            )

        # Constraint: total hours worked hours worked
        # for worker in self.workers:
        #    self.model.constraints.add(
        #        1200 >= sum(shift_lengths[shift] * self.model.works[worker, week, day, shift] for week in weeks for day in days for shift in days_shifts[day])
        #    )

        # for week in weeks:
        #    for worker in self.workers:
        #        self.model.constraints.add(
        #            48 >= sum(shift_lengths[shift] * self.model.works[worker, week, day, shift] for day in days for shift in days_shifts[day])
        #        )

        def maxHoursPerWeekRule(model, wid, week):
            return 72 >= sum(
                shift.length * model.works[wid, week, day, shift.name]
                for day, shift in self.get_day_shiftclass_products()
            )

        self.model.max_hours_per_week_constraint = Constraint(
            [worker.id for worker in self.workers], self.weeks, rule=maxHoursPerWeekRule
        )

        def maxHoursRule(model, wid):
            return 8000 >= sum(
                shift.length * model.works[wid, week, day, shift.name]
                for week, day, shift in self.get_all_shiftclass_combinations()
            )

        self.model.max_hours_constraint = Constraint(
            [worker.id for worker in self.workers], rule=maxHoursRule
        )

        # self should not be allocated on nwds (excepting night shifts)
        def nwdRule(model, worker):
            print(worker, worker.id, worker.nwd)
            wid = worker.id
            nwd = worker.nwd
            return 0 == sum(
                model.works[wid, week, day, shift.name]
                for week, day, shift in self.get_all_shiftclass_combinations()
                if (not shift.rota_on_nwds and day in nwd)
            )

        # self.model.nwd_constraint = Constraint([(worker) for worker in self.workers if worker.nwd], rule=nwdRule)

        # Set a rule to even twilight shifts:
        # total_twilight_shifts = len(weeks) * 5 * 1 # One person per site

        # def trurotwilightShiftNumberRule(self.model, wid, fte_adj, site, shift):
        #    #site, shift = site_and_shift
        #    max_shifts = total_twilight_shifts / self.full_time_equivalent_sites[site] * fte_adj + twilight_balance_offset
        #    min_shifts = total_twilight_shifts / self.full_time_equivalent_sites[site] * fte_adj - twilight_balance_offset

        #    return inequality(min_shifts, sum(self.model.works[wid, week, day, shift.GetShift()] for week in weeks for day in days), max_shifts)
        # self.model.twilight_shifts_truro_constraint = Constraint([(worker.id, worker.fte_adj, ss[0], ss[1]) for ss in site_and_shifts for worker in self.workers if worker.site == ss[0]], rule=trurotwilightShiftNumberRule)

        # def plymouthtwilightShiftNumberRule(self.model, wid, fte_adj):
        #    max_shifts = total_twilight_shifts / self.full_time_equivalent_sites["plymouth"] * fte_adj + twilight_balance_offset
        #    min_shifts = total_twilight_shifts / self.full_time_equivalent_sites["plymouth"] * fte_adj - twilight_balance_offset
        #
        #    return inequality(min_shifts, sum(self.model.works[wid, week, day, "plymouth_twilight"] for week in weeks for day in days), max_shifts)
        # self.model.twilight_shifts_plymouth_constraint = Constraint([(worker.id, worker.fte_adj) for worker in self.workers if worker.site == "plymouth"], #rule=plymouthtwilightShiftNumberRule)

        # Limit to 1 ST2 (or below) on a night shift
        def nightShiftMaxSTRule(model, week, block):
            """Limits to 1 ST2 (or below) per night"""
            single_workers = [w for w in self.workers if w.grade < 3]
            if not single_workers:
                return Constraint.Skip
            return sum(model.nights[w.id, week, block] for w in single_workers) <= 1

        self.model.night_shifts_max_st_constraint = Constraint(
            [week for week in self.weeks],
            [block for block in ["weekday", "weekend"]],
            rule=nightShiftMaxSTRule,
        )

        # Enusre at least 1 ST4+ on a night shift
        def nightShiftMinST4Rule(model, week, block):
            single_workers = [w for w in self.workers if w.grade > 3]
            if not single_workers:
                return Constraint.Skip
            return sum(model.nights[w.id, week, block] for w in single_workers) >= 1

        self.model.night_shifts_min_st4_constraint = Constraint(
            [week for week in self.weeks],
            [block for block in ["weekday", "weekend"]],
            rule=nightShiftMinST4Rule,
        )

        # Constraint (def of self.model.needed)
        # for worker in self.workers:
        #    self.model.constraints.add(
        #        10000 * self.model.needed[worker.id] >= sum(self.model.works[worker.id, week, day, shift] for week in weeks for day in days for shift in days_shifts[day])
        #    )  # if any self.model.works[worker, ·, ·] non-zero, self.model.needed[worker] must be one; else is zero to reduce the obj function
        # 10000 is to remark, but 5 was enough since max of 40 hours yields max of 5 shifts, the maximum possible sum

        # Constraint (def of self.model.no_pref)
        # for worker in self.workers:
        #    for week in weeks:
        #        self.model.constraints.add(
        #            self.model.no_pref[worker.id] >= sum(self.model.works[worker.id, week, 'Sat', shift] for shift in days_shifts['Sat'])
        #            - sum(self.model.works[worker.id, week, 'Sun', shift] for shift in days_shifts['Sun'])
        #        )  # if not working on sunday but working saturday self.model.needed must be 1; else will be zero to reduce the obj function

        # def balanceSiteShiftNumberRule(model, wid, fte_adj, shiftClass):
        #     # TODO: make helper to retrieve
        #     total_shifts = len(self.weeks) * len(shiftClass.shift_days) * shiftClass.self.workers_required # One person per site

        #     full_time_equivalent_joined = sum([self.full_time_equivalent_sites[i] for i in shiftClass.site])

        #     max_shifts = total_shifts / full_time_equivalent_joined * fte_adj + twilight_balance_offset
        #     min_shifts = total_shifts / full_time_equivalent_joined * fte_adj - twilight_balance_offset

        #     return inequality(min_shifts, sum(self.model.works[wid, week, day, shiftClass.name] for week, day in self.get_week_day_combinations()), max_shifts)

        # #self.model.balance_site_shift_constraint = Constraint([(worker.id, worker.fte_adj, shift) for shift in self.GetShiftOptions() for worker in self.workers if worker.site in shift.site and shift.balance_by_site], rule=balanceSiteShiftNumberRule)

        # # set a rule to even night shifts (could also be done in blocks, which may be quicker?)
        # # takes into accound both LTFT, OOP and leaving the training scheme
        # #total_night_shifts = self.rota_days_length * self.GetShiftByName("night").self.workers_required

        # def nightShiftNumberRule(model, wid, fte_adj): # NOTE no longer used - see generic function below
        #     target_shifts = total_night_shifts / self.full_time_equivalent * fte_adj
        #     max_shifts = target_shifts + night_balance_offset
        #     min_shifts = target_shifts - night_balance_offset

        #     return inequality(min_shifts, sum(model.works[wid, week, day, "night"] for week, day in self.get_week_day_combinations()), max_shifts)
        # #self.model.night_shifts_constraint = Constraint([(worker.id, worker.fte_adj) for worker in self.workers], rule=nightShiftNumberRule)
        # twilight_constraints = {}

        # Balance shifts
        for worker in self.workers:
            for shift in self.GetShiftOptions():
                if (
                    worker.site in shift.site
                ):  # Each site specfies which sites self.workers can fullfill it

                    # total_shifts = self.rota_days_length * shift.workers_required # One person per site
                    total_shifts = (
                        len(self.weeks)
                        * len(shift.shift_days)
                        * shift.workers_required
                    )

                    # if shift.balance_by_site:
                    full_time_equivalent_joined = sum(
                        [self.full_time_equivalent_sites[i] for i in shift.site]
                    )

                    target_shifts = (
                        total_shifts / full_time_equivalent_joined * worker.fte_adj
                    )

                    worker.shift_target_number[shift.name] = target_shifts

                    max_shifts = target_shifts + shift.balance_offset
                    min_shifts = target_shifts - shift.balance_offset
                    # print(max_shifts, min_shifts)

                    self.model.constraints.add(
                        inequality(
                            min_shifts,
                            sum(
                                self.model.works[worker.id, week, day, shift.name]
                                for week, day in self.get_week_day_combinations()
                            ),
                            max_shifts,
                        )
                    )
                    # return inequality(min_shifts, sum(self.model.works[wid, week, day, shiftClass.name] for week, day in self.get_week_day_combinations()), max_shifts)

        for worker in self.workers:
            # Define shift_count_t1 and shift_count_t2 constraints for the object
            # This bypassing the need for a quadratic solver
            # t1-t2 is the target
            # As the objective is to minimise t1+t2 and t1 and t2 are positive reals
            # t1+t2 approximates the absolute target (which otherwise requires a quadratic solver)
            self.model.constraints.add(
                # self.model.shift_count_t3[worker.id] == abs(sum(self.model.shift_count[worker.id, shift.name] - worker.shift_target_number[shift.name] for shift in self.GetShiftOptions()))
                self.model.shift_count_t1[worker.id]
                - self.model.shift_count_t2[worker.id]
                == sum(
                    (
                        self.model.shift_count[worker.id, shift.name]
                        - worker.shift_target_number[shift.name]
                    )
                    * shift.balance_weighting
                    for shift in self.GetShiftOptions()
                )
            )
            for shift in self.GetShiftOptions():
                self.model.constraints.add(
                    self.model.shift_count[worker.id, shift.name]
                    == sum(
                        self.model.works[worker.id, week, day, shift.name]
                        for week, day in self.get_week_day_combinations()
                    )
                )
                # self.model.constraints.add(
                #    #self.model.shift_count_t1[worker.id, shift.name]-self.model.shift_count_t2[worker.id, shift.name] == self.model.shift_count[worker.id, shift.name]
                # self.model.shift_count_t1[worker.id, shift.name]-self.model.shift_count_t2[worker.id, shift.name] == self.model.shift_count[worker.id, shift.name] - worker.shift_target_number[shift.name]
                # )

            if worker.nwd:
                for week, day, shift in self.get_all_shiftclass_combinations():
                    if not shift.rota_on_nwds and day in worker.nwd:
                        self.model.constraints.add(
                            0 == self.model.works[worker.id, week, day, shift.name]
                        )

        # Constraint: rest between two shifts is of 12 hours (i.e., at least two shifts)
        weeks_days = self.get_week_day_combinations()
        for worker in self.workers:
            if "nights" in self.get_shift_names():

                for week_blocks in self.get_week_block_iterator(2):
                    # Prevent nights more than once every n weeks
                    self.model.constraints.add(
                        1
                        >= sum(
                            self.model.nights[worker.id, week, block]
                            for week in week_blocks
                            for block in ["weekday", "weekend"]
                        )
                    )

            for week in self.weeks:
                if "nights" in self.get_shift_names():
                    # Force nights to be assigned in blocks
                    self.model.constraints.add(
                        1
                        == sum(
                            self.model.nights[worker.id, week, block]
                            for block in self.night_blocks
                        )
                    )

                    # if night block is weekday make sure Mon - Thurs is assigned as nights
                    self.model.constraints.add(
                        self.model.nights[worker.id, week, "weekday"] * 4
                        == sum(
                            self.model.works[worker.id, week, day, "night"]
                            for day in days[:4]
                        )
                    )
                    self.model.constraints.add(
                        self.model.nights[worker.id, week, "weekend"] * 3
                        == sum(
                            self.model.works[worker.id, week, day, "night"]
                            for day in days[4:]
                        )
                    )

                #

            for n in range(len(weeks_days)):

                week, day = weeks_days[n]

                p1 = 1
                if n > 0:
                    pweek, pday = weeks_days[n - 1]
                else:
                    p1 = 0
                    pweek, pday = weeks_days[n]

                n1 = 1
                n2 = 1
                try:
                    nweek, nday = weeks_days[n + 1]
                except IndexError:
                    # print(worker.id, "shit")
                    nweek, nday = weeks_days[n]
                    n1 = 0

                try:

                    n2week, n2day = weeks_days[n + 2]
                except IndexError:
                    # print(worker.id, "shit2")
                    n2week, n2day = weeks_days[n]
                    n2 = 0

                # print(week, day, nweek, nday, n2week, n2day, n1, n2)
                # Unable to work (hard constraint not preference)
                self.model.constraints.add(
                    self.model.available[worker.id, week, day]
                    >= sum(
                        self.model.works[worker.id, week, day, shift]
                        for shift in self.get_shift_namesByDay(day)
                    )
                )

                if "nights" in self.get_shift_names():
                    self.model.constraints.add(
                        self.model.available[worker.id, week, day]
                        >= self.model.works[worker.id, pweek, pday, "night"]
                    )

                # single shift per day
                self.model.constraints.add(
                    1
                    >= sum(
                        self.model.works[worker.id, week, day, shift]
                        for shift in self.get_shift_namesByDay(day)
                    )
                )
                ## if working in evening, until next evening (note that after sunday comes next monday)
                # self.model.constraints.add(
                #    1 >= sum(self.model.works[worker.id, week, days[j], shift] for shift in ['evening', 'night']) +
                #    self.model.works[worker.id, nweek, days[(j + 1) % 7], 'truro_twilight']
                # )

                # if working a night ensure preceeding (1) or subsequent (2) shifts can only be nights

                if "nights" in self.get_shift_names():

                    self.model.constraints.add(
                        1
                        >= self.model.works[worker.id, week, day, "night"]
                        + sum(
                            n1 * self.model.works[worker.id, nweek, nday, shift]
                            for shift in self.get_shift_namesByDay(nday)
                            if shift != "night"
                        )
                        + sum(
                            n2 * self.model.works[worker.id, n2week, n2day, shift]
                            for shift in self.get_shift_namesByDay(n2day)
                            if shift != "night"
                        )
                        + sum(
                            p1 * self.model.works[worker.id, pweek, pday, shift]
                            for shift in self.get_shift_namesByDay(pday)
                            if shift != "night"
                        )
                    )

        self.define_objectives()

    def define_objectives(self):
        # Define an objective function with model as input, to pass later
        def obj_rule(m):
            
            #c = len(workers)

            #return 1
            return sum(self.model.shift_count_t1[(worker.id)] + self.model.shift_count_t2[(worker.id)] for worker in self.workers)
            #return sum(model.shift_count_t3[worker.id] for worker in workers)
            


        # add objective function to the model. rule (pass function) or expr (pass expression directly)
        self.model.obj = Objective(rule=obj_rule, sense=minimize)

    def add_worker(self, worker: Worker):
        """Add a worker to the rota

        Args:
            worker (Worker):
        """
        self.workers.append(worker)

    def add_workers(self, workers: List):
        """Add multiple worker to the rota

        Args:
            workers (List(Worker)):
        """
        self.workers.extend(workers)

    def build_workers(self):
        """Process loaded workers

        Must be called prior to attempting to solve
        """
        self.workers = sorted(self.workers)

        self.full_time_equivalent = sum([w.fte_adj for w in self.workers])
        self.full_time_equivalent_sites = {}

        for site in sites:
            self.full_time_equivalent_sites[site] = sum(
                [w.fte_adj for w in self.workers if w.site == site]
            )

    def add_shift(self, shift):
        """Add a shift to the collection

        :param SingleShift shift: Shift object

        """
        self.shifts.append(shift)
        self.BuildShift()

    def add_shifts(self, *shifts: SingleShift):
        """Add multiple shifts
        
        Returns:
            None
        """
        self.shifts.extend(shifts)
        self.BuildShift()

    def get_shift_namesByDay(self, day: DayStr):
        """Returns the shifts required for a specific day

        Returns:
            set: Set of ShiftName

        """
        return self.day_shifts_dict[day]

    def GetShiftByName(self, name):
        """

        :param name: 

        """
        return self.shifts_by_name[name]

    def GetShiftLengthByName(self, name):
        """Get the length of a shift
        
        Args:
            name (str): Name of a shift
        
        Returns:
            int: Length of the shift
        """
        return self.shifts_by_name[name].length

    def BuildShift(self):
        """ """
        self.shifts_by_name = {}
        self.shift_names = []  # type: List[ShiftName]

        self.day_shifts_dict = {day: set() for day in days}

        for s in self.shifts:
            self.shifts_by_name[s.name] = s
            self.shift_names.append(s.name)

            for day in s.shift_days:
                self.day_shifts_dict[day].add(s.name)

        self.week_day_shift_product = []
        self.week_day_shiftclass_product = []
        for week, day in self.weeks_days_product:
            for s in self.shifts:
                if day in s.shift_days:
                    self.week_day_shift_product.append((week, day, s.name))
                    self.week_day_shiftclass_product.append((week, day, s))

        self.day_shift_product = []
        self.day_shiftclass_product = []
        for s in self.shifts:
            for day in days:
                if day in s.shift_days:
                    self.day_shift_product.append((day, s.name))
                    self.day_shiftclass_product.append((day, s))

    def GetAllDayShiftAsNames(self):
        """Returns a list of all required day / shift combinations
        
        Returns:
            list: list of in the following format (str->day, str->shift)
        """
        return self.day_shift_product

    def get_day_shiftclass_products(self) -> List[Tuple[DayStr, SingleShift]]:
        """Returns a list of all required day / shift combinations
        
        Returns:
            list: list of in the following format (str->day, shiftObject->shift)
        """
        return self.day_shiftclass_product

    def GetAllShiftAsNames(self) -> List[Tuple[WeekInt, DayStr, ShiftName]]:
        """Returns a list of all possible week / day / shift combinations
        
        Returns:
            list: contains tuple of all possible week / day / shift combinations
        """
        return self.week_day_shift_product

    def get_all_shiftclass_combinations(self) -> List[Tuple[WeekInt, DayStr, SingleShift]]:
        """Returns a list of all possible week / day / shift combinations
        
        Returns:
            list: contains tuple of all possible week / day / shift combinations
        """
        return self.week_day_shiftclass_product

    def get_week_day_combinations(self) -> list:
        """Returns a list of all week / day tuple combinations

        Returns:
            list: list of possible week day combinations
                    [(1, "Mon"), (1, "Tue"), ... (n, "Fri")]
        """
        return self.weeks_days_product

    def GetShiftOptions(self) -> List[SingleShift]:
        """Returns a list of all the registered shifts
        
        Returns:
            List[SingleShift]: list of registered shifts (as SingleShift)
        """
        return self.shifts

    def get_shift_names(self) -> List[ShiftName]:
        """Returns a list of all the registered shift names
        
        Returns:
            List[ShiftName]: List of names of all available shifts
        """ """ """
        return self.shift_names

    def get_week_block_iterator(self, block_length):
        """Gets a two dimensional list of week blocks in specified length

        e.g. block_length = 4
            [
                [ 1, 2, 3, 4],
                [ 2, 3, 4, 5],
                ...
            ]
        
        Args:
            block_length (int): length of blocks to create
        
        Returns:
            list: two dimensional list containing weeks in blocks
        """
        blocks = []
        for i in range(len(self.weeks)):
            blocks.append(self.weeks[i : i + block_length])
        return blocks

    def get_required_workers_and_site_combinations(self):
        """Returns a list of all possible shifts, the workers required and site
        
        Returns:
            list: list (week, day, shift name, workers required, shift site)
        """
        l = []
        for week, day, shift in self.week_day_shiftclass_product:
            if day in shift.shift_days:
                l.append(
                    (week, day, shift.name, shift.workers_required, shift.site)
                )

        return l

    def get_not_required_shifts(self):
        """Returns a set of all possible shifts combinations that are not required

        Includes those on days
        
        Returns:
            list: (week, day, shift)
        """
        l = set()

        for week in self.weeks:
            for day in days:
                for shift in self.shifts:
                    l.add((week, day, shift))

        return l - set(self.get_all_shiftclass_combinations())