text:
using Test, LinearAlgebra, SpecialFunctions
using Statistics: mean
using SparseArrays: nnz
using Random
@everywhere using SparseArrays: sprandn

@testset "test distribute and other constructors" begin
    A = rand(1:100, (100,100))

    @testset "test default distribute" begin
        DA = distribute(A)
        @test length(procs(DA)) == nworkers()
        @test sum(DA) == sum(A)
        close(DA)
    end

    @testset "test distribute with procs arguments" begin
        DA = distribute(A, procs = procs())
        @test length(procs(DA)) == nprocs()
        @test sum(DA) == sum(A)
        close(DA)
    end

    @testset "test distribute with procs and dist arguments" begin
        DA = distribute(A, procs = [1, 2], dist = [1,2])
        @test size(procs(DA)) == (1,2)
        @test sum(DA) == sum(A)
        close(DA)
    end

    @testset "Create darray with unconventional distribution and distibute like it" begin
        block = 10
        Y = nworkers() * block
        X = nworkers() * block
        remote_parts = map(workers()) do wid
            remotecall(rand, wid, block, Y)
        end
        DA1 = DArray(reshape(remote_parts, (length(remote_parts), 1)))
        A = rand(X, Y)
        DA2 = distribute(A, DA1)

        @test size(DA1) == size(DA2)

        close(DA1)
        close(DA2)
    end

    @testset "Global DArray serialization issue #134" begin
        global A134 = drandn(1)
        D2 = DArray(I -> DistributedArrays.localpart(A134), A134)
        @test D2 == A134
        close(A134)
        close(D2)
    end

    @testset "empty_localpart should work when only constructor (not conversion is defined)" begin
        @test DistributedArrays.empty_localpart(Float64,2,LowerTriangular{Float64,Matrix{Float64}}) isa
                LowerTriangular
    end

    @testset "Consistent Uneven Distribution issue #166" begin
        DA = drand((2+length(OTHERIDS),), [MYID, OTHERIDS])
        @test fetch(@spawnat MYID length(localpart(DA)) == 2)
        @test fetch(@spawnat OTHERIDS length(localpart(DA)) == 1)
        close(DA)
        @test DistributedArrays.defaultdist(50,4) == [1,14,27,39,51]
    end

    @testset "Inhomogenous typeof(localpart)" begin
        block = 10
        Y = nworkers() * block
        X = nworkers() * block

        @assert nworkers() > 1
        @test_throws ErrorException DArray((X, Y)) do I
            eltype = first(CartesianIndices(I)) == CartesianIndex(1, 1) ? Int64 : Float64
            zeros(eltype, map(length, I))
        end
    end
end

check_leaks()

@testset "test DArray equality/copy/deepcopy" begin
    D = drand((200,200), [MYID, OTHERIDS])

    @testset "test isequal(::DArray, ::DArray)" begin
        DC = copy(D)
        @test D == DC
        close(DC)
    end

    @testset "test [deep]copy(::DArray) does a copy of each localpart" begin
        DC = copy(D)
        @spawnat OTHERIDS localpart(DC)[1] = 0
        @test fetch(@spawnat OTHERIDS localpart(D)[1] != 0)
        DD = deepcopy(D)
        @spawnat OTHERIDS localpart(DD)[1] = 0
        @test fetch(@spawnat OTHERIDS localpart(D)[1] != 0)
        close(DC)
        close(DD)
    end

    @testset "test copy(::DArray) is shallow" begin
        DA = @DArray [rand(100) for i=1:10]
        DC = copy(DA)
        id = procs(DC)[1]
        @test DA == DC
        fetch(@spawnat id localpart(DC)[1] .= -1.0)
        @test DA == DC
        @test fetch(@spawnat id all(localpart(DA)[1] .== -1.0))
        close(DA)
        close(DC)
    end

    @testset "test deepcopy(::DArray) is not shallow" begin
        DA = @DArray [rand(100) for i=1:10]
        DC = deepcopy(DA)
        id = procs(DC)[1]
        @test DA == DC
        fetch(@spawnat id localpart(DC)[1] .= -1.0)
        @test DA != DC
        @test fetch(@spawnat id all(localpart(DA)[1] .>= 0.0))
        close(DA)
        close(DC)
    end

    close(D)
end

check_leaks()

@testset "test DArray similar" begin
    D = drand((200,200), [MYID, OTHERIDS])
    DS = similar(D,Float16)

    @testset "test eltype of a similar" begin
        @test eltype(DS) == Float16
    end

    @testset "test dims of a similar" begin
        @test size(D) == size(DS)
    end
    close(D)
    close(DS)
end

check_leaks()

@testset "test DArray reshape" begin
    D = drand((200,200), [MYID, OTHERIDS])

    @testset "Test error-throwing in reshape" begin
        @test_throws DimensionMismatch reshape(D,(100,100))
    end

    DR = reshape(D,(100,400))
    @testset "Test reshape" begin
        @test size(DR) == (100,400)
    end
    close(D)
end

check_leaks()

@testset "test @DArray comprehension constructor" begin

    @testset "test valid use of @DArray" begin
        D = @DArray [i+j for i=1:10, j=1:10]
        @test D == [i+j for i=1:10, j=1:10]
        close(D)
    end

    @testset "test invalid use of @DArray" begin
        #@test_throws ArgumentError eval(:((@DArray [1,2,3,4])))
        @test_throws LoadError eval(:((@DArray [1,2,3,4])))
    end
end

check_leaks()

@testset "test DArray / Array conversion" begin
    D = drand((200,200), [MYID, OTHERIDS])

    @testset "test construct Array from (Sub)DArray" begin
        S = Matrix{Float64}(D[1:150, 1:150])
        A = Matrix{Float64}(D)

        @test A[1:150,1:150] == S
        D2 = DArray{Float64,2,Matrix{Float64}}(A)
        @test D2 == D
        DistributedArrays.allowscalar(true)
        @test fetch(@spawnat MYID localpart(D)[1,1]) == D[1,1]
        @test fetch(@spawnat OTHERIDS localpart(D)[1,1]) == D[1,101]
        DistributedArrays.allowscalar(false)
        close(D2)

        S2 = Vector{Float64}(D[4, 23:176])
        @test A[4, 23:176] == S2

        S3 = Vector{Float64}(D[23:176, 197])
        @test A[23:176, 197] == S3

        S4 = zeros(4)
        setindex!(S4, D[3:4, 99:100], :)
        # FixMe! Hitting the AbstractArray fallback here is extremely unfortunate but vec() becomes a ReshapedArray which makes it diffuclt to hit DArray methods. Unless this can be fixed in Base, we might have to add special methods for ReshapedArray{DArray}
        DistributedArrays.allowscalar(true)
        @test S4 == vec(D[3:4, 99:100])
        @test S4 == vec(A[3:4, 99:100])
        DistributedArrays.allowscalar(false)

        S5 = zeros(2,2)
        setindex!(S5, D[1,1:4], :, 1:2)
        # FixMe! Hitting the AbstractArray fallback here is extremely unfortunate but vec() becomes a ReshapedArray which makes it diffuclt to hit DArray methods. Unless this can be fixed in Base, we might have to add special methods for ReshapedArray{DArray}
        DistributedArrays.allowscalar(true)
        @test vec(S5) == D[1, 1:4]
        @test vec(S5) == A[1, 1:4]
        DistributedArrays.allowscalar(false)
    end
    close(D)
end

check_leaks()

@testset "test copy!" begin
    D1 = dzeros((10,10))
    r1 = remotecall_wait(() -> randn(3,10), workers()[1])
    r2 = remotecall_wait(() -> randn(7,10), workers()[2])
    D2 = DArray(reshape([r1; r2], 2, 1))
    copyto!(D2, D1)
    @test D1 == D2
    close(D1)
    close(D2)
end

check_leaks()

@testset "test DArray reduce" begin
    D = DArray(id->fill(myid(), map(length,id)), (10,10), [MYID, OTHERIDS])

    @testset "test reduce" begin
        @test reduce(+, D) == ((50*MYID) + (50*OTHERIDS))
    end

    @testset "test map / reduce" begin
        D2 = map(x->1, D)
        @test D2 isa DArray
        @test reduce(+, D2) == 100
        close(D2)
    end

    @testset "test map! / reduce" begin
        map!(x->1, D, D)
        @test reduce(+, D) == 100
    end
    close(D)
end

check_leaks()

@testset "test rmul" begin
    A = randn(100,100)
    DA = distribute(A)
    @test rmul!(DA, 2) == rmul!(A, 2)
    close(DA)
end

check_leaks()

@testset "test rmul!(Diagonal, A)" begin
    A = randn(100, 100)
    b = randn(100)
    D = Diagonal(b)
    DA = distribute(A)
    @test lmul!(D, A) == lmul!(D, DA)
    close(DA)
    A = randn(100, 100)
    b = randn(100)
    DA = distribute(A)
    @test rmul!(A, D) == rmul!(DA, D)
    close(DA)
end

check_leaks()

@testset "test mapreduce on DArrays" begin
    for _ = 1:25, f = [x -> Int128(2x), x -> Int128(x^2), x -> Int128(x^2 + 2x - 1)], opt = [+, *]
        A = rand(1:5, rand(2:30))
        DA = distribute(A)
        @test DA isa DArray
        @test mapreduce(f, opt, DA) - mapreduce(f, opt, A) == 0
        close(DA)
    end
end

check_leaks()

@testset "test mapreducedim on DArrays" begin
    D = DArray(I->fill(myid(), map(length,I)), (73,73), [MYID, OTHERIDS])
    D2 = map(x->1, D)
    @test D2 isa DArray
    @test mapreduce(t -> t*t, +, D2, dims=1) == mapreduce(t -> t*t, +, convert(Array, D2), dims=1)
    @test mapreduce(t -> t*t, +, D2, dims=2) == mapreduce(t -> t*t, +, convert(Array, D2), dims=2)
    @test mapreduce(t -> t*t, +, D2, dims=(1,2)) == mapreduce(t -> t*t, +, convert(Array, D2), dims=(1,2))

    # Test non-regularly chunked DArrays
    r1 = DistributedArrays.remotecall(() -> sprandn(3, 10, 0.1), workers()[1])
    r2 = DistributedArrays.remotecall(() -> sprandn(7, 10, 0.1), workers()[2])
    D = DArray(reshape([r1; r2], (2,1)))
    @test Array(sum(D, dims=2)) == sum(Array(D), dims=2)

    # close(D)
    # close(D2)
    d_closeall()   # temp created by the mapreduce above
end

check_leaks()

@testset "test mapreducdim, reducedim on DArrays" begin
    dims = (20,20,20)
    DA = drandn(dims)
    A = convert(Array, DA)

    @testset "dimension $dms" for dms in (1, 2, 3, (1,2), (1,3), (2,3), (1,2,3))
        @test mapreduce(t -> t*t, +, A, dims=dms) ≈ mapreduce(t -> t*t, +, DA, dims=dms)
        @test mapreduce(t -> t*t, +, A, dims=dms, init=1.0) ≈ mapreduce(t -> t*t, +, DA, dims=dms, init=1.0)
        @test reduce(*, A, dims=dms) ≈ reduce(*, DA, dims=dms)
        @test reduce(*, A, dims=dms, init=2.0) ≈ reduce(*, DA, dims=dms, init=2.0)
    end
    close(DA)
    d_closeall()   # temp created by the mapreduce above
end

check_leaks()

@testset "test statistical functions on DArrays" begin
    dims = (20,20,20)
    DA = drandn(dims)
    A = Array(DA)

    @testset "test $f for dimension $dms" for f in (mean, ), dms in (1, 2, 3, (1,2), (1,3), (2,3), (1,2,3))
        # std is pending implementation
        @test f(DA, dims=dms) ≈ f(A, dims=dms)
    end

    close(DA)
    d_closeall()   # temporaries created above
end

check_leaks()

unpack(ex::Base.CapturedException) = unpack(ex.ex)
unpack(ex::Distributed.RemoteException) = unpack(ex.captured)
if VERSION >= v"1.3.0-alpha.110"
unpack(ex::Base.TaskFailedException) = unpack(ex.task.exception)
end
unpack(ex) = ex

@testset "test sum on DArrays" begin
    A = randn(100,100)
    DA = distribute(A)

    # sum either throws an ArgumentError or a CompositeException of ArgumentErrors
    try
        sum(DA, dims=-1)
    catch err
        if isa(err, CompositeException)
            @test !isempty(err.exceptions)
            for excep in err.exceptions
                # Unpack the remote exception
                orig_err = unpack(excep)
                @test isa(orig_err, ArgumentError)
            end
        else
            @test isa(err, ArgumentError)
        end
    end
    try
        sum(DA, dims=0)
    catch err
        if isa(err, CompositeException)
            @test !isempty(err.exceptions)
            for excep in err.exceptions
                # Unpack the remote exception
                orig_err = unpack(excep)
                @test isa(orig_err, ArgumentError)
            end
        else
            @test isa(err, ArgumentError)
        end
    end

    @test sum(DA) ≈ sum(A)
    @test sum(DA, dims=1) ≈ sum(A, dims=1)
    @test sum(DA, dims=2) ≈ sum(A, dims=2)
    @test sum(DA, dims=3) ≈ sum(A, dims=3)
    close(DA)
    d_closeall()   # temporaries created above
end

check_leaks()

@testset "test size on DArrays" begin

    A = randn(100,100)
    DA = distribute(A)

    @test_throws BoundsError size(DA, 0)
    @test size(DA,1) == size(A,1)
    @test size(DA,2) == size(A,2)
    @test size(DA,3) == size(A,3)
    close(DA)
end

check_leaks()

# test length / lastindex
@testset "test collections API" begin
    A = randn(23,23)
    DA = distribute(A)

    @testset "test length" begin
        @test length(DA) == length(A)
    end

    @testset "test lastindex" begin
        @test lastindex(DA) == lastindex(A)
    end
    close(DA)
end

check_leaks()

@testset "test max / min / sum" begin
    a = map(x -> Int(round(rand() * 100)) - 50, Array{Int}(undef,100,1000))
    d = distribute(a)

    @test sum(d)          == sum(a)
    @test maximum(d)      == maximum(a)
    @test minimum(d)      == minimum(a)
    @test maximum(abs, d) == maximum(abs, a)
    @test minimum(abs, d) == minimum(abs, a)
    @test sum(abs, d)     == sum(abs, a)
    @test sum(abs2, d)    == sum(abs2, a)
    @test extrema(d)      == extrema(a)
    close(d)
