// 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"
)

// Convex quadratic programming by the primal active-set method on the
// house two-sided rows:
//
//	min ½xᵀHx + c·x  subject to  l ≤ A·x ≤ u.
//
// The method walks the active sets: each iterate solves the
// equality-constrained subproblem that holds the working rows exactly
// at their walls, moves along the solution until a new row blocks, and
// adds the blocker to the working set. When the step vanishes, the
// working rows' multipliers are inspected: a row whose multiplier
// turned negative was pushed into the set by an earlier wall and is
// released. The subproblem is the KKT system
//
//	[H   Ãᵀ] [p]   [−g]
//	[Ã   0] [ν] = [ 0],
//
// with g = Hx + c the gradient at the iterate and à the working rows
// each carrying the sign of the wall it holds, written as the
// constraint's own gradient: a row held at its upper wall enters as
// +a, one at its lower wall as −a (the constraint reads l − a·x ≤ 0
// there), and an equality as +a. The ν that come out are then the
// multipliers themselves: non-negative on the inequality rows, signed
// on the equalities. The system is factored by the same dense LU the
// revised simplex in simplex.go uses: one factorisation per iterate.
//
// Termination is the classic result for the method (Nocedal and
// Wright, Numerical Optimization, chapter 16.5): with H positive
// definite each subproblem has a unique solution, the objective
// strictly decreases on every non-zero step, the objective is the same
// on the zero steps that release a row, and the working set never
// repeats with a lower objective, so the loop reaches the unique KKT
// point in finitely many iterations on a nondegenerate problem.
// Degenerate ties (several rows blocking at one step, several rows
// tied at the most negative multiplier) are broken by the lowest row
// index, in the spirit of Bland's rule; a degenerate configuration
// that still cannot progress is stopped by the iteration budget as an
// error, never returned as a solution. Without positive definiteness
// none of that holds, so the Hessian is factorised at entry and a
// matrix that refuses the Cholesky factorisation is refused.

// QPOptions tunes MinimiseQP. MaxIterations ≤ 0 means 1000 active-set
// rounds, Tolerance ≤ 0 means 1e-10. The tolerance is the threshold on
// the KKT step (below it the iterate is stationary on its working set),
// on the multiplier that releases a row, and on the symmetry of H; it
// is absolute, like the other tolerances in the package.
type QPOptions struct {
	MaxIterations int
	Tolerance     float64
}

