
using Manifolds, ManifoldDiff, Test, ManifoldsBase

using ManifoldDiff:
    adjoint_differential_exp_argument,
    adjoint_differential_exp_argument!,
    adjoint_differential_exp_basepoint,
    adjoint_differential_exp_basepoint!,
    adjoint_differential_log_argument,
    adjoint_differential_log_argument!,
    adjoint_differential_log_basepoint,
    adjoint_differential_log_basepoint!,
    adjoint_differential_shortest_geodesic_endpoint,
    adjoint_differential_shortest_geodesic_endpoint!,
    adjoint_differential_shortest_geodesic_startpoint,
    adjoint_differential_shortest_geodesic_startpoint!


@testset "Differentials (on Sphere(2))" begin
    # The Adjoint Differentials test using the same variables as the differentials
    # test
    p = [1.0, 0.0, 0.0]
    q = [0.0, 1.0, 0.0]
    M = Sphere(2)
    X = log(M, p, q)
    # Tept differentials (1) Dp of Log_pq
    Y = similar(X)
    @test adjoint_differential_log_basepoint(M, p, p, X) == -X
    adjoint_differential_log_basepoint!(M, Y, p, p, X)
    @test Y == -X
    @test adjoint_differential_log_argument(M, p, p, X) == X
    adjoint_differential_log_argument!(M, Y, p, p, X)
    @test Y == X
    @test adjoint_differential_exp_basepoint(M, p, zero_vector(M, p), X) == X
    adjoint_differential_exp_basepoint!(M, Y, p, zero_vector(M, p), X)
    @test Y == X
    @test adjoint_differential_exp_argument(M, p, zero_vector(M, p), X) == X
    adjoint_differential_exp_argument!(M, Y, p, zero_vector(M, p), X)
    @test Y == X
    for t in [0, 0.15, 0.33, 0.66, 0.9]
        @test adjoint_differential_shortest_geodesic_startpoint(M, p, p, t, X) ==
              (1 - t) * X
        adjoint_differential_shortest_geodesic_startpoint!(M, Y, p, p, t, X)
        @test Y == (1 - t) * X
        @test norm(
            M,
            p,
            adjoint_differential_shortest_geodesic_endpoint(M, p, p, t, X) - t * X,
        ) ≈ 0 atol = 10.0^(-16)
        adjoint_differential_shortest_geodesic_endpoint!(M, Y, p, p, t, X)
        @test norm(M, p, Y - t * X) ≈ 0 atol = 10.0^(-16)
    end
    Mp = PowerManifold(M, NestedPowerRepresentation(), 3)
    pP = [p, q, p]
    qP = [p, p, q]
    XP = [X, zero_vector(M, p), -X]
    YP = similar.(XP)
    ZP = [[0.0, π / 2, 0.0], [0.0, 0.0, 0.0], [π / 2, 0.0, 0.0]]
    @test adjoint_differential_log_argument(Mp, pP, qP, XP) == ZP
    adjoint_differential_log_argument!(Mp, YP, pP, qP, XP)
    @test ZP == YP

    Mprod = ProductManifold(M, M, M)
    pP = ArrayPartition(p, q, p)
    qP = ArrayPartition(p, p, q)
    XP = ArrayPartition(X, zero_vector(M, p), -X)
    YP = allocate(XP)
    ZP = ArrayPartition([0.0, π / 2, 0.0], [0.0, 0.0, 0.0], [π / 2, 0.0, 0.0])
    @test adjoint_differential_log_argument(Mprod, pP, qP, XP) == ZP
    adjoint_differential_log_argument!(Mprod, YP, pP, qP, XP)
    @test ZP == YP

end

@testset "Adjoint differentials on the Circle" begin
    M = Circle()
    p = 0
    q = π / 4
    X = π / 8
    @test adjoint_differential_log_argument(M, p, q, X) == X
end
@testset "Adjoint differentials on Euclidean()" begin
    M = Euclidean()
    p = 0
    q = π / 4
    X = π / 8
    @test adjoint_differential_log_argument(M, p, q, X) == X
end
