text:
# ----------------------------------------------------------
# -------------  @recycle optimization macro  -------------
# ----------------------------------------------------------

# Get all variable names from an expression
function getSymbols(expr)
    var_names = []
    MacroTools.postwalk(expr) do x
        #x isa Expr && [push!(var_names, s)
        #    for s in x.args[(x.head in (:call, :macrocall) ? 2 : 1):end] if s isa Symbol]
        x isa Symbol && Base.isidentifier(x) && push!(var_names, x)
        x
    end
    unique(var_names)
end

# Transform a number to subscript number (E.g.: 13 -> ₁₃)
toSubscript(num::Int64) = join([Char(0x2080 + parse(Int, d)) for d in split(string(num), "")])

function newExtractedFunOp(op, funops, extracted_ops)
    newSymbol = gensym("🔝$(toSubscript(length(extracted_ops)+1))")
    extracted_ops[newSymbol] = op
    push!(funops, newSymbol)
    newSymbol
end

# Receive a list of multiplication terms (sides), and replace all series of FunOps / any
# compatible types (numbers, scalings) with a unique identifier assigned to the replaced FunOp
# series before the optimized expression
# E.g.: Op₁ * Op₂ * array -> 🔝₁ = Op₁ * Op₂ (before the entire expression), 🔝₁ * array
# Replaced/extracted series cannot contain any variable modified inside the optimized expr
function extractFunOp(sides, funops, numbers, scalings, extracted_ops, modified)
    next_idx = 0
    new_sides = []
    temp = []
    for e in sides
        if (e in funops || e in numbers || e in scalings) && !(e in modified)
            push!(temp, e)
        else
            if length(temp) > 1 || length(temp) > 0 && temp[1] isa Expr
                op = length(temp) > 1 ? Expr(:call, :*, temp...) : temp[1]
                temp = [newExtractedFunOp(op, funops, extracted_ops)]
            end
            push!(new_sides, temp..., e)
            temp = []
        end
    end
    new_sides
end

# Get type of expressions and extract combined/adjoint operators applied to an array
# Eg: a * b is a number, if both a and b are numbers, but a * B is an array if B is array
# Eg: A * B * C -> A * B is extracted (and replaced by 🔝₁), if A, B are operators, and C is an array
function preprocess(expr_str, arrays, funops, numbers, scalings, extracted_ops, modified)
    MacroTools.postwalk(Meta.parse(expr_str)) do x
        if @capture(x, lhs_ + rhs__) || @capture(x, lhs_ .+ rhs__) ||
                @capture(x, lhs_ - rhs__) || @capture(x, lhs_ .- rhs__)
            any(e -> e in funops, [lhs, rhs...]) && !any(e -> e in modified, [lhs, rhs...]) &&
                return newExtractedFunOp(x, funops, extracted_ops)
            any(e -> e in arrays, [lhs, rhs...]) && push!(arrays, x)
            all(e -> e in numbers, [lhs, rhs...]) && push!(numbers, x)
        elseif @capture(x, lhs_ * rhs__)
            sides = [lhs, rhs...]
            if any(e -> e in arrays, sides)
                any(e -> e in funops, sides) &&
                    (sides = extractFunOp(sides, funops, numbers, scalings, extracted_ops, modified))
                push!(arrays, Expr(:call, :*, sides...))
                return Expr(:call, :*, sides...)
            elseif all(x -> (x in numbers || x in scalings), sides) &&
                    any(x -> x in scalings, sides)
                push!(scalings, x)
            elseif all(e -> e in numbers, sides)
                push!(numbers, x)
            end
        elseif @capture(x, lhs_ .* rhs__)
            push!(arrays, x)
        elseif @capture(x, lhs_ / rhs_) || @capture(x, lhs_ ./ rhs_) ||
                @capture(x, lhs_ \ rhs_) || @capture(x, lhs_ .\ rhs_) ||
                @capture(x, lhs_ ÷ rhs_) || @capture(x, lhs_ .÷ rhs_) ||
                @capture(x, lhs_ % rhs_) || @capture(x, lhs_ .% rhs_) ||
                @capture(x, lhs_ ^ rhs_) || @capture(x, lhs_ .^ rhs_)
            any(e -> e in arrays, [lhs, rhs]) && push!(arrays, x)
            all(e -> e in numbers, [lhs, rhs]) && push!(numbers, x)
        elseif x isa Expr && x.head in (:(=), :+=, :-=, :*=, :/=, :\=, :÷=, :%=, :^=, :&=, :|=, :⊻=, :>>>=, :>>=, :<<=, :(.=), :.+=, :.-=, :.*=, :./=, :.\=, :.÷=, :.%=, :.^=, :.&=, :.|=, :.⊻=, :.>>>=, :.>>=, :.<<=)
            lhs, rhs = x.args
            push!(modified, lhs)
            rhs in funops && push!(funops, lhs)
            rhs in arrays && push!(arrays, lhs)
            rhs in numbers && push!(numbers, lhs)
            rhs in scalings && push!(scalings, lhs)
        elseif x isa Expr && string(x.head) == "'"                        # adjoint
            x.args[1] in funops && push!(funops, x)
            x.args[1] in arrays && push!(arrays, x)
        elseif x isa Expr && x.head == :call && x.args[1] in (:-, :+, :√) # unary operators
            x.args[2] in numbers && push!(numbers, x)
            x.args[2] in arrays && push!(arrays, x)
        elseif x isa Number
            push!(numbers, x)
        end
        x
    end
end

# Interestingly, a*b*c is parsed to *(a,b,c), but a.*b.*c is parsed to .*(.*(a,b),c),
# and, accordingly, .*(a,b,c) doesn't work.
# This function creates a correct expression from and operator and multiple terms
# E.g.: multiDot(:(.*), [a,b,c]) => .*(.*(a,b),c)
function multiDot(op, expr_list)
    length(expr_list) > 2 ?
        Expr(:call, op, multiDot(op, expr_list[1:end-1]), expr_list[end]) :
        Expr(:call, op, expr_list[1], expr_list[2])
end

function fixUniformScalingForDotify(terms, scalings)
    new_terms = []
    for term in terms
        push!(new_terms, term in scalings ? :(($term).λ) : term)
    end
    new_terms
end

# Attempts to replace +, -, *, and / operators with their respective "dot" operations
# whenever it is possible
function dotify(expr, arrays, scalings, numbers)
    MacroTools.postwalk(expr) do x
        if @capture(x, lhs_ + rhs__)
            terms = [lhs, rhs...]
            if any(x -> x in arrays, terms)
                terms = fixUniformScalingForDotify(terms, scalings)
                new_expr = length(terms) > 1 ? multiDot(:.+, terms) : Expr(:call, :.+, terms[1])
                push!(arrays, new_expr)
                return new_expr
            end
        elseif @capture(x, lhs_ - rhs_)
            if lhs in arrays || rhs in arrays
                lhs, rhs = fixUniformScalingForDotify([lhs, rhs], scalings)
                new_expr = :($lhs .- $rhs)
                push!(arrays, new_expr)
                return new_expr
            end
        elseif @capture(x, - rhs_)
            if rhs in arrays
                new_expr = :( .- $rhs)
                push!(arrays, new_expr)
                return new_expr
            end
        elseif @capture(x, lhs_ / rhs_)
            if lhs in arrays && rhs in numbers || lhs in numbers && rhs in arrays
                lhs, rhs = fixUniformScalingForDotify([lhs, rhs], scalings)
                new_expr = :($lhs ./ $rhs)
                push!(arrays, new_expr)
                return new_expr
            end
        elseif @capture(x, lhs_ * rhs__)
            terms = [lhs, rhs...]
            if any(x -> x in arrays, terms)
                left_last_number = 0
                for (i,e) in enumerate(terms)
                    e in numbers ? (left_last_number = i) : break
                end
                right_last_number = length(terms) + 1
                for (i,e) in enumerate(reverse(terms))
                    e in numbers ? (right_last_number = length(terms) - i + 1) : break
                end
                if left_last_number > 0 || right_last_number < length(terms) + 1
                    terms = fixUniformScalingForDotify(terms, scalings)
                    new_expr = left_last_number + 1 == right_last_number - 1 ||
                                    right_last_number == length(terms) ?
                        multiDot(:(.*), terms) :
                        multiDot(:(.*), vcat(terms[1:left_last_number],
                            Expr(:call, :*, terms[left_last_number+1:right_last_number-1]...),
                            terms[right_last_number:end]))
                    push!(arrays, new_expr)
                    return new_expr
                end
            end
        end
        return x
    end
end

# Tell if the given expression's top level operation is a "dot" operation
# E.g.: a .+ (b * c) => true, a + (b .* c) => false
isDotExpr(expr) = expr isa Expr && expr.head == :call &&
    (startswith(string(expr.args[1]), ".") || endswith(string(expr.args[1]), "."))

# Create a new identifier for a buffer and "register" it (i.e. add to "buffers" list,
# so it will be initialized to nothing before the entire expression)
function newTempBuffer(buffers)
    newSymbol = gensym("🔃$(toSubscript(length(buffers)+1))")
    push!(buffers, newSymbol)
    newSymbol
end

# Just to make sure that my helper function calls doesn't conflict with functions
# created by the user. Actually, these function that are formally called are no functions at all,
# but rather they are used to mark to spots where further replacements needed later
const ifFirstEqual = gensym("ifFirstEqual")
const ifFirstMult = gensym("ifFirstMult")

# Add a buffer to store the result of a "dot" expression, if applicable
# E.g.: A .+ (B * C) -> ifFirstEqual(🔃ₓ, A .+ (B * C), :(=), :(.=))
#   which will be later (in postprocess function) transformed to:
# if is_first_runₓ
#    is_first_runₓ
#    ...
#    🔃ₓ = A .+ (B * C)
#    ...
# else
#    ...
#    🔃ₓ .= A .+ (B * C)
#    ...
# end
# ...🔃ₓ...
function bufferIfDotExpr(expr, buffers)
    if isDotExpr(expr)
        buffer = newTempBuffer(buffers)
        push!(buffers, buffer)
        Expr(:call, ifFirstEqual, buffer, expr, :(=), :(.=))
    else
        expr
    end
end

# Try to add a buffer to store result of "dot" expressions, and replace array-array and funop-array
# multiplications with a call to mul!, and also add a buffer to store result of mul!
function buffer(expr, arrays, array_symbols, funops, buffers)
    # Capture the expressions where buffer is already given (i.e., as left side of assignment)
    # E.g.: A = B * C -> is_first_run ? A = B * C : mul!(A, B, C)
    expr = MacroTools.prewalk(expr) do x
        if x isa Expr && x.head in (:(=), :+=, :-=, :*=, :/=, :(.=), :.+=, :.-=, :.*=, :./=)
            lhs, rhs = x.args[1], x.args[2]
            if isDotExpr(x.args[2])
                dotOp = Symbol(".", x.head)
                if !(lhs in array_symbols)
                    push!(buffers, lhs)
                    Expr(:call, ifFirstEqual, lhs, rhs, x.head, dotOp)
                else
                    Expr(dotOp, lhs, rhs)
                end
            elseif @capture(rhs, rhs1_ * rhs2_) && any(x -> x in arrays, [lhs, rhs1, rhs2])
                if lhs in array_symbols # if lhs is already assigned
                    :(mul!($lhs, $rhs1, $rhs2))
                else
                    Expr(:call, ifFirstMult, lhs, rhs1, rhs2)
                end
            elseif @capture(rhs, rhs1_ * rhs2__) && (rhs1 in arrays || any(x -> x in arrays, rhs2))
                if lhs in array_symbols # if lhs is already assigned
                    Expr(:call, :mul!, lhs, rhs1, Expr(:call, :*, rhs2...))
                else
                    Expr(:call, ifFirstMult, lhs, rhs1, Expr(:call, :*, rhs2...))
                end
            else
                x
            end
        else
            x
        end
    end
    # Capture all other possible cases
    # ("dot" expressions, and array-array, funop-array multiplications)
    MacroTools.postwalk(expr) do x
        if @capture(x, rhs1_ * rhs2_) && (rhs1 in arrays || rhs2 in arrays)
            lhs = newTempBuffer(buffers)
            rhs1, rhs2 = bufferIfDotExpr(rhs1, buffers), bufferIfDotExpr(rhs2, buffers)
            Expr(:call, ifFirstMult, lhs, rhs1, rhs2)
        elseif @capture(x, rhs1_ * rhs2__) && (rhs1 in arrays || any(x -> x in arrays, rhs2))
            lhs = newTempBuffer(buffers)
            rhs1 = bufferIfDotExpr(rhs1, buffers)
            rhs2 = map(x -> bufferIfDotExpr(x, buffers), rhs2)
            Expr(:call, ifFirstMult, lhs, rhs1, Expr(:call, :*, rhs2...))
        elseif x isa Expr && x.head == :call && x.args[1] != ifFirstEqual &&
                !(startswith(string(x.args[1]), ".") || endswith(string(x.args[1]), "."))
            Expr(:call, x.args[1], map(x -> bufferIfDotExpr(x, buffers), x.args[2:end])...)
        else
            x
        end
    end
end

