text:
using ExtensibleMCMC
using Test, StaticArrays, Distributions
const eMCMC = ExtensibleMCMC

@testset "schedule" begin
    schedule_params = (
        num_mcmc_iter = 10,
        num_params = 4,
        exclude_params = [(1,3:8), (2,4:2:10)],
    )
    schedule = eMCMC.MCMCSchedule(schedule_params...)

    expected = (
        (1,1),  (1,2),  (1,3),  (1,4),
        (2,1),  (2,2),  (2,3),  (2,4),
                (3,2),  (3,3),  (3,4),
                        (4,3),  (4,4),
                (5,2),  (5,3),  (5,4),  (5,5),  (5,6),
                        (6,3),          (6,5),  (6,6),
                (7,2),  (7,3),          (7,5),  (7,6),
                        (8,3),                  (8,6),
        (9,1),  (9,2),  (9,3),                  (9,6),
        (10,1),         (10,3),         (10,5), (10,6),
    )

    for (i,s) in enumerate(schedule)
        @test expected[i] == (s.mcmciter, s.pidx)
        s.mcmciter == 5 && s.pidx == 3 && eMCMC.reschedule!(
            schedule, 2, [4], [(5, 8:9)]
        )
    end
end

@testset "adaptation for random walk" begin
    template = eMCMC.AdaptationUnifRW{Float64}(0.234,100,1.0,1.0e-12,1e7,1e2,1)
    @test template == eMCMC.AdaptationUnifRW(1.0)
    @test template == eMCMC.AdaptationUnifRW([2.0])
    @test template == eMCMC.AdaptationUnifRW(SVector{1}(3.0))

    longerθ = eMCMC.AdaptationUnifRW([1.0, 2.0])
    @test template != longerθ
    @test eMCMC.isequal_except(template, longerθ, :N)

    longer_staticθ = eMCMC.AdaptationUnifRW(SVector{3}(1.0, 2.0, 3.0))
    @test template != longer_staticθ
    @test eMCMC.isequal_except(template, longer_staticθ, :N)

    new_scale = eMCMC.AdaptationUnifRW(1.0; scale=3.0)
    @test template != new_scale
    @test eMCMC.isequal_except(template, new_scale, :scale)

    new_params = eMCMC.AdaptationUnifRW(1.0; scale=3.0, target_accpt_rate=0.111,min=10.0)
    @test template != new_scale
    @test eMCMC.isequal_except(template, new_scale, :scale, :target_accpt_rate, :min)
    @test new_params.scale == 3.0
    @test new_params.target_accpt_rate == 0.111
    @test new_params.min == 10.0

    ar_vec = eMCMC.AdaptationUnifRW([1.0, 2.0]; scale=[3.0, 4.0], target_accpt_rate=0.111, min=10.0)
    @test ar_vec.target_accpt_rate == 0.111
    @test ar_vec.min == [10.0, 10.0]
    @test ar_vec.max == [1e7, 1e7]
    @test ar_vec.scale == [3.0, 4.0]
    @test ar_vec.offset == [100.0, 100.0]
    @test ar_vec.N == 2
    @test ar_vec.adapt_every_k_steps == 100

    ar_vec2 = eMCMC.AdaptationUnifRW([1.0, 2.0]; scale=[3.0, 4.0], target_accpt_rate=0.111, min=[10.0, 10.0])
    @test ar_vec == ar_vec2

    ar_svec = eMCMC.AdaptationUnifRW(SVector{2}([1.0, 2.0]); scale=[3.0, 4.0], target_accpt_rate=0.111, min=10.0)
    @test ar_svec.target_accpt_rate == 0.111
    @test ar_svec.min == @SVector [10.0, 10.0]
    @test ar_svec.max == @SVector [1e7, 1e7]
    @test ar_svec.scale == @SVector [3.0, 4.0]
    @test ar_svec.offset == @SVector [100.0, 100.0]
    @test ar_svec.N == 2
    @test ar_svec.adapt_every_k_steps == 100

    ar_svec2 = eMCMC.AdaptationUnifRW(SVector{2}([1.0, 2.0]); scale=[3.0, 4.0], target_accpt_rate=0.111, min=SVector{2}([10.0, 10.0]))
    @test ar_svec == ar_svec2

    ar_svec3 = eMCMC.AdaptationUnifRW(SVector{2}([1.0, 2.0]); scale=[3.0, 4.0], target_accpt_rate=0.111, min=[10.0, 10.0])
    @test ar_svec2 == ar_svec3
end

@testset "mcmc" begin
    # estiate mean of a normal distribution
    μ = [1.0, 2.0]
    Σ = [1.0 0.5; 0.5 1.0]

    #Random.seed!(10)
    trgt = MvNormal(μ, Σ)
    num_obs = 10

    mcmc_params = (
        mcmc = MCMC(
            [
                RandomWalkUpdate(UniformRandomWalk([1.0]), [1]),
                RandomWalkUpdate(UniformRandomWalk([1.0]), [2]),
            ];
            backend=GenericMCMCBackend(), # this is a default anyway
        ),
        num_mcmc_steps = Integer(1e3),
        data = (
            P = GsnTargetLaw(μ, Σ),
            obs = [rand(trgt) for _ in 1:num_obs],
        ),
        θinit = [0.0, 0.0],
        callbacks = eMCMC.Callback[],
    )

    ws = run!(mcmc_params...)
end


text:
using StaticArrays
using LoopVectorization

N = 100000000

function transit(m::SVector{1, Float64}, s::SVector{1, Float64}, x::SVector{2, Float64}, M::SVector{1, Float64}, p::SVector{9, Float64})
    z_m1_ = m[1]
    k_m1_ = s[1]
    n_m1_ = x[1]
    i_m1_ = x[2]
    delta_ = p[5]
    alpha_ = p[6]
    y_m1_ = (exp(z_m1_) * k_m1_ ^ alpha_) * n_m1_ ^ (1.0 - alpha_)
    k__0_ = (1.0 - delta_) * k_m1_ + i_m1_ + y_m1_

    return k__0_
    # oo_1 = SVector(k__0_)
    # res_ = (oo_1,)
    # return res_
end

@fastmath function transit_fm(m::SVector{1, Float64}, s::SVector{1, Float64}, x::SVector{2, Float64}, M::SVector{1, Float64}, p::SVector{9, Float64})
    z_m1_ = m[1]
    k_m1_ = s[1]
    n_m1_ = x[1]
    i_m1_ = x[2]
    delta_ = p[5]
    alpha_ = p[6]
    y_m1_ = (exp(z_m1_) * k_m1_ ^ alpha_) * n_m1_ ^ (1.0 - alpha_)
    k__0_ = (1.0 - delta_) * k_m1_ + i_m1_ + y_m1_

    return k__0_
    # oo_1 = SVector(k__0_)
    # res_ = (oo_1,)
    # return res_
end

@fastmath function transit_loop(m, s, x, M , p)
    N = length(m)
    out = zeros(SVector{1,Float64}, N)

    @simd  for n ∈ 1:N

    delta_ = p[n][5]
    alpha_ = p[n][6]
    z_m1_ = m[n][1]
    k_m1_ = s[n][1]
    n_m1_ = x[n][1]
    i_m1_ = x[n][2]
    y_m1_ = (exp(z_m1_) * k_m1_ ^ alpha_) * n_m1_ ^ (1.0 - alpha_)
    k__0_ = (1.0 - delta_) * k_m1_ + i_m1_ + y_m1_
    oo = SVector(k__0_)
    out[n] = oo
    end

    return out
    # oo_1 = SVector(k__0_)
    # res_ = (oo_1,)
    # return res_
end

@fastmath function transit_loop_multi(m, s, x, M , p)
    N = length(m)
    out = zeros(SVector{1,Float64}, N)

    Threads.@threads  for n ∈ 1:N

    delta_ = p[n][5]
    alpha_ = p[n][6]
    z_m1_ = m[n][1]
    k_m1_ = s[n][1]
    n_m1_ = x[n][1]
    i_m1_ = x[n][2]
    y_m1_ = (exp(z_m1_) * k_m1_ ^ alpha_) * n_m1_ ^ (1.0 - alpha_)
    k__0_ = (1.0 - delta_) * k_m1_ + i_m1_ + y_m1_
    oo = SVector(k__0_)
    out[n] = oo
    end

    return out
    # oo_1 = SVector(k__0_)
    # res_ = (oo_1,)
    # return res_
end

m0 = SVector(0.0)
s0 = SVector(9.5)
x0 = SVector(0.33, 0.23)
p = SVector(0.96, 2.00, 0.4, 0.5,  0.1, 0.3, 0.9, 0.0, 0.01)

# @code_warntype transit(m0, s0, x0, m0, p)

mv = [m0 for i=1:N]
sv = [s0 for i=1:N]
xv = [x0 for i=1:N]
pv = [p  for i=1:N]

println("Warmup")
@time transit.(mv, sv, xv, mv, pv);
@time transit_fm.(mv, sv, xv, mv, pv);

println("Real Test")

@time transit.(mv, sv, xv, mv, pv);
@time transit_fm.(mv, sv, xv, mv, pv);

println("Loop Test")

@time transit_loop(mv, sv, xv, mv, pv);
@time transit_loop(mv, sv, xv, mv, pv);

println("Loop Test2")

@time transit_loop_multi(mv, sv, xv, mv, pv);
@time transit_loop_multi(mv, sv, xv, mv, pv);

ext: jl


text:
using QuAlgorithmZoo, Yao,YaoExtensions
using BitBasis: log2i
using Test
using Random, LinearAlgebra

"""
Quantum singular value decomposition algorithm.

    * `reg`, input register (A, B) as the target matrix to decompose,
    * `circuit_a`, U matrix applied on register A,
    * `circuit_b`, V matrix applied on register B,
    * `optimizer`, the optimizer, normally we use `Adam(lr=0.1)`,
    * `Nc`, log2 number of singular values kept,
    * `maxiter`, the maximum number of iterations.
"""
function train_qsvd!(reg, circuit_a::AbstractBlock{Na}, circuit_b::AbstractBlock{Nb}, optimizer; Nc::Int=min(Na, Nb), maxiter::Int=100) where {Na, Nb}
    nbit = Na+Nb
    c = circuit_qsvd(circuit_a, circuit_b, Nc)

    obs = -mapreduce(i->put(nbit, i=>Z), +, (1:Na..., Na+Nc+1:Na+Nb...))
    params = parameters(c)
    for i = 1:maxiter
        grad = expect'(obs, reg => c).second
        QuAlgorithmZoo.update!(params, grad, optimizer)
        println("Iter $i, Loss = $(Na+expect(obs, copy(reg) |> c))")
        dispatch!(c, params)
    end
end

"""build the circuit for quantum SVD training."""
function circuit_qsvd(circuit_a::AbstractBlock{Na}, circuit_b::AbstractBlock{Nb}, Nc::Int) where {Na, Nb}
    nbit = Na+Nb
    cnots = chain(control(nbit, i+Na, i=>X) for i=1:Nc)
    c = chain(concentrate(nbit, circuit_a, 1:Na), concentrate(nbit, circuit_b, Na+1:nbit), cnots)
end

"""read QSVD results"""
function readout_qsvd(reg::AbstractRegister, circuit_a::AbstractBlock{Na}, circuit_b::AbstractBlock{Nb}, Nc::Int) where {Na, Nb}
    reg = copy(reg) |> concentrate(Na+Nb, circuit_a, 1:Na) |> concentrate(Na+Nb, circuit_b, Na+1:Na+Nb)
    _S = [select(reg, b|b<<Na).state[] for b in basis(Nc)]
    S = abs.(_S)
    order = sortperm(S, rev=true)
    S, _S = S[order], _S[order]
    mat(circuit_a)[order,:]'.*transpose(_S./S), S, transpose(mat(circuit_b)[order,:])
end

"""
    QuantumSVD(M; kwargs...)
Quantum SVD.
    * `M`, the matrix to decompose, size should be (2^Na × 2^Nb), the sum of squared spectrum must be 1.
kwargs includes
    * `Nc`, log2 number of singular values kept,
    * `circuit_a` and `circuit_b`, the circuit ansatz for `U` and `V` matrices,
    * `maxiter`, maximum number of iterations,
    * `optimizer`, default is `Adam(lr=0.1)`.
"""
function QuantumSVD(M::AbstractMatrix; Nc::Int=log2i(min(size(M)...)),
        circuit_a=variational_circuit(log2i(size(M, 1))),
        circuit_b=variational_circuit(log2i(size(M, 2))),
        maxiter=200, optimizer=Adam(lr=0.1))

    dispatch!(circuit_a, :random)
    dispatch!(circuit_b, :random)
    reg = ArrayReg(vec(M))
    train_qsvd!(reg, circuit_a, circuit_b, optimizer, Nc=Nc, maxiter=maxiter)
    readout_qsvd(reg, circuit_a, circuit_b, Nc)
end

