using SparseArrays, LinearSolve, LinearAlgebra

struct JacobiInstance{T} 
    invdiag::Vector{T}
end

jacobi(A::AbstractMatrix{T}) where T =JacobiInstance([ 1/A[i,i] for i=1:size(A,2)])

LinearAlgebra.ldiv!(u, M::JacobiInstance, v) =  u.=M.invdiag.*v

num_prec_calls::Int=0

struct JacobiPreconditioner end

function (prec::JacobiPreconditioner)(A, p)
    global num_prec_calls
    num_prec_calls+=1
    (prec(A), I)
end



function test_jacobi(;n=10,linsolver = KrylovJL_CG(precs=JacobiPreconditioner()))
    global num_prec_calls

    num_prec_calls=0
    
    A=spdiagm( -1 => -ones(n-1), 0 => fill(10.0,n), 1=>-ones(n-1))
    b=rand(n)
    
    p = LinearProblem(A, b)
    x0=solve(p,linsolver)
    cache = x0.cache
    x0 = copy(x0)
    for i = 4:(n - 3)
        A[i, i + 3] -= 1.0e-4
        A[i - 3, i] -= 1.0e-4
    end
    LinearSolve.reinit!(cache;A)
    x1 = copy(solve!(cache, reuse_precs=true))

    # This must be true
    all(x0 .< x1) && num_prec_calls==1
end

@show test_jacobi()