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

from rota_generator.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 TestWorkerRequests:
    def test_preferences_not_to_work(self):
        Rota = generate_basic_rota()

        Rota.add_worker(
            Worker(
                name="worker3", site="group1", grade=1,
                pref_not_to_work=[{"date": d} for d in date_generator(datetime.date(2022, 3, 7), 5)],
            ),
        )

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

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

        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)

            shift_string = Rota.get_worker_shift_list_string(worker)

            if worker.name == "worker3":
                assert shift_string.startswith("-----")
            else:
                assert shift_string.startswith("aaaaa")

    def test_preferences_not_to_work2(self):
        """_summary_"""
        Rota = generate_basic_rota()

        Rota.add_worker(
            Worker(
                name="worker3", site="group1", grade=1,
                pref_not_to_work=[{"date": d} for d in date_generator(datetime.date(2022, 3, 7), 55)],
            ),
        )

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

        Rota.build_and_solve(options={"ratio": 0.000})
        Rota.export_rota_to_html("preferences_not_to_work2", timestamp_filename=False)

        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)

            shift_string = Rota.get_worker_shift_list_string(worker)

            if worker.name == "worker3":
                assert shift_string.endswith("a" * 15)

    def test_preferences_not_to_work_all_shifts(self):
        """_summary_"""
        Rota = generate_basic_rota()

        Rota.add_worker(
            Worker(
                name="worker3", site="group1", grade=1,
                pref_not_to_work=[{"date": d} for d in date_generator(datetime.date(2022, 3, 7), 70)],
            ),
        )

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

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

            # shift_string = Rota.get_worker_shift_list_string(worker)

            # if worker.name == "worker3":
            #    assert shift_string.startswith("-----")
            # else:
            #    assert shift_string.startswith("aaaaa")

    def test_unavailable_to_work(self):
        """_summary_"""
        Rota = generate_basic_rota()

        Rota.add_worker(
            Worker(
                name="worker3", site="group1", grade=1,
                not_available_to_work=(
                    [
                        {"date": i, "reason": "NOT AROUND"}
                        for i in list(date_generator(datetime.date(2022, 3, 7), 23))
                    ]
                ),
            ),
        )

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

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

        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 (approx(46), approx(47))

            # shift_string = Rota.get_worker_shift_list_string(worker)

            # if worker.name == "worker3":
            #    assert shift_string.startswith("-----")
            # else:
            #    assert shift_string.startswith("aaaaa")

    def test_unavailable_to_work2(self):
        """_summary_"""
        Rota = generate_basic_rota()

        Rota.add_worker(
            Worker(
                name="worker3", site="group1", grade=1,
                not_available_to_work=(
                    [
                        {"date": i, "reason": "NOT AROUND"}
                        for i in list(date_generator(datetime.date(2022, 3, 7), 25))
                    ]
                ),
            ),
        )

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

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

        assert Rota.results.solver.status in ("warning", "error")
        assert Rota.results.solver.termination_condition == "infeasible"

    def test_unavailable_to_work3(self):
        """_summary_"""
        Rota = generate_basic_rota()

        Rota.constraint_options_model.balance_bank_holidays = False

        Rota.add_worker(
            Worker(
                name="worker3", site="group1", grade=1,
                not_available_to_work=(
                    [
                        {"date":i, "reason":"NOT AROUND"}
                        for i in list(date_generator(datetime.date(2022, 3, 7), 23))
                    ]
                ),
                pref_not_to_work=[{"date":datetime.date(2022, 5, 15)}],
            ),
        )

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

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

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

        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)

            shift_string = Rota.get_worker_shift_list_string(worker)

            if worker.name == "worker3":
               assert shift_string.startswith("-"*23)
               assert shift_string.endswith("-")

    def test_unavailable_to_work2(self):
        """_summary_"""
        Rota = generate_basic_rota()

        Rota.add_workers(
            [Worker(
                name="worker3", site="group1", grade=1,
                not_available_to_work=(
                    [
                        {"date": i, "reason": "NOT AROUND"}
                        for i in list(date_generator(datetime.date(2022, 3, 7), 25))
                    ]
                ),
            ),
            Worker(
                name="worker4", site="group1", grade=1,
                not_available_to_work=(
                    [
                        {"date": i, "reason": "NOT AROUND"}
                        for i in list(date_generator(datetime.date(2022, 3, 7), 25))
                    ]
                ),
            ),]
        )

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

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

        assert Rota.results.solver.status in ("warning", "error")
        assert Rota.results.solver.termination_condition == "infeasible"

    def test_work_requests(self):
        """_summary_"""
        Rota = generate_basic_rota()
        Rota.constraint_options_model.balance_bank_holidays = False

        Rota.add_workers(
            [Worker(
                name="worker3", site="group1", grade=1,
                work_requests=(
                    [
                        {"date": i, "shift": "a"}
                        for i in list(date_generator(datetime.date(2022, 3, 7), 35))
                    ]
                ),
            ),
            Worker(
                name="worker4", site="group1", grade=1,
                work_requests=(
                    [
                        {"date": i, "shift": "a"}
                        for i in list(date_generator(datetime.date(2022, 3, 7), 35))
                    ]
                ),
            ),]
        )

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

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

        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"] == 35

            shift_string = Rota.get_worker_shift_list_string(worker)

            if worker.name in ("worker3", "worker4"):
               assert shift_string.startswith("a"*35)

    def test_work_requests_invalid(self):
        """_summary_"""
        Rota = generate_basic_rota()
        Rota.constraint_options_model.balance_bank_holidays = False

        Rota.add_worker(
            Worker(
                name="worker3", site="group1", grade=1,
                work_requests=(
                    [
                        {"date": i, "shift": "b"}
                        for i in list(date_generator(datetime.date(2022, 3, 7), 35))
                    ]
                ),
            ),
        )

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


        with pytest.raises(InvalidShift):
            Rota.build_and_solve(options={"ratio": 0.000})
        
        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"), name="c", length= 12.5, days=days,
                workers_required=1,
                balance_offset=1,
                force_as_block=False,
            ),
        )

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

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

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

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

        Rota.add_worker(
            Worker(
                name="worker4", site="group1", grade=1,
                work_requests=(
                    [
                        {"date": i, "shift": "d"}
                        for i in list(date_generator(datetime.date(2022, 3, 7), 35))
                    ]
                ),
            ),
        )

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

    def test_work_requests_generic(self):
        """_summary_"""
        Rota = generate_basic_rota()
        Rota.constraint_options_model.balance_bank_holidays = False

        Rota.add_workers(
            [Worker(
                name="worker3", site="group1", grade=1,
                work_requests=(
                    [
                        {"date": i, "shift": "*"}
                        for i in list(date_generator(datetime.date(2022, 3, 7), 35))
                    ]
                ),
            ),
            Worker(
                name="worker4", site="group1", grade=1,
                work_requests=(
                    [
                        {"date": i, "shift": "*"}
                        for i in list(date_generator(datetime.date(2022, 3, 7), 35))
                    ]
                ),
            ),]
        )

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

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

        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"] == 35

            shift_string = Rota.get_worker_shift_list_string(worker)

            if worker.name in ("worker3", "worker4"):
               assert shift_string.startswith("a"*35)

    def test_work_requests_generic2(self):
        """_summary_"""
        Rota = generate_basic_rota()
        Rota.constraint_options_model.balance_bank_holidays = False

        Rota.add_workers(
            [Worker(
                name="worker3", site="group1", grade=1,
                work_requests=(
                    [
                        {"date": i, "shift": "*"}
                        for i in list(date_generator(datetime.date(2022, 3, 7), 20))
                    ]
                ),
            ),
            ]
        )

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"), name="a", length= 12.5, days=days[3:],
                workers_required=2,
                balance_offset=1,
                force_as_block=False,
            ),
            SingleShift(
                sites=("group1", "group2"), name="b", length= 12.5, days=days[:3],
                workers_required=1,
                balance_offset=1,
                force_as_block=False,
            ),
        )

        Rota.build_and_solve(options={"ratio": 0.1})
        Rota.export_rota_to_html("work_requests")

        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 (26, 27)
            assert worker_shifts["b"] == approx(10)

            shift_string = Rota.get_worker_shift_list_string(worker)

            if worker.name in ("worker3",):
               assert "-" not in shift_string[:20]