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

package core

import (
	"math"
	"math/big"
	mrand "math/rand/v2"
	"testing"
)

// The central-sum evidence file: the accessor-chain Covariance,
// Correlation and Integrate walks against their canonical-block +
// treeSum candidates, both measured against exact big.Float referents
// (512 bits) and timed in one binary.

// covarianceLegacy is the walk the entry point keeps: two accessor
// reads per element folded into one chain.
func covarianceLegacy(a, b []float64, ma, mb float64) float64 {
	var cov float64
	for i := range a {
		cov += (a[i] - ma) * (b[i] - mb)
	}
	return cov
}

// covarianceBlocks is the candidate: the same per-element arithmetic
// cut into the canonical blocks, one chain partial per block, the
// partials combined through the balanced tree.
func covarianceBlocks(a, b []float64, ma, mb float64) float64 {
	n := len(a)
	parts := foldParts(n)
	if parts == 1 {
		return covarianceLegacy(a, b, ma, mb)
	}
	partials := make([]float64, parts)
	for c := range parts {
		lo, hi := c*n/parts, (c+1)*n/parts
		var acc float64
		for i := lo; i < hi; i++ {
			acc += (a[i] - ma) * (b[i] - mb)
		}
		partials[c] = acc
	}
	return treeSum(partials)
}

// correlationLegacy is the accessor-chain correlation walk.
func correlationLegacy(a, b []float64, ma, mb float64) (num, da2, db2 float64) {
	for i := range a {
		da := a[i] - ma
		dbv := b[i] - mb
		num += da * dbv
		da2 += da * da
		db2 += dbv * dbv
	}
	return num, da2, db2
}

// correlationBlocks is the candidate: canonical blocks and a tree over
// the partials of each of the three sums.
func correlationBlocks(a, b []float64, ma, mb float64) (num, da2, db2 float64) {
	n := len(a)
	parts := foldParts(n)
	if parts == 1 {
		return correlationLegacy(a, b, ma, mb)
	}
	pn, px, py := make([]float64, parts), make([]float64, parts), make([]float64, parts)
	for c := range parts {
		lo, hi := c*n/parts, (c+1)*n/parts
		var sn, sx, sy float64
		for i := lo; i < hi; i++ {
			da := a[i] - ma
			dbv := b[i] - mb
			sn += da * dbv
			sx += da * da
			sy += dbv * dbv
		}
		pn[c], px[c], py[c] = sn, sx, sy
	}
	return treeSum(pn), treeSum(px), treeSum(py)
}

// integrateLegacy is the plain trapezoid chain.
func integrateLegacy(y []float64) float64 {
	var total float64
	for i := 1; i < len(y); i++ {
		total += (y[i-1] + y[i]) / 2
	}
	return total
}

// integrateBlocks is the candidate: the trapezoid areas cut into the
// canonical blocks over the area indices (n−1 of them), the block
// chains combined through the tree.
func integrateBlocks(y []float64) float64 {
	m := len(y) - 1
	parts := foldParts(m)
	if parts == 1 {
		return integrateLegacy(y)
	}
	partials := make([]float64, parts)
	for c := range parts {
		lo, hi := c*m/parts, (c+1)*m/parts
		var acc float64
		for i := lo; i < hi; i++ {
			acc += (y[i] + y[i+1]) / 2
		}
		partials[c] = acc
	}
	return treeSum(partials)
}

// spreadPair builds a deterministic correlated pair whose magnitudes
// span thirty orders, the shape that cancels a central sum fastest.
func spreadPair(n int, seed uint64) (x, y []float64) {
	rng := mrand.New(mrand.NewPCG(seed, seed))
	x = make([]float64, n)
	y = make([]float64, n)
	for i := range x {
		mag := math.Pow(10, -15+30*rng.Float64())
		if rng.Float64() < 0.5 {
			mag = -mag
		}
		x[i] = mag
		y[i] = 3*mag + math.Pow(10, -15+30*rng.Float64())
	}
	return x, y
}

