text:
using AMDGPU
using CUDA
using ExaTronKernels
using KernelAbstractions
using LinearAlgebra
using Random
using Test

"""
Test ExaTron's internal routines written for GPU.

The current implementation assumes the following:
  - A thread block takes a matrix structure, (tx,ty), with
    n <= blockDim().x = blockDim().y <= 32 and blockDim().x is even.
  - Arguments passed on to these routines are assumed to be
    of size at least n. This is to prevent multiple thread
    divergence when we call a function with n_hat < n.
    Such a case occurs when we fix active variables.

We test the following routines, where [O] indicates if the routine
is checked if n < blockDim().x is OK.
  - dicf     [O][O]: this routine also tests dnsol and dtsol.
  - dicfs    [O][T]
  - dcauchy  [O][T]
  - dtrpcg   [O][T]
  - dprsrch  [O][T]
  - daxpy    [O][O]
  - dssyax   [O][O]: we do shuffle using blockDim().x.
  - dmid     [O][O]: we could use a single thread only to multiple divergences.
  - dgpstep  [O][O]
  - dbreakpt [O][O]: we use the existing ExaTron implementation as is.
  - dnrm2    [O][O]: we do shuffle using blockDim().x.
  - nrm2     [O][O]: we do shuffle using blockDim().x.
  - dcopy    [O][O]
  - ddot     [O][O]
  - dscal    [O][O]
  - dtrqsol  [O][O]
  - dspcg    [O][T]: we use a single thread to avoid multiple divergences.
  - dgpnorm  [O][O]
  - dtron    [O]
  - driver_kernel
"""

Random.seed!(0)
itermax = 10
n = 4
nblk = 4

if has_cuda_gpu()
    using CUDAKernels
    device = CUDADevice()
    AT = CuArray
elseif has_rocm_gpu()
    using ROCKernels
    device = ROCDevice()
    AT = ROCArray
else
    device = CPU()
    error("CPU KA implementation is currently broken for nested functions")
end

@testset "dicf" begin
@kernel function dicf_test(::Val{n}, d_in,
                    d_out) where {n}
    tx = @index(Local, Linear)
    bx = @index(Group, Linear)

    L = @localmem Float64 (n,n)
    for i in 1:n
        L[tx,i] = d_in[tx,i]
    end
    @synchronize

    # Test Cholesky factorization.
    ExaTronKernels.dicf(n,L,tx)
    if bx == 1
        for i in 1:n
            d_out[tx,i] = L[tx,i]
        end
    end
    @synchronize
end

    for i=1:itermax
        L = tril(rand(n,n))
        A = L*transpose(L)
        A .= tril(A) .+ (transpose(tril(A)) .- Diagonal(A))
        tron_A = ExaTronKernels.TronDenseMatrix{Array{Float64,2}}(n)
        tron_A.vals .= A

        d_in = AT{Float64,2}(undef, (n,n))
        d_out = AT{Float64,2}(undef, (n,n))
        copyto!(d_in, tron_A.vals)
        wait(dicf_test(device, n)(Val{n}(), d_in, d_out, ndrange=(n,nblk), dependencies=Event(device)))
        h_L = d_out |> Array

        tron_L = ExaTronKernels.TronDenseMatrix{Array{Float64,2}}(n)
        tron_L.vals .= tron_A.vals
        indr = zeros(Int, n)
        indf = zeros(n)
        list = zeros(n)
        w = zeros(n)
        ExaTronKernels.dicf(n, n^2, tron_L, 5, indr, indf, list, w)

        @test norm(tron_A.vals .- tril(h_L)*transpose(tril(h_L))) <= 1e-10
        @test norm(tril(h_L) .- transpose(triu(h_L))) <= 1e-10
        @test norm(tril(tron_L.vals) .- tril(h_L)) <= 1e-9
    end
end

@testset "dicfs" begin
    @kernel function dicfs_test(::Val{n}, alpha::Float64,
                        dA,
                        d_out) where {n}
        tx = @index(Local, Linear)
        bx = @index(Group, Linear)

        wa1 = @localmem Float64 (n,)
        wa2 = @localmem Float64 (n,)
        A = @localmem Float64 (n,n)
        L = @localmem Float64 (n,n)

        for i in 1:n
            A[tx,i] = dA[tx,i]
        end
        @synchronize

        ExaTronKernels.dicfs(n, alpha, A, L, wa1, wa2, tx)
        if bx <= 1
            for i in 1:n
                d_out[tx,i] = L[tx,i]
            end
        end
        @synchronize
    end

    for i=1:itermax
        L = tril(rand(n,n))
        A = L*transpose(L)
        A .= tril(A) .+ (transpose(tril(A)) .- Diagonal(A))
        tron_A = ExaTronKernels.TronDenseMatrix{Array{Float64,2}}(n)
        tron_L = ExaTronKernels.TronDenseMatrix{Array{Float64,2}}(n)
        tron_A.vals .= A

        dA = AT{Float64,2}(undef, (n,n))
        d_out = AT{Float64,2}(undef, (n,n))
        alpha = 1.0
        copyto!(dA, tron_A.vals)
        wait(dicfs_test(device, n)(Val{n}(),alpha,dA,d_out,ndrange=(n, nblk),dependencies=Event(device)))
        h_L = d_out |> Array
        iwa = zeros(Int, 3*n)
        wa1 = zeros(n)
        wa2 = zeros(n)
        ExaTronKernels.dicfs(n, n*n, tron_A, tron_L, 5, alpha, iwa, wa1, wa2)

        @test norm(tril(h_L) .- transpose(triu(h_L))) <= 1e-10
        @test norm(tril(tron_L.vals) .- tril(h_L)) <= 1e-9

        # Make it negative definite.
        for j=1:n
            tron_A.vals[j,j] = -tron_A.vals[j,j]
        end
        copyto!(dA, tron_A.vals)
        wait(dicfs_test(device, n)(Val{n}(),alpha,dA,d_out,ndrange=(n,nblk),dependencies=Event(device)))
        copyto!(h_L, d_out)
        ExaTronKernels.dicfs(n, n^2, tron_A, tron_L, 5, alpha, iwa, wa1, wa2)

        @test norm(tril(h_L) .- transpose(triu(h_L))) <= 1e-10
        @test norm(tril(tron_L.vals) .- tril(h_L)) <= 1e-10
    end
end

@testset "dcauchy" begin
    @kernel function dcauchy_test(::Val{n},dx,
                            dl,
                            du,
                            dA,
                            dg,
                            delta::Float64,
                            alpha::Float64,
                            d_out1,
                            d_out2
                            ) where {n}
        tx = @index(Local, Linear)
        bx = @index(Group, Linear)

        x = @localmem Float64 (n,)
        xl = @localmem Float64 (n,)
        xu = @localmem Float64 (n,)
        g = @localmem Float64 (n,)
        s = @localmem Float64 (n,)
        wa = @localmem Float64 (n,)
        A = @localmem Float64 (n,n)
        for i in 1:n
            A[tx,i] = dA[tx,i]
        end
        x[tx] = dx[tx]
        xl[tx] = dl[tx]
        xu[tx] = du[tx]
        g[tx] = dg[tx]

        alpha = ExaTronKernels.dcauchy(n,x,xl,xu,A,g,delta,alpha,s,wa,tx)
        if bx == 1
            d_out1[tx] = s[tx]
            d_out2[tx] = alpha
        end
        @synchronize
    end

    for i=1:itermax
        L = tril(rand(n,n))
        A = ExaTronKernels.TronDenseMatrix{Array{Float64,2}}(n)
        A.vals .= L*transpose(L)
        A.vals .= tril(A.vals) .+ (transpose(tril(A.vals)) .- Diagonal(A.vals))
        x = rand(n)
        xl = x .- abs.(rand(n))
        xu = x .+ abs.(rand(n))
        g = A.vals*x .+ rand(n)
        s = zeros(n)
        wa = zeros(n)
        alpha = 1.0
        delta = 2.0*norm(g)

        dx = AT{Float64}(undef, n)
        dl = AT{Float64}(undef, n)
        du = AT{Float64}(undef, n)
        dg = AT{Float64}(undef, n)
        dA = AT{Float64,2}(undef, (n,n))
        d_out1 = AT{Float64}(undef, n)
        d_out2 = AT{Float64}(undef, n)
        copyto!(dx, x)
        copyto!(dl, xl)
        copyto!(du, xu)
        copyto!(dg, g)
        copyto!(dA, A.vals)
        wait(dcauchy_test(device, n)(Val{n}(),dx,dl,du,dA,dg,delta,alpha,d_out1,d_out2,ndrange=(n,nblk),dependencies=Event(device)))
        h_s = zeros(n)
        h_alpha = zeros(n)
        copyto!(h_s, d_out1)
        copyto!(h_alpha, d_out2)

        alpha = ExaTronKernels.dcauchy(n, x, xl, xu, A, g, delta, alpha, s, wa)

        @test norm(s .- h_s) <= 1e-10
        @test norm(alpha .- h_alpha) <= 1e-10
    end
end

@testset "dtrpcg" begin
    @kernel function dtrpcg_test(::Val{n}, delta::Float64, tol::Float64,
                            stol::Float64, d_in,
                            d_g,
                            d_out_L,
                            d_out
                            ) where {n}
        tx = @index(Local, Linear)
        bx = @index(Group, Linear)

        A = @localmem Float64 (n,n)
        L = @localmem Float64 (n,n)

        g = @localmem Float64 (n,)
        w = @localmem Float64 (n,)
        p = @localmem Float64 (n,)
        q = @localmem Float64 (n,)
        r = @localmem Float64 (n,)
        t = @localmem Float64 (n,)
        z = @localmem Float64 (n,)
        for i in 1:n
            A[tx,i] = d_in[tx,i]
            L[tx,i] = d_in[tx,i]
        end
        g[tx] = d_g[tx]
        @synchronize

        ExaTronKernels.dicf(n,L,tx)
        info, iters = ExaTronKernels.dtrpcg(n,A,g,delta,L,tol,stol,n,w,p,q,r,t,z,tx)
        if bx == 1
            d_out[tx] = w[tx]
            for i in 1:n
                d_out_L[tx,i] = L[tx,i]
            end
        end
        @synchronize
    end

    delta = 100.0
    tol = 1e-6
    stol = 1e-6
    tron_A = ExaTronKernels.TronDenseMatrix{Array{Float64,2}}(n)
    tron_L = ExaTronKernels.TronDenseMatrix{Array{Float64,2}}(n)
    for i=1:itermax
        L = tril(rand(n,n))
        A = L*transpose(L)
        A .= tril(A) .+ (transpose(tril(A)) .- Diagonal(A))
        g = 0.1*ones(n)
        w = zeros(n)
        p = zeros(n)
        q = zeros(n)
        r = zeros(n)
        t = zeros(n)
        z = zeros(n)
        tron_A.vals .= A
        tron_L.vals .= A
        d_in = AT{Float64,2}(undef, (n,n))
        d_g = AT{Float64}(undef, n)
        d_out_L = AT{Float64,2}(undef, (n,n))
        d_out = AT{Float64}(undef, n)
        copyto!(d_in, A)
        copyto!(d_g, g)
        wait(dtrpcg_test(device, n)(Val{n}(),delta,tol,stol,d_in,d_g,d_out_L,d_out,ndrange=(n,nblk),dependencies=Event(device)))
        h_w = zeros(n)
        h_L = zeros(n,n)
        copyto!(h_L, d_out_L)
        copyto!(h_w, d_out)

        indr = zeros(Int, n)
        indf = zeros(n)
        list = zeros(n)
        ExaTronKernels.dicf(n, n^2, tron_L, 5, indr, indf, list, w)
        ExaTronKernels.dtrpcg(n, tron_A, g, delta, tron_L, tol, stol, n, w, p, q, r, t, z)

        @test norm(tril(h_L) .- tril(tron_L.vals)) <= tol
        @test norm(h_w .- w) <= tol
    end
end

