text:
"""
The @knet macro defines a new knet function:

    @knet function wbf(x; f=relu, o...)
        y = wdot(x; o...)
        z = bias(y; o...)
        return f(z; o...)
    end

Once a knet function is defined it can be:

    (1) Used as an operator in other knet function definitions
    (2) Compiled into a knet model (Net) using the compile function
"""
macro knet(f)
    (fname, fargs, fpars, fbody) = _comp_parse_def(f)
    esc(Expr(:call,Kfun.kdef,Expr(:quote,fname),Expr(:quote,f)))
end

"""
compile(fname::Symbol; o...) compiles the knet function given by fname
into a knet model, i.e. a Net object.  fname should be a symbol
(i.e. use compile(:lstm), not compile(lstm)).  Optional keyword
arguments (o...) can be used to pass initialization parameters to the
compiler.  Knet function names in keyword args should also be escaped
with a colon.  For example:

    net = compile(:wbf; f=:sigm, out=100)
"""
function compile(fname::Symbol; o...)
    @dbg println((:compile,:fname,fname,:o,o))
    isdefined(Kfun, fname) || error("$fname not defined as a knet function, use @knet in function declaration")
    prog = _comp(Kfun.(fname); o...)
    @dbg println((:compile,:prog,prog))
    Net(prog)
    # inst = _comp_inst(prog)
    # @dbg println((:compile,:inst,inst))
    # Net(inst)
end

function _comp(f::Expr; o...)
    @dbg println((:_comp1,:f,f,:o,o))
    (fname, fargs, fpars, fbody) = _comp_parse_def(f)
    expr = Expr(:block)
    for x in fargs; push!(expr.args, :($x=input())); end
    append!(expr.args, fbody.args)
    locals = [x=>x for x in _comp_locals(fbody)]
    @dbg println((:_comp1,:locals,locals))
    fpars = _comp_fpars(f, o)
    cond = Expr(:&&)  # A 0-arg && evaluates to true
    prog = _comp(expr, locals, fpars, cond)
    @dbg println((:_comp1,:return,prog))
    return prog
end

function _comp{T<:Op}(f::Type{T}; o...)
    @dbg println((:_comp2,:f,f,:o,o))
    feval = f(;o...)
    fargs = [ symbol("x$i") for i in 1:ninputs(feval) ]
    fcall = Expr(:call, feval, fargs...)
    expr = Expr(:block)
    for x in fargs; push!(expr.args, :($x=input())); end
    push!(expr.args, Expr(:return, fcall))
    locals = [x=>x for x in _comp_locals(expr)]
    @dbg println((:_comp2,:locals,locals))
    prog = _comp(expr, locals, Dict(), Expr(:&&))
    @dbg println((:_comp2,:return,prog))
    return prog
end

# _comp compiles expr in the context defined by name, value, and cond.
# It returns a Vector{Reg} array of instructions.
# name is a Dict{Symbol,Symbol} that provides name substitution rules.
# value is a Dict{Symbol,Any} that gives values for variables.
# cond is an Expr that gives the current condition.
#
# Example:
#
# @knet function wf(x; f=relu, o...)
#     y = wdot(x; o...)
#     return f(y; o...)
# end
#
# compile(:wf; out=10, f=:sigm)
#
# calls _comp with:
# expr: body of wf prepended with :(x=input())
# name: {:x=>:x, :y=>:y, :return=>:return} preserving the local variable names of the top level call
#       name is used for renaming variables that are lhs in assignments, and positional arguments.
#       any local variable not found in name will be replaced with a gensym when compiling the body.
# value: {:f=>:sigm, :o=>[(:out,10)]} these values should be used while compiling the function body
#        value is for evaluating variables that are function names and keyword arguments.
# cond: :($(Expr(:&&))) the initial cond is an empty conjunction which evaluates to true

function _comp(expr,name::Dict,value::Dict,cond::Expr)
    @dbg println((:_comp,:expr,expr,:name,name,:value,value,:cond,cond))
    p = (isa(expr,LineNumberNode) ? Any[] :
         !isa(expr,Expr) ? error("expecting expression got $expr") :
         expr.head == :block ? mapreduce(x->_comp(x,name,value,cond), append!, expr.args) :
         expr.head == :if ? _comp_if(expr, name, value, cond) :
         expr.head == :return ? _comp(Expr(:(=), :return, expr.args[1]),name,value,cond) : # return y -> return = y
         expr.head != :(=) ? error("expecting assignment expression got $expr") :
         !isa(expr.args[1], Symbol) ? error("lhs should be a symbol in $expr") :
         isa(expr.args[2], Symbol) ? _comp(Expr(:(=), expr.args[1], Expr(:call,:copy,expr.args[2])),name,value,cond) : # y=x -> y=copy(x)
         !isa(expr.args[2], Expr) || expr.args[2].head != :call ? error("rhs should be a function call in $expr") :
         _comp_assignment(expr, name, value, cond) # y=f(x...;o...)
         )
    @dbg println((:_comp, :return, p))
    return p
end

# Continuing the last example, since the body is a block, _comp calls
# itself on each line recursively appending the results.  At some
# point we get to a typical line like :(y=wdot(x;o...))  which calls
# _comp_assignment with the name, value, cond given above.

# _comp_assignment :(y=wdot(x;o...)) has to create a new name2, value2, cond2 to compile the body of wdot.
# (1) eval the name of the function (wdot) in the value dict for possible replacement (it could be a variable f), then the Kfun global environment (all knet functions must be defined there).
# (2) lookup the name of lhs (y) in the name dict, creating a gensym for it if not found.  Set name2[:return] to this name.
# (3) compile any arg expressions, adding their instructions to the program, and generating gensyms for their return variables. Add these as {par=>arg} mappings to name2.
# (4) lookup the names of the args (x) in the name dict, creating gensyms for them if not found, adding {par=>arg} mappings to name2.
# (5) evaluate keyword argument values in the call using (a) value dict, (b) Kfun table, (c) julia environment.
# (6) construct the value dict considering the function definition of wdot and the output of #5.

function _comp_assignment(expr::Expr,name::Dict,value::Dict,cond::Expr)
    @dbg println((:_comp_assignment,:expr,expr,:name,name,:value,value,:cond,cond))
    prog = Reg[]
    (f,fargs,fpars) = _comp_parse_call(expr.args[2])

    haskey(value, f) && (f=value[f])
    # haskey(_KENV, f) || error("$f is not a knet function.")
    # feval = _KENV[f]
    # feval = _comp_eval(f, value)
    feval = eval(Kfun, f)
    @dbg println((:_comp_assignment,:feval,feval))

    y = expr.args[1]
    yname = get!(name, y, gensym(y))
    @dbg println((:_comp_assignment,:yname,yname))

    xname = map(fargs) do x
        if isa(x,Symbol)
            get!(name,x,gensym(x))
        elseif isa(x,Expr)
            # TODO: use the same reg if possible for register sharing
            tmp = gensym(:tmp)
            name[tmp] = tmp
            append!(prog, _comp(Expr(:(=), tmp, x), name, value, cond))
            tmp
        else
            error("Malformed positional argument $x in $expr")
        end
    end
    xname = convert(Vector{Symbol}, xname)
    @dbg println((:_comp_assignment,:xname,xname))

    odict = Dict{Symbol,Any}()
    for k in fpars              # this is the fpars from the call, not the definition
        if k.head == :kw
            odict[k.args[1]] = _comp_eval(k.args[2], value)
        elseif k.head == :(...)
            for (a,b) in _comp_eval(k.args[1], value)
                odict[a] = b
            end
        else
            error("Malformed keyword argument $k in $expr")
        end
    end
    @dbg println((:_comp_assignment,:odict,odict))

    if isa(feval, Op)
        # This happens when compiling a primitive Op directly, kwargs already taken into account in comp2
        push!(prog, Reg(feval, yname, xname, cond))
    elseif isa(feval, DataType) && (feval <: Op)
        op = feval(; odict...)
        push!(prog, Reg(op, yname, xname, cond))
    else
        isa(feval, Function) && (feval = feval(; odict...)) # This allows macros like repeat
        isa(feval, Expr) || error("expecting Op or Expr got $f")
        name2 = _comp_fargs(feval, xname, yname)
        value2 = _comp_fpars(feval, odict)
        fbody = feval.args[2]
        append!(prog, _comp(fbody, name2, value2, cond))
    end
    @dbg println((:_comp_assignment,:return,prog))
    return prog
end

function _comp_if(expr::Expr,name::Dict,value::Dict,cond::Expr)
    @dbg println((:_comp_if,:expr,expr,:name,name,:value,value,:cond,cond))
    p = append!(_comp(expr.args[2],name,value,Expr(cond.head,cond.args...,expr.args[1])),
                length(expr.args) == 2 ? Any[] :
                _comp(expr.args[3],name,value,Expr(cond.head,cond.args...,Expr(:call,:!,expr.args[1]))))
    @dbg println((:_comp_if,:return,p))
    return p
end

function _comp_locals(ex)       # TODO: should we ignore conditionals, which are global?
    @dbg println((:_comp_locals,:ex,ex))
    l = (isa(ex, Symbol) ? Any[ex] :
         isa(ex, Number) ? Any[] :
         isa(ex,LineNumberNode) ? Any[] :
         !isa(ex, Expr) ? error("Expected Expr got $ex") :
         ex.head == :parameters ? Any[] :
         ex.head == :kw ? Any[] :
         ex.head == :return ? mapreduce(_comp_locals, append!, Any[:return], ex.args) :
         ex.head == :call ? mapreduce(_comp_locals, append!, Any[], ex.args[2:end]) :
         mapreduce(_comp_locals, append!, Any[], ex.args))
    @dbg println((:_comp_locals,:return,l))
    return l
end

# Given an expression or symbol and a dictionary of symbols, evaluate the
# expression in an environment defined by the dict, and Julia.
function _comp_eval(s,d::Dict)
    @dbg println((:_comp_eval,:s,s,:d,d))
    a = Expr(:let,Expr(:block,s))
    for (k,v) in d
        isa(v,Symbol) && (v=Expr(:quote,v)) # this is to prevent over-evaluation of (f=:sigm)
        push!(a.args, Expr(:(=),k,v))
    end
    ev = eval(current_module(), a)
    @dbg println((:_comp_eval,:return,ev))
    return ev
end

# Given a function definition, an array of argument symbols, x, and a
# return symbol, y, return a name dictionary that should be used when
# compiling the body of the function.
function _comp_fargs(f::Expr, x::Array, y::Symbol)
    @dbg println((:_comp_fargs,:f,f,:x,x,:y,y))
    (fname, fargs, fpars) = _comp_parse_call(f.args[1])
    length(fargs)==length(x) || error("parameter mismatch [$fargs] [$x]")
    fdict = Dict{Symbol,Symbol}()
    for i=1:length(x)
        (isa(x[i],Symbol) && isa(fargs[i],Symbol)) || error("parameter not symbol [$fargs] [$x]")
        fdict[fargs[i]] = x[i]
    end
    fdict[:return] = y
    @dbg println((:_comp_fargs,:return,fdict))
    return fdict
end

# Given a function definition and an array of keyword arguments found
# in the function call, returns a dictionary of keyword arguments that
# should be passed to the body of the function.  I tried writing my
# own function for this, but it is difficult to be consistent with the
# Julia implementation, e.g. having a kwarg value be in scope for a
# later kwarg etc.  It is easier to just construct a Julia function
# and see what it does.  Unfortunately anonymous functions do not
# currently accept kwargs
# (https://github.com/JuliaLang/julia/issues/2773), so we construct a
# named function.

function _comp_fpars(f::Expr, o)
    @dbg println((:_comp_fpars,:f,f,:o,o))
    (fname, fargs, fpars) = _comp_parse_call(f.args[1]) # these are from the original function definition
    ftemp = :_comp_fpars_tmp                            # constructing a new function definition
    fhead = Expr(:call, ftemp)
    !isempty(fpars) && push!(fhead.args, Expr(:parameters, fpars...)) # including fpars, excluding fargs
    fvars = map(s->Expr(:(=>), QuoteNode(s.args[1]), s.args[1]), fpars)
    fbody = Expr(:block, Expr(:call, :Dict, fvars...))
    fdefn = Expr(:function, fhead, fbody) # the new function returns a dictionary of fpars and their values
    # If we use Kfun instead of current_module to capture knet
    # functions, things go awry, e.g. a kwarg (k=a*b) ends up calling
    # Dot instead of regular multiply.  It is best not to mix Kfun and
    # Julia.
    eval(current_module(),fdefn)
    fdict = eval(current_module(),ftemp)(;o...)
    @dbg println((:_comp_fpars,:return,fdict))
    return fdict
end

function _comp_parse_def(f)
    @dbg println((:_comp_parse_def,:f,f))
    (isa(f,Expr) &&
     (f.head == :function || f.head == :(=)) &&
     (length(f.args) == 2) &&
     (f.args[1].head == :call) &&
     (f.args[2].head == :block)) ||
    error("Expected function definition got $f")
    (fhead, fbody) = f.args
    (fname, fargs, fpars) = _comp_parse_call(fhead)
    @dbg println((:_comp_parse_def,:return,:fname,fname,:fargs,fargs,:fpars,fpars,:fbody,fbody))
    (fname, fargs, fpars, fbody)
end

function _comp_parse_call(s)
    @dbg println((:_comp_parse_call,:s,s))
    (isa(s,Expr) && s.head == :call) ||
    error("Expected function call got $s")
    f = s.args[1]
    a = s.args[2:end]           # Use [:] notation to create a copy
    f == :+  ? _comp_parse_add(a...) :
    f == :.+ ? _comp_parse_badd(a...) :
    f == :-  ? _comp_parse_sub(a...) :
    f == :.- ? _comp_parse_bsub(a...) :
    f == :*  ? _comp_parse_mul(a...) :
    f == :.* ? _comp_parse_bmul(a...) :
    f == :/  ? _comp_parse_div(a...) :
    f == :./ ? _comp_parse_bdiv(a...) :
    _comp_parse_fcall(f,a)
end

# what if we have expression instead of symbol: define a compound type SymEx
# what if we have symbol instead of number (from compiler kwargs): can't happen, value subst does not take place with variables, only op args and func names

