// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package iterable

import (
	"container/vector";
	"testing";
)

func TestArrayTypes(t *testing.T) {
	// Test that conversion works correctly.
	bytes := ByteArray([]byte{1, 2, 3});
	if x := Data(bytes)[1].(byte); x != 2 {
		t.Error("Data(bytes)[1].(byte) = %v, want 2", x)
	}
	ints := IntArray([]int{1, 2, 3});
	if x := Data(ints)[2].(int); x != 3 {
		t.Error("Data(ints)[2].(int) = %v, want 3", x)
	}
	floats := FloatArray([]float{1, 2, 3});
	if x := Data(floats)[0].(float); x != 1 {
		t.Error("Data(floats)[0].(float) = %v, want 1", x)
	}
	strings := StringArray([]string{"a", "b", "c"});
	if x := Data(strings)[1].(string); x != "b" {
		t.Error(`Data(strings)[1].(string) = %q, want "b"`, x)
	}
}

var (
	oneToFive	= IntArray{1, 2, 3, 4, 5};
	sixToTen	= IntArray{6, 7, 8, 9, 10};
	elevenToTwenty	= IntArray{11, 12, 13, 14, 15, 16, 17, 18, 19, 20};
)

func isNegative(n interface{}) bool	{ return n.(int) < 0 }
func isPositive(n interface{}) bool	{ return n.(int) > 0 }
func isAbove3(n interface{}) bool	{ return n.(int) > 3 }
func isEven(n interface{}) bool		{ return n.(int)%2 == 0 }
func doubler(n interface{}) interface{}	{ return n.(int) * 2 }
func addOne(n interface{}) interface{}	{ return n.(int) + 1 }
func adder(acc interface{}, n interface{}) interface{} {
	return acc.(int) + n.(int)
}

// A stream of the natural numbers: 0, 1, 2, 3, ...
type integerStream struct{}

func (i integerStream) Iter() <-chan interface{} {
	ch := make(chan interface{});
	go func() {
		for i := 0; ; i++ {
			ch <- i
		}
	}();
	return ch;
}

func TestAll(t *testing.T) {
	if !All(oneToFive, isPositive) {
		t.Error("All(oneToFive, isPositive) == false")
	}
	if All(oneToFive, isAbove3) {
		t.Error("All(oneToFive, isAbove3) == true")
	}
}

func TestAny(t *testing.T) {
	if Any(oneToFive, isNegative) {
		t.Error("Any(oneToFive, isNegative) == true")
	}
	if !Any(oneToFive, isEven) {
		t.Error("Any(oneToFive, isEven) == false")
	}
}

func assertArraysAreEqual(t *testing.T, res []interface{}, expected []int) {
	if len(res) != len(expected) {
		t.Errorf("len(res) = %v, want %v", len(res), len(expected));
		goto missing;
	}
	for i := range res {
		if v := res[i].(int); v != expected[i] {
			t.Errorf("res[%v] = %v, want %v", i, v, expected[i]);
			goto missing;
		}
	}
	return;
missing:
	t.Errorf("res = %v\nwant  %v", res, expected);
}

func TestFilter(t *testing.T) {
	ints := integerStream{};
	moreInts := Filter(ints, isAbove3).Iter();
	res := make([]interface{}, 3);
	for i := 0; i < 3; i++ {
		res[i] = <-moreInts
	}
	assertArraysAreEqual(t, res, []int{4, 5, 6});
}

func TestFind(t *testing.T) {
	ints := integerStream{};
	first := Find(ints, isAbove3);
	if first.(int) != 4 {
		t.Errorf("Find(ints, isAbove3) = %v, want 4", first)
	}
}

func TestInject(t *testing.T) {
	res := Inject(oneToFive, 0, adder);
	if res.(int) != 15 {
		t.Errorf("Inject(oneToFive, 0, adder) = %v, want 15", res)
	}
}

func TestMap(t *testing.T) {
	res := Data(Map(Map(oneToFive, doubler), addOne));
	assertArraysAreEqual(t, res, []int{3, 5, 7, 9, 11});
}

func TestPartition(t *testing.T) {
	ti, fi := Partition(oneToFive, isEven);
	assertArraysAreEqual(t, Data(ti), []int{2, 4});
	assertArraysAreEqual(t, Data(fi), []int{1, 3, 5});
}

func TestTake(t *testing.T) {
	res := Take(oneToFive, 2);
	assertArraysAreEqual(t, Data(res), []int{1, 2});
	assertArraysAreEqual(t, Data(res), []int{1, 2});	// second test to ensure that .Iter() returns a new channel

	// take none
	res = Take(oneToFive, 0);
	assertArraysAreEqual(t, Data(res), []int{});

	// try to take more than available
	res = Take(oneToFive, 20);
	assertArraysAreEqual(t, Data(res), oneToFive);
}

