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

package optim

import (
	"math"
	"strings"
	"testing"

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

// jacobianBuilds counts the central-difference Jacobian builds in a
// recorded residual trace. A build is n consecutive ± pairs, column j
// perturbed first, each pair one stencil width √ε·max(1, |xⱼ|) about
// its base point: the exact pattern FindRootSystem's sweep produces,
// which backtracking trials (single points, several moving
// coordinates at once) never match.
func jacobianBuilds(points [][]float64) int {
	n := len(points[0])
	stencil := math.Sqrt(base.EpsF)
	matchPair := func(p, q []float64, col int) bool {
		diff := -1
		for k := range n {
			if p[k] != q[k] {
				if diff != -1 {
					return false
				}
				diff = k
			}
		}
		if diff != col {
			return false
		}
		mid := (p[col] + q[col]) / 2
		eps := math.Abs(p[col]-q[col]) / 2
		want := stencil * math.Max(1, math.Abs(mid))
		return math.Abs(eps-want) <= 1e-6*want
	}
	count := 0
	i := 0
	for i+2*n <= len(points) {
		built := true
		for col := range n {
			if !matchPair(points[i+2*col], points[i+2*col+1], col) {
				built = false
				break
			}
		}
		if built {
			count++
			i += 2 * n
			continue
		}
		i++
	}
	return count
}

// traceResidual wraps a residual so every evaluation's point is
// recorded, for the stencil counter to walk.
func traceResidual(t *testing.T, trace *[][]float64, n int, r func(x []float64) []float64) func(*core.Array) (*core.Array, error) {
	return func(x *core.Array) (*core.Array, error) {
		v := make([]float64, n)
		for i := range n {
			v[i] = x.FloatAt(i)
		}
		*trace = append(*trace, v)
		return mustFloats(t, r(v)), nil
	}
}

// TestFindRootSystemBroydenOneJacobian pins the option's promise on
// the analytic systems: with UseBroyden the same roots are reached
// within tolerance and exactly one numerical Jacobian is built, the
// one at the start.
func TestFindRootSystemBroydenOneJacobian(t *testing.T) {
	cases := []struct {
		name  string
		n     int
		start []float64
		r     func(x []float64) []float64
		want  []float64
	}{
		{"circle-line", 2, []float64{0.5, 0.5},
			func(x []float64) []float64 { return []float64{x[0]*x[0] + x[1]*x[1] - 4, x[0] - x[1]} },
			[]float64{math.Sqrt2, math.Sqrt2}},
		{"circle-hyperbola", 2, []float64{0.4, 2.2},
			func(x []float64) []float64 { return []float64{x[0]*x[0] + x[1]*x[1] - 5, x[0]*x[1] - 2} },
			nil},
		{"trig", 2, []float64{0.3, 0.1},
			func(x []float64) []float64 { return []float64{math.Cos(x[0]) - x[1], math.Sin(x[0]) - x[1]} },
			[]float64{math.Pi / 4, math.Sqrt2 / 2}},
	}
	for _, tc := range cases {
		var trace [][]float64
		residual := traceResidual(t, &trace, tc.n, tc.r)
		x, res, err := FindRootSystem(residual, mustFloats(t, tc.start), RootSystemOptions{UseBroyden: true})
		if err != nil {
			t.Fatalf("%s: FindRootSystem(UseBroyden): %v", tc.name, err)
		}
		if res > 1e-10 {
			t.Fatalf("%s: residual %g, want ≤ 1e-10", tc.name, res)
		}
		if tc.want != nil {
			for i := range tc.n {
				if math.Abs(x.FloatAt(i)-tc.want[i]) > 1e-9 {
					t.Fatalf("%s: x[%d] = %.12g, want %.12g", tc.name, i, x.FloatAt(i), tc.want[i])
				}
			}
		} else {
			// The hyperbola's two roots are (1, 2) and (2, 1).
			s1 := math.Abs(x.FloatAt(0)-1) < 1e-9 && math.Abs(x.FloatAt(1)-2) < 1e-9
			s2 := math.Abs(x.FloatAt(0)-2) < 1e-9 && math.Abs(x.FloatAt(1)-1) < 1e-9
			if !s1 && !s2 {
				t.Fatalf("%s: solution = (%.12g, %.12g), want (1, 2) or (2, 1)",
					tc.name, x.FloatAt(0), x.FloatAt(1))
			}
		}
		if got := jacobianBuilds(trace); got != 1 {
			t.Fatalf("%s: %d numerical Jacobian builds, want 1", tc.name, got)
		}
	}
}

// TestFindRootSystemBroydenEightUnknowns pins the option on a harder
// system: eight coupled nonlinear equations with the known root
// xᵢ = i+1, converged under the default budget with the single
// starting Jacobian.
func TestFindRootSystemBroydenEightUnknowns(t *testing.T) {
	const n = 8
	var trace [][]float64
	residual := traceResidual(t, &trace, n, func(x []float64) []float64 {
		r := make([]float64, n)
		for i := range n {
			r[i] = x[i]*x[i] - float64(i+1)*float64(i+1)
			for j := range n {
				if j != i {
					r[i] += 0.05 * (x[j] - float64(j+1))
				}
			}
		}
		return r
	})
	start := make([]float64, n)
	for i := range n {
		start[i] = 0.5 * float64(i+1)
	}
	x, res, err := FindRootSystem(residual, mustFloats(t, start), RootSystemOptions{UseBroyden: true})
	if err != nil {
		t.Fatalf("FindRootSystem(UseBroyden, 8 unknowns): %v", err)
	}
	if res > 1e-10 {
		t.Fatalf("residual %g, want ≤ 1e-10", res)
	}
	for i := range n {
		if math.Abs(x.FloatAt(i)-float64(i+1)) > 1e-9 {
			t.Fatalf("x[%d] = %.12g, want %d", i, x.FloatAt(i), i+1)
		}
	}
	if got := jacobianBuilds(trace); got != 1 {
		t.Fatalf("%d numerical Jacobian builds, want 1", got)
	}
}

// TestFindRootSystemBroydenSingularJacobianDescends pins the singular
// escape under the option. The duplicated equation x² = 1 has a rank-
// one Jacobian at every point, so no Newton solve ever succeeds and
// the steepest-descent fallback carries the iteration. From (2, 2) the
// damped descent lands on a root at once; from (1.5, 1.5) the descent
// cannot reach one within the budget, and the run refuses with the
// budget error while the trace shows the Jacobian was rebuilt every
// single round, the same restart path a degraded update takes.
func TestFindRootSystemBroydenSingularJacobianDescends(t *testing.T) {
	rankOne := func(x []float64) []float64 {
		return []float64{x[0]*x[0] - 1, x[0]*x[0] - 1}
	}
	var trace [][]float64
	x, res, err := FindRootSystem(traceResidual(t, &trace, 2, rankOne),
		mustFloats(t, []float64{2, 2}), RootSystemOptions{UseBroyden: true})
	if err != nil {
		t.Fatalf("FindRootSystem(UseBroyden, singular Jacobian): %v", err)
	}
	if res > 1e-10 {
		t.Fatalf("residual %g, want ≤ 1e-10", res)
	}
	if math.Abs(math.Abs(x.FloatAt(0))-1) > 1e-8 {
		t.Fatalf("x[0] = %.12g, want a root of x² = 1", x.FloatAt(0))
	}
	// The hopeless start: the refusal is honest and the rebuilds are
	// visible in the trace, one per round while the inverse stays
	// unfit.
	trace = nil
	if _, _, err := FindRootSystem(traceResidual(t, &trace, 2, rankOne),
		mustFloats(t, []float64{1.5, 1.5}), RootSystemOptions{UseBroyden: true}); err == nil {
		t.Fatal("a stalling singular system: want the budget refusal")
	}
	if got := jacobianBuilds(trace); got < 50 {
		t.Fatalf("%d numerical Jacobian builds, want one per round: a singular inverse rebuilds every time", got)
	}
}

// TestFindRootSystemBroydenBudgetStillRefused pins that the shared
// error contract survives the option: an impossible tolerance under
// UseBroyden is a budget refusal, not a silent answer.
func TestFindRootSystemBroydenBudgetStillRefused(t *testing.T) {
	residual := func(x *core.Array) (*core.Array, error) {
		cx, cy := x.FloatAt(0), x.FloatAt(1)
		return mustFloats(t, []float64{cx*cx + cy*cy - 4, cx - cy}), nil
	}
	_, _, err := FindRootSystem(residual, mustFloats(t, []float64{1, 1}),
		RootSystemOptions{UseBroyden: true, MaxIterations: 1, Tolerance: 1e-20})
	if err == nil || !strings.Contains(err.Error(), "MaxIterations=1") {
		t.Fatalf("err = %v, want the budget refusal", err)
	}
}

// TestBroydenMaintainSecantCondition pins the update formula itself:
// after the rank-one correction the maintained inverse satisfies the
// secant equation H·y = s exactly to rounding, which is what makes the
// following steps quasi-Newton at all.
func TestBroydenMaintainSecantCondition(t *testing.T) {
	h := [][]float64{{2, 0.5}, {-1, 3}}
	s := []float64{0.3, -0.7}
	y := []float64{1.1, 0.4}
	r := []float64{0.9, -0.2}
	ok, stalled := broydenMaintain(h, s, y, r, 0.1, 1.0, 0)
	if !ok || stalled != 0 {
		t.Fatalf("healthy update rejected: ok = %v, stalled = %d", ok, stalled)
	}
	for i := range 2 {
		hy := h[i][0]*y[0] + h[i][1]*y[1]
		if math.Abs(hy-s[i]) > 1e-12 {
			t.Fatalf("secant equation violated: (H·y)[%d] = %.17g, want %.17g", i, hy, s[i])
		}
	}
}

// TestBroydenMaintainRestartTriggers pins the documented restart
// triggers of the rank-one maintenance: a degenerate denominator
// orders a rebuild at once, a rounding-level residual change likewise,
// and two consecutive steps without a fall of the residual infinity
// norm do what one cannot. A refused update leaves the inverse
// untouched, so the rebuild starts from a Jacobian and not from a
// half-updated one.
func TestBroydenMaintainRestartTriggers(t *testing.T) {
	h := [][]float64{{1, 0}, {0, 1}}
	// y all zero: the denominator trigger.
	if ok, _ := broydenMaintain(h, []float64{1, 1}, []float64{0, 0}, []float64{1, 1}, 1, 2, 0); ok {
		t.Fatal("a zero residual change was folded into the inverse")
	}
	// y at rounding level against the residual's own scale.
	if ok, _ := broydenMaintain(h, []float64{1, 1}, []float64{1e-20, 0}, []float64{1, 1}, 1, 2, 0); ok {
		t.Fatal("a rounding-level residual change was folded into the inverse")
	}
	// One stalled step keeps the inverse fit but counts the stall.
	ok, stalled := broydenMaintain(h, []float64{0.1, 0}, []float64{0.5, 0.5}, []float64{1, 1}, 2, 2, 0)
	if !ok || stalled != 1 {
		t.Fatalf("first stall: ok = %v, stalled = %d, want the inverse kept and the stall counted", ok, stalled)
	}
	// A falling step resets the count.
	ok, stalled = broydenMaintain(h, []float64{0.1, 0}, []float64{0.5, 0.5}, []float64{1, 1}, 1, 2, stalled)
	if !ok || stalled != 0 {
		t.Fatalf("falling step: ok = %v, stalled = %d, want the count reset", ok, stalled)
	}
	// The second consecutive stall orders a rebuild and leaves the
	// inverse untouched.
	before := [2][2]float64{{h[0][0], h[0][1]}, {h[1][0], h[1][1]}}
	ok, stalled = broydenMaintain(h, []float64{0.1, 0}, []float64{0.5, 0.5}, []float64{1, 1}, 2, 2, 1)
	if ok || stalled != 0 {
		t.Fatalf("second stall: ok = %v, stalled = %d, want a rebuild ordered", ok, stalled)
	}
	for i := range 2 {
		for k := range 2 {
			if h[i][k] != before[i][k] {
				t.Fatalf("a refused update moved h[%d][%d] from %g to %g", i, k, before[i][k], h[i][k])
			}
		}
	}
}