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

import datetime

from rota.workers import Worker


def weeks_from_list(lst):
    """Yield successive n-sized chunks from a lst."""
    for i in range(0, len(lst), 7):
        yield lst[i : i + 7]


class TestDemoRota:
    def test_nwd(self):
        # Set up rota
        weeks_to_rota = 10
        start_date = datetime.date(2022, 3, 7)

        self.Rota = RotaBuilder(
            start_date,
            weeks_to_rota=weeks_to_rota,
        )

        # Add a few workers
        worker1 = Worker(name="worker1", site="group1", grade=1, nwds=[{"day": "Mon"}])
        worker2 = Worker(name="worker2", site="group1", grade=1, nwds=[{"day": "Fri"}])

        self.Rota.add_workers((worker1, worker2))

        # Add a weekday and weekend shift
        self.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=2,
            ),
        )

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

        self.Rota.export_rota_to_html("nwd1")

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

        for worker in self.Rota.workers:
            shifts = self.Rota.get_worker_shift_list(worker)

            # Convert shift to a string representation
            shifts_string = "".join([i if i != "" else "-" for i in shifts])

            if worker.name == "worker1":
                assert "wwd" not in shifts_string

            if worker.name == "worker2":
                assert "dww" not in shifts_string

    def test_nwd_simple_fail(self):
        # Set up rota
        weeks_to_rota = 10
        start_date = datetime.date(2022, 3, 7)

        self.Rota = RotaBuilder(
            start_date,
            weeks_to_rota=weeks_to_rota,
        )

        # Add a few workers
        worker1 = Worker(name="worker1", site="group1", grade=1, nwds=[{"day": "Mon"}])
        worker2 = Worker(name="worker2", site="group1", grade=1, nwds=[{"day": "Mon"}])

        self.Rota.add_workers((worker1, worker2))

        # Add a weekday and weekend shift
        self.Rota.add_shifts(
            SingleShift(
                sites=("group1",), name="d", length= 12.5, days=days[:5],
                balance_offset=10,
                workers_required=1,
            ),
        )

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

        self.Rota.export_rota_to_html("nwd")

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

    def test_nwd_partial_rota(self):
        # Set up rota
        weeks_to_rota = 10
        start_date = datetime.date(2022, 3, 7)

        self.Rota = RotaBuilder(
            start_date,
            weeks_to_rota=weeks_to_rota,
        )

        d = start_date + datetime.timedelta(weeks=5)
        # Add a few workers
        worker1 = Worker(
            name="worker1",
            site="group1",
            grade=1,
            nwds=[{"day": "Mon", "start_date": start_date, "end_date": d}],
        )
        worker2 = Worker(
            name="worker2",
            site="group1",
            grade=1,
            nwds=[{"day": "Mon", "start_date": d, "end_date": self.Rota.rota_end_date}],
        )

        self.Rota.add_workers((worker1, worker2))

        # Add a weekday and weekend shift
        self.Rota.add_shifts(
            SingleShift(
                sites=("group1",), name="d", length= 12.5, days=days[:5],
                balance_offset=10,
                workers_required=1,
            ),
        )

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

        self.Rota.export_rota_to_html("nwd")

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

        for worker in self.Rota.workers:
            shifts = self.Rota.get_worker_shift_list(worker)

            # Convert shift to a string representation
            shifts_string = "".join([i if i != "" else "-" for i in shifts])

            if worker.name == "worker1":
                for n in range(0,4):
                    assert shifts_string[0 + 7*n] == "-"

            if worker.name == "worker2":
                for n in range(5,9):
                    assert shifts_string[0 + 7*n] == "-"


    def test_nwd_partial_overlay(self):
        # Set up rota
        weeks_to_rota = 10
        start_date = datetime.date(2022, 3, 7)

        self.Rota = RotaBuilder(
            start_date,
            weeks_to_rota=weeks_to_rota,
        )

        d = start_date + datetime.timedelta(weeks=5)
        d2 = start_date + datetime.timedelta(weeks=6)
        # Add a few workers
        worker1 = Worker(
            name="worker1",
            site="group1",
            grade=1,
            nwds=[{"day":"Mon", "start_date":start_date, "end_date":d2}],
        )
        worker2 = Worker(
            name="worker2",
            site="group1",
            grade=1,
            nwds=[{"day":"Mon", "start_date":d, "end_date":self.Rota.rota_end_date}],
        )

        self.Rota.add_workers((worker1, worker2))

        # Add a weekday and weekend shift
        self.Rota.add_shifts(
            SingleShift(
                sites=("group1",), name="d", length= 12.5, days=days[:5],
                balance_offset=10,
                workers_required=1,
            ),
        )

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

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

    def test_nwd_force_as_block(self):
        # Set up rota
        weeks_to_rota = 10
        start_date = datetime.date(2022, 3, 7)

        self.Rota = RotaBuilder(
            start_date,
            weeks_to_rota=weeks_to_rota,
        )

        d = start_date + datetime.timedelta(weeks=4)
        worker1 = Worker(
            name="worker1",
            site="group1",
            grade=1,
            nwds=[{"day": "Mon", "start_date":start_date, "end_date":d}],
        )
        worker2 = Worker(name="worker2", site="group1", grade=1,
        )
        worker3 = Worker(name="worker3", site="group1", grade=1,
        )
        worker4 = Worker(name="worker4", site="group1", grade=1,
        )

        self.Rota.add_workers((worker1, worker2, worker3, worker4))

        # Add a weekday and weekend shift
        self.Rota.add_shifts(
            SingleShift(
                sites=("group1",), name="d", length= 12.5, days=days[:5],
                balance_offset=10,
                workers_required=2,
                force_as_block_unless_nwd=True,
                assign_as_block=True,
            ),
        )

        self.Rota.build_and_solve(options={"ratio": 0.00})
        self.Rota.export_rota_to_html("nwd")

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

        for worker in self.Rota.workers:
            shifts = self.Rota.get_worker_shift_list(worker)

            # Convert shift to a string representation
            shifts_string = "".join([i if i != "" else "-" for i in shifts])

            for week in weeks_from_list(shifts_string):
                assert week in ("-------", "ddddd--")

            assert shifts_string.count("d") == 25

    def test_nwd_force_as_block_force_split(self):
        # Set up rota
        weeks_to_rota = 10
        start_date = datetime.date(2022, 3, 7)

        self.Rota = RotaBuilder(
            start_date,
            weeks_to_rota=weeks_to_rota,
        )

        d = start_date + datetime.timedelta(weeks=5)
        worker1 = Worker(
            name="worker1",
            site="group1",
            grade=1,
            nwds=[{"day": "Mon", "start_date":start_date, "end_date":d}],
        )
        worker2 = Worker(name="worker2", site="group1", grade=1,
        )
        worker3 = Worker(name="worker3", site="group1", grade=1,
        )
        worker4 = Worker(name="worker4", site="group1", grade=1,
        )

        self.Rota.add_workers((worker1, worker2, worker3, worker4))

        # Add a weekday and weekend shift
        self.Rota.add_shifts(
            SingleShift(
                sites=("group1",), name="d", length= 12.5, days=days[:5],
                balance_offset=10,
                workers_required=2,
                force_as_block_unless_nwd=True,
                assign_as_block=True,
            ),
            SingleShift(
                sites=("group1",), name="w", length= 12.5, days=days[5:],
                balance_offset=10,
                workers_required=4,
                force_as_block_unless_nwd=True,
                assign_as_block=True,
            ),
        )

        self.Rota.build_and_solve(options={"ratio": 0.00})
        self.Rota.export_rota_to_html("nwd_block_force_split")

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

        for worker in self.Rota.workers:
            shifts = self.Rota.get_worker_shift_list(worker)

            # Convert shift to a string representation
            shifts_string = "".join([i if i != "" else "-" for i in shifts])

            for week in weeks_from_list(shifts_string[7 * 5 :]):
                assert week in ("-----ww", "dddddww")

            if worker.name == "worker1":
                assert shifts_string[: 7 * 5].count("ww-") == 4

            assert shifts_string.count("d") == 25
            assert shifts_string.count("w") == 20

    def test_nwd_testing(self):
        # Set up rota
        weeks_to_rota = 10
        start_date = datetime.date(2022, 3, 7)

        self.Rota = RotaBuilder(
            start_date,
            weeks_to_rota=weeks_to_rota,
        )

        d = start_date + datetime.timedelta(weeks=5)
        worker1 = Worker(name="worker1", site="group1", grade=1,
            nwds=[
                {"day": "Mon", "start_date":d, "end_date":self.Rota.rota_end_date},
                {"day": "Tue", "start_date":d, "end_date":self.Rota.rota_end_date},
                {"day": "Wed", "start_date":d, "end_date":self.Rota.rota_end_date},
                {"day": "Thu", "start_date":d, "end_date":self.Rota.rota_end_date},
                {"day": "Fri", "start_date":d, "end_date":self.Rota.rota_end_date},
            ],
        )
        worker2 = Worker(name="worker2", site="group1", grade=1,
        )

        self.Rota.add_workers((worker1, worker2))

        # Add a weekday and weekend shift
        self.Rota.add_shifts(
            SingleShift(
                sites=("group1",), name="d", length= 12.5, days=days[:5],
                balance_offset=10,
                workers_required=1,
                force_as_block_unless_nwd=True,
                assign_as_block=True,
            ),
            SingleShift(
                sites=("group1",), name="w", length= 12.5, days=days[5:],
                balance_offset=10,
                workers_required=2,
                force_as_block_unless_nwd=True,
                assign_as_block=True,
            ),
        )

        self.Rota.build_and_solve(options={"ratio": 0.00})
        self.Rota.export_rota_to_html("nwd")

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

        for worker in self.Rota.workers:
            shifts = self.Rota.get_worker_shift_list(worker)

            # Convert shift to a string representation
            shifts_string = "".join([i if i != "" else "-" for i in shifts])

            for week in weeks_from_list(shifts_string[7 * 5 :]):
                if worker.name == "worker1":
                    assert week == "-----ww"
                else:
                    assert week == "dddddww"

            for week in weeks_from_list(shifts_string[: 7 * 5]):
                if worker.name == "worker1":
                    assert week == "dddddww"
                else:
                    assert week == "-----ww"