import numpy as np

from tools.ab_test_planner import fixed_sample_size, main, parse_args, plan, top_two_allocation


def test_top_two_allocation_properties():
    p = np.array([[0.98, 0.01, 0.01], [1 / 3, 1 / 3, 1 / 3]])
    psi = top_two_allocation(p, beta=0.5)
    np.testing.assert_allclose(psi.sum(axis=1), 1)
    # when one arm is almost surely best, it still gets ~beta of the traffic
    assert 0.45 < psi[0, 0] < 0.55
    np.testing.assert_allclose(psi[1], 1 / 3)


def test_sample_size_textbook_value():
    # 5% -> 6% (20% relative), alpha .05 two-sided, power .8: ~8,155 per arm
    n = fixed_sample_size(0.05, 0.20, 1)
    assert 8000 < n < 8300


def test_planner_end_to_end(tmp_path):
    args = parse_args(["--baseline", "0.1", "--lifts", "0.3", "--traffic", "2000",
                       "--days", "14", "--sims", "40", "--draws", "300"])
    rates, out = plan(args)
    # the fixed test is designed for 80% power at this lift: it ships the
    # winner about 80% of the time; the bandits keep learning and do better
    assert 0.6 < out["fixed"]["correct"].mean() < 0.95
    assert out["ts"]["correct"].mean() > 0.9 and out["ttts"]["correct"].mean() > 0.9
    assert out["ts"]["lost"].sum(axis=1).mean() < out["fixed"]["lost"].sum(axis=1).mean()
    assert main(["--sims", "20", "--draws", "200", "--days", "7",
                 "--plot", str(tmp_path / "p.png")]) == 0
    assert (tmp_path / "p.png").exists()