using DiffEqFlux, DelayDiffEq, OrdinaryDiffEq, StochasticDiffEq, Test, Random

mp = Float32[0.1,0.1]
x = Float32[2.; 0.]
xs = Float32.(hcat([0.; 0.], [1.; 0.], [2.; 0.]))
tspan = (0.0f0,1.0f0)
dudt = Flux.Chain(Flux.Dense(2,50,tanh),Flux.Dense(50,2))
fastdudt = FastChain(FastDense(2,50,tanh),FastDense(50,2))
fastcdudt = FastChain(FastDense(2,50,tanh,precache=true,numcols=size(xs)[2]),FastDense(50,2,precache=true,numcols=size(xs)[2]))
staticdudt = FastChain(StaticDense(2,50,tanh),StaticDense(50,2))

NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false)(x)
NeuralODE(dudt,tspan,Tsit5(),saveat=0.1)(x)
NeuralODE(dudt,tspan,Tsit5(),saveat=0.1,sensealg=TrackerAdjoint())(x)

NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false)(xs)
NeuralODE(dudt,tspan,Tsit5(),saveat=0.1)(xs)
NeuralODE(dudt,tspan,Tsit5(),saveat=0.1,sensealg=TrackerAdjoint())(xs)

@info "Test some gradients"

node = NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false)
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test ! iszero(grads[xs])
@test ! iszero(grads[node.p])

node = NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=TrackerAdjoint())
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test ! iszero(grads[xs])
@test ! iszero(grads[node.p])

node = NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=BacksolveAdjoint())
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test ! iszero(grads[xs])
@test ! iszero(grads[node.p])

node = NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=BacksolveAdjoint(autojacvec=ZygoteVJP()))
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test ! iszero(grads[xs])
@test ! iszero(grads[node.p])

## Fast

@info "Test some fast layers"

node = NeuralODE(fastdudt,tspan,Tsit5(),save_everystep=false,save_start=false)
pd = Flux.params(node)[1]
gradsnc = Zygote.gradient(()->sum(node(x)),Flux.params(x,node)) # no cache
@test ! iszero(gradsnc[x])
@test ! iszero(gradsnc[node.p])

gradsnc2 = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test ! iszero(gradsnc2[xs])
@test ! iszero(gradsnc2[node.p])

nodec = NeuralODE(fastcdudt,tspan,Tsit5(),save_everystep=false,save_start=false,p=pd)
gradsc = Zygote.gradient(()->sum(nodec(x)),Flux.params(x,nodec)) #with cache
@test ! iszero(gradsc[x])
@test ! iszero(gradsc[nodec.p])
@test gradsnc[x] ≈ gradsc[x] rtol=1e-6
@test gradsnc[node.p] ≈ gradsc[nodec.p] rtol=1e-6

gradsc2 = Zygote.gradient(()->sum(nodec(xs)),Flux.params(xs,nodec))
@test ! iszero(gradsc2[xs])
@test ! iszero(gradsc2[nodec.p])
@test gradsnc2[xs] ≈ gradsc2[xs] rtol=1e-6
@test gradsnc2[node.p] ≈ gradsc2[nodec.p] rtol=1e-6
#test with low tolerance ode solver
node = NeuralODE(fastdudt, tspan, Tsit5(), abstol=1e-12, reltol=1e-12, save_everystep=false, save_start=false)
pd = Flux.params(node)[1]
gradsnc = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(gradsnc[x])
@test ! iszero(gradsnc[node.p])

nodec = NeuralODE(fastcdudt, tspan, Tsit5(), abstol=1e-12, reltol=1e-12, save_everystep=false, save_start=false,p=pd)
gradsc = Zygote.gradient(()->sum(nodec(x)),Flux.params(x,nodec))
@test gradsnc[x] ≈ gradsc[x] rtol=1e-3
@test gradsnc[node.p] ≈ gradsc[nodec.p] rtol=1e-3

node = NeuralODE(fastdudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=TrackerAdjoint())
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test ! iszero(grads[xs])
@test ! iszero(grads[node.p])

goodgrad = grads[node.p]
p = node.p

node = NeuralODE(fastcdudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=TrackerAdjoint())
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test ! iszero(grads[xs])
@test ! iszero(grads[node.p])

goodgradc = grads[node.p]
pc = node.p

node = NeuralODE(fastdudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=BacksolveAdjoint(),p=p)
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test ! iszero(grads[xs])
@test ! iszero(grads[node.p])
goodgrad2 = grads[node.p]
@test goodgrad ≈ goodgrad2 # Make sure adjoint overloads are correct

node = NeuralODE(fastcdudt,tspan,Tsit5(),save_everystep=false,save_start=false, sensealg=BacksolveAdjoint(),p=pc)
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test !iszero(grads[xs])
@test ! iszero(grads[node.p])
goodgrad2 = grads[node.p]
@test goodgradc ≈ goodgrad2 rtol=1e-6

node = NeuralODE(staticdudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=BacksolveAdjoint(autojacvec=false),p=p)
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
goodgrad2 = grads[node.p]
@test goodgrad ≈ goodgrad2 rtol = 1e-6

node = NeuralODE(staticdudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=BacksolveAdjoint(autojacvec=ZygoteVJP()),p=p)
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test !iszero(grads[x])
@test !iszero(grads[node.p])

@test_throws ErrorException grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))

@info "Test some adjoints"

# Adjoint
@testset "adjoint mode" begin
    node = NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false)
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
    @test ! iszero(grads[xs])
    @test ! iszero(grads[node.p])

    node = NeuralODE(dudt,tspan,Tsit5(),saveat=0.0:0.1:1.0)
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
    @test ! iszero(grads[xs])
    @test ! iszero(grads[node.p])

    node = NeuralODE(dudt,tspan,Tsit5(),saveat=0.1)
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
    @test ! iszero(grads[xs])
    @test ! iszero(grads[node.p])

    node = NeuralODE(fastdudt,tspan,Tsit5(),save_everystep=false,save_start=false)
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    @test_broken ! iszero(grads[xs])
    @test_broken ! iszero(grads[node.p])

    node = NeuralODE(fastcdudt,tspan,Tsit5(),save_everystep=false,save_start=false)
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    @test_broken ! iszero(grads[xs])
    @test_broken ! iszero(grads[node.p])

    node = NeuralODE(fastdudt,tspan,Tsit5(),saveat=0.0:0.1:1.0)
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    @test_broken ! iszero(grads[xs])
    @test_broken ! iszero(grads[node.p])

    node = NeuralODE(fastcdudt,tspan,Tsit5(),saveat=0.0:0.1:1.0)
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    @test_broken ! iszero(grads[xs])
    @test_broken ! iszero(grads[node.p])

    node = NeuralODE(fastdudt,tspan,Tsit5(),saveat=0.1)
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    @test_broken ! iszero(grads[xs])
    @test_broken ! iszero(grads[node.p])

    node = NeuralODE(fastcdudt,tspan,Tsit5(),saveat=0.1)
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    @test_broken ! iszero(grads[xs])
    @test_broken ! iszero(grads[node.p])
end

@info "Test Tracker"

# RD
@testset "Tracker mode" begin
    node = NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=TrackerAdjoint())
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
    @test ! iszero(grads[xs])
    @test ! iszero(grads[node.p])

    node = NeuralODE(dudt,tspan,Tsit5(),saveat=0.0:0.1:1.0,sensealg=TrackerAdjoint())
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
    @test ! iszero(grads[xs])
    @test ! iszero(grads[node.p])

    node = NeuralODE(dudt,tspan,Tsit5(),saveat=0.1,sensealg=TrackerAdjoint())
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
    @test ! iszero(grads[xs])
    @test ! iszero(grads[node.p])

    node = NeuralODE(fastdudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=TrackerAdjoint())
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    @test_broken ! iszero(grads[xs])
    @test_broken ! iszero(grads[node.p])

    node = NeuralODE(fastcdudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=TrackerAdjoint())
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    @test_broken ! iszero(grads[xs])
    @test_broken ! iszero(grads[node.p])

    node = NeuralODE(fastdudt,tspan,Tsit5(),saveat=0.0:0.1:1.0,sensealg=TrackerAdjoint())
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    @test_broken ! iszero(grads[xs])
    @test_broken ! iszero(grads[node.p])

    node = NeuralODE(fastcdudt,tspan,Tsit5(),saveat=0.0:0.1:1.0,sensealg=TrackerAdjoint())
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    @test_broken ! iszero(grads[xs])
    @test_broken ! iszero(grads[node.p])

    node = NeuralODE(fastdudt,tspan,Tsit5(),saveat=0.1,sensealg=TrackerAdjoint())
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    @test_broken ! iszero(grads[xs])
    @test_broken ! iszero(grads[node.p])

    node = NeuralODE(fastcdudt,tspan,Tsit5(),saveat=0.1,sensealg=TrackerAdjoint())
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    @test_broken ! iszero(grads[xs])
    @test_broken ! iszero(grads[node.p])
end

@info "Test non-ODEs"

dudt2 = Flux.Chain(Flux.Dense(2,50,tanh),Flux.Dense(50,2))
fastdudt2 = FastChain(FastDense(2,50,tanh),FastDense(50,2))
fastcdudt2 = FastChain(FastDense(2,50,tanh,numcols=size(xs)[2],precache=true),FastDense(50,2,numcols=size(xs)[2],precache=true))
NeuralDSDE(dudt,dudt2,(0.0f0,.1f0),SOSRI(),saveat=0.1)(x)
sode = NeuralDSDE(dudt,dudt2,(0.0f0,.1f0),SOSRI(),saveat=0.0:0.01:0.1)

grads = Zygote.gradient(()->sum(sode(x)),Flux.params(x,sode))
@test ! iszero(grads[x])
@test ! iszero(grads[sode.p])
@test ! iszero(grads[sode.p][end])

grads = Zygote.gradient(()->sum(sode(xs)),Flux.params(xs,sode))
@test ! iszero(grads[xs])
@test ! iszero(grads[sode.p])
@test ! iszero(grads[sode.p][end])

sode = NeuralDSDE(fastdudt,fastdudt2,(0.0f0,.1f0),SOSRI(),saveat=0.0:0.01:0.1)
pd = Flux.params(sode)[1]
Random.seed!(1234)
gradsnc = Zygote.gradient(()->sum(sode(x)),Flux.params(x,sode)) #no cacahe
@test ! iszero(gradsnc[x])
@test ! iszero(gradsnc[sode.p])
@test ! iszero(gradsnc[sode.p][end])

gradsnc2 = Zygote.gradient(()->sum(sode(xs)),Flux.params(xs,sode))
@test_broken gradsnc2 isa Tuple
@test ! iszero(gradsnc2[xs])
@test ! iszero(gradsnc2[sode.p])
@test ! iszero(gradsnc2[sode.p][end])

sodec = NeuralDSDE(fastcdudt,fastcdudt2,(0.0f0,.1f0),SOSRI(),saveat=0.0:0.01:0.1,p=pd)
Random.seed!(1234)
gradsc = Zygote.gradient(()->sum(sodec(x)),Flux.params(x,sodec))
@test ! iszero(gradsc[x])
@test ! iszero(gradsc[sodec.p])
@test ! iszero(gradsc[sodec.p][end])
@test gradsnc[x] ≈ gradsc[x] rtol=1e-6
@test gradsnc[sode.p] ≈ gradsc[sodec.p] rtol=1e-6

gradsc2 = Zygote.gradient(()->sum(sodec(xs)),Flux.params(xs,sodec))
@test_broken gradsc2 isa Tuple
@test ! iszero(gradsc2[xs])
@test ! iszero(gradsc2[sodec.p])
@test ! iszero(gradsc2[sodec.p][end])

dudt22 = Flux.Chain(Flux.Dense(2,50,tanh),Flux.Dense(50,4),x->reshape(x,2,2))
fastdudt22 = FastChain(FastDense(2,50,tanh),FastDense(50,4),(x,p)->reshape(x,2,2))
fastcdudt22 = FastChain(FastDense(2,50,tanh,numcols=size(xs)[2],precache=true),FastDense(50,4,numcols=size(xs)[2],precache=true),(x,p)->reshape(x,2,2))
NeuralSDE(dudt,dudt22,(0.0f0,.1f0),2,LambaEM(),saveat=0.01)(x)

sode = NeuralSDE(dudt,dudt22,(0.0f0,0.1f0),2,LambaEM(),saveat=0.0:0.01:0.1)

grads = Zygote.gradient(()->sum(sode(x)),Flux.params(x,sode))
@test ! iszero(grads[x])
@test ! iszero(grads[sode.p])
@test ! iszero(grads[sode.p][end])

@test_broken grads = Zygote.gradient(()->sum(sode(xs)),Flux.params(xs,sode))
@test_broken ! iszero(grads[xs])
@test ! iszero(grads[sode.p])
@test ! iszero(grads[sode.p][end])

sode = NeuralSDE(fastdudt,fastdudt22,(0.0f0,0.1f0),2,LambaEM(),saveat=0.0:0.01:0.1)
pd = Flux.params(sode)[1]
Random.seed!(1234)
gradsnc = Zygote.gradient(()->sum(sode(x)),Flux.params(x,sode))
@test ! iszero(gradsnc[x])
@test ! iszero(gradsnc[sode.p])
@test ! iszero(gradsnc[sode.p][end])

