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

// TestDESphere pins the global minimum of the sphere from a wide box.
func TestDESphere(t *testing.T) {
	lower, _ := core.FromFloats([]float64{-10, -10, -10}, 3)
	upper, _ := core.FromFloats([]float64{10, 10, 10}, 3)
	x, fv, err := MinimiseDifferentialEvolution(func(a *core.Array) (float64, error) {
		s := 0.0
		for i := range 3 {
			s += a.FloatAt(i) * a.FloatAt(i)
		}
		return s, nil
	}, lower, upper, DifferentialEvolutionOptions{Generations: 300})
	if err != nil {
		t.Fatalf("MinimiseDifferentialEvolution: %v", err)
	}
	if fv > 1e-10 {
		t.Fatalf("sphere minimum = %.3e, want 0", fv)
	}
	for i := range 3 {
		if math.Abs(x.FloatAt(i)) > 1e-5 {
			t.Fatalf("x[%d] = %g, want 0", i, x.FloatAt(i))
		}
	}
}

// TestDERosenbrockRastrigin pins two multimodal classics: Rosenbrock's
// valley and Rastrigin's minefield of local minima.
func TestDERosenbrockRastrigin(t *testing.T) {
	lower, _ := core.FromFloats([]float64{-2.5, -2.5}, 2)
	upper, _ := core.FromFloats([]float64{2.5, 2.5}, 2)
	_, fvR, err := MinimiseDifferentialEvolution(func(a *core.Array) (float64, error) {
		x, y := a.FloatAt(0), a.FloatAt(1)
		return (1-x)*(1-x) + 100*(y-x*x)*(y-x*x), nil
	}, lower, upper, DifferentialEvolutionOptions{Generations: 600})
	if err != nil {
		t.Fatalf("Rosenbrock: %v", err)
	}
	if fvR > 1e-6 {
		t.Fatalf("Rosenbrock minimum = %.3e, want 0", fvR)
	}

	lowerR, _ := core.FromFloats([]float64{-5.12, -5.12, -5.12}, 3)
	upperR, _ := core.FromFloats([]float64{5.12, 5.12, 5.12}, 3)
	x, fv, err := MinimiseDifferentialEvolution(func(a *core.Array) (float64, error) {
		s := 0.0
		for i := range 3 {
			z := a.FloatAt(i)
			s += z*z - 10*math.Cos(2*math.Pi*z) + 10
		}
		return s, nil
	}, lowerR, upperR, DifferentialEvolutionOptions{Generations: 800})
	if err != nil {
		t.Fatalf("Rastrigin: %v", err)
	}
	if fv > 1e-6 {
		t.Fatalf("Rastrigin minimum = %.3e, want 0 (global basin found at %v)", fv, x)
	}
}

// TestDEBoundClamping pins that the search stays inside the box.
func TestDEBoundClamping(t *testing.T) {
	lower, _ := core.FromFloats([]float64{-1, -1}, 2)
	upper, _ := core.FromFloats([]float64{1, 1}, 2)
	x, _, err := MinimiseDifferentialEvolution(func(a *core.Array) (float64, error) {
		// Push toward a corner outside the box.
		return -a.FloatAt(0) - 2*a.FloatAt(1), nil
	}, lower, upper, DifferentialEvolutionOptions{Generations: 200})
	if err != nil {
		t.Fatalf("MinimiseDifferentialEvolution: %v", err)
	}
	for i := range 2 {
		if x.FloatAt(i) < lower.FloatAt(i)-1e-12 || x.FloatAt(i) > upper.FloatAt(i)+1e-12 {
			t.Fatalf("x[%d] = %g left the box", i, x.FloatAt(i))
		}
	}
}

// mustBound builds one bound vector.
func mustBound(vs ...float64) *core.Array {
	a, _ := core.FromFloats(vs, len(vs))
	return a
}

// TestDEErrors pins the input gates.
func TestDEErrors(t *testing.T) {
	lower, _ := core.FromFloats([]float64{1, 1}, 2)
	upper, _ := core.FromFloats([]float64{0, 2}, 2)
	if _, _, err := MinimiseDifferentialEvolution(func(a *core.Array) (float64, error) { return 0, nil },
		lower, upper, DifferentialEvolutionOptions{}); err == nil {
		t.Error("degenerate bound accepted")
	}
	lo, _ := core.FromFloats([]float64{0, 0}, 2)
	hi, _ := core.FromFloats([]float64{1, 1, 1}, 3)
	if _, _, err := MinimiseDifferentialEvolution(func(a *core.Array) (float64, error) { return 0, nil },
		lo, hi, DifferentialEvolutionOptions{}); err == nil {
		t.Error("mismatched bounds accepted")
	}
	if _, _, err := MinimiseDifferentialEvolution(func(a *core.Array) (float64, error) {
		return math.NaN(), nil
	}, lo, mustBound(1, 1), DifferentialEvolutionOptions{}); err == nil {
		t.Error("NaN objective accepted")
	}
}