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

// TestNonlinearEqualityCircle pins the hand-solved circle case: min x
// subject to x² + y² = 1. The constrained minimum is (−1, 0) with
// value −1, and the KKT stationarity ∇f + λ∇g = 0 there reads
// (1, 0) + λ(−2, 0) = 0, so the multiplier converges to the analytic
// 0.5. The row's own gradient is differentiated, so the tolerance on
// the multiplier is the finite-difference one.
func TestNonlinearEqualityCircle(t *testing.T) {
	cons := NonlinearConstraints{
		Equalities: []func(*core.Array) (float64, error){
			func(p *core.Array) (float64, error) {
				x, y := p.FloatAt(0), p.FloatAt(1)
				return x*x + y*y - 1, nil
			},
		},
	}
	f := func(p *core.Array) (float64, error) { return p.FloatAt(0), nil }
	x, value, multipliers, err := MinimiseNonlinearConstrained(f, nil, mustFloats(t, []float64{-2, 0.5}), cons,
		LBFGSOptions{Tolerance: 1e-10})
	if err != nil {
		t.Fatalf("MinimiseNonlinearConstrained: %v", err)
	}
	if math.Abs(x.FloatAt(0)+1) > 1e-3 || math.Abs(x.FloatAt(1)) > 1e-3 {
		t.Fatalf("point = (%.10g, %.10g), want (−1, 0)", x.FloatAt(0), x.FloatAt(1))
	}
	if math.Abs(value+1) > 1e-4 {
		t.Fatalf("value = %.10g, want −1", value)
	}
	if len(multipliers) != 1 {
		t.Fatalf("multipliers = %v, want one entry for the equality row", multipliers)
	}
	if math.Abs(multipliers[0]-0.5) > 1e-3 {
		t.Fatalf("multiplier = %.10g, want the analytic 0.5", multipliers[0])
	}
}

// TestNonlinearInequalityCircle pins the active inequality: min
// −(x + y) subject to x² + y² ≤ 1. The unconstrained minimum runs to
// infinity, the constrained one sits on the circle at (1/√2, 1/√2)
// with value −√2, and the stationarity (−1, −1) + μ(√2, √2) = 0 fixes
// the multiplier at 1/√2.
func TestNonlinearInequalityCircle(t *testing.T) {
	cons := NonlinearConstraints{
		Inequalities: []func(*core.Array) (float64, error){
			func(p *core.Array) (float64, error) {
				x, y := p.FloatAt(0), p.FloatAt(1)
				return x*x + y*y - 1, nil
			},
		},
	}
	f := func(p *core.Array) (float64, error) { return -(p.FloatAt(0) + p.FloatAt(1)), nil }
	x, value, multipliers, err := MinimiseNonlinearConstrained(f, nil, mustFloats(t, []float64{0.5, 0.5}), cons,
		LBFGSOptions{Tolerance: 1e-10})
	if err != nil {
		t.Fatalf("MinimiseNonlinearConstrained: %v", err)
	}
	root := 1 / math.Sqrt2
	if math.Abs(x.FloatAt(0)-root) > 1e-3 || math.Abs(x.FloatAt(1)-root) > 1e-3 {
		t.Fatalf("point = (%.10g, %.10g), want (%g, %g)", x.FloatAt(0), x.FloatAt(1), root, root)
	}
	if math.Abs(value+math.Sqrt2) > 1e-4 {
		t.Fatalf("value = %.10g, want −√2", value)
	}
	if len(multipliers) != 1 || math.Abs(multipliers[0]-math.Sqrt2/2) > 5e-3 {
		t.Fatalf("multipliers = %v, want [1/√2] within the finite-difference tolerance", multipliers)
	}
	if multipliers[0] < 0 {
		t.Fatalf("multiplier = %.10g, non-negative on an active inequality", multipliers[0])
	}
}

// TestNonlinearInequalitySlackMultiplier pins the complementary
// slackness the returned slice carries: a row that sits strictly slack
// at the answer has KKT multiplier exactly zero, not whatever the
// rounds its row was violated in left behind. The quartic's minimum
// over x ≤ 0.5 from x₀ = 3 is the unconstrained x = −2, the row deep
// in its slack.
func TestNonlinearInequalitySlackMultiplier(t *testing.T) {
	f := func(p *core.Array) (float64, error) {
		v := p.FloatAt(0)
		return 100 * (v + 2) * (v + 2) * (v - 3) * (v - 3), nil
	}
	cons := NonlinearConstraints{
		Inequalities: []func(*core.Array) (float64, error){
			func(p *core.Array) (float64, error) { return p.FloatAt(0) - 0.5, nil },
		},
	}
	x, _, multipliers, err := MinimiseNonlinearConstrained(f, nil, mustFloats(t, []float64{3}), cons,
		LBFGSOptions{})
	if err != nil {
		t.Fatalf("MinimiseNonlinearConstrained: %v", err)
	}
	if math.Abs(x.FloatAt(0)+2) > 1e-3 {
		t.Fatalf("x = %.10g, want the unconstrained −2", x.FloatAt(0))
	}
	if len(multipliers) != 1 || multipliers[0] != 0 {
		t.Fatalf("multipliers = %v, want the slack row's KKT zero", multipliers)
	}
}

