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

package linalg

import (
	"sourcedock.dev/petrbalvin/tensor/internal/base"
	"sourcedock.dev/petrbalvin/tensor/internal/core"
)

// Sparse symmetric positive-definite solve. The dense `Solve` factors
// A once and back-substitutes, which costs O(n³) time and O(n²) memory
// and is the right answer for a modest, dense system. When A is large
// and sparse that factorisation is both unaffordable and wasteful: the
// fill-in alone can exceed the memory the non-zeros needed. SpSolve
// runs the preconditioned conjugate gradient instead, whose every step
// is one sparse-matrix-vector product plus a handful of vector
// operations, so the cost tracks the non-zero count rather than the
// dimension.
//
// Conjugate gradient applies only to a symmetric positive-definite
// matrix, and it is not a drop-in for `Solve`: the answer is an
// approximation within a stated residual tolerance, not the exact
// solution a factorisation returns. Symmetry is verified the same way
// `SpEigen` verifies it. Positive-definiteness cannot be checked
// cheaply up front, so it is caught in flight: a non-positive search
// curvature is the matrix declaring itself indefinite.

// spSolveTol is the default relative residual when the caller passes
// zero. The iteration default is the dimension rather than a constant:
// an exact arithmetic conjugate gradient terminates in at most n
// steps, so a larger budget cannot buy convergence, only round-off
// work.
const spSolveTol = 1e-10

// checkSparseSquare validates the inputs every real sparse solver
// shares: a is a square 2-D matrix of non-complex dtype and, when b is
// not nil, b is a real right-hand side vector of the matrix dimension.
// It returns the dimension n.
func checkSparseSquare(name string, a *core.SparseCOO, b *core.Array) (int, error) {
	if a.Values.Dtype() == core.Complex {
		return 0, base.Errf("%s: complex sparse matrices are not supported", name)
	}
	if len(a.Shape) != 2 || a.Shape[0] != a.Shape[1] {
		return 0, base.Errf("%s: needs a square 2-D sparse matrix, got shape %v", name, a.Shape)
	}
	n := a.Shape[0]
	if n == 0 {
		return 0, base.Errf("%s: zero-sized matrix, got shape %v", name, a.Shape)
	}
	if b != nil {
		if b.NDim() != 1 || b.Shape()[0] != n {
			return 0, base.Errf("%s: right-hand side must be a vector of length %d, got shape %s",
				name, n, base.ShapeText(b.Shape()))
		}
		if b.Dtype() == core.Complex {
			return 0, base.Errf("%s: complex right-hand side is not supported", name)
		}
	}
	return n, nil
}

// pickPreconditioner resolves the optional ILU argument of the Krylov
// solvers, falling back to the Jacobi diagonal of c when none was
// given. An ILU built for another dimension is refused: applying it
// would index past its vectors, or precondition with the wrong
// factorisation without a word.
func pickPreconditioner(name string, c *sparseCSR, precond []*SparseILU) (*SparseILU, []float64, error) {
	if len(precond) > 0 {
		if precond[0] == nil {
			return nil, nil, base.Errf("%s: the preconditioner is nil", name)
		}
		if precond[0].n != c.n {
			return nil, nil, base.Errf("%s: the preconditioner was built for dimension %d, the system is %d",
				name, precond[0].n, c.n)
		}
		return precond[0], nil, nil
	}
	diag, err := c.diagonal(name)
	if err != nil {
		return nil, nil, err
	}
	return nil, diag, nil
}

// vectorF64 flattens a rank-1 array into a plain float64 working
// vector of the expected length n.
func vectorF64(b *core.Array, n int) []float64 {
	r := make([]float64, n)
	for i := range n {
		r[i] = b.FloatAt(i)
	}
	return r
}