# Create a new identifier for a bool variable and "register" it (i.e. add to "first_bools" list,
# so it will be initialized to true before the entire expression)
function newFirstBool(first_bools)
    newSymbol = gensym("is_first_run$(toSubscript(length(first_bools)+1))")
    push!(first_bools, newSymbol)
    newSymbol
end

# Create an if-expression that holds all calculations of intermediate values for a given line
# For example:
# if is_first_runₓ
#    is_first_runₓ
#    ...
#    🔃ₓ = A .+ (B * C)
#    ...
# else
#    ...
#    🔃ₓ .= A .+ (B * C)
#    ...
# end
function createIfBlock(x, first, second, first_bools)
    first_bool = newFirstBool(first_bools)
    if_block = Expr(:if, first_bool,
        Expr(:block, :($first_bool = false), first...), Expr(:block, second...))
    if x isa Expr && x.head == :macrocall && length(x.args) == 3
        x.args[3] isa Symbol ?
            Expr(:macrocall, x.args[1], x.args[2], if_block) :
            Expr(:macrocall, x.args[1], x.args[2], Expr(:block, if_block, x.args[3]))
    elseif x isa Expr
        Expr(:block, if_block, x)
    else
        if_block
    end
end

# Remove virtual calls to ifFirstEqual and ifFirstMult "functions", and add the respective if-block
function postprocess(expr, first_bools)
    # Dicts holding intermediate calculations
    first, otherwise = Dict(), Dict() # key: line number, value: expressions
    current_line = 0
    # Just gather all intermediate calculations (for both first and later runs)
    expr = MacroTools.postwalk(expr) do x
        if @capture(x, fun_name_(lhs_, rhs_, op_, dotOp_)) && fun_name == ifFirstEqual
            push!(get!(first, current_line, []), Expr(op, lhs, rhs))
            push!(get!(otherwise, current_line, []), Expr(dotOp, lhs, rhs))
            x = lhs
        elseif @capture(x, fun_name_(lhs_, rhs1_, rhs2_)) && fun_name == ifFirstMult
            push!(get!(first, current_line, []),
                Expr(:(=), lhs, Expr(:call, :*, rhs1, rhs2)))
            push!(get!(otherwise, current_line, []),
                Expr(:call, :mul!, lhs, rhs1, rhs2))
            x = lhs
        end
        if x isa LineNumberNode
            current_line = x.line
        end
        x
    end
    current_line_changed = current_line == 0
    # Create if-blocks from gathered intermediate calculations
    expr = MacroTools.prewalk(expr) do x
        if x isa LineNumberNode && current_line != x.line
            current_line = x.line
            current_line_changed = true
            x
        elseif current_line_changed && # top level = we see an entire line
                current_line in keys(first) && length(first[current_line]) > 0
            current_line_changed = false
            createIfBlock(x, first[current_line], otherwise[current_line], first_bools)
        else
            x
        end
    end
    # Expand blocks nested directly in other blocks
    # E.g:
    # begin
    #     a = 3
    #     begin
    #         b = 4
    #         c = 7
    #     end
    # end
    # ...transformed to
    # begin
    #     a = 3
    #     b = 4
    #     c = 7
    # end
    MacroTools.postwalk(expr) do x
        if x isa Expr && x.head == :block
            new_block_content = []
            for item in x.args
                if item isa Expr && item.head == :block
                    for item2 in item.args
                        push!(new_block_content, item2)
                    end
                else
                    push!(new_block_content, item)
                end
            end
            Expr(:block, new_block_content...)
        else
            x
        end
    end
end

# It is not a necessary function, but might help user to debug transformed code
# Add line numbers matching the lines of transformed code (the call stack will be more
# informative then if an error occures inside the transformed code)
function addLineNumbers(expr)
    # Remove previous line numbers and add correct ones instead
    result = Meta.parse(string(MacroTools.striplines(expr)))
    # fix LineNumberNodes to be more informative
    MacroTools.postwalk(result) do x
        x isa LineNumberNode ?
            LineNumberNode(x.line, Symbol("generated_code_by_recycle")) :
            x
    end
end

# Replace gensyms with their names and add line numbers in the beginning of each line
function prettify(expr)
    expr = MacroTools.postwalk(
        x -> MacroTools.isgensym(x) ? Symbol(MacroTools.gensymname(x)) : x, expr)
    txt = string(MacroTools.striplines(expr))
    lines = split(txt, "\n")
    lines_with_numbers = [@sprintf("\n%3d | %s", i, line) for (i, line) in enumerate(lines)]
    join(lines_with_numbers)
end

# This function basicly does the entire job by calling all functions above
function transform(expr_str, arrays, funops, numbers, scalings)
    array_symbols = copy(arrays)
    extracted_ops, buffers, modified, first_bools = Dict(), [], [], []
    # Optimize by replacements
    result = preprocess(expr_str, arrays, funops, numbers, scalings, extracted_ops, modified) |>
        x -> dotify(x, arrays, scalings, numbers) |>
        x -> buffer(x, arrays, array_symbols, funops, buffers) |>
        x -> postprocess(x, first_bools)
    # Add initialization of extracted operaors, buffers and "first bool"s
    init_🔝 = length(extracted_ops) > 0 ?
        [Expr(:(=), l, r) for (l,r) in sort(collect(extracted_ops), by=x->x[1])] : []
    init_🔃 = Expr(:(=), Expr(:tuple, buffers...), Expr(:call, :fill, :nothing, length(buffers)))
    init_first_bools =
            Expr(:(=), Expr(:tuple, first_bools...), Expr(:call, :fill, :true, length(first_bools)))
    result = length(buffers) > 0 ?
        Expr(:block, init_🔝..., init_🔃, init_first_bools, result) : Expr(:block, init_🔝..., result)
    # Self-explanatory...
    result = addLineNumbers(result)
    # Wrap it in a try-catch block to print the generated code before the error message is rethrown
    extended = Expr(:try, result, :e,
        Expr(:block,
            Expr(:call, :println, "\ngenerated_code_by_recycle:", prettify(result)),
            :(rethrow(e))))
    # return...
    result, extended, modified
end

macro string_varname_assign(str,val)
    s = Symbol(str)
    esc(:($s = $val))
end

"""
Speed up iteratively executed code fragments with many matrix operations by transforming code in such a way that preserves arrays allocated for intermediate results, and re-use them for subsequent iterations.

First variant:
```julia
@recycle <code to be optimized>
```
Second variant:
```julia
@recycle(arrays = [<list of array variables>], funops = [<list of funop variables], numbers = [<list of number variables>], <code to be optimized>)
```

The **first variant** is the more convenient one that tries to guess the type of variables (the other variant requires its user to declare explicitly the list of variables which are type of Array, FunOp, and Number. As a tradeoff, this variant fails when the optimized code contains either a closure or non-const global variable.

The **second variant** is the more flexible (and also more verbose) one one that requires its user to declare explicitly the list of variables which are type of Array, FunOp, and Number. The other variant tries to guess the type of variables, thus it is more convenient, but as a tradeoff, that variant fails when the optimized code contains either a closure or non-const global variable. On the other hand, this (more verbose) variant is free from these limitations.
*Note: All of the "keyword arguments" are optional, and also their order is arbitrary.*

An example to first variant:
This function
```julia
function foo()
    A = rand(100,100)
    B = rand(100,100)
    @recycle for i = 1:5
        A += A / 2 + B
        C = A * B + 5
    end
end
```
is turned into the following:
```julia
function foo()
    A = rand(100,100)
    B = rand(100,100)
    (C, 🔃₂) = fill(nothing, 2)
    (is_first_run₁,) = fill(true, 1)
    for i = 1:5
        A .+= A ./ 2 .+ B
        if is_first_run₁
            is_first_run₁ = false
            🔃₂ = A * B
            C = 🔃₂ .+ 5
        else
            mul!(🔃₂, A, B)
            C .= 🔃₂ .+ 5
        end
    end
end
```

Another example showing what second variant can do (and the first can't):
```julia
bar = @recycle(arrays=[A,B], (A, B) -> begin
    A += A / 2 + B
    B = A * B .+ 5
end)
function baz()
    A = rand(100,100)
    B = rand(100,100)
    for i = 1:5
        bar(A,B)
    end
end
```
is turned into the following:
```julia
bar = begin
    (🔃₁,) = fill(nothing, 1)
    (is_first_run₁,) = fill(true, 1)
    (A, B)->begin
            A .+= A ./ 2 .+ B
            begin
                if is_first_run₁
                    is_first_run₁ = false
                    🔃₁ = A * B
                else
                    mul!(🔃₁, A, B)
                end
                B .= 🔃₁ .+ 5
            end
        end
end
function baz()
    A = rand(100,100)
    B = rand(100,100)
    for i = 1:5
        bar(A,B)
    end
end
```
"""
macro recycle(expr)
    generated_func_name = Symbol("recycle_function_@",__source__.file, "_", __source__.line)
    vars = getSymbols(expr)
    !(:mul! in vars) && push!(vars, :mul!)
    build_array_of_array_varnames =
        [:($var <: AbstractArray && push!(arrays, Symbol($(string(var))))) for var in vars]
    build_array_of_FunOp_varnames =
        [:($var <: FunOp && push!(funops, Symbol($(string(var))))) for var in vars]
    build_array_of_Number_varnames =
        [:($var <: Number && push!(numbers, Symbol($(string(var))))) for var in vars]
    build_array_of_outer_vars =  # Variables defined outside of the optimized code
        [:(push!(variables, $var <: Nothing ? nothing : Symbol($(string(var))))) for var in vars]
    expr_str = string(expr)
    print_vars = [:(Core.println($(Expr(:., :Main, var)))) for var in vars]
    eval(quote
        @generated function $generated_func_name($(vars...))
            arrays, funops, numbers, variables = [], [], [], []
            # Get type of atoms (symbols)
            $(build_array_of_array_varnames...)
            $(build_array_of_FunOp_varnames...)
            $(build_array_of_Number_varnames...)
            $(build_array_of_outer_vars...)
            quote
                @recycle(arrays=[$((arrays)...)], funops=[$((funops)...)], numbers=[$((numbers)...)], variables=[$((variables)...)], $$expr_str)
            end
        end
    end)
    if_params = [Expr(:if, Expr(:isdefined, var), var, :nothing) for var in vars]
    var_names = map(var -> string(var), vars)
    new_values = gensym("new_values")
    assign_list = [:($var_name in keys($new_values) &&
        FunctionOperators.@string_varname_assign $var_name $new_values[$var_name])
            for var_name in var_names if !isconst(Main, Symbol(var_name))]
    esc(quote
        $new_values = FunctionOperators.$generated_func_name( $(if_params...) )
        $(assign_list...)
        $new_values["##@recycle return value##"]
    end)
end

macro recycle(args...)
    arrays, funops, numbers, variables = [], [], [], []
    expr = nothing
    for a in args
        if a isa Expr && a.head == :(=)
            @assert a isa Expr "All keywords arguments of @recycle must be an expression"
            if a.args[1] == :arrays
                @assert a.args[2].head == :vect "Parameter \"arrays\" must be a list of arrays variables"
                arrays = a.args[2].args
            elseif a.args[1] == :funops
                @assert a.args[2].head == :vect "Parameter \"funops\" must be a list of arrays variables"
                funops = a.args[2].args
            elseif a.args[1] == :numbers
                @assert a.args[2].head == :vect "Parameter \"numbers\" must be a list of arrays variables"
                numbers = a.args[2].args
            elseif a.args[1] == :variables
                @assert a.args[2].head == :vect "Parameter \"variables\" must be a list of arrays variables"
                variables = a.args[2].args
            else
                @assert false "Unrecognized keyword argument: $(a.args[1])"
            end
        else
            @assert a isa Expr || a isa String "@recycle must receive the expression to be optimized as an expression or a string"
            @assert expr == nothing "@recycle can accept only one positional argument (the expression to be optimized)"
            expr = a
        end
    end
    scalings = @isdefined(I) && I isa UniformScaling ? Array{Any}([:I]) : []
    expr_str = expr isa String ? expr : string(expr)
    result, extended, modified = transform(expr_str, arrays, funops, numbers, scalings)
    FunctionOperators_global_settings.macro_verbose && Core.println(prettify(result))
    if length(variables) > 0
        to_be_returned = intersect(modified, variables)
        return_pairs = [Expr(:call, :(=>), string(var), var) for var in to_be_returned]
        return_val = gensym("return_val")
        push!(return_pairs, Expr(:call, :(=>), "##@recycle return value##", return_val))
        return_dict = Expr(:call, :Dict, return_pairs...)
        extended.args[1] = Expr(:block, Expr(:(=), return_val, result), return_dict)
    end
    esc(extended)
end


text:
module TestBlocking

using ..TestCancellableContainers: Blocking, TreiberStack, taking
using Reagents
using Test

function test_unblock()
    s1 = Blocking(TreiberStack{Int}())
    a2b, b2a = Reagents.channel(Int, Int)
    t = @task (taking(s1) ⨟ a2b)()
    yield(t)
    put!(s1, 111)
    istaskdone(t) && wait(t)
    @test b2a(222) == 111
    @test fetch(t) == 222
end

function test_both()
    s1 = Blocking(TreiberStack{Int}())
    s2 = Blocking(TreiberStack{Int}())
    t = @task (taking(s1) & taking(s2))()
    yield(t)
    put!(s1, 111)
    put!(s2, 222)
    @test fetch(t) == (111, 222)
