# SPDX-FileCopyrightText: 2015-2026 Jukka Aho
# SPDX-License-Identifier: MIT

"""
Lightweight, matrix-free preconditioners for Krylov solves built on top
of `apply_K!` + `matrix_free_op`. Concrete types include:

  * `JacobiPreconditioner`: stores `1 ./ diag(K_op)` and implements the
    `LinearAlgebra.ldiv!` protocol consumed by `IterativeSolvers.cg!`,
    `Krylov.jl`, `LinearSolve.jl`, etc.

The diagonal is extracted matrix-free via `compute_diagonal!` (same
DOF-row traversal as `apply_K!`), so building the preconditioner has
the same complexity as a single `apply_K!` call — no extra storage
beyond the `ndofs`-vector of inverse diagonal entries.

For symmetric-indefinite mixed `u`–`p` problems (`MixedUPKernel` with
`inv_bulk = 0`), scalar Jacobi is only a baseline: Krylov solvers should
use GMRES or MINRES, not CG, and a field-split or approximate Schur
preconditioner is usually required for robust convergence (see
`test/domains/continuum/test_mixed_up_incompressible_solve.jl`).
[`ApproxSchurDiagBlockPreconditioner`](@ref) applies diagonal scaling on the primal block
and an entry-wise diagonal approximation of the pressure Schur complement from an assembled
sparse `K` (useful baseline for symmetric saddle-point systems; not a general-purpose
solver for hard indefinite problems).
[`ApproxSchurICholDiagBlockPreconditioner`](@ref) replaces the primal diagonal by IC(0) on
the `A` block when that slice is SPD enough for an incomplete factorisation.

# Usage

```julia
using IterativeSolvers, LinearOperators, JuliaFEM

bc = PenaltyDirichlet(fixed_dofs, vals; penalty = 1e8)
op = matrix_free_op(cache, asm, mesh; dirichlet = bc)
P  = JacobiPreconditioner(cache, asm, mesh; dirichlet = bc)

linop = LinearOperator(Float64, n, n, true, true, op)

T = zeros(n)
cg!(T, linop, b; Pl = P, abstol = 1e-12, reltol = 1e-12, maxiter = 4n)
```

For `PenaltyDirichlet` the Jacobi diagonal includes the `λ`
contribution, so CG sees a roughly `O(1)`-conditioned system on the
fixed DOFs and the original conditioning on the free block — closing
the gap that motivated `EliminatedDirichlet` in the first place.
"""

import LinearAlgebra
import LinearAlgebra: ldiv!
using LinearAlgebra: diag, inv, LowerTriangular, UpperTriangular
using SparseArrays: SparseMatrixCSC, sparse, nnz, rowvals, nonzeros, getcolptr

"""
    JacobiPreconditioner(inv_diag::Vector{Float64})

Diagonal preconditioner with stored *inverse* diagonal so that
`ldiv!(P, x)` is a single multiplication per entry. Direct constructor;
prefer the `(cache, asm, mesh; dirichlet)` factory for the common matrix-free path.
"""
struct JacobiPreconditioner
    inv_diag::Vector{Float64}
end

"""
    JacobiPreconditioner(cache, asm, mesh; dirichlet = nothing, mpc = nothing)

Build a Jacobi preconditioner from the matrix-free operator. Internally:

  1. extract `diag(K)` matrix-free via `compute_diagonal!`,
  2. apply the constraint's diagonal hook (`apply_constraint_diag!`) so
     the preconditioner matches the operator returned by
     `matrix_free_op(...; dirichlet, mpc)`. Both Dirichlet and MPC
     constraints are added if supplied (additive: the MPC adds
     `λ · diag(Cᵀ C)` on top of the Dirichlet diagonal),
  3. invert entry-wise (zeros are kept as `1.0` so the preconditioner
     stays well-defined; in practice every DOF row of `K` has a positive
     diagonal in any well-posed FEM problem).

    JacobiPreconditioner(cache, asm, kernel, mesh; …)

Four-argument backward-compatibility overload: `kernel` is ignored; volume
physics come from `cache.kernel_column` only (same as [`compute_diagonal!`](@ref)).
Emits `Base.depwarn` once per session (shared with other redundant-kernel overloads).
"""
function JacobiPreconditioner(cache::DOFBasedCOOCache,
                              asm::DOFBasedCOOAssembler,
                              mesh::AbstractMesh;
                              dirichlet::Union{AbstractDirichletConstraint, Nothing} = nothing,
                              mpc::Union{AbstractMultipointConstraint, Nothing} = nothing)
    n = cache.ndofs
    d = zeros(Float64, n)
    compute_diagonal!(d, cache, asm, mesh)
    if dirichlet !== nothing
        apply_constraint_diag!(d, dirichlet)
    end
    if mpc !== nothing
        apply_constraint_diag!(d, mpc)
    end

    inv_d = similar(d)
    @inbounds for i in eachindex(d)
        inv_d[i] = d[i] != 0.0 ? 1.0 / d[i] : 1.0
    end
    return JacobiPreconditioner(inv_d)
