text:
module Infiltrator

using REPL, UUIDs, InteractiveUtils
using REPL.LineEdit: getproperty
using REPL.LineEdit

export @infiltrate, @exfiltrate, @withstore, safehouse, exfiltrated, infiltrate

REPL_HOOKED = Ref{Bool}(false)
function __init__()
  clear_store!(store)
  if VERSION >= v"1.5.0-DEV.282"
    if isdefined(Base, :active_repl_backend)
        pushfirst!(Base.active_repl_backend.ast_transforms, ast_transformer(gensym()))
        REPL_HOOKED[] = true
    else
      atreplinit() do repl
        @async begin
          iter = 0
          # wait for active_repl_backend to exist
          while !isdefined(Base, :active_repl_backend) && iter < 20
            sleep(0.05)
            iter += 1
          end
          if isdefined(Base, :active_repl_backend)
            pushfirst!(Base.active_repl_backend.ast_transforms, ast_transformer(gensym()))
            REPL_HOOKED[] = true
          end
        end
      end
    end
  end
end

"""
    @infiltrate cond = true

`@infiltrate` sets an infiltration point (or breakpoint).

When the infiltration point is hit, it will drop you into an interactive REPL session that
lets you inspect local variables and the call stack as well as execute aribtrary statements
in the context of the current functions module.

This macro also accepts an optional argument `cond` that must evaluate to a boolean,
and then this macro will serve as a "conditinal breakpoint", which starts inspections only
when its condition is `true`.
```
"""
macro infiltrate(cond = true)
  quote
    if $(esc(cond))
      $(start_prompt)(@__MODULE__, Base.@locals, $(String(__source__.file)), $(__source__.line))
    end
  end
end

"""
    infiltrate(mod, locals, file, line)

Functional form of `@infiltrate`. Use this to conditionally infiltrate package code without
using e.g. Revise (because this version is valid during precompilation).

Typically used as
```julia
if isdefined(Main, :Infiltrator)
  Main.infiltrate(@__MODULE__, Base.@locals, @__FILE__, @__LINE__)
end
```
"""
function infiltrate(mod, locals, file, line)
  start_prompt(mod, locals, string(file), line)
end

"""
    @exfiltrate

Assigns all local variables into the global storage.
"""
macro exfiltrate()
  quote
    for (k, v) in Base.@locals
      try
        Core.eval(getfield(getfield($(@__MODULE__), :store), :store), Expr(:(=), k, QuoteNode(v)))
      catch err
        println(stderr, "Assignment to store variable failed.")
        Base.display_error(stderr, err, catch_backtrace())
      end
    end
  end
end

const prompt_color = if Sys.iswindows()
  "\e[33m"
else
  "\e[38;5;166m"
end

# these are used for easier testing
const TEST_TERMINAL_REF = Ref{Any}(nothing)
const TEST_REPL_REF = Ref{Any}(nothing)
const TEST_NOSTACK = Ref{Any}(false)

const CHECK_TASK = Ref{Bool}(true)

mutable struct Session
  store::Module
  exiting::Bool
  disabled::Set
end
function Base.show(io::IO, s::Session)
  n = length(get_store_names(s))
  print(io, "Safehouse ($(n) variable$(n == 1 ? "" : "s"))")
end
function Base.show(io::IO, ::MIME"text/plain", s::Session)
  names = get_store_names(s)
  n = length(names)
  println(io, "Safehouse with $(n) stored variable$(n == 1 ? "" : "s"):")
  for (i, (var, val)) in enumerate(names)
    if i > 6
      println(io, "...")
      break
    end
    one_line_show(io, var, val)
  end
end
function Base.getproperty(sp::Session, s::Symbol)
  m = getfield(sp, :store)
  if isdefined(m, s)
    getproperty(m, s)
  else
    throw(UndefVarError(s))
  end
end
Base.propertynames(s::Session) = keys(get_store_names(s))

"""
    store
    safehouse
    exfiltrated

Global storage for storing values while `@infiltrate`ing or `@exfiltrate`ing.
"""
const store = Session(Module(), false, Set())
const safehouse = store
const exfiltrated = store

"""
    @withstore ex

Evaluates the expression `ex` in the context of the global store.

Mainly intended for interactive use, as changes to the store's state will not
propagate into the returned expression.
"""
macro withstore(ex)
  withstore(ex)
end

function withstore(ex)
  ns = get_store_names(store)
  return Expr(:let,
    Expr(:block, map(x->Expr(:(=), x...), [(k, maybe_quote(v)) for (k, v) in ns])...),
    Expr(:block, ex)
  )
end

"""
    clear_disabled!(s = safehouse)

Clear all disabled infiltration points.
"""
function clear_disabled!(s = store)
  empty!(getfield(s, :disabled))
  return nothing
end

"""
    end_session!(s = safehouse)

End this infiltration session (reverts the effect of `@exit` in the `debug>` REPL).

Only needs to be manually called on Julia versions prior to 1.5.
"""
function end_session!(s::Session = store)
  setfield!(s, :exiting, false)
  return nothing
end

"""
    clear_store!(s = safehouse)

Reset the store used for global symbols.
"""
clear_store!(s::Session = store) = set_store!(s, Module())

"""
    set_store!(s = safehouse, m::Module)

Set the module backing the store `s`.
"""
set_store!(m::Module) = set_store!(store, m)
function set_store!(s::Session, m::Module)
  setfield!(s, :store, m)
  nothing
end

function start_prompt(mod, locals, file, fileline;
                        terminal = TEST_TERMINAL_REF[],
                        repl = TEST_REPL_REF[],
                        nostack = TEST_NOSTACK[]
                      )
  getfield(store, :exiting) && return
  (file, fileline) in getfield(store, :disabled) && return

  if terminal === nothing || repl === nothing
    if isdefined(Base, :active_repl) && isdefined(Base.active_repl, :t) && isdefined(Base, :active_repl_backend)
      repl = Base.active_repl
      terminal = Base.active_repl.t
    else
      println("Infiltrator.jl needs a fully-functional Julia REPL.")
      return
    end
  end

  io = Base.pipe_writer(terminal)

  trace = stacktrace()
  start = something(findlast(x -> x.func === Symbol("start_prompt"), trace), 0) + 2
  last = something(findfirst(x -> x.func === StackTraces.top_level_scope_sym, trace),
                   length(trace))
  trace = trace[start:last]

  if CHECK_TASK[] && !Base.active_repl_backend.in_eval
    println(io, "Cannot infiltrate while the REPL is idle. This typically happens when an async task is running in the background.")
    if length(trace) > 0
      println(io, "Disabling infiltration point at $(trace[1]).")
    else
      println(io, "Cannot infiltrate while the REPL is idle. Disabling this infiltration point.")
    end
    push!(getfield(store, :disabled), ((file, fileline)))
    return
  end

  if length(trace) > 0
    if nostack
      println(io, "Infiltrating <unknown>:")
    else
      print(io, "Infiltrating ")
      print_verbose_stackframe(io, trace[1])
    end
  else
    println(io, "Infiltrating top-level frame:")
  end
  println(io)
  debugprompt(mod, locals, trace, terminal, repl, nostack, file=file, fileline=fileline)
  println(io)
end

function show_help(io)
  println(io, """
    Code entered is evaluated in the current functions module. Note that you cannot change local
    variables, but can assign to globals in a permanent store module.

    The following commands are special cased:
      - `?`: Print this help text.
      - `@trace`: Print the current stack trace.
      - `@locals`: Print local variables. `@locals x y` only prints `x` and `y`.
      - `@exfiltrate`: Save all local variables into the store. `@exfiltrate x y` saves `x` and `y`;
        this variant can also exfiltrate variables defined in the `infil>` REPL.
      - `@toggle`: Toggle infiltrating at this `@infiltrate` spot (clear all with `Infiltrator.clear_disabled!()`).
      - `@continue`: Continue to the next infiltration point or exit (shortcut: Ctrl-D).
      - `@doc symbol`: Get help for `symbol` (same as in the normal Julia REPL).
      - `@exit`: Stop infiltrating for the remainder of this session and exit (on Julia versions prior to
        1.5 this needs to be manually cleared with `Infiltrator.end_session!()`).
  """)
end

function strlimit(str, limit = 30)
  strwidth = 0
  for (i, c) in enumerate(str)
    strwidth += textwidth(c)
    strwidth >= limit && return string(first(str, i - 1), "…")
  end
  return str
end

function show_locals(io, locals, selected::AbstractSet = Set())
  for (var, val) in locals
    if isempty(selected) || var in selected
      one_line_show(io, var, val)
    end
  end
  println(io)
end

get_pmacro_args(mname, line) =
  Set(Symbol.(filter!(x -> x != mname, strip.(filter!(!isempty, split(line, ' '))))))

function show_locals(io, locals, sline::AbstractString)
  selected = get_pmacro_args("@locals", sline)
  show_locals(io, locals, selected)
end

function exfiltrate_locals(io, evalmod, locals, sline::AbstractString)
  selected = get_pmacro_args("@exfiltrate", sline)
  locals_only = isempty(selected)

  iter = locals_only ? keys(locals) : selected
  n = length(iter)

  println(io, "Exfiltrating $(n) local variable$(n == 1 ? "" : "s") into the safehouse.\n")

  for k in iter
    if locals_only
      v = locals[k]
    else
      string(k) in ("@__MODULE__", "#@__MODULE__", "anonymous") && continue

      if !isdefined(evalmod, k)
        printstyled(io, "Warning: "; bold = true, color = :red)
        println(io, "Binding `$k` not defined in this context. Refusing to exfiltrate.")
        continue
      end

      v = getfield(evalmod, k)
    end
    try
      Core.eval(getfield(store, :store), Expr(:(=), k, QuoteNode(v)))
    catch err
      println(io, "Assignment to store variable failed.")
      Base.display_error(io, err, catch_backtrace())
    end
  end
end

function one_line_show(io::IO, var, val)
  width = max(displaysize(io)[2], 40)
  prefix = string("- ", strlimit(string(var), width ÷ 3), "::", strlimit(string(typeof(val)), width ÷ 3), " = ")
  println(io, prefix, strlimit(repr(val), max(width - textwidth(prefix), 10)))
end

function is_complete(s)
  ast = Base.parse_input_line(String(take!(copy(LineEdit.buffer(s)))), depwarn=false)
  return !(isa(ast, Expr) && ast.head === :incomplete)
end

function init_transient_eval_module(mod, locals)
  modns = get_module_names(mod)

  newmod = Module()
  # assign source code module name
  Core.eval(newmod, Expr(:(=), Symbol(mod), mod))
  # spoof @__MODULE__
  Core.eval(newmod, Expr(:macro,
    Expr(:call, Symbol("__MODULE__")),
    Expr(:block, mod))
  )
  # insert local variables into current scope
  Core.eval(newmod, Expr(:block, map(x->Expr(:(=), x...), [(k, maybe_quote(v)) for (k, v) in locals])...))
  # insert variables in safehouse
  Core.eval(newmod, Expr(:block, map(x->Expr(:(=), x...), [(k, maybe_quote(v)) for (k, v) in get_store_names() if !isdefined(newmod, k)])...))
  # insert all bindings from the source module that aren't already defined in the eval module
  Core.eval(newmod, Expr(:block, map(x->Expr(:(=), x...), [(k, maybe_quote(v)) for (k, v) in modns if !isdefined(newmod, k)])...))

  return newmod
end

const PROMPT = Ref{Any}()