end

check_leaks()

@testset "test all / any" begin
    a = map(x->Int(round(rand() * 100)) - 50, Array{Int}(undef,100,1000))
    a = [true for i in 1:100]
    d = distribute(a)

    @test all(d)
    @test any(d)

    close(d)

    a[50] = false
    d = distribute(a)
    @test !all(d)
    @test any(d)

    close(d)

    a = [false for i in 1:100]
    d = distribute(a)
    @test !all(d)
    @test !any(d)

    close(d)

    d = dones(10,10)
    @test !all(x-> x>1.0, d)
    @test all(x-> x>0.0, d)

    close(d)

    a = ones(10,10)
    a[10] = 2.0
    d = distribute(a)
    @test any(x-> x == 1.0, d)
    @test any(x-> x == 2.0, d)
    @test !any(x-> x == 3.0, d)

    close(d)
end

check_leaks()

@testset "test count"  begin
    a = ones(10,10)
    a[10] = 2.0
    d = distribute(a)

    @test count(x-> x == 2.0, d) == 1
    @test count(x-> x == 1.0, d) == 99
    @test count(x-> x == 0.0, d) == 0

    close(d)
end

check_leaks()

@testset "test prod" begin
    a = fill(2, 10);
    d = distribute(a);
    @test prod(d) == 2^10

    close(d)
end

check_leaks()

@testset "test zeros" begin
    @testset "1D dzeros default element type" begin
        A = dzeros(10)
        @test A == zeros(10)
        @test eltype(A) == Float64
        @test size(A) == (10,)
        close(A)
    end

    @testset "1D dzeros with specified element type" begin
        A = dzeros(Int, 10)
        @test A == zeros(10)
        @test eltype(A) == Int
        @test size(A) == (10,)
        close(A)
    end

    @testset "2D dzeros default element type, Dims constuctor" begin
        A = dzeros((10,10))
        @test A == zeros((10,10))
        @test eltype(A) == Float64
        @test size(A) == (10,10)
        close(A)
    end

    @testset "2D dzeros specified element type, Dims constructor" begin
        A = dzeros(Int, (10,10))
        @test A == zeros(Int, (10,10))
        @test eltype(A) == Int
        @test size(A) == (10,10)
        close(A)
    end

    @testset "2D dzeros, default element type" begin
        A = dzeros(10,10)
        @test A == zeros(10,10)
        @test eltype(A) == Float64
        @test size(A) == (10,10)
        close(A)
    end

    @testset "2D dzeros, specified element type" begin
        A = dzeros(Int, 10, 10)
        @test A == zeros(Int, 10, 10)
        @test eltype(A) == Int
        @test size(A) == (10,10)
        close(A)
    end
end

check_leaks()

@testset "test dones" begin
    @testset "1D dones default element type" begin
        A = dones(10)
        @test A == ones(10)
        @test eltype(A) == Float64
        @test size(A) == (10,)
        close(A)
    end

    @testset "1D dones with specified element type" begin
        A = dones(Int, 10)
        @test eltype(A) == Int
        @test size(A) == (10,)
        close(A)
    end

    @testset "2D dones default element type, Dims constuctor" begin
        A = dones((10,10))
        @test A == ones((10,10))
        @test eltype(A) == Float64
        @test size(A) == (10,10)
        close(A)
    end

    @testset "2D dones specified element type, Dims constructor" begin
        A = dones(Int, (10,10))
        @test A == ones(Int, (10,10))
        @test eltype(A) == Int
        @test size(A) == (10,10)
        close(A)
    end

    @testset "2D dones, default element type" begin
        A = dones(10,10)
        @test A == ones(10,10)
        @test eltype(A) == Float64
        @test size(A) == (10,10)
        close(A)
    end

    @testset "2D dones, specified element type" begin
        A = dones(Int, 10, 10)
        @test A == ones(Int, 10, 10)
        @test eltype(A) == Int
        @test size(A) == (10,10)
        close(A)
    end
end

check_leaks()

@testset "test drand" begin
    @testset "1D drand" begin
        A = drand(100)
        @test eltype(A) == Float64
        @test size(A) == (100,)
        @test all(x-> x >= 0.0 && x <= 1.0, A)
        close(A)
    end

    @testset "1D drand, specified element type" begin
        A = drand(Int, 100)
        @test eltype(A) == Int
        @test size(A) == (100,)
        close(A)
    end

    @testset "1D drand, UnitRange" begin
        A = drand(1:10, 100)
        @test eltype(A) == Int
        @test size(A) == (100,)
        close(A)
    end

    @testset "1D drand, Array" begin
        A = drand([-1,0,1], 100)
        @test eltype(A) == Int
        @test size(A) == (100,)
        close(A)
    end

    @testset "2D drand, Dims constructor" begin
        A = drand((50,50))
        @test eltype(A) == Float64
        @test size(A) == (50,50)
        @test all(x-> x >= 0.0 && x <= 1.0, A)
        close(A)
    end

    @testset "2D drand" begin
        A = drand(100,100)
        @test eltype(A) == Float64
        @test size(A) == (100,100)
        @test all(x-> x >= 0.0 && x <= 1.0, A)
        close(A)
    end

    @testset "2D drand, Dims constructor, specified element type" begin
        A = drand(Int, (100,100))
        @test eltype(A) == Int
        @test size(A) == (100,100)
        close(A)
    end

    @testset "2D drand, specified element type" begin
        A = drand(Int, 100, 100)
        @test eltype(A) == Int
        @test size(A) == (100,100)
        close(A)
    end
end

check_leaks()

@testset "test randn" begin
    @testset "1D drandn" begin
        A = drandn(100)
        @test eltype(A) == Float64
        @test size(A) == (100,)
        close(A)
    end

    @testset "2D drandn, Dims constructor" begin
        A = drandn((50,50))
        @test eltype(A) == Float64
        @test size(A) == (50,50)
        close(A)
    end

    @testset "2D drandn" begin
        A = drandn(100,100)
        @test eltype(A) == Float64
        @test size(A) == (100,100)
        close(A)
    end
end

check_leaks()

@testset "test transpose/adjoint" begin
    @testset "test transpose real" begin
        A = drand(Float64, 100, 200)
        @test copy(transpose(A)) == transpose(Array(A))
        close(A)
    end
    @testset "test transpose complex" begin
        A = drand(ComplexF64, 200, 100)
        @test copy(transpose(A)) == transpose(Array(A))
        close(A)
    end
    @testset "test adjoint real" begin
        A = drand(Float64, 200, 100)
        @test copy(adjoint(A)) == adjoint(Array(A))
        close(A)
    end
    @testset "test adjoint complex" begin
        A = drand(ComplexF64, 100, 200)
        @test copy(adjoint(A)) == adjoint(Array(A))
        close(A)
    end

    d_closeall()  # close the temporaries created above
end

check_leaks()

@testset "makelocal" begin
    A = randn(5*nprocs(), 5*nprocs())
    dA = distribute(A, procs=procs())
    for i in 1:size(dA, 2)
        a = DistributedArrays.makelocal(dA, :, i)
        @test all(Array(view(dA, :, i)) .== a)
        @test all(      view( A, :, i) .== a)
    end
    for i in 1:size(dA, 1)
        a = DistributedArrays.makelocal(dA, i, :)
        @test all(Array(view(dA, i:i, :)) .== a)
        @test all(      view( A, i:i, :) .== a)
    end
    a = DistributedArrays.makelocal(dA, 1:5, 1:5)
    @test all(Array(view(dA, 1:5, 1:5)) .== a)
    @test all(      view( A, 1:5, 1:5) .== a)
    close(dA)
end

@testset "test convert from subdarray" begin
    a = drand(20, 20);

    s = view(a, 1:5, 5:8)
    @test isa(s, SubDArray)
    @test s == DArray(s)

    s = view(a, 6:5, 5:8)
    @test isa(s, SubDArray)
    @test s == DArray(s)
    close(a)
    d_closeall()  # close the temporaries created above
end

check_leaks()

@testset "test scalar math" begin
    a = drand(20, 20);
    b = convert(Array, a)
    @testset "$f" for f in (-, abs, abs2, acos, acosd, acot,
              acotd, acsch, angle, asech, asin,
              asind, asinh, atan, atand, atanh,
              big, cbrt, ceil, cis, complex, conj,
              cos, cosc, cosd, cosh, cospi, cot,
              cotd, coth, csc, cscd, csch, dawson,
              deg2rad, digamma, erf, erfc, erfcinv,
              erfcx, erfi, erfinv, exp, exp10, exp2,
              expm1, exponent, float, floor, gamma, imag,
              invdigamma, isfinite, isinf, isnan,
              lgamma, log, log10, log1p, log2, rad2deg, real,
              sec, secd, sech, sign, sin, sinc, sind,
              sinh, sinpi, sqrt, tan, tand, tanh, trigamma)
        @test f.(a) == f.(b)
    end
    a = a .+ 1
    b = b .+ 1
    @testset "$f" for f in (asec, asecd, acosh, acsc, acscd, acoth)
        @test f.(a) == f.(b)
    end
    close(a)
    d_closeall()  # close the temporaries created above
end

check_leaks()

@testset "test mapslices" begin
    A = randn(5,5,5)
    D = distribute(A, procs = workers(), dist = [1, 1, min(nworkers(), 5)])
    @test mapslices(svdvals, D, dims=(1,2)) ≈ mapslices(svdvals, A, dims=(1,2))
    @test mapslices(svdvals, D, dims=(1,3)) ≈ mapslices(svdvals, A, dims=(1,3))
    @test mapslices(svdvals, D, dims=(2,3)) ≈ mapslices(svdvals, A, dims=(2,3))
    @test mapslices(sort, D, dims=(1,)) ≈ mapslices(sort, A, dims=(1,))
    @test mapslices(sort, D, dims=(2,)) ≈ mapslices(sort, A, dims=(2,))
    @test mapslices(sort, D, dims=(3,)) ≈ mapslices(sort, A, dims=(3,))

    # issue #3613
    B = mapslices(sum, dones(Float64, (2,3,4), workers(), [1,1,min(nworkers(),4)]), dims=[1,2])
    @test size(B) == (1,1,4)
    @test all(B.==6)

    # issue #5141
    C1 = mapslices(x-> maximum(-x), D, dims=[])
    @test C1 == -D

    # issue #5177
    c = dones(Float64, (2,3,4,5), workers(), [1,1,1,min(nworkers(),5)])
    m1 = mapslices(x-> ones(2,3), c, dims=[1,2])
    m2 = mapslices(x-> ones(2,4), c, dims=[1,3])
    m3 = mapslices(x-> ones(3,4), c, dims=[2,3])
    @test size(m1) == size(m2) == size(m3) == size(c)

    n1 = mapslices(x-> ones(6), c, dims=[1,2])
    n2 = mapslices(x-> ones(6), c, dims=[1,3])
    n3 = mapslices(x-> ones(6), c, dims=[2,3])
    n1a = mapslices(x-> ones(1,6), c, dims=[1,2])
    n2a = mapslices(x-> ones(1,6), c, dims=[1,3])
    n3a = mapslices(x-> ones(1,6), c, dims=[2,3])
    @test (size(n1a) == (1,6,4,5) && size(n2a) == (1,3,6,5) && size(n3a) == (2,1,6,5))
    @test (size(n1) == (6,1,4,5) && size(n2) == (6,3,1,5) && size(n3) == (2,6,1,5))
    close(D)
    close(c)
    d_closeall()  # close the temporaries created above
end

check_leaks()

@testset "test scalar ops" begin
    a = drand(20,20)
    b = convert(Array, a)
    c = drand(20,20)
    d = convert(Array, c)

    @testset "$f" for f in (:+, :-, :*, :/, :%)
        x = rand()
        @test @eval ($f).($a, $x) == ($f).($b, $x)
        @test @eval ($f).($x, $a) == ($f).($x, $b)
        @test @eval ($f).($a, $c) == ($f).($b, $d)
    end

    close(a)
    close(c)

    a = dones(Int, 20, 20)
    b = convert(Array, a)
    @testset "$f" for f in (:<<, :>>)
        @test @eval ($f).($a, 2)  == ($f).($b, 2)
        @test @eval ($f).(2, $a)  == ($f).(2, $b)
        @test @eval ($f).($a, $a) == ($f).($b, $b)
    end

    @testset "$f" for f in (:rem,)
        x = rand()
        @test @eval ($f).($a, $x) == ($f).($b, $x)
    end
    close(a)
    close(c)
    d_closeall()  # close the temporaries created above
end

check_leaks()

@testset "test broadcast ops" begin
    wrkrs = workers()
    nwrkrs = length(wrkrs)
    nrows = 20 * nwrkrs
    ncols = 10 * nwrkrs
    a = drand((nrows,ncols), wrkrs, (1, nwrkrs))
    m = mean(a, dims=1)
    c = a .- m
    d = convert(Array, a) .- convert(Array, m)
    @test c == d
    e = @DArray [ones(10) for i=1:4]
    f = 2 .* e
    @test Array(f) == 2 .* Array(e)
    @test Array(map(x -> sum(x) .+ 2, e)) == map(x -> sum(x) .+ 2, e)

    @testset "test nested broadcast" begin
       g = a .- m .* sin.(c)
       @test Array(g) == Array(a) .- Array(m) .* sin.(Array(c))
    end

    # @testset "lazy wrapped broadcast" begin
    #    l = similar(a)
    #    l[1:10, :] .= view(a, 1:10, : )
    # end
    d_closeall()
end

check_leaks()

