package arbo

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"math/big"
	"testing"
	"time"

	qt "github.com/frankban/quicktest"
	"github.com/iden3/go-merkletree/db/memory"
)

var debug = true

func debugTime(descr string, time1, time2 time.Duration) {
	if debug {
		fmt.Printf("%s was %f times faster than without AddBatch\n",
			descr, float64(time1)/float64(time2))
	}
}

func testInit(c *qt.C, n int) (*Tree, *Tree) {
	tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
	c.Assert(err, qt.IsNil)
	defer tree1.db.Close()

	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
	c.Assert(err, qt.IsNil)
	defer tree2.db.Close()

	// add the initial leafs to fill a bit the trees before calling the
	// AddBatch method
	for i := 0; i < n; i++ {
		k := BigIntToBytes(big.NewInt(int64(i)))
		v := BigIntToBytes(big.NewInt(int64(i * 2)))
		if err := tree1.Add(k, v); err != nil {
			c.Fatal(err)
		}
		if err := tree2.Add(k, v); err != nil {
			c.Fatal(err)
		}
	}
	return tree1, tree2
}

func TestAddBatchCaseA(t *testing.T) {
	c := qt.New(t)

	nLeafs := 1024

	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
	c.Assert(err, qt.IsNil)
	defer tree.db.Close()

	start := time.Now()
	for i := 0; i < nLeafs; i++ {
		k := BigIntToBytes(big.NewInt(int64(i)))
		v := BigIntToBytes(big.NewInt(int64(i * 2)))
		if err := tree.Add(k, v); err != nil {
			t.Fatal(err)
		}
	}
	time1 := time.Since(start)

	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
	c.Assert(err, qt.IsNil)
	defer tree2.db.Close()

	var keys, values [][]byte
	for i := 0; i < nLeafs; i++ {
		k := BigIntToBytes(big.NewInt(int64(i)))
		v := BigIntToBytes(big.NewInt(int64(i * 2)))
		keys = append(keys, k)
		values = append(values, v)
	}
	start = time.Now()
	indexes, err := tree2.AddBatch(keys, values)
	c.Assert(err, qt.IsNil)
	time2 := time.Since(start)
	debugTime("CASE A, AddBatch", time1, time2)
	c.Check(len(indexes), qt.Equals, 0)

	// check that both trees roots are equal
	c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
}

func TestAddBatchCaseANotPowerOf2(t *testing.T) {
	c := qt.New(t)

	nLeafs := 1027

	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
	c.Assert(err, qt.IsNil)
	defer tree.db.Close()

	for i := 0; i < nLeafs; i++ {
		k := BigIntToBytes(big.NewInt(int64(i)))
		v := BigIntToBytes(big.NewInt(int64(i * 2)))
		if err := tree.Add(k, v); err != nil {
			t.Fatal(err)
		}
	}

	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
	c.Assert(err, qt.IsNil)
	defer tree2.db.Close()

	var keys, values [][]byte
	for i := 0; i < nLeafs; i++ {
		k := BigIntToBytes(big.NewInt(int64(i)))
		v := BigIntToBytes(big.NewInt(int64(i * 2)))
		keys = append(keys, k)
		values = append(values, v)
	}
	indexes, err := tree2.AddBatch(keys, values)
	c.Assert(err, qt.IsNil)
	c.Check(len(indexes), qt.Equals, 0)

	// check that both trees roots are equal
	c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
}

func randomBytes(n int) []byte {
	b := make([]byte, n)
	_, err := rand.Read(b)
	if err != nil {
		panic(err)
	}
	return b
}

