using CUDA
using Enzyme
using LinearAlgebra: mul!, dot, Symmetric
using Test

@testset "CUDA memory copies" begin
    #= Exercise the reverse copy rules across CUDA's memory types. A host<->device
    roundtrip gradient must recover `2x`. =#
    grad_roundtrip = function (to_gpu)
        x = Float32[1, 2, 3]
        dx = zeros(Float32, 3)
        Enzyme.autodiff(
            Reverse,
            x -> sum(abs2, Array(to_gpu(x))),
            Active,
            Duplicated(x, dx),
        )
        return dx
    end

    @testset "device memory" begin
        @test grad_roundtrip(x -> cu(x)) == Float32[2, 4, 6]
    end
    # Unified/host memory: `pointer(::CuArray{…,Unified/Host})` is inferred as `Union{CuPtr,Ptr}`, so the `pointer` rule cannot yet return a concrete
    @testset "unified memory" begin
        @test grad_roundtrip(x -> cu(x; unified = true)) == Float32[2, 4, 6]
    end
    @testset "host memory" begin
        @test grad_roundtrip(x -> cu(x; host = true)) == Float32[2, 4, 6]
    end

    x = CuArray(Float32[1, 2, 3])
    original = copy(x)
    dx = CUDA.zeros(Float32, 3)
    Enzyme.autodiff(
        Reverse,
        x -> sum(abs2, Array(x)),
        Active,
        Duplicated(x, dx),
    )
    @test Array(x) == Array(original)
    @test Array(dx) == 2 .* Array(x)

    fill!(dx, 0)
    Enzyme.autodiff(
        Reverse,
        x -> sum(abs2, Array(copy(x))),
        Active,
        Duplicated(x, dx),
    )
    @test Array(x) == Array(original)
    @test Array(dx) == 2 .* Array(x)
end

function mul_kernel(A)
    i = threadIdx().x
    if i <= length(A)
        A[i] *= A[i]
    end
    return nothing
end

function grad_mul_kernel(A, dA)
    autodiff_deferred(Reverse, Const(mul_kernel), Const, Duplicated(A, dA))
    return nothing
end

@testset "mul_kernel" begin
    A = CUDA.ones(64,)
    @cuda threads=length(A) mul_kernel(A)
    A = CUDA.ones(64,)
    dA = similar(A)
    dA .= 1
    @cuda threads=length(A) grad_mul_kernel(A, dA)
    @test all(dA .== 2)
end

function exp_kernel(A)
    i = threadIdx().x
    if i <= length(A)
        A[i] = exp(A[i])
    end
    return nothing
end

function grad_exp_kernel(A, dA)
    autodiff_deferred(Reverse, Const(exp_kernel), Const, Duplicated(A, dA))
    return nothing
end

@testset "exp_kernel" begin
    A = CUDA.ones(64,)
    @cuda threads=length(A) exp_kernel(A)
    A = CUDA.ones(64,)
    dA = similar(A)
    dA .= 1
    @cuda threads=length(A) grad_exp_kernel(A, dA)
    @test all(dA .== exp(1.f0))
end

function cos_kernel(A)
    i = threadIdx().x
    if i <= length(A)
        A[i] = cos(A[i])
    end
    return nothing
end

function grad_cos_kernel(A, dA)
    autodiff_deferred(Reverse, Const(cos_kernel), Const, Duplicated(A, dA))
    return nothing
end

@testset "cos_kernel" begin
    A = CUDA.ones(64,)
    @cuda threads=length(A) cos_kernel(A)
    A = CUDA.ones(64,)
    dA = similar(A)
    dA .= 1
    @cuda threads=length(A) grad_cos_kernel(A, dA)
    @test all(dA .≈ -sin(1.f0))
end

function val_kernel!(_, ::Val{N}) where N
    return nothing
end

function dval_kernel!(du, ::Val{N}) where {N}
    autodiff_deferred(Reverse, Const(val_kernel!), Const, du, Const(Val(N)))
    return nothing
