import pytest
from rota_generator.shifts import PostShiftConstraint, PreShiftConstraint, RotaBuilder, SingleShift, days
from rota_generator.workers import Worker
import datetime


def generate_basic_rota(
    weeks_to_rota=10, workers=0, start_date=datetime.date(2022, 3, 7)
):
    Rota = RotaBuilder(
        start_date,
        weeks_to_rota=weeks_to_rota,
    )

    # Add a few workers
    if workers:
        Rota.add_workers(
            [
                Worker(name=f"worker{i:02}", site="group1", grade=1)
                for i in range(1, workers + 1)
            ]
        )

    return Rota


@pytest.fixture
def demo_rota_clear():
    weeks_to_rota = 10
    start_date = datetime.date(2022, 3, 7)

    Rota = RotaBuilder(
        start_date,
        weeks_to_rota=weeks_to_rota,
        balance_offset_modifier=1,
    )
    Rota.constraint_options_model.max_shifts_per_month = 20

    workers = [
        Worker(name="worker1", site="group1", grade=1),
        Worker(name="worker2", site="group1", grade=1),
        Worker(name="worker3", site="group1", grade=1),
        Worker(name="worker4", site="group1", grade=1),
        Worker(name="worker5", site="group2", grade=1, fte=60),
        Worker(name="worker6", site="group2", grade=1, fte=40),
        Worker(name="worker7", site="group2", grade=1),
        Worker(name="worker8", site="group2", grade=1),
    ]
    Rota.add_workers(workers)
    return Rota


def test_pre(demo_rota_clear):
    Rota = demo_rota_clear
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="a",
            length=12.5,
            days=days[:5],
            workers_required=2,
            constraints=[PreShiftConstraint(days=1)],
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="b",
            length=12.5,
            days=days[:5],
            workers_required=2,
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="c",
            length=12.5,
            days=days[3],
            workers_required=1,
            constraints=[PreShiftConstraint(days=1), PostShiftConstraint(days=1)],
        ),
    )

    Rota.build_shifts()
    Rota.build_workers()
    Rota.build_model()

    solver_options = {"seconds": 1000, "threads": 10}
    Rota.solve_model(options=solver_options)
    Rota.export_rota_to_html("test4")

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

        # check for patterns that should not occur
        assert "ba" not in shifts_string
        assert "ca" not in shifts_string
        assert "ac" not in shifts_string
        assert "ca" not in shifts_string
        assert "bc" not in shifts_string
        assert "cb" not in shifts_string
        # Uncomment if you want to check for these as well:
        # assert "a-c" not in shifts_string
        # assert "c-a" not in shifts_string
        # assert "b-c" not in shifts_string
        # assert "c-b" not in shifts_string


def test_pre_post_clear(demo_rota_clear):
    Rota = generate_basic_rota(workers=4, weeks_to_rota=16)

    workers = Rota.get_workers()

    # for worker in workers:
    #    worker.add_force_assign_with("a", "b")

    Rota.add_shifts(
        SingleShift(
            sites=("group1",),
            name="a",
            length=8,
            days=days[:5],
            workers_required=2,
            assign_as_block=False,  # global setting is off
            constraints=[
                PreShiftConstraint(days=1),
                PostShiftConstraint(days=1),
            ],
        ),
        SingleShift(
            sites=("group1",),
            name="b",
            length=8,
            days=days[5:],
            workers_required=2,
            assign_as_block=False,  # global setting is off
            constraints=[
                PreShiftConstraint(days=1),
                PostShiftConstraint(days=1),
            ],
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.0})
    Rota.export_rota_to_html("test_pre_post_clear", 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])

        # check for patterns that should not occur
        assert "ab" not in shifts_string
        assert "ba" not in shifts_string
        # assert "a-b" not in shifts_string
        # assert "b-a" not in shifts_string


def test_pre_post_clear2(demo_rota_clear):
    Rota = generate_basic_rota(workers=4, weeks_to_rota=16)

    workers = Rota.get_workers()

    # for worker in workers:
    #    worker.add_force_assign_with("a", "b")

    Rota.add_shifts(
        SingleShift(
            sites=("group1",),
            name="a",
            length=8,
            days=days[:5],
            workers_required=2,
            assign_as_block=False,  # global setting is off
            constraints=[
                PreShiftConstraint(days=2),
                PostShiftConstraint(days=2),
            ],
        ),
        SingleShift(
            sites=("group1",),
            name="b",
            length=8,
            days=days[5:],
            workers_required=2,
            assign_as_block=False,  # global setting is off
            constraints=[
                PreShiftConstraint(days=2),
                PostShiftConstraint(days=2),
            ],
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.0})
    Rota.export_rota_to_html("test_pre_post_clear", 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])

        # check for patterns that should not occur
        assert "ab" not in shifts_string
        assert "ba" not in shifts_string
        assert "a-b" not in shifts_string
        assert "b-a" not in shifts_string


