using Adapt
using Test


# custom array type

struct CustomArray{T,N} <: AbstractArray{T,N}
    arr::Array{T,N}
end

Adapt.adapt_storage(::Type{<:CustomArray}, xs::Array) = CustomArray(xs)

Base.size(x::CustomArray, y...) = size(x.arr, y...)
Base.getindex(x::CustomArray, y...) = getindex(x.arr, y...)
Base.count(x::CustomArray) = count(x.arr)

const mat = CustomArray{Float64,2}(rand(2,2))
const vec = CustomArray{Float64,1}(rand(2))

const mat_bools = CustomArray{Bool,2}(rand(Bool,2,2))

# test that adaption of `src_expr` to `to` matches `dst_expr`:
# - using ==
# - comparing the types
#
# if `typ` is set, assert that `src` (or `dst`, as they have been asserted to be identical)
# is a subtype of `typ`. this is useful to check that complex unions cover what's needed.
macro test_adapt(to, src_expr, dst_expr, typ=nothing)
    quote
        src = $(esc(src_expr))
        dst = $(esc(dst_expr))

        res = @inferred(adapt($(esc(to)), src))
        @test res == dst
        @test typeof(res) == typeof(dst)

        if $(esc(typ)) !== nothing
            @test typeof(dst) <: $(esc(typ))
        end
    end
end

AnyCustomArray{T,N} = Union{CustomArray,WrappedArray{T,N,CustomArray,CustomArray{T,N}}}


# basic adaption
@test_adapt CustomArray mat.arr mat

# idempotency
@test_adapt CustomArray mat mat

# to array
@test_adapt Array mat mat.arr
@test_adapt Array{Float32} mat Float32.(mat.arr)

# custom wrapper
struct Wrapper{T}
    arr::T
end
Adapt.adapt_structure(to, xs::Wrapper) = Wrapper(adapt(to, xs.arr))
@test_adapt CustomArray Wrapper(mat.arr) Wrapper(mat)


@testset "base structures" begin

@test @inferred(adapt(nothing, NamedTuple())) == NamedTuple()
@test_adapt CustomArray (mat.arr,) (mat,)
@test @allocated(adapt(nothing, ())) == 0
@test @allocated(adapt(nothing, (1,))) == 0
@test @allocated(adapt(nothing, (1,2,3,4,5,6,7,8,9,10))) == 0

@test_adapt CustomArray (a=mat.arr,) (a=mat,)

end

# NOTE: if we put this in the preceding testset, unrelated tests start to allocate
@testset "closures" begin

# basic capture of a variable
function closure1(x)
    function foobar()
        x
    end
    foobar
end

f = closure1(mat.arr)
@test f() === mat.arr
f′ = adapt(CustomArray, f)
@test f′() === mat
@test @inferred(adapt(nothing, f)()) === f()

# closure with sparams
function closure2(A::AbstractArray{T}=zeros(1), b::Number=0) where {T}
    function f()
        convert(T, 0)
        A, A[1] == b
    end
    f
end

f = closure2(mat.arr)
@test f() === (mat.arr, false)
f′ = adapt(CustomArray, f)
@test f′() === (mat, false)

# closure with box
x1 = CustomArray{Int,1}([1])
x2 = CustomArray{Int,1}([2])
x3 = CustomArray{Int,1}([3])
function closure3(a)
    b = "whatever"
    f(c) = (a, b, c)
    b = x2.arr
    return f
end

f = closure3(x1.arr)
@test f(x3.arr) === (x1.arr, x2.arr, x3.arr)
f′ = adapt(CustomArray, f)
@test f′(x3.arr) === (x1, x2, x3.arr)
# NOTE: actual arguments should be adapted by the caller

end


@testset "array wrappers" begin

@test_adapt CustomArray view(mat.arr,:,:) view(mat,:,:) AnyCustomArray
inds = CustomArray{Int,1}([1,2])
@test_adapt CustomArray view(mat.arr,inds.arr,:) view(mat,inds,:) AnyCustomArray

# NOTE: manual creation of PermutedDimsArray because permutedims collects
@test_adapt CustomArray PermutedDimsArray(mat.arr,(2,1)) PermutedDimsArray(mat,(2,1)) AnyCustomArray

