// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package core

import (
	"math"
	"testing"
)

// Benchmarks for the walks whose per-element body changed: the binary
// arithmetic maps across every dtype, the mask comparisons and gathers,
// the scan and product kernels with short lines, the Lp norm, the
// broadcast fill and the magnitude map. Inputs are fixed literals, so
// every run feeds the same bytes.

// opFloats builds an n-element float array whose i-th element is
// deterministic and nonzero.
func opFloats(b *testing.B, n int) *Array {
	b.Helper()
	a, _ := FromFloats(make([]float64, n), n)
	for i := range a.RawFloats() {
		a.RawFloats()[i] = float64(i%97) - 48
	}
	return a
}

// opFloat32s is opFloats in single precision, built through the
// float32 payload so the values are the same ones the dtype round-trips.
func opFloat32s(b *testing.B, n int) *Array {
	b.Helper()
	f64 := opFloats(b, n)
	a, _ := Astype(f64, Float32)
	return a
}

// opInts is opFloats over an int payload; no element is zero, so the
// division benchmarks never take the error path.
func opInts(b *testing.B, n int) *Array {
	b.Helper()
	a, _ := FromInts(make([]int64, n), n)
	for i := range a.RawInts() {
		a.RawInts()[i] = int64(i%97) + 1
	}
	return a
}

// --- binary arithmetic maps ---

