include(joinpath(@__DIR__, "shared", "test_setup.jl"))

using SparseArrays: sparse

struct ExternalADType <: AbstractADType end
ADTypes.mode(::ExternalADType) = ForwardMode()

struct ExternalSparsityDetector <: ADTypes.AbstractSparsityDetector end
ADTypes.jacobian_sparsity(f, x, ::ExternalSparsityDetector) = trues(length(f(x)), length(x))
ADTypes.jacobian_sparsity(f!, y, x, ::ExternalSparsityDetector) = trues(length(y), length(x))
ADTypes.hessian_sparsity(f, x, ::ExternalSparsityDetector) = trues(length(x), length(x))

struct ExternalColoringAlgorithm <: ADTypes.AbstractColoringAlgorithm end
ADTypes.column_coloring(M, ::ExternalColoringAlgorithm) = 1:size(M, 2)
ADTypes.row_coloring(M, ::ExternalColoringAlgorithm) = 1:size(M, 1)
ADTypes.symmetric_coloring(M, ::ExternalColoringAlgorithm) = 1:size(M, 1)

@testset "External public interface contracts" begin
    ad = ExternalADType()
    @test mode(ad) isa ForwardMode

    x = rand(2)
    y = similar(x)
    f(x) = vcat(x, x)
    f!(y, x) = y .= x
    detector = ExternalSparsityDetector()
    @test size(jacobian_sparsity(f, x, detector)) == (4, 2)
    @test size(jacobian_sparsity(f!, y, x, detector)) == (2, 2)
    @test size(hessian_sparsity(sum, x, detector)) == (2, 2)

    colors = ExternalColoringAlgorithm()
    M = trues(2, 3)
    @test column_coloring(M, colors) == 1:3
    @test row_coloring(M, colors) == 1:2
    @test symmetric_coloring(trues(2, 2), colors) == 1:2
end

@testset "Subtyping" begin
    for ad in every_ad()
        sparse_ad = AutoSparse(ad)
        @test sparse_ad isa AbstractADType
        if mode(ad) isa ForwardMode
            @test mode(sparse_ad) isa ForwardMode
        elseif mode(ad) isa ForwardOrReverseMode
            @test mode(sparse_ad) isa ForwardOrReverseMode
        elseif mode(ad) isa ReverseMode
            @test mode(sparse_ad) isa ReverseMode
        elseif mode(ad) isa SymbolicMode
            @test mode(sparse_ad) isa SymbolicMode
        end
        @test dense_ad(ad) == ad
        @test dense_ad(sparse_ad) == ad
        @test sparsity_detector(sparse_ad) == NoSparsityDetector()
        @test coloring_algorithm(sparse_ad) == NoColoringAlgorithm()
    end
end

