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

def generate_basic_rota(weeks_to_rota=10):
    start_date = datetime.date(2022, 3, 7)
    Rota = RotaBuilder(
        start_date,
        weeks_to_rota=weeks_to_rota,
    )
    Rota.add_workers(
        [
            Worker(name="worker1", site="group1", grade=1),
            Worker(name="worker2", site="group1", grade=1),
        ]
    )
    return Rota

def date_generator(from_date, days):
    n = 0
    while True:
        yield from_date
        n = n + 1
        if n >= days:
            break
        from_date = from_date + datetime.timedelta(days=1)

def test_nights():
    Rota = generate_basic_rota()
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="n",
            length=12.5,
            days=days[5:],
            force_as_block=True,
            constraints=[NightConstraint(), PreShiftConstraint(days=2)],
            workers_required=2,
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="b",
            length=12.5,
            days=days[:3],
            force_as_block=False,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("night", folder="tests")
    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"

def test_nights_pre3():
    Rota = generate_basic_rota()
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="n",
            length=12.5,
            days=days[5:],
            force_as_block=True,
            constraints=[NightConstraint(), PreShiftConstraint(days=3)],
            workers_required=2,
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="b",
            length=12.5,
            days=days[:2],
            force_as_block=False,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("night_pre3", folder="tests")
    assert Rota.results.solver.status == "ok"

def test_nights_pre3_2():
    Rota = generate_basic_rota()
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="n",
            length=12.5,
            days=days[5:],
            force_as_block=True,
            constraints=[NightConstraint(), PreShiftConstraint(days=3)],
            workers_required=1,
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="b",
            length=12.5,
            days=days[:3],
            force_as_block=False,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("night_pre3", folder="tests")
    assert Rota.results.solver.status == "ok"

def test_nights_fail():
    Rota = generate_basic_rota()
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="n",
            length=12.5,
            days=days[5:],
            force_as_block=True,
            constraints=[NightConstraint(), PreShiftConstraint(days=8), PostShiftConstraint(days=2)],
            workers_required=2,
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="b",
            length=12.5,
            days=days[:3],
            force_as_block=False,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("night fail", folder="tests")
    assert Rota.results.solver.status in ("warning", "error")

def test_nights2():
    Rota = generate_basic_rota()
    Rota.constraint_options["max_night_frequency"] = 0
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="n",
            length=12.5,
            days=days[5:],
            force_as_block=True,
            constraints=[NightConstraint()],
            workers_required=2,
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="m",
            length=12.5,
            days=days[:5],
            force_as_block=True,
            workers_required=2,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("night", folder="tests")
    assert Rota.results.solver.status == "ok"

def test_nights_pre_wrap_around():
    Rota = generate_basic_rota(weeks_to_rota=8)
    Rota.constraint_options["max_night_frequency"] = 0
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="n",
            length=12.5,
            days=days[:2],
            force_as_block=True,
            constraints=[
                NightConstraint(),
                PreShiftConstraint(days=5),
                PostShiftConstraint(days=5),
            ],
            workers_required=3,
        ),
    )
    Rota.add_workers(
        [
            Worker(
                name="worker3",
                site="group1",
                grade=1,
                not_available_to_work=[
                    {
                        "date": Rota.rota_end_date - datetime.timedelta(days=1),
                        "reason": "no wanna",
                    },
                ],
            ),
        ]
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("night pre wrap around", folder="tests")
    assert Rota.results.solver.status == "ok"

def test_nights_max_frequency():
    Rota = generate_basic_rota()
    Rota.constraint_options["max_night_frequency"] = 1
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="n",
            length=12.5,
            days=days[5:],
            force_as_block=True,
            constraints=[NightConstraint()],
            workers_required=2,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("night max frequency", folder="tests")
    assert Rota.results.solver.status == "ok"

def test_nights_max_frequency_fail():
    Rota = generate_basic_rota()
    Rota.constraint_options["max_night_frequency"] = 2
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="n",
            length=12.5,
            days=days[5:],
            force_as_block=True,
            constraints=[NightConstraint()],
            workers_required=2,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    assert Rota.results.solver.status in ("warning", "error")

def test_nights_max_frequency3():
    Rota = generate_basic_rota()
    Rota.constraint_options["max_night_frequency"] = 2
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="n",
            length=12.5,
            days=days[5:],
            force_as_block=True,
            constraints=[NightConstraint()],
            workers_required=1,
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="m",
            length=12.5,
            days=days[:5],
            force_as_block=True,
            constraints=[NightConstraint()],
            workers_required=1,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("night max frequency 3", folder="tests")
    assert Rota.results.solver.status in ("warning", "error")

def test_nights_max_frequency4():
    Rota = generate_basic_rota()
    Rota.constraint_options["max_night_frequency"] = 2
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="n",
            length=12.5,
            days=days[5:],
            force_as_block=True,
            constraints=[NightConstraint()],
            workers_required=1,
        ),
        SingleShift(
            sites=("group1", "group2"),
            name="m",
            length=12.5,
            days=days[:5],
            force_as_block=True,
            constraints=[NightConstraint()],
            workers_required=1,
        ),
    )
    Rota.add_workers(
        [
            Worker(name="worker3", site="group1", grade=1),
            Worker(name="worker4", site="group1", grade=1),
        ]
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("night max frequency 4", folder="tests")
    assert Rota.results.solver.status == "ok"

def test_nights_max_frequency_exclusions():
    Rota = generate_basic_rota()
    Rota.constraint_options["max_night_frequency"] = 2
    Rota.constraint_options["max_night_frequency_week_exclusions"] = list(range(10))
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="n",
            length=12.5,
            days=days[5:],
            force_as_block=True,
            constraints=[NightConstraint()],
            workers_required=2,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("night max frequency exclusions", folder="tests")
    assert Rota.results.solver.status == "ok"

def test_nights_max_frequency_exclusions2():
    Rota = generate_basic_rota(weeks_to_rota=8)
    Rota.constraint_options["max_night_frequency"] = 2
    Rota.constraint_options["max_night_frequency_week_exclusions"] = list(range(1,4))
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="n",
            length=12.5,
            days=days[5:],
            force_as_block=True,
            constraints=[NightConstraint()],
            workers_required=2,
        ),
    )
    Rota.add_workers(
        [
            Worker(
                name="worker3",
                site="group1",
                grade=1,
                start_date=Rota.start_date + datetime.timedelta(weeks=4)
            ),
            Worker(
                name="worker4",
                site="group1",
                grade=1,
                start_date=Rota.start_date + datetime.timedelta(weeks=4)
            ),
        ]
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("night max frequency exclusions 2", folder="tests")
    assert Rota.results.solver.status == "ok"

def test_nights_max_frequency_exclusions3():
    Rota = generate_basic_rota(weeks_to_rota=9)
    Rota.constraint_options["max_night_frequency"] = 2
    Rota.constraint_options["max_night_frequency_week_exclusions"] = list(range(1,5))
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="n",
            length=12.5,
            days=days[5:],
            force_as_block=True,
            balance_offset=10,
            constraints=[NightConstraint()],
            workers_required=2,
        ),
    )
    Rota.add_workers(
        [
            Worker(
                name="worker3",
                site="group1",
                grade=1,
                start_date=Rota.start_date + datetime.timedelta(weeks=4)
            ),
            Worker(
                name="worker4",
                site="group1",
                grade=1,
                start_date=Rota.start_date + datetime.timedelta(weeks=4)
            ),
        ]
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("night max frequency exclusions 3", folder="tests")
    assert Rota.results.solver.status == "ok"