type BatesModel{CalibratedModelType} <: Model{CalibratedModelType}
  modT::CalibratedModelType
  hestonModel::HestonModel
  nu::ConstantParameter
  delta::ConstantParameter
  lambda::ConstantParameter
  process::BatesProcess
  common::ModelCommon
end

function BatesModel(process::BatesProcess)
  hm = HestonModel(process.hestonProcess)
  nu = ConstantParameter([process.nu], NoConstraint())
  delta = ConstantParameter([process.delta], PositiveConstraint())
  lambda = ConstantParameter([process.lambda], PositiveConstraint())

  bm = BatesModel(CalibratedModelType(), hm, nu, delta, lambda, process, ModelCommon())

  generate_arguments!(bm)

  return bm
end

get_theta(bm::BatesModel) = get_theta(bm.hestonModel)
get_kappa(bm::BatesModel) = get_kappa(bm.hestonModel)
get_sigma(bm::BatesModel) = get_sigma(bm.hestonModel)
get_rho(bm::BatesModel) = get_rho(bm.hestonModel)
get_v0(bm::BatesModel) = get_v0(bm.hestonModel)
get_nu(bm::BatesModel) = bm.nu.data[1]
get_delta(bm::BatesModel) = bm.delta.data[1]
get_lambda(bm::BatesModel) = bm.lambda.data[1]

generate_arguments!(bm::BatesModel) =
  bm.process = BatesProcess(get_risk_free_rate(bm.process), get_dividend_yield(bm.process), get_s0(bm.process),
              get_v0(bm), get_kappa(bm), get_theta(bm), get_sigma(bm), get_rho(bm),
              get_nu(bm), get_delta(bm), get_lambda(bm))


###########################
### Auxiliary Functions ###
###########################

"""
    consword(w::Integer)

Create a word with n consecutive set bits as least significant.
"""
consword(n::T) where {T <: Integer} = (T(1) << n) - T(1)

#= Counts the number of bits in the integer w =#
nbits(w::T) where {T <: Integer} = 8*sizeof(T) - leading_zeros(w)
nbits(w::BigInt)                 = (w == 0) ? Int64(0) : ndigits(w, base=2)

#= Counting the number of set bits in the integer w =#
nsetbits(w::T) where {T <: Integer} = count_ones(w)

#= Counting the maximum number of consecutive set bits in an integer w =#
function nconsbits(w::T) where {T <: Integer}
    n = 0
    while w != Miss
        w &= (w << T(1))
        n += 1
    end # while
    n
end # function

#################
### Functions ###
#################

"""
    shift(w::Integer, c)

Left shifts bits once and adds new bit as LSB (according to character c).
Removes the MSB if character is a Hit.
"""
function shift(w::T, c) where {T <: Integer}
    if c == Miss
        return w << T(1)
    else
        return ((w << T(1)) | T(1)) & consword(T(nbits(w)))
    end
end

"""
    shift_rowhit(w::Integer, c, n::Integer)

Left shifts bits once and adds new bit as LSB (according to character c).
Returns the word n if the bits in nw coincide with the ones in n.
"""
function shift_rowhit(w::T, c, n::T) where {T <: Integer}
    #= Shift bits according to charcter c and default to n if last few char are equal to n =#
    if c == Miss
        w <<= T(1)
    else
        w = (w << T(1)) | T(1)
    end
    return (w & n == n) ? n : w
end


using VortexDistributions