function debugprompt(mod, locals, trace, terminal, repl, nostack = false; file, fileline)
  io = Base.pipe_writer(terminal)

  evalmod = init_transient_eval_module(mod, locals)

  try
    if isassigned(PROMPT)
      panel = PROMPT[]
      panel.complete = InfiltratorCompletionProvider(mod, evalmod)
    else
      panel = PROMPT[] = REPL.LineEdit.Prompt("infil> ";
                prompt_prefix = prompt_color,
                prompt_suffix = Base.text_colors[:normal],
                complete = InfiltratorCompletionProvider(mod, evalmod),
                on_enter = is_complete)
      panel.hist = REPL.REPLHistoryProvider(Dict{Symbol,Any}(:Infiltrator => panel))
      REPL.history_reset_state(panel.hist)
    end
    search_prompt, skeymap = LineEdit.setup_search_keymap(panel.hist)
    search_prompt.complete = REPL.LatexCompletions()

    panel.on_done = (s, buf, ok) -> begin
      if !ok
        LineEdit.transition(s, :abort)
        REPL.LineEdit.reset_state(s)
        return false
      end

      line = String(take!(buf))

      if isempty(line)
        println(io)
        LineEdit.reset_state(s)
        return true
      end

      sline = strip(line)
      if sline == "?"
        show_help(io)
        LineEdit.reset_state(s)
        return true
      elseif sline == "@trace"
        print_verbose_stacktrace(io, trace, nostack ? 0 : 100)
        LineEdit.reset_state(s)
        return true
      elseif startswith(sline, "@locals")
        show_locals(io, locals, sline)
        LineEdit.reset_state(s)
        return true
      elseif startswith(sline, "@exfiltrate")
        exfiltrate_locals(io, evalmod, locals, sline)
        LineEdit.reset_state(s)
        return true
      elseif sline == "@toggle"
        spot = (file, fileline)
        ds = getfield(store, :disabled)
        if spot in ds
          delete!(ds, spot)
          println(io, "Enabled infiltration at this infiltration point.\n")
        else
          push!(ds, spot)
          println(io, "Disabled infiltration at this infiltration point.\n")
        end
        LineEdit.reset_state(s)
        return true
      elseif sline == "@exit"
        setfield!(store, :exiting, true)
        if !REPL_HOOKED[]
          printstyled(
            io,
            ">> Restore infiltration capabilities with `Infiltrator.end_session!()` or you will not be able to enter a new session!\n";
            color=:red,
            bold=true
          )
        end
        LineEdit.transition(s, :abort)
        LineEdit.reset_state(s)
        return true
      elseif sline == "@continue"
        LineEdit.transition(s, :abort)
        LineEdit.reset_state(s)
        return true
      end
      ok = true
      result = nothing

      expr = try
        Base.parse_input_line(line)
      catch err
        if get(io, :color, false)
          println(io, "Parsing error.")
        else
          printstyled(io, "Parsing error.", color = Base.error_color())
        end

        LineEdit.reset_state(s)
        return false
      end

      @static if VERSION >= v"1.2.0-DEV.253"
        try
          result = interpret(expr, evalmod)
        catch err
          ok = false
          result = if VERSION >= v"1.7"
            Base.current_exceptions(current_task(); backtrace = true)
          else
            Base.catch_stack()
          end
          if nostack
            result = map(result) do (err, bt)
              if VERSION >= v"1.8-"
                return (exception=err, backtrace=crop_backtrace(bt))
              else
                return err, []
              end
            end
          else
            result = map(result) do (err, bt)
              if VERSION >= v"1.8-"
                return (exception=err, backtrace=crop_backtrace(bt))
              else
                return err, crop_backtrace(bt)
              end
            end
          end
          if isdefined(Base, :ExceptionStack)
            result = Base.ExceptionStack(result)
          end
        end
        if !ok || !REPL.ends_with_semicolon(line)
          REPL.print_response(repl, (result, !ok), true, true)
        end
      else
        try
          result = interpret(expr, evalmod)
        catch err
          ok = false
          result = (err, nostack ? Any[] : crop_backtrace(catch_backtrace()))
        end
        if !ok || !REPL.ends_with_semicolon(line)
          REPL.print_response(repl, ok ? result : result[1], ok ? nothing : result[2], true, true)
        end
      end
      println(io)
      LineEdit.reset_state(s)

      return true
    end

    panel.keymap_dict = LineEdit.keymap(Dict{Any,Any}[skeymap, LineEdit.history_keymap, LineEdit.default_keymap, LineEdit.escape_defaults])

    REPL.run_interface(terminal, REPL.LineEdit.ModalInterface([panel, search_prompt]))
  catch e
    e isa InterruptException || rethrow(e)
  end
end

get_store_names(s = store) = get_module_names(getfield(s, :store), m -> names(m, all = true))

function get_module_names(m::Module, get_names = all_names)
  ns = get_names(m)
  out = Dict()
  for n in ns
    if isdefined(m, n) && !Base.isdeprecated(m, n) && n !== :eval && n !== :include && n !== :anonymous
      out[n] = getfield(m, n)
    end
  end
  out
end

function print_verbose_stacktrace(io, st, limit = 100)
  len = length(st)
  for (i, sf) in enumerate(st)
    if i > limit
      limit > 0 && println(io, "\n<< omitted $(len - limit) stack frames >>")
      break
    end

    print_verbose_stackframe(io, sf, i, len)
  end
  println(io)
end

function print_verbose_stackframe(io, sf::StackTraces.StackFrame, i, len)
  num_padding = ceil(Int, log10(len))
  padding = num_padding + 3
  print(io, "[", lpad(i, num_padding), "] ")
  print_verbose_stackframe(io, sf, padding)
end

function print_verbose_stackframe(io, sf::StackTraces.StackFrame, padding = 2)
  Base.StackTraces.show_spec_linfo(IOContext(io, :backtrace=>true), sf)
  println(io)

  sf.func == StackTraces.top_level_scope_sym && return

  print(io, ' '^padding)
  printstyled(io, "at ", color=:light_black)

  file, line = string(sf.file), sf.line
  file = fixup_stdlib_path(file)
  file = something(Base.find_source_file(file), file)
  printstyled(io, normpath(file), ":", line, '\n', color=:light_black)
end

# https://github.com/JuliaLang/julia/blob/1fb28ad8768cfdc077e968df7adf5716ae8eb9ab/base/methodshow.jl#L134-L148
function fixup_stdlib_path(path::String)
  BUILD_STDLIB_PATH = isdefined(Sys, :BUILD_STDLIB_PATH) ?
    Sys.BUILD_STDLIB_PATH :
    dirname(abspath(joinpath(String((@which uuid1()).file), "..", "..", "..")))

  STDLIB = Sys.STDLIB::String
  if BUILD_STDLIB_PATH != STDLIB
      # BUILD_STDLIB_PATH gets defined in sysinfo.jl
      npath = normpath(path)
      npath′ = replace(npath, normpath(BUILD_STDLIB_PATH) => normpath(STDLIB))
      return npath == npath′ ? path : npath′
  end
  return path
end

function all_names(m)
  symbols = Set(Symbol[])
  seen = Set(Module[])

  for mod in values(Base.loaded_modules)
    all_names(mod, (x, v) -> isdefined(m, x) && !Base.isdeprecated(m, x) && getfield(m, x) === v, symbols, seen)
  end

  return symbols
end
function all_names(m, pred, symbols = Set(Symbol[]), seen = Set(Module[]))
    push!(seen, m)
    ns = names(m; all = true, imported = true)

    for n in ns
        isdefined(m, n) || continue
        Base.isdeprecated(m, n) && continue
        val = getfield(m, n)
        val isa Core.IntrinsicFunction && continue
        if val isa Module && !(val in seen)
            all_names(val, pred, symbols, seen)
        end
        if pred(n, val)
            push!(symbols, n)
        end
    end
    return symbols
end

maybe_quote(x) = (isa(x, Expr) || isa(x, Symbol)) ? QuoteNode(x) : x

function interpret(expr, evalmod)
  out = Core.eval(evalmod, :(ans = $(expr)))

  return out
end

function ast_transformer(sym)
  return function (ex)
    return quote
      let $(sym) = $(ex)
        $(@__MODULE__).end_session!()
        $(sym)
      end
    end
  end
end

# backtraces

function find_first_topelevel_scope(bt::Vector{<:Union{Base.InterpreterIP,Ptr{Cvoid}}})
  for (i, ip) in enumerate(bt)
      st = Base.StackTraces.lookup(ip)
      ind = findfirst(st) do frame
          linfo = frame.linfo
          if linfo isa Core.CodeInfo
              linetable = linfo.linetable
              if isa(linetable, Vector) && length(linetable) ≥ 1
                  lin = first(linetable)
                  if isa(lin, Core.LineInfoNode) && lin.method === StackTraces.top_level_scope_sym
                      return true
                  end
              end
          else
              return frame.func === StackTraces.top_level_scope_sym
          end
      end
      ind === nothing || return i
  end
  return
end

function crop_backtrace(bt)
  i = find_first_topelevel_scope(bt)
  return bt[1:(i === nothing ? end : i)]
end

# completions

mutable struct InfiltratorCompletionProvider <: REPL.CompletionProvider
  mod::Module
  localmod::Module
end

function LineEdit.complete_line(c::InfiltratorCompletionProvider, s)
  partial = REPL.beforecursor(s.input_buffer)
  full = LineEdit.input_string(s)

  ret, range, should_complete = completions(c, full, partial)
  return ret, partial[range], should_complete
end

function completions(c::InfiltratorCompletionProvider, full, partial)
  # repl backend completions
  comps, range, should_complete = REPL.REPLCompletions.completions(full, lastindex(partial), c.mod)
  ret = map(REPL.REPLCompletions.completion_text, comps)

  # completions for local variables
  comps, range, should_complete = REPL.REPLCompletions.completions(full, lastindex(partial), c.localmod)
  prepend!(ret, map(REPL.REPLCompletions.completion_text, comps))

  # completions for safehouse variables
  comps, range, should_complete = REPL.REPLCompletions.completions(full, lastindex(partial), getfield(store, :store))
  prepend!(ret, map(REPL.REPLCompletions.completion_text, comps))

  # Infiltrator commands completions
  commands = ["?", "@trace", "@locals", "@toggle", "@exit", "@continue", "@exfiltrate"]
  prepend!(ret, filter!(c -> startswith(c, partial), commands))

  unique!(ret), range, should_complete
end

end # module


text:
module Classes

using DataStructures
using MacroTools
using InteractiveUtils: subtypes

export @class, @method, Class, AbstractClass, isclass, classof, superclass, superclasses, issubclass, subclasses, absclass

abstract type AbstractClass end            # supertype of all shadow class types
abstract type Class <: AbstractClass end   # superclass of all concrete classes

abs_symbol(cls::Symbol) = Symbol("Abstract", cls)

# Return info about a class in a named tuple
function _class_info(::Type{T}) where {T <: AbstractClass}
    ivars = (isabstracttype(T) ? Expr[] : [:($vname::$vtype) for (vname, vtype) in zip(fieldnames(T), T.types)])
    return (modname=T.name.module, mutable=T.mutable, parameters=T.parameters, ivars=ivars, super=superclass(T))
end

"""
    superclass(t::Type{Class})

Returns the type of the concrete superclass of the given class, or `nothing`
for `Class`, which is the root of the class hierarchy.
"""
superclass(::Type{Class}) = nothing

"""
    superclasses(::Type{T}) where {T <: AbstractClass}

Returns a vector of superclasses from the superclass of the current class
to `Class`, in order.
"""
superclasses(::Type{Class}) = []

function superclasses(::Type{T}) where {T <: AbstractClass}
    super = superclass(T)
    [super, superclasses(super)...]
end

"""
    isclass(X)

Return `true` if `X` is a concrete subclass of `AbstractClass`, or is `Class`, which is abstract.
"""
isclass(any) = false
isclass(::Type{T}) where {T <: AbstractClass} = (T === Class || isconcretetype(T))

"""
    issubclass(t1::DataType, t2::DataType)

Returns `true` if `t1` is a subclass of `t2`, else false.
"""
# identity
issubclass(::Type{T}, ::Type{T}) where {T <: AbstractClass} = true

issubclass(::Type{T1}, ::Type{T2}) where {T1 <: AbstractClass, T2 <: AbstractClass} = T1 in Set(superclasses(T2))

"""
    classof(::Type{T}) where {T <: AbstractClass}

Compute the concrete class associated with abstract class `T`, which must
be a subclass of `AbstractClass`.
"""
function classof(::Type{T}) where {T <: AbstractClass}
    if isclass(T)
        return T
    end

    # Abstract types should have only one concrete subtype
    concrete = filter(isconcretetype, subtypes(T))

    if length(concrete) == 1
        return concrete[1]
    end

    # Should never happen unless user manually creates errant subtypes
    error("Abstract class supertype $T has multiple concrete subtypes: $concrete")
end

"""
    subclasses(::Type{T}) where {T <: AbstractClass}

Compute the vector of subclasses for a given class.
"""
function subclasses(::Type{T}) where {T <: AbstractClass}
    # immediate supertype is "our" entry in the type hierarchy
    super = supertype(T)

    # collect immediate subclasses
    subs = [classof(t) for t in subtypes(super) if isabstracttype(t)]

    # recurse on subclasses
    return [subs; [subclasses(t) for t in subs]...]
end

"""
    absclass(::Type{T}) where {T <: AbstractClass}

Returns the abstract type associated with the concrete class `T`.
"""
function absclass(::Type{T}) where {T <: AbstractClass}
    isclass(T) ? supertype(T) : error("absclass(T) must be called on concrete classes; $T is abstract.")
end

function _argnames(fields)
    return [sym for (sym, arg_type, slurp, default) in map(splitarg, fields)]
end

# We generate two initializer functions: one takes all fields, cumulative through superclasses,
# and another initializes only locally-defined fields. This function produces either, depending
# on the fields passed by _constructors().
function _initializer(class, fields, wheres)
    args = _argnames(fields)
    assigns = [:(_self.$arg = $arg) for arg in args]

    funcdef = :(
        function $class(_self::T, $(fields...)) where {T <: $(abs_symbol(class)), $(wheres...)}
            $(assigns...)
            _self
        end
    )

    return funcdef
end

function _constructors(clsname, super, local_fields, all_fields, wheres)
    args = _argnames(all_fields)
    params = [clause.args[1] for clause in wheres]  # extract parameter names from where clauses

    dflt = length(params) > 0 ? :(
        function $clsname{$(params...)}($(all_fields...)) where {$(wheres...)}
            new{$(params...)}($(args...))
        end) : :(
        function $clsname($(all_fields...))
            new($(args...))
        end)

    init_all = _initializer(clsname, all_fields, wheres)
    methods = [dflt, init_all]

    # If clsname is a direct subclasses of Classes.Class, it has no fields
    # other than those defined locally, so the two methods would be identical.
    # In this case, we emit only one of them.
    if all_fields != local_fields
        init_local = _initializer(clsname, local_fields, wheres)
        push!(methods, init_local)
    end

    # Primarily for immutable classes, we emit a constructor that takes an instance
    # of the direct superclass and copies values when creating a new object.
    super_info = _class_info(super)
    super_fields = super_info.ivars
    if length(super_fields) != 0
        super_args = [:(_super.$arg) for arg in _argnames(super_fields)]
        local_args = _argnames(local_fields)
        all_args = [super_args; local_args]

        immut_init = length(params) > 0 ? :(
            function $clsname{$(params...)}(_super::$super, $(local_fields...)) where {$(wheres...)}
                new{$(params...)}($(all_args...))
            end) : :(
            function $clsname(_super::$super, $(local_fields...))
                new($(all_args...))
            end)
        push!(methods, immut_init)
    end

    return methods
