text:
module Mixtape

using MacroTools: isexpr
using IRTools
using Core.Compiler
using Core.Compiler: MethodInstance, NativeInterpreter, CodeInfo, CodeInstance, WorldView,
                     OptimizationState, Const, widenconst, MethodResultPure, CallMeta
import Core.Compiler.abstract_call_known
using LLVM
using LLVM.Interop

resolve(x) = x
resolve(gr::GlobalRef) = getproperty(gr.mod, gr.name)
resolve(c::Core.Const) = c.val

# Cache
using Core.Compiler: WorldView

# Interpreter
import Core.Compiler: AbstractInterpreter, InferenceResult, InferenceParams, InferenceState,
                      OptimizationParams
import Core.Compiler: get_world_counter, get_inference_cache, code_cache, lock_mi_inference,
                      unlock_mi_inference
import Core.Compiler: retrieve_code_info, validate_code_in_debug_mode

# Optimizer
import Core.Compiler: optimize
#import Core.Compiler: CodeInfo, convert_to_ircode, copy_exprargs, slot2reg, compact!, coverage_enabled, adce_pass!
#import Core.Compiler: ssa_inlining_pass!, getfield_elim_pass!, type_lift_pass!, verify_ir, verify_linetable

# JIT
using GPUCompiler: GPUCompiler
import GPUCompiler: FunctionSpec

#####
##### Cache
#####

struct CodeCache
    dict::Dict{MethodInstance,Vector{CodeInstance}}
    callback::Function
    CodeCache(callback) = new(Dict{MethodInstance,Vector{CodeInstance}}(), callback)
end

function Base.show(io::IO, ::MIME"text/plain", cc::CodeCache)
    print(io, "CodeCache: ")
    for (mi, cis) in cc.dict
        println(io)
        print(io, "- ")
        show(io, mi)

        for ci in cis
            println(io)
            print(io, "  - ")
            print(io, (ci.min_world, ci.max_world))
        end
    end
end

function Core.Compiler.setindex!(cache::CodeCache, ci::CodeInstance, mi::MethodInstance)
    if !isdefined(mi, :callbacks)
        mi.callbacks = Any[cache.callback]
    else
        # Check if callback is present
        if all(cb -> cb !== cache.callback, mi.callbacks)
            push!(mi.callbacks, cache.callback)
        end
    end
    cis = get!(cache.dict, mi, CodeInstance[])
    return push!(cis, ci)
end

const CACHE = Dict{DataType,CodeCache}()
get_cache(ai::DataType) = CACHE[ai]

#####
##### Setup execution engine
#####

# We have one global JIT and TM
const orc = Ref{LLVM.OrcJIT}()
const tm = Ref{LLVM.TargetMachine}()

function __init__()
    CACHE[NativeInterpreter] = CodeCache(cpu_invalidate)

    opt_level = Base.JLOptions().opt_level
    if opt_level < 2
        optlevel = LLVM.API.LLVMCodeGenLevelNone
    elseif opt_level == 2
        optlevel = LLVM.API.LLVMCodeGenLevelDefault
    else
        optlevel = LLVM.API.LLVMCodeGenLevelAggressive
    end

    tm[] = LLVM.JITTargetMachine(; optlevel=optlevel)
    LLVM.asm_verbosity!(tm[], true)

    orc[] = LLVM.OrcJIT(tm[]) # takes ownership of tm
    atexit() do
        return LLVM.dispose(orc[])
    end
end

#####
##### Cache world view
#####

function Core.Compiler.haskey(wvc::WorldView{CodeCache}, mi::MethodInstance)
    return Core.Compiler.get(wvc, mi, nothing) !== nothing
end

function Core.Compiler.get(wvc::WorldView{CodeCache}, mi::MethodInstance, default)
    cache = wvc.cache
    for ci in get!(cache.dict, mi, CodeInstance[])
        if ci.min_world <= wvc.worlds.min_world && wvc.worlds.max_world <= ci.max_world
            # TODO: if (code && (code == jl_nothing || jl_ir_flag_inferred((jl_array_t*)code)))
            return ci
        end
    end

    return default
end

function Core.Compiler.getindex(wvc::WorldView{CodeCache}, mi::MethodInstance)
    r = Core.Compiler.get(wvc, mi, nothing)
    r === nothing && throw(KeyError(mi))
    return r::CodeInstance
end

function Core.Compiler.setindex!(wvc::WorldView{CodeCache}, ci::CodeInstance,
                                 mi::MethodInstance)
    return Core.Compiler.setindex!(wvc.cache, ci, mi)
end

#####
##### Cache invalidation
#####

function invalidate(cache::CodeCache, replaced::MethodInstance, max_world, depth)
    cis = get(cache.dict, replaced, nothing)
    if cis === nothing
        return
    end
    for ci in cis
        if ci.max_world == ~0 % Csize_t
            @assert ci.min_world - 1 <= max_world "attempting to set illogical constraints"
            ci.max_world = max_world
        end
        @assert ci.max_world <= max_world
    end

    # recurse to all backedges to update their valid range also
    if isdefined(replaced, :backedges)
        backedges = replaced.backedges
        # Don't touch/empty backedges `invalidate_method_instance` in C will do that later
        # replaced.backedges = Any[]
        for mi in backedges
            invalidate(cache, mi, max_world, depth + 1)
        end
    end
end

#####
##### MixtapeInterpreter
#####

abstract type CompilationContext end
struct Fallback <: CompilationContext end
allow_transform(f::C, args...) where C <: CompilationContext = false
show_after_inference(f::CompilationContext) = false
show_after_optimization(f::CompilationContext) = false
debug(f::CompilationContext) = false
transform(ctx::CompilationContext, ir) = ir

function check(ctx::CompilationContext, mod::Module, fn, args...)
    allow_transform(ctx, mod) || allow_transform(ctx, fn, args...)
end

export CompilationContext, transform, allow_transform, show_after_inference,
       show_after_optimization, debug

struct MixtapeInterpreter{Ctx<:CompilationContext,Inner<:AbstractInterpreter} <:
       AbstractInterpreter
    ctx::Ctx
    inner::Inner
    errors::Any
    function MixtapeInterpreter(ctx::Ctx,
                                interp::Inner) where {Ctx<:CompilationContext,Inner}
        return new{Ctx,Inner}(ctx, interp, Any[])
    end
end
Base.push!(mxi::MixtapeInterpreter, e) = push!(mxi.errors, e)

get_world_counter(mxi::MixtapeInterpreter) = get_world_counter(mxi.inner)
get_inference_cache(mxi::MixtapeInterpreter) = get_inference_cache(mxi.inner)
InferenceParams(mxi::MixtapeInterpreter) = InferenceParams(mxi.inner)
OptimizationParams(mxi::MixtapeInterpreter) = OptimizationParams(mxi.inner)
Core.Compiler.may_optimize(mxi::MixtapeInterpreter) = true
Core.Compiler.may_compress(mxi::MixtapeInterpreter) = true
Core.Compiler.may_discard_trees(mxi::MixtapeInterpreter) = true
function Core.Compiler.add_remark!(mxi::MixtapeInterpreter, sv::InferenceState, msg)
    return Core.Compiler.add_remark!(mxi.inner, sv, msg)
end
lock_mi_inference(mxi::MixtapeInterpreter, mi::MethodInstance) = nothing
unlock_mi_inference(mxi::MixtapeInterpreter, mi::MethodInstance) = nothing

#####
##### Codegen/inference integration
#####

function code_cache(mxi::MixtapeInterpreter)
    return WorldView(get_cache(typeof(mxi.inner)), get_world_counter(mxi))
end

function cpu_invalidate(replaced, max_world)
    cache = get_cache(NativeInterpreter)
    invalidate(cache, replaced, max_world, 0)
    return nothing
end

function cpu_cache_lookup(mi, min_world, max_world)
    wvc = WorldView(get_cache(NativeInterpreter), min_world, max_world)
    return Core.Compiler.get(wvc, mi, nothing)
end

function cpu_infer(ctx, mi, min_world, max_world)
    wvc = WorldView(get_cache(NativeInterpreter), min_world, max_world)
    interp = MixtapeInterpreter(ctx, NativeInterpreter(min_world))
    ret = infer(wvc, mi, interp)
    if debug(ctx) && !isempty(interp.errors)
        println("Encountered the following non-critical errors during interception:")
        for k in interp.errors
            display(k)
        end
        println("This may imply that your transform failed to operate on some calls.")
    end
    return ret
end

function update!(ci::CodeInfo, ir::Core.Compiler.IRCode)
    Core.Compiler.replace_code_newstyle!(ci, ir, length(ir.argtypes) - 1)
    ci.inferred = false
    ci.ssavaluetypes = length(ci.code)
    IRTools.slots!(ci)
    fill!(ci.slotflags, 0)
    return ci
end

function update!(ci::CodeInfo, ir::IRTools.IR)
    if ir.meta isa IRTools.Inner.Meta
        ci.method_for_inference_limit_heuristics = ir.meta.method
        if isdefined(ci, :edges)
            ci.edges = Core.MethodInstance[ir.meta.instance]
        end
    end
    return update!(ci, Core.Compiler.IRCode(IRTools.slots!(ir)))
end

function infer(wvc, mi, interp)
    src = Core.Compiler.typeinf_ext_toplevel(interp, mi)
    try
        fn = resolve(GlobalRef(mi.def.module, mi.def.name))
        as = mi.specTypes.parameters[2 : end]
        if check(interp.ctx, mi.def.module, fn, as...) && show_after_inference(interp.ctx)
            println("(Inferred) $fn in $(mi.def.module)")
            display(src)
        end
    catch e
        push!(interp, e)
    end
    @assert Core.Compiler.haskey(wvc, mi)
    ci = Core.Compiler.getindex(wvc, mi)
    if ci !== nothing && ci.inferred === nothing
        ci.inferred = src
    end
    return
end

# Workaround for what appears to be a Base bug
untvar(t::TypeVar) = t.ub
untvar(x) = x

function prepare_ir!(ir::IRTools.IR; type=Any)
    for (v, st) in ir
        isexpr(st.expr) || continue
        ir[v] = IRTools.stmt(Expr(st.expr.head, map(resolve, st.expr.args)...); type=type)
    end
    return ir
end

function set_argtypes!(ir::IRTools.IR, as)
    for (ind, t) in enumerate(as)
        IRTools.argtypes(ir)[ind] = t
    end
    return ir
end

# Replace usage sited of `retrieve_code_info`, OptimizationState is one such, but in all interesting use-cases
# it is derived from an InferenceState. There is a third one in `typeinf_ext` in case the module forbids inference.
function InferenceState(result::InferenceResult, cached::Bool, interp::MixtapeInterpreter)
    src = retrieve_code_info(result.linfo)
    try
        fn = resolve(GlobalRef(result.linfo.def.module, result.linfo.def.name))
        as = map(resolve, result.argtypes[2 : end])
        m = IRTools.Meta(result.linfo.def,
                         result.linfo,
                         src,
                         result.linfo.def.nargs,
                         result.linfo.sparam_vals)
        if !=(m, nothing) && check(interp.ctx, result.linfo.def.module, fn, as...)
            ir = prepare_ir!(IRTools.IR(m))
            ir = set_argtypes!(ir, result.linfo.specTypes.parameters)
            ir = transform(interp.ctx, ir)
            update!(src, ir)
        end
    catch e
        push!(interp, e)
    end
    src === nothing && return nothing
    validate_code_in_debug_mode(result.linfo, src, "lowered")
    return InferenceState(result, src, cached, interp)
end

function cpu_compile(ctx, mi, world)
    params = Base.CodegenParams(; track_allocations=false, code_coverage=false,
                                prefer_specsig=true,
                                lookup=@cfunction(cpu_cache_lookup, Any, (Any, UInt, UInt)))

    # generate IR
    # TODO: Instead of extern policy integrate with Orc JIT

    # populate the cache
    if cpu_cache_lookup(mi, world, world) === nothing
        cpu_infer(ctx, mi, world, world)
    end

    native_code = ccall(:jl_create_native, Ptr{Cvoid},
                        (Vector{Core.MethodInstance}, Base.CodegenParams, Cint), [mi],
                        params, 1) #=extern policy=#
    @assert native_code != C_NULL
    llvm_mod_ref = ccall(:jl_get_llvm_module, LLVM.API.LLVMModuleRef, (Ptr{Cvoid},),
                         native_code)
    @assert llvm_mod_ref != C_NULL
    llvm_mod = LLVM.Module(llvm_mod_ref)

    # get the top-level code
    code = cpu_cache_lookup(mi, world, world)

    # get the top-level function index
    llvm_func_idx = Ref{Int32}(-1)
    llvm_specfunc_idx = Ref{Int32}(-1)
    ccall(:jl_get_function_id, Nothing, (Ptr{Cvoid}, Any, Ptr{Int32}, Ptr{Int32}),
          native_code, code, llvm_func_idx, llvm_specfunc_idx)
    @assert llvm_func_idx[] != -1
    @assert llvm_specfunc_idx[] != -1

    # get the top-level function)
    llvm_func_ref = ccall(:jl_get_llvm_function, LLVM.API.LLVMValueRef,
                          (Ptr{Cvoid}, UInt32), native_code, llvm_func_idx[] - 1)
    @assert llvm_func_ref != C_NULL
    llvm_func = LLVM.Function(llvm_func_ref)
    llvm_specfunc_ref = ccall(:jl_get_llvm_function, LLVM.API.LLVMValueRef,
                              (Ptr{Cvoid}, UInt32), native_code, llvm_specfunc_idx[] - 1)
    @assert llvm_specfunc_ref != C_NULL
    llvm_specfunc = LLVM.Function(llvm_specfunc_ref)

    return llvm_specfunc, llvm_func, llvm_mod
