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

// Precision pins for the variance fold: the squared deviations sum
// through the canonical partition the core reductions keep, and the pin
// holds the fold against an exact big.Rat referent on the data the
// single chain loses on, beside the chain it replaced. The acceptance
// bar is the chain's own error: on every dataset the partition must sit
// no further from the referent than the chain it replaced, and on the
// descending-magnitude data, where every later add rounds against an
// accumulator already near the total, strictly closer.

package stats

import (
	"math"
	"math/big"
	"testing"

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

// variancePrecDatasets builds the samples the fold is held on, at
// n = 2^20: a wide-magnitude square wave on a large offset, and a
// geometric spread of deviations walked in descending and ascending
// order of their squares, the shapes the chain loses the most and the
// least on.
func variancePrecDatasets() ([][]float64, []string) {
	const n = 1 << 20
	descending := make([]float64, n)
	ascending := make([]float64, n)
	squareWave := make([]float64, n)
	for i := range squareWave {
		w := 1.0
		if (i/64)%2 == 1 {
			w = -1
		}
		squareWave[i] = 1e9 + w + 0.125*float64(i%7-3)
		devD := math.Pow(2, 10-20*float64(i)/n)
		devA := math.Pow(2, -10+20*float64(i)/n)
		descending[i] = 1e6 + devD
		ascending[i] = 1e6 + devA
	}
	return [][]float64{squareWave, descending, ascending},
		[]string{"wide-magnitude square wave", "descending squares", "ascending squares"}
}

// variancePrecExact sums the squared deviations from the exact mean
// exactly in big.Rat, the quantity both folds answer.
func variancePrecExact(vals []float64) *big.Rat {
	mean := new(big.Rat)
	for _, v := range vals {
		mean.Add(mean, new(big.Rat).SetFloat64(v))
	}
	mean.Quo(mean, new(big.Rat).SetInt64(int64(len(vals))))
	total := new(big.Rat)
	for _, v := range vals {
		d := new(big.Rat).SetFloat64(v)
		d.Sub(d, mean)
		total.Add(total, new(big.Rat).Mul(d, d))
	}
	return total
}

// variancePrecChainFold is the single chain the partition replaced: the
// accuracy bar the fold may not drop below.
func variancePrecChainFold(vals []float64, mean float64) float64 {
	sum := 0.0
	for _, v := range vals {
		d := v - mean
		sum += d * d
	}
	return sum
}

// variancePrecFoldRelErr reports the folded sum of squares' relative
// error against the exact referent.
func variancePrecFoldRelErr(t *testing.T, label string, folded float64, exact *big.Rat) *big.Rat {
	t.Helper()
	got := new(big.Rat).SetFloat64(folded)
	got.Sub(got, exact)
	got.Abs(got)
	if exact.Sign() != 0 {
		got.Quo(got, new(big.Rat).Abs(exact))
	}
	f, _ := got.Float64()
	t.Logf("%s: relative error %.3g", label, f)
	return got
}

// TestVarianceFoldPrecisionAgainstBigRat holds the canonical partition
// against the exact referent on the three datasets, at the same mean
// both folds read, so the comparison isolates the fold. The partition
// must never lose to the chain, and must beat it on the descending data.
func TestVarianceFoldPrecisionAgainstBigRat(t *testing.T) {
	datasets, names := variancePrecDatasets()
	for k, vals := range datasets {
		t.Run(names[k], func(t *testing.T) {
			exact := variancePrecExact(vals)
			mean := 0.0
			for _, v := range vals {
				mean += v
			}
			mean /= float64(len(vals))
			chain := variancePrecFoldRelErr(t, "single chain", variancePrecChainFold(vals, mean), exact)
			partition := variancePrecFoldRelErr(t, "canonical partition", sqDeviations(vals, mean), exact)
			if partition.Cmp(chain) > 0 {
				t.Fatalf("the canonical partition's error %.3g exceeds the chain's %.3g",
					ratFloat(partition), ratFloat(chain))
			}
			if k == 1 && partition.Cmp(chain) == 0 {
				t.Fatalf("the canonical partition failed to beat the chain on the descending data")
			}
		})
	}
}

func ratFloat(r *big.Rat) float64 {
	f, _ := r.Float64()
	return f
}

// TestVarianceFoldPathAgreement pins the two routes to one answer: a
// float64 sample rides the raw payload, a float32 sample of the same
// dyadic values rides the accessor, and every widening and every mean
// fold sees identical operands, so the two routes must answer identical
// bits, below and above the one-block partition length alike.
func TestVarianceFoldPathAgreement(t *testing.T) {
	build := func(n int) []float64 {
		vals := make([]float64, n)
		for i := range vals {
			// Exact in float32 and in float64 alike, with a magnitude
			// spread so the fold has something to chew.
			vals[i] = float64(i%97-48) / 4 * math.Pow(2, float64(i%11-5))
		}
		return vals
	}
	for _, n := range []int{3000, 1<<16 + 7} {
		vals := build(n)
		wide, err := core.FromFloats(vals, n)
		if err != nil {
			t.Fatalf("FromFloats: %v", err)
		}
		narrow := core.New(core.Float32, n)
		for i, v := range vals {
			narrow.RawFloat32s()[i] = float32(v)
		}
		varW, err := Var(wide)
		if err != nil {
			t.Fatalf("Var: %v", err)
		}
		varN, err := Var(narrow)
		if err != nil {
			t.Fatalf("Var narrow: %v", err)
		}
		stdW, err := Std(wide)
		if err != nil {
			t.Fatalf("Std: %v", err)
		}
		stdN, err := Std(narrow)
		if err != nil {
			t.Fatalf("Std narrow: %v", err)
		}
		vsW, err := VarSample(wide)
		if err != nil {
			t.Fatalf("VarSample: %v", err)
		}
		vsN, err := VarSample(narrow)
		if err != nil {
			t.Fatalf("VarSample narrow: %v", err)
		}
		if math.Float64bits(varW) != math.Float64bits(varN) {
			t.Fatalf("n = %d: Var answers %.17g on the accessor path and %.17g on the payload path", n, varN, varW)
		}
		if math.Float64bits(stdW) != math.Float64bits(stdN) {
			t.Fatalf("n = %d: Std answers %.17g on the accessor path and %.17g on the payload path", n, stdN, stdW)
		}
		if math.Float64bits(vsW) != math.Float64bits(vsN) {
			t.Fatalf("n = %d: VarSample answers %.17g on the accessor path and %.17g on the payload path", n, vsN, vsW)
		}
	}
}