text:
# Copyright (c) 2015, Intel Corporation
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#     * Redistributions of source code must retain the above copyright notice,
#       this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of Intel Corporation nor the names of its contributors
#       may be used to endorse or promote products derived from this software
#       without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

export Net, init, forward, backward, clear_∇, clear_values, add_ensemble,
       add_connections, copy_from_mic, copy_to_mic, get_buffer,
       get_value, get_gradient
importall ParallelAccelerator
import ParallelAccelerator.CGen
import ParallelAccelerator.CGen.mk_parallel_loophead
import CompilerTools
CGen.set_include_blas()
# using ParallelAccelerator.J2CArray
ENV["KMP_AFFINITY"] = "granularity=fine,compact"

num_threads = get(ENV, "OMP_NUM_THREADS", nothing)
if num_threads == nothing
    num_threads = Base.CPU_CORES
    ENV["OMP_NUM_THREADS"] = "$num_threads"
else
    num_threads = parse(Int, num_threads)
end
LOSSY_GRADIENTS = false
const TILE_SIZE = 4
const MICRO_BATCH_SIZE = num_threads
NOFUSE = 0
LATTE_DISABLE_TILING = false
LATTE_DISABLE_TILE_FUSION = false

include("transforms/util.jl")
include("transforms/fixed_dims.jl")
include("transforms/distribute.jl")
include("transforms/tree_cleaners.jl")
include("transforms/neuron.jl")

include("optimizers/tiling.jl")
include("optimizers/fusion.jl")
include("optimizers/gemm_pattern_match.jl")
include("optimizers/array_expr_inline.jl")
include("optimizers/wrap_for_loops.jl")
include("optimizers/parallelize.jl")

"""
Optimize a function `fn`

** Params **

- `args`                -- an ordered vector of arguments (buffers)
- `tile_fusion_factors` -- factors for tile fusions determined by connection
                           structure
- `fn`                  -- the ast of the function to be optimized
"""
function optimize(args::Vector, tile_fusion_factors::Vector, fn::Expr)
    ast = macroexpand(fn)
    ast = remove_line_nodes(ast)
    arg_name_map = build_arg_name_map(args, ast)
    debugp(2, "----- Begin : Function Before Optimization -----")
    debugp(2, ast)
    debugp(2, "----- End   : Function Before Optimization -----")
    for (var1, var2, var3) in simple_ijk_orders
        fn = symbol(:pattern_match_gemm, var1, var2, var3)
        ast = AstWalk(ast, eval(fn), arg_name_map)
    end
    if !LATTE_DISABLE_TILING
        map!(get_inner_loop_tiler(tile_fusion_factors), ast.args[2].args)
    end
    ast = clean_tree(ast)
    ast = AstWalk(ast, pattern_match_gemm6, arg_name_map)
    ast = AstWalk(ast, pattern_match_gemm5, arg_name_map)
    ast = AstWalk(ast, pattern_match_gemm4, arg_name_map)
    # TODO: Above pattern matches should be cleaned up to ingest the "clean"
    # ast like fuse_loops below
    ast = clean_tree(ast)
    ast.args[2].args = fuse_loops(ast.args[2].args)
    debugp(1, "----- Begin : Function After Optimization  -----")
    debugp(1, ast)
    debugp(1, "----- End   : Function After Optimization  -----")
    return ast
end

mk_clamp(idx, _max) = :($idx > 0 && $idx <= $_max)

"""
Synthesize a loopnest to copy value or ∇ to the appropriate buffer
"""
function gen_copy_block(net::Net, ensemble::AbstractEnsemble,
                        connection::Connection, index::Int; ∇::Bool=false)
    if ∇
        shape = size(get_buffer(net, ensemble, :∇))
        sink_name  = symbol(ensemble.name, :∇inputs, index)
        source_name = symbol(connection.source.name, :∇)
    else
        shape = size(get_buffer(net, ensemble, :value))
        sink_name   = symbol(ensemble.name,:inputs,index)
        source_name = symbol(connection.source.name,:value)
    end
    N = length(shape)
    source = get_buffer(net, source_name)

    # connection.is_dim_fixed is a Bool[], we use the inverse to select non
    # fixed indices
    non_fixed_indices = [collect(1:N-1)[!connection.is_dim_fixed]..., N]

    sink_idx = [symbol(:_neuron_index_,d) for d in non_fixed_indices]

    mapping_args = [symbol(:_neuron_index_,d) for d in 1:N-1]
    mapping_inlined = inline(connection.mapping, mapping_args)
    src_idx = get_src_idx(mapping_inlined)

    body = Expr(:block)
    idx = [symbol(:_u__, i) for i = 1:length(src_idx)]
    for_block = gen_loop_nest(body, idx, src_idx)

    if ∇
        assign = :($source_name[$(idx...), $(symbol(:_neuron_index_,N))] +=
                       $sink_name[count, $(sink_idx...)])
        if connection.padding > 0
            cond = mk_clamp(idx[1], size(source, 1))
            for (i, d) in zip(idx[2:end], size(source)[2:end-1])
                cond = :($cond && $(mk_clamp(i, d)))
            end
            assign = :(if $cond
                $assign
            end)
        end
        statements = [
            assign
        ]
    else
        rhs = :($source_name[$(idx...), $(symbol(:_neuron_index_,N))])
        if connection.padding > 0
            cond = mk_clamp(idx[1], size(source, 1))
            for (i, d) in zip(idx[2:end], size(source)[2:end-1])
                cond = :($cond && $(mk_clamp(i, d)))
            end
            rhs = :(ifelse($cond, $rhs, 0.0f0))
        end
        statements = [:($sink_name[count, $(sink_idx...)] = $rhs)]
    end
    append!(body.args, (quote
        $(statements...)
        count += 1
    end).args)
    copy_block = quote
        $(mapping_inlined.args[1:end-1]...)
        count = 1
        $for_block
    end
    if ∇
        push!(copy_block.args, quote
            for i = 1:count-1
                $sink_name[i, $(sink_idx...)] = 0.0f0
            end
        end)
    end
    vars   = [symbol(:_neuron_index_,i) for i in 1:N]
    ranges = [connection.is_dim_fixed[i] ? :(1:1) : :(1:$(shape[i])) for i in 1:N-1]
    push!(ranges, :(1:$(shape[N])))
    gen_loop_nest(copy_block, vars, ranges)
end

"""
Extract the final expression of the mapping function to be used as an indexing
expression.  Handles cases where the final expression can be a Tuple or a
single value.

** Params **

- `mapping` -- an ast for a mapping function
"""
function get_src_idx(mapping::Expr)
    if isa(mapping.args[end], Expr) &&
            mapping.args[end].head == :call &&
            isa(mapping.args[end].args[1], TopNode)
        return mapping.args[end].args[2:end]
    else
        return [mapping.args[end]]
    end
end

"""
Synthesize the code to execute forward for an ensemble

** Params **

- `ensemble`     -- the `AbstractEnsemble` that is having its forward synthesized
- `net`          -- the `Net` containing `ensemble`
- `compute_body` -- the body that the synthesized code will be added to
- `compute_args` -- a set of arguments reference in `compute_body`.  Any values
                    referenced in the synthesized code should be added to
                    `compute_args`
"""
function gen_neuron_forward(ensemble::AbstractEnsemble, net::Net,
                            compute_body::Dict{Phase, Vector},
                            compute_args::ArgSet)
    forward_tasks = net.forward_tasks

    neuron_forward = get_forward(ensemble)
    value = get_buffer(net, ensemble, :value)
    N = ndims(value)
    body = []
    args = []
    arg_dim_info = ensemble.arg_dim_info
    for (index, connection) in enumerate(ensemble.connections)
        # Store argument info
        sink_name = symbol(ensemble.name,:inputs,index)
        arg_dim_info[sink_name] = connection.is_dim_fixed

        # If copying, add source buffer to args and generate the copy block
        if connection.copy
            source_name = symbol(connection.source.name,:value)
            push!(args, source_name)
            push!(body, gen_copy_block(net, ensemble, connection, index))
        end
    end
    # Transform body of neuron function
    ast, _args = transform_neuron_fn(neuron_forward, ensemble)
    map((a) -> push!(args, a), _args)  # FIXME: append! doesn't support sets

    ast = remove_line_nodes(ast)

    # Drop indices for shared value dimensions
    ast = drop_fixed_dims(ast, arg_dim_info)

    # Synthesize loop nest and append to body
    append_neuron_loop!(body, ast.args, value)

    # Collect arguments
    args = collect(args)

    fn = symbol(ensemble.name,:forward)
    arg_bufs = map((arg) -> get_buffer(net, arg), args)
    params = [:($arg::$(typeof(buf))) for (arg, buf) in zip(args, arg_bufs)]
    tile_fusion_factors = []
    if length(ensemble.connections) == 1 && ndims(ensemble) > 2
        for d in ensemble.connections[1].shape[2:end]
            push!(tile_fusion_factors, (d, 0))
        end
    else
        for d in ensemble.connections[1].shape[2:end]
            push!(tile_fusion_factors, (0, 0))
        end
    end
    defn = optimize(arg_bufs, tile_fusion_factors, :(function $fn($(params...))
        $(body...)
    end))
    body = defn.args[2].args
    # Disable fusion across ensemble code if there exists a dependency
    is_dependence(conn) = length(conn.shape) > 2 && any(collect(conn.shape[3:end]) .> 1)
    if any(map(is_dependence, ensemble.connections))
        unshift!(body, :NOFUSE)
    end
    if ensemble.phase in [Train, TrainTest]
        for param in ensemble.params
            push!(compute_body[Train], :(update_param($(object_id(param)))))
        end
        append!(compute_body[Train], deepcopy(body))
        union!(compute_args[Train], args)
    end
    if ensemble.phase in [Test, TrainTest]
        append!(compute_body[Test], deepcopy(body))
        union!(compute_args[Test], args)
    end
end

"""
TODO: doc
"""
function gen_neuron_backward(ensemble::AbstractEnsemble, net::Net,
                             compute_body::Vector,
                             compute_args::Set)
    if !(ensemble.phase in [Train, TrainTest])
        return
    end
    backward_tasks = net.backward_tasks

    neuron_backward = get_backward(ensemble)
    ∇ = get_buffer(net, ensemble, :∇)
    N = ndims(∇)

    # Transform neuron backward definition and collect arguments
    ast, args = transform_neuron_fn(neuron_backward, ensemble)
    ast = remove_line_nodes(ast)

    body = []
    arg_dim_info = ensemble.arg_dim_info

    # Collect source buffer info
    for (index, connection) in enumerate(ensemble.connections)
        arg_dim_info[symbol(ensemble.name, :inputs, index)] =
            connection.is_dim_fixed
        sink_name = symbol(ensemble.name, :∇inputs, index)
        arg_dim_info[sink_name] = connection.is_dim_fixed
        # Copy inputs if needed
        # inputs$index will be added to args by transform_neuron_fn if it is
        # used in the function
        inputs_name = symbol(ensemble.name, :inputs, index)
        if connection.copy && inputs_name in args
            source_name = symbol(connection.source.name, :value)
            push!(args, source_name)
            push!(body, gen_copy_block(net, ensemble, connection, index))
        end
    end

    # Replace indexing expressions for shared values
    ast = drop_fixed_dims(ast, arg_dim_info)

    # Synthesize loop nest and append to body
    append_neuron_loop!(body, ast.args, ∇)
    fn = symbol(ensemble.name,:backward)
    arg_bufs = map((arg) -> get_buffer(net, arg), args)
    params = [:($arg::$(typeof(buf))) for (arg, buf) in zip(args, arg_bufs)]
    tile_fusion_factors = []
    if length(ensemble.connections) == 1 && body[1] != :NOFUSE && ndims(ensemble) > 2
        for d in ensemble.connections[1].shape[2:end]
            push!(tile_fusion_factors, (0, d))
        end
    else
        for d in ensemble.connections[1].shape[2:end]
            push!(tile_fusion_factors, (0, 0))
        end
    end
    for (index, connection) in enumerate(ensemble.connections)
        if !isa(connection.source, DataEnsemble) && connection.copy
            source_name = symbol(connection.source.name, :∇)
            push!(args, source_name)
            push!(body, gen_copy_block(net, ensemble, connection, index; ∇=true))
        end
    end
    defn = optimize(arg_bufs, tile_fusion_factors, :(function $fn($(params...))
        $(body...)
    end))
    body = defn.args[2].args
    for (index, connection) in enumerate(ensemble.connections)
        if length(connection.shape) > 2 && any(collect(connection.shape[3:end]) .> 1)
            push!(body, :NOFUSE)
            break
        end
    end
    prepend!(compute_body, body)
    union!(compute_args, args)
end