end

end  # module

ext: jl


text:
module Rivers

using Unitful
using Unitful: m, s, km, yr
import Unitful: uconvert
using Distances: Haversine
using NearestNeighbors: BruteTree, knn
using OceanGrids
import OceanGrids: regrid

# These river discharge volumetric flow rates and locations were taken from the Dai and Trenberth dataset
# located on UCAR's Research Data Archive (RDA)
# (URL: https://rda.ucar.edu/datasets/ds551.0/index.html#!description)
# and were converted to Julia/AIBECS with permission from Aiguo Dai.
#
# This module only contains the annual mean river discharge of the 200 largest rivers, taken specifically from files `runoff-table-A.txt` and `runoff-table2-top50r.txt` of the original dataset.
#
# Then I copied these files here and wrote them up into a Julia format by hand
# because it was the most straightforward approach for me here...

struct River{T}
    name::String
    lon::Float64
    lat::Float64
    VFR::T # volumetric flow rate
end

# Show method for rivers
function Base.show(io::IO, r::River)
    lat, lon, name, v = r.lat, r.lon, r.name, r.VFR
    print(io, name, " ", (lat,lon), " ", round(Int, ustrip(v)) * unit(v))
end
Base.show(io::IO, m::MIME"text/plain", r::River) = show(io, r)

const RIVERS = [
    River("Amazon",                          -55.5,  -2.0, 6642.0km^3/yr),
    River("Congo",                            15.3,  -4.3, 1308.0km^3/yr),
    River("Orinoco",                         -63.6,   8.1, 1129.0km^3/yr),
    River("Changjiang",                      117.6,  30.8,  944.0km^3/yr),
    River("Brahmaputra",                      89.7,  25.2,  628.0km^3/yr),
    River("Mississippi",                     -90.9,  32.3,  610.0km^3/yr),
    River("Yenisey",                          86.5,  67.4,  599.0km^3/yr),
    River("Paraná",                          -60.7, -32.7,  568.0km^3/yr),
    River("Lena",                            127.4,  70.7,  531.0km^3/yr),
    River("Mekong",                          105.8,  15.1,  525.0km^3/yr),
    River("Tocantins",                       -49.7,  -3.8,  511.0km^3/yr),
    River("Tapajos",                         -56.8,  -5.2,  415.0km^3/yr),
    River("Ob",                               66.6,  66.6,  412.0km^3/yr),
    River("Ganges",                           88.1,  24.5,  404.0km^3/yr),
    River("Irrawaddy",                        96.0,  21.9,  393.0km^3/yr),
    River("St Lawrence",                     -74.7,  45.0,  363.0km^3/yr),
    River("Amur",                            137.0,  50.5,  354.0km^3/yr),
    River("Xingu",                           -52.2,  -3.2,  302.0km^3/yr),
    River("Mackenzie",                      -133.7,  67.5,  290.0km^3/yr),
    River("Xijiang",                         111.3,  23.5,  270.0km^3/yr),
    River("Columbia",                       -121.2,  45.6,  252.0km^3/yr),
    River("Magdalena",                       -74.9,  10.2,  231.0km^3/yr),
    River("Uruguay",                         -58.0, -31.4,  228.0km^3/yr),
    River("Yukon",                          -162.9,  61.9,  212.0km^3/yr),
    River("Atrato",                          -76.7,   6.2,  204.0km^3/yr),
    River("Danube",                           28.7,  45.2,  202.0km^3/yr),
    River("Niger",                             3.5,  11.9,  193.0km^3/yr),
    River("Ogooué",                           10.2,  -0.7,  186.0km^3/yr),
    River("Essequibo",                       -58.6,   5.8,  154.0km^3/yr),
    River("Fraser",                         -121.4,  49.4,  144.0km^3/yr),
    River("Pechora",                          52.2,  67.6,  140.0km^3/yr),
    River("Nelson",                          -97.9,  54.8,  126.0km^3/yr),
    River("Khatanga",                        102.5,  72.0,  124.0km^3/yr),
    River("Sepik",                           142.2,  -4.2,  123.0km^3/yr),
    River("Kolyma",                          158.7,  68.7,  118.0km^3/yr),
    River("Zambeze",                          33.6, -16.1,  117.0km^3/yr),
    River("Severnaya Dvina",                  41.9,  64.1,  112.0km^3/yr),
    River("Indus",                            68.3,  25.4,  104.0km^3/yr),
    River("Sanaga",                           10.1,   3.8,   99.0km^3/yr),
    River("Godavari",                         81.8,  16.9,   97.0km^3/yr),
    River("Rajang",                          112.9,   2.0,   93.0km^3/yr),
    River("Sao Francisco",                   -37.0, -10.0,   90.0km^3/yr),
    River("Usumacinta",                      -91.5,  17.4,   89.0km^3/yr),
    River("Maroni",                          -54.5,   5.0,   86.0km^3/yr),
    River("Rhine",                             6.1,  51.8,   75.0km^3/yr),
    River("Purari",                          145.1,  -7.0,   74.0km^3/yr),
    River("Caniapiscau",                     -69.2,  57.4,   73.0km^3/yr),
    River("Mahanadi",                         84.0,  20.8,   73.0km^3/yr),
    River("Sacramento",                     -121.5,  38.6,   69.0km^3/yr),
    River("Jacui",                           -51.5, -30.0,   69.0km^3/yr),
    River("San Juan CO",                     -77.2,   4.2,   65.0km^3/yr),
    River("Bénoué",                           13.4,   9.3,   64.0km^3/yr),
    River("Kuskokwim",                      -158.1,  61.9,   57.0km^3/yr),
    River("Albany",                          -83.9,  51.3,   57.0km^3/yr),
    River("Huai",                            117.4,  32.9,   57.0km^3/yr),
    River("La Grande",                       -78.6,  53.7,   56.0km^3/yr),
    River("Tigris (Dicle)",                   44.4,  33.3,   56.0km^3/yr),
    River("Krishna",                          80.6,  16.5,   55.0km^3/yr),
    River("Ottawa",                          -76.2,  45.5,   55.0km^3/yr),
    River("Chindwin",                         95.7,  26.0,   55.0km^3/yr),
    River("Taz",                              82.3,  66.6,   55.0km^3/yr),
    River("Po",                               11.6,  44.9,   55.0km^3/yr),
    River("Courantyne",                      -57.7,   5.0,   54.0km^3/yr),
    River("Indigirka",                       147.5,  69.6,   54.0km^3/yr),
    River("Rhone",                             4.6,  43.8,   54.0km^3/yr),
    River("Saguenay",                        -71.6,  48.6,   53.0km^3/yr),
    River("Alabama",                         -87.5,  31.5,   51.0km^3/yr),
    River("Stikine",                        -132.1,  56.7,   51.0km^3/yr),
    River("Eastmain",                        -78.1,  52.2,   49.0km^3/yr),
    River("Brahmani",                         85.0,  21.5,   48.0km^3/yr),
    River("Dnepr",                            35.2,  47.9,   47.0km^3/yr),
    River("Huanghe",                         114.9,  35.2,   47.0km^3/yr),
    River("Jamanxim",                        -55.8,  -5.5,   47.0km^3/yr),
    River("Susquehanna",                     -76.9,  40.2,   46.0km^3/yr),
    River("Beijiang",                        112.9,  23.6,   46.0km^3/yr),
    River("Don",                              40.7,  47.5,   45.0km^3/yr),
    River("Susitna",                        -150.5,  61.5,   45.0km^3/yr),
    River("Narmada",                          73.7,  21.9,   44.0km^3/yr),
    River("Santa Cruz",                      -71.9, -50.3,   41.0km^3/yr),
    River("Oyapock (Oia)",                   -51.9,   3.8,   40.0km^3/yr),
    River("George",                          -65.8,  58.2,   40.0km^3/yr),
    River("Rufiji",                           37.9,  -7.8,   40.0km^3/yr),
    River("Nile",                             31.3,  29.7,   40.0km^3/yr),
    River("Nottaway",                        -77.4,  50.1,   39.0km^3/yr),
    River("Volta",                             0.1,   6.2,   37.0km^3/yr),
    River("Willamette",                     -123.0,  44.9,   37.0km^3/yr),
    River("Skeena",                         -128.4,  54.6,   36.0km^3/yr),
    River("Garonne",                           0.2,  44.4,   36.0km^3/yr),
    River("Red",                             -92.4,  31.3,   36.0km^3/yr),
    River("Pahang",                          102.4,   3.4,   35.0km^3/yr),
    River("Moose",                           -81.3,  50.8,   34.0km^3/yr),
    River("Chao Phraya",                     100.1,  15.3,   34.0km^3/yr),
    River("Cuyuni",                          -58.8,   6.4,   34.0km^3/yr),
    River("Vistula (Wisła)",                  18.8,  54.1,   34.0km^3/yr),
    River("Copper",                         -144.5,  61.5,   34.0km^3/yr),
    River("Saint John",                      -66.8,  46.0,   34.0km^3/yr),
    River("Olenek",                          123.7,  71.9,   34.0km^3/yr),
    River("Rupert",                          -76.9,  51.5,   33.0km^3/yr),
    River("Vuoksi",                           28.8,  61.2,   33.0km^3/yr),
    River("Kamchatka",                       161.6,  56.3,   33.0km^3/yr),
    River("Manicouagan",                     -68.3,  49.2,   33.0km^3/yr),
    River("Esmeraldas",                      -79.4,   0.5,   32.0km^3/yr),
    River("Doce",                            -40.1, -19.4,   32.0km^3/yr),
    River("Yana",                            136.1,  70.8,   32.0km^3/yr),
    River("Jari",                            -52.6,  -0.6,   32.0km^3/yr),
    River("Bío Bío",                         -73.1, -36.8,   32.0km^3/yr),
    River("Severn-CA",                       -88.3,  55.4,   32.0km^3/yr),
    River("Nushagak",                       -157.5,  59.3,   31.0km^3/yr),
    River("Tsiribihina",                      45.0, -19.7,   31.0km^3/yr),
    River("Baker",                           -72.9, -47.3,   30.0km^3/yr),
    River("Araguari",                        -51.4,   0.7,   30.0km^3/yr),
    River("Pur",                              78.2,  67.0,   29.0km^3/yr),
    River("Nass",                           -129.1,  55.2,   29.0km^3/yr),
    River("Douro (Duero)",                    -7.8,  41.2,   29.0km^3/yr),
    River("Negro-AR",                        -63.7, -40.5,   29.0km^3/yr),
    River("Kouilou",                          12.1,  -4.1,   29.0km^3/yr),
    River("Loire",                            -0.8,  47.4,   28.0km^3/yr),
    River("Puelo",                           -72.2, -41.5,   28.0km^3/yr),
    River("Dongjiang",                       114.3,  23.2,   27.0km^3/yr),
    River("Tista",                            89.5,  25.8,   27.0km^3/yr),
    River("Elbe (Labe)",                      11.8,  53.0,   27.0km^3/yr),
    River("Alsek",                          -138.1,  59.4,   27.0km^3/yr),
    River("Tombigbee",                       -87.9,  32.5,   27.0km^3/yr),
    River("Paraiba do Sul",                  -41.3, -21.8,   26.0km^3/yr),
    River("Parnaiba",                        -42.4,  -3.5,   26.0km^3/yr),
    River("Karun",                            48.7,  31.3,   25.0km^3/yr),
    River("Curua-2",                         -54.5,  -5.6,   25.0km^3/yr),
    River("Mezen",                            45.6,  65.0,   25.0km^3/yr),
    River("Kinabatangan",                    117.6,   5.3,   25.0km^3/yr),
    River("Saint-Maurice",                   -72.7,  46.6,   24.0km^3/yr),
    River("Nyanga",                           10.7,  -2.7,   23.0km^3/yr),
    River("Capim",                           -47.8,  -2.5,   23.0km^3/yr),
    River("Grijalva",                        -93.2,  18.0,   22.0km^3/yr),
    River("Senegal",                         -15.5,  16.5,   22.0km^3/yr),
    River("Papaloapan",                      -96.1,  18.2,   22.0km^3/yr),
    River("Euphrates (F",                     44.3,  32.7,   21.0km^3/yr),
    River("Glomma (Glåma)",                   11.1,  59.6,   21.0km^3/yr),
    River("Apalachicola",                    -84.9,  30.7,   20.0km^3/yr),
    River("Nadym",                            72.7,  65.6,   20.0km^3/yr),
    River("Odra (Oder)",                      14.3,  52.8,   20.0km^3/yr),
    River("Daugava (Západnaya Dviná)",        25.9,  56.5,   19.0km^3/yr),
    River("Hayes, MB",                       -92.8,  56.4,   19.0km^3/yr),
    River("Baleine (Great Whale River)",     -77.0,  55.2,   19.0km^3/yr),
    River("Mangoky",                          43.9, -21.8,   19.0km^3/yr),
    River("Feuilles (Rivière aux Feuilles)", -70.4,  58.6,   19.0km^3/yr),
    River("Melezes (Rivière aux Mélèzes)",   -69.6,  57.7,   19.0km^3/yr),
    River("Churchill",                       -94.6,  58.1,   18.0km^3/yr),
    River("Kelantan",                        102.2,   5.8,   18.0km^3/yr),
    River("Clutha",                          169.7, -46.2,   17.0km^3/yr),
    River("Cross",                             9.3,   5.8,   17.0km^3/yr),
    River("Kemi",                             24.7,  65.9,   17.0km^3/yr),
    River("Göta",                             12.4,  58.4,   17.0km^3/yr),
    River("Han",                             127.0,  37.5,   17.0km^3/yr),
    River("Back",                            -96.5,  66.1,   17.0km^3/yr),
    River("Onega",                            38.5,  63.8,   16.0km^3/yr),
    River("Tapi (Tapti)",                     72.9,  21.3,   16.0km^3/yr),
    River("Anabar",                          114.1,  72.0,   16.0km^3/yr),
    River("Jequitinhonh",                    -39.5, -15.9,   14.0km^3/yr),
    River("Kuban",                            38.2,  45.2,   14.0km^3/yr),
    River("Pánuco",                          -98.6,  22.0,   14.0km^3/yr),
    River("Limpopo",                          33.0, -24.5,   14.0km^3/yr),
    River("Winisk",                          -87.2,  54.5,   13.0km^3/yr),
    River("Ebro",                              0.5,  40.8,   13.0km^3/yr),
    River("Narva",                            28.2,  59.4,   12.0km^3/yr),
    River("Colorado-AR",                    -114.6,  32.7,   12.0km^3/yr),
    River("Seine",                             1.2,  49.3,   11.0km^3/yr),
    River("Coopermine",                     -115.4,  67.7,   11.0km^3/yr),
    River("Dnestr",                           29.5,  46.8,   11.0km^3/yr),
    River("Tagus (Tejo)",                     -8.4,  39.5,   11.0km^3/yr),
    River("Itapecuru",                       -44.4,  -3.6,    9.9km^3/yr),
    River("Burdekin",                        147.2, -19.8,   10.0km^3/yr),
    River("Murray",                          142.8, -34.6,    9.4km^3/yr),
    River("Bolshoy Anyuy",                   161.2,  68.2,    8.4km^3/yr),
    River("Grande de Santiago",             -105.1,  21.8,    8.4km^3/yr),
    River("Bandama",                          -4.8,   5.9,    8.3km^3/yr),
    River("Cauvery",                          78.8,  10.8,    7.7km^3/yr),
    River("Juba",                             42.5,   3.8,    7.5km^3/yr),
    River("Fitzroy-QX",                      150.1, -23.1,    7.4km^3/yr),
    River("Brazos",                          -95.8,  29.6,    7.1km^3/yr),
    River("Comoé (Komoé)",                    -3.7,   6.6,    6.8km^3/yr),
    River("Jaguaribe",                       -38.2,  -5.2,    6.6km^3/yr),
    River("Daly",                            130.7, -13.8,    6.3km^3/yr),
    River("Ouémé",                             2.5,   6.9,    5.4km^3/yr),
    River("Guadiana",                         -7.6,  37.8,    5.2km^3/yr),
    River("Anderson",                       -128.4,  68.6,    5.0km^3/yr),
    River("Victoria",                        130.9, -15.5,    4.9km^3/yr),
    River("Roper",                           134.4, -14.7,    4.8km^3/yr),
    River("Liao",                            123.5,  42.2,    4.6km^3/yr),
    River("Orange (Senqu)",                   17.6, -28.8,    4.6km^3/yr),
    River("Paraguacu",                       -39.0, -12.6,    4.3km^3/yr),
    River("Contas [de]",                     -39.3, -14.3,    3.4km^3/yr),
    River("Guadalquivir",                     -6.0,  37.5,    3.3km^3/yr),
    River("Yaqui",                          -109.5,  29.2,    3.3km^3/yr),
    River("Colorado-TX",                     -96.1,  29.3,    2.6km^3/yr),
    River("Rio Grande",                      -97.4,  25.9,    1.5km^3/yr),
    River("de Grey",                         119.2, -20.3,    1.4km^3/yr),
    River("Gascoyne",                        113.8, -24.8,    0.7km^3/yr),
    River("Avon",                            116.1, -31.8,    0.7km^3/yr),
    River("Fortescue",                       116.2, -21.3,    0.3km^3/yr),
    River("Murchison",                       114.5, -27.9,    0.2km^3/yr)
]

