import datetime
import pytest
from pytest import approx
from rota.shifts import PostShiftConstraint, PreShiftConstraint, RotaBuilder, SingleShift, days
from rota.workers import Worker

def generate_basic_rota(weeks_to_rota=10):
    start_date = datetime.date(2022, 3, 7)
    Rota = RotaBuilder(
        start_date,
        weeks_to_rota=weeks_to_rota,
    )
    Rota.constraint_options["balance_weekends"] = True
    Rota.add_workers(
        [
            Worker(name="worker1", site="group1", grade=1),
            Worker(name="worker2", site="group1", grade=1),
            Worker(name="worker3", site="group2", grade=1),
        ]
    )
    return Rota

def test_basic_balancing():
    Rota = generate_basic_rota()
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="a",
            length=12.5,
            days=days,
            balance_offset=2,
            workers_required=1,
            force_as_block=False,
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="b",
            length=12.5,
            days=days,
            balance_offset=2,
            workers_required=1,
            force_as_block=False,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.001})
    Rota.export_rota_to_html("basic_balancing")
    shift_summary = Rota.get_shift_summary_dict()
    for worker_name in shift_summary:
        worker_shifts = shift_summary[worker_name]
        assert worker_shifts["a"] in (approx(22), approx(23), approx(24), approx(25))
        assert worker_shifts["b"] in (approx(22), approx(23), approx(24), approx(25))

def test_weighted_shift_balancing():
    Rota = generate_basic_rota(20)
    Rota.add_shifts(
        SingleShift(
            sites=["group1", "group2"],
            name="a",
            length=12.5,
            days=days,
            balance_weighting=10,
            workers_required=1,
            force_as_block=False,
            constraints=[PreShiftConstraint(days=2), PostShiftConstraint(days=2)]
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="b",
            length=12.5,
            days=days,
            workers_required=1,
            force_as_block=False,
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="c",
            length=12.5,
            days=days[0],
            workers_required=1,
            force_as_block=False,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.0001})
    Rota.export_rota_to_html("basic_balancing_weighted_shifts")
    shift_summary = Rota.get_shift_summary_dict()
    for worker_name in shift_summary:
        worker_shifts = shift_summary[worker_name]
        assert worker_shifts["a"] in (approx(46), approx(47), approx(48))

def test_weighted_shift_balancing2():
    Rota = generate_basic_rota(23)
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="a",
            length=12.5,
            days=days,
            workers_required=1,
            force_as_block=False,
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="b",
            length=12.5,
            days=days,
            balance_weighting=8,
            workers_required=1,
            force_as_block=False,
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="c",
            length=12.5,
            days=days[0],
            workers_required=1,
            force_as_block=False,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.0001})
    Rota.export_rota_to_html("basic_balancing_weighted_shifts2")
    shift_summary = Rota.get_shift_summary_dict()
    for worker_name in shift_summary:
        worker_shifts = shift_summary[worker_name]
        assert worker_shifts["b"] in (approx(53), approx(54), approx(55))

def test_weighted_shift_balancing3():
    Rota = generate_basic_rota(23)
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="a",
            length=12.5,
            days=days,
            workers_required=1,
            force_as_block=False,
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="b",
            length=12.5,
            days=days,
            balance_weighting=1,
            workers_required=1,
            force_as_block=False,
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="c",
            length=12.5,
            days=days[0],
            balance_weighting=8,
            workers_required=1,
            force_as_block=False,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.0001})
    Rota.export_rota_to_html("basic_balancing_weighted_shifts3")
    shift_summary = Rota.get_shift_summary_dict()
    for worker_name in shift_summary:
        worker_shifts = shift_summary[worker_name]
        assert worker_shifts["c"] in (approx(7), approx(8))

def test_weighted_shift_balancing4():
    Rota = generate_basic_rota(10)
    Rota.add_worker(
        Worker(name="worker4", site="group3", grade=1),
    )
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="a",
            length=12.5,
            days=days,
            balance_offset=99,
            balance_weighting=8,
            workers_required=2,
            force_as_block=False,
        ),
        SingleShift(
            sites=("group2", "group3"),
            name="b",
            length=12.5,
            days=days,
            balance_offset=99,
            workers_required=1,
            force_as_block=False,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("basic_balancing_weighted_shifts4")
    shift_summary = Rota.get_shift_summary_dict()
    for worker_name in shift_summary:
        worker_shifts = shift_summary[worker_name]
        assert worker_shifts["a"] in (0, 46, 47, 48)

def test_weighted_shift_balancing5():
    Rota = generate_basic_rota(10)
    Rota.add_worker(
        Worker(name="worker4", site="group3", grade=1),
    )
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="a",
            length=12.5,
            days=days,
            balance_offset=99,
            workers_required=2,
            force_as_block=False,
        ),
        SingleShift(
            sites=("group2", "group3"),
            name="b",
            length=12.5,
            days=days,
            balance_offset=99,
            balance_weighting=4,
            workers_required=1,
            force_as_block=False,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("basic_balancing_weighted_shifts5")
    shift_summary = Rota.get_shift_summary_dict()
    for worker_name in shift_summary:
        worker_shifts = shift_summary[worker_name]
        assert worker_shifts["b"] in (0, 35)