text:
export hypercubicI, copyvertex, copytensor, copygate
export Gcp, Greset
export MaskedSquareLattice, rand_maskedsquare

struct MaskedSquareLattice <: Viznet.AbstractSquareLattice
    mask::Matrix{Bool}
end

Base.size(lt::MaskedSquareLattice) = size(lt.mask)
Base.size(lt::MaskedSquareLattice, i::Int) = size(lt.mask, i)
Viznet.unit(lt::MaskedSquareLattice) = 1/(max(size(lt)...))
Viznet.vertices(lt::MaskedSquareLattice) = [i for i in 1:length(lt.mask) if lt.mask[i]]
function rand_maskedsquare(Nx::Int, Ny::Int, ρ::Real; seed=nothing)
    if seed !== nothing
        Random.seed!(seed)
    end
    MaskedSquareLattice(rand(Nx, Ny) .< ρ)
end
function Base.getindex(lt::MaskedSquareLattice, i::Real, j::Real)
    step = unit(lt)
    (i-0.5)*step, (j-0.5)*step
end

function Viznet.isconnected(sq::MaskedSquareLattice, i::Int, j::Int)
    u = unit(sq)
    d = distance(sq[i], sq[j])
    sq.mask[i] && sq.mask[j] && d > 0.999u && d < (sqrt(2)+0.001)*u
end

Viznet.bonds(lt::MaskedSquareLattice) = sgbonds(lt)

function solve(::Type{TT}, sg::AbstractSpinglass{LT}; usecuda=false) where {TT, LT<:MaskedSquareLattice}
    lt = sg.lattice
    Lx, Ly = size(lt)
    nbit = Ly + 2
    reg = _init_reg(TT, lt, Val(usecuda))
    LI = LinearIndices(lt)
    k = 0
    l = 0
    _c(a, b) = isconnected(lt, LI[a...], LI[b...])
    for i=1:Lx
        println("Layer $i/$Lx")
        for j=1:Ly-1
            _c((i,j), (i,j+1)) && (k += 1; reg |> put(nbit, (j,j+1)=>Gvb(bondtensor(TT, sg, k))))
        end
        for j=1:Ly
            if lt.mask[i,j]
                l += 1
                reg |> put(nbit, j=>Gh(vertextensor(TT, sg, l)))
            elseif i!=Lx && lt.mask[i+1,j] && any(lt.mask[1:i,j])
                reg |> put(nbit, j=>Gcut(TT))
            end
        end
        (i!=Lx) && for j=1:Ly
            ancpre = nbit-(j-1)%2
            ancthis = nbit-(j)%2
            # store the information in qubit `j` to ancilla `nbit-j%2`
            j!=Ly && _c((i,j), (i+1,j+1)) && (reg |> put(nbit, (j, ancthis)=>Gcp(TT)))
            # interact with j-1 th qubit (a)
            j!=1 && _c((i+1,j-1), (i,j)) && (k += 1; reg |> put(nbit, (j-1,j)=>Gvb(bondtensor(TT, sg, k))))
            # onsite term (b)
            _c((i,j), (i+1,j)) && (k+=1; reg |> put(nbit, j=>Ghb(bondtensor(TT, sg, k))))
            if j!=1 && _c((i,j-1), (i+1,j))
                # interact with cached j-1 th qubit (c)
                k += 1
                reg |> put(nbit, (ancpre,j)=>Gvb(bondtensor(TT, sg, k)))
                # erease the information in previous ancilla
                reg |> put(nbit, ancpre=>Gcut(TT))
            end
        end
    end
    sum(state(reg))
end

function sgvertices(lt::MaskedSquareLattice)
    [v for v in sgvertices(SquareLattice(size(lt)...)) if lt.mask[v]]
end

function sgbonds(lt::MaskedSquareLattice)
    edges = Tuple{Int, Int}[]
    LI = LinearIndices(size(lt))
    function trypush!(i, j)
        if lt.mask[i...] && lt.mask[j...]
            push!(edges, (LI[i...], LI[j...]))
        end
    end
    Lx, Ly = size(lt)
    for i=1:Lx
        for j=1:Ly-1
            trypush!([i,j], [i,j+1])
        end
        (i!=Lx) && for j=1:Ly
            j!=1 && trypush!([i+1,j-1], [i,j])
            trypush!([i,j], [i+1,j])
            j!=1 && trypush!([i,j-1],[i+1,j])
        end
    end
    edges
end

regsize(lt::MaskedSquareLattice) = size(lt,2)+2

function assign_Js_hs(lt::MaskedSquareLattice, grad_Js::AbstractVector{T}, grad_hs) where T
    grid = zeros(size(lt.mask))
    vorder = sgvertices(lt)
    for (i, hg) in enumerate(grad_hs)
        grid[vorder[i]] = hg
    end
    for ((i, j), Jg) in zip(sgbonds(lt), grad_Js)
        assign_one!(grid, i, j, Jg)
    end
    return grid
end

function _init_reg(::Type{T}, lt::MaskedSquareLattice, ::Val{:false}) where T
    nbit = size(lt, 2) + 2
    state = ones(T, 1<<nbit)
    ArrayReg(state)
end


text:
# This tutorial uses the TMAR Filter from [Light2016](@cite)
#
# to reproduce the tutorial in section 4b.  It is a shear swirling
# flow deformation of a transported quantity from LeVeque 1996.  The exact
# solution at the final time is the same as the initial condition.

using MPI
using Test
using ClimateMachine
ClimateMachine.init()
using Logging
using ClimateMachine.Mesh.Topologies
using ClimateMachine.Mesh.Grids
using ClimateMachine.Mesh.Filters
using ClimateMachine.DGMethods
using ClimateMachine.DGMethods.NumericalFluxes
using ClimateMachine.DGMethods.FVReconstructions: FVConstant, FVLinear
using ClimateMachine.MPIStateArrays
using ClimateMachine.ODESolvers
using LinearAlgebra
using Printf
using Dates
using ClimateMachine.GenericCallbacks:
    EveryXWallTimeSeconds, EveryXSimulationSteps
using ClimateMachine.VTK: writevtk, writepvtu

using ClimateMachine
const clima_dir = dirname(dirname(pathof(ClimateMachine)))

if !@isdefined integration_testing
    const integration_testing = parse(
        Bool,
        lowercase(get(ENV, "JULIA_CLIMA_INTEGRATION_TESTING", "false")),
    )
end
const output = parse(Bool, lowercase(get(ENV, "JULIA_CLIMA_OUTPUT", "false")))

include("advection_diffusion_model.jl")

Base.@kwdef struct SwirlingFlow{FT} <: AdvectionDiffusionProblem
    period::FT = 5
end

init_velocity_diffusion!(::SwirlingFlow, aux::Vars, geom::LocalGeometry) =
    nothing

function initial_condition!(::SwirlingFlow, state, aux, localgeo, t)
    FT = eltype(state)
    x, y, _ = aux.coord
    x0, y0 = FT(1 // 3), FT(1 // 3)
    τ = 6hypot(x - x0, y - y0)
    state.ρ = exp(-τ^2)
end

has_variable_coefficients(::SwirlingFlow) = true
function update_velocity_diffusion!(
    problem::SwirlingFlow,
    ::AdvectionDiffusion,
    state::Vars,
    aux::Vars,
    t::Real,
)
    x, y, _ = aux.coord
    sx, cx = sinpi(x), cospi(x)
    sy, cy = sinpi(y), cospi(y)
    ct = cospi(t / problem.period)

    u = 2 * sx^2 * sy * cy * ct
    v = -2 * sy^2 * sx * cx * ct
    aux.advection.u = SVector(u, v, 0)
end

function do_output(
    mpicomm,
    vtkdir,
    vtkstep,
    dg,
    Q,
    model,
    testname;
    number_sample_points = 0,
)
    ## Name of the file that this MPI rank will write
    filename = @sprintf(
        "%s/%s_mpirank%04d_step%04d",
        vtkdir,
        testname,
        MPI.Comm_rank(mpicomm),
        vtkstep
    )

    statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))

    writevtk(
        filename,
        Q,
        dg,
        statenames;
        number_sample_points = number_sample_points,
    )

    ## Generate the pvtu file for these vtk files
    if MPI.Comm_rank(mpicomm) == 0
        ## Name of the pvtu file
        pvtuprefix = @sprintf("%s/%s_step%04d", vtkdir, testname, vtkstep)

        ## Name of each of the ranks vtk files
        prefixes = ntuple(MPI.Comm_size(mpicomm)) do i
            @sprintf("%s_mpirank%04d_step%04d", testname, i - 1, vtkstep)
        end

        writepvtu(pvtuprefix, prefixes, (statenames...,), eltype(Q))

        @info "Done writing VTK: $pvtuprefix"
    end
end

function test_run(
    mpicomm,
    ArrayType,
    fvmethod,
    topl,
    problem,
    dt,
    N,
    timeend,
    FT,
    vtkdir,
    outputtime,
)
    grid = DiscontinuousSpectralElementGrid(
        topl,
        FloatType = FT,
        DeviceArray = ArrayType,
        polynomialorder = N,
    )

    bcs = (HomogeneousBC{0}(),)
    model = AdvectionDiffusion{2}(problem, bcs, diffusion = false)

    dg = DGFVModel(
        model,
        grid,
        fvmethod,
        RusanovNumericalFlux(),
        CentralNumericalFluxSecondOrder(),
        CentralNumericalFluxGradient(),
    )

    Q = init_ode_state(dg, FT(0))

    ## We integrate so that the final solution is equal to the initial solution
    Qe = copy(Q)
    odesolver = LSRK54CarpenterKennedy(dg, Q; dt = dt, t0 = 0)

    # Set up the information callback
    starttime = Ref(Dates.now())
    cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)
        if s
            starttime[] = Dates.now()
        else
            energy = norm(Q)
            @info @sprintf(
                """Update
                simtime = %.16e
                runtime = %s
                norm(Q) = %.16e""",
                gettime(odesolver),
                Dates.format(
                    convert(Dates.DateTime, Dates.now() - starttime[]),
                    Dates.dateformat"HH:MM:SS",
                ),
                energy
            )
        end
    end
    callbacks = (cbinfo,)
    if ~isnothing(vtkdir)
        # Create vtk dir
        if MPI.Comm_rank(mpicomm) == 0
            mkpath(vtkdir)
        end
        MPI.Barrier(mpicomm)

        vtkstep = 0
        # Output initial step
        do_output(
            mpicomm,
            vtkdir,
            vtkstep,
            dg,
            Q,
            model,
            "swirling_flow";
            number_sample_points = N[1] + 1,
        )

        # Setup the output callback
        cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do
            vtkstep += 1
            Qe = init_ode_state(dg, gettime(odesolver))
            do_output(
                mpicomm,
                vtkdir,
                vtkstep,
                dg,
                Q,
                model,
                "swirling_flow";
                number_sample_points = N[1] + 1,
            )
        end
        callbacks = (callbacks..., cbvtk)
    end

    solve!(Q, odesolver; timeend = timeend, callbacks = callbacks)

    error = euclidean_distance(Q, Qe)

    # Print some end of the simulation information
    eng0 = norm(Qe)
    engf = norm(Q)

    engfe = norm(Qe)
    errf = euclidean_distance(Q, Qe)
    @info @sprintf """Finished
    norm(Q)                 = %.16e
    norm(Q) / norm(Q₀)      = %.16e
    norm(Q) - norm(Q₀)      = %.16e
    norm(Q - Qe)            = %.16e
    norm(Q - Qe) / norm(Qe) = %.16e
    """ engf engf / eng0 engf - eng0 errf errf / engfe
    errf
end

