proc-rota/gen_proc.py

import datetime
import os
import sys
import time
from rota_generator.shifts import (
    NoWorkers,
    RotaBuilder,
    SingleShift,
    WorkerRequirement,
    days,
    PreShiftConstraint,
    PostShiftConstraint,
    NightConstraint,
    RequireRemoteSitePresenceConstraint,
    LimitGradeNumberConstraint,
    MinimumGradeNumberConstraint,
    MaxShiftsPerWeekConstraint,
)
import typer
import subprocess

app = typer.Typer()


sites = (
    "truro",
    "exeter",
    "torbay",
    "barnstaple",
    "plymouth",
    "taunton",
    "proc",
    "plymouth ir",
    "truro ir",
    "exeter ir",
    "torbay ir",
)

from rota_generator.workers import (
    Worker,
    NotAvailableToWork,
    NonWorkingDays,
    WorkRequests,
    PreferenceNotToWork,
    OutOfProgramme,
)


@app.command()
def main(
    suspend: bool = False,
    solve: bool = True,
    time_to_run: int = 60 * 60 * 10,
    ratio: float = 0.1,
    start_date: datetime.datetime = "2026-09-07",
    weeks: int = 26,
    bom: float = 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="proc_rota_sep_2026",
    )
    # Rota = RotaBuilder(start_date, weeks_to_rota=20, balance_offset_modifier=1)
    Rota.constraint_options["balance_shifts"] = False
    Rota.constraint_options["balance_shifts_quadratic"] = True
    Rota.constraint_options["balance_weekends"] = True
    Rota.constraint_options["max_night_frequency"] = 3  # 3
    #Rota.constraint_options["max_night_frequency_week_exclusions"] = []
    Rota.constraint_options["max_weekend_frequency"] = 3

    Rota.constraint_options["max_days_per_week_block"] = [(8,4)]

    # wr = [
    #    WorkerRequirement(
    #        end_date=datetime.datetime.strptime("2025-06-02", "%Y-%m-%d").date(),
    #        number=3,
    #    ),
    #    WorkerRequirement(
    #        start_date=datetime.datetime.strptime("2025-06-02", "%Y-%m-%d").date(),
    #        number=4,
    #    ),
    # ]

    Rota.add_shifts(
        SingleShift(
            sites=(
                "exeter",
                "exeter twilights and weekends",
                "exeter no nights",
                "exeter twilights",
                "exeter ir",
            ),
            name="exeter_twilight",
            length=12.5,
            days=days[:5],
            balance_offset=2,
            constraints=[MaxShiftsPerWeekConstraint(max_shifts=1)],
        ),
        SingleShift(
            sites=(
                "truro",
                "truro no nights",
                "truro twilights",
                "truro twilights and weekends",
                "truro twilights and weekend nights",
                "truro twilights, weekends and weekend nights",
                "truro ir",
            ),
            name="truro_twilight",
            length=12.5,
            days=days[:4],
            balance_offset=4,
            constraints=[MaxShiftsPerWeekConstraint(max_shifts=1)],
        ),
        SingleShift(
            sites=(
                "torbay",
                "torbay twilights",
                "torbay twilights and weekends",
                "torbay ir",
            ),
            name="torbay_twilight",
            length=12.5,
            days=days[:4],
            balance_offset=2,
            assign_as_block=True,
        ),
        SingleShift(
            sites=(
                "plymouth",
                "plymouth twilights",
                "plymouth twilights and weekends",
                "plymouth twilights, weekends and weekend nights",
                "plymouth ir",
            ),
            name="plymouth_twilight",
            length=12.5,
            days=days[:5],
            balance_offset=2,
            workers_required=1,
            constraints=[MaxShiftsPerWeekConstraint(max_shifts=2)],
        ),
        SingleShift(
            sites=(
                "exeter",
                "exeter twilights and weekends",
                "exeter no nights",
                "exeter weekends",
                # "exeter twilights",
            ),
            name="weekend_exeter",
            length=12.5,
            days=days[5:],
            # balance_offset=2,
            rota_on_nwds=True,
            force_as_block=True,
            constraints=[PostShiftConstraint(days=2), PreShiftConstraint(days=2)],
            # constraint=[{"name": "pre", "options": 2}, {"name": "post", "options": 2},],
        ),
        SingleShift(
            sites=(
                "truro",
                "truro no nights",
                "truro twilights and weekends",
                "truro weekends",
                "truro twilights, weekends and weekend nights",
            ),
            name="weekend_truro",
            length=12.5,
            days=days[4:],
            # balance_offset=3,
            rota_on_nwds=True,
            force_as_block=True,
            # assign_as_block=True,
            constraints=[PreShiftConstraint(days=2), PostShiftConstraint(days=2)],
            # force_as_block_unless_nwd=True
        ),
        SingleShift(
            sites=("torbay", "torbay twilights and weekends"),
            name="weekend_torbay",
            length=12.5,
            days=days[4:],
            # balance_offset=2,
            rota_on_nwds=True,
            force_as_block=True,
            constraints=[PreShiftConstraint(days=2), PostShiftConstraint(days=2)],
            # force_as_block_unless_nwd=True
        ),
        SingleShift(
            sites=(
                "plymouth",
                "plymouth twilights and weekends",
                "plymouth first on weekends only",
                "plymouth twilights, weekends and weekend nights",
            ),
            name="weekend_plymouth1",
            length=8,
            days=days[5:],
            # balance_offset=2.9,
            workers_required=1,
            rota_on_nwds=True,
            force_as_block=True,
            constraints=[PreShiftConstraint(days=2), PostShiftConstraint(days=2)],
        ),
        SingleShift(
            sites=(
                "plymouth",
                "plymouth twilights and weekends",
                "plymouth twilights, weekends and weekend nights",
            ),
            name="weekend_plymouth2",
            length=8,
            days=days[5:],
            # balance_offset=2.9,
            workers_required=1,
            rota_on_nwds=True,
            force_as_block=True,
            constraints=[PreShiftConstraint(days=2), PostShiftConstraint(days=2)],
        ),
        SingleShift(
            sites=(
                "plymouth",
                "plymouth twilights",
                "plymouth twilights and weekends",
                "plymouth twilights, weekends and weekend nights",
                "plymouth ir",
            ),
            name="plymouth_bank_holidays",
            length=8,
            days=days[5:],
            balance_offset=1,
            workers_required=1,
            rota_on_nwds=True,
            force_as_block=False,
            bank_holidays_only=True,
        ),
        SingleShift(
            sites=[
                *sites,
            ],
            name="night_weekday",
            length=12.25,
            days=days[:4],
            # balance_offset=3.9,
            balance_weighting=1,
            # hard_constrain_shift=False,
            workers_required=4,
            force_as_block=True,
            rota_on_nwds=True,
            constraints=[
                NightConstraint(),
                PreShiftConstraint(days=2),
                PostShiftConstraint(days=2),
                RequireRemoteSitePresenceConstraint(site="plymouth", required_number=1),
                LimitGradeNumberConstraint(grade=2, max_number=2),
                MinimumGradeNumberConstraint(grade=4, min_number=1),
            ],
            # end_date=datetime.datetime.strptime("2025-06-01", "%Y-%m-%d").date(),
        ),
        SingleShift(
            sites=[
                *sites,
                "plymouth twilights, weekends and weekend nights",
                "truro twilights and weekend nights",
                "weekend nights",
                "truro twilights, weekends and weekend nights",
            ],
            name="night_weekend",
            length=12.25,
            days=days[4:],
            # balance_offset=2.9,
            balance_weighting=1,
            # hard_constrain_shift=False,
            workers_required=4,
            force_as_block=True,
            rota_on_nwds=True,
            constraints=[
                NightConstraint(),
                PreShiftConstraint(days=2),
                PostShiftConstraint(days=3),
                RequireRemoteSitePresenceConstraint(site="plymouth", required_number=1),
                LimitGradeNumberConstraint(grade=2, max_number=2),
                MinimumGradeNumberConstraint(grade=4, min_number=1),
            ],
            # end_date=datetime.datetime.strptime("2025-06-01", "%Y-%m-%d").date(),
        ),
    )

    # Prevent shifts for ST2s for 2 weeks
    # Rota.add_grade_constraint_by_week([2], [1, 2], [shift.name for shift in Rota.shifts])

    Rota.add_min_summed_grade_by_shifts_per_day_constraint(
        ["weekend_plymouth1", "weekend_plymouth2"], 6
    )
    Rota.add_min_summed_grade_by_shifts_per_day_constraint(
        ["plymouth_twilight", "plymouth_bank_holidays"], 6
    )

    Rota.add_min_summed_grade_by_shifts_per_day_constraint(["night_weekday"], 12)
    Rota.add_min_summed_grade_by_shifts_per_day_constraint(["night_weekend"], 12)

    Rota.terminate_on_warning.remove("Worker/no valid shifts")

    load_leave = True
    Rota.build_shifts()
    # Rota.add_grade_constraint_by_week([2], [1], Rota.get_shift_names())

    if load_leave:
        import leave

        # Internal-only option: remove leave requests between two dates.
        # Configure by editing the two variables below (not exposed via CLI).
        # Use strings in 'YYYY-MM-DD' or datetime.date objects. Set either to None to disable.
        REMOVE_LEAVE_FROM = ""
        REMOVE_LEAVE_TO = ""

        workers = leave.load_leave(Rota)

        # If the internal remove range is set, filter out NotAvailableToWork entries
        if REMOVE_LEAVE_FROM and REMOVE_LEAVE_TO:
            # Coerce to date objects if strings provided
            def _to_date(v):
                if isinstance(v, str):
                    for fmt in ("%Y-%m-%d", "%d/%m/%Y", "%d/%m/%y"):
                        try:
                            return datetime.datetime.strptime(v, fmt).date()
                        except Exception:
                            continue
                    raise ValueError(f"Cannot parse date: {v}")
                if isinstance(v, datetime.datetime):
                    return v.date()
                return v

            _from = _to_date(REMOVE_LEAVE_FROM)
            _to = _to_date(REMOVE_LEAVE_TO)

            for name, w in list(workers.items()):
                original = w.get("leave", [])
                before_count = len(original)
                filtered = [lr for lr in original if not (_from <= lr.date <= _to)]
                after_count = len(filtered)
                if after_count != before_count:
                    print(
                        f"Removed {before_count - after_count} leave entries for '{name}' between {_from} and {_to}"
                    )
                w["leave"] = filtered

        n = 0
        for worker in workers:
            worker_name = worker

            if worker_name == "Name":
                continue
            n = n + 1
            w = workers[worker]
            try:
                site = w["site"]
            except KeyError:
                print(f"Worker {worker} has no site")
                print(w)
                sys.exit(1)
            grade = w["grade"]
            try:
                fte = float(w["fte"]) * 100
            except ValueError:
                print(f"{worker} has invalid fte: {fte}")
                raise ValueError
            except KeyError:
                print(f"Worker {worker} has no FTE")
                print(w)
                sys.exit(1)

            nwd = w["nwd"]
            end_date = w["end_date"]
            start_date = w["start_date"]
            oop = w["oop"]
            grade = w["grade"]
            pair = w["pair"]
            # name, site, grade, fte, nwd, end_date, oop, twi, twi_we, night, night_we, pw1, pw2, extra = row

            # Ignore trainees if fte == 0 (what are they doing here anyway)
            if fte == 0 or fte == "0" or fte == "":
                continue

            nwds = []
            if nwd:
                for i in nwd.split(","):
                    nwd_start_date = Rota.start_date
                    nwd_end_date = Rota.rota_end_date

                    if "[" in i:
                        # print("Split nwd", worker_name, i)
                        a, b = i.split("[")[1][:-1].split("-")
                        # print(f"{a=} {b=}")
                        nwd_start_date = datetime.datetime.strptime(
                            a, "%d/%m/%Y"
                        ).date()
                        nwd_end_date = datetime.datetime.strptime(b, "%d/%m/%Y").date()

                    nwds.append(
                        NonWorkingDays(
                            day=i.split("[")[0].strip()[:3].capitalize(),
                            start_date=nwd_start_date,
                            end_date=nwd_end_date,
                        )
                    )
            # nwds = [i.strip()[:3].capitalize() for i in nwd.split("/")] if nwd else None
            end_date = end_date if end_date else None

            if oop:
                formatted_oops = []
                for dates in oop.split(","):
                    raw = dates.strip()
                    # strip surrounding brackets/parentheses/quotes
                    raw = raw.strip("()[]\"' ")
                    if "-" in raw:
                        s, e = raw.split("-", 1)
                    elif " to " in raw:
                        s, e = raw.split(" to ", 1)
                    elif "to" in raw:
                        s, e = raw.split("to", 1)
                    else:
                        raise ValueError(
                            f"Cannot parse OOP dates: '{oop}' for {worker}"
                        )

                    s = s.strip().strip("()[]\"' ")
                    e = e.strip().strip("()[]\"' ")

                    parsed = False
                    for fmt in ("%d/%m/%Y", "%d/%m/%y"):
                        try:
                            start_dt = datetime.datetime.strptime(s, fmt).date()
                            end_dt = datetime.datetime.strptime(e, fmt).date()
                            formatted_oops.append(
                                {"start_date": start_dt, "end_date": end_dt}
                            )
                            parsed = True
                            break
                        except ValueError:
                            continue

                    if not parsed:
                        print("WORKER", worker)
                        print("DATES", s, e)
                        raise ValueError(
                            f"Cannot parse OOP dates: '{oop}' for {worker}"
                        )

                oop = formatted_oops
            else:
                oop = []

            # print(f"{worker} OOP {oop}")

            previous_shifts = {}
            # if twi:
            #    worked = twi.split(" ")[0]
            #    allocated = twi.split(" ")[1][1:-1]
            #    previous_shifts["{}_twilight".format(site.lower())] = (worked, allocated)
            # if twi_we:
            #    worked = twi_we.split(" ")[0]
            #    allocated = twi_we.split(" ")[1][1:-1]
            #    previous_shifts["weekend_{}".format(site.lower())] = (worked, allocated)
            # if night:
            #    worked = night.split(" ")[0]
            #    allocated = night.split(" ")[1][1:-1]
            #    previous_shifts["night_weekday".format(site.lower())] = (worked, allocated)
            # if night_we:
            #    worked = night_we.split(" ")[0]
            #    allocated = night_we.split(" ")[1][1:-1]
            #    previous_shifts["night_weekend".format(site.lower())] = (worked, allocated)
            # if pw1:
            #    worked = pw1.split(" ")[0]
            #    allocated = pw1.split(" ")[1][1:-1]
            #    previous_shifts["weekend_plymouth1"] = (worked, allocated)
            # if pw2:
            #    worked = pw2.split(" ")[0]
            #    allocated = pw2.split(" ")[1][1:-1]
            #    previous_shifts["weekend_plymouth2"] = (worked, allocated)

            leave_requests = w["leave"]
            work_requests = w["requests"]
            site_pref = w["site_pref"]

            shift_fte_overrides = {}

            exact_shifts = {}

            if worker in ["hannah.lewis29@nhs.net", "shruti.bodapati@nhs.net", "a.abouelatta@nnhs.net"]:
                exact_shifts = {
                    "night_weekday": 4,
                    "night_weekend": 3,
                }
            elif worker in ["g.abdelhalim@nhs.net", "nang.thiriphoo@nhs.net"]:
                exact_shifts = {
                    "night_weekend": 3,
                    "night_weekday": 0,
                }

            # if worker_name == "Hadi Mohamed":
            #    shift_fte_overrides = {
            #        "weekend_exeter": 50,
            #    }
            # elif worker_name == "Joel Lim":
            #    shift_fte_overrides = {
            #        "plymouth_twilight": 100,
            #        "weekend_exeter": 50,
            #    }
            # if worker_name == "Nang Thiriphoo":
            #    shift_fte_overrides = {
            #        "weekend_exeter": 40,
            #    }

            w = Worker(
                name=worker_name,
                site=site.lower(),
                grade=int(grade),
                id=n,
                fte=int(fte),
                nwds=nwds,
                start_date=start_date,
                end_date=end_date,
                oop=oop,
                not_available_to_work=leave_requests,
                work_requests=work_requests,
                remote_site=site_pref,
                previous_shifts=previous_shifts,
                pair=pair,
                shift_balance_extra=w["shift_balance_extra"],
                bank_holiday_extra=w["bank_holiday_extra"],
                shift_fte_overrides=shift_fte_overrides,
                exact_shifts=exact_shifts,
            )

            # print(w)

            Rota.add_worker(w)
        leave.load_academy(Rota)

    # 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", export_with_timestamp=True)

    # 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()