// Copyright 2026 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.

//go:build goexperiment.simd

package test_helpers

import (
	"math/bits"
	"testing"
	"unsafe"
)

func ReduceSum[S number](s []S) S {
	var x S
	for _, v := range s {
		x += v
	}
	return x
}

func Neg[S number](x S) S {
	return -x
}

func Abs[S, S2 number](x S) S2 {
	if x == 0 {
		return 0 // this is zero not -0 if S is float
	}
	if x < 0 {
		return S2(-x)
	}
	return S2(x)
}

func Add[S number](x, y S) S {
	return x + y
}

func Sub[S number](x, y S) S {
	return x - y
}

func Mul[S number](x, y S) S {
	return x * y
}

func Div[S float](x, y S) S {
	return x / y
}

func And[S integer](x, y S) S {
	return x & y
}

func Or[S integer](x, y S) S {
	return x | y
}

func Xor[S integer](x, y S) S {
	return x ^ y
}

func AndNot[S integer](x, y S) S {
	return x &^ y
}

func Xnor[S integer](x, y S) S {
	return ^(x ^ y)
}

func Not[S integer](x S) S {
	return ^x
}

func Max[S number](x, y S) S {
	return max(x, y)
}

func Min[S number](x, y S) S {
	return min(x, y)
}

func ShiftLeft[S integer](x S, s uint64) S {
	return x << s
}

func ShiftRight[S integer](x S, s uint64) S {
	return x >> s
}

func Equal[S number](x, y S) bool {
	return x == y
}

func NotEqual[S number](x, y S) bool {
	return x != y
}

func Less[S number](x, y S) bool {
	return x < y
}

func LessEqual[S number](x, y S) bool {
	return x <= y
}

func Greater[S number](x, y S) bool {
	return x > y
}

func GreaterEqual[S number](x, y S) bool {
	return x >= y
}

func RotateLeft[S integer](x S, dist uint64) S {
	size := uint64(unsafe.Sizeof(x)) * 8
	mask := (uint64(1) << size) - 1
	dist = dist & (size - 1)
	if dist == 0 {
		return x
	}
	return (x << dist) | S((mask&uint64(x))>>(size-dist))
}

func RotateRight[S integer](x S, dist uint64) S {
	size := uint64(unsafe.Sizeof(x)) * 8
	mask := (uint64(1) << size) - 1
	dist = dist & (size - 1)
	if dist == 0 {
		return x
	}
	return S((mask&uint64(x))>>dist) | (x << (size - dist))
}

func OnesCount[T integer](x T) T {
	size := uint64(unsafe.Sizeof(x)) * 8
	return T(bits.OnesCount64(uint64(x) & ((1 << size) - 1)))
}

func TrailingZeros[T integer](x T) T {
	size := int(unsafe.Sizeof(x)) * 8
	tz := bits.TrailingZeros64(uint64(x))
	if tz > size {
		tz = size
	}
	return T(tz)
}

// Vec is any vector of any number (int, uint, or float)
type Vec[E number] interface {
	Store(s []E)
	Len() int
}

// Mask is any mask
type Mask interface {
	String() string
}

// IVec is a vector of signed integers, that will also have a ToMask() method.
type IVec[E signed, M Mask] interface {
	Store(s []E)
	Len() int
	ToMask() M
}

// The Test-XX-2-YY-[e] functions test various unary and binary vector operations.
// V = 1 vector, VV = 2 vectors, VR = vector and rotation.
// S = scalar.
// M = mask.
// A trailing "e" means that the emulation function is elementwise, or scalar.

// TestV2S tests the simd Vector-To-Scalar method f against the expected behavior generated by want
func TestV2S[S number, V Vec[S]](t *testing.T, slice []S, load func([]S) V, f func(_ V) S, want func([]S) S) {
	var v V
	n := v.Len()
	t.Helper()
	forSlice(t, slice, n, func(x []S) bool {
		t.Helper()
		a := load(x)
		g := f(a)
		w := want(x)
		return CheckScalarsLogInput(t, g, w, 0.000001, func() { t.Helper(); t.Logf("x=%v", x) })
	})
}

// TestV2V tests the simd Vector-To-Vector method f against the expected behavior generated by want
func TestV2V[S number, V Vec[S]](t *testing.T, slice []S, load func([]S) V, f func(_ V) V, want func([]S) []S) {
	var v V
	n := v.Len()
	t.Helper()
	forSlice(t, slice, n, func(x []S) bool {
		t.Helper()
		a := load(x)
		g := make([]S, n)
		f(a).Store(g)
		w := want(x)
		return CheckSlicesLogInput(t, g, w, 0.000001, func() { t.Helper(); t.Logf("x=%v", x) })
	})
}

// TestV2Ve tests the simd unary Vector-To-Vector method f against the expected behavior generated by per-Element function want
func TestV2Ve[S, S2 number, V Vec[S], V2 Vec[S2]](t *testing.T, slice []S, load func([]S) V, f func(_ V) V2, want func(S) S2) {
	var v V
	n := v.Len()
	t.Helper()
	forSlice(t, slice, n, func(x []S) bool {
		t.Helper()
		a := load(x)
		g := make([]S2, n)
		f(a).Store(g)
		w := ImpliedDo(n, func(i int) S2 { return want(x[i]) })
		return CheckSlicesLogInput(t, g, w, 0.000001, func() { t.Helper(); t.Logf("x=%v", x) })
	})
}

