import datetime
import pytest
from pydantic import ValidationError
from rota_generator.shifts import (
    InvalidShift, MaxShiftsPerWeekBlockConstraint, MaxShiftsPerWeekConstraint,
    NoWorkers, PostShiftConstraint, PreShiftConstraint, RotaBuilder, SingleShift,
    WarningTermination, days, WorkerRequirement
)
from rota_generator.workers import (
    MaxDaysPerWeekBlockConstraint,
    MaxUniqueShiftsPerWeekBlockConstraint,
    NotAvailableToWork,
    WorkRequests,
    Worker,
    generate_not_available_to_works,
)

def generate_basic_rota(weeks_to_rota=10, workers=2, start_date=datetime.date(2022, 3, 7)):
    Rota = RotaBuilder(
        start_date,
        weeks_to_rota=weeks_to_rota,
    )
    Rota.add_workers(
        [Worker(name=f"worker{i}", site="group1", grade=1) for i in range(1, workers+1)]
    )
    return Rota

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

def test_duplicate_shifts():
    Rota = generate_basic_rota()
    Rota.add_worker(Worker(name="worker3", site="group1", grade=1))
    Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days))
    Rota.build_and_solve(options={"ratio": 0.000})
    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"
    Rota.add_shifts(SingleShift(sites=("group1",), name="a", length=12.5, days=days))
    with pytest.raises(InvalidShift):
        Rota.build_and_solve(options={"ratio": 0.000})

def test_max_shifts():
    Rota = generate_basic_rota()
    Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
                                constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)]))
    Rota.build_and_solve(options={"ratio": 0.000})
    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"

def test_max_shifts_fail():
    Rota = generate_basic_rota()
    Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
                                constraints=[MaxShiftsPerWeekConstraint(max_shifts=3)]))
    Rota.build_and_solve(options={"ratio": 0.000})
    assert Rota.results.solver.termination_condition == "infeasible"
    Rota.add_worker(Worker(name="worker3", site="group1", grade=1))
    Rota.build_and_solve(options={"ratio": 0.000})
    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"

def test_max_shifts_per_week_block():
    Rota = generate_basic_rota()
    Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
                                constraints=[MaxShiftsPerWeekBlockConstraint(week_block=2)]))
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("max_shifts_per_week_block")
    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"

def test_max_shifts_per_week_block_fail():
    Rota = generate_basic_rota()
    Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
                                constraints=[MaxShiftsPerWeekBlockConstraint(week_block=3)]))
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("max_shifts_per_week_block")
    assert Rota.results.solver.status in ("warning", "error")

def test_max_shifts_per_week_block_shift_number():
    Rota = generate_basic_rota()
    Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
                                constraints=[MaxShiftsPerWeekBlockConstraint(week_block=2, max_shifts=7)]))
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("max_shifts_per_week_block")
    assert Rota.results.solver.status == "ok"

def test_max_shifts_per_week_block_shift_number_fail():
    Rota = generate_basic_rota()
    Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
                                constraints=[MaxShiftsPerWeekBlockConstraint(week_block=2, max_shifts=6)]))
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("max_shifts_per_week_block")
    assert Rota.results.solver.status in ("warning", "error")


