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

package optim

import (
	"math"

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

// Differential evolution: the global optimiser for the
// landscapes the local methods cannot be trusted with: multimodal,
// discontinuous, derivative-free. The classic rand/1/bin scheme:
// every generation, each population member is challenged by a mutant
// built from three distinct others, mixed by binomial crossover, and
// kept only if it beats the incumbent. No gradient, no assumptions
// beyond the bounds.

// DifferentialEvolutionOptions tunes MinimiseDifferentialEvolution.
// Population defaults to 15·d when unset (at least 4, the scheme's
// minimum); F is the differential weight (default 0.7), CR the
// crossover probability (default 0.9), Generations the budget
// (default 1000), Seed the generator seed (zero is replaced by 42, so
// unset runs are reproducible; every other value, negatives included,
// seeds the xoshiro stream directly).
type DifferentialEvolutionOptions struct {
	Population  int
	F           float64
	CR          float64
	Generations int
	Seed        int64
}

// MinimiseDifferentialEvolution returns the point and value of the
// global minimum of f over the box [lower, upper] by differential
// evolution (rand/1/bin with reflection-free clamping to the bounds).
// f receives candidate points as rank-1 arrays; a non-finite value is
// an error, mismatched or degenerate bounds are errors, and the result
// is the best point ever evaluated, fresh array the caller owns.
// The generation budget is a tuning parameter, not a convergence
// budget: differential evolution keeps improving the population as
// long as it runs, so the budget running out returns the best point
// found without an error, unlike the local solvers whose AllowBudgetExit
// default refuses a budget stop.
func MinimiseDifferentialEvolution(f func(*core.Array) (float64, error),
	lower, upper *core.Array, opts DifferentialEvolutionOptions) (*core.Array, float64, error) {
	const name = "MinimiseDifferentialEvolution"
	n := lower.Len()
	if lower.NDim() != 1 || upper.NDim() != 1 || upper.Len() != n {
		return nil, 0, base.Errf("%s: lower and upper must be equal-length rank-1 bounds", name)
	}
	if n == 0 {
		return nil, 0, base.Errf("%s: the bounds must not be empty", name)
	}
	if lower.Dtype() == core.Complex || upper.Dtype() == core.Complex {
		return nil, 0, base.Errf("%s: complex bounds are not supported", name)
	}
	for i := range n {
		if !(upper.FloatAt(i) > lower.FloatAt(i)) {
			return nil, 0, base.Errf("%s: bound %d runs from %g to %g", name, i, lower.FloatAt(i), upper.FloatAt(i))
		}
	}
	pop := opts.Population
	if pop <= 0 {
		pop = max(15*n, 4)
	}
	// The scheme draws three distinct others besides the target, so a
	// population below four has no admissible draw: the search would
	// never leave the picking loop.
	if pop < 4 {
		return nil, 0, base.Errf("%s: the population must be at least 4 to draw three distinct others, got %d",
			name, pop)
	}
	fw := opts.F
	if fw <= 0 {
		fw = 0.7
	}
	cr := opts.CR
	if cr <= 0 {
		cr = 0.9
	}
	gens := opts.Generations
	if gens <= 0 {
		gens = 1000
	}
	seed := opts.Seed
	if seed == 0 {
		seed = 42
	}
	g := core.NewGenerator(seed)

	// The bounds hoisted into plain slices: the generation loops read
	// them twice per coordinate per trial, and the accessor walk would
	// pay the dtype dispatch that many times. The elements are the
	// ones FloatAt returned, so the run is bit for bit the same.
	lo := make([]float64, n)
	hi := make([]float64, n)
	for i := range n {
		lo[i] = lower.FloatAt(i)
		hi[i] = upper.FloatAt(i)
	}

	eval := func(x []float64) (float64, error) {
		arr, err := core.FromFloats(x, n)
		if err != nil {
			return 0, base.Errf("%s: %w", name, err)
		}
		v, err := f(arr)
		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
	}

	// Latin-square-ish start: uniform draws inside the box.
	popX := make([][]float64, pop)
	popF := make([]float64, pop)
	for p := range pop {
		x := make([]float64, n)
		for i := range n {
			x[i] = lo[i] + g.Unit()*(hi[i]-lo[i])
		}
		v, err := eval(x)
		if err != nil {
			return nil, 0, err
		}
		popX[p], popF[p] = x, v
	}

	// pick draws population indices until one avoids the excluded
	// candidates; the excludes arrive as plain values, so the hot draw
	// loop allocates nothing. A negative exclude never matches, which
	// is how the callers drop the slots they do not need.
	pick := func(avoid1, avoid2, avoid3 int) int {
		for {
			c := int(g.Unit() * float64(pop))
			if c >= 0 && c < pop && c != avoid1 && c != avoid2 && c != avoid3 {
				return c
			}
		}
	}

	trial := make([]float64, n)
	for gen := 0; gen < gens; gen++ {
		for target := range pop {
			r1 := pick(target, -1, -1)
			r2 := pick(target, r1, -1)
			r3 := pick(target, r1, r2)
			jr := int(g.Unit()*float64(n)) % n
			for i := range n {
				if g.Unit() < cr || i == jr {
					m := popX[r1][i] + fw*(popX[r2][i]-popX[r3][i])
					// Clamp, never reflect: the bounds are the contract.
					m = math.Min(math.Max(m, lo[i]), hi[i])
					trial[i] = m
				} else {
					trial[i] = popX[target][i]
				}
			}
			v, err := eval(trial)
			if err != nil {
				return nil, 0, err
			}
			if v <= popF[target] {
				copy(popX[target], trial)
				popF[target] = v
			}
		}
	}
	best := 0
	for p := 1; p < pop; p++ {
		if popF[p] < popF[best] {
			best = p
		}
	}
	out, err := core.FromFloats(popX[best], n)
	if err != nil {
		return nil, 0, base.Errf("%s: %w", name, err)
	}
	return out, popF[best], nil
}