// MinimiseQP returns the point, the value ½xᵀHx + c·x and the row
// multipliers of the minimum of the strictly convex quadratic over the
// two-sided rows. The multipliers hold one entry per row of cons: an
// equality row carries its signed multiplier, an inequality row the
// non-negative multiplier of its active wall and zero when the row is
// slack, so complementary slackness reads directly off the slice.
//
// H must be symmetric and positive definite: anything else is refused
// at entry with the failed pivot named, because the termination
// guarantee and the uniqueness of the answer both rest on it. The
// starting point x0 may be nil, in which case a feasible point is
// found by running MinimiseLinearRows on the same rows with a zero
// cost, whose phase-1 refusal is the honest answer for an infeasible
// set; a supplied x0 must be feasible within 1e-9 and is refused with
// the worst violation otherwise. A nil or empty constraint set is the
// unconstrained quadratic and solves in one Newton step.
func MinimiseQP(h, c *core.Array, cons LinearConstraints, x0 *core.Array, opts QPOptions) (*core.Array, float64, []float64, error) {
	const name = "MinimiseQP"
	tol := opts.Tolerance
	if tol <= 0 {
		tol = 1e-10
	}
	maxIter := opts.MaxIterations
	if maxIter <= 0 {
		maxIter = 1000
	}
	if c.NDim() != 1 || c.Len() == 0 {
		return nil, 0, nil, base.Errf("%s: c must be a non-empty rank-1 cost vector", name)
	}
	if c.Dtype() == core.Complex {
		return nil, 0, nil, base.Errf("%s: complex costs are not supported", name)
	}
	n := c.Len()
	qc := make([]float64, n)
	for j := range n {
		v := c.FloatAt(j)
		if math.IsNaN(v) || math.IsInf(v, 0) {
			return nil, 0, nil, base.Errf("%s: the cost carries a non-finite entry at %d", name, j+1)
		}
		qc[j] = v
	}
	if h == nil {
		return nil, 0, nil, base.Errf("%s: the Hessian matrix is nil", name)
	}
	if err := requireReal(name, "Hessian matrices", h); err != nil {
		return nil, 0, nil, err
	}
	if h.NDim() != 2 || h.Shape()[0] != n || h.Shape()[1] != n {
		return nil, 0, nil, base.Errf("%s: the Hessian is %s, want %d×%d", name, base.ShapeText(h.Shape()), n, n)
	}
	hm := make([]float64, n*n)
	for i := range n {
		for j := range n {
			v := h.FloatAt(i*n + j)
			if math.IsNaN(v) || math.IsInf(v, 0) {
				return nil, 0, nil, base.Errf("%s: the Hessian carries a non-finite entry at (%d, %d)", name, i+1, j+1)
			}
			hm[i*n+j] = v
		}
	}
	for i := range n {
		for j := i + 1; j < n; j++ {
			if d := math.Abs(hm[i*n+j] - hm[j*n+i]); d > tol*math.Max(1, math.Abs(hm[i*n+j])) {
				return nil, 0, nil, base.Errf("%s: the Hessian is not symmetric at (%d, %d): %g against %g",
					name, i+1, j+1, hm[i*n+j], hm[j*n+i])
			}
		}
	}
	if err := requirePositiveDefinite(hm, n, name); err != nil {
		return nil, 0, nil, err
	}
	rows, lower, upper, equal, err := qpRows(cons, n, name)
	if err != nil {
		return nil, 0, nil, err
	}
	r := len(rows)

	// The start: a feasible point from the linear wrapper when none is
	// supplied, or the caller's own point checked against the rows.
	x := make([]float64, n)
	if x0 == nil {
		if r > 0 {
			// The zero cost makes the wrapper return any feasible
			// point; its phase-1 refusal is the honest answer for an
			// infeasible set.
			feas, _, err := MinimiseLinearRows(core.New(core.Float, n), cons, LinearProgramOptions{Tolerance: math.Max(tol, 1e-9)})
			if err != nil {
				return nil, 0, nil, base.Errf("%s: %w", name, err)
			}
			for i := range n {
				x[i] = feas.FloatAt(i)
			}
		}
	} else {
		if x0.Dtype() == core.Complex {
			return nil, 0, nil, base.Errf("%s: complex starting points are not supported", name)
		}
		if x0.Len() != n {
			return nil, 0, nil, base.Errf("%s: the starting point holds %d elements for %d variables", name, x0.Len(), n)
		}
		for i := range n {
			x[i] = x0.FloatAt(i)
		}
		if worst, row := worstViolation(rows, lower, upper, x); worst > 1e-9 {
			return nil, 0, nil, base.Errf("%s: the starting point violates row %d by %g", name, row+1, worst)
		}
	}

	// The working set: every equality row starts pinned to its wall;
	// the inequality rows join as the steps meet them. wallSign is the
	// KKT column sign: +1 on an upper wall or an equality, -1 on a
	// lower wall.
	active := make([]int, 0, r)
	wallSign := make([]float64, 0, r)
	inW := make([]bool, r)
	for i := range r {
		if equal[i] {
			active = append(active, i)
			wallSign = append(wallSign, 1)
			inW[i] = true
		}
	}

	g := make([]float64, n)
	kkt := make([]float64, (n+r)*(n+r))
	rhs := make([]float64, n+r)
	sol := make([]float64, n+r)
	// The KKT system is refactored once per iterate, so one factor
	// serves the whole run: the factorisation workspace is rewritten
	// from kkt every time.
	var fac lu

	for range maxIter {
		// The gradient at the iterate and the KKT system for the step
		// that stays stationary on the working set.
		matVec(hm, x, n, n, g)
		for j := range n {
			g[j] += qc[j]
		}
		w := len(active)
		k := n + w
		clear(rhs[:k])
		// Only the working rows' own block survives from the previous
		// iteration: the Hessian block and the two constraint blocks are
		// written whole just below, the w×w zero block of the KKT system
		// is not written at all.
		for i := n; i < k; i++ {
			for j := n; j < k; j++ {
				kkt[i*k+j] = 0
			}
		}
		for i := range n {
			for j := range n {
				kkt[i*k+j] = hm[i*n+j]
			}
			rhs[i] = -g[i]
		}
		for t := range w {
			row := rows[active[t]]
			for i := range n {
				v := wallSign[t] * row[i]
				kkt[i*k+n+t] = v
				kkt[(n+t)*k+i] = v
			}
		}
		if err := fac.factor(kkt[:k*k], k); err != nil {
			return nil, 0, nil, base.Errf("%s: the working set has lost rank: %w", name, err)
		}
		fac.solve(rhs, sol)
		pNorm := 0.0
		for j := range n {
			pNorm = math.Max(pNorm, math.Abs(sol[j]))
		}
		if pNorm <= tol*math.Max(1, maxAbs(x)) {
			// Stationary on the working set: an inequality row whose
			// multiplier turned negative belongs off the set. The most
			// negative multiplier leaves, ties to the lowest row
			// index; with none left the KKT point is reached.
			tie, worst := -1, 0.0
			for t := range w {
				if equal[active[t]] {
					continue
				}
				mu := sol[n+t]
				if mu >= -tol {
					continue
				}
				switch {
				case tie == -1 || mu < worst-1e-12:
					worst, tie = mu, t
				case mu <= worst+1e-12 && active[t] < active[tie]:
					tie = t
				}
			}
			if tie == -1 {
				var multipliers []float64
				if r > 0 {
					multipliers = make([]float64, r)
					for t := range w {
						multipliers[active[t]] = sol[n+t]
					}
				}
				if worst, row := worstViolation(rows, lower, upper, x); worst > 1e-9 {
					return nil, 0, nil, base.Errf("%s: the solution violates row %d by %g", name, row+1, worst)
				}
				value := 0.5 * dot(x, matVecR(hm, x, n))
				for j := range n {
					value += qc[j] * x[j]
				}
				out, fv := packResult(x, value)
				return out, fv, multipliers, nil
			}
			inW[active[tie]] = false
			active = append(active[:tie], active[tie+1:]...)
			wallSign = append(wallSign[:tie], wallSign[tie+1:]...)
			continue
		}
		// The largest step before an off-set row blocks: a row can only
		// be crossed towards a finite wall, and the blocker hit first,
		// ties to the lowest row index, joins the working set there.
		alpha := 1.0
		blocker, blockSign := -1, 0.0
		for i := range r {
			if inW[i] || equal[i] {
				continue
			}
			d := dot(rows[i], sol[:n])
			var ai float64
			ai = math.Inf(1)
			switch {
			case d > 0 && upper[i] < math.Inf(1):
				ai = (upper[i] - dot(rows[i], x)) / d
			case d < 0 && lower[i] > math.Inf(-1):
				ai = (lower[i] - dot(rows[i], x)) / d
			}
			if ai < 0 {
				ai = 0
			}
			if ai < 1 {
				if ai < alpha-1e-12 {
					alpha, blocker, blockSign = ai, i, wallKKTSign(d)
				} else if ai <= alpha+1e-12 && (blocker == -1 || i < blocker) {
					if ai < alpha {
						alpha = ai
					}
					blocker, blockSign = i, wallKKTSign(d)
				}
			}
		}
		for j := range n {
			x[j] += alpha * sol[j]
		}
		if blocker >= 0 && alpha < 1 {
			active = append(active, blocker)
			wallSign = append(wallSign, blockSign)
			inW[blocker] = true
		}
	}
	return nil, 0, nil, base.Errf("%s: the active-set budget of %d rounds ran out without reaching the KKT point", name, maxIter)
}

