"""Pins every figure the commentary quotes.

    python3 test_mlfq.py
"""
from mlfq import Job, MLFQ, DEMOTE_BURST, DEMOTE_ALLOT

Q = (10, 20, 40)
TESTS = []


def test(fn):
    TESTS.append(fn)
    return fn


def hogs(policy, quanta=Q, boost=0, yield_at=9, n=2, horizon=1000,
         switch=0, preempt=True, block=0, allots=None):
    jobs = [Job(0, 0, 10 ** 9, yield_after=yield_at, block=block)]
    jobs += [Job(i, 0, 10 ** 9) for i in range(1, n + 1)]
    sim = MLFQ(jobs, quanta=quanta, allots=allots, policy=policy, boost=boost,
               switch=switch, preempt=preempt)
    sim.run(horizon)
    return jobs, sim


def mixed(policy, boost, horizon=2000):
    js = [Job(0, 0, 10 ** 9, yield_after=9, block=0)]
    js += [Job(1, 0, 10 ** 9), Job(2, 0, 10 ** 9)]
    js += [Job(3, 0, 10 ** 9, yield_after=2, block=30),
           Job(4, 0, 10 ** 9, yield_after=2, block=30)]
    MLFQ(js, quanta=Q, policy=policy, boost=boost).run(horizon)
    lat = js[3].lat + js[4].lat
    return (js[0].cpu, js[1].cpu + js[2].cpu,
            js[3].cpu + js[4].cpu, round(sum(lat) / len(lat), 1))


# ---------------------------------------------------------------- panel A
@test
def t_headline_98x():
    j, _ = hogs(DEMOTE_BURST)
    assert [x.cpu for x in j] == [980, 10, 10], [x.cpu for x in j]
    assert j[0].cpu // j[1].cpu == 98


@test
def t_headline_levels():
    j, _ = hogs(DEMOTE_BURST)
    assert [x.lvl for x in j] == [0, 1, 1]


@test
def t_headline_trace_first_40():
    _, sim = hogs(DEMOTE_BURST)
    row = "".join(str(t[1]) for t in sim.trace[:40])
    assert row == "0000000001111111111222222222200000000000", row


@test
def t_no_ticks_lost():
    j, sim = hogs(DEMOTE_BURST)
    assert sum(x.cpu for x in j) == 1000
    assert sim.idle == 0 and sim.overhead == 0


@test
def t_closed_form_h_minus_nq():
    for n in (1, 2, 3, 5, 9):
        for q0 in (10, 50):
            j, _ = hogs(DEMOTE_BURST, quanta=(q0, 2 * q0, 4 * q0),
                        yield_at=q0 - 1, n=n)
            assert j[0].cpu == 1000 - n * q0, (n, q0, j[0].cpu)


# ---------------------------------------------------------------- panel B
def dispatch_input(block, horizon=2000):
    js = [Job(0, 0, 10 ** 9, yield_after=1, block=block)]
    js += [Job(i, 0, 10 ** 9) for i in (1, 2)]
    MLFQ(js, quanta=Q, policy=DEMOTE_BURST).run(horizon)
    return [(b, s) for (b, s, _) in js[0].disp], js[0]


@test
def t_rule_input_is_identical():
    g, gj = dispatch_input(0)
    i, ij = dispatch_input(20)
    n = min(len(g), len(i))
    assert g[:n] == i[:n]
    assert n == 96, n


@test
def t_rule_input_is_all_ones():
    g, _ = dispatch_input(0)
    assert set(g) == {(1, False)}, set(g)


@test
def t_identical_input_different_outcome():
    _, gj = dispatch_input(0)
    _, ij = dispatch_input(20)
    assert (gj.cpu, ij.cpu) == (1980, 96)
    assert round(gj.cpu / ij.cpu, 1) == 20.6


@test
def t_interactive_latency_is_zero_under_r4_old():
    _, ij = dispatch_input(20)
    assert max(ij.lat) == 0


# ---------------------------------------------------------------- panel C
@test
def t_mixed_matrix():
    assert mixed(DEMOTE_BURST, 0) == (1780, 20, 200, 8.4)
    assert mixed(DEMOTE_BURST, 100) == (1440, 400, 160, 18.1)
    assert mixed(DEMOTE_ALLOT, 0) == (216, 1660, 124, 32.2)
    assert mixed(DEMOTE_ALLOT, 100) == (600, 1200, 200, 8.1)


@test
def t_only_the_last_row_is_fair():
    """600 for the gamer is exactly what each CPU hog gets."""
    js = [Job(0, 0, 10 ** 9, yield_after=9, block=0)]
    js += [Job(1, 0, 10 ** 9), Job(2, 0, 10 ** 9)]
    js += [Job(3, 0, 10 ** 9, yield_after=2, block=30),
           Job(4, 0, 10 ** 9, yield_after=2, block=30)]
    MLFQ(js, quanta=Q, policy=DEMOTE_ALLOT, boost=100).run(2000)
    assert js[0].cpu == js[1].cpu == js[2].cpu == 600


@test
def t_interactive_unloaded_ceiling():
    js = [Job(3, 0, 10 ** 9, yield_after=2, block=30),
          Job(4, 0, 10 ** 9, yield_after=2, block=30)]
    MLFQ(js, quanta=Q, policy=DEMOTE_BURST).run(2000)
    assert [x.cpu for x in js] == [126, 126]


