import pytest
import datetime
from rota_generator.shifts import MinimumGradeNumberConstraint, RequireRemoteSitePresenceConstraint, RotaBuilder, SingleShift, days, LimitGradeNumberConstraint
from rota_generator.workers import Worker

@pytest.fixture
def limit_constraint_rota():
    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=2, remote_site="group1")
    worker2 = Worker(name="worker2", site="group1", grade=2, remote_site="group1")
    worker3 = Worker(name="worker3", site="group1", grade=2, remote_site="group1")
    worker4 = Worker(name="worker4", site="group3", grade=3, remote_site="group2")
    worker5 = Worker(name="worker5", site="group3", grade=3, remote_site="group2")
    worker6 = Worker(name="worker6", site="group3", grade=3, remote_site="group2")
    Rota.add_workers((worker1, worker2, worker3, worker4, worker5, worker6))
    Rota.shifts = []
    return Rota

def test_constraint_limit_grades(limit_constraint_rota):
    Rota = limit_constraint_rota
    Rota.shifts = []
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2", "group3"),
            name="weekday_night",
            length=12.5,
            days=days,
            balance_offset=60,
            workers_required=4,
            force_as_block=True,
            constraints=[LimitGradeNumberConstraint(grade=2, max_number=1)]
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.01})
    Rota.export_rota_to_html("constraint_limit_grades", folder="tests")
    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"
    grade_workers = Rota.get_workers_by_grade()
    for grade in grade_workers:
        shift_patterns = [Rota.get_worker_shift_list(w) for w in grade_workers[grade]]
        zipped_lists = list(zip(*shift_patterns))
        limit = 1 if grade == 2 else 3
        for day_shifts in zipped_lists:
            assert day_shifts.count("weekday_night") == limit

def test_constraint_limit_grades2(limit_constraint_rota):
    Rota = limit_constraint_rota
    Rota.shifts = []
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2", "group3"),
            name="weekday_night",
            length=12.5,
            days=days,
            balance_offset=40,
            workers_required=4,
            force_as_block=True,
            constraints=[LimitGradeNumberConstraint(grade=3, max_number=1)],
        ),
    )
    Rota.constraint_options["balance_shifts_quadratic"] = True
    Rota.build_and_solve(options={"ratio": 0.1})
    Rota.export_rota_to_html("test9", folder="tests")
    grade_workers = Rota.get_workers_by_grade()
    for grade in grade_workers:
        shift_patterns = [Rota.get_worker_shift_list(w) for w in grade_workers[grade]]
        zipped_lists = list(zip(*shift_patterns))
        limit = 1 if grade == 3 else 3
        for day_shifts in zipped_lists:
            assert day_shifts.count("weekday_night") == limit
    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"

def test_constraint_limit_grades3(limit_constraint_rota):
    Rota = limit_constraint_rota
    Rota.workers = []
    worker1 = Worker(name="worker1", site="group1", grade=2, remote_site="group1")
    worker2 = Worker(name="worker2", site="group1", grade=2, remote_site="group1")
    worker3 = Worker(name="worker3", site="group1", grade=2, remote_site="group1")
    worker4 = Worker(name="worker4", site="group3", grade=3, remote_site="group2")
    worker5 = Worker(name="worker5", site="group3", grade=3, remote_site="group2")
    worker6 = Worker(name="worker6", site="group3", grade=3, remote_site="group2")
    worker7 = Worker(name="worker7", site="group3", grade=2, remote_site="group2")
    Rota.add_workers((worker1, worker2, worker3, worker4, worker5, worker6, worker7))
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2", "group3"),
            name="weekday_night",
            length=12.5,
            days=days,
            balance_offset=60,
            workers_required=4,
            force_as_block=True,
            constraints=[LimitGradeNumberConstraint(grade=2, max_number=3), LimitGradeNumberConstraint(grade=3, max_number=1)],
        ),
    )
    Rota.constraint_options["balance_shifts_quadratic"] = True
    Rota.build_and_solve(options={"ratio": 0.01})
    Rota.export_rota_to_html("test9")
    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"
    grade_workers = Rota.get_workers_by_grade()
    for grade in grade_workers:
        shift_patterns = [Rota.get_worker_shift_list(w) for w in grade_workers[grade]]
        zipped_lists = list(zip(*shift_patterns))
        limit = 3 if grade == 2 else 1
        for day_shifts in zipped_lists:
            assert day_shifts.count("weekday_night") == limit

def test_constraint_limit_grades4(limit_constraint_rota):
    Rota = limit_constraint_rota
    Rota.shifts = []
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2", "group3"),
            name="weekday_night",
            length=12.5,
            days=days,
            balance_offset=40,
            workers_required=4,
            force_as_block=True,
            constraints=[LimitGradeNumberConstraint(grade=2, max_number=4), LimitGradeNumberConstraint(grade=3, max_number=0)],
        ),
    )
    Rota.constraint_options["balance_shifts_quadratic"] = True
    Rota.build_and_solve(options={"ratio": 0.1})
    assert Rota.results.solver.termination_condition == "infeasible"