typealias SymEx Union{Symbol,Expr}

_comp_parse_add(x)=error("+$x")
_comp_parse_add(x,y)=_comp_parse_badd(x,y)
_comp_parse_add(x...)=_comp_parse_call(_comp_binarize(:+,x...))
_comp_parse_badd(x::SymEx,y::SymEx)=(:.+, Any[x,y], Any[])
_comp_parse_badd(x::Number,y::SymEx)=(:axpb, Any[y], Any[Expr(:kw,:b,x)])
_comp_parse_badd(x::SymEx,y::Number)=(:axpb, Any[x], Any[Expr(:kw,:b,y)])
_comp_parse_badd(x...)=error(".+$x")

_comp_parse_sub(x...)=_comp_parse_bsub(x...)
_comp_parse_bsub(x::SymEx,y::SymEx)=(:.+, Any[x,:(axpb($y;a=-1))], Any[])
_comp_parse_bsub(x::Number,y::SymEx)=(:axpb, Any[y], Any[Expr(:kw,:a,-1), Expr(:kw,:b,x)])
_comp_parse_bsub(x::SymEx,y::Number)=(:axpb, Any[x], Any[Expr(:kw,:b,-y)])
_comp_parse_bsub(x::SymEx)=(:axpb, Any[x], Any[Expr(:kw,:a,-1)])
_comp_parse_bsub(x...)=error(".-$x")

_comp_parse_mul(x::SymEx,y::SymEx)=(:*, Any[x,y], Any[])
_comp_parse_mul(x,y)=_comp_parse_bmul(x,y)
_comp_parse_mul(x...)=_comp_parse_call(_comp_binarize(:*,x...))
_comp_parse_mul(x)=error("*$x")
_comp_parse_bmul(x::SymEx,y::SymEx)=(:.*, Any[x,y], Any[])
_comp_parse_bmul(x::Number,y::SymEx)=(:axpb, Any[y], Any[Expr(:kw,:a,x)])
_comp_parse_bmul(x::SymEx,y::Number)=(:axpb, Any[x], Any[Expr(:kw,:a,y)])
_comp_parse_bmul(x...)=error(".*$x")

_comp_parse_div(x::SymEx,y::SymEx)=error("$x/$y")
_comp_parse_div(x,y)=_comp_parse_bdiv(x,y)
_comp_parse_div(x...)=error("/$x")
_comp_parse_bdiv(x::SymEx,y::SymEx)=(:.*, Any[x,:(axpb($y;p=-1))], Any[])
_comp_parse_bdiv(x::Number,y::SymEx)=(:axpb, Any[y], Any[Expr(:kw,:a,x),Expr(:kw,:p,-1)])
_comp_parse_bdiv(x::SymEx,y::Number)=(:axpb, Any[x], Any[Expr(:kw,:a,1/y)])
_comp_parse_bdiv(x...)=error("./$x")

_comp_binarize(f,x,y)=Expr(:call,f,x,y)
_comp_binarize(f,x,y,z,t...)=_comp_binarize(f,Expr(:call,f,x,y),z,t...)

function _comp_parse_fcall(f, a)
    x = Any[]
    o = Any[]
    for ai in a
        !isa(ai,Expr) ? push!(x,ai) :
        ai.head==:parameters ? append!(o, ai.args) :
        ai.head==:kw ? push!(o, ai) :
        push!(x, ai)
    end
    @dbg println((:_comp_parse_fcall,:return,:fname,f,:fargs,x,:fpars,o))
    (f, x, o)
end

function test_compiler()
    for n in names(Kfun,true)
        isa(Kfun.(n),DataType) || isa(Kfun.(n),Expr) || continue
        println(n)
        compile(n; out=10, f=:tanh)
    end
end

### DEAD CODE:

# # The compiler:
# _knet(x::Expr)=(x.head == :block ? _knet_bloc(x.args) : x.head == :(=) ? _knet_assn(x.args...) : error())
# _knet_bloc(x::Array)=mapreduce(_knet, append!, x)
# _knet_assn(s::Symbol, x::Expr)=(x.head == :call ? _knet_call(x.args...,s) : error())
# _knet_call(f, p::Expr, o...)=(p.head == :parameters ? _knet(eval(current_module(), Expr(:call, f, p, map(QuoteNode, o)...))) : error())
# _knet_call(f, o...)=_knet(eval(current_module(),Expr(:call,f,map(QuoteNode, o)...)))
# _knet(x::Tuple)=(isa(x[1],Op) ? Any[x] : error())
# _knet(::LineNumberNode)=Any[]

# ### Structure of an example function expression:
# #
# # function drop(x)
# #     if training
# #         r = rnd(; rgen=Bernoulli(0.5))
# #         return mul(r,x)
# #     elseif testing
# #         x = r
# #         return x
# #     else
# #         return x
# #     end
# # end
# #
# # (function[2], call[2], block[2])		# call for function header, block for body
# #   (call[2], :drop, :x)			# this is more complicated with keyword args
# #   (block[2], line, if[3])			# need to ignore line number elements
# #     (if[3], :training, block[4], block[2])	# if will have two args (cond,consequent) if no else part
# #       (block[4], line, =[2], line, return[1])
# #         (=[2], :r, call[2])
# #           (call[2],:rnd,parameters[1])
# #             (parameters[1],kw[2])		# keyword arguments following semicolon, see without semicolon
# #               (kw[2],:rgen,call[2])		# kw takes place of =
# #                 (call[2],:Bernoulli,0.5)
# #         (return[1], call[3])
# #           (call[3],:mul,:r,:x)
# #       (block[2], line, if[3])			# elseif turns into a simple if inside the else component
# #         (if, :testing, block[4], block[2])
# #           (block[4],line,=[2],line,return[1])
# #             (=[2],:x,:r)
# #             (return[1],:x)
# #           (block[2],line,return[1])
# #             (return[1],:x)

# # Julia parser output for function parameters:
# #
# # foo(x)	:call(foo,x)
# # foo(x,y)	:call(foo,x,y)
# # foo(x,y=0)	:call(foo,x,:kw(y,0))
# # foo(y=0)	:call(foo,:kw(y,0))
# # foo(;y=0)	:call(foo,:parameters(:kw(y,0)))
# # foo(x;y=0)	:call(foo,:parameters(:kw(y,0)),x)
# # foo(x,y=0;z=0):call(foo,:parameters(:kw(z,0)),x,:kw(y,0))
# # foo(x,y...)	:call(foo,x,:(...)(y))
# # foo(x;o...)	:call(foo,:parameters(:(...)(o)),x)
# # foo(x,y=0,z...;a,b=0,c...)	:call(foo,:parameters(a,:kw(b,0),:(...)(c)),x,:kw(y,0),:(...)(z))

# macro knet_old(f)
#     @assert f.head == :function "@knet macro requires a function definition as input."
#     @assert length(f.args) == 2
#     (fhead,fbody) = f.args

#     @assert fhead.head == :call
#     @assert length(fhead.args) >= 1
#     fname = fhead.args[1]
#     fargs = fhead.args[2:end]
#     if !isempty(fargs) && isa(fargs[1],Expr) && fargs[1].head == :parameters
#         fparameters = shift!(fargs)
#         fpars = map(fparameters.args) do k
#             @assert (isa(k,Expr) && isa(k.args[1], Symbol) && (k.head == :kw || (k.head == :(...) && k === fparameters.args[end]))) "Malformed keyword argument $k"
#             k.args[1]
#         end
#     else
#         fparameters = nothing
#         fpars = Any[]
#     end
#     @assert all(x->isa(x,Symbol), fargs) "Positional arguments must be symbols only, knet does not support default values or ... slurping before the semicolon."
#     @assert unique(fargs)==fargs "Positional arguments must be unique."

#     @assert fbody.head == :block
#     @show fvars = setdiff(_knet_vars(fbody), fargs)

#     freturn = gensym()
#     push!(fargs, freturn)
#     map!(s->Expr(:(::), s, :Symbol), fargs)
#     @show fargs

#     _knet_return(fbody, freturn)

#     # @show name = f.args[1].args[1]
#     # @show args = f.args[1].args[2:end]
#     # @show pars = _knet_pars(f)
#     # @show vars = _knet_vars(f)
#     # @assert !isempty(vars)

#     # @show head = Expr(:call, name, map(_knet_sym, [args; vars[end]])...)
#     # @show m = Expr(:macrocall, symbol("@gensym"), vars[1:end-1]...)
#     # @show q = Expr(:quote, _knet_esc(f.args[2], vcat(pars, vars)))
#     # body = Expr(:block, m, q)
#     # newf = Expr(:function, head, body)
#     # dump(newf,100)
#     # esc(newf)
# end

# "replace returns with assignment statements"
# function _knet_return(f::Expr, freturn::Symbol)
#     for a in f.args
#         if !isa(a, Expr)
#             continue
#         elseif a.head == :return
#             a.head = :(=)
#             unshift!(a.args, freturn)
#         else
#             _knet_return(a, freturn)
#         end
#     end
# end

# "find local variables that are assignment targets"
# function _knet_vars(f::Expr)
#     vars = Any[]
#     for a in f.args
#         if !isa(a, Expr)
#             continue
#         elseif a.head == :(=)
#             push!(vars, a.args[1])
#         else
#             append!(vars, _knet_vars(a))
#         end
#     end
#     return vars
# end

# function _knet_esc(x::Expr,v::Array)
#     if x.head == :line
#         x
#     elseif x.head == :kw
#         Expr(x.head, x.args[1], map(a->_knet_esc(a,v), x.args[2:end])...)
#     else
#         Expr(x.head, map(a->_knet_esc(a,v), x.args)...)
#     end
# end

# # TODO: add checks when s is not an element of v
# _knet_esc(s::Symbol,v::Array)=(in(s,v) ? Expr(:$, s) : s)
# _knet_esc(x,v::Array)=x

# _knet_sym(s::Symbol)=Expr(:(::), s, :Symbol)
# _knet_sym(x)=x

# # "positional and keyword arguments"
# # function _knet_pars(f::Expr)
# #     a = f.args[1].args[2:end]
# #     all(x->isa(x,Symbol), a) && return a
# #     @assert isa(a[1],Expr) && a[1].head == :parameters && all(x->isa(x,Symbol), a[2:end])
# #     vcat(a[2:end], map(kw->kw.args[1], a[1].args))
# # end

# # The compiler:
# _knet(x::Expr)=(x.head == :block ? _knet_bloc(x.args) : x.head == :(=) ? _knet_assn(x.args...) : error())
# _knet_bloc(x::Array)=mapreduce(_knet, append!, x)
# _knet_assn(s::Symbol, x::Expr)=(x.head == :call ? _knet_call(x.args...,s) : error())
# _knet_call(f, p::Expr, o...)=(p.head == :parameters ? _knet(eval(current_module(), Expr(:call, f, p, map(QuoteNode, o)...))) : error())
# _knet_call(f, o...)=_knet(eval(current_module(),Expr(:call,f,map(QuoteNode, o)...)))
# _knet(x::Tuple)=(isa(x[1],Op) ? Any[x] : error())
# _knet(::LineNumberNode)=Any[]

# """
# _knet(x) is a compiler that takes the expressions generated by a
# @knet function and outputs a list of machine instructions of the form:
# ```
# (Op, in1, in2, ..., out)
# ```
# where Op is a primitive operator, in1, in2, ..., out are symbols
# representing input and output registers.

# The whole compiler is 7 lines long!

# Example:

# (1) The user defines a @knet function:
# ```
# @knet function wb(x; o...)
#     y = wdot(x; o...)
#     z = bias(y; o...)
# end
# ```

# (2) The @knet macro turns this into a Julia function which takes
# input/output symbols and generates an expression sequence using these
# symbols.  The locals are replaced by gensyms:

# ```
# function wb(x::Symbol, z::Symbol; o...)
#     @gensym y
#     quote
#         \$y = wdot(\$x; \$o...)
#         \$z = bias(\$y; \$o...)
#     end
# end
# ```

# (3) The user calls a Net constructor, e.g. FNN(wb; out=100).  The net
# constructor (or recursively the compiler) runs the Julia function `wb`
# to get:

# ```
# julia> prog = wb(:a,:b; out=100)
# quote
#     ##y#8260 = wdot(a; Any[(:out,100)]...)
#     b = bias(##y#8260; Any[(:out,100)]...)
# end
# ```

# (4) This gets passed to the _knet compiler to get:
# ```
# julia> _knet(prog)
#  (Par((100,0),Gaussian(0,0.01),...),symbol("##w#8267"))
#  (Dot(),symbol("##w#8267"),:a,symbol("##y#8262"))
#  (Par((0,),Constant(0),...),symbol("##b#8270"))
#  (Add(),symbol("##b#8270"),symbol("##y#8262"),:b)
# ```

# (5) The Net constructor adds input expressions and replaces symbols with Ints to get:
# ```
# net.op = [ Input(), Par((100,0),...), Dot(), Par((0,),...), Add() ]
# net.inputs = [ [], [], [2,1], [], [4,3] ]
# ```
# """
# _knet

# """
# @knet macro -- This is what the user types:
# ```
# @knet function wdot(x; out=0, winit=Gaussian(0,.01), o...)
#     w = par(out,0; init=winit, o...)
#     y = dot(w, x)
# end
# ```
# This is what the @knet macro turns it into:
# ```
# function wdot(x::Symbol, y::Symbol; out=0, winit=Gaussian(0,.01), o...)
#     @gensym w
#     quote
#         \$w = par(\$out,0; init=\$winit, \$o...)
#         \$y = dot(\$w, \$x)
#     end
# end
# ```
# """
# :@knet

# :ok

    # fparameters contain the parameters in the function call
    # their values might be unevaluated
    # o... contains the parameters from the compiler call
    # the op constructor should be called with ?

# pass down pars with appropriate overrides and evaluation
# compiler options example:
# comp(f; hidden=100, lr=0.1)
# f = :(function f(x; hidden=0,o...); w = par(;out=hidden,o...); return dot(w,x); end)
# par should we called with out=100, lr=0.1
# if f(x; hidden=0) then par should be called with out=100 but no lr
# so the pars that get passed to the body of f is determined by the header of f
# if the header contains o..., then all pars are passed, overriding any defaults in f header

