proc-rota/shifts.py

import datetime
import itertools
from typing import List, Tuple

import datetime

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

from workers import Worker

import json

import csv

from collections import defaultdict

from io import StringIO
import sys

ShiftName = str
DayStr = str
WeekInt = int

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

full_weekend_count = 2

from govuk_bank_holidays.bank_holidays import BankHolidays

bank_holidays = BankHolidays()
bank_holiday_map = {}
for bank_holiday in bank_holidays.get_holidays():
    bank_holiday_map[bank_holiday['date']] = bank_holiday['title']


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=2,  # this could be generated dynamically
        balance_weighting=1,
        workers_required=1,
        rota_on_nwds=False,
        assign_as_block=False,
        force_as_block=False,
        force_as_block_unless_nwd=False,
        hard_constrain_shift=True,
        bank_holidays_only=False,
        constraints=[],
    ):
        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.hard_constrain_shift = hard_constrain_shift

        self.workers_required = workers_required
        self.rota_on_nwds = rota_on_nwds
        self.assign_as_block = assign_as_block
        self.force_as_block = force_as_block
        self.force_as_block_unless_nwd = force_as_block_unless_nwd
        self.constraints = constraints
        
        self.bank_holidays_only = bank_holidays_only

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

    def __str__(self):
        return f"name: {self.name}, site: {self.site}, length: {self.length}, balance weighting: {self.balance_weighting}, balance offset: {self.balance_offset}, hard_constrain: {self.hard_constrain_shift}, workers required : {self.workers_required}, rota on nwds {self.rota_on_nwds}, assign as block: {self.assign_as_block}, force as block: {self.force_as_block}, constraints: {self.constraints}"#, total shifts: {self.total_shifts}"

    def get_shift_summary(self):
        """
        Returns a text summary of the shift
        """

        return "{}: {} workers for sites ({}) on days ({})".format(
            self.name, self.workers_required, ", ".join(self.site),
            ", ".join(self.shift_days))


