import datetime
import pytest
from rota_generator.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, days, MinSummedGradeByShiftsPerDayConstraint

from rota_generator.workers import Worker

import itertools


def generate_basic_rota(weeks_to_rota=9):
    start_date = datetime.date(2022, 3, 7)

    Rota = RotaBuilder(
        start_date,
        weeks_to_rota=weeks_to_rota,
    )

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

    return Rota


def test_summed_shift_grade_constraints():
    Rota = generate_basic_rota(5)

    Rota.add_shifts(
        SingleShift(
            sites=("group1",), name="a", length=12.5, days=days,
            workers_required=1,
            force_as_block=False,
        ),
        SingleShift(
            sites=("group1",), name="b", length=12.5, days=days,
            workers_required=1,
            force_as_block=False,
        ),
    )

    # Add summed shift grade constraints
    Rota.constraint_options_model.min_summed_grade_by_shifts_per_day.append(
        MinSummedGradeByShiftsPerDayConstraint(
            shifts=["a", "b"],
            min_grade_sum=4,
        ),
    )

    Rota.build_and_solve(options={"ratio": 0.000})

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

    Rota.export_rota_to_html("summed_shift_grade_constraints", folder="tests")

    a_grades = []
    b_grades = []

    for shift, shift_timetable in Rota.get_shift_timetable_by_week().items():
        for week, week_timetable in shift_timetable.items():
            for day, workers in week_timetable.items():
                if shift == "a":
                    for worker in workers:
                        a_grades.append(worker.grade)
                elif shift == "b":
                    for worker in workers:
                        b_grades.append(worker.grade)

    # Check that the minimum summed grade is respected
    summed_grades = [a + b for a, b in zip(a_grades, b_grades)]
    for summed_grade in summed_grades:
        assert summed_grade >= 4, f"Summed grade {summed_grade} is less than 4"

    Rota.constraint_options_model.min_summed_grade_by_shifts_per_day = [
        MinSummedGradeByShiftsPerDayConstraint(
            shifts=["a", "b"],
            min_grade_sum=5,
        ),
    ]

    Rota.build_and_solve(options={"ratio": 0.000})
    assert Rota.results.solver.status == "ok"

    for worker in Rota.get_workers():
        if worker.grade == 1:
            assert Rota.get_worker_shift_list(worker).count("a") == 0, f"Worker {worker.name} with grade 1 should not have shift 'a'"
            assert Rota.get_worker_shift_list(worker).count("b") == 0, f"Worker {worker.name} with grade 1 should not have shift 'b'"