# function call example
# > w = par(; out=10, o...)
# function definition example
# > function wdot(x; out=10, o...)
# compiler call example
# > comp(wdot; out=10, o...)
# Differences:
# we can have o... in the middle in function call: par(; o..., out=10)
# we can have more than one splat: par(; o..., p..., out=10)
# Difference in terms of what should be passed as kwargs?
# Going from a comp call to function definition we excluded vars not in function kwarg list
# We treated the comp call as if it were a function call for wdot
# So it should be an error to have variables wdot does not expect

# 1. comp(wdot; hidden=10, lr=1) is called
# .. wdot(x; hidden=0, winit=Gaussian(0,.01), o...) is the signature of wdot
# 2. we construct a dictionary of keyword args for wdot from the comp call and wdot signature
# .. we pass this dictionary down to _comp...
# 3. in the body we have: w = par(; out=hidden, init=winit, o...)
# 4. we need to call par with (out=10,lr=1)
# this means the rhs of par keywords or par slurps need to be evaluated in an environment
# extended by our dict.

# comp(f; o...) => o: compiler options
# forw(f,x; o...) => o: runtime options
# so we can do forw(f,x; training=true, encoding=false) etc.
# we already do this!
# so forw has to evaluate conditions with respect to o... and we are all set!
# do we compile with the @knet macro or the comp function?

# function _comp_old(s::Expr, r::Symbol, c::Expr, d::Dict)
#     if s.head == :block
#         _comp_block(s, r, c, d)
#     elseif s.head == :if
#         _comp_if(s, r, c, d)
#     elseif s.head == :return
#         _comp_return(s, r, c, d)
#     elseif s.head == :(=)
#         _comp_assign(s, r, c, d)
#     else
#         @show (:unknown, s)
#     end
# end

# function _comp_block(s::Expr,r::Symbol,c::Expr,d::Dict)
#     mapreduce(x->_comp(x,r,c,d), append!, s.args)
# end

# function _comp_if(s::Expr,r::Symbol,c::Expr,d::Dict)
#     append!(_comp(s.args[2],r,Expr(c.head,c.args...,s.args[1]), d),
#             length(s.args) == 2 ? Any[] :
#             _comp(s.args[3],r,Expr(c.head,c.args...,Expr(:!,s.args[1])), d))
# end

# function _comp_return(s::Expr,r::Symbol,c::Expr,d::Dict)
#     _comp(Expr(:(=), r, s.args[1]),r,c,d)
# end

# function _comp_assign(s::Expr,r::Symbol,c::Expr,d::Dict)
#     @assert length(s.args)==2
#     (lhs,rhs) = s.args
#     @assert isa(lhs, Symbol)
#     if isa(rhs, Symbol)
#         @show (:assign_symbol, s,r,c,d); Any[]
#     else
#         @assert isa(rhs, Expr)
#         @assert rhs.head == :call
#         func = eval(current_module(), rhs.args[1])
#         if isa(func, Expr)
#             _comp_sub(s,r,c,d)
#         elseif isa(func, DataType) && func <: Op
#             _comp_op(s,r,c,d)
#         else
#             @show (:assign_unknown, s,r,c,d); Any[]
#         end
#     end
# end

# "compile: w = par(; out=10, o...)"
# function _comp_op(s::Expr,r::Symbol,c::Expr,d::Dict)
#     (lhs,rhs) = s.args
#     a = rhs.args
#     if length(a) > 1 && isa(a[2],Expr) && a[2].head == :parameters
#         fcall = a[1:2]
#         fargs = a[3:end]
#     else
#         fcall = a[1:1]
#         fargs = a[2:end]
#     end
#     @assert all(x->isa(x,Symbol), fargs) "Positional arguments must be symbols only, knet does not support default values or ... slurping before the semicolon."
#     @assert unique(fargs)==fargs "Positional arguments must be unique."
#     op = _comp_eval(Expr(:call,fcall...), d)
#     Any[tuple(c, lhs, op, fargs...)]
# end

# f (wdot) is a function definition.
# compile the body of f, but:
# returnvar needs to be lhs
# args in function def need to be replaced with fargs
# local vars in function def need to be gensym
# fpar dict needs to be handled
# we just need the instructions back.
# _comp the body
# we need two dictionaries: symbol replacement dict, and eval dict
# when we see y=wdot(x; o...) in the body we need to replace y, x, and o.
# when we see z = f(y; o...) find y and z in vars, f and o in dict
# so target and input vars in vars, all else eval
# initial: comp gets fname and kwargs
# kwargs matches kwpars in f definition, may have extras if o...

# problem symbols like conv and dot are defined in base and cannot be overwritten
# solution: make the compiler look them up in a special primitives table
# that special table can also include operators like '+' or '*'

# However if we do use a primitives table and do not define symbols
# like dot in Julia then we can't pass them using comp(wbf;f=dot)

# """

# Compiling a single instruction::

#     y=f(x1,x2,...;o1=v1,o2=v2,on...)

# when f is a primitive operator or
# when the function definition for f is::

#     function f(z1,z2,...;p1=w1,p2=w2,pn...)

# """
# _comp

# The core of the compiler is the compilation of a function call.  If
# the function is a primitive, this gets compiled into a single
# instruction.  If it is a compound, then the body of the function
# with certain transformations gets compiled into an array of
# instructions.

# The compilation of the function call takes place in a lexical
# environment.  There are several components of this environment:

# 1. The target variable.  The function call can be the rhs of an
# assignment statement, an internal call in a compound statement, or
# we may be compiling the top level function as a model.  In each case
# a target symbol should be provided to _comp_call.  Any return
# statement in the body of the function will be turned into an
# assignment to this target symbol.  But see #3.

# 2. The function name.  This could be a symbol in _kenv, which is a
# table of knet functions and their names, or a kwarg symbol from the
# parent function, in which case it should be evaluated.  The kwarg
# takes precedence.

# 3. The positional arguments.  These are either symbols or other
# calls.  To compile other calls we gensym target variables and insert
# their instructions preceding the current function call.  Once each
# argument turns into a symbol, these symbols will need to replace the
# parameters of the function call in the body.  Recursively this means
# we need a symbol map passed to _comp_call.  The target variable can
# also be in this map as the value of :return.

# 4. The keyword arguments.  These come as symbol=expr pairs,
# optionally followed by a o... splat.  The expr can be a kenv symbol,
# or a julia expression which should be evaluated in that order.
# These symbols should go into a symbol table and be used while
# compiling the body of the function.

# What is the simplest way to do this?  How many separate symtables do
# we need?  Can we combine symbol lookup and expr eval?

# eval args until all are symbols => we need a real _comp_call?
# should we do syntactic xforms?

# fargs or fpars could be empty depending on whether the user uses semicolon
# they may also have plain symbols or splats, which we should warn the user about
# fname could be a primitive op or a kfun, but it should be stored in _kenv in either case
# f = get(_kenv, fname, nothing)
# f == nothing && error("$fname not defined as a knet function")
# # in either case we should be able to compile it
# @show global tmp =
# fcomp = eval(current_module(), tmp)

# user defines kfun using @knet function foo...
# user compiles kfun calling @knet kfun(a=1,b=2...)
# this should return a net.
# user uses kfun in other code using kfun(x,y;a=1,b=2...)
# the two usages may be confusing.
# we could just provide compile(:kfun; args...) instead of the knet macro

# function _knet(f)
#     (fname, fargs, fpars) = _comp_parse_call(f)
#     haskey(_kenv, fname) || error("$fname not defined as a knet function")
#     isempty(fpars) && ((fargs,fpars)=(fpars,fargs))
#     isempty(fargs) || error("Usage: @knet f(a=1,b=2,...)")
#     ex = Expr(:call,:comp,Expr(:parameters,fpars...),Expr(:ref,:_kenv,QuoteNode(fname)))
#     fcomp = eval(current_module(), ex)
# end

# macro kdef(f)
#     _kdef(f)
# end

# # """
# # The @knet macro compiles a knet function as a model:

# #     net = @knet foo(a=1,b=2)

# # compiles the knet function foo as a model with initialization
# # parameters `a=1` and `b=2` and assigns the result to the Julia
# # variable `net`.
# # """
# macro knet(f)
#     if f.head==:function
#         _kdef(f)
#     elseif f.head == :(=)
#         _kdef(f)
#     elseif f.head == :call
#         _knet(f)
#     end
# end

    # inputs = filter(x->isa(x.op,Input), op)
    # params = filter(x->isa(x.op,Par), op)

# """
# _kops lists primitive operations and their knet names.  The primitive
# operations (capitalized in the code) are represented by types that are
# subtypes of Op.
# """
# _kops = [(:add,Add),(:+,Add),   # TODO: catch add with alpha/beta
#          (:arr,Arr),
#          (:axpb,Axpb),          # TODO: catch axpb
#          (:conv,Conv),
#          (:dot,Dot),(:*,Dot),
#          (:input,Input),
#          (:logp,Logp),
#          (:mul,Mul),(:.*,Mul),  # TODO: define div,sub
#          (:nce,Nce),
#          (:par,Par),
#          (:pool,Pool),
#          (:relu,Relu),
#          (:rnd,Rnd),
#          (:sigm,Sigm),
#          (:soft,Soft),
#          (:tanh,Tanh)]

#     feval = _comp_func(f, value)
# # TODO: not sure if we need this special treatment:
# # TODO: this forces quotation f=:relu, is there another way?
# function _comp_func(s,d::Dict)
#     @dbg println((:_comp_func,:s,s,:d,d))
#     f = s
#     haskey(d, f) && (f = d[f])
#     haskey(_KENV, f) && (f = _KENV[f])
#     @dbg println((:_comp_func,:return,f))
#     return f
# end

# macro knet(f)
#     (fname, fargs, fpars, fbody) = _comp_parse_def(f)
#     esc(Expr(:(=),Expr(:ref,:_KENV,Expr(:quote,fname)),Expr(:quote,f)))
# end

# function _comp_fpars(f::Expr, o)
#     @dbg println((:_comp_fpars,:f,f,:o,o))
#     (fname, fargs, fpars) = _comp_parse_call(f.args[1]) # these are from the function definition
#     slurp = (!isempty(fpars) && isa(fpars[end],Expr) && fpars[end].head == :(...) ? pop!(fpars).args[1] : nothing)
#     fdict = Dict{Symbol,Any}()
#     for k in fpars
#         (isa(k, Expr) && isa(k.args[1], Symbol) && k.head == :kw) || error("Malformed keyword argument $k")
#         fdict[k.args[1]] = eval(Kfun, k.args[2])
#         # (isa(k.args[2],Symbol) && haskey(_KENV,k.args[2]) ? _KENV[k.args[2]] :
#         #  eval(current_module(), k.args[2]))
#     end
#     if slurp != nothing
#         fdict[slurp] = Any[]
#     end
#     for (k,v) in o
#         if haskey(fdict, k)
#             fdict[k] = v
#         elseif slurp != nothing
#             push!(fdict[slurp], (k,v))
#         else
#             error("Unrecognized keyword argument \"$k\"")
#         end
#     end
#     @dbg println((:_comp_fpars,:return,fdict))
#     return fdict
# end

# the compiler turns expressions into low level instruction sequences.
# + variable renaming in function calls.
# + inheriting conditions
# + inheriting keyword arguments
# + replacing returns with variables
# + recording of arg names at top level
# + preserve local and arg variable names of top level function
# + representation of primitive ops
# + ninputs problem: look up the definition or ask the Op
# + can we compile primitive operators into net: why not
# + separate env table for knet ops and funs: cleaner but cannot pass name to Net or comp, or f=relu type kwarg unless we write macros.
# + adapt to the new _kenv structure: need next pointers?
# + compound operations
# + arithmetic operators
# + multiple functions sharing registers vs one with conditionals
# + instead of lots of arrays in net, have an Instruction type with the necessary fields: fix the rest of the code in net/
# + wf and f=:sigm not working. the dict that comes to _comp_eval should splat :o.
# + it works in recursive call from lstm to add2

:ok


text:
# Even though shapes are specified with real geometric parameters, the size of voxels can be
# complex when the underlying space is stretched by complex factors.  So, we allow complex
# `vxl` below.  Even in such a case, r₀ and nout are still real.  (Maybe this condition
# should be relaxed later.)

export volfrac

const X, Y, Z = 1, 2, 3
const XYZ, YZX, ZXY = (X,Y,Z), (Y,Z,X), (Z,X,Y)
const UVW = YZX, ZXY, XYZ
const N, P = 1, 2  # negative, positive
const NP = (N, P)

corner(vxl::NTuple{2,SVector{3,<:Number}}, sx::Integer, sy::Integer, sz::Integer) =
    SVector(vxl[sx][X], vxl[sy][Y], vxl[sz][Z])

# Calculate the bit array that indicates corner contained-ness.
#
# The bit array is 8 bits (= 1 byte).  The kth bit is 1 if the kth corner is contained in
# the half space defined by the plane, an 0 otherwise.  The 8 corners of the voxel are
# indexed 1 through 8 in the order of (---), (+--), (-+-), (++-), (--+), (+-+), (-++), (+++),
# where, e.g., (+--) indicates the corner at the +x, -y, -z corner.
function corner_bits(vxl::NTuple{2,SVector{3,<:Number}},  # two ends of solid diagonal of voxel
                     nout::SVector{3,<:Real}, # unit outward normal of plane
                     nr₀::Real)  # equation of plane: nout⋅(r - r₀) = 0, or nout⋅r = nr₀
    cbits = 0x00
    bit = 0x01
    n_on = 0
    for sz = NP, sy = NP, sx = NP
        r = corner(vxl, sx, sy, sz)
        nr = nout⋅r
        if nr ≤ nr₀  # corner is contained and not on boundary (nout is outward normal)
            cbits |= bit
            n_on += nr==nr₀  # corner is on boundary
        end
        bit <<= 1
    end

    return cbits, n_on
end

