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, ) 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 if date < datetime.date(2026, 2, 16): 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 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 if week_start < datetime.date(2026, 2, 16): skipped_rows += 1 continue # skip rows before rota start 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-skipped_rows, # 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/CONSULTANT TWILIGHT & ONCALL ROTA.ods" # 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 = "cons/Rota B Feb -July 2026.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].upper() 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 # --- 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) 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 = {} rota_data[name] = f"rota {str(row.iloc[2]).strip().lower()}" for i, day in enumerate(weekday_cols): val = row.iloc[3 + 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 = "2026-02-16" # 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, previous_shifts=priors_map.get(initial, {}), ) if initial == "LB": w.start_date = datetime.date(2026, 6, 15) #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",): w.add_hard_day_dependency({"Fri", "Sat", "Sun"}, ignore_dates=[datetime.date(2026, 5, 2)]) else: w.add_hard_day_dependency({"Fri", "Sun"}) w.add_hard_day_exclusion({"Sat", "Sun"}) #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 = "2026-02-16", 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, use_previous_shifts=True, name="cons rota feb 2026 run 5", 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=[(4, 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/"] ) if suspend_on_finish: os.system("systemctl suspend") if __name__ == "__main__": app()