@testset "Sparsity detector" begin
    f_jac1(x) = vcat(x, x)
    f_jac2(x) = hcat(x, x)
    function f_jac! end
    f_hess(x) = sum(x)

    @testset "NoSparsityDetector" begin
        sd = NoSparsityDetector()

        for x in (rand(2), rand(2, 3)), f in (f_jac1, f_jac2)

            y = f(x)
            Js = jacobian_sparsity(f, x, sd)
            @test Js isa AbstractMatrix
            @test size(Js) == (length(y), length(x))
            @test all(isone, Js)
        end

        for x in (rand(2), rand(2, 3)), y in (rand(5), rand(5, 6))

            Js = jacobian_sparsity(f_jac!, y, x, sd)
            @test Js isa AbstractMatrix
            @test size(Js) == (length(y), length(x))
            @test all(isone, Js)
        end

        for x in (rand(2), rand(2, 3))
            Hs = hessian_sparsity(f_hess, x, sd)
            @test Hs isa AbstractMatrix
            @test size(Hs) == (length(x), length(x))
            @test all(isone, Hs)
        end
    end

    @testset "KnownJacobianSparsityDetector" begin
        @testset "Pattern element types" begin
            # Consumers treat any nonzero entry as a structural nonzero, so patterns
            # written as integer or float ones and zeroes are stored and returned as-is.
            Jbool = rand(Bool, 4, 2)
            x = rand(2)
            y = rand(4)
            for J in (
                    Jbool, BitMatrix(Jbool), Int.(Jbool), UInt8.(Jbool),
                    Float64.(Jbool), Float32.(Jbool), sparse(Float64.(Jbool)),
                    Any[Int(v) for v in Jbool],
                )
                sd = KnownJacobianSparsityDetector(J)
                @test sd isa KnownJacobianSparsityDetector{typeof(J)}
                @test jacobian_sparsity(f_jac1, x, sd) === J
                @test jacobian_sparsity(f_jac!, y, x, sd) === J
                @test eltype(jacobian_sparsity(f_jac1, x, sd)) === eltype(J)
            end
        end

        @testset "Jacobian sparsity detection" begin
            @testset "Out-of-place functions" begin
                for sx in ((2,), (2, 3)), f in (f_jac1, f_jac2)

                    x = rand(sx...)
                    nx = length(x)
                    Jref = rand(Bool, 2 * nx, nx)
                    sd = KnownJacobianSparsityDetector(Jref)
                    Js = jacobian_sparsity(f, x, sd)
                    @test Js isa AbstractMatrix
                    @test size(Js) == (2 * nx, nx)
                    @test Js === Jref
                end
            end
            @testset "In-place functions" begin
                for sx in ((2,), (2, 3)), sy in ((5,), (5, 6))

                    x, y = rand(sx...), rand(sy...)
                    nx, ny = length(x), length(y)
                    Jref = rand(Bool, ny, nx)
                    sd = KnownJacobianSparsityDetector(Jref)
                    Js = jacobian_sparsity(f_jac!, y, x, sd)
                    @test Js isa AbstractMatrix
                    @test size(Js) == (ny, nx)
                    @test Js === Jref
                end
            end
        end
        @testset "Exceptions: Hessian sparsity detection not supported" begin
            for sx in ((2,), (2, 3))
                x = rand(sx...)
                nx = length(x)
                Href = rand(Bool, nx, nx)
                sd = KnownJacobianSparsityDetector(Href)
                @test_throws ArgumentError hessian_sparsity(f_hess, x, sd)
            end
        end
        @testset "Exceptions: DimensionMismatch" begin
            # wrong Jacobian size
            for J in (rand(Bool, 6, 7), Int.(rand(Bool, 6, 7)), Float64.(rand(Bool, 6, 7)))
                sd = KnownJacobianSparsityDetector(J)
                for x in (rand(2), rand(2, 3)), f in (f_jac1, f_jac2)

                    @test_throws DimensionMismatch jacobian_sparsity(f, x, sd)
                end
                for x in (rand(2), rand(2, 3)), y in (rand(5), rand(5, 6))

                    @test_throws DimensionMismatch jacobian_sparsity(f_jac!, y, x, sd)
                end
            end
        end
    end

    @testset "KnownHessianSparsityDetector" begin
        @testset "Pattern element types" begin
            Hbool = rand(Bool, 2, 2)
            x = rand(2)
            for H in (
                    Hbool, BitMatrix(Hbool), Int.(Hbool), UInt8.(Hbool),
                    Float64.(Hbool), Float32.(Hbool), sparse(Float64.(Hbool)),
                    Any[Int(v) for v in Hbool],
                )
                sd = KnownHessianSparsityDetector(H)
                @test sd isa KnownHessianSparsityDetector{typeof(H)}
                @test hessian_sparsity(f_hess, x, sd) === H
                @test eltype(hessian_sparsity(f_hess, x, sd)) === eltype(H)
            end
        end

        @testset "Hessian sparsity detection" begin
            for sx in ((2,), (2, 3))
                x = rand(sx...)
                nx = length(x)
                Href = rand(Bool, nx, nx)
                sd = KnownHessianSparsityDetector(Href)

                Hs = hessian_sparsity(f_hess, x, sd)
                @test Hs isa AbstractMatrix
                @test size(Hs) == (nx, nx)
                @test Hs === Href
            end
        end
        @testset "Exceptions: Jacobian sparsity detection not supported" begin
            @testset "Out-of-place functions" begin
                for sx in ((2,), (2, 3)), f in (f_jac1, f_jac2)

                    x = rand(sx...)
                    nx = length(x)
                    Jref = rand(Bool, 2 * nx, nx)
                    sd = KnownHessianSparsityDetector(Jref)
                    @test_throws ArgumentError jacobian_sparsity(f, x, sd)
                end
            end
            @testset "In-place functions" begin
                for sx in ((2,), (2, 3)), sy in ((5,), (5, 6))

                    x, y = rand(sx...), rand(sy...)
                    nx, ny = length(x), length(y)
                    Jref = rand(Bool, ny, nx)
                    sd = KnownHessianSparsityDetector(Jref)
                    @test_throws ArgumentError jacobian_sparsity(f_jac!, y, x, sd)
                end
            end
        end
        @testset "Exceptions: DimensionMismatch" begin
            # wrong Hessian size
            for H in (rand(Bool, 2, 3), Int.(rand(Bool, 2, 3)), Float64.(rand(Bool, 2, 3)))
                sd = KnownHessianSparsityDetector(H)
                for x in (rand(2), rand(2, 3))
                    @test_throws DimensionMismatch hessian_sparsity(f_hess, x, sd)
                end
            end
        end
    end
end

@testset "Coloring algorithm" begin
    ca = NoColoringAlgorithm()

    for M in (rand(2, 3), rand(3, 2))
        cv = column_coloring(M, ca)
        @test cv isa AbstractVector{<:Integer}
        @test length(cv) == size(M, 2)
        @test allunique(cv)

        rv = row_coloring(M, ca)
        @test rv isa AbstractVector{<:Integer}
        @test length(rv) == size(M, 1)
        @test allunique(rv)
    end

    M = rand(3, 3)
    sv = symmetric_coloring(M, ca)
    @test sv isa AbstractVector{<:Integer}
    @test length(sv) == size(M, 1) == size(M, 2)
    @test allunique(sv)
end