end

function _defclass(clsname, supercls, mutable, wheres, exprs)
    wheres   = (wheres === nothing ? [] : wheres)

    # partition expressions into constructors and field defs
    ctors  = Vector{Expr}()
    fields = Vector{Expr}()
    for ex in exprs
        try
            splitdef(ex)        # throws AssertionError if not a func def
            push!(ctors, ex)
        catch
            push!(fields, ex)
        end
    end

    superinfo = _class_info(supercls)
    all_fields = copy(superinfo.ivars)
    append!(all_fields, fields)

    # add default constructors
    append!(ctors, _constructors(clsname, supercls, fields, all_fields, wheres))

    abs_class = abs_symbol(clsname)
    abs_super = absclass(supercls)

    struct_def = :(
        struct $clsname{$(wheres...)} <: $abs_class
            $(all_fields...)
            $(ctors...)
        end
    )

    # set mutability flag
    struct_def.args[1] = mutable

    result = quote
        abstract type $abs_class <: $abs_super end
        $struct_def

        Classes.superclass(::Type{$clsname}) = $supercls
        $clsname    # return the struct type
    end

    return result
end

macro class(elements...)
    if (mutable = (elements[1] == :mutable))
        elements = elements[2:end]
    end

    if (len = length(elements)) == 1                       # no fields defined
        name_expr = elements[1]
        definition = quote end
    elseif len == 2
        (name_expr, definition) = elements
    else
        error("Unrecognized form for @class definition: $elements")
    end

    # initialize the "captured" vars to avoid "unknown var" warnings
    cls = clsname = exprs = wheres = nothing

    @capture(definition, begin exprs__ end)

    # allow for optional type params and supertype
    if ! (@capture(name_expr, ((cls_{wheres__} | cls_) <: supername_) | (cls_{wheres__} | cls_)) && cls isa Symbol)
        error("Unrecognized class name expression: `$name_expr`")
    end

    # The explicit eval forces supername to be eval'd in calling environment
    supername = (supername === nothing ? :Class : supername)

    # __module__ is a "hidden" arg passed to macros with the caller's Module
    expr = _defclass(cls, __module__.eval(supername), mutable, wheres, exprs)
    return esc(expr)
end

"""
    @method(funcdef)

Translates a function whose first argument is a concrete subclass of Class
the same function but with type of the first argument changed to the abstract
supertype of the class, thereby allowing it to be called on subclasses as well.
"""
macro method(funcdef)
    parts = splitdef(funcdef)
    name = parts[:name]
    args = parts[:args]
    whereparams = parts[:whereparams]

    if ! @capture(args[1], arg1_::T_)
        error("First argument of method $name must be explicitly typed")
    end

    type_symbol = gensym("T")  # gensym avoids conflict with user's type params
    abs_super = abs_symbol(T)

    # Redefine the function to accept any first arg that's a subclass of abstype
    parts[:whereparams] = (:($type_symbol <: $abs_super), whereparams...)
    args[1] = :($arg1::$type_symbol)
    expr = MacroTools.combinedef(parts)
    return esc(expr)
end

end # module

ext: jl


text:
""" Tests the functionalities of BinaryNode. """
function test_binarynode()
    @info "Testing BinaryNode..."
    bn = BinaryNode(1)
    @test_throws ErrorException get_parent(bn)
    leftchild(bn, BinaryNode(2))
    rightchild(bn, BinaryNode(3))
    rightchild(bn.right, BinaryNode(4))
    @test all([child.idx for child in children(bn)] .== [2,3]) &&
        isempty(children(bn.right.right)) && length(alloffspring(bn)) == 3

    delete_children!(bn.right)
    @test_throws ErrorException get_lower_child(bn.right)
    @test_throws ErrorException get_upper_child(bn.right)
    @test isnothing(bn.parent) && isnothing(bn.right.right) &&
        isempty(children(bn.right)) && length(alloffspring(bn)) == 2

    @test_throws ErrorException set_split_values!(bn.right, [1,2,3], 4)
    @test_throws ErrorException set_classification_label!(bn, 5)
end

function test_data_processing()
    @info "Testing data processing..."
    Y = transpose(MLDatasets.BostonHousing.targets())
    X = transpose(MLDatasets.BostonHousing.features())

    @test_throws ErrorException split_data(X, Y, bins = 3, sample_proportion = [0.1, 0.2, 0.7])
    @test_throws ErrorException split_data(X, Y, sample_proportion = [0.1, 0.2, 0.3])
    @test_throws ErrorException split_data(X, Y, sample_count = [10, 20, 30])

    data = split_data(X, Y)
    @test length(data) == 2 && length(data[1][2]) + length(data[2][2]) == length(Y)
    data = split_data(X, Y, bins = 9)
    @test length(data) == 9 && sum(length(dat[2]) for dat in data) == length(Y) && all(isapprox(length(dat[2]), length(Y)/9, atol = 1) for dat in data)
    data = split_data(X, Y, sample_count = [100, 200, 206])
    @test length(data) == 3 && sum(length(dat[2]) for dat in data) == length(Y)
end

""" Tests full MIO solution functionalities of MIOTree. """
function test_miotree()
    @info "Testing MIOTree..."
    d = MIOTree_defaults()
    d = MIOTree_defaults(max_depth =  4, cp = 1e-5)
    @test d[:max_depth] == 4
    @test get_param(d, :cp) == 1e-5
    mt = MIOTree(SOLVER_SILENT; max_depth = 4, minbucket = 0.03)
    df = load_irisdata()
    df = df[1:Int(floor(size(df, 1)/2)),:]
    X = Matrix(df[:,1:4])
    Y =  Array(df[:, "class"])
    md = 2
    set_params!(mt, max_depth = md, minbucket = 0.001)
    generate_binary_tree(mt)
    generate_MIO_model(mt, X, Y)
    @test length(allleaves(mt)) == 2^md

    set_optimizer(mt, SOLVER_SILENT)
    optimize!(mt)
    @test !check_if_trained(mt) # Trees must be populated and pruned before they qualify!

    m = mt.model
    as, d, Lt, ckt, Nkt = [getvalue.(m[i]) for i in [:a, :d, :Lt, :ckt, :Nkt]];
    populate_nodes!(mt)
    @test length(allleaves(mt)) == 2^md
    @test sum(!isnothing(node.a) for node in allnodes(mt)) == 2^md-1 - sum(all(isapprox.(as[i,:], 0, atol = 1e-10)) for i = 1:2^md-1)
    @test OT.complexity(mt) == sum(as .!= 0)
    prune!(mt)
    @test check_if_trained(mt)
    @test length(allleaves(mt)) == sum(isapprox.(ckt, 1, atol = 1e-4)) == count(!isnothing(node.label) for node in allnodes(mt))
    @test all(getproperty.(apply(mt, X), :label) .== predict(mt, X))

    # Testing cloning
    new_mt = clone(mt)
    @test all(get_split_values(mt.root) .== (get_split_weights(new_mt.root), get_split_threshold(new_mt.root)))

    # # Plotting results for debugging
    # using Plots
    # colors = ["green", "red", "blue"]
    # labels = unique(Y)
    # plt = plot()
    # for i = 1:length(labels)
    #     label = labels[i]
    #     idxs = findall(x -> x == label, Y)
    #     plt = scatter!(X[idxs,1], X[idxs, 2], color = colors[i], label = labels[i])
    # end
    # display(plt)

    # # Plotting correct predictions
    # colors = ["purple", "yellow"]
    # labels = [0, 1]
    # plt = plot()
    # for i=1:length(labels)
    #     label = labels[i]
    #     idxs = findall(x -> x == label, predict(mt, X) .== Y)
    #     plt = scatter!(X[idxs,1], X[idxs, 2], color = colors[i], label = labels[i])
    # end
    # display(plt)

    # Checking data extraction
    ud, ld = trust_region_data(mt)
    @test all(sum(length(ud[lf.idx]) + length(ld[lf.idx])) == depth(lf) for lf in allleaves(mt)) # The right number of splits
end

function test_hyperplanecart()
    @info "Testing hyperplane CART..."
    mt = MIOTree(SOLVER_SILENT, max_depth = 5)
    df = binarize(load_irisdata())
    X = Matrix(df[:,1:4])
    Y =  Array(df[:, "class"])
    fit!(mt, "cart", X, Y)
    @test check_if_trained(mt)
    @test score(mt, X, Y) == 1.
    nds = apply(mt, X)
    @test all(getproperty.(nds, :label) .== Y)
    @test all(isnothing.(getproperty.(allleaves(mt), :label)) .== false)
    leaves = allleaves(mt)
    not_leaves = [nd for nd in allnodes(mt) if !is_leaf(nd)]
    @test all(is_leaf.(not_leaves) .== false)
    @test all(isnothing.(getproperty.(not_leaves, :a)) .== false)
    @test all(isnothing.(getproperty.(leaves, :label)) .== false)

    # Check that pruning does nothing
    prune!(mt)
    @test check_if_trained(mt)

    # Checking data extraction
    ud, ld = trust_region_data(mt)
    @test all(sum(length(ud[lf.idx]) + length(ld[lf.idx])) == depth(lf) for lf in allleaves(mt)) # The right number of splits

    # Checking deepening trees, and warmstarting
    deepen_to_max_depth!(mt)
    clean_model!(mt)
    generate_MIO_model(mt, X, Y)
    warmstart(mt)
    optimize!(mt)
    populate_nodes!(mt)
    prune!(mt)

    @test isapprox(score(mt, X, Y), 1, atol = 0.05)
    leaves = allleaves(mt)
    not_leaves = [nd for nd in allnodes(mt) if !is_leaf(nd)]
    @test all(is_leaf.(not_leaves) .== false)
    @test all(isnothing.(getproperty.(not_leaves, :a)) .== false)
    # @test all(isnothing.(getproperty.(leaves, :label)) .== false) # TODO: fix this test.
end

function test_sequential()
    @info "Testing sequential training..."
    mt = MIOTree(SOLVER_SILENT; max_depth = 2, minbucket = 0.03)
    df = load_irisdata()
    X = Matrix(df[:,1:4])
    Y =  Array(df[:, "class"])
    sequential_train!(mt, X, Y, pruning = true)
    @test true
end

function test_regression()
    @info "Testing regression..."
    feature_names = MLDatasets.BostonHousing.feature_names()
    n_samples = 20
    X_all = Matrix(transpose(MLDatasets.BostonHousing.features()))
    Y_all = Array(transpose(MLDatasets.BostonHousing.targets()))
    X = X_all[1:n_samples, :]
    Y = Y_all[1:n_samples, :]
    mt = MIOTree(SOLVER_SILENT, max_depth = 2, regression = true)
    fit!(mt, "mio", X, Y)
    @test all(label isa Tuple for label in get_classification_label.(allleaves(mt)))
    @test all(label isa Real for label in get_regression_constant.(allleaves(mt)))
    @test all(label isa AbstractArray for label in get_regression_weights.(allleaves(mt)))

    @test check_if_trained(mt)
    score1 = score(mt, X_all, Y_all)

    # Upping number of samples, and warmstart procedure
    n_samples = 30
    X = Matrix(transpose(MLDatasets.BostonHousing.features()))[1:n_samples, :]
    Y = Array(transpose(MLDatasets.BostonHousing.targets()))[1:n_samples, :]
    clean_model!(mt)
    generate_MIO_model(mt, X, Y)
    warmstart(mt)
    optimize!(mt)
    populate_nodes!(mt)
    prune!(mt)
    @test check_if_trained(mt)
    score2 = score(mt, X_all, Y_all)
    @test score1 <= 1 && score2 <= 1
end

function test_ensemblereg()
    @info "Testing ensemble regression..."
    feature_names = MLDatasets.BostonHousing.feature_names()
    Y = Array(transpose(MLDatasets.BostonHousing.targets()))
    # Shuffled data
    shuffle_idxs = shuffle(1:Int(length(Y)/2))
    Y = Y[shuffle_idxs]
    X = Matrix(transpose(MLDatasets.BostonHousing.features()))[shuffle_idxs, :]

    te = TreeEnsemble(SOLVER_SILENT; regression = true, max_depth = 1)
    plant_trees(te, 10)
    fit!(te, "mio", X, Y)
    @test all(check_if_trained.(te.trees))
    weigh_trees(te, X, Y)
    @test isapprox(sum(te.weights), 1, atol = 1e-5)
    @test score(te, X, Y) >= 0.5
end