let
    ArrayType = ClimateMachine.array_type()

    mpicomm = MPI.COMM_WORLD

    expected_result = Dict()
    expected_result[1, FVConstant()] = 1.2437314516997458e-01
    expected_result[2, FVConstant()] = 9.8829388561567838e-02
    expected_result[3, FVConstant()] = 7.2096312912198937e-02
    expected_result[4, FVConstant()] = 4.7720226730379636e-02
    expected_result[1, FVLinear()] = 4.3807967239640234e-02
    expected_result[2, FVLinear()] = 1.4913083518239653e-02
    expected_result[3, FVLinear()] = 4.3666055848495802e-03
    expected_result[4, FVLinear()] = 1.2749240022458719e-03

    @testset "$(@__FILE__)" begin
        numlevels =
            integration_testing || ClimateMachine.Settings.integration_testing ?
            4 : 1
        FT = Float64
        for fvmethod in (FVConstant(), FVLinear())
            result = zeros(FT, numlevels)
            for l in 1:numlevels
                Ne = 20 * 2^(l - 1)
                polynomialorder = (4, 0)

                problem = SwirlingFlow()

                brickrange = (
                    range(FT(0); length = Ne + 1, stop = 1),
                    range(
                        FT(0);
                        length = polynomialorder[1] * Ne + 1,
                        stop = 1,
                    ),
                )

                topology = StackedBrickTopology(
                    mpicomm,
                    brickrange,
                    boundary = ((1, 1), (1, 1)),
                    connectivity = :face,
                )

                maxvelocity = 2
                elementsize = 1 / Ne
                dx = elementsize / polynomialorder[1]^2
                CFL = 1
                dt = CFL * dx / maxvelocity

                vtkdir = abspath(joinpath(
                    ClimateMachine.Settings.output_dir,
                    "fvm_swirl_lvl_$l",
                ))

                timeend = problem.period

                outputtime = timeend / 10
                dt = outputtime / ceil(Int64, outputtime / dt)

                @info @sprintf """Starting
                FT               = %s
                ArrayType        = %s
                FV Reconstuction = %s
                dim              = %d
                Ne               = %d
                polynomial order = %d
                final time       = %.16e
                time step        = %.16e
                """ FT ArrayType fvmethod 2 Ne polynomialorder[1] timeend dt

                result[l] = test_run(
                    mpicomm,
                    ArrayType,
                    fvmethod,
                    topology,
                    problem,
                    dt,
                    polynomialorder,
                    timeend,
                    FT,
                    output ? vtkdir : nothing,
                    outputtime,
                )
                @test result[l] ≈ FT(expected_result[l, fvmethod])
            end
            @info begin
                msg = ""
                for l in 1:(numlevels - 1)
                    rate = log2(result[l]) - log2(result[l + 1])
                    msg *= @sprintf("\n  rate for level %d = %e\n", l, rate)
                end
                msg
            end
        end
    end
end

ext: jl


text:
# To-DO: For hyperband and bayasian, do a version where every k jobs, the results are already merged

import Base.push!, Base.getindex, Base.show
using JLD2, DataFrames

Base.@kwdef struct HyperoptResults{S,T,V}
    pars::S
    res::T
    additonal_res::V = nothing
end

res(r::HyperoptResults) = r.res
pars(r::HyperoptResults) = r.pars

"""
    SlurmHyperoptimizer(ho::Hyperoptimizer; slurm_params::SlurmParams=nothing)

Creates a Slurm Job Array submission file and a Hyperoptimizer instance for your Hyperoptimizer needs. Takes the same arguments as Hyperoptimizer from Hyperopt + the keyword argument slurm_params.
"""
mutable struct SlurmHyperoptimizer
    sampler::AbstractHyperparameterSampler
    results::Vector{HyperoptResults}
    N_samples::Integer
    slurm_params::SlurmParams
    temp_dir
end

function SlurmHyperoptimizer(N_samples::Integer, sampler::AbstractHyperparameterSampler, slurm_params::SlurmParams, temp_dir="temp-hyperopt/")
    generate_slurm_file(slurm_params, N_samples)
    mkpath(temp_dir)
    results = Vector{HyperoptResults}(undef, N_samples)
    SlurmHyperoptimizer(sampler, results, N_samples, slurm_params, temp_dir)
end

"""
    save_result(sho::SlurmHyperoptimizer, res::HyperoptResults, i::Integer)

(Temporally) saves the results `res` of job `i`. Results are later merged with `merge_results!`.
"""
function save_result(sho::SlurmHyperoptimizer, res::HyperoptResults, i::Integer)
    mkpath(sho.temp_dir) # make sure the folder is really there
    save_path = string(sho.temp_dir,"res-",i,".jld2")
    JLD2.@save save_path res
end

"""
    merge_results!(sho::SlurmHyperoptimizer)

Merges the results and store them in the `sho.results` field. Also deletes the temporal files.
"""
function merge_results!(sho::SlurmHyperoptimizer)
    for i ∈ 1:sho.N_samples
        save_path = string(sho.temp_dir,"res-",i,".jld2")
        JLD2.@load save_path res
        sho.results[i] = res
    end
    rm(sho.temp_dir, recursive=true)
    return sho
end

get_results(sho::SlurmHyperoptimizer) = res.(sho.results)
get_params(sho::SlurmHyperoptimizer) = pars.(sho.results)

function get_param(sho::SlurmHyperoptimizer, key::Symbol)
    if key ∈ sho.sampler.par_names
        params = pars.(sho.results)
        return [i_param[key] for i_param ∈ params]
    else
        error("Parameter name not found.")
    end
end

"""
    DataFrames(sho::SlurmHyperoptimizer)

Returns the results and parameters as a DataFrame
"""
function DataFrames.DataFrame(sho::SlurmHyperoptimizer)
    DataFrame(Dict([key => get_param(sho, key) for key in sho.sampler.par_names]...,:results => get_results(sho)))
end

"""
    get_index(ho::SlurmHyperoptimizer, i::Integer)

Very hacky way how to index a Hyperoptimizer object, that respect its history, in case the individual samples are not independend (as for Hyperband)
"""
function Base.getindex(ho::SlurmHyperoptimizer, i::Integer)
    @assert i <= ho.N_samples
    return ho.sampler(ho.results)
end

function Base.iterate(iter::SlurmHyperoptimizer, state=1)
    if state>iter.N
        return nothing
    else
        return (iter.sampler(ho.results), state+1)
    end
end

Base.show(io::IO, ho::SlurmHyperoptimizer) = "SlurmHyperoptimizer"

# reimplement RandomSampler
abstract type AbstractHyperparameterSampler end
struct RandomSampler <: AbstractHyperparameterSampler
    par_dic
    par_names
end

RandomSampler(;kwargs...) = RandomSampler(kwargs, keys(kwargs))

function (samp::RandomSampler)(results)
    Dict([key => rand(samp.par_dic[key]) for key in samp.par_names]...)
end


text:
export seqassim

"""
		seqassim(F::StateSpace, data::SyntheticData, J::Int64, ϵx::InflationType, algo::SeqFilter, X, Ny, Nx, t0::Float64)

Generic API for sequential data assimilation for any sequential filter of parent type `SeqFilter`.
"""
function seqassim(F::StateSpace, data::SyntheticData, J::Int64, ϵx::InflationType, algo::SeqFilter, X, Ny, Nx, t0::Float64)

Ne = size(X, 2)

step = ceil(Int, algo.Δtobs/algo.Δtdyn)
statehist = Array{Float64,2}[]
push!(statehist, deepcopy(X[Ny+1:Ny+Nx,:]))

n0 = ceil(Int64, t0/algo.Δtobs) + 1
Acycle = n0:n0+J-1
tspan = (t0, t0 + algo.Δtobs)
# Define the dynamical system
prob = ODEProblem(F.f, zeros(Nx), tspan)

# Run filtering algorithm
@showprogress for i=1:length(Acycle)
    # Forecast
	tspan = (t0+(i-1)*algo.Δtobs, t0+i*algo.Δtobs)
	prob = remake(prob; tspan=tspan)

	prob_func(prob,i,repeat) = ODEProblem(prob.f, X[Ny+1:Ny+Nx,i],prob.tspan)

	ensemble_prob = EnsembleProblem(prob,output_func = (sol,i) -> (sol[end], false),
	prob_func=prob_func)
	sim = solve(ensemble_prob, Tsit5(), dt = algo.Δtdyn, adaptive = false, EnsembleThreads(),trajectories = Ne,
				dense = false, save_everystep=false);

	@inbounds for i=1:Ne
	    X[Ny+1:Ny+Nx, i] .= deepcopy(sim[i])
	end

    # Assimilation # Get real measurement # Fix this later # Things are shifted in data.yt
    ystar = data.yt[:,Acycle[i]]
	# Replace at some point by realobserve(model.h, t0+i*model.Δtobs, ens)
	# Perform inflation for each ensemble member
	ϵx(X, Ny+1, Ny+Nx)

	# Compute measurements
	observe(F.h, X, t0+i*algo.Δtobs, Ny, Nx)

    # Generate posterior samples.
	# Note that the additive inflation of the observation is applied within the sequential filter.
    X = algo(X, ystar, t0+i*algo.Δtobs-t0)

	# Filter state
	if algo.isfiltered == true
		for i=1:Ne
			statei = view(X, Ny+1:Ny+Nx, i)
			statei .= algo.G(statei)
		end
	end

    push!(statehist, copy(X[Ny+1:Ny+Nx,:]))
	end

return statehist
end

#Version with Localization metric

"""
		seqassim(F::StateSpace, data::SyntheticData, J::Int64, ϵx::InflationType, algo::SeqFilter, X, Ny, Nx, t0::Float64, Loc::Localization)

Generic API for localized sequential data assimilation for any sequential filter of parent type `SeqFilter` with localization scheme `Loc`.
"""
function seqassim(F::StateSpace, data::SyntheticData, J::Int64, ϵx::InflationType, algo::SeqFilter, X, Ny, Nx,t0::Float64, Loc::Localization)

Ne = size(X, 2)

step = ceil(Int, algo.Δtobs/algo.Δtdyn)
statehist = Array{Float64,2}[]
push!(statehist, deepcopy(X[Ny+1:Ny+Nx,:]))

n0 = ceil(Int64, t0/algo.Δtobs) + 1;
Acycle = n0:n0+J-1
tspan = (t0, t0 + algo.Δtobs)
prob = ODEProblem(F.f, zeros(Nx), tspan)

# Run filtering algorithm
@showprogress for i=1:length(Acycle)
    # Forecast
	tspan = (t0+(i-1)*algo.Δtobs, t0+i*algo.Δtobs)
	prob = remake(prob; tspan=tspan)

	prob_func(prob,i,repeat) = ODEProblem(prob.f,view(X,Ny+1:Ny+Nx,i),prob.tspan)

	ensemble_prob = EnsembleProblem(prob, output_func = (sol,i) -> (sol[end], false),
	prob_func=prob_func)
	sim = solve(ensemble_prob, Tsit5(), dt = algo.Δtdyn, adaptive = false, EnsembleThreads(), trajectories = Ne,
				dense = false, save_everystep=false);

	@inbounds for i=1:Ne
		view(X, Ny+1:Ny+Nx, i) .= sim[i]
	end

    # Assimilation # Get real measurement # Fix this later # Things are shifted in data.yt
    ystar = deepcopy(view(data.yt,:,Acycle[i]))
	# Replace at some point by realobserve(model.h, t0+i*model.Δtobs, ens)
	# Perform inflation for each ensemble member
	ϵx(X, Ny+1, Ny+Nx)
	# Compute measurements
	observe(F.h, X, t0+i*algo.Δtobs, Ny, Nx)
    # Generate posterior samples
    X = algo(X, ystar, t0+i*algo.Δtobs-t0, Loc)

	# Filter state
	if algo.isfiltered == true
		for i=1:Ne
			statei = view(X, Ny+1:Ny+Nx, i)
			statei .= algo.G(statei)
		end
	end

    push!(statehist, copy(X[Ny+1:Ny+Nx,:]))
end

return statehist
end


text:
struct HamiltonianProblem{iip} <: DiffEqBase.AbstractDynamicalODEProblem end