func TestBuildTreeBottomUpSingleThread(t *testing.T) {
	c := qt.New(t)
	tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
	c.Assert(err, qt.IsNil)
	defer tree1.db.Close()

	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
	c.Assert(err, qt.IsNil)
	defer tree2.db.Close()

	testvectorKeys := []string{
		"1c7c2265e368314ca58ed2e1f33a326f1220e234a566d55c3605439dbe411642",
		"2c9f0a578afff5bfa4e0992a43066460faaab9e8e500db0b16647c701cdb16bf",
		"1c45cb31f2fa39ec7b9ebf0fad40e0b8296016b5ce8844ae06ff77226379d9a5",
		"d8af98bbbb585129798ae54d5eabbc9d0561d583faf1663b3a3724d15bda4ec7",
	}
	var keys, values [][]byte
	for i := 0; i < len(testvectorKeys); i++ {
		key, err := hex.DecodeString(testvectorKeys[i])
		c.Assert(err, qt.IsNil)
		keys = append(keys, key)
		values = append(values, []byte{0})
	}

	for i := 0; i < len(keys); i++ {
		if err := tree1.Add(keys[i], values[i]); err != nil {
			t.Fatal(err)
		}
	}

	kvs, err := tree2.keysValuesToKvs(keys, values)
	c.Assert(err, qt.IsNil)
	sortKvs(kvs)

	tree2.tx, err = tree2.db.NewTx()
	c.Assert(err, qt.IsNil)
	// indexes, err := tree2.buildTreeBottomUpSingleThread(kvs)
	indexes, err := tree2.buildTreeBottomUp(4, kvs)
	c.Assert(err, qt.IsNil)
	// tree1.PrintGraphviz(nil)
	// tree2.PrintGraphviz(nil)

	c.Check(len(indexes), qt.Equals, 0)

	// check that both trees roots are equal
	c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
	// 15b6a23945ae6c81342b7eb14e70fff50812dc8791cb15ec791eb08f91784139
}

func TestAddBatchCaseATestVector(t *testing.T) {
	c := qt.New(t)
	tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
	c.Assert(err, qt.IsNil)
	defer tree1.db.Close()

	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
	c.Assert(err, qt.IsNil)
	defer tree2.db.Close()

	// leafs in 2nd level subtrees: [ 6, 0, 1, 1]
	testvectorKeys := []string{
		"1c7c2265e368314ca58ed2e1f33a326f1220e234a566d55c3605439dbe411642",
		"2c9f0a578afff5bfa4e0992a43066460faaab9e8e500db0b16647c701cdb16bf",
		"1c45cb31f2fa39ec7b9ebf0fad40e0b8296016b5ce8844ae06ff77226379d9a5",
		"d8af98bbbb585129798ae54d5eabbc9d0561d583faf1663b3a3724d15bda4ec7",
	}
	var keys, values [][]byte
	for i := 0; i < len(testvectorKeys); i++ {
		key, err := hex.DecodeString(testvectorKeys[i])
		c.Assert(err, qt.IsNil)
		keys = append(keys, key)
		values = append(values, []byte{0})
	}

	for i := 0; i < len(keys); i++ {
		if err := tree1.Add(keys[i], values[i]); err != nil {
			t.Fatal(err)
		}
	}

	indexes, err := tree2.AddBatch(keys, values)
	c.Assert(err, qt.IsNil)
	// tree1.PrintGraphviz(nil)
	// tree2.PrintGraphviz(nil)

	c.Check(len(indexes), qt.Equals, 0)

	// check that both trees roots are equal
	// fmt.Println(hex.EncodeToString(tree1.Root()))
	c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())

	//////

	// tree1, err = NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
	// c.Assert(err, qt.IsNil)
	// defer tree1.db.Close()
	//
	// tree2, err = NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
	// c.Assert(err, qt.IsNil)
	// defer tree2.db.Close()
	//
	// // leafs in 2nd level subtrees: [ 6, 0, 1, 1]
	// testvectorKeys = []string{
	//         "1c7c2265e368314ca58ed2e1f33a326f1220e234a566d55c3605439dbe411642",
	//         "2c9f0a578afff5bfa4e0992a43066460faaab9e8e500db0b16647c701cdb16bf",
	//         "9cb87ec67e875c61390edcd1ab517f443591047709a4d4e45b0f9ed980857b8e",
	//         "9b4e9e92e974a589f426ceeb4cb291dc24893513fecf8e8460992dcf52621d4d",
	//         "1c45cb31f2fa39ec7b9ebf0fad40e0b8296016b5ce8844ae06ff77226379d9a5",
	//         "d8af98bbbb585129798ae54d5eabbc9d0561d583faf1663b3a3724d15bda4ec7",
	//         "3cd55dbfb8f975f20a0925dfbdabe79fa2d51dd0268afbb8ba6b01de9dfcdd3c",
	//         "5d0a9d6d9f197c091bf054fac9cb60e11ec723d6610ed8578e617b4d46cb43d5",
	// }
	// keys = [][]byte{}
	// values = [][]byte{}
	// for i := 0; i < len(testvectorKeys); i++ {
	//         key, err := hex.DecodeString(testvectorKeys[i])
	//         c.Assert(err, qt.IsNil)
	//         keys = append(keys, key)
	//         values = append(values, []byte{0})
	// }
	//
	// for i := 0; i < len(keys); i++ {
	//         if err := tree1.Add(keys[i], values[i]); err != nil {
	//                 t.Fatal(err)
	//         }
	// }
	//
	// indexes, err = tree2.AddBatch(keys, values)
	// c.Assert(err, qt.IsNil)
	// tree1.PrintGraphviz(nil)
	// tree2.PrintGraphviz(nil)
	//
	// c.Check(len(indexes), qt.Equals, 0)
	//
	// // check that both trees roots are equal
	// // c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
}

