text:
"""
    AbstractCircuitGate

Abstract unitary quantum circuit gate.
"""
abstract type AbstractCircuitGate end

"""
    CircuitGate{M,G} <: AbstractCircuitGate

Unitary quantum circuit gate. `M` is the number of wires affected by the CircuitGate, and `G` is the basic gate used to construct the CircuitGate.
"""
struct CircuitGate{M,G} <: AbstractCircuitGate
    "ordered wire indices which this gate acts on"
    iwire::NTuple{M,Int64}
    "actual gate"
    gate::G

    @doc """
        CircuitGate{M,G}(iwire::NTuple{M,<:Integer}, gate::G) where {M,G}

    Creates a `CircuitGate{M,G}` object. `M` is the number of wires affected by the CircuitGate, and `G` is the basic gate used to construct the CircuitGate.
    """
    function CircuitGate{M,G}(iwire::NTuple{M,<:Integer}, gate::G) where {M,G <: AbstractGate}
        M == num_wires(gate) || error("$G affects $(num_wires(gate)) wires but $M wires, $iwire, were passed.")
        M ≥ 1 || error("Need at least one wire to act on.")
        length(unique(iwire)) == M || error("Wire indices must be unique.")
        minimum(iwire) ≥ 1 || error("Wire index cannot be smaller than 1.")

        new{M,G}(iwire, gate)
    end
end

"""
    CircuitGate(iwire::NTuple{M,<:Integer}, gate::G)

Create a `CircuitGate{M,G}`.
"""
function CircuitGate(iwire::NTuple{M,<:Integer}, gate::G) where {M,G <: AbstractGate}
    CircuitGate{M,G}(iwire, gate)
end

"""
    req_wires(cg::CircuitGate{M,G})

Minimum number of required qubit wires in a circuit to host the circuit gate `cg`.
```jldoctest
julia> cg = circuit_gate(3, X, (1,2,5));
julia> req_wires(cg)
5
```
"""
req_wires(cg::CircuitGate{M,G}) where {M,G} = maximum(cg.iwire)

"""
    Base.isapprox(cg1::CircuitGate{M,G}, cg2::CircuitGate{M,G})

Compares two circuit gates of basic type `G`.
"""
function Base.isapprox(cg1::CircuitGate{M,G}, cg2::CircuitGate{M,G}) where {M,G}
# for parametric gates, return true only if parameters are approximately equal
    if cg1.iwire != cg2.iwire
        return false
    end
    return cg1.gate ≈ cg2.gate
end

Base.isapprox(cg1::CircuitGate{M,G}, cg2::CircuitGate{M,H}) where {M,G,H} = false

LinearAlgebra.ishermitian(cg::CircuitGate) = LinearAlgebra.ishermitian(cg.gate)

"""
    sparse_matrix(cg::CircuitGate{M,G}, N::Integer=0) where {M,G <: AbstractGate}

returns matrix representation of a [CircuitGate](@ref) object that can applied to a state vector of `N` qubits. `N` can be

# Examples

```jldoctest
julia> cg = circuit_gate(2, Y, 1);
julia> sparse_matrix(cg)
4×4 SparseArrays.SparseMatrixCSC{Complex{Float64},Int64} with 4 stored entries:
  [1, 1]  =  1.0+0.0im
  [4, 2]  =  0.0+1.0im
  [3, 3]  =  1.0+0.0im
  [2, 4]  =  0.0-1.0im
```
"""
function sparse_matrix(cg::CircuitGate{M,G}, N::Integer=0) where {M,G <: AbstractGate}
    # convert to array
    iwire = collect(cg.iwire)

    if N == 0
        N = req_wires(cg)
    else
        N >= maximum(iwire) || error("Circuit size `$N` too small for CircuitGate applied to wires `$iwire`.")
    end

    # TODO: handle sparse matrices efficiently
    gmat = sparse_matrix(cg.gate)

    distribute_to_wires(gmat, iwire, N, M)
end

"""
    sparse_matrix(cgs::Vector{<:CircuitGate}, N::Integer=0)

returns matrix representation of a `Vector{<:CircuitGate}` object that can applied to a state vector of `N` qubits.

# Examples

```jldoctest
julia> cgs = CircuitGate[
                circuit_gate(2, Y, 1),
                circuit_gate(2, Z),
                circuit_gate(1, HadamardGate())
                ];
julia> sparse_matrix(cgs)
4×4 SparseArrays.SparseMatrixCSC{Complex{Float64},Int64} with 8 stored entries:
  [1, 1]  =  0.707107+0.0im
  [2, 1]  =  0.707107+0.0im
  [3, 2]  =  0.0-0.707107im
  [4, 2]  =  0.0+0.707107im
  [3, 3]  =  -0.707107+0.0im
  [4, 3]  =  -0.707107+0.0im
  [1, 4]  =  0.0-0.707107im
  [2, 4]  =  0.0+0.707107im
```
"""
function sparse_matrix(cgs::Vector{<:CircuitGate}, N::Integer=0)
    Nmin = maximum(req_wires.(cgs))
    if N == 0
        N = Nmin
    else
        N >= Nmin || error("Circuit size `$N` too small; vector of CircuitGate requires $Nmin wires.")
    end

    gmat = sparse(one(ComplexF64)*I, 2^N, 2^N)
    for cg in cgs
        iwire = collect(cg.iwire)
        gmat = distribute_to_wires(sparse_matrix(cg.gate), iwire, N, num_wires(cg.gate)) * gmat
    end

    return gmat
end

function distribute_to_wires(gmat::SparseMatrixCSC{Complex{Float64},Int}, iwire::Vector{Int}, N::Int, M::Int)

    # complementary wires
    iwcompl = setdiff(1:N, iwire)
    @assert length(iwire) + length(iwcompl) == N
    @assert size(gmat) == (2^M, 2^M)

    # Note: following the ordering convention of `kron` here, i.e.,
    # last qubit corresponds to fastest varying index
    strides = Int[2^j for j in 0:N-1]
    wstrides = strides[iwire]
    cstrides = strides[iwcompl]
    b = BitArray{1}(undef, M)

    values = Vector{ComplexF64}(undef, length(gmat.nzval) * 2^(N - M))
    nnz = length(gmat.nzval)
    rowind = Vector{Int}(undef, length(values))
    colind = Vector{Int}(undef, length(values))

    for j in 1:length(gmat.colptr)-1
        index = gmat.colptr[j]:gmat.colptr[j + 1] - 1
        colvalue = dot(binary_digits!(b, j - 1), wstrides) + 1
        for i in index
            @inbounds rowind[i] = dot(binary_digits!(b, gmat.rowval[i] - 1), wstrides) + 1
            @inbounds colind[i] = colvalue
        end
    end

    r = Base.view(rowind, 1:nnz)
    c = Base.view(colind, 1:nnz)
    @inbounds values[1:nnz] = gmat.nzval

    b = BitArray(undef, N - M)
    for k in 2:2^(N - M)
        koffset = dot(binary_digits!(b, k - 1), cstrides)
        @inbounds rowind[nnz*(k-1)+1:nnz*k] = r .+ koffset
        @inbounds colind[nnz*(k-1)+1:nnz*k] = c .+ koffset
        @inbounds values[nnz*(k-1)+1:nnz*k] = gmat.nzval
    end

    return sparse(rowind, colind, values, 2^N, 2^N)
end

"""
Base.adjoint(cg::CircuitGate{M,G})

Construct a `CircuitGate{M,H}` object where `H` is the adjoint of `AbstractGate` `G`
"""
function Base.adjoint(cg::CircuitGate{M,G}) where {M,G}
    adj_gate = adjoint(cg.gate)
    CircuitGate{M,typeof(adj_gate)}(cg.iwire, adj_gate)
end

function Base.adjoint(cg::Vector{<:CircuitGate})
    reverse(adjoint.(cg))
end

"""
    circuit_gate

helper function to construct a [`CircuitGate`](@ref) object from basic gate types.

```jldoctest
julia> circuit_gate(1, X)
CircuitGate{1,XGate}((1,), XGate())

julia> circuit_gate(1, X, 2)
CircuitGate{2,ControlledGate{XGate}}((1, 2), ControlledGate{XGate}(XGate(), 1))

julia> circuit_gate((1,2), SwapGate(), (3,4))
CircuitGate{4,ControlledGate{SwapGate}}((1, 2, 3, 4), ControlledGate{SwapGate}(SwapGate(), 2))
```
"""
function circuit_gate(iwire::Integer, gate::AbstractGate, control::Integer...)
    circuit_gate(iwire, gate, control)
end

function circuit_gate(iwire1::Integer, iwire2::Integer, gate::AbstractGate, control::NTuple{M,Integer}=NTuple{0,Integer}()) where {M}
    if isempty(control)
        return CircuitGate((iwire1, iwire2), gate)
    end
    C = length(control)
    return CircuitGate((iwire1, iwire2, control...), ControlledGate(gate, M))
end

function circuit_gate(iwire1::NTuple{L,Integer}, gate::AbstractGate, control::NTuple{M,Integer}=NTuple{0,Integer}()) where {L,M}
    if isempty(control)
        return CircuitGate(iwire1, gate)
    end
    return CircuitGate((iwire1..., control...), ControlledGate(gate, M))
end

function circuit_gate(iwire1::NTuple{M,Integer}, gate::AbstractGate, control::Integer...) where {M}
    circuit_gate(iwire1, gate, control)
end

function circuit_gate(iwire::Integer, gate::AbstractGate, control::NTuple{M,Integer}=NTuple{0,Integer}()) where {M}
    if isempty(control)
        return CircuitGate((iwire,), gate)
    end
    return CircuitGate((iwire, control...), ControlledGate(gate, M))
end

circuit_gate(iwire1::Integer, iwire2::Integer, gate::AbstractGate, control::Integer...) = circuit_gate(iwire1, iwire2, gate, control)


text:
"""
    StatLearn(p, args...; rate=LearningRate())

Fit a model that is linear in the parameters.

The (offline) objective function that StatLearn approximately minimizes is

``(1/n) ∑ᵢ f(yᵢ, xᵢ'β) + ∑ⱼ λⱼ g(βⱼ),``

where ``fᵢ`` are loss functions of a single response and linear predictor, ``λⱼ``s are
nonnegative regularization parameters, and ``g`` is a penalty function.

# Arguments

- `loss = .5 * L2DistLoss()`
- `penalty = NoPenalty()`
- `algorithm = SGD()`
- `rate = LearningRate(.6)` (keyword arg)

# Example

    x = randn(1000, 5)
    y = x * range(-1, stop=1, length=5) + randn(1000)

    o = fit!(StatLearn(5, MSPI()), (x, y))
    coef(o)
"""
mutable struct StatLearn{A<:Algorithm, L<:Loss, P<:Penalty, W} <: OnlineStat{XY}
    β::Vector{Float64}
    λ::Vector{Float64}
    gx::Vector{Float64}
    loss::L
    penalty::P
    alg::A
    rate::W
    n::Int
end
function StatLearn(p::Int, args...; rate=LearningRate())
    λ, loss, pen, alg = zeros(p), .5*L2DistLoss(), PenaltyFunctions.NoPenalty(), SGD()
    for a in args
        a isa AbstractVector && (λ = a)
        a isa Float64        && (λ = fill(a, 1))
        a isa Loss           && (loss = a)
        a isa Penalty        && (pen = a)
        a isa Algorithm      && (alg = a)
    end
    init!(alg, p)
    StatLearn(zeros(p), λ, zeros(p), loss, pen, alg, rate, 0)
end

function Base.show(io::IO, o::StatLearn)
    print(io, "StatLearn: ")
    print(io, name(o.alg, false, false))
    print(io, " | mean(λ)=", mean(o.λ))
    print(io, " | ", o.loss)
    print(io, " | ", o.penalty)
    print(io, " | nobs=", nobs(o))
    print(io, " | nvars=", length(o.β))
end
coef(o::StatLearn) = value(o)

function gradient!(o::StatLearn, x, y)
    d_dη = LearnBase.deriv(o.loss, y, predict(o, x))
    for j in eachindex(o.gx)
        o.gx[j] = x[j] * d_dη
    end
end
function _fit!(o::StatLearn{<:SGAlgorithm}, xy)
    x, y = xy
    o.n += 1
    gradient!(o, x, y)
    update!(o.alg, o.gx)
    updateβ!(o, o.rate(o.n))
end
function _merge!(o::StatLearn, o2::StatLearn)
    o.n += o2.n
    γ = nobs(o2) / nobs(o)
    smooth!(o.β, o2.β, γ)
    merge!(o.alg, o2.alg, γ)
    smooth!(o.λ, o2.λ, γ)
end

predict(o::StatLearn, x::VectorOb) = dot(x, o.β)
predict(o::StatLearn, x::AbstractMatrix) = x * o.β
classify(o::StatLearn, x) = sign.(predict(o, x))

function objective(o::StatLearn, x::AbstractMatrix, y::VectorOb)
    value(o.loss, y, predict(o, x), AggMode.Mean()) + value(o.penalty, o.β, o.λ)
end

#-----------------------------------------------------------------------# updateβ!
function updateβ!(o::StatLearn{SGD}, γ)
    for j in eachindex(o.β)
        o.β[j] = prox(o.penalty, o.β[j] - γ * o.gx[j], γ * o.λ[j])
    end
end
function updateβ!(o::StatLearn{T}, γ) where {T<:Union{ADAGRAD, RMSPROP}}
    for j in eachindex(o.β)
        s = γ / sqrt(o.alg.h[j] + ϵ)
        o.β[j] = prox(o.penalty, o.β[j] - s * o.gx[j], s * o.λ[j])
    end
end
function updateβ!(o::StatLearn{ADAM}, γ)
    for j in eachindex(o.β)
        s = γ / sqrt(o.alg.v[j] + ϵ)
        o.β[j] = prox(o.penalty, o.β[j] - s * o.alg.m[j], s * o.λ[j])
    end
end
function updateβ!(o::StatLearn{ADAMAX}, γ)
    for j in eachindex(o.β)
        s = γ / ((1 - o.alg.β1^nobs(o)) * o.alg.v[j])
        o.β[j] = prox(o.penalty, o.β[j] - s * o.alg.m[j], s * o.λ[j])
    end
end
function updateβ!(o::StatLearn{ADADELTA}, γ)
    for j in eachindex(o.β)
        s = o.alg.δ[j]
        o.β[j] = prox(o.penalty, o.β[j] - s * o.gx[j], s * o.λ[j])
    end
end

