n, p = 50, 5
((X, y, θ), (X1, y1, θ1)) = generate_continuous(n, p; seed=1212224)
maskint = vcat(ones(p), 0.0)

@testset "GH> Ridge" begin
    rng = StableRNG(551551)
    λ = 0.5
    # without fit_intercept
    s = R.scratch(X; i=false)
    # >> Hv!
    r = RidgeRegression(λ; fit_intercept=false, scale_penalty_with_samples = false)
    hv! = R.Hv!(r, X, y, s)
    v = randn(rng, length(θ))
    hv = similar(v)
    hv!(hv, v)
    @test hv ≈              X' * (X * v) .+ λ * v
    # ------------------
    # >> Hv! scale with samples
    r = RidgeRegression(λ; fit_intercept=false, scale_penalty_with_samples = true)
    hv! = R.Hv!(r, X, y, s)
    v = randn(rng, length(θ))
    hv = similar(v)
    hv!(hv, v)
    @test hv ≈              X' * (X * v) .+ λ * size(X, 1) * v
    # ------------------
    # with fit_intercept
    s = R.scratch(X; i=true)
    r = RidgeRegression(λ; penalize_intercept=true, scale_penalty_with_samples = false)
    hv! = R.Hv!(r, X, y, s)
    v = randn(rng, p+1)
    hv = similar(v)
    hv!(hv, v)
    @test hv ≈              X1' * (X1 * v) .+ λ * v
    # ------------------
    # with fit_intercept but no penalty
    r = RidgeRegression(λ, scale_penalty_with_samples = false)
    hv! = R.Hv!(r, X, y, s)
    v = randn(rng, p+1)
    hv = similar(v)
    hv!(hv, v)
    @test hv ≈              X1' * (X1 * v) .+ λ * v .* vcat(ones(p), 0.0)
end

@testset "GH> EN/Lasso" begin
    # without fit_intercept
    s = R.scratch(X; i=false)
    λ = 6.2
    γ = 0.7
    r = LassoRegression(λ; fit_intercept=false, scale_penalty_with_samples = false)
    fg! = R.smooth_fg!(r, X, y, s)
    g = similar(θ)
    f = fg!(g, θ)
    @test f ≈               sum(abs2.(X*θ .- y))/2
    @test g ≈               X' * (X * θ .- y)
    # with fit_intercept
    s = R.scratch(X; i=true)
    λ = 3.4
    γ = 2.7
    r = LassoRegression(λ; penalize_intercept=true, scale_penalty_with_samples = false)
    fg! = R.smooth_fg!(r, X, y1, s)
    g = similar(θ1)
    f = fg!(g, θ1)
    @test f ≈               sum(abs2.(X1*θ1 .- y1))/2
    @test g ≈               X1' * (X1 * θ1 .- y1)

    # elasticnet (with intercept)
    r = ElasticNetRegression(λ, γ; penalize_intercept=true, scale_penalty_with_samples = false)
    fg! = R.smooth_fg!(r, X, y1, s)
    g = similar(θ1)
    f = fg!(g, θ1)
    @test f ≈               sum(abs2.(X1*θ1 .- y1))/2 + λ .* norm(θ1)^2/2
    @test g ≈               X1' * (X1*θ1 .- y1) .+ λ .* θ1

    # elasticnet with intercept but no penalty of intercept
    r = ElasticNetRegression(λ, γ; scale_penalty_with_samples = false)
    fg! = R.smooth_fg!(r, X, y1, s)
    g = similar(θ1)
    f = fg!(g, θ1)
    @test f ≈               sum(abs2.(X1*θ1 .- y1))/2 + λ .* norm(θ1 .* maskint)^2/2
    @test g ≈               X1' * (X1*θ1 .- y1) .+ λ .* θ1 .* vcat(ones(p), 0)
end

n, p = 50, 5
((X, y, θ), (X1, y1, θ1)) = generate_binary(n, p; seed=1212224)
maskint = vcat(ones(p), 0.0)

@testset "GH> LogitL2" begin
    rng = StableRNG(551551)
    # fgh! without fit_intercept
    s = R.scratch(X; i=false)
    λ = 0.5
    lr = LogisticRegression(λ; fit_intercept=false, scale_penalty_with_samples = false)
    fgh! = R.fgh!(lr, X, y, s)
    θ = randn(rng, p)
    J = objective(lr, X, y)
    f = 0.0
    g = similar(θ)
    H = zeros(p, p)
    f = fgh!(f, g, H, θ)

    wminus = R.σ.(-y .* (X * θ))
    wplus  = R.σ.(y .* (X * θ))

    @test f == J(θ)
    @test g ≈ -X' * (y .* wminus) .+ λ .* θ
    @test H ≈  X' * (Diagonal(wplus .* (1 .- wplus)) * X) + λ * I

    # Use ForwardDiff
    logsigmoid(x)  = -log1p(exp(-x))
    objective2(θ_) = -sum(logsigmoid.(y .* (X*θ_))) + λ * sum(abs2, θ_) / 2
    fd_grad = ForwardDiff.gradient(θ_ -> objective2(θ_), θ)
    fd_hess = ForwardDiff.hessian(θ_ -> objective2(θ_), θ)
    @test g ≈ fd_grad
    @test H ≈ fd_hess

    # Hv! without  fit_intercept
    s = R.scratch(X; i=false)
    Hv! = R.Hv!(lr, X, y, s)
    v   = randn(rng, p)
    Hv  = similar(v)
    Hv!(Hv, θ, v)
    @test Hv ≈ H * v

    # fgh! with fit_intercept
    s = R.scratch(X; i=true)
    λ = 0.5
    lr1 = LogisticRegression(λ; penalize_intercept=true, scale_penalty_with_samples = false)
    fgh! = R.fgh!(lr1, X, y, s)
    θ1 = randn(rng, p+1)
    J  = objective(lr1, X, y)
    f1 = 0.0
    g1 = similar(θ1)
    H1 = zeros(p+1, p+1)
    f1 = fgh!(f1, g1, H1, θ1)

    wminus = R.σ.(-y .* (X1 * θ1))
    wplus  = R.σ.(y .* (X1 * θ1))

    @test f1 ≈ J(θ1)
    @test g1 ≈ -X1' * (y .* wminus) .+ λ .* θ1
    @test H1 ≈ X1' * (Diagonal(wplus .* (1 .- wplus)) * X1) + λ * I

    # Use ForwardDiff
    logsigmoid(x)  = -log1p(exp(-x))
    objective2(θ_) = -sum(logsigmoid.(y .* (X1*θ_))) + λ * sum(abs2, θ_) / 2
    fd_grad = ForwardDiff.gradient(θ_ -> objective2(θ_), θ1)
    fd_hess = ForwardDiff.hessian(θ_ -> objective2(θ_), θ1)
    @test g1 ≈ fd_grad
    @test H1 ≈ fd_hess

    # Hv! with fit_intercept
    Hv! = R.Hv!(lr1, X, y, s)
    v   = randn(rng, p+1)
    Hv  = similar(v)
    Hv!(Hv, θ1, v)
    @test Hv ≈ H1 * v

    # fgh! with fit intercept and no penalty on intercept
    lr1 = LogisticRegression(λ, scale_penalty_with_samples = false)
    fgh! = R.fgh!(lr1, X, y, s)
    θ1 = randn(rng, p+1)
    J  = objective(lr1, X, y)
    f1 = 0.0
    g1 = similar(θ1)
    H1 = zeros(p+1, p+1)
    f1 = fgh!(f1, g1, H1, θ1)

    wminus = R.σ.(-y .* (X1 * θ1))
    wplus  = R.σ.(y .* (X1 * θ1))

    @test f1 ≈ J(θ1)
    @test g1 ≈ -X1' * (y .* wminus) .+ λ .* θ1 .* maskint
    @test H1 ≈  X1' * (Diagonal(wplus .* (1.0 .- wplus)) * X1) + λ * Diagonal(maskint)
    Hv! = R.Hv!(lr1, X, y, s)
    v   = randn(rng, p+1)
    Hv  = similar(v)
    Hv!(Hv, θ1, v)
    @test Hv ≈               H1 * v
end

# Comparison with sklearn
@testset "GH> MultinL2" begin
    X = [ 0.78843 -0.28336;
         -0.75568  0.22546;
         -0.09012  0.68069;
         -0.34437 -0.98773;
          1.09285 -0.37161 ]
    n, p = size(X)
    c = 3
    y = [1, 2, 3, 1, 3]
    # comparison sklearn // no intercept
    θ = [-0.04843, 0.99519, -0.67237, 1.08812, 0.13362, 0.77136]
    g_sk = [-0.12941349639677957, 1.033822503077806, 0.6025709048825946, -0.3233237353163467,
            -0.47315740848581506, -0.7104987677614594]
    v = [-0.28802, -0.90018, 1.48613, 1.77976, -1.06333, 1.36275]
    Hv_sk = [0.18500344954627015, -0.9109175305006267, 0.8560420396655112, 0.4034408910070668,
            -1.0410454892117813, 0.5074766394935599]
    s = R.scratch(n, p, c; i=false)
    mnr = MultinomialRegression(0.0; fit_intercept=false)
    fg! = R.fg!(mnr, X, y, s)
    f = fg!(0.0, nothing, θ)
    mnl = MultinomialLoss()
    @test f ≈ mnl(y, X*reshape(θ, 2, 3))
    g = similar(θ)
    fg!(nothing, g, θ)
    @test g ≈ g_sk
    hv! = R.Hv!(mnr, X, y, s)
    Hv = similar(θ)
    hv!(Hv, θ, v)
    @test Hv ≈ Hv_sk

    # -- with intercept
    θ1 = [0.17905, 1.91598, 1.30329, -1.03438, -1.26994, -0.38288, -0.96238, -0.47912, 0.70813]
    g_sk1 = [0.26979349788581053, 1.5036100824512861, 0.7042460191768531, 0.7282930655963799,
            -0.6529621156985143, -0.2504598040370489, -0.9980865634821908, -0.8506479667527724,
            -0.45378621513980455]
    v1 = [-1.05534, -0.12016, 0.28984, 0.81068, 0.1617, -0.82832, -0.60992, 0.74291, -0.14735]
    Hv1_sk = [-0.20815994906492608, -0.03987039776759468, 0.38689109859835424, 0.3516398288195569,
               0.0703034889451111, -0.3054162498173413, -0.14347987975463106, -0.03043309117751647,
               -0.08147484878101276]
    s = R.scratch(n, p, c; i=true)
    mnr1 = MultinomialRegression(0.0; penalize_intercept=true)
    fg! = R.fg!(mnr1, X, y, s)
    f = fg!(0.0, nothing, θ1)
    @test  f ≈ mnl(y, R.apply_X(X, θ1, 3))
    g = similar(θ1)
    fg!(nothing, g, θ1)
    @test g ≈ g_sk1
    hv! = R.Hv!(mnr1, X, y, s)
    Hv1 = similar(θ1)
    hv!(Hv1, θ1, v1)
    @test Hv1 ≈ Hv1_sk
end

@testset "GH> Huber" begin
    rng = StableRNG(12341412)
    δ, λ  = 0.5, 3.4

    # without intercept
    s = R.scratch(X; i=false)
    hlr  = HuberRegression(δ, λ; fit_intercept=false, scale_penalty_with_samples = false)
    fgh! = R.fgh!(hlr, X, y, s)
    θ_   = randn(rng, p) # otherwise the residuals are too small and everything is in the l1-ball

    g = similar(θ)
    H = zeros(p, p)

    f = fgh!(0.0, g, H, θ_)
    r = X*θ_ .- y
    @test f == hlr.loss(r) + hlr.penalty(θ_)
    mask = abs.(r) .<= δ
    @test g ≈               (X' * (r .* mask)) .+ (X' * (δ .* sign.(r) .* .!mask)) .+ λ .* θ_
    @test H ≈                X' * (mask .* X) + λ*I

    Hv! = R.Hv!(hlr, X, y, s)
    Hv = similar(θ_)
    v = randn(rng, p)
    Hv!(Hv, θ_, v)

    @test Hv ≈               H * v

    # with intercept
    s = R.scratch(X; i=true)
    hlr1  = HuberRegression(δ, λ; penalize_intercept=true, scale_penalty_with_samples = false)
    fgh1! = R.fgh!(hlr1, X, y1, s)
    θ1_   = randn(rng, p+1)

    g1 = similar(θ1)
    H1 = zeros(p+1, p+1)

    Hv! = R.Hv!(hlr, X, y, s)
    Hv = similar(θ_)
    v = randn(rng, p)
    Hv!(Hv, θ_, v)

    @test Hv ≈ H * v

    # with intercept
    hlr1  = HuberRegression(δ, λ; penalize_intercept=true, scale_penalty_with_samples = false)
    fgh1! = R.fgh!(hlr1, X, y1, s)
    θ1_   = randn(rng, p+1)

    g1 = similar(θ1)
    H1 = zeros(p+1, p+1)

    f1 = fgh1!(0.0, g1, H1, θ1_)
    r1 = X1*θ1_ .- y1
    mask = abs.(r1) .<= δ
    @test f1 ≈ hlr1.loss(r1) + hlr1.penalty(θ1_)
    @test g1 ≈              (X1' * (r1 .* mask)) .+ (X1' * (δ .* sign.(r1) .* .!mask)) .+ λ .* θ1_
    @test H1 ≈               X1' * (mask .* X1) + λ*I

    Hv1! = R.Hv!(hlr1, X, y1, s)
    Hv1 = similar(θ1_)
    v1 = randn(rng, p+1)
    Hv1!(Hv1, θ1_, v1)

    @test Hv1 ≈              H1 * v1

    # with intercept and no penalty on intercept
    hlr1  = HuberRegression(δ, λ, scale_penalty_with_samples = false)
    fgh1! = R.fgh!(hlr1, X, y1, s)
    θ1_   = randn(rng, p+1)

    g1 = similar(θ1)
    H1 = zeros(p+1, p+1)

    f1 = fgh1!(0.0, g1, H1, θ1_)
    r1 = X1*θ1_ .- y1
    mask = abs.(r1) .<= δ
    @test f1 ≈ hlr1.loss(r1) + hlr1.penalty(θ1_ .* maskint)
    @test g1 ≈              (X1' * (r1 .* mask)) .+ (X1' * (δ .* sign.(r1) .* .!mask)) .+
                                λ .* θ1_ .* maskint
    @test H1 ≈               X1' * (mask .* X1) + λ * Diagonal(maskint)

    Hv1! = R.Hv!(hlr1, X, y1, s)
    Hv1 = similar(θ1_)
    v1 = randn(rng, p+1)
    Hv1!(Hv1, θ1_, v1)

    @test Hv1 ≈              H1 * v1
end


#!/usr/bin/env julia

using Seq

fibs = cat([0, 1], @lazyseq map(+, rest(fibs), fibs))

for k in take(10, fibs)
    println(k)
end


type Generator
  # specified
  t::Vector{Float64}
  n::Int64
  "number of dimensions."
  k::Int64

  # computed
  d::Tuple
  y::Matrix{Float64}
  dt::Vector{Float64}
  sdt::Vector{Float64}
  b::Matrix{Float64}
  db::Matrix{Float64}

  function Generator(t, s, k::Int)
    g      = new(t, length(t), k)
    g.d    = (s..., k)
    g.dt   = diff(t)
    g.sdt  = sqrt(g.dt)
    g.b    = zeros(g.d)
    g.db   = zeros(g.d)
    g.y    = zeros(g.d)
    return g
  end
end

