package corpus

import (
	"math/rand"
	"testing"
)

func naivePrefixSum(a []int64, r int) int64 {
	var s int64
	for i := 1; i <= r; i++ {
		s += a[i]
	}
	return s
}

func TestFenwickBasic(t *testing.T) {
	var ft fenwikTree

	values := []int64{0, 5, -3, 7, 10}
	for _, v := range values[1:] {
		ft.addNew(v)
	}

	if ft.size != len(values)-1 {
		t.Fatalf("size mismatch: got %d want %d", ft.size, len(values)-1)
	}

	for i := 1; i < ft.size; i++ {
		got := ft.prefixSumTo(i)
		want := naivePrefixSum(values, i)
		if got != want {
			t.Fatalf("prefixSumTo(%d) = %d, want %d", i, got, want)
		}
	}
}

func TestFenwickExtend(t *testing.T) {
	var ft fenwikTree

	n := 1000
	ref := make([]int64, 1)

	for i := 1; i <= n; i++ {
		ft.addNew(int64(i))
		ref = append(ref, int64(i))

		if ft.prefixSumTo(i) != naivePrefixSum(ref, i) {
			t.Fatalf("extend failed at i=%d, fw=%d and naive=%d", i, ft.prefixSumTo(i), naivePrefixSum(ref, i))
		}
	}
}

func TestFenwickRandom(t *testing.T) {
	var ft fenwikTree
	ft.extendAndRebuild()

	ref := make([]int64, 1)

	rng := rand.New(rand.NewSource(1))

	for step := range 10_000 {
		if rng.Intn(2) == 0 || ft.size <= 2 {
			v := int64(rng.Intn(1000) - 500)
			ft.addNew(v)
			ref = append(ref, v)
		} else {
			i := rng.Intn(ft.size-1) + 1
			v := int64(rng.Intn(1000) - 500)
			ft.add(i, v)
			ref[i] += v
		}

		i := rng.Intn(ft.size)
		if ft.prefixSumTo(i) != naivePrefixSum(ref, i) {
			t.Fatalf("mismatch at step %d, i=%d", step, i)
		}

		if ft.prefixSumTo(ft.size) != ft.fullSum {
			t.Fatalf("fullSum mismatch %d, sm=%d, fullSum=%d, naive=%d", step, ft.prefixSumTo(ft.size), ft.fullSum, naivePrefixSum(ref, ft.size))
		}
	}
}