end

@inline function JacobiPreconditioner(cache::DOFBasedCOOCache,
                                      asm::DOFBasedCOOAssembler,
                                      ::AbstractKernel,
                                      mesh::AbstractMesh;
                                      kwargs...)
    _depwarn_redundant_kernel_arg!(:JacobiPreconditioner)
    return JacobiPreconditioner(cache, asm, mesh; kwargs...)
end

# `ldiv!(y, P, x)` is the contract IterativeSolvers.cg! / Krylov.jl call
# for left-preconditioning: compute `y = P^{-1} * x`. For a diagonal
# `P`, this is just an entry-wise multiplication by the stored inverse.
function ldiv!(y::AbstractVector{Float64}, P::JacobiPreconditioner,
               x::AbstractVector{Float64})
    @inbounds @simd for i in eachindex(y)
        y[i] = x[i] * P.inv_diag[i]
    end
    return y
end

# In-place variant — mutates `x` directly.
function ldiv!(P::JacobiPreconditioner, x::AbstractVector{Float64})
    @inbounds @simd for i in eachindex(x)
        x[i] *= P.inv_diag[i]
    end
    return x
end


# ============================================================================
# BlockJacobiPreconditioner{N} — N×N block-diagonal preconditioner
# ============================================================================
#
# Scalar Jacobi is weak for vector problems where the off-diagonal entries
# *within a node block* are O(diag) — typical for elasticity where the
# 3×3 nodal block has full off-diagonal couplings. Inverting the 3×3
# block as a unit recovers a much better approximation of `K^{-1}` at
# the same matrix-free cost (one extra `evaluate_entry` call per off-
# diagonal entry of each block).
#
# The block layout assumes the DOFHandler's standard interleaved
# numbering: `dof = N · (block_id − 1) + comp_in_block`, with
# `comp_in_block ∈ 1..N`. This matches how `create_elements!` numbers
# vector-field DOFs (and the same convention generalizes to N=4 for
# `Displacement + Temperature`).

"""
    compute_block_diagonal!(blocks::Array{Float64,3},
                            cache::DOFBasedCOOCache,
                            asm::DOFBasedCOOAssembler,
                            mesh::AbstractMesh) -> blocks

Assemble the `N × N` block-diagonal of the stiffness matrix into
`blocks`, a `(N, N, n_blocks)` array. Block `b` covers global DOFs
`N · (b − 1) + 1 : N · b` (interleaved numbering, the standard layout
for vector fields).

Same DOF-by-element traversal as `assemble!` but only the entries with
both DOFs in the same block contribute. Allocation-free after warmup.

    compute_block_diagonal!(blocks, cache, asm, kernel, mesh)

Five-argument overload: `kernel` is ignored; volume kernels are read from
`cache.kernel_column`. Emits `Base.depwarn` once per session when used.
"""
function compute_block_diagonal!(
    blocks::Array{Float64,3},
    cache::DOFBasedCOOCache{T,B,IPS,E,GC,Buf,FieldType,StateType,KS},
    asm::DOFBasedCOOAssembler,
    mesh::AbstractMesh,
) where {T,B,IPS,E<:AbstractElement,GC,Buf,FieldType,StateType,KS}
    N, N2, n_blocks = size(blocks)
    @assert N == N2 (
        "compute_block_diagonal!: blocks must be (N,N,n_blocks); got $(size(blocks))")
    @assert N * n_blocks == cache.ndofs (
        "compute_block_diagonal!: N · n_blocks = $(N * n_blocks) must equal " *
        "cache.ndofs = $(cache.ndofs)")

    fill!(blocks, 0.0)
    _prepare_caches!(cache, mesh)

    elements         = cache.elements
    element_caches   = cache.element_caches
    geometry_caches  = cache.geometry_caches
    qp_buffers       = cache.qp_buffers

    layout     = local_dof_layout(E)
    ndofs_elem = length(layout)

    @inbounds for elem_idx in 1:length(elements)
        k_e = kernel_at(cache, elem_idx)
        ec = element_caches[elem_idx]
        gc = geometry_caches[elem_idx]
        qp = view(qp_buffers, :, elem_idx)
        dofs_elem = ec.dofs

        @inbounds for li in 1:ndofs_elem
            entry_i  = layout[li]
            dof_i    = Int(dofs_elem[li])
            blk_i    = div(dof_i - 1, N) + 1
            comp_i   = mod(dof_i - 1, N) + 1

            @inbounds for lj in 1:ndofs_elem
                dof_j = Int(dofs_elem[lj])
                blk_j = div(dof_j - 1, N) + 1
                if blk_i != blk_j
                    continue
                end
                comp_j  = mod(dof_j - 1, N) + 1
                entry_j = layout[lj]

                K_ij = evaluate_entry(k_e, gc, qp, entry_i, entry_j, elem_idx)
                blocks[comp_i, comp_j, blk_i] += K_ij
            end
        end
    end

    return blocks
