# Test based on https://julianlsolvers.github.io/LineSearches.jl/stable/examples/generated/customoptimizer.html
# Note: Enzyme tests are in a separate test group (test/enzyme/), LineSearches.jl tests in test/linesearchesjl/
using LineSearch, Test

module CustomOptimizer
    using LinearAlgebra, SciMLBase, LineSearch, SciMLJacobianOperators

    function gradient_descent(
            prob, alg; g_atol::Real = 1.0e-5, maxiters::Int = 10000, autodiff = nothing
        )
        u = copy(prob.u0)
        if SciMLBase.isinplace(prob)
            fu = similar(u)
            prob.f(fu, u, prob.p)
        else
            fu = prob.f(u, prob.p)
        end

        ls_cache = init(prob, alg, fu, u)
        vjp_op = VecJacOperator(prob, fu, u; autodiff)

        alphas = Float64[]
        iter = 0
        for _ in 1:maxiters
            iter += 1
            svjp_op = StatefulJacobianOperator(vjp_op, u, prob.p)
            gs = svjp_op * fu .* 2
            δu = -gs

            ls_sol = solve!(ls_cache, u, δu)

            push!(alphas, ls_sol.step_size)
            @. u = u + ls_sol.step_size * δu
            gnorm = norm(gs)

            if SciMLBase.isinplace(prob)
                fu = similar(u)
                prob.f(fu, u, prob.p)
            else
                fu = prob.f(u, prob.p)
            end

            LineSearch.callback_into_cache!(ls_cache, fu)

            gnorm < g_atol && break
        end

        return fu, u, iter, alphas
    end

    export gradient_descent

end

using .CustomOptimizer

@testset "Native Line Search: Custom Optimizer" begin
    using SciMLBase
    using ADTypes, Tracker, ForwardDiff, Zygote, ReverseDiff, FiniteDiff

    @testset "OOP Problem" begin
        nlf(x, p) = [p[1] - x[1], 10.0 * (x[2] - x[1]^2)]
        nlp = NonlinearProblem(nlf, [-1.0, 1.0], [1.0])

        @testset for autodiff in (
                AutoTracker(), AutoForwardDiff(), AutoZygote(),
                AutoReverseDiff(), AutoFiniteDiff(),
            )
            @testset "method: $(nameof(typeof(method)))" for method in (
                    LiFukushimaLineSearch(),
                    NoLineSearch(0.001),
                    GoldenSection(; tol = 1.0e-4),
                    BackTracking(; order = Val(3), autodiff),
                    BackTracking(; order = Val(2), autodiff),
                    StrongWolfeLineSearch(; autodiff),
                )
                fu, u, iter, alphas = gradient_descent(nlp, method; autodiff)

                @test fu ≈ [0.0, 0.0] atol = 1.0e-1
                @test u ≈ [1.0, 1.0] atol = 1.0e-1
                @test !all(isone, alphas)
            end
        end
    end

    @testset "In-Place Problem" begin
        nlf(dx, x, p) = (dx .= [p[1] - x[1], 10.0 * (x[2] - x[1]^2)])
        nlp = NonlinearProblem(nlf, [-1.0, 1.0], [1.0])

        @testset for autodiff in (
                AutoForwardDiff(), AutoReverseDiff(), AutoFiniteDiff(),
            )
            @testset "method: $(nameof(typeof(method)))" for method in (
                    LiFukushimaLineSearch(),
                    NoLineSearch(0.001),
                    GoldenSection(; tol = 1.0e-4),
                    BackTracking(; order = Val(3), autodiff),
                    BackTracking(; order = Val(2), autodiff),
                    StrongWolfeLineSearch(; autodiff),
                )
                fu, u, iter, alphas = gradient_descent(nlp, method; autodiff)

                @test fu ≈ [0.0, 0.0] atol = 1.0e-1
                @test u ≈ [1.0, 1.0] atol = 1.0e-1
                @test !all(isone, alphas)
            end
        end
    end
end