function test_ensemblecls()
    @info "Testing ensemble classification... "
    feature_names = MLDatasets.BostonHousing.feature_names()
    Y = Array(transpose(MLDatasets.BostonHousing.targets()))
    shuffle_idxs = shuffle(1:Int(length(Y)))
    Y = Array(Y[shuffle_idxs] .>= 20)
    X = Matrix(transpose(MLDatasets.BostonHousing.features()))[shuffle_idxs, :]
    te = TreeEnsemble(SOLVER_SILENT; max_depth = 3)
    plant_trees(te, 11)
    fit!(te, "cart", X, Y)
    @test all(check_if_trained.(te.trees))
    @test score(te, X, Y) >= 0.82
end

function test_cluster_heuristic()
    @info "Testing clustering heuristic in classification..."
    Y = Array(transpose(MLDatasets.BostonHousing.targets()))
    shuffle_idxs = shuffle(1:Int(length(Y)))
    Y = Array(Y[shuffle_idxs])
    X = Matrix(transpose(MLDatasets.BostonHousing.features()))[shuffle_idxs, :]

    X_norm, X_bounds = normalize(X)
    Y_norm, Y_bounds = normalize(Y)
    dists = pairwise_distances(X_norm)
    clust = hclust(dists)
    max_depth = 5
    n_clust = 2^(max_depth-1)
    idxs = cutree(clust; k = n_clust, h = max_depth)
    cluster_bins = Dict(i => [] for i in unique(idxs))
    [push!(cluster_bins[idxs[i]], i) for i = 1:length(idxs)]
    subset_idxs = []
    sample_proportion = 0.125
    for (key, val) in cluster_bins # Picking just one-eighth of all samples in clusters
        append!(subset_idxs, val[1:Int(ceil(sample_proportion*length(val)))])
    end
    @test sample_proportion <= length(subset_idxs)/length(Y) <= 2*sample_proportion

    mt = MIOTree(SOLVER_SILENT, max_depth = max_depth)
    generate_binary_tree(mt)
    generate_MIO_model(mt, X_norm[subset_idxs, :], Array(Y_norm .>= 0.3)[subset_idxs])
    optimize!(mt)
    populate_nodes!(mt)
    prune!(mt)
    @test score(mt, X_norm, Array(Y_norm .>= 0.3)) >= 0.8
end

test_binarynode()

test_data_processing()

test_miotree()

test_hyperplanecart()

test_sequential()

test_regression()

test_ensemblereg()

test_ensemblecls()

test_cluster_heuristic()


text:
@doc read(normpath(dirname(@__DIR__), "README.md"), String)
module JET

const CC = Core.Compiler

@static if !isdefined(CC, :AbstractInterpreter)
    throw(ErrorException("JET.jl only works with Julia versions 1.6 and higher"))
elseif :backedges ∉ fieldnames(Core.MethodInstance)
    @warn """with your Julia version, JET.jl may not be able to update analysis result
    correctly after refinement of a method in deeper call sites
    """
end

# imports
# =======

# `JETInterpreter`
import .CC:
    # abstractinterpreterinterface.jl
    InferenceParams,
    OptimizationParams,
    get_world_counter,
    get_inference_cache,
    lock_mi_inference,
    unlock_mi_inference,
    add_remark!,
    may_optimize,
    may_compress,
    may_discard_trees,
    # jetcache.jl
    code_cache,
    cache_result!,
    # tfuncs.jl
    builtin_tfunction,
    return_type_tfunc,
    # abstractinterpretation.jl
    abstract_call_gf_by_type,
    abstract_call_method_with_const_args,
    abstract_eval_special_value,
    abstract_eval_value,
    abstract_eval_statement,
    # typeinfer.jl
    typeinf,
    _typeinf,
    typeinf_edge

# `ConcreteInterpreter`
import JuliaInterpreter:
    # virtualprocess.jl
    step_expr!,
    evaluate_call_recurse!,
    handle_err

# usings
# ======
# TODO: really use `using` instead

import Core:
    Builtin,
    MethodMatch,
    LineInfoNode,
    SimpleVector,
    svec

import .CC:
    AbstractInterpreter,
    NativeInterpreter,
    InferenceState,
    InferenceResult,
    CodeInfo,
    InternalCodeCache,
    CodeInstance,
    WorldRange,
    WorldView,
    MethodInstance,
    Bottom,
    NOT_FOUND,
    MethodMatchInfo,
    UnionSplitInfo,
    MethodLookupResult,
    Const,
    VarTable,
    SSAValue,
    SlotNumber,
    Slot,
    slot_id,
    GlobalRef,
    GotoIfNot,
    ReturnNode,
    widenconst,
    ⊑,
    CallMeta,
    is_throw_call,
    tmerge,
    argtypes_to_type,
    _methods_by_ftype,
    specialize_method,
    add_backedge!,
    compute_basic_blocks

import Base:
    parse_input_line,
    unwrap_unionall,
    rewrap_unionall,
    Fix1,
    Fix2,
    IdSet

import Base.Meta:
    isexpr,
    _parse_string,
    lower

using LoweredCodeUtils, JuliaInterpreter

import LoweredCodeUtils:
    istypedef,
    ismethod,
    callee_matches

import JuliaInterpreter:
    bypass_builtins,
    maybe_evaluate_builtin,
    collect_args,
    is_return,
    is_quotenode_egal

using FileWatching, Requires

using InteractiveUtils

# common
# ======

# hooks
# -----

const INIT_HOOKS = Function[]
push_inithook!(f) = push!(INIT_HOOKS, f)
__init__() = foreach(f->f(), INIT_HOOKS)

# macros
# ------

# TODO: upstream these macros into Base

"""
    @invoke f(arg::T, ...; kwargs...)

Provides a convenient way to call [`invoke`](@ref);
`@invoke f(arg1::T1, arg2::T2; kwargs...)` will be expanded into `invoke(f, Tuple{T1,T2}, arg1, arg2; kwargs...)`.
When an argument's type annotation is omitted, it's specified as `Any` argument, e.g.
`@invoke f(arg1::T, arg2)` will be expanded into `invoke(f, Tuple{T,Any}, arg1, arg2)`.

This could be used to call down to `NativeInterpreter`'s abstract interpretation method of
  `f` while passing `JETInterpreter` so that subsequent calls of abstract interpretation
  functions overloaded against `JETInterpreter` can be called from the native method of `f`;
e.g. calls down to `NativeInterpreter`'s `abstract_call_gf_by_type` method:
```julia
@invoke abstract_call_gf_by_type(interp::AbstractInterpreter, f, argtypes::Vector{Any}, atype, sv::InferenceState,
                                 max_methods::Int)
```
"""
macro invoke(ex)
    f, args, kwargs = destructure_callex(ex)
    arg2typs = map(args) do x
        isexpr(x, :(::)) ? (x.args...,) : (x, GlobalRef(Core, :Any))
    end
    args, argtypes = first.(arg2typs), last.(arg2typs)
    return if isempty(kwargs)
        :($(GlobalRef(Core, :invoke))($(f), Tuple{$(argtypes...)}, $(args...))) # might not be necessary
    else
        :($(GlobalRef(Core, :invoke))($(f), Tuple{$(argtypes...)}, $(args...); $(kwargs...)))
    end |> esc
end

function destructure_callex(ex)
    @assert isexpr(ex, :call) "call expression f(args...; kwargs...) should be given"

    f = first(ex.args)
    args = []
    kwargs = []
    for x in ex.args[2:end]
        if isexpr(x, :parameters)
            append!(kwargs, x.args)
        elseif isexpr(x, :kw)
            push!(kwargs, x)
        else
            push!(args, x)
        end
    end

    return f, args, kwargs
end

"""
    @invokelatest f(args...; kwargs...)

Provides a convenient way to call [`Base.invokelatest`](@ref).
`@invokelatest f(args...; kwargs...)` will simply be expanded into
`Base.invokelatest(f, args...; kwargs...)`.
"""
macro invokelatest(ex)
    f, args, kwargs = destructure_callex(ex)
    return if isempty(kwargs) # eliminates dispatch to kwarg methods, might unnecessary to be special cased
        :($(GlobalRef(Base, :invokelatest))($(f), $(args...)))
    else
        :($(GlobalRef(Base, :invokelatest))($(f), $(args...); $(kwargs...)))
    end |> esc
end

# inference frame
# ---------------

is_constant_propagated(frame::InferenceState) = CC.any(frame.result.overridden_by_const)

# # XXX: should sync with the `haveconst` check within `abstract_call_method_with_const_args` ?
# is_constant_propagated(frame::InferenceState) = is_constant_propagated(frame.result)
# function is_constant_propagated(result::InferenceResult)
#     for a in result.argtypes
#         if CC.has_nontrivial_const_info(a) && CC.const_prop_profitable(a)
#             return true
#         end
#     end
#     return false
# end

istoplevel(linfo::MethodInstance) = linfo.def == __toplevel__
istoplevel(sv::InferenceState)    = istoplevel(sv.linfo)

prewalk_inf_frame(@nospecialize(f), ::Nothing) = return
function prewalk_inf_frame(@nospecialize(f), frame::InferenceState)
    ret = f(frame)
    prewalk_inf_frame(f, frame.parent)
    return ret
end

postwalk_inf_frame(@nospecialize(f), ::Nothing) = return
function postwalk_inf_frame(@nospecialize(f), frame::InferenceState)
    postwalk_inf_frame(f, frame.parent)
    return f(frame)
end

# NOTE: these methods assume `frame` is not inlined
get_cur_pc(frame::InferenceState) = return frame.currpc
get_cur_stmt(frame::InferenceState) = frame.src.code[get_cur_pc(frame)]
get_cur_loc(frame::InferenceState) = frame.src.codelocs[get_cur_pc(frame)]
get_cur_linfo(frame::InferenceState) = frame.src.linetable[get_cur_loc(frame)]::LineInfoNode
get_cur_varstates(frame::InferenceState) = frame.stmt_types[get_cur_pc(frame)]
get_result(frame::InferenceState) = frame.result.result

# includes
# ========

include("reports.jl")
include("abstractinterpreterinterface.jl")
include("jetcache.jl")
include("tfuncs.jl")
include("abstractinterpretation.jl")
include("typeinfer.jl")
include("print.jl")
include("virtualprocess.jl")
include("watch.jl")

# entry
# =====

function profile_file(io::IO,
                      filename::AbstractString,
                      mod::Module = Main;
                      kwargs...)
    text = read(filename, String)
    return profile_text(io, text, filename, mod; kwargs...)
end
profile_file(args...; kwargs...) = profile_file(stdout::IO, args...; kwargs...)

function profile_text(io::IO,
                      text::AbstractString,
                      filename::AbstractString = "top-level",
                      mod::Module = Main;
                      profiling_logger::Union{Nothing,IO} = nothing,
                      kwargs...)
    included_files, reports, postprocess = collect_reports(profiling_logger,
                                                           mod,
                                                           text,
                                                           filename;
                                                           kwargs...,
                                                           )
    return included_files, print_reports(io, reports, postprocess; kwargs...)
end
profile_text(args...; kwargs...) = profile_text(stdout::IO, args...; kwargs...)