# Determine if the corner contained-ness bit array corresponds to the case where the plane
# crosses one of the three sets of four parallel edges of the voxel.
function isquadsect(cbits::UInt8)
    return cbits==0x0F || cbits==~0x0F || cbits==0x33 || cbits==~0x33 || cbits==0x55 || cbits==~0x55

    # Equivalent to
    # n = count_ones(cbits)
    # m = count_ones(cbits & 0x0F)
    # return n==4 && iseven(m)
end

# For the cases where the plane crosses a set of four parallel edges of the voxel, determine
# which direction those edges lie.
function edgedir_quadsect(cbits::UInt8)
    if cbits==0x0F || cbits==~0x0F
        dir = Z
    elseif cbits==0x33 || cbits==~0x33
        dir = Y
    else
        @assert cbits==0x55 || cbits==~0x55
        dir = X
    end

    return dir
end

# Return the volume fraction when the plane croses a set of four parallel edges.
function rvol_quadsect(vxl::NTuple{2,SVector{3,<:Number}}, nout::SVector{3,<:Real}, nr₀, cbits::UInt8)
    w = edgedir_quadsect(cbits)
    ∆w = vxl[P][w] - vxl[N][w]

    u, v, _w = UVW[w]
    nu, nv, nw = nout[u], nout[v], nout[w]
    mean_cepts = 4nr₀
    for sv = NP, su = NP
        mean_cepts -= nu*vxl[su][u] + nv*vxl[sv][v]
    end
    mean_cepts /=  nw * 4∆w

    sw = nw>0 ? N : P  # nw cannot be 0
    return abs(mean_cepts - vxl[sw][w]/∆w)
end

# Calculate the volume fraction for the most general cases, by cutting out corners from a
# triangular pyramid.
#
# Assume count_ones(cbits) ≤ 4.  Othewise, call this function with flipped nout, nr₀,
# cbits.
function rvol_gensect(vxl::NTuple{2,SVector{3,<:Number}}, nout::SVector{3,<:Real}, nr₀::Real, cbits::UInt8)
    s = (nout.<0) .+ 1
    c = corner(vxl, s[X], s[Y], s[Z])  # corner coordinates
    ∆ = vxl[P] - vxl[N]  # vxl edges
    nc = nout .* c
    rmax, rmid, rmin =  abs.(((nr₀-sum(nc)) .+ nc) ./ nout - c) ./ ∆ # (lengths from corner to intercetps) / (voxel edges)

    # nx, ny, nz = nout
    # sx, sy, sz = ((nx≥0 ? N : P), (ny≥0 ? N : P), (nz≥0 ? N : P))  # signs of corner
    # cx, cy, cz = vxl[nX][sx], vxl[nY][sy], vxl[nZ][sz]  # corner coordinates
    # ∆x, ∆y, ∆z = vxl[nX][nP]-vxl[nX][nN], vxl[nY][nP]-vxl[nY][nN], vxl[nZ][nP]-vxl[nZ][nN]  # voxel edges
    # nxcx, nycy, nzcz = nx*cx, ny*cy, nz*cz
    # rmax, rmid, rmin =  # (lengths from corner to intercetps) / (voxel edges)
    #     abs((nr₀-nycy-nzcz)/nx-cx)/∆x, abs((nr₀-nzcz-nxcx)/ny-cy)/∆y, abs((nr₀-nxcx-nycy)/nz-cz)/∆z

    # Sort rmax, rmin, rmin properly.
    if rmax < rmid; rmax, rmid = rmid, rmax; end
    if rmid < rmin; rmid, rmin = rmin, rmid; end
    if rmax < rmid; rmax, rmid = rmid, rmax; end

    # Calculate the volume of the triangular pyramid, and cut off appropriate corners.
    tmax = 1 - 1/rmax
    rvol_core = 1 + tmax + tmax^2

    # Below, if rmax == Inf, rmid and rmin must be ≤ 1 in exact arithmetic and subtraction
    # must not occur, but they can be > 1 in floating-point arithmetic and lead to large
    # subtraction.  Prevent this by subtracting only for rmax ≠ Inf.
    isfin_rmax = !isinf(rmax)
    if rmid > 1 && isfin_rmax
        tmid = 1 - 1/rmid
        rvol_core -= rmax * tmid^3
    end
    if rmin > 1 && isfin_rmax
        tmin = 1 - 1/rmin
        rvol_core -= rmax * tmin^3
    end

    return rvol_core * rmin * rmid / 6
end

"""
    volfrac(voxel, nout, r₀)

Returns the volume fraction `rvol = vol(voxel ⋂ half-space) / vol(voxel)` that indicates how
much portion of the volume of a voxel is included the half-space defined by a plane.  The
result is between 0.0 and 1.0, inclusive.  Inside and outside of the half space is
determined by `nout`, which is the outward normal.  (If a point is on the `nout` side of the
plane, then it is outside.)

The voxel needs to be aligned with the Cartesian axes, and is described by `v =
([xn,yn,zn], [xp,yp,zp])` that specifies the x-, y-, z-boundaries of the voxel.  The
half-space is described by the boundary plane.  The boundary plane is described by the
outward normal vector `nout = [nx, ny, nz]` and a point `r₀ = [rx, ry, rz]` on the plane.
`nout` does not have to be normalized.
"""
function volfrac(vxl::NTuple{2,SVector{3,<:Number}}, nout::SVector{3,<:Real}, r₀::SVector{3,<:Real})
    nr₀ = nout⋅r₀
    cbits, n_on = corner_bits(vxl, nout, nr₀)
    n_in = count_ones(cbits)  # number of corners contained

    if n_in == 8  # voxel is inside half-space
        rvol = 1.
    elseif n_in - n_on == 0  # voxel is outside half-space
        rvol = 0.
    elseif isquadsect(cbits) # plane crosses a set of four parallel edges of voxel
        rvol = rvol_quadsect(vxl, nout, nr₀, cbits)
    elseif n_in ≤ 4 # general cases with n_in ≤ 4
        rvol = rvol_gensect(vxl, nout, nr₀, cbits)
    else  # general cases with n_in ≥ 5
        @assert n_in ≥ 5
        rvol = 1. - rvol_gensect(vxl, .-nout, -nr₀, ~cbits)
    end

    return rvol
end

# volfrac() for a 2D voxel (= pixel).
# - Turn the nout and r₀ into 3D vectors on the xy-plane.
# - Turn the pixel into a 3D voxel whose base is the pixel and the height is 1.
# - Pass these 3D objects to volfrac() for 3D.
volfrac(vxl::NTuple{2,SVector{2,<:Number}}, nout::SVector{2,<:Real}, r₀::SVector{2,<:Real}) =
    volfrac((SVector(vxl[N][1], vxl[N][2], 0), SVector(vxl[P][1], vxl[P][2], 1)),
            SVector(nout[1], nout[2], 0),
            SVector(r₀[1], r₀[2], 0))

# volfrac() for 1D voxel (= line segment).
# - Turn the nout and r₀ into 3D vectors along the z-axis
# - Turn the line segment into a 3D voxel whose height is the line segment and the base is a 1×1 pixel.
# - Pass these 3D objects to volfrac() for 3D.
#
# volfrac() for 1D can be implemented from scratch without calling volfrac() for 3D, but
# let's keep it this way for now.
volfrac(vxl::NTuple{2,SVector{1,<:Number}}, nout::SVector{1,<:Real}, r₀::SVector{1,<:Real}) =
    volfrac((SVector(0, 0, vxl[N][1]), SVector(1, 1, vxl[P][1])),
            SVector(0, 0, nout[1]),
            SVector(0, 0, r₀[1]))

ext: jl


text:
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    DirectionPartition(direction; tol=1e-6)

