using Sundials, Test, LinearAlgebra, IncompleteLU, SparseArrays
import AlgebraicMultigrid
import SparseConnectivityTracer, DifferentiationInterface, ADTypes, ForwardDiff

const N = 32
const xyd_brusselator = range(0; stop = 1, length = N)
brusselator_f(x, y, t) = (((x - 0.3)^2 + (y - 0.6)^2) <= 0.1^2) * (t >= 1.1) * 5.0
limit(a, N) = a == N + 1 ? 1 : a == 0 ? N : a

function brusselator_2d_loop(du, u, p, t)
    A, B, alpha, dx = p
    alpha = alpha / dx^2
    return @inbounds for I in CartesianIndices((N, N))
        i, j = Tuple(I)
        x, y = xyd_brusselator[I[1]], xyd_brusselator[I[2]]
        ip1, im1, jp1,
            jm1 = limit(i + 1, N), limit(i - 1, N), limit(j + 1, N),
            limit(j - 1, N)
        du[i, j, 1] = alpha * (
            u[im1, j, 1] + u[ip1, j, 1] + u[i, jp1, 1] + u[i, jm1, 1] -
                4u[i, j, 1]
        ) +
            B + u[i, j, 1]^2 * u[i, j, 2] - (A + 1) * u[i, j, 1] +
            brusselator_f(x, y, t)
        du[i, j, 2] = alpha * (
            u[im1, j, 2] + u[ip1, j, 2] + u[i, jp1, 2] + u[i, jm1, 2] -
                4u[i, j, 2]
        ) +
            A * u[i, j, 1] - u[i, j, 1]^2 * u[i, j, 2]
    end
end

# vector arguments to simplify Jacobian setup
function brusselator_2d_vec(du_vec, u_vec, p, t)
    du = reshape(du_vec, N, N, 2)
    u = reshape(u_vec, N, N, 2)
    return brusselator_2d_loop(du, u, p, t)
end

p = (3.4, 1.0, 10.0, step(xyd_brusselator))

function init_brusselator_2d(xyd)
    N = length(xyd)
    u = zeros(N, N, 2)
    for I in CartesianIndices((N, N))
        x = xyd[I[1]]
        y = xyd[I[2]]
        u[I, 1] = 22 * (y * (1 - y))^(3 / 2)
        u[I, 2] = 27 * (x * (1 - x))^(3 / 2)
    end
    return u
end
u0 = vec(init_brusselator_2d(xyd_brusselator))
du = similar(u0)

prob_ode_brusselator_2d = ODEProblem(
    brusselator_2d_vec,
    u0, (0.0, 11.5), p
)

detector = SparseConnectivityTracer.TracerSparsityDetector()
brus_uf = (du, u) -> brusselator_2d_vec(du, u, p, 0.1)
const jaccache = similar(SparseConnectivityTracer.jacobian_sparsity(brus_uf, du, u0, detector), Float64)
const W = I - 1.0 * jaccache

# setup sparse AD for Jacobian
# Setup sparse jacobian computation using DifferentiationInterface
# NOTE: AutoSparse doesn't support in-place functions with ForwardDiff in DI v0.7
# (requires pushforward performance which isn't defined for this combination)
# So we use regular AutoForwardDiff() and compute dense then copy to sparse
const backend = DifferentiationInterface.AutoForwardDiff()
# Don't prepare with a specific function - we'll prepare fresh in each callback
# since the function changes with different p and t values

prectmp = ilu(W; τ = 50.0)
const preccache = Ref(prectmp)

function psetupilu(p, t, u, du, jok, jcurPtr, gamma)
    return if jok
        # Compute jacobian using DifferentiationInterface with ForwardDiff backend
        # Create a wrapper that captures p and t
        f_wrapper! = (y, x) -> brusselator_2d_vec(y, x, p, t)
        # Prepare jacobian for this specific function
        prep_local = DifferentiationInterface.prepare_jacobian(f_wrapper!, du, backend, u)
        # Compute dense jacobian first, then copy to sparse
        jac_dense = Matrix(jaccache)
        DifferentiationInterface.jacobian!(f_wrapper!, du, jac_dense, prep_local, backend, u)
        # Copy non-zero entries back to sparse matrix
        for (i, j, v) in zip(findnz(jaccache)...)
            jaccache[i, j] = jac_dense[i, j]
        end
        jcurPtr[] = true

        # W = I - gamma*J
        @. W = -gamma * jaccache
        idxs = diagind(W)
        @. @view(W[idxs]) = @view(W[idxs]) + 1

        # Build preconditioner on W
        preccache[] = ilu(W; τ = 5.0)
    end