end

# Test for https://github.com/EnzymeAD/Enzyme.jl/issues/358
@testset "Test val kernel" begin
    n = 10
    u = CUDA.rand(n)
    dzdu = CUDA.rand(n)
    @cuda threads=4 dval_kernel!(Duplicated(u, dzdu), Val(n))
end

# https://github.com/EnzymeAD/Enzyme.jl/issues/367

relu(x) = ifelse(x<0, zero(x), x)
## Define CUDA kernel for dense layer
function dense!(
  feats_out, feats_in, W, b,
  ::Val{nfeat_out}, ::Val{nfeat_in}, ::Val{ndof}
) where {nfeat_out, nfeat_in, ndof}

  ## Each thread will update features for a single sample
  idof = (blockIdx().x - 1) * blockDim().x + threadIdx().x

  ## Prevent out-of-bounds array access
  (idof > ndof) && return nothing

  ## Compute `feats_out`
  for i = 1:nfeat_out
    for k = 1:nfeat_in
      feats_out[i, idof] += W[i, k] * feats_in[k, idof]
    end

    feats_out[i, idof] = relu(feats_out[i, idof] + b[i])
  end

  return nothing
end # dense!

## Wrapper for Enzyme to differentiate `dense!`
function ddense!(
  dfeats_out, dfeats_in, dW, db,
  ::Val{nfeat_out}, ::Val{nfeat_in}, ::Val{ndof}
) where {nfeat_out, nfeat_in, ndof}

  autodiff_deferred(
    Reverse,
    Const(dense!),
    Const,
    dfeats_out, dfeats_in, dW, db,
    Const(Val(nfeat_out)), Const(Val(nfeat_in)), Const(Val(ndof))
  )
  return nothing

end # ddense!

function call_ddense()
  nthread      = 32
  ndof         = 32
  nblock       = ceil(Int, ndof / nthread)
  nfeat_out    = 32
  nfeat_in     = 16
  feats_out    = zeros(nfeat_out, ndof)    |> cu
  feats_in     = rand(nfeat_in, ndof)      |> cu
  W            = rand(nfeat_out, nfeat_in) |> cu
  b            = rand(nfeat_out)           |> cu
  dzdfeats_out = rand(nfeat_out, ndof)     |> cu
  dfeats_out   = Duplicated(feats_out, dzdfeats_out)
  dzdfeats_in  = zero(feats_in)
  dfeats_in    = Duplicated(feats_in, dzdfeats_in)
  dzdW         = zero(W)
  dW           = Duplicated(W, dzdW)
  dzdb         = zero(b)
  db           = Duplicated(b, dzdb)

  @cuda threads=nthread blocks=nblock ddense!(
    dfeats_out, dfeats_in, dW, db,
    Val(nfeat_out), Val(nfeat_in), Val(ndof)
  )

end # call_ddense

@testset "DDense" begin
    call_ddense()
end

function square_kernel!(x)
    i = threadIdx().x
    x[i] *= x[i]
    sync_threads()
    return nothing
end

# basic squaring on GPU
function square!(x)
    @cuda blocks = 1 threads = length(x) square_kernel!(x)
    return nothing
end

@testset "Reverse Kernel" begin
    A = CUDA.rand(64)
    dA = CUDA.ones(64)
    A .= (1:1:64)
    dA .= 1
    Enzyme.autodiff(Reverse, square!, Duplicated(A, dA))
    @test all(dA .≈ (2:2:128))

    A = CUDA.rand(32)
    dA = CUDA.ones(32)
    dA2 = CUDA.ones(32)
    A .= (1:1:32)
    dA .= 1
    dA2 .= 3
    Enzyme.autodiff(Reverse, square!, BatchDuplicated(A, (dA, dA2)))
    @test all(dA .≈ (2:2:64))
    @test all(dA2 .≈ 3*(2:2:64))
