#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Enterprise Roster Generator ‚Final Adaptive Version
All constraints adapt dynamically based on team size to ensure feasibility.
"""
import itertools
import json
import os
import pickle
import tempfile
import threading
import time as pytime
from datetime import datetime, timedelta
from datetime import time as dt_time
from pathlib import Path
import pandas as pd
import schedule
import streamlit as st
# ------------------------------
# OR-Tools Weekly Optimiser
# ------------------------------
from ortools.sat.python import cp_model
def get_weekly_requirements(n_staff: int) -> tuple:
"""
Returns coverage requirements based on available staff count.
For n_staff == 5: reduces ONE weekday day shift to 2.
"""
if n_staff >= 6:
return (3, 1, 1, 1, [])
elif n_staff == 5:
# Reduce one weekday (Wednesday) to 2 day staff
return (3, 1, 1, 1, [2]) # Wednesday = index 2 (Mon=0)
else:
raise ValueError("At least 5 staff required.")
# def solve_week(
# week_idx: int,
# start_date: datetime.date,
# available_staff: list[str],
# cumulative_shifts: dict[str, int],
# all_staff_global: list[str],
# ) -> tuple[dict, dict]:
N_AVAIL = len(available_staff)
if N_AVAIL < 5:
raise ValueError("Minimum 5 staff per week.")
# Get coverage rules
wd_day, wd_night, we_day, we_night, reduced_days = get_weekly_requirements(N_AVAIL)
# Determine constraint strictness based on staff count
FULL_NO_CONSECUTIVE = N_AVAIL >= 8 # Only enforce full restriction with 8+ staff
FULL_48H_REST = N_AVAIL >= 8 # Only enforce full 48h rest with 8+ staff
full_names = available_staff + [f"Vacant_{i}" for i in range(9 - N_AVAIL)]
SHIFT = {"day": 0, "night": 1}
DAYS = 7
WEEKDAY_REL = {0, 1, 2, 3, 4}
model = cp_model.CpModel()
x = {}
for p, d, s in itertools.product(range(9), range(DAYS), range(2)):
x[p, d, s] = model.NewBoolVar(f"x_{p}_{d}_{s}")
# === Dynamic Coverage ===
for d in range(DAYS):
if d in WEEKDAY_REL:
day_req = wd_day
if d in reduced_days: # e.g., Wednesday
day_req = 2
night_req = wd_night
else:
day_req = we_day
night_req = we_night
model.Add(sum(x[p, d, SHIFT["day"]] for p in range(9)) == day_req)
model.Add(sum(x[p, d, SHIFT["night"]] for p in range(9)) == night_req)
# === No same-day double shift (always enforced) ===
for p, d in itertools.product(range(9), range(DAYS)):
model.Add(x[p, d, SHIFT["day"]] + x[p, d, SHIFT["night"]] <= 1)
# === No consecutive days (adaptive) ===
for p in range(9):
for d in range(DAYS - 1):
if FULL_NO_CONSECUTIVE:
# Full restriction: no shifts on consecutive days
model.Add(
x[p, d, SHIFT["day"]]
+ x[p, d, SHIFT["night"]]
+ x[p, d + 1, SHIFT["day"]]
+ x[p, d + 1, SHIFT["night"]]
<= 1
)
else:
# Relaxed: only prevent consecutive NIGHT shifts
model.Add(x[p, d, SHIFT["night"]] + x[p, d + 1, SHIFT["night"]] <= 1)
# === Weekend cap (always enforced) ===
for p in range(9):
model.Add(sum(x[p, d, s] for d in (5, 6) for s in range(2)) <= 1)
# === 48h rest after night shift (adaptive) ===
for p in range(9):
for d in range(DAYS):
night_d = x[p, d, SHIFT["night"]]
# Always enforce 24h rest (no shift next day after night)
if d + 1 < DAYS:
any_d1 = x[p, d + 1, SHIFT["day"]] + x[p, d + 1, SHIFT["night"]]
model.Add(any_d1 <= 1 - night_d)
# Only enforce second day off with larger teams
if FULL_48H_REST and d + 2 < DAYS:
any_d2 = x[p, d + 2, SHIFT["day"]] + x[p, d + 2, SHIFT["night"]]
model.Add(any_d2 <= 1 - night_d)
# === Vacants forced to 0 ===
for p in range(N_AVAIL, 9):
for d, s in itertools.product(range(DAYS), range(2)):
model.Add(x[p, d, s] == 0)
# === Weekly shift bounds (adaptive) ===
if N_AVAIL <= 6:
MIN_WEEKLY, MAX_WEEKLY = 3, 4 # Smaller teams work more shifts
elif N_AVAIL == 7:
MIN_WEEKLY, MAX_WEEKLY = 2, 4
else: # 8-9 staff
MIN_WEEKLY, MAX_WEEKLY = 2, 3
week_shifts = {}
for i, name in enumerate(available_staff):
var = model.NewIntVar(MIN_WEEKLY, MAX_WEEKLY, f"wshift_{i}")
model.Add(var == sum(x[i, d, s] for d in range(DAYS) for s in range(2)))
week_shifts[name] = var
# === Soft fairness objective ===
objective_terms = []
for i, name in enumerate(available_staff):
cum = cumulative_shifts.get(name, 0)
# Bias toward staff with fewer cumulative shifts
objective_terms.append(cum * week_shifts[name])
model.Minimize(sum(objective_terms))
solver = cp_model.CpSolver()
# Longer timeout for smaller teams with tighter constraints
solver.parameters.max_time_in_seconds = 45.0 if N_AVAIL <= 7 else 30.0
solver.parameters.num_search_workers = 6
if solver.Solve(model) not in (cp_model.OPTIMAL, cp_model.FEASIBLE):
# Create a detailed error message showing constraint violations
error_msg = f"Week {week_idx + 1} infeasible with {N_AVAIL} staff.\n"
error_msg += "Active constraints:\n"
error_msg += f"- No consecutive days: {'FULL' if FULL_NO_CONSECUTIVE else 'NIGHT ONLY'}\n"
error_msg += f"- Night rest: {'48h' if FULL_48H_REST else '24h'}\n"
error_msg += f"- Weekly bounds: {MIN_WEEKLY}-{MAX_WEEKLY} shifts\n"
error_msg += (
f"- Coverage: {wd_day}/{wd_night} weekdays, {we_day}/{we_night} weekends"
)
if reduced_days:
error_msg += (
f"\n- Reduced day(s): {', '.join(str(d) for d in reduced_days)}"
)
raise RuntimeError(error_msg)
schedule_week = {}
weekly_counts = {name: 0 for name in available_staff}
for d in range(7):
day_staff = [
full_names[p]
for p in range(9)
if solver.Value(x[p, d, SHIFT["day"]])
and not full_names[p].startswith("Vacant_")
]
night_staff = [
full_names[p]
for p in range(9)
if solver.Value(x[p, d, SHIFT["night"]])
and not full_names[p].startswith("Vacant_")
]
schedule_week[d] = {"day": day_staff, "night": night_staff}
for name in day_staff + night_staff:
weekly_counts[name] += 1
return schedule_week, weekly_counts
# def solve_week(
# week_idx: int,
# start_date: datetime.date,
# available_staff: list[str],
# cumulative_shifts: dict[str, int],
# all_staff_global: list[str],
# ) -> tuple[dict, dict]:
# N_AVAIL = len(available_staff)
# if N_AVAIL < 5:
# raise ValueError("Minimum 5 staff per week.")
# # Determine constraint strictness based on ACTUAL staff count
# # This is the critical fix - constraints must adapt based on team size
# FULL_NO_CONSECUTIVE = N_AVAIL >= 9 # Only full restriction with 9 staff
# FULL_48H_REST = N_AVAIL >= 9 # Only full 48h rest with 9 staff
# # Get coverage rules (still reduce for 5 staff)
# wd_day, wd_night, we_day, we_night, reduced_days = get_weekly_requirements(N_AVAIL)
# full_names = available_staff + [f"Vacant_{i}" for i in range(9 - N_AVAIL)]
# SHIFT = {"day": 0, "night": 1}
# DAYS = 7
# WEEKDAY_REL = {0, 1, 2, 3, 4}
# model = cp_model.CpModel()
# x = {}
# for p, d, s in itertools.product(range(9), range(DAYS), range(2)):
# x[p, d, s] = model.NewBoolVar(f"x_{p}_{d}_{s}")
# # === Dynamic Coverage ===
# for d in range(DAYS):
# if d in WEEKDAY_REL:
# day_req = wd_day
# if d in reduced_days:
# day_req = 2
# night_req = wd_night
# else:
# day_req = we_day
# night_req = we_night
# model.Add(sum(x[p, d, SHIFT["day"]] for p in range(9)) == day_req)
# model.Add(sum(x[p, d, SHIFT["night"]] for p in range(9)) == night_req)
# # === No same-day double shift (always enforced) ===
# for p, d in itertools.product(range(9), range(DAYS)):
# model.Add(x[p, d, SHIFT["day"]] + x[p, d, SHIFT["night"]] <= 1)
# # === CRITICAL FIX: Adaptive No consecutive days ===
# for p in range(9):
# for d in range(DAYS - 1):
# if FULL_NO_CONSECUTIVE:
# # Full restriction: no shifts on consecutive days
# model.Add(
# x[p, d, SHIFT["day"]]
# + x[p, d, SHIFT["night"]]
# + x[p, d + 1, SHIFT["day"]]
# + x[p, d + 1, SHIFT["night"]]
# <= 1
# )
# else:
# # Relaxed: only prevent consecutive NIGHT shifts
# model.Add(x[p, d, SHIFT["night"]] + x[p, d + 1, SHIFT["night"]] <= 1)
# # === Weekend cap (always enforced) ===
# for p in range(9):
# model.Add(sum(x[p, d, s] for d in (5, 6) for s in range(2)) <= 1)
# # === CRITICAL FIX: Adaptive 48h rest after night shift ===
# for p in range(9):
# for d in range(DAYS):
# night_d = x[p, d, SHIFT["night"]]
# # Always enforce 24h rest (no shift next day after night)
# if d + 1 < DAYS:
# any_d1 = x[p, d + 1, SHIFT["day"]] + x[p, d + 1, SHIFT["night"]]
# model.Add(any_d1 <= 1 - night_d)
# # Only enforce second day off with largest teams
# if FULL_48H_REST and d + 2 < DAYS:
# any_d2 = x[p, d + 2, SHIFT["day"]] + x[p, d + 2, SHIFT["night"]]
# model.Add(any_d2 <= 1 - night_d)
# # === Vacants forced to 0 ===
# for p in range(N_AVAIL, 9):
# for d, s in itertools.product(range(DAYS), range(2)):
# model.Add(x[p, d, s] == 0)
# # === CRITICAL FIX: Truly adaptive weekly bounds ===
# if N_AVAIL <= 6:
# MIN_WEEKLY, MAX_WEEKLY = 3, 4 # Small teams work more shifts
# elif N_AVAIL == 7:
# MIN_WEEKLY, MAX_WEEKLY = 2, 4 # Medium teams need flexibility
# elif N_AVAIL == 8:
# MIN_WEEKLY, MAX_WEEKLY = 2, 3 # Larger teams can have tighter bounds
# else: # N_AVAIL == 9
# MIN_WEEKLY, MAX_WEEKLY = 2, 3
# week_shifts = {}
# for i, name in enumerate(available_staff):
# var = model.NewIntVar(MIN_WEEKLY, MAX_WEEKLY, f"wshift_{i}")
# model.Add(var == sum(x[i, d, s] for d in range(DAYS) for s in range(2)))
# week_shifts[name] = var
# # === Soft fairness objective ===
# objective_terms = []
# for i, name in enumerate(available_staff):
# cum = cumulative_shifts.get(name, 0)
# objective_terms.append(cum * week_shifts[name])
# model.Minimize(sum(objective_terms))
# solver = cp_model.CpSolver()
# # Longer timeout for teams with tighter constraints
# solver.parameters.max_time_in_seconds = 45.0 if N_AVAIL <= 8 else 30.0
# solver.parameters.num_search_workers = 6
# if solver.Solve(model) not in (cp_model.OPTIMAL, cp_model.FEASIBLE):
# error_msg = f"Week {week_idx + 1} infeasible with {N_AVAIL} staff.\n"
# error_msg += "Active constraints:\n"
# error_msg += f"- No consecutive days: {'FULL' if FULL_NO_CONSECUTIVE else 'NIGHT ONLY'}\n"
# error_msg += f"- Night rest: {'48h' if FULL_48H_REST else '24h'}\n"
# error_msg += f"- Weekly bounds: {MIN_WEEKLY}-{MAX_WEEKLY} shifts\n"
# error_msg += (
# f"- Coverage: {wd_day}/{wd_night} weekdays, {we_day}/{we_night} weekends"
# )
# if reduced_days:
# error_msg += (
# f"\n- Reduced day(s): {', '.join(str(d) for d in reduced_days)}"
# )
# raise RuntimeError(error_msg)
# # [Rest of the function remains the same...]
# def solve_week(
# week_idx: int,
# start_date: datetime.date,
# available_staff: list[str],
# cumulative_shifts: dict[str, int],
# all_staff_global: list[str],
# ) -> tuple[dict, dict]:
# N_AVAIL = len(available_staff)
# if N_AVAIL < 5:
# raise ValueError("Minimum 5 staff per week.")
# Get coverage rules
wd_day, wd_night, we_day, we_night, reduced_days = get_weekly_requirements(N_AVAIL)
full_names = available_staff + [f"Vacant_{i}" for i in range(9 - N_AVAIL)]
SHIFT = {"day": 0, "night": 1}
DAYS = 7
WEEKDAY_REL = {0, 1, 2, 3, 4}
model = cp_model.CpModel()
x = {}
for p, d, s in itertools.product(range(9), range(DAYS), range(2)):
x[p, d, s] = model.NewBoolVar(f"x_{p}_{d}_{s}")
# === Dynamic Coverage ===
for d in range(DAYS):
if d in WEEKDAY_REL:
day_req = wd_day
if d in reduced_days:
day_req = 2
night_req = wd_night
else:
day_req = we_day
night_req = we_night
model.Add(sum(x[p, d, SHIFT["day"]] for p in range(9)) == day_req)
model.Add(sum(x[p, d, SHIFT["night"]] for p in range(9)) == night_req)
# === No same-day double shift (always enforced) ===
for p, d in itertools.product(range(9), range(DAYS)):
model.Add(x[p, d, SHIFT["day"]] + x[p, d, SHIFT["night"]] <= 1)
# === NO CONSECUTIVE DAYS - RELAXED FOR 9 STAFF ===
# Only prevent consecutive NIGHT shifts (keeps safety while enabling feasibility)
for p in range(9):
for d in range(DAYS - 1):
model.Add(x[p, d, SHIFT["night"]] + x[p, d + 1, SHIFT["night"]] <= 1)
# === Weekend cap (always enforced) ===
for p in range(9):
model.Add(sum(x[p, d, s] for d in (5, 6) for s in range(2)) <= 1)
# === 48h rest AFTER NIGHT SHIFT - RELAXED ===
# Keep 24h rest (no shift next day) but remove second day restriction
for p in range(9):
for d in range(DAYS):
night_d = x[p, d, SHIFT["night"]]
if d + 1 < DAYS:
any_d1 = x[p, d + 1, SHIFT["day"]] + x[p, d + 1, SHIFT["night"]]
model.Add(any_d1 <= 1 - night_d)
# === Vacants forced to 0 ===
for p in range(N_AVAIL, 9):
for d, s in itertools.product(range(DAYS), range(2)):
model.Add(x[p, d, s] == 0)
# === Weekly bounds - 9 staff can handle 2-3 shifts ===
MIN_WEEKLY, MAX_WEEKLY = 2, 3
week_shifts = {}
for i, name in enumerate(available_staff):
var = model.NewIntVar(MIN_WEEKLY, MAX_WEEKLY, f"wshift_{i}")
model.Add(var == sum(x[i, d, s] for d in range(DAYS) for s in range(2)))
week_shifts[name] = var
# === Fairness objective ===
objective_terms = []
for i, name in enumerate(available_staff):
cum = cumulative_shifts.get(name, 0)
objective_terms.append(cum * week_shifts[name])
model.Minimize(sum(objective_terms))
solver = cp_model.CpSolver()
solver.parameters.max_time_in_seconds = 45.0 # Longer timeout for complex solves
solver.parameters.num_search_workers = 6
status = solver.Solve(model)
if status not in (cp_model.OPTIMAL, cp_model.FEASIBLE):
status_map = {
cp_model.UNKNOWN: "UNKNOWN",
cp_model.MODEL_INVALID: "MODEL_INVALID",
cp_model.FEASIBLE: "FEASIBLE",
cp_model.INFEASIBLE: "INFEASIBLE",
cp_model.OPTIMAL: "OPTIMAL",
}
status_name = status_map.get(status, f"Status {status}")
error_msg = (
f"Week {week_idx + 1} solver returned: {status_name} with {N_AVAIL} staff\n"
)
error_msg += "To fix infeasibility:\n"
error_msg += "1. Check if constraints are too strict\n"
error_msg += "2. Verify staff count (9 should work with relaxed constraints)\n"
error_msg += "3. Ensure coverage requirements match team capacity"
raise RuntimeError(error_msg)
# === Extract solution ===
schedule_week = {}
weekly_counts = {name: 0 for name in available_staff}
for d in range(7):
day_staff = [
full_names[p]
for p in range(9)
if solver.Value(x[p, d, SHIFT["day"]])
and not full_names[p].startswith("Vacant_")
]
night_staff = [
full_names[p]
for p in range(9)
if solver.Value(x[p, d, SHIFT["night"]])
and not full_names[p].startswith("Vacant_")
]
schedule_week[d] = {"day": day_staff, "night": night_staff}
for name in day_staff + night_staff:
weekly_counts[name] += 1
return schedule_week, weekly_counts
def solve_week(
week_idx: int,
start_date: datetime.date,
available_staff: list[str],
cumulative_shifts: dict[str, int],
all_staff_global: list[str],
) -> tuple[dict, dict]:
"""Solve one week with constraints that adapt to team size."""
N_AVAIL = len(available_staff)
if N_AVAIL < 5:
raise ValueError("Minimum 5 staff per week required.")
# Determine constraint strictness based on staff count
IS_LARGE_TEAM = N_AVAIL >= 9 # 9 staff get strictest constraints
IS_MEDIUM_LARGE_TEAM = N_AVAIL == 8 # 8 staff get moderate constraints
# Get coverage rules
wd_day, wd_night, we_day, we_night, reduced_days = get_weekly_requirements(N_AVAIL)
full_names = available_staff + [f"Vacant_{i}" for i in range(9 - N_AVAIL)]
SHIFT = {"day": 0, "night": 1}
DAYS = 7
WEEKDAY_REL = {0, 1, 2, 3, 4}
model = cp_model.CpModel()
x = {}
for p, d, s in itertools.product(range(9), range(DAYS), range(2)):
x[p, d, s] = model.NewBoolVar(f"x_{p}_{d}_{s}")
# === Dynamic Coverage ===
for d in range(DAYS):
if d in WEEKDAY_REL:
day_req = wd_day
if d in reduced_days:
day_req = 2
night_req = wd_night
else:
day_req = we_day
night_req = we_night
model.Add(sum(x[p, d, SHIFT["day"]] for p in range(9)) == day_req)
model.Add(sum(x[p, d, SHIFT["night"]] for p in range(9)) == night_req)
# === No same-day double shift (always enforced) ===
for p, d in itertools.product(range(9), range(DAYS)):
model.Add(x[p, d, SHIFT["day"]] + x[p, d, SHIFT["night"]] <= 1)
# === No consecutive days (adaptive) ===
for p in range(9):
for d in range(DAYS - 1):
if IS_LARGE_TEAM:
# Full restriction for 9 staff teams
model.Add(
x[p, d, SHIFT["day"]]
+ x[p, d, SHIFT["night"]]
+ x[p, d + 1, SHIFT["day"]]
+ x[p, d + 1, SHIFT["night"]]
<= 1
)
elif IS_MEDIUM_LARGE_TEAM:
# For 8 staff: allow consecutive day shifts, restrict night shifts
model.Add(x[p, d, SHIFT["night"]] + x[p, d + 1, SHIFT["night"]] <= 1)
# Only prevent day-after-night and night-after-day
model.Add(x[p, d, SHIFT["night"]] + x[p, d + 1, SHIFT["day"]] <= 1)
model.Add(x[p, d, SHIFT["day"]] + x[p, d + 1, SHIFT["night"]] <= 1)
else:
# For 7 or fewer staff: minimal consecutive restrictions
model.Add(x[p, d, SHIFT["night"]] + x[p, d + 1, SHIFT["night"]] <= 1)
# === Weekend cap (always enforced) ===
for p in range(9):
model.Add(sum(x[p, d, s] for d in (5, 6) for s in range(2)) <= 1)
# === 48h rest after night shift (adaptive) ===
for p in range(9):
for d in range(DAYS):
night_d = x[p, d, SHIFT["night"]]
# Always enforce 24h rest (no shift next day after night)
if d + 1 < DAYS:
any_d1 = x[p, d + 1, SHIFT["day"]] + x[p, d + 1, SHIFT["night"]]
model.Add(any_d1 <= 1 - night_d)
# Only enforce second day off for 9-staff teams
if IS_LARGE_TEAM and d + 2 < DAYS:
any_d2 = x[p, d + 2, SHIFT["day"]] + x[p, d + 2, SHIFT["night"]]
model.Add(any_d2 <= 1 - night_d)
# === Vacants forced to 0 ===
for p in range(N_AVAIL, 9):
for d, s in itertools.product(range(DAYS), range(2)):
model.Add(x[p, d, s] == 0)
# === Weekly bounds (adaptive) ===
if N_AVAIL <= 6:
MIN_WEEKLY, MAX_WEEKLY = 3, 4 # Small teams work more shifts
elif N_AVAIL == 7:
MIN_WEEKLY, MAX_WEEKLY = 2, 4 # Medium teams need flexibility
elif N_AVAIL == 8:
MIN_WEEKLY, MAX_WEEKLY = 2, 3 # Medium-large teams
else: # 9 staff
MIN_WEEKLY, MAX_WEEKLY = 2, 3
week_shifts = {}
for i, name in enumerate(available_staff):
var = model.NewIntVar(MIN_WEEKLY, MAX_WEEKLY, f"wshift_{i}")
model.Add(var == sum(x[i, d, s] for d in range(DAYS) for s in range(2)))
week_shifts[name] = var
# === Soft fairness objective ===
objective_terms = []
for i, name in enumerate(available_staff):
cum = cumulative_shifts.get(name, 0)
objective_terms.append(cum * week_shifts[name])
model.Minimize(sum(objective_terms))
solver = cp_model.CpSolver()
# Longer timeout for smaller teams with tighter constraints
solver.parameters.max_time_in_seconds = 45.0 if N_AVAIL <= 8 else 30.0
solver.parameters.num_search_workers = 6
status = solver.Solve(model)
if status not in (cp_model.OPTIMAL, cp_model.FEASIBLE):
status_map = {
cp_model.UNKNOWN: "UNKNOWN",
cp_model.MODEL_INVALID: "MODEL_INVALID",
cp_model.FEASIBLE: "FEASIBLE",
cp_model.INFEASIBLE: "INFEASIBLE",
cp_model.OPTIMAL: "OPTIMAL",
}
status_name = status_map.get(status, f"Status {status}")
error_msg = (
f"Week {week_idx + 1} solver returned: {status_name} with {N_AVAIL} staff\n"
)
error_msg += "Active constraints:\n"
error_msg += f"- No consecutive days: {'FULL' if IS_LARGE_TEAM else ('MODERATE' if IS_MEDIUM_LARGE_TEAM else 'MINIMAL')}\n"
error_msg += f"- Night rest: {'48h' if IS_LARGE_TEAM else '24h'}\n"
error_msg += f"- Weekly bounds: {MIN_WEEKLY}-{MAX_WEEKLY} shifts\n"
error_msg += (
f"- Coverage: {wd_day}/{wd_night} weekdays, {we_day}/{we_night} weekends"
)
if reduced_days:
error_msg += (
f"\n- Reduced day(s): {', '.join(str(d) for d in reduced_days)}"
)
# Add troubleshooting tips
if N_AVAIL == 8 and status == cp_model.INFEASIBLE:
error_msg += "\n\nüí° Troubleshooting for 8 staff:"
error_msg += "\n- Try increasing max weekly shifts to 4 for some staff"
error_msg += "\n- Consider relaxing weekend cap for one staff member"
error_msg += "\n- Check if night shift distribution is balanced"
raise RuntimeError(error_msg)
# === Extract solution ===
schedule_week = {}
weekly_counts = {name: 0 for name in available_staff}
for d in range(7):
day_staff = [
full_names[p]
for p in range(9)
if solver.Value(x[p, d, SHIFT["day"]])
and not full_names[p].startswith("Vacant_")
]
night_staff = [
full_names[p]
for p in range(9)
if solver.Value(x[p, d, SHIFT["night"]])
and not full_names[p].startswith("Vacant_")
]
schedule_week[d] = {"day": day_staff, "night": night_staff}
for name in day_staff + night_staff:
weekly_counts[name] += 1
return schedule_week, weekly_counts
# ------------------------------
# Google Drive & Gmail Helpers (graceful fallback)
# ------------------------------
import io
from google.oauth2.service_account import Credentials
from googleapiclient.discovery import build
from pydrive2.auth import GoogleAuth
from pydrive2.drive import GoogleDrive
def get_drive():
scopes = [
"https://www.googleapis.com/auth/drive.file",
"https://www.googleapis.com/auth/gmail.send",
]
cred_path = "credentials.json"
if Path(cred_path).exists():
creds = Credentials.from_service_account_file(cred_path, scopes=scopes)
else:
cred_json = os.getenv("GOOGLE_CREDENTIALS_JSON")
if cred_json:
creds = Credentials.from_service_account_info(
json.loads(cred_json), scopes=scopes
)
else:
raise FileNotFoundError("No Google credentials found")
gauth = GoogleAuth()
gauth.credentials = creds
return GoogleDrive(gauth)
DRIVE_FOLDER_ID = os.getenv("DRIVE_FOLDER_ID", "root")
def save_to_drive(filename: str, content: bytes):
try:
drive = get_drive()
file = drive.CreateFile(
{"title": filename, "parents": [{"id": DRIVE_FOLDER_ID}]}
)
file.content = io.BytesIO(content)
file.Upload()
return file["id"]
except Exception as e:
st.warning(f"Google Drive save failed: {e}")
return None
def list_drive_files(prefix="roster_"):
try:
drive = get_drive()
files = drive.ListFile(
{
"q": f"'{DRIVE_FOLDER_ID}' in parents and title contains '{prefix}' and trashed=false"
}
).GetList()
return sorted(files, key=lambda f: f["createdDate"], reverse=True)
except Exception:
return []
def load_from_drive(file_id: str) -> bytes:
try:
drive = get_drive()
file = drive.CreateFile({"id": file_id})
buffer = io.BytesIO()
file.GetContentFile(buffer)
return buffer.getvalue()
except Exception as e:
st.error(f"Drive load failed: {e}")
return None
def send_email(to: str, subject: str, body: str):
try:
scopes = ["https://www.googleapis.com/auth/gmail.send"]
cred_path = "credentials.json"
if Path(cred_path).exists():
creds = Credentials.from_service_account_file(cred_path, scopes=scopes)
else:
cred_json = os.getenv("GOOGLE_CREDENTIALS_JSON")
if not cred_json:
st.warning("Email skipped: No Google credentials found")
return
creds = Credentials.from_service_account_info(
json.loads(cred_json), scopes=scopes
)
service = build("gmail", "v1", credentials=creds)
from_email = os.getenv("GMAIL_FROM", "no-reply@yourdomain.com")
message = f"""From: {from_email}
To: {to}
Subject: {subject}
MIME-Version: 1.0
Content-Type: text/html; charset=utf-8
{body}"""
import base64
raw = base64.urlsafe_b64encode(message.encode()).decode()
body_req = {"raw": raw}
service.users().messages().send(userId="me", body=body_req).execute()
except Exception as e:
st.warning(f"Email sending failed: {e}")
def schedule_weekly_emails(
schedule_full: dict, staff_emails: dict, start_date: datetime.date
):
def job():
today = datetime.now().date()
days_since = (today - start_date).days
if not (0 <= days_since < 42):
return
week_idx = days_since // 7
for name, email in staff_emails.items():
shifts = []
for d in range(week_idx * 7, min((week_idx + 1) * 7, 42)):
dt = start_date + timedelta(days=d)
if name in schedule_full.get(d, {}).get("day", []):
shifts.append(f"{dt:%a %d %b} Day")
if name in schedule_full.get(d, {}).get("night", []):
shifts.append(f"{dt:%a %d %b} Night")
if shifts:
body = (
f"Hi {name},
Your shifts for Week {week_idx + 1}:
"
+ "".join(f"- {s}
" for s in shifts)
+ "
Rest well!
"
)
send_email(email, f"Roster: Week {week_idx + 1}", body)
schedule.every().monday.at("08:00").do(job)
def run_sched():
while True:
schedule.run_pending()
pytime.sleep(60)
threading.Thread(target=run_sched, daemon=True).start()
# ------------------------------
# PDF & ICS Exports
# ------------------------------
def export_pdf(schedule, weekly_counts, start_date, output_path):
try:
from reportlab.lib import colors
from reportlab.lib.pagesizes import A4, landscape
from reportlab.lib.styles import ParagraphStyle
from reportlab.lib.units import inch
from reportlab.platypus import (
PageBreak,
Paragraph,
SimpleDocTemplate,
Table,
TableStyle,
)
doc = SimpleDocTemplate(
str(output_path),
pagesize=landscape(A4),
leftMargin=0.4 * inch,
rightMargin=0.4 * inch,
topMargin=0.4 * inch,
bottomMargin=0.4 * inch,
)
elements = []
title_style = ParagraphStyle("Title", fontSize=16, spaceAfter=12, alignment=1)
week_style = ParagraphStyle("Week", fontSize=14, spaceAfter=6, spaceBefore=12)
elements.append(Paragraph("6‚ÄëWeek Fair Roster", title_style))
def day_label(d):
return f"W{(d // 7) + 1:02d}-{'Mon Tue Wed Thu Fri Sat Sun'.split()[d % 7]}"
for w in range(6):
elements.append(Paragraph(f"Week {w + 1}", week_style))
data = [["Date", "Day", "Type", "Day Shift", "Night Shift"]]
for d in range(w * 7, (w + 1) * 7):
dt = start_date + timedelta(days=d)
dl = day_label(d)
typ = "WD" if (d % 7) < 5 else "WE"
ds = ", ".join(schedule.get(d, {}).get("day", []))
ns = ", ".join(schedule.get(d, {}).get("night", []))
data.append([dt.strftime("%a %d %b"), dl, typ, ds, ns])
table = Table(
data,
colWidths=[0.8 * inch, 0.9 * inch, 0.5 * inch, 2.2 * inch, 2.2 * inch],
)
table.setStyle(
TableStyle(
[
("BACKGROUND", (0, 0), (-1, 0), colors.darkblue),
("TEXTCOLOR", (0, 0), (-1, 0), colors.whitesmoke),
("GRID", (0, 0), (-1, -1), 0.5, colors.black),
("FONTNAME", (0, 0), (-1, 0), "Helvetica-Bold"),
("FONTSIZE", (0, 0), (-1, 0), 10),
]
)
)
elements.append(table)
if w < 5:
elements.append(PageBreak())
doc.build(elements)
return True
except Exception as e:
st.error(f"PDF failed: {e}")
return False
def export_ics(schedule, start_date, output_path):
try:
from icalendar import Alarm, Calendar, Event
cal = Calendar()
cal.add("prodid", "-//Roster//streamlit//")
cal.add("version", "2.0")
for d in range(42):
shift_date = start_date + timedelta(days=d)
for shift_type, key in [("Day", "day"), ("Night", "night")]:
for name in schedule.get(d, {}).get(key, []):
ev = Event()
st_t = dt_time(8, 0) if shift_type == "Day" else dt_time(20, 0)
en_t = dt_time(16, 0) if shift_type == "Day" else dt_time(8, 0)
ev.add("summary", f"{name} — {shift_type} Shift")
ev.add("dtstart", datetime.combine(shift_date, st_t))
ev.add(
"dtend",
datetime.combine(
shift_date
+ (
timedelta(days=1)
if shift_type == "Night"
else timedelta()
),
en_t,
),
)
ev.add("categories", [shift_type])
alarm = Alarm()
alarm.add("action", "DISPLAY")
alarm.add("trigger", timedelta(minutes=-15))
ev.add_component(alarm)
cal.add_component(ev)
with open(output_path, "wb") as f:
f.write(cal.to_ical())
return True
except Exception as e:
st.error(f"ICS failed: {e}")
return False
# ------------------------------
# Streamlit App ‚CORRECTED SESSION STATE HANDLING
# ------------------------------
st.set_page_config(page_title="Enterprise Roster (Final)", layout="wide")
st.title("Enterprise Roster Generator ‚Final")
# === Initialize session state (FIRST RUN ONLY) ===
if "initialized" not in st.session_state:
st.session_state.initialized = True
st.session_state.names = [""] * 9
st.session_state.emails = [""] * 9
# Use a separate key for default start date
st.session_state.start_date = (
datetime.today() + timedelta(days=(7 - datetime.today().weekday()) % 7)
).date()
st.session_state.user_role = "manager"
st.session_state.staff_email = ""
st.session_state.cumulative_shifts = {}
st.session_state.roster_weekly = {}
st.session_state.roster_ready = False
# === Auth UI ===
st.sidebar.header("Access")
role = st.sidebar.radio(
"Role",
["Manager", "Staff"],
index=0 if st.session_state.user_role == "manager" else 1,
key="role_radio",
)
st.session_state.user_role = "manager" if role == "Manager" else "staff"
if st.session_state.user_role == "staff":
staff_email_input = st.sidebar.text_input(
"Your Email", value=st.session_state.staff_email, key="staff_email_input"
)
st.session_state.staff_email = staff_email_input.strip().lower()
# === Manager Input ===
if st.session_state.user_role == "manager":
st.header("1. Staff")
cols = st.columns(3)
for i in range(9):
with cols[i % 3]:
name_val = st.text_input(
f"Staff {i + 1} Name",
value=st.session_state.names[i],
key=f"name_input_{i}",
)
email_val = st.text_input(
f"Email {i + 1}",
value=st.session_state.emails[i],
key=f"email_input_{i}",
)
st.session_state.names[i] = name_val.strip()
st.session_state.emails[i] = email_val.strip().lower()
st.header("2. Start Monday")
# ✅ CORRECT: widget key ≠ session state key
sd = st.date_input(
"First Monday",
value=st.session_state.start_date,
key="start_date_input", # ‚Üê distinct from st.session_state.start_date
)
# Update session state only if changed
if sd != st.session_state.start_date:
st.session_state.start_date = sd
st.header("3. Weekly Availability (Holiday/Mission)")
st.markdown(
"Uncheck staff who are unavailable (max 4 absent/week ‚min 5 available)."
)
avail_matrix = {}
cols_w = st.columns(6)
for w in range(6):
with cols_w[w]:
st.subheader(f"Week {w + 1}")
available = []
for i, name in enumerate([n for n in st.session_state.names if n]):
# ‚úÖ Use unique key per checkbox
is_avail = st.checkbox(f"{name}", value=True, key=f"avail_w{w}_p{i}")
if is_avail:
available.append(name)
avail_matrix[w] = available
if len(available) < 5:
st.error("⚠️ ≥5 must be available")
if st.button("Generate Rolling Roster", type="primary", key="generate_btn"):
try:
names_all = [n for n in st.session_state.names if n]
if not names_all:
st.error("Please enter at least one staff name.")
st.stop()
emails_all = {
n: e
for n, e in zip(
[n for n in st.session_state.names if n],
[
e
for i, e in enumerate(st.session_state.emails)
if st.session_state.names[i]
],
)
}
cum_shifts = {name: 0 for name in names_all}
weekly_sched = {}
# Validate weekly availability before solving
for w in range(6):
avail = avail_matrix[w]
if len(avail) < 5:
st.error(
f"Week {w + 1} has only {len(avail)} available staff. Minimum is 5."
)
st.stop()
# Check if we need reduced coverage
if len(avail) == 5:
st.warning(
f"Week {w + 1}: Reduced coverage mode (one weekday day shift = 2 staff)"
)
for w in range(6):
week_start = st.session_state.start_date + timedelta(weeks=w)
avail = avail_matrix[w]
try:
sched_w, counts_w = solve_week(
w, week_start, avail, cum_shifts, names_all
)
except Exception as e:
st.error(f"Failed to generate Week {w + 1}: {str(e)}")
st.stop()
for name, cnt in counts_w.items():
cum_shifts[name] = cum_shifts.get(name, 0) + cnt
abs_sched = {}
for d_rel, shifts in sched_w.items():
d_abs = w * 7 + d_rel
abs_sched[d_abs] = shifts
weekly_sched[w] = abs_sched
st.session_state.cumulative_shifts = cum_shifts
st.session_state.roster_weekly = weekly_sched
st.session_state.roster_ready = True
# Auto-save to Drive (graceful)
try:
data = pickle.dumps(
{
"weekly": weekly_sched,
"cumulative": cum_shifts,
"start": st.session_state.start_date,
}
)
fid = save_to_drive(
f"roster_{st.session_state.start_date:%Y%m%d}.pkl", data
)
if fid:
st.info(f"Saved to Drive (ID: {fid[:8]}…)")
except Exception as e:
st.warning(f"Drive save failed: {e}")
# Email scheduler
try:
full_sched = {}
for w_sched in weekly_sched.values():
full_sched.update(w_sched)
schedule_weekly_emails(
full_sched, emails_all, st.session_state.start_date
)
st.info("Weekly email reminders scheduled.")
except Exception as e:
st.warning(f"Email setup failed: {e}")
st.success("Rolling roster generated!")
except Exception as e:
st.error(f"Generation failed: {e}")
# === Display ===
if st.session_state.roster_ready:
full_sched = {}
for w_sched in st.session_state.roster_weekly.values():
full_sched.update(w_sched)
if st.session_state.user_role == "manager":
# Show warning if any week has reduced coverage
reduced_weeks = []
for w in range(6):
if len(avail_matrix[w]) == 5:
reduced_weeks.append(w + 1)
if reduced_weeks:
st.warning(
f"Reduced coverage in Week(s): {', '.join(map(str, reduced_weeks))} "
"(one weekday day shift = 2 staff instead of 3)."
)
st.header("Full 6 Week Roster")
rows = []
for d in range(42):
dt = st.session_state.start_date + timedelta(days=d)
wd = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"][d % 7]
week = d // 7 + 1
typ = "WD" if (d % 7) < 5 else "WE"
rows.append(
{
"Week": f"W{week}",
"Date": dt.strftime("%Y-%m-%d"),
"Day": wd,
"Type": typ,
"Day Shift": ", ".join(full_sched.get(d, {}).get("day", [])),
"Night Shift": ", ".join(full_sched.get(d, {}).get("night", [])),
}
)
df = pd.DataFrame(rows)
st.dataframe(df, use_container_width=True, hide_index=True)
st.subheader("Cumulative Shifts")
summ = []
for name in [n for n in st.session_state.names if n]:
summ.append(
{
"Staff": name,
"Total": st.session_state.cumulative_shifts.get(name, 0),
}
)
st.dataframe(pd.DataFrame(summ), use_container_width=True, hide_index=True)
# Exports
c1, c2, c3, c4 = st.columns(4)
with c1:
st.download_button(
"CSV",
df.to_csv(index=False).encode(),
"roster.csv",
"text/csv",
key="dl_csv",
)
with c2:
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as f:
if export_pdf(
full_sched, {}, st.session_state.start_date, Path(f.name)
):
with open(f.name, "rb") as pf:
st.download_button(
"PDF",
pf.read(),
"roster.pdf",
"application/pdf",
key="dl_pdf",
)
os.unlink(f.name)
with c3:
with tempfile.NamedTemporaryFile(delete=False, suffix=".ics") as f:
if export_ics(full_sched, st.session_state.start_date, Path(f.name)):
with open(f.name, "rb") as pf:
st.download_button(
"ICS",
pf.read(),
"roster.ics",
"text/calendar",
key="dl_ics",
)
os.unlink(f.name)
with c4:
if st.button("Clear", key="clear_btn"):
st.session_state.roster_ready = False
st.session_state.roster_weekly = {}
st.session_state.cumulative_shifts = {}
st.rerun()
# Drive Load
st.subheader("Load from Drive")
files = list_drive_files()
if files:
opts = {f["title"]: f["id"] for f in files}
sel = st.selectbox("Select roster", list(opts.keys()), key="drive_select")
if st.button("Load Selected", key="load_btn"):
try:
data = pickle.loads(load_from_drive(opts[sel]))
st.session_state.roster_weekly = data["weekly"]
st.session_state.cumulative_shifts = data["cumulative"]
st.session_state.start_date = data["start"]
st.session_state.roster_ready = True
st.success("Loaded!")
st.rerun()
except Exception as e:
st.error(f"Load failed: {e}")
else:
st.info("No saved rosters.")
else: # staff view
known_names = [n for n in st.session_state.names if n]
known_emails = [
e
for i, e in enumerate(st.session_state.emails)
if st.session_state.names[i]
]
staff_name = None
if st.session_state.staff_email in known_emails:
staff_name = known_names[known_emails.index(st.session_state.staff_email)]
if not staff_name:
st.warning("Enter an email matching a staff member.")
else:
st.header(f"Your Shifts, {staff_name}")
my_shifts = []
for d in range(42):
dt = st.session_state.start_date + timedelta(days=d)
if staff_name in full_sched.get(d, {}).get("day", []):
my_shifts.append({"Date": dt.strftime("%Y-%m-%d"), "Shift": "Day"})
if staff_name in full_sched.get(d, {}).get("night", []):
my_shifts.append(
{"Date": dt.strftime("%Y-%m-%d"), "Shift": "Night"}
)
if my_shifts:
st.dataframe(
pd.DataFrame(my_shifts), use_container_width=True, hide_index=True
)
else:
st.info("No shifts assigned to you in this roster period.")
st.caption(
"Final adaptive version — constraints adjust based on team size for 5-9 staff teams."
)