def test_constraint_minimum_grades(limit_constraint_rota):
    Rota = limit_constraint_rota
    Rota.shifts = []
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2", "group3"),
            name="weekday_night",
            length=12.5,
            days=days,
            balance_offset=40,
            workers_required=1,
            force_as_block=True,
            constraints=[MinimumGradeNumberConstraint(grade=3, min_number=1)],
        ),
    )
    Rota.constraint_options["balance_shifts_quadratic"] = True
    Rota.build_and_solve(options={"ratio": 0.1})
    Rota.export_rota_to_html("test9")
    assert Rota.results.solver.status == "ok"
    grade_workers = Rota.get_workers_by_grade()
    for grade in grade_workers:
        shift_patterns = [Rota.get_worker_shift_list(w) for w in grade_workers[grade]]
        zipped_lists = list(zip(*shift_patterns))
        limit = 1 if grade > 2 else 0
        for day_shifts in zipped_lists:
            assert day_shifts.count("weekday_night") == limit

def test_constraint_minimum_grades2(limit_constraint_rota):
    Rota = limit_constraint_rota
    Rota.shifts = []
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2", "group3"),
            name="weekday_night",
            length=12.5,
            days=days,
            balance_offset=40,
            workers_required=4,
            force_as_block=True,
            constraints=[MinimumGradeNumberConstraint(grade=3, min_number=3)],
        ),
    )
    Rota.constraint_options["balance_shifts_quadratic"] = True
    Rota.build_and_solve(options={"ratio": 0.1})
    Rota.export_rota_to_html("test9")
    grade_workers = Rota.get_workers_by_grade()
    for grade in grade_workers:
        shift_patterns = [Rota.get_worker_shift_list(w) for w in grade_workers[grade]]
        zipped_lists = list(zip(*shift_patterns))
        limit = 3 if grade > 2 else 1
        for day_shifts in zipped_lists:
            assert day_shifts.count("weekday_night") == limit
    assert Rota.results.solver.status == "ok"

def test_constraint_minimum_grades3(limit_constraint_rota):
    Rota = limit_constraint_rota
    Rota.shifts = []
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2", "group3"),
            name="weekday_night",
            length=12.5,
            days=days,
            balance_offset=40,
            workers_required=4,
            force_as_block=True,
            constraints=[MinimumGradeNumberConstraint(grade=3, min_number=0)],
        ),
    )
    Rota.constraint_options["balance_shifts_quadratic"] = True
    Rota.build_and_solve(options={"ratio": 0.1})
    Rota.export_rota_to_html("test9")
    assert Rota.results.solver.status == "ok"

def test_constraint_minimum_grades_no_valid_worker(limit_constraint_rota):
    Rota = limit_constraint_rota
    Rota.shifts = []
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2", "group3"),
            name="weekday_night",
            length=12.5,
            days=days,
            balance_offset=40,
            workers_required=4,
            force_as_block=True,
            constraints=[MinimumGradeNumberConstraint(grade=4, min_number=1)],
        ),
    )
    Rota.constraint_options["balance_shifts_quadratic"] = True
    with pytest.raises(ValueError):
        Rota.build_and_solve(options={"ratio": 0.1})

def test_constraint_require_remote_site_presence_week(limit_constraint_rota):
    Rota = limit_constraint_rota
    Rota.shifts = []
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2", "group3"),
            name="night",
            length=12.5,
            days=days,
            balance_offset=40,
            workers_required=2,
            force_as_block=True,
            constraints=[
                RequireRemoteSitePresenceConstraint(site="group1", required_number=2)
            ],
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.00})
    Rota.export_rota_to_html("remote1")
    group_workers = Rota.get_workers_by_group()
    for group in group_workers:
        shift_patterns = [Rota.get_worker_shift_list(w) for w in group_workers[group]]
        zipped_lists = list(zip(*shift_patterns))
        limit = 2 if group == "group1" else 0
        for day_shifts in zipped_lists:
            assert day_shifts.count("night") == limit
    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"

def test_constraint_require_remote_site_presence_week2(limit_constraint_rota):
    Rota = limit_constraint_rota
    Rota.constraint_options["minimise_shift_diffs"] = False
    Rota.constraint_options["balance_shifts_quadratic"] = False
    Rota.constraint_options["balance_shifts"] = False
    Rota.constraint_options["balance_nights_across_sites"] = True
    Rota.shifts = []
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2", "group3"),
            name="night",
            length=12.5,
            days=days,
            balance_offset=20,
            workers_required=3,
            force_as_block=True,
            constraints=[
                RequireRemoteSitePresenceConstraint(site="group1", required_number=2)
            ],
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.00})
    Rota.export_rota_to_html("remote2")
    group_workers = Rota.get_workers_by_remote_group()
    for group in group_workers:
        shift_patterns = [Rota.get_worker_shift_list(w) for w in group_workers[group]]
        zipped_lists = list(zip(*shift_patterns))
        if group == "group1":
            for day_shifts in zipped_lists:
                assert day_shifts.count("night") >= 2
    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"

def test_constraint_require_remote_site_presence_week3(limit_constraint_rota):
    Rota = limit_constraint_rota
    Rota.shifts = []
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2", "group3"),
            name="night",
            length=12.5,
            days=days,
            balance_offset=20,
            workers_required=3,
            force_as_block=True,
            constraints=[
                RequireRemoteSitePresenceConstraint(site="group2", required_number=1)
            ],
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.00})
    Rota.export_rota_to_html("remote3")
    group_workers = Rota.get_workers_by_remote_group()
    for group in group_workers:
        shift_patterns = [Rota.get_worker_shift_list(w) for w in group_workers[group]]
        zipped_lists = list(zip(*shift_patterns))
        if group == "group2":
            for day_shifts in zipped_lists:
                assert day_shifts.count("night") >= 1
    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"