def test_worker_max_days_per_week_block_fail():
    Rota = RotaBuilder(datetime.date(2022, 3, 7), weeks_to_rota=2)
    Rota.add_worker(
        Worker(
            name="worker1",
            site="group1",
            grade=1,
            max_days_per_week_block=[
                MaxDaysPerWeekBlockConstraint(max_days=1, week_block=1)
            ],
        )
    )
    Rota.add_shifts(
        SingleShift(
            sites=("group1",),
            name="a",
            length=12.5,
            days=days[:3],
            workers_required=1,
        )
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    assert Rota.results.solver.termination_condition == "infeasible"


def test_worker_max_unique_shifts_per_week_block_fail():
    Rota = RotaBuilder(datetime.date(2022, 3, 7), weeks_to_rota=2)
    Rota.add_worker(
        Worker(
            name="worker1",
            site="group1",
            grade=1,
            max_unique_shifts_per_week_block=[
                MaxUniqueShiftsPerWeekBlockConstraint(max_unique_shifts=1, week_block=1)
            ],
        )
    )
    Rota.add_shifts(
        SingleShift(
            sites=("group1",),
            name="a",
            length=12.5,
            days=["Mon"],
            workers_required=1,
        ),
        SingleShift(
            sites=("group1",),
            name="b",
            length=12.5,
            days=["Tue"],
            workers_required=1,
        ),
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    assert Rota.results.solver.termination_condition == "infeasible"


def test_worker_block_constraint_validation():
    with pytest.raises(ValidationError):
        MaxDaysPerWeekBlockConstraint(max_days=0, week_block=1)
    with pytest.raises(ValidationError):
        MaxUniqueShiftsPerWeekBlockConstraint(max_unique_shifts=0, week_block=1)


def test_worker_max_days_per_week_block_pass():
    """Positive test: constraint allowing reasonable days."""
    Rota = generate_basic_rota()
    # Add constraint allowing up to 4 days per week
    Rota.workers[0].max_days_per_week_block = [
        MaxDaysPerWeekBlockConstraint(max_days=4, week_block=1)
    ]
    # Add shifts to the rota
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="a",
            length=12.5,
            days=days,
        )
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"


def test_worker_max_unique_shifts_per_week_block_pass():
    """Positive test: allowing multiple shifts within week block."""
    Rota = generate_basic_rota()
    # Add constraint allowing up to 2 unique shifts per week
    Rota.workers[0].max_unique_shifts_per_week_block = [
        MaxUniqueShiftsPerWeekBlockConstraint(max_unique_shifts=2, week_block=1)
    ]
    # Add shifts to the rota
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="a",
            length=12.5,
            days=days,
        )
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"


def test_worker_max_days_per_week_block_with_ignore_dates():
    """Test max_days_per_week_block with ignored dates works without error."""
    Rota = generate_basic_rota()
    rota_start = Rota.start_date
    ignored_date = rota_start  # Ignore first Monday
    # Add constraint with ignore_dates
    Rota.workers[0].max_days_per_week_block = [
        MaxDaysPerWeekBlockConstraint(
            max_days=3,
            week_block=1,
            ignore_dates=[ignored_date],
        )
    ]
    # Add shifts to the rota
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="a",
            length=12.5,
            days=days,
        )
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    assert Rota.results.solver.status == "ok"


def test_worker_max_unique_shifts_per_week_block_with_ignore_dates():
    """Test max_unique_shifts_per_week_block with ignored dates works without error."""
    Rota = generate_basic_rota()
    rota_start = Rota.start_date
    ignored_date = rota_start + datetime.timedelta(days=1)  # Ignore Tuesday
    # Add constraint with ignore_dates
    Rota.workers[0].max_unique_shifts_per_week_block = [
        MaxUniqueShiftsPerWeekBlockConstraint(
            max_unique_shifts=1,
            week_block=1,
            ignore_dates=[ignored_date],
        )
    ]
    # Add shifts to the rota
    Rota.add_shifts(
        SingleShift(
            sites=("group1", "group2"),
            name="a",
            length=12.5,
            days=days,
        )
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    # This should still be feasible with only one shift type in main rota
    assert Rota.results.solver.status == "ok"

def test_pre_shift_constraint():
    Rota = generate_basic_rota()
    for i, d in enumerate(days[:6]):
        Rota.add_shifts(SingleShift(sites=("group1", "group2"), name=chr(97+i), length=12.5, days=d,
                                    constraints=[PreShiftConstraint(days=1)]))
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("pre")
    assert Rota.results.solver.status == "ok"
    for worker in Rota.get_workers():
        shift_string = Rota.get_worker_shift_list_string(worker)
        for s in ["ab", "bc", "cd", "de", "ef", "fg"]:
            assert s not in shift_string
    Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="g", length=12.5, days=days[6],
                                constraints=[PreShiftConstraint(days=2)]))
    Rota.build_and_solve(options={"ratio": 0.000})
    assert Rota.results.solver.status in ("warning", "error")

def test_pre_shift_constraint2():
    Rota = generate_basic_rota()
    Rota.add_shifts(
        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[0],
                    constraints=[PreShiftConstraint(days=1)], workers_required=2)
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("pre")
    assert Rota.results.solver.status == "ok"

def test_pre_shift_constraint3():
    Rota = generate_basic_rota()
    Rota.add_shifts(
        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[0],
                    constraints=[PreShiftConstraint(days=3)], workers_required=2),
        SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[4],
                    constraints=[PreShiftConstraint(days=3)], workers_required=2),
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("pre")
    assert Rota.results.solver.status in ("warning", "error")