func TestTakeWhile(t *testing.T) {
	// take some
	res := TakeWhile(oneToFive, func(v interface{}) bool { return v.(int) <= 3 });
	assertArraysAreEqual(t, Data(res), []int{1, 2, 3});
	assertArraysAreEqual(t, Data(res), []int{1, 2, 3});	// second test to ensure that .Iter() returns a new channel

	// take none
	res = TakeWhile(oneToFive, func(v interface{}) bool { return v.(int) > 3000 });
	assertArraysAreEqual(t, Data(res), []int{});

	// take all
	res = TakeWhile(oneToFive, func(v interface{}) bool { return v.(int) < 3000 });
	assertArraysAreEqual(t, Data(res), oneToFive);
}

func TestDrop(t *testing.T) {
	// drop none
	res := Drop(oneToFive, 0);
	assertArraysAreEqual(t, Data(res), oneToFive);
	assertArraysAreEqual(t, Data(res), oneToFive);	// second test to ensure that .Iter() returns a new channel

	// drop some
	res = Drop(oneToFive, 2);
	assertArraysAreEqual(t, Data(res), []int{3, 4, 5});
	assertArraysAreEqual(t, Data(res), []int{3, 4, 5});	// second test to ensure that .Iter() returns a new channel

	// drop more than available
	res = Drop(oneToFive, 88);
	assertArraysAreEqual(t, Data(res), []int{});
}

func TestDropWhile(t *testing.T) {
	// drop some
	res := DropWhile(oneToFive, func(v interface{}) bool { return v.(int) < 3 });
	assertArraysAreEqual(t, Data(res), []int{3, 4, 5});
	assertArraysAreEqual(t, Data(res), []int{3, 4, 5});	// second test to ensure that .Iter() returns a new channel

	// test case where all elements are dropped
	res = DropWhile(oneToFive, func(v interface{}) bool { return v.(int) < 100 });
	assertArraysAreEqual(t, Data(res), []int{});

	// test case where none are dropped
	res = DropWhile(oneToFive, func(v interface{}) bool { return v.(int) > 1000 });
	assertArraysAreEqual(t, Data(res), oneToFive);
}

func TestCycle(t *testing.T) {
	res := Cycle(oneToFive);
	exp := []int{1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4};

	// read the first nineteen values from the iterable
	out := make([]interface{}, 19);
	for i, it := 0, res.Iter(); i < 19; i++ {
		out[i] = <-it
	}
	assertArraysAreEqual(t, out, exp);

	res2 := Cycle(sixToTen);
	exp2 := []int{6, 7, 8, 9, 10, 6, 7, 8, 9, 10, 6, 7, 8, 9, 10, 6, 7, 8, 9};
	for i, it := 0, res2.Iter(); i < 19; i++ {
		out[i] = <-it
	}
	assertArraysAreEqual(t, out, exp2);

	// ensure first iterator was not harmed
	for i, it := 0, res.Iter(); i < 19; i++ {
		out[i] = <-it
	}
	assertArraysAreEqual(t, out, exp);
}

func TestChain(t *testing.T) {

	exp := []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20};
	res := Chain([]Iterable{oneToFive, sixToTen, elevenToTwenty});
	assertArraysAreEqual(t, Data(res), exp);

	// reusing the same iterator should produce the same result again
	assertArraysAreEqual(t, Data(res), exp);

	// test short read from Chain
	i := 0;
	out := make([]interface{}, 4);
	for v := range res.Iter() {
		out[i] = v;
		i++;
		if i == len(out) {
			break
		}
	}
	assertArraysAreEqual(t, out, exp[0:4]);

	// test zero length array
	res = Chain([]Iterable{});
	assertArraysAreEqual(t, Data(res), []int{});
}

func TestZipWith(t *testing.T) {
	exp := []int{7, 9, 11, 13, 15};

	// f with 2 args and 1 return value
	f := func(a, b interface{}) interface{} { return a.(int) + b.(int) };
	res := ZipWith2(f, oneToFive, sixToTen);
	assertArraysAreEqual(t, Data(res), exp);

	// test again to make sure returns new iter each time
	assertArraysAreEqual(t, Data(res), exp);

	// test a function with 3 args
	f2 := func(a, b, c interface{}) interface{} { return a.(int) + b.(int) + c.(int) };
	res = ZipWith3(f2, oneToFive, sixToTen, oneToFive);
	exp = []int{8, 11, 14, 17, 20};
	assertArraysAreEqual(t, Data(res), exp);

	// test a function with multiple values returned
	f3 := func(a, b interface{}) interface{} { return ([]interface{}{a.(int) + 1, b.(int) + 1}) };
	res = ZipWith2(f3, oneToFive, sixToTen);

	exp2 := [][]int{[]int{2, 7}, []int{3, 8}, []int{4, 9}, []int{5, 10}, []int{6, 11}};
	i := 0;
	for v := range res.Iter() {
		out := v.([]interface{});
		assertArraysAreEqual(t, out, exp2[i]);
		i++;
	}

	// test different length iterators--should stop after shortest is exhausted
	res = ZipWith2(f, elevenToTwenty, oneToFive);
	exp = []int{12, 14, 16, 18, 20};
	assertArraysAreEqual(t, Data(res), exp);
}