class RotaBuilder(object):
    """Class to hold and manipulate shifts"""
    def __init__(self,
                 start_date,
                 weeks_to_rota=26,
                 balance_offset_modifier=1,
                 ltft_balance_offset=2,
                 max_night_frequency=2,
                 max_weekend_frequency=2,
                 use_previous_shifts=False,
                 ):

        print("Start time: {}".format(datetime.datetime.now()))
        print("Weeks to rota: {}".format(weeks_to_rota))
        print("Use previous shifts {}".format(use_previous_shifts))

        self.shifts = []  # type List[SingleShift]

        self.days = days

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

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

        self.use_previous_shifts = use_previous_shifts

        if start_date.weekday() != 0:
            raise ValueError("Start date {} must be a Mon (not a {})".format(
                start_date.isoformat(), days[start_date.weekday()]))

        self.start_date = start_date
        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.unavailable_to_work_reason = {}
        self.pref_not_to_work = {}
        self.work_requests = set()
        self.work_requests_map = {}

        self.workers = []
        self.worker_pairs = []

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

        self.sites = None

        self.balance_offset_modifier = balance_offset_modifier
        self.ltft_balance_offset = ltft_balance_offset

        # Don't assign multiple shifts every (n) weeks
        self.max_night_frequency = max_night_frequency
        self.max_weekend_frequency = max_weekend_frequency

        self.constraint_options = {
            "limit_to_1_st2_on_nights": True,
            "ensure_1_st4_plus_on_nights": True,
            "ensure_derriford_reg_for_nights": True,
            "balance_nights": True,
            "constrain_time_off_after_nights": True,
            "balance_nights_across_sites": True,
            "balance_bank_holidays": True,
            "balance_blocks": True,
            "balance_shifts": True,
            "balance_weekends": True,
            "max_weekends": False,
            "max_shifts_per_week": 4,
            "max_shifts_per_month": 12,
            "prevent_monday_after_full_weekends": [],
            "prevent_monday_and_tuesday_after_full_weekends": [],
            "prevent_fridays_before_full_weekends": [],
            "prevent_thursdays_before_full_weekends": [],
            "avoid_st2_first_month": False,
            "hard_constrain_pair_separation": False,
        }

        # Generate a map for week, day combinations to a given date
        self.week_day_date_map = {}

        n = 0
        for week, day in self.get_week_day_combinations():
            d = self.start_date + datetime.timedelta(n)
            self.week_day_date_map[(week, day)] = d
            n = n + 1

    def solve_model(self, solver='cbc', use_neos=False, options={}):
        print("Setting up solver")
        self.opt = SolverFactory(solver)   

        print("Solving")
        if use_neos:
            solver_manager = SolverManagerFactory('neos')  # Solve in neos server
            #results = solver_manager.solve(Rota.model, opt=opt, logfile="test.log")
            results = solver_manager.solve(self.model, keepfiles=True, tee=True, opt=self.opt, logfile="test.log")
        else:
            results = self.opt.solve(
                self.model,
                tee=True,
                options=options,
                #options={
                #    "threads": 10,
                #},
                logfile="test.log"
            )

        self.results = results

        print(results)
        results.solver.status
        if not results.solver.status:
            sys.exit(0)


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

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

        self.model.worker_block_shift_split = Var(
            ((week, shift) for week in self.weeks
             for shift in self.shifts_to_assign_or_force_as_blocks()),
            within=Binary,
            initialize=0,
        )

        self.model.blocks_worker_shift_assigned = Var(
            ((worker.id, week, shift) for worker in self.workers
             for week in self.weeks
             for shift in self.shifts_to_assign_or_force_as_blocks()),
            within=Binary,
            initialize=0,
        )
        

        self.model.blocks_assigned = Var(
            ((week, shift) for week in self.weeks
             for shift in self.shifts_to_assign_or_force_as_blocks()),
            within=NonNegativeIntegers,
            initialize=0,
        )

        if self.constraint_options["balance_bank_holidays"]:
            # Bank holidays
            self.model.bank_holiday_count = Var(
                ((worker.id) for worker in self.workers),
                within=NonNegativeIntegers,
                initialize=0,
            )
            self.model.bank_holiday_count_w = Var(
                ((worker.id) for worker in self.workers),
                within=NonNegativeIntegers,
                initialize=0,
            )
            # self.model.bank_holiday_count = Var(
            #     ((worker.id, week, day) for worker in self.workers
            #     for week, day in self.get_week_day_combinations() if self.week_day_date_map[(week, day)] in bank_holiday_map
            #     ),
            #     within=Binary,
            #     initialize=0,
            # )

        # Used to limit number of workers on night shift per site
        # if we force a binary it will in effect hard constrain to < 2
        # from a single site
        if self.constraint_options["balance_nights_across_sites"]:
            self.model.night_per_site = Var(
                ((week, shift.name, site) for site in self.sites
                 for week in self.weeks
                 for shift in self.get_shifts_with_constraint("night")),
                #within=Binary,
                within=NonNegativeIntegers,
                initialize=0,
            )

            self.model.night_per_site_t1 = Var(
                ((week, shift.name, site) for site in self.sites
                 for week in self.weeks
                 for shift in self.get_shifts_with_constraint("night")),
                domain=NonNegativeIntegers,
                initialize=0,
            )
            self.model.night_per_site_t2 = Var(
                ((week, shift.name, site) for site in self.sites
                 for week in self.weeks
                 for shift in self.get_shifts_with_constraint("night")),
                domain=NonNegativeIntegers,
                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,
        )

        if self.constraint_options["balance_shifts"]:

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

        if self.constraint_options["balance_nights"]:
            # We also try to even out the night shifts seperately
            self.model.night_shift_count = Var(
                ((worker.id) for worker in self.workers),
                domain=NonNegativeReals,
                initialize=0,
            )

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

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

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

        # self.model.worker_weekend_assigned = Var(
        #     ((worker.id, week) for worker in self.workers
        #      for week in self.weeks),
        #     domain=NonNegativeReals,
        #     initialize=0,
        # )

        # self.model.works_weekend_count = Var(
        #     ((worker.id, week) for worker in self.workers
        #      for week in self.weeks),
        #     domain=NonNegativeIntegers,
        #     initialize=0,
        # )
        # self.model.works_saturday = Var(
        #     ((worker.id, week) for worker in self.workers
        #      for week in self.weeks),
        #     domain=Binary,
        #     initialize=0,
        # )
        # self.model.works_sunday = Var(
        #     ((worker.id, week) for worker in self.workers
        #      for week in self.weeks),
        #     domain=Binary,
        #     initialize=0,
        # )

        # self.model.works_whole_weekend = Var(
        #     ((worker.id, week) for worker in self.workers
        #      for week in self.weeks),
        #     domain=Binary,
        #     initialize=0,
        # )

        if self.constraint_options["balance_weekends"]:

            self.model.works_weekend = Var(
                ((worker.id, week) for worker in self.workers
                 for week in self.weeks),
                domain=Binary,
                initialize=0,
            )

            self.model.worker_weekend_count = Var(
                ((worker.id) for worker in self.workers),
                domain=NonNegativeIntegers,
                initialize=0,
            )

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

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

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

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

        # self.model.weekend_count_t2 = Var(
        #     ((worker.id) for worker in self.workers),
        #     domain=NonNegativeReals,
        #     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,
        )

        def work_request_init(model, wid, week, day, shift):
            if (wid, week, day, shift) in self.work_requests:
                self.work_requests_map[(wid, week, day)] = shift
                return 1
            return 0

        self.model.work_requests = Param(
            ((worker.id, week, day, shift) for worker in self.workers
             for week, day, shift in self.get_all_shiftname_combinations()),
            initialize=work_request_init,
        )

        def pref_init(model, wid, week, day):
            if (wid, week, day) in self.pref_not_to_work:
                return self.pref_not_to_work[(wid, week, day)]
            return 0

        # It would also be possible to assign sequential days as blocks
        # which would allow more weight to be given if the whole block
        # is not allocated
        self.model.pref_not_to_work = Param(
            ((worker.id, week, day) for worker in self.workers
             for week, day in self.get_week_day_combinations()),
            initialize=pref_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

        # Ensure each shift has the requisit number of workers assigned
        for (
                week,
                day,
                shift,
                workers_required,
                site_required,
        ) in self.get_required_workers_and_site_combinations():
            print(week, day, shift, workers_required, site_required)
            self.model.constraints.add(workers_required == sum(
                self.model.works[worker.id, week, day, shift]
                for worker in self.workers if worker.site in site_required))
            # And it is not assigned if the worker is from the wrong site
            self.model.constraints.add(0 == sum(
                self.model.works[worker.id, week, day, shift]
                for worker in self.workers if worker.site not in site_required))
            # # Ensure shifts are only assigned by workers from the allowed sites
            # self.model.constraints.add(0 == sum(
            #     self.model.works[worker.id, week, day, shift]
            #     for worker in self.workers if worker.site not in site_required))

        # Ensure no workers are assigned to shifts that are not 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])
        #        )

        # Define a maximum number hours that can be worked per week
        # NOTE: this works on a calander week not a 7 day period
        
        # TODO: consider fixing
        #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_combinations())

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

        # Define an upper limit for hours worked over the rota period
        #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)

        # Limit to 1 ST2 (or below) on a night shift
        def nightShiftMaxSTRule(model, week, shift):
            """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.shift_week_worker_assigned[shift, week, w.id]
                       for w in single_workers) <= 1

        if self.constraint_options["limit_to_1_st2_on_nights"]:
            self.model.night_shifts_max_st_constraint = Constraint(
                [week for week in self.weeks],
                [
                    shift.name
                    for shift in self.get_shifts_with_constraint("night")
                ],
                rule=nightShiftMaxSTRule,
            )

        def nightShiftMinST4Rule(model, week, shift):
            single_workers = [w for w in self.workers if w.grade >= 4]
            if not single_workers:
                return Constraint.Skip
            return sum(model.shift_week_worker_assigned[shift, week, w.id]
                       for w in single_workers) >= 1

        if self.constraint_options["ensure_1_st4_plus_on_nights"]:
            self.model.night_shifts_min_st4_constraint = Constraint(
                [week for week in self.weeks],
                [
                    shift.name
                    for shift in self.get_shifts_with_constraint("night")
                ],
                rule=nightShiftMinST4Rule,
            )

        def nightShiftDerrifordRule(model, week, shift):
            derriford_workers = [w for w in self.workers if w.night_at_derriford >= 1]
            #if not derriford_workers:
                #return Constraint.Skip
            return sum(model.shift_week_worker_assigned[shift, week, w.id]
                       for w in derriford_workers) >= 1

        if self.constraint_options["ensure_derriford_reg_for_nights"]:
            self.model.night_shifts_derriford_rule = Constraint(
                [week for week in self.weeks],
                [
                    shift.name
                    for shift in self.get_shifts_with_constraint("night")
                ],
                rule=nightShiftDerrifordRule,
            )

        # # Count the number of workers from each site on each night shift
        # # As 1 or 0 is optimum we can simply subtract 1 from the number
        if self.constraint_options["balance_nights_across_sites"]:
            for site in self.sites:
                for shift in self.get_shifts_with_constraint("night"):
                    block = shift.name
                    for week in self.weeks:
                        self.model.constraints.add(
                            self.model.night_per_site[week, block, site] ==
                            sum(self.model.shift_week_worker_assigned[
                                block, week, worker.id]
                                for worker in self.workers
                                if worker.site == site)
                            #1 >= sum(self.model.shift_week_worker_assigned[shift.name, week, worker.id] for worker in self.workers if worker.site == site)
                        )
                        self.model.constraints.add(
                            self.model.night_per_site_t1[week, block, site] -
                            self.model.night_per_site_t2[week, block, site] ==
                            self.model.night_per_site[week, block, site] - 1)

        weeks_days = self.get_week_day_combinations()

        for week in self.weeks:
            for shift_name in self.shifts_to_assign_or_force_as_blocks():
                shift = self.get_shift_by_name(shift_name)
                for worker in self.get_workers_for_shift(shift):
                    self.model.constraints.add(
                        8 * self.model.blocks_worker_shift_assigned[
                            worker.id, week, shift_name] >= sum(
                                self.model.works[worker.id, week, day,
                                                 shift_name]
                                for day in shift.shift_days))
                self.model.constraints.add(
                    self.model.blocks_assigned[week, shift_name] == sum(
                        self.model.blocks_worker_shift_assigned[worker.id,
                                                                week,
                                                                shift_name]
                        for worker in self.workers))

                if shift.force_as_block:

                    
                    if shift.force_as_block_unless_nwd:
                        workers = self.get_workers_for_shift(shift)
                        # Get workers who have a nwd on the shift
                        nwd_workers = []
                        full_workers = []
                        
                        for w in workers:
                            if w.nwd is not None:
                                l = []
                                # This should take into account dates!
                                for nwd, d1, d2 in w.nwd:
                                    if nwd in shift.shift_days:
                                        l.append(w)

                                if l:
                                    nwd_workers.append(w)
                                else:
                                    full_workers.append(w)
                                    
                            else:
                                full_workers.append(w)
                        print(shift.name, [n.name for n in nwd_workers])
                                
                        if not nwd_workers: # Just do the usual
                            self.model.constraints.add(
                            sum(self.model.blocks_worker_shift_assigned[worker.id,
                                                                        week,
                                                                        shift.name]
                                for worker in self.workers) <= 1)
                                
                        else:
                            print(shift.name)
                            self.model.constraints.add(
                                sum(self.model.blocks_worker_shift_assigned[worker.id,
                                                                        week,
                                                                        shift.name]
                                for worker in full_workers) <= 1)
                                    
                            self.model.constraints.add(
                                sum(self.model.blocks_worker_shift_assigned[worker.id,
                                                                        week,
                                                                        shift.name]
                                for worker in workers) <= 2)
                                
                            
                            #self.model.constraints.add(
                            #    self.model.worker_block_shift_split[week, shift.name] <= 2 - sum(self.model.blocks_worker_shift_assigned[worker.id,
                            #                                            week,
                            #                                            shift.name]
                            #    for worker in workers))
                            #  
                            #self.model.constraints.add(
                            #    sum(self.model.works[w.id, week, day,shift.name]
                            #    for w in nwd_workers for day in shift.shift_days) - sum(self.model.works[w.id, week, day,shift.name]
                            #    for w in full_workers for day in shift.shift_days) >= self.model.worker_block_shift_split[week, shift.name] )
                                    
                    else:
                        self.model.constraints.add(
                            sum(self.model.blocks_worker_shift_assigned[worker.id,
                                                                        week,
                                                                        shift.name]
                                for worker in self.workers) <= 1)

        # Most of our constraints apply per worker
        for worker in self.workers:


            for week_blocks in self.get_week_block_iterator(4):
                # Prevent weekend shifts more than once every n weeks
                self.model.constraints.add(
                    self.constraint_options["max_shifts_per_month"] >= sum(
                        self.model.works[worker.id, week, day, shift.name] for week in week_blocks
                        for day in self.days for shift in self.get_shifts()))


            if self.constraint_options[
                    "avoid_st2_first_month"] and worker.grade == 2:
                # Avoid ST1s on the first month
                try:
                    self.model.constraints.add(0 == sum(
                        self.model.works[worker.id, week, day, shift] for week,
                        day, shift in self.get_all_shiftname_combinations()
                        if week in (1, 2, 3, 4) and shift in ("night_weekday", "night_weekend")))
                except ValueError as e:
                    print("Failure setting constraint", "avoid_st2_first_month")
                    print(e)
            # Count number of weekends an worker works
            if self.constraint_options["balance_weekends"]:
                self.model.constraints.add(
                    self.model.worker_weekend_count[worker.id] == sum(
                        self.model.works_weekend[worker.id, week]
                        for week in self.weeks))

            if self.constraint_options["max_weekends"]:
                self.model.constraints.add(
                    self.constraint_options["max_weekends"] >=
                    self.model.worker_weekend_count[worker.id])

            # Balance shifts within required limits (set by balance_offset)
            # If balance_offset is too restrictive (for the rota length)
            # no solution will be possible
            for shift in self.get_shifts():
                if (
                        worker.site in shift.site
                ):  # Each site specfies which sites self.workers can fullfill it

                    # calaculate the total number of shifts that need assigning over the rota
                    #total_shifts = (len(self.weeks) * len(shift.shift_days) *
                    total_shifts = self.shift_counts[shift] * shift.workers_required

                    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)

                    if self.use_previous_shifts:
                        if shift.name in worker.previous_shifts:
                            worked, allocated = worker.previous_shifts[shift.name]
                            target_shifts = target_shifts + float(allocated) - float(worked)

                    #print(worker.name, shift.name, target_shifts)
                    worker.shift_target_number[shift.name] = target_shifts

                    if shift.hard_constrain_shift:
                        adjusted_balance_offset = shift.balance_offset
                        if worker.fte_adj < 100:
                            adjusted_balance_offset = adjusted_balance_offset * self.ltft_balance_offset

                        max_shifts = target_shifts + adjusted_balance_offset * self.balance_offset_modifier
                        min_shifts = target_shifts - adjusted_balance_offset * self.balance_offset_modifier

                        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,
                            ))
                else:
                    # Make sure shifts aren't assigned to those from other sites
                    self.model.constraints.add(0 == sum(
                        self.model.works[worker.id, week, day, shift.name]
                        for week, day in self.get_week_day_combinations()))

                # model.shift_count stores the number of shifts that a worker is assigned to worker by shift
                # this is used (by the objective) to balance the total number of shifts worked
                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()))

            # Define shift_count_t1 and shift_count_t2 constraints for the object
            # Thus 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)
            if self.constraint_options["balance_shifts"]:
                self.model.constraints.add(
                    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.get_shifts()))
            #if "night" not in shift.constraints))

            if self.constraint_options["balance_bank_holidays"]:
                self.model.constraints.add(
                    self.model.bank_holiday_count[worker.id] == sum(
                        self.model.works[worker.id, week, day, shift] for week,
                        day, shift in self.get_all_shiftname_combinations()
                        if self.week_day_date_map[(week,
                                                   day)] in bank_holiday_map))

                xU = len(self.get_bank_holiday_week_days()) + 1
                xL = 1
                self.model.constraints.add(
                    inequality(
                        xL,
                        self.model.bank_holiday_count[worker.id] + 1,
                        xU,
                    ))

                self.model.constraints.add(
                    self.model.bank_holiday_count_w[worker.id] >= xL *
                    (self.model.bank_holiday_count[worker.id] + 1) * 2 -
                    xL * xL)

                self.model.constraints.add(
                    self.model.bank_holiday_count_w[worker.id] >= xU *
                    (self.model.bank_holiday_count[worker.id] + 1) * 2 -
                    xU * xU)

            if self.constraint_options["balance_nights"]:
                self.model.constraints.add(
                    self.model.night_shift_count[worker.id] == sum(
                        self.model.works[worker.id, week, day, shift.name]
                        for week, day in self.get_week_day_combinations()
                        for shift in self.get_shifts_with_constraint("night")))

                night_shift_target_number = sum(
                    worker.shift_target_number[shift.name]
                    for shift in self.get_shifts_with_constraint("night"))

                # min_shifts = night_shift_target_number - 20
                # max_shifts = night_shift_target_number + 20

                # self.model.constraints.add(
                #     inequality(
                #         min_shifts,
                #         self.model.night_shift_count[worker.id],
                #         max_shifts,
                #     ))

                self.model.constraints.add(
                    self.model.night_shift_count_t1[worker.id] -
                    self.model.night_shift_count_t2[worker.id] ==
                    self.model.night_shift_count[worker.id] -
                    night_shift_target_number)

                # This may need to be updated
                xU = 6
                xL = 1
                self.model.constraints.add(
                    inequality(
                        xL,
                        self.model.night_shift_count_t1[worker.id] +
                        self.model.night_shift_count_t2[worker.id] + 1,
                        xU,
                    ))

                self.model.constraints.add(
                    self.model.night_shift_count_w[worker.id] >= xL *
                    (self.model.night_shift_count_t1[worker.id] +
                     self.model.night_shift_count_t2[worker.id] + 1) * 2 -
                    xL * xL)

                self.model.constraints.add(
                    self.model.night_shift_count_w[worker.id] >= xU *
                    (self.model.night_shift_count_t1[worker.id] +
                     self.model.night_shift_count_t2[worker.id] + 1) * 2 -
                    xU * xU)

                # self.model.constraints.add(
                #     self.model.night_shift_count_w[worker.id] >= 0)

            # We use a similar method to balance the number of weekends worked
            # This works as long as weekend shifts are assigned as blocks!
            if self.constraint_options["balance_weekends"]:
                weekend_shift_target_number = sum(
                    worker.shift_target_number[shift.name] /
                    len(shift.shift_days) for shift in self.get_shifts()
                    if "Sat" in shift.shift_days or "Sun" in shift.shift_days)

                worker.weekend_shift_target_number = weekend_shift_target_number

                # min_shifts = weekend_shift_target_number - 20
                # max_shifts = weekend_shift_target_number + 20

                # self.model.constraints.add(
                #     inequality(
                #         min_shifts,
                #         self.model.worker_weekend_count[worker.id],
                #         max_shifts,
                #     ))

                self.model.constraints.add(
                    self.model.weekend_shift_count_t1[worker.id] -
                    self.model.weekend_shift_count_t2[worker.id] ==
                    self.model.worker_weekend_count[worker.id] -
                    weekend_shift_target_number)

                xU = self.constraint_options["max_weekends"]

                xL = 1
                self.model.constraints.add(
                    inequality(
                        xL,
                        self.model.weekend_shift_count_t1[worker.id] +
                        self.model.weekend_shift_count_t2[worker.id] + 1,
                        xU,
                    ))

                self.model.constraints.add(
                    self.model.weekend_shift_count_w[worker.id] >= xL *
                    (self.model.weekend_shift_count_t1[worker.id] +
                     self.model.weekend_shift_count_t2[worker.id] + 1) * 2 -
                    xL * xL)

                self.model.constraints.add(
                    self.model.weekend_shift_count_w[worker.id] >= xU *
                    (self.model.weekend_shift_count_t1[worker.id] +
                     self.model.weekend_shift_count_t2[worker.id] + 1) * 2 -
                    xU * xU)

            # Ensure worker is not allocated shifts on non working days
            if worker.nwd:
                for week, day, shift in self.get_all_shiftclass_combinations():
                    for n, start, end in worker.nwd:
                        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:
                                self.model.constraints.add(
                                    0 == self.model.works[worker.id, week, day,
                                                        shift.name])

            if self.constraint_options["balance_blocks"]:
                for week_blocks in self.get_week_block_iterator(
                        self.max_night_frequency):
                    if self.get_shifts_with_constraint("night"):
                        # Prevent nights more than once every n weeks
                        self.model.constraints.add(1 >= sum(
                            self.model.shift_week_worker_assigned[shift.name,
                                                                  week,
                                                                  worker.id]
                            for week in week_blocks
                            for shift in self.get_shifts_with_constraint(
                                "night")))

                for week_blocks in self.get_week_block_iterator(
                        self.max_weekend_frequency):
                    # Prevent weekend shifts more than once every n weeks
                    self.model.constraints.add(
                        1 >= sum(self.model.works_weekend[worker.id, week]
                                 for week in week_blocks))

            for week in self.weeks:

                self.model.constraints.add(self.constraint_options["max_shifts_per_week"] >= sum(
                    self.model.works[worker.id, week, day, shift.name]
                    for day in self.days for shift in self.get_shifts()))

                for shift in self.get_shifts_with_constraint(
                        "max_2_shifts_per_week"):
                    self.model.constraints.add(2 >= sum(
                        self.model.works[worker.id, week, day, shift.name]
                        for day in self.days))

                if worker.site in self.constraint_options[
                        "prevent_monday_after_full_weekends"]:
                    # Ignore last week
                    if week + 1 in self.weeks:
                        self.model.constraints.add(
                            full_weekend_count >= sum(self.model.works[worker.id, week, "Sat",
                                                      shift.name]
                                     for shift in self.get_shifts()) +
                            sum(self.model.works[worker.id, week, "Sun",
                                                 shift.name]
                                for shift in self.get_shifts()) +
                            sum(self.model.works[worker.id, week + 1, "Mon",
                                                 shift.name]
                                for shift in self.get_shifts()))
                if worker.site in self.constraint_options[
                        "prevent_monday_and_tuesday_after_full_weekends"]:
                    # Ignore last week
                    if week + 1 in self.weeks:
                        self.model.constraints.add(
                            full_weekend_count >= sum(self.model.works[worker.id, week, "Sat",
                                                      shift.name]
                                     for shift in self.get_shifts()) +
                            sum(self.model.works[worker.id, week, "Sun",
                                                 shift.name]
                                for shift in self.get_shifts()) +
                            sum(self.model.works[worker.id, week + 1, "Mon",
                                                 shift.name]
                                for shift in self.get_shifts()) +
                            sum(self.model.works[worker.id, week + 1, "Tue",
                                                 shift.name]
                                for shift in self.get_shifts()))
                if worker.site in self.constraint_options[
                        "prevent_fridays_before_full_weekends"]:
                    self.model.constraints.add(
                        full_weekend_count >= sum(self.model.works[worker.id, week, "Sat",
                                                  shift.name]
                                 for shift in self.get_shifts()
                                 if "night" not in shift.constraints) +
                        sum(self.model.works[worker.id, week, "Sun",
                                             shift.name]
                            for shift in self.get_shifts()
                            if "night" not in shift.constraints) +
                        sum(self.model.works[worker.id, week, "Fri",
                                             shift.name]
                            for shift in self.get_shifts()
                            if "night" not in shift.constraints))
                if worker.site in self.constraint_options[
                        "prevent_thursdays_before_full_weekends"]:
                    self.model.constraints.add(
                        full_weekend_count >= sum(self.model.works[worker.id, week, "Sat",
                                                  shift.name]
                                 for shift in self.get_shifts()
                                 if "night" not in shift.constraints) +
                        sum(self.model.works[worker.id, week, "Sun",
                                             shift.name]
                            for shift in self.get_shifts()
                            if "night" not in shift.constraints) +
                        sum(self.model.works[worker.id, week, "Thu",
                                             shift.name]
                            for shift in self.get_shifts()
                            if "night" not in shift.constraints))

                # # model.weekend_count stores the number of weekend shifts that a worker is assigned to wor
                # # this is used (by the objective) to balance the total number of weekend shifts worked
                # self.model.constraints.add(
                #     # We could use a helper to get required shifts to check
                #     self.model.works_whole_weekend[worker.id, week] >= sum(
                #         self.model.works[worker.id, week, day, shift.name]
                #         for day in self.days[5:]
                #         for shift in self.get_shifts()) - 1)
                # self.model.constraints.add(
                #     self.model.works_whole_weekend[worker.id, week] <= sum(
                #         self.model.works[worker.id, week, day, shift.name]
                #         for day in self.days[5:]
                #         for shift in self.get_shifts()) / 2)
                if self.constraint_options["balance_weekends"]:
                    self.model.constraints.add(
                        self.model.works_weekend[worker.id, week] >= sum(
                            self.model.works[worker.id, week, day, shift.name]
                            for day in self.days[5:]
                            for shift in self.get_shifts()) / 2)
                    self.model.constraints.add(
                        self.model.works_weekend[worker.id, week] <= sum(
                            self.model.works[worker.id, week, day, shift.name]
                            for day in self.days[5:]
                            for shift in self.get_shifts()))
                # self.model.constraints.add(
                #     self.model.works_saturday[worker.id, week] == sum(
                #         self.model.works[worker.id, week, "Sat", shift.name]
                #         for shift in self.get_shifts()))
                # self.model.constraints.add(
                #     self.model.works_sunday[worker.id, week] == sum(
                #         self.model.works[worker.id, week, "Sun", shift.name]
                #         for shift in self.get_shifts()))

                # self.model.constraints.add(1 == sum(
                #     self.model.shift_week_worker_assigned[s.name, week, worker.id]
                #     for s in self.get_shifts_with_constraint("night")))

                for shift in self.get_shifts_with_constraint("night"):
                    # Force nights to be assigned in blocks
                    # self.model.constraints.add(8* self.model.shift_week_worker_assigned[shift.name, week, worker.id] >= sum(self.model.works[worker.id, week, day, shift.name] for day in self.days)
                    #     )

                    # Force night shifts to be assigned in blocks
                    self.model.constraints.add(
                        self.model.shift_week_worker_assigned[shift.name, week,
                                                              worker.id] *
                        len(shift.shift_days) == sum(
                            self.model.works[worker.id, week, day, shift.name]
                            for day in shift.shift_days))

                #

            for n in range(len(weeks_days)):

                week, day = weeks_days[n]

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

                try:
                    p2week, p2day = weeks_days[n - 2]
                except IndexError:
                    p2week, p2day = weeks_days[n]
                    p2 = 0

                n1 = 1
                n2 = 1
                try:
                    nweek, nday = weeks_days[n + 1]
                except IndexError:
                    nweek, nday = weeks_days[n]
                    n1 = 0

                try:
                    n2week, n2day = weeks_days[n + 2]
                except IndexError:
                    n2week, n2day = weeks_days[n]
                    n2 = 0

                try:
                    n3week, n3day = weeks_days[n + 3]
                except IndexError:
                    n3week, n3day = weeks_days[n]
                    n3 = 0

                if self.get_shift_names_by_week_day(week, day):
                    # 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_names_by_week_day(week, day)))

                    # single shift per day
                    self.model.constraints.add(1 >= sum(
                        self.model.works[worker.id, week, day, shift]
                        for shift in self.get_shift_names_by_week_day(week, day)))

                    if self.constraint_options["hard_constrain_pair_separation"]:
                        for worker_pairs in self.worker_pairs:
                            if worker_pairs[0] == worker:
                                self.model.constraints.add(1 >= sum(self.model.works[w.id, week, day, shift] for shift in self.get_shift_names_by_week_day(week, day) for w in worker_pairs))

                # if working a night ensure preceeding (1) or subsequent (2) shifts can only be nights
                if self.constraint_options["constrain_time_off_after_nights"]:
                    for constraint_shift in self.get_shifts_with_constraint(
                            "night"):
                        # Ensure night prior to unavalibity is not assigned
                        self.model.constraints.add(
                            self.model.available[worker.id, week, day] >=
                            self.model.works[worker.id, pweek, pday,
                                             constraint_shift.name])

                        # IF paired we check the following against both workers
                        workers = [worker]
                        if self.constraint_options["hard_constrain_pair_separation"]:
                            for worker_pairs in self.worker_pairs:
                                if worker_pairs[0] == worker:
                                    workers = worker_pairs

                        #print("Workers", workers)


                        self.model.constraints.add(
                            1 >= self.model.works[worker.id, week, day,
                                                  constraint_shift.name] +
                            sum(n1 *
                                self.model.works[w.id, nweek, nday, shift]
                                for shift in self.get_shift_names_by_week_day(nweek, nday)
                                if shift != constraint_shift.name for w in workers))

                        self.model.constraints.add(
                            1 >= self.model.works[worker.id, week, day,
                                                  constraint_shift.name] +
                            sum(n2 * self.model.works[w.id, n2week, n2day,
                                                      shift]
                                for shift in self.get_shift_names_by_week_day(n2week, n2day)
                                if shift != constraint_shift.name for w in workers))
                        # self.model.constraints.add(
                        #     1 >= self.model.works[worker.id, week, day,
                        #                           constraint_shift.name] +
                        #     sum(n3 * self.model.works[worker.id, n3week, n3day,
                        #                               shift]
                        #         for shift in self.get_shift_names_by_week_day(week, n3day)
                        #         if shift != constraint_shift.name))
                        self.model.constraints.add(
                            1 >= self.model.works[worker.id, week, day,
                                                  constraint_shift.name] +
                            sum(p1 *
                                self.model.works[w.id, pweek, pday, shift]
                                for shift in self.get_shift_names_by_week_day(pweek, pday)
                                if shift != constraint_shift.name for w in workers))
                        self.model.constraints.add(
                            1 >= self.model.works[worker.id, week, day,
                                                  constraint_shift.name] +
                            sum(p2 * self.model.works[w.id, p2week, p2day,
                                                      shift]
                                for shift in self.get_shift_names_by_week_day(p2week, p2day)
                                if shift != constraint_shift.name for w in workers))

        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)

            balance_modifier_constant = 1

            if self.constraint_options["balance_shifts"]:
                shift_balancing = sum(balance_modifier_constant *
                                      (self.model.shift_count_t1[(worker.id)] +
                                       self.model.shift_count_t2[(worker.id)])
                                      for worker in self.workers)
            else:
                shift_balancing = 0

            if self.constraint_options["balance_nights"]:
                night_balance_modifier_constant = 10
                night_shift_balancing = sum(
                    night_balance_modifier_constant *
                    self.model.night_shift_count_w[(worker.id)]
                    # (self.model.night_shift_count_t1[(worker.id)] +
                    #  self.model.night_shift_count_t2[(worker.id)])
                    for worker in self.workers)
            else:
                night_shift_balancing = 0

            if self.constraint_options["balance_bank_holidays"]:
                bank_holiday_balance_modifier_constant = 1000
                bank_holiday_balancing = sum(
                    bank_holiday_balance_modifier_constant *
                    (self.model.bank_holiday_count_w[(worker.id)] - 1)
                    # (self.model.night_shift_count_t1[(worker.id)] +
                    #  self.model.night_shift_count_t2[(worker.id)])
                    for worker in self.workers)
            else:
                bank_holiday_balancing = 0

            if self.constraint_options["balance_weekends"]:
                weekend_balance_modifier_constant = 5
                weekend_shift_balancing = sum(
                    weekend_balance_modifier_constant *
                    self.model.weekend_shift_count_w[(worker.id)]
                    for worker in self.workers)
            else:
                weekend_shift_balancing = 0

            preference_constant = 10
            # Preferences (not to work)
            preferences = sum(
                self.model.pref_not_to_work[worker.id, week, day] *
                self.model.works[worker.id, week, day, shift] *
                preference_constant for worker in self.workers
                for week, day, shift in self.get_all_shiftname_combinations())

            # Work requsets
            work_request_constant = 10000  # Needs to be very big to override loss from bank holiday requests
            work_requests = sum(
                self.model.work_requests[worker.id, week, day, shift] *
                self.model.works[worker.id, week, day, shift] *
                work_request_constant for worker in self.workers
                for week, day, shift in self.get_all_shiftname_combinations())

            # # Spread nights

            if self.constraint_options["balance_nights_across_sites"]:
                nights_site_balancing = sum(
                    (self.model.night_per_site_t1[week, shift.name, site] +
                     self.model.night_per_site_t2[week, shift.name, site]) * 20
                    #self.model.night_per_site2[week, block, site]
                    for week in self.weeks
                    for shift in self.get_shifts_with_constraint("night")
                    for site in self.sites)
            else:
                nights_site_balancing = 0

            if self.constraint_options["balance_blocks"]:
                blocks_balancing = sum(
                    1 * self.model.blocks_assigned[week, shift]
                    for week in self.weeks
                    for shift in self.shifts_to_assign_as_blocks())
            else:
                blocks_balancing = 0

            # Quadratic :(
            #worker_pairs_balancing = 0
            #worker_pairs_constant = 100000
            #if self.worker_pairs:
            #    for worker_a, worker_b in self.worker_pairs:
            #        print(worker_a, worker_b)
            #        for week, day in self.get_week_day_combinations():
            #            worker_pairs_balancing = worker_pairs_balancing + sum(self.model.works[worker_a.id, week, day, shift] for shift in self.get_shift_names_by_week_day(week, day)) * sum(self.model.works[worker_b.id, week, day, shift] for shift in self.get_shift_names_by_week_day(week, day)) * worker_pairs_constant 


            #return weekend_shift_balancing + blocks_balancing
            #return bank_holiday_balancing
            #return shift_balancing + preferences + blocks_balancing
            #return shift_balancing + preferences + nights_site_balancing + blocks_balancing
            return weekend_shift_balancing + shift_balancing + night_shift_balancing + preferences + nights_site_balancing + bank_holiday_balancing + blocks_balancing - work_requests

        # 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_shifts_and_workers(self):
        """Process loaded shifts and workers

        Must be called prior to attempting to solve
        """
        #self.build_shifts()

        self.workers = sorted(self.workers)

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

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

        pairs = defaultdict(list)
        for w in self.workers:
            if w.pair > 0:
                pairs[w.pair].append(w)

        for p in pairs:
            self.worker_pairs.append(tuple(pairs[p]))


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

        :param SingleShift shift: Shift object

        """
        self.shifts.append(shift)

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

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

        Returns:
            set: Set of ShiftName

        """
        return self.week_day_shifts_dict[(week, day)]

    def get_shift_by_name(self, name):
        """

        :param name: 

        """
        return self.shifts_by_name[name]

    def get_shift_length_by_name(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 build_shifts(self):
        """ """
        self.shifts_by_name = {}
        self.shift_names = []  # type: List[ShiftName]

        #self.week_day_shifts_dict = {(week, day): set() for week in self.weeks for day in days}
        self.week_day_shifts_dict = {}
        self.sites = set()


        self.week_day_shift_product = []
        self.week_day_shiftclass_product = []
        
        self.shift_counts = defaultdict(int)
        for week, day in self.weeks_days_product:
            self.week_day_shifts_dict[(week, day)] = set()
            for s in self.shifts:
                if s.bank_holidays_only:
                    if self.week_day_date_map[(week, day)] in bank_holiday_map:
                        self.week_day_shift_product.append((week, day, s.name))
                        self.week_day_shiftclass_product.append((week, day, s))
                        self.week_day_shifts_dict[(week, day)].add(s.name)
                        self.shift_counts[s] = self.shift_counts[s] + 1
                else:
                    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.week_day_shifts_dict[(week, day)].add(s.name)
                        self.shift_counts[s] = self.shift_counts[s] + 1
        #print(self.week_day_shift_product)   
        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.week_day_shifts_dict[(week, day)].add(s.name)

            for site in s.site:
                self.sites.add(site)             
                
        # Todo replace with week_day.....
        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 get_day_shiftname_combinations(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_combinations(
            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 get_all_shiftname_combinations(
            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 get_shifts(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_shifts_with_constraint(self, constraint) -> List[SingleShift]:
        return [
            shift for shift in self.shifts if constraint in shift.constraints
        ]

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

    def shifts_to_assign_as_blocks(self):
        return [
            shift.name for shift in self.get_shifts() if shift.assign_as_block
        ]

    def shifts_to_force_as_blocks(self):
        return [
            shift.name for shift in self.get_shifts() if shift.force_as_block
        ]

    def shifts_to_assign_or_force_as_blocks(self):
        s = self.shifts_to_assign_as_blocks()
        s.extend(self.shifts_to_force_as_blocks())
        return s
        
    def get_workers_for_shift(self, shift):
        return [worker for worker in self.workers if worker.site in shift.site]

    def get_workers_total_fte(self):
        return self.full_time_equivalent

    def get_worker_details(self):

        w = defaultdict(list)
        for worker in self.workers:
            #print(worker)
            w[worker.site].append(worker)

        l = []
        for site in w:
            l.append("{}: {}".format(
                site,
                sum([worker.fte_adj for worker in w[site]]) / 100))
        t = "\n".join(l)
        return "Full time equivalent trainees by site:\n{}".format(t)

    def get_bank_holiday_week_days(self):
        return ([(week, day) for week, day in self.get_week_day_combinations()
                 if self.week_day_date_map[(week, day)] in bank_holiday_map])


class RotaResults(object):
    def __init__(self, rota, results):
        self.rota = rota
        self.results = results

    def export_rota_to_html(self, filename="rota"):
        with open("{}.html".format(filename), "w") as f:
            f.write(self.get_worker_timetable_html(True))

    def export_rota_to_csv(self, filename="rota"):
        works = self.rota.model.works
        with open("{}.csv".format(filename), 'w', newline='') as f:
            wr = csv.writer(f, quoting=csv.QUOTE_ALL)
            l = ["Name"]
            l.extend([worker.name for worker in self.rota.workers])
            wr.writerow(l)

            l2 = ["Site"]
            l2.extend([worker.site for worker in self.rota.workers])
            wr.writerow(l2)

            l3 = ["Grade"]
            l3.extend([worker.grade for worker in self.rota.workers])
            wr.writerow(l3)

            l4 = ["FTE"]
            l4.extend([worker.fte for worker in self.rota.workers])
            wr.writerow(l4)

            for week, day in self.rota.get_week_day_combinations():
                d = ["Week {} Day {}".format(week, day)]

                for worker in self.rota.workers:
                    i = ""
                    for shift in self.rota.get_shift_names_by_week_day(week, day):
                        if works[worker.id, week, day, shift].value == 1:
                            i = shift

                    d.append(i)
                wr.writerow(d)

    def get_work_table(self):
        """Build a timetable of the week as a dictionary from the model's optimal solution."""
        works = self.rota.model.works
        week_table = {
            week: {
                day: {shift: []
                      for shift in self.rota.get_shift_names()}
                for day in days
            }
            for week in self.rota.weeks
        }
        for week in self.rota.weeks:
            for worker in self.rota.workers:
                for day, shift in self.rota.get_day_shiftname_combinations():
                    if works[worker.id, week, day, shift].value == 1:
                        week_table[week][day][shift].append(
                            worker.get_details())
        return week_table

    def get_worker_timetable(self):
        works = self.rota.model.works
        timetable = {
            worker.name:
            {week: {day: ""
                    for day in days}
             for week in self.rota.weeks}
            for worker in self.rota.workers
        }
        for worker in self.rota.workers:
            for week in self.rota.weeks:
                for day, shift in self.rota.get_day_shiftname_combinations():
                    if works[worker.id, week, day, shift].value == 1:
                        timetable[worker.name][week][day] = shift
        return timetable

    def get_worker_timetable_brief(self,
                                   show_prefs=False,
                                   marker_every=30,
                                   show_unavailable=False):
        model = self.rota.model
        week_string = "{:20}".format("-Week-") + "".join(
            [7 * str("{}".format(w))[-1:] for w in self.rota.weeks])
        days_string = "{:20}".format("-Day-") + "".join(
            "MTWTFSS" * len(self.rota.weeks))
        timetable = []
        for worker in self.rota.workers:
            shifts = []
            w = ["{:20}".format(worker.name)]
            for week, day in self.rota.get_week_day_combinations():
                a = "-"
                for shift in self.rota.get_shift_names_by_week_day(week, day):
                    if model.works[worker.id, week, day, shift].value > 0:
                        shifts.append(shift)
                        a = shift[0]
                w.append(a)

            shift_count = ""
            for s in set(shifts):
                shift_count = shift_count + "{}: {}, ".format(
                    s, shifts.count(s))

            shift_count = shift_count + "#weekends_worked: {}\\#".format(
                model.worker_weekend_count[worker.id].value)
            timetable.append("{} {}".format("".join(w), shift_count))

            if show_prefs:
                # prefs
                w = ["{:20}".format("Preferences")]
                for week, day in self.rota.get_week_day_combinations():
                    if model.pref_not_to_work[worker.id, week, day] > 0:
                        w.append("Y")
                    else:
                        w.append("N")
                timetable.append("".join(w))

            if show_unavailable:
                # prefs
                w = ["{:20}".format("Unavailable")]
                for week, day in self.rota.get_week_day_combinations():
                    if model.available[worker.id, week, day] > 0:
                        w.append("A")
                    else:
                        w.append("U")
                timetable.append("".join(w))

        i = marker_every
        while i < len(timetable):
            timetable.insert(i, week_string)
            timetable.insert(i + 1, days_string)
            i = i + marker_every + 2

        timetable.append(week_string)
        timetable.insert(0, week_string)
        timetable.insert(0, days_string)
        return "\n".join(timetable)

    def get_worker_timetable_html(self,
                                  include_html_tag=False,
                                  table_name="rota-table"):
        model = self.rota.model

        timetable = []

        timetable.append("<h2>Rota start date: {} ({} weeks)</h2>".format(self.rota.start_date.isoformat(), self.rota.weeks[-1]))

        date_row = ["<th class='worker'></th>"]

        n = 0
        for week, day in self.rota.get_week_day_combinations():
            d = self.rota.start_date + datetime.timedelta(n)
            date_row.append("<th title='{}'>Week {}: {}</th>".format(
                d, week, day))
            n = n + 1
        timetable.append("<tr class='data-row'>{}</tr>".format(
            "".join(date_row)))

        current_site = ""

        for worker in self.rota.workers:
            if worker.site != current_site:
                try:
                    timetable.append("<tr><th class='site-title'>{} n={} fte={}</th><tr>".format(worker.site, len(self.rota.workers_at_sites[worker.site]), self.rota.full_time_equivalent_sites[worker.site]))
                except KeyError as e:
                    print(e)
                current_site = worker.site

            shifts = []
            if worker.nwd is not None:
                # TODO: limit to dates
                nwds = ", ".join([i[0] for i in worker.nwd])
            else:
                nwds = None

            worker_targets = json.dumps(worker.shift_target_number)

            tds = []

            n = 0

            shift_tds = []
            for week, day in self.rota.get_week_day_combinations():
                d = self.rota.start_date + datetime.timedelta(n)

                n = n + 1
                a = "-"
                shift_name = ""

                for shift in self.rota.get_shift_names_by_week_day(week, day):
                    if model.works[worker.id, week, day, shift].value > 0:
                        shifts.append(shift)
                        a = shift[0]
                        shift_name = shift
                        break
                title = "{} ({})".format(shift_name, d)
                css_class = day
                unavailable_reason = ""
                if model.available[worker.id, week, day] > 0:
                    available = True
                else:
                    available = False
                    css_class = "unavailable"
                    try:
                        unavailable_reason = self.rota.unavailable_to_work_reason[(worker.id, week, day)]
                    except KeyError:
                        print("Error getting reason")
                        print(worker.id, worker.name)
                        print("Week {}, Day {}".format(week, day))
                        unavailable_reason = "????"
                    title = "{} / {}".format(title, unavailable_reason)

                bank_holiday = ""
                if d in bank_holiday_map:
                    css_class = " ".join((css_class, "bank-holiday"))
                    bank_holiday = " data-bank-holiday='{}'".format(
                        bank_holiday_map[d])

                requests = ""
                if (worker.id, week, day) in self.rota.work_requests_map:
                    css_class = " ".join((css_class, "shift-requested"))
                    title = " ".join((title, "[REQUESTED]"))
                    requests = " data-shift-request='{}'".format(
                        self.rota.work_requests_map[(worker.id, week, day)])

                shift_tds.append(
                    "<td title='{}' class='rota-day {}' data-shift='{}' data-available='{}' data-unavailable_reason='{}' data-date='{}' data-week='{}' data-day='{}'{}{}>{}</td>"
                    .format(title, css_class, shift_name, available, unavailable_reason, d, week,
                            day, requests, bank_holiday, a))

            shift_count = ""
            shift_count_dict = {}
            for s in set(shifts):
                c = shifts.count(s)
                shift_count_dict[s] = c
                shift_count = shift_count + "{}: {}, ".format(s, c)

            worker_td = "<td title='Site: {site}' class='worker {site}' data-nwds='{nwds}' data-site='{site}' data-worker='{name}' data-fte='{fte}' data-fte_adj='{fte_adj}' data-end_date='{end_date}' data-worker-targets='{targets}' data-shift-counts='{shift_counts}' data-weekend-target='{weekend_target}' data-night-at-derriford='{nights_at_derriford}' data-pair='{pair}'><span class='name' title='{name}'>{name}</span> ({grade}) [{fte}]</td>".format(
                site=worker.site, nwds=nwds, name=worker.name,
                fte=worker.fte, fte_adj=worker.fte_adj, end_date=worker.end_date, targets=worker_targets,
                shift_counts=json.dumps(shift_count_dict),
                weekend_target=worker.weekend_shift_target_number, nights_at_derriford=worker.night_at_derriford, grade=worker.grade, pair=worker.pair
                )

            shift_count = shift_count + "#weekends_worked: {}\\#".format(
                model.worker_weekend_count[worker.id].value)

            bank_holiday_count = model.bank_holiday_count[worker.id].value
            bank_holiday_count_w = model.bank_holiday_count_w[
                worker.id].value - 1
            #print(worker.name, bank_holiday_count, bank_holiday_count_w)
            timetable.append("<tr class='worker-row'>{}{}</tr>".format(
                worker_td, "".join(shift_tds)))
            #timetable.append("<tr>{}</tr>".format("".join(w), shift_count))

            # if show_prefs:
            #     # prefs
            #     w = ["{:20}".format("Preferences")]
            #     for week, day in self.rota.get_week_day_combinations():
            #         if model.pref_not_to_work[worker.id, week, day] > 0:
            #             w.append("Y")
            #         else:
            #             w.append("N")
            #     timetable.append("".join(w))

            # if show_unavailable:
            #     # prefs
            #     w = ["{:20}".format("Unavailable")]
            #     for week, day in self.rota.get_week_day_combinations():
            #         if model.available[worker.id, week, day] > 0:
            #             w.append("A")
            #         else:
            #             w.append("U")
            #     timetable.append("".join(w))

        
        result_stream = StringIO()

        self.results.write(ostream=result_stream)
        self.results.write_json(ostream=result_stream)
        result_string = result_stream.getvalue()

        html = """ 
        <body>
        <div class="table-div" id="{}">
        <table>{}</table>
        </div>
        <details>
        <summary><h2>Rota settings</h2></summary>
        <div>
        <pre>
        {}
        </pre>
        </details>
        </div>
        <details>
        <summary><h2>Shifts settings</h2></summary>
        <div>
        {}
        </div>
        </details>
        <details>
        <summary><h2>Output</h2></summary>
        <pre>
        {}
        </pre>
        <div>
        </body>
        """.format(table_name, "\n".join(timetable), json.dumps(self.rota.constraint_options, indent=4), "<br/>".join(str(i) for i in self.rota.shifts), result_string)

        if include_html_tag:
            html = """<html>
        <head>
        <link rel="stylesheet" type="text/css" href="timetable.css">
        <script src="https://ajax.googleapis.com/ajax/libs/jquery/3.5.1/jquery.min.js"></script>
        <script src="timetable.js" defer></script>
        </head>
            {}</html>""".format(html)

        return html

    def get_shift_summary(self):
        works = self.rota.model.works
        timetable = {
            worker.get_details():
            {shift[0]: ""
             for shift in self.rota.get_shift_names()}
            for worker in self.rota.workers
        }
        #timetable = { worker.get_details() : { shift : "" for shift in ["truro_twilight"] } for worker in workers }
        t = []
        for worker in self.rota.workers:
            l = []

            total_shifts = 0
            for shift in self.rota.get_shift_names():
                #for shift in ["truro_twilight"]:
                c = [
                    works[worker.id, week, day, shift]
                    for week in self.rota.weeks for day in days
                ].count(1)
                if c > 0:
                    l.append("{} ({})".format(shift, c))
                    total_shifts = total_shifts + c
                #print(worker.id, shift)
                timetable[worker.get_details()][shift[0]] = c
            t.append("{} [{}]: {}".format(worker.get_full_details(),
                                          total_shifts, ", ".join(l)))
        return "\n".join(t)

    def get_shift_summary_html(self):
        works = self.rota.model.works
        timetable = {
            worker.get_details():
            {shift[0]: ""
             for shift in self.rota.get_shift_names()}
            for worker in self.rota.workers
        }
        #timetable = { worker.get_details() : { shift : "" for shift in ["truro_twilight"] } for worker in workers }
        t = []
        for worker in self.rota.workers:
            l = []

            total_shifts = 0
            for shift in self.rota.get_shift_names():
                #for shift in ["truro_twilight"]:
                c = [
                    works[worker.id, week, day, shift]
                    for week in self.rota.weeks for day in days
                ].count(1)
                if c > 0:
                    l.append("{} ({})".format(shift, c))
                    total_shifts = total_shifts + c
                #print(worker.id, shift)
                timetable[worker.get_details()][shift[0]] = c
            t.append("{} [{}]: {}".format(worker.get_full_details(),
                                          total_shifts, ", ".join(l)))
        return "\n".join(t)

    #def get_no_preference(no_pref):
    #    """Extract to a list the workers not satisfied with their weekend preference."""
    #    return [worker.id for worker in workers if no_pref[worker.id].value == 1]

    # def get_night_blocks(self):
    #     nights = self.rota.model.nights
    #     timetable = {
    #         worker.name: {
    #             week: {block: ""
    #                    for block in self.rota.night_blocks}
    #             for week in self.rota.weeks
    #         }
    #         for worker in self.rota.workers
    #     }
    #     for worker in self.rota.workers:
    #         for week in self.rota.weeks:
    #             for block in self.rota.night_blocks:
    #                 if nights[worker.id, week, block].value == 1:
    #                     timetable[worker.name][week][block] = "true"
    #     return timetable

    def get_multinight(self):
        for site in self.rota.sites:
            for shift in self.get_shifts_with_constraint("night"):
                block = shift.name
                for week in self.rota.weeks:
                    print(
                        site, week, block,
                        self.rota.model.night_per_site[week, block,
                                                       site].value,
                        self.rota.model.night_per_site_t1[week, block,
                                                          site].value,
                        self.rota.model.night_per_site_t2[week, block,
                                                          site].value)
    """
    def get_night_details(self):
        for worker in self.rota.workers:
            print(
                "{:20}".format(worker.name),
                "worked: {},".format(
                    self.rota.model.night_shift_count[(worker.id)].value),
                    worker.shift_target_number["night_weekday"] +
                    worker.shift_target_number["night_weekend"]),
                "target_diff: {},".format(
                    self.rota.model.night_shift_count_t1[(worker.id)].value +
                    self.rota.model.night_shift_count_t2[(worker.id)].value +
                    1),
                "balance: {},".format(
                    self.rota.model.night_shift_count_w[worker.id].value),
            )

    def get_weekend_details(self):
        for worker in self.rota.workers:
            print(
                "{:20}".format(worker.name),
                "worked: {},".format(
                    self.rota.model.worker_weekend_count[(worker.id)].value),
                "target: {},".format(
                    sum([
                        worker.shift_target_number[shift.name]
                        for shift in self.rota.get_shifts()
                        if not set(("Sat", "Sun")).isdisjoint(shift.shift_days)
                    ])),
                "target_diff: {},".format(
                    self.rota.model.weekend_shift_count_t1[(worker.id)].value +
                    self.rota.model.weekend_shift_count_t2[(worker.id)].value +
                    1),
                "balance: {},".format(
                    self.rota.model.weekend_shift_count_w[worker.id].value),
            )
            
    """