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

package optim

import (
	"sourcedock.dev/petrbalvin/tensor/internal/core"
	"strings"
	"testing"
)

// TestLBFGSBudgetExitIsNotConvergence pins the last silent exit of the
// optimiser. A run that spends its iteration budget has not converged:
// on a stiff objective it stops with a projected gradient orders of
// magnitude above the tolerance, and reporting that point as the answer
// is the same silent wrongness the stall and direction exits refuse.
// AllowBudgetExit is the documented escape hatch, and it is what the
// augmented Lagrangian's inexact inner solves use.
func TestLBFGSBudgetExitIsNotConvergence(t *testing.T) {
	// Mixed units: the y direction is 1e10 times stiffer, so five
	// iterations cannot reach the default 1e-8 tolerance.
	stiff := func(p *core.Array) (float64, error) {
		dx := p.FloatAt(0) - 3
		dy := p.FloatAt(1) - 5
		return dx*dx + 1e10*dy*dy, nil
	}
	start := mustFloats(t, []float64{0, 0}, 2)

	point, value, err := MinimiseLBFGS(stiff, nil, start, LBFGSOptions{MaxIterations: 5})
	if err == nil {
		t.Fatalf("a budget stop was reported as convergence: point = %v, value = %g", floatsOf(point), value)
	}
	if !strings.Contains(err.Error(), "iteration budget") {
		t.Fatalf("error = %v, want the iteration-budget refusal", err)
	}
	if point != nil {
		t.Fatalf("the refused run returned the point %v", floatsOf(point))
	}

	// The escape hatch is what the constrained wrapper relies on: the
	// point comes back with no error, and the outer loop's feasibility
	// check is what judges it.
	point, value, err = MinimiseLBFGS(stiff, nil, start, LBFGSOptions{MaxIterations: 5, AllowBudgetExit: true})
	if err != nil {
		t.Fatalf("AllowBudgetExit: %v", err)
	}
	if point == nil {
		t.Fatal("AllowBudgetExit returned no point")
	}
	t.Logf("best effort after 5 iterations: %v, value %g", floatsOf(point), value)

	// A run that does converge is unaffected by either setting.
	for _, allow := range []bool{false, true} {
		point, _, err = MinimiseLBFGS(func(p *core.Array) (float64, error) {
			dx := p.FloatAt(0) - 3
			return dx * dx, nil
		}, nil, mustFloats(t, []float64{0}, 1), LBFGSOptions{AllowBudgetExit: allow})
		if err != nil {
			t.Fatalf("AllowBudgetExit=%v on a converging run: %v", allow, err)
		}
		if math := point.FloatAt(0); math < 2.999999999 || math > 3.000000001 {
			t.Fatalf("AllowBudgetExit=%v: point %g, want 3", allow, math)
		}
	}
}

// TestMinimiseAndLMBudgetExitIsNotConvergence pins the same refusal on
// Minimise and LevenbergMarquardt: a run that spends its iteration
// budget must not publish its last point as a converged answer, and
// AllowBudgetExit is the documented escape hatch.
func TestMinimiseAndLMBudgetExitIsNotConvergence(t *testing.T) {
	stiff := func(p *core.Array) (float64, error) {
		dx := p.FloatAt(0) - 3
		dy := p.FloatAt(1) - 5
		return dx*dx + 1e10*dy*dy, nil
	}
	start := mustFloats(t, []float64{0, 0}, 2)

	point, _, err := Minimise(stiff, start, MinimiseOptions{MaxIterations: 5})
	if err == nil || !strings.Contains(err.Error(), "iteration budget") {
		t.Fatalf("Minimise: a budget stop was reported as convergence (err = %v)", err)
	}
	if point != nil {
		t.Fatal("Minimise: the refused run returned a point")
	}
	point, _, err = Minimise(stiff, start, MinimiseOptions{MaxIterations: 5, AllowBudgetExit: true})
	if err != nil || point == nil {
		t.Fatalf("Minimise AllowBudgetExit: err = %v, point = %v", err, point)
	}

	residual := func(p *core.Array) (*core.Array, error) {
		r := core.New(core.Float, 2)
		r.RawFloats()[0] = p.FloatAt(0) - 3
		r.RawFloats()[1] = p.FloatAt(1) - 5
		return r, nil
	}
	fp, _, err := LevenbergMarquardt(residual, start, LMOptions{MaxIterations: 1})
	if err == nil || !strings.Contains(err.Error(), "iteration budget") {
		t.Fatalf("LevenbergMarquardt: a budget stop was reported as convergence (err = %v)", err)
	}
	if fp != nil {
		t.Fatal("LevenbergMarquardt: the refused run returned a point")
	}
	fp, _, err = LevenbergMarquardt(residual, start, LMOptions{MaxIterations: 1, AllowBudgetExit: true})
	if err != nil || fp == nil {
		t.Fatalf("LevenbergMarquardt AllowBudgetExit: err = %v, point = %v", err, fp)
	}
}