from collections import defaultdict import datetime import os import sys import time from rota.shifts import MaxShiftsPerWeekConstraint, NoWorkers, PostShiftConstraint, PreShiftConstraint, RotaBuilder, SingleShift, WorkerRequirement, days import typer import subprocess import pandas as pd from rota.workers import ( Worker, NotAvailableToWork, NonWorkingDays, WorkRequests, PreferenceNotToWork, OutOfProgramme, ) from loguru import logger 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, }) 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 = [] 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 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": 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 = "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, ) #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 test", allow_force_assignment_with_leave_conflict=True, ) # 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.constraint_options["max_days_per_week_block"] = [(4, 2), (6, 3)] # (max_days, week_block) #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")), PostShiftConstraint(days=1, start_date="2025-11-17", exclude_days=("Fri", "Sat", "Sun"), allow_self=True), #{ # "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, constraints=[ #{"name": "pre", "options": 2}, #{"name": "post", "options": 1}, 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", workers_required=1, length=12.5, days=days[5:], balance_offset=2, #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, constraints=[ #PreShiftConstraint(days=1, start_date="2025-11-17", ignore_shifts=["oncall"]), PreShiftConstraint(days=1, start_date="2025-11-17", ignore_shifts=["oncall"], exclude_days=("Sat", "Sun")), #PostShiftConstraint(days=1, start_date="2025-11-17", ignore_shifts=["oncall", "weekend a", "weekend b"], include_weekends=False), MaxShiftsPerWeekConstraint(max_shifts=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, 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/"] ) if suspend_on_finish: os.system("systemctl suspend") if __name__ == "__main__": app()