@test
def t_boost_alone_makes_latency_worse():
    assert mixed(DEMOTE_BURST, 0)[3] < mixed(DEMOTE_BURST, 100)[3]


# ---------------------------------------------------------------- panel D
@test
def t_yield_point_is_inert_when_free():
    assert [hogs(DEMOTE_BURST, yield_at=y)[0][0].cpu
            for y in range(1, 10)] == [980] * 9


@test
def t_yield_point_matters_once_switching_costs():
    row = [hogs(DEMOTE_BURST, yield_at=y, switch=1)[0][0].cpu
           for y in (1, 3, 5, 7, 9)]
    assert row == [489, 733, 815, 855, 880], row
    assert row == sorted(row)


@test
def t_quantum_minus_one_is_optimal_under_switch_cost():
    for sw in (1, 2, 4):
        row = [hogs(DEMOTE_BURST, yield_at=y, switch=sw)[0][0].cpu
               for y in range(1, 10)]
        assert row.index(max(row)) == 8, (sw, row)


# ---------------------------------------------------------------- panel E
@test
def t_one_level_kills_it():
    j, _ = hogs(DEMOTE_BURST, quanta=(10,))
    assert [x.cpu for x in j] == [315, 345, 340]
    assert j[0].cpu < min(x.cpu for x in j[1:])


@test
def t_two_levels_is_enough():
    for nl in (2, 3, 4, 5):
        q = tuple(10 * 2 ** i for i in range(nl))
        j, _ = hogs(DEMOTE_BURST, quanta=q)
        assert [x.cpu for x in j] == [980, 10, 10], nl


@test
def t_boost_below_39_makes_the_rules_agree():
    for b in range(1, 39):
        o = [x.cpu for x in hogs(DEMOTE_BURST, boost=b)[0]]
        w = [x.cpu for x in hogs(DEMOTE_ALLOT, boost=b)[0]]
        assert o == w, (b, o, w)
    o = [x.cpu for x in hogs(DEMOTE_BURST, boost=39)[0]]
    w = [x.cpu for x in hogs(DEMOTE_ALLOT, boost=39)[0]]
    assert o != w and o[0] == 484 and w[0] == 459, (o, w)


@test
def t_boost_alone_never_fixes_the_exploit():
    for b in (100, 200, 500):
        j, _ = hogs(DEMOTE_BURST, boost=b)
        assert j[0].cpu > 2 * (1000 / 3), (b, j[0].cpu)


# --------------------------------------------------------- counterfactuals
@test
def t_load_bearing_line():
    assert [x.cpu for x in hogs(DEMOTE_BURST)[0]] == [980, 10, 10]
    assert [x.cpu for x in hogs(DEMOTE_ALLOT)[0]] == [135, 435, 430]


@test
def t_preemption_flips_the_sign():
    on, _ = hogs(DEMOTE_BURST, block=2, preempt=True)
    off, _ = hogs(DEMOTE_BURST, block=2, preempt=False)
    assert on[0].cpu == 804 and off[0].cpu == 198
    assert on[0].cpu > max(x.cpu for x in on[1:])
    assert off[0].cpu < min(x.cpu for x in off[1:])


@test
def t_everyone_games_and_nobody_wins():
    got = []
    for k in range(0, 5):
        js = [Job(i, 0, 10 ** 9, yield_after=(9 if i < k else None))
              for i in range(4)]
        MLFQ(js, quanta=Q, policy=DEMOTE_BURST).run(1000)
        got.append([x.cpu for x in js])
    assert got[0] == [270, 270, 230, 230]
    assert got[1] == [970, 10, 10, 10]
    assert got[4] == [252, 252, 252, 244]
    assert max(got[4]) < max(got[0])


@test
def t_allotment_sweep_tops_out_at_round_robin():
    for a0 in (500, 1000):
        j, _ = hogs(DEMOTE_ALLOT, allots=(a0, 2 * a0, 4 * a0))
        assert [x.cpu for x in j] == [315, 345, 340], a0


@test
def t_both_policies_agree_when_nobody_yields():
    """The allotment defaults to one quantum, so a job that never releases
    the CPU is demoted at the same instant under both rules."""
    for pol in (DEMOTE_BURST, DEMOTE_ALLOT):
        js = [Job(i, 0, 10 ** 9) for i in range(3)]
        MLFQ(js, quanta=Q, policy=pol).run(1000)
        assert [x.cpu for x in js] == [350, 340, 310], (pol, [x.cpu for x in js])


@test
def t_deterministic_across_runs():
    a = [x.cpu for x in hogs(DEMOTE_BURST, boost=100)[0]]
    b = [x.cpu for x in hogs(DEMOTE_BURST, boost=100)[0]]
    assert a == b
    assert mixed(DEMOTE_ALLOT, 100) == mixed(DEMOTE_ALLOT, 100)


@test
def t_conservation_no_tick_invented():
    """Every tick is executed, idle, or spent switching -- nothing vanishes."""
    for sw in (0, 1, 2, 4):
        j, sim = hogs(DEMOTE_BURST, switch=sw)
        assert sum(x.cpu for x in j) + sim.idle + sim.overhead == 1000
        assert len(sim.trace) == 1000


if __name__ == "__main__":
    for fn in TESTS:
        fn()
    print("%d tests passed" % len(TESTS))