// 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/core"
)

// Callback-hygiene pins: callbacks that were handed arrays aliasing
// reused scratch, gradients that divided by a stencil clamped to the
// same wall by equal bounds, and NaN objectives that outran their
// diagnosis.

// TestFindRootSystemCallbackImmutability: the residual callback
// received an aliasing view of the reused iterate, so a callback that
// retained its argument observed the solver's later writes into it.
func TestFindRootSystemCallbackImmutability(t *testing.T) {
	var kept *core.Array
	f := func(x *core.Array) (*core.Array, error) {
		if kept == nil {
			kept = x // retain the first argument, as a caching callback would
		}
		return core.FromFloats([]float64{x.FloatAt(0) - 1, x.FloatAt(1) - 2}, 2)
	}
	x0, err := core.FromFloats([]float64{0, 0}, 2)
	if err != nil {
		t.Fatal(err)
	}
	root, _, err := FindRootSystem(f, x0, RootSystemOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if root.FloatAt(0) != 1 || root.FloatAt(1) != 2 {
		t.Fatalf("root = %v, %v", root.FloatAt(0), root.FloatAt(1))
	}
	if kept.FloatAt(0) != 0 || kept.FloatAt(1) != 0 {
		t.Fatalf("the retained argument was mutated to %v, %v", kept.FloatAt(0), kept.FloatAt(1))
	}
}

// TestLBFGSEqualBoundsFiniteDifferences: a coordinate pinned by equal
// bounds clamped both stencil points to the same wall and the central
// difference divided 0/0 into a NaN gradient.
func TestLBFGSEqualBoundsFiniteDifferences(t *testing.T) {
	f := func(x *core.Array) (float64, error) {
		return (x.FloatAt(0)-3)*(x.FloatAt(0)-3) + (x.FloatAt(1)-1)*(x.FloatAt(1)-1), nil
	}
	x0, err := core.FromFloats([]float64{0, 0}, 2)
	if err != nil {
		t.Fatal(err)
	}
	opts := LBFGSOptions{Lower: []float64{1, -10}, Upper: []float64{1, 10}}
	pt, val, err := MinimiseLBFGS(f, nil, x0, opts)
	if err != nil {
		t.Fatalf("MinimiseLBFGS with an equality-pinned coordinate: %v", err)
	}
	if pt.FloatAt(0) != 1 {
		t.Fatalf("pinned coordinate = %g, want 1", pt.FloatAt(0))
	}
	if math.Abs(pt.FloatAt(1)-1) > 1e-6 || math.Abs(val-4) > 1e-9 {
		t.Fatalf("free coordinate = %g, value = %g (want 1 and 4)", pt.FloatAt(1), val)
	}
}

// TestLBFGSRejectsNaNGradient: a NaN gradient coordinate skipped the
// projected-gradient max and could read as converged.
func TestLBFGSRejectsNaNGradient(t *testing.T) {
	f := func(x *core.Array) (float64, error) { return x.FloatAt(0) * x.FloatAt(0), nil }
	grad := func(x *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{math.NaN()}, 1)
	}
	x0, err := core.FromFloats([]float64{1}, 1)
	if err != nil {
		t.Fatal(err)
	}
	if _, _, err := MinimiseLBFGS(f, grad, x0, LBFGSOptions{}); err == nil || !strings.Contains(err.Error(), "non-finite") {
		t.Fatalf("MinimiseLBFGS with a NaN gradient: err = %v", err)
	}
}

// TestMinimiseRejectsNaNObjective: a NaN vertex burned the budget and
// was reported as a budget problem, or with AllowBudgetExit came back
// as the answer.
func TestMinimiseRejectsNaNObjective(t *testing.T) {
	f := func(x *core.Array) (float64, error) { return math.NaN(), nil }
	x0, err := core.FromFloats([]float64{1}, 1)
	if err != nil {
		t.Fatal(err)
	}
	if _, _, err := Minimise(f, x0, MinimiseOptions{}); err == nil || !strings.Contains(err.Error(), "non-finite") {
		t.Fatalf("Minimise with a NaN objective: err = %v", err)
	}
}

// TestMinimiseConstrainedReturnsObjective: the returned value was the
// augmented Lagrangian the inner solver minimised, not f at the
// answer.
func TestMinimiseConstrainedReturnsObjective(t *testing.T) {
	f := func(x *core.Array) (float64, error) {
		return (x.FloatAt(0) - 2) * (x.FloatAt(0) - 2), nil
	}
	grad := func(x *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{2 * (x.FloatAt(0) - 2)}, 1)
	}
	x0, err := core.FromFloats([]float64{0}, 1)
	if err != nil {
		t.Fatal(err)
	}
	cons := LinearConstraints{
		A:     mustFloats2D(t, []float64{1}, 1, 1),
		Lower: []float64{1},
		Upper: []float64{1},
	}
	pt, val, err := MinimiseConstrained(f, grad, x0, cons, LBFGSOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if math.Abs(pt.FloatAt(0)-1) > 1e-6 {
		t.Fatalf("constrained minimum at %g, want 1", pt.FloatAt(0))
	}
	direct, err := f(pt)
	if err != nil {
		t.Fatal(err)
	}
	if val != direct {
		t.Fatalf("returned value = %.17g, want f at the answer = %.17g exactly", val, direct)
	}
}

func mustFloats2D(t *testing.T, vals []float64, rows, cols int) *core.Array {
	t.Helper()
	a, err := core.FromFloats(vals, rows, cols)
	if err != nil {
		t.Fatal(err)
	}
	return a
}