"""
    struct CylindricalImpurityDensity{T <: SSDFloat} <: AbstractImpurityDensity{T}

Impurity density model which assumes a linear gradient in impurity density in each spatial dimension
of a cylindrical coordinate system.

## Fields
* `offsets::NTuple{3,T}`: impurity density values at the origin of each dimension.
* `gradients::NTuple{3,T}`: linear slopes in `r` and `z` direction.

## Definition in Configuration File

A `CylindricalImpurityDensity` is defined in the configuration file through the field `impurity_density`
(of a `passive` or `surrounding`) with `name: cylindrical` and optional fields `r` and `z`
that can each contain `init` for initial values at 0 and `gradient` for gradients in that dimension.

An example definition of a cylindrical impurity density looks like this:
```yaml
impurity_density:
  name: cylindrical
  r:  # impurity profile with linear gradient in r
    init: 1.0e10     # 1/m³
    gradient: 1.0e11 # 1/m⁴
```
"""
struct CylindricalImpurityDensity{T <: SSDFloat} <: AbstractImpurityDensity{T}
    offsets::NTuple{3, T}
    gradients::NTuple{3, T}
end

function ImpurityDensity(T::DataType, t::Val{:cylindrical}, dict::Union{Dict{String, Any}, Dict{Any, Any}}, input_units::NamedTuple)
    offsets, gradients = zeros(T,3), zeros(T,3)
    density_unit = input_units.length^(-3)
    density_gradient_unit = input_units.length^(-4)
    if prod(map(k -> k in ["r","z"], collect(keys(dict)))) @warn "Only r and z are supported in the cylindrical impurity density model.\nChange the impurity density model in the config file or remove all other entries." end
    if haskey(dict, "r")
        if haskey(dict["r"], "init")     offsets[1]   = _parse_value(T, dict["r"]["init"], density_unit) end
        if haskey(dict["r"], "gradient") gradients[1] = _parse_value(T, dict["r"]["gradient"], density_gradient_unit) end
    end
    if haskey(dict, "z")
        if haskey(dict["z"], "init")     offsets[3]   = _parse_value(T, dict["z"]["init"], density_unit) end
        if haskey(dict["z"], "gradient") gradients[3] = _parse_value(T, dict["z"]["gradient"], density_gradient_unit) end
    end
    CylindricalImpurityDensity{T}( NTuple{3, T}(offsets), NTuple{3, T}(gradients) )
end

function get_impurity_density(lcdm::CylindricalImpurityDensity{T}, pt::AbstractCoordinatePoint{T})::T where {T <: SSDFloat}
    pt::CylindricalPoint{T} = CylindricalPoint(pt)
    ρ::T = 0
    for i in eachindex(lcdm.offsets)
        ρ += (lcdm.offsets[i] + pt[i] * lcdm.gradients[i])
    end
    return ρ
end


mutable struct CompSumAcc{T, EFT}
    s :: T
    e :: T
end

@inline compSumAcc(EFT) = T->compSumAcc(EFT, T)
@inline compSumAcc(EFT, T) = CompSumAcc{T, EFT}(vzero(T), vzero(T))

@inline function add!(acc::CompSumAcc{T, EFT}, x::T) where {T, EFT}
    SIMDops.@explicit

    acc.s, e = EFT(acc.s, x)
    acc.e += e
end

@inline function add!(acc::A, x::A) where {A<:CompSumAcc{T, EFT}} where {T, EFT}
    SIMDops.@explicit

    acc.s, e = EFT(acc.s, x.s)
    acc.e += x.e + e
end

@inline function Base.sum(acc::CompSumAcc{T, EFT}) where {T<:Vec, EFT}
    acc_r = compSumAcc(EFT, fptype(T))
    acc_r.e = vsum(acc.e)
    for xi in acc.s
        acc_r.s, ei = EFT(acc_r.s, xi.value)
        acc_r.e += ei
    end
    acc_r
end

@inline function Base.sum(acc::CompSumAcc{T, EFT}) where {T<:Real, EFT}
    acc.s + acc.e
end


using Distributions
using Base.LinAlg.BLAS

function sample(::Type{Bernoulli}, means::Mat{Float64})
    s = zeros(means)
    r = rand(size(means))
    @simd for i=1:length(means)
        @inbounds s[i] = r[i] < means[i] ? 1.0 : 0.0
    end

    return s
end

function sample(::Type{Gaussian}, means::Mat{Float64})
    sigma2 = 1                   # using fixed standard diviation
    samples = zeros(size(means))
    for j=1:size(means, 2), i=1:size(means, 1)
        samples[i, j] = rand(Normal(means[i, j], sigma2))
    end
    return samples
end

function sample_hiddens{V,H}(rbm::RBM{V,H}, vis::Mat{Float64})
    means = ProbHidCondOnVis(rbm, vis)
    return sample(H, means), means
end

function sample_visibles{V,H}(rbm::RBM{V,H}, hid::Mat{Float64})
    means = ProbVisCondOnHid(rbm, hid)
    return sample(V, means), means
end

function gibbs(rbm::RBM, vis::Mat{Float64}; n_times=1)
    v_pos = vis
    h_samp, h_pos = sample_hiddens(rbm, v_pos)
    h_neg = Array(Float64,0,0)::Mat{Float64}
    v_neg = Array(Float64,0,0)::Mat{Float64}
    if n_times > 0
    # Save computation by setting `n_times=0` in the case
    # of persistent CD.
        v_neg = sample_visibles(rbm, h_samp)
        h_samp, h_neg = sample_hiddens(rbm, v_neg)
        for i=1:n_times-1
            v_neg = sample_visibles(rbm, h_samp)
            h_samp, h_neg = sample_hiddens(rbm, v_neg)
        end
    end
    return v_pos, h_pos, v_neg, h_neg
end

function MCMC(rbm::RBM, init::Mat{Float64}; iterations=1, StartMode="visible")
    if StartMode == "visible"
    # In this first mode we assume that we are starting from the visible samples. E.g. in
    # the case of binary RBM, we should be starting with binary samples.
        vis_samples = copy(init)                                    # Start from the visible samples
        vis_means   = copy(init)                                    # Giving a starting point for the means
        hid_samples, hid_means = sample_hiddens(rbm,vis_samples)     # Get the first hidden means [NMF-ish]
    end

    if StartMode == "hidden"
    # In this second mode we assume that we are starting from a set of hidden
    # samples. Because of this, we increment the iteration count by 1
        hid_samples = copy(init)
        hid_means = copy(init)
        iterations+=1
    end

    for i=1:iterations-1
        vis_samples, vis_means = sample_visibles(rbm,hid_samples)          # Sample the visible units from true distribution
        hid_samples, hid_means = sample_hiddens(rbm,vis_samples)           # Update the hidden unit means, a NMF-ish approach
    end

    return vis_samples, hid_samples, vis_means, hid_means
end


# test whether prox satisfies some conditions

using LinearAlgebra
using SparseArrays
using ProximalOperators
using Test