@testset "dprsrch" begin
    @kernel function dprsrch_test(::Val{n},d_x,
                            d_xl,
                            d_xu,
                            d_g,
                            d_w,
                            d_A,
                            d_out1,
                            d_out2
                            ) where {n}
        tx = @index(Local, Linear)
        bx = @index(Group, Linear)

        x = @localmem Float64 (n,)
        xl = @localmem Float64 (n,)
        xu = @localmem Float64 (n,)
        g = @localmem Float64 (n,)
        w = @localmem Float64 (n,)
        wa1 = @localmem Float64 (n,)
        wa2 = @localmem Float64 (n,)
        A = @localmem Float64 (n,n)

        for i in 1:n
            A[tx,i] = d_A[tx,i]
        end
        x[tx] = d_x[tx]
        xl[tx] = d_xl[tx]
        xu[tx] = d_xu[tx]
        g[tx] = d_g[tx]
        w[tx] = d_w[tx]
        @synchronize

        ExaTronKernels.dprsrch(n, x, xl, xu, A, g, w, wa1, wa2, tx)
        if bx == 1
            d_out1[tx] = x[tx]
            d_out2[tx] = w[tx]
        end
        @synchronize

    end

    for i=1:itermax
        L = tril(rand(n,n))
        A = ExaTronKernels.TronDenseMatrix{Array{Float64,2}}(n)
        A.vals .= L*transpose(L)
        A.vals .= tril(A.vals) .+ (transpose(tril(A.vals)) .- Diagonal(A.vals))
        x = rand(n)
        xl = x .- abs.(rand(n))
        xu = x .+ abs.(rand(n))
        g = A.vals*x .+ rand(n)
        w = -g
        wa1 = zeros(n)
        wa2 = zeros(n)

        dx = AT{Float64}(undef, n)
        dl = AT{Float64}(undef, n)
        du = AT{Float64}(undef, n)
        dg = AT{Float64}(undef, n)
        dw = AT{Float64}(undef, n)
        dA = AT{Float64,2}(undef, (n,n))
        d_out1 = AT{Float64}(undef, n)
        d_out2 = AT{Float64}(undef, n)
        copyto!(dx, x)
        copyto!(dl, xl)
        copyto!(du, xu)
        copyto!(dg, g)
        copyto!(dw, w)
        copyto!(dA, A.vals)
        wait(dprsrch_test(device, n)(Val{n}(),dx,dl,du,dg,dw,dA,d_out1,d_out2,ndrange=(n,nblk),dependencies=Event(device)))
        h_x = zeros(n)
        h_w = zeros(n)
        copyto!(h_x, d_out1)
        copyto!(h_w, d_out2)

        ExaTronKernels.dprsrch(n,x,xl,xu,A,g,w,wa1,wa2)

        @test norm(x .- h_x) <= 1e-10
        @test norm(w .- h_w) <= 1e-10
    end
end

@testset "daxpy" begin
    @kernel function daxpy_test(::Val{n}, da, d_in,
                        d_out
                        ) where {n}
        tx = @index(Local, Linear)
        bx = @index(Group, Linear)

        x = @localmem Float64 (n,)
        y = @localmem Float64 (n,)
        x[tx] = d_in[tx]
        y[tx] = d_in[tx + n]
        @synchronize

        ExaTronKernels.daxpy(n, da, x, 1, y, 1, tx)
        if bx == 1
            d_out[tx] = y[tx]
        end
        @synchronize

    end

    for i=1:itermax
        da = rand(1)[1]
        h_in = rand(2*n)
        h_out = zeros(n)
        d_in = AT{Float64}(undef, 2*n)
        d_out = AT{Float64}(undef, n)
        copyto!(d_in, h_in)
        wait(daxpy_test(device,n)(Val{n}(),da,d_in,d_out,ndrange=(n,nblk),dependencies=Event(device)))
        copyto!(h_out, d_out)

        @test norm(h_out .- (h_in[n+1:2*n] .+ da.*h_in[1:n])) <= 1e-12
    end
end

@testset "dssyax" begin
    @kernel function dssyax_test(::Val{n},d_z,
                            d_in,
                            d_out
                            ) where {n}
        tx = @index(Local, Linear)
        bx = @index(Group, Linear)

        z = @localmem Float64 (n,)
        q = @localmem Float64 (n,)
        A = @localmem Float64 (n,n)

        for i in 1:n
            A[tx,i] = d_in[tx,i]
        end
        z[tx] = d_z[tx]
        @synchronize

        ExaTronKernels.dssyax(n, A, z, q, tx)
        if bx == 1
            d_out[tx] = q[tx]
        end
        @synchronize

    end

    for i=1:itermax
        z = rand(n)
        h_in = rand(n,n)
        h_out = zeros(n)
        d_z = AT{Float64}(undef, n)
        d_in = AT{Float64,2}(undef, (n,n))
        d_out = AT{Float64}(undef, n)
        copyto!(d_z, z)
        copyto!(d_in, h_in)
        wait(dssyax_test(device,n)(Val{n}(),d_z,d_in,d_out,ndrange=(n,nblk),dependencies=Event(device)))
        copyto!(h_out, d_out)

        @test norm(h_out .- h_in*z) <= 1e-12
    end
end

@testset "dmid" begin
    @kernel function dmid_test(::Val{n}, dx,
                        dl,
                        du,
                        d_out
                        ) where {n}
        tx = @index(Local, Linear)
        bx = @index(Group, Linear)

        x = @localmem Float64 (n,)
        xl = @localmem Float64 (n,)
        xu = @localmem Float64 (n,)
        x[tx] = dx[tx]
        xl[tx] = dl[tx]
        xu[tx] = du[tx]
        @synchronize

        ExaTronKernels.dmid(n, x, xl, xu, tx)
        if bx == 1
            d_out[tx] = x[tx]
        end
        @synchronize

    end

    for i=1:itermax
        x = rand(n)
        xl = x .- abs.(rand(n))
        xu = x .+ abs.(rand(n))

        # Force some components to go below or above bounds
        # so that we can test all cases.
        for j=1:n
            k = rand(1:3)
            if k == 1
                x[j] = xl[j] - 0.1
            elseif k == 2
                x[j] = xu[j] + 0.1
            end
        end
        x_out = zeros(n)
        dx = AT{Float64}(undef, n)
        dl = AT{Float64}(undef, n)
        du = AT{Float64}(undef, n)
        d_out = AT{Float64}(undef, n)
        copyto!(dx, x)
        copyto!(dl, xl)
        copyto!(du, xu)
        wait(dmid_test(device,n)(Val{n}(),dx,dl,du,d_out,ndrange=(n,nblk),dependencies=Event(device)))
        copyto!(x_out, d_out)

        ExaTronKernels.dmid(n, x, xl, xu)
        @test !(false in (x .== x_out))
    end
end

@testset "dgpstep" begin
    @kernel function dgpstep_test(::Val{n},dx,
                            dl,
                            du,
                            alpha::Float64,
                            dw,
                            d_out
                            ) where {n}
        tx = @index(Local, Linear)
        bx = @index(Group, Linear)

        x = @localmem Float64 (n,)
        xl = @localmem Float64 (n,)
        xu = @localmem Float64 (n,)
        w = @localmem Float64 (n,)
        s = @localmem Float64 (n,)
        x[tx] = dx[tx]
        xl[tx] = dl[tx]
        xu[tx] = du[tx]
        w[tx] = dw[tx]
        @synchronize

        ExaTronKernels.dgpstep(n, x, xl, xu, alpha, w, s, tx)
        if bx == 1
            d_out[tx] = s[tx]
        end
        @synchronize

    end

    for i=1:itermax
        x = rand(n)
        xl = x .- abs.(rand(n))
        xu = x .+ abs.(rand(n))
        w = rand(n)
        alpha = rand(1)[1]
        s = zeros(n)
        s_out = zeros(n)

        # Force some components to go below or above bounds
        # so that we can test all cases.
        for j=1:n
            k = rand(1:3)
            if k == 1
                if x[j] + alpha*w[j] >= xl[j]
                    w[j] = (xl[j] - x[j]) / alpha - 0.1
                end
            elseif k == 2
                if x[j] + alpha*w[j] <= xu[j]
                    w[j] = (xu[j] - x[j]) / alpha + 0.1
                end
            end
        end

        dx = AT{Float64}(undef, n)
        dl = AT{Float64}(undef, n)
        du = AT{Float64}(undef, n)
        dw = AT{Float64}(undef, n)
        d_out = AT{Float64}(undef, n)
        copyto!(dx, x)
        copyto!(dl, xl)
        copyto!(du, xu)
        copyto!(dw, w)
        wait(dgpstep_test(device,n)(Val{n}(),dx,dl,du,alpha,dw,d_out,ndrange=(n,nblk),dependencies=Event(device)))
        copyto!(s_out, d_out)

        ExaTronKernels.dgpstep(n, x, xl, xu, alpha, w, s)
        @test !(false in (s .== s_out))
    end
end

@testset "dbreakpt" begin
    @kernel function dbreakpt_test(::Val{n},dx,
                            dl,
                            du,
                            dw,
                            d_nbrpt,
                            d_brptmin,
                            d_brptmax
                            ) where {n}
        tx = @index(Local, Linear)
        bx = @index(Group, Linear)

        x = @localmem Float64 (n,)
        xl = @localmem Float64 (n,)
        xu = @localmem Float64 (n,)
        w = @localmem Float64 (n,)
        x[tx] = dx[tx]
        xl[tx] = dl[tx]
        xu[tx] = du[tx]
        w[tx] = dw[tx]
        @synchronize

        nbrpt, brptmin, brptmax = ExaTronKernels.dbreakpt(n, x, xl, xu, w, tx)
        if bx == 1
            for i in 1:n
                d_nbrpt[tx,i] = nbrpt
                d_brptmin[tx,i] = brptmin
                d_brptmax[tx,i] = brptmax
            end
        end
        @synchronize

    end

    for i=1:itermax
        x = rand(n)
        xl = x .- abs.(rand(n))
        xu = x .+ abs.(rand(n))
        w = 2.0*rand(n) .- 1.0     # (-1,1]
        h_nbrpt = zeros((n,n))
        h_brptmin = zeros((n,n))
        h_brptmax = zeros((n,n))

        dx = AT{Float64}(undef, n)
        dl = AT{Float64}(undef, n)
        du = AT{Float64}(undef, n)
        dw = AT{Float64}(undef, n)
        d_nbrpt = AT{Float64,2}(undef, (n,n))
        d_brptmin = AT{Float64,2}(undef, (n,n))
        d_brptmax = AT{Float64,2}(undef, (n,n))
        copyto!(dx, x)
        copyto!(dl, xl)
        copyto!(du, xu)
        copyto!(dw, w)
        wait(dbreakpt_test(device,n)(Val{n}(),dx,dl,du,dw,d_nbrpt,d_brptmin,d_brptmax,ndrange=(n,nblk),dependencies=Event(device)))
        copyto!(h_nbrpt, d_nbrpt)
        copyto!(h_brptmin, d_brptmin)
        copyto!(h_brptmax, d_brptmax)

        nbrpt, brptmin, brptmax = ExaTronKernels.dbreakpt(n, x, xl, xu, w)
        @test !(false in (nbrpt .== h_nbrpt))
        @test !(false in (brptmin .== h_brptmin))
        @test !(false in (brptmax .== h_brptmax))
    end
end

@testset "dnrm2" begin
    @kernel function dnrm2_test(::Val{n}, d_in,
                        d_out
                        ) where {n}
        tx = @index(Local, Linear)
        bx = @index(Group, Linear)

        x = @localmem Float64 (n,)
        x[tx] = d_in[tx]
        @synchronize

        v = ExaTronKernels.dnrm2(n, x, 1, tx)
        if bx == 1
            for i in 1:n
                d_out[tx,i] = v
            end
        end
        @synchronize

    end

    @inbounds for i=1:itermax
        h_in = rand(n)
        h_out = zeros((n,n))
        d_in = AT{Float64}(undef, n)
        d_out = AT{Float64,2}(undef, (n,n))
        copyto!(d_in, h_in)
        wait(dnrm2_test(device,n)(Val{n}(),d_in,d_out,ndrange=(n,nblk),dependencies=Event(device)))
        copyto!(h_out, d_out)
        xnorm = norm(h_in, 2)

        @test norm(xnorm .- h_out) <= 1e-10
    end
end

@testset "nrm2" begin
    @kernel function nrm2_test(::Val{n}, d_A, d_out) where {n}
        tx = @index(Local, Linear)
        bx = @index(Group, Linear)

        wa = @localmem Float64 (n,)
        A = @localmem Float64 (n,n)
        for i in 1:n
            A[tx,i] = d_A[tx,i]
        end
        @synchronize

        ExaTronKernels.nrm2!(wa, A, n, tx)
        if bx == 1
            d_out[tx] = wa[tx]
        end
        @synchronize
    end

    @inbounds for i=1:itermax
        L = tril(rand(n,n))
        A = L*transpose(L)
        A .= tril(A) .+ (transpose(tril(A)) .- Diagonal(A))
        wa = zeros(n)
        tron_A = ExaTronKernels.TronDenseMatrix{Array{Float64,2}}(n)
        tron_A.vals .= A
        ExaTronKernels.nrm2!(wa, tron_A, n)

        d_A = AT{Float64,2}(undef, (n,n))
        d_out = AT{Float64}(undef, n)
        h_wa = zeros(n)
        copyto!(d_A, A)
        wait(nrm2_test(device,n)(Val{n}(),d_A,d_out,ndrange=(n,nblk),dependencies=Event(device)))
        copyto!(h_wa, d_out)

        @test norm(wa .- h_wa) <= 1e-10
    end
end

