proc-rota/gen_cons.py

import datetime
import os
import sys
import time
from rota.shifts import NoWorkers, RotaBuilder, SingleShift, WorkerRequirement, days
import typer
import subprocess

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


app = typer.Typer()


sites = ("rota a", "rota b", "rota c")


@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 = 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,
        name="cons rota",
    )
    # Rota = RotaBuilder(start_date, weeks_to_rota=20, balance_offset_modifier=1)
    Rota.constraint_options["balance_weekends"] = True
    Rota.constraint_options["max_weekend_frequency"] = 4
    Rota.constraint_options["max_shifts_per_week"] = 20
    # Rota.constraint_options["avoid_st2_first_month"] = True

    Rota.add_shifts(
        SingleShift(
            sites=("rota a",),
            name="oncall",
            length=12.5,
            days=days,
            balance_offset=2,
            assign_as_block=False,
            constraint=[
                #{
                #    "name": "max_shifts_per_week",
                #    "options": 3,
                #}
            ],
        ),
        SingleShift(
            sites=("rota a",),
            name="weekend",
            workers_required=1,
            length=12.5,
            days=days[5:],
            balance_offset=2,
            constraint=[
                #{"name": "pre", "options": 2},
                #{"name": "post", "options": 1},
            ],
            display_char="a",
            #force_assign_with=["oncall"]
        ),
        SingleShift(
            sites=("rota b",),
            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",
            ),
            name="twilight",
            workers_required=1,
            length=12.5,
            days=days[:5],
            balance_offset=5,
            #constraint=[
            #    {"name": "pre", "options": 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())

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