"""
TODO: doc
"""
function generate_c_function(func::Function, signature::Tuple,
                             run_where::Int, signal::Array{Cint,1},
                             buffers::Dict)
    m = methods(func, signature)
    def = m[1].func.code
    func_ref = GlobalRef(def.module, symbol(string(func)))
    ast = code_typed(func, signature)[1]

    ast = remove_temp_nodes(ast)
    ast = transform_to_raw_array(ast)
    ast = collect_parallel_loop_private_vars(ast)

    # dir_ast = Driver.toDomainIR(func_ref, ast, signature)
    # pir_ast = Driver.toParallelIR(func_ref, dir_ast, signature)

    function_name_string = CGen.canonicalize(string(func_ref.name))
    for varinfo in ast.args[2][1]
        if varinfo[2] <: Array
            varinfo[2] = CGen.RawArray{eltype(varinfo[2]), ndims(varinfo[2])}
        end
    end
    s = ""
    @latte_mpi(s *= "#include \"comm.h\"\n")
    s *= CGen.from_root_entry(ast, function_name_string)
    proxy_name = string("_", function_name_string, "_j2c_proxy")
    proxy_sym = symbol(proxy_name)
    j2c_name = string("_", function_name_string, "_unaliased_")
    args = [:($(symbol(:__arg_,d))) for d in 1:length(signature)]
    if run_where >= 0
        @assert false
    end
    outfile_name = CGen.writec(s)
    cflags = []
    # push!(cflags, "-qopt-report")
    @latte_mpi push!(cflags, "-I$latte_library_path/communication")
    @latte_mpi ENV["CGEN_COMPILER"] = "mpiicpc"
    CGen.compile(outfile_name; flags=cflags)

    lflags = []
    @latte_mpi push!(lflags, "-L$latte_library_path")
    @latte_mpi push!(lflags, "-lLatteComm")
    dyn_lib = CGen.link(outfile_name; flags=lflags)

    proxy_params = [:($arg::$typ) for (arg, typ) in zip(args, signature)]
    sig_types = [:(Ptr{$(eltype(typ))}) for typ in signature]
    if run_where < 0
        j2c_typs = Expr(:tuple, Cint, sig_types..., Ptr{Cint})
        return @eval function $proxy_sym($(proxy_params...))
            ccall(($j2c_name, $dyn_lib), Void, $j2c_typs, $(run_where), $(args...), &(0))
        end
    else
        @assert false
    end
end

"""
TODO: doc
FIXME: My, my this function is ugly, clean this up some day...
"""
function push_compute_tasks!(tasks::TaskSet, buffers::Dict,
                             compute_body::Dict, compute_args::ArgSet,
                             run_where::Int, signal::Array{Cint, 1},
                             phases::Vector{Phase};
                             distribute::Bool=false)
    for phase in phases
        args = collect(compute_args[phase])
        if length(args) == 0
            continue
        end
        # BEGIN REPLACE ALIASED ARRAYS
        # TODO: Abstract this into a function
        # When two arrays x and y are aliased such that x === y, we replace all
        # references to y with x to assist the compiler with vectorization
        filtered_args = []
        arg_map = Dict{Symbol, Any}()
        for arg in args
            index = find((x) -> buffers[x] === buffers[arg], filtered_args)
            if length(index) == 0
                push!(filtered_args, arg)
            else
                @assert length(index) == 1
                arg_map[arg] = filtered_args[index[1]]
            end
        end
        compute_body[phase] = replace_symbols(compute_body[phase], arg_map)
        # END REPLACE ALIASED ARRAYS

        arg_bufs = [buffers[arg] for arg in filtered_args]
        signature = tuple([typeof(buf) for buf in arg_bufs]...)
        params = [:($arg::$typ) for (arg, typ) in zip(filtered_args, signature)]
        compute_body[phase] = clean_for_loops(compute_body[phase])
        # println(:(function $compute_task($(params...))
        #     $(compute_body[phase]...)
        # end))
        compute_body[phase] = fuse_loops(compute_body[phase])
        # println(:(function $compute_task($(params...))
        #     $(compute_body[phase]...)
        # end))
        # throw("Err")
        compute_body[phase] = inline_tile_size(compute_body[phase])
        compute_body[phase] = distribute_batch_loops(compute_body[phase])
        compute_body[phase] = parallelize_batch_tile_loops(compute_body[phase])
        compute_body[phase] = fuse_loops(compute_body[phase])

        is_update_param_call = (node) -> isa(node, Expr) && node.head == :call && node.args[1] == :update_param
        to_push = []
        for statement in compute_body[phase]
            if is_update_param_call(statement)
                if length(to_push) > 0
                    compute_task = gensym("compute_task")
                    func = :(function $compute_task($(params...))
                        $(to_push...)
                        return 0
                    end)
                    func = wrap_for_loops(func)
                    func = inline_array_expressions(func, buffers)
                    func = @eval $func
                    # println(code_typed(func, signature))
                    # throw("Err")

                    proxy_func = generate_c_function(func, signature, run_where,
                                                     signal, buffers)
                    proxy_args = arg_bufs

                    push!(tasks[phase], JuliaTask(proxy_func, tuple(filtered_args...)))
                end
                to_push = []
                push!(tasks[phase], UpdateTask(eval(statement.args[2])))
            else
                push!(to_push, statement)
            end
        end
        if length(to_push) > 0
            compute_task = gensym("compute_task")
            func = :(function $compute_task($(params...))
                $(to_push...)
                return 0
            end)
            func = wrap_for_loops(func)
            func = inline_array_expressions(func, buffers)
            func = @eval $func
            # println(code_typed(func, signature))
            # throw("Err")

            proxy_func = generate_c_function(func, signature, run_where,
                                             signal, buffers)
            proxy_args = arg_bufs

            push!(tasks[phase], JuliaTask(proxy_func, tuple(filtered_args...)))
        end
    end
end

"""
TODO: doc
"""
function add_forward_data_tasks(ensemble::DataEnsemble, tasks::TaskSet, net::Net)
    test_args = (ensemble, symbol(ensemble.name, :value), net, Test)
    train_args = (ensemble, symbol(ensemble.name, :value), net, Train)
    push!(tasks[Train], JuliaTask(forward, train_args))
    push!(tasks[Test], JuliaTask(forward, test_args))
end

"""
TODO: doc
"""
function add_forward_julia_tasks(ensemble::JuliaEnsemble, tasks::TaskSet, net::Net)
    inputs = []
    for connection in ensemble.connections
        push!(inputs, symbol(connection.source.name, :value))
    end
    test_args = (ensemble, inputs..., symbol(ensemble.name, :value), Test)
    train_args = (ensemble, inputs..., symbol(ensemble.name, :value), Train)
    push!(tasks[Train], JuliaTask(forward, train_args))
    push!(tasks[Test], JuliaTask(forward, test_args))
end

"""
TODO: doc
"""
function init_forward(ensemble::ReshapeEnsemble, net::Net,
                      compute_args::ArgSet, compute_body::Dict{Phase, Vector})
end

"""
TODO: doc
"""
function init_backward(ensemble::ReshapeEnsemble, net::Net,
                       compute_args::ArgSet, compute_body::Dict{Phase, Vector})
end

"""
TODO: doc
"""
function init_forward(ensemble::ConcatEnsemble, net::Net, compute_args::ArgSet,
                      compute_body::Dict{Phase, Vector})
    asts = []
    output_name = symbol(ensemble.name, :value)
    push!(compute_args[Train], output_name)
    push!(compute_args[Test], output_name)
    for (index, input) in enumerate(ensemble.inputs)
        input_name = symbol(input.name, :value)
        push!(compute_args[Train], input_name)
        push!(compute_args[Test], input_name)
        loopvars = []
        loopranges = []
        for i in size(input)
            push!(loopvars, gensym("loopvar"))
            push!(loopranges, :(1:$i))
        end
        # Add batch loop
        push!(loopvars, gensym("loopvar"))
        push!(loopranges, :(1:$(net.batch_size)))
        index_vars = Any[loopvars[1:end-1]...]
        if index > 1
            index_vars[end] = :($(index_vars[end]) + $((index - 1) * ensemble.inner_size))
        end

        push!(asts, gen_loop_nest(
            :($output_name[$(index_vars...), $(loopvars[end])] =
                $input_name[$(loopvars...)]),
            loopvars, loopranges))
    end
    for phase in [Train, Test]
        append!(compute_body[phase], asts)
    end
end

"""
TODO: doc
"""
function init_backward(ensemble::ConcatEnsemble, net::Net,
                       compute_args::ArgSet, compute_body::Dict{Phase, Vector})
    asts = []
    output_name = symbol(ensemble.name, :∇)
    push!(compute_args[Train], output_name)
    for (index, input) in enumerate(ensemble.inputs)
        input_name = symbol(input.name, :∇)
        push!(compute_args[Train], input_name)
        loopvars = []
        loopranges = []
        for i in size(input)
            push!(loopvars, gensym("loopvar"))
            push!(loopranges, :(1:$i))
        end
        # Add batch loop
        push!(loopvars, gensym("loopvar"))
        push!(loopranges, :(1:$(net.batch_size)))
        index_vars = Any[loopvars[1:end-1]...]
        if index > 1
            index_vars[end] = :($(index_vars[end]) + $((index - 1) * ensemble.inner_size))
        end

        push!(asts, gen_loop_nest(
            :($input_name[$(loopvars...)] =
              $output_name[$(index_vars...), $(loopvars[end])]),
            loopvars, loopranges))
    end
    append!(compute_body[Train], asts)
end

"""
TODO: doc
"""
function init_forward(ensemble::NormalizationEnsemble, net::Net,
                      compute_args::ArgSet, compute_body::Dict{Phase, Vector})
    args = get_forward_args(ensemble)
    union!(compute_args[ensemble.phase], args)

    for connection in ensemble.connections
        arg = symbol(connection.source.name,:value)
        push!(args, arg)
        push!(compute_args[ensemble.phase], arg)
    end

    ast = forward(ensemble)
    symbol_map = Dict{Symbol, Symbol}()
    for (param, arg) in zip(ast.args[1].args[2:end], args)
        if isa(param, Expr) && param.head == :(::)
            param = param.args[1]
        end
        symbol_map[param] = arg
    end
    body = AstWalk(ast.args[2], symbol_replacer, symbol_map)
    push!(compute_body[ensemble.phase], body)
end

"""
TODO: doc
"""
function init_backward(ensemble::NormalizationEnsemble, net::Net,
                       compute_args::ArgSet, compute_body::Dict{Phase, Vector})
    args = get_backward_args(ensemble)
    union!(compute_args[ensemble.phase], args)
    for connection in ensemble.connections
        # Don't backprop to label's ∇, but pass it's value
        if :label == connection.source.name
            arg = symbol(connection.source.name,:value)
        else
            arg = symbol(connection.source.name,:∇)
        end
        push!(args, arg)
        push!(compute_args[ensemble.phase], arg)
    end
    ast = backward(ensemble)
    symbol_map = Dict{Symbol, Symbol}()
    for (param, arg) in zip(ast.args[1].args[2:end], args)
        if isa(param, Expr) && param.head == :(::)
            param = param.args[1]
        end
        symbol_map[param] = arg
    end
    body = AstWalk(ast.args[2], symbol_replacer, symbol_map)
    unshift!(compute_body[ensemble.phase], body)
end

"""
TODO: doc
"""
function add_recv_expr(net::Net, source::AbstractEnsemble,
                       ensemble::AbstractEnsemble,
                       compute_body::Dict{Phase, Vector}, compute_args::ArgSet)
    key = symbol(source.name, :value)
    source_subgroup = source.net_subgroup - 1  # -1 for zero based indexing with MPI ranks
    tag = find(x -> x == ensemble, net.ensembles)[1]
    expr = :(ccall((:recv_intra, $libComm), Void,
                   (Ptr{Float32}, Cint, Cint, Cint),
                   pointer($key), length($key), $tag, $source_subgroup))
    if ensemble.phase in [Train, TrainTest] &&
            connection.source.phase in [Train, TrainTest]
        push!(compute_body[Train], expr)
        push!(compute_args[Train], key)
    end
    if ensemble.phase in [Test, TrainTest] &&
            connection.source.phase in [Train, TrainTest]
        push!(compute_body[Test], expr)
        push!(compute_args[Test], key)
    end
end

"""
TODO: doc
"""
function add_send_exprs(net::Net, ensemble::AbstractEnsemble,
                        compute_body::Dict{Phase, Vector},
                        compute_args::ArgSet)
    for (target, tag) in net.ensemble_send_list[ensemble.name]
        target = target - 1 # 0-based indexing for MPI
        target_phase = net.ensembles[tag].phase
        target_buf = symbol(ensemble.name, :value)
        expr = :(ccall((:send_intra, $libComm), Void,
                       (Ptr{Float32}, Cint, Cint, Cint),
                       pointer($target_buf), length($target_buf), $tag, $target))
        if ensemble.phase in [Train, TrainTest] &&
                target_phase in [Train, TrainTest]
            push!(compute_body[Train], expr)
            push!(compute_args[Train], target_buf)
        end
        if ensemble.phase in [Test, TrainTest] &&
                target_phase in [Test, TrainTest]
            push!(compute_body[Test], expr)
            push!(compute_args[Test], target_buf)
        end
    end
end

