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

package stats

import (
	"fmt"
	"math"
	"slices"
	"testing"

	"sourcedock.dev/petrbalvin/tensor/internal/core"
)

// Benchmarks for the estimation kernels: the mixture expectation
// maximisation sweep, the kernel-density sweep, the windowed extrema
// and the histogram passes. Every input is a fixed arithmetic
// progression, so two runs of a benchmark measure the same work.

// benchKernelsSeries builds a deterministic vector of length n: a few
// incommensurate frequencies plus a modular jitter, so the sorts and the
// window folds meet a spread without a generator. The phase separates
// two series of the same length.
func benchKernelsSeries(n int, phase float64) []float64 {
	vals := make([]float64, n)
	for i := range vals {
		x := float64(i)
		vals[i] = math.Sin(0.0017*x+phase)*4 + math.Cos(0.071*x+phase)*2 + float64((i*37)%101)/101
	}
	return vals
}

// benchKernelsVector wraps the series as a rank-1 array.
func benchKernelsVector(b *testing.B, n int, phase float64) *core.Array {
	b.Helper()
	a, err := core.FromFloats(benchKernelsSeries(n, phase), n)
	if err != nil {
		b.Fatal(err)
	}
	return a
}

// benchKernelsCloud builds an n-by-d sample of k well-separated clusters
// whose jitter follows a fixed congruence: the same rows on every run,
// and a cloud the mixture's covariances can factor.
func benchKernelsCloud(n, d, k int) []float64 {
	vals := make([]float64, n*d)
	for i := range vals {
		row, col := i/d, i%d
		jitter := float64((row*2654435761+col*40503)%1000)/1000 - 0.5
		vals[i] = float64(row%k)*4 + jitter*1.5
	}
	return vals
}

// BenchmarkKernelsGaussianMixture measures one full fit: the k-means++
// seeding, the Lloyd sweeps and the expectation maximisation sweeps.
func BenchmarkKernelsGaussianMixture(b *testing.B) {
	const n, d, k = 3000, 5, 5
	data := benchKernelsCloud(n, d, k)
	x, err := core.FromFloats(data, n, d)
	if err != nil {
		b.Fatal(err)
	}
	g := core.NewGenerator(7)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := GaussianMixture(g, x, k); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkKernelsGMMSweep measures the expectation maximisation sweeps
// alone: the seeding runs once outside the timed loop, and every
// iteration re-copies the seeded parameters, which gmmEM updates in
// place.
func BenchmarkKernelsGMMSweep(b *testing.B) {
	const n, d, k = 3000, 5, 5
	data := benchKernelsCloud(n, d, k)
	weights, means, covs, err := gmmSeed(core.NewGenerator(7), data, n, d, k)
	if err != nil {
		b.Fatal(err)
	}
	b.ReportAllocs()
	for b.Loop() {
		w := slices.Clone(weights)
		m := make([][]float64, k)
		c := make([][]float64, k)
		for j := range k {
			m[j] = slices.Clone(means[j])
			c[j] = slices.Clone(covs[j])
		}
		res, err := gmmEM(data, n, d, k, w, m, c)
		if err != nil {
			b.Fatal(err)
		}
		b.ReportMetric(float64(res.Iterations), "sweeps")
	}
}

// BenchmarkKernelsKernelDensity measures the O(n·points) kernel-density
// sweep with a fixed bandwidth, so no Silverman pre-pass is timed.
func BenchmarkKernelsKernelDensity(b *testing.B) {
	sample := benchKernelsVector(b, 2048, 0)
	points := benchKernelsVector(b, 512, 0.5)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := KernelDensity(sample, 0.35, points); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkKernelsRollingMaxWide measures the windowed maximum at half
// the series length, the shape that separates a rescan from a deque.
func BenchmarkKernelsRollingMaxWide(b *testing.B) {
	a := benchKernelsVector(b, 16384, 0)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := RollingMax(a, 8192); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkKernelsRollingMinWide is the minimum's counterpart.
func BenchmarkKernelsRollingMinWide(b *testing.B) {
	a := benchKernelsVector(b, 16384, 0)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := RollingMin(a, 8192); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkKernelsHistogram measures the fused scan and the counted
// bins over a long sample.
func BenchmarkKernelsHistogram(b *testing.B) {
	a := benchKernelsVector(b, 262144, 0)
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := Histogram(a, 256); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkKernelsHistogram2D measures the paired binning over a long
// sample on both axes.
func BenchmarkKernelsHistogram2D(b *testing.B) {
	const n = 65536
	x := benchKernelsVector(b, n, 0)
	y := benchKernelsVector(b, n, 0.5)
	b.ReportAllocs()
	for b.Loop() {
		if _, _, _, err := Histogram2D(x, y, 64, 64); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkKernelsGaussianProcessFit measures one fit over a design the
// size a Gaussian-process regression usually runs on: the Gram matrix,
// its factorisation, the weights and the posterior at the test points.
func BenchmarkKernelsGaussianProcessFit(b *testing.B) {
	const n, m = 150, 60
	train := make([]float64, n)
	test := make([]float64, m)
	for i := range train {
		train[i] = float64(i) / float64(n-1)
	}
	for i := range test {
		test[i] = float64(i) / float64(m-1)
	}
	trainX, err := core.FromFloats(train, n, 1)
	if err != nil {
		b.Fatal(err)
	}
	trainY, err := core.FromFloats(benchKernelsSeries(n, 0), n)
	if err != nil {
		b.Fatal(err)
	}
	testX, err := core.FromFloats(test, m, 1)
	if err != nil {
		b.Fatal(err)
	}
	kernel, err := SquaredExponentialKernel(0.2)
	if err != nil {
		b.Fatal(err)
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := GaussianProcessRegression(kernel, trainX, trainY, 0.05, testX); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkKernelsRollingMaxScaling reports the monotonic deque against
// the series length at a window of half the series: the linear growth
// the deque replaced the window rescan with.
func BenchmarkKernelsRollingMaxScaling(b *testing.B) {
	for _, n := range []int{4096, 16384, 65536} {
		a := benchKernelsVector(b, n, 0)
		b.Run(fmt.Sprintf("n=%d/w=%d", n, n/2), func(b *testing.B) {
			b.ReportAllocs()
			for b.Loop() {
				if _, err := RollingMax(a, n/2); err != nil {
					b.Fatal(err)
				}
			}
		})
	}
}