using AbstractDifferentiation
using Test, LinearAlgebra
using Random
Random.seed!(1234)

fder(x, y) = exp(y) * x + y * log(x)
dfderdx(x, y) = exp(y) + y * 1 / x
dfderdy(x, y) = exp(y) * x + log(x)
dfderdxdx(x, y) = -y / x^2

fgrad(x, y) = prod(x) + sum(y ./ (1:length(y)))
dfgraddx(x, y) = prod(x) ./ x
dfgraddy(x, y) = one(eltype(y)) ./ (1:length(y))
dfgraddxdx(x, y) = prod(x) ./ (x * x') - Diagonal(diag(prod(x) ./ (x * x')))
dfgraddydy(x, y) = zeros(length(y), length(y))

function fjac(x, y)
    return x + -3 * y + [y[2:end]; zero(y[end])] / 2# Bidiagonal(-ones(length(y)) * 3, ones(length(y) - 1) / 2, :U) * y
end
dfjacdx(x, y) = I(length(x))
dfjacdy(x, y) = Bidiagonal(-ones(length(y)) * 3, ones(length(y) - 1) / 2, :U)

# Jvp
jxvp(x, y, v) = dfjacdx(x, y) * v
jyvp(x, y, v) = dfjacdy(x, y) * v

# vJp
vJxp(x, y, v) = dfjacdx(x, y)' * v
vJyp(x, y, v) = dfjacdy(x, y)' * v

const xscalar = rand()
const yscalar = rand()

const xvec = rand(5)
const yvec = rand(5)

# to check if vectors get mutated
xvec2 = deepcopy(xvec)
yvec2 = deepcopy(yvec)

function test_higher_order_backend(backends...)
    ADbackends = AD.HigherOrderBackend(backends)
    @test backends[end] == AD.lowest(ADbackends)
    @test backends[end - 1] == AD.second_lowest(ADbackends)

    for i in length(backends):-1:1
        @test backends[i] == AD.lowest(ADbackends)
        ADbackends = AD.reduce_order(ADbackends)
    end
    return backends[1] == AD.reduce_order(ADbackends)
end

function test_derivatives(backend; multiple_inputs=true, test_types=true)
    # test with respect to analytical solution
    der_exact = (dfderdx(xscalar, yscalar), dfderdy(xscalar, yscalar))
    if multiple_inputs
        der1 = AD.derivative(backend, fder, xscalar, yscalar)
        @test minimum(isapprox.(der_exact, der1, rtol=1e-10))
        valscalar, der2 = AD.value_and_derivative(backend, fder, xscalar, yscalar)
        @test valscalar == fder(xscalar, yscalar)
        @test der2 .- der1 == (0, 0)
    end
    # test if single input (no tuple works)
    valscalara, dera = AD.value_and_derivative(backend, x -> fder(x, yscalar), xscalar)
    valscalarb, derb = AD.value_and_derivative(backend, y -> fder(xscalar, y), yscalar)
    if test_types
        @test valscalara isa Float64
        @test valscalarb isa Float64
        @test dera[1] isa Float64
        @test derb[1] isa Float64
    end
    @test fder(xscalar, yscalar) == valscalara
    @test fder(xscalar, yscalar) == valscalarb
    @test isapprox(dera[1], der_exact[1], rtol=1e-10)
    @test isapprox(derb[1], der_exact[2], rtol=1e-10)
end

function test_gradients(backend; multiple_inputs=true, test_types=true)
    # test with respect to analytical solution
    grad_exact = (dfgraddx(xvec, yvec), dfgraddy(xvec, yvec))
    if multiple_inputs
        grad1 = AD.gradient(backend, fgrad, xvec, yvec)
        @test minimum(isapprox.(grad_exact, grad1, rtol=1e-10))
        valscalar, grad2 = AD.value_and_gradient(backend, fgrad, xvec, yvec)
        if test_types
            @test valscalar isa Float64
            @test grad1[1] isa AbstractVector{Float64}
            @test grad1[2] isa AbstractVector{Float64}
            @test grad2[1] isa AbstractVector{Float64}
            @test grad2[2] isa AbstractVector{Float64}
        end
        @test valscalar == fgrad(xvec, yvec)
        @test norm.(grad2 .- grad1) == (0, 0)
    end
    # test if single input (no tuple works)
    valscalara, grada = AD.value_and_gradient(backend, x -> fgrad(x, yvec), xvec)
    valscalarb, gradb = AD.value_and_gradient(backend, y -> fgrad(xvec, y), yvec)
    if test_types
        @test valscalara isa Float64
        @test valscalarb isa Float64
        @test grada[1] isa AbstractVector{Float64}
        @test gradb[1] isa AbstractVector{Float64}
    end
    @test fgrad(xvec, yvec) == valscalara
    @test fgrad(xvec, yvec) == valscalarb
    @test isapprox(grada[1], grad_exact[1], rtol=1e-10)
    @test isapprox(gradb[1], grad_exact[2], rtol=1e-10)
    @test xvec == xvec2
    @test yvec == yvec2
