import datetime
import pytest
from rota.shifts import NoWorkers, 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

    # Add a few workers
    Rota.add_workers([
        Worker(name="worker1", site="group1", grade=1),
        Worker(name="worker2", site="group1", grade=1),
        Worker(name="worker3", site="group2", grade=1),
        #Worker(name="worker4", site="group2", grade=1),
        #Worker(name="worker5", site="group2", grade=1),
        #Worker(name="worker6", site="group2", grade=1, fte=50),
    ])

    return Rota


class TestBalancing:
    def test_basic_balancing(self):
        Rota = generate_basic_rota()

        Rota.add_shifts(
            SingleShift(
                ("group1", "group2"),
                "a",
                12.5,
                days,
                balance_offset=2,
                workers_required=1,
                force_as_block=False
            ),
            SingleShift(
                ("group1", "group2"),
                "b",
                12.5,
                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]

            worker = Rota.get_worker_by_name(worker_name)
            assert worker_shifts["a"] in (22, 23, 24, 25)
            assert worker_shifts["b"] in (22, 23, 24, 25)

    def test_weighted_shift_balancing(self):
        Rota = generate_basic_rota(20)

        Rota.add_shifts(
            SingleShift(
                ("group1", "group2"),
                "a",
                12.5,
                days,
                balance_weighting=4,
                workers_required=1,
                force_as_block=False,
                constraint=["preclear2", "postclear2"],
            ),
            SingleShift(
                ("group1", "group2"),
                "b",
                12.5,
                days,
                workers_required=1,
                force_as_block=False
            ),
            SingleShift(
                ("group1", "group2"),
                "c",
                12.5,
                days[0],
                workers_required=1,
                force_as_block=False
            ),
        )

        Rota.build_and_solve(options={"ratio": 0.001})
        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]

            worker = Rota.get_worker_by_name(worker_name)
            assert worker_shifts["a"] in (46,47,48)
            #assert worker_shifts["b"] in (22, 23, 24, 25)

    def test_weighted_shift_balancing2(self):
        Rota = generate_basic_rota(23)

        Rota.add_shifts(
            SingleShift(
                ("group1", "group2"),
                "a",
                12.5,
                days,
                #balance_weighting=4,
                workers_required=1,
                force_as_block=False,
                #constraint=["preclear2", "postclear2"],
            ),
            SingleShift(
                ("group1", "group2"),
                "b",
                12.5,
                days,
                balance_weighting=4,
                workers_required=1,
                force_as_block=False
            ),
            SingleShift(
                ("group1", "group2"),
                "c",
                12.5,
                days[0],
                workers_required=1,
                force_as_block=False
            ),
        )

        Rota.build_and_solve(options={"ratio": 0.001})
        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]

            worker = Rota.get_worker_by_name(worker_name)
            #assert worker_shifts["a"] in (46,47,48)
            #assert worker_shifts["c"] in (46,47,48)
            assert worker_shifts["b"] in (52,53,54,55)

    def test_weighted_shift_balancing3(self):
        Rota = generate_basic_rota(23)

        Rota.add_shifts(
            SingleShift(
                ("group1", "group2"),
                "a",
                12.5,
                days,
                #balance_weighting=4,
                workers_required=1,
                force_as_block=False,
                #constraint=["preclear2", "postclear2"],
            ),
            SingleShift(
                ("group1", "group2"),
                "b",
                12.5,
                days,
                balance_weighting=4,
                workers_required=1,
                force_as_block=False
            ),
            SingleShift(
                ("group1", "group2"),
                "c",
                12.5,
                days[0],
                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_shifts3")

        shift_summary = Rota.get_shift_summary_dict()
        for worker_name in shift_summary:
            worker_shifts = shift_summary[worker_name]

            worker = Rota.get_worker_by_name(worker_name)
            #assert worker_shifts["a"] in (46,47,48)
            assert worker_shifts["c"] in (7,8)
            assert worker_shifts["b"] in (52,53,54,55)

    def test_weighted_shift_balancing4(self):
        Rota = generate_basic_rota(10)

        Rota.add_worker(
            Worker(name="worker4", site="group3", grade=1),
        )

        Rota.add_shifts(
            SingleShift(
                ("group1", "group2"),
                "a",
                12.5,
                days,
                balance_offset=99,
                balance_weighting=8,
                workers_required=2,
                force_as_block=False,
                #constraint=["preclear2", "postclear2"],
            ),
            SingleShift(
                ("group2","group3"),
                "b",
                12.5,
                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_shifts4")

        shift_summary = Rota.get_shift_summary_dict()
        for worker_name in shift_summary:
            worker_shifts = shift_summary[worker_name]

            worker = Rota.get_worker_by_name(worker_name)
            assert worker_shifts["a"] in (0, 46,47,48)

    def test_weighted_shift_balancing5(self):
        Rota = generate_basic_rota(10)

        Rota.add_worker(
            Worker(name="worker4", site="group3", grade=1),
        )

        Rota.add_shifts(
            SingleShift(
                ("group1", "group2"),
                "a",
                12.5,
                days,
                balance_offset=99,
                #balance_weighting=8,
                workers_required=2,
                force_as_block=False,
                #constraint=["preclear2", "postclear2"],
            ),
            SingleShift(
                ("group2","group3"),
                "b",
                12.5,
                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]

            worker = Rota.get_worker_by_name(worker_name)
            assert worker_shifts["b"] in (0, 35)