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

import datetime

from rota_generator.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_balance_weekend(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,
        )
        self.Rota.constraint_options["balance_weekends"] = True

        # Add a few workers
        worker1 = Worker(name="worker1", site="group1", grade=1)
        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="a", length= 12.5, days=days,
                balance_offset=1,
                workers_required=1,
            ),
            SingleShift(
                sites=("group1",), name="b", length= 12.5, days=days[5:],
                balance_offset=0,
                workers_required=1
            ),
        )

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

        self.Rota.export_rota_to_html("balance_weekend", folder="tests")

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

            assert shifts_string.count("b") == 5
            assert shifts_string.count("a") in (17,18)

            weekend_count = 0
            for week in weeks_from_list(shifts_string):
                if week[5:] == "--":
                    continue
                weekend_count = weekend_count + 1

            assert weekend_count == 5 == self.Rota.model.worker_weekend_count[worker.id].value

    def test_balance_weekend2(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,
        )
        self.Rota.constraint_options["balance_weekends"] = True

        # Add a few workers
        worker1 = Worker(name="worker1", site="group1", grade=1)
        worker2 = Worker(name="worker2", site="group1", grade=1)
        worker3 = Worker(name="worker3", site="group1", grade=1)
        worker4 = Worker(name="worker4", site="group1", grade=1)
        worker5 = Worker(name="worker5", site="group1", grade=1)

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

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

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

        self.Rota.export_rota_to_html("balance_weekend2")

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

            assert shifts_string.count("b") == 8
            assert shifts_string.count("a") == 14

            weekend_count = 0
            for week in weeks_from_list(shifts_string):
                if week[5:] == "--":
                    continue
                weekend_count = weekend_count + 1

            assert weekend_count == 6 == self.Rota.model.worker_weekend_count[worker.id].value

    def test_post_mon_tue(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,
        )
        self.Rota.constraint_options["prevent_monday_and_tuesday_after_full_weekends"] = ["group1",]

        # Add a few workers
        worker1 = Worker(name="worker1", site="group1", grade=1)
        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,
            ),
            SingleShift(
                sites=("group1",), name="w", length= 12.5, days=days[5:],
                workers_required=1
            ),
        )

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

        #self.Rota.export_rota_to_html("weekend_padding")

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

            assert "wd" not in shifts_string
            assert "w-d" not in shifts_string

    def test_post_mon_tue_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,
        )
        self.Rota.constraint_options["prevent_monday_and_tuesday_after_full_weekends"] = ["group1",]

        # Add a few workers
        worker1 = Worker(name="worker1", site="group1", grade=1)
        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,
            ),
            SingleShift(
                sites=("group1",), name="c", length= 12.5, days=days[1],
                balance_offset=1,
                workers_required=1,
            ),
            SingleShift(
                sites=("group1",), name="w", length= 12.5, days=days[5:],
                workers_required=1
            ),
        )

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

        assert self.Rota.results.solver.status in ("warning", "error")

    def test_post_mon(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,
        )
        self.Rota.constraint_options["prevent_monday_after_full_weekends"] = ["group1",]

        # Add a few workers
        worker1 = Worker(name="worker1", site="group1", grade=1)
        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=1,
                workers_required=1,
            ),
            SingleShift(
                sites=("group1",), name="c", length= 12.5, days=days[1],
                balance_offset=1,
                workers_required=1,
            ),
            SingleShift(
                sites=("group1",), name="w", length= 12.5, days=days[5:],
                workers_required=1
            ),
        )

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

        self.Rota.export_rota_to_html("weekend_padding")

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

            assert "wd" not in shifts_string

    def test_post_mon_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,
        )
        self.Rota.constraint_options["prevent_monday_after_full_weekends"] = ["group1",]

        # Add a few workers
        worker1 = Worker(name="worker1", site="group1", grade=1)
        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=1,
                workers_required=1,
            ),
            SingleShift(
                sites=("group1",), name="c", length= 12.5, days=days[0],
                balance_offset=1,
                workers_required=1,
            ),
            SingleShift(
                sites=("group1",), name="w", length= 12.5, days=days[5:],
                workers_required=1
            ),
        )

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

        self.Rota.export_rota_to_html("weekend_padding")

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

    def test_post_mon_prevent_split_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,
        )
        self.Rota.constraint_options["prevent_monday_after_full_weekends"] = ["group1",]

        # Add a few workers
        worker1 = Worker(name="worker1", site="group1", grade=1)
        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=1,
                workers_required=1,
            ),
            SingleShift(
                sites=("group1",), name="c", length= 12.5, days=days[0],
                balance_offset=1,
                workers_required=1,
            ),
            SingleShift(
                sites=("group1",), name="w", length= 12.5, days=days[5:],
                workers_required=1,
                force_as_block=True
            ),
        )

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

        self.Rota.export_rota_to_html("weekend_padding")

        assert self.Rota.results.solver.status in ("warning", "error")