// SpSolve returns the vector x solving A·x = b for a real symmetric
// positive-definite sparse A, by preconditioned conjugate gradient
// with a Jacobi (diagonal) preconditioner. The right-hand side b is a
// rank-1 vector of length n.
//
// The stopping rule is the relative residual ‖b − A·x‖₂ ≤ tol·‖b‖₂.
// A tol ≤ 0 uses 1e-10, and maxIter ≤ 0 uses n steps. x starts at
// zero, so the initial residual is b itself and the loop needs no
// separate first matvec. The default preconditioner is Jacobi
// scaling; an ILU(0) factorisation from NewSparseILU may be passed to
// replace it.
//
// An unconverged solve is an error naming the residual achieved, with
// no estimate returned, rather than a silent approximation.
// The matrix must be symmetric, and a zero diagonal entry is refused:
// the Jacobi preconditioner divides by it, and for a symmetric
// positive-definite matrix the diagonal is necessarily positive.
func SpSolve(a *core.SparseCOO, b *core.Array, tol float64, maxIter int, precond ...*SparseILU) (*core.Array, error) {
	const name = "SpSolve"
	n, err := checkSparseSquare(name, a, b)
	if err != nil {
		return nil, err
	}
	c, err := symmetricCSR(a, name)
	if err != nil {
		return nil, err
	}
	ilu, diag, err := pickPreconditioner(name, c, precond)
	if err != nil {
		return nil, err
	}

	if tol <= 0 {
		tol = spSolveTol
	}
	if maxIter <= 0 {
		maxIter = n
	}

	// x starts at zero, so r = b − A·x is just b and the loop needs no
	// opening matvec.
	r := vectorF64(b, n)
	bNorm := norm2F64(r)
	if bNorm == 0 {
		// The exact solution of A·0 = 0 is the zero vector; a relative
		// stopping rule would otherwise divide by zero.
		return core.Zeros(core.Float, n)
	}

	z := make([]float64, n)
	iluPrecondition(z, r, ilu, diag)
	p := append([]float64(nil), z...)
	rz := dotF64(r, z)
	x := make([]float64, n)
	ap := make([]float64, n)

	for iter := range maxIter {
		c.matVec(p, ap)
		pAp := dotF64(p, ap)
		if !finiteF64(pAp) || pAp <= 0 {
			// A non-positive or non-finite curvature means the matrix
			// is not positive-definite, or has overflowed the
			// representable range; either way conjugate gradient
			// cannot proceed. A non-finite value must be an error,
			// not a NaN that slips through the <= 0 test.
			return nil, base.Errf("%s: matrix is not positive-definite (curvature %.3g at step %d)",
				name, pAp, iter+1)
		}
		alpha := rz / pAp
		for i := range n {
			x[i] += alpha * p[i]
			r[i] -= alpha * ap[i]
		}
		// norm2F64 skips NaN entries, so an all-NaN residual reads as a
		// zero norm; a non-finite state is a breakdown before the
		// convergence test can mistake it for an exact solve.
		if !vecFinite(r) {
			return nil, base.Errf("%s: non-finite residual at step %d", name, iter+1)
		}
		if norm2F64(r) <= tol*bNorm {
			return floatsToArray(x, []int{n}), nil
		}
		iluPrecondition(z, r, ilu, diag)
		rzNext := dotF64(r, z)
		if rz == 0 {
			return nil, base.Errf("%s: breakdown at step %d (preconditioned residual vanished)", name, iter+1)
		}
		beta := rzNext / rz
		for i := range n {
			p[i] = z[i] + beta*p[i]
		}
		rz = rzNext
	}
	return nil, base.Errf("%s: no convergence in %d steps, residual %.3g (tolerance %.3g)",
		name, maxIter, norm2F64(r), tol*bNorm)
}

// diagonal returns the main diagonal, which the Jacobi preconditioner
// divides by. A missing or zero entry makes the preconditioner
// undefined and, for a symmetric positive-definite matrix, signals
// that the input is not one.
func (c *sparseCSR) diagonal(name string) ([]float64, error) {
	d := make([]float64, c.n)
	for i := range c.n {
		v, ok := c.at(i, i)
		if !ok || v == 0 {
			return nil, base.Errf("%s: zero or missing diagonal entry at %d", name, i)
		}
		d[i] = v
	}
	return d, nil
}