# #-----------------------------------------------------------------------# NSGD
# function fit!(o::StatLearn{NSGD}, t::Tuple{VectorOb, Real}, γ::Float64)
#     U = o.updater
#     x, y = t
#     for j in eachindex(o.β)
#         U.θ[j] = o.β[j] - U.α * U.v[j]
#     end
#     ŷ = dot(x, U.θ)
#     for j in eachindex(o.β)
#         U.v[j] = U.α * U.v[j] + deriv(o.loss, y, ŷ) * x[j]
#         @inbounds o.β[j] = prox(o.penalty, o.β[j] - γ * U.v[j], γ * o.λfactor[j])
#     end
# end
# #-----------------------------------------------------------------------# NADAM
# function update!(o::StatLearn{NADAM}, γ)
#     U = o.updater
#     β1 = U.β1
#     β2 = U.β2
#     U.nups += 1
#     @inbounds for j in eachindex(o.β)
#         gx = o.gx[j] + deriv(o.penalty, o.β[j], o.λfactor[j])
#         U.M[j] = smooth(gx, U.M[j], U.β1)
#         U.V[j] = smooth(gx ^ 2, U.V[j], U.β2)
#         mt = U.M[j] / (1 - U.β1 ^ U.nups)
#         vt = U.V[j] / (1 - U.β2 ^ U.nups)
#         Δ = γ / (sqrt(vt + ϵ)) * (U.β1 * mt + (1 - U.β1) / (1 - U.β1^U.nups) * gx)
#         o.β[j] -= Δ
#     end
# end

#------------------------------------------------------------------# Majorization-based
const L2Scaled{N} = LossFunctions.ScaledDistanceLoss{L2DistLoss, N}

# lipschitz_constant (L): f(θ) ≤ f(θₜ) + ∇f(θₜ)'(θ - θₜ) + (L / 2) ||θ - θₜ||^2
lconst(o::StatLearn, x, y) = lconst(o.loss, x, y)

lconst(o::Loss, x, y) = error("No defined Lipschitz constant for $o")
lconst(o::L2Scaled{N}, x, y) where {N} = 2N * dot(x, x)
lconst(o::LossFunctions.L2DistLoss, x, y) = 2 * dot(x, x)
lconst(o::LossFunctions.LogitMarginLoss, x, y) = .25 * dot(x, x)
lconst(o::LossFunctions.DWDMarginLoss, x, y) = (o.q + 1)^2 / o.q * dot(x, x)

#-----------------------------------------------------------------------# OMAS
# L stored in o.alg.a[1]
function _fit!(o::StatLearn{OMAS}, xy)
    γ = o.rate(o.n += 1)
    x, y = xy
    b = o.alg.b
    gradient!(o, x, y)
    ht = lconst(o.loss, x, y)
    L = (o.alg.a[1] = smooth(o.alg.a[1], ht, γ))
    for j in eachindex(o.β)
        b[j] = smooth(b[j], ht * o.β[j] - o.gx[j], γ)
        o.β[j] = prox(o.penalty, b[j] / L, o.λ[j] / L)
    end
end
#-----------------------------------------------------------------------# OMAP
function _fit!(o::StatLearn{OMAP}, xy)
    γ = o.rate(o.n += 1)
    x, y = xy
    gradient!(o, x, y)
    h_inv = inv(lconst(o, x, y))
    for j in eachindex(o.β)
        o.β[j] -= γ * h_inv * o.gx[j]
    end
end
#-----------------------------------------------------------------------# MSPI
function _fit!(o::StatLearn{MSPI}, xy)
    γ = o.rate(o.n += 1)
    x, y = xy
    gradient!(o, x, y)
    γ2 = γ / (1 + γ * lconst(o, x, y))
    for j in eachindex(o.β)
        @inbounds o.β[j] = prox(o.penalty, o.β[j] - γ2 * o.gx[j], γ2 * o.λ[j])
    end
end

ext: jl


text:
# This script won't work until we've released the new version of DSGE.jl (within next week).

## Add number of workers of one's choosing:
using DSGE, ModelConstructors, SMC
#addprocs_frbny(40)
@everywhere using DSGE, ModelConstructors, SMC, HDF5
@everywhere import ModelConstructors: ParameterVector
@everywhere import SMC: smc

## Initialize model object
fix_gamma = false
subspec = fix_gamma ? "ss1" : "ss0"
m = SmetsWoutersOrig(subspec)

saveroot = ""

m <= Setting(:saveroot, saveroot, "Output data directory path")
m <= Setting(:data_vintage,"041231")
m <= Setting(:date_presample_start, quartertodate("1965-Q4"))
m <= Setting(:date_mainsample_start, quartertodate("1966-Q4"))
m <= Setting(:date_forecast_start, quartertodate("2004-Q4"))
m <= Setting(:date_conditional_end, quartertodate("2004-Q4"))

# Load data
data = h5read("data/sw_orig_smc.h5", "data")

m <= Setting(:n_particles, 400)
m <= Setting(:n_Φ, 500)
m <= Setting(:λ, 2.1)
m <= Setting(:n_smc_blocks, 3)
m <= Setting(:n_mh_steps_smc, 1)
m <= Setting(:step_size_smc, 0.4)
m <= Setting(:use_parallel_workers, true)
m <= Setting(:resampler_smc, :multinomial)
m <= Setting(:target_accept, 0.25)

m <= Setting(:mixture_proportion, .9)
m <= Setting(:resampling_threshold, 0.5)
m <= Setting(:adaptive_tempering_target_smc, false)

# If want adaptive, just set the tempering target you want
#m <= Setting(:adaptive_tempering_target_smc, 0.95)

m <= Setting(:use_chand_recursion, true)

## Specify settings
verbose = :low
use_chand_recursion = true
filestring_addl     = Vector{String}()

## Specify output paths
loadpath            = rawpath(m, "estimate", "smc_cloud.jld2", filestring_addl)
savepath            = rawpath(m, "estimate", "smc_cloud.jld2", filestring_addl)
particle_store_path = rawpath(m, "estimate", "smcsave.h5",     filestring_addl)

DSGE.sendto(workers(), m = m)
@everywhere function my_likelihood(parameters::ParameterVector, data::Matrix{Float64})::Float64
    DSGE.update!(m, parameters)
    DSGE.likelihood(m, data; sampler = false, catch_errors = true,
                    use_chand_recursion = true, verbose = :low)
end

smc(my_likelihood, m.parameters, data; data_vintage = data_vintage(m),
    n_parts = 1_000, n_Φ = 100, λ = 2.0, parallel = true, save_intermediate = false,
    particle_store_path = particle_store_path, loadpath = loadpath,
    savepath = savepath)


text:
precompile = @timed using PowerSimulationsDynamics

using PowerSimulationsDynamics
const PSID = PowerSimulationsDynamics
using Sundials
using PowerSystems
const PSY = PowerSystems
using OrdinaryDiffEq
using Logging

open("precompile_time.txt", "a") do io
    write(io, "| $(ARGS[1]) | $(precompile.time) |\n")
end

sys = System("test/data_tests/WECC_240_dynamic.json"; runchecks = false)
for l in get_components(PSY.PowerLoad, sys)
    PSY.set_model!(l, PSY.LoadModels.ConstantImpedance)
end

# First runs
try
    for m in [ResidualModel, MassMatrixModel]
        Simulation(
            m,
            sys,
            pwd(),
            (0.0, 20.0), #time span
            BranchTrip(1.0, Line, "CORONADO    -1101-PALOVRDE    -1401-i_10");
            console_level = Logging.Error,
        )
    end
catch e
    @error exception = (e, catch_backtrace())
end

try
    sim_ida, time_build_ida, _, _ = @timed Simulation(
        ResidualModel,
        sys,
        pwd(),
        (0.0, 20.0), #time span
        BranchTrip(1.0, Line, "CORONADO    -1101-PALOVRDE    -1401-i_10");
        console_level = Logging.Error,
    )
    status = execute!(sim_ida, IDA(), dtmax = 0.01, enable_progress_bar = false)
    if status == PSID.SIMULATION_FINALIZED
        res_ida = read_results(sim_ida)
        solve_time = res_ida.time_log[:timed_solve_time]
        open("execute_time.txt", "a") do io
            write(io, "| $(ARGS[1])-Build ResidualModel | $(time_build_ida) |\n")
            write(io, "| $(ARGS[1])-Execute ResidualModel | $(solve_time) |\n")
        end
    else
        error("FAILED TO SOLVE")
    end
catch e
    @error exception = (e, catch_backtrace())
    open("execute_time.txt", "a") do io
        write(io, "| $(ARGS[1])- ResidualModel | FAILED TO TEST |\n")
    end
end

try
    sim_rodas, time_build_rodas, _, _ = @timed Simulation(
        MassMatrixModel,
        sys, #system
        pwd(),
        (0.0, 20.0), #time span
        BranchTrip(1.0, Line, "CORONADO    -1101-PALOVRDE    -1401-i_10"),
        #console_level = Logging.Error,
    ) #Type of Fault
    status = execute!(sim_rodas, Rodas4(), rtol = 1e-9, atol = 1e-9)
    if status == PSID.SIMULATION_FINALIZED
        res_rodas = read_results(sim_rodas)
        solve_time = res_rodas.time_log[:timed_solve_time]
        open("execute_time.txt", "a") do io
            write(io, "| $(ARGS[1])-Build MassMatrixModel | $(time_build_rodas) |\n")
            write(io, "| $(ARGS[1])-Execute MassMatrixModel | $(solve_time) |\n")
        end
    else
        error("FAILED TO SOLVE")
    end
catch e
    @error exception = (e, catch_backtrace())
    open("execute_time.txt", "a") do io
        write(io, "| $(ARGS[1])-MassMatrixModel | FAILED TO TEST |\n")
    end
end


text:
export CuTensor

using LinearAlgebra

mutable struct CuTensor{T, N}
    data::CuArray{T, N}
    inds::Vector{Cwchar_t}
    function CuTensor{T, N}(data::CuArray{T, N}, inds::Vector{Cwchar_t}) where {T<:Number, N}
        new(data, inds)
    end
    function CuTensor{T, N}(data::CuArray{N, T}, inds::Vector{<:AbstractChar}) where {T<:Number, N}
        new(data, Cwchar_t.(inds))
    end
end
CuTensor(data::CuArray{T, N}, inds::Vector{<:AbstractChar}) where {T<:Number, N} =
    CuTensor{T, N}(data, convert(Vector{Cwchar_t}, inds))
CuTensor(data::CuArray{T, N}, inds::Vector{Cwchar_t}) where {T<:Number, N} =
    CuTensor{T, N}(data, inds)

Base.size(T::CuTensor) = size(T.data)
Base.size(T::CuTensor, i) = size(T.data, i)
Base.length(T::CuTensor) = length(T.data)
Base.ndims(T::CuTensor) = length(T.inds)
Base.strides(T::CuTensor) = strides(T.data)
Base.eltype(T::CuTensor) = eltype(T.data)
Base.similar(T::CuTensor{Tv, N}) where {Tv, N} = CuTensor{Tv, N}(similar(T.data), copy(T.inds))
Base.collect(T::CuTensor) = (collect(T.data), T.inds)

function Base.:(+)(A::CuTensor, B::CuTensor)
    α = convert(eltype(A), 1.0)
    γ = convert(eltype(B), 1.0)
    C = similar(B)
    elementwiseBinary!(α, A, CUTENSOR_OP_IDENTITY, γ, B, CUTENSOR_OP_IDENTITY, C, CUTENSOR_OP_ADD)
end

function Base.:(-)(A::CuTensor, B::CuTensor)
    α = convert(eltype(A), 1.0)
    γ = convert(eltype(B), -1.0)
    C = similar(B)
    elementwiseBinary!(α, A, CUTENSOR_OP_IDENTITY, γ, B, CUTENSOR_OP_IDENTITY, C, CUTENSOR_OP_ADD)
end

LinearAlgebra.axpy!(a, X::CuTensor, Y::CuTensor) = elementwiseBinary!(a, X, CUTENSOR_OP_IDENTITY, one(eltype(Y)), Y, CUTENSOR_OP_IDENTITY, similar(Y), CUTENSOR_OP_ADD)
LinearAlgebra.axpby!(a, X::CuTensor, b, Y::CuTensor) = elementwiseBinary!(a, X, CUTENSOR_OP_IDENTITY, b, Y, CUTENSOR_OP_IDENTITY, similar(Y), CUTENSOR_OP_ADD)

LinearAlgebra.mul!(C::CuTensor, A::CuTensor, B::CuTensor) = contraction!(one(eltype(C)), A, CUTENSOR_OP_IDENTITY, B, CUTENSOR_OP_IDENTITY, zero(eltype(C)), C, CUTENSOR_OP_IDENTITY, CUTENSOR_OP_IDENTITY)

function Base.:(*)(A::CuTensor, B::CuTensor)
    tC = promote_type(eltype(A), eltype(B))
    A_uniqs = [(idx, i) for (idx, i) in enumerate(A.inds) if !(i in B.inds)]
    B_uniqs = [(idx, i) for (idx, i) in enumerate(B.inds) if !(i in A.inds)]
    A_sizes = map(x->size(A,x[1]), A_uniqs)
    B_sizes = map(x->size(B,x[1]), B_uniqs)
    A_inds = map(x->Cwchar_t(x[2]), A_uniqs)
    B_inds = map(x->Cwchar_t(x[2]), B_uniqs)
    C = CuTensor(CuArrays.zeros(tC, Dims(vcat(A_sizes, B_sizes))), vcat(A_inds, B_inds))
    return mul!(C, A, B)
end


text:
#Lazy loading of stitched high-speed experiment (applies both a stitching tfm and a slicewise calibration)
function lazy_load(cam1::AbstractArray, cam2::AbstractArray, octr, stitch_tfm::T, fwd_tfms, bck_tfms;
                        sqrt_tfm=false, correct_bias=true) where{T<:CoordinateTransformations.Transformation}
    #the match_axisspacing call works around a breakage in AxisArrays in this case
    cam1fwd = match_axisspacing(view(cam1.data.data,:,:,:,1:2:size(cam1,4)), cam1.data)
    cam2fwd = match_axisspacing(view(cam2.data.data,:,:,:,1:2:size(cam2,4)), cam2.data)
    cam1bck = match_axisspacing(view(cam1.data.data,:,:,:,2:2:size(cam1,4)), cam1.data)
    cam2bck = match_axisspacing(view(cam2.data.data,:,:,:,2:2:size(cam2,4)), cam2.data)

    #(This workaround isn't ideal because it doesn't create a StepRangeLen, so the notion of uniform axis spacing is lost
    #cam1fwd = view(cam1,:,:,:,[1:2:size(cam1,4)...])
    #cam2fwd = view(cam2,:,:,:,[1:2:size(cam2,4)...])
    #cam1bck = view(cam1,:,:,:,[2:2:size(cam1,4)...])
    #cam2bck = view(cam2,:,:,:,[2:2:size(cam2,4)...])

    ssfwd = TeamedImaging.stitch_slicewise(cam1fwd, cam2fwd, stitch_tfm, fwd_tfms, Float64;
        octr=octr, correct_bias=correct_bias, sqrt_tfm=sqrt_tfm, flip_y_bottom=true)
    ssbck = TeamedImaging.stitch_slicewise(cam1bck, cam2bck, stitch_tfm, bck_tfms, Float64;
        octr=octr, correct_bias=correct_bias, sqrt_tfm=sqrt_tfm, flip_y_bottom=true)
    ssi = InterlacedStackSeries(ssfwd, ssbck)
    return ssi