"""
    HamiltonianProblem(H, p0, q0, tspan, param=nothing; kwargs...)
    HamiltonianProblem((dp, dq), p0, q0, tspan, param=nothing; kwargs...)

Define a physical system by its Hamiltonian function `H(p, q, param)` or the function pair
`dp = -∂H/∂q` and `dq = ∂H/∂p`.

The equations of motion are then given by `q̇ = ∂H/∂p = dq` and `ṗ  = -∂H/∂q = dp`.

The initial values for canonical impulses `p0` and coordinates `q0` may be scalars, `SArray`s or
other `AbstractArray`s. Their type determines the type of functions `dp` and `dq`.

In the first two cases the derivative functions require the signatures
`dp(p, q, param, t)` and `dq(p, q, param, t)` while
in the latter case the partial derivatives use mutating signatures
`dp!(Δp, p, q, param, t)` and `dq!(Δq, p, q, param, t)` with predefined arrays `Δp` and `Δq`.

If the Hamiltonian function is given, `dp` and `dq` are calculated automatically using
AD (`ForwardDiff`).

!!! note
`H` may be defined with or without time as fourth argument. If both methods are defined,
that with 4 arguments is used.
"""
function HamiltonianProblem(H, p0::S, q0::T, tspan, param=nothing; kwargs...) where {S,T}
  iip = T <: AbstractArray && !(T <: SArray) && S <: AbstractArray && !(S <: SArray)
  HamiltonianProblem{iip}(H, p0, q0, tspan, param; kwargs...)
end

struct PhysicsTag end

function generic_derivative(q0, hami, x)
    ForwardDiff.gradient(hami, x)
end

function generic_derivative(q0::Number, hami, x)
    ForwardDiff.derivative(hami, x)
end

function HamiltonianProblem{false}((dp, dq)::Tuple{Any,Any}, p0, q0, tspan, param=nothing; kwargs...)
    return ODEProblem(DynamicalODEFunction{false}(dp, dq), ArrayPartition(p0, q0), tspan, param; kwargs...)
end
function HamiltonianProblem{true}((dp, dq)::Tuple{Any,Any}, p0, q0, tspan, param=nothing; kwargs...)
    return ODEProblem(DynamicalODEFunction{true}(dp, dq), ArrayPartition(p0, q0), tspan, param; kwargs...)
end

function HamiltonianProblem{false}(H, p0, q0, tspan, param=nothing; kwargs...)
    if DiffEqBase.numargs(H) == 4
        dp = (p, q, param, t) -> generic_derivative(q0, q -> -H(p, q, param, t), q)
        dq = (p, q, param, t) -> generic_derivative(q0, p -> H(p, q, param, t), p)
    else
        issue_depwarn()
        dp = (p, q, param, t) -> generic_derivative(q0, q -> -H(p, q, param), q)
        dq = (p, q, param, t) -> generic_derivative(q0, p -> H(p, q, param), p)
    end
    return HamiltonianProblem{false}((dp, dq), p0, q0, tspan, param; kwargs...)
end

function HamiltonianProblem{true}(H, p0, q0, tspan, param=nothing; kwargs...)
    let cp = ForwardDiff.GradientConfig(PhysicsTag(), p0),
        cq = ForwardDiff.GradientConfig(PhysicsTag(), q0),
        vfalse = Val(false)

        if DiffEqBase.numargs(H) == 4
            dp = (Δp, p, q, param, t) -> ForwardDiff.gradient!(Δp, q->-H(p, q, param, t), q, cq, vfalse)
            dq = (Δq, p, q, param, t) -> ForwardDiff.gradient!(Δq, p-> H(p, q, param, t), p, cp, vfalse)
        else
            issue_depwarn()
            dp = (Δp, p, q, param, t) -> ForwardDiff.gradient!(Δp, q->-H(p, q, param), q, cq, vfalse)
            dq = (Δq, p, q, param, t) -> ForwardDiff.gradient!(Δq, p-> H(p, q, param), p, cp, vfalse)
        end
        return HamiltonianProblem{true}((dp, dq), p0, q0, tspan, param; kwargs...)
    end
end

function issue_depwarn()
    Base.depwarn("Hamiltonians with 3 arguments are deprecated; please use `H(p, q, params, t)`",
                 :HamiltonianProblem)
end


text:
using Base.Test
using AGRN

function get_diff(parent::GRN, child::GRN)
    diffs = 0
    for i in 1:100
        inputs = rand(parent.nin)
        parent_outs = step!(parent, inputs)
        child_outs = step!(parent, inputs)
        if any(parent_outs .!= child_outs)
            diffs += 1
        end
    end
    diffs
end

@testset "Mutation tests" begin
    config = AGRN.Config()
    nin = rand(1:10)
    nout = rand(1:10)
    nreg = rand(1:10)
    parent = GRN(nin, nout, nreg, config)
    @testset "Add mutate" begin
        child = AGRN.mutate_add(parent)
        @test child.nin == parent.nin
        @test child.nout == parent.nout
        @test child.beta == parent.beta
        @test child.delta == parent.delta
        @test length(child.ids) == length(parent.ids) + 1
        @test length(child.enh) == length(parent.enh) + 1
        @test length(child.inh) == length(parent.inh) + 1
    end
    @testset "Delete mutate" begin
        child = AGRN.mutate_delete(parent)
        @test child.nin == parent.nin
        @test child.nout == parent.nout
        @test child.beta == parent.beta
        @test child.delta == parent.delta
        @test length(child.ids) == length(parent.ids) - 1
        @test length(child.enh) == length(parent.enh) - 1
        @test length(child.inh) == length(parent.inh) - 1
    end
    @testset "Modify mutate" begin
        child = AGRN.mutate_modify(parent, config)
        @test child.nin == parent.nin
        @test child.nout == parent.nout
        @test length(child.ids) == length(parent.ids)
        @test length(child.enh) == length(parent.enh)
        @test length(child.inh) == length(parent.inh)
        @test get_diff(parent, child) > 0
    end
    @testset "General mutation" begin
        child = mutate(parent, config)
        @test child.nin == parent.nin
        @test child.nout == parent.nout
        @test get_diff(parent, child) > 0
    end
end


text:
TruncatedInverseGaussian(μ, λ, l, u)
rand(PolyaGamma(1, 2))
using Plots
using StatsPlots
using Distributions

a = rand(LogNormal(0, 0.25), 500)
b = rand(Normal(0, 1), 500)
theta = rand(Normal(0, 1), 10_000)
x = abs.([a[i] * (theta[e] - b[i]) for e = 1:10_000, i = 1:500])
@benchmark [rand(PolyaGamma(1, x[e, i]), 100_000) for e = 1:1, i = 1:50]


text:
# ------------ Choice sites ------------ #

@inline function (ctx::UpdateContext)(call::typeof(trace),
                                      addr::T,
                                      d::Distribution{K}) where {T <: Address, K}
    visit!(ctx, addr)

    # Check if in previous trace's choice map.
    in_prev_chm = has_value(ctx.prev, addr)
    in_prev_chm && begin
        prev = get_sub(ctx.prev, addr)
        prev_ret = get_value(prev)
        prev_score = get_score(prev)
    end

    # Check if in target.
    in_target = has_value(ctx.target, addr)

    # Ret.
    if in_target
        ret = getindex(ctx.target, addr)
        in_prev_chm && begin
            set_sub!(ctx.discard, addr, prev)
        end
    elseif in_prev_chm
        ret = prev_ret
    else
        ret = rand(d)
    end

    # Update.
    score = logpdf(d, ret)
    if in_prev_chm
        increment!(ctx, score - prev_score)
    elseif in_target
        increment!(ctx, score)
    end
    add_choice!(ctx, addr, score, ret)

    return ret
end

# ------------ Learnable ------------ #

@inline function (ctx::UpdateContext)(fn::typeof(learnable), addr::Address)
    visit!(ctx, addr)
    haskey(ctx.params, addr) && return getindex(ctx.params, addr)
    error("(learnable): parameter not provided at address $addr.")
end

# ------------ Black box call sites ------------ #

@inline function (ctx::UpdateContext)(c::typeof(trace),
                                      addr::T,
                                      call::Function,
                                      args...) where {T <: Address, D <: Diff}
    visit!(ctx, addr)
    ps = get_sub(ctx.params, addr)
    ss = get_sub(ctx.target, addr)
    if has_sub(ctx.prev, addr)
        prev = get_prev(ctx, addr)
        ret, cl, w, rd, d = update(ss, ps, prev, args...)
    else
        ret, cl, w = generate(ss, ps, call, args...)
    end
    add_call!(ctx, addr, cl)
    increment!(ctx, w)
    return ret
end

@inline function (ctx::UpdateContext)(c::typeof(trace),
                                      addr::T,
                                      call::G,
                                      args...) where {G <: GenerativeFunction,
                                                      T <: Address}
    visit!(ctx, addr)
    ps = get_sub(ctx.params, addr)
    ss = get_sub(ctx.target, addr)
    if has_sub(ctx.prev, addr)
        prev = get_prev(ctx, addr)
        ret, cl, w, rd, d = update(ss, ps, prev, args...)
    else
        ret, cl, w = generate(ss, ps, call.fn, args...)
    end
    add_call!(ctx, addr, cl)
    increment!(ctx, w)
    return ret
end

# ------------ Utilities ------------ #

# TODO: re-write with new projection.
function update_projection_walk(tr::DynamicTrace,
                                visited::Visitor)
    weight = 0.0
    for (k, v) in shallow_iterator(tr)
        k in visited && continue
        weight += projection(v, SelectAll())[1]
    end
    weight
end

function update_discard_walk!(d::DynamicDiscard,
                              visited::Visitor,
                              prev::DynamicTrace)
    for (k, v) in shallow_iterator(prev)
        if !(k in visited)
            ss = get_sub(visited, k)
            if isempty(ss)
                set_sub!(d, k, v)
            else
                sd = get_sub(d, k)
                sd = isempty(sd) ? Empty() : sd
                update_discard_walk!(sd, ss, v)
                set_sub!(d, k, sd)
            end
        end
    end
end

# ------------ Convenience ------------ #

function update(ctx::UpdateContext, cs::DynamicCallSite, args::NTuple{N, Diffed}) where N
    cs.fn isa JFunction ? func = get_fn(cs.fn) : func = cs.fn
    ret = ctx(func, tupletype(args...), args...)
    adj_w = update_projection_walk(ctx.tr, ctx.visited)
    update_discard_walk!(ctx.discard, ctx.visited, unwrap(get_trace(cs)))
    return ret, DynamicCallSite(ctx.tr, ctx.score - adj_w, cs.fn, map(a -> unwrap(a), args), ret), ctx.weight, UnknownChange(), ctx.discard
end

# Explicitly parametrize the compilation pipeline.
function update(opt::J, sel::L, ps::P, cs::DynamicCallSite, args::Diffed...) where {J <: CompilationOptions, L <: AddressMap, P <: AddressMap, N}
    ctx = Update(opt, sel, ps, cs, DynamicTrace(), DynamicDiscard())
    return update(ctx, cs, args)
end

# Explicitly parametrize the compilation pipeline.
function update(opt::J, sel::L, ps::P, cs::DynamicCallSite, args...) where {J <: CompilationOptions, L <: AddressMap, P <: AddressMap, N}
    ctx = Update(opt, sel, ps, cs, DynamicTrace(), DynamicDiscard())
    args = map(args) do a
        a isa Diffed ? a : Δ(a, NoChange())
    end
    return update(ctx, cs, args)
end

function update(sel::L, ps::P, cs::DynamicCallSite, args...) where {L <: AddressMap, P <: AddressMap, N}
    ctx = Update(DefaultPipeline(), sel, ps, cs, DynamicTrace(), DynamicDiscard())
    args = map(args) do a
        a isa Diffed ? a : Δ(a, NoChange())
    end
    return update(ctx, cs, args)
end

function update(sel::L, cs::DynamicCallSite, args...) where L <: AddressMap
    ctx = Update(DefaultPipeline(), sel, Empty(), cs, DynamicTrace(), DynamicDiscard())
    if isempty(args)
        args = get_args(cs)
    end
    args = map(args) do a
        a isa Diffed ? a : Δ(a, NoChange())
    end
    return update(ctx, cs, args)
