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

// rotEllipsoid builds the test landscape: f(x) = Σ i·(Rx)ᵢ² for a
// fixed orthonormal R (the Householder reflection through the
// normalised all-ones vector), a rotated ellipsoid with condition 6ⁿ
// along axes the coordinate system does not see. It returns the
// objective and the apply function so tests can check coordinates.
func rotEllipsoid(t *testing.T, n int) (func(*core.Array) (float64, error), func(x []float64) []float64) {
	t.Helper()
	v := make([]float64, n)
	for i := range n {
		v[i] = float64(i + 1)
	}
	norm := 0.0
	for _, a := range v {
		norm += a * a
	}
	norm = math.Sqrt(norm)
	apply := func(x []float64) []float64 {
		dot := 0.0
		for i := range n {
			dot += v[i] * x[i]
		}
		dot *= 2 / (norm * norm)
		rx := make([]float64, n)
		for i := range n {
			rx[i] = x[i] - dot*v[i]
		}
		return rx
	}
	f := func(a *core.Array) (float64, error) {
		total := 0.0
		for i, rx := range apply(floatsOf(a)) {
			total += float64(i+1) * rx * rx
		}
		return total, nil
	}
	return f, apply
}

// TestJacobiEigen pins the eigendecomposition the CMA-ES loop leans
// on: a known symmetric matrix with irrational eigenvalues, verified
// through the reconstruction C = B·D·Bᵀ and the orthogonality of B.
func TestJacobiEigen(t *testing.T) {
	// [[2, 1], [1, 3]]: eigenvalues (5 ± √5)/2.
	vals, vecs := jacobiEigen([]float64{2, 1, 1, 3}, 2)
	want1 := (5 + math.Sqrt(5)) / 2
	want2 := (5 - math.Sqrt(5)) / 2
	if math.Abs(vals[0]-want1) > 1e-12 || math.Abs(vals[1]-want2) > 1e-12 {
		t.Fatalf("eigenvalues = (%g, %g), want (%g, %g)", vals[0], vals[1], want1, want2)
	}
	// Orthonormality: B·Bᵀ = I.
	for i := range 2 {
		for j := range 2 {
			s := 0.0
			for k := range 2 {
				s += vecs[k*2+i] * vecs[k*2+j]
			}
			want := 0.0
			if i == j {
				want = 1
			}
			if math.Abs(s-want) > 1e-12 {
				t.Fatalf("B·Bᵀ[%d][%d] = %g, want %g", i, j, s, want)
			}
		}
	}
	// Reconstruction: Σ_j vals[j]·vec_j·vec_jᵀ = C.
	for i := range 2 {
		for j := range 2 {
			s := 0.0
			for k := range 2 {
				s += vals[k] * vecs[k*2+i] * vecs[k*2+j]
			}
			want := 1.0
			if i == j {
				want = float64(2 + i)
			}
			if math.Abs(s-want) > 1e-12 {
				t.Fatalf("reconstruction[%d][%d] = %g, want %g", i, j, s, want)
			}
		}
	}
}

// TestMinimiseCMAESRotatedEllipsoid pins the strategy on a rotated
// ellipsoid, the landscape its covariance adaptation exists for: seed
// 7, sigma0 0.5 and a budget of 300 generations carry the run to the
// origin within 1e−6 per coordinate and 1e−12 in value.
func TestMinimiseCMAESRotatedEllipsoid(t *testing.T) {
	f, _ := rotEllipsoid(t, 6)
	start := mustFloats(t, []float64{1, -1, 0.5, 2, 0, -0.5})
	x, fv, err := MinimiseCMAES(f, start, CMAESOptions{Seed: 7, Sigma0: 0.5, Generations: 300})
	if err != nil {
		t.Fatalf("MinimiseCMAES: %v", err)
	}
	if fv > 1e-12 {
		t.Fatalf("value = %.3e, want <= 1e-12", fv)
	}
	for i := range 6 {
		if math.Abs(x.FloatAt(i)) > 1e-5 {
			t.Fatalf("x[%d] = %.3e, want within 1e-5 of the origin", i, x.FloatAt(i))
		}
	}
}

