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

from rota.workers import Worker

import itertools


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,
    )

    # Add a few workers
    Rota.add_workers(
        [
            Worker(name="worker1", site="group1", grade=1),
            Worker(name="worker2", site="group1", grade=1),
        ]
    )

    return Rota


def date_generator(from_date, days):
    n = 0
    while True:
        yield from_date

        n = n + 1
        if n >= days:
            break
        from_date = from_date + datetime.timedelta(days=1)


class TestShifts:
    def test_duplicate_shifts(self):
        Rota = generate_basic_rota()

        Rota.add_worker(
            Worker(
                name="worker3", site="group1", grade=1,
            ),
        )

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"), name="a", length= 12.5, days=days,
            ),
        )

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

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

        Rota.add_shifts(
            SingleShift(
                sites=("group1",), name="a", length= 12.5, days=days,
            ),
        )

        with pytest.raises(InvalidShift):
            Rota.build_and_solve(options={"ratio": 0.000})


class TestShiftConstraints:

    def test_max_shifts(self):
        Rota = generate_basic_rota()

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"), name="a", length= 12.5, days=days,
                constraint=[{"name": "max_shifts_per_week", "options": 4}],
            ),
        )

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

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

    def test_max_shifts_fail(self):
        Rota = generate_basic_rota()

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"), name="a", length= 12.5, days=days,
                constraint=[{"name": "max_shifts_per_week", "options": 3}],
            ),
        )

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

        assert Rota.results.solver.termination_condition == "infeasible"

        Rota.add_worker(
            Worker(
                name="worker3", site="group1", grade=1,
            ),
        )

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

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


    def test_max_shifts_per_week_block(self):
        Rota = generate_basic_rota()

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"), name="a", length= 12.5, days=days,
                constraint=[{"name": "max_shifts_per_week_block", "options": {"weeks": 2}}],
            ),
        )

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

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



    def test_max_shifts_per_week_block_fail(self):
        Rota = generate_basic_rota()

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"), name="a", length= 12.5, days=days,
                constraint=[{"name": "max_shifts_per_week_block", "options": {"weeks": 3}}],
            ),
        )

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

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

    def test_max_shifts_per_week_block_shift_number(self):
        Rota = generate_basic_rota()

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"), name="a", length= 12.5, days=days,
                constraint=[{"name": "max_shifts_per_week_block", "options": {"weeks": 2, "shift_number": 7}}],
            ),
        )

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

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

    def test_max_shifts_per_week_block_shift_number_fail(self):
        Rota = generate_basic_rota()

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"), name="a", length= 12.5, days=days,
                constraint=[{"name": "max_shifts_per_week_block", "options": {"weeks": 2, "shift_number": 6}}],
            ),
        )

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

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

    def test_pre_shift_constraint(self):
        Rota = generate_basic_rota()

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"), name="a", length= 12.5, days=days[0],
                constraint=[{"name": "pre", "options": 1}],
            ),
            SingleShift(
                sites=("group1", "group2"), name="b", length= 12.5, days=days[1],
                constraint=[{"name": "pre", "options": 1}],
            ),
            SingleShift(
                sites=("group1", "group2"), name="c", length= 12.5, days=days[2],
                constraint=[{"name": "pre", "options": 1}],
            ),
            SingleShift(
                sites=("group1", "group2"), name="d", length= 12.5, days=days[3],
                constraint=[{"name": "pre", "options": 1}],
            ),
            SingleShift(
                sites=("group1", "group2"), name="e", length= 12.5, days=days[4],
                constraint=[{"name": "pre", "options": 1}],
            ),
            SingleShift(
                sites=("group1", "group2"), name="f", length= 12.5, days=days[5],
                constraint=[{"name": "pre", "options": 1}],
            ),
        )

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

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

        for worker in Rota.get_workers():
            shift_string = Rota.get_worker_shift_list_string(worker)

            print(shift_string)
            for s in ["ab", "bc", "cd", "de", "ef", "fg"]:
                assert s not in shift_string

            
        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"), name="g", length= 12.5, days=days[6],
                constraint=[{"name": "pre", "options": 2}],
            ),
        )

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

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

    def test_pre_shift_constraint2(self):
        Rota = generate_basic_rota()

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"), name="a", length= 12.5, days=days[0],
                constraint=[{"name": "pre", "options": 2}],
                workers_required=2
            ),
            SingleShift(
                sites=("group1", "group2"), name="b", length= 12.5, days=days[4],
                constraint=[{"name": "pre", "options": 3}],
                workers_required=2
            ),

        )

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

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

    def test_pre_shift_constraint3(self):
        Rota = generate_basic_rota()

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"), name="a", length= 12.5, days=days[0],
                constraint=[{"name": "pre", "options": 3}],
                workers_required=2
            ),
            SingleShift(
                sites=("group1", "group2"), name="b", length= 12.5, days=days[4],
                constraint=[{"name": "pre", "options": 3}],
                workers_required=2
            ),
        )

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

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

    def test_post_shift_constraint(self):
        Rota = generate_basic_rota()

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"), name="a", length= 12.5, days=days[0],
                constraint=[{"name": "post", "options": 3}],
                workers_required=2
            ),
            SingleShift(
                sites=("group1", "group2"), name="b", length= 12.5, days=days[4],
                constraint=[{"name": "post", "options": 2}],
                workers_required=2
            ),
        )

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

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

    def test_post_shift_constraint2(self):
        Rota = generate_basic_rota()

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"), name="a", length= 12.5, days=days[0],
                constraint=[{"name": "post", "options": 3}],
                workers_required=2
            ),
            SingleShift(
                sites=("group1", "group2"), name="b", length= 12.5, days=days[4],
                constraint=[{"name": "post", "options": 3}],
                workers_required=2
            ),
        )

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

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