// exactMeans returns the exact arithmetic means at the given precision.
func exactMeans(x, y []float64, prec uint) (*big.Float, *big.Float) {
	sx := new(big.Float).SetPrec(prec)
	sy := new(big.Float).SetPrec(prec)
	one := new(big.Float).SetPrec(prec)
	for i := range x {
		sx.Add(sx, new(big.Float).SetPrec(prec).SetFloat64(x[i]))
		sy.Add(sy, new(big.Float).SetPrec(prec).SetFloat64(y[i]))
	}
	one.SetInt(new(big.Int).SetInt64(int64(len(x))))
	return new(big.Float).SetPrec(prec).Quo(sx, one), new(big.Float).SetPrec(prec).Quo(sy, one)
}

// TestCentralSumsAccuracy measures the legacy and candidate walks of
// the three central-sum kernels against exact referents.
func TestCentralSumsAccuracy(t *testing.T) {
	const prec = 512
	n := 1 << 20
	x, y := spreadPair(n, 0xBEEF)
	exma, exmb := exactMeans(x, y, prec)
	maf, _ := exma.Float64()
	mbf, _ := exmb.Float64()

	// Covariance referent: Σ(x−x̄)(y−ȳ)/(n−1) at full precision.
	ref := new(big.Float).SetPrec(prec)
	for i := range x {
		dx := new(big.Float).SetPrec(prec).Sub(new(big.Float).SetPrec(prec).SetFloat64(x[i]), exma)
		dy := new(big.Float).SetPrec(prec).Sub(new(big.Float).SetPrec(prec).SetFloat64(y[i]), exmb)
		ref.Add(ref, new(big.Float).SetPrec(prec).Mul(dx, dy))
	}
	ref.Quo(ref, new(big.Float).SetPrec(prec).SetInt(new(big.Int).SetInt64(int64(n-1))))
	refF, _ := ref.Float64()

	oldCov := covarianceLegacy(x, y, maf, mbf) / float64(n-1)
	newCov := covarianceBlocks(x, y, maf, mbf) / float64(n-1)
	scale := math.Max(math.Abs(refF), 1)
	t.Logf("covariance n=2^20: exact=%.17g legacy=%.17g (%.3e) blocks=%.17g (%.3e)",
		refF, oldCov, math.Abs(oldCov-refF)/scale, newCov, math.Abs(newCov-refF)/scale)

	// Correlation referent.
	var rnum, rda2, rdb2 = new(big.Float).SetPrec(prec), new(big.Float).SetPrec(prec), new(big.Float).SetPrec(prec)
	for i := range x {
		dx := new(big.Float).SetPrec(prec).Sub(new(big.Float).SetPrec(prec).SetFloat64(x[i]), exma)
		dy := new(big.Float).SetPrec(prec).Sub(new(big.Float).SetPrec(prec).SetFloat64(y[i]), exmb)
		rnum.Add(rnum, new(big.Float).SetPrec(prec).Mul(dx, dy))
		rda2.Add(rda2, new(big.Float).SetPrec(prec).Mul(dx, dx))
		rdb2.Add(rdb2, new(big.Float).SetPrec(prec).Mul(dy, dy))
	}
	root := new(big.Float).SetPrec(prec).Mul(rda2, rdb2)
	root.Sqrt(root)
	refR, _ := new(big.Float).SetPrec(prec).Quo(rnum, root).Float64()

	on, oda2, odb2 := correlationLegacy(x, y, maf, mbf)
	nn, nda2, ndb2 := correlationBlocks(x, y, maf, mbf)
	oldR := on / math.Sqrt(oda2*odb2)
	newR := nn / math.Sqrt(nda2*ndb2)
	t.Logf("correlation n=2^20: exact=%.17g legacy=%.17g (%.3e) blocks=%.17g (%.3e)",
		refR, oldR, math.Abs(oldR-refR), newR, math.Abs(newR-refR))

	// Integrate referent: the exact trapezoid sum of the same sample.
	rngI := mrand.New(mrand.NewPCG(0xDADA, 99))
	yv := make([]float64, n)
	for i := range yv {
		yv[i] = math.Pow(10, -15+30*rngI.Float64())
	}
	tri := new(big.Float).SetPrec(prec)
	for i := 1; i < n; i++ {
		a := new(big.Float).SetPrec(prec).SetFloat64(yv[i-1])
		b := new(big.Float).SetPrec(prec).SetFloat64(yv[i])
		tri.Add(tri, new(big.Float).SetPrec(prec).Quo(new(big.Float).SetPrec(prec).Add(a, b), new(big.Float).SetPrec(prec).SetFloat64(2)))
	}
	refI, _ := tri.Float64()
	oldI := integrateLegacy(yv)
	newI := integrateBlocks(yv)
	scaleI := math.Max(math.Abs(refI), 1)
	t.Logf("integrate n=2^20: exact=%.17g legacy=%.17g (%.3e) blocks=%.17g (%.3e)",
		refI, oldI, math.Abs(oldI-refI)/scaleI, newI, math.Abs(newI-refI)/scaleI)
}