// TestMinimiseCMAESDeterministic pins reproducibility: two runs on one
// seed walk the same landscape with the same draws and must return
// bit-identical trajectories.
func TestMinimiseCMAESDeterministic(t *testing.T) {
	f, _ := rotEllipsoid(t, 4)
	start := mustFloats(t, []float64{0.5, -0.5, 1, -1})
	// Bit-identical trajectories under one seed: 80 generations do not
	// exhaust to convergence, so the runs use the escape hatch and the
	// pin is on the trajectories themselves.
	xa, fa, errA := MinimiseCMAES(f, start, CMAESOptions{Seed: 7, Sigma0: 0.5, Generations: 80, AllowBudgetExit: true})
	xb, fb, errB := MinimiseCMAES(f, start, CMAESOptions{Seed: 7, Sigma0: 0.5, Generations: 80, AllowBudgetExit: true})
	if errA != nil || errB != nil {
		t.Fatalf("MinimiseCMAES: %v, %v", errA, errB)
	}
	if math.Float64bits(fa) != math.Float64bits(fb) {
		t.Fatalf("values differ across identical runs: %.20g vs %.20g", fa, fb)
	}
	for i := range 4 {
		if math.Float64bits(xa.FloatAt(i)) != math.Float64bits(xb.FloatAt(i)) {
			t.Fatalf("x[%d] differs across identical runs", i)
		}
	}
}

// TestMinimiseCMAESBudgetAndGates pins the honest budget refusal, the
// escape hatch, the divergence refusal on an unbounded-below
// objective, and the input and objective gates.
func TestMinimiseCMAESBudgetAndGates(t *testing.T) {
	f, _ := rotEllipsoid(t, 4)
	start := mustFloats(t, []float64{0.5, -0.5, 1, -1})
	// Two generations cannot collapse the distribution: the budget
	// stop is refused with the evidence, not reported as an answer.
	_, _, err := MinimiseCMAES(f, start, CMAESOptions{Seed: 7, Sigma0: 0.5, Generations: 2})
	if err == nil || !strings.Contains(err.Error(), "budget") {
		t.Fatalf("error = %v, want the budget refusal", err)
	}
	// The escape hatch returns the best point with no error.
	xb, _, err := MinimiseCMAES(f, start, CMAESOptions{Seed: 7, Sigma0: 0.5, Generations: 2, AllowBudgetExit: true})
	if err != nil || xb == nil {
		t.Fatalf("AllowBudgetExit: err = %v, x = %v", err, xb)
	}
	// An objective unbounded below drives sigma past the guard.
	_, _, err = MinimiseCMAES(func(a *core.Array) (float64, error) {
		s := 0.0
		for i := range a.Len() {
			s += a.FloatAt(i) * a.FloatAt(i)
		}
		return -math.Sqrt(s), nil
	}, start, CMAESOptions{Seed: 7, Sigma0: 0.5, Generations: 200})
	if err == nil || !strings.Contains(err.Error(), "diverged") {
		t.Fatalf("error = %v, want the divergence refusal", err)
	}
	// Empty and complex starts are refused.
	if _, _, err := MinimiseCMAES(f, core.New(core.Float, 0), CMAESOptions{}); err == nil {
		t.Fatal("an empty starting point was accepted")
	}
	if _, _, err := MinimiseCMAES(f, mustComplexPoint(t), CMAESOptions{}); err == nil {
		t.Fatal("a complex starting point was accepted")
	}
	// A non-finite objective is fatal.
	if _, _, err := MinimiseCMAES(func(*core.Array) (float64, error) { return math.NaN(), nil },
		start, CMAESOptions{Seed: 7, Generations: 5}); err == nil {
		t.Fatal("a NaN objective was accepted")
	}
	// An objective's own error propagates.
	if _, _, err := MinimiseCMAES(func(*core.Array) (float64, error) {
		return 0, base.Errf("the model exploded")
	}, start, CMAESOptions{Seed: 7, Generations: 5}); err == nil {
		t.Fatal("the objective's error did not propagate")
	}
	// The defaults carry the run when the caller passes nothing but
	// the escape hatch: sigma0 0.3, 500 generations, tolerance
	// 1e-12 and seed 42.
	_, _, err = MinimiseCMAES(f, start, CMAESOptions{AllowBudgetExit: true})
	if err != nil {
		t.Fatalf("MinimiseCMAES with defaults: %v", err)
	}
}

