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

// Nonsymmetric sparse solve. SpSolve's conjugate gradient leans on
// symmetry for its short recurrences; the convection terms that make
// systems nonsymmetric break that structure, and the fix is van der
// Vorst's BiCGSTAB: two nested three-term recurrences whose second
// half smooths the erratic convergence the plain biconjugate gradient
// shows. Every step costs two sparse products and two preconditioner
// applications, the same order as one CG step pair.
//
// The Jacobi preconditioner needs a nonzero diagonal, which for the
// systems this solver targets (shifted Laplacians with convection,
// discretised transport) is present. Breakdowns (the recurrences
// collapsing) are reported as errors with the step they happened at,
// never swallowed.

// SpSolveBiCGSTAB returns the vector x solving A·x = b for a general
// real square sparse A, by preconditioned BiCGSTAB. Unlike SpSolve it
// asks for no symmetry; like SpSolve the stopping rule is the
// relative residual ‖b − A·x‖₂ ≤ tol·‖b‖₂ (tol ≤ 0 means 1e-10),
// maxIter ≤ 0 means n steps, and an unconverged solve is an error
// naming the residual achieved. The default preconditioner is Jacobi
// scaling, which divides by the diagonal and refuses a zero or
// missing entry; passing an ILU(0) factorisation from NewSparseILU
// replaces it and usually cuts the step count on hard systems.
func SpSolveBiCGSTAB(a *core.SparseCOO, b *core.Array, tol float64, maxIter int, precond ...*SparseILU) (*core.Array, error) {
	const name = "SpSolveBiCGSTAB"
	n, err := checkSparseSquare(name, a, b)
	if err != nil {
		return nil, err
	}
	c, err := cooToCSR(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
	}

	r := vectorF64(b, n)
	bNorm := norm2F64(r)
	if bNorm == 0 {
		return core.Zeros(core.Float, n)
	}
	rHat := append([]float64(nil), r...)
	x := make([]float64, n)
	v := make([]float64, n)
	p := make([]float64, n)
	s := make([]float64, n)
	t := make([]float64, n)
	z := make([]float64, n)
	y := make([]float64, n)
	mt := make([]float64, n)

	rho := 1.0
	alpha := 1.0
	omega := 1.0
	for iter := range maxIter {
		rhoNext := dotF64(rHat, r)
		if !finiteF64(rhoNext) || rhoNext == 0 {
			return nil, base.Errf("SpSolveBiCGSTAB: breakdown at step %d (rho vanished)", iter+1)
		}
		if iter > 0 {
			beta := rhoNext / rho * (alpha / omega)
			for i := range n {
				p[i] = r[i] + beta*(p[i]-omega*v[i])
			}
		} else {
			copy(p, r)
		}
		iluPrecondition(z, p, ilu, diag)
		c.matVec(z, v)
		rhatV := dotF64(rHat, v)
		if rhatV == 0 {
			return nil, base.Errf("SpSolveBiCGSTAB: breakdown at step %d (direction orthogonal to shadow)", iter+1)
		}
		alpha = rhoNext / rhatV
		for i := range n {
			s[i] = r[i] - alpha*v[i]
		}
		if !vecFinite(s) {
			return nil, base.Errf("SpSolveBiCGSTAB: breakdown at step %d (non-finite residual)", iter+1)
		}
		// The half step may already be exact. The check is affirmative so
		// a NaN norm (which fails every comparison) can never read as
		// converged.
		if norm2F64(s) > tol*bNorm {
			iluPrecondition(y, s, ilu, diag)
			c.matVec(y, t)
			iluPrecondition(mt, t, ilu, diag)
			mtS := dotF64(mt, s)
			mtT := dotF64(mt, t)
			if mtT == 0 {
				return nil, base.Errf("SpSolveBiCGSTAB: breakdown at step %d (stabiliser vanished)", iter+1)
			}
			omega = mtS / mtT
			// The finiteness half of the guard belongs before the
			// update: a NaN omega poisons x and r below, and the
			// all-NaN residual then reads as an exact solve through
			// the NaN-skipping norm.
			if !finiteF64(omega) {
				return nil, base.Errf("SpSolveBiCGSTAB: breakdown at step %d (stabiliser not finite)", iter+1)
			}
			for i := range n {
				x[i] += alpha*z[i] + omega*y[i]
				r[i] = s[i] - omega*t[i]
			}
			if !vecFinite(r) {
				return nil, base.Errf("SpSolveBiCGSTAB: non-finite residual at step %d", iter+1)
			}
			if norm2F64(r) <= tol*bNorm {
				return floatsToArray(x, []int{n}), nil
			}
			if omega == 0 {
				return nil, base.Errf("SpSolveBiCGSTAB: stagnation at step %d (omega zero)", iter+1)
			}
		} else {
			for i := range n {
				x[i] += alpha * z[i]
			}
			return floatsToArray(x, []int{n}), nil
		}
		rho = rhoNext
	}
	return nil, base.Errf("SpSolveBiCGSTAB: no convergence in %d steps, residual %.3g (tolerance %.3g)",
		maxIter, norm2F64(r), tol*bNorm)
}