@testset "QSVD" begin
    Random.seed!(2)
    # define a matrix of size (2^Na, 2^Nb)
    Na = 2
    Nb = 2

    # the exact result
    M = reshape(rand_state(Na+Nb).state, 1<<Na, 1<<Nb)
    U_exact, S_exact, V_exact = svd(M)

    U, S, V = QuantumSVD(M; maxiter=400)

    @test isapprox(U*Diagonal(S)*V', M, atol=1e-2)
    @test isapprox(abs.(S), S_exact, atol=1e-2)
    @test isapprox(U'*U_exact .|> abs2, I, atol=1e-2)
    @test isapprox(V'*V_exact .|> abs2, I, atol=1e-2)
end


text:
#= Custom calibration configuration file. =#

using Distributions
using StatsBase
using LinearAlgebra
using CalibrateEDMF
using CalibrateEDMF.ModelTypes
using CalibrateEDMF.DistributionUtils
using CalibrateEDMF.ReferenceModels
using CalibrateEDMF.ReferenceStats
using CalibrateEDMF.LESUtils
using CalibrateEDMF.TurbulenceConvectionUtils
using CalibrateEDMF.HelperFuncs
# Import EKP modules
using JLD2
using EnsembleKalmanProcesses.ParameterDistributions

struct SCT1Train end
struct SCT1Val end

function get_config()
    config = Dict()
    # Flags for saving output data
    config["output"] = get_output_config()
    # Define regularization of inverse problem
    config["regularization"] = get_regularization_config()
    # Define reference used in the inverse problem
    config["reference"] = get_reference_config(SCT1Train())
    # Define reference used for validation
    config["validation"] = get_reference_config(SCT1Val())
    # Define the parameter priors
    config["prior"] = get_prior_config()
    # Define the kalman process
    config["process"] = get_process_config()
    # Define the SCM static configuration
    config["scm"] = get_scm_config()
    return config
end

function get_output_config()
    config = Dict()
    config["outdir_root"] = pwd()
    config["overwrite_scm_file"] = false # Flag for overwritting SCM input file
    return config
end

function get_regularization_config()
    config = Dict()
    # Regularization of observations: mean and covariance
    config["perform_PCA"] = true # Performs PCA on data
    config["variance_loss"] = 1.0e-2 # Variance truncation level in PCA
    config["normalize"] = true  # whether to normalize data by pooled variance
    config["tikhonov_mode"] = "relative" # Tikhonov regularization
    config["tikhonov_noise"] = 1.0e-6 # Tikhonov regularization
    config["dim_scaling"] = true # Dimensional scaling of the loss

    # Parameter regularization: L2 regularization with respect to prior mean.
    #  - Set to `nothing` to use prior covariance as regularizer,
    #  - Set to a float for isotropic parameter regularization.
    #  - Pass a dictionary of lists similar to config["prior_mean"] for
    #       anisotropic regularization. The dictionary must be complete.
    #       If you want to avoid regularizing a certain parameter, set the entry
    #       to [0].
    # To turn off regularization, set config["process"]["augmented"] to false.
    # config["l2_reg"] = 1.0
    return config
end

function get_process_config()
    config = Dict()
    config["N_iter"] = 50
    config["N_ens"] = 100 # Must be 2p+1 when algorithm is "Unscented"
    config["algorithm"] = "Inversion" # "Sampler", "Unscented", "Inversion"
    config["noisy_obs"] = false # Choice of covariance in evaluation of y_{j+1} in EKI. True -> Γy, False -> 0
    # Artificial time stepper of the EKI.
    config["Δt"] = 1.0
    # Whether to augment the outputs with the parameters for regularization
    config["augmented"] = false
    config["failure_handler"] = "sample_succ_gauss" #"high_loss" #"sample_succ_gauss"
    return config
end

function get_reference_config(::SCT1Train)
    config = Dict()

    # AMIP data: October
    cfsite_numbers = (17, 18, 20, 22, 23)
    les_kwargs = (forcing_model = "HadGEM2-A", month = 10, experiment = "amip")
    ref_dirs = [get_cfsite_les_dir(cfsite_number; les_kwargs...) for cfsite_number in cfsite_numbers]
    suffixes = [get_gcm_les_uuid(cfsite_number; les_kwargs...) for cfsite_number in cfsite_numbers]

    n_repeat = length(ref_dirs)
    config["case_name"] = repeat(["LES_driven_SCM"], n_repeat)
    # Flag to indicate whether reference data is from a perfect model (i.e. SCM instead of LES)
    config["y_reference_type"] = LES()
    config["Σ_reference_type"] = LES()
    config["y_names"] =
        repeat([["s_mean", "ql_mean", "qt_mean", "total_flux_qt", "total_flux_s", "u_mean", "v_mean"]], n_repeat)
    config["y_dir"] = ref_dirs
    config["scm_suffix"] = suffixes
    config["scm_parent_dir"] = repeat(["scm_init"], n_repeat)
    config["t_start"] = repeat([3.0 * 3600], n_repeat)
    config["t_end"] = repeat([6.0 * 3600], n_repeat)
    # Use full LES timeseries for covariance
    config["Σ_t_start"] = repeat([-5.75 * 24 * 3600], n_repeat)
    config["Σ_t_end"] = repeat([6.0 * 3600], n_repeat)
    config["write_full_stats"] = false
    return config
end

function get_reference_config(::SCT1Val)
    config = Dict()

    # AMIP4K data: July
    cfsite_numbers = (3, 5, 8, 11, 14)
    les_kwargs = (forcing_model = "HadGEM2-A", month = 7, experiment = "amip4K")
    ref_dirs = [get_cfsite_les_dir(cfsite_number; les_kwargs...) for cfsite_number in cfsite_numbers]
    suffixes = [get_gcm_les_uuid(cfsite_number; les_kwargs...) for cfsite_number in cfsite_numbers]

    n_repeat = length(ref_dirs)
    config["case_name"] = repeat(["LES_driven_SCM"], n_repeat)
    # Flag to indicate whether reference data is from a perfect model (i.e. SCM instead of LES)
    config["y_reference_type"] = LES()
    config["Σ_reference_type"] = LES()
    config["y_names"] =
        repeat([["s_mean", "ql_mean", "qt_mean", "total_flux_qt", "total_flux_s", "u_mean", "v_mean"]], n_repeat)
    config["y_dir"] = ref_dirs
    config["scm_suffix"] = suffixes
    config["scm_parent_dir"] = repeat(["scm_init"], n_repeat)
    config["t_start"] = repeat([3.0 * 3600], n_repeat)
    config["t_end"] = repeat([6.0 * 3600], n_repeat)
    # Use full LES timeseries for covariance
    config["Σ_t_start"] = repeat([-5.75 * 24 * 3600], n_repeat)
    config["Σ_t_end"] = repeat([6.0 * 3600], n_repeat)
    config["write_full_stats"] = false
    return config
end

function get_prior_config()
    config = Dict()
    config["constraints"] = Dict(
        # entrainment parameters
        "general_ent_params" => [repeat([no_constraint()], 69)...],
        "turbulent_entrainment_factor" => [bounded(0.0, 10.0)],
    )

    # TC.jl prior mean
    config["prior_mean"] = Dict(
        # entrainment parameters
        "general_ent_params" => 0.1 .* (rand(69) .- 0.5),
        "turbulent_entrainment_factor" => [0.075],
    )

    config["unconstrained_σ"] = 1.0
    return config
end

function get_scm_config()
    config = Dict()
    config["namelist_args"] = [
        ("time_stepping", "dt_min", 1.0),
        ("time_stepping", "dt_max", 2.0),
        ("stats_io", "frequency", 60.0),
        ("grid", "dz", 50.0),
        ("grid", "nz", 80),
        ("turbulence", "EDMF_PrognosticTKE", "entrainment", "NN"),
        ("turbulence", "EDMF_PrognosticTKE", "area_limiter_power", 0.0),
        ("turbulence", "EDMF_PrognosticTKE", "entr_dim_scale", "none"),
    ]
    return config
end


text:
function _precompile_()
    ccall(:jl_generating_output, Cint, ()) == 1 || return nothing

    precompile(Tuple{Type{LoopVectorization.ArrayRefStruct},LoopVectorization.LoopSet,LoopVectorization.ArrayReferenceMeta,Vector{Symbol},Vector{Int64}})
    precompile(Tuple{Type{LoopVectorization.LoopSet},Expr,Symbol})
    precompile(Tuple{Type{LoopVectorization.LoopSet},Symbol})
    precompile(Tuple{Type{Vector{Tuple{Int64,LoopVectorization.NumberType}}},Vector{Int64}})
    # precompile(Tuple{typeof(Base._mapreduce),Base.ComposedFunction{Type{Float64},typeof(length)},typeof(Base.mul_prod),IndexLinear,Vector{LoopVectorization.Loop}})
    precompile(Tuple{typeof(LoopVectorization._avx_loopset),Core.SimpleVector,Core.SimpleVector,Core.SimpleVector,Core.SimpleVector,Core.SimpleVector,Any})
    precompile(Tuple{typeof(LoopVectorization.add_broadcast!),LoopVectorization.LoopSet,Symbol,Symbol,Vector{Symbol},Type{Vector{Float64}},Int64})
    precompile(Tuple{typeof(LoopVectorization.add_compute!),LoopVectorization.LoopSet,Symbol,Expr,Int64,Int64,LoopVectorization.ArrayReferenceMetaPosition})
    precompile(Tuple{typeof(LoopVectorization.add_compute!),LoopVectorization.LoopSet,Symbol,Expr,Int64,Int64,Nothing})
    precompile(Tuple{typeof(LoopVectorization.add_constant!),LoopVectorization.LoopSet,Float64,Vector{Symbol},Symbol,Int64})
    precompile(Tuple{typeof(LoopVectorization.add_if!),LoopVectorization.LoopSet,Symbol,Expr,Int64,Int64,LoopVectorization.ArrayReferenceMetaPosition})
    precompile(Tuple{typeof(LoopVectorization.add_load!),LoopVectorization.LoopSet,LoopVectorization.Operation,Bool,Bool})
    precompile(Tuple{typeof(LoopVectorization.add_mref!),LoopVectorization.LoopSet,LoopVectorization.ArrayReferenceMeta,Int64,Type{VectorizationBase.PermutedDimsStridedPointer{(4, 2, 3, 1),(4, 2, 3, 1),Float64,VectorizationBase.PermutedDimsStridedPointer{(3, 1, 4, 2),(2, 4, 1, 3),Float64,VectorizationBase.PackedStridedPointer{Float64,3}}}},Symbol})
    precompile(Tuple{typeof(LoopVectorization.add_pow!),LoopVectorization.LoopSet,Symbol,Expr,Float64,Int64,Int64})
    precompile(Tuple{typeof(LoopVectorization.add_pow!),LoopVectorization.LoopSet,Symbol,Expr,Int64,Int64,Int64})
    precompile(Tuple{typeof(LoopVectorization.argmeta_and_consts_description),LoopVectorization.LoopSet,Vector{Symbol}})
    precompile(Tuple{typeof(LoopVectorization.array_reference_meta!),LoopVectorization.LoopSet,Symbol,SubArray{Any,1,Vector{Any},Tuple{UnitRange{Int64}},true},Int64,Nothing})
    precompile(Tuple{typeof(LoopVectorization.avx_loopset),Vector{LoopVectorization.Instruction},Vector{LoopVectorization.OperationStruct},Vector{LoopVectorization.ArrayRefStruct},Core.SimpleVector,Core.SimpleVector,Core.SimpleVector,Any})
    precompile(Tuple{typeof(LoopVectorization.check_macro_kwarg),Expr})
    precompile(Tuple{typeof(LoopVectorization.cost_vec_buf),LoopVectorization.LoopSet})
    precompile(Tuple{typeof(LoopVectorization.evaluate_cost_tile),LoopVectorization.LoopSet,Vector{Symbol},LoopVectorization.UnrollSymbols})
    precompile(Tuple{typeof(LoopVectorization.evaluate_cost_unroll),LoopVectorization.LoopSet,Vector{Symbol},Symbol,Float64})
    precompile(Tuple{typeof(LoopVectorization.generate_call),LoopVectorization.LoopSet,Tuple{Int8,Int8,Int8},Bool})
    precompile(Tuple{typeof(LoopVectorization.pushop!),LoopVectorization.LoopSet,LoopVectorization.Operation,Symbol})
    precompile(Tuple{typeof(LoopVectorization.reject_candidate),LoopVectorization.Operation,Symbol,Symbol})
    precompile(Tuple{typeof(LoopVectorization.stride_penalty),LoopVectorization.LoopSet,Vector{Symbol}})
    precompile(Tuple{typeof(LoopVectorization.substitute_broadcast),Expr,Symbol})
    precompile(Tuple{typeof(println),IOBuffer,Vector{LoopVectorization.Operation}})
    precompile(Tuple{typeof(push!),LoopVectorization.LoopSet,Expr,Int64,Int64})
end


text:
module MultivariateTest
using OnlineStats, Base.Test

@testset "Multivariate" begin
@testset "KMeans" begin
    x = vcat(randn(1000, 5), 10 + randn(1000, 5))
    o = KMeans(x, 2)
    o = KMeans(x, 2, 10)

    @test nobs(o)       == 2000
    @test size(o.value) == (5, 2)
end
end
end#module


text:
# Autogenerated wrapper script for Wayland_jll for x86_64-linux-musl
export libwayland_client, libwayland_cursor, libwayland_egl, libwayland_server, wayland_scanner

using Expat_jll
using Libffi_jll
using XML2_jll
JLLWrappers.@generate_wrapper_header("Wayland")
JLLWrappers.@declare_library_product(libwayland_client, "libwayland-client.so.0")
JLLWrappers.@declare_library_product(libwayland_cursor, "libwayland-cursor.so.0")
JLLWrappers.@declare_library_product(libwayland_egl, "libwayland-egl.so.1")
JLLWrappers.@declare_library_product(libwayland_server, "libwayland-server.so.0")
JLLWrappers.@declare_executable_product(wayland_scanner)
function __init__()
    JLLWrappers.@generate_init_header(Expat_jll, Libffi_jll, XML2_jll)
    JLLWrappers.@init_library_product(
        libwayland_client,
        "lib/libwayland-client.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libwayland_cursor,
        "lib/libwayland-cursor.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libwayland_egl,
        "lib/libwayland-egl.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libwayland_server,
        "lib/libwayland-server.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_executable_product(
        wayland_scanner,
        "bin/wayland-scanner",
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
using BinaryBuilder, Pkg

name = "libgeotiff"
version = v"1.7.1"

# Collection of sources required to complete build
sources = [
    ArchiveSource("https://github.com/OSGeo/libgeotiff/releases/download/$version/libgeotiff-$version.tar.gz",
                  "05ab1347aaa471fc97347d8d4269ff0c00f30fa666d956baba37948ec87e55d6"),
]

# Bash recipe for building across all platforms
script = raw"""

cd $WORKSPACE/srcdir/libgeotiff-*/

mkdir build && cd build

cmake -DCMAKE_INSTALL_PREFIX=$prefix \
      -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TARGET_TOOLCHAIN} \
      -DCMAKE_BUILD_TYPE=Release \
      -DBUILD_SHARED_LIBS=ON \
      ..

make -j${nproc}
make install
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = supported_platforms()

# The products that we will ensure are always built
products = [
    LibraryProduct("libgeotiff", :libgeotiff),
    ExecutableProduct("makegeo", :makegeo),
    ExecutableProduct("geotifcp", :geotifcp),
    ExecutableProduct("listgeo", :listgeo),
    ExecutableProduct("applygeo", :applygeo),
]

# Dependencies that must be installed before this package can be built
dependencies = [
    Dependency("PROJ_jll"; compat="~900.0"),
    Dependency("Libtiff_jll"; compat="4.3"),
    Dependency("LibCURL_jll"),
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies; julia_compat="1.6")


text:
# Load packages
using Dojo
using Plots
using Random
using MeshCat

# Open visualizer
vis= Visualizer()
open(vis)

# Include new files
include( "../utils.jl")

mech = get_mechanism(:block2d, timestep=0.05, g=-9.81, friction_coefficient=0.2, radius = 0.05, side = 0.50);
initialize!(mech, :block2d, x=[-1,1.], v=[2,1.], θ=0.1, ω=2.)
storage = simulate!(mech, 5.0, record=true,
    opts=SolverOptions(btol=1e-6, rtol=1e-6, verbose=false))
visualize(mech, storage, vis=vis, show_contact = true)

################################################################################
# Generate & Save Dataset
################################################################################
init_kwargs = Dict(:xlims => [[0,0.2], [1,0.4]],
				   :vlims => [-2ones(2), ones(2)],
				   :θlims => [-π, π],
				   :ωlims => [-10, 10])
mech_kwargs = Dict(:friction_coefficient => 0.1, :radius => 0.05, :side => 0.5)
generate_dataset(:block2d, H=0.75, N=15,
	opts=SolverOptions(btol=3e-4, rtol=3e-4),
	init_kwargs=init_kwargs,
	mech_kwargs=mech_kwargs)

################################################################################
# Load Dataset
################################################################################
params0, trajs0, pairs0 = open_dataset(:block2d; N = 15, mech_kwargs...)

data0 = params0[:data]

################################################################################
# Optimization Objective: Evaluation & Gradient
################################################################################
clean_loss(:block2d, pairs0, data0, opts=SolverOptions(btol=3e-4, rtol=3e-4))

[clean_loss(:block2d, pairs0, data0 + [i;zeros(27)], opts=SolverOptions(btol=3e-4, rtol=3e-4))
	for i in Vector(-0.10:0.01:0.1)]
[clean_loss(:block2d, pairs0, data0 + [0;i;zeros(26)], opts=SolverOptions(btol=3e-4, rtol=3e-4))
	for i in Vector(-0.10:0.01:0.1)]

plot(hcat([p[1][1:3] for p in pairs0]...)')
plot(hcat([p[1][4:6] for p in pairs0]...)')
plot(hcat([p[1][7:10] for p in pairs0]...)')
plot(hcat([p[1][11:13] for p in pairs0]...)')

################################################################################
# Optimization Algorithm: L-BFGS:
# We learn a single coefficient of friction and a 4 contact locations [y,z] -> 9 params in total
################################################################################

include("../quasi_newton.jl")
function d2data(d)
	friction_coefficient = d[1]
	data = [friction_coefficient; 0;0; 0.05; 0; +d[2]; +d[3];
			friction_coefficient; 0;0; 0.05; 0; +d[4]; +d[5];
			friction_coefficient; 0;0; 0.05; 0; +d[6]; +d[7];
			friction_coefficient; 0;0; 0.05; 0; +d[8]; +d[9];
			]
	return data
end

∇d2data = ForwardDiff.jacobian(d -> d2data(d), zeros(9))

d0 = [0.40, +0.50, +0.50, +0.50, -0.50, -0.50, +0.50, -0.50, -0.50]
lower = [0.00, +0.05, +0.05, +0.05, -1.00, -1.00, +0.05, -1.00, -1.00]
upper = [0.80, +1.00, +1.00, +1.00, -0.05, -0.05, +1.00, -0.05, -0.05]

function f0(d; rot=0)
	return clean_loss(:block2d, pairs0, d2data(d), n_sample=15, rot=rot, opts=SolverOptions(btol=3e-4, rtol=3e-4))[1]
end

function fgH0(d; rot=0)
	f, g, H = clean_loss(:block2d, pairs0, d2data(d), n_sample=15, rot=rot, opts=SolverOptions(btol=3e-4, rtol=3e-4))
	return f, ∇d2data' * g, ∇d2data' * H * ∇d2data
end
# dsol = bfgs_solve(f0, g0, H0, d0, iter=200, gtol=1e-8, ftol=3e-5,
dsol = quasi_newton_solve(f0, fgH0, d0, iter=200, gtol=1e-8, ftol=1e-6,
	lower=lower, upper=upper, reg=1e-6)

d0
f0(d0)
f0(dsol)

# We can learn the coefficient of friction and the side dimenson of the cube
# form 15*0.75 seconds of recording. We use the simulator to evaluate the loss
# and its gradients by differentiating through the simulator. With gradient
# information we can use L-BFGS

################################################################################
# Optimization Algorithm: L-BFGS:
# We learn a single coefficient of friction and a single side length.
################################################################################
include("../quasi_newton.jl")
function d2data(d)
	friction_coefficient, side = d
	data = [friction_coefficient; 0;0; 0.05; 0; +side; +side;
			friction_coefficient; 0;0; 0.05; 0; +side; -side;
			friction_coefficient; 0;0; 0.05; 0; -side; +side;
			friction_coefficient; 0;0; 0.05; 0; -side; -side;
			]
	return data
end
∇d2data = ForwardDiff.jacobian(d -> d2data(d), zeros(2))

d0 = [0.40, 0.75]
lower = [0.00, 0.05]
upper = [0.80, 1.00]

function f0(d; rot=0)
	return clean_loss(:block2d, pairs0, d2data(d), n_sample=15, rot=rot, opts=SolverOptions(btol=3e-4, rtol=3e-4))[1]
end

function fgH0(d; rot=0)
	f, g, H = clean_loss(:block2d, pairs0, d2data(d), n_sample=15, rot=rot, opts=SolverOptions(btol=3e-4, rtol=3e-4))
	return f, ∇d2data' * g, ∇d2data' * H * ∇d2data
end
dsol = quasi_newton_solve(f0, fgH0, d0, iter=200, gtol=1e-8, ftol=1e-6,
	lower=lower, upper=upper, reg=1e-9)

d0
f0(d0)
f0(dsol)
f0([0.1, 0.5])

# We can learn the coefficient of friction and the side dimenson of the cube
# form 15*0.75 seconds of recording. We use the simulator to evaluate the loss
# and its gradients by differentiating through the simulator. With gradient
# information we can use L-BFGS

################################################################################
# Visualization
################################################################################


text:
######################################################################################################
# Copy from https://bitbucket.org/wmkoolen/tidnabbil/src/master/purex_games_paper/ by Wouter Koolen  #
######################################################################################################
# binary search for zero of continuous function f
# with one zero-crossing in [lo, hi]
function binary_search(f, lo, hi; ϵ = 1e-10, maxiter = 100)
    for i in 1:maxiter
        mid = (lo+hi)/2;
        if mid == lo || mid == hi
            return mid;
        end
        fmid = f(mid);
        if fmid < -ϵ
            lo = mid
        elseif fmid > ϵ
            hi = mid;
        else
            return mid;
        end
    end
    @warn "binary_search did not reach tolerance $ϵ in $maxiter iterations.\nf($(lo)) = $(f(lo))\nf($(hi)) = $(f(hi)),\nmid would be $((lo+hi)/2)";
    return (lo+hi)/2;
end


text:
struct QuadrupedV2{I, T} <: Model{I, T}
    n::Int
    m::Int
    d::Int

    g
    μ

	mb # body mass
	mf # foot mass
	Ix # body inertia (x)
	Iy # body inertia (y)
	Iz # body inertia (z)

	l_torso
	w_torso

    # joint limits
    qL
    qU

    # torque limits
    uL
    uU

    nq
    nu
    nc
    nf
    nb
    ns

    idx_u
    idx_λ
    idx_b
    idx_ψ
    idx_η
    idx_s

	joint_friction
end

# Dimensions
nq = 3 + 3 + 3 * 4        # configuration dimension
nu = 4 * 3                # control dimension
nc = 4                    # number of contact points
nf = 4                    # number of parameters for friction cone
nb = nc * nf
ns = 1

# World parameters
μ = 0.5      # coefficient of friction
g = 9.81     # gravity

# # ~Unitree A1
# mf = 0.01
#
# # TRUCK ONLY TODO: parallel axis theorem to add shoulders
# mb = 4.713
# Ix = 0.01683993
# Iy = 0.056579028
# Iz = 0.064713601
#
# l_torso = 0.5 * 0.267 # dimension from com
# w_torso = 0.5 * 0.194 # dimension from com

## Mini Cheetah
mb = 9.0
# mf = 0.1 * mb
mf = 0.025 * mb
Ix = 0.07
Iy = 0.26
Iz = 0.242

l_torso = 0.5 * 0.38 # dimension from com
w_torso = 0.5 * 0.203 # dimension from com

n = 2 * nq
m = nu + nc + nb + nc + nb + ns
d = 0

idx_u = (1:nu)
idx_λ = nu .+ (1:nc)
idx_b = nu + nc .+ (1:nb)
idx_ψ = nu + nc + nb .+ (1:nc)
idx_η = nu + nc + nb + nc .+ (1:nb)
idx_s = nu + nc + nb + nc + nb .+ (1:ns)

qL = -Inf * ones(nq)
qU = Inf * ones(nq)

uL = -100.0 * ones(nu)
uU = 100.0 * ones(nu)

function lagrangian_derivatives(model, q, v)
	D1L = -1.0 * C_func(model, q, v)
    D2L = M_func(model, q) * v
	return D1L, D2L
end

# Methods
function M_func(model::QuadrupedV2, q)
	M = Diagonal([model.mb * ones(3) ;
				  model.Ix; model.Iy; model.Iz;
				  model.mf * ones(12)])
	return M
end

function C_func(model::QuadrupedV2, q, q̇)
	SVector{18}([0.0, 0.0, model.g * model.mb,
				 cross(q̇[4:6], Diagonal([model.Ix, model.Iy, model.Iz]) * q̇[4:6])...,
				 0.0, 0.0, model.g * model.mf,
				 0.0, 0.0, model.g * model.mf,
				 0.0, 0.0, model.g * model.mf,
				 0.0, 0.0, model.g * model.mf])
end

function kinematics(model::QuadrupedV2, q)
	p_torso = q[1:3]

	rot = MRP(q[4:6]...)

	pw1 = q[6 .+ (1:3)]
	pw2 = q[9 .+ (1:3)]
	pw3 = q[12 .+ (1:3)]
	pw4 = q[15 .+ (1:3)]

	SVector{12}([pw1; pw2; pw3; pw4])
end

function ϕ_func(model::QuadrupedV2, q)
	k = kinematics(model, q)

	@SVector [k[3], k[6], k[9], k[12]]
end

function skew(x)
	SMatrix{3,3}([0.0 -x[3] x[2];
	               x[3] 0.0 -x[1];
				   -x[2] x[1] 0.0])
end

function B_func(model::QuadrupedV2, q)
	p_torso = q[1:3]

	rot = MRP(q[4:6]...)

	# r in world frame
	r1 = q[6 .+ (1:3)] - p_torso
	r2 = q[9 .+ (1:3)] - p_torso
	r3 = q[12 .+ (1:3)] - p_torso
	r4 = q[15 .+ (1:3)] - p_torso

	z3 = zeros(3, 3)

	SMatrix{18, 12}([I I I I;
	                transpose(rot) * skew(r1) transpose(rot) * skew(r2) transpose(rot) * skew(r3) transpose(rot) * skew(r4);
					-I z3 z3 z3;
					z3 -I z3 z3;
					z3 z3 -I z3;
					z3 z3 z3 -I])
end

function N_func(model::QuadrupedV2, q)
	# p_torso = q[1:3]
	#
	# rot = MRP(q[4:6]...)
	#
	# # r in world frame
	# r1 = q[6 .+ (1:3)] - p_torso
	# r2 = q[9 .+ (1:3)] - p_torso
	# r3 = q[12 .+ (1:3)] - p_torso
	# r4 = q[15 .+ (1:3)] - p_torso

	z3 = zeros(3, 3)

	N = [zeros(6, 12);
	     I z3 z3 z3;
	     z3 I z3 z3;
	     z3 z3 I z3;
	     z3 z3 z3 I]

	idx = collect([3, 6, 9, 12])
	transpose(N[:, idx])
end

function _P_func(model::QuadrupedV2, q)
	# p_torso = q[1:3]
	#
	# rot = MRP(q[4:6]...)
	#
	# # r in world frame
	# r1 = q[6 .+ (1:3)] - p_torso
	# r2 = q[9 .+ (1:3)] - p_torso
	# r3 = q[12 .+ (1:3)] - p_torso
	# r4 = q[15 .+ (1:3)] - p_torso

	z3 = zeros(3, 3)

	P = [zeros(6, 12);
		 I z3 z3 z3;
		 z3 I z3 z3;
		 z3 z3 I z3;
		 z3 z3 z3 I]

	idx = collect([1, 2, 4, 5, 7, 8, 10, 11])
	transpose(P[:, idx])
end

function P_func(model::QuadrupedV2, q)
	_P = _P_func(model, q)

	map = [1.0 0.0;
		   0.0 1.0;
		   -1.0 0.0;
		   0.0 -1.0]

	return [map * _P[1:2, :];
	        map * _P[3:4, :];
			map * _P[5:6, :];
			map * _P[7:8, :]]
end

function friction_cone(model::QuadrupedV2, u)
	λ = view(u, model.idx_λ)
	b = view(u, model.idx_b)

	return @SVector [model.μ * λ[1] - sum(view(b, 1:4)),
					 model.μ * λ[2] - sum(view(b, 5:8)),
					 model.μ * λ[3] - sum(view(b, 9:12)),
					 model.μ * λ[4] - sum(view(b, 13:16))]
end

function maximum_dissipation(model::QuadrupedV2{Discrete, FixedTime}, x⁺, u, h)
	q3 = view(x⁺, model.nq .+ (1:model.nq))
	q2 = view(x⁺, 1:model.nq)
	ψ = view(u, model.idx_ψ)
	ψ_stack = [ψ[1] * ones(4); ψ[2] * ones(4); ψ[3] * ones(4); ψ[4] * ones(4)]
	η = view(u, model.idx_η)
	return P_func(model, q3) * (q3 - q2) / h + ψ_stack - η
end

function fd(model::QuadrupedV2{Discrete, FixedTime}, x⁺, x, u, w, h, t)
	q3 = view(x⁺, model.nq .+ (1:model.nq))
	q2⁺ = view(x⁺, 1:model.nq)
	q2⁻ = view(x, model.nq .+ (1:model.nq))
	q1 = view(x, 1:model.nq)
	u_ctrl = view(u, model.idx_u)
	λ = view(u, model.idx_λ)
	b = view(u, model.idx_b)

	qm1 = 0.5 * (q1 + q2⁺)
    vm1 = (q2⁺ - q1) / h
    qm2 = 0.5 * (q2⁺ + q3)
    vm2 = (q3 - q2⁺) / h

	D1L1, D2L1 = lagrangian_derivatives(model, qm1, vm1)
	D1L2, D2L2 = lagrangian_derivatives(model, qm2, vm2)

	[q2⁺ - q2⁻;
     (0.5 * h * D1L1 + D2L1 + 0.5 * h * D1L2 - D2L2
     + B_func(model, qm2) * SVector{12}(u_ctrl)
     + transpose(N_func(model, q3)) * SVector{4}(λ)
     + transpose(P_func(model, q3)) * SVector{16}(b))]
end

function fd(model::QuadrupedV2{Discrete, FreeTime}, x⁺, x, u, w, h, t)
	q3 = view(x⁺, model.nq .+ (1:model.nq))
	q2⁺ = view(x⁺, 1:model.nq)
	q2⁻ = view(x, model.nq .+ (1:model.nq))
	q1 = view(x, 1:model.nq)
	u_ctrl = view(u, model.idx_u)
	λ = view(u, model.idx_λ)
	b = view(u, model.idx_b)
	h = u[end]

	qm1 = 0.5 * (q1 + q2⁺)
    vm1 = (q2⁺ - q1) / h
    qm2 = 0.5 * (q2⁺ + q3)
    vm2 = (q3 - q2⁺) / h

	D1L1, D2L1 = lagrangian_derivatives(model, qm1, vm1)
	D1L2, D2L2 = lagrangian_derivatives(model, qm2, vm2)

	[q2⁺ - q2⁻;
     (0.5 * h * D1L1 + D2L1 + 0.5 * h * D1L2 - D2L2
     + B_func(model, qm2) * SVector{12}(u_ctrl)
     + transpose(N_func(model, q3)) * SVector{4}(λ)
     + transpose(P_func(model, q3)) * SVector{16}(b))]
end

function maximum_dissipation(model::QuadrupedV2{Discrete, FreeTime}, x⁺, u, h)
	q3 = view(x⁺, model.nq .+ (1:model.nq))
	q2 = view(x⁺, 1:model.nq)
	ψ = view(u, model.idx_ψ)
	ψ_stack = [ψ[1] * ones(4); ψ[2] * ones(4); ψ[3] * ones(4); ψ[4] * ones(4)]
	η = view(u, model.idx_η)
	h = u[end]
	return P_func(model, q3) * (q3 - q2) / h + ψ_stack - η
end

model = QuadrupedV2{Discrete, FixedTime}(n, m, d,
			  g,
			  μ,
			  mb,
			  mf,
			  Ix,
			  Iy,
			  Iz,
			  l_torso,
			  w_torso,
			  qL, qU,
			  uL, uU,
			  nq,
			  nu,
			  nc,
			  nf,
			  nb,
			  ns,
			  idx_u,
			  idx_λ,
			  idx_b,
			  idx_ψ,
			  idx_η,
			  idx_s,
			  nothing)

function visualize!(vis, model::QuadrupedV2, q;
      r = 0.025, Δt = 0.1)

	default_background!(vis)

	setobject!(vis["torso"],
    	Rect(Vec(-model.l_torso, -model.w_torso, -0.025),Vec(2.0 * model.l_torso, 2.0 * model.w_torso, 0.05)),
    	MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0)))

	feet1 = setobject!(vis["feet1"], Sphere(Point3f0(0),
        convert(Float32, r)),
        MeshPhongMaterial(color = RGBA(1.0, 165.0 / 255.0, 0, 1.0)))

	feet2 = setobject!(vis["feet2"], Sphere(Point3f0(0),
		convert(Float32, r)),
		MeshPhongMaterial(color = RGBA(1.0, 165.0 / 255.0, 0, 1.0)))

	feet3 = setobject!(vis["feet3"], Sphere(Point3f0(0),
		convert(Float32, r)),
		MeshPhongMaterial(color = RGBA(1.0, 165.0 / 255.0, 0, 1.0)))

	feet4 = setobject!(vis["feet4"], Sphere(Point3f0(0),
        convert(Float32, r)),
        MeshPhongMaterial(color = RGBA(1.0, 165.0 / 255.0, 0, 1.0)))

	anim = MeshCat.Animation(convert(Int, floor(1.0 / Δt)))

	T = length(q)
	p_shift = [0.0, 0.0, r]

	for t = 1:T
		MeshCat.atframe(anim, t) do
			rot = MRP(q[t][4:6]...)

			p_torso = [q[t][1]; q[t][2]; q[t][3]] + p_shift
			p_foot1 = q[t][6 .+ (1:3)] + p_shift
			p_foot2 = q[t][9 .+ (1:3)] + p_shift
			p_foot3 = q[t][12 .+ (1:3)] + p_shift
			p_foot4 = q[t][15 .+ (1:3)] + p_shift

			settransform!(vis["torso"], compose(Translation(p_torso), LinearMap(MRP(q[t][4:6]...))))
			settransform!(vis["feet1"], Translation(p_foot1))
			settransform!(vis["feet2"], Translation(p_foot2))
			settransform!(vis["feet3"], Translation(p_foot3))
			settransform!(vis["feet4"], Translation(p_foot4))
		end
	end

	# settransform!(vis["/Cameras/default"],
	#     compose(Translation(0.0, 0.0, -1.0), LinearMap(RotZ(-pi / 2.0))))

	MeshCat.setanimation!(vis, anim)
end

function initial_configuration(model::QuadrupedV2)
	q1 = zeros(model.nq)

	# position
	q1[3] = 0.5

	# orientation
	mrp = MRP(RotZ(0.0))
	q1[4:6] = [mrp.x; mrp.y; mrp.z]

	# feet positions (in body frame)
	q1[7:9] = [model.l_torso; model.w_torso; 0.0]
	q1[10:12] = [model.l_torso; -model.w_torso; 0.0]
	q1[13:15] = [-model.l_torso; model.w_torso; 0.0]
	q1[16:18] = [-model.l_torso; -model.w_torso; 0.0]

	return q1
end

max_issues_repo_name: thowell/DirectMotionPlanning


text:
using ProgressMeter

function generate_trades(strat::Strategy; verbose::Bool=true)::Dict{String,TS}
    all_trades = Dict{String,TS}()
    verbose ? progress = Progress(length(strat.universe.assets), 1, "Generating Trades") : nothing
    for asset in strat.universe.assets
        verbose ? next!(progress) : nothing
        trades = TS(falses(size(strat.universe.data[asset],1), length(strat.rules)),
                    strat.universe.data[asset].index)
        local indicator_data = calculate(strat.indicator, strat.universe.data[asset])
        for (i,rule) in enumerate(strat.rules);
            trades[:,i] = rule.trigger.fun(indicator_data)
        end
        all_trades[asset] = trades
    end
    return all_trades
end

function generate_trades!(strat::Strategy; args...)::Nothing
    strat.backtest.trades = generate_trades(strat; args...)
    return nothing
end

#TODO: generalize this logic to incorporate order types
function backtest(strat::Strategy; px_trade::Symbol=:Open, px_close::Symbol=:Settle, verbose::Bool=true)::Dict{String,TS{Float64}}
    if isempty(strat.backtest.trades)
        generate_trades!(strat, verbose=verbose)
    end
    result = Dict{String,TS}()
    verbose ? progress = Progress(length(strat.universe.assets), 1, "Running Backtest") : nothing
    for asset in strat.universe.assets
        verbose ? next!(progress) : nothing
        trades = strat.backtest.trades[asset].values
        N = size(trades, 1)
        summary_ts = strat.universe.data[asset]
        trade_price = summary_ts[px_trade].values
        close_price = summary_ts[px_close].values
        pos = zeros(Float64, N)
        pnl = zeros(Float64, N)
        do_trade = false
        for t in 2:N
            for (i,rule) in enumerate(strat.rules)
                if trades[t-1,i] != 0
                    do_trade = true
                    order_side = rule.action in (long,buy) ? 1 : rule.action in (short,sell) ? -1 : 0
                    (order_qty,) = rule.args
                    pos[t] = order_qty * order_side
                    pnl[t] = pos[t] * (close_price[t] - trade_price[t])
                end
            end
            if !do_trade
                pos[t] = pos[t-1]
                pnl[t] = pos[t] * (close_price[t]-close_price[t-1])
            end
            do_trade = false
        end
        summary_ts = [summary_ts TS([pos pnl cumsum(pnl)], summary_ts.index, [:Pos,:PNL,:CumPNL])]
        result[asset] = summary_ts
    end
    return result
end

function backtest!(strat::Strategy; args...)::Nothing
    strat.backtest.backtest = backtest(strat; args...)
    return nothing
end


text:
modularity{V}(graph::AbstractGraph{V}, membership::Vector{Int}) = modularity(graph, membership, Float64[])

function modularity{V,T<:Real}(graph::AbstractGraph{V}, membership::Vector{Int}, weights::Vector{T})

    types = maximum(membership)
    no_of_edges = num_edges(graph)
    m = 0.0

    if length(membership) < num_vertices(graph)
        error("cannot calculate modularity, membership vector too short")
    end

    e = zeros(types)
    a = zeros(types)

    if !isempty(weights)
        if length(weights) < no_of_edges
            error("cannot caculate modularity, weight vector too short")
        end
        m = sum(weights)
        for i in edges(graph)
            w = weights[edge_index(i, graph)]
            w >= 0 || error("negative weight in weight vector")
            c1 = membership[vertex_index(source(graph, i), graph)]
            c2 = membership[vertex_index(target(graph, i), graph)]
            if c1 == c2
                e[c1] += 2w
            end
            a[c1] += w
            a[c2] += w
        end
    else
        m = no_of_edges
        for i in edges(graph)
            c1 = membership[vertex_index(source(i, graph), graph)]
            c2 = membership[vertex_index(target(i, graph), graph)]
            if c1 == c2
                e[c1] += 2
            end
            a[c1] += 1
            a[c2] += 1
        end
    end

    modularity_value = 0.0
    if m > 0
        for i=1:types
            tmp = a[i]/2/m
            modularity_value += e[i]/2/m
            modularity_value -= tmp*tmp
        end
    end
    modularity_value
end

max_stars_repo_name: afternone/LPA

max_issues_repo_name: afternone/LPA

max_forks_repo_name: afternone/LPA


text:
using StanSample, DiffEqBayes, OrdinaryDiffEq, ParameterizedFunctions,
      ModelingToolkit, RecursiveArrayTools, Distributions, Random, Test

# Uncomment for local testing only, make sure MCMCChains and StatsPlots are available
using MCMCChains, StatsPlots

#Random.seed!(123)
cd(@__DIR__)
isdir("tmp") && rm("tmp", recursive=true)

include("../../../stan_string.jl")
include("../../../stan_inference.jl")

println("\nOne parameter case 1\n")
f1 = @ode_def begin
  dx = a*x - x*y
  dy = -3y + x*y
end a

u0 = [1.0,1.0]
tspan = (0.0,10.0)
pars = [1.5]

prob1 = ODEProblem(f1,u0,tspan,pars)
sol = solve(prob1,Tsit5())
t = collect(range(1,stop=10,length=10))
randomized = VectorOfArray([(sol(t[i]) + .5randn(2)) for i in 1:length(t)])
data = convert(Array,randomized)
priors = [truncated(Normal(0.7,1),0.1,2)]

bayesian_result = stan_inference(prob1,t,data,priors;num_samples=2000, nchains=4,
  num_warmup=1000,likelihood=Normal)

if success(bayesian_result.return_code)
  sdf = read_summary(bayesian_result.model)
  sdf |> display

  p  = read_samples(bayesian_result.model, output_format=:particles)
  @test mean(p.theta1) ≈ 1.5 atol=5e-1

  # Uncomment for local chain inspection
  p1 = Vector{Plots.Plot{Plots.GRBackend}}(undef, length(pars))
  p2 = Vector{Plots.Plot{Plots.GRBackend}}(undef, length(pars))
  for i in 1:length(pars)
    p1[i] = plot(title="theta.$i")
    plot!(p1[i], p[Symbol("theta.$i")].particles)
    p2[i] = plot(title="theta.$i")
    density!(p2[i], p[Symbol("theta.$i")].particles)
  end
  plot(p1..., p2..., layout=(length(pars),2))
  savefig("$(@__DIR__)/case_1_chains.png")

  v1 = [mean(p[Symbol("u_hat.$i.1")]) for i in 1:10]
  v2 = [mean(p[Symbol("u_hat.$i.2")]) for i in 1:10]
  qa_prey = zeros(10, 3);
  qa_pred = zeros(10, 3);

  df = read_samples(bayesian_result.model, output_format=:dataframe)
  achn = Array(df);
  for i in 1:10
    qa_prey[i, :] = quantile(achn[:, 4+i], [0.055, 0.5, 0.945])
    qa_pred[i, :] = quantile(achn[:, 14+i], [0.055, 0.5, 0.945])
  end
  p1 = plot(1:10, v1; ribbon=(qa_prey[:, 1], qa_prey[:, 3]), color=:lightgrey, leg=false)
  title!("Prey u_hat 89% quantiles")
  plot!(v1, lab="prey", xlab="time", ylab="prey", color=:darkred)

  p2 = plot(1:10, v2; ribbon=(qa_pred[:, 1], qa_pred[:, 3]), color=:lightgrey, leg=false)
  title!("Preditors u_hat 89% quantiles")
  plot!(v2, lab="pred", xlab="time", ylab="preditors", color=:darkblue)

  plot(p1, p2, layout=(2,1))
  savefig("$(@__DIR__)/case_1_pred_prey.png")
end


text:
export kdownload

const KAGGLE_DOMAIN = "www.kaggle.com"
const KAGGLE_API = "https://www.kaggle.com/api/v1/datasets/download"

struct Auth
    username::String
    key::String
end

StructTypes.StructType(::Type{Auth}) = StructTypes.Struct()

function gen_auth_key()
    auth_file = joinpath(homedir(), ".kaggle", "kaggle.json")

    if isfile(auth_file)
        f = open(auth_file, "r")
        auth = JSON3.read(f, Auth)
        close(f)
    else
        auth = Auth(ENV["KAGGLE_USERNAME"], ENV["KAGGLE_KEY"])
    end

    auth_str = Base64.base64encode("$(auth.username):$(auth.key)")

    return "Basic $(auth_str)"
end

function gen_kaggle_url(dataset)
    return "$KAGGLE_API/$dataset"
end

function kaggle_url2dataset(url_or_dataset)
    if contains(url_or_dataset, KAGGLE_DOMAIN)
        user_name, dataset_name = match(Regex("$KAGGLE_DOMAIN/([^/]+)/([^/]+)"), url_or_dataset).captures
        dataset = "$user_name/$dataset_name"
    else
        dataset = url_or_dataset
    end

    @assert HTTP.request("HEAD", "https://$KAGGLE_DOMAIN/$dataset").status == 200

    return dataset
end

function kdownload(url_or_dataset, localdir)
    dataset = kaggle_url2dataset(url_or_dataset)

    url = gen_kaggle_url(dataset)
    filepath = joinpath(localdir, "$(replace(dataset, '/'=>'_')).zip")

    HTTP.open("GET", url, ["Authorization"=>gen_auth_key()]) do stream
        HTTP.startread(stream)
        eof(stream) && return

        total_bytes = tryparse(Int64, HTTP.header(stream, "Content-Length"))
        (total_bytes === nothing) && (total_bytes = NaN)
        println("Total: $total_bytes bytes")

        downloaded_bytes = progress = 0
        print("Downloaded:\e[s")
        Base.open(filepath, "w") do f
            while !eof(stream)
                downloaded_bytes += write(f, readavailable(stream))
                new_progress = 100downloaded_bytes ÷ total_bytes
                (new_progress > progress) && print("\e[u $downloaded_bytes bytes ($new_progress%)")
                progress = new_progress
            end
        end
    end

    return filepath
end


text:
using JLD2
using PyPlot
using Statistics

function make_ekp_plots(ekp_path::ST, param_names::Vector{ST}) where {ST <: AbstractString}
    n_param = length(param_names)

    # Load data
    data = load(joinpath(ekp_path, "ekp.jld2"))

    # Mean
    phi_m = mean(data["phi_params"], dims = 3)[:, :, 1]
    # Variance
    _ustd = std.(data["ekp_u"], dims = 2)
    n_iter = length(_ustd)
    n_param = length(_ustd[1])
    ustd = zeros((n_iter, n_param))
    for i in 1:n_iter
        ustd[i, :] = _ustd[i]
    end

    # plot parameter evolution
    fig, axs = subplots(nrows = n_param, sharex = true, figsize = (15, 4 * n_param))
    axs = n_param == 1 ? [axs] : axs
    x = 0:(n_iter - 1)
    for (i, ax) in enumerate(axs)
        ax.plot(x, phi_m[:, i])
        ax.fill_between(x, phi_m[:, i] .- 2 * ustd[:, i], phi_m[:, i] .+ 2 * ustd[:, i], alpha = 0.5)
        ax.set_ylabel(param_names[i])
    end

    axs[1].set_xlim(0, n_iter - 1)
    axs[1].set_title("Parameter evolution (mean ±2 std)")
    axs[end].set_xlabel("iteration")
    savefig(joinpath(ekp_path, "param_evol.png"))

    # Error plot
    x = 1:(n_iter - 1)
    err = data["ekp_err"]
    fig, ax = subplots()
    ax.plot(x, err)
    ax.set_xlim(1, n_iter - 1)
    ax.set_ylabel("Error")
    ax.set_xlabel("iteration")
    ax.set_title("Error evolution")
    savefig(joinpath(ekp_path, "error_evol.png"))
end


text:
#=
% Example: Determine the total amount of remineralized phosphate   %
=#
using Revise
using TMI, PyPlot, PyCall

# A, Alu, γ, inputfile = config(url,inputdir)
# ΔPO₄ = readtracer(inputfile,"qpo4")
TMIversion = "modern_90x45x33_GH10_GH12"
#TMIversion = "modern_180x90x33_GH10_GH12"
A, Alu, γ, TMIfile, L, B = config_from_nc(TMIversion)

PO₄ᴿ = regeneratedphosphate(TMIversion,Alu,γ)

# plot a section at 330 east longitude (i.e., 30 west)
lon_section = 330 # only works if exact
lims = 0:0.1:2.0
tlabel = "Regenerated Phosphate [μmol/kg]"
sectionplot(PO₄ᴿ,lon_section,lims,titlelabel = tlabel)


text:
using Revise
using Test

include("common.jl")

@testset "Polling" begin
    @test Revise.polling_files[]

    testdir = randtmp()
    mkdir(testdir)
    push!(LOAD_PATH, testdir)
    dn = joinpath(testdir, "Polling", "src")
    mkpath(dn)
    srcfile = joinpath(dn, "Polling.jl")
    joinpath(dn, "Polling.jl")
    open(srcfile, "w") do io
        println(io, """
__precompile__(false)

module Polling

f() = 1

end
""")
    end
    sleep(0.5) # let the source file age a bit
    @eval using Polling
    @test Polling.f() == 1
    # I'm not sure why 2 sleeps are better than one, but here it seems to make a difference
    sleep(0.1)
    sleep(0.1)
    open(srcfile, "w") do io
        println(io, """
__precompile__(false)

module Polling

f() = 2

end
""")
    end
    # Wait through the polling interval
    yry()
    sleep(7)
    yry()
    @test Polling.f() == 2

    rm(testdir; force=true, recursive=true)
end


text:
module Code

using StaticArrays, LabelledArrays, SparseArrays

export toexpr, Assignment, (←), Let, Func, DestructuredArgs, LiteralExpr,
       SetArray, MakeArray, MakeSparseArray, MakeTuple, AtIndex,
       SpawnFetch, Multithreaded

import ..SymbolicUtils
import SymbolicUtils: @matchable, Sym, Term, istree, operation, arguments

##== state management ==##

struct NameState
    symbolify::Dict{Any, Symbol}
end
NameState() = NameState(Dict{Any, Symbol}())
function union_symbolify!(n, ts)
    for t in ts
        n[t] = Symbol(string(t))
    end
end

struct LazyState
    ref::Ref{Any}
end
LazyState() = LazyState(Ref{Any}(nothing))

function Base.get(st::LazyState)
    s = getfield(st, :ref)[]
    s === nothing ? getfield(st, :ref)[] = NameState() : s
end

@inline Base.getproperty(st::LazyState, f::Symbol) = getproperty(get(st), f)

##========================##

"""
    toexpr(ex, [st,])

Convert a symbolic expression into an `Expr`, suitable to be passed into `eval`.

For example,

```julia
julia> @syms a b
(a, b)

julia> toexpr(a+b)
:((+)(a, b))

julia> toexpr(a+b) |> dump
Expr
  head: Symbol call
  args: Array{Any}((3,))
    1: + (function of type typeof(+))
    2: Symbol a
    3: Symbol b
```

Note that the function is an actual function object.

For more complex expressions, see other code-related combinators,

Namely `Assignment`, `Let`, `Func`, `SetArray`, `MakeArray`, `MakeSparseArray` and
`MakeTuple`.

To make your own type convertible to Expr using `toexpr` define `toexpr(x, st)` and
forward the state `st` in internal calls to `toexpr`. `st` is state used to know
when to leave something like `y(t)` as it is or when to make it `var"y(t)"`. E.g.
when `y(t)` is itself the argument of a function rather than `y`.

"""
toexpr(x) = toexpr(x, LazyState())
toexpr(s::Sym, st) = nameof(s)

@matchable struct Assignment
    lhs
    rhs
end

"""
    Assignment(lhs, rhs)

An assignment expression. Shorthand `lhs ← rhs` (`\\leftarrow`)
"""
Assignment

lhs(a::Assignment) = a.lhs
rhs(a::Assignment) = a.rhs

const (←) = Assignment

Base.convert(::Type{Assignment}, p::Pair) = Assignment(pair[1], pair[2])

toexpr(a::Assignment, st) = :($(toexpr(a.lhs, st)) = $(toexpr(a.rhs, st)))

function_to_expr(op, args, st) = nothing

function function_to_expr(::typeof(^), O, st)
    args = arguments(O)
    if length(args) == 2 && args[2] isa Number && args[2] < 0
        ex = args[1]
        if args[2] == -1
            return toexpr(Term{Any}(inv, [ex]), st)
        else
            return toexpr(Term{Any}(^, [Term{Any}(inv, [ex]), -args[2]]), st)
        end
    end
    return nothing
end

function function_to_expr(::typeof(SymbolicUtils.ifelse), O, st)
    args = arguments(O)
    :($(toexpr(args[1], st)) ? $(toexpr(args[2], st)) : $(toexpr(args[3], st)))
end

function_to_expr(::Sym, O, st) = get(st.symbolify, O, nothing)

function toexpr(O, st)
    !istree(O) && return O
    op = operation(O)
    expr′ = function_to_expr(op, O, st)
    if expr′ !== nothing
        return expr′
    else
        haskey(st.symbolify, O) && return st.symbolify[O]
        args = arguments(O)
        return Expr(:call, toexpr(op, st), map(x->toexpr(x, st), args)...)
    end
end

# Call elements of vector arguments by their name.
@matchable struct DestructuredArgs
    elems
    inds
    name
end

function DestructuredArgs(elems, name=gensym("arg"); inds=eachindex(elems))
    DestructuredArgs(elems, inds, name)
end

"""
    DestructuredArgs(elems, [name=gensym("arg")])

`elems` is a vector of symbols or call expressions.  When it appears as an argument in
`Func`, it expects a vector of the same length and de-structures the vector into its named
components. See example in `Func` for more information.

`name` is the name to be used for the argument in the generated function Expr.
"""
DestructuredArgs

toexpr(x::DestructuredArgs, st) = toexpr(x.name, st)
get_symbolify(args::DestructuredArgs) = ()
function get_symbolify(args::Union{AbstractArray, Tuple})
    cflatten(map(get_symbolify, args))
end
get_symbolify(x) = istree(x) ? (x,) : ()
cflatten(x) = Iterators.flatten(x) |> collect

function get_assignments(d::DestructuredArgs, st)
    name = toexpr(d, st)
    map(d.inds, d.elems) do i, a
        a ← (i isa Symbol ? :($name.$i) : :($name[$i]))
    end
end

@matchable struct Let
    pairs::Vector{Union{Assignment,DestructuredArgs}} # an iterator of pairs, ordered
    body
end

"""
    Let(assignments, body)

A Let block.

- `assignments` is a vector of `Assignment`s
- `body` is the body of the let block
"""
Let

function toexpr(l::Let, st)
    if all(x->x isa Assignment && !(x.lhs isa DestructuredArgs), l.pairs)
        dargs = l.pairs
    else
        dargs = map(l.pairs) do x
            if x isa DestructuredArgs
                get_assignments(x, st)
            elseif x isa Assignment && x.lhs isa DestructuredArgs
                [x.lhs.name ← x.rhs, get_assignments(x.lhs, st)...]
            else
                (x,)
            end
        end |> cflatten
        # expand and come back
        return toexpr(Let(dargs, l.body), st)
    end

    funkyargs = get_symbolify(map(lhs, dargs))
    union_symbolify!(st.symbolify, funkyargs)

    Expr(:let,
         Expr(:block, map(p->toexpr(p, st), dargs)...),
         toexpr(l.body, st))
end

@matchable struct Func
    args::Vector
    kwargs
    body
end

"""
    Func(args, kwargs, body)

A function.

- `args` is a vector of expressions
- `kwargs` is a vector of `Assignment`s
- `body` is the body of the function

**Special features in `args`**:

- args can contain `DestructuredArgs`
- call expressions

For example,

```julia

julia> @syms a b c t f(d) x(t) y(t) z(t)
(a, b, c, t, f(::Number)::Number, x(::Number)::Number, y(::Number)::Number, z(::Number)::Number)

julia> func = Func([a,x(t), DestructuredArgs([b, y(t)]), f], # args
                   [c ← 2, z(t) ← 42], # kwargs
                   f((a + b + c) / x(t) + y(t) + z(t)));

julia> toexpr(func)
:(function (a, var"x(t)", var"##arg#255", f; c = 2, var"z(t)" = 42)
      let b = var"##arg#255"[1], var"y(t)" = var"##arg#255"[2]
          f((+)(var"y(t)", var"z(t)", (*)((+)(a, b, c), (inv)(var"x(t)"))))
      end
  end)
```

- the second argument is a `DestructuredArgs`, in the `Expr` form, it is given a random name, and is expected to receive a vector or tuple of size 2 containing the values of `c` and `y(t)`. The let block that is automatically generated "destructures" these arguments.
- `x(t)` and `y(t)` have been replaced with `var"x(t)"` and `var"y(t)"` symbols throughout
the generated Expr. This makes sure that we are not actually calling the expressions `x(t)` or `y(t)` but instead passing the right values in place of the whole expression.
- `f` is also a function-like symbol, same as `x` and `y`, but since the `args` array contains `f` as itself rather than as say, `f(t)`, it does not become a `var"f(t)"`. The generated function expects a function of one argument to be passed in the position of `f`.

An example invocation of this function is:

```julia
julia> executable = eval(toexpr(func))
#10 (generic function with 1 method)

julia> exec(1, 2.0, [2,3.0], x->string(x); var"z(t)" = sqrt(42))
"11.98074069840786"
```
"""
Func

toexpr_kw(f, st) = Expr(:kw, toexpr(f, st).args...)

function toexpr(f::Func, st)
    funkyargs = get_symbolify(vcat(f.args, map(lhs, f.kwargs)))
    union_symbolify!(st.symbolify, funkyargs)
    dargs = filter(x->x isa DestructuredArgs, f.args)
    if !isempty(dargs)
        body = Let(dargs, f.body)
    else
        body = f.body
    end
    if isempty(f.kwargs)
        :(function ($(map(x->toexpr(x, st), f.args)...),)
              $(toexpr(body, st))
          end)
    else
        :(function ($(map(x->toexpr(x, st), f.args)...),;
                    $(map(x->toexpr_kw(x, st), f.kwargs)...))
              $(toexpr(body, st))
          end)
    end
end

@matchable struct SetArray
    inbounds::Bool
    arr
    elems  # Either iterator of Pairs or just an iterator
end

"""
    SetArray(inbounds, arr, elems)

An expression representing setting of elements of `arr`.

By default, every element of `elems` is copied over to `arr`,

but if `elems` contains `AtIndex(i, val)` objects, then `arr[i] = val`
is performed in its place.

`inbounds` is a boolean flag, `true` surrounds the resulting expression
in an `@inbounds`.
"""
SetArray

@matchable struct AtIndex
    i
    elem
end

function toexpr(a::AtIndex, st)
    toexpr(a.elem, st)
end

function toexpr(s::SetArray, st)
    ex = quote
        $([:($(toexpr(s.arr, st))[$(ex isa AtIndex ? ex.i : i)] = $(toexpr(ex, st)))
           for (i, ex) in enumerate(s.elems)]...)
        nothing
    end
    s.inbounds ? :(@inbounds $ex) : ex
end

@matchable struct MakeArray
    elems
    similarto # Must be either a reference to an array or a concrete type
    output_eltype
end

"""
    MakeArray(elems, similarto, [output_eltype=nothing])

An expression which constructs an array.

- `elems` is the output array
- `similarto` can either be a type, or some symbol that is an array whose type needs to
   be emulated. If `similarto` is a StaticArrays.SArray, then the output array is also
   created as an `SArray`, similarly, an `Array` will result in an `Array`, and a
   `LabelledArrays.SLArray` will result in a labelled static array.
- `output_eltype`: if set, then forces the element type of the output array to be this.
   by default, the output type is inferred automatically.

You can define:
```
@inline function create_array(A::Type{<:MyArray},a
                              ::Nothing, d::Val{dims}, elems...) where dims

# and

@inlline function create_array(::Type{<:MyArray}, T, ::Val{dims}, elems...) where dims
```

which creates an array of size `dims` using the elements `elems` and eltype `T`, to allow
`MakeArray` to create arrays similarto `MyArray`s.

"""
MakeArray

MakeArray(elems, similarto) = MakeArray(elems, similarto, nothing)

function toexpr(a::MakeArray, st)
    similarto = toexpr(a.similarto, st)
    T = similarto isa Type ? similarto : :(typeof($similarto))
    elT = a.output_eltype
    quote
        $create_array($T,
                     $elT,
                     Val{$(size(a.elems))}(),
                     $(map(x->toexpr(x, st), a.elems)...),)
    end
end

## Array
@inline function _create_array(::Type{<:Array}, T, ::Val{dims}, elems...) where dims
    arr = Array{T}(undef, dims)
    @assert prod(dims) == nfields(elems)
    @inbounds for i=1:prod(dims)
        arr[i] = elems[i]
    end
    arr
end

@inline function create_array(A::Type{<:Array}, T, d::Val, elems...)
    _create_array(A, T, d, elems...)
end

@inline function create_array(A::Type{<:Array}, ::Nothing, d::Val{dims}, elems...) where dims
    T = promote_type(map(typeof, elems)...)
    _create_array(A, T, d, elems...)
end

## Matrix

@inline function create_array(::Type{<:Matrix}, ::Nothing, ::Val{dims}, elems...) where dims
    Base.hvcat(dims, elems...)
end

@inline function create_array(::Type{<:Matrix}, T, ::Val{dims}, elems...) where dims
    Base.typed_hvcat(T, dims, elems...)
end

## SArray
@inline function create_array(::Type{<:SArray}, ::Nothing, ::Val{dims}, elems...) where dims
    SArray{Tuple{dims...}}(elems...)
end

@inline function create_array(::Type{<:SArray}, T, ::Val{dims}, elems...) where dims
    SArray{Tuple{dims...}, T}(elems...)
end

## LabelledArrays
@inline function create_array(A::Type{<:SLArray}, T, d::Val{dims}, elems...) where {dims}
    a = create_array(SArray, T, d, elems...)
    if nfields(dims) === ndims(A)
        similar_type(A, eltype(a), Size(dims))(a)
    else
        a
    end
end

@inline function create_array(A::Type{<:LArray}, T, d::Val{dims}, elems...) where {dims}
    data = create_array(Array, T, d, elems...)
    if nfields(dims) === ndims(A)
        LArray{eltype(data),nfields(dims),typeof(data),LabelledArrays.symnames(A)}(data)
    else
        data
    end
end

## We use a separate type for Sparse Arrays to sidestep the need for
## iszero to be defined on the expression type
@matchable struct MakeSparseArray{S<:AbstractSparseArray}
    array::S
end

"""
    MakeSpaseArray(array)

An expression which creates a `SparseMatrixCSC` or a `SparseVector`.

The generated expression contains the sparsity information of `array`,

it only creates the `nzval` field at run time.
"""
MakeSparseArray

function toexpr(a::MakeSparseArray{<:SparseMatrixCSC}, st)
    sp = a.array
    :(SparseMatrixCSC($(sp.m), $(sp.n),
                      $(copy(sp.colptr)), $(copy(sp.rowval)),
                      [$(toexpr.(sp.nzval, (st,))...)]))
end

function toexpr(a::MakeSparseArray{<:SparseVector}, st)
    sp = a.array
    :(SparseVector($(sp.n),
                   $(copy(sp.nzind)),
                   [$(toexpr.(sp.nzval, (st,))...)]))
end

@matchable struct MakeTuple
    elems
end

"""
    MakeTuple(tup)

Make a Tuple from a tuple of expressions.
"""
MakeTuple

function toexpr(a::MakeTuple, st)
    :(($(toexpr.(a.elems, (st,))...),))
end

struct Multithreaded end
"""
    SpawnFetch{ParallelType}(funcs [, args], reduce)

Run every expression in `funcs` in its own task, the expression
should be a `Func` object and is passed to `Threads.Task(f)`.
If `Func` takes arguments, then the arguments must be passed in as `args`--a vector of vector of arguments to each function in `funcs`. We don't use `@spawn` in order to support RuntimeGeneratedFunctions which disallow closures, instead we interpolate these functions or closures as smaller RuntimeGeneratedFunctions.

`reduce` function is used to combine the results of executing `exprs`. A SpawnFetch expression returns the reduced result.

Use `Symbolics.MultithreadedForm` ParallelType from the Symbolics.jl package to get the RuntimeGeneratedFunction version SpawnFetch.

`ParallelType` can be used to define more parallelism types
SymbolicUtils supports `Multithreaded` type. Which spawns
threaded tasks.
"""
struct SpawnFetch{Typ}
    exprs::Vector
    args::Union{Nothing, Vector}
    combine
end

(::Type{SpawnFetch{T}})(exprs, combine) where {T} = SpawnFetch{T}(exprs, nothing, combine)

function toexpr(p::SpawnFetch{Multithreaded}, st)
    args = isnothing(p.args) ? Iterators.repeated((), length(p.exprs)) : p.args
    spawns = map(p.exprs, args) do thunk, xs
        :(Base.Threads.@spawn $(toexpr(thunk, st))($(toexpr.(xs, (st,))...)))
    end
    quote
        $(toexpr(p.combine, st))(map(fetch, ($(spawns...),))...)
    end
end

"""
    LiteralExpr(ex)

Literally `ex`, an `Expr`. `toexpr` on `LiteralExpr` recursively calls
`toexpr` on any interpolated symbolic expressions.
"""
struct LiteralExpr
    ex
end

recurse_expr(ex::Expr, st) = Expr(ex.head, recurse_expr.(ex.args, (st,))...)
recurse_expr(ex, st) = toexpr(ex, st)

function toexpr(exp::LiteralExpr, st)
    recurse_expr(exp.ex, st)
end

end


text:
@testset "String" begin
    ex1 = intern("ex")
    @test ex1 == "ex"
    @test !addr_eq(ex1, "ex")
    ex2 = intern("ex")
    @test addr_eq(ex1, ex2)
    ex3 = intern(String, "ex")
    @test addr_eq(ex1, ex3)

    @testset "type inference" begin
        @test ex1 isa String
        @test ex2 isa String
        @inferred intern("ex")
        @inferred intern(String, "ex")
    end
end

@testset "SubString" begin
    aa1, bb1, cc1 = intern.(split("aa bb cc"))
    aa2, bb2, cc2 = intern.(split("aa bb cc"))
    aa3, bb3, cc3 = intern.(String, split("aa bb cc"))

    @test bb1 == "bb"
    @test !addr_eq(bb1, "bb")
    @test addr_eq(bb1, bb2)
    @test addr_eq(bb1, bb3)

    @testset "type inference" begin
        @test intern(split("aa bb cc")[1]) isa String
        @inferred intern(split("aa bb cc")[1])
        @test intern(String, split("aa bb cc")[1]) isa String
        @inferred intern(String, split("aa bb cc")[1])
    end
end

@testset "WeakRefString" begin
    using WeakRefStrings
    s1 = "ex"
    s2 = "ex"
    ex1 = @inferred intern(String, WeakRefString(unsafe_wrap(Vector{UInt8}, s1)))
    @test ex1=="ex"
    @test !addr_eq(ex1, s1)
    @test ex1 isa String
    ex2 = @inferred intern(String, WeakRefString(unsafe_wrap(Vector{UInt8}, s2)))
    @test addr_eq(ex1, ex2)
end

@testset "BigFloat" begin
    let
        pi1 = intern(BigFloat(π))
        @test pi1 == BigFloat(π)
        @test !addr_eq(pi1, BigFloat(π))

        pi2 = intern(BigFloat(π))
        @test addr_eq(pi1, pi2)

        @testset "type inference" begin
            @test pi1 isa BigFloat
            @inferred intern(BigFloat(π))
        end
    end
end


text:
module TestMacros

using YaoCompiler
using YaoCompiler.Intrinsics
using Test

@device function qft(n::Int)
    1 => H
    for k in 2:n
        @ctrl k 1 => shift(2π / 2^k)
    end

    if n > 1
        2:n => qft(n - 1)
    end
end

function YaoCompiler.routine_stub(::typeof(GenericRoutine{:qft}), n::Int)
    gate(H, 1)
    for k in 2:n
        ctrl(shift(2π / 2^k), k, 1)
    end

    if n > 1
        ctrl(qft(n - 1), 2:n)
    end
end

const qft = GenericRoutine{:qft}

@device function qft4()
    1 => H
    @ctrl 2 1 => shift(π / 2)
    @ctrl 3 1 => shift(π / 4)
    @ctrl 4 1 => shift(π / 8)

    2 => H
    @ctrl 3 2 => shift(π / 2)
    @ctrl 4 2 => shift(π / 4)

    3 => H
    @ctrl 4 3 => shift(π / 2)
    4 => H
end

@device function hadamard()
    1 => H
end

@testset "tracing" begin
    tape1 = YaoCompiler.@trace qft(4)
    tape2 = YaoCompiler.@trace qft4()
    @test tape1 == tape2
end

struct Foo
    a::Int
    b::Int
end

@device function (b::Foo)(theta)
    @ctrl b.a b.b => shift(theta)
    return theta
end

@testset "callable" begin
    m = Foo(1, 2)
    tape = YaoCompiler.@trace m(0.1)
    @test length(tape) == 1
    @test tape[1] ==
          Expr(:call, YaoCompiler.Semantic.ctrl, shift(0.1), Locations(2), CtrlLocations(1))
end

qasm"""OPENQASM 2.0;
include "qelib1.inc";
"""

@device function adder3()
    @ctrl 1 4 => X
    @ctrl 2 4 => X
    @ctrl 3 4 => X
    (1, 2, 5) => ccx()
    (1, 3, 5) => ccx()
    (2, 3, 5) => ccx()

    ctrl(X, 1, 4)
    gate(ccx(), (1, 2, 5))
end

@testset "YaoLang/#66" begin
    @code_qasm adder3()
end

end


text:
# Abstract, helpers ====================================================================== #
"""
    AbstractManifold{D, C}

A manifold of dimension `D` and codimension `C`.

Should be callable with `D` arguments in [0, 1] and should return a `SVector{D+C}`.
"""
abstract type AbstractManifold{D, C} end

dim(::AbstractManifold{D}) where D = D
codim(::AbstractManifold{<:Any, C}) where C = C

Base.rand(man::AbstractManifold) = rand(Random.GLOBAL_RNG, man)
Base.rand(rng::AbstractRNG, man::AbstractManifold{D}) where D =
    man(rand(Random.GLOBAL_RNG, D)...)
Base.rand(man::AbstractManifold, n::Integer) = rand(Random.GLOBAL_RNG, man, n)
Base.rand(rng::AbstractRNG, man::AbstractManifold{D}, n::Integer) where D =
    [rand(rng, man) for _ in 1:n]

# Point ================================================================================== #
"""
    PointSpace()

Single point, unit for `×`.
"""
struct PointSpace <: AbstractManifold{0, 1}
end

Base.show(io::IO, ::PointSpace) = print(io, "{⋆}")
(ps::PointSpace)() = SVector(0.0)

# Spheres ================================================================================ #
"""
    Sphere{D}

A D-Sphere.
"""
struct Sphere{D} <: AbstractManifold{D, 1}
end

Sphere(D) = Sphere{D}()
Base.show(io::IO, ::Sphere{D}) where D = print(io, "𝕊^$D")

function (::Sphere{D})(varargs::Vararg{T, D}) where {D, T}
    args = 2 .* varargs
    SVector((prod(sinpi.(args[1:i-1])) * cospi(args[i]) for i in 1:D)...,
            prod(sinpi.(args[1:end])))
end

Base.rand(rng::AbstractRNG, ::Sphere{D}) where D =
    normalize(SVector{D + 1, Float64}(randn(rng, D + 1)))

# Ball =================================================================================== #
"""
    Ball{D}

A D-Ball.
"""
struct Ball{D} <: AbstractManifold{D, 0}
end

Ball(D) = Ball{D}()
Base.show(io::IO, ::Ball{D}) where D = print(io, "𝔹^$D")

function (::Ball{D})(varargs::Vararg{T, D}) where {D, T}
    args = 2 .* varargs
    (args[1]/2)^(1/D) .*
        SVector((prod(sinpi.(args[2:i-1])) * cospi(args[i]) for i in 2:D)...,
                prod(sinpi.(args[2:end])))
end

function Base.rand(rng::AbstractRNG, ::Ball{D}) where D
    res = @SVector fill(2.0, D)
    while norm(res) > 1
        res = SVector{D, Float64}(2rand(rng, D) .- 1)
    end
    res
end

# Cube =================================================================================== #
"""
    Cube{D}

A D-Cube.
"""
struct Cube{D} <: AbstractManifold{D, 0}
end

Cube(D) = Cube{D}()
Base.show(io::IO, ::Cube{D}) where D = print(io, "𝕀", D ≠ 1 ? "^$D" : "")
(::Cube{D})(args::Vararg{T, D}) where {D, T} = SVector(args)

# ParametricManifold ===================================================================== #
struct ParametricManifold{D, C, F} <: AbstractManifold{D, C}
    fun::F
end

ParametricManifold{D, C}(f::F) where {D, C, F} =
    ParametricManifold{D, C, F}(f)
Base.show(io::IO, pm::ParametricManifold{D, C}) where {D, C} =
    print(io, "ParametricManifold{$D, $C}($(pm.fun))")
(pm::ParametricManifold{D, C})(args::Vararg{T, D}) where {D, C, T} =
    SVector{D+C, T}(pm.fun(args...))


text:
immutable FLSynapseParameter
end

@withkw type FLSynapse
    param::FLSynapseParameter = FLSynapseParameter()
    colptr::Vector{Int} # column pointer of sparse W
    I::Vector{Int}      # postsynaptic index of W
    W::Vector{Float}  # synaptic weight
    rI::Vector{Float} # postsynaptic rate
    rJ::Vector{Float} # presynaptic rate
    g::Vector{Float}  # postsynaptic conductance
    P::Vector{Float}  # <rᵢrⱼ>⁻¹
    q::Vector{Float}  # P * r
    u::Vector{Float} # force weight
    w::Vector{Float} # output weight
    f::Float = 0 # postsynaptic traget
    z::Float = 0.5randn()  # output z ≈ f
    records::Dict = Dict()
end

function FLSynapse(pre, post; σ = 1.5, p = 0.0, α = 1)
    w = σ * 1 / √(p * pre.N) * sprandn(post.N, pre.N, p)
    rowptr, colptr, I, J, index, W = dsparse(w)
    rI, rJ, g = post.r, pre.r, post.g
    P = α .* (I .== J)
    q = zeros(post.N)
    u = 2rand(post.N) - 1
    w = 1 / √post.N * (2rand(post.N) - 1)
    FLSynapse(;@symdict(colptr, I, W, rI, rJ, g, P, q, u, w)...)
end

@replace function forward!(c::FLSynapse, param::FLSynapseParameter)
    z = dot(w, rI)
    g .= z .* u
    fill!(q, zero(Float))
    @inbounds for j in 1:(length(colptr) - 1)
        rJj = rJ[j]
        for s = colptr[j]:(colptr[j+1] - 1)
            i = I[s]
            q[i] += P[s] * rJj
            g[i] += W[s] * rJj
        end
    end
end

@replace function plasticity!(c::FLSynapse, param::FLSynapseParameter, dt::Float, t::Float)
    C = 1 / (1 + dot(q, rI))
    BLAS.axpy!(C * (f - z), q, w)
    @inbounds for j in 1:(length(colptr) - 1)
        for s in colptr[j]:(colptr[j+1] - 1)
            P[s] += -C * q[I[s]] * q[j]
        end
    end
end


text:
function ApsideDetector(arg0::Orbit)
    return ApsideDetector((Orbit,), arg0)
end

function ApsideDetector(arg0::jdouble, arg1::Orbit)
    return ApsideDetector((jdouble, Orbit), arg0, arg1)
end

function g(obj::ApsideDetector, arg0::SpacecraftState)
    return jcall(obj, "g", jdouble, (SpacecraftState,), arg0)
end


text:
load_bits(s::String, skipbytes = 0) = load_bits(UInt, s, skipbytes)

function load_bits(::Type{T}, s::String, skipbytes = 0) where T<:Unsigned
    n = sizeof(s)
    # if n < skipbytes
    #     return zero(T)
    # else
        ns = (sizeof(T) - min(8, n - skipbytes))*8
        h = unsafe_load(Ptr{T}(pointer(s, skipbytes+1)))
        # h = unsafe_load(Ptr{T}(pointer(s)+skipbytes))
        h = h << ns
        h = h >> ns
        return h
    # end
end

load_bits(s::String, skipbytes = 0) = load_bits(UInt, s, skipbytes)

function load_bits(::Type{T}, s::String, skipbytes = 0) where T
    n = sizeof(s)
    if n < skipbytes
        return zero(T)
    elseif n - skipbytes >= sizeof(T)
        return unsafe_load(Ptr{T}(pointer(s, skipbytes+1)))
    else
        ns = (sizeof(T) - min(sizeof(T), n - skipbytes))*8
        h = unsafe_load(Ptr{T}(pointer(s, skipbytes+1)))
        # h = unsafe_load(Ptr{T}(pointer(s)+skipbytes))
        h = h << ns
        h = h >> ns
        return h
    end
end

function load_bits_fast(::Type{T}, s::String) where T
    n = sizeof(s)
    ns = (sizeof(T) - n)*8
    h = unsafe_load(Ptr{T}(pointer(s)))
    h = h << ns
    h = h >> ns
    return h
end

function load_bits_fast_ntoh(::Type{T}, s::String) where T
    n = sizeof(s)
    ns = (sizeof(T) - n)*8
    h = unsafe_load(Ptr{T}(pointer(s)))
    h = h << ns
    h = h >> ns
    return ntoh(h)
end

function roughhash(s::String)
    n = sizeof(s)
    if n >= 8
        return unsafe_load(Ptr{UInt64}(pointer(s)))
    else
        h = zero(UInt64)
        for i = 1:n
            @inbounds h = (h << 8) | codeunit(s, i)
        end
        return h
    end
end

function radixsort_ntoh!(svec::Vector{String}, skipbytes = 0, pointer_type = UInt)
    lens = reduce((x,y) -> max(x,sizeof(y)),0, svec)
    iters = ceil(lens*8/SortingAlgorithms.RADIX_SIZE)
    skipbytes = 0
    for i = iters:-1:1
        skipbytes += sizeof(pointer_type)
        x = ntoh.(unsafe_load.(Ptr{pointer_type}.(pointer.(svec) .+ skipbytes)))
        sorttwo!(x, svec)
    end
end


text:
fold(op::OP, xs; kw...) where {OP} = mapfold(identity, op, xs; kw...)

function mapfold(
    f::F,
    op::OP,
    xs::AbstractArray;
    basesize = cld(length(xs), Threads.nthreads()),
) where {F,OP}
    basesize = max(3, basesize)  # for length(left) + length(right) >= 4
    if length(xs) < basesize
        return mapfoldl(f, op, xs)
    end
    return _mapfold(f, op, xs, basesize)
end

function _mapfold(f::F, op::OP, xs, basesize) where {F,OP}
    if length(xs) <= basesize
        acc = @inbounds op(f(xs[begin]), f(xs[begin+1]))
        @simd for i in eachindex(xs)[3:end]
            acc = op(acc, f(@inbounds xs[i]))
        end
        return acc
    else
        left = @inbounds @view xs[begin:(end-begin+1)÷2]
        right = @inbounds @view xs[(end-begin+1)÷2+1:end]
        ref = Tapir.Output()
        local y
        Tapir.@sync begin
            Tapir.@spawn ref[] = _mapfold(f, op, right, basesize)
            y = _mapfold(f, op, left, basesize)
        end
        return op(y, ref[])
    end
end

function append!!(a, b)
    ys::Vector = a isa Vector ? a : collect(a)
    if eltype(b) <: eltype(ys)
        zs = append!(ys, b)
    else
        zs = similar(ys, promote_type(eltype(ys), eltype(b)), (length(ys) + length(b)))
        copyto!(zs, 1, ys, 1, length(ys))
        zs[length(ys)+1:end] .= b
    end
    return zs
end

function tmap(f::F, xs; kw...) where {F}
    ys = mapfold(tuple ∘ f, append!!, xs; kw...)
    if ys isa Tuple
        return collect(ys)
    else
        return ys
    end
end

function tforeach(f::F, xs; kw...) where {F}
    mapfold(f, (_, _) -> nothing, xs; kw...)
    return
end

sum(f::F, xs) where {F} = mapfold(f, +, xs)
sum(xs) = sum(identity, xs)
mean(f::F, xs) where {F} = sum(f, xs) / length(xs)
mean(xs) = mean(identity, xs)


text:
abstract type MHV1RuleModel <: MixingRule end

struct MHV1Rule{γ} <: MHV1RuleModel
    components::Array{String,1}
    activity::γ
    references::Array{String,1}
end

@registermodel MHV1Rule
export MHV1Rule
function MHV1Rule(components::Vector{String}; activity = Wilson, userlocations::Vector{String}=String[],activity_userlocations::Vector{String}=String[], verbose::Bool=false)
    init_activity = activity(components;userlocations = activity_userlocations,verbose)
    references = ["10.1016/0378-3812(90)85053-D"]
    model = MHV1Rule(components, init_activity,references)
    return model
end

MHV1q(::PRModel) = 0.53
MHV1q(::RKModel) = 0.593
function mixing_rule(model::Union{RKModel,PRModel},V,T,z,mixing_model::MHV1RuleModel,α,a,b,c)
    n = sum(z)
    invn = (one(n)/n)
    invn2 = invn^2
    g_E = excess_gibbs_free_energy(mixing_model.activity,1e5,T,z) / n
    b̄ = dot(z,Symmetric(b),z) * invn2
    c̄ = dot(z,c)/n
    q = MHV1q(model)
    Σlogb = sum(z[i]*log(b̄/b[i,i]) for i ∈ @comps)*invn
    Σab = invn*sum(z[i]*a[i,i]*α[i]/b[i,i]/(R̄*T) for i ∈ @comps)
    ā = b̄*R̄*T*(Σab-1/q*(g_E/(R̄*T)+Σlogb))
    return ā,b̄,c̄
end


text:
abstract type JMLCompilerError <: Exception end
struct Positioned <: JMLCompilerError
    lineno :: Int
    colno  :: Int
    err    :: JMLCompilerError
end

struct NameUnresolved <: JMLCompilerError
    name :: String
end

struct SimpleMessage <: JMLCompilerError
    msg :: String
end

struct ModulePathNotFound <: JMLCompilerError
    modname :: String
end

struct ComposeExc <: JMLCompilerError
    first :: JMLCompilerError
    next  :: JMLCompilerError
end

print_exc(io, a :: Positioned) =
     begin
        println(io, "line $(a.lineno), column $(a.colno)")
        print_exc(io, a.err)
     end
print_exc(io, a :: NameUnresolved) =
     begin
        println(io, "Name $(a.name) not found")
     end

print_exc(io, a :: ModulePathNotFound) =
     begin
        println(io, "Path of module $(a.modname) didn't exist in current JMLPATH")
     end
print_exc(io, a :: SimpleMessage) =
     begin
        println(io, a.msg)
     end
print_exc(io, a :: ComposeExc) =
     begin
        print_exc(io, a.first)
        print_exc(io, a.next)
     end

Base.show(io::IO, a::JMLCompilerError) = print_exc(io, a)

max_forks_repo_name: thautwarm/Rupy


text:
# SPDX-License-Identifier: X11
# 2020-11-15
# ATM Queue (5pt, 8pt)

function main()
  T = parse(Int, readline())
  for i ∈ 1:T
    print("Case #$i: ")
    solve()
  end
end

function solve()
  N, X = parse.(Int, (split ∘ readline)())
  A = (split ∘ readline)()
  K = Vector{Tuple{Int, Int, Int}}()
  for (i, s) ∈ enumerate(A)
    x = parse(Int, s)
    push!(K, (i, x, ceil(Int, x/X)))
  end

  sort!(K; alg=MergeSort, by = x -> x[3])

  for (i, x, k) ∈ K[1:end - 1]
    print("$i ")
  end
  println(K[end][1])
end

main()

max_stars_repo_name: jthvai/competitive-programming

max_issues_repo_name: jthvai/competitive-programming

max_forks_repo_name: jthvai/competitive-programming


text:
using Luxor

width  = 1000;
height = 1000;
demo = Movie(width, height, "test")
global x = rand(1:width);
global y = rand(1:height);

points = [];

t = Turtle(
0,#X
0,#y
true,#pendown
0,#orientation
(1., 0., 0.)#pencolor, rgb, r = 1, so it's red
);

function backdrop(scene, framenumber)
    background("yellow")
end

function frame(scene, framenumber)
    sethue("black")

    #str = string("frame $framenumber of $(scene.framerange.stop) x=$(x) y=$(y)  points=$(points)")
    #text(str, Point(O.x, O.y-190), halign=:center)

    for l in 1:150

        v = rand(1:3);
        if v == 1
            global x /= 2;
            global y /= 2;
        elseif v == 2
            global x = width/2 + (width/2 - x)/2;
            global y = height - (height - y)/2;
        else
            global x = width - (width - x)/2;
            global y = y / 2;
        end

        push!(points, (x,y));
    end

    for pt in points
        Reposition(t, pt[1]-width/2, height-pt[2]);

        Circle(t, 3);
    end
end

animate(demo, [
    Scene(demo, backdrop, 0:20),
    Scene(demo, frame, 0:20, easingfunction=easeinoutcubic)
    ],
    creategif=false)

max_stars_repo_name: tomaszkot/SierpinskiTriangle

max_issues_repo_name: tomaszkot/SierpinskiTriangle

max_forks_repo_name: tomaszkot/SierpinskiTriangle


text:
mutable struct RAPID <: Sampler
    threshold_strat::ThresholdStrategy
    threshold::Union{Real, Nothing}
    kde_cache::Union{KDECache, Nothing}
    gamma::Union{Symbol, Real}
    function RAPID(threshold_strat::ThresholdStrategy, kde_cache::Union{KDECache, Nothing}, gamma::Union{Symbol, Real})
        return new(threshold_strat, nothing, kde_cache, gamma)
    end
end

function RAPID(threshold_strat::ThresholdStrategy, kde_cache::KDECache)
    return RAPID(threshold_strat, kde_cache, :precomputed)
end

function RAPID(threshold_strat::ThresholdStrategy, gamma::Union{Symbol, Real})
    return RAPID(threshold_strat, nothing, gamma)
end

function sample(sampler::RAPID, data::DataSet, labels::Vector{Symbol})::BitArray{1}
    if sampler.kde_cache === nothing || sampler.kde_cache.data != data
        sampler.kde_cache = KDECache(data, sampler.gamma)
    end
    sampler.threshold, inlier_mask, outlier_mask = split_masks(sampler.threshold_strat,
                                                               sampler.kde_cache,
                                                               labels)

    remove_order = zeros(Int, size(data, 2))
    sample_mask = copy(inlier_mask)
    d_x = kde(sampler.kde_cache, sample_mask, trues(size(data, 2)))
    i = 1
    num_iter = count(sample_mask) - 2
    while i < num_iter
        idx = findall(sample_mask)[findmax(d_x[sample_mask])[2]]
        sample_mask[idx] = false
        d_x -= sampler.kde_cache.K[:, idx]
        θ = minimum(d_x[sample_mask])
        if !all(d_x[inlier_mask] .>= θ)
            break
        end
        remove_order[i] = idx
        i += 1
    end
    sample_mask = inlier_mask
    sample_mask[remove_order[1:i-1]] .= false
    sampler.threshold = any(sample_mask) ? minimum(kde(sampler.kde_cache, sample_mask, sample_mask)) : 0
    return sample_mask
end


text:
function spatial_arrangement(model; atol = 1e-6, outerCell = false)
    split_model = CAGD.facesplitting(model, atol = atol)

    # Topological Invariants Check #3: [Δ_1] × [Δ_0] = 0
    @assert isempty(dropzeros(split_model.T[2]*split_model.T[1]).nzval) "Invariant 3 Error"

    congr_model = CAGD.mergeModelVertices(split_model, err=atol, signed_merge=true)

    # Topological Invariants Check #4: [δ_1] × [δ_0] = 0
    @assert isempty(dropzeros(congr_model.T[2]*congr_model.T[1]).nzval) "Invariant 4 Error"

    #bico_models = CAGD.split_biconnected_models(congr_model)
    FCcycles, bicon_comps = CAGD.tgw(congr_model, 3, atol=atol)

    externals = [bc[CAGD.getExternalBoundary(congr_model, 3, FCcycles, bc)] for bc in bicon_comps]
    internals = [setdiff(bicon_comps[i], externals[i]) for i = 1 : length(externals)]
    # containment_graph = CAGD.buildContainmentGraph(gift_model, 3, externals, internals)

    # ?? REARRANGE OUTER CELLS

    # assembled_model = CAGD.assembleModel(model, 3, externals, containment_graph, outerCell)

    # WILL BE REMOVED
    # Sums all of the component external in a single component
    assembled_model = deepcopy(congr_model)

    FC = FCcycles[:, internals...]
    if outerCell
        JK = ones(Int8, length(externals))
        FC = [FCcycles * SparseArrays.sparse(externals, JK, JK, size(FCcycles, 2), 1) FC]
    end
    CAGD.addModelCells!(assembled_model, 3, convert(Lar.ChainOp, FC'))

    # Topological Invariants Check #7: [δ_1] × [δ_0] = 0
    @assert isempty(dropzeros(assembled_model.T[2]*assembled_model.T[1]).nzval) "Invariant 7 Error"
    # Topological Invariants Check #8: [δ_2] × [δ_1] = 0
    @assert isempty(dropzeros(assembled_model.T[3]*assembled_model.T[2]).nzval) "Invariant 8 Error"

    return assembled_model

end


text:
export ortho_latin, check_ortho

"""
`A,B = ortho_latin(n)` returns a pair of orthogonal `n`-by-`n`
Latin squares.

`A,B = ortho_latin(n,true)` returns a pair of orthogonal Latin squares
that are transposes of each other.

`A,B = ortho_latin(n,r,s)` builds the Latin squares `latin(n,r)`
and `latin(n,s)` and, if they are orthogonal, returns them as the
answer. (Otherwise, throws an error.) See: `find_ortho_parameters`.
"""
function ortho_latin(n::Int, self::Bool=false)
    if !self
        try
            r,s = find_ortho_parameters(n)
            return ortho_latin(n,r,s)
        catch
        end
    end
    println("No quick solution. Using integer programming.")

    return ortho_latin_IP(n,self)
end

function ortho_latin_IP(n::Int,self::Bool=false)
    MOD = Model(get_solver())
    # Z[i,j,k,l] is an indicator that there is a k in A[i,j] and
    # an l in B[i,j]
    @variable(MOD,Z[1:n,1:n,1:n,1:n], Bin)

    # one entry per cell constraint
    for i=1:n
        for j=1:n
            @constraint(MOD, sum(Z[i,j,k,l] for k=1:n for l=1:n) == 1)
        end
    end

    # Top row 11 22 33 ... nn
    if !self
        for i=1:n
            @constraint(MOD, Z[1,i,i,i]==1)  # A[1,i] = B[1,i] = i
        end
    end

    # orthogonality constraint
    for k=1:n
        for l=1:n
            @constraint(MOD, sum(Z[i,j,k,l] for i=1:n for j=1:n) == 1)
        end
    end

    # Row constraints

    for i=1:n
        for k=1:n
            @constraint(MOD, sum(Z[i,j,k,l] for j=1:n for l=1:n) == 1)
        end
    end

    for i=1:n
        for l=1:n
            @constraint(MOD, sum(Z[i,j,k,l] for j=1:n for k=1:n) == 1)
        end
    end

    # Col constraints
    for j=1:n
        for k=1:n
            @constraint(MOD, sum(Z[i,j,k,l] for i=1:n for l=1:n) == 1)
        end
    end

    for j=1:n
        for l=1:n
            @constraint(MOD, sum(Z[i,j,k,l] for i=1:n for k=1:n) == 1)
        end
    end

    if self   # force A^T==B
        for i=1:n
            for j=1:n
                for k=1:n
                    for l=1:n
                        @constraint(MOD, Z[i,j,k,l] == Z[j,i,l,k])
                    end
                end
            end
        end
    end

    optimize!(MOD)
    status = Int(termination_status(MOD))

    if status != 1
        error("No pair of orthogonal Latin squares of order $n can be found.")
    end

    ZZ = value.(Z)
    A = zeros(Int,n,n)
    B = zeros(Int,n,n)

    for i=1:n
        for j=1:n
            for k=1:n
                for l=1:n
                    if ZZ[i,j,k,l]>0
                        A[i,j] = k
                        B[i,j] = l
                    end
                end
            end
        end
    end

    return A,B
end

function ortho_latin(n::Int, r::Int, s::Int)
    A = latin(n,r)
    B = latin(n,s)
    @assert check_ortho(A,B) "Parameters n=$n, r=$r, and s=$s do not generate a pair of orthogonal Latin squares"
    return A,B
end

"""
`find_ortho_parameters(n)` tries to find parameters `r` and `s`
so that `ortho_latin(n,r,s)` will succeed. Returns `(r,s)` if
successful or throws an error if not.
"""
function find_ortho_parameters(n::Int)
    for r=1:n-1
        for s=1:n-1
            if gcd(n,r)==1 && gcd(n,s)==1 && gcd(n,r-s)==1
                return r,s
            end
        end
    end
    error("No parameters for n=$n found")
end

"""
`check_ortho(A,B)` checks that matrices `A` and `B` are a pair of
orthogonal Latin squares.
"""
function check_ortho(A::Matrix{Int},B::Matrix{Int})::Bool
    if size(A) != size(B)
        return false
    end
    if !check_latin(A) || !check_latin(B)
        return false
    end
    n,r = size(A)

    vals = unique((n+1)*A + B)
    return length(vals) == n*n
end


text:
const OrbitProjection = Equilibrium.Boundary

const in_orbit = in_boundary

function orbit_projection(M::AbstractEquilibrium, gcp; dx = 0.01, dy =0.01)
    hc = HamiltonianCoordinate(M,gcp)
    rlims, zlims = limits(M)

    x = range(rlims...,step=dx)
    y = range(zlims...,step=dy)

    mu = zeros(length(x), length(y))
    for ix in eachindex(x), iy in eachindex(y)
        B = norm(Bfield(M, x[ix], y[iy]))
        ir, iphi, iz = cylindrical_cocos_indices(cocos(M))
        g = poloidal_current(M, M(x[ix], y[iy])) # RBphi
        energy = 1e3*hc.energy*e0
        mu[ix,iy] = max(energy/B - (B/(2*hc.m))*((hc.p_phi - hc.q*e0*M(x[ix],y[iy]))/g)^2,0.0)
    end

    cl = Contour.contour(x,y,mu,hc.mu)
    ls = Contour.lines(cl)
    oproj = OrbitProjection([(gcp.r, gcp.z)])
    if length(ls) == 0
        return oproj
    end
    omin = Inf
    for (i,l) in enumerate(ls)
        r, z = Contour.coordinates(l)
        omin_tmp = minimum((r .- gcp.r).^2 .+ (z .- gcp.z).^2)
        if omin_tmp < omin
            omin = omin_tmp
            oproj = OrbitProjection(collect(zip(r,z)))
        end
    end

    return oproj
end


text:
# Generate the data sets used in the comparisons of mix-SQP and the
# first-order methods (EM and projected gradient descent).

# SCRIPT PARAMETERS
# -----------------
n        = 20000;
m        = 20;
filename = "simdata-n=20000-m=20.csv";

# SET UP ENVIRONMENT
# ------------------
using Random
using Distributions
using DelimitedFiles
include("../code/datasim.jl");
include("../code/likelihood.jl");

# Initialize the sequence of pseudorandom numbers.
Random.seed!(1);

# SIMULATE DATA
# -------------
# Generate the n x m conditional likelihood matrix.
z  = normtmixdatasim(n);
sd = autoselectmixsd(z,nv = m);
L  = normlikmatrix(z,sd = sd);

# Save the matrix to a CSV file.
writedlm(filename,L,',');

max_stars_repo_name: stephenslab/mixsqp-paper

max_issues_repo_name: stephenslab/mixsqp-paper

max_forks_repo_name: stephenslab/mixsqp-paper


text:
sparse_options = (; conic_solver = sparse_hypatia)

insts = OrderedDict()

insts["test"] = [
    ((2, 2, false, true), sparse_options),
    ((2, 2, false, false), sparse_options),
    ((2, 2, true, true), sparse_options),
    ((2, 2, true, false), sparse_options),
    ((3, 3, false, true), sparse_options),
    ((3, 3, false, false), sparse_options),
    ((3, 3, true, true), sparse_options),
    ((10, 3, false, true), sparse_options),
    ((5, 2, true, true), sparse_options),
]

function twostagestochastic_insts(use_nat::Bool, p_max::Int)
    return [((3, 2^p, false, use_nat), sparse_options) for p in vcat(3, 3:p_max)]
end

insts["nat"] = twostagestochastic_insts(true, 10)
insts["ext"] = twostagestochastic_insts(false, 5)

return (TwoStageStochastic, insts)


text:
# slt_read.jl
#
# Jeremy Rüffer
# Thünen Institut
# Institut für Agrarklimaschutz
# Junior Research Group NITROSPHERE
# Julia 0.7
# 09.12.2016
# Last Edit: 18.04.2019

"""# slt_read

Load a single SLT file

`slt_read(F,a_inputs,sample_rate)`\n
* **F**::String = File name, including path
* **a_inputs**::Int = Number of analog channels in the SLT file
* **sample_rate**::Int or String = Sampling rate of the sonic"""
function slt_read(F::String,AnalogIn::Int,freq::String)
	return slt_read(F,AnallogIn,parse(freq))
end # slt_read(F::String,AnalogIn::Int,freq::String)

function slt_read(F::String,a_inputs::Int,sample_rate::Int)
	# Arrays to output
	t = [] # Time
	d = [] # Data

	# Parsing information for SLT files can be found in the EddySoft PDF on page 15 (actual page 23).

	fid = open(F,"r")
	try
		# Header
		bpr = Int(read!(fid,Array{Int8}(undef,1))[1]) # Bytes Per Record
		eddymeasver = read!(fid,Array{Int8}(undef,1))
		dom = Int(read!(fid,Array{Int8}(undef,1))[1]) # Day of Month
		m = Int(read!(fid,Array{Int8}(undef,1))[1]) # Month number
		yr = 100*Int(read!(fid,Array{Int8}(undef,1))[1]) + Int(read!(fid,Array{Int8}(undef,1))[1]) # Year
		h = Int(read!(fid,Array{Int8}(undef,1))[1]) # Hour
		minut = Int(read!(fid,Array{Int8}(undef,1))[1]) # Minute
		t0 = DateTime(yr,m,dom,h,minut,0) # Initial file time

		fs = stat(F).size # File size in bytes
		l = (fs - 8 - 2*a_inputs)/(8 + 2*a_inputs) # Number of rows of data, trusting the user inputs
		if l - floor(l) != 0
			warn(F * " either has a different column count than " * string(a_input + 4) * " or was terminated early.")
		end
		l = Int(floor(l))

		# Bit Masks and Channels
		bm = read!(fid,Array{Int8}(undef,2*a_inputs))
		ch = bm[2:2:end] # Channels
		bm = bm[1:2:end] # Bit masks

		## Preallocate Output Array
		t_offset = fill!(Array{Dates.Millisecond}(undef,floor(l)),Dates.Millisecond(1000/sample_rate)) # dt between every sample
		t_offset[1] = Dates.Millisecond(0) # Correction, the first sample shouldn't have an offset
		t_offset = cumsum(t_offset) # Time offset from the start for every sample
		t = DateTime(yr,m,dom,h,minut,0) + t_offset # Time
		d = NaN.*Array{Float64}(undef,(Int(l),4 + a_inputs)) # Data Array

		for i=1:1:l
			d[i,1] = Float64(read!(fid,Array{Int16}(undef,1))[1])/100 # u
			d[i,2] = Float64(read!(fid,Array{Int16}(undef,1))[1])/100 # v
			d[i,3] = Float64(read!(fid,Array{Int16}(undef,1))[1])/100 # w
			d[i,4] = Float64(read!(fid,Array{Int16}(undef,1))[1])/50 # c, # Speed of Sound

			for j=1:1:a_inputs
				V = Float64(read!(fid,Array{Int16}(undef,1))[1]) # V1
				if bm[j] & 2^(1-1) > 0 # If the first bit is high
					# If the bit mask is high, use the following formula to convert the binary value to mV
					V = (V[1] + 25000)/10 # Binary to mV
				end
				d[i,j+4] = V
			end
		end
	catch
		println("Error loading file: " * F)
	end
	close(fid)

	return t,d
end # End of slt_read(F::String,a_inputs::Int,sample_rate::Int)

max_stars_repo_name: JeremyRueffer/ClimateDataIO

max_issues_repo_name: JeremyRueffer/ClimateDataIO