@testset "dcopy" begin
    @kernel function dcopy_test(::Val{n}, d_in,
                        d_out
                        ) where {n}
        tx = @index(Local, Linear)
        bx = @index(Group, Linear)

        x = @localmem Float64 (n,)
        y = @localmem Float64 (n,)

        x[tx] = d_in[tx]
        @synchronize

        ExaTronKernels.dcopy(n, x, 1, y, 1, tx)

        if bx == 1
            d_out[tx] = y[tx]
        end
        @synchronize

    end

    @inbounds for i=1:itermax
        h_in = rand(n)
        h_out = zeros(n)
        d_in = AT{Float64}(undef, n)
        d_out = AT{Float64}(undef, n)
        copyto!(d_in, h_in)
        wait(dcopy_test(device,n)(Val{n}(),d_in,d_out,ndrange=(n,nblk),dependencies=Event(device)))
        copyto!(h_out, d_out)

        @test !(false in (h_in .== h_out))
    end
end

@testset "ddot" begin
    @kernel function ddot_test(::Val{n}, d_in,
                        d_out
                        ) where {n}
        tx = @index(Local, Linear)
        bx = @index(Group, Linear)

        x = @localmem Float64 (4,)
        y = @localmem Float64 (4,)
        x[tx] = d_in[tx]
        y[tx] = d_in[tx]
        @synchronize

        v = ExaTronKernels.ddot(n, x, 1, y, 1)

        if bx == 1
            for i in 1:n
                d_out[tx,i] = v
            end
        end
        @synchronize

    end

    @inbounds for i=1:itermax
        h_in = rand(n)
        h_out = zeros((n,n))
        d_in = AT{Float64}(undef, n)
        d_out = AT{Float64,2}(undef, (n,n))
        copyto!(d_in, h_in)
        wait(ddot_test(device, n)(Val{n}(),d_in,d_out,ndrange=(n,nblk),dependencies=Event(device)))
        copyto!(h_out, d_out)

        @test norm(dot(h_in,h_in) .- h_out, 2) <= 1e-10
    end
end

@testset "dscal" begin
    @kernel function dscal_test(::Val{n}, da::Float64,
                        d_in,
                        d_out
                        ) where {n}
        tx = @index(Local, Linear)
        bx = @index(Group, Linear)

        x = @localmem Float64 (n,)
        if tx <= n
            x[tx] = d_in[tx]
        end
        @synchronize

        ExaTronKernels.dscal(n, da, x, 1, tx)
        if bx == 1
            d_out[tx] = x[tx]
        end
        @synchronize

    end

    for i=1:itermax
        h_in = rand(n)
        h_out = zeros(n)
        da = rand(1)[1]
        d_in = AT{Float64}(undef, n)
        d_out = AT{Float64}(undef, n)
        copyto!(d_in, h_in)
        wait(dscal_test(device,n)(Val{n}(),da,d_in,d_out,ndrange=(n,nblk),dependencies=Event(device)))
        copyto!(h_out, d_out)

        @test norm(h_out .- (da.*h_in)) <= 1e-12
    end
end

@testset "dtrqsol" begin
    @kernel function dtrqsol_test(::Val{n}, d_x,
                            d_p,
                            d_out,
                            delta::Float64
                            ) where {n}
        tx = @index(Local, Linear)
        bx = @index(Group, Linear)

        x = @localmem Float64 (n,)
        p = @localmem Float64 (n,)

        x[tx] = d_x[tx]
        p[tx] = d_p[tx]
        @synchronize

        sigma = ExaTronKernels.dtrqsol(n, x, p, delta, tx)
        if bx == 1
            for i in 1:n
                d_out[tx,i] = sigma
            end
        end
        @synchronize
    end

    for i=1:itermax
        x = rand(n)
        p = rand(n)
        sigma = abs(rand(1)[1])
        delta = norm(x .+ sigma.*p)

        d_x = AT{Float64}(undef, n)
        d_p = AT{Float64}(undef, n)
        d_out = AT{Float64,2}(undef, (n,n))
        copyto!(d_x, x)
        copyto!(d_p, p)
        wait(dtrqsol_test(device,n)(Val{n}(),d_x,d_p,d_out,delta,ndrange=(n,nblk),dependencies=Event(device)))

        d_out = d_out |> Array
        @test norm(sigma .- d_out) <= 1e-10
    end
end

@testset "dspcg" begin
    @kernel function dspcg_test(::Val{n}, delta::Float64, rtol::Float64,
                        cg_itermax::Int, dx,
                        dxl,
                        dxu,
                        dA,
                        dg,
                        ds,
                        d_out
                        ) where {n}
        tx = @index(Local, Linear)
        bx = @index(Group, Linear)

        x = @localmem Float64 (n,)
        xl = @localmem Float64 (n,)
        xu = @localmem Float64 (n,)
        g = @localmem Float64 (n,)
        s = @localmem Float64 (n,)
        w = @localmem Float64 (n,)
        wa1 = @localmem Float64 (n,)
        wa2 = @localmem Float64 (n,)
        wa3 = @localmem Float64 (n,)
        wa4 = @localmem Float64 (n,)
        wa5 = @localmem Float64 (n,)
        gfree = @localmem Float64 (n,)
        indfree = @localmem Int (n,)
        iwa = @localmem Int (n,)

        A = @localmem Float64 (n,n)
        B = @localmem Float64 (n,n)
        L = @localmem Float64 (n,n)

        for i in 1:n
            A[i,tx] = dA[i,tx]
        end
        x[tx] = dx[tx]
        xl[tx] = dxl[tx]
        xu[tx] = dxu[tx]
        g[tx] = dg[tx]
        s[tx] = ds[tx]
        @synchronize

        ExaTronKernels.dspcg(n, delta, rtol, cg_itermax, x, xl, xu,
                        A, g, s, B, L, indfree, gfree, w, iwa,
                        wa1, wa2, wa3, wa4, wa5, tx)

        if bx == 1
            d_out[tx] = x[tx]
        end
        @synchronize

    end

    for i=1:itermax
        L = tril(rand(n,n))
        A = L*transpose(L)
        A .= tril(A) .+ (transpose(tril(A)) .- Diagonal(A))
        tron_A = ExaTronKernels.TronDenseMatrix{Array{Float64,2}}(n)
        tron_A.vals .= A
        tron_B = ExaTronKernels.TronDenseMatrix{Array{Float64,2}}(n)
        tron_L = ExaTronKernels.TronDenseMatrix{Array{Float64,2}}(n)
        x = rand(n)
        xl = x .- abs.(rand(n))
        xu = x .+ abs.(rand(n))
        g = A*x .+ rand(n)
        s = rand(n)
        delta = 2.0*norm(g)
        rtol = 1e-6
        cg_itermax = n
        w = zeros(n)
        wa = zeros(5*n)
        gfree = zeros(n)
        indfree = zeros(Int, n)
        iwa = zeros(Int, 3*n)

        dx = AT{Float64}(undef, n)
        dxl = AT{Float64}(undef, n)
        dxu = AT{Float64}(undef, n)
        dA = AT{Float64,2}(undef, (n,n))
        dg = AT{Float64}(undef, n)
        ds = AT{Float64}(undef, n)
        d_out = AT{Float64}(undef, n)

        copyto!(dx, x)
        copyto!(dxl, xl)
        copyto!(dxu, xu)
        copyto!(dA, tron_A.vals)
        copyto!(dg, g)
        copyto!(ds, s)

        wait(dspcg_test(device, n)(Val{n}(),delta,rtol,cg_itermax,dx,dxl,dxu,dA,dg,ds,d_out,ndrange=(n,1),dependencies=Event(device)))
        h_x = zeros(n)
        copyto!(h_x, d_out)

        ExaTronKernels.dspcg(n, x, xl, xu, tron_A, g, delta, rtol, s, 5, cg_itermax,
                        tron_B, tron_L, indfree, gfree, w, wa, iwa)

        @test norm(x .- h_x) <= 1e-10
    end
end

@testset "dgpnorm" begin
    @kernel function dgpnorm_test(::Val{n}, dx, dxl, dxu, dg, d_out) where {n}
        tx = @index(Local, Linear)
        bx = @index(Group, Linear)

        x = @localmem Float64 (n,)
        xl = @localmem Float64 (n,)
        xu = @localmem Float64 (n,)
        g = @localmem Float64 (n,)

        x[tx] = dx[tx]
        xl[tx] = dxl[tx]
        xu[tx] = dxu[tx]
        g[tx] = dg[tx]
        @synchronize

        v = ExaTronKernels.dgpnorm(n, x, xl, xu, g, tx)
        if bx == 1
            d_out[tx] = v
        end
        @synchronize

    end

    for i=1:itermax
        x = rand(n)
        xl = x .- abs.(rand(n))
        xu = x .+ abs.(rand(n))
        g = 2.0*rand(n) .- 1.0

        dx = AT{Float64}(undef, n)
        dxl = AT{Float64}(undef, n)
        dxu = AT{Float64}(undef, n)
        dg = AT{Float64}(undef, n)
        d_out = AT{Float64}(undef, n)

        copyto!(dx, x)
        copyto!(dxl, xl)
        copyto!(dxu, xu)
        copyto!(dg, g)

        wait(dgpnorm_test(device, n)(Val{n}(), dx, dxl, dxu, dg, d_out, ndrange=(n,n*nblk), dependencies=Event(device)))
        h_v = zeros(n)
        copyto!(h_v, d_out)

        v = ExaTronKernels.dgpnorm(n, x, xl, xu, g)
        @test norm(h_v .- v) <= 1e-10
    end
end

@testset "dtron" begin
    @kernel function dtron_test(::Val{n}, f::Float64, frtol::Float64, fatol::Float64, fmin::Float64,
                        cgtol::Float64, cg_itermax::Int, delta::Float64, task::Int,
                        disave, ddsave,
                        dx, dxl,
                        dxu, dA,
                        dg, d_out
                        ) where {n}
        tx = @index(Local, Linear)
        bx = @index(Group, Linear)

        x = @localmem Float64 (n,)
        xl = @localmem Float64 (n,)
        xu = @localmem Float64 (n,)
        g = @localmem Float64 (n,)
        xc = @localmem Float64 (n,)
        s = @localmem Float64 (n,)
        wa = @localmem Float64 (n,)
        wa1 = @localmem Float64 (n,)
        wa2 = @localmem Float64 (n,)
        wa3 = @localmem Float64 (n,)
        wa4 = @localmem Float64 (n,)
        wa5 = @localmem Float64 (n,)
        gfree = @localmem Float64 (n,)
        indfree = @localmem Int (n,)
        iwa = @localmem Int (2*n,)

        A = @localmem Float64 (n,n)
        B = @localmem Float64 (n,n)
        L = @localmem Float64 (n,n)

        for i in 1:n
            A[i,tx] = dA[i,tx]
        end
        x[tx] = dx[tx]
        xl[tx] = dxl[tx]
        xu[tx] = dxu[tx]
        g[tx] = dg[tx]
        @synchronize

        ExaTronKernels.dtron(n, x, xl, xu, f, g, A, frtol, fatol, fmin, cgtol,
                        cg_itermax, delta, task, B, L, xc, s, indfree, gfree,
                        disave, ddsave, wa, iwa, wa1, wa2, wa3, wa4, wa5, tx)
        if bx == 1
            d_out[tx] = x[tx]
        end
        @synchronize

    end

    for i=1:itermax
        L = tril(rand(n,n))
        A = L*transpose(L)
        A .= tril(A) .+ (transpose(tril(A)) .- Diagonal(A))
        x = rand(n)
        xl = x .- abs.(rand(n))
        xu = x .+ abs.(rand(n))
        c = rand(n)
        g = A*x .+ c
        xc = zeros(n)
        s = zeros(n)
        wa = zeros(7*n)
        gfree = zeros(n)
        indfree = zeros(Int, n)
        iwa = zeros(Int, 3*n)
        isave = zeros(Int, 3)
        dsave = zeros(3)
        task = 0
        fatol = 0.0
        frtol = 1e-12
        fmin = -1e32
        cgtol = 0.1
        cg_itermax = n
        delta = 2.0*norm(g)
        f = 0.5*transpose(x)*A*x .+ transpose(x)*c

        tron_A = ExaTronKernels.TronDenseMatrix{Array{Float64,2}}(n)
        tron_A.vals .= A
        tron_B = ExaTronKernels.TronDenseMatrix{Array{Float64,2}}(n)
        tron_L = ExaTronKernels.TronDenseMatrix{Array{Float64,2}}(n)

        dx = AT{Float64}(undef, n)
        dxl = AT{Float64}(undef, n)
        dxu = AT{Float64}(undef, n)
        dA = AT{Float64,2}(undef, (n,n))
        dg = AT{Float64}(undef, n)
        disave = AT{Int}(undef, n)
        ddsave = AT{Float64}(undef, n)
        d_out = AT{Float64}(undef, n)

        copyto!(dx, x)
        copyto!(dxl, xl)
        copyto!(dxu, xu)
        copyto!(dA, tron_A.vals)
        copyto!(dg, g)

        wait(dtron_test(device,n)(Val{n}(),f,frtol,fatol,fmin,cgtol,cg_itermax,delta,task,disave,ddsave,dx,dxl,dxu,dA,dg,d_out,ndrange=(n,n*nblk),dependencies=Event(device)))
        h_x = zeros(n)
        copyto!(h_x, d_out)

        task_str = Vector{UInt8}(undef, 60)
        for (i,s) in enumerate("START")
            task_str[i] = UInt8(s)
        end

        ExaTronKernels.dtron(n, x, xl, xu, f, g, tron_A, frtol, fatol, fmin, cgtol,
                        cg_itermax, delta, task_str, tron_B, tron_L, xc, s, indfree,
                        isave, dsave, wa, iwa)
        @test norm(x .- h_x) <= 1e-10
    end