"""
Initialize the network `net`.

This function begins by initializing each `Ensemble` in the Network.  This is
done by calling `init(ens)` which will be dispatched to the appropriate
initialization routine.  These initialization routines are responsible for
adding buffers to the network to contain neuron fields and output values.  See
the specific initialization functions for different Ensemble types for more
information (TODO: Reference these).

After initializaing the local fields and output values for each ensemble, we
then initialize the input buffers for each ensemble.  This is done after all
output values have been initialized so that we can analyze recurrent
connections.  This is done by calling the `init_inputs(ens)` routine.
"""
function init(net::Net)
    log_info("Initializing net...")
    forward_tasks = net.forward_tasks
    backward_tasks = net.backward_tasks
    forward_compute_body = Dict{Phase, Vector}(Train => [], Test => [])
    forward_compute_args = ArgSet()
    forward_data_tasks   = TaskSet()
    backward_compute_body = Dict{Phase, Vector}(Train => [], Test => [])
    backward_compute_args = ArgSet()
    seen_names = Set()
    log_info("  Initializing ensembles.")
    for ensemble in net.ensembles
        # Check for duplicate ensemble names
        if ensemble.name in seen_names
            throw("Error: Found duplicate ensemble name: $(ensemble.name)")
        end
        push!(seen_names, ensemble.name)

        # If in MPI mode skip ensembles not assigned to this subrank
        @latte_mpi(if ensemble.net_subgroup != get_net_subrank(net) + 1
            continue  # skip
        end)
        net.ensemble_send_list[ensemble.name] = Tuple{Int, Int}[]

        log_info("    $(ensemble.name) size=$(size(ensemble))")

        # Initialize the ensemble
        init(ensemble, net)
        init_params(ensemble, net)
    end

    for (index, ensemble) in enumerate(net.ensembles)
        # If in MPI mode, populate the send_list for connected ensembles not in
        # this subrank and skip initializations
        @latte_mpi if ensemble.net_subgroup != get_net_subrank(net) + 1
            for connection in ensemble.connections
                if connection.source.net_subgroup == get_net_subrank(net) + 1
                    push!(net.ensemble_send_list[connection.source.name],
                          (ensemble.net_subgroup, index))
                end
            end
            continue  # skip
        end

        # Initialize the inputs for the ensemble, this is done after all
        # ensembles have initialized their outputs (necessary for rnns)
        init_inputs(ensemble, net)
    end
    log_info("  Finished initializing ensembles.")

    log_info("  Synthesizing forward functions.")
    # Generate forward tasks
    for ensemble in net.ensembles
        @latte_mpi if ensemble.net_subgroup != get_net_subrank(net) + 1
            continue  # skip
        end

        @latte_mpi for connection in ensemble.connections
            if connection.source.net_subgroup != ensemble.net_subgroup
                add_recv_expr(net, connection.source, ensemble,
                              forward_compute_body, forward_compute_args)

            end
        end
        # Skip code generation for data and loss ensembles
        if typeof(ensemble) <: DataEnsemble
            add_forward_data_tasks(ensemble, forward_data_tasks, net)
        elseif typeof(ensemble) <: JuliaEnsemble
            add_forward_julia_tasks(ensemble, forward_data_tasks, net)
        elseif isa(ensemble, Union{NormalizationEnsemble, ConcatEnsemble, ReshapeEnsemble})
            init_forward(ensemble, net, forward_compute_args, forward_compute_body)
        elseif isa(ensemble, Union{Ensemble, ActivationEnsemble})
            gen_neuron_forward(ensemble, net, forward_compute_body, forward_compute_args)
        else
            throw("Latte Error: Encountered unsupported ensemble type $(typeof(ensemble)).")
        end

        @latte_mpi add_send_exprs(net, ensemble, forward_compute_body,
                                  forward_compute_args)
    end
    append!(net.forward_tasks, forward_data_tasks)

    push_compute_tasks!(forward_tasks, net.buffers[1][1],
                        forward_compute_body, forward_compute_args,
                        net.run_where, net.signal, [Train, Test]; distribute=true)

    log_info("  Synthesizing backward functions.")
    # Backward tasks
    for ensemble in net.ensembles
        @latte_mpi (if ensemble.net_subgroup != get_net_subrank(net) + 1
            continue
        end)
        if typeof(ensemble) <: DataEnsemble || ensemble.phase == Test
            continue  # skip
        end
        for (target, tag) in net.ensemble_send_list[ensemble.name]
            target_phase = net.ensembles[tag].phase
            if !(target_phase in [Train, TrainTest])
                continue
            end
            target = target - 1 # 0-based indexing for MPI
            target_buf = symbol(ensemble.name, :∇)
            tag += length(net.ensembles)
            expr = :(ccall((:recv_intra, $libComm), Void,
                           (Ptr{Float32}, Cint, Cint, Cint),
                           pointer($target_buf), length($target_buf), $tag, $target))
            push!(backward_compute_body[Train], expr)
            push!(backward_compute_args[Train], target_buf)
        end
        if isa(ensemble, Union{NormalizationEnsemble, ConcatEnsemble, ReshapeEnsemble})
            init_backward(ensemble, net, backward_compute_args, backward_compute_body)
        elseif typeof(ensemble) <: JuliaEnsemble
            # Skip
            # throw("NotImplementedError")
        elseif isa(ensemble, Union{Ensemble, ActivationEnsemble})
            for param in ensemble.params
                @latte_mpi(begin
                    if LOSSY_GRADIENTS
                        gradient_length = length(param.gradient)
                        reduce_num = -1
                    else
                        gradient_length = length(param.gradient) / num_threads
                        reduce_num = num_threads
                    end
                    unshift!(backward_compute_body[Train], quote
                        ccall((:sync_gradients, $libComm), Void,
                              (Ptr{Float32}, Cint, Cint),
                              pointer($(param.gradient_name)),
                              $(gradient_length),
                              $(param.request), $reduce_num)
                    end)
                end)
                push!(net.params, param)
            end
            gen_neuron_backward(ensemble, net, backward_compute_body[Train], backward_compute_args[Train])
        else
            throw("Latte Error: Encountered unsupported ensemble type $(typeof(ensemble)).")
        end
        for connection in ensemble.connections
            if :label == connection.source.name
                # Skip
                continue
            end
            if connection.source.net_subgroup != ensemble.net_subgroup
                key = symbol(connection.source.name, :∇)
                source_subgroup = connection.source.net_subgroup - 1  # 0-based indexing for MPI
                tag = find(x -> x == ensemble, net.ensembles)[1] + length(net.ensembles)
                push!(backward_compute_args[Train], key)
                push!(backward_compute_body[Train],
                    :(ccall((:send_intra, $libComm), Void,
                            (Ptr{Float32}, Cint, Cint, Cint),
                            pointer($key), length($key), $tag, $source_subgroup)))
            end
        end
    end

    push_compute_tasks!(backward_tasks, net.buffers[1][1], backward_compute_body,
                        backward_compute_args, net.run_where, net.signal, [Train])
    if net.run_where >= 0
        @assert false
    end
    log_info("Initialization finished.")
end

function get_task_args(net, task_args, t)
    args = []
    for arg in task_args
        if isa(arg, Symbol)
            push!(args, get_buffer(net, arg, t))
        else
            push!(args, arg)
        end
    end
    args
end

# Use metaprogramming to generate single and multi versions of forward
# and backward.
for direction in [:forward, :backward]
    tasks = symbol(direction,:_tasks)
    @eval function $direction(net::Net; phase=Train, solver=nothing)
        for t = 1:net.time_steps
            net.curr_time_step = t
            for task in net.$tasks[phase]
                if isa(task, JuliaTask)
                    task.func(get_task_args(net, task.args, t)...)
                elseif isa(task, UpdateTask)
                    @assert phase == Train && solver != nothing
                    update(solver, net, task.param_id)
                else
                    throw("Unsupported task type $(typeof(task))")
                end
            end
        end
    end
end

"""
Add an ensemble to the network `net`
"""
function add_ensemble(net::Net, ens::AbstractEnsemble)
    push!(net.ensembles, ens)
    net.ensembles_map[ens.name] = ens
end

"""
TODO: doc
"""
function check_dimensions_fixed(mapping::Function, sink_shape::Tuple)
    n = length(sink_shape)
    is_dim_fixed = [true for _ in 1:n]
    first = mapping(ones(Int, n)...)
    if !all(map((x) -> isa(x, UnitRange) || isa(x, Colon), first))
        is_dim_fixed = [false for _ in 1:n]
    else
        for d in 1:n
            for i in 1:sink_shape[d]
                idx = ones(Int, n)
                idx[d] = i
                if first != mapping(idx...)
                    is_dim_fixed[d] = false
                    break
                end
            end
        end
    end
    is_dim_fixed
end

"""
TODO: doc
"""
function check_one_to_one(mapping::Function, shape::Tuple)
    is_one_to_one = true
    for i in CartesianRange(shape)
        if mapping(i.I...) != i.I
            is_one_to_one = false
            break
        end
    end
    is_one_to_one
end

"""
Connect neurons in `source` to neurons in `sink` using the function `mapping`.
`mapping` should be a function with a parameter for the index in each dimension
of sink.
For example, if sink is a 3-d ensemble, mapping = f(i, j, k) -> ...

`mapping` should return a tuple of continuous ranges corresponding to the
indices of neurons in source that should be connected to the neuron at the
current index
"""
function add_connections(net::Net, source::AbstractEnsemble,
                         sink::AbstractEnsemble, mapping::Function; padding=0,
                         recurrent=false)
    n = ndims(sink)

    # Compute the size and shape of the connection
    range_size = 1
    range_shape = []
    for (index, d) in enumerate(mapping(ones(Int, n)...))
        if isa(d, Colon)
            range_size *= size(source)[index]
            push!(range_shape, size(source)[index])
        else
            range_size *= length(d)
            push!(range_shape, length(d))
        end
    end

    # Determine if any dimensions are fixed
    is_dim_fixed = check_dimensions_fixed(mapping, size(sink))
    is_one_to_one = false
    if !all(is_dim_fixed)
        is_one_to_one = check_one_to_one(mapping, size(sink))
    end
    push!(sink.connections, Connection(source, mapping, tuple(range_shape...),
                                       range_size, true, is_dim_fixed,
                                       is_one_to_one, padding, recurrent))
end

"""
Test `net` for one epoch
"""
function test(net::Net)
    curr_epoch = net.test_epoch
    accuracy = 0.0f0
    num_batches = 0.0f0
    while net.test_epoch == curr_epoch
        forward(net, phase=Test)
        num_batches += 1.0f0
        if haskey(net.buffers[net.curr_buffer_set][1], :accuracyvalue)
            accuracy += get_buffer(net, :accuracyvalue)[1]
        end
        clear_values(net)
        @latte_mpi sync_intra_test_epoch(net)
    end
    accuracy / num_batches * 100.0f0
end

export save_snapshot, load_snapshot

"""
Save a snapshot of `net` to `file`
"""
function save_snapshot(net::Net, file::AbstractString)
    param_dict = Dict{Symbol, Vector{Param}}()
    for ens in net.ensembles
        if :params in fieldnames(ens) && length(ens.params) > 0
            param_dict[ens.name] = ens.params
        end
    end
    save(file, "param_dict", param_dict)
end

"""
Load a network snapshot from `file`.

TODO: Can we save the structure of `net` in the snapshot?
"""
function load_snapshot(net::Net, file::AbstractString)
    param_dict = load(file, "param_dict")
    for ens in net.ensembles
        if :params in fieldnames(ens) && length(ens.params) > 0
            for (param, saved) in zip(ens.params, param_dict[ens.name])
                param.value[:] = saved.value[:]
            end
        end
    end
end

"""
Get the current loss for `net`

TODO: This is not general, assumes one :loss ensemble
"""
function get_loss(net::Net)
    @latte_mpi(if haskey(net.buffers[net.curr_buffer_set][1], :lossvalue)
        loss = get_buffer(net, :lossvalue)[1]
        sync_intra_loss(net, loss)
        return loss
    else
        return sync_intra_loss(net, 0.0f0)
    end,
    return get_buffer(net, :lossvalue)[1])
end

"""
Fill buffers with names containing `∇` with zeros.
"""
function clear_∇(net::Net)
    i = net.curr_buffer_set
    for t in 1:net.time_steps
        for name in keys(net.buffers[i][t])
            if contains(string(name), "∇")
                fill!(net.buffers[i][t][name], 0.0)
            end
        end
    end
end

"""
Fill buffers with names containing `value` with zeros.
"""
function clear_values(net::Net)
    i = net.curr_buffer_set
    for t in 1:net.time_steps
        for name in keys(net.buffers[i][t])
            if contains(string(name), "value")
                fill!(net.buffers[i][t][name], 0.0)
            end
        end
    end
end

"""
Fill buffers with names containing `randval` with random values
"""
function rand_values(net::Net)
    i = net.curr_buffer_set
    for t in 1:net.time_steps
        for name in keys(net.buffers[i][t])
            if contains(string(name), "randval")
                rand!(net.buffers[i][t][name])
            end
        end
    end
end

"""
Get a buffer associated with an ensemble

** Params **

- `net`   -- network to get buffer
- `ens`   -- the ensemble
- `name`  -- name of buffer associated with `ens`
"""
function get_buffer(net::Net, ens::AbstractEnsemble, name::Symbol)
    return net.buffers[net.curr_buffer_set][1][symbol(ens.name, name)]
end

"""
Get a buffer at the time_step `t`

** Params **

- `net`   -- network to get buffer
- `name`  -- name of the buffer
- `t`     -- time step
"""
function get_buffer(net::Net, name::Symbol, t::Int=1)
    return net.buffers[net.curr_buffer_set][t][name]
end

"""
Add or update a buffer at a particular time step `t`

** Params **

- `net`   -- network to add/update buffer
- `name`  -- name of the buffer
- `arr`   -- buffer
- `t`     -- time step to add buffer
"""
function set_buffer(net::Net, name::Symbol, arr::Array, t::Int)
    net.buffers[1][t][name] = arr
    net.buffers[2][t][name] = arr
end