end

# ------------ Documentation ------------ #

@doc(
"""
```julia
mutable struct UpdateContext{C <: CallSite,
                             T <: AddressMap,
                             K <: AddressMap,
                             P <: AddressMap,
                             D <: Diff} <: ExecutionContext
    prev::C
    tr::T
    select::K
    weight::Float64
    score::Float64
    discard::DynamicAddressMap
    visited::Visitor
    params::P
    argdiffs::D
end
```

Inner constructor:

```julia
UpdateContext(cl::C, select::K, argdiffs::D) where {C <: CallSite, K <: AddressMap, D <: Diff} = new{C, typeof(cl.trace), K, EmptyAddressMap, D}(cl, typeof(cl.trace)(), select, 0.0, 0.0, AddressMap(), Visitor(), AddressMap(), argdiffs)
UpdateContext(cl::C, select::K, ps::P, argdiffs::D) where {C <: CallSite, K <: AddressMap, P <: AddressMap, D <: Diff} = new{C, typeof(cl.trace), K, EmptyAddressMap, D}(cl, typeof(cl.trace)(), select, 0.0, 0.0, AddressMap(), Visitor(), ps, argdiffs)
```

`UpdateContext` is an execution context used for updating the value of random choices in an existing recorded call site. This context will perform corrective updates to the likehood weights and scores so that this operation produces the correct weights and scores for the original model program constrained with the `select` selection in the `UpdateContext`.
""", UpdateContext)

@doc(
"""
```julia
ret, cl, w, retdiff, d = update(ctx::UpdateContext, cs::DynamicCallSite, args...) where D <: Diff
ret, cl, w, retdiff, d = update(sel::L, cs::DynamicCallSite) where L <: AddressMap
ret, cl, w, retdiff, d = update(sel::L, cs::DynamicCallSite, argdiffs::D, new_args...) where {L <: AddressMap, D <: Diff}
```

`update` provides an API to the `UpdateContext` execution context. You can use this function on any of the matching signatures above - it will return the return value `ret`, the updated `RecordSite` instance `cl` or `v_cl`, the updated weight `w`, a `Diff` instance for the return value `retdiff`, and a structure which contains any changed (i.e. discarded) record sites `d`.
""", update)


text:
include("PARAMETERS_quinto.jl")
include("locations_quinto.jl")
include("logic_quinto.jl")

"""
    quinto_sim([seed])

Run the random walk simulation of Yoojin/Eleanor's experiments.

Simulation controls:
    set_NUMBER_OF_WALKS()
    set_MAX_STEPS_PER_WAL()
"""
function quinto_sim(seed::Int=randseed())
    quinto_print(seed)
    data = Dict{String,Integer}()
    data["escape"] = 0
    data["spawn sensor"] = 0
    data["second sensor"] = 0
    data["third sensor"] = 0
    data["fourth sensor"] = 0
    data["particles unresolved"] = 0
    output = run_sim!(data, seed)
    present_quinto_data(output)
    nothing
end

"Helper function to start, not exported in Arlett module"
function quinto_print(seed::Int64)
    println("############################")
    println("      Quinto Simulation     ")
    println("############################")
    println("_________Parameters_________")
    println("# of trials:\t\t", string(NUMBER_OF_WALKS))
    println("# of steps (max):\t", string(MAX_STEPS_PER_WALK))
    println("Step size, water:\t", string(WATER_STEP_SIZE))
    println("Step size, enz:\t\t", string(ENZ_STEP_SIZE))
    println("Step size, ppd:\t\t", string(PPD_STEP_SIZE))
    println("Random seed:\t\t$seed")
    nothing
end

"Helper function to end, not exported in Arlett module"
function present_quinto_data(data::Dict{String, Integer})
    println("_________Results____________")
    println("# escaped:\t\t", data["escape"])
    println("*Sensors from left to right*")
    println("# in first (spawn):\t", data["spawn sensor"])
    println("# in second:\t\t", data["second sensor"])
    println("# in third:\t\t", data["third sensor"])
    println("# in fourth:\t\t", data["fourth sensor"])
    println("# unresolved:\t\t", data["particles unresolved"])
    println("Average steps taken:\t", data["avg steps taken"])
    nothing
end

max_stars_repo_name: s-y-wu/HMCResearchRandomWalks

max_issues_repo_name: s-y-wu/HMCResearchRandomWalks

max_forks_repo_name: s-y-wu/HMCResearchRandomWalks


text:
"""
    latest_value(ts::SortedDataFrame, at_times::Vector{<:Real};
                 assets::Vector{Symbol} = get_assets(ts))

Get the latest price at a each input time.
### Inputs
* `ts` - The tick data.
* `at_times` - The times you want the latest prices for.
* `assets` - The assets you want latest prices for.
### Returns
* A `DataFrame`. Rows are for each time specified in at_times. Columns are for each asset.
"""
function latest_value(ts::SortedDataFrame, at_times::Vector{<:Real};
                      assets::Vector{Symbol} = get_assets(ts))
    dd = DataFrame(Time = at_times)
    for a in assets
        indx = searchsortedlast_both(ts.df[ts.groupingrows[a],ts.time], at_times)
        dd[!,a] = ts.df[ts.groupingrows[a][indx]  , ts.value]
    end
    return dd
end

"""
If you have two sorted vectors. You want to find the index of each "indices" in
the vector "reference". then this is much faster than doing something like
map( x -> searchsortedlast(reference, x), indices) as it only goes once through
the reference vector.
"""
function searchsortedlast_both(reference::Vector, indices::Vector)
    len = length(indices)
    lenj = length(reference)
    inds = Array{Int64,1}(repeat([lenj], len))
    inds[1] = max(1,searchsortedlast(reference, indices[1]))
    if len == 1 return inds end
    j = inds[1]
    i = 2
    while true
        if reference[j+1] > indices[i]
            inds[i] = j
            i += 1
        else
            j += 1
        end
        if (i > len) | (j == lenj) return inds end
    end
end

"""
    random_value_in_interval(ts::SortedDataFrame, at_times::Vector{<:Real};
                             assets::Vector{Symbol} = get_assets(ts),
                             twister_arb_value_in_interval::MersenneTwister = MersenneTwister(2604))

Get a random value in an interval. So if you input times 1,7,8 then for the second entry it will pick a random update (if any exist) between times 1 and 7.
### Inputs
* `ts` - The tick data.
* `at_times` - The times that seperate the intervals of interest.
* `assets` - The assets of interest.
* `twister_arb_value_in_interval` - The RNG used in selecting the random interval.
### Returns
* A `DataFrame` with prices for each asset from random ticks in each interval.
"""
function random_value_in_interval(ts::SortedDataFrame, at_times::Vector{<:Real};
                                  assets::Vector{Symbol} = get_assets(ts),
                                  twister_arb_value_in_interval::MersenneTwister = MersenneTwister(2604))
    dd = DataFrame(Time = at_times)
    for a in assets
        indx = [1, searchsortedlast_both(ts.df[ts.groupingrows[a],ts.time], at_times)...]
        rand_indx = map(i -> rand(twister_arb_value_in_interval, min(indx[i+1],indx[i]+1):indx[i+1],1)[1]  ,1:(length(indx)-1))
        dd[!,a] = ts.df[ts.groupingrows[a][rand_indx]  , ts.value]
    end
    return dd
end

function next_tick(ts::SortedDataFrame, from_index::I; assets::Vector{Symbol} = get_assets(ts)) where R<:Real where I<:Integer
    inds = Array{I,1}()
    for a in assets
        ind = searchsortedfirst(ts.groupingrows[a], from_index)
        if (ind > length(ts.groupingrows[a])) return missing, -10 end
        push!(inds, ts.groupingrows[a][ind])
    end
    refresh_index = maximum(inds)
    refresh_time = ts.df[refresh_index, ts.time]
    return refresh_time, refresh_index
end

"""
    get_all_refresh_times(ts::SortedDataFrame, assets::Vector{Symbol} = get_assets(ts);
                          start_time::R = minimum(ts.df[:,ts.time])) where R<:Real

Get a vector of all refresh times when all assets have an updated price.
So if there are assets A and B that trade at times (1,5,6,7,10) and (2,5,7,9)
then the refresh times are (2,5,7,10) as at these four times there are updated
prices for all assets that have happened since the previous refresh time.
### Inputs
* `ts` - The tick data.
* `assets` - The assets of interest.
* `start_time` - From what time should we start looking for updated prices.
### Returns
* A `Vector` of refresh times.
"""
function get_all_refresh_times(ts::SortedDataFrame, assets::Vector{Symbol} = get_assets(ts);
                               start_time::R = minimum(ts.df[:,ts.time])) where R<:Real
    ticks = Vector{R}()
    start_ind = searchsortedfirst(ts.df[:,ts.time], start_time) - 1
    while true
        new_tick, start_ind = next_tick(ts, start_ind + 1; assets = assets)
        if ismissing(new_tick) return ticks end
        push!(ticks, new_tick)
    end
end


text:
using NNLS
using Test
# import NonNegLeastSquares
using PyCall
using ECOS
using JuMP
using Random
using LinearAlgebra

import Libdl

const pyopt = pyimport_conda("scipy.optimize", "scipy")

macro wrappedallocs(expr)
    argnames = [gensym() for a in expr.args]
    quote
        function g($(argnames...))
            @allocated $(Expr(expr.head, argnames...))
        end
        $(Expr(:call, :g, [esc(a) for a in expr.args]...))
    end
end

include("nnls.jl")
include("qp.jl")