// TestCentralSumsProductionPins pins the shipped entry points on the
// spread sample against the exact referent, at tolerances the block
// walk clears by two orders and the plain chain fails: the covariance
// chain erred 1.111e-11 relative here, the block walk 1.642e-14; the
// correlation chain 1.858e-11, the block walk 2.032e-14; the integrate
// chain 5.341e-13, the block walk 3.046e-14.
func TestCentralSumsProductionPins(t *testing.T) {
	const prec = 512
	n := 1 << 20
	x, y := spreadPair(n, 0xBEEF)
	xa, _ := FromFloats(x, n)
	ya, _ := FromFloats(y, n)
	exma, exmb := exactMeans(x, y, prec)

	ref := new(big.Float).SetPrec(prec)
	for i := range x {
		dx := new(big.Float).SetPrec(prec).Sub(new(big.Float).SetPrec(prec).SetFloat64(x[i]), exma)
		dy := new(big.Float).SetPrec(prec).Sub(new(big.Float).SetPrec(prec).SetFloat64(y[i]), exmb)
		ref.Add(ref, new(big.Float).SetPrec(prec).Mul(dx, dy))
	}
	ref.Quo(ref, new(big.Float).SetPrec(prec).SetInt(new(big.Int).SetInt64(int64(n-1))))
	refF, _ := ref.Float64()
	cov, err := Covariance(xa, ya)
	if err != nil {
		t.Fatal(err)
	}
	if rel := math.Abs(cov-refF) / math.Abs(refF); rel > 1e-12 {
		t.Errorf("covariance relative error %g exceeds 1e-12 (got %v, want %v)", rel, cov, refF)
	}

	rnum := new(big.Float).SetPrec(prec)
	rda2 := new(big.Float).SetPrec(prec)
	rdb2 := new(big.Float).SetPrec(prec)
	for i := range x {
		dx := new(big.Float).SetPrec(prec).Sub(new(big.Float).SetPrec(prec).SetFloat64(x[i]), exma)
		dy := new(big.Float).SetPrec(prec).Sub(new(big.Float).SetPrec(prec).SetFloat64(y[i]), exmb)
		rnum.Add(rnum, new(big.Float).SetPrec(prec).Mul(dx, dy))
		rda2.Add(rda2, new(big.Float).SetPrec(prec).Mul(dx, dx))
		rdb2.Add(rdb2, new(big.Float).SetPrec(prec).Mul(dy, dy))
	}
	root := new(big.Float).SetPrec(prec).Mul(rda2, rdb2)
	root.Sqrt(root)
	refR, _ := new(big.Float).SetPrec(prec).Quo(rnum, root).Float64()
	corr, err := Correlation(xa, ya)
	if err != nil {
		t.Fatal(err)
	}
	if math.Abs(corr-refR) > 1e-12 {
		t.Errorf("correlation error %g exceeds 1e-12 (got %v, want %v)", math.Abs(corr-refR), corr, refR)
	}

	yv := make([]float64, n)
	rngI := mrand.New(mrand.NewPCG(0xDADA, 99))
	for i := range yv {
		yv[i] = math.Pow(10, -15+30*rngI.Float64())
	}
	tri := new(big.Float).SetPrec(prec)
	for i := 1; i < n; i++ {
		a := new(big.Float).SetPrec(prec).SetFloat64(yv[i-1])
		b := new(big.Float).SetPrec(prec).SetFloat64(yv[i])
		tri.Add(tri, new(big.Float).SetPrec(prec).Quo(new(big.Float).SetPrec(prec).Add(a, b), new(big.Float).SetPrec(prec).SetFloat64(2)))
	}
	refI, _ := tri.Float64()
	ya2, _ := FromFloats(yv, n)
	gotI, err := Integrate(ya2, 1)
	if err != nil {
		t.Fatal(err)
	}
	if rel := math.Abs(gotI-refI) / math.Abs(refI); rel > 1e-13 {
		t.Errorf("integrate relative error %g exceeds 1e-13 (got %v, want %v)", rel, gotI, refI)
	}
}