end

@testset "driver_kernel" begin
    @inline function eval_f(n, x, dA, dc)
        f = 0
        @inbounds for i=1:n
            @inbounds for j=1:n
                f += x[i]*dA[i,j]*x[j]
            end
        end
        f = 0.5*f
        @inbounds for i=1:n
            f += x[i]*dc[i]
        end
        @synchronize
        return f
    end

    @inline function eval_g(n, x, g, dA, dc)
        @inbounds for i=1:n
            gval = 0
            @inbounds for j=1:n
                gval += dA[i,j]*x[j]
            end
            g[i] = gval + dc[i]
        end
        @synchronize
        return
    end

    @inline function eval_h(n, scale, x, A, dA, tx)
        for i in 1:n
            A[i,tx] = dA[i,tx]
        end
        @synchronize
        return
    end

    @inline function driver_kernel(n, max_feval::Int, max_minor::Int,
                           x, xl,
                           xu, dA,
                           dc,
                           tx)
        # We start with a shared memory allocation.
        # The first 3*n*sizeof(Float64) bytes are used for x, xl, and xu.
        g = @localmem Float64 (n,)
        xc = @localmem Float64 (n,)
        s = @localmem Float64 (n,)
        wa = @localmem Float64 (n,)
        wa1 = @localmem Float64 (n,)
        wa2 = @localmem Float64 (n,)
        wa3 = @localmem Float64 (n,)
        wa4 = @localmem Float64 (n,)
        wa5 = @localmem Float64 (n,)
        gfree = @localmem Float64 (n,)
        dsave = @localmem Float64 (n,)
        indfree = @localmem Int (n,)
        iwa = @localmem Int (2*n,)
        isave = @localmem Int (n,)

        A = @localmem Float64 (n,n)
        B = @localmem Float64 (n,n)
        L = @localmem Float64 (n,n)

        task = 0
        status = 0

        delta = 0.0
        fatol = 0.0
        frtol = 1e-12
        fmin = -1e32
        gtol = 1e-6
        cgtol = 0.1
        cg_itermax = n

        f = 0.0
        nfev = 0
        ngev = 0
        nhev = 0
        minor_iter = 0
        search = true

        while search

            # [0|1]: Evaluate function.

            if task == 0 || task == 1
                f = eval_f(n, x, dA, dc)
                nfev += 1
                if nfev >= max_feval
                    search = false
                end
            end

            # [2] G or H: Evaluate gradient and Hessian.

            if task == 0 || task == 2
                eval_g(n, x, g, dA, dc)
                eval_h(n, 1.0, x, A, dA, tx)
                ngev += 1
                nhev += 1
                minor_iter += 1
            end

            # Initialize the trust region bound.

            if task == 0
                gnorm0 = ExaTronKernels.dnrm2(n, g, 1, tx)
                delta = gnorm0
            end

            # Call Tron.

            if search
                delta, task = ExaTronKernels.dtron(n, x, xl, xu, f, g, A, frtol, fatol, fmin, cgtol,
                                            cg_itermax, delta, task, B, L, xc, s, indfree, gfree,
                                            isave, dsave, wa, iwa, wa1, wa2, wa3, wa4, wa5, tx)
            end

            # [3] NEWX: a new point was computed.

            if task == 3
                gnorm_inf = ExaTronKernels.dgpnorm(n, x, xl, xu, g, tx)
                if gnorm_inf <= gtol
                    task = 4
                end

                if minor_iter >= max_minor
                    status = 1
                    search = false
                end
            end

            # [4] CONV: convergence was achieved.

            if task == 4
                search = false
            end
        end

        return status, minor_iter
    end

    @kernel function driver_kernel_test(::Val{n}, max_feval, max_minor,
                                dx, dxl, dxu, dA, dc, d_out) where {n}
        tx = @index(Local, Linear)
        bx = @index(Group, Linear)

        x = @localmem Float64 (n,)
        xl = @localmem Float64 (n,)
        xu = @localmem Float64 (n,)

        x[tx] = dx[tx]
        xl[tx] = dxl[tx]
        xu[tx] = dxu[tx]
        @synchronize

        status, minor_iter = driver_kernel(n, max_feval, max_minor, x, xl, xu, dA, dc,tx)

        if bx == 1
            d_out[tx] = x[tx]
        end
        @synchronize
    end

    max_feval = 500
    max_minor = 100

    tron_A = ExaTronKernels.TronDenseMatrix{Array{Float64,2}}(n)
    tron_B = ExaTronKernels.TronDenseMatrix{Array{Float64,2}}(n)
    tron_L = ExaTronKernels.TronDenseMatrix{Array{Float64,2}}(n)

    dx = AT{Float64}(undef, n)
    dxl = AT{Float64}(undef, n)
    dxu = AT{Float64}(undef, n)
    dA = AT{Float64,2}(undef, (n,n))
    dc = AT{Float64}(undef, n)
    d_out = AT{Float64}(undef, n)

    for i=1:itermax
        L = tril(rand(n,n))
        A = L*transpose(L)
        A .= tril(A) .+ (transpose(tril(A)) .- Diagonal(A))
        x = rand(n)
        xl = x .- abs.(rand(n))
        xu = x .+ abs.(rand(n))
        c = rand(n)

        tron_A.vals .= A

        copyto!(dx, x)
        copyto!(dxl, xl)
        copyto!(dxu, xu)
        copyto!(dA, tron_A.vals)
        copyto!(dc, c)

        function eval_f_cb(x)
            f = 0.5*(transpose(x)*A*x) + transpose(c)*x
            return f
        end

        function eval_g_cb(x, g)
            g .= A*x .+ c
        end

        function eval_h_cb(x, mode, rows, cols, scale, lambda, values)
            nz = 1
            if mode == :Structure
                @inbounds for j=1:n
                    @inbounds for i=j:n
                        rows[nz] = i
                        cols[nz] = j
                        nz += 1
                    end
                end
            else
                @inbounds for j=1:n
                    @inbounds for i=j:n
                        values[nz] = A[i,j]
                        nz += 1
                    end
                end
            end
        end

        nele_hess = div(n*(n+1), 2)
        tron = ExaTronKernels.createProblem(n, xl, xu, nele_hess, eval_f_cb, eval_g_cb, eval_h_cb; :matrix_type=>:Dense, :max_minor=>max_minor)
        copyto!(tron.x, x)
        status = ExaTronKernels.solveProblem(tron)

        wait(driver_kernel_test(device,n)(Val{n}(),max_feval,max_minor,dx,dxl,dxu,dA,dc,d_out,ndrange=(n,nblk),dependencies=Event(device)))
        h_x = zeros(n)
        copyto!(h_x, d_out)

        @test norm(h_x .- tron.x) <= 1e-10
    end
end


text:
import Base.Broadcast: Broadcasted, AbstractArrayStyle, DefaultArrayStyle, flatten
toRanges(r::Tuple) = r
toRanges(r::CartesianIndices) = r.indices

struct ChunkStyle{N} <: Base.Broadcast.AbstractArrayStyle{N} end

Base.BroadcastStyle(::ChunkStyle{N}, ::ChunkStyle{M}) where {N,M}= ChunkStyle{max(N,M)}()
Base.BroadcastStyle(::ChunkStyle{N}, ::DefaultArrayStyle{M}) where {N,M}= ChunkStyle{max(N,M)}()
Base.BroadcastStyle(::DefaultArrayStyle{M},::ChunkStyle{N}) where {N,M}= ChunkStyle{max(N,M)}()

struct BroadcastDiskArray{T,N,BC<:Broadcasted{<:ChunkStyle{N}}} <: AbstractDiskArray{T,N}
  bc::BC
end
function BroadcastDiskArray(bcf::B) where {B<:Broadcasted{<:ChunkStyle{N}}} where N
  ElType = Base.Broadcast.combine_eltypes(bcf.f, bcf.args)
  BroadcastDiskArray{ElType,N,B}(bcf)
end
Base.size(bc::BroadcastDiskArray) = size(bc.bc)
function DiskArrays.readblock!(a::BroadcastDiskArray,aout,i::OrdinalRange...)
  argssub = map(arg->subsetarg(arg,i),a.bc.args)
  aout .= a.bc.f.(argssub...)
end
Base.broadcastable(bc::BroadcastDiskArray) = bc.bc
haschunks(a::BroadcastDiskArray) = Chunked()
function eachchunk(a::BroadcastDiskArray)
  cs,off = common_chunks(size(a.bc),a.bc.args...)
  GridChunks(a.bc,cs,offset=off)
end
function Base.copy(bc::Broadcasted{ChunkStyle{N}}) where N
  BroadcastDiskArray(flatten(bc))
end
Base.copy(a::BroadcastDiskArray) = copyto!(zeros(eltype(a),size(a)),a.bc)
function Base.copyto!(dest::AbstractArray, bc::Broadcasted{ChunkStyle{N}}) where N
  bcf = flatten(bc)
  cs,off = common_chunks(size(bcf),dest,bcf.args...)
  gcd = GridChunks(bcf,cs,offset=off)
  foreach(gcd) do cnow
    #Possible optimization would be to use a LRU cache here, so that data has not
    #to be read twice in case of repeating indices
    argssub = map(i->subsetarg(i,cnow.indices),bcf.args)
    dest[cnow.indices...] .= bcf.f.(argssub...)
  end
  dest
end
function common_chunks(s,args...)
  N = length(s)
  chunkedars = filter(i->haschunks(i)===Chunked(),collect(args))
  all(ar->isa(eachchunk(ar),GridChunks), chunkedars) || error("Currently only chunks of type GridChunks can be merged by broadcast")
  if isempty(chunkedars)
    totalsize = sum(sizeof ∘ eltype, args)
    return (estimate_chunksize(s,totalsize),ntuple(zero,N))
  else

    allcs = map(ar->(eachchunk(ar).chunksize,eachchunk(ar).offset),chunkedars)
    tt = ntuple(N) do n
      csnow = filter(cs->length(cs[1])>=n && cs[1][n]>1,allcs)
      isempty(csnow) && return (1, 0)
      cs = (csnow[1][1][n],csnow[1][2][n])
      all(s->(s[1][n],s[2][n]) == cs,csnow) || error("Chunks do not align in dimension $n")
      return cs
    end
    return map(i->i[1],tt),map(i->i[2],tt)
  end
end
subsetarg(x, a) = x
function subsetarg(x::AbstractArray,a)
  ashort = maybeonerange(size(x),a)
  view(x,ashort...) #Maybe making a copy here would be faster, need to check...
end
repsingle(s,r) = s==1 ? (1:1) : r
function maybeonerange(out,sizes,ranges)
  s1,sr = splittuple(sizes...)
  r1,rr = splittuple(ranges...)
  maybeonerange((out...,repsingle(s1,r1)),sr,rr)
end
maybeonerange(out,::Tuple{},ranges) = out
maybeonerange(sizes,ranges) = maybeonerange((),sizes,ranges)
splittuple(x1,xr...) = x1,xr

macro implement_mapreduce(t)
  quote
    function Base._mapreduce(f,op,v::$t)
      mapreduce(op,eachchunk(v)) do cI
        a = v[toRanges(cI)...]
        mapreduce(f,op,a)
      end
    end
    function Base.mapreducedim!(f, op, R::AbstractArray, A::$t)
      sR = size(R)
      foreach(eachchunk(A)) do cI
        a = A[toRanges(cI)...]
        ainds = map((cinds, arsize)->arsize==1 ? Base.OneTo(1) : cinds, toRanges(cI),size(R))
        #Maybe the view into R here is problematic and a copy would be faster
        Base.mapreducedim!(f,op,view(R,ainds...),a)
      end
      R
    end

    function Base.mapfoldl_impl(f, op, nt::NamedTuple{()}, itr::$t)
      cc = eachchunk(itr)
      isempty(cc) && return Base.mapreduce_empty_iter(f, op, itr, IteratorEltype(itr))
      Base.mapfoldl_impl(f,op,nt,itr,cc)
    end
    function Base.mapfoldl_impl(f, op, nt::NamedTuple{()}, itr::$t, cc)
      y = first(cc)
      a = itr[toRanges(y)...]
      init = mapfoldl(f,op,a)
      Base.mapfoldl_impl(f,op,(init=init,),itr,Iterators.drop(cc,1))
    end
    function Base.mapfoldl_impl(f, op, nt::NamedTuple{(:init,)}, itr::$t, cc)
      init = nt.init
      for y in cc
        a = itr[toRanges(y)...]
        init = mapfoldl(f,op,a,init=init)
      end
      init
    end
  end