// qpRows extracts the two-sided rows into plain slices and validates
// them with the same gates MinimiseConstrained applies to its matrix.
// A nil or empty set is no rows at all.
func qpRows(cons LinearConstraints, n int, name string) (rows [][]float64, lower, upper []float64, equal []bool, err error) {
	if cons.A == nil {
		return nil, nil, nil, nil, nil
	}
	if err := requireReal(name, "constraint matrices", cons.A); err != nil {
		return nil, nil, nil, nil, err
	}
	if cons.A.NDim() != 2 || cons.A.Shape()[1] != n {
		return nil, nil, nil, nil, base.Errf("%s: the constraint matrix is %s, want r×%d", name, base.ShapeText(cons.A.Shape()), n)
	}
	r := cons.A.Shape()[0]
	if r == 0 {
		return nil, nil, nil, nil, nil
	}
	if len(cons.Lower) != r || len(cons.Upper) != r {
		return nil, nil, nil, nil, base.Errf("%s: the bounds hold %d and %d entries for %d rows",
			name, len(cons.Lower), len(cons.Upper), r)
	}
	rows = make([][]float64, r)
	lower = make([]float64, r)
	upper = make([]float64, r)
	equal = make([]bool, r)
	// One backing block for every row: the rows are read-only after
	// this build, so one allocation carries the whole block.
	back := make([]float64, r*n)
	for i := range r {
		row := back[i*n : (i+1)*n]
		for j := range n {
			v := cons.A.FloatAt(i*n + j)
			if math.IsNaN(v) || math.IsInf(v, 0) {
				return nil, nil, nil, nil, base.Errf("%s: row %d carries a non-finite coefficient", name, i+1)
			}
			row[j] = v
		}
		lo, up := cons.Lower[i], cons.Upper[i]
		if math.IsNaN(lo) || math.IsNaN(up) || lo > up {
			return nil, nil, nil, nil, base.Errf("%s: row %d has bounds [%g, %g]", name, i+1, lo, up)
		}
		rows[i], lower[i], upper[i], equal[i] = row, lo, up, lo == up
	}
	return rows, lower, upper, equal, nil
}

