Files
proc-rota/gen_cons.py
T
2026-05-20 21:12:42 +01:00

707 lines
25 KiB
Python

from pyomo.opt.results.container import ignore
from collections import defaultdict
import datetime
import os
import sys
import time
from rota_generator.shifts import MaxShiftsPerWeekConstraint, NoWorkers, PostShiftConstraint, PreShiftConstraint, RotaBuilder, SingleShift, WorkerRequirement, days, RotaConstraintOptions, MaxShiftsPerWeekBlockConstraint
import typer
import subprocess
import pandas as pd
from rota_generator.workers import (
Worker,
NotAvailableToWork,
NonWorkingDays,
WorkRequests,
PreferenceNotToWork,
OutOfProgramme, MaxUniqueShiftsPerWeekBlockConstraint,
)
from loguru import logger
app = typer.Typer()
ROTA_REQUESTS = "/home/ross/Sync/cons/requests.ods"
ROTA_B_PATH = "/home/ross/Sync/cons/rota_b.xlsx"
ROTA_START_DATE = "2026-08-03"
sites = ("rota a", "rota b", "rota c")
def _parse_sheet_date(value):
if pd.isna(value):
return None
parsed = pd.to_datetime(value, errors="coerce")
if pd.isna(parsed):
return None
return parsed.date()
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 = _parse_sheet_date(row.iloc[1])
if date is None:
continue
if date < datetime.date.fromisoformat(ROTA_START_DATE):
continue # skip rows before rota start 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,
})
elif leave_val != "" and leave_val != "nan":
leave_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 = []
skipped_rows = 0
rota_start = datetime.date.fromisoformat(ROTA_START_DATE)
for idx, row in rota_df.iloc[2:].iterrows():
week_start = _parse_sheet_date(row.iloc[0])
if week_start is None:
print(f"Invalid date format in row {idx}: {row.iloc[0]}")
continue # skip rows with invalid date
if week_start < rota_start:
skipped_rows += 1
continue # skip rows before rota start date
week = ((week_start - rota_start).days // 7) + 1
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 = 2 + i * 2 + j # 2 for Mon twilight, 3 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 = 12 + (i - 5) # 12 for Sat, 13 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 weekend:
initial = initial.split("/")[0].strip() # In case of "initial1/initial2", take the first one
if initial and initial.lower() != "nan" and initial.lower() != "nat":
if weekend:
shift = "weekend"
assignments.append({
"week_start": week_start,
"week": week,
"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 = ROTA_REQUESTS
# Read all sheets
xls = pd.read_excel(ods_path, engine="odf", sheet_name=None)
# Extract sheets
days_df = xls["Days"]
calender_df = xls["Calendar"]
rota_df = xls["Rota"]
rota_b_path = ROTA_B_PATH
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():
logger.debug(f"{idx}: {row}")
date_val = row.iloc[0]
try:
worker_val = row.iloc[1].upper() if len(row) > 1 else None
except AttributeError:
# end of file
break
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
# --- Load prior allocations (targets) from CSV ---
priors_path = "cons/sep2025priors.csv"
priors_map = {}
try:
priors_df = pd.read_csv(priors_path)
for _, prow in priors_df.iterrows():
initial = str(prow.get("Worker", "")).strip()
if not initial or initial.lower() == "nan":
continue
prev = {}
# expected columns: oncall, twilight, weekend
for col in ("oncall", "twilight", "weekend"):
raw = prow.get(col, "")
if pd.isna(raw) or raw == "" or str(raw).strip() in ("-", "nan"):
continue
s = str(raw).strip()
# Expect formats like: '4 (2.50)' or '3 (2.86)'
worked = 0.0
allocated = 0.0
if "(" in s:
try:
worked_part = s.split("(")[0].strip()
worked = float(worked_part)
except Exception:
worked = 0.0
try:
inner = s.split("(", 1)[1].split(")", 1)[0]
allocated = float(inner.strip())
except Exception:
allocated = worked
else:
try:
allocated = float(s)
except Exception:
continue
worked = 0.0
# store as tuple (worked, allocated) as expected by RotaBuilder
prev[col] = (worked, allocated)
if prev:
priors_map[initial] = prev
logger.debug("priors_map: {}", priors_map)
except FileNotFoundError:
print(f"Prior allocations file not found: {priors_path}")
except Exception as e:
print(f"Error reading prior allocations: {e}")
# --- 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 = []
rota_data = {}
for idx, row in days_df.iloc[7:].iterrows():
# Extract initials from the name column (assumed to be in brackets at the end)
name = str(row.iloc[1]).strip()
initial = str(row.iloc[2]).strip()
if not name or name.lower() == "nan":
continue
availability = {}
requests = {}
rota_data[name] = f"rota {str(row.iloc[3]).strip().lower()}"
for i, day in enumerate(weekday_cols):
val = row.iloc[4 + 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() or str(val).strip().lower() == "yes*")
availability[day] = available
workers_days.append({"initial": initial, "name": name, "availability": availability, "requests": requests})
leave_requests, work_requests = extract_leave_and_rota_from_calender(calender_df)
assignments, assignments_by_worker = extract_shift_assignments_from_rota(rota_df)
print(assignments)
# --- Load per-day prior worked data (Sat/Sun) from sep2025priordays.csv for Rota A workers ---
# Build set of initials who are on Rota A
rota_a_initials = set()
for wd in workers_days:
name = wd["name"]
initial = wd["initial"]
if rota_data.get(name, "") == "rota a":
rota_a_initials.add(initial)
priordays_path = "cons/sep2025priordays.csv"
try:
import csv
perday_counts = defaultdict(lambda: defaultdict(int))
with open(priordays_path, newline="", encoding="utf-8") as f:
reader = csv.reader(f)
for row in reader:
if len(row) < 4:
continue
day = row[2].strip()
if day not in ("Sat", "Sun"):
continue
# fields from index 3 onwards are assignments; some may be empty
for field in row[3:]:
if not field:
continue
val = field.strip()
if not val:
continue
# entries may be like: RK (oncall, weekend) or CK (weekend b)
import re
m = re.match(r"\s*([A-Za-z0-9]+)\s*\(([^)]*)\)", val)
if not m:
# sometimes multiple assignments are concatenated in a single cell separated by commas
parts = [p.strip() for p in re.split(r",\s*", val) if p.strip()]
for p in parts:
m2 = re.match(r"\s*([A-Za-z0-9]+)\s*\(([^)]*)\)", p)
if m2:
initial = m2.group(1).strip()
paren = m2.group(2).lower()
if initial in rota_a_initials and "weekend" in paren:
perday_counts[initial][day] += 1
continue
initial = m.group(1).strip()
paren = m.group(2).lower()
if initial in rota_a_initials and "weekend" in paren:
perday_counts[initial][day] += 1
# Merge per-day counts into priors_map under 'weekend' per-day keys
for initial, daymap in perday_counts.items():
prev = priors_map.get(initial, {})
existing = prev.get("weekend")
# If existing is a tuple, convert to dict preserving overall tuple under '__all__'
if existing and not isinstance(existing, dict):
prev["weekend"] = {"__all__": existing}
existing = prev["weekend"]
for day, worked_count in daymap.items():
# Determine allocated value for this per-day entry
allocated_day = worked_count
if existing:
# prefer an explicit per-day allocated if present
if isinstance(existing, dict) and day in existing:
try:
allocated_day = float(existing[day][1])
except Exception:
allocated_day = worked_count
elif isinstance(existing, dict) and "__all__" in existing:
try:
overall_alloc = float(existing["__all__"][1])
allocated_day = overall_alloc / 2.0
except Exception:
allocated_day = worked_count
elif isinstance(existing, tuple):
try:
allocated_day = float(existing[1]) / 2.0
except Exception:
allocated_day = worked_count
# ensure structure
if "weekend" not in prev or not isinstance(prev["weekend"], dict):
prev.setdefault("weekend", {})
prev["weekend"][day] = (float(worked_count), float(allocated_day))
if prev:
priors_map[initial] = prev
logger.debug("priors_map after per-day merge: {}", priors_map)
except FileNotFoundError:
print(f"Per-day prior file not found: {priordays_path}")
except Exception as e:
print(f"Error reading per-day priors: {e}")
workers = []
for worker_day in workers_days:
worker = worker_day["name"]
initial = worker_day["initial"]
start_date = ROTA_START_DATE
#if initial == "ND":
# start_date = "2025-09-29" # Non-working day worker starts later
#if initial == "MF":
# start_date = "2026-01-15" # 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-17" # 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,
previous_shifts=priors_map.get(initial, {}),
max_days_worked_per_week=2
)
#if initial == "L1":
# w.oop = [OutOfProgramme(
# start_date="2025-09-15",
# end_date="2025-10-14",
# reason="even out shifts",
# )]
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","RG"):
w.add_hard_day_dependency({"Fri", "Sat", "Sun"}, ignore_dates=[datetime.date(2026, 5, 2)])
w.max_days_worked_per_week = 3
else:
w.add_hard_day_dependency({"Fri", "Sun"})
w.add_hard_day_exclusion({"Sat", "Sun"})
else:
w.max_unique_shifts_per_week_block = [MaxUniqueShiftsPerWeekBlockConstraint(week_block=1, max_unique_shifts=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.1,
start_date: datetime.datetime = ROTA_START_DATE,
weeks: int = 22,
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,
use_previous_shifts=True,
name="cons rota 2026 july",
allow_force_assignment_with_leave_conflict=True,
constraint_options=RotaConstraintOptions(
balance_weekends=True,
max_shifts_per_week=6,
max_weekend_frequency=3,
max_days_per_week_block=[(3, 2), (4, 4)],
#balance_weekend_days=True,
balance_weekend_days_quadratic=True,
weekend_day_hard_max_deviation=2,
balance_weekend_days_use_previous_shifts=True,
balance_weekend_days_weight=500000
),
)
#Rota.enable_unavailabilities = False
Rota.add_shifts(
SingleShift(
sites=("rota a",),
name="oncall",
length=12.5,
days=days,
balance_offset=2,
assign_as_block=False,
constraints=[
#PreShiftConstraint(days=2, start_date="2025-11-17", exclude_days=("Fri", "Sat", "Sun"), exclude_dates=["2026-04-06", "2026-05-04", "2026-05-25"]),
#PreShiftConstraint(days=1, start_date="2025-11-17", exclude_days=("Fri", "Sat", "Sun"), allow_self=False),
MaxShiftsPerWeekConstraint(max_shifts=1, days=days[:5]),
],
),
SingleShift(
sites=("rota a",),
name="weekend",
workers_required=1,
length=12.5,
days=days[5:],
balance_offset=1,
constraints=[
PostShiftConstraint(days=1, start_date="2025-11-17", ignore_shifts=["oncall"]),
],
display_char="a",
#force_assign_with=["oncall"]
),
SingleShift(
sites=("rota b", "rota c"),
name="weekend b",
#end_date="2026-07-20",
workers_required=1,
length=12.5,
days=days[5:],
balance_offset=2,
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,
constraints=[
PreShiftConstraint(days=1, start_date="2025-11-17", ignore_shifts=["oncall"], exclude_days=("Sat", "Sun")),
MaxShiftsPerWeekConstraint(max_shifts=1),
MaxShiftsPerWeekBlockConstraint(week_block=3,max_shifts=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, export_with_timestamp=False, 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/"]
)
subprocess.run(
["scp", "output/timetable.js", "ross@46.101.13.46:proc/proc-rota/output/"]
)
if suspend_on_finish:
os.system("systemctl suspend")
if __name__ == "__main__":
app()