end


text:
using NHST

@assert isapprox(p_value(BinomialTest, 0, 10, 0.5), 0.001953125)
@assert isapprox(p_value(BinomialTest, 5, 10, 0.5), 1.0)
@assert isapprox(p_value(BinomialTest, 10, 10, 0.5), 0.001953125)
@assert isapprox(p_value(BinomialTest, 4, 10, 0.2), 0.1208739)
@assert isapprox(p_value(BinomialTest, 2, 10, 0.8), 7.79264e-05)
@assert isapprox(p_value(BinomialTest, 3, 100, 0.5), 2.630867e-25)
@assert isapprox(p_value(BinomialTest, 97, 100, 0.5), 2.630867e-25)
@assert isapprox(p_value(BinomialTest, 3, 10, 0.79), 0.001178273)
@assert isapprox(p_value(BinomialTest, 9, 10, 0.79), 0.6989298)

p_value(BinomialTest, 0, 10, 0.5)
p_value(BinomialTest, 0.0, 10, 0.5)
p_value(BinomialTest, 0, 10.0, 0.5)
p_value(BinomialTest, 0.0, 10.0, 0.5)
p_value(BinomialTest, zeros(10), 0.5)

run(BinomialTest, 0, 10, p = 0.5)
run(BinomialTest, 0.0, 10, p = 0.5)
run(BinomialTest, 0, 10.0, p = 0.5)
run(BinomialTest, 0.0, 10.0, p = 0.5)
run(BinomialTest, zeros(10), p = 0.5)

@assert isapprox(p_value(run(BinomialTest, 0, 10, p = 0.5)), 0.001953125)
@assert isapprox(p_value(run(BinomialTest, 5, 10, p = 0.5)), 1.0)
@assert isapprox(p_value(run(BinomialTest, 10, 10, p = 0.5)), 0.001953125)
@assert isapprox(p_value(run(BinomialTest, 4, 10, p = 0.2)), 0.1208739)
@assert isapprox(p_value(run(BinomialTest, 2, 10, p = 0.8)), 7.79264e-05)


text:
include("distributions.jl")
include("control.jl")

function surface_mesh(
    dist::SurfaceDistribution, x::AbstractVector, y::AbstractVector; mg_level::Int=0
)::Vector{MVector{2,Float64}}
    # check inputs
    @assert length(x) == length(y)
    @assert mg_level in 0:6
    if y[end - 1] < y[2]
        reverse!(x)
        reverse!(y)
    end
    x_te = (x[end] + x[1]) / 2
    y_te = (y[end] + y[1]) / 2
    t_te = sqrt((y[end] - y[1])^2 + (x[end] - x[1])^2)
    closed_te = t_te < 1e-3
    @assert abs(x_te - 1) < 0.01
    if closed_te
        y[end] = y[1] = y_te
        x[end] = x[1] = x_te
        np_te = 0
    end
    i_le = argmin(x)
    @assert issorted(x[i_le:end]) && issorted(x[1:i_le]; rev=true)

    s = Vector{Float64}(undef, length(x))
    s[1] = 0
    @inbounds for i in 2:length(s)
        s[i] = s[i - 1] + sqrt((x[i] - x[i - 1])^2 + (y[i] - y[i - 1])^2)
    end

    smax = s[end]
    s_le = s[i_le]
    sx = Spline1D(s, x)
    sy = Spline1D(s, y)
    xus = Spline1D(x[i_le:end], s[i_le:end])
    xls = Spline1D(x[i_le:-1:1], s[i_le:-1:1])

    s_nodes = map(n -> to_s(n, s_le, smax, xus, xls), dist.node)
    #
    sm = nodes(dist, s_nodes; mg_level=mg_level)
    xsurf, ysurf = sx(sm), sy(sm)
    if !closed_te
        ds1 = sqrt((xsurf[1]-xsurf[2])^2 + (ysurf[1]-ysurf[2])^2)
        ds2 = sqrt((xsurf[end]-xsurf[end-1])^2 + (ysurf[end]-ysurf[end-1])^2)
        ds_te = sqrt(ds1*ds2)
        ns_te = round(Int, t_te/ds_te)
        mg_te = 2^max(mg_level,1)
        np_te = ceil(Int, ns_te//mg_te)*mg_te + 1
        halfte = (np_te + 1) ÷ 2
        xsurf = vcat(
            LinRange(x_te, x[1], halfte),
            xsurf[2:(end - 1)],
            LinRange(x[end], x_te, halfte),
        )
        ysurf = vcat(
            LinRange(y_te, y[1], halfte),
            ysurf[2:(end - 1)],
            LinRange(y[end], y_te, halfte),
        )
    else
        xsurf[1] = xsurf[end] = x_te
        ysurf[1] = ysurf[end] = y_te
    end
    points = map(zip(xsurf, ysurf)) do (x, y)
        MVector{2,Float64}(x, y)
    end
    points[end] = points[1]
    return points
end


text:
using Test, TestSetExtensions, SafeTestsets, Logging, Pkg
using HTTP, Genie, Pluto

include("html.jl")

function testnotebook(input)
    session = Pluto.ServerSession();
    notebook = Pluto.SessionActions.open(session, input; run_async=false)
    html_contents = body(notebook)

    # split input string by '--', '.' and '_' to array of strings
    # e.g. '18--a_b_c.jl' -> ['18', 'a', 'b', 'c', 'jl']
    # html_input string
    first_input = last(split(input, "/"))
    new_input = replace(last(split(split(first_input, ".")[1], "--")), r"_" => s"-")

    write(joinpath(@__DIR__, "../build/tutorials/","$(new_input).html"), html_contents)

    errored=false
    for c in notebook.cells
        if c.errored
            errored=true
            @error "Error in  $(c.cell_id): $(c.output.body[:msg])\n $(c.code)"
        end
        @assert !c.errored
        @test !c.errored
    end
    !errored
end

@testset "testing tutorials" begin
    notebooks= [
        "0--Index/0--index.jl",
        "1--Overview/1--Overview.jl",
        "2--Installing_Genie/2--Installing_Genie.jl",
        "3--Getting_Started/3--Getting_Started.jl",
        "4--Developing_Web_Services/4--Developing_Web_Services.jl",
        "41--Developing_MVC_Web_Apps/41--Developing_MVC_Web_Apps.jl",
        "5--Handling_Query_Params/5--Handling_Query_Params.jl",
        "6--Working_with_POST_Payloads/6--Working_with_POST_Payloads.jl",
        "8--Handling_File_Uploads/8--Handling_File_Uploads.jl",
        "9--Publishing_Your_Julia_Code_Online_With_Genie_Apps/9--Publishing_Your_Julia_Code_Online_With_Genie_Apps.jl",
        "10--Loading_Genie_Apps/10--Loading_Genie_Apps.jl",
        "11--Managing_External_Packages/11--Managing_External_Packages.jl",
        "13--Initializers/13--Initializers.jl",
        "14--The_Secrets_File/14--The_Secrets_File.jl",
        "15--The_Lib_Folder/15--The_Lib_Folder.jl",
        "16--Using_Genie_With_Docker/16--Using_Genie_With_Docker.jl",
        "17--Working_with_Web_Sockets/17--Working_with_Web_Sockets.jl",
        "80--Force_Compiling_Routes/80--Force_Compiling_Routes.jl",
        "90--Deploying_With_Heroku_Buildpacks/90--Deploying_With_Heroku_Buildpacks.jl",
        "92--Deploying_Genie_Server_Apps_with_Nginx/92--Deploying_Genie_Server_Apps_with_Nginx.jl"
    ]

    for notebook in notebooks
        input=joinpath(@__DIR__,"PlutoNotebooks/",notebook)
        @show input
        @info "notebook: $(input)"
        @test testnotebook(input)
    end
end


text: hero(bullets, dragons) = dragons ≤ bullets ÷ 2

max_stars_repo_name: UlrichBerntien/Codewars-Katas

max_issues_repo_name: UlrichBerntien/Codewars-Katas

max_forks_repo_name: UlrichBerntien/Codewars-Katas


text:
# This could be factored out into its own package sometime,
# there is some overlap with ProgressMeter.jl though.
# Don't depend on PackageCompiler.jl to use this :P
module TerminalSpinners

using Printf

export Spinner, autospin, spin!, @spin

const CSI = "\e["

show_cursor!(io::IO) = print(io, CSI, "?25h")
hide_cursor!(io::IO) = print(io, CSI, "?25l")
cursor_up!(io::IO, n=1) = print(io, CSI, n, 'A')
cursor_down!(io::IO, n=1) = print(io, CSI, n, 'B')
cursor_horizontal_absolute!(io::IO, n=1) = print(io, CSI, n, 'G')

@enum EraseLineMode begin
    CURSOR_TO_END = 0
    CURSOR_TO_BEGINNING = 1
    ENTIRE_LINE = 2
end
erase_line!(io::IO, mode::EraseLineMode=ENTIRE_LINE) = print(io, CSI, Int(mode), 'K')

function remove_ansi_characters(str::String)
    r = r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])"
    return replace(str, r => "")
end

function count_lines_without_ansi(width::Int, str::String)
    n_lines = 0
    for line in readlines(IOBuffer(remove_ansi_characters(str)))
        w, f = divrem(textwidth(line), width)
        n_lines += w + (f > 0)
    end
    return n_lines
end

tostring(msg) = msg
tostring(f::Function) = f()

Base.@kwdef mutable struct Spinner{IO_t <: IO}
    frames::Vector{String} =  ["⠋", "⠙", "⠸", "⢰", "⣠", "⣄", "⡆", "⠇"]
    freq::Float64 = 10.0 # [1/s]
    msg::Any = ""
    stream::IO_t = stderr
    timer::Union{Nothing, Timer} = nothing
    hidecursor::Bool = true
    silent::Bool=false
    enabled::Bool= silent ? false : stream isa Base.TTY && !haskey(ENV, "CI")
    frame_idx::Int=1
    start = time()
    color::Symbol = Base.info_color()
    nlines::Int=0
    first::Bool=false
end

getframe(s::Spinner) = s.frames[s.frame_idx]

function advance_frame!(s::Spinner)
    frame = s.frames[s.frame_idx]
    s.frame_idx = s.frame_idx == length(s.frames) ? 1 : s.frame_idx + 1
    return frame
end

function erase_and_reset(io::IO, n_lines::Int)
    erase_line!(io)
    for _ in 1:n_lines-1
        cursor_up!(io, 1)
        erase_line!(io)
    end
    cursor_horizontal_absolute!(io, 1)
end

function getline(s::Spinner, spinner, color)
    ioc = IOContext(s.stream, :displaysize=>displaysize(s.stream)) # https://github.com/JuliaLang/julia/issues/42649
    return sprint(; context=ioc) do io
        s.first || erase_and_reset(io, s.nlines)
        printstyled(io, spinner; color, bold=true)

        elapsed = time() - s.start
        (minutes, seconds) = fldmod(elapsed, 60)
        (hours, minutes) = fldmod(minutes, 60)

        if hours == 0
            printstyled(io, @sprintf(" [%02dm:%02ds]", minutes, seconds); color, bold=true)
        else
            printstyled(io, @sprintf(" [%02dh:%02dm:%02ds]", hours, minutes, seconds); color, bold=true)
        end

        msg = tostring(s.msg)
        print(io, " ", msg)
    end
end

function render(s::Spinner, spinner=getframe(s), color=s.color)
    s.silent && return
    str = getline(s, spinner, color)
    s.nlines = count_lines_without_ansi(displaysize(s.stream)[2], str)
    print(s.stream, str)
    return
end

function start!(s::Spinner)
    s.silent && return
    s.frame_idx = 1
    s.first = true
    s.start = time()

    if !s.enabled
        println(s.stream, "- ", s.msg)
        return
    end

    s.hidecursor && hide_cursor!(s.stream)

    t = Timer(0.0; interval=1/s.freq) do timer
        try
            render(s)
            s.first = false
            advance_frame!(s)
        catch e
            @error "internal error in spinner" exception=(e, catch_backtrace())
            stop!(s)
        end
    end
    s.timer = t
    return s
end

function stop!(s::Spinner)
    if s.timer !== nothing
        close(s.timer)
    end
    if !s.enabled || s.silent
        return
    end
    s.hidecursor && show_cursor!(s.stream)
    println(s.stream)
end

function success!(s)
    if s.enabled && !s.silent
        render(s, "✔", :light_green)
    end
    stop!(s)
end
function fail!(s)
    if s.enabled && !s.silent
        render(s, "✖", :light_red)
    end
    stop!(s)
end

macro spin(s, work)
    return quote
        spin(() -> $(esc(work)), $(esc(s)))
    end
end

function spin(f, s::Spinner)
    start!(s)
    try
        v = f()
        success!(s)
        v
    catch
        fail!(s)
        rethrow()
    end
end

end # module


text:
# old test file

include("../src/read_solomon_data.jl")
include("../src/pulse.jl")

using ElasticArrays
# using Gurobi, JuMP
using Random
# Random.seed!(0)

dataset_name = "C102_025"
data_file_path = dataset_name * ".xml"
data_file_path = joinpath(@__DIR__, "..", "solomon-1987", data_file_path)
dataset, data_name, nodes, fleet, requests = read_solomon_data(data_file_path)

n_nodes = length(nodes)
n_vehicles = fleet.number
n_requests = length(requests)

# Add a dummy node for depot at the end
push!(nodes, nodes[n_nodes])
depot0 = n_requests + 1
depot_dummy = n_requests + 2

n_nodes = length(nodes)
@assert n_nodes - 2 == n_requests
V = 1:n_vehicles
N = 1:n_nodes
C = 1:n_requests

d = [requests[i].quantity for i in C]
q = fleet.capacity
a = [requests[i].start_time for i in C]
b = [requests[i].end_time for i in C]

t = calculate_cost(nodes)
t[depot0, depot_dummy] = Inf
t[depot_dummy, depot0] = Inf
for i in N, j in N
    if i in C
        t[i, j] = t[i, j] + requests[i].service_time
    else
        t[i, j] = t[i, j]
    end
end

origin = depot0
destination = depot_dummy
capacity = fleet.capacity
cost_mtx = t .* (rand(size(t)...) .- 0.50)
time_mtx = t

resrc_mtx = zeros(n_nodes, n_nodes)
for i in N, j in N
    if j in C
        resrc_mtx[i, j] = requests[j].quantity
    end