# NOTE: manual creation of ReshapedArray because Base.Array has an optimized `reshape`
@test_adapt CustomArray Base.ReshapedArray(mat.arr,(2,2),()) reshape(mat,(2,2)) AnyCustomArray

@test_adapt CustomArray Base.LogicalIndex(mat_bools.arr) Base.LogicalIndex(mat_bools) AnyCustomArray

@test_adapt CustomArray reinterpret(Int64,mat.arr) reinterpret(Int64,mat) AnyCustomArray

@test_adapt CustomArray reinterpret(reshape,Int64,mat.arr) reinterpret(reshape,Int64,mat) AnyCustomArray


## doubly-wrapped

@test_adapt CustomArray reinterpret(Int64,view(mat.arr,:,:)) reinterpret(Int64,view(mat,:,:)) AnyCustomArray

@test_adapt CustomArray reshape(view(mat.arr,:,:), (2,2)) reshape(view(mat,:,:), (2,2)) AnyCustomArray
@test_adapt CustomArray reshape(reinterpret(Int64,mat.arr), (2,2)) reshape(reinterpret(Int64,mat), (2,2)) AnyCustomArray
@test_adapt CustomArray reshape(reinterpret(Int64,view(mat.arr,:,:)), (2,2)) reshape(reinterpret(Int64,view(mat,:,:)), (2,2)) AnyCustomArray

@test_adapt CustomArray view(reinterpret(Int64,mat.arr), :, :) view(reinterpret(Int64,mat), :, :) AnyCustomArray
@test_adapt CustomArray view(reinterpret(Int64,view(mat.arr,:,:)), :, :) view(reinterpret(Int64,view(mat,:,:)), :, :) AnyCustomArray
@test_adapt CustomArray view(Base.ReshapedArray(mat.arr,(2,2),()), :, :) view(reshape(mat, (2,2)), :, :) AnyCustomArray
@test_adapt CustomArray view(reshape(view(mat.arr,:,:), (2,2)), :, :) view(reshape(view(mat,:,:), (2,2)), :, :) AnyCustomArray
@test_adapt CustomArray view(reshape(reinterpret(Int64,mat.arr), (2,2)), :, :) view(reshape(reinterpret(Int64,mat), (2,2)), :, :) AnyCustomArray
@test_adapt CustomArray view(reshape(reinterpret(Int64,view(mat.arr,:,:)), (2,2)), :, :) view(reshape(reinterpret(Int64,view(mat,:,:)), (2,2)), :, :) AnyCustomArray

@test_adapt CustomArray reinterpret(reshape,Int64,view(mat.arr,:,:)) reinterpret(reshape,Int64,view(mat,:,:)) AnyCustomArray
@test_adapt CustomArray view(reinterpret(reshape,Int64,mat.arr), :, :) view(reinterpret(reshape,Int64,mat), :, :) AnyCustomArray
@test_adapt CustomArray view(reinterpret(reshape,Int64,view(mat.arr,:,:)), :, :) view(reinterpret(reshape,Int64,view(mat,:,:)), :, :) AnyCustomArray


using LinearAlgebra

@test_adapt CustomArray mat.arr' mat' AnyCustomArray

@test_adapt CustomArray transpose(mat.arr) transpose(mat) AnyCustomArray

@test_adapt CustomArray LowerTriangular(mat.arr) LowerTriangular(mat) AnyCustomArray
@test_adapt CustomArray UnitLowerTriangular(mat.arr) UnitLowerTriangular(mat) AnyCustomArray
@test_adapt CustomArray UpperTriangular(mat.arr) UpperTriangular(mat) AnyCustomArray
@test_adapt CustomArray UnitUpperTriangular(mat.arr) UnitUpperTriangular(mat) AnyCustomArray
@test_adapt CustomArray Symmetric(mat.arr) Symmetric(mat) AnyCustomArray
@test_adapt CustomArray Hermitian(mat.arr) Hermitian(mat) AnyCustomArray
@test_adapt CustomArray Symmetric(mat.arr, :L) Symmetric(mat, :L) AnyCustomArray
@test_adapt CustomArray Hermitian(mat.arr, :L) Hermitian(mat, :L) AnyCustomArray
    
@test_adapt CustomArray Diagonal(vec.arr) Diagonal(vec) AnyCustomArray