func TestSlice(t *testing.T) {
	out := Data(Slice(elevenToTwenty, 2, 6));
	exp := []int{13, 14, 15, 16};
	assertArraysAreEqual(t, out, exp);

	// entire iterable
	out = Data(Slice(elevenToTwenty, 0, len(elevenToTwenty)));
	exp = []int{11, 12, 13, 14, 15, 16, 17, 18, 19, 20};
	assertArraysAreEqual(t, out, exp);

	// empty slice at offset 0
	exp = []int{};
	out = Data(Slice(elevenToTwenty, 0, 0));
	assertArraysAreEqual(t, out, exp);

	// slice upper bound exceeds length of iterable
	exp = []int{1, 2, 3, 4, 5};
	out = Data(Slice(oneToFive, 0, 88));
	assertArraysAreEqual(t, out, exp);

	// slice upper bounce is lower than lower bound
	exp = []int{};
	out = Data(Slice(oneToFive, 93, 4));
	assertArraysAreEqual(t, out, exp);

	// slice lower bound is greater than len of iterable
	exp = []int{};
	out = Data(Slice(oneToFive, 93, 108));
	assertArraysAreEqual(t, out, exp);
}

func TestRepeat(t *testing.T) {
	res := Repeat(42);
	i := 0;
	for v := range res.Iter() {
		if v.(int) != 42 {
			t.Fatal("Repeat returned the wrong value")
		}
		if i == 9 {
			break
		}
		i++;
	}
}

func TestRepeatTimes(t *testing.T) {
	res := RepeatTimes(84, 9);
	exp := []int{84, 84, 84, 84, 84, 84, 84, 84, 84};
	assertArraysAreEqual(t, Data(res), exp);
	assertArraysAreEqual(t, Data(res), exp);	// second time to ensure new iter is returned

	// 0 repeat
	res = RepeatTimes(7, 0);
	exp = []int{};
	assertArraysAreEqual(t, Data(res), exp);

	// negative repeat
	res = RepeatTimes(7, -3);
	exp = []int{};
	assertArraysAreEqual(t, Data(res), exp);
}

// a type that implements Key for ints
type intkey struct{}

func (v intkey) Key(a interface{}) interface{} {
	return a
}
func (v intkey) Equal(a, b interface{}) bool	{ return a.(int) == b.(int) }

func TestGroupBy(t *testing.T) {
	in := IntArray{1, 2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 5};
	exp := [][]int{[]int{1}, []int{2, 2}, []int{3, 3, 3}, []int{4, 4, 4, 4}, []int{5, 5, 5, 5, 5}};
	i := 0;
	for x := range GroupBy(in, intkey{}).Iter() {
		gr := x.(Group);
		if gr.Key.(int) != i+1 {
			t.Fatal("group key wrong; expected", i+1, "but got", gr.Key.(int))
		}
		vals := Data(gr.Vals);
		assertArraysAreEqual(t, vals, exp[i]);
		i++;
	}
	if i != 5 {
		t.Fatal("did not return expected number of groups")
	}

	// test 0 length Iterable
	for _ = range GroupBy(IntArray([]int{}), &intkey{}).Iter() {
		t.Fatal("iterator should be empty")
	}

	// test case with only uniques
	var out vector.Vector;
	for x := range GroupBy(elevenToTwenty, intkey{}).Iter() {
		out.Push(x.(Group).Key)
	}
	assertArraysAreEqual(t, out.Data(), elevenToTwenty);
}

func TestUnique(t *testing.T) {
	in := IntArray([]int{1, 2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 5});
	exp := []int{1, 2, 3, 4, 5};
	res := Unique(in, intkey{});
	assertArraysAreEqual(t, Data(res), exp);
	assertArraysAreEqual(t, Data(res), exp);	// second time to ensure new iter is returned

	// test case with only uniques
	res = Unique(elevenToTwenty, intkey{});
	assertArraysAreEqual(t, Data(res), elevenToTwenty);
}