end

function test_jacobians(backend; multiple_inputs=true, test_types=true)
    # test with respect to analytical solution
    jac_exact = (dfjacdx(xvec, yvec), dfjacdy(xvec, yvec))
    if multiple_inputs
        jac1 = AD.jacobian(backend, fjac, xvec, yvec)
        @test minimum(isapprox.(jac_exact, jac1, rtol=1e-10))
        valvec, jac2 = AD.value_and_jacobian(backend, fjac, xvec, yvec)
        if test_types
            @test valvec isa Vector{Float64}
            @test jac1[1] isa Matrix{Float64}
            @test jac1[2] isa Matrix{Float64}
            @test jac2[1] isa Matrix{Float64}
            @test jac2[2] isa Matrix{Float64}
        end
        @test valvec == fjac(xvec, yvec)
        @test norm.(jac2 .- jac1) == (0, 0)
    end

    # test if single input (no tuple works)
    valveca, jaca = AD.value_and_jacobian(backend, x -> fjac(x, yvec), xvec)
    valvecb, jacb = AD.value_and_jacobian(backend, y -> fjac(xvec, y), yvec)
    if test_types
        @test valveca isa Vector{Float64}
        @test valvecb isa Vector{Float64}
        @test jaca[1] isa Matrix{Float64}
        @test jacb[1] isa Matrix{Float64}
    end
    @test fjac(xvec, yvec) == valveca
    @test fjac(xvec, yvec) == valvecb
    @test isapprox(jaca[1], jac_exact[1], rtol=1e-10)
    @test isapprox(jacb[1], jac_exact[2], rtol=1e-10)
    @test xvec == xvec2
    @test yvec == yvec2
end

function test_second_derivatives(backend; test_types=true)
    # explicit test that AbstractDifferentiation throws an error
    # don't support tuple of second derivatives
    @test_throws ArgumentError AD.second_derivative(
        backend, x -> fder(x, yscalar), (xscalar, yscalar)
    )
    @test_throws MethodError AD.second_derivative(
        backend, x -> fder(x, yscalar), xscalar, yscalar
    )

    # test if single input (no tuple works)
    dder1 = AD.second_derivative(backend, x -> fder(x, yscalar), xscalar)
    if test_types
        @test only(dder1) isa Float64
    end
    @test dfderdxdx(xscalar, yscalar) ≈ only(dder1) atol = 1e-8
    valscalar, dder2 = AD.value_and_second_derivative(
        backend, x -> fder(x, yscalar), xscalar
    )
    if test_types
        @test valscalar isa Float64
        @test only(dder2) isa Float64
    end
    @test valscalar == fder(xscalar, yscalar)
    @test dder2 == dder1
    valscalar, der, dder3 = AD.value_derivative_and_second_derivative(
        backend, x -> fder(x, yscalar), xscalar
    )
    if test_types
        @test valscalar isa Float64
        @test only(der) isa Float64
        @test only(dder3) isa Float64
    end
    @test valscalar == fder(xscalar, yscalar)
    @test der == AD.derivative(backend, x -> fder(x, yscalar), xscalar)
    @test dder3 == dder1
end

