import pytest
import datetime
from rota_generator.shifts import BalanceAcrossGroupsConstraint, RotaBuilder, SingleShift, days
from rota_generator.workers import Worker

@pytest.fixture
def demo_rota_balance_sites():
    weeks_to_rota = 8
    start_date = datetime.date(2022, 3, 7)
    Rota = RotaBuilder(
        start_date,
        weeks_to_rota=weeks_to_rota,
    )
    worker1 = Worker(name="worker1", site="group1", grade=1)
    worker2 = Worker(name="worker2", site="group1", grade=1)
    worker3 = Worker(name="worker3", site="group2", grade=1)
    worker4 = Worker(name="worker4", site="group2", grade=1)
    worker5 = Worker(name="worker5", site="group3", grade=1)
    worker6 = Worker(name="worker6", site="group3", grade=1)
    Rota.add_workers((worker1, worker2, worker3, worker4, worker5, worker6))
    return Rota

def test_balance_blocks_across_2_groups(demo_rota_balance_sites):
    Rota = demo_rota_balance_sites
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="weekday_night",
            length=12.5,
            days=days[:4],
            balance_offset=40,
            workers_required=2,
            force_as_block=True,
            constraints=[BalanceAcrossGroupsConstraint()],
        ),
    )
    Rota.constraint_options["balance_nights_across_sites"] = False
    Rota.constraint_options["balance_nights"] = False
    Rota.constraint_options["constrain_time_off_after_nights"] = False

    # Worker 5 and 6 don't have any valid shifts so this should fail
    #with pytest.raises(WarningTermination):
    #    Rota.build_and_solve(options={"ratio": 0.0})

    # Unless we remove the warning
    Rota.terminate_on_warning.remove("Worker/no valid shifts")
    Rota.build_and_solve(options={"ratio": 0.0})

    assert len(Rota.get_warnings("Worker/no valid shifts")) == 2

    Rota.export_rota_to_html("test5")

    group_workers = Rota.get_workers_by_group()
    for group in group_workers:
        shift_patterns = []
        for w in group_workers[group]:
            shift_patterns.append(Rota.get_worker_shift_list(w))

        zipped_lists = list(zip(*shift_patterns))

        for day_shifts in zipped_lists:
            assert day_shifts.count("weekday_night") <= 1

    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"

def test_balance_blocks_across_2_groups_unbalanced(demo_rota_balance_sites):
    Rota = demo_rota_balance_sites
    Rota.shifts = []
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="weekday_night",
            length=12.5,
            days=days[:4],
            balance_offset=40,
            workers_required=3,
            force_as_block=True,
            constraints=[BalanceAcrossGroupsConstraint()],
        ),
    )
    Rota.constraint_options["balance_nights_across_sites"] = False
    Rota.constraint_options["balance_nights"] = False
    Rota.constraint_options["constrain_time_off_after_nights"] = False

    try:
        Rota.terminate_on_warning.remove("Worker/no valid shifts")
    except ValueError:
        # Seems to happen with sequential test running
        pass
    Rota.build_and_solve(options={"ratio": 0.1})

    group_workers = Rota.get_workers_by_group()
    for group in group_workers:
        shift_patterns = []
        for w in group_workers[group]:
            shift_patterns.append(Rota.get_worker_shift_list(w))

        zipped_lists = list(zip(*shift_patterns))

        for day_shifts in zipped_lists:
            assert day_shifts.count("weekday_night") <= 2

    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"

def test_balance_blocks_across_3_groups(demo_rota_balance_sites):
    Rota = demo_rota_balance_sites
    Rota.shifts = []
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2", "group3"),
            name="weekday_night",
            length=12.5,
            days=days[:4],
            balance_offset=40,
            workers_required=3,
            force_as_block=True,
            constraints=[BalanceAcrossGroupsConstraint()],
        ),
    )
    Rota.constraint_options["balance_nights_across_sites"] = False
    Rota.constraint_options["balance_nights"] = False
    Rota.constraint_options["constrain_time_off_after_nights"] = False
    Rota.build_and_solve(options={"ratio": 0.1})

    group_workers = Rota.get_workers_by_group()
    for group in group_workers:
        shift_patterns = []
        for w in group_workers[group]:
            shift_patterns.append(Rota.get_worker_shift_list(w))

        zipped_lists = list(zip(*shift_patterns))

        for day_shifts in zipped_lists:
            assert day_shifts.count("weekday_night") <= 1

    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"