// TestVR2Ve tests the simd binary (Vector,Rotation)-To-Vector method f against the expected behavior generated by per-Element function want
func TestVR2Ve[S number, V Vec[S]](t *testing.T, slice []S, shifts []uint64, load func([]S) V, f func(V, uint64) V, want func(S, uint64) S) {
	var v V
	n := v.Len()
	t.Helper()
	for _, shift := range shifts {
		forSlice(t, slice, n, func(x []S) bool {
			t.Helper()
			a := load(x)
			g := make([]S, n)
			f(a, shift).Store(g)
			w := ImpliedDo(n, func(i int) S { return want(x[i], shift) })
			return CheckSlicesLogInput(t, g, w, 0.000001, func() { t.Helper(); t.Logf("x=%v, dist=%v", x, shift) })
		})
	}
}

// TestVV2V tests the simd binary (vector,vector)-to-vector method f against the expected behavior generated by want
func TestVV2V[S number, V Vec[S]](t *testing.T, slice []S, load func([]S) V, f func(_, _ V) V, want func(_, _ []S) []S) {
	var v V
	n := v.Len()
	t.Helper()
	forSlicePair(t, slice, n, func(x, y []S) bool {
		t.Helper()
		a := load(x)
		b := load(y)
		g := make([]S, n)
		f(a, b).Store(g)
		w := want(x, y)
		return CheckSlicesLogInput(t, g, w, 0.000001, func() { t.Helper(); t.Logf("x=%v, y=%v", x, y) })
	})
}

// TestVV2Ve tests the simd binary (Vector,Vector)-To-Vector method f against the expected behavior generated by per-Element function want
func TestVV2Ve[S number, V Vec[S]](t *testing.T, slice []S, load func([]S) V, f func(_, _ V) V, want func(_, _ S) S) {
	var v V
	n := v.Len()
	t.Helper()
	forSlicePair(t, slice, n, func(x, y []S) bool {
		t.Helper()
		a := load(x)
		b := load(y)
		g := make([]S, n)
		f(a, b).Store(g)
		w := ImpliedDo(n, func(i int) S { return want(x[i], y[i]) })
		return CheckSlicesLogInput(t, g, w, 0.000001, func() { t.Helper(); t.Logf("x=%v, y=%v", x, y) })
	})
}

// TestVV2Me tests the simd comparison method f against the expected behavior generated by per-Element function want
func TestVV2Me[S number, V Vec[S], M Mask](t *testing.T, slice []S, load func([]S) V, f func(_, _ V) M, want func(_, _ S) bool) {
	var v V
	n := v.Len()
	t.Helper()
	forSlicePair(t, slice, n, func(x, y []S) bool {
		t.Helper()
		a := load(x)
		b := load(y)
		g := f(a, b).String()
		w := ImpliedBools(n, func(i int) bool { return want(x[i], y[i]) })
		return CheckStringsLogInput(t, g, w, 0.000001, func() { t.Helper(); t.Logf("x=%v, y=%v", x, y) })
	})
}

// TestM2Me tests the simd mask unary method f against the expected behavior generated by per-Element function want
func TestM2Me[S signed, M Mask, V IVec[S, M]](t *testing.T, load func([]S) V, f func(_ M) M, want func(_ S) S) {
	var v V
	n := v.Len()
	t.Helper()
	b2s := func(b bool) S {
		if b {
			return -1
		}
		return 0
	}

	bload := func(s []bool) V {
		sl := make([]S, n)
		for i := range sl {
			sl[i] = b2s(s[i])
		}
		return load(sl)
	}
	forBoolSlice(t, Bools(), n, func(x []bool) bool {
		t.Helper()
		a := bload(x)
		g := f(a.ToMask()).String()
		w := ImpliedBools(n, func(i int) bool { return 0 != want(b2s(x[i])) })
		return CheckStringsLogInput(t, g, w, 0.000001, func() { t.Helper(); t.Logf("x=%v", x) })
	})
}

// TestMM2Me tests the simd mask binary method f against the expected behavior generated by per-Element function want
func TestMM2Me[S signed, M Mask, V IVec[S, M]](t *testing.T, load func([]S) V, f func(_, _ M) M, want func(_, _ S) S) {
	var v V
	n := v.Len()
	t.Helper()
	b2s := func(b bool) S {
		if b {
			return -1
		}
		return 0
	}

	bload := func(s []bool) V {
		sl := make([]S, n)
		for i := range sl {
			sl[i] = b2s(s[i])
		}
		return load(sl)
	}
	forBoolSlicePair(t, Bools(), n, func(x, y []bool) bool {
		t.Helper()
		a := bload(x)
		b := bload(y)
		g := f(a.ToMask(), b.ToMask()).String()
		w := ImpliedBools(n, func(i int) bool { return 0 != want(b2s(x[i]), b2s(y[i])) })
		return CheckStringsLogInput(t, g, w, 0.000001, func() { t.Helper(); t.Logf("x=%v, y=%v", x, y) })
	})
}