citation() = """
- For the dataset itself: Dai, A. 2017. Dai and Trenberth Global River Flow and Continental Discharge Dataset. Research Data Archive at the National Center for Atmospheric Research, Computational and Information Systems Laboratory. https://doi.org/10.5065/D6V69H1T.
- For the formal publication that describe the dataset: Dai, A., and K. E. Trenberth, 2002: Estimates of freshwater discharge from continents: Latitudinal and seasonal variations. J. Hydrometeorol., 3, 660–687.
- For descriptions of updates to the dataset:
    - Dai, A., T. Qian, K. E. Trenberth, and J. D Milliman, 2009: Changes in continental freshwater discharge from 1948-2004. J. Climate, 22, 2773–2791.
    - Dai, A., 2016: Historical and future changes in streamflow and continental runoff: A review, in Terrestrial Water Cycle and Climate Change: Natural and Human-Induced Impacts, Geophysical Monograph 221. Ed. Qiuhong Tang and Taikan Oki, AGU, John Wiley & Sons, 17–37 (DOI: 10.1002/9781118971772).
"""

uconvert(u, r::River) = River(r.name, r.lon, r.lat, uconvert(u, r.VFR))

function load(unit=m^3/s)
    @info """You are about to use the Dai and Trenberth river discharge dataset.
          If you use it for research, please cite:

          $(citation())

          You can find the corresponding BibTeX entries in the CITATION.bib file
          at the root of the AIBECS.jl package repository.
          (Look for the "Dai_2017", "Dai_Trenberth_2002", "Dai_etal_2009", and "Dai_2016" keys.)
          """
    return uconvert.(unit, RIVERS)
end

function regrid(R::Vector{River{T}}, grd) where T <: Quantity
    lats = [r.lat for r in R]
    lons = [r.lon for r in R]
    depths = zeros(length(R))
    vs = [r.VFR for r in R]
    return regrid(vs, lats, lons, depths, grd)
end
export regrid

end # module

export Rivers


text:
using AxisKeys, NamedTupleTools
using PrettyTables, ParetoSmooth
using StanSample
using StatisticalRethinking
using StatisticalRethinkingPlots
using Test

df = CSV.read(sr_datadir("WaffleDivorce.csv"), DataFrame);
scale!(df, [:Marriage, :MedianAgeMarriage, :Divorce])
data = (N=size(df, 1), D=df.Divorce_s, A=df.MedianAgeMarriage_s,
    M=df.Marriage_s)

stan5_1 = "
data {
    int < lower = 1 > N; // Sample size
    vector[N] D; // Outcome
    vector[N] A; // Predictor
}
parameters {
    real a; // Intercept
    real bA; // Slope (regression coefficients)
    real < lower = 0 > sigma;    // Error SD
}
transformed parameters {
    vector[N] mu;               // mu is a vector
    for (i in 1:N)
        mu[i] = a + bA * A[i];
}
model {
    a ~ normal(0, 0.2);         //Priors
    bA ~ normal(0, 0.5);
    sigma ~ exponential(1);
    D ~ normal(mu , sigma);     // Likelihood
}
generated quantities {
    vector[N] log_lik;
    for (i in 1:N)
        log_lik[i] = normal_lpdf(D[i] | mu[i], sigma);
}
";

stan5_2 = "
data {
    int N;
    vector[N] D;
    vector[N] M;
}
parameters {
    real a;
    real bM;
    real<lower=0> sigma;
}
transformed parameters {
    vector[N] mu;
    for (i in 1:N)
        mu[i]= a + bM * M[i];

}
model {
    a ~ normal( 0 , 0.2 );
    bM ~ normal( 0 , 0.5 );
    sigma ~ exponential( 1 );
    D ~ normal( mu , sigma );
}
generated quantities {
    vector[N] log_lik;
    for (i in 1:N)
        log_lik[i] = normal_lpdf(D[i] | mu[i], sigma);
}
";

stan5_3 = "
data {
  int N;
  vector[N] D;
  vector[N] M;
  vector[N] A;
}
parameters {
  real a;
  real bA;
  real bM;
  real<lower=0> sigma;
}
transformed parameters {
    vector[N] mu;
    for (i in 1:N)
        mu[i] = a + bA * A[i] + bM * M[i];
}
model {
  a ~ normal( 0 , 0.2 );
  bA ~ normal( 0 , 0.5 );
  bM ~ normal( 0 , 0.5 );
  sigma ~ exponential( 1 );
  D ~ normal( mu , sigma );
}
generated quantities{
    vector[N] log_lik;
    for (i in 1:N)
        log_lik[i] = normal_lpdf(D[i] | mu[i], sigma);
}
";

m5_1s = SampleModel("m5.1s", stan5_1)
rc5_1s = stan_sample(m5_1s; data)

m5_2s = SampleModel("m5.2s", stan5_2)
rc5_2s = stan_sample(m5_2s; data)

m5_3s = SampleModel("m5.3s", stan5_3)
rc5_3s = stan_sample(m5_3s; data)

struct LooCompare1
    psis::Vector{PsisLoo}
    table::KeyedArray
end

function to_paretosmooth(ll, pd = [3,1, 2])
    permutedims(ll, [3, 1, 2])
end

function loo_compare1(models::Vector{SampleModel};
    loglikelihood_name="log_lik",
    model_names=nothing,
    sort_models=true,
    show_psis=true)

    nmodels = length(models)
    mnames = [models[i].name for i in 1:nmodels]

    chains_vec = read_samples.(models)
    ll_vec = Array.(matrix.(chains_vec, loglikelihood_name))
    ll_vecp = map(to_paretosmooth, ll_vec)
    psis_vec = psis_loo.(ll_vecp)

    if show_psis
        for i in 1:nmodels
            psis_vec[i] |> display
        end
    end

    psis_values = Vector{Float64}(undef, nmodels)
    se_values = Vector{Float64}(undef, nmodels)
    loos = Vector{Vector{Float64}}(undef, nmodels)

    for i in 1:nmodels
        psis_values[i] = psis_vec[i].estimates(:cv_elpd, :total)
        se_values[i] = psis_vec[i].estimates(:cv_elpd, :se_total)
        loos[i] = psis_vec[i].pointwise(:cv_elpd)
    end

    if sort_models
        ind = sortperm([psis_values[i][1] for i in 1:nmodels]; rev=true)
        psis_vec = psis_vec[ind]
        psis_values = psis_values[ind]
        se_values = se_values[ind]
        loos = loos[ind]
        mnames = mnames[ind]
    end

    # Setup comparison vectors

    elpd_diff = zeros(nmodels)
    se_diff = zeros(nmodels)
    weight = ones(nmodels)

    # Compute comparison values

    for i in 2:nmodels
        elpd_diff[i] = psis_values[i] - psis_values[1]
        diff = loos[1] - loos[i]
        se_diff[i] = √(length(loos[i]) * var(diff; corrected=false))
    end
    data = elpd_diff
    data = hcat(data, se_diff)

    sumval = sum([exp(psis_values[i]) for i in 1:nmodels])
    @. weight = exp(psis_values) / sumval
    data = hcat(data, weight)

    # Create KeyedArray object

    table = KeyedArray(
        data,
        model = mnames,
        statistic = [:cv_elpd, :se_diff, :weight],
    )

    # Return LooCompare object

    LooCompare1(psis_vec, table)

end

function Base.show(io::IO, ::MIME"text/plain", loo_compare::LooCompare1)
    table = loo_compare.table
    return pretty_table(
        table;
        compact_printing=false,
        header=table.statistic,
        row_names=table.model,
        formatters=ft_printf("%5.2f"),
        alignment=:r,
    )
end

if success(rc5_1s) && success(rc5_2s) && success(rc5_3s)

    nt5_1s = read_samples(m5_1s, :particles)
    NamedTupleTools.select(nt5_1s, (:a, :bA, :sigma)) |> display
    nt5_2s = read_samples(m5_2s, :particles)
    NamedTupleTools.select(nt5_2s, (:a, :bM, :sigma)) |> display
    nt5_3s = read_samples(m5_3s, :particles)
    NamedTupleTools.select(nt5_3s, (:a, :bA, :bM, :sigma)) |> display
    println()

    models = [m5_1s, m5_2s, m5_3s]
    for i in 1:length(models)
        pngfname = joinpath(@__DIR__, "m5.$(i)s.png")
        isfile(pngfname) && rm(pngfname)
    end

    loo_comparison = loo_compare1(models)
    println()

    if isinteractive()

        for i in 1:length(models)
            pw = loo_comparison.psis[i].pointwise
            pk_plot(pw(:pareto_k))
            savefig(joinpath(@__DIR__, "m5.$(i)s.png"))
        end
    end

    loo_comparison |> display
end
#=
With SR/ulam():
```
       PSIS    SE dPSIS  dSE pPSIS weight
m5.1u 126.0 12.83   0.0   NA   3.7   0.67
m5.3u 127.4 12.75   1.4 0.75   4.7   0.33
m5.2u 139.5  9.95  13.6 9.33   3.0   0.00
```
=#


