function use_linear_search(time_series)
    return length(time_series) < 15
end

function univariate_search_bounds(p0)
    return p0[1] - 100.0, p0[1] + 100.0
end

# create diagonal matrix given x0, t, x1 s.t. x1 = e^{At}*x0 (if possible)
function guess_dyn_linear(time_series, x0)
    n = length(time_series[1]) - 1
    t_x1 = time_series[end]
    t = t_x1[1] - time_series[1][1]
    diag = Vector{Float64}(undef, n)
    for i in 1:n
        if iszero(x0[i]) || t_x1[i+1] / x0[i] <= 1e-5
            diag[i] = 0.0  # not solvable
        else
            @assert t > 0.0
            diag[i] = log(t_x1[i+1] / x0[i]) / t
        end
    end
    A = Diagonal{Float64}(diag)
    return A
end

# use constant dynamics
function guess_dyn(time_series, x0)
    n = length(time_series[1]) - 1
    t_x1 = time_series[end]
    t = t_x1[1] - time_series[1][1]
    A = zeros(n, n)
    b = [(t_x1[i+1] - x0[i]) / t for i in 1:n]
    return AffDyn(A, b)
end

function guess_x0(time_series)
    x0 = collect(time_series[1][2:end])
    return x0
end

# convert parameter array to symbolic representation
function parametric_to_symbolic(p, n; first::Bool=false, linear::Bool=false)
    A = Matrix{Float64}(undef, n, n)
    k = 0
    for j in 1:n
        for i in 1:n
            A[i, j] = p[k + i]
        end
        k += n
    end

    if linear
        dyn = A
    else
        b = p[(n*n+1):(n*(n+1))]
        dyn = AffDyn(A, b)
    end

    if first
        x0 = p[(end-n+1):end]
        return dyn, x0
    else
        return dyn
    end
end

# solve ODE problem in continuous time
function solve_ode_continuous(dyn, x0, time)
    f(u, p, t) = dyn.A * u + dyn.b
    problem = DifferentialEquations.ODEProblem(f, x0, time)
    sol = DifferentialEquations.solve(problem)
    return sol
end

# solve ODE problem at discrete points in time
function solve_ode_discrete(dyn, x0, time)
    return (t) -> successor(dyn, t, x0)
end

# distance between ODE solution and time series
function distance(dyn, x0, time_series)
    d = -Inf
    t0 = time_series[1][1]
    for t_p_ts in time_series
        t = t_p_ts[1] - t0
        p_ts = t_p_ts[2:end]
        p_ode = successor(dyn, t, x0)
        if p_ode == nothing
            return Inf  # exact solution failed
        end
        d_raw = map(-, p_ts, p_ode)
        d_new = norm(d_raw, Inf)  # infinity norm of the distance
        d = max(d, d_new)
    end
    return d
end

# distance between parametric representation and time series
function parametric_problem(p, time_series, n, x0; linear::Bool=false)
    first = x0 == nothing

    # convert to symbolic representation
    result = parametric_to_symbolic(p, n; first=first, linear=linear)
    if first
        dyn, x0 = result
    else
        dyn = result
    end

    # compute distance between ODE solution and time series
    d = distance(dyn, x0, time_series)

    return d
end

# convert symbolic representation to parameter array
function symbolic_to_parametric(dyn, n, x0=nothing; linear::Bool=linear)
    first = x0 != nothing
    b_factor = linear ? 0 : 1
    A = linear ? dyn : dyn.A

    p = Vector{Float64}(undef,
                        first ? n * (n + 1 + b_factor) : n * (n + b_factor))
    k = 0
    for j in 1:n
        for i in 1:n
            p[k + i] = A[i, j]
        end
        k += n
    end
    if !linear
        p[(n*n+1):(n*(n+1))] = dyn.b
    end
    if first
        p[(end-n+1):end] = x0
    end
    return p
end

function time_series_to_symbolic_one_piece(time_series, dyn_init, x0_init;
                                           first::Bool=false,
                                           linear::Bool=false)
    # problem dimension
    n = length(time_series[1]) - 1

    # distance function to optimize
    x0_init_use = first ? nothing : x0_init
    dist = (p) -> parametric_problem(p, time_series, n, x0_init_use;
                                     linear=linear)

    # convert initial guess to parametric representation
    if first
        p0 = symbolic_to_parametric(dyn_init, n, x0_init; linear=linear)
    else
        p0 = symbolic_to_parametric(dyn_init, n; linear=linear)
    end

    # optimize
    if length(p0) == 1
        bounds = univariate_search_bounds(p0)
        sol = optimize(dist, bounds[1], bounds[2])
    else
        sol = optimize(dist, p0)
    end
    p = sol.minimizer
    min_d = sol.minimum

    result = parametric_to_symbolic(p, n; first=first, linear=linear)
    if first
        dyn, x0 = result
        return dyn, x0, min_d
    else
        dyn = result
        return dyn, min_d
    end
end

function time_series_to_symbolic_longest_piece(time_series, ε, dyn_init,
                                               x0_init;
                                               linear_search::Bool=true,
                                               first::Bool=false)
    guess_dyn_first = dyn_init == nothing
    solution_found = false
    i = length(time_series)
    binary_search_range = [1, i+1]  # range [l, u] where l works and u does not
    dyn_best = nothing
    x0_best = nothing
    while i > 1
        time_series_shortened = time_series[1:i]

        for linear in [true, false]
            if length(time_series_shortened) == 2 && linear
                # skipping linear fit
                continue
            end

            if guess_dyn_first
                if linear
                    dyn_init_use = guess_dyn_linear(time_series_shortened,
                                                    x0_init)
                else
                    dyn_init_use = guess_dyn(time_series_shortened, x0_init)
                end
            else
                if linear
                    dyn_init_use = dyn_init.A
                else
                    dyn_init_use = dyn_init
                end
            end

            result = time_series_to_symbolic_one_piece(time_series_shortened,
                                                       dyn_init_use, x0_init;
                                                       first=first,
                                                       linear=linear)
            if first
                dyn, x0, min_d = result
            else
                dyn, min_d = result
            end
            if min_d <= ε
                if linear
                    @assert dyn isa AbstractMatrix
                    n = length(time_series[1]) - 1
                    dyn = AffDyn(dyn, zeros(n))
                else
                    @assert dyn isa AffDyn
                end
                dyn_best = dyn
                if first
                    x0_best = x0
                end
                solution_found = true
                break
            end
        end
        guess_dyn_first = true

        if linear_search
            # linear search
            if !solution_found
                i -= 1
                solution_found = false
                continue
            end
        else
            # binary search, rounding up
            if solution_found
                binary_search_range[1] = i
            else
                binary_search_range[2] = i
            end
            if binary_search_range[1] < binary_search_range[2] - 1
                i = div(binary_search_range[1] + binary_search_range[2] + 1, 2)
                solution_found = false
                continue
            end
        end

        # found the largest solution
        if !linear_search
            i = binary_search_range[1]
        end
        @assert dyn_best != nothing
        if first
            @assert x0_best != nothing
            return dyn_best, x0_best, i
        else
            return dyn_best, i
        end
    end
    @assert false "it should always be possible to fit a function to two points"
end

function time_series_to_symbolic(time_series, ε, dyn_init, x0_init;
                                 linear_search::Bool=use_linear_search(time_series),
                                 log::Bool=false)
    @assert iszero(time_series[1][1]) "the time series should start with t = 0 but started with " *
        time_series[1][1]

    if x0_init == nothing
        x0_init = guess_x0(time_series)
    end

    dyn, x0, l = time_series_to_symbolic_longest_piece(time_series, ε, dyn_init,
        x0_init; linear_search=linear_search, first=true)
    t = time_series[l][1]
    x1 = successor(dyn, t, x0)
    dyns = [dyn]
    times = [t]
    indices = [l]
    while l < length(time_series)
        time_series_shortened = time_series[l:end]
        dyn, Δl = time_series_to_symbolic_longest_piece(time_series_shortened,
            ε, nothing, x1; linear_search=linear_search, first=false)
        l_new = l + Δl - 1
        t = time_series[l_new][1] - time_series[l][1]
        x1 = successor(dyn, t, x1)
        l = l_new
        push!(dyns, dyn)
        push!(times, time_series[l_new][1])
        push!(indices, l)
    end
    if log
        println("found a function with $(length(dyns)) pieces")
    end
    return dyns, x0, times, indices
end


using InvertibleNetworks, Test, LinearAlgebra, Random

Random.seed!(11)

##############################################################################################################################
# 4D Tensor
X = randn(Float32, 28, 28, 2, 4)

# Squeeze and unsqueeze
@test isapprox(norm(X - unsqueeze(squeeze(X; pattern="column"); pattern="column")), 0f0; atol=1f-6)
@test isapprox(norm(X - unsqueeze(squeeze(X; pattern="patch"); pattern="patch")), 0f0; atol=1f-6)
@test isapprox(norm(X - unsqueeze(squeeze(X; pattern="checkerboard"); pattern="checkerboard")), 0f0; atol=1f-6)

#test norm preservation of orthogonal transforms
@test norm(X) ≈ norm(Haar_squeeze(X))

# Wavelet transform invertibility
Y = wavelet_squeeze(X)
X_ = wavelet_unsqueeze(Y)
@test isapprox(norm(X - X_)/norm(X), 0f0; atol=1f-5)
@test X ≈ invHaar_unsqueeze(Haar_squeeze(X))

# Wavelet transform adjoint test
X = randn(Float32, 28, 28, 2, 4)
Y = randn(Float32, 14, 14, 8, 4)
a = dot(Y, wavelet_squeeze(X))
b = dot(X, wavelet_unsqueeze(Y))
@test isapprox(a/b - 1f0, 0f0; atol=1f-5)
a = dot(Y, Haar_squeeze(X))
b = dot(X, invHaar_unsqueeze(Y))
@test isapprox(a/b - 1f0, 0f0; atol=1f-5)

# Split and cat
@test isapprox(norm(X - tensor_cat(tensor_split(X))), 0f0; atol=1f-5)

# Unsqueeze works for any size as long as channel size divisible by 4 (for 4D Tensor)
x_size  = 27
y_size  = 27
ch_size = 8
batch_size = 4
X = randn(Float32, x_size, y_size, ch_size, batch_size)
@test isequal(size(unsqueeze(X)), (x_size*2, y_size*2, div(ch_size,4), batch_size))

##############################################################################################################################
# 5D Tensor
X = randn(Float32, 16, 16, 16, 2, 4)

# Squeeze and unsqueeze
@test isapprox(norm(X - unsqueeze(squeeze(X; pattern="column"); pattern="column")), 0f0; atol=1f-6)
@test isapprox(norm(X - unsqueeze(squeeze(X; pattern="patch"); pattern="patch")), 0f0; atol=1f-6)

#test norm preservation of orthogonal transforms
@test norm(X) ≈ norm(Haar_squeeze(X))

# Wavelet transform invertibility
Y = wavelet_squeeze(X)
X_ = wavelet_unsqueeze(Y)
@test isapprox(norm(X - X_)/norm(X), 0f0; atol=1f-5)
@test X ≈ invHaar_unsqueeze(Haar_squeeze(X))

# Wavelet transform adjoint test
X = randn(Float32, 16, 16, 16, 2, 4)
Y = randn(Float32, 8, 8, 8, 16, 4)
a = dot(Y, wavelet_squeeze(X))
b = dot(X, wavelet_unsqueeze(Y))
@test isapprox(a/b - 1f0, 0f0; atol=1f-5)
a = dot(Y, Haar_squeeze(X))
b = dot(X, invHaar_unsqueeze(Y))
@test isapprox(a/b - 1f0, 0f0; atol=1f-5)

# Split and cat
@test isapprox(norm(X - tensor_cat(tensor_split(X))), 0f0; atol=1f-5)

# Unsqueeze works for any size as long as channel size divisible by 8 (for 5D Tensor)
x_size  = 27
y_size  = 27
z_size  = 27
ch_size = 8
batch_size = 4
X = randn(Float32, x_size, y_size, z_size, ch_size, batch_size)
@test isequal(size(unsqueeze(X)), (x_size*2, y_size*2, z_size*2, div(ch_size,8), batch_size))


mutable struct body_motion_data_t <: LCMType
    timestamp::Int64
    body_or_frame_name::String
    spline::piecewise_polynomial_t
    quat_task_to_world::SVector{4, Float64}
    translation_task_to_world::SVector{3, Float64}
    xyz_kp_multiplier::SVector{3, Float64}
    xyz_damping_ratio_multiplier::SVector{3, Float64}
    expmap_kp_multiplier::Float64
    expmap_damping_ratio_multiplier::Float64
    weight_multiplier::SVector{6, Float64}
    in_floating_base_nullspace::Bool
    control_pose_when_in_contact::Bool
end

@lcmtypesetup(body_motion_data_t)


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENCE in the project root.
# ------------------------------------------------------------------

"""
    IDW(var₁=>param₁, var₂=>param₂, ...)

Inverse distance weighting estimation solver.

## Parameters

* `neighbors` - Number of neighbors (default to all the data)
* `distance`  - A distance defined in Distances.jl (default to `Euclidean()`)
* `power`     - Power of the distances (default to `1`)

### References

Shepard 1968. *A two-dimensional interpolation function for irregularly-spaced data.*
"""
@estimsolver IDW begin
  @param neighbors = nothing
  @param distance = Euclidean()
  @param power = 1
end

function solve(problem::EstimationProblem, solver::IDW)
  # retrieve problem info
  pdata = data(problem)
  pdomain = domain(problem)

  mactypeof = Dict(name(v) => mactype(v) for v in variables(problem))

  # result for each variable
  μs = []; σs = []

  for covars in covariables(problem, solver)
    for var in covars.names
      # get user parameters
      varparams = covars.params[(var,)]

      # determine value type
      V = mactypeof[var]

      # retrieve non-missing data
      locs = findall(!ismissing, pdata[var])
      𝒟 = view(pdata, locs)
      n = nelements(𝒟)

      # determine number of nearest neighbors to use
      k = isnothing(varparams.neighbors) ? n : varparams.neighbors

      # determine distance type
      D = varparams.distance

      # determine power of distances
      p = varparams.power

      @assert n > 0 "estimation requires data"

      @assert k ≤ n "invalid number of neighbors"

      @assert p > 0 "power must be positive"

      # fit search tree
      X = [coordinates(centroid(𝒟, i)) for i in 1:nelements(𝒟)]
      if D isa NearestNeighbors.MinkowskiMetric
        tree = KDTree(X, D)
      else
        tree = BallTree(X, D)
      end

      # lookup non-missing values
      z = 𝒟[var]

      # pre-allocate memory for results
      varμ = Vector{V}(undef, nelements(pdomain))
      varσ = Vector{V}(undef, nelements(pdomain))

      # estimation loop
      for loc in traverse(pdomain, LinearPath())
        x = coordinates(centroid(pdomain, loc))

        is, ds = knn(tree, x, k)
        ws = one(V) ./ ds.^p
        Σw = sum(ws)

        if isinf(Σw) # some distance is zero?
          j = findfirst(iszero, ds)
          μ = z[is[j]]
          σ = zero(V)
        else
          ws /= Σw
          vs  = view(z, is)
          μ = sum(ws[i]*vs[i] for i in eachindex(vs))
          σ = minimum(ds)
        end

        varμ[loc] = μ
        varσ[loc] = σ
      end

      push!(μs, var => varμ)
      push!(σs, Symbol(var,"_distance") => varσ)
    end
  end

  georef((; μs..., σs...), pdomain)