end

macro implement_broadcast(t)
  quote
    #Broadcasting with a DiskArray on LHS
    function Base.copyto!(dest::$t, bc::Broadcasted{Nothing})
      foreach(eachchunk(dest)) do c
        ar = [bc[i] for i in CartesianIndices(toRanges(c))]
        dest[toRanges(c)...] = ar
      end
      dest
    end
    Base.BroadcastStyle(T::Type{<:$t}) = ChunkStyle{ndims(T)}()
    function DiskArrays.subsetarg(x::$t,a)
      ashort = maybeonerange(size(x),a)
      x[ashort...]
    end

    #This is a heavily allocating implementation, but working for all cases.
    #As performance optimization one might:
    #Allocate the array only once if all chunks have the same size
    #Use FillArrays, if the DiskArray accepts these
    function Base.fill!(dest::$t, value)
      foreach(eachchunk(dest)) do c
        ar = fill(value,length.(toRanges(c)))
        dest[toRanges(c)...] = ar
      end
      dest
    end
  end
end

ext: jl


text:
function attach_dirichlet(cellmatvec,cellvals,cellmask=Fill(true,length(cellvals)))
  k = AttachDirichletKernel()
  apply(k,cellmatvec,cellvals,cellmask)
end

struct AttachDirichletKernel <: Kernel
  muladd::MulAddKernel{Int}
  AttachDirichletKernel() = new(MulAddKernel(-1,1))
end

function Arrays.kernel_cache(k::AttachDirichletKernel,matvec::Tuple,vals,mask)
  mat, vec = matvec
  kernel_cache(k.muladd,mat,vals,vec)
end

@inline function Arrays.apply_kernel!(cache,k::AttachDirichletKernel,matvec::Tuple,vals,mask)
  if mask
    mat, vec = matvec
    vec_with_bcs = apply_kernel!(cache,k.muladd,mat,vals,vec)
    (mat, vec_with_bcs)
  else
    matvec
  end
end

function Arrays.kernel_cache(k::AttachDirichletKernel,mat::AbstractMatrix,vals,mask)
  cm = kernel_cache(MulKernel(),mat,vals)
  cv = CachedArray(mat*vals)
  fill!(cv.array,zero(eltype(cv)))
  (cm,cv)
end

@inline function Arrays.apply_kernel!(cache,k::AttachDirichletKernel,mat::AbstractMatrix,vals,mask)
  cm, cv = cache
  if mask
    vec_with_bcs = apply_kernel!(cm,MulKernel(),mat,vals)
    scale_entries!(vec_with_bcs,-1)
    (mat, vec_with_bcs)
  else
    if size(mat,1) != size(cv,1)
      m = axes(mat,1)
      setaxes!(cv,(m,))
      fill!(cv.array,zero(eltype(cv)))
    end
    (mat, cv.array)
  end
end


text:
using ITensors
using ITensorVisualizationBase: visualize

_fill_or_identity(v::Vector, n) = v
_fill_or_identity(v, n) = fill(v, n)

function n_edges(A::Vector{ITensor})
  n_edges = 0
  for i in 1:length(A)
    for j in (i + 1):length(A)
      if hascommoninds(A[i], A[j])
        n_edges += 1
      end
    end
  end
  return n_edges
end

function _visualize_kwargs(A::Vector{ITensor};
  vertex_labels="",
  tags=true,
  plevs=true,
  textsize=60,
  vertex_size=100,
  edge_widths=10,
)
  # Number of tensors
  nt = length(A)

  # Number of dangling
  nd = length(noncommoninds(A...))

  # Number of internal edges
  ne_t = n_edges(A)

  # Number of total vertices, including the
  # tensors and dangling
  nv = nt + nd

  # Number of total edges, including between
  # tensors and dangling
  ne = ne_t + nd

  vertex_labels = _fill_or_identity(vertex_labels, nt)
  vertex_textsize = _fill_or_identity(textsize, nt)
  vertex_size = _fill_or_identity(vertex_size, nt)
  edge_textsize = _fill_or_identity(textsize, ne)
  edge_widths = _fill_or_identity(edge_widths, ne)
  return (;
    vertex_labels=[n in 1:nt ? vertex_labels[n] : "" for n in 1:nv],
    vertex_textsize=[n in 1:nt ? vertex_textsize[n] : 0 for n in 1:nv],
    vertex_size=[n in 1:nt ? vertex_size[n] : 5 for n in 1:nv],
    edge_labels=(; plevs=plevs, dims=false, tags=tags),
    edge_textsize=edge_textsize,
    edge_widths=edge_widths,
  )
end

function visualize_dangling(A::Vector{ITensor}; kwargs...)
  is = noncommoninds(A...)
  vs = [ITensor(i) for i in is]
  Av = [A; vs]
  visualize_kwargs = _visualize_kwargs(A; kwargs...)
  return visualize(Av; visualize_kwargs...)
end

visualize_dangling(A::ITensor; kwargs...) = visualize_dangling([A]; kwargs...)


text:
#=
    METAPROGRAMMING INFIX ASSIGNMENT OPERATOR
    -----------------------------------------
    streamlines the creation of objects via string-valued left-hand-sides:
    the line
        "handle" ∷ object
    is equivalent to
        handle = object
    without the quotes, and so object naming can be neatly scripted
=#

∷(handle::String, value::Any) = :( $(Symbol(handle)) = $value ) |> eval

max_stars_repo_name: mll-gh/scratch

max_issues_repo_name: mll-gh/scratch

max_forks_repo_name: mll-gh/scratch


text:
include("bils_chang_kim_2011/steady_state_equilibrium.jl")

#using Main.SSEquilibrium

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# Test out HHBellmanMap()

para  = ModelParams()
W_old = ones(length(para.agrid),length(para.xgrid))
U_old = ones(length(para.agrid))
wage  = ones(length(para.agrid),length(para.xgrid))
W_new, U_new, emp_policy, unemp_policy = HHBellmanMap(para, wage, W_old, U_old)

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# Test out SolveHHBellman()

para  = ModelParams()
W_old = ones(length(para.agrid),length(para.xgrid))
U_old = ones(length(para.agrid))
#wage  = 10*randn(length(para.agrid),length(para.xgrid)).^2
wage = ones(length(para.agrid),length(para.xgrid))
for i in 1:length(para.agrid)
    for j in 1:length(para.xgrid)
        wage[i,j] = i*j
    end
end

W, U, emp_policy, unemp_policy = SolveHHBellman(para, wage*0.9, W_old, U_old)

using Plots

fig = plot();
for i in 1:9
    plot!(fig, 1:length(W[:,i]), W[:,i]);
end
plot(fig)
fig2 = plot();
plot!(fig2, 1:length(U), U);
plot(fig2)

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# Test out UpdateWage()

wage_old = wage
wage_new, J = UpdateWage(para, W, U, emp_policy, wage_old)
wage_old = wage_new

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# Test out SolveWage()

# Load functions
include("bils_chang_kim_2011/steady_state_equilibrium.jl")

# Approximate value, policy, and wage functions
para = ModelParams(θ=1.0)
para.agrid = LinRange(para.amin, para.amax, 50)
W, U, J, emp_policy, unemp_policy, wage = SolveWage(para)

# Plot approximated functions
using Plots

## Plot value of employment for each x
fig = plot();
for i in 1:9
    plot!(fig, para.agrid, W[:,i]);
end
plot(fig)

## Plot value of unemployment
fig2 = plot();
plot!(fig2, para.agrid, U);
plot(fig2)

## Plot value of vacancy for each x
fig3 = plot();
for i in 1:9
    plot!(fig3, para.agrid, J[:,i]);
end
plot(fig3)

## Plot wage for each x
fig4 = plot();
for i in 1:9
    plot!(fig4, para.agrid, wage[:,i]);
end
plot(fig4)

# Save objects externally
using DelimitedFiles
writedlm( "bils_chang_kim_2011/W.csv",  W, ',')
writedlm( "bils_chang_kim_2011/U.csv",  U, ',')
writedlm( "bils_chang_kim_2011/J.csv",  J, ',')
writedlm( "bils_chang_kim_2011/emp_policy.csv",  emp_policy, ',')
writedlm( "bils_chang_kim_2011/unemp_policy.csv",  unemp_policy, ',')
writedlm( "bils_chang_kim_2011/wage.csv",  wage, ',')

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# Analyze exported objects from SolveWage()

#####################################
# Import saved objects
using DataFrames
using CSV

## Value of employment
W = CSV.read("bils_chang_kim_2011/W.csv", DataFrame, header = false)
W = Matrix(W)

# Value of unemployment
U = CSV.read("bils_chang_kim_2011/U.csv", DataFrame, header = false)
U = Matrix(U)

# Value of a matched job
J = CSV.read("bils_chang_kim_2011/J.csv", DataFrame, header = false)
J = Matrix(J)

# Asset policy function under employment
emp_policy      = CSV.read("bils_chang_kim_2011/emp_policy.csv", DataFrame, header = false)
emp_policy      = Matrix(emp_policy)

# Asset policy function under unemployment
unemp_policy    = CSV.read("bils_chang_kim_2011/unemp_policy.csv", DataFrame, header = false)
unemp_policy    = Matrix(unemp_policy)

# Wage function
wage            = CSV.read("bils_chang_kim_2011/wage.csv", DataFrame, header = false)
wage            = Matrix(wage)

#####################################
# Plot approximated functions
using Plots
using LaTeXStrings

## Save asset grid for x-axes
agrid = ModelParams().agrid
xgrid = ModelParams().xgrid

## Plot value of employment for each x
fig1 = plot();
for i in 1:9
    plot!(fig1, agrid, emp_policy[:,i]);
end
plot(fig1)

## Plot value of unemployment
fig2 = plot();
plot!(fig2, agrid, unemp_policy);
plot(fig2)

## Plot value of vacancy for each x
fig3 = plot(title = "Matched Firm Value");
for i in 1:9
    plot!(fig3, agrid, J[:,i], label="x=$(round(xgrid[i],digits=2))");
end
xlabel!("Level of Asset Holdings");
ylabel!(L"J");
plot(fig3)

## Plot wages for each x
fig4 = plot(title = "Wage Schedule");
for i in 1:9
    plot!(fig4, agrid, wage[:,i], label="x=$(round(xgrid[i],digits=2))");
end
xlabel!("Level of Asset Holdings");
ylabel!("Wage");
plot(fig4)

## Plot W - U for each x
fig5 = plot(title = "Employment vs. Unemployment");
for i in 1:9
    plot!(fig5, agrid, W[:,i] - U, label="x=$(round(xgrid[i],digits=2))");
end
xlabel!("Level of Asset Holdings");
ylabel!(L"W-U");
plot(fig5)

# All plots together
plot(fig1)
plot(fig2)
plot(fig3)
plot(fig4)
plot(fig5)

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# Test out ReservationProductivity()

# Load functions
include("bils_chang_kim_2011/steady_state_equilibrium.jl")

# Test
x_star = ReservationProductivity(para, W, U)
plot(1:length(x_star),x_star)

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# Test out ConstructTransitionMatrices()

# Load functions
include("bils_chang_kim_2011/steady_state_equilibrium.jl")

# Test
para = ModelParams()
para.agrid = LinRange(para.amin, para.amax, 50)
H_ee, H_eu, H_ue, H_uu = InitializeTransitionMatrices(para, emp_policy, unemp_policy)
π_emp, π_unemp = FindStationaryMeasures(para,  emp_policy, unemp_policy, x_star)

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# Test out AdjustTheta()
θ_new = AdjustTheta(para, J, π_unemp)

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# Test SolveEquilibrium()
para = ModelParams()
para, θ, W, U, J, emp_policy, unemp_policy, wage, x_star, π_emp, π_unemp = SolveEquilibrium(para::ModelParams)

# Save objects externally
using DelimitedFiles
writedlm( "bils_chang_kim_2011/data/theta.csv",  θ, ',')
writedlm( "bils_chang_kim_2011/data/W.csv",  W, ',')
writedlm( "bils_chang_kim_2011/data/U.csv",  U, ',')
writedlm( "bils_chang_kim_2011/data/J.csv",  J, ',')
writedlm( "bils_chang_kim_2011/data/emp_policy.csv",  emp_policy, ',')
writedlm( "bils_chang_kim_2011/data/unemp_policy.csv",  unemp_policy, ',')
writedlm( "bils_chang_kim_2011/data/wage.csv",  wage, ',')
writedlm( "bils_chang_kim_2011/data/x_star.csv",  x_star, ',')
writedlm( "bils_chang_kim_2011/data/pi_emp.csv",  π_emp, ',')
writedlm( "bils_chang_kim_2011/data/pi_unemp.csv",  π_unemp, ',')


