from pamle.communication import protocol_accounting, sufficient_exponent, sufficient_exponent_stale, sufficient_exponent_single


def test_exponents_table():
    assert sufficient_exponent(1, 0) == 0.5 and sufficient_exponent(1, 1) == 0.75 and sufficient_exponent(1, 2) == 0.875
    assert sufficient_exponent(2, 0) == 0.75 and sufficient_exponent(2, 1) == 0.875 and sufficient_exponent(3, 1) == 11 / 12
    assert abs(sufficient_exponent_stale(2, 1) - 5 / 6) < 1e-15 and abs(sufficient_exponent_single(1, 1) - 2 / 3) < 1e-15
    assert sufficient_exponent_stale(1, 1) == sufficient_exponent(1, 1)   # r=1: both 3/4


def test_accounting_exact_vs_stale():
    K, p = 100, 10
    e = protocol_accounting(K, p, r=2, q=1, hessian_mode="exact")
    s = protocol_accounting(K, p, r=2, q=1, hessian_mode="stale")
    assert e["uplink_numbers"] == K * (p + p + p * (p + 1) // 2)
    assert s["uplink_numbers"] == K * (p + p * (p + 1) // 2 + p)
    e2 = protocol_accounting(K, p, r=1, q=2, hessian_mode="exact")
    s2 = protocol_accounting(K, p, r=1, q=2, hessian_mode="stale")
    assert s2["uplink_numbers"] < e2["uplink_numbers"]           # stale Hessian saves bytes only when q >= 2
    assert e["rounds"] == 2 and e2["rounds"] == 3
    assert protocol_accounting(K, p, r=3, q=0)["local_fits_per_worker"] == 7
    assert protocol_accounting(K, p, r=1, q=0)["uplink_numbers"] == K * p