"""
Add or update a buffer

** Params **

- `net`   -- network to add/update buffer
- `name`  -- name of the buffer
- `arr`   -- buffer
- `_copy` -- whether to copy the buffer
"""
function set_buffer(net::Net, name::Symbol, arr::Array; _copy=true)
    for t in 1:net.time_steps
        if _copy
            net.buffers[1][t][name] = copy(arr)
            net.buffers[2][t][name] = copy(arr)
        else
            net.buffers[1][t][name] = arr
            net.buffers[2][t][name] = arr
        end
    end
end

"""
Initialize a buffer in `net`

** Params **

- `net`   -- network to receive initialized buffer
- `name`  -- name of the buffer
- `shape` -- shape of the buffer
- `func`  -- function used to initialize the buffer, should return an Array and
             have a signature (Float32, dims...)
"""
function init_buffer(net::Net, name::Symbol, shape; func=zeros)
    for t in 1:net.time_steps
        net.buffers[1][t][name] = func(Float32, shape...)
        net.buffers[2][t][name] = func(Float32, shape...)
    end
end


text:
function projected_gradient(
    nlp,
    uk;
    max_iter=1000,
    ls_itermax=30,
    α0=1.0,
    β=0.4,
    τ=1e-4,
    tol=1e-5,
    bfgs=false,
    verbose_it=100,
)

    u_prev = copy(uk)
    grad = copy(uk)
    wk = copy(uk)
    u♭ = nlp.inner.u_min
    u♯ = nlp.inner.u_max
    if bfgs
        H = InverseLBFGSOperator(Float64, length(uk), 50, scaling=true)
    else
        H = I
    end
    norm_grad = Inf
    n_iter = 0

    for i in 1:max_iter
        n_iter += 1
        # solve power flow and compute gradients
        ExaPF.update!(nlp, uk)

        # evaluate cost
        c = ExaPF.objective(nlp, uk)
        # Evaluate cost of problem without penalties
        c_ref = nlp.scaler.scale_obj * ExaPF.objective(nlp.inner, uk)
        ExaPF.gradient!(nlp, grad, uk)

        # compute control step
        # Armijo line-search (Bertsekas, 1976)
        dk = H * grad
        step = α0
        for j_ls in 1:ls_itermax
            step *= β
            ExaPF.project!(wk, uk .- step .* dk, u♭, u♯)
            ExaPF.update!(nlp, wk)
            ft = ExaPF.objective(nlp, wk)
            if ft <= c - τ * dot(dk, wk .- uk)
                break
            end
        end

        # step = αi
        wk .= uk .- step * dk
        ExaPF.project!(uk, wk, u♭, u♯)

        # Stopping criteration: uₖ₊₁ - uₖ
        ## Dual infeasibility
        norm_grad = norm(uk .- u_prev, Inf)
        ## Primal infeasibility
        inf_pr = ExaPF.primal_infeasibility(nlp.inner, nlp.cons)

        # check convergence
        if (i % verbose_it == 0)
            @printf("%6d %.6e %.3e %.2e %.2e %.2e\n", i, c, c - c_ref, norm_grad, inf_pr, step)
        end

        if bfgs
            push!(H, uk .- u_prev, grad .- grad_prev)
            grad_prev .= grad
        end
        u_prev .= uk
        # Check whether we have converged nicely
        if (norm_grad < tol)
            converged = true
            break
        end
    end
    return uk, norm_grad, n_iter
end

function ls_reference(costs::AbstractVector, lm_size::Int)
    lag = max(1, length(costs) - lm_size + 1)
    return maximum(costs[lag:end])
end
# Non-monotone gradient projection algorithm
function ngpa(
    nlp,
    uk;
    max_iter=1000,
    ls_itermax=30,
    α_bb=1.0,
    α♭=0.0,
    α♯=1.0,
    β=0.4,
    δ=1e-4,
    tol=1e-5,
    lm_size=1,
    verbose_it=100,
)

    u_prev = copy(uk)
    grad = copy(uk)
    wk = copy(uk)
    u♭ = nlp.inner.u_min
    u♯ = nlp.inner.u_max

    # Initial evaluation
    ExaPF.update!(nlp, uk)
    c = ExaPF.objective(nlp, uk)
    ExaPF.gradient!(nlp, grad, uk)

    # Memory
    grad_prev = copy(grad)
    inf_pr = ExaPF.primal_infeasibility(nlp.inner, nlp.cons ./ nlp.scaler.scale_cons)
    costs_history = Float64[c]
    grad_history = Float64[norm(grad, Inf)]
    feas_history = Float64[inf_pr]

    norm_grad = Inf
    n_iter = 0
    ## Line-search params
    j_bb = 0
    flag_bb = 1
    θ_bb = 0.975
    m_bb = 10
    ## Reference function params
    Δ_ref = 0.1
    A_ref = 0
    L_ref = 0
    a_ref = 0
    l_ref = 0
    fᵣ = c

    n_up = 0
    for i in 1:max_iter
        n_iter += 1

        ExaPF.project!(wk, uk .- α_bb .* grad, u♭, u♯)
        # Feasible direction
        dk = wk .- uk
        # Armijo line-search
        step = 1.0
        for j_ls in 1:ls_itermax
            ExaPF.project!(wk, uk .+ step .* dk, u♭, u♯)
            conv = ExaPF.update!(nlp, wk)
            ft = ExaPF.objective(nlp, wk)
            if ft <= fᵣ + step * δ * dot(dk, grad)
                break
            end
            step *= β
        end

        uk .= wk
        # Objective
        c = ExaPF.objective(nlp, uk)
        c_ref = ExaPF.inner_objective(nlp, uk)
        # Gradient
        ExaPF.gradient!(nlp, grad, uk)

        # Stopping criteration: uₖ₊₁ - uₖ
        ## Dual infeasibility
        norm_grad = norm(uk .- u_prev, Inf)
        ## Primal infeasibility
        inf_pr = ExaPF.primal_infeasibility(nlp.inner, nlp.cons ./ nlp.scaler.scale_cons)

        # check convergence
        if (i % verbose_it == 0)
            @printf("%6d %.6e %.3e %.2e %.2e %.2e\n", i, c, c - c_ref, norm_grad, inf_pr, step)
        end

        ##################################################
        ## Update parameters
        sk = uk - u_prev
        yk = grad - grad_prev
        ##################################################
        ## Update Barzilai-Borwein step
        flag_bb = 0
        if !isnothing(findfirst(0.0 .< abs.(dk) .< α_bb .* abs.(grad)))
            flag_bb = 1
        end
        if step == 1.0
            j_bb += 1
        else
            flag_bb = 1
        end
        θ = dot(sk, yk) / (norm(sk) * norm(yk))
        if (j_bb >= m_bb) || (θ >= θ_bb) || (flag_bb == 1)
            # Non-convexity detected
            if dot(sk, yk) <= 0.0
                if j_bb >= 1.5 * m_bb
                    t_bb = min(norm(uk, Inf), 1) / norm(dk, Inf)
                    α_bb = max(t, step)
                    j_bb = 0
                end
            else
                # Everything went ok. Set new step to Barzilai-Borwein value
                α_bb = dot(sk, sk) / dot(sk, yk)
                j_bb = 0
            end
        end
        α_bb = min(α♯, α_bb)
        # Update history
        u_prev .= uk
        grad_prev .= grad
        push!(costs_history, c)
        push!(grad_history, norm_grad)
        push!(feas_history, inf_pr)
        ##################################################
        ## Update reference value
        fr_max = ls_reference(costs_history, lm_size)
        fᵣ = .5 * (fr_max + c)
        # if j > 0
        #     fᵣ = min(fr_max, fᵣ)
        # end
        # a_ref = (step < 1) ? 0 : a_ref + 1
        # if kkkk

        # Check whether we have converged nicely
        if (norm_grad < tol)
            converged = true
            break
        end
    end
    return uk, norm_grad, n_iter, costs_history, grad_history, feas_history
end

ext: jl


text:
function t_components()

########################################
# Test adjacency-related functionality #
########################################

# Define adjacency data
A = [							# adjacency matrix
 0.0  0.0  1.0  1.0  0.0;
 0.0  0.0  1.0  0.0  0.0;
 1.0  1.0  0.0  1.0  1.0;
 1.0  0.0  1.0  0.0  0.0;
 0.0  0.0  1.0  0.0  0.0;
]

Ai = Int.(A)
G = Graph(Ai)

Gw = SimpleWeightedGraph(A)
I = findall(!iszero, A)
rows, cols = getindex.(I, 1), getindex.(I, 2)

data = [(i,j) for (i,j) in zip(rows,cols)]		# edges of A
n = 5							# number of vertices
f(n::Int, T=Int) = (data::Vector{Tuple{Int,Int}})->begin
	m = zeros(T,n,n)
	for t in data
		m[t[1],t[2]] = one(T)
	end
	return m
end

# Test functionality
E_Adj = adjacency()
Test.@test E_Adj isa EmptyAdjacency
Test.@test adjacency(nothing) isa EmptyAdjacency

A_Adj = adjacency(Ai)
Test.@test A_Adj isa MatrixAdjacency
Test.@test A_Adj.am == A

G_Adj = adjacency(G)
Test.@test G_Adj isa GraphAdjacency
Test.@test G_Adj.ag == G

Gw_Adj = adjacency(Gw)
Test.@test Gw_Adj isa GraphAdjacency
Test.@test Gw_Adj.ag == Gw

C_Adj = adjacency(f(n), data)
Test.@test C_Adj isa ComputableAdjacency
Test.@test C_Adj.f == f(n)
Test.@test C_Adj.data == data

C2_Adj = adjacency((f(n), data))
Test.@test C2_Adj isa ComputableAdjacency
Test.@test C2_Adj.f == f(n)
Test.@test C2_Adj.data == data

P_Adj = adjacency(f(n))
Test.@test P_Adj isa PartialAdjacency
Test.@test P_Adj.f == f(n)
Test.@test adjacency(P_Adj, data).am == A

A_Adj2 = adjacency(A_Adj)
Test.@test A_Adj2 == A_Adj

Test.@test adjacency(strip_adjacency(A_Adj),A).am == A_Adj.am
Test.@test adjacency(strip_adjacency(G_Adj),A).ag == G_Adj.ag
Test.@test adjacency(strip_adjacency(C_Adj),data).am == A_Adj.am
Test.@test adjacency(strip_adjacency(P_Adj),data).am == A_Adj.am
Test.@test adjacency(strip_adjacency(E_Adj),A).am == A_Adj.am

Test.@test adjacency_graph(A_Adj) == G
Test.@test adjacency_graph(G_Adj) == G
Test.@test adjacency_graph(Gw_Adj) == Gw
Test.@test adjacency_graph(C_Adj) == G

Test.@test adjacency_matrix(A_Adj) == Ai
Test.@test adjacency_matrix(G_Adj) == Ai
Test.@test adjacency_matrix(Gw_Adj) == A
Test.@test adjacency_matrix(C_Adj) == A

r=1:1
As = sparse(A)
Test.@test adjacency_obs(A, r, ObsDim.Constant{1}()) == view(A,r,:)
Test.@test adjacency_obs(A, r, ObsDim.Constant{2}()) == view(A,:,r)
Test.@test adjacency_obs(As, r, ObsDim.Constant{1}()) == As[r,:]
Test.@test adjacency_obs(As, r, ObsDim.Constant{2}()) == As[:,r]

# Test show methods
buf = IOBuffer()
Test.@test try
	for Adj in [A_Adj, G_Adj, C_Adj, P_Adj, E_Adj]
		show(buf,Adj)
	end
	true
catch
	false
end

# Test update_adjacency
A = [0 1 0; 1 0 0; 0 0 0]
Adj_m = adjacency(A)
update_function_m!(X,x,y) = begin
	X[x,y] += 1
	X[y,x] += 1
	return X
end
f_update_m(x,y) = X->update_function_m!(X,x,y)
for i in 1:3
	update_adjacency!(Adj_m, f_update_m(1,3))
end
Test.@test adjacency_matrix(Adj_m) == [0 1 3; 1 0 0; 3 0  0]

Adj_g = adjacency(Graph(A))
f_update_g(x,y) = G->add_edge!(G,x,y)
update_adjacency!(Adj_g, f_update_g(1,3))
Test.@test Matrix(adjacency_matrix(Adj_g)) == [0 1 1; 1 0 0; 1 0  0]

# Test errors
for Adj in [P_Adj, E_Adj]
	Test.@test try
		update_adjacency(Adj)
		false
	catch
		true # must fail because of the eror
	end

	Test.@test try
		adjacency_graph(Adj)
		false
	catch
		true # must fail because of the eror
	end
	Test.@test try
		adjacency_matrix(Adj)
		false
	catch
		true # must fail because of the eror
	end
end

######################################################
# Test fit/transform methods for relational learners #
######################################################
LEARNER = [SimpleRN,
	   WeightedRN,
	   BayesRN,
	   ClassDistributionRN]
N = 5							# number of observations
C = 2; 							# number of classes
A = [							# adjacency matrix
 0.0  0.0  1.0  1.0  0.0;
 0.0  0.0  1.0  0.0  0.0;
 1.0  1.0  0.0  1.0  1.0;
 1.0  0.0  1.0  0.0  0.0;
 0.0  0.0  1.0  0.0  0.0;
]
Ad = adjacency(A);

X = [							# local model estimates (2 classes)
 1.0  1.0  1.0  0.0  0.0;
 0.5  1.0  0.0  1.5  0.0
]

y = [1, 1, 1, 2, 2]					# labels