function test_hessians(backend; multiple_inputs=false, test_types=true)
    if multiple_inputs
        # ... but 
        error("multiple_inputs=true is not supported.")
    else
        # explicit test that AbstractDifferentiation throws an error
        # don't support tuple of Hessians
        @test_throws ArgumentError AD.hessian(backend, fgrad, (xvec, yvec))
        @test_throws MethodError AD.hessian(backend, fgrad, xvec, yvec)
    end

    # @test dfgraddxdx(xvec,yvec) ≈ H1[1] atol=1e-10
    # @test dfgraddydy(xvec,yvec) ≈ H1[2] atol=1e-10

    # test if single input (no tuple works)
    fhess = x -> fgrad(x, yvec)
    hess1 = AD.hessian(backend, fhess, xvec)
    if test_types
        @test hess1[1] isa Matrix{Float64}
    end
    # test with respect to analytical solution
    @test dfgraddxdx(xvec, yvec) ≈ hess1[1] atol = 1e-10

    valscalar, hess2 = AD.value_and_hessian(backend, fhess, xvec)
    if test_types
        @test valscalar isa Float64
        @test hess2[1] isa Matrix{Float64}
    end
    @test valscalar == fgrad(xvec, yvec)
    @test norm.(hess2 .- hess1) == (0,)
    valscalar, grad, hess3 = AD.value_gradient_and_hessian(backend, fhess, xvec)
    if test_types
        @test valscalar isa Float64
        @test grad[1] isa AbstractVector{Float64}
        @test hess3[1] isa Matrix{Float64}
    end
    @test valscalar == fgrad(xvec, yvec)
    @test norm.(grad .- AD.gradient(backend, fhess, xvec)) == (0,)
    @test norm.(hess3 .- hess1) == (0,)

    @test xvec == xvec2
    @test yvec == yvec2
    fhess2 = x -> dfgraddx(x, yvec)
    hess4 = AD.jacobian(backend, fhess2, xvec)
    if test_types
        @test hess4[1] isa Matrix{Float64}
    end
    @test minimum(isapprox.(hess4, hess1, atol=1e-10))
end

function test_jvp(
    backend; multiple_inputs=true, vaugmented=false, rng=Random.GLOBAL_RNG, test_types=true
)
    v = (rand(rng, length(xvec)), rand(rng, length(yvec)))

    if multiple_inputs
        if vaugmented
            identity_like = AD.identity_matrix_like(v)
            vaug = map(identity_like) do identity_like_i
                identity_like_i .* v
            end

            pf1 = map(v -> AD.pushforward_function(backend, fjac, xvec, yvec)(v), vaug)
            ((valvec1, pf2x), (valvec2, pf2y)) = map(
                v -> AD.value_and_pushforward_function(backend, fjac, xvec, yvec)(v), vaug
            )
        else
            pf1 = AD.pushforward_function(backend, fjac, xvec, yvec)(v)
            valvec, pf2 = AD.value_and_pushforward_function(backend, fjac, xvec, yvec)(v)
            ((valvec1, pf2x), (valvec2, pf2y)) = (valvec, pf2[1]), (valvec, pf2[2])
        end

        if test_types
            @test valvec1 isa Vector{Float64}
            @test valvec2 isa Vector{Float64}
            @test pf1[1] isa Vector{Float64}
            @test pf1[2] isa Vector{Float64}
            @test pf2x isa Vector{Float64}
            @test pf2y isa Vector{Float64}
        end
        @test valvec1 == fjac(xvec, yvec)
        @test valvec2 == fjac(xvec, yvec)
        @test norm.((pf2x, pf2y) .- pf1) == (0, 0)
        # test with respect to analytical solution
        @test minimum(
            isapprox.(pf1, (jxvp(xvec, yvec, v[1]), jyvp(xvec, yvec, v[2])), atol=1e-10)
        )
        @test xvec == xvec2
        @test yvec == yvec2
    end

    valvec1, pf1 = AD.value_and_pushforward_function(backend, x -> fjac(x, yvec), xvec)(
        v[1]
    )
    valvec2, pf2 = AD.value_and_pushforward_function(backend, y -> fjac(xvec, y), yvec)(
        v[2]
    )

    if test_types
        @test valvec1 isa Vector{Float64}
        @test valvec2 isa Vector{Float64}
        @test pf1[1] isa Vector{Float64}
        @test pf2[1] isa Vector{Float64}
    end
    @test valvec1 == fjac(xvec, yvec)
    @test valvec2 == fjac(xvec, yvec)
    @test minimum(
        isapprox.(
            (pf1[1], pf2[1]), (jxvp(xvec, yvec, v[1]), jyvp(xvec, yvec, v[2])), atol=1e-10
        ),
    )
end

function test_j′vp(backend; multiple_inputs=true, rng=Random.GLOBAL_RNG, test_types=true)
    # test with respect to analytical solution
    w = rand(rng, length(fjac(xvec, yvec)))
    if multiple_inputs
        pb1 = AD.pullback_function(backend, fjac, xvec, yvec)(w)
        valvec, pbf2 = AD.value_and_pullback_function(backend, fjac, xvec, yvec)
        pb2 = pbf2(w)

        if test_types
            @test valvec isa Vector{Float64}
            @test pb1[1] isa AbstractVector{Float64}
            @test pb1[2] isa AbstractVector{Float64}
            @test pb2[1] isa AbstractVector{Float64}
            @test pb2[2] isa AbstractVector{Float64}
        end
        @test valvec == fjac(xvec, yvec)
        @test norm.(pb2 .- pb1) == (0, 0)
        @test minimum(
            isapprox.(pb1, (vJxp(xvec, yvec, w), vJyp(xvec, yvec, w)), atol=1e-10)
        )
        @test xvec == xvec2
        @test yvec == yvec2
    end

    valvec1, pbf1 = AD.value_and_pullback_function(backend, x -> fjac(x, yvec), xvec)
    pb1 = pbf1(w)
    valvec2, pbf2 = AD.value_and_pullback_function(backend, y -> fjac(xvec, y), yvec)
    pb2 = pbf2(w)
    if test_types
        @test valvec1 isa Vector{Float64}
        @test valvec2 isa Vector{Float64}
        @test pb1[1] isa AbstractVector{Float64}
        @test pb2[1] isa AbstractVector{Float64}
    end
    @test valvec1 == fjac(xvec, yvec)
    @test valvec2 == fjac(xvec, yvec)
    @test minimum(
        isapprox.((pb1[1], pb2[1]), (vJxp(xvec, yvec, w), vJyp(xvec, yvec, w)), atol=1e-10)
    )
end

function test_lazy_derivatives(backend; multiple_inputs=true)
    # single input function
    der1 = AD.derivative(backend, x -> fder(x, yscalar), xscalar)
    lazyder = AD.LazyDerivative(backend, x -> fder(x, yscalar), xscalar)

    # multiplication with scalar
    @test lazyder * yscalar == der1 .* yscalar
    @test lazyder * yscalar isa Tuple

    @test yscalar * lazyder == yscalar .* der1
    @test yscalar * lazyder isa Tuple

    # multiplication with array
    @test lazyder * yvec == (der1 .* yvec,)
    @test lazyder * yvec isa Tuple

    @test yvec * lazyder == (yvec .* der1,)
    @test yvec * lazyder isa Tuple

    # multiplication with tuple
    @test lazyder * (yscalar,) == lazyder * yscalar
    @test lazyder * (yvec,) == lazyder * yvec

    @test (yscalar,) * lazyder == yscalar * lazyder
    @test (yvec,) * lazyder == yvec * lazyder

    # two input function
    if multiple_inputs
        der1 = AD.derivative(backend, fder, xscalar, yscalar)
        lazyder = AD.LazyDerivative(backend, fder, (xscalar, yscalar))

        # multiplication with scalar
        @test lazyder * yscalar == der1 .* yscalar
        @test lazyder * yscalar isa Tuple

        @test yscalar * lazyder == yscalar .* der1
        @test yscalar * lazyder isa Tuple

        # multiplication with array
        @test lazyder * yvec == (der1[1] * yvec, der1[2] * yvec)
        @test lazyder * yvec isa Tuple

        @test yvec * lazyder == (yvec * der1[1], yvec * der1[2])
        @test lazyder * yvec isa Tuple

        # multiplication with tuple
        @test_throws AssertionError lazyder * (yscalar,)
        @test_throws AssertionError lazyder * (yvec,)

        @test_throws AssertionError (yscalar,) * lazyder
        @test_throws AssertionError (yvec,) * lazyder
    end
end

function test_lazy_gradients(backend; multiple_inputs=true)
    # single input function
    grad1 = AD.gradient(backend, x -> fgrad(x, yvec), xvec)
    lazygrad = AD.LazyGradient(backend, x -> fgrad(x, yvec), xvec)

    # multiplication with scalar
    @test norm.(lazygrad * yscalar .- grad1 .* yscalar) == (0,)
    @test lazygrad * yscalar isa Tuple

    @test norm.(yscalar * lazygrad .- yscalar .* grad1) == (0,)
    @test yscalar * lazygrad isa Tuple

    # multiplication with tuple
    @test lazygrad * (yscalar,) == lazygrad * yscalar
    @test (yscalar,) * lazygrad == yscalar * lazygrad

    # two input function
    if multiple_inputs
        grad1 = AD.gradient(backend, fgrad, xvec, yvec)
        lazygrad = AD.LazyGradient(backend, fgrad, (xvec, yvec))

        # multiplication with scalar
        @test norm.(lazygrad * yscalar .- grad1 .* yscalar) == (0, 0)
        @test lazygrad * yscalar isa Tuple

        @test norm.(yscalar * lazygrad .- yscalar .* grad1) == (0, 0)
        @test yscalar * lazygrad isa Tuple

        # multiplication with tuple
        @test_throws AssertionError lazygrad * (yscalar,) == lazygrad * yscalar
        @test_throws AssertionError (yscalar,) * lazygrad == yscalar * lazygrad
    end
end

function test_lazy_jacobians(
    backend; multiple_inputs=true, vaugmented=false, rng=Random.GLOBAL_RNG
)
    # single input function
    jac1 = AD.jacobian(backend, x -> fjac(x, yvec), xvec)
    lazyjac = AD.LazyJacobian(backend, x -> fjac(x, yvec), xvec)

    # multiplication with scalar
    @test norm.(lazyjac * yscalar .- jac1 .* yscalar) == (0,)
    @test lazyjac * yscalar isa Tuple

    @test norm.(yscalar * lazyjac .- yscalar .* jac1) == (0,)
    @test yscalar * lazyjac isa Tuple

    w = rand(rng, length(fjac(xvec, yvec)))
    v = (rand(rng, length(xvec)), rand(rng, length(xvec)))

    # vjp
    pb1 = (vJxp(xvec, yvec, w),)
    res = w' * lazyjac
    @test minimum(isapprox.(pb1, res, atol=1e-10))
    @test res isa Tuple

    # jvp
    pf1 = (jxvp(xvec, yvec, v[1]),)
    res = lazyjac * v[1]
    @test minimum(isapprox.(pf1, res, atol=1e-10))
    @test res isa Tuple

    # two input function
    if multiple_inputs
        jac1 = AD.jacobian(backend, fjac, xvec, yvec)
        lazyjac = AD.LazyJacobian(backend, fjac, (xvec, yvec))

        # multiplication with scalar
        @test norm.(lazyjac * yscalar .- jac1 .* yscalar) == (0, 0)
        @test lazyjac * yscalar isa Tuple

        @test norm.(yscalar * lazyjac .- yscalar .* jac1) == (0, 0)
        @test yscalar * lazyjac isa Tuple

        # vjp
        pb1 = (vJxp(xvec, yvec, w), vJyp(xvec, yvec, w))
        res = w'lazyjac
        @test minimum(isapprox.(pb1, res, atol=1e-10))
        @test res isa Tuple

        # jvp
        pf1 = (jxvp(xvec, yvec, v[1]), jyvp(xvec, yvec, v[2]))

        if vaugmented
            identity_like = AD.identity_matrix_like(v)
            vaug = map(identity_like) do identity_like_i
                identity_like_i .* v
            end

            res = map(v -> (lazyjac * v)[1], vaug)
        else
            res = lazyjac * v
        end
        @test minimum(isapprox.(pf1, res, atol=1e-10))
        @test res isa Tuple
    end
end

function test_lazy_hessians(backend; multiple_inputs=true, rng=Random.GLOBAL_RNG)
    # fdm_backend not used here yet..
    # single input function
    fhess = x -> fgrad(x, yvec)
    hess1 = (dfgraddxdx(xvec, yvec),)
    lazyhess = AD.LazyHessian(backend, fhess, xvec)

    # multiplication with scalar
    @test minimum(isapprox.(lazyhess * yscalar, hess1 .* yscalar, atol=1e-10))
    @test lazyhess * yscalar isa Tuple

    # multiplication with scalar
    @test minimum(isapprox.(yscalar * lazyhess, yscalar .* hess1, atol=1e-10))
    @test yscalar * lazyhess isa Tuple

    w = rand(rng, length(xvec))
    v = rand(rng, length(xvec))

    # Hvp
    Hv = map(h -> h * v, hess1)
    res = lazyhess * v
    @test minimum(isapprox.(Hv, res, atol=1e-10))
    @test res isa Tuple

    # H′vp
    wH = map(h -> h' * w, hess1)
    res = w' * lazyhess
    @test minimum(isapprox.(wH, res, atol=1e-10))
    @test res isa Tuple
end