# simple test field
Nx = 400; Ny = 400
Lx = 200; Ly = Lx
x = LinRange(-Lx/2,Lx/2, Nx+1)[1:end-1]; y = LinRange(-Ly/2,Ly/2, Ny+1)[1:end-1];
psi0 = one.(x*y') |> complex
psi = Torus(copy(psi0),x,y)
xp = 1.133
yp = 2.787
vp = PointVortex(xp,yp,1)
sp = ScalarVortex(vp)
vortex!(psi,sp)

ψ = psi.ψ
phase = angle.(ψ)

diffx,diffy = phase_jumps(phase, 1),phase_jumps(phase, 2)
a = zero(phase)
b = zero(phase)
phase_jumps!(a, phase)
@test a == diffx
phase_jumps!(b, phase, 2)
@test b == diffy


using QuantumTensors
using Base.Test

function test_to_sparse()
    tol = 1e-10
    N = 5
    @test norm(((destroy(N) |> sparse) - spdiagm(sqrt(1:(N-1)),1,N,N))|>full )< tol
    @test norm(((create(N) |> sparse) - spdiagm(sqrt(1:(N-1)),-1,N,N))|>full )< tol
    @test norm(((jz(N) |> sparse) - spdiagm(-N:N,0,2N+1,2N+1))|>full) < tol

    jzm = jz(N) |> sparse
    jp = jplus(N) |> sparse
    jm1 = jp'
    jm2 = jminus(N) |> sparse
    @test norm((jm1-jm2)|> full) < tol

    comm_rel = jp*jm1 - jm1*jp
    @test norm((comm_rel-2jzm)|>full) < tol
end

function test_to_full()
    tol = 1e-10
    N = 5
    @test norm((destroy(N) |> full) - diagm(sqrt(1:(N-1)),1))< tol
    @test norm((create(N) |> full) - diagm(sqrt(1:(N-1)),-1))< tol
    @test norm((sigma(2,1,2) |> full) - [0 1; 0 0])< tol
    @test norm((sigma(2,1,1) |> full) - [1 0; 0 0])< tol
    @test norm((sigma(2,2,1) |> full) - [0 0; 1 0])< tol
    @test norm((sigma(2,2,2) |> full) - [0 0; 0 1])< tol
end

function test_transpose()
    tol = 1e-10
    @test norm((sigma(2,1,2) |>full) - (sigma(2,2,1) |> transpose |>full)) < tol
    @test norm((destroy(5) |>full) - (create(5) |> transpose |>full)) < tol
end

function test_convert_promote_ops()
  tol = 1e-10
  N = 5
  # a = destroy(N)
  # ac = convert(DiagonalOp{Complex128}, a)
  idt = IdentityOp{N}()
  idtf = convert(DiagonalOp{Float64}, idt)
  idtz1 = convert(DiagonalOp{Complex128}, idt)
  idtz2 = convert(DiagonalOp{Complex128}, idtf)
  idtz3, idtz4 = promote(idtf, idtz1)
  @test idtz4 === idtz1
  @test isa(idtf, DiagonalOp{Float64})
  @test isa(idtz1, DiagonalOp{Complex128})
  @test isa(idtz2, DiagonalOp{Complex128})
  @test isa(idtz3, DiagonalOp{Complex128})
  @test idtf.jkm == (0,0,0)
  @test idtz1.jkm == (0,0,0)
  @test idtz2.jkm == (0,0,0)
  @test idtz3.jkm == (0,0,0)
  @test norm(idtf.vals - ones(N)) < tol
  @test norm(idtz1.vals - ones(N)) < tol
  @test norm(idtz2.vals - ones(N)) < tol
  @test norm(idtz3.vals - ones(N)) < tol

  @test all_close(idt, idtf)
  @test all_close(idt, idtz1)
  @test all_close(idt, idtz2)
  @test all_close(idt, idtz3)
  @test all_close(idt, idtz4)
  @test all_close(idtf, idtz1)
  @test all_close(idtf, idtz2)
  @test all_close(idtf, idtz3)
  @test all_close(idtf, idtz4)
end

function test_convert_promote_tops()
  tol = 1e-10
  N = 5
  a = destroy(N)
  bc = convert(DiagonalOp{Complex128}, destroy(N))

  ac, bc = promote(a, bc)
  @test bc === bc
  @test all_close(ac, a)
  abc = tensor(a, bc)
  acbc = tensor(ac, bc)

  @test typeof(abc) === typeof(acbc)
  @test all_close(abc, acbc)
end

function test_convert_promote_topsum()
  tol = 1e-10
  N = 5
  a = destroy(N)
  bc = convert(DiagonalOp{Complex128}, destroy(N))
  ii = IdentityOp{N}()

  ac, bc = promote(a, bc)
  @test bc === bc
  @test all_close(ac, a)
  abc = tensor(a, bc)
  acbc = tensor(ac, bc)

  iibc = tensor(ii, bc)
  aii = tensor(a, ii)

  tops = TensorOpProductSum(abc, iibc, aii)
end

function test_mul()
    tol = 1e-10
    a = destroy(5)
    @test norm((a*a |> full) -(a|> full)*(a|> full)) < tol
    @test norm((a*a' |> full) -(a|> full)*(a|> full)')<tol
end

function test_apply_serial(N=20,tol=1e-8)
  println("Testing serial application of operators")
  a = destroy(N)
  b = destroy(N)
  c = destroy(N)
  II = IdentityOp{N}()
  tops = (tensor(a', b, II), tensor(a, b', II), tensor(II, II, c'*c))
  topsum = TensorOpProductSum(tops...)

  B = randn(N,N,N)
  C = zeros(N,N,N)
  coeffs = [1.,1.,1.]
  apply_serial!(coeffs, topsum, B, C)
  # Profile.clear()
  println("apply_serial!: TensorOpProductSum ")
  gc()
  @time for k=1:10
      apply_serial!(coeffs, topsum, B,C)
  end

  topskarr = collect([TensorOpProductKernel(top) for top in tops]);
  topsarr = collect(tops);

  println("apply_serial!: [TensorOpProductKernel]")
  Cpp = zeros(C)
  apply_serial!(coeffs, topskarr, B, Cpp)
  gc()
  @time for k=1:10
      apply_serial!(coeffs, topskarr, B, Cpp)
  end

  @test (norm((Cpp-C)[:])) < tol

  println("apply_serial!: [TensorOpProduct]")
  Cpp = zeros(C)
  apply_serial!(coeffs, topsarr, B, Cpp)
  gc()
  @time for k=1:10
      apply_serial!(coeffs, topsarr, B, Cpp)
  end
  @test norm((Cpp-C)[:]) < tol

  Cp = zeros(C)
  Cpf = sub(Cp,1:length(C));
  Bf = sub(B,1:length(B));
  top_sps = [sparse(top) for top = tops]
  combinedsp = sum(top_sps)
  println("combined CSR matrix")
  A_mul_B!(Cpf, combinedsp, Bf)
  gc()
  @time for k=1:10
      A_mul_B!(Cpf, combinedsp, Bf)
  end
  @test norm((Cp-C)[:]) < tol

  println("apply_serial!: [CSR-mat]")
  apply_serial!(coeffs, top_sps, Bf, Cpf)
  gc()
  @time for k=1:10
      apply_serial!(coeffs, top_sps, Bf, Cpf)
  end
  @test norm((Cp-C)[:]) < tol
end

module SymbolicTest
  const sqrt_arr1 = sqrt(1:1000000)
  const sqrt_arr2 = sqrt(1:1000000)
  function a1ad2_symbolic!(y, x, α=0, β=1)
    N1, N2=size(x)

    if α != 1
      if α == 0
        fill!(y, 0)
      else
        scale!(y, α)
      end
    end

    for l=1:(N2-1)
      @simd for k=1:(N1-1)
        @inbounds y[k,l+1] +=  β * sqrt_arr1[k]*sqrt_arr2[l]*x[k+1,l]
      end
    end
    y
  end
  function a1ad2_symbolic2!(y, x, α=0, β=1)
    N1, N2=size(x)

    if α != 1
      if α == 0
        fill!(y, 0)
      else
        scale!(y, α)
      end
    end

    offset1 = -1
    offset2 = 1

    oIDX = 1
    iIDX = 1 - stride(x,1)*offset1 - stride(x,2)*offset2
    for l=1:N2
      lflag = 2 <= l <=N2
      @simd for k=1:N1
        kflag = 1 <= k <= N1-1
        if kflag && lflag
          @inbounds y[oIDX] +=  β * sqrt_arr1[k]*sqrt_arr2[l-1]*x[iIDX]
        end
        oIDX +=1
        iIDX +=1
      end
    end
    y
  end
  function a1ad2_symbolic3!(y, x, α=0, β=1)
    N1, N2=size(x)

    if α != 1
      if α == 0
        fill!(y, 0)
      else
        scale!(y, α)
      end
    end

    offset1 = -1
    offset2 = 1

    oIDX = 1
    iIDX = 1 - stride(x,1)*offset1 - stride(x,2)*offset2
    for l=1:N2
      lflag = 2 <= l <=N2
      for k=1:N1
        kflag = 1 <= k <= N1-1
        if kflag && lflag
          @inbounds y[oIDX] +=  β * sqrt_arr1[k]*sqrt_arr2[l-1]*x[iIDX]
        end
        oIDX +=1
        iIDX +=1
      end
    end
    y
  end
end

function test_A_mul_B(N1=10,N2=10)
    println("Testing application of single tensor product operators to vectors")
    tol = 1e-5
    a1 = destroy(N1)
    ad2 = create(N2)
    a1ad2 = tensor(a1,ad2)
    p1 = randn(N1)
    p2 = randn(N2)
    x = tensor(p1, p2)

    println("Tensor prod with allocation:")
    y1 = zeros(N1,N2)
    A_mul_B!(y1, a1ad2, x)
    gc()
    @time begin
        y1 = A_mul_B!(zeros(N1,N2), a1ad2, x)
    end

    println("Tensor prod (inplace):")
    y5 = zeros(N1,N2)
    gc()
    @time begin
        A_mul_B!(y5, a1ad2, x)
    end
    @test norm(y1[:] - y5[:]) < tol

    println("Direct usage of TensorOpProductKernel:")
    y10 = zeros(N1,N2)
    A_mul_B!(2., a1ad2.kernel, x, 3., y10)
    y10 = zeros(N1,N2)
    gc()
    @time begin
        A_mul_B!(2., a1ad2.kernel, x, 3., y10)
    end
    @test norm(2*y1[:] - y10[:]) < tol

    println("'Symbolic' application of specialized function:")
    y11 = zeros(N1,N2)
    SymbolicTest.a1ad2_symbolic!(y11, x, 3., 2.)
    y11 *= 0
    gc()
    @time begin
        SymbolicTest.a1ad2_symbolic!(y11, x, 3., 2.)
    end
    @test norm(2*y1[:] - y11[:]) < tol

    println("'Symbolic' application of second specialized function:")
    y12 = zeros(N1,N2)
    SymbolicTest.a1ad2_symbolic2!(y12, x, 3., 2.)
    y12 *= 0
    gc()
    @time begin
        SymbolicTest.a1ad2_symbolic2!(y12, x, 3., 2.)
    end
    @test norm(2*y1[:] - y12[:]) < tol

    println("'Symbolic' application of third specialized function:")
    y12 = zeros(N1,N2)
    SymbolicTest.a1ad2_symbolic3!(y12, x, 3., 2.)
    y12 *= 0
    gc()
    @time begin
        SymbolicTest.a1ad2_symbolic3!(y12, x, 3., 2.)
    end
    @test norm(2*y1[:] - y12[:]) < tol

    println("Tensorproduct of individual products:")
    y2 = tensor(((a1|>sparse)*p1),((ad2|>sparse)*p2))
    gc()
    @time y2 = tensor(((a1|>sparse)*p1),((ad2|>sparse)*p2))
    @test norm(y1[:] - y2[:]) < tol

    println("CSR mat-vec:")

    xv = x[:]

    a1ad2sp = (a1ad2 |> sparse)
    y3 = a1ad2sp * xv
    gc()
    @time y3 = a1ad2sp * xv
    @test norm(y1[:] - y3[:]) < tol

    println("CSR mat-vec (inplace):")
    A_mul_B!(2., a1ad2sp, xv, 1., y3)
    y3[:] = y1[:]
    gc()
    @time A_mul_B!(2., a1ad2sp, xv, 1., y3)
    @test norm(3*y1[:] - y3[:]) < tol

    println("Size of vectors: ", Base.summarysize(x)/1024^2, " MB")
    println("Size of sparse matrix rep: ", Base.summarysize(a1ad2sp)/1024^2, " MB")
    println("Size of tensor matrix rep: ", Base.summarysize(a1ad2)/1024^2, " MB")
end

# addprocs(3)
versioninfo()
test_to_sparse()
test_to_full()
test_transpose()
test_mul()
test_convert_promote_ops()
test_convert_promote_tops()
test_convert_promote_topsum()
test_apply_serial(20)
test_A_mul_B(100,100)


import RayTracing: AzimuthalQuadrature, nazim2

"""
    Quadrature{A<:AzimuthalQuadrature,P<:PolarQuadrature,T<:Real}

Holds information regarding both the azimuthal and polar quadrature and the total weights.
"""
struct Quadrature{T<:Real,A<:AzimuthalQuadrature,P<:PolarQuadrature}
    azimuthal::A
    polar::P
    ω::Matrix{T}
end

function Quadrature(azimuthal::AzimuthalQuadrature{Na,T}, polar::PolarQuadrature{Np,T}) where {Na,Np,T}
    @unpack δs, ωₐ = azimuthal
    @unpack sinθs, ωₚ = polar
    n_azim_2 = nazim2(azimuthal)
    n_polar_2 = npolar2(polar)

    # IDEA: I think we can use n_azim_4 because the matrix has repeated values
    ω = Matrix{T}(undef, n_azim_2, n_polar_2)

    for i in 1:n_azim_2, j in 1:n_polar_2
        ω[i, j] = 4π * ωₐ[i] * ωₚ[j] * δs[i] * sinθs[j]
    end
    # ω .*= 4π

    return Quadrature(azimuthal, polar, ω)
end


using DiffEqFlux, OrdinaryDiffEq, Flux, Optim, Plots, AdvancedHMC
using JLD, StatsPlots, Distributions

u0 = [2.0; 0.0]
datasize = 40
tspan = (0.0, 1)
tsteps = range(tspan[1], tspan[2], length = datasize)

function trueODEfunc(du, u, p, t)
    true_A = [-0.1 2.0; -2.0 -0.1]
    du .= ((u.^3)'true_A)'
end

prob_trueode = ODEProblem(trueODEfunc, u0, tspan)
mean_ode_data = Array(solve(prob_trueode, Tsit5(), saveat = tsteps))
ode_data = mean_ode_data .+ 0.1 .* randn(size(mean_ode_data)..., 30)

####DEFINE THE NEURAL ODE#####
dudt2 = FastChain((x, p) -> x.^3,
                  FastDense(2, 50, relu),
                  FastDense(50, 2))
prob_neuralode = NeuralODE(dudt2, tspan, Tsit5(), saveat = tsteps)

function predict_neuralode(p)
    Array(prob_neuralode(u0, p))
end
function loss_neuralode(p)
    pred = predict_neuralode(p)
    loss = sum(abs2, ode_data .- pred)
    return loss, pred
end

function l(θ)
    lp = logpdf(MvNormal(zeros(length(θ) - 1), 1.0), θ[1:end-1])
    ll = sum(logpdf.(Normal.(ode_data, θ[end]), predict_neuralode(θ[1:end-1])))
    return lp + ll
end
function lp(θ)
    return logpdf(MvNormal(zeros(length(θ) - 1), 1.0), θ[1:end-1])
end
function ll(θ)
    return sum(logpdf.(Normal.(ode_data, θ[end]), predict_neuralode(θ[1:end-1])))
end

function dldθ(θ)
    x, lambda = Flux.Zygote.pullback(l,θ)
    grad = first(lambda(1))
    return x, grad
end
function dlpdθ(θ)
    x, lambda = Flux.Zygote.pullback(lp,θ)
    grad = first(lambda(1))
    return x, grad
end
function dlldθ(θ)
    x, lambda = Flux.Zygote.pullback(ll,θ)
    grad = first(lambda(1))
    return x, grad
end

init = [Float64.(prob_neuralode.p); 1.0]

opt = DiffEqFlux.sciml_train(x -> -l(x), init, ADAM(0.05), maxiters = 1500)
pmin = opt.minimizer;
metric  = DiagEuclideanMetric(length(pmin))
h = Hamiltonian(metric, l, dldθ)
integrator = Leapfrog(find_good_stepsize(h, pmin))
prop = AdvancedHMC.NUTS{MultinomialTS, GeneralisedNoUTurn}(integrator, 10)
adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.5, prop.integrator))

samples, stats = sample(h, prop, pmin, 500, adaptor, 500; progress=true)

using MCMCTempering
tempered_samples = sample()


# module for calculating standard deviation parameters
# robot rotates at same position
# including random noise and bias
# this is iterated for 100 times

module MotionTestRotBias
  using DataFrames, Statistics

  include(joinpath(split(@__FILE__, "src")[1], "src/model/world/world.jl"))
  include(joinpath(split(@__FILE__, "src")[1], "src/model/robot/differential_wheeled_robot/differential_wheeled_robot.jl"))

  function main(time_interval=0.1)
    w = World(-5.0, 5.0, -5.0, 5.0)

    rs = []
    init_pose = [0.0, 0.0, 0.0]
    rotate = Agent(0.0, 0.1)
    for i in 1:100
      r = DifferentialWheeledRobot(init_pose, 0.2, "black",
                                   rotate, time_interval,
                                   noise_per_meter=5,
                                   noise_std=pi/60,
                                   bias_rate_stds=[0.1, 0.1])
      push!(rs, r)
    end

    for t in 0:time_interval:40
      draw(w)
      annotate!(-3.5, 4.5, "t = $(time_interval)", "black")
      for r in rs
        draw!(r)
      end
    end

    poses = DataFrame(range=[sqrt(r.pose[1]^2 + r.pose[2]^2) for r in rs],
                      theta=[r.pose[3] for r in rs])
    theta_var = var(poses.theta, corrected=false)
    theta_mean = mean(poses.theta)
    println("Variance of range and theta: ", poses)
    println("Variance of theta: ", theta_var)
    println("Mean of theta: ", theta_mean)
    println("Standard deviation per deg of theta: ", sqrt(theta_var/theta_mean))
  end
end


correct = sqrt(π/sqrt(ℯ)) / 2
let
    x, w = hermite(10)
    Q = sum(w .* ( x .- 1.0) .* sin.(x))
    @test Q ≈ correct
end


@polly function kernel_atax(A, x, y, tmp)
    m,n = size(A)

    for i = 1:n
        y[i] = 0
    end

    for i = 1:m
        tmp[i] = zero(eltype(tmp))

        for j = 1:n
            tmp[i] = tmp[i] + A[i,j] * x[j]
        end

        for j = 1:n
            y[j] = y[j] + A[i,j] * tmp[i]
        end
    end
end

let
    m = 1900
    n = 2100

    A = zeros(Float32, m, n)
    x = zeros(Float32, n)
    y = zeros(Float32, n)
    tmp = zeros(Float32, m)

    for i = 1:n
        x[i] = 1 + (i / n)
    end
    for i = 1:m, j = 1:n
        A[i,j] = ((i+j) % n) / (5*m)
    end

    SUITE["atax"] = @benchmarkable kernel_atax(A, x, y, tmp) setup = (A = copy($A); x = copy($x); y = copy($y); tmp = copy($tmp))
end


using DataFrames

function df2mm(df::DataFrame)
    n = size(df, 1)
    mm_raw = [fill(1.0, n, 1)]
    mm_name = ["const"]
    for (name, value) in eachcol(df)
        if eltype(value) <: Real
            push!(mm_raw, hcat(Float64.(value)))
            push!(mm_name, string(name))
        else
            uvalue = unique(value)
            length(uvalue) == 1 && continue
            dvalue = Dict(v=>i for (i, v) in enumerate(uvalue))
            mvalue = zeros(n, length(uvalue))
            for i in 1:n
                mvalue[i, dvalue[value[i]]] = 1.0
            end
            push!(mm_raw, mvalue[:, 2:end])
            append!(mm_name, string.(name, "_", uvalue[2:end]))
        end
    end
    (data=hcat(mm_raw...), names=mm_name)
end

function lm(df, y, xs)
    yv = Float64.(df[y])
    xv, xn = df2mm(df[[xs;]])
    params = (transpose(xv)*xv)\(transpose(xv)*yv)
    DataFrame(name = xn, estimate=params)
end

Chapter04/09. Estimating linear regression/lm.jl


using Statistics

@testsuite "statistics" (AT, eltypes)->begin
    for ET in eltypes
        if !(ET in [Float16, Float32, Float64])
            continue
        end
        @testset "std" begin
            @test compare(std, AT, rand(ET, 10))
            @test compare(std, AT, rand(ET, 10,1,2))
            @test compare(std, AT, rand(ET, 10,1,2); corrected=true)
            @test compare(std, AT, rand(ET, 10,1,2); dims=1)
        end

        @testset "var" begin
            @test compare(var, AT, rand(ET, 10))
            @test compare(var, AT, rand(ET, 10,1,2))
            @test compare(var, AT, rand(ET, 10,1,2); corrected=true)
            @test compare(var, AT, rand(ET, 10,1,2); dims=1)
            @test compare(var, AT, rand(ET, 10,1,2); dims=[1])
            @test compare(var, AT, rand(ET, 10,1,2); dims=(1,))
            @test compare(var, AT, rand(ET, 10,1,2); dims=[2,3])
            @test compare(var, AT, rand(ET, 10,1,2); dims=(2,3))
        end

        @testset "mean" begin
            @test compare(mean, AT, rand(ET, 2, 2))
            @test compare(mean, AT, rand(ET, 2, 2); dims=2)
            @test compare(mean, AT, rand(ET, 2, 2, 2); dims=[1,3])
            @test compare(x->mean(sin, x), AT, rand(ET, 2,2))
            @test compare(x->mean(sin, x; dims=2), AT, rand(ET, 2,2))
            @test compare(x->mean(sin, x; dims=[1,3]), AT, rand(ET, 2,2,2))
        end
    end

    for ET in eltypes
        # Doesn't work with ComplexF32 in oneAPI for some reason.
        if !(ET in [Float32, Float64, Float16, ComplexF64])
            continue
        end
        @testset "cov" begin
            s = 100
            @test compare(cov, AT, rand(ET, s))
            @test compare(cov, AT, rand(ET, s, 2))
            @test compare(cov, AT, rand(ET, s, 2); dims=2)
            if ET <: Real
                @test compare(cov, AT, rand(ET(1):ET(100), s))
            end
        end

        @testset "cor" begin
            s = 100
            @test compare(cor, AT, rand(ET, s))
            @test compare(cor, AT, rand(ET, s, 2))
            @test compare(cor, AT, rand(ET, s, 2); dims=2)
            if ET <: Real
                @test compare(cor, AT, rand(ET(1):ET(100), s))
            end
        end
    end
end


#hellwarththeory.jl

"""
Implementation of Hellwarth's extended treatment of Osaka's finite temperature version of Feynman's variational technique applied to the Polaron model; generalising it from the case of a single phonon mode to the case of multiple phonon modes.

And Hellwarth 1999:
https://doi.org/10.1103/PhysRevB.60.299

"""

using Optim
using QuadGK

function singlemode_variation(α, β; v = 7.2, w = 6.5, verbose::Bool = false)

    initial = [v, w]
    lower = [0.0, 0.0]
    upper = [Inf, v]

    f(x) = F(x[1], x[2], β, α)

    solution = Optim.optimize(
        Optim.OnceDifferentiable(f, initial; autodiff = :forward),
        lower,
        upper,
        initial,
        Fminbox(BFGS()),
    )

    if Optim.converged(solution) == false
        print("\tWARNING: Failed to converge to v,w soln? : ", Optim.converged(solution))
    end

    if verbose
        println()
        show(solution)
    end

    v, w = Optim.minimizer(solution)
    return v, w
end


# This file is a part of AstroLib.jl. License is MIT "Expat".
# Copyright (C) 2016 Mosè Giordano.

function _mag2geo{T<:AbstractFloat}(lat::T, long::T, pole_lat::T, pole_long::T)
    r       = 1.0 # Distance from planet center.  Value unimportant -- just need
                  # a length for conversion to rectangular coordinates
    lat  = deg2rad(lat)
    long = deg2rad(long)
    alt  = lat*0.0 + r

    # convert to rectangular coordinates
    #   x-axis: defined by the vector going from Earth's center towards
    #        the intersection of the equator and Greenwich's meridian.
    #   z-axis: axis of the geographic poles
    #   y-axis: defined by y=z^x
    x = alt*cos(lat)*cos(long)
    y = alt*cos(lat)*sin(long)
    z = alt*sin(lat)

    # First rotation: in the plane of the current meridian from magnetic pole to
    # geographic pole.
    togeolat = Array(T, 3, 3)
    togeolat[:,1] = [cos(pi/2 - pole_lat), 0.0, -sin(pi/2 - pole_lat)]
    togeolat[:,2] = [                 0.0, 1.0,                   0.0]
    togeolat[:,3] = [sin(pi/2 - pole_lat), 0.0,  cos(pi/2 - pole_lat)]
    out = togeolat * [x, y, z]

    # Second rotation matrix: rotation around plane of the equator, from the
    # meridian containing the magnetic poles to the Greenwich meridian.
    maglong2geolong = Array(T, 3, 3)
    maglong2geolong[:,1] = [ cos(pole_long), sin(pole_long), 0.0]
    maglong2geolong[:,2] = [-sin(pole_long), cos(pole_long), 0.0]
    maglong2geolong[:,3] = [            0.0,            0.0, 1.0]
    out = maglong2geolong * out

    geolat  = rad2deg(atan2(out[3], hypot(out[1], out[2])))
    geolong = rad2deg(atan2(out[2], out[1]))
    # I don't care about that one...just put it there for completeness' sake
    # magalt  = vecnorm(out) - r
    return geolat, geolong
end

"""
    mag2geo(latitude, longitude[, year]) -> geographic_latitude, geographic_longitude

### Purpose ###

Convert from geomagnetic to geographic coordinates.

### Explanation ###

Converts from geomagnetic (latitude, longitude) to geographic (latitude,
longitude).  Altitude is not involved in this function.

### Arguments ###

* `latitude`: geomagnetic latitude (North), in degrees.
* `longitude`: geomagnetic longitude (East), in degrees.
* `year` (optional numerical argument): the year in which to perform conversion.
  If omitted, defaults to current year.

The coordinates can be passed as arrays of the same length.

### Output ###

The 2-tuple of geographic (latitude, longitude) coordinates, in degrees.

If geomagnetic coordinates are given as arrays, a 2-tuple of arrays of the same
length is returned.

### Example ###

Find position of North Magnetic Pole in 2016

``` julia
mag2geo(90, 0, 2016)
# => (86.395,-166.29000000000002)
```

### Notes ###

This function uses list of North Magnetic Pole positions provided by World
Magnetic Model (https://www.ngdc.noaa.gov/geomag/data/poles/NP.xy).

`geo2mag` converts geographic coordinates to geomagnetic coordinates.

Code of this function is based on IDL Astronomy User's Library.
"""
mag2geo(lat::Real, long::Real, year::Real=Dates.year(Dates.now())) =
    _mag2geo(promote(float(lat), float(long),
                     deg2rad(POLELATLONG[year][1]::Float64),
                     deg2rad(POLELATLONG[year][2]::Float64))...)

function mag2geo{LA<:Real, LO<:Real}(lat::AbstractArray{LA},
                                     long::AbstractArray{LO},
                                     year::Real=Dates.year(Dates.now()))
    @assert length(lat) == length(long)
    typela   = typeof(float(one(LA)))
    geolat   = similar(lat, typela)
    geolong  = similar(lat, typela)
    polelat  = deg2rad(POLELATLONG[year][1])
    polelong = deg2rad(POLELATLONG[year][2])
    for i in eachindex(lat)
        geolat[i], geolong[i] =
            _mag2geo(promote(float(lat[i]), float(long[i]),
                             polelat, polelong)...)
    end
    return geolat, geolong
end


abstract type TrainingMethods end

struct Batch <: TrainingMethods end
struct BatchL2 <: TrainingMethods
    lambda::AbstractFloat
end

struct BatchCM <: TrainingMethods end
struct BatchCML2 <: TrainingMethods
    lambda::AbstractFloat
end

function train_core(W::AbstractMatrix{T}, b::AbstractMatrix{T}, X::AbstractMatrix{T}, Y::AbstractMatrix{T},
                    mf::MappingFunctions, tm::Batch) where {T<:BaseType}
    H = mapping(W, b, X, mf)
    return Y*pinv(H);
end

function train_core(W::AbstractMatrix{T}, b::AbstractMatrix{T}, X::AbstractMatrix{T}, Y::AbstractMatrix{T},
                    mf::MappingFunctions, tm::BatchL2) where {T<:BaseType}
    H = mapping(W, b, X, mf)
    L = size(H,1)
    N = size(H,2)
    tH = ctranspose(H)
    if N > L
        return Y*tH*pinv(H*tH + tm.lambda*eye(T,L))
    else
        return Y*(pinv(tH*H + tm.lambda*eye(T,N))*tH)
    end
end

function train_core(W::AbstractMatrix{T}, b::AbstractMatrix{T}, X::AbstractMatrix{T}, Y::AbstractMatrix{T},
                    mf::MappingFunctions, tm::BatchCM) where {T<:BaseType}
    A, B = cormat(W, b, X, Y, mf)
    return B*pinv(A)'
end

function train_core(W::AbstractMatrix{T}, b::AbstractMatrix{T}, X::AbstractMatrix{T}, Y::AbstractMatrix{T},
                    mf::MappingFunctions, tm::BatchCML2) where {T<:BaseType}
    A, B = cormat(W, b, X, Y, mf)
    return B*pinv(A+tm.lambda*eye(T,size(A,1)))'
end

function cormat(W::AbstractMatrix{T}, b::AbstractMatrix{T}, X::AbstractMatrix{T}, Y::AbstractMatrix{T},
                mf::MappingFunctions) where {T<:BaseType}
    P = size(X,2)
    Nh = size(W, 1)
    No = size(Y, 1)
    Hj = zeros(T, Nh, 1)
    A = zeros(T, Nh, Nh)
    B = zeros(T, No, Nh)
    for i=1:P
        for j=1:Nh
            Hj[j] = mapping(W[j,:], b[j], X[:,i], mf);
        end
        tHj = ctranspose(Hj)
        A += Hj * tHj
        B += Y[:,i] * tHj
    end

    return A, B
end


# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    DirectionPartitioner(direction; tol=1e-6)

A method for partitioning spatial objects along a given `direction`
with bandwidth tolerance `tol`.
"""
struct DirectionPartitioner{T,N} <: AbstractSpatialPredicatePartitioner
  direction::SVector{N,T}
  tol::Float64

  function DirectionPartitioner{T,N}(direction, tol) where {N,T}
    new(normalize(direction), tol)
  end
end

DirectionPartitioner(direction::SVector{N,T}; tol=1e-6) where {T,N} =
  DirectionPartitioner{T,N}(direction, tol)

DirectionPartitioner(direction::NTuple{N,T}; tol=1e-6) where {T,N} =
  DirectionPartitioner(SVector(direction), tol=tol)

(p::DirectionPartitioner)(x, y) = begin
  δ = x - y
  d = p.direction
  norm(δ - (δ⋅d)*d) < p.tol
end


# -----------------------------------------------
# Minimal Distributions routines
# -----------------------------------------------

# TEST:
# using Rmath

#=
macro QD2(A, B) quote
    esc(:(sum((A - B).^2)))
end
=#

abstract type AbstractDistributionType end

# #########   NORMAL DISTRIBUTION   #########

abstract type Normal<:AbstractDistributionType end

function PDF{T<:Number}(x::T, dtype::Type{Normal})
    return dnorm(x, 0, 1)
end

function PDF{Tx<:Number, Tmu<:Number, Tsig<:Number}(x::Tx,  mu::Tmu, sigma²::Tsig, dtype::Type{Normal})
    return dnorm(x, mu, sqrt(sigma²))
end

function PDF{Tx<:Number, Tmu<:Number, Tsig<:Number}(X::Array{Tx,1}, mu::Tmu, sigma²::Tsig, dtype::Type{Normal})
    # return exp.(-(X - mu).^2 / (2*sigma²)) / sqrt(2*pi*sigma²);
    return [ dnorm(x, mu, sqrt(sigma²)) for x in X ]
end

function PDF{T<:Number}(X::Array{T}, mu::Array{T,1}, Σ::Array{T},  dtype::Type{Normal})
    m, n = size(X)
    if length(dtype.mu) == n && size(dtype.Σ)[1] == n
        return sqrt(2*pi*det(Σ)) * exp.(-0.5 * (X' - row(mu)') * inv(Σ) * (X - row(mu)));
    else
        error("X, mu and Σ have not valid dimensions")
    end
end

function CDF{T<:Number}(x::T, dtype::Type{Normal})
    return pnorm(x, 0, 1)
end

function CDF{Tx<:Number, Tmu<:Number, Tsig<:Number}(x::Tx,  mu::Tmu, sigma²::Tsig, dtype::Type{Normal})
    # return CDF((sqrt(sigma²)*x + mu), Normal)
    return pnorm(x, mu, sqrt(sigma²))
end

function CDF{Tx<:Number, Tmu<:Number, Tsig<:Number}(X::Array{Tx,1}, mu::Tmu, sigma²::Tsig, dtype::Type{Normal})
    # return exp.(-(X - mu).^2 / (2*sigma²)) / sqrt(2*pi*sigma²);
    return [ pnorm(x, mu, sqrt(sigma²)) for x in X ]
end

#=
function CDF{T<:Number}(X::Array{T}, mu::Array{T,1}, Σ::Array{T},  dtype::Normal)
    return 0;
end
=#

function logLike{T<:Number}(X::Array{T,1}, mu::T, sigma²::T, dtype::Type{Normal})
    N = length(X)
    return - N/2*log(sigma²) - (0.5/sigma²)*sum((X - mu).^2)
end

function logLike{T<:Number}(X::Array{T,1}, mu::Array{T,1}, Σ::Array{T}, dtype::Type{Normal})
    N = length(X)
    return - N/2*log(det(Σ)) - 0.5 * sum(X' - row(mu)') * inv(Σ) * (X - row(mu));
end

# println(@time CDF([-1.0, 1], 1.0, 4.0, Normal))

# #########   STUDENT DISTRIBUTION   #########

abstract type Student<:AbstractDistributionType end

function PDF{Tx<:Number, Tnu<:Number}(x::Tx, ν::Tnu, dtype::Type{Student})
    # dt(x, ν) == gamma((ν+1)/2) / (sqrt(ν*pi) * gamma(ν/2)) * (1 + x.^2/ν).^(-(ν+1)/2)
    return dt(x, ν)
end

function PDF{Tx<:Number, Tnu<:Number}(X::Array{Tx,1}, ν::Tnu, dtype::Type{Student})
    return [ dt(x, ν) for x in X ]
end

function CDF{Tx<:Number, Tnu<:Number}(x::Tx, ν::Tnu, dtype::Type{Student})
    return pt(x, ν)
end

function CDF{Tx<:Number, Tnu<:Number}(X::Array{Tx,1}, ν::Tnu, dtype::Type{Student})
    return [ pt(x, ν) for x in X ]
end

# println(@time CDF(1.0, 1.0, Student))

# #########   LOGISTIC   #########

abstract type Logistic<:AbstractDistributionType end

logistic{T<:Number}(X::Array{T,1}) = 1 ./ (1 + exp.(X));

function logLike{T<:Number}(X::Array{T}, Y::Array{T}, theta::Array{T}, dtype::Type{Logistic})
    N = length(Y)
    return sum((logistic(X * theta)).^Y .* (1 - logistic(X * theta)).^(1 - Y)) / N
end


export
    Projectile,
    propagate,
    closest_time_of_approach,
    closest_approach_distance,
    closest_time_of_approach_and_distance,
    get_intersection_time

struct Projectile
    pos::VecSE2 # position (x,y,θ)
    v::Float64  # speed
end

propagate(P::Projectile, Δt::Float64) = Projectile(P.pos + polar(P.v*Δt, P.pos.θ), P.v)
"""
The time at which two projectiles are closest to one another
A negative value indicates that this occurred in the past.
see: http://geomalgorithms.com/a07-_distance.html
"""
function closest_time_of_approach(A::Projectile, B::Projectile)

    W = convert(VecE2, A.pos - B.pos)
    Δ = polar(A.v, A.pos.θ) - polar(B.v, B.pos.θ)
    aΔ = normsquared(Δ)

    if aΔ ≈ 0.0
        return 0.0
    else
        return -(W⋅Δ) / aΔ
    end
end
function closest_approach_distance(A::Projectile, B::Projectile, t_CPA::Float64=closest_time_of_approach(A, B))
    Aₜ = convert(VecE2, propagate(A, t_CPA).pos)
    Bₜ = convert(VecE2, propagate(B, t_CPA).pos)
    return LinearAlgebra.norm(Aₜ - Bₜ)
end
function closest_time_of_approach_and_distance(A::Projectile, B::Projectile)
    t_CPA = closest_time_of_approach(A, B)
    d_CPA = closest_approach_distance(A, B, t_CPA)
    return (t_CPA, d_CPA)
end

function get_intersection_time(A::Projectile, seg::LineSegment)

    o = VecE2(A.pos)
    v₁ = o - seg.A
    v₂ = seg.B - seg.A
    v₃ = polar(1.0, A.pos.θ + π/2)

    denom = (v₂⋅v₃)

    if !isapprox(denom, 0.0, atol=1e-10)
        d₁ = (v₂×v₁) / denom # time for projectile (0 ≤ t₁)
        t₂ = (v₁⋅v₃) / denom # time for segment (0 ≤ t₂ ≤ 1)
        if 0.0 ≤ d₁ && 0.0 ≤ t₂ ≤ 1.0
            return d₁/A.v
        end
    else
        # denom is zero if the segment and the projectile are parallel
        # only collide if they are perfectly aligned
        if are_collinear(A.pos, seg.A, seg.B)
            dist_a = normsquared(seg.A - o)
            dist_b = normsquared(seg.B - o)
            return sqrt(min(dist_a, dist_b)) / A.v
        end
    end

    return Inf
end
function closest_time_of_approach_and_distance(P::Projectile, seg::LineSegment, if_no_col_skip_eval::Bool=false)

    o = VecE2(P.pos)
    v₁ = o - seg.A
    v₂ = seg.B - seg.A
    v₃ = polar(1.0, P.pos.θ + π/2)

    denom = (v₂⋅v₃)

    if !isapprox(denom, 0.0, atol=1e-10)

        d₁ = (v₂×v₁) / denom # time for projectile (0 ≤ d₁)
        t₂ = (v₁⋅v₃) / denom # time for segment (0 ≤ t₂ ≤ 1)

        if 0.0 ≤ d₁ && 0.0 ≤ t₂ ≤ 1.0
            t = d₁/P.v
            return (t, 0.0)
        else
            # no collision, get time of closest approach to each endpoint
            if !if_no_col_skip_eval
                r = polar(1.0, P.pos.θ)
                projA = proj(seg.A - o, r, VecE2)
                projB = proj(seg.B - o, r, VecE2)
                tA = max(LinearAlgebra.norm(projA)/P.v * sign(r⋅projA), 0.0)
                tB = max(LinearAlgebra.norm(projB)/P.v * sign(r⋅projB), 0.0)
                pA = VecE2(propagate(P, tA).pos)
                pB = VecE2(propagate(P, tB).pos)
                distA = normsquared(seg.A - pA)
                distB = normsquared(seg.B - pB)
                if distA < distB
                    return (tA, sqrt(distA))
                else
                    return (tB, sqrt(distB))
                end
            else
                return (Inf, Inf)
            end
        end
    else
        # denom is zero if the segment and the projectile are parallel
        # only collide if they are perfectly aligned
        if are_collinear(P.pos, seg.A, seg.B)
            # either closest now, will be to A, or will be to B
            dist_a = normsquared(seg.A - o)
            dist_b = normsquared(seg.B - o)
            t = sqrt(min(dist_a, dist_b)) / P.v
            return (t, 0.0)
        else
            # not colinear
            if !if_no_col_skip_eval
                # either closest now, will be to A, or will be to B
                r = polar(1.0, P.pos.θ)
                projA = proj(seg.A - o, r, VecE2)
                projB = proj(seg.B - o, r, VecE2)
                tA = max(LinearAlgebra.norm(projA)/P.v * sign(r⋅projA), 0.0)
                tB = max(LinearAlgebra.norm(projB)/P.v * sign(r⋅projB), 0.0)
                pA = VecE2(propagate(P, tA).pos)
                pB = VecE2(propagate(P, tB).pos)
                distA = normsquared(seg.A - pA)
                distB = normsquared(seg.B - pB)
                if distA < distB
                    return (tA, distA)
                else
                    return (tB, distB)
                end
            else
                return (Inf,Inf)
            end
        end
    end
end


using LinearAlgebra

#
# KF
#

s = Statevector([0, 0, 0], Matrix(1.0I, 3, 3))
KalmanFilter.update!(s, [1, 1, 1], Matrix(1.0I, 3, 3), Matrix(1.0I, 3, 3))
@test s.xhat == [.5, .5, .5]
@test s.P == Matrix(0.5I, 3, 3)
KalmanFilter.predict!(s, [1 1 0;0 1 1;0 0 1], Matrix(0.5I, 3, 3))
@test s.xhat == [1, 1, .5]

#
# EKF
#

s = Statevector([0, 0, 0], Matrix(1.0I, 3, 3))
KalmanFilter.update!(s, [1, 1, 1], x -> [x[1], x[2], x[3]], x -> Matrix(1.0I, 3, 3), Matrix(1.0I, 3, 3))
@test s.xhat == [.5, .5, .5]
@test s.P == Matrix(0.5I, 3, 3)
KalmanFilter.predict!(s, x -> [1, 1, .5], x -> [1 1 0;0 1 1;0 0 1], Matrix(0.5I, 3, 3))
@test s.xhat == [1, 1, .5]

#
# EKF, numeric Jacobian
#

s = Statevector([0., 0., 0.], Matrix(1.0I, 3, 3))
KalmanFilter.update!(s, [1., 1., 1.], x -> [x[1], x[2], x[3]], Matrix(1.0I, 3, 3))
@test s.xhat == [.5, .5, .5]
@test s.P == Matrix(0.5I, 3, 3)
KalmanFilter.predict!(s, x -> [1, 1, .5], Matrix(0.5I, 3, 3))
@test s.xhat == [1, 1, .5]


using StatsPlots

include("simulation_helpers.jl")
include("../helper_functions.jl")

begin
    begin
        total_simulation_length = 1000
        interval_length = 1000
        job_size = 1
        p1_arrival_rate = 1
        p2_arrival_rate = 1
        private_service_rate = 1 / 1.5
        public_service_rate = 1 / 0.7
        breakin = 3
        pst, qst = equilibrium_probability(
        p1_arrival_rate,
        p2_arrival_rate,
        job_size,
        α -> 1 / (1.5 * (α)),
        α -> 1 / (0.7 * (α)))
    end

event_data, interval_data = simulate_queue(
    total_simulation_length,
    interval_length,
    Exponential(1 / p1_arrival_rate),
    Exponential(1 / p2_arrival_rate),
    Exponential(job_size / public_service_rate),
    Exponential(job_size / private_service_rate),
    pst,
    qst,
)

transform!(event_data, :eventtype => ByRow(x -> split(x, "_")) => [:source, :event])
select!(event_data, Not(:eventtype))

event_data_wide = vcat([match_events(gdf) for gdf in groupby(event_data, :source)]...)
select!(event_data_wide, [:source, :arrival_time, :departure_time])
transform!(event_data_wide,
    [:arrival_time, :departure_time] => ((x, y) -> (y - x)) => :response_time,
    :arrival_time => (x -> cut(x, 0:interval_length:total_simulation_length + interval_length, extend=true)) => :interval)
performance_data = combine(groupby(event_data_wide, [:interval, :source]), :response_time => (x -> mean(skipmissing(x))) => :mean_response_time)

transform!(performance_data, :source => ByRow(x -> split(x, "-")) => [:player, :dummy])
select!(performance_data, Not(:dummy))
groupby(performance_data, [:player, :interval])
interval_data.cut = cut((interval_data.interval .- 1) .* interval_length, 0:interval_length:total_simulation_length + interval_length)
performance_data = leftjoin(performance_data, select(interval_data, Not(:interval)), on = [:interval => :cut])

cost_data = combine(groupby(performance_data, [:player, :interval])) do sdf
    evaluate_cost(sdf)
end
        rename!(cost_data, :x1 => :cost)

begin
    p2 = plot(title = "Cost by Communication Interval", xlabel = "Interval", ylabel = "Cost")
    cost_data = leftjoin(sort!(cost_data, :player), interval_data[:, [:interval, :cut]], on = (:interval => :cut), makeunique=true)
    for g in groupby(cost_data, :player)
        plot!(p2, g.interval_1, g.cost, label = first(g.player), lw = 2)
    end
    p2
end
p2
end

println(median(skipmissing(event_data_wide.response_time)))
println(mean(skipmissing(event_data_wide.response_time)))
histogram(log10.(collect(skipmissing(event_data_wide.response_time))),
    nbins = 50,
    label = false,
    xlabel = "log10(Response Time)",
    ylabel = "Frequency",
    title = "Histogram of Response Times")

ARs = 0.1:0.1:2
total_volume = 2
JSs = 1 / 2 * total_volume ./ ARs
n_reps = 50
dot_every = Int(floor(n_reps / 5))

begin
mean_resval = Float64[]
var_resval = Float64[]

for (AR, JS) in zip(ARs, JSs)
    mean_temp = Float64[]
    var_temp = Float64[]
    private_service_rate = 1 / (1.5 * JS)
    public_service_rate = 1 / (0.7 * JS)

    print("AR = ", AR, " JS = ", JS, " ")
    for i in 1:n_reps
        if mod(i, dot_every) == 0
            print(".")
        end
        event_data, interval_data = simulate_queue(
            total_simulation_length,
            total_simulation_length,
            AR,
            AR,
            1 / 0.7,
            1 / 1.5,
            JS
        )
        transform!(event_data, :eventtype => ByRow(x -> split(x, "_")) => [:source, :event])
        select!(event_data, Not(:eventtype))
        event_data_wide = vcat([match_events(gdf) for gdf in groupby(event_data, :source)]...)
        select!(event_data_wide, [:source, :arrival_time, :departure_time])
            transform!(event_data_wide,
                    [:arrival_time, :departure_time] => ((x, y) -> (y - x)) => :response_time)
        mean_and_var = combine(event_data_wide,
        :response_time => (
            x -> (mean_RT = mean(skipmissing(x)), var_RT = var(skipmissing(x)))
            ) => AsTable
        )

        append!(mean_temp, mean_and_var[1, 1])
        append!(var_temp, mean_and_var[1, 2])
    end
    println()
    append!(mean_resval, mean(mean_temp))
    append!(var_resval, mean(var_temp))
end
end

plot(ARs, JSs, xlabel = "Arrival Rate", ylabel = "Job Size")
mean_plot = plot(JSs, mean_resval, xlabel = "Job Size", yaxis = :log10, ylabel = "Mean Response Time", legend = :none)
sd_plot = plot(JSs, sqrt.(var_resval), xlabel = "Job Size", yaxis = :log10, ylabel = "SD Response Time", legend = :none)

begin
    total_simulation_length = 10000
    interval_length = total_simulation_length / 10
    job_size = 0.5
    p1_arrival_rate = 2
    p2_arrival_rate = 2
    private_service_rate = 1 / (1.5 * job_size)
    public_service_rate = 1 / (0.7 * job_size)
    breakin = 3
    pst, qst = equilibrium_probability(
        p1_arrival_rate,
        p2_arrival_rate,
        job_size,
        α -> 1 / (1.5 * (α)),
        α -> 1 / (0.7 * (α)))
end

# to ensure correctness: calculate_distributions

event_data, interval_data = simulate_queue(
    1000,
    1000,
    1,
    1,
    1 / 0.7,
    1 / 1.5,
    1
)

function evaluate_response_time(event_data)
    transform!(event_data, :eventtype => ByRow(x -> split(x, "_")) => [:source, :event])
    select!(event_data, Not(:eventtype))
    event_data_wide = vcat([match_events(gdf) for gdf in groupby(event_data, :source)]...)
    select!(event_data_wide, [:source, :arrival_time, :departure_time])
    transform!(event_data_wide,
                [:arrival_time, :departure_time] => ((x, y) -> (y - x)) => :response_time)
    df = combine(event_data_wide,
        :response_time => (
            x -> (mean_RT = mean(skipmissing(x)), sd_RT = sqrt(var(skipmissing(x))))
        ) => AsTable
    )
    return df[1, 1], df[1, 2]
end

evaluate_response_time(event_data)

prepared_data = prepare_event_data_for_plot(event_data)
@df prepared_data plot(
    :time,
    :effect_cumsum,
    group = :groupcolour,
)


using OrdinaryDiffEq, Test, LinearAlgebra, Statistics

# TODO: clean up
@testset "Mass Matrix Accuracy Tests" begin
  function make_prob(mm_A; iip)
    mm_b = mm_A*ones(3)
    # iip
    function mm_f_iip(du,u,p,t)
      mul!(du,mm_A,u)
      tmp = t*mm_b
      du .+= tmp
    end
    mm_g_iip(du,u,p,t) = @. du = u + t
    # oop
    function mm_f_oop(u,p,t)
      du = u./t
      mm_f_iip(du,u,p,t)
      du
    end
    function mm_g_oop(u,p,t)
      du = u./t
      mm_g_iip(du,u,p,t)
      du
    end
    mm_analytic(u0,p,t) = @. 2u0*exp(t) - t - 1
    f = ((mm_f_oop, mm_g_oop), (mm_f_iip, mm_g_iip))[iip+1]
    prob = ODEProblem(ODEFunction(f[1],analytic=mm_analytic,mass_matrix=mm_A),ones(3),
                                     (0.0,1.0))
    prob2 = ODEProblem(ODEFunction(f[2],analytic=mm_analytic),ones(3),(0.0,1.0))
    return prob, prob2
  end
  for iip in (false, true)
    mm_A = [-2.0 1 4
             4 -2 1
             2 1 3]
    prob, prob2 = make_prob(mm_A; iip=iip)

    ######################################### Test each method for exactness

    if iip
      sol = solve(prob  ,Rosenbrock23())
      sol2 = solve(prob2,Rosenbrock23())

      @test norm(sol .- sol2) ≈ 0 atol=1e-11

      sol = solve(prob, Rosenbrock32())
      sol2 = solve(prob2,Rosenbrock32())

      @test norm(sol .- sol2) ≈ 0 atol=1e-11

      sol = solve(prob,  ROS3P())
      sol2 = solve(prob2,ROS3P())

      @test norm(sol .- sol2) ≈ 0 atol=1e-11

      sol = solve(prob,  Rodas3())
      sol2 = solve(prob2,Rodas3())

      @test norm(sol .- sol2) ≈ 0 atol=1e-11

      sol = solve(prob,  RosShamp4())
      sol2 = solve(prob2,RosShamp4())

      @test norm(sol .- sol2) ≈ 0 atol=1e-10

      sol = solve(prob,  Rodas4())
      sol2 = solve(prob2,Rodas4())

      @test norm(sol .- sol2) ≈ 0 atol=1e-9

      sol = solve(prob,  Rodas5())
      sol2 = solve(prob2,Rodas5())

      @test norm(sol .- sol2) ≈ 0 atol=1e-7
    end

    sol = solve(prob,  ImplicitEuler(),dt=1/10,adaptive=false)
    sol2 = solve(prob2,ImplicitEuler(),dt=1/10,adaptive=false)

    @test norm(sol .- sol2) ≈ 0 atol=1e-7

    sol = solve(prob,  RadauIIA5(),dt=1/10,adaptive=false)
    sol2 = solve(prob2,RadauIIA5(),dt=1/10,adaptive=false)
    @test norm(sol .- sol2) ≈ 0 atol=1e-12

    sol = solve(prob,  ImplicitMidpoint(extrapolant = :constant),dt=1/10)
    sol2 = solve(prob2,ImplicitMidpoint(extrapolant = :constant),dt=1/10)

    @test norm(sol .- sol2) ≈ 0 atol=1e-7

    # Anderson method
    sol = solve(prob,  ImplicitEuler(nlsolve=NLAnderson()),dt=1/10,adaptive=false)
    sol2 = solve(prob2,ImplicitEuler(nlsolve=NLAnderson()),dt=1/10,adaptive=false)

    @test_broken norm(sol .- sol2) ≈ 0 atol=1e-7

    sol = solve(prob,  ImplicitMidpoint(nlsolve=NLAnderson(),extrapolant = :constant),dt=1/10)
    sol2 = solve(prob2,ImplicitMidpoint(nlsolve=NLAnderson(),extrapolant = :constant),dt=1/10)

    @test_broken norm(sol .- sol2) ≈ 0 atol=1e-7
    # Functional iteration
    prob, prob2 = make_prob(Matrix{Float64}(1.01I, 3, 3); iip=iip)
    sol = solve(prob,ImplicitEuler(
                          nlsolve=NLFunctional(κ=2000.,tol=1e-7,min_iter=10,max_iter=100)),dt=1/10,adaptive=false)
    sol2 = solve(prob2,ImplicitEuler(),dt=1/10,adaptive=false)
    @test norm(sol .- sol2) ≈ 0 atol=1e-7

    sol = solve(prob, ImplicitMidpoint(extrapolant = :constant,
                          nlsolve=NLFunctional(κ=2000.,tol=1e-7,min_iter=10,max_iter=100)),dt=1/10,adaptive=false)
    sol2 = solve(prob2,ImplicitMidpoint(extrapolant = :constant),dt=1/10)
    @test norm(sol .- sol2) ≈ 0 atol=1e-7

    sol = solve(prob,ImplicitEuler(nlsolve=NLAnderson()),dt=1/10,adaptive=false)
    sol2 = solve(prob2,ImplicitEuler(nlsolve=NLAnderson()),dt=1/10,adaptive=false)
    @test norm(sol .- sol2) ≈ 0 atol=1e-7
    @test norm(sol[end] .- sol2[end]) ≈ 0 atol=1e-7

    sol = solve(prob, ImplicitMidpoint(extrapolant = :constant, nlsolve=NLAnderson()),dt=1/10,adaptive=false)
    sol2 = solve(prob2,ImplicitMidpoint(extrapolant = :constant, nlsolve=NLAnderson()),dt=1/10,adaptive=false)
    @test norm(sol .- sol2) ≈ 0 atol=1e-7
    @test norm(sol[end] .- sol2[end]) ≈ 0 atol=1e-7
  end
end

#=

sol = solve(prob,  Trapezoid())
sol2 = solve(prob2,Trapezoid())

@test norm(sol .- sol2) ≈ 0 atol=1e-7

sol = solve(prob,  TRBDF2())
sol2 = solve(prob2,TRBDF2())

@test norm(sol .- sol2) ≈ 0 atol=1e-9

#sol = solve(prob, SDIRK2())
#sol2 = solve(prob2, SDIRK2())

#@test norm(sol .- sol2) ≈ 0 atol=1e-11

sol = solve(prob,   TRBDF2(),adaptive=false,dt=1/10)
sol2 = solve(prob2, TRBDF2(),adaptive=false,dt=1/10)

=#

# Singular mass matrices

@testset "Mass Matrix Tests with Singular Mass Matrices" begin
  function f!(du, u, p, t)
      du[1] = u[2]
      du[2] = u[2] - 1.
      return
  end

  u0 = [0.,1.]
  tspan = (0.0, 1.0)

  M = fill(0., 2,2)
  M[1,1] = 1.

  m_ode_prob = ODEProblem(ODEFunction(f!;mass_matrix=M), u0, tspan)
  @test_nowarn sol = solve(m_ode_prob, Rosenbrock23())

  M = [0.637947  0.637947
       0.637947  0.637947]

  inv(M) # not caught as singular

  function f2!(du, u, p, t)
      du[1] = u[2]
      du[2] = u[1]
      return
  end
  u0 = fill(0., 2)

  m_ode_prob = ODEProblem(ODEFunction(f2!;mass_matrix=M), u0, tspan)
  @test_nowarn sol = solve(m_ode_prob, Rosenbrock23())
end


## Least squares solutions
a = [ones(20) 1:20 1:20]
b = reshape(eye(8, 5), 20, 2)
for elty in (Float32, Float64, Complex64, Complex128)
    a = convert(Matrix{elty}, a)
    b = convert(Matrix{elty}, b)

    # Vector rhs
    x = a[:,1:2]\b[:,1]
    @test_approx_eq ((a[:,1:2]*x-b[:,1])'*(a[:,1:2]*x-b[:,1]))[1] convert(elty, 2.546616541353384)

    # Matrix rhs
    x = a[:,1:2]\b
    @test_approx_eq det((a[:,1:2]*x-b)'*(a[:,1:2]*x-b)) convert(elty, 4.437969924812031)

    # Rank deficient
    x = a\b
    @test_approx_eq det((a*x-b)'*(a*x-b)) convert(elty, 4.437969924812031)

    # Underdetermined minimum norm
    x = convert(Matrix{elty}, [1 0 0; 0 1 -1]) \ convert(Vector{elty}, [1,1])
    @test_approx_eq x convert(Vector{elty}, [1, 0.5, -0.5])

    # symmetric, positive definite
    @test_approx_eq inv(convert(Matrix{elty}, [6. 2; 2 1])) convert(Matrix{elty}, [0.5 -1; -1 3])

    # symmetric, indefinite
    @test_approx_eq inv(convert(Matrix{elty}, [1. 2; 2 1])) convert(Matrix{elty}, [-1. 2; 2 -1]/3)
end

## Test Julia fallbacks to BLAS routines

# matrices with zero dimensions
@test ones(0,5)*ones(5,3) == zeros(0,3)
@test ones(3,5)*ones(5,0) == zeros(3,0)
@test ones(3,0)*ones(0,4) == zeros(3,4)
@test ones(0,5)*ones(5,0) == zeros(0,0)
@test ones(0,0)*ones(0,4) == zeros(0,4)
@test ones(3,0)*ones(0,0) == zeros(3,0)
@test ones(0,0)*ones(0,0) == zeros(0,0)

# 2x2
A = [1 2; 3 4]
B = [5 6; 7 8]
@test A*B == [19 22; 43 50]
@test At_mul_B(A, B) == [26 30; 38 44]
@test A_mul_Bt(A, B) == [17 23; 39 53]
@test At_mul_Bt(A, B) == [23 31; 34 46]
Ai = A+(0.5*im).*B
Bi = B+(2.5*im).*A[[2,1],[2,1]]
@test Ai*Bi == [-21+53.5im -4.25+51.5im; -12+95.5im 13.75+85.5im]
@test Ac_mul_B(Ai, Bi) == [68.5-12im 57.5-28im; 88-3im 76.5-25im]
@test A_mul_Bc(Ai, Bi) == [64.5+5.5im 43+31.5im; 104-18.5im 80.5+31.5im]
@test Ac_mul_Bc(Ai, Bi) == [-28.25-66im 9.75-58im; -26-89im 21-73im]

# 3x3
A = [1 2 3; 4 5 6; 7 8 9].-5
B = [1 0 5; 6 -10 3; 2 -4 -1]
@test A*B == [-26 38 -27; 1 -4 -6; 28 -46 15]
@test Ac_mul_B(A, B) == [-6 2 -25; 3 -12 -18; 12 -26 -11]
@test A_mul_Bc(A, B) == [-14 0 6; 4 -3 -3; 22 -6 -12]
@test Ac_mul_Bc(A, B) == [6 -8 -6; 12 -9 -9; 18 -10 -12]
Ai = A+(0.5*im).*B
Bi = B+(2.5*im).*A[[2,1,3],[2,3,1]]
@test Ai*Bi == [-44.75+13im 11.75-25im -38.25+30im; -47.75-16.5im -51.5+51.5im -56+6im; 16.75-4.5im -53.5+52im -15.5im]
@test Ac_mul_B(Ai, Bi) == [-21+2im -1.75+49im -51.25+19.5im; 25.5+56.5im -7-35.5im 22+35.5im; -3+12im -32.25+43im -34.75-2.5im]
@test A_mul_Bc(Ai, Bi) == [-20.25+15.5im -28.75-54.5im 22.25+68.5im; -12.25+13im -15.5+75im -23+27im; 18.25+im 1.5+94.5im -27-54.5im]
@test Ac_mul_Bc(Ai, Bi) == [1+2im 20.75+9im -44.75+42im; 19.5+17.5im -54-36.5im 51-14.5im; 13+7.5im 11.25+31.5im -43.25-14.5im]

# Generic integer matrix multiplication
A = [1 2 3; 4 5 6] .- 3
B = [2 -2; 3 -5; -4 7]
@test A*B == [-7 9; -4 9]
@test At_mul_Bt(A, B) == [-6 -11 15; -6 -13 18; -6 -15 21]
A = ones(Int, 2, 100)
B = ones(Int, 100, 3)
@test A*B == [100 100 100; 100 100 100]
A = rand(1:20, 5, 5) .- 10
B = rand(1:20, 5, 5) .- 10
@test At_mul_B(A, B) == A'*B
@test A_mul_Bt(A, B) == A*B'

# Preallocated
C = Array(Int, size(A, 1), size(B, 2))
@test A_mul_B!(C, A, B) == A*B
@test At_mul_B!(C, A, B) == A'*B
@test A_mul_Bt!(C, A, B) == A*B'
@test At_mul_Bt!(C, A, B) == A'*B'

# matrix algebra with subarrays of floats (stride != 1)
A = reshape(float64(1:20),5,4)
Aref = A[1:2:end,1:2:end]
Asub = sub(A, 1:2:5, 1:2:4)
b = [1.2,-2.5]
@test (Aref*b) == (Asub*b)
@test At_mul_B(Asub, Asub) == At_mul_B(Aref, Aref)
@test A_mul_Bt(Asub, Asub) == A_mul_Bt(Aref, Aref)
Ai = A .+ im
Aref = Ai[1:2:end,1:2:end]
Asub = sub(Ai, 1:2:5, 1:2:4)
@test Ac_mul_B(Asub, Asub) == Ac_mul_B(Aref, Aref)
@test A_mul_Bc(Asub, Asub) == A_mul_Bc(Aref, Aref)

# syrk & herk
A = reshape(1:1503, 501, 3).-750.0
res = float64([135228751 9979252 -115270247; 9979252 10481254 10983256; -115270247 10983256 137236759])
@test At_mul_B(A, A) == res
@test A_mul_Bt(A',A') == res
cutoff = 501
A = reshape(1:6*cutoff,2*cutoff,3).-(6*cutoff)/2
Asub = sub(A, 1:2:2*cutoff, 1:3)
Aref = A[1:2:2*cutoff, 1:3]
@test At_mul_B(Asub, Asub) == At_mul_B(Aref, Aref)
Ai = A .- im
Asub = sub(Ai, 1:2:2*cutoff, 1:3)
Aref = Ai[1:2:2*cutoff, 1:3]
@test Ac_mul_B(Asub, Asub) == Ac_mul_B(Aref, Aref)

# Matrix exponential
for elty in (Float32, Float64, Complex64, Complex128)
        A1  = convert(Matrix{elty}, [4 2 0; 1 4 1; 1 1 4])
        eA1 = convert(Matrix{elty}, [147.866622446369 127.781085523181  127.781085523182;
        183.765138646367 183.765138646366  163.679601723179;
        71.797032399996  91.8825693231832 111.968106246371]')
        @test_approx_eq expm(A1) eA1

        A2  = convert(Matrix{elty},
            [29.87942128909879    0.7815750847907159 -2.289519314033932;
            0.7815750847907159 25.72656945571064    8.680737820540137;
            -2.289519314033932   8.680737820540137  34.39400925519054])
        eA2 = convert(Matrix{elty},
            [  5496313853692458.0 -18231880972009236.0 -30475770808580460.0;
             -18231880972009252.0  60605228702221920.0 101291842930249760.0;
             -30475770808580480.0 101291842930249728.0 169294411240851968.0])
        @test_approx_eq expm(A2) eA2

        A3  = convert(Matrix{elty}, [-131 19 18;-390 56 54;-387 57 52])
        eA3 = convert(Matrix{elty}, [-1.50964415879218 -5.6325707998812  -4.934938326092;
        0.367879439109187 1.47151775849686  1.10363831732856;
        0.135335281175235 0.406005843524598 0.541341126763207]')
        @test_approx_eq expm(A3) eA3

        # issue 5116
        A4  = [0 10 0 0; -1 0 0 0; 0 0 0 0; -2 0 0 0]
        eA4 = [-0.999786072879326  -0.065407069689389   0.0   0.0
                0.006540706968939  -0.999786072879326   0.0   0.0
                0.0                 0.0                 1.0   0.0
                0.013081413937878  -3.999572145758650   0.0   1.0]
        @test_approx_eq expm(A4) eA4

        # issue 5116
        A5  = [ 0. 0. 0. 0. ; 0. 0. -im 0.; 0. im 0. 0.; 0. 0. 0. 0.]
        eA5 = [ 1.0+0.0im   0.0+0.0im                 0.0+0.0im                0.0+0.0im
                0.0+0.0im   1.543080634815244+0.0im   0.0-1.175201193643801im  0.0+0.0im
                0.0+0.0im   0.0+1.175201193643801im   1.543080634815243+0.0im  0.0+0.0im
                0.0+0.0im   0.0+0.0im                 0.0+0.0im                1.0+0.0im]
        @test_approx_eq expm(A5) eA5

        # Hessenberg
        @test_approx_eq hessfact(A1)[:H] convert(Matrix{elty},
                        [4.000000000000000  -1.414213562373094  -1.414213562373095
                        -1.414213562373095   4.999999999999996  -0.000000000000000
                                         0  -0.000000000000002   3.000000000000000])
end

# Hermitian matrix exponential
A1 = randn(4,4) + im*randn(4,4)
A2 = A1 + A1'
@test_approx_eq expm(A2) expm(Hermitian(A2))

# matmul for types w/o sizeof (issue #1282)
A = Array(Complex{Int},10,10)
A[:] = complex(1,1)
A2 = A^2
@test A2[1,1] == 20im

# test sparse matrix norms
Ac = sprandn(10,10,.1) + im* sprandn(10,10,.1)
Ar = sprandn(10,10,.1)
Ai = int(ceil(Ar*100))
@test_approx_eq norm(Ac,1)     norm(full(Ac),1)
@test_approx_eq norm(Ac,Inf)   norm(full(Ac),Inf)
@test_approx_eq vecnorm(Ac)    vecnorm(full(Ac))
@test_approx_eq norm(Ar,1)     norm(full(Ar),1)
@test_approx_eq norm(Ar,Inf)   norm(full(Ar),Inf)
@test_approx_eq vecnorm(Ar)    vecnorm(full(Ar))
@test_approx_eq norm(Ai,1)     norm(full(Ai),1)
@test_approx_eq norm(Ai,Inf)   norm(full(Ai),Inf)
@test_approx_eq vecnorm(Ai)    vecnorm(full(Ai))

# scale real matrix by complex type
@test_throws ErrorException scale!([1.0], 2.0im)
@test isequal(scale([1.0], 2.0im),             Complex{Float64}[2.0im])
@test isequal(scale(Float32[1.0], 2.0f0im),    Complex{Float32}[2.0im])
@test isequal(scale(Float32[1.0], 2.0im),      Complex{Float64}[2.0im])
@test isequal(scale(Float64[1.0], 2.0f0im),    Complex{Float64}[2.0im])
@test isequal(scale(Float32[1.0], big(2.0)im), Complex{BigFloat}[2.0im])
@test isequal(scale(Float64[1.0], big(2.0)im), Complex{BigFloat}[2.0im])
@test isequal(scale(BigFloat[1.0], 2.0im),     Complex{BigFloat}[2.0im])
@test isequal(scale(BigFloat[1.0], 2.0f0im),   Complex{BigFloat}[2.0im])

# issue #6450
@test dot({1.0,2.0},{3.5,4.5}) === 12.5

# Issue 7181
A = reshape([1:9], (3, 3))
@test_throws BoundsError diag(A, 3)
@test [7] == diag(A, 2)
@test [4, 8] == diag(A, 1)
@test [1,5,9] == diag(A, 0)
@test [2,6] == diag(A, -1)
@test [3] == diag(A, -2)
@test_throws BoundsError diag(A, -3)


struct Sinus{T} <: Distances.SemiMetric
# struct Sinus{T} <: Distances.UnionSemiMetric
    r::Vector{T}
end

Distances.parameters(d::Sinus) = d.r
@inline Distances.eval_op(::Sinus, a::Real, b::Real, p::Real) = abs2(sinpi(a - b) / p)
@inline (dist::Sinus)(a::AbstractArray, b::AbstractArray) = Distances._evaluate(dist, a, b)
@inline (dist::Sinus)(a::Number, b::Number) = abs2(sinpi(a - b) / first(dist.r))

Distances.result_type(::Sinus{T}, Ta::Type, Tb::Type) where {T} = promote_type(T, Ta, Tb)

@inline function Distances._evaluate(d::Sinus, a::AbstractVector, b::AbstractVector) where {T}
    @boundscheck if (length(a) != length(b)) || length(a) != length(d.r)
        throw(DimensionMismatch("Dimensions of the inputs are not matching : a = $(length(a)), b = $(length(b)), r = $(length(d.r))"))
    end
    return sum(abs2, sinpi.(a - b) ./ d.r)
end