A method for partitioning spatial objects along a given `direction`
with bandwidth tolerance `tol`.
"""
struct DirectionPartition{Dim,T} <: SPredicatePartitionMethod
  direction::Vec{Dim,T}
  tol::Float64

  function DirectionPartition{Dim,T}(direction, tol) where {Dim,T}
    new(normalize(direction), tol)
  end
end

DirectionPartition(direction::Vec{Dim,T}; tol=1e-6) where {Dim,T} =
  DirectionPartition{Dim,T}(direction, tol)

DirectionPartition(direction::NTuple{Dim,T}; tol=1e-6) where {Dim,T} =
  DirectionPartition(Vec(direction), tol=tol)

function (p::DirectionPartition)(x, y)
  δ = x - y
  d = p.direction
  norm(δ - (δ⋅d)*d) < p.tol
end


text:
struct PoissonVirtualElement{dim} <: AbstractVirtualElement
    degree::Int
end

function PoissonVirtualElement(dim = 2, degree=1)
    PoissonVirtualElement{dim}(degree)
end

get_degree(el::PoissonVirtualElement) = el.degree

struct LocalPoissonVirtualElement
    degree::Int
    Pk_basis::Monomials
    edge_quad::QuadratureRule
end

function LocalPoissonVirtualElement(dim, degree, centroid, diameter)
    Pk_basis = Monomials(dim,degree,centroid, diameter)
    if degree == 1
        points = Vector{Float64}(); weights = Vector{Float64}()
    else
        points, weights = FastGaussQuadrature.gausslobatto(2*degree-1)
        # Shift interval from (-1,1) to (0,1)
        weights *= 0.5
        points = points .+ 1.0; points /= 2.0
    end
    quad = QuadratureRule{Segment,1,Float64}(weights, [Tensors.Tensor{1,1}([x]) for x in points])
    LocalPoissonVirtualElement(degree, Pk_basis, quad)
end

get_degree(el::LocalPoissonVirtualElement) = el.degree

function gettopology(mesh::AbstractPolytopalMesh{2},cell, element::PoissonVirtualElement)
    V = getnvertices(mesh,cell)
    deg = get_degree(element)
    Dict(0=> V, 1=> V*(deg-1), 2=> Int((deg-1)*deg/2))
end

function getnlocaldofs(element::PoissonVirtualElement{2}, cell::Cell{2,V,E,M}) where {V,E,M}
    k = get_degree(element)
    V*k+(k-1)*k/2
end

function spatial_nodal_coordinate(mesh, ci::Int,element::PoissonVirtualElement{2},i::Int) where {V,E,M}
    if i <= getnvertices(mesh,getcells(mesh)[ci]) #Vertices coordinates
        return getvertexcoords(mesh,getcells(mesh)[ci],i)
    else     #Edge points coordinates
        cell = getcells(mesh)[ci]
        nv = getnvertices(mesh,cell)
        cell_topology = gettopology(mesh,cell, element)
        dofs_per_edge = Int(cell_topology[1]/getnedges(mesh,cell))
        ei = div((i - nv)-1,dofs_per_edge)+1
        val = 1/(dofs_per_edge+1)*mod1((i-nv),dofs_per_edge)
        map_unit_to_segment(Tensors.Vec{1}((val,)), getverticescoords(mesh,EdgeIndex(ci,ei)))
    end

end

max_stars_repo_name: Paulms/jFEMToools

max_issues_repo_name: Paulms/jFEMToools

max_forks_repo_name: Paulms/jFEMToools


text:
# The Bellman Ford algorithm for single-source shortest path

###################################################################
#
#   The type that capsulates the states of Bellman Ford algorithm
#
###################################################################

type NegativeCycleError <: Exception end

type BellmanFordStates{V,D<:Number}
    parents::Vector{V}
    dists::Vector{D}
end

# create Bellman Ford states

function create_bellman_ford_states{V,D<:Number}(g::AbstractGraph{V}, ::Type{D})
    n = num_vertices(g)
    parents = Array(V, n)
    dists = fill(typemax(D), n)

    BellmanFordStates(parents, dists)
end

function bellman_ford_shortest_paths!{V,D}(
    graph::AbstractGraph{V},
    edge_dists::AbstractEdgePropertyInspector{D},
    sources::AbstractVector{V},
    state::BellmanFordStates{V,D})

    @graph_requires graph incidence_list vertex_map vertex_list
    edge_property_requirement(edge_dists, graph)

    active = Set{V}()
    for v in sources
        i = vertex_index(v, graph)
        state.dists[i] = 0
        state.parents[i] = v
        push!(active, v)
    end
    no_changes = false
    for i in 1:num_vertices(graph)
        no_changes = true
        new_active = Set{V}()
        for u in active
            uind = vertex_index(u, graph)
            for e in out_edges(u, graph)
                v = target(e, graph)
                vind = vertex_index(v, graph)
                edist = edge_property(edge_dists, e, graph)
                if state.dists[vind] > state.dists[uind] + edist
                    state.dists[vind] = state.dists[uind] + edist
                    state.parents[vind] = u
                    no_changes = false
                    push!(new_active, v)
                end
            end
        end
        if no_changes
            break
        end
        active = new_active
    end
    if !no_changes
        throw(NegativeCycleError())
    end
    state
end

function bellman_ford_shortest_paths{V,D}(
    graph::AbstractGraph{V},
    edge_dists::AbstractEdgePropertyInspector{D},
    sources::AbstractVector{V})
    state = create_bellman_ford_states(graph, D)
    bellman_ford_shortest_paths!(graph, edge_dists, sources, state)
end

function bellman_ford_shortest_paths{V,D}(
    graph::AbstractGraph{V},
    edge_dists::Vector{D},
    sources::AbstractVector{V})
    edge_inspector = VectorEdgePropertyInspector{D}(edge_dists)
    bellman_ford_shortest_paths(graph, edge_inspector, sources)
end

function has_negative_edge_cycle{V, D}(
    graph::AbstractGraph{V},
    edge_dists::AbstractEdgePropertyInspector{D})
    try
        bellman_ford_shortest_paths(graph, edge_dists, vertices(graph))
    catch e
        if isa(e, NegativeCycleError)
            return true
        end
    end
    return false
end

function has_negative_edge_cycle{V, D}(
    graph::AbstractGraph{V},
    edge_dists::Vector{D})
    edge_inspector = VectorEdgePropertyInspector{D}(edge_dists)
    has_negative_edge_cycle(graph, edge_inspector)
end


text:
"""
Server example that takes after the simple server, however,
handles dealing with CORS preflight headers when dealing with more
than just a simple request
"""
using HTTP

# modified Animal struct to associate with specific user
mutable struct Animal
    id::Int
    userId::Base.UUID
    type::String
    name::String
end

# use a plain `Dict` as a "data store"
const ANIMALS = Dict{Int, Animal}()
const NEXT_ID = Ref(0)
function getNextId()
    id = NEXT_ID[]
    NEXT_ID[] += 1
    return id
end

# CORS headers that show what kinds of complex requests are allowed to API
headers = [
    "Access-Control-Allow-Origin" => "*",
    "Access-Control-Allow-Headers" => "*",
    "Access-Control-Allow-Methods" => "POST, GET, OPTIONS"
]

#=
JSONHandler minimizes code by automatically converting the request body
to JSON to pass to the other service functions automatically. JSONHandler
recieves the body of the response from the other service funtions and sends
back a success response code
=#
function JSONHandler(req::HTTP.Request)
    # first check if there's any request body
    body = IOBuffer(HTTP.payload(req))
    if eof(body)
        # no request body
        response_body = handle(ANIMAL_ROUTER, req)
    else
        # there's a body, so pass it on to the handler we dispatch to
        response_body = handle(ANIMAL_ROUTER, req, JSON3.read(body, Animal))
    end
    return HTTP.Response(200, JSON3.write(response_body))
end

#= CorsHandler: handles preflight request with the OPTIONS flag
If a request was recieved with the correct headers, then a response will be
sent back with a 200 code, if the correct headers were not specified in the request,
then a CORS error will be recieved on the client side

Since each request passes throught the CORS Handler, then if the request is
not a preflight request, it will simply go to the JSONHandler to be passed to the
correct service function =#
function CorsHandler(req)
    if HTTP.hasheader(req, "OPTIONS")
        return HTTP.Response(200, headers = headers)
    else
        return JSONHandler(req)
    end

# **simplified** "service" functions
function createAnimal(req::HTTP.Request, animal)
    animal.id = getNextId()
    ANIMALS[animal.id] = animal
    return animal
end

function getAnimal(req::HTTP.Request)
    animalId = HTTP.URIs.splitpath(req.target)[5] # /api/zoo/v1/animals/10, get 10
    return ANIMALS[animalId]
end

function updateAnimal(req::HTTP.Request, animal)
    ANIMALS[animal.id] = animal
    return animal
end

function deleteAnimal(req::HTTP.Request)
    animalId = HTTP.URIs.splitpath(req.target)[5] # /api/zoo/v1/animals/10, get 10
    delete!(ANIMALS, animal.id)
    return ""
end

# add an additional endpoint for user creation
HTTP.@register(ANIMAL_ROUTER, "POST", "/api/zoo/v1/users", createUser)
# modify service endpoints to have user pass UUID in
HTTP.@register(ANIMAL_ROUTER, "GET", "/api/zoo/v1/users/*/animals/*", getAnimal)
HTTP.@register(ANIMAL_ROUTER, "DELETE", "/api/zoo/v1/users/*/animals/*", deleteAnimal)

HTTP.serve(CorsHandler, ip"127.0.0.1", 8080)


text:
module TestExamples3DCurved

using Test
using Trixi

include("test_trixi.jl")

# pathof(Trixi) returns /path/to/Trixi/src/Trixi.jl, dirname gives the parent directory
EXAMPLES_DIR = joinpath(pathof(Trixi) |> dirname |> dirname, "examples", "3d")

@testset "Curved mesh" begin
  @trixi_testset "elixir_advection_basic_curved.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_basic_curved.jl"),
      l2   = [0.00013446460962856976],
      linf = [0.0012577781391462928])
  end

  @trixi_testset "elixir_advection_free_stream_curved.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_free_stream_curved.jl"),
      l2   = [1.830875777528287e-14],
      linf = [7.491784970170556e-13],
      atol = 8e-13, # required to make tests pass on Windows
      )
  end

  @trixi_testset "elixir_advection_nonperiodic_curved.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_nonperiodic_curved.jl"),
      l2   = [6.522004549411137e-5],
      linf = [0.005554857853361295])
  end

  @trixi_testset "elixir_advection_restart_curved.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_restart_curved.jl"),
      l2   = [0.0281388160824776],
      linf = [0.08740635193023694])
  end

  @trixi_testset "elixir_euler_source_terms_curved.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_source_terms_curved.jl"),
      l2   = [0.01032310150257373, 0.009728768969448439, 0.009728768969448494, 0.009728768969448388, 0.015080412597559597],
      linf = [0.034894790428615874, 0.033835365548322116, 0.033835365548322116, 0.03383536554832034, 0.06785765131417065])
  end

  @trixi_testset "elixir_euler_free_stream_curved.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_free_stream_curved.jl"),
      l2   = [2.8815700334367128e-15, 9.361915278236651e-15, 9.95614203619935e-15, 1.6809941842374106e-14, 1.4815037041566735e-14],
      linf = [4.1300296516055823e-14, 2.0444756998472258e-13, 1.0133560657266116e-13, 3.1896707497480747e-13, 6.092903959142859e-13])
  end

  @trixi_testset "elixir_euler_free_stream_curved.jl with FluxRotated(flux_lax_friedrichs)" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_free_stream_curved.jl"),
      surface_flux=FluxRotated(flux_lax_friedrichs),
      l2   = [2.8815700334367128e-15, 9.361915278236651e-15, 9.95614203619935e-15, 1.6809941842374106e-14, 1.4815037041566735e-14],
      linf = [4.1300296516055823e-14, 2.0444756998472258e-13, 1.0133560657266116e-13, 3.1896707497480747e-13, 6.092903959142859e-13])
  end

  @trixi_testset "elixir_euler_nonperiodic_curved.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_nonperiodic_curved.jl"),
      l2   = [0.0018268326813744103, 0.001745601029521995, 0.001745601029521962, 0.0017456010295218891, 0.003239834454817457],
      linf = [0.014660503198892005, 0.01506958815284798, 0.01506958815283821, 0.015069588152864632, 0.02700205515651044])
  end

  @trixi_testset "elixir_euler_ec_curved.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_ec_curved.jl"),
      l2   = [0.011367083018614027, 0.007022020327490176, 0.006759580335962235, 0.006820337637760632, 0.02912659127566544],
      linf = [0.2761764220925329, 0.20286331858055706, 0.18763944865434593, 0.19313636558790004, 0.707563913727584],
      tspan = (0.0, 0.25))
  end

  @trixi_testset "elixir_mhd_ec_curved.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_mhd_ec_curved.jl"),
      l2   = [0.009004021452004856, 0.007143957956917952, 0.006973755636325691, 0.0069016516181158,
              0.033006366599547275, 0.0031996515755530752, 0.0030752688029571135, 0.0030711742765310865,
              4.411700983795491e-5],
      linf = [0.2935375181753366, 0.25906662938283426, 0.2607906900919801, 0.24554085013509105,
              1.1920299902550622, 0.12964835514332484, 0.13226019826981572, 0.14207453734982423,
              0.00362554151305806],
      tspan = (0.0, 0.25))
  end

  @trixi_testset "elixir_mhd_alfven_wave_curved.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_mhd_alfven_wave_curved.jl"),
      l2   = [0.003015476175153681, 0.00145499403283373, 0.0009125744757935803, 0.0017703080480578979,
              0.0013046447673965966, 0.0014564863387645508, 0.0013332311430907598, 0.001647832598455728,
              0.0013647609788548722],
      linf = [0.027510637768610846, 0.02797062834945721, 0.01274249949295704, 0.038940694415543736,
              0.02200825678588325, 0.03167600959583505, 0.021420957993862344, 0.03386589835999665,
              0.01888303191983353])
  end

  @trixi_testset "elixir_mhd_alfven_wave_curved.jl with flux_lax_friedrichs" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_mhd_alfven_wave_curved.jl"),
      l2   = [0.003096205134732006, 0.0014592894240925266, 0.000907526173067474, 0.001821970949880564,
              0.001289067735865978, 0.0014595578872467106, 0.0013269375784248857, 0.0016998139907659089,
              0.0013685168042016343],
      linf = [0.027618475633440998, 0.027093787212065318, 0.012584560784257667, 0.039456640084648914,
              0.020759073985165077, 0.031771018340953416, 0.02059036404759229, 0.03456102393654076,
              0.019663511833857894],
      surface_flux = flux_lax_friedrichs)
  end
end

end # module


text:
export Shamanskii_MA57

function Shamanskii_MA57(nlp_stop,
					     PData :: PDataFact,
					     dₙ :: Vector,
						 gt :: Vector)
	nlp_at_x = nlp_stop.current_state
	x = nlp_at_x.x
	xdemi = x + dₙ
	gtemp = grad(nlp_stop.pb, xdemi)
	ϵ2 =  sqrt(eps(eltype(gt)))
	Γ = max.(abs.(PData.Δ), ϵ2)

	sg = PData.s .* gtemp
	ĝ = PData.L \ (sg[PData.pp])
	g̃ = PData.Q' * ĝ
	d̃ = -(PData.Δ)  .\ g̃
	d = AInv(PData, d̃)

	dₕₒ = dₙ .+ d
	return dₕₒ
end


text:
using SparseArrays
"""
```
krylov_reduction(Γ0, Γ1, Ψ, Π, C; F = x -> x, n_Z = 0)
valuef_reduction(grids, params, n_g_red, Γ0, Γ1, Ψ, Π, C, n_v)
```

Perform Krylov reduction on the state space system defined by Γ0 and Γ1
and/or value function reduction using a spline approximation across some
subset of the state space.

Returns the new matrices, the change of basis matrix and its inverse,
and the new size of the state space.

