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]

def setup_rota(weeks_to_rota=10, start_date=None):
    if start_date is None:
        start_date = datetime.date(2022, 3, 7)
    return RotaBuilder(start_date, weeks_to_rota=weeks_to_rota)

def test_nwd():
    Rota = setup_rota()
    worker1 = Worker(name="worker1", site="group1", grade=1, nwds=[{"day": "Mon"}])
    worker2 = Worker(name="worker2", site="group1", grade=1, nwds=[{"day": "Fri"}])
    Rota.add_workers((worker1, worker2))
    Rota.add_shifts(
        SingleShift(
            sites=("group1",), name="d", length=12.5, days=days[:5], workers_required=1,
        ),
        SingleShift(
            sites=("group1",), name="w", length=12.5, days=days[5:], workers_required=2,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.10})
    Rota.export_rota_to_html("nwd1", folder="tests")
    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"
    for worker in Rota.workers:
        shifts = Rota.get_worker_shift_list(worker)
        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():
    Rota = setup_rota()
    worker1 = Worker(name="worker1", site="group1", grade=1, nwds=[{"day": "Mon"}])
    worker2 = Worker(name="worker2", site="group1", grade=1, nwds=[{"day": "Mon"}])
    Rota.add_workers((worker1, worker2))
    Rota.add_shifts(
        SingleShift(
            sites=("group1",), name="d", length=12.5, days=days[:5], balance_offset=10, workers_required=1,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.00})
    Rota.export_rota_to_html("nwd", folder="tests")
    assert Rota.results.solver.status in ("warning", "error")

def test_nwd_partial_rota():
    Rota = setup_rota()
    start_date = Rota.start_date
    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,
        nwds=[{"day": "Mon", "start_date": d, "end_date": Rota.rota_end_date}],
    )
    Rota.add_workers((worker1, worker2))
    Rota.add_shifts(
        SingleShift(
            sites=("group1",), name="d", length=12.5, days=days[:5], balance_offset=10, workers_required=1,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.00})
    Rota.export_rota_to_html("nwd", folder="tests")
    assert Rota.results.solver.status == "ok"
    for worker in Rota.workers:
        shifts = Rota.get_worker_shift_list(worker)
        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():
    Rota = setup_rota()
    start_date = Rota.start_date
    d = start_date + datetime.timedelta(weeks=5)
    d2 = start_date + datetime.timedelta(weeks=6)
    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": Rota.rota_end_date}],
    )
    Rota.add_workers((worker1, worker2))
    Rota.add_shifts(
        SingleShift(
            sites=("group1",), name="d", length=12.5, days=days[:5], balance_offset=10, workers_required=1,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.00})
    assert Rota.results.solver.status in ("warning", "error")

def test_nwd_force_as_block():
    Rota = setup_rota()
    start_date = Rota.start_date
    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)
    Rota.add_workers((worker1, worker2, worker3, worker4))
    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,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.00})
    Rota.export_rota_to_html("nwd", folder="tests")
    assert Rota.results.solver.status == "ok"
    for worker in Rota.workers:
        shifts = Rota.get_worker_shift_list(worker)
        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():
    Rota = setup_rota()
    start_date = Rota.start_date
    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)
    Rota.add_workers((worker1, worker2, worker3, worker4))
    Rota.add_shifts(
        SingleShift(
            sites=("group1",), name="d", length=12.5, days=days[:5],
            workers_required=2, force_as_block_unless_nwd=True, assign_as_block=True,
        ),
        SingleShift(
            sites=("group1",), name="w", length=12.5, days=days[5:],
            workers_required=4, force_as_block_unless_nwd=True, assign_as_block=True,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.00})
    Rota.export_rota_to_html("nwd_block_force_split", folder="tests")

    assert Rota.results.solver.status == "ok"
    for worker in Rota.workers:
        shifts = Rota.get_worker_shift_list(worker)
        shifts_string = "".join([i if i != "" else "-" for i in shifts])
        for week in weeks_from_list(shifts_string[7 * 6 :]):
            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():
    Rota = setup_rota()
    start_date = Rota.start_date
    d = start_date + datetime.timedelta(weeks=5)
    worker1 = Worker(
        name="worker1", site="group1", grade=1,
        nwds=[
            {"day": "Mon", "start_date": d, "end_date": Rota.rota_end_date},
            {"day": "Tue", "start_date": d, "end_date": Rota.rota_end_date},
            {"day": "Wed", "start_date": d, "end_date": Rota.rota_end_date},
            {"day": "Thu", "start_date": d, "end_date": Rota.rota_end_date},
            {"day": "Fri", "start_date": d, "end_date": Rota.rota_end_date},
        ],
    )
    worker2 = Worker(name="worker2", site="group1", grade=1)
    Rota.add_workers((worker1, worker2))
    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,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.00})
    Rota.export_rota_to_html("nwd", folder="tests")
    assert Rota.results.solver.status == "ok"
    for worker in Rota.workers:
        shifts = Rota.get_worker_shift_list(worker)
        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"


