text:
"""
    AbstractCallback

Abstract type of callback functions used in training.
"""
abstract AbstractCallback

"""
    AbstractBatchCallback

Abstract type of callbacks to be called every mini-batch.
"""
abstract AbstractBatchCallback <: AbstractCallback

"""
    AbstractEpochCallback

Abstract type of callbacks to be called every epoch.
"""
abstract AbstractEpochCallback <: AbstractCallback

type BatchCallback <: AbstractBatchCallback
  frequency :: Int
  call_on_0 :: Bool
  callback  :: Function
end

"""
    every_n_batch(callback :: Function, n :: Int; call_on_0 = false)

A convenient function to construct a callback that runs every `n` mini-batches.

# Arguments
* `call_on_0::Bool`: keyword argument, default false. Unless set, the callback
          will *not* be run on batch 0.

For example, the [`speedometer`](@ref) callback is defined as

```julia
every_n_iter(frequency, call_on_0=true) do state :: OptimizationState
  if state.curr_batch == 0
    # reset timer
  else
    # compute and print speed
  end
end
```

See also [`every_n_epoch`](@ref) and [`speedometer`](@ref).
"""
function every_n_batch(callback :: Function, n :: Int; call_on_0 :: Bool = false)
  BatchCallback(n, call_on_0, callback)
end
@compat function (cb :: BatchCallback)(state :: OptimizationState)
  if state.curr_batch == 0
    if cb.call_on_0
      cb.callback(state)
    end
  elseif state.curr_batch % cb.frequency == 0
    cb.callback(state)
  end
end

"""
    speedometer(; frequency=50)

Create an `AbstractBatchCallback` that measure the training speed
   (number of samples processed per second) every k mini-batches.

# Arguments
* `frequency::Int`: keyword argument, default 50. The frequency (number of
          min-batches) to measure and report the speed.
"""
function speedometer(;frequency::Int=50)
  cl_tic = 0
  every_n_batch(frequency, call_on_0=true) do state :: OptimizationState
    if state.curr_batch == 0
      # reset timer
      cl_tic = time()
    else
      speed = frequency * state.batch_size / (time() - cl_tic)
      info(format("Speed: {1:>6.2f} samples/sec", speed))
      cl_tic = time()
    end
  end
end

type EpochCallback <: AbstractEpochCallback
  frequency :: Int
  call_on_0 :: Bool
  callback  :: Function
end

"""
    every_n_epoch(callback :: Function, n :: Int; call_on_0 = false)

A convenient function to construct a callback that runs every `n` full data-passes.

* `call_on_0::Int`: keyword argument, default false. Unless set, the callback
          will *not* be run on epoch 0. Epoch 0 means no training has been performed
          yet. This is useful if you want to inspect the randomly initialized model
          that has not seen any data yet.

See also [`every_n_iter`](@ref).
"""
function every_n_epoch(callback :: Function, n :: Int; call_on_0 :: Bool = false)
  EpochCallback(n, call_on_0, callback)
end
@compat function (cb :: EpochCallback){T<:Real}(model :: Any, state :: OptimizationState, metric :: Vector{Tuple{Base.Symbol, T}})
  if state.curr_epoch == 0
    if cb.call_on_0
      cb.callback(model, state, metric)
    end
  elseif state.curr_epoch % cb.frequency == 0
    cb.callback(model, state, metric)
  end
end

"""
    do_checkpoint(prefix; frequency=1, save_epoch_0=false)

Create an `AbstractEpochCallback` that save checkpoints of the model to disk.
The checkpoints can be loaded back later on.

# Arguments
* `prefix::AbstractString`: the prefix of the filenames to save the model. The model
          architecture will be saved to prefix-symbol.json, while the weights will be saved
          to prefix-0012.params, for example, for the 12-th epoch.
* `frequency::Int`: keyword argument, default 1. The frequency (measured in epochs) to
          save checkpoints.
* `save_epoch_0::Bool`: keyword argument, default false. Whether we should save a
          checkpoint for epoch 0 (model initialized but not seen any data yet).
"""
function do_checkpoint(prefix::AbstractString; frequency::Int=1, save_epoch_0=false)
  mkpath(dirname(prefix))
  every_n_epoch(frequency, call_on_0=save_epoch_0) do model, state, metric
    save_checkpoint(model, prefix, state)
  end
end


text:
int main() {
   int !a;
   return 1;
}

ext: jl

max_stars_repo_name: carlostome/Javalette-Compiler

max_issues_repo_name: carlostome/Javalette-Compiler

max_forks_repo_name: carlostome/Javalette-Compiler


text:
module HSA

include("discovery.jl")

# Some things are always available
include("emulation/emulation.jl")

if has_libhsa()
info("HSA library found.")
include("binding/binding.jl")

if HSA.has_ext_finalization()
info("HSA Finalization Extension is available.")
include("binding/ext_finalization.jl")

if has_hsa_codegen()
info("Experimental HSA CodeGen found, enabling.")
include("codegen/codegen.jl")
end
end
end
end


text:
module CUDA

using ..CuArrays

using CuArrays: CUDNN
include("curnn.jl")
include("cudnn.jl")

end


text:
using BenchmarkTools
using FinancialMonteCarlo, CUDA
CUDA.allowscalar(false)
S0 = 100.0f0;
K = 100.0f0;
r = 0.02f0;
T = 1.0f0;
d = 0.01f0;
D = 90.0f0;

Nsim = 10000;
Nstep = 30;
sigma = 0.2f0;
mc = MonteCarloConfiguration(Nsim, Nstep);
mc_2 = MonteCarloConfiguration(Nsim, Nstep, FinancialMonteCarlo.CudaMode());
toll = 1e-3;

rfCurve = ZeroRate(r);

FwdData = Forward(T)
EUData = EuropeanOption(T, K)
AMData = AmericanOption(T, K)
BarrierData = BarrierOptionDownOut(T, K, D)
AsianFloatingStrikeData = AsianFloatingStrikeOption(T)
AsianFixedStrikeData = AsianFixedStrikeOption(T, K)
Model = BlackScholesProcess(sigma, Underlying(S0, d));

@show "CUDA_2 fwd"
@btime FwdPrice = pricer(Model, rfCurve, mc_2, FwdData);
@show "CUDA_2 eu"
@btime EuPrice = pricer(Model, rfCurve, mc_2, EUData);


text:
export scan_file

"""
    scan_file(file::String, keys::Array{String, 1}, blocksize::Int;
              chunksize::Int = CHUNKSIZE,
              verbosity::Int = 1)

Scan `file` for header fields in `keys`, and return a SeisCon object containing
the metadata summaries in `blocksize` groups of traces. Load `chunksize` MB of `file`
into memory at a time.

If the number of traces in `file` are not divisible by `blocksize`, the last block will
summarize the remaining traces.

`verbosity` set to 0 silences updates on the current file being scanned.

# Example

    s = scan_file('testdata.segy', ["SourceX", "SourceY"], 300)
"""
function scan_file(file::String, keys::Array{String, 1}, blocksize::Int;
                   chunksize::Int = CHUNKSIZE,
                   verbosity::Int = 1)

    # Put fileheader in memory and read
    verbosity==1 && println("Scanning ... $file")
    s = open(file)
    fh = read_fileheader(s)

    # Calc number of blocks
    fsize = position(seekend(s))
    mem_trace = 240 + fh.bfh.ns*4
    mem_block = blocksize*mem_trace
    ntraces_file = Int((fsize - 3600)/mem_trace)
    nblocks_file = Int(ceil(ntraces_file/blocksize))
    scan = Array{BlockScan,1}(undef, nblocks_file)
    count = 1

    # Blocks to load per chunk
    max_blocks_per_chunk = Int(floor(chunksize*MB2B/mem_block))

    # Read at most one full chunk into buffer
    seek(s, 3600)

    # For each chunk
    for c in 1:max_blocks_per_chunk:nblocks_file

        count = scan_chunk!(s, max_blocks_per_chunk, mem_block, mem_trace,
                            keys, file, scan, count)

    end # c

    return SeisCon(fh.bfh.ns, fh.bfh.DataSampleFormat, scan)

end

"""
    scan_file(file::String, keys::Array{String, 1};
              chunksize::Int = CHUNKSIZE,
              verbosity::Int = 1)

Scan `file` for header fields in `keys`, and return a SeisCon object containing
the metadata summaries in single-source groups of traces. Load `chunksize` MB of `file`
into memory at a time.

# Example

    s = scan_file('testdata.segy', ["SourceX", "SourceY"])

"""
function scan_file(file::String, keys::Array{String, 1};
                   chunksize::Int = 10*CHUNKSIZE,
                   verbosity::Int = 1)

    # Put fileheader in memory and read
    verbosity==1 && println("Scanning ... $file")
    s = open(file)
    fh = read_fileheader(s)

    # Add src keys if necessary
    "SourceX" in keys ? nothing : push!(keys, "SourceX")
    "SourceY" in keys ? nothing : push!(keys, "SourceY")

    # Calc number of blocks
    fsize = filesize(file)
    mem_trace = 240 + fh.bfh.ns*4
    ntraces_file = Int((fsize - 3600)/mem_trace)
    scan = Array{BlockScan,1}(undef, 0)
    seek(s, 3600)
    traces_per_chunk = Int(floor(chunksize*MB2B/mem_trace))

    mem_chunk = traces_per_chunk*mem_trace
    fl_eof = false
    while !eof(s)
        scan_shots!(s, mem_chunk, mem_trace, keys, file, scan, fl_eof)
    end

    close(s)
    return SeisCon(fh.bfh.ns, fh.bfh.DataSampleFormat, scan[1:end])

end


text:
### Following R. Nelsen 'An Introduction to Copulas', Springer Science & Business Media, 1999 - 216,
### for bivariate Archimedean copulas `C(u₁,u₂)` data can be generated as follow:
### * draw `u₁ = rand()`,
### * define `w = ∂C(u₁, u₂)\∂u₁` and inverse `u₂ = f(w, u₁)`,
### * draw  `w = rand()`
### * return a pair u₁, u₂.

### This method is applied here for Clayton, Frank and Ali-Mikhail-Haq copula. If we use
### this method recursively, we can get `n`-variate data with uniform marginals on
### `[0,1]`, where each neighbour pair
### of marginals `uᵢ uⱼ` for `j = i+1` are draw form a bivariate subcopula with
### parameter `θᵢ`, the only condition for `θᵢ`
### is such as for a corresponding bivariate copula.

"""
  rand2cop(u1::Vector{Float64}, θ::Union{Int, Float64}, copula::String)

Returns vector of data generated using copula::String given vector of uniformly
distributed u1 and copula parameter θ.
"""
function rand2cop(u1::Vector{Float64}, θ::Union{Int, Float64}, copula::String)
  w = rand(length(u1))
  copula in ["clayton", "amh", "frank"] || throw(AssertionError("$(copula) copula is not supported"))
  if copula == "clayton"
    return (u1.^(-θ).*(w.^(-θ/(1+θ)) .-1) .+1).^(-1/θ)
  elseif copula == "frank"
    return -1/θ*log.(1 .+(w*(1-exp(-θ)))./(w.*(exp.(-θ*u1) .-1) .-exp.(-θ*u1)))
  elseif copula == "amh"
    a = 1 .-u1
    b = 1 .-θ .*(1 .+2 .*a .*w)+2*θ^2*a.^2 .*w
    c = 1 .-θ .*(2 .-4*w .+4 .*a .*w)+θ.^2 .*(1 .-4 .*a .*w+4 .*a.^2 .*w)
    return 2*w .*(a .*θ .-1).^2 ./(b+sqrt.(c))
  end
end

