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

package optim

import (
	"math"
	"strings"
	"testing"

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

// qpBowl builds the canonical data of ½xᵀHx + c·x for the squared
// distance to centre: H = 2I and c = −2·centre, the shape most of the
// hand-solved pins below use.
func qpBowl(t *testing.T, cx, cy float64) (*core.Array, *core.Array) {
	t.Helper()
	h := mustFloats(t, []float64{2, 0, 0, 2}, 2, 2)
	c := mustFloats(t, []float64{-2 * cx, -2 * cy})
	return h, c
}

// TestMinimiseQPActiveUpperWall pins the analytic case: the bowl's
// unconstrained minimum (2, 2) lies beyond x + y ≤ 2, the constrained
// optimum is the wall point (1, 1) with value −6 in the canonical
// form, and the row's multiplier is the hand-solved 2, positive as the
// KKT conditions demand for an active upper wall.
func TestMinimiseQPActiveUpperWall(t *testing.T) {
	h, c := qpBowl(t, 2, 2)
	cons := LinearConstraints{A: mustFloats(t, []float64{1, 1}, 1, 2),
		Lower: []float64{math.Inf(-1)}, Upper: []float64{2}}
	for _, x0 := range []*core.Array{nil, mustFloats(t, []float64{0, 0})} {
		x, value, multipliers, err := MinimiseQP(h, c, cons, x0, QPOptions{})
		if err != nil {
			t.Fatalf("MinimiseQP: %v", err)
		}
		if math.Abs(x.FloatAt(0)-1) > 1e-8 || math.Abs(x.FloatAt(1)-1) > 1e-8 {
			t.Fatalf("point = (%.10g, %.10g), want (1, 1)", x.FloatAt(0), x.FloatAt(1))
		}
		if math.Abs(value+6) > 1e-8 {
			t.Fatalf("value = %.12g, want −6", value)
		}
		if len(multipliers) != 1 || math.Abs(multipliers[0]-2) > 1e-7 {
			t.Fatalf("multipliers = %v, want [2]", multipliers)
		}
	}
}

// TestMinimiseQPActiveLowerWall pins the lower-wall case: the bowl
// around (−3, −3) with y ≥ 0 bottoms out at (−3, 0) and the row's
// multiplier is the hand-solved 6.
func TestMinimiseQPActiveLowerWall(t *testing.T) {
	h, c := qpBowl(t, -3, -3)
	cons := LinearConstraints{A: mustFloats(t, []float64{0, 1}, 1, 2),
		Lower: []float64{0}, Upper: []float64{math.Inf(1)}}
	x, value, multipliers, err := MinimiseQP(h, c, cons, mustFloats(t, []float64{-3, 2}), QPOptions{})
	if err != nil {
		t.Fatalf("MinimiseQP: %v", err)
	}
	if math.Abs(x.FloatAt(0)+3) > 1e-8 || math.Abs(x.FloatAt(1)) > 1e-8 {
		t.Fatalf("point = (%.10g, %.10g), want (−3, 0)", x.FloatAt(0), x.FloatAt(1))
	}
	if math.Abs(value+9) > 1e-8 {
		t.Fatalf("value = %.12g, want −9", value)
	}
	if len(multipliers) != 1 || math.Abs(multipliers[0]-6) > 1e-7 {
		t.Fatalf("multipliers = %v, want [6]", multipliers)
	}
}

// TestMinimiseQPEqualityAndSlack pins complementary slackness on a
// problem with an equality row and a slack inequality: min x² + y² on
// x + y = 2 sits at (1, 1) with the signed equality multiplier −2, and
// the inactive wall y ≤ 3 must report exactly zero.
func TestMinimiseQPEqualityAndSlack(t *testing.T) {
	h := mustFloats(t, []float64{2, 0, 0, 2}, 2, 2)
	c := mustFloats(t, []float64{0, 0})
	cons := LinearConstraints{
		A:     mustFloats(t, []float64{1, 1, 0, 1}, 2, 2),
		Lower: []float64{2, math.Inf(-1)},
		Upper: []float64{2, 3},
	}
	x, value, multipliers, err := MinimiseQP(h, c, cons, nil, QPOptions{})
	if err != nil {
		t.Fatalf("MinimiseQP: %v", err)
	}
	if math.Abs(x.FloatAt(0)-1) > 1e-8 || math.Abs(x.FloatAt(1)-1) > 1e-8 {
		t.Fatalf("point = (%.10g, %.10g), want (1, 1)", x.FloatAt(0), x.FloatAt(1))
	}
	if math.Abs(value-2) > 1e-8 {
		t.Fatalf("value = %.12g, want 2", value)
	}
	if len(multipliers) != 2 {
		t.Fatalf("multipliers = %v, want one entry per row", multipliers)
	}
	if math.Abs(multipliers[0]+2) > 1e-7 {
		t.Fatalf("equality multiplier = %.12g, want −2", multipliers[0])
	}
	if multipliers[1] != 0 {
		t.Fatalf("slack inequality multiplier = %.12g, want 0", multipliers[1])
	}
}

// TestMinimiseQPRotated pins a genuinely coupled Hessian: min ½xᵀHx +
// c·x with H = [[2, 1], [1, 2]] and c = (−2, −2) on x − y ≥ ½. The
// hand-solved KKT point is (11/12, 5/12) with multiplier ¼: the
// stationarity Hx + c = ν(1, −1) and the row give three linear
// equations with exactly that solution.
func TestMinimiseQPRotated(t *testing.T) {
	h := mustFloats(t, []float64{2, 1, 1, 2}, 2, 2)
	c := mustFloats(t, []float64{-2, -2})
	cons := LinearConstraints{A: mustFloats(t, []float64{1, -1}, 1, 2),
		Lower: []float64{0.5}, Upper: []float64{math.Inf(1)}}
	x, _, multipliers, err := MinimiseQP(h, c, cons, mustFloats(t, []float64{0.5, 0}), QPOptions{})
	if err != nil {
		t.Fatalf("MinimiseQP: %v", err)
	}
	if math.Abs(x.FloatAt(0)-11.0/12.0) > 1e-8 || math.Abs(x.FloatAt(1)-5.0/12.0) > 1e-8 {
		t.Fatalf("point = (%.10g, %.10g), want (11/12, 5/12)", x.FloatAt(0), x.FloatAt(1))
	}
	if len(multipliers) != 1 || math.Abs(multipliers[0]-0.25) > 1e-7 {
		t.Fatalf("multipliers = %v, want [0.25]", multipliers)
	}
}

// TestMinimiseQPUnconstrained pins the nil-or-empty constraint set:
// one Newton step to −H⁻¹c with no multipliers.
func TestMinimiseQPUnconstrained(t *testing.T) {
	h := mustFloats(t, []float64{2, 0, 0, 4}, 2, 2)
	c := mustFloats(t, []float64{-2, -8})
	x, value, multipliers, err := MinimiseQP(h, c, LinearConstraints{}, nil, QPOptions{})
	if err != nil {
		t.Fatalf("MinimiseQP: %v", err)
	}
	if math.Abs(x.FloatAt(0)-1) > 1e-8 || math.Abs(x.FloatAt(1)-2) > 1e-8 {
		t.Fatalf("point = (%.10g, %.10g), want (1, 2)", x.FloatAt(0), x.FloatAt(1))
	}
	// ½(2·1 + 4·4) + (−2 −16) = 9 − 18 = −9.
	if math.Abs(value+9) > 1e-8 {
		t.Fatalf("value = %.12g, want −9", value)
	}
	if multipliers != nil {
		t.Fatalf("multipliers = %v, want nil", multipliers)
	}
	// The unconstrained quadratic in three variables walks the
	// Cholesky test past the first pivot: H = diag(2, 4, 6) and
	// c = (−2, −8, −18) give the analytic minimum (1, 2, 3) with
	// value −36.
	h3 := mustFloats(t, []float64{2, 0, 0, 0, 4, 0, 0, 0, 6}, 3, 3)
	c3 := mustFloats(t, []float64{-2, -8, -18})
	x3, v3, mult3, err := MinimiseQP(h3, c3, LinearConstraints{}, nil, QPOptions{})
	if err != nil {
		t.Fatalf("MinimiseQP: %v", err)
	}
	for i, w := range []float64{1, 2, 3} {
		if math.Abs(x3.FloatAt(i)-w) > 1e-8 {
			t.Fatalf("x3[%d] = %.10g, want %g", i, x3.FloatAt(i), w)
		}
	}
	if math.Abs(v3+36) > 1e-8 {
		t.Fatalf("value = %.10g, want −36", v3)
	}
	if mult3 != nil {
		t.Fatalf("multipliers = %v, want nil", mult3)
	}
	// A constraint matrix with zero rows is no constraints: valid.
	cons0 := LinearConstraints{A: core.New(core.Float, 0, 2)}
	_, _, _, err = MinimiseQP(h, c, cons0, nil, QPOptions{})
	if err != nil {
		t.Fatalf("MinimiseQP with zero rows: %v", err)
	}
}

// TestMinimiseQPCornerRows pins a two-row corner: min ½x² + 2y² −
// 3x − 3y on x + 2y ≤ 2, y ≥ 0.5. The unconstrained minimum (3, ¾)
// violates the first row, the A-face minimiser (1.75, 0.125) violates
// the second, so the answer is the corner (1, 0.5) with the
// hand-solved multipliers 2 and 3, both positive as an active corner
// requires.
func TestMinimiseQPCornerRows(t *testing.T) {
	h := mustFloats(t, []float64{1, 0, 0, 4}, 2, 2)
	c := mustFloats(t, []float64{-3, -3})
	cons := LinearConstraints{
		A:     mustFloats(t, []float64{1, 2, 0, 1}, 2, 2),
		Lower: []float64{math.Inf(-1), 0.5},
		Upper: []float64{2, math.Inf(1)},
	}
	x, value, multipliers, err := MinimiseQP(h, c, cons, mustFloats(t, []float64{0, 1}), QPOptions{})
	if err != nil {
		t.Fatalf("MinimiseQP: %v", err)
	}
	if math.Abs(x.FloatAt(0)-1) > 1e-7 || math.Abs(x.FloatAt(1)-0.5) > 1e-7 {
		t.Fatalf("point = (%.10g, %.10g), want (1, 0.5)", x.FloatAt(0), x.FloatAt(1))
	}
	canonical := 0.5*(x.FloatAt(0)*x.FloatAt(0)+4*x.FloatAt(1)*x.FloatAt(1)) - 3*x.FloatAt(0) - 3*x.FloatAt(1)
	if math.Abs(value-canonical) > 1e-12 {
		t.Fatalf("value %.12g disagrees with the canonical form %.12g", value, canonical)
	}
	if len(multipliers) != 2 {
		t.Fatalf("multipliers = %v, want one entry per row", multipliers)
	}
	if math.Abs(multipliers[0]-2) > 1e-7 || math.Abs(multipliers[1]-3) > 1e-7 {
		t.Fatalf("multipliers = (%.12g, %.12g), want (2, 3)", multipliers[0], multipliers[1])
	}
}

// TestMinimiseQPReleaseRow drives a working-set release end to end:
// min ½(x² + 10y²) − 4x − 40y on x ≤ 1, x + y ≤ 2 from the origin. The
// Newton pull (4, 4) blocks x ≤ 1 first, the face step meets x + y = 2
// at a zero-length step, and at that corner the first row's multiplier
// is negative, so it is released and the KKT point lands on the second
// row alone: x = (−16/11, 38/11) with multiplier 60/11 there and zero
// on the released row, all three figures hand-solved.
func TestMinimiseQPReleaseRow(t *testing.T) {
	h := mustFloats(t, []float64{1, 0, 0, 10}, 2, 2)
	c := mustFloats(t, []float64{-4, -40})
	cons := LinearConstraints{
		A:     mustFloats(t, []float64{1, 0, 1, 1}, 2, 2),
		Lower: []float64{math.Inf(-1), math.Inf(-1)},
		Upper: []float64{1, 2},
	}
	x, value, multipliers, err := MinimiseQP(h, c, cons, mustFloats(t, []float64{0, 0}), QPOptions{})
	if err != nil {
		t.Fatalf("MinimiseQP: %v", err)
	}
	if math.Abs(x.FloatAt(0)+16.0/11.0) > 1e-7 || math.Abs(x.FloatAt(1)-38.0/11.0) > 1e-7 {
		t.Fatalf("point = (%.10g, %.10g), want (−16/11, 38/11)", x.FloatAt(0), x.FloatAt(1))
	}
	canonical := 0.5*(x.FloatAt(0)*x.FloatAt(0)+10*x.FloatAt(1)*x.FloatAt(1)) - 4*x.FloatAt(0) - 40*x.FloatAt(1)
	if math.Abs(value-canonical) > 1e-9 {
		t.Fatalf("value %.12g disagrees with the canonical form %.12g", value, canonical)
	}
	if len(multipliers) != 2 {
		t.Fatalf("multipliers = %v, want one entry per row", multipliers)
	}
	if multipliers[0] != 0 {
		t.Fatalf("released row's multiplier = %.12g, want 0", multipliers[0])
	}
	if math.Abs(multipliers[1]-60.0/11.0) > 1e-6 {
		t.Fatalf("active row's multiplier = %.12g, want 60/11", multipliers[1])
	}
}

// TestMinimiseQPRefusals checks the entry gates: indefinite or
// singular Hessians, asymmetry, shape mismatches, infeasible rows and
// an infeasible starting point.
func TestMinimiseQPRefusals(t *testing.T) {
	goodH := mustFloats(t, []float64{2, 0, 0, 2}, 2, 2)
	goodC := mustFloats(t, []float64{-2, -2})
	cases := []struct {
		name string
		run  func() error
	}{
		{"nil Hessian", func() error {
			_, _, _, err := MinimiseQP(nil, goodC, LinearConstraints{}, nil, QPOptions{})
			return err
		}},
		{"indefinite Hessian", func() error {
			h := mustFloats(t, []float64{2, 0, 0, -2}, 2, 2)
			_, _, _, err := MinimiseQP(h, goodC, LinearConstraints{}, nil, QPOptions{})
			return err
		}},
		{"singular Hessian", func() error {
			h := mustFloats(t, []float64{2, 0, 0, 0}, 2, 2)
			_, _, _, err := MinimiseQP(h, goodC, LinearConstraints{}, nil, QPOptions{})
			return err
		}},
		{"asymmetric Hessian", func() error {
			h := mustFloats(t, []float64{2, 1, 0, 2}, 2, 2)
			_, _, _, err := MinimiseQP(h, goodC, LinearConstraints{}, nil, QPOptions{})
			return err
		}},
		{"Hessian shape", func() error {
			h := mustFloats(t, []float64{2, 0, 0, 2, 0, 0, 0, 2, 0}, 3, 3)
			_, _, _, err := MinimiseQP(h, goodC, LinearConstraints{}, nil, QPOptions{})
			return err
		}},
		{"empty cost", func() error {
			_, _, _, err := MinimiseQP(goodH, core.New(core.Float, 0), LinearConstraints{}, nil, QPOptions{})
			return err
		}},
		{"complex cost", func() error {
			_, _, _, err := MinimiseQP(goodH, mustComplexPoint(t), LinearConstraints{}, nil, QPOptions{})
			return err
		}},
		{"NaN cost", func() error {
			_, _, _, err := MinimiseQP(goodH, mustFloats(t, []float64{math.NaN(), -2}), LinearConstraints{}, nil, QPOptions{})
			return err
		}},
		{"matrix shape", func() error {
			cons := LinearConstraints{A: mustFloats(t, []float64{1, 1}), Lower: []float64{0}, Upper: []float64{1}}
			_, _, _, err := MinimiseQP(goodH, goodC, cons, nil, QPOptions{})
			return err
		}},
		{"crossed bounds", func() error {
			cons := LinearConstraints{A: mustFloats(t, []float64{1, 1}, 1, 2), Lower: []float64{3}, Upper: []float64{1}}
			_, _, _, err := MinimiseQP(goodH, goodC, cons, nil, QPOptions{})
			return err
		}},
		{"NaN coefficient", func() error {
			cons := LinearConstraints{A: mustFloats(t, []float64{math.NaN(), 1}, 1, 2), Lower: []float64{0}, Upper: []float64{1}}
			_, _, _, err := MinimiseQP(goodH, goodC, cons, nil, QPOptions{})
			return err
		}},
		{"complex Hessian", func() error {
			ch, _ := core.FromComplexes([]complex128{2, 0, 0, 2}, 2, 2)
			_, _, _, err := MinimiseQP(ch, goodC, LinearConstraints{}, nil, QPOptions{})
			return err
		}},
		{"NaN Hessian entry", func() error {
			h := mustFloats(t, []float64{2, 0, 0, math.NaN()}, 2, 2)
			_, _, _, err := MinimiseQP(h, goodC, LinearConstraints{}, nil, QPOptions{})
			return err
		}},
		{"duplicate equality rows", func() error {
			// Two identical equalities pin the same wall twice: the
			// KKT system loses rank and the refusal is honest.
			cons := LinearConstraints{A: mustFloats(t, []float64{1, 1, 1, 1}, 2, 2),
				Lower: []float64{2, 2}, Upper: []float64{2, 2}}
			_, _, _, err := MinimiseQP(goodH, goodC, cons, mustFloats(t, []float64{1, 1}), QPOptions{})
			return err
		}},
		{"short bounds with rows", func() error {
			cons := LinearConstraints{A: mustFloats(t, []float64{1, 1, 0, 1}, 2, 2),
				Lower: []float64{0}, Upper: []float64{1}}
			_, _, _, err := MinimiseQP(goodH, goodC, cons, nil, QPOptions{})
			return err
		}},
		{"complex constraint matrix", func() error {
			ca, _ := core.FromComplexes([]complex128{1, 1, 1, 1}, 2, 2)
			_, _, _, err := MinimiseQP(goodH, goodC, LinearConstraints{A: ca, Lower: []float64{0, 0}, Upper: []float64{1, 1}}, nil, QPOptions{})
			return err
		}},
		{"infeasible rows", func() error {
			cons := LinearConstraints{
				A:     mustFloats(t, []float64{1, 0, 1, 0}, 2, 2),
				Lower: []float64{1, math.Inf(-1)},
				Upper: []float64{math.Inf(1), 0},
			}
			_, _, _, err := MinimiseQP(goodH, goodC, cons, nil, QPOptions{})
			return err
		}},
		{"infeasible start", func() error {
			cons := LinearConstraints{A: mustFloats(t, []float64{1, 1}, 1, 2),
				Lower: []float64{math.Inf(-1)}, Upper: []float64{1}}
			_, _, _, err := MinimiseQP(goodH, goodC, cons, mustFloats(t, []float64{1, 1}), QPOptions{})
			return err
		}},
		{"complex start", func() error {
			_, _, _, err := MinimiseQP(goodH, goodC, LinearConstraints{}, mustComplexPoint(t), QPOptions{})
			return err
		}},
		{"start length", func() error {
			_, _, _, err := MinimiseQP(goodH, goodC, LinearConstraints{}, mustFloats(t, []float64{1}), QPOptions{})
			return err
		}},
	}
	for _, c := range cases {
		if err := c.run(); err == nil {
			t.Fatalf("%s accepted", c.name)
		}
	}
}

// TestMinimiseQPBudget pins the honest refusal when the active-set
// budget is too small for the two rounds the wall case needs, and the
// infeasibility message the linear phase propagates.
func TestMinimiseQPBudget(t *testing.T) {
	h, c := qpBowl(t, 2, 2)
	cons := LinearConstraints{A: mustFloats(t, []float64{1, 1}, 1, 2),
		Lower: []float64{math.Inf(-1)}, Upper: []float64{2}}
	_, _, _, err := MinimiseQP(h, c, cons, mustFloats(t, []float64{0, 0}), QPOptions{MaxIterations: 1})
	if err == nil || !strings.Contains(err.Error(), "budget") {
		t.Fatalf("error = %v, want the budget refusal", err)
	}
	_, _, _, err = MinimiseQP(h, c, cons, mustFloats(t, []float64{0, 0}), QPOptions{})
	if err != nil {
		t.Fatalf("MinimiseQP: %v", err)
	}
}