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

package optim

import (
	"cmp"
	"math"
	"slices"

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

// Two derivative-free global searchers beside the differential
// evolution in devolution.go, both drawing from the house xoshiro
// generator so a fixed seed fixes the whole trajectory.
//
// CMA-ES (the covariance matrix adaptation evolution strategy) is the
// choice for smooth continuous landscapes where a local method cannot
// be trusted to find the right basin: it adapts a full covariance
// matrix from the successful offspring, which turns it along the
// valley whatever orientation the valley has. The implementation is a
// bounded single run of the standard algorithm with the default
// parameters of N. Hansen, The CMA Evolution Strategy: A Tutorial,
// arXiv:1604.00772 (2016): rank-one and rank-mu updates, the evolution
// paths p_sigma and p_c with their tag for the stalled-path case, and
// the tutorial's equations (48) to (53) for the weights and rates.
// No restarts: the restart schemes (IPOP and friends) are the
// caller's loop, and this entry reports one run honestly.
//
// Simulated annealing is the choice for landscapes too rough, too
// discrete-like or too deceptive for covariance adaptation: a random
// walk that accepts uphill steps with the Metropolis probability and
// cools the acceptance threshold geometrically. It is a basin finder,
// not a precision optimiser.

// CMAESOptions tunes MinimiseCMAES. Sigma0 ≤ 0 means 0.3 (the
// tutorial's typical starting step for problems scaled to O(1)),
// Generations ≤ 0 means 500, Tolerance ≤ 0 means 1e-12, Seed 0 is
// replaced by 42 as in MinimiseDifferentialEvolution (any other
// value, negatives included, seeds the xoshiro stream directly).
//
// The tolerance ends the run when either the distribution has
// collapsed or the landscape has gone flat: the largest principal axis
// of the search distribution, sigma·sqrt(max C_ii), has fallen to
// Tolerance·max(1, ‖mean‖∞), or the objective spread over one
// generation has fallen to Tolerance·max(1, |best|). The axis is the
// diagonal of C, a proxy for the largest eigenvalue that avoids a
// second decomposition; a run that stops on the spread criterion on a
// genuinely flat landscape reports what it saw.
type CMAESOptions struct {
	Sigma0      float64
	Generations int
	Tolerance   float64
	Seed        int64
	// AllowBudgetExit makes a run that exhausts Generations report its
	// best point instead of an error. The default is false, so a
	// budget stop is never mistaken for a converged answer; the flag
	// mirrors LBFGSOptions.AllowBudgetExit.
	AllowBudgetExit bool
}

// MinimiseCMAES returns the point and value of the best solution f the
// strategy found. f receives candidate points as rank-1 arrays; a
// non-finite value or an error is fatal for the run. There are no
// bounds: CMA-ES is an unconstrained method, and the caller who needs
// a box should reparametrise (or use the differential-evolution entry,
// which clamps).
//
// Termination, budget and honesty: a run ends when the tolerance
// criteria above are met (converged) or when the generation budget is
// spent. Following the house budget policy, an exhausted budget is an
// error naming the best value reached and the tolerance it fell short
// of, never a silent answer; AllowBudgetExit is the documented escape
// hatch.
func MinimiseCMAES(f func(*core.Array) (float64, error), x0 *core.Array, opts CMAESOptions) (*core.Array, float64, error) {
	const name = "MinimiseCMAES"
	if x0.Dtype() == core.Complex {
		return nil, 0, base.Errf("%s: complex starting points are not supported", name)
	}
	n := x0.Len()
	if n == 0 {
		return nil, 0, base.Errf("%s: the starting point must have at least one element", name)
	}
	sigma := opts.Sigma0
	if sigma <= 0 {
		sigma = 0.3
	}
	generations := opts.Generations
	if generations <= 0 {
		generations = 500
	}
	tol := opts.Tolerance
	if tol <= 0 {
		tol = 1e-12
	}
	seed := opts.Seed
	if seed == 0 {
		seed = 42
	}
	g := core.NewGenerator(seed)

	// The tutorial's default parameters, equations (48) to (53).
	lambda := 4 + int(3*math.Log(float64(n)))
	mu := lambda / 2
	weights := make([]float64, mu)
	wSum := 0.0
	for i := range mu {
		weights[i] = math.Log(float64(lambda)/2+0.5) - math.Log(float64(i+1))
		wSum += weights[i]
	}
	for i := range mu {
		weights[i] /= wSum
	}
	muEff := 0.0
	for _, w := range weights {
		muEff += w * w
	}
	muEff = 1 / muEff
	cSigma := (muEff + 2) / (float64(n) + muEff + 5)
	dSigma := 1 + 2*math.Max(0, math.Sqrt((muEff-1)/(float64(n)+1))-1) + cSigma
	cC := (4 + muEff/float64(n)) / (4 + float64(n) + 2*muEff/float64(n))
	c1 := 2 / ((float64(n)+1.3)*(float64(n)+1.3) + muEff)
	cMu := math.Min(1-c1, 2*(muEff-2+1/muEff)/((float64(n)+2)*(float64(n)+2)+muEff))
	chiN := math.Sqrt(float64(n)) * (1 - 1/(4*float64(n)) + 1/(21*float64(n)*float64(n)))

	mean := cloneDense(x0)
	cov := make([]float64, n*n)
	for i := range n {
		cov[i*n+i] = 1
	}
	pathC := make([]float64, n)
	pathS := make([]float64, n)

	eval := func(p []float64) (float64, error) {
		v, err := f(linalg.ArrayFromFloatsSafe(p, len(p)))
		if err != nil {
			return 0, base.Errf("%s: %w", name, err)
		}
		if math.IsNaN(v) || math.IsInf(v, 0) {
			return 0, base.Errf("%s: the objective is non-finite (%g)", name, v)
		}
		return v, nil
	}

	bestF, bestX := math.Inf(1), make([]float64, n)
	offX := make([]float64, lambda*n)
	offY := make([]float64, lambda*n)
	offF := make([]float64, lambda)
	order := make([]int, lambda)
	yW := make([]float64, n)
	yC := make([]float64, n)
	newCov := make([]float64, n*n)
	// Per-generation scratch, owned by the run: the sampling scale per
	// principal axis, the unit normal drawn per offspring, the
	// eigenvector projection rootInverse accumulates in, the rank-mu
	// outer-product accumulator with its weighted axis vector, and the
	// diagonalisation's own working set. Each is fully overwritten
	// before it is read.
	sd := make([]float64, n)
	z := make([]float64, n)
	yInv := make([]float64, n)
	rankMu := make([]float64, n*n)
	wy := make([]float64, n)
	var eig jacobiScratch

	for gen := range generations {
		vals, vecs := eig.eigen(cov, n)
		// Sampling scale per principal axis, floored at zero: a
		// numerically degenerate axis contributes nothing rather than
		// a complex square root.
		for i := range n {
			sd[i] = math.Sqrt(math.Max(vals[i], 0))
		}
		worstF := math.Inf(-1)
		bestGen := math.Inf(1)
		for k := range lambda {
			// One unit normal per principal direction: the draw is
			// B·diag(sd)·z, whose covariance is exactly C. Sharing a
			// single scalar across the directions would sample a
			// diagonal distribution scaled by one fixed vector and
			// leave the adapted covariance unused.
			for i := range n {
				z[i] = g.NormalUnit()
			}
			cmaDraw(vecs, sd, z, offY[k*n:k*n+n])
			for i := range n {
				offX[k*n+i] = mean[i] + sigma*offY[k*n+i]
			}
			v, err := eval(offX[k*n : k*n+n])
			if err != nil {
				return nil, 0, err
			}
			offF[k] = v
			worstF = math.Max(worstF, v)
			bestGen = math.Min(bestGen, v)
		}
		for k := range lambda {
			order[k] = k
		}
		slices.SortStableFunc(order, func(a, b int) int { return cmp.Compare(offF[a], offF[b]) })
		if offF[order[0]] < bestF {
			bestF = offF[order[0]]
			copy(bestX, offX[order[0]*n:(order[0]+1)*n])
		}
		clear(yW)
		for k := range mu {
			w := weights[k]
			for i := range n {
				yW[i] += w * offY[order[k]*n+i]
			}
		}
		for i := range n {
			mean[i] += sigma * yW[i]
		}

		// Conjugate evolution path: C^{-1/2} y_w through the
		// eigendecomposition, then the sigma update from its length.
		rootInverse(yC, yW, vals, vecs, yInv, n)
		psNorm := 0.0
		for i := range n {
			pathS[i] = (1-cSigma)*pathS[i] + math.Sqrt(cSigma*(2-cSigma)*muEff)*yC[i]
			psNorm += pathS[i] * pathS[i]
		}
		psNorm = math.Sqrt(psNorm)
		sigma *= math.Exp((cSigma / dSigma) * (psNorm/chiN - 1))
		hs := 0.0
		if psNorm/math.Sqrt(1-math.Pow(1-cSigma, 2*float64(gen+1))) < (1.4+2/(float64(n)+1))*chiN {
			hs = 1
		}
		pcNormScale := math.Sqrt(cC * (2 - cC) * muEff)
		for i := range n {
			pathC[i] = (1-cC)*pathC[i] + hs*pcNormScale*yW[i]
		}
		// The rank-one and rank-mu updates; delta(hs) keeps the
		// covariance from growing along p_c across a stall of the
		// conjugate path. The rank-mu sum accumulates as contiguous
		// outer products, one per selected offspring walked in order:
		// every cell still sums (w·yᵢ)·yⱼ over ascending k, so the
		// bits are the per-cell walk's and the inner loop stays on
		// unit stride.
		clear(rankMu)
		for k := range mu {
			base := order[k] * n
			w := weights[k]
			for i := range n {
				wy[i] = w * offY[base+i]
			}
			for i := range n {
				wi := wy[i]
				row := rankMu[i*n : i*n+n]
				yk := offY[base : base+n]
				for j := range n {
					row[j] += wi * yk[j]
				}
			}
		}
		factor := 1 + c1*(1-hs) - c1 - cMu
		for i := range n {
			ci := c1 * pathC[i]
			off := i * n
			for j := range n {
				newCov[off+j] = factor*cov[off+j] + ci*pathC[j] + cMu*rankMu[off+j]
			}
		}
		for i := range n {
			for j := i + 1; j < n; j++ {
				avg := (newCov[i*n+j] + newCov[j*n+i]) / 2
				newCov[i*n+j], newCov[j*n+i] = avg, avg
			}
		}
		copy(cov, newCov)

		axis := 0.0
		for i := range n {
			axis = math.Max(axis, sigma*math.Sqrt(math.Max(cov[i*n+i], 0)))
		}
		if math.IsNaN(axis) || sigma > 1e12 {
			return nil, 0, base.Errf("%s: the search diverged (sigma %g, largest axis %g) at f = %g", name, sigma, axis, bestF)
		}
		if axis <= tol*math.Max(1, maxAbs(mean)) || worstF-bestGen <= tol*math.Max(1, math.Abs(bestF)) {
			out, fv := packResult(bestX, bestF)
			return out, fv, nil
		}
	}
	if !opts.AllowBudgetExit {
		return nil, 0, base.Errf("%s: the generation budget of %d ran out at f = %g, above the tolerance %g", name, generations, bestF, tol)
	}
	out, fv := packResult(bestX, bestF)
	return out, fv, nil
}

// rootInverse writes C^{-1/2} y into dst through the eigendecomposition
// the caller already paid for: B·diag(1/sqrt(d))·Bᵀ·y with the
// eigenvalues floored away from zero so a numerically flat direction
// cannot divide by nothing. tmp is scratch of length at least n, fully
// overwritten before it is read.
func rootInverse(dst, y, vals, vecs, tmp []float64, n int) {
	scale := 0.0
	for i := range n {
		scale = math.Max(scale, vals[i])
	}
	floor := 1e-20 * math.Max(1, scale)
	for j := range n { // tmp = Bᵀ y
		s := 0.0
		for i := range n {
			s += vecs[j*n+i] * y[i]
		}
		tmp[j] = s / math.Sqrt(math.Max(vals[j], floor))
	}
	clear(dst)
	for j := range n { // dst = B tmp
		c := tmp[j]
		if c == 0 {
			continue
		}
		for i := range n {
			dst[i] += vecs[j*n+i] * c
		}
	}
}

// eigenPair is one eigenvalue with the column its eigenvector occupies
// in the accumulator the diagonalisation carries.
type eigenPair struct {
	val   float64
	index int
}

// jacobiScratch is one diagonalisation's working set: the cyclically
// rotated copy of the matrix, the eigenvector accumulator, the
// value/index pairs the descending sort walks and the two blocks handed
// back. The driver keeps one instance for the whole run, so the
// per-generation decomposition allocates nothing; every buffer is fully
// overwritten before it is read.
type jacobiScratch struct {
	work   []float64
	acc    []float64
	pairs  []eigenPair
	vals   []float64
	sorted []float64
	evecs  []float64
}

// jacobiEigen diagonalises the symmetric row-major n×n matrix into
// freshly allocated blocks, the allocating entry over
// (*jacobiScratch).eigen.
func jacobiEigen(a []float64, n int) (vals, vecs []float64) {
	var s jacobiScratch
	return s.eigen(a, n)
}

// eigen diagonalises the symmetric row-major n×n matrix by cyclic
// Jacobi rotations: the eigenvalues come back in descending order and
// vecs[j*n+i] is component i of the eigenvector that belongs to
// vals[j]. The sweeps stop once the off-diagonal mass has fallen to
// the working precision of the matrix's own Frobenius norm. Both
// returned blocks belong to the scratch and stay valid until its next
// call.
func (s *jacobiScratch) eigen(a []float64, n int) (vals, vecs []float64) {
	if cap(s.work) < n*n {
		s.work = make([]float64, n*n)
	}
	work := s.work[:n*n]
	copy(work, a)
	if cap(s.acc) < n*n {
		s.acc = make([]float64, n*n)
	}
	acc := s.acc[:n*n]
	clear(acc)
	for i := range n {
		acc[i*n+i] = 1
	}
	frob := 0.0
	for _, v := range work {
		frob += v * v
	}
	const maxSweeps = 60
	for range maxSweeps {
		off := 0.0
		for i := range n {
			for j := i + 1; j < n; j++ {
				off += work[i*n+j] * work[i*n+j]
			}
		}
		if off <= 1e-30*math.Max(frob, 1) {
			break
		}
		for p := range n {
			for q := p + 1; q < n; q++ {
				apq := work[p*n+q]
				if apq == 0 {
					continue
				}
				theta := (work[q*n+q] - work[p*n+p]) / (2 * apq)
				t := 1 / (math.Abs(theta) + math.Sqrt(theta*theta+1))
				if theta < 0 {
					t = -t
				}
				c := 1 / math.Sqrt(t*t+1)
				s := t * c
				for k := range n {
					akp, akq := work[k*n+p], work[k*n+q]
					work[k*n+p] = c*akp - s*akq
					work[k*n+q] = s*akp + c*akq
				}
				for k := range n {
					apk, aqk := work[p*n+k], work[q*n+k]
					work[p*n+k] = c*apk - s*aqk
					work[q*n+k] = s*apk + c*aqk
				}
				work[p*n+q], work[q*n+p] = 0, 0
				for k := range n {
					vkp, vkq := acc[k*n+p], acc[k*n+q]
					acc[k*n+p] = c*vkp - s*vkq
					acc[k*n+q] = s*vkp + c*vkq
				}
			}
		}
	}
	if cap(s.vals) < n {
		s.vals = make([]float64, n)
	}
	vals = s.vals[:n]
	for i := range n {
		vals[i] = work[i*n+i]
	}
	// Sort the pairs descending by value. The accumulator carries
	// eigenvector k in its column k, and the documented layout here is
	// row j for vector j, so the extraction transposes.
	if cap(s.pairs) < n {
		s.pairs = make([]eigenPair, n)
	}
	pairs := s.pairs[:n]
	for i := range n {
		pairs[i] = eigenPair{vals[i], i}
	}
	slices.SortFunc(pairs, func(a, b eigenPair) int { return cmp.Compare(b.val, a.val) })
	if cap(s.sorted) < n {
		s.sorted = make([]float64, n)
	}
	if cap(s.evecs) < n*n {
		s.evecs = make([]float64, n*n)
	}
	sortedVals := s.sorted[:n]
	sortedVecs := s.evecs[:n*n]
	for j, pr := range pairs {
		sortedVals[j] = pr.val
		for i := range n {
			sortedVecs[j*n+i] = acc[i*n+pr.index]
		}
	}
	return sortedVals, sortedVecs
}

// SimulatedAnnealingOptions tunes MinimiseSimulatedAnnealing.
// Steps ≤ 0 means 20000 proposals, Temperature0 ≤ 0 means 1,
// CoolingRate ≤ 0 means 0.9995 (the geometric factor per proposal),
// StepScale ≤ 0 means 0.1 (the relative Gaussian proposal scale),
// Tolerance ≤ 0 means 1e-6, Seed 0 is replaced by 42 as in
// MinimiseDifferentialEvolution.
//
// The schedule is the algorithm: the temperature runs from
// Temperature0 down by the factor CoolingRate at every proposal, so
// the chain freezes exponentially and the tail of the schedule is a
// local polish. Tolerance is the convergence test on that tail: the
// best value may not improve by more than Tolerance·max(1, |best|)
// over the final quarter of the schedule. A run that is still
// improving at the end of the schedule has not converged, and the
// budget refusal says so with both figures, exactly as the local
// solvers refuse an unfinished run; AllowBudgetExit is the escape
// hatch that reports the best point anyway.
type SimulatedAnnealingOptions struct {
	Steps        int
	Temperature0 float64
	CoolingRate  float64
	StepScale    float64
	Seed         int64
	Tolerance    float64
	// AllowBudgetExit makes a run whose schedule ended while the best
	// value was still improving report its best point instead of an
	// error. The flag mirrors LBFGSOptions.AllowBudgetExit.
	AllowBudgetExit bool
}

// MinimiseSimulatedAnnealing returns the best point and value the
// chain found: a Metropolis walk over the landscape with a Gaussian
// proposal of the given relative scale (coordinate i moves by
// StepScale·max(1, |x_i|) standard normals) and the geometric cooling
// schedule above. f receives candidate points as rank-1 arrays; a
// non-finite value or an error is fatal for the run. Simulated
// annealing finds basins, not minima to full precision: run
// MinimiseLBFGS from the returned point when a polished answer is
// wanted, which is the standard composition.
func MinimiseSimulatedAnnealing(f func(*core.Array) (float64, error), x0 *core.Array, opts SimulatedAnnealingOptions) (*core.Array, float64, error) {
	const name = "MinimiseSimulatedAnnealing"
	if x0.Dtype() == core.Complex {
		return nil, 0, base.Errf("%s: complex starting points are not supported", name)
	}
	n := x0.Len()
	if n == 0 {
		return nil, 0, base.Errf("%s: the starting point must have at least one element", name)
	}
	steps := opts.Steps
	if steps <= 0 {
		steps = 20000
	}
	t0 := opts.Temperature0
	if t0 <= 0 {
		t0 = 1
	}
	cooling := opts.CoolingRate
	if cooling > 1 {
		return nil, 0, base.Errf("%s: the cooling rate %g would heat the chain; want a factor in (0, 1]", name, cooling)
	}
	if cooling <= 0 {
		cooling = 0.9995
	}
	scale := opts.StepScale
	if scale <= 0 {
		scale = 0.1
	}
	tol := opts.Tolerance
	if tol <= 0 {
		tol = 1e-6
	}
	seed := opts.Seed
	if seed == 0 {
		seed = 42
	}
	g := core.NewGenerator(seed)

	x := cloneDense(x0)
	cur, err := f(linalg.ArrayFromFloatsSafe(x, n))
	if err != nil {
		return nil, 0, base.Errf("%s: %w", name, err)
	}
	if math.IsNaN(cur) || math.IsInf(cur, 0) {
		return nil, 0, base.Errf("%s: the objective is non-finite (%g) at the start", name, cur)
	}
	bestF, bestX := cur, slices.Clone(x)
	proposal := make([]float64, n)
	quarterF := math.Inf(1)
	quarterAt := 3 * steps / 4

	for k := range steps {
		temperature := t0 * math.Pow(cooling, float64(k))
		for i := range n {
			proposal[i] = x[i] + scale*math.Max(1, math.Abs(x[i]))*g.NormalUnit()
		}
		fv, ferr := f(linalg.ArrayFromFloatsSafe(proposal, n))
		if ferr != nil {
			return nil, 0, base.Errf("%s: %w", name, ferr)
		}
		if math.IsNaN(fv) || math.IsInf(fv, 0) {
			return nil, 0, base.Errf("%s: the objective is non-finite (%g) at proposal %d", name, fv, k+1)
		}
		delta := fv - cur
		if delta <= 0 || g.Unit() < math.Exp(-delta/temperature) {
			copy(x, proposal)
			cur = fv
		}
		if cur < bestF {
			bestF = cur
			copy(bestX, x)
		}
		if k == quarterAt {
			quarterF = bestF
		}
	}
	// The schedule ended: convergence is the frozen tail, judged by
	// the improvement the final quarter still bought.
	if quarterF-bestF > tol*math.Max(1, math.Abs(bestF)) && !opts.AllowBudgetExit {
		return nil, 0, base.Errf("%s: the schedule of %d steps ended with the best value still improving (%g to %g); raise Steps or set AllowBudgetExit",
			name, steps, quarterF, bestF)
	}
	out, fv := packResult(bestX, bestF)
	return out, fv, nil
}

// cmaDraw fills y with B·diag(sd)·z: one unit normal per principal
// direction transformed into the covariance's own axes, the draw the
// tutorial's sampling equation defines, whose covariance is the
// adapted C itself. The walk goes one principal direction at a time so
// the eigenvector rows are read on unit stride; the per-element
// grouping (v·sdⱼ)·zⱼ and the ascending j order are the i-outer walk's,
// so the draw is bit for bit the same.
func cmaDraw(vecs, sd, z, y []float64) {
	clear(y)
	n := len(y)
	for j := range z {
		sdJ, zJ := sd[j], z[j]
		row := vecs[j*n : j*n+n]
		for i := range n {
			y[i] += row[i] * sdJ * zJ
		}
	}
}