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

// TestLBFGSRejectsNonFiniteObjectiveAndGradient pins that a
// non-finite objective value on the finite-difference path is an
// error: the NaN used to slip past the projected-gradient max and read
// as converged.
func TestLBFGSRejectsNonFiniteObjectiveAndGradient(t *testing.T) {
	f := func(x *core.Array) (float64, error) {
		if x.FloatAt(0) < 0 {
			return math.NaN(), nil
		}
		return (x.FloatAt(0) - 3) * (x.FloatAt(0) - 3), nil
	}
	x0 := mustFloats(t, []float64{0}, 1)
	if _, _, err := MinimiseLBFGS(f, nil, x0, LBFGSOptions{}); err == nil {
		t.Fatal("expected an error for a NaN objective on the FD path")
	}
	// A NaN objective at the start is refused before any step.
	f2 := func(*core.Array) (float64, error) { return math.Inf(1), nil }
	if _, _, err := MinimiseLBFGS(f2, nil, x0, LBFGSOptions{}); err == nil {
		t.Fatal("expected an error for an Inf objective at the start")
	}
}

// TestLBFGSConvergesOnTheLastIteration pins that a tolerance
// met exactly on the final permitted iteration is a success, not a
// budget error.
func TestLBFGSConvergesOnTheLastIteration(t *testing.T) {
	f := func(x *core.Array) (float64, error) {
		return (x.FloatAt(0) - 1.5) * (x.FloatAt(0) - 1.5), nil
	}
	g := func(x *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{2 * (x.FloatAt(0) - 1.5)}, 1)
	}
	x0 := mustFloats(t, []float64{0}, 1)
	out, _, err := MinimiseLBFGS(f, g, x0, LBFGSOptions{MaxIterations: 1, Tolerance: 1e-10})
	if err != nil {
		t.Fatalf("MinimiseLBFGS on a quadratic with one exact step: %v", err)
	}
	if math.Abs(out.FloatAt(0)-1.5) > 1e-9 {
		t.Fatalf("minimum = %g, want 1.5", out.FloatAt(0))
	}
}

// TestFindRootSystemRejectsNonFiniteResidual pins that a NaN
// residual at an adopted state is an error: the infinity norm used to
// read it as a converged root with residual zero.
func TestFindRootSystemRejectsNonFiniteResidual(t *testing.T) {
	f := func(*core.Array) (*core.Array, error) {
		return mustFloats(t, []float64{math.NaN(), math.NaN()}), nil
	}
	x0 := mustFloats(t, []float64{1, 1}, 2)
	_, _, err := FindRootSystem(f, x0, RootSystemOptions{})
	if err == nil || !strings.Contains(err.Error(), "non-finite") {
		t.Fatalf("err = %v, want the non-finite refusal", err)
	}
	// One prefix only: the message carries the entry point once.
	if strings.Count(err.Error(), "FindRootSystem") != 1 {
		t.Fatalf("err = %v, want exactly one entry-point prefix", err)
	}
}

// TestFindRootSystemSurvivesOverflowingTrials pins the trial
// softening: log(x) is NaN on the negative half-line, the first
// Newton step from afar overshoots straight into it, and the damping
// halves past the rejected trial instead of dying. The start itself
// stays finite, so only the trials ever see NaN.
func TestFindRootSystemSurvivesOverflowingTrials(t *testing.T) {
	f := func(x *core.Array) (*core.Array, error) {
		return mustFloats(t, []float64{math.Log(x.FloatAt(0)) - 1}), nil
	}
	x0 := mustFloats(t, []float64{1000}, 1)
	sol, _, err := FindRootSystem(f, x0, RootSystemOptions{})
	if err != nil {
		t.Fatalf("FindRootSystem through NaN trials: %v", err)
	}
	if math.Abs(sol.FloatAt(0)-math.E) > 1e-8 {
		t.Fatalf("root = %.12g, want e", sol.FloatAt(0))
	}
}

// TestFindRootSystemVanishedStepIsNotARoot pins that a step
// collapsing to zero under a numerically zero Jacobian is a stall the
// iteration reports, never a root published beside a live residual.
func TestFindRootSystemVanishedStepIsNotARoot(t *testing.T) {
	f := func(x *core.Array) (*core.Array, error) {
		return mustFloats(t, []float64{1e-320*x.FloatAt(0)*x.FloatAt(0) - 10}), nil
	}
	x0 := mustFloats(t, []float64{0}, 1)
	sol, res, err := FindRootSystem(f, x0, RootSystemOptions{})
	if err == nil {
		t.Fatalf("a vanished step published root %g with residual %g", sol.FloatAt(0), res)
	}
}

// TestFindRootSystemConvergesOnTheLastIteration pins the
// exact-final-step convergence of the root solver.
func TestFindRootSystemConvergesOnTheLastIteration(t *testing.T) {
	f := func(x *core.Array) (*core.Array, error) {
		return mustFloats(t, []float64{x.FloatAt(0) - 2, x.FloatAt(1) + 1}), nil
	}
	x0 := mustFloats(t, []float64{0, 0}, 2)
	out, _, err := FindRootSystem(f, x0, RootSystemOptions{MaxIterations: 1, Tolerance: 1e-10})
	if err != nil {
		t.Fatalf("FindRootSystem on a linear system with one exact step: %v", err)
	}
	if out.FloatAt(0) != 2 || out.FloatAt(1) != -1 {
		t.Fatalf("root = (%g, %g), want (2, -1)", out.FloatAt(0), out.FloatAt(1))
	}
}

// TestLevenbergMarquardtPerfectStart pins that a start whose
// residual cancels exactly is the answer, not the damping collapse the
// strict improvement test used to die in.
func TestLevenbergMarquardtPerfectStart(t *testing.T) {
	f := func(p *core.Array) (*core.Array, error) {
		return mustFloats(t, []float64{p.FloatAt(0) - 2}), nil
	}
	p0 := mustFloats(t, []float64{2}, 1)
	out, chi2, err := LevenbergMarquardt(f, p0, LMOptions{})
	if err != nil {
		t.Fatalf("LevenbergMarquardt at the exact solution: %v", err)
	}
	if chi2 != 0 || out.FloatAt(0) != 2 {
		t.Fatalf("fit = (%g, %g), want (2, 0)", out.FloatAt(0), chi2)
	}
}

// TestMinimiseConvergesOnTheLastIteration exercises the simplex
// re-test after the final move; the convergence here lands well inside
// the budget, so the case is a smoke of the returned pair, not a pin of
// the exact final move (a deterministic last-move budget is not stable
// across the portable and vector builds).
func TestMinimiseConvergesOnTheLastIteration(t *testing.T) {
	f := func(x *core.Array) (float64, error) {
		return math.Abs(x.FloatAt(0)-1) + 0.5*math.Abs(x.FloatAt(1)+1), nil
	}
	x0 := mustFloats(t, []float64{0, 0}, 2)
	out, _, err := Minimise(f, x0, MinimiseOptions{MaxIterations: 200, Tolerance: 1e-8})
	if err != nil {
		t.Fatalf("Minimise: %v", err)
	}
	if math.Abs(out.FloatAt(0)-1) > 1e-4 || math.Abs(out.FloatAt(1)+1) > 1e-4 {
		t.Fatalf("minimum = (%g, %g), want (1, -1)", out.FloatAt(0), out.FloatAt(1))
	}
}