result = [
 [0.5  1.0  0.5  1.0  1.0; 				# validation data for SimpleRN
  0.5  0.0  0.5  0.0  0.0]
,
 [0.5  1.0  0.5   1.0   1.0;				# validation data for WeightedRN
  0.75 0.0  0.75  0.25  0.0]
,
 [1.60199  1.51083  1.60199  1.51083  1.51083;		# validation data for BayesRN
  1.14384  1.28768  1.14384  1.28768  1.28768]
,
 [0.300463  0.600925  0.300463  0.416667  0.600925; 	# validation data for ClassDistributionRN
  0.800391  0.125     0.800391  0.125     0.125]
]

Xo = zeros(size(X));					# ouput (same size as X)
Xon = zeros(size(X));					# normalized ouput (same size as X)

tol = 1e-5
for li in 1:length(LEARNER)
	rl = fit(LEARNER[li], Ad, X, y; priors=ones(length(unique(y))),normalize=false)
	rln = fit(LEARNER[li], Ad, X, y; priors=ones(length(unique(y))),normalize=true)
        transform!(Xo, rl, Ad, X, y);
        transform!(Xon, rln, Ad, X, y);
	Test.@test all(abs.(Xo - result[li]) .<= tol);	# external validation
	Test.@test Xon ≈ (Xo./sum(Xo, dims=1))		# normalization validation
end