end

@inline function compute_block_diagonal!(
    blocks::Array{Float64,3},
    cache::DOFBasedCOOCache{T,B,IPS,E,GC,Buf,FieldType,StateType,KS},
    asm::DOFBasedCOOAssembler,
    ::AbstractKernel,
    mesh::AbstractMesh,
) where {T,B,IPS,E<:AbstractElement,GC,Buf,FieldType,StateType,KS}
    _depwarn_redundant_kernel_arg!(:compute_block_diagonal!)
    return compute_block_diagonal!(blocks, cache, asm, mesh)
end

"""
    apply_constraint_block_diag!(blocks::Array{Float64,3}, c) -> blocks

Apply a Dirichlet constraint to the block-diagonal storage so the
preconditioner stays consistent with the operator returned by
`matrix_free_op(...; dirichlet = c)`. Default: no-op. Override for new
constraint types parallel to `apply_constraint_diag!`.
"""
@inline apply_constraint_block_diag!(blocks::Array{Float64,3},
                                     ::AbstractDirichletConstraint) = blocks

@inline function apply_constraint_block_diag!(blocks::Array{Float64,3},
                                              c::PenaltyDirichlet)
    N, _, _ = size(blocks)
    λ = c.penalty
    @inbounds for fdof in c.fixed_dofs
        blk = div(fdof - 1, N) + 1
        cmp = mod(fdof - 1, N) + 1
        blocks[cmp, cmp, blk] += λ
    end
    return blocks
end

@inline function apply_constraint_block_diag!(blocks::Array{Float64,3},
                                              c::EliminatedDirichlet)
    N, _, _ = size(blocks)
    @inbounds for fdof in c.fixed_dofs
        blk = div(fdof - 1, N) + 1
        cmp = mod(fdof - 1, N) + 1
        # Zero the row/column of the eliminated DOF inside the block,
        # then identity on the diagonal — mirrors the eliminated
        # operator on the corresponding block row.
        @inbounds for k in 1:N
            blocks[cmp, k, blk] = 0.0
            blocks[k, cmp, blk] = 0.0
        end
        blocks[cmp, cmp, blk] = 1.0
    end
    return blocks
end

"""
    BlockJacobiPreconditioner{N}

`N × N` block-diagonal preconditioner. Stores the *inverse* of every
block (a single `(N, N, n_blocks)` array) so `ldiv!` is a single dense
`N`-by-`N` matrix-vector multiply per block — `O(N² · n_blocks)` work,
no allocations.

For elasticity (`Displacement{3}`) use `N = 3`; for thermo-elastic
(displacement + temperature) use `N = 4`. The block grouping assumes
DOFs are numbered `dof = N · (block − 1) + comp`, which matches the
`create_elements!` layout for any `@DOFSet` whose total per-vertex DOF
count is `N`.

Mixed displacement (`Vertex`) plus cell-centered pressure (`Cell`) places all
velocity DOFs first, then all pressures ([`global_field_ranges`](@ref)); that layout
does **not** match fixed-width `N`-DOF blocks over the full unknown vector, so
[`BlockJacobiPreconditioner`](@ref) is not appropriate for the whole saddle-point
system — use scalar [`JacobiPreconditioner`](@ref), a field-split / Schur-based
preconditioner on extracted blocks ([`saddle_point_blocks`](@ref)), or a direct solver.
"""
struct BlockJacobiPreconditioner{N}
    inv_blocks::Array{Float64,3}     # (N, N, n_blocks)
end

