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

package linalg

import (
	"math"

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

// ILU(0): the incomplete LU factorisation whose L and U factors keep
// exactly the sparsity pattern of A, dropping every fill-in entry the
// complete factorisation would create. The result is not a solver but
// a preconditioner: applying it approximates A⁻¹ at the cost of one
// forward and one backward substitution over the stored pattern, and
// the Krylov solvers converge in far fewer steps than under plain
// Jacobi scaling because the approximation carries the matrix's local
// coupling. For a tridiagonal matrix the complete factorisation
// creates no fill anyway, so ILU(0) equals the complete LU there; the
// wider the stencil, the bigger the accuracy gap and the cheaper the
// factorisation stays.
//
// The factorisation runs the row-oriented IKJ elimination: row i's
// left part scales by the pivots above it, and each elimination
// updates only the positions the pattern already stores, intersecting
// row i with the pivot row by a two-pointer walk over the sorted
// column indices. Cost tracks the non-zero count, not n³.

// SparseILU holds the ILU(0) factorisation of a square sparse matrix
// in its CSR pattern: entries left of the diagonal are L multipliers
// with the unit diagonal implied, entries at or right of it are U
// values including the pivots on the diagonal.
type SparseILU struct {
	vals     []float64
	colIdx   []int
	rowStart []int
	n        int
}

// NewSparseILU factors a square sparse matrix with ILU(0). Column
// indices within each row must be sorted ascending, which the COO to
// CSR conversions guarantee. A missing diagonal entry or a zero pivot
// is refused: the factorisation divides by both, and an incomplete
// factorisation cannot repair a structurally singular matrix.
// Non-finite entries are refused, exactly as NewSparseLU and
// NewSparseCholesky refuse them.
func NewSparseILU(a *core.SparseCOO) (*SparseILU, error) {
	if a.Values.Dtype() == core.Complex {
		return nil, base.Errf("NewSparseILU: complex sparse matrices are not supported")
	}
	if len(a.Shape) != 2 || a.Shape[0] != a.Shape[1] {
		return nil, base.Errf("NewSparseILU: needs a square 2-D sparse matrix, got shape %v", a.Shape)
	}
	c, err := cooToCSR(a, "NewSparseILU")
	if err != nil {
		return nil, err
	}
	for i := range c.n {
		for p := c.rowStart[i]; p < c.rowStart[i+1]; p++ {
			v := c.vals[p]
			if math.IsNaN(v) || math.IsInf(v, 0) {
				return nil, base.Errf("NewSparseILU: entry [%d,%d] is not finite", i, c.colIdx[p])
			}
		}
	}
	n := c.n
	// Every row needs a pivot: the forward substitution divides by the
	// pivots above the diagonal and the backward substitution by every
	// diagonal, so a row without one is structurally singular. The
	// elimination loop below only inspects the diagonals a later row
	// actually pivots on, which never includes the last row (and misses
	// any other row nothing below it references), so the scan is what
	// makes the documented refusal hold for the whole matrix.
	for i := range n {
		if c.atDiagonal(i) == 0 {
			return nil, base.Errf("NewSparseILU: missing or zero diagonal at row %d", i)
		}
	}
	for i := 1; i < n; i++ {
		rs, re := c.rowStart[i], c.rowStart[i+1]
		// Walk row i's left part; every pivot row k < i is final by
		// the time row i is processed.
		for p := rs; p < re && c.colIdx[p] < i; p++ {
			k := c.colIdx[p]
			pivot := c.atDiagonal(k)
			if pivot == 0 {
				return nil, base.Errf("NewSparseILU: zero pivot at row %d", k)
			}
			mult := c.vals[p] / pivot
			c.vals[p] = mult
			// Two-pointer intersection of row i right of column k
			// with row k right of column k. The elimination can
			// overflow on legal finite input (the siblings in LU and
			// Cholesky refuse the same), and Apply has no error
			// return, so a poisoned factor must never leave here.
			if math.IsInf(mult, 0) || math.IsNaN(mult) {
				return nil, base.Errf("NewSparseILU: the elimination overflowed at row %d, column %d (multiplier %g)", i, k, mult)
			}
			q := c.rowStart[k]
			kre := c.rowStart[k+1]
			for q < kre && c.colIdx[q] <= k {
				q++
			}
			t := p + 1
			for t < re && q < kre {
				switch {
				case c.colIdx[t] == c.colIdx[q]:
					c.vals[t] -= mult * c.vals[q]
					if math.IsInf(c.vals[t], 0) || math.IsNaN(c.vals[t]) {
						return nil, base.Errf("NewSparseILU: the elimination overflowed at row %d, column %d", i, c.colIdx[t])
					}
					t++
					q++
				case c.colIdx[t] < c.colIdx[q]:
					t++
				default:
					q++
				}
			}
		}
	}
	return &SparseILU{
		vals:     c.vals,
		colIdx:   c.colIdx,
		rowStart: c.rowStart,
		n:        n,
	}, nil
}

// atDiagonal returns the diagonal entry of row i, or 0 when absent.
func (c *sparseCSR) atDiagonal(i int) float64 {
	for p := c.rowStart[i]; p < c.rowStart[i+1] && c.colIdx[p] <= i; p++ {
		if c.colIdx[p] == i {
			return c.vals[p]
		}
	}
	return 0
}

// Apply solves (L·U)·x = r over the stored pattern: a forward
// substitution with unit diagonal, then a backward substitution with
// the pivots. The result approximates A⁻¹·r to the accuracy the
// dropped fill allows.
func (ilu *SparseILU) Apply(r []float64) []float64 {
	dst := make([]float64, ilu.n)
	ilu.applyTo(dst, r)
	return dst
}

// applyTo writes (L·U)⁻¹·r into dst, which must be a buffer of the
// factor's dimension distinct from r. Both substitutions run in place:
// the forward pass overwrites each entry only after reading the ones
// below it, and the backward pass reads dst[i] before overwriting it
// with y[i], so an application allocates nothing.
func (ilu *SparseILU) applyTo(dst, r []float64) {
	n := ilu.n
	copy(dst, r)
	for i := range n {
		s := dst[i]
		for p := ilu.rowStart[i]; p < ilu.rowStart[i+1] && ilu.colIdx[p] < i; p++ {
			s -= ilu.vals[p] * dst[ilu.colIdx[p]]
		}
		dst[i] = s
	}
	for i := n - 1; i >= 0; i-- {
		s := dst[i]
		var diag float64
		for p := ilu.rowStart[i]; p < ilu.rowStart[i+1]; p++ {
			j := ilu.colIdx[p]
			if j == i {
				diag = ilu.vals[p]
			} else if j > i {
				s -= ilu.vals[p] * dst[j]
			}
		}
		if diag == 0 {
			// NewSparseILU refuses zero pivots, so this is defensive.
			diag = 1
		}
		dst[i] = s / diag
	}
}

// iluPrecondition writes the preconditioned vector of r into dst: the
// caller's ILU factor when one was given, otherwise the Jacobi scaling
// by the diagonal.
func iluPrecondition(dst, r []float64, ilu *SparseILU, diag []float64) {
	if ilu != nil {
		ilu.applyTo(dst, r)
		return
	}
	for i := range r {
		dst[i] = r[i] / diag[i]
	}
}