func TestAddBatchCaseARandomKeys(t *testing.T) {
	c := qt.New(t)

	nLeafs := 8

	tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
	c.Assert(err, qt.IsNil)
	defer tree1.db.Close()

	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
	c.Assert(err, qt.IsNil)
	defer tree2.db.Close()

	var keys, values [][]byte
	for i := 0; i < nLeafs; i++ {
		keys = append(keys, randomBytes(32))
		// values = append(values, randomBytes(32))
		values = append(values, []byte{0})
		// fmt.Println("K", hex.EncodeToString(keys[i]))
	}

	// TMP:
	keys[0], _ = hex.DecodeString("1c7c2265e368314ca58ed2e1f33a326f1220e234a566d55c3605439dbe411642")
	keys[1], _ = hex.DecodeString("2c9f0a578afff5bfa4e0992a43066460faaab9e8e500db0b16647c701cdb16bf")
	keys[2], _ = hex.DecodeString("9cb87ec67e875c61390edcd1ab517f443591047709a4d4e45b0f9ed980857b8e")
	keys[3], _ = hex.DecodeString("9b4e9e92e974a589f426ceeb4cb291dc24893513fecf8e8460992dcf52621d4d")
	keys[4], _ = hex.DecodeString("1c45cb31f2fa39ec7b9ebf0fad40e0b8296016b5ce8844ae06ff77226379d9a5")
	keys[5], _ = hex.DecodeString("d8af98bbbb585129798ae54d5eabbc9d0561d583faf1663b3a3724d15bda4ec7")
	keys[6], _ = hex.DecodeString("3cd55dbfb8f975f20a0925dfbdabe79fa2d51dd0268afbb8ba6b01de9dfcdd3c")
	keys[7], _ = hex.DecodeString("5d0a9d6d9f197c091bf054fac9cb60e11ec723d6610ed8578e617b4d46cb43d5")

	for i := 0; i < len(keys); i++ {
		if err := tree1.Add(keys[i], values[i]); err != nil {
			t.Fatal(err)
		}
	}

	indexes, err := tree2.AddBatch(keys, values)
	c.Assert(err, qt.IsNil)
	// tree1.PrintGraphviz(nil)
	// tree2.PrintGraphviz(nil)

	c.Check(len(indexes), qt.Equals, 0)

	// check that both trees roots are equal
	c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
}

func TestAddBatchCaseB(t *testing.T) {
	c := qt.New(t)

	nLeafs := 1024
	initialNLeafs := 99 // TMP TODO use const minLeafsThreshold-1 once ready

	tree1, tree2 := testInit(c, initialNLeafs)

	start := time.Now()
	for i := initialNLeafs; i < nLeafs; i++ {
		k := BigIntToBytes(big.NewInt(int64(i)))
		v := BigIntToBytes(big.NewInt(int64(i * 2)))
		if err := tree1.Add(k, v); err != nil {
			t.Fatal(err)
		}
	}
	time1 := time.Since(start)

	// prepare the key-values to be added
	var keys, values [][]byte
	for i := initialNLeafs; i < nLeafs; i++ {
		k := BigIntToBytes(big.NewInt(int64(i)))
		v := BigIntToBytes(big.NewInt(int64(i * 2)))
		keys = append(keys, k)
		values = append(values, v)
	}
	start = time.Now()
	indexes, err := tree2.AddBatch(keys, values)
	c.Assert(err, qt.IsNil)
	time2 := time.Since(start)
	debugTime("CASE B, AddBatch", time1, time2)
	c.Check(len(indexes), qt.Equals, 0)

	// check that both trees roots are equal
	c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
}