"""
  chaincopulagen(t::Int, θ::Union{Vector{Float64}, Vector{Int}}, copula::String;
                                              rev::Bool = false, cor::String = "")

Returns: t x n Matrix{Float}, t realisations of n variate data, where n = length(θ)+1.
To generate data uses chain of bivariate Archimedean one parameter copulas.

Following copula families are supported: clayton, frank and amh -- Ali-Mikhail-Haq.

If rev == true, reverse the copula output i.e. u → 1-u (we call it reversed copula).
It cor == pearson, kendall, uses correlation coeficient as a parameter

```jldoctest
julia> Random.seed!(43);

julia> chaincopulagen(10, [4., 11.], "frank")
10×3 Array{Float64,2}:
 0.180975  0.386303   0.879254
 0.775377  0.247895   0.144803
 0.888934  0.426854   0.772457
 0.924876  0.395564   0.223155
 0.408278  0.139002   0.142997
 0.912603  0.901252   0.949828
 0.828727  0.0295759  0.0897796
 0.400537  0.0337673  0.27872
 0.429437  0.462771   0.425435
 0.955881  0.953623   0.969038
```
"""
VFI = Union{Vector{Float64}, Vector{Int}}
chaincopulagen(t::Int, θ::VFI, copula::String; rev::Bool = false, cor::String = "") =
chaincopulagen(t, θ, [fill(copula, length(θ))...] ; rev = rev, cor= cor)

function chaincopulagen(t::Int, θ::VFI, copula::Vector{String}; rev::Bool = false, cor::String = "")
  if ((cor == "Spearman") | (cor == "Kendall"))
    θ = map(i -> usebivρ(θ[i], copula[i], cor), 1:length(θ))
  else
    map(i -> testbivθ(θ[i], copula[i]), 1:length(θ))
  end
  u = rand(t,1)
  for i in 1:length(θ)
    u = hcat(u, rand2cop(u[:, i], θ[i], copula[i]))
  end
  rev ? 1 .-u : u
end

"""
  testbivθ(θ::Union{Float64}, copula::String)

Tests bivariate copula parameter

clayton bivariate sub-copulas with parameters (θᵢ ≥ -1) ^ ∧ (θᵢ ≠ 0).
amh -- Ali-Mikhail-Haq bivariate sub-copulas with parameters -1 ≥ θᵢ ≥ 1
frank bivariate sub-copulas with parameters (θᵢ ≠ 0)
"""
function testbivθ(θ::Union{Float64, Int}, copula::String)
  !(0. in θ)|(copula == "amh") || throw(DomainError("not supported for θ = 0"))
  if copula == "clayton"
    θ >= -1 || throw(DomainError("not supported for θ < -1"))
  elseif copula == "amh"
    -1 <= θ <= 1|| throw(DomainError("amh biv. copula supported only for -1 ≤ θ ≤ 1"))
  end
  Nothing
end

"""
Returns Float64, a copula parameter given a pearson or kendall correlation

For clayton or frank copula correlation fulfulling (-1 > ρᵢ > 1) ∧ (ρᵢ ≠ 0)
For amh copula pearson correlation fulfilling -0.2816 > ρᵢ >= .5. while kendall -0.18 < τ < 1/3
"""
function usebivρ(ρ::Float64, copula::String, cor::String)
  if copula == "amh"
      -0.2816 < ρ <= 0.5 || throw(DomainError("correlation coeficiant must fulfill -0.2816 < ρ <= 0.5"))
    if cor == "Kendall"
      -0.18 < ρ < 1/3 || throw(DomainError("correlation coeficiant must fulfill -0.2816 < ρ <= 1/3"))
    end
  else
    -1 < ρ < 1 || throw(DomainError("correlation coeficiant must fulfill -1 < ρ < 1"))
    !(0. in ρ) || throw(DomainError("not supported for ρ = 0"))
  end
  (cor == "Kendall") ? τ2θ(ρ, copula) : ρ2θ(ρ, copula)
end

# chain frechet copulas

"""
  chainfrechetcopulagen(t::Int, α::Vector{Float64}, β::Vector{Float64} = zeros(α))

Retenares data from a chain of bivariate two parameter frechet copuls with parameters
vectors α and β, such that ∀ᵢ 0 α[i] + β[i] ≤1 α[i] > 0, and β[i] > 0 |α| = |β|

```jldoctest
julia> Random.seed!(43)

julia> julia> chainfrechetcopulagen(10, [0.6, 0.4], [0.3, 0.5])
10×3 Array{Float64,2}:
 0.996764  0.996764  0.996764
 0.204033  0.795967  0.204033
 0.979901  0.979901  0.0200985
 0.120669  0.879331  0.120669
 0.453027  0.453027  0.453027
 0.800909  0.199091  0.800909
 0.54892   0.54892   0.54892
 0.933832  0.933832  0.0661679
 0.396943  0.396943  0.396943
 0.804096  0.851275  0.955881
```
"""
function chainfrechetcopulagen(t::Int, α::Vector{Float64}, β::Vector{Float64} = zero(α))
  length(α) == length(β) || throw(AssertionError("different lengths of parameters"))
  minimum(α) >= 0 || throw(DomainError("negative α parameter"))
  minimum(β) >= 0 || throw(DomainError("negative β parameter"))
  maximum(α+β) <= 1 || throw(DomainError("α[i] + β[i] > 0"))
  fncopulagen(α, β, rand(t, length(α)+1))
end

"""

  fncopulagen(α::Vector{Float64}, β::Vector{Float64}, u::Matrix{Float64})

```jldoctest

julia> fncopulagen(2, [0.2, 0.4], [0.1, 0.1], [0.2 0.4 0.6; 0.3 0.5 0.7])
2×3 Array{Float64,2}:
 0.6  0.4  0.2
 0.7  0.5  0.3

```
"""
function fncopulagen(α::Vector{Float64}, β::Vector{Float64}, u::Matrix{Float64})
  p = invperm(sortperm(u[:,1]))
  u = u[:,end:-1:1]
  lx = floor.(Int, size(u,1).*α)
  li = floor.(Int, size(u,1).*β) + lx
  for j in 1:size(u, 2)-1
    u[p[1:lx[j]],j+1] = u[p[1:lx[j]], j]
    r = p[lx[j]+1:li[j]]
    u[r,j+1] = 1 .-u[r,j]
  end
  u
end


text:
pushfirst!(LOAD_PATH, joinpath(@__DIR__, ".."))
using Syncer
Syncer.command_main()


text:
#---------------------------------------------------------------------#
#This function computes the Norm of two vectors
#Written by F.X. Giraldo on April 19, 2019
#           Department of Applied Mathematics
#           Naval Postgraduate School
#           Monterey; CA 93943-5216
#---------------------------------------------------------------------#
using LinearAlgebra

function compute_norms(qn,qe,npoin,DFloat)

    norms=zeros(DFloat,3)
    norms[1]=norm(qn-qe,1)/norm(qe,1)
    norms[2]=norm(qn-qe,2)/norm(qe,2)
    norms[3]=norm(qn-qe,Inf)/norm(qe,Inf)
    return (norms)

end


text:
using ArviZ, LinearAlgebra

# https://arviz-devs.github.io/ArviZ.jl/stable/quickstart/#Convert-to-InferenceData
function to_arviz(prior_model::Turing.Model, conditioned_model::Turing.Model, posterior_samples::MCMCChains.Chains)
    # Sample from the prior of the latent variables, i.e. everything
    # except `complaints`.
    prior_samples = sample(conditioned_model, Prior(), 1_000, progress=false);
    # Sample "predictions" of `complaints` under `prior_samples`.
    prior_predictive = predict(prior_model, prior_samples);
    # Sample predictions of `complains` under `posterior_samples`.
    posterior_predictive = predict(prior_model, posterior_samples);
    # Compute p(constraints[i] ∣ params...) where `params` are given
    # by samples form `posterior_samples`.
    loglikelihoods = Turing.pointwise_loglikelihoods(
        conditioned_model,
        MCMCChains.get_sections(posterior_samples, :parameters)
    );

    # Ensure the ordering of the loglikelihoods matches the ordering of `posterior_predictive`.
    ynames = string.(keys(posterior_predictive));
    loglikelihoods_vals = getindex.(Ref(loglikelihoods), ynames);
    # Reshape into `(nchains, nsamples, size(y)...)`
    loglikelihoods_arr = permutedims(cat(loglikelihoods_vals...; dims=3), (2, 1, 3));

    # Finally convert into `InferenceData`.
    return from_mcmcchains(
        posterior_samples;
        posterior_predictive=posterior_predictive,
        log_likelihood=Dict("complaints" => loglikelihoods_arr),
        prior=prior_samples,
        prior_predictive=prior_predictive,
        observed_data=Dict("complaints" => conditioned_model.args.complaints),
        library="Turing",
    )
end


text:
include("helper.jl")
using Libdl
# Locate and/or download Yeppp!
if isfile("$(pkgdir)deps/downloads/yeppp-1.0.0.tar.bz2") || (Libdl.find_library(["libyeppp"]) != "")
else
    @info("====== Installing Yeppp! into local directory ======")
    trycmd(`mkdir $(pkgdir)deps/downloads`, err="Unable to create Vectorize/deps/downloads")
    trycmd(`mkdir $(pkgdir)deps/src`, err="Unable to create Vectorize/src/")
    trycmd(`mkdir $(pkgdir)deps/src/yeppp`, err="Unable to create Vectorize/src/yeppp")
    yeppploc = "http://bitbucket.org/MDukhan/yeppp/downloads/yeppp-1.0.0.tar.bz2"
    @static if Sys.isapple()
        run(pipeline(`curl -L $(yeppploc)`, stdout="$(pkgdir)deps/downloads/yeppp-1.0.0.tar.bz2"))
    else
        trycmd(`wget -P downloads $(yeppploc)`, err="Unable to download Yeppp!")
    end
    trycmd(`tar -xjvf $(pkgdir)deps/downloads/yeppp-1.0.0.tar.bz2 -C $(pkgdir)deps/src/yeppp --strip-components=1`)
    @info("====== Successfully installed Yeppp! ======")
end


text:
using BatchedRoutines
using Test

@testset "Testing BLAS" begin
    include("blas.jl")
end


text:
# PREAMBLE

# PKG_SETUP

# ## Setup

using ContactImplicitMPC
using LinearAlgebra
using Quaternions
using BenchmarkTools

# ## Visualizer
vis = ContactImplicitMPC.Visualizer()
ContactImplicitMPC.open(vis)

@show Threads.nthreads()

# include("continuous_policy.jl")
# include("continuous_policy_v2.jl")
include("continuous_policy_v3.jl")
# include("continuous_policy_gains.jl")

# ## Simulation
s = get_simulation("centroidal_quadruped", "flat_3D_lc", "flat")
model = s.model
env = s.env

# ## Reference Trajectory
ref_traj = deepcopy(get_trajectory(s.model, s.env,
	joinpath(@__DIR__, "reference/inplace_crawl_v10.jld2"),
    load_type = :split_traj_alt));

H = ref_traj.H
h = ref_traj.h

# ## MPC setup
N_sample = 5
H_mpc = 5
h_sim = h / N_sample
H_sim = 200
κ_mpc = 1.0e-3

v0 = 0.00
function get_stride(model::CentroidalQuadruped, traj::ContactTraj; v0=0.2*v0)
	stride = zeros(model.nq)
	stride[[1,7,10,13,16]] .+= 1.0 * v0 * traj.h * traj.H
	stride[1] = 1.0 * v0 * traj.h * traj.H
	return stride
end
obj = TrackingVelocityObjective(model, env, H_mpc,
    # v = [Diagonal(1e-3 * [[1,1,1]; 1e+3*[1,1,1]; fill([1,1,1], 4)...]) for t = 1:H_mpc],
	v = [Diagonal([2e+1*[1,1,1]; 1e+3*[1,1,1]; fill(1e0*[1,1,0.3], 4)...]) for t = 1:H_mpc],
	q = [Diagonal([1e-0*[1e-1,1e-1,1]; 3e-0*[1,1,1]; fill(1e+0*[0.2,0.2,20], 4)...]) for t = 1:H_mpc],
	u = [Diagonal(1e-4 * vcat(fill([1,1,0.1], 4)...)) for t = 1:H_mpc],
	v_target = [ref_traj.h * [v0;0;0; 0;0;0; v0;0;0; v0;0;0; v0;0;0; v0;0;0] for t = 1:H_mpc],
	)

