include("../../testutils.jl")

using ADTypes: ADTypes
using DifferentiationInterface, DifferentiationInterfaceTest
import DifferentiationInterfaceTest as DIT
using Enzyme: Enzyme
using LinearAlgebra
using StaticArrays
using Test

using ExplicitImports
check_no_implicit_imports(DifferentiationInterface)

backends = [
    AutoEnzyme(; mode = nothing),
    AutoEnzyme(; mode = Enzyme.Forward),
    AutoEnzyme(; mode = Enzyme.Reverse),
    AutoEnzyme(; mode = nothing, function_annotation = Enzyme.Const),
]

duplicated_backends = [
    AutoEnzyme(; mode = Enzyme.Forward, function_annotation = Enzyme.Duplicated),
    AutoEnzyme(; mode = Enzyme.Reverse, function_annotation = Enzyme.Duplicated),
]

@testset "Checks" begin
    @testset "Check $(typeof(backend))" for backend in backends
        @test check_available(backend)
        @test check_inplace(backend)
        test_counterparts(backend)
    end
end;

@testset "Counterpart mode attributes" begin
    # a reverse mode with non-default attributes, to check they survive the flip
    rev_mode = Enzyme.set_runtime_activity(Enzyme.ReverseWithPrimal)
    rev = AutoEnzyme(; mode = rev_mode, function_annotation = Enzyme.Const)
    fwd = DifferentiationInterface.forward_counterpart(rev)
    @test ADTypes.mode(fwd) isa ADTypes.ForwardMode
    @test fwd isa AutoEnzyme{<:Any, Enzyme.Const}  # function annotation preserved
    # ReturnPrimal and RuntimeActivity must carry over (ForwardMode{ReturnPrimal,...,RuntimeActivity,...})
    @test fwd.mode isa Enzyme.EnzymeCore.ForwardMode{true, <:Any, <:Any, true}
    # round-trip back to reverse recovers the original mode
    @test DifferentiationInterface.reverse_counterpart(fwd) === rev
    # split reverse mode also maps to forward mode, carrying over its attributes
    split_mode = Enzyme.set_runtime_activity(Enzyme.ReverseSplitWithPrimal)
    split = AutoEnzyme(; mode = split_mode, function_annotation = Enzyme.Const)
    fwd_split = DifferentiationInterface.forward_counterpart(split)
    @test ADTypes.mode(fwd_split) isa ADTypes.ForwardMode
    @test fwd_split isa AutoEnzyme{<:Any, Enzyme.Const}
    @test fwd_split.mode isa Enzyme.EnzymeCore.ForwardMode{true, <:Any, <:Any, true}
end;

@testset "First order" begin
    test_differentiation(
        backends, default_scenarios(); excluded = SECOND_ORDER, logging = LOGGING
    )

    test_differentiation(
        backends[1:3],
        default_scenarios(; include_normal = false, include_constantified = true);
        excluded = SECOND_ORDER,
        logging = LOGGING,
    )

    test_differentiation(
        backends[2:3],
        default_scenarios(;
            include_normal = false,
            include_cachified = true,
            include_constantorcachified = true,
            use_tuples = true,
        );
        excluded = SECOND_ORDER,
        logging = LOGGING,
    )

    test_differentiation(
        duplicated_backends,
        default_scenarios(; include_normal = false, include_closurified = true);
        excluded = SECOND_ORDER,
        logging = LOGGING,
    )
end

@testset "Second order" begin
    test_differentiation(
        AutoEnzyme(),
        default_scenarios(; include_constantified = true, include_cachified = true);
        excluded = vcat(FIRST_ORDER, :hvp, :hessian),
        logging = LOGGING,
    )

    test_differentiation(
        AutoEnzyme(),
        default_scenarios(; include_constantified = true);
        excluded = vcat(FIRST_ORDER, :second_derivative),
        logging = LOGGING,
    )

    test_differentiation(
        # TODO: simplify when https://github.com/EnzymeAD/Enzyme.jl/issues/2854 and https://github.com/EnzymeAD/Enzyme.jl/issues/2925 are fixed
        if VERSION >= v"1.11"
            SecondOrder(
                AutoEnzyme(; mode = Enzyme.set_runtime_activity(Enzyme.Forward)),
                AutoEnzyme(; mode = Enzyme.set_runtime_activity(Enzyme.Reverse))
            )
        else
            AutoEnzyme()
        end,
        default_scenarios(; include_normal = false, include_constantified = false, include_cachified = true);
        excluded = vcat(FIRST_ORDER, :second_derivative),
        logging = LOGGING,
    )

    test_differentiation(
        [
            SecondOrder(
                AutoEnzyme(; mode = Enzyme.Reverse), AutoEnzyme(; mode = Enzyme.Forward)
            ),
        ],
        default_scenarios(; include_constantified = true, include_cachified = true);
        excluded = FIRST_ORDER,
        logging = LOGGING,
    )
end

@testset "Sparse" begin
    test_differentiation(
        MyAutoSparse.(AutoEnzyme(; function_annotation = Enzyme.Const)),
        if VERSION < v"1.11"
            sparse_scenarios()
        else
            filter(s -> s.x isa AbstractVector, sparse_scenarios())
        end;
        sparsity = true,
        logging = LOGGING,
    )
end

@testset "Static" begin
    filtered_static_scenarios = filter(static_scenarios()) do s
        DIT.operator_place(s) == :out && DIT.function_place(s) == :out
    end

    test_differentiation(
        [AutoEnzyme(; mode = Enzyme.Forward), AutoEnzyme(; mode = Enzyme.Reverse)],
        filtered_static_scenarios;
        excluded = SECOND_ORDER,
        logging = LOGGING,
    )
end

@testset "Coverage" begin
    # ConstantOrCache without cache
    f_nocontext(x, p) = x
    @test I == DifferentiationInterface.jacobian(
        f_nocontext, AutoEnzyme(; mode = Enzyme.Forward), rand(10), ConstantOrCache(nothing)
    )
    @test I == DifferentiationInterface.jacobian(
        f_nocontext, AutoEnzyme(; mode = Enzyme.Reverse), rand(10), ConstantOrCache(nothing)
    )
end

@testset "Hints" begin
    @testset "MutabilityError" begin
        f = let
            cache = [0.0]
            x -> sum(copyto!(cache, x))
        end

        e = nothing
        try
            gradient(f, AutoEnzyme(), [1.0])
        catch e
        end
        msg = sprint(showerror, e)
        @test occursin("AutoEnzyme", msg)
        @test occursin("function_annotation", msg)
        @test occursin("ADTypes", msg)
    end

    @testset "RuntimeActivityError" begin
        function g(active_var, constant_var, cond)
            if cond
                return active_var
            else
                return constant_var
            end
        end

        function h(active_var, constant_var, cond)
            return [g(active_var, constant_var, cond), g(active_var, constant_var, cond)]
        end

        e = nothing
        try
            pushforward(
                h,
                AutoEnzyme(; mode = Enzyme.Forward),
                [1.0],
                ([1.0],),
                Constant([1.0]),
                Constant(true),
            )
        catch e
        end
        msg = sprint(showerror, e)
        @test_broken occursin("AutoEnzyme", msg)
        @test_broken occursin("ADTypes", msg)
    end
end

@testset "Empty arrays" begin
    test_differentiation(
        [AutoEnzyme(; mode = Enzyme.Forward), AutoEnzyme(; mode = Enzyme.Reverse)],
        empty_scenarios();
        excluded = [:jacobian],
    )
end;
