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 = 16,
    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_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"] = 4
    Rota.constraint_options["hard_constrain_pair_separation"] = True
    Rota.constraint_options["max_shifts_per_week"] = 3
    # Rota.constraint_options["avoid_st2_first_month"] = True

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

    # 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 w in range(1, 12):
        w = Worker(name=f"A{w:02d}", site="rota a", grade=1)

        print(w)

        Rota.add_worker(w)

    for w in range(12, 21):
        w = Worker(
            name=f"B{w}",
            site="rota b",
            grade=1,
            nwds=[
                NonWorkingDays(
                    day="Fri",
                )
            ],
        )

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