@test_broken gradsnc = Zygote.gradient(()->sum(sode(xs)),Flux.params(xs,sode))
@test_broken ! iszero(gradsnc[xs])
@test ! iszero(gradsnc[sode.p])
@test ! iszero(gradsnc[sode.p][end])

sodec = NeuralSDE(fastcdudt,fastcdudt22,(0.0f0,0.1f0),2,LambaEM(),saveat=0.0:0.01:0.1,p=pd)
Random.seed!(1234)
gradsc = Zygote.gradient(()->sum(sodec(x)),Flux.params(x,sodec))
@test ! iszero(gradsc[x])
@test ! iszero(gradsc[sodec.p])
@test ! iszero(gradsc[sodec.p][end])
@test gradsnc[x] ≈ gradsc[x] rtol=1e-6
@test gradsnc[sode.p] ≈ gradsc[sodec.p] rtol=1e-6

@test_broken gradsc = Zygote.gradient(()->sum(sodec(xs)),Flux.params(xs,sodec))
@test_broken ! iszero(gradsc[xs])
@test ! iszero(gradsc[sodec.p])
@test ! iszero(gradsc[sodec.p][end])

ddudt = Flux.Chain(Flux.Dense(6,50,tanh),Flux.Dense(50,2))
NeuralCDDE(ddudt,(0.0f0,2.0f0),(p,t)->zero(x),(1f-1,2f-1),MethodOfSteps(Tsit5()),saveat=0.1)(x)
dode = NeuralCDDE(ddudt,(0.0f0,2.0f0),(p,t)->zero(x),(1f-1,2f-1),MethodOfSteps(Tsit5()),saveat=0.0:0.1:2.0)

grads = Zygote.gradient(()->sum(dode(x)),Flux.params(x,dode))
@test ! iszero(grads[x])
@test ! iszero(grads[dode.p])

@test_broken grads = Zygote.gradient(()->sum(dode(xs)),Flux.params(xs,dode)) isa Tuple
@test_broken ! iszero(grads[xs])
@test ! iszero(grads[dode.p])

fastddudt = FastChain(FastDense(6,50,tanh),FastDense(50,2))
NeuralCDDE(fastddudt,(0.0f0,2.0f0),(p,t)->zero(x),(1f-1,2f-1),MethodOfSteps(Tsit5()),saveat=0.1)(x)
dode = NeuralCDDE(fastddudt,(0.0f0,2.0f0),(p,t)->zero(x),(1f-1,2f-1),MethodOfSteps(Tsit5()),saveat=0.0:0.1:2.0)
pd = Flux.params(dode)[1]
gradsnc = Zygote.gradient(()->sum(dode(x)),Flux.params(x,dode))
@test ! iszero(gradsnc[x])
@test ! iszero(gradsnc[dode.p])

@test_broken gradsnc = Zygote.gradient(()->sum(dode(xs)),Flux.params(xs,dode)) isa Tuple
@test_broken ! iszero(gradsnc[xs])
@test ! iszero(gradsnc[dode.p])

fastcddudt = FastChain(FastDense(6,50,tanh,numcols=size(xs)[2],precache=true),FastDense(50,2,numcols=size(xs)[2],precache=true))
NeuralCDDE(fastcddudt,(0.0f0,2.0f0),(p,t)->zero(x),(1f-1,2f-1),MethodOfSteps(Tsit5()),saveat=0.1)(x)
dodec = NeuralCDDE(fastcddudt,(0.0f0,2.0f0),(p,t)->zero(x),(1f-1,2f-1),MethodOfSteps(Tsit5()),saveat=0.0:0.1:2.0,p=pd)

gradsc = Zygote.gradient(()->sum(dodec(x)),Flux.params(x,dodec))
@test ! iszero(gradsc[x])
@test ! iszero(gradsc[dodec.p])
@test gradsnc[x] ≈ gradsc[x] rtol=1e-6
@test gradsnc[dode.p] ≈ gradsc[dodec.p] rtol=1e-6

@test_broken gradsc = Zygote.gradient(()->sum(dodec(xs)),Flux.params(xs,dodec)) isa Tuple
@test_broken ! iszero(gradsc[xs])
@test ! iszero(gradsc[dodec.p])


# The precision of forward-differiation is accurate at Float64 but not at Float32/16.
# Anyway the backward-differentiation is accurate at Float64/32/16.
# Just check at Float64/32/16 while keeping the same random seed.
# For the below example the gradients for CPU/GPU on my device are
# Float64 is 100.724249419
# Float32 is 100.72425
# Float16 is 100.75
# ochannels = 1

@testset "check gradient for PlainConv1d block" begin
    clear()
    # using Random
    # Random.seed!(1234);
    TYPE = Array{Float64};

    # [0] prepare model
    ichannels = 1;
    ochannels = 2;
    c1 = PlainConv1d(ichannels,4,3; type=TYPE);
    c2 = PlainConv1d(4,ochannels,2; type=TYPE);

    # [1] prepare input data and its label
    timeSteps = 128;
    batchsize = 32;
    x = Variable(rand(ichannels,timeSteps,batchsize); type=TYPE);
    l = Variable(rand(ochannels,      125,batchsize); type=TYPE);

    # [2] forward and backward propagation
    o1 = forward(c1,  x);
    o2 = forward(c2, o1);
    COST1 = mseCost(o2, l);
    backward();

    # [3] with a samll change of a weight
    GRAD = c1.w.delta[1];
    DELTA = 1e-6;
    c1.w.value[1] += DELTA;

    # [4] forward and backward propagation
    o1 = forward(c1,  x);
    o2 = forward(c2, o1);
    COST2 = mseCost(o2, l);
    backward();

    # [5] check if the auto-grad is true or not
    dLdW = (COST2 - COST1)/DELTA;   # numerical gradient
    err  = abs((dLdW-GRAD)/(GRAD+eps(Float64)))*100;  # relative error in %
    err  = err < 1e-1 ? 0.0 : err;
    @test err < 0.1
end


using HPAT
using HPAT.API.LinearRegression
#HPAT.CaptureAPI.set_debug_level(3)
#HPAT.DomainPass.set_debug_level(3)
#HPAT.DataTablePass.set_debug_level(3)
#using CompilerTools
#CompilerTools.OptFramework.set_debug_level(3)

#using ParallelAccelerator
#ParallelAccelerator.DomainIR.set_debug_level(3)

@acc hpat function q05(category, education, gender, file_name)
    web_clickstreams = DataSource(DataTable{:wcs_item_sk=Int64, :wcs_user_sk=Int64}, HDF5, file_name)
    item = DataSource(DataTable{:i_item_sk=Int64,:i_category_id=Int64,:i_category=Int64}, HDF5, file_name)
    customer = DataSource(DataTable{:c_customer_sk=Int64,:c_current_cdemo_sk=Int64}, HDF5, file_name)
    customer_demographics = DataSource(DataTable{:cd_demo_sk=Int64,:cd_gender=Int64,:cd_education_status=Int64}, HDF5, file_name)
    # TODO: following two filters are added by spark as an optimization;
    # I am adding manually. Need to add optimization in datatable pass
    web_clickstreams = web_clickstreams[:wcs_user_sk<typemax(Int32)]
    customer = customer[:c_current_cdemo_sk<typemax(Int32)]
    # typemax(Int32) used for NULL; typemin(Int32) was not working
    web_clickstreams = web_clickstreams[:wcs_item_sk<typemax(Int32)]

    user_items = join(web_clickstreams, item, :wcs_item_sk==:i_item_sk, :user_items_sk)

    user_clicks_in_cat = aggregate(user_items, :wcs_user_sk, :clicks_in_category = sum(:i_category==category),
                                                         :clicks_in_1 = sum(:i_category_id==1),
                                                         :clicks_in_2 = sum(:i_category_id==2),
                                                         :clicks_in_3 = sum(:i_category_id==3),
                                                         :clicks_in_4 = sum(:i_category_id==4),
                                                         :clicks_in_5 = sum(:i_category_id==5),
                                                         :clicks_in_6 = sum(:i_category_id==6),
                                                         :clicks_in_7 = sum(:i_category_id==7))
    customer_clicks = join(user_clicks_in_cat, customer, :wcs_user_sk==:c_customer_sk, :customer_clicks_sk)
    customer_demo_clicks = join(customer_clicks, customer_demographics, :c_current_cdemo_sk==:cd_demo_sk , :customer_demo_clicks_sk)
    # To make int array from bool
    college_education = (customer_demo_clicks[:cd_education_status].==education)
    male = (customer_demo_clicks[:cd_gender].==gender)
    responses = 1.0 .*(customer_demo_clicks[:clicks_in_category])
    data = transpose(typed_hcat(Float64, college_education,
            male,
            customer_demo_clicks[:clicks_in_1],
            customer_demo_clicks[:clicks_in_2],
            customer_demo_clicks[:clicks_in_3],
            customer_demo_clicks[:clicks_in_4],
            customer_demo_clicks[:clicks_in_5],
            customer_demo_clicks[:clicks_in_6],
            customer_demo_clicks[:clicks_in_7]))
    #pointsF = convert(Matrix{Float64},data)
    model = LinearRegression(data, responses)
    #return customer_demo_clicks[:customer_demo_clicks_sk], customer_demo_clicks[:cd_gender]
    return model
end

println(q05(3, 8303423, 1, "test_q05.hdf5"))


# Simple ridge regression based upon the sindy-mpc
# repository, see https://arxiv.org/abs/1711.05501
# and https://github.com/eurika-kaiser/SINDY-MPC/blob/master/LICENSE

mutable struct STRRidge{T} <: AbstractOptimizer
    λ::T
end

"""
    STRRidge(λ = 0.1)

`STRRidge` is taken from the [original paper on SINDY](https://www.pnas.org/content/113/15/3932) and implements a
sequentially thresholded least squares iteration. `λ` is the threshold of the iteration.
It is based upon [this matlab implementation](https://github.com/eurika-kaiser/SINDY-MPC/utils/sparsifyDynamics.m).

# Example
```julia
opt = STRRidge()
opt = STRRidge(1e-1)
```
"""
STRRidge() = STRRidge(0.1)

function set_threshold!(opt::STRRidge, threshold)
    opt.λ = threshold
end

get_threshold(opt::STRRidge) = opt.λ

init(o::STRRidge, A::AbstractArray, Y::AbstractArray) = A \ Y
init!(X::AbstractArray, o::STRRidge, A::AbstractArray, Y::AbstractArray) =  ldiv!(X, qr(A, Val(true)), Y)

function fit!(X::AbstractArray, A::AbstractArray, Y::AbstractArray, opt::STRRidge; maxiter::Int64 = 1, convergence_error::T = eps()) where T <: Real
    smallinds = abs.(X) .<= opt.λ
    biginds = @. ! smallinds[:, 1]

    x_i = similar(X)
    x_i .= X

    iters = 0

    for i in 1:maxiter
        iters += 1

        smallinds .= abs.(X) .<= opt.λ
        X[smallinds] .= zero(eltype(X))
        for j in 1:size(Y, 2)
            biginds .= @. ! smallinds[:, j]
            X[biginds, j] .= A[:, biginds] \ Y[:,j]
        end

        if norm(x_i - X, 2) < convergence_error
            break
        else
            x_i .= X
        end

    end

    X[abs.(X) .< get_threshold(opt)] .= zero(eltype(X))
    return iters
end


using naumer_Dimensionality_2022, Test

@testset "StaticSystem" begin

    dim = 2
    system = StaticSystem(dim)
    x = [1.0, 1.0]
    τ = 1.0

    @testset "dimension" begin
        @test dimension(system) == dim
    end

    @testset "differential" begin
        @test differential(system, x) ≈ [0.0, 0.0]
    end

    @testset "flow" begin
        @test flow(x, τ, system) ≈ x
    end

    @testset "flowJacobian" begin
        @test flowJacobian(x, τ, system) ≈ [1.0 0.0; 0.0 1.0]
    end
end

@testset "LinearSystem" begin

    dynamic_mat = [2.0 0.0; 0.0 1.0]
    system = LinearSystem(dynamic_mat)
    x = [1.0, 1.0]
    τ = 1.0

    @testset "dimension" begin
        @test dimension(system) == 2
    end

    @testset "differential" begin
        @test differential(system, x) ≈ [2.0, 1.0]
    end

    @testset "flow" begin
        @test flow(x, τ, system) ≈ [exp(2.0), exp(1.0)]
    end

    @testset "flowJacobian" begin
        @test flowJacobian(x, τ, system) ≈ [exp(2.0) 0; 0 exp(1.0)]
    end
end

@testset "VanDerPolSystem" begin

    μ = 1
    system = VanDerPolSystem(μ)
    x = [0.0, 1.0]
    τ = 1

    @testset "dimension" begin
        @test dimension(system) == 2
    end

    @testset "differential" begin
        @test isapprox(differential(system, x), [1.0, 1.0], atol=1e-5)
    end

    @testset "flow" begin
        # Very simple test, assuming correct functionality on May 04, 2022 to check for later issues
        @test isapprox(flow(x, τ, system), [1.250701247938, 0.9793760435906557], atol=1e-5)
    end

    @testset "flowJacobian" begin
        # Very simple test, assuming correct functionality on May 04, 2022 to check for later issues
        @test isapprox(flowJacobian(x, τ, system), [0.0435376 0.935839; -1.72123 -0.0820911], atol=1e-5)

    end
end