stuff = [
  Dict( "constr" => LeastSquares,
        "params" => ( (randn(20, 10), randn(20)), (randn(15, 40), randn(15), rand()), (rand(Complex{Float64}, 15, 40), rand(Complex{Float64}, 15), rand()), (sprandn(100,1000,0.05), randn(100), rand()) ),
        "args"   => ( randn(10), randn(40), rand(Complex{Float64}, 40), randn(1000) ),
        "gammas" => ( rand(), rand(), rand(), rand() ),
        # test first order optimality condition of prox
        "test"   => (f, x, gamma, y) -> norm(y + gamma*f.lambda*(f.A'*(f.A*y-f.b)) - x) <= 1e-10
      ),

  Dict( "constr" => IndSimplex,
        "params" => ( (), (2,), (5.0,), (rand(),) ),
        "args"   => ( randn(10), randn(10), randn(10), randn(10) ),
        "gammas" => ( rand(), rand(), rand(), rand() ),
        # test y belonging to the simplex
        "test"   => (f, x, gamma, y) -> all(y .>= 0.0) && abs(sum(y)-f.a) <= (1+f.a)*1e-12
      ),

  Dict( "constr" => IndBallL1,
        "params" => ( (), (1.7,), (5.0,), (rand(), ) ),
        "args"   => ( randn(10), randn(10), randn(10), randn(10) ),
        "gammas" => ( rand(), rand(), rand(), rand() ),
        # test y belonging to the L1 ball
        "test"   => (f, x, gamma, y) -> norm(y, 1) <= (1+1e-12)*f.r
      ),

  Dict( "constr" => HuberLoss,
        "params" => ( (), (rand(),), (rand(), rand()), (rand(), rand()) ),
        "args"   => ( randn(10), randn(8, 10), randn(20), rand(Complex{Float64}, 12, 15) ),
        "gammas" => ( rand(), rand(), rand(), rand() ),
        # test optimality condition of prox
        "test"   => (f, x, gamma, y) -> isapprox((x-y)/gamma, (norm(y) <= f.rho ? f.mu*y : f.rho*f.mu*y/norm(y)))
      ),

  Dict( "constr" => SqrHingeLoss,
        "params" => ( (randn(5),), (randn(5), 0.1+rand()), (randn(3, 5), 0.1+rand()) ),
        "args"   => ( randn(5), randn(5), randn(3, 5) ),
        "gammas" => ( 0.1+rand(), 0.1+rand(), 0.1+rand() ),
        # test optimality condition of prox
        "test"   => (f, x, gamma, z) -> isapprox((x-z)/gamma, -2 .* f.mu.*f.y.*max.(0, 1 .- f.y.*z))
      ),
]

@testset "$(d["constr"])" for d in stuff
      constr = d["constr"]
      params = d["params"]
      args = d["args"]
      gammas = d["gammas"]
      test = d["test"]
      for i = 1:length(params)
            f = constr(params[i]...)
            y, fy = prox(f, args[i], gammas[i])
            @test test(f, args[i], gammas[i], y)
      end
end


#xor-ml.jl - set up to test the xor function in julia.

using SigmoidNumbers
using GenML
using Base.Test

PType = Float32#Posit{12,0}

#assign sigmoid to the pseudologistic
GenML.TF.sigmoid(x::Posit) = SigmoidNumbers.pseudologistic(x)
GenML.TF.dasysigmoid(y::Posit) = SigmoidNumbers.delta_psl(y)
GenML.CF.logloss(x::Posit) = SigmoidNumbers.pseudologcost(x)

xornet = GenML.MLP.MultilayerPerceptron{PType, (2,2,1)}()

#hand-written transition matrices.
xornet.layers[1].transition[:] = PType[5.0 5.0; -10.0 -10.0]
xornet.layers[1].bias[:] = PType[-7.5, 7.5]
xornet.layers[2].transition[:] = PType[-10.0 -10.0]
xornet.layers[2].bias[:] = PType[5.0]

#test the possible input/output pairs.

println("attempting premade xor neural net, with posit $PType")

@test xornet([true, true])[1] < PType(0.5)
@test xornet([true, false])[1] > PType(0.5)
@test xornet([false, true])[1] > PType(0.5)
@test xornet([false, false])[1] < PType(0.5)

println("execution of premade xor neural net successful with $PType")

println("attempting to train xor neural net, with posit $PType")

Base.randn(::Type{PType}, dims::Integer...) = map(PType, randn(dims...))
Base.rand(::Type{PType}, dims::Integer...) = map(PType, rand(dims...))

function xortrain()

  println("working on unreliable xor data set with backpropagation")

  input_matrix = rand(Bool, 10, 500)
  training_results = Array{Bool,2}(1,500)
  training_results[:] = [input_matrix[1, col] $ input_matrix[2, col] $ (rand() < 0.05) for col in 1:size(input_matrix,2)]

  xornet = GenML.MLP.MultilayerPerceptron{PType,(10,2,1)}(randn)

  for rounds = 1:40
  for idx = 1:500
    input_column = input_matrix[:, idx]
    training_column = training_results[:, idx]

    GenML.Optimizers.backpropagationoptimize(xornet, input_column, training_column, GenML.CF.crossentropy)
  end
  end

  #verify that the optimization has resulted in a good data set.
  wrongcount = 0
  for x = 1:50
    input_vector = rand(Bool, 10)
    wrongcount += (xornet(input_vector)[1] > PType(0.5)) != (input_vector[1] $ input_vector[2])
  end
  println("incorrect responses for noisy xor: $wrongcount / 50 => $(2 * (50 - wrongcount)) % correct")

  return wrongcount
end

xortrain() == 0 || xortrain() == 0 || xortrain() == 0 || xortrain() == 0 || xortrain() == 0 || xortrain() == 0 || xortrain() == 0 || xortrain() == 0


using LinearAlgebra
using LightGraphs
using SimpleWeightedGraphs
using GraphHierarchy
using SparseArrays
using Test

# (1,1)-(2,2)

@testset "GraphHierarchy.jl" begin
    # 3-chain
    gw = SimpleWeightedDiGraph(3)
    add_edge!(gw,1,2,2)
    add_edge!(gw,2,3,2)
    # forward: not weighted, not big
    HS = forward_hierarchical_structure(gw, weighted = false)
    @test norm(HS[1] - [-1,0,1]) < 1.0e-12
    @test norm(HS[2] - [1,0,0]) < 1.0e-12
    @test norm(HS[3] - [1,1,0]) < 1.0e-12
    @test norm(HS[4] - [ 0 1 0 ; 0 0 1 ; 0 0 0 ]) < 1.0e-12
    HC = hierarchical_coefficients(gw, HS, weighted = false)
    @test abs(HC[1]) < 1.0e-12
    @test abs(HC[2]) < 1.0e-12
    # forward: weighted, not big
    HS = forward_hierarchical_structure(gw)
    @test norm(HS[1] - [-1,0,1]) < 1.0e-12
    @test norm(HS[2] - [1,0,0]) < 1.0e-12
    @test norm(HS[3] - [1,1,0]) < 1.0e-12
    @test norm(HS[4] - [ 0 1 0 ; 0 0 1 ; 0 0 0 ]) < 1.0e-12
    HC = hierarchical_coefficients(gw, HS)
    @test abs(HC[1]) < 1.0e-12
    @test abs(HC[2]) < 1.0e-12
    # forward: not weighted, big
    HS = forward_hierarchical_structure(gw, weighted = false, big = true)
    @test norm(HS[1] - [-1,0,1]) < 1.0e-70
    @test norm(HS[2] - [1,0,0]) < 1.0e-70
    @test norm(HS[3] - [1,1,0]) < 1.0e-70
    @test norm(HS[4] - [ 0 1 0 ; 0 0 1 ; 0 0 0 ]) < 1.0e-70
    HC = hierarchical_coefficients(gw, HS, weighted = false)
    @test abs(HC[1]) < 1.0e-70
    @test abs(HC[2]) < 1.0e-70
    # forward: weighted, big
    HS = forward_hierarchical_structure(gw, big = true)
    @test norm(HS[1] - [-1,0,1]) < 1.0e-70
    @test norm(HS[2] - [1,0,0]) < 1.0e-70
    @test norm(HS[3] - [1,1,0]) < 1.0e-70
    @test norm(HS[4] - [ 0 1 0 ; 0 0 1 ; 0 0 0 ]) < 1.0e-70
    HC = hierarchical_coefficients(gw, HS)
    @test abs(HC[1]) < 1.0e-70
    @test abs(HC[2]) < 1.0e-70
    # backward: not weighted, not big
    HS = backward_hierarchical_structure(gw, weighted = false)
    @test norm(HS[1] - [1,0,-1]) < 1.0e-12
    @test norm(HS[2] - [0,0,1]) < 1.0e-12
    @test norm(HS[3] - [0,1,1]) < 1.0e-12
    @test norm(HS[4] - [ 0 1 0 ; 0 0 1 ; 0 0 0 ]) < 1.0e-12
    HC = hierarchical_coefficients(gw, HS, weighted = false)
    @test abs(HC[1]) < 1.0e-12
    @test abs(HC[2]) < 1.0e-12
    # backward: weighted, not big
    HS = backward_hierarchical_structure(gw)
    @test norm(HS[1] - [1,0,-1]) < 1.0e-12
    @test norm(HS[2] - [0,0,1]) < 1.0e-12
    @test norm(HS[3] - [0,1,1]) < 1.0e-12
    @test norm(HS[4] - [ 0 1 0 ; 0 0 1 ; 0 0 0 ]) < 1.0e-12
    HC = hierarchical_coefficients(gw, HS)
    @test abs(HC[1]) < 1.0e-12
    @test abs(HC[2]) < 1.0e-12
    # backward: not weighted, big
    HS = backward_hierarchical_structure(gw, weighted = false, big = true)
    @test norm(HS[1] - [1,0,-1]) < 1.0e-70
    @test norm(HS[2] - [0,0,1]) < 1.0e-70
    @test norm(HS[3] - [0,1,1]) < 1.0e-70
    @test norm(HS[4] - [ 0 1 0 ; 0 0 1 ; 0 0 0 ]) < 1.0e-70
    HC = hierarchical_coefficients(gw, HS, weighted = false)
    @test abs(HC[1]) < 1.0e-70
    @test abs(HC[2]) < 1.0e-70
    # backward: weighted, big
    HS = backward_hierarchical_structure(gw, big = true)
    @test norm(HS[1] - [1,0,-1]) < 1.0e-70
    @test norm(HS[2] - [0,0,1]) < 1.0e-70
    @test norm(HS[3] - [0,1,1]) < 1.0e-70
    @test norm(HS[4] - [ 0 1 0 ; 0 0 1 ; 0 0 0 ]) < 1.0e-70
    HC = hierarchical_coefficients(gw, HS)
    @test abs(HC[1]) < 1.0e-70
    @test abs(HC[2]) < 1.0e-70

    # 3-cycle
    gw = SimpleWeightedDiGraph(3)
    add_edge!(gw,1,2,1)
    add_edge!(gw,2,3,1)
    add_edge!(gw,3,1,2)
    # forward: not weighted, not big
    HS = forward_hierarchical_structure(gw, weighted = false)
    @test norm(HS[1]) < 1.0e-12
    @test norm(HS[2] - [1,1,1]) < 1.0e-12
    @test norm(HS[3]) < 1.0e-12
    @test norm(HS[4]) < 1.0e-12
    HC = hierarchical_coefficients(gw, HS, weighted = false)
    @test abs(HC[1] - 1) < 1.0e-12
    @test abs(HC[2]) < 1.0e-12
    # forward: weighted, not big
    HS = forward_hierarchical_structure(gw)
    @test norm(HS[1] - [1//3,0,-1//3]) < 1.0e-12
    @test norm(HS[2] - [1//3, 4//3, 4//3]) < 1.0e-12
    @test norm(HS[3] - [-1//3,-1//3,2//3]) < 1.0e-12
    @test norm(HS[4] - [ 0 -1//3 0 ; 0 0 -1//3 ; 2//3 0 0 ]) < 1.0e-12
    HC = hierarchical_coefficients(gw, HS)
    @test abs(HC[1] - 5//6) < 1.0e-12
    @test abs(HC[2] - 1//2) < 1.0e-12
    # forward: not weighted, big
    HS = forward_hierarchical_structure(gw, weighted = false, big = true)
    @test norm(HS[1]) < 1.0e-70
    @test norm(HS[2] - [1,1,1]) < 1.0e-70
    @test norm(HS[3]) < 1.0e-70
    @test norm(HS[4]) < 1.0e-70
    HC = hierarchical_coefficients(gw, HS, weighted = false)
    @test abs(HC[1] - 1) < 1.0e-70
    @test abs(HC[2]) < 1.0e-70
    # forward: weighted, big
    HS = forward_hierarchical_structure(gw, big = true)
    @test norm(HS[1] - [1//3,0,-1//3]) < 1.0e-70
    @test norm(HS[2] - [1//3, 4//3, 4//3]) < 1.0e-70
    @test norm(HS[3] - [-1//3,-1//3,2//3]) < 1.0e-70
    @test norm(HS[4] - [ 0 -1//3 0 ; 0 0 -1//3 ; 2//3 0 0 ]) < 1.0e-70
    HC = hierarchical_coefficients(gw, HS)
    @test abs(HC[1] - 5//6) < 1.0e-70
    @test abs(HC[2] - 1//2) < 1.0e-70
    # backward: not weighted, not big
    HS = backward_hierarchical_structure(gw, weighted = false)
    @test norm(HS[1]) < 1.0e-12
    @test norm(HS[2] - [1,1,1]) < 1.0e-12
    @test norm(HS[3]) < 1.0e-12
    @test norm(HS[4]) < 1.0e-12
    HC = hierarchical_coefficients(gw, HS, weighted = false)
    @test abs(HC[1] - 1) < 1.0e-12
    @test abs(HC[2]) < 1.0e-12
    # # backward: weighted, not big
    HS = backward_hierarchical_structure(gw)
    @test norm(HS[1] - [-1//3,0,1//3]) < 1.0e-12
    @test norm(HS[2] - [4//3,4//3,1//3]) < 1.0e-12
    @test norm(HS[3] - [2//3,-1//3,-1//3]) < 1.0e-12
    @test norm(HS[4] - [ 0 -1//3 0 ; 0 0 -1//3 ; 2//3 0 0 ]) < 1.0e-12
    HC = hierarchical_coefficients(gw, HS)
    @test abs(HC[1] - 5//6) < 1.0e-12
    @test abs(HC[2] - 1//2) < 1.0e-12
    # backward: not weighted, big
    HS = backward_hierarchical_structure(gw, weighted = false, big = true)
    @test norm(HS[1]) < 1.0e-70
    @test norm(HS[2] - [1,1,1]) < 1.0e-70
    @test norm(HS[3]) < 1.0e-70
    @test norm(HS[4]) < 1.0e-70
    HC = hierarchical_coefficients(gw, HS, weighted = false)
    @test abs(HC[1] - 1) < 1.0e-70
    @test abs(HC[2]) < 1.0e-70
    # backward: weighted, big
    HS = backward_hierarchical_structure(gw, big = true)
    @test norm(HS[1] - [-1//3,0,1//3]) < 1.0e-70
    @test norm(HS[2] - [4//3,4//3,1//3]) < 1.0e-70
    @test norm(HS[3] - [2//3,-1//3,-1//3]) < 1.0e-70
    @test norm(HS[4] - [ 0 -1//3 0 ; 0 0 -1//3 ; 2//3 0 0 ]) < 1.0e-70
    HC = hierarchical_coefficients(gw, HS)
    @test abs(HC[1] - 5//6) < 1.0e-70
    @test abs(HC[2] - 1//2) < 1.0e-70

    # pair with stem
    gw = SimpleWeightedDiGraph(3)
    add_edge!(gw,1,2,2)
    add_edge!(gw,2,3,1)
    add_edge!(gw,2,1,1)
    # forward: not weighted, not big
    HS = forward_hierarchical_structure(gw, weighted = false)
    @test norm(HS[1] - [-1//3, -1//3, 2//3]) < 1.0e-12
    @test norm(HS[2] - [1,1,0]) < 1.0e-12
    @test norm(HS[3] - [0,1//2,0]) < 1.0e-12
    @test norm(HS[4] - [ 0 0 0 ; 0 0 1 ; 0 0 0 ]) < 1.0e-12
    HC = hierarchical_coefficients(gw, HS, weighted = false)
    @test abs(HC[1] - 2//3) < 1.0e-12
    @test abs(HC[2] - sqrt(2)/3) < 1.0e-12
    # forward: weighted, not big
    HS = forward_hierarchical_structure(gw)
    @test norm(HS[1] - [-11//15,-2//15,13//15]) < 1.0e-12
    @test norm(HS[2] - [8//5,2//5,0]) < 1.0e-12
    @test norm(HS[3] - [3//5,1//5,0]) < 1.0e-12
    @test norm(HS[4] - [ 0 3//5 0 ; -3//5 0 1 ; 0 0 0 ]) < 1.0e-12
    HC = hierarchical_coefficients(gw, HS)
    @test abs(HC[1] - 3//5) < 1.0e-12
    @test abs(HC[2] - 3//5) < 1.0e-12
    # forward: not weighted, big
    HS = forward_hierarchical_structure(gw, weighted = false, big = true)
    @test norm(HS[1] - [-1//3, -1//3, 2//3]) < 1.0e-70
    @test norm(HS[2] - [1,1,0]) < 1.0e-70
    @test norm(HS[3] - [0,1//2,0]) < 1.0e-70
    @test norm(HS[4] - [ 0 0 0 ; 0 0 1 ; 0 0 0 ]) < 1.0e-70
    HC = hierarchical_coefficients(gw, HS, weighted = false)
    @test abs(HC[1] - 2//3) < 1.0e-70
    @test abs(HC[2] - sqrt(big(2))/3) < 1.0e-70
    # forward: weighted, big
    HS = forward_hierarchical_structure(gw, big = true)
    @test norm(HS[1] - [-11//15,-2//15,13//15]) < 1.0e-70
    @test norm(HS[2] - [8//5,2//5,0]) < 1.0e-70
    @test norm(HS[3] - [3//5,1//5,0]) < 1.0e-70
    @test norm(HS[4] - [ 0 3//5 0 ; -3//5 0 1 ; 0 0 0 ]) < 1.0e-70
    HC = hierarchical_coefficients(gw, HS)
    @test abs(HC[1] - 3//5) < 1.0e-70
    @test abs(HC[2] - 3//5) < 1.0e-70
    # backward: not weighted, not big
    HS = backward_hierarchical_structure(gw, weighted = false)
    @test norm(HS[1] - [5//3,2//3,-7//3]) < 1.0e-12
    @test norm(HS[2] - [0,0,1]) < 1.0e-12
    @test norm(HS[3] - [-1,1,3]) < 1.0e-12
    @test norm(HS[4] - [ 0 1 0 ; -1 0 3 ; 0 0 0 ]) < 1.0e-12
    HC = hierarchical_coefficients(gw, HS, weighted = false)
    @test abs(HC[1]) < 1.0e-12
    @test abs(HC[2] - 2*sqrt(2/3)) < 1.0e-12
    # backward: weighted, not big
    HS = backward_hierarchical_structure(gw)
    @test norm(HS[1] - [5//3,2//3,-7//3]) < 1.0e-12
    @test norm(HS[2] - [0,0,1]) < 1.0e-12
    @test norm(HS[3] - [-1,1,3]) < 1.0e-12
    @test norm(HS[4] - [ 0 1 0 ; -1 0 3 ; 0 0 0 ]) < 1.0e-12
    HC = hierarchical_coefficients(gw, HS)
    @test abs(HC[1]) < 1.0e-12
    @test abs(HC[2] - sqrt(2)) < 1.0e-12
    # backward: not weighted, big
    HS = backward_hierarchical_structure(gw, weighted = false, big = true)
    @test norm(HS[1] - [5//3,2//3,-7//3]) < 1.0e-70
    @test norm(HS[2] - [0,0,1]) < 1.0e-70
    @test norm(HS[3] - [-1,1,3]) < 1.0e-70
    @test norm(HS[4] - [ 0 1 0 ; -1 0 3 ; 0 0 0 ]) < 1.0e-70
    HC = hierarchical_coefficients(gw, HS, weighted = false)
    @test abs(HC[1]) < 1.0e-70
    @test abs(HC[2] - 2*sqrt(big(2)/3)) < 1.0e-70
    # backward: weighted, big
    HS = backward_hierarchical_structure(gw, big = true)
    @test norm(HS[1] - [5//3,2//3,-7//3]) < 1.0e-70
    @test norm(HS[2] - [0,0,1]) < 1.0e-70
    @test norm(HS[3] - [-1,1,3]) < 1.0e-70
    @test norm(HS[4] - [ 0 1 0 ; -1 0 3 ; 0 0 0 ]) < 1.0e-70
    HC = hierarchical_coefficients(gw, HS)
    @test abs(HC[1]) < 1.0e-70
    @test abs(HC[2] - sqrt(big(2))) < 1.0e-70

    # unbalanced triangle
    gw = SimpleWeightedDiGraph(3)
    add_edge!(gw,1,2,0.1)
    add_edge!(gw,2,3,1)
    add_edge!(gw,3,1,1)
    add_edge!(gw,2,1,1)
    # forward: not weighted, not big
    HS = forward_hierarchical_structure(gw, weighted = false)
    @test norm(HS[1] - [7//18,-5//18,-1//9]) < 1.0e-12
    @test norm(HS[2] - [5//12,5//3,5//6]) < 1.0e-12
    @test norm(HS[3] - [-2//3,5//12,1//2]) < 1.0e-12
    @test norm(HS[4] - [ 0 -2//3 0 ; 2//3 0 1//6 ; 1//2 0 0 ]) < 1.0e-12
    HC = hierarchical_coefficients(gw, HS, weighted = false)
    @test abs(HC[1] - 5//6) < 1.0e-12
    @test abs(HC[2] - sqrt(19/2)/6) < 1.0e-12
    # forward: weighted, not big
    HS = forward_hierarchical_structure(gw)
    @test norm(HS[1] - [799//1206,-995//1206,98//603]) < 1.0e-12
    @test norm(HS[2] - [5//804,500//201,5//402]) < 1.0e-12
    @test norm(HS[3] - [-299//201,995//804,1//2]) < 1.0e-12
    @test norm(HS[4] - [ 0 -299//201 0 ; 299//201 0 397//402 ; 1//2 0 0 ]) < 1.0e-12
    HC = hierarchical_coefficients(gw, HS)
    @test abs(HC[1] - 550//6231) < 1.0e-12
    @test abs(HC[2] - sqrt(13557755)/6231) < 1.0e-12
    # forward: not weighted, big
    HS = forward_hierarchical_structure(gw, weighted = false, big = true)
    @test norm(HS[1] - [7//18,-5//18,-1//9]) < 1.0e-70
    @test norm(HS[2] - [5//12,5//3,5//6]) < 1.0e-70
    @test norm(HS[3] - [-2//3,5//12,1//2]) < 1.0e-70
    @test norm(HS[4] - [ 0 -2//3 0 ; 2//3 0 1//6 ; 1//2 0 0 ]) < 1.0e-70
    HC = hierarchical_coefficients(gw, HS, weighted = false)
    @test abs(HC[1] - 5//6) < 1.0e-70
    @test abs(HC[2] - sqrt(big(19)/2)/6) < 1.0e-70
    # forward: weighted, big  ---  Note that since the weights are given as Float64, accuracy is low even when using BigFloat
    HS = forward_hierarchical_structure(gw, big = true)
    @test norm(HS[1] - [799//1206,-995//1206,98//603]) < 1.0e-16
    @test norm(HS[2] - [5//804,500//201,5//402]) < 1.0e-16
    @test norm(HS[3] - [-299//201,995//804,1//2]) < 1.0e-16
    @test norm(HS[4] - [ 0 -299//201 0 ; 299//201 0 397//402 ; 1//2 0 0 ]) < 1.0e-16
    HC = hierarchical_coefficients(gw, HS)
    @test abs(HC[1] - 550//6231) < 1.0e-16
    @test abs(HC[2] - sqrt(13557755)/6231) < 1.0e-16
    # backward: not weighted, not big
    HS = backward_hierarchical_structure(gw, weighted = false)
    @test norm(HS[1] - [-5//18,7//18,-1//9]) < 1.0e-12
    @test norm(HS[2] - [5//3,5//12,5//6]) < 1.0e-12
    @test norm(HS[3] - [5//12,-2//3,1//2]) < 1.0e-12
    @test norm(HS[4] - [ 0 -2//3 0 ; 2//3 0 1//2 ; 1//6 0 0 ]) < 1.0e-12
    HC = hierarchical_coefficients(gw, HS, weighted = false)
    @test abs(HC[1] - 5//6) < 1.0e-12
    @test abs(HC[2] - sqrt(19/2)/6) < 1.0e-12
    # # backward: weighted, not big
    HS = backward_hierarchical_structure(gw)
    @test norm(HS[1] - [-995//1206,799//1206,98//603]) < 1.0e-12
    @test norm(HS[2] - [500//201,5//804,5//402]) < 1.0e-12
    @test norm(HS[3] - [995//804,-299//201,1//2]) < 1.0e-12
    @test norm(HS[4] - [ 0 -299//201 0 ; 299//201 0 1//2 ; 397//402 0 0 ]) < 1.0e-12
    HC = hierarchical_coefficients(gw, HS)
    @test abs(HC[1] - 550//6231) < 1.0e-12
    @test abs(HC[2] - sqrt(13557755)/6231) < 1.0e-12
    # backward: not weighted, big
    HS = backward_hierarchical_structure(gw, weighted = false, big = true)
    @test norm(HS[1] - [-5//18,7//18,-1//9]) < 1.0e-70
    @test norm(HS[2] - [5//3,5//12,5//6]) < 1.0e-70
    @test norm(HS[3] - [5//12,-2//3,1//2]) < 1.0e-70
    @test norm(HS[4] - [ 0 -2//3 0 ; 2//3 0 1//2 ; 1//6 0 0 ]) < 1.0e-70
    HC = hierarchical_coefficients(gw, HS, weighted = false)
    @test abs(HC[1] - 5//6) < 1.0e-70
    @test abs(HC[2] - sqrt(big(19)/2)/6) < 1.0e-70
    # backward: weighted, big
    HS = backward_hierarchical_structure(gw, big = true)
    @test norm(HS[1] - [-995//1206,799//1206,98//603]) < 1.0e-16
    @test norm(HS[2] - [500//201,5//804,5//402]) < 1.0e-16
    @test norm(HS[3] - [995//804,-299//201,1//2]) < 1.0e-16
    @test norm(HS[4] - [ 0 -299//201 0 ; 299//201 0 1//2 ; 397//402 0 0  ]) < 1.0e-16
    HC = hierarchical_coefficients(gw, HS)
    @test abs(HC[1] - 550//6231) < 1.0e-16
    @test abs(HC[2] - sqrt(big(13557755))/6231) < 1.0e-16
end


using CmdStan

ProjDir = dirname(@__FILE__);
cd(ProjDir) do

  isdir("tmp") && rm("tmp", recursive=true);

  ar1 = "
  data {
    int T;
    real y0;
    real phi;
    real sigma;
  }
  parameters{}
  model{}
  generated quantities {
    vector[T] yhat;
    yhat[1] = y0;
    for (t in 2:T)
      yhat[t] = normal_rng(phi*yhat[t-1],sigma);
  }
  "

  #
  T = 100;
  y0 = 1;
  phi = 0.95;
  sigma = 1;

  # make the data dictionary
  dat = Dict("T"=>T,"y0"=>y0,"phi"=>phi,"sigma"=>sigma);

  stanmodel= Stanmodel(name = "ar1", model = ar1,  Sample(algorithm=CmdStan.Fixed_param()));
  rc, sim1 = stan(stanmodel, [dat], ProjDir, CmdStanDir=CMDSTAN_HOME);

  isdir("tmp") && rm("tmp", recursive=true);

end


"""
# Provide implementation of default layers

Provided layers
- DenseLayer
- DenseNoBiasLayer
- VectorFunctionLayer
"""

#using Random, Zygote
#import ..Utils
#import Base.size

# ------------------------------------------------------------------------------
# DenseLayer layer

"""
   DenseLayer

Representation of a layer in the network

# Fields:
* `w`:  Weigths matrix with respect to the input from previous layer or data (n x n pr. layer)
* `wb`: Biases (n)
* `f`:  Activation function
* `df`: Derivative of the activation function
"""
mutable struct DenseLayer <: AbstractLayer
     w::Array{Float64,2}
     wb::Array{Float64,1}
     f::Function
     df::Union{Function,Nothing}
     """
        DenseLayer(f,n,nₗ;w,wb,df)

     Instantiate a new DenseLayer

     # Positional arguments:
     * `nₗ`: Number of nodes of the previous layer
     * `n`:  Number of nodes
     # Keyword arguments:
     * `w`:   Initial weigths with respect to input [default: Xavier initialisation, dims = (n,nₗ)]
     * `wb`:  Initial weigths with respect to bias [default: Xavier initialisation, dims = (n)]
     * `f`:   Activation function [def: `identity`]
     * `df`:  Derivative of the activation function [default: `nothing` (i.e. use AD)]
     * `rng`: Random Number Generator (see [`FIXEDSEED`](@ref)) [deafult: `Random.GLOBAL_RNG`]

     # Notes:
     - Xavier initialization = `rand(Uniform(-sqrt(6)/sqrt(nₗ+n),sqrt(6)/sqrt(nₗ+n))`

     """
     function DenseLayer(nₗ,n;rng = Random.GLOBAL_RNG,w=rand(rng, Uniform(-sqrt(6)/sqrt(nₗ+n),sqrt(6)/sqrt(nₗ+n)),n,nₗ),wb=rand(rng, Uniform(-sqrt(6)/sqrt(nₗ+n),sqrt(6)/sqrt(nₗ+n)),n),f=identity,df=nothing)
         # To be sure w is a matrix and wb a column vector..
         w  = reshape(w,n,nₗ)
         wb = reshape(wb,n)
         return new(w,wb,f,df)
     end
end

function _zComp(layer::DenseLayer,x)
    w  = layer.w
    wb = layer.wb
    z  = zeros(eltype(x),size(w,1))
    @inbounds for n in axes(w,1)
        zn = zero(eltype(x))
        @simd for nl in axes(x,1)
            zn += w[n,nl] * x[nl]
        end
        zn   += wb[n]
        z[n]  = zn
    end
    return z
end

function forward(layer::DenseLayer,x)
  z =  _zComp(layer,x) #@avx layer.w * x + layer.wb #_zComp(layer,x) #layer.w * x + layer.wb # _zComp(layer,x) #   layer.w * x + layer.wb # testd @avx
  return layer.f.(z)
end

function backward(layer::DenseLayer,x,nextGradient)
   z = _zComp(layer,x) #@avx layer.w * x + layer.wb #_zComp(layer,x) # layer.w * x + layer.wb # _zComp(layer,x) # @avx layer.w * x + layer.wb               # tested @avx
   if layer.df != nothing
       dϵ_dz =  layer.df.(z) .* nextGradient # tested @avx
    else
       dϵ_dz = layer.f'.(z) .* nextGradient # using AD
    end
   dϵ_dI =  layer.w' * dϵ_dz # @avx
end

function getParams(layer::DenseLayer)
  return Learnable((layer.w,layer.wb))
end

function getGradient(layer::DenseLayer,x,nextGradient)
   z      =  _zComp(layer,x) #@avx layer.w * x + layer.wb #  _zComp(layer,x) #layer.w * x + layer.wb # @avx
   if layer.df != nothing
       dϵ_dz = layer.df.(z) .* nextGradient
    else
       dϵ_dz = layer.f'.(z) .* nextGradient # using AD
    end
   dϵ_dw  = dϵ_dz * x' # @avx
   dϵ_dwb = dϵ_dz
   return Learnable((dϵ_dw,dϵ_dwb))
end

function setParams!(layer::DenseLayer,w)
   layer.w  = w.data[1]
   layer.wb = w.data[2]
end
function size(layer::DenseLayer)
    return size(layer.w')
end

# ------------------------------------------------------------------------------
# DenseNoBiasLayer layer

"""
   DenseNoBiasLayer

Representation of a layer without bias in the network

# Fields:
* `w`:  Weigths matrix with respect to the input from previous layer or data (n x n pr. layer)
* `f`:  Activation function
* `df`: Derivative of the activation function
"""
mutable struct DenseNoBiasLayer <: AbstractLayer
     w::Array{Float64,2}
     f::Function
     df::Union{Function,Nothing}
     """
        DenseNoBiasLayer(f,nₗ,n;w,df)

     Instantiate a new DenseNoBiasLayer

        DenseNoBiasLayer(f,nₗ,n;w,df)

     Instantiate a new DenseNoBiasLayer

     # Positional arguments:
     * `nₗ`:  Number of nodes of the previous layer
     * `n`:   Number of nodes
     # Keyword arguments:
     * `w`:   Initial weigths with respect to input [default: Xavier initialisation, dims = (nₗ,n)]
     * `f`:   Activation function [def: `identity`]
     * `df`:  Derivative of the activation function [def: `nothing` (i.e. use AD)]
     * `rng`: Random Number Generator (see [`FIXEDSEED`](@ref)) [deafult: `Random.GLOBAL_RNG`]
     # Notes:
     - Xavier initialization = `rand(Uniform(-sqrt(6)/sqrt(nₗ+n),sqrt(6)/sqrt(nₗ,n))`
     """
     function DenseNoBiasLayer(nₗ,n;rng = Random.GLOBAL_RNG,w=rand(rng,Uniform(-sqrt(6)/sqrt(nₗ+n),sqrt(6)/sqrt(nₗ+n)),n,nₗ),f=identity,df=nothing)
         # To be sure w is a matrix and wb a column vector..
         w  = reshape(w,n,nₗ)
         return new(w,f,df)
     end
end

function _zComp(layer::DenseNoBiasLayer,x)
    w  = layer.w
    z  = zeros(eltype(x),size(w,1))
    @inbounds for n in axes(w,1)
        zn = zero(eltype(x))
        @simd for nl in axes(x,1)
            zn += w[n,nl] * x[nl]
        end
        z[n]  = zn
    end
    return z
end

function forward(layer::DenseNoBiasLayer,x)
  z = _zComp(layer,x)
  return layer.f.(z)
end

function backward(layer::DenseNoBiasLayer,x,nextGradient)
   z = _zComp(layer,x)
   if layer.df != nothing
       dϵ_dz = layer.df.(z) .* nextGradient
    else
       dϵ_dz = layer.f'.(z) .* nextGradient # using AD
    end
   dϵ_dI = layer.w' * dϵ_dz
end

function getParams(layer::DenseNoBiasLayer)
  return Learnable((layer.w,))
end

function getGradient(layer::DenseNoBiasLayer,x,nextGradient)
   z      = _zComp(layer,x)
   if layer.df != nothing
       dϵ_dz = layer.df.(z) .* nextGradient
    else
       dϵ_dz = layer.f'.(z) .* nextGradient # using AD
    end
   dϵ_dw  =  dϵ_dz * x'
   return Learnable((dϵ_dw,))
end

function setParams!(layer::DenseNoBiasLayer,w)
   layer.w = w.data[1]
end

function size(layer::DenseNoBiasLayer)
    return size(layer.w')
end

# ------------------------------------------------------------------------------
# VectorFunction layer

"""
   VectorFunctionLayer

Representation of a (weightless) VectorFunction layer in the network. Vector
function layer expects a vector activation function, i.e. a function taking the
whole output of the previous layer an input rather than working on a single node
as "normal" activation functions would do.
Useful for example with the SoftMax function in classification or with the `pool1D`
function to implement a "pool" layer in 1 dimensions.
As it is weightless, it doesn't apply any transformation to the output coming
from the previous layer. It means that the number of nodes must be set to the same
as in the previous layer (and if you are using this for classification, to the
number of classes, i.e. the _previous_ layer must be set equal to the number of
classes in the predictions).

# Fields:
* `nₗ`: Number of nodes in input (i.e. length of previous layer)
* `n`:  Number of nodes in output (automatically inferred in the constructor)
* `f`:  Activation function (vector)
* `df`: Derivative of the (vector) activation function

# Notes:
* The output `size` of this layer is given by the size of the output function,
that not necessarily is the same as the previous layers.
"""
mutable struct VectorFunctionLayer <: AbstractLayer
     nₗ::Int64
     n::Int64
     f::Function
     df::Union{Function,Nothing}
     """
        VectorFunctionLayer(f,n;df)

     Instantiate a new VectorFunctionLayer

     # Positional arguments:
     * `nₗ`: Number of nodes (must be same as in the previous layer)
     # Keyword arguments:
     * `f`:  Activation function [def: `softMax`]
     * `df`: Derivative of the activation function [default: `nothing` (i.e. use AD)]
     * `dummyDataToTestOutputSize`: Dummy data to test the output size [def: `ones(nₗ)`]

     # Notes:
     - If the derivative is provided, it should return the gradient as a (n,n) matrix (i.e. the Jacobian)
     - To avoid recomputing the activation function just to determine its output size,
       we compute the output size once here in the layer constructor by calling the
       activation function with `dummyDataToTestOutputSize`. Feel free to change
       it if it doesn't match with the activation function you are setting
     """
     function VectorFunctionLayer(nₗ;f=softMax,df=nothing,dummyDataToTestOutputSize=ones(nₗ))
         n = length(f(dummyDataToTestOutputSize))
         return new(nₗ,n,f,df)
     end
end

function forward(layer::VectorFunctionLayer,x)
  return layer.f(x)
end

function backward(layer::VectorFunctionLayer,x,nextGradient)
   if layer.df != nothing
       dϵ_dI = layer.df(x)' * nextGradient
    else  # using AD
      j = autoJacobian(layer.f,x; nY=layer.n)
      dϵ_dI = j' * nextGradient
    end
   return dϵ_dI
end

function getParams(layer::VectorFunctionLayer)
  return Learnable(())
end

function getGradient(layer::VectorFunctionLayer,x,nextGradient)
   return Learnable(()) # parameterless layer
end

function setParams!(layer::VectorFunctionLayer,w)
   return nothing
end
function size(layer::VectorFunctionLayer)
    # Output size for the VectorFunctionLayer is given by its activation function
    # We test its length with dummy values
    return (layer.nₗ,layer.n)
end


function max_dt(u, t, mesh::TreeMesh{3},
                constant_speed::Val{false}, equations, dg::DG, cache)
  # to avoid a division by zero if the speed vanishes everywhere,
  # e.g. for steady-state linear advection
  max_scaled_speed = nextfloat(zero(t))

  for element in eachelement(dg, cache)
    max_λ1 = max_λ2 = max_λ3 = zero(max_scaled_speed)
    for k in eachnode(dg), j in eachnode(dg), i in eachnode(dg)
      u_node = get_node_vars(u, equations, dg, i, j, k, element)
      λ1, λ2, λ3 = max_abs_speeds(u_node, equations)
      max_λ1 = max(max_λ1, λ1)
      max_λ2 = max(max_λ2, λ2)
      max_λ3 = max(max_λ3, λ3)
    end
    inv_jacobian = cache.elements.inverse_jacobian[element]
    max_scaled_speed = max(max_scaled_speed, inv_jacobian * (max_λ1 + max_λ2 + max_λ3))
  end

  return 2 / (nnodes(dg) * max_scaled_speed)
end

function max_dt(u, t, mesh::TreeMesh{3},
                constant_speed::Val{true}, equations, dg::DG, cache)
  # to avoid a division by zero if the speed vanishes everywhere,
  # e.g. for steady-state linear advection
  max_scaled_speed = nextfloat(zero(t))

  for element in eachelement(dg, cache)
    max_λ1, max_λ2, max_λ3 = max_abs_speeds(equations)
    inv_jacobian = cache.elements.inverse_jacobian[element]
    max_scaled_speed = max(max_scaled_speed, inv_jacobian * (max_λ1 + max_λ2 + max_λ3))
  end

  return 2 / (nnodes(dg) * max_scaled_speed)
end

function max_dt(u, t, mesh::CurvedMesh{3},
                constant_speed::Val{false}, equations, dg::DG, cache)
  # to avoid a division by zero if the speed vanishes everywhere,
  # e.g. for steady-state linear advection
  max_scaled_speed = nextfloat(zero(t))

  @unpack contravariant_vectors = cache.elements

  for element in eachelement(dg, cache)
    max_λ1 = max_λ2 = max_λ3 = zero(max_scaled_speed)
    for k in eachnode(dg), j in eachnode(dg), i in eachnode(dg)
      u_node = get_node_vars(u, equations, dg, i, j, k, element)
      λ1, λ2, λ3 = max_abs_speeds(u_node, equations)

      Ja11, Ja12, Ja13 = get_contravariant_vector(1, contravariant_vectors, i, j, k, element)
      λ1_transformed   = abs(Ja11 * λ1 + Ja12 * λ2 + Ja13 * λ3)
      Ja21, Ja22, Ja23 = get_contravariant_vector(2, contravariant_vectors, i, j, k, element)
      λ2_transformed   = abs(Ja21 * λ1 + Ja22 * λ2 + Ja23 * λ3)
      Ja31, Ja32, Ja33 = get_contravariant_vector(3, contravariant_vectors, i, j, k, element)
      λ3_transformed   = abs(Ja31 * λ1 + Ja32 * λ2 + Ja33 * λ3)

      inv_jacobian = abs(cache.elements.inverse_jacobian[i, j, k, element])

      max_λ1 = max(max_λ1, inv_jacobian * λ1_transformed)
      max_λ2 = max(max_λ2, inv_jacobian * λ2_transformed)
      max_λ3 = max(max_λ3, inv_jacobian * λ3_transformed)
    end

    max_scaled_speed = max(max_scaled_speed, max_λ1 + max_λ2 + max_λ3)
  end

  return 2 / (nnodes(dg) * max_scaled_speed)
end

function max_dt(u, t, mesh::CurvedMesh{3},
                constant_speed::Val{true}, equations, dg::DG, cache)
  # to avoid a division by zero if the speed vanishes everywhere,
  # e.g. for steady-state linear advection
  max_scaled_speed = nextfloat(zero(t))

  @unpack contravariant_vectors = cache.elements

  max_λ1, max_λ2, max_λ3 = max_abs_speeds(equations)

  for element in eachelement(dg, cache)
    for k in eachnode(dg), j in eachnode(dg), i in eachnode(dg)
      Ja11, Ja12, Ja13 = get_contravariant_vector(1, contravariant_vectors, i, j, k, element)
      λ1_transformed   = abs(Ja11 * max_λ1 + Ja12 * max_λ2 + Ja13 * max_λ3)
      Ja21, Ja22, Ja23 = get_contravariant_vector(2, contravariant_vectors, i, j, k, element)
      λ2_transformed   = abs(Ja21 * max_λ1 + Ja22 * max_λ2 + Ja23 * max_λ3)
      Ja31, Ja32, Ja33 = get_contravariant_vector(3, contravariant_vectors, i, j, k, element)
      λ3_transformed   = abs(Ja31 * max_λ1 + Ja32 * max_λ2 + Ja33 * max_λ3)

      inv_jacobian = abs(cache.elements.inverse_jacobian[i, j, k, element])

      max_scaled_speed = max(max_scaled_speed,
                             inv_jacobian * (λ1_transformed + λ2_transformed + λ3_transformed))
    end
  end

  return 2 / (nnodes(dg) * max_scaled_speed)
end


include("../../../../std/util/src/util_module.jl")

module dtest
using ..util
using ..distributions
using PyPlot
PyPlot.pygui(true)

function test_uniform()

    distr = distributions.UniformDistribution(2.0, 3.0)

    n = 1000
    rnb = distributions.random(distr, n)
    # show(rnb)

    (cdf_x, cdf_y) = distributions.emp_cdf(rnb)

    PyPlot.figure()
    PyPlot.grid()
    PyPlot.plot(cdf_x, cdf_y, marker = :o)

end

function test_normal()

    distr = distributions.NormalDistribution(2.0, 3.0)

    n = 10000
    rnb = distributions.random(distr, n)
    # show(rnb)

    (cdf_x, cdf_y) = distributions.emp_cdf(rnb)

    PyPlot.figure()
    PyPlot.grid()
    PyPlot.plot(cdf_x, cdf_y, marker = :o)

end

function test_cauchy()

    distr = distributions.CauchyDistribution(2.0, 3.0)

    n = 10000
    rnb = distributions.random(distr, n)
    # show(rnb)

    (cdf_x, cdf_y) = distributions.emp_cdf(rnb)

    PyPlot.figure()
    PyPlot.grid()
    PyPlot.plot(cdf_x, cdf_y, marker = :o)

end

function test_pdf_uniform()

    gauss_distr = distributions.NormalDistribution(0.0, 1.0)
    cauchy_distr = distributions.CauchyDistribution(0.0, 1.0)

    x = collect( -3.0:0.1:3.0 )
    yc = distributions.pdf(cauchy_distr, x)
    yg = distributions.pdf(gauss_distr, x)

    PyPlot.figure()
    PyPlot.grid()
    PyPlot.plot(x, yc)
    PyPlot.plot(x, yg)

end

function test_histo()

    gauss_distr = distributions.NormalDistribution(0.0, 2.0)
    n_rand = 10000
    rnb = distributions.random(gauss_distr, n_rand)

    edges = [0.0, 1.0, 1.1, 1.2, 1.5, 2.0]
    # (bin_edges, freq, bin_idx) = distributions.histogram(rnb, edges)
    (bin_edges, freq) = distributions.histogram(rnb, 50)

    println( " The sum of freqs: $(sum(freq)) " )

    x_mids = 0.5 * (bin_edges[1:end-1] + bin_edges[2:end])

    x_min = gauss_distr.mu - 5 * gauss_distr.sigma
    x_max = gauss_distr.mu + 5 * gauss_distr.sigma

    pdf_x = collect( range(x_min, x_max, 100) )
    pdf_y = distributions.pdf(gauss_distr, pdf_x)

    # PyPlot.figure()
    # PyPlot.plot(pdf_x, pdf_y)

    # PyPlot.figure()
    # PyPlot.plot(x_mids, freq, marker = :o)

    widths = diff(bin_edges)
    PyPlot.figure()
    PyPlot.bar(bin_edges[1:end-1], freq, align=:edge, width = widths) #

end

function test_rejection_sampling()

    gauss_distr = distributions.NormalDistribution(0.0, 1.0)
    gauss_distr2 = distributions.NormalDistribution(5.0, 2.0)

    pdf_fcn(x) = 0.5*(
        1 * distributions.pdf(gauss_distr2, x) +
        distributions.pdf(distributions.NormalDistribution(3.0, 1.0), x) )

    x_min = -4.0
    x_max = 8.0

    n_rand = 10000
    rnb = distributions.rejection_sampling_1d(pdf_fcn, x_min, x_max, n_rand)
    # rnb = distributions.random(gauss_distr2, n_rand)
    (cdf_x, cdf_y) = distributions.emp_cdf(rnb)

    # rnb_g = distributions.random(gauss_distr, n_rand)
    # (cdf_x_g, cdf_y_g) = distributions.emp_cdf(rnb_g)

    n_discr = 100
    x_vec = collect( range(x_min, x_max, n_discr) )
    y_vec = pdf_fcn.(x_vec)

    # n_bin = max( 10, convert(Int, ceil(sqrt( n_rand ))) )

    sigma = distributions.std(rnb)
    h = distributions.bin_size_scott(sigma, n_rand)

    n_bin = convert(Int, ceil( (x_max - x_min) / h ))

    (bin_edges, freq) = distributions.histogram(rnb, n_bin)
    x_mids = 0.5 .* (bin_edges[1:end-1] .+ bin_edges[2:end])
    pdf_freq = 1 / (bin_edges[2] - bin_edges[1]) * freq[2:end-1] ./ n_rand

    PyPlot.figure()
    PyPlot.grid()
    PyPlot.plot(x_vec, y_vec)

    # PyPlot.plot(cdf_x_g, cdf_y_g)
    PyPlot.plot(cdf_x, cdf_y)
    PyPlot.plot( x_mids, pdf_freq )

end

function test_multi_gauss()

    mu = [1.0, 3.0]
    sigma = [0.5 0.8; 0.8 3.0]
    multi_normal = distributions.MultiNormalDistribution(mu, sigma)

    x_vec = collect( range(-5., 5, 150) )
    y_vec = collect( range(-5., 11, 150) )

    (X, Y) = util.mesh_grid(x_vec, y_vec)

    Lx = length(x_vec)
    Ly = length(y_vec)

    pdf_s = zeros( Lx, Ly )
    for ii = 1:Lx
        for jj = 1:Ly
            pdf_s[ii, jj] = distributions.pdf( multi_normal, (x_vec[ii], y_vec[jj]) )
        end
    end

    PyPlot.figure()
    PyPlot.grid()
    # PyPlot.surf(X, Y, pdf_s, alpha = 0.7)
    PyPlot.contour(X, Y, pdf_s, levels = [1e-4, 0.001, 0.01, 0.025, 0.05,  0.1, 0.125, 0.16, 0.17])

end

function test_multi_gauss_random()

    mu = [1.0, 3.0]
    sigma = [0.5 0.8; 0.8 3.0]
    multi_normal = distributions.MultiNormalDistribution(mu, sigma)

    n_rand = 1000
    rnb = distributions.random(multi_normal, n_rand)

    n_bin_x = 29
    n_bin_y = 29

    (x_bin_edges, y_bin_edges, freq) = distributions.histogram_2d(rnb, n_bin_x, n_bin_y)

    # show(x_bin_edges)
    # show(y_bin_edges)
    # show(freq)
    (X_bin, Y_bin) = util.mesh_grid(x_bin_edges[1:end-1], y_bin_edges[1:end-1])

    x_vec = collect( range(-5., 5, 150) )
    y_vec = collect( range(-5., 11, 150) )

    (X, Y) = util.mesh_grid(x_vec, y_vec)

    Lx = length(x_vec)
    Ly = length(y_vec)

    pdf_s = zeros( Lx, Ly )
    for ii = 1:Lx
        for jj = 1:Ly
            pdf_s[ii, jj] = distributions.pdf( multi_normal, (x_vec[ii], y_vec[jj]) )
        end
    end

    PyPlot.figure()
    PyPlot.grid()
    PyPlot.contour(X, Y, pdf_s, levels = [1e-4, 0.001, 0.01, 0.025, 0.05,  0.1, 0.125, 0.16, 0.17])
    PyPlot.plot( rnb[1, :], rnb[2, :], linestyle = :none, marker = :., markersize = 10 )

    width = x_bin_edges[2] - x_bin_edges[1]
    depth = y_bin_edges[2] - y_bin_edges[1]
    height = freq[:]

    xxx = X_bin[:]
    yyy = Y_bin[:]
    zzz = similar(xxx) .* 0

    PyPlot.bar3d( xxx, yyy, zzz, width, depth, height)

end

function test_rejection_sampling_2d()

    multi_gauss = distributions.MultiNormalDistribution( [1.0, 3.0], [0.5 0.0; 0.0 1.0] )
    pdf_fcn(x) = distributions.pdf(multi_gauss, x)

    factor = 5.0
    x_min = multi_gauss.mu[1] - factor * multi_gauss.sigma[1,1]
    x_max = multi_gauss.mu[1] + factor * multi_gauss.sigma[1,1]

    y_min = multi_gauss.mu[2] - factor * multi_gauss.sigma[2,2]
    y_max = multi_gauss.mu[2] + factor * multi_gauss.sigma[2,2]

    x_vec = collect( range(x_min, x_max, 50) )
    y_vec = collect( range(y_min, y_max, 50) )
    (X, Y) = util.mesh_grid(x_vec, y_vec)

    Lx = length(x_vec)
    Ly = length(y_vec)

    pdf_s = zeros( Lx, Ly )
    for ii = 1:Lx
        for jj = 1:Ly
            pdf_s[ii, jj] = distributions.pdf( multi_gauss, (x_vec[ii], y_vec[jj]) )
        end
    end

    n_rand = 100

    (x_rnb, y_rnb) = distributions.rejection_sampling_2d(pdf_fcn, x_min, x_max, y_min, y_max, n_rand)

    PyPlot.figure()
    PyPlot.contour(X, Y, pdf_s, levels = [1e-4, 0.001, 0.01, 0.025, 0.05,  0.1, 0.125, 0.16, 0.17])
    PyPlot.plot( x_rnb, y_rnb , linestyle = :none, marker = :., markersize = 10)

end

end # dtest


## Medical Diagnosis Benchmark
## ==========================

using Sigma
using DataFrames
disease_priors = [0.025851035,0.058152077,0.020436158,0.040025596,
                  0.032941836,0.078530717,0.005548213,0.024689777,
                  0.029336553,0.020641227,0.074125078,0.028544798,
                  0.054423213,0.077353640,0.078099579,0.065297382,
                  0.039513058,0.030245903,0.128596589,0.095504584]
diseases = RandArray(map(p->flip(p),disease_priors))

datapath = joinpath(benchdir, "diseasediagnosis","problem-2-edges.csv")
edges = readtable(datapath)
nfindings = size(edges,2)
nfindings = 10
ndiseases = size(edges,1)

findings = RandVar{Bool}[]
for i = 1:nfindings
  ands = RandVar{Bool}[]
  for j = 1:ndiseases
    if edges[j,i] != 0.0
      push!(ands, diseases[j] & flip(edges[j,i]))
    end
  end
  push!(findings, (|)(ands...))
end

findings_randvar = RandArray(findings)

## Benchmarking

samples = cond_sample_mh(diseases,findings_randvar == [randbool() for i = 1:nfindings],10)
@show samples


# ==========================================================================
# Main interface for linear Support Vector Machines

"""
`svm(X, y⃗; nargs...) → SVM`

Description
============

Fits a Support Vector Machine (SVM) to the given training data `X` and `y⃗`.
Depending on the specified parameters in `nargs`, the SVM can be trained to
perform various forms of regression or classification tasks.
There are a number of parameters that can be configured to the individual needs,
such as the type of lossfunction, the type of regularization, or the algorithm
used for training the SVM in the primal/dual.

Usage
======

    svm(X, y; solver = DualCD(), kernel = ScalarProductKernel(), loss = HingeLoss(), regtype = L2Penalty, bias = 1, C = 1, dual = false, maxiter = 1000, ftol = 1.0e-6, xtol = 1.0e-6, grtol = 1.0e-9, callback = nothing, show_trace = false, nargs...)

    svm(X, Y; nargs...) do t, alpha, v, g
      # this is the callback function
    end

Arguments
==========

- **`X`** : The input matrix of the training set, with the observations as columns.
It can be either a dense or sparse matrix of floatingpoint numbers.
Note that it is generally advised to center and scale your data before training a SVM.

- **`y⃗`** : The target vector of the training set. It has to have the same length as `X` has columns.
In the case of classification, it can be either a vector with elements yᵢ ∈ {-1.0, 1.0}
for two class prediction, or a vector with elements yᵢ ∈ N for multiclass prediction.

- **`solver`** : The concrete algorithm to train the SVM. Depending on the algorithm
this will either solve the linear SVM in the primal or dual formulation.
In contrast to other frameworks, the algorithm does not determine if the solution
of the primal or dual problem is returned (see Details).

- **`kernel`** : The kernel that should be used. The ScalarProductKernel denotes the special
case of a linear SVM and allows for additional functionality and solvers that are specialized
for the linear case.

- **`loss`** : The utilized loss function. The typical loss functions for SVMs are
`L1HingeLoss()` (L1-SVM), `L2HingeLoss()` (L2-SVM), or `SmoothedL1HingeLoss(h)` for classification,
and `EpsilonInsLoss(e)` or `L2EpsilonInsLoss(e)` for support vector regression.
Note that in general each solvers is only able to deal with a subset of those
loss functions and some combinations might not be supported.
For example the implementation of `DualCD` does only support
`HingeLoss` and `L2HingeLoss`.

- **`regtype`** : The type of regularization that should be used. In general this can either
be `L1Penalty` or `L2Penalty`. Note that not all solver support `L1Penalty`.

- **`bias`** : The scaling factor of the bias. If set to 0, no intercept will be fitted.

- **`C`** : The penalty parameter. It plays a role similar to λ⁻¹ in Empirical Risk Minimization.
Conceptually, it denotes the trade off between model complexity and the training error.
Larger C will increase the penalization of training errors and thus lead to behavior
more similar to a hard margin classifier.

- **`dual`** : Boolean that specifies if the solution to the dual problem should be returned
instead of the solution to the primal (default = false).

- **`iterations`** : Specifies the maximum number of iterations before the solver exits with the
current (and probably suboptimal) solution.

- **`ftol`** : Specifies the tolerance for the change of the objective function value.

- **`xtol`** : Specifies the tolerance for the change of the solution.

- **`grtol`** : Specifies the tolerance for the change of the gradient norm.

- **`callback`** : The optional callback function with signature `f(i, w, v, G)`. If a callback is specified, then it will be called
every iteration with the following four parameters in order:
1) `i` … the current iteration number,
2) `w` … the current coefficients (i.e. α if dual=true, or θ if dual=false),
3) `v` … the current objective value
4) `G` … the current gradient.
Note: Some solver support early stopping if the callback function returns the symbol `:Exit`.

- **`show_trace`** : Instead of (or additional to) the callback function, the user can enable the output of the training information to STDOUT.

- **`nargs...`** : Additional named arguments that are passed unchecked to the specified solver.
This functionality can be used to pass around special arguments that the library does not natively implement.

Details
========

In contrast to other frameworks, the `dual` parameter does not influence whether the dual or
the primal problem is solved (this is instead specified implicitly by the `solver` parameter),
but the result the user is interested in.
It is generally recommended to leave `dual = false`, unless the user is either explicitly
interested in the support vectors, or the dataset is sparse and high-dimensional.
Note: Not all solver are able to natively provide both solutions. If the algorithm is unable
to provide the desired solution type then the solution will be converted accordingly, if feasible.

Returns
========

Either a subtype of `DualSVM` or `PrimalSVM`

Author(s)
==========

- Christof Stocker (email: c.stocker@hci-kdd.org)

Examples
=========

### Iris Dataset (adapted example from SVM.jl by John Myles White)

    using KSVM
    using RDatasets
    iris = dataset("datasets", "iris")
    X = convert(Array, iris[1:100, 1:2])'  # The observations have to be in the columns
    y = iris[1:100, :Species]
    train = bitrand(size(X,2))             # Split training and testset
    model = svm(X[:, train], y[train])     # Fit the linear SVM
    acc = accuracy(model, X[:, ~train], y[~train])

### Random Data (adapted example from Regression.jl by Dahua Lin)

    using KSVM
    d = 2             # Number of input dimensions
    n = 100           # Number of pbservations
    w = randn(d+1)    # generate the weight vector
    X = randn(d, n)   # generate input features
    y = sign(X'w[1:d] + w[d+1] + 0.2 * randn(n))  # generate (noisy) response
    model = svm(X, y) # Fit the linear SVM
    y_hat = classify(model, X)

References
===========

- Hsieh, Cho-Jui, et al. "A dual coordinate descent method for large-scale
linear SVM." Proceedings of the 25th international conference on Machine learning.
ACM, 2008. DOI=10.1145/1390156.1390208, http://doi.acm.org/10.1145/1390156.1390208
"""
@inline function svm{TReg<:Penalty}(
    X::AbstractMatrix, y⃗::AbstractVector;
    kernel::Kernel = ScalarProductKernel(),
    solver::SvmSolver = DualCD(),
    loss::Loss = L2HingeLoss(),
    regtype::Type{TReg} = L2Penalty,
    C::Real = 1.,
    nargs...)
  spec = CSVM(kernel = kernel,
              loss = loss,
              regtype = regtype,
              C = C)
  fit(spec, X, y⃗; solver = solver, nargs...)
end

# ==========================================================================
# Special method overload to support the "do" notation for the callback
# res = svm(X, Y) do t, alpha, v, g
#   println("Iteration $t: $v")
# end

@inline function svm(callback::Function,
             X::AbstractMatrix, y⃗::AbstractVector;
             nargs...)
  svm(X, y⃗; callback = callback, nargs...)
end


using Units
using Base.Test

l = Quantity(Micro, Meter, 380)
@test prefix(l) == Micro
@test base(l) == Meter
pshow(l)
println()
lshow(l)
println()
fshow(l)
println()
lcm = convert(Unit(Centi,Meter), l)
@test lcm.value == l.value/10^4

temp = parse_quantity("32F")
tempC = convert(Unit(SINone,Celsius), temp)
println(temp, " = ", tempC)

q1 = parse_quantity("3.2 μm")
q2 = parse_quantity("5.7 m")
@test q2/q1 == (5.7/3.2)*1e6
@test 5*q1 == Quantity(Micro, Meter, 5*3.2)
@test q2/7 == Quantity(SINone, Meter, 5.7/7)

q1 = parse_quantity("7 s")
q2 = parse_quantity("4 d")
@test q1/q2 == 7/4/86400

q1 = parse_quantity("1 g")
q2 = parse_quantity("1 Da")
@test q2/q1 == 1.66053892173e-24

q1 = parse_quantity("3.2 μm")
q2 = parse_quantity("4 d")
@test_throws q1/q2


using CGMethod1D
using Ferrite

grid = generate_grid(Line, (10,))

dim = 1
order = 3
ip = Lagrange{dim, RefLine, order}()
qr = QuadratureRule{dim, RefLine}(order)
# cellvalues = CellScalarValues(qr, ip)

# dh = DofHandler(grid)
# push!(dh, :p, 1, ip)
# # push!(dh, :p, 1)
# close!(dh)

# K = create_sparsity_pattern(dh);

# using UnicodePlots
# fill!(K.nzval, 1.0)
# spy(K; height = 15)

# ch = ConstraintHandler(dh);

# ∂Ω = union(getfaceset.((grid, ), ["left", "right", "top", "bottom"])...);

# dbc = Dirichlet(:u, ∂Ω, (x, t) -> 0)
# add!(ch, dbc);

# close!(ch)
# update!(ch, 0.0);

# function doassemble(cellvalues::CellScalarValues{dim}, K::SparseMatrixCSC, dh::DofHandler) where {dim}

#     n_basefuncs = getnbasefunctions(cellvalues)
#     Ke = zeros(n_basefuncs, n_basefuncs)
#     fe = zeros(n_basefuncs)

#     f = zeros(ndofs(dh))
#     assembler = start_assemble(K, f)

#     @inbounds for cell in CellIterator(dh)

#         fill!(Ke, 0)
#         fill!(fe, 0)

#         reinit!(cellvalues, cell)

#         for q_point in 1:getnquadpoints(cellvalues)
#             dΩ = getdetJdV(cellvalues, q_point)

#             for i in 1:n_basefuncs
#                 v  = shape_value(cellvalues, q_point, i)
#                 ∇v = shape_gradient(cellvalues, q_point, i)
#                 fe[i] += v * dΩ
#                 for j in 1:n_basefuncs
#                     ∇u = shape_gradient(cellvalues, q_point, j)
#                     Ke[i, j] += (∇v ⋅ ∇u) * dΩ
#                 end
#             end
#         end

#         assemble!(assembler, celldofs(cell), fe, Ke)
#     end
#     return K, f
# end

# K, f = doassemble(cellvalues, K, dh);

# apply!(K, f, ch)
# u = K \ f;


module GravitationalAcceleration

# This implements the gravitational acceleration

using Parameters, LaTeXStrings, Unitful
using ..Units
using GeoParams: AbstractMaterialParam
import Base.show

abstract type AbstractGravity <: AbstractMaterialParam end

export  ComputeGravity,        # calculation routines
        ConstantGravity        # constant

# Constant Gravity -------------------------------------------------------
"""
    GravityConstant(g=9.81m/s^2)

Set a constant value for the gravitational acceleration:
```math
    g  = 9.81 m s^{-2}
```
"""
@with_kw_noshow mutable struct ConstantGravity <: AbstractGravity
    equation::LaTeXString   =   L"g = 9.81 m s^{-2}"
    g::GeoUnit              =   9.81m/s^2               # gravitational acceleration
end

# Calculation routine
function ComputeGravity(s::ConstantGravity)
    @unpack g   = s

    return g*1.0   # multiply with 1.0, to return Float64
end

# Print info
function show(io::IO, d::ConstantGravity)
    print(io, "Gravitational acceleration: g=$(d.g.val)")
end
#-------------------------------------------------------------------------

# Help info for the calculation routines
"""
ComputeGravity(s:<AbstractGravity)

Returns the gravitational acceleration

"""
ComputeGravity

end


"""
Finite difference calculation & machine learning of electric potential & electric field from charge
"""

using LinearAlgebra
using Plots
using Random
using Flux
Random.seed!(1)
include("../src/operators.jl")
include("../src/plotutils.jl")
# using Main.EquivariantOperators # end users should omit Main.

# make grid
dims = 2
dx = 0.1
cell = dx * Matrix(I, dims, dims)
rmax = 1.0
grid = Grid(cell, rmax)

# make operators
rmin = 1e-9
ϕ = Op(r -> 1 / (4π * r), rmin, rmax, cell)
E = Op(r -> 1 / (4π * r^2), rmin, rmax, cell; l = 1)
▽ = Op(:▽, cell)

# put dipole charges
ρf = zeros(size(grid))
put!(ρf, grid, [0.5, 0.0],1)
put!(ρf, grid, [-0.5, 0.0],-1)

# calculate fields
Ef = E(ρf)
ϕf = ϕ(ρf)

# test
rvec = [0, 0]
@show get(ϕf, grid, rvec), [0.0]
@show get(Ef, grid, rvec), get(-▽(ϕf), grid, rvec), [-2 / (4π * 0.5^2), 0]

p = []
push!(p, heatmap(ρf',title = "dipole charge"), )
push!(p, heatmap(ϕf', title = "dipole potential"),)
plot(p..., layout = length(p))
vfplot(Ef,grid; title = "dipole electric field")

##
# make neural operators
ϕ_ = Op(Radfunc(), rmin, rmax, cell)
E_ = Op(Radfunc(), rmin, rmax, cell; l = 1)

ps = Flux.params(ϕ_, E_)
function loss()
    remake!(E_)
    remake!(ϕ_)
    global E_f = E_(ρf)
    global ϕ_f = ϕ_(ρf)
    @show l = (nae(E_f, Ef) + nae(ϕ_f, ϕf)) / 2
    l
end

data = [()]
loss()
opt = ADAM(0.1)
Flux.@epochs 10 Flux.train!(loss, ps, data, opt)

## plot
p = []
push!(p, heatmap(ϕ_.kernel, title = "learned potential kernel"))
r = 0:0.01:1
push!(
    p,
    plot(r, ϕ_.radfunc.(r), title = "learned potential kernel radial function"),
)
plot(p..., layout = length(p))

vfplot(E_.kernel,grid; title = "learned E field kernel ")


"""
    NumericalContinuationShootingProblems

This module implements a basic single shooting method for use with ordinary
differential equations (ODEs). The solvers are provided by OrdinaryDiffEq.jl.

This assumes that the ODE takes a `Vector` or `SVector` input and produces a
corresponding output. For small ODEs (i.e., fewer than around 20 dimensions)
`SVector` inputs are often significantly faster.

See [`add_shootingproblem!`](@ref) for details.
"""
module NumericalContinuationShootingProblems

using DocStringExtensions
using NumericalContinuation: add_var!, add_func!, add_pars!
using OrdinaryDiffEq: solve, remake, ODEProblem, Tsit5

export add_shootingproblem!, periodic

struct ShootingProblem{T, O, S, B}
    abstol::T
    reltol::T
    odeprob::O
    solver::S
    bc::B
end

function (shoot::ShootingProblem)(res, u, p, tspan)
    sol = solve(remake(shoot.odeprob, u0=u, p=p, tspan=(tspan[1], tspan[2])),
        shoot.solver, save_everystep=false, save_start=false, abstol=shoot.abstol, reltol=shoot.reltol)
    shoot.bc(res, u, sol[end], p, tspan)
end

"""
    $SIGNATURES

Periodic boundary conditions for numerical shooting.
"""
periodic(res, u0, u1, p, tspan) = (res .= u1 .- u0)

"""
$SIGNATURES

Construct an ODE shooting problem (boundary value problem) of the form

```math
   u' = f(u, p, t),
```
for `t∈[t₀, t₁]` where `u` is the state and `p` are the parameters, and add it
to the problem structure. The function `f` is assumed to be an appropriate
input to `ODEProblem` of the OrdinaryDiffEq.jl package.

# Parameters

* `prob` : the underlying continuation problem.
* `name::String` : the name of the algebraic zero problem.
* `f` : the function to use for the zero problem. It takes either three
  arguments (`u`, `p`, and `t`) or four arguments for an in-place version
  (`dudt`, `u`, `p`, and `t`).
* `u0` : the initial state value (typically either a Vector or an SVector).
* `p0` : the initial parameter values.
* `tspan` : the integration time span as a vector or tuple.
* `bc` : (optional) the boundary conditions to use, specified as a function
  with five arguments (`res`, `u0`, `u1`, `p`, `tspan`) where `res` is the
  residual. If not specified, periodic boundary conditions will be used.
* `pnames` : (keyword, optional) the names of the parameters (default:
  auto-generated names).
* `solver` : (keyword, optional) the IVP solver to use (default: Tsit5).
* `reltol` : (keyword, optional) the absolute tolerance for the IVP solver to
  use (default: 1e-6).
* `abstol` : (keyword, optional) the absolute tolerance for the IVP solver to
  use (default: 1e-6).

# Example

The Hopf bifurcation normal form
```
hopf = (out, u, p, t) -> (out[1] = p[1]*u[1] - u[2] + p[2]*u[1]*(u[1]^2 + u[2]^2);
                          out[2] = u[1] + p[1]*u[2] + p[2]*u[2]*(u[1]^2 + u[2]^2))
prob = ProblemStructure()
add_shootingproblem!(prob, "hopf", hopf, [1.0, 0.0], [1.0, -1.0], [0, 2π], periodic)
```
"""
function add_shootingproblem!(prob, name::String, f, u0, p0, tspan, bc=periodic; pnames=nothing, solver=Tsit5(), reltol=1e-6, abstol=1e-8)
    _tspan = length(tspan) == 1 ? [zero(tspan[1]), tspan[1]] : [tspan[1], tspan[2]]
    odeprob = ODEProblem(f, u0, (_tspan[1], _tspan[2]), p0)
    shooting = ShootingProblem(reltol, abstol, odeprob, solver, bc)
    # Check for parameter names
    _pnames = pnames !== nothing ? [string(pname) for pname in pnames] : ["$(name).p$i" for i in 1:length(p0)]
    if length(_pnames) != length(p0)
        throw(ArgumentError("Length of parameter vector does not match number of parameter names"))
    end
    # Create the necessary continuation variables and add the function
    uidx = add_var!(prob, "$(name).u", length(u0), u0=u0)
    pidx = add_var!(prob, "$(name).p", length(p0), u0=p0)
    tidx = add_var!(prob, "$(name).tspan", 2, u0=_tspan)
    fidx = add_func!(prob, name, length(u0), shooting, [uidx, pidx, tidx])
    add_pars!(prob, _pnames, pidx, active=false)
    add_pars!(prob, ("$(name).t0", "$(name).t1"), tidx, active=false)
    return prob
end

end # module


"""
    rewire_edges(g, p)

Rewire each edge of `g` with probability `p`.

When `g` is a cyclic k-lattice, this essentially produces Watts-Strogatz networks.
"""
function rewire_edges(g::T, p::Real) where T <: AbstractGraph
    h = typeof(g)(nv(g))
    for e in edges(g)
        (u,v) = src(e), dst(e)
        # this is somewhat inefficient, but at least it's correct
        rewirings = append!(
            [(u,w) for w in setdiff(vertices(g), outneighbors(h, u)) if w != v],
            [(w,v) for w in setdiff(vertices(g), inneighbors(h, v)) if w != u]
            )
        if !isempty(rewirings) && rand() < p
            (u, v) = rand(rewirings)
        end
        add_edge(h, u, v)
    end
    return h
end

rewire_edges(p::Real) = g -> rewire_edges(g,p)

"""
    drop_edges!(g, p)

Remove each edge in `g` with probability `p`.
"""
function drop_edges!(g, p::Real)
    for e in collect(edges(g))
        if rand() < p
            rem_edge(g,p)
        end
    end
    return g
end

"""
    drop_edges(g, p)

Remove each edge in `g` with probability `p`.

This function creates a copy of `g`. See [drop_edges!](@ref) for the in-place version.
"""
drop_edges(g, p::Real) = drop_edges!(copy(g), p)

drop_edges!(p::Real) = g -> drop_edges!(g,p)
drop_edges(p::Real) = g -> drop_edges!(g,p)


@testset "$TEST $G" begin

if !@isdefined(readcentrality)
    function readcentrality(f::AbstractString)
        f = open(f,"r")
        c = Vector{Float64}()
        while !eof(f)
            line = chomp(readline(f))
            push!(c, parse(Float64,line))
        end
        return c
    end
end

g = readgraph(joinpath(testdir,"testdata","graph-50-500.gml"),:gml)

c = readcentrality(joinpath(testdir,"testdata","graph-50-500-bc.txt"))
z = betweenness_centrality(g)

@test map(Float32, z) == map(Float32, c)

y = betweenness_centrality(g, endpoints=true, normalize=false)
@test round.(y[1:3], digits=4) ==
    round.([122.10760591498584, 159.0072453120582, 176.39547945994505], digits=4)
x = betweenness_centrality(g; approx=3)
@test length(x) == 50

# self loops
s2 = DG(3)
add_edge!(s2,1,2); add_edge!(s2,2,3); add_edge!(s2,3,3)
s1 = graph(s2)

@test betweenness_centrality(s1) == [0, 1, 0]
@test betweenness_centrality(s2) == [0, 0.5, 0]

g = G(2)
add_edge!(g,1,2)
z = betweenness_centrality(g; normalize=true)
@test z[1] == z[2] == 0.0

g3 = PathGraph(5, G)
z = betweenness_centrality(g3; normalize=false)
@test z[1] == z[5] == 0.0

end # testset


#!/usr/bin/env julia
using Printf
using Distributed
using TimerOutputs
using CUDAdrv
using CUDAnative
using CuArrays

function ζ_gen_range(mstart,mend,e)
    s = 0.0
    for i=mstart:-1:mend
      s += (1.0/i)^e
    end
    s
end

function ζ_gen_thr_range(mstart,mend,e)
    s = 0
    cls  = 8  # cache line size in units of Float64 storage size
              # so we can avoid false sharing
    Nthr = Threads.nthreads()
    psum = zeros(Float64,Nthr*cls)
    ms   = zeros(Float64,Nthr*cls)
    me   = zeros(Float64,Nthr*cls)

    # preparing each thread, start, end, and partial sums
    for i=1:Nthr
        ms[i*cls]   = (N/Nthr)*i
        me[i*cls]   = 1 + (N/Nthr)*(i-1)
        psum[i*cls] = 0.0
    end

    # actual function we will execute per thread
    # psums will be updated per thread id.  This is why
    # the psums are spaced 1 cache line apart, versus adjacent indices
    function inner!(psum,ms,me,x)
        ndx     = Threads.threadid()
        for i=ms[ndx*cls]:-1:me[ndx*cls]
            psum[ndx*cls] += (1.0/i)^e
        end
	nothing
    end

    # Threaded inner loop, each thread has no dependence upon others
    Threads.@threads for t=1:Nthr
        inner!(psum,ms,me,2)
    end

    # now sum up the partial sums and return s
    s = sum(psum)
    s
end

function ζ_gen_gpu_range(N)

    Nthr = 512
    st   = 8

    psum_d = CuArrays.fill(0.0::Float64, Nthr*st)

    # actual function we will execute per thread
    # psums will be updated per thread id.  This is why
    # the psums are spaced 1 cache line apart, versus adjacent indices
    function innergpu!(N,ps)
        ndx     = threadIdx().x
        stride  = blockDim().x

        for i=ndx:stride:N
		#for i=N:-stride:ndx
	    	r = 1.0/i
            ps[ndx] += r*r
        end
        return nothing
    end

    # Threaded inner loop, each thread has no dependence upon others
    @cuda threads=Nthr  innergpu!(N,psum_d)

    # now sum up the partial sums and return s
    s = sum(psum_d)
    s
end

function ζ_2_pow2_range(mstart,mend)
    s = 0.0
    for i=mstart:-1:mend
      s += (1.0/i)^2
    end
    s
end

function ζ_2_sqr_2_range(mstart,mend)
    s = 0.0
    for i=mstart:-1:mend
      s += (1.0/i)*(1.0/i)
    end
    s
end

function ζ_gen_simd_range(mstart,mend,e)
    s = 0.0
    @simd for i=mstart:-1:mend
      s += (1.0/i)^e
    end
    s
end

function ζ_2_pow2_simd_range(mstart,mend)
    s = 0.0
    @simd for i=mstart:-1:mend
      s += (1.0/i)^2
    end
    s
end

function ζ_2_sqr_simd_range(mstart,mend)
    s = 0.0
    @simd for i=mstart:-1:mend
      s += (1.0/i)*(1.0/i)
    end
    s
end

to = TimerOutput()

N = 16000

# first bit of code to handle the compilation stage, otherwise you
# see this in the time output
s = ζ_gen_range(N,1,2) + ζ_2_pow2_range(N,1) + ζ_2_sqr_2_range(N,1) +
    ζ_gen_simd_range(N,1,2) + ζ_2_pow2_simd_range(N,1) +
    ζ_2_sqr_simd_range(N,1) + ζ_gen_thr_range(N,1,2) + ζ_gen_gpu_range(N)

s /= 8.0
@printf("s=%f\n",s)
s=0.0

@timeit to "ζ_gen_range"            s+=ζ_gen_range(N,1,2)
@timeit to "ζ_2_pow2_range"         s+=ζ_2_pow2_range(N,1)
@timeit to "ζ_2_sqr_2_range"        s+=ζ_2_sqr_2_range(N,1)
@timeit to "ζ_gen_simd_range"       s+=ζ_gen_simd_range(N,1,2)
@timeit to "ζ_2_pow2_simd_range"    s+=ζ_2_pow2_simd_range(N,1)
@timeit to "ζ_2_sqr_simd_range"     s+=ζ_2_sqr_simd_range(N,1)
@timeit to "ζ_gen_thr_range"        s+=ζ_gen_thr_range(N,1,2)
@timeit to "ζ_gen_gpu_range"        s+=ζ_gen_gpu_range(N)

s /= 8.0
@printf("s=%f\n",s)

show(to)
@printf("\n");


@testset "GPCP coorShift" begin
	# convert from -180to180 to "0to360"
	lon = collect(-179.5:1:179.5);
	datain = repeat(-179.5:1:179.5,1,180);
	lonout,dataout = coorShift_lon(lon,datain,to="0to360")
	@test size(datain) == size(dataout)
	@test lonout == collect(0.5:1:359.5)
	@test dataout[1,1] == 0.5
	@test dataout[end,1] == -0.5
	@test dataout[181,2] == -179.5

	# convert from 0to360 to "-180to180" + 3D
	lon = collect(0.125:0.25:359.875);
	datain = Array{Float64}(undef,720,length(lon),2);
	datain[:,:,1] = transpose(repeat(0.125:0.25:359.875,1,720));
	datain[:,:,2] = transpose(repeat(0.125:0.25:359.875,1,720));
	lonout,dataout = coorShift_lon(lon,datain,to="-180to180")
	@test size(datain) == size(dataout)
	@test lonout == collect(-179.875:0.25:179.875)
	@test dataout[1,1,1] == 180.125
	@test dataout[end,1,2] == 180.125
	@test dataout[2,181,1] == dataout[2,181,2]
	@test dataout[2,end,1] == 179.875
end

@testset "GPCP read (head, lon/lat, time, data)" begin
	header_read = readgpcpd_head(joinpath(dirname(@__DIR__),"test","input","gpcpd_data"));
	@test length(header_read) == 1440
	lon,lat = readgpcpd_lonlat(joinpath(dirname(@__DIR__),"test","input","gpcpd_data"))
	@test lon == collect(0.5:1:359.5)
	@test lat == collect(89.5:-1:-89.5)
	timeout = readgpcpd_time(joinpath(dirname(@__DIR__),"test","input","gpcpd_data"));
	@test timeout == collect(DateTime(2011,6,1):Dates.Day(1):DateTime(2011,6,30))
	dataout = readgpcpd(joinpath(dirname(@__DIR__),"test","input","gpcpd_data"));
	@test size(dataout) == (360,180,30)
end


using SigmaRidgeRegression
using Test
using LinearAlgebra
using Random
import StatsBase

@testset "Woodbury and Cholesky" begin
    Random.seed!(1)
    σ = 5.0
    grp = GroupedFeatures([30;30;30;30;30;30])
    n = 400
    p = grp.p
    X = randn(n, p)
    αs = sqrt.(range(4.0,12.0,length=ngroups(grp)))#r#ange(2.0, 2.75, 3.5; length=3)#1.0:5.0
    β = random_betas(grp, αs)
    Y = X*β .+ σ .* randn(n)

    tmp_chol = BasicGroupRidgeWorkspace(X=X, Y=Y, groups=grp)

    tmp_woodb = BasicGroupRidgeWorkspace(X=X, Y=Y, groups=grp,
                               XtXpΛ_chol = WoodburyRidgePredictor(X))

    @test SigmaRidgeRegression.λωλας_λ(tmp_chol) ≈ SigmaRidgeRegression.λωλας_λ(tmp_woodb)
    beta_chol = tmp_chol.XtXpΛ_chol \ tmp_woodb.XtY
    beta_wdb =  tmp_woodb.XtXpΛ_chol \ tmp_woodb.XtY
    @test beta_chol ≈ beta_wdb

    tmp_chol.XtXpΛ_chol \ tmp_chol.X' ≈ tmp_woodb.XtXpΛ_chol \ tmp_woodb.X'

    ldiv_chol = ldiv!(tmp_chol.XtXpΛ_div_Xt, tmp_chol.XtXpΛ_chol, tmp_chol.X')
    ldiv_wdb = ldiv!(tmp_woodb.XtXpΛ_div_Xt, tmp_woodb.XtXpΛ_chol, tmp_woodb.X')

    @test ldiv_chol ≈ ldiv_wdb

    loo_chol = StatsBase.fit!(tmp_chol, SigmaRidgeRegression.λωλας_λ(tmp_chol))
    loo_wdb = StatsBase.fit!(tmp_woodb, SigmaRidgeRegression.λωλας_λ(tmp_chol))

    @test tmp_woodb.β_curr ≈ tmp_chol.β_curr
end

@testset "MMI" begin
    include("test_mmi.jl")
end
#include("test_grouplasso.jl")
include("test_covariance.jl")
include("test_simulations.jl")

@testset "Theoretical risks" begin
    include("test_theoretical_risks.jl")
end
# TO ADD
# Compare: - GroupLasso AGAINST R
#          - Check group lasso lambda_max
#          - test the TunedRidgeRegressor type
#          - test σ_max  (probably won't increase coverage but...)


rhoxb(x::Real, b::Real) = 2*b*b + 2.5 - sqrt(4*b^4 + 6*b*b+2.25 - x*x - x/b)

function multiply!(des::RealVector, x::RealVector, y::Real, n::Int=length(x))
    for i in 1:n
        @inbounds des[i] = x[i]*y
    end
end
multiply!(x::RealVector, y::Real) = multiply!(x, x, y)
function divide!(des::RealVector, x::RealVector, y::Real, n::Int=length(x))
    for i in 1:n
        @inbounds des[i] = x[i]/y
    end
end
divide!(x::RealVector, y::Real) = divide!(x, x, y)
function minus!(des::RealVector, y::Float64, x::RealVector, n::Int64=length(x))
   for i in 1:n
       @inbounds des[i] = y - x[i]
   end
   nothing
end
function add!(x::Vector{Float64}, y::Float64, n::Int64=length(x))
   for i in 1:n
       @inbounds x[i] = x[i] + y
   end
   nothing
end
function abs2!(des::RealVector, x::RealVector, n::Int64=length(x))
   for i in 1:n
       @inbounds des[i] = abs2(x[i])
   end
   nothing
end
function betakernel(x::Real, xdata::RealVector, h::Real, w::Vector, n::Int)
    a = x / h - 1
    b = (1 - x) / h - 1
    if (x < 0) || (x > 1)
        fill!(w, 0.0)
        return nothing
    elseif x < 2h
        a =  rhoxb(x, h) - 1
    elseif x>1-2*h
        b = rhoxb(1-x, h) - 1
    end

    minus!(w, 1.0, xdata, n)
    w .= log.(w)
    wtmp = log.(xdata)
    multiply!(w, b)
    multiply!(wtmp, a)
    w .= w .+ wtmp

    # for ind in 1:n
    #     @inbounds w[ind] = a * log(xdata[ind]) + b * log(1 - xdata[ind])
    # end

    add!(w, -lbeta(a+1, b+1))
    w .= exp.(w)
    nothing
end
function betakernel(x::Real, logxdata::RealVector, log1_xdata::RealVector, h::Real, w::Vector, n::Int)
    a = x / h - 1
    b = (1 - x) / h - 1
    if (x < 0) || (x > 1)
        fill!(w, 0.0)
        return nothing
    elseif x < 2h
        a =  rhoxb(x, h) - 1
    elseif x>1-2*h
        b = rhoxb(1-x, h) - 1
    end

    for ind in 1:n
        @inbounds w[ind] = a * logxdata[ind] + b * log1_xdata[ind]
    end
    add!(w, -lbeta(a+1, b+1))
    w .= exp.(w)
    nothing
end
#f̂(x) = 1/n ∑ᵢ K(xᵢ;x /b+1, b )
#xdata should be positive, or domain error will be raised.
function gammakernel(x::Real, xdata::RealVector, h::Real, w::Vector, n::Int)
    rhob = x/h
    if x <= 0
        fill!(w, 0.0)
        return nothing
    elseif x < 2*h
        rhob = 0.25 * rhob * rhob + 1.0
    end

    w .= log.(xdata)
    multiply!(w, rhob-1.0)
    tmp = -rhob*log(h)-lgamma(rhob)
    add!(w, tmp)
    h1 = 1/h
    for ind in 1:n
        @inbounds w[ind] -= xdata[ind] * h1
    end
    w .= exp.(w)
    nothing
end
function gammakernel(x::Real, xdata::RealVector, logxdata::RealVector, h::Real, w::Vector, n::Int)
    rhob = x/h
    if x <= 0
        fill!(w, 0.0)
        return nothing
    elseif x < 2*h
        rhob = 0.25 * rhob * rhob + 1.0
    end

    # Yeppp.log!(w, xdata)
    multiply!(w, logxdata, rhob-1.0)
    tmp = -rhob*log(h)-lgamma(rhob)
    add!(w, tmp)
    h1 = 1/h
    for ind in 1:n
        @inbounds w[ind] -= xdata[ind] * h1
    end
    w .= exp.(w)
    nothing
end

function gaussiankernel(x::Real, xdata::RealVector, h::Real, w::Vector, n::Int)

    # for ind in 1:n
    #     @inbounds w[ind] = normlogpdf(xdata[ind], h, x)
    # end
    h1= 1.0/h
    # minus!(w, x, xdata, n)
    # multiply!(w, w, h1, n)
    # abs2!(w, w, n)
    # multiply!(w, w, -0.5, n)
    tmp = log(h) + log2π/2
    for ind in 1:n
        @inbounds w[ind]=-0.5*abs2((x - xdata[ind])*h1) - tmp
    end
    # add!(w, tmp, n)
    w .= exp.(w)

    nothing
end
function ekernel(x::Real, xdata::RealVector, h::Real, w::Vector, n::Int)
    ind = 1
    ind_end = 1+n
    @inbounds while ind < ind_end
        u = (x - xdata[ind]) / h
        w[ind] = ifelse(abs(u)>=1.0, 0.0, 1-u*u)
        ind += 1
    end
    multiply!(w, 0.75 / h)
    nothing
end

# #univariate normal kernel

# gaussiankernel{T<:FloatingPoint}(xi::T, x::T, h::T)=exp(-0.5*abs2((x - xi)/h)) / h * invsqrt2π
# gaussiankernel{T<:Real}(xi::T, x::T, h::T) = gkernel(float(xi), float(x),float(h))
# gaussiankernel(xi::Real, x::Real, h::Real)=gkernel(promote(xi, x, h)...)
# # type Gkernel <: Functor{3} end
# # NumericExtensions.evaluate(::Gkernel, xi, x, h) = gkernel(xi, x, h)

# function ekernel{T<:FloatingPoint}(xi::T, xeval::T, h::T)
#     u=(xeval - xi) / h
#     abs(u)>=1.0 ? 0.0 : 0.75*(1-u*u)
# end
# ekernel{T<:Real}(xi::T,xeval::T, h::T) = ekernel(float(xi),float(xeval), float(h))
# ekernel(xi::Real, xeval::Real, h::Real)=ekernel(promote(xi, xeval, h)...)
# # type Ekernel <: Functor{3} end
# # NumericExtensions.evaluate(::Ekernel, xi, xeval, h) = ekernel(xi, xeval, h)

# # # data need to be transformed on [0, ∞]
# # # K_{ρₕ(x), h}(xᵢ) = (t/h)^(ρₕ(x)-1) exp(-t/h)/gamma(ρ)/h
# function gammakernel{T<:FloatingPoint}(xi::T, x::T, b::T)
#     rhob = x/b
#     if x < 0
#         return(0)
#     elseif x < 2*b
#         rhob = 0.25 * rhob * rhob + 1
#     else
#         nothing
#     end
#     xi_b = xi/b
#     xi_b^(rhob-1) * exp(-xi_b) / b / gamma(rhob)
# end
# gammakernel{T<:Real}(xi::T,xeval::T, h::T) = gammakernel(float(xi),float(xeval), float(h))
# gammakernel(xi::Real, xeval::Real, h::Real)=gammakernel(promote(xi, xeval, h)...)
# # type Gammakernel <: Functor{3} end
# # NumericExtensions.evaluate(::Gammakernel, xi, xeval, h) = gammakernel(xi, xeval, h)

# # #data need be transformed on [0, 1]
# function betakernel{T<:FloatingPoint}(xi::T, x::T, h::T)
#     a = x / h
#     b = (1 - x) / h
#     if (x < 0) | (x > 1)
#         return(0)
#     elseif x < 2h
#         a =  rhoxb(x, h)
#     elseif x>1-2*h
#         b = rhoxb(1-x, h)
#     else
#         nothing
#     end
#     xi ^ (a-1) * (1 - xi) ^ (b-1) #/ beta(a, b)
# end
# betakernel{T<:Real}(xi::T,xeval::T, h::T) = betakernel(float(xi),float(xeval), float(h))
# betakernel(xi::Real, xeval::Real, h::Real)=betakernel(promote(xi, xeval, h)...)
# # type Betakernel <: Functor{3} end
# NumericExtensions.evaluate(::Betakernel, xi, xeval, h) = betakernel(xi, xeval, h)

# # #MultiVariate Normal Kernel
# # function GaussianKernel(xeval::Vector{Float64}, xi::Vector{Float64}, h::Vector{Float64})
# #   if any(h .<= 0.0)
# #     return Inf
# #   end
# #   p=length(xi)

# #   exp(-sumsq((xeval .- xi) ./ h)) / (2*pi)^(p/2) / prod(h)

# # end
# # # GaussianKernel2(xdiff::Float64, h::Float64) = GaussianKernel(xdiff, 0.0, sqrt(2)*h)
# # # GaussianKernel2(xdiff::Vector{Float64}, h::Vector{Float64}) = GaussianKernel(xdiff, zeros(length(xdiff)), sqrt(2).* h)

# # # Gaussian=KernelType(GaussianKernel, GaussianKernel2)

# # function EKernel(xeval::Float64, xi::Float64,h::Float64)
# #     if h <= 0.0
# #       return Inf
# #     end
# #     u=(xeval - xi) / h
# #     if abs(u) > 1.0
# #         return 0.0
# #     end
# #     return (1.0-u*u)*3.0/4.0
# # end

# # function EKernel(xeval::Vector{Float64}, xi::Vector{Float64}, h::Vector{Float64})

# #     if any(h .<= 0.0)
# #         return Inf
# #     end
# #     p=length(xi)
# #     tmp=1.0
# #     for i in 1:p
# #         tmp *= EKernel(xeval[i],xi[i],h[i])
# #     end
# #     tmp
# # end

# # function EKernel2(xdiff::Float64, h::Float64)
# #     ax=abs(xdiff / h)
# #     if ax > 2.0
# #         return 0.0
# #     end

# #     (2 - ax)^3*(ax^2+6*ax+4) * 3 / 160 / h

# # end
# # function EKernel2(xdiff::Vector{Float64}, h::Vector{Float64})
# #     ax=abs(xdiff ./ h)
# #     if any(ax .> 2.0)
# #         return 0.0
# #     end
# #     p=length(ax)
# #     tmp=1.0
# #     for i in 1:p
# #         tmp *= EKernel2(ax[i], h[i])
# #     end
# #     tmp
# # end
# # Epanechnikov=KernelType(EKernel, EKernel2)


#=
  Random walk model

  from Foundations of Probabilistic Programming,
  page 3

  The distribution of u and v are about the same as this
ulia> show_var_dist_pct(rand(DiscreteUniform(-100,100),10000),10)

Distributions of variable (num:10)
-76.00000 =>      67  0.006700
-38.00000 =>      65  0.006500
17.00000 =>      64  0.006400
59.00000 =>      64  0.006400
81.00000 =>      63  0.006300
24.00000 =>      62  0.006200
74.00000 =>      61  0.006100
-6.00000 =>      61  0.006100
-100.00000 =>      60  0.006000
-95.00000 =>      60  0.006000

=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

data = []
@model function random_walk()
    u ~ DiscreteUniform(-10,10)
    v ~ DiscreteUniform(-10,10)
    function step(tu=0,tv=0)
        x ~ flip()
        y ~ flip()
        tu = tu + (x-y)
        tv = tv + (x+y-1)
        return tu,tv
    end
    # println("\nSTART")
    datad = [[0,0]]
    u,v = step(0,0)
    push!(datad,[u,v])
    c = 0
    while u != 0 || v != 0
        u,v = step(u,v)
        push!(datad,[u,v])
        c += 1
        # Note: We must restrict the walk, otherwise
        #       it'll blow up...
        if u == 0 && v == 0
            println("BACK TO 0,0!")
        end
        if c < 10
            push!(data,datad)
        else
            break
        end

    end
end

model = random_walk()

num_chns = 4
chns = sample(model, MH(), MCMCThreads(), 1000, num_chns)
# chns = sample(model, MH(), 10_000)

display(chns)
show_var_dist_pct(chns, :x)
show_var_dist_pct(chns, :y)
show_var_dist_pct(chns, :u,20)
show_var_dist_pct(chns, :v,20)

# for d in data
#    println(d)
# end


#using Revise, Test
using BifurcationKit, LinearAlgebra, Setfield
const BK = BifurcationKit

function test_newton(x0)
	Ty = eltype(x0)

	F(x, p) = x.^3 .- 1
	Jac(x, p) = diagm(0 => 3 .* x.^2)

	opts = NewtonPar( tol = Ty(1e-8))
	sol, hist, flag, _ = newton(F, Jac, x0, nothing, opts, normN = x->norm(x,Inf))
end
####################################################################################################
# we test the regular newton algorithm
# simple case
test_newton(ones(10) .+ rand(10) * 0.1)

# test types for newton
sol, = test_newton(Float16.(ones(10) .+ rand(10) * 0.1))
@test eltype(sol) == Float16
####################################################################################################
function test_newton_palc(x0, p0)
	Ty = eltype(x0)
	N = length(x0)

	θ = Ty(0.2)
	dotθ = BK.DotTheta(dot)

	F(x, p) = x.^3 .- p
	Jac(x, p) = diagm(0 => 3 .* x.^2)

	z0 = BorderedArray(x0, p0)
	τ0 = BorderedArray(rand(Ty, N), convert(typeof(p0), 0.2))
	zpred = BorderedArray(x0, convert(typeof(p0), 0.3))
	optn = NewtonPar{Ty, DefaultLS, DefaultEig}()
	optc = ContinuationPar{Ty, DefaultLS, DefaultEig}(newtonOptions = optn)
	sol, hist, flag, _ = newtonPALC(F, Jac, p0, (@lens _), z0, τ0, zpred, Ty(0.02), θ, optc, dotθ)
end

sol, = test_newton_palc(ones(10) .+ rand(10) * 0.1, 1.)

#test type
sol, = test_newton_palc(Float32.(ones(10) .+ rand(10) * 0.1), Float32(1.))
@test typeof(sol) == BorderedArray{Vector{Float32}, Float32}
####################################################################################################
# test of  deflated problems
_T = Float32
F4def = (x,p) -> (x-1) * (x-2)
deflationOp = DeflationOperator(_T(2), dot, _T(1), [[_T(1)]])
@test eltype(deflationOp) == _T
@test deflationOp(rand(_T,1)) isa _T


using DelimitedFiles

function read_vector(doc, n)
    vector = Array{Int, 1}()
    data = readdlm(doc, skipstart=n)
    for i in data[1,:]
        append!(vector, i)
    end
    return(n+1, vector)
end

function read_matrix(P, Q, doc, n)
    matrix = zeros(Int, P, Q)
    data = readdlm(doc, skipstart=n, use_mmap=true)
    for i = 1:P
        for j = 1:Q
            matrix[i,j] = data[i,j]
        end
    end
    return(n+P, matrix)
end

function read_data(doc)
    open(doc, "r") do f
        n = 7
        C = parse(Int, readline(f)) #number of containers
        P = parse(Int, readline(f)) #number of positions
        CP = parse(Int, readline(f)) #number of compatible combinations
        Q = parse(Int, readline(f)) #number of quay cranes
        J = parse(Int, readline(f)) #number of bays
        tt = parse(Int, readline(f)) #quay crane travel time per bay
        d = parse(Int, readline(f)) #clearance, in number of bays

        n, ci = read_vector(doc, n) #container class vector
        n, pj = read_vector(doc, n) #position class vector
        n, bj = read_vector(doc, n) #position class vector

        #n, qj = read_matrix(P,Q,doc,n) #quay cranes position matrix
        #n, sj = read_matrix(P,S,doc,n) #transfer vehicle position matrix
        n, prejj = read_matrix(P, P, doc, n) #position precedence matrix
        n, cpij = read_matrix(CP, 3, doc ,n) #combination c, p, time(cp)
        #return(C,P,CP,S,Q,J,tt,d,mt,beta,EFT,ci,pj,qj,sj,prejj,cpij)
        return(C,P,CP,Q,J,tt,d,ci,pj,bj,prejj,cpij)
    end
end
#doc = "../data/5C_3Type_2QC.txt"
#doc = "../data/QCSP_example.txt"
#doc = "../../data/10C_4Type_2QC.txt"
#doc = "./data/Benchmark/60C_10Type_LessDense_2QC.txt"
#doc = "./data/Benchmark/240C_20Type_LessDense_2QC.txt"
doc = "../data/Benchmark/240C_20Type_LessDense_2QC.txt"
#C,P,CP,S,Q,J,tt,d,mt,beta,EFT,ci,pj,qj,sj,prejj,cpij = read_data(doc)
C,P,CP,Q,J,tt,d,ci,pj,bj,prejj,cpij = read_data(doc)


# This file was generated, do not modify it.

using MLJ
using PrettyPrinting
import DataFrames: DataFrame, select!, Not, describe
import Statistics
using Dates
using PyPlot
ioff() # hide
using UrlDownload

MLJ.color_off() # hide

df = DataFrame(urldownload("https://raw.githubusercontent.com/tlienart/DataScienceTutorialsData.jl/master/data/kc_housing.csv", true))
describe(df)

select!(df, Not([:id, :date]))
schema(df)

coerce!(df, :zipcode => Multiclass)
df.isrenovated  = @. !iszero(df.yr_renovated)
df.has_basement = @. !iszero(df.sqft_basement)
schema(df)

coerce!(df, :isrenovated => OrderedFactor, :has_basement => OrderedFactor);

unique(df.waterfront)

df.waterfront = (df.waterfront .!= "FALSE")
coerce!(df, :waterfront => OrderedFactor);

coerce!(df, autotype(df, :few_to_finite))
schema(df)

df.price = df.price ./ 1000;

select!(df, Not([:yr_renovated, :sqft_basement, :zipcode]));

plt.figure(figsize=(8,6))
plt.hist(df.price, color = "blue", edgecolor = "white", bins=50,
         density=true, alpha=0.5)
plt.xlabel("Price", fontsize=14)
plt.ylabel("Frequency", fontsize=14)
plt.savefig(joinpath(@OUTPUT, "hist_price.svg")) # hide

plt.figure(figsize=(8,6))
plt.hist(df.price[df.isrenovated .== true], color="blue", density=true,
        edgecolor="white", bins=50, label="renovated", alpha=0.5)
plt.hist(df.price[df.isrenovated .== false], color="red", density=true,
        edgecolor="white", bins=50, label="unrenovated", alpha=0.5)
plt.xlabel("Price", fontsize=14)
plt.ylabel("Frequency", fontsize=14)
plt.legend(fontsize=12)
plt.savefig(joinpath(@OUTPUT, "hist_price2.svg")) # hide

@load DecisionTreeRegressor

y, X = unpack(df, ==(:price), col -> true)
train, test = partition(eachindex(y), 0.7, shuffle=true, rng=5)

tree = machine(DecisionTreeRegressor(), X, y)

fit!(tree, rows=train);

rms(y[test], predict(tree, rows=test))

@load RandomForestRegressor pkg=ScikitLearn

coerce!(X, Finite => Count);

rf_mdl = RandomForestRegressor()
rf = machine(rf_mdl, X, y)
fit!(rf, rows=train)

rms(y[test], predict(rf, rows=test))

cv3 = CV(; nfolds=3)
res = evaluate(rf_mdl, X, y, resampling=CV(shuffle=true),
               measure=rms, verbosity=0)

@load XGBoostRegressor

coerce!(X, Count => Continuous)

xgb  = XGBoostRegressor()
xgbm = machine(xgb, X, y)
fit!(xgbm, rows=train)

rms(y[test], predict(xgbm, rows=test))

r1 = range(xgb, :max_depth, lower=3, upper=10)
r2 = range(xgb, :num_round, lower=1, upper=25);

tm = TunedModel(model=xgb, tuning=Grid(resolution=7),
                resampling=CV(rng=11), ranges=[r1,r2],
                measure=rms)
mtm = machine(tm, X, y)
fit!(mtm, rows=train)

rms(y[test], predict(mtm, rows=test))

PyPlot.close_figs() # hide


# This file defines the hidden dynamics system problem
# with (1) background noise and sparse large
# deviations, (2) "broken sensor" noise, and
# (3) an added "bump". L2, robust (huber), and
# robust (trimming) fits are
# run for several background noise levels
# and the quality of the resulting eigenvalues is
# returned

using Optim

################################################################
################################################################
################################################################
# Define some helper routines

function hidden_dynamics_system_data(;nt=2^9,
                                     nx=300,k1=1.0,
                                     omega1=1.0,
                                     gamma1=1.0,
                                     k2=0.4,
                                     omega2=3.7,
                                     gamma2=-0.2)
################################################
# Generate data for hidden dynamics system
#
# z(t,x) = exp(gamma1*t)*sin(k1*x-omega1*t) +
#      exp(gamma2*t)*sin(k2*x-omega2*t) +
#
# as gamma1 is positive and gamma2 is negative,
# the mode for gamma2 becomes hidden, especially
# with significant noise.
#
################################################

    evals = [gamma1 + im*omega1; gamma1 - im*omega1;
             gamma2 + im*omega2; gamma2 - im*omega2]

    x = collect(range(0,stop=15.0,length=nx))
    t = complex(collect(range(0,stop=2*pi,length=nt)))

    z = complex(exp.(gamma1*t) .* sin.(k1*transpose(x) .- omega1*t) +
         exp.(gamma2*t) .* sin.(k2*transpose(x) .- omega2*t))

    return z, evals, t, x
end

function make_corrupt_data_sparse_spikes(xclean,sigma,mu,p)
    #------------------------------------------
    # Generate Corrupted Data - Spikes
    #------------------------------------------
    #
    # sigma - background noise
    # mu - spike size
    # p - spike frequency

    m, n = size(xclean)

    noise = sigma*randn(m,n) .+ mu*randn(m,n).*( rand(m,n) .< p )
    xdat = xclean .+ noise

    return xdat
end

function make_corrupt_data_bump(xclean,x,t,sigma,A,w)
    #------------------------------------------
    # Generate Corrupted Data - bump
    #------------------------------------------
    #
    # sigma - background noise
    # A - bump amplitude
    # w - bump width

    m, n = size(xclean)

    i1 = rand(1:m)
    j1 = rand(1:n)
    x1 = x[j1]
    t1 = t[i1]

    dx = x[2]-x[1]
    dt = t[2]-t[1]

    noise = (sigma*randn(m,n) .+
             A*exp.( real(-(t .- t1).^2/(dt*w)^2 .- (transpose(x) .- x1).^2/(dx*w)^2 )))
    xdat = xclean .+ noise

    return xdat
end

function make_corrupt_data_broken_sensors(xclean,sigma,mu,p)
    #---------------------------------------------------------
    # Generate Corrupted Data - "broken sensors"
    #---------------------------------------------------------
    #
    # sigma - background noise
    # mu - spike size
    # p - spike frequency

    m, n = size(xclean)

    ncol = Integer(round(n*p))
    icols = sample(1:n,ncol,replace=false)

    noise = sigma*randn(m,n)
    noise[:,icols] .+= mu*randn(m,length(icols))
    xdat = xclean .+ noise

    return xdat
end

function dofits(xdat,t,eigs,kappa,optionsl2,optionshuber,optionstrim,
                nkeep)
    # this wraps up the calls to the DMD solver
    # Here we use the robust but generally slower
    # proximal gradient routine for both L2 and Huber
    # fits. On return, we provide the error for these
    # fits and the exact DMD error.

    k = length(eigs)

    #-------------------------------------------------------------
    # Initial Guess
    #-------------------------------------------------------------
    m,n = size(xdat);

    a0, B0 = dmdexactestimate(m, n, k, xdat, t);

    ainit, Binit = dmdexactestimate(m,n,k,xdat,t,dmdtype="trap")

    #----------------------------------------------------
    # l2 Trial
    #----------------------------------------------------
    #println("running l2 experiment...");
    # loss functions
    lossf = (z) -> l2_func(z);
    lossg = (z) -> l2_grad!(z);

    l2params = DMDParams(k, xdat, t, lossf, lossg,
                         inner_directl2=true);
    copyto!(l2params.a,ainit)

    # apply solver
    obj_his1, err_his1 = solveDMD_withPG(l2params, optionsl2);
    res = solveDMD_withOptimBFGS(l2params)

    #----------------------------------------------------
    # huber Trial
    #----------------------------------------------------
    #println("running huber experiment...");
    # loss functions
    lossf = (z) -> huber_func(z,kappa);
    lossg = (z) -> huber_grad!(z,kappa);

    huberparams = DMDParams(k, xdat, t, lossf, lossg);
    copyto!(huberparams.a,ainit)

    # apply solver
    obj_his2, err_his2 = solveDMD_withPG(huberparams, optionshuber);
    res = solveDMD_withOptimBFGS(huberparams)
    #----------------------------------------------------
    # l2 trimming trial
    #----------------------------------------------------
    #println("running l2 trimming experiment...");
    # loss functions
    lossf = (z) -> l2_func(z);
    lossg = (z) -> l2_grad!(z);

    trimparams = DMDParams(k, xdat, t, lossf, lossg,
                           inner_directl2=true,nkeep=nkeep);

    copyto!(trimparams.a,ainit)

    # apply solver
    obj_his3, err_his3 = solveDMD_withPG(trimparams, optionstrim);
    res = solveDMD_withOptimBFGS(trimparams)
    #--------------------------------------------------------------------
    # Evaluate Errors and Get Best Eigenvalue Permutation
    #--------------------------------------------------------------------

    if any(isnan,l2params.a)
        err1 = Inf
        p1 = 1:k
    else
        err1, p1 = besterrperm_wi(l2params.a, eigs);
    end
    if any(isnan,huberparams.a)
        err2 = Inf
        p2 = 1:k
    else
        err2, p2 = besterrperm_wi(huberparams.a, eigs);
    end
    if any(isnan,trimparams.a)
        err3 = Inf
        p3 = 1:k
    else
        err3, p3 = besterrperm_wi(trimparams.a, eigs);
    end
    err4, p4 = besterrperm_wi(a0, eigs);

    return (err1, err2, err3, err4, l2params.a[p1], huberparams.a[p2],
            trimparams.a[p3], a0[p4])

end


using IterTools, ProgressBars
using JLD
using LaTeXStrings, PyPlot

include("synthetic.jl")
qrng_names = vcat(collect(keys(qrngs)), ["RR", "SO"])

RUNS    = 1
EPOCHS  = 1000
T       = n * EPOCHS

SEED    = 42
Random.seed!(SEED)

function gen_syn_data_all()
    # generate the training set consisting of n examples
    # sampled iid uniformly from the underlying distribution,
    # and the same training set will be used for all runs,
    gaussian_samples = gen_syn_data(n, "IID Uniform", d+1, SEED)
    return gaussian_samples
end

println("generating data...")
samples_offline = gen_syn_data_all()

# compute wopt, which is different from wgen since the objective is a finite sum
# so wopt should be the solution to the lstsq problem

X = hcat([pair[1] for pair in samples_offline]...)'
y = [pair[2] for pair in samples_offline]

wopt = (X'X)\(X'y)

println("begin training...")
result = Dict()

for qrng_name in qrng_names
    subopt_list = zeros(RUNS, T)

    for run in 1:RUNS
        # for offline mode we first need to determine sample ordering of the training set
        if qrng_name in collect(keys(qrngs))
            distrib = qrngs[qrng_name]
            # generate QMC indices
            sample_indices = [Int(ceil(i)) for i in QuasiMonteCarlo.sample(T, 0, n, distrib)]
        else
            if qrng_name == "SO"
                Random.seed!(SEED + run)
                perm = randperm!(collect(1:n))
                sample_indices = repeat(perm, EPOCHS)

            elseif qrng_name == "RR"
                sample_indices = zeros(Int64, T)
                for epoch in 1:EPOCHS
                    Random.seed!(SEED + run + epoch)
                    perm = randperm!(collect(1:n))
                    sample_indices[(epoch-1)*n+1 : epoch*n] = perm
                end
            else
                ArgumentError("SGD variant does not exist.")
            end
        end

        # run sgd for offline mode
        subopt_list[run, :] = sgd(T, w0, wopt, samples_offline, sample_indices)
    end
    result[qrng_name] = subopt_list
end

println("training completed")
println("plotting...")

# final plotting
figure(figsize = (6,6))
colors = ["#1f77bf", "#ff7f0e", "#17becf", "#7f7f7f"]
for (q, qrng_name) in enumerate(["IID Uniform", "Sobol", "RR", "SO"])
    subopt_list = result[qrng_name]
    subopt_mean = mean(subopt_list, dims=1)

    # use moving average to smooth out the curves as RR and SO are highly oscillatory
    # towards the end
    using RollingFunctions
    ws = 2000
    subopt_mean = runmean(reshape(subopt_mean, T), ws)

    loglog(collect(100:100:T), subopt_mean[100:100:T], label=qrng_name, color=colors[q])
end

legend(loc="lower left", fontsize=18)

xlabel(string("iterations ", L"t"), fontsize=20)
ylabel(L"\Vert w_t-w^*_n\Vert^2", fontsize=20)
title("Offline", fontsize=24)
xticks(fontsize=20)
yticks(fontsize=20)

ax = gca()
locmin = matplotlib.ticker.LogLocator(base=10.0, subs=(0.25,0.5,0.75), numticks=50)
ax.yaxis.set_minor_locator(locmin)
ax.yaxis.set_minor_formatter(matplotlib.ticker.NullFormatter())

tight_layout()
plot_fn = joinpath(pwd(),"figs", "qmc-offline.pdf")
savefig(plot_fn)
println("plot saved at $plot_fn")