def test_post_shift_constraint():
    Rota = generate_basic_rota()
    Rota.add_shifts(
        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[0],
                    constraints=[PostShiftConstraint(days=3)], workers_required=2),
        SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[4],
                    constraints=[PostShiftConstraint(days=2)], workers_required=2),
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("post")
    assert Rota.results.solver.status == "ok"

def test_post_shift_constraint2():
    Rota = generate_basic_rota()
    Rota.add_shifts(
        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[0],
                    constraints=[PostShiftConstraint(days=3)], workers_required=2),
        SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[4],
                    constraints=[PostShiftConstraint(days=3)], workers_required=2),
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("post")
    assert Rota.results.solver.status in ("warning", "error")

def test_shift_constraint_week_limit():
    Rota = generate_basic_rota()
    Rota.add_shifts(
        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[0],
                    constraints=[PostShiftConstraint(days=1, weeks=[1,2,3])], workers_required=2),
        SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[1], workers_required=2),
    )
    Rota.build_and_solve(options={"ratio": 0.001})
    Rota.export_rota_to_html("post")
    assert Rota.results.solver.status in ("warning", "error")
    Rota.add_workers([
        Worker(name="worker3", end_date=Rota.start_date + datetime.timedelta(weeks=3), site="group1", grade=1),
        Worker(name="worker4", end_date=Rota.start_date + datetime.timedelta(weeks=3), site="group1", grade=1)
    ])
    Rota.build_and_solve(options={"ratio": 0.001})
    assert Rota.results.solver.status == "ok"

def test_shift_constraint_week_limit2():
    Rota = generate_basic_rota()
    Rota.add_shifts(
        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[0],
                    constraints=[PostShiftConstraint(days=1, start_date=Rota.start_date + datetime.timedelta(weeks=2), end_date=Rota.start_date + datetime.timedelta(weeks=4))], workers_required=2),
        SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[1], workers_required=2),
    )
    Rota.build_and_solve(options={"ratio": 0.001})
    assert Rota.results.solver.status in ("warning", "error")
    Rota.add_workers([
        Worker(name="worker3", start_date=Rota.start_date + datetime.timedelta(weeks=2), end_date=Rota.start_date + datetime.timedelta(weeks=4), site="group1", grade=1),
        Worker(name="worker4", start_date=Rota.start_date + datetime.timedelta(weeks=2), end_date=Rota.start_date + datetime.timedelta(weeks=4), site="group1", grade=1)
    ])
    Rota.build_and_solve(options={"ratio": 0.001})
    Rota.export_rota_to_html("post")
    assert Rota.results.solver.status == "ok"

def test_shift_start_date():
    Rota = generate_basic_rota(workers=2)
    Rota.add_shifts(
        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)],
                    start_date=(Rota.start_date + datetime.timedelta(days=7))),
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"
    assert Rota.get_workers_total_shifts()["worker1"] in (31, 32)
    assert Rota.get_workers_total_shifts()["worker2"] in (31, 32)

def test_shift_start_date_end_date():
    Rota = generate_basic_rota(workers=2)
    Rota.add_shifts(
        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)],
                    start_date=(Rota.start_date + datetime.timedelta(days=7)),
                    end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1)))),
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"
    assert Rota.get_workers_total_shifts()["worker1"] in (17, 18)
    assert Rota.get_workers_total_shifts()["worker2"] in (17, 18)
    for worker in Rota.get_workers():
        assert Rota.get_worker_shift_list_string(worker).startswith("-"*7)
        assert Rota.get_worker_shift_list_string(worker).endswith("-"*28)

def test_shift_invalid_start_date():
    Rota = generate_basic_rota(workers=2)
    Rota.add_shifts(
        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)],
                    start_date=(Rota.start_date - datetime.timedelta(days=6))),
    )
    with pytest.raises(WarningTermination):
        Rota.build_and_solve(options={"ratio": 0.000})

def test_shift_invalid_start_date2():
    Rota = generate_basic_rota(workers=2)
    Rota.add_shifts(
        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)],
                    start_date=(Rota.rota_end_date + datetime.timedelta(days=1))),
    )
    with pytest.raises(WarningTermination):
        Rota.build_and_solve(options={"ratio": 0.000})

def test_shift_invalid_end_date():
    Rota = generate_basic_rota(workers=2)
    Rota.add_shifts(
        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)],
                    end_date=(Rota.rota_end_date + datetime.timedelta(days=1))),
    )
    with pytest.raises(WarningTermination):
        Rota.build_and_solve(options={"ratio": 0.000})