text:
# show method for html output of question types
function _markdown_to_html(x)
    length(x) == 0 && return("")
    x = Markdown.parse(x)
    x = sprint(io -> Markdown.html(io, x))
    x = replace(x, r"^
"=>"", r"
$"=>"")
    return x
end

function Base.show(io::IO, m::MIME"text/html", x::Question)
    ID = randstring()

    FORM, GRADING_SCRIPT = prepare_question(x, ID)

    Mustache.render(io, html_templates["question_tpl"];
                    ID = ID,
                    STATUS = "",
                    LABEL=_markdown_to_html(x.label),
                    HINT = length(x.label) > 0 ? x.hint : "",
                    EXPLANATION = _markdown_to_html(x.explanation),
                    FORM = FORM,
                    GRADING_SCRIPT = GRADING_SCRIPT
                    )
end

function prepare_question(x::Numericq, ID)

    FORM = Mustache.render(html_templates["inputq_form"];
                           ID=ID,
                           PLACEHOLDER = isnothing(x.placeholder) ? "Numeric answer" : x.placeholder,
                           UNITS=x.units,
                           TYPE="number",
                           HINT = length(x.label) == 0 ? x.hint : ""
                           )

    GRADING_SCRIPT =
        Mustache.render(html_templates["input_grading_script"];
                        ID = ID,
                        CORRECT_ANSWER = "(Math.abs(this.value - $(x.val)) <= $(x.atol))",
                        INCORRECT = "Incorrect",
                        CORRECT = "Correct"
                        )
    (FORM, GRADING_SCRIPT)
end

function prepare_question(x::Stringq, ID)

    FORM = Mustache.render(html_templates["inputq_form"];
                           ID=ID,
                           PLACEHOLDER = isnothing(x.placeholder) ? "Text answer" : x.placeholder,
                           TYPE="text",
                           HINT = length(x.label) == 0 ? x.hint : ""
                           )

    GRADING_SCRIPT =
        Mustache.render(html_templates["input_grading_script"];
                        ID = ID,
                        CORRECT_ANSWER = """RegExp('$(x.re.pattern)').test(this.value)""",
                        INCORRECT = "Incorrect",
                        CORRECT = "Correct"
                        )
    (FORM, GRADING_SCRIPT)

end

function _make_item(i, choice)
    choice′ = sprint(io -> Markdown.html(io, Markdown.parse(choice)))
    # strip
 tag
    choice′ = chomp(choice′)
    choice′ = replace(choice′, r"^
" => "")
    choice′ = replace(choice′, r"
$" => "")

    return (NO=i, LABEL=choice′, VALUE=i)
end

function prepare_question(x::Radioq, ID)

    choices = string.(x.choices)
    items = [_make_item(i, choice) for (i,choice) ∈ enumerate(choices)]

    GRADING_SCRIPT = Mustache.render(html_templates["radio_grading_script"];
                             ID = ID,
                                     CORRECT_ANSWER = x.answer,
                                     INCORRECT = "Incorrect",
                                     CORRECT = "Correct"
                                     )
    FORM = Mustache.render(html_templates["Radioq"];
                           ID = ID,
                           ITEMS = items,
                           INLINE = x.inline ? " inline" : ""
                           )

    (FORM, GRADING_SCRIPT)

end

function prepare_question(x::Buttonq, ID)

    n = length(x.choices)
    choices = _markdown_to_html.(x.choices)
    buttons = [(i=i, TEXT=_markdown_to_html(x.choices[i]), ANSWER=x.answer[i]) for i ∈ 1:n]

    GRADING_SCRIPT = nothing
    FORM = Mustache.render(html_templates["Buttonq"];
                           ID = ID,
                           BUTTONS = buttons,
                           GREEN = x.colors.GREEN,
                           RED = x.colors.RED,
                           BLUE = x.colors.BLUE,
                           CORRECT = "✓",
                           INCORRECT="⨉"
                           )
    (FORM, GRADING_SCRIPT)

end

function prepare_question(x::Multiq, ID)

    choices = string.(x.choices)
    items = [_make_item(i, choice) for (i,choice) ∈ enumerate(choices)]

    GRADING_SCRIPT = Mustache.render(html_templates["multi_grading_script"];
                                     ID = ID,
                                     CORRECT_ANSWER = x.answer,
                                     INCORRECT = "Not yet",
                                     CORRECT = "Correct",
                             )
    FORM = Mustache.render(html_templates["Multiq"];
                           ID = ID,
                           ITEMS = items,
                           INLINE = x.inline ? " inline" : ""
                           )
    (FORM, GRADING_SCRIPT)

end

function prepare_question(x::MultiButtonq, ID)

    n = length(x.choices)
    choices = _markdown_to_html.(x.choices)
    buttons = [(i=i, TEXT=_markdown_to_html(x.choices[i])) for i ∈ 1:n]
    BLUE = "#0033CC11"

    GRADING_SCRIPT = Mustache.render(html_templates["multi_button_grading_script"];
                                     ID = ID,
                                     CORRECT_ANSWER = length(x.answer) > 0 ? x.answer : "[]",
                                     SELECTED_COLOR = BLUE,
                                     INCORRECT = "Something isn't correct",
                                     CORRECT = "Correct",
                                     CORRECT_flag = "✓ ",
                                     INCORRECT_flag ="⨉ "

                                     )

    FORM = Mustache.render(html_templates["MultiButtonq"];
                           ID = ID,
                           BUTTONS = buttons,
                           )
    (FORM, GRADING_SCRIPT)

end

function prepare_question(x::Matchq, ID)

    BLANK = "Choose..."
    ANSWER_CHOICES = [(INDEX=i, LABEL=_markdown_to_html(label)) for (i, label) in enumerate(x.choices)]
    items = [(ID=ID, NO=i, BLANK=BLANK,  QUESTION=_markdown_to_html(question), ANSWER_CHOICES=ANSWER_CHOICES) for (i,question) ∈ enumerate(x.questions)]
    GRADING_SCRIPT = Mustache.render(html_templates["matchq_grading_script"];
                                     ID = ID,
                                     CORRECT_ANSWER = collect(string.(x.answer)),
                                     INCORRECT = "Not yet",
                                     CORRECT = "Correct",
                                     )
    FORM = Mustache.render(html_templates["Matchq"];
                           ID = ID,
                           ITEMS = items,

                           )
    (FORM, GRADING_SCRIPT)
end

## ----
## blank creates BLANK, GRADING_SCRIPT
function blank(x::FillBlankChoiceQ, ID)
    _blank = "Choose..."

    ANSWER_CHOICES = [(INDEX=i, LABEL=_markdown_to_html(choice)) for
                      (i,choice) ∈ enumerate(x.choices)]
    BLANK = Mustache.render(html_templates["fill_in_blank_select"],
                            ID = ID,
                            BLANK=_blank,
                            ANSWER_CHOICES = ANSWER_CHOICES)

    GRADING_SCRIPT = Mustache.render(html_templates["matchq_grading_script"];
                                     ID = ID,
                                     CORRECT_ANSWER = collect(string.(x.answer)),
                                     INCORRECT = "Not yet",
                                     CORRECT = "Correct",
                                     )
    (BLANK, GRADING_SCRIPT)
end

function blank(x::FillBlankStringQ, ID)
    PLACEHOLDER = isnothing(x.placeholder) ? "answer here..." : x.placeholder
    BLANK = Mustache.render(html_templates["fill_in_blank_input"],
                            ID = ID,
                            TYPE="text",
                            PLACEHOLDER=PLACEHOLDER)

    GRADING_SCRIPT =
        Mustache.render(html_templates["input_grading_script"];
                        ID = ID,
                        CORRECT_ANSWER = """RegExp('$(x.re.pattern)').test(this.value)""",
                        INCORRECT = "Incorrect",
                        CORRECT = "Correct"
                        )
        (BLANK, GRADING_SCRIPT)
end

function blank(x::FillBlankNumericQ, ID)
    PLACEHOLDER = isnothing(x.placeholder) ? "answer here..." : x.placeholder
    BLANK = Mustache.render(html_templates["fill_in_blank_input"],
                            ID = ID,
                            TYPE="number",
                            PLACEHOLDER=PLACEHOLDER)
    GRADING_SCRIPT =
        Mustache.render(html_templates["input_grading_script"];
                        ID = ID,
                        CORRECT_ANSWER = "(Math.abs(this.value - $(x.val)) <= $(x.atol))",
                        INCORRECT = "Incorrect",
                        CORRECT = "Correct"
                        )
        (BLANK, GRADING_SCRIPT)
end

function prepare_question(x::FillBlankQ, ID)

    BLANK, GRADING_SCRIPT = blank(x, ID)

    question = _markdown_to_html(x.question)
    question = replace(question, r"_{4,}" => "{{{:BLANK}}}")
    FORM = Mustache.render(question; BLANK=BLANK)

    (FORM, GRADING_SCRIPT)
end

## ----
function prepare_question(x::HotspotQ, ID)

    x₀,x₁ = extrema(x.xs)
    y₀,y₁ = extrema(x.ys)

    CORRECT_ANSWER = isnothing(x.correct_answer) ?
        "(x >= $(x₀) && x <= $(x₁) && y >= $(y₀) && y <= $(y₁));" :
        x.correct_answer

    GRADING_SCRIPT =
        Mustache.render(html_templates["hotspot_grading_script"];
                        ID = ID,
                        CORRECT_ANSWER = CORRECT_ANSWER,
                        INCORRECT = "Incorrect",
                        CORRECT = "Correct"
                        )

    IMG = base64encode(read(x.imgfile, String))
    FORM = Mustache.render(html_templates["hotspot"]; ID=ID, IMG=IMG)

    (FORM, GRADING_SCRIPT)
end

handle_inf(x) = x == Inf ? "Infinity" : x == -Inf ? "-Infinity" : x
function prepare_question(x::PlotlyLightQ, ID)
    p = x.p
    p.id = ID

    x₀, x₁ = handle_inf.(x.xs)
    y₀, y₁ = handle_inf.(x.ys)
    CORRECT_ANSWER = isnothing(x.correct_answer) ?
        "(x >= $(x₀) && x <= $(x₁) && y >= $(y₀) && y <= $(y₁));" :
        x.correct_answer

    GRADING_SCRIPT =
        Mustache.render(html_templates["plotlylight_grading_script"];
                        ID = ID,
                        CORRECT_ANSWER = CORRECT_ANSWER,
                        INCORRECT = "Incorrect",
                        CORRECT = "Correct"
                        )

    FORM = sprint(io -> show(io, "text/html", p))
    FORM = "\n" * FORM

    (FORM, GRADING_SCRIPT)
end


text:
"""Localization for the Hijri month and day names."""
module locales
using Parameters

#@doc   """A Hijri locale object represents locale-specific data and functionality.""" ->
@with_kw struct Locale
    language_tag::String
    month_names::Array{String}
    gregorian_month_names::Array{String}
    day_names::Array{String}
    notation::String
    gregorian_notation::String
end

EnglishLocale = Locale(
    language_tag = "en",
    month_names = [
        "Muharram",
        "Safar",
        "Rabi’ al-Awwal",
        "Rabi’ al-Thani",
        "Jumada al-Ula",
        "Jumada al-Akhirah",
        "Rajab",
        "Sha’ban",
        "Ramadhan",
        "Shawwal",
        "Dhu al-Qi’dah",
        "Dhu al-Hijjah",
    ],
    gregorian_month_names = [
        "January",
        "February",
        "March",
        "April",
        "May",
        "June",
        "July",
        "August",
        "September",
        "October",
        "November",
        "December",
    ],
    day_names = [
        "Monday",
        "Tuesday",
        "Wednesday",
        "Thursday",
        "Friday",
        "Saturday",
        "Sunday",
    ],
    notation = "AH",
    gregorian_notation = "CE")

ArabicLocale = Locale(
    language_tag = "ar",
    month_names = [
        "محرم",
        "صفر",
        "ربيع الأول",
        "ربيع الثاني",
        "جمادى الأولى",
        "جمادى الآخرة",
        "رجب",
        "شعبان",
        "رمضان",
        "شوال",
        "ذو القعدة",
        "ذو الحجة",
    ],
    gregorian_month_names = [
        "يناير",
        "فبراير",
        "مارس",
        "أبريل",
        "مايو",
        "يونيو",
        "يوليو",
        "أغسطس",
        "سبتمبر",
        "أكتوبر",
        "نوفمبر",
        "ديسمبر",
    ],
    day_names = [
        "الإثنين",
        "الثلاثاء",
        "الأربعاء",
        "الخميس",
        "الجمعة",
        "السبت",
        "الأحد",
    ],
    notation = "هـ",
    gregorian_notation = "م"
)

AzerbaijaniLocale = Locale(
    language_tag = "az",
    month_names = [
        "məhərrəm",
        "səfər",
        "rəbiüləvvəl",
        "rəbiülaxır",
        "cəmadiyələvvəl",
        "cəmadiyəlaxır",
        "rəcəb",
        "şaban",
        "ramazan",
        "şəvval",
        "zilqədə",
        "zilhiccə",
    ],
    gregorian_month_names = [
        "yanvar",
        "fevral",
        "mart",
        "aprel",
        "may",
        "iyun",
        "iyul",
        "avqust",
        "sentyabr",
        "oktyabr",
        "noyabr",
        "dekabr",
    ],
    day_names = [
        "bazar ertəsi",
        "çərşənbə axşamı",
        "çərşənbə",
        "cümə axşamı",
        "cümə",
        "şənbə",
        "bazar",
    ],
    notation = "h.s.", # TODO: Verify this
    gregorian_notation = "b.e." # TODO: Verify this
    )

BengaliLocale = Locale(
    language_tag = "bn",
    month_names = [
        "মুহাররম",
        "সফর",
        "রবিউল আউয়াল",
        "রবিউস সানী",
        "জুমাদাল উলা",
        "জুমাদাস সানী",
        "রজব",
        "শাবান",
        "রমজান",
        "শাওয়াল",
        "জিলক্বদ",
        "জিলহজ",
    ],
    gregorian_month_names = [
        "জানুয়ারি",
        "ফেব্রুয়ারি",
        "মার্চ",
        "এপ্রিল",
        "মে",
        "জুন",
        "ঞ্জুলাই",
        "আগস্ট",
        "সেপ্টেম্বর",
        "অক্টোবর",
        "নভেম্বর",
        "ডিসেম্বর",
    ],
    day_names = [
        "সোমবার",
        "মঙ্গলবার",
        "বুধবার",
        "বৃহস্পতিবার",
        "শুক্রবার",
        "শনিবার",
        "রবিবার",
    ],
    notation = "হিজরি",
    gregorian_notation = "খ্রিস্টাব্দ"
    )

TurkishLocale = Locale(
    language_tag = "tr",
    month_names = [
        "Muharrem",
        "Safer",
        "Rebiülevvel",
        "Rebiülahir",
        "Cemaziyelevvel",
        "Cemaziyelahir",
        "Recep",
        "Şaban",
        "Ramazan",
        "Şevval",
        "Zilkade",
        "Zilhicce",
    ],
    gregorian_month_names = [
        "Ocak",
        "Şubat",
        "Mart",
        "Nisan",
        "Mayıs",
        "Haziran",
        "Temmuz",
        "Ağustos",
        "Eylül",
        "Ekim",
        "Kasım",
        "Aralık",
    ],
    day_names = [
        "Pazartesi",
        "Salı",
        "Çarşamba",
        "Perşembe",
        "Cuma",
        "Cumartesi",
        "Pazar",
    ],
    notation = "H.S.", # TODO: Verify this
    gregorian_notation = "M.S."
    )

_locale_map = Dict{String, Locale}("en" => EnglishLocale, "ar" => ArabicLocale, "az" => AzerbaijaniLocale, "bn" => BengaliLocale, "tr" => TurkishLocale)

@doc  """Return an appropriate :obj:`Locale` corresponding to a locale name.

      :param name: name of the locale.
      :type name: str
      """
function get_locale(name::String)
    language_tag = lowercase(name)[1:2]
    locale_cls = get(_locale_map, language_tag, nothing)

    if locale_cls == nothing
        throw(error(string("Unsupported language: ", language_tag))) #TODO: replace with ValueError
    end
    return locale_cls
end

locale_union = Union{Locale} #TODO: Automate this list
function __init_subclass__(locale::locale_union)
        if cls.language_tag in _locale_map
            throw(error(string("Duplicated language tag: ", cls.language_tag))) #TODO: replace with Lookup Error
        end
        _locale_map[locale.language_tag] = locale
end

@doc    """Return the month name for a specified Hijri month of the year.

        :param month: month of year, in range 1-12.
        :type month: int
        """
function month_name(locale::locale_union, month::Int)

        return locale.month_names[month]
end

@doc     """Return the month name for a specified Gregorian month of the year.

        :param month: month of year, in range 1-12.
        :type month: int
        """
function gregorian_month_name(locale::locale_union, month::Int)

        return locale.gregorian_month_names[month]
end

@doc    """Return the day name for a specified day of the week.

        :param day: day of week, where Monday is 1 and Sunday is 7.
        :type day: int
        """
function day_name(locale::locale_union, day::Int)  #TODO: locale should be one of Locale structs

        return locale.day_names[day]
end
end

max_stars_repo_name: ceferisbarov/hijri-converter

max_issues_repo_name: ceferisbarov/hijri-converter


text:
module PyCMA

using PyCall

const cma = PyNULL()

function __init__()
    copy!(cma, pyimport("cma"))
    py"""
    import numpy as np
    """
end

end # module


text:
function pattern(n::Int)::String
    output = IOBuffer()
    for i = 1:n
        str = string(i)
        for j = 1:i write(output,str) end
        if i<n println(output) end
    end
    return String(take!(output))
end

max_stars_repo_name: UlrichBerntien/Codewars-Katas

max_issues_repo_name: UlrichBerntien/Codewars-Katas

max_forks_repo_name: UlrichBerntien/Codewars-Katas


text:
# Based on this guide: https://ase.tufts.edu/chemistry/lin/images/FortranMD_TeachersGuide.pdf

using Atomistic
using InteratomicPotentials
using NBodySimulator
using Plots

N = 864
box_size = 3.47786u"nm"
reference_temp = 94.4u"K"
thermostat_prob = 0.1 # this number was chosen arbitrarily
Δt = 1e-2u"ps"

initial_system = generate_atoms_in_cubic_cell(N, :Ar, box_size, reference_temp)

eq_steps = 2000
eq_thermostat = NBodySimulator.AndersenThermostat(austrip(reference_temp), thermostat_prob / austrip(Δt))
eq_simulator = NBSimulator(Δt, eq_steps, thermostat=eq_thermostat)
potential = InteratomicPotentials.LennardJones(1.657e-21u"J", 0.34u"nm", 0.765u"nm", [:Ar])

eq_result = @time simulate(initial_system, eq_simulator, potential)

temp = plot_temperature(eq_result, 10)
energy = plot_energy(eq_result, 10)

prod_steps = 5000
prod_simulator = NBSimulator(Δt, prod_steps, t₀=get_time(eq_result))

prod_result = @time simulate(get_system(eq_result), prod_simulator, potential)

display(plot_temperature!(temp, prod_result, 10))
display(plot_energy!(energy, prod_result, 10))

rdf = plot_rdf(prod_result, potential.σ, Int(0.95 * prod_steps))
display(rdf)
savefig(rdf, "artifacts/argon_lj_nbs_ip_rdf.svg")

;


text:
cd(@__DIR__); allNGWPs = false; include("setups/toronto.jl")
gr(dpi = 200)

## (a) a smooth spatial distribution of the street intersections graph signal
f = zeros(N); for i in 1:N; f[i] = length(findall(dist_X[:,i] .< 1/minimum(edge_weight))); end
G_Sig.f = reshape(f, (N, 1))
gplot(A, X; width=1); plot!(size = (600, 500))
signal_plt = scatter_gplot!(X; marker = f, plotOrder = :s2l, ms = 3)
savefig(signal_plt, "../figs/Toronto_fdensity.png")

## (b) spatial distribution signal relative l2 approximation error by various methods
# uncomment the following if `allNGWPs` = true
# DVEC = getall_expansioncoeffs2(G_Sig, GP_dual, GP_dual_Lsym, VM_NGWP, PC_NGWP, LP_NGWP,
#                                VM_NGWP_Lsym, PC_NGWP_Lsym, LP_NGWP_Lsym, 𝚽, 𝚽sym)
# use precomputed results
DVEC = load("../datasets/Toronto_fdensity_DVEC.jld", "DVEC")
approx_error_plot2(DVEC);
plt = plot!(xguidefontsize = 14, yguidefontsize = 14, legendfontsize = 11)
savefig(plt, "../figs/Toronto_fdensity_DAG_approx.png")


text:
# these functions are related to calculating priors on the model parameters
using SpecialFunctions
using LinearAlgebra
using Unitful
using UnitfulAstro
import GPLinearODEMaker; GLOM = GPLinearODEMaker

# Keplerian priors references
# https://arxiv.org/abs/astro-ph/0608328
# table 1

const prior_K_min = 1e-4#u"m/s"  # m/s
const prior_K_max = 2129#u"m/s"  # m/s, corresponds to a maximum planet-star mass ratio of 0.01
const prior_γ_min = -prior_K_max  # m/s
const prior_γ_max = prior_K_max  # m/s
const prior_P_min = 1#u"d"  # days
const prior_P_max = 1e3 * convert_and_strip_units(u"d", 1u"yr")#u"d"
const prior_K0 = 0.3#u"m/s"  # * sqrt(50 / n_meas)  # m/s
const prior_e_min = 0
const prior_e_max = 1
const prior_ω_min = 0  # radians
const prior_ω_max = 2 * π  # radians
const prior_M0_min = 0  # radians
const prior_M0_max = 2 * π  # radians
const characteristic_P = 10  # days

function logprior_K(K::Unitful.Velocity; d::Integer=0, P::Unitful.Time=characteristic_P * u"d")
    return GLOM.log_loguniform(convert_and_strip_units(u"m/s", K), [prior_K_min, cbrt(prior_P_min / convert_and_strip_units(u"d", P)) * prior_K_max]; d=d, shift=prior_K0)
end

function logprior_P(P::Unitful.Time; d::Integer=0)
    # return GLOM.log_loguniform(convert_and_strip_units(u"d", P), [prior_P_min, prior_P_max]; d=d)
    return GLOM.log_uniform(convert_and_strip_units(u"d", P); min_max=[prior_P_min, prior_P_max], d=d)
end

function logprior_M0(M0::Real; d::Integer=0)
    return GLOM.log_uniform(M0; min_max=[prior_M0_min, prior_M0_max], d=d)
end

function logprior_e(e::Real; d::Integer=0)
    # return log_uniform(e; min_max=[prior_e_min, prior_e_max], d=d)
    return GLOM.log_Rayleigh(e, 1/5; d=d, cutoff=1)
end

function logprior_ω(ω::Real; d::Integer=0)
    return GLOM.log_uniform(ω; min_max=[prior_ω_min, prior_ω_max], d=d)
end

function logprior_hk(h::Real, k::Real; d::Vector{<:Integer}=[0,0])
    # return log_quad_cone([h, k]; d=d)
    return GLOM.log_rot_Rayleigh([h, k], 1/5; d=d, cutoff=1)
end

function logprior_γ(γ::Unitful.Velocity; d::Integer=0)
    return GLOM.log_uniform(convert_and_strip_units(u"m/s", γ); min_max=[prior_γ_min, prior_γ_max], d=d)
end

function logprior_kepler(
    K::Unitful.Velocity,
    P::Unitful.Time,
    M0::Real,
    e_or_h::Real,
    ω_or_k::Real,
    γ::Unitful.Velocity;
    d::Vector{<:Integer}=[0,0,0,0,0,0],
    use_hk::Bool=false)

    @assert all(0 .<= d .<= 2)
    @assert sum(d) <= 2

    if use_hk
        if any(d[4:5] .!= 0) && all(d[[1,2,3,6]] .== 0); return logprior_hk(e_or_h, ω_or_k; d=d[4:5]) end
        if sum(d .!= 0) > 1; return 0 end
    else
        if sum(d .!= 0) > 1; return 0 end
        if d[4] != 0; return logprior_e(e_or_h; d=d[4]) end
        if d[5] != 0; return logprior_ω(ω_or_k; d=d[5]) end
    end
    if d[1] != 0; return logprior_K(K; d=d[1]) end
    if d[2] != 0; return logprior_P(P; d=d[2]) end
    if d[3] != 0; return logprior_M0(M0; d=d[3]) end
    if d[6] != 0; return logprior_γ(γ; d=d[6]) end

    logP = logprior_K(K)
    # logP += logprior_K(K; P=P)
    logP += logprior_P(P)
    logP += logprior_M0(M0)
    use_hk ? logP += logprior_hk(e_or_h, ω_or_k) : logP += logprior_e(e_or_h; d=d[4]) + logprior_ω(ω_or_k; d=d[5])
    logP += logprior_γ(γ)

    return logP

end
logprior_kepler(ks::Union{kep_signal, kep_signal_epicyclic, kep_signal_wright}; d::Vector{<:Integer}=zeros(Int, n_kep_parms), use_hk::Bool=false) =
    use_hk ? logprior_kepler(ks.K, ks.P, ks.M0, ks.h, ks.k, ks.γ; d=d, use_hk=use_hk) : logprior_kepler(ks.K, ks.P, ks.M0, ks.e, ks.ω, ks.γ; d=d, use_hk=use_hk)
function logprior_kepler_tot(ks::Union{kep_signal, kep_signal_epicyclic, kep_signal_wright}; d_tot::Integer=0, use_hk::Bool=false)
    @assert 0 <= d_tot <= 2
    if d_tot == 0
        return logprior_kepler(ks; d=zeros(Int, n_kep_parms), use_hk=use_hk)
    elseif d_tot == 1
        G = zeros(n_kep_parms)
        d = zeros(Int, n_kep_parms)
        for i in 1:n_kep_parms
            d[i] = 1
            G[i] = logprior_kepler(ks; use_hk=use_hk, d=d)
            d[i] = 0
        end
        return G
    else
        H = zeros(n_kep_parms, n_kep_parms)
        d = zeros(Int, n_kep_parms)
        for i in 1:n_kep_parms
            for j in 1:n_kep_parms
                if i <= j
                    d[i] += 1
                    d[j] += 1
                    # println(typeof(d))
                    H[i, j] = logprior_kepler(ks; d=d, use_hk=use_hk)
                    d[i] = 0
                    d[j] = 0
                end
            end
        end
        return Symmetric(H)
    end
end


text:
using Printf, Statistics

function export_markdown_mcells(filename, results)
    rows = ["2D", "3D"]
    columns= ["compress", "decompress"]

    μ_cells = zeros(length(rows), length(columns))
    σ_cells = zeros(length(rows), length(columns))

    μ_bytes = zeros(length(rows), length(columns))
    σ_bytes = zeros(length(rows), length(columns))

    for (irow,row) in enumerate(rows), (icolumn,column) in enumerate(columns)
        ncells = results.benchmarkgroup["compression, $row"].tags[1]["ncells"]
        benchmark = results.benchmarkgroup["compression, $row"]["$column, F32"]
        x = (ncells / 1_000_000) ./ (benchmark.times .* 1e-9) # Mega-Cells per second
        μ_cells[irow,icolumn] = mean(x)
        σ_cells[irow,icolumn] = std(x)
        x = (4 * ncells / 1_000_000) ./ (benchmark.times .* 1e-9) # MB per second
        μ_bytes[irow,icolumn] = mean(x)
        σ_bytes[irow,icolumn] = std(x)
    end

    io = open(filename, "w")

    write(io, "# CvxCompress Throughput, Mega-Cells per second\n\n")

    write(io, "|    ")
    for column in columns
        write(io, " | $column")
    end
    write(io, "|\n")
    write(io, "|------")
    for column in columns
        write(io, "| ------ ")
    end
    write(io, "|\n")
    for (irow,row) in enumerate(rows)
        write(io, "| $row")
        for icolumn = 1:length(columns)
            _μ = @sprintf("%.2f", μ_cells[irow,icolumn])
            _σ = @sprintf("%.2f", 100*(σ_cells[irow,icolumn] / μ_cells[irow,icolumn]))
            write(io, " | $_μ MC/s ($_σ %)")
        end
        write(io, "|\n")
    end

    write(io, "# CvxCompress Throughput, MB per second\n\n")

    write(io, "|    ")
    for column in columns
        write(io, " | $column")
    end
    write(io, "|\n")
    write(io, "|------")
    for column in columns
        write(io, "| ------ ")
    end
    write(io, "|\n")
    for (irow,row) in enumerate(rows)
        write(io, "| $row")
        for icolumn = 1:length(columns)
            _μ = @sprintf("%.2f", μ_bytes[irow,icolumn])
            _σ = @sprintf("%.2f", 100*(σ_bytes[irow,icolumn] / μ_bytes[irow,icolumn]))
            write(io, " | $_μ MB/s ($_σ %)")
        end
        write(io, "|\n")
    end

    write(io, """
    ## Julia versioninfo
    ```
    $(results.vinfo)

    Environment Info:
      $(results.benchmarkconfig.env)

    Date:
      $(results.date)
    ```
    """)
    close(io)
end


text:
"""
`init_subspec!(m::HetDSGELag)`

Initializes a model subspecification by overwriting parameters from
the original model object with new parameter objects. This function is
called from within the model constructor.
"""
function init_subspec!(m::HetDSGELag)
    if subspec(m) == "ss0"
        return
    elseif subspec(m) == "ss1"
        return ss1!(m)
    else
        error("This subspec should be a 0")
    end
end

"""
```
ss1!(m::HetDSGELag)
```

Initializes subspec 1 of `HetDSGELag`.
This shuts down all shocks except the technology process (z shock).
"""
function ss1!(m::HetDSGELag)

end


text:
# This is just a playground file

using DataFrames
using PkgMining

df = pkg_mining()

df[df.owner .== "JuliaImages", :name]

df[df.name .== "DifferentialEquations", :owner]

df[occursin.("holy", df.owner), :].name

gd = groupby(df, :owner)

owners = Dict(first(g.owner) => nrow(g) for g in gd)
owners = [first(g.owner) => nrow(g) for g in gd]

# Ratio of owners with 10 or more packages:
length(filter(x -> last(x) ≥ 10, owners))/length(owners)

sort!(owners, by = x -> last(x), rev = true)

[nrow(df[df.owner .== owner, :]) for owner in df.owner]

# Create DataFrame with owner and the number of packages they own
df_owners = combine(groupby(df, :owner), nrow => :pkgs_owned)

# Copy and/or extend the packages DataFrame to contain column with the number of packages
# owned by the onwer of that package
df_owners2 = combine(groupby(df, :owner), :owner, :owner => length => :pkgs_owned)
df_owners3 = copy(df)
transform(groupby(df_owners3, :owner), :owner => length => :pkgs_owned)

# Check total number of packages
@info sum(df_owners.pkgs_owned) == nrow(df)


text:
using Statistics
using NCDatasets

using ArgParse
using JSON

include("nanop.jl")

function parse_commandline()

    s = ArgParseSettings()
    @add_arg_table s begin

        "--data-file"
            help = "PSLA data file. New variable will be appended."
            arg_type = String
            required = true

        "--domain-file"
            help = "Domain file."
            arg_type = String
            required = true

    end

    return parse_args(ARGS, s)
end

parsed = parse_commandline()
print(json(parsed, 4))

Dataset(parsed["data-file-SSTA"], "r") do ds

    global SSTA  = replace(ds["SSTA"][:], missing=>NaN)
    global months  = ds.dim["time"]

end

Dataset(parsed["domain-file"], "r") do ds

    global lat  = ds["yc"][:]
    global lon  = ds["xc"][:]

end

# El nino 3.4
# 170W - 120W  => [190 - 240]
# 5S - 5N

EN34idx = (abs.(lat) .<= 5) .& (lon .>= 190) .& (lon .<= 240)

EN34 = zeros(Float64, months)
for i = 1:months

    SSTA_slice = view(SSTA, :, :, i)
    EN34[i] = nanmean(SSTA_slice[EN34idx])

end

Dataset(parsed["data-file-SSTA"], "a") do ds

    for (varname, vardata, vardim, attrib) in [
        ("EN34",  EN34, ("time",), Dict()),
    ]

        if ! haskey(ds, varname)
            var = defVar(ds, varname, Float64, vardim)
            var.attrib["_FillValue"] = 1e20
        end

        var = ds[varname]

        for (k, v) in attrib
            var.attrib[k] = v
        end

        rng = []
        for i in 1:length(vardim)-1
            push!(rng, Colon())
        end
        push!(rng, 1:size(vardata)[end])
        var[rng...] = vardata

    end

end

max_stars_repo_name: meteorologytoday/CESM-diagnostic

max_issues_repo_name: meteorologytoday/CESM-diagnostic

max_forks_repo_name: meteorologytoday/CESM-diagnostic


text:
### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ e077ab16-22e5-11eb-15ed-5f0654bf7328
using Pkg, DrWatson

# ╔═╡ e75a9a06-22e5-11eb-2e36-8d4f62b830ed
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Random, Distributions, Plots, LaTeXStrings
	Random.seed!(0)
end;

# ╔═╡ ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
md"## Listing5.03"

# ╔═╡ fa68607e-22e5-11eb-0558-c9a4d9f77426
function statPair(dist,n)
    sample = rand(dist,n)
    [mean(sample),var(sample)]
end

# ╔═╡ 1111d8c4-252c-11eb-3813-81de856f7fba
begin
	stdUni = Uniform(-sqrt(3),sqrt(3))
	n, N = 3, 10^4

	dataUni     = [statPair(stdUni,n) for _ in 1:N]
	dataUniInd  = [[mean(rand(stdUni,n)),var(rand(stdUni,n))] for _ in 1:N]
	dataNorm    = [statPair(Normal(),n) for _ in 1:N]
	dataNormInd = [[mean(rand(Normal(),n)),var(rand(Normal(),n))] for _ in 1:N]
end;

# ╔═╡ 11120d28-252c-11eb-1c1e-0dba64d8ae5c
begin
	p1 = scatter(first.(dataUni), last.(dataUni),
		c=:blue, ms=1, msw=0, label="Same group")
	p1 = scatter!(first.(dataUniInd), last.(dataUniInd),
		c=:red, ms=0.8, msw=0, label="Separate group",
		xlabel=L"\overline{X}", ylabel=L"S^2")

	p2 = scatter(first.(dataNorm), last.(dataNorm),
		c=:blue, ms=1, msw=0, label="Same group")
	p2 = scatter!(first.(dataNormInd), last.(dataNormInd),
		c=:red, ms=0.8, msw=0, label="Separate group",
		xlabel=L"\overline{X}", ylabel=L"$S^2$")

	plot(p1, p2, ylims=(0,5), size=(800, 400))
end

# ╔═╡ 475ff888-22e6-11eb-2354-09f8f40a8e12
md"## End of listing5.03"

# ╔═╡ Cell order:
# ╟─ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
# ╠═e077ab16-22e5-11eb-15ed-5f0654bf7328
# ╠═e75a9a06-22e5-11eb-2e36-8d4f62b830ed
# ╠═fa68607e-22e5-11eb-0558-c9a4d9f77426
# ╠═1111d8c4-252c-11eb-3813-81de856f7fba
# ╠═11120d28-252c-11eb-1c1e-0dba64d8ae5c
# ╟─475ff888-22e6-11eb-2354-09f8f40a8e12


text:
using Pkg
if isfile("Project.toml")
    Pkg.activate(".")
end

max_stars_repo_name: RodrigoZepeda/docker-julia

max_issues_repo_name: RodrigoZepeda/docker-julia

max_forks_repo_name: RodrigoZepeda/docker-julia


text:
function get_mutated_genotype(genotype)
    if MUTATION_PROBABILITY*length(genotype) < rand() return genotype end

    choices = [i for (i, _) in enumerate(genotype)]
    choice = rand(choices) # determine which bit to flip
    return genotype[1:(choice-1)] * ((genotype[choice] == '1') ? "0" : "1" ) * genotype[(choice+1):end]
end

max_stars_repo_name: tochsner/ALifeBenchmark

max_issues_repo_name: tochsner/ALifeBenchmark

max_forks_repo_name: tochsner/ALifeBenchmark


text:
function test_okl_reduction(backend::TestBackend)
    linfo=get_info_string(backend);
    entries=10000;
    nred   =256;

    reducedentries = ceil(Int64,(entries+nred-1)/nred);
    a=Array(Float32,(entries,)); a[:]=1.0;
    ared=Array(Float32,(reducedentries,)); ared[:]=0.0;

    kernelinfo = OCCA.KernelInfo();
    OCCA.adddefine!(kernelinfo,"p_Nred",@sprintf("%d",nred));

    device  = OCCA.Device(linfo);
    o_a     = OCCA.malloc(device, a);
    o_ared  = OCCA.malloc(device, ared);

    reduction = OCCA.buildkernel(device,"reduction.okl","reduction",kernelinfo);

    OCCA.runkernel!(reduction,entries,o_a,o_ared);

    OCCA.memcpy!(ared,o_ared);

    val=sum(ared);

    relerr= abs(sum(a)-val)/sum(a);

    return relerr<1e-5;

end


text:
using DataFrames
using Test

for testscen in 1:2
    valdir, scenario, use_permafrost, use_seaice = get_scenario(testscen)
    println(scenario)

    m = test_page_model()
    include("../src/components/SeaLevelRise.jl")

    SLR = add_comp!(m, SeaLevelRise)

    set_param!(m, :SeaLevelRise, :rt_g_globaltemperature, readpagedata(m, "test/validationdata/$valdir/rt_g_globaltemperature.csv"))
    set_param!(m, :SeaLevelRise, :y_year_0, 2015.)
    set_param!(m, :SeaLevelRise, :y_year, [2020, 2030, 2040, 2050, 2075, 2100, 2150, 2200, 2250, 2300])

    run(m)

    es_equilibriumSL = m[:SeaLevelRise, :es_equilibriumSL]
    es_equilibriumSL_compare = readpagedata(m, "test/validationdata/$valdir/es_equilibriumSL.csv")
    @test ones(10) ≈ es_equilibriumSL ./ es_equilibriumSL rtol = .01

    s_sealevel = m[:SeaLevelRise, :s_sealevel]
    s_sealevel_compare = readpagedata(m, "test/validationdata/$valdir/s_sealevel.csv")
    @test ones(10) ≈ s_sealevel ./ s_sealevel_compare rtol = .01

    expfs_exponential = m[:SeaLevelRise, :expfs_exponential]
    expfs_exponential_compare = readpagedata(m, "test/validationdata/expfs_exponential.csv")
    @test ones(10) ≈ expfs_exponential ./ expfs_exponential_compare rtol = .01
end


text:
# Exercise 15-1
struct Point
    x
    y
end

mutable struct MPoint
    x
    y
end

function distancebetweenpoints(p1::Point, p2::Point)
    dx = (p1.x - p2.x)^2
    dy = (p1.y - p2.y)^2
    return sqrt(dx + dy)
end

p1 = Point(3.0, 4.5)
p2 = Point(-1, 7.6)
# println(distancebetweenpoints(p1, p2))

# Exercise 15-2
# println(p1 == p2)
# println(p1 === p2)

# Exercise 15-3
struct Circle
    center::Point
    radius::Number
end

struct Rectangle
    width::Number
    height::Number
    corner
end

c = Circle(Point(150.0, 100.0), 75.0)
# println(c)

function pointincircle(c::Circle, p::Point)
    return distancebetweenpoints(c.center, p) <= c.radius
end

# function movepoint!(p, dx, dy)
#     p.x += dx
#     p.y += dy
#     nothing
# end

function rectincircle(c::Circle, r::Rectangle)
    w, h = r.width, r.height
    xchange = [0, 0, w, 0]
    ychange = [0, h, 0 , h]
    px, py = r.corner.x, r.corner.y
    condition = true
    for (i, dx) in enumerate(xchange)
        dy = ychange[i]
        if !pointincircle(c, Point(px + dx, py + dy))
            condition &= false
        end
    end
    return condition
end

r = Rectangle(3, 7.1, MPoint(0.0, 0.0))
c2 = Circle(Point(0.0, 0.0), 10)
# println(rectincircle(c, r))
# println(rectincircle(c2, r))

function rectcircleoverlap(c::Circle, r::Rectangle)
    w, h = r.width, r.height
    xchange = [0, 0, w, 0]
    ychange = [0, h, 0 , h]
    px, py = r.corner.x, r.corner.y
    for (i, dx) in enumerate(xchange)
        dy = ychange[i]
        if pointincircle(c, Point(px + dx, py + dy))
            return true
        end
    end
    return false
end

# println(rectcircleoverlap(c, r))
# println(rectcircleoverlap(c2, r))

# Exercise 15-4
using ThinkJulia

function drawrect(t::Turtle, r::Rectangle)
    px, py = rect.corner.x, rect.corner.y
    w, h = r.width, r.height
    dist = [h, w, h, w]
    ThinkJulia.Luxor.Reposition(t, px, py)
    for i in 1:length(dist)
        forward(t, dist[i])
        turn(t, -90)
    end
end

function ex15_4_1(rect)
    turtle = Turtle()
    @svg begin
        drawrect(turtle, rect)
    end
end

rect = Rectangle(20, 55, MPoint(-3.0, 15.0))
# ex15_4_1(rect)

function drawcircle(t::Turtle, c::Circle)
    center = c.center
    ThinkJulia.Luxor.Reposition(t, center.x, center.y)
    ThinkJulia.Luxor.Circle(t, c.radius)
end

function ex15_4_2(circ)
    turtle = Turtle()
    @svg begin
        drawcircle(turtle, circ)
    end
end

circ = Circle(Point(-3.0, 15.0), 30)
# ex15_4_2(circ)


text:
using LinearAlgebra

function f_model(f_x, grad_f_x, res, L)
    return f_x - real(dot(grad_f_x, res)) + (L / 2) * norm(res)^2
end

function lower_bound_smoothness_constant(f, A, x, grad_f_Ax)
    R = real(eltype(x))
    xeps = x .+ 1
    grad_f_Axeps, _ = gradient(f, A * xeps)
    return norm(A' * (grad_f_Axeps - grad_f_Ax)) / R(sqrt(length(x)))
end

function lower_bound_smoothness_constant(f, A, x)
    Ax = A * x
    grad_f_Ax, _ = gradient(f, Ax)
    return lower_bound_smoothness_constant(f, A, x, grad_f_Ax)
end

_mul!(y, L, x) = mul!(y, L, x)
_mul!(y, ::Nothing, x) = return

_gradient!(y, f, x) = gradient!(y, f, x)
_gradient!(::Nothing, f, x) = f(x)

function backtrack_stepsize!(
    gamma::R, f, A, g, x, f_Ax::R, At_grad_f_Ax, y, z, g_z::R, res, Az, grad_f_Az=nothing;
    alpha = 1, minimum_gamma = 1e-7
) where R
    f_Az_upp = f_model(f_Ax, At_grad_f_Ax, res, alpha / gamma)
    _mul!(Az, A, z)
    f_Az = _gradient!(grad_f_Az, f, Az)
    tol = 10 * eps(R) * (1 + abs(f_Az))
    while f_Az > f_Az_upp + tol && gamma >= minimum_gamma
        gamma /= 2
        y .= x .- gamma .* At_grad_f_Ax
        g_z = prox!(z, g, y, gamma)
        res .= x .- z
        f_Az_upp = f_model(f_Ax, At_grad_f_Ax, res, alpha / gamma)
        _mul!(Az, A, z)
        f_Az = _gradient!(grad_f_Az, f, Az)
        tol = 10 * eps(R) * (1 + abs(f_Az))
    end
    if gamma < minimum_gamma
        @warn "stepsize `gamma` became too small ($(gamma))"
    end
    return gamma, g_z, f_Az, f_Az_upp
end

function backtrack_stepsize!(
    gamma, f, A, g, x; alpha = 1, minimum_gamma = 1e-7
)
    Ax = A * x
    grad_f_Ax, f_Ax = gradient(f, Ax)
    At_grad_f_Ax = A' * grad_f_Ax
    y = x - gamma .* At_grad_f_Ax
    z, g_z = prox(g, y, gamma)
    return backtrack_stepsize!(
        gamma, f, A, g, x, f_Ax, At_grad_f_Ax, y, z, g_z, x - z, Ax, grad_f_Ax;
        alpha = alpha, minimum_gamma = minimum_gamma
    )
end


text:
struct DepParsingState
    words::Vector{Int}
    postags::Vector{Int}
    golds::Vector{Int}
    preds::Vector{Int}
    target::Int
end

function DepParsingState(words::Vector{Int}, postags::Vector{Int})
    @assert length(words) == length(postags)
    heads = zeros(Int, length(words))
    DepParsingState(words, postags, heads, 0)
end

function Base.isequal(x::DepParsingState, y::DepParsingState)
    isequal(x.target,y.target) && isequal(x.heads,y.heads)
end
Base.hash(x::DepParsingState) = hash(x.heads, hash(x.target))

function legalactions(state::DepParsingState)
    all(x -> x == 0, state.heads) && return Int[]

    actions = Int[]
    if state.target > 0
        for i = 1:length(state.heads)
            state.heads[i] == 0 && i != state.target && push!(actions,i)
        end
    else
        for i = 1:length(state.heads)
            state.heads[i] == 0 && push!(actions,i)
        end
    end
    actions
end

function reward(state::DepParsingState)
    c = 0
    for i = 1:length(golds)
        golds[i] == preds[i] && (c += 1)
    end
    c / length(golds)
end

function transition(state::DepParsingState, action::Int)
    heads = copy(state.heads)
    if state.target > 0
        @assert heads[state.target] == 0
        heads[state.target] = action
        target = 0
    else
        target = action
    end
    DepParsingState(state.words, state.postags, heads, target)
end


text:
"""
    Orbital(name::String, size::Int, l::Int, mr::Int)

Wannier90 orbital as defined in the [Wannier90 User Guide](https://github.com/wannier-developers/wannier90/raw/v3.1.0/doc/compiled_docs/user_guide.pdf). The `name` will be written in the `projections` block of the Wannier90 input.
"""
struct Orbital
    name :: String
    size :: Int
    l    :: Int
    mr   :: Int
end
Orbital() = Orbital("none", 0, 0, 0)
StructTypes.StructType(::Type{Orbital}) = StructTypes.Struct()

const ORBITALS = [Orbital("s", 1, 0, 1), Orbital("p", 3, 1, 0), Orbital("px", 1, 1, 2),
                  Orbital("py", 1, 1, 3), Orbital("pz", 1, 1, 1), Orbital("d", 5, 2, 0),
                  Orbital("dz2", 1, 2, 1), Orbital("dxz", 1, 2, 2), Orbital("dyz", 1, 2, 3),
                  Orbital("dx2-y2", 1, 2, 4), Orbital("dxy", 1, 2, 5),
                  Orbital("f", 7, 3, 0), Orbital("fz3", 1, 3, 1), Orbital("fxz2", 1, 3, 2),
                  Orbital("fyz2", 1, 3, 3), Orbital("fz(x2-y2)", 1, 3, 4),
                  Orbital("fxyz", 1, 3, 5), Orbital("fx(x2-3y2)", 1, 3, 6),
                  Orbital("fy(3x2-y2)", 1, 3, 7), Orbital("sp", 2, -1, 0),
                  Orbital("sp2", 3, -2, 0), Orbital("sp3", 4, -3, 0),
                  Orbital("sp3d2", 6, -5, 0)]

"""
    orbital(s::String)

Returns the [Orbital](@ref Orbitals) identified by `s`.
"""
orbital(s::String) = getfirst(x -> x.name == s, ORBITALS)
orbital(l::Number) = getfirst(x -> x.l == l, ORBITALS)

Base.convert(::Type{String}, x::Orbital) = x.name
Base.length(o::Orbital) = o.size

Base.hash(orbital::Orbital, h::UInt) = hash(orbital.mr, hash(orbital.l, h))
Base.:(==)(o1::Orbital, o2::Orbital) = o1.l == o2.l && o1.mr == o2.mr

"""
    Projection(orbital::Orbital, start::Int, last::Int)

A Wannier90 `Projection`, representing an `Orbital` with indices from `start` to `last`.
"""
mutable struct Projection
    orbital::Orbital
    start::Int
    last::Int
end
Projection() = Projection(Orbital(), 0, 0)
Projection(o::Orbital) = Projection(o, 0, 0)
Projection(o::String) = Projection(orbital(o), 0, 0)
Projection(o::Orbital, start::Int) = Projection(o, start, start + o.size - 1)
Projection(ostr::String, start::Int) = (o = orbital(ostr); Projection(o, start))
StructTypes.StructType(::Type{Projection}) = StructTypes.Struct()

function sanitize!(projs::Vector{Projection}, soc::Bool)
    id = 1
    for proj in projs
        size = soc ? 2 * proj.orbital.size : proj.orbital.size
        proj.start = id
        proj.last = id + size - 1
        id += id
    end
    return projs
end

Base.range(proj::Projection) = proj.start:proj.last
Base.length(proj::Projection) = length(proj.orbital)

function Base.hash(data::Projection, h::UInt)
    for f in fieldnames(Projection)
        h = hash(getfield(data, f), h)
    end
    return h
end


text:
using ImageEdgeDetection

using FileIO
using ImageDraw
using ImageFiltering
using ImageCore
using OffsetArrays
using ReferenceTests
using StaticArrays
using Test
using TestImages

include("testutils.jl")

@testset "ImageEdgeDetection.jl" begin
    include("util.jl")

    include("algorithms/canny.jl")
    include("algorithms/thin_edges.jl")
    include("algorithms/thin_subpixel_edges.jl")
    include("algorithms/nonmaxima_suppression.jl")
    include("algorithms/subpixel_nonmaxima_suppression.jl")
    include("algorithms/gradient_orientation.jl")
end

nothing