@inline function next{P <: AProcess}(process::P, g::Generator, i::Int64)
  if i == 1
    return init!(process, g)
  end

  randn!(g.db)
  cmul!(g.db, g.sdt[i-1])
  result = next!(process, g, i)
  # Important to use previous b for process computation not to have bias.
  add!(g.b, g.db)
  result
end

function cumsim{P <: AProcess}(process::P, t, k::Int=1)
  s=size(vwrap(initial(process)))
  gen = Generator(vec(t), s, k)
  result = zeros(gen.n, s..., k)
  for i in 1:gen.n
    y = next(process, gen, i)
    result[i,:,:] .= y[:,:]
  end
  CumsimResult(collapse_but_first(result), t)
end

function sim{P <: AProcess}(process::P, t, k::Int=1)
  s=size(vwrap(initial(process)))
  gen = Generator(t, s, k)
  for i in 1:(gen.n-1)
    next(process, gen, i)
  end
  SimResult(collapse_but_first(next(process, gen, gen.n)), t[end], gen.b)
end


"""
First-order MAML (FOMAML) from
Model-Agnostic Meta-Learning for Fast Adapdation of Deep Networks [1].

[1] https://arxiv.org/abs/1703.03400
"""
function fomaml(model; meta_opt=Descent(0.01), inner_opt=Descent(0.02), epochs=30_000,
              n_tasks=3, train_batch_size=10, eval_batch_size=10, eval_interval=1000)

    weights = params(model)
    dist = Uniform(-5, 5)
    testx = Float32.(range(-5, stop=5, length=50))

    for i in 1:epochs
        prev_weights = deepcopy(Flux.data.(weights))

        for _ in 1:n_tasks
            task = SineWave()

            x = Float32.(rand(dist, train_batch_size))
            y = task(x)
            grad = Flux.Tracker.gradient(() -> Flux.mse(model(x'), y'), weights)

            for w in weights
                w.data .-= Flux.Optimise.apply!(inner_opt, w.data, grad[w].data)
            end

            testy = task(testx)
            grad = Flux.Tracker.gradient(() -> Flux.mse(model(testx'), testy'), weights)

            # reset weights and accumulate gradients
            for (w1, w2) in zip(weights, prev_weights)
                w1.data .= w2
                w1.grad .+= grad[w1].data
            end

        end

        Flux.Optimise._update_params!(meta_opt, weights)

        if i % eval_interval == 0
            @printf("Iteration %d, evaluating model on random task...\n", i)
            evalx = Float32.(rand(dist, eval_batch_size))
            eval_model(model, evalx, testx, SineWave())
        end

    end
end


import Base.&

#=
&,And : Logical & of A and B
=#
function Base.&(A::Assoc, B::Assoc)
#First, create the row and col of the intersection
ABrow = intersect(A.row,B.row)
ABcol = intersect(A.col,B.col)
#Filling the sparse matrix with

AA = spzeros(size(ABrow,1), size(ABcol,1))
rowMapping = searchsortedmapping(ABrow,A.row)
colMapping = searchsortedmapping(ABcol,A.col)
AA = spones(A.A[rowMapping,colMapping])
AA = round(Int64,AA)

BB = spzeros(size(ABrow,1), size(ABcol,1))
rowMapping = searchsortedmapping(ABrow,B.row)
colMapping = searchsortedmapping(ABcol,B.col)
BB = spones(B.A[rowMapping,colMapping])
BB = round(Int64,BB)

ABA = AA & BB
ABA = ABA * 1.0

return Assoc(ABrow,ABcol,promote([1.0],A.val)[1],ABA)
end

########################################################
# D4M: Dynamic Distributed Dimensional Data Model
# Architect: Dr. Jeremy Kepner (kepner@ll.mit.edu)
# Software Engineer: Alexander Chen (alexc89@mit.edu)
########################################################


using Revise # reduce recompilation time
using Plots
# using Documenter
using LinearAlgebra
using SparseArrays
using BenchmarkTools
using UnPack
using ToeplitzMatrices

push!(LOAD_PATH, "./src")
using CommonUtils
using Basis1D
using Basis2DQuad
using UniformQuadMesh

using SetupDG

# TODO: refactor the code
function fS(uL, uR)
    return 1/6*(uL^2+uR^2+uL*uR)
end

function burgers_exact_sol_2D(u0,x,y,T,dt)
    Nsteps = ceil(Int,T/dt)
    dt = T/Nsteps
    u = u0(x,y) # computed at input points
    for i = 1:Nsteps
        t = i*dt
        u .= @. u0(x-u*t,y) # evolve solution at quadrature points using characteristics
    end
    return u
end

S_N(x) = @. sin(pi*x/h)/(2*pi/h)/tan(x/2)
"""
Vandermonde matrix of sinc basis functions determined by h,
evaluated at r
"""
function vandermonde_Sinc(h,r)
    N = convert(Int, 2*pi/h)
    V = zeros(length(r),N)
    for n = 1:N
        V[:,n] = S_N(r.-n*h)
    end
    V[1,1] = 1
    return V
end

"Constants"
const sp_tol = 1e-12
# u0(x,y) = @. -sin(pi*x)*sin(y)
u0(x,y) = @. exp(-1*((y-pi)^2+2*pi*x^2))
# u0(x,y) = @. exp(-1*((y-pi)^2))
# u0(x,y) = @. exp(-10*(x^2))

"Program parameters"
compute_L2_err = false

"Approximation Parameters"
N_P = 8;    # The order of approximation in polynomial dimension
Np_P = N_P+1;
Np_F = 16;    # The order of approximation in Fourier dimension
CFL  = 0.2;
T    = 3.0;  # End time

"Time integration Parameters"
rk4a,rk4b,rk4c = rk45_coeffs()
CN = (max(N_P,Np_F)+1)*(max(N_P,Np_F)+2)/2  # estimated trace constant for CFL
dt = CFL / CN
Nsteps = convert(Int,ceil(T/dt))
dt = T/Nsteps

"Initialize Reference Element in Fourier dimension"
h = 2*pi/Np_F
column = [0; .5*(-1).^(1:Np_F-1).*cot.((1:Np_F-1)*h/2)]
column2 = [-pi^2/3/h^2-1/6; -((-1).^(1:Np_F-1)./(2*(sin.((1:Np_F-1)*h/2)).^2))]
Ds = Toeplitz(column,column[[1;Np_F:-1:2]])
D2s = Toeplitz(column2,column2[[1;Np_F:-1:2]])
s = LinRange(h,2*pi,Np_F)

"Initialize Reference Element in polynomial dimension"
rd = init_reference_interval(N_P);
@unpack r,rq,wq,rf,Dr,VDM,Vq,Vf,Vp,M,Pq,LIFT,nrJ = rd
wf = [1;1]
Nq_P = length(rq)
rp = LinRange(-1,1,100)
Vp = vandermonde_1D(N_P,rp)/VDM # Redefine Vp
Lq = LIFT

"Initialize Mesh"
s,r = meshgrid(s,r)
Nfp = 2*Np_F

"scale by Fourier dimension"
M = h*M
wq = h*wq
wf = h*wf

"Hybridized operators"
Ef = Vf*Pq
Br = diagm(wf.*nrJ)
Qr = Pq'*M*Dr*Pq
Qs = Ds
Drh = [Vq*Dr*Pq-1/2*Vq*Lq*diagm(nrJ)*Vf*Pq  1/2*Vq*Lq*diagm(nrJ);
       -1/2*diagm(nrJ)*Vf*Pq                1/2*diagm(nrJ)]
Qrh = 1/2*[Qr-Qr' Ef'*Br;
           -Br*Ef Br]
Qrh_skew = 1/2*(Qrh-Qrh')
Qsh = Qs # Not the SBP operator, weighted when flux differencing

ops = (Vq,Vf,wq,wf,Pq,Lq,Qrh,Qsh,D2s,nrJ,h)
function rhs(u,ops,compute_rhstest)
    Vq,Vf,wq,wf,Pq,Lq,Qrh,Qsh,D2s,nrJ,h = ops
    # Volume term
    uq,uf = (A->A*u).((Vq,Vf))
    uh = [uq;uf]
    w = [wq;wf]
    ∇fh = zeros(size(uh))

    # Flux differencing in x direction
    for k = 1:size(uh,2)
        ∇fh[:,k] += 2*Qrh.*[fS(uL,uR) for uL in uh[:,k], uR in uh[:,k]]*ones(size(uh,1),1)
    end

    # Flux differencing in y direction
    for k = 1:size(uh,1)
        ∇fh[k,:] += 2*pi*[wq;0;0][k]*Qsh.*[fS(uL,uR) for uL in uh[k,:], uR in uh[k,:]]*ones(size(uh,2),1)
    end

    # Spatial term
    ∇f = [Pq Lq]*diagm(1 ./ w)*∇fh

    # Flux term
    uM = Vf*u
    uP = [uM[2,:] uM[1,:]]'
    LF = @. max(abs(uP),abs(uM))*(uP-uM)
    uflux = diagm(nrJ)*(@. fS(uM,uP)-uM*uM/2)-LF

    # Artificial Viscosity
    # ϵ = 0.05
    ϵ = 0.0
    Δf = ϵ*u*D2s'
    rhsu = -(∇f+Lq*uflux-Δf)

    rhstest = 0
    if compute_rhstest
        rhstest += sum(diagm(wq)*uq.*(Vq*rhsu))
    end

    return rhsu,rhstest
end

u = @. u0(r,s)
# Test single element
resu = zeros(size(u))
for i = 1:Nsteps
    rhstest = 0
    for INTRK = 1:5
        compute_rhstest = INTRK==5
        rhsu,rhstest = rhs(u,ops,compute_rhstest)
        @. resu = rk4a[INTRK]*resu + dt*rhsu
        @. u += rk4b[INTRK]*resu
    end

    if i%10==0 || i==Nsteps
        println("Time step: $i out of $Nsteps with rhstest = $rhstest")
    end
end

# Plotting
gr(aspect_ratio=1/pi, legend=false,
   markerstrokewidth=0, markersize=2.5,
   xlims=(-1,1),ylims=(h,2*pi),showaxis=false)
plot()

sp = LinRange(h,2*pi,100)
VDM_F = vandermonde_Sinc(h,sp)
V2 = vandermonde_1D(1,sp)/vandermonde_1D(1,LinRange(h,2*pi,Np_F))

plt = scatter(Vp*r*V2',Vp*s*V2',Vp*u*VDM_F',zcolor=Vp*u*VDM_F',camera=(0,90))
display(plt)
# exact_sol = burgers_exact_sol_2D(u0,Vp*r*V2',Vp*s*V2',T,dt/100)
# scatter(Vp*r*V2',Vp*s*V2',exact_sol,zcolor=exact_sol,camera=(0,90))

# L2 error
if compute_L2_err
    rq2,wq2 = gauss_quad(0,0,N_P+2)
    wq2 = wq2*h
    Vq2 = vandermonde_1D(N_P,rq2)/VDM
    xq2 = Vq2*r
    yq2 = Vq2*s

    exact_sol = burgers_exact_sol_2D(u0,xq2,yq2,T,dt/100)
    L2_err = sum(diagm(wq2)*(Vq2*u-exact_sol).^2)
    @show L2_err
end


# he linearized forward modeling operator `F`, such that
# `Fx` can be computed without explicitly storing the operator matrix (
#  see `LinearMaps.jl`);
# the imaging/migration operator `F*`;

# These maps are the building blocks of iterative optimization schemes.

for scenario in [:downhole, :pizza]
	println("@@@@@@@@@@@@TESTING ", scenario)
	for rfields in [[:P], [:Vx], [:Vz]]
		pa=JG.SeisInvExpt(scenario, born_flag=true, rfields=rfields)

		F=JF.operator_Born(pa);

		x1=randn(size(F,2))
		x2=randn(size(F,2))
		x12=x1.+x2

		d12=F*x12
		d1=F*x1
		δmodtt1=copy(pa.paf.c.δmodtt)
		d2=F*x2
		δmodtt2=copy(pa.paf.c.δmodtt)

		d12new=d1.+d2

		f=Misfits.error_squared_euclidean!(nothing, d12, d12new, nothing, norm_flag=true)

		@test f<1e-25

		function adjtest()
			x=randn(size(F,2))
			y=randn(size(F,1))
			a=LinearAlgebra.dot(y,F*x)
			b=LinearAlgebra.dot(x,adjoint(F)*y)
			c=LinearAlgebra.dot(x, transpose(F)*F*x)
			println("adjoint test: ", a, "\t", b)
			@test isapprox(a,b,rtol=1e-5)
			println("must be positive: ", c)
			@test c>0.0
		end

		for i in 1:3
			adjtest()
		end
	end
end


using SimJulia
using Distributions

function visit(customer::Process,
 			time_in_bank::Float64,
 			counter::Resource,
 			max_in_queue::Int)
  arrive = now(customer);
  @printf("%8.4f %s: Here I am\n", arrive, customer);
  if length(counter.wait_queue) < max_in_queue
      request(customer, counter);
      wait = now(customer) – arrive;
      @printf("%8.4f %s: Waited %6.3f\n",
 			now(customer), customer, wait);
      hold(customer, time_in_bank)
      release(customer, counter);
      @printf("%8.4f %s: Finished\n", now(customer), customer)
  else
      @printf("%8.4f %s: Waiting\n", now(customer), customer)
  end
end

function generate(source::Process, number::Int,
 			mean_time_between_arrivals::Float64,
 			mean_time_in_bank::Float64,
 			counter::Resource,
 			max_in_queue::Int)
  d_tba = Exponential(mean_time_of_arrivals);
  d_tib = Exponential(mean_time_in_bank);
  for i = 1:number
      c = Process(simulation(source),
 			@sprintf("Customer%02d", i));
    tib = rand(d_tib);
    activate(c, now(source), visit, tib,
    			counter, max_in_queue);
    tba = rand(d_tba);
    hold(source, tba);
  end
end

function number_in_system(counter::Resource)
  length(counter.active_set) + length(counter.wait_queue);
end

# Set simulation data
max_number = 8;
max_time = 3900.0;
max_queue = 1;
mean_arrival_time = 10.0;
mean_service_time = 12.0;
number_of_tellers = 1;
seed = randomize();	# Use function defined early or pick
# a value to make run reproducible

# Model/Experiment
srand(seed);
sim = Simulation(uint(16));
k = Resource(sim, "Counter", uint(1), false);
s = Process(sim, "Source");
activate(s, 0.0,  generate,
max_number, mean_arrival_time,
 	mean_service_time, k, max_queue);
run(sim, max_time);


using Revise
using NNFEM
using PyPlot

NT = 200
T = 1.0
Δt = T/NT

m, n =  15, 15
h = 1/m

# Create a very simple mesh
elements = []
prop = Dict("name"=> "PlaneStrain", "rho"=> 1.0, "E"=> 2.0, "nu"=> 0.35)
coords = zeros((m+1)*(n+1), 2)
for j = 1:n
    for i = 1:m
        idx = (m+1)*(j-1)+i
        elnodes = [idx; idx+1; idx+1+m+1; idx+m+1]
        ngp = 3
        nodes = [
            (i-1)*h (j-1)*h
            i*h (j-1)*h
            i*h j*h
            (i-1)*h j*h
        ]
        coords[elnodes, :] = nodes
        push!(elements, SmallStrainContinuum(nodes, elnodes, prop, ngp))
    end
end

# fixed on the bottom, push on the right
EBC = zeros(Int64, (m+1)*(n+1), 2)
FBC = zeros(Int64, (m+1)*(n+1), 2)
g = zeros((m+1)*(n+1), 2)
f = zeros((m+1)*(n+1), 2)

for i = 2:m
    idx = n*(m+1)+i
    EBC[idx,:] .= -2
    EBC[i,:] .= -2
end

for j = 1:n+1
    idx = (j-1)*(m+1)+1
    EBC[idx,:] .= -2
    idx = (j-1)*(m+1)+m+1
    EBC[idx,:] .= -2
end

ndims = 2
domain = Domain(coords, elements, ndims, EBC, g, FBC, f)

EBC_idx = findall(EBC[:,1] .== -2)
EBC_X = domain.nodes[EBC_idx,:]
function gt_func(t)
    x, y = EBC_X[:,1], EBC_X[:,2]
    u1 = @. (1-y^2)*(x^2+y^2)*exp(-t) * 0.1
    u2 = @. (1-y^2)*(x^2-y^2)*exp(-t) * 0.1
    return [u1;u2], -[u1;u2], [u1;u2]
end

FBC_idx = findall(FBC[:,1] .== -2)
FBC_X = domain.nodes[FBC_idx,:]
# function ft(t)
#     x, y = FBC_X[:,1], FBC_X[:,2]
#     T1 = @. -1.48148148148148*x*exp(-t) - 1.48148148148148*y*exp(-t)
#     T2 = @. -3.45679012345679*x*exp(-t) - 6.41975308641975*y*exp(-t)
#     return [T1;T2] * h * 0.1
# end

function bt(x, y, t)
    f1 = @. 0.987654320987654*x*y*exp(-t) + 0.592592592592593*y^2*exp(-t) + (0.1 - 0.1*y^2)*(x^2 + y^2)*exp(-t) - (0.148148148148148 - 0.148148148148148*y^2)*exp(-t) - (0.641975308641975 - 0.641975308641975*y^2)*exp(-t) - (-0.148148148148148*x^2 - 0.148148148148148*y^2)*exp(-t)
    f2 = @. 0.987654320987654*x*y*exp(-t) - 2.5679012345679*y^2*exp(-t) + (0.1 - 0.1*y^2)*(x^2 - y^2)*exp(-t) - (0.148148148148148 - 0.148148148148148*y^2)*exp(-t) - (-0.641975308641975*x^2 + 0.641975308641975*y^2)*exp(-t) - (0.641975308641975*y^2 - 0.641975308641975)*exp(-t)
    [f1 f2]
end

c = domain.nodes[domain.dof_to_eq]
nn = div(length(c), 2)
x = c[1:nn]
y = c[nn+1:end]

Dstate = zeros(domain.neqs) # d at last step
state = [(@. (1-y^2)*(x^2+y^2)); (@. (1-y^2)*(x^2-y^2))] * 0.1
velo = -[(@. (1-y^2)*(x^2+y^2)); (@. (1-y^2)*(x^2-y^2))] * 0.1
acce = [(@. (1-y^2)*(x^2+y^2)); (@. (1-y^2)*(x^2-y^2))] * 0.1
globdat = GlobalData(state, Dstate, velo, acce, domain.neqs, gt_func, nothing, bt)

assembleMassMatrix!(globdat, domain)
updateTimeDependentEssentialBoundaryCondition!(domain,globdat)
updateStates!(domain, globdat)

# #solve for a0
# SolverInitial!(Δt, globdat, domain)

ω = EigenMode(Δt, globdat, domain)
@show "stable time step is ", 0.8 * 2/ω, " current time step is ", Δt

for i = 1:NT
    @info i
    #global globdat, domain = GeneralizedAlphaSolverStep(globdat, domain, Δt)
    global globdat, domain = ExplicitSolverStep(globdat, domain, Δt)
end
# # # visualize_displacement(domain)
# plot(hcat(domain.history["state"]...)[n*(m+1)+1,:])
# plot(exp.(-LinRange(0,1,NT+1)))

ts = ExplicitSolverTime(Δt, NT)
d_ = hcat(domain.history["state"]...)'

using Random; Random.seed!(233)
for k = 1:5
    i = rand(1:m+1)
    j = rand(1:n+1)
    if k==1
        plot(d_[:,(j-1)*(m+1)+i], color = "C$k", label="Computed")
    else
        plot(d_[:,(j-1)*(m+1)+i], color = "C$k")
    end
    x0 = (i-1)*h
    y0 = (j-1)*h

    if k==1
        plot((@. (1-y0^2)*(x0^2+y0^2)*exp(-ts))*0.1 ,"--", color="C$k", label="Reference")
    else
        plot((@. (1-y0^2)*(x0^2+y0^2)*exp(-ts))*0.1 ,"--", color="C$k")
    end
end
legend()


using IterativelyHypergraphEmbedding
using SimpleHypergraphs
using LinearAlgebra
using Test

@testset "IterativelyHypergraphEmbedding.jl" begin
  h = Hypergraph{Int, Int}(4, 0)
  add_hyperedge!(h, vertices=Dict(1=>1, 2=>1, 3=>1))
  add_hyperedge!(h, vertices=Dict(3=>1, 4=>1))
  dims = 10
  hn_vec, he_vec = iteratively_hypergraph_embedding(h, dims)
  @test size(hn_vec) == (dims, nhv(h))
  @test size(he_vec) == (dims, nhe(h))

  # Test that the hypernodes and hyperedges are embedded in the hypersphere.
  for vec in eachcol(hn_vec)
    @test abs(norm(vec) - 1) < 1e-8
  end
  for vec in eachcol(he_vec)
    @test abs(norm(vec) - 1) < 1e-8
  end

end


module TestRandomAdditions

using ConicHulls, ConicHulls.RefHull

function test(NC::Int, ngen::Int, r::Int)
    eval(ConicHulls.Hulls,:(numfacets=0))
    eval(ConicHulls.Primitives,:(numgenerators=0))
    test(create_hull(NC), ngen, r)
end

function test(hull::ConicHull, ngen::Int, r::Int)
    NC, F, G = nconic(hull), ftype(hull), gtype(hull)

    gs = init_hull!(hull, eye(Int, NC))

    for k=1:ngen
        verify(hull)
        verify_hull(hull, gs)

        x = rand(-r:r, NC-1)
        push!(x, 1-sum(x))
        g = G(x)

        add!(hull, g)
        push!(gs, g)
    end
end

srand(673826715664)
for NC in 3:5
    @time for r=1:8, ngen=1:8, k=1:(1<<(8-NC))/ngen; test(NC, ngen, r); end
end

end # module


eachindex(arg::MultiValue) = eachindex(1:prod(size(arg)))

CartesianIndices(arg::MultiValue) = CartesianIndices(size(arg))

LinearIndices(arg::MultiValue) = LinearIndices(size(arg))

getindex(arg::VectorValue, i::Integer) = arg.data[i]

function getindex(arg::TensorValue{D},i::Integer,j::Integer) where D
  index = _2d_tensor_linear_index(D,i,j)
  arg.data[index]
end

function getindex(arg::SymTensorValue{D},i::Integer,j::Integer) where D
  index = _2d_sym_tensor_linear_index(D,i,j)
  arg.data[index]
end

function getindex(arg::SymFourthOrderTensorValue{D},i::Integer,j::Integer,k::Integer,l::Integer) where D
  index = _4d_sym_tensor_linear_index(D,i,j,k,l)
  arg.data[index]
end

function getindex(arg::ThirdOrderTensorValue{D1,D2},i::Integer,j::Integer,k::Integer) where {D1,D2}
  index = _3d_tensor_linear_index(D1,D2,i,j,k)
  arg.data[index]
end

getindex(arg::VectorValue, ci::CartesianIndex{1})              = getindex(arg,ci[1])
getindex(arg::TensorValue,ci::CartesianIndex{2})               = getindex(arg,ci[1],ci[2])
getindex(arg::SymTensorValue,ci::CartesianIndex{2})            = getindex(arg,ci[1],ci[2])
getindex(arg::ThirdOrderTensorValue,ci::CartesianIndex{3})     = getindex(arg,ci[1],ci[2],ci[3])
getindex(arg::SymFourthOrderTensorValue,ci::CartesianIndex{4}) = getindex(arg,ci[1],ci[2],ci[3],ci[4])

getindex(arg::MultiValue, i::Integer) = getindex(arg, CartesianIndices(arg)[i])
getindex(arg::TensorValue, i::Integer) = arg.data[i]
getindex(arg::ThirdOrderTensorValue, i::Integer) = arg.data[i]

@inline iterate(arg::MultiValue)        = iterate(arg.data)
@inline iterate(arg::MultiValue, state) = iterate(arg.data, state)

data_index(::Type{<:VectorValue},i) = i
data_index(::Type{<:TensorValue{D}},i,j) where D = _2d_tensor_linear_index(D,i,j)
data_index(::Type{<:SymTensorValue{D}},i,j) where D = _2d_sym_tensor_linear_index(D,i,j)
data_index(::Type{<:ThirdOrderTensorValue{D1,D2}},i,j,k) where {D1,D2} = _3d_tensor_linear_index(D1,D2,i,j,k)
data_index(::Type{<:SymFourthOrderTensorValue{D}},i,j,k,l) where D = _4d_sym_tensor_linear_index(D,i,j,k,l)

_symmetric_index_gaps(i::Integer) = i*(i-1)÷2

_2d_tensor_linear_index(D,i,j) = ((j-1)*D)+i

_3d_tensor_linear_index(D1,D2,i,j,k) = (k-1)*D1*D2+(j-1)*D1+i

function _2d_sym_tensor_linear_index(D,i,j)
  _j=min(i,j)
  _i=max(i,j)
  index=_2d_tensor_linear_index(D,_i,_j)-_symmetric_index_gaps(_j)
  index
end

#function _4d_sym_tensor_linear_index(D,i,j,k,l)
#  _j=min(i,j)
#  _i=max(i,j)
#  _l=min(k,l)
#  _k=max(k,l)
#  block_length=_symmetric_index_gaps(D+1)
#  element_index=_2d_tensor_linear_index(D,_i,_j)-_symmetric_index_gaps(_j)
#  block_index=_2d_tensor_linear_index(D,_l,_k)-_symmetric_index_gaps(_l)
#  index=(block_index-1)*block_length+element_index
#  index
#end

function _4d_sym_tensor_linear_index(D,i,j,k,l)
  block_length = (D*(D+1))÷2
  block_index = _2d_sym_tensor_linear_index(D,i,j)
  element_index = _2d_sym_tensor_linear_index(D,k,l)
  index=(block_index-1)*block_length+element_index
  index
end


using Test, DSAModels

function read_dsa_data(filename)
    f = open(filename, "r")
    N_mach = read(f, Int64)
    M = read!(f, Vector{Float64}(undef, N_mach))
    η = read!(f, Vector{Float64}(undef, N_mach))
    η_r = read!(f, Vector{Float64}(undef, N_mach))
    close(f)

    return N_mach, M, η, η_r
end

@testset "DSAModels" begin

    @testset "Mach Models" begin

        @testset "Kang et. al. 2007" begin
            η_model = Kang07()
            N_mach, M, η, η_r = read_dsa_data("Kang07.dat")

            @testset "η" begin

                @test η_Ms(η_model, 0.5, 0.0) == 0.0

                for i = 1:N_mach
                    @test η_Ms(η_model, M[i], 0.0) ≈ η[i]
                end
            end

            @testset "η_r" begin
                @test η_Ms(η_model, 0.5, 1.0) == 0.0

                for i = 1:N_mach
                    @test η_Ms(η_model, M[i], 1.0) ≈ η_r[i]
                end
            end

        end

        @testset "Kang & Ryu 2013" begin
            η_model = KR13()
            N_mach, M, η, η_r = read_dsa_data("KR13.dat")

            @testset "η" begin
                for i = 1:N_mach
                    @test η_Ms(η_model, M[i], 0.0) ≈ η[i]
                end
            end

            @testset "η_r" begin
                for i = 1:N_mach
                    @test η_Ms(η_model, M[i], 1.0) ≈ η_r[i]
                end
            end

        end

        @testset "Caprioli & Spitkovsky 2014" begin
            η_model = CS14()
            N_mach, M, η, η_r = read_dsa_data("CS14.dat")

            @testset "η" begin
                for i = 1:N_mach
                    @test η_Ms(η_model, M[i], 0.0) ≈ η[i]
                end
            end

            @testset "η_r" begin
                for i = 1:N_mach
                    @test η_Ms(η_model, M[i], 1.0) ≈ η_r[i]
                end
            end

        end

        @testset "Ryu et. al. 2019" begin
            η_model = Ryu19()
            N_mach, M, η, η_r = read_dsa_data("Ryu19.dat")

            @testset "η" begin
                for i = 1:N_mach
                    @test η_Ms(η_model, M[i], 0.0) ≈ η[i]
                end
            end

            @testset "η_r" begin
                for i = 1:N_mach
                    @test η_Ms(η_model, M[i], 1.0) ≈ η_r[i]
                end
            end

        end

        @testset "Pfrommer et. al. 2016" begin
            η_model = P16()

            @test η_Ms(η_model, 10.0, 0.0) == 0.5
            @test η_Ms(η_model, 10.0, 1.0) == 0.5
        end
    end

    @testset "B-field Models" begin

        @testset "Electrons" begin
            # read reference data
            N_θ, θ_B, η, η_r = read_dsa_data("theta_B_e.dat")

            @testset "η" begin
                for i = 1:N_θ
                    @test ηB_acc_e(θ_B[i]) ≈ η[i]
                end
            end
            @testset "η_r" begin
                for i = 1:N_θ
                    @test ηB_reacc_e(θ_B[i]) ≈ η_r[i]
                end
            end
        end

        @testset "Protons" begin
            # read reference data
            N_θ, θ_B, η, η_r = read_dsa_data("theta_B_p.dat")

            @testset "η" begin
                for i = 1:N_θ
                    @test ηB_acc_p(θ_B[i]) ≈ η[i]
                end
            end
            @testset "η_r" begin
                for i = 1:N_θ
                    @test ηB_reacc_p(θ_B[i]) ≈ η_r[i]
                end
            end
        end
    end

end


#Sets
set_generators = 1:2
set_nodes = 1:3
set_nodes_ref=3
set_nodes_noref=1:2
set_nodes_g= 1:2
set_nodes_d= 2:3
set_demands = 1:2

#Parameters
generators=[100      12;
            80       20]

demands=[100  40;
         50   35]

capacity=1
varcost=2
load_node=1
utility_node=2

network=[ 1  2  100 500;
          1  3   40 500;
          2  3  100 500;
          2  1  100 500;
          3  1  100 500;
          3  2  100 500]

start_node=network[:,1]
end_node=network[:,2]
F_max=network[:,3]
B=network[:,4]

n_node=3
n_link=6

set_links=1:n_link
#create a n*m Tuple Array with ‘undefined’ contents
links = Array{Tuple{Int64, Int64},2}(undef, n_link, 1)

for i=1:n_link
  links[i] = (start_node[i], end_node[i])
end

links_rev = Array{Tuple{Int64, Int64},2}(undef, n_link, 1)

for i=1:n_link
  links_rev[i] = (end_node[i], start_node[i])
end

F_max_dict=Dict()
B_dict=Dict()

for i=1:n_link
  F_max_dict[(start_node[i], end_node[i])] = F_max[i]
  B_dict[(start_node[i], end_node[i])] = B[i]
end

MapG=[(1,1), (2,2), (0,0)]
MapD=[(1,2),(2,3),(0,0)]


export Degree, Variable, show, iscomplex, isreal, isbool

"""
    Degree(explvar::Int, conjvar::Int)

Define a mathematical degree, which is used to define a variable exponent to
`explvar` and its conjugate exponent to `conjvar`.

### Attributes
- `explvar`: Int
- `conjvar`: Int
"""
mutable struct Degree
    explvar::Int
    conjvar::Int
end

#############################
## Equality, hash
#############################
==(d1::Degree, d2::Degree) = (d1.explvar==d2.explvar) && (d1.conjvar==d2.conjvar)
!=(d1::Degree, d2::Degree) = !(d1 == d2)
hash(deg::Degree, h::UInt) = hash(deg.explvar, hash(deg.conjvar, h))

"""
Variable(varname::String, Complex)

Define a mathematical variable `varname` of a certain mathematical kind,
`Complex` here.

### Attributes
- `name`: String
- `kind`: a type among `Complex`, `Real` and `Bool`
"""
mutable struct Variable <: AbstractPolynomial
    name::String
    kind::Type

    function Variable(name, kind)
        validtypes = Set([Complex, Real, Bool])
        if true ∉ map(x-> kind <: x, validtypes) || typeof(name) ∉ Set([String, SubString{String}])
            error(LOGGER, "Variable() : attempting to define a variable $name of type $kind, supported types are {Complex, Real, Bool}")
        end
        return new(String(name), kind)
    end
end

#############################
## Equality, hash
#############################
==(x::Variable, y::Variable) = (x.name==y.name) && (x.kind==y.kind)
!=(x::Variable, y::Variable) = !(x == y)
hash(var::Variable, h::UInt) = hash(var.name, hash(var.kind, h))

#############################
## Type detection utils
#############################
iscomplex(x::Variable) = x.kind <: Complex
isreal(x::Variable) = x.kind <: Real
isbool(x::Variable) = x.kind <: Bool

#############################
## Show
#############################
function Base.show(io::IO, x::Variable)
    print(io, x.name)
end

function Base.show(io::IO, d::Degree)
    print(io, "(", d.explvar, ",", d.conjvar, ")")
end


# logbeta
using CairoMakie, LaTeXStrings, SpecialFunctions
let
    x = 0.1:0.1:15
    cycle = Cycle([:color, :linestyle, :marker], covary = true)
    set_theme!(Lines = (cycle = cycle,), Scatter = (cycle = cycle,))
    fig = Figure(resolution = (600, 400), font = "CMU Serif", backgroundcolor = :honeydew) # probably you need to install this font in your system
    ax = Axis(fig[1, 1], xlabel = L"x", ylabel = L"Y_{\nu}(x)", ylabelsize = 22,
        xlabelsize = 22, xgridstyle = :dash, ygridstyle = :dash, xtickalign = 1,
        xticksize = 10, ytickalign = 1, yticksize = 10, xlabelpadding = -10)
    for ν in 0:4
        CairoMakie.lines!(ax, x, logbeta.(ν, x), label = latexstring("Y_{$(ν)}(x)"), linewidth = 3)
    end
    axislegend(; position = :rb, nbanks = 2, framecolor = (:grey, 0.5))
    CairoMakie.ylims!(2.5, -10)
    display(fig)
end;


X = collect(Float64, reshape(1:40, 10, 4))
x = rand(10) * 10

D = DataFrame(A=rand(10), B=collect(1:10), C=[hex(x) for x in 11:20])
D_NA = deepcopy(D)
D_NA[1, :A] = NA

@testset "Array" begin
    XX = deepcopy(X)
    mu = center!(XX, obsdim=1)
    @test sum(abs.(mean(XX, 1))) == 0
    @test all(std(XX, 1) .== std(X, 1))
    @test all(mu .== vec(mean(X, 1)))

    XX = deepcopy(X)
    mu = center!(XX, ObsDim.First())
    @test sum(abs.(mean(XX, 1))) == 0
    @test all(std(XX, 1) .== std(X, 1))
    @test all(mu .== vec(mean(X, 1)))

    XX = deepcopy(X)
    mu = center!(XX, ObsDim.Last())
    @test sum(abs.(mean(XX, 2))) == 0
    @test all(std(XX, 2) .== std(X, 2))
    @test all(mu .== vec(mean(X, 2)))

    XX = deepcopy(X)
    mu = center!(XX)
    @test sum(abs.(mean(XX, 2))) == 0
    @test all(std(XX, 2) .== std(X, 2))
    @test all(mu .== vec(mean(X, 2)))

    XX = deepcopy(X)
    mu = vec(mean(X, 1))
    center!(XX, mu, obsdim=1)
    @test sum(abs.(mean(XX, 1))) == 0
    @test all(std(XX, 1) .== std(X, 1))

    XX = deepcopy(X)
    mu = vec(mean(X, 1))
    center!(XX, mu, ObsDim.First())
    @test sum(abs.(mean(XX, 1))) == 0
    @test all(std(XX, 1) .== std(X, 1))

    XX = deepcopy(X)
    mu = vec(mean(XX, 2))
    center!(XX, mu, obsdim=2)
    @test sum(abs.(mean(XX, 2))) == 0
    @test all(std(XX, 2) .== std(X, 2))

    XX = deepcopy(X)
    mu = vec(mean(XX, 2))
    center!(XX, mu, ObsDim.Last())
    @test sum(abs.(mean(XX, 2))) == 0
    @test all(std(XX, 2) .== std(X, 2))

    XX = deepcopy(X)
    mu = vec(mean(X[:,[1,3]], 1))
    center!(XX, mu, obsdim=1, operate_on=[1, 3])
    @test sum(abs.(mean(XX[:,[1,3]], 1))) == 0
    @test all(XX[:,2] .== X[:,2])
    @test all(std(XX, 1) .== std(X, 1))

    XX = deepcopy(X)
    mu = vec(mean(X[[1,3],:], 2))
    center!(XX, mu, obsdim=2, operate_on=[1, 3])
    @test sum(abs.(mean(XX[[1,3],:], 2))) == 0
    @test all(XX[2,:] .== X[2,:])
    @test all(std(XX, 2) .== std(X, 2))
    println()

    xx = deepcopy(x)
    center!(xx)
    @test mean(xx) <= 10e-10

    xx = deepcopy(x)
    mu = mean(xx)
    center!(xx, mu)
    @test mean(xx) <= 10e-10

    xx = deepcopy(x)
    mu = ones(xx)
    center!(xx, mu)
    @test mean(xx) - mean(x) ≈ -1

    xx = deepcopy(x)
    mu = ones(xx)
    center!(xx, mu)
    @test mean(xx) - mean(x) ≈ -1
end

@testset "DataFrame" begin
    # Center DataFrame
    DD = deepcopy(D)
    center!(DD)
    @test abs.(mean(DD[:A])) <= 10e-10
    @test abs.(mean(DD[:B])) <= 10e-10
    @test all(DD[:C] .== D[:C])

    DD = deepcopy(D)
    center!(DD, operate_on=[:B])
    @test all(DD[:A] .== D[:A])
    @test abs.(mean(DD[:B])) <= 10e-10
    @test all(DD[:C] .== D[:C])

    DD = deepcopy(D)
    mu = center!(DD, operate_on=[:A, :B])
    @test abs.(mean(DD[:A])) <= 10e-10
    @test abs.(mean(DD[:B])) <= 10e-10
    @test all(DD[:C] .== D[:C])
    @test all(mu .== [mean(D[:A]), mean(D[:B])])

    DD = deepcopy(D_NA)
    mu = center!(DD, [:A, :B, :C])
    @test abs.(mean(DD[:B])) <= 10e-10
    @test all(DD[2:end, :A] .== D[2:end, :A])
    @test all(DD[:C] .== D[:C])

    DD = deepcopy(D_NA)
    mu = [0.0, mean(DD[:B]), 0.0]
    mu = center!(DD, mu, [:A, :B, :C])
    @test abs.(mean(DD[:B])) <= 10e-10
    @test all(DD[2:end, :A] .== D[2:end, :A])
    @test all(DD[:C] .== D[:C])

    DD = deepcopy(D_NA)
    mu = 0.0
    mu = center!(DD, mu, :A)
    @test all(DD[2:end, :A] .== D[2:end, :A])
    @test all(DD[:B] .== D[:B])
    @test all(DD[:C] .== D[:C])

    DD = deepcopy(D)
    mu =  [mean(D[:A]), mean(D[:B])]
    @test all(center!(DD, mu, operate_on=[:A, :B]) .== mu)
    @test abs.(mean(DD[:A])) <= 10e-10
    @test abs.(mean(DD[:B])) <= 10e-10
    @test all(DD[:C] .== D[:C])

    DD = deepcopy(D_NA)
    center!(DD)
    @test all(DD[2:end, :A] .== D[2:end, :A])
    @test abs.(mean(DD[:B])) <= 10e-10
    @test all(DD[:C] .== D[:C])
    @test isna(DD[1, :A])
end


using OrdinaryDiffEq, DiffEqBase
using DynamicalSystemsBase: DEFAULT_DIFFEQ_KWARGS, _get_solver
using Roots: find_zero, A42, AlefeldPotraShi, Brent
export poincaresos, produce_orbitdiagram

const ROOTS_ALG = A42()

#####################################################################################
#                               Poincare Section                                    #
#####################################################################################
"""
    poincaresos(ds::ContinuousDynamicalSystem, plane, tfinal = 1000.0; kwargs...)
Calculate the Poincaré surface of section (also called Poincaré map) [1, 2]
of the given system with the given `plane`.
The system is evolved for total time of `tfinal`.

If the state of the system is ``\\mathbf{u} = (u_1, \\ldots, u_D)`` then the
equation intersecting the hyperplane is
```math
a_1u_1 + \\dots + a_Du_D = \\mathbf{a}\\cdot\\mathbf{u}=b
```
where ``\\mathbf{a}, b`` are the parameters that define the hyperplane.

In code, `plane` can be either:

* A `Tuple{Int, <: Number}`, like `(j, r)` : the planecrossing is defined
  as when the `j` variable of the system crosses the value `r`.
* An `AbstractVector` of length `D+1`. The first `D` elements of the
  vector correspond to ``\\mathbf{a}`` while the last element is ``b``.

Returns a [`Dataset`](@ref) of the points that are on the surface of section.

## Keyword Arguments
* `direction = 1` : Only crossings with `sign(direction)` are considered to belong to
  the surface of section. Positive direction means going from less than ``b``
  to greater than ``b``.
* `idxs = 1:dimension(ds)` : Optionally you can choose which variables to save.
  Defaults to the entire state.
* `Ttr = 0.0` : Transient time to evolve the system before starting
  to compute the PSOS.
* `warning = true` : Throw a warning if the Poincaré section was empty.
* `rootkw = (xrtol = 1e-6, atol = 1e-6)` : A `NamedTuple` of keyword arguments
  passed to `find_zero` from [Roots.jl](https://github.com/JuliaMath/Roots.jl).
* `diffeq...` : All other extra keyword arguments are propagated into `init`
  of DifferentialEquations.jl. See [`trajectory`](@ref) for examples.

## References
[1] : H. Poincaré, *Les Methods Nouvelles de la Mécanique Celeste*,
Paris: Gauthier-Villars (1892)

[2] : M. Tabor, *Chaos and Integrability in Nonlinear Dynamics: An Introduction*,
§4.1, in pp. 118-126, New York: Wiley (1989)

See also [`orbitdiagram`](@ref), [`produce_orbitdiagram`](@ref).
"""
function poincaresos(ds::CDS{IIP, S, D}, plane, tfinal = 1000.0;
    direction = +1, Ttr::Real = 0.0, warning = true, idxs = 1:D,
    rootkw = (xrtol = 1e-6, atol = 1e-6), diffeq...) where {IIP, S, D}

    _check_plane(plane, D)
    integ = integrator(ds; diffeq...)
    planecrossing = PlaneCrossing{D}(plane, direction > 0 )
    f = (t) -> planecrossing(integ(t))

    i = typeof(idxs) <: Int ? i : SVector{length(idxs), Int}(idxs...)

    data = _initialize_output(get_state(ds), i)

    poincare_cross!(data, integ,
                     f, planecrossing, tfinal, Ttr, i, rootkw)
    warning && length(data) == 0 && @warn PSOS_ERROR

    return Dataset(data)
end

function _initialize_output(u::S, i::Int) where {S}
    output = eltype(S)[]
end
function _initialize_output(u::S, i::SVector) where {S}
    s = u[i]
    output = typeof(s)[]
end

const PSOS_ERROR =
"the Poincaré surface of section did not have any points!"

struct PlaneCrossing{D, P}
    plane::P
    dir::Bool
end
PlaneCrossing{D}(p::P, b) where {P, D} = PlaneCrossing{D, P}(p, b)
function (hp::PlaneCrossing{D, P})(u::AbstractVector) where {D, P<:Tuple}
    @inbounds x = u[hp.plane[1]] - hp.plane[2]
    hp.dir ? x : -x
end
function (hp::PlaneCrossing{D, P})(u::AbstractVector) where {D, P<:AbstractVector}
    x = zero(eltype(u))
    @inbounds for i in 1:D
        x += u[i]*hp.plane[i]
    end
    @inbounds x -= hp.plane[D+1]
    hp.dir ? x : -x
end

function poincare_cross!(data, integ,
                          f, planecrossing, tfinal, Ttr, j, rootkw)

    Ttr != 0 && step!(integ, Ttr)

    side = planecrossing(integ.u)

    while integ.t < tfinal + Ttr
        while side < 0
            integ.t > tfinal + Ttr && break
            step!(integ)
            side = planecrossing(integ.u)
        end
        while side ≥ 0
            integ.t > tfinal + Ttr && break
            step!(integ)
            side = planecrossing(integ.u)
        end
        integ.t > tfinal + Ttr && break

        # I am now guaranteed to have `t` in negative and `tprev` in positive
        tcross = find_zero(f, (integ.tprev, integ.t), ROOTS_ALG; rootkw...)
        ucross = integ(tcross)

        push!(data, ucross[j])
    end
    return data
end

function _check_plane(plane, D)
    P = typeof(plane)
    L = length(plane)
    if P <: AbstractVector
        if L != D + 1
            throw(ArgumentError(
            "The plane for the `poincaresos` must be either a 2-Tuple or a vector of "*
            "length D+1 with D the dimension of the system."
            ))
        end
    elseif P <: Tuple
        if !(P <: Tuple{Int, Number})
            throw(ArgumentError(
            "If the plane for the `poincaresos` is a 2-Tuple then "*
            "it must be subtype of `Tuple{Int, Number}`."
            ))
        end
    else
        throw(ArgumentError(
        "Unrecognized type for the `plane` argument."
        ))
    end
end

#####################################################################################
#                            Produce Orbit Diagram                                  #
#####################################################################################
"""
    produce_orbitdiagram(ds::ContinuousDynamicalSystem, plane, i::Int,
                         p_index, pvalues; kwargs...)
Produce an orbit diagram (also called bifurcation diagram)
for the `i` variable(s) of the given continuous
system by computing Poincaré surfaces of section using `plane`
for the given parameter values (see [`poincaresos`](@ref)).

`i` can be `Int` or `AbstractVector{Int}`.
If `i` is `Int`, returns a vector of vectors. Else
it returns vectors of vectors of vectors.
Each entry are the points at each parameter value.

## Keyword Arguments
* `printparams::Bool = false` : Whether to print the parameter used
  during computation in order to keep track of running time.
* `direction, warning, Ttr, rootkw, diffeq...` :
  Propagated into [`poincaresos`](@ref).
* `u0 = get_state(ds)` : Initial condition. Besides a vector you can also give
  a vector of vectors such that `length(u0) == length(pvalues)`. Then each parameter
  has a different initial condition.

## Description
For each parameter, a PSOS reduces the system from a flow to a map. This then allows
the formal computation of an "orbit diagram" for the `i` variable
of the system, just like it is done in [`orbitdiagram`](@ref).

The parameter change is done as `p[p_index] = value` taking values from `pvalues`
and thus you must use a parameter container that supports this
(either `Array`, `LMArray`, dictionary or other).

See also [`poincaresos`](@ref), [`orbitdiagram`](@ref).
"""
function produce_orbitdiagram(
    ds::CDS{IIP, S, D},
    plane,
    idxs,
    p_index,
    pvalues;
    tfinal::Real = 100.0,
    direction = +1,
    printparams = false,
    warning = true,
    Ttr = 0.0,
    u0 = get_state(ds),
    rootkw = (xrtol = 1e-6, atol = 1e-6),
    diffeq...) where {IIP, S, D}

    i = typeof(idxs) <: Int ? idxs : SVector{length(idxs), Int}(idxs...)

    _check_plane(plane, D)
    typeof(u0) <: Vector{<:AbstractVector} && @assert length(u0)==length(p)

    integ = integrator(ds; diffeq...)
    planecrossing = PlaneCrossing{D}(plane, direction > 0 )
    f = (t) -> planecrossing(integ(t))
    p0 = ds.p[p_index]

    output = [_initialize_output(ds.u0, i) for k in 1:length(pvalues)]

    for (n, p) in enumerate(pvalues)
        integ.p[p_index] = p
        printparams && println("parameter = $p")

        if typeof(u0) <: Vector{<:AbstractVector}
            st = u0[n]
        else
            st = u0
        end

        reinit!(integ, st)
        poincare_cross!(output[n], integ,
                         f, planecrossing, tfinal, Ttr, i, rootkw)

        warning && length(output[n]) == 0 && @warn "For parameter $p $PSOS_ERROR"

    end
    # Reset the parameter of the system:
    ds.p[p_index] = p0
    return output
end


# ProductSpace
#--------------
# Tensor product of several ElementarySpace objects
immutable ProductSpace{S<:ElementarySpace,N} <: TensorSpace{S,N}
    spaces::NTuple{N, S}
end

# Additional constructors
ProductSpace{S<:ElementarySpace}(V::S,Vlist::S...) = ProductSpace(tuple(V,Vlist...))
ProductSpace(P::ProductSpace) = P

# This one is needed for unambiguous treatment of zero-dimensional ProductSpaces
ProductSpace{S<:ElementarySpace,N}(::Type{S}, spaces::NTuple{N,S}) = ProductSpace{S,N}(spaces)

# Corresponding methods
dim(P::ProductSpace) = (d::Int=1;for V in P;d*=dim(V);end;return d)
iscnumber(P::ProductSpace) = all(iscnumber,P)

dim{S<:UnitaryRepresentationSpace,G<:Sector,N}(P::ProductSpace{S,N},s::NTuple{N,G})=_dim(P.spaces,s)
sectors{S<:UnitaryRepresentationSpace}(P::ProductSpace{S})=_sectors(P.spaces)
sectortype{S<:UnitaryRepresentationSpace}(P::ProductSpace{S})=sectortype(S)
sectortype{S<:UnitaryRepresentationSpace,N}(::Type{ProductSpace{S,N}})=sectortype(S)

# Convention on dual, conj, transpose and ctranspose of tensor product spaces
dual{S,N}(P::ProductSpace{S,N}) = ProductSpace{S,N}(ntuple(n->dual(P[n]),N))
Base.conj{S,N}(P::ProductSpace{S,N}) = ProductSpace{S,N}(ntuple(n->conj(P[n]),N))

Base.transpose(P::ProductSpace) = reverse(P)
Base.ctranspose(P::ProductSpace) = reverse(conj(P))

# Default construction from product of spaces:
⊗{S<:ElementarySpace}(V1::S, V2::S) = ProductSpace(S, (V1, V2))
⊗{S<:ElementarySpace}(P1::ProductSpace{S}, V2::S) = ProductSpace(S, tuple(P1.spaces..., V2))
⊗{S<:ElementarySpace}(V1::S, P2::ProductSpace{S}) = ProductSpace(S, tuple(V1, P2.spaces...))
⊗{S<:ElementarySpace}(P1::ProductSpace{S}, P2::ProductSpace{S}) = ProductSpace(S, tuple(P1.spaces..., P2.spaces...))

# Functionality for extracting and iterating over spaces
Base.length{S,N}(P::ProductSpace{S,N}) = N
Base.endof(P::ProductSpace) = length(P)
Base.getindex(P::ProductSpace, n::Integer) = P.spaces[n]
Base.getindex{S<:ElementarySpace}(P::ProductSpace{S}, r) = ProductSpace(S, P.spaces[r])

Base.reverse(P::ProductSpace)=ProductSpace(reverse(P.spaces))
Base.map(f::Base.Callable,P::ProductSpace) = map(f,P.spaces) # required to make map(dim,P) efficient

Base.start(P::ProductSpace) = start(P.spaces)
Base.next(P::ProductSpace, state) = next(P.spaces, state)
Base.done(P::ProductSpace, state) = done(P.spaces, state)

# Comparison
==(P1::ProductSpace,P2::ProductSpace) = P1.spaces==P2.spaces
==(P::ProductSpace,V::ElementarySpace) = length(P)==1 && P[1] == V
==(V::ElementarySpace,P::ProductSpace) = length(P)==1 && P[1] == V

# Show method
function Base.show(io::IO, P::ProductSpace)
    print(io,"ProductSpace(")
    for i in 1:length(P)
        i==1 || print(io," ⊗ ")
        show(io, P[i])
    end
    print(io, ")")
end

# # Promotion and conversion
# Base.convert{S<:ElementarySpace}(::Type{ProductSpace{S,1}}, V::S) = ProductSpace(V)
# Base.convert{S<:ElementarySpace}(::Type{ProductSpace{S}}, V::S) = ProductSpace(V)
# Base.convert(::Type{ProductSpace}, V::ElementarySpace) = ProductSpace(V)
#
# Base.promote_rule{S<:ElementarySpace,N}(::Type{ProductSpace{S,N}},::Type{S}) = ProductSpace{S}
# Base.promote_rule{S<:ElementarySpace}(::Type{ProductSpace{S}},::Type{S}) = ProductSpace{S}
#

# # basis and basisvector
# typealias ProductBasisVector{S,N} BasisVector{ProductSpace{S,N},Int} # use integer from 1 to dim as identifier
# typealias ProductBasis{S,N} Basis{ProductSpace{S,N}}
#
# Base.length{S,N}(B::ProductBasis{S,N}) = dim(space(B))
# Base.start{S,N}(B::ProductBasis{S,N}) = 1
# Base.next{S,N}(B::ProductBasis{S,N}, state::Int) = (ProductBasisVector{S,N}(space(B),state),state+1)
# Base.done{S,N}(B::ProductBasis{S,N}, state::Int) = state>length(B)
#
# Base.to_index{S,N}(b::ProductBasisVector{S,N}) = b.identifier # use linear indexing as long as we cannot efficiently generate a cartesian iterator
# Base.show{S,N}(io::IO,b::ProductBasisVector{S,N}) = print(io, "BasisVector($(b.space),$(ind2sub(map(dim,b.space),b.identifier)))")


#Matrices related to stochastic processes

import Base: iterate
export AiryProcess, BrownianProcess, WhiteNoiseProcess, next!

abstract type StochasticProcess{T<:Real} end

struct WhiteNoiseProcess{T<:Real} <: StochasticProcess{T}
    dt::T
end

struct BrownianProcess{T<:Real} <: StochasticProcess{T}
    dt::T
end

struct AiryProcess{S<:Real, T<:Real} <: StochasticProcess{T}
    dt::T
    beta::S
end

done(p::StochasticProcess{T}, x...) where {T} = false #Processes can always go on forever

#######################
# White noise process #
#######################

iterate(p::WhiteNoiseProcess, state=()) = (randn()*sqrt(p.dt), ())

####################
# Brownian process #
####################

function iterate(p::BrownianProcess{T}, state=zero(T)) where {T}
    newx = state + randn()*sqrt(p.dt)
    return newx, newx
end

################
# Airy process #
################

"""
Like next, but update only the state of the AiryProcess

Skip the eigenvalue computation, which gets expensive
"""
function next!(p::AiryProcess{T}, S::SymTridiagonal{T}=SymTridiagonal(T[-(2/p.dt^2)], T[])) where {T} # TODO maybe change this to iterate!
    t = (size(S, 1)-1)*p.dt

    #Discredited Airy operator plus diagonal noise
    x = inv(p.dt^2)
    push!(S.dv, -2x - t + 2/sqrt(p.dt*p.beta)*randn())
    push!(S.ev, x)

    return S
end
function iterate(p::AiryProcess{T}, S::SymTridiagonal{T}= SymTridiagonal(T[-(2/p.dt^2)], T[])) where {T}
    t = (size(S, 1)-1)*p.dt

    #Discredited Airy operator plus diagonal noise
    x = inv(p.dt^2)
    push!(S.dv, -2x - t + 2/sqrt(p.dt*p.beta)*randn())
    push!(S.ev, x)

    return (eigmax(S), S)
end


using DiffEqBase, DiffEqJump
using Test, Statistics

# using Plots; plotlyjs()
doplot = false
# using BenchmarkTools
# dobenchmark = false

dotestmean   = true
doprintmeans = false

# SSAs to test
SSAalgs = (Direct(), SortingDirect(), NRM(), RSSA(), DirectCR()) #, DirectFW(), FRM(), FRMFW()),

# numerical parameters
Nsims        = 8000
tf           = 1000.0
u0           = [1,0,0,0]
expected_avg = 5.926553750000000e+02
reltol       = .01

# average number of proteins in a simulation
function runSSAs(jump_prob)
    Psamp = zeros(Int, Nsims)
    for i in 1:Nsims
        sol = solve(jump_prob, SSAStepper())
        Psamp[i] = sol[3,end]
    end
    mean(Psamp)
end

# MODEL SETUP

# DNA repression model DiffEqBiological
# using DiffEqBiological
# rs = @reaction_network dtype begin
#     k1, DNA --> mRNA + DNA
#     k2, mRNA --> mRNA + P
#     k3, mRNA --> 0
#     k4, P --> 0
#     k5, DNA + P --> DNAR
#     k6, DNAR --> DNA + P
# end k1 k2 k3 k4 k5 k6

# model using mass action jumps
# ids: DNA=1, mRNA = 2, P = 3, DNAR = 4
reactstoch =
[
    [1 => 1],
    [2 => 1],
    [2 => 1],
    [3 => 1],
    [1 => 1, 3 => 1],
    [4 => 1]
]
netstoch =
[
    [2 => 1],
    [3 => 1],
    [2 => -1],
    [3 => -1],
    [1 => -1, 3 => -1, 4 => 1],
    [1 => 1, 3 => 1, 4 => -1]
]
spec_to_dep_jumps = [[1,5],[2,3],[4,5],[6]]
jump_to_dep_specs = [[2],[3],[2],[3],[1,3,4],[1,3,4]]
rates = [.5, (20*log(2.)/120.), (log(2.)/120.), (log(2.)/600.), .025, 1.]
majumps = MassActionJump(rates, reactstoch, netstoch)

# TESTING:
prob = DiscreteProblem(u0, (0.0, tf), rates)

# plotting one full trajectory
if doplot
    plothand = plot(reuse=false)
    for alg in SSAalgs
        jump_prob = JumpProblem(prob, alg, majumps, vartojumps_map=spec_to_dep_jumps, jumptovars_map=jump_to_dep_specs)
        sol = solve(jump_prob, SSAStepper())
        plot!(plothand, sol.t, sol[3,:], seriestype=:steppost)
    end
    display(plothand)
end

# test the means
if dotestmean
    means = zeros(Float64,length(SSAalgs))
    for (i,alg) in enumerate(SSAalgs)
        jump_prob = JumpProblem(prob, alg, majumps, save_positions=(false,false), vartojumps_map=spec_to_dep_jumps, jumptovars_map=jump_to_dep_specs)
        means[i]  = runSSAs(jump_prob)
        relerr = abs(means[i] - expected_avg) / expected_avg
        if doprintmeans
            println("Mean from method: ", typeof(alg), " is = ", means[i], ", rel err = ", relerr)
        end

        # if dobenchmark
        #     @btime (runSSAs($jump_prob);)
        # end

        @test abs(means[i] - expected_avg) < reltol*expected_avg
    end
end

# # benchmark performance
# if dobenchmark
#     # exact methods
#     for alg in SSAalgs
#         println("Solving with method: ", typeof(alg), ", using SSAStepper")
#         jump_prob = JumpProblem(prob, alg, majumps, vartojumps_map=spec_to_dep_jumps, jumptovars_map=jump_to_dep_specs)
#         @btime solve($jump_prob, SSAStepper())
#     end
#     println()
# end


using Test
using IndividualDisplacements, MeshArrays, DifferentialEquations

@testset "MeshArrays tests:" begin

mygrid=gcmgrid("flt_example/","ll",1,[(80,42)], [80 42], Float32, read, write)
XC=MeshArray(mygrid,Float32); XC[1]=vec(2500.:5000.:397500.0)*ones(1,42);
XG=MeshArray(mygrid,Float32); XG[1]=vec(0.:5000.:395000.0)*ones(1,42);
YC=MeshArray(mygrid,Float32); YC[1]=ones(80,1)*transpose(vec(2500.:5000.:207500.0));
YG=MeshArray(mygrid,Float32); YG[1]=ones(80,1)*transpose(vec(0.:5000.:205000.0));
GridVariables=Dict("XC" => XC,"YC" => YC,"XG" => XG,"YG" => YG,"dx" => 5000.0);

t0=0.0; t1=18001.0*3600.0
u0=-(YG.-YC[1][40,21])/2000000.; u1=u0
v0=(XG.-XC[1][40,21])/2000000.; v1=v0

u0=u0./GridVariables["dx"]
u1=u1./GridVariables["dx"]
v0=v0./GridVariables["dx"]
v1=v1./GridVariables["dx"]

uvt = Dict("u0" => u0, "u1" => u1, "v0" => v0, "v1" => v1, "t0" => t0, "t1" => t1)
uvetc=merge(uvt,GridVariables);

uInit=[200000.0;0.0]./uvetc["dx"]
nSteps=3000-2
du=fill(0.0,2);
tspan = (0.0,nSteps*3600.0)
prob = ODEProblem(IndividualDisplacements.VelComp,uInit,tspan,uvetc)
sol = solve(prob,Tsit5(),reltol=1e-8,abstol=1e-8)

@test isapprox(sol[1,end],117237.0./uvetc["dx"]; atol=100.)
@test isapprox(sol[2,end],40448.0./uvetc["dx"]; atol=100.)

end


# """
#    abstract type Source has two subtype: SingleSource and SimultaneousSource
# """
# abstract type Source end

"""
Defines an explosive source for time domain finite difference modelling.

# Fields

- `isz::Ti`: vertical index location
- `isx::Ti`: horizontal index location
- `src2spt::Ti` : index mapping source location to snapshot
- `src2wfd::Ti` : index mapping source location to wavefield
- `dt::Tv`      : time sampling
- `it_min::Ti`  : source lower time index
- `it_max::Ti`  : source upper time index
- `p::Vector{Tv}` :  Vector of Source wavelets
- `hello`
"""
struct Source{Ti<:Int64, Tv<:AbstractFloat}
    isz     :: Ti  # vertical index location
    isx     :: Ti  # horizontal index location
    src2spt :: Ti  # index mapping to snapshot
    src2wfd :: Ti  # index mapping to wavefield
    dt      :: Tv
    it_min  :: Ti  # the index lower bound in time axis
    it_max  :: Ti  # the index upper bound in time axis
    p       :: Vector{Tv} # source wavelet
end

"""
   create a ricker wavelet
"""
function ricker(fdom, dt)

 	  nw = 2.2 / fdom / dt         # decide number of samples
	  nw = 2*floor(Int64, nw/2)+1  # gaurantee the number of samples is odd

	  nc = floor(Int64, nw/2)
	  w  = zeros(nw)               # allocate space for wavelet
	  k  = collect(1:nw)

    for i = 1 : nw
        alpha = (nc-i+1) * fdom * dt * pi
        beta  = alpha^2
        w[i]  = (1 - beta * 2) * exp(-beta)
    end
	  return w
end

"""
  convert wavelet to its corresponding minimum phase version
"""
function convert2miniphase(w::Vector{Tv}) where {Tv<:AbstractFloat}

    nw = length(w)

    # number of samples
    nf = 8*nextpow(2, nw)

    W = fft(vcat(w,zeros(nf-nw)))
    A = log.(abs.(W) .+ 1.e-8)
    a = 2*ifft(A)
    n2 = floor(Int, nf/2)
    a[n2+2:nf] .= 0.0
    a[1] = a[1]/2
    A = exp.(fft(a))
    a = real(ifft(A))
    wmin = real(a[1:nw])

    return wmin
end

"""
   fc is the high cut frequency of low-pass filter, hl is half length of filter
and dt is sampling rate
"""
function tapered_sinc(fc, hl, dt) where {Tv <: AbstractFloat}

    # length of filter
    l = 2*hl
    filter = zeros(l+1)
    ff = fc / (1.0/dt)

    for i = 0 : l
        idx = i - hl
        if idx == 0
           filter[i+1] = 2*pi*ff
        else
           filter[i+1] = sin(2*pi*ff*idx)/idx
        end
        taper = 0.42 - 0.5 * cos(2*pi*i/l) + 0.08*cos(4*pi*i/l)
        filter[i+1] = filter[i+1] * taper
    end

    # normalize filter
    filter = filter / sum(filter)
    return filter
end

"""
   scaled sinc function without tapering, so the decaying of sinc function is slow
"""
function scaled_sinc(fc, hl, dt::Tv) where {Tv<:AbstractFloat}

    l = 2*hl
    filter = zeros(l+1)
    # high cut frequency fraction
    ff = fc / (1.0/dt)

    for i = 0 : l
        idx = i - hl
        if idx == 0
           filter[i+1] = 2*pi*ff
        else
           filter[i+1] = sin(2*pi*ff*idx)/idx
        end
    end

    filter = filter / sum(filter)
    return filter
end

"""
   S=Source(sz, sx, params; <keyword arguments>)

Constructs an explosive Source structure.

Please see the extended help with `?? Source`

# Arguments

-`sz::`: source coordinate
-`sx::`: soirce cordinate
-`params::TdParams`: Parameters for time domain modelling

# Keyword Arguments

-`location_flag`: Source(s) positioning by `"index"` or `"distance"`
-`ot`: Source(s) firing time
-`amp`: Source(s) amplitude
-`fdom`: Dominant frequency of wavelet
-`type_flag`: Type of source wavelet.
-`hl`:Window lenght for sinc interpolation (for type_flag="sinc")
-`p`: Either a user provided source wavelet or a vector of source wavelets.

# Extended help

   S=Source(sz, sx, params; location_flag="index", ot=0.0, amp=1.0, fdom=20.0, hl=128, type_flag="ricker", p::Vector{Float32}=p)

The structure defines the source location by either its position or relative index in the model through the `location_flag`.

If the source wavelet (or a vector of source wavelets) `p` is not provieded, the user can specify it with "type_flag".

-Default Ricker wavelet: `type_flag="ricker"`
-Sinc wavelet: `type_flag="sinc"`
-Minimum phase wavelet: `type_flag=miniphase`
"""
function Source(sz, sx, params::TdParams; location_flag="index", ot=0.0, amp=1.0, fdom=20.0, hl=128,
                type_flag="ricker", p=Vector{Float32}(undef,0)) where {Tv <: Real}

    # source location given index
    if location_flag == "index"
       isz = round(Int64, sz)
       isx = round(Int64, sx)

    # source location given as distance
    elseif location_flag == "distance"
       isz = round(Int64, sz/params.dz) + 1
       isx = round(Int64, sx/params.dx) + 1
    else
       error("wrong specification of source location")
    end

    # error checking
    if isz > params.nz || isx > params.nx || isz < 1 || isx < 1
       error("source located outside of modeling area")
    end

    # can't put on the free-surface
    if params.free_surface && isz == 1
       error("can't inject source on the surface")
    end

    # the auxillary index mapping the location of source to snapshot or wavefield
    src2spt = (isx+params.npml-1) * params.Nz + isz + params.ntop
    src2wfd = (isx            -1) * params.nz + isz

    # non-user provided source wavelet
    if length(p) == 0
       if type_flag == "ricker"
          p = ricker(fdom, params.dt)
          p = amp * p
       elseif type_flag == "miniphase"
          p = ricker(fdom, params.dt)
          p = convert2miniphase(p)
          p = amp * p
       elseif type_flag == "sinc"
          fc = params.data_format(fdom)
          hl = round(Int64, hl)
          p  = tapered_sinc(fc, hl, params.dt)
          p  = amp * p
       end
    end

    # make the data type consistent
    p = convert(Vector{params.data_format}, p)
    nt= length(p)

    # the index time range of source wavelet
    it_min = floor(Int64, ot/params.dt) + 1
    it_max = it_min + nt - 1

    # call the default constructor
    return Source(isz, isx, src2spt, src2wfd, params.dt, it_min, it_max, p)
end

"""
   source_array= get_multi_sources(sz::Vector, sx::Vector, params::TdParams; <keyword arguments>)

Get an array of sources with different positions.

Every element of the source array contains an initialized Source structure.

# Arguments:

-`sz::Vector`: source z coordinates.
-`sx::Vector`: source x coordinates.
-`params::TdParams`: Time domain parameters for modeling.

# Extended help

This function extends the function constructor `Source` to a vector of source.

-`location_flag::String`: Define the type of location for each source.
-`ot::Vector`: Vector of firing times for each source.
-`amp::Float64`: Amplitude for the array of sources.
-`fdom::Float64`: Dominant frequency the array of sources.
-`hl::`:
-`typef_flag::String`: Type of source (default: "ricker")
-`p::Vector{Float32}`: Vector of wavelets for the array of soruces.
"""
function get_multi_sources(sz::Vector, sx::Vector, params::TdParams; location_flag="index", ot=0.0, amp=1.0,
                          fdom=20.0, hl=128, type_flag="ricker", p=Vector{Float32}(undef,0))

    # number of source
    ns = length(sz)
    if length(sx) != ns
       error("length of source coordinate mismatch")
    end

    # source location given index
    isz = zeros(Int64, ns)
    isx = zeros(Int64, ns)
    if location_flag == "index"
       for i = 1 : ns
           isz[i] = round(Int64, sz[i])
           isx[i] = round(Int64, sx[i])
       end

    # source location given as distance
    elseif location_flag == "distance"
       isz[i] = round(Int64, sz[i]/params.dz) + 1
       isx[i] = round(Int64, sx[i]/params.dx) + 1
    else
       error("wrong specification of source location")
    end

    # error checking
    for i = 1 : ns
        if isz[i] > params.nz || isx[i] > params.nx
           error("source located outside of modeling area")
        end
    end

    # starting time
    if typeof(ot) <: Real  # all sources starting at same time
       ot = ot * ones(ns)
    end
    if length(ot) != ns
       error("length of starting time ot mismatch")
    end

    # amplitude for each source
    if typeof(amp) <: Real
       amp = amp * ones(ns)
    end
    if length(amp) != ns
       error("length of amplitude is wrong")
    end

    # dominant frequency for source wavelet
    if typeof(fdom) <: Real
       fdom = fdom * ones(ns)
    end
    if length(fdom) != ns
       error("length of dominant frequency is wrong")
    end

    wavelet = Vector{Vector{params.data_format}}(undef, ns)
    if length(p) == 0  # without user provided wavelet
       if typeof(type_flag) == String               # all sources share same wavelet
          wavelet_type = Vector{String}(undef, ns)
          wavelet_type[:] .= type_flag
          type_flag    = wavelet_type
       elseif typeof(type_flag) == Vector{String}   # use different wavelet
          if length(type_flag) != ns
             error("length of type_flag mismatch")
          end
       end

       # generate the specific source wavelet
       for i = 1 : ns
           if type_flag[i] == "ricker"
              tmp = amp[i] * ricker(fdom[i], params.dt)
              wavelet[i] = convert(Vector{params.data_format}, tmp)
           elseif type_flag[i] == "miniphase"
              tmp = ricker(fdom[i], params.dt)
              tmp = amp[i] * convert2miniphase(tmp)
              wavelet[i] = convert(Vector{params.data_format}, tmp)
           elseif type_flag[i] == "sinc"
              hl   = round(Int64, hl)
              tmp  = amp[i] * tapered_sinc(fdom[i], hl, params.dt)
              wavelet[i] = convert(Vector{params.data_format}, tmp)
           end
       end

    elseif eltype(p) <: AbstractFloat   # all sources share same user-given wavelet
       for i = 1 : ns
           wavelet[i] = convert(Vector{params.data_format}, p)
       end

    else                                # sources have different user-given wavelet
       for i = 1 : ns
           wavelet[i] = convert(Vector{params.data_format}, p[i])
       end
    end

    # allocate memory for vector of sources
    srcs = Vector{Source}(undef, ns)
    for i = 1 : ns
        srcs[i] = Source(isz[i], isx[i], params; p=wavelet[i], ot=ot[i])
    end
    return srcs
end

"""
   Add source to snapshot (we add source to pressure terms, so the wavelet is caled by dt)
"""
function add_source!(spt::Snapshot, src::Source, it::Int64)

    # add source
    if src.it_min <= it <= src.it_max

       idx_t = it - src.it_min + 1
       idx_p = src.src2spt

       # equvialent to adding to pressure field
       spt.px[idx_p] = spt.px[idx_p] + src.p[idx_t] * src.dt
    end

    return nothing
end

"""
   add source to wavefield, used for forward wavefield reconstruction
"""
function add_source!(wfd::Wavefield, src::Source, it::Int64)

    if src.it_min <= it <= src.it_max

       idx_t = it - src.it_min + 1
       idx_p = src.src2wfd

       wfd.p[idx_p] = wfd.p[idx_p] + src.p[idx_t] * src.dt
    end
    return nothing
end

"""
   add multi-sources simultaneously to snapshot
"""
function add_source!(spt::Snapshot, srcs::Vector{Source}, it::Int64)

    # number of sources
    ns = length(srcs)

    # loop over each source
    for i = 1 : ns
        add_source!(spt, srcs[i], it)
    end
    return nothing
end

"""
   add multi-sources simultaneously to snapshot
"""
function add_source!(wfd::Wavefield, srcs::Vector{Source}, it::Int64)

    # number of sources
    ns = length(srcs)

    # loop over each source
    for i = 1 : ns
        add_source!(wfd, srcs[i], it)
    end
    return nothing
end

"""
   subtract source from wavefield, used for backward wavefield reconstruction
"""
function subtract_source!(wfd::Wavefield, src::Source, it::Int64)

    if src.it_min <= it <= src.it_max
       idx_t = it - src.it_min + 1
       idx_p = src.src2wfd

       wfd.p[idx_p] = wfd.p[idx_p] - src.p[idx_t] * src.dt
    end
    return nothing
end

"""
   subtract multiple sources from wavefield, used for backward wavefield reconstruction
"""
function subtract_source!(wfd::Wavefield, srcs::Vector{Source}, it::Int64)

    # number of sources
    ns = length(srcs)

    # loop over sources
    for i = 1 : ns
        if srcs[i].it_min <= it <= srcs[i].it_max
           idx_t = it - srcs[i].it_min + 1
           idx_p = srcs[i].src2wfd

           wfd.p[idx_p] = wfd.p[idx_p] - srcs[i].p[idx_t] * srcs[i].dt
        end
    end

    return nothing
end

"""
   find the time range of multiple sources
"""
function time_range_source(srcs::Vector{Source})

    tmax = 0.0
    tmin = 0.0
    ns   = length(srcs)

    for i = 1 : ns
        tmp1 = (srcs[i].it_min-1) * srcs[i].dt
        tmp2 = (srcs[i].it_max-1) * srcs[i].dt

        if tmp1 < tmin
           tmin = tmp1
        end
        if tmp2 > tmax
           tmax = tmp2
        end
    end

    return tmin, tmax
end

function write_source(path_src::String, src::Source)

    fid = open(path_src, "w")

    write(fid, src.isz)
    write(fid, src.isx)
    write(fid, src.src2spt)
    write(fid, src.src2wfd)
    write(fid, src.it_min)
    write(fid, src.it_max)

    typeof(src.dt) == Float32 && write(fid, 1)
    typeof(src.dt) == Float64 && write(fid, 2)

    write(fid, src.dt)
    write(fid, src.p )

    close(fid)
    return nothing
end

function read_source(path_src::String)

    fid     = open(path_src, "r")
    isz     = read(fid, Int64)
    isx     = read(fid, Int64)
    src2spt = read(fid, Int64)
    src2wfd = read(fid, Int64)
    it_min  = read(fid, Int64)
    it_max  = read(fid, Int64)
    nt      = it_max - it_min + 1
    format_code = read(fid, Int64)

    if format_code == 1
       dt = read(fid, Float32)
       p  = zeros(Float32, nt)
       read!(fid, p)

    elseif format_code == 2
       dt = read(fid, Float64)
       p  = zeros(Float64, nt)
       read!(fid, p)

    else
       error("wrong data format")
    end

    return Source(isz, isx, src2spt, src2wfd, dt, it_min, it_max, p)
end

function source_isequal(src1::Source, src2::Source)

    src1.isz     != src2.isz && (return false)
    src1.isx     != src2.isx && (return false)
    src1.src2spt != src2.src2spt && (return false)
    src1.src2wfd != src2.src2wfd && (return false)
    src1.it_min  != src2.it_min  && (return false)
    src1.it_max  != src2.it_max  && (return false)
    src1.dt      != src2.dt      && (return false)
    norm(src1.p-src2.p)/norm(src1.p) > 10*eps(eltype(src1.p)) && (return false)

    return true
end

# ========================need to be updated ===================================
# """
#    randomly source starting time shift and make the starting time is the integer
# times the time interval.
# """
# function random_start_time(ns::Ti, dt::Tv, s::Tv) where {Ti<:Int64, Tv<:AbstractFloat}
#
#     ot = s * rand(ns)
#     indt = round(Int64, ot/dt)
#     ot[:] = indt*dt
#
#     return convert(Vector{typeof(dt)}, ot)
# end

# """
#    generate a snapshot via multiple sources
# """
# function srcs2spt(srcs::Array{Source,1}, it::Int64)
#
#     nz = srcs[1].nz ; nx = srcs[1].nx;
#     ext= srcs[1].ext; iflag=srcs[1].iflag;
#     dt = srcs[1].dt
#     spt = InitSnapShot(nz, nx, ext, iflag, dt, it)
#
#     AddMultiSources!(spt, srcs)
#     return spt
# end
#
# """
#    generate wavefield via multiple sources and write the snapshots to hard drive
# """
# function WriteSrcs2SnapShots(path::String , srcs::Array{Source,1})
#     nz = srcs[1].nz
#     (tl, tu) = SrcRange(srcs)
#     nt  = round(Int64, (tu-tl)/dt) + 1
#     fid = open(path, "w")
#     write(fid, nz, nx, ext, iflag, dt)
#     for it = 1 : nt
#         it1 = round(Int64, (tl+(it-1)*dt)/dt+1)
#         spt = srcs2spt(srcs, it1)
#         writeSnapShot(fid, spt)
#     end
#     close(fid)
#     return nothing
# end
#
# """
#    extract a wavelet from the stress field
# """
# function multiStress2wlet(path::String , iz::Int64, ix::Int64)
#     fid = open(path, "r")
#     (nz, nx, ext, iflag, dt, nt) = InfoStress(path)
#     w = zeros(nt)
#     for it = 1 : nt
#         sts= readStress(path, it)
#         w[it] = sts.p[iz,ix]
#     end
#     return w
# end


function test_l1_regularizer(backend::Backend, T, eps)
  println("-- Testing L1 regularizer on $(typeof(backend)){$T}...")

  coef = rand()
  param = rand(T, 2,3,4,5) - convert(T, 0.5)
  param_blob = make_blob(backend, param)
  regu = L1Regu(coef)

  loss = forward(backend, regu, 1.0, param_blob)
  expected_loss = coef * sum(abs.(param))
  @test -eps < loss - expected_loss < eps

  grad_blob = make_zero_blob(backend, T, size(param))
  backward(backend, regu, 1.0, param_blob, grad_blob)
  grad = zeros(T, size(param))
  copy!(grad, grad_blob)

  @test all(-eps .< grad - coef*sign.(param) .< eps)
end

function test_l1_regularizer(backend::Backend)
  test_l1_regularizer(backend, Float32, 1e-3)
  test_l1_regularizer(backend, Float64, 1e-4)
end

if test_cpu
  test_l1_regularizer(backend_cpu)
end
if test_gpu
  test_l1_regularizer(backend_gpu)
end


export
    AbstractDiophantineSolution,
    AbstractDiophantineSolutions,
    AbstractDiophantineSolutionsXY,
    DiophantineSolutionXY,
    DiophantineSolutions,
    DiophantineSolutionsNoneXNoneY,
    DiophantineSolutionsAnyXAnyY,
    DiophantineSolutionsOneXAnyY,
    DiophantineSolutionsAnyXOneY,
    DiophantineSolutionsLinearXLinearY,
    DiophantineSolutionsQuadraticXQuadraticY,
    diophantine_solution,
    diophantine_solutions,
    diophantine_nonex_noney,
    diophantine_anyx_anyy,
    diophantine_onex_anyy,
    diophantine_anyx_oney,
    diophantine_linearx_lineary,
    diophantine_quadraticx_quadraticy

abstract type AbstractDiophantineSolution{T<:Integer} end

@auto_hash_equals struct DiophantineSolutionXY{T<:Integer} <: AbstractDiophantineSolution{T}
    x::T
    y::T
end

abstract type AbstractDiophantineSolutions{T<:AbstractDiophantineSolution} end

@auto_hash_equals struct DiophantineSolutions{T<:AbstractDiophantineSolution} <: AbstractDiophantineSolutions{T}
    solutions::Array{T,1}
end

abstract type AbstractDiophantineSolutionsXY{T<:Integer} <: AbstractDiophantineSolutions{DiophantineSolutionXY{T}} end

struct DiophantineSolutionsNoneXNoneY{T<:Integer} <: AbstractDiophantineSolutionsXY{T}
end

struct DiophantineSolutionsAnyXAnyY{T<:Integer} <: AbstractDiophantineSolutionsXY{T}
end

@auto_hash_equals struct DiophantineSolutionsOneXAnyY{T<:Integer} <: AbstractDiophantineSolutionsXY{T}
    x::T
end

@auto_hash_equals struct DiophantineSolutionsAnyXOneY{T<:Integer} <: AbstractDiophantineSolutionsXY{T}
    y::T
end

@auto_hash_equals struct DiophantineSolutionsLinearXLinearY{T<:Integer} <: AbstractDiophantineSolutionsXY{T}
    mx::T
    nx::T
    my::T
    ny::T
end

@auto_hash_equals struct DiophantineSolutionsQuadraticXQuadraticY{T<:Integer} <: AbstractDiophantineSolutionsXY{T}
    ax::T
    bx::T
    cx::T
    ay::T
    by::T
    cy::T
end

const SolutionXY            = DiophantineSolutionXY
const Solutions             = DiophantineSolutions
const NoneX_NoneY           = DiophantineSolutionsNoneXNoneY
const AnyX_AnyY             = DiophantineSolutionsAnyXAnyY
const OneX_AnyY             = DiophantineSolutionsOneXAnyY
const AnyX_OneY             = DiophantineSolutionsAnyXOneY
const LinearXLinearY        = DiophantineSolutionsLinearXLinearY
const QuadraticXQuadraticY  = DiophantineSolutionsQuadraticXQuadraticY

diophantine_solution(x::T, y::T) where {T<:Integer} = SolutionXY(x, y)

diophantine_solutions(s::Array{T,1}) where {T<:AbstractDiophantineSolution} = Solutions{T}(s)
diophantine_solutions(s::Tuple{T,T}...) where {T<:Integer} = diophantine_solutions([diophantine_solution(x, y) for (x, y) in s])
diophantine_solutions(s::AbstractArray{Tuple{T,T},1}) where {T<:Integer} = diophantine_solutions(s...)

diophantine_nonex_noney(T::Type = Int) = NoneX_NoneY{T}()
diophantine_anyx_anyy(T::Type = Int) = AnyX_AnyY{T}()
diophantine_onex_anyy(x::T) where {T<:Integer} = OneX_AnyY{T}(x)
diophantine_anyx_oney(y::T) where {T<:Integer} = AnyX_OneY{T}(y)
diophantine_linearx_lineary(mx::T, nx::T, my::T, ny::T) where {T<:Integer} = LinearXLinearY{T}(mx, nx, my, ny)
diophantine_quadraticx_quadraticy(ax::T, bx::T, cx::T, ay::T, by::T, cy::T) where {T<:Integer} = QuadraticXQuadraticY{T}(ax, bx, cx, ay, by, cy)

Base.show(io::IO, sol::SolutionXY)              = print(io, "x=$(sol.x) ∧ y=$(sol.y)")
Base.show(io::IO, sol::Solutions)               = print(io, sol.solutions)    # writemime(io, MIME("text/plain"), sol.solutions)
Base.show(io::IO, ::NoneX_NoneY)                = print(io, "∅")
Base.show(io::IO, ::AnyX_AnyY)                  = print(io, "∀t,s∈ℤ: x=t ∧ y=s")
Base.show(io::IO, sol::OneX_AnyY)               = print(io, "∀t∈ℤ: x=$(sol.x) ∧ y=t")
Base.show(io::IO, sol::AnyX_OneY)               = print(io, "∀t∈ℤ: x=t ∧ y=$(sol.y)")
Base.show(io::IO, sol::LinearXLinearY)          = print(io, "∀t∈ℤ: x=$(formatlinear(sol.mx, sol.nx)) ∧ y=$(formatlinear(sol.my, sol.ny))")
Base.show(io::IO, sol::QuadraticXQuadraticY)    = print(io, "∀t∈ℤ: x=$(formatquadratic(sol.ax, sol.bx, sol.cx)) ∧ y=$(formatquadratic(sol.ay, sol.by, sol.cy))")

formatlinear(m::T, n::T) where {T<:Integer} = begin
    nstr = string(n)
    m == 0 && return nstr
    n > 0 && (nstr = "+" * nstr)

    mstr = abs(m) == 1 ? (signbit(m) ? "-" : "") * "t" : "$(m)t"
    n == 0 && return mstr

    mstr * nstr
end

formatquadratic(a::T, b::T, c::T) where {T<:Integer} = begin
    linstr = formatlinear(b, c)
    a == 0 && return linstr
    ((b > 0) || (b == 0 && c > 0)) && (linstr = "+" * linstr)

    astr = abs(a) == 1 ? (signbit(a) ? "-" : "") * "t²" : "$(a)t²"
    b == 0 && c == 0 && return astr

    astr * linstr
end

Base.eltype(it::AbstractDiophantineSolutions) = eltype(typeof(it))
Base.eltype(::Type{Solutions{T}}) where {T} = T
Base.eltype(::Type{NoneX_NoneY{T}}) where {T<:Integer} = DiophantineSolutionXY{T}
Base.eltype(::Type{AnyX_AnyY{T}}) where {T<:Integer} = DiophantineSolutionXY{T}
Base.eltype(::Type{OneX_AnyY{T}}) where {T<:Integer} = DiophantineSolutionXY{T}
Base.eltype(::Type{AnyX_OneY{T}}) where {T<:Integer} = DiophantineSolutionXY{T}
Base.eltype(::Type{LinearXLinearY{T}}) where {T<:Integer} = DiophantineSolutionXY{T}
Base.eltype(::Type{QuadraticXQuadraticY{T}}) where {T<:Integer} = DiophantineSolutionXY{T}

Base.start(it::Solutions) = start(it.solutions)
Base.next(it::Solutions, state) = next(it.solutions, state)
Base.done(it::Solutions, state) = done(it.solutions, state)
Base.length(it::Solutions) = length(it.solutions)

Base.start(::NoneX_NoneY) = Void
Base.next(::NoneX_NoneY, _) = Void
Base.done(::NoneX_NoneY, _) = true
Base.length(::NoneX_NoneY) = 0

Base.start(it::OneX_AnyY) = zero(it.x)
Base.next(it::OneX_AnyY, state) = begin
    s = diophantine_solution(it.x, state)
    s, nextstate(state)
end
Base.done(::OneX_AnyY, _) = false
Base.iteratorsize(::OneX_AnyY) = Base.IsInfinite()

Base.start(it::AnyX_OneY) = zero(it.y)
Base.next(it::AnyX_OneY, state) = begin
    s = diophantine_solution(state, it.y)
    s, nextstate(state)
end
Base.done(::AnyX_OneY, _) = false
Base.iteratorsize(::AnyX_OneY) = Base.IsInfinite()

Base.start(it::LinearXLinearY) = zero(it.mx)
Base.next(it::LinearXLinearY, state) = begin
    s = diophantine_solution(it.mx * state + it.nx, it.my * state + it.ny)
    s, nextstate(state)
end
Base.done(::LinearXLinearY, _) = false
Base.iteratorsize(::LinearXLinearY) = Base.IsInfinite()

Base.start(it::QuadraticXQuadraticY) = zero(it.ax)
Base.next(it::QuadraticXQuadraticY, state) = begin
    s = diophantine_solution(it.ax * state^2 + it.bx * state + it.cx, it.ay * state^2 + it.by * state + it.cy)
    s, nextstate(state)
end
Base.done(::QuadraticXQuadraticY, _) = false
Base.iteratorsize(::QuadraticXQuadraticY) = Base.IsInfinite()

nextstate(state::Integer) = state <= zero(state) ? abs(state) + one(state) : -state


module TestDual

using Compat
using Compat.Test
using IntervalArithmetic
using EAGODomainReduction

Y = Interval{Float64}[Interval(1.0,4.0),Interval(1.0,4.0),Interval(1.0,4.0),Interval(1.0,4.0)]
ymult_lo = [50, 1.0, 2.0, 3.0]
ymult_hi = [0, 1.0, 2.0, 3.0]
yLBD = 1.0
yUBD = 3.0
Variable_DR!(Y,ymult_lo,ymult_hi,yLBD,yUBD)
println("Y post full duality-based bound tighteing: ", Y)
@test 3.95999-1E-4 <= Y[1].lo <= 3.95999+1E-4
@test 4.0-1E-4 <= Y[1].hi <= 4.0+1E-4
@test 2.0-1E-4 <= Y[2].lo <= 2.0+1E-4
@test 4.0-1E-4 <= Y[2].hi <= 4.0+1E-4
@test 3.0-1E-4 <= Y[3].lo <= 3.0+1E-4
@test 4.0-1E-4 <= Y[3].hi <= 4.0+1E-4
@test 3.33333-1E-4 <= Y[4].lo <= 3.33333+1E-4
@test 4.0-1E-4 <= Y[4].hi <= 4.0+1E-4

X = Interval{Float64}[Interval(1.0,4.0),Interval(1.0,4.0),Interval(1.0,4.0),Interval(1.0,4.0)]
mult_lo = [50, 1.0]
mult_hi = [-1.0, 0]
LBD = 1.0
UBD = 2.0
nx = 2
Variable_DR_Imp!(X,mult_lo,mult_hi,LBD,UBD,nx)
println("X post implicit duality-based bound tighteing: ", X)
@test 1.0-1E-4 <= X[1].lo <= 1.0+1E-4
@test 4.0-1E-4 <= X[1].hi <= 4.0+1E-4
@test 1.0-1E-4 <= X[2].lo <= 1.0+1E-4
@test 4.0-1E-4 <= X[2].hi <= 4.0+1E-4
@test 3.97999-1E-4 <= X[3].lo <= 3.97999+1E-4
@test 4.0-1E-4 <= X[3].hi <= 4.0+1E-4
@test 3.0-1E-4 <= X[4].lo <= 3.0+1E-4
@test 4.0-1E-4 <= X[4].hi <= 4.0+1E-4

end


#Struct to hold information on wavelet-grid (u-component).
struct UWavelets{T <: Real, N <: Integer}
          nxuw :: N                                       # Number of grid points in x in UWavelets (equal to UGrid)
          nyuw :: N                                       # Number of grid points in y in UWavelets (equal to UGrid)
        mask :: Array{Bool,2}                           # Model domain on the U grid
           n :: Base.RefValue{N}                        # Number of grid points in domain
        crop :: Base.RefValue{Diagonal{T,Array{T,1}}}   # Crop matrix: diagonal matrix with mask entries on diag
        samp :: Base.RefValue{SparseMatrixCSC{T,N}}     # Sampling matrix: take full domain to model domain
      spread :: Base.RefValue{SparseMatrixCSC{T,N}}     # Spread matrix: take model domain to full domain
      levels :: N                                       # Number of wavelet levels
        idwt :: KronType{T,N}                           # Wavelet matrix cross produce
    wavelets :: Array{T,2}                              # Wavelet matrix
end

"""
    UWavelets(;
            nxuw,
            nyuw,
            mask = trues(nxuw,nyuw),
            levels)

Construct a WAVI.jl UWavelets structure with size (nxuw,nyuw)
UWavelets stores wavelet matrix and associated fields for velocities in x direction

Keyword arguments
=================
    - 'nxuw': (required) Number of grid cells in x-direction in UWavelts
    - 'nyuw': (required) Number of grid cells in y-direction in UWavelets
    - 'mask': Mask specifying the model domain with respect to UWavelet grid (UGrid)
    - levels: (required) Number of levels in the preconditioner
"""
function UWavelets(;
                    nxuw,
                    nyuw,
                    mask = trues(nxuw,nyuw),
                    levels)

    (size(mask) == (nxuw,nyuw)) || throw(DimensionMismatch("Sizes of inputs to UWavelets must all be equal to nxuw x nyuw (i.e. $nxuw x $nyuw)"))

    #compute non-inputs
    n = Ref(count(mask));  @assert n[] == count(mask)
    crop = Ref(Diagonal(float(mask[:]))); @assert crop[] == Diagonal(float(mask[:]));
    samp  = Ref(crop[][mask[:],:]); @assert samp[] == crop[][mask[:],:]
    spread = Ref(sparse(samp[]')); @assert spread[] == sparse(samp[]')
    idwt =  wavelet_matrix(nyuw,levels,"reverse" ) ⊗ wavelet_matrix(nxuw,levels,"reverse")
    wavelets = zeros(nxuw,nyuw); @assert size(wavelets)==(nxuw,nyuw)

    #make sure boolean type rather than bitarray
    mask = convert(Array{Bool,2}, mask)

    return UWavelets(
                    nxuw,
                    nyuw,
                    mask,
                    n,
                    crop,
                    samp,
                    spread,
                    levels,
                    idwt,
                    wavelets)
end


@testset "mSGNHT" begin
    #-------------------
    # test inplace addition vs reference
    #-------------------
    function sample_reference!(s::mSGNHTState,grad::Function)
        # increment iterations counter
        s.t += 1

        # improve readability
        nparams = length(s.x)
        stepsize = s.stepsize(s.t)
        x = s.x
        p = s.p
        xi = s.xi
        D = s.D
        mu = s.mu

        # second order symmetric splitting scheme from Leimkuhler et al.
        p[:] += .5*stepsize*grad(x)
        x[:] += .5*stepsize*p
        xi[:] +=.5*stepsize/mu* (p.*p - 1)
        ids = abs(stepsize*xi) .> 1e-6
        p[ids] = exp(-stepsize*xi[ids]).*p[ids] + sqrt(.5.*(1-exp(-2.0*stepsize*xi[ids]))./xi[ids] * D) .* randn(sum(ids))
        p[!ids] += sqrt(stepsize*D)*randn(sum(!ids))
        xi[:] +=.5*stepsize/mu* (p.*p - 1)
        x[:] += .5*stepsize*p
        p[:] += .5*stepsize*grad(x)
        s
    end

    srand(1)
    s = mSGNHTState(ones(2),0.01)
    srand(1)
    ref_s = mSGNHTState(ones(2),0.01)
    msgnhtgrad(x) = x.*x

    srand(2)
    sample!(s,msgnhtgrad)

    srand(2)
    sample_reference!(ref_s,msgnhtgrad)

    @test s.x == ref_s.x

    #------------------
    # Gaussian sanity check.
    #------------------
    dm = GaussianMixtureModel([0.,],[1.],[1.])

    gradient = DataModel.getgrad(dm)
    s = mSGNHTState(zeros(1),0.1)

    x_avg = [0.]
    x2_avg = [0.]
    for i in 1:50000
        sample!(s,gradient)
        x_avg =(i-1)/i*x_avg + s.x/i
        x2_avg = (i-1)/i*x2_avg + s.x.^2/i
    end

    @test abs(x_avg-0)[1] < 0.1
    @test abs(x2_avg-1)[1] < 0.2
end


struct Wave

    amps::Vector
    time
    framerate::Number

    signal::Signal
    #isnoise::Bool

end

function Wave(s::Signal; t = 0, framerate = 44_100)

    if t == 0
        freq = s.freq
        if !(s.freq |> typeof <: Number)
            freq = minimum(s.freq)
        end

        t = inv(freq) * 5

    end

    dt = inv(framerate)

    r = range(0, t, step = dt)

    amps = s.(r) # evaluating the signal

    if haskey(s.info, "WhiteNoise")
        amps = amps + WhiteNoise(s.info["WhiteNoise"])
    end

    if haskey(s.info, "GaussianNoise")
        amps = amps + GaussianNoise(s.info["GaussianNoise"])
    end

    if haskey(s.info, "BrownianNoise")
        amps = amps + BrownianNoise(s.info["BrownianNoise"])
    end

    if haskey(s.info, "PinkNoise")
        amps = amps + PinkNoise(s.info["PinkNoise"]...)
    end

    Wave(amps, r, framerate, s)

end

function wavplay(w::Wave)

    wavplay(w.amps, w.framerate)

end

@recipe function plot(w::Wave)

    if haskey(w.signal.info, "name")
        label --> (w.signal.info["name"] * " framerate $(w.framerate)")
    else
        label --> "Framerate $(w.framerate)"
    end
    yguide --> "Amplitude"
    xguide --> "Time in s"

    w.time, w.amps

end

function filter(w::Wave, type::Symbol, args...)

    amps = deepcopy(spec.amps)

    if type == :band
        @assert length(args) == 2
        @assert args[1] < args[2]
        amps[spec.fs .< args[1] ] .= 0
        amps[spec.fs .> args[2] ] .= 0
    end

    if type == :low
        @assert length(args) == 1
        amps[spec.fs .> args[1] ] .= 0
    end

    if type == :high
        @assert length(args) == 1
        amps[spec.fs .< args[1] ] .= 0
    end

    return Spectrum(amps, spec.fs, abs2.(amps), spec.signal, spec.framerate, spec.d)

end


mutable struct TspLibProblem
    name::String
    numcities::Int
    weights::Matrix{Float64}
end

size(t::TspLibProblem) = t.numcities

function resetweights!(t::TspLibProblem, nc::Int)
    t.numcities = nc
    t.weights = zeros(Float64, nc, nc)
end

function cost(p::TspLibProblem, visitorder::Vector{Int})
    startcity = city = first(visitorder)
    cost = 0.0
    for i in 2:length(visitorder)
        nextcity = visitorder[i]
        cost += p.weights[city, nextcity]
        city = nextcity
    end
    cost += p.weights[city, startcity]
    return cost
end

function whenmatch(fn, re, s)
    m = match(re, s)
    if !isnothing(m)
        fn(m)
        return true
    end
    return false
end

function parse_edge_weights!(p, lines, lineindex)
    row = col = 1
    while !isnothing(match(r"^(\s*(\d+))+\s*$", lines[lineindex]))
        map(split(lines[lineindex], r"\s+")) do dstr
            if length(dstr) > 0
                p.weights[row, col] = p.weights[col, row] = parse(Int, strip(dstr))
                col += 1
                if col > row
                    row += 1
                    col = 1
                end
            end
        end
        lineindex += 1
    end
    return lineindex - 1
end

filename = "dantzig42.tsp"

function read_tsplib_file(filename)
    p = TspLibProblem("", 1, zeros(Float64, 1, 1))
    lines = readlines(filename)
    lineindex = 0
    while lineindex < length(lines)
        lineindex += 1
        l = lines[lineindex]
        whenmatch(m -> p.name = strip(m[1]),
            r"\s*NAME\s*:\s*(.+)$", l) && continue
        whenmatch(m -> resetweights!(p, parse(Int, m[1])),
            r"\s*DIMENSION\s*:\s*(\d+)$", l) && continue
        whenmatch(m -> lineindex = parse_edge_weights!(p, lines, lineindex+1),
            r"\s*EDGE_WEIGHT_SECTION\s*$", l)
    end
    return p
end


"""
A [logarithmic number system](https://en.wikipedia.org/wiki/Logarithmic_number_system).

Provides the signed [`Logarithmic`](@ref) and unsigned [`ULogarithmic`](@ref) types, which represent real numbers and positive real numbers respectively.
"""
module LogarithmicNumbers

using Requires, Random

import Base:
	exp, log, *, /, ^, inv, +, -, prod, sum, angle, show,
	write, read, float, big, unsigned, signed, widen,
	typemin, typemax, zero, one, iszero, isone, isinf,
	isfinite, isnan, sign, signbit, abs, ==, isequal, <, ≤,
	cmp, isless, nextfloat, prevfloat, rand, promote_rule,
	conj, real, imag

export ULogarithmic, Logarithmic, CLogarithmic

include("types.jl")
include("constructors.jl")
include("ulogarithmic.jl")
include("logarithmic.jl")
include("clogarithmic.jl")
include("promotion.jl")
include("arithmetic.jl")
include("random.jl")
include("io.jl")
include("init.jl")

# type aliases
for T in (Float64, Float32, Float16, BigFloat)
  ULogT = Symbol(:ULog, T)
  LogT = Symbol(:Log, T)
  CLogT = Symbol(:CLog, T)
  @eval begin
      const $(ULogT) = $(ULogarithmic{T})
      const $(LogT) = $(Logarithmic{T})
      const $(CLogT) = $(CLogarithmic{T})
      export $(ULogT), $(LogT), $(CLogT)
  end
end

end # module


function twin01()

   k = 4
   p = poly([-0.2*ones(k); 0.39*ones(k); 0.40*ones(k)])
   z = [-0.2 k; 0.39 k; 0.4 k];

   p, PolyZeros(z)
end


### Override for Graphs' version, copied from https://github.com/JuliaGraphs/Graphs.jl/blob/master/src/shortestpaths/dijkstra.jl with some adjustments
function dijkstra_shortest_paths(g::AbstractIndexedGraph,
    srcs::Vector{U},
    distvec::AbstractVector{T}=ones(Int, ne(g));
    allpaths=false,
    trackvertices=false
    ) where T <: Real where U <: Integer

    issubset(1:ne(g), eachindex(distvec)) || throw(ArgumentError("Bad indices in distvec"))
    nvg = nv(g)
    dists = fill(typemax(T), nvg)
    parents = zeros(U, nvg)

    pathcounts = zeros(nvg)
    preds = fill(Vector{U}(), nvg)
    H::TrackingHeap{Int64, T, 2, MinHeapOrder, NoTrainingWheels} = TrackingHeap(T; S=NoTrainingWheels)
    sizehint!(H, nvg)

    for src in srcs
        dists[src] = zero(T)
        pathcounts[src] = one(Float64)
        H[src] = zero(T)
    end

    closest_vertices = Vector{U}()  # Maintains vertices in order of distances from source

    trackvertices && sizehint!(closest_vertices, nvg)

    @inbounds while !isempty(H)
        u, d = pop!(H)

        if trackvertices
            push!(closest_vertices, u)
        end

        for e in outedges(g, u)
            alt = d + distvec[idx(e)]
            v = dst(e)

            visited = dists[v] != typemax(T)
            if !visited
                dists[v] = alt
                parents[v] = u

                pathcounts[v] += pathcounts[u]
                if allpaths
                    preds[v] = [u;]
                end
                H[v] = alt
            elseif alt < dists[v]
                dists[v] = alt
                parents[v] = u
                #615
                pathcounts[v] = pathcounts[u]
                if allpaths
                    resize!(preds[v], 1)
                    preds[v][1] = u
                end
                H[v] = alt
            elseif alt == dists[v]
                pathcounts[v] += pathcounts[u]
                if allpaths
                    push!(preds[v], u)
                end
            end
        end
    end

    if trackvertices
        for s in vertices(g)
            visited = dists[s] != typemax(T)
            if !visited
                push!(closest_vertices, s)
            end
        end
    end

    for src in srcs
        pathcounts[src] = one(Float64)
        parents[src] = 0
        empty!(preds[src])
    end

    return DijkstraState{T,U}(parents, dists, preds, pathcounts, closest_vertices)
end

function dijkstra_shortest_paths(g::AbstractIndexedGraph, src::Integer,
        distvec::AbstractVector=ones(Int, ne(g)); kw...)
    dijkstra_shortest_paths(g, [src;], distvec; kw...)
end

function dijkstra_shortest_paths(g::AbstractIndexedGraph, src,
        distmx::AbstractMatrix; kw...)
    Wt = convert(SparseMatrixCSC, transpose(distmx))
    distvec = nonzeros(Wt)
    dijkstra_shortest_paths(g, src, distvec; kw...)
end

### Homemade minimal version
function dijkstra(g::AbstractIndexedGraph{T}, s::Integer,
        w::AbstractVector{U}=ones(Int, ne(g))) where {T<:Integer, U<:Real}
    n = nv(g)
    inf = typemax(U)
    D = fill(inf, n); D[s] = 0
    parents = fill(0, n); parents[s] = 0
    Q::TrackingHeap{Int64, U, 2, MinHeapOrder, NoTrainingWheels} = TrackingHeaps.TrackingHeap(U; S=TrackingHeaps.NoTrainingWheels)
    sizehint!(Q, n)
    Q[s] = zero(U)
    while !isempty(Q)
        v, d = TrackingHeaps.pop!(Q)
        for e in outedges(g, v)
            i = dst(e)
            dd = d + w[idx(e)]
            if dd < D[i]
                Q[i] = D[i] = dd
                parents[i] = v
            end
        end
    end
    return D, parents
end


# Internal facilities for fast random number generation

immutable RandIntSampler  # for generating Int samples in [0, K-1]
    a::Int
    Ku::UInt
    U::UInt

    @compat RandIntSampler(K::Int) = (Ku = UInt(K); new(1, Ku, div(typemax(UInt), Ku) * Ku))
    @compat RandIntSampler(a::Int, b::Int) = (Ku = UInt(b-a+1); new(a, Ku, div(typemax(UInt), Ku) * Ku))
end

function rand(s::RandIntSampler)
    x = rand(UInt)
    while x >= s.U
        x = rand(UInt)
    end
    @compat s.a + Int(rem(x, s.Ku))
end

randi(K::Int) = rand(RandIntSampler(K))
randi(a::Int, b::Int) = rand(RandIntSampler(a, b))


pos2box(x::Float64; w=2) = floor(Int,x/w)
pos2box(p::Vector{Float64}, dimv::Val; w=2) = ntuple(i->pos2box(p[i];w=w), dimv)

function find_neighbors(cellbox, root::Int; s=2)
	bools = isneighbor.(cellbox; rootbox=cellbox[root], s=s)
	bools[root] = false
	return bools
end
isneighbor(neighbox::NTuple{dim,Int}; rootbox::NTuple{dim,Int}, s=2) where dim = all(-s < neighbox[i] - rootbox[i] < s for i=1:dim)

function pushing!(tumor, root, cellbox, dimv::Val)
    queue = [root]
    while !isempty(queue)
        root = popfirst!(queue)
        r1 = tumor[root].position

        for n in shuffle!(findall(find_neighbors(cellbox, root; s=2)))
            r2 = tumor[n].position
            d = r2.-r1
            if norm(d) < 2.
                r2 .+= d*(2.2/norm(d) - 1)#/2
                #r1 .-= d*(2.2/norm(d) - 1)/2
                cellbox[n] = pos2box(r2, dimv)
                #cellbox[root] = pos2box(r1;dimv=Val(dim))
                push!(queue, n)
				#push!(queue, n)
            end
        end
    end
end

function pushing_recursive!(tumor, root, cellbox, dimv::Val)
    r1 = tumor[root].position

    N=length(tumor)
    for n in shuffle!(findall(find_neighbors(cellbox, root)))
        r2 = tumor[n].position
        d = r2.-r1
        if norm(d) < 2.
            r2 .+= d*(2.2/norm(d) - 1)
            cellbox[n] = pos2box(r2, dimv)

            pushing_recursive!(tumor, n, cellbox)
        end
    end
end


# error bars

#makie

using CairoMakie

f = Figure()
Axis(f[1, 1])

xs = 0:0.5:10
ys = 0.5 .* sin.(xs)

lowerrors = fill(0.1, length(xs))
higherrors = LinRange(0.1, 0.4, length(xs))

errorbars!(xs, ys, lowerrors, higherrors,
    color = range(0, 1, length = length(xs)),
    whiskerwidth = 10, linewidth=2, colormap = :tableau_hue_circle)

# plot position scatters so low and high errors can be discriminated
CairoMakie.scatter!(xs, ys, markersize = 4, color = :black)