def test_shift_invalid_end_date2():
    Rota = generate_basic_rota(workers=2)
    Rota.add_shifts(
        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)],
                    end_date=(Rota.start_date - datetime.timedelta(days=1))),
    )
    with pytest.raises(WarningTermination):
        Rota.build_and_solve(options={"ratio": 0.000})

def test_shift_start_date_end_date_block():
    Rota = generate_basic_rota(workers=2)
    Rota.add_shifts(
        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
                    assign_as_block=True,
                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)],
                    start_date=(Rota.start_date + datetime.timedelta(days=7)),
                    end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1)))),
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"
    assert Rota.get_workers_total_shifts()["worker1"] in (17, 18)
    assert Rota.get_workers_total_shifts()["worker2"] in (17, 18)
    for worker in Rota.get_workers():
        assert Rota.get_worker_shift_list_string(worker).startswith("-"*7)
        assert Rota.get_worker_shift_list_string(worker).endswith("-"*28)

def test_shift_start_date_end_date_block_more_workers():
    Rota = generate_basic_rota(workers=10)
    Rota.add_shifts(
        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[:4],
                    assign_as_block=True, force_as_block=True, workers_required=4,
                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)]),
        SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[4:],
                    assign_as_block=True, force_as_block=True, workers_required=4,
                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4),
                                 PreShiftConstraint(days=3), PostShiftConstraint(days=3)],
                    start_date=(Rota.start_date + datetime.timedelta(days=7)),
                    end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1)))),
    )
    Rota.build_and_solve(options={"ratio": 0.000})
    Rota.export_rota_to_html("test_shift_start_date")
    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"
    total_shifts = Rota.get_workers_total_shifts()
    for worker in Rota.get_workers():
        assert total_shifts[worker.name] == 22
        shift_string = Rota.get_worker_shift_list_string(worker)
        assert "ab" not in shift_string
        assert "ba" not in shift_string

def test_shift_start_date_end_date_block_with_requests():
    Rota = generate_basic_rota(workers=7)
    request_date = Rota.start_date + datetime.timedelta(days=(7*6)+4)
    request = [WorkRequests(date=request_date, shift="c")]
    Rota.add_workers([
        Worker(name=f"worker-wr-{i}", site="group2", grade=1, work_requests=request) for i in range(1, 4)
    ])
    Rota.add_shifts(
        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[:4],
                    assign_as_block=True, force_as_block=True, workers_required=4,
                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)]),
        SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[4:],
                    assign_as_block=True, force_as_block=True, workers_required=4,
                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4),
                                 PreShiftConstraint(days=3), PostShiftConstraint(days=3)],
                    start_date=(Rota.start_date + datetime.timedelta(days=7)),
                    end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1)))),
        SingleShift(sites=("group1", "group2"), name="c", length=12.5, days=days[4:],
                    assign_as_block=True, force_as_block=True, workers_required=4,
                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4),
                                 PreShiftConstraint(days=3), PostShiftConstraint(days=3)],
                    start_date=(Rota.start_date + datetime.timedelta(days=(7*6)))),
    )
    Rota.build_and_solve(options={"ratio": 0.1})
    Rota.export_rota_to_html("test_shift_start_date")
    assert Rota.results.solver.status == "ok"
    assert Rota.results.solver.termination_condition == "optimal"
    total_shifts = Rota.get_workers_total_shifts()
    for worker in Rota.get_workers():
        assert total_shifts[worker.name] in range(24, 30)
        shift_string = Rota.get_worker_shift_list_string(worker)
        assert "ab" not in shift_string
        assert "ba" not in shift_string
        assert shift_string.count("aaaa") == shift_string.count("a") / 4
        assert shift_string.count("bbb") == shift_string.count("b") / 3
        assert shift_string.count("ccc") == shift_string.count("c") / 3
    for worker in Rota.get_workers_by_group()["group2"]:
        assert Rota.get_worker_shifts_by_date(worker)[request_date] == "c"

def test_worker_requirement_start_and_end_date():
    WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 14))

def test_worker_requirement_end_date_before_start_date():
    with pytest.raises(ValueError):
        WorkerRequirement(start_date=datetime.date(2022, 3, 14), end_date=datetime.date(2022, 3, 7))