p = ci_mpc_policy(ref_traj, s, obj,
    H_mpc = H_mpc,
    N_sample = N_sample,
    κ_mpc = κ_mpc,
	mode = :configuration,
	ip_opts = InteriorPointOptions(
					undercut = 5.0,
					κ_tol = κ_mpc,
					r_tol = 1.0e-4, # TODO maybe relax this parameter
					diff_sol = true,
					solver = :empty_solver,
					max_time = 1e5),
    n_opts = NewtonOptions(
        r_tol = 3e-6,
        max_time=5e-2,
		solver=:ldl_solver,
        threads=false,
        verbose=false,
        max_iter = 5),
    mpc_opts = CIMPCOptions(
		# live_plotting=true
		gains=true,
		control_gain_scaling=100.0,
		velocity_gain_scaling=100.0,
		));

# ## Disturbances
w = [ref_traj.h * [00; 00; -00.0] for i=1:H_sim/N_sample]
d = open_loop_disturbances(w, N_sample)

# ## Initial conditions
q1_sim, v1_sim = initial_conditions(ref_traj);

# ## Simulator
sim = simulator(s, H_sim, h=h_sim, policy=p, dist=d);

# ## Simulate
q1_sim0 = deepcopy(q1_sim)
RoboDojo.simulate!(sim, q1_sim0, v1_sim)

# ## Visualize
set_light!(vis)
set_floor!(vis, grid=true)
set_background!(vis)
# anim = visualize!(vis, model, ref_traj.q; Δt=h)
anim = visualize!(vis, model, sim.traj.q; Δt=h_sim)

# # ## Timing result
# # Julia is [JIT-ed](https://en.wikipedia.org/wiki/Just-in-time_compilation) so re-run the MPC setup through Simulate for correct timing results.
process!(sim.stats, N_sample) # Time budget
H_sim * h_sim / sum(sim.stats.policy_time) # Speed ratio
plot(sim.stats.policy_time, xlabel="timestep", ylabel="mpc time (s)",
	ylims=[-0.001, 0.1],
	label="", linetype=:steppost)