func TestAddBatchCaseBRepeatedLeafs(t *testing.T) {
	c := qt.New(t)

	nLeafs := 1024
	initialNLeafs := 99 // TMP TODO use const minLeafsThreshold-1 once ready

	tree1, tree2 := testInit(c, initialNLeafs)

	for i := initialNLeafs; i < nLeafs; i++ {
		k := BigIntToBytes(big.NewInt(int64(i)))
		v := BigIntToBytes(big.NewInt(int64(i * 2)))
		if err := tree1.Add(k, v); err != nil {
			t.Fatal(err)
		}
	}

	// prepare the key-values to be added, including already existing keys
	var keys, values [][]byte
	for i := 0; i < nLeafs; i++ {
		k := BigIntToBytes(big.NewInt(int64(i)))
		v := BigIntToBytes(big.NewInt(int64(i * 2)))
		keys = append(keys, k)
		values = append(values, v)
	}
	indexes, err := tree2.AddBatch(keys, values)
	c.Assert(err, qt.IsNil)
	c.Check(len(indexes), qt.Equals, initialNLeafs)

	// check that both trees roots are equal
	c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
}

func TestCombineInKVSet(t *testing.T) {
	c := qt.New(t)

	var a, b, expected []kv
	for i := 0; i < 10; i++ {
		k := BigIntToBytes(big.NewInt(int64(i)))
		kv := kv{k: k}
		if i < 7 {
			a = append(a, kv)
		}
		if i >= 4 {
			b = append(b, kv)
		}
		expected = append(expected, kv)
	}

	r, invalids := combineInKVSet(a, b)
	for i := 0; i < len(r); i++ {
		c.Assert(r[i].k, qt.DeepEquals, expected[i].k)
	}
	c.Assert(len(invalids), qt.Equals, 7-4)
}

func TestGetKeysAtLevel(t *testing.T) {
	c := qt.New(t)

	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
	c.Assert(err, qt.IsNil)
	defer tree.db.Close()

	for i := 0; i < 32; i++ {
		k := BigIntToBytes(big.NewInt(int64(i)))
		v := BigIntToBytes(big.NewInt(int64(i * 2)))
		if err := tree.Add(k, v); err != nil {
			t.Fatal(err)
		}
	}

	keys, err := tree.getKeysAtLevel(2)
	c.Assert(err, qt.IsNil)
	expected := []string{
		"a5d5f14fce7348e40751496cf25d107d91b0bd043435b9577d778a01f8aa6111",
		"e9e8dd9b28a7f81d1ff34cb5cefc0146dd848b31031a427b79bdadb62e7f6910",
	}
	for i := 0; i < len(keys); i++ {
		c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
	}

	keys, err = tree.getKeysAtLevel(3)
	c.Assert(err, qt.IsNil)
	expected = []string{
		"9f12c13e52bca96ad4882a26558e48ab67ddd63e062b839207e893d961390f01",
		"16d246dd6826ec7346c7328f11c4261facf82d4689f33263ff6e207956a77f21",
		"4a22cc901c6337daa17a431fa20170684b710e5f551509511492ec24e81a8f2f",
		"470d61abcbd154977bffc9a9ec5a8daff0caabcf2a25e8441f604c79daa0f82d",
	}
	for i := 0; i < len(keys); i++ {
		c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
	}

	keys, err = tree.getKeysAtLevel(4)
	c.Assert(err, qt.IsNil)
	expected = []string{
		"7a5d1c81f7b96318012de3417e53d4f13df5b1337718651cd29d0cb0a66edd20",
		"3408213e4e844bdf3355eb8781c74e31626812898c2dbe141ed6d2c92256fc1c",
		"dfd8a4d0b6954a3e9f3892e655b58d456eeedf9367f27dfdd9bc2dd6a5577312",
		"9e99fbec06fb2a6725997c12c4995f62725eb4cce4808523a5a5e80cca64b007",
		"0befa1e070231dbf4e8ff841c05878cdec823e0c09594c24910a248b3ff5a628",
		"b7131b0a15c772a57005a4dc5d0d6dd4b3414f5d9ee7408ce5e86c5ab3520e04",
		"6d1abe0364077846a56bab1deb1a04883eb796b74fe531a7676a9a370f83ab21",
		"4270116394bede69cf9cd72069eca018238080380bef5de75be8dcbbe968e105",
	}
	for i := 0; i < len(keys); i++ {
		c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
	}
}