// TestNonlinearLinearRowComposition composes an affine row as a
// function: min (x−1)² + (y−1)² subject to x + y − 2 = 0. The plane
// passes through the unconstrained minimum, so the answer is (1, 1)
// with value 0 and the multiplier converging to zero, the linear
// entry's TestConstrainedEquality through the nonlinear door.
func TestNonlinearLinearRowComposition(t *testing.T) {
	cons := NonlinearConstraints{
		Equalities: []func(*core.Array) (float64, error){
			func(p *core.Array) (float64, error) {
				return p.FloatAt(0) + p.FloatAt(1) - 2, nil
			},
		},
	}
	f := constrainedBowl(1, 1)
	// The bowl's gradient, supplied in one of the two runs so both the
	// chained-gradient path and the slope-zero skip are exercised.
	grad := func(p *core.Array) (*core.Array, error) {
		out := core.New(core.Float, 2)
		out.RawFloats()[0] = 2 * (p.FloatAt(0) - 1)
		out.RawFloats()[1] = 2 * (p.FloatAt(1) - 1)
		return out, nil
	}
	for _, grad := range []func(*core.Array) (*core.Array, error){nil, grad} {
		x, value, multipliers, err := MinimiseNonlinearConstrained(f, grad, mustFloats(t, []float64{5, -3}), cons,
			LBFGSOptions{Tolerance: 1e-10})
		if err != nil {
			t.Fatalf("MinimiseNonlinearConstrained: %v", err)
		}
		if math.Abs(x.FloatAt(0)-1) > 1e-4 || math.Abs(x.FloatAt(1)-1) > 1e-4 {
			t.Fatalf("point = (%.8g, %.8g), want (1, 1)", x.FloatAt(0), x.FloatAt(1))
		}
		if math.Abs(value) > 1e-6 {
			t.Fatalf("value = %.8g, want 0", value)
		}
		if len(multipliers) != 1 || math.Abs(multipliers[0]) > 1e-3 {
			t.Fatalf("multipliers = %v, want [≈0]", multipliers)
		}
	}
}

// TestNonlinearStalledInnerRefused pins the honest refusal of a
// stalled inner solve: min 10⁶x² + y² subject to xy = 2 with one
// L-BFGS iteration per round crawls toward the hyperbola, the outer
// feasibility check stays unsatisfied through the 40 rounds, and the
// run is refused with the remaining violation, never returned as a
// solution.
func TestNonlinearStalledInnerRefused(t *testing.T) {
	cons := NonlinearConstraints{
		Equalities: []func(*core.Array) (float64, error){
			func(p *core.Array) (float64, error) {
				return p.FloatAt(0)*p.FloatAt(1) - 2, nil
			},
		},
	}
	f := func(p *core.Array) (float64, error) {
		x, y := p.FloatAt(0), p.FloatAt(1)
		return 1e6*x*x + y*y, nil
	}
	_, _, _, err := MinimiseNonlinearConstrained(f, nil, mustFloats(t, []float64{5, 5}), cons,
		LBFGSOptions{MaxIterations: 1})
	if err == nil {
		t.Fatal("a stalled inner solve was returned as a solution")
	}
	if !strings.Contains(err.Error(), "violation") {
		t.Fatalf("error = %v, want the outer feasibility refusal", err)
	}
}

// TestNonlinearDelegatesUnconstrained pins the empty-set composition:
// no rows at all is MinimiseLBFGS with nil multipliers.
func TestNonlinearDelegatesUnconstrained(t *testing.T) {
	x, value, multipliers, err := MinimiseNonlinearConstrained(constrainedBowl(2, -1), nil,
		mustFloats(t, []float64{0, 0}), NonlinearConstraints{}, LBFGSOptions{})
	if err != nil {
		t.Fatalf("MinimiseNonlinearConstrained: %v", err)
	}
	if math.Abs(x.FloatAt(0)-2) > 1e-4 || math.Abs(x.FloatAt(1)+1) > 1e-4 {
		t.Fatalf("point = (%.8g, %.8g), want (2, −1)", x.FloatAt(0), x.FloatAt(1))
	}
	if value > 1e-8 {
		t.Fatalf("value = %.8g, want ≈ 0", value)
	}
	if multipliers != nil {
		t.Fatalf("multipliers = %v, want nil", multipliers)
	}
}