// TestCentralSumsShortInputsPinEqualBits pins the contract that the
// canonical partition leaves short inputs' bits alone: at and below one
// fold block the block walk is the chain walk, sample for sample.
func TestCentralSumsShortInputsPinEqualBits(t *testing.T) {
	rng := mrand.New(mrand.NewPCG(5, 5))
	n := 3000
	x := make([]float64, n)
	y := make([]float64, n)
	for i := range x {
		x[i] = rng.NormFloat64()
		y[i] = rng.NormFloat64()
	}
	ma, mb := 0.37, -1.2
	if covarianceLegacy(x, y, ma, mb) != covarianceBlocks(x, y, ma, mb) {
		t.Error("covariance: short-input bits moved")
	}
	on, oda, odb := correlationLegacy(x, y, ma, mb)
	nn, nda, ndb := correlationBlocks(x, y, ma, mb)
	if on != nn || oda != nda || odb != ndb {
		t.Error("correlation: short-input bits moved")
	}
	yy := make([]float64, n+1)
	for i := range yy {
		yy[i] = rng.NormFloat64()
	}
	if integrateLegacy(yy) != integrateBlocks(yy) {
		t.Error("integrate: short-input bits moved")
	}
}

// BenchmarkCentralSumsTimes the two walks of each kernel over one
// pair of 2^20 samples.
func BenchmarkCentralSums(b *testing.B) {
	n := 1 << 20
	x, y := spreadPair(n, 0xBEEF)
	ma, mb := 0.37, -1.2
	b.Run("covariance/legacy", func(b *testing.B) {
		for b.Loop() {
			covarianceLegacy(x, y, ma, mb)
		}
		b.SetBytes(int64(n) * 8)
	})
	b.Run("covariance/blocks", func(b *testing.B) {
		for b.Loop() {
			covarianceBlocks(x, y, ma, mb)
		}
		b.SetBytes(int64(n) * 8)
	})
	b.Run("correlation/legacy", func(b *testing.B) {
		for b.Loop() {
			correlationLegacy(x, y, ma, mb)
		}
		b.SetBytes(int64(n) * 8)
	})
	b.Run("correlation/blocks", func(b *testing.B) {
		for b.Loop() {
			correlationBlocks(x, y, ma, mb)
		}
		b.SetBytes(int64(n) * 8)
	})
	b.Run("integrate/legacy", func(b *testing.B) {
		yy := make([]float64, n)
		for i := range yy {
			yy[i] = float64(i % 977)
		}
		for b.Loop() {
			integrateLegacy(yy)
		}
		b.SetBytes(int64(n) * 8)
	})
	b.Run("integrate/blocks", func(b *testing.B) {
		yy := make([]float64, n)
		for i := range yy {
			yy[i] = float64(i % 977)
		}
		for b.Loop() {
			integrateBlocks(yy)
		}
		b.SetBytes(int64(n) * 8)
	})
}