func TestSplitInBuckets(t *testing.T) {
	c := qt.New(t)

	nLeafs := 16
	kvs := make([]kv, nLeafs)
	for i := 0; i < nLeafs; i++ {
		k := BigIntToBytes(big.NewInt(int64(i)))
		v := BigIntToBytes(big.NewInt(int64(i * 2)))
		keyPath := make([]byte, 32)
		copy(keyPath[:], k)
		kvs[i].pos = i
		kvs[i].keyPath = k
		kvs[i].k = k
		kvs[i].v = v
	}

	// check keyToBucket results for 4 buckets & 8 keys
	c.Assert(keyToBucket(kvs[0].k, 4), qt.Equals, 0)
	c.Assert(keyToBucket(kvs[1].k, 4), qt.Equals, 2)
	c.Assert(keyToBucket(kvs[2].k, 4), qt.Equals, 1)
	c.Assert(keyToBucket(kvs[3].k, 4), qt.Equals, 3)
	c.Assert(keyToBucket(kvs[4].k, 4), qt.Equals, 0)
	c.Assert(keyToBucket(kvs[5].k, 4), qt.Equals, 2)
	c.Assert(keyToBucket(kvs[6].k, 4), qt.Equals, 1)
	c.Assert(keyToBucket(kvs[7].k, 4), qt.Equals, 3)

	// check keyToBucket results for 8 buckets & 8 keys
	c.Assert(keyToBucket(kvs[0].k, 8), qt.Equals, 0)
	c.Assert(keyToBucket(kvs[1].k, 8), qt.Equals, 4)
	c.Assert(keyToBucket(kvs[2].k, 8), qt.Equals, 2)
	c.Assert(keyToBucket(kvs[3].k, 8), qt.Equals, 6)
	c.Assert(keyToBucket(kvs[4].k, 8), qt.Equals, 1)
	c.Assert(keyToBucket(kvs[5].k, 8), qt.Equals, 5)
	c.Assert(keyToBucket(kvs[6].k, 8), qt.Equals, 3)
	c.Assert(keyToBucket(kvs[7].k, 8), qt.Equals, 7)

	buckets := splitInBuckets(kvs, 4)

	expected := [][]string{
		{
			"00000000", // bucket 0
			"08000000",
			"04000000",
			"0c000000",
		},
		{
			"02000000", // bucket 1
			"0a000000",
			"06000000",
			"0e000000",
		},
		{
			"01000000", // bucket 2
			"09000000",
			"05000000",
			"0d000000",
		},
		{
			"03000000", // bucket 3
			"0b000000",
			"07000000",
			"0f000000",
		},
	}

	for i := 0; i < len(buckets); i++ {
		sortKvs(buckets[i])
		c.Assert(len(buckets[i]), qt.Equals, len(expected[i]))
		for j := 0; j < len(buckets[i]); j++ {
			c.Check(hex.EncodeToString(buckets[i][j].k[:4]), qt.Equals, expected[i][j])
		}
	}
}

func TestAddBatchCaseC(t *testing.T) {
	c := qt.New(t)

	nLeafs := 1024
	initialNLeafs := 101 // TMP TODO use const minLeafsThreshold+1 once ready

	tree1, tree2 := testInit(c, initialNLeafs)

	start := time.Now()
	for i := initialNLeafs; i < nLeafs; i++ {
		k := BigIntToBytes(big.NewInt(int64(i)))
		v := BigIntToBytes(big.NewInt(int64(i * 2)))
		if err := tree1.Add(k, v); err != nil {
			t.Fatal(err)
		}
	}
	time1 := time.Since(start)

	// prepare the key-values to be added
	var keys, values [][]byte
	for i := initialNLeafs; i < nLeafs; i++ {
		k := BigIntToBytes(big.NewInt(int64(i)))
		v := BigIntToBytes(big.NewInt(int64(i * 2)))
		keys = append(keys, k)
		values = append(values, v)
	}
	start = time.Now()
	indexes, err := tree2.AddBatch(keys, values)
	c.Assert(err, qt.IsNil)
	time2 := time.Since(start)
	debugTime("CASE C, AddBatch", time1, time2)
	c.Check(len(indexes), qt.Equals, 0)

	// check that both trees roots are equal
	c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
}