text:
######################################################################################################
# Copy from https://github.com/xuedong/LinBAI.jl  by Xuedong Shang                                   #
# We replace Xuedong's `glrt` with our elegant `glrt` that is based on Envelope Theorem.             #
######################################################################################################
using Random;
using CPUTime;
using LinearAlgebra;
include("../utilities/peps.jl");
include("../utilities/expfam.jl");
include("samplingrules.jl");

# Run the learning algorithm, paramterised by a sampling rule
# The stopping and recommendation rules are common
# βs must be a list of thresholds *in increasing order*

function play!(i, k, rng, pep, µ, S, N, Vinv)
    arm = pep.arms[k]
    Y = sample(rng, getexpfam(pep, 1), arm'µ)
    S .+= Y .* arm
    Vinv .= sherman_morrison(Vinv, arm)
    N[k] += 1
end

function runit(seed, sr, μ, pep::LinearBestArm, βs, δs)
    # seed: random seed. UInt.
    # sr: sampling rule.
    # μ: mean vector.
    # pep: pure exploration problem.
    # βs: list of thresholds.
    convex_sr = (typeof(sr) == ConvexGame) || (typeof(sr) == LearnerK);
    xya_sr = typeof(sr) == XYAdaptive;
    qbc_sr = typeof(sr) == FWSampling;
    lztas_sr = typeof(sr) == LazyTrackAndStop;
    gap_sr = typeof(sr) == LinGapE;

    βs = collect(βs) # mutable copy
    δs = collect(δs); # mutable copy
    rng = MersenneTwister(seed)
    K = narms(pep, μ);
    nb_I = nanswers(pep, μ)
    dim = length(μ)
    N = zeros(Int64, K) # allocations
    S = zeros(dim) # sum of samples
    Vinv = Matrix{Float64}(I, dim, dim)  # inverse of the design matrix
    (convex_sr) ? P = ones(Int64, (nb_I, K)) : nothing; # P is used only for (Degenne et al. 2020)

    if xya_sr
        ρ = 1; ρ_old = 1; Xactive = copy(pep.arms); α = 0.1; # setting by Soare'14
    elseif qbc_sr # initialize Vxinv matrix
        Vxinv = zeros(Float64, dim,dim);
        for k in 1:K
            Vxinv += pep.arms[k]*(pep.arms[k]')/K;
        end
    end

    # Force exploration for initialization
    baseline = CPUtime_us();
    if lztas_sr # the strategy in Lazy T&S (Jedra and Proutiere, 2020)
        lztas_A = sum([N[k]*pep.arms[k]*(pep.arms[k]') for k=1:K]);
        lztas_A0 = zeros(Int64, dim);
        r = 0;
        while r < dim
            k = rand(rng, 1:K);
            play!(1, k, rng, pep, μ, S, N, Vinv);
            if rank(lztas_A + pep.arms[k]*(pep.arms[k]')) > r
                lztas_A += pep.arms[k]*(pep.arms[k]');
                lztas_A0[r+1] = k; # P represents A0
                r += 1;
            end
        end
        lztas_c0 = minimum(eigvals(lztas_A)) / sqrt(dim);
        lztas_c1 = minimum(eigvals(lztas_A));
        lztas_c2 = 1.1; # (1+u) * (sigma^2), where u=0.1, sigma=1
        lztas_c3 = dim*log(sqrt(11)) # dim * log(sqrt(u^-1 + 1)), used by DesignType="Heuristic"
        lztas_i0 = 0;
        lztas_w = [1.0/K for i=1:K];
    else # pull each arm once.
        for k = 1:K
            play!(1, k, rng, pep, μ, S, N, Vinv);
        end
    end

    # start sampling rules
    if qbc_sr
        state = start(sr, N, Vxinv)
    elseif convex_sr # keep as the implementations by Xuedong
        state = start(sr, N, P);
    else
        state = start(sr, N);
    end

    t_old = sum(N) # required by XY and RAGE
    R = Tuple{Int64,Array{Int64,1},UInt64}[] # collect return values
    while true
        t = sum(N)
        hµ = Vinv * S  # emp. estimates
        # test stopping criterion
        if xya_sr
            # invoke sampling rule
            i, k, Xactive, ρ, ρ_old, t_old = nextsample(state, pep, N, S, Vinv, Xactive, α, ρ, ρ_old, t_old, δs[1]);
            # test stopping criterion
            while length(Xactive) <= 1
                popfirst!(δs)
                push!(R, (i, copy(N), CPUtime_us() - baseline))
                if isempty(δs)
                    return R
                end
            end
        elseif gap_sr
            _, star, ξ = glrt(pep, N, hμ, Vinv)
            # invoke sampling rule
            i, k, ucb = nextsample(state, pep, star, ξ, N, S, Vinv, βs[1](t))
            while ucb <= 0
                popfirst!(βs)
                push!(R, (star, copy(N), CPUtime_us() - baseline))
                if isempty(βs)
                    return R
                end
            end
        else
            Z, star, ξ = glrt(pep, N, hμ, Vinv);
            while Z > βs[1](t)
                popfirst!(βs)
                push!(R, (star, copy(N), CPUtime_us() - baseline))
                if isempty(βs)
                    return R
                end
            end
            # invoke sampling rule
            if convex_sr
                i, k = nextsample(state, pep, star, ξ, N, P, S, Vinv);
            elseif lztas_sr
                i, k, lztas_i0 = nextsample(state, pep, N, S, Vinv, lztas_A, lztas_A0, lztas_i0, lztas_c0);
                lztas_A += pep.arms[k]*(pep.arms[k]');
            else
                i, k = nextsample(state, pep, N, S, Vinv)
            end
        end
        # play the choosen arm
        play!(i, k, rng, pep, μ, S, N, Vinv);
        convex_sr ? P[i, k] += 1 : nothing;
        t += 1;
    end
end


text:
using Distributions
using Random
using BenchmarkTools
using Plots

function sellke(u0, p)
    (S, I, R) = u0
    N = S + I + R
    (β, c, γ) = p
    λ = β*c/N

    Q = rand(Exponential(), S)
    sort!(Q)

    T0 = rand(Exponential(1/γ), I)
    T = rand(Exponential(1/γ), S)

    ST0 = sum(T0)
    Y = [ST0; ST0 .+ cumsum(T[1:end-1])]

    Z = findfirst(Q .> Y*λ)

    if isnothing(Z)
        Z = S + I # entire population infected
    else
        Z = Z + I - 1
    end

    TT = [T0; T] # all infectious periods
    QQ = [T0 * 0; Q] # all thresholds
    R = T[1:I] # recovery times of the initial infectives

    # max num of events possible
    max = I + 2*S

    t = zeros(max)
    St = zeros(max)
    It = zeros(max)
    It[1:I] = (1:I)
    St[1:I] = N .- (1:I)

    tt = 0
    La = 0
    j = I+1
    k = j

    while It[k-1] > 0
        (minR, i) = findmin(R)
        dtprop = minR-tt
        Laprop = La + (λ * It[k-1] * dtprop)
        if j > length(QQ) # only recoveries remain
            R = R[setdiff(1:length(R), i)]
            tt = minR
            t[k] = minR
            It[k] = It[k-1]-1
            St[k] = St[k-1]
            La = Laprop
            k = k+1
        else # infections
            if QQ[j] > Laprop
                R = R[setdiff(1:length(R), i)]
                tt = minR
                t[k] = minR
                It[k] = It[k-1]-1
                St[k] = St[k-1]
                La = Laprop
            else
                tt = tt + ((QQ[j]-La)/(Laprop-La))*dtprop
                La = QQ[j]
                t[k] = tt
                It[k] = It[k-1]+1
                St[k] = St[k-1]-1
                R = [R; tt+TT[j]]
                j = j+1
            end
            k = k+1
        end
    end

    trajectory = hcat(t, St, It, N .- (St + It))
    trajectory = trajectory[I:k-1, :]

    return trajectory
end

u0 = [990,10,0]; # S,I,R

p = [0.05,10.0,0.25]; # β,c,γ

Random.seed!(1234);

out = sellke(u0, p)

plot(out[:, 1], out[:, 2:end],
    label=["S" "I" "R"],
    xlabel="Time",
    ylabel="Number")

@benchmark sellke(u0, p)


text:
module CDSAPI

## Base Modules Used
using Base64
using Dates
using Logging
using Printf

## Modules Used
using HTTP
using JSON

export retrieve, cdskeys

"""
    retrieve(
        dataset::AbstractString,
        params::AbstractDict,
        fname::AbstractString,
        keys::AbstractDict = cdskeys()
    )

Retrieves datasets from the Climate Data Store, with options specified in a Julia Dictionary and saves it into a specified file.

Arguments:
    * `dataset::AbstractString` : string specifies the name of the dataset within the Climate Data Store that the `retrieve` function is attempting to retrieve data from
    * `params::AbstractDict` : dictionary that contains the keywords that specify the properties (e.g. date, resolution, grid) of the data being retrieved
    * `fname::AbstractString` : string that contains the path and name of the file that the data is to be saved into
    * `keys::AbstractDict` : dictionary that contains API Key information read from the .cdsapirc file in the home directory (optional)
"""
function retrieve(
    dataset::AbstractString,
    params::AbstractDict,
    fname::AbstractString,
    keys::AbstractDict = cdskeys();
    verbose::Bool=true
)

    if verbose
        @info "$(now()) - Welcome to the Climate Data Store"
    else
        @info """$(now()) - Requesting $(uppercase(dataset)) data from the Climate Data Store ...
        """
    end
    apikey = string("Basic ", base64encode(keys["key"]))

    if verbose
        @info "$(now()) - Sending request to https://cds.climate.copernicus.eu/api/v2/resources/$(dataset) ..."
    end
    response = HTTP.request(
        "POST", keys["url"] * "/resources/$(dataset)",
        ["Authorization" => apikey],
        body=JSON.json(params),
        verbose=0
    )
    resp_dict = JSON.parse(String(response.body))
    data = Dict("state" => "queued")

    if verbose; @info "$(now()) - Request is queued" end
    while data["state"] == "queued"
        parserequest!(data)
    end

    if verbose; @info "$(now()) - Request is running" end
    while data["state"] == "running"
        parserequest!(data)
    end

    if data["state"] == "completed"

        if verbose; @info "$(now()) - Request is completed" end

        @info """$(now()) - Downloading $(uppercase(dataset)) data ...
          URL:         $(data["location"])
          Destination: $(fnc)
        """

        dt1 = now()
        HTTP.download(data["location"],fname,update_period=Inf)
        dt2 = now()

        if verbose
            @info "$(now()) - Downloaded $(@sprintf("%.1f",data["content_length"]/1e6)) MB in $(@sprintf("%.1f",Dates.value(dt2-dt1)/1000)) seconds (Rate: $(@sprintf("%.1f",data["content_length"]/1e3/Dates.value(dt2-dt1))) MB/s)"
        end

    elseif data["state"] == "failed"

        @error "$(now()) - Request failed"

    end

    return

end

"""
    cdskeys() -> Dict{Any,Any}

Retrieves the CDS API Keys from the .cdsapirc file in the home directory
"""
function cdskeys()

    keys = Dict(); cdsapirc = joinpath(homedir(),".cdsapirc")

    @info "$(now()) - Loading CDSAPI credentials from $(cdsapirc) ..."
    open(cdsapirc) do f
        for line in readlines(f)
            key,val = strip.(split(line,':',limit=2))
            keys[key] = val
        end
    end

    return keys

end

"""
    parserequest(data::AbstractDict) -> Dict{Any,Any}

Get info on HTTP request, and parse the information and update the dictionary
"""
function parserequest!(data::AbstractDict)

    data = HTTP.request(
        "GET", keys["url"] * "/tasks/" * string(resp_dict["request_id"]),
        ["Authorization" => apikey]
    )
    data = JSON.parse(String(data.body))

    return

end

end # module


text:
module TestConnectance
    using Test
    using EcologicalNetworks
    using LinearAlgebra

    # Generate some data
    N = BipartiteProbabilisticNetwork([0.0 0.1 0.0; 0.2 0.0 0.2; 0.4 0.5 0.0])

    @test links(N) ≈ 1.4
    @test links_var(N) ≈ 0.9
    @test connectance(N) ≈ 1.4 / 9.0
    @test connectance_var(N) ≈ 0.011111111111111111111

    # Once more with a deterministic network
    N = BipartiteNetwork([false true false; false false true; true true true])

    @test links(N) ≈ 5
    @test connectance(N) ≈ 5/9
    @test linkage_density(N) ≈ 5/6

    # Once more with a deterministic unipartite network
    N = UnipartiteNetwork([false true false; false false true; true true true])

    @test links(N) ≈ 5
    @test connectance(N) ≈ 5/9
    @test linkage_density(N) ≈ 5/3

    # Links with quantitative networks
    K = BipartiteQuantitativeNetwork(rand(Float64, (3,3)))
    @test links(K) == 9

    # Convert to adjacency
    @test connectance(UnipartiteNetwork(Matrix(I, (10,10)))) ≈ connectance(UnipartiteQuantitativeNetwork(Matrix{Int64}(I, (10,10)).*2))

end


text:
using Documenter: doctest
using IndexedDims
using Test

@testset "IndexedDims.jl" begin
    # The doctests fail on x86, so only run them on 64-bit hardware
    Sys.WORD_SIZE == 64 && doctest(IndexedDims)

    @testset "constructor" begin
        @test ==(
            IndexedDimsArray([1 2; 3 4; 5 6]),
            IndexedDimsArray([1 2; 3 4; 5 6], Base.OneTo(2), Base.OneTo(3))
        )
    end

    @testset "standard fallback indexing" begin
        data = [1 2 3; 4 5 6; 7 8 9]
        a = IndexedDimsArray(data, 2:4, 2:4)

        @test a[1] == data[1]
        @test a[3, 2] == data[3, 2]
        @test a == data
        @test a[:, :] == data[:, :]
        @test a[:, :] isa IndexedDimsArray
        @test a[:] == data[:]
    end

    @testset "fancy indexing" begin
        @testset "UnitRange" begin
            data = [1 2 3; 4 5 6; 7 8 9]
            a = IndexedDimsArray(data, 2:4, 2:4)

            @test a[:, :] == a
            @test typeof(a[:, :]) == typeof(a)

            @test a[3:4, 2:3] isa IndexedDimsArray
            @test a[3:4, 2:3] == data[2:3, 1:2]
            @test a[3:4, :] == data[2:3, :]
            @test a[3:4, bypass(1:2)] == data[2:3, 1:2]
        end

        @testset "Strings" begin
            data = [1 2 3; 4 5 6; 7 8 9]
            a = IndexedDimsArray(data, ["a", "b", "c"], ["b", "a", "b"])

            @test a[:, :] == a
            @test typeof(a[:, :]) == typeof(a)

            @test a[1] == data[1]
            @test a[3, 2] == data[3, 2]
            @test a[["b", "c"], "a"] == data[[2, 3], 2]
            @test a[["b", "c"], "b"] in (data[[2, 3], 1], data[[2, 3], 3])
            @test a[["b", "c"], ["b"]] == data[[2, 3], [1, 3]]
            @test a["b", ["b", "a"]] == data[2, 1:3]
        end
    end
end


text:
#########################################################################
#
#   zeronode() : Builds the graph generating the diff accumulator starting node
#
#########################################################################
#
#   - Generated graph contains an External node :tv that should be mapped to the source node
#
#########################################################################

function zeronode(n)
    v = n.val

    if isa(v, Union{Real, Symbol, DataType, TypeConstructor, Function, Module})
        return tograph( :(0.) )

    elseif isa(v, Range)
        return tograph( :( zeros(2) ) )

    elseif isa(v, Array) && (eltype(v) <: Real)  # is it an array of Reals ?
        return tograph( :( zeros(size(tv)) ) )

    elseif isa(v, BitArray)  # is it an array of bits ?
        return tograph( :( zeros(size(tv)) ) )

    elseif isa(v, Tuple) && all( map( x -> typeof(x) <: Real, v ) ) # is it a Tuple of Reals ?
        return tograph( :( zeros(length(tv)) ) )

    elseif (isa(v, Array) && isleaftype(eltype(v))) ||  # array of concrete type ?
           (isa(v, Array) && (eltype(v) == Any)) && (length(v) > 20) # large cell Array
           # IMPORTANT : if longer than 20, we will presume that elements
           #  of cell arrays are identical in structure

        # build element constructor
        n2 = NConst(:abcd, [], [], v[1], false)
        ge = zeronode(n2)

        # build loop sub-graph
        fg = ExGraph()
        ni  = addnode!(fg, NExt(:i))  ; fg.exti[ni] = :i    # loop index
        nv  = addnode!(fg, NExt(:v))  ; fg.exti[nv] = :v    # array to fill
        nt  = addnode!(fg, NExt(:tv)) ; fg.exti[nt] = :tv   # ref node
        nt2 = addnode!(fg, NRef(:getidx, [nt, ni]))         # ref node[i]

        nr = addgraph!(ge, fg, Dict( :tv => nt2))
        ns = addnode!(fg, NSRef(:setidx, [nv, nr, ni]) )
        fg.seti[ns] = :v

        # build final graph
        g  = tograph( :( cell( size(tv)) ) )
        nv = getnode(g.seti, nothing) ; fg.exto[nv] = :v
        nt = addnode!(g, NExt(:tv))   ; fg.exto[nt] = :tv ; g.exti[nt] = :tv
        nr = addgraph!( :( 1:length(nv) ), g, Dict( :nv => nv) )
        nf = addnode!(g, NFor( Any[:i, fg], [ nr, nv, nt ]) )
        ns = addnode!(g, NIn( :v, [ nf ]) ) ; fg.seto[ns] = :v
        g.seti[ns] = nothing

        return g

    elseif (isa(v, Array) && (eltype(v) == Any)) ||  # small cell array (presumably coming from a type)
            isa(v, Tuple) # or tuple
        g  = tograph( :( cell( $(length(v)) ) ) )
        nv = addnode!(g, NExt(:tv)) ; g.exti[nv] = :tv
        # TODO : optimize to an array{Float64} instead of array{Any} if all fields are Reals

        for i in 1:length(v)
            ni      = addnode!(g, NConst(i))
            nf      = addnode!(g, NRef(:getidx, [ getnode(g.exti, :tv), ni ], [], v[i], false) )

            ng      = zeronode( nf )
            nn      = addgraph!(ng, g, Dict( :tv => nf ))
            ns      = addnode!(g, NSRef(:setidx, [getnode(g.seti, nothing), nn, ni]))
            g.seti[ns] = nothing
        end

        return g

    elseif isleaftype(typeof(v)) # composite type
        g  = tograph( :( cell( $(length(fieldnames(v))) ) ) )
        nv = addnode!(g, NExt(:tv)) ; g.exti[nv] = :tv
        # TODO : optimize to an array{Float64} instead of array{Any} if all fields are Reals

        for (i, n2) in enumerate(fieldnames(typeof(v)))  # i, n2 = 1, :val
            # create node for holding field value
            nf      = addnode!(g, NDot(QuoteNode(n2), [ getnode(g.exti, :tv) ], [], getfield(v, n2), false) )

            ng      = zeronode( nf )
            nn      = addgraph!(ng, g, Dict( :tv => nf ))
            ni      = addnode!(g, NConst(i))
            ns      = addnode!(g, NSRef(:setidx, [getnode(g.seti, nothing), nn, ni]))
            g.seti[ns] = nothing
        end

        return g

    else
        error("[zeronode] Unable to build diff accumulator for node $(repr(n)[1:min(40, end)])")

    end
end


text:
using Documenter, GaussianRandomFields

DocMeta.setdocmeta!(GaussianRandomFields, :DocTestSetup, :(using GaussianRandomFields); recursive=true)

makedocs(
    sitename="GaussianRandomFields.jl",
    modules = [GaussianRandomFields],
    format = Documenter.HTML(
        prettyurls = get(ENV, "CI", nothing) == "true"
    ),
    pages = [
             "Home" => "index.md",
             "Tutorial" => "tutorial.md",
             "API" => "API.md"
    ],
    doctest=false,
    checkdocs=:exports
)

deploydocs(
    repo = "github.com/PieterjanRobbe/GaussianRandomFields.jl.git",
)


text:
using ModelConstructors, Test

@testset "Prior computation" begin
    p = ParameterVector{Float64}(undef, 3)
    p[1] = ModelConstructors.parameter(:a, 0.1, (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Normal(0.5, 1.); fixed = false)
    p[2] = ModelConstructors.parameter(:b, 0., (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Normal(0.5, 1.); fixed = true)
    p[3] = ModelConstructors.parameter(:c, 0.4, (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Normal(0.25, 1.); fixed = false)

    @test ModelConstructors.prior(p) == logpdf(p[1]) + logpdf(p[3])
end

@testset "Posterior computation" begin
    p = ParameterVector{Float64}(undef, 3)
    p[1] = ModelConstructors.parameter(:a, 0.1, (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Normal(0.5, 1.); fixed = false)
    p[2] = ModelConstructors.parameter(:b, 0., (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Normal(0.5, 1.); fixed = true)
    p[3] = ModelConstructors.parameter(:c, 0.4, (0., 1.), (0., 1.), ModelConstructors.Untransformed(), Normal(0.25, 1.); fixed = false)

    loglh = (x, data) -> 2.
    data = rand(2,2)

    @test posterior(loglh, p, data, ϕ_smc = 3.104) == 2. * 3.104 + ModelConstructors.prior(p)

    oldvals = map(x -> x.value, p)
    oldpost = 2. * 3.104 + ModelConstructors.prior(p)
    @test posterior!(loglh, p, oldvals, data, ϕ_smc = 3.104) == oldpost

    newvals = copy(oldvals)
    newvals[1] += 0.4 # moves value to the mean
    newvals[3] -= 0.15  # moves value to the mean
    @test posterior!(loglh, p, newvals, data, ϕ_smc = 3.104) > oldpost
end


text:
@testset "Quick" begin
    qm = QuickMDP(statetype=Int, actiontype=Int)
    @test @inferred(statetype(qm)) == Int
    @test @inferred(actiontype(qm)) == Int
    @test @inferred(discount(qm)) == 1.0
    @test @inferred(isterminal(qm, 1)) == false
    @test_throws MissingQuickArgument transition(qm, 1, 1)
    @test_throws MissingQuickArgument initialstate(qm)
    @test_throws MissingQuickArgument reward(qm, 1, 1, 1)
    @test_throws MissingQuickArgument initialstate(qm)
    @test_throws MissingQuickArgument states(qm)
    @test_throws MissingQuickArgument actions(qm)
    @test_throws MissingQuickArgument stateindex(qm, 1)
    @test_throws MissingQuickArgument actionindex(qm, 1)

    qp = QuickPOMDP(obstype=Int)
    @test @inferred(obstype(qp)) == Int
    @test @inferred(statetype(qp)) == Any
    @test @inferred(actiontype(qp)) == Any
    @test @inferred(discount(qp)) == 1.0
    @test @inferred(isterminal(qp, 1)) == false
    @test_throws MissingQuickArgument transition(qp, 1, 1)
    @test_throws MissingQuickArgument observation(qp, 1, 1)
    @test_throws MissingQuickArgument initialstate(qp)
    @test_throws MissingQuickArgument reward(qp, 1, 1, 1)
    @test_throws MissingQuickArgument initialstate(qp)
    @test_throws MissingQuickArgument initialobs(qp, 1)
    @test_throws MissingQuickArgument states(qp)
    @test_throws MissingQuickArgument actions(qp)
    @test_throws MissingQuickArgument observations(qp)
    @test_throws MissingQuickArgument stateindex(qp, 1)
    @test_throws MissingQuickArgument actionindex(qp, 1)
    @test_throws MissingQuickArgument obsindex(qp, 1)

    @test_throws MissingQuickArgument @gen(:sp,:o,:r)(qp, 1, 1, MersenneTwister(2))

    qp = QuickMDP(states=[3,2,1], actions=[1,2,3])
    @test ordered_states(qp) == [3,2,1]
    @test ordered_actions(qp) == [1,2,3]
end

@testset "Mountan Car" begin
    mountaincar = begin
        include("../examples/mountaincar.jl")
    end

    @show statetype(mountaincar)

    energize = FunctionPolicy(
        function (s)
            if s[2] < 0.0
                return minimum(actions(mountaincar))
            else
                return maximum(actions(mountaincar))
            end
        end
    )
    sim = HistoryRecorder(max_steps=1000)
    # below should work with a new version of POMDPSimulators
    @test_skip simulate(sim, mountaincar, energize) != nothing
    hist = simulate(sim, mountaincar, energize, rand(initialstate(mountaincar)))
    @test maximum(hist[:r]) == 100.0
    @test all(isa.(hist[:s], statetype(mountaincar)))
    @test last(collect(hist[:sp]))[1] > 0.5
end

@testset "Mountan Car Visualization" begin
    m = begin
        include("../examples/mountaincar_with_visualization.jl")
    end

    @inferred @gen(:sp,:r)(m, (0.2, 0.0), 0.0, MersenneTwister(2))

    energize = FunctionPolicy(
        function (s)
            if s[2] < 0.0
                return minimum(actions(m))
            else
                return maximum(actions(m))
            end
        end
    )
    sim = HistoryRecorder(max_steps=1000)
    # below should work when POMDPSimulators is fixed, but it can be deleted then
    @test_skip simulate(sim, m, energize) != nothing
    hist = simulate(sim, m, energize, rand(initialstate(m)))
    @test maximum(hist[:r]) == 100.0
    @test all(isa.(hist[:s], statetype(mountaincar)))
    @test last(collect(hist[:sp]))[1] > 0.5
    draw(SVG("test_render.svg", 5cm, 4cm), render(m, (s=(0.2, 0.0),)))
    @test isfile("test_render.svg")
end

@testset "Simple Light-Dark" begin
    m = begin
        include("../examples/lightdark.jl")
    end
    p = RandomPolicy(m, rng = MersenneTwister(2))
    sim = HistoryRecorder(max_steps=1000)
    up = DiscreteUpdater(m)
    hist = simulate(sim, m, p, up)
    @test all(-100 .<= hist[:r] .<= 100)
    @test all(isa.(hist[:s], statetype(m)))
    @test last(collect(hist[:sp])) == 61
end

struct A end
@testset "preprocess" begin
    QuickPOMDPs.preprocess(::Type{A}) = Char
    qm = QuickMDP(statetype=A, actiontype=Int)
    @test statetype(qm) == Char
    qp = QuickPOMDP(statetype=A, actiontype=Int, obstype=Int)
    @test statetype(qp) == Char
end

@testset "state action reward" begin
    qm = QuickMDP(states=1:3, actions=1:3, reward=(s, a)->0.0)
    @test StateActionReward(qm) isa FunctionSAR
    qm = QuickMDP(states=1:3, actions=1:3, reward=(s, a, sp)->0.0)
    @test StateActionReward(qm) isa LazyCachedSAR

    qp = QuickPOMDP(states=1:3, actions=1:3, observations=1:3, reward=(s, a)->0.0)
    @test StateActionReward(qp) isa FunctionSAR
    qp = QuickPOMDP(states=1:3, actions=1:3, observations=1:3, reward=(s, a, sp, o)->0.0)
    @test StateActionReward(qp) isa LazyCachedSAR
end

@testset "variable args" begin
    m = QuickMDP(states=1:3, actions=1:3, reward=(s, a)->0.0)
    @test @inferred(reward(m, 1, 2)) == 0.0
    @test @inferred(reward(m, 1, 2, 3)) == 0.0
    m = QuickPOMDP(states=1:3, actions=1:3, observations=1:3, reward=(s, a)->0.0)
    @test @inferred(reward(m, 1, 2)) == 0.0
    @test @inferred(reward(m, 1, 2, 3)) == 0.0
    @test @inferred(reward(m, 1, 2, 3, 1)) == 0.0
    m = QuickPOMDP(states=1:3, actions=1:3, observations=1:3, observation=(a, sp)->Deterministic(sp))
    @test rand(@inferred(observation(m, 2, 3))) == 3
    @test rand(@inferred(observation(m, 1, 2, 3))) == 3
end

@testset "#23" begin
    @test_throws ArgumentError QuickPOMDP(initialstate=:test)
end


text:
module TestLearningNetworks

# using Revise
using Test
using MLJBase
using ..Models
using CategoricalArrays
import Random.seed!
seed!(1234)

@load KNNRegressor

@testset "network #1" begin

    N =100
    X = (x1=rand(N), x2=rand(N), x3=rand(N))
    y = 2X.x1  - X.x2 + 0.05*rand(N)

    knn_ = KNNRegressor(K=7)

    # split the rows:
    allrows = eachindex(y);
    train, valid, test = partition(allrows, 0.7, 0.15);
    @test vcat(train, valid, test) == allrows

    Xtrain = selectrows(X, train)
    ytrain = y[train]

    Xs = source(Xtrain)
    ys = source(ytrain)

    knn1 = machine(knn_, Xs, ys)
    @test_logs (:info, r"Training") fit!(knn1, verbosity=3)
    knn_.K = 5
    @test_logs((:info, r"Training"),
               fit!(knn1, rows=train[1:end-10], verbosity=2))
    @test_logs (:info, r"Training") fit!(knn1, verbosity=2)
    yhat = predict(knn1, Xs)
    yhat(selectrows(X, test))
    @test rms(yhat(selectrows(X, test)), y[test]) < 0.3
    @test MLJBase.is_stale(knn1) == false

    fit!(yhat)
    pred = yhat()

    # test serialization of NodalMachine:
    io = IOBuffer()
    MLJBase.save(io, knn1)
    seekstart(io)
    mach = fit!(machine(io, Xtrain, ytrain))
    @test predict(mach) ≈ pred

end

@testset "network #2" begin

    N =100
    X = (x1=rand(N),
         x2=rand(N),
         x3=categorical(rand("yn",N)),
         x4=categorical(rand("yn",N)))

    y = 2X.x1  - X.x2 + 0.05*rand(N)
    X = source(X)
    y = source(y)

    hot = OneHotEncoder()
    hotM = machine(hot, X)
    W = transform(hotM, X)
    knn = KNNRegressor()
    knnM = machine(knn, W, y)
    yhat = predict(knnM, W)

    # should get "Training" for both:
    @test_logs (:info, r"^Training") (:info, r"^S") (:info, r"^S") (:info, r"^Training") fit!(yhat)

    # should get "Not retraining" for both:
    @test_logs (:info, r"^Not retraining") (:info, r"^Not retraining") fit!(yhat)

    # should get "Updating" for first, "Training" for second:
    hot.drop_last = true
    @test_logs (:info, r"^Updating")  (:info, r"^S") (:info, r"^S") (:info, r"^Training") fit!(yhat)

    # should get "Not retraining" for both:
    @test_logs (:info, r"^Not retraining") (:info, r"^Not retraining") fit!(yhat)

    # should get "Not retraining" for first, "Updating for second":
    knn.K = 17
    @test_logs (:info, r"^Not retraining") (:info, r"^Updating") fit!(yhat)

    # should get "Training" for both:
    @test_logs (:info, r"^Training")  (:info, r"^S") (:info, r"^S") (:info, r"^Training") fit!(yhat, rows=1:100)

    # should get "Training" for both"
    @test_logs (:info, r"^Training")  (:info, r"^S") (:info, r"^S") (:info, r"^Training") fit!(yhat)

    # should get "Training" for both"
    @test_logs (:info, r"^Training")  (:info, r"^S") (:info, r"^S") (:info, r"^Training") fit!(yhat, force=true)

end

@testset "network #3" begin

    N =100
    X = (x1=rand(N), x2=rand(N), x3=rand(N))
    y = 2X.x1  - X.x2 + 0.05*rand(N)

    XX = source(X)
    yy = source(y)

    # construct a transformer to standardize the target:
    uscale_ = UnivariateStandardizer()
    uscale = machine(uscale_, yy)

    # get the transformed inputs, as if `uscale` were already fit:
    z = transform(uscale, yy)

    # construct a transformer to standardize the inputs:
    scale_ = Standardizer()
    scale = machine(scale_, XX) # no need to fit

    # get the transformed inputs, as if `scale` were already fit:
    Xt = transform(scale, XX)

    # do nothing to the DataFrame
    Xa = node(identity, Xt)

    # choose a learner and make it machine:
    knn_ = KNNRegressor(K=7) # just a container for hyperparameters
    knn = machine(knn_, Xa, z) # no need to fit

    # get the predictions, as if `knn` already fit:
    zhat = predict(knn, Xa)

    # inverse transform the target:
    yhat = inverse_transform(uscale, zhat)

    # fit-through training:
    @test_logs((:info, r"Training"),
               (:info, r"Features standarized: "),
               (:info, r" *:x1"),
               (:info, r" *:x2"),
               (:info, r" *:x3"),
               (:info, r"Training"),
               (:info, r"Training"),
               fit!(yhat, rows=1:50, verbosity=2))
    @test_logs(
        (:info, r"Not retraining"),
        (:info, r"Not retraining"),
        (:info, r"Not retraining"),
        fit!(yhat, rows=1:50, verbosity=1))
    @test_logs(
        (:info, r"Training"),
        (:info, r"Training"),
        (:info, r"Training"),
        fit!(yhat, verbosity=1))
    knn_.K =67
    @test_logs(
        (:info, r"Not retraining"),
        (:info, r"Not retraining"),
        (:info, r"Updating"),
        fit!(yhat, verbosity=1))

end

@testset "overloading methods for AbstractNode" begin
    A  = rand(3,7)
    As = source(A)
    @test MLJBase.matrix(MLJBase.table(As))() == A

    X = (x1 = [1,2,3], x2=[10, 20, 30], x3=[100, 200, 300])
    Xs = source(X)
    @test selectrows(Xs, 1)() == selectrows(X, 1)
    @test selectrows(Xs, 2:3)() == selectrows(X, 2:3)
    @test selectcols(Xs, 1)() == selectcols(X, 1)
    @test selectcols(Xs, 2:3)() == selectcols(X, 2:3)
    @test selectcols(Xs, :x1)() == selectcols(X, :x1)
    @test selectcols(Xs, [:x1, :x3])() == selectcols(X, [:x1, :x3])

    y = rand(4)
    ys = source(y)
    @test vcat(ys, ys)() == vcat(y, y)
    @test hcat(ys, ys)() == hcat(y, y)
    @test log(ys)() == log.(y)
    @test exp(ys)() == exp.(y)

    Z = (rand(4), rand(4), rand(4))
    Zs = source(Z)
    @test mean(Zs)() == mean.(Z)
    @test mode(Zs)() == mode.(Z)
    @test median(Zs)() == median.(Z)

    a, b, λ = rand(4), rand(4), rand()
    as, bs = source(a), source(b)
    @test (as + bs)() == a + b
    @test (λ * bs)() == λ * b
end

struct MyTransformer1 <: Static
    ftr::Symbol
end

MLJBase.transform(transf::MyTransformer1, verbosity, X) =
    selectcols(X, transf.ftr)

@testset "nodal machine constructor for static transformers" begin
    X = (x1=rand(3), x2=[1, 2, 3])
    mach = machine(MyTransformer1(:x2))
    fit!(mach)
    @test transform(mach, X) == [1, 2, 3]
end

## MORE STATIC TRANSFORMER TESTS

x1 = rand(30)
x2 = rand(30)
x3 = rand(30)
y = exp.(x1 - x2 -2x3 + 0.1*rand(30))
X = (x1=x1, x2=x2, x3=x3)

f(X) = (a=selectcols(X, :x1), b=selectcols(X, :x2))

knn = @load KNNRegressor

# 1. function in a pipeline:
comp1 = @pipeline Comp1(f, std=Standardizer(),
                        rgs=knn,
                        target=UnivariateBoxCoxTransformer())
e = evaluate(comp1, X, y, measure=mae, resampling=Holdout(), verbosity=0)

# 2. function with parameters in a pipeline:
mutable struct MyTransformer <: Static
    ftr::Symbol
end
MLJBase.transform(transf::MyTransformer, verbosity, X) =
    (a=selectcols(X, transf.ftr), b=selectcols(X, :x2))
comp2 = @pipeline Comp2(transf=MyTransformer(:x3), std=Standardizer(),
                      rgs=knn,
                      target=UnivariateBoxCoxTransformer())

comp2.transf.ftr = :x1 # change the parameter
e2 =  evaluate(comp2, X, y, measure=mae, resampling=Holdout(), verbosity=0)
@test e2.measurement[1] ≈ e.measurement[1]

# 3. function in an exported learning network:
Xs = source(X)
ys = source(y, kind=:target)
f(X::AbstractNode) = node(f, X)
# or, without arrow syntax:
# W = Xs |> f |> Standardizer()
# z = ys |> UnivariateBoxCoxTransformer()
# zhat = (W, z) |> knn
# yhat = zhat |> inverse_transform(z)
X2 = f(Xs)
W = transform(machine(Standardizer(), X2), X2)
box_mach = machine(UnivariateBoxCoxTransformer(), ys)
z = transform(box_mach, ys)
knn_mach = machine(knn, W, z)
zhat = predict(knn_mach, W)
yhat = inverse_transform(box_mach, zhat)
comp3 = @from_network Comp3(rgs=knn) <= yhat
e3 =  evaluate(comp3, X, y, measure=mae, resampling=Holdout(), verbosity=0)
@test e2.measurement[1] ≈ e.measurement[1]

# 4. function with parameters in exported learning network:
inserter = MyTransformer(:x3)
# W = Xs |> inserter |> Standardizer()
# z = ys |> UnivariateBoxCoxTransformer()
# zhat = (W, z) |> knn
# yhat = zhat |> inverse_transform(z)
inserter_mach = machine(inserter)
X2 = transform(inserter_mach, Xs)
W = transform(machine(Standardizer(), X2), X2)
box_mach = machine(UnivariateBoxCoxTransformer(), ys)
z = transform(box_mach, ys)
knn_mach = machine(knn, W, z)
zhat = predict(knn_mach, W)
yhat = inverse_transform(box_mach, zhat)
comp4 = @from_network Comp4(transf=inserter, rgs=knn) <= yhat
comp4.transf.ftr = :x1 # change the parameter
e4 =  evaluate(comp4, X, y, measure=mae, resampling=Holdout(), verbosity=0)
@test e4.measurement[1] ≈ e.measurement[1]

end

true