end

function method_instance(@nospecialize(f), @nospecialize(tt), world)
    # get the method instance
    meth = which(f, tt)
    sig = Base.signature_type(f, tt)::Type
    (ti, env) = ccall(:jl_type_intersection_with_env, Any, (Any, Any), sig,
                      meth.sig)::Core.SimpleVector
    meth = Base.func_for_method_checked(meth, ti, env)
    return ccall(:jl_specializations_get_linfo, Ref{Core.MethodInstance},
                 (Any, Any, Any, UInt), meth, ti, env, world)
end

function codegen(ctx::CompilationContext, @nospecialize(f), @nospecialize(tt);
                 world=Base.get_world_counter())
    return cpu_compile(ctx, method_instance(f, tt, world), world)
end

#####
##### Optimize
#####

import Core.Compiler: optimize

function optimize(interp::MixtapeInterpreter, opt::OptimizationState,
                  params::OptimizationParams, @nospecialize(result))
    nargs = Int(opt.nargs) - 1
    ir = run_passes(opt.src, nargs, opt)
    if show_after_optimization(interp.ctx)
        fn = resolve(GlobalRef(opt.linfo.def.module, opt.linfo.def.name))
        println("(Optimized) $fn in $(opt.linfo.def.module)")
        display(ir)
    end
    return Core.Compiler.finish(opt, params, ir, result)
end

import Core.Compiler: CodeInfo, convert_to_ircode, copy_exprargs, slot2reg, compact!,
                      coverage_enabled, adce_pass!
import Core.Compiler: ssa_inlining_pass!, getfield_elim_pass!, type_lift_pass!, verify_ir,
                      verify_linetable

function run_passes(ci::CodeInfo, nargs::Int, sv::OptimizationState)
    preserve_coverage = coverage_enabled(sv.mod)
    ir = convert_to_ircode(ci, copy_exprargs(ci.code), preserve_coverage, nargs, sv)
    ir = slot2reg(ir, ci, nargs, sv)
    # TODO: Domsorting can produce an updated domtree - no need to recompute here
    ir = compact!(ir)
    ir = ssa_inlining_pass!(ir, ir.linetable, sv.inlining, ci.propagate_inbounds)
    ir = compact!(ir)
    ir = getfield_elim_pass!(ir) # SROA
    ir = adce_pass!(ir)
    ir = type_lift_pass!(ir)
    ir = compact!(ir)
    verify_ir(ir)
    verify_linetable(ir.linetable)
    return ir
end

#####
##### LLVM optimization pipeline
#####

# https://github.com/JuliaLang/julia/blob/2eb5da0e25756c33d1845348836a0a92984861ac/src/aotcompile.cpp#L603
function addTargetPasses!(pm, tm)
    add_library_info!(pm, LLVM.triple(tm))
    return add_transform_info!(pm, tm)
end

# TODO (Missing C-API):
#  - https://reviews.llvm.org/D86764 adds InstSimplify
#  - createDivRemPairs
#  - createLoopLoadEliminationPass
#  - createVectorCombinePass
# TODO (Missing LLVM.jl)
#  - AggressiveInstCombinePass

# https://github.com/JuliaLang/julia/blob/2eb5da0e25756c33d1845348836a0a92984861ac/src/aotcompile.cpp#L620
function addOptimizationPasses!(pm, tm, opt_level, lower_intrinsics, dump_native)
    constant_merge!(pm)
    if opt_level < 2
        error("opt_level less than 2 not supported")
        return
    end

    propagate_julia_addrsp!(pm)
    scoped_no_alias_aa!(pm)
    type_based_alias_analysis!(pm)
    if opt_level >= 3
        basic_alias_analysis!(pm)
    end
    cfgsimplification!(pm)
    dce!(pm)
    scalar_repl_aggregates!(pm)

    # mem_cpy_opt!(pm)

    always_inliner!(pm) # Respect always_inline

    # Running `memcpyopt` between this and `sroa` seems to give `sroa` a hard time
    # merging the `alloca` for the unboxed data and the `alloca` created by the `alloc_opt`
    # pass.

    alloc_opt!(pm)
    # consider AggressiveInstCombinePass at optlevel > 2

    instruction_combining!(pm)
    cfgsimplification!(pm)
    if dump_native
        error("dump_native not supported")
        # TODO: createMultiversoningPass
    end
    scalar_repl_aggregates!(pm)
    instruction_combining!(pm) # TODO: createInstSimplifyLegacy
    jump_threading!(pm)

    reassociate!(pm)

    early_cse!(pm)

    # Load forwarding above can expose allocations that aren't actually used
    # remove those before optimizing loops.
    alloc_opt!(pm)
    loop_rotate!(pm)
    # moving IndVarSimplify here prevented removing the loop in perf_sumcartesian(10:-1:1)
    loop_idiom!(pm)

    # TODO: Polly (Quo vadis?)

    # LoopRotate strips metadata from terminator, so run LowerSIMD afterwards
    lower_simdloop!(pm) # Annotate loop marked with "loopinfo" as LLVM parallel loop
    licm!(pm)
    julia_licm!(pm)
    # Subsequent passes not stripping metadata from terminator
    instruction_combining!(pm) # TODO: createInstSimplifyLegacy
    ind_var_simplify!(pm)
    loop_deletion!(pm)
    loop_unroll!(pm) # TODO: in Julia createSimpleLoopUnroll

    # Run our own SROA on heap objects before LLVM's
    alloc_opt!(pm)
    # Re-run SROA after loop-unrolling (useful for small loops that operate,
    # over the structure of an aggregate)
    scalar_repl_aggregates!(pm)
    instruction_combining!(pm) # TODO: createInstSimplifyLegacy

    gvn!(pm)
    mem_cpy_opt!(pm)
    sccp!(pm)

    # Run instcombine after redundancy elimination to exploit opportunities
    # opened up by them.
    # This needs to be InstCombine instead of InstSimplify to allow
    # loops over Union-typed arrays to vectorize.
    instruction_combining!(pm)
    jump_threading!(pm)
    dead_store_elimination!(pm)

    # More dead allocation (store) deletion before loop optimization
    # consider removing this:
    alloc_opt!(pm)

    # see if all of the constant folding has exposed more loops
    # to simplification and deletion
    # this helps significantly with cleaning up iteration
    cfgsimplification!(pm)
    loop_deletion!(pm)
    instruction_combining!(pm)
    loop_vectorize!(pm)
    # TODO: createLoopLoadEliminationPass
    cfgsimplification!(pm)
    slpvectorize!(pm)
    # might need this after LLVM 11:
    # TODO: createVectorCombinePass()

    aggressive_dce!(pm)

    if lower_intrinsics
        # LowerPTLS removes an indirect call. As a result, it is likely to trigger
        # LLVM's devirtualization heuristics, which would result in the entire
        # pass pipeline being re-exectuted. Prevent this by inserting a barrier.
        barrier_noop!(pm)
        lower_exc_handlers!(pm)
        gc_invariant_verifier!(pm, false)
        # Needed **before** LateLowerGCFrame on LLVM < 12
        # due to bug in `CreateAlignmentAssumption`.
        remove_ni!(pm)
        late_lower_gc_frame!(pm)
        final_lower_gc!(pm)
        # We need these two passes and the instcombine below
        # after GC lowering to let LLVM do some constant propagation on the tags.
        # and remove some unnecessary write barrier checks.
        gvn!(pm)
        sccp!(pm)
        # Remove dead use of ptls
        dce!(pm)
        lower_ptls!(pm, dump_native)
        instruction_combining!(pm)
        # Clean up write barrier and ptls lowering
        cfgsimplification!(pm)
    else
        remove_ni!(pm)
    end
    return combine_mul_add!(pm)
    # TODO: createDivRemPairs[]
end

function addMachinePasses!(pm, tm)
    demote_float16!(pm)
    return gvn!(pm)
end

function run_pipeline!(mod::LLVM.Module)
    LLVM.ModulePassManager() do pm
        addTargetPasses!(pm, tm[])
        addOptimizationPasses!(pm, tm[], 3, true, false)
        addMachinePasses!(pm, tm[])
        return run!(pm, mod)
    end
end

#####
##### JIT
#####

using GPUCompiler: GPUCompiler
import GPUCompiler: FunctionSpec

struct Entry{F,TT}
    f::F
    specfunc::Ptr{Cvoid}
    func::Ptr{Cvoid}
end

# Slow ABI
function __call(entry::Entry{F,TT}, args::TT) where {F,TT}
    args = Any[args...]
    return ccall(entry.func, Any, (Any, Ptr{Any}, Int32), entry.f, args, length(args))
end

(entry::Entry)(args...) = __call(entry, args)

const compiled_cache = Dict{UInt,Any}()

function jit(ctx, f::F, tt::TT=Tuple{}) where {F,TT<:Type}
    fspec = FunctionSpec(f, tt, false, nothing) #=name=#
    return GPUCompiler.cached_compilation(compiled_cache, fspec -> _jit(ctx, fspec), _link,
                                          fspec)::Entry{F,tt}
end

function _link(@nospecialize(fspec::FunctionSpec), (llvm_mod, func_name, specfunc_name))
    # Now invoke the JIT
    jitted_mod = compile!(orc[], llvm_mod)

    specfunc_addr = addressin(orc[], jitted_mod, specfunc_name)
    specfunc_ptr = pointer(specfunc_addr)

    func_addr = addressin(orc[], jitted_mod, func_name)
    func_ptr = pointer(func_addr)

    if specfunc_ptr === C_NULL || func_ptr === C_NULL
        @error "Compilation error" fspec specfunc_ptr func_ptr
    end

    return Entry{typeof(fspec.f),fspec.tt}(fspec.f, specfunc_ptr, func_ptr)
end

# actual compilation
function _jit(ctx, @nospecialize(fspec::FunctionSpec))
    llvm_specfunc, llvm_func, llvm_mod = codegen(ctx, fspec.f, fspec.tt)

    specfunc_name = LLVM.name(llvm_specfunc)
    func_name = LLVM.name(llvm_func)

    # set linkage to extern visible
    # otherwise `addressin` won't find them
    linkage!(llvm_func, LLVM.API.LLVMExternalLinkage)
    linkage!(llvm_specfunc, LLVM.API.LLVMExternalLinkage)

    run_pipeline!(llvm_mod)

    return (llvm_mod, func_name, specfunc_name)
end

end # module


text:
circ_guess(y,z,r,y₀,z₀) = (r>sqrt((y-y₀)^2 + (z-z₀)^2))
function circ_error(data, circ_guess_data, yb, zb, r, y₀, z₀)
    circ_guess_in(y,z) = circ_guess(y,z,r,y₀,z₀)
    for (i, yy) in enumerate(yb), (j, zz) in enumerate(zb)
        circ_guess_data[i,j] = circ_guess_in(yy, zz)
        circ_guess_data[i,j] ⊻= data[i,j]
    end
    return reduce(+, circ_guess_data)
end

ext: jl


text:
#Data Loading specifics
data_path = "/path/to/data/"
data_prefix = "data_prefix"  #If the data file name is bob.npy, this should be 'bob'

#Model Parameters
iterations = 100
hard_clustering = false  #Soft or hard assignments
initial_clusters = 1
argmax_sample_stop = 5 #Change to hard assignment from soft at iterations - argmax_sample_stop
split_stop  = 5#Stop split/merge moves at  iterations - split_stop

random_seed = nothing #When nothing, a random seed will be used.

max_split_iter = 20
burnout_period = 20
max_clusters = Inf

#Model hyperparams
α = 10.0 #Concetration Parameter
hyper_params = niw_hyperparams(1.0,
    zeros(Float32,2),
    5,
    Matrix{Float32}(I, 2, 2)*1.0)

outlier_mod = 0.05 #Concetration Parameter
outlier_hyper_params = niw_hyperparams(1.0,
    zeros(Float32,2),
    5,
    Matrix{Float32}(I, 2, 2)*1.0)

#Saving specifics:
enable_saving = true
model_save_interval = 1000
save_path = "/path/to/save/dir/"
overwrite_prec = false
save_file_prefix = "checkpoint_"


text:
""" This function generates an authentication url needed to get the tokens """
function authurl(id::AbstractString)
    AUTH * "?client_id=$id&response_type=code&redirect_uri=$(get_redirect_uri())&scope=https://outlook.office.com/mail.read%20offline_access"
end