func TestAddBatchCaseD(t *testing.T) {
	c := qt.New(t)

	nLeafs := 4096
	initialNLeafs := 900

	tree1, tree2 := testInit(c, initialNLeafs)

	start := time.Now()
	for i := initialNLeafs; i < nLeafs; i++ {
		k := BigIntToBytes(big.NewInt(int64(i)))
		v := BigIntToBytes(big.NewInt(int64(i * 2)))
		if err := tree1.Add(k, v); err != nil {
			t.Fatal(err)
		}
	}
	time1 := time.Since(start)

	// prepare the key-values to be added
	var keys, values [][]byte
	for i := initialNLeafs; i < nLeafs; i++ {
		k := BigIntToBytes(big.NewInt(int64(i)))
		v := BigIntToBytes(big.NewInt(int64(i * 2)))
		keys = append(keys, k)
		values = append(values, v)
	}
	start = time.Now()
	indexes, err := tree2.AddBatch(keys, values)
	c.Assert(err, qt.IsNil)
	time2 := time.Since(start)
	debugTime("CASE D, AddBatch", time1, time2)
	c.Check(len(indexes), qt.Equals, 0)

	// check that both trees roots are equal
	c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
}

func TestAddBatchCaseE(t *testing.T) {
	c := qt.New(t)

	nLeafs := 4096
	initialNLeafs := 900

	tree1, _ := testInit(c, initialNLeafs)

	start := time.Now()
	for i := initialNLeafs; i < nLeafs; i++ {
		k := BigIntToBytes(big.NewInt(int64(i)))
		v := BigIntToBytes(big.NewInt(int64(i * 2)))
		if err := tree1.Add(k, v); err != nil {
			t.Fatal(err)
		}
	}
	time1 := time.Since(start)

	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
	c.Assert(err, qt.IsNil)
	defer tree2.db.Close()

	var keys, values [][]byte
	// add the initial leafs to fill a bit the tree before calling the
	// AddBatch method
	for i := 0; i < initialNLeafs; i++ {
		k := BigIntToBytes(big.NewInt(int64(i)))
		v := BigIntToBytes(big.NewInt(int64(i * 2)))
		// use only the keys of one bucket, store the not used ones for
		// later
		if i%4 != 0 {
			keys = append(keys, k)
			values = append(values, v)
			continue
		}
		if err := tree2.Add(k, v); err != nil {
			t.Fatal(err)
		}
	}

	for i := initialNLeafs; i < nLeafs; i++ {
		k := BigIntToBytes(big.NewInt(int64(i)))
		v := BigIntToBytes(big.NewInt(int64(i * 2)))
		keys = append(keys, k)
		values = append(values, v)
	}
	start = time.Now()
	indexes, err := tree2.AddBatch(keys, values)
	c.Assert(err, qt.IsNil)
	time2 := time.Since(start)
	debugTime("CASE E, AddBatch", time1, time2)
	c.Check(len(indexes), qt.Equals, 0)

	// check that both trees roots are equal
	c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
}

func TestFlp2(t *testing.T) {
	c := qt.New(t)
	c.Assert(flp2(31), qt.Equals, 16)
	c.Assert(flp2(32), qt.Equals, 32)
	c.Assert(flp2(33), qt.Equals, 32)
	c.Assert(flp2(63), qt.Equals, 32)
	c.Assert(flp2(64), qt.Equals, 64)
	c.Assert(flp2(9000), qt.Equals, 8192)
}

// func printLeafs(name string, t *Tree) {
//         w := bytes.NewBufferString("")
//
//         err := t.Iterate(func(k, v []byte) {
//                 if v[0] != PrefixValueLeaf {
//                         return
//                 }
//                 leafK, _ := readLeafValue(v)
//                 fmt.Fprintf(w, hex.EncodeToString(leafK[:4])+"\n")
//         })
//         if err != nil {
//                 panic(err)
//         }
//         err = ioutil.WriteFile(name, w.Bytes(), 0644)
//         if err != nil {
//                 panic(err)
//         }
//
// }

// func TestComputeCosts(t *testing.T) {
//         fmt.Println(computeSimpleAddCost(10))
//         fmt.Println(computeBottomUpAddCost(10))
//
//         fmt.Println(computeSimpleAddCost(1024))
//         fmt.Println(computeBottomUpAddCost(1024))
// }

// TODO test tree with nLeafs > minLeafsThreshold, but that at level L, there is
// less keys than nBuckets (so CASE C could be applied if first few leafs are
// added to balance the tree)

// TODO test adding batch with repeated keys in the batch
// TODO test adding batch with multiple invalid keys