package grad

import (
	"math"
	"testing"

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

// Backward-coverage pins for the graph kernels: every test here checks
// a committed leaf gradient against a closed-form identity or a central
// difference of the summed loss, on combinations the op-level tests do
// not build.

// totalLoss sums a non-scalar loss so a central difference matches the
// all-ones seed Backward uses.
func totalLoss(l *Tensor) float64 {
	s := 0.0
	d := l.Data()
	for i := range d.Len() {
		s += d.FloatAt(i)
	}
	return s
}

func abs2c(z complex128) float64 { return real(z)*real(z) + imag(z)*imag(z) }

func mustRecover(vals []float64, sh ...int) *core.Array {
	a, _ := core.FromFloats(vals, sh...)
	return a
}

func mustRecoverComplex(zs []complex128, shape ...int) *core.Array {
	if len(shape) == 0 {
		shape = []int{len(zs)}
	}
	w := append([]complex128(nil), zs...)
	a, _ := core.ComplexFromArray(w, shape...)
	return a
}

// probeCheckCentral builds the loss from base, backpropagates once and
// compares every element of the committed gradient against central
// differences of the summed loss.
func probeCheckCentral(t *testing.T, base *Tensor, name string, build func() (*Tensor, error), tol float64) {
	t.Helper()
	base.ZeroGrad()
	loss, err := build()
	if err != nil {
		t.Fatal(err)
	}
	if err := loss.Backward(); err != nil {
		t.Fatal(err)
	}
	grads := flatFloats(base.Grad())
	orig := flatFloats(base.Data())
	sh := base.Data().Shape()
	for i := range orig {
		const h = 1e-6
		plus := append([]float64(nil), orig...)
		plus[i] += h
		base.ReplaceWith(mustRecover(plus, sh...))
		lp, err := build()
		if err != nil {
			t.Fatal(err)
		}
		minus := append([]float64(nil), orig...)
		minus[i] -= h
		base.ReplaceWith(mustRecover(minus, sh...))
		lm, err := build()
		if err != nil {
			t.Fatal(err)
		}
		num := (totalLoss(lp) - totalLoss(lm)) / (2 * h)
		if math.Abs(grads[i]-num) > tol*math.Max(1, math.Abs(num)) {
			t.Errorf("%s grad[%d] = %g, central %g", name, i, grads[i], num)
		}
		base.ReplaceWith(mustRecover(orig, sh...))
	}
}

// TestFFTParsevalGradient pins the FFT backward through Parseval's
// theorem: L = sum |FFT(x)|^2 has gradient n·2x for real x, because the
// unnormalised forward DFT scales the energy by n.
func TestFFTParsevalGradient(t *testing.T) {
	const n = 16
	vals := make([]float64, n)
	for i := range vals {
		vals[i] = math.Sin(0.7*float64(i)) + 0.3*float64(i%4)
	}
	x, _ := FromFloat64s(vals, true, n)
	f, _ := x.FFT()
	abs2, _ := f.Abs2()
	loss, _ := abs2.Sum()
	if err := loss.Backward(); err != nil {
		t.Fatal(err)
	}
	g := x.Grad()
	for i := range n {
		want := 2 * float64(n) * vals[i]
		if got := g.FloatAt(i); math.Abs(got-want) > 1e-8*math.Max(1, math.Abs(want)) {
			t.Fatalf("Parseval: g[%d] = %g, want %g", i, got, want)
		}
	}
}

// TestFFT2ParsevalGradient is Parseval's identity for the rank-2
// transform, where the backward's scale is the full element count.
func TestFFT2ParsevalGradient(t *testing.T) {
	vals := []float64{1, 2, -1, 0.5, 3, -2, 0.25, 1.5, -0.5, 2, 1, -3}
	x, _ := FromFloat64s(vals, true, 3, 4)
	f, err := x.FFT2()
	if err != nil {
		t.Fatal(err)
	}
	a, err := f.Abs2()
	if err != nil {
		t.Fatal(err)
	}
	loss, err := a.Sum()
	if err != nil {
		t.Fatal(err)
	}
	if err := loss.Backward(); err != nil {
		t.Fatal(err)
	}
	for i := range vals {
		want := 2 * float64(12) * vals[i]
		if got := x.Grad().FloatAt(i); math.Abs(got-want) > 1e-8*math.Max(1, math.Abs(want)) {
			t.Fatalf("FFT2 Parseval g[%d] = %g, want %g", i, got, want)
		}
	}
}

// TestSpectralHalfSpectrumBackward pins the RFFT and IRFFT backwards
// against central differences, the half-spectrum combinatorics
// included: the doubled real part on the mirrored bins and the halved
// self-mirrored bins of the inverse.
func TestSpectralHalfSpectrumBackward(t *testing.T) {
	x, _ := FromFloat64s([]float64{1, -2, 0.5, 3, -1.5, 0.25, 2, -0.75}, true, 8)
	buildF := func() (*Tensor, error) {
		h, err := x.RFFT()
		if err != nil {
			return nil, err
		}
		a, err := h.Abs2()
		if err != nil {
			return nil, err
		}
		return a.Sum()
	}
	probeCheckCentral(t, x, "rfft", buildF, 1e-4)

	// IRFFT: the leaf is the half spectrum; the loss is the summed
	// square of the real signal. The complex leaf's gradient is checked
	// against a Wirtinger central difference of the real loss:
	// dL/dzbar_j = (dL/dRe_j + i dL/dIm_j)/2, each part probed by
	// perturbing that component.
	zr := []float64{3, -1, 0.5, 2, 1.5}
	zi := []float64{0, 1, -0.5, 0.25, -1}
	zs := make([]complex128, 5)
	for i := range zs {
		zs[i] = complex(zr[i], zi[i])
	}
	za, _ := core.ComplexFromArray(zs, 5)
	z := FromArray(za, true)
	buildI := func() (*Tensor, error) {
		sig, err := z.IRFFT(8)
		if err != nil {
			return nil, err
		}
		sq, err := sig.Pow(2)
		if err != nil {
			return nil, err
		}
		return sq.Sum()
	}
	z.ZeroGrad()
	loss, err := buildI()
	if err != nil {
		t.Fatal(err)
	}
	if err := loss.Backward(); err != nil {
		t.Fatal(err)
	}
	const h = 1e-6
	for j := range 5 {
		perturb := func(dr, di float64) float64 {
			ws := append([]complex128(nil), zs...)
			ws[j] = complex(zr[j]+dr, zi[j]+di)
			z.ReplaceWith(mustRecoverComplex(ws, 5))
			l, err := buildI()
			if err != nil {
				t.Fatal(err)
			}
			return totalLoss(l)
		}
		dr := (perturb(h, 0) - perturb(-h, 0)) / (2 * h)
		di := (perturb(0, h) - perturb(0, -h)) / (2 * h)
		want := complex(0.5*dr, 0.5*di)
		if got := z.Grad().ComplexAt(j); abs2c(got-want) > 1e-4*abs2c(want) {
			t.Errorf("irfft grad[%d] = %g, want %g", j, got, want)
		}
		z.ReplaceWith(mustRecoverComplex(zs, 5))
	}
}

// TestCompositeGraphBackward walks a graph spanning slice, matmul,
// tanh, broadcast, abs2 and an axis reduction, and checks both leaves'
// gradients against central differences of the summed loss.
func TestCompositeGraphBackward(t *testing.T) {
	a, _ := FromFloat64s([]float64{1, -2, 3, -4, 5, -6}, true, 2, 3)
	w, _ := FromFloat64s([]float64{0.5, -0.25, 0.125, 1, -2, 0.75}, true, 2, 3)

	build := func() (*Tensor, error) {
		s, err := a.Slice(1, 1, 3) // (2,2)
		if err != nil {
			return nil, err
		}
		mm, err := s.MatMul(w)
		if err != nil {
			return nil, err
		}
		th, err := mm.Tanh()
		if err != nil {
			return nil, err
		}
		bc, err := th.BroadcastTo(2, 2, 3)
		if err != nil {
			return nil, err
		}
		sq, err := bc.Abs2()
		if err != nil {
			return nil, err
		}
		return sq.SumAxis(2)
	}
	numCheck := func(base *Tensor, name string, grads []float64) {
		t.Helper()
		orig := flatFloats(base.Data())
		sh := base.Data().Shape()
		for i := range orig {
			const h = 1e-6
			plus := append([]float64(nil), orig...)
			plus[i] += h
			base.ReplaceWith(mustRecover(plus, sh...))
			lp, err := build()
			if err != nil {
				t.Fatal(err)
			}
			minus := append([]float64(nil), orig...)
			minus[i] -= h
			base.ReplaceWith(mustRecover(minus, sh...))
			lm, err := build()
			if err != nil {
				t.Fatal(err)
			}
			num := (totalLoss(lp) - totalLoss(lm)) / (2 * h)
			if math.Abs(grads[i]-num) > 1e-5*math.Max(1, math.Abs(num)) {
				t.Fatalf("%s[%d]: backward %g, central difference %g", name, i, grads[i], num)
			}
			base.ReplaceWith(mustRecover(orig, sh...))
		}
	}
	// Backward accumulates, so every read starts from a clean slate.
	a.ZeroGrad()
	w.ZeroGrad()
	loss, err := build()
	if err != nil {
		t.Fatal(err)
	}
	if err := loss.Backward(); err != nil {
		t.Fatal(err)
	}
	numCheck(a, "a", flatFloats(a.Grad()))
	a.ZeroGrad()
	w.ZeroGrad()
	loss, err = build()
	if err != nil {
		t.Fatal(err)
	}
	if err := loss.Backward(); err != nil {
		t.Fatal(err)
	}
	numCheck(w, "w", flatFloats(w.Grad()))
}

// TestMatMulBatchedBackward checks the stacked product rule of
// MatMulBatched against central differences on both operands.
func TestMatMulBatchedBackward(t *testing.T) {
	a, _ := FromFloat64s([]float64{1, 2, 3, 4, 5, 6, 7, 8}, true, 2, 2, 2)
	b, _ := FromFloat64s([]float64{0.5, -1, 2, 0.25, 1.5, -0.5, -2, 1}, true, 2, 2, 2)

	build := func() (*Tensor, error) {
		c, err := a.MatMulBatched(b)
		if err != nil {
			return nil, err
		}
		sq, err := c.Abs2()
		if err != nil {
			return nil, err
		}
		return sq.Sum()
	}
	a.ZeroGrad()
	b.ZeroGrad()
	loss, err := build()
	if err != nil {
		t.Fatal(err)
	}
	if err := loss.Backward(); err != nil {
		t.Fatal(err)
	}
	for _, tc := range []struct {
		tensor *Tensor
		name   string
	}{{a, "a"}, {b, "b"}} {
		grads := flatFloats(tc.tensor.Grad())
		orig := flatFloats(tc.tensor.Data())
		sh := tc.tensor.Data().Shape()
		for i := range orig {
			const h = 1e-6
			plus := append([]float64(nil), orig...)
			plus[i] += h
			tc.tensor.ReplaceWith(mustRecover(plus, sh...))
			lp, err := build()
			if err != nil {
				t.Fatal(err)
			}
			minus := append([]float64(nil), orig...)
			minus[i] -= h
			tc.tensor.ReplaceWith(mustRecover(minus, sh...))
			lm, err := build()
			if err != nil {
				t.Fatal(err)
			}
			num := (totalLoss(lp) - totalLoss(lm)) / (2 * h)
			if math.Abs(grads[i]-num) > 1e-4*math.Max(1, math.Abs(num)) {
				t.Fatalf("%s[%d]: backward %g, central difference %g", tc.name, i, grads[i], num)
			}
			tc.tensor.ReplaceWith(mustRecover(orig, sh...))
		}
	}
}

// TestStridedOperandGradientMatchesDenseTwin pins the graph's stride
// discipline: a MatMul over a sliced view and the same product over a
// dense twin must commit identical gradients. The view's leaf receives
// the dense twin's gradient scattered into the sliced span, zeros
// elsewhere.
func TestStridedOperandGradientMatchesDenseTwin(t *testing.T) {
	wVals := []float64{0.5, -0.25, 0.125, 1, -2, 0.75}

	// The strided run: leaf (2,4) -> slice -> matmul.
	full, _ := FromFloat64s([]float64{9, 1, -2, 3, 9, 5, -6, 9}, true, 2, 4)
	s, err := full.Slice(1, 1, 4) // (2,3) strided view holding 1,-2,3 / 5,-6,9
	if err != nil {
		t.Fatal(err)
	}
	w, _ := FromFloat64s(wVals, true, 3, 2)
	mm, err := s.MatMul(w)
	if err != nil {
		t.Fatal(err)
	}
	sq, err := mm.Abs2()
	if err != nil {
		t.Fatal(err)
	}
	loss, err := sq.Sum()
	if err != nil {
		t.Fatal(err)
	}
	if err := loss.Backward(); err != nil {
		t.Fatal(err)
	}
	if full.Grad() == nil || w.Grad() == nil {
		t.Fatal("both leaves must receive a gradient")
	}
	gs := flatFloats(full.Grad())
	wStrided := flatFloats(w.Grad())
	full.ZeroGrad()
	w.ZeroGrad()

	// The dense twin.
	a2, _ := FromFloat64s(flatFloats(s.Data()), true, 2, 3)
	w2, _ := FromFloat64s(wVals, true, 3, 2)
	mm2, err := a2.MatMul(w2)
	if err != nil {
		t.Fatal(err)
	}
	sq2, err := mm2.Abs2()
	if err != nil {
		t.Fatal(err)
	}
	loss2, err := sq2.Sum()
	if err != nil {
		t.Fatal(err)
	}
	if err := loss2.Backward(); err != nil {
		t.Fatal(err)
	}
	aDense := flatFloats(a2.Grad())
	wDense := flatFloats(w2.Grad())

	for i := range wStrided {
		if wStrided[i] != wDense[i] {
			t.Fatalf("w gradient %v over the strided operand, want the dense twin's %v", wStrided, wDense)
		}
	}
	if gs[0] != 0 || gs[4] != 0 {
		t.Fatalf("columns outside the slice carry %g and %g, want 0", gs[0], gs[4])
	}
	if gs[1] != aDense[0] || gs[2] != aDense[1] || gs[3] != aDense[2] ||
		gs[5] != aDense[3] || gs[6] != aDense[4] || gs[7] != aDense[5] {
		t.Fatalf("leaf gradient %v does not scatter the dense gradient %v into the slice", gs, aDense)
	}
}

// TestHessianMatchesAnalyticForm differentiates
// f(x, y) = x^3 y + exp(x) log(y+2) twice and compares the answer with
// the closed form:
//
//	dxx = 6xy + exp(x)log(y+2), dxy = 3x^2 + exp(x)/(y+2),
//	dyy = -exp(x)/(y+2)^2.
func TestHessianMatchesAnalyticForm(t *testing.T) {
	f := func(v *Tensor) (*Tensor, error) {
		x, err := v.Slice(0, 0, 1)
		if err != nil {
			return nil, err
		}
		y, err := v.Slice(0, 1, 2)
		if err != nil {
			return nil, err
		}
		x3, err := x.Pow(3)
		if err != nil {
			return nil, err
		}
		t1, err := x3.Mul(y)
		if err != nil {
			return nil, err
		}
		e, err := x.Exp()
		if err != nil {
			return nil, err
		}
		y2, err := y.Add(FromArray(mustRecover([]float64{2}, 1), false))
		if err != nil {
			return nil, err
		}
		lg, err := y2.Log()
		if err != nil {
			return nil, err
		}
		t2, err := e.Mul(lg)
		if err != nil {
			return nil, err
		}
		return t1.Add(t2)
	}
	x, _ := FromFloat64s([]float64{0.4, 1.7}, true, 2)
	h, err := Hessian(f, x, HessianOptions{Step: 1e-5})
	if err != nil {
		t.Fatal(err)
	}
	xx, yy := 0.4, 1.7
	ex, ly := math.Exp(xx), math.Log(yy+2)
	want := [4]float64{
		6*xx*yy + ex*ly, 3*xx*xx + ex/(yy+2),
		3*xx*xx + ex/(yy+2), -ex / ((yy + 2) * (yy + 2)),
	}
	for i := range 4 {
		if math.Abs(h.FloatAt(i)-want[i]) > 1e-4*math.Max(1, math.Abs(want[i])) {
			t.Fatalf("Hessian[%d] = %g, want %g", i, h.FloatAt(i), want[i])
		}
	}
}

// TestTripleUseLeafGradient folds a leaf through three nodes and pins
// the accumulated gradient: dL/dz = 2(v^2+v)(2v+1) for
// L = (v^2+v)^2.
func TestTripleUseLeafGradient(t *testing.T) {
	for _, v := range []float64{0.5, -1.25, 2} {
		z, _ := FromFloat64s([]float64{v}, true, 1)
		z2, _ := z.Mul(z)
		s, err := z2.Add(z)
		if err != nil {
			t.Fatal(err)
		}
		sq, err := s.Pow(2)
		if err != nil {
			t.Fatal(err)
		}
		loss, err := sq.Sum()
		if err != nil {
			t.Fatal(err)
		}
		if err := loss.Backward(); err != nil {
			t.Fatal(err)
		}
		want := 2 * (v*v + v) * (2*v + 1)
		if got := z.Grad().FloatAt(0); math.Abs(got-want) > 1e-12*math.Max(1, math.Abs(want)) {
			t.Fatalf("triple use at z=%g: g = %g, want %g", v, got, want)
		}
	}
}

// TestComplexAbs2LeafGradient pins the Wirtinger gradient of a summed
// |z|^2 loss: the leaf receives z itself.
func TestComplexAbs2LeafGradient(t *testing.T) {
	zr := []float64{0.3, -1.2, 0.7}
	zi := []float64{-0.4, 0.8, 1.1}
	zs := make([]complex128, 3)
	for i := range zs {
		zs[i] = complex(zr[i], zi[i])
	}
	za, _ := core.ComplexFromArray(zs, 3)
	z := FromArray(za, true)
	abs2, err := z.Abs2()
	if err != nil {
		t.Fatal(err)
	}
	loss, err := abs2.Sum()
	if err != nil {
		t.Fatal(err)
	}
	if err := loss.Backward(); err != nil {
		t.Fatal(err)
	}
	g := z.Grad()
	for i := range 3 {
		want := complex(zr[i], zi[i])
		if got := g.ComplexAt(i); got != want {
			t.Fatalf("|z|^2 leaf grad[%d] = %g, want %g", i, got, want)
		}
	}
}

// TestComplexMatMulBackward checks the Wirtinger adjoint of MatMul
// against component-wise central differences of a real loss. The leaf
// gradient is dL/dzbar = (dL/dRe + i dL/dIm)/2 under the package
// convention.
func TestComplexMatMulBackward(t *testing.T) {
	az := []complex128{complex(1, 0.5), complex(-0.5, 2), complex(0.25, -1), complex(2, 0.25)}
	bz := []complex128{complex(0.5, 1), complex(-1, 0.25), complex(1.5, -0.5), complex(0.75, 1.25)}
	aa, _ := core.ComplexFromArray(az, 2, 2)
	ba, _ := core.ComplexFromArray(bz, 2, 2)
	a := FromArray(aa, true)
	b := FromArray(ba, true)
	build := func() (*Tensor, error) {
		m, err := a.MatMul(b)
		if err != nil {
			return nil, err
		}
		r, err := m.Real()
		if err != nil {
			return nil, err
		}
		sq, err := r.Abs2()
		if err != nil {
			return nil, err
		}
		return sq.Sum()
	}
	a.ZeroGrad()
	b.ZeroGrad()
	loss, err := build()
	if err != nil {
		t.Fatal(err)
	}
	if err := loss.Backward(); err != nil {
		t.Fatal(err)
	}
	const h = 1e-6
	check := func(base *Tensor, vals []complex128, name string) {
		t.Helper()
		for j := range len(vals) {
			perturb := func(dr, di float64) float64 {
				ws := append([]complex128(nil), vals...)
				ws[j] = complex(real(vals[j])+dr, imag(vals[j])+di)
				base.ReplaceWith(mustRecoverComplex(ws, 2, 2))
				l, err := build()
				if err != nil {
					t.Fatal(err)
				}
				return totalLoss(l)
			}
			dr := (perturb(h, 0) - perturb(-h, 0)) / (2 * h)
			di := (perturb(0, h) - perturb(0, -h)) / (2 * h)
			want := complex(0.5*dr, 0.5*di)
			if got := base.Grad().ComplexAt(j); abs2c(got-want) > 1e-4*abs2c(want) {
				t.Errorf("complex matmul %s grad[%d] = %g, want %g", name, j, got, want)
			}
			base.ReplaceWith(mustRecoverComplex(vals, 2, 2))
		}
	}
	check(a, az, "a")
	check(b, bz, "b")
}

// TestConcatMixedDtypeBackward checks the join's backward on a real and
// a complex side at once: the real side receives 2 Re g through the
// narrowing, the complex side its Wirtinger gradient.
func TestConcatMixedDtypeBackward(t *testing.T) {
	r, _ := FromFloat64s([]float64{1, -2, 3}, true, 3)
	zr := []complex128{complex(0.5, 1), complex(-1.5, 0.25)}
	za, _ := core.ComplexFromArray(zr, 2)
	z := FromArray(za, true)
	build := func() (*Tensor, error) {
		c, err := r.Concat(z, 0)
		if err != nil {
			return nil, err
		}
		re, err := c.Real()
		if err != nil {
			return nil, err
		}
		sq, err := re.Abs2()
		if err != nil {
			return nil, err
		}
		return sq.Sum()
	}
	r.ZeroGrad()
	z.ZeroGrad()
	loss, err := build()
	if err != nil {
		t.Fatal(err)
	}
	if err := loss.Backward(); err != nil {
		t.Fatal(err)
	}
	grads := flatFloats(r.Grad())
	orig := flatFloats(r.Data())
	for i := range orig {
		const h = 1e-6
		plus := append([]float64(nil), orig...)
		plus[i] += h
		r.ReplaceWith(mustRecover(plus, 3))
		lp, _ := build()
		minus := append([]float64(nil), orig...)
		minus[i] -= h
		r.ReplaceWith(mustRecover(minus, 3))
		lm, _ := build()
		num := (totalLoss(lp) - totalLoss(lm)) / (2 * h)
		if math.Abs(grads[i]-num) > 1e-5*math.Max(1, math.Abs(num)) {
			t.Errorf("concat real side grad[%d] = %g, central %g", i, grads[i], num)
		}
		r.ReplaceWith(mustRecover(orig, 3))
	}
	for j := range 2 {
		perturb := func(dr, di float64) float64 {
			ws := append([]complex128(nil), zr...)
			ws[j] = complex(real(zr[j])+dr, imag(zr[j])+di)
			z.ReplaceWith(mustRecoverComplex(ws, 2))
			l, _ := build()
			return totalLoss(l)
		}
		dr := (perturb(1e-6, 0) - perturb(-1e-6, 0)) / (2e-6)
		di := (perturb(0, 1e-6) - perturb(0, -1e-6)) / (2e-6)
		want := complex(0.5*dr, 0.5*di)
		if got := z.Grad().ComplexAt(j); abs2c(got-want) > 1e-4*abs2c(want) {
			t.Errorf("concat complex side grad[%d] = %g, want %g", j, got, want)
		}
		z.ReplaceWith(mustRecoverComplex(zr, 2))
	}
}

// TestTransposeAxesBackward reverses an axis permutation under a
// non-linear loss and checks the gradient against central differences.
func TestTransposeAxesBackward(t *testing.T) {
	x, _ := FromFloat64s([]float64{1, 2, 3, 4, 5, 6, 7, 8}, true, 2, 2, 2)
	build := func() (*Tensor, error) {
		p, err := x.TransposeAxes(2, 0, 1)
		if err != nil {
			return nil, err
		}
		sq, err := p.Abs2()
		if err != nil {
			return nil, err
		}
		w, err := sq.Mul(p)
		if err != nil {
			return nil, err
		}
		return w.Sum()
	}
	probeCheckCentral(t, x, "transposeaxes", build, 1e-5)
}

// TestUnaryKernelChainBackward pushes Sqrt, Log, Sigmoid, Scale, Neg,
// Mean and Pow through one graph and checks the committed gradient.
func TestUnaryKernelChainBackward(t *testing.T) {
	x, _ := FromFloat64s([]float64{0.5, 1.5, 2.5, 3.5}, true, 4)
	build := func() (*Tensor, error) {
		s, err := x.Sqrt()
		if err != nil {
			return nil, err
		}
		l, err := s.Log()
		if err != nil {
			return nil, err
		}
		sg, err := x.Sigmoid()
		if err != nil {
			return nil, err
		}
		m, err := l.Mul(sg)
		if err != nil {
			return nil, err
		}
		sc, err := m.Scale(2.5)
		if err != nil {
			return nil, err
		}
		ng, err := sc.Neg()
		if err != nil {
			return nil, err
		}
		mn, err := ng.Mean()
		if err != nil {
			return nil, err
		}
		return mn.Pow(2)
	}
	probeCheckCentral(t, x, "unary-chain", build, 1e-4)
}

// TestBackwardAccumulatesUntilZeroGrad pins the accumulation contract
// end to end: two Backward calls double the gradient, ZeroGrad resets.
func TestBackwardAccumulatesUntilZeroGrad(t *testing.T) {
	x, _ := FromFloat64s([]float64{2}, true, 1)
	build := func() (*Tensor, error) {
		s, err := x.Pow(2)
		if err != nil {
			return nil, err
		}
		return s.Sum()
	}
	l1, err := build()
	if err != nil {
		t.Fatal(err)
	}
	if err := l1.Backward(); err != nil {
		t.Fatal(err)
	}
	l2, err := build()
	if err != nil {
		t.Fatal(err)
	}
	if err := l2.Backward(); err != nil {
		t.Fatal(err)
	}
	if got := x.Grad().FloatAt(0); got != 8 {
		t.Fatalf("accumulated g = %g, want 8", got)
	}
	x.ZeroGrad()
	l3, err := build()
	if err != nil {
		t.Fatal(err)
	}
	if err := l3.Backward(); err != nil {
		t.Fatal(err)
	}
	if got := x.Grad().FloatAt(0); got != 4 {
		t.Fatalf("post-zero g = %g, want 4", got)
	}
}

// TestFloat32LeafKeepsGradientDtype pins the dtype contract on a
// float32 leaf: the gradient narrows to the leaf's width.
func TestFloat32LeafKeepsGradientDtype(t *testing.T) {
	v := []float32{1.5, -2.5}
	a, _ := core.FromFloat32Slice(v, 2)
	x := FromArray(a, true)
	y, err := x.Pow(2)
	if err != nil {
		t.Fatal(err)
	}
	l, err := y.Sum()
	if err != nil {
		t.Fatal(err)
	}
	if err := l.Backward(); err != nil {
		t.Fatal(err)
	}
	g := x.Grad()
	if g.Dtype() != core.Float32 {
		t.Fatalf("grad dtype = %s, want float32", g.Dtype())
	}
	for i := range 2 {
		want := 2 * float64(v[i])
		if got := g.FloatAt(i); math.Abs(got-want) > 1e-6 {
			t.Fatalf("g[%d] = %g, want %g", i, got, want)
		}
	}
}

// TestBroadcastToBackwardMatchesCentralDifferences and its SumAxis
// neighbour cover the reduction/expansion pair in isolation.
func TestBroadcastToBackwardMatchesCentralDifferences(t *testing.T) {
	x, _ := FromFloat64s([]float64{1, -2, 3, -4, 5, -6}, true, 2, 3)
	build := func() (*Tensor, error) {
		bc, err := x.BroadcastTo(2, 2, 3)
		if err != nil {
			return nil, err
		}
		sq, err := bc.Abs2()
		if err != nil {
			return nil, err
		}
		return sq.Sum()
	}
	probeCheckCentral(t, x, "broadcast", build, 1e-5)
}

// TestSumAxisBackwardMatchesCentralDifferences reduces the middle axis
// with a non-scalar loss, so the backward must scatter into both rows.
func TestSumAxisBackwardMatchesCentralDifferences(t *testing.T) {
	x, _ := FromFloat64s([]float64{1, -2, 3, -4, 5, -6}, true, 2, 3)
	build := func() (*Tensor, error) {
		sq, err := x.Abs2()
		if err != nil {
			return nil, err
		}
		return sq.SumAxis(1)
	}
	probeCheckCentral(t, x, "sumaxis", build, 1e-5)
}

// TestMatMulTanhBackwardMatchesCentralDifferences checks the matmul
// product rule under a tanh on top, non-scalar loss included.
func TestMatMulTanhBackwardMatchesCentralDifferences(t *testing.T) {
	a, _ := FromFloat64s([]float64{1, -2, 3, -4}, true, 2, 2)
	w, _ := FromFloat64s([]float64{0.5, -0.25, 1, -2}, true, 2, 2)
	build := func() (*Tensor, error) {
		mm, err := a.MatMul(w)
		if err != nil {
			return nil, err
		}
		return mm.Tanh()
	}
	a.ZeroGrad()
	w.ZeroGrad()
	loss, err := build()
	if err != nil {
		t.Fatal(err)
	}
	if err := loss.Backward(); err != nil {
		t.Fatal(err)
	}
	check := func(base *Tensor, name string) {
		t.Helper()
		grads := flatFloats(base.Grad())
		orig := flatFloats(base.Data())
		sh := base.Data().Shape()
		for i := range orig {
			const h = 1e-6
			plus := append([]float64(nil), orig...)
			plus[i] += h
			base.ReplaceWith(mustRecover(plus, sh...))
			lp, _ := build()
			minus := append([]float64(nil), orig...)
			minus[i] -= h
			base.ReplaceWith(mustRecover(minus, sh...))
			lm, _ := build()
			var num float64
			for j := range lp.Data().Len() {
				num += lp.Data().FloatAt(j) - lm.Data().FloatAt(j)
			}
			num /= 2 * h
			if math.Abs(grads[i]-num) > 1e-5*math.Max(1, math.Abs(num)) {
				t.Errorf("%s[%d] backward %g, central %g", name, i, grads[i], num)
			}
			base.ReplaceWith(mustRecover(orig, sh...))
		}
	}
	check(a, "a")
	check(w, "w")
}

// TestNewtonCGSolvesLeastSquares pins the truncated Newton method on a
// small full-rank least-squares problem whose solution is A^-1 b.
func TestNewtonCGSolvesLeastSquares(t *testing.T) {
	ar := []float64{2, 0.5, 1, 3}
	br := []float64{1, -1}
	A, _ := FromFloat64s(ar, false, 2, 2)
	B, _ := FromFloat64s(br, false, 2)
	f := func(x *Tensor) (*Tensor, error) {
		ax, err := A.MatMul(x)
		if err != nil {
			return nil, err
		}
		d, err := ax.Sub(B)
		if err != nil {
			return nil, err
		}
		s, err := d.Abs2()
		if err != nil {
			return nil, err
		}
		return s.Sum()
	}
	x0 := mustRecover([]float64{0, 0}, 2)
	x, fv, err := MinimiseNewtonCG(f, x0, NewtonCGOptions{})
	if err != nil {
		t.Fatal(err)
	}
	det := 2*3 - 0.5*1
	wantX := (3*1 - 0.5*(-1)) / det
	wantY := (2*(-1) - 1*1) / det
	if math.Abs(x.FloatAt(0)-wantX) > 1e-5 || math.Abs(x.FloatAt(1)-wantY) > 1e-5 {
		t.Fatalf("NewtonCG point (%g, %g), want (%g, %g)", x.FloatAt(0), x.FloatAt(1), wantX, wantY)
	}
	if fv > 1e-10 {
		t.Fatalf("NewtonCG f = %g, want ~0", fv)
	}
}