"""
    BlockJacobiPreconditioner{N}(cache, asm, mesh; dirichlet = nothing)

Build a block-Jacobi preconditioner from the matrix-free operator.
Internally:

  1. assemble the `N × N` diagonal blocks via `compute_block_diagonal!`,
  2. apply `apply_constraint_block_diag!` so the preconditioner matches
     the operator returned by `matrix_free_op(...; dirichlet)`,
  3. invert each block in place (`LinearAlgebra.inv` via `LU`); on
     singular blocks, fall back to identity to keep the preconditioner
     well-defined (a singular nodal block typically signals an ill-posed
     problem, so this is intentionally permissive).

    BlockJacobiPreconditioner{N}(cache, asm, kernel, mesh; …)

Five-argument overload: `kernel` is ignored (emits `Base.depwarn` once per session).
"""
function BlockJacobiPreconditioner{N}(
    cache::DOFBasedCOOCache,
    asm::DOFBasedCOOAssembler,
    mesh::AbstractMesh;
    dirichlet::Union{AbstractDirichletConstraint, Nothing} = nothing,
) where {N}
    @assert N >= 1 "BlockJacobiPreconditioner: N must be >= 1 (got $N)"
    n_blocks = div(cache.ndofs, N)
    @assert N * n_blocks == cache.ndofs (
        "BlockJacobiPreconditioner{N=$N}: ndofs = $(cache.ndofs) is not divisible by N")

    blocks = zeros(Float64, N, N, n_blocks)
    compute_block_diagonal!(blocks, cache, asm, mesh)
    if dirichlet !== nothing
        apply_constraint_block_diag!(blocks, dirichlet)
    end

    inv_blocks = similar(blocks)
    block_view = zeros(Float64, N, N)
    @inbounds for b in 1:n_blocks
        @inbounds for j in 1:N, i in 1:N
            block_view[i, j] = blocks[i, j, b]
        end
        invblk = try
            inv(block_view)
        catch
            Matrix{Float64}(LinearAlgebra.I, N, N)
        end
        @inbounds for j in 1:N, i in 1:N
            inv_blocks[i, j, b] = invblk[i, j]
        end
    end

    return BlockJacobiPreconditioner{N}(inv_blocks)
end

@inline function BlockJacobiPreconditioner{N}(
    cache::DOFBasedCOOCache,
    asm::DOFBasedCOOAssembler,
    ::AbstractKernel,
    mesh::AbstractMesh;
    kwargs...,
) where {N}
    _depwarn_redundant_kernel_arg!(:BlockJacobiPreconditioner)
    return BlockJacobiPreconditioner{N}(cache, asm, mesh; kwargs...)
end

# ldiv!(y, P, x): y = P^{-1} x, block by block.
function ldiv!(y::AbstractVector{Float64},
               P::BlockJacobiPreconditioner{N},
               x::AbstractVector{Float64}) where {N}
    n_blocks = size(P.inv_blocks, 3)
    @inbounds for b in 1:n_blocks
        base = N * (b - 1)
        @inbounds for i in 1:N
            yi = 0.0
            @inbounds for j in 1:N
                yi += P.inv_blocks[i, j, b] * x[base + j]
            end
            y[base + i] = yi
        end
    end
    return y
end

# In-place: x ← P^{-1} x.
function ldiv!(P::BlockJacobiPreconditioner{N},
               x::AbstractVector{Float64}) where {N}
    n_blocks = size(P.inv_blocks, 3)
    tmp = zeros(Float64, N)
    @inbounds for b in 1:n_blocks
        base = N * (b - 1)
        @inbounds for j in 1:N
            tmp[j] = x[base + j]
        end
        @inbounds for i in 1:N
            yi = 0.0
            @inbounds for j in 1:N
                yi += P.inv_blocks[i, j, b] * tmp[j]
            end
            x[base + i] = yi
        end
    end
    return x
end

# ============================================================================
# ApproxSchurDiagBlockPreconditioner — field-split diagonal + diag(Schur) estimate
# ============================================================================
#
# For K ≈ [ A  B ; Bt C ] with A (velocity block) square and SPD-ish, a classical
# cheap preconditioner uses inv(diag(A)) on the u block and inv(diag(S)) on p,
# where S ≈ C - B' diag(A)^{-1} B is the Schur complement. We approximate diag(S)
# entry-wise as  Cjj - sum_i B_ij^2 / A_ii  (using B and Bt consistently only when
# K is symmetric). For non-symmetric K the same diagonal formula is still used as a
# heuristic when building from `B = K[r_u,r_p]` and `diag(A)`.

function _inv_diag_A_and_schur(blks)
    A = blks.A
    B = blks.B
    C = blks.C
    nu = size(A, 1)
    np = size(C, 1)
    dA = diag(A)
    inv_dA = Vector{Float64}(undef, nu)
    @inbounds for i in 1:nu
        inv_dA[i] = dA[i] != 0.0 ? 1.0 / dA[i] : 1.0
    end
    dC = diag(C)
    schur_diag = Vector{Float64}(undef, np)
    @inbounds for j in 1:np
        s = dC[j]
        @inbounds for i in 1:nu
            bij = B[i, j]
            s -= bij * bij * inv_dA[i]
        end
        schur_diag[j] = s
    end
    inv_schur = Vector{Float64}(undef, np)
    @inbounds for j in 1:np
        sd = schur_diag[j]
        inv_schur[j] = abs(sd) > 1e-30 ? 1.0 / sd : 1.0
    end
    return inv_dA, inv_schur