# Test show methods
for li in 1:length(LEARNER)
	rl = fit(LEARNER[li], Ad, X', y; obsdim=ObsDim.Constant{1}())
	Test.@test try
		show(IOBuffer(), rl)
		true
	catch
		false
	end
	rl = fit(LEARNER[li], Ad, X, y; obsdim=ObsDim.Constant{2}())
	Test.@test try
		show(IOBuffer(), rl)
		true
	catch
		false
	end
end

###################################
# Tests for the utility functions #
###################################

# Test observation-related functions i.e. oppdim, intdim, matrix_prealloc
r_oppdim = [ObsDim.Constant{2}(), ObsDim.First(), ObsDim.Constant{2}(), ObsDim.Constant{1}()]
r_intdim = [1,2,1,2]

tv = [ObsDim.First(), ObsDim.Last(), ObsDim.Constant{1}(), ObsDim.Constant{2}()]
for (i,o) in enumerate(tv)
	Test.@test oppdim(o) == r_oppdim[i]
	Test.@test intdim(o) == r_intdim[i]
end

Test.@test matrix_prealloc(10,2,ObsDim.Constant{1}(),1) == ones(10,2)
Test.@test matrix_prealloc(10,2,ObsDim.Constant{2}(),1) == ones(2,10)
Test.@test matrix_prealloc(10,2,ObsDim.First(),1) == ones(10,2)
Test.@test matrix_prealloc(10,2,ObsDim.Last(),1) == ones(2,10)
Test.@test try
	matrix_prealloc(10,2,ObsDim.Undefined(),1)
	false
catch
	true
end

# Test the rest ...
cited =  [1,1,2,2,3,3,4,6,5,6]
citing = [2,3,1,5,2,6,7,1,2,3]
useidx = [1,2,6]

# 1 and 2 cite each other (edge weight of 2), 1 cites 6 one time (edge weight of 1)
Test.@test NetworkLearning.generate_partial_adjacency(cited,citing,useidx) == [0.0 2 1; 2 0 0; 1 0 0]

Test.@test NetworkLearning.get_size_out([1.,2,3]) == 1
Test.@test NetworkLearning.get_size_out([1,2,3]) == 3
Test.@test try
	NetworkLearning.get_size_out(rand(2,1))
	false
catch
	true
end

Test.@test NetworkLearning.getpriors([1.,2,3]) == [1.0]
Test.@test NetworkLearning.getpriors([1,2,3]) == 1/3*ones(3)
Test.@test try
	NetworkLearning.getpriors(rand(2,1))
	false
catch
	true
end

end


text:
int_rules_6_5_11 = @theory begin

    #= ::Subsection::Closed:: =#
    #= 6.5.11*(e*x)^m*(a+b*sech(c+d*x^n))^p =#
    @apply_utils Antiderivative((~(a') + ~(b') * sech(~(c') + ~(d') * (~x) ^ ~n)) ^ ~(p'), ~x) => (1 / ~n) * Subst(Antiderivative((~x) ^ (1 / ~n - 1) * (~a + ~b * sech(~c + ~d * ~x)) ^ ~p, ~x), ~x, (~x) ^ ~n) <-- FreeQ([~a, ~b, ~c, ~d, ~p], ~x) && (IGtQ(1 / ~n, 0) && IntegerQ(~p))
    @apply_utils Antiderivative((~(a') + ~(b') * csch(~(c') + ~(d') * (~x) ^ ~n)) ^ ~(p'), ~x) => (1 / ~n) * Subst(Antiderivative((~x) ^ (1 / ~n - 1) * (~a + ~b * csch(~c + ~d * ~x)) ^ ~p, ~x), ~x, (~x) ^ ~n) <-- FreeQ([~a, ~b, ~c, ~d, ~p], ~x) && (IGtQ(1 / ~n, 0) && IntegerQ(~p))
    @apply_utils Antiderivative((~(a') + ~(b') * sech(~(c') + ~(d') * (~x) ^ ~n)) ^ ~(p'), ~x) => Unintegrable((~a + ~b * sech(~c + ~d * (~x) ^ ~n)) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~n, ~p], ~x)
    @apply_utils Antiderivative((~(a') + ~(b') * csch(~(c') + ~(d') * (~x) ^ ~n)) ^ ~(p'), ~x) => Unintegrable((~a + ~b * csch(~c + ~d * (~x) ^ ~n)) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~n, ~p], ~x)
    @apply_utils Antiderivative((~(a') + ~(b') * sech(~(c') + ~(d') * (~u) ^ ~n)) ^ ~(p'), ~x) => (1 / Coefficient(~u, ~x, 1)) * Subst(Antiderivative((~a + ~b * sech(~c + ~d * (~x) ^ ~n)) ^ ~p, ~x), ~x, ~u) <-- FreeQ([~a, ~b, ~c, ~d, ~n, ~p], ~x) && (LinearQ(~u, ~x) && NeQ(~u, ~x))
    @apply_utils Antiderivative((~(a') + ~(b') * csch(~(c') + ~(d') * (~u) ^ ~n)) ^ ~(p'), ~x) => (1 / Coefficient(~u, ~x, 1)) * Subst(Antiderivative((~a + ~b * csch(~c + ~d * (~x) ^ ~n)) ^ ~p, ~x), ~x, ~u) <-- FreeQ([~a, ~b, ~c, ~d, ~n, ~p], ~x) && (LinearQ(~u, ~x) && NeQ(~u, ~x))
    @apply_utils Antiderivative((~(a') + ~(b') * sech(~u)) ^ ~(p'), ~x) => Antiderivative((~a + ~b * sech(ExpandToSum(~u, ~x))) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~p], ~x) && (BinomialQ(~u, ~x) && Not(BinomialMatchQ(~u, ~x)))
    @apply_utils Antiderivative((~(a') + ~(b') * csch(~u)) ^ ~(p'), ~x) => Antiderivative((~a + ~b * csch(ExpandToSum(~u, ~x))) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~p], ~x) && (BinomialQ(~u, ~x) && Not(BinomialMatchQ(~u, ~x)))
    @apply_utils Antiderivative((~x) ^ ~(m') * (~(a') + ~(b') * sech(~(c') + ~(d') * (~x) ^ ~n)) ^ ~(p'), ~x) => (1 / ~n) * Subst(Antiderivative((~x) ^ (Simplify((~m + 1) / ~n) - 1) * (~a + ~b * sech(~c + ~d * ~x)) ^ ~p, ~x), ~x, (~x) ^ ~n) <-- FreeQ([~a, ~b, ~c, ~d, ~m, ~n, ~p], ~x) && (IGtQ(Simplify((~m + 1) / ~n), 0) && IntegerQ(~p))
    @apply_utils Antiderivative((~x) ^ ~(m') * (~(a') + ~(b') * csch(~(c') + ~(d') * (~x) ^ ~n)) ^ ~(p'), ~x) => (1 / ~n) * Subst(Antiderivative((~x) ^ (Simplify((~m + 1) / ~n) - 1) * (~a + ~b * csch(~c + ~d * ~x)) ^ ~p, ~x), ~x, (~x) ^ ~n) <-- FreeQ([~a, ~b, ~c, ~d, ~m, ~n, ~p], ~x) && (IGtQ(Simplify((~m + 1) / ~n), 0) && IntegerQ(~p))
    @apply_utils Antiderivative((~x) ^ ~(m') * (~(a') + ~(b') * sech(~(c') + ~(d') * (~x) ^ ~n)) ^ ~(p'), ~x) => Unintegrable((~x) ^ ~m * (~a + ~b * sech(~c + ~d * (~x) ^ ~n)) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~m, ~n, ~p], ~x)
    @apply_utils Antiderivative((~x) ^ ~(m') * (~(a') + ~(b') * csch(~(c') + ~(d') * (~x) ^ ~n)) ^ ~(p'), ~x) => Unintegrable((~x) ^ ~m * (~a + ~b * csch(~c + ~d * (~x) ^ ~n)) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~m, ~n, ~p], ~x)
    @apply_utils Antiderivative((~e * ~x) ^ ~(m') * (~(a') + ~(b') * sech(~(c') + ~(d') * (~x) ^ ~n)) ^ ~(p'), ~x) => (((~e) ^ IntPart(~m) * (~e * ~x) ^ FracPart(~m)) / (~x) ^ FracPart(~m)) * Antiderivative((~x) ^ ~m * (~a + ~b * sech(~c + ~d * (~x) ^ ~n)) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~e, ~m, ~n, ~p], ~x)
    @apply_utils Antiderivative((~e * ~x) ^ ~(m') * (~(a') + ~(b') * csch(~(c') + ~(d') * (~x) ^ ~n)) ^ ~(p'), ~x) => (((~e) ^ IntPart(~m) * (~e * ~x) ^ FracPart(~m)) / (~x) ^ FracPart(~m)) * Antiderivative((~x) ^ ~m * (~a + ~b * csch(~c + ~d * (~x) ^ ~n)) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~e, ~m, ~n, ~p], ~x)
    @apply_utils Antiderivative((~e * ~x) ^ ~(m') * (~(a') + ~(b') * sech(~u)) ^ ~(p'), ~x) => Antiderivative((~e * ~x) ^ ~m * (~a + ~b * sech(ExpandToSum(~u, ~x))) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~e, ~m, ~p], ~x) && (BinomialQ(~u, ~x) && Not(BinomialMatchQ(~u, ~x)))
    @apply_utils Antiderivative((~e * ~x) ^ ~(m') * (~(a') + ~(b') * csch(~u)) ^ ~(p'), ~x) => Antiderivative((~e * ~x) ^ ~m * (~a + ~b * csch(ExpandToSum(~u, ~x))) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~e, ~m, ~p], ~x) && (BinomialQ(~u, ~x) && Not(BinomialMatchQ(~u, ~x)))
    @apply_utils Antiderivative((~x) ^ ~(m') * sech(~(a') + ~(b') * (~x) ^ ~(n')) ^ ~p * sinh(~(a') + ~(b') * (~x) ^ ~(n')), ~x) => (-((~x) ^ ((~m - ~n) + 1)) * sech(~a + ~b * (~x) ^ ~n) ^ (~p - 1)) / (~b * ~n * (~p - 1)) + (((~m - ~n) + 1) / (~b * ~n * (~p - 1))) * Antiderivative((~x) ^ (~m - ~n) * sech(~a + ~b * (~x) ^ ~n) ^ (~p - 1), ~x) <-- FreeQ([~a, ~b, ~p], ~x) && (IntegerQ(~n) && (GeQ(~m - ~n, 0) && NeQ(~p, 1)))
    @apply_utils Antiderivative((~x) ^ ~(m') * csch(~(a') + ~(b') * (~x) ^ ~(n')) ^ ~p * cosh(~(a') + ~(b') * (~x) ^ ~(n')), ~x) => (-((~x) ^ ((~m - ~n) + 1)) * csch(~a + ~b * (~x) ^ ~n) ^ (~p - 1)) / (~b * ~n * (~p - 1)) + (((~m - ~n) + 1) / (~b * ~n * (~p - 1))) * Antiderivative((~x) ^ (~m - ~n) * csch(~a + ~b * (~x) ^ ~n) ^ (~p - 1), ~x) <-- FreeQ([~a, ~b, ~p], ~x) && (IntegerQ(~n) && (GeQ(~m - ~n, 0) && NeQ(~p, 1)))
end

max_stars_repo_path: src/rules/6 Hyperbolic functions/6.5 Hyperbolic secant/6.5.11 (e x)^m (a+b sech(c+d x^n))^p.jl


text:
# Library Path Discovery
if haskey(ENV, "HSA_RUNTIME_PATH")
    const hsa_runtime_path = ENV["HSA_RUNTIME_PATH"]
else
    #warn("environment variable HSA_RUNTIME_PATH is not set, using default path")
    const hsa_runtime_path = "/opt/hsa"
end

function dlopen_or_warn_if_not_found(libraryname, path, optional = false)
    if path == ""
        msg = "HSA library '$libraryname' not found, add its location to LD_LIBRARY_PATH"

        if optional
            info(msg)
        else
            warn(msg)
        end
    end
end
const hsa_libpath = joinpath(hsa_runtime_path, "lib")
const hsa_paths = [hsa_libpath]

const libhsa = Libdl.find_library(["libhsa-runtime64"], hsa_paths)
const libhsa_dl = dlopen_or_warn_if_not_found("libhsa-runtime64", libhsa)

const hsa_include_path = joinpath(dirname(libhsa), "../include")

# Capability Checks

function has_libhsa()
    return libhsa != ""
end

function has_hsa_codegen()
    try
        return ccall(:jl_has_device_target, Bool, (Any,), :hsail)
    catch err
        # Something went wrong, probably the method was not found
        # meaning we are running on a version of julia without device target support
        return false
    end
end


text:
using Statistics
################################################################################
# Defining Functions
function get_x(e)
    x_t = Array{Float64}(undef, M+T, S)

    x_t[1,:] = e[1,:]
    for i = 2:(M+T)
        x_t[i,:] = 0.5*x_t[i-1,:] + e[i,:] + 0.3*e[i-1,:]
    end
    return x_t
end

function get_ψ(d)
    ψ = Array{Float64}(undef, M+T)
    ψ[1] = 1

    for j=1:M+T-1
        ψ[j+1] = ((j-1-d)/j) * ψ[j]
    end
    return ψ
end

function get_data(x, ψ)

    y_final = Array{Float64}(undef, M+T, S)

    # Generating matrix to sum over
    xmat = Array{Float64}(undef, M+T, M+T)
    ymat = similar(xmat)

    y = Array{Float64}(undef, M+T)
    x_used = Array{Float64}(undef, M+T)
    x_tmp = zeros(M+T)

    for s=1:S
        # Running sample s
        x_used = x[:,s]

        for j=0:M+T-1
            x_tmp = zeros(M+T)
            x_tmp[j+1:end] = x_used[1:end-j]
            xmat[j+1,:] = x_tmp
        end

        ymat = xmat.*ψ

        for col=1:size(ymat)[2]
            y[col] = sum(ymat[:,col])
        end
        y_final[:,s] = y
        println(s)
    end
    return y_final
end

function Iy(λj, x)
    bigT = length(x)
    prefactor = (1/2*π*bigT)

    firstpart = Array{Float64}(undef, bigT-1)
    secondpart = similar(firstpart)

    for t=1:bigT-1
        firstpart[t] = x[t] * cos(λj*(t-1))
        secondpart[t] = x[t] * sin(λj*(t-1))
    end
    value = (sum(firstpart))^2 + (sum(secondpart))^2
    return value
end

function get_λs(α, tslength)
    little_m = floor(Int,tslength^α)
    λs = Array{Float64}(undef, little_m)

    for m=1:little_m
        λs[m] = (2*π/tslength) * m
    end
    return λs, little_m
end

function objf(α, x, d)
    λ_list, little_m = get_λs(α, length(x))

    first_sum = Array{Float64}(undef, little_m)
    second_sum = similar(first_sum)

    for (m, λ) in enumerate(λ_list)
        first_sum[m] = Iy(λ, x) / λ^(-2*d)
        second_sum[m] = log(λ)
    end

    firstpart = (1/little_m) * sum(first_sum)
    secondpart = (2*d/little_m) * sum(second_sum)

    value = log(firstpart) - secondpart
    return value
end

function optimize_overd(α, x, ds)
    values = Array{Float64}(undef, length(ds))

    for (i, d) in enumerate(ds)
        values[i] = objf(α, x, d)
    end
    min_d = d_list[argmin(values)]
    return min_d
end

function optimize_overα(αs, x, ds)
    values_LW = Array{Float64}(undef, size(x)[2], length(αs))
    values_PR = similar(values_LW)

    for (i, α) in enumerate(αs)
        for sample=1:size(x)[2]
            values_LW[sample, i] = optimize_overd(α, x[:,sample], ds)
            values_PR[sample, i] = PR(α, x[:,sample])
        end
        println(α)
    end
    return values_LW, values_PR
end

function Ij(λj, x)
    value = Iy(λj, x)
    return log(value)
end

function Rj(λj)
    r = 4 * sin((λj/2))^2
    return -log(r)
end

function OLSestimator(x, y)
    res = inv(x'*x)*(x'*y)
    return res
end

function PR(α, x)
    λ_list, little_m = get_λs(α, length(x))

    x_values = Array{Float64}(undef, little_m, 2)
    y_values = Array{Float64}(undef, little_m, 1)

    x_values[:,1] = ones(little_m)

    for (m, λ) in enumerate(λ_list)
        x_values[m,2] = Rj(λ)
        y_values[m] = Ij(λ, x)
    end

    estimtates = OLSestimator(x_values, y_values)
    return estimtates[2]
end

function get_cis(n, α, type)
    """
    Type 1 = LW, Type 2 = PR
    """

    m = floor(T^α)
    z = 1.96
    if type == 1
        sigma = sqrt(m/4)
    else
        sigma = sqrt((m*π^2/24))
    end

    ci_min = d - (z*sigma)/sqrt(n)
    ci_max = d + z*sigma/sqrt(n)

    return (ci_min, ci_max)
end

function calc_share(x, α, type)
    """
    Type 1 = LW, Type 2 = PR
    """
    n = length(x)
    ci_low, ci_high = get_cis(n, α, type)

    positives = sum([(element > ci_low) & (element < ci_high) for element in x])
    value = positives/n
    return value
end

function calc_bias(x)
    bias = d - mean(x)
    return bias
end

function calc_sigma2(x)
    σ2 = var(x)
    return σ2
end

function cov_probs(LW_values, PR_values, αs)
    probs = Array{Float64}(undef, 2, size(LW_values)[2])
    biases = similar(probs)
    vars = similar(probs)

    for (index, df) in enumerate([LW_values, PR_values])
        for i=1:size(LW_values)[2]
            α = αs[i]
            probs[index,i] = calc_share(df[:,i], α, index)
            biases[index,i] = calc_bias(df[:,i])
            vars[index,i] = calc_sigma2(df[:,i])
        end
    end

    return probs, biases, vars
end
################################################################################
const M=2000
const T=1000
const S=100
const d=0.25
const d_list = LinRange(-0.45, 0.45, 19)
const α_list = LinRange(0.2, 0.8, 7)

e_t = randn(Float64, M+T+1, S)
e_t[1,:] = zeros(S)

x_one = get_x(e_t)
ψ_one = get_ψ(d)

y = get_data(x_one, ψ_one)
y_use = y[M+1:end,:]

# Optimizing
df_LW, df_PR = optimize_overα(α_list, y_use, d_list)
# Calculating final results
co_prob, b, sigma = cov_probs(df_LW, df_PR, α_list)

max_stars_repo_name: christianhilscher/TSA-Assignments

max_issues_repo_name: christianhilscher/TSA-Assignments

max_forks_repo_name: christianhilscher/TSA-Assignments


text:
@testset "Table" begin
    m = Model(optimizer_with_attributes(CS.Optimizer, "no_prune" => true, "logging" => []))
    @variable(m, -5 <= x <= 5, Int)
    @variable(m, -5 <= y <= 5, Int)
    @variable(m, -5 <= z <= 5, Int)

    tab = [
        1 2 3
        1 3 2
        2 1 3
        2 3 1
        3 1 2
        3 2 1
        4 4 4
    ]

    @constraint(m, [x, y, z] in CS.TableSet(tab))
    optimize!(m)
    com = CS.get_inner_model(m)

    constraint = com.constraints[1]

    # check if impossible values got removed
    constr_indices = constraint.indices
    for ind in constr_indices
        @test sort(CS.values(com.search_space[ind])) == 1:4
    end

    @test !CS.is_constraint_solved(constraint, constraint.fct, constraint.set, [1, 2, 4])
    @test CS.is_constraint_solved(constraint, constraint.fct, constraint.set, [1, 2, 3])

    @test CS.still_feasible(
        com,
        constraint,
        constraint.fct,
        constraint.set,
        constr_indices[1],
        4,
    )
    @test CS.fix!(com, com.search_space[constr_indices[2]], 2)
    @test !CS.still_feasible(
        com,
        constraint,
        constraint.fct,
        constraint.set,
        constr_indices[1],
        4,
    )

    # need to create a backtrack_vec to reverse pruning
    dummy_backtrack_obj = CS.BacktrackObj(com)
    push!(com.backtrack_vec, dummy_backtrack_obj)
    # reverse previous fix
    CS.reverse_pruning!(com, 1)
    com.c_backtrack_idx = 1

    @test CS.still_feasible(
        com,
        constraint,
        constraint.fct,
        constraint.set,
        constr_indices[1],
        4,
    )

    m = Model(optimizer_with_attributes(CS.Optimizer, "no_prune" => true, "logging" => []))
    @variable(m, -5 <= x <= 5, Int)
    @variable(m, -5 <= y <= 5, Int)
    @variable(m, -5 <= z <= 5, Int)

    tab = [
        1 2 3
        1 3 2
        2 1 3
        2 3 1
        3 1 2
        3 2 1
        4 4 4
    ]

    @constraint(m, [x, y, z] in CS.TableSet(tab))
    optimize!(m)
    com = CS.get_inner_model(m)

    constraint = com.constraints[1]
    constr_indices = constraint.indices

    @test CS.fix!(com, com.search_space[constr_indices[2]], 2)
    @test CS.prune_constraint!(com, constraint, constraint.fct, constraint.set)
    for ind in constr_indices[[1, 3]]
        @test sort(CS.values(com.search_space[ind])) == [1, 3]
    end

    m = Model(optimizer_with_attributes(CS.Optimizer, "no_prune" => true, "logging" => []))
    @variable(m, -5 <= x <= 5, Int)
    @variable(m, -5 <= y <= 5, Int)
    @variable(m, -5 <= z <= 5, Int)

    tab = [
        1 2 3
        1 3 2
        2 1 3
        2 3 1
        3 1 2
        3 2 1
        4 4 4
    ]

    @constraint(m, [x, y, z] in CS.TableSet(tab))
    optimize!(m)
    com = CS.get_inner_model(m)

    constraint = com.constraints[1]
    constr_indices = constraint.indices

    @test CS.fix!(com, com.search_space[constr_indices[2]], 4)
    @test CS.prune_constraint!(com, constraint, constraint.fct, constraint.set)
    for ind in constr_indices
        @test CS.isfixed(com.search_space[ind])
        @test CS.value(com.search_space[ind]) == 4
    end
end

@testset "table is_constraint_violated test" begin
    m = Model(optimizer_with_attributes(CS.Optimizer, "no_prune" => true, "logging" => []))
    @variable(m, 0 <= x <= 5, Int)
    @variable(m, 0 <= y <= 3, Int)
    @constraint(m, [x, y] in CS.TableSet([
        1 2
        3 4
        2 1
        2 3
    ]))
    optimize!(m)
    com = CS.get_inner_model(m)

    constraint = com.constraints[1]

    variables = com.search_space
    @test CS.fix!(com, variables[constraint.indices[1]], 1; check_feasibility = false)
    @test CS.fix!(com, variables[constraint.indices[2]], 3; check_feasibility = false)
    @test CS.is_constraint_violated(com, constraint, constraint.fct, constraint.set)

    m = Model(optimizer_with_attributes(CS.Optimizer, "no_prune" => true, "logging" => []))
    @variable(m, 0 <= x <= 5, Int)
    @variable(m, 0 <= y <= 3, Int)
    @constraint(m, [x, y] in CS.TableSet([
        1 2
        3 4
        2 1
        2 3
    ]))
    optimize!(m)
    com = CS.get_inner_model(m)

    constraint = com.constraints[1]

    variables = com.search_space
    @test CS.fix!(com, variables[constraint.indices[1]], 1; check_feasibility = false)
    @test CS.fix!(com, variables[constraint.indices[2]], 2; check_feasibility = false)
    @test !CS.is_constraint_violated(com, constraint, constraint.fct, constraint.set)
end


text:
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
# ==============================================================================
#
#   Tests related to ECEF to ECI transformations using satellite state vector.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
# ==============================================================================
#
#   [1] Vallado, D. A (2013). Fundamentals of Astrodynamics and Applications.
#       Microcosm Press, Hawthorn, CA, USA.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

# Get the current EOP Data.
#
# TODO: The EOP data obtained from IERS website does not match the values in the
# examples in [1]. However, this should be enough, because 1) the individual
# functions at the low level are tested using the same values of [1], and 2) the
# difference is smaller than 30 cm.

eop_iau1980  = read_iers_eop("./eop_IAU1980.txt",  Val(:IAU1980))
eop_iau2000a = read_iers_eop("./eop_IAU2000A.txt", Val(:IAU2000A))

# File: ./src/transformations/sv_ecef_to_eci.jl
# =============================================

# Functions: sv_ecef_to_eci
# -------------------------

# The rotations functions were already heavily tested in `ecef_to_eci.jl`,
# `fk5/fk5.jl`, and `iau2006/iau2006.jl`. Hence, here we will do only some minor
# testing involving the following transformations:
#
#   ITRF => GCRF  (FK5)
#   PEF  => J2000 (FK5, no corrections)
#   ITRF => GCRF  (IAU-2010)
#   TIRS => GCRF  (IAU-2010, no corrections)

################################################################################
#                                  IAU-76/FK5
################################################################################

# ITRF to GCRF
# ============

################################################################################
#                                 Test Results
################################################################################
#
# Scenario 01
# ===========
#
# Example 3-15: Performing IAU-76/FK5 reduction [1, p. 230].
#
# According to this example and Table 3-6, using:
#
#   UTC    = April 6, 2004, 07:51:28.386009
#   r_itrf = -1033.4793830    i + 7901.2952754    j + 6380.3565958    k [km]
#   v_itrf =    -3.225636520  i -    2.872451450  j +    5.531924446  k [km/s]
#
# one gets:
#
#   r_gcrf =  5102.50895790   i + 6123.01140070   j + 6378.13692820   k [km]
#   v_gcrf =    -4.7432201570 i +    0.7905364970 j +    5.533755270  k [km/s]
#
################################################################################

@testset "Function sv_ecef_to_eci ITRF => GCRF" begin
    JD_UTC = date_to_jd(2004, 4, 6, 7, 51, 28.386009)

    r_itrf  = [-1033.4793830; 7901.2952754; 6380.3565958]
    v_itrf  = [-3.225636520; -2.872451450; +5.531924446]
    sv_itrf = orbsv(JD_UTC, r_itrf, v_itrf)
    sv_gcrf = sv_ecef_to_eci(sv_itrf, ITRF(), GCRF(), JD_UTC, eop_iau1980)

    @test sv_gcrf.t === JD_UTC

    @test sv_gcrf.r[1] ≈ +5102.50895790 atol=3e-4
    @test sv_gcrf.r[2] ≈ +6123.01140070 atol=3e-4
    @test sv_gcrf.r[3] ≈ +6378.13692820 atol=3e-4

    @test sv_gcrf.v[1] ≈ -4.7432201570  atol=8e-7
    @test sv_gcrf.v[2] ≈ +0.7905364970  atol=8e-7
    @test sv_gcrf.v[3] ≈ +5.5337552700  atol=8e-7
end

# PEF to J2000
# ============

################################################################################
#                                 Test Results
################################################################################
#
# Scenario 01
# ===========
#
# Example 3-15: Performing IAU-76/FK5 reduction [1, p. 230].
#
# According to this example and Table 3-6, using:
#
#   UTC    = April 6, 2004, 07:51:28.386009
#   UT1    = UTC - 0.4399619/86400
#   r_pef  = -1033.47503130   i + 7901.30558560   j + 6380.34453270   k [km]
#   v_pef  =    -3.2256327470 i -    2.8724425110 j +    5.5319312880 k [km/s]
#
# one gets:
#
#   r_j2000 = 5102.50960000  i + 6123.01152000   j + 6378.13630000   k [km]
#   v_j2000 =  -4.7432196000 i + 0.7905366000    j + 5.5337561900    k [km/s]
#
################################################################################

@testset "Function sv_ecef_to_eci PEF => J2000" begin
    JD_UT1 = date_to_jd(2004,4,6,7,51,28.386009) - 0.4399619/86400

    r_pef    = [-1033.47503130; 7901.30558560; 6380.34453270]
    v_pef    = [-3.2256327470; -2.8724425110; +5.5319312880]
    sv_pef   = orbsv(JD_UT1, r_pef, v_pef)
    sv_j2000 = sv_ecef_to_eci(sv_pef, PEF(), J2000(), JD_UT1)

    @test sv_j2000.t === JD_UT1

    @test sv_j2000.r[1] ≈ +5102.50960000 atol=1e-7
    @test sv_j2000.r[2] ≈ +6123.01152000 atol=1e-7
    @test sv_j2000.r[3] ≈ +6378.13630000 atol=1e-7

    @test sv_j2000.v[1] ≈ -4.7432196000  atol=1e-7
    @test sv_j2000.v[2] ≈ +0.7905366000  atol=1e-7
    @test sv_j2000.v[3] ≈ +5.5337561900  atol=1e-7
end

################################################################################
#                                IAU-2006/2010
################################################################################

# ITRF to GCRF
# ============

################################################################################
#                                 Test Results
################################################################################
#
# Scenario 01
# ===========
#
# Example 3-15: Performing IAU-2010 reduction [1, p. 220].
#
# According to this example and Table 3-6, using:
#
#   UTC    = April 6, 2004, 07:51:28.386009
#   r_itrf = -1033.4793830    i + 7901.2952754    j + 6380.3565958    k [km]
#   v_itrf =    -3.225636520  i -    2.872451450  j +    5.531924446  k [km/s]
#
# one gets:
#
#   r_gcrf = 5102.50895290   i + 6123.01139910   j + 6378.13693380   k [km]
#   v_gcrf =   -4.7432201610 i +    0.7905364950 j +    5.533755240  k [km/s]
#
################################################################################

@testset "Function sv_ecef_to_eci ITRF => GCRF" begin
    JD_UTC = date_to_jd(2004, 4, 6, 7, 51, 28.386009)

    r_itrf  = [-1033.4793830; 7901.2952754; 6380.3565958]
    v_itrf  = [-3.225636520; -2.872451450; +5.531924446]
    sv_itrf = orbsv(JD_UTC, r_itrf, v_itrf)
    sv_gcrf = sv_ecef_to_eci(sv_itrf, ITRF(), GCRF(), JD_UTC, eop_iau2000a)

    @test sv_gcrf.t === JD_UTC

    @test sv_gcrf.r[1] ≈ +5102.50895290 atol=3e-4
    @test sv_gcrf.r[2] ≈ +6123.01139910 atol=3e-4
    @test sv_gcrf.r[3] ≈ +6378.13693380 atol=3e-4

    @test sv_gcrf.v[1] ≈ -4.7432201610  atol=8e-7
    @test sv_gcrf.v[2] ≈ +0.7905364950  atol=8e-7
    @test sv_gcrf.v[3] ≈ +5.5337557240  atol=8e-7
end

# TIRS to GCRF
# ============

################################################################################
#                                 Test Results
################################################################################
#
# Scenario 01
# ===========
#
# Example 3-15: Performing IAU-2010 reduction [1, p. 220].
#
# According to this example and Table 3-6, using:
#
#   UTC    = April 6, 2004, 07:51:28.386009
#   UT1    = UTC - 0.4399619/86400
#   r_tirs  = -1033.47503120   i + 7901.30558560   j + 6380.34453270   k [km]
#   v_tirs  =    -3.2256327470 i -    2.8724425110 j +    5.5319312880 k [km/s]
#
# one gets:
#
#   r_gcrf = 5102.50895290   i + 6123.01139910   j + 6378.13693380   k [km]
#   v_gcrf =   -4.7432201610 i +    0.7905364950 j +    5.533755240  k [km/s]
#
################################################################################

@testset "Function sv_ecef_to_eci TIRS => GCRF" begin
    JD_UT1 = date_to_jd(2004, 4, 6, 7, 51, 28.386009) - 0.4399619/86400

    r_tirs  = [-1033.47503120; 7901.30558560; 6380.34453270]
    v_tirs  = [-3.2256327470; -2.8724425110; +5.5319312880]
    sv_tirs = orbsv(JD_UT1, r_tirs, v_tirs)
    sv_gcrf = sv_ecef_to_eci(sv_tirs, TIRS(), GCRF(), JD_UT1)

    @test sv_gcrf.t === JD_UT1

    @test sv_gcrf.r[1] ≈ +5102.50895290 atol=3e-4
    @test sv_gcrf.r[2] ≈ +6123.01139910 atol=3e-4
    @test sv_gcrf.r[3] ≈ +6378.13693380 atol=3e-4

    @test sv_gcrf.v[1] ≈ -4.7432201610  atol=8e-7
    @test sv_gcrf.v[2] ≈ +0.7905364950  atol=8e-7
    @test sv_gcrf.v[3] ≈ +5.5337557240  atol=8e-7
end


text:
@testset "Type Assertions in Ematcher" begin
    some_theory = @theory begin
        a * b => b * a
        a::$Number * b::$Number => matched(a,b)
        a::$Int64 * b::$Int64 => specific(a,b)
        a * (b * c) => (a * b) * c
    end

    G = EGraph(:(2*3))
    display(G.M); println()

    @test true == areequal(G, some_theory, :(2 * 3), :(matched(2,3)))
    display(G.M); println()
    @test true == areequal(G, some_theory, :(matched(2,3)), :(specific(3,2)))
end

# TODO removed by now! Text Taine Zhao
# @testset "Type Variables in Ematcher" begin
#     some_theory = @theory begin
#         a * b => b * a
#         a::T * b::T => sametype(T)
#         a * (b * c) => (a * b) * c
#     end
#
#     G = EGraph(:(2*3))
#     display(G.M); println()
#     res = areequal(G, some_theory, :(2 * 3), :(sametype($Int64)))
#     display(G.M); println()
#     @test res
#
#     G = EGraph(:(2*"ciao"))
#     display(G.M); println()
#     res = !areequal(G, some_theory, :(2 * "ciao"), :(sametype($Int64)))
#     display(G.M); println()
#     @test res
#
#     G = EGraph(:("ciaoz"*"ciao"))
#     display(G.M); println()
#     res = areequal(G, some_theory, :("ciaoz" * "ciao"), :(sametype($String)))
#     display(G.M); println()
#     @test res
#     # @test true == areequal(G, some_theory, :(matched(2,3)), :(specific(3,2)))
# end


text:
"""
    DeltaRCWA

An approximate RCWA-based solver for the scattering of electromagnetic fields in
layered, periodic metasurfaces via a simplified model with infinitesimal sheets.

Developed by Luke Qi, Lorenzo Van Muñoz at Steven G. Johnson's group in 2021
"""
module DeltaRCWA

using LinearAlgebra: I, inv, Diagonal

using BlockArrays
using FFTW: fftfreq, fft, ifft, bfft, Frequencies
using RecipesBase
using KrylovKit: linsolve
using IterativeSolvers: gmres
using LinearMaps: LinearMap, BlockMap

export smatrix
"""
    smatrix([T=Matrix], scatterer, modes, pol)

Form the scattering matrix of a given return type `T` for the specified
scatterers, incident modes, and polarizations.
"""
smatrix(args...) = smatrix(Matrix, args...)
smatrix(::Type{Matrix}, args...) = _sMatrix(args...)
smatrix(::Type{BlockMatrix}, args...) = _sBlockMatrix(args...)
smatrix(::Type{LinearMap}, args...) = _sLinearMap(args...)
smatrix(::Type{BlockMap}, args...) = _sBlockMap(args...)

include("modes.jl")
include("media.jl")
include("slabs.jl")
include("gstcs.jl")
include("sheets.jl")
include("stacks.jl")
include("DeltaRCWAProblem.jl")
include("star_products.jl")
include("plotting.jl")

end


text:
using NFFT, DataFrames, LinearAlgebra, LaTeXStrings, DelimitedFiles
using BenchmarkTools
using Plots; pgfplotsx()

include("../Wrappers/NFFT3.jl")
include("../Wrappers/FINUFFT.jl")

const packagesCtor = [NFFTPlan, NFFTPlan, NFFTPlan, NFFT3Plan, FINUFFTPlan]
const packagesStr = [ "NFFT.jl/TENSOR", "NFFT.jl/POLY", "NFFT.jl/LINEAR", "NFFT3/TENSOR", "FINUFFT"]
const precomp = [NFFT.TENSOR, NFFT.POLYNOMIAL, NFFT.LINEAR, NFFT.TENSOR, NFFT.LINEAR]
const blocking = [true, true, true, true, true]

const benchmarkTime = [4, 4]

NFFT.FFTW.set_num_threads(Threads.nthreads())
ccall(("omp_set_num_threads",NFFT3.lib_path_nfft),Nothing,(Int64,),convert(Int64,Threads.nthreads()))
@info ccall(("nfft_get_num_threads",NFFT3.lib_path_nfft),Int64,())
NFFT._use_threads[] = (Threads.nthreads() > 1)

const σs = [2.0]
const ms = 3:8
const NBase = [4*4096, 128, 32]
const Ds = 1:3

function nfft_accuracy_comparison(Ds=1:3)
  println("\n\n ##### nfft_performance vs accuracy ##### \n\n")

  df = DataFrame(Package=String[], D=Int[], M=Int[], N=Int[], m = Int[], σ=Float64[],
                   ErrorTrafo=Float64[], ErrorAdjoint=Float64[],
                   TimeTrafo=Float64[], TimeAdjoint=Float64[] )

  for D in Ds
    @info "### Dimension D=$D ###"
    N = ntuple(d->NBase[D], D)
    M = prod(N)

    x = rand(D,M) .- 0.5
    fHat = randn(ComplexF64, M)
    fApprox = randn(ComplexF64, N)
    gHatApprox = randn(ComplexF64, M)

    # ground truth (numerical)
    pNDFT = NDFTPlan(x, N)
    f = adjoint(pNDFT) * fHat
    gHat = pNDFT * f

    for σ in σs
      for m in ms
        @info "m=$m D=$D σ=$σ "

        for pl = 1:length(packagesStr)
          planner = packagesCtor[pl]
          p = planner(x, N; m, σ, precompute=precomp[pl], blocking=blocking[pl])

          @info "Adjoint accuracy: $(packagesStr[pl])"
          mul!(fApprox, adjoint(p), fHat)
          eadjoint = norm(f[:] - fApprox[:]) / norm(f[:])

          @info "Adjoint benchmark: $(packagesStr[pl])"
          BenchmarkTools.DEFAULT_PARAMETERS.seconds = benchmarkTime[1]
          b = @benchmark mul!($fApprox, $(adjoint(p)), $fHat)
          tadjoint = minimum(b).time / 1e9

          @info "Trafo accuracy: $(packagesStr[pl])"
          mul!(gHatApprox, p, f)
          etrafo = norm(gHat[:] - gHatApprox[:]) / norm(gHat[:])

          @info "Trafo benchmark: $(packagesStr[pl])"
          BenchmarkTools.DEFAULT_PARAMETERS.seconds = benchmarkTime[2]
          b = @benchmark mul!($gHatApprox, $p, $f)
          ttrafo = minimum(b).time / 1e9

          push!(df, (packagesStr[pl], D, M, N[D], m, σ, etrafo, eadjoint, ttrafo, tadjoint))

        end
      end
    end
  end
  return df
end

function plot_accuracy(df, D=1)

  df1_ = df[df.σ.==2.0 .&& df.D.==D,:]
  Plots.scalefontsizes()
  Plots.scalefontsizes(1.5)

  titleTrafo = L"\textrm{NFFT}, \textrm{%$(D)D}"
  titleAdjoint = L"\textrm{NFFT}^H, \textrm{%$(D)D}"

  colors = [:black, :orange, :blue, :green, :brown, :gray, :blue, :purple, :yellow ]
  ls = [:solid, :dashdot, :solid, :solid, :solid, :dash, :solid, :dash, :solid]
  shape = [:circle, :circle, :circle, :xcross, :circle, :xcross, :xcross, :circle]

  maxTimeTrafo = maximum(df1_[:,:TimeTrafo])
  maxTimeAdjoint = maximum(df1_[:,:TimeAdjoint])

  p1 = plot(df1_[df1_.Package.==packagesStr[1],:ErrorTrafo],
            df1_[df1_.Package.==packagesStr[1],:TimeTrafo], ylims=(0.0,maxTimeTrafo),
            xscale = :log10, label=packagesStr[1], lw=2, xlabel = "Relative Error", ylabel="Runtime / s",
            legend = (:topright), title=titleTrafo, shape=:circle, c=:black)

  for p=2:length(packagesStr)
    plot!(p1, df1_[df1_.Package.==packagesStr[p],:ErrorTrafo],
          df1_[df1_.Package.==packagesStr[p],:TimeTrafo],
            xscale = :log10, label=packagesStr[p], lw=2, shape=shape[p], ls=ls[p],
            c=colors[p], msc=colors[p], mc=colors[p], ms=4, msw=2)
  end

  p2 = plot(df1_[df1_.Package.==packagesStr[1],:ErrorAdjoint],
            df1_[df1_.Package.==packagesStr[1],:TimeAdjoint], ylims=(0.0,maxTimeAdjoint),
            xscale = :log10,  lw=2, xlabel = "Relative Error", ylabel="Runtime / s", #label=packagesStr[1],
            legend = nothing, title=titleAdjoint, shape=:circle, c=:black)

  for p=2:length(packagesStr)
    plot!(p2, df1_[df1_.Package.==packagesStr[p],:ErrorAdjoint],
          df1_[df1_.Package.==packagesStr[p],:TimeAdjoint],
            xscale = :log10,  lw=2, shape=shape[p], ls=ls[p], #label=packagesStr[p],
            c=colors[p], msc=colors[p], mc=colors[p], ms=4, msw=2)
  end

  p = plot(p1, p2, layout=(1,2), size=(800,300), dpi=200)

  savefig(p, "../docs/src/assets/performanceVsAccuracy_D$(D).svg")
  return p
end

df = nfft_accuracy_comparison(Ds)
writedlm("performanceVsAccuracy.csv", Iterators.flatten(([names(df)], eachrow(df))), ',')

data, header = readdlm("performanceVsAccuracy.csv", ',', header=true);
df = DataFrame(data, vec(header))

for d=1:3
  plot_accuracy(df, d)
end


text:
@doc doc"""

    Taylens(ϕ, order)

`Taylens` is a lensing operator which lenses a map with a
nearest-pixel permute step followed by power series expansion in the
residual displacement, to any order. This is the algorithm from
[Næss&Louis 2013](https://arxiv.org/abs/1307.0719).
"""
struct Taylens{L<:PowerLens} <: ImplicitOp{Bottom}

    # nearest pixel remapping
    i :: Matrix{Int}
    j :: Matrix{Int}

    # residual lensing (via PowerLens)
    residual :: L

end

Taylens(order) = x -> Taylens(x, order)
Taylens(ϕ::FlatS0, order) = Taylens(∇*ϕ, order)
function Taylens(d::FieldVector, order)

    require_unbatched(d[1])
    @unpack ℓx,ℓy,Nx,Ny,Δx,metadata = d[1]

    # total displacement
    dx, dy = d[1][:Ix], d[2][:Ix]

    # nearest pixel displacement
    indexwrap(i, N) = mod(i - 1, N) + 1
    di = round.(Int,dy/Δx)
    dj = round.(Int,dx/Δx)
    i = indexwrap.(di .+ (1:Ny) , Ny)
    j = indexwrap.(dj .+ (1:Nx)', Nx)

    # residual displacement
    r = @SVector[
        FlatMap(dx .- dj .* Δx, metadata),
        FlatMap(dy .- di .* Δx, metadata),
    ]
    residual = PowerLens(r, order)

    Taylens(i, j, residual)

end

function (*)(Lϕ::Taylens, f::FlatField)
    require_unbatched(f)
    function nearest_pixel_remapping(Łf)
        Łf.arr .= getindex.(Ref(copy(Łf.arr)), Lϕ.i, Lϕ.j, (reshape(1:Npol, 1, 1, :) for Npol in size(Łf.arr)[3:end])...)
        Łf
    end
    @unpack ∇1ϕᵖ, ∇2ϕᵖ, order = Lϕ.residual
    Ðf = Ð(f)
    # nearest pixel remapping
    f̃ = nearest_pixel_remapping(Ł(f))
    # add in residual Taylor correction
    for n in 1:order, (a,b) in zip(0:n,n:-1:0)
        @. f̃ += ∇1ϕᵖ[a] * ∇2ϕᵖ[b] * $(nearest_pixel_remapping(Ł(∇[1]^a * ∇[2]^b * Ðf))) / factorial(a) / factorial(b)
    end
    f̃
end


text:
mutable struct GFAssessState
    choices::ChoiceMap
    weight::Float64
    visitor::AddressVisitor
    params::Dict{Symbol,Any}
end

function GFAssessState(choices, params::Dict{Symbol,Any})
    GFAssessState(choices, 0., AddressVisitor(), params)
end

function traceat(state::GFAssessState, dist::Distribution{T},
              args, key) where {T}
    local retval::T

    # check that key was not already visited, and mark it as visited
    visit!(state.visitor, key)

    # get return value
    retval = get_value(state.choices, key)

    # update weight
    state.weight += logpdf(dist, retval, args...)

    retval
end

function traceat(state::GFAssessState, gen_fn::GenerativeFunction{T,U},
              args, key) where {T,U}
    local retval::T

    # check that key was not already visited, and mark it as visited
    visit!(state.visitor, key)

    # get constraints for this call
    choices = get_submap(state.choices, key)

    # get return value and weight increment
    (weight, retval) = assess(gen_fn, args, choices)

    # update score
    state.weight += weight

    retval
end

function splice(state::GFAssessState, gen_fn::DynamicDSLFunction, args::Tuple)
    prev_params = state.params
    state.params = gen_fn.params
    retval = exec(gen_fn, state, args)
    state.params = prev_params
    retval
end

function assess(gen_fn::DynamicDSLFunction, args::Tuple, choices::ChoiceMap)
    state = GFAssessState(choices, gen_fn.params)
    retval = exec(gen_fn, state, args)

    unvisited = get_unvisited(get_visited(state.visitor), choices)
    if !isempty(unvisited)
        error("Assess did not visit the following constraint addresses:\n$unvisited")
    end

    (state.weight, retval)
end


text:
import Statistics, StatsModels, GLM

"""
    r2_tjur(model::StatsModels.DataFrameRegressionModel{<:GLM.GeneralizedLinearModel})

Compute Tjur's (2009) D (R²).

The Coefficients of Determination (D), also referred to as Tjur's R² ([Tjur, 2009](https://amstat.tandfonline.com/doi/abs/10.1198/tast.2009.08210#.W5eJjOgzYuU)), is asymptotically equivalent to the classical version of R² for linear models.

# Arguments
- `model`: A `GeneralizedLinearModel`.

# Examples
```jldoctest
using GLM, DataFrames

model = glm(@formula(y ~ Var1), DataFrame(y=[0, 0, 1, 1], Var1=[1, 2, 2, 4]), GLM.Binomial())
r2_tjur(model)

# output
0.5
```

# References
- Tjur, T. (2009). Coefficients of determination in logistic regression models—A new proposal: The coefficient of discrimination. The American Statistician, 63(4), 366-372.
"""
function r2_tjur(model::StatsModels.DataFrameRegressionModel{<:GLM.GeneralizedLinearModel})
    data = model_data(model)
    outcome = values(model_description(model))["Outcome"]

    Average_1 = Statistics.mean(round.(GLM.predict(model, data[data[outcome] .== 1.0, :])))
    Average_0 = Statistics.mean(round.(GLM.predict(model, data[data[outcome] .== 0.0, :])))
    D = abs(Average_1 - Average_0)
    return D
end

# http://rsif.royalsocietypublishing.org/content/14/134/20170213
# https://github.com/cran/MuMIn/blob/master/R/r.squaredGLMM2.R
# https://github.com/strengejacke/sjstats/blob/master/R/r2.R#L445
# r2glmm


text:
using BasisFunctions: DictionarySolverOperator
import BasisFunctions: operator, apply!

export AZSolver
"""
Fast implementation of the AZ Algorithm. Returns an operator that is approximately
`pinv(A)`, based on a suitable `Zt` so that `(A*Zt-I)*A` is low rank.

Steps when applying to right hand side b:

1. (I-A*Zt)*A*x2=(I-A*Zt)*b, solve for x2
This is by default done using a low rank decomposition of (I-A*Zt)*A (RandomizedSvdSolver).

2. x1 = Zt*(b-A*x2)
Can be done fast if Zt and A are fast.

3. x = x1+x2
This is the solution.
"""
struct AZSolver{T} <: DictionarySolverOperator{T}
    A           ::  DictionaryOperator # Store for application in step 2
    Zt          ::  DictionaryOperator # Store for application in step 2
    plunge_op   ::  DictionaryOperator # (A*Zt-I), store because it allocates memory
    psolver     ::  DictionaryOperator # The low rank plunge solver for (A*Zt-I)*A
    Pb                                 # Scratch space for the plunge opreator applied to the right hand size
    x1                                 # Scratch space for the intermediate result vector x1

    function AZSolver{T}(A::DictionaryOperator, Zt::DictionaryOperator, plunge_op::DictionaryOperator,
            psolver::DictionaryOperator) where {T}
        # Allocate scratch space
        Pb = zeros(T, src(psolver))
        x1 = zeros(T, src(A))
        new(A, Zt, plunge_op, psolver, Pb, x1)
    end
end

AZSolver(A::DictionaryOperator{T}, Zt::DictionaryOperator{T}, plunge_op::DictionaryOperator{T}, psolver::DictionaryOperator{T}) where {T} =
    AZSolver{eltype(A)}(A, Zt, plunge_op, psolver)

operator(s::AZSolver) = s.A

default_regularization(A; options...) = pSVD_solver(A; options...)

function default_regularization(A::DictionaryOperator{BigFloat}; options...)
    @warn "Our own implementation of a randomized SVD solver is used with `BigFloat`s.
        This implementation is not as robust as the partial SVD solver in `LowRankApprox.jl`.
        Concider using `DoubleFloats.jl` instead."
    rSVD_solver(A; options...)
end
function default_regularization(A::DictionaryOperator{Complex{BigFloat}}; options...)
    @warn "Our own implementation of a randomized SVD solver is used with `BigFloat`s.
        This implementation is not as robust as the partial SVD solver in `LowRankApprox.jl`.
        Concider using `DoubleFloats.jl` instead."
    rSVD_solver(A; options...)
end

function AZSolver(A::DictionaryOperator, Zt::DictionaryOperator;
            REG = default_regularization,
            rankestimate = 40,
            threshold = default_threshold(A),
            verbose = false,
            weightedAZ = false,
            options...)

    plunge_op = I-A*Zt
    if weightedAZ
        verbose && @info "Weighted AZ"
        AZ_Cweight = haskey(options,:AZ_Cweight) ? options[:AZ_Cweight] : error("No options `AZ_Cweight`")
        psolver = AZ_Cweight*REG(plunge_op*A*AZ_Cweight; verbose=verbose, threshold = threshold, rankestimate = rankestimate, options...)
    else
        psolver = REG(plunge_op*A; verbose=verbose, threshold = threshold, rankestimate = rankestimate, options...)
    end
    AZSolver(A, Zt, plunge_op, psolver)
end

default_threshold(A::DictionaryOperator) = regularization_threshold(eltype(A))

apply!(s::AZSolver, coef_dest, coef_src) = _apply!(s, coef_dest, coef_src,
        s.plunge_op, s.A, s.Zt, s.psolver, s.Pb, s.x1)

function _apply!(s::AZSolver, coef_dest, coef_src, plunge_op::DictionaryOperator, A, Zt, psolver, Pb, x1)
    # Step 1: Solve (I-A*Zt)*A * x1 = (I-A*Zt)*b
    apply!(plunge_op, Pb, coef_src)
    apply!(psolver, x1, Pb)

    # Step 2: Compute x2 =  Zt*(b-A*x1)
    apply!(A, Pb, x1)
    Pb .= coef_src .- Pb
    apply!(Zt, coef_dest, Pb)

    # Step 3: x = x1 + x2
    coef_dest .+= x1
end