// worstViolation measures the rows at x and returns the largest
// violation with the row that carries it.
func worstViolation(rows [][]float64, lower, upper []float64, x []float64) (float64, int) {
	worst, row := 0.0, 0
	for i := range rows {
		ax := dot(rows[i], x)
		v := math.Max(ax-upper[i], lower[i]-ax)
		if v > worst {
			worst, row = v, i
		}
	}
	return worst, row
}

// requirePositiveDefinite runs a Cholesky factorisation as the test:
// a pivot at or below the relative floor means the matrix is singular
// or indefinite, and both refuse the problem.
func requirePositiveDefinite(hm []float64, n int, name string) error {
	scale := 1.0
	for i := range n {
		scale = math.Max(scale, math.Abs(hm[i*n+i]))
	}
	work := make([]float64, n*n)
	copy(work, hm)
	for k := range n {
		p := work[k*n+k]
		for j := range k {
			p -= work[k*n+j] * work[k*n+j]
		}
		if p <= 1e-13*scale {
			return base.Errf("%s: the Hessian is not positive definite (pivot %g at %d)", name, p, k+1)
		}
		p = math.Sqrt(p)
		work[k*n+k] = p
		for i := k + 1; i < n; i++ {
			s := work[i*n+k]
			for j := range k {
				s -= work[i*n+j] * work[k*n+j]
			}
			work[i*n+k] = s / p
		}
	}
	return nil
}

// matVec writes A·x into dst for a row-major m×n A.
func matVec(a []float64, x []float64, m, n int, dst []float64) {
	for i := range m {
		dst[i] = dot(a[i*n:i*n+n], x)
	}
}

// matVecR returns A·x as a fresh slice, the value form of matVec.
func matVecR(a []float64, x []float64, n int) []float64 {
	dst := make([]float64, n)
	matVec(a, x, n, n, dst)
	return dst
}

// dot is the plain inner product of two equal-length slices.
func dot(a, b []float64) float64 {
	s := 0.0
	for i := range a {
		s += a[i] * b[i]
	}
	return s
}

// wallKKTSign maps a step's slope along a row to the KKT column sign
// of the wall it is heading for: a positive slope meets the upper wall
// (the constraint reads a·x − u ≤ 0 there), a negative one the lower
// wall (the constraint reads l − a·x ≤ 0).
func wallKKTSign(d float64) float64 {
	if d > 0 {
		return 1
	}
	return -1
}