end
early_time = [a; 0; 0]
late_time = [b; 0; fleet.max_travel_time]

# pulse_path = solveESPPRC(origin, destination, resrc_cap, cost_mtx, time_mtx, resrc_mtx, early_time, late_time)
# @show pulse_path.path, pulse_path.cost, pulse_path.time

pg = PulseGraph(
    origin,
    destination,
    capacity,
    cost_mtx,
    time_mtx,
    resrc_mtx,
    early_time,
    late_time,
)

##############################################################################################################
# cost_mtx[origin, destination] = 10000
# cost_mtx[destination, origin] = 10000
# time_mtx[origin, destination] = 10000
# time_mtx[destination, origin] = 10000

# bigM = maximum(b) + maximum(t)
# bigQ = q + maximum(d)

# n_vehicles = 1
# m = Model(Gurobi.Optimizer)
# @variable(m, x[N, N], Bin)
# @variable(m, s[C] >= 0)
# @variable(m, qq[N] >= 0)

# @objective(m, Min, sum(cost_mtx[i,j] * x[i,j] for i in N, j in N))

# @constraint(m, [i in N, j in C], qq[j] >= qq[i] + d[j] - bigQ * (1 - x[i,j]))
# @constraint(m, [i in N], 0 <= qq[i] <= q)

# @constraint(m, sum(x[depot0,j] for j in N) == n_vehicles )
# @constraint(m, [h in C], sum(x[i,h] for i in N) - sum(x[h,j] for j in N) == 0 )
# @constraint(m, sum(x[i,depot_dummy] for i in N) == sum(x[depot0,j] for j in N) )

# @constraint(m, [i in C, j in C], s[j] >= s[i] + time_mtx[i,j] - bigM * (1 - x[i,j]))
# @constraint(m, [i in C], a[i] <= s[i] <= b[i])

# @constraint(m, [i in N], x[i, depot0] == 0)
# @constraint(m, [i in N], x[depot_dummy, i] == 0)
# @constraint(m, [i in N], x[i, i] == 0)

# optimize!(m)
# @show raw_status(m)
# @show objective_value(m)

# xx = zeros(size(x))
# for i in 1:size(xx,1)
#     for j in 1:size(xx,2)
#         xx[i,j] = JuMP.value.(x[i,j])
#     end
# end

# function to_path(xx, origin, destination)
#     current = origin
#     path = [current]
#     while current != destination
#         j = findall(x -> x==1, xx[current,:])
#         push!(path, j[1])
#         current = j[1]
#     end
#     return path
# end

# path = to_path(xx, origin, destination)
# @show path

##############################################################################################################

@time pulse_path = solveESPPRCpulse(pg)
@show pulse_path.path, pulse_path.cost, pulse_path.time

# @assert isapprox(objective_value(m), pulse_path.cost; atol=1e-6)
# @assert path == pulse_path.path


text:
"""
Adapted from
https://github.com/sandreza/OceanConvectionUQSupplementaryMaterials/blob/master/src/ForwardMap/fm.jl
updated for latest version of OceanTurb
"""

# include("../les/custom_avg.jl")
# include("../les/output_to_data.jl")
# include("../les/get_les_data.jl")
# include("../data/scalings.jl")

"""
closure_free_convection(N, Δt, les::OceananigansData; subsample = 1, grid = 1)
# Description
- constructs forward map. Assumes initial conditions and boundary conditions are taken from les data.
# Arguments
- `N`: number of gridpoints to output to
- `Δt`: time step size in seconds
- `les`: les data of the OceananigansData type
# Keyword Arguments
- `subsample`: indices to subsample in time,
- `grid`: in case one wants to save the model grid
# Output
- The forward map. A function that takes parameters and outputs temperature profiles
-   `𝑪`: parameters in KPP, assumes that \n
    𝑪[1]: Surface Layer Fraction  \n
    𝑪[2]: Nonlocal Flux Amplitude \n
    𝑪[3]: Diffusivity Amplitude \n
    𝑪[4]: Shear Constant \n
"""
function closure_free_convection_kpp_full_evolution(parameters, D, Δt, les::OceananigansData;
                                     subsample = 1, grid = 1)
     # # set parameters
     # parameters = KPP.Parameters( CSL = 𝑪[1], CNL = 𝑪[2], Cb_T = 𝑪[3], CKE = 𝑪[4])
     # Build the model with a Backward Euler timestepper
     constants = Constants(Float64; α = les.α , β = les.β, ρ₀= les.ρ, cP=les.cᵖ, f=les.f⁰, g=les.g)
     model = KPP.Model(N=D, H=les.L, stepper=:BackwardEuler, constants = constants, parameters = parameters)
     # Get grid if necessary
     if grid != 1
         zp = collect(model.grid.zc)
         @. grid  = zp
     end

     # Set boundary conditions
     model.bcs.T.top = FluxBoundaryCondition(les.top_T)
     model.bcs.T.bottom = GradientBoundaryCondition(les.bottom_T)

    # define the closure
    function free_convection()
        # get average of initial condition of LES
        T⁰ = custom_avg(les.T⁰, D)
        # set equal to initial condition of parameterization
        model.solution.T[1:D] = copy(T⁰)
        # # Set boundary conditions
        # model.bcs.T.top = FluxBoundaryCondition(les.top_T)
        # model.bcs.T.bottom = GradientBoundaryCondition(les.bottom_T)
        # set aside memory
        if subsample != 1
            time_index = subsample
        else
            time_index = 1:length(les.t)
        end
        Nt = length(les.t[time_index])
        𝒢 = zeros(D, Nt)

        # loop the model
        ti = collect(time_index)
        for i in 1:Nt
            t = les.t[ti[i]]
            run_until!(model, Δt, t)
            @. 𝒢[:,i] = model.solution.T[1:D]
        end
        return 𝒢
    end
    return free_convection
end

function closure_free_convection_kpp(parameters, D, Δt, les::OceananigansData;
                                     subsample = 1, grid = 1, n_steps=1)

     # # set parameters
     # parameters = KPP.Parameters( CSL = 𝑪[1], CNL = 𝑪[2], Cb_T = 𝑪[3], CKE = 𝑪[4])
     # Build the model with a Backward Euler timestepper
     constants = Constants(Float64; α = les.α , β = les.β, ρ₀= les.ρ, cP=les.cᵖ, f=les.f⁰, g=les.g)
     model = KPP.Model(N=D, H=les.L, stepper=:BackwardEuler, constants = constants, parameters = parameters)
     # Get grid if necessary
     if grid != 1
         zp = collect(model.grid.zc)
         @. grid  = zp
     end
     # Set boundary conditions
     model.bcs.T.top = FluxBoundaryCondition(les.top_T)
     model.bcs.T.bottom = GradientBoundaryCondition(les.bottom_T)

     # set aside memory
     if subsample != 1
         time_index = subsample
     else
         time_index = 1:length(les.t)
     end

     Nt = length(les.t[time_index])

     # loop the model
     ti = collect(time_index)
     ts = [les.t[ti[i]] for i in 1:n_steps+1]
     𝒢 = zeros(D, n_steps+1)

    # define the closure
    function evolve_forward(; T⁰=T⁰)

        # average the initial condition
        T⁰ = custom_avg(T⁰, D)

        # set equal to initial condition of parameterization
        model.solution.T[1:D] = T⁰
        # # Set boundary conditions
        # model.bcs.T.top = FluxBoundaryCondition(les.top_T)
        # model.bcs.T.bottom = GradientBoundaryCondition(les.bottom_T)

        for i in 1:n_steps+1
            # t = les.t[ti[i]]
            run_until!(model, Δt, ts[i])
            @. 𝒢[:,i] = model.solution.T[1:D]
        end
        return 𝒢
    end
    return evolve_forward
end

"""
closure_free_convection_flexible(N, Δt, les::OceananigansData; subsample = 1, grid = 1)
# Description
- constructs forward map. Assumes initial conditions and boundary conditions are taken from les data. Adds two more parameters compared to previous one
# Arguments
- `N`: number of gridpoints to output to
- `Δt`: time step size in seconds
- `les`: les data of the OceananigansData type
# Keyword Arguments
- `subsample`: indices to subsample in time,
- `grid`: in case one wants to save the model grid
- `power`: a method of modifying the scaling law, default is 1
# Output
- The forward map. A function that takes parameters and outputs temperature profiles
-   `𝑪`: parameters in KPP, assumes that \n
    𝑪[1]: Surface Layer Fraction  \n
    𝑪[2]: Nonlocal Flux Amplitude \n
    𝑪[3]: Diffusivity Amplitude \n
    𝑪[4]: Shear Constant, w1 * w2 \n
    𝑪[5] = Shear Constant 2
    𝑪[6]: Shear Constant 3,\n
- The extra parameters correspond to  C[5] * w1^(1-power) * w2^(1+power)
- where w1 = max(0, h * Qb)^(1/3), w2 = sqrt(𝑪[6] * h^2), are two velocity scales.
"""
function closure_free_convection_flexible(D, Δt, les::OceananigansData;
                                 subsample = 1, grid = 1, power = 0.0)
     # set parameters
     parameters = KPP.Parameters( CSL = 𝑪[1], CNL = 𝑪[2], Cb_T = 𝑪[3], CKE = 𝑪[4], CKE2 = 𝑪[5], CKE3 = 𝑪[6], CKE4 = power)
     # Build the model with a Backward Euler timestepper
     constants = Constants(Float64; α = les.α , β = les.β, ρ₀= les.ρ, cP=les.cᵖ, f=les.f⁰, g=les.g)
     model = KPP.Model(N=D, H=les.L, stepper=:BackwardEuler, constants = constants, parameters = parameters)
     # Get grid if necessary
     if grid != 1
         zp = collect(model.grid.zc)
         @. grid  = zp
     end

    # define the closure
    function free_convection(𝑪)
        # get average of initial condition of LES
        T⁰ = custom_avg(les.T⁰, N)
        # set equal to initial condition of parameterization
        model.solution.T[1:N] = copy(T⁰)
        # Set boundary conditions
        model.bcs.T.top = FluxBoundaryCondition(les.top_T)
        model.bcs.T.bottom = GradientBoundaryCondition(les.bottom_T)
        # set aside memory
        if subsample != 1
            time_index = subsample
        else
            time_index = 1:length(les.t)
        end
        Nt = length(les.t[time_index])
        𝒢 = zeros(N, Nt)

        # loop the model
        ti = collect(time_index)
        for i in 1:Nt
            t = les.t[ti[i]]
            run_until!(model, Δt, t)
            @. 𝒢[:,i] = model.solution.T[1:N]
        end
        return 𝒢
    end
    return free_convection
end

max_stars_repo_name: adelinehillier/LearnConvection

max_issues_repo_name: adelinehillier/LearnConvection

max_forks_repo_name: adelinehillier/LearnConvection


text:
################################################################################
#
#	INTERFACE FOR DEFAULT POINTS IN RECIPROCAL SPACE
#	(implementations see below)
#
################################################################################

# general interface
function getReciprocalPoint(
            :: Type{R},
            unitcell   :: U,
            identifier :: Symbol,
            instance   :: Int64 = 1
        ) :: R where {D,N,LS,LB, R<:AbstractReciprocalPoint{D}, S<:AbstractSite{LS,D}, B<:AbstractBond{LB,N}, U<:AbstractUnitcell{S,B}}

    # maybe instance is passed from another source in some notation as e.g. X_2
    if instance == 1 && occursin('_', string(identifier))
        return getReciprocalPoint(R, unitcell, Val(Symbol(split(string(identifier),'_')[1])), Val(Meta.eval(Meta.parse(split(string(identifier),'_')[2]))))
    end
    # return the specific val typed function
    return getReciprocalPoint(R, unitcell, Val(identifier), Val(instance))
end
function getReciprocalPoint(
            :: Type{R},
            unitcell   :: U,
            identifier :: Tuple{Symbol, Int64}
        ) :: R where {D,N,LS,LB, R<:AbstractReciprocalPoint{D}, S<:AbstractSite{LS,D}, B<:AbstractBond{LB,N}, U<:AbstractUnitcell{S,B}}

    # return the specific val typed function
    return getReciprocalPoint(R, unitcell, Val(identifier[1]), Val(identifier[2]))
end

# interface for own concrete reciprocal point type
function getReciprocalPoint(
            unitcell   :: U,
            identifier :: Symbol,
            instance   :: Int64 = 1
        ) :: ReciprocalPoint{D} where {D,N,LS,LB, S<:AbstractSite{LS,D}, B<:AbstractBond{LB,N}, U<:AbstractUnitcell{S,B}}

    # maybe instance is passed from another source in some notation as e.g. X_2
    if instance == 1 && occursin('_', string(identifier))
        return getReciprocalPoint(ReciprocalPoint{D}, unitcell, Val(Symbol(split(string(identifier),'_')[1])), Val(Meta.eval(Meta.parse(split(string(identifier),'_')[2]))))
    end
    # return the specific val typed function
    return getReciprocalPoint(ReciprocalPoint{D}, unitcell, Val(identifier), Val(instance))
end

# Fallback for all functions
function getReciprocalPoint(
            :: Type{R},
            unitcell   :: U,
            identifier :: Val{I1},
            instance   :: Val{I2}
        ) :: R where {I1,I2,D,N,LS,LB, R<:AbstractReciprocalPoint{D}, S<:AbstractSite{LS,D}, B<:AbstractBond{LB,N},U<:AbstractUnitcell{S,B}}

    # fallback / fail due to missing implementation
    error("Implementation of reciprocal point " * string(I1) * " (instance " * string(I2) * ") not implemented yet " *
    "for concrete reciprocal point type " * string(R) * " and unitcells of dimension D=" * string(D) * " and N="*string(N))
end

# export the general interface function
export getReciprocalPoint

################################################################################
#
#	IMPLEMENTATIONS OF DEFAULT POINTS IN RECIPROCAL SPACE
#
################################################################################

#########################
#
#	GAMMA POINT
#
#########################

# Gamma point in arbitrary dimension (Fallback)
function getReciprocalPoint(
            :: Type{R},
            unitcell   :: U,
            identifier :: Val{:Gamma},
            instance   :: Val{I}
        ) :: R where {I,D,N,LS,LB, R<:AbstractReciprocalPoint{D}, S<:AbstractSite{LS,D}, B<:AbstractBond{LB,N},U<:AbstractUnitcell{S,B}}

    # fallback / fail due to missing implementation
    error("Implementation of Gamma point in " * string(D) * " spatial dimensions missing " *
    "for concrete reciprocal point type " * string(R))
end