@testset "test matrix multiplication" begin
    A = drandn(20,20)
    b = drandn(20)
    B = drandn(20,20)

    @test norm(convert(Array, A*b) - convert(Array, A)*convert(Array, b), Inf) < sqrt(eps())
    @test norm(convert(Array, A*B) - convert(Array, A)*convert(Array, B), Inf) < sqrt(eps())
    @test norm(convert(Array, A'*b) - convert(Array, A)'*convert(Array, b), Inf) < sqrt(eps())
    @test norm(convert(Array, A'*B) - convert(Array, A)'*convert(Array, B), Inf) < sqrt(eps())
    close(A)
    close(b)
    close(B)
    d_closeall()  # close the temporaries created above
end

check_leaks()

@testset "dot product" begin
    A = drandn(20,20)
    b = drandn(20)
    c = A * b

    @test dot(c, b) ≈ dot(convert(Array, c), convert(Array, b))
    close(A)
    close(b)
    close(c)
end

check_leaks()

@testset "test norm" begin
    x = drandn(20)

    @test abs(norm(x) - norm(convert(Array, x))) < sqrt(eps())
    @test abs(norm(x, 1) - norm(convert(Array, x), 1)) < sqrt(eps())
    @test abs(norm(x, 2) - norm(convert(Array, x), 2)) < sqrt(eps())
    @test abs(norm(x, Inf) - norm(convert(Array, x), Inf)) < sqrt(eps())
    close(x)
end

check_leaks()

@testset "test axpy!" begin
    for (x, y) in ((drandn(20), drandn(20)),
                   (drandn(20, 2), drandn(20, 2)))

        @test Array(axpy!(2.0, x, copy(y))) ≈ axpy!(2.0, Array(x), Array(y))
        @test_throws DimensionMismatch axpy!(2.0, x, zeros(length(x) + 1))
        close(x)
        close(y)
    end

    d_closeall()  # close the temporaries created above
end

check_leaks()

@testset "test ppeval" begin
    A = drandn((10, 10, nworkers()), workers(), [1, 1, nworkers()])
    B = drandn((10, nworkers()), workers(), [1, nworkers()])

    R = zeros(10, nworkers())
    for i = 1:nworkers()
        R[:, i] = convert(Array, A)[:, :, i]*convert(Array, B)[:, i]
    end
    @test convert(Array, ppeval(*, A, B)) ≈ R
    @test sum(ppeval(eigvals, A)) ≈ sum(ppeval(eigvals, A, Matrix{Float64}(I,10,10)))
    close(A)
    close(B)
    d_closeall()  # close the temporaries created above
end

check_leaks()

@testset "test nnz" begin
    A = sprandn(10, 10, 0.5)
    @test nnz(distribute(A)) == nnz(A)
end

@testset "test matmatmul" begin
    A = drandn(30, 30)
    B = drandn(30, 20)
    a = convert(Array, A)
    b = convert(Array, B)

    AB = A * B
    AtB = transpose(A) * B
    AcB = A' * B

    ab = a * b
    atb = transpose(a) * b
    acb = a' * b

    @test AB ≈ ab
    @test AtB ≈ atb
    @test AcB ≈ acb
    d_closeall()  # close the temporaries created above
end

@testset "sort, T = $T, 10^$i elements" for i in 0:6, T in [Int, Float64]
    d = DistributedArrays.drand(T, 10^i)
    @testset "sample = $sample" for sample in Any[true, false, (minimum(d),maximum(d)), rand(T, 10^i>512 ? 512 : 10^i)]
        d2 = DistributedArrays.sort(d; sample=sample)
        a  = convert(Array, d)
        a2 = convert(Array, d2)
        @test length(d) == length(d2)
        @test sort(a) == a2
    end
    d_closeall()  # close the temporaries created above
end

check_leaks()

@testset "ddata" begin
    d = ddata(;T=Int, init=I->myid())
    for p in workers()
        @test p == remotecall_fetch(d->d[:L], p, d)
    end
    @test Int[workers()...] == gather(d)

    close(d)

    d = ddata(;T=Int, data=workers())
    for p in workers()
        @test p == remotecall_fetch(d->d[:L], p, d)
    end
    @test Int[workers()...] == gather(d)

    close(d)

    d = ddata(;T=Any, init=I->"Hello World!")
    for p in workers()
        @test "Hello World!" == remotecall_fetch(d->d[:L], p, d)
    end
    Any["Hello World!" for p in workers()] == gather(d)

    close(d)
end

@testset "rand!" begin
    d = dzeros(30, 30)
    rand!(d)

    close(d)
end

check_leaks()

d_closeall()

@testset "test for any leaks" begin
    sleep(1.0)     # allow time for any cleanup to complete
    allrefszero = Bool[remotecall_fetch(()->length(DistributedArrays.refs) == 0, p) for p in procs()]
    @test all(allrefszero)

    allregistrieszero = Bool[remotecall_fetch(()->length(DistributedArrays.registry) == 0, p) for p in procs()]
    @test all(allregistrieszero)
end


text:
# Dynamic binding is invaluable for creating certain kinds of
# complex control flow in an intuitive way. Declare and access
# a dynamic var:

#     @dynamic a = 1
#     a[] #> 1

# Binding `a` alters its value for everywhere further down in
# the stack.

#     showa() = @show a[]
#     @dynamic let a = 2
#       showa()
#     end
#     #> a[] => 2

# Including for sub tasks which may continue to run:

#     function showa()
#       @schedule for i = 1:10
#         @show a[]
#         sleep(1)
#       end
#     end

#     @dynamic let a = 3
#       showa()
#     end

# This will print `a[] => 3` repeatedly, while evluating `a[]`
# will return `1`.

# You can also use the syntax

#     a[] = 5

# which will modify either the most recent binding value or the
# root of the var, as applicable.

# See also Clojure's dynamic vars for the inspiration:
# http://clojure.org/vars

import Base: getindex, setindex!

type Binding{T}
  root::T
end

root(b::Binding) = b.root

function dynamic_eq(def)
  name, val = def.args
  if isexpr(name, :(::))
    name, T = name.args
    :(const $(esc(name)) = Binding{$(esc(T))}($(esc(val))))
  else
    :(const $(esc(name)) = Binding($(esc(val))))
  end
end

function dynamic_let(ex)
  bindings = [:(bind($(esc(b.args[1])), $(esc(b.args[2])), t)) for b in ex.args[2:end]]
  quote
    t = @task $(esc(ex.args[1]))
    $(bindings...)
    schedule(t)
    wait(t)
  end
end

macro dynamic(def)
  isexpr(def, :(=)) ? dynamic_eq(def) :
  isexpr(def, :let) ? dynamic_let(def) :
  error("Unsupported @dynamic expression")
end

parent(t::Task) = t.parent

storage(t::Task) =
  t.storage == nothing ? (t.storage = ObjectIdDict()) :
  t.storage

isroot(task::Task) =
  task.parent == task

bindings(task::Task) =
  get!(storage(task), :bindings, Dict())

bind{T}(b::Binding{T}, v::T, t::Task) =
  bindings(t)[b] = v

function bind{T}(f::Function, b::Binding{T}, v::T)
  t = Task(f)
  bind(b, v, t)
  schedule(t)
  return wait(t)
end

binding{T}(b::Binding{T}, t::Task = current_task()) =
  haskey(bindings(t), b) ? bindings(t)[b]::T :
  isroot(t) ? root(b) : binding(b, parent(t))

set!{T}(b::Binding{T}, v::T, t::Task = current_task()) =
  haskey(bindings(t), b) ? (bindings(t)[b] = v) :
  isroot(t) ? (b.root = v) : set!(b, v, parent(t))

getindex(b::Binding) = binding(b)
setindex!{T}(b::Binding{T}, v::T) = set!(b, v)

ext: jl


text:
module MartaCT

using Reexport
import
    LinearAlgebra,
    IterTools,
    YAML,
    FFTW,
    Statistics,
    StatsBase,
    Distributions,
    IntervalSets

include("Utils.jl")
include("Applicative.jl")
include("Monads.jl")
include("TypeDict.jl")
include("CalibrationBase.jl")
include("Interpolation.jl")
include("Coordinates.jl")
include("Geometry.jl")
include("CTImages.jl")
include("FanBeam.jl")
include("AbstractAlgorithms.jl")
include("RadonAlgorithm.jl")
include("Simulations.jl")
include("CTIO.jl")
include("Info.jl")
include("TestImages.jl")
include("Calibration.jl")
include("CTScan.jl")
include("CTPlots.jl")

@reexport using
    .Utils,
    .Applicative,
    .Monads,
    .Info,
    .CalibrationBase,
    .Coordinates,
    .Geometry,
    .CTIO,
    .CTImages,
    .FanBeam,
    .AbstractAlgorithms,
    .RadonAlgorithm,
    .TestImages,
    .Calibration,
    .CTScan,
    .CTPlots

include("precompile_includer.jl")

end # module


text:
import DifferentiableBackwardEuler
import NLsolve
import NonlinearEquations
using OrdinaryDiffEq
import Random
import Zygote

include("utilities.jl")
include("inputdeck.jl")
include("ex.jl")

@NonlinearEquations.equations exclude=(dirichletnodes, neighbors, areasoverlengths) function groundwater(h, Ks, neighbors, areasoverlengths, dirichletnodes, dirichleths, Qs, specificstorage, volumes)
	NonlinearEquations.setnumequations(length(h))
	for i = 1:length(h)
		if i in dirichletnodes
			NonlinearEquations.addterm(i, h[i] - dirichleths[i])
		else
			NonlinearEquations.addterm(i, Qs[i] / (specificstorage[i] * volumes[i]))
		end
	end
	for (i, (node_a, node_b)) in enumerate(neighbors)
		for (node1, node2) in [(node_a, node_b), (node_b, node_a)]
			if !(node1 in dirichletnodes) && !(node2 in dirichletnodes)
				NonlinearEquations.addterm(node1, Ks[i] * (h[node2] - h[node1]) * areasoverlengths[i] / (specificstorage[node1] * volumes[node1]))
			elseif !(node1 in dirichletnodes) && node2 in dirichletnodes
				NonlinearEquations.addterm(node1, Ks[i] * (dirichleths[node2] - h[node1]) * areasoverlengths[i] / (specificstorage[node1] * volumes[node1]))
			end
		end
	end
end

groundwater_residuals(us, Ks) = groundwater_residuals(us, Ks, neighbors, areasoverlengths, dirichletnodes, dirichleths, Qs, specificstorage, volumes)
groundwater_h(us, Ks) = groundwater_h(us, Ks, neighbors, areasoverlengths, dirichletnodes, dirichleths, Qs, specificstorage, volumes)
groundwater_Ks(us, Ks) = groundwater_Ks(us, Ks, neighbors, areasoverlengths, dirichletnodes, dirichleths, Qs, specificstorage, volumes)

function unpack(p)
	@assert length(p) == length(neighbors)
	Ks = p[1:length(neighbors)]
	return Ks
end

function f(u, p, t)
	Ks = unpack(p)
	return groundwater_residuals(u, Ks)
end

function f_u(u, p, t)
	Ks = unpack(p)
	return groundwater_h(u, Ks, neighbors, areasoverlengths, dirichletnodes, dirichleths, Qs, specificstorage, volumes)
end

function f_p(u, p, t)
	Ks = unpack(p)
	return groundwater_Ks(u, Ks, neighbors, areasoverlengths, dirichletnodes, dirichleths, Qs, specificstorage, volumes)
end

f_t(u, p, t) = zeros(length(u))

h0 = zeros(size(coords, 2))
h0[dirichletnodes] = dirichleths[dirichletnodes]
p = Ks
tspan = [0, 1e4]
odef = ODEFunction(f; jac=f_u, jac_prototype=f_u(h0, p, 0.0))
prob = ODEProblem(odef, h0, tspan, p)
@time soln_diffeq = solve(prob, ImplicitEuler())
@time soln_dbe = DifferentiableBackwardEuler.steps(h0, f, f_u, f_p, f_t, p, soln_diffeq.t; ftol=1e-8)
sum((soln_diffeq[:, :] .- soln_dbe) .^ 2)
@test isapprox(soln_diffeq[:, :], soln_dbe)
g(p) = DifferentiableBackwardEuler.steps(h0, f, f_u, f_p, f_t, p, soln_diffeq.t; ftol=1e-8)[obsnode, end]
@time dgdp_zygote = Zygote.gradient(g, p)[1]
@test isapprox(dgdp_zygote, grad)#this checks the zygote gradient against the steady state gradient that doesn't use zygote


text:
## To Lower Case ##

function lowercase!{T<:ASCIIChar}(s::AbstractArray{T})
  for i = 1:length(s)
    if 'A' <= char(s[i]) <= 'Z'
      s[i] += 32
    end
  end
  return s
end

lowercase{T<:ASCIIChar}(s::AbstractArray{T}) = lowercase!(copy(s))

## To Upper Case ##

function uppercase!{T<:ASCIIChar}(s::AbstractArray{T})
  for i = 1:length(s)
    if 'a' <= char(s[i]) <= 'z'
      s[i] -= 32
    end
  end
  return s
end

uppercase{T<:ASCIIChar}(s::AbstractArray{T}) = uppercase!(copy(s))

## Upper and Lower for ASCIIChar ##

uppercase{T<:ASCIIChar}(x::T) = convert(T,uppercase(uint8(x)))
lowercase{T<:ASCIIChar}(x::T) = convert(T,lowercase(uint8(x)))

## Matching ##

ismatch{T<:ASCIIChar}(r::Regex, s::Vector{T}) = ismatch(r, bytestring(s))

contains{T<:ASCIIChar}(s::Vector{T}, r::Regex) = ismatch(r,s)

match{T<:ASCIIChar}(re::Regex, s::Vector{T}, idx::Integer) = match(re,bytestring(s),idx)
match{T<:ASCIIChar}(r::Regex, s::Vector{T}) = match(r, bytestring(s), 1)

search{T<:ASCIIChar}(s::Vector{T}, r::Regex, idx::Integer) = search(bytestring(s),r,idx)
search{T<:ASCIIChar}(s::Vector{T}, r::Regex) = search(bytestring(s),r,1)

memchr{T<:ASCIIChar}(s::Vector{T}, c, i::Integer) = memchr(reinterpret(Uint8,s),c,i)
memchr{T<:ASCIIChar}(s::Vector{T}, c) = memchr(reinterpret(Uint8,s),c,1)

search{T<:ASCIIChar}(s::Vector{T}, c::Integer, i::Integer) = (res = memchr(s,c,i); (res,res+1))
search{T<:ASCIIChar}(s::Vector{T}, c::Integer) = search(s,c,1)

each_match{T<:ASCIIChar}(re::Regex, s::Vector{T}, ovr::Bool) = RegexMatchIterator(re,bytestring(s),ovr)
each_match{T<:ASCIIChar}(re::Regex, s::Vector{T}) = RegexMatchIterator(re,bytestring(s),false)

## Some String Function ##

repeat{T<:ASCIIChar}(s::Vector{T}, n::Int) = vec(repmat(s,n,1))

repeat{T<:ASCIIChar}(x::T, n::Int) = fill(x,n)

function swap!{T<:ASCIIChar,Tp<:Integer,Tc<:Integer}(x::AbstractArray{T},p::Tp,c::Tc)
  pat = convert(T,p)
  cha = convert(T,c)
  for i in 1:length(x)
    if x[i] == pat
      x[i] = cha
    end
  end
  return x
end

swap{T<:ASCIIChar,Tp<:Integer,Tc<:Integer}(x::AbstractArray{T},p::Tp,c::Tc) = swap!(copy(x),p,c)

function replace{T<:ASCIIChar}(s::Vector{T}, pattern, r, limit::Integer)
    n = 1
    i = a = start(s)
    j, k  = search(s,pattern,i)
    res::Vector{T} = s[1:(j-1)]
    while j != 0
        if i < k
           if i!=a; append!(res,s[i:(j-1)]) end
	   isa(r,Vector) ? append!(res,r) : push!(res,r)
	   i = k
        end
        if k <= j; k = nextind(s,j) end
        j, k = search(s,pattern,k)
        if n == limit break end
        n += 1
    end
    if i<=length(s); append!(res,s[i:end]) end
    res
end

replace{T<:ASCIIChar}(s::Vector{T}, pat, r) = replace(s, pat, r, 0)


text:
"""
    abstract type AbstractBlockTask <: LearningTask

Abstract supertype for learning tasks that derive their
functionality from [`Block`](#)s and [`Encoding`](#)s.

These learning tasks require you only to specify blocks and
encodings by defining which blocks of data show up at which
stage of the pipeline. Generally, a subtype will have a field
`blocks` of type `NamedTuple` that contains this information
and a field `encodings` of encodings that are applied to samples.
They can be accessed with `getblocks` and `getencodings`
respectively. For example, [`SupervisedTask`](#) represents a
learning task where each sample consists of an input and a target.

{cell=main}
```julia
using FastAI
task = SupervisedTask(
    (Image{2}(), Label(["cat", "dog"])),
    (ImagePreprocessing(), OneHot(),)
)
FastAI.getblocks(task)
```

To implement a new `AbstractBlockTask` either

- use the helper [`BlockTask`](#) (simpler)
- or subtype [`AbstractBlockTask`](#) (allows customization through
    dispatch)

## Blocks and interfaces

To support different learning task interfaces, a `AbstractBlockTask`'s
blocks need to contain different blocks. Below we list first block names
with descriptions, and afterwards relevant interface functions and which
blocks are required to use them.

### Blocks

Each name corresponds to a key of the named tuple
`blocks = getblocks(task)`). A block is referred to with `blocks.\$name`
and an instance of data from a block is referred to as `\$name`.

- `blocks.sample`: The most important block, representing one full
    observation of unprocessed data. Data containers used with a learning
    task should have compatible observations, i.e.
    `checkblock(blocks.sample, data[i])`.
- `blocks.x`: Data that will be fed into the model, i.e. (neglecting batching)
    `model(x)` should work
- `blocks.ŷ`: Data that is output by the model, i.e. (neglecting batching)
    `checkblock(blocks.ŷ, model(x))`
- `blocks.y`: Data that is compared to the model output using a loss function,
    i.e. `lossfn(ŷ, y)`
- `blocks.encodedsample`: An encoded version of `blocks.sample`. Will usually
    correspond to `encodedblockfilled(getencodings(task), blocks.sample)`.

### Interfaces/functionality and required blocks:

Core:
- [`encode`](#)`(task, ctx, sample)` requires `sample`. Also enables use of
    [`taskdataset`](#), [`taskdataloaders`](#)
- [`decode`](#)`(task, ctx, encodedsample)` requires `encodedsample`
- [`decodeypred`](#)`(task, ctx, ŷ)` requires `ŷ`
- [`decodey`](#)`(task, ctx, y)` requires `y`

Training:
- [`taskmodel`](#)`(task)` requires `x`, `ŷ`
- [`tasklossfn`](#)`(task)` requires `y`, `ŷ`

Visualization:
- [`showsample`](#), [`showsamples`](#) require `sample`
- [`showencodedsample`](#), [`showencodedsamples`](#), [`showbatch`](#)
    require `encodedsample`
- [`showsample`](#), [`showsamples`](#) require `sample`
- [`showoutput`](#), [`showoutputs`](#), [`showoutputbatch`](#) require
    `ŷ`, `encodedsample`

Testing:
- [`mockmodel`](#)`(task)` requires `x`, `ŷ`
- [`mocksample`](#)`(task)` requires `sample`

"""
abstract type AbstractBlockTask <: LearningTask end

getblocks(task::AbstractBlockTask) = task.blocks
getencodings(task::AbstractBlockTask) = task.encodings

# Core interface

function encodesample(task::AbstractBlockTask, context, sample)
    encode(getencodings(task), context, getblocks(task).sample, sample)
end

function encodeinput(task::AbstractBlockTask, context, input)
    encode(getencodings(task), context, getblocks(task).input, input)
end

function encodetarget(task::AbstractBlockTask, context, target)
    encode(getencodings(task), context, getblocks(task).target, target)
end

function decode(task::AbstractBlockTask, context, encodedsample)
    decode(getencodings(task), context, getblocks(task).encodedsample, encodedsample)
end

function decodeypred(task::AbstractBlockTask, context, ŷ)
    decode(getencodings(task), context, getblocks(task).ŷ, ŷ)
end

function decodey(task::AbstractBlockTask, context, y)
    decode(getencodings(task), context, getblocks(task).y, y)
end

# Training interface

function taskmodel(task::AbstractBlockTask, backbone)
    return blockmodel(getblocks(task).x, getblocks(task).ŷ, backbone)
end

function taskmodel(task::AbstractBlockTask)
    backbone = blockbackbone(getblocks(task).x)
    return blockmodel(getblocks(task).x, getblocks(task).ŷ, backbone)
end

function tasklossfn(task::AbstractBlockTask)
    return blocklossfn(getblocks(task).ŷ, getblocks(task).y)
end

# Testing interface

mocksample(task::AbstractBlockTask) = mockblock(task, :sample)
mockblock(task::AbstractBlockTask, name::Symbol) = mockblock(getblocks(task)[name])

mockmodel(task::AbstractBlockTask) =
    mockmodel(getblocks(task).x, getblocks(task).ŷ)

"""
    mockmodel(xblock, ŷblock)
    mockmodel(task::AbstractBlockTask)

Create a fake model that maps batches of block `xblock` to batches of block
`ŷblock`. Useful for testing.
"""
function mockmodel(_, ŷblock)
    return function mockmodel_block(xs)
        return MLUtils.batch([mockblock(ŷblock) for _ in 1:Datasets.batchsize(xs)])
    end
end

# ## Block task
#
# `BlockTask` is a helper to create anonymous block tasks.

"""
    BlockTask(blocks, encodings)

Create an [`AbstractBlockTask`](#) directly, passing in a named tuple `blocks`
and `encodings`. See [`SupervisedTask`](#) for supervised training tasks.
"""
struct BlockTask{B<:NamedTuple,E} <: AbstractBlockTask
    blocks::B
    encodings::E
end

Base.show(io::IO, task::BlockTask) = print(io,
    "BlockTask(blocks=", keys(getblocks(task)), ")")

# ## Supervised learning task

"""
    SupervisedTask((inputblock, targetblock), encodings)

A [`AbstractBlockTask`](#) learning task for the supervised
task of learning to predict a `target` given an `input`. `encodings`
are applied to samples before being input to the model. Model outputs
are decoded using those same encodings to get a target prediction.

In addition, to the blocks defined by [`AbstractBlockTask`](#),
`getblocks(::SupervisedTask)` defines the following blocks:

By default the model output is assumed to be an encoded target, but the
`ŷblock` keyword argument to overwrite this.

- `blocks.input`: An unencoded input and the first element in the tuple
    `sample = (input, target)`
- `blocks.target`: An unencoded target and the second element in the tuple
    `sample = (input, target)`
- `blocks.pred`: A prediction. Usually the same as `blocks.target` but may
    differ if a custom `ŷblock` is specified.

A `SupervisedTask` also enables some additional functionality:

- [`encodeinput`](#)
- [`encodetarget`](#)
- [`showprediction`](#), [`showpredictions`](#)
"""
struct SupervisedTask{B<:NamedTuple,E} <: AbstractBlockTask
    blocks::B
    encodings::E
end

function SupervisedTask(blocks::Tuple{Any,Any}, encodings; ŷblock = nothing)
    sample = input, target = blocks
    x, y = encodedsample = encodedblockfilled(encodings, sample)
    ŷ = isnothing(ŷblock) ? y : ŷblock
    pred = decodedblockfilled(encodings, ŷ)
    blocks = (; input, target, sample, encodedsample, x, y, ŷ, pred)
    SupervisedTask(blocks, encodings)
end

function Base.show(io::IO, task::SupervisedTask)
    print(
        io,
        "SupervisedTask(",
        summary(getblocks(task).input),
        " -> ",
        summary(getblocks(task).target),
        ")",
    )
end

# ## Deprecations

Base.@deprecate BlockTask(blocks::Tuple{Any, Any}, encodings; kwargs...) SupervisedTask(blocks, encodings; kwargs...)


text:
using Grid
using Base.Test

x = -1.0:0.1:1.0
z = sin(x)

cg = CoordInterpGrid(x,z,BCnil,InterpQuadratic)

@test abs(cg[0.0] - sin(0.0)) < 100*eps()
@test abs(cg[0.3] - sin(0.3)) < 100*eps()

v,g = valgrad(cg,0.0)
@test abs(g - cos(0.0)) < 0.001

v,g = valgrad(cg,0.3)
@test abs(g - cos(0.3)) < 0.001

y = -2.0:0.1:2.0
z = Float64[sin(i+j) for i in x, j in y]

@test_throws DimensionMismatch cg = CoordInterpGrid((y,x),z,BCnil,InterpQuadratic)
cg = CoordInterpGrid((x,y),z,BCnil,InterpQuadratic)

@test abs(cg[0.0,0.0] - sin(0.0)) < 0.0001
@test abs(cg[0.0,0.3] - sin(0.3)) < 0.0001
@test abs(cg[0.3,0.0] - sin(0.3)) < 0.0001

v,g = valgrad(cg,0.0,0.0)
@test abs(g[1] - cos(0.0)) < 0.001
@test abs(g[2] - cos(0.0)) < 0.001

v = valgrad(g,cg,0.3,0.0)
@test abs(g[1] - cos(0.3)) < 0.001
@test abs(g[2] - cos(0.3)) < 0.001

# Issue #29
xg = -10:.5:10
yg = -3:.15:3.5
zg = Float64[sin(x)*cos(y) for x in xg, y in yg]
ig = CoordInterpGrid((xg,yg),zg, 0, InterpLinear)
@test size(ig) == size(zg)
@test ig[xg[1],yg[1]] == zg[1,1]


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct CheckNameAvailabilityParameters <: SwaggerModel
    name::Any # spec type: Union{ Nothing, String } # spec name: name
    type::Any # spec type: Union{ Nothing, String } # spec name: type

    function CheckNameAvailabilityParameters(;name=nothing, type=nothing)
        o = new()
        validate_property(CheckNameAvailabilityParameters, Symbol("name"), name)
        setfield!(o, Symbol("name"), name)
        validate_property(CheckNameAvailabilityParameters, Symbol("type"), type)
        setfield!(o, Symbol("type"), type)
        o
    end
end # type CheckNameAvailabilityParameters

const _property_map_CheckNameAvailabilityParameters = Dict{Symbol,Symbol}(Symbol("name")=>Symbol("name"), Symbol("type")=>Symbol("type"))
const _property_types_CheckNameAvailabilityParameters = Dict{Symbol,String}(Symbol("name")=>"String", Symbol("type")=>"String")
Base.propertynames(::Type{ CheckNameAvailabilityParameters }) = collect(keys(_property_map_CheckNameAvailabilityParameters))
Swagger.property_type(::Type{ CheckNameAvailabilityParameters }, name::Symbol) = Union{Nothing,eval(Meta.parse(_property_types_CheckNameAvailabilityParameters[name]))}
Swagger.field_name(::Type{ CheckNameAvailabilityParameters }, property_name::Symbol) =  _property_map_CheckNameAvailabilityParameters[property_name]

const _allowed_CheckNameAvailabilityParameters_type = ["Microsoft.DataLakeStore/accounts"]

function check_required(o::CheckNameAvailabilityParameters)
    (getproperty(o, Symbol("name")) === nothing) && (return false)
    true
end

function validate_property(::Type{ CheckNameAvailabilityParameters }, name::Symbol, val)
    if name === Symbol("type")
        Swagger.validate_param(name, "CheckNameAvailabilityParameters", :enum, val, _allowed_CheckNameAvailabilityParameters_type)
    end
end


text:
module TestStat
  include("../src/stat.jl")
  using DataFrames
  using StatsBase
  using Base.Test

  nona = @data [1, 2, 3, 4, 5]
  onena = @data [1, NA, 3, 4, 5]
  twona = @data [1, NA, 3, NA, 5]

  # Using 0 NA
  # Track StatsBase zscore results. They use sample stddev as default, so we
  # should too for consistency. Throw an error if this changes.
  statsbase_zscore = @data zscore(nona)
  statsbase_zscore_pop = @data zscore(nona, 3., √2)
  @test isequal(nazscore(nona), statsbase_zscore)
  @test isequal(nazscore_pop(nona), statsbase_zscore_pop)
  # Give the objectively correct results as a sanity check in case StatsBase
  # gets whacky.
  expected_nona = @data [
    -1.2649110640673518,
    -0.6324555320336759,
    0.0,
    0.6324555320336759,
    1.2649110640673518
  ]
  expected_nona_pop = @data [
    -1.414213562373095,
    -0.7071067811865475,
    0.0,
    0.7071067811865475,
    1.414213562373095
  ]
  @test isequal(nazscore(nona), expected_nona)
  @test isequal(nazscore_pop(nona), expected_nona_pop)

  # Using 1 NA
  expected_onena = @data [
    -1.3174650984805198,
    NA,
    -0.14638501094227999,
    0.43915503282683993,
    1.02469507659596
  ]
  expected_onena_pop = @data [
    -1.52127765851133,
    NA,
    -0.1690308509457033,
    0.50709255283711,
    1.1832159566199232
  ]

  @test isequal(nazscore(onena), expected_onena)
  @test isequal(nazscore_pop(onena), expected_onena_pop)

  # Using 2 NA
  expected_twona = @data [-1., NA, 0., NA, 1.]
  expected_twona_pop = @data [
    -1.224744871391589,
    NA,
    0.0,
    NA,
    1.224744871391589
  ]

  @test isequal(nazscore(twona), expected_twona)
  @test isequal(nazscore_pop(twona), expected_twona_pop)
end

max_stars_repo_name: jnu/FantasyFootball

max_issues_repo_name: jnu/FantasyFootball

max_forks_repo_name: jnu/FantasyFootball


text:
## GaussianMixtures.jl  Some functions for potentially large Gaussian Mixture Models
## (c) 2013--2015 David A. van Leeuwen

__precompile__()

module GaussianMixtures

using Distributions
using PDMats
using Clustering
using JLD
using Compat

include("compat.jl")
include("gmmtypes.jl")
include("bayestypes.jl")

include("gmms.jl")
include("train.jl")
include("bayes.jl")
include("io.jl")
include("stats.jl")
include("rand.jl")
include("data.jl")
include("recognizer.jl")

include("distributions.jl")
include("scikitlearn.jl")

export GMM, VGMM, GMMprior, CSstats, Cstats, History, Data, DataOrMatrix,
gmmsplit, em!, maxapost, llpg, avll, gmmposterior, sanitycheck!,
history, show, stats, nparams, means, covars, weights, setmem, vec, rand, kind, dmap

## some init code.  Turn off subnormal computation, as it is slow.  This is a global setting...
set_zero_subnormals(true)

end


text:
"""
    function addUpdateLocal_!(p::Program_,platform::String)

Generate the functions related with Local Updates.
"""
function addUpdateLocal_!(p::Program_,platform::String)

    if "UpdateLocal" in keys(p.agent.declaredUpdates)

        push!(p.execInloop.args,
            :(@platformAdapt locStep_!(ARGS_))
        )

        code = p.agent.declaredUpdates["UpdateLocal"]

        #Add events
        code = addEventAddAgent_(code, p, platform)
        code = addEventRemoveAgent_(code, p, platform)

        #Construct functions
        f = simpleFirstLoopWrapInFunction_(platform,:locStep_!,code)
        f = vectorize_(p.agent,f,p)

        #Updates to localV
        f = postwalk(x->@capture(x,localVCopy[g__] = localVCopy[g2__]) ? :(localV[$(g...)] = localVCopy[$(g2...)]) : x, f)
        f = postwalk(x->@capture(x,identityVCopy[g__] = identityVCopy[g2__]) ? :(identityV[$(g...)] = identityVCopy[$(g2...)]) : x, f)

        push!(p.declareF.args,
            f)

    end

    return nothing
end


text:
@testset "Constant" begin
    F1 = Constant(randn(Float64, (2,2)))
    F2 = Constant(randn(ComplexF64, (2,2)))

    p = 3*rand()
    a = rand()
    ab = (a, p+a)
    x = range(ab..., length = 11)

    α, β = randn(), randn(ComplexF64)
    A = randn(2,2)

    @test (F1 + F2).(x) ≈ F1.(x) .+ F2.(x)
    @test (F1 - F2).(x) ≈ F1.(x) .- F2.(x)
    @test (F2 - F1).(x) ≈ F2.(x) .- F1.(x)
    @test (α*F1).(x) ≈ α .* F1.(x)
    @test (β*F1).(x) ≈ β .* F1.(x)
    @test (β*F1 + α*F2).(x) ≈ β .* F1.(x) .+ α .* F2.(x)
    @test (β*F1 - α\F2).(x) ≈ β .* F1.(x) .- α .\ F2.(x)
    @test (A*F2*A).(x) ≈ Ref(A) .* F2.(x) .* Ref(A)
    @test (F1 * F2).(x) ≈ F1.(x) .* F2.(x)
    @test (transpose(F2)).(x) ≈ transpose.(F2.(x))
    @test (conj(F2)).(x) ≈ conj.(F2.(x))
    @test (adjoint(F2)).(x) ≈ adjoint.(F2.(x))
    @test (real(F2)).(x) ≈ real.(F2.(x))
    @test (imag(F2)).(x) ≈ imag.(F2.(x))
    @test localdot(F1, F2).(x) ≈ dot.(F1.(x), F2.(x))

    @test rmul!(copy(F1), α).(x) ≈ α .* F1.(x)
    @test rmul!(copy(F2), α).(x) ≈ α .* F2.(x)
    @test rmul!(copy(F2), β).(x) ≈ β .* F2.(x)

    @test lmul!(α, copy(F1)).(x) ≈ α .* F1.(x)
    @test lmul!(α, copy(F2)).(x) ≈ α .* F2.(x)
    @test lmul!(β, copy(F2)).(x) ≈ β .* F2.(x)

    @test mul!(zero(F2), β, F1).(x) ≈ β .* F1.(x)

    @test axpy!(β, F1, copy(F2)).(x) ≈ β .* F1.(x) .+ F2.(x)
    @test axpby!(β, F1, α, copy(F2)).(x) ≈ β .* F1.(x) .+ α .* F2.(x)

    @test rmul!(copy(F1), α) ≈ F1 * α
    @test lmul!(β, copy(F2)) ≈ β * F2

    @test fit(F1, Constant, ab) ≈ F1
    @test fit(F2, Constant, ab) ≈ F2
end

@testset "FourierSeries" begin
    p = 3*rand()
    a = rand()
    ab = (a, a+p)
    x = range(ab..., length = 201)

    F1 = FourierSeries([randn(Float64, (2,2)) for k = 1:5], p)
    F2 = FourierSeries([randn(ComplexF64, (2,2)) for k = 1:3], p)

    α, β = randn(), randn(ComplexF64)
    A = randn(2,2)

    @test (F1 + F2).(x) ≈ F1.(x) .+ F2.(x)
    @test (F1 - F2).(x) ≈ F1.(x) .- F2.(x)
    @test (F2 - F1).(x) ≈ F2.(x) .- F1.(x)
    @test (α*F1).(x) ≈ α .* F1.(x)
    @test (β*F1).(x) ≈ β .* F1.(x)
    @test (β*F1 + α*F2).(x) ≈ β .* F1.(x) .+ α .* F2.(x)
    @test (β*F1 - α\F2).(x) ≈ β .* F1.(x) .- α .\ F2.(x)
    @test (A*F2*A).(x) ≈ Ref(A) .* F2.(x) .* Ref(A)
    @test (F1 * F2).(x) ≈ F1.(x) .* F2.(x)
    @test (transpose(F2)).(x) ≈ transpose.(F2.(x))
    @test (conj(F2)).(x) ≈ conj.(F2.(x))
    @test (adjoint(F2)).(x) ≈ adjoint.(F2.(x))
    @test (real(F2)).(x) ≈ real.(F2.(x))
    @test (imag(F2)).(x) ≈ imag.(F2.(x))
    @test localdot(F1, F2).(x) ≈ dot.(F1.(x), F2.(x))

    @test rmul!(copy(F1), α).(x) ≈ α .* F1.(x)
    @test rmul!(copy(F2), α).(x) ≈ α .* F2.(x)
    @test rmul!(copy(F2), β).(x) ≈ β .* F2.(x)

    @test lmul!(α, copy(F1)).(x) ≈ α .* F1.(x)
    @test lmul!(α, copy(F2)).(x) ≈ α .* F2.(x)
    @test lmul!(β, copy(F2)).(x) ≈ β .* F2.(x)

    @test mul!(zero(F2), β, F1).(x) ≈ β .* F1.(x)

    @test axpy!(β, F1, copy(F2)).(x) ≈ β .* F1.(x) .+ F2.(x)
    @test axpby!(β, F1, α, copy(F2)).(x) ≈ β .* F1.(x) .+ α .* F2.(x)

    @test rmul!(copy(F1), α) ≈ F1 * α
    @test lmul!(β, copy(F2)) ≈ β * F2

    @test fit(F1, FourierSeries, ab) ≈ F1
    @test fit(F2, FourierSeries, ab) ≈ F2
end

@testset "TaylorSeries" begin
    p = 3*rand()
    a = rand()
    ab = (a, a+p)
    x = range(ab..., length = 201)

    F1 = TaylorSeries([randn(Float64, (2,2)) for k = 1:5], a)
    F2 = TaylorSeries([randn(ComplexF64, (2,2)) for k = 1:3], a)

    α, β = randn(), randn(ComplexF64)
    A = randn(2,2)

    @test (F1 + F2).(x) ≈ F1.(x) .+ F2.(x)
    @test (F1 - F2).(x) ≈ F1.(x) .- F2.(x)
    @test (F2 - F1).(x) ≈ F2.(x) .- F1.(x)
    @test (α*F1).(x) ≈ α .* F1.(x)
    @test (β*F1).(x) ≈ β .* F1.(x)
    @test (β*F1 + α*F2).(x) ≈ β .* F1.(x) .+ α .* F2.(x)
    @test (β*F1 - α\F2).(x) ≈ β .* F1.(x) .- α .\ F2.(x)
    @test (A*F2*A).(x) ≈ Ref(A) .* F2.(x) .* Ref(A)
    @test (F1 * F2).(x) ≈ F1.(x) .* F2.(x)
    @test (transpose(F2)).(x) ≈ transpose.(F2.(x))
    @test (conj(F2)).(x) ≈ conj.(F2.(x))
    @test (adjoint(F2)).(x) ≈ adjoint.(F2.(x))
    @test (real(F2)).(x) ≈ real.(F2.(x))
    @test (imag(F2)).(x) ≈ imag.(F2.(x))
    @test localdot(F1, F2).(x) ≈ dot.(F1.(x), F2.(x))

    @test rmul!(copy(F1), α).(x) ≈ α .* F1.(x)
    @test rmul!(copy(F2), α).(x) ≈ α .* F2.(x)
    @test rmul!(copy(F2), β).(x) ≈ β .* F2.(x)

    @test lmul!(α, copy(F1)).(x) ≈ α .* F1.(x)
    @test lmul!(α, copy(F2)).(x) ≈ α .* F2.(x)
    @test lmul!(β, copy(F2)).(x) ≈ β .* F2.(x)

    @test mul!(zero(F2), β, F1).(x) ≈ β .* F1.(x)

    @test axpy!(β, F1, copy(F2)).(x) ≈ β .* F1.(x) .+ F2.(x)
    @test axpby!(β, F1, α, copy(F2)).(x) ≈ β .* F1.(x) .+ α .* F2.(x)

    @test rmul!(copy(F1), α).(x) ≈ (F1 * α).(x)
    @test lmul!(β, copy(F2)).(x) ≈ (β * F2).(x)

    @test fit(F1, TaylorSeries, ab).(x) ≈ F1.(x)
    @test fit(F2, TaylorSeries, ab).(x) ≈ F2.(x)
end


text:
"""
6, [1, 4] => [2, 3, 5, 6]
"""
function complement(N::Int, indices::Vector{Int})
    L = length(indices)
    x = Vector{Int}(undef, N - L)
    i_ = 1 # Position in the x vector
    j = 1 # Position in indices vector
    for i=1:N
        if j > L || indices[j]!=i
            x[i_] = i
            i_ += 1
        else
            j += 1
        end
    end
    x
end

"""
[1, 4, 5], [2, 4, 7] => [1, 5]
"""
function remove(ind1::Vector{Int}, ind2::Vector{Int})
    x = Int[]
    for i in ind1
        if i ∉ ind2
            push!(x, i)
        end
    end
    x
end

"""
[1, 4, 5], [2, 4, 7] => [1, 3]
"""
function shiftremove(ind1::Vector{Int}, ind2::Vector{Int})
    x = Int[]
    for i in ind1
        if i ∉ ind2
            counter = 0
            for i2 in ind2
                if i2 < i
                    counter += 1
                else
                    break
                end
            end
            push!(x, i-counter)
        end
    end
    x
end

function reducedindices(I_::Vector{Int}, I::Vector{Int})
    N = length(I_)
    x = Vector{Int}(undef, N)
    for n in 1:N
        x[n] = findfirst(isequal(I_[n]), I)
    end
    x
end

function reducedindices!(I_::Vector{Int}, I::Vector{Int})
    for n in 1:length(I_)
        I_[n] = findfirst(isequal(I_[n]), I)
    end
end

"""
Check if all indices are unique and smaller than or equal to imax.
"""
function check_indices(imax::Int, indices::Vector{Int})
    N = length(indices)
    for n=1:N
        i = indices[n]
        @assert 0 < i <= imax
        for m in n+1:N
            @assert i != indices[m]
        end
    end
end

"""
Check if the indices are sorted, unique and smaller than or equal to imax.
"""
function check_sortedindices(imax::Int, indices::Vector{Int})
    N = length(indices)
    if N == 0
        return nothing
    end
    i_ = indices[1]
    @assert 0 < i_ <= imax
    for i in indices[2:end]
        @assert 0 < i <= imax
        @assert i > i_
    end
end

"""
    check_embed_indices(indices::Array)

Determine whether a collection of indices, written as a list of (integers or lists of integers) is unique.
This assures that the embedded operators are in non-overlapping subspaces.
"""
function check_embed_indices(indices::Array)
    # Check whether `indices` is empty.
    (length(indices) == 0) ? (return true) : nothing

    err_str = "Variable `indices` comes in an unexpected form. Expecting `Array{Union{Int, Array{Int, 1}}, 1}`"
    @assert mapreduce(x -> typeof(x)<:Array || typeof(x)<:Int, &, indices) err_str

    isunique = true
    # Check that no sub-list contains duplicates.
    for i in filter(x -> typeof(x) <: Array, indices)
        isunique &= (length(Set(i)) == length(i))
    end
    # Check that there are no duplicates across `indices`
    for (idx, i) in enumerate(indices[1:end-1])
        for j in indices[idx+1:end]
            isunique &= (length(Base.intersect(i, j)) == 0)
        end
    end
    return isunique
end


text:
# previous
previous_configuration(state::State) = (state.x1, state.q1)
previous_configuration_velocity(state::State) = (state.x1, state.v15, state.q1, state.ϕ15)

# current
current_position(state::State) = state.x2
current_orientation(state::State) = state.q2
current_configuration(state::State) = (current_position(state), current_orientation(state))
current_velocity(state::State) = (state.vsol[2], state.ϕsol[2])
current_configuration_velocity(state::State) = (current_position(state), state.vsol[2], current_orientation(state), state.ϕsol[2])
initial_configuration_velocity(state::State) = (current_position(state), state.v15, current_orientation(state), state.ϕ15)

# next
next_position(x2::SVector{3,T}, v25::SVector{3,T}, timestep::T) where T = x2 + v25 * timestep
next_orientation(q2::Quaternion{T}, ϕ25::SVector{3,T}, timestep::T) where T = q2 * quaternion_map(ϕ25, timestep) * timestep / 2
next_position(state::State, timestep) = next_position(state.x2, state.vsol[2], timestep)
next_orientation(state::State, timestep) = next_orientation(state.q2, state.ϕsol[2], timestep)
next_configuration(state::State, timestep) = (next_position(state, timestep), next_orientation(state, timestep))
next_configuration_velocity(state::State, timestep) = (next_position(state, timestep), state.vsol[2], next_orientation(state, timestep), state.ϕsol[2])

# angular velocity
function angular_velocity(q1::Quaternion, q2::Quaternion, timestep)
    2.0 / timestep  * Vmat() * Lᵀmat(q1) * vector(q2)
end

function ∂angular_velocity∂q1(q1::Quaternion, q2::Quaternion, timestep)
    2.0 / timestep  * Vmat() * Rmat(q2) * Tmat()
end

function ∂angular_velocity∂q2(q1::Quaternion, q2::Quaternion, timestep)
    2.0 / timestep  * Vmat() * Lᵀmat(q1)
end

# Jacobians
function integrator_jacobian_velocity(q2::Quaternion{T}, ϕ25::SVector{3,T}, timestep::T) where T
    V = [linear_integrator_jacobian_velocity(timestep) szeros(T,3,3)]
    Ω = [szeros(T,4,3) rotational_integrator_jacobian_velocity(q2, ϕ25, timestep)]
    return [V; Ω] # 7x6
end

function integrator_jacobian_configuration(q2::Quaternion{T}, ϕ25::SVector{3,T}, timestep::T;
    attjac::Bool=true) where T

    Z = attjac ? szeros(T,3,3) : szeros(T,3,4)
    X = [linear_integrator_jacobian_position() Z]
    Q = [szeros(T,4,3) rotational_integrator_jacobian_orientation(q2, ϕ25, timestep; attjac=attjac)]
    return [X; Q] # 7x6 or 7x7
end

function linear_integrator_jacobian_position(; T=Float64)
    return SMatrix{3,3,T,9}(Diagonal(sones(T,3)))
end

function linear_integrator_jacobian_velocity(timestep::T) where T
    return timestep * SMatrix{3,3,T,9}(Diagonal(sones(T,3)))
end

function rotational_integrator_jacobian_orientation(q2::Quaternion{T}, ϕ25::SVector{3,T}, timestep::T;
    attjac::Bool = true) where T
    M = Rmat(quaternion_map(ϕ25, timestep) * timestep / 2)
    attjac && (M *= LVᵀmat(q2))
    return M
end

function rotational_integrator_jacobian_velocity(q2::Quaternion{T}, ϕ25::SVector{3,T}, timestep::T) where T
    return Lmat(q2) * quaternion_map_jacobian(ϕ25, timestep) * timestep / 2
end


text:
include("../src/mcmc.jl")
include("../src/truth.jl")
include("../src/eki.jl")
include("../src/transforms.jl")

using Plots
using Cubature
using Plots.PlotMeasures
using LinearAlgebra
using ColorSchemes
using .MCMC
using .EKI
using Statistics
using Distributions
using JLD2
using Random
using StatsPlots, KernelDensity
using Interpolations

"""
Plotting script for the Parameter Ensemble and the GP,

produces two plots,

Requires:
--------
- MCMC object
- (if eki_opt_flag = True) EKP Object
- truthobject
- (transformations and inverse tranformations real -> prior space (if required)
- units transform to add units (if required)
A KNOWN ISSUE (in certain settings)
-----------------------------------
The contour calculates an integral of g(x) = max(0, f(x)>C ) over a box x in D using an integrator.
If g(x) is if small support relative to the size of D the adaptive refinement of the integrator can miss it,
temporary solve -> reduce size of D. Otherwise one needs to adjust/replace the blackbox integrator
"""
######################################################
function main()

    Random.seed!(100)

    @load "truth.jld2" truthobj
    @load "mcmc.jld2" mcmcobj

    ekp_opt_flag = false
    label_flag = false

    if ekp_opt_flag:
        @load "ekp.jld2" ekpobj
    end

        #create an Array of the MCMC stored
    homedir=split(pwd(),"/test")[1]*"/"
    outdir=homedir*"output/"

    hours=3600.0

    true_params_raw = [0,1]
    true_unitless_params_raw = remove_units_transform(true_params_raw)
    true_params = transform_real_to_prior(true_unitless_params_raw)

    #load the ensemble optimized params
    if ekp_opt_flag
        @load datdir*"eki.jld2" ekiobj

        U=ekpobj.u[end]
        ekp_params = mean(U,dims=1)
        ekp_params_raw = transform_prior_to_real(ekp_params)
        ekp_params_raw = apply_units_transform(ekp_params_raw)
    end

    #load MCMC results get prior
    posterior_samples=mcmc.posterior

    #provide contours? (NB exclude min and max value of posterior - these will be included)
    # if true, then put in the contour values
    # if false then put in desired percentiles and the code will try to find contours which bound the desired percentiles
    contour_provided=false

    if contour_provided #provide contour values
        input_contours=[]
    else #provide percentiles
        integral_target = [0.5,0.75,0.99]
    end

    #plots
    gr(size=(500,500))
    Plots.scalefontsizes(1.5)

    #Box size
    param1bd = [-1.8,-0.3]
    param2bd = [6,10]

    xmin = [param1bd[1],param2bd[1]]
    xmax = [param1bd[2],param2bd[2]]

    #Kernel density estimator
    postd=kde(posterior_samples, boundary=((param1bd[1],param1bd[2]),(param2bd[1],param2bd[2])) )
    interp_postd = InterpKDE(postd)
    min_cval = 0.0
    max_cval = maximum(maximum(postd.density))
    println("contour upper bound: ", max_cval)

    # Stage 1
    # Find the contours if required

    if contour_provided
        contour_values = input_contours
    else
        #
        # if the code hangs then either due to small tolerances (where hcubature will hang, or
        # due to domain being too large, so initial h-grid misses the posterior distribution
        # this is likely to happen when max(post-c,0) for c large and the distribution becomes more spiked

        function posterior_mass(x,pd)
            return pdf(pd,x...)
        end

        target_tol = 1e-3
        int_tol = target_tol/10
        max_evals=Int64(0) #default
        (tot_integral, err) = hcubature(x -> posterior_mass(x,interp_postd), xmin, xmax, abstol=int_tol, maxevals=max_evals) #should be 1 but may be less
        println("full density value, error:(", tot_integral, " ", err,")")
        #contour finder
        println("Finding contours:")
        flush(stdout)
        contour_values = zeros(size(integral_target))
        for (i,target) in enumerate(integral_target)
            count = 0
            cval = max_cval / 2

            cval_tmpmin = min_cval
            cval_tmpmax = max_cval
            integral_tmpmin = 0
            integral_tmpmax = 1
            #calc integral of midpoint firstly
            (integral,err) = hcubature(x-> max(posterior_mass(x,interp_postd) - cval, 0), xmin, xmax, abstol=int_tol, maxevals=max_evals)
            integral = integral / tot_integral

            while abs(integral - target)>target_tol
                # update the window
                if integral > target
                    cval_tmpmin = cval # increase the tmp minimum to the value
                    integral_tmpmax = integral # record the new max integral
                elseif integral < target
                    cval_tmpmax = cval # decrease the tmp maximum to the value
                    integral_tmpmin = integral #record the new min integral
                end

                # update linearly
                stepsize = (target - integral_tmpmin) / (integral_tmpmax - integral_tmpmin)
                cval = cval_tmpmin + (1-stepsize)*(cval_tmpmax - cval_tmpmin)

                #recalculate integral
                (integral, err) = hcubature(x -> max(posterior_mass(x,interp_postd) - cval,0), xmin, xmax, abstol=int_tol, maxevals=max_evals)
                integral = integral /tot_integral
                count = count+1
                println("Target ", target, ", current integral, error (", integral, " ", err,") after ", count, " iterations. Contour value = ", cval)
                flush(stdout)
                @assert count < 100
            end

            println("Target ", target, " achieved with integral ", integral, " after ", count, " iterations. Contour value = ", cval)
            flush(stdout)
            contour_values[i] = cval

        end
    end

    #apply to low are upper bounds
    param1bd_raw,param2bd_raw = transform_prior_to_real.([param1bd,param2bd])
    param1bd_raw,param2bd_raw = apply_units_transform.([param1bd_raw,param2bd_raw])

    c_values = sort([contour_values... , max_cval, 0.0])

    #plot in transformed coordinates
    plot(postd,
         seriestype=:contour,
         levels=c_values,
         seriescolor=cgrad(:tempo),
         #colorbar=:right,
         fill=true,
         dpi=300,
         legend=false,
         grid=false,
         xlims=param1bd,
         ylims=param2bd,
         framestyle=:box,
         left_margin=50px,
         bottom_margin=50px)

    plot!([true_params[1]],[true_params[2]], markercolor=:blue, markershape=:circle,msw=0,markersize=5)#need the extra [ ]
    if ekp_opt_flag
        plot!([ekp_params[1]],[ekp_params[2]], markercolor=:red, markershape=:cross,markersize=8)#need the extra [ ]
    end
    if label_flag
        xlabel!("transformed param1")
        ylabel!("transformed param2")
    end
    savefig(outdir*"posterior_density_transformed.pdf")

    #plot in untransformed coordinates
    posterior_samples_raw = zeros(size(posterior_samples))
    posterior_samples_raw[:,1] = transform_prior_to_real.(posterior_samples)
    posterior_samples_raw[:,2] = apply_units_transform.(posterior_samples_raw)

    postd_raw = kde(posterior_samples_raw, boundary = ((param1bd_raw[1],param1bd_raw[2]),(param2bd_raw[1],param2bd_raw[2])) )
    max_cval_raw = maximum(maximum(postd_raw.density))
    c_values_raw = sort([contour_values... , max_cval_raw, 0.0])
    plot(postd_raw,
         seriestype=:contour,
         levels=c_values_raw,
         seriescolor=cgrad(:tempo),
         #colorbar=:right,
         fill=true,
         dpi=300,
         grid=false,
         legend=false,
         xlims=param1bd_raw,
         ylims=param2bd_raw,
         framestyle=:box,
         left_margin=50px,
         bottom_margin=50px)

    plot!([true_params_raw[1]],[true_params_raw[2]], markercolor=:blue, markershape=:circle,msw=0,markersize=5)#need the extra [ ]
    if ekp_opt_flag
        plot!([ekp_params_raw[1]],[ekp_params_raw[2]], markercolor=:red, markershape=:cross,markersize=8)#need the extra [ ]
    end
    if label_flag
    xlabel!("untransformed param1")
    ylabel!("untransformed param2")
    end
    savefig(outdir*"posterior_density_untransformed.pdf")

end

main()


text:
using Test

@testset "Package SSGE" begin
    include("SSGE.jl")
end


text:
using Documenter, FixedPolynomials

makedocs(
    format = :html,
    sitename = "FixedPolynomials.jl",
    pages = [
        "Introduction" => "index.md",
        "Polynomial" => "reference.md",
        "Fast Evaluation" => "performance.md"
        ]
)

deploydocs(
    repo   = "github.com/saschatimme/FixedPolynomials.jl.git",
    target = "build",
    julia = "0.6",
    osname = "linux",
    deps   = nothing,
    make   = nothing
)


text:
@testset "Collision operator" begin
    tolerances = Dict(Float16 => 1e-4, Float32 => 1e-5, Float64 => 1e-7, Real => 1e-7)
    @testset "One spatial dimension" begin
        @testset "Collision -forcing -temp" begin
            N = 8
            typelist = [Float16 Float32 Float64]
            for type in typelist
                testmom = JuSwalbe.Macroquant(height = ones(type, N), velocity=zeros(type, N), pressure=zeros(type, N), energy=zeros(type, N))
                testforce = JuSwalbe.Forces(slip = zeros(type, N), h∇p=zeros(type, N), bathymetry=zeros(type, N), thermal=zeros(type, N))
                testtempdist = JuSwalbe.DistributionD1Q3(f0 = zeros(type, N), f1 = zeros(type, N), f2 = zeros(type, N))
                input = JuSwalbe.Inputconstants(τ = type(1.0), gravity=type(0.0))
                @test isa(testmom, JuSwalbe.Macroquant{Vector{type}, Vector{type}})
                @test isa(testtempdist, JuSwalbe.DistributionD1Q3{Vector{type}})
                @test isa(input, JuSwalbe.Inputconstants)
                testequi = calc_equilibrium_distribution(testmom, gravity=type(input.gravity))
                testdist = collisionBGK(testmom, testforce, testtempdist, input)
                @test isa(testdist, JuSwalbe.DistributionD1Q3{Vector{type}})
                for i in 1:N
                    @test testdist.f0[i] -  testequi.f0[i] ≈ type(0.0) atol=tolerances[type]
                    @test testdist.f1[i] -  testequi.f1[i] ≈ type(0.0) atol=tolerances[type]
                    @test testdist.f2[i] -  testequi.f2[i] ≈ type(0.0) atol=tolerances[type]
                end
            end
        end
        @testset "Collision +forcing -temp" begin
            N = 8
            typelist = [Float16 Float32 Float64]
            for type in typelist
                testmom = JuSwalbe.Macroquant(height = ones(type, N), velocity=zeros(type, N), pressure=zeros(type, N), energy=zeros(type, N))
                testforce = JuSwalbe.Forces(slip = ones(type, N), h∇p=ones(type, N), bathymetry=fill(type(0.1), N), thermal=zeros(type, N))
                testtempdist = JuSwalbe.DistributionD1Q3(f0 = zeros(type, N), f1 = zeros(type, N), f2 = zeros(type, N))
                input = JuSwalbe.Inputconstants(τ = type(1.0), gravity=type(0.1))
                @test isa(testmom, JuSwalbe.Macroquant{Vector{type}, Vector{type}})
                @test isa(testtempdist, JuSwalbe.DistributionD1Q3{Vector{type}})
                @test isa(input, JuSwalbe.Inputconstants)
                testequi = calc_equilibrium_distribution(testmom, gravity=type(input.gravity))
                testdist = collisionBGK(testmom, testforce, testtempdist, input)
                @test isa(testdist, JuSwalbe.DistributionD1Q3{Vector{type}})
                for i in 1:N
                    @test testdist.f0[i] -  testequi.f0[i] ≈ type(0.0) atol=tolerances[type]
                    @test (testdist.f1[i] -  testequi.f1[i] - type(3/6) * type(2.1)) ≈ type(0.0) atol=tolerances[type]
                    @test (testdist.f2[i] -  testequi.f2[i] + type(3/6) * type(2.1)) ≈ type(0.0) atol=tolerances[type]
                end
            end
        end
        @testset "Collision +forcing +temp" begin
            N = 8
            typelist = [Float16 Float32 Float64]
            for type in typelist
                testmom = JuSwalbe.Macroquant(height = ones(type, N), velocity=zeros(type, N), pressure=zeros(type, N), energy=zeros(type, N))
                testforce = JuSwalbe.Forces(slip = ones(type, N), h∇p=ones(type, N), bathymetry=fill(type(0.1), N), thermal=zeros(type, N))
                testtempdist = JuSwalbe.DistributionD1Q3(f0 = fill(type(0.1), N), f1 = fill(type(0.2), N), f2 = fill(type(-0.2), N))
                input = JuSwalbe.Inputconstants(τ = type(0.75), gravity=type(0.1))
                @test isa(testmom, JuSwalbe.Macroquant{Vector{type}, Vector{type}})
                @test isa(testtempdist, JuSwalbe.DistributionD1Q3{Vector{type}})
                @test isa(input, JuSwalbe.Inputconstants)
                testequi = calc_equilibrium_distribution(testmom, gravity=type(input.gravity))
                testdist = collisionBGK(testmom, testforce, testtempdist, input)
                @test isa(testdist, JuSwalbe.DistributionD1Q3{Vector{type}})
                for i in 1:N
                    @test testdist.f0[i] - (type(1-1/input.τ)*testtempdist.f0[i] + 1/input.τ * testequi.f0[i]) ≈ type(0.0) atol=tolerances[type]
                    @test (testdist.f1[i] -  (type(1-1/input.τ)*testtempdist.f1[i] + 1/input.τ * testequi.f1[i] + type(3/6) * type(2.1))) ≈ type(0.0) atol=tolerances[type]
                    @test (testdist.f2[i] -  (type(1-1/input.τ)*testtempdist.f2[i] + 1/input.τ * testequi.f2[i] - type(3/6) * type(2.1))) ≈ type(0.0) atol=tolerances[type]
                end
            end
        end
        @testset "Collision -forcing +temp" begin
            N = 8
            typelist = [Float16 Float32 Float64]
            for type in typelist
                testmom = JuSwalbe.Macroquant(height = ones(type, N), velocity=zeros(type, N), pressure=zeros(type, N), energy=zeros(type, N))
                testforce = JuSwalbe.Forces(slip=zeros(type, N), h∇p=zeros(type, N), bathymetry=zeros(type, N), thermal=zeros(type, N))
                testtempdist = JuSwalbe.DistributionD1Q3(f0 = fill(type(0.1), N), f1 = fill(type(0.2), N), f2 = fill(type(-0.2), N))
                input = JuSwalbe.Inputconstants(τ = type(0.75), gravity=type(0.0))
                @test isa(testmom, JuSwalbe.Macroquant{Vector{type}, Vector{type}})
                @test isa(testtempdist, JuSwalbe.DistributionD1Q3{Vector{type}})
                @test isa(input, JuSwalbe.Inputconstants)
                testequi = calc_equilibrium_distribution(testmom, gravity=type(input.gravity))
                testdist = collisionBGK(testmom, testforce, testtempdist, input)
                @test isa(testdist, JuSwalbe.DistributionD1Q3{Vector{type}})
                for i in 1:N
                    @test testdist.f0[i] - (type(1-1/input.τ)*testtempdist.f0[i] + 1/input.τ * testequi.f0[i]) ≈ type(0.0) atol=tolerances[type]
                    @test (testdist.f1[i] -  (type(1-1/input.τ)*testtempdist.f1[i] + 1/input.τ * testequi.f1[i])) ≈ type(0.0) atol=tolerances[type]
                    @test (testdist.f2[i] -  (type(1-1/input.τ)*testtempdist.f2[i] + 1/input.τ * testequi.f2[i])) ≈ type(0.0) atol=tolerances[type]
                end
            end
        end
    end
    @testset "Two spatial dimensions" begin
        @testset "Collision -forcing -temp" begin
            N, M = 8, 6
            typelist = [Float16 Float32 Float64]
            for type in typelist
                testvel = JuSwalbe.Twovector(x=fill(type(0.1),(N,M)), y=fill(type(0.1),(N,M)))
                zerovec = JuSwalbe.Twovector(x=zeros(type, (N,M)), y=zeros(type, (N,M)))
                testmom = JuSwalbe.Macroquant(height = ones(type, (N,M)), velocity=zerovec, pressure=zeros(type, (N,M)), energy=zeros(type, (N,M)))
                testforce = JuSwalbe.Forces(slip = zerovec, h∇p=zerovec, bathymetry=zerovec, thermal=zerovec)
                testtempdist = JuSwalbe.DistributionD2Q9(f0 = zeros(type, (N,M)),
                                                         f1 = zeros(type, (N,M)),
                                                         f2 = zeros(type, (N,M)),
                                                         f3 = zeros(type, (N,M)),
                                                         f4 = zeros(type, (N,M)),
                                                         f5 = zeros(type, (N,M)),
                                                         f6 = zeros(type, (N,M)),
                                                         f7 = zeros(type, (N,M)),
                                                         f8 = zeros(type, (N,M)))

                input = JuSwalbe.Inputconstants(τ = type(1.0), gravity=type(0.0))
                @test isa(testmom, JuSwalbe.Macroquant{Matrix{type}, JuSwalbe.Twovector{Matrix{type}}})
                @test isa(testtempdist, JuSwalbe.DistributionD2Q9{Matrix{type}})
                @test isa(input, JuSwalbe.Inputconstants)
                testequi = calc_equilibrium_distribution(testmom, gravity=type(input.gravity))
                testdist = collisionBGK(testmom, testforce, testtempdist, input)
                @test isa(testdist, JuSwalbe.DistributionD2Q9{Matrix{type}})
                testdict = struct2dict(testdist)
                equidict = struct2dict(testequi)
                for key in keys(testdict)
                    for i in 1:N, j in 1:M
                        @test testdict[key][i,j] -  equidict[key][i,j] ≈ type(0.0) atol=tolerances[type]
                    end
                end
            end
        end
        @testset "Collision +forcing -temp" begin
            N, M = 5, 5
            typelist = [Float16 Float32 Float64]
            for type in typelist
                testvel = JuSwalbe.Twovector(x=fill(type(0.1),(N,M)), y=fill(type(0.1),(N,M)))
                zerovec = JuSwalbe.Twovector(x=zeros(type, (N,M)), y=zeros(type, (N,M)))
                testmom = JuSwalbe.Macroquant(height = ones(type, (N,M)), velocity=zerovec, pressure=zeros(type, (N,M)), energy=zeros(type, (N,M)))
                testforce = JuSwalbe.Forces(slip = testvel, h∇p=zerovec, bathymetry=zerovec, thermal=zerovec)
                testtempdist = JuSwalbe.DistributionD2Q9(f0 = zeros(type, (N,M)), f1 = zeros(type, (N,M)), f2 = zeros(type, (N,M)),
                                                         f3 = zeros(type, (N,M)), f4 = zeros(type, (N,M)), f5 = zeros(type, (N,M)),
                                                         f6 = zeros(type, (N,M)), f7 = zeros(type, (N,M)), f8 = zeros(type, (N,M)))

                input = JuSwalbe.Inputconstants(τ = type(1.0), gravity=type(0.0))
                @test isa(testmom, JuSwalbe.Macroquant{Matrix{type}, JuSwalbe.Twovector{Matrix{type}}})
                @test isa(testtempdist, JuSwalbe.DistributionD2Q9{Matrix{type}})
                @test isa(testforce, JuSwalbe.Forces{JuSwalbe.Twovector{Matrix{type}}})
                @test isa(input, JuSwalbe.Inputconstants)
                testequi = calc_equilibrium_distribution(testmom, gravity=type(input.gravity))
                testdist = collisionBGK(testmom, testforce, testtempdist, input)
                @test isa(testdist, JuSwalbe.DistributionD2Q9{Matrix{type}})
                testdict = struct2dict(testdist)
                equidict = struct2dict(testequi)
                # Solution is just the force: 3w_i/c_i^2 * c_i * F
                solutiondict = Dict(:f0 => type(0), :f1 => type(1/30), :f2 => type(1/30), :f3 => type(-1/30), :f4 => type(-1/30),
                                    :f5 => type(1/120), :f6 => type(0), :f7 => type(-1/120), :f8 => type(0))
                for key in keys(testdict)
                    for i in 1:N, j in 1:M
                        @test testdict[key][i,j] -  (equidict[key][i,j] + solutiondict[key])  ≈ type(0.0) atol=tolerances[type]
                    end
                end
            end
        end
        @testset "Collision +forcing +temp" begin
            N, M = 5, 5
            typelist = [Float16 Float32 Float64]
            for type in typelist
                testvel = JuSwalbe.Twovector(x=fill(type(0.1),(N,M)), y=fill(type(0.1),(N,M)))
                zerovec = JuSwalbe.Twovector(x=zeros(type, (N,M)), y=zeros(type, (N,M)))
                testmom = JuSwalbe.Macroquant(height = ones(type, (N,M)), velocity=zerovec, pressure=zeros(type, (N,M)), energy=zeros(type, (N,M)))
                testforce = JuSwalbe.Forces(slip = testvel, h∇p=zerovec, bathymetry=zerovec, thermal=zerovec)
                testtempdist = JuSwalbe.DistributionD2Q9(f0 = fill(type(0.1), (N,M)), f1 = fill(type(0.2), (N,M)), f2 = fill(type(-0.2), (N,M)),
                                                         f3 = zeros(type, (N,M)), f4 = zeros(type, (N,M)), f5 = zeros(type, (N,M)),
                                                         f6 = zeros(type, (N,M)), f7 = zeros(type, (N,M)), f8 = zeros(type, (N,M)))

                input = JuSwalbe.Inputconstants(τ = type(0.75), gravity=type(0.0))
                @test isa(testmom, JuSwalbe.Macroquant{Matrix{type}, JuSwalbe.Twovector{Matrix{type}}})
                @test isa(testtempdist, JuSwalbe.DistributionD2Q9{Matrix{type}})
                @test isa(testforce, JuSwalbe.Forces{JuSwalbe.Twovector{Matrix{type}}})
                @test isa(input, JuSwalbe.Inputconstants)
                testequi = calc_equilibrium_distribution(testmom, gravity=type(input.gravity))
                testdist = collisionBGK(testmom, testforce, testtempdist, input)
                @test isa(testdist, JuSwalbe.DistributionD2Q9{Matrix{type}})
                testdict = struct2dict(testdist)
                equidict = struct2dict(testequi)
                tempdict = struct2dict(testtempdist)
                soldict = Dict(:f0 => type(0), :f1 => type(1/30), :f2 => type(1/30), :f3 => type(-1/30), :f4 => type(-1/30),
                               :f5 => type(1/120), :f6 => type(0), :f7 => type(-1/120), :f8 => type(0))
                for key in keys(testdict)
                    for i in 1:N, j in 1:M
                        @test testdict[key][i,j] - (type(1-1/input.τ)*tempdict[key][i,j] + 1/input.τ * equidict[key][i,j] + soldict[key]) ≈ type(0.0) atol=tolerances[type]
                    end
                end
            end
        end
        @testset "Collision -forcing +temp" begin
            N, M = 5, 5
            typelist = [Float16 Float32 Float64]
            for type in typelist
                testvel = JuSwalbe.Twovector(x=fill(type(0.1),(N,M)), y=fill(type(0.1),(N,M)))
                zerovec = JuSwalbe.Twovector(x=zeros(type, (N,M)), y=zeros(type, (N,M)))
                testmom = JuSwalbe.Macroquant(height = ones(type, (N,M)), velocity=zerovec, pressure=zeros(type, (N,M)), energy=zeros(type, (N,M)))
                testforce = JuSwalbe.Forces(slip = zerovec, h∇p=zerovec, bathymetry=zerovec, thermal=zerovec)
                testtempdist = JuSwalbe.DistributionD2Q9(f0 = fill(type(0.1), (N,M)), f1 = fill(type(0.2), (N,M)), f2 = fill(type(-0.2), (N,M)),
                                                        f3 = zeros(type, (N,M)), f4 = zeros(type, (N,M)), f5 = zeros(type, (N,M)),
                                                        f6 = zeros(type, (N,M)), f7 = zeros(type, (N,M)), f8 = zeros(type, (N,M)))

                input = JuSwalbe.Inputconstants(τ = type(0.75), gravity=type(0.0))
                @test isa(testmom, JuSwalbe.Macroquant{Matrix{type}, JuSwalbe.Twovector{Matrix{type}}})
                @test isa(testtempdist, JuSwalbe.DistributionD2Q9{Matrix{type}})
                @test isa(testforce, JuSwalbe.Forces{JuSwalbe.Twovector{Matrix{type}}})
                @test isa(input, JuSwalbe.Inputconstants)
                testequi = calc_equilibrium_distribution(testmom, gravity=type(input.gravity))
                testdist = collisionBGK(testmom, testforce, testtempdist, input)
                @test isa(testdist, JuSwalbe.DistributionD2Q9{Matrix{type}})
                testdict = struct2dict(testdist)
                equidict = struct2dict(testequi)
                tempdict = struct2dict(testtempdist)

                for key in keys(testdict)
                    for i in 1:N, j in 1:M
                        @test testdict[key][i,j] - (type(1-1/input.τ)*tempdict[key][i,j] + 1/input.τ * equidict[key][i,j]) ≈ type(0.0) atol=tolerances[type]
                    end
                end
            end
        end
    end
end

@testset "Streaming function" begin
    @testset "One dimension periodic" begin
        N = 8
        typelist = [Float16 Float32 Float64]
        for type in typelist
            testtempdist = JuSwalbe.DistributionD1Q3(f0 = fill(type(0.1), N), f1 = fill(type(0.2), N), f2 = fill(type(-0.2), N))
            testtempdist.f1[3] = type(-0.1)
            testtempdist.f2[3] = type(0.1)

            @test isa(testtempdist, JuSwalbe.DistributionD1Q3{Vector{type}})
            streamdistperiodic!(testtempdist)
            @test isa(testtempdist, JuSwalbe.DistributionD1Q3{Vector{type}})
            @test length(testtempdist.f0) == 8
            @test length(testtempdist.f1) == 8
            @test length(testtempdist.f2) == 8
            for i in testtempdist.f0
                @test i == type(0.1)
            end
            @test testtempdist.f1[4] == type(-0.1)
            @test testtempdist.f1[3] == type(0.2)
            @test testtempdist.f1[5] == type(0.2)
            @test testtempdist.f2[2] == type(0.1)
            @test testtempdist.f2[3] == type(-0.2)
            @test testtempdist.f2[4] == type(-0.2)

        end
    end
    @testset "Two dimensions periodic" begin
        N, M = 8, 6
        typelist = [Float16 Float32 Float64]
        for type in typelist
            testtempdist = JuSwalbe.DistributionD2Q9(f0 = fill(type(0.1), (N,M)), f1 = fill(type(0.2), (N,M)), f2 = fill(type(-0.2), (N,M)),
                                                     f3 = fill(type(0.3), (N,M)), f4 = fill(type(0.1), (N,M)), f5 = fill(type(-0.1), (N,M)),
                                                     f6 = fill(type(0.6), (N,M)), f7 = fill(type(0.5), (N,M)), f8 = fill(type(-0.5), (N,M)))
            testtempdist.f1[3,4] = type(1)
            testtempdist.f2[3,3] = type(1)
            testtempdist.f3[3,4] = type(1)
            testtempdist.f4[3,3] = type(1)
            testtempdist.f5[3,4] = type(1)
            testtempdist.f6[3,3] = type(1)
            testtempdist.f7[3,4] = type(1)
            testtempdist.f8[3,3] = type(1)

            @test isa(testtempdist, JuSwalbe.DistributionD2Q9{Matrix{type}})
            testdist = streamdistperiodic!(testtempdist)
            @test isa(testtempdist, JuSwalbe.DistributionD2Q9{Matrix{type}})
            @test size(testdist.f0) == (N,M)
            @test size(testdist.f1) == (N,M)
            @test size(testdist.f2) == (N,M)
            @test size(testdist.f3) == (N,M)
            @test size(testdist.f4) == (N,M)
            @test size(testdist.f5) == (N,M)
            @test size(testdist.f6) == (N,M)
            @test size(testdist.f7) == (N,M)
            @test size(testdist.f8) == (N,M)
            for i in testdist.f0
                @test i == type(0.1)
            end
            for i in 1:N, j in 1:M
                if i == 4 && j == 4
                    @test testdist.f1[i,j] == type(1)
                else
                    @test testdist.f1[i,j] == type(0.2)
                end
            end
            for i in 1:N, j in 1:M
                if i == 3 && j == 4
                    @test testdist.f2[i,j] == type(1)
                else
                    @test testdist.f2[i,j] == type(-0.2)
                end
            end
            for i in 1:N, j in 1:M
                if i == 2 && j == 4
                    @test testdist.f3[i,j] == type(1)
                else
                    @test testdist.f3[i,j] == type(0.3)
                end
            end
            for i in 1:N, j in 1:M
                if i == 3 && j == 2
                    @test testdist.f4[i,j] == type(1)
                else
                    @test testdist.f4[i,j] == type(0.1)
                end
            end
            for i in 1:N, j in 1:M
                if i == 4 && j == 5
                    @test testdist.f5[i,j] == type(1)
                else
                    @test testdist.f5[i,j] == type(-0.1)
                end
            end
            for i in 1:N, j in 1:M
                if i == 2 && j == 4
                    @test testdist.f6[i,j] == type(1)
                else
                    @test testdist.f6[i,j] == type(0.6)
                end
            end
            for i in 1:N, j in 1:M
                if i == 2 && j == 3
                    @test testdist.f7[i,j] == type(1)
                else
                    @test testdist.f7[i,j] == type(0.5)
                end
            end
            for i in 1:N, j in 1:M
                if i == 4 && j == 2
                    @test testdist.f8[i,j] == type(1)
                else
                    @test testdist.f8[i,j] == type(-0.5)
                end
            end
        end
    end
end

max_stars_repo_name: Zitzeronion/JuSwalbe

max_issues_repo_name: Zitzeronion/JuSwalbe


text:
module Torch

using PyCall

using ChainRulesCore
using DLPack
using Functors: @functor
using Adapt

using ..PyCallChainRules: PyAdaptor, fmap

const inspect = PyNULL()
const torch = PyNULL()
const functorch = PyNULL()
const dlpack = PyNULL()
const ispysetup = Ref{Bool}(false)

pyto_dlpack(x) = @pycall dlpack.to_dlpack(x)::PyObject
pyfrom_dlpack(x) = @pycall dlpack.from_dlpack(x)::PyObject

struct TorchModuleWrapper
    torch_stateless_module::PyObject
    dtype::PyObject
    params::Tuple
    buffers::Tuple
end

@functor TorchModuleWrapper (params,)

Base.show(io::IO, f::TorchModuleWrapper) = print(io, f.torch_stateless_module, " ", f.dtype, " ", " params size=", size.(f.params))

Base.length(f::TorchModuleWrapper) = length(f.params)
Base.iterate(f::TorchModuleWrapper) = iterate(f.params)
Base.iterate(f::TorchModuleWrapper, state) = iterate(f.params, state)

function TorchModuleWrapper(torch_module)
    pybuiltin("isinstance")(torch_module, torch.nn.Module) || error("Not a torch.nn.Module")
    funmod, params, buffers = functorch.make_functional_with_buffers(torch_module)
    dtype = params[1].dtype
    jlparams = map(params) do x
        DLPack.wrap(x, pyto_dlpack)
    end
    return TorchModuleWrapper(funmod, dtype, jlparams, buffers)
end

function (wrap::TorchModuleWrapper)(args...; kwargs...)
    params = wrap.params
    out = wrap.torch_stateless_module(Tuple(map(x -> DLPack.share(x, PyObject, pyfrom_dlpack).requires_grad_(true), params)),
        wrap.buffers, fmap(x -> DLPack.share(x, PyObject, pyfrom_dlpack), args)...; kwargs...)
    res = fmap(x->DLPack.wrap(x, pyto_dlpack), out)
    return res
end

function ChainRulesCore.rrule(wrap::TorchModuleWrapper, args...; kwargs...)
    T = typeof(first(wrap.params))
    params = wrap.params
    pyparams = Tuple(map(x -> DLPack.share(x, PyObject, pyfrom_dlpack).requires_grad_(true), params))
    pyargs = fmap(x -> DLPack.share(x, PyObject, pyfrom_dlpack).requires_grad_(true), args)

    torch_primal, torch_vjpfun = functorch.vjp(py"buffer_implicit"(wrap.torch_stateless_module, wrap.buffers), pyparams, pyargs...; kwargs...)
    project = ProjectTo(args)
    function TorchModuleWrapper_pullback(Δ)
        cΔ = fmap(x->Adapt.adapt(PyAdaptor{T}(), x), Δ)
        pycΔ = fmap(x->DLPack.share(x, PyObject, pyfrom_dlpack), cΔ)
        torch_tangent_vals = torch_vjpfun(pycΔ)
        jlparams_tangents = map(x -> DLPack.wrap(x, pyto_dlpack), torch_tangent_vals[1])
        args_tangents = project(fmap(x -> DLPack.wrap(x, pyto_dlpack), torch_tangent_vals[2:end]))
        return (Tangent{TorchModuleWrapper}(; torch_stateless_module = NoTangent(), dtype = NoTangent(), params = jlparams_tangents, buffers = NoTangent()), args_tangents...)
    end
    res = fmap(x->DLPack.wrap(x, pyto_dlpack), torch_primal)
    return res, TorchModuleWrapper_pullback
end

function __init__()
    try
        copy!(torch, pyimport("torch"))
        copy!(dlpack, pyimport("torch.utils.dlpack"))
        copy!(functorch, pyimport("functorch"))
        copy!(inspect, pyimport("inspect"))
        ispysetup[] = true
        py"""
        def buffer_implicit(fn, buffers):
            def newfn(params, *inputs, **kwargs):
                return fn(params, buffers, *inputs, **kwargs)

            return newfn
        """
    catch err
        @warn """PyCallChainRules.jl has failed to import torch and functorch from Python.
                 Please make sure these are installed.
        """
        @debug err
        ispysetup[] = false
        #rethrow(err)
    end
end

end


text:
using Test, ApproxInference

const VecI  = AbstractVector # Input  Vector
const VecIO = AbstractVector # In/Out Vector
const MatI  = AbstractMatrix # Input  Matrix
const MatIO = AbstractMatrix # In/Out Matrix

print_head(s::String) = println("\033[1m\033[32mTesting Task\033[0m \033[1m\033[33m$s\033[0m")
print_body(s::String) = println("             $s")

##############################
#   NIST Dataset: BoxBOD     #
##############################
@objective struct BoxBOD
    x::Vector{Float64}; y::Vector{Float64}

    BoxBOD() = new([1., 2., 3., 5., 7., 10.], [109., 149., 149., 191., 213., 224.])
end

function ApproxInference.get_residual!(r::VecIO, θ::VecI, f::BoxBOD)
    x, y = f.x, f.y
    @inbounds θ1 = θ[1]
    @inbounds θ2 = θ[2]
    @simd for i in eachindex(x)
        @inbounds r[i] = y[i] - θ1 * (1.0 - exp(-θ2 * x[i]))
    end
    return nothing
end

function ApproxInference.get_jacobian!(J::MatIO, θ::VecI, f::BoxBOD)
    x = f.x
    @inbounds θ1 = θ[1]
    @inbounds θ2 = θ[2]
    @inbounds for i in eachindex(x)
        tmp = exp(-θ2 * x[i])
        J[i,1] = 1.0 - tmp
        J[i,2] = θ1 * x[i] * tmp
    end
    return nothing
end

include("./test-BoxBOD-LM.jl")
include("./test-BoxBOD-VI.jl")

##############################
#   NIST Dataset: Eckerle4   #
##############################
@objective struct Eckerle4
    x::Vector{Float64}; y::Vector{Float64}

    Eckerle4() = new(
        # x vector
        [400.0, 405.0, 410.0, 415.0, 420.0, 425.0, 430.0,
         435.0, 436.5, 438.0, 439.5, 441.0, 442.5, 444.0,
         445.5, 447.0, 448.5, 450.0, 451.5, 453.0, 454.5,
         456.0, 457.5, 459.0, 460.5, 462.0, 463.5, 465.0,
         470.0, 475.0, 480.0, 485.0, 490.0, 495.0, 500.0],
        # y vector
        [1.575000e-04, 1.699000e-04, 2.350000e-04, 3.102000e-04, 4.917000e-04,
         8.710000e-04, 1.741800e-03, 4.640000e-03, 6.589500e-03, 9.730200e-03,
         1.490020e-02, 2.373100e-02, 4.016830e-02, 7.125590e-02, 1.264458e-01,
         2.073413e-01, 2.902366e-01, 3.445623e-01, 3.698049e-01, 3.668534e-01,
         3.106727e-01, 2.078154e-01, 1.164354e-01, 6.167640e-02, 3.372000e-02,
         1.940230e-02, 1.178310e-02, 7.435700e-03, 2.273200e-03, 8.800000e-04,
         4.579000e-04, 2.345000e-04, 1.586000e-04, 1.143000e-04, 7.100000e-05]
    )
end

function ApproxInference.get_residual!(r::VecIO, θ::VecI, f::Eckerle4)
    x, y = f.x, f.y
    @inbounds θ1 = θ[1]
    @inbounds θ2 = θ[2]
    @inbounds θ3 = θ[3]
    @simd for i in eachindex(x)
        @inbounds r[i] = y[i] - θ1 * exp(-0.5 * abs2((x[i] - θ3) / θ2)) / θ2
    end
    return nothing
end

function ApproxInference.get_jacobian!(J::MatIO, θ::VecI, f::Eckerle4)
    x = f.x
    @inbounds θ1 = θ[1]
    @inbounds θ2 = θ[2]
    @inbounds θ3 = θ[3]

    θ2² = abs2(θ2)
    θ2³ = θ2 * θ2 * θ2

    @inbounds for i in eachindex(x)
        tmp1 = abs2(x[i] - θ3) / θ2²
        tmp2 = exp(-0.5 * tmp1)
        J[i,1] = tmp2 / θ2
        J[i,2] = θ1 * tmp2 * (tmp1 - 1.0) / θ2²
        J[i,3] = θ1 * (x[i] - θ3) * tmp2 / θ2³
    end
    return nothing
end

include("./test-Eckerle4-LM.jl")
include("./test-Eckerle4-VI.jl")