end

"""
    ApproxSchurDiagBlockPreconditioner(inv_diag_A, inv_schur_diag, r_u, r_p)

Block-diagonal preconditioner in field-split ordering:

  * `y[r_u] ← inv_diag_A .* x[r_u]`  (diagonal scaling on the primal block)
  * `y[r_p] ← inv_schur_diag .* x[r_p]`  (diagonal Schur approximation on the second field)

Ranges must partition the relevant global indices (typical two-field mixed layout from
[`global_field_ranges`](@ref)). Construct with the `(K, r_u, r_p)` factory on an assembled
[`SparseMatrixCSC`](@ref).
"""
struct ApproxSchurDiagBlockPreconditioner
    inv_diag_A::Vector{Float64}
    inv_schur_diag::Vector{Float64}
    r_u::UnitRange{Int}
    r_p::UnitRange{Int}
end

"""
    ApproxSchurDiagBlockPreconditioner(K::SparseMatrixCSC, r_u, r_p)

Build from sparse `K` using [`saddle_point_matrix_blocks`](@ref): diagonal of `A`,
then `diag(S) ≈ diag(C) - sum_i B[i,j]^2 / A[i,i]` for each pressure column `j`.
Zeros on `diag(A)` fall back to unit inverse entries; tiny `diag(S)` uses `1.0`.
"""
function ApproxSchurDiagBlockPreconditioner(
    K::SparseMatrixCSC{Float64, Ti},
    r_u::AbstractRange{Int},
    r_p::AbstractRange{Int},
) where {Ti<:Integer}
    ru = UnitRange(first(r_u), last(r_u))
    rp = UnitRange(first(r_p), last(r_p))
    blks = saddle_point_matrix_blocks(K, ru, rp)
    nu = size(blks.A, 1)
    np = size(blks.C, 1)
    length(ru) == nu || error("ApproxSchurDiagBlockPreconditioner: |r_u| mismatch")
    length(rp) == np || error("ApproxSchurDiagBlockPreconditioner: |r_p| mismatch")

    inv_dA, inv_schur = _inv_diag_A_and_schur(blks)
    return ApproxSchurDiagBlockPreconditioner(inv_dA, inv_schur, ru, rp)
end

function ldiv!(
    y::AbstractVector{Float64},
    P::ApproxSchurDiagBlockPreconditioner,
    x::AbstractVector{Float64},
)
    length(y) == length(x) || throw(DimensionMismatch("ldiv!"))
    ru, rp = P.r_u, P.r_p
    @views @. y[ru] = P.inv_diag_A * x[ru]
    @views @. y[rp] = P.inv_schur_diag * x[rp]
    return y
end

function ldiv!(P::ApproxSchurDiagBlockPreconditioner, x::AbstractVector{Float64})
    ru, rp = P.r_u, P.r_p
    @views @. x[ru] = P.inv_diag_A * x[ru]
    @views @. x[rp] = P.inv_schur_diag * x[rp]
    return x
end

# ============================================================================
# ICholPreconditioner — Incomplete Cholesky with no fill-in (IC(0))
# ============================================================================
#
# `IC(0)` is the canonical "industrial-strength" preconditioner for SPD
# Krylov solves: it computes a sparse lower-triangular `L` such that
# `L L^T ≈ K`, with `L`'s sparsity pattern *equal* to the lower
# triangle of `K` (no fill). For most well-shaped FE meshes this gives
# CG convergence in `O(√cond(K))` iterations — compared to scalar
# Jacobi's `O(cond(K))` — at the cost of one sparse triangular solve
# per Krylov iteration plus a one-time factorisation.
#
# Implementation notes
# --------------------
# * IC(0) is built directly on the user's assembled `K` (via
#   `extract_system(cache)`) — the matrix-free operator stays
#   matrix-free for the *Krylov product* (`apply_K!`) but uses the
#   assembled K *only* for the preconditioner. This hybrid is the
#   standard pattern: PCG hits `K` once per `apply_K!` call but `L`
#   twice per iteration via `ldiv!`.
# * The factorisation is a Crout left-looking IC(0) over CSC
#   storage. Per-column workspace is a dense `Vector{Float64}` of
#   length `n` plus an integer `pos` array of the same length;
#   total `O(n)` extra memory, completely independent of the matrix.
# * After factorisation, `ldiv!(P, x)` runs SparseArrays' optimised
#   forward / backward substitution via `LowerTriangular(L)` and
#   `UpperTriangular(L')`. Both are zero-allocation hot paths.
# * Apply Dirichlet BCs to `K` *before* building `IC(0)` (typical
#   pattern: the penalty / elimination has already added `λ` to the
#   diagonal entries on the constrained DOFs). The preconditioner
#   then matches the operator the Krylov solver sees.

