package optim

import (
	"math"
	"testing"

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

// Solver-coverage pins: every test drives an optimiser on a problem
// class with a closed-form answer the existing fixtures do not build,
// from corner optima over degenerate active sets to mixed nonlinear
// constraints.

func mustMatrix(t *testing.T, vals []float64, r, c int) *core.Array {
	t.Helper()
	a, err := core.FromFloats(vals, r, c)
	if err != nil {
		t.Fatal(err)
	}
	return a
}

// TestLBFGSCornerOptimumWithEqualityBound minimises
// f = (x-3)^2 + (y+1)^2 + (z-2)^2 over x <= 1 with z frozen at 2: the
// constrained optimum (1, -1, 2) keeps the wall term at 4, so 4 is the
// answer the projected-gradient machinery must report.
func TestLBFGSCornerOptimumWithEqualityBound(t *testing.T) {
	f := func(a *core.Array) (float64, error) {
		x, y, z := a.FloatAt(0), a.FloatAt(1), a.FloatAt(2)
		return (x-3)*(x-3) + (y+1)*(y+1) + (z-2)*(z-2), nil
	}
	grad := func(a *core.Array) (*core.Array, error) {
		x, y, z := a.FloatAt(0), a.FloatAt(1), a.FloatAt(2)
		return core.FromFloats([]float64{2 * (x - 3), 2 * (y + 1), 2 * (z - 2)}, 3)
	}
	x0, _ := core.FromFloats([]float64{0, 0, 2}, 3)
	opts := LBFGSOptions{
		Tolerance: 1e-12,
		Lower:     []float64{math.Inf(-1), math.Inf(-1), 2},
		Upper:     []float64{1, math.Inf(1), 2},
	}
	p, fv, err := MinimiseLBFGS(f, grad, x0, opts)
	if err != nil {
		t.Fatal(err)
	}
	if math.Abs(fv-4) > 1e-10 {
		t.Fatalf("f = %g, want 4", fv)
	}
	if math.Abs(p.FloatAt(0)-1) > 1e-9 || math.Abs(p.FloatAt(1)+1) > 1e-9 || p.FloatAt(2) != 2 {
		t.Fatalf("point (%g, %g, %g), want (1, -1, 2)", p.FloatAt(0), p.FloatAt(1), p.FloatAt(2))
	}
}

// TestLBFGSCornerOptimumFiniteDifference repeats the corner problem
// without a gradient callback, so the one-sided difference stencils at
// the walls carry the run.
func TestLBFGSCornerOptimumFiniteDifference(t *testing.T) {
	f := func(a *core.Array) (float64, error) {
		x, y := a.FloatAt(0), a.FloatAt(1)
		return (x-3)*(x-3) + (y+1)*(y+1), nil
	}
	x0, _ := core.FromFloats([]float64{0, 0}, 2)
	opts := LBFGSOptions{
		Tolerance:     1e-10,
		Lower:         []float64{-2, math.Inf(-1)},
		Upper:         []float64{1, math.Inf(1)},
		MaxIterations: 500,
	}
	p, fv, err := MinimiseLBFGS(f, nil, x0, opts)
	if err != nil {
		t.Fatal(err)
	}
	if math.Abs(fv-4) > 1e-8 {
		t.Fatalf("f = %g, want 4", fv)
	}
	if math.Abs(p.FloatAt(0)-1) > 1e-5 || math.Abs(p.FloatAt(1)+1) > 1e-5 {
		t.Fatalf("point (%g, %g), want (1, -1)", p.FloatAt(0), p.FloatAt(1))
	}
}

// TestQPDegenerateRatioTies builds a problem whose ratio test ties
// three rows at one point and whose release cycle must then free a
// slack row: min 1/2(x^2+y^2) - x - y over x+y >= 1, x >= 1/2,
// y >= 1/2. All rows block at (1/2, 1/2); the optimum is (1, 1) with
// the first row slack and a zero multiplier.
func TestQPDegenerateRatioTies(t *testing.T) {
	h, _ := core.FromFloats([]float64{1, 0, 0, 1}, 2, 2)
	c, _ := core.FromFloats([]float64{-1, -1}, 2)
	cons := LinearConstraints{
		A:     mustMatrix(t, []float64{1, 1, 1, 0, 0, 1}, 3, 2),
		Lower: []float64{1, 0.5, 0.5},
		Upper: []float64{math.Inf(1), math.Inf(1), math.Inf(1)},
	}
	x, fv, multipliers, err := MinimiseQP(h, c, cons, nil, QPOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if math.Abs(x.FloatAt(0)-1) > 1e-8 || math.Abs(x.FloatAt(1)-1) > 1e-8 {
		t.Fatalf("point (%g, %g), want (1, 1)", x.FloatAt(0), x.FloatAt(1))
	}
	if math.Abs(fv+1) > 1e-10 {
		t.Fatalf("f = %g, want -1", fv)
	}
	if len(multipliers) != 3 {
		t.Fatalf("multipliers %v", multipliers)
	}
	if multipliers[0] > 1e-10 {
		t.Fatalf("slack row multiplier %g, want 0", multipliers[0])
	}
}

// TestCMAESRotatedIllConditionedQuadratic runs the strategy on a valley
// rotated 45 degrees with a conditioning of 100, where a diagonal
// sampler cannot turn and the full covariance must.
func TestCMAESRotatedIllConditionedQuadratic(t *testing.T) {
	f := func(a *core.Array) (float64, error) {
		x, y := a.FloatAt(0), a.FloatAt(1)
		u, v := (x+y)/math.Sqrt2, (x-y)/math.Sqrt2
		return u*u + 100*v*v, nil
	}
	x0, _ := core.FromFloats([]float64{5, -5}, 2)
	p, fv, err := MinimiseCMAES(f, x0, CMAESOptions{Generations: 2000, Tolerance: 1e-10})
	if err != nil {
		t.Fatal(err)
	}
	if fv > 1e-8 {
		t.Fatalf("f = %g, want ~0", fv)
	}
	if math.Abs(p.FloatAt(0)) > 1e-3 || math.Abs(p.FloatAt(1)) > 1e-3 {
		t.Fatalf("point (%g, %g), want ~(0, 0)", p.FloatAt(0), p.FloatAt(1))
	}
}

// TestDifferentialEvolutionBowl drives the rand/1/bin scheme on a
// separable bowl to full precision against the seeded bounds.
func TestDifferentialEvolutionBowl(t *testing.T) {
	f := func(a *core.Array) (float64, error) {
		s := 0.0
		for i := range a.Len() {
			d := a.FloatAt(i) - float64(i+1)
			s += d * d
		}
		return s, nil
	}
	lower, _ := core.FromFloats([]float64{-10, -10, -10}, 3)
	upper, _ := core.FromFloats([]float64{10, 10, 10}, 3)
	p, fv, err := MinimiseDifferentialEvolution(f, lower, upper, DifferentialEvolutionOptions{
		Generations: 600, Seed: 7,
	})
	if err != nil {
		t.Fatal(err)
	}
	if fv > 1e-10 {
		t.Fatalf("f = %g, want ~0", fv)
	}
	for i := range 3 {
		if math.Abs(p.FloatAt(i)-float64(i+1)) > 1e-4 {
			t.Fatalf("point[%d] = %g, want %g", i, p.FloatAt(i), float64(i+1))
		}
	}
}

// TestLinearRowsRedundantEqualityExpelled adds a third equality that is
// the sum of the first two, so phase 1 must expel a redundant
// artificial through the row-drop path before phase 2 prices.
func TestLinearRowsRedundantEqualityExpelled(t *testing.T) {
	c, _ := core.FromFloats([]float64{1, 1}, 2)
	cons := LinearConstraints{
		A:     mustMatrix(t, []float64{1, 1, 1, -1, 2, 0}, 3, 2),
		Lower: []float64{2, 0, 2},
		Upper: []float64{2, 0, 2},
	}
	x, fv, err := MinimiseLinearRows(c, cons, LinearProgramOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if math.Abs(x.FloatAt(0)-1) > 1e-7 || math.Abs(x.FloatAt(1)-1) > 1e-7 {
		t.Fatalf("point (%g, %g), want (1, 1)", x.FloatAt(0), x.FloatAt(1))
	}
	if math.Abs(fv-2) > 1e-7 {
		t.Fatalf("value %g, want 2", fv)
	}
}

// TestSimplexRosenbrock walks the derivative-free simplex down the
// banana valley to its floor at (1, 1).
func TestSimplexRosenbrock(t *testing.T) {
	f := func(a *core.Array) (float64, error) {
		x, y := a.FloatAt(0), a.FloatAt(1)
		return (1-x)*(1-x) + 100*(y-x*x)*(y-x*x), nil
	}
	x0, _ := core.FromFloats([]float64{-1.2, 1}, 2)
	p, fv, err := Minimise(f, x0, MinimiseOptions{MaxIterations: 20000, Tolerance: 1e-12})
	if err != nil {
		t.Fatal(err)
	}
	if fv > 1e-12 {
		t.Fatalf("f = %g, want ~0", fv)
	}
	if math.Abs(p.FloatAt(0)-1) > 1e-4 || math.Abs(p.FloatAt(1)-1) > 1e-4 {
		t.Fatalf("point (%g, %g), want (1, 1)", p.FloatAt(0), p.FloatAt(1))
	}
}

// TestRootSystemBroyden solves a mildly nonlinear system with the
// maintained inverse and verifies both equations at the answer.
func TestRootSystemBroyden(t *testing.T) {
	f := func(x *core.Array) (*core.Array, error) {
		a, b := x.FloatAt(0), x.FloatAt(1)
		return core.FromFloats([]float64{
			a*a + b*b - 4,
			math.Exp(a) + b - 3,
		}, 2)
	}
	x0, _ := core.FromFloats([]float64{1, 1}, 2)
	x, res, err := FindRootSystem(f, x0, RootSystemOptions{Tolerance: 1e-11, UseBroyden: true, MaxIterations: 200})
	if err != nil {
		t.Fatal(err)
	}
	if res > 1e-11 {
		t.Fatalf("residual %g", res)
	}
	a, b := x.FloatAt(0), x.FloatAt(1)
	if math.Abs(a*a+b*b-4) > 1e-9 || math.Abs(math.Exp(a)+b-3) > 1e-9 {
		t.Fatalf("solution (%g, %g) does not satisfy the system", a, b)
	}
}

// TestLevenbergMarquardtWeightedCovariance fits a line under per-point
// variances and checks the parameters and the reported covariance.
func TestLevenbergMarquardtSigmaWeightsAndCovariance(t *testing.T) {
	xs := []float64{0, 1, 2, 3, 4}
	ys := []float64{0.5, 2.49, 4.52, 6.48, 8.51}
	residual := func(p *core.Array) (*core.Array, error) {
		out := core.New(core.Float, len(xs))
		v := out.RawFloats()
		for i, x := range xs {
			v[i] = p.FloatAt(0)*x + p.FloatAt(1) - ys[i]
		}
		return out, nil
	}
	sigma, _ := core.FromFloats([]float64{1, 1, 4, 1, 4}, 5)
	p0, _ := core.FromFloats([]float64{0, 0}, 2)
	res, err := LevenbergMarquardtFit(residual, p0, LMOptions{Sigma: sigma, RequestCovariance: true})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != FitConverged {
		t.Fatalf("status %v", res.Status)
	}
	if math.Abs(res.Parameters.FloatAt(0)-2) > 0.05 || math.Abs(res.Parameters.FloatAt(1)-0.5) > 0.05 {
		t.Fatalf("parameters (%g, %g), want ~(2, 0.5)",
			res.Parameters.FloatAt(0), res.Parameters.FloatAt(1))
	}
	if res.Covariance == nil {
		t.Fatal("covariance missing")
	}
}

// TestNonlinearConstrainedMixedRows minimises a bowl over an equality
// and an inequality at once, with the inequality active at the answer:
// min (x-2)^2 + (y-2)^2 over x + y = 1 and x <= 1/2 sits at
// (1/2, 1/2) with f = 4.5 and a non-negative inequality multiplier.
func TestNonlinearConstrainedMixedRows(t *testing.T) {
	f := func(a *core.Array) (float64, error) {
		x, y := a.FloatAt(0), a.FloatAt(1)
		return (x-2)*(x-2) + (y-2)*(y-2), nil
	}
	grad := func(a *core.Array) (*core.Array, error) {
		x, y := a.FloatAt(0), a.FloatAt(1)
		return core.FromFloats([]float64{2 * (x - 2), 2 * (y - 2)}, 2)
	}
	cons := NonlinearConstraints{
		Equalities: []func(*core.Array) (float64, error){func(a *core.Array) (float64, error) {
			return a.FloatAt(0) + a.FloatAt(1) - 1, nil
		}},
		Inequalities: []func(*core.Array) (float64, error){func(a *core.Array) (float64, error) {
			return a.FloatAt(0) - 0.5, nil
		}},
	}
	x0, _ := core.FromFloats([]float64{0, 0}, 2)
	p, fv, mult, err := MinimiseNonlinearConstrained(f, grad, x0, cons, LBFGSOptions{Tolerance: 1e-10})
	if err != nil {
		t.Fatal(err)
	}
	if math.Abs(p.FloatAt(0)-0.5) > 1e-4 || math.Abs(p.FloatAt(1)-0.5) > 1e-4 {
		t.Fatalf("point (%g, %g), want (0.5, 0.5)", p.FloatAt(0), p.FloatAt(1))
	}
	if fv < 4.49 || fv > 4.51 {
		t.Fatalf("f = %g, want 4.5", fv)
	}
	if len(mult) != 2 || mult[1] < 0 {
		t.Fatalf("multipliers %v", mult)
	}
}

// TestSimulatedAnnealingTwoWell puts a deep left well against a
// shallower right one, starting in the right: the Metropolis walk must
// cross the barrier and report the global basin.
func TestSimulatedAnnealingTwoWell(t *testing.T) {
	f := func(a *core.Array) (float64, error) {
		x := a.FloatAt(0)
		w1 := (x + 2) * (x + 2)
		w2 := (x-3)*(x-3) + 0.5
		return math.Min(w1, w2), nil
	}
	x0, _ := core.FromFloats([]float64{3}, 1)
	p, fv, err := MinimiseSimulatedAnnealing(f, x0, SimulatedAnnealingOptions{
		Steps: 40000, Seed: 5, StepScale: 0.5, AllowBudgetExit: true,
	})
	if err != nil {
		t.Fatal(err)
	}
	if fv > 0.1 {
		t.Fatalf("f = %g, want the left well (~0)", fv)
	}
	if p.FloatAt(0) > 0 {
		t.Fatalf("point %g, want the left well", p.FloatAt(0))
	}
}