text:
export RegionGraph, random_region_graph, region_graph_2_pc
import Random: shuffle

"Root of region graph node hierarchy"
abstract type RegionGraph <: Tree end

######
## Partition
#####
mutable struct Partition
    partition::AbstractVector{T} where T <: RegionGraph
    Partition(n::AbstractVector{T}) where T <: RegionGraph = begin
        new(n)
    end
end

import Base: getindex, size, length
@inline Base.size(n::Partition) = size(n.partition)
@inline Base.length(n::Partition) = length(n.partition)
@inline Base.getindex(n::Partition, i::Int) = n.partition[i]

mutable struct RegionGraphInnerNode <: RegionGraph
    partitions::AbstractVector{Partition}
    variables::BitSet
    parent::Union{Nothing, RegionGraphInnerNode}

    RegionGraphInnerNode(partitions::AbstractVector{Partition}) = begin
        for partition in partitions
            for i = 1:length(partition), j = i+1 : length(partition)
                @assert isdisjoint(variables(partition[i]), variables(partition[j]))
            end
        end
        scope = variables(partitions[1][1])
        for ind = 2:size(partitions[1])[1]
            scope = scope ∪ variables(partitions[1][ind])
        end

        this = new(partitions, scope, nothing)
        for partition in partitions
            for i = 1:size(partition)[1]
                @assert isnothing(partition[i].parent)
                partition[i].parent = this
            end
        end
        this
    end
end

mutable struct RegionGraphLeafNode <: RegionGraph
    variables::BitSet
    parent::Union{Nothing, RegionGraphInnerNode}
    RegionGraphLeafNode(v) = new(v, nothing)
end

#####################
# Constructors
#####################

RegionGraph(v::AbstractVector{Var}) = RegionGraphLeafNode(BitSet(v))
RegionGraph(partitions::AbstractVector{Partition}) = RegionGraphInnerNode(partitions)

#####################
# Traits
#####################

import LogicCircuits.NodeType
@inline LogicCircuits.NodeType(::RegionGraphInnerNode) = Inner()
@inline LogicCircuits.NodeType(::RegionGraphLeafNode) = Leaf()

###################
## Methods
#################

import LogicCircuits.Utils: children, variables

@inline children(n::RegionGraphInnerNode) = n.partitions
@inline variables(n::RegionGraphInnerNode) = n.variables
@inline variables(n::RegionGraphLeafNode) = n.variables

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

"""
    random_region_graph(X::AbstractVector{Var}, depth::Int = 5, replicas::Int = 2, num_splits::Int = 2)

- `X`: Vector of all variables to include; for the root region
- `depth`: how many layers to do splits
- `replicas`: number of replicas or paritions (replicas only used for the root region; for other regions only 1 parition (inner nodes), or 0 parition for leaves)
- `num_splits`: number of splits for each parition; split variables into random equaly sized regions
"""
function random_region_graph(X::AbstractVector{Var};
    depth::Int = 5, replicas::Int = 2, num_splits::Int = 2)::RegionGraph
    partitions = Vector{Partition}()
    for repeat = 1 : replicas
        cur_rg = split_rg(X, depth; num_splits=num_splits, return_partition_only = true)
        push!(partitions, cur_rg)
    end

    # Each parition should include the same set of variables
    prev_scope = nothing
    for cur_partition in partitions
        cur_scope = variables(cur_partition[1])
        for i = 2:length(cur_partition)
            cur_scope = cur_scope ∪ variables(cur_partition[i])
        end

        if !isnothing(prev_scope)
            @assert prev_scope == cur_scope "All paritions should include the same set of variables."
        else
            prev_scope = cur_scope
        end
    end

    RegionGraphInnerNode(partitions)
end

function split_rg(variables::AbstractVector{Var}, depth::Int; num_splits::Int = 2, return_partition_only = false)
    if length(variables) < 2 || depth == 0
        # Cannot/shoud not split anymore
        return RegionGraph(variables)
    end

    shuffle_variables = shuffle(variables)

    partition(x, n) = begin
        ## TODO; might not work if num_splits > 2
        sz = ceil(Int, (length(x) / n))
        [ x[i:min(i+sz-1, length(x))] for i = 1:sz:length(x) ]
    end
    splits = partition(shuffle_variables, num_splits)

    answer = Vector{RegionGraph}()
    for split in splits
        push!(answer, split_rg(split, depth-1; num_splits = num_splits))
    end

    if return_partition_only
        Partition(answer)
    else
        RegionGraphInnerNode([Partition(answer)])
    end
end

########################################
### Rat-SPNs in Juice PC data structure
########################################

"""

- num_nodes_root: number of sum nodes in the root region
- num_nodes_leaf: number of sum nodes per leaf region
- num_nodes_region: number of in each region except root and leaves
"""
function region_graph_2_pc(node::RegionGraph;
    num_nodes_root::Int = 4,  num_nodes_region::Int= 3, num_nodes_leaf::Int = 2)::Vector{ProbCircuit}
    sum_nodes = Vector{ProbCircuit}()

    if isleaf(node)
        for i = 1:num_nodes_leaf
            # reIndex_bijection = [(i, v) for (i,v) in enumerate(variables(node))]
            # leaf_circuit = fully_factorized_circuit(ProbCircuit, length(variables(node)); reIndex_bijection)

            vtree = Vtree([PlainVtreeLeafNode(x) for x ∈ variables(node)], :balanced)
            leaf_circuit = PlainProbCircuit(fully_factorized_circuit(ProbCircuit, vtree))
            push!(sum_nodes, leaf_circuit)
        end

    else
        mul_nodes = Vector{ProbCircuit}()

        # Foreach replication; usually only > 1 at root
        for partition in node.partitions
            @assert length(partition) == 2 "Only supporting partitions of size 2 at the moment"
            lefts = region_graph_2_pc(partition[1]; num_nodes_root, num_nodes_region, num_nodes_leaf)
            rights = region_graph_2_pc(partition[2]; num_nodes_root, num_nodes_region, num_nodes_leaf)
            @assert all([issum(l) for l in lefts])
            @assert all([issum(r) for r in rights])

            for l in lefts, r in rights
                mul_node = PlainMulNode([l, r])
                push!(mul_nodes, mul_node)
            end
        end

        # Repeat Sum nodes nodes based on where in Region Graph
        if isnothing(node.parent)
            # Root region
            for i = 1:num_nodes_root
                sum_node = PlainSumNode(mul_nodes)
                push!(sum_nodes, sum_node)
            end
        else
            # Inner region
            for i = 1:num_nodes_region
                sum_node = PlainSumNode(mul_nodes)
                push!(sum_nodes, sum_node)
            end
        end
    end

    sum_nodes
end


text:
# AUTO GENERATED FILE - DO NOT EDIT

export vtk_dataarray

"""
    vtk_dataarray(;kwargs...)

A DataArray component.
DataArray is creating a vtkDataArray for the container fields
It takes the following set of properties:
  - type: 'Float32Array', 'Float64Array', 'Uint16Array', ...
  - values: [number, number, ...]
  - numberOfComponents: 1,
  - registration: 'addArray', 'setScalars', ...
Keyword arguments:
- `id` (String; optional): The ID used to identify this component.
- `type` (String; optional): Typed array name
- `name` (String; optional): Field name
- `values` (Array of Reals; optional): Actual values to use inside our array
- `numberOfComponents` (Real; optional): Number of components / Tuple size
- `registration` (String; optional): Name of the method to call on the fieldData (addArray, setScalars, setVectors...)
"""
function vtk_dataarray(; kwargs...)
        available_props = Symbol[:id, :type, :name, :values, :numberOfComponents, :registration]
        wild_props = Symbol[]
        return Component("vtk_dataarray", "DataArray", "dash_vtk", available_props, wild_props; kwargs...)
end

max_stars_repo_name: ericbdaniels/dash-vtk

max_issues_repo_name: ericbdaniels/dash-vtk

max_forks_repo_name: ericbdaniels/dash-vtk


text:
#!/bin/bash
# -*- mode: julia -*-
#=
JULIA="${JULIA:-julia}"
JULIA_CMD="${JULIA_CMD:-$JULIA --color=yes --startup-file=no --compile=min}"
exec $JULIA_CMD "${BASH_SOURCE[0]}" "$@"
=#

generate_mergify(path::AbstractString) = open(generate_mergify, path, write = true)
function generate_mergify(io::IO)
    header = """
    # [NOTE] This setting relies on "required status check":
    # https://docs.github.com/en/github/administering-a-repository/about-required-status-checks

    pull_request_rules:
      - name: remove outdated reviews
        conditions:
          - base=master
        actions:
          dismiss_reviews: {}
    """
    print(io, header)

    for (method_type, method_name) in [
            # https://doc.mergify.io/actions.html#merge
            ("squash", "squash-merge"),
            ("rebase", "rebase-merge"),
            ("merge", "merge"),
        ],
        (condition, name_suffix) in [
            # https://doc.mergify.io/conditions.html
            ("\"#approved-reviews-by>=1\"", ""),
            ("author=tkf", " (@tkf)"),
        ]

        rule = """
        - name: automatic $(method_name) when CI passes$(name_suffix)
          conditions:
            - base=master
            - $(condition)
            - label=ready-to-merge:$(method_type)
            - label!=work-in-progress
          actions:
            merge:
              method: $(method_type)
        """

        for ln in eachline(IOBuffer(rule))
            println(io, "  ", ln)
        end
    end
end

if abspath(PROGRAM_FILE) == @__FILE__
    generate_mergify(joinpath(@__DIR__, ".mergify.yml"))
end

generate_mergify


text:
using ..Test

using HTTP.Messages
import HTTP.Messages.appendheader
import HTTP.URI
import HTTP.request
import HTTP: bytes

using HTTP: StatusError

using HTTP.MessageRequest: bodylength
using HTTP.MessageRequest: bodybytes
using HTTP.MessageRequest: unknown_length

using JSON

@testset "HTTP.Messages" begin
    req = Request("GET", "/foo", ["Foo" => "Bar"])
    res = Response(200, ["Content-Length" => "5"]; body="Hello", request=req)

    protocols = ["http", "https"]
    http_reads = ["GET", "HEAD", "OPTIONS"]
    http_writes = ["POST", "PUT", "DELETE", "PATCH"]

    @testset "Body Length" begin
        @test bodylength(7) == unknown_length
        @test bodylength(UInt8[1,2,3]) == 3
        @test bodylength(view(UInt8[1,2,3], 1:2)) == 2
        @test bodylength("Hello") == 5
        @test bodylength(SubString("World!",1,5)) == 5
        @test bodylength(["Hello", " ", "World!"]) == 12
        @test bodylength(["Hello", " ", SubString("World!",1,5)]) == 11
        @test bodylength([SubString("Hello", 1,5), " ", SubString("World!",1,5)]) == 11
        @test bodylength([UInt8[1,2,3], UInt8[4,5,6]]) == 6
        @test bodylength([UInt8[1,2,3], view(UInt8[4,5,6],1:2)]) == 5
        @test bodylength([view(UInt8[1,2,3],1:2), view(UInt8[4,5,6],1:2)]) == 4
        @test bodylength(IOBuffer("foo")) == 3
        @test bodylength([IOBuffer("foo"), IOBuffer("bar")]) == 6
    end

    @testset "Body Bytes" begin
        @test bodybytes(7) == UInt8[]
        @test bodybytes(UInt8[1,2,3]) == UInt8[1,2,3]
        @test bodybytes(view(UInt8[1,2,3], 1:2)) == UInt8[1,2]
        @test bodybytes("Hello") == bytes("Hello")
        @test bodybytes(SubString("World!",1,5)) == bytes("World")
        @test bodybytes(["Hello", " ", "World!"]) == UInt8[]
        @test bodybytes([UInt8[1,2,3], UInt8[4,5,6]]) == UInt8[]
    end

    @testset "Request" begin
        @test req.method == "GET"
        @test res.request.method == "GET"

        @test String(req) == "GET /foo HTTP/1.1\r\nFoo: Bar\r\n\r\n"
        @test String(res) == "HTTP/1.1 200 OK\r\nContent-Length: 5\r\n\r\nHello"

        @test header(req, "Foo") == "Bar"
        @test header(res, "Content-Length") == "5"

        setheader(req, "X" => "Y")
        @test header(req, "X") == "Y"
    end

    @testset "Header Append" begin
        append_header(m, h) = appendheader(m, SubString(h[1]) => SubString(h[2]))

        append_header(req, "X" => "Z")
        @test header(req, "X") == "Y, Z"
        @test hasheader(req, "X", "Y, Z")
        @test headercontains(req, "X", "Y")
        @test headercontains(req, "X", "Z")
        @test !headercontains(req, "X", "more")

        append_header(req, "X" => "more")
        @test header(req, "X") == "Y, Z, more"
        @test hasheader(req, "X", "Y, Z, more")
        @test headercontains(req, "X", "Y")
        @test headercontains(req, "X", "Z")
        @test headercontains(req, "X", "more")

        append_header(req, "Set-Cookie" => "A")
        append_header(req, "Set-Cookie" => "B")
        @test filter(x->first(x) == "Set-Cookie", req.headers) == ["Set-Cookie" => "A", "Set-Cookie" => "B"]
    end

    @testset "HTTP Version" begin
        @test Messages.httpversion(req) == "HTTP/1.1"
        @test Messages.httpversion(res) == "HTTP/1.1"

        raw = String(req)
        req = parse(Request,raw)
        @test String(req) == raw

        req = parse(Request, raw * "xxx")
        @test String(req) == raw

        raw = String(res)
        res = parse(Response,raw)
        @test String(res) == raw

        res = parse(Response,raw * "xxx")
        @test String(res) == raw
    end

    @testset "Read methods" for protocol in protocols, method in http_reads
        @test request(method, "$protocol://httpbin.org/ip", verbose=1).status == 200
    end

    @testset "Body - Response Stream" for protocol in protocols, method in http_writes
        uri = "$protocol://httpbin.org/$(lowercase(method))"
        r = request(method, uri, verbose=1)
        @test r.status == 200
        r1 = JSON.parse(String(r.body))

        io = Base.BufferStream()
        r = request(method, uri, response_stream=io, verbose=1)
        @test r.status == 200
        r2 = JSON.parse(IOBuffer(read(io)))
        for (k, v) in r1
            if k == "headers"
                for (k2, v2) in r1[k]
                    if k2 != "X-Amzn-Trace-Id"
                        @test r1[k][k2] == r2[k][k2]
                    end
                end
            else
                @test r1[k] == r2[k]
            end
        end
    end

    @testset "Body - JSON Parse" for protocol in protocols, method in http_writes
        uri = "$protocol://httpbin.org/$(lowercase(method))"
        io = Base.BufferStream()
        r = request(method, uri, response_stream=io, verbose=1)
        @test r.status == 200

        r = request("POST",
            "$protocol://httpbin.org/post",
            ["Expect" => "100-continue"],
            "Hello",
            verbose=1)

        @test r.status == 200
        r = JSON.parse(String(r.body))
        @test r["data"] == "Hello"
    end

    @testset "Write to file" begin
        cd(mktempdir()) do

            line_count = 0
            num_lines = 50

            open("result_file", "w") do io
                r = request("GET",
                    "http://httpbin.org/stream/$num_lines",
                    response_stream=io,
                    verbose=1)
            end

            for line in readlines("result_file")
                line_parsed = JSON.parse(line)
                @test line_count == line_parsed["id"]
                line_count += 1
            end

            @test line_count == num_lines
        end
    end

    @testset "Display" begin
        @test repr(Response(200, []; body="Hello world.")) == "Response:\n\"\"\"\nHTTP/1.1 200 OK\r\n\r\nHello world.\"\"\""

        # truncation of long bodies
        for body_show_max in (Messages.body_show_max, 100)
            Messages.set_show_max(body_show_max)
            @test repr(Response(200, []; body="Hello world.\n"*'x'^10000)) == "Response:\n\"\"\"\nHTTP/1.1 200 OK\r\n\r\nHello world.\n"*'x'^(body_show_max-13)*"\n⋮\n10013-byte body\n\"\"\""
        end

        # don't display raw binary (non-Unicode) data:
        @test repr(Response(200, []; body=String([0xde,0xad,0xc1,0x71,0x1c]))) == "Response:\n\"\"\"\nHTTP/1.1 200 OK\r\n\r\n\n⋮\n5-byte body\n\"\"\""
    end