end

#=
The matmul/dot rules on a real backend, where they route through CUBLAS instead of
the GPUArrays fallback kernels. This file also covers the allocating `A * B` and
reductions over the result, neither of which works on JLArray (see
test/ext/jlarrays.jl for why, and for the rest of the coverage).
=#
@testset "GPUArrays linalg rules" begin
    cu(x) = CuArray(x)

    # sum(A·B) seeds dC with ones, so dA = ones·B' and dB = A'·ones.
    @testset "matmul reverse" begin
        A0 = randn(Float32, 3, 4)
        B0 = randn(Float32, 4, 2)
        dA = cu(zero(A0))
        dB = cu(zero(B0))
        Enzyme.autodiff(Reverse, (A, B) -> sum(A * B), Active, Duplicated(cu(A0), dA), Duplicated(cu(B0), dB))
        @test Array(dA) ≈ ones(Float32, 3, 2) * B0'
        @test Array(dB) ≈ A0' * ones(Float32, 3, 2)
    end

    # `X` plain and transposed in one loss, so both cotangents accumulate into the
    # same shadow. Checked against central differences.
    @testset "composed transpose matmul" begin
        X0 = randn(Float32, 6, 3)
        β0 = randn(Float32, 3)
        dX = cu(zero(X0))
        dβ = cu(zero(β0))
        Enzyme.autodiff(Reverse, (X, β) -> sum(transpose(X) * (X * β)), Active, Duplicated(cu(X0), dX), Duplicated(cu(β0), dβ))
        ϵ = 1.0f-3
        fdβ = map(eachindex(β0)) do i
            βp = copy(β0); βp[i] += ϵ
            βm = copy(β0); βm[i] -= ϵ
            (sum(transpose(X0) * (X0 * βp)) - sum(transpose(X0) * (X0 * βm))) / (2ϵ)
        end
        @test Array(dβ) ≈ fdβ rtol = 1.0f-2
    end

    # 3-arg `mul!` on an explicit buffer, the path the JLArray tests use throughout.
    @testset "mul! reverse (preallocated)" begin
        A0 = randn(Float32, 3, 4)
        B0 = randn(Float32, 4, 2)
        dA = cu(zero(A0))
        dB = cu(zero(B0))
        function matmul!(C, A, B)
            mul!(C, A, B)
            return nothing
        end
        C = cu(zeros(Float32, 3, 2))
        dC = cu(ones(Float32, 3, 2))
        Enzyme.autodiff(Reverse, matmul!, Const, Duplicated(C, dC), Duplicated(cu(A0), dA), Duplicated(cu(B0), dB))
        @test Array(dA) ≈ ones(Float32, 3, 2) * B0'
        @test Array(dB) ≈ A0' * ones(Float32, 3, 2)
    end

    @testset "dot reverse" begin
        a0 = randn(Float32, 8)
        b0 = randn(Float32, 8)
        da = cu(zero(a0))
        db = cu(zero(b0))
        Enzyme.autodiff(Reverse, (a, b) -> dot(a, b), Active, Duplicated(cu(a0), da), Duplicated(cu(b0), db))
        @test Array(da) ≈ b0
        @test Array(db) ≈ a0
    end

    # Symmetric/Hermitian operands aren't supported yet; the rule rejects them
    # instead of returning the full matrix's cotangent.
    @testset "Symmetric operand is rejected" begin
        n = 4
        X0 = randn(Float32, n, n)
        dX = cu(zero(X0))
        function matmul!(C, A, B)
            mul!(C, A, B)
            return nothing
        end
        @test_throws ArgumentError Enzyme.autodiff(
            Reverse, matmul!, Const,
            Duplicated(cu(zeros(Float32, n, 3)), cu(ones(Float32, n, 3))),
            Duplicated(Symmetric(cu(X0), :U), Symmetric(dX, :U)), Const(cu(randn(Float32, n, 3))),
        )
    end
end
