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

package optim_test

import (
	"fmt"
	"math"

	"sourcedock.dev/petrbalvin/tensor"
	"sourcedock.dev/petrbalvin/tensor/optim"
)

// ExampleFindRoot brackets the root of a cubic between 2 and 3. The
// bracket changes sign, which is all Brent's method needs.
func ExampleFindRoot() {
	f := func(x float64) float64 { return x*x*x - 2*x - 5 }
	root, err := optim.FindRoot(f, 2, 3, 1e-12)
	if err != nil {
		fmt.Println("root finding failed:", err)
		return
	}
	fmt.Printf("root = %.10f\n", root)
	fmt.Printf("residual = %.3e\n", math.Abs(f(root)))
	// Output:
	// root = 2.0945514815
	// residual = 3.553e-15
}

// ExampleFindRootSystem solves a coupled pair of equations,
// x0 + x1 = 3 and x0² − x1 = 1, from the starting guess (1.5, 1.5).
func ExampleFindRootSystem() {
	system := func(x *tensor.Array) (*tensor.Array, error) {
		x0, x1 := x.FloatAt(0), x.FloatAt(1)
		return tensor.FromFloats([]float64{x0 + x1 - 3, x0*x0 - x1 - 1}, 2)
	}
	x0, _ := tensor.FromFloats([]float64{1.5, 1.5}, 2)
	x, residual, err := optim.FindRootSystem(system, x0, optim.RootSystemOptions{})
	if err != nil {
		fmt.Println("system solve failed:", err)
		return
	}
	fmt.Printf("x = (%.6f, %.6f)\n", x.FloatAt(0), x.FloatAt(1))
	fmt.Printf("residual = %.3e\n", residual)
	// Output:
	// x = (1.561553, 1.438447)
	// residual = 4.841e-14
}

// ExampleLevenbergMarquardt fits the exponential decay y = a·exp(−b·t)
// to seven perturbed samples, recovering both parameters from the
// samples alone. The generator is seeded, so the data and the fit are
// reproduced exactly on every run.
func ExampleLevenbergMarquardt() {
	const wantA, wantB = 2.5, 1.3
	g := tensor.NewGenerator(7)
	ts := make([]float64, 7)
	ys := make([]float64, 7)
	for i := range ts {
		ts[i] = 0.5 * float64(i)
		ys[i] = wantA*math.Exp(-wantB*ts[i]) + 0.02*g.NormalUnit()
	}

	// The residual the fit drives to zero holds one entry per
	// observation: the model at the parameters minus the sample.
	residual := func(p *tensor.Array) (*tensor.Array, error) {
		out := make([]float64, len(ts))
		for i, t := range ts {
			out[i] = p.FloatAt(0)*math.Exp(-p.FloatAt(1)*t) - ys[i]
		}
		return tensor.FromFloats(out, len(out))
	}

	p0, _ := tensor.FromFloats([]float64{1, 1}, 2)
	p, chi2, err := optim.LevenbergMarquardt(residual, p0, optim.LMOptions{})
	if err != nil {
		fmt.Println("fit failed:", err)
		return
	}
	fmt.Printf("a = %.4f (want %.4f)\n", p.FloatAt(0), wantA)
	fmt.Printf("b = %.4f (want %.4f)\n", p.FloatAt(1), wantB)
	fmt.Printf("chi2 = %.6f\n", chi2)
	// Output:
	// a = 2.5325 (want 2.5000)
	// b = 1.2978 (want 1.3000)
	// chi2 = 0.000404
}

// ExampleMinimiseLBFGS minimises the Rosenbrock function inside a box,
// from a start that lies outside it: the iterate is projected onto the
// box rather than refused.
func ExampleMinimiseLBFGS() {
	rosenbrock := func(x *tensor.Array) (float64, error) {
		xx, yy := x.FloatAt(0), x.FloatAt(1)
		return 100*(yy-xx*xx)*(yy-xx*xx) + (1-xx)*(1-xx), nil
	}
	gradient := func(x *tensor.Array) (*tensor.Array, error) {
		xx, yy := x.FloatAt(0), x.FloatAt(1)
		return tensor.FromFloats([]float64{
			-400*xx*(yy-xx*xx) - 2*(1-xx),
			200 * (yy - xx*xx),
		}, 2)
	}
	x0, _ := tensor.FromFloats([]float64{-3, 5}, 2)
	opts := optim.LBFGSOptions{
		Lower: []float64{-2, -1},
		Upper: []float64{2, 3},
	}
	x, f, err := optim.MinimiseLBFGS(rosenbrock, gradient, x0, opts)
	if err != nil {
		fmt.Println("minimisation failed:", err)
		return
	}
	fmt.Printf("x = (%.4f, %.4f)\n", x.FloatAt(0), x.FloatAt(1))
	fmt.Printf("f = %.3e\n", f)
	// Output:
	// x = (1.0000, 1.0000)
	// f = 3.369e-21
}

