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

package optim

import (
	"math"
	"math/rand/v2"
	"testing"

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

// Benchmarks for the solver iterations whose cost is dominated by
// repeated dense linear algebra or by a difference stencil: the revised
// simplex, the active-set QP, CMA-ES, L-BFGS and the finite-difference
// paths. Every input comes from a fixed seed, so each run walks one
// deterministic trajectory.

// solversRNG returns a generator with a fixed stream: the inputs are
// identical on every machine and every run.
func solversRNG() *rand.Rand { return rand.New(rand.NewPCG(0x5eed, 0x1234)) }

// solverArray builds a float array, panicking on a bad shape: every
// caller passes a literal shape.
func solverArray(vals []float64, shape ...int) *core.Array {
	a, err := core.FromFloats(vals, shape...)
	if err != nil {
		panic("solverArray: " + err.Error())
	}
	return a
}

// lpProblem builds a standard-form LP with m rows and n = 2m columns:
// A = [I | R] with every entry of R at least 0.5, b = 1 and a random
// cost. The identity block makes x = (1, …, 1, 0, …, 0) feasible, and
// the feasible set is bounded: the slack block forces R·x_R ≤ 1, whose
// positive coefficients bound the 1-norm of x_R by 2, and x_I = 1 −
// R·x_R is bounded with it. The run therefore ends at a vertex rather
// than on an unbounded ray.
func lpProblem(m int, rng *rand.Rand) (c, a, b *core.Array) {
	n := 2 * m
	av := make([]float64, m*n)
	for i := range m {
		av[i*n+i] = 1
		for j := range m {
			av[i*n+m+j] = 0.5 + rng.Float64()
		}
	}
	cv := make([]float64, n)
	for j := range n {
		cv[j] = 2*rng.Float64() - 1
	}
	bv := make([]float64, m)
	for i := range m {
		bv[i] = 1
	}
	return solverArray(cv, n), solverArray(av, m, n), solverArray(bv, m)
}

// qpProblem builds a strictly convex quadratic in n variables with r
// two-sided rows that all admit the origin, so the start point is
// feasible and the benchmark measures the active-set iteration itself.
func qpProblem(n, r int, rng *rand.Rand) (*core.Array, *core.Array, LinearConstraints) {
	hv := make([]float64, n*n)
	for i := range n {
		hv[i*n+i] = 1 + rng.Float64()
		for j := i + 1; j < n; j++ {
			v := 0.25 * (2*rng.Float64() - 1)
			hv[i*n+j], hv[j*n+i] = v, v
		}
	}
	cv := make([]float64, n)
	for j := range n {
		cv[j] = 2*rng.Float64() - 1
	}
	av := make([]float64, r*n)
	lower := make([]float64, r)
	upper := make([]float64, r)
	for i := range r {
		for j := range n {
			av[i*n+j] = 2*rng.Float64() - 1
		}
		lower[i], upper[i] = math.Inf(-1), 0.05+0.1*rng.Float64()
	}
	cons := LinearConstraints{A: solverArray(av, r, n), Lower: lower, Upper: upper}
	return solverArray(hv, n, n), solverArray(cv, n), cons
}

// solverBowl returns a coupled bowl around (0.5, …, 0.5) and a start
// point away from it.
func solverBowl(n int, rng *rand.Rand) (f func(*core.Array) (float64, error), x0 *core.Array) {
	weights := make([]float64, n)
	for i := range n {
		weights[i] = 1 + 2*rng.Float64()
	}
	f = func(p *core.Array) (float64, error) {
		total := 0.0
		for i := range n {
			d := p.FloatAt(i) - 0.5
			total += weights[i] * d * d
			if i+1 < n {
				total += 0.3 * d * (p.FloatAt(i+1) - 0.5)
			}
		}
		return total, nil
	}
	start := make([]float64, n)
	for i := range start {
		start[i] = 1.5 + 0.1*float64(i)
	}
	return f, solverArray(start, n)
}

// BenchmarkSolversLinearProgram measures the two-phase revised simplex
// on a bounded LP with 60 rows and 120 columns, the shape the per-pivot
// refactorisation pays for.
func BenchmarkSolversLinearProgram(b *testing.B) {
	rng := solversRNG()
	c, a, rhs := lpProblem(60, rng)
	opts := LinearProgramOptions{}
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := MinimiseLinear(c, a, rhs, opts); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSolversLinearProgramRows measures an LP of the same family
// through the two-sided-row wrapper, whose conversion builds the
// standard form before the same simplex runs. The equality rows carry
// the LP itself and every variable is boxed, so the feasible set the
// wrapper sees is bounded.
func BenchmarkSolversLinearProgramRows(b *testing.B) {
	rng := solversRNG()
	const m = 8
	c, a, rhs := lpProblem(m, rng)
	n := c.Len()
	rows := m + n
	av := make([]float64, rows*n)
	for i := range m {
		for j := range n {
			av[i*n+j] = a.FloatAt(i*n + j)
		}
	}
	lower := make([]float64, rows)
	upper := make([]float64, rows)
	for i := range m {
		lower[i], upper[i] = rhs.FloatAt(i), rhs.FloatAt(i)
	}
	for i := range n {
		av[(m+i)*n+i] = 1
		lower[m+i], upper[m+i] = -3, 3
	}
	cons := LinearConstraints{A: solverArray(av, rows, n), Lower: lower, Upper: upper}
	opts := LinearProgramOptions{}
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := MinimiseLinearRows(c, cons, opts); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSolversQuadraticProgram measures the active-set QP on a
// strictly convex quadratic with 10 variables and 24 two-sided rows,
// tight enough that the working set moves several times.
func BenchmarkSolversQuadraticProgram(b *testing.B) {
	rng := solversRNG()
	h, c, cons := qpProblem(12, 30, rng)
	x0 := solverArray(make([]float64, 12), 12)
	opts := QPOptions{}
	b.ReportAllocs()
	for b.Loop() {
		if _, _, _, err := MinimiseQP(h, c, cons, x0, opts); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSolversCMAES measures the covariance-adaptation strategy on
// a four-dimensional bowl, one generation pair being cheap at that size.
func BenchmarkSolversCMAES(b *testing.B) {
	rng := solversRNG()
	f, x0 := solverBowl(4, rng)
	opts := CMAESOptions{Sigma0: 0.5, Generations: 60, Tolerance: 1e-12, Seed: 7, AllowBudgetExit: true}
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := MinimiseCMAES(f, x0, opts); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSolversLBFGSLeastSquares fits a six-term cosine series to
// samples of a fixed function with a consistent analytic gradient, the
// path whose history buffer rotates once the memory is full.
func BenchmarkSolversLBFGSLeastSquares(b *testing.B) {
	const nPar, nObs = 6, 64
	truth := make([]float64, nPar)
	for j := range truth {
		truth[j] = 1 / float64(j+1)
	}
	tSamples := make([]float64, nObs)
	y := make([]float64, nObs)
	for i := range nObs {
		t := 0.5 * float64(i) / float64(nObs)
		tSamples[i] = t
		for j := range nPar {
			y[i] += truth[j] * math.Cos(float64(j)*t)
		}
	}
	f := func(p *core.Array) (float64, error) {
		total := 0.0
		for i := range nObs {
			model := 0.0
			for j := range nPar {
				model += p.FloatAt(j) * math.Cos(float64(j)*tSamples[i])
			}
			d := model - y[i]
			total += d * d
		}
		return total, nil
	}
	grad := func(p *core.Array) (*core.Array, error) {
		g := make([]float64, nPar)
		for i := range nObs {
			model := 0.0
			for j := range nPar {
				model += p.FloatAt(j) * math.Cos(float64(j)*tSamples[i])
			}
			d := 2 * (model - y[i])
			for j := range nPar {
				g[j] += d * math.Cos(float64(j)*tSamples[i])
			}
		}
		return core.FromFloats(g, nPar)
	}
	x0 := solverArray(make([]float64, nPar), nPar)
	opts := LBFGSOptions{MaxIterations: 300, Memory: 4}
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := MinimiseLBFGS(f, grad, x0, opts); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSolversLBFGSFiniteDiff measures the central-difference
// gradient path: two objective evaluations per coordinate per step.
func BenchmarkSolversLBFGSFiniteDiff(b *testing.B) {
	rng := solversRNG()
	f, x0 := solverBowl(24, rng)
	opts := LBFGSOptions{MaxIterations: 60}
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := MinimiseLBFGS(f, nil, x0, opts); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSolversLevenbergFiniteDiff measures Levenberg-Marquardt on a
// small polynomial fit with the Jacobian by central differences.
func BenchmarkSolversLevenbergFiniteDiff(b *testing.B) {
	const nObs, nPar = 40, 6
	t := make([]float64, nObs)
	obs := make([]float64, nObs)
	truth := make([]float64, nPar)
	for j := range truth {
		truth[j] = 0.5 + 0.1*float64(j)
	}
	for i := range nObs {
		t[i] = float64(i) / 8
		v := 0.0
		for j := range nPar {
			v += truth[j] * math.Pow(t[i], float64(j))
		}
		obs[i] = v
	}
	residual := func(p *core.Array) (*core.Array, error) {
		r := make([]float64, nObs)
		for i := range nObs {
			v := 0.0
			for j := range nPar {
				v += p.FloatAt(j) * math.Pow(t[i], float64(j))
			}
			r[i] = v - obs[i]
		}
		return core.FromFloats(r, nObs)
	}
	p0 := solverArray(make([]float64, nPar), nPar)
	opts := LMOptions{MaxIterations: 20}
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := LevenbergMarquardt(residual, p0, opts); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSolversRootSystemFiniteDiff measures the damped Newton
// iteration with a per-step central-difference Jacobian.
func BenchmarkSolversRootSystemFiniteDiff(b *testing.B) {
	const n = 8
	r := func(x *core.Array) (*core.Array, error) {
		out := make([]float64, n)
		for i := range n {
			v := x.FloatAt(i)
			out[i] = v*v + 0.1*v - float64(i+1)
			if i+1 < n {
				out[i] += 0.05 * x.FloatAt(i+1)
			}
		}
		return core.FromFloats(out, n)
	}
	start := make([]float64, n)
	for i := range start {
		start[i] = 1 + 0.1*float64(i)
	}
	x0 := solverArray(start, n)
	opts := RootSystemOptions{MaxIterations: 20}
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := FindRootSystem(r, x0, opts); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSolversNonlinearConstrained measures the augmented-Lagrangian
// outer loop over one equality and two inequality rows, whose stencil is
// the per-row, per-coordinate hot path.
func BenchmarkSolversNonlinearConstrained(b *testing.B) {
	const n = 6
	f := func(p *core.Array) (float64, error) { return p.FloatAt(0), nil }
	cons := NonlinearConstraints{
		Equalities: []func(*core.Array) (float64, error){
			func(p *core.Array) (float64, error) {
				s := -1.0
				for i := range n {
					s += p.FloatAt(i) * p.FloatAt(i)
				}
				return s, nil
			},
		},
		Inequalities: []func(*core.Array) (float64, error){
			func(p *core.Array) (float64, error) {
				s := 0.0
				for i := range n {
					s += p.FloatAt(i)
				}
				return s - 2, nil
			},
			func(p *core.Array) (float64, error) { return -p.FloatAt(0) - 2, nil },
		},
	}
	start := make([]float64, n)
	for i := range start {
		start[i] = 0.5
	}
	x0 := solverArray(start, n)
	b.ReportAllocs()
	for b.Loop() {
		if _, _, _, err := MinimiseNonlinearConstrained(f, nil, x0, cons, LBFGSOptions{}); err != nil {
			b.Fatal(err)
		}
	}
}