end


text:
using KUparser, KUnet, HDF5, JLD
KUnet.atype(Array)
@date @load "zn11oparse.jld"
for y in (:ytrn, :ydev, :ytst); @eval $y = full($y); end
@show net = Layer[Perceptron(size(ydev,1))]
for i=1:100
    @date train(net, xtrn, ytrn; shuffle=true)
    @date println((i,
                   accuracy(ytrn, predict(net, xtrn)),
                   accuracy(ydev, predict(net, xdev)),
                   accuracy(ytst, predict(net, xtst)),
                   ))
end
:ok


text:
# enumerated types
export GridType, Sign
export NP, nN, nP, nNP
export PD, nPR, nDL, nPD

# functions
export gt_w, numel, alter

# Signs
const nN, nP = 1, 2  # negative, positive directions
const nNP = SVec(nN,nP)  # tuple allocation is more efficient than array allocation
@enum Sign NEG=nN POS
const NP = SVec(NEG, POS)
for ins in instances(Sign); @eval export $(Symbol(ins)); end  # export all instances
Base.string(ins::Sign) = ins==NEG ? "negative" : "positive"
alter(ins::Sign) = ins==NEG ? POS : NEG

# Grid types
const nPR, nDL = 1, 2  # primal, dual grids
const nPD = SVec(nPR,nDL)  # tuple allocation is more efficient than array allocation
@enum GridType PRIM=nPR DUAL
const PD = SVec(PRIM, DUAL)
for ins in instances(GridType); @eval export $(Symbol(ins)); end  # export all instances
Base.string(ins::GridType) = ins==PRIM ? "primal" : "dual"
alter(ins::GridType) = ins==PRIM ? DUAL : PRIM

# Return grid types of the w-component of the field F, given the grid types gt₀ of the
# corners of the voxels whose edges are the field lines.
# If w is not one of the valid axes, return the voxel corner grid types gt₀.
gt_w(nw::Int,  # component index of field
     gt₀::SVec{K,GridType}  # grid type of corners of voxel whose edges are field lines
     ) where {K} =
    broadcast((a,w,g)->(a==w ? alter(g) : g), SVec(ntuple(identity, Val(K))), nw, gt₀)  # no change from gt₀ for nw ∉ {1,...,K}

# Functions for enumerated types
numel(::Type{T}) where {T<:Enum} = length(instances(T))
alter(n::Int) = 3-n


text:
ROS.@include std_srvs: SetBool

function testServiceClient()
    ROS.init("testSC")
    nh = ROS.NodeHandle()

    srvc = ROS.serviceClient(nh, "test_srv", ROS.std_srvs_SetBool)

    @test ROS.getService(srvc) == "/test_srv"
    @test ROS.isValid(srvc)

    loop = 100
    called_success = false
    srv = ROS.std_srvs_SetBool()
    while loop > 0 && !called_success
        println("Service client at loop: $loop")
        if ROS.exists(srvc)
            called_success = ROS.call(srvc, srv)
        end
        loop -= 1
        ROS.sleep(ROS.Rate(10))
        ROS.spinOnce()
    end
    @test called_success
    @test !ROS.isPersistent(srvc)

    ROS.shutdown(nh)

    println("All $(basename(@__FILE__)) tests passed.")
end


text:
#### ExportFuncs

# Provides a set of export functions that operate on StateVec.

using Printf, DelimitedFiles
export export_state, import_state!

include("ExportFuncsPrimitive.jl")

"""
    export_state(sv::StateVec, grid::Grid, dir, filename)

Export StateVec to a human-readable file `filename` in directory `dir`.
"""
function export_state(sv::StateVec, grid::Grid, dir, filename)
  domain = over_elems(grid)
  vn = var_names(sv)
  FT = eltype(grid)

  headers = ["z", [var_string(sv, ϕ, i) for ϕ in vn for i in over_sub_domains(sv, ϕ)]...]

  data = zeros(FT,length(domain), length(headers)-1)
  for ϕ in vn, i in over_sub_domains(sv, ϕ), k in domain
    i_eff = get_i_state_vec(sv.var_mapper, sv.a_map[ϕ], ϕ, i)
    data[k, i_eff] = sv[ϕ, k, i]
  end

  z = grid.zc[domain]
  data_all = hcat(z, data)
  mkpath(dir)

  file = joinpath(dir, filename)
  max_len = 18 # max_len = max(length.(headers)...) # Must match fmt in export_data
  headers = [repeat(" ", max_len-length(x))*x for x in headers]
  export_data(data_all, file, headers)
end

"""
    import_state!(sv::StateVec, grid::Grid, dir, filename)

Import StateVec from a human-readable file `filename` in directory `dir`.
"""
function import_state!(sv::StateVec, grid::Grid, dir, filename)
  domain = over_elems(grid)
  vn = var_names(sv)
  FT = eltype(grid)

  headers = ["z", [var_string(sv, ϕ, i) for ϕ in vn for i in over_sub_domains(sv, ϕ)]...]
  data_all = zeros(FT, length(domain), length(headers))
  mkpath(dir)

  file = joinpath(dir, filename)
  max_len = 18 # max_len = max(length.(headers)...) # Must match fmt in export_data
  headers = [repeat(" ", max_len-length(x))*x for x in headers]
  import_data!(data_all, file, headers)

  data = data_all[:,2:end] # remove z
  for ϕ in vn, i in over_sub_domains(sv, ϕ), k in domain
    i_eff = get_i_state_vec(sv.var_mapper, sv.a_map[ϕ], ϕ, i)
    sv[ϕ, k, i] = data[k, i_eff]
  end

end

max_stars_repo_name: akshaysridhar/CLIMA

max_issues_repo_name: akshaysridhar/CLIMA

max_forks_repo_name: akshaysridhar/CLIMA


