using Flux, CuArrays, Test
using Flux: gpu
using Zygote

@info "Testing GPU Support"

@testset "CuArrays" begin

x = param(randn(5, 5))
cx = gpu(x)
@test cx isa CuArray

@test Flux.onecold(param(gpu([1.,2.,3.]))) == 3

x = Flux.onehotbatch([1, 2, 3], 1:3)
cx = gpu(x)
@test cx isa Flux.OneHotMatrix && cx isa CuArray
@test (cx .+ 1) isa CuArray

m = Chain(Dense(10, 5, tanh), Dense(5, 2), softmax)
cm = gpu(m)

@test all(p isa CuArray for p in params(cm))
@test cm(gpu(rand(10, 10))) isa CuArray{Float32,2}

x = [1,2,3]
cx = gpu(x)
@test Flux.crossentropy(x,x) ≈ Flux.crossentropy(cx,cx)

xs = param(rand(5,5))
ys = Flux.onehotbatch(1:5,1:5)
@test collect(cu(xs) .+ cu(ys)) ≈ collect(xs .+ ys)

c = gpu(Conv((2,2),3=>4))
l = c(gpu(rand(10,10,3,2)))
fwd, back = Zygote.forward(sum, l)
back(one(Float64))

c = gpu(CrossCor((2,2),3=>4))
l = c(gpu(rand(10,10,3,2)))
fwd, back = Zygote.forward(sum, l)
back(one(Float64))

end

@testset "onecold gpu" begin
  y = Flux.onehotbatch(ones(3), 1:10) |> gpu;
  @test Flux.onecold(y) isa CuArray
  @test y[3,:] isa CuArray
end

if CuArrays.libcudnn != nothing
    @info "Testing Flux/CUDNN"
    include("cudnn.jl")
    if !haskey(ENV, "CI_DISABLE_CURNN_TEST")
      include("curnn.jl")
    end
end