// BenchmarkSub1M and BenchmarkMul1M measure the two relations the
// add-size sweep in math_bench_extra_test.go does not cover.
func BenchmarkSub1M(b *testing.B) {
	a, c := opFloats(b, 1<<20), opFloats(b, 1<<20)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Sub(a, c); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkMul1M(b *testing.B) {
	a, c := opFloats(b, 1<<20), opFloats(b, 1<<20)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Mul(a, c); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkPow1M walks the float branch of Pow, whose body holds the
// math.Pow call instead of a func value.
func BenchmarkPow1M(b *testing.B) {
	a, c := opFloats(b, 1<<20), opFloats(b, 1<<20)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Pow(a, c); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkAddInt1M pins the int payload walk: the int64 sum wraps, so
// the loop body is the add itself.
func BenchmarkAddInt1M(b *testing.B) {
	a, c := opInts(b, 1<<20), opInts(b, 1<<20)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Add(a, c); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkAddFloat32_1M pins the narrow float walk, whose widening is
// exact and whose result narrows once.
func BenchmarkAddFloat32_1M(b *testing.B) {
	a, c := opFloat32s(b, 1<<20), opFloat32s(b, 1<<20)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Add(a, c); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkAddFloat16_1M pins the half walk: both payloads widen to
// float64 and the sum narrows once per element.
func BenchmarkAddFloat16_1M(b *testing.B) {
	f64a, f64b := opFloats(b, 1<<20), opFloats(b, 1<<20)
	a, err := Astype(f64a, Float16)
	if err != nil {
		b.Fatal(err)
	}
	c, err := Astype(f64b, Float16)
	if err != nil {
		b.Fatal(err)
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Add(a, c); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkMulComplex100k measures the widest payload: a complex product
// is four multiplies and two adds per element over 16 bytes each side.
func BenchmarkMulComplex100k(b *testing.B) {
	ac, _ := FromComplexes(make([]complex128, 100_000), 100_000)
	bc, _ := FromComplexes(make([]complex128, 100_000), 100_000)
	for i := range ac.RawComplexes() {
		v := complex(float64(i%97)-48, float64(i%31)-15)
		ac.RawComplexes()[i] = v
		bc.RawComplexes()[i] = complex(float64(i%53)-26, float64(i%17)-8)
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Mul(ac, bc); err != nil {
			b.Fatal(err)
		}
	}
}

// --- mask comparisons, selection and gather ---

// BenchmarkLtFloat64_1M measures the array-with-array comparison at the
// size where the walk has to fan out.
func BenchmarkLtFloat64_1M(b *testing.B) {
	a, c := opFloats(b, 1<<20), opFloats(b, 1<<20)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Lt(a, c); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkGeIFloat64_1M measures the scalar comparison, which compares
// exactly against the int and hoists the mirrored relation.
func BenchmarkGeIFloat64_1M(b *testing.B) {
	a := opFloats(b, 1<<20)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := GeI(a, 3); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkEqComplex1M pins the complex equality walk.
func BenchmarkEqComplex1M(b *testing.B) {
	ac, _ := FromComplexes(make([]complex128, 1<<20), 1<<20)
	bc, _ := FromComplexes(make([]complex128, 1<<20), 1<<20)
	for i := range ac.RawComplexes() {
		ac.RawComplexes()[i] = complex(float64(i%97)-48, float64(i%31)-15)
		bc.RawComplexes()[i] = complex(float64(i%53)-26, float64(i%17)-8)
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Eq(ac, bc); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkWhereFloat64_1M measures the dense int-conditioned select
// over float payloads: one condition read and one gathered value per
// element, the payload-slots fast path.
func BenchmarkWhereFloat64_1M(b *testing.B) {
	n := 1 << 20
	cond, _ := FromInts(make([]int64, n), n)
	for i := range cond.RawInts() {
		cond.RawInts()[i] = int64(i % 2)
	}
	x, y := opFloats(b, n), opFloats(b, n)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Where(cond, x, y); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkWhereMixed100k measures the accessor fallback: the bool mask
// a comparison answers, an int operand against a float one promoting to
// float, so the walk reads through the accessors.
func BenchmarkWhereMixed100k(b *testing.B) {
	n := 100_000
	cond, err := GeI(opFloats(b, n), 0)
	if err != nil {
		b.Fatal(err)
	}
	x := opInts(b, n)
	y := opFloats(b, n)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Where(cond, x, y); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSelectHalf1M gathers every other element: the count and the
// compaction carry half a million hits each.
func BenchmarkSelectHalf1M(b *testing.B) {
	n := 1 << 20
	a := opFloats(b, n)
	m, _ := FromInts(make([]int64, n), n)
	for i := range m.RawInts() {
		m.RawInts()[i] = int64(i % 2)
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Select(a, m); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSelectSparse1M gathers one element in sixty-four: the count
// pass dominates, since the compaction writes almost nothing.
func BenchmarkSelectSparse1M(b *testing.B) {
	n := 1 << 20
	a := opFloats(b, n)
	m, _ := FromInts(make([]int64, n), n)
	for i := range m.RawInts() {
		if i%64 == 0 {
			m.RawInts()[i] = 1
		}
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Select(a, m); err != nil {
			b.Fatal(err)
		}
	}
}

// --- scans and products with short lines ---

// BenchmarkCumSumShortLines scans a 4096x2 array along the two-element
// trailing dimension: there are many lines but each carries almost no
// work, so the fan-out must not be granted per line.
func BenchmarkCumSumShortLines(b *testing.B) {
	a, _ := FromFloats(make([]float64, 4096*2), 4096, 2)
	for i := range a.RawFloats() {
		a.RawFloats()[i] = float64(i%97) - 48
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := CumSum(a, 1); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkCumProdShortLines is the multiply twin at the same shape.
func BenchmarkCumProdShortLines(b *testing.B) {
	a, _ := FromFloats(make([]float64, 4096*2), 4096, 2)
	for i := range a.RawFloats() {
		a.RawFloats()[i] = float64(i%7) + 1
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := CumProd(a, 1); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkProdShortLines reduces a 4096x8 array along the trailing
// dimension, the product twin of the same short-line shape.
func BenchmarkProdShortLines(b *testing.B) {
	a, _ := FromFloats(make([]float64, 4096*8), 4096, 8)
	for i := range a.RawFloats() {
		a.RawFloats()[i] = float64(i%7) + 1
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Prod(a, 1, false); err != nil {
			b.Fatal(err)
		}
	}
}

// --- the Lp norm ---

// BenchmarkNormL1_512 walks the branch whose sum multiplies instead of
// calling Pow.
func BenchmarkNormL1_512(b *testing.B) {
	a, _ := FromFloats(make([]float64, 512*512), 512, 512)
	for i := range a.RawFloats() {
		a.RawFloats()[i] = float64(i%97) - 48
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Norm(a, 1, 1, false); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkNormP3_512 walks the general exponent, which calls math.Pow
// and math.Abs per element and therefore carries the low fan-out floor.
func BenchmarkNormP3_512(b *testing.B) {
	a, _ := FromFloats(make([]float64, 512*512), 512, 512)
	for i := range a.RawFloats() {
		a.RawFloats()[i] = float64(i%97) - 48
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Norm(a, 3, 1, false); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkNormInfShortLines is the infinity norm over 64 lines of 4096
// elements: few lines, each carrying a long max-abs fold.
func BenchmarkNormInfShortLines(b *testing.B) {
	a, _ := FromFloats(make([]float64, 64*4096), 64, 4096)
	for i := range a.RawFloats() {
		a.RawFloats()[i] = float64(i%97) - 48
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Norm(a, math.Inf(1), 1, false); err != nil {
			b.Fatal(err)
		}
	}
}

// --- broadcast ---

// BenchmarkBroadcastToTiny expands a 16x16 target from a 16x1 source:
// 256 elements, far below any spawn floor.
func BenchmarkBroadcastToTiny(b *testing.B) {
	src, _ := FromFloats(make([]float64, 16), 16, 1)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := BroadcastTo(src, 16, 16); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkBroadcastToRun1M expands 256x1x1 to 256x64x64: one constant
// run of 4096 elements per outer step.
func BenchmarkBroadcastToRun1M(b *testing.B) {
	src, _ := FromFloats(make([]float64, 256), 256, 1, 1)
	for i := range src.RawFloats() {
		src.RawFloats()[i] = float64(i%97) - 48
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := BroadcastTo(src, 256, 64, 64); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkBroadcastToRunCached expands 4x1 to 4x1024: the target is a
// 32k working set, so the run fill dominates the walk instead of the
// allocation and the first touch of a fresh target.
func BenchmarkBroadcastToRunCached(b *testing.B) {
	src, _ := FromFloats(make([]float64, 4), 4, 1)
	for i := range src.RawFloats() {
		src.RawFloats()[i] = float64(i) + 0.5
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := BroadcastTo(src, 4, 1024); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkBroadcastToSerial and BenchmarkBroadcastToParallel fix the
// crossover the spawn floor sits on: the same 1x256 to 256x256
// expansion, walked by one worker and by every worker.
func BenchmarkBroadcastToSerial(b *testing.B) {
	SetNumCPU(1)
	defer SetNumCPU(0)
	opBroadcastNC(b)
}

func BenchmarkBroadcastToParallel(b *testing.B) { opBroadcastNC(b) }

func opBroadcastNC(b *testing.B) {
	src, _ := FromFloats(make([]float64, 256), 1, 256)
	for i := range src.RawFloats() {
		src.RawFloats()[i] = float64(i%97) - 48
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := BroadcastTo(src, 256, 256); err != nil {
			b.Fatal(err)
		}
	}
}

// --- magnitude, mean and integer quotient ---

// BenchmarkAbs1k pins the small end of the magnitude map: the walk is
// cheaper than the spawn it used to pay for.
func BenchmarkAbs1k(b *testing.B) {
	a := opFloats(b, 1_000)
	b.ReportAllocs()
	for b.Loop() {
		Abs(a)
	}
}

// BenchmarkAbs1M pins the large end, where the fan-out pays.
func BenchmarkAbs1M(b *testing.B) {
	a := opFloats(b, 1<<20)
	b.ReportAllocs()
	for b.Loop() {
		Abs(a)
	}
}

// BenchmarkAbsComplex100k walks the magnitude of a complex payload: a
// math.Hypot per element, so the low fan-out floor.
func BenchmarkAbsComplex100k(b *testing.B) {
	ac, _ := FromComplexes(make([]complex128, 100_000), 100_000)
	for i := range ac.RawComplexes() {
		ac.RawComplexes()[i] = complex(float64(i%97)-48, float64(i%31)-15)
	}
	b.ReportAllocs()
	for b.Loop() {
		Abs(ac)
	}
}

// BenchmarkMean1M measures the overflow guard and the fold behind it.
func BenchmarkMean1M(b *testing.B) {
	a := opFloats(b, 1<<20)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Mean(a); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkQuo1M measures the integer quotient, whose zero check now
// rides the division walk instead of scanning the divisors first.
func BenchmarkQuo1M(b *testing.B) {
	a, c := opInts(b, 1<<20), opInts(b, 1<<20)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Quo(a, c); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSum1M is the global sum over a payload large enough that the
// fold is worth partitioning.
func BenchmarkSum1M(b *testing.B) {
	a := benchFloatsShape(1 << 20)
	b.ReportAllocs()
	for b.Loop() {
		_ = Sum(a)
	}
}

// BenchmarkSum1MInt is the integer fold, which is exact under any
// partition and therefore needs no interleaved partials.
func BenchmarkSum1MInt(b *testing.B) {
	a := benchIntsShape(1<<20, 1<<20)
	b.ReportAllocs()
	for b.Loop() {
		_ = Sum(a)
	}
}