"""
    ICholPreconditioner

Lower-triangular Incomplete-Cholesky preconditioner. Build it with
either `ICholPreconditioner(K::SparseMatrixCSC)` (factor an existing
SPD matrix) or `ICholPreconditioner(cache, asm, mesh; dirichlet)`
(assemble + factor in one shot for the matrix-free workflow).

`ldiv!(P, x)` solves `L L^T y = x` via two sparse triangular solves
(forward then backward) — the canonical PCG preconditioner action.
"""
struct ICholPreconditioner
    L::SparseMatrixCSC{Float64,Int}
    # Inner constructor used by the IC(0) factorisation pipeline.
    # The `::Val{:ichol_internal}` tag prevents the outer
    # `ICholPreconditioner(K::SparseMatrixCSC...)` constructor from
    # recursing into itself when it tries to wrap the freshly-built
    # `L` factor (both have the same `SparseMatrixCSC{Float64,Int}`
    # signature).
    ICholPreconditioner(L::SparseMatrixCSC{Float64,Int}, ::Val{:ichol_internal}) = new(L)
end

"""
    ICholPreconditioner(K::SparseMatrixCSC{Float64,Int})

Build the IC(0) factor of the SPD matrix `K` and wrap it as a
preconditioner. Sparsity pattern of `L` = lower triangle of `K`
(no fill).

# Robustness

If a Cholesky update produces a non-positive diagonal (numerically not
SPD — e.g. the user passed an indefinite or near-singular matrix), the
factorisation falls back to a small diagonal-shift retry: `K + α · I`
with `α = max(0, 1e-12 · ‖K‖∞)`. This recovers in practice for
slightly-ill-conditioned PenaltyDirichlet systems while keeping a
well-defined preconditioner.
"""
function ICholPreconditioner(K::SparseMatrixCSC{Float64,Int})
    n = size(K, 1)
    @assert n == size(K, 2) "ICholPreconditioner: K must be square (got $(size(K)))"

    # Try the unshifted IC(0) first. On breakdown, fall back to a
    # progressively larger diagonal shift `K + α·I`. This is the
    # standard Manteuffel-style robustification for elasticity-grade
    # matrices, where IC(0) frequently produces a non-positive pivot
    # because of dropped fill — even though `K` itself is SPD. The
    # operator the Krylov solver sees is still the unshifted `K`
    # (we just need a usable preconditioner).
    L = try
        _ic0_factor(K)
    catch err
        if err isa ErrorException && occursin("IC(0)", err.msg)
            kmax  = maximum(abs.(nonzeros(K)))
            shift_levels = (1e-12, 1e-9, 1e-6, 1e-3, 1e-1, 1.0)
            local Lα = nothing
            for s in shift_levels
                α   = max(s * kmax, 1e-300)
                Kα  = K + α * sparse(LinearAlgebra.I, n, n)
                Lα = try
                    _ic0_factor(Kα)
                catch e
                    if e isa ErrorException && occursin("IC(0)", e.msg)
                        nothing
                    else
                        rethrow(e)
                    end
                end
                Lα === nothing || break
            end
            Lα === nothing ?
                error("ICholPreconditioner: IC(0) breakdown even after " *
                      "diagonal shift up to ‖K‖∞ = $(kmax). Matrix may not " *
                      "be SPD or is too far from diagonally dominant for " *
                      "no-fill incomplete Cholesky.") :
                Lα
        else
            rethrow(err)
        end
    end
    return ICholPreconditioner(L, Val(:ichol_internal))
end

"""
    ICholPreconditioner(cache::DOFBasedCOOCache, asm, mesh;
                        dirichlet = nothing) -> ICholPreconditioner

Convenience constructor for the matrix-free workflow:

    1. assemble K via `assemble!(cache, asm, mesh)` and extract it with
       `extract_system(cache)`,
    2. apply the Dirichlet constraint (if given) to `K`,
    3. build `IC(0)` of the modified `K`.

This *does* materialise the sparse `K` once — the price of IC(0). The
returned preconditioner then plugs into the matrix-free Krylov solve.

    ICholPreconditioner(cache, asm, kernel, mesh; …)

Four-argument overload: `kernel` is ignored (same as [`assemble!`](@ref)).
"""
function ICholPreconditioner(cache::DOFBasedCOOCache,
                             asm::DOFBasedCOOAssembler,
                             mesh::AbstractMesh;
                             dirichlet::Union{AbstractDirichletConstraint, Nothing} = nothing)
    assemble!(cache, asm, mesh)
    K, _ = extract_system(cache)
    if dirichlet !== nothing
        apply_constraint!(K, dirichlet)
    end
    return ICholPreconditioner(K)