@testset "VanDerPolSystem_expanded" begin

    μ = 1
    dim = 1
    system = VanDerPolSystem_expanded(μ, dim)
    x = [0.0, 1.0, 1.0]
    τ = 1

    @testset "dimension" begin
        @test dimension(system) == dim + 2
    end

    @testset "differential" begin

        @test isapprox(differential(system, x), [1.0, 1.0, -x[3]], atol=1e-5)
    end

    @testset "flow" begin
        # Very simple test, assuming correct functionality on May 04, 2022 to check for later issues
        @test isapprox(flow(x, τ, system), [1.250701247938, 0.9793760435906557, exp(-1.0)*x[3]], atol=1e-5)
    end

    @testset "flowJacobian" begin
        # Very simple test, assuming correct functionality on May 04, 2022 to check for later issues
        @test isapprox(flowJacobian(x, τ, system), [0.0435376 0.935839 0; -1.72123 -0.0820911 0; 0 0 exp(-1.0)], atol=1e-5)

    end
end

@testset "HopfSystem" begin

    λ = 1
    b = -1
    system = HopfSystem(λ, b)
    x = [0.0, 1.0]
    τ = 2π

    @testset "dimension" begin
        @test dimension(system) == 2
    end

    @testset "differential" begin
        @test isapprox(differential(system, x), [-1.0, 0.0], atol=1e-5)
    end

    @testset "flow" begin
        @test isapprox(flow(x, τ, system), [0.0, 1.0], atol=1e-5)
    end

    @testset "flowJacobian" begin
        jacobian = flowJacobian(x, τ, system)
        @test isapprox(jacobian, [1.0 0; 0 0.0], atol=1e-4)
    end
end

@testset "LorenzSystem" begin

    σ = 10.0
    ρ = 28.0
    β = 8.0/3.0

    system = LorenzSystem(σ,ρ,β)
    x = [0.0, 0.0, 1.0]
    τ = 1

    @testset "dimension" begin
        @test dimension(system) == 3
    end

    @testset "differential" begin
        @test isapprox(differential(system, x), [0.0, 0.0, -8.0/3.0], atol=1e-5)
    end

    @testset "flow" begin
        # Very simple test, assuming correct functionality on May 04, 2022 to check for later issues
        @test isapprox(flow(x, τ, system), [0.0, 0.0, 0.06948364385801184], atol=1e-5)
    end

    @testset "flowJacobian" begin
        # Very simple test, assuming correct functionality on May 04, 2022 to check for later issues
        @test isapprox(flowJacobian(x, τ, system), [45246.6 36030.0 0; 98664.2 78567.0 0; 0 0 0.0694835], atol=1e-5, rtol=1e-4)
    end
end


module helpers

using Dates

"""Helper methods for Hijri conversion."""

_DAYS_BEFORE_MONTH = [-1]
dbm = 0
for dim in Dates.DAYSINMONTH
    global dbm
    append!(_DAYS_BEFORE_MONTH, dbm)
    dbm += dim
end

function _days_before_year(year)
    "year -> number of days before January 1st of year."
    y = year - 1
    return y*365 + y÷4 - y÷100 + y÷400
end

function _days_before_month(year, month)
    "year, month -> number of days in year preceding first day of month."
    @assert 1 <= month && month <= 12 "month must be in 1..12"
    return _DAYS_BEFORE_MONTH[month + 1] + (month > 2 && _is_leap(year))
end
#TODO: The functions above should probably be implemented in stdlib

_DI400Y = _days_before_year(401)    # number of days in 400 years
_DI100Y = _days_before_year(101)    #    "    "   "   " 100   "
_DI4Y   = _days_before_year(5)      #    "    "   "   "   4   "

# A 4-year cycle has an extra leap day over what we'd get from pasting
# together 4 single years.
@assert _DI4Y == 4 * 365 + 1 _DI4Y

# Similarly, a 400-year cycle has an extra leap day over what we'd get from
# pasting together 4 100-year cycles.
@assert _DI400Y == 4 * _DI100Y + 1 _DI400Y

# OTOH, a 100-year cycle has one fewer leap day than we'd get from
# pasting together 25 4-year cycles.
@assert _DI100Y == 25 * _DI4Y - 1 _DI100Y

function bisect_left(a, x, lo = 1, hi = nothing)
    if lo < 1
        throw(BoundsError(a, lo))
    end
    if hi === nothing
        hi = length(a) + 1  # It's not `length(a)`!
    end
    while lo < hi
        mid = (lo + hi) ÷ 2
        a[mid] < x ? lo = mid + 1 : hi = mid
    end
    return lo
end

@doc  """Convert Julian day number (JDN) to date ordinal number.

      :param jdn: Julian day number (JDN).
      :type jdn: int
      """
function jdn_to_ordinal(jdn::Int)

    return jdn - 1721425
end

@doc  """Convert date ordinal number to Julian day number (JDN).

      :param n: Date ordinal number.
      :type n: int
      """
function ordinal_to_jdn(n::Int)

    return n + 1721425
end

@doc  """Return Reduced Julian Day (RJD) number from Julian day number (JDN).

      :param jdn: Julian day number (JDN).
      :type jdn: int
      """
function jdn_to_rjd(jdn::Int)

    return jdn - 2400000
end

@doc  """Return Julian day number (JDN) from Reduced Julian Day (RJD) number.

      :param rjd: Reduced Julian Day (RJD) number.
      :type rjd: int
      """
function rjd_to_jdn(rjd::Int)

    return rjd + 2400000
end

@doc   "year, month, day -> ordinal, considering 01-Jan-0001 as day 1."
function _ymd2ord(year, month, day)
    @assert 1 <= month && month <= 12 "month must be in 1..12"
    dim = Dates.daysinmonth(Date(year, month))
    @assert 1 <= day && day <= dim "day must be in 1..%d' % dim"
    return floor(Int, _days_before_year(year) +
            _days_before_month(year, month) +
            day) - 1 # TODO: -1 is not necessary; without it, we should get .5s
end

function _is_leap(year)
    "year -> 1 if leap year, else 0."
    return year % 4 == 0 && (year % 100 != 0 || year % 400 == 0)
end

@doc   "ordinal -> (year, month, day), considering 01-Jan-0001 as day 1."
function _ord2ymd(n)

    # n is a 1-based index, starting at 1-Jan-1.  The pattern of leap years
    # repeats exactly every 400 years.  The basic strategy is to find the
    # closest 400-year boundary at or before n, then work with the offset
    # from that boundary to n.  Life is much clearer if we subtract 1 from
    # n first -- then the values of n at 400-year boundaries are exactly
    # those divisible by _DI400Y:
    #
    #     D  M   Y            n              n-1
    #     -- --- ----        ----------     ----------------
    #     31 Dec -400        -_DI400Y       -_DI400Y -1
    #      1 Jan -399         -_DI400Y +1   -_DI400Y      400-year boundary
    #     ...
    #     30 Dec  000        -1             -2
    #     31 Dec  000         0             -1
    #      1 Jan  001         1              0            400-year boundary
    #      2 Jan  001         2              1
    #      3 Jan  001         3              2
    #     ...
    #     31 Dec  400         _DI400Y        _DI400Y -1
    #      1 Jan  401         _DI400Y +1     _DI400Y      400-year boundary
    n -= 1
    n400, n = divrem(n, _DI400Y)
    year = n400 * 400 + 1   # ..., -399, 1, 401, ...

    # Now n is the (non-negative) offset, in days, from January 1 of year, to
    # the desired date.  Now compute how many 100-year cycles precede n.
    # Note that it's possible for n100 to equal 4!  In that case 4 full
    # 100-year cycles precede the desired day, which implies the desired
    # day is December 31 at the end of a 400-year cycle.
    n100, n = divrem(n, _DI100Y)

    # Now compute how many 4-year cycles precede it.
    n4, n = divrem(n, _DI4Y)

    # And now how many single years.  Again n1 can be 4, and again meaning
    # that the desired day is December 31 at the end of the 4-year cycle.
    n1, n = divrem(n, 365)

    year += n100 * 100 + n4 * 4 + n1
    if n1 == 4 || n100 == 4
        @assert n == 0
        return year-1, 12, 31
    end

    # Now the year is correct, and n is the offset from January 1.  We find
    # the month via an estimate that's either exact or one too large.
    leapyear = n1 == 3 && (n4 != 24 || n100 == 3)
    @assert leapyear == _is_leap(year)
    month = (n + 50) >> 5
    preceding = _DAYS_BEFORE_MONTH[month+1] + (month > 2 && leapyear)
    if preceding > n  # estimate is too large
        month -= 1
        preceding -= Dates.DAYSINMONTH[month] + (month == 2 && leapyear)
    end

    n -= preceding
    @assert 0 <= n && n < Dates.daysinmonth(Date(year, month)) string("n: ", n, ", days in month: ", Dates.daysinmonth((year, month)))

    # Now the year and month are correct, and n is the offset from the
    # start of that month:  we're done!
    return year, month, n+1
end

end