""" This function sends a request that gets us the access & refresh tokens """
function get_tokens(client::OutlookClient, authcode::AbstractString)
    data = Dict(
        "client_id"     => clientid(client),
        "scope"         => "https://outlook.office.com/mail.read offline_access",
        "code"          => authcode,
        "redirect_uri"  => get_redirect_uri(),
        "grant_type"    => "authorization_code",
        "client_secret" => clientsecret(client),
    )

    response = post(TOKEN, data=data)
    tokens = readall(response)
    tokens_parsed = JSON.parse(tokens)

    open("credentials.json", "w+") do file
        write(file, tokens)
        close(file)
    end

    access_token = tokens_parsed["access_token"]
    refresh_token = tokens_parsed["refresh_token"]

    auth = OutlookAuthentication(access_token, refresh_token)
    client.authentication = auth
    auth
end

function refresh_tokens(client::OutlookClient)
    data = Dict(
        "client_id"     => clientid(client),
        "redirect_uri"  => get_redirect_uri(),
        "grant_type"    => "refresh_token",
        "client_secret" => clientsecret(client),
        "refresh_token" => refreshtoken(client)
    )

    response = post(TOKEN, data=data)
    if statuscode(response) != 200
        println("Couldn't authenticate. I removed your credentials.json file, restart your application and authenticate again!")
        rm("credentials.json", force=true)
        exit()
    else
        tokens = readall(response)
        tokens_parsed = JSON.parse(tokens)

        open("credentials.json", "w+") do file
            write(file, tokens)
            close(file)
        end

        access_token = tokens_parsed["access_token"]
        refresh_token = tokens_parsed["refresh_token"]

        auth = OutlookAuthentication(access_token, refresh_token)
        client.authentication = auth
    end
end

build_token(client::OutlookClient) = "Bearer " * accesstoken(client)


text:
module Problems
export gen_matching_pennies, manh_dist, hide_and_seek, nemirovski1, nemirovski2, policeman_and_burglar_matrix

import LinearAlgebra
const LA = LinearAlgebra
using Distributions
using Random

"""
   Two test problems from Nemirovski et al. "Robust stochastic approximation approach to stochastic programming"
"""
function nemirovski1(n, α=1)
    A = zeros(n, n)
    for i in 1:n
        for j in 1:n
            A[i, j] = ((i + j - 1) / ( 2n - 1 ))^α
        end
    end
    return A
end

function nemirovski2(n, α=1)
    A = zeros(n, n)
    for i in 1:n
        for j in 1:n
            A[i, j] = ((abs(i - j) + 1) / ( 2n - 1 ))^α
        end
    end
    return A
end

"""
    From Juditski & Nemirovki tutorial. Problem Policemen and Burglar.
"""
function policeman_and_burglar_matrix(n, th=0.8; seed="false")
    if seed != "false"
        Random.seed!(parse(Int, seed))
    end
    w = abs.(randn(n))
    th = 0.8
    C = reshape(abs.([i - j for i in 1:n, j in 1:n]), (n, n))
    A = w .*(1 .- exp.(-th .* C))
    return A
end

function randunif(m, n; seed="false")
    if seed != "false"
        Random.seed!(parse(Int, seed))
    end

    A = rand(m, n)
    return A
end

"""
Problem 2.24
"""
function manh_dist(m=100, n=100)
    A = zeros(m, n)
    for i in 1:m
        for j in 1:n
            y = (i - 1) / (m - 1)
            x = (j - 1) / (n - 1)
            A[i,j] = abs(x - y)
        end
    end
    return A
end

"""
    'Hide and Seek' game
"""
function hide_and_seek(m, n, param=0.2; seed="false")
    if seed != "false"
        Random.seed!(parse(Int, seed))
    end
    A = 1.0 * reshape(rand(Bernoulli(param),  m * n), (m,n))
    return A
end

end


text:
# complement = -1*
function cmpl(x::CachedResult)
    return CachedResult(-1*x.value, -1*x.signint, x.index)
end

# union = min
# Base.min(x::CachedResult, y::CachedResult) = ifelse(isless(x.value, y.value), x, y)
uni(x::CachedResult, y::CachedResult) = ifelse(isless(x.value, y.value), x, y)

# intersection = max
# Base.max(x::CachedResult, y::CachedResult) = ifelse(isless(y.value, x.value), x, y)
inters(x::CachedResult, y::CachedResult) = ifelse(isless(y.value, x.value), x, y)

# subtraction f,g = f union complement(g)
subtr(x::CachedResult, y::CachedResult) = intersection(x, complement(y))

const opDictCode = Dict(:complement=>cmpl,
                :intersection=>inters,
                :subtraction=>subtr,
                :union=>uni)

function evalf(surface::CachedSurface, cachedvalues, ind)
    val = cachedvalues[ind][surface.index]
    return CachedResult(val, 1, surface.index)
end

function leaf2code(l::CachedCSGNode, _cvals, _ind)
    ex = Expr(Symbol("="), l.sym, Expr(:call, evalf, l.data, _cvals, _ind))
    return ex
end

function node2code(l::CachedCSGNode)
    arg1 = l.children[1].sym
    op = l.data
    if op === :complement
        ex = Expr(Symbol("="), l.sym, Expr(:call, opDictCode[:complement], arg1))
        return ex
    else
        arg2 = l.children[2].sym
        ex = Expr(Symbol("="), l.sym, Expr(:call, opDictCode[op], arg1, arg2))
        return ex
    end
end

function tree2code(tree::CachedCSGNode)
    _cachedv = gensym()
    _inds = gensym()
    exprs = Array{Expr,1}(undef,0)
    for l in Set{CachedCSGNode}(Leaves(tree))
        push!(exprs, leaf2code(l, _cachedv, _inds))
    end
    for n in PostOrderDFS(tree)
        isempty(n.children) && continue
        push!(exprs, node2code(n))
    end
    return CodeWrap(exprs, [_cachedv, _inds])
end

struct CodeWrap
    Core::Array{Expr,1}
    Params::Array{Symbol,1}
end

function code2func(cwrap)
    expr = Expr(:function,
        Expr(:tuple,
        cwrap.Params[1],
        cwrap.Params[2]),
        Expr(:block,cwrap.Core...))
    return eval(expr)
end

tree2func(tree::CachedCSGNode) = code2func(tree2code(tree))


text:
module RobRecSolver

using Base.Distributed
using JuMP
using CPLEX

using MicroLogging
using ConfParser

export
    #Modules
    Experiments,

    # Types
    ProblemDescriptor,
    KnapsackProblemDescriptor,
    AssignmentProblemDescriptor,

    # Functions
    minimumAssignmentProblem,
    getAssignmentConstraints,

    minimumKnapsackProblem,
    getKnapsackConstraints,

    initialScenario,

    adversarialProblem,

    evaluationProblem,

    incrementalProblem,
    recoverableProblem,
    selectionLowerBound,

    lagrangianLowerBound,
    relaxedIncrementalProblem,

    # Utils
    loadProperties,
    getProperty,
    getProblemSize,
    getSaneComputationLimit,
    hasEqualCardinalityProperty,
    getCardinality

files = [
        "minimum_knapsack_problem",
        "minimum_assignment_problem",
        "initial_scenario",
        "evaluation_problem",
        "recoverable_problem",
        "adversarial_problem",
        "incremental_problem",
        "selection_lower_bound",
        "lagrangian_lower_bound",
        "problem_descriptors",
        "properties" # Add more files here
    ]

    for file in files
        include("$(file).jl")
    end

    # Import experiments
    include(joinpath("experiments", "Experiments.jl"))
    using .Experiments

    configure_logging(min_level=:debug)

end # module