end

@inline function ICholPreconditioner(cache::DOFBasedCOOCache,
                                     asm::DOFBasedCOOAssembler,
                                     ::AbstractKernel,
                                     mesh::AbstractMesh;
                                     kwargs...)
    _depwarn_redundant_kernel_arg!(:ICholPreconditioner)
    return ICholPreconditioner(cache, asm, mesh; kwargs...)
end

# Crout left-looking IC(0) on CSC storage. Returns the factor `L`
# such that `L * L^T ≈ K` and `nnz(L) == nnz(tril(K))`.
function _ic0_factor(K::SparseMatrixCSC{Float64,Int})
    n = size(K, 1)
    # Initialise `L` with the lower-triangle pattern of `K`. We must
    # deep-copy `nzval` explicitly: in some Julia / SparseArrays
    # versions `sparse(LowerTriangular(K))` can alias `K.nzval`, so
    # writing IC(0) factor entries into `L.nzval` would silently
    # corrupt the user's `K`. Found the hard way during the
    # diagonal-shift fallback path.
    Ltmp   = sparse(LowerTriangular(K))
    L      = SparseMatrixCSC(size(Ltmp, 1), size(Ltmp, 2),
                             copy(getcolptr(Ltmp)),
                             copy(rowvals(Ltmp)),
                             copy(nonzeros(Ltmp)))
    rowval = rowvals(L)
    nzval  = nonzeros(L)
    colptr = getcolptr(L)

    # Workspace: dense column buffer + position lookup. `pos[i] != 0`
    # iff row `i` is in column `k`'s pattern (i.e. there's an entry
    # `L[i,k]`).
    work = zeros(Float64, n)
    pos  = zeros(Int,     n)

    @inbounds for k in 1:n
        # 1. Splat column k of L into the dense workspace.
        kpos = 0
        for jp in colptr[k]:(colptr[k+1] - 1)
            i = rowval[jp]
            work[i] = nzval[jp]
            pos[i]  = jp
            if i == k
                kpos = jp
            end
        end
        if kpos == 0
            error("IC(0): missing diagonal entry K[$k,$k] — matrix is not " *
                  "structurally SPD.")
        end

        # 2. For each previous column j < k that has L[k, j] ≠ 0,
        #    subtract `L[k,j] * L[i,j]` from `work[i]` for every row
        #    `i ≥ k` already in the pattern of column k.
        for j in 1:(k-1)
            # Look up L[k, j] (row k in column j). Bail out fast if
            # row k isn't in column j's pattern.
            jp_kj = _find_row_in_col(rowval, colptr, j, k)
            if jp_kj == 0
                continue
            end
            Lkj = nzval[jp_kj]

            # Walk down column j from row k onwards; each row i with
            # `pos[i] ≠ 0` (i.e. in column k's pattern) contributes
            # `-Lkj * L[i,j]` to `work[i]`. Other rows fall on
            # would-be fill, dropped by IC(0).
            for jp in jp_kj:(colptr[j+1] - 1)
                i = rowval[jp]
                if pos[i] != 0
                    work[i] -= Lkj * nzval[jp]
                end
            end
        end

        # 3. Diagonal — must be positive for the matrix to be SPD.
        d = work[k]
        if d <= 0.0
            error("IC(0): non-positive diagonal at row $k (got $d). " *
                  "Matrix may not be SPD.")
        end
        sqrt_d = sqrt(d)
        nzval[kpos] = sqrt_d

        # 4. Sub-diagonals of column k are the updated work entries
        #    divided by the new diagonal.
        for jp in (kpos + 1):(colptr[k+1] - 1)
            i = rowval[jp]
            nzval[jp] = work[i] / sqrt_d
        end

        # 5. Clear the dense workspace for the next column. We only
        #    touched entries in column k's pattern, so this is cheap.
        for jp in colptr[k]:(colptr[k+1] - 1)
            i = rowval[jp]
            work[i] = 0.0
            pos[i]  = 0
        end
    end

    return L
end

# Binary search column j (sorted row indices) for row r. Returns the
# nzval index, or `0` if not found.
@inline function _find_row_in_col(rowval, colptr, j::Int, r::Int)
    lo = colptr[j]
    hi = colptr[j+1] - 1
    while lo <= hi
        mid = (lo + hi) >> 1
        rv  = rowval[mid]
        if rv == r
            return mid
        elseif rv < r
            lo = mid + 1
        else
            hi = mid - 1
        end
    end
    return 0