# Gamma point in 2D
function getReciprocalPoint(
            :: Type{R},
            unitcell   :: U,
            identifier :: Val{:Gamma},
            instance   :: Val{I}
        ) :: R where {I,N,LS,LB, R<:AbstractReciprocalPoint{2}, S<:AbstractSite{LS,2}, B<:AbstractBond{LB,N},U<:AbstractUnitcell{S,B}}

    # return the specific type
    return newReciprocalPoint(
        R,
        [0.0, 0.0],
        "Gamma",
        "\\Gamma"
    )
end

# Gamma point in 3D
function getReciprocalPoint(
            :: Type{R},
            unitcell   :: U,
            identifier :: Val{:Gamma},
            instance   :: Val{I}
        ) :: R where {I,N,LS,LB, R<:AbstractReciprocalPoint{3}, S<:AbstractSite{LS,3}, B<:AbstractBond{LB,N},U<:AbstractUnitcell{S,B}}

    # return the specific type
    return newReciprocalPoint(
        R,
        [0.0, 0.0, 0.0],
        "Gamma",
        "\\Gamma"
    )
end

#########################
#
#	K POINT (and K')
#
#########################

# K point (Based on BZ object) in 2D
function getReciprocalPoint(
            :: Type{R},
            unitcell   :: U,
            identifier :: Val{:K},
            instance   :: Val{I}
        ) :: R where {D,I,N,LS,LB, R<:AbstractReciprocalPoint{2}, S<:AbstractSite{LS,2}, B<:AbstractBond{LB,N},U<:AbstractUnitcell{S,B}}

    # get the bz
    bz = getBrillouinZone(getReciprocalUnitcell(unitcell))

    # return the specific type
    return newReciprocalPoint(
        R,
        corners(bz)[faces(bz)[1][mod(I,length(corners(bz))) + 1]],
        "K",
        "K"
    )
end

# K' point (Based on BZ object) in 2D
function getReciprocalPoint(
            :: Type{R},
            unitcell   :: U,
            identifier :: Val{:Kp},
            instance   :: Val{I}
        ) :: R where {D,I,N,LS,LB, R<:AbstractReciprocalPoint{2}, S<:AbstractSite{LS,2}, B<:AbstractBond{LB,N},U<:AbstractUnitcell{S,B}}

    # get the bz
    bz = getBrillouinZone(getReciprocalUnitcell(unitcell))

    # return the specific type
    return newReciprocalPoint(
        R,
        corners(bz)[faces(bz)[1][mod(I+1,length(corners(bz))) + 1]],
        "K\'",
        "K\'"
    )
end

# K point (Based on BZ object) in 3D
function getReciprocalPoint(
            :: Type{R},
            unitcell   :: U,
            identifier :: Val{:K},
            instance   :: Val{I}
        ) :: R where {D,I,N,LS,LB, R<:AbstractReciprocalPoint{3}, S<:AbstractSite{LS,3}, B<:AbstractBond{LB,N},U<:AbstractUnitcell{S,B}}

    # get the bz
    bz = getBrillouinZone(getReciprocalUnitcell(unitcell))

    # get the faceindex
    faceindex = mod(div(I, 10), length(faces(bz))) + 1
    face = faces(bz)[faceindex]
    # first and second index
    cornerindex = mod(I, 10)
    i = face[mod(cornerindex, length(face))+1]

    # return the specific type
    return newReciprocalPoint(
        R,
        corners(bz)[i],
        "K",
        "K"
    )
end

# K' point (Based on BZ object) in 3D
function getReciprocalPoint(
            :: Type{R},
            unitcell   :: U,
            identifier :: Val{:Kp},
            instance   :: Val{I}
        ) :: R where {D,I,N,LS,LB, R<:AbstractReciprocalPoint{3}, S<:AbstractSite{LS,3}, B<:AbstractBond{LB,N},U<:AbstractUnitcell{S,B}}

    # get the bz
    bz = getBrillouinZone(getReciprocalUnitcell(unitcell))

    # get the faceindex
    faceindex = mod(div(I, 10), length(faces(bz))) + 1
    face = faces(bz)[faceindex]
    # first and second index
    cornerindex = mod(I, 10)
    i = face[mod(cornerindex+1, length(face))+1]

    # return the specific type
    return newReciprocalPoint(
        R,
        corners(bz)[i],
        "K\'",
        "K\'"
    )
end

#########################
#
#	M POINT (and M')
#
#########################

# M point (Based on BZ object) in 2D
function getReciprocalPoint(
            :: Type{R},
            unitcell   :: U,
            identifier :: Val{:M},
            instance   :: Val{I}
        ) :: R where {D,I,N,LS,LB, R<:AbstractReciprocalPoint{2}, S<:AbstractSite{LS,2}, B<:AbstractBond{LB,N},U<:AbstractUnitcell{S,B}}

    # get the bz
    bz = getBrillouinZone(getReciprocalUnitcell(unitcell))

    # first and second index
    i1 = faces(bz)[1][mod(I-1,length(faces(bz)[1])) + 1]
    i2 = faces(bz)[1][mod(I,length(faces(bz)[1])) + 1]
    # find the first and next corner
    c1 = corners(bz)[i1]
    c2 = corners(bz)[i2]

    # return the specific type
    return newReciprocalPoint(
        R,
        (c1.+c2)./2,
        "M",
        "M"
    )
end

# M' point (Based on BZ object) in 2D
function getReciprocalPoint(
            :: Type{R},
            unitcell   :: U,
            identifier :: Val{:Mp},
            instance   :: Val{I}
        ) :: R where {D,I,N,LS,LB, R<:AbstractReciprocalPoint{2}, S<:AbstractSite{LS,2}, B<:AbstractBond{LB,N},U<:AbstractUnitcell{S,B}}

    # get the bz
    bz = getBrillouinZone(getReciprocalUnitcell(unitcell))

    # first and second index
    i1 = faces(bz)[1][mod(I,length(faces(bz)[1])) + 1]
    i2 = faces(bz)[1][mod(I+1,length(faces(bz)[1])) + 1]
    # find the first and next corner
    c1 = corners(bz)[i1]
    c2 = corners(bz)[i2]

    # return the specific type
    return newReciprocalPoint(
        R,
        (c1.+c2)./2,
        "M\'",
        "M\'"
    )
end

# M point (Based on BZ object) in 3D
function getReciprocalPoint(
            :: Type{R},
            unitcell   :: U,
            identifier :: Val{:M},
            instance   :: Val{I}
        ) :: R where {D,I,N,LS,LB, R<:AbstractReciprocalPoint{3}, S<:AbstractSite{LS,3}, B<:AbstractBond{LB,N},U<:AbstractUnitcell{S,B}}

    # get the bz
    bz = getBrillouinZone(getReciprocalUnitcell(unitcell))

    # get the faceindex
    faceindex = mod(div(I, 10), length(faces(bz))) + 1
    face = faces(bz)[faceindex]
    # first and second index
    cornerindex = mod(I, 10)
    i1 = face[mod(cornerindex-1, length(face)) + 1]
    i2 = face[mod(cornerindex,   length(face)) + 1]
    # find the first and next corner
    c1 = corners(bz)[i1]
    c2 = corners(bz)[i2]

    # return the specific type
    return newReciprocalPoint(
        R,
        (c1.+c2)./2,
        "M",
        "M"
    )
end

# M' point (Based on BZ object) in 3D
function getReciprocalPoint(
            :: Type{R},
            unitcell   :: U,
            identifier :: Val{:Mp},
            instance   :: Val{I}
        ) :: R where {D,I,N,LS,LB, R<:AbstractReciprocalPoint{3}, S<:AbstractSite{LS,3}, B<:AbstractBond{LB,N},U<:AbstractUnitcell{S,B}}

    # get the bz
    bz = getBrillouinZone(getReciprocalUnitcell(unitcell))

    # get the faceindex
    faceindex = mod(div(I, 10), length(faces(bz))) + 1
    face = faces(bz)[faceindex]
    # first and second index
    cornerindex = mod(I, 10)
    i1 = face[mod(cornerindex, length(face)) + 1]
    i2 = face[mod(cornerindex+1,   length(face)) + 1]
    # find the first and next corner
    c1 = corners(bz)[i1]
    c2 = corners(bz)[i2]

    # return the specific type
    return newReciprocalPoint(
        R,
        (c1.+c2)./2,
        "M\'",
        "M\'"
    )
end

#########################
#
#	X POINT (center of face)
#
#########################

# X point (Based on BZ object) in 3D
function getReciprocalPoint(
            :: Type{R},
            unitcell   :: U,
            identifier :: Val{:X},
            instance   :: Val{I}
        ) :: R where {D,I,N,LS,LB, R<:AbstractReciprocalPoint{3}, S<:AbstractSite{LS,3}, B<:AbstractBond{LB,N},U<:AbstractUnitcell{S,B}}

    # get the bz
    bz = getBrillouinZone(getReciprocalUnitcell(unitcell))

    # get the faceindex
    faceindex = mod(I-1, length(faces(bz))) + 1
    face = faces(bz)[faceindex]
    # find the center
    center = [0.0, 0.0, 0.0]
    for c in face
        center .+= corners(bz)[c]
    end
    center ./= length(face)

    # return the specific type
    return newReciprocalPoint(
        R,
        center,
        "X",
        "X"
    )
end


text:
using Test
using LoopVectorization

const var"@_avx" = LoopVectorization.var"@_turbo"

using LinearAlgebra
function clenshaw(x, coeff)
    len_c = length(coeff)
    tmp = zero(x)
    ret = zero(x)
    for i in len_c:-1:2
        ret     = muladd(x,2tmp,coeff[i]-ret)
        ret,tmp = tmp,ret
    end
    ret = muladd(x,tmp,coeff[1]-ret)
    return ret
end

"""
Causes `check_args` to fail.
"""
struct FallbackArrayWrapper{T,N} <: AbstractArray{T,N}
    data::Array{T,N}
end
Base.size(A::FallbackArrayWrapper) = size(A.data)
Base.@propagate_inbounds Base.getindex(A::FallbackArrayWrapper, i::Vararg{Int, N}) where {N} = getindex(A.data, i...)
Base.@propagate_inbounds Base.setindex!(A::FallbackArrayWrapper, v, i::Vararg{Int, N}) where {N} = setindex!(A.data, v, i...)
Base.IndexStyle(::Type{<:FallbackArrayWrapper}) = IndexLinear()

const RUN_SLOW_TESTS = LoopVectorization.register_count() ≤ 16 || !parse(Bool, get(ENV, "GITHUB_ACTIONS", "false"))


text:
using SparseArrays
using Distributions
using StatsBase
using GVFHordes
import GVFHordes: AbstractHorde
using MinimalRLCore

mutable struct Agent{#IR<:IntrinsicReward,
                     # H,
                     # BH,
                     # DL<:Learner,
                     # BL<:Learner,
                     O,
                     Φ,
                     SC} <: AbstractAgent

    behaviour_learner::Learner
    behaviour_demons::Union{AbstractHorde,Nothing} # Round Robin learners have behaviour demons set to nothing
    behaviour_gamma::Float64

    demons::AbstractHorde
    demon_learner::Learner

    last_obs::O
    last_state::Φ

    last_action::Int
    num_actions::Int

    intrinsic_reward::IntrinsicReward
    state_constructor::SC
    use_external_reward::Bool
    exploration::ExplorationStrategy

    random_first_action::Bool

end

function Agent(horde,
               behaviour_feature_size::Int,
               behaviour_learner,
               behaviour_horde,
               behaviour_gamma,
               demon_learner,
               observation_size::Int,
               num_actions::Int,
               intrinsic_reward_type,
               state_constructor,
               use_external_reward,
               exploration_strategy,
               random_first_action)

    behaviour_weight_dims = (num_actions, behaviour_feature_size)

    intrinsic_reward = if intrinsic_reward_type == "weight_change"
        #TODO: The intrinsic reward is defined by how the components of the agent are put together. For example, an intrinsic reward
        # could be the model error, which would then require different components that are assembled in the agent
        # Not sure how the construction of intrinisic reward could be abstracted out of this constructor
        WeightChange(demon_learner)
    elseif intrinsic_reward_type == "no_reward"
        NoReward()
    else
        throw(ArgumentError("Not a valid intrinsic reward"))
    end

    Agent(behaviour_learner,
          behaviour_horde,
          behaviour_gamma,

          horde,
          demon_learner,

          zeros(observation_size),
          spzeros(behaviour_feature_size),

          0, # last_action
          num_actions,

          intrinsic_reward,
          state_constructor,

          use_external_reward,
          exploration_strategy,
          random_first_action)

end

function proc_input(agent, obs)
    return agent.state_constructor(obs)
end

function get_action(agent, state, obs)
    action_probs = if agent.behaviour_learner isa OneDTMazeUtils.RoundRobinPolicy || agent.behaviour_learner isa TwoDGridWorldUtils.RoundRobinPolicy
        get_action_probs(agent.behaviour_learner, state, obs)
    elseif agent.behaviour_learner.update isa TabularRoundRobin
        get_action_probs(agent.behaviour_learner.update, state, obs)
    elseif agent.behaviour_learner isa BaselineUtils.FollowDemon
        agent.behaviour_learner(obs)
    elseif agent.behaviour_learner isa BaselineUtils.RandomDemons
        qs = agent.behaviour_learner(obs)
        agent.exploration(qs)
    else
        qs = agent.behaviour_learner(state)
        agent.exploration(qs)
    end
    action = sample(1:agent.num_actions, Weights(action_probs))
    return action, action_probs
end

function assign_horde!(agent, horde)
    agent.demons = horde
end

MinimalRLCore.end!(agent, obs, reward, is_terminal) =
    MinimalRLCore.step!(agent, obs, reward, is_terminal)

function MinimalRLCore.start!(agent::Agent, obs, args...)
    next_state = proc_input(agent, obs)
    #Always exploring starts
    step!(agent.exploration)
    next_action = if agent.random_first_action
        sample(1:agent.num_actions, Weights(ones(agent.num_actions)))
    else
        a, probs = get_action(agent,next_state, obs)
        a
    end

    # NOTE: IS THIS RIGHT, seems wierd!?
    _, discounts, _ = get(agent.demons; state_t = obs, action_t = next_action, state_tp1 = obs, action_tp1 = next_action)

    agent.last_state = next_state
    agent.last_action = next_action
    agent.last_obs = obs
    zero_eligibility_traces!(agent.demon_learner)

    return next_action
end

function MinimalRLCore.step!(agent::Agent, obs, r, is_terminal, args...)
    next_state = proc_input(agent, obs)
    step!(agent.exploration)
    next_action, next_action_probs = get_action(agent, next_state, obs)

    update_demons!(agent,
                   agent.last_obs,
                   obs,
                   agent.last_state,
                   agent.last_action,
                   next_state,
                   next_action,
                   is_terminal,
                   r)

    #get intrinssic reward
    r_int = update_reward!(agent.intrinsic_reward, agent)
    total_reward = agent.use_external_reward ? r_int + r : r_int
    update_behaviour!(agent,
                      agent.last_obs,
                      obs,
                      agent.last_state,
                      agent.last_action,
                      next_state,
                      next_action,
                      is_terminal,
                      total_reward)

    agent.last_state = next_state
    agent.last_action = next_action
    agent.last_obs = obs
    return next_action