def test_force_as_block_unless_nwd_with_subblocks():
    # Test that a Fri/Sat/Sun shift can be split into Fri and Sat/Sun blocks
    # if a worker does not work on Fridays.
    Rota = setup_rota(weeks_to_rota=2)
    start_date = Rota.start_date
    
    # worker1 does not work on Fridays (Fri is NWD)
    worker1 = Worker(
        name="worker1", site="group1", grade=1,
        nwds=[{"day": "Fri", "start_date": start_date, "end_date": start_date + datetime.timedelta(weeks=2)}],
    )
    # worker2 works normally
    worker2 = Worker(name="worker2", site="group1", grade=1)
    
    Rota.add_workers((worker1, worker2))
    
    # Fri/Sat/Sun shift, requiring 1 worker, forced as block unless NWD
    # We specify sub-blocks: [["Fri"], ["Sat", "Sun"]]
    Rota.add_shifts(
        SingleShift(
            sites=("group1",), name="weekend", length=12.5,
            days=["Fri", "Sat", "Sun"],
            workers_required=1,
            force_as_block_unless_nwd=[["Fri"], ["Sat", "Sun"]],
        ),
    )
    
    # We remove no valid shifts warning
    Rota.terminate_on_warning.remove("Worker/no valid shifts")
    
    Rota.build_and_solve()
    assert Rota.results.solver.status == "ok"
    
    # Check assignments:
    # Since worker1 has NWD on Friday:
    # - worker1 cannot work Friday.
    # - But worker1 can work Sat and Sun as a block!
    # So on Saturday and Sunday, worker1 should be assigned.
    # On Friday, worker2 should be assigned.
    for week in Rota.weeks:
        w1_fri = Rota.model.works[worker1.id, week, "Fri", "weekend"].value
        w1_sat = Rota.model.works[worker1.id, week, "Sat", "weekend"].value
        w1_sun = Rota.model.works[worker1.id, week, "Sun", "weekend"].value
        
        w2_fri = Rota.model.works[worker2.id, week, "Fri", "weekend"].value
        w2_sat = Rota.model.works[worker2.id, week, "Sat", "weekend"].value
        w2_sun = Rota.model.works[worker2.id, week, "Sun", "weekend"].value
        
        assert w1_fri == 0
        
        # Either worker1 works Sat/Sun and worker2 works Fri (split block)
        # OR worker2 works Fri/Sat/Sun (full block) and worker1 works nothing
        is_split = (w1_sat == 1 and w1_sun == 1 and w2_fri == 1 and w2_sat == 0 and w2_sun == 0)
        is_full_w2 = (w1_sat == 0 and w1_sun == 0 and w2_fri == 1 and w2_sat == 1 and w2_sun == 1)
        assert is_split or is_full_w2