end

# `ldiv!(y, P, x)`: y = (L L^T)^{-1} x. Two sparse triangular solves.
function ldiv!(y::AbstractVector{Float64},
               P::ICholPreconditioner,
               x::AbstractVector{Float64})
    L = P.L
    # Forward solve: L * z = x  ⇒  z = L \ x. We use `y` as the
    # output buffer; SparseArrays' triangular solve is allocation-
    # free when both args are concrete `Vector{Float64}`.
    copyto!(y, x)
    LinearAlgebra.ldiv!(LowerTriangular(L), y)
    # Backward solve: L^T * y = z. Reuses the same storage.
    LinearAlgebra.ldiv!(UpperTriangular(L'), y)
    return y
end

# In-place: x ← (L L^T)^{-1} x.
function ldiv!(P::ICholPreconditioner, x::AbstractVector{Float64})
    L = P.L
    LinearAlgebra.ldiv!(LowerTriangular(L), x)
    LinearAlgebra.ldiv!(UpperTriangular(L'), x)
    return x
end

# ============================================================================
# ApproxSchurICholDiagBlockPreconditioner — IC(0) on A + diagonal Schur on p
# ============================================================================
#
# Stronger than [`ApproxSchurDiagBlockPreconditioner`](@ref) on the primal block:
# applies [`ICholPreconditioner`](@ref) to the sparse `A = K[r_u,r_u]` slice (requires that
# block to be SPD enough for IC(0)), and keeps the same diagonal Schur approximation on
# the pressure block.

"""
    ApproxSchurICholDiagBlockPreconditioner(Pa, inv_schur_diag, r_u, r_p, tmp_u)

Field-split preconditioner: [`ldiv!`](@ref) with [`ICholPreconditioner`](@ref) on `x[r_u]`
(stored in `tmp_u`) and diagonal Schur scaling on `x[r_p]`.
"""
struct ApproxSchurICholDiagBlockPreconditioner
    Pa::ICholPreconditioner
    inv_schur_diag::Vector{Float64}
    r_u::UnitRange{Int}
    r_p::UnitRange{Int}
    tmp_u::Vector{Float64}
end

"""
    ApproxSchurICholDiagBlockPreconditioner(K::SparseMatrixCSC, r_u, r_p)

Factor `A = K[r_u,r_u]` with IC(0) and build the same `inv_schur_diag` as
[`ApproxSchurDiagBlockPreconditioner`](@ref).
"""
function ApproxSchurICholDiagBlockPreconditioner(
    K::SparseMatrixCSC{Float64, Ti},
    r_u::AbstractRange{Int},
    r_p::AbstractRange{Int},
) where {Ti<:Integer}
    ru = UnitRange(first(r_u), last(r_u))
    rp = UnitRange(first(r_p), last(r_p))
    blks = saddle_point_matrix_blocks(K, ru, rp)
    nu = size(blks.A, 1)
    np = size(blks.C, 1)
    length(ru) == nu || error("ApproxSchurICholDiagBlockPreconditioner: |r_u| mismatch")
    length(rp) == np || error("ApproxSchurICholDiagBlockPreconditioner: |r_p| mismatch")

    Aint = convert(SparseMatrixCSC{Float64, Int}, blks.A)
    Pa = ICholPreconditioner(Aint)
    _, inv_schur = _inv_diag_A_and_schur(blks)
    tmp_u = zeros(Float64, nu)
    return ApproxSchurICholDiagBlockPreconditioner(Pa, inv_schur, ru, rp, tmp_u)
end

function ldiv!(
    y::AbstractVector{Float64},
    P::ApproxSchurICholDiagBlockPreconditioner,
    x::AbstractVector{Float64},
)
    length(y) == length(x) || throw(DimensionMismatch("ldiv!"))
    ru, rp = P.r_u, P.r_p
    @views copyto!(P.tmp_u, x[ru])
    ldiv!(P.Pa, P.tmp_u)
    @views y[ru] .= P.tmp_u
    @views @. y[rp] = P.inv_schur_diag * x[rp]
    return y
end

function ldiv!(P::ApproxSchurICholDiagBlockPreconditioner, x::AbstractVector{Float64})
    ru, rp = P.r_u, P.r_p
    @views copyto!(P.tmp_u, x[ru])
    ldiv!(P.Pa, P.tmp_u)
    @views x[ru] .= P.tmp_u
    @views @. x[rp] = P.inv_schur_diag * x[rp]
    return x
end
