|
|
from datetime import date, datetime, time, timedelta |
|
|
from itertools import product |
|
|
from enum import Enum |
|
|
from random import Random |
|
|
from typing import Generator |
|
|
from dataclasses import dataclass, field |
|
|
|
|
|
from .domain import Employee, EmployeeSchedule, Shift |
|
|
|
|
|
|
|
|
class DemoData(Enum): |
|
|
SMALL = 'SMALL' |
|
|
LARGE = 'LARGE' |
|
|
|
|
|
|
|
|
@dataclass(frozen=True, kw_only=True) |
|
|
class CountDistribution: |
|
|
count: int |
|
|
weight: float |
|
|
|
|
|
|
|
|
def counts(distributions: tuple[CountDistribution, ...]) -> tuple[int, ...]: |
|
|
return tuple(distribution.count for distribution in distributions) |
|
|
|
|
|
|
|
|
def weights(distributions: tuple[CountDistribution, ...]) -> tuple[float, ...]: |
|
|
return tuple(distribution.weight for distribution in distributions) |
|
|
|
|
|
|
|
|
@dataclass(kw_only=True) |
|
|
class DemoDataParameters: |
|
|
locations: tuple[str, ...] |
|
|
required_skills: tuple[str, ...] |
|
|
optional_skills: tuple[str, ...] |
|
|
days_in_schedule: int |
|
|
employee_count: int |
|
|
optional_skill_distribution: tuple[CountDistribution, ...] |
|
|
shift_count_distribution: tuple[CountDistribution, ...] |
|
|
availability_count_distribution: tuple[CountDistribution, ...] |
|
|
random_seed: int = field(default=37) |
|
|
|
|
|
|
|
|
demo_data_to_parameters: dict[DemoData, DemoDataParameters] = { |
|
|
DemoData.SMALL: DemoDataParameters( |
|
|
locations=("Ambulatory care", "Critical care", "Pediatric care"), |
|
|
required_skills=("Doctor", "Nurse"), |
|
|
optional_skills=("Anaesthetics", "Cardiology"), |
|
|
days_in_schedule=14, |
|
|
employee_count=15, |
|
|
optional_skill_distribution=( |
|
|
CountDistribution(count=1, weight=3), |
|
|
CountDistribution(count=2, weight=1) |
|
|
), |
|
|
shift_count_distribution=( |
|
|
CountDistribution(count=1, weight=0.9), |
|
|
CountDistribution(count=2, weight=0.1) |
|
|
), |
|
|
availability_count_distribution=( |
|
|
CountDistribution(count=1, weight=4), |
|
|
CountDistribution(count=2, weight=3), |
|
|
CountDistribution(count=3, weight=2), |
|
|
CountDistribution(count=4, weight=1) |
|
|
), |
|
|
random_seed=37 |
|
|
), |
|
|
|
|
|
DemoData.LARGE: DemoDataParameters( |
|
|
locations=("Ambulatory care", |
|
|
"Neurology", |
|
|
"Critical care", |
|
|
"Pediatric care", |
|
|
"Surgery", |
|
|
"Radiology", |
|
|
"Outpatient"), |
|
|
required_skills=("Doctor", "Nurse"), |
|
|
optional_skills=("Anaesthetics", "Cardiology", "Radiology"), |
|
|
days_in_schedule=28, |
|
|
employee_count=50, |
|
|
optional_skill_distribution=( |
|
|
CountDistribution(count=1, weight=3), |
|
|
CountDistribution(count=2, weight=1) |
|
|
), |
|
|
shift_count_distribution=( |
|
|
CountDistribution(count=1, weight=0.5), |
|
|
CountDistribution(count=2, weight=0.3), |
|
|
CountDistribution(count=3, weight=0.2) |
|
|
), |
|
|
availability_count_distribution=( |
|
|
CountDistribution(count=5, weight=4), |
|
|
CountDistribution(count=10, weight=3), |
|
|
CountDistribution(count=15, weight=2), |
|
|
CountDistribution(count=20, weight=1) |
|
|
), |
|
|
random_seed=37 |
|
|
) |
|
|
} |
|
|
|
|
|
|
|
|
FIRST_NAMES = ("Amy", "Beth", "Carl", "Dan", "Elsa", "Flo", "Gus", "Hugo", "Ivy", "Jay") |
|
|
LAST_NAMES = ("Cole", "Fox", "Green", "Jones", "King", "Li", "Poe", "Rye", "Smith", "Watt") |
|
|
SHIFT_LENGTH = timedelta(hours=8) |
|
|
MORNING_SHIFT_START_TIME = time(hour=6, minute=0) |
|
|
DAY_SHIFT_START_TIME = time(hour=9, minute=0) |
|
|
AFTERNOON_SHIFT_START_TIME = time(hour=14, minute=0) |
|
|
NIGHT_SHIFT_START_TIME = time(hour=22, minute=0) |
|
|
|
|
|
SHIFT_START_TIMES_COMBOS = ( |
|
|
(MORNING_SHIFT_START_TIME, AFTERNOON_SHIFT_START_TIME), |
|
|
(MORNING_SHIFT_START_TIME, AFTERNOON_SHIFT_START_TIME, NIGHT_SHIFT_START_TIME), |
|
|
(MORNING_SHIFT_START_TIME, DAY_SHIFT_START_TIME, AFTERNOON_SHIFT_START_TIME, NIGHT_SHIFT_START_TIME), |
|
|
) |
|
|
|
|
|
|
|
|
location_to_shift_start_time_list_map = dict() |
|
|
|
|
|
|
|
|
def earliest_monday_on_or_after(target_date: date): |
|
|
""" |
|
|
Returns the date of the next given weekday after |
|
|
the given date. For example, the date of next Monday. |
|
|
|
|
|
NB: if it IS the day we're looking for, this returns 0. |
|
|
consider then doing onDay(foo, day + 1). |
|
|
""" |
|
|
days = (7 - target_date.weekday()) % 7 |
|
|
return target_date + timedelta(days=days) |
|
|
|
|
|
|
|
|
def generate_demo_data(demo_data_or_parameters: DemoData | DemoDataParameters) -> EmployeeSchedule: |
|
|
global location_to_shift_start_time_list_map, demo_data_to_parameters |
|
|
if isinstance(demo_data_or_parameters, DemoData): |
|
|
parameters = demo_data_to_parameters[demo_data_or_parameters] |
|
|
else: |
|
|
parameters = demo_data_or_parameters |
|
|
|
|
|
start_date = earliest_monday_on_or_after(date.today()) |
|
|
random = Random(parameters.random_seed) |
|
|
shift_template_index = 0 |
|
|
for location in parameters.locations: |
|
|
location_to_shift_start_time_list_map[location] = SHIFT_START_TIMES_COMBOS[shift_template_index] |
|
|
shift_template_index = (shift_template_index + 1) % len(SHIFT_START_TIMES_COMBOS) |
|
|
|
|
|
name_permutations = [f'{first_name} {last_name}' |
|
|
for first_name, last_name in product(FIRST_NAMES, LAST_NAMES)] |
|
|
random.shuffle(name_permutations) |
|
|
|
|
|
employees = [] |
|
|
for i in range(parameters.employee_count): |
|
|
count, = random.choices(population=counts(parameters.optional_skill_distribution), |
|
|
weights=weights(parameters.optional_skill_distribution)) |
|
|
skills = [] |
|
|
skills += random.sample(parameters.optional_skills, count) |
|
|
skills += random.sample(parameters.required_skills, 1) |
|
|
employees.append( |
|
|
Employee(name=name_permutations[i], |
|
|
skills=set(skills)) |
|
|
) |
|
|
|
|
|
shifts: list[Shift] = [] |
|
|
|
|
|
def id_generator(): |
|
|
current_id = 0 |
|
|
while True: |
|
|
yield str(current_id) |
|
|
current_id += 1 |
|
|
|
|
|
ids = id_generator() |
|
|
|
|
|
for i in range(parameters.days_in_schedule): |
|
|
count, = random.choices(population=counts(parameters.availability_count_distribution), |
|
|
weights=weights(parameters.availability_count_distribution)) |
|
|
employees_with_availabilities_on_day = random.sample(employees, count) |
|
|
current_date = start_date + timedelta(days=i) |
|
|
for employee in employees_with_availabilities_on_day: |
|
|
rand_num = random.randint(0, 2) |
|
|
if rand_num == 0: |
|
|
employee.unavailable_dates.add(current_date) |
|
|
elif rand_num == 1: |
|
|
employee.undesired_dates.add(current_date) |
|
|
elif rand_num == 2: |
|
|
employee.desired_dates.add(current_date) |
|
|
shifts += generate_shifts_for_day(parameters, current_date, random, ids) |
|
|
|
|
|
shift_count = 0 |
|
|
for shift in shifts: |
|
|
shift.id = str(shift_count) |
|
|
shift_count += 1 |
|
|
|
|
|
return EmployeeSchedule( |
|
|
employees=employees, |
|
|
shifts=shifts |
|
|
) |
|
|
|
|
|
|
|
|
def generate_shifts_for_day(parameters: DemoDataParameters, current_date: date, random: Random, |
|
|
ids: Generator[str, any, any]) -> list[Shift]: |
|
|
global location_to_shift_start_time_list_map |
|
|
shifts = [] |
|
|
for location in parameters.locations: |
|
|
shift_start_times = location_to_shift_start_time_list_map[location] |
|
|
for start_time in shift_start_times: |
|
|
shift_start_date_time = datetime.combine(current_date, start_time) |
|
|
shift_end_date_time = shift_start_date_time + SHIFT_LENGTH |
|
|
shifts += generate_shifts_for_timeslot(parameters, shift_start_date_time, shift_end_date_time, |
|
|
location, random, ids) |
|
|
|
|
|
return shifts |
|
|
|
|
|
|
|
|
def generate_shifts_for_timeslot(parameters: DemoDataParameters, timeslot_start: datetime, timeslot_end: datetime, |
|
|
location: str, random: Random, ids: Generator[str, any, any]) -> list[Shift]: |
|
|
shift_count, = random.choices(population=counts(parameters.shift_count_distribution), |
|
|
weights=weights(parameters.shift_count_distribution)) |
|
|
|
|
|
shifts = [] |
|
|
for i in range(shift_count): |
|
|
if random.random() >= 0.5: |
|
|
required_skill = random.choice(parameters.required_skills) |
|
|
else: |
|
|
required_skill = random.choice(parameters.optional_skills) |
|
|
shifts.append(Shift( |
|
|
id=next(ids), |
|
|
start=timeslot_start, |
|
|
end=timeslot_end, |
|
|
location=location, |
|
|
required_skill=required_skill)) |
|
|
|
|
|
return shifts |
|
|
|
