package flow_test

import (
	"bytes"
	"encoding/json"
	"testing"

	"flow"

	vecasm "github.com/gohxs/vec-benchmark/asm"

	"flow/internal/assert"
	"flow/registry"
)

func init() {
	assert.Quiet = true
}

func TestSerialize(t *testing.T) {
	// Does not text yet
	f := flow.New()
	var1 := f.Var("var1", []float32{4, 4, 4})

	c1 := f.Const([]float32{1, 2, 3})
	c2 := f.Const([]float32{2, 2, 2})

	op1 := f.Op("vecmul", // op:0 - expected: [12,16,20,24]
		f.Var("vec1", []float32{4, 4, 4, 4}),
		f.Op("vecadd", // op:1 - expected: [3,4,5,6]
			f.Const([]float32{1, 2, 3, 4}),
			f.Const([]float32{2, 2, 2, 2}),
		),
	)
	mul1 := f.Op("vecmul", c1, op1)       // op:2 - expected 12, 32, 60, 0
	mul2 := f.Op("vecmul", mul1, var1)    // op:3 - expected 48, 128, 240, 0
	mul3 := f.Op("vecmul", c2, mul2)      // op:4 - expected 96, 256, 480, 0
	mul4 := f.Op("vecmul", mul3, f.In(0)) // op:5 - expected 96, 512, 1440,0

	s := bytes.NewBuffer(nil)
	f.Analyse(s, []float32{1, 2, 3, 4})
	t.Log(s)
	res := mul4.Process([]float32{1, 2, 3, 4})

	t.Log("Res:", res)
	t.Log("Flow:\n", f)

	ret := bytes.NewBuffer(nil)
	e := json.NewEncoder(ret)
	e.Encode(f)

	t.Log("Flow:", ret)

}
func TestConst(t *testing.T) {
	assert := assert.A(t)
	f := flow.New()

	c := f.Const(1)
	res := c.Process()
	//Expect(t).Eq(res, 1)
	assert.Eq(res, 1, "It should be one")
}
func TestOp(t *testing.T) {
	assert := assert.A(t)
	f := flow.New()

	add := f.Op("vecadd",
		f.Op("vecmul",
			[]float32{1, 2, 3},
			[]float32{2, 2, 2},
		),
		[]float32{1, 2, 3},
	)
	res, err := f.Run(add)
	assert.Eq(err, nil)

	test := []float32{3, 6, 9}
	assert.Eq(test, res)
}

func TestVariable(t *testing.T) {
	f := flow.New()
	v := f.Var("v1", 1)

	res := v.Process()
	assert.A(t).Eq(res, 1)

	v.Set(2)
	res = v.Process()
	assert.Eq(t, res, 2)
}

func TestCache(t *testing.T) {
	f := flow.New()
	{
		r := f.Op("inc")
		var res interface{}
		for i := 1; i < 5; i++ {
			res = r.Process()
			assert.Eq(t, res, i)
		}
	}
	{
		var res flow.Data
		inc := f.Op("inc")
		add := f.Op("add", inc, inc)
		res = add.Process() // 1+1
		assert.Eq(t, res, 2)
		res = add.Process() // 2+2
		assert.Eq(t, res, 4)
	}
}
func TestDefOp(t *testing.T) {
	assert := assert.A(t)
	f := flow.New()

	f.DefOp("2", "vecadd", []float32{1, 1, 1}, []float32{2, 2, 2}) // r:3 3 3
	assert.Eq(f.Err(), nil, "doing DefOp")

	f.DefOp("1", "vecadd", []float32{1, 2, 3}, f.Res("2")) // r: 4 5 6
	assert.Eq(f.Err(), nil, "doing DefOp")

	op := f.Op("vecmul", f.Res("1"), []float32{2, 2, 2}) //r:8 10 12
	assert.Eq(f.Err(), nil, "mul operation")
	assert.NotEq(op, nil, "operation not nil")

	desired := []float32{8, 10, 12}
	res := op.Process()
	assert.Eq(desired, res, "vector result should match")
}

func TestHandler(t *testing.T) {
	f, op := prepareComplex()
	f.Handle(func(name string, payLoad map[string]flow.Data) {
		t.Log("Something happened:", name, payLoad)
	})
	op.Process()
}

func TestLocalRegistry(t *testing.T) {
	assert := assert.A(t)

	r := registry.New()
	e := r.Register("test", func() string { return "" })
	assert.NotEq(e, nil, "registered in a local register")

	f := flow.New()
	f.SetRegistry(r)
	op := f.Op("test")
	assert.NotEq(op, nil, "operation should be valid")

	op = f.Op("none")
	assert.NotEq(f.Err(), nil, "flow should contain an error")
}

func BenchmarkComplex(b *testing.B) {
	f, op := prepareComplex()

	b.Run("OP", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			op.Process()
		}
	})
	b.Run("RUN", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			f.Run(op)
		}
	})
}

func init() {
	flow.Register("vecmul", VecMul)
	flow.Register("vecadd", VecAdd)
	flow.Register("vecdiv", VecDiv)
	flow.Register("inc", Inc)
	flow.Register("add", Add)
}

func prepareComplex() (*flow.Flow, flow.Operation) {
	vecsize := 5
	v1 := make([]float32, vecsize)
	v2 := make([]float32, vecsize)
	for i := range v1 {
		v1[i], v2[i] = float32(i+1), 2
	}

	f := flow.New()
	f1 := f.Var("f1", v1)
	f2 := f.Var("f2", v2)

	mul := f.Op("vecmul", f1, f2)      // Doubles 2,4,6,8...
	add := f.Op("vecadd", mul, f2)     // Sum 4,8,10,12...
	mul2 := f.Op("vecmul", mul, add)   // mul again
	mul3 := f.Op("vecmul", mul2, f1)   // mul with f1
	div1 := f.Op("vecdiv", mul3, mul2) // div

	return f, div1
}

func VecMul(a, b []float32) []float32 {

	sz := Min(len(a), len(b))

	out := make([]float32, sz)
	vecasm.VecMulf32x8(a, b, out)
	return out
}
func VecAdd(a, b []float32) []float32 {
	sz := Min(len(a), len(b))
	out := make([]float32, sz)
	for i := 0; i < sz; i++ {
		out[i] = a[i] + b[i]
	}
	return out
}
func VecDiv(a, b []float32) []float32 {
	sz := Min(len(a), len(b))
	out := make([]float32, sz)
	for i := 0; i < sz; i++ {
		out[i] = a[i] / b[i]
	}
	return out
}

// ScalarInt
// Every time this operator is called we increase 1
func Inc() func() int {
	i := 0
	return func() int {
		i++
		return i
	}
}
func Add(a, b int) int {
	return a + b
}

// Utils
func Min(p ...int) int {
	min := p[0]
	for _, v := range p[1:] {
		if min < v {
			min = v
		}
	}
	return min
}