def test_worker_requirement():
    Rota = generate_basic_rota(workers=2)
    wr = [WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 21))]
    Rota.add_shifts(
        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[:4],
                    assign_as_block=True, force_as_block=True, workers_required=wr),
    )
    Rota.build_and_solve(options={"ratio": 0.1})
    Rota.export_rota_to_html("test_worker_requirement")
    assert Rota.results.solver.status == "ok"
    total_shifts = Rota.get_workers_total_shifts()
    for worker in Rota.get_workers():
        assert total_shifts[worker.name] == 4
        assert "aaaa" in Rota.get_worker_shift_list_string(worker)

def test_worker_requirement2():
    Rota = generate_basic_rota(workers=2)
    wr = [
        WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 21)),
        WorkerRequirement(start_date=datetime.date(2022, 3, 28))
    ]
    Rota.add_shifts(
        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[:4],
                    assign_as_block=True, workers_required=wr),
    )
    Rota.build_and_solve(options={"ratio": 0.1})
    Rota.export_rota_to_html("test_worker_requirement")
    assert Rota.results.solver.status == "ok"
    total_shifts = Rota.get_workers_total_shifts()
    for worker in Rota.get_workers():
        assert total_shifts[worker.name] == 18
        assert Rota.get_worker_shift_list_string(worker)[14:21] == "-------"

def test_worker_requirement_increase_workers():
    Rota = generate_basic_rota(workers=2)
    wr = [
        WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 14)),
        WorkerRequirement(start_date=datetime.date(2022, 3, 28), number=2)
    ]
    Rota.add_shifts(
        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[:4],
                    assign_as_block=True, workers_required=wr),
    )
    Rota.build_and_solve(options={"ratio": 0.1})
    Rota.export_rota_to_html("test_worker_requirement")
    assert Rota.results.solver.status == "ok"
    total_shifts = Rota.get_workers_total_shifts()
    for worker in Rota.get_workers():
        assert total_shifts[worker.name] == 30
        assert Rota.get_worker_shift_list_string(worker)[7:21] == "-" * 14
        assert Rota.get_worker_shift_list_string(worker).endswith("aaaa---"*7)

def test_worker_requirement_balancing():
    Rota = generate_basic_rota(workers=2)
    natws = generate_not_available_to_works(datetime.date(2022,4,18), days=21, reason="holiday")
    Rota.add_worker(Worker(name="worker3", site="group1", grade=1, not_available_to_work=natws))
    wr = [
        WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 14)),
        WorkerRequirement(start_date=datetime.date(2022, 3, 28), number=2)
    ]
    Rota.add_shifts(
        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[:4],
                    assign_as_block=True, workers_required=wr),
    )
    Rota.build_and_solve(options={"ratio": 0.0})
    Rota.export_rota_to_html("test_worker_requirement")
    assert Rota.results.solver.status == "ok"
    total_shifts = Rota.get_workers_total_shifts()
    for worker in Rota.get_workers():
        assert total_shifts[worker.name] == 20
        assert Rota.get_worker_shift_list_string(worker)[7:21] == "-" * 14
        if worker.name == "worker3":
            assert Rota.get_worker_shift_list_string(worker).endswith("aaaa---")
            assert Rota.get_worker_shift_list_string(worker).startswith("aaaa---")

def test_worker_requirement_multiple_shifts():
    Rota = generate_basic_rota(workers=4)
    wr = [
        WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 14)),
        WorkerRequirement(start_date=datetime.date(2022, 3, 28), number=2)
    ]
    Rota.add_shifts(
        SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[4:],
                    assign_as_block=True, workers_required=2),
        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[:4],
                    assign_as_block=True, workers_required=wr),
    )
    Rota.build_and_solve(options={"ratio": 0.1})
    Rota.export_rota_to_html("test_worker_requirement")
    assert Rota.results.solver.status == "ok"
    total_shifts = Rota.get_workers_total_shifts()
    for worker in Rota.get_workers():
        assert total_shifts[worker.name] == 30

def test_display_char_in_output(monkeypatch):
    Rota = generate_basic_rota(workers=1)
    shift = SingleShift(
        sites=("group1",),
        name="testshift",
        length=8,
        days=days[:1],
        display_char="Z"
    )
    Rota.add_shifts(shift)
    assert shift.display_char == "Z"
    Rota.build_and_solve(options={"ratio": 0.1})
    assert Rota.results.solver.status == "ok"
    for worker in Rota.get_workers():
        assert Rota.get_worker_shift_list_string(worker).startswith("Z------" * 5)