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

package optim

import (
	"math"
	"testing"

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

// constrainedBowl returns the separable bowl shifted to the given
// centre, the standard workhorse for verifying constraint handling:
// the unconstrained answer is known and the constrained one is a
// projection of it that admits an analytic check.
func constrainedBowl(cx, cy float64) func(*core.Array) (float64, error) {
	return func(p *core.Array) (float64, error) {
		dx, dy := p.FloatAt(0)-cx, p.FloatAt(1)-cy
		return dx*dx + dy*dy, nil
	}
}

// TestConstrainedEquality drives the augmented Lagrangian through an
// equality hyperplane that cuts the bowl's minimum: the constrained
// optimum is the projection of the unconstrained one onto the plane.
func TestConstrainedEquality(t *testing.T) {
	// min (x-1)² + (y-1)² subject to x + y = 2: the plane passes
	// through the unconstrained minimum, so the answer is (1, 1) with
	// value 0 and the multiplier converges to zero.
	A, _ := core.FromFloats([]float64{1, 1}, 1, 2)
	point, value, err := MinimiseConstrained(constrainedBowl(1, 1), nil,
		mustLin(t, []float64{5, -3}), LinearConstraints{A: A, Lower: []float64{2}, Upper: []float64{2}},
		LBFGSOptions{Tolerance: 1e-10})
	if err != nil {
		t.Fatalf("MinimiseConstrained: %v", err)
	}
	if math.Abs(point.FloatAt(0)-1) > 1e-4 || math.Abs(point.FloatAt(1)-1) > 1e-4 {
		t.Fatalf("point = (%.8g, %.8g), want (1, 1)", point.FloatAt(0), point.FloatAt(1))
	}
	if value > 1e-6 {
		t.Fatalf("value = %.8g, want 0", value)
	}
}

// TestConstrainedInequality forces an active inequality: the bowl's
// minimum at (2, -1) lies beyond x + y = 0, so the constrained
// optimum sits on the wall at (1.5, -1.5) with value 0.5.
func TestConstrainedInequality(t *testing.T) {
	A, _ := core.FromFloats([]float64{1, 1}, 1, 2)
	gradFn := func(p *core.Array) (*core.Array, error) {
		out := core.New(core.Float, 2)
		out.RawFloats()[0] = 2 * (p.FloatAt(0) - 2)
		out.RawFloats()[1] = 2 * (p.FloatAt(1) + 1)
		return out, nil
	}
	for _, grad := range []func(*core.Array) (*core.Array, error){nil, gradFn} {
		point, value, err := MinimiseConstrained(constrainedBowl(2, -1), grad,
			mustLin(t, []float64{4, 4}), LinearConstraints{A: A, Lower: []float64{math.Inf(-1)}, Upper: []float64{0}},
			LBFGSOptions{Tolerance: 1e-10})
		if err != nil {
			t.Fatalf("MinimiseConstrained: %v", err)
		}
		if math.Abs(point.FloatAt(0)-1.5) > 1e-4 || math.Abs(point.FloatAt(1)+1.5) > 1e-4 {
			t.Fatalf("point = (%.8g, %.8g), want (1.5, -1.5)", point.FloatAt(0), point.FloatAt(1))
		}
		if math.Abs(value-0.5) > 1e-6 {
			t.Fatalf("value = %.8g, want 0.5", value)
		}
	}
}

// TestConstrainedTwoSided clamps the bowl between two parallel walls:
// 1 ≤ x ≤ 2 drags the first coordinate to the upper wall and leaves
// the second coordinate free.
func TestConstrainedTwoSided(t *testing.T) {
	A, _ := core.FromFloats([]float64{1, 0}, 1, 2)
	point, _, err := MinimiseConstrained(constrainedBowl(3, 1), nil,
		mustLin(t, []float64{0, 0}), LinearConstraints{A: A, Lower: []float64{1}, Upper: []float64{2}},
		LBFGSOptions{})
	if err != nil {
		t.Fatalf("MinimiseConstrained: %v", err)
	}
	if math.Abs(point.FloatAt(0)-2) > 1e-4 {
		t.Fatalf("first coordinate = %.8g, want 2 on the upper wall", point.FloatAt(0))
	}
	if math.Abs(point.FloatAt(1)-1) > 1e-4 {
		t.Fatalf("second coordinate = %.8g, want 1", point.FloatAt(1))
	}
}

// TestConstrainedWithBox composes linear rows with the box walls from
// the options: x + y = 0 pins the pair to a line and y ≥ -0.5 cuts
// the line at x = 0.5. The bowl around (1, 0) then bottoms out at
// (0.5, -0.5) with value 0.5, both terms contributing equally.
func TestConstrainedWithBox(t *testing.T) {
	A, _ := core.FromFloats([]float64{1, 1}, 1, 2)
	point, value, err := MinimiseConstrained(constrainedBowl(1, 0), nil,
		mustLin(t, []float64{2, 2}),
		LinearConstraints{A: A, Lower: []float64{0}, Upper: []float64{0}},
		LBFGSOptions{Tolerance: 1e-10, Lower: []float64{math.Inf(-1), -0.5}})
	if err != nil {
		t.Fatalf("MinimiseConstrained: %v", err)
	}
	if math.Abs(point.FloatAt(0)-0.5) > 1e-4 || math.Abs(point.FloatAt(1)+0.5) > 1e-4 {
		t.Fatalf("point = (%.8g, %.8g), want (0.5, -0.5)", point.FloatAt(0), point.FloatAt(1))
	}
	if math.Abs(value-0.5) > 1e-5 {
		t.Fatalf("value = %.8g, want 0.5", value)
	}
}

// TestConstrainedRefusals checks the loud rejections: nil matrix,
// wrong shape, bound or coefficient nonsense.
func TestConstrainedRefusals(t *testing.T) {
	start := mustLin(t, []float64{0, 0})
	A, _ := core.FromFloats([]float64{1, 1}, 1, 2)
	cases := []struct {
		name string
		cons LinearConstraints
	}{
		{"nil matrix", LinearConstraints{Lower: []float64{0}, Upper: []float64{1}}},
		{"wrong shape", LinearConstraints{A: mustLin(t, []float64{1, 1}), Lower: []float64{0}, Upper: []float64{1}}},
		{"short bounds", LinearConstraints{A: A, Lower: []float64{0}, Upper: []float64{}}},
		{"crossed bounds", LinearConstraints{A: A, Lower: []float64{1}, Upper: []float64{0}}},
	}
	for _, c := range cases {
		if _, _, err := MinimiseConstrained(constrainedBowl(0, 0), nil, start, c.cons, LBFGSOptions{}); err == nil {
			t.Fatalf("%s accepted", c.name)
		}
	}
	nan, _ := core.FromFloats([]float64{math.NaN(), 1}, 1, 2)
	if _, _, err := MinimiseConstrained(constrainedBowl(0, 0), nil, start,
		LinearConstraints{A: nan, Lower: []float64{0}, Upper: []float64{1}}, LBFGSOptions{}); err == nil {
		t.Fatal("NaN coefficient accepted")
	}
}

// mustLin builds a float vector for the constraint tests.
func mustLin(t *testing.T, vals []float64, shape ...int) *core.Array {
	t.Helper()
	if len(shape) == 0 {
		shape = []int{len(vals)}
	}
	a, err := core.FromFloats(vals, shape...)
	if err != nil {
		t.Fatalf("FromFloats(%v, %v): %v", vals, shape, err)
	}
	return a
}