collect_reports(::Nothing, args...; kwargs...) = collect_reports(args...; kwargs...)
function collect_reports(logger::IO, args...; kwargs...)
    print(logger, "profiling from ", #= filename =# last(args), " ...")
    s = time()

    ret = collect_reports(args...; kwargs...)

    sec = round(time() - s; digits = 3)

    print(logger, '\b'^3)
    println(logger, "(finished in $(sec) sec)")

    return ret
end

function collect_reports(actualmod, text, filename; kwargs...)
    virtualmod = gen_virtual_module(actualmod)

    interp = JETInterpreter(; # dummy
                              inf_params      = gen_inf_params(; kwargs...),
                              opt_params      = gen_opt_params(; kwargs...),
                              analysis_params = AnalysisParams(; kwargs...),
                              kwargs...)
    ret, interp = virtual_process!(text,
                                   filename,
                                   virtualmod,
                                   Symbol(actualmod),
                                   interp,
                                   )

    return ret.included_files,
           # non-empty `ret.toplevel_error_reports` means critical errors happened during
           # the AST transformation, so they always have precedence over `ret.inference_error_reports`
           !isempty(ret.toplevel_error_reports) ? ret.toplevel_error_reports : ret.inference_error_reports,
           gen_postprocess(virtualmod, actualmod)
end

gen_virtual_module(actualmod = Main) =
    return Core.eval(actualmod, :(module $(gensym(:JETVirtualModule)) end))::Module

# fix virtual module printing based on string manipulation; the "actual" modules may not be
# loaded into this process
function gen_postprocess(virtualmod, actualmod)
    virtual = string(virtualmod)
    actual  = string(actualmod)
    return actualmod == Main ?
        Fix2(replace, "Main." => "") ∘ Fix2(replace, virtual => actual) :
        Fix2(replace, virtual => actual)
end

# this dummy module will be used by `istoplevel` to check if the current inference frame is
# created by `profile_toplevel!` or not (i.e. `@generated` function)
module __toplevel__ end

function profile_toplevel!(interp::JETInterpreter, mod::Module, src::CodeInfo)
    # construct toplevel `MethodInstance`
    mi = ccall(:jl_new_method_instance_uninit, Ref{Core.MethodInstance}, ());
    mi.uninferred = src
    mi.specTypes = Tuple{}
    mi.def = mod

    result = InferenceResult(mi);
    frame = InferenceState(result, src, #= cached =# true, interp);

    mi.def = __toplevel__ # set to the dummy module

    return profile_frame!(interp, frame)
end

# TODO `profile_call_builtin!` ?
function profile_gf_by_type!(interp::JETInterpreter,
                             @nospecialize(tt::Type{<:Tuple}),
                             world::UInt = get_world_counter(interp),
                             )
    mms = _methods_by_ftype(tt, InferenceParams(interp).MAX_METHODS, world)
    @assert mms !== false "unable to find matching method for $(tt)"

    filter!(mm::MethodMatch->mm.spec_types===tt, mms)
    @assert length(mms) == 1 "unable to find single target method for $(tt)"

    mm = first(mms)::MethodMatch

    return profile_method_signature!(interp, mm.method, mm.spec_types, mm.sparams)
end

function profile_method!(interp::JETInterpreter,
                         m::Method,
                         world::UInt = get_world_counter(interp),
                         )
    return profile_method_signature!(interp, m, m.sig, sparams_from_method_signature(m), world)
end

function sparams_from_method_signature(m)
    s = TypeVar[]
    sig = m.sig
    while isa(sig, UnionAll)
        push!(s, sig.var)
        sig = sig.body
    end
    return svec(s...)
end

function profile_method_signature!(interp::JETInterpreter,
                                   m::Method,
                                   @nospecialize(atype),
                                   sparams::SimpleVector,
                                   world::UInt = get_world_counter(interp),
                                   )
    mi = specialize_method(m, atype, sparams)

    result = InferenceResult(mi)

    frame = InferenceState(result, #= cached =# true, interp)

    return profile_frame!(interp, frame)
end

function profile_frame!(interp::JETInterpreter, frame::InferenceState)
    typeinf(interp, frame)

    # report `throw` calls "appropriately";
    # if the final return type here is `Bottom`-annotated, it _may_ mean the control flow
    # didn't catch some of the `UncaughtExceptionReport`s stashed within `interp.uncaught_exceptions`,
    if get_result(frame) === Bottom
        if !isempty(interp.uncaught_exceptions)
            append!(interp.reports, interp.uncaught_exceptions)
        end
    end

    return interp, frame
end

# test, interactive
# =================

# profiles from call expression
macro profile_call(ex0...)
    return InteractiveUtils.gen_call_with_extracted_types_and_kwargs(__module__, :profile_call, ex0)
end

function profile_call(@nospecialize(f), @nospecialize(types = Tuple{}); kwargs...)
    ft = Core.Typeof(f)
    if isa(types, Type)
        u = unwrap_unionall(types)
        tt = rewrap_unionall(Tuple{ft, u.parameters...}, types)
    else
        tt = Tuple{ft, types...}
    end

    interp = JETInterpreter(; inf_params      = gen_inf_params(; kwargs...),
                              opt_params      = gen_opt_params(; kwargs...),
                              analysis_params = AnalysisParams(; kwargs...),
                              kwargs...)
    return profile_gf_by_type!(interp, tt)
end

# collects and prints reports from call expression
macro report_call(ex0...)
    return InteractiveUtils.gen_call_with_extracted_types_and_kwargs(__module__, :report_call, ex0)
end

function report_call(@nospecialize(f), @nospecialize(types = Tuple{}); kwargs...)
    interp, frame = profile_call(f, types; kwargs...)
    print_reports(interp.reports; kwargs...)
    return get_result(frame)
end

print_reports(args...; kwargs...) = print_reports(stdout::IO, args...; kwargs...)

# for inspection
macro lwr(ex) QuoteNode(lower(__module__, ex)) end
macro src(ex) QuoteNode(first(lower(__module__, ex).args)) end

# exports
# =======

export
    profile_file,
    profile_and_watch_file,
    profile_text,
    @profile_call,
    profile_call,
    @report_call,
    report_call

end


text:
# ---
# title: 436. Find Right Interval
# id: problem436
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Binary Search
# link: <https://leetcode.com/problems/find-right-interval/description/>
# hidden: true
# ---
#
# You are given an array of `intervals`, where `intervals[i] = [starti, endi]`
# and each `starti` is **unique**.
#
# The **r** **ight** **  interval** for an interval `i` is an interval `j` such
# that `startj`` >= endi` and `startj` is  **minimized**.
#
# Return  _an array of   **right interval**  indices for each interval `i`_. If
# no  **right interval**  exists for interval `i`, then put `-1` at index `i`.
#
#
#
# **Example 1:**
#
#
#
#     Input: intervals = [[1,2]]
#     Output: [-1]
#     Explanation: There is only one interval in the collection, so it outputs -1.
#
#
# **Example 2:**
#
#
#
#     Input: intervals = [[3,4],[2,3],[1,2]]
#     Output: [-1,0,1]
#     Explanation: There is no right interval for [3,4].
#     The right interval for [2,3] is [3,4] since start0 = 3 is the smallest start that is >= end1 = 3.
#     The right interval for [1,2] is [2,3] since start1 = 2 is the smallest start that is >= end2 = 2.
#
#
# **Example 3:**
#
#
#
#     Input: intervals = [[1,4],[2,3],[3,4]]
#     Output: [-1,2,-1]
#     Explanation: There is no right interval for [1,4] and [3,4].
#     The right interval for [2,3] is [3,4] since start2 = 3 is the smallest start that is >= end1 = 3.
#
#
#
#
# **Constraints:**
#
#   * `1 <= intervals.length <= 2 * 104`
#   * `intervals[i].length == 2`
#   * `-106 <= starti <= endi <= 106`
#   * The start point of each interval is **unique**.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
import UMAP

function plot_noise_estimation_diagnostics(edge_length::Vector{Float64}, confidences::Vector{Float64}, d1::T, d2::T; title::String="Noise estimation",
        confidence_nn_id::Union{Int, String}="k", linewidth::Float64=2.0, bins::Int=200, kwargs...) where T <: Distributions.UnivariateDistribution
    x_max = quantile(edge_length, 0.99)
    Plots.histogram(edge_length[edge_length .< x_max]; bins=bins, normalize=true, label="Observed",
        xlabel="Distance to $(confidence_nn_id)'th nearest neighbor", ylabel="Density", title=title,
        linewidth=linewidth, legendtitle="Distribution", kwargs...)
    x_vals = range(0, x_max, length=1000)
    n1 = sum(confidences)
    n2 = length(confidences) - n1
    Plots.plot!(x_vals, n1 / (n1 + n2) .* pdf.(d1, x_vals), label="Intracellular", linewidth=linewidth)
    Plots.plot!(x_vals, n2 / (n1 + n2) .* pdf.(d2, x_vals), label="Background", linewidth=linewidth)
end

function plot_num_transcript_overview(genes::Vector{Int}, confidences::Vector{Float64}, gene_names::Vector; size=(800, 300), kwargs...)
    order = sortperm(gene_names)
    plot_expression_vectors(
        count_array(genes[confidences .>= 0.5], max_value=length(gene_names))[order],
        count_array(genes[confidences .< 0.5], max_value=length(gene_names))[order],
        gene_names=gene_names[order]; labels=["Real", "Noise"], xlabel="Gene ID", ylabel="#Transcripts",
        min_expr_frac=0.01, legend=:topleft, alpha=0.3, xlims=(0, length(gene_names)), size=size, kwargs...)
end

function plot_gene_structure(df_spatial::DataFrame, gene_names::Vector, confidence::Vector{Float64}=df_spatial.confidence; kwargs...)
    adjacent_points, adjacent_weights = build_molecule_graph(df_spatial, filter=false)[1:2];
    cor_mat = pairwise_gene_spatial_cor(df_spatial.gene, df_spatial.confidence, adjacent_points, adjacent_weights);
    cor_vals = vec(cor_mat)
    min_cor = quantile(cor_vals[cor_vals .> 0], 0.001) ./ 5;
    max_cor = quantile(cor_vals, 0.99);
    p_dists = 1 .- max.(min.(cor_mat, max_cor), min_cor) ./ max_cor;
    p_dists[diagind(p_dists)] .= 0.0;

    embedding = UMAP.umap(p_dists, 2; metric=:precomputed, spread=1.0, min_dist=0.1, n_epochs=5000, n_neighbors=max(min(15, length(gene_names) ÷ 2), 2));
    marker_sizes = log.(count_array(df_spatial.gene));
    marker_sizes ./= median(marker_sizes) ./ 2;

    Plots.scatter(embedding[1,:], embedding[2,:]; ms=marker_sizes, legend=false, alpha=0.1, size=(800, 800),
        ticks=false, xlabel="UMAP-1", ylabel="UMAP-2", title="Gene local structure", kwargs...)

    Plots.annotate!(collect(zip(embedding[1,:], embedding[2,:], Plots.text.(gene_names, ceil.(Int, marker_sizes .* 1.5)))), color=Colors.RGBA(0.0, 0.0, 0.0, 0.75))
end

function estimate_panel_plot_size(df_spatial::DataFrame, min_molecules_per_cell::Int, min_pixels_per_cell::Int=7)
    n_cells_per_side = sqrt(size(df_spatial, 1) / min_molecules_per_cell)
    plot_size = min_pixels_per_cell * n_cells_per_side
    x_rng = val_range(df_spatial.x)
    y_rng = val_range(df_spatial.y)
    y_ratio = (y_rng[2] - y_rng[1]) / (x_rng[2] - x_rng[1])

    y_ratio = y_ratio^0.5;
    x_ratio = 1 / y_ratio^0.5;
    plot_size = (x_ratio * plot_size, y_ratio * plot_size)
    return plot_size, n_cells_per_side
end

plot_dataset_colors(df_spatial::DataFrame, colors::Symbol; kwargs...) =
    plot_dataset_colors(df_spatial, df_spatial[!, colors]; kwargs...)

function plot_dataset_colors(df_spatial::DataFrame, color::Union{Vector, Symbol, String}; min_molecules_per_cell::Int, min_pixels_per_cell::Int=7, size::Float64=-1., alpha::Union{Float64, Vector{Float64}}=0.25,
        prior_polygons::Array{Matrix{Float64}, 1}=Matrix{Float64}[], polygons::Array{Matrix{Float64}, 1}=Matrix{Float64}[], ticks::Bool=true, axis_kwargs::KWArgT=nothing, kwargs...)

    axis_kwargs = update_args((xticklabelsize=12, yticklabelsize=12), axis_kwargs)
    plot_size = estimate_panel_plot_size(df_spatial, min_molecules_per_cell, min_pixels_per_cell)[1]
    if size < 0
        size = 5 / min_molecules_per_cell
    end

    fig = MK.Figure(resolution=plot_size)
    fig[1, 1] = MK.Axis(fig; xticksvisible=ticks, yticksvisible=ticks, axis_kwargs...);

    if length(prior_polygons) > 0
        MK.poly!([MK.Point2.(eachrow(p)) for p in prior_polygons]; strokecolor="darkred", color=Colors.RGBA(1, 0.65, 0, 0.25), strokewidth=0.5)
    end

    return plot_molecules!(df_spatial, polygons; color=color, point_size=size, alpha=alpha, kwargs...)
end

max_stars_repo_name: YinuoJin/Baysor

max_issues_repo_name: YinuoJin/Baysor

max_forks_repo_name: YinuoJin/Baysor


text:
export UCBSelector

using .Utils: findallmax

"""
    UCBSelector(na; c=2.0, ϵ=1e-10)

# Arguments
- `na` is the number of actions used to create a internal counter.
- `t` is used to store current time step.
- `c` is used to control the degree of exploration.
"""
mutable struct UCBSelector <: AbstractDiscreteActionSelector
    c::Float64
    actioncounts::Vector{Float64}
    step::Int
    UCBSelector(na; c = 2.0, ϵ = 1e-10, step = 1) = new(c, fill(ϵ, na), 1)
end

@doc raw"""
    (ucb::UCBSelector)(values::AbstractArray)

Unlike [`EpsilonGreedySelector`](@ref), uncertaintyies are considered in UCB.

!!! note
    If multiple values with the same maximum value are found.
    Then a random one will be returned!

```math
A_t = \underset{a}{\arg \max} \left[ Q_t(a) + c \sqrt{\frac{\ln t}{N_t(a)}} \right]
```

See more details at Section (2.7) on Page 35 of the book *Sutton, Richard S., and Andrew G. Barto. Reinforcement learning: An introduction. MIT press, 2018.*
""" function (p::UCBSelector)(values::AbstractArray)
    action = findallmax(@. values + p.c * sqrt(log(p.step + 1) / p.actioncounts))[2] |>
             sample
    p.actioncounts[action] += 1
    p.step += 1
    action
end


text:
using DataStructures: SortedDict, OrderedDict

_empty_range(t::Type{<:Signed}) = zero(t):-one(t)
_empty_range(t::Type{<:Unsigned}) = one(t):zero(t)

"""
    AbstractBags{T}

Supertype for structures storing indices of type `T` of bags' instances in [`BagNode`](@ref)s.
"""
abstract type AbstractBags{T} end

"""
    AlignedBags{T <: Integer} <: AbstractBags{T}

[`AlignedBags`](@ref) struct stores indices of bags' instances in one or more `UnitRange{T}`s.
This is only possible if instances of every bag are stored in one contiguous block.

See also: [`ScatteredBags`](@ref).
"""
struct AlignedBags{T <: Integer} <: AbstractBags{T}
    bags::Vector{UnitRange{T}}
end

Flux.@forward AlignedBags.bags Base.getindex, Base.setindex!, Base.firstindex, Base.lastindex,
    Base.eachindex, Base.first, Base.last, Base.iterate, Base.eltype, Base.length

"""
    AlignedBags()

Construct a new [`AlignedBags`](@ref) struct containing no bags.

# Examples
```jldoctest
julia> AlignedBags()
AlignedBags{Int64}(UnitRange{Int64}[])
```
"""
AlignedBags() = AlignedBags(UnitRange{Int}[])

"""
    AlignedBags(bags::UnitRange{<:Integer}...)

Construct a new [`AlignedBags`](@ref) struct from bags in arguments.

# Examples
```jldoctest
julia> AlignedBags(1:3, 4:8)
AlignedBags{Int64}(UnitRange{Int64}[1:3, 4:8])
```
"""
AlignedBags(bags::UnitRange{<:Integer}...) = AlignedBags(collect(bags))

"""
    AlignedBags(k::Vector{<:Integer})

Construct a new [`AlignedBags`](@ref) struct from `Vector` `k` specifying the index of the bag each instance belongs to.
Throws `ArgumentError` if this is not possible.

# Examples
```jldoctest
julia> AlignedBags([2, 2, 1, 1, 1, 3])
AlignedBags{Int64}(UnitRange{Int64}[1:2, 3:5, 6:6])
```
"""
function AlignedBags(k::Vector{T}) where T <: Integer
    b = AlignedBags()
    !isempty(k) || return b
    a, v = 1, k[1]
    s = Set{T}(v)
    for (i, x) in enumerate(k[2:end])
        if x != v
            !(x in s) || ArgumentError("Scattered bags") |> throw
            push!(b.bags, a:i)
            v = x
            a = i+1
            push!(s, x)
        end
    end
    push!(b.bags, a:length(k))
    b
end

"""
    length2bags(ls::Vector{<:Integer})

Convert lengths of bags given in `ls` to [`AlignedBags`](@ref) with contiguous blocks.

# Examples
```jldoctest
julia> length2bags([1, 3, 2])
AlignedBags{Int64}(UnitRange{Int64}[1:1, 2:4, 5:6])
```

See also: [`AlignedBags`](@ref).
"""
function length2bags(ls::Vector{T}) where T <: Integer
    ls = vcat([zero(T)], cumsum(ls))
    bags = map(i -> i[1] + 1:i[2], zip(ls[1:end-1], ls[2:end]))
    bags = map(b -> isempty(b) ? _empty_range(T) : b, bags)
    AlignedBags(bags)
end

ChainRulesCore.@non_differentiable length2bags(ls)

"""
    ScatteredBags{T <: Integer} <: AbstractBags{T}

[`ScatteredBags`](@ref) struct stores indices of bags' instances that are not necessarily contiguous.

See also: [`AlignedBags`](@ref).
"""
struct ScatteredBags{T <: Integer} <: AbstractBags{T}
    bags::Vector{Vector{T}}
end

Flux.@forward ScatteredBags.bags Base.getindex, Base.setindex!, Base.firstindex, Base.lastindex,
    Base.eachindex, Base.first, Base.last, Base.iterate, Base.eltype, Base.length

"""
    ScatteredBags()

Construct a new [`ScatteredBags`](@ref) struct containing no bags.

# Examples
```jldoctest
julia> ScatteredBags()
ScatteredBags{Int64}(Vector{Int64}[])
```
"""
ScatteredBags() = ScatteredBags(Vector{Vector{Int}}())

"""
    ScatteredBags(k::Vector{<:Integer})

Construct a new [`ScatteredBags`](@ref) struct from `Vector` `k` specifying the index of the bag each instance belongs to.

# Examples
```jldoctest
julia> ScatteredBags([2, 2, 1, 1, 1, 3])
ScatteredBags{Int64}([[3, 4, 5], [1, 2], [6]])
```
"""
function ScatteredBags(k::Vector{T}) where T <: Integer
    !isempty(k) || return ScatteredBags()
    d = SortedDict{T, Vector{T}}()
    for (i, x) in enumerate(k)
        if !(x in keys(d))
            d[x] = Int[]
        end
        push!(d[x], i)
    end
    ScatteredBags(collect(values(d)))
end

"""
    bags(k::Vector{<:Integer})
    bags(k::Vector{T}) where T <: UnitRange{<:Integer}
    bags(b::AbstractBags)

Construct an [`AbstractBags`](@ref) structure that is most suitable for the input
([`AlignedBags`](@ref) if possible, [`ScatteredBags`](@ref) otherwise).

# Examples
```jldoctest
julia> bags([2, 2, 3, 1])
AlignedBags{Int64}(UnitRange{Int64}[1:2, 3:3, 4:4])

julia> bags([2, 3, 1, 2])
ScatteredBags{Int64}([[3], [1, 4], [2]])

julia> bags([1:3, 4:5])
AlignedBags{Int64}(UnitRange{Int64}[1:3, 4:5])

julia> bags(ScatteredBags())
ScatteredBags{Int64}(Vector{Int64}[])
```

See also: [`AlignedBags`](@ref), [`ScatteredBags`](@ref).
"""
bags(b::AbstractBags) = b
bags(b::Vector{T}) where T <: UnitRange{<:Integer} = AlignedBags(b)
function bags(k::Vector{<:Integer})
    try
        return AlignedBags(k)
    catch ErrorException
        return ScatteredBags(k)
    end
end

"""
    remapbags(b::AbstractBags, idcs::VecOrRange{<:Integer}) -> (rb, I)

Select a subset of bags in `b` corresponding to indices `idcs` and remap instance indices appropriately.
Return new bags `rb` as well as a `Vector` of remapped instances `I`.

# Examples
```jldoctest
julia> remapbags(AlignedBags([1:1, 2:3, 4:5]), [1, 3])
(AlignedBags{Int64}(UnitRange{Int64}[1:1, 2:3]), [1, 4, 5])

julia> remapbags(ScatteredBags([[1,3], [2], Int[]]), [2])
(ScatteredBags{Int64}([[1]]), [2])
```
"""
function remapbags(b::AlignedBags{T}, idcs::VecOrRange{<:Integer}) where T
    rb = AlignedBags(Vector{UnitRange{T}}(undef, length(idcs)))
    offset = one(T)
    for (i, j) in enumerate(idcs)
        rb[i] = (b[j] == _empty_range(T)) ? b[j] : b[j] .- b[j].start .+ offset
        offset += length(b[j])
    end
    rb, Array{T}(reduce(vcat, [collect(b[i]) for i in idcs]; init=AlignedBags[]))
end

function remapbags(b::ScatteredBags{T}, idcs::VecOrRange{<:Int}) where T
    rb = ScatteredBags(Vector{Vector{T}}(undef, length(idcs)))
    m = OrderedDict{T, Int}((v => i for (i, v) in enumerate(unique(vcat(b.bags[idcs]...)))))
    for (i, j) in enumerate(idcs)
        rb[i] = [m[v] for v in b[j]]
    end
    rb, collect(keys(m))
end

"""
    adjustbags(b::AlignedBags, mask::AbstractVector{Bool})

Remove indices of instances brom bags `b` and remap the remaining instances accordingly.

# Examples
```jldoctest
julia> adjustbags(AlignedBags([1:2, 0:-1, 3:4]), [false, false, true, true])
AlignedBags{Int64}(UnitRange{Int64}[0:-1, 0:-1, 1:2])
```
"""
adjustbags(b::AlignedBags, mask::AbstractVector{Bool}) = length2bags(map(b -> sum(@view mask[b]), b))

Base.vcat(bs::AbstractBags{T}...) where T = _catbags(collect(bs))
Base.reduce(::typeof(vcat), bs::Vector{T}) where T <: AbstractBags = _catbags(bs)

function _catbags(bs::Vector{AlignedBags{T}}) where T <: Integer
    nbs = Vector{UnitRange{T}}(undef, sum(length.(bs)))
    o, no, i = zero(T), zero(T), 1
    @inbounds for b in bs
        for bb in b
            if isempty(bb)
                nbs[i] = bb
            else
                nbs[i] = bb .+ o
                no = max(no, last(nbs[i]))
            end
            i += 1
        end
        o = no
    end
    AlignedBags(nbs)
end

function _catbags(bs::Vector{ScatteredBags{T}}) where T <: Integer
    nbs = Vector{Vector{T}}(undef, sum(length.(bs)))
    o, no, i = zero(T), zero(T), 1
    for b in bs
        for bb in b
            nbs[i] = bb .+ o
            no = max(no, maximum(nbs[i]; init=0))
            i += 1
        end
        o = no
    end
    ScatteredBags(nbs)
end

Base.hash(e::AlignedBags, h::UInt) where {A,C} = hash(e.bags, h)
e1::AlignedBags == e2::AlignedBags = e1.bags == e2.bags
Base.hash(e::ScatteredBags, h::UInt) where {A,C} = hash(e.bags, h)
e1::ScatteredBags == e2::ScatteredBags = e1.bags == e2.bags


text:
struct OperatorSum{T,L<:AbstractOperator,R<:AbstractOperator} <: AbstractOperator{T}
    left::L
    right::R
end
function OperatorSum(left::L, right::R) where {L,R}
    T = promote_type(eltype(left), eltype(right))
    return OperatorSum{T,L,R}(left, right)
end

Base.show(io::IO, op::OperatorSum) = print(io, op.left, " + ", op.right)

Base.:+(left::AbstractOperator, right::AbstractOperator) = OperatorSum(left, right)

LOStructure(op::OperatorSum) = combine_structure(LOStructure(op.left), LOStructure(op.right))
combine_structure(::IsHermitian, ::IsHermitian) = IsHermitian()
combine_structure(::AdjointKnown, ::IsHermitian) = AdjointKnown()
combine_structure(::IsHermitian, ::AdjointKnown) = AdjointKnown()
combine_structure(::AdjointKnown, ::AdjointKnown) = AdjointKnown()
combine_structure(_, _) = AdjointUnknown()

Base.adjoint(op::OperatorSum) = adjoint(op.left) + adjoint(op.right)

struct OperatorSumColumn{
    A,T,L<:AbstractColumn{A,T},R<:AbstractColumn{A,T}
} <: AbstractColumn{A,T}
    left::L
    right::R
    n_left::Int
    n_right::Int
end

Base.show(io::IO, col::OperatorSumColumn) = print(io, col.left, " + ", col.right)

function column(op::OperatorSum, add::A) where {A}
    left = column(op.left, add)
    right = column(op.right, add)
    return OperatorSumColumn(left, right, num_offdiagonals(left), num_offdiagonals(right))
end
function column(op::OperatorSum, add::A) where {N,A<:CompositeFS{N}}
    left = column(op.left, add)
    right = column(op.right, add)
    return OperatorSumColumn(left, right, num_offdiagonals(left), num_offdiagonals(right))
end

num_offdiagonals(col::OperatorSumColumn) = col.n_left + col.n_right
function get_offdiagonal(col::OperatorSumColumn, chosen)
    orig_chosen = chosen
    if chosen ≤ col.n_left
        return get_offdiagonal(col.left, chosen)
    elseif chosen ≤ col.n_left + col.n_right
        chosen -= col.n_left
        return get_offdiagonal(col.right, chosen)
    else
        # TODO removing this throw makes things faster. Hide it under @boundscheck?
        throw(BoundsError(col, orig_chosen))
    end
end

function diagonal_element(col::OperatorSumColumn)
    return diagonal_element(col.left) + diagonal_element(col.right)
end


text:
export rule

@rule typeof(dot)(:out, Marginalisation, symmetrical = [:in1, :in2]) (m_in1::PointMass{ <: AbstractVector }, m_in2::NormalDistributionsFamily) = begin
    NormalMeanVariance(mean(m_in1)'*mean(m_in2), first(mean(m_in1)'*cov(m_in2)*mean(m_in1)))
end


text:
__precompile__()

module STCA
import Base.show
import Base.print
import Base.println

export CellSpace, Rule
export save, load, update!, get_state

macro NBITS(typ)
    :(sizeof($typ)*8)
end

typealias Index UInt64
typealias Count UInt64
typealias Partition UInt8     # unit state
typealias Cell UInt32         # a cell has 4 partitions
typealias Transition UInt64   # target cell with its Neumann-neighbor cells

const NBITS_P = @NBITS(Partition)
const NBITS_C = @NBITS(Cell)
const NBITS_T = @NBITS(Transition)
const BMASK_P = typemax(Partition)
const BMASK_C = typemax(Cell)

# 0x11223344 <-> 0xNNEESSWW
const BSHIFT_CELL  = NBITS_C

# 0xNNEESSWW
const BSHIFT_WEST  = NBITS_P*0
const BSHIFT_SOUTH = NBITS_P*1
const BSHIFT_EAST  = NBITS_P*2
const BSHIFT_NORTH = NBITS_P*3

const BMASK_WEST  = Cell(Partition(BMASK_P)) << BSHIFT_WEST   # 0x0000_00ff
const BMASK_SOUTH = Cell(Partition(BMASK_P)) << BSHIFT_SOUTH  # 0x0000_ff00
const BMASK_EAST  = Cell(Partition(BMASK_P)) << BSHIFT_EAST   # 0x00ff_0000
const BMASK_NORTH = Cell(Partition(BMASK_P)) << BSHIFT_NORTH  # 0xff00_0000

const QWEST_IDX  = 2  # 3x3 SubArray indices
const QSOUTH_IDX = 6  # [1][4][7]  [  ][qn][  ]
const QEAST_IDX  = 8  # [2][5][8]  [qw][tg][qe]
const QNORTH_IDX = 4  # [3][6][9]  [  ][qs][  ]
const TARGET_IDX = 5  #

include("rule.jl")
include("cellspace.jl")
include("common.jl")

end # module


text:
# ScalarField functions related to Images.jl

@info "Loading RoME.jl tools related to both ImageCore.jl and ImageIO.jl"

using .ImageCore
using .ImageIO

export generateField_CanyonDEM

"""
    $SIGNATURES

Loads a sample DEM (as if simulated) on a regular grid... It's the Grand Canyon, 18x18km, 17m
"""
function generateField_CanyonDEM( scale=1, N=100;
                                  x_is_north=true,
                                  x_min::Real=-9000, x_max::Real=9000,
                                  y_min::Real=-9000, y_max::Real=9000)
  #
  filepath = joinpath(dirname(dirname(@__DIR__)), "data","CanyonDEM.png")
  img_ = load(filepath) .|> Gray
  img_ = scale.*Float64.(img_)

  N_ = minimum([N; size(img_)...])
  img = img_[1:N_, 1:N_]

  # flip image so x-axis in plot is North and y-axis is West (ie img[-north,west], because top left is 0,0)
  _img_ = if x_is_north
    _img = collect(img')
    reverse(_img, dims=2)
  else
    # flip so north is down along with Images.jl [i,j] --> (x,y)
    reverse(img_, dims=1)
  end

  x = range(x_min, x_max, length = size(_img_,1)) # North
  y = range(y_min, y_max, length = size(_img_,2)) # East

  return (x, y, _img_)
end

#


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

using Markdown
using InteractiveUtils

# ╔═╡ a4fef4fe-f36a-11ea-244e-cdb543d14fd8
using Pkg, DrWatson

# ╔═╡ a4ff2640-f36a-11ea-062f-2d35275f2b12
begin
	@quickactivate "StatisticalRethinkingTuring"
	using StatisticalRethinking
end

# ╔═╡ ebda0e64-f369-11ea-072e-239cd87efc28
md"## Clip-03-06-10t.jl"

# ╔═╡ a4ff9ce2-f36a-11ea-2729-0705057ce156
md"### snippet 3.2"

# ╔═╡ a50ce9d8-f36a-11ea-3ef1-b3256e84a761
begin
	p_grid = range(0, step=0.001, stop=1)
	prior = ones(length(p_grid))
	likelihood = pdf.(Binomial.(9, p_grid), 6)
	posterior = likelihood .* prior
	posterior = posterior / sum(posterior)
end;

# ╔═╡ a5112304-f36a-11ea-00e8-c7a181813a3c
md"### snippet 3.3"

# ╔═╡ a51b0cd4-f36a-11ea-15c3-67e40bc741f0
md"##### Draw 10000 samples from this posterior distribution."

# ╔═╡ a51b8858-f36a-11ea-1e51-395eb0b3d22f
begin
	N = 10000
	samples = sample(p_grid, Weights(posterior), N)
end;

# ╔═╡ a5246036-f36a-11ea-06ba-e15630b3b040
md"##### Store samples in an MCMCChains.Chains object."

# ╔═╡ a524d2fa-f36a-11ea-0113-11b1e9ffe306
chn = MCMCChains.Chains(reshape(samples, N, 1, 1), [:p]);

# ╔═╡ a52f7ba6-f36a-11ea-16cd-5de42e29f3b4
md"##### Describe the chain."

# ╔═╡ a53460f8-f36a-11ea-13f9-af8cf26fdf2c
chn

# ╔═╡ 4ddc489c-f36b-11ea-121b-e3e288b12fdf
md"##### Plot the chain."

# ╔═╡ 5c369870-f36b-11ea-3eae-3b11cc0614fa
plot(chn)

# ╔═╡ a53cf4ac-f36a-11ea-075d-ffcf6547d4b0
md"### snippet 3.6"

# ╔═╡ 59dc40c4-f6d2-11ea-266b-3f329f9ee1f4
v1 = sum(posterior[filter(i -> p_grid[i] < 0.5, 1:length(p_grid))])

# ╔═╡ a5452fe4-f36a-11ea-012f-3b9363b0bc55
md"### snippet 3.7"

# ╔═╡ a54b93d6-f36a-11ea-034d-7f9a0b257a2e
b1 = mapreduce(p -> p < 0.5 ? 1 : 0, +, samples) / N

# ╔═╡ a55285c4-f36a-11ea-2efd-a1eaa158dd93
md"### snippet 3.8"

# ╔═╡ a558b1ec-f36a-11ea-1f56-b506f4939cdc
b2 = mapreduce(p -> (p > 0.5 && p < 0.75) ? 1 : 0, +, samples) / N

# ╔═╡ a55ef4ee-f36a-11ea-3ff6-8737718563c7
md"### snippet 3.9"

# ╔═╡ a5695a74-f36a-11ea-1476-a9425b2b38a5
b3 = quantile(samples, 0.8)

# ╔═╡ a56c4216-f36a-11ea-20c0-0fd3bb1e32a0
md"### snippet 3.10"

# ╔═╡ a5730a24-f36a-11ea-3d02-df3631450b3b
b4 = quantile(samples, [0.1, 0.9])

# ╔═╡ 42631124-f6d3-11ea-2155-9703136d8189
begin
	p1 = plot_density_interval(samples, [0.0, 0.5],
	  xlab="Proportion water (p)");
	p2 = plot_density_interval(samples, [0.5, 0.75],
	  xlab="Proportion water (p)");
	p3 = plot_density_interval(samples, [0.0, b3],
	  xlab="Proportion water (p)");
	p4 = plot_density_interval(samples, b4,
	  xlab="Proportion water (p)");
	plot(p1, p2, p3, p4, layout=(2, 2))
end

# ╔═╡ a57a49ec-f36a-11ea-167c-e1d90e858751
md"## End of clip-03-06-10t.jl"

# ╔═╡ Cell order:
# ╟─ebda0e64-f369-11ea-072e-239cd87efc28
# ╠═a4fef4fe-f36a-11ea-244e-cdb543d14fd8
# ╠═a4ff2640-f36a-11ea-062f-2d35275f2b12
# ╟─a4ff9ce2-f36a-11ea-2729-0705057ce156
# ╠═a50ce9d8-f36a-11ea-3ef1-b3256e84a761
# ╟─a5112304-f36a-11ea-00e8-c7a181813a3c
# ╟─a51b0cd4-f36a-11ea-15c3-67e40bc741f0
# ╠═a51b8858-f36a-11ea-1e51-395eb0b3d22f
# ╟─a5246036-f36a-11ea-06ba-e15630b3b040
# ╠═a524d2fa-f36a-11ea-0113-11b1e9ffe306
# ╟─a52f7ba6-f36a-11ea-16cd-5de42e29f3b4
# ╠═a53460f8-f36a-11ea-13f9-af8cf26fdf2c
# ╟─4ddc489c-f36b-11ea-121b-e3e288b12fdf
# ╠═5c369870-f36b-11ea-3eae-3b11cc0614fa
# ╟─a53cf4ac-f36a-11ea-075d-ffcf6547d4b0
# ╠═59dc40c4-f6d2-11ea-266b-3f329f9ee1f4
# ╟─a5452fe4-f36a-11ea-012f-3b9363b0bc55
# ╠═a54b93d6-f36a-11ea-034d-7f9a0b257a2e
# ╟─a55285c4-f36a-11ea-2efd-a1eaa158dd93
# ╠═a558b1ec-f36a-11ea-1f56-b506f4939cdc
# ╟─a55ef4ee-f36a-11ea-3ff6-8737718563c7
# ╠═a5695a74-f36a-11ea-1476-a9425b2b38a5
# ╟─a56c4216-f36a-11ea-20c0-0fd3bb1e32a0
# ╠═a5730a24-f36a-11ea-3d02-df3631450b3b
# ╠═42631124-f6d3-11ea-2155-9703136d8189
# ╟─a57a49ec-f36a-11ea-167c-e1d90e858751


text:
using FastAI
import FastAI: createhandle, showblock!

"""
    ShowMakie([; kwargs...]) <: ShowBackend

A backend for showing block data that uses
[Makie.jl](https://github.com/JuliaPlots/Makie.jl) figures for
visualization.

Keyword arguments are passed through to the constructed `Figure`s.
"""
struct ShowMakie <: ShowBackend
    size::Tuple{Int,Int}
    kwargs
end
ShowMakie(sz=(500, 500); kwargs...) = ShowMakie(sz, kwargs)

function createhandle(backend::ShowMakie; kwargs...)
    fig = Figure(;kwargs..., backend.kwargs...)
    grid = fig[1, 1] = GridLayout()
    return grid
end

function showblock(backend::ShowMakie, block, data)
    # Calculate resolution based on number of blocks and size for one block
    # A margin of 10% is added for titles
    width = _nblocks(block) * backend.size[2]
    height = round(Int, 1.1 * backend.size[1])

    grid = createhandle(backend, resolution=(width, height))
    fig = grid.parent.parent
    showblock!(grid, backend, block, data)
    return fig
end

_nblocks(t::Tuple) = sum(_nblocks, t)
_nblocks(b::AbstractBlock) = 1
_nblocks((_, block)::Pair) = _nblocks(block)

function showblock!(grid, backend::ShowMakie, (title, block)::Pair, data)
    showblock!(grid[1, 1], backend, block, data)
end

function showblock!(grid, backend::ShowMakie, blocks::Tuple, datas::Tuple)
    header = [block isa Pair ? first(block) : "" for block in blocks]
    blocks = Tuple(block isa Pair ? last(block) : block for block in blocks)

    rowsize!(grid, 1, Makie.Fixed(backend.size[2]))

    # Show blocks in a row
    col = 1
    for (i, (block, data)) in enumerate(zip(blocks, datas))
        w = _nblocks(block)
        subgrid = grid[1, col:col+w-1] = GridLayout(tellheight=false)
        showblock!(subgrid, backend, block, data)
        for j in col:col+w-1
            colsize!(grid, j, Makie.Fixed(backend.size[1]))
        end
        col += w
    end

    # Add titles to named blocks
    Makie.Label(grid[0, 1], "", tellwidth=false, textsize=25)
    for (i, title) in enumerate(header)
        Makie.Label(grid[1, i], title, tellwidth=false, textsize=25)
    end

end

function showblocks(backend::ShowMakie, block, datas)
    width = _nblocks(block) * backend.size[2]
    height = round(Int, length(datas) * 1.1 * backend.size[1])
    res = (width*1.2, height*1.1)

    grid = createhandle(backend, resolution=res)
    fig = grid.parent.parent

    showblocks!(grid, backend, block, datas)
    return fig
end

function showblocks!(grid, backend::ShowMakie, blocks::Tuple, datas::AbstractVector)
    header = [block isa Pair ? first(block) : "" for block in blocks]
    blocks = Tuple(block isa Pair ? last(block) : block for block in blocks)
    n = _nblocks(blocks)

    # Show each sample in one row
    for (i, data) in enumerate(datas)
        subgrid = grid[i, 1:n] = GridLayout(tellheight=false)
        rowsize!(grid, i, Makie.Fixed(backend.size[2]))
        showblock!(subgrid, backend, blocks, data)
    end

    for i in 1:n
        colsize!(grid, i, Makie.Fixed(backend.size[1]))
    end

    # Add titles to named blocks
    Makie.Label(grid[0, 1], "", textsize=25)
    col = 1
    for (i, (title, block)) in enumerate(zip(header, blocks))
        w = _nblocks(block)
        Makie.Label(grid[1, col:col+w-1], title, tellwidth=false, textsize=25)
        col += w
    end
end
showblocks!(grid, backend::ShowMakie, block, datas::AbstractVector) =
    showblocks!(grid, backend, (block,), map(data -> (data,), datas))

## Block definitions

function showblock!(grid, ::ShowMakie, block::Label, data)
    ax = imageaxis(grid[1, 1])
    text!(ax, string(data), space=:data)
end

function showblock!(grid, ::ShowMakie, block::LabelMulti, data)
    ax = imageaxis(grid[1, 1])
    text!(ax, join(string.(data), "\n"), space=:data)
end

function showblock!(grid, ::ShowMakie, block::Union{<:OneHotTensor{0},<:OneHotTensorMulti{0}}, data)
    if !(sum(data) ≈ 1)
        data = softmax(data)
    end
    ax = Axis(grid[1, 1], yticks=(1:length(block.classes), string.(block.classes),))
    barplot!(
		ax,
		data,
		direction=:x,
	)
    hidespines!(ax)
end

function showblock!(grid, ::ShowMakie, block::Image{2}, data)
    ax = imageaxis(grid[1, 1])
    plotimage!(ax, data)
end
function showblock!(grid, ::ShowMakie, block::Mask{2}, data)
    ax = imageaxis(grid[1, 1])
    plotmask!(ax, data, block.classes)
end

function showblock!(grid, ::ShowMakie, block::Keypoints{2}, data)
    ax = imageaxis(grid[1, 1])
    h = maximum(first.(data))
    ks = [SVector(x, h-y) for (y, x) in data]
    scatter!(ax, ks)
end

function showblock!(grid, ::ShowMakie, block::Bounded{2, <:Keypoints{2}}, data)
    ax = imageaxis(grid[1, 1])
    h, w = block.size
    ks = [SVector(x, h-y) for (y, x) in data]
    xlims!(ax, 0, w)
    ylims!(ax, 0, h)
    scatter!(ax, ks)
end

function default_showbackend()
    if ismissing(Makie.current_backend[])
        return ShowText()
    else
        return ShowMakie()
    end
end


text:
using nctools
# using Glob
# using NCDatasets
# using CFTime

# rh = rand(3, 3, 10) * 100
# tair = rand(3, 3, 10)
# heat_index(tair, rh)

# GFDL-CM4, r1i1p1f1_gr2, 2.5, 2.0
# GFDL-CM4, r1i1p1f1_gr1, 1.25, 1.0

# files = glob("k:/Researches/CMIP6/CMIP6_ChinaHW_mergedFiles/HItasmax/historical/*.nc")
# d = CMIPFiles_info(files)
# files = glob("k:/Researches/CMIP6/CMIP6_ChinaHW_mergedFiles/HItasmax/historical/")
# dir_prj  = "$pan/Researches/CMIP6/ChinaHW_cluster"

pan = path_mnt("K:")
dir_data = "$pan/Researches/CMIP6/CMIP6_ChinaHW_mergedFiles"

dirs = dir(dir_data)
files = dir(dirs[1], "nc", recursive = true)
d = CMIPFiles_info(files)

# file = files[1]
# fid = nc_open(file)


text:
using Revise
module BayesMNIST

using Statistics, LinearAlgebra, StatsBase, Flux, Flux.Data.MNIST, DistributionsAD
using Flux: softmax, onehotbatch#, params
using DynamicHMC, LogDensityProblems, TransformVariables, DynamicHMC.Diagnostics, Parameters
using CuArrays, Random, Distributions, Zygote, BenchmarkTools, InteractiveUtils#, Plots
using CUDAapi
if has_cuda()
    @info "CUDA is on"
    import CuArrays
    CuArrays.allowscalar(false)
end
import ForwardDiff

using Revise

#load the data
#the purpose is to explore flux in a Bayesian context

Random.seed!(11)

const D=784
const K=10

#Lasso Logistic Categorical regression
struct LCRegression{TY, TX, Tατ, Tθτ}
  Y::TY
  X::TX

  ατβ::Tατ
  θτβ::Tθτ

end

function (problem::LCRegression)(θ)
    #unpack parameters and data
    @unpack Y, X, ατβ, θτβ = problem
    @unpack βs, τβs = θ
    Ncovariates, Npoints = size(X,1),size(X,2)
    Nclasses = size(βs,2)

    @assert size(βs,1) == Ncovariates
    @assert size(βs,2) == length(τβs) == Nclasses

    #p(τβᵢ)
    logτβpri = sum([logpdf(Gamma(ατβ, θτβ), τβs[i]) for i ∈ 1:Nclasses])

    #p(βᵢ|τβᵢ)
    logβpri = sum([logpdf(MvNormal(zeros(Ncovariates), 1/τβs[i]), βs[:,i]) for i ∈ 1:Nclasses])

    μs = softmax(X' * βs, dims=2)
    #μs = (xᵢ->softmax(βs'*xᵢ)).(X)
    #  display(μs[1])
    #p(Y|β,τβ)
    #loglik = sum([logpdf(Categorical(μs[n, :]), Y[n]) for n ∈ 1:Npoints])
    #println(size(μs[1]))
    #println(size(Y[1]))
    #display(logpdf(Multinomial(1, μs[1]), Y[1]))
    loglik = sum([logpdf(Multinomial(1,μs[n,:]), Y[:,n]) for n ∈ 1:Npoints])

    return logτβpri + logβpri + loglik
end

#hyperpriors
const Πmwe = Dict{Symbol, Float64}(
  :ατβ=>1.0,
  :θτβ=>1.0
)

function LCtransform(p::LCRegression)
    return as((βs=as(Array, D+1, K),
      τβs=as(Array, asℝ₊, K)))
end

function dynamicmnistmwe(Npoints::Int = 1000,
    ::Type{V} = AbstractVector, ::Type{T}=Float32) where {V,T}
  imgs = MNIST.images()
  labels = MNIST.labels()

  #make each image into a feature vector
  #concatenate each feature vector into a matrix
  preX =hcat(float.(reshape.(imgs,:))...)

  #add an intercept
  preX = vcat(preX, ones(1,size(preX,2)))[:,1:Npoints]
  #X = (Vector).(eachcol(preX)) #|>gpu
  X = preX
  preY = onehotbatch(labels[1:Npoints], 0:9)
  Y = Matrix(preY)

  #Y=labels[1:5_000]

  #make the problem
  p = LCRegression(Y, X, Πmwe[:ατβ], Πmwe[:θτβ])

  #transform the support
  t = LCtransform(p)
  P = TransformedLogDensity(t, p)

  #differentiate via AD
  ∇P =  ADgradient(:Zygote, P)

  @info "beginning MCMC"

  testβs = rand(D+1,K)# |>gpu
  testτβs = rand(K)

  #InteractiveUtils.@code_warntype p((βs=testβs, τβs=testτβs))
  #@btime $p((βs=$testβs, τβs=$testτβs))
  results = mcmc_with_warmup(Random.GLOBAL_RNG, ∇P, 1000, reporter = LogProgressReport())
  serialize("results.jls", results)

  posterior = transform.(t, results.chain)

  β_posterior = first.(transform.(t, results.chain))
  display(mean(β_posterior))

end

#dynamicmnistmwe()

end # module

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
using PkgTemplates

t = Template(;
        user="csimal",
        license="MIT",
        authors=["Cedric Simal"],
        plugins=[
            TravisCI(),
            Codecov(),
            Coveralls(),
            AppVeyor(),
            GitHubPages(),
           ],
       )

generate("TestPackage", t)

max_stars_repo_name: csimal/metaplex-networks

max_issues_repo_name: csimal/metaplex-networks

max_forks_repo_name: csimal/metaplex-networks


text:
module Octree

using GRU.Shapes
using GRU.Shapes: intersect, isvalid

export Tree, add, remove, for_overlapping, getbound

abstract type OptionalNode{T} end

mutable struct NullNode{T} <: OptionalNode{T}
end

mutable struct Node{T} <: OptionalNode{T}
	subNodes::Array{OptionalNode{T}, 3}
	objects::Array{T, 1}

	Node{T}() where T = new(fill(NullNode{T}(), 2, 2, 2), T[])
end

mutable struct Tree{T, F}
	bound::AABB{F}
	root::Node{T}
	minNodeSize::F
	canExpand::Bool
end

Tree(::Type{T}, bound::AABB{F}; minNodeSize::F = F(16), canExpand::Bool = false) where {T, F} = Tree{T, F}(bound, Node{T}(), minNodeSize, canExpand)

function add(octree::Tree{T, F}, obj::T, objBound::AABB{F}) where {T, F}
	@assert isvalid(objBound)
	@assert !any(x->isinf(x), objBound.p)
	if !inside(octree.bound, objBound)
		if octree.canExpand
			expand(octree, objBound)
			@assert inside(octree.bound, objBound)
		else
			error("Octree: Trying to add an object outside of tree bounds")
		end
	end
	add(octree.root, octree.bound, octree, obj, objBound)
end
add(octree::Tree{T, F}, obj::T) where {T, F} = add(octree, obj, getbound(obj))

remove(octree::Tree{T, F}, obj::T, objBound::AABB{F}) where {T, F} = remove(octree.root, octree.bound, octree, obj, objBound)
remove(octree::Tree{T, F}, obj::T) where {T, F} = remove(octree, obj, getbound(obj))

for_overlapping(f::Function, octree::Tree{T, F}, bound::AABB{F}) where {T, F} = for_overlapping(f, octree.root, octree.bound, bound)

getbound(::Any) = nothing # this needs to be overriden by the user

function expand(octree::Tree{T, F}, objBound::AABB{F}) where {T, F}
	ind = Int[0, 0, 0]
	while true
		for i = 1:3
			if objBound.p[i, 1] < octree.bound.p[i, 1]
				ind[i] = 2
			elseif objBound.p[i, 2] > octree.bound.p[i, 2]
				ind[i] = 1
			else
				ind[i] = 0
			end
		end
		if all(x->x==0, ind)
			break;
		end
		for i = 1:3
			if ind[i] == 0 # In case the bound we want to fit doesn't exit a dimension, grow symmetrically around 0
				ind[i] = abs(octree.bound.p[i, 1]) < abs(octree.bound.p[i, 2]) ? 2 : 1;
			end
			sz = octree.bound.p[i, 2] - octree.bound.p[i, 1]
			if ind[i] == 1
				octree.bound.p[i, 2] += sz
			else
				@assert ind[i] == 2
				octree.bound.p[i, 1] -= sz
			end
		end
		newRoot = Node{T}()
		newRoot.subNodes[ind...] = octree.root
		octree.root = newRoot
	end
end

cansubdivide(nodeBound::AABB{F}, octree::Tree{T, F}) where {T, F} = minimum(nodeBound.p[:, 2] - nodeBound.p[:, 1]) / 2 >= octree.minNodeSize

function getsubbound(nodeBound::AABB{F}, ind::Vector{Int}) where F
	subBound = AABB{F}()
	for i=1:3
		mid = (nodeBound.p[i, 1] + nodeBound.p[i, 2]) / 2
		subBound.p[i, ind[i]] = nodeBound.p[i, ind[i]]
		subBound.p[i, 3-ind[i]] = mid
	end
	return subBound
end

function getsubindex(nodeBound::AABB{F}, objBound::AABB{F}) where F
	@assert inside(nodeBound, objBound)
	ind = Int[]
	for i=1:3
		mid = (nodeBound.p[i, 1] + nodeBound.p[i, 2]) / 2
		if objBound.p[i, 2] <= mid
			push!(ind, 1)
		elseif objBound.p[i, 1] >= mid
			push!(ind, 2)
		else
			return ind
		end
	end
	return ind
end

function getsubnode(node::Node{T}, nodeBound::AABB{F}, octree::Tree{T, F}, objBound::AABB{F}, create::Bool) where {T, F}
	if create && !cansubdivide(nodeBound, octree)
		return NullNode{T}(), nodeBound
	end
	ind = getsubindex(nodeBound, objBound)
	if length(ind) != 3
		return NullNode{T}(), nodeBound
	end
	subNode = node.subNodes[ind...]
	if isa(subNode, NullNode{T})
		if !create
			return NullNode{T}(), nodeBound
		end
		subNode = Node{T}()
		node.subNodes[ind...] = subNode
	end
	return subNode, getsubbound(nodeBound, ind)
end

function add(node::Node{T}, nodeBound::AABB{F}, octree::Tree{T, F}, obj::T, objBound::AABB{F}) where {T, F}
	subNode, subBound = getsubnode(node, nodeBound, octree, objBound, true)
	if isa(subNode, NullNode{T})
		@assert findfirst(isequal(obj), node.objects) == nothing
		push!(node.objects, obj)
	else
		add(subNode, subBound, octree, obj, objBound)
	end
end

function deleteat_unordered!(v::Vector{T}, ind::Int) where T
	v[ind] = v[end]
	resize!(v, length(v)-1)
end

function remove(node::Node{T}, nodeBound::AABB{F}, octree::Tree{T, F}, obj::T, objBound::AABB{F}) where {T, F}
	subNode, subBound = getsubnode(node, nodeBound, octree, objBound, false)
	if isa(subNode, NullNode{T})
		deleteat_unordered!(node.objects, findfirst(node.objects, obj))
	else
		remove(subNode, subBound, octree, obj, objBound)
	end
end

function for_overlapping(f::Function, node::Node{T}, nodeBound::AABB{F}, bound::AABB{F}) where {T, F}
	for obj in node.objects
		if intersect(bound, getbound(obj))
			f(obj)
		end
	end
	for z = 1:2, y = 1:2, x = 1:2
		if !isa(node.subNodes[x, y, z], NullNode{T})
			subBound = getsubbound(nodeBound, [x, y, z])
			if intersect(bound, subBound)
				for_overlapping(f, node.subNodes[x, y, z], subBound, bound)
			end
		end
	end
end

#=
ab = AABB([-32.0, -32, -32], [32.0, 32, 32])
ab1 = AABB([-2.0, -2, -2], [-1.0, -1, -1])
ab5 = AABB([-5.0, -5, -5], [5.0, 5, 5])
ab10 = AABB([10.0, 10, 10], [12.0, 12, 12])
ab50 = AABB([50.0, 0, 10], [60.0, 10, 20])

getbound(ab::AABB) = ab

oc = Tree(typeof(ab5), ab; canExpand=true)
add(oc, ab5)
add(oc, ab1)
add(oc, ab10)
add(oc, ab50)
for_overlapping(oc, ab5) do ab
	println(ab)
end
remove(oc, ab5)
println()
for o in @async for_overlapping(x->produce(x), oc, oc.bound)
	println(o)
end
=#

end