end


using Test
using OpenQuantumSystems
using SparseArrays, LinearAlgebra

@testset "metrics" begin

    b1 = SpinBasis(1 // 2)
    b2 = FockBasis(6)

    psi1 = spinup(b1) ⊗ coherentstate(b2, 0.1)
    psi2 = spindown(b1) ⊗ fockstate(b2, 2)

    rho = tensor(psi1, dagger(psi1))
    sigma = tensor(psi2, dagger(psi2))
    sigma_s = Commutator(sigma)
    rho_s = Commutator(rho)

    # tracenorm
    @test tracenorm(0 * rho_s) ≈ 0.0
    @test tracenorm_h(0 * rho_s) ≈ 0.0
    @test tracenorm_nh(0 * rho_s) ≈ 0.0

    @test tracenorm(rho_s) ≈ 26.0 # 1 * 6 * 4
    @test tracenorm_h(rho_s) ≈ 26.0
    @test tracenorm_nh(rho_s) ≈ 26.0

    # tracedistance
    @test tracedistance(rho_s, sigma_s) ≈ 26.0
    @test tracedistance_h(rho_s, sigma_s) ≈ 26.0
    @test tracedistance_nh(rho_s, sigma_s) ≈ 26.0

    @test tracedistance(rho_s, rho_s) ≈ 0.0
    @test tracedistance_h(rho_s, rho_s) ≈ 0.0
    @test tracedistance_nh(rho_s, rho_s) ≈ 0.0

    @test tracedistance(sigma_s, sigma_s) ≈ 0.0
    @test tracedistance_h(sigma_s, sigma_s) ≈ 0.0
    @test tracedistance_nh(sigma_s, sigma_s) ≈ 0.0

    # tracedistance
    @test tracedistance(rho_s, sigma_s) ≈ 26.0
    @test tracedistance(rho_s, rho_s) ≈ 0.0
    @test tracedistance(sigma_s, sigma_s) ≈ 0.0

    @test tracedistance_nh(rho_s, rho_s) ≈ 0.0

end # testset


"""
    fS(eqn::Euler{d},UL,UR) where {d}

Entropy conservative fluxes for the compressible Euler equations, where UL,UR are
tuples of left and right states in terms of the conservative variables Q = (ρ,ρuvw...,E)
"""
@inline function fS(eqn::Euler{d},UL,UR) where {d}
    QL = cons_to_prim_beta(eqn,UL)
    QR = cons_to_prim_beta(eqn,UR)
    return fS_prim(eqn,QL,QR)
end

"""
    fS_prim(eqn::Euler{d},UL,UR)

Entropy conservative fluxes for the compressible Euler equations, where
UL,UR are tuples of left and right solutions in terms of β-primitive variables
Q = (rho,uvw...,β) (where uvw = velocity components and β is inverse temperature).
"""
@inline function fS_prim(eqn::Euler{d},UL,UR) where {d}
    logL = map(x->log.(x),(first(UL),last(UL)))
    logR = map(x->log.(x),(first(UR),last(UR)))

    return fS_prim_log(eqn,(UL...,logL...),(UR...,logR...))
end

# manually unroll for d = 1,2,3 to avoid type instability
@inline function fS_prim(eqn::Euler{1},UL,UR)
    UlogL = (UL[1],UL[2],UL[3],log.(UL[1]),log.(UL[3]))
    UlogR = (UR[1],UR[2],UR[3],log.(UR[1]),log.(UR[3]))
    return fS_prim_log(eqn,UlogL,UlogR)
end

@inline function fS_prim(eqn::Euler{2},UL,UR)
    UlogL = (UL[1],UL[2],UL[3],UL[4],log.(UL[1]),log.(UL[4]))
    UlogR = (UR[1],UR[2],UR[3],UR[4],log.(UR[1]),log.(UR[4]))
    return fS_prim_log(eqn,UlogL,UlogR)
end

@inline function fS_prim(eqn::Euler{3},UL,UR)
    UlogL = (UL[1],UL[2],UL[3],UL[4],UL[5],log.(UL[1]),log.(UL[5]))
    UlogR = (UR[1],UR[2],UR[3],UR[4],UR[5],log.(UR[1]),log.(UR[5]))
    return fS_prim_log(eqn,UlogL,UlogR)
end

"""
    fS_prim_log(eqn::Euler{d},UL,UR)

Entropy conservative fluxes for the compressible Euler equations, where
UL,UR are tuples of left and right solutions (in β-primitive variables) and logs
of rho, β, e.g., U = (rho,uvw...,β,log.(rho),log.(β)).
"""
@inline function fS_prim_log(eqn::Euler{1},UL,UR)

    @unpack γ = eqn

    rhoL,uL,betaL,rhologL,betalogL = UL
    rhoR,uR,betaR,rhologR,betalogR = UR

    rholog  = logmean.(rhoL,rhoR,rhologL,rhologR)
    betalog = logmean.(betaL,betaR,betalogL,betalogR)

    # arithmetic avgs
    rhoavg = (@. .5*(rhoL+rhoR))
    uavg   = (@. .5*(uL+uR))

    unorm = (@. uL*uR)
    pa    = (@. rhoavg/(betaL+betaR))
    f4aux = (@. rholog/(2*(γ-1)*betalog) + pa + .5*rholog*unorm)

    FxS1 = (@. rholog*uavg)
    FxS2 = (@. FxS1*uavg + pa)
    FxS3 = (@. f4aux*uavg)

    return SVector(FxS1,FxS2,FxS3)
end

@inline function fS_prim_log(eqn::Euler{2},UL,UR)

    @unpack γ = eqn

    rhoL,uL,vL,betaL,rhologL,betalogL = UL
    rhoR,uR,vR,betaR,rhologR,betalogR = UR

    rholog = logmean.(rhoL,rhoR,rhologL,rhologR)
    betalog = logmean.(betaL,betaR,betalogL,betalogR)

    # arithmetic avgs
    rhoavg = (@. .5*(rhoL+rhoR))
    uavg   = (@. .5*(uL+uR))
    vavg   = (@. .5*(vL+vR))

    unorm = (@. uL*uR + vL*vR)
    pa    = (@. rhoavg/(betaL+betaR))
    f4aux = (@. rholog/(2*(γ-1)*betalog) + pa + .5*rholog*unorm)

    FxS1 = (@. rholog*uavg)
    FxS2 = (@. FxS1*uavg + pa)
    FxS3 = (@. FxS1*vavg)
    FxS4 = (@. f4aux*uavg)

    FyS1 = (@. rholog*vavg)
    FyS2 = (@. FxS3)
    FyS3 = (@. FyS1*vavg + pa)
    FyS4 = (@. f4aux*vavg)

    return SVector(FxS1,FxS2,FxS3,FxS4),SVector(FyS1,FyS2,FyS3,FyS4)
end

@inline function fS_prim_log(eqn::Euler{3},UL,UR)

    @unpack γ = eqn

    rhoL,uL,vL,wL,betaL,rhologL,betalogL = UL
    rhoR,uR,vR,wR,betaR,rhologR,betalogR = UR

    rholog = logmean.(rhoL,rhoR,rhologL,rhologR)
    betalog = logmean.(betaL,betaR,betalogL,betalogR)

    # arithmetic avgs
    rhoavg = (@. .5*(rhoL+rhoR))
    uavg   = (@. .5*(uL+uR))
    vavg   = (@. .5*(vL+vR))
    wavg   = (@. .5*(wL+wR))

    unorm = (@. uL*uR + vL*vR + wL*wR)
    pa    = (@. rhoavg/(betaL+betaR))
    E_plus_p  = (@. rholog/(2*(γ-1)*betalog) + pa + .5*rholog*unorm)

    FxS1 = (@. rholog*uavg)
    FxS2 = (@. FxS1*uavg + pa)
    FxS3 = (@. FxS1*vavg) # rho * u * v
    FxS4 = (@. FxS1*wavg) # rho * u * w
    FxS5 = (@. E_plus_p*uavg)

    FyS1 = (@. rholog*vavg)
    FyS2 = (@. FxS3) # rho * u * v
    FyS3 = (@. FyS1*vavg + pa)
    FyS4 = (@. FyS1*wavg) # rho * v * w
    FyS5 = (@. E_plus_p*vavg)

    FzS1 = (@. rholog*wavg)
    FzS2 = (@. FxS4) # rho * u * w
    FzS3 = (@. FyS4) # rho * v * w
    FzS4 = (@. FzS1*wavg + pa) # rho * w^2 + p
    FzS5 = (@. E_plus_p*wavg)

    Fx = SVector(FxS1,FxS2,FxS3,FxS4,FxS5)
    Fy = SVector(FyS1,FyS2,FyS3,FyS4,FyS5)
    Fz = SVector(FzS1,FzS2,FzS3,FzS4,FzS5)
    return Fx,Fy,Fz
end

### didn't add coordinates b/c timing doesn't seem to make a diff.
# julia> Qlog = (1.0, 0.1, 0.2, 2.5250000000000004, 0.0, 0.9262410627273233)
# julia> @btime fS_prim_log($Euler{2}(),$Qlog,$Qlog);
#   2.963 μs (0 allocations: 0 bytes)
#
# julia> @btime fS_prim_log($Euler{2}(),$Qlog,$Qlog,$1);
#   3.044 μs (0 allocations: 0 bytes)


using Test, SpecialFunctions
import TruncatedNormal as TN

@test TN.tnmom1c(0, -Inf, Inf) == 0
@test TN.tnmom1c(0, -Inf, 0) ≈ TN.tnmean(-Inf, 0)
@test TN.tnmom1c(0, 0, +Inf) ≈ TN.tnmean(0, +Inf)

@test TN.tnmom1c(0, +Inf, +Inf) == +Inf
@test TN.tnmom1c(0, -Inf, -Inf) == -Inf
@test TN.tnmom1c(1, +Inf, +Inf) == +Inf
@test TN.tnmom1c(1, -Inf, -Inf) == -Inf
@test TN.tnmom1c(-Inf, +Inf, +Inf) == +Inf
@test TN.tnmom1c(+Inf, -Inf, -Inf) == -Inf

@test isnan(TN.tnmom1c(+Inf, +Inf, +Inf))
@test isnan(TN.tnmom1c(-Inf, -Inf, -Inf))

@test isnan(TN.tnmom1c(0, 1, 0))
@test isnan(TN.tnmom1c(0, +Inf, 0))
@test isnan(TN.tnmom1c(0, 0, -Inf))
@test isnan(TN.tnmom1c(0, +Inf, -Inf))
@test isnan(TN.tnmom1c(1, 1, 0))
@test isnan(TN.tnmom1c(1, +Inf, 0))
@test isnan(TN.tnmom1c(1, 0, -Inf))

for c = (0, 1), x = (0, NaN), y = (0, NaN)
    if isnan(x) || isnan(y)
        @test isnan(TN.tnmom1c(c, x, y))
    end
end

for x = exp.(-100:100), c = -100:100
    @test TN.tnmom1c(c, -x, x) ≈ -c
end

for x = -10:10, c = -10:10
    @test TN.tnmom1c(c, x, +Inf) ≈ TN.tnmean(x, +Inf) - c
    @test TN.tnmom1c(c, -Inf, x) ≈ TN.tnmean(-Inf, x) - c
end

#= TN.tnmom1c(c, a, b) = 0 when c == tnmean(a,b) =#
for a = exp.(-10:100), b = exp.(-10:100)
    a ≤ b || continue
    c = TN.tnmean(a, b)
    @test abs(TN.tnmom1c(c, a, b)) ≤ max(1e-10, c * 1e-10)
    @test TN.tnmom1c(a, a, b) ≥ 0
    @test TN.tnmom1c(b, a, b) ≤ 0
end

@test TN.tnmom1c(0, -Inf, Inf, 0, 1) ≈ 0

for x = -10:10
    @test TN.tnmom1c(0, x, +Inf, 0, 1) ≈ TN.tnmean(x, +Inf)
    @test TN.tnmom1c(0, -Inf, x, 0, 1) ≈ TN.tnmean(-Inf, x)
end

@test TN.tnmom1c(10, 0, 10) ≈ -9.20211543919713464412024961258
@test TN.tnmom1c(1, 3, 4) ≈ 2.26045428559002115390313889722
@test TN.tnmom1c(-1, 3, 4) ≈ 4.26045428559002115390313889722
@test TN.tnmom1c(2, 3, 4) ≈ 1.26045428559002115390313889722
@test TN.tnmom1c(-2, 3, 4) ≈ 5.26045428559002115390313889722
@test TN.tnmom1c(0, -1e200, 1e200) ≈ 0
@test TN.tnmom1c(1, -1e200, 1e200) ≈ -1
@test TN.tnmom1c(1e200, -1e200, 1e200) ≈ -1e200
@test TN.tnmom1c(0, -1e6, -999000) ≈ TN.tnmean(-1e6, -999000)
@test TN.tnmom1c(0, +1e6, +Inf) ≈ TN.tnmean(+1e6, +Inf)
@test TN.tnmom1c(0, -Inf, -1e6) ≈ TN.tnmean(-Inf, -1e6)

@test TN.tnmom1c(999000, 999000, 1e6) ≈ 1.00100100099899498898098100910e-6
@test TN.tnmom1c(999500, 999000, 1e6) ≈ -499.999998998998999001005011019
@test TN.tnmom1c(1e6, 999000, 1e6) ≈ -999.999998998998999001005011019


using Test
using OnlineStats: Series, fit!
using LyapunovExponents: VecMean, VecVariance

@time @testset "OnlineStats" begin
    max_dim = 10
    len = 5
    rng = MersenneTwister(0)
    @testset "VecMean dim=$dim" for dim in 1:max_dim
        stat = VecMean(dim)
        os = Series(stat)
        xs = randn(rng, dim, len)
        for i in 1:len
            fit!(os, xs[:, i])
            @test mean(stat) ≈ mean(xs[:, 1:i], 2)[:, 1]
        end
    end
    @testset "VecVariance dim=$dim" for dim in 1:max_dim
        stat = VecVariance(dim)
        os = Series(stat)
        xs = randn(rng, dim, len)
        for i in 1:len
            fit!(os, xs[:, i])
            @test mean(stat) ≈ mean(xs[:, 1:i], 2)[:, 1]
            @test var(stat) ≈ var(xs[:, 1:i], 2)[:, 1] nans=true
        end
    end
end


# poisson processes
#
struct InhomogeneousPoisson end
struct InhomogeneousPoissonPiecewiseDecreasing end

"""
next_event_time(rate::Function, max_rate::Float64, rng::AbstractRNG)::Float64

Generate a new event from an inhomogeneous poisson process with rate Lambda(t).
Based on (Ogata’s Modified Thinning Algorithm: Ogata,  1981,  p.25,  Algorithm  2)
see also https://www.math.fsu.edu/~ychen/research/Thinning%20algorithm.pdf

# Arguments
- `rate`: rate(dt) has to be defined outside (e..g t-> rate(t+t0,args))
- `max_rate`: maximal rate in near future (has to be evaluated externally)
- `rng`: random number generator

API - output
* returns the next event time

"""
function next_time(alg::InhomogeneousPoisson, rng::AbstractRNG, rate::Function, max_rate::Float64)::Float64
    dt = 0.0
    theta = 1.0 / max_rate
    while true
        # generate next event from bounding homogeneous Poisson process with max_rate
        dt += randexp(rng) * theta
        # accept next event with probability rate(t)/rate_max [Thinning algorithm]
        if rand(rng) < rate(dt) / max_rate
            return dt
        end
    end
end

"""
Generate a new event from an inhomogeneous poisson process with rate Lambda(t) under
the assumption that rate(dt) is monotonically decreasing.
Based on (Ogata’s Modified Thinning Algorithm: Ogata,  1981,  p.25,  Algorithm  3)
see also https://www.math.fsu.edu/~ychen/research/Thinning%20algorithm.pdf

# Arguments
- `rate`: rate(dt) has to be defined outside (e..g t-> rate(t+t0,args))
- `rng`: random number generator

API - output
* returns the next event time

"""
function next_time(alg::InhomogeneousPoissonPiecewiseDecreasing, rng::AbstractRNG, rate::Function)::Float64
    dt = 0.0
    while true
        # future rate can only be smaller than current rate
        max_rate = rate(dt)
        # generate next event from bounding homogeneous Poisson process with max_rate
        dt += randexp(rng) / max_rate
        # accept next event with probability rate(t)/rate_max [Thinning algorithm]
        if rand(rng) < rate(dt) / max_rate
            return dt
        end
    end
end

"""
Generate a new event id from a collection of inhomogeneous poisson processes with
rates Lambda(t).
# Arguments
- `rates`: rates(dt); Float -> [Float]
- `max_rate`: maximal rate in near future (has to be evaluated externally)
- `rng`: random number generator

API - output
* returns the next event id
"""
function next_event(alg::InhomogeneousPoisson, rng::AbstractRNG, rates::Array{Float64}, max_rate::Float64)::Int
    next_index = 1
    cumulated_rates = cumsum(rates)
    sum_rate = cumulated_rates[end]

    theta = rand(rng) * sum_rate
    id = 1
    # catch lower-bound case that cannot be reached by binary search
    if theta >= cumulated_rates[1]
        # binary search
        index_l = 1
        index_r = length(cumulated_rates)
        while index_l < index_r - 1
            # index_m = floor(Int,(index_l+index_r)/2)
            index_m = fld(index_l + index_r, 2)
            if cumulated_rates[index_m] < theta
                index_l = index_m
            else
                index_r = index_m
            end
        end
        id = index_r
    end

    return id
end

"""
Generate a new event from a collection of inhomogeneous poisson processes with
rates Lambda(t).
# Arguments
- `rates`: rates(dt); Float -> [Float]
- `max_rate`: maximal rate in near future (has to be evaluated externally)
- `rng`: random number generator

API - output
* returns the next event time and event id as tuple (dt, id)
"""
function next(alg::InhomogeneousPoisson, rng::AbstractRNG, rates::Function, max_rate::Float64)::Tuple{Float64,Int}
    rate(t) = sum(rates(t))
    dt = next_time(alg, rng, rate, max_rate)
    id = next_event(alg, rng, rates(dt), max_rate)
    return dt, id
end


module psmp

export weibel_streaming
using Sobol

function uniform_6d(Np::Int, nskip=0,T=Float64)
  x1n=Array{T}(Np)
  x2n=Array{T}(Np)
  x3n=Array{T}(Np)
  v1n=Array{T}(Np)
  v2n=Array{T}(Np)
  v3n=Array{T}(Np)

  sob = Sobol.SobolSeq(6)
  Sobol.skip(sob, 4+nskip) # Skip some entries
  for i=1:Np
    x1n[i],x2n[i],x3n[i],v1n[i],v2n[i],v3n[i]=Sobol.next(sob)
  end

  return x1n,x2n,x3n,v1n,v2n,v3n
end

import SpecialFunctions: erfinv
# Sampling

function maxwellian_6d{T}(Np::Int, nskip::Int=0,  L::Array{T,1}=ones(T,3),
            sigma::Array{T,1}=ones(T,3),
            mu::Array{T,1}=zeros(T,3) )

            x1n,x2n,x3n,v1n,v2n,v3n=uniform_6d(Np,nskip)

            p=maxwellian_6d(x1n,x2n,x3n,v1n,v2n,v3n,
                            L, sigma,mu)
            return x1n,x2n,x3n,v1n,v2n,v3n,p
end

function maxwellian_6d{T}(x1n,x2n,x3n,v1n,v2n,v3n,
               L::Array{T,1}=ones(T,3),
              sigma::Array{T,1}=ones(T,3),
              mu::Array{T,1}=zeros(T,3) )
              x1n.*=L[1]
              x2n.*=L[2]
              x3n.*=L[3]

              norminv=p->sqrt(2.).*erfinv.(2.*p - 1.0)

              v1n.=norminv(v1n).*sigma[1].+mu[1]
              v2n.=norminv(v2n).*sigma[2].+mu[2]
              v3n.=norminv(v3n).*sigma[3].+mu[3]

  #Probability density
  p(x1,x2,x3,v1,v2,v3)=exp.(- (v1.-mu[1]).^2./sigma[1]^2/2.
                        .-(v2.-mu[2]).^2./sigma[2]^2/2.
                       .-(v3.-mu[3]).^2./sigma[3]^2/2.  ).*
                   (2*pi)^(-1.5)/(prod(sigma)*prod(L))

  return p
end

# Rotation of initial density for testing
# function rotate3d
#
# end

# input uniformly distributed
function maxwellian_stream_3d{T}(v1n,v2n,v3n, alpha::Array{T,1},
                  sigma::Array{T,2},mu::Array{T,2})
  Ns::Int=length(alpha) # number of streams
  Np::Int=length(v1n)
  norminv=p->sqrt(2.).*erfinv.(2.*p - 1.0)
  v1n.=norminv(v1n)
  v2n.=norminv(v2n)
  v3n.=norminv(v3n)

  nidx=Array{Int}(Ns+1)
  nidx.=[1;ceil(Int,alpha[1:Ns-1]*Np);Np]

  for ndx=1:Ns
    # use fused multiply add fma
    v1n[nidx[ndx]:nidx[ndx+1]].=fma.(v1n[nidx[ndx]:nidx[ndx+1]],sigma[1,ndx],mu[1,ndx])
    v2n[nidx[ndx]:nidx[ndx+1]].=fma.(v2n[nidx[ndx]:nidx[ndx+1]],sigma[2,ndx],mu[2,ndx])
    v3n[nidx[ndx]:nidx[ndx+1]].=fma.(v3n[nidx[ndx]:nidx[ndx+1]],sigma[3,ndx],mu[3,ndx])
  end

  function p(v1,v2,v3)

    fun=similar(v1)
    fun.=0
     for ndx=1:Ns
       fun+=exp.(- (v1-mu[1,ndx]).^2./sigma[1,ndx]^2/2.
      -(v2-mu[2,ndx]).^2./sigma[2,ndx]^2/2.
     -(v3-mu[3,ndx]).^2./sigma[3,ndx]^2/2.  )*
         (2*pi)^(-1.5)/prod(sigma[:,ndx])
     end
    return fun
  end

  return p
end

d = Dict("a"=>1, "b"=>2, "c"=>3);

"""
Linear Landau Damping (Langmuir Wave)
# Arguments
- epsilon: default 0.5 for strong landau damping, set 0.05 for weak (linear)
"""
function landau{T}(x1n::Array{T,1},x2n,x3n,v1n,v2n,v3n; epsilon=0.5,
                   sigma=ones(T,3),
                   k::Array{T,1}=[0.5,0.5,0.5], q=-1., m=1.)

    L=ones(T,3)*2*pi./k;
    g=maxwellian_6d(x1n,x2n,x3n,v1n,v2n,v3n,L,sigma, zeros(T,3))

    f(x1,x2,x3,v1,v2,v3)=(1. + epsilon*cos.(k[1]*x1)).*
                          g(x1,x2,x3,v1,v2,v3)*prod(L)

    # Default parameters
    params=Dict("dt"=>0.1,"tmax"=>25, "q"=>q, "m"=>m, "c"=>1.,
                 "sigma"=>sigma, "k"=>k, "epsilon"=>epsilon,"L"=>L,
                  "f0"=>f, "g0"=>g)
    return params
end

"""
Jeans instability
"""
function jeans{T}(x1n::Array{T,1},x2n,x3n,v1n,v2n,v3n; epsilon=0.1,
                   sigma=ones(T,3),
                   k::Array{T,1}=[0.5,0.5,0.5], q=-1.,m=-1.)
    @assert m<0
    landau(x1n,x2n,x3n,v1n,v2n,v3n, epsilon=epsilon,sigma=sigma,k=k,q=q,m=m)
end

"Weibel instability in six dimensions"
function weibel{T}(x1n,x2n,x3n,v1n,v2n,v3n;
                   epsilon=1e-3, beta=-1e-3 ,
                   sigma= [ 0.02/sqrt(2), 0.02*sqrt(12/2),0.02*sqrt(12/2)],
                   k::Array{T,1}=[1.25,1.25,1.25])
  L=ones(T,3)*2*pi./k;

  g=maxwellian_6d(x1n,x2n,x3n,v1n,v2n,v3n,L,sigma, zeros(T,3))

  f(x1,x2,x3,v1,v2,v3)=(1. + epsilon*sin.(k[1]*x1)).*g(x1,x2,x3,v1,v2,v3)*prod(L)

  B3(x1,x2,x3)=real(beta)*cos.(k[1]*x1)+imag(beta)*sin.(k[1]*x1);

  # Default parameters
  params=Dict("dt"=>0.2,"tmax"=>150, "q"=>-1., "m"=>1., "c"=>1.,
               "sigma"=>sigma, "k"=>k, "epsilon"=>epsilon,"L"=>L,
                "B"=>(nothing,nothing,B3 ),
                "f0"=>f, "g0"=>g)

  return params
end

"Weibel streaming instability in six dimensions"
function weibel_streaming{T}(x1n::Array{T,1},x2n,x3n,v1n,v2n,v3n;
                   epsilon::T=T(0), beta=-im*1e-3,
                   sigma::Array{T,1}= 0.1./sqrt(2).*ones(T,3),
                   k::Array{T,1}=[0.2;0.2;0.2],
                   delta::T=(1./6.0) )
  L=ones(T,3)*2*pi./k;
  x1n.*=L[1]
  x2n.*=L[2]
  x3n.*=L[3]

  v01=0.5;
  v02=-0.1;

  gv=maxwellian_stream_3d(v1n,v2n,v3n, [delta;1.-delta],
                       sigma.*ones(1,2),
                        [ [0. 0.]; [v01 v02]; [0. 0.]  ]  )

  g(x1,x2,x3,v1,v2,v3)=(1. .+ epsilon*sin.(k[1]*x1)).*gv(v1,v2,v3)./prod(L)
  f(x1,x2,x3,v1,v2,v3)=(1. .+ epsilon*sin.(k[1]*x1)).*gv(v1,v2,v3)

  B3(x1,x2,x3)=real(beta)*cos.(k[1]*x1)+imag(beta)*sin.(k[1]*x1);

  # Default parameters
  params=Dict("dt"=>0.1,"tmax"=>200, "q"=>-1., "m"=>1., "c"=>1.,
               "sigma"=>sigma, "k"=>k, "epsilon"=>epsilon,"L"=>L,
                "B"=>(nothing,nothing,B3 ),
                "f0"=>f, "g0"=>g)

  return params
end

"""
Two dimensional periodic Kelvin Helmholtz instability

# Arguments
 k - k=1 is the stable mode in periodic domain
     chose k<1 for instability in corresponding dimension
     domain stability in 2-3
"""
function KelvinHelmholtz{T}(x1n::Array{T,1},x2n,x3n,v1n,v2n,v3n;
                   epsilon::T=0.015,
                   k::Array{T,1}=[0.4;1.;1.],
                   q=-1.,m=1., B0::Array{T,1}=[0.;0.;10.])

  L=ones(T,3)*2*pi./k;

  g=maxwellian_6d(x1n,x2n,x3n,v1n,v2n,v3n,L)

f(x1,x2,x3,v1,v2,v3)=(1+ sin.(k[2].*x2)+epsilon.*cos.(k[1].*x1) ).*
                        exp.(-0.5.*(v1.^2.+v2.^2.+v3.^2))./sqrt(2*pi).^3;

# Default parameters
params=Dict("dt"=>0.1/norm(B0),"tmax"=>20*norm(B0), "q"=>q, "m"=>m, "c"=>1.,
              "k"=>k, "epsilon"=>epsilon,"L"=>L,
              "B"=>((x1,x2,x3)->B0[1],(x1,x2,x3)->B0[2],(x1,x2,x3)->B0[3]),
              "f0"=>f, "g0"=>g)

      return params

end

end


module ScoreAD

using ForwardDiff: Dual, value, partials

"""
    score_AD(p)

When `p` is real, return `1` (with the same type).

When `p` is a primitive for automatic differentiation, return ``log ∂p``, ie the
score function.

This function can be used for gradient estimation using simulation. See the
references in the README.
"""
score_AD(p::Real) = one(p)

score_AD(p::Dual{T,V}) where {T,V} = Dual{T}(one(V), partials(log(p)))

"""
    score_AD_log(ℓ)

Similar to [`score_AD`](@ref), but uses the log likelihood, ie
```julia
score_AD(p) == score_AD_log(log(p))
```
"""
score_AD_log(ℓ::Real) = one(ℓ)

score_AD_log(ℓ::Dual{T,V}) where {T,V} = Dual{T}(one(V), partials(ℓ))

"""
    reject_nonfinite(ℓ)

When the likelihood `ℓ` (or its derivative, when applicable) is not finite,
replace the value with `-Inf`.

Always return the same type as `ℓ`.

# Usage

This useful when calculations using automatic differentiation of a score
function in a (finite) sample, which may deliver non-finite derivatives for
extreme values. In this case, reject the parameter.
"""
@inline reject_nonfinite(ℓ::V) where {V <: Real} =
    ifelse(isfinite(ℓ), ℓ, V(-Inf))

@inline function reject_nonfinite(ℓ::Dual{T,V}) where {T,V}
    p = partials(ℓ)
    if isfinite(value(ℓ)) && all(isfinite, p)
        ℓ
    else
        Dual{T}(V(-Inf), p)
    end
end

end # module


abstract type SKCE <: CalibrationErrorEstimator end

"""
    unsafe_skce_eval(k, p, y, p̃, ỹ)

Evaluate
```math
k((p, y), (p̃, ỹ)) - E_{z ∼ p}[k((p, z), (p̃, ỹ))] - E_{z̃ ∼ p̃}[k((p, y), (p̃, z̃))] + E_{z ∼ p, z̃ ∼ p̃}[k((p, z), (p̃, z̃))]
```
for kernel `k` and predictions `p` and `p̃` with corresponding targets `y` and `ỹ`.

This method assumes that `p`, `p̃`, `y`, and `ỹ` are valid and specified correctly, and
does not perform any checks.
"""
function unsafe_skce_eval end

# default implementation for classification
# we do not use the symmetry of `kernel` since it seems unlikely that `(p, y) == (p̃, ỹ)`
function unsafe_skce_eval(
    kernel::Kernel,
    p::AbstractVector{<:Real},
    y::Integer,
    p̃::AbstractVector{<:Real},
    ỹ::Integer,
)
    # precomputations
    n = length(p)

    @inbounds py = p[y]
    @inbounds p̃ỹ = p̃[ỹ]
    pym1 = py - 1
    p̃ỹm1 = p̃ỹ - 1

    tuple_p_y = (p, y)
    tuple_p̃_ỹ = (p̃, ỹ)

    # i = y, j = ỹ
    result = kernel((p, y), (p̃, ỹ)) * (1 - py - p̃ỹ + py * p̃ỹ)

    # i < y
    for i in 1:(y - 1)
        @inbounds pi = p[i]
        tuple_p_i = (p, i)

        # j < ỹ
        @inbounds for j in 1:(ỹ - 1)
            result += kernel(tuple_p_i, (p̃, j)) * pi * p̃[j]
        end

        # j = ỹ
        result += kernel(tuple_p_i, tuple_p̃_ỹ) * pi * p̃ỹm1

        # j > ỹ
        @inbounds for j in (ỹ + 1):n
            result += kernel(tuple_p_i, (p̃, j)) * pi * p̃[j]
        end
    end

    # i = y, j < ỹ
    @inbounds for j in 1:(ỹ - 1)
        result += kernel(tuple_p_y, (p̃, j)) * pym1 * p̃[j]
    end

    # i = y, j > ỹ
    @inbounds for j in (ỹ + 1):n
        result += kernel(tuple_p_y, (p̃, j)) * pym1 * p̃[j]
    end

    # i > y
    for i in (y + 1):n
        @inbounds pi = p[i]
        tuple_p_i = (p, i)

        # j < ỹ
        @inbounds for j in 1:(ỹ - 1)
            result += kernel(tuple_p_i, (p̃, j)) * pi * p̃[j]
        end

        # j = ỹ
        result += kernel(tuple_p_i, tuple_p̃_ỹ) * pi * p̃ỹm1

        # j > ỹ
        @inbounds for j in (ỹ + 1):n
            result += kernel(tuple_p_i, (p̃, j)) * pi * p̃[j]
        end
    end

    return result
end

# for binary classification with probabilities (corresponding to parameters of Bernoulli
# distributions) and boolean targets the expression simplifies to
# ```math
# k((p, y), (p̃, ỹ)) = (y(1-p) + (1-y)p)(ỹ(1-p̃) + (1-ỹ)p̃)(k((p, y), (p̃, ỹ)) - k((p, 1-y), (p̃, ỹ)) - k((p, y), (p̃, 1-ỹ)) + k((p, 1-y), (p̃, 1-ỹ)))
# ```
function unsafe_skce_eval(kernel::Kernel, p::Real, y::Bool, p̃::Real, ỹ::Bool)
    noty = !y
    notỹ = !ỹ
    z =
        kernel((p, y), (p̃, ỹ)) - kernel((p, noty), (p̃, ỹ)) -
        kernel((p, y), (p̃, notỹ)) + kernel((p, noty), (p̃, notỹ))
    return (y ? 1 - p : p) * (ỹ ? 1 - p̃ : p̃) * z
end

# evaluation for tensor product kernels
function unsafe_skce_eval(kernel::KernelTensorProduct, p, y, p̃, ỹ)
    κpredictions, κtargets = kernel.kernels
    return κpredictions(p, p̃) * unsafe_skce_eval_targets(κtargets, p, y, p̃, ỹ)
end

# resolve method ambiguity
function unsafe_skce_eval(
    kernel::KernelTensorProduct,
    p::AbstractVector{<:Real},
    y::Integer,
    p̃::AbstractVector{<:Real},
    ỹ::Integer,
)
    κpredictions, κtargets = kernel.kernels
    return κpredictions(p, p̃) * unsafe_skce_eval_targets(κtargets, p, y, p̃, ỹ)
end
function unsafe_skce_eval(kernel::KernelTensorProduct, p::Real, y::Bool, p̃::Real, ỹ::Bool)
    κpredictions, κtargets = kernel.kernels
    return κpredictions(p, p̃) * unsafe_skce_eval_targets(κtargets, p, y, p̃, ỹ)
end

function unsafe_skce_eval_targets(
    κtargets::Kernel,
    p::AbstractVector{<:Real},
    y::Integer,
    p̃::AbstractVector{<:Real},
    ỹ::Integer,
)
    # ensure that y ≤ ỹ (simplifies the implementation)
    y > ỹ && return unsafe_skce_eval_targets(κtargets, p̃, ỹ, p, y)

    # precomputations
    n = length(p)

    @inbounds begin
        py = p[y]
        pỹ = p[ỹ]
        p̃y = p̃[y]
        p̃ỹ = p̃[ỹ]
    end
    pym1 = py - 1
    pỹm1 = pỹ - 1
    p̃ym1 = p̃y - 1
    p̃ỹm1 = p̃ỹ - 1

    # i = y, j = ỹ
    result = κtargets(y, ỹ) * (1 - py - p̃ỹ + py * p̃ỹ)

    # i < y
    for i in 1:(y - 1)
        @inbounds pi = p[i]
        @inbounds p̃i = p̃[i]

        # i = j < y ≤ ỹ
        result += κtargets(i, i) * pi * p̃i

        # i < j < y ≤ ỹ
        @inbounds for j in (i + 1):(y - 1)
            result += κtargets(i, j) * (pi * p̃[j] + p[j] * p̃i)
        end

        # i < y < j < ỹ
        @inbounds for j in (y + 1):(ỹ - 1)
            result += κtargets(i, j) * (pi * p̃[j] + p[j] * p̃i)
        end

        # i < y ≤ ỹ < j
        @inbounds for j in (ỹ + 1):n
            result += κtargets(i, j) * (pi * p̃[j] + p[j] * p̃i)
        end
    end

    # y < i < ỹ
    for i in (y + 1):(ỹ - 1)
        @inbounds pi = p[i]
        @inbounds p̃i = p̃[i]

        # y < i = j < ỹ
        result += κtargets(i, i) * pi * p̃i

        # y < i < j < ỹ
        @inbounds for j in (i + 1):(ỹ - 1)
            result += κtargets(i, j) * (pi * p̃[j] + p[j] * p̃i)
        end

        # y < i < ỹ < j
        @inbounds for j in (ỹ + 1):n
            result += κtargets(i, j) * (pi * p̃[j] + p[j] * p̃i)
        end
    end

    # ỹ < i
    for i in (ỹ + 1):n
        @inbounds pi = p[i]
        @inbounds p̃i = p̃[i]

        # ỹ < i = j
        result += κtargets(i, i) * pi * p̃i

        # ỹ < i < j
        @inbounds for j in (i + 1):n
            result += κtargets(i, j) * (pi * p̃[j] + p[j] * p̃i)
        end
    end

    # handle special case y = ỹ
    if y == ỹ
        # i < y = ỹ, j = y = ỹ
        @inbounds for i in 1:(y - 1)
            result += κtargets(i, y) * (p[i] * p̃ym1 + pym1 * p̃[i])
        end

        # i = y = ỹ, j > y = ỹ
        @inbounds for j in (y + 1):n
            result += κtargets(y, j) * (pym1 * p̃[j] + p[j] * p̃ym1)
        end
    else
        # i < y
        for i in 1:(y - 1)
            @inbounds pi = p[i]
            @inbounds p̃i = p̃[i]

            # j = y < ỹ
            result += κtargets(i, y) * (pi * p̃y + pym1 * p̃i)

            # y < j = ỹ
            result += κtargets(i, ỹ) * (pi * p̃ỹm1 + pỹ * p̃i)
        end

        # i = y = j < ỹ
        result += κtargets(y, y) * pym1 * p̃y

        # i = y < j < ỹ and y < i < j = ỹ
        for ij in (y + 1):(ỹ - 1)
            @inbounds pij = p[ij]
            @inbounds p̃ij = p̃[ij]

            # i = y < j < ỹ
            result += κtargets(y, ij) * (pym1 * p̃ij + pij * p̃y)

            # y < i < j = ỹ
            result += κtargets(ij, ỹ) * (pij * p̃ỹm1 + pỹ * p̃ij)
        end

        # i = ỹ = j
        result += κtargets(ỹ, ỹ) * pỹ * (p̃ỹ - 1)

        # i = y < ỹ < j and i = ỹ < j
        for j in (ỹ + 1):n
            @inbounds pj = p[j]
            @inbounds p̃j = p̃[j]

            # i = y < ỹ < j
            result += κtargets(y, j) * (pym1 * p̃j + pj * p̃y)

            # i = ỹ < j
            result += κtargets(ỹ, j) * (p̃ỹm1 * pj + p̃j * pỹ)
        end
    end

    return result
end

function unsafe_skce_eval_targets(
    ::WhiteKernel,
    p::AbstractVector{<:Real},
    y::Integer,
    p̃::AbstractVector{<:Real},
    ỹ::Integer,
)
    @inbounds res = (y == ỹ) - p[ỹ] - p̃[y] + dot(p, p̃)
    return res
end

function unsafe_skce_eval_targets(κtargets::Kernel, p::Real, y::Bool, p̃::Real, ỹ::Bool)
    noty = !y
    notỹ = !ỹ
    z = κtargets(y, ỹ) - κtargets(noty, ỹ) - κtargets(y, notỹ) + κtargets(noty, notỹ)
    return (y ? 1 - p : p) * (ỹ ? 1 - p̃ : p̃) * z
end
function unsafe_skce_eval_targets(::WhiteKernel, p::Real, y::Bool, p̃::Real, ỹ::Bool)
    return 2 * (y - p) * (ỹ - p̃)
end


using FileIO: load
import LibSndFile
using SampledSignals
using FFTW
using ArgParse

function parse_commandline()
    s = ArgParseSettings()
    @add_arg_table s begin
        "--within"
            help = "video/audio as input to be aligned with --with"
            required = true
        "--find-offset-of"
            help = "another video/audio file"
            required = true
    end

    return parse_args(s)
end

function main()
    # Parse the program parameters
    parsed_args = parse_commandline()
    input_filepath = parsed_args["within"]
    find_offest_input_filepath = parsed_args["find-offset-of"]

    # Convert the input files to ogg, samplerate of 16KHz
    convert_to_audio(input_filepath)
    convert_to_audio(find_offest_input_filepath)

    # Load the audio files
    audio_1 = load(audio_filepath(input_filepath))
    audio_2 = load(audio_filepath(find_offest_input_filepath))

    # Pad the signals
    audio_max_size = 2 * max(length(audio_2), length(audio_1))

    audio_1_padded = padded(audio_1, audio_max_size)
    audio_2_padded = padded(audio_2, audio_max_size)
    audio_2_padded_reversed = reverse(audio_2_padded, dims=1)

    # Perform the convolutions
    convolution = irfft(rfft(audio_1_padded) .* rfft(audio_2_padded_reversed), length(audio_1_padded))

    # Find the lag in seconds
    peak_bin, peak_value = argmax(convolution[:, 1]), maximum(convolution[:, 1])
    lag_seconds = peak_bin / audio_1.samplerate
    # The FFT performs a ciruclar convolution, so the peak_value can be on the second half of the padded audio signal
    # which would be longer than the audio file itself. Ensure the offset is lower than the length of the audio file.
    convolution_length_seconds =  length(convolution[:, 1]) / audio_1.samplerate
    lag_seconds = min(lag_seconds, convolution_length_seconds - lag_seconds)

    # Improve the lag
    p = interpolate(convolution[:, 1], peak_bin)
    lag_seconds_improved = lag_seconds + p / audio_1.samplerate

    println(lag_seconds_improved)
end

function convert_command(filepath)
    output_filepath = audio_filepath(filepath)
    return `ffmpeg -i $filepath -vn -ar 16000 -ac 1 $output_filepath -n`
end

function audio_filepath(filepath)
    # Keep the actual filename, but not the folders
    parts = rsplit(filepath, "/"; limit=2)
    return "$(tempdir())/$(parts[length(parts)]).ogg"
end

function convert_to_audio(filepath)
    if (!isfile(audio_filepath(filepath)))
        run(convert_command(filepath))
    end
end

function padded(audio, expected_size)
    pad_size = expected_size - length(audio)
    audio_pad = SampledSignals.SampleBuf(zeros(Float32, pad_size, 1), audio.samplerate)
    audio_padded = vcat(audio, audio_pad)

    return audio_padded
end

function interpolate(signal, peak_bin)
    # Quadratic Interpolation of Spectral Peaks
    # https://ccrma.stanford.edu/~jos/sasp/Quadratic_Interpolation_Spectral_Peaks.html
    α = signal[peak_bin - 1]
    β = signal[peak_bin]
    γ = signal[peak_bin + 1]

    p = 1 / 2 * (α - γ) / (α - 2 * β + γ)
return p
end

main()


using Convex, MultiConvex, Compat, SCS
set_default_solver(SCSSolver(verbose=false))
using Base.Test
import Convex.hash

x = Variable()
y = Variable()
z = Variable()
x.value = 1
y.value = 2
z.value = 3

# variables in
@assert length(variablesin(x*y*z)) == 3
@assert length(variablesin(x*y+y*z+x+y+z)) == 3
@assert length(variablesin(log(x) + 7*x + 5)) == 1

# partitioning sets of integers
n = Dict{Int, Set{Int}}({1=>Set{Int}({2,3,4}),
	                             2=>Set{Int}({1,3}),
	                             3=>Set{Int}({1,2}),
	                             4=>Set{Int}({1}),
	                             5=>Set{Int}({6,7}),
	                             6=>Set{Int}({5,7}),
	                             7=>Set{Int}({6,5})})
p = stablepartition!(n)
@assert length(p) == 3
@assert sum(map(length, p)) == 7
@assert length(union(p...)) == 7

# partitioning given initial sets
p = stablepartition!(n, Set{Int}[Set{Int}({1,5})])
@assert length(p) == 3
@assert sum(map(length, p)) == 7
@assert length(union(p...)) == 7
@assert p[1] == Set({1,5})

# multivexity: (convex) addition
conflicts = Set{@compat Tuple{Uint64, Uint64}}()
@assert multivexity(x+y, conflicts) == AffineVexity()
@assert length(conflicts) == 0

# multivexity: (biconvex) multiplication
conflicts = Set{@compat Tuple{Uint64, Uint64}}()
v = multivexity(x*y, conflicts)
@assert v == AffineVexity()
@assert length(conflicts) == 1
@assert pop!(conflicts) == (x.id_hash, y.id_hash)
# note hash(x::Variable) is not defined to grab id_hash except *internally* in the multiconvex module

# multivexity: addition and multiplication
conflicts = Set{@compat Tuple{Uint64, Uint64}}()
v = multivexity(x*y + y*z + z*x, conflicts)
@assert v == AffineVexity()
@assert length(conflicts) == 3

# partitioning variables given conflict graphs
conflicts = Set{@compat Tuple{Uint64, Uint64}}()
m = x*y + y*z + z*x
v = multivexity(m, conflicts)
vars = variablesin(m)
parts = stablepartition!(conflicts)
@assert length(parts) == 3
@assert all(map(x -> length(x) == 1, parts))

# multivexity: problems with no constraints
x = Variable()
y = Variable()
z = Variable()
problem = minimize(exp(x*y*z) + x^2 + y^2 + z^2)
vex, stablesetvars, vars = multivexity(problem)
@assert length(vars) == 3
@assert length(stablesetvars) == 3
@assert sum(map(length, stablesetvars)) == 3
@assert vex == ConvexVexity()

# multivexity: problems with constraints
x = Variable()
y = Variable()
z = Variable()
w = Variable()
problem = minimize(exp(x*z) + exp(w*y) + x^2 + y^2 + z^2, x+y >= 0)
vex, stablesetvars, vars = multivexity(problem)
@assert length(vars) == 4
@assert length(stablesetvars) == 2
@assert map(length, stablesetvars) == [2, 2]
@assert vex == ConvexVexity()

# solving a simple problem
x = Variable()
y = Variable()
z = Variable()
# solution to this problem should be (0,0,0)
problem = minimize(exp(x*y*z) + x^2 + y^2 + z^2, x>=0, y>=0, z>=0)
altmin!(problem)
@show t1 = (x.value, y.value, z.value)
# solution to this problem should be closer to (some permutation of) (0,1,1)
problem = minimize(exp(x*y*z) + 5*(1-x)^2 + 5*(1-y)^2 + 5*(1-z)^2, x>=0, y>=0, z>=0)
altmin!(problem, warmstart = false, maxAMiters=100)
@show t2 = (x.value[1], y.value[1], z.value[1])
@assert sort([t2...])[2] > maximum(t1)


abstract type EGOAquisition end

update_parameters!(::EGOAquisition, args...) = error("NOT IMPLEMENTED")
objective(::EGOAquisition, m::Kriging, x) = error("NOT IMPLEMENTED")

import StatsFuns

const φ = StatsFuns.normpdf
const Φ = StatsFuns.normcdf

include("single_obj_aquisitions.jl")
include("constrained_aquisition.jl")


import GaussQuadrature.legendre
using FractionalTimeDG
using ArgCheck
using OffsetArrays

#include("FEM1D.jl")

const L = 2.0
const max_t = 2.0
const C0 = 1.0
const Cf = 2.0
α = 0.6

u0(x) = C0 * x * (L-x)
f(x,t) = Cf * t * exp(-t)

function refsoln(x, t, α, N)
    return Bromich_integral(t, z -> uhat(x, z, α), N)
end

function uhat(x::T, z::Complex{T}, α::T) where T <: AbstractFloat
    ω = z^(α/2)
    ρ1(x) = ( ω * x * (L-x) - 2/ω ) * cosh(ω*x) + (2x-L) * sinh(ω*x) + 2 / ω
    ρ2(x) = cosh(ω*x) - 1
    return ( (C0/ω) * ( ρ1(x) * sinh(ω*(L-x)) + ρ1(L-x) * sinh(ω*x) )
           + Cf/(z+1)^2 * ( ρ2(x) * sinh(ω*(L-x)) + ρ2(L-x) * sinh(ω*x) )
          ) / ( z * sinh(ω*L) )
end

"""
    Bromich_integral(t, F, N)

Evaluates `1/(2πi)` times the integral of `exp(zt)F(z) dz` along a
Hankel contour, assuming `F(z)` is analytic in the cut plane
`|arg(z)|<π` and `F(conj(z)) = conj(F(z))`.
"""
function Bromich_integral(t::Float64, F::Function, N::Integer)
    μ = 4.492075 * N / t
    ϕ = 1.172104
    h = 1.081792 / N
    z0, w0 = integration_contour(0.0, μ, ϕ)
    Σ = real( w0 * exp(z0*t) * F(z0) ) / 2
    for n = 1:N
        zn, wn = integration_contour(n*h, μ, ϕ)
        Σ += real( wn * exp(zn*t) * F(zn) )
    end
    return h * Σ
end

function integration_contour(u, μ, ϕ)
    z = μ * ( 1 + sin(Complex(-ϕ, u)) )
    w = (μ/π) * cos(Complex(-ϕ, u))
    return z, w
end

function L2error(U::Array{Float64}, u::Function, t::OffsetArray{Float64},
                 τ::AbstractVector{Float64},
                 grid::SpatialGrid{Float64}, Mx::Integer)
    ξ, w = legendre(Mx)
    xvals, pcwise_t, pcwise_U = evaluate_pcwise_poly(U, t, grid, τ, ξ)
    pcwise_err = evaluate_pcwise_error(pcwise_U, u, pcwise_t, xvals)

    x = grid.x
    pts_per_time_interval = length(τ)
    Nt = length(t) - 1
    Nx = length(x) - 1
    err = Array{Float64}(undef, pts_per_time_interval, Nt)
    for nt = 1:Nt
        for i = 1:pts_per_time_interval
            m = 0
            s = 0.0
            for nx = 1:Nx
                hx = x[nx] - x[nx-1]
                for j = 1:Mx
                    m += 1
                    s += (hx/2) * w[j] * pcwise_err[m,i,nt]^2
                end
            end
            err[i,nt] = sqrt(s)
        end
    end
    return pcwise_t, err
end

function evaluate_jumps(U::Array{Float64}, U0::Vector{Float64},
                        t::OffsetVector{Float64},
                        grid::SpatialGrid{Float64}, Mx::Integer)
    x = grid.x
    Nx = length(x) - 1
    ξ, w = legendre(Mx)
    τ = [-1.0, 1.0]
    xvals, pcwise_t, pcwise_U = evaluate_pcwise_poly(U, t, grid, τ, ξ)
    U0vals = similar(xvals)
    evaluate_fem1d_soln!(U0vals, U0, grid, ξ)
    Nt = length(t) - 1
    jump = OffsetArray{Float64}(undef, 0:Nt-1)
    m = 0
    s = 0.0
    for nx = 1:Nx
        hx = x[nx] - x[nx-1]
        for j = 1:Mx
            m += 1
            s += (hx/2) * w[j] * ( pcwise_U[m,1,1] - U0vals[m] )^2
        end
    end
    jump[0] = sqrt(s)
    for nt = 1:Nt-1
        m = 0
        s = 0.0
        for nx = 1:Nx
            hx = x[nx] - x[nx-1]
            for j = 1:Mx
                m += 1
                s += (hx/2) * w[j] * ( pcwise_U[m,1,nt+1] - pcwise_U[m,2,nt] )^2
            end
        end
        jump[nt] = sqrt(s)
    end
    return jump
end

function reconstruction(U::Array{Float64}, U0::Vector{Float64},
                        store::Store{Float64})
    NgDoF, rt, Nt = size(U)
    Uhat = Array{Float64}(undef, NgDoF, rt+1, Nt)
    jumpU = OffsetArray{Float64}(undef, 1:NgDoF, 0:Nt-1)
    rmax = store.rmax
    pow = OffsetArray{Float64}(undef, 0:rmax)
    pow[0:2:end] .=  1.0
    pow[1:2:end] .= -1.0
    @argcheck rt+1 ≤ rmax
    U_left  = zeros(NgDoF)
    U_right = zeros(NgDoF)
    for j = 1:rt
        U_left .+= pow[j-1] * U[:,j,1]
        U_right .+= U[:,j,1]
    end
    jumpU[:,0] .= U_left - U0
    Uhat[:,1:rt,1] .= U[:,1:rt,1]
    Uhat[:,rt,1] .+= pow[rt] * jumpU[:,0] / 2
    Uhat[:,rt+1,1] .= - pow[rt] * jumpU[:,0] / 2
    for n = 2:Nt
        fill!(U_left, 0.0)
        for j = 1:rt
            U_left .+= pow[j-1] * U[:,j,n]
        end
        jumpU[:,n-1] .= U_left - U_right
        Uhat[:,1:rt,n] .= U[:,1:rt,n]
        Uhat[:,rt,n]  .+=  pow[rt] * jumpU[:,n-1] / 2
        Uhat[:,rt+1,n] .= -pow[rt] * jumpU[:,n-1] / 2
        fill!(U_right, 0.0)
        for j = 1:rt
            U_right .+= U[:,j,n]
        end
    end
    return Uhat, jumpU
end


push!(LOAD_PATH, string(pwd(),"/src"))
using ModifiedStokesSolver
using Base.Test

@testset "FMM" begin
    ## define source and target points
    ns = 100
    nt = 102

    srand(1)
    src = 1 - 2*rand(2,ns)
    fvec = 1 - 2*rand(2,ns)
    nvec = 1 - 2*rand(2,ns)
    targ = 1 - 2*rand(2,nt)

    # Compute using the kernels from ModifiedStokesSolver
    function mss_stokeslet_direct(src, targ, fvec, alpha)
        ns = size(src, 2)
        nt = size(targ, 2)
        u = zeros(2, nt)
        for i=1:ns
            for j=1:nt
                f = fvec[:, i]
                r = src[:,i]-targ[:,j]
                u[:,j] += stokeslet(r, f, alpha)
            end
        end
        return u
    end
    function mss_stresslet_direct(src, targ, fvec, nvec, alpha)
        ns = size(src, 2)
        nt = size(targ, 2)
        u = zeros(2, nt)
        for i=1:ns
            for j=1:nt
                r = src[:,i]-targ[:,j]
                u[:,j] += stresslet(r, fvec[:,i], nvec[:,i], alpha)
            end
        end
        return u
    end
    function mss_stresslet_direct_self(src, fvec, nvec, alpha)
        ns = size(src, 2)
        u = zeros(2, ns)
        for i=1:ns
            for j=1:ns
                i==j && continue
                r = src[:,i]-src[:,j]
                u[:,j] += stresslet(r, fvec[:,i], nvec[:,i], alpha)
            end
        end
        return u
    end

    ## Test with large alpha
    alpha = 20.0

    ## Stokeslet
    u1 = mss_stokeslet_direct(src, targ, fvec, alpha)
    u2 = fmm_stokeslet_direct(src, targ, fvec, alpha)
    u3 = fmm_stokeslet_targ(src, targ, fvec, alpha)
    println(" * Stokeslet, alpha=", alpha)
    Gerrdirect = norm(u1[:]-u2[:]) / norm(u1[:], Inf)
    Gerrfmm = norm(u1[:]-u3[:]) / norm(u1[:], Inf)
    @test Gerrdirect < 5e-14
    @test Gerrfmm < 5e-14

    ## Stresslet
    ## Test also with small alpha
    ## (Stokeslet has no special eval for small args)
    for alpha=[alpha, 1e-15]
        println(" * Stresslet, alpha=", alpha)
        # Stresslet targ
        d1 = mss_stresslet_direct(src, targ, fvec, nvec, alpha)
        d2 = fmm_stresslet_direct(src, targ, fvec, nvec, alpha)
        d3 = fmm_stresslet_targ(src, targ, fvec, nvec, alpha)
        # Stresslet self
        ds1 = mss_stresslet_direct_self(src, fvec, nvec, alpha)
        ds2 = fmm_stresslet_direct(src, src, fvec, nvec, alpha; self=true)
        ds3 = fmm_stresslet_self(src, fvec, nvec, alpha)
        # Test
        Terrdirect = norm(d1[:]-d2[:]) / norm(d1[:], Inf)
        Terrdirect_self = norm(ds1[:]-ds2[:]) / norm(d1[:], Inf)
        Terrfmm = norm(d1[:]-d3[:]) / norm(d1[:], Inf)
        Terrfmm_self = norm(ds1[:]-ds3[:]) / norm(ds1[:], Inf)
        @test Terrdirect < 5e-14
        @test Terrdirect_self < 5e-14
        @test Terrfmm < 5e-14
        @test Terrfmm_self < 5e-14
    end
end


# we generate code in this module, so precompile where possible
__precompile__(true)

module Rotations

using LinearAlgebra
using StaticArrays
using Random

import Statistics

include("util.jl")
include("core_types.jl")
include("unitquaternion.jl")
include("mrps.jl")
include("euler_types.jl")
include("angleaxis_types.jl")
# TODO: Add method `mean_rotation` instead of `mean`
# include("mean.jl")
include("derivatives.jl")
include("principal_value.jl")

include("rodrigues_params.jl")
include("error_maps.jl")
include("rotation_error.jl")
include("log.jl")
include("eigen.jl")
include("deprecated.jl")

export
    Rotation, RotMatrix, RotMatrix2, RotMatrix3,
    Angle2d,
    Quat, UnitQuaternion,
    AngleAxis, RodriguesVec, RotationVec, RodriguesParam, MRP,
    RotX, RotY, RotZ,
    RotXY, RotYX, RotZX, RotXZ, RotYZ, RotZY,
    RotXYX, RotYXY, RotZXZ, RotXZX, RotYZY, RotZYZ,
    RotXYZ, RotYXZ, RotZXY, RotXZY, RotYZX, RotZYX,

    # Quaternion math ops
    logm, expm, ⊖, ⊕,

    # Quaternion maps
    ExponentialMap, QuatVecMap, CayleyMap, MRPMap, IdentityMap,

    # check validity of the rotation (is it close to unitary?)
    isrotation,

    # angle and axis introspection
    rotation_angle,
    rotation_axis,

    # quaternion from two vectors
    rotation_between,

    # principal value of a rotation
    principal_value

    # derivatives (names clash with ForwarDiff?)
    #jacobian, hessian

end # module


using Images, Luxor, BenchmarkTools

function calc(dims)
    img = fill(RGB(0,0,0), (dims, dims))
    a = 2.24; b = 0.43; c = -0.65; d = -2.43
    x = y = z = .0
    for _ in 1:dims^2
        x, y, z = sin(a * y) - z * cos(b * x), z * sin(c * x) - cos(d * y), sin(x)
        xx = trunc(Int, rescale(x, -2., 2., 0, dims))
        yy = trunc(Int, rescale(y, -2., 2., 0, dims))
        zz = rescale(z, -1, 1, .3, .8)
        img[yy,xx] = RGB(.7xx / dims, .3yy / dims + .2, zz)
    end
    return img
end

@time save("strange.png",calc(1000))


using Test
using AtiyahBott

# first test
println( "Testing O1" )
P = (g,c,w,s,m) -> O1(g,c,w,s,m)^2
res = AtiyahBottFormula(2,2,6-1,P,true,true,false)
if res[1] == 1//1
    @test true
else
    @test false
end

#second test
println( "Testing Incidency and Contact" )
d = 4
P = Vector{Any}(undef, d+1)
for a in 0:d
    P[a+1] = (g,c,w,s,m) -> Incidency(g,c,w,s,3)^a*
        Incidency(g,c,w,s,2)^(2*(d-a)+1)*
        Contact(g,c,w,s)
end
res = AtiyahBottFormula(3,d,0,P,true,true,false);

if res == [ 1089024, 96512, 9408, 1024, 128 ]
    @test true
else
    @test false
end

#third test
println( "Testing Hypersurface" )
d = 4 #for other values of d, change this line
P = (g,c,w,s,m) -> Hypersurface(g,c,w,s,5)
res = AtiyahBottFormula(4,d,0,P,true,true,false);

if res == [ 15517926796875//64 ]
    @test true
else
    @test false
end

#fourth test
println( "Testing R1" )
d = 4 #for other values of d, change this line
P = (g,c,w,s,m) -> R1(g,c,w,s,1)^2
res = AtiyahBottFormula(1,d,0,P,true,true,false)

if res == [ 1//64 ]
    @test true
else
    @test false
end


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

using JuliaFEM, Test

# Example of 2d nonlinear elasticity with non-homogeneous boundary contitions

# Geometry of nodes
X = Dict(1 => [0.0, 0.0], 2 => [1.0, 0.0], 3 => [1.0, 1.0], 4 => [0.0, 1.0])

block = Problem(Elasticity, "block", 2)
block.properties.formulation = :plane_stress
block.properties.finite_strain = true
block.properties.geometric_stiffness = true

element = Element(Quad4, (1, 2, 3, 4))
update!(element, "geometry", X)
update!(element, "youngs modulus", 288.0)
update!(element, "poissons ratio", 1/3)
add_element!(block, element)

# boundary conditions
bc = Problem(Dirichlet, "bc", 2, "displacement")
bel1 = Element(Seg2, (1, 2))
bel2 = Element(Seg2, (3, 4))
bel3 = Element(Seg2, (4, 1))
update!((bel1, bel2, bel3), "geometry", X)
update!(bel1, "displacement 2", 0.0)
update!(bel2, "displacement 2", 0.5)
update!(bel3, "displacement 1", 0.0)
add_elements!(bc, bel1, bel2, bel3)

analysis = Analysis(Nonlinear)
add_problems!(analysis, block, bc)
run!(analysis)

# results are verified using Code Aster
eps_expected = [-2.08333312468287E-01, 6.25000000000000E-01, 0.0]
sig_expected = [ 4.50685020821470E-06, 4.62857140373777E+02, 0.0]
u3_expected =  [-2.36237356855269E-01, 5.00000000000000E-01]

u3 = block("displacement", 0.0)[3]
@test isapprox(u3, u3_expected, atol=1.0e-5)


function projection_onto_bounds_ipopt_style( nlp::Class_CUTEst, pars::Class_parameters, x::Array{Float64,1} )
    if pars.output_level >= 4
        println("projecting onto bounds ...")
    end
    ifree = _i_not_fixed(nlp.nlp)
    uvar = deepcopy(nlp.nlp.meta.uvar[ifree])
    lvar =  deepcopy(nlp.nlp.meta.lvar[ifree])

    x = deepcopy(x)
    b_L = zeros(length(x))
    b_U = zeros(length(x))

    κ_1 = 1e-2;
    κ_2 = 1e-2;

    threshold_too_close = 1e-8

    for i = 1:length(x)
      if (uvar[i] < lvar[i] + threshold_too_close)
        @show uvar[i] - lvar[i]
        println("ERROR: bounds too close!")
        error("bounds too close!")
      end

      p_L = min(κ_1 * max(1, abs(lvar[i])), κ_2 * (uvar[i] - lvar[i]) )
      p_U = min(κ_1 * max(1, abs(uvar[i])), κ_2 * (uvar[i] - lvar[i]) )

      if lvar[i] > -Inf
        b_L[i] = lvar[i] + p_L
      else
        b_L[i] = -Inf
      end

      if uvar[i] < Inf
        b_U[i] = uvar[i] - p_U
      else
        b_U[i] = Inf
      end

      x[i] = x[i] #+ rand() * 100.0

      if b_L[i] < b_U[i]
        x[i] = min(max(b_L[i], x[i]), b_U[i])
      else
        @show b_L[i], b_U[i]
        error("this shouldn't happen!")
        #x[i] = (uvar[i] + lvar[i]) / 2.0
      end

      if(x[i] <= lvar[i] || x[i] >= uvar[i])
         @show (b_L[i], b_U[i]), (lvar[i], uvar[i]) , x[i]
         error("Error in projection_onto_bounds")
      end
    end

    return x
end


using Plots;
gr();
using LaTeXStrings

savepath = "/home/user/cameronperot.com/images/2019-12-29-simulating-magnetism-part-1";
save_bool = true
if save_bool
    dpi = 300
else
    dpi = 144
end

# %% energy plot

βs = range(0, stop = 5, length = 100)
e = -tanh.(βs);
plot_e = plot(
    βs,
    e,
    dpi = dpi,
    label = L"-\tanh\beta",
    xaxis = (L"\beta"),
    yaxis = (L"e"),
    line = (2, :dodgerblue, 0.8),
);

# %% heat capacity plot

c = βs .^ 2 ./ cosh.(βs) .^ 2;
plot_c = plot(
    βs,
    c,
    dpi = dpi,
    label = L"\beta^2 / \cosh^2 \beta",
    xaxis = (L"\beta"),
    yaxis = (L"c"),
    line = (2, :dodgerblue, 0.8),
);

# %% combined plot

plot_combined = plot(plot_e, plot_c, layout = (2, 1), size = (800, 600))

if save_bool
    savefig(plot_combined, joinpath(savepath, "1d_ising_e_c.png"))
end

plot_combined


"""
    importASCII(file::AbstractString [, outfile])

Imports an ESRI Ascii grid to a .grd file and returns a new VerySimpleRaster
object. Specify `outfile` to keep the newly created file.
"""
function importASCII(x::AbstractString, outfile = "")
    file = open(x, "r")
    nc = parse(Int, match(r"NCOLS (.+)", readline(file)).captures[1])
    nr = parse(Int, match(r"NROWS (.+)", readline(file)).captures[1])
    xll = parse(Float64, match(r"XLLCORNER (.+)", readline(file)).captures[1])
    yll = parse(Float64, match(r"YLLCORNER (.+)", readline(file)).captures[1])
    cell = parse(Float64, match(r"CELLSIZE (.+)", readline(file)).captures[1])
    NA = parse(Float64, match(r"NODATA_value (.+)", readline(file)).captures[1])
    tmp = Vector{Float64}(undef, nc)

    outfile == "" && (outfile = tempname())
    outfile[end-3:end] ∈ (".gri", "grd") && (outfile = outfile[1:end-4])
    IO = open(outfile*".gri", "w")

    for row in nr:-1:1
        tmp .= parse.(Float64, split(readline(file), " "))
        for r in tmp
            write(IO, r)
        end
    end
    close(file)
    close(IO)
    mat = Mmap.mmap(outfile*".gri", Matrix{Float64}, (nc, nr))
    vsr = VerySimpleRaster(mat, NA, xll, yll, cell, cell, "", outfile*".gri")
    write_grd_header(outfile*".grd", vsr)
    vsr
end


# By default, Julia/LLVM does not use fused multiply-add operations (FMAs).
# Since these FMAs can increase the performance of many numerical algorithms,
# we need to opt-in explicitly.
# See https://ranocha.de/blog/Optimizing_EC_Trixi for further details.
@muladd begin

@doc raw"""
    LatticeBoltzmannEquations2D(; Ma, Re, collision_op=collision_bgk,
                               c=1, L=1, rho0=1, u0=nothing, nu=nothing)

The Lattice-Boltzmann equations
```math
\partial_t u_\alpha + v_{\alpha,1} \partial_1 u_\alpha + v_{\alpha,2} \partial_2 u_\alpha = 0
```
in two space dimensions for the D2Q9 scheme.

The characteristic Mach number and Reynolds numbers are specified as `Ma` and `Re`. By the
default, the collision operator `collision_op` is set to the BGK model. `c`, `L`, and `rho0`
specify the mean thermal molecular velocity, the characteristic length, and the reference density,
respectively. They can usually be left to the default values. If desired, instead of the Mach
number, one can set the macroscopic reference velocity `u0` directly (`Ma` needs to be set to
`nothing` in this case). Likewise, instead of the Reynolds number one can specify the kinematic
viscosity `nu` directly (in this case, `Re` needs to be set to `nothing`).

The nine discrete velocity directions of the D2Q9 scheme are sorted as follows [4]:
```
  6     2     5       y
    ┌───┼───┐         │
    │       │         │
  3 ┼   9   ┼ 1        ──── x
    │       │        ╱
    └───┼───┘       ╱
  7     4     8    z
```
Note that usually the velocities are numbered from `0` to `8`, where `0` corresponds to the zero
velocity. Due to Julia using 1-based indexing, here we use indices from `1` to `9`, where `1`
through `8` correspond to the velocity directions in [4] and `9` is the zero velocity.

The corresponding opposite directions are:
* 1 ←→ 3
* 2 ←→ 4
* 3 ←→ 1
* 4 ←→ 2
* 5 ←→ 7
* 6 ←→ 8
* 7 ←→ 5
* 8 ←→ 6
* 9 ←→ 9

The main sources for the base implementation were
1. Misun Min, Taehun Lee, **A spectral-element discontinuous Galerkin lattice Boltzmann method for
   nearly incompressible flows**, J Comput Phys 230(1), 2011
   [doi:10.1016/j.jcp.2010.09.024](https://doi.org/10.1016/j.jcp.2010.09.024)
2. Karsten Golly, **Anwendung der Lattice-Boltzmann Discontinuous Galerkin Spectral Element Method
   (LB-DGSEM) auf laminare und turbulente nahezu inkompressible Strömungen im dreidimensionalen
   Raum**, Master Thesis, University of Cologne, 2018.
3. Dieter Hänel, **Molekulare Gasdynamik**, Springer-Verlag Berlin Heidelberg, 2004
   [doi:10.1007/3-540-35047-0](https://doi.org/10.1007/3-540-35047-0)
4. Dieter Krüger et al., **The Lattice Boltzmann Method**, Springer International Publishing, 2017
   [doi:10.1007/978-3-319-44649-3](https://doi.org/10.1007/978-3-319-44649-3)
"""
struct LatticeBoltzmannEquations2D{RealT<:Real, CollisionOp} <: AbstractLatticeBoltzmannEquations{2, 9}
  c::RealT    # mean thermal molecular velocity
  c_s::RealT  # isothermal speed of sound
  rho0::RealT # macroscopic reference density

  Ma::RealT   # characteristic Mach number
  u0::RealT   # macroscopic reference velocity

  Re::RealT   # characteristic Reynolds number
  L::RealT    # reference length
  nu::RealT   # kinematic viscosity

  weights::SVector{9, RealT}  # weighting factors for the equilibrium distribution
  v_alpha1::SVector{9, RealT} # discrete molecular velocity components in x-direction
  v_alpha2::SVector{9, RealT} # discrete molecular velocity components in y-direction

  collision_op::CollisionOp   # collision operator for the collision kernel
end

function LatticeBoltzmannEquations2D(; Ma, Re, collision_op=collision_bgk,
                                       c=1, L=1, rho0=1, u0=nothing, nu=nothing)
  # Sanity check that exactly one of Ma, u0 is not `nothing`
  if isnothing(Ma) && isnothing(u0)
    error("Mach number `Ma` and reference speed `u0` may not both be `nothing`")
  elseif !isnothing(Ma) && !isnothing(u0)
    error("Mach number `Ma` and reference speed `u0` may not both be set")
  end

  # Sanity check that exactly one of Re, nu is not `nothing`
  if isnothing(Re) && isnothing(nu)
    error("Reynolds number `Re` and visocsity `nu` may not both be `nothing`")
  elseif !isnothing(Re) && !isnothing(nu)
    error("Reynolds number `Re` and visocsity `nu` may not both be set")
  end

  # Calculate isothermal speed of sound
  # The relation between the isothermal speed of sound `c_s` and the mean thermal molecular velocity
  # `c` depends on the used phase space discretization, and is valid for D2Q9 (and others). For
  # details, see, e.g., [3] in the docstring above.
  c_s = c / sqrt(3)

  # Calculate missing quantities
  if isnothing(Ma)
    Ma = u0 / c_s
  elseif isnothing(u0)
    u0 = Ma * c_s
  end
  if isnothing(Re)
    Re = u0 * L / nu
  elseif isnothing(nu)
    nu = u0 * L / Re
  end

  # Promote to common data type
  Ma, Re, c, L, rho0, u0, nu = promote(Ma, Re, c, L, rho0, u0, nu)

  # Source for weights and speeds: [4] in the docstring above
  weights  = SVector(1/9, 1/9, 1/9, 1/9, 1/36, 1/36, 1/36, 1/36, 4/9)
  v_alpha1 = SVector( c,   0,  -c,   0,   c,   -c,   -c,    c,    0 )
  v_alpha2 = SVector( 0,   c,   0,  -c,   c,    c,   -c,   -c,    0 )

  LatticeBoltzmannEquations2D(c, c_s, rho0, Ma, u0, Re, L, nu,
                             weights, v_alpha1, v_alpha2,
                             collision_op)
end

varnames(::typeof(cons2cons), equations::LatticeBoltzmannEquations2D) = ntuple(v -> "pdf"*string(v), nvariables(equations))
varnames(::typeof(cons2prim), equations::LatticeBoltzmannEquations2D) = varnames(cons2cons, equations)

# Convert conservative variables to macroscopic
@inline function cons2macroscopic(u, equations::LatticeBoltzmannEquations2D)
  rho    = density(u, equations)
  v1, v2 = velocity(u, equations)
  p      = pressure(u, equations)
  return SVector(rho, v1, v2, p)
end
varnames(::typeof(cons2macroscopic), ::LatticeBoltzmannEquations2D) = ("rho", "v1", "v2", "p")

# Set initial conditions at physical location `x` for time `t`
"""
    initial_condition_constant(x, t, equations::LatticeBoltzmannEquations2D)

A constant initial condition to test free-stream preservation.
"""
function initial_condition_constant(x, t, equations::LatticeBoltzmannEquations2D)
  @unpack u0 = equations
  rho = pi
  v1 = u0
  v2 = u0

  return equilibrium_distribution(rho, v1, v2, equations)
end

"""
    boundary_condition_noslip_wall(u_inner, orientation, direction, x, t,
                                   surface_flux_function,
                                   equations::LatticeBoltzmannEquations2D)

No-slip wall boundary condition using the bounce-back approach.
"""
function boundary_condition_noslip_wall(u_inner, orientation, direction, x, t,
                                        surface_flux_function,
                                        equations::LatticeBoltzmannEquations2D)
  # For LBM no-slip wall boundary conditions, we set the boundary state to
  # - the inner state for outgoing particle distribution functions
  # - the *opposite* inner state for all other particle distribution functions
  # See the list of (opposite) directions in the docstring of `LatticeBoltzmannEquations2D`.
  if direction == 1 # boundary in -x direction
    pdf1 = u_inner[3]
    pdf2 = u_inner[4]
    pdf3 = u_inner[3] # outgoing
    pdf4 = u_inner[2]
    pdf5 = u_inner[7]
    pdf6 = u_inner[6] # outgoing
    pdf7 = u_inner[7] # outgoing
    pdf8 = u_inner[6]
    pdf9 = u_inner[9]
  elseif direction == 2 # boundary in +x direction
    pdf1 = u_inner[1] # outgoing
    pdf2 = u_inner[4]
    pdf3 = u_inner[1]
    pdf4 = u_inner[2]
    pdf5 = u_inner[5] # outgoing
    pdf6 = u_inner[8]
    pdf7 = u_inner[5]
    pdf8 = u_inner[8] # outgoing
    pdf9 = u_inner[9]
  elseif direction == 3 # boundary in -y direction
    pdf1 = u_inner[3]
    pdf2 = u_inner[4]
    pdf3 = u_inner[1]
    pdf4 = u_inner[4] # outgoing
    pdf5 = u_inner[7]
    pdf6 = u_inner[8]
    pdf7 = u_inner[7] # outgoing
    pdf8 = u_inner[8] # outgoing
    pdf9 = u_inner[9]
  else # boundary in +y direction
    pdf1 = u_inner[3]
    pdf2 = u_inner[2] # outgoing
    pdf3 = u_inner[1]
    pdf4 = u_inner[2]
    pdf5 = u_inner[5] # outgoing
    pdf6 = u_inner[6] # outgoing
    pdf7 = u_inner[5]
    pdf8 = u_inner[6]
    pdf9 = u_inner[9]
  end
  u_boundary = SVector(pdf1, pdf2, pdf3, pdf4, pdf5, pdf6, pdf7, pdf8, pdf9)

  # Calculate boundary flux
  if iseven(direction) # u_inner is "left" of boundary, u_boundary is "right" of boundary
    flux = surface_flux_function(u_inner, u_boundary, orientation, equations)
  else # u_boundary is "left" of boundary, u_inner is "right" of boundary
    flux = surface_flux_function(u_boundary, u_inner, orientation, equations)
  end

  return flux
end

function boundary_condition_moving_wall_ypos(u_inner, orientation, direction, x, t,
                                             surface_flux_function,
                                             equations::LatticeBoltzmannEquations2D)
  @assert direction == 4 "moving wall assumed in +y direction"

  @unpack rho0, u0, weights, c_s = equations
  cs_squared = c_s^2

  pdf1 = u_inner[3] + 2 * weights[1] * rho0 * u0 / cs_squared
  pdf2 = u_inner[2] # outgoing
  pdf3 = u_inner[1] + 2 * weights[3] * rho0 * (-u0) / cs_squared
  pdf4 = u_inner[2]
  pdf5 = u_inner[5] # outgoing
  pdf6 = u_inner[6] # outgoing
  pdf7 = u_inner[5] + 2 * weights[7] * rho0 * (-u0) / cs_squared
  pdf8 = u_inner[6] + 2 * weights[8] * rho0 * u0 / cs_squared
  pdf9 = u_inner[9]

  u_boundary = SVector(pdf1, pdf2, pdf3, pdf4, pdf5, pdf6, pdf7, pdf8, pdf9)

  # Calculate boundary flux (u_inner is "left" of boundary, u_boundary is "right" of boundary)
  flux = surface_flux_function(u_inner, u_boundary, orientation, equations)

  return flux
end

"""
    initial_condition_lid_driven_cavity(x, t, equations::LatticeBoltzmannEquations2D)

Initial state for a lid-driven cavity flow setup. To be used in combination with
[`boundary_condition_lid_driven_cavity`](@ref) and [`boundary_condition_noslip_wall`](@ref).
"""
function initial_condition_lid_driven_cavity(x, t, equations::LatticeBoltzmannEquations2D)
  @unpack L, u0, nu = equations

  rho = 1
  v1 = 0
  v2 = 0

  return equilibrium_distribution(rho, v1, v2, equations)
end

"""
    boundary_condition_lid_driven_cavity(u_inner, orientation, direction, x, t,
                                         surface_flux_function,
                                         equations::LatticeBoltzmannEquations2D)

Boundary condition for a lid-driven cavity flow setup, where the top lid (+y boundary) is a moving
no-slip wall. To be used in combination with [`initial_condition_lid_driven_cavity`](@ref).
"""
function boundary_condition_lid_driven_cavity(u_inner, orientation, direction, x, t,
                                              surface_flux_function,
                                              equations::LatticeBoltzmannEquations2D)
  return boundary_condition_moving_wall_ypos(u_inner, orientation, direction, x, t,
                                             surface_flux_function, equations)
end

"""
    initial_condition_couette_unsteady(x, t, equations::LatticeBoltzmannEquations2D)

Initial state for an *unsteady* Couette flow setup, which is also the exact solution for the
incompressible Navier-Stokes equations. To be used in combination with
[`boundary_condition_couette`](@ref) and [`boundary_condition_noslip_wall`](@ref). In the limit,
this setup will converge to the state set in [`initial_condition_couette_steady`](@ref).
"""
function initial_condition_couette_unsteady(x, t, equations::LatticeBoltzmannEquations2D)
  @unpack L, u0, rho0, nu = equations

  x1, x2 = x
  v1 = u0*x2/L
  for m in 1:100
    lambda_m = m * pi / L
    v1 += 2 * u0 * (-1)^m/(lambda_m * L) * exp(-nu * lambda_m^2 * t) * sin(lambda_m * x2)
  end

  rho = 1
  v2 = 0

  return equilibrium_distribution(rho, v1, v2, equations)
end

"""
    initial_condition_couette_unsteady(x, t, equations::LatticeBoltzmannEquations2D)

Initial state for a *steady* Couette flow setup. To be used in combination with
[`boundary_condition_couette`](@ref) and [`boundary_condition_noslip_wall`](@ref).
"""
function initial_condition_couette_steady(x, t, equations::LatticeBoltzmannEquations2D)
  @unpack L, u0, rho0 = equations

  rho = rho0
  v1 = u0 * x[2] / L
  v2 = 0

  return equilibrium_distribution(rho, v1, v2, equations)
end

"""
    boundary_condition_couette(u_inner, orientation, direction, x, t,
                               surface_flux_function,
                               equations::LatticeBoltzmannEquations2D)

Moving *upper* wall boundary condition for a Couette flow setup. To be used in combination with
[`boundary_condition_noslip_wall`](@ref) for the lower wall and
[`boundary_condition_periodic`](@ref) for the lateral boundaries.
"""
function boundary_condition_couette(u_inner, orientation, direction, x, t,
                                    surface_flux_function,
                                    equations::LatticeBoltzmannEquations2D)
  return boundary_condition_moving_wall_ypos(u_inner, orientation, direction, x, t,
                                             surface_flux_function, equations)
end

# Pre-defined source terms should be implemented as
# function source_terms_WHATEVER(u, x, t, equations::LatticeBoltzmannEquations2D)

# Calculate 1D flux in for a single point
@inline function flux(u, orientation::Integer, equations::LatticeBoltzmannEquations2D)
  if orientation == 1
    v_alpha = equations.v_alpha1
  else
    v_alpha = equations.v_alpha2
  end
  return v_alpha .* u
end

# Calculate maximum wave speed for local Lax-Friedrichs-type dissipation
# @inline function max_abs_speed_naive(u_ll, u_rr, orientation::Integer, equations::LatticeBoltzmannEquations2D)
#   λ_max =
# end

@inline function flux_godunov(u_ll, u_rr, orientation::Integer, equations::LatticeBoltzmannEquations2D)
  if orientation == 1
    v_alpha = equations.v_alpha1
  else
    v_alpha = equations.v_alpha2
  end
  return 0.5 * ( v_alpha .* (u_ll + u_rr) - abs.(v_alpha) .* (u_rr - u_ll) )
end

"""
    density(p::Real, equations::LatticeBoltzmannEquations2D)
    density(u, equations::LatticeBoltzmannEquations2D)

Calculate the macroscopic density from the pressure `p` or the particle distribution functions `u`.
"""
@inline density(p::Real, equations::LatticeBoltzmannEquations2D) = p / equations.c_s^2
@inline density(u, equations::LatticeBoltzmannEquations2D) = sum(u)

"""
    velocity(u, orientation, equations::LatticeBoltzmannEquations2D)

Calculate the macroscopic velocity for the given `orientation` (1 -> x, 2 -> y) from the
particle distribution functions `u`.
"""
@inline function velocity(u, orientation::Integer, equations::LatticeBoltzmannEquations2D)
  if orientation == 1
    v_alpha = equations.v_alpha1
  else
    v_alpha = equations.v_alpha2
  end

  return dot(v_alpha, u)/density(u, equations)
end

"""
    velocity(u, equations::LatticeBoltzmannEquations2D)

Calculate the macroscopic velocity vector from the particle distribution functions `u`.
"""
@inline function velocity(u, equations::LatticeBoltzmannEquations2D)
  @unpack v_alpha1, v_alpha2 = equations
  rho = density(u, equations)

  return SVector(dot(v_alpha1, u)/rho,
                 dot(v_alpha2, u)/rho)
end

"""
    pressure(rho::Real, equations::LatticeBoltzmannEquations2D)
    pressure(u, equations::LatticeBoltzmannEquations2D)

Calculate the macroscopic pressure from the density `rho` or the  particle distribution functions
`u`.
"""
@inline pressure(rho::Real, equations::LatticeBoltzmannEquations2D) = rho * equations.c_s^2
@inline pressure(u, equations::LatticeBoltzmannEquations2D) = pressure(density(u, equations), equations)

"""
    equilibrium_distribution(alpha, rho, v1, v2, equations::LatticeBoltzmannEquations2D)

Calculate the local equilibrium distribution for the distribution function with index `alpha` and
given the macroscopic state defined by `rho`, `v1`, `v2`.
"""
@inline function equilibrium_distribution(alpha, rho, v1, v2, equations::LatticeBoltzmannEquations2D)
  @unpack weights, c_s, v_alpha1, v_alpha2 = equations

  va_v = v_alpha1[alpha]*v1 + v_alpha2[alpha]*v2
  cs_squared = c_s^2
  v_squared = v1^2 + v2^2

  return weights[alpha] * rho * (1 + va_v/cs_squared
                                   + va_v^2/(2*cs_squared^2)
                                   - v_squared/(2*cs_squared))
end

@inline function equilibrium_distribution(rho, v1, v2, equations::LatticeBoltzmannEquations2D)
  return SVector(equilibrium_distribution(1, rho, v1, v2, equations),
                 equilibrium_distribution(2, rho, v1, v2, equations),
                 equilibrium_distribution(3, rho, v1, v2, equations),
                 equilibrium_distribution(4, rho, v1, v2, equations),
                 equilibrium_distribution(5, rho, v1, v2, equations),
                 equilibrium_distribution(6, rho, v1, v2, equations),
                 equilibrium_distribution(7, rho, v1, v2, equations),
                 equilibrium_distribution(8, rho, v1, v2, equations),
                 equilibrium_distribution(9, rho, v1, v2, equations))
end

function equilibrium_distribution(u, equations::LatticeBoltzmannEquations2D)
  rho = density(u, equations)
  v1, v2 = velocity(u, equations)

  return equilibrium_distribution(rho, v1, v2, equations)
end

"""
    collision_bgk(u, dt, equations::LatticeBoltzmannEquations2D)

Collision operator for the Bhatnagar, Gross, and Krook (BGK) model.
"""
@inline function collision_bgk(u, dt, equations::LatticeBoltzmannEquations2D)
  @unpack c_s, nu = equations
  tau = nu / (c_s^2 * dt)
  return -(u - equilibrium_distribution(u, equations))/(tau + 1/2)
end

@inline have_constant_speed(::LatticeBoltzmannEquations2D) = Val(true)

@inline function max_abs_speeds(equations::LatticeBoltzmannEquations2D)
  @unpack c = equations

  return c, c
end

# Convert conservative variables to primitive
@inline cons2prim(u, equations::LatticeBoltzmannEquations2D) = u

# Convert conservative variables to entropy variables
@inline cons2entropy(u, equations::LatticeBoltzmannEquations2D) = u

end # @muladd


using JuMP
using Gurobi

mutable struct node
  xUb
  xLb
  zUb
  zLb
  modelo::JuMP.Model
  status::Symbol
end

mutable struct list
  L::Vector{node}
  Zsup
  Zinf
  xOt
  isup
  solint
end

### meu prob é de max

function mudaparamax(m::Model)
  #muda o c da função objetivo
  if m.objSense == :Min
    @objective(m,:Max,-m.obj)
  end
  return m
end

function testabin(v::Vector,vtype::Vector)
    bin=1
    tam=length(v)
    for i in 1:tam
      if vtype[i] == :Bin
        if v[i]!=0 && v[i]!=1
            bin=0
        end
      end
    end
    return bin
end

function podainv(no::node)
  ###retorna 1 caso positivo (relaxação inviavel, realize a poda) e 0 caso a relaxação seja viavel e prossiga para outros testes
  if no.status == :Optimal
    return 0
  end
  return 1
end

function podaotim(no::node,lista::list)
  vtype = no.modelo.colCat
  if testabin(no.xUb,vtype) == 1
    no.xLb=no.xUb
    no.zLb=no.zUb
    lista.solint = lista.solint + 1
    if lista.Zinf < no.zUb
      lista.Zinf = no.zUb
      lista.xOt = no.xUb
    end
    return 1
  end
  return 0
end

function podalimit(no::node,lista::list)
  if no.zUb < lista.Zinf
    return 1
  end
  return 0
end

function atualizaLupper(lista::list)
  tam=length(lista.L)
  maior=-1e10
  ind=0
  for i in 1:tam
    if lista.L[i].zUb > maior
      maior=lista.L[i].zUb
      ind=i
    end
  end
  lista.Zsup=maior
  lista.isup=ind
end

function bound(no::node,lista::list)
  #retorno 1 se fiz poda, 0 c.c
  if podainv(no) == 1
    return 1
  elseif podaotim(no,lista) == 1 || podalimit(no,lista) == 1
    return 1
  end
  return 0
end

function preencheno(mod::JuMP.Model)
  m=deepcopy(mod)
  status=solve(m, relaxation=true)
  xrel=copy(m.colVal)
  zub=copy(m.objVal)
  no1=node(xrel,[],zub,-1e10,m,status)
  return no1
end

function achamaisfrac(no::node)
  tam=length(no.xUb)
  f=copy(no.xUb)
  vtype = no.modelo.colCat
  for i in 1:tam
    if vtype[i] == :Bin
      teto=ceil(no.xUb[i])
      partefrac=teto-no.xUb[i]
      f[i]=min(partefrac,1-partefrac)
    else
      f[i]=0
    end
  end
  #max,ind=findmax(f)
  return findmax(f)
end

function inserefilhos(no::node,S::list)

  max,ind=achamaisfrac(no)

  mod1=deepcopy(no.modelo)
  mod1.colLower[ind]=1
  mod1.colUpper[ind]=1
  filho1=preencheno(mod1)

  mod2=deepcopy(no.modelo)
  mod2.colLower[ind]=0
  mod2.colUpper[ind]=0
  filho2=preencheno(mod2)

  push!(S.L,filho1,filho2)

  atualizaLupper(S)

  return S
end

function criaretornos(m,obj,xotim,status,tempo,bound,cont,nos,solint)
  m.objVal = copy(obj)
  m.colVal = copy(xotim)
  m.objBound = bound
  m.ext[:time] = tempo
  m.ext[:status] = status
  m.ext[:nodes] = nos
  m.ext[:iter] = cont
  m.ext[:intsols] = solint
  return m
end

function solveMIP(mod)

  tic()

  m=deepcopy(mod)

  #antes de tudo devemos mudar nosso problema para Max
  m=mudaparamax(m)

  #crio o primeio no e dentro dele relaxo o problema
  no1=preencheno(m)

  #vejo se o problema é inviavel
  if podainv(no1) == 1
    model=criaretornos(mod,no1.zLb,no1.xLb,no1.status,0,0,0,0,0)
    if no1.status == :Unbounded
      println("Problema Unbounded")
      return model
    end
    println("Problema inviavel")
    return model
  end

  vtype=no1.modelo.colCat
  #vejo se o problema ja deu a resposta inteira mesmo com a relaxação
  if testabin(no1.xUb,vtype) == 1
    if mod.objSense == :Min
      obj=-no1.zUb
    else
      obj=no1.zUb
    end
    model=criaretornos(mod,obj,no1.xUb,no1.status,0,no1.zUb,0,0,1)
    println("Solução otima encontrada")
    return model
  end

  #como nao é bin nem inviavel podamos (branch), escolho o x mais fracionario para fixar
  #usamos uma lista vazia para inicialização
  listaini=list([],1e10,-1e10,[],0,0)

  #criamos a lista com dois nós
  # esses nós foram originados a partir da fixação do x mais fracionario (em um no fixamos em 0 e no outro o fixamos em 1)

  S=inserefilhos(no1,listaini)

  #agora começa o momento de branch e de bound, perceba que ja fiz o branch então dentro do while começamos com o bound
  #paro quando encontrei uma solução viavel suficientemente proxima da solução da relaxação

  ϵ = abs(S.Zsup - S.Zinf)
  cont=0
  jg_time0 = time_ns()

  while ϵ > 1e-5 && cont < 5000 && length(S.L)!=0

    if (time_ns()-jg_time0)/1e9 > 180
      status = :UserLimit
      break
    end

    cont=cont+1

    #seleciono um no para testar o bound
    #a cada 8 vezes escolhendo o no do mesmo ramo (sempre retirando do final da lista), mudo para a esolha do que tem maior upperbound
    #faço isso com o objetivo de não ficar preso em um nó que pode não ser o otimo escolho o que tem maior upper bound pq neste não corro
    #o risco de analisar um nó que possui um upperbound inferior ao da função objetivo
    h=(cont/30)

    if floor(h) == h && S.isup != 0
      no=S.L[S.isup]
      deleteat!(S.L,S.isup)
    else
      no=pop!(S.L)
    end

    devobound=bound(no,S)
    #a função bound atualiza o lowerbound no caso de poda por otimalidade

    if devobound == 1
      #atualizo o upperbound com os novos elementos e volto pro loop para escolher outro nó para analisar
      atualizaLupper(S)

    else
      #como não consegui podar o ramo que contém esse nó devo fazer o branch do no o separando em 2 e fixando na variavel mais fracionaria
      S=inserefilhos(no,S)
      #a inserefilhos ja atualiza os bounds
    end

    ϵ = abs(S.Zsup - S.Zinf)
  end

  time=toc()

  if cont >= 5000
    status = :UserLimit
  elseif (S.Zinf < -1e5)
    status = :Infeasible
  else
    status = :Optimal
  end

  if mod.objSense == :Min
    obj=-(S.Zinf)
  else
    obj=S.Zinf
  end
  xotim=S.xOt
  nos=length(S.L)
  model=criaretornos(mod,obj,xotim,status,time,S.Zinf,cont,nos,S.solint)

  return model

end


# Spectra and time series plots from virtual satellite data.
#
# t rho vx vy p bz ex ey
#
# Hongyang Zhou, hyzhou@umich.edu

using SignalAnalysis, DelimitedFiles, Glob
using SignalAnalysis.Units
using Plots, Measures

files = glob("satellite*.csv")

samplerate = 2Hz
color = :dense
left_margin = 5mm

const ax = fill(plot(),10,1)

for file in files
   println("file: ", file)
   data = readdlm(file, ','; header=true)
   t0 = data[1][1]
   p = signal(data[1][:,5] .* 1e9, samplerate) # [nPa]
   bz = signal(data[1][:,6] .* 1e9, samplerate)# [nT]
   ex = signal(data[1][:,7] .* 1e3, samplerate)# [mV/m]
   ey = signal(data[1][:,8] .* 1e3, samplerate)# [mV/m]

   l = @layout[[a1;a2;a3;a4;a5] grid(2,2){0.7w}]
   #ax[1] = plot(title="Power Spectra at [$(file[end-5:end-4]), 0, 0]Re",
   #   framestyle=nothing,showaxis=false,xticks=false,yticks=false,margin=0mm)

   x = parse(Int, file[end-5:end-4])
   ax[2] = plot(data[1][:,1], data[1][:,2] ./ 1e6; label="density", title="[$x, 0, 0]Re",
      ylabel="[amu/cc]", left_margin)
   ax[3] = plot(data[1][:,1], data[1][:,3:4] ./ 1e3; label=["Vx" "Vy"],
      ylabel="[km/s]", left_margin)
   ax[4] = plot(data[1][:,1], samples(p); label="p", ylabel="[nPa]", left_margin)
   ax[5] = plot(data[1][:,1], samples(bz); label="bz", ylabel="[nT]", left_margin)
   ax[6] = plot(data[1][:,1], data[1][:,7:8] .* 1e3; label=["Ex" "Ey"],
      ylabel="[mV/m]", xlabel="Time [s]", left_margin)

   # magnitude units in [db], or 10log10(power)
   ax[7] = specgram(p; nfft=128, t0, crange=100, color ,title="P")
   ax[8] = specgram(bz; nfft=128, t0, crange=100, color, title="Bz")
   ax[9] = specgram(ex; nfft=128, t0, crange=100, color, title="Ex")
   ax[10] = specgram(ey; nfft=128, t0, crange=100, color, title="Ey")
   plot(ax[2:10]..., layout=l, size=(1500,1000), legend=true)
   savefig("spectrum_"*file[end-5:end-4]*"Re.png")
end


using Test, Statistics
using Knet.Ops20: batchnorm, bnmoments, bnparams
using Knet.KnetArrays: KnetArray
using AutoGrad: gradcheck
using CUDA: CUDA, functional

@testset "batchnorm" begin

    # Random.seed!(42)

    # utils
    std2(x) = let x_mu = x .- mean(x)
        mean(x_mu .* x_mu)
    end

    # gradcheck functions
    bn3(a) = batchnorm(a[1], nothing, a[2]; training=true)
    bn1(a) = batchnorm(a; training=true)
    bn3ts(a) = batchnorm(a[1], bnmoments(), a[2]; training=false)
    bn1ts(a) = batchnorm(a, bnmoments(); training=false)

    types = (Float64,) #TODO: [Float32, Float64]
    sizes = Dict([2=>(5,10), 4=>(3,4,5,3), 5=>(4,3,4,5,2)])
    dims = [2, 4, 5]
    gpu_av = CUDA.functional()

    for d in dims
        for et in types
            sz = sizes[d]
            tn(str) = string(str, d, "d")
            # CPU buffers
            ax = 5rand(et, sz...);
            C = sz[end-1]
            aw = bnparams(et, C)
            # GPU buffers
            kax = gpu_av ? KnetArray{et}(5rand(sz...)) : nothing
            kaw = gpu_av ? KnetArray{et}(bnparams(C)) : nothing

            # Array types
            at = typeof(ax)
            kat = typeof(kax)

            @testset "{$et, $d}" begin
                @testset "cpu-stat" begin
                    # The unit gaussian output
                    @test abs(mean(batchnorm(ax; training=true))) < 1e-3
                    @test isapprox(std2(batchnorm(ax; training=true)), 1.0; rtol=1e-2)
                end

                @testset "cpu-grads" begin
                    @test gradcheck(bn1, ax)
                    @test gradcheck(bn3, (ax, aw))
                end

                if gpu_av
                    @testset "gpu-stats" begin
                        # isapprox 0 didn't work
                        @test abs(mean(batchnorm(kax; training=true))) < 1e-3
                        @test isapprox(std2(batchnorm(kax; training=true)), 1.0; rtol=1e-2)
                    end

                    @testset "gpu-grads" begin
                        @test gradcheck(bn1, kax)
                        @test gradcheck(bn3, (kax, kaw))
                    end

                    @testset "dev-consistency" begin
                        mc = bnmoments() #cpu
                        mg = bnmoments() #gpu
                        y1 = batchnorm(ax, mc)
                        y2 = batchnorm(kat(ax), mg) #use the same array
                        @test isapprox(mc.mean, at(mg.mean))
                        @test isapprox(mc.var, at(mg.var))
                        @test isapprox(y1, at(y2))
                    end
                end

                # Moments should be updates
                @testset "test-moments" begin
                    mc = bnmoments()
                    batchnorm(ax, mc)
                    m, v = copy(mc.mean), copy(mc.var)
                    batchnorm(.1ax.+3, mc)
                    @test isapprox(m, mc.mean; rtol=1e-20)
                    @test isapprox(v, mc.var; rtol=1e-20)
                    if gpu_av
                        mg = bnmoments()
                        batchnorm(kax, mg)
                        m, v = copy(mg.mean), copy(mg.var)
                        batchnorm(.1kax.+3, mg)
                        @test isapprox(m, mg.mean; rtol=1e-20)
                        @test isapprox(v, mg.var; rtol=1e-20)
                    end
                end

                @testset "training-moments" begin
                    mc = bnmoments()
                    batchnorm(ax, mc; training=true)
                    m, v = copy(mc.mean), copy(mc.var)
                    batchnorm(.1ax.+3, mc; training=true)
                    @test ~isapprox(m, mc.mean; rtol=1e-20)
                    @test ~isapprox(v, mc.var; rtol=1e-20)
                    if gpu_av
                        mg = bnmoments()
                        batchnorm(kax, mg; training=true)
                        m, v = copy(mg.mean), copy(mg.var)
                        batchnorm(.1kax.+3, mg; training=true)
                        @test ~isapprox(m, mg.mean; rtol=1e-20)
                        @test ~isapprox(v, mg.var; rtol=1e-20)
                    end
                end
                # TODO: Remove this if after 2d mode test backward supported
                if d > 2
                    @testset "cpu-grads-testing" begin
                        m1 = bnmoments()
                        @test gradcheck(bn1ts, ax)
                        @test gradcheck(bn3ts, (ax, aw))
                    end

                    if gpu_av
                        @testset "gpu-grads-testing" begin
                            @test gradcheck(bn1ts, kax)
                            @test gradcheck(bn3ts, (kax, kaw))
                        end
                    end
                end
            end #end of {dim, type} testset
        end
    end
end #end of testset

# suppress the return
nothing


function read_trees(filename)

    open(filename) do f
        data = readlines(f)
        return transpose(map(x->(x=='#'), hcat(collect.(data)...)))
    end

end

function hit_count(trees, down, right)
    nrow, ncol = size(trees)
    hit = 0
    j = 1
    for i in 1:down:nrow
        hit += trees[i,j]
        j = mod1(j+right,ncol)
    end
    return hit
end

trees = read_trees("input03.txt")
println(hit_count(trees, 1, 3))
println(hit_count(trees, 1, 1) *
        hit_count(trees, 1, 3) *
        hit_count(trees, 1, 5) *
        hit_count(trees, 1, 7) *
        hit_count(trees, 2, 1) )


@testset "check tree isomorphy" begin
    @testset "fails on non-trees" begin
        nodes = [Node(i,i) for i in 1:3]
        arcs = [Arc(1,2), Arc(1, 3), Arc(2, 3)]

        triangle = Topology(nodes, arcs)
        @test_throws AssertionError are_isomorphic(triangle, triangle)

        forest = Topology(nodes, arcs[1:1])
        @test_throws AssertionError are_isomorphic(forest, forest)
    end

    @testset "test on trees" begin
        nodes = [Node(i,i) for i in 1:5]
        path1 = Topology(nodes, [Arc(i,i+1) for i=1:4])
        path2 = Topology(nodes, [Arc(1,3), Arc(3,5), Arc(5,2), Arc(2,4)])
        star1 = Topology(nodes, [Arc(1,i) for i=2:5])
        star2 = Topology(nodes, [Arc(5,i) for i=1:4])

        @test are_isomorphic(path1, path1) == true
        @test are_isomorphic(path1, path2) == true
        @test are_isomorphic(star1, star1) == true
        @test are_isomorphic(star1, star2) == true

        @test are_isomorphic(path1, star1) == false
        @test are_isomorphic(star1, path2) == false
    end
end

@testset "test steiner tree enumeration" begin
    nclass3, repr3 = enumerate_steiner(3)
    @test length(nclass3) == 1
    @test nclass3[1] == 1
    @test length(repr3) == sum(nclass3)

    nclass4, repr4 = enumerate_steiner(4)
    @test length(nclass4) == 2
    @test nclass4 == [1,1]
    @test length(repr4) == sum(nclass4)

    nclass5, repr5 = enumerate_steiner(5)
    @test length(nclass5) == 3
    @test nclass5 == [1,1,1]
    @test length(repr5) == sum(nclass5)

    nclass6, repr6 = enumerate_steiner(6)
    @test length(nclass6) == 4
    @test nclass6 == [1,1,1,2]
    @test length(repr6) == sum(nclass6)
end

@testset "test FST labeling" begin
    nodes = [Node(i,i) for i in 1:6]

    # single steiner node
    Y = Topology(nodes[1:4], [Arc(1,4), Arc(2,4), Arc(3,4)])
    @test label_fst(Y) == (2, 3)

    # two steiner nodes
    H = Topology(nodes, [Arc(1,5), Arc(2,5), Arc(3,6),Arc(4,6), Arc(5,6)])
    @test label_fst(H) == (2, (3, 4))
end

@testset "test FST enumeration" begin
    terminals = [Node(2,3), Node(4,5), Node(6,7), Node(8,9)]

    @test_throws AssertionError enumerate_fst([])
    @test_throws AssertionError enumerate_fst(terminals[1:1])
    @test_throws AssertionError enumerate_fst(terminals[1:2])

    trees3 = enumerate_fst(terminals[1:3])
    @test length(trees3) == 1

    trees4 = enumerate_fst(terminals[1:4])
    @test length(trees4) == 1 + 2
end