end

function precilu(z, r, p, t, y, fy, gamma, delta, lr)
    return ldiv!(z, preccache[], r)
end

# AlgebraicMultigrid can fail with LAPACK errors on some systems
prectmp2 = try
    AlgebraicMultigrid.aspreconditioner(
        AlgebraicMultigrid.ruge_stuben(
            W;
            presmoother = AlgebraicMultigrid.Jacobi(
                rand(
                    size(
                        W,
                        1
                    )
                )
            ),
            postsmoother = AlgebraicMultigrid.Jacobi(
                rand(
                    size(
                        W,
                        1
                    )
                )
            )
        )
    )
catch e
    @warn "AlgebraicMultigrid setup failed, using identity preconditioner as fallback" exception = e
    nothing
end
const preccache2 = Ref(prectmp2)
function psetupamg(p, t, u, du, jok, jcurPtr, gamma)
    if preccache2[] === nothing
        return  # Skip setup if AMG failed initially
    end

    return if jok
        # Compute jacobian using DifferentiationInterface with ForwardDiff backend
        # Create a wrapper that captures p and t
        f_wrapper! = (y, x) -> brusselator_2d_vec(y, x, p, t)
        # Prepare jacobian for this specific function
        prep_local = DifferentiationInterface.prepare_jacobian(f_wrapper!, du, backend, u)
        # Compute dense jacobian first, then copy to sparse
        jac_dense = Matrix(jaccache)
        DifferentiationInterface.jacobian!(f_wrapper!, du, jac_dense, prep_local, backend, u)
        # Copy non-zero entries back to sparse matrix
        for (i, j, v) in zip(findnz(jaccache)...)
            jaccache[i, j] = jac_dense[i, j]
        end
        jcurPtr[] = true

        # W = I - gamma*J
        @. W = -gamma * jaccache
        idxs = diagind(W)
        @. @view(W[idxs]) = @view(W[idxs]) + 1

        # Build preconditioner on W
        try
            preccache2[] = AlgebraicMultigrid.aspreconditioner(
                AlgebraicMultigrid.ruge_stuben(
                    W;
                    presmoother = AlgebraicMultigrid.Jacobi(
                        rand(
                            size(
                                W,
                                1
                            )
                        )
                    ),
                    postsmoother = AlgebraicMultigrid.Jacobi(
                        rand(
                            size(
                                W,
                                1
                            )
                        )
                    )
                )
            )
        catch e
            @warn "AlgebraicMultigrid update failed in psetupamg" exception = e
        end
    end
end

function precamg(z, r, p, t, y, fy, gamma, delta, lr)
    return if preccache2[] === nothing
        # Identity preconditioner fallback
        z .= r
    else
        ldiv!(z, preccache2[], r)
    end
end

sol1 = solve(
    prob_ode_brusselator_2d, CVODE_BDF(; linear_solver = :GMRES);
    save_everystep = false
);
sol2 = solve(
    prob_ode_brusselator_2d,
    CVODE_BDF(;
        linear_solver = :GMRES, prec = precilu, psetup = psetupilu,
        prec_side = 1
    ); save_everystep = false
);
sol3 = solve(
    prob_ode_brusselator_2d,
    CVODE_BDF(;
        linear_solver = :GMRES, prec = precamg, psetup = psetupamg,
        prec_side = 1
    ); save_everystep = false
);
@test sol1.stats.nf > sol2.stats.nf
# AlgebraicMultigrid can fail with LAPACK errors - mark as broken if it failed
if preccache2[] === nothing
    @test_broken sol1.stats.nf > sol3.stats.nf
else
    @test sol1.stats.nf > sol3.stats.nf
end