// TestNonlinearWithAnalyticGradient runs the same hand-solved circle
// case with f's gradient supplied: the row terms are then chained onto
// it analytically, the row's own gradient differentiated at each
// measured point, and the answer must match the finite-difference run.
func TestNonlinearWithAnalyticGradient(t *testing.T) {
	cons := NonlinearConstraints{
		Equalities: []func(*core.Array) (float64, error){
			func(p *core.Array) (float64, error) {
				x, y := p.FloatAt(0), p.FloatAt(1)
				return x*x + y*y - 1, nil
			},
		},
	}
	f := func(p *core.Array) (float64, error) { return p.FloatAt(0), nil }
	grad := func(p *core.Array) (*core.Array, error) {
		out := core.New(core.Float, 2)
		out.RawFloats()[0] = 1
		return out, nil
	}
	x, value, multipliers, err := MinimiseNonlinearConstrained(f, grad, mustFloats(t, []float64{-2, 0.5}), cons,
		LBFGSOptions{Tolerance: 1e-10})
	if err != nil {
		t.Fatalf("MinimiseNonlinearConstrained: %v", err)
	}
	if math.Abs(x.FloatAt(0)+1) > 1e-3 || math.Abs(x.FloatAt(1)) > 1e-3 {
		t.Fatalf("point = (%.10g, %.10g), want (−1, 0)", x.FloatAt(0), x.FloatAt(1))
	}
	if math.Abs(value+1) > 1e-4 {
		t.Fatalf("value = %.10g, want −1", value)
	}
	if len(multipliers) != 1 || math.Abs(multipliers[0]-0.5) > 1e-3 {
		t.Fatalf("multipliers = %v, want [0.5]", multipliers)
	}
	// A gradient of the wrong length is refused.
	if _, _, _, err := MinimiseNonlinearConstrained(f, func(*core.Array) (*core.Array, error) {
		return core.New(core.Float, 3), nil
	}, mustFloats(t, []float64{-2, 0.5}), cons, LBFGSOptions{}); err == nil {
		t.Fatal("a wrong-length gradient was accepted")
	}
	// A complex gradient payload is refused.
	if _, _, _, err := MinimiseNonlinearConstrained(f, func(*core.Array) (*core.Array, error) {
		return mustComplexPoint(t), nil
	}, mustFloats(t, []float64{-2, 0.5}), cons, LBFGSOptions{}); err == nil {
		t.Fatal("a complex gradient was accepted")
	}
	// A gradient callback's own error propagates.
	if _, _, _, err := MinimiseNonlinearConstrained(f, func(*core.Array) (*core.Array, error) {
		return nil, base.Errf("the gradient exploded")
	}, mustFloats(t, []float64{-2, 0.5}), cons, LBFGSOptions{}); err == nil {
		t.Fatal("the gradient's error did not propagate")
	}
	// The objective's own error propagates out of the inner solve.
	if _, _, _, err := MinimiseNonlinearConstrained(func(*core.Array) (float64, error) {
		return 0, base.Errf("the objective exploded")
	}, nil, mustFloats(t, []float64{-2, 0.5}), cons, LBFGSOptions{}); err == nil {
		t.Fatal("the objective's error did not propagate")
	}
}

// TestNonlinearRefusals checks the loud rejections: nil functions,
// empty and complex starts, a non-finite row value and an objective
// error.
func TestNonlinearRefusals(t *testing.T) {
	good := NonlinearConstraints{
		Equalities: []func(*core.Array) (float64, error){
			func(p *core.Array) (float64, error) { return p.FloatAt(0) - 1, nil },
		},
	}
	f := constrainedBowl(0, 0)
	start := mustFloats(t, []float64{0, 0})
	// Nil equality function.
	nilEq := NonlinearConstraints{Equalities: []func(*core.Array) (float64, error){nil}}
	if _, _, _, err := MinimiseNonlinearConstrained(f, nil, start, nilEq, LBFGSOptions{}); err == nil {
		t.Fatal("a nil equality function was accepted")
	}
	nilIn := NonlinearConstraints{Inequalities: []func(*core.Array) (float64, error){nil}}
	if _, _, _, err := MinimiseNonlinearConstrained(f, nil, start, nilIn, LBFGSOptions{}); err == nil {
		t.Fatal("a nil inequality function was accepted")
	}
	// Empty start.
	if _, _, _, err := MinimiseNonlinearConstrained(f, nil, core.New(core.Float, 0), good, LBFGSOptions{}); err == nil {
		t.Fatal("an empty starting point was accepted")
	}
	// Complex start.
	if _, _, _, err := MinimiseNonlinearConstrained(f, nil, mustComplexPoint(t), good, LBFGSOptions{}); err == nil {
		t.Fatal("a complex starting point was accepted")
	}
	// A row function that returns NaN is fatal.
	nan := NonlinearConstraints{
		Equalities: []func(*core.Array) (float64, error){
			func(*core.Array) (float64, error) { return math.NaN(), nil },
		},
	}
	if _, _, _, err := MinimiseNonlinearConstrained(f, nil, start, nan, LBFGSOptions{}); err == nil {
		t.Fatal("a NaN row value was accepted")
	}
	// A row function's own error propagates.
	failing := NonlinearConstraints{
		Inequalities: []func(*core.Array) (float64, error){
			func(*core.Array) (float64, error) { return 0, base.Errf("the row exploded") },
		},
	}
	if _, _, _, err := MinimiseNonlinearConstrained(f, nil, start, failing, LBFGSOptions{}); err == nil {
		t.Fatal("the row function's error did not propagate")
	}
}