end

get_behaviour_pis(agent::Agent, state, obs) =
    get_action(agent, state, obs)[2]

function μ_π(agent::Agent, obs)
    state = proc_input(agent, obs)
    get_behaviour_pis(agent,state,obs)
end
function update_demons!(agent,obs, next_obs, state, action, next_state, next_action, is_terminal, env_reward)

    update!(agent.demon_learner,
            agent.demons,
            obs,
            next_obs,
            state,
            action,
            next_state,
            next_action,
            is_terminal,
            (state, obs) -> get_behaviour_pis(agent, state, obs),
            env_reward)

end

function update_behaviour!(agent, obs, next_obs, state, action, next_state, next_action, is_terminal, reward)

    # behaviour_pis = get_behaviour_pis(agent, state, obs)
    # next_behaviour_pis = get_behaviour_pis(agent, next_state, next_obs)

        #NOTE: Different call than demon updates as the reward and environment pseudotermination function
    update!(agent.behaviour_learner,
            agent.behaviour_demons,
            obs,
            next_obs,
            state,
            action,
            next_state,
            next_action,
            is_terminal,
            (state, obs) -> get_behaviour_pis(agent, state, obs),
            reward)
end

function get_demon_prediction(agent::Agent, obs, action)
    ϕ = agent.state_constructor(obs)
    agent.demon_learner(ϕ, action)
end

max_stars_repo_name: MatthewMcLeod/curiosity

max_issues_repo_name: MatthewMcLeod/curiosity

max_forks_repo_name: MatthewMcLeod/curiosity


text:
# calculate probability by bayes theorem

module ProbsBayes
    using DataFrames, CSV, FreqTables

    function main()
        data_path = joinpath(split(@__FILE__, "src")[1], "data/sensor_data_600.txt")
        df_org = CSV.read(data_path, DataFrame,
                      header=["date", "time", "ir", "lidar"],
                      delim=' ')

        hours = [Int64(floor(e/10000)) for e in df_org.time]

        # add hour array to df
        df_h = DataFrame(hour=hours)
        df_new = hcat(df_org, df_h)

        # P(t)
        freqs = freqtable(df_new, :lidar, :hour)
        probs = freqs / length(df_org.lidar)
        keys = names(probs, 2)
        p_t = [sum(probs[begin:end, Name(key)]) for key in keys]
        println("P(t) = $(p_t)")

        # P(z)
        probs_trans = transpose(probs)
        keys_trans = names(probs_trans, 2)
        p_z = [sum(probs_trans[begin:end, Name(key)]) for key in keys_trans]
        println("P(z) = $(p_z)")

        # P(t|z)
        for (i, key) in enumerate(keys_trans)
            probs_trans[begin:end, Name(key)] /= p_z[i]
        end
        cond_t_z = probs_trans

        p_z_630 = p_z[findfirst(keys_trans .== 630)]
        p_t_13 = p_t[findfirst(keys .== 13)]
        p_t_13_z_630 = cond_t_z[Name(13), Name(630)]

        bayes = p_t_13_z_630 * p_z_630 /p_t_13

        cond_z_t = probs / p_t[findfirst(keys .== 13)]
        answer = cond_z_t[Name(630), Name(13)]

        println("P(z=630) = $(p_z_630)")
        println("P(t=13) = $(p_t_13)")
        println("P(t=13|z=630) = $(p_t_13_z_630)")
        println("Bayes P(z=630|t=13) = $(bayes)")
        println("Answer P(z=630|t=13) = $(answer)")
    end
end

max_stars_repo_name: ShisatoYano/JuliaAutonomy

max_issues_repo_name: ShisatoYano/JuliaAutonomy

max_forks_repo_name: ShisatoYano/JuliaAutonomy


text:
using ExtensibleUnions
using Test

struct S5 end
function f1 end
@test !isextensibleunion(S5)
@test !isextensiblefunction(f1)
extensiblefunction!(f1)
@test isextensiblefunction(f1)
@test_throws ArgumentError extensiblefunction!(f1, S5)
@test !isextensibleunion(S5)
extensibleunion!(S5)
@test isextensibleunion(S5)
extensiblefunction!(f1, S5)
@test isextensiblefunction(f1)
extensiblefunction!(f1, S5)
@test isextensiblefunction(f1)
extensiblefunction!(f1, S5)
@test isextensiblefunction(f1)


text:
module NumericError

import Base: convert, promote_rule, +, -, *, /, ^

export Measure

type Measure <: Number
  x::Float64
  σ::Float64
  Measure(x::Real) = new(Float64(x), 0)
  Measure(x::Real, σ::Real) = new(Float64(x), σ)
end

Base.show(io::IO, x::Measure) = print(io, string(x.x, " ± ", x.σ))

Base.convert(::Type{Measure}, x::Real) = Measure(Float64(x), 0.0)
Base.promote_rule(::Type{Float64}, ::Type{Measure}) = Measure
Base.promote_rule(::Type{Int64}, ::Type{Measure}) = Measure

+(a::Measure, b::Measure) = Measure(a.x + b.x, sqrt(a.σ ^ 2 + b.σ ^ 2))
-(a::Measure, b::Measure) = Measure(a.x - b.x, sqrt(a.σ ^ 2 + b.σ ^ 2))
-(a::Measure) = Measure(-a.x, a.σ)

*(a::Measure, b::Measure) = begin
  x = a.x * b.x
  σ = sqrt(x ^ 2 * ((a.σ / a.x) ^ 2 + (b.σ / b.x) ^ 2))
  Measure(x, σ)
end
/(a::Measure, b::Measure) = begin
  x = a.x / b.x
  σ = sqrt(x ^ 2 * ((a.σ / a.x) ^ 2 + (b.σ / b.x) ^ 2))
  Measure(x, σ)
end

^(a::Measure, b::Float64) = begin
  x = a.x ^ b
  σ = abs(x * b * a.σ / a.x)
  Measure(x, σ)
end

Base.sqrt(a::Measure) = a ^ .5

end  # module NumericError

using NumericError

# x1 = 100 ± 1.1
# y1 = 50 ± 1.2
# x2 = 200 ± 2.2
# y2 = 100 ± 2.3

x1 = Measure(100, 1.1)
y1 = Measure(50,  1.2)
x2 = Measure(200, 2.2)
y2 = Measure(100, 2.3)

d = sqrt((x1 - x2) ^ 2 + (y1 - y2) ^ 2)

@show x1 y1 x2 y2 sqrt((x1 - x2) ^ 2 + (y1 - y2) ^ 2)

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
struct Graph{N}
    directchildren::Vector{Vector{Int64}}
    loopchildren::Vector{Vector{Int64}}
    ineqchildren::Vector{Vector{Int64}}
    successors::Vector{Vector{Int64}} # contains direct and loop children
    predecessors::Vector{Vector{Int64}}
    connections::Vector{Vector{Int64}}

    dfslist::SVector{N,Int64}
    rdfslist::SVector{N,Int64}

    dict::UnitDict{Base.OneTo{Int64},Int64}
    rdict::UnitDict{Base.OneTo{Int64},Int64}

    function Graph(origin::Origin,bodies::Vector{<:Body},
        eqconstraints::Vector{<:EqualityConstraint},ineqconstraints::Vector{<:InequalityConstraint})

        oid = origin.id
        adjacency, dict = adjacencyMatrix(eqconstraints, bodies)
        dfsgraph, dfslist, loops = dfs(adjacency, dict, oid)
        pat = pattern(dfsgraph, dict, loops)
        fil, originals = fillins(dfsgraph, pat, dict, loops)

        adjacency = deleteat(adjacency, dict[oid])
        dfsgraph = deleteat(dfsgraph, dict[oid])
        pat = deleteat(pat, dict[oid])
        fil = deleteat(fil, dict[oid])
        originals = deleteat(originals, dict[oid])
        dfslist = StaticArrays.deleteat(dfslist, length(dfslist))

        for (id, ind) in dict
            ind > dict[oid] && (dict[id] = ind - 1)
        end
        pop!(dict, oid)
        rdict = Dict(ind => id for (id, ind) in dict)

        for constraint in eqconstraints
            constraint.pid == oid && (constraint.pid = nothing)
        end

        N = length(dict)

        adjacency = convert(Vector{SVector{N,Bool}}, adjacency)
        dfsgraph = convert(Vector{SVector{N,Bool}}, dfsgraph)
        pat = convert(Vector{SVector{N,Bool}}, pat)
        fil = convert(Vector{SVector{N,Bool}}, fil)
        originals = convert(Vector{SVector{N,Bool}}, originals)

        dirs = directchildren(dfslist, originals, dict)
        loos = loopchildren(dfslist, fil, dict)
        ineqs = ineqchildren(dfslist, bodies, ineqconstraints, dict)
        sucs = successors(dfslist, pat, dict)
        preds = predecessors(dfslist, pat, dict)
        cons = connections(dfslist, adjacency, dict)

        dict = UnitDict(dict)
        rdict = UnitDict(rdict)

        new{N}(dirs, loos, ineqs, sucs, preds, cons, dfslist, reverse(dfslist), dict, rdict)
    end
end

function adjacencyMatrix(eqconstraints::Vector{<:EqualityConstraint}, bodies::Vector{<:Body})
    A = zeros(Bool, 0, 0)
    dict = Dict{Int64,Int64}()
    n = 0

    for constraint in eqconstraints
        cid = constraint.id
        A = [A zeros(Bool, n, 1); zeros(Bool, 1, n) zero(Bool)]
        dict[cid] = n += 1
        for bodyid in unique([constraint.pid;constraint.bodyids])
            if !haskey(dict, bodyid)
                A = [A zeros(Bool, n, 1); zeros(Bool, 1, n) zero(Bool)]
                dict[bodyid] = n += 1
            end
            A[dict[cid],dict[bodyid]] = true
        end
    end
    for body in bodies # add unconnected bodies
        if !haskey(dict, body.id)
            A = [A zeros(Bool, n, 1); zeros(Bool, 1, n) zero(Bool)]
            dict[body.id] = n += 1
        end
    end

    A = A .| A'
    convert(Matrix{Bool}, A), dict
end

function dfs(adjacency::Matrix, dict::Dict, originid::Int64)
    N = size(adjacency)[1]
    dfsgraph = zeros(Bool, N, N)
    dfslist = zeros(Int64, N)
    visited = zeros(Bool, N)
    loops = Vector{Int64}[]
    index = N

    dfslist[index] = originid
    visited[dict[originid]] = true
    dfs!(adjacency, dfsgraph, dict, dfslist, visited, loops, N, originid, -1)
    loops = loops[sortperm(sort.(loops))[1:2:length(loops)]] # removes double entries of loop connections and keeps the first found pair

    return dfsgraph, convert(SVector{N}, dfslist), loops
end

function dfs!(A::Matrix, Adfs::Matrix, dict::Dict, list::Vector, visited::Vector, loops::Vector{Vector{Int64}}, index::Int64, currentid::Int64, parentid::Int64)
    i = dict[currentid]
    for (childid, j) in dict
        if A[i,j] && parentid != childid # connection from i to j in adjacency && not a direct connection back to the parent
            if visited[j]
                push!(loops, [childid,currentid]) # childid is actually a predecessor of currentid since it's a loop
            else
                index -= 1
                list[index] = childid
                visited[j] = true
                Adfs[i,j] = true
                index = dfs!(A, Adfs, dict, list, visited, loops, index, childid, currentid)
            end
        end
    end
    return index
end

function pattern(dfsgraph::Matrix, dict::Dict, loops::Vector{Vector{Int64}})
    pat = deepcopy(dfsgraph)

    for loop in loops # loop = [index of starting constraint, index of last body in loop]
        startid = loop[1]
        endid = loop[2]
        currentid = endid
        nextid = parent(dfsgraph, dict, currentid)
        while nextid != startid
            pat[dict[startid],dict[currentid]] = true
            currentid = nextid
            nextid = parent(dfsgraph, dict, nextid)
        end
    end

    return pat
end

function fillins(dfsgraph::Matrix, pattern::Matrix, dict::Dict, loops::Vector{Vector{Int64}})
    fil = deepcopy(dfsgraph .⊻ pattern) # xor so only fillins remain (+ to be removed loop closure since this is a child)
    originals = deepcopy(dfsgraph)

    for loop in loops # loop = [index of starting constraint, index of last body in loop]
        startid = loop[1]
        endid = loop[2]
        fil[dict[startid],dict[endid]] = false
        originals[dict[startid],dict[endid]] = true
    end

    return convert(Matrix{Bool}, fil), convert(Matrix{Bool}, originals)
end

function parent(dfsgraph::Matrix, dict::Dict, childid::Int64)
    j = dict[childid]
    for (parentid, i) in dict
        dfsgraph[i,j] && (return parentid)
    end
    return -1
end

# this is done in order!
function successors(dfslist, pattern, dict::Dict)
    N = length(dfslist)
    sucs = [Int64[] for i = 1:N]
    for i = 1:N
        for cid in dfslist
            pattern[i][dict[cid]] && push!(sucs[i], cid)
        end
    end

    return sucs
end

# this is done in order!
function directchildren(dfslist, dfsgraph, dict::Dict)
    N = length(dfslist)
    dirs = [Int64[] for i = 1:N]
    for i = 1:N
        for cid in dfslist
            dfsgraph[i][dict[cid]] && push!(dirs[i], cid)
        end
    end

    return dirs
end

# this is done in order!
function loopchildren(dfslist, fillins, dict::Dict)
    N = length(dfslist)
    loos = [Int64[] for i = 1:N]
    for i = 1:N
        for cid in dfslist
            fillins[i][dict[cid]] && push!(loos[i], cid)
        end
    end

    return loos
end

function ineqchildren(dfslist, bodies, ineqconstraints, dict::Dict)
    N = length(dfslist)
    ineqs = [Int64[] for i = 1:N]
    for body in bodies
        for c in ineqconstraints
            c.pid == body.id && push!(ineqs[dict[body.id]], c.id)
        end
    end

    return ineqs
end

# this is done in reverse order (but this is not really important for predecessors)
function predecessors(dfslist, pattern, dict::Dict)
    N = length(dfslist)
    preds = [Int64[] for i = 1:N]
    for i = 1:N
        for cid in reverse(dfslist)
            pattern[dict[cid]][i] && push!(preds[i], cid)
        end
    end

    return preds
end

