import random
from fenwick import Fenwick


def test_basic_update_prefix():
    ft = Fenwick(5)
    for i in range(1, 6):
        ft.update(i, i)
    assert ft.prefix(0) == 0
    assert ft.prefix(1) == 1
    assert ft.prefix(3) == 6
    assert ft.prefix(5) == 15
    assert ft.total() == 15


def test_range_sum():
    ft = Fenwick(6)
    arr = [3, 1, 4, 1, 5, 9]
    for i, v in enumerate(arr, start=1):
        ft.update(i, v)
    assert ft.range_sum(2, 4) == 1 + 4 + 1
    assert ft.range_sum(1, 6) == sum(arr)
    assert ft.range_sum(3, 3) == 4


def test_from_array():
    arr = [2, 5, 1, 3, 7]
    ft = Fenwick.from_array(arr)
    for i in range(1, len(arr) + 1):
        assert ft.prefix(i) == sum(arr[:i])
    # update works on built tree
    ft.update(3, 10)
    assert ft.prefix(3) == 2 + 5 + 1 + 10


def test_lower_bound():
    ft = Fenwick(5)
    for i in range(1, 6):
        ft.update(i, 1)  # all ones
    # prefix: 1,2,3,4,5
    assert ft.lower_bound(1) == 1
    assert ft.lower_bound(2) == 2
    assert ft.lower_bound(5) == 5
    assert ft.lower_bound(6) == 6  # none


def test_lower_bound_weighted():
    ft = Fenwick(4)
    ft.update(1, 5)
    ft.update(2, 2)
    ft.update(3, 7)
    ft.update(4, 1)
    # prefixes: 5,7,14,15
    assert ft.lower_bound(5) == 1
    assert ft.lower_bound(6) == 2
    assert ft.lower_bound(7) == 2
    assert ft.lower_bound(8) == 3
    assert ft.lower_bound(15) == 4
    assert ft.lower_bound(16) == 5


def test_random_vs_brute():
    n = 50
    arr = [0] * n
    ft = Fenwick(n)
    rng = random.Random(42)
    for _ in range(2000):
        op = rng.random()
        if op < 0.5:
            i = rng.randrange(1, n + 1)
            d = rng.randrange(-5, 6)
            arr[i - 1] += d
            ft.update(i, d)
        else:
            # check a random prefix
            i = rng.randrange(0, n + 1)
            assert ft.prefix(i) == sum(arr[:i]), (i, ft.prefix(i), sum(arr[:i]))
    # lower_bound against non-negative only: rebuild with non-negative
    arr2 = [rng.randrange(0, 6) for _ in range(n)]
    ft2 = Fenwick.from_array(arr2)
    # verify lower_bound for every target 1..total
    total = sum(arr2)
    for t in range(1, total + 1):
        # brute
        s = 0
        brute = None
        for idx, v in enumerate(arr2, start=1):
            s += v
            if s >= t:
                brute = idx
                break
        if brute is None:
            brute = n + 1
        assert ft2.lower_bound(t) == brute, (t, ft2.lower_bound(t), brute)


def test_bounds():
    ft = Fenwick(3)
    try:
        ft.update(0, 1)
    except IndexError:
        pass
    else:
        raise AssertionError("expected IndexError")
    try:
        ft.update(4, 1)
    except IndexError:
        pass
    else:
        raise AssertionError("expected IndexError")
    try:
        ft.prefix(4)
    except IndexError:
        pass
    else:
        raise AssertionError("expected IndexError")
    try:
        ft.range_sum(2, 1)
    except ValueError:
        pass
    else:
        raise AssertionError("expected ValueError")


def test_empty():
    ft = Fenwick(0)
    assert ft.total() == 0
    assert ft.prefix(0) == 0


if __name__ == "__main__":
    # run all test_ functions
    import sys
    g = globals()
    for name in sorted(g):
        if name.startswith("test_"):
            g[name]()
    print("ALL TESTS PASSED")