Files
proc-rota/gen_cons.py
T

523 lines
18 KiB
Python

from collections import defaultdict
import datetime
import os
import sys
import time
from rota.shifts import NoWorkers, RotaBuilder, SingleShift, WorkerRequirement, days
import typer
import subprocess
import pandas as pd
from rota.workers import (
Worker,
NotAvailableToWork,
NonWorkingDays,
WorkRequests,
PreferenceNotToWork,
OutOfProgramme,
)
app = typer.Typer()
sites = ("rota a", "rota b", "rota c")
def extract_leave_and_rota_from_calender(calender_df):
"""
Extract leave and rota assignments from calender_df.
- The first row contains worker names, starting from the 5th column (index 4).
- The second column (index 1) contains dates in dd/mm/yyyy format.
- The second column (index 1) also contains the rota/leave code for that date.
Returns a list of dicts: {"date": ..., "worker": ..., "rota": ...}
"""
# Get worker names from the first row, starting from column 5 (index 4)
worker_names = [str(x).strip() for x in calender_df.columns[5:]]
# Iterate over each row (skip the header row)
# Get rota data from the second row (index 1), starting from column 5 (index 4)
# Extract the rota data for each worker from the first data row (index 0), columns 5 onwards
rota_row = calender_df.iloc[0]
rota_data = {}
for i, worker in enumerate(worker_names):
assignment = str(rota_row.iloc[5 + i]).strip() if len(rota_row) > 5 + i else ""
rota_data[worker] = f"rota {assignment.lower()}"
leave_requests = []
work_requests = []
# Process leave data from row 3 onwards (index 2 and up)
for idx, row in calender_df.iloc[3:].iterrows():
date_str = str(row.iloc[1]).strip().split(" ")[0] # Get the date string from the second column (index 1)
if not date_str or date_str.lower() == "nan":
continue
try:
date = datetime.datetime.strptime(date_str, "%Y-%m-%d").date()
except Exception:
continue # skip rows with invalid date
for i, worker in enumerate(worker_names):
if not worker or worker.lower() == "nan":
continue
leave_val = str(row.iloc[5 + i]).strip().lower() if len(row) > 5 + i else ""
if "no" in leave_val:
leave_requests.append({
"date": date,
"worker": worker,
})
elif "yes" in leave_val:
work_requests.append({
"date": date,
"worker": worker,
})
#print(leave_requests)
return leave_requests, work_requests, rota_data
def extract_shift_assignments_from_rota(rota_df):
"""
Extract shift assignments from rota_df.
- Data starts from the 3rd row (index 2).
- The first column is the week start date.
- For each day (Mon-Sun), two columns: twilight initials, night initials.
Returns a list of dicts: {"week_start": ..., "day": ..., "shift": ..., "initial": ...}
"""
days = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]
shifts = ["twilight", "oncall"]
assignments = []
for idx, row in rota_df.iloc[2:].iterrows():
week_start_str = str(row.iloc[0]).strip().split(" ")[0] # Get the week start date from the first column (index 0)
try:
week_start = datetime.datetime.strptime(week_start_str, "%Y-%m-%d").date()
except Exception:
print(f"Invalid date format in row {idx}: {week_start_str}")
continue # skip rows with invalid date
for i, day in enumerate(days):
for j, shift in enumerate(shifts):
if i < 5:
# Mon-Fri: two columns per day (twilight, oncall)
col_idx = 1 + i * 2 + j # 1 for Mon twilight, 2 for Mon oncall, etc.
weekend = False
else:
# Sat/Sun: only one column for "weekend" shift (twilight), skip "night"
weekend = True
if j == 0:
col_idx = 11 + (i - 5) # 11 for Sat, 12 for Sun
else:
continue # No night shift column for Sat/Sun
if col_idx >= len(row):
continue
initial = str(row.iloc[col_idx]).strip()
if initial and initial.lower() != "nan" and initial.lower() != "nat":
if weekend:
shift = "weekend"
assignments.append({
"week_start": week_start,
"week": idx-1, # 1-indexed week number
"day": day,
"date": week_start + datetime.timedelta(days=i),
"shift": shift,
"initial": initial,
})
# Build a mapping of assignments by worker initial
assignments_by_worker = {}
for assignment in assignments:
initial = assignment["initial"]
if initial not in assignments_by_worker:
assignments_by_worker[initial] = []
assignments_by_worker[initial].append(assignment)
return assignments, assignments_by_worker
def load_workers():
# Path to the ODS file
ods_path = "cons/Oncall rota requests + New rota (first draft) - September 2025.ods"
# Read all sheets
xls = pd.read_excel(ods_path, engine="odf", sheet_name=None)
# Extract sheets
days_df = xls["Days"]
calender_df = xls["Calender"]
rota_df = xls["Rota"]
rota_b_path = "cons/Draft Rota B Sept 25-Feb 26.xlsx"
rota_b_df = pd.read_excel(rota_b_path, engine="openpyxl", sheet_name="Sheet1", header=None)
# Extract rota_b assignments: date in column A, worker in column B
rota_b_assignments = defaultdict(list)
for idx, row in rota_b_df.iterrows():
date_val = row.iloc[0]
worker_val = row.iloc[1] if len(row) > 1 else None
if pd.isna(date_val) or pd.isna(worker_val):
continue
try:
date = pd.to_datetime(date_val).date()
except Exception:
continue
worker = str(worker_val).strip()
if not worker or worker.lower() == "nan":
continue
rota_b_assignments[worker].append(date)
print(rota_b_assignments)
# You can now use rota_b_assignments as needed
# --- Example: Print the first few rows of each sheet ---
#print("Days sheet:")
#print(days_df.head())
weekday_cols = ["Mon", "Tue", "Wed", "Thu", "Fri"]
workers_days = []
for idx, row in days_df.iloc[7:].iterrows():
# Extract initials from the name column (assumed to be in brackets at the end)
text = str(row.iloc[1]).strip()
if "(" in text and ")" in text:
initial = text.split("(")[-1].split(")")[0].strip()
else:
raise ValueError(f"Invalid format in row {idx}: {text}")
name = text.split("(")[0].strip()
if not name or name.lower() == "nan":
continue
availability = {}
requests = {}
for i, day in enumerate(weekday_cols):
val = row.iloc[2 + i]
# You can adjust the logic below depending on your marking scheme (e.g. "Y", "Yes", "1", etc.)
available = not ("no" in str(val).strip().lower())
availability[day] = available
workers_days.append({"initial": initial, "name": name, "availability": availability, "requests": requests})
leave_requests, work_requests, rota_data = extract_leave_and_rota_from_calender(calender_df)
assignments, assignments_by_worker = extract_shift_assignments_from_rota(rota_df)
print(assignments)
workers = []
for worker_day in workers_days:
worker = worker_day["name"]
initial = worker_day["initial"]
start_date = "2025-09-01" # Default start date, can be adjusted later
if initial == "ND":
start_date = "2025-09-29" # Non-working day worker starts later
if initial == "MF":
start_date = "2025-10-06" # Medical leave worker starts earlier
if initial == "AA":
start_date = "2025-11-17" # Absent worker starts later
end_date = None
if initial == "L1":
end_date = "2025-11-06" # L1 worker has a leave request
elif initial == "L2":
end_date = "2025-09-29" # L2 worker has a leave request
non_working_days = []
for day, available in worker_day["availability"].items():
if not available:
non_working_days.append(NonWorkingDays(day=day, start_date=datetime.datetime(2025, 11, 17)))
w = Worker(
name=worker_day["initial"],
site=rota_data[worker],
grade=1,
start_date=start_date,
end_date=end_date,
# Add leave requests for this worker
not_available_to_work=[
NotAvailableToWork(
date=req["date"],
reason="leave",
) for req in leave_requests if req["worker"] == worker
],
work_requests=[
WorkRequests(
date=req["date"],
shift="*"
) for req in work_requests if req["worker"] == worker
],
nwds=non_working_days,
)
if initial == "L1":
w.oop = [OutOfProgramme(
start_date="2025-09-15",
end_date="2025-10-14",
reason="even out shifts",
)]
if initial == "ND":
print(w.not_available_to_work)
if initial in assignments_by_worker:
# Add shift assignments for this worker
#if initial in ("RK", "HB", "DS", "L1", "L2", "L3", "GFM", "TS", "AB", "IR", "MF", "ND", "RG", "AA", "BR", "MS", "HA", "COD", "JL", "TB", "NHO"):
for assignment in assignments_by_worker[initial]:
#if initial == "NHO":
# pass
# #if assignment["shift"] == "oncall":
# # continue
# #if assignment["shift"] == "twilight":
# # continue
# #1if assignment["shift"] == "weekend":
# #1 continue
#print(assignment)
w.force_assign_shift(
week=assignment["week"],
day=assignment["day"],
shift_name=assignment["shift"],
)
if initial in rota_b_assignments:
for date in rota_b_assignments.get(initial, []):
w.force_assign_shift_by_date(
date=date,
shift_name="weekend b",
)
if w.site == "rota a":
w.prefer_multi_shift_together = 10
w.allow_shifts_together("oncall", "twilight")
#w.add_force_assign_with("twilight", "oncall")
w.add_force_assign_with("weekend", "oncall")
if w.name in ("DS",):
w.add_hard_day_dependency({"Fri", "Sat", "Sun"}, start_date="2025-11-17")
else:
w.add_hard_day_dependency({"Fri", "Sun"}, start_date="2025-11-17")
w.add_hard_day_exclusion({"Sat", "Sun"}, start_date="2025-11-17")
#if w.name == "TB":
# w.add_hard_day_exclusion({"Sat", "Sun"}, start_date="2025-11-17")
#w.set_max_shifts
#w.set_max_shifts_per_week("weekend", 1)
workers.append(w)
return workers
@app.command()
def main(
suspend: bool = False,
solve: bool = True,
time_to_run: int = 60 * 60,
ratio: float = 0.001,
start_date: datetime.datetime = "2025-09-01",
weeks: int = 24,
bom: int = 1,
):
rota_start_date = start_date.date()
suspend_on_finish = suspend
Rota = RotaBuilder(
rota_start_date,
weeks_to_rota=weeks,
balance_offset_modifier=bom,
name="cons rota",
)
# Rota = RotaBuilder(start_date, weeks_to_rota=20, balance_offset_modifier=1)
Rota.constraint_options["balance_weekends"] = True
#Rota.constraint_options["max_weekend_frequency"] = 2
Rota.constraint_options["max_shifts_per_week"] = 6
Rota.constraint_options["max_days_worked_per_week"] = 3
# Rota.constraint_options["avoid_st2_first_month"] = True
Rota.add_shifts(
SingleShift(
sites=("rota a",),
name="oncall",
length=12.5,
days=days,
balance_offset=2,
assign_as_block=False,
constraint=[
#{
# "name": "max_shifts_per_week",
# "options": 3,
#}
],
),
SingleShift(
sites=("rota a",),
name="weekend",
workers_required=1,
length=12.5,
days=days[5:],
balance_offset=1,
constraint=[
#{"name": "pre", "options": 2},
{"name": "post", "options": 1},
],
display_char="a",
#force_assign_with=["oncall"]
),
SingleShift(
sites=("rota b", "rota c"),
name="weekend b",
workers_required=1,
length=12.5,
days=days[5:],
balance_offset=1,
#constraint=[
# {"name": "pre", "options": 1},
# {"name": "post", "options": 1},
#],
display_char="b",
),
SingleShift(
sites=(
"rota a",
"rota b",
"rota d"
),
name="twilight",
workers_required=1,
length=12.5,
days=days[:5],
balance_offset=1,
constraint=[
#{"name": "pre", "options": 1},
#{"name": "post", "options": 1},
],
),
)
#Rota.allow_shifts_together_for_all_workers(
# "oncall",
# "twilight",
#)
#Rota.allow_shifts_together_for_all_workers(
# "oncall",
# "twilight",
#)
# Rota.add_grade_constraint_by_week([2], [1, 2], ["night_weekday", "night_weekend"])
load_leave = True
Rota.build_shifts()
# Rota.add_grade_constraint_by_week([2], [1], Rota.get_shift_names())
# Build workers
workers = load_workers()
Rota.add_workers(workers)
#for n in range(1, 12):
# w = Worker(name=f"A{n:02d}", site="rota a", grade=1)
# #w.prefer_multi_shift_together = 10
# #w.allow_shifts_together("oncall", "twilight")
# #w.allow_shifts_together("oncall", "weekend")
# #w.add_hard_day_dependency("Fri", "Sat", "Sun")
# w.add_force_assign_with("twilight", "oncall")
# w.add_force_assign_with("weekend", "oncall")
# w.set_max_shifts_per_week("twilight", 1)
# if n <= 2:
# pass
# w.add_hard_day_dependency("Fri", "Sat", "Sun")
# else:
# w.add_hard_day_dependency("Fri", "Sun")
# #w.set_max_shifts_per_week("weekend", 1)
# # #w.add_hard_day_dependency("Sat", "Sun")
# print(w)
# Rota.add_worker(w)
#w13 = Worker(name=f"A13", site="rota a", grade=1, start_date="2025-10-01")
##w.prefer_multi_shift_together = 10
##w.allow_shifts_together("oncall", "twilight")
##w.allow_shifts_together("oncall", "weekend")
##w.add_hard_day_dependency("Fri", "Sat", "Sun")
#w13.add_force_assign_with("twilight", "oncall")
#w13.add_force_assign_with("weekend", "oncall")
##w.set_max_shifts_per_week("twilight", 1)
#Rota.add_worker(w13)
#w13 = Worker(name=f"A14", site="rota a", grade=1, start_date="2025-12-01")
##w.prefer_multi_shift_together = 10
##w.allow_shifts_together("oncall", "twilight")
##w.allow_shifts_together("oncall", "weekend")
##w.add_hard_day_dependency("Fri", "Sat", "Sun")
#w13.add_force_assign_with("twilight", "oncall")
#w13.add_force_assign_with("weekend", "oncall")
##w.set_max_shifts_per_week("twilight", 1)
#Rota.add_worker(w13)
#for w in range(14, 24):
# w = Worker(
# name=f"B{w}",
# site="rota b",
# grade=1,
# nwds=[
# NonWorkingDays(
# day="Fri",
# )
# ],
# )
# #w.add_hard_day_exclusion("Sat", "Sun")
# w.set_max_shifts_per_week("twilight", 2)
# w.set_max_shifts_per_week("weekend b", 1)
# print(w)
# Rota.add_worker(w)
# Rota.build_workers()
# Rota.build_model()
solver_options = {"ratio": ratio, "seconds": time_to_run, "threads": 10}
# solver_options = {"seconds": time_to_run, "threads": 10}
# start_time = time.time()
Rota.build_and_solve(solver_options, export=True, solve=solve, solver="appsi_highs")
# Rota.solve_shifts_by_block(solver_options, block_length=13)
# Rota.solve_shifts_individually(solver_options)
# Rota.solve_model(options=solver_options)
# end_time = time.time()
# print(f"Time taken {end_time-start_time}")
# print(Rota.get_worker_details())
# optimizer = SolverFactory('cbc')
# result = optimizer.solve(prob,tee=True)
# result.Solver.Status = SolverStatus.warning
# prob.solutions.load_from(result)
# ResultsHolder = RotaResults(Rota)
# worker_timetable_brief = ResultsHolder.get_worker_timetable_brief(
# show_prefs=False, show_unavailable=False
# )
#
# print(worker_timetable_brief)
# Rota.export_rota_to_html("proc_rota")
# Rota.export_rota_to_csv("rota")
subprocess.run(
["scp", Rota.exported_rota_file, "ross@46.101.13.46:proc/proc-rota/output/"]
)
if suspend_on_finish:
os.system("systemctl suspend")
if __name__ == "__main__":
app()