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

import datetime
from rota.workers import Worker


def setup_basic_rota():
    weeks_to_rota = 10
    start_date = datetime.date(2022, 3, 7)

    return RotaBuilder(
        start_date,
        weeks_to_rota=weeks_to_rota,
    )


class TestWorkers:
    def test_no_workers(self):
        """Check error raised if attempt to solve with no workers"""
        Rota = setup_basic_rota()

        Rota.add_shifts(
            SingleShift(
                sites=("group1",),
                name="d",
                length=12.5,
                days=days[:5],
                balance_offset=10,
                workers_required=1,
            ),
            SingleShift(
                sites=("group1",),
                name="w",
                length=12.5,
                days=days[5:],
                workers_required=1,
            ),
        )

        with pytest.raises(NoWorkers):
            Rota.build_and_solve(solve=False)

    def test_workers_duplicate_name(self):

        Rota = setup_basic_rota()

        workers = []
        workers.append(Worker(name="worker1", site="group1", grade=1, fte=100))
        workers.append(Worker(name="worker1", site="group1", grade=1, fte=100))
        workers.append(Worker(name="worker3", site="group2", grade=1, fte=50))
        workers.append(Worker(name="worker4", site="group2", grade=1, fte=50))

        Rota.add_workers(workers)

        Rota.add_shifts(
            SingleShift(
                sites=("group1",),
                name="d",
                length=12.5,
                days=days[:5],
                balance_offset=10,
                workers_required=1,
            ),
        )

        with pytest.raises(ValueError):
            Rota.build_and_solve(solve=False)

    def test_worker_ftes(self):

        weeks_to_rota = 10
        start_date = datetime.date(2022, 3, 7)

        Rota = RotaBuilder(
            start_date,
            weeks_to_rota=weeks_to_rota,
        )

        worker1_end_date = start_date + datetime.timedelta(weeks=weeks_to_rota / 2)

        worker2_start_date = start_date + datetime.timedelta(weeks=weeks_to_rota / 2)

        worker3_start_date = start_date + datetime.timedelta(weeks=weeks_to_rota / 4)
        worker3_end_date = start_date + datetime.timedelta(weeks=3 * weeks_to_rota / 4)

        workers = []
        workers.append(
            Worker(
                name="worker1",
                site="group1",
                grade=1,
                fte=100,
                end_date=worker1_end_date,
            )
        )
        workers.append(
            Worker(
                name="worker2",
                site="group1",
                grade=1,
                fte=100,
                start_date=worker2_start_date,
            )
        )
        workers.append(
            Worker(
                name="worker3",
                site="group2",
                grade=1,
                fte=100,
                start_date=worker3_start_date,
                end_date=worker3_end_date,
            )
        )
        workers.append(Worker(name="worker4", site="group2", grade=1, fte=50))
        workers.append(
            Worker(name="worker5", site="group2", grade=1, fte=50, nwd=["Mon", "Tue"])
        )
        workers.append(
            Worker(
                name="worker6",
                site="group3",
                grade=1,
                fte=100,
                start_date=start_date + datetime.timedelta(weeks=weeks_to_rota),
            )
        )
        workers.append(
            Worker(name="worker7", site="group3", grade=1, fte=50, end_date=start_date)
        )

        Rota.add_workers(workers)

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"),
                name="d",
                length=12.5,
                days=days[:5],
                balance_offset=10,
                workers_required=1,
            ),
        )

        Rota.build_and_solve()

        # All workers should have an adjusted fte of 50 %
        for worker in Rota.get_workers():

            # Except those in group3
            if worker.site == "group3":
                assert worker.fte_adj == 0
                assert Rota.get_worker_shift_list(worker).count("d") == 0
                continue

            assert worker.fte_adj == 50

            if worker.fte == 100:
                assert worker.proportion_rota_to_work == 0.5
                assert worker.days_to_work == 35

            assert Rota.get_worker_shift_list(worker).count("d") == 10

            if worker.name == "worker1":
                assert Rota.get_worker_shift_list_string(worker).endswith(
                    "-" * (Rota.rota_days_length // 2)
                )
            if worker.name == "worker2":
                assert Rota.get_worker_shift_list_string(worker).startswith(
                    "-" * (Rota.rota_days_length // 2)
                )
            if worker.name == "worker3":
                assert Rota.get_worker_shift_list_string(worker).startswith(
                    "-" * (Rota.rota_days_length // 4)
                )
                assert Rota.get_worker_shift_list_string(worker).endswith(
                    "-" * (Rota.rota_days_length // 4)
                )

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

        Rota.export_rota_to_html("worker_ftes")

    def test_worker_oops(self):
        weeks_to_rota = 10
        start_date = datetime.date(2022, 3, 7)

        Rota = RotaBuilder(
            start_date,
            weeks_to_rota=weeks_to_rota,
        )

        worker1_oop = [
            {
                "start_date": start_date,
                "end_date": start_date + datetime.timedelta(weeks=weeks_to_rota / 2),
                "reason": "First half OOP",
            }
        ]

        worker2_oop = [
            {
                "start_date": start_date + datetime.timedelta(weeks=weeks_to_rota / 2),
                "end_date": Rota.rota_end_date,
                "reason": "First half OOP",
            }
        ]

        worker3_start_date = start_date + datetime.timedelta(weeks=weeks_to_rota / 4)
        worker3_oop = [
            {
                "start_date": start_date + datetime.timedelta(weeks=weeks_to_rota / 2),
                "end_date": start_date
                + datetime.timedelta(weeks=3 * weeks_to_rota / 4),
                "reason": "Quarter oop",
            }
        ]

        worker4_oop = [
            {
                "start_date": start_date,
                "end_date": start_date + datetime.timedelta(weeks=weeks_to_rota / 4),
                "reason": "Quarter oop",
            },
            {
                "start_date": start_date + datetime.timedelta(weeks=weeks_to_rota / 2),
                "end_date": start_date
                + datetime.timedelta(weeks=weeks_to_rota / 2)
                + datetime.timedelta(weeks=weeks_to_rota / 4),
                "reason": "Quarter oop 2",
            },
        ]

        workers = []
        workers.append(
            Worker(name="worker1", site="group1", grade=1, fte=100, oop=worker1_oop)
        )
        workers.append(
            Worker(name="worker2", site="group1", grade=1, fte=100, oop=worker2_oop)
        )
        workers.append(
            Worker(
                name="worker3",
                site="group2",
                grade=1,
                fte=100,
                start_date=worker3_start_date,
                oop=worker3_oop,
            )
        )
        workers.append(Worker(name="worker4", site="group2", grade=1, oop=worker4_oop))
        workers.append(
            Worker(name="worker5", site="group2", grade=1, fte=50, nwd=["Mon", "Tue"])
        )

        worker6_oop = [
            {
                "start_date": start_date,
                "end_date": start_date + datetime.timedelta(weeks=weeks_to_rota / 6),
                "reason": "oop",
            },
            {
                "start_date": start_date + datetime.timedelta(weeks=weeks_to_rota / 2),
                "end_date": start_date
                + datetime.timedelta(weeks=weeks_to_rota / 2)
                + datetime.timedelta(weeks=weeks_to_rota / 3),
                "reason": " oop 2",
            },
        ]

        worker7_oop = [
            {
                "start_date": start_date - datetime.timedelta(weeks=1),
                "end_date": Rota.rota_end_date,
                "reason": "oop",
            },
        ]

        workers.append(
            Worker(name="worker6", site="group2", grade=1, fte=100, oop=worker6_oop)
        )
        workers.append(
            Worker(name="worker7", site="group3", grade=1, fte=100, oop=worker7_oop)
        )

        Rota.add_workers(workers)

        Rota.add_shifts(
            SingleShift(
                sites=("group1", "group2"),
                name="d",
                length=12.5,
                days=days[:5],
                balance_offset=10,
                workers_required=1,
            ),
        )

        Rota.build_and_solve()

        # All workers should have an adjusted fte of 50 %
        for worker in Rota.get_workers():
            if worker.site == "group3":
                assert worker.fte_adj == 0
                continue

            assert worker.fte_adj == pytest.approx(50, 2)

            if worker.fte == 100:
                assert worker.proportion_rota_to_work == pytest.approx(0.5, 0.2)
                assert worker.days_to_work == pytest.approx(35, 1)

            assert Rota.get_worker_shift_list(worker).count("d") in (8, 9)

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

        Rota.export_rota_to_html("worker_oops")