# this is done in order (but this is not really important for connections)
function connections(dfslist, adjacency, dict::Dict)
    N = length(dfslist)
    cons = [Int64[] for i = 1:N]
    for i = 1:N
        for cid in dfslist
            adjacency[i][dict[cid]] && push!(cons[i], cid)
        end
    end

    return cons
end

@inline directchildren(graph, id::Int64) = graph.directchildren[graph.dict[id]]
@inline loopchildren(graph, id::Int64) = graph.loopchildren[graph.dict[id]]
@inline ineqchildren(graph, id::Int64) = graph.ineqchildren[graph.dict[id]]
@inline successors(graph, id::Int64) = graph.successors[graph.dict[id]]
@inline predecessors(graph, id::Int64) = graph.predecessors[graph.dict[id]]
@inline connections(graph, id::Int64) = graph.connections[graph.dict[id]]

@inline function hassuccessor(graph::Graph{N}, id, cid) where N
    for val in graph.successors[graph.dict[id]]
        val == cid && (return true)
    end
    return false
end

@inline function haspredecessor(graph::Graph{N}, id, pid) where N
    for val in graph.predecessor[graph.dict[id]]
        val == pid && (return true)
    end
    return false
end

@inline function hasdirectchild(graph::Graph{N}, id, cid) where N
    for val in graph.directchildren[graph.dict[id]]
        val == cid && (return true)
    end
    return false
end


text:
using Libdl, LLVM

function show_inttoptr(mod)
    for fun in LLVM.functions(mod),
        blk in LLVM.blocks(fun),
        instr in LLVM.instructions(blk)

        s = string(instr)
        if occursin("inttoptr", s) && occursin(r"[0-9]{8,30}", s)
            println(LLVM.name(fun), "  ---------------------------")
            @show instr
            println()
        end
    end
end

"""
Compiles function call provided and calls it with `ccall` using the shared library that was created.
"""
macro jlrun(e)

    fun = e.args[1]
    efun = esc(fun)
    args = length(e.args) > 1 ? e.args[2:end] : Any[]
    libpath = abspath("test.o")
    dylibpath = abspath("test.so")
    tt = Tuple{(typeof(eval(a)) for a in args)...}
    if length(e.args) > 1
        ct = code_typed(Base.eval(__module__, fun), tt)
    else
        ct = code_typed(Base.eval(__module__, fun))
    end
    rettype = ct[1][2]
    bindir = joinpath(dirname(Sys.BINDIR), "tools")
    libdir = joinpath(dirname(Sys.BINDIR), "lib")

    runCommand = :(run(
        $(`$shellcmd -shared -fPIC -o test.so -L$libdir test.o -ljulia`),
        wait = true,
    ))

    quote
        m = irgen($efun, $tt)
        # m = irgen($efun, $tt, overdub = false)
        LLVM.verify(m)
        # show_inttoptr(m)
        write(m, "test.bc")
        write_object(m, "test.o")
        $runCommand
        dylib = Libdl.dlopen($dylibpath)
        ccall(Libdl.dlsym(dylib, "jl_init_globals"), Cvoid, ())
        res = ccall(
            Libdl.dlsym(dylib, $(Meta.quot(fun))),
            $rettype,
            ($((typeof(eval(a)) for a in args)...),),
            $(eval.(args)...),
        )
        Libdl.dlclose(dylib)
        res
    end
end


text:
#! /usr/bin/julia

# Rosetta Code, Hostname

for i in [0, 5, 50, 9000]
    println(i, " => ", bin(i))
end

max_stars_repo_name: MichaeLeroy/rosetta-code

max_issues_repo_name: MichaeLeroy/rosetta-code

max_forks_repo_name: MichaeLeroy/rosetta-code


text:
#=
NOTES:
- Proposal will be updated after the end of each cycle in Warmup and Adaptionˢˡᵒʷ.
- Proposal is kept constant in Adaptionᶠᵃˢᵗ, where final stepsize is tuned.
- Proposal is kept constant in Exploration.
=#
############################################################################################
"""
$(TYPEDEF)

Choice for Posterior Covariance Matrix adaption.

# Fields
$(TYPEDFIELDS)
"""
abstract type MatrixMetric end
struct MDense <: MatrixMetric end
struct MDiagonal <: MatrixMetric end
struct MUnit <: MatrixMetric end

############################################################################################
"""
$(TYPEDEF)

Default configuration for posterior Covariance Matrix adaption.

# Fields
$(TYPEDFIELDS)
"""
struct ConfigProposal{A<:BaytesCore.UpdateBool, M<:MatrixMetric}
    "Boolean if Posterior covariance for proposal steps is adapted."
    proposaladaption::A
    "Covariance estimate metric: MDense(), MDiagonal(), MUnit()"
    metric::M
    "Shrinkage parameter towards Diagonal Matrix with equal variance"
    shrinkage::Float64
    function ConfigProposal(;
        proposaladaption = BaytesCore.UpdateTrue(),
        metric = MDiagonal(),
        shrinkage = 0.05
        )
        @argcheck 0.0 <= shrinkage <= 1.0 "Shrinkage not bounded between 0 and 1"
        return new{typeof(proposaladaption), typeof(metric)}(
            proposaladaption, metric, shrinkage
            )
    end
end

############################################################################################
#Wrapper to generate Covariance buffers
function init(type::Type{T}, metric::MDense, param_length::Int64) where {T<:Real}
    Σ = LinearAlgebra.Symmetric(zeros(T, param_length, param_length) + LinearAlgebra.I)
    Σ⁻¹ᶜʰᵒˡ = LinearAlgebra.cholesky(LinearAlgebra.inv(Σ)).L
    return Σ, Σ⁻¹ᶜʰᵒˡ
end
function init(
    type::Type{T}, metric::M, param_length::Int64
) where {T<:Real,M<:Union{MDiagonal,MUnit}}
    Σ = LinearAlgebra.Diagonal(fill(T(1.0), param_length))
    Σ⁻¹ᶜʰᵒˡ = Diagonal(.√LinearAlgebra.inv.(LinearAlgebra.diag(Σ)))
    return Σ, Σ⁻¹ᶜʰᵒˡ
end

############################################################################################
"""
$(TYPEDEF)

Mass and Covariance Matrix specification for MCMC sampler, relevant for Euclidean Metric.

# Fields
$(TYPEDFIELDS)
"""
struct MatrixTune{A<:MatrixMetric}
    "Dense, Diagonal or Unit Mass Matrix."
    metric::A
    "Shrinkage parameter for Covariance estimation."
    shrinkage::Float64
    function MatrixTune(metric::A, shrinkage::Float64) where {A<:MatrixMetric}
        return new{A}(metric, shrinkage)
    end
end

############################################################################################
"""
$(SIGNATURES)
Calculate regularized Covariance Matrix.

# Examples
```julia
```

"""
function get_Σ(Σ::D, shrinkage::T) where {D<:Union{Diagonal,Symmetric},T<:Real}
    return (1 - shrinkage) * Σ + shrinkage * UniformScaling(max(1e-3, median(diag(Σ))))
end

############################################################################################
"""
$(TYPEDEF)

Proposal distribution container.

# Fields
$(TYPEDFIELDS)
"""
mutable struct Proposal{
    A<:BaytesCore.UpdateBool,T<:AbstractFloat,P<:AbstractMatrix,C<:AbstractMatrix,M<:MatrixTune
}
    "Check if adaption true in current iteration"
    adaption::A
    "θᵤ samples in current MCMC Phase, used for Σ estimation"
    chain::Matrix{T}
    "Posterior Covariance estimate"
    Σ::P
    "Cholesky decomposition of Inverse Posterior Covariance matrix Σ, s.t. Σ⁻¹ᶜʰᵒˡ*Σ⁻¹ᶜʰᵒˡ'=Σ⁻¹. Used to generate random draws in HMC/NUTS sampler."
    Σ⁻¹ᶜʰᵒˡ::C
    "Tuning parameter for Σ estimation"
    matrixtune::M
    function Proposal(
        type::Type{T}, adaption::A, tune::M, param_length::Int64, warmup_length::Int64
    ) where {T<:AbstractFloat,A<:BaytesCore.UpdateBool,M<:MatrixTune}
        Σ, Σ⁻¹ᶜʰᵒˡ = init(T, tune.metric, param_length)
        chain = zeros(T, param_length, warmup_length)
        return new{typeof(adaption),eltype(chain),typeof(Σ),typeof(Σ⁻¹ᶜʰᵒˡ),M}(
            adaption, chain, Σ, Σ⁻¹ᶜʰᵒˡ, tune
        )
    end
end