// TestMinimiseAnnealRotatedEllipsoid pins simulated annealing on the
// same rotated ellipsoid at basin precision: 30000 proposals reach the
// origin's basin, which for this landscape means every coordinate
// within 0.5 and a value below 6.
func TestMinimiseAnnealRotatedEllipsoid(t *testing.T) {
	f, _ := rotEllipsoid(t, 6)
	start := mustFloats(t, []float64{1, -1, 0.5, 2, 0, -0.5})
	x, fv, err := MinimiseSimulatedAnnealing(f, start, SimulatedAnnealingOptions{Seed: 7, Steps: 30000})
	if err != nil {
		t.Fatalf("MinimiseSimulatedAnnealing: %v", err)
	}
	if fv > 6 {
		t.Fatalf("value = %.3e, want <= 6 (the basin of the origin)", fv)
	}
	for i := range 6 {
		if math.Abs(x.FloatAt(i)) > 0.5 {
			t.Fatalf("x[%d] = %.3e, want within the unit basin", i, x.FloatAt(i))
		}
	}
}

// TestMinimiseAnnealDeterministicAndBudget pins reproducibility, the
// budget refusal of a chain that is still improving, the escape hatch
// and the input gates.
func TestMinimiseAnnealDeterministicAndBudget(t *testing.T) {
	f, _ := rotEllipsoid(t, 4)
	start := mustFloats(t, []float64{0.5, -0.5, 1, -1})
	// Bit-identical trajectories under one seed: 500 proposals end
	// mid-polish, so the runs use the escape hatch and the pin is on
	// the trajectories themselves.
	xa, fa, errA := MinimiseSimulatedAnnealing(f, start, SimulatedAnnealingOptions{Seed: 7, Steps: 500, AllowBudgetExit: true})
	xb, fb, errB := MinimiseSimulatedAnnealing(f, start, SimulatedAnnealingOptions{Seed: 7, Steps: 500, AllowBudgetExit: true})
	if errA != nil || errB != nil {
		t.Fatalf("MinimiseSimulatedAnnealing: %v, %v", errA, errB)
	}
	if math.Float64bits(fa) != math.Float64bits(fb) {
		t.Fatalf("values differ across identical runs: %.20g vs %.20g", fa, fb)
	}
	for i := range 4 {
		if math.Float64bits(xa.FloatAt(i)) != math.Float64bits(xb.FloatAt(i)) {
			t.Fatalf("x[%d] differs across identical runs", i)
		}
	}
	// A linear objective keeps producing new bests to the last
	// proposal: the schedule ends unfinished and is refused.
	line := func(a *core.Array) (float64, error) {
		s := 0.0
		for i := range a.Len() {
			s -= a.FloatAt(i)
		}
		return s, nil
	}
	_, _, err := MinimiseSimulatedAnnealing(line, start, SimulatedAnnealingOptions{Seed: 7, Steps: 2000, Tolerance: 1e-300})
	if err == nil || !strings.Contains(err.Error(), "still improving") {
		t.Fatalf("error = %v, want the unfinished-schedule refusal", err)
	}
	// The escape hatch reports the best point anyway.
	xbest, fv, err := MinimiseSimulatedAnnealing(line, start, SimulatedAnnealingOptions{Seed: 7, Steps: 2000, Tolerance: 1e-300, AllowBudgetExit: true})
	if err != nil || xbest == nil {
		t.Fatalf("AllowBudgetExit: err = %v, x = %v", err, xbest)
	}
	if fv >= 0 {
		t.Fatalf("value = %.3e, want the linear objective's negative value", fv)
	}
	// Empty and complex starts are refused.
	if _, _, err := MinimiseSimulatedAnnealing(f, core.New(core.Float, 0), SimulatedAnnealingOptions{}); err == nil {
		t.Fatal("an empty starting point was accepted")
	}
	if _, _, err := MinimiseSimulatedAnnealing(f, mustComplexPoint(t), SimulatedAnnealingOptions{}); err == nil {
		t.Fatal("a complex starting point was accepted")
	}
	// A non-finite objective at the start is refused.
	if _, _, err := MinimiseSimulatedAnnealing(func(*core.Array) (float64, error) { return math.Inf(-1), nil },
		start, SimulatedAnnealingOptions{}); err == nil {
		t.Fatal("a non-finite start value was accepted")
	}
	// An objective's own error propagates.
	if _, _, err := MinimiseSimulatedAnnealing(func(*core.Array) (float64, error) {
		return 0, base.Errf("the model exploded")
	}, start, SimulatedAnnealingOptions{}); err == nil {
		t.Fatal("the objective's error did not propagate")
	}
	// A proposal that wanders where the objective is undefined is
	// fatal, and so is one that errors: the chain climbs a linear
	// slope until it crosses the model's domain.
	climb := func(limit float64, verdict func() (float64, error)) func(*core.Array) (float64, error) {
		return func(a *core.Array) (float64, error) {
			s := 0.0
			for i := range a.Len() {
				s -= a.FloatAt(i)
			}
			if s < limit {
				return verdict()
			}
			return s, nil
		}
	}
	if _, _, err := MinimiseSimulatedAnnealing(climb(-30, func() (float64, error) { return math.NaN(), nil }),
		start, SimulatedAnnealingOptions{Seed: 7, Steps: 4000}); err == nil {
		t.Fatal("a NaN proposal value was accepted")
	}
	if _, _, err := MinimiseSimulatedAnnealing(climb(-30, func() (float64, error) {
		return 0, base.Errf("the proposal exploded")
	}), start, SimulatedAnnealingOptions{Seed: 7, Steps: 4000}); err == nil {
		t.Fatal("the proposal's error did not propagate")
	}
}