plt = plot()
plot!(plt, hcat(Vector.([(sim.traj.q[i+1][1:1] - sim.traj.q[i][1:1]) / sim.h for i=1:H_sim])...)')
plot!(plt, hcat(Vector.([(sim.traj.q[i+1][7:7] - sim.traj.q[i][7:7]) / sim.h for i=1:H_sim])...)')
plot!(plt, hcat(Vector.([(sim.traj.q[i+1][10:10] - sim.traj.q[i][10:10]) / sim.h for i=1:H_sim])...)')
plot!(plt, hcat(Vector.([(sim.traj.q[i+1][13:13] - sim.traj.q[i][13:13]) / sim.h for i=1:H_sim])...)')
plot!(plt, hcat(Vector.([(sim.traj.q[i+1][16:16] - sim.traj.q[i][16:16]) / sim.h for i=1:H_sim])...)')


text:
"""
    pdf(
        f::AbstractArray{<:AbstractFloat, N ∈ (3, 4)},
        mask::AbstractArray{Bool, 3},
        vsz::NTuple{3, Real};
        W::Union{Nothing, AbstractArray{<:AbstractFloat, M ∈ (3, N)}} = nothing,
        pad::NTuple{3, Integer} = (0, 0, 0),
        bdir::NTuple{3, Real} = (0, 0, 1),
        Dkernel::Symbol = :i,
        lambda::Real = 0,
        tol::Real = 1e-5,
        maxit::Integer = ceil(sqrt(numel(mask))),
        verbose::Bool = false
    ) typeof(similar(f))

Projection onto dipole fields (PDF) [1].

### Arguments
- `f::AbstractArray{<:AbstractFloat, N ∈ (3, 4)}`: unwrapped (multi-echo) field/phase
- `mask::AbstractArray{Bool, 3}`: binary mask of region of interest
- `vsz::NTuple{3, Real}`: voxel size for dipole kernel

### Keywords
- `W::Union{Nothing, AbstractArray{<:AbstractFloat, M ∈ (3, N)}} = nothing`:
    data fidelity weights
- `pad::NTuple{3, Integer} = (0, 0, 0)`: zero padding array
    - `< 0`: no padding
    - `≥ 0`: minimum padding to fast fft size
- `bdir::NTuple{3, Real} = (0, 0, 1)`: unit vector of B field direction
- `Dkernel::Symbol = :i`: dipole kernel method
- `lambda::Real = 0`: regularization parameter
- `tol::Real = 1e-5`: stopping tolerance for iterative solver
- `maxit::Integer = ceil(sqrt(length(mask)))`: maximum number of iterations for
    iterative solver
- `verbose::Bool = false`: print convergence information

### Returns
- `typeof(similar(f))`: background corrected local field/phase

### References
[1] Liu T, Khalidov I, de Rochefort L, Spincemaille P, Liu J, Tsiouris AJ,
    Wang Y. A novel background field removal method for MRI using projection
    onto dipole fields. NMR in Biomedicine. 2011 Nov;24(9):1129-36.
"""
function pdf(
    f::AbstractArray{<:AbstractFloat, N},
    mask::AbstractArray{Bool, 3},
    vsz::NTuple{3, Real};
    W::Union{Nothing, AbstractArray{<:AbstractFloat}} = nothing,
    pad::NTuple{3, Integer} = (0, 0, 0),
    bdir::NTuple{3, Real} = (0, 0, 1),
    Dkernel::Symbol = :i,
    lambda::Real = 0,
    tol::Real = 1e-5,
    maxit::Integer = ceil(Int, sqrt(length(mask))),
    verbose::Bool = false
) where {N}
    N ∈ (3, 4) || throw(ArgumentError("arrays must be 3d or 4d, got $(N)d"))
    return _pdf!(
        tzero(f), f, mask, vsz, W, pad, bdir, Dkernel, lambda, tol, maxit, verbose
    )
end

function _pdf!(
    fl::AbstractArray{<:AbstractFloat, N},
    f::AbstractArray{T, N},
    mask::AbstractArray{Bool, 3},
    vsz::NTuple{3, Real},
    W::Union{Nothing, AbstractArray{<:AbstractFloat}},
    pad::NTuple{3, Integer},
    bdir::NTuple{3, Real},
    Dkernel::Symbol,
    lambda::Real,
    tol::Real,
    maxit::Integer,
    verbose::Bool = false
) where {T<:AbstractFloat, N}
    N ∈ (3, 4) || throw(ArgumentError("arrays must be 3d or 4d, got $(N)d"))

    size(fl) == size(f) || throw(DimensionMismatch())
    size(mask) == size(f)[1:3] || throw(DimensionMismatch())

    if W !== nothing
        ndims(W) ∈ (3, 4) || throw(ArgumentError("arrays must be 3d or 4d, got $(M)d"))
        size(W)[1:3] == size(f)[1:3] || throw(DimensionMismatch())
        ndims(W) == 4 && size(W, 4) == size(f, 4) || throw(DimensionMismatch())
    end

    Dkernel ∈ (:k, :kspace, :i, :ispace) ||
        throw(ArgumentError("Dkernel must be one of :k, :kspace, :i, :ispace, got :$(Dkernel)"))

    # pad to fast fft size
    fp = padfastfft(@view(f[:,:,:,1]), pad, rfft=true)
    m = padfastfft(mask, pad, rfft=true)

    # init vars and fft plans
    sz0 = size(mask)
    sz = size(m)
    sz_ = (sz[1]>>1 + 1, sz[2], sz[3])

    M̃ = similar(m)
    MW = similar(fp)

    D = Array{T, 3}(undef, sz_)
    F̂ = Array{complex(T), 3}(undef, sz_)

    FFTW.set_num_threads(num_cores())
    P = plan_rfft(fp)
    iP = inv(P)

    # LSMR setup
    # move up here so we can use LSMR's `b` as a tmp variable
    A = LinearMap{T}(
        (Av, v) -> _A_pdf!(Av, v, MW, iP, D, F̂, P, M̃),
        (Atu, u) -> _At_pdf!(Atu, u, M̃, iP, D, F̂, P, MW),
        length(fp),
        issymmetric = false,
        ishermitian = false,
        isposdef = true,
        ismutating = true
    )

    WS = LSMRWorkspace(A)
    b  = reshape(WS.b, sz)

    # get dipole kernel
    D = _dipole_kernel!(D, F̂, b, sz0, vsz, bdir, P, Dkernel, :rfft)

    # background mask
    @inbounds @batch for I in eachindex(m)
        M̃[I] = !m[I]
    end

    # pre-compute mask*weights
    if W === nothing
        # no weights
        MW = _tcopyto!(MW, m)

    elseif ndims(W) == 3
        # same weights for all echoes
        b = padarray!(b, W)
        @inbounds @batch for I in eachindex(MW)
            MW[I] = m[I] * b[I]
        end

    elseif ndims(W) == 4
        # weights for each echo. compute inside for loop
        # VOID
    end

    @inbounds for t in axes(f, 4)
        if verbose && size(f, 4) > 1
            @printf("Echo: %d/%d\n", t, size(f, 4))
        end

        if t > 1
            fp = padarray!(fp, @view(f[:,:,:,t]))
        end

        # mask * weight if we can't precompute
        if W !== nothing
            if ndims(W) == 4
                b = padarray!(b, @view(W[:,:,:,t]))
                @batch for I in eachindex(MW)
                    MW[I] = m[I] * b[I]
                end
            end
        end

        # rhs
        @batch for I in eachindex(b)
            b[I] = MW[I] * fp[I]
        end

        # solve
        lsmr!(WS; lambda=lambda, atol=tol, btol=tol, maxit=maxit, verbose=verbose)

        # compute background fields
        xb = reshape(WS.x, sz)
        @batch for I in eachindex(xb)
            xb[I] *= M̃[I]
        end

        F̂ = mul!(F̂, P, xb)
        @batch for I in eachindex(F̂)
            F̂[I] *= D[I]
        end

        xb = mul!(xb, iP, F̂)
        @batch for I in eachindex(fp)
            fp[I] = m[I] * (fp[I] - xb[I])
        end

        if verbose
            println()
        end

        unpadarray!(@view(fl[:,:,:,t]), fp)
    end

    return fl
end

function _A_pdf!(Av, v, W, iP, D, F̂, P, M̃)
    v = reshape(v, size(W))
    x = reshape(Av, size(W))

    @inbounds @batch for I in eachindex(x)
        x[I] = M̃[I] * v[I]
    end

    F̂ = mul!(F̂, P, x)
    @inbounds @batch for I in eachindex(F̂)
        F̂[I] *= D[I]
    end

    x = mul!(x, iP, F̂)
    @inbounds @batch for I in eachindex(x)
        x[I] *= W[I]
    end

    return Av
end

function _At_pdf!(Atu, u, M̃, iP, D, F̂, P, W)
    u = reshape(u, size(W))
    x = reshape(Atu, size(W))

    @inbounds @batch for I in eachindex(x)
        x[I] = W[I] * u[I]
    end

    F̂ = mul!(F̂, P, x)
    @inbounds @batch for I in eachindex(F̂)
        F̂[I] *= D[I] # conj(D), D is real
    end

    x = mul!(x, iP, F̂)
    @inbounds @batch for I in eachindex(x)
        x[I] *= M̃[I]
    end

    return Atu
end


text:
using Test
import Forecast: Δt, nΔt

@testset "utils_datetime" begin

    #Date
    ts = [Date("2000-01-01"), Date("2000-01-2")]
    @test_throws AssertionError Δt(ts)

    ts = [Date("2000-01-01"), Date("2000-01-02"),Date("2000-01-03")]
    @test Δt(ts) == Day(1)

    ts = [Date("2000-01-01"), Date("2000-01-11"),Date("2000-01-21")]
    @test Δt(ts) == Day(10)

    ts = [Date("2000-01-04"), Date("2000-01-11"),Date("2000-01-18")]
    @test Δt(ts) == Week(1)

    ts = [Date("2000-01-01"), Date("2000-02-01"),Date("2000-03-01")]
    @test Δt(ts) == Month(1)

    ts = [Date("2000-01-01"), Date("2000-03-01"),Date("2000-05-01")]
    @test Δt(ts) == Month(2)

    ts = [Date("2000-01-01"), Date("2001-01-01"),Date("2002-01-01")]
    @test Δt(ts) == Year(1)

    #DateTime
    ts = [DateTime("2000-01-01"), DateTime("2000-01-2")]
    @test_throws AssertionError Δt(ts)

    ts = [DateTime("2000-01-01T00:00:00"),
          DateTime("2000-01-02T00:00:00"),
          DateTime("2000-01-03T00:00:00")]
    @test Δt(ts) == Day(1)

    ts = [DateTime("2000-01-04T00:00:00"),
          DateTime("2000-01-11T00:00:00"),
          DateTime("2000-01-18T00:00:00")]
    @test Δt(ts) == Week(1)

    ts = [DateTime("2000-01-01T00:00:00"),
          DateTime("2000-02-01T00:00:00"),
          DateTime("2000-03-01T00:00:00")]
    @test Δt(ts) == Month(1)

    ts = [DateTime("2000-01-01T00:00:00"),
          DateTime("2001-01-01T00:00:00"),
          DateTime("2002-01-01T00:00:00")]
    @test Δt(ts) == Year(1)

    ts = [DateTime("2000-01-01T01:00:00"),
          DateTime("2000-01-01T02:00:00"),
          DateTime("2000-01-01T03:00:00")]
    @test Δt(ts) == Hour(1)

    ts = [DateTime("2000-01-01T01:00:00"),
          DateTime("2000-01-01T03:00:00"),
          DateTime("2000-01-01T05:00:00")]
    @test Δt(ts) == Hour(2)

    ts = [DateTime("2000-01-01T00:01:00"),
          DateTime("2000-01-01T00:02:00"),
          DateTime("2000-01-01T00:03:00")]
    @test Δt(ts) == Minute(1)

    ts = [DateTime("2000-01-01T00:00:01"),
          DateTime("2000-01-01T00:00:02"),
          DateTime("2000-01-01T00:00:03")]
    @test Δt(ts) == Second(1)

    ts = [DateTime("2000-01-01T00:00:00.001"),
          DateTime("2000-01-01T00:00:00.002"),
          DateTime("2000-01-01T00:00:00.003")]
    @test Δt(ts) == Millisecond(1)

    # Time
    # Time(h, [mi, s, ms, us, ns]) -> Time
    ts = [Time("00:00:00"), Time("01:00:00")]
    @test_throws AssertionError Δt(ts)

    ts = [Time(1,0,0,0,0,0),
          Time(3,0,0,0,0,0),
          Time(5,0,0,0,0,0)]
    @test Δt(ts) == Hour(2)

    ts = [Time(1,0,0,0,0,0),
          Time(2,0,0,0,0,0),
          Time(3,0,0,0,0,0)]
    @test Δt(ts) == Hour(1)

    ts = [Time(0,1,0,0,0,0),
          Time(0,2,0,0,0,0),
          Time(0,3,0,0,0,0)]
    @test Δt(ts) == Minute(1)

    ts = [Time(0,0,1,0,0,0),
          Time(0,0,2,0,0,0),
          Time(0,0,3,0,0,0)]
    @test Δt(ts) == Second(1)

    ts = [Time(0,0,0,1,0,0),
          Time(0,0,0,2,0,0),
          Time(0,0,0,3,0,0)]
    @test Δt(ts) == Millisecond(1)

    ts = [Time(0,0,0,0,1,0),
          Time(0,0,0,0,2,0),
          Time(0,0,0,0,3,0)]
    @test Δt(ts) == Microsecond(1)

    ts = [Time(0,0,0,0,0,0),
          Time(0,0,0,0,0,1),
          Time(0,0,0,0,0,2)]
    @test Δt(ts) == Nanosecond(1)

    # Generic Δt
    ts = [1,2,3]
    @test Δt(ts) == 1

    ts = [1.,3,5]
    @test Δt(ts) == 2.0

    # nΔt
    ts = [Date("2000-01-01"), Date("2000-01-02"), Date("2000-01-03")]
    @test nΔt(ts,  2) == reshape([Date("2000-01-04"),  Date("2000-01-05")],:,1)
    @test nΔt(ts, -2) == reshape([Date("1999-12-30")  Date("1999-12-31")],:,1)
    @test nΔt(ts,  0) == reshape(ts,:,1)

    ts = [DateTime("2000-01-01"), DateTime("2000-01-02"), DateTime("2000-01-03")]
    @test nΔt(ts,  2) == reshape([DateTime("2000-01-04"); DateTime("2000-01-05")],:,1)
    @test nΔt(ts, -2) == reshape([DateTime("1999-12-30"); DateTime("1999-12-31")],:,1)
    @test nΔt(ts,  0) == reshape(ts,:,1)

    ts = [Time("00:00:03"), Time("00:00:04"), Time("00:00:05")]
    @test nΔt(ts,  2) == reshape([Time(0, 0, 6), Time(0, 0, 7)],:,1)
    @test nΔt(ts, -2) == reshape([Time(0, 0, 1), Time(0, 0, 2)],:,1)
    @test nΔt(ts,  0) == reshape(ts,:,1)

    ts = [0,1,2]
    @test nΔt(ts,  2) == reshape([3,4],:,1)
    @test nΔt(ts, -2) == reshape([-2,-1],:,1)
    @test nΔt(ts,  0) == reshape(ts,:,1)

end


text:
using RandomDotOrg
using Test

@testset "RandomDotOrg.jl" begin
    # Write your tests here.
end


text:
"""
    ExplicitEmbeddedBackend{TF<:NamedTuple} <: AbstractDiffBackend

A backend to use with the [`RiemannianProjectionBackend`](@ref) or the [`TangentDiffBackend`](@ref),
when you have explicit formulae for the gradient in the embedding available.

# Constructor
    ExplicitEmbeddedBackend(M::AbstractManifold; kwargs)

Construct an [`ExplicitEmbeddedBackend`](@ref) in the embedding `M`,
where currently the following keywords may be used

* `gradient` for a(n allocating) gradient function `gradient(M, p)` defined in the embedding
* `gradient!` for a mutating gradient function `gradient!(M, X, p)`.

Note that the gradient functions are defined on the embedding manifold `M` passed to the Backend as well
"""
struct ExplicitEmbeddedBackend{TM<:AbstractManifold,TF<:NamedTuple} <: AbstractDiffBackend
    manifold::TM
    functions::TF
end
function ExplicitEmbeddedBackend(M::TM; kwargs...) where {TM<:AbstractManifold}
    return ExplicitEmbeddedBackend{TM,typeof(values(kwargs))}(M, values(kwargs))
end

function _gradient(f, p, e::ExplicitEmbeddedBackend)
    g = get(e.functions, :gradient, Missing())
    g === missing &&
        throw(MissingException("The provided Embedded backend does not provide a gradient"))
    return g(e.manifold, p)
end

function _gradient!(f, X, p, e::ExplicitEmbeddedBackend)
    g! = get(e.functions, :gradient!, Missing())
    g! === missing && throw(
        MissingException(
            "The provided Embedded backend does not provide a mutating gradient",
        ),
    )
    g!(e.manifold, X, p)
    return X
end


text:
# Test if tf listner works correctly
using RobotOS
@rosimport std_msgs.msg: Header, Bool
@rosimport geometry_msgs.msg: Transform, TransformStamped, Vector3, Quaternion
rostypegen(@__MODULE__)
import .std_msgs.msg: Header, BoolMsg
import .geometry_msgs.msg: Transform, TransformStamped, Vector3, Quaternion

using ROSTransforms.tf2_ros
using Test

# Initialize node
init_node("listner", anonymous=true) # delete

# Setup the transform listener
buffer = Buffer()
sleep(10)

# Wait and collect the transform
transform = lookupTransform(buffer, "child_link", "parent_link", RobotOS.now(),
                            Duration(10))
display()
display(transform)
display()
@testset "Buffer" begin
    tf_correct = TransformStamped(Header(0, RobotOS.now(), "parent_link"),
                                  "child_link",
                                  Transform(Vector3(1., 2., 3.),
                                            Quaternion(0, 0, 0, 1.))
                                  )
    @test transform == tf_correct
end

# Publsih success
pub = Publisher{BoolMsg}("/success", queue_size=10)
for i ∈ 1:100
    publish(pub, BoolMsg(true))
    rossleep(.1)
end


text:
# FreeType DSL

module DSL

macro ok(block...)
    :(ret = $(esc(block...)); 0 == ret ? nothing : error($(esc(block))[1]))
end

export create_library, create_typeface, create_glyph, outline_decompose
export close_library, close_typeface
export pt, px

using FreeType

function FT_Outline_Decompose(outline::FreeType.FT_Outline,outline_funcs::FreeType.FT_Outline_Funcs,user)
    ccall((:FT_Outline_Decompose,FreeType.freetype),FreeType.FT_Error,(Ptr{FreeType.FT_Outline},Ptr{FreeType.FT_Outline_Funcs},Ptr{Void}),pointer_from_objref(outline),pointer_from_objref(outline_funcs),user)
end

struct Library
    ref::Ref{FT_Library}
end

struct TypeFace
    ref::Ref{FT_Face}
end

abstract type PathOperation end
struct Move <: PathOperation
    to
end

struct Line <: PathOperation
    to
end

struct Conic <: PathOperation
    control
    to
end

struct Cubic <: PathOperation
    control1
    control2
    to
end

function create_library()::Library
    reflibrary = Ref{FreeType.FT_Library}()
    @ok FT_Init_FreeType(reflibrary)
    Library(reflibrary)
end

function close_typeface(face::TypeFace)
    FT_Done_Face(face.ref[])
end

function close_library(library::Library)
    FT_Done_FreeType(library.ref[])
end

function create_typeface(library::Library, pathname::String)::TypeFace
    refface = Ref{FreeType.FT_Face}()
    @ok FT_New_Face(library.ref[], pathname, 0, refface)
    TypeFace(refface)
end

abstract type LengthUnit end
abstract type RelativeUnit <: LengthUnit end
abstract type AbsoluteUnit <: LengthUnit end

struct Point <: AbsoluteUnit
    val::Real
end

struct Pixel <: RelativeUnit
    val::Real
end

struct Glyph
    face::FreeType.FT_FaceRec
    slot::FreeType.FT_GlyphSlotRec
    char::Char
    unit::LengthUnit
end

const pt = Point
const px = Pixel

Base.:*(n::Real, unit::Type{T}) where T <: LengthUnit = T(n)

Base.promote_rule(::Type{Pixel}, ::Type{Point}) = Point
Base.convert(::Type{Pixel}, pt::Point) = Pixel(4/3 * pt.val)
Base.convert(::Type{Point}, px::Pixel) = Pixel(3/4 * px.val)

function create_glyph(typeface::TypeFace, char::Char, unit::T) where T <: LengthUnit
    glyph_index = FT_Get_Char_Index(typeface.ref[], char)
    if isa(unit, RelativeUnit)
        pixel_width = pixel_height = UInt32(unit.val)
        @ok FT_Set_Pixel_Sizes(typeface.ref[], pixel_width, pixel_height)
    elseif isa(unit, AbsoluteUnit)
        char_width = char_height = 64unit.val
        horz_resolution = vert_resolution = 72
        @ok FT_Set_Char_Size(typeface.ref[], char_width, char_height, horz_resolution, vert_resolution)
    end
    @ok FT_Load_Glyph(typeface.ref[], glyph_index, FT_LOAD_DEFAULT)
    face = unsafe_load(typeface.ref[])
    slot = unsafe_load(face.glyph)
    Glyph(face, slot, char, unit)
end

function outline_decompose(glyph::Glyph)
    function pos(p::Ptr{FreeType.FT_Vector})
        vec = unsafe_load(p)
        (vec.x, vec.y)
    end

    move_f = cfunction(Cint,(Ptr{FreeType.FT_Vector},Ptr{Void})) do to, user_f
        C_NULL != user_f && ccall(user_f, Void, (Any,), Move(pos.([to])...))
        Cint(0)
    end

    line_f = cfunction(Cint,(Ptr{FreeType.FT_Vector},Ptr{Void})) do to, user_f
        C_NULL != user_f && ccall(user_f, Void, (Any,), Line(pos.([to])...))
        Cint(0)
    end

    conic_f = cfunction(Cint,(Ptr{FreeType.FT_Vector},Ptr{FreeType.FT_Vector},Ptr{Void})) do control, to, user_f
        C_NULL != user_f && ccall(user_f, Void, (Any,), Conic(pos.([control, to])...))
        Cint(0)
    end

    cubic_f = cfunction(Cint,(Ptr{FreeType.FT_Vector},Ptr{FreeType.FT_Vector},Ptr{FreeType.FT_Vector},Ptr{Void})) do control1, control2, to, user_f
        C_NULL != user_f && ccall(user_f, Void, (Any,), Cubic(pos.([control1, control2, to])...))
        Cint(0)
    end

    global paths = Vector{PathOperation}()
    user_f = cfunction(Void, (PathOperation,)) do op
        push!(paths, op)
        nothing
    end

    outline_funcs = FreeType.FT_Outline_Funcs(move_f, line_f, conic_f, cubic_f, 0, 0)
    @ok FT_Outline_Decompose(glyph.slot.outline, outline_funcs, user_f)
    paths::Vector{PathOperation}
end

end


text:
module MLDatasets

using Requires
using FixedPointNumbers, ColorTypes

bytes_to_type(::Type{UInt8}, A::Array{UInt8}) = A
bytes_to_type(::Type{N0f8}, A::Array{UInt8}) = reinterpret(N0f8, A)
bytes_to_type(::Type{T}, A::Array{UInt8}) where T<:Integer = convert(Array{T}, A)
bytes_to_type(::Type{T}, A::Array{UInt8}) where T<:AbstractFloat = A ./ T(255)
bytes_to_type(::Type{T}, A::Array{UInt8}) where T<:Number  = convert(Array{T}, reinterpret(N0f8, A))

global __images_supported__ = false
global __matfiles_supported__ = false

function _channelview(image::AbstractArray{<:Colorant})
    __images_supported__ ||
        error("Converting to/from image requires `using ImageCore`")
    ImageCore.channelview(image)
end

function _colorview(::Type{T}, array::AbstractArray{<:Number}) where T <: Color
    __images_supported__ ||
        error("Converting to image requires `using ImageCore`")
    ImageCore.colorview(T, array)
end

function _matopen(f::Function, path::AbstractString)
    __matfiles_supported__ ||
        error("Reading .mat files requires `using MAT`")
    MAT.matopen(f, path)
end

function _matread(path::AbstractString)
    __matfiles_supported__ ||
        error("Reading .mat files requires `using MAT`")
    MAT.matread(path)
end

include("download.jl")
include("CoNLL.jl")

include("CIFAR10/CIFAR10.jl")
include("CIFAR100/CIFAR100.jl")
include("MNIST/MNIST.jl")
include("FashionMNIST/FashionMNIST.jl")
include("SVHN2/SVHN2.jl")
include("PTBLM/PTBLM.jl")
include("UD_English/UD_English.jl")

function __init__()
    # initialize optional dependencies
    @require ImageCore="a09fc81d-aa75-5fe9-8630-4744c3626534" begin
        global __images_supported__ = true
    end
    @require MAT="23992714-dd62-5051-b70f-ba57cb901cac" begin
        global __matfiles_supported__ = true
    end
end

end


text:
#= Julia (v1.0+) function to retrieve the collocates of a node word
Earl K. Brown, ekbrown byu edu (add appropriate characters to create email)
=#

using DataFrames, Languages

function find_collocates(dir_with_txt, node_wd, stop_words = [], span = 4, side = "both", min_freq = 2, sort_by = "freq")
    #=
    dir_with_txt: a string indicating the directory with the TXT files.
    node_wd: a string with the node word whose collocates are desired.
    stop_words: an array with words to exclude as possible collocates.
    span: an integer with the width in words of the span around the node word (default is 4).
    side: a string indicating which side, or both, of the node word to look for collocates,
        from among "both", "left", "right" (default is "both").
    min_freq: an integer giving the minimum frequency of the collocates (default is 2).
    sort_by: a string indicating the metric to order the collocates, from among
        "freq" (frequency), "t_score", "mi" (mutual information), "log_dice" (default is "freq").
    =#

    ### verify arguments given by user
    if !isa(node_wd, String)
       error("In the call to find_collocates, you need to supply a String to the argument 'node_wd'.")
    end
    node_wd = uppercase(node_wd)

    if span <= 0 || !isa(span, Integer)
       error("In the call to find_collocates(), you need to supply a positive integer to the argument 'span'.")
    end

    if lowercase(side) == "both"
        span_to_search = -span:span
    elseif lowercase(side) == "left"
        span_to_search = -span:-1
    elseif lowercase(side) == "right"
        span_to_search = 1:span
    else
        error("In the call to find_collocates(), you need to specify 'side' as either 'left', 'right', or 'both'.")
    end

    stop_words = uppercase.(stop_words)
    ### end data verification

    # create collector dictionaries
    freqs_collocates = Dict{String,Int64}()
    freqs_wds = Dict{String, Int64}()

    # get TXT filenames
    cd(dir_with_txt)
    filenames = filter(x -> occursin(r"\.txt$"i, x), readdir())

    # create collector for total words in all files
    count_wds = 0

    # loop over files
    for f in filenames

        # open connection to current file
        open(f) do fin

            # read whole file as string and make it uppercase
            whole_file = uppercase(read(fin, String))

            # split up current file into words
            wds = split(whole_file, r"[^-'a-z]"i, keepempty = false)

            # loop over words in current file
            for j in 1:length(wds)

                # increment frequency counter for current word and total word counter
                freqs_wds[wds[j]] = get(freqs_wds, wds[j], 0) + 1
                count_wds += 1

                # checks whether the current word is the node word
                if node_wd == wds[j]

                    # loop over the collocates within the span
                    for k in span_to_search

                        # if the current span word is the node word itself
                        if k == 0
                            continue
                        end

                        # try to get the next collocate word, if it doesn't fall outside the range of the words in the current file
                        try
                            collocate_wd = wds[j + k]

                            # if the collocate word is not a stopword, add to collocate collector
                            if !in(collocate_wd, stop_words)
                                freqs_collocates[collocate_wd] = get(freqs_collocates, collocate_wd, 0) + 1
                            end
                        catch BoundsError
                            continue
                        end  # try catch block
                    end  # next collocate word
                end  # if match is found in current file
            end  # next index over words
        end  # close connection to current file
    end  # next file in directory

    # push dictionary to data frame
    freqs_df = DataFrame(collocate = String[], freq = Int64[], t_score = Float64[], mi = Float64[], log_dice = Float64[])
    for (k, v) in freqs_collocates

        # add collocate word is above minimum frequency specified by the user
        if v >= min_freq

            # calculate word association metrics and push to collector DataFrame
            t_score = (v - ((freqs_wds[node_wd] * freqs_wds[k]) / count_wds)) / sqrt(v)
            mi = log2((v * count_wds) / (freqs_wds[node_wd] * freqs_wds[k]))
            log_dice = 14 + log2((2 * v) / (freqs_wds[node_wd] + freqs_wds[k]))
            push!(freqs_df, [k, v, t_score, mi, log_dice])
        end
    end

    # sort in descending order by association metric, then in ascending order by collocate
    if sort_by == "freq"
        sort!(freqs_df, [order(:freq, rev = true), order(:collocate)])
    elseif sort_by == "t_score"
        sort!(freqs_df, [order(:t_score, rev = true), order(:collocate)])
    elseif sort_by == "mi"
        sort!(freqs_df, [order(:mi, rev = true), order(:collocate)])
    elseif sort_by == "log_dice"
        sort!(freqs_df, [order(:log_dice, rev = true), order(:collocate)])
    end

    return freqs_df

end  # end function definition

### test the function
dir_with_txt = "/Users/ekb5/Corpora/gen_conf_trunc"
node_wd = "prophet"  # as a string
# stop_words = stopwords(Languages.English())
stop_words = []
span = 4
side = "both"
min_freq = 1
sort_by = "log_dice"

@time results = find_collocates(dir_with_txt, node_wd, stop_words, span, side, min_freq, sort_by)
println(results[1:10,:])


text:
"""
spin30.jl - unscented transform robustness for the δfq problem
"""

WDIR = joinpath(@__DIR__, "../../")
include(joinpath(WDIR, "src", "spin", "spin.jl"))

using Altro
using Distributions
using HDF5
using LinearAlgebra
using Random
using RobotDynamics
using StaticArrays
using TrajectoryOptimization
const RD = RobotDynamics
const TO = TrajectoryOptimization

# paths
const EXPERIMENT_META = "spin"
const EXPERIMENT_NAME = "spin30"
const SAVE_PATH = joinpath(WDIR, "out", EXPERIMENT_META, EXPERIMENT_NAME)

# problem
const CONTROL_COUNT = 1
const ACONTROL_SIZE = CONTROL_COUNT
const STATE_COUNT = 3
# const STATE_COUNT = 4
const ASTATE_SIZE_BASE = STATE_COUNT * HDIM_ISO + 3 * CONTROL_COUNT
const SAMPLES_PER_STATE = 10
const PENALTY_SIZE = 1
const CHUNK_SIZE = SAMPLES_PER_STATE * HDIM_ISO + PENALTY_SIZE
# state indices
const STATE1_IDX = SVector{HDIM_ISO}(1:HDIM_ISO)
const STATE2_IDX = SVector{HDIM_ISO}(STATE1_IDX[end] + 1:STATE1_IDX[end] + HDIM_ISO)
const STATE3_IDX = SVector{HDIM_ISO}(STATE2_IDX[end] + 1:STATE2_IDX[end] + HDIM_ISO)
# const STATE4_IDX = SVector{HDIM_ISO}(STATE3_IDX[end] + 1:STATE3_IDX[end] + HDIM_ISO)
# const INTCONTROLS_IDX = SVector{CONTROL_COUNT}(STATE4_IDX[end] + 1:STATE4_IDX[end] + CONTROL_COUNT)
const INTCONTROLS_IDX = SVector{CONTROL_COUNT}(STATE3_IDX[end] + 1:STATE3_IDX[end] + CONTROL_COUNT)
const CONTROLS_IDX = SVector{CONTROL_COUNT}(INTCONTROLS_IDX[end] + 1:
                                            INTCONTROLS_IDX[end] + CONTROL_COUNT)
const DCONTROLS_IDX = SVector{CONTROL_COUNT}(CONTROLS_IDX[end] + 1:
                                             CONTROLS_IDX[end] + CONTROL_COUNT)
# control indices
const D2CONTROLS_IDX = SVector{CONTROL_COUNT}(1:CONTROL_COUNT)
# sample indices
const S1_IDX = SVector{HDIM_ISO}(HDIM_ISO * 0 + 1:HDIM_ISO * 1)
const S2_IDX = SVector{HDIM_ISO}(HDIM_ISO * 1 + 1:HDIM_ISO * 2)
const S3_IDX = SVector{HDIM_ISO}(HDIM_ISO * 2 + 1:HDIM_ISO * 3)
const S4_IDX = SVector{HDIM_ISO}(HDIM_ISO * 3 + 1:HDIM_ISO * 4)
const S5_IDX = SVector{HDIM_ISO}(HDIM_ISO * 4 + 1:HDIM_ISO * 5)
const S6_IDX = SVector{HDIM_ISO}(HDIM_ISO * 5 + 1:HDIM_ISO * 6)
const S7_IDX = SVector{HDIM_ISO}(HDIM_ISO * 6 + 1:HDIM_ISO * 7)
const S8_IDX = SVector{HDIM_ISO}(HDIM_ISO * 7 + 1:HDIM_ISO * 8)
const S9_IDX = SVector{HDIM_ISO}(HDIM_ISO * 8 + 1:HDIM_ISO * 9)
const S10_IDX = SVector{HDIM_ISO}(HDIM_ISO * 9 + 1:HDIM_ISO * 10)
const STATE_IDX = SVector{HDIM_ISO}(1:HDIM_ISO)

# model
struct Model <: RD.AbstractModel
    S::Diagonal{Float64,SVector{HDIM_ISO,Float64}}
    nominal_idxs::Array{SVector{HDIM_ISO,Int},1}
    fq_cov::Float64
    alpha::Float64
    sample_state_count::Int
end
@inline RD.state_dim(model::Model) = (
    ASTATE_SIZE_BASE + model.sample_state_count * CHUNK_SIZE
)
@inline RD.control_dim(model::Model) = ACONTROL_SIZE
@inline astate_sample_inds(sample_state_index::Int, sample_index::Int) = (
    SVector{HDIM_ISO}((
        ASTATE_SIZE_BASE + (sample_state_index - 1) * CHUNK_SIZE
        + (sample_index - 1) * HDIM_ISO + 1
    ):(
        ASTATE_SIZE_BASE + (sample_state_index - 1) * CHUNK_SIZE
        + sample_index * HDIM_ISO
    ))
)
@inline sample_idxs(model::Model) = [
    astate_sample_inds(i, j)
    for i = 1:model.sample_state_count
    for j = 1:SAMPLES_PER_STATE
]

function unscented_transform(model::Model, astate::AbstractVector,
                             negi_hc::AbstractMatrix, h_prop::AbstractMatrix,
                             dt::Real, i::Int)
    # get states
    offset = ASTATE_SIZE_BASE + (i - 1) * CHUNK_SIZE
    s1 = astate[offset + S1_IDX]
    s2 = astate[offset + S2_IDX]
    s3 = astate[offset + S3_IDX]
    s4 = astate[offset + S4_IDX]
    s5 = astate[offset + S5_IDX]
    s6 = astate[offset + S6_IDX]
    s7 = astate[offset + S7_IDX]
    s8 = astate[offset + S8_IDX]
    s9 = astate[offset + S9_IDX]
    s10 = astate[offset + S10_IDX]
    # grab chol info
    fq_chol5 = sqrt(model.fq_cov)
    # propagate states
    s1 = h_prop * s1
    s2 = h_prop * s2
    s3 = h_prop * s3
    s4 = h_prop * s4
    s5 = exp(dt * ((FQ + fq_chol5) * NEGI_H0_ISO + negi_hc)) * s5
    s6 = h_prop * s6
    s7 = h_prop * s7
    s8 = h_prop * s8
    s9 = h_prop * s9
    s10 = exp(dt * ((FQ - fq_chol5) * NEGI_H0_ISO + negi_hc)) * s10
    # compute state mean
    sm = 1//SAMPLES_PER_STATE .* (
        s1 + s2 + s3 + s4 + s5
        + s6 + s7 + s8 + s9 + s10
    )
    # compute state covariance
    d1 = s1 - sm
    d2 = s2 - sm
    d3 = s3 - sm
    d4 = s4 - sm
    d5 = s5 - sm
    d6 = s6 - sm
    d7 = s7 - sm
    d8 = s8 - sm
    d9 = s9 - sm
    d10 = s10 - sm
    s_cov = 1 / (2 * model.alpha^2) .* (
        d1 * d1' + d2 * d2' + d3 * d3' + d4 * d4'
        + d5 * d5' + d6 * d6' + d7 * d7' + d8 * d8'
        + d9 * d9' + d10 * d10'
    )
    # perform cholesky decomposition on joint covariance
    cov = zeros(eltype(s_cov), HDIM_ISO + 1, HDIM_ISO + 1)
    cov[1:HDIM_ISO, 1:HDIM_ISO] .= s_cov
    cov[HDIM_ISO + 1, HDIM_ISO + 1] = model.fq_cov
    # TOOD: cholesky! requires writing zeros in upper triangle
    cov_chol = model.alpha * cholesky(Symmetric(cov)).L
    # resample states
    s_chol1 = cov_chol[STATE_IDX, 1]
    s_chol2 = cov_chol[STATE_IDX, 2]
    s_chol3 = cov_chol[STATE_IDX, 3]
    s_chol4 = cov_chol[STATE_IDX, 4]
    s1 = sm + s_chol1
    s2 = sm + s_chol2
    s3 = sm + s_chol3
    s4 = sm + s_chol4
    s5 = sm
    s6 = sm - s_chol1
    s7 = sm - s_chol2
    s8 = sm - s_chol3
    s9 = sm - s_chol4
    s10 = sm
    # normalize
    s1 = s1 ./sqrt(s1's1)
    s2 = s2 ./sqrt(s2's2)
    s3 = s3 ./sqrt(s3's3)
    s4 = s4 ./sqrt(s4's4)
    s5 = s5 ./sqrt(s5's5)
    s6 = s6 ./sqrt(s6's6)
    s7 = s7 ./sqrt(s7's7)
    s8 = s8 ./sqrt(s8's8)
    s9 = s9 ./sqrt(s9's9)
    s10 = s10 ./sqrt(s10's10)
    # compute penalty
    dnom = sm - astate[model.nominal_idxs[i]]
    penalty = SVector{1}(tr(s_cov * model.S) + dnom' * model.S * dnom)

    samples = [s1; s2; s3; s4; s5; s6; s7; s8; s9; s10; penalty]

    return samples
end

# dynamics
function RD.discrete_dynamics(::Type{RK3}, model::Model, astate::StaticVector,
                              acontrol::StaticVector, time::Real, dt::Real) where {n}
    # base dynamics
    negi_hc = astate[CONTROLS_IDX[1]] * NEGI_H1_ISO
    h_prop = exp(dt * (FQ_NEGI_H0_ISO + negi_hc))
    state1 = h_prop * astate[STATE1_IDX]
    state2 = h_prop * astate[STATE2_IDX]
    state3 = h_prop * astate[STATE3_IDX]
    # state4 = h_prop * astate[STATE4_IDX]
    intcontrols = astate[INTCONTROLS_IDX] + dt * astate[CONTROLS_IDX]
    controls = astate[CONTROLS_IDX] + dt * astate[DCONTROLS_IDX]
    dcontrols = astate[DCONTROLS_IDX] + dt * acontrol[D2CONTROLS_IDX]

    astate_ = [
        # state1; state2; state3; state4; intcontrols; controls; dcontrols;
        state1; state2; state3; intcontrols; controls; dcontrols;
    ]

    # unscented transform
    for i = 1:model.sample_state_count
        sample_states = unscented_transform(model, astate, negi_hc, h_prop, dt, i)
        astate_ = [astate_; sample_states]
    end

    return astate_
end

# main
function run_traj(;gate_type=xpiby2, evolution_time=60., solver_type=altro,
                  sqrtbp=false, integrator_type=rk3, qs=[1e0, 1e0, 1e0, 1e-1, 1e0, 1e-1],
                  dt_inv=Int64(1e1), smoke_test=false, constraint_tol=1e-8, al_tol=1e-4,
                  pn_steps=2, max_penalty=1e11, verbose=true, save=true,
                  fq_cov=FQ * 1e-2, max_iterations=Int64(2e5), gradient_tol_int=1,
                  dJ_counter_limit=Int(1e2), state_cov=1e-2, seed=0, alpha=1.,
                  sample_states=[IS1_ISO_], nominal_idxs=[STATE1_IDX], static=true,
                  benchmark=false)
    Random.seed!(seed)
    (sample_state_count,) = size(sample_states)
    S = Diagonal(SVector{HDIM_ISO}(fill(qs[5], HDIM_ISO)))
    model = Model(S, nominal_idxs, fq_cov, alpha, sample_state_count)
    n = RD.state_dim(model)
    m = RD.control_dim(model)
    t0 = 0.

    # initial state, target state
    x0 = zeros(n)
    xf = zeros(n)
    gate = GT_GATE_ISO[gate_type]
    x0[STATE1_IDX] = IS1_ISO_
    x0[STATE2_IDX] = IS2_ISO_
    x0[STATE3_IDX] = IS3_ISO_
    # x0[STATE4_IDX] = IS4_ISO_
    xf[STATE1_IDX] = gate * IS1_ISO_
    xf[STATE2_IDX] = gate * IS2_ISO_
    state_dist = Distributions.Normal(0., state_cov)
    for i = 1:sample_state_count
        sample_state = sample_states[i]
        for j = 1:SAMPLES_PER_STATE
            sample_idx = astate_sample_inds(i, j)
            sample = sample_state .+ rand(state_dist, HDIM_ISO)
            x0[sample_idx] = sample ./ sqrt(sample'sample)
        end
    end
    if static
        x0 = SVector{n}(x0)
        xf = SVector{n}(xf)
    end

    # control amplitude constraint
    x_max = fill(Inf, n)
    x_max[CONTROLS_IDX] .= MAX_CONTROL_NORM_0
    x_min = fill(-Inf, n)
    x_min[CONTROLS_IDX] .= -MAX_CONTROL_NORM_0
    if static
        x_max = SVector{n}(x_max)
        x_min = SVector{n}(x_min)
    end

    # control amplitude constraint at boundary
    x_max_boundary = fill(Inf, n)
    x_max_boundary[CONTROLS_IDX] .= 0
    x_min_boundary = fill(-Inf, n)
    x_min_boundary[CONTROLS_IDX] .= 0
    if static
        x_max_boundary = SVector{n}(x_max_boundary)
        x_min_boundary = SVector{n}(x_min_boundary)
    end

    # initial trajectory
    dt = dt_inv^(-1)
    N = Int(floor(evolution_time * dt_inv)) + 1
    if static
        U0 = [SVector{m}(fill(1e-4, CONTROL_COUNT)) for k = 1:N-1]
        X0 = [SVector{n}(fill(NaN, n)) for k = 1:N]
    else
        U0 = [fill(1e-4, CONTROL_COUNT) for k = 1:N-1]
        X0 = [fill(NaN, n) for k = 1:N]
    end
    Z = Traj(X0, U0, dt * ones(N))

    # cost function
    Q = zeros(n)
    Q[STATE1_IDX] = Q[STATE2_IDX] = fill(qs[1], HDIM_ISO)
    Q[INTCONTROLS_IDX] = fill(qs[2], CONTROL_COUNT)
    Q[CONTROLS_IDX] = fill(qs[3], CONTROL_COUNT)
    Q[DCONTROLS_IDX] = fill(qs[4], CONTROL_COUNT)
    # penalty has unit value, its cost can be increased through qs[5],
    # which modifies the model.S matrix
    for i = 1:sample_state_count
        Q[ASTATE_SIZE_BASE + i * CHUNK_SIZE] = 1
    end
    R = zeros(m)
    R[D2CONTROLS_IDX] = fill(qs[6], CONTROL_COUNT)
    if static
        Q = SVector{n}(Q)
        R = SVector{m}(R)
    end
    Q = Diagonal(Q)
    Qf = Q * N
    R = Diagonal(R)
    objective = LQRObjective(Q, R, Qf, xf, N)

    # constraints
    # must satisfy control amplitude bound
    control_bnd = BoundConstraint(n, m, x_max=x_max, x_min=x_min)
    # must statisfy conrols start and end at 0
    control_bnd_boundary = BoundConstraint(n, m, x_max=x_max_boundary, x_min=x_min_boundary)
    # must reach target state, must have zero net flux
    target_astate_constraint = GoalConstraint(xf, [STATE1_IDX; STATE2_IDX; INTCONTROLS_IDX])
    # must obey unit norm
    norm_idxs = sample_idxs(model)
    push!(norm_idxs, STATE1_IDX)
    push!(norm_idxs, STATE2_IDX)
    push!(norm_idxs, STATE3_IDX)
    # push!(norm_idxs, STATE4_IDX)
    norm_constraints = [NormConstraint(n, m, 1, TO.Equality(), idx) for idx in norm_idxs]
    constraints = ConstraintList(n, m, N)
    add_constraint!(constraints, control_bnd, 2:N-2)
    add_constraint!(constraints, control_bnd_boundary, N-1:N-1)
    add_constraint!(constraints, target_astate_constraint, N:N)
    for norm_constraint in norm_constraints
        add_constraint!(constraints, norm_constraint, 2:N-1)
    end

    # solve problem
    prob = Problem{IT_RDI[integrator_type]}(model, objective, constraints,
                                            x0, xf, Z, N, t0, evolution_time)
    solver = ALTROSolver(prob)
    verbose_pn = verbose ? true : false
    verbose_ = verbose ? 2 : 0
    projected_newton = solver_type == altro ? true : false
    constraint_tolerance = solver_type == altro ? constraint_tol : al_tol
    iterations_inner = smoke_test ? 1 : 300
    iterations_outer = smoke_test ? 1 : 30
    n_steps = smoke_test ? 1 : pn_steps
    static_bp = static ? true : false
    set_options!(
        solver, square_root=sqrtbp, constraint_tolerance=constraint_tolerance,
        projected_newton_tolerance=al_tol, n_steps=n_steps,
        penalty_max=max_penalty, verbose_pn=verbose_pn, verbose=verbose_,
        projected_newton=projected_newton, iterations_inner=iterations_inner,
        iterations_outer=iterations_outer, iterations=max_iterations,
        gradient_tolerance_intermediate=gradient_tol_int,
        dJ_counter_limit=dJ_counter_limit, static_bp=static_bp
    )
    if benchmark
        benchmark_result = Altro.benchmark_solve!(solver)
    else
        benchmark_result = nothing
        Altro.solve!(solver)
    end

    # Post-process.
    acontrols_raw = TO.controls(solver)
    acontrols_arr = permutedims(reduce(hcat, map(Array, acontrols_raw)), [2, 1])
    astates_raw = TO.states(solver)
    astates_arr = permutedims(reduce(hcat, map(Array, astates_raw)), [2, 1])
    cidx_arr = Array(CONTROLS_IDX)
    d2cidx_arr = Array(D2CONTROLS_IDX)
    cmax = TO.max_violation(solver)
    cmax_info = TO.findmax_violation(TO.get_constraints(solver))
    iterations_ = iterations(solver)

    result = Dict(
        "acontrols" => acontrols_arr,
        "controls_idx" => cidx_arr,
        "d2controls_dt2_idx" => d2cidx_arr,
        "evolution_time" => evolution_time,
        "astates" => astates_arr,
        "qs" => qs,
        "cmax" => cmax,
        "cmax_info" => cmax_info,
        "dt" => dt,
        "solver_type" => Integer(solver_type),
        "sqrtbp" => Integer(sqrtbp),
        "max_penalty" => max_penalty,
        "constraint_tol" => constraint_tol,
        "al_tol" => al_tol,
        "gradient_tolerance_intermediate" => gradient_tol_int,
        "dJ_counter_limit" => dJ_counter_limit,
        "integrator_type" => Integer(integrator_type),
        "gate_type" => Integer(gate_type),
        "save_type" => Integer(jl),
        "iterations" => iterations_,
        "seed" => seed,
        "fq_cov" => fq_cov,
        "max_iterations" => max_iterations,
    )

    # save
    if save
        save_file_path = generate_file_path("h5", EXPERIMENT_NAME, SAVE_PATH)
        println("Saving this optimization to $(save_file_path)")
        h5open(save_file_path, "cw") do save_file
            for key in keys(result)
                write(save_file, key, result[key])
            end
        end
        result["save_file_path"] = save_file_path
    end

    result = benchmark ? benchmark_result : result

    return result
end

max_stars_repo_name: SchusterLab/rbqoc

max_issues_repo_name: SchusterLab/rbqoc

max_forks_repo_name: SchusterLab/rbqoc


text:
println("Db) Test of the PhotoExcitation module with ASF from an internally generated initial- and final-state multiplet.")
#
setDefaults("print summary: open", "zzz-PhotoExcitation.sum")
wa = Atomic.Computation(Atomic.Computation(), name="xx", grid=JAC.Radial.Grid(true), nuclearModel=Nuclear.Model(3.),
                        initialConfigs=[Configuration("1s^2 2s"), Configuration("1s^2 2p")],
                        finalConfigs  =[Configuration("1s^2 2s"), Configuration("1s^2 2p")],
                        processSettings=PhotoExcitation.Settings([E1, M1], [UseCoulomb, UseBabushkin], true, true, true, true,
                                                                 LineSelection(), 0., 0., 1.0e6, ExpStokes(0., 0., 0.) ) )

wb = perform(wa)
setDefaults("print summary: close", "")

#     [Configuration("[Ne] 3s^2 3p^5"), Configuration("[Ne] 3s 3p^6"), Configuration("[Ne] 3s^2 3p^4 3d")],


text:
using Distributed
if haskey(ENV, "BENCHMARK_PROCS")
    const np, nt = parse.(Ref(Int), split(ENV["BENCHMARK_PROCS"], ":"))
    for i in workers()
        addprocs(np; exeflags="-t $nt")
    end
else
    const np = 2
    const nt = 1
end
if haskey(ENV, "BENCHMARK_REMOTES")
    const remotes = split(ENV["BENCHMARK_REMOTES"], ":")
    if !isempty(remotes)
        for i in 1:np
            addprocs(remotes; exeflags="-t $nt")
        end
    end
end
@everywhere using Dagger
import Dagger: Computation, reduceblock
using Dates, Random, Statistics, LinearAlgebra, InteractiveUtils

const output_format = get(ENV, "BENCHMARK_OUTPUT_FORMAT", "jls")
if output_format == "jld"
    using JLD
elseif output_format == "jls"
    using Serialization
else
    error("Unknown output format: $output_format")
end

const render = get(ENV, "BENCHMARK_RENDER", "")
if render == "live"
    const live = true
    using Luxor, ProfileSVG
    using Mux
elseif render == "webdash"
    const live = true
    using DaggerWebDash
    import DaggerWebDash: LinePlot, GanttPlot, GraphPlot, ProfileViewer
    using DataFrames
elseif render == "offline"
    const live = false
    using Luxor, ProfileSVG
    using FFMPEG, FileIO, ImageMagick
end
const RENDERS = Dict{Int,Dict}()
const live_port = parse(Int, get(ENV, "BENCHMARK_LIVE_PORT", "8000"))

const graph = parse(Bool, get(ENV, "BENCHMARK_GRAPH", "0"))
if graph && render == "webdash"
    @warn "BENCHMARK_GRAPH=1 is not compatible with BENCHMARK_RENDER=webdash; disabling graphing"
end
const profile = parse(Bool, get(ENV, "BENCHMARK_PROFILE", "0"))

_benches = get(ENV, "BENCHMARK", "cpu,cpu+dagger")
const benches = []
for bench in split(_benches, ',')
    if endswith(bench, "+dagger")
        accel = split(bench, '+')[1]
        dagger = true
    else
        accel = bench
        dagger = false
    end
    push!(benches, (name=bench, accel=accel, dagger=dagger))
end

if any(x->x.accel == "cuda", benches)
    @everywhere using DaggerGPU, CUDA
elseif any(x->x.accel == "amdgpu", benches)
    @everywhere using DaggerGPU, AMDGPU
end

const scales = eval(Meta.parse(get(ENV, "BENCHMARK_SCALE", "1:5:50")))

using BenchmarkTools

AmulB_compatible(x,y) =
    try; size(x), size(y), x * y; true; catch err false end

AmulB_compatible(x::Computation,y::Computation) = AmulB_compatible(domain(x), domain(y))

function AcmulB_compatible(x,y)
    AmulB_compatible(domain(x)', domain(y))
end

function array_suite_inner(ctx, X, Y, f=x->x)
    suite = BenchmarkGroup(["array"])

    suite["alloc"] = @benchmarkable (compute($f($X)); GC.gc())

    Y = compute(ctx, f(X))
    suite["X.+1.0"]  = @benchmarkable (compute($ctx, $f($Y.+1)); GC.gc())
    suite["X+X"]     = @benchmarkable (compute($ctx, $f($Y+$Y)); GC.gc())
    suite["X'"]      = @benchmarkable (compute($ctx, $f($Y')); GC.gc())
    suite["sin.(X)"] = @benchmarkable (compute($ctx, $f(sin.($Y))); GC.gc())

    if ndims(Y) == 2
        suite["X+X'"] = @benchmarkable (compute($ctx, $f($Y+$Y')); GC.gc())
        suite["X*X"]  = @benchmarkable (compute($ctx, $f($Y*$Y)); GC.gc())
        suite["X'*X"] = @benchmarkable (compute($ctx, $f($Y'*$Y)); GC.gc())
    end

    suite
end

function nnmf(X, W, H)
    # H update
    H = (H .* (W' * (X ./ (W * H)))
         ./ (sum(W; dims=1))')
    # W update
    W = (W .* ((X ./ (W * H)) * (H'))
         ./ (sum(H; dims=2)'))
    # error estimate
    X - W * H
end

theory_flops(nrow, ncol, nfeatures) = 11 * ncol * nrow * nfeatures + 2 * (ncol + nrow) * nfeatures

function nmf_suite(ctx; dagger, accel)
    suite = BenchmarkGroup()

    #= TODO: Re-enable
    pairs = Set{Tuple}()
    for i in (1, 2, 5, 10, 64, 100, 1028)
        for bscale in (1, 0.5, 0.3, 0.25, 0.1)
            b = ceil(Int, i*bscale)
            if !((b,i) in pairs)
                push!(pairs, (b,i))
                X = zeros(Blocks(b), i)
                Y = compute(ctx, f(X))
                suite["vector $i/$b"] = array_suite_inner(ctx, X, Y, f)
            end
            if !(((b,b),(i,i)) in pairs)
                push!(pairs, ((b,b),(i,i)))
                X = zeros(Blocks(b,b), i, i)
                Y = compute(ctx, f(X))
                suite["matrix ($i,$i)/$b"] = array_suite_inner(ctx, X, Y, f)
            end
        end
    end
    =#

    X = Ref{Any}()
    W = Ref{Any}()
    H = Ref{Any}()

    for scale in scales
        ncol = 2001 * scale
        nrow = 1002
        nfeatures = 12

        if !dagger
            suite["NNMF scaled by: $scale"] = @benchmarkable begin
                nnmf($X[], $W[], $H[])
            end setup=begin
                _scale = $scale
                @info "Starting non-Dagger NNMF (scale by $_scale)"
                if $accel == "cuda"
                    $X[] = CUDA.rand(Float32, $nrow, $ncol)
                    $W[] = CUDA.rand(Float32, $nrow, $nfeatures)
                    $H[] = CUDA.rand(Float32, $nfeatures, $ncol)
                elseif $accel == "amdgpu"
                    $X[] = ROCArray(rand(Float32, $nrow, $ncol))
                    $W[] = ROCArray(rand(Float32, $nrow, $nfeatures))
                    $H[] = ROCArray(rand(Float32, $nfeatures, $ncol))
                elseif $accel == "cpu"
                    $X[] = rand(Float32, $nrow, $ncol)
                    $W[] = rand(Float32, $nrow, $nfeatures)
                    $H[] = rand(Float32, $nfeatures, $ncol)
                end
            end teardown=begin
                $X[] = nothing
                $W[] = nothing
                $H[] = nothing
                @everywhere GC.gc()
            end
        else
            RENDERS[scale] = Dict{Int,Vector}()
            nw = length(workers())
            nsuite = BenchmarkGroup()
            while nw > 0
                opts = if accel == "cuda"
                    Dagger.Sch.SchedulerOptions(;proctypes=[
                        DaggerGPU.CuArrayDeviceProc
                    ])
                elseif accel == "amdgpu"
                    Dagger.Sch.SchedulerOptions(;proctypes=[
                        DaggerGPU.ROCArrayProc
                    ])
                elseif accel == "cpu"
                    Dagger.Sch.SchedulerOptions()
                else
                    error("Unknown accelerator $accel")
                end
                p = sum([length(Dagger.get_processors(OSProc(id))) for id in 2:(nw+1)])
                #bsz = ncol ÷ length(workers())
                bsz = ncol ÷ 64
                nsuite["Workers: $nw"] = @benchmarkable begin
                    _ctx = Context($ctx, workers()[1:$nw])
                    compute(_ctx, nnmf($X[], $W[], $H[]); options=$opts)
                end setup=begin
                    _nw, _scale = $nw, $scale
                    @info "Starting $_nw worker Dagger NNMF (scale by $_scale)"
                    if $accel == "cuda"
                        # FIXME: Allocate with CUDA.rand if possible
                        $X[] = Dagger.mapchunks(CUDA.cu, compute(rand(Blocks($bsz, $bsz), Float32, $nrow, $ncol); options=$opts))
                        $W[] = Dagger.mapchunks(CUDA.cu, compute(rand(Blocks($bsz, $bsz), Float32, $nrow, $nfeatures); options=$opts))
                        $H[] = Dagger.mapchunks(CUDA.cu, compute(rand(Blocks($bsz, $bsz), Float32, $nfeatures, $ncol); options=$opts))
                    elseif $accel == "amdgpu"
                        $X[] = Dagger.mapchunks(ROCArray, compute(rand(Blocks($bsz, $bsz), Float32, $nrow, $ncol); options=$opts))
                        $W[] = Dagger.mapchunks(ROCArray, compute(rand(Blocks($bsz, $bsz), Float32, $nrow, $nfeatures); options=$opts))
                        $H[] = Dagger.mapchunks(ROCArray, compute(rand(Blocks($bsz, $bsz), Float32, $nfeatures, $ncol); options=$opts))
                    elseif $accel == "cpu"
                        $X[] = compute(rand(Blocks($bsz, $bsz), Float32, $nrow, $ncol); options=$opts)
                        $W[] = compute(rand(Blocks($bsz, $bsz), Float32, $nrow, $nfeatures); options=$opts)
                        $H[] = compute(rand(Blocks($bsz, $bsz), Float32, $nfeatures, $ncol); options=$opts)
                    end
                end teardown=begin
                    if render != "" && !live
                        Dagger.continue_rendering[] = false
                        for i in 1:5
                            isready(Dagger.render_results) && break
                            sleep(1)
                        end
                        if isready(Dagger.render_results)
                            video_paths = take!(Dagger.render_results)
                            try
                                video_data = Dict(key=>read(video_paths[key]) for key in keys(video_paths))
                                push!(get!(()->[], RENDERS[$scale], $nw), video_data)
                            catch err
                                @error "Failed to process render results" exception=(err,catch_backtrace())
                            end
                        else
                            @warn "Failed to fetch render results"
                        end
                    end
                    $X[] = nothing
                    $W[] = nothing
                    $H[] = nothing
                    @everywhere GC.gc()
                end
                nw ÷= 2
            end
            suite["NNMF scaled by: $scale"] = nsuite
        end
    end

    suite
end

function main()
    nw = length(workers())
    output_prefix = "result-$(np)workers-$(nt)threads-$(Dates.now())"

    suites = Dict()
    opts = (;profile=profile)
    if render == "live"
        opts = merge(opts, (;log_sink=Dagger.LocalEventLog()))
        if graph
            opts = merge(opts, (;log_file=output_prefix*".dot"))
        end
    elseif render == "webdash"
        ml = Dagger.MultiEventLog()
        ml[:core] = Dagger.Events.CoreMetrics()
        ml[:id] = Dagger.Events.IDMetrics()
        ml[:timeline] = Dagger.Events.TimelineMetrics()
        profile && (ml[:profile] = DaggerWebDash.ProfileMetrics())
        ml[:wsat] = Dagger.Events.WorkerSaturation()
        ml[:loadavg] = Dagger.Events.CPULoadAverages()
        ml[:bytes] = Dagger.Events.BytesAllocd()
        ml[:mem] = Dagger.Events.MemoryFree()
        ml[:esat] = Dagger.Events.EventSaturation()
        ml[:psat] = Dagger.Events.ProcessorSaturation()
        lw = Dagger.Events.LogWindow(5*10^9, :core)
        df = DataFrame([key=>[] for key in keys(ml.consumers)]...)
        ts = Dagger.Events.TableStorage(df)
        push!(lw.creation_handlers, ts)
        d3r = DaggerWebDash.D3Renderer(live_port; seek_store=ts)
        push!(lw.creation_handlers, d3r)
        push!(lw.deletion_handlers, d3r)
        push!(d3r, GanttPlot(:core, :id, :timeline, :esat, :psat, "Overview"))
        # TODO: push!(d3r, ProfileViewer(:core, :profile, "Profile Viewer"))
        push!(d3r, LinePlot(:core, :wsat, "Worker Saturation", "Running Tasks"))
        push!(d3r, LinePlot(:core, :loadavg, "CPU Load Average", "Average Running Threads"))
        push!(d3r, LinePlot(:core, :bytes, "Allocated Bytes", "Bytes"))
        push!(d3r, LinePlot(:core, :mem, "Available Memory", "% Free"))
        #push!(d3r, GraphPlot(:core, :id, :timeline, :profile, "DAG"))
        ml.aggregators[:logwindow] = lw
        ml.aggregators[:d3r] = d3r
        opts = merge(opts, (;log_sink=ml))
    end
    ctx = Context(collect((1:nw) .+ 1); opts...)
    for bench in benches
        name = bench.name
        println("creating $name benchmarks")
        suites[name] = nmf_suite(ctx; dagger=bench.dagger, accel=bench.accel)
    end
    if render == "live" || render == "offline"
        Dagger.show_gantt(ctx; width=1800, window_length=5, delay=2, port=live_port, live=live)
        if live
            # Make sure server code is compiled
            sleep(1)
            run(pipeline(`curl -s localhost:$live_port/`; stdout=devnull))
            run(pipeline(`curl -s localhost:$live_port/profile`; stdout=devnull))
            @info "Rendering started on port $live_port"
        end
    elseif render == "webdash"
        # Kick the webserver into gear
        collect(ctx, delayed(identity)(1))
        run(pipeline(`curl -s localhost:$live_port/index.html`; stdout=devnull))
        @info "Rendering started on port $live_port"
    end
    res = Dict()
    for bench in benches
        name = bench.name
        println("running $name benchmarks")
        res[name] = try
            run(suites[name]; samples=3, seconds=10*60, gcsample=true)
        catch err
            @error "Error running $name benchmarks" exception=(err,catch_backtrace())
            nothing
        end
    end
    for bench in benches
        println("benchmark results for $(bench.name): $(minimum(res[bench.name]))")
    end

    println("saving results in $output_prefix.$output_format")
    if output_format == "jld"
        JLD.save(output, "results", res, "peakflops", peakflops(), "renders", RENDERS)
    elseif output_format == "jls"
        outdict = Dict("results"=>res, "peakflops"=>peakflops(), "renders"=>RENDERS)
        open(output_prefix*".jls", "w") do io
            serialize(io, outdict)
        end
    end

    if parse(Bool, get(ENV, "BENCHMARK_VISUALIZE", "0"))
        run(`$(Base.julia_cmd()) $(joinpath(pwd(), "visualize.jl")) -- $(output_prefix*"."*output_format)`)
    end

    println("Done.")

    # TODO: Compare with multiple results
    if length(ARGS) == 1
        compare_file = ARGS[1]
        println("Comparing with results in $compare_file")
        reference_results = JLD.load(compare_file, "result")
        @show judge(mean(reference_results), mean(res))
    end
end

main()


text:
module AdaptiveDesignOptimization
    using Distributions, Parameters, StatsFuns
    export Optimizer, Model, Optimize, Randomize
    export Dynamic, Static
    export update!, mean_post, std_post, get_best_design

    include("structs.jl")
    include("functions.jl")
end


text:
using NewtonsMethod
using Test

@testset "NewtonsMethod.jl" begin
    # Write your tests here.
    f_1(x) = x^2 - 3x + 2
    f_2(x) = sin(x)
    f_3(x) = log(x)
    f_4(x) = 1/x - 4
    f_5(x) = x+3
    f_6(x) = 3
    f_7(x) = 0
    f_8(x) = exp(x) - 5
    f_prime_1(x) = 2x - 3
    f_prime_2(x) = cos(x)
    f_prime_3(x) = 1/x
    f_prime_4(x) = -1/(x^2)
    f_prime_5(x) = 1
    f_prime_6(x) = 0
    f_prime_7(x) = 0
    f_prime_8(x) = exp(x)
    @test newtonroot(f_1, f_prime_1; x_0 = 0) ≈ 1
    @test newtonroot(f_2, f_prime_2; x_0 = 3) ≈ π
    @test newtonroot(f_3, f_prime_3; x_0 = 0.7) ≈ 1
    @test newtonroot(f_4, f_prime_4; x_0 = 0.4) ≈ 0.25
    @test newtonroot(f_5, f_prime_5; x_0 = 0) ≈ -3
    @test newtonroot(f_6, f_prime_6; x_0 = 1) === nothing
    @test newtonroot(f_7, f_prime_7; x_0 = 4) == 4
    @test newtonroot(f_8, f_prime_8; x_0 = 3) ≈ log(5)
    @test newtonroot(f_1; x_0 = 0) ≈ 1
    @test newtonroot(f_2; x_0 = 3) ≈ π
    @test newtonroot(f_3; x_0 = 0.7) ≈ 1
    @test newtonroot(f_4; x_0 = 0.4) ≈ 0.25
    @test newtonroot(f_5; x_0 = 0) ≈ -3
    @test newtonroot(f_6; x_0 = 1) === nothing
    @test newtonroot(f_7; x_0 = 4) == 4
    @test newtonroot(f_8; x_0 = 3) ≈ log(5)
end

@testset "non convergence" begin
    f(x) = x^2+2
    f_prime(x) = 2x
    @test newtonroot(f, f_prime; x_0 = 0) === nothing
    @test newtonroot(f, f_prime; x_0 = 0.5) === nothing
    @test newtonroot(f;x_0 = 0) === nothing
    @test newtonroot(f;x_0 = 0.5) === nothing
end

@testset "Big Float" begin
    f(x) = x^2 - 3x + 2
    f_prime(x) = 2x - 3
    x_0 = parse(BigFloat, "0.34")
    @test convert(Float64, newtonroot(f, f_prime; x_0 = x_0)) ≈ 1
    @test convert(Float64, newtonroot(f; x_0 = x_0)) ≈ 1
end

@testset "max iter" begin
    f(x) = exp(x) - 5
    f_prime(x) = exp(x)
    @test newtonroot(f, f_prime; x_0 = 0, maxiter = 5) === nothing
    @test newtonroot(f; x_0 = 0, maxiter = 5) === nothing
end

@testset "tolerance" begin
    f(x) = 3
    f_prime(x) = 3
    @test newtonroot(f, f_prime; x_0 = 4, tol = 3) == 4
    @test newtonroot(f; x_0 = 4, tol = 3) == 4
end