using Test
using SparseArrays
using OrdinaryDiffEqCore: find_algebraic_vars_eqs, get_differential_vars
using LinearAlgebra
using SciMLBase: ODEFunction
using SciMLOperators: IdentityOperator, ScalarOperator

@testset "Sparse Algebraic Detection Performance" begin
    # Test 1: Correctness - results should match between sparse and dense methods
    @testset "Correctness" begin
        # Small mixed matrix
        M1 = sparse([1.0 0.0 0.0; 0.0 0.0 0.0; 0.0 0.0 2.0])
        vars1, eqs1 = find_algebraic_vars_eqs(M1)
        @test vars1 == [false, true, false]  # col 2 is zero (algebraic)
        @test eqs1 == [false, true, false]   # row 2 is zero (algebraic)

        # Diagonal matrix (no algebraic)
        M2 = sparse(1:10, 1:10, ones(10))
        vars2, eqs2 = find_algebraic_vars_eqs(M2)
        @test all(.!vars2)  # no algebraic vars
        @test all(.!eqs2)   # no algebraic eqs

        # Zero matrix (all algebraic)
        M3 = spzeros(5, 5)
        vars3, eqs3 = find_algebraic_vars_eqs(M3)
        @test all(vars3)  # all algebraic vars
        @test all(eqs3)   # all algebraic eqs

        # Matrix with explicit zeros in sparsity pattern
        M4 = sparse([1, 2], [1, 2], [1.0, 0.0], 3, 3)  # (2,2) is stored but zero
        vars4, eqs4 = find_algebraic_vars_eqs(M4)
        @test vars4 == [false, true, true]   # cols 2,3 are effectively zero
        @test eqs4 == [false, true, true]    # rows 2,3 are effectively zero
    end

    # Test 2: Complexity verification
    # For sparse matrices, time should scale with nnz, not n²
    # We verify this by checking that processing a large sparse diagonal matrix
    # is much faster than it would be with O(n²) complexity
    @testset "Complexity O(nnz) not O(n²)" begin
        # Create diagonal matrices of increasing size
        # For a diagonal matrix: nnz = n, so O(nnz) = O(n)
        # If algorithm were O(n²), doubling n would 4x the time
        # With O(nnz) = O(n), doubling n should ~2x the time

        n_small = 10_000
        n_large = 40_000

        M_small = sparse(1:n_small, 1:n_small, ones(n_small))
        M_large = sparse(1:n_large, 1:n_large, ones(n_large))

        # Warmup
        find_algebraic_vars_eqs(M_small)
        find_algebraic_vars_eqs(M_large)

        # Measure using best-of-N to reduce scheduler / GC noise on shared runners.
        # Ratio still reflects algorithmic complexity: O(n^2) would yield a ratio
        # near 16 even with the minimum, while O(nnz) = O(n) yields ~4.
        function _best_time(f, M)
            best = Inf
            for _ in 1:5
                t = @elapsed for _ in 1:10
                    f(M)
                end
                best = min(best, t)
            end
            return best
        end
        t_small = _best_time(find_algebraic_vars_eqs, M_small)
        t_large = _best_time(find_algebraic_vars_eqs, M_large)

        ratio = t_large / t_small
        size_ratio = n_large / n_small  # 4x size increase

        # With O(n²), ratio would be ~16
        # With O(nnz) = O(n), ratio should be ~4
        # Allow some tolerance for measurement noise
        @test ratio < size_ratio * 2  # Should be closer to linear than quadratic

        # Absolute sanity check: processing 40k diagonal shouldn't take > 100ms.
        # Use minimum of a few runs for the same reason (single @elapsed can be
        # dominated by a stray GC pause on shared CI runners).
        t_single = minimum(@elapsed(find_algebraic_vars_eqs(M_large)) for _ in 1:5)
        @test t_single < 0.1  # Should be < 100ms, typically < 1ms
    end

    # Test 3: Verify fallback for dense matrices still works
    @testset "Dense matrix fallback" begin
        M_dense = [1.0 0.0; 0.0 0.0]
        vars, eqs = find_algebraic_vars_eqs(M_dense)
        @test vars == [false, true]
        @test eqs == [false, true]
    end

    # Test 4: Test Diagonal case
    @testset "Test Diagonal cast" begin
        M_diag = Diagonal([1.0, 0.0, 1.0, 1.0, 0.0])
        vars, eqs = find_algebraic_vars_eqs(M_diag)
        @test vars == [false, true, false, false, true]
        @test eqs == [false, true, false, false, true]
        # compare to dense
        @test find_algebraic_vars_eqs(M_diag) == find_algebraic_vars_eqs(collect(M_diag))
    end
end

@testset "get_differential_vars: operator mass matrix" begin
    f!(du, u, p, t) = (du .= -u; nothing)
    u = ones(3)
    ff_id = ODEFunction(f!, mass_matrix = IdentityOperator(3))
    @test get_differential_vars(ff_id, u) === nothing
    ff_scalar = ODEFunction(f!, mass_matrix = ScalarOperator(2.0))
    @test get_differential_vars(ff_scalar, u) === nothing
    ff_singular = ODEFunction(f!, mass_matrix = ScalarOperator(0.0))
    @test get_differential_vars(ff_singular, u) == falses(3)
end

@testset "get_differential_vars: Diagonal mass matrix" begin
    f!(du, u, p, t) = (du .= -u; nothing)
    u = zeros(4)
    ff_mixed = ODEFunction(f!, mass_matrix = Diagonal([1.0, 1.0, 0.0, 0.0]))
    @test get_differential_vars(ff_mixed, u) == Bool[1, 1, 0, 0]
    ff_alldiff = ODEFunction(f!, mass_matrix = Diagonal([1.0, 2.0, 3.0, 4.0]))
    @test get_differential_vars(ff_alldiff, u) == Bool[1, 1, 1, 1]
    ff_allalg = ODEFunction(f!, mass_matrix = Diagonal(zeros(4)))
    @test get_differential_vars(ff_allalg, u) == Bool[0, 0, 0, 0]
end

# Mimics GPU device-array semantics on the CPU: elementwise scalar `getindex`
# is forbidden (like CUDA's default `allowscalar(false)`), while whole-array
# broadcast still works. Any code path that walks a `Diagonal` of this via
# `getindex` -- `all(!iszero, mm)` or a `diag(mm)` copy -- errors, so the test
# below catches a regression back to the whole-matrix check without needing a
# GPU. Must be top-level: structs cannot be defined inside a `@testset`.
struct NoScalarIndexVector{T} <: AbstractVector{T}
    data::Vector{T}
end
Base.size(v::NoScalarIndexVector) = size(v.data)
function Base.getindex(::NoScalarIndexVector, ::Int)
    error("scalar getindex disallowed on NoScalarIndexVector")
end
Base.Broadcast.broadcastable(v::NoScalarIndexVector) = v.data

@testset "get_differential_vars: no scalar indexing of Diagonal mass matrix" begin
    f!(du, u, p, t) = (du .= -u; nothing)
    u = zeros(4)
    mm = Diagonal(NoScalarIndexVector([1.0, 1.0, 0.0, 0.0]))
    ff = ODEFunction(f!, mass_matrix = mm)
    @test get_differential_vars(ff, u) == Bool[1, 1, 0, 0]
end