dl = CustomArray{Float64,1}(rand(2))
du = CustomArray{Float64,1}(rand(2))
d = CustomArray{Float64,1}(rand(3))
@test_adapt CustomArray Tridiagonal(dl.arr, d.arr, du.arr) Tridiagonal(dl, d, du) AnyCustomArray
@test_adapt CustomArray Bidiagonal(d.arr, du.arr, 'U') Bidiagonal(d, du, 'U') AnyCustomArray
end


@testset "type information" begin
    # single wrapping
    @test parent_type(Transpose{Int,Array{Int,1}}) == Array{Int,1}
    @test parent_type(Transpose{Int,Transpose{Int,Array{Int,1}}}) == Transpose{Int,Array{Int,1}}

    # double wrapping
    @test unwrap_type(Transpose{Int,Array{Int,1}}) == Array{Int,1}
    @test unwrap_type(Transpose{Int,Transpose{Int,Array{Int,1}}}) == Array{Int,1}
end


struct MyStruct{A,B}
    a::A
    b::B
end

@testset "@adapt_structure" begin
    Adapt.@adapt_structure MyStruct

    u = ones(3)
    v = zeros(5)

    @test_adapt CustomArray MyStruct(u,v) MyStruct(CustomArray(u), CustomArray(v))
    @test_adapt CustomArray MyStruct(u,1.0) MyStruct(CustomArray(u), 1.0)

    # Bug 102; type defined in the same block as @adapt_structure
    struct LocalStruct
        x::Int
    end
    Adapt.@adapt_structure LocalStruct
    @test adapt(CustomArray, LocalStruct(1)) === LocalStruct(1)

    # The macro must work in modules that only import @adapt_structure,
    # not the Adapt module itself (e.g. `using Adapt: @adapt_structure`).
    @eval module ModuleOnlyImportingMacro
        using Adapt: @adapt_structure
        struct S
            x::Int
        end
        @adapt_structure S
    end
    @test adapt(CustomArray, ModuleOnlyImportingMacro.S(1)) === ModuleOnlyImportingMacro.S(1)
end


@testset "Broadcast" begin
    @test_adapt CustomArray Base.broadcasted(identity, mat.arr) Base.broadcasted(identity, mat)

    f = (x)->((y)->(y,x))
    bc = Base.broadcasted(f(mat.arr), (mat.arr,))
    @test typeof(copy(adapt(CustomArray, bc))) == typeof(broadcast(f(mat), (mat,)))
end

@testset "SparseArrays" begin
    using SparseArrays
    m = sparse([1, 2], [2, 1], [1, 2])
    @test_adapt Array m m
    @test_adapt Array{Float64} m SparseMatrixCSC{Float64}(m)

    v = sparsevec([1, 3], [1, 2])
    @test_adapt Array v v
    @test_adapt Array{Float64} v SparseVector{Float64}(v)
end

@testset "StaticArrays" begin
    using StaticArrays
    @test_adapt SArray{Tuple{3}} [1,2,3] SArray{Tuple{3}}([1,2,3])

    # can't possibly infer this one, so not using @test_adapt
    #@test_adapt SArray           [1,2,3] SArray{Tuple{3}}([1,2,3])
    @test adapt(SArray, [1,2,3]) === SArray{Tuple{3}}([1,2,3])

    # adapt_structure recursively adapts SArray elements
    struct F64toF32 end
    Adapt.adapt_storage(::F64toF32, x::Float64) = Float32(x)
    @test adapt(F64toF32(), SVector{2}(1.0, 2.0)) === SVector{2}(1.0f0, 2.0f0)
    @test adapt(F64toF32(), SVector{2}(1, 2)) === SVector{2}(1, 2)  # no-op for Int
end

@testset "Ranges" begin
    # normally these fall back to `convert(Array, r)`, so we only need to test
    # that the type matches

    @test adapt(Array, 1:10) === 1:10
    @test adapt(Array, Base.OneTo(10)) === Base.OneTo(10)
    @test adapt(Array, 1:2:10) === 1:2:10
    @test adapt(Array, 1.:2.:10.) === 1.:2.:10.
    @test adapt(Array, Base.Slice(1:10)) === Base.Slice(1:10)
    @test adapt(Array, LinRange(1,2,10)) === LinRange(1,2,10)
end

@testset "CartesianIndices" begin
    @test adapt(Array, CartesianIndices((1,2,3))) === CartesianIndices((1,2,3))
end