// TestCMADrawSamplesTheAdaptedCovariance pins the sampling equation:
// the candidates are B·diag(sd)·z with one unit normal per principal
// direction, so the empirical covariance of many draws reproduces the
// adapted C even after a rotation that leaves no axis aligned. A
// shared scalar across the directions sampled a diagonal distribution
// instead and failed this probe loudly.
func TestCMADrawSamplesTheAdaptedCovariance(t *testing.T) {
	n := 4
	// C = Q·D·Qᵀ for a Householder reflection through (1, 2, 3, 4) and
	// a spread of eigenvalues, so no axis survives the rotation.
	q := make([]float64, n*n)
	{
		v := []float64{1, 2, 3, 4}
		norm := 0.0
		for _, x := range v {
			norm += x * x
		}
		norm = math.Sqrt(norm)
		for i := range n {
			for j := range n {
				h := 0.0
				if i == j {
					h = 1
				}
				q[i*n+j] = h - 2*v[i]*v[j]/(norm*norm)
			}
		}
	}
	eigen := []float64{4, 1, 0.5, 0.25}
	sd := make([]float64, n)
	for i := range n {
		sd[i] = math.Sqrt(eigen[i])
	}
	c := make([]float64, n*n)
	for i := range n {
		for j := range n {
			for k := range n {
				c[i*n+j] += q[k*n+i] * eigen[k] * q[k*n+j]
			}
		}
	}
	g := core.NewGenerator(9)
	draws := 300000
	z := make([]float64, n)
	var y []float64
	sum := make([]float64, n)
	cov := make([]float64, n*n)
	for range draws {
		for i := range n {
			z[i] = g.NormalUnit()
		}
		y = make([]float64, n)
		cmaDraw(q, sd, z, y)
		for i := range n {
			sum[i] += y[i]
			for j := range n {
				cov[i*n+j] += y[i] * y[j]
			}
		}
	}
	froC, froErr := 0.0, 0.0
	for i := range n {
		for j := range n {
			cov[i*n+j] = cov[i*n+j]/float64(draws) - sum[i]*sum[j]/(float64(draws)*float64(draws))
			d := cov[i*n+j] - c[i*n+j]
			froErr += d * d
			froC += c[i*n+j] * c[i*n+j]
		}
	}
	if math.Sqrt(froErr/froC) > 0.02 {
		t.Fatalf("empirical covariance off the adapted C by %.3g relative Frobenius", math.Sqrt(froErr/froC))
	}
}