@testset "csr_to_sparse" begin
    ei = [1;2;3]
    ej = [3;4;4]
    ev = [8;9;10]
    (rp,ci,ai,m) = sparse_to_csr(ei,ej,ev)
    (nzi,nzj,nzv) = csr_to_sparse(rp,ci,ai)
    A = sparse(nzi,nzj,nzv,length(rp)-1,maximum(ci))

    # more tests added here
    # clique to sparse test
    rp = collect(1:5:26)
    #ci = vec(reshape(repmat(1:5,5,1)',25,1))
    ci = copy(vec(repeat(1:5,5,1)'))
    ai = ones(Int64,25)
    A = csr_to_sparse_matrix(rp,ci,ai,5,5)
    @test Matrix(A) == ones(5,5)

    A = csr_to_sparse_matrix(rp,ci,ai)
    @test Matrix(A) == ones(5,5)

    # 100 random trials
    for t = 1:100
        A = sprand(100,80,0.01)
        (rp,ci,ai) = sparse_to_csr(A)
        A2 = csr_to_sparse_matrix(rp,ci,ai,100,80)
        @test A == A2
    end
end


# bounds.jl

"""
   bound_absrem(f::Function, polf::Taylor1, polfI::Taylor1, x0::Interval, I::Interval)

Bound the absolute remainder of the polynomial approximation of `f` given
by the Taylor polynomial `polf` around `x0` on the interval `I`. It requires
the interval extension `polfI` of the polynomial that approximates `f` for
the whole interval `I`, in order to compute the Lagrange remainder.

If `polfI[end]` has a definite sign, then it is monotonic in the intervals
[I.lo, x0] and [x0.hi, I.hi], which is exploited; otherwise, it is used
to compute the Lagrange remainder.

"""
function bound_absrem(f::Function, polf::Taylor1, polfI::Taylor1,
        x0::Interval, I::Interval)

    _order = get_order(polf) + 1
    fTIend = polfI[_order]
    if (sup(fTIend) < 0 || inf(fTIend) > 0)
        # Absolute remainder is monotonic
        a = interval(I.lo)
        b = interval(I.hi)
        Δlo = f(a) - polf(a-x0)
        # Δlo = f(a) - bound_taylor1(polf, a-x0)
        Δhi = f(b) - polf(b-x0)
        # Δhi = f(b) - bound_taylor1(polf, b-x0)
        Δx0 = f(x0) - polf[0]
        Δ = hull(hull(Δlo, Δx0), Δhi)
    else
        # Lagrange bound
        Δ = fTIend * (I-x0)^_order
    end
    return Δ
end

"""
   bound_relrem(f::Function, polf::Taylor1, polfI::Taylor1, x0::Interval, I::Interval)

Bound the relative remainder of the polynomial approximation of `f` given
by the Taylor polynomial `polf` around `x0` on the interval `I`. It requires
an the interval extension `polfI` of a polynomial that approximates `f` for
the whole interval `I`, in order to compute the Lagrange remainder.

If `polfI[end]` has a definite sign, then it is monotonic in the interval `I`,
which is exploited; otherwise, the last coefficients bounds the relative
remainder.

"""
function bound_relrem(f::Function, polf::Taylor1, polfI::Taylor1,
        x0::Interval, I::Interval)

    _order = get_order(polf) + 1
    fTIend = polfI[_order+1]
    a = interval(I.lo)
    b = interval(I.hi)
    if (sup(fTIend) < 0 || inf(fTIend) > 0) && isempty(a ∩ x0) && isempty(b ∩ x0)
        # Error is monotonic
        denom_lo = (a-x0)^_order
        Δlo = f(a) - polf(a-x0)
        # Δlo = f(a) - bound_taylor1(polf, a-x0)
        Δlo = Δlo / denom_lo
        denom_hi = (b-x0)^_order
        Δhi = f(b) - polf(b-x0)
        # Δhi = f(b) - bound_taylor1(polf, b-x0)
        Δhi = Δhi / denom_hi
        Δ = hull(Δlo, Δhi)
    else
        # Lagrange coefficient
        Δ = polfI[_order]
    end
    return Δ
end

"""
    bound_taylor1(fT::Taylor1, I::Interval)

Compute a *tight* polynomial bound for the Taylor polynomial `fT`
in the interval `I`.

Note: Algorithm 2.1.1 corresponds to `evaluate(fT, I)` or simply `fT(I).
This function uses the roots of the derivative of `ft` to obtain a
tighter bound.

"""
function bound_taylor1(fT::Taylor1, I::Interval)

    # Check if the fT is monotonous (the derivative has a given sign)
    fTd  = TaylorSeries.derivative(fT)
    range_deriv = fTd(I)
    (sup(range_deriv) ≤ 0 || inf(range_deriv) ≥ 0) && return bound_taylor1(fT, fTd, I)

    # Compute roots of the derivative using the second derivative
    # Fix some sort of relative tolerance for Newton root search
    fTd2 = TaylorSeries.derivative(fTd)
    rootsder = roots(fTd, fTd2, I, Newton, 1.0e-5*mag(I))

    # Bound the range of fT using the roots and end points
    num_roots = length(rootsder)
    num_roots == 0 && return fT(I)
    rangepoly = hull( fT(interval(I.lo)), fT(interval(I.hi)) )
    for ind in 1:num_roots
        rangepoly = hull(rangepoly, fT(rootsder[ind].interval))
    end

    return rangepoly
end

"""
    bound_taylor1(fT::Taylor1, fTd::Taylor1, I::Interval)

Compute a *tight* polynomial bound for the Taylor polynomial `fT`
in the interval `I`, considering whether its derivative `ftd` has
a definite sign.

"""
function bound_taylor1(fT::Taylor1{T}, fTd::Taylor1{T}, I::Interval{T}) where {T}
    #
    if inf(fTd(I)) ≥ 0
        return @interval(fT(I.lo), fT(I.hi))
    elseif sup(fTd(I)) ≤ 0
        return @interval(fT(I.hi), fT(I.lo))
    end
    return fT(I)
end
function bound_taylor1(fT::Taylor1{Interval{T}}, fTd::Taylor1{Interval{T}},
        I::Interval{S}) where {T,S}
    #
    if inf(fTd(I)) ≥ 0
        return hull(fT(I.lo), fT(I.hi))
    elseif sup(fTd(I)) ≤ 0
        return hull(fT(I.hi), fT(I.lo))
    end
    return fT(I)
end


#!/usr/bin/env julia
#
# Jerzy Wroczyński (nr. indeksu 250075)
# Obliczenia Naukowe, Laboratorium
# Lista 3, Zadanie 4.1.
#

include("iterative-methods.jl")
using .IterativeMethods

fl = IterativeMethods.FL

r, fr, k, err =  IterativeMethods.mbisekcji(
    x -> sin(x) - (x/2)^2,
    fl(1.5),
    fl(2),
    fl((10^(-5))/2),
    fl((10^(-5))/2)
)

println(
    "r = ", r, "\n",
    "f(r) = ", fr, "\n",
    "k = ", k, "\n",
    err == 0 ? "" : string("error code: ", err)
)


using HeartRateVariability
using Test

n=HeartRateVariability.infile("e1304.txt")
td=HeartRateVariability.time_domain(n)
fd=HeartRateVariability.frequency(n)
nl=HeartRateVariability.nonlinear(n)
g=HeartRateVariability.geometric(n)

@testset "HeartRateVariability.jl" begin

    @testset "HeartRateVariability.infile" begin
        @test HeartRateVariability.infile("e1304","atr")==n
    end

    @testset "HeartRateVariability.time_domain" begin
        @test td.mean≈917.24 atol=0.1
        @test td.sdnn≈137.19 atol=0.1
        @test td.rmssd≈27.85 atol=0.1
        @test td.sdsd≈20.35 atol=0.1
        @test td.nn50≈342 atol=1
        @test td.pnn50≈4.41 atol=0.1
        @test td.nn20≈2831 atol=1
        @test td.pnn20≈36.53 atol=0.1
        @test td.rRR≈2.67 atol=0.1
    end

    @testset "HeartRateVariability.frequency" begin
        @test fd.vlf≈1317.96 atol=0.01*fd.vlf
        @test fd.lf≈90.36 atol=0.01*fd.lf
        @test fd.hf≈176.05 atol=0.01*fd.hf
        @test fd.lfhf_ratio≈0.51 atol=0.01*fd.lfhf_ratio
        @test fd.tp≈1584.35 atol=0.01*fd.tp
    end

    @testset "HeartRateVariability.nonlinear" begin
        @test nl.apen≈2.16 atol=0.1
        @test nl.sampen≈2.16 atol=0.1
        @test nl.hurst≈0.37 atol=0.1
        @test nl.renyi0≈-6.82 atol=0.1
        @test nl.renyi1≈-6.83 atol=0.1
        @test nl.renyi2≈-6.84 atol=0.1

        #testing if get_rs from module Nonlinear returns 0 when S or R is 0
        @test HeartRateVariability.Nonlinear.get_rs(ones(100))==0

        #testing if warning is thwown
        @test_logs (:warn,"To obtain a valid value for the hurst coefficient, the length of the data series must be greater than or equal to 100.") HeartRateVariability.nonlinear([1.0,2.0,1.0,2.0,1.0,2.0,1.0,2.0,1.0,2.0])
    end

    @testset "HeartRateVariability.geometric" begin
        @test g.poincare!=nothing
        @test g.recurrence!=nothing
        @test_throws ErrorException HeartRateVariability.geometric(n,"error")
    end
end


################### stepwise regression with replacement (SRR) #################
# one-lookahead greedy two-step algorithm, based on stepwise regression algorithms
function srr(A::AbstractMatOrFac, b::AbstractVector, k::Int, δ::Real = 1e-12,
                                            x = spzeros(eltype(A), size(A, 2));
                                            maxiter = 4k, initialization::Int = 1,
                                            l::Int = 1)
    P = StepwiseRegression(A, b, x.nzind)
    # initialize support with k atoms using
    if initialization == 1 # oblivious algorithm
        oblivious_acquisition!(P, x, k-nnz(x)) # initialize with k largest inner products
    elseif initialization == 2 # forward regression
        for _ in 1:k
            update!(P, x)
        end
    elseif initialization == 3 # random indices
        random_acquisition!(P, x, k-nnz(x)) # initialize with k largest inner products
    end
    resnorm = norm(residual!(P, x))
    for i in 1:maxiter
        oldnorm = resnorm
        for _ in 1:l
            forward_step!(P, x, 0, 0) || break # add unless we found a solution
        end
        while nnz(x) > k # delete until we have an active set of size k
            backward_step!(P, x, Inf, Inf)
        end
        resnorm = norm(residual!(P, x))
        if resnorm ≤ δ || oldnorm ≤ resnorm # until convergence or stationarity
            break
        end
    end
    return x
end

@inline function argmaxinner!(P::StepwiseRegression, k::Int)
    Ar = P.A' * P.r
    @. Ar = abs(Ar)
    partialsortperm(Ar, 1:k, rev = true)
end

################################## Subspace Pursuit ############################
struct SubspacePursuit{T, AT<:AbstractMatOrFac{T}, B<:AbstractVector{T}} <: AbstractMatchingPursuit{T}
    A::AT
    b::B
    k::Int # maximum number of non-zeros

    # temporary storage
    r::B # residual
    Ar::B # inner products between measurement matrix and residual
    Ai::AT # space for A[:, x.nzind] and its qr factorization
end
const SP = SubspacePursuit

function SP(A::AbstractMatOrFac, b::AbstractVector, k::Integer)
    2k > length(b) && error("2k = $(2k) > $(length(b)) = length(b) is invalid for Subspace Pursuit")
    n, m = size(A)
    T = eltype(A)
    r, Ar = copy(b), zeros(T, m)
    Ai = zeros(T, (n, 2k))
    SP(A, b, k, r, Ar, Ai)
end

# adds indices of k atoms with largest inner products with residual
# solves least-squares problem on expanded support
# could use threshold on P.Ar for adaptive stopping
# this is oblivious algorithm applied to residual
function sp_acquisition!(P::SP, x, k::Int = P.k)
    residual!(P, x)
    i = argmaxinner!(P, k)
    @. x[i] = NaN
    solve!(P, x)
end
# IDEA: could pre-allocate nz arrays to be of length 2K
# IDEA: could add efficient qr updating
function update!(P::SP, x::AbstractVector = spzeros(size(P.A, 2)), ε::Real = 0.)
    nnz(x) == P.k || throw("nnz(x) = $(nnz(x)) ≠ $(P.k) = k")
    sp_acquisition!(P, x)
    i = partialsortperm(abs.(x.nzval), 1:nnz(x)-P.k) # find the smallest atoms
    sort!(i) # for deleteat! to work, indices need to be sorted
    deleteat!(x.nzind, i)
    deleteat!(x.nzval, i)
    solve!(P, x) # optimize all active atoms
end

# calculates k-sparse approximation to Ax = b via subspace pursuit
# could also stop if indices are same between iterations
function sp(A::AbstractMatOrFac, b::AbstractVector, k::Int, δ::Real = 1e-12; maxiter = 16k)
    P = SP(A, b, k)
    x = spzeros(size(A, 2))
    sp_acquisition!(P, x, P.k)
    resnorm = norm(residual!(P, x))
    for i in 1:maxiter
        oldnorm = resnorm
        update!(P, x)
        resnorm = norm(residual!(P, x))
        if resnorm ≤ δ || oldnorm ≤ resnorm
            break
        end
    end
    return x
end

# ordinary least squares solve
@inline function solve!(P::SP, x::AbstractVector, b::AbstractVector = P.b)
    F = factorize!(P, x)
    ldiv!(x.nzval, F, b)     # optimize all active atoms
end

############################# OMP with replacement #############################
struct OMPR{T, AT<:AbstractMatOrFac{T}, B<:AbstractVector{T}, FT} <: AbstractMatchingPursuit{T}
    A::AT
    b::B
    k::Int
    # l # number of atoms to replace

    # temporary storage
    r::B # residual
    Ar::B # inner products between measurement matrix and residual
    QA::AT # (k, m)
    AiQR::FT # updatable QR factorization of Ai
end

# similar to RMP, or OMP?
function OMPR(A::AbstractMatOrFac, b::AbstractVector, k::Int)
    n, m = size(A)
    T = eltype(A)
    r, Ar = zeros(T, n), zeros(T, m)
    AiQR = UpdatableQR(T, n, n) # initializing empty qr factorization with maximum rank n
    QA = zeros(T, (n, m))
    OMPR(A, b, k, r, Ar, QA, AiQR)
end

# η is stepsize
function update!(P::OMPR, x::AbstractVector, η::Real = 1.)
    nnz(x) == P.k || throw("nnz(x) = $(nnz(x)) ≠ $(P.k) = k")
    residual!(P, x)
    copy!(P.Ar, x)
    mul!(P.Ar, P.A', P.r, η, 1)
    begin # argmax of abs(P.Ar) over i not in x.nzind     # argmax( [abs(ar) for ar in P.Ar]
        m = 0.
        i = 0
        for (j, Arj) in enumerate(P.Ar)
            if j ∉ x.nzind
                fArj = abs(Arj)
                if fArj > m
                    m = fArj
                    i = j
                end
            end
        end
    end

    if i == 0
        return x
    end

    # add non-zero index to active set
    x[i] = NaN
    qr_i = findfirst(==(i), x.nzind) # index in qr where new atom should be added

    # least-squares solve for active atoms (this would be two-stage algorithm)
    # ldiv!(x.nzval, P.AiQR, P.b)

    # gradient descent for active atoms
    @. x.nzval = P.Ar[x.nzind]

    j = argmin(abs, x.nzval) # delete index
    deleteat!(x.nzind, j) # delete the value from array
    deleteat!(x.nzval, j)

    if qr_i ≠ j # update qr factorization using Givens rotations
        a = @view P.A[:, i]
        add_column!(P.AiQR, a, qr_i)
        remove_column!(P.AiQR, j)
    end

    # least-squares solve for active atoms
    ldiv!!(x.nzval, P.AiQR, P.b, P.r)
    return x
end

# k is desired sparsity level
# l is cardinality of maximum replacement per iteration, l = k corresponds to sp
function ompr(A::AbstractMatOrFac, b::AbstractVector, k::Int, δ::Real,
                                x = spzeros(size(A, 2)); maxiter = size(A, 1))
    P = OMPR(A, b, k)
    if nnz(x) < P.k # make sure support set is of size k
        @. x = 0
        dropzeros!(x)
        oblivious_acquisition!(P, x, P.k)
    end
    resnorm = norm(residual!(P, x))
    for i in 1:maxiter
        oldnorm = resnorm
        update!(P, x)
        resnorm = norm(residual!(P, x))
        if resnorm ≤ δ || oldnorm ≤ resnorm
            break
        end
    end
    return x
end


using Gadfly
using Statistics
using Clustering
using DataFrames

import Distances: pairwise, Euclidean
import StatsBase: cov2cor
import RLEVectors: rep

#include("helpers.jl")

export plot_corr
function plot_corr(chain, m::Int64, labs::Array{String,1}; reorder = true, key = false, min_value = -1, max_value = 1, linkage = :ward, reds = false)
"""
    Plot various correlation matrices
    1. Obtain some covariance chain and average over the estimates
    2. Convert to a correlation matrix
    3. Cluster using hierarchical clustering, and reorder the dimensions based on the output
    4. Convert to a molten data frame and plot
"""
    Schain = get_Sigma_chain(chain, m)
    covmat = mapslices(x -> mean(x), Schain; dims = 3)[:,:,1]
    cormat = cov2cor(covmat, sqrt.(diag(covmat)))

    if reorder
            D = pairwise(Euclidean(), cormat)
            try
                    cl = Clustering.hclust(D, linkage=linkage)
                    ordering = getproperty(cl, :order)
            catch
                    ordering = 1:size(cormat,1)
            end
    else
            ordering = 1:size(cormat,1)
    end

    df = DataFrame(cormat[ordering, ordering], Symbol.(labs[ordering]))
    molten = melt(df)
    molten.v2 = rep(labs[ordering], times = size(df,1))
    if reds
            Gadfly.plot(molten, x="variable", y="v2", color="value",
                Geom.rectbin,
                Coord.cartesian(fixed = true),
                Guide.xlabel(nothing), Guide.ylabel(nothing),
                Scale.color_continuous(minvalue=min_value, maxvalue=max_value, colormap = Scale.lab_gradient("white", "red")),
                Guide.colorkey(title=""),
                key ? Theme(key_position = :right) : Theme(key_position = :none))
        else
                Gadfly.plot(molten, x="variable", y="v2", color="value",
                    Geom.rectbin,
                    Coord.cartesian(fixed = true),
                    Guide.xlabel(nothing), Guide.ylabel(nothing),
                    Scale.color_continuous(minvalue=min_value, maxvalue=max_value),
                    Guide.colorkey(title=""),
                    key ? Theme(key_position = :right) : Theme(key_position = :none))
        end
end

export plot_eigvec
function plot_eigvec(chain, m::Int64, labs::Array{String,1}, eig_num::Int64; pal=missing, reorder=missing)
        Schain = get_Sigma_chain(chain, m)
        covmat = mapslices(x -> mean(x), Schain; dims = 3)[:,:,1]

        if !ismissing(reorder)
                covmat = covmat[reorder,reorder]
                labs = labs[reorder]
                pal = pal[reorder]
        end

        evals, evecs = (eigvals(covmat), eigvecs(covmat))

        # Make names dataframe so plot has names
        # The kth ROW is the kth eigenvector here
        edf = DataFrame(evecs', Symbol.(labs))
        edf.vec = 1:size(covmat, 1)
        molten = melt(edf)

        # Plot first eigenvector
        if ismissing(pal)
                p = Gadfly.plot(molten[molten.vec .== eig_num, :],
                        x=:variable,
                        y=:value,
                        Guide.xlabel(nothing), Guide.ylabel(nothing),
                        Geom.bar)#,
                        #yintercept = 0,
                        #Geom.hline(color="black", style = :dash))
        else
                p = Gadfly.plot(molten[molten.vec .== eig_num, :],
                        x=:variable,
                        y=:value,
                        color=pal,
                        Guide.xlabel(nothing), Guide.ylabel(nothing),
                        Geom.bar)#,
                        #yintercept = 0,
                        #Geom.hline(color="black", style = :dash))
        end
        (p, sort(evals; rev=true))
end

export plot_effects
function plot_effects(chain, obsidx::Int64, labs::Array{String,1}, M::Int64, mu_ests; pal = missing)
        z_chain = get_z_chain(chain)
        MAPs = mapslices(x -> [ mean(x .== y) for y in 1:M ], z_chain; dims = 2)
        emap = MAPs[obsidx,:]

        # Effect size estimates in each cell type,
        #       weighted by the posterior probability the this site is in any given cluster
        #       plus the standard errors
        emeans = mapslices(x -> sum(x), hcat([ mu_ests[m,:] .* emap[m] for m in 1:M ]...)'; dims = 1)[1,:]
        eses = mapslices(y -> sum(y), hcat([ mapslices(x -> std(x), get_mu_chain(chain, m);
                                dims = 2) .* emap[m] for m in 1:M ]...)'; dims = 1)[1,:]

        # Plot first eigenvector
        if ismissing(pal)
                Gadfly.plot(y = labs, x = emeans,
                                xmin = emeans .- eses,
                                xmax = emeans .+ eses,
                                Geom.point,
                                Guide.xlabel(nothing), Guide.ylabel(nothing),
                                Geom.xerrorbar)
        else
                Gadfly.plot(y = labs, x = emeans,
                                xmin = emeans .- eses,
                                xmax = emeans .+ eses,
                                Geom.point,
                                Geom.xerrorbar,
                                Guide.xlabel(nothing), Guide.ylabel(nothing),
                                color = pal)
        end
end


"""
Function to shrink acceptbw by weighted quantile of previously accepted distances

# Arguments
* `distances::Array{Float64, 1}`: vector of distances or vector of distance vectors
* `weights::WeightVec`: weight vector must be in WeightVec format
* `acceptbwQuantile::Float64`: quantile (0,1) to shrink to

# Value
acceptbw acceptance acceptbws
"""
function quantile_acceptbw(d::Array{<:AbstractFloat, 1}, w::StatsBase.AnalyticWeights, q::AbstractFloat)
    return StatsBase.quantile(d, w, q)
end

function quantile_acceptbw{A <: SingleMeasureAbcPmc}(abcpmc::A, q::Float64)
    return StatsBase.quantile(abcpmc.distances, abcpmc.weights, q)
end

function independentquantile_acceptbw(d::Array{<:AbstractFloat, 2}, w::StatsBase.AnalyticWeights, q::AbstractFloat)
    m = size(d, 2)
    acceptbws = Array{Float64}(m)
    for jj in 1:m
        acceptbws[jj] = quantile(d[:, jj], w, q)
    end
    return acceptbws
end

function independentquantile_acceptbw(abcpmc::A, q::Float64) where A <: MultiMeasureAbcPmc
    return independentquantile_acceptbw(abcpmc.distances, abcpmc.weights, q)
end

function jointquantile_acceptbw(d::Array{<:AbstractFloat, 2}, w::StatsBase.AnalyticWeights, q::AbstractFloat)
    m = size(d, 2)
    return independentquantile_acceptbw(d, w, 1.0 - (1.0 - q) / m)
end

function jointquantile_acceptbw(abcpmc::A, q::Float64) where A <: MultiMeasureAbcPmc
    m = size(abcpmc.distances, 2)
    return independentquantile_acceptbw(abcpmc, 1.0 - (1.0 - q) / m)
end


export indicator

"""
    indicator(y, b; indregion=:above, ineq=:nonstrict)

Indicator variable of default form ``\\mathbf{1}\\{y ≥ b\\}``. The inequality can be made strict and the direction of inequality can be reversed.

# Arguments
- `ineq=:nonstrict` (default):  calculates ``\\mathbf{1}\\{y ≥ b\\}`` or ``\\mathbf{1}\\{y ≤ b\\}``.
- `ineq=:strict`: calculates ``\\mathbf{1}\\{y > b\\}`` or ``\\mathbf{1}\\{y < b\\}``.
- `indregion=:above` (default): calculates ``\\mathbf{1}\\{y ≥ b\\}`` or ``\\mathbf{1}\\{y > b\\}``.
- `indregion=:below`: calculates ``\\mathbf{1}\\{y ≤ b\\}`` or ``\\mathbf{1}\\{y < b\\}``.
"""
function indicator(y, b; indregion = :above, ineq=:nonstrict)
    # Indicator function. The default mode ('above') returns 𝟏{y ≥ b}. The other available mode, 'below', returns 𝟏{y ≤ b}. Set 'ineq' to 'strict' to get 𝟏{y > b} or 𝟏{y < b}.
    indregion = Symbol(indregion)
    ineq = Symbol(ineq)
    if ineq == :nonstrict
        if indregion == :above
            y ≥ b ? (return 1) : (return 0)
        end
        if indregion == :below
            y ≤ b ? (return 1) : (return 0)
        end
    end
    if ineq == :strict
        if indregion == :above
            y > b ? (return 1) : (return 0)
        end
        if indregion == :below
            y < b ? (return 1) : (return 0)
        end
    end
    regionlist = [:above, :below]
    if (indregion .== regionlist) == zeros(2)
        error("Unknown indicator region: $indregion")
    end
    ineqlist = [:strict, :nonstrict]
    if (indregion .== ineqlist) == zeros(2)
        error("Unknown inequality type: $ineq")
    end
end


@doc raw"""
    NelderMead(M, F [, p])
perform a nelder mead minimization problem for the cost funciton `F` on the
manifold `M`. If the initial population `p` is not given, a random set of
points is chosen.

This algorithm is adapted from the Euclidean Nelder-Mead method, see
[https://en.wikipedia.org/wiki/Nelder–Mead_method](https://en.wikipedia.org/wiki/Nelder–Mead_method)
and
[http://www.optimization-online.org/DB_FILE/2007/08/1742.pdf](http://www.optimization-online.org/DB_FILE/2007/08/1742.pdf).

# Input

* `M` – a manifold $\mathcal M$
* `F` – a cost function $F\colon\mathcal M\to\mathbb R$ to minimize
* `population` – (n+1 `random_point(M)`) an initial population of $n+1$ points, where $n$
  is the dimension of the manifold `M`.

# Optional

* `stopping_criterion` – ([`StopAfterIteration`](@ref)`(2000)`) a [`StoppingCriterion`](@ref)
* `retraction` – (`exp`) a `retraction(M,x,ξ)` to use.
* `α` – (`1.`) reflection parameter ($\alpha > 0$)
* `γ` – (`2.`) expansion parameter ($\gamma$)
* `ρ` – (`1/2`) contraction parameter, $0 < \rho \leq \frac{1}{2}$,
* `σ` – (`1/2`) shrink coefficient, $0 < \sigma \leq 1$

and the ones that are passed to [`decorate_options`](@ref) for decorators.

# Output
* either `x` the last iterate or the complete options depending on the optional
  keyword `return_options`, which is false by default (hence then only `x` is
  returned).
"""
function NelderMead(M::MT,
    F::Function,
    population = [random_point(M) for i=1:(manifoldDimension(M)+1) ];
    stopping_criterion::StoppingCriterion = StopAfterIteration(200000),
    α = 1., γ = 2., ρ=1/2, σ = 1/2,
    return_options=false,
    kwargs... #collect rest
  ) where {MT <: Manifold,T}
    p = CostProblem(M,F)
    o = NelderMeadOptions(population, stopping_criterion;
    α = α, γ = γ, ρ = ρ, σ = σ)
    o = decorate_options(o; kwargs...)
    resultO = solve(p,o)
    if return_options
        return resultO
    else
        return get_solver_result(resultO)
    end
end
#
# Solver functions
#
function initialize_solver!(p::P,o::O) where {P <: CostProblem, O <: NelderMeadOptions}
    # init cost and x
    o.costs = get_cost.(Ref(p), o.population )
    o.x = o.population[argmin(o.costs)] # select min
end
function step_solver!(p::P,o::O,iter) where {P <: CostProblem, O <: NelderMeadOptions}
    m = mean(p.M, o.population)
    ind = sortperm(o.costs) # reordering for cost and p, i.e. minimizer is at ind[1]
    ξ =log( p.M, m, o.population[last(ind)])
    # reflect last
    xr = exp(p.M, m, - o.α * ξ )
    Costr = get_cost(p,xr)
    # is it better than the worst but not better than the best?
    if Costr >= o.costs[first(ind)] && Costr < o.costs[last(ind)]
        # store as last
        o.population[last(ind)] = xr
        o.costs[last(ind)] = Costr
    end
    # --- Expansion ---
    if Costr < o.costs[first(ind)] # reflected is better than fist -> expand
        xe = exp(p.M, m, - o.γ * o.α * ξ)
        Coste = get_cost(p,xe)
        if Coste < Costr # expanded successful
            o.population[last(ind)] = xe
            o.costs[last(ind)] = Coste
        else # expansion failed but xr is still quite good -> store
            o.population[last(ind)] = xr
            o.costs[last(ind)] = Costr
        end
    end
    # --- Contraction ---
    if Costr > o.costs[ind[end-1]] # even worse than second worst
        if Costr < o.costs[last(ind)] # but at least better tham last
            # outside contraction
            xc = exp(p.M, m, - o.ρ*ξ)
            Costc = get_cost(p,xc)
            if Costc < Costr # better than reflected -> store as last
                o.population[last(ind)] = xr
                o.costs[last(ind)] = Costr
            end
        else # even worse than last -> inside contraction
            # outside contraction
            xc = exp(p.M, m, o.ρ*ξ)
            Costc = get_cost(p,xc)
            if Costc < o.costs[last(ind)] # better than last ? -> store
                o.population[last(ind)] = xr
                o.costs[last(ind)] = Costr
            end
        end
    end
    # --- Shrink ---
    for i=2:length(ind)
        o.population[ind[i]] = shortest_geodesic(p.M, o.population[ind[1]], o.population[ind[i]], o.σ)
        # update cost
        o.costs[ind[i]] = get_cost(p, o.population[ind[i]])
    end
    # store best
    o.x = o.population[ argmin(o.costs) ]
end
get_solver_result(p::P,o::O) where {P <: CostProblem, O <: NelderMeadOptions} = o.x


import SigDigits.count_digits
import SigDigits.count_decimalplaces
import SigDigits.count_nonzero_decimalplaces

@testset "Total digits" begin
    @test count_digits(0.0000007) == 8
    @test count_digits(5.034001) == 7
    @test count_digits(7000) == 4
    @test count_digits(1.000) == 1
    @test count_digits(6.983410) == 6
    @test count_digits(14080.093149) == 11
end

@testset "Decimal places" begin
    @test count_decimalplaces(0.0000007) == 7
    @test count_decimalplaces(5.034001) == 6
    @test count_decimalplaces(7000) == 0
    @test count_decimalplaces(1.000) == 0
    @test count_decimalplaces(6.983410) == 5
    @test count_decimalplaces(14080.093149) == 6
end

@testset "Nonzero decimal places" begin
    @test count_nonzero_decimalplaces(0.0000007) == 1
    @test count_nonzero_decimalplaces(5.034001) == 3
    @test count_nonzero_decimalplaces(7000) == 0
    @test count_nonzero_decimalplaces(1.000) == 0
    @test count_nonzero_decimalplaces(6.983410) == 5
    @test count_nonzero_decimalplaces(14080.093149) == 5
end


using GeometricIntegrators.Integrators
using GeometricIntegrators.Tableaus
using GeometricIntegrators.Utils
using Test

include("coefficients_tests.jl")
include("rk_tableaus_tests.jl")
include("splitting_tableaus_tests.jl")


## squared_exponential.jl : implementation of squared exponential covariance function

## SquaredExponential ##
struct SquaredExponential{T} <: IsotropicCovarianceStructure{T}
    λ::T
    σ::T
    p::T
end

"""
    SquaredExponential(λ, σ=1, p=2)

Create a squared exponential (Gaussian) covariance structure with correlation length `λ`, (optional) marginal standard deviation `σ` and (optional) `p`-norm.

# Examples
```jldoctest
julia> s1 = SquaredExponential(0.1)
Gaussian (λ=0.1, σ=1.0, p=2.0)

julia> s2 = SquaredExponential(0.1, σ=2.)
Gaussian (λ=0.1, σ=2.0, p=2.0)

```
"""
function SquaredExponential(λ::T where {T<:Real}; σ=1.0::T where {T<:Real}, p=2::T where {T<:Real})
    λ > 0 || throw(ArgumentError("correlation length λ of squared exponential covariance cannot be negative or zero"))
    σ > 0 || throw(ArgumentError("marginal standard deviation σ of squared exponential covariance cannot be negative or zero"))
    p >= 1 || throw(ArgumentError("in p-norm, p must be greater than or equal to 1"))
	isinf(p) && throw(ArgumentError("in p-norm, p cannot be infinity"))
    SquaredExponential{promote_type(typeof(λ),typeof(σ),typeof(p))}(promote(λ,σ,p)...)
end

"""
    Gaussian(λ, σ=1, p=2)

Create a Gaussian (squared exponential) covariance structure with correlation length `λ`, (optional) marginal standard deviation `σ` and (optional) `p`-norm.

# Examples
```jldoctest
julia> g1 = Gaussian(0.1)
Gaussian (λ=0.1, σ=1.0, p=2.0)

julia> g2 = Gaussian(0.1, σ=2.)
Gaussian (λ=0.1, σ=2.0, p=2.0)

```
"""
const Gaussian = SquaredExponential

# evaluate Gaussian covariance
function apply(s::SquaredExponential,x::T) where {T<:Real}
    exp(-(x/s.λ)^2)
end

show(io::IO, s::SquaredExponential) = print(io, "Gaussian (λ=$(s.λ), σ=$(s.σ), p=$(s.p))")


"""
compute the kronecker product of a and b.\n
used as a shorthand notation.\n
"""
⊗(a, b) = kron(a, b)

"""
AToE = _electric_vector_potential_to_electric_field_operator(number_of_cells, resolution, k_b)\n
This operators is used to compute the total electric field on the FDTD grid (Yee cell) given the vector potential.\n
k_b is the wavenumber in the background medium.\n
"""
function _electric_vector_potential_to_electric_field_operator(
    number_of_cells,
    resolution,
    k_b
)
    T = ComplexF32
    p = number_of_cells[1]
    q = number_of_cells[2]
    r = number_of_cells[3]
    Δx = Float32(resolution[1])
    Δy = Float32(resolution[2])
    Δz = Float32(resolution[3])

    k_b = Float32(k_b)
    # needed -> Ex : (P+2) x (Q+1) : (R+1), wanted -> Ex : P x (Q+1) : (R+1)
    # needed -> Ey : (P+1) x (Q+2) : (R+1), wanted -> Ey : (P+1) x Q : (R+1)
    # needed -> Ez : (P+1) x (Q+1) : (R+2), wanted -> Ez : (P+1) x (Q+1) : R
    #operators for Ex
    Rp = spdiagm(p, p+2, 1=>ones(T, p))
    Rq = spdiagm(q+1, q+1, 0=>ones(T, q+1))
    Rr = spdiagm(r+1, r+1, 0=>ones(T, r+1))

    Ix   = (Rr ⊗ (Rq ⊗ Rp))
    ∂x   = 1/Δx .*spdiagm(p, p+1, 0 =>-one(T) .*ones(T, p), 1 =>ones(T, p))
    ∂y   = 1/Δy .*spdiagm(q+1, q+2, 0 =>-one(T) .*ones(T, q+1), 1 =>ones(T, q+1))
    ∂z   = 1/Δz .*spdiagm(r+1, r+2, 0 =>-one(T) .*ones(T, r+1), 1 =>ones(T, r+1))
    ∂x²  = 1/(Δx)^2 .*spdiagm(p, p+2, 0 => ones(T, p),1=> -2 .*ones(T, p), 2 => ones(T, p))

    xx = Ix'*((Rr ⊗ (Rq ⊗ ∂x²)) + k_b^2 .*Ix)
    yx = Ix'*((Rr ⊗ (∂y ⊗ ∂x)))
    zx = Ix'*((∂z ⊗ (Rq ⊗ ∂x)))

    #operators for Ey
    Rp = spdiagm(p+1, p+1, 0=>ones(T, p+1))
    Rq = spdiagm(q, q+2, 1=>ones(T, q))
    Rr = spdiagm(r+1, r+1, 0=>ones(T, r+1))

    Iy   = (Rr ⊗ (Rq ⊗ Rp))
    ∂x   = 1/Δx .*spdiagm(p+1, p+2, 0 =>-one(T) .*ones(T, p+1), 1 =>ones(T, p+1))
    ∂y   = 1/Δy .*spdiagm(q, q+1, 0 =>-one(T) .*ones(T, q), 1 =>ones(T, q))
    ∂z   = 1/Δz .*spdiagm(r+1, r+2, 0 =>-one(T) .*ones(T, r+1), 1 =>ones(T, r+1))
    ∂y²  = 1/(Δy)^2 .*spdiagm(q, q+2,0 =>ones(T, q),1=> -2 .*ones(T, q), 2 =>ones(T, q))

    xy = Iy'*((Rr ⊗ (∂y ⊗ ∂x)))
    yy = Iy'*((Rr ⊗ (∂y² ⊗ Rp)) + k_b^2 .*Iy)
    zy = Iy'*((∂z ⊗ (∂y ⊗ Rp)))

    #operators for Ez
    Rp = spdiagm(p+1, p+1, 0=>ones(T, p+1))
    Rq = spdiagm(q+1, q+1, 0=>ones(T, q+1))
    Rr = spdiagm(r, r+2, 1=>ones(T, r))

    Iz   = (Rr ⊗ (Rq ⊗ Rp))
    ∂x   = 1/Δx .*spdiagm(p+1, p+2, 0 =>-one(T) .*ones(T, p+1), 1 =>ones(T, p+1))
    ∂y   = 1/Δy .*spdiagm(q+1, q+2, 0 =>-one(T) .*ones(T, q+1), 1 =>ones(T, q+1))
    ∂z   = 1/Δz .*spdiagm(r, r+1, 0 =>-one(T) .*ones(T, r), 1 =>ones(T, r))
    ∂z²  = 1/(Δz)^2 .*spdiagm(r, r+2, 0 =>ones(T, r),1=> -2 .*ones(T, r), 2 =>ones(T, r))

    xz = Iz'*((∂z ⊗ (Rq ⊗ ∂x)))
    yz = Iz'*((∂z ⊗ (∂y ⊗ Rp)))
    zz = Iz'*((∂z² ⊗ (Rq ⊗ Rp)) + k_b^2 .*Iz)

    return hcat(vcat(xx, xy, xz), vcat(yx, yy, yz), vcat(zx, zy, zz));
end

"""
Ig = _create_Greens_functions_restriction_operators(number_of_cells)\n
These operators are used to restrict the output of ∫GχEdV.\n
This integral is computed using an FFT on a domain that is larger the number of cells.\n
Therefore, the output of these should be restricted to the correct Yee cell grid.\n
"""
function _create_Greens_functions_restriction_operators(number_of_cells)
    T = ComplexF32;
    p, q, r = number_of_cells
    dx, dy, dz = 2 .^(ceil.(Int, log2.(number_of_cells.+2).+1))

    Ix = spdiagm(p+2, dx, 0=>ones(T, p+2))
    Iy = spdiagm(q+1, dy, 0=>ones(T, q+1))
    Iz = spdiagm(r+1, dz, 0=>ones(T, r+1))
    x = (Iz ⊗ (Iy ⊗ Ix))

    Ix = spdiagm(p+1, dx, 0=>ones(T, p+1))
    Iy = spdiagm(q+2, dy, 0=>ones(T, q+2))
    y = (Iz ⊗ (Iy ⊗ Ix))

    Iy = spdiagm(q+1, dy, 0=>ones(T, q+1))
    Iz = spdiagm(r+2, dz, 0=>ones(T, r+2))
    z = (Iz ⊗ (Iy ⊗ Ix))

    return blockdiag(x, y, z)
end

"""
Ig = _create_Greens_functions_restriction_operators_gpu(number_of_cells)\n
These operators are used to restrict the output of ∫GχEdV.\n
This integral is computed using an FFT on a domain that is larger the number of 2*cells.\n
Therefore, the output of these should be restricted to the correct Yee cell grid.\n
"""
function _create_Greens_functions_restriction_operators_gpu(number_of_cells)
    T = ComplexF32;
    p, q, r = number_of_cells
    dx, dy, dz = 2 .^(ceil.(Int, log2.(number_of_cells.+2).+1))

    Ix = spdiagm(p+2, dx, 0=>ones(T, p+2))
    Iy = spdiagm(q+1, dy, 0=>ones(T, q+1))
    Iz = spdiagm(r+1, dz, 0=>ones(T, r+1))
    x = (Iz ⊗ (Iy ⊗ Ix))

    Ix = spdiagm(p+1, dx, 0=>ones(T, p+1))
    Iy = spdiagm(q+2, dy, 0=>ones(T, q+2))
    y = (Iz ⊗ (Iy ⊗ Ix))

    Iy = spdiagm(q+1, dy, 0=>ones(T, q+1))
    Iz = spdiagm(r+2, dz, 0=>ones(T, r+2))
    z = (Iz ⊗ (Iy ⊗ Ix))

    Ig = blockdiag(x, y, z) |> CUDA.CUSPARSE.CuSparseMatrixCSC

    return Ig
end

#=
some code not directly required for the method that this package was developed for.
=#
"""
FToE = _magnetic_vector_potential_to_electric_field_operator(number_of_cells, resolution, k_b)\n
This operators is used to compute the total electric field on the FDTD grid (Yee cell) given the magnetic vector potential.\n
k_b is the wavenumber in the background medium.\n
"""
function _magnetic_vector_potential_to_electric_field_operator(
    number_of_cells,
    resolution,
    k_b
)
    T = ComplexF32
    p = number_of_cells[1]
    q = number_of_cells[2]
    r = number_of_cells[3]
    Δx = Float32(resolution[1])
    Δy = Float32(resolution[2])
    Δz = Float32(resolution[3])

    k_b = Float32(k_b)
    # needed -> Hx : (P+1) x (Q+2) : (R+2), wanted -> Ex : P x (Q+1) : (R+1)
    # needed -> Hy : (P+2) x (Q+1) : (R+2), wanted -> Ey : (P+1) x Q : (R+1)
    # needed -> Hz : (P+2) x (Q+2) : (R+1), wanted -> Ez : (P+1) x (Q+1) : R
    # from the new dimension we can map back to the larger E field and get the old solver to work

    #operators for Kx to E

    # an operator to go from the smaller E field domain to the larger one
    Rp = spdiagm(p, p+2, 1 => ones(T, p))
    Rq = spdiagm(q+1, q+1, 0 => ones(T, q+1))
    Rr = spdiagm(r+1, r+1, 0 => ones(T, r+1))

    Ix   = (Rr ⊗ (Rq ⊗ Rp))

    #Kx to E
    ∂y   = 1/Δy .*spdiagm(q+1, q+2, 0 => -one(T) .*ones(T, q+1), 1 => ones(T, q+1))
    ∂z   = 1/Δz .*spdiagm(r+1, r+2, 0 => -one(T) .*ones(T, r+1), 1 => ones(T, r+1))

    xx = spzeros(T, (p+2)*(q+1)*(r+1),(p+1)*(q+2)*(r+2))
    yx = Ix'*((∂z ⊗ (Rq ⊗ Rp))) #f.y -> e.x
    zx = Ix'*((Rr ⊗ (∂y ⊗ Rp))) #f.z -> e.x

    #operators for Ey
    Rp = spdiagm(p+1, p+1, 0 => ones(T, p+1))
    Rq = spdiagm(q, q+2, 1 => ones(T, q))
    Rr = spdiagm(r+1, r+1, 0 => ones(T, r+1))

    Iy   = (Rr ⊗ (Rq ⊗ Rp))

    ∂x   = 1/Δx .*spdiagm(p+1, p+2, 0 => -one(T) .*ones(T, p+1), 1 => ones(T, p+1))
    ∂z   = 1/Δz .*spdiagm(r+1, r+2, 0 => -one(T) .*ones(T, r+1), 1 => ones(T, r+1))

    yy = spzeros(T, (p+1)*(q+2)*(r+1),(p+2)*(q+1)*(r+2))
    xy = Iy'*((∂z ⊗ (Rq ⊗ Rp))) #f.x -> e.y
    zy = Iy'*((Rr ⊗ (Rq ⊗ ∂x))) #f.z -> e.y

    #operators for Ez
    Rp = spdiagm(p+1, p+1, 0 => ones(T, p+1))
    Rq = spdiagm(q+1, q+1, 0=> ones(T, q+1))
    Rr = spdiagm(r, r+2, 1 => ones(T, r))

    Iz   = (Rr ⊗ (Rq ⊗ Rp))

    ∂x   = 1/Δx .*spdiagm(p+1, p+2, 0 => -one(T) .*ones(T, p+1), 1 => ones(T, p+1))
    ∂y   = 1/Δy .*spdiagm(q+1, q+2, 0 => -one(T) .*ones(T, q+1), 1 => ones(T, q+1))

    zz = spzeros(T, (p+1)*(q+1)*(r+2),(p+2)*(q+2)*(r+1))
    xz = Iz'*((Rr ⊗ (∂y ⊗ Rp)))
    yz = Iz'*((Rr ⊗ (Rq ⊗ ∂x)))

    return hcat(vcat(xx, xy, xz), vcat(yx, yy, yz), vcat(zx, zy, zz));
end

"""
AToH = _electric_vector_potential_to_magnetic_field_operator(number_of_cells, resolution, k_b)\n
This operators is used to compute the total magnetic field on the FDTD grid (Yee cell) given the electric vector potential.\n
k_b is the wavenumber in the background medium.\n
"""
function _electric_vector_potential_to_magnetic_field_operator(
    number_of_cells,
    resolution,
    k_b
)
    T = ComplexF32
    p = number_of_cells[1]
    q = number_of_cells[2]
    r = number_of_cells[3]
    Δx = Float32(resolution[1])
    Δy = Float32(resolution[2])
    Δz = Float32(resolution[3])

    k_b = Float32(k_b)
    # needed -> Ex : (P+2) x (Q+1) : (R+1), wanted -> Hx : (P+1) x (Q+2) X (R+2)
    # needed -> Ey : (P+1) x (Q+2) : (R+1), wanted -> Hy : (P+2) x (Q+1) : (R+2)
    # needed -> Ez : (P+1) x (Q+1) : (R+2), wanted -> Hz : (P+2) x (Q+2) : (R+1)
    # from the new dimension we can map back to the larger E field and get the old solver to work

    #operators for Kx to E

    # an operator to go from the smaller E field domain to the larger one
    Rp = spdiagm(p+1, p+1, 0 => ones(T, p+1))
    Rq = spdiagm(q, q+2, 1 => ones(T, q))
    Rr = spdiagm(r, r+2, 1 => ones(T, r))

    Ix   = (Rr ⊗ (Rq ⊗ Rp))

    #Kx to E
    ∂y   = 1/Δy .*spdiagm(q, q+1, 0 => -one(T) .*ones(T, q), 1 => ones(T, q))
    ∂z   = 1/Δz .*spdiagm(r, r+1, 0 => -one(T) .*ones(T, r), 1 => ones(T, r))

    xx = spzeros(T, (p+1)*(q+2)*(r+2),(p+2)*(q+1)*(r+1))
    yx = Ix'*((∂z ⊗ (Rq ⊗ Rp))) #A.y -> h.x
    zx = Ix'*((Rr ⊗ (∂y ⊗ Rp))) #A.z -> h.x

    #operators for Ey
    Rp = spdiagm(p, p+2, 1 => ones(T, p))
    Rq = spdiagm(q+1, q+1, 0 => ones(T, q+1))
    Rr = spdiagm(r, r+2, 1 => ones(T, r))

    Iy   = (Rr ⊗ (Rq ⊗ Rp))

    ∂x   = 1/Δx .*spdiagm(p, p+1, 0 => -one(T) .*ones(T, p), 1 => ones(T, p))
    ∂z   = 1/Δz .*spdiagm(r, r+1, 0 => -one(T) .*ones(T, r), 1 => ones(T, r))

    yy = spzeros(T, (p+2)*(q+1)*(r+2),(p+1)*(q+2)*(r+1))
    xy = Iy'*((∂z ⊗ (Rq ⊗ Rp))) #A.x -> h.y
    zy = Iy'*((Rr ⊗ (Rq ⊗ ∂x))) #A.z -> h.y

    #operators for Ez
    Rp = spdiagm(p, p+2, 1 => ones(T, p))
    Rq = spdiagm(q, q+2, 1 => ones(T, q))
    Rr = spdiagm(r+1, r+1, 0 => ones(T, r+1))

    Iz   = (Rr ⊗ (Rq ⊗ Rp))

    ∂x   = 1/Δx .*spdiagm(p, p+1, 0 => -one(T) .*ones(T, p), 1 => ones(T, p))
    ∂y   = 1/Δy .*spdiagm(q, q+1, 0 => -one(T) .*ones(T, q), 1 => ones(T, q))

    zz = spzeros(T, (p+2)*(q+2)*(r+1),(p+1)*(q+1)*(r+2))
    xz = Iz'*((Rr ⊗ (∂y ⊗ Rp)))
    yz = Iz'*((Rr ⊗ (Rq ⊗ ∂x)))

    return hcat(vcat(xx, xy, xz), vcat(yx, yy, yz), vcat(zx, zy, zz));
end

"""
FToH = _electric_vector_potential_to_magnetic_field_operator(number_of_cells, resolution, k_b)\n
This operators is used to compute the total magnetic field on the FDTD grid (Yee cell) given the magnetic vector potential.\n
k_b is the wavenumber in the background medium.\n
"""
function _magnetic_vector_potential_to_magnetic_field_operator(
    number_of_cells,
    resolution,
    k_b
)
    T = ComplexF32

    p = number_of_cells[1]
    q = number_of_cells[2]
    r = number_of_cells[3]
    Δx = Float32(resolution[1])
    Δy = Float32(resolution[2])
    Δz = Float32(resolution[3])

    k_b = Float32(k_b)
    # needed -> Hx : (P+1) x (Q+2) : (R+2), wanted -> Hx : (P-1) x (Q) : (R)
    # needed -> Hy : (P+2) x (Q+1) : (R+2), wanted -> Hy : (P) x (Q-1) : (R)
    # needed -> Hz : (P+2) x (Q+2) : (R+1), wanted -> Hz : (P) x (Q) : (R-1)
    #operators for Hx
    Rp = spdiagm(p-1, p+1, 1 => ones(T, p-1))
    Rq = spdiagm(q, q+2, 1 => ones(T, q))
    Rr = spdiagm(r, r+2, 1 => ones(T, r))

    Ix   = (Rr ⊗ (Rq ⊗ Rp))
    ∂x   = 1/Δx .*spdiagm(p-1, p+2, 1 => -one(T) .*ones(T, p-1), 2 => ones(T, p-1)) # this could be broken AF ?

    ∂y   = 1/Δy .*spdiagm(q, q+1, 0 => -one(T) .*ones(T, q), 1 => ones(T, q))
    ∂z   = 1/Δz .*spdiagm(r, r+1, 0 => -one(T) .*ones(T, r), 1 => ones(T, r))
    ∂x²  = 1/(Δx)^2 .*spdiagm(p-1, p+1, 0 => ones(T, p-1), 1 => -2 .*ones(T, p-1), 2 => ones(T, p-1))

    xx = Ix'*((Rr ⊗ (Rq ⊗ ∂x²)) + k_b^2 .*Ix) #f.x -> h.x
    yx = Ix'*((Rr ⊗ (∂y ⊗ ∂x))) # f.y -> h.x (P+2) x (Q+1) : (R+2) -> (P-1) x (Q) : (R)
    zx = Ix'*((∂z ⊗ (Rq ⊗ ∂x))) # f.z -> h.x (P+2) x (Q+2) : (R+1) -> (P-1) x (Q) : (R)

    #operators for Ey
    Rp = spdiagm(p, p+2, 1 => ones(T, p))
    Rq = spdiagm(q-1, q+1, 1 => ones(T, q-1))
    Rr = spdiagm(r, r+2, 1 => ones(T, r))

    Iy   = (Rr ⊗ (Rq ⊗ Rp))
    ∂x   = 1/Δx .*spdiagm(p, p+1, 0 => -one(T) .*ones(T, p), 1 => ones(T, p))
    ∂y   = 1/Δy .*spdiagm(q-1, q+2, 1 => -one(T) .*ones(T, q-1), 2 => ones(T, q-1))
    ∂z   = 1/Δz .*spdiagm(r, r+1, 0 => -one(T) .*ones(T, r), 1 => ones(T, r))
    ∂y²  = 1/(Δy)^2 .*spdiagm(q-1, q+1, 0 => ones(T, q-1), 1 => -2 .*ones(T, q-1), 2 => ones(T, q-1))

    xy = Iy'*((Rr ⊗ (∂y ⊗ ∂x)))
    yy = Iy'*((Rr ⊗ (∂y² ⊗ Rp)) + k_b^2 .*Iy)
    zy = Iy'*((∂z ⊗ (∂y ⊗ Rp)))

    #operators for Ez
    Rp = spdiagm(p, p+2, 1 => ones(T, p))
    Rq = spdiagm(q, q+2, 1 => ones(T, q))
    Rr = spdiagm(r-1, r+1, 1 => ones(T, r-1))

    Iz   = (Rr ⊗ (Rq ⊗ Rp))
    ∂x   = 1/Δx .*spdiagm(p, p+1, 0 => -one(T) .*ones(T, p), 1 => ones(T, p))
    ∂y   = 1/Δy .*spdiagm(q, q+1, 0 => -one(T) .*ones(T, q), 1 => ones(T, q))
    ∂z   = 1/Δz .*spdiagm(r-1, r+2, 1 => -one(T) .*ones(T, r-1), 2 => ones(T, r-1))
    ∂z²  = 1/(Δz)^2 .*spdiagm(r-1, r+1, 0 => ones(T, r-1), 1 => -2 .*ones(T, r-1), 2 => ones(T, r-1))

    xz = Iz'*((∂z ⊗ (Rq ⊗ ∂x)))
    yz = Iz'*((∂z ⊗ (∂y ⊗ Rp)))
    zz = Iz'*((∂z² ⊗ (Rq ⊗ Rp)) + k_b^2 .*Iz)

    return hcat(vcat(xx, xy, xz), vcat(yx, yy, yz), vcat(zx, zy, zz));
end

"""
Ig = _create_Greens_functions_restriction_operators_magnetic_field(number_of_cells)\n
These operators are used to restrict the output of ∫GKdV.\n
This integral is computed using an FFT on a domain that is larger the number of 2*cells.\n
Therefore, the output of these should be restricted to the correct Yee cell grid.\n
"""
function _create_Greens_functions_restriction_operators_magnetic_field(number_of_cells)
    T = ComplexF32
    p, q, r = number_of_cells
    dx, dy, dz = 2 .^(ceil.(Int, log2.(number_of_cells.+2).+1))

    Ix = spdiagm(p+1, dx, 0 => ones(T, p+1))
    Iy = spdiagm(q+2, dy, 0 => ones(T, q+2))
    Iz = spdiagm(r+2, dz, 0 => ones(T, r+2))
    x = (Iz ⊗ (Iy ⊗ Ix))

    Ix = spdiagm(p+2, dx, 0 => ones(T, p+2))
    Iy = spdiagm(q+1, dy, 0 => ones(T, q+1))
    y = (Iz ⊗ (Iy ⊗ Ix))

    Iy = spdiagm(q+2, dy, 0 => ones(T, q+2))
    Iz = spdiagm(r+1, dz, 0 => ones(T, r+1))
    z = (Iz ⊗ (Iy ⊗ Ix))

    return blockdiag(x, y, z)
end

"""
Ig = _create_Greens_functions_restriction_operators_magnetic_field_gpu(number_of_cells)\n
These operators are used to restrict the output of ∫GKdV.\n
This integral is computed using an FFT on a domain that is larger the number of 2*cells.\n
Therefore, the output of these should be restricted to the correct Yee cell grid.\n
"""
function _create_Greens_functions_restriction_operators_magnetic_field_gpu(number_of_cells)
    T = ComplexF32
    p, q, r = number_of_cells
    dx, dy, dz = 2 .^(ceil.(Int, log2.(number_of_cells.+2).+1))

    Ix = spdiagm(p+1, dx, 0 => ones(T, p+1))
    Iy = spdiagm(q+2, dy, 0 => ones(T, q+2))
    Iz = spdiagm(r+2, dz, 0 => ones(T, r+2))
    x = (Iz ⊗ (Iy ⊗ Ix))

    Ix = spdiagm(p+2, dx, 0 => ones(T, p+2))
    Iy = spdiagm(q+1, dy, 0 => ones(T, q+1))
    y = (Iz ⊗ (Iy ⊗ Ix))

    Iy = spdiagm(q+2, dy, 0 => ones(T, q+2))
    Iz = spdiagm(r+1, dz, 0 => ones(T, r+1))
    z = (Iz ⊗ (Iy ⊗ Ix))

    Ig = blockdiag(x, y, z) |> CUDA.CUSPARSE.CuSparseMatrixCSC

    return Ig
end


export RadialFlow

import Base: eltype, size, length

@doc raw"""
The RadialFlow function is defined as

```math
f({\bf{x}}) = {\bf{x}} + \frac{\beta({bf{z}} - {\bf{z}}_0)}{\alpha + |{\bf{z}} - {\bf{z}}_0|}
```

with input and output dimension ``D``. Here ``{\bf{x}}\in \mathbb{R}^D`` represents the input of the function. Furthermore ``{\bf{z}}_0\in \mathbb{R}^D``, ``\alpha\in \mathbb{R}`` and ``\beta\in\mathbb{R}`` represent the parameters of the function. The function contracts and expands the input space.

This function has been introduced in:

Rezende, Danilo, and Shakir Mohamed. "Variational inference with normalizing flows." _International conference on machine learning._ PMLR, 2015.
"""
mutable struct RadialFlow{T1, T2 <: Real} <: AbstractCouplingFlow
    z0      :: T1
    α       :: T2
    β       :: T2
    function RadialFlow(z0::T1, α::T2, β::T2) where { T1, T2 <: Real}
        @assert α > 0 "The parameter α should be larger than 0."
        return new{T1,T2}(z0, float(α), float(β))
    end
end
struct RadialFlowEmpty{N} <: AbstractCouplingFlowEmpty end
RadialFlowEmpty(dim::Int) = RadialFlowEmpty{dim}()
struct RadialFlowPlaceholder <: AbstractCouplingFlowPlaceholder end
@doc raw"""
The `RadialFlow()` function creates a planar flow object. Its dimensionality is automatically set when wrapping this object inside a model. Its parameters are initialized during compilation.
"""
RadialFlow() = RadialFlowPlaceholder()
function prepare(dim::Int, flow::RadialFlowPlaceholder)
    return RadialFlowEmpty(dim)
end

# compile placeholder
compile(f::RadialFlowEmpty{1})          = RadialFlow(randn(), rand(), randn())
compile(f::RadialFlowEmpty)             = RadialFlow(randn(getdim(f)), rand(), randn())
compile(f::RadialFlowEmpty{1}, params)  = RadialFlow(params[1], params[2], params[3])
compile(f::RadialFlowEmpty, params)     = RadialFlow(params[1:getdim(f)], params[getdim(f)+1], params[getdim(f)+2])

@doc raw"""
The `RadialFlow(dim::Int64)` function creates a mutable `RadialFlow` structure with parameters corresponding to input of dimensions `dim`. The parameters are each random sampled from a standard (multivariate) normal distribution.
"""
function RadialFlow(dim::Int64)
    return RadialFlow(randn(dim), rand(), randn())
end

# number of parameters
nr_params(flow::RadialFlow)      = 2 + length(flow.z0)
nr_params(flow::RadialFlowEmpty) = 2 + getdim(flow)

# get-functions for the RadialFlow structure.
getz0(f::RadialFlow)                      = return f.z0
getα(f::RadialFlow)                       = return f.α
getβ(f::RadialFlow)                       = return f.β
getall(f::RadialFlow)                     = return f.z0, f.α, f.β
getdim(f::RadialFlowEmpty{N}) where { N } = return N

# set-functions for the RadialFlow structure
function setz0!(f::RadialFlow{T1,T2}, z0::T1) where { T1, T2 <: Real}
    @assert length(f.z0) == length(z0) "The dimensionality of the current value of z0 and its new value do not match."
    f.z0 = z0
end

function setα!(f::RadialFlow{T1,T2}, α::T2) where { T1, T2 <: Real}
    f.α = α
end

function setβ!(f::RadialFlow{T1,T2}, β::T2) where { T1, T2 <: Real }
    f.β = β
end

# custom Base function for the RadialFlow structure
eltype(f::RadialFlow{T1,T2}) where { T1 <: Real, T2 <: Real}                 = promote_type(T1, T2)
eltype(f::RadialFlow{T1,T2}) where { T1 <: AbstractVector, T2 <: Real}       = promote_type(eltype(T1), T2)
eltype(::Type{RadialFlow{T1,T2}}) where { T1 <: Real, T2 <: Real}            = promote_type(T1, T2)
eltype(::Type{RadialFlow{T1,T2}}) where { T1 <: AbstractVector, T2 <: Real}  = promote_type(eltype(T1), T2)

size(f::RadialFlow{T1,T2}) where { T1 <: Real, T2 <: Real}                   = 1
size(f::RadialFlow{T1,T2}) where { T1 <: AbstractVector, T2 <: Real}         = length(f.z0)
size(f::RadialFlowEmpty{N}) where { N }                                      = return N

length(f::RadialFlow{T1,T2}) where { T1 <: Real, T2 <: Real}                 = 1
length(f::RadialFlow{T1,T2}) where { T1 <: AbstractVector, T2 <: Real}       = length(f.z0)
length(f::RadialFlowEmpty{N}) where { N }                                    = return N

# forward pass through the RadialFlow function (multivariate input)
function _forward(f::RadialFlow{T1,T2}, input::T1) where { T1, T2 <: Real }

    # fetch values
    z0, α, β = getall(f)

    # calculate result
    denominator = α + norm(input - z0) # Not sure whether this is the correct norm
    denominator /= β

    result = copy(input)
    result .-= z0
    result ./= denominator
    result .+= input

    # return result
    return result

end
forward(f::RadialFlow{T1,T2}, input::T1) where { T1, T2 <: Real } = _forward(f, input)
Broadcast.broadcasted(::typeof(forward), f::RadialFlow{T1,T2}, input::AbstractVector{T1}) where { T1, T2 <: Real } = broadcast(_forward, Ref(f), input)

# forward pass through the RadialFlow function (univariate input)
function _forward(f::RadialFlow{T1,T2}, input::T3) where { T1 <: Real, T2 <: Real, T3 <: Real }

    # fetch values
    z0, α, β = getall(f)

    # calculate result
    denominator = α + norm(input - z0) # Not sure whether this is the correct norm
    denominator /= β

    result = copy(input)
    result -= z0
    result /= denominator
    result += input

    # return result
    return result

end
forward(f::RadialFlow{T1,T2}, input::T3) where { T1 <: Real, T2 <: Real, T3 <: Real } = _forward(f, input)
Broadcast.broadcasted(::typeof(forward), f::RadialFlow{T1,T2}, input::AbstractVector{ <: Real }) where { T1 <: Real, T2 <: Real } = broadcast(_forward, Ref(f), input)

function _forward(f::RadialFlow{T1,T2}, input) where { T1 <: Real, T2 <: Real}
    # function when the input is an array with 1 element
    @assert length(input) == 1 "Something is wrong with the dimensionality of the input to the RadialFlow flow."
    return forward(f, input[1])
end
forward(f::RadialFlow{T1,T2}, input) where { T1 <: Real, T2 <: Real } = _forward(f, input)
Broadcast.broadcasted(::typeof(forward), f::RadialFlow{T1,T2}, input::AbstractVector) where { T1 <: Real, T2 <: Real } = broadcast(_forward, Ref(f), input)

# inplace forward pass through the RadialFlow function (multivariate input)
function forward!(output::T1, f::RadialFlow{T1,T2}, input::T1) where { T1, T2 <: Real }

    # check dimensionality
    @assert length(output) == length(input) "The length of the preallocated vector does not seem to match the length of the input vector."

    # fetch values
    z0, α, β = getall(f)

    # calculate result
    denominator = α + norm(input - z0) # Not sure whether this is the correct norm
    denominator /= β

    output .= input
    output .-= z0
    output ./= denominator
    output .+= input

end

# jacobian of the RadialFlow function (multivariate input)
function _jacobian(f::RadialFlow{T1,T2}, input::T1) where { T1, T2 <: Real}

    # fetch values
    z0, α, β = getall(f)

    # # calculate result
    diff = input - z0
    result = diff*diff'
    hi = α + norm(diff)
    βh = β/hi
    ζ  = -β / hi / hi / norm(diff)
    result .*= ζ
    @inbounds for k = 1:length(input)
        result[k,k] += 1
        result[k,k] += βh
    end

    # return result
    return result

end
jacobian(f::RadialFlow{T1,T2}, input::T1) where { T1, T2 <: Real } = _jacobian(f, input)
Broadcast.broadcasted(::typeof(jacobian), f::RadialFlow{T1,T2}, input::AbstractVector{T1}) where { T1, T2 <: Real } = broadcast(_jacobian, Ref(f), input)

function _jacobian(f::RadialFlow{T1,T2}, input) where { T1 <: Real, T2 <: Real }
    # function when the input is an array with 1 element
    @assert length(input) == 1 "Something is wrong with the dimensionality of the input to the RadialFlow flow."
    return jacobian(f, input[1])
end
jacobian(f::RadialFlow{T1,T2}, input) where { T1 <: Real, T2 <: Real } = _jacobian(f, input)
Broadcast.broadcasted(::typeof(jacobian), f::RadialFlow{T1,T2}, input::AbstractVector) where { T1 <: Real, T2 <: Real } = broadcast(_jacobian, Ref(f), input)

# jacobian of the RadialFlow function (univariate input)
function _jacobian(f::RadialFlow{T1,T2}, input::T3) where { T1 <: Real, T2 <: Real, T3 <: Real }

    # fetch values
    z0, α, β = getall(f)

    # calculate result (optimized)
    diff = input - z0
    h = 1 / (α + norm(diff))
    result = 1 + β*h - β * h * h /norm(diff) * diff * diff

    # return result
    return result

end
jacobian(f::RadialFlow{T1,T2}, input::T3) where { T1 <: Real, T2 <: Real, T3 <: Real } = _jacobian(f, input)
Broadcast.broadcasted(::typeof(jacobian), f::RadialFlow{T1,T2}, input::AbstractVector{ <: Real }) where { T1 <: Real, T2 <: Real } = broadcast(_jacobian, Ref(f), input)

# inplace jacobian of the RadialFlow function (multivariate input)
function jacobian!(output::AbstractMatrix{T2}, f::RadialFlow{T1,T2}, input::T1) where { T1, T2 <: Real}

    # check whether the dimensionality is correct
    @assert size(output) == (length(input), length(f.z0)) "The dimensionality of the preallocated jacobian matrix seems incorrect."

    # fetch values
    z0, α, β = getall(f)

    # # calculate result
    diff = input - z0
    for ku = 1:length(diff)
        for kw = 1:length(diff)
            output[ku,kw] = diff[ku]*diff[kw]
        end
    end
    hi = α + norm(diff)
    βh = β/hi
    ζ  = -β / hi / hi / norm(diff)
    output .*= ζ
    @inbounds for k = 1:length(input)
        output[k,k] += 1
        output[k,k] += βh
    end

end

# determinant of the jacobian of the RadialFlow function (multivariate input)
det_jacobian(f::RadialFlow{T1,T2}, input::T1) where { T1, T2 <: Real} = det(jacobian(f, input))

# determinant of the jacobian of the RadialFlow function (univariate input)
function det_jacobian(f::RadialFlow{T, T}, input::T) where { T <: Real }

    # fetch values
    z0, α, β = getall(f)

    # calculate result
    r = norm(input - z0)
    h = 1 / (α + r)
    result = ( 1 + β*h - β*h^2*r ) * ( 1 + β*h )^(length(input)-1)

    # return result
    return result

end

# extra utility function (multivariate)
inv_jacobian(f::RadialFlow{T1,T2}, input::T1) where { T1, T2 <: Real }       = inv(jacobian(f, input))
absdet_jacobian(f::RadialFlow{T1,T2}, input::T1) where { T1, T2 <: Real }    = abs(det_jacobian(f, input))
logabsdet_jacobian(f::RadialFlow{T1,T2}, input::T1) where { T1, T2 <: Real}  = logabsdet(jacobian(f, input))

# extra utility functions (univariate)
inv_jacobian(f::RadialFlow{T,T}, input::T) where { T <: Real }               = 1.0 / jacobian(f, input)
absdet_jacobian(f::RadialFlow{T,T}, input::T) where { T <: Real }            = abs(det_jacobian(f, input))
logabsdet_jacobian(f::RadialFlow{T,T}, input::T) where { T <: Real}          = log(absdet_jacobian(f, input))

function inv_jacobian(f::RadialFlow{T1,T2}, input::T3) where { T1 <: Real, T2 <: Real, T3 }
    # function when the input is an array with 1 element
    @assert length(input) == 1 "Something is wrong with the dimensionality of the input to the RadialFlow flow."
    return inv_jacobian(f, input[1])
end


using Turing, CSV

ProjDir = @__DIR__

include("../src/quap_turing.jl")

df = CSV.read(joinpath(@__DIR__, "..", "data", "Howell1.csv"))

# Use only adults and center the weight observations

df2 = filter(row -> row.age >= 18, df)
mean_weight = mean(df2.weight)
df2.weight_c = df2.weight .- mean_weight
first(df2, 5)

@model height(heights) = begin
    μ ~ Normal(178, 20)
    σ ~ Uniform(0, 50)
    heights .~ Normal(μ, σ)
end

m = height(df2.height)
res = quap(m)
MvNormal(res.coef, res.vcov)