def test_pre_post_clear_multi_group(demo_rota_clear):
    Rota = generate_basic_rota(workers=2, weeks_to_rota=16)

    workers = Rota.get_workers()

    Rota.add_workers(
        (
            Worker(name="worker3", site="group2", grade=1),
            Worker(name="worker4", site="group2", grade=1),
        )
    )

    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="a",
            length=8,
            days=days[:5],
            workers_required=2,
            assign_as_block=False,  # global setting is off
            constraints=[
                PreShiftConstraint(days=2),
                PostShiftConstraint(days=2),
            ],
        ),
        SingleShift(
            sites=("group1",),
            name="b",
            length=8,
            days=days[5:],
            workers_required=1,
            assign_as_block=False,  # global setting is off
            constraints=[
                PreShiftConstraint(days=2),
                PostShiftConstraint(days=2),
            ],
        ),
        SingleShift(
            sites=("group2",),
            name="c",
            length=8,
            days=days[5:],
            workers_required=1,
            assign_as_block=False,  # global setting is off
            constraints=[
                PreShiftConstraint(days=2),
                PostShiftConstraint(days=2),
            ],
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.0})
    Rota.export_rota_to_html("test_pre_post_clear", 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])

        # check for patterns that should not occur
        assert "ab" not in shifts_string
        assert "ba" not in shifts_string
        assert "a-b" not in shifts_string
        assert "b-a" not in shifts_string

        assert "ac" not in shifts_string
        assert "ca" not in shifts_string
        assert "a-c" not in shifts_string
        assert "c-a" not in shifts_string

def test_pre_post_clear_force_assign(demo_rota_clear):
    Rota = generate_basic_rota(workers=6, weeks_to_rota=16)

    workers = Rota.get_workers()

    for worker in workers:
        worker.add_force_assign_with("a", "o")

    Rota.add_workers(
        (
            Worker(name="workera", site="group2", grade=1),
            Worker(name="workerb", site="group2", grade=1),
        )
    )

    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="a",
            length=8,
            days=days[:5],
            workers_required=1,
            assign_as_block=False,  # global setting is off
            constraints=[
                PreShiftConstraint(days=1),
                PostShiftConstraint(days=1),
            ],
        ),
        SingleShift(
            sites=("group1",),
            name="o",
            length=8,
            days=days[:7],
            workers_required=1,
            assign_as_block=False,  # global setting is off
        ),
        SingleShift(
            sites=("group1",),
            name="b",
            length=8,
            days=days[5:],
            workers_required=1,
            assign_as_block=False,  # global setting is off
            #constraint=[
            #    {
            #        "name": "pre",
            #        "options": 2,
            #    },
            #    {"name": "post", "options": 2},
            #],
        ),
        SingleShift(
            sites=("group2",),
            name="c",
            length=8,
            days=days[5:],
            workers_required=1,
            assign_as_block=False,  # global setting is off
            #constraint=[
            #    {
            #        "name": "pre",
            #        "options": 1,
            #    },
            #    {"name": "post", "options": 1},
            #],
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.01})
    Rota.export_rota_to_html("test_pre_post_clear", folder="tests")

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

def test_pre_post_clear_force_assign2(demo_rota_clear):
    Rota = generate_basic_rota(workers=6, weeks_to_rota=16)

    workers = Rota.get_workers()

    for worker in workers:
        worker.add_force_assign_with("a", "o")
        worker.add_force_assign_with("b", "o")

    Rota.add_workers(
        (
            Worker(name="workera", site="group2", grade=1),
            Worker(name="workerb", site="group2", grade=1),
        )
    )

    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="a",
            length=8,
            days=days[:5],
            workers_required=1,
            assign_as_block=False,  # global setting is off
            constraints=[
            ],
        ),
        SingleShift(
            sites=("group1",),
            name="o",
            length=8,
            days=days[:7],
            workers_required=1,
            assign_as_block=False,  # global setting is off
        ),
        SingleShift(
            sites=("group1",),
            name="b",
            length=8,
            days=days[5:],
            workers_required=1,
            assign_as_block=False,  # global setting is off
            constraints=[
            ],
        ),
        SingleShift(
            sites=("group2",),
            name="c",
            length=8,
            days=days[5:],
            workers_required=1,
            assign_as_block=False,  # global setting is off
            constraints=[
                PreShiftConstraint(days=1),
                PostShiftConstraint(days=1),
            ],
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.01})
    Rota.export_rota_to_html("test_pre_post_clear", 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])

        # check for patterns that should not occur
        assert "ac" not in shifts_string
        assert "ca" not in shifts_string