// ExampleMinimiseConstrained minimises the distance to (1, 1) subject
// to the equality row x0 + x1 = 1, which the unconstrained answer
// violates. The augmented Lagrangian pulls the iterate onto the row.
func ExampleMinimiseConstrained() {
	objective := func(x *tensor.Array) (float64, error) {
		d0, d1 := x.FloatAt(0)-1, x.FloatAt(1)-1
		return d0*d0 + d1*d1, nil
	}
	gradient := func(x *tensor.Array) (*tensor.Array, error) {
		return tensor.FromFloats([]float64{2 * (x.FloatAt(0) - 1), 2 * (x.FloatAt(1) - 1)}, 2)
	}
	a, _ := tensor.FromFloats([]float64{1, 1}, 1, 2)
	cons := optim.LinearConstraints{
		A:     a,
		Lower: []float64{1},
		Upper: []float64{1},
	}
	x0, _ := tensor.FromFloats([]float64{0, 0}, 2)
	x, f, err := optim.MinimiseConstrained(objective, gradient, x0, cons, optim.LBFGSOptions{})
	if err != nil {
		fmt.Println("minimisation failed:", err)
		return
	}
	fmt.Printf("x = (%.4f, %.4f)\n", x.FloatAt(0), x.FloatAt(1))
	fmt.Printf("x0+x1 = %.4f\n", x.FloatAt(0)+x.FloatAt(1))
	fmt.Printf("f = %.4f\n", f)
	// Output:
	// x = (0.5000, 0.5000)
	// x0+x1 = 1.0000
	// f = 0.5000
}

// ExampleMinimiseLinearRows solves a linear program in the house
// two-sided form: maximise 3x0 + 2x1, written as the minimisation of
// its negation, subject to x0 + x1 ≤ 4, x0 + 3x1 ≤ 6 and x0, x1 ≥ 0.
func ExampleMinimiseLinearRows() {
	cost, _ := tensor.FromFloats([]float64{-3, -2}, 2)
	a, _ := tensor.FromFloats([]float64{
		1, 1,
		1, 3,
		1, 0,
		0, 1,
	}, 4, 2)
	cons := optim.LinearConstraints{
		A: a,
		Lower: []float64{
			math.Inf(-1), // -∞ ≤ x0 + x1
			math.Inf(-1), // -∞ ≤ x0 + 3x1
			0,            // 0 ≤ x0
			0,            // 0 ≤ x1
		},
		Upper: []float64{
			4, 6, // x0 + x1 ≤ 4, x0 + 3x1 ≤ 6
			math.Inf(1), math.Inf(1),
		},
	}
	x, value, err := optim.MinimiseLinearRows(cost, cons, optim.LinearProgramOptions{})
	if err != nil {
		fmt.Println("linear program failed:", err)
		return
	}
	fmt.Printf("x = (%.4f, %.4f)\n", x.FloatAt(0), x.FloatAt(1))
	fmt.Printf("c·x = %.4f\n", value)
	// Output:
	// x = (4.0000, 0.0000)
	// c·x = -12.0000
}

// ExampleMinimiseDifferentialEvolution finds the global minimum of the
// two-dimensional Rastrigin function, a landscape of many local minima
// around one global minimum at the origin. The seed makes the run
// reproducible.
func ExampleMinimiseDifferentialEvolution() {
	rastrigin := func(x *tensor.Array) (float64, error) {
		sum := 20.0
		for i := range 2 {
			v := x.FloatAt(i)
			sum += v*v - 10*math.Cos(2*math.Pi*v)
		}
		return sum, nil
	}
	lower, _ := tensor.FromFloats([]float64{-5.12, -5.12}, 2)
	upper, _ := tensor.FromFloats([]float64{5.12, 5.12}, 2)
	opts := optim.DifferentialEvolutionOptions{Seed: 7, Population: 60, Generations: 400}
	x, f, err := optim.MinimiseDifferentialEvolution(rastrigin, lower, upper, opts)
	if err != nil {
		fmt.Println("global search failed:", err)
		return
	}
	// The search reaches the origin. The residual coordinate is a
	// round-off of the population's spread, so its sign is not fixed
	// by the algorithm; the distance to the minimum is what the run
	// guarantees.
	fmt.Printf("f = %.6f\n", f)
	fmt.Printf("|x| < 1e-6: %v\n", math.Hypot(x.FloatAt(0), x.FloatAt(1)) < 1e-6)
	// Output:
	// f = 0.000000
	// |x| < 1e-6: true
}