text:
function row_row_prime(dd, i, type::Type = eltype(dd[1,1]))
    r = Array{type}(dd[i,:])
    return Hermitian(r * r')
end
function b_bar(dd,S)
    etype = eltype(dd[1,1])
    N = nrow(dd)
    total = 0.0
    for k in 1:N
        total += sqrt(squared_frobenius_distance(row_row_prime.(Ref(dd), k, etype), S))
    end
    return total / (N^2)
end

"""
    identity_regularisation(mat::Hermitian, identity_weight::Real)

Regularisation of the correlation matrix by mixing with the identity matrix.
### Inputs
* `mat` - A matrix to be regularised.
* `identity_weight` - How much weight to give to the identity matrix. Should be between 0 and 1.
### Returns
* A `Hermitian`.

    identity_regularisation(mat::Hermitian, asset_returns::DataFrame) where R<:Real

Regularisation of the correlation matrix by mixing with the identity matrix as per Ledoit & Wolf 2003.
### Inputs
* `mat` - A matrix to be regularised.
* `ts` - Tick data.
### Returns
* A `Hermitian`.

    identity_regularisation(mat::Hermitian, ts::SortedDataFrame,  mat_labels::Vector;
                            spacing::Union{Missing,<:Real} = missing)

Regularisation of the correlation matrix by mixing with the identity matrix as per Ledoit & Wolf 2003.
### Inputs
* `mat` - A matrix to be regularised.
* `ts` - Tick data.
* `mat_labels` - The labels for each asset in the matrix.
* `spacing` A spacing to use to estimate returns. This is used in determining the optimal weight to give to the identity matrix.
### Returns
* A `Hermitian`.

    identity_regularisation(covariance_matrix::CovarianceMatrix, ts::SortedDataFrame;
                            spacing::Union{Missing,<:Real} = missing, apply_to_covariance::Bool = true)

Regularisation of the correlation matrix by mixing with the identity matrix as per Ledoit & Wolf 2003.
### Inputs
* `covariance_matrix` - The `CovarianceMatrix` to be regularised.
* `ts` - Tick data.
* `spacing` A spacing to use to estimate returns. This is used in determining the optimal weight to give to the identity matrix.
* `apply_to_covariance` Should regularisation be applied to the covariance matrix or the correlation matrix.
### Returns
* A `CovarianceMatrix`.

    identity_regularisation(covariance_matrix::CovarianceMatrix, identity_weight::Real;
                            apply_to_covariance = false)

Regularisation of the correlation matrix by mixing with the identity matrix.
### Inputs
* `covariance_matrix` - The `CovarianceMatrix` to be regularised.
* `identity_weight` - How much weight to give to the identity matrix. Should be between 0 and 1.
* `apply_to_covariance` Should regularisation be applied to the covariance matrix or the correlation matrix.
### Returns
* A `CovarianceMatrix`.

### References
Ledoit, O. , Wolf, M. 2003. Improved Estimation of the Covariance Matrix of Stock Returns with an application to portfolio selection. Journal of empirical finance. 10. 603-621.
"""
function identity_regularisation(mat::Hermitian, identity_weight::Real)
    II = I(size(mat)[1])
    # This is a bit different to the paper as we don't have m_n. This is to ensure the diagonal stays as ones.
    mat_prime = (identity_weight .* II) + (1 - identity_weight) .* mat
    return Hermitian(mat_prime)
end
function identity_regularisation(mat::Hermitian, asset_returns::DataFrame) where R<:Real
    II   = I(size(mat)[1])
    m    = squared_frobenius_distance(mat, II) # Lemma 3.2
    d    = sqrt(squared_frobenius_distance(mat, m*II))      # Lemma 3.3
    bbar = b_bar(asset_returns, mat)
    b    = min(bbar, d)
    return identity_regularisation(mat, b/d)
end
function identity_regularisation(mat::Hermitian, ts::SortedDataFrame,  mat_labels::Vector;
                                 spacing::Union{Missing,<:Real} = missing)
    at_times = ismissing(spacing) ? get_all_refresh_times(ts, mat_labels) : collect(0:spacing:maximum(ts.df[:,ts.time]))
    dd_compiled = latest_value(ts, at_times; assets = mat_labels)
    asset_returns = get_returns(dd_compiled; rescale_for_duration = true)
    return identity_regularisation(mat, asset_returns)
end
function identity_regularisation(covariance_matrix::CovarianceMatrix, ts::SortedDataFrame;
                                 spacing::Union{Missing,<:Real} = missing, apply_to_covariance::Bool = true)
     if apply_to_covariance
         actual_covariance = covariance(covariance_matrix)
         regularised_covariance = identity_regularisation(actual_covariance,
                                           ts, covariance_matrix.labels; spacing = spacing)
         corr, vols = cov2cor_and_vol(regularised_covariance, 1)
         return CovarianceMatrix(corr, vols, covariance_matrix.labels, covariance_matrix.time_period_per_unit)
     else
         return CovarianceMatrix(Hermitian(identity_regularisation(covariance_matrix.correlation, ts, covariance_matrix.labels; spacing = spacing)),
                     covariance_matrix.volatility, covariance_matrix.labels, covariance_matrix.time_period_per_unit)
     end
end
function identity_regularisation(covariance_matrix::CovarianceMatrix, identity_weight::Real; apply_to_covariance = false)
     if apply_to_covariance
         regularised_covariance = identity_regularisation(covariance(covariance_matrix,1),identity_weight)
         corr, vols = cov2cor_and_vol(regularised_covariance, 1)
         return CovarianceMatrix(corr, vols, covariance_matrix.labels, covariance_matrix.time_period_per_unit)
     else
         return CovarianceMatrix(Hermitian(identity_regularisation(covariance_matrix.correlation, identity_weight)),
                       covariance_matrix.volatility, covariance_matrix.labels, covariance_matrix.time_period_per_unit)
     end
end


text:
struct FiveNumberSummary
    minimum::Float64
    lower_quartile::Float64
    median::Float64
    upper_quartile::Float64
    maximum::Float64
end

mutable struct BoxplotGraphics{R<:Number} <: GraphicsArea
    data::Vector{FiveNumberSummary}
    color::ColorType
    char_width::Int
    visible::Bool
    min_x::R
    max_x::R

    function BoxplotGraphics{R}(
        data::AbstractVector{R},
        char_width::Int,
        visible::Bool,
        color::UserColorType,
        min_x::R,
        max_x::R
    ) where R
        char_width = max(char_width, 10)
        if min_x == max_x
            min_x = min_x - 1
            max_x = max_x + 1
        end
        new{R}(
            [FiveNumberSummary(
                minimum(data),
                percentile(data, 25),
                percentile(data, 50),
                percentile(data, 75),
                maximum(data)
            )],
            crayon_256_color(color), char_width, visible, min_x, max_x
        )
    end
end

nrows(c::BoxplotGraphics) = 3*length(c.data)
ncols(c::BoxplotGraphics) = c.char_width

BoxplotGraphics(
    data::AbstractVector{R},
    char_width::Int;
    visible::Bool = true,
    color::UserColorType = :green,
    min_x::Number = minimum(data),
    max_x::Number = maximum(data)
) where {R <: Number} = BoxplotGraphics{R}(
    data, char_width, visible, crayon_256_color(color), R(min_x), R(max_x)
)

function addseries!(c::BoxplotGraphics, data::AbstractVector{R}) where {R <: Number}
    mi, ma = extrema(data)
    push!(c.data, FiveNumberSummary(
        mi,
        percentile(data, 25),
        percentile(data, 50),
        percentile(data, 75),
        ma
    ))
    c.min_x = min(mi, c.min_x)
    c.max_x = max(ma, c.max_x)
    c
end

function printrow(io::IO, c::BoxplotGraphics, row::Int)
    0 < row <= nrows(c) || throw(ArgumentError("Argument row out of bounds: $row"))
    series = c.data[ceil(Int, row/3)]

    transform(value) = clamp(
        round(Int, (value - c.min_x) / (c.max_x - c.min_x) * c.char_width), 1, c.char_width
    )

    series_row = Int((row-1) % 3) + 1

    min_char = ['╷', '├' , '╵'][series_row]
    line_char = [' ', '─' , ' '][series_row]
    left_box_char = ['┌', '┤' , '└'][series_row]
    line_box_char = ['─', ' ' , '─'][series_row]
    median_char = ['┬', '│' , '┴'][series_row]
    right_box_char = ['┐', '├' , '┘'][series_row]
    max_char = ['╷', '┤' , '╵'][series_row]

    line = [' ' for _ in 1:c.char_width]

    # Draw shapes first - this is most important,
    # so they'll always be drawn even if there's not enough space
    line[transform(series.minimum)] = min_char
    line[transform(series.lower_quartile)] = left_box_char
    line[transform(series.median)] = median_char
    line[transform(series.upper_quartile)] = right_box_char
    line[transform(series.maximum)] = max_char

    # Fill in gaps with lines
    for i in (transform(series.minimum)+1):(transform(series.lower_quartile)-1)
        line[i] = line_char
    end
    for i in (transform(series.lower_quartile)+1):(transform(series.median)-1)
        line[i] = line_box_char
    end
    for i in (transform(series.median)+1):(transform(series.upper_quartile)-1)
        line[i] = line_box_char
    end
    for i in (transform(series.upper_quartile)+1):(transform(series.maximum)-1)
        line[i] = line_char
    end

    print_color(c.color, io, join(line))
    nothing
end


text:
"""
An AbstractDimension tags the dimensions in an AbstractArray.

It can also contain spatial coordinates and their metadata. For simplicity,
the same types are used both for storing dimension information and for indexing.
"""
abstract type AbstractDimension{T,G,M} end

ConstructionBase.constructorof(d::Type{<:AbstractDimension}) = basetypeof(d)

const AbDim = AbstractDimension
const AbDimType = Type{<:AbDim}
const AbDimTuple = Tuple{Vararg{<:AbDim,N}} where N
const AbDimVector = Vector{<:AbDim}
const DimOrDimType = Union{AbDim,AbDimType}
const AllDimensions = Union{AbDim,AbDimTuple,AbDimType,
                            Tuple{Vararg{AbDimType}},
                            Vector{<:AbDim}}

# Getters
val(dim::AbDim) = dim.val
grid(dim::AbDim) = dim.grid
metadata(dim::AbDim) = dim.metadata

order(dim::AbDim) = order(grid(dim))
indexorder(dim::AbDim) = indexorder(order(dim))
arrayorder(dim::AbDim) = arrayorder(order(dim))

# DimensionalData interface methods
rebuild(dim::D, val, grid=grid(dim)) where D <: AbDim =
    rebuild(dim, val, grid, metadata(dim))
rebuild(dim::D, val, grid, metadata) where D <: AbDim =
    constructorof(D)(val, grid, metadata)

dims(x::AbDim) = x
dims(x::AbDimTuple) = x
name(dim::AbDim) = name(typeof(dim))
shortname(d::AbDim) = shortname(typeof(d))
shortname(d::Type{<:AbDim}) = name(d)
units(dim::AbDim) = metadata(dim) == nothing ? nothing : get(metadata(dim), :units, nothing)

bounds(A, args...) = bounds(dims(A), args...)
bounds(dims::AbDimTuple, lookupdims::Tuple) = bounds(dims[[dimnum(dims, lookupdims)...]]...)
bounds(dims::AbDimTuple, dim::DimOrDimType) = bounds(dims[dimnum(dims, dim)])
bounds(dims::AbDimTuple) = (bounds(dims[1]), bounds(tail(dims))...)
bounds(dims::Tuple{}) = ()
bounds(dim::AbDim) = bounds(indexorder(dim), dim)
# TODO bounds should include the span of the last cell
bounds(::Forward, dim::AbDim) = first(val(dim)), last(val(dim)) # + val(span(dim))
bounds(::Reverse, dim::AbDim) = last(val(dim)), first(val(dim)) # + val(span(dim))
# TODO bounds for irregular grids

# Base methods
Base.eltype(dim::Type{<:AbDim{T}}) where T = T
Base.length(dim::AbDim) = length(val(dim))
Base.show(io::IO, dim::AbDim) = begin
    printstyled(io, "\n", name(dim), ": "; color=:red)
    show(io, typeof(dim))
    printstyled(io, "\nval: "; color=:green)
    show(io, val(dim))
    printstyled(io, "\ngrid: "; color=:yellow)
    show(io, grid(dim))
    printstyled(io, "\nmetadata: "; color=:blue)
    show(io, metadata(dim))
    print(io, "\n")
end

# AbstractArray methods where dims are the dispatch argument

@inline rebuildsliced(A, data, I) = rebuild(A, data, slicedims(A, I)...)

Base.@propagate_inbounds Base.getindex(A::AbstractArray, dims::Vararg{<:AbDim{<:Number}}) =
    getindex(A, dims2indices(A, dims)...)
Base.@propagate_inbounds Base.getindex(A::AbstractArray, dims::Vararg{<:AbstractDimension}) =
    getindex(A, dims2indices(A, dims)...)

Base.@propagate_inbounds Base.setindex!(A::AbstractArray, x, dims::Vararg{<:AbstractDimension}) =
    setindex!(A, x, dims2indices(A, dims)...)

Base.@propagate_inbounds Base.view(A::AbstractArray, dims::Vararg{<:AbstractDimension}) =
    view(A, dims2indices(A, dims)...)

@inline Base.axes(A::AbstractArray, dims::DimOrDimType) = axes(A, dimnum(A, dims))
@inline Base.size(A::AbstractArray, dims::DimOrDimType) = size(A, dimnum(A, dims))

#= SplitApplyCombine methods?
Should allow groupby using dims lookup to make this worth the dependency
Like group by time/lattitude/height band etc.

SplitApplyCombine.splitdims(A::AbstractArray, dims::AllDimensions) =
    SplitApplyCombine.splitdims(A, dimnum(A, dims))

SplitApplyCombine.splitdimsview(A::AbstractArray, dims::AllDimensions) =
    SplitApplyCombine.splitdimsview(A, dimnum(A, dims))
=#

"""
Dimensions with user-set type paremeters
"""
abstract type AbstractParametricDimension{X,T,G,M} <: AbstractDimension{T,G,M} end

"""
A generic dimension. For use when custom dims are required when loading
data from a file. The sintax is ugly and verbose to use for indexing,
ie `Dim{:lat}(1:9)` rather than `Lat(1:9)`. This is the main reason
they are not the only type of dimension availabile.
"""
struct Dim{X,T,G,M} <: AbstractParametricDimension{X,T,G,M}
    val::T
    grid::G
    metadata::M
    Dim{X}(val, grid, metadata) where X =
        new{X,typeof(val),typeof(grid),typeof(metadata)}(val, grid, metadata)
end

Dim{X}(val=:; grid=AllignedGrid(), metadata=nothing) where X = Dim{X}(val, grid, metadata)
name(::Type{<:Dim{X}}) where X = "Dim $X"
shortname(::Type{<:Dim{X}}) where X = "$X"
basetypeof(::Type{<:Dim{X}}) where {X} = Dim{X}

"""
    @dim typ [name=string(typ)] [shortname=string(typ)]

Macro to easily define specific dimensions.

Example:
```julia
@dim Lat "Lattitude"
```
"""
macro dim(typ, name=string(typ), shortname=string(typ))
    esc(quote
        struct $typ{T,G,M} <: AbstractDimension{T,G,M}
            val::T
            grid::G
            metadata::M
        end
        $typ(val=:; grid=AllignedGrid(), metadata=nothing) = $typ(val, grid, metadata)
        DimensionalData.name(::Type{<:$typ}) = $name
        DimensionalData.shortname(::Type{<:$typ}) = $shortname
    end)
end

# Define some common dimensions
@dim Time
@dim X
@dim Y
@dim Z


text:
#
# crop-tests.jl -
#
# Tests for cropping and zero-padding.
#
module TestingLazyAlgebraCrop

using Test
using LazyAlgebra
using LazyAlgebra.Foundations

@testset "Cropping and zero-padding" begin
    #
    # Private methods for testing.
    #

    offset(outer::NTuple{N,Int}, inner::NTuple{N,Int}) where {N} =
        ntuple(i -> (outer[i]>>1) - (inner[i]>>1), Val(N))

    function crop(x::AbstractArray{T,N}, siz::NTuple{N,Int}) where {T,N}
        return crop!(Array{T,N}(undef, siz), x)
    end

    function crop(x::AbstractArray{T,N}, siz::NTuple{N,Int},
                  off::NTuple{N,Int}) where {T,N}
        return crop!(Array{T,N}(undef, siz), x, off)
    end

    function zeropad(x::AbstractArray{T,N}, siz::NTuple{N,Int}) where {T,N}
        return zeropad!(Array{T,N}(undef, siz), x)
    end

    function zeropad(x::AbstractArray{T,N}, siz::NTuple{N,Int},
                     off::NTuple{N,Int}) where {T,N}
        return zeropad!(Array{T,N}(undef, siz), x, off)
    end

    function subregionindices(sub::AbstractArray{<:Any,N},
                              big::AbstractArray{<:Any,N},
                              off::NTuple{N,Int}) where {N}
        I = CartesianIndices(sub) # indices in smallest region
        J = CartesianIndices(big) # indices in largest region
        k = CartesianIndex(off)   # offset index
        (first(J) ≤ first(I) + k && last(I) + k ≤ last(J)) ||
            error("out of range sub-region")
        return I, J, k
    end

    function crop!(y::AbstractArray{T,N},
                   x::AbstractArray{<:Any,N},
                   off::NTuple{N,Int} = offset(size(x), size(y))) where {T,N}
        I, J, k = subregionindices(y, x, off)
        @inbounds @simd for i ∈ I
            y[i] = x[i + k]
        end
        return y
    end

    function zeropad!(y::AbstractArray{T,N},
                      x::AbstractArray{<:Any,N},
                      off::NTuple{N,Int} = offset(size(y), size(x)),
                      init::Bool = false) where {T,N}
        I, J, k = subregionindices(x, y, off)
        init || fill!(y, zero(T))
        @inbounds @simd for i ∈ I
            y[i+k] = x[i]
        end
        return y
    end

    #
    # Miscellaneous tests.
    #
    @test_throws ErrorException CroppingOperator((2,3,), (2,3,4,))
    @test_throws ErrorException CroppingOperator((2,3,), (2,3,4,), (0,1,))

    #
    # Tests for different sizes and element types.
    #
    for (osz, isz) in (((3,), (8,)),
                       ((4,), (8,)),
                       ((3,), (7,)),
                       ((4,), (7,)),
                       ((4,5), (6,7))),
        T in (Float64, Complex{Float32})

        # Basics methods.
        R = real(T)
        C = CroppingOperator(osz, isz)
        off = ntuple(i->Int16(isz[i] - osz[i])>>1, length(isz))
        @test ZeroPaddingOperator(isz, osz) === Adjoint(C)
        @test ZeroPaddingOperator(isz, osz, off) === Adjoint(CroppingOperator(osz, isz, off))
        @test input_ndims(C) == length(isz)
        @test input_size(C) == isz
        @test all(i -> input_size(C,i) == isz[i], 1:length(isz))
        @test output_ndims(C) == length(osz)
        @test output_size(C) == osz
        @test all(i -> output_size(C,i) == osz[i], 1:length(osz))

        # Compare result of opertaor with reference implementation.
        x = rand(T, isz)
        xsav = vcopy(x)
        Cx = C*x
        @test x == xsav
        @test Cx == crop(x, osz)
        y = rand(T, osz)
        ysav = vcopy(y)
        Cty = C'*y
        @test y == ysav
        @test Cty == zeropad(y, isz)

        # Test various possibilities for apply!
        atol = 0
        rtol = eps(R)
        for α in (0, 1, -1,  2.71, π),
            β in (0, 1, -1, -1.33, Base.MathConstants.φ),
            scratch in (false, true)
            # Test operator.
            @test apply!(α, Direct, C, x, scratch, β, vcopy(y)) ≈
                R(α)*Cx + R(β)*y  atol=atol rtol=rtol
            if scratch
                vcopy!(x, xsav)
            else
                @test x == xsav
            end
            # Test  adjoint.
            @test apply!(α, Adjoint, C, y, scratch, β, vcopy(x)) ≈
                R(α)*Cty + R(β)*x  atol=atol rtol=rtol
            if scratch
                vcopy!(y, ysav)
            else
                @test y == ysav
            end
        end
    end
end
nothing

end # module


text:
using InteractiveChaos, Makie

ds = Systems.henonheiles()

potential(x, y) = 0.5(x^2 + y^2) + (x^2*y - (y^3)/3)
energy(x,y,px,py) = 0.5(px^2 + py^2) + potential(x,y)
const E = energy(get_state(ds)...)

function complete(y, py, x)
    V = potential(x, y)
    Ky = 0.5*(py^2)
    Ky + V ≥ E && error("Point has more energy!")
    px = sqrt(2(E - V - Ky))
    ic = [x, y, px, py]
    return ic
end

chaotic = get_state(ds)
stable = [0., 0.1, 0.5, 0.]

plane = (1, 0.0)

#does this work?

psos = interactive_poincaresos(ds, plane, (2, 4), complete; markersizes = (-5, -1))


text:
# This is a batched matrix multiplication test, with the batch index coming first,
# from this thread:
# https://discourse.julialang.org/t/non-matching-indices-error-using-tensoroperations/35136
# Some improvements to OMEinsum since then make it much quicker.

julia> using Einsum, OMEinsum, TensorCast, Tullio, LoopVectorization

julia> f_ome(a,b) = @ein c[k, n] := a[k, n, c] * conj(b)[c,k]; # permutedims + batched_mul

julia> f_cast(a,b) = @reduce c[k, n] := sum(l) a[k, n, l] * conj(b[l, k]); # broadcasting

julia> f_ein(a,b) = @einsum c_[k, n] := a[k, n, c] * conj(b[c,k]); # naiive loops

julia> f_viel(a,b) = @vielsum c_[k, n] := a[k, n, c] * conj(b[c,k]); # plus threads

julia> f_tul(a,b) = @tullio c[k, n] := a[k, n, c] * conj(b[c,k]); # less naiive loops?

julia> a = randn(ComplexF64, 300, 400, 500); b = randn(ComplexF64, 500, 300);

julia> f_ome(a,b) ≈ f_cast(a,b) ≈ f_ein(a,b) ≈ f_tul(a,b)
true

julia> @btime f_ome($a, $b);
  342.619 ms (104 allocations: 921.49 MiB)

julia> @btime f_cast($a, $b);
  557.428 ms (26 allocations: 919.65 MiB)

julia> @btime f_ein($a, $b);
  264.429 ms (2 allocations: 1.83 MiB)

julia> @btime f_viel($a, $b);
  147.739 ms (25 allocations: 1.83 MiB)

julia> @btime f_tul($a, $b);
  191.859 ms (836 allocations: 1.86 MiB)

# But LoopVectorization isn't yet in play, as the arrays have complex elements.
# On real numbers, it makes a big difference!

julia> ar = real(a); br = real(b);

julia> @btime f_ome($ar, $br);
  190.672 ms (102 allocations: 459.60 MiB)

julia> @btime f_ein($ar, $br);
  122.616 ms (2 allocations: 937.58 KiB)

julia> @btime f_viel($ar, $br);
  82.922 ms (25 allocations: 940.33 KiB)

julia> @btime f_tul($ar, $br);
  39.105 ms (836 allocations: 967.17 KiB)

# Can we get there with StructArrays?
# (Perhaps this could be made automatic, with yet more macrology...)

julia> using StructArrays, LoopVectorization

julia> @time sa = StructArray(a); sb = StructArray(b);
  4.160518 seconds (120.00 M allocations: 3.576 GiB, 29.61% gc time)

julia> function f_tul(a::StructArray, b::StructArray)
           a_re, a_im = a.re, a.im
           b_re, b_im = b.re, b.im
           @tullio c_re[k, n] := a_re[k, n, c] * b_re[c,k] + a_im[k, n, c] * b_im[c,k]
           @tullio c_im[k, n] := a_im[k, n, c] * b_re[c,k] - a_re[k, n, c] * b_im[c,k]
           StructArray{eltype(a)}((c_re, c_im))
       end
f_tul (generic function with 2 methods)

julia> f_ein(sa, sb) ≈ f_tul(sa, sb)
true

julia> @btime f_tul($sa, $sb);
  136.801 ms (1670 allocations: 1.89 MiB)

# That's worse than twice the real calculation, but I guess each line is harder.
# Converting to StructArrays is really slow though!

julia> @btime StructArray($a);
  3.068 s (119999498 allocations: 3.58 GiB)

julia> @btime real($a);
  248.114 ms (2 allocations: 457.76 MiB)

# Compare Einsum on StructArrays: not great, however you do it.

julia> @btime f_viel($sa, $sb);
  479.932 ms (25 allocations: 1.83 MiB)

julia> typeof(f_viel(sa, sb))
Array{Complex{Float64},2}

julia> similar(sa, 1,3)
1×3 StructArray(::Array{Float64,2}, ::Array{Float64,2}) with eltype Complex{Float64}:
 2.29092e-314+2.27751e-314im  2.29092e-314+3.7518e-314im  3.7518e-314+3.75184e-314im

julia> similar(sa, ComplexF64, 1,3)
1×3 Array{Complex{Float64},2}:
 2.66668e-314+2.66668e-314im  2.66668e-314+2.29246e-314im  2.29246e-314+2.27751e-314im

julia> f_viel!(c_,a,b) = @vielsum c_[k, n] = a[k, n, c] * conj(b[c,k]);

julia> sc = similar(f_tul(sa, sb)); typeof(sc)
StructArray{Complex{Float64},2,NamedTuple{(:re, :im),Tuple{Array{Float64,2},Array{Float64,2}}},Int64}

julia> @btime f_viel!($sc, $sa, $sb);
  478.993 ms (240023 allocations: 5.50 MiB)

# Try KernelAbstractions? This should handle threading instead of my code.

julia> using KernelAbstractions, CuArrays # CuArrays just to trigger it!

julia> ENV["JULIA_DEBUG"] = Main;

julia> f_tul2(a,b) = @tullio c[k, n] := a[k, n, c] * conj(b[c,k])  threads=false;

julia> f_tul2(a, b);
┌ Debug: KernelAbstractions CPU actor:
│   typeof.(tuple(ℛ::AbstractArray{𝒯}, a, b, 𝒶_n, 𝒶_k, 𝒶_c)) = (Array{Complex{Float64},2}, Array{Complex{Float64},3}, Array{Complex{Float64},2}, UnitRange{Int64}, UnitRange{Int64}, UnitRange{Int64})
└ @ Main ~/.julia/dev/Tullio/src/macro.jl:724

julia> ENV["JULIA_DEBUG"] = "none";

julia> @btime f_tul2($a, $b);
  599.003 ms (64 allocations: 1.84 MiB)

# Note that if you just run f_tul threads=false without KernelAbstractions, it takes 1.3sec,
# which is much much slower than @einsum, so perhaps that's another bug.

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

julia> using StructArrays

julia> a = randn(ComplexF64, 300, 400, 500);

julia> @time StructArray(a);
  4.967282 seconds (120.00 M allocations: 3.576 GiB, 21.80% gc time)

julia> @time StructArray{ComplexF64}((real(a), imag(a)));
  0.630163 seconds (7 allocations: 915.528 MiB, 10.91% gc time)

julia> @code_warntype StructArray(a)
Variables
  #self#::Type{StructArray}
  v::Array{Complex{Float64},3}
  #46::StructArrays.var"#46#48"

Body::StructArray{Complex{Float64},3,NamedTuple{(:re, :im),Tuple{Array{Float64,3},Array{Float64,3}}},Int64}
1 ─      (#46 = %new(StructArrays.:(var"#46#48")))
│   %2 = #46::Core.Compiler.Const(StructArrays.var"#46#48"(), false)
│   %3 = StructArrays.:(var"#StructArray#45")(%2, #self#, v)::StructArray{Complex{Float64},3,NamedTuple{(:re, :im),Tuple{Array{Float64,3},Array{Float64,3}}},Int64}
└──      return %3


text:
using JuMP
using Clp

CARBS = 50
PROTEIN = 40
FAT = 60

m = Model(Clp.Optimizer)

@variable(m, steak ≥ 0)
@variable(m, potatoes ≥ 0)

@objective(m, Min, 8 * steak + 4 * potatoes)

@constraint(m, 5 * steak + 15 * potatoes ≥ CARBS)
@constraint(m, 20 * steak + 5 * potatoes ≥ PROTEIN)
@constraint(m, 15 * steak + 2 * potatoes ≤ FAT)

print(m)

optimize!(m)

status = termination_status(m)

println("Solution status: ", status)

println("Objective value: ", objective_value(m))
println("potatoes = ", value(steak))
println("steak = ", value(potatoes))


text:
export MDString, MDNode, operands, Metadata

@checked struct MetadataAsValue <: Value
    ref::API.LLVMValueRef
end
identify(::Type{Value}, ::Val{API.LLVMMetadataAsValueValueKind}) = MetadataAsValue

# NOTE: the C API doesn't allow us to differentiate between MD kinds,
#       all are wrapped by the opaque MetadataAsValue...

const MDString = MetadataAsValue

MDString(val::String) = MDString(API.LLVMMDString(val, length(val)))

MDString(val::String, ctx::Context) =
    MDString(API.LLVMMDStringInContext(ctx, val, length(val)))

function Base.string(md::MDString)
    len = Ref{Cuint}()
    ptr = API.LLVMGetMDString(md, len)
    ptr == C_NULL && throw(ArgumentError("invalid metadata, not a MDString?"))
    return unsafe_string(convert(Ptr{Int8}, ptr), len[])
end

const MDNode = MetadataAsValue

MDNode(vals::Vector{<:Value}) =
    MDNode(API.LLVMMDNode(vals, length(vals)))

MDNode(vals::Vector{<:Value}, ctx::Context) =
    MDNode(API.LLVMMDNodeInContext(ctx, vals, length(vals)))

function operands(md::MDNode)
    nops = API.LLVMGetMDNodeNumOperands(md)
    ops = Vector{API.LLVMValueRef}(undef, nops)
    API.LLVMGetMDNodeOperands(md, ops)
    return Value[Value(op) for op in ops]
end

@checked struct Metadata
    ref::API.LLVMMetadataRef
end

Base.unsafe_convert(::Type{API.LLVMMetadataRef}, md::Metadata) = md.ref

function Metadata(val::Value)
    return Metadata(API.LLVMValueAsMetadata(val))
end

function Value(md::Metadata, ctx::Context)
    return MetadataAsValue(API.LLVMMetadataAsValue(ctx, md))
end