############################################################################################
"""
$(SIGNATURES)
Calculate regularized Covariance Matrix and Cholesky decomposition.

# Examples
```julia
```

"""
function update!(proposal::P, metric::MDense) where {P<:Proposal}
    @unpack Σ, Σ⁻¹ᶜʰᵒˡ, matrixtune, chain = proposal
    @unpack metric, shrinkage = matrixtune
    ## Compute new estimate
    Σ .= get_Σ(LinearAlgebra.Symmetric(Statistics.cov(chain')), shrinkage)
    Σ⁻¹ᶜʰᵒˡ .= LinearAlgebra.cholesky(LinearAlgebra.inv(Σ)).L
    ## Pack container and return matrices
    @pack! proposal = Σ, Σ⁻¹ᶜʰᵒˡ
    return nothing
end
function update!(proposal::P, metric::MDiagonal) where {P<:Proposal}
    @unpack Σ, Σ⁻¹ᶜʰᵒˡ, matrixtune, chain = proposal
    @unpack metric, shrinkage = matrixtune
    ## Compute new estimate
    Σ .= get_Σ(LinearAlgebra.Diagonal(vec(var(chain; dims=2))), shrinkage)
    Σ⁻¹ᶜʰᵒˡ .= Diagonal(.√LinearAlgebra.inv.(LinearAlgebra.diag(Σ)))
    ## Pack container and return matrices
    @pack! proposal = Σ, Σ⁻¹ᶜʰᵒˡ
    return nothing
end
function update!(proposal::P, metric::MUnit) where {P<:Proposal}
    return nothing
end
update!(proposal::P) where {P<:Proposal} = update!(proposal, proposal.matrixtune.metric)

############################################################################################
"""
$(SIGNATURES)
Assign new chain buffer with dedicated size.

# Examples
```julia
```

"""
function chain!(
    proposal::P, phasename::N, Niterations::Integer
) where {P<:Proposal,N<:Union{Warmup,Adaptionˢˡᵒʷ}}
    proposal.chain = zeros(eltype(proposal.chain), size(proposal.chain, 1), Niterations)
    return nothing
end
function chain!(
    proposal::P, phasename::N, Niterations::Integer
) where {P<:Proposal,N<:Union{Adaptionᶠᵃˢᵗ}}#, Exploration}}
    #!NOTE: Do not create a new chain for Adaptionᶠᵃˢᵗ (before Exploration cycle), as in this phase Σ is kept constant. (But still have to calculate get_Σ!(proposal) for current cycle)
    return nothing
end

############################################################################################
"""
$(SIGNATURES)
Update Proposal with new parameter θᵤ.

# Examples
```julia
```

"""
function update!(
    proposal::P,
    phaseupdate::Val{true},
    phasename::N,
    phase::PhaseTune,
    θᵤ::AbstractVector{T},
) where {P<:Proposal,N<:Union{Warmup,Adaptionˢˡᵒʷ,Adaptionᶠᵃˢᵗ},T<:Real}
    #NOTE!: Proposal will be updated at the beginning of Adaptionˢˡᵒʷ and Adaptionᶠᵃˢᵗ, but not at the beginning of Exploration (from Adaptionᶠᵃˢᵗ to Exploration, only local tuning parameter will be updated)
    ## Add current sample to chain
    @inbounds for iter in eachindex(θᵤ)
        proposal.chain[iter, phase.iter.current] = θᵤ[iter]
    end
    ## Estimate new covariance matrix
    update!(proposal)
    ## Preallocate new buffer for chain in next phase
    chain!(proposal, phasename, phase.iterations[phase.counter.current])
    return nothing
end
function update!(
    proposal::P,
    phaseupdate::Val{true},
    phasename::Exploration,
    phase::PhaseTune,
    θᵤ::AbstractVector{T},
) where {P<:Proposal,T<:Real}
    return nothing
end

function update!(
    proposal::P,
    phaseupdate::Val{false},
    phasename::N,
    phase::PhaseTune,
    θᵤ::AbstractVector{T},
) where {P<:Proposal,N<:Union{Warmup,Adaptionˢˡᵒʷ},T<:Real}
    ## Add current sample to chain
    @inbounds for iter in eachindex(θᵤ)
        proposal.chain[iter, phase.iter.current] = θᵤ[iter]
    end
    return nothing
end
function update!(
    proposal::P,
    phaseupdate::Val{false},
    phasename::N,
    phase::PhaseTune,
    θᵤ::AbstractVector{T},
) where {P<:Proposal,N<:Union{Adaptionᶠᵃˢᵗ,Exploration},T<:Real}
    #NOTE!: chain will be updated during Warmup and each Adaptionˢˡᵒʷ cycles, so can be used to estimate Σ for each Adaptionˢˡᵒʷ cycle and Adaptionᶠᵃˢᵗ cycle.
    #!NOTE: NO need to record during Adaptionᶠᵃˢᵗ, as during Exploration no Σ update.
    return nothing
end

############################################################################################
"""
$(SIGNATURES)
Update Proposal according to current tuning phase.

# Examples
```julia
```

"""
function update!(
    proposal::P, proposalupdate::BaytesCore.UpdateFalse, θᵤ, phasename, phase
) where {P<:Proposal}
    return nothing
end
function update!(
    proposal::P,
    proposalupdate::BaytesCore.UpdateTrue,
    θᵤ::AbstractVector{T},
    phasename::N,
    phase::PhaseTune,
) where {P<:Proposal,T<:Real,N<:Union{Warmup,Adaptionˢˡᵒʷ,Adaptionᶠᵃˢᵗ,Exploration}}
    ## Update Proposal chain and covariance estimate
    update!(proposal, Val(phase.update.current), phasename, phase, θᵤ)
    return nothing
end
function update!(
    proposal::P, θᵤ::AbstractVector{T}, phase::PhaseTune
) where {P<:Proposal,T<:Real}
    return update!(
        proposal, proposal.adaption, θᵤ, phase.name[phase.counter.current], phase
    )
end

############################################################################################
# Export
export MatrixMetric, MDense, MDiagonal, MUnit, ConfigProposal, MatrixTune, Proposal, update!


text:
using Images, FileIO, Colors, FixedPointNumbers, Statistics
# using ColorBrewer
#
# definition of system:
# x_{n+1} = sin(a y_n) + c cos(a x_n)
# y_{n+1} = sin(b x_n) + d cos(b y_n)

#size of the rendered canvas in pixels. it is always a square canvas
CANVAS_SIZE = 1500
#parameters governing the iterations
A=2.2
B=1.8
C=1.2
D=1.3
#number of points to render.
NUM_POINTS=Int(1e9)
#initial value
X_INIT=0
Y_INIT=0
#proportion of the canvas to be kept as a margin
MARGIN=0.025
#file path to save to
OUT_FILE_NAME="./image_6.jpg"
#occupancy root. smaller values lead to a more even image. higher values
#increase the rendered contrast. smaller values lead to an even image.
OCCUPANCY_ROOT = 1/4
#what percentile of the clifford attractor should not be completely saturated
#by color?
PERCENTILE_CLIP = 0.998
#what color should the attractor be? it will always be on a white background.
#you can invert it to black with the INVERSE parameter below.  grabbing some
#some sample color palettes from the ColorBrewer package which look nice.
#COLORS=((0.4,0.761,0.647),(0.988,0.553,0.384),(0.553,0.627,0.796))
COLOR=(0.553,0.627,0.796)
#the default is the above color (in rgb) on a white background. you can invert
#all colors using this parameter.
INVERSE=true

#these functions might be called a few hundred million times, so an if statement has
#been brought outside the function definitions to try and speed things up.
if D < C
    function create_plot_scalers(c, d, margin)
        x_range = 1 + c
        y_range = 1 + d
        norm = (2 * x_range * (1/(1-2*margin)))
        other_max = 2*y_range / norm
        less_margin = (1-other_max) / 2
        (x_range, y_range, norm, margin, less_margin)
    end
    function scale_point((x, y), (x_range, y_range, norm, margin, less_margin))
        trans_x, trans_y = x + x_range, y + y_range
        trans_x /= norm
        trans_y /= norm
        trans_x += margin
        trans_y += less_margin
        (trans_x, trans_y)
    end
else
    function create_plot_scalers(c, d, margin)
        x_range = 1 + c
        y_range = 1 + d
        norm = (2 * y_range * (1/(1-2*margin)))
        other_max = 2*x_range / norm
        less_margin = (1-other_max) / 2
        (x_range, y_range, norm, margin, less_margin)
    end
    function scale_point((x, y), (x_range, y_range, norm, margin, less_margin))
        trans_x, trans_y = x + x_range, y + y_range
        trans_x /= norm
        trans_y /= norm
        trans_y += margin
        trans_x += less_margin
        (trans_x, trans_y)
    end
end

function increment_canvas!(canvas, state, canvas_size, plot_scale_params)
    plot_x, plot_y = scale_point(state, plot_scale_params)
    plot_x *= canvas_size
    plot_y *= canvas_size
    plot_x = floor(Int, plot_x)
    plot_y = floor(Int, plot_y)
    plot_x = canvas_size - plot_x + 1
    plot_y = canvas_size - plot_y + 1
    canvas[plot_y, plot_x] += 1
end

function update_state(a,b,c,d,(x, y))
    new_x = sin(a*y) + c*cos(a*x)
    new_y = sin(b*x) + d*cos(b*y)
    (new_x, new_y)
end

function clifford_attractor(a, b, c, d, num_points, canvas_size, margin, x_init, y_init, percentile_clip, occupancy_root)
    plot_scale_params = create_plot_scalers(c, d, margin)
    canvas = zeros(canvas_size, canvas_size)
    state = (x_init, y_init)
    increment_canvas!(canvas, state, canvas_size, plot_scale_params)
    for step in 2:num_points
        state = update_state(a,b,c,d,state)
        increment_canvas!(canvas, state, canvas_size, plot_scale_params)
    end
    canvas = canvas.^occupancy_root
    scale = quantile(canvas[canvas.>0], percentile_clip)
    canvas/scale
end

function pixel_color_transform(pixel,(r,g,b))
    if pixel <= 0
        output = RGB{N0f8}(0.,0.,0.)
    elseif pixel <= 1
        output = RGB{N0f8}((1-r)*pixel,(1-g)*pixel,(1-b)*pixel)
    else
        output = RGB{N0f8}((1-r),(1-g),1-b)
    end
    output
end

function render_img(data,color)
    x,y = size(data)
    canvas = zeros(RGB{N0f8},x,y)
    for i in 1:x
        for j in 1:y
            canvas[i,j] = pixel_color_transform(data[i,j], color)
        end
    end
    canvas
end

function main(a=A, b=B, c=C, d=D,
              num_points=NUM_POINTS, x_init=X_INIT, y_init=Y_INIT,
              canvas_size=CANVAS_SIZE, margin=MARGIN,
              out_file_name=OUT_FILE_NAME,
              percentile_clip=PERCENTILE_CLIP, occupancy_root=OCCUPANCY_ROOT, color=COLOR, inverse=INVERSE)
    output = clifford_attractor(a, b, c, d, num_points, canvas_size, margin, x_init, y_init, percentile_clip, occupancy_root)
    img = render_img(output,color)
    if !inverse
        img = RGB{N0f8}(1.,1.,1.) .- img
    end
    save(out_file_name, img)
end

main()


text:
using Augmentor

img, imgw, flow = gen_init(:chess, flow_args=[15]);
showflow(flow)

pl = FlipX() * FlipY() |> Zoom(0.9:0.1:1.2)

Augmentor.augment_impl(pl)

pl = Rotate(-10:101)
pl = ElasticDistortion(10, 10, 0.3, 4, 3, true)

img_processed = augment(img, pl); imgshow(img_processed)

img_processed = augment((img, positions), pl); imgshow(img_processed)
imgshow(img)

flow_processed = augment(flow, pl); showflow(flow_processed)

A = [1 2 3; 4 5 6; 7 8 9]

Juno.@enter augment(A, pl)

pl = Rotate(90)

B = augment(A, pl)
imgshow(A)
imgshow(B)

pl = ElasticDistortion(3, 3, 0.25, 2, 3, true, true)

img_processed = augment(img, plborder)
imgshow(img_processed)

imgshow(img)

mutable struct Pos
    x::Float64
    y::Float64
end

Base.:+(a::Pos, b::Pos) = Pos(a.x + b.x, a.y + b.y)
Base.:*(a::Any, b::Pos) = Pos(a * b.x, a * b.y)
Base.:*(b::Pos, a::Any) = Pos(a * b.x, a * b.y)
Base.zero(Pos) = Pos(0, 0)
Pos(a::Pos) = a

Base.show(io::IO, m::Pos) = print(io, "(", m.x, ", ", m.y, ")")

zero(Pos)

1/2 * Pos(1, 2) + 1/2 * Pos(1, 1)

function ind_to_pos(ind)
    return Pos(ind[1], ind[2])
end

positions = map(ind_to_pos, CartesianIndices(img))

function make_flow(positions::Array{Pos,2})
    fun = (ind, pos) -> pos.x - ind[1] + im * (pos.y - ind[2])
    foo = (ind, pos) -> ind[1] - pos.x + im * (ind[2] - pos.y)
    return map(foo, CartesianIndices(positions), positions) .* (-1)
end

inds = CartesianIndices(positions)
inds[1,1][1]
((ind, pos) -> ind[1] - pos.x + im * (ind[2] - pos.y))(CartesianIndex(1,1), Pos(2,2))

augflow = make_flow(end_positions)
showflow(augflow)
imgshow(img_processed)

pl = Rotate(9)
begin
    img_processed, end_positions = augment((img, positions), pl);
    augflow = make_flow(end_positions);
    imgshowflow(img_processed, augflow)
end

augflow
imgshowflow(img_processed, augflow)

img, imgw, flow = gen_init(:chess, flow_args=[20, 130]);
imgshowflow(imgw, flow)
addpoints([(50, 0)])

imgshow(imgw)

showflow(ones(3,3) + im * ones(3,3))

addpoints([(50, 0)])

imgw[54,1]

end_positions

## profiling
const to = TimerOutput();

reset_timer!();

timeit_debug_enabled()

@timeit "section" sleep(0.02)
@timeit "section2" sleep(0.1)
@timeit_debug "debug time" sleep(1)

TimerOutputs.enable_debug_timings(LAP_julia)

print_timer()

img, imgw, flow = gen_init()

u_est, source_reg = pflap(img, imgw, display=false);

showflow(flow)
showflow(u_est)
imgshow(source_reg)
imgshow(img)

@profiler pflap(img, imgw, display=false);

## single lap

function lap(img, imgw, fhs, window_size)
    classic_estim = single_lap(img, imgw, fhs, window_size)
    LAP_julia.inpaint_nans!(classic_estim)
    LAP_julia.smooth_with_gaussian(classic_estim, window_size)
    LAP_julia.smooth_with_gaussian(classic_estim, window_size)
    return classic_estim
end

fhs = 23
window = [47, 47]
u_est = lap(img, imgw, fhs, window);

showflow(u_est)

@time lap(img, imgw, fhs, window)

# profile
Profile.init()
Profile.init(10^10, 0.00001)
Juno.@profiler lap(img, imgw, fhs, window);

## test A
C = [1 2 3; 4 5 6; 7 8 19]
C[3, :]

k, l = size(C)

A = rand(2, 3, 4)

A[:, :, 1][1,2]

A[2,1,:]

C = [1 2 3; 4 5 6; 7 8 19]

function gem2d(A, b)
    for k in axes(A,1)
        for l in k+1:size(A, 2)
            ratio = A[l, k] / A[k, k]
            println(ratio)
            A[l, :] .-= ratio .* A[k, :]
            b[l] -= ratio * b[k]
        end
    end
    return A, b
end

gem2d(C, [1, 2, 3])

ratios = zeros(size(A, 3))

function gem3d!(A, b)
    ratios = zeros(size(A, 3))
    for k in axes(A,1)
        for l in k+1:size(A, 2)
            ratios .= A[l, k, :] ./ A[k, k, :]
            A[l, :, :] .-= ratios' .* A[k, :, :]
            b[l, :] .-= ratios .* b[k, :]
        end
    end
end

C = [1 2; 3 4]
C[:,1]

A = Float64.(cat([1 2; 3 4], [5 6; 7 8], [1 2; 3 4], dims=3))
b = Float64.([1 3 5; 2 4 6])

A,b

A[:, :, 1]

gem3d!(A, b)

A,b

# TODO: add and test
function tt(v, A)
    r = similar(A)
    @inbounds for j = 1:size(A,2)
        @simd for i = 1:size(A,1)
            r[i,j] = v[j] * A[i,j] # fixed a typo here!
        end
    end
    r
end

N = 3
coeffs = zeros(size(A, 3), size(A, 1))

for k in (N-1):-1:1
    coeffs[:,k] = (b[k,:])'
    for m in (k+1):(N-1)
        coeffs[:,k] = coeffs[:,k]-A[k,m,:] .* coeffs[:,m]
    end
    coeffs[:,k]=coeffs[:,k]./A[k,k,:];
end

coeffs

function back_substitution3d(A, b)
    row_count = size(A, 1)
    coeffs = zeros(size(A, 3), size(A, 1))
    for k in row_count:-1:1
        coeffs[:,k] = (b[k,:])'
        for m in (k+1):row_count
            coeffs[:,k] = coeffs[:,k]-A[k,m,:] .* coeffs[:,m]
        end
        coeffs[:,k]=coeffs[:,k]./A[k,k,:];
    end
    return coeffs
end

function multi_mat_div_gem(A, b)
    gem3d!(A, b)
    return back_substitution3d(A, b)
end

coeffs = multi_mat_div_gem(A, b)

coeffs[:, 1]

big = rand(1234, 1234)

##
using TimerOutputs
to = TimerOutput()

function foo()
   global to
   @timeit to "sleepytime" out = 1+1
   println(out)
   return nothing
end

foo()

print_timer(to)

@info "With debug timings disabled:"
TimerOutputs.disable_debug_timings(LAP_julia)
foo()

@info "With debug timings enabled:"
TimerOutputs.enable_debug_timings(LAP_julia)
foo()

## make spaghetti
using LAP_julia

img = gen_spaghetti((256,256), 0.15, 70, spread=20)
imgshow(img)

img = gen_one_spaghetti((200,200), 0.02, 10)
imgshow(img)

minimum(img)
maximum(img)

pl = Rotate(35);
img_processed = augment(img, pl);
imgshow(img_processed)

img_size = (200,200)
X = ones(img_size[1]) * collect(range(-1,1,length=img_size[2]))'
X_rot = augment(X, pl)

imgshow(X)
imgshow(X_rot)

img = rand(200,200)
pl = Either(1=>FlipX(), 1=>FlipY(), 2=>NoOp()) |>
            Rotate(0:360) |>
            ShearX(-5:5) * ShearY(-5:5) |>
            CropSize(165, 165) |>
            Zoom(1:0.05:1.2) |>
            Resize(64, 64)

img_new = augment(img, pl)


text:
using Unitful
import Unitful: ms, s, kHz, Hz
export ms, s, kHz, Hz, samples, uconvert, ustrip, inseconds, inHz, insamples

const TimeDim = Unitful.Dimensions{(Unitful.Dimension{:Time}(1//1),)}
const FreqDim = Unitful.Dimensions{(Unitful.Dimension{:Time}(-1//1),)}
const Time{N} = Quantity{N,TimeDim}
const Freq{N} = Quantity{N,FreqDim}

@dimension 𝐒 "𝐒" Sample
@refunit samples "samples" Samples 𝐒 false
const SampDim = Unitful.Dimensions{(Unitful.Dimension{:Sample}(1//1),)}
const SampleQuant{N} = Quantity{N,SampDim}

insamples{N <: Integer}(time::SampleQuant{N},rate) = ustrip(time)
insamples{N}(time::SampleQuant{N},rate) = error("Cannot use non-integer samples.")
function insamples(time,rate)
  r = inHz(rate)
  floor(Int,ustrip(inseconds(time,r)*r))
end
function insamples{N,M}(time::Time{N},rate::Freq{M})
  floor(Int,ustrip(uconvert(s,time)*uconvert(Hz,rate)))
end

inHz(x::Quantity) = uconvert(Hz,x)
inHz(typ::Type{N},x) where {N <: Number} = floor(N,ustrip(inHz(x)))*Hz
inHz(typ::Type{N},x::N) where {N <: Number} = inHz(x)
function inHz(x::Number)
  warn("Unitless value, assuming Hz. Append Hz or kHz to avoid this warning",
       " (e.g. 1kHz instead of 1).")
  x*Hz
end

inseconds(x::SampleQuant{N},R) where N = (ustrip(x) / R)*s
inseconds(x::Time) = uconvert(s,x)
inseconds(x::Quantity,R) = uconvert(s,x)
inseconds(x::Number,R) = inseconds(x)
inseconds(x::Quantity) = error("Expected second argument, specifying sample rate.")
function inseconds(x::Number)
  warn("Unitless value, assuming seconds. Append s, ms or samples to avoid",
       " this warning (e.g. 500ms instead of 500)")
  x*s
end