text:
@testset "TriMesh Metrics" begin

    _mesh = load_trimesh(
        joinpath(@__DIR__, "assets/teapot.obj"),
        joinpath(@__DIR__, "assets/sphere.obj"),
    )

    @testset "Laplacian loss" begin

        verts1 = [
            0.1 0.3 0.5
            0.5 0.2 0.1
            0.6 0.8 0.7
        ]
        verts2 = [
            0.1 0.3 0.3
            0.6 0.7 0.8
            0.2 0.3 0.4
            0.1 0.5 0.3
        ]
        verts3 = [
            0.7 0.3 0.6
            0.2 0.4 0.8
            0.9 0.5 0.2
            0.2 0.3 0.4
            0.9 0.3 0.8
        ]

        faces1 = [1 2 3]
        faces2 = [
            1 2 3
            2 3 4
        ]
        faces3 = [
            2 3 1
            1 2 4
            3 4 2
            5 4 3
            5 1 2
            5 4 2
            5 3 2
        ]

        T, R = Float32, UInt32
        verts_list = [T.(verts1'), T.(verts2'), T.(verts3')]
        faces_list = [R.(faces1'), R.(faces2'), R.(faces3')]
        m = TriMesh(verts_list, faces_list)

        _edges = get_edges_packed(m)
        V = size(get_verts_packed(m), 2)
        L = zeros(V, V)
        deg = zeros(V, V)
        for i = 1:size(_edges, 1)
            L[_edges[i, 1], _edges[i, 2]] = 1
            L[_edges[i, 2], _edges[i, 1]] = 1
        end
        deg = sum(L; dims = 2)
        deg = map(x -> (x > 0 ? 1 / x : x), deg)
        for i = 1:V, j = 1:V
            if i == j
                L[i, j] = -1
            elseif L[i, j] == 1
                L[i, j] = deg[i]
            end
        end

        verts = get_verts_packed(m)
        L = L * transpose(verts)
        L = Flux3D._norm(L; dims = 2)
        @test isapprox(mean(L), laplacian_loss(m))
        @test gradient(x -> laplacian_loss(x), m) isa Tuple
    end

    @testset "Edge Loss" begin
        m = deepcopy(_mesh)
        verts = get_verts_packed(m)
        edges = get_edges_packed(m)
        v1 = verts[:, edges[:, 1]]
        v2 = verts[:, edges[:, 2]]
        loss = mean((Flux3D._norm(v1 - v2; dims = 1)) .^ 2)
        @test edge_loss(m) == loss
        @test gradient(x -> edge_loss(x), m) isa Tuple
    end

    @testset "Chamfer Loss" begin
        m = deepcopy(_mesh)
        loss = chamfer_distance(m, m)
        @test all(isapprox.(loss, 0, rtol = 1e-5, atol = 1e-2))
        @test gradient(x -> chamfer_distance(x, x), m) isa Tuple

        # naive chamfer distance

        function naive_chamfer(x, y)
            xx = sum(x .^ 2, dims = 1)
            yy = sum(y .^ 2, dims = 1)
            zz = Flux3D.Flux.batched_mul(permutedims(x, (2, 1, 3)), y)
            rx = reshape(xx, size(xx, 2), 1, :)
            ry = reshape(yy, 1, size(yy, 2), :)
            P = (rx .+ ry) .- (2 .* zz)
            nn_for_x = minimum(P; dims = 2)
            nn_for_y = minimum(P; dims = 1)
            nn_for_x = mean(nn_for_x; dims = 2)
            nn_for_y = mean(nn_for_y; dims = 2)
            loss = mean(nn_for_x) + mean(nn_for_y)
            return loss
        end

        x = rand(Float32, 3, 1000, 2)
        y = rand(Float32, 3, 500, 2)
        @test isapprox(chamfer_distance(x, y), naive_chamfer(x, y))
        grad1 = gradient(x -> chamfer_distance(x, y), x)[1]
        grad2 = gradient(x -> naive_chamfer(x, y), x)[1]
        @test isapprox(grad1, grad2, atol = 1e-2, rtol = 1e-3)
    end
end


text:
export estimate_parameters_sgd!

using LogicCircuits: num_nodes
using StatsFuns: logaddexp
using CUDA

"""
    estimate_parameters_sgd!(pc::ProbCircuit, data; lr::Float64 = 0.01, reuse_values = nothing, reuse_flows = nothing, reuse = (nothing, nothing))

One epoch SGD to learn paramters of the pc. It will run on CPU/GPU depending on where the data is located.
- `lr`: Learning Rate
- `pc`: The ProbCircuit
- 'data': training data
- 'resue_*`: Allows reusing of allocated memory to avoid extra memory allocation between epochs. If training for only one epoch don't need to specify them.
"""
function estimate_parameters_sgd!(pc::ProbCircuit, data; lr::Float64 = 0.01,
                                reuse_values = nothing, reuse_flows = nothing,
                                reuse = (nothing, nothing))
    # Construct the low-level circuit representation
    pbc = ParamBitCircuit(pc, data)
    if isgpu(data)
        pbc = to_gpu(pbc)
    end

    # Main training loop
    @assert isbatched(data)
    for idx = 1 : length(data)
        estimate_parameters_sgd!(pbc, data[idx]; lr, reuse_values, reuse_flows, reuse)
    end

    # Store the updated parameters back to `pc`
    update_pc_params_from_pbc!(pc, pbc)

    pbc.params
end

"""
    estimate_parameters_sgd!(pc::ProbCircuit, data; lr::Float64 = 0.01, reuse_values = nothing, reuse_flows = nothing, reuse = (nothing, nothing))

One epoch SGD to learn paramters of a ParamBitCircuit. It will run on CPU/GPU depending on where the data is located.
- `lr`: Learning Rate
- `pbc`: The ParamBitCircuit
- 'data': training data
- 'resue_*`: Allows reusing of allocated memory to avoid extra memory allocation between epochs. If training for only one epoch don't need to specify them.
"""
function estimate_parameters_sgd!(pbc::ParamBitCircuit, data; lr::Float64 = 0.01,
                                reuse_values = nothing, reuse_flows = nothing,
                                reuse = (nothing, nothing))

    weights = nothing
    # Extract weights from dataset
    if isweighted(data)
        # `data' is weighted according to its `weight' column
        data, weights = split_sample_weights(data)
    end

    if isgpu(data)
        if !isgpu(pbc)
            pbc.bitcircuit = to_gpu(pbc.bitcircuit)
            pbc.params = to_gpu(pbc.params)
        end
        if !isnothing(weights) && !isgpu(weights)
            weights = to_gpu(weights)
        end

        estimate_parameters_sgd_gpu!(pbc, data; lr, weights, reuse_values, reuse_flows, reuse)
    else
        estimate_parameters_sgd_cpu!(pbc, data; lr, weights, reuse_values, reuse_flows, reuse)
    end

    pbc.params
end

function estimate_parameters_sgd_cpu!(pbc::ParamBitCircuit, data; lr::Float64 = 0.01,
                                    weights = nothing, reuse_values = nothing, reuse_flows = nothing,
                                    reuse = (nothing, nothing))
    bc = pbc.bitcircuit
    params = pbc.params

    # Forward/upward pass
    values = marginal_all(pbc, data, reuse_values)

    # Compute the gradient (this part can be customized)
    log_grads = similar!(reuse[1], typeof(params), num_examples(data))
    @inbounds @views log_grads[:] .= zero(eltype(values)) .- values[:, num_nodes(pbc.bitcircuit)]

    param_grads = similar!(reuse[2], typeof(params), size(params)...)

    # Backward/downward pass
    flows = backprop_flows_down_cpu(pbc, log_grads, values, param_grads, reuse_flows; weights)

    # Apply gradients
    apply_gradients_cpu!(pbc, param_grads; lr)

    pbc.params, values, flows, (log_grads, param_grads)
end

function estimate_parameters_sgd_gpu!(pbc::ParamBitCircuit, data; lr::Float64 = 0.01,
                                    weights = nothing, reuse_values = nothing, reuse_flows = nothing,
                                    reuse = (nothing, nothing))
    bc = pbc.bitcircuit
    params = pbc.params

    # Forward/upward pass
    values = marginal_all(pbc, data, reuse_values)

    # Compute the gradient (this part can be customized)
    log_grads = similar!(reuse[1], typeof(params), num_examples(data))
    @inbounds @views log_grads[:] .= zero(eltype(values)) .- values[:, num_nodes(pbc.bitcircuit)]

    param_grads = similar!(reuse[2], typeof(params), size(params)...)

    # Backward/downward pass
    flows = backprop_flows_down_gpu(pbc, log_grads, values, param_grads, reuse_flows; weights)

    # Apply gradients
    apply_gradients_gpu!(pbc, param_grads; lr)

    pbc.params, values, flows, (log_grads, param_grads)
end

######################################
# Bit circuit backprop downward pass #
######################################

function backprop_flows_down_cpu(pbc::ParamBitCircuit, log_grads, values, param_grads, reuse = nothing; weights = nothing)
    flows = similar!(reuse, typeof(values), size(values)...)

    if weights !== nothing
        @inbounds @views log_grads .+= log.(weights)
    end

    backprop_flows_down_layers_cpu(pbc, log_grads, flows, values, param_grads)

    flows # Return for reuse
end

function backprop_flows_down_layers_cpu(pbc::ParamBitCircuit, log_grads::Vector, flows::Matrix, values::Matrix,
                                    param_grads::Vector)
    @assert flows !== values
    bc::BitCircuit = pbc.bitcircuit
    els = bc.elements
    for layer in Iterators.reverse(bc.layers)
        ## Compute gradients
        Threads.@threads for dec_id in layer
            par_start = @inbounds bc.nodes[3, dec_id]
            if iszero(par_start)
                if dec_id == num_nodes(bc)
                    # Assign (log) gradient to the root node
                    @inbounds @views flows[:, dec_id] .= log_grads
                end
                # no parents, ignore (can happen for false/true node and root)
            else
                par_end = @inbounds bc.nodes[4, dec_id]
                for j = par_start : par_end
                    par = @inbounds bc.parents[j] # Parent #j of `dec_id`
                    grandpa = @inbounds els[1, par] # Parent of `par`
                    sib_id = sibling(els, par, dec_id) # Sibling of `dec_id` w.r.t. `par`

                    single_child = has_single_child(bc.nodes, grandpa)
                    if single_child
                        if j == par_start
                            @inbounds @views @. flows[:, dec_id] = flows[:, grandpa] + values[:, sib_id]
                        else
                            @inbounds @simd for i = 1 : size(flows, 1)
                                flows[i, dec_id] = logaddexp(flows[i, dec_id], flows[i, grandpa] + values[i, sib_id])
                            end
                        end
                    else
                        θ = eltype(flows)(pbc.params[par])
                        if j == par_start
                            @inbounds @views @. flows[:, dec_id] = flows[:, grandpa] + values[:, sib_id] + θ
                        else
                            @inbounds @simd for i = 1 : size(flows, 1)
                                flows[i, dec_id] = logaddexp(flows[i, dec_id], flows[i, grandpa] + values[i, sib_id] + θ)
                            end
                        end
                    end
                    # Compute gradient only once
                    if sib_id > dec_id
                        param_grads[par] = sum(exp.(flows[:, grandpa] + values[:, dec_id] + values[:, sib_id]))
                    end
                end
            end
        end
    end
end

function apply_gradients_cpu!(pbc::ParamBitCircuit, param_grads::Vector; lr::Float64 = 0.01)
    bc::BitCircuit = pbc.bitcircuit
    nodes = bc.nodes
    params = pbc.params
    for layer in Iterators.reverse(bc.layers)
        ## Apply gradients
        Threads.@threads for dec_id in layer
            if !has_single_child(nodes, dec_id)
                ele_start_id = nodes[1, dec_id]
                ele_end_id = nodes[2, dec_id]

                sum_grads = 0.0
                @inbounds for ele_id = ele_start_id : ele_end_id
                    sum_grads += param_grads[ele_id]
                end
                sum_params = -Inf
                @inbounds for ele_id = ele_start_id : ele_end_id
                    params[ele_id] += lr * param_grads[ele_id] / (sum_grads + 1e-8)
                    sum_params = logaddexp(sum_params, params[ele_id])
                end
                @inbounds for ele_id = ele_start_id : ele_end_id
                    params[ele_id] -= sum_params
                end
            end
        end
    end
end

function backprop_flows_down_gpu(pbc::ParamBitCircuit, log_grads::CuVector, values::CuMatrix,
                                 param_grads::CuVector, reuse =  nothing; dec_per_thread = 8,
                                 log2_threads_per_block = 7, weights = nothing)
    flows = similar!(reuse, typeof(values), size(values)...)
    if weights !== nothing
        @inbounds @views log_grads .+= CUDA.log.(weights)
    end

    bc::BitCircuit = pbc.bitcircuit
    @inbounds @views param_grads .= zero(eltype(param_grads))
    param_grads_device = CUDA.cudaconvert(param_grads)

    CUDA.@sync for layer in Iterators.reverse(bc.layers)
        num_examples = size(values, 1)
        num_decision_sets = length(layer)/dec_per_thread

        kernel = @cuda name="backprop_flows_down_layers_cuda" launch=false backprop_flows_down_layers_cuda(layer, bc.nodes, bc.elements, bc.parents, pbc.params, log_grads, flows, values, param_grads_device)
        config = launch_configuration(kernel.fun)
        threads, blocks = balance_threads_2d(num_examples, num_decision_sets, config.threads)
        kernel(layer, bc.nodes, bc.elements, bc.parents, pbc.params, log_grads, flows, values, param_grads_device
            ; threads, blocks)
    end

    flows # Return for reuse
end

function backprop_flows_down_layers_cuda(layer, nodes, elements, parents, params, log_grads, flows, values, param_grads)
    index_x = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    index_y = (blockIdx().y - 1) * blockDim().y + threadIdx().y
    stride_x = blockDim().x * gridDim().x
    stride_y = blockDim().y * gridDim().y

    for k = index_x : stride_x : size(values, 1)
        for i = index_y : stride_y : length(layer)
            dec_id = @inbounds layer[i]
            if dec_id == size(nodes, 2)
                # Assign (log) gradient to the root node
                flow = log_grads[k]
            else
                par_start = @inbounds nodes[3, dec_id]
                flow = typemin(eltype(flows)) # log(0)
                if !iszero(par_start)
                    par_end = @inbounds nodes[4, dec_id]
                    for j = par_start : par_end
                        par = @inbounds parents[j]
                        grandpa = @inbounds elements[1, par]
                        sib_id = sibling(elements, par, dec_id)

                        g_flow = @inbounds flows[k, grandpa]
                        d_value = @inbounds values[k, dec_id]
                        s_value = @inbounds values[k, sib_id]
                        if has_single_child(nodes, grandpa)
                            edge_flow = g_flow + s_value
                        else
                            θ = eltype(flows)(params[par])
                            edge_flow = g_flow + s_value + θ
                        end
                        flow = logsumexp_cuda(flow, edge_flow)

                        # Compute gradient only once
                        if sib_id > dec_id
                            grad::Float64 = CUDA.exp(g_flow + d_value + s_value)
                            CUDA.@atomic param_grads[par] += grad
                        end
                    end
                end
            end
            @inbounds flows[k, dec_id] = flow
        end
    end
    return nothing
end

function apply_gradients_gpu!(pbc::ParamBitCircuit, param_grads::CuVector; lr::Float64 = 0.01)
    bc::BitCircuit = pbc.bitcircuit

    CUDA.@sync for layer in Iterators.reverse(bc.layers)
        kernel = @cuda name="apply_gradients_cuda" launch=false apply_gradients_cuda(layer, bc.nodes, param_grads, pbc.params, lr::Float64)
        config = launch_configuration(kernel.fun)
        threads =  Base.min(length(layer), config.threads)
        blocks = cld(length(layer), threads)
        kernel(layer, bc.nodes, param_grads, pbc.params, lr::Float64; threads, blocks)
    end
end

function apply_gradients_cuda(layer, nodes, param_grads, params, lr::Float64)
    index_x = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    stride_x = blockDim().x * gridDim().x

    for i = index_x : stride_x : length(layer)
        dec_id = @inbounds layer[i]
        single_child = has_single_child(nodes, dec_id)
        if !single_child
            ele_start_id = nodes[1, dec_id]
            ele_end_id = nodes[2, dec_id]

            sum_grads = zero(eltype(param_grads))
            @inbounds for ele_id = ele_start_id : ele_end_id
                sum_grads += param_grads[ele_id]
            end
            sum_params = -Inf
            @inbounds for ele_id = ele_start_id : ele_end_id
                params[ele_id] += lr * param_grads[ele_id] / (sum_grads + 1e-8)
                sum_params = logsumexp_cuda(sum_params, params[ele_id])
            end
            @inbounds for ele_id = ele_start_id : ele_end_id
                params[ele_id] -= sum_params
            end
        end
    end
end


text:
module TimeSeriesA5P1

#Use this to print
#=
weave(Pkg.dir("$(pwd())\\FamaMacbeth.jl"),
  informat="script",
  out_path = "$(pwd())\\FamaMacbethOUT.html",
  doctype = "md2html")
=#

if pwd() ∉ LOAD_PATH
  push!(LOAD_PATH,pwd())
end

using DataFrames, Distributions, StatsBase, GZip, JLD, ForwardDiff,
  Gadfly, JuMP, NLopt, HCubature, StaticArrays, CTNumerical,
  Measures, Formatting, ParallelAccelerator

importall CTIO, CTReg, CTStat

const DATA_PATH = pwd() * "\\data" #path where we store the data
const OUTPUT_PATH = pwd() * "\\output"
const YAHOO_DATE_FORMAT = "m/d/yyyy"
const FRED_DATE_FORMAT = "m/d/yyyy"
const PPI_DATE_FORMAT = "m/d/yyyy"
const SP500_NAME = "GSPC"
const DBV_NAME = "DBV"
const FRED_NAME = "CPI"
const PPI_NAME = "PPI"

const FOOTER_NAME = "footer.tex" #these are just headers for the tex tables
const HEADER_NAME = "header.tex"
const DECIMALS = 2
const NUM_FORMAT = "{:.$(DECIMALS)f}"
const PlotContainer = Union{Plot,Gadfly.Compose.Context,Vector{Plot}} #holds graphs
const OVERRIDE_STAR_STRINGS = ["\\ostar","\\dstar","\\tstar"]
const COLORS = ["DodgerBlue", "OrangeRed", "Green",
  "Brown", "Black", "BlueViolet"]

n2s(x::T where T<:Real) = num2Str(x, 4, Ints=true)

#holds parameters for a measurement error simulation
mutable struct AR1Measurement
  c::Vector{T} where T<:Real
  ϕ::T where T<:Real
  dist::T where T<:MultivariateDistribution
  ARMA11::ARMA
end

#main constructor, also creates an ARMA11
AR1Measurement(ϕ::Float64 ;c::Vector{Float64}=zeros(2), Σ::Matrix{Float64}=eye(2),
  dist::T where T<:MultivariateDistribution = MvNormal(2,Σ))::AR1Measurement =
    AR1Measurement(c,ϕ,dist,
    ARMA(ϕ=[ϕ], ϑ=[-ϕ^2/2],σ²=ones(2)'*Σ*ones(2)))

mutable struct AR1MeasurementRet
  X::Vector{T} where T<:Real
  Y::Vector{T} where T<:Real
  ε::Matrix{T} where T<:Real
end

#simulates returns
function AR1MeasurementRet(Θ::AR1Measurement, N::Int;
    x₀::T=0.0,  Y::Vector{T} = Vector{typeof(x₀)}(N), X::Vector{T} = similar(Y),
    ε::Matrix{Float64} = rand(Θ.dist, N))::AR1MeasurementRet where T<:Real

  X[1] = Θ.ϕ*x₀ + ε[1,1]
  Y[1] = X[1] + ε[2,1]

  for i ∈ 2:N
      X[i] = Θ.ϕ*X[i-1] + ε[1,i]
      Y[i] = X[i] + ε[2,i]
  end

  return AR1MeasurementRet(X,Y,ε)
end

#calculate the analytical auto correlations
ARMeasurementAC(Θ::AR1Measurement, k::Int)::Float64 = Θ.ϕ^k*(2-Θ.ϕ^2.)/(1.0-Θ.ϕ^2)
ARMeasurementAC(Θ::AR1Measurement, r::Range)::Vector{Float64} =
  ((k::Int)->ARMeasurementAC.(Θ,k)).(r)

#problem specific graphs
function visualizeA5P1(Θ::AR1Measurement,
  ACRange::Range)::Void

  plots::Vector{PlotContainer} = Vector{PlotContainer}()
  plotNames::Vector{String} = Vector{String}()

  Gadfly.push_theme(:default)

  push!(plotNames, "ACF Graphs")
  push!(plots, plot(
    layer(x=collect(ACRange), y=ARMeasurementAC(Θ,ACRange),  Geom.line, Geom.point,
      Theme(default_color=COLORS[1])),
    #Guide.manual_color_key("Legend", ["Actual", "AR1"], COLORS[1:2]),
    Guide.xlabel("Lag"),Guide.ylabel("Autocorrelation"),
    Guide.title("AutoCorrelation")))

  #write the graphs
  for i::Int ∈ 1:length(plots)
    draw(SVG("$OUTPUT_PATH\\$(plotNames[i]).svg", 9inch, 7inch), plots[i])
    #display(plots[i])
  end

  return nothing
end

#script for the problem
function A5P1Script(;refreshData::Bool = true)::Void

  ϕ::Float64 = 0.85 #Process parameter
  Σ::Matrix{Float64} = eye(2) #covariance matrix for shocks
  ACRange::Range = 1:25 #range of autocorrelations
  N::Int = 50_000

  Θ::AR1Measurement = AR1Measurement(ϕ, c=zeros(2), Σ=Σ, dist = MvNormal(Σ))
  returns::AR1MeasurementRet = AR1MeasurementRet(Θ,N)
  println("Analytical AC: $(ARMeasurementAC(Θ,ACRange))")
  #println("Y: $(returns.ε)")
  println("Mean ε: $(mean(returns.ε))")
  println("Σ: $(cov(returns.X,returns.Y))")
  println("σ^2: $(var(returns.Y))")
  println("Analytical: $((2-Θ.ϕ^2)/(1-Θ.ϕ^2))")

  #Algorithms: :LD_LBFGS* , :LN_COBYLA, :LN_BOBYQA
  #  :LD_SLSQPx , :LD_TNEWTON*, :LD_VAR2/:LD_VAR1 ,
  #  :LD_MMA , :GN_ESCH, :GN_DIRECT , GN_DIRECT_L ,
  # GN_ISRES, :LN_NELDERMEAD*, :LN_SBPLX, :LD_LBFGS , :GD_STOGO
  #GN_CRS2_LM

  startingARMA = ARMA(c=0.0, ϕ=[0.85], ϑ=[-0.35])

  est::ARMAEstimate =
      ARMAEstimate(ARMA(1,1), returns.Y,
        alg = :LD_TNEWTON, verbose = true, bound=5.0, boundσ=5.0, maxTime=120., suppressC=false)
  println("σ²: $(est.Θ.σ²)")
  println("c: $(est.Θ.c)")
  println("ϕ: $(est.Θ.ϕ)")
  println("ϑ: $(est.Θ.ϑ)")
  println("VParam: $((2est.Θ.ϑ[1]^2+4est.Θ.ϕ[1]*est.Θ.ϑ[1]+2)/(1-est.Θ.ϕ[1]^2))")
  println("Unc: $(var(est.YHat))")

  #visualizeA5P1(Θ,ACRange)

  #visualizeA4P2!()

  return nothing
end

@time begin
  srand(11)
  A5P1Script(refreshData = true)
end

end

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
module CornerPoints

using Check3DPackings

export compute_free_cps

"""
  compute_free_cps(xyzpqr, LWH)
  compute_free_cps(x, y, z, p, q, r, L, W, H)

# Arguments
- `xyzpqr`: a sextuple of collections, the first three are the position of the
  boxes in the x/y/z axis, the last three are the sizes of the boxes in th
  x/y/z axis.
- `LWH`: is a triple of the size of the container in the x/y/z axis.
# Returns
- A nonuple of vectors. The first three vectors are the x/y/z positions of
  the corner points. The fourth, fifth, and sixth vectors are the minimum
  length/width/height a box needs to have to be placed in that corner point.
  The last three vectors have the indexes of the boxes which provide the
  x/y/z-planes that form the corner point. The n+1, n+2, and n+3 values
  in those last three vectors denote, respectively, the back, left, and
  bottom faces of the container (that do not pertain to any box but are
  also used to form corner points).
# Examples

"""
function compute_free_cps(xyzpqr, LWH)
  compute_free_cps(xyzpqr..., LWH...)
end

function compute_free_cps(x, y, z, p, q, r, L, W, H)
  xyzpqr = (x, y, z, p, q, r)
  @assert (isone ∘ length ∘ unique ∘ map)(length, xyzpqr)
  @assert all(map(a -> all(v -> v >= 1, a), xyzpqr))
  for i = 1:length(x), j = (i+1):length(x)
    @assert !has_overlap(
      x[i], y[i], z[i], p[i], q[i], r[i],
      x[j], y[j], z[j], p[j], q[j], r[j]
    )
  end

  # Create the container walls.
  map(a->push!(a...), zip(xyzpqr, (0, 1, 1, 1, W, H)))
  map(a->push!(a...), zip(xyzpqr, (1, 0, 1, L, 1, H)))
  map(a->push!(a...), zip(xyzpqr, (1, 1, 0, L, W, 1)))
  n = length(x)

  # Change of notation for arrays with container walls.
  l = p; w = q; h = r

  # returned data structure
  xyzpqr′ijks = (Vector{typeof(L)}(), Vector{typeof(L)}(), Vector{typeof(L)}(),
    Vector{typeof(L)}(), Vector{typeof(L)}(), Vector{typeof(L)}(),
    Vector{Int64}(), Vector{Int64}(), Vector{Int64}())

  for k = 1:n
    z′ = z[k] + h[k]
    for j = 1:n
      y′ = y[j] + w[j]
      (z[j] + h[j] <= z′ || y[k] + w[k] <= y′) && continue
      for i = 1:n
        x′ = x[i] + l[i]
        ( x[k] + l[k] <= x′ ||
          x[j] + l[j] <= x′ ||
          z[i] + h[i] <= z′ ||
          y[i] + w[i] <= y′) && continue
        p′ = max(x′, x[k], x[j]) - x′ + 1
        q′ = max(y′, y[k], y[i]) - y′ + 1
        r′ = max(z′, z[j], z[i]) - z′ + 1
        for o = 1:(n-3) # the walls can never block a corner point
          has_overlap(
            x[o], y[o], z[o], l[o], w[o], h[o], x′, y′, z′, p′, q′, r′
          ) && @goto after_pushing_cp
        end
        xyzpqr′ijk = (x′, y′, z′, p′, q′, r′, i, j, k)
        map(a->push!(a...), zip(xyzpqr′ijks, xyzpqr′ijk))
        @label after_pushing_cp
      end
    end
  end

  # Remove the container walls from the list of boxes.
  map(a -> resize!(a, n - 3), xyzpqr)

  xyzpqr′ijks
end

end


text:
@testset "Embedding" begin
    ex1 = embedding_example(10, 3, 1)

    @testset "Dimensions are correct" begin
        @test size(ex1)[1] == 8
        @test size(ex1)[2] == 3
    end

    @testset "Does embedding form invariant set?" begin
        @test isa(invariantset(ex1), Bool)
    end

end

max_stars_repo_name: kahaaga/SimplexSplitting

max_issues_repo_name: kahaaga/SimplexSplitting

max_forks_repo_name: kahaaga/SimplexSplitting


text:
# multi-point crossover operator
struct MultiPointCrossover <: CrossoverOperator{2,2}
    points::UInt64

    MultiPointCrossover(numpoints::Integer=2) = new(numpoints)
end
numpoints(xover::MultiPointCrossover) = xover.points

function apply!(xover::MultiPointCrossover,
                targets::AbstractVector{<:AbstractVector},
                target_indices::AbstractVector{<:Integer},
                pop,
                parent_indices::AbstractVector{<:Integer})

    p1, p2 = map(i -> view(pop.individuals, :, i), parent_indices)
    c1, c2 = map(t -> view(t, :), targets)

    @inbounds copyto!(c1, p1)
    @inbounds copyto!(c2, p2)

    # sample xover positions uniformly at random and swap
    positions= sample(1:numdims(pop), numpoints(xover), replace=false)
    @inbounds c1[positions] = p2[positions]
    @inbounds c2[positions] = p1[positions]

    return targets
end

# mutation operator
struct BinaryFlipMutation <: MutationOperator
    ρ::Float64 # mutation rate
    BinaryFlipMutation(ρ::Float64=.01) = new(ρ)
end
ρ(mo::BinaryFlipMutation) = mo.ρ

function apply!(mo::BinaryFlipMutation,
                            target::AbstractVector{<:Integer},
                            target_index::Integer)

    amount_positions = ceil(Int, ρ(mo) * length(target))
    positions = sample(1:length(target), amount_positions, replace=false)
    @inbounds foreach(p -> target[p] = target[p] > 0 ? 0 : 1, positions)

    return target
end


text:
################################################################################
#
#             EllCrv/Heights.jl : Height functions on elliptic curves
#
# This file is part of Hecke.
#
# Copyright (c) 2015, 2016: Claus Fieker, Tommy Hofmann
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# (C) 2016 Tommy Hofmann
# (C) 2016 Robin Ammon
# (C) 2016 Sofia Brenner
# (C) 2022 Jeroen Hanselman
#
################################################################################
export local_height, real_height, canonical_height, naive_height

################################################################################
#
#  Naive Height
#
################################################################################

@doc Markdown.doc"""
    naive_height(P::EllCrvPt{fmpq}, prec) -> ArbField

return the naive height of a point $P$ on an elliptic curve defined over $\mathbb{Q}$.
"""
function naive_height(P::EllCrvPt{fmpq}, prec = 100)
  attempt = 1

  while true
    R = ArbField(attempt*prec)
    x = P.coordx
    p = numerator(x)
    q = denominator(x)
    result = log(R(max(abs(p),abs(q))))
    if radiuslttwopower(result, -prec)
      return result
    end
    attempt = 2*attempt
  end
end

################################################################################
#
#  Local Height at finite prime
#
################################################################################

# Equal to Magma command with Renormalization flag. In Magma the default is to add a factor of (1/6)log Δv at every place.

#TODO: Fine-tune precision

@doc Markdown.doc"""
    local_height(P::EllCrvPt{fmpq}, p::Int, prec::Int) -> ArbField

Computes the local height of a point $P$ on an elliptic curve defined over
$\mathbb{Q}$ at the prime $p$.
"""
function local_height(P::EllCrvPt{fmpq}, p, prec = 100)

  E = P.parent
  F = minimal_model(E)
  phi = isomorphism(E,F)

  P = phi(P)

  d = ceil(Int, prec*log(10,2))

  p = FlintZZ(p)

  x = P.coordx
  y = P.coordy

  a1, a2, a3, a4, a6 = map(numerator,(a_invars(F)))

  b2,b4,b6,b8,c4,c6 = get_b_c_integral(F)

  delta = discriminant(E)

  A = 3*x^2 + 2*a2*x + a4 - a1*y
  B = 2*y + a1*x + a3 # = psi2(P)
  C = 3*x^4 + b2 * x^3 + 3*b4*x^2 + 3*b6*x + b8 # = psi3(P)
  L = 0

  attempt = 2
  while true
    R = ArbField(attempt*prec)

    if !is_finite(P)
      return R(0)
    end

    if (A != 0 && valuation(A, p) <= 0) || (B != 0 && valuation(B, p) <= 0)
      if x != 0
        L = max(0, -valuation(x, p))
      else
        L = 0
      end
    elseif (c4 != 0 && valuation(c4, p) == 0)
      N = valuation(delta, p)
      if B == 0
        M = N//2
        L = M*(M - N)//N
      else
        M = min(valuation(B, p), N//2)
        L = M*(M - N)//N
      end
    elseif ( C == 0 || ( C != 0 && B != 0 && valuation(C, p) >= 3*valuation(B, p)))
      L = -2*valuation(B, p)//3
    else
      L = -valuation(C, p)//4
    end
    result = L*log(R(p))
    if Hecke.radiuslttwopower(result,-prec)
      return result
    end
      attempt = 2*attempt
  end
end

################################################################################
#
#  Real Height
#
################################################################################

#Precision is given in bits (as Real Field also works this way), but maybe this should be changed. In Magma precision is given in decimals

@doc Markdown.doc"""
    real_height(P::EllCrvPt{fmpq}, prec::Int) -> ArbField

Computes the real height of a point $P$ on an elliptic curve defined
over $\mathbb{Q}$.
"""
function real_height(P::EllCrvPt{fmpq}, prec = 100)

  attempt = 3
  d = ceil(Int, prec*log(10,2))

  E = P.parent
  F = minimal_model(E)
  phi = isomorphism(E,F)

  P = phi(P)

  a1, a2, a3, a4, a6 = map(numerator,(a_invars(F)))

  b2, b4, b6, b8, c4, c6 = get_b_c_integral(F)
  H = max(4, abs(b2), 2*abs(b4), 2*abs(b6), abs(b8))
  _b2 = b2-12
  _b4 = b4-b2+6
  _b6 = b6-2*b4+b2-4
  _b8 = b8-3*b6+3*b4-b2+3

  N = ceil(Int, (5/3)*d + 1/2 + (3/4)*log(7+(4/3)*log(H)))

  while true

    R = ArbField(attempt*prec)
    x = R(P.coordx)
    y = R(P.coordy)

    if abs(x)<0.5
      t = 1/(x+1)
      beta = 0
    else
      t = 1/x
      beta = 1
    end

    mu = -log(abs(t))
    f = 1

    for n in (0:N)
      f = f/4
      if beta==1
        w = b6*t^4+2*b4*t^3+b2*t^2+4*t
        z = 1-b4*t^2-2*b6*t^3-b8*t^4
        zw = z+w
      else
        w = _b6*t^4+2*_b4*t^3+_b2*t^2+4*t
        z = 1-_b4*t^2-2*_b6*t^3-_b8*t^4
        zw = z+w
      end
      if abs(w) <= 2*abs(z)
        mu = mu+f*log(abs(z))
        t = w/z
      else
        mu = mu+f*log(abs(zw))
        t = w/zw
        beta = 1-beta
      end
    end
    if isfinite(mu) & (radius(mu)<R(10)^(-d))
    # Algorithm is only precise up to d decimals
      add_error!(mu,R(10)^(-d))
      return(mu)
    else
      attempt = 2*attempt
    end
  end
end

################################################################################
#
#  Néron-Tate Height
#
################################################################################
@doc Markdown.doc"""
    neron_tate_height(P::EllCrvPt{fmpq}, prec::Int) -> ArbField

Compute the Néron-Tate height (or canonical height) of a point $P$ on an
elliptic curve defined over $\mathbb{Q}$.
"""
function neron_tate_height(P::EllCrvPt{fmpq}, prec = 100)
  return canonical_height(P,prec)
end

@doc Markdown.doc"""
    canonical_height(P::EllCrvPt{fmpq}, prec::Int) -> ArbField

Compute the Néron-Tate height (or canonical height) of a point $P$ on an
elliptic curve defined over $\mathbb{Q}$.
"""
function canonical_height(P::EllCrvPt{fmpq}, prec = 100)
  attempt = 1

  while true
    R = Nemo.RealField(attempt*prec)
    E = P.parent
    disc = discriminant(E)
    d = (denominator(P.coordx))
    h = real_height(P, attempt*prec) + log(R(d))
    plist = bad_primes(E)

    for p in plist
      if !divides(d,p)[1]
       h = h + local_height(P,p, attempt*prec)
      end
    end
    if Hecke.radiuslttwopower(h,-prec)
      return h
    else
      attempt = 2*attempt
    end
  end
end


text:
module TestCore
using Test
using Setfield
using Setfield: compose, get_update_op
using ConstructionBase: ConstructionBase
using StaticNumbers: static

struct T
    a
    b
end

struct TT{A,B}
    a::A
    b::B
end

@testset "get_update_op" begin
    @test get_update_op(:(&=)) === :(&)
    @test get_update_op(:(^=)) === :(^)
    @test get_update_op(:(-=)) === :(-)
    @test get_update_op(:(%=)) === :(%)
    @test_throws ArgumentError get_update_op(:(++))
    @test_throws ArgumentError get_update_op(:(<=))
end

@testset "@set!" begin
    a = 1
    @set a = 2
    @test a === 1
    @set! a = 2
    @test a === 2

    t = T(1, T(2,3))
    @set t.b.a = 20
    @test t === T(1, T(2,3))

    @set! t.b.a = 20
    @test t === T(1,T(20,3))

    a = 1
    @set! a += 10
    @test a === 11
    nt = (a=1,)
    @set! nt.a = 5
    @test nt === (a=5,)
end

@testset "@set" begin

    t = T(1, T(2, T(T(4,4),3)))
    s = @set t.b.b.a.a = 5
    @test t === T(1, T(2, T(T(4,4),3)))
    @test s === T(1, T(2, T(T(5, 4), 3)))
    @test_throws ArgumentError @set t.b.b.a.a.a = 3

    t = T(1,2)
    @test T(1, T(1,2)) === @set t.b = T(1,2)
    @test_throws ArgumentError @set t.c = 3

    t = T(T(2,2), 1)
    s = @set t.a.a = 3
    @test s === T(T(3, 2), 1)

    t = T(1, T(2, T(T(4,4),3)))
    s = @set t.b.b = 4
    @test s === T(1, T(2, 4))

    t = T(1,2)
    s = @set t.a += 1
    @test s === T(2,2)

    t = T(1,2)
    s = @set t.b -= 2
    @test s === T(1,0)

    t = T(10, 20)
    s = @set t.a *= 10
    @test s === T(100, 20)

    t = T(2,1)
    s = @set t.a /= 2
    @test s === T(1.0,1)

    t = T(1, 2)
    s = @set t.a <<= 2
    @test s === T(4, 2)

    t = T(8, 2)
    s = @set t.a >>= 2
    @test s === T(2, 2)

    t = T(1, 2)
    s = @set t.a &= 0
    @test s === T(0, 2)

    t = T(1, 2)
    s = @set t.a |= 2
    @test s === T(3, 2)

    t = T((1,2),(3,4))
    @set t.a[1] = 10
    s1 = @set t.a[1] = 10
    @test s1 === T((10,2),(3,4))
    i = 1
    si = @set t.a[i] = 10
    @test s1 === si
    se = @set t.a[end] = 20
    @test se === T((1,20),(3,4))
    se1 = @set t.a[end-1] = 10
    @test s1 === se1

    s1 = @set t.a[static(1)] = 10
    @test s1 === T((10,2),(3,4))
    i = 1
    si = @set t.a[static(i)] = 10
    @test s1 === si

    t = @set T(1,2).a = 2
    @test t === T(2,2)

    t = (1, 2, 3, 4)
    @test (@set t[length(t)] = 40) === (1, 2, 3, 40)
    @test (@set t[length(t) ÷ 2] = 20) === (1, 20, 3, 4)
end

struct UserDefinedLens <: Lens end

struct LensWithTextPlain <: Lens end
Base.show(io::IO, ::MIME"text/plain", ::LensWithTextPlain) =
    print(io, "I define text/plain.")

@testset "show it like you build it " begin
    i = 3
    @testset for item in [
            @lens _.a
            @lens _[1]
            @lens _[:a]
            @lens _["a"]
            @lens _[static(1)]
            @lens _[static(1), static(1 + 1)]
            @lens _.a.b[:c]["d"][2][static(3)]
            @lens _
            @lens first(_)
            @lens last(first(_))
            @lens last(first(_.a))[1]
            UserDefinedLens()
            (@lens _.a) ∘ UserDefinedLens()
            UserDefinedLens() ∘ (@lens _.b)
            (@lens _.a) ∘ UserDefinedLens() ∘ (@lens _.b)
            (@lens _.a) ∘ LensWithTextPlain() ∘ (@lens _.b)
        ]
        buf = IOBuffer()
        show(buf, item)
        item2 = eval(Meta.parse(String(take!(buf))))
        @test item === item2
    end
end

function test_getset_laws(lens, obj, val1, val2)

    # set ∘ get
    val = get(obj, lens)
    @test set(obj, lens, val) == obj

    # get ∘ set
    obj1 = set(obj, lens, val1)
    @test get(obj1, lens) == val1

    # set idempotent
    obj12 = set(obj1, lens, val2)
    obj2 = set(obj, lens, val2)
    @test obj12 == obj2
end

function test_modify_law(f, lens, obj)
    obj_modify = modify(f, obj, lens)
    old_val = get(obj, lens)
    val = f(old_val)
    obj_setfget = set(obj, lens, val)
    @test obj_modify == obj_setfget
end

@testset "lens laws" begin
    obj = T(2, T(T(3,(4,4)), 2))
    i = 2
    for lens ∈ [
            @lens _.a
            @lens _.b
            @lens _.b.a
            @lens _.b.a.b[2]
            @lens _.b.a.b[i]
            @lens _.b.a.b[static(2)]
            @lens _.b.a.b[static(i)]
            @lens _.b.a.b[end]
            @lens _.b.a.b[identity(end) - 1]
            @lens _
        ]
        val1, val2 = randn(2)
        f(x) = (x,x)
        test_getset_laws(lens, obj, val1, val2)
        test_modify_law(f, lens, obj)
    end
end

@testset "equality & hashing" begin
    # singletons (identity and property lens) are egal
    for (l1, l2) ∈ [
        @lens(_) => @lens(_),
        @lens(_.a) => @lens(_.a)
    ]
        @test l1 === l2
        @test l1 == l2
        @test hash(l1) == hash(l2)
    end

    # composite and index lenses are structurally equal
    for (l1, l2) ∈ [
        @lens(_[1]) => @lens(_[1])
        @lens(_.a[2]) => @lens(_.a[2])
        @lens(_.a.b[3]) => @lens(_.a.b[3])
    ]
        @test l1 == l2
        @test hash(l1) == hash(l2)
    end

    # inequality
    for (l1, l2) ∈ [
        @lens(_[1]) => @lens(_[2])
        @lens(_.a[1]) => @lens(_.a[2])
        @lens(_.a[1]) => @lens(_.b[1])
    ]
        @test l1 != l2
    end

    # Hash property: equality implies equal hashes, or in other terms:
    # lenses either have equal hashes or are unequal
    # Because collisions can occur theoretically (though unlikely), this is a property test,
    # not a unit test.
    random_lenses = (@lens(_.a[rand(Int)]) for _ in 1:1000)
    @test all((hash(l2) == hash(l1)) || (l1 != l2)
              for (l1, l2) in zip(random_lenses, random_lenses))

    # Lenses should hash differently from the underlying tuples, to avoid confusion.
    # To account for potential collisions, we check that the property holds with high
    # probability.
    @test count(hash(@lens(_[i])) != hash((i,)) for i = 1:1000) > 900

    # Same for tuples of tuples (√(1000) ≈ 32).
    @test count(hash(@lens(_[i][j])) != hash(((i,), (j,))) for i = 1:32, j = 1:32) > 900
end

@testset "type stability" begin
    o1 = 2
    o22 = 2
    o212 = (4,4)
    o211 = 3
    o21 = TT(o211, o212)
    o2 = TT(o21, o22)
    obj = TT(o1, o2)
    @assert obj === TT(2, TT(TT(3,(4,4)), 2))
    i = 1
    for (lens, val) ∈ [
          ((@lens _.a           ),   o1 ),
          ((@lens _.b           ),   o2 ),
          ((@lens _.b.a         ),   o21),
          ((@lens _.b.a.b[2]    ),   4  ),
          ((@lens _.b.a.b[i+1]  ),   4  ),
          ((@lens _.b.a.b[static(2)]   ),   4  ),
          ((@lens _.b.a.b[static((i+1))]),  4  ),
          ((@lens _.b.a.b[static(2)]   ),   4.0),
          ((@lens _.b.a.b[static((i+1))]),  4.0),
          ((@lens _.b.a.b[end]),     4.0),
          ((@lens _.b.a.b[end÷2+1]), 4.0),
          ((@lens _             ),   obj),
          ((@lens _             ),   :xy),
        ]
        @inferred get(obj, lens)
        @inferred set(obj, lens, val)
        @inferred modify(identity, obj, lens)
    end
end

@testset "IndexLens" begin
    l = @lens _[]
    @test l isa Setfield.IndexLens
    x = randn()
    obj = Ref(x)
    @test get(obj, l) == x

    l = @lens _[][]
    @test l.outer isa Setfield.IndexLens
    @test l.inner isa Setfield.IndexLens
    inner = Ref(x)
    obj = Base.RefValue{typeof(inner)}(inner)
    @test get(obj, l) == x

    obj = (1,2,3)
    l = @lens _[1]
    @test l isa Setfield.IndexLens
    @test get(obj, l) == 1
    @test set(obj, l, 6) == (6,2,3)

    l = @lens _[1:3]
    @test l isa Setfield.IndexLens
    @test get([4,5,6,7], l) == [4,5,6]
end

@testset "DynamicIndexLens" begin
    l = @lens _[end]
    @test l isa Setfield.DynamicIndexLens
    obj = (1,2,3)
    @test get(obj, l) == 3
    @test set(obj, l, true) == (1,2,true)

    l = @lens _[end÷2]
    @test l isa Setfield.DynamicIndexLens
    obj = (1,2,3)
    @test get(obj, l) == 1
    @test set(obj, l, true) == (true,2,3)

    two = 2
    plusone(x) = x + 1
    l = @lens _.a[plusone(end) - two].b
    obj = (a=(1, (a=10, b=20), 3), b=4)
    @test get(obj, l) == 20
    @test set(obj, l, true) == (a=(1, (a=10, b=true), 3), b=4)

    if Setfield.HAS_BEGIN_INDEXING
        # Need to keep this in a separate file since `begin` won't parse
        # on older Julia versions.
        include("dynamiclens_begin.jl")
    end
end

@testset "StaticNumbers" begin
    obj = (1, 2.0, '3')
    l = @lens _[static(1)]
    @test (@inferred get(obj, l)) === 1
    @test (@inferred set(obj, l, 6.0)) === (6.0, 2.0, '3')
    l = @lens _[static(1 + 1)]
    @test (@inferred get(obj, l)) === 2.0
    @test (@inferred set(obj, l, 6)) === (1, 6, '3')
    n = 1
    l = @lens _[static(3n)]
    @test (@inferred get(obj, l)) === '3'
    @test (@inferred set(obj, l, 6)) === (1, 2.0, 6)

    l = @lens _[static(1):static(3)]
    @test get([4,5,6,7], l) == [4,5,6]

    @testset "complex example (sweeper)" begin
        sweeper_with_const = (
            model = (1, 2.0, 3im),
            axis = (@lens _[static(2)]),
        )

        sweeper_with_noconst = @set sweeper_with_const.axis = @lens _[2]

        function f(s)
            a = sum(set(s.model, s.axis, 0))
            for i in 1:10
                a += sum(set(s.model, s.axis, i))
            end
            return a
        end

        @test (@inferred f(sweeper_with_const)) == 66 + 33im
        @test_broken (@inferred f(sweeper_with_noconst)) == 66 + 33im
    end
end

mutable struct M
    a
    b
end

@testset "IdentityLens" begin
    id = @lens _
    @test compose(id, id) === id
    obj1 = M(1,1)
    obj2 = M(2,2)
    @test obj2 === set(obj1, id, obj2)
    la = @lens _.a
    @test compose(id, la) === la
    @test compose(la, id) === la
end

struct ABC{A,B,C}
    a::A
    b::B
    c::C
end

@testset "type change during @set (default constructorof)" begin
    obj = TT(2,3)
    obj2 = @set obj.b = :three
    @test obj2 === TT(2, :three)
end

# https://github.com/tkf/Reconstructables.jl#how-to-use-type-parameters
struct B{T, X, Y}
    x::X
    y::Y
    B{T}(x::X, y::Y = 2) where {T, X, Y} = new{T, X, Y}(x, y)
end
ConstructionBase.constructorof(::Type{<: B{T}}) where T = B{T}

@testset "type change during @set (custom constructorof)" begin
    obj = B{1}(2,3)
    obj2 = @set obj.y = :three
    @test obj2 === B{1}(2, :three)
end

@testset "text/plain show" begin
    @testset for lens in [
        LensWithTextPlain()
        (@lens _.a) ∘ LensWithTextPlain()
        LensWithTextPlain() ∘ (@lens _.b)
        (@lens _.a) ∘ LensWithTextPlain() ∘ (@lens _.b)
    ]
        @test occursin("I define text/plain.", sprint(show, "text/plain", lens))
    end

    @testset for lens in [
        UserDefinedLens()
        (@lens _.a) ∘ UserDefinedLens()
        UserDefinedLens() ∘ (@lens _.b)
        (@lens _.a) ∘ UserDefinedLens() ∘ (@lens _.b)
    ]
        @test sprint(show, lens) == sprint(show, "text/plain", lens)
    end
end

@testset "Named Tuples" begin
    t = (x=1, y=2)
    @test (@set t.x =2) === (x=2, y=2)
    @test (@set t.x += 2) === (x=3, y=2)
    @test (@set t.x =:hello) === (x=:hello, y=2)
    l = @lens _.x
    @test get(t, l) === 1

    # do we want this to throw an error?
    @test_throws ArgumentError (@set t.z = 3)
end

struct CustomProperties
    _a
    _b
end

function ConstructionBase.setproperties(o::CustomProperties, patch::NamedTuple)
    CustomProperties(get(patch, :a, getfield(o, :_a)),
                     get(patch, :b, getfield(o, :_b)))

end

ConstructionBase.constructorof(::Type{CustomProperties}) = error()

@testset "setproperties overloading" begin
    o = CustomProperties("A", "B")
    o2 = @set o.a = :A
    @test o2 == CustomProperties(:A, "B")
    o3 = @set o.b = :B
    @test o3 == CustomProperties("A", :B)
end

@testset "issue #83" begin
    @test_throws ArgumentError Setfield.lensmacro(identity, :(_.[:a]))
end

@testset "@lens and ∘" begin
    @test @lens(∘()) === @lens(_)
    @test @lens(∘(_.a)) === @lens(_.a)
    @test @lens(∘(_.a, _.b)) === @lens(_.a) ∘ @lens(_.b)
    @test @lens(∘(_.a, _.b, _.c)) === Setfield.compose(@lens(_.a), @lens(_.b), @lens(_.c))

    @test @lens(∘(_[1])) === @lens(_[1])
    @test @lens(∘(_[1], _[2])) === @lens(_[1]) ∘ @lens(_[2])
    @test @lens(∘(_[1], _[2], _[3])) === Setfield.compose(@lens(_[1]), @lens(_[2]), @lens(_[3]))

    @test @lens(_ ∘ (_[1] ∘ _.a) ∘ first(_)) == @lens(_) ∘ (@lens(_[1]) ∘ @lens(_.a)) ∘ @lens(first(_))
end

@testset "@lens ∘ and \$" begin
    lbc = @lens _.b.c
    @test @lens($lbc)== lbc
    @test @lens(_.a ∘ $lbc) == @lens(_.a) ∘ lbc
    @test @lens(_.a ∘ $lbc ∘ _[1] ∘ $lbc) == @lens(_.a) ∘ lbc ∘ @lens(_[1]) ∘ lbc
end

end


text:
function
initial_condition(
    x::Real,
    ε::Real     = 10^(-1.5),
    x_tp::Real  = 0.1,
)
    ξ = (x - x_tp) / ε;
    return ((x + 1) + (x - 5) * exp(-3 * ξ)) / (1 + exp(-3 * ξ))
end

function
default_parameters()

    exp = quote
        a       = 0.0;
        b       = 1.0;
        Tₑ      = 0.5;
        t₀      = 0.0;
        N       = 100;
        M       = 250;
        h       = (b - a) / N;
        τ       = (Tₑ - t₀) / M;
        Xₙ      = [h * n for n in 0:N];
        Tₘ      = [τ * m for m in 0:M];
        ε       = 10^(-1.5);
        x_tp    = 0.1;
        u₀      = [IPIL8.initial_condition(x, ε, x_tp) for x in Xₙ];
        ulₘ     = [-5.0 for m in 0:M];
        urₘ     = [ 2.0 for m in 0:M];
        u       = zeros(N+1, M+1);    # Выделим память заблаговременно
    end;

    return exp;
end

function θ(x)
    return x < 0 ? 0.0 : 1.0;
end

function δ(x, w = 0.001)
    out = 0.0;
    p = abs(x / w);

    if p <= 1/2
        out = 2 / w * (1 - p - 4 * p^2 + 4 * p^3)
    elseif 1/2 < p <= 1
        out = 2 / w * (1 - 11/3 * p + 4 * p^2 - 4/3 * p^3)
    end

    return out;
end


text:
ENV["GKS_WSTYPE"] = "100"

using GR
using Printf
using BenchmarkTools

w, h = (800, 600)

image = rand(UInt32, w, h)  # allocate the memory
mem = Printf.@sprintf("%p", pointer(image))
conid = Printf.@sprintf("!%dx%d@%s.mem",  w, h, mem[3:end])

@btime begin
    beginprint(conid)
    plot(rand(1000), size=(w, h))
    endprint()
end

emergencyclosegks()

delete!(ENV, "GKS_WSTYPE")

mwidth, mheight, width, height = inqdspsize()
size = 400 * mwidth / width

setwsviewport(0, size, 0, 3/4 * size)
setwswindow(0, 1, 0, 3/4)

setviewport(0, 1, 0, 3/4)
setwindow(0, 1, 0, 3/4)

drawimage(0, 1, 0, 3/4, w, h, image)
updatews()


text:
include("../emi.jl")
using EMI
using EMI.Draw, EMI.Gmsh

loop = Loop([CircleArc(Point(1, 0), Point(0, 0), Point(0, 1)),
             CircleArc(Point(0, 1), Point(0, 0), Point(-1, 0)),
             CircleArc(Point(-1, 0), Point(0, 0), Point(0, -1)),
             CircleArc(Point(0, -1), Point(0, 0), Point(1, 0))])

canvas = Canvas()
canvas = canvas + loop

set_bbox!(canvas, 0.2, 0.3)
gmsh_script(canvas, 0.2)

max_stars_repo_name: MiroK/emi-cylinders

max_issues_repo_name: MiroK/emi-cylinders

max_forks_repo_name: MiroK/emi-cylinders


text:
# length: interfaces.jl
# endof: interfaces.jl
# size: interfaces.jl
# getindex: interfaces.jl
# start: interfaces.jl
# done: interfaces.jl
# next: interfaces.jl
# indexed_next: Not exported
# eltype: interfaces.jl
# ntuple
# map
# heads: Not exported
# tails: Not exported
# isequal: interfaces.jl
# ==: interfaces.jl
# hash: works for Node
# isless: interfaces.jl
# isempty: interfaces.jl
# revargs: Not exported
# reverse
# sum
# prod
# all
# any


text:
# NOTE: Do not construct a PT such that PTs can be fused together or used as-is
# in a way such that there aren't functions for splitting and merging them in
# pt_lib.jl. Note that each splitting and merging function in pt_lib.jl is for
# specific locations and so, for example, a Div should not be used on a value
# with CSV location unless there is a splitting function for that.

# Block() = PartitionType(Dict("name" => "Block"))
# Block(dim) = PartitionType(Dict("name" => "Block", "dim" => dim))
# BlockBalanced() = PartitionType(Dict("name" => "Block", "balanced" => true))
# BlockBalanced(dim) =
#     PartitionType(Dict("name" => "Block", "dim" => dim, "balanced" => true))
# BlockUnbalanced() = PartitionType(Dict("name" => "Block", "balanced" => false))
# BlockUnbalanced(dim) =
#     PartitionType(Dict("name" => "Block", "dim" => dim, "balanced" => false))

# Div() = PartitionType(Dict("name" => "Replicate", "dividing" => true))
# Replicated() = PartitionType(Dict("name" => "Replicate", "replicated" => true))
# Reducing(op) = PartitionType(Dict("name" => "Replicate", "replicated" => false, "reducer" => to_jl_value(op)))

# TODO: Generate AtMost and ScaledBy constraints in handling filters and joins
# that introduce data skew and in other operations that explicitly don't

Replicating() = PartitionType("name" => "Replicating", f->ScaleBy(f, 1.0))
Replicated() = Replicating() & PartitionType("replication" => "all", "reducer" => nothing)
# TODO: Add Replicating(f) to the below if needed for reducing operations on
# large objects such as unique(df::DataFrame)

# TODO: Determine whether the `"reducer" => nothing` should be there
Divided() = Replicating() & PartitionType("dividing" => true)
Syncing() = Replicating() & PartitionType("replication" => "one", "reducer" => nothing) # TODO: Determine whether this is really needed
Reducing(op) = Replicating() & PartitionType("replication" => nothing, "reducer" => to_jl_value(op), "with_key" => false)
ReducingWithKey(op) = Replicating() & PartitionType("replication" => nothing, "reducer" => to_jl_value(op), "with_key" => true)
# TODO: Maybe replace banyan_reduce_size_by_key with an anonymous function since that actually _can_ be ser/de-ed
# or instead make there be a reducing type that passes in the key to the reducing functions so it can reduce by that key
# ReducingSize() = PartitionType("replication" => "one", "reducer" => "banyan_reduce_size_by_key")

Distributing() = PartitionType("name" => "Distributing")
Blocked(; along = nothing) =
    if isnothing(along)
        PartitionType("name" => "Distributing", "distribution" => "blocked")
    else
        PartitionType(
            "name" => "Distributing",
            "distribution" => "blocked",
            "key" => along,
        )
    end
Grouped() = PartitionType("name" => "Distributing", "distribution" => "grouped")
# Blocked(;balanced) = PartitionType("name" => "Distributing", "distribution" => "blocked", "balanced" => balanced)
# Grouped(;balanced) = PartitionType("name" => "Distributing", "distribution" => "grouped", "balanced" => balanced)

ScaledBySame(;as) = PartitionType(f -> ScaleBy(f, 1.0, as))
Drifted() = Distributing() & PartitionType("id" => "!")
Balanced() =
    Distributing() & PartitionType("balanced" => true, f -> ScaleBy(f, 1.0))
Unbalanced(; scaled_by_same_as = nothing) =
    if isnothing(scaled_by_same_as)
        Distributing() & PartitionType("balanced" => false)
    else
        Unbalanced() & ScaledBySame(as = scaled_by_same_as)
    end

# These functions (along with `keep_sample_rate`) allow for managing memory
# usage in annotated code. `keep_sample_rate` allows for setting the sample
# rate as it changes from value to value. Some operations such as joins
# actually require a change in sample rate so propagating this information is
# important and must be done before partition annotations are applied (in
# `partitioned_using`). In the partition annotation itself, we sometimes want
# to set constraints on how we scale the memory usage based on how much skew
# is introduced by an operation. Some operations not only change the sample
# rate but also introduce skew and so applying these constraints is important.
# FilteredTo and FilteredFrom help with constraining skew when it is introduced
# through data filtering operations while MutatedTo and MutatedFrom allow for
# propagatng skew for operations where the skew is unchanged. Balanced data
# doesn't have any skew and Balanced and balanced=true help to make this clear.
# TODO: Remove this if we don't need
# MutatedRelativeTo(f, mutated_relative_to) = PartitionType(ScaleBy(1.0, f, mutated_relative_to))
# MutatedTo(f, mutated_to) = MutatedRelativeTo(f, mutated_to)
# MutatedFrom(f, mutated_from) = MutatedRelativeTo(f, mutated_from)

Distributed(args...; kwargs...) = Blocked(args...; kwargs...) | Grouped(args...; kwargs...)
Partitioned(args...; kwargs...) = Distributed(args...; kwargs...) | Replicated()

function Blocked(
    f::AbstractFuture;
    along = :,
    balanced = nothing,
    filtered_from = nothing,
    filtered_to = nothing,
    scaled_by_same_as = nothing,
)
    parameters = Dict()
    constraints = PartitioningConstraints()

    # Prepare `along`
    if along isa Colon
        along = sample(f, :axes)
        # TODO: Ensure that axes returns [1] for DataFrame and axes for Array
        # while keys returns keys for DataFrame and axes for Array
    end
    along = to_vector(along)
    # TODO: Maybe assert that along isa Vector{String} or Vector{Symbol}

    # Create PTs for each axis that can be used to block along
    pts::Vector{PartitionType} = []
    for axis in first(along, 4)
        # Handle combinations of `balanced` and `filtered_from`/`filtered_to`
        for b in (isnothing(balanced) ? [true, false] : [balanced])
            # Initialize parameters
            parameters = Dict("key" => axis, "balanced" => b)
            constraints = PartitioningConstraints()

            # Create `ScaleBy` constraints
            if b
                push!(constraints.constraints, ScaleBy(f, 1.0))
                # TODO: Add an AtMost constraint in the case that there are very few rows.
                # That AtMost constraint would only go here in Blocked
            else
                if !isnothing(filtered_from)
                    filtered_from = to_vector(filtered_from)
                    factor, from = maximum(filtered_from) do ff
                        # If 100 elements get filtered to 20 elements and the
                        # original data was block-partitioned in a balanced
                        # way, the result may all be on one partition in the
                        # msot extreme case (not balanced at all) and so we
                        # should adjust the memory usage of the result by
                        # multiplying it by the size of the original / the size
                        # of the result (100 / 20 = 5).
                        (sample(ff, :memory_usage) / sample(f, :memory_usage), filtered_from)
                    end

                    # The factor will be infinite or NaN if either what we are
                    # filtering from or filtering to has an empty sample. In
                    # that case, it wouldn't make sense to have a `ScaleBy`
                    # constraint. A value with an empty sample must either be
                    # replicated (if it is from an empty dataset) or
                    # grouped/blocked but not balanced (since if it were
                    # balanced, we might try to use its divisions - which would
                    # be empty - for other PTs and think that they to are balanced).
                    if factor != Inf && factor != NaN
                        push!(constraints.constraints, ScaleBy(f, factor, from))
                    end
                elseif !isnothing(filtered_to)
                    filtered_to = to_vector(filtered_to)
                    factor, to = maximum(filtered_to) do ft
                        (sample(f, :memory_usage) / sample(ft, :memory_usage), filtered_to)
                    end

                    # The factor will be infinite or NaN if either what we are
                    # filtering from or filtering to has an empty sample. In
                    # that case, it wouldn't make sense to have a `ScaleBy`
                    # constraint. A value with an empty sample must either be
                    # replicated (if it is from an empty dataset) or
                    # grouped/blocked but not balanced (since if it were
                    # balanced, we might try to use its divisions - which would
                    # be empty - for other PTs and think that they to are balanced).
                    if factor != Inf && factor != NaN
                        push!(constraints.constraints, ScaleBy(f, factor, to))
                    end
                elseif !isnothing(scaled_by_same_as)
                    push!(constraints.constraints, ScaleBy(f, 1.0, scaled_by_same_as))
                end
            end

            # Append new PT to PT union being produced
            push!(pts, PartitionType(parameters, constraints))
        end
    end

    # Return the resulting PT union that can then be passed into a call to `pt`
    # which would in turn result in a PA union
    Blocked() & pts
end

# NOTE: A reason to use Grouped for element-wise computation (with no
# filtering) is to allow for the input to be re-balanced. If you just use
# Any then there wouldn't be any way to re-balance right before the
# computation. Grouped allows the input to have either balanced=true or
# balanced=false and if balanced=true is chosen then a cast may be applied.

function Grouped(
    f::AbstractFuture;
    # Parameters for splitting into groups
    by = nothing,
    balanced = nothing,
    rev = nothing,
    # Options to deal with skew
    filtered_from = nothing,
    filtered_to = nothing,
    scaled_by_same_as = nothing,
)
    # Prepare `by`
    by = if isnothing(by)
        sample(f, :groupingkeys)
    elseif by isa Colon
        sample(f, :keys)
    else
        by
    end
    by = to_vector(by)
    if !(by isa Colon)
        intersect!(by, sample(f, :keys))
    end

    # Create PTs for each key that can be used to group by
    pts::Vector{PartitionType} = []
    for (i, key) in enumerate(first(by, 8))
        # Handle combinations of `balanced` and `filtered_from`/`filtered_to`
        for b in (isnothing(balanced) ? [true, false] : [balanced])
            parameters = Dict("key" => key, "balanced" => b)
            constraints = PartitioningConstraints()

            # Create `ScaleBy` constraint and also compute `divisions` and
            # `AtMost` constraint if balanced
            if b
                # Set divisions
                # TODO: Change this if `divisions` is not a `Vector{Tuple{Any,Any}}`
                parameters["divisions"] = to_jl_value(sample(f, :statistics, key, :divisions))
                max_ngroups = sample(f, :statistics, key, :max_ngroups)

                # Set flag for reversing the order of the groups
                if !isnothing(rev)
                    parameters["rev"] = rev
                end

                # Add constraints
                # In the future if the element size can be really big (like a multi-dimensional array
                # that is very wide or data frame with many columns), we may want to have a constraint
                # where the partition size must be larger than that minimum element size.

                # Note that if the sample is empty, the maximum # of groups
                # will be zero and so this AtMost constraint will cause the PA
                # to fail to be used. This is expected. You can't have a PA
                # where empty data is balanced. If the empty data arises
                # because of an empty dataset being queried/processed, we
                # should be using replication. If the empty data arises because
                # of highly selective filtering, we will filter from some data
                # that _is_ balanced.
                push!(constraints.constraints, AtMost(max_ngroups, f))
                push!(constraints.constraints, ScaleBy(f, 1.0))

                # TODO: Make AtMost only accept a value (we can support PT references in the future if needed)
                # TODO: Make scheduler check that the values in AtMost or ScaledBy are actually present to ensure
                # that the constraint can be satisfied for this PT to be used
            else
                # TODO: Support joins
                if !isnothing(filtered_from)
                    filtered_from = to_vector(filtered_from)
                    filtered_from_futures = [ff isa Pair ? first(ff) : ff for ff in filtered_from]
                    factor, from = maximum(filtered_from) do ff
                        # Get key to use for filtering
                        ff, fby = ff isa Pair ? ff : (ff => by)
                        fby = fby isa Colon ? sample(ff, :keys) : by
                        fby = to_vector(Symbol.(fby))
                        fkey = fby[i]

                        # IF what wea re filtering to is empty, we don't know
                        # anything about the skew of data being filtered.
                        # Everything could be in a single partition or evenly
                        # distributed and the result would be the same. If this
                        # PT is ever fused with some matching PT that _is_
                        # balanced and has divisions specified, then that will
                        # be becasue of some filtering going on in which case a
                        # `ScaleBy` constraint will be enforced.

                        # Handling empty data:
                        # If empty data is considered balanced, we should
                        # replicate everything. Otherwise, there should be some
                        # other data in the filtering pipeline that is
                        # balanced.
                        # - disallow balanced grouping of empty dataset
                        # - filtering to an empty dataset - no ScaleBy constraint
                        # - filtering from an empty dataset - no ScaleBy constraint
                        # - maybe add in a ScaleBy(-1j) to prevent usage of an empty data

                        # Compute the amount to scale memory usage by based on data skew
                        min_filtered_to = sample(f, :statistics, key, :min)
                        max_filtered_to = sample(f, :statistics, key, :max)
                        # divisions_filtered_from = sample(ff, :statistics, key, :divisions)
                        ff_percentile = sample(ff, :statistics, fkey, :percentile, min_filtered_to, max_filtered_to)
                        (1 / ff_percentile, filtered_from_futures)
                    end

                    # If the sample of what we are filtering into is empty, the
                    # factor will be infinity. In that case, we shouldn't be
                    # creating a ScaleBy constraint.
                    if factor != Inf
                        push!(constraints.constraints, ScaleBy(f, factor, from))
                    end
                elseif !isnothing(filtered_to)
                    filtered_to = to_vector(filtered_to)
                    # TODO: Revisit this and ensure it's okay to just take the
                    # maximum and we don't have to actually multiply all of the
                    # factors by which this fails to scale.
                    filtered_to_futures = [ft isa Pair ? first(ft) : ft for ft in filtered_to]
                    factor, to = maximum(filtered_to) do ft
                        # Get key to use for filtering
                        ft, fby = ft isa Pair ? ft : (ft => by)
                        fby = fby isa Colon ? sample(ff, :keys) : by
                        fby = to_vector(Symbol.(fby))
                        fkey = fby[i]

                        # Compute the amount to scale memory usage by based on data skew
                        min_filtered_to = sample(ft, :statistics, fkey, :min)
                        max_filtered_to = sample(ft, :statistics, fkey, :max)
                        # f_divisions = sample(f, :statistics, key, :divisions)
                        f_percentile = sample(f, :statistics, key, :percentile, min_filtered_to, max_filtered_to)
                        (1 / f_percentile, filtered_to_futures)
                    end
                    # TODO: Return all filtered_from/filtered_to but with the appropriate factor and without pairs

                    # If the sample of what we are filtering into is empty, the
                    # factor will be infinity. In that case, we shouldn't be
                    # creating a ScaleBy constraint.
                    if factor != Inf
                        push!(constraints.constraints, ScaleBy(f, factor, to))
                    end
                elseif !isnothing(scaled_by_same_as)
                    push!(constraints.constraints, ScaleBy(f, 1.0, scaled_by_same_as))
                end
            end

            push!(pts, PartitionType(parameters, constraints))
        end
    end

    # If there are no PTs, ensure that we at least have one impossible PT. This
    # PT doesn't need to have divisions and can be unbalanced but the important
    # thing is that we assign an AtMost-zero constraint which will prevent a PA
    # containing this from being used. This is important because we can't group
    # data on keys that don't belong to it.
    if isempty(pts)
        push!(pts, PartitionType("key" => nothing, "balanced" => false, f -> AtMost(0, f)))
    end

    Grouped() & pts
end

max_stars_repo_name: banyan-team/banyan-julia

max_issues_repo_name: banyan-team/banyan-julia

max_forks_repo_name: banyan-team/banyan-julia


text:
"""
    read(A::AbstractGeoArray)

`read` will move any `AbstractGeoArray` completely to memory, as a `GeoArray`.
```
"""
Base.read(A::AbstractGeoArray) = GeoArray(A)

"""
    read(A::AbstractGeoStack)

`read` will move any `AbstractGeoStack` completely to memory, as a `GeoStack` of `GeoArray`.
"""
Base.read(st::AbstractGeoStack) = GeoStack(st)

"""
    read(A::AbstractGeoSeries)

`read` will move any `AbstractGeoSeries` completely to memory, as a `GeoSeries`
of `GeoStack` or `GeoArray`.
"""
Base.read(ser::AbstractGeoSeries) = rebuild(ser; data=[read(ser[i]) for i in eachindex(ser)])