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

// Benchmarks for the optimiser hot paths: the L-BFGS two-loop recursion
// with analytic and finite-difference gradients, the simplex method,
// Levenberg-Marquardt's normal equations and the damped Newton system
// solver.

// benchQuadratic builds a separable convex quadratic
// f(x) = Σ (xᵢ − cᵢ)² + 0.01·Σ xᵢ² with cᵢ = i/n, whose minimum and
// gradient are closed form, so the L-BFGS runs are identical every
// iteration.
func benchQuadratic(n int) (f func(*core.Array) (float64, error), grad func(*core.Array) (*core.Array, error), x0 []float64) {
	c := make([]float64, n)
	for i := range c {
		c[i] = float64(i) / float64(n)
	}
	f = func(a *core.Array) (float64, error) {
		total := 0.0
		for i := range n {
			d := a.FloatAt(i) - c[i]
			total += d*d + 0.01*a.FloatAt(i)*a.FloatAt(i)
		}
		return total, nil
	}
	grad = func(a *core.Array) (*core.Array, error) {
		out := core.New(core.Float, n)
		vals := out.RawFloats()
		for i := range n {
			vals[i] = 2*(a.FloatAt(i)-c[i]) + 0.02*a.FloatAt(i)
		}
		return out, nil
	}
	x0 = make([]float64, n)
	for i := range x0 {
		x0[i] = 1
	}
	return f, grad, x0
}

func benchVector(b *testing.B, vals []float64) *core.Array {
	b.Helper()
	a, err := core.FromFloats(vals, len(vals))
	if err != nil {
		b.Fatal(err)
	}
	return a
}

func BenchmarkMinimiseLBFGS(b *testing.B) {
	f, grad, x0 := benchQuadratic(64)
	start := benchVector(b, x0)
	opts := LBFGSOptions{MaxIterations: 200}
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := MinimiseLBFGS(f, grad, start, opts); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkMinimiseLBFGSFiniteDiff(b *testing.B) {
	f, _, x0 := benchQuadratic(64)
	start := benchVector(b, x0)
	opts := LBFGSOptions{MaxIterations: 200}
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := MinimiseLBFGS(f, nil, start, opts); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkMinimiseLBFGSBounded(b *testing.B) {
	f, grad, x0 := benchQuadratic(64)
	lower := make([]float64, 64)
	upper := make([]float64, 64)
	for i := range upper {
		lower[i] = -2
		upper[i] = 2
	}
	start := benchVector(b, x0)
	opts := LBFGSOptions{MaxIterations: 200, Lower: lower, Upper: upper}
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := MinimiseLBFGS(f, grad, start, opts); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkMinimiseSimplex(b *testing.B) {
	f, _, x0 := benchQuadratic(8)
	start := benchVector(b, x0)
	opts := MinimiseOptions{MaxIterations: 500}
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := Minimise(f, start, opts); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkLevenbergMarquardt(b *testing.B) {
	// Fit y = p0·exp(−p1·t) on 40 noisy-free samples.
	const nObs = 40
	t := make([]float64, nObs)
	y := make([]float64, nObs)
	for i := range nObs {
		t[i] = float64(i) / 4
		y[i] = 2.5 * math.Exp(-0.7*t[i])
	}
	residual := func(p *core.Array) (*core.Array, error) {
		out := core.New(core.Float, nObs)
		vals := out.RawFloats()
		for i := range nObs {
			vals[i] = p.FloatAt(0)*math.Exp(-p.FloatAt(1)*t[i]) - y[i]
		}
		return out, nil
	}
	p0 := benchVector(b, []float64{1, 0.2})
	opts := LMOptions{MaxIterations: 30}
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := LevenbergMarquardt(residual, p0, opts); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkFindRootSystem(b *testing.B) {
	n := 6
	r := func(x *core.Array) (*core.Array, error) {
		out := core.New(core.Float, n)
		vals := out.RawFloats()
		for i := range n {
			vals[i] = x.FloatAt(i)*x.FloatAt(i) - float64(i+1)
		}
		return out, nil
	}
	start := make([]float64, n)
	for i := range start {
		start[i] = float64(i) + 1.5
	}
	x0 := benchVector(b, start)
	opts := RootSystemOptions{MaxIterations: 40}
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := FindRootSystem(r, x0, opts); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkMinimiseDifferentialEvolution(b *testing.B) {
	f, _, _ := benchQuadratic(4)
	lower := benchVector(b, []float64{-5, -5, -5, -5})
	upper := benchVector(b, []float64{5, 5, 5, 5})
	opts := DifferentialEvolutionOptions{Generations: 20}
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := MinimiseDifferentialEvolution(f, lower, upper, opts); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkFindRoot guards the scalar Brent iteration the root
// benchmarks above surround.
func BenchmarkFindRoot(b *testing.B) {
	f := math.Cos
	b.ReportAllocs()
	for b.Loop() {
		if _, err := FindRoot(f, 0.5, 2, 0); err != nil {
			b.Fatal(err)
		}
	}
}