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

package optim

import (
	"math"
	"testing"

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

// TestLBFGSBoundedQuadratic pins a coordinate onto each wall kind: the
// optimum of the separable bowl sits at (1, 1), the box drags the
// first coordinate to its lower wall and leaves the second free, which
// is the shape every constrained fit with physical parameter ranges
// takes.
func TestLBFGSBoundedQuadratic(t *testing.T) {
	f := func(p *core.Array) (float64, error) {
		total := 0.0
		for i := range p.Len() {
			d := p.FloatAt(i) - 1
			total += d * d
		}
		return total, nil
	}
	start, _ := core.FromFloats([]float64{0, 0}, 2)
	point, value, err := MinimiseLBFGS(f, nil, start, LBFGSOptions{
		Lower: []float64{2, math.Inf(-1)},
		Upper: []float64{math.Inf(1), 9},
	})
	if err != nil {
		t.Fatalf("MinimiseLBFGS: %v", err)
	}
	if math.Abs(point.FloatAt(0)-2) > 1e-6 {
		t.Fatalf("first coordinate = %.10g, want 2 on the lower wall", point.FloatAt(0))
	}
	if math.Abs(point.FloatAt(1)-1) > 1e-6 {
		t.Fatalf("second coordinate = %.10g, want 1 free", point.FloatAt(1))
	}
	if value > 1+1e-6 {
		t.Fatalf("minimum value = %.10g, want 1", value)
	}
}

// TestLBFGSBoundedRosenbrock is the analytic case: with x forced past
// 1.5, the valley's unconstrained neck at (1, 1) is infeasible and the
// constrained minimum sits exactly on the wall at (1.5, 2.25) with
// value 0.25. The wall coordinate's gradient pushes outward, which is
// the KKT signature the optimiser must respect rather than project it
// away.
func TestLBFGSBoundedRosenbrock(t *testing.T) {
	rosenbrock := func(p *core.Array) (float64, error) {
		x, y := p.FloatAt(0), p.FloatAt(1)
		return (1-x)*(1-x) + 100*(y-x*x)*(y-x*x), nil
	}
	gradFn := func(p *core.Array) (*core.Array, error) {
		x, y := p.FloatAt(0), p.FloatAt(1)
		out := core.New(core.Float, 2)
		out.RawFloats()[0] = -2*(1-x) - 400*x*(y-x*x)
		out.RawFloats()[1] = 200 * (y - x*x)
		return out, nil
	}
	start, _ := core.FromFloats([]float64{-1.2, 1}, 2)
	point, value, err := MinimiseLBFGS(rosenbrock, gradFn, start, LBFGSOptions{
		Lower: []float64{1.5, math.Inf(-1)},
	})
	if err != nil {
		t.Fatalf("MinimiseLBFGS: %v", err)
	}
	if math.Abs(point.FloatAt(0)-1.5) > 1e-6 {
		t.Fatalf("first coordinate = %.10g, want 1.5 on the wall", point.FloatAt(0))
	}
	if math.Abs(point.FloatAt(1)-2.25) > 1e-4 {
		t.Fatalf("second coordinate = %.10g, want 2.25", point.FloatAt(1))
	}
	if math.Abs(value-0.25) > 1e-6 {
		t.Fatalf("value = %.10g, want 0.25", value)
	}
}

// TestLBFGSBoundedDomain proves the finite-difference gradient turns
// one-sided at a wall: log is undefined below the wall, so a central
// stencil would make the objective return an error and the run would
// fail outright.
func TestLBFGSBoundedDomain(t *testing.T) {
	f := func(p *core.Array) (float64, error) {
		x := p.FloatAt(0)
		if x < 0.5 {
			return 0, base.Errf("the objective is undefined below 0.5")
		}
		return math.Log(x), nil
	}
	start, _ := core.FromFloats([]float64{2}, 1)
	point, _, err := MinimiseLBFGS(f, nil, start, LBFGSOptions{
		Lower: []float64{0.5},
	})
	if err != nil {
		t.Fatalf("MinimiseLBFGS: %v", err)
	}
	if math.Abs(point.FloatAt(0)-0.5) > 1e-6 {
		t.Fatalf("coordinate = %.10g, want 0.5 on the wall", point.FloatAt(0))
	}
}

// TestLBFGSBoundProjection checks that an infeasible start is
// projected onto the box and still converges, and that hostile bounds
// are refused rather than silently swapped or clamped.
func TestLBFGSBoundProjection(t *testing.T) {
	f := func(p *core.Array) (float64, error) {
		d := p.FloatAt(0) - 3
		return d * d, nil
	}
	start, _ := core.FromFloats([]float64{-10}, 1)
	point, _, err := MinimiseLBFGS(f, nil, start, LBFGSOptions{Lower: []float64{1}})
	if err != nil {
		t.Fatalf("MinimiseLBFGS: %v", err)
	}
	if math.Abs(point.FloatAt(0)-3) > 1e-6 {
		t.Fatalf("coordinate = %.10g, want 3", point.FloatAt(0))
	}
	bad, _ := core.FromFloats([]float64{0}, 1)
	if _, _, err := MinimiseLBFGS(f, nil, bad, LBFGSOptions{Lower: []float64{2}, Upper: []float64{1}}); err == nil {
		t.Fatal("crossed walls accepted")
	}
	if _, _, err := MinimiseLBFGS(f, nil, bad, LBFGSOptions{Lower: []float64{1, 2}}); err == nil {
		t.Fatal("bounds of the wrong length accepted")
	}
	if _, _, err := MinimiseLBFGS(f, nil, bad, LBFGSOptions{Lower: []float64{math.NaN()}}); err == nil {
		t.Fatal("NaN wall accepted")
	}
}

// TestLBFGSBoundedAgainstEvolution cross-checks the local method
// against the global one on the same box: differential evolution
// clamps its population into the bounds too, and on a convex problem
// both must land on the same value.
func TestLBFGSBoundedAgainstEvolution(t *testing.T) {
	f := func(p *core.Array) (float64, error) {
		x, y := p.FloatAt(0), p.FloatAt(1)
		return (x-2)*(x-2) + 10*(y+1)*(y+1), nil
	}
	start, _ := core.FromFloats([]float64{0, 0}, 2)
	point, value, err := MinimiseLBFGS(f, nil, start, LBFGSOptions{
		Lower: []float64{3, math.Inf(-1)},
		Upper: []float64{math.Inf(1), 0.5},
	})
	if err != nil {
		t.Fatalf("MinimiseLBFGS: %v", err)
	}
	if math.Abs(point.FloatAt(0)-3) > 1e-6 || math.Abs(point.FloatAt(1)+1) > 1e-6 {
		t.Fatalf("minimiser = (%.10g, %.10g), want (3, -1)", point.FloatAt(0), point.FloatAt(1))
	}
	lower, _ := core.FromFloats([]float64{3, -1}, 2)
	upper, _ := core.FromFloats([]float64{10, 5}, 2)
	_, evoValue, err := MinimiseDifferentialEvolution(f, lower, upper,
		DifferentialEvolutionOptions{Seed: 7, Generations: 200})
	if err != nil {
		t.Fatalf("MinimiseDifferentialEvolution: %v", err)
	}
	if value > evoValue+1e-6 || evoValue > value+1e-4 {
		t.Fatalf("L-BFGS value %.10g and evolution value %.10g disagree", value, evoValue)
	}
}