# This file is a part of ChangesOfVariables.jl, licensed under the MIT License (MIT).

using ChangesOfVariables
using Test

using LinearAlgebra

using ChangesOfVariables
using ChangesOfVariables: test_with_logabsdet_jacobian

include("getjacobian.jl")


@testset "with_logabsdet_jacobian" begin
    _bc_func(f) = Base.Broadcast.BroadcastFunction(f)

    @test with_logabsdet_jacobian(sum, rand(5)) === NoLogAbsDetJacobian(sum, rand(5))
    @test with_logabsdet_jacobian(log ∘ sum, 5.0f0) === NoLogAbsDetJacobian(log ∘ sum, 5.0f0)
    @test_throws MethodError _, _ = with_logabsdet_jacobian(sum, rand(5))

    @test with_logabsdet_jacobian(sin, 4.9) === NoLogAbsDetJacobian{typeof(sin), Float64}()
    @test with_logabsdet_jacobian(String, 4.9) === NoLogAbsDetJacobian{Type{String}, Float64}()
    @test with_logabsdet_jacobian(String, Float64) === NoLogAbsDetJacobian{Type{String}, Type{Float64}}()
    @test with_logabsdet_jacobian(sin, Float64) === NoLogAbsDetJacobian{typeof(sin), Type{Float64}}()

    @test with_logabsdet_jacobian(sin ∘ log, 4.9) === NoLogAbsDetJacobian{typeof(sin ∘ log), Float64}()
    @test with_logabsdet_jacobian(log ∘ sin, 4.9) === NoLogAbsDetJacobian{typeof(log ∘ sin), Float64}()

    @test with_logabsdet_jacobian(Base.Fix1(broadcast, sin), 4.9) === NoLogAbsDetJacobian(Base.Fix1(broadcast, sin), 4.9)

    function ChangesOfVariables.with_logabsdet_jacobian(::typeof(foo), x)
        y = foo(x)
        ladj = -x + 2 * log(y)
        (y, ladj)
    end

    x = 4.2
    X = rand(10)
    A = rand(5, 5)
    CA = Complex.(rand(5, 5), rand(5, 5))

    isaprx(a, b; kwargs...) = isapprox(a,b; kwargs...)
    isaprx(a::NTuple{N,Any}, b::NTuple{N,Any}; kwargs...) where N = all(map((a,b) -> isaprx(a, b; kwargs...), a, b))


    test_with_logabsdet_jacobian(foo, x, getjacobian)

    test_with_logabsdet_jacobian(log ∘ foo, x, getjacobian)

    @testset "with_logabsdet_jacobian on mapped and broadcasted" begin
        for f in (_bc_func(foo), Base.Fix1(map, foo), Base.Fix1(broadcast, foo))
            for arg in (x, fill(x,), Ref(x), (x,), (x, 2*x, 3*x), X)
                test_with_logabsdet_jacobian(f, arg, getjacobian, compare = isaprx)
            end
        end
    end

    @testset "with_logabsdet_jacobian on mapped and broadcasted without ladj" begin
        for f in (_bc_func(sin), Base.Fix1(map, sin), Base.Fix1(broadcast, sin))
            for arg in (x, (x,), (x, x), X)
                @test with_logabsdet_jacobian(f, arg) === NoLogAbsDetJacobian(f, arg)
            end
        end
    end

    @testset "with_logabsdet_jacobian on mapped and broadcasted with mixed ladj" begin
        @test with_logabsdet_jacobian(_bc_func(bar), (1.0, 2.0)) == ((1.0, 2.0), 0.0)
        @test with_logabsdet_jacobian(_bc_func(bar), [1.0, 2.0]) == ([1.0, 2.0], 0.0)
        @test with_logabsdet_jacobian(_bc_func(bar), (1.0, 2)) === NoLogAbsDetJacobian(_bc_func(bar), (1.0, 2))
        @test with_logabsdet_jacobian(Base.Fix1(map, bar), (1, 2.0)) === NoLogAbsDetJacobian(Base.Fix1(map, bar), (1, 2.0))
        @test with_logabsdet_jacobian(_bc_func(bar), Real[1.0, 2]) === NoLogAbsDetJacobian(_bc_func(bar), Real[1.0, 2])
        @test with_logabsdet_jacobian(Base.Fix1(map, bar), Real[1, 2.0]) === NoLogAbsDetJacobian(Base.Fix1(map, bar), Real[1, 2.0])
        @test with_logabsdet_jacobian(Base.Fix1(broadcast, bar), Real[1, 2]) === NoLogAbsDetJacobian(Base.Fix1(broadcast, bar), Real[1, 2])
    end

    @testset "with_logabsdet_jacobian on identity, adjoint and transpose" begin
        for f in (identity, adjoint, transpose)
            for arg in (x, A)
                test_with_logabsdet_jacobian(f, arg, getjacobian)
            end
        end
    end

    @testset "with_logabsdet_jacobian on inv" begin
        for arg in (x, A, CA)
            test_with_logabsdet_jacobian(inv, arg, getjacobian)
        end
    end

    @testset "with_logabsdet_jacobian on log and exp functions" begin
        for f in (+, -, exp, log, exp2, log2, exp10, log10, expm1, log1p)
            test_with_logabsdet_jacobian(f, x, getjacobian)
        end
    end
end