"""
function krylov_reduction(m::AbstractCTModel{Float64}, Γ0::Matrix{Float64}, Γ1::Matrix{Float64},
                          Ψ::Matrix{Float64}, Π::Matrix{Float64}, C::Array{Float64})

    # Get Krylov subspace function
    if !in(:F, keys(m.settings))
        m <= Setting(:F, identity, "Function applied during Krylov reduction")
    end
    F = get_setting(m, :F)

    # Grab Dimensions
    n_state_vars_unreduce = Int64(get_setting(m, :n_state_vars_unreduce))
    n_state_vars          = Int64(get_setting(m, :n_state_vars))
    n_jump_vars           = Int64(get_setting(m, :n_jump_vars))
    krylov_dim            = Int64(get_setting(m, :krylov_dim))
    n_total               = n_jump_vars + n_state_vars
    n_vars                = n_states(m)
    n_state_vars          = n_state_vars - n_state_vars_unreduce

    # Slice Dynamic Equation Γ1 into different parts. See Why Inequality matters paper.
    B_pv = -Γ1[n_total .+ 1:n_vars, n_total .+ 1:n_vars] \ Γ1[n_total .+ 1:n_vars, 1:n_jump_vars]

    B_pg = -Γ1[n_total .+ 1:n_vars, n_total .+ 1:n_vars] \
        Γ1[n_total .+ 1:n_vars, n_jump_vars .+ 1:n_jump_vars + n_state_vars]
    B_pZ = -Γ1[n_total .+ 1:n_vars, n_total .+ 1:n_vars] \
        Γ1[n_total .+ 1:n_vars, n_jump_vars + n_state_vars .+ 1:n_jump_vars + n_state_vars+n_state_vars_unreduce]
    B_gg = Γ1[n_jump_vars .+ 1:n_jump_vars + n_state_vars, n_jump_vars .+ 1:n_jump_vars + n_state_vars]
    B_gv = Γ1[n_jump_vars .+ 1:n_jump_vars + n_state_vars, 1:n_jump_vars]
    B_gp = Γ1[n_jump_vars .+ 1:n_jump_vars + n_state_vars, n_total .+ 1:n_vars]

    @show size(B_pv), size(B_pg), size(B_pZ), size(B_gg), size(B_gv), size(B_gp)
    @show n_total .+ 1:n_vars
    @show n_jump_vars .+ 1:n_jump_vars + n_state_vars
    @show n_jump_vars + n_state_vars .+ 1:n_jump_vars + n_state_vars+n_state_vars_unreduce
    @show n_jump_vars .+ 1:n_jump_vars + n_state_vars

    # Drop redundant equations in B_pg
    obs        = B_pg
    ~, d0, V_g = svd(obs)
    aux        = d0/d0[1]
    n_Bpg      = Int64(sum(aux .> 10*eps()))
    V_g        = V_g[:, 1:n_Bpg] .* aux[1:n_Bpg]'

    # Compute Krylov subspace
    A(x::Matrix{Float64}) = (F == identity ? 0. : F(B_gv' * x)) + B_gg' * x + B_pg' * (B_gp' * x)
    V_g, ~, ~             = deflated_block_arnoldi(A, V_g, krylov_dim)
    n_state_vars_red      = size(V_g, 2) # number of state variables after reduction

    # Build state space reduction transform
    reduced_basis = spzeros(Float64, n_jump_vars+n_state_vars_red,n_vars)
    reduced_basis[1:n_jump_vars,1:n_jump_vars] = speye(n_jump_vars)
    reduced_basis[n_jump_vars .+ 1:n_jump_vars+n_state_vars_red,n_jump_vars .+ 1:n_jump_vars+n_state_vars] = V_g'
    reduced_basis[n_jump_vars+n_state_vars_red .+ 1:n_jump_vars+n_state_vars_red+n_state_vars_unreduce,n_jump_vars+n_state_vars .+ 1:n_jump_vars+n_state_vars+n_state_vars_unreduce] = eye(n_state_vars_unreduce)

    # Build inverse transform
    inv_reduced_basis = spzeros(n_vars,n_jump_vars+n_state_vars_red)
    inv_reduced_basis[1:n_jump_vars,1:n_jump_vars] = speye(n_jump_vars)
    inv_reduced_basis[n_jump_vars+1:n_jump_vars+n_state_vars,n_jump_vars+1:n_state_vars_red+n_jump_vars] = V_g
    inv_reduced_basis[n_total .+ 1:n_vars,1:n_jump_vars] = B_pv
    inv_reduced_basis[n_total .+ 1:n_vars,n_jump_vars .+ 1:n_jump_vars+n_state_vars_red] = B_pg*V_g
    inv_reduced_basis[n_total .+ 1:n_vars,n_jump_vars+n_state_vars_red .+ 1:n_jump_vars+n_state_vars_red+n_state_vars_unreduce] = B_pZ
    inv_reduced_basis[n_jump_vars+n_state_vars .+ 1:n_total,n_jump_vars+n_state_vars_red .+ 1:n_jump_vars+n_state_vars_red+n_state_vars_unreduce] = speye(n_state_vars_unreduce)

    m <= Setting(:n_state_vars_red, n_state_vars_red + n_state_vars_unreduce,
                 "Number of state variables after reduction")

    # change basis
    Γ0_kry, Γ1_kry, Ψ_kry, Π_kry, C_kry = change_basis(reduced_basis, inv_reduced_basis, Γ0, Γ1, Ψ, Π, C; ignore_Γ0 = true)

    return Γ0_kry, Γ1_kry, Ψ_kry, Π_kry, C_kry, reduced_basis, inv_reduced_basis
end

function krylov_reduction(m::AbstractModel{Float64}, JJ::Matrix{Float64}) #A::Matrix{Float64}, B::Matrix{Float64})

    # Get Krylov subspace function
    if !in(:F, keys(m.settings))
        m <= Setting(:F, identity, "Function applied during Krylov reduction")
    end
    F = get_setting(m, :F)

    # Grab Dimensions
    n_state_vars_unreduce = 0 #Int64(get_setting(m, :n_model_states))
    n_state_vars          = Int64(get_setting(m, :n_backward_looking_states)) #number of backward looking states
    n_jump_vars           = Int64(get_setting(m, :n_jumps)) #number of jumps
    krylov_dim            = Int64(get_setting(m, :krylov_dim))
    n_total               = n_jump_vars + n_state_vars #number of jumps+states (should be number of model dates)
    n_vars                = Int64(get_setting(m, :n_model_states)) #number of model states
    n_state_vars          = n_state_vars - n_state_vars_unreduce #number of backward looking states -> 302

    @show n_state_vars, n_jump_vars, krylov_dim, n_total, n_vars, n_state_vars

    endo = DSGE.augment_model_states(m.endogenous_states_unnormalized, DSGE.n_model_states_unnormalized(m))
    eq = m.equilibrium_conditions
    exo_states_functions = [first(eq[:eq_TFP]), first(eq[:eq_monetary_policy]), first(eq[:eq_markup])]
    exo_states =           [first(endo[:z′_t]), first(endo[:mon′_t]),           first(endo[:mkp′_t])]

    g_inds = [endo[:μ_t], endo[:l_t], endo[:w_t], endo[:R_t], endo[:t_t]]
    #v_inds = [

    # Slice Dynamic Equation JJ into different parts. See Why Inequality matters paper.
    B_pv = -JJ[n_total .+ 1:n_vars, n_total .+ 1:n_vars] #\ JJ[n_total .+ 1:n_vars, 1:n_jump_vars]
    B_pg = -JJ[n_total .+ 1:n_vars, n_total .+ 1:n_vars] #\ JJ[n_total .+ 1:n_vars, n_jump_vars .+ 1:n_jump_vars + n_state_vars]

    B_pZ = -JJ[n_total .+ 1:n_vars, n_total .+ 1:n_vars] #\JJ[n_total .+ 1:n_vars, n_jump_vars + n_state_vars .+ 1:n_jump_vars + n_state_vars+n_state_vars_unreduce]

    B_gg = JJ[n_jump_vars .+ 1:n_jump_vars + n_state_vars, n_jump_vars .+ 1:n_jump_vars + n_state_vars]

    B_gv = JJ[n_jump_vars .+ 1:n_jump_vars + n_state_vars, 1:n_jump_vars]

    B_gp = JJ[n_jump_vars .+ 1:n_jump_vars + n_state_vars, n_total .+ 1:n_vars]

    @show size(B_pv), size(B_pg), size(B_pZ), size(B_gg), size(B_gv), size(B_gp)
    @show n_total .+ 1:n_vars
    @show n_jump_vars .+ 1:n_jump_vars + n_state_vars
    @show n_jump_vars + n_state_vars .+ 1:n_jump_vars + n_state_vars+n_state_vars_unreduce
    @show n_jump_vars .+ 1:n_jump_vars + n_state_vars

    # Slice Dynamic Equation JJ into different parts. See Why Inequality matters paper.
#    @show -JJ[exo_states_functions, exo_states]
#    @show JJ[exo_states_functions, 1:n_jump_vars]
#=    @show "A"
    B_pv = -JJ[exo_states_functions, exo_states] \ JJ[exo_states, 1:n_jump_vars]
    @show "here"
    B_pg = -JJ[exo_states_functions, exo_states] \
        JJ[exo_states_functions, (n_state_vars+1):(n_state_vars+n_jump_vars)]
    B_pZ = -JJ[exo_states_functions, exo_states] \
        JJ[exo_states, 2:1] #2:1 returns 0. this is what was being returned anyways
    B_gg = JJ[(n_state_vars+1):(n_state_vars+n_jump_vars), (n_state_vars+1):(n_state_vars+n_jump_vars)]
    B_gv = JJ[(n_state_vars+1):(n_state_vars+n_jump_vars), 1:n_state_vars]
    B_gp = JJ[(n_state_vars+1):(n_state_vars+n_jump_vars), exo_states]

    @show size(B_pv), size(B_pg), size(B_pZ), size(B_gg), size(B_gv), size(B_gp)
=#
    # Drop redundant equations in B_pg
    obs        = B_pg
    F_svd      = svd(obs)
    d0, V_g    = F_svd.S, F_svd.V
    @show typeof(d0), d0, typeof(V_g), size(V_g), B_pg
    aux        = d0/d0[1]
    n_Bpg      = Int64(sum(aux .> 10*eps()))
    V_g        = V_g[:, 1:n_Bpg] .* aux[1:n_Bpg]'

    # Compute Krylov subspace
    A(x::Matrix{Float64}) = (F == identity ? 0. : F(B_gv' * x)) + B_gg' * x + B_pg' * (B_gp' * x)
    V_g, ~, ~             = deflated_block_arnoldi(A, V_g, krylov_dim)
    n_state_vars_red      = size(V_g, 2) # number of state variables after reduction

    # Build state space reduction transform
    reduced_basis = spzeros(Float64, n_jump_vars+n_state_vars_red,n_vars)
    reduced_basis[1:n_jump_vars,1:n_jump_vars] = speye(n_jump_vars)
    reduced_basis[n_jump_vars .+ 1:n_jump_vars+n_state_vars_red,n_jump_vars .+ 1:n_jump_vars+n_state_vars] = V_g'
    reduced_basis[n_jump_vars+n_state_vars_red .+ 1:n_jump_vars+n_state_vars_red+n_state_vars_unreduce,n_jump_vars+n_state_vars .+ 1:n_jump_vars+n_state_vars+n_state_vars_unreduce] = eye(n_state_vars_unreduce)

    # Build inverse transform
    inv_reduced_basis = spzeros(n_vars,n_jump_vars+n_state_vars_red)
    inv_reduced_basis[1:n_jump_vars,1:n_jump_vars] = speye(n_jump_vars)
    inv_reduced_basis[n_jump_vars+1:n_jump_vars+n_state_vars,n_jump_vars+1:n_state_vars_red+n_jump_vars] = V_g
    inv_reduced_basis[exo_states,1:n_jump_vars] = B_pv
    inv_reduced_basis[exo_states,n_jump_vars .+ 1:n_jump_vars+n_state_vars_red] = B_pg*V_g
    inv_reduced_basis[exo_states,n_jump_vars+n_state_vars_red .+ 1:n_jump_vars+n_state_vars_red+n_state_vars_unreduce] = B_pZ
    inv_reduced_basis[n_jump_vars+n_state_vars .+ 1:n_total,n_jump_vars+n_state_vars_red .+ 1:n_jump_vars+n_state_vars_red+n_state_vars_unreduce] = speye(n_state_vars_unreduce)

    m <= Setting(:n_state_vars_red, n_state_vars_red + n_state_vars_unreduce,
                 "Number of state variables after reduction")

    # change basis
    JJ_kry = change_basis(reduced_basis, inv_reduced_basis, JJ; ignore_Γ0 = true)

    return JJ_kry, reduced_basis, inv_reduced_basis
end

function valuef_reduction(m::AbstractModel,
                          Γ0::Array{Float64,2}, Γ1::Array{Float64,2}, Ψ::Array{Float64,2},
                          Π::Array{Float64,2}, C::Array{Float64})

    n_jump_vars      = Int64(get_setting(m, :n_jump_vars))
    n_state_vars_red = Int64(get_setting(m, :n_state_vars_red))
    spline_grid      = get_setting(m, :spline_grid)
    knots_dict       = Dict{Int64,Vector{Float64}}()

    merge!(knots_dict, get_setting(m, :knots_dict))
    n_prior = Int64(get_setting(m, :n_prior))
    n_post  = Int64(get_setting(m, :n_post))

    # Function calls to create basis reduction
    knots_dim       = Int64(sum(map(i -> length(knots_dict[i]) + 1, 1:length(keys(knots_dict)))))
    spline_grid_dim = Int64(prod(size(spline_grid)))
    from_spline     = spzeros(Float64, spline_grid_dim, knots_dim)
    to_spline       = spzeros(Float64, knots_dim, spline_grid_dim)

    from_spline[:,:], to_spline[:,:] = spline_basis(spline_grid, knots_dict, 2) # create spline basis

    # Extend to other dimensions along which we did not use a spline approximation.
    # In this case, we computed a spline basis for wealth but not for income
    from_spline, to_spline = extend_to_nd(from_spline, to_spline, n_prior, n_post)

    # extra jump variables besides value function
    extra_jv = Int64(n_jump_vars - length(m.endogenous_states[:value_function]))

    if extra_jv > 0
        # Add additional dimensions for any jump variables we did not reduce b/c
        # spline reduction only for value function
        dim1_from, dim2_from = size(from_spline)
        from_spline = sparse([from_spline spzeros(dim1_from, extra_jv)]) # add extra columns
        from_spline = sparse([from_spline; spzeros(extra_jv, dim2_from + extra_jv)])  # add extra zeros
        to_spline = sparse([to_spline spzeros(dim2_from, extra_jv)])     # same as above
        to_spline = sparse([to_spline; spzeros(1, dim1_from + extra_jv)])
        for i = 0:(extra_jv - 1)
            from_spline[end - i, end - i] = 1 # add 1 along diagonal of added dimensions
            to_spline[end - i, end - i] = 1
        end
    end
    n_splined = Int64(size(from_spline, 2))

    # Create projection matrix that projects value function onto spline basis
    from_spline, to_spline = projection_for_subset(from_spline, to_spline, 0, n_state_vars_red)

    # Reduce the decision vector
    Γ0_spl, Γ1_spl, Ψ_spl, ~, C_spl = change_basis(to_spline, from_spline, Γ0, Γ1, Ψ, Π, C; ignore_Γ0 = true)
    Π_spl = zeros(Float64, Int64(n_states(m)), Int64(n_shocks_expectational(m)))

    Π_spl = to_spline * Π * from_spline[1:n_jump_vars, 1:n_splined]
    m <= Setting(:n_splined, n_splined, "Dimension of jump variables after spline basis reduction")

    return Γ0_spl, Γ1_spl, Ψ_spl, Π_spl, C_spl, to_spline, from_spline
 end

# Look up the algorithm for this method.
function deflated_block_arnoldi(A::Function, B::Matrix{Float64}, m::Int64)
    F = qr(B)
    Q = Matrix{Float64}(F.Q * Matrix{Float64}(I, size(B)))
    basis = Matrix{Float64}(undef, size(Q,1), 0)
    realsmall = sqrt(eps())

    if m == 1
        basis = Q
    else
        for i in 1:m-1
            # Manual Gram-Schmidt
            basis = hcat(basis, Q)
            aux = A(Q)::Matrix{Float64}
            for j in 1:size(basis,2)
                aux -= basis[:,j] .* basis[:,j]' * aux
            end

            # Check for potential deflation
            Q = Matrix{Float64}(undef, size(Q,1), 0)
            for j in 1:size(aux,2)
                weight = sqrt(sum(aux[:,j].^2))
                if weight > realsmall
                    Q = hcat(Q, aux[:,j] / weight)
                    for k in j+1:size(aux,2)
                        aux[:,k] -= Q[:,end] * (Q[:,end]' * aux[:,k])
                    end
                end
            end

            # More Gram-Schmidt
            for j in 1:size(basis,2)
                Q -= basis[:,j] .* basis[:,j]' * Q
            end
            Q = Q ./ sqrt.(sum(Q.^2, dims=1))
        end
    end
    err = qr(A(Q) - basis * basis' * A(Q)).R::Matrix{Float64}
    return basis, Q, err
end

# Just think through how matrices can be used to do a change of basis
function change_basis(basis, inv_basis, Γ0::Matrix{Float64}, Γ1::Matrix{Float64},
                      Ψ::Matrix{Float64}, Π::Matrix{Float64}, C::Array{Float64};
                      ignore_Γ0::Bool = false)
    g1 = basis * Γ1 * inv_basis
    if ignore_Γ0
        g0 = eye(g1)
    else
        g0 = basis * Γ0 * inv_basis
    end
    c    = basis * C
    psi  = basis * Ψ
    Pi   = basis * Π

    return g0, g1, psi, Pi, c
end

# Just think through how matrices can be used to do a change of basis
function change_basis(basis, inv_basis, JJ::Matrix{Float64};
                      ignore_Γ0::Bool = false)
    g1 = basis * JJ * inv_basis

    return g1
end

# Creates a quadratic polynomial based spline basis reduction
# QUASI - DEPRECATED: translation from SeHyoun's code but not used b/c not general, better
# packages available already in Julia
# Parameters:
#    x = fine grid points
#    knots = coarser grid points onto which we reduce
# Outputs:
#    from_knots = change of basis from spline basis to finer grid points
#    to_knots = change of basis from finer grid points to spline basis
function oneDquad_spline(x::Vector{Float64}, knots::Vector{Float64})
    n_a = length(x)
    n_knots = length(knots)

    # Preallocate auxiliary matrices
    first_interp_mat = zeros(n_a, n_knots + 1)
    aux_mat = zeros(n_a, n_knots)

    # Find to where each points corresponds on knots
    for i = 1:n_a
        loc = sum(knots .<= x[i])
        if loc == n_knots
            loc = n_knots - 1
        end
        first_interp_mat[i, loc] = 1 - (x[i] - knots[loc])^2/(knots[loc + 1] - knots[loc])^2
        first_interp_mat[i, loc + 1] = (x[i] - knots[loc])^2/(knots[loc + 1] - knots[loc])^2
        aux_mat[i, loc] = (x[i] - knots[loc]) - (x[i] - knots[loc])^2/(knots[loc + 1] - knots[loc])
    end

    aux_mat2 = spdiagm(ones(n_knots), 0, n_knots, n_knots) + spdiagm(ones(n_knots-1), 1, n_knots, n_knots)
    aux_mat2[end, end] = 0
    aux_mat2[n_knots, 1] = 1
    aux_mat3 = spdiagm([-2 ./ diff(knots); 0], 0, n_knots, n_knots + 1) +
        spdiagm([2 ./ diff(knots); 1], 1, n_knots, n_knots + 1)

    # Return values
    from_knots = first_interp_mat + aux_mat*(Matrix(aux_mat2)\Matrix(aux_mat3))
    to_knots = from_knots'*from_knots\(from_knots'*eye(n_a))
    return sparse(from_knots), sparse(to_knots)
end

# Creates the spline basis matrix and its inverse, assuming the same degree of splines
# in each direction of approximation
function spline_basis(x::Array{Float64}, knots::Dict{Int64,Vector{Float64}}, degree::Int64)

    # Create spline approximations along each dimension
    approx_dim = length(size(x))
    basis = Basis(map(i -> Basis(SplineParams(knots[i], 0, degree)), 1:approx_dim)...)

    # Create spline basis matrix and its inverse. The spline basis matrix transforms
    # the coefficient vector of knot points to the coefficient vector of the provided
    # state space array x, i.e. S*knots = x, where S is the spline basis matrix.
    from_knots = BasisMatrix(basis, Expanded(), x, 0).vals[1]
    to_knots = Matrix{Float64}(from_knots'*from_knots) \ Matrix{Float64}(from_knots')
    return from_knots, sparse(to_knots)
end

# Creates the spline basis matrix and its inverse, assuming a vector specifying the desired degrees.
function spline_basis(x::Array{Float64}, knots::Dict{Int64,Vector{Float64}}, degree::Vector{Int64})
    # Create spline approximations along each dimension
    approx_dim = length(size(x))
    basis = Basis(SplineParams(knots[1], 0, degree))
    for i = 2:approx_dim
        basis = Basis(basis, SplineParams(knots[i], 0, degree)) # UNSATISFIED WITH THIS; FASTER WAY?
    end

    # Create spline basis matrix and its inverse. The spline basis matrix transforms
    # the coefficient vector of knot points to the coefficient vector of the provided
    # state space array x, i.e. S*knots = x, where S is the spline basis matrix.
    from_knots = BasisMatrix(basis, Expanded(), x, 0).vals[1]
    to_knots = Matrix(from_knots'*from_knots) \ Matrix(from_knots')
    return from_knots, sparse(to_knots)
end

# Parameters:
#    from_small = projection to approximation
#    to_small = approximation to projection
#    n_prior = number of grid points stacked prior to reduction dimension
#    n_post = number of grid points stacked after reduction dimension
#
# Outputs:
#    from_approx = basis change from small points to large
#    to_approx = basis change from large number of points to small
function extend_to_nd(from_small::SparseMatrixCSC{Float64,Int64}, to_small::SparseMatrixCSC{Float64,Int64},
                      n_prior::Int64, n_post::Int64)
    from_approx = kron(kron(SparseMatrixCSC{Float64}(I, n_post, n_post), from_small),
                       SparseMatrixCSC{Float64}(I, n_prior, n_prior))
    to_approx = kron(kron(SparseMatrixCSC{Float64}(I, n_post, n_post), to_small),
                     SparseMatrixCSC{Float64}(I, n_prior, n_prior))
    return from_approx, to_approx
end

# Create transformations when we only transform subset of the points.
#
# Arguments
#
# from_small: (nxk) matrix transformation from approximated k points into full grids
# to_small: (kxn) matrix transformation into approximation with k points
# n_pre: number of points preceding the transformation
# n_post: number of points following the transformation
#
# Outputs:
# from_approx: transformation from smaller approximation to larger grid points
# to_approx: reverse of from_approx
function projection_for_subset(from_small::SparseMatrixCSC{Float64,Int64},
                               to_small::SparseMatrixCSC{Float64,Int64},
                               n_pre::Int64, n_post::Int64)

    n_full, n_red = size(from_small)
    I, J, V = SparseArrays.spdiagm_internal(0 => ones(n_red+n_pre)); from_approx = sparse(I, J, V, n_full + n_pre, n_pre + n_red)
    I, J, V = SparseArrays.spdiagm_internal(0 => ones(n_red+n_pre)); to_approx = sparse(I, J, V, n_pre + n_red, n_full + n_pre)
    from_approx[1+n_pre:n_pre+n_full, n_pre+1:n_pre+n_red] = from_small
    to_approx[n_pre+1:n_pre+n_red, n_pre+1:n_pre+n_full] = to_small

    # Expand matrices and add needed values
    dim1_from_approx, dim2_from_approx = size(from_approx)
    from_approx = [from_approx spzeros(dim1_from_approx, n_post)]
    from_approx = [from_approx; spzeros(n_post, dim2_from_approx + n_post)]
    to_approx = [to_approx spzeros(dim2_from_approx, n_post)]
    to_approx = [to_approx; spzeros(n_post, dim1_from_approx + n_post)]
    from_approx[dim1_from_approx+1:end, dim2_from_approx+1:end] = speye(n_post)
    to_approx[dim2_from_approx+1:end, dim1_from_approx+1:end] = speye(n_post)

    return from_approx, to_approx
end

# Cleans Γ0 matrix to be identity by solving out static conditions
function solve_static_conditions(Γ0::Matrix{Float64},
                                 Γ1::Matrix{Float64},
                                 Ψ::Matrix{Float64},
                                 Π::Matrix{Float64},
                                 C::Array{Float64})

    redundant = maximum(abs.([Γ0 Ψ]), 2) .== 0  # Find rows of both Γ0 & Ψ with only zeros
    redundant_list = Array{Int64,1}(0)
    for row in 1:length(redundant)
        if redundant[row]
            push!(redundant_list, row)
        end
    end
    inv_state_red = nullspace(Γ1[redundant_list,:]) # Compute orthonormal basis of null space of Γ1
    state_red = inv_state_red'

    g0  = state_red * Γ0 * inv_state_red # adjust Γ0 to new basis
    g1  = state_red * Γ1 * inv_state_red # Zero out redundant rows/columns in Γ1
    g1  = g0 \ g1                        # Force Γ1 and other matrices
    Psi = g0 \ (state_red * Ψ)           # to be the ones corresponding to when
    Pi  = g0 \ (state_red * Π)           # Γ0 is the identity
    c   = g0 \ (state_red * C)

    return g0, g1, c, Pi, Psi, state_red, inv_state_red
end

# When we have sparse matrices
function solve_static_conditions(Γ0::SparseMatrixCSC{Float64, Int64},
                                 Γ1::SparseMatrixCSC{Float64, Int64},
                                 Ψ::SparseMatrixCSC{Float64, Int64},
                                 Π::SparseMatrixCSC{Float64, Int64},
                                 C::SparseMatrixCSC{Float64, Int64})

    redundant = maximum(abs.([Γ0 Ψ]), 2) .== 0 # R: I removed the comma here, given the spirit of what
                                               # it appears we're trying to do.
    redundant_list = Array{Int64,1}(0)
    for row in 1:length(redundant)
        if redundant[row]
            push!(redundant_list, row)
        end
    end
    inv_state_red = nullspace(full(Γ1[redundant_list,:]))
    state_red = inv_state_red'

    g0 = state_red * Matrix(Γ0) * inv_state_red
    g1 = state_red * Matrix(Γ1) * inv_state_red
    g1 = g0 \ g1
    Psi = g0 \ (state_red * Matrix(Ψ))
    Pi = g0 \ (state_red * Matrix(Π))
    c = g0 \ (state_red * Vector(C))

    return g0, g1, c, Pi, Psi, state_red, inv_state_red
end

# When we have sparse matrices and a vector for C
function solve_static_conditions(Γ0::SparseMatrixCSC{Float64, Int64},
                                 Γ1::SparseMatrixCSC{Float64, Int64},
                                 Ψ::SparseMatrixCSC{Float64, Int64},
                                 Π::SparseMatrixCSC{Float64, Int64},
                                 C::SparseVector{Float64, Int64})

    redundant = maximum(abs.([Γ0, Ψ]), 2) .== 0
    inv_state_red = null(full(Γ1(redundant,:)))
    state_red = inv_state_red'

    g0 = state_red * Matrix(Γ0) * inv_state_red
    g1 = state_red * Matrix(Γ1) * inv_state_red
    g1 = g0 \ g1
    Psi = g0 \ (state_red * Matrix(Ψ))
    Pi = g0 \ (state_red * Matrix(Π))
    c = g0 \ (state_red * Vector(C))

    return g0, g1, c, Pi, Psi, state_red, inv_state_red
end


text:
using DataDeps

register(DataDep(
    "breast-cancer-wisconsin",
    """
    Dataset: breast-cancer-wisconsin
    Author: Dr. WIlliam H. Wolberg
    Donors: Olvi Mangasarian
        Received by David W. Aha
    Website: https://archive.ics.uci.edu/ml/datasets/breast+cancer+wisconsin+(original)

    Observations: 699 (16 missing)
    Features:     9
    Classes:      2
    """,
    "https://archive.ics.uci.edu/ml/machine-learning-databases/breast-cancer-wisconsin/breast-cancer-wisconsin.data",
    "402c585309c399237740f635ef9919dc512cca12cbeb20de5e563a4593f22b64",
    post_fetch_method = (path -> begin
        MVDA.process_dataset(path,
            missingstrings=["?"],
            header=false,
            target_index=11,
            feature_indices=2:10,
            ext=".csv",
        )
    end),
))


text:
function densesystem(mechanism::Mechanism{T,Nn,Nb}) where {T,Nn,Nb}
    bodies = mechanism.bodies
    eqcs = mechanism.eqconstraints
    graph = mechanism.graph
    ldu = mechanism.ldu

    n = 6 * Nb
    for eqc in eqcs
        n += length(eqc)
    end

    A = zeros(T,n,n)
    x = zeros(T,n)
    b = zeros(T,n)

    rangeDict = Dict{Int64,UnitRange}()
    ind1 = 1
    ind2 = 0

    for id in graph.dfslist
        component = getcomponent(mechanism, id)
        ind2 += length(component)
        range = ind1:ind2
        rangeDict[id] = range

        # A
        diagonal = getentry(ldu,id)
        A[range,range] = diagonal.D

        for childid in successors(graph, id)
            offdiagonal = getentry(ldu, (id, childid))
            nc1 = first(rangeDict[childid])
            nc2 = last(rangeDict[childid])

            A[range,nc1:nc2] = offdiagonal.L
            A[nc1:nc2,range] = offdiagonal.U
        end

        # x
        x[range] = diagonal.Δs

        # b
        if component isa Body
            b[range] = dynamics(mechanism, component)
        else
            b[range] = g(mechanism, component)
        end

        ind1 = ind2+1
    end

    return A, x, b
end

# TODO only works for 1DOF active constraints (eqcids)
# Maximal Coordinates
function linearsystem(mechanism::Mechanism{T,Nn,Nb}, xd, vd, qd, ωd, Fτd, bodyids, eqcids) where {T,Nn,Nb}
    statesold = [State{T}() for i=1:Nb]

    # store old state and set new initial state
    for (i,id) in enumerate(bodyids)
        stateold = settempvars!(getbody(mechanism, id), xd[i], vd[i], zeros(T,3), qd[i], ωd[i], zeros(T,3), zeros(T,6))
        statesold[i] = stateold
    end
    for (i,id) in enumerate(eqcids)
        setForce!(mechanism, geteqconstraint(mechanism, id), Fτd[i])
    end

    A, Bu, Bλ, G = lineardynamics(mechanism, eqcids)

    # restore old state
    for (i,id) in enumerate(bodyids)
        body = getbody(mechanism, id)
        body.state = statesold[i]
    end

    return A, Bu, Bλ, G
end
# Minimal Coordinates
function linearsystem(mechanism::Mechanism{T,Nn,Nb}, xθd, vωd, Fτd, controlledids, controlids) where {T,Nn,Nb}
    statesold = [State{T}() for i=1:Nb]
    xd = [szeros(T,3) for i=1:Nb]
    vd = [szeros(T,3) for i=1:Nb]
    qd = [one(UnitQuaternion{T}) for i=1:Nb]
    ωd = [szeros(T,3) for i=1:Nb]

    # store old state and set new initial state
    for (id,body) in pairs(mechanism.bodies)
        statesold[id] = deepcopy(body.state)
    end
    for (i,id) in enumerate(controlledids)
        setPosition!(mechanism, geteqconstraint(mechanism, id), [xθd[i]])
    end
    for (id,body) in pairs(mechanism.bodies)
        state = body.state
        xd[id] = state.xc
        vd[id] = state.vc
        qd[id] = state.qc
        ωd[id] = state.ωc
    end
    for (i,id) in enumerate(controlids)
        setForce!(mechanism, geteqconstraint(mechanism, id), [Fτd[i]])
    end

    A, Bu, Bλ, G = lineardynamics(mechanism, controlids)

    # restore old state
    for (id,body) in pairs(mechanism.bodies)
        body.state = statesold[id]
    end

    return A, Bu, Bλ, G, xd, vd, qd, ωd
end

function lineardynamics(mechanism::Mechanism{T,Nn,Nb}, eqcids) where {T,Nn,Nb}
    Δt = mechanism.Δt
    bodies = mechanism.bodies
    eqcs = mechanism.eqconstraints
    graph = mechanism.graph

    nc = 0
    for eqc in eqcs
        isinactive(eqc) && continue
        nc += length(eqc)
    end

    nu = 0
    for id in eqcids
        eqc = geteqconstraint(mechanism, id)
        nu += 6-length(eqc)
    end

    # calculate next state
    discretizestate!(mechanism)
    foreach(setsolution!, mechanism.bodies)
    foreach(applyFτ!, eqcs, mechanism, false)
    newton!(mechanism)

    # get state linearization
    Fz = zeros(T,Nb*13,Nb*12)
    Fu = zeros(T,Nb*13,Nb*6)
    Fλ = zeros(T,Nb*13,nc)
    Ffz = zeros(T,Nb*13,Nb*13)
    invFfzquat = zeros(T,Nb*12,Nb*13)

    Bcontrol = zeros(T,Nb*6,nu)

    for (id,body) in pairs(bodies)
        col6 = offsetrange(id,6)
        col12 = offsetrange(id,12)
        col13 = offsetrange(id,13)

        Fzi = ∂F∂z(body, Δt)
        Fui = ∂F∂u(body, Δt)
        Ffzi, invFfzquati = ∂F∂fz(body, Δt)

        Fz[col13,col12] = Fzi
        Fu[col13,col6] = Fui
        Ffz[col13,col13] = Ffzi
        invFfzquat[col12,col13] = invFfzquati
    end

    Ffz += linearconstraintmapping(mechanism)
    Fz += linearforcemapping(mechanism)

    n1 = 1
    n2 = 0

    for id in eqcids
        eqc = geteqconstraint(mechanism, id)
        n2 += 6-length(eqc)

        parentid = eqc.parentid
        if parentid !== nothing
            col6 = offsetrange(parentid,6)
            Bcontrol[col6,n1:n2] = ∂Fτ∂ua(mechanism, eqc, parentid)
        end
        for childid in eqc.childids
            col6 = offsetrange(childid,6)
            Bcontrol[col6,n1:n2] = ∂Fτ∂ub(mechanism, eqc, childid)
        end

        n1 = n2+1
    end

    G, Fλ = linearconstraints(mechanism)

    invFfz = invFfzquat * inv(Ffz)
    A = -invFfz * Fz
    Bu = -invFfz * Fu * Bcontrol
    Bλ = -invFfz * Fλ

    return A, Bu, Bλ, G
end

function linearconstraints(mechanism::Mechanism{T,Nn,Nb}) where {T,Nn,Nb}
    Δt = mechanism.Δt
    eqcs = mechanism.eqconstraints

    nc = 0
    for eqc in eqcs
        isinactive(eqc) && continue
        nc += length(eqc)
    end

    Gl = zeros(T,nc,Nb*12)
    Gr = zeros(T,nc,Nb*13)

    oneindc = 0
    for eqc in eqcs
        isinactive(eqc) && continue
        ind1 = 1
        ind2 = 0

        parentid = eqc.parentid
        if parentid !== nothing
            pbody = getbody(mechanism,parentid)
            pstate = pbody.state

            for (i,childid) in enumerate(eqc.childids)
                cbody = getbody(mechanism,childid)
                cstate = cbody.state

                ind2 += length(eqc.constraints[i])
                range = oneindc+ind1:oneindc+ind2

                pcol3a12 = offsetrange(parentid,3,12,1)
                pcol3b12 = offsetrange(parentid,3,12,2)
                pcol3c12 = offsetrange(parentid,3,12,3)
                pcol3d12 = offsetrange(parentid,3,12,4)

                pcol3b = offsetrange(parentid,3,13,2)
                pcol3d = offsetrange(parentid,3,13,4)
                pcol3d = first(pcol3d)+1:last(pcol3d)+1

                ccol3a12 = offsetrange(childid,3,12,1)
                ccol3b12 = offsetrange(childid,3,12,2)
                ccol3c12 = offsetrange(childid,3,12,3)
                ccol3d12 = offsetrange(childid,3,12,4)

                ccol3b = offsetrange(childid,3,13,2)
                ccol3d = offsetrange(childid,3,13,4)
                ccol3d = first(ccol3d)+1:last(ccol3d)+1

                pXl, pQl = ∂g∂posa(eqc.constraints[i], posargsnext(pstate, Δt)..., posargsnext(cstate, Δt)...) # x3
                pGr = ∂g∂ʳposa(eqc.constraints[i], posargssol(pstate)..., posargssol(cstate)...) # x2
                pXr, pQr = pGr[:,1:3], pGr[:,4:6]
                cXl, cQl =  ∂g∂posb(eqc.constraints[i], posargsnext(pstate, Δt)..., posargsnext(cstate, Δt)...) # x3
                cGr =  ∂g∂ʳposb(eqc.constraints[i], posargssol(pstate)..., posargssol(cstate)...) # x2
                cXr, cQr = cGr[:,1:3], cGr[:,4:6]

                mat = constraintmat(eqc.constraints[i])
                pGlx = mat * pXl
                pGlq = mat * pQl
                pGrx = mat * pXr
                pGrq = mat * pQr
                cGlx = mat * cXl
                cGlq = mat * cQl
                cGrx = mat * cXr
                cGrq = mat * cQr

                Gl[range,pcol3a12] = pGlx
                Gl[range,pcol3b12] = pGlx*Δt
                Gl[range,pcol3c12] = pGlq*Rmat(ωbar(pstate.ωsol[2],Δt))*LVᵀmat(pstate.qsol[2])
                Gl[range,pcol3d12] = pGlq*Lmat(pstate.qsol[2])*derivωbar(pstate.ωsol[2],Δt)
                Gr[range,pcol3b] = pGrx
                Gr[range,pcol3d] = pGrq

                Gl[range,ccol3a12] = cGlx
                Gl[range,ccol3b12] = cGlx*Δt
                Gl[range,ccol3c12] = cGlq*Rmat(ωbar(cstate.ωsol[2],Δt))*LVᵀmat(cstate.qsol[2])
                Gl[range,ccol3d12] = cGlq*Lmat(cstate.qsol[2])*derivωbar(cstate.ωsol[2],Δt)
                Gr[range,ccol3b] = cGrx
                Gr[range,ccol3d] = cGrq

                ind1 = ind2+1
            end
        else
            for (i,childid) in enumerate(eqc.childids)
                cbody = getbody(mechanism,childid)
                cstate = cbody.state

                ind2 += length(eqc.constraints[i])
                range = oneindc+ind1:oneindc+ind2

                ccol3a12 = offsetrange(childid,3,12,1)
                ccol3b12 = offsetrange(childid,3,12,2)
                ccol3c12 = offsetrange(childid,3,12,3)
                ccol3d12 = offsetrange(childid,3,12,4)

                ccol3b = offsetrange(childid,3,13,2)
                ccol3d = offsetrange(childid,3,13,4)
                ccol3d = first(ccol3d)+1:last(ccol3d)+1

                cXl, cQl =  ∂g∂posb(eqc.constraints[i], posargsnext(cstate, Δt)...) # x3
                cGr =  ∂g∂ʳposb(eqc.constraints[i], posargssol(cstate)...) # x2
                cXr, cQr = cGr[:,1:3], cGr[:,4:6]

                mat = constraintmat(eqc.constraints[i])
                cGlx = mat * cXl
                cGlq = mat * cQl
                cGrx = mat * cXr
                cGrq = mat * cQr

                Gl[range,ccol3a12] = cGlx
                Gl[range,ccol3b12] = cGlx*Δt
                Gl[range,ccol3c12] = cGlq*Rmat(ωbar(cstate.ωsol[2],Δt))*LVᵀmat(cstate.qsol[2])
                Gl[range,ccol3d12] = cGlq*Lmat(cstate.qsol[2])*derivωbar(cstate.ωsol[2],Δt)
                Gr[range,ccol3b] = cGrx
                Gr[range,ccol3d] = cGrq

                ind1 = ind2+1
            end
        end

        oneindc += length(eqc)
    end

    return Gl, -Gr'
end

function linearconstraintmapping(mechanism::Mechanism{T,Nn,Nb}) where {T,Nn,Nb}
    FfzG = zeros(T,Nb*13,Nb*13)

    K = zeros(T,9,9)
    K[1,1] = K[2,4] = K[3,7] = K[4,2] = K[5,5] = K[6,8] = K[7,3] = K[8,6] = K[9,9] = 1
    E = SMatrix{3,3,T,9}(I)

    for eqc in mechanism.eqconstraints
        parentid = eqc.parentid

        if parentid !== nothing
            body1 = getbody(mechanism, parentid)
            state1 = body1.state
            pcol13 = offsetrange(parentid,13)

            for (i, childid) in enumerate(eqc.childids)
                body2 = getbody(mechanism, childid)
                state2 = body2.state
                constraint = eqc.constraints[i]
                ccol13 = offsetrange(childid,13)

                n1 = 1
                n2 = 0
                for j=1:i-1
                    n1 += length(eqc.constraints[j])
                    n2 += length(eqc.constraints[j])
                end
                n2 += length(eqc.constraints[i])
                λ = eqc.λsol[2][n1:n2]

                Aaa = zeros(T,13,13)
                Aab = zeros(T,13,13)
                Aba = zeros(T,13,13)
                Abb = zeros(T,13,13)

                kronproduct = -kron(λ'*constraintmat(constraint),E)*K

                XX, XQ, QX, QQ = ∂2g∂posaa(constraint, state1.xsol[2], state1.qsol[2], state2.xsol[2], state2.qsol[2])
                Aaa[4:6,1:3] = kronproduct*XX
                Aaa[4:6,7:10] = kronproduct*XQ
                Aaa[11:13,1:3] = kronproduct*QX
                Aaa[11:13,7:10] = kronproduct*QQ

                XX, XQ, QX, QQ = ∂2g∂posab(constraint, state1.xsol[2], state1.qsol[2], state2.xsol[2], state2.qsol[2])
                Aab[4:6,1:3] = kronproduct*XX
                Aab[4:6,7:10] = kronproduct*XQ
                Aab[11:13,1:3] = kronproduct*QX
                Aab[11:13,7:10] = kronproduct*QQ

                XX, XQ, QX, QQ = ∂2g∂posba(constraint, state1.xsol[2], state1.qsol[2], state2.xsol[2], state2.qsol[2])
                Aba[4:6,1:3] = kronproduct*XX
                Aba[4:6,7:10] = kronproduct*XQ
                Aba[11:13,1:3] = kronproduct*QX
                Aba[11:13,7:10] = kronproduct*QQ

                XX, XQ, QX, QQ = ∂2g∂posbb(constraint, state1.xsol[2], state1.qsol[2], state2.xsol[2], state2.qsol[2])
                Abb[4:6,1:3] = kronproduct*XX
                Abb[4:6,7:10] = kronproduct*XQ
                Abb[11:13,1:3] = kronproduct*QX
                Abb[11:13,7:10] = kronproduct*QQ

                FfzG[pcol13,pcol13] += Aaa
                FfzG[pcol13,ccol13] += Aab
                FfzG[ccol13,pcol13] += Aba
                FfzG[ccol13,ccol13] += Abb
            end
        else
            for (i, childid) in enumerate(eqc.childids)
                body2 = getbody(mechanism, childid)
                state2 = body2.state
                constraint = eqc.constraints[i]
                ccol13 = offsetrange(childid,13)

                n1 = 1
                n2 = 0
                for i=1:i-1
                    n1 += length(eqc.constraints[i])
                    n2 += length(eqc.constraints[i])
                end
                n2 += length(eqc.constraints[i])
                λ = eqc.λsol[2][n1:n2]

                Abb = zeros(T,13,13)

                kronproduct = -kron(λ'*constraintmat(constraint),E)*K

                XX, XQ, QX, QQ = ∂2g∂posbb(constraint, state2.xsol[2], state2.qsol[2])
                Abb[4:6,1:3] = kronproduct*XX
                Abb[4:6,7:10] = kronproduct*XQ
                Abb[11:13,1:3] = kronproduct*QX
                Abb[11:13,7:10] = kronproduct*QQ

                FfzG[ccol13,ccol13] += Abb
            end
        end
    end

    return FfzG
end

function linearforcemapping(mechanism::Mechanism{T,Nn,Nb}) where {T,Nn,Nb}
    Fzu = zeros(T,Nb*13,Nb*12)

    for eqc in mechanism.eqconstraints
        parentid = eqc.parentid
        for (i,childid) in enumerate(eqc.childids)
            if parentid !== nothing
                FaXa, FaQa, τaXa, τaQa, FbXa, FbQa, τbXa, τbQa = ∂Fτ∂posa(eqc.constraints[i], getbody(mechanism, parentid).state, getbody(mechanism, childid).state)
                FaXb, FaQb, τaXb, τaQb, FbXb, FbQb, τbXb, τbQb = ∂Fτ∂posb(eqc.constraints[i], getbody(mechanism, parentid).state, getbody(mechanism, childid).state)

                cola6 = offsetrange(parentid,6)
                colb6 = offsetrange(childid,6)
                rowav = offsetrange(parentid,3,13,2)
                rowaω = offsetrange(parentid,3,13,4).+1
                rowbv = offsetrange(childid,3,13,2)
                rowbω = offsetrange(childid,3,13,4).+1

                Fzu[rowav,cola6] = [FaXa FaQa]
                Fzu[rowaω,cola6] = [τaXa τaQa]
                Fzu[rowbv,cola6] = [FbXa FbQa]
                Fzu[rowbω,cola6] = [τbXa τbQa]

                Fzu[rowav,colb6] = [FaXb FaQb]
                Fzu[rowaω,colb6] = [τaXb τaQb]
                Fzu[rowbv,colb6] = [FbXb FbQb]
                Fzu[rowbω,colb6] = [τbXb τbQb]
            else
                FaXb, FaQb, τaXb, τaQb, FbXb, FbQb, τbXb, τbQb = ∂Fτ∂posb(eqc.constraints[i], getbody(mechanism, childid).state)

                colb6 = offsetrange(childid,6)
                rowbv = offsetrange(childid,3,13,2)
                rowbω = offsetrange(childid,3,13,4).+1

                Fzu[rowbv,colb6] = [FbXb FbQb]
                Fzu[rowbω,colb6] = [τbXb τbQb]
            end

        end
    end

    return Fzu
end