export NodeInterface, IndexedNodeInterface, name, tag, messageout, messagein
export FactorNode, functionalform, interfaces, factorisation, localmarginals, localmarginalnames, metadata
export iscontain, isfactorised, getinterface
export clusters, clusterindex
export deps, connect!, activate!
export make_node, on_make_node, AutoVar
export ValidNodeFunctionalForm, UndefinedNodeFunctionalForm, as_node_functional_form
export sdtype, Deterministic, Stochastic, isdeterministic, isstochastic
export MeanField, FullFactorisation, collect_factorisation
export DefaultFunctionalDependencies, RequireInboundFunctionalDependencies, RequireEverythingFunctionalDependencies
export @node

using Rocket
using TupleTools

import Base: show, +
import Base: getindex, setindex!, firstindex, lastindex

## Node traits

struct ValidNodeFunctionalForm end
struct UndefinedNodeFunctionalForm end

as_node_functional_form(::Function) = ValidNodeFunctionalForm()
as_node_functional_form(some)       = UndefinedNodeFunctionalForm()

## Node types

"""
    Deterministic

`Deterministic` object used to parametrize factor node object with determinstic type of relationship between variables.

See also: [`Stochastic`](@ref), [`isdeterministic`](@ref), [`isstochastic`](@ref)
"""
struct Deterministic end

"""
    Stochastic

`Stochastic` object used to parametrize factor node object with stochastic type of relationship between variables.

See also: [`Deterministic`](@ref), [`isdeterministic`](@ref), [`isstochastic`](@ref)
"""
struct Stochastic end

"""
    isdeterministic(node)

Function used to check if factor node object is deterministic or not. Returns true or false.

See also: [`Deterministic`](@ref), [`Stochastic`](@ref), [`isstochastic`](@ref)
"""
function isdeterministic end

"""
    isstochastic(node)

Function used to check if factor node object is stochastic or not. Returns true or false.

See also: [`Deterministic`](@ref), [`Stochastic`](@ref), [`isdeterministic`](@ref)
"""
function isstochastic end

isdeterministic(::Deterministic)         = true
isdeterministic(::Type{ Deterministic }) = true
isdeterministic(::Stochastic)            = false
isdeterministic(::Type{ Stochastic })    = false

isstochastic(::Stochastic)            = true
isstochastic(::Type{ Stochastic })    = true
isstochastic(::Deterministic)         = false
isstochastic(::Type{ Deterministic }) = false

"""
    sdtype(object)

Returns either `Deterministic` or `Stochastic` for a given object (if defined).

See also: [`Deterministic`](@ref), [`Stochastic`](@ref), [`isdeterministic`](@ref), [`isstochastic`](@ref)
"""
function sdtype end

sdtype(::Function) = Deterministic()

## Generic factorisation constraints

"""
    MeanField

Generic factorisation constraint used to specify a mean-field factorisation for recognition distribution `q`.

See also: [`FullFactorisation`](@ref)
"""
struct MeanField end

"""
    FullFactorisation

Generic factorisation constraint used to specify a full factorisation for recognition distribution `q`.

See also: [`MeanField`](@ref)
"""
struct FullFactorisation end

"""
    collect_factorisation(nodetype, factorisation)

This function converts given factorisation to a correct internal factorisation representation for a given node.

See also: [`MeanField`](@ref), [`FullFactorisation`](@ref)
"""
function collect_factorisation end

"""
    collect_meta(nodetype, meta)

This function converts given meta object to a correct internal meta representation for a given node.
Fallbacks to `default_meta` in case if meta is `nothing`.

See also: [`default_meta`](@ref), [`FactorNode`](@ref)
"""
function collect_meta end

collect_meta(T::Any, ::Nothing) = default_meta(T)
collect_meta(T::Any, meta::Any) = meta

"""
    default_meta(nodetype)

Returns default meta object for a given node type.

See also: [`collect_meta`](@ref), [`FactorNode`](@ref)
"""
function default_meta end

default_meta(any) = nothing

## NodeInterface

"""
    NodeInterface

`NodeInterface` object represents a single node-variable connection.

See also: [`name`](@ref), [`tag`](@ref), [`messageout`](@ref), [`messagein`](@ref)
"""
mutable struct NodeInterface
    name               :: Symbol
    m_out              :: LazyObservable{AbstractMessage}
    m_in               :: LazyObservable{AbstractMessage}
    connected_variable :: Union{Nothing, AbstractVariable}
    connected_index    :: Int

    NodeInterface(name::Symbol) = new(name, lazy(AbstractMessage), lazy(AbstractMessage), nothing, 0)
end

Base.show(io::IO, interface::NodeInterface) = print(io, string("Interface(", name(interface), ")"))

"""
    name(interface)

Returns a name of the interface.

See also: [`NodeInterface`](@ref), [`tag`](@ref)
"""
name(symbol::Symbol)              = symbol
name(interface::NodeInterface)    = name(interface.name)

"""
    tag(interface)

Returns a tag of the interface in the form of `Val{ name(interface) }`.
The major difference between tag and name is that it is possible to dispath on interface's tag in message computation rule.

See also: [`NodeInterface`](@ref), [`name`](@ref)
"""
tag(interface::NodeInterface)     = Val{ name(interface) }

"""
    messageout(interface)

Returns an outbound messages stream from the given interface.

See also: [`NodeInterface`](@ref), [`messagein`](@ref)
"""
messageout(interface::NodeInterface) = interface.m_out

"""
    messagein(interface)

Returns an inbound messages stream from the given interface.

See also: [`NodeInterface`](@ref), [`messageout`](@ref)
"""
messagein(interface::NodeInterface)  = interface.m_in

"""
    connectvariable!(interface, variable, index)

Connects a variable with the interface and given index. Index is used to distinguish this connection from others in case if variable is connected to multiple interfaces.

See also: [`NodeInterface`](@ref), [`connectedvar`](@ref), [`connectedvarindex`](@ref)
"""
function connectvariable!(interface::NodeInterface, variable, index)
    interface.connected_variable = variable
    interface.connected_index    = index
end

"""
    connectedvar(interface)

Returns connected variable for the interface.

See also: [`NodeInterface`](@ref), [`connectvariable!`](@ref), [`connectedvarindex`](@ref)
"""
connectedvar(interface::NodeInterface)      = interface.connected_variable

"""
    connectedvarindex(interface)

Returns an index of connected variable for the interface.

See also: [`NodeInterface`](@ref), [`connectvariable!`](@ref), [`connectedvar`](@ref)
"""
connectedvarindex(interface::NodeInterface) = interface.connected_index

"""
    get_pipeline_stages(interface)

Returns an instance of pipeline stages of connected variable for the given interface

See also: [`NodeInterface`](@ref), [`connectvariable!`](@ref), [`connectedvar`](@ref), [`add_inbound_pipeline_stage!`](@ref)
"""
get_pipeline_stages(interface::NodeInterface) = get_pipeline_stages(connectedvar(interface))

"""
    IndexedNodeInterface

`IndexedNodeInterface` object represents a repetative node-variable connection.
Used in cases when node may connect different number of random variables with the same name, e.g. means and precisions of Gaussian Mixture node.

See also: [`name`](@ref), [`tag`](@ref), [`messageout`](@ref), [`messagein`](@ref)
"""
struct IndexedNodeInterface
    index     :: Int
    interface :: NodeInterface
end

Base.show(io::IO, interface::IndexedNodeInterface) = print(io, string("IndexedInterface(", name(interface), ",", index(interface), ")"))

name(interface::IndexedNodeInterface)  = name(interface.interface)
index(interface::IndexedNodeInterface) = interface.index
tag(interface::IndexedNodeInterface)   = (Val{ name(interface) }(), index(interface))

messageout(interface::IndexedNodeInterface) = messageout(interface.interface)
messagein(interface::IndexedNodeInterface)  = messagein(interface.interface)

connectvariable!(interface::IndexedNodeInterface, variable, index) = connectvariable!(interface.interface, variable, index)
connectedvar(interface::IndexedNodeInterface)                      = connectedvar(interface.interface)
connectedvarindex(interface::IndexedNodeInterface)                 = connectedvarindex(interface.interface)
get_pipeline_stages(interface::IndexedNodeInterface)               = get_pipeline_stages(interface.interface)

## FactorNodeLocalMarginals

"""
    FactorNodeLocalMarginal

This object represents local marginals for some specific factor node.
Local marginal can be joint in case of structured factorisation.
Local to factor node marginal also can be shared with a corresponding marginal of some random variable.

See also: [`FactorNodeLocalMarginals`](@ref)
"""
mutable struct FactorNodeLocalMarginal
    index  :: Int
    name   :: Symbol
    stream :: Union{Nothing, MarginalObservable}

    FactorNodeLocalMarginal(index::Int, name::Symbol) = new(index, name, nothing)
end

index(localmarginal::FactorNodeLocalMarginal) = localmarginal.index
name(localmarginal::FactorNodeLocalMarginal)  = localmarginal.name

getstream(localmarginal::FactorNodeLocalMarginal) = localmarginal.stream
setstream!(localmarginal::FactorNodeLocalMarginal, observable::MarginalObservable) = localmarginal.stream = observable

"""
    FactorNodeLocalMarginals

This object acts as an iterable and indexable proxy for local marginals for some node.
"""
struct FactorNodeLocalMarginals{M}
    marginals :: M
end

function FactorNodeLocalMarginals(variablenames, factorisation)
    marginal_names = map(fcluster -> clustername(map(i -> variablenames[i], fcluster)), factorisation)
    index          = 0 # its better not to use zip here to preserve tuple-like structure
    marginals      = map((mname) -> begin index += 1; FactorNodeLocalMarginal(index, mname) end, marginal_names)
    return FactorNodeLocalMarginals(marginals)
end

## AbstractFactorNode

abstract type AbstractFactorNode end

isstochastic(factornode::AbstractFactorNode)    = isstochastic(sdtype(factornode))
isdeterministic(factornode::AbstractFactorNode) = isdeterministic(sdtype(factornode))

## Generic Factor Node

struct FactorNode{F, I, C, M, A, P} <: AbstractFactorNode
    fform          :: F
    interfaces     :: I
    factorisation  :: C
    localmarginals :: M
    metadata       :: A
    pipeline       :: P
end

function FactorNode(fform::Type{F}, interfaces::I, factorisation::C, localmarginals::M, metadata::A, pipeline::P) where { F, I, C, M, A, P }
    return FactorNode{Type{F}, I, C, M, A, P}(fform, interfaces, factorisation, localmarginals, metadata, pipeline)
end

function FactorNode(fform, varnames::NTuple{N, Symbol}, factorisation, metadata, pipeline) where N
    interfaces     = map(varname -> NodeInterface(varname), varnames)
    localmarginals = FactorNodeLocalMarginals(varnames, factorisation)
    return FactorNode(fform, interfaces, factorisation, localmarginals, metadata, pipeline)
end

function Base.show(io::IO, factornode::FactorNode)
    println(io, "FactorNode:")
    println(io, string(" form            : ", functionalform(factornode)))
    println(io, string(" sdtype          : ", sdtype(factornode)))
    println(io, string(" interfaces      : ", interfaces(factornode)))
    println(io, string(" factorisation   : ", factorisation(factornode)))
    println(io, string(" local marginals : ", localmarginalnames(factornode)))
    println(io, string(" metadata        : ", metadata(factornode)))
    println(io, string(" pipeline        : ", getpipeline(factornode)))
end

functionalform(factornode::FactorNode)            = factornode.fform
sdtype(factornode::FactorNode)                    = sdtype(functionalform(factornode))
interfaces(factornode::FactorNode)                = factornode.interfaces
factorisation(factornode::FactorNode)             = factornode.factorisation
localmarginals(factornode::FactorNode)            = factornode.localmarginals.marginals
localmarginalnames(factornode::FactorNode)        = map(name, localmarginals(factornode))
metadata(factornode::FactorNode)                  = factornode.metadata
getpipeline(factornode::FactorNode)               = factornode.pipeline

clustername(cluster) = mapreduce(v -> name(v), (a, b) -> Symbol(a, :_, b), cluster)

# Cluster is reffered to a tuple of node interfaces
clusters(factornode::FactorNode) = map(factor -> map(i -> begin return @inbounds interfaces(factornode)[i] end, factor), factorisation(factornode))

clusterindex(factornode::FactorNode, v::Symbol)                              = clusterindex(factornode, (v, ))
clusterindex(factornode::FactorNode, vindex::Int)                            = clusterindex(factornode, (vindex, ))
clusterindex(factornode::FactorNode, vars::NTuple{N, NodeInterface}) where N = clusterindex(factornode, map(v -> name(v), vars))
clusterindex(factornode::FactorNode, vars::NTuple{N, Symbol})        where N = clusterindex(factornode, map(v -> interfaceindex(factornode, v), vars))

function interfaceindex(factornode::FactorNode, iname::Symbol)
    iindex = findfirst(interface -> name(interface) === iname, interfaces(factornode))
    return iindex !== nothing ? iindex : error("Unknown interface ':$(iname)' for $(functionalform(factornode)) node")
end

function clusterindex(factornode::FactorNode, vars::NTuple{N, Int}) where N
    cindex = findfirst(cluster -> all(v -> v ∈ cluster, vars), factorisation(factornode))
    return cindex !== nothing ? cindex : error("Unknown cluster '$(vars)' for $(functionalform(factornode)) node")
end

function varclusterindex(cluster, iindex::Int)
    vcindex = findfirst(index -> index === iindex, cluster)
    return vcindex !== nothing ? vcindex : error("Invalid cluster '$(vars)' for a given interface index '$(iindex)'")
end

getinterface(factornode::FactorNode, iname::Symbol)   = @inbounds interfaces(factornode)[ interfaceindex(factornode, iname) ]

getclusterinterfaces(factornode::FactorNode, cindex::Int) = @inbounds map(i -> interfaces(factornode)[i], factorisation(factornode)[ cindex ])

iscontain(factornode::FactorNode, iname::Symbol)      = findfirst(interface -> name(interface) === iname, interfaces(factornode)) !== nothing
isfactorised(factornode::FactorNode, factor)          = findfirst(f -> f == factor, factorisation(factornode)) !== nothing

function connect!(factornode::FactorNode, iname::Symbol, variable)
    return connect!(factornode::FactorNode, iname::Symbol, variable, getlastindex(variable))
end

function connect!(factornode::FactorNode, iname::Symbol, variable, index)
    vinterface = getinterface(factornode, iname)
    connectvariable!(vinterface, variable, index)
    setmessagein!(variable, index, messageout(vinterface))
end

## Node pipeline

abstract type AbstractNodeFunctionalDependenciesPipeline end

struct FactorNodePipeline{F <: AbstractNodeFunctionalDependenciesPipeline, S <: AbstractPipelineStage}
    functional_dependencies :: F
    extra_stages            :: S
end

get_pipeline_dependencies(pipeline::FactorNodePipeline) = pipeline.functional_dependencies
get_pipeline_stages(pipeline::FactorNodePipeline)       = pipeline.extra_stages

function Base.show(io::IO, pipeline::FactorNodePipeline)
    print(io, "FactorNodePipeline(functional_dependencies = $(pipeline.functional_dependencies), extra_stages = $(pipeline.extra_stages)")
end

function collect_pipeline end

collect_pipeline(T::Any, ::Nothing)                                       = FactorNodePipeline(default_functional_dependencies_pipeline(T), EmptyPipelineStage())
collect_pipeline(T::Any, stage::AbstractPipelineStage)                    = FactorNodePipeline(default_functional_dependencies_pipeline(T), stage)
collect_pipeline(T::Any, fdp::AbstractNodeFunctionalDependenciesPipeline) = FactorNodePipeline(fdp, EmptyPipelineStage())
collect_pipeline(T::Any, pipeline::FactorNodePipeline)                    = pipeline

## Functional Dependencies

function message_dependencies end
function marginal_dependencies end

Base.:+(left::AbstractNodeFunctionalDependenciesPipeline, right::AbstractPipelineStage) = FactorNodePipeline(left, right)
Base.:+(left::FactorNodePipeline, right::AbstractPipelineStage)                         = FactorNodePipeline(left.functional_dependencies, left.extra_stages + right)

### Default

"""
    DefaultFunctionalDependencies
"""
struct DefaultFunctionalDependencies <: AbstractNodeFunctionalDependenciesPipeline end

function message_dependencies(::DefaultFunctionalDependencies, nodeinterfaces, varcluster, iindex)
    # First we remove current edge index from the list of dependencies
    vdependencies = TupleTools.deleteat(varcluster, varclusterindex(varcluster, iindex))
    # Second we map interface indices to the actual interfaces
    return map(inds -> map(i -> begin return @inbounds nodeinterfaces[i] end, inds), vdependencies)
end

function marginal_dependencies(::DefaultFunctionalDependencies, nodelocalmarginals, varcluster, cindex)
    return TupleTools.deleteat(nodelocalmarginals, cindex)
end

### With inbound

struct RequireInboundFunctionalDependencies{I, S} <: AbstractNodeFunctionalDependenciesPipeline
    indices    :: I
    start_with :: S
end

struct InterfacePluginStartWithMessage{M, S}
    msg        :: M
    start_with :: S
end

name(p::InterfacePluginStartWithMessage)      = name(p.msg)
messagein(p::InterfacePluginStartWithMessage) = messagein(p.start_with, p)

messagein(::Nothing, p::InterfacePluginStartWithMessage) = messagein(p.msg)
messagein(something, p::InterfacePluginStartWithMessage) = messagein(p.msg) |> start_with(Message(something, false, true))

function message_dependencies(dependencies::RequireInboundFunctionalDependencies, nodeinterfaces, varcluster, iindex)

    # First we find dependency index in `indices`, we use it later to find `start_with` distribution
    depindex = findfirst((i) -> i === iindex, dependencies.indices)

    # If we have `depindex` in our `indices` we include it in our list of functional dependencies. It effectively forces rule to require inbound message
    if depindex !== nothing
        # `mapindex` is a lambda function here
        mapindex = let nodeinterfaces = nodeinterfaces, depindex = depindex
            (i) -> begin
                interface = @inbounds nodeinterfaces[i]
                # InterfacePluginStartWithMessage is a proxy structure for `name` and `messagein` method for an interface
                # It returns the same name but modifies `messagein` to return an observable with `start_with` operator
                return i === iindex ? InterfacePluginStartWithMessage(interface, dependencies.start_with[depindex]) : interface
            end
        end
        return map(inds -> map(mapindex, inds), varcluster)
    else
        return message_dependencies(DefaultFunctionalDependencies(), nodeinterfaces, varcluster, iindex)
    end
end

function marginal_dependencies(::RequireInboundFunctionalDependencies, nodelocalmarginals, varcluster, cindex)
    return marginal_dependencies(DefaultFunctionalDependencies(), nodelocalmarginals, varcluster, cindex)
end

### Everything

struct RequireEverythingFunctionalDependencies <: AbstractNodeFunctionalDependenciesPipeline end

function ReactiveMP.message_dependencies(::RequireEverythingFunctionalDependencies, nodeinterfaces, varcluster, iindex)
    # Return all node interfaces including the edge we are trying to compuate a message on
    return nodeinterfaces
end

function ReactiveMP.marginal_dependencies(::RequireEverythingFunctionalDependencies, nodelocalmarginals, varcluster, cindex)
    # Returns only local marginals based on local q factorisation, it does not return all possible combinations of all joint posterior marginals
    return nodelocalmarginals
end

###

default_functional_dependencies_pipeline(_) = DefaultFunctionalDependencies()

### Generic

function functional_dependencies(dependencies, factornode::FactorNode, iname::Symbol)
    return functional_dependencies(dependencies, factornode, interfaceindex(factornode, iname))
end

function functional_dependencies(dependencies, factornode::FactorNode, iindex::Int)
    cindex  = clusterindex(factornode, iindex)

    nodeinterfaces     = interfaces(factornode)
    nodeclusters       = factorisation(factornode)
    nodelocalmarginals = localmarginals(factornode)

    varcluster = @inbounds nodeclusters[ cindex ]

    messages  = message_dependencies(dependencies, nodeinterfaces, varcluster, iindex)
    marginals = marginal_dependencies(dependencies, nodelocalmarginals, varcluster, cindex)

    return tuple(messages...), tuple(marginals...)
end

function get_messages_observable(factornode, messages)
    if !isempty(messages)
        msgs_names      = Val{ map(name, messages) }
        msgs_observable = combineLatest(map(m -> messagein(m), messages), PushNew())
        return msgs_names, msgs_observable
    else
        return nothing, of(nothing)
    end
end

function get_marginals_observable(factornode, marginals)
    if !isempty(marginals)
        marginal_names       = Val{ map(name, marginals) }
        marginals_streams    = map(marginal -> getmarginal!(factornode, marginal, IncludeAll()), marginals)
        marginals_observable = combineLatestUpdates(marginals_streams, PushNew())
        return marginal_names, marginals_observable
    else
        return nothing, of(nothing)
    end
end

# Variational Message Passing
apply_mapping(msgs_observable, marginals_observable, mapping) = (dependencies) -> VariationalMessage(dependencies[1], dependencies[2], mapping)

# Fallback for Belief Propagation
apply_mapping(msgs_observable, marginals_observable::SingleObservable{Nothing}, mapping) = mapping

function activate!(model, factornode::AbstractFactorNode)
    fform = functionalform(factornode)
    meta  = metadata(factornode)
    node_pipeline = getpipeline(factornode)

    node_pipeline_dependencies = get_pipeline_dependencies(node_pipeline)
    node_pipeline_extra_stages = get_pipeline_stages(node_pipeline)

    for (iindex, interface) in enumerate(interfaces(factornode))
        message_dependencies, marginal_dependencies = functional_dependencies(node_pipeline_dependencies, factornode, iindex)

        msgs_names, msgs_observable          = get_messages_observable(factornode, message_dependencies)
        marginal_names, marginals_observable = get_marginals_observable(factornode, marginal_dependencies)

        vtag        = tag(interface)
        vconstraint = local_constraint(connectedvar(interface))

        vmessageout = combineLatest((msgs_observable, marginals_observable), PushNew()) # TODO check PushEach
        vmessageout = apply_pipeline_stage(get_pipeline_stages(interface), factornode, vtag, vmessageout)

        mapping = MessageMapping(fform, vtag, vconstraint, msgs_names, marginal_names, meta, factornode)
        mapping = apply_mapping(msgs_observable, marginals_observable, mapping)

        vmessageout = vmessageout |> map(AbstractMessage, mapping)
        vmessageout = apply_pipeline_stage(get_pipeline_stages(getoptions(model)), factornode, vtag, vmessageout)
        vmessageout = apply_pipeline_stage(node_pipeline_extra_stages, factornode, vtag, vmessageout)
        vmessageout = vmessageout |> schedule_on(global_reactive_scheduler(getoptions(model)))

        set!(messageout(interface), vmessageout |> share_recent())
        set!(messagein(interface), messageout(connectedvar(interface), connectedvarindex(interface)))
    end
end

function setmarginal!(factornode::FactorNode, cname::Symbol, marginal)
    lindex = findnext(lmarginal -> name(lmarginal) === cname, localmarginals(factornode), 1)
    @assert lindex !== nothing "Invalid local marginal id: $cname"
    lmarginal = @inbounds localmarginals(factornode)[ lindex ]
    setmarginal!(getstream(lmarginal), marginal)
end

function getmarginal!(factornode::FactorNode, localmarginal::FactorNodeLocalMarginal)
    return getmarginal!(factornode, localmarginal, IncludeAll())
end

function getmarginal!(factornode::FactorNode, localmarginal::FactorNodeLocalMarginal, skip_strategy::MarginalSkipStrategy)
    cached_stream = getstream(localmarginal)

    if cached_stream !== nothing
        return as_marginal_observable(cached_stream, skip_strategy)
    end

    clusterindex = index(localmarginal)

    marginalname = name(localmarginal)
    marginalsize = @inbounds length(factorisation(factornode)[ clusterindex ])

    if marginalsize === 1
        # Cluster contains only one variable, we can take marginal over this variable
        vmarginal = getmarginal(connectedvar(getinterface(factornode, marginalname)), IncludeAll())
        setstream!(localmarginal, vmarginal)
        return as_marginal_observable(vmarginal, skip_strategy)
    else
        cmarginal = MarginalObservable()

        message_dependencies  = tuple(getclusterinterfaces(factornode, clusterindex)...)
        marginal_dependencies = tuple(TupleTools.deleteat(localmarginals(factornode), clusterindex)...)

        msgs_names, msgs_observable          = get_messages_observable(factornode, message_dependencies)
        marginal_names, marginals_observable = get_marginals_observable(factornode, marginal_dependencies)

        fform       = functionalform(factornode)
        vtag        = Val{ name(localmarginal) }
        meta        = metadata(factornode)

        mapping = let fform = fform, vtag = vtag, msgs_names = msgs_names, marginal_names = marginal_names, meta = meta, factornode = factornode
            (dependencies) -> begin
                messages  = dependencies[1]
                marginals = getrecent(dependencies[2])

                # Marginal is clamped if all of the inputs are clamped
                is_marginal_clamped = __check_all(is_clamped, messages) && __check_all(is_clamped, marginals)

                # Marginal is initial if it is not clamped and all of the inputs are either clamped or initial
                is_marginal_initial = !is_marginal_clamped && (__check_all(m -> is_clamped(m) || is_initial(m), messages) && __check_all(m -> is_clamped(m) || is_initial(m), marginals))

                return Marginal(marginalrule(fform, vtag, msgs_names, messages, marginal_names, marginals, meta, factornode), is_marginal_clamped, is_marginal_initial)
            end
        end

        # TODO: discontinue operater is needed for loopy belief propagation? Check
        marginalout = combineLatest((msgs_observable, marginals_observable), PushNew()) |> discontinue() |> map(Marginal, mapping)

        connect!(cmarginal, marginalout) # MarginalObservable has RecentSubject by default, there is no need to share_recent() here

        setstream!(localmarginal, cmarginal)

        return as_marginal_observable(cmarginal, skip_strategy)
    end
end

## TODO
function conjugate_type end

## make_node

struct AutoVar
    name :: Symbol
end

getname(autovar::AutoVar) = autovar.name

function make_node end

function interface_get_index end
function interface_get_name end

function interface_get_index(::Type{ Val{ Node } }, ::Type{ Val{ Interface } }) where { Node, Interface }
    error("Node $Node has no interface named $Interface")
end

function interface_get_name(::Type{ Val{ Node } }, ::Type{ Val{ Interface } }) where { Node, Interface }
    error("Node $Node has no interface named $Interface")
end

make_node(fform, ::AutoVar, ::Vararg{ <: AbstractVariable }; kwargs...) = error("Unknown functional form '$(fform)' used for node specification.")
make_node(fform, args::Vararg{ <: AbstractVariable }; kwargs...)        = error("Unknown functional form '$(fform)' used for node specification.")

function make_node(fform::Function, autovar::AutoVar, args::Vararg{ <: ConstVariable }; kwargs...)
    var  = constvar(getname(autovar), fform(map((d) -> getconst(d), args)...))
    return nothing, var
end

function make_node(::Type{ T }, autovar::AutoVar, args::Vararg{ <: ConstVariable }; kwargs...) where T
    var  = constvar(getname(autovar), T(map((d) -> getconst(d), args)...))
    return nothing, var
end

function make_node(fform::Function, autovar::AutoVar, args::Vararg{ <: DataVariable{ <: PointMass } }; kwargs...)
    # TODO
    message_cb = let fform = fform
        (d::Tuple) -> Message(fform(d...), false, false)
    end

    subject = combineLatest(tuple(map((a) -> messageout(a, getlastindex(a)) |> map(Any, (d) -> mean(getdata(d))), args)...), PushNew()) |> map(Message, message_cb)
    var     = datavar(getname(autovar), Any, subject = subject)

    return nothing, var
end

# end

## macro helpers

import .MacroHelpers

"""
    @node(fformtype, sdtype, interfaces_list)

`@node` macro creates a node for a `fformtype` type object

# Examples
```julia

struct MyNormalDistribution
    mean :: Float64
    var  :: Float64
end

@node MyNormalDistribution Stochastic [ out, mean, var ]
```

```julia

@node typeof(+) Deterministic [ out, in1, in2 ]
```

See also: [`make_node`](@ref), [`Stochastic`](@ref), [`Deterministic`](@ref)
"""
macro node(fformtype, sdtype, interfaces_list)

    fbottomtype = MacroHelpers.bottom_type(fformtype)
    fuppertype  = MacroHelpers.upper_type(fformtype)

    @assert sdtype ∈ [ :Stochastic, :Deterministic ] "Invalid sdtype $(sdtype). Can be either Stochastic or Deterministic."

    @capture(interfaces_list, [ interfaces_args__ ]) ||
        error("Invalid interfaces specification.")

    interfaces = map(interfaces_args) do arg
        if @capture(arg, name_Symbol)
            return (name, [])
        elseif @capture(arg, (name_Symbol, aliases = [ aliases__ ]))
            @assert all(a -> a isa Symbol && !isequal(a, name), aliases)
            return (name, aliases)
        else
            error("Interface specification should have a 'name' or (name, aliases = [ alias1, alias2,... ]) signature.")
        end
    end

    @assert length(interfaces) !== 0 "Node should have at least one interface."

    names = map(d -> first(d), interfaces)

    names_quoted_tuple     = Expr(:tuple, map(name -> Expr(:quote, name), names)...)
    names_indices          = Expr(:tuple, map(i -> i, 1:length(names))...)
    names_splitted_indices = Expr(:tuple, map(i -> Expr(:tuple, i), 1:length(names))...)

    interface_args        = map(name -> :($name::AbstractVariable), names)
    interface_connections = map(name -> :(ReactiveMP.connect!(node, $(Expr(:quote, name)), $name)), names)

    # Here we create helpers function for GraphPPL.jl interfacing
    # They are used to convert interface names from `where { q = q(x, y)q(z) }` to an equivalent tuple respresentation, e.g. `((1, 2), (3, ))`
    # The general recipe to get a proper index is to call `interface_get_index(Val{ :NodeTypeName }, interface_get_name(Val{ :NodeTypeName }, Val{ :name_expression }))`
    interface_name_getters = map(enumerate(interfaces)) do (index, interface)
        name    = first(interface)
        aliases = last(interface)

        index_name_getter  = :(ReactiveMP.interface_get_index(::Type{ Val{ $(Expr(:quote, fbottomtype)) } }, ::Type{ Val{ $(Expr(:quote, name)) } }) = $(index))
        name_symbol_getter = :(ReactiveMP.interface_get_name(::Type{ Val{ $(Expr(:quote, fbottomtype)) } }, ::Type{ Val{ $(Expr(:quote, name)) } }) = $(Expr(:quote, name)))
        name_index_getter  = :(ReactiveMP.interface_get_name(::Type{ Val{ $(Expr(:quote, fbottomtype)) } }, ::Type{ Val{ $index } }) = $(Expr(:quote, name)))

        alias_getters = map(aliases) do alias
            return :(ReactiveMP.interface_get_name(::Type{ Val{ $(Expr(:quote, fbottomtype)) } }, ::Type{ Val{ $(Expr(:quote, alias)) } }) = $(Expr(:quote, name)))
        end

        return quote
            $index_name_getter
            $name_symbol_getter
            $name_index_getter
            $(alias_getters...)
        end
    end

    # By default every argument passed to a factorisation option of the node is transformed by
    # `collect_factorisation` function to have a tuple like structure.
    # The default recipe is simple: for stochastic nodes we convert `FullFactorisation` and `MeanField` objects
    # to their tuple of indices equivalents. For deterministic nodes any factorisation is replaced by a FullFactorisation equivalent
    factorisation_collectors = if sdtype === :Stochastic
        quote
            ReactiveMP.collect_factorisation(::$fuppertype, factorisation::Tuple)           = factorisation
            ReactiveMP.collect_factorisation(::$fuppertype, ::ReactiveMP.FullFactorisation) = ($names_indices, )
            ReactiveMP.collect_factorisation(::$fuppertype, ::ReactiveMP.MeanField)         = $names_splitted_indices
        end

    elseif sdtype === :Deterministic
        quote
            ReactiveMP.collect_factorisation(::$fuppertype, factorisation::Tuple)           = ($names_indices, )
            ReactiveMP.collect_factorisation(::$fuppertype, ::ReactiveMP.FullFactorisation) = ($names_indices, )
            ReactiveMP.collect_factorisation(::$fuppertype, ::ReactiveMP.MeanField)         = ($names_indices, )
        end
    else
        error("Unreachable in @node macro.")
    end

    make_node_const_mapping = if sdtype === :Stochastic
        quote
            function ReactiveMP.make_node(fform::$fuppertype, autovar::ReactiveMP.AutoVar, args::Vararg{ <: ReactiveMP.ConstVariable{ <: ReactiveMP.PointMass } }; kwargs...)
                var  = ReactiveMP.randomvar(ReactiveMP.getname(autovar))
                node = ReactiveMP.make_node(fform, var, args...; kwargs...)
                return node, var
            end
        end
    elseif sdtype === :Deterministic
        quote
            function ReactiveMP.make_node(fform::$fuppertype, autovar::ReactiveMP.AutoVar, args::Vararg{ <: ReactiveMP.ConstVariable{ <: ReactiveMP.PointMass } }; kwargs...)
                var  = ReactiveMP.constvar(ReactiveMP.getname(autovar), fform(map((d) -> ReactiveMP.getconst(d), args)...))
                return nothing, var
            end
        end
    else
        error("Unreachable in @node macro.")
    end

    res = quote

        ReactiveMP.as_node_functional_form(::$fuppertype) = ReactiveMP.ValidNodeFunctionalForm()

        ReactiveMP.sdtype(::$fuppertype) = (ReactiveMP.$sdtype)()

        function ReactiveMP.make_node(::$fuppertype; factorisation = ($names_indices, ), meta = nothing, pipeline = nothing)
            return ReactiveMP.FactorNode($fbottomtype, $names_quoted_tuple, ReactiveMP.collect_factorisation($fbottomtype, factorisation), ReactiveMP.collect_meta($fbottomtype, meta), ReactiveMP.collect_pipeline($fbottomtype, pipeline))
        end

        function ReactiveMP.make_node(::$fuppertype, $(interface_args...); factorisation = ($names_indices, ), meta = nothing, pipeline = nothing)
            node = ReactiveMP.make_node($fbottomtype, factorisation = factorisation, meta = meta, pipeline = pipeline)
            $(interface_connections...)
            return node
        end

        function ReactiveMP.make_node(fform::$fuppertype, autovar::ReactiveMP.AutoVar, args::Vararg{ <: ReactiveMP.AbstractVariable }; kwargs...)
            var  = ReactiveMP.randomvar(ReactiveMP.getname(autovar))
            node = ReactiveMP.make_node(fform, var, args...; kwargs...)
            return node, var
        end

        $(make_node_const_mapping)
        $(interface_name_getters...)

        $factorisation_collectors

    end

    return esc(res)
end


import Base: show, copy

"""
$(TYPEDEF)
    options for Solver
"""
mutable struct SolverOptions

    ##########################
    ## Satellite Quaternion ##
    ##########################

    sat_att::Bool #if this is true, the 4:7 numbers in the state vector are a quaternion

    ###################
    ## Functionality ##
    ###################
    verbose::Bool # Display solve statistics at each iteration
    benchmark::Bool # Run benchmarks on forward and backward passes
    live_plotting::Bool # Plot state and control trajectories during solve

    ##########################
    ## Convergence Criteria ##
    ##########################
    cost_tolerance::Float64 # dJ < ϵ, cost convergence criteria for unconstrained solve or to enter outerloop for constrained solve
    cost_tolerance_intermediate::Float64 # dJ < ϵ_int, intermediate cost convergence criteria to enter outerloop of constrained solve
    constraint_tolerance::Float64 # max(constraint) < ϵ, constraint convergence criteria
    constraint_tolerance_intermediate::Float64 # max(constraint) < ϵ_int, intermediate constraint convergence criteria
    constraint_tolerance_infeasible::Float64  # infeasible control constraint tolerance
    gradient_norm_tolerance::Float64 # gradient_norm < ϵ, gradient norm convergence criteria
    gradient_norm_tolerance_intermediate::Float64 # gradient_norm_int < ϵ, gradient norm intermediate convergence criteria
    gradient_type::Symbol # type of gradient to evaluate: Augmented Lagrangian ∂L/∂u :AuLa, feedfoward Σ|d|^2 :feedforward, Todorov normalized feedforward :todorov

    ######################
    ## Solve Parameters ##
    ######################
    iterations::Int64 # caps cumulative number of iLQR iterations (see iterations_innerloop, iterations_outerloop)
    dJ_counter_limit::Int64 # restricts the total number of times a forward pass fails, resulting in regularization, before entering an outerloop update

    ###################
    ## Iterative LQR ##
    ###################
    square_root::Bool # use square root method backward pass for numerical conditioning
    iterations_innerloop::Int64 # maximum iterations for inner loop (iLQR)
    line_search_lower_bound::Float64 # forward pass approximate line search lower bound, 0 < line_search_lower_bound < line_search_upper_bound
    line_search_upper_bound::Float64 # forward pass approximate line search upper bound, 0 < line_search_lower_bound < line_search_upper_bound < Inf
    iterations_linesearch::Int64 # maximum number of backtracking steps during forward pass line search

    # Regularization
    bp_reg_initial::Float64 # initial regularization
    bp_reg_increase_factor::Float64 # regularization scaling factor
    bp_reg_max::Float64 # maximum regularization value
    bp_reg_min::Float64 # minimum regularization value
    bp_reg_type::Symbol # type of regularization- control: () + ρI, state: (S + ρI); see "Synthesis and Stabilization of Complex Behaviors through Online Trajectory Optimization"
    bp_reg_fp::Float64 # additive regularization when forward pass reaches max iterations
    eigenvalue_scaling::Float64 # add this multiple of the magnitude of the most negative eigenvalue to Quu decomposition to make positive definite
    eigenvalue_threshold::Float64 # eigenvalues less than this threshold will be increased using the additive eigenvalue scaling
    bp_sqrt_inv_type::Symbol # type of fix for potential inverse failure, :pseudo -pseudo inverse of Wxx, :reg - regularize Wxx to make invertible
    bp_reg_sqrt_initial::Float64 # initial regularization for square root method
    bp_reg_sqrt_increase_factor::Float64 # regularization scaling factor for square root method

    ##########################
    ## Augmented Lagrangian ##
    ##########################
    iterations_outerloop::Int64 # maximum outerloop updates
    dual_min::Float64 # minimum Lagrange multiplier
    dual_max::Float64 # maximum Lagrange multiplier
    penalty_max::Float64 # maximum penalty term
    penalty_initial::Float64 # initial penalty term
    penalty_scaling::Float64 # penalty update multiplier; penalty_scaling > 0
    penalty_scaling_no::Float64 # penalty update multiplier when μ should not be update, typically 1.0 (or 1.0 + ϵ)
    constraint_decrease_ratio::Float64 # update term; 0 < constraint_decrease_ratio < 1
    outer_loop_update_type::Symbol # type of outer loop update (default, momentum, individual, accelerated)
    use_second_order_dual_update::Bool # second order update for Lagrange multipliers once sqrt(cost tolerance | gradient) < desired tolerance
    penalty_update_frequency::Int  # determines how many iterations should pass before the penalty is updated (1 is every iteration)
    constraint_tolerance_second_order_dual_update::Float64 # constraint tolerance for switching to second order dual update
    active_constraint_tolerance::Float64 # numerical tolerance for constraint violation
    use_nesterov::Bool # accelerated gradient descent for dual variables
    use_penalty_burnin::Bool # perform only penalty updates (no dual updates) until constraint_tolerance_intermediate < ϵ_int
    al_type::Symbol  # [:default, :algencan] pick which Augmented Lagrangian formulation to use. ALGENCAN uses a slightly different check for inactive constraints.

    ######################
    ## Infeasible Start ##
    ######################
    R_infeasible::Float64 # regularization term for infeasible controls
    resolve_feasible::Bool # resolve feasible problem after infeasible solve
    feasible_projection::Bool # project infeasible solution into feasible space w/ BP, rollout
    penalty_initial_infeasible::Float64 # initial penalty term for infeasible controls
    penalty_scaling_infeasible::Float64 # penalty update rate for infeasible controls

    ##################
    ## Minimum Time ##
    ##################
    R_minimum_time::Float64 # regularization term for dt
    max_dt::Float64 # maximum allowable dt
    min_dt::Float64 # minimum allowable dt
    minimum_time_tf_estimate::Float64 # initial guess for the length of the minimum time problem (in seconds)
    minimum_time_dt_estimate::Float64 # initial guess for dt of the minimum time problem (in seconds)
    penalty_initial_minimum_time_inequality::Float64 # initial penalty term for minimum time bounds constraints
    penalty_initial_minimum_time_equality::Float64 # initial penalty term for minimum time equality constraints
    penalty_scaling_minimum_time_inequality::Float64 # penalty update rate for minimum time bounds constraints
    penalty_scaling_minimum_time_equality::Float64 # penalty update rate for minimum time equality constraints

    #############################
    ## Solver Numerical Limits ##
    #############################
    max_cost_value::Float64 # maximum cost value, if exceded solve will error
    max_state_value::Float64 # maximum state value, evaluated during rollout, if exceded solve will error
    max_control_value::Float64 # maximum control value, evaluated during rollout, if exceded solve will error

    function SolverOptions(;sat_att=false,verbose=false,benchmark=false,
        live_plotting=false,
        cost_tolerance=1.0e-4,
        cost_tolerance_intermediate=1.0e-3,
        constraint_tolerance=1.0e-3,
        constraint_tolerance_intermediate=sqrt(constraint_tolerance),
        constraint_tolerance_infeasible=1.0e-5,
        gradient_norm_tolerance=1.0e-5,
        gradient_norm_tolerance_intermediate=1.0e-5,
        gradient_type=:todorov,
        iterations=500,
        dJ_counter_limit=10,
        square_root=false,
        iterations_innerloop=250,
        line_search_lower_bound=1.0e-8,
        line_search_upper_bound=10.0,
        iterations_linesearch=20,
        bp_reg_initial=0.0,
        bp_reg_increase_factor=1.6,
        bp_reg_max=1.0e8,
        bp_reg_min=1.0e-8,
        bp_reg_type=:control,
        bp_reg_fp=10.0,
        eigenvalue_scaling=2.0,
        eigenvalue_threshold=1e-8,
        bp_sqrt_inv_type=:pseudo,
        bp_reg_sqrt_initial=1.0e-6,
        bp_reg_sqrt_increase_factor=10.0,
        iterations_outerloop=30,
        dual_min=-1.0e8,
        dual_max=1.0e8,
        penalty_max=1.0e8,
        penalty_initial=1.0,
        penalty_scaling=10.0,
        penalty_scaling_no=1.0,
        constraint_decrease_ratio=0.25,
        outer_loop_update_type=:default,
        use_second_order_dual_update=false,
        penalty_update_frequency=1,
        constraint_tolerance_second_order_dual_update=sqrt(constraint_tolerance),
        active_constraint_tolerance=0.0,
        use_nesterov=false,
        use_penalty_burnin=false,
        al_type=:default,
        R_infeasible=1.0,
        resolve_feasible=true,
        feasible_projection=true,
        penalty_initial_infeasible=1.0,
        penalty_scaling_infeasible=10.0,
        R_minimum_time=1.0,
        max_dt=1.0,
        min_dt=1.0e-3,
        minimum_time_tf_estimate=0.0,
        minimum_time_dt_estimate=0.0,
        penalty_initial_minimum_time_inequality=1.0,
        penalty_initial_minimum_time_equality=1.0,
        penalty_scaling_minimum_time_inequality=1.0,
        penalty_scaling_minimum_time_equality=1.0,
        max_cost_value=1.0e8,
        max_state_value=1.0e8,
        max_control_value=1.0e8)

        new(sat_att,
            verbose,
            benchmark,
            live_plotting,
            cost_tolerance,
            cost_tolerance_intermediate,
            constraint_tolerance,
            constraint_tolerance_intermediate,
            constraint_tolerance_infeasible,
            gradient_norm_tolerance,
            gradient_norm_tolerance_intermediate,
            gradient_type,
            iterations,
            dJ_counter_limit,
            square_root,
            iterations_innerloop,
            line_search_lower_bound,
            line_search_upper_bound,
            iterations_linesearch,
            bp_reg_initial,
            bp_reg_increase_factor,
            bp_reg_max,
            bp_reg_min,
            bp_reg_type,
            bp_reg_fp,
            eigenvalue_scaling,
            eigenvalue_threshold,
            bp_sqrt_inv_type,
            bp_reg_sqrt_initial,
            bp_reg_sqrt_increase_factor,
            iterations_outerloop,
            dual_min,
            dual_max,
            penalty_max,
            penalty_initial,
            penalty_scaling,
            penalty_scaling_no,
            constraint_decrease_ratio,
            outer_loop_update_type,
            use_second_order_dual_update,
            penalty_update_frequency,
            constraint_tolerance_second_order_dual_update,
            active_constraint_tolerance,
            use_nesterov,
            use_penalty_burnin,
            al_type,
            R_infeasible,
            resolve_feasible,
            feasible_projection,
            penalty_initial_infeasible,
            penalty_scaling_infeasible,
            R_minimum_time,
            max_dt,
            min_dt,
            minimum_time_tf_estimate,
            minimum_time_dt_estimate,
            penalty_initial_minimum_time_inequality,
            penalty_initial_minimum_time_equality,
            penalty_scaling_minimum_time_inequality,
            penalty_scaling_minimum_time_equality,
            max_cost_value,
            max_state_value,
            max_control_value)
    end

end

copy(opts::SolverOptions) = SolverOptions(;[name=>getfield(opts,name) for name in fieldnames(typeof(opts))]...)

function Base.:(==)(A::SolverOptions, B::SolverOptions)
    for name in fieldnames(A)
        if getfield(A,name) != getfield(B,name)
            return false
        end
    end
    return true
end


module MHDFlows

using
  CUDA,
  Statistics,
  SpecialFunctions,
  Reexport,
  DocStringExtensions

@reexport using FourierFlows

using LinearAlgebra: mul!, ldiv!
import Base: show, summary

"Abstract supertype for problem."
abstract type AbstractProblem end
abstract type MHDVars <: AbstractVars end

include("DyeModule.jl")
include("Problems.jl")
include("pgen.jl")
include("HDSolver.jl")
include("MHDSolver.jl")
include("DiagnosticWrapper.jl")
include("integrator.jl")
include("datastructure.jl")
include("utils/VectorCalculus.jl")
include("utils/MHDAnalysis.jl")
include("utils/GeometryFunction.jl")

export Problem,
       TimeIntegrator!,
       Curl,
       Div,
       LaplaceSolver,
       Crossproduct,
       Dotproduct,
       xy_to_polar,
       ScaleDecomposition,
       h_k,
       VectorPotential,
       LaplaceSolver,
       getL

end


############################################################
## joAbstractLinearOperator(s) - overloaded Base functions

# eltype(jo)
eltype(A::joAbstractLinearOperator{DDT,RDT}) where {DDT,RDT} = promote_type(DDT,RDT)

# deltype(jo)
deltype(A::joAbstractLinearOperator{DDT,RDT}) where {DDT,RDT} = DDT

# reltype(jo)
reltype(A::joAbstractLinearOperator{DDT,RDT}) where {DDT,RDT} = RDT

# show(jo)
function show(A::joAbstractLinearOperator)
    println("Type: $(typeof(A).name)")
    println("Name: $(A.name)")
    println("Size: $(size(A))")
    println(" DDT: $(deltype(A))")
    println(" RDT: $(reltype(A))")
    return nothing
end

# display(jo)
display(A::joAbstractLinearOperator) = println((typeof(A),(A.m,A.n),A.name))

# size(jo)
size(A::joAbstractLinearOperator) = A.m,A.n

# size(jo,1/2)
function size(A::joAbstractLinearOperator,ind::Integer)
    if ind==1
		return A.m
	elseif ind==2
		return A.n
	else
		throw(joAbstractLinearOperatorException("invalid index"))
	end
end

# length(jo)
length(A::joAbstractLinearOperator) = A.m*A.n

# jo_full(jo)
jo_full(A::joAbstractLinearOperator{DDT,RDT}) where {DDT,RDT} = A*jo_eye(DDT,A.n)

# norm(jo)
norm(A::joAbstractLinearOperator,p::Real=2) = norm(elements(A),p)

# real(jo)
#real{DDT<:Real,RDT<:Real}(A::joAbstractLinearOperator{DDT,RDT}) = A
function real(A::joAbstractLinearOperator{DDT,RDT}) where {DDT,RDT}
    throw(joAbstractLinearOperatorException("real(jo) not implemented"))
end
joReal(A::joAbstractLinearOperator{DDT,RDT}) where {DDT,RDT} = real(A)

# imag(jo)
#imag{DDT<:Real,RDT<:Real}(A::joAbstractLinearOperator{DDT,RDT}) = joZeros(A.m,A.n,DDT,RDT)
function imag(A::joAbstractLinearOperator{DDT,RDT}) where {DDT,RDT}
    throw(joAbstractLinearOperatorException("imag(jo) not implemented"))
end
joImag(A::joAbstractLinearOperator{DDT,RDT}) where {DDT,RDT} = imag(A)

# conj(jo)
joConj(A::joAbstractLinearOperator) = conj(A)
conj(A::joLinearOperator{DDT,RDT}) where {DDT,RDT} =
    joLinearOperator{DDT,RDT}("conj("*A.name*")",A.m,A.n,
        get(A.fop_C), A.fop_A, A.fop_T, A.fop,
        A.iop_C, A.iop_A, A.iop_T, A.iop
        )
conj(A::joMatrix{DDT,RDT}) where {DDT,RDT} =
    joMatrix{DDT,RDT}("conj("*A.name*")",A.m,A.n,
        A.fop_C, A.fop_A, A.fop_T, A.fop,
        A.iop_C, A.iop_A, A.iop_T, A.iop
        )
conj(A::joLinearFunction{DDT,RDT}) where {DDT,RDT} =
    joLinearFunction{DDT,RDT}("conj("*A.name*")",A.m,A.n,
        get(A.fop_C), A.fop_A, A.fop_T, A.fop, A.fMVok,
        A.iop_C, A.iop_A, A.iop_T, A.iop, A.iMVok
        )

# transpose(jo)
transpose(A::joLinearOperator{DDT,RDT}) where {DDT,RDT} =
    joLinearOperator{RDT,DDT}("transpose("*A.name*")",A.n,A.m,
        get(A.fop_T), A.fop, A.fop_C, A.fop_A,
        A.iop_T, A.iop, A.iop_C, A.iop_A
        )
transpose(A::joMatrix{DDT,RDT}) where {DDT,RDT} =
    joMatrix{RDT,DDT}("transpose("*A.name*")",A.n,A.m,
        A.fop_T, A.fop, A.fop_C, A.fop_A,
        A.iop_T, A.iop, A.iop_C, A.iop_A
        )
transpose(A::joLinearFunction{DDT,RDT}) where {DDT,RDT} =
    joLinearFunction{RDT,DDT}("transpose("*A.name*")",A.n,A.m,
        get(A.fop_T), A.fop, A.fop_C, A.fop_A, A.fMVok,
        A.iop_T, A.iop, A.iop_C, A.iop_A, A.iMVok
        )

# adjoint(jo)
adjoint(A::joLinearOperator{DDT,RDT}) where {DDT,RDT} =
    joLinearOperator{RDT,DDT}("adjoint("*A.name*")",A.n,A.m,
        get(A.fop_A), A.fop_C, A.fop, A.fop_T,
        A.iop_A, A.iop_C, A.iop, A.iop_T
        )
adjoint(A::joMatrix{DDT,RDT}) where {DDT,RDT} =
    joMatrix{RDT,DDT}("adjoint("*A.name*")",A.n,A.m,
        A.fop_A, A.fop_C, A.fop, A.fop_T,
        A.iop_A, A.iop_C, A.iop, A.iop_T
        )
adjoint(A::joLinearFunction{DDT,RDT}) where {DDT,RDT} =
    joLinearFunction{RDT,DDT}("adjoint("*A.name*")",A.n,A.m,
        get(A.fop_A), A.fop_C, A.fop, A.fop_T, A.fMVok,
        A.iop_A, A.iop_C, A.iop, A.iop_T, A.iMVok
        )

# isreal(jo)
isreal(A :: joAbstractLinearOperator{DDT,RDT}) where {DDT,RDT} = (DDT<:Real && RDT<:Real)

# issymmetric(jo)
issymmetric(A::joAbstractLinearOperator{DDT,RDT}) where {DDT,RDT} =
    (A.m == A.n && (norm(elements(A)-elements(transpose(A))) < joTol))

# ishermitian(jo)
ishermitian(A::joAbstractLinearOperator{DDT,RDT}) where {DDT,RDT} =
    (A.m == A.n && (norm(elements(A)-elements(adjoint(A))) < joTol))

# getindex(jo,...)
getindex(A::joAbstractLinearOperator{DDT,RDT},::Colon,i::Integer) where {DDT,RDT} =
    joHelpers_etc.elements_column_helper(A,i)
getindex(A::joAbstractLinearOperator{DDT,RDT},::Colon,r::UnitRange{URT}) where {DDT,RDT,URT<:Integer}  =
    joHelpers_etc.elements_column_helper(A,r)

############################################################
## overloaded Base *(...jo...)

# *(jo,jo)
function *(A::joAbstractLinearOperator{CDT,ARDT},B::joAbstractLinearOperator{BDDT,CDT}) where {ARDT,BDDT,CDT}
    size(A,2) == size(B,1) || throw(joAbstractLinearOperatorException("shape mismatch"))
    return joLinearOperator{BDDT,ARDT}("("*A.name*"*"*B.name*")",size(A,1),size(B,2),
        v1->A*(B*v1),
        v2->transpose(B)*(transpose(A)*v2),
        v3->adjoint(B)*(adjoint(A)*v3),
        v4->conj(A)*(conj(B)*v4),
        @joNF, @joNF, @joNF, @joNF
        )
end
# in joAbstractSAparallelLinearOperator/base_functions.jl
#function *(A::joSAdistributingLinearOperator{CDT,ARDT,2},B::joAbstractLinearOperator{BDDT,CDT}) where {ARDT,BDDT,CDT}
#function *(A::joAbstractLinearOperator{CDT,ARDT},B::joSAgatheringLinearOperator{BDDT,CDT,2}) where {ARDT,BDDT,CDT}
# in joAbstractDAparallelLinearOperator/base_functions.jl
#function *(A::joDAdistributingLinearOperator{CDT,ARDT,2},B::joAbstractLinearOperator{BDDT,CDT}) where {ARDT,BDDT,CDT}
#function *(A::joAbstractLinearOperator{CDT,ARDT},B::joDAgatheringLinearOperator{BDDT,CDT,2}) where {ARDT,BDDT,CDT}

# *(jo,mvec)
function *(A::joLinearOperator{ADDT,ARDT},mv::LocalMatrix{mvDT}) where {ADDT,ARDT,mvDT<:Number}
    A.n == size(mv,1) || throw(joLinearOperatorException("shape mismatch"))
    jo_check_type_match(ADDT,mvDT,join(["DDT for *(jo,mvec):",A.name,typeof(A),mvDT]," / "))
    MV = A.fop(mv)
    jo_check_type_match(ARDT,eltype(MV),join(["RDT from *(jo,mvec):",A.name,typeof(A),eltype(MV)]," / "))
    return MV
end
function *(A::joMatrix{ADDT,ARDT},mv::LocalMatrix{mvDT}) where {ADDT,ARDT,mvDT<:Number}
    A.n == size(mv,1) || throw(joMatrixException("shape mismatch"))
    jo_check_type_match(ADDT,mvDT,join(["DDT for *(jo,mvec):",A.name,typeof(A),mvDT]," / "))
    MV = A.fop(mv)
    jo_check_type_match(ARDT,eltype(MV),join(["RDT from *(jo,mvec):",A.name,typeof(A),eltype(MV)]," / "))
    return MV
end
function *(A::joLinearFunction{ADDT,ARDT},mv::LocalMatrix{mvDT}) where {ADDT,ARDT,mvDT<:Number}
    A.n == size(mv,1) || throw(joLinearFunctionException("shape mismatch"))
    jo_check_type_match(ADDT,mvDT,join(["DDT for *(jo,mvec):",A.name,typeof(A),mvDT]," / "))
    if A.fMVok
        MV=A.fop(mv)
        jo_check_type_match(ARDT,eltype(MV),join(["RDT from *(jo,mvec):",A.name,typeof(A),eltype(MV)]," / "))
    else
        MV=Matrix{ARDT}(undef,A.m,size(mv,2))
        for i=1:size(mv,2)
            V=A.fop(mv[:,i])
            i==1 && jo_check_type_match(ARDT,eltype(V),join(["RDT from *(jo,mvec):",A.name,typeof(A),eltype(V)]," / "))
            MV[:,i]=V
        end
    end
    return MV
end

# *(mvec,jo)

# *(jo,vec)
function *(A::joLinearOperator{ADDT,ARDT},v::LocalVector{vDT}) where {ADDT,ARDT,vDT<:Number}
    A.n == size(v,1) || throw(joLinearOperatorException("shape mismatch"))
    jo_check_type_match(ADDT,vDT,join(["DDT for *(jo,vec):",A.name,typeof(A),vDT]," / "))
    V=A.fop(v)
    jo_check_type_match(ARDT,eltype(V),join(["RDT from *(jo,vec):",A.name,typeof(A),eltype(V)]," / "))
    return V
end
function *(A::joMatrix{ADDT,ARDT},v::LocalVector{vDT}) where {ADDT,ARDT,vDT<:Number}
    A.n == size(v,1) || throw(joMatrixException("shape mismatch"))
    jo_check_type_match(ADDT,vDT,join(["DDT for *(jo,vec):",A.name,typeof(A),vDT]," / "))
    V=A.fop(v)
    jo_check_type_match(ARDT,eltype(V),join(["RDT from *(jo,vec):",A.name,typeof(A),eltype(V)]," / "))
    return V
end
function *(A::joLinearFunction{ADDT,ARDT},v::LocalVector{vDT}) where {ADDT,ARDT,vDT<:Number}
    A.n == size(v,1) || throw(joLinearFunctionException("shape mismatch"))
    jo_check_type_match(ADDT,vDT,join(["DDT for *(jo,vec):",A.name,typeof(A),vDT]," / "))
    V=A.fop(v)
    jo_check_type_match(ARDT,eltype(V),join(["RDT from *(jo,vec):",A.name,typeof(A),eltype(V)]," / "))
    return V
end

# *(vec,jo)

# *(num,jo)
function *(a::Number,A::joAbstractLinearOperator{ADDT,ARDT}) where {ADDT,ARDT}
    return joLinearOperator{ADDT,ARDT}("(N*"*A.name*")",A.m,A.n,
        v1->jo_convert(ARDT,a*(A*v1),false),
        v2->jo_convert(ADDT,a*(transpose(A)*v2),false),
        v3->jo_convert(ADDT,conj(a)*(adjoint(A)*v3),false),
        v4->jo_convert(ARDT,conj(a)*(conj(A)*v4),false),
        @joNF, @joNF, @joNF, @joNF
        )
end
function *(a::joNumber{ADDT,ARDT},A::joAbstractLinearOperator{ADDT,ARDT}) where {ADDT,ARDT}
    return joLinearOperator{ADDT,ARDT}("(N*"*A.name*")",A.m,A.n,
        v1->jo_convert(ARDT,a.rdt*(A*v1),false),
        v2->jo_convert(ADDT,a.ddt*(transpose(A)*v2),false),
        v3->jo_convert(ADDT,conj(a.ddt)*(adjoint(A)*v3),false),
        v4->jo_convert(ARDT,conj(a.rdt)*(conj(A)*v4),false),
        @joNF, @joNF, @joNF, @joNF
        )
end

# *(jo,num)
*(A::joAbstractLinearOperator{ADDT,ARDT},a::Number) where {ADDT,ARDT} = a*A
*(A::joAbstractLinearOperator{ADDT,ARDT},a::joNumber{ADDT,ARDT}) where {ADDT,ARDT} = a*A

############################################################
## overloaded Base \(...jo...)

# \(jo,jo)

# \(jo,mvec)
function \(A::joLinearOperator{ADDT,ARDT},mv::LocalMatrix{mvDT}) where {ADDT,ARDT,mvDT<:Number}
    A.m == size(mv,1) || throw(joLinearOperatorException("shape mismatch"))
    jo_check_type_match(ARDT,mvDT,join(["RDT for \\(jo,mvec):",A.name,typeof(A),mvDT]," / "))
    if hasinverse(A)
        MV=get(A.iop)(mv)
        jo_check_type_match(ADDT,eltype(MV),join(["DDT from \\(jo,mvec):",A.name,typeof(A),eltype(MV)]," / "))
    elseif issquare(A)
        MV=Matrix{ADDT}(undef,A.n,size(mv,2))
        for i=1:size(mv,2)
            V=jo_convert(ADDT,jo_iterative_solver4square(A,mv[:,i]))
            i==1 && jo_check_type_match(ADDT,eltype(V),join(["DDT from \\(jo,mvec):",A.name,typeof(A),eltype(V)]," / "))
            MV[:,i]=V
        end
    elseif (istall(A) && !isnull(jo_iterative_solver4tall))
        MV=Matrix{ADDT}(undef,A.n,size(mv,2))
        for i=1:size(mv,2)
            V=jo_convert(ADDT,jo_iterative_solver4tall(A,mv[:,i]))
            i==1 && jo_check_type_match(ADDT,eltype(V),join(["DDT from \\(jo,mvec):",A.name,typeof(A),eltype(V)]," / "))
            MV[:,i]=V
        end
    elseif (iswide(A) && !isnull(jo_iterative_solver4wide))
        MV=Matrix{ADDT}(undef,A.n,size(mv,2))
        for i=1:size(mv,2)
            V=jo_convert(ADDT,jo_iterative_solver4wide(A,mv[:,i]))
            i==1 && jo_check_type_match(ADDT,eltype(V),join(["DDT from \\(jo,mvec):",A.name,typeof(A),eltype(V)]," / "))
            MV[:,i]=V
        end
    else
        throw(joLinearOperatorException("\\(jo,MultiVector) not supplied"))
    end
    return MV
end
function \(A::joMatrix{ADDT,ARDT},mv::LocalMatrix{mvDT}) where {ADDT,ARDT,mvDT<:Number}
    A.m == size(mv,1) || throw(joMatrixException("shape mismatch"))
    jo_check_type_match(ARDT,mvDT,join(["RDT for *(jo,mvec):",A.name,typeof(A),mvDT]," / "))
    if hasinverse(A)
        MV=get(A.iop)(mv)
        jo_check_type_match(ADDT,eltype(MV),join(["DDT from *(jo,mvec):",A.name,typeof(A),eltype(MV)]," / "))
    else
        throw(joMatrixException("\\(jo,Vector) not supplied"))
    end
    return MV
end
function \(A::joLinearFunction{ADDT,ARDT},mv::LocalMatrix{mvDT}) where {ADDT,ARDT,mvDT<:Number}
    A.m == size(mv,1) || throw(joLinearFunctionException("shape mismatch"))
    jo_check_type_match(ARDT,mvDT,join(["RDT for \\(jo,mvec):",A.name,typeof(A),mvDT]," / "))
    if hasinverse(A)
        if A.iMVok
            MV = get(A.iop)(mv)
        else
            MV=Matrix{ADDT}(undef,A.n,size(mv,2))
            for i=1:size(mv,2)
                V=get(A.iop)(mv[:,i])
                i==1 && jo_check_type_match(ADDT,eltype(V),join(["DDT from \\(jo,mvec):",A.name,typeof(A),eltype(V)]," / "))
                MV[:,i]=V
            end
        end
    elseif issquare(A)
        MV=Matrix{ADDT}(undef,A.n,size(mv,2))
        for i=1:size(mv,2)
            V=jo_convert(ADDT,jo_iterative_solver4square(A,mv[:,i]))
            i==1 && jo_check_type_match(ADDT,eltype(V),join(["DDT from \\(jo,mvec):",A.name,typeof(A),eltype(V)]," / "))
            MV[:,i]=V
        end
    elseif (istall(A) && !isnull(jo_iterative_solver4tall))
        MV=Matrix{ADDT}(undef,A.n,size(mv,2))
        for i=1:size(mv,2)
            V=jo_convert(ADDT,jo_iterative_solver4tall(A,mv[:,i]))
            i==1 && jo_check_type_match(ADDT,eltype(V),join(["DDT from \\(jo,mvec):",A.name,typeof(A),eltype(V)]," / "))
            MV[:,i]=V
        end
    elseif (iswide(A) && !isnull(jo_iterative_solver4wide))
        MV=Matrix{ADDT}(undef,A.n,size(mv,2))
        for i=1:size(mv,2)
            V=jo_convert(ADDT,jo_iterative_solver4wide(A,mv[:,i]))
            i==1 && jo_check_type_match(ADDT,eltype(V),join(["DDT from \\(jo,mvec):",A.name,typeof(A),eltype(V)]," / "))
            MV[:,i]=V
        end
    else
        throw(joLinearFunctionException("\\(jo,MultiVector) not supplied"))
    end
    return MV
end

# \(mvec,jo)

# \(jo,vec)
function \(A::joLinearOperator{ADDT,ARDT},v::LocalVector{vDT}) where {ADDT,ARDT,vDT<:Number}
    A.m == size(v,1) || throw(joLinearOperatorException("shape mismatch"))
    jo_check_type_match(ARDT,vDT,join(["RDT for \\(jo,vec):",A.name,typeof(A),vDT]," / "))
    if hasinverse(A)
        V=get(A.iop)(v)
        jo_check_type_match(ADDT,eltype(V),join(["DDT from \\(jo,vec):",A.name,typeof(A),eltype(V)]," / "))
    elseif issquare(A)
        V=jo_convert(ADDT,jo_iterative_solver4square(A,v))
    elseif (istall(A) && !isnull(jo_iterative_solver4tall))
        V=jo_convert(ADDT,jo_iterative_solver4tall(A,v))
    elseif (iswide(A) && !isnull(jo_iterative_solver4wide))
        V=jo_convert(ADDT,jo_iterative_solver4wide(A,v))
    else
        throw(joLinearOperatorException("\\(jo,Vector) not supplied"))
    end
    return V
end
function \(A::joMatrix{ADDT,ARDT},v::LocalVector{vDT}) where {ADDT,ARDT,vDT<:Number}
    A.m == size(v,1) || throw(joMatrixException("shape mismatch"))
    jo_check_type_match(ARDT,vDT,join(["RDT for *(jo,vec):",A.name,typeof(A),vDT]," / "))
    if hasinverse(A)
        V=get(A.iop)(v)
        jo_check_type_match(ADDT,eltype(V),join(["DDT from *(jo,vec):",A.name,typeof(A),eltype(V)]," / "))
    else
        throw(joMatrixException("\\(jo,Vector) not supplied"))
    end
    return V
end
function \(A::joLinearFunction{ADDT,ARDT},v::LocalVector{vDT}) where {ADDT,ARDT,vDT<:Number}
    A.m == size(v,1) || throw(joLinearFunctionException("shape mismatch"))
    jo_check_type_match(ARDT,vDT,join(["RDT for \\(jo,vec):",A.name,typeof(A),vDT]," / "))
    if hasinverse(A)
        V=get(A.iop)(v)
        jo_check_type_match(ADDT,eltype(V),join(["DDT from \\(jo,vec):",A.name,typeof(A),eltype(V)]," / "))
    elseif issquare(A)
        V=jo_convert(ADDT,jo_iterative_solver4square(A,v))
    elseif (istall(A) && !isnull(jo_iterative_solver4tall))
        V=jo_convert(ADDT,jo_iterative_solver4tall(A,v))
    elseif (iswide(A) && !isnull(jo_iterative_solver4wide))
        V=jo_convert(ADDT,jo_iterative_solver4wide(A,v))
    else
        throw(joLinearFunctionException("\\(jo,Vector) not supplied"))
    end
    return V
end

# \(vec,jo)

# \(num,jo)

# \(jo,num)
#\{ADDT,ARDT}(A::joAbstractLinearOperator{ADDT,ARDT},a::Number) = inv(a)*A
#\{ADDT,ARDT}(A::joAbstractLinearOperator{ADDT,ARDT},a::joNumber{ADDT,ARDT}) = inv(a)*A

############################################################
## overloaded Base +(...jo...)

# +(jo)
+(A::joAbstractLinearOperator) = A

# +(jo,jo)
function +(A::joAbstractLinearOperator{DDT,RDT},B::joAbstractLinearOperator{DDT,RDT}) where {DDT,RDT}
    size(A) == size(B) || throw(joAbstractLinearOperatorException("shape mismatch"))
    return joLinearOperator{DDT,RDT}("("*A.name*"+"*B.name*")",size(A,1),size(B,2),
        v1->A*v1+B*v1,
        v2->transpose(A)*v2+transpose(B)*v2,
        v3->adjoint(A)*v3+adjoint(B)*v3,
        v4->conj(A)*v4+conj(B)*v4,
        @joNF, @joNF, @joNF, @joNF
        )
end

# +(jo,mvec)

# +(mvec,jo)

# +(jo,vec)

# +(vec,jo)

# +(jo,num)
function +(A::joAbstractLinearOperator{ADDT,ARDT},b::Number) where {ADDT,ARDT}
    return joLinearOperator{ADDT,ARDT}("("*A.name*"+N)",size(A,1),size(A,2),
        v1->A*v1+joConstants(A.m,A.n,b;DDT=ADDT,RDT=ARDT)*v1,
        v2->transpose(A)*v2+joConstants(A.n,A.m,b;DDT=ARDT,RDT=ADDT)*v2,
        v3->adjoint(A)*v3+joConstants(A.n,A.m,conj(b);DDT=ARDT,RDT=ADDT)*v3,
        v4->conj(A)*v4+joConstants(A.m,A.n,conj(b);DDT=ADDT,RDT=ARDT)*v4,
        @joNF, @joNF, @joNF, @joNF
        )
end
function +(A::joAbstractLinearOperator{ADDT,ARDT},b::joNumber{ADDT,ARDT}) where {ADDT,ARDT}
    return joLinearOperator{ADDT,ARDT}("("*A.name*"+N)",size(A,1),size(A,2),
        v1->A*v1+joConstants(A.m,A.n,b.rdt;DDT=ADDT,RDT=ARDT)*v1,
        v2->transpose(A)*v2+joConstants(A.n,A.m,b.ddt;DDT=ARDT,RDT=ADDT)*v2,
        v3->adjoint(A)*v3+joConstants(A.n,A.m,conj(b.ddt);DDT=ARDT,RDT=ADDT)*v3,
        v4->conj(A)*v4+joConstants(A.m,A.n,conj(b.rdt);DDT=ADDT,RDT=ARDT)*v4,
        @joNF, @joNF, @joNF, @joNF
        )
end

# +(num,jo)
+(b::Number,A::joAbstractLinearOperator{ADDT,ARDT}) where {ADDT,ARDT} = A+b
+(b::joNumber{ADDT,ARDT},A::joAbstractLinearOperator{ADDT,ARDT}) where {ADDT,ARDT} = A+b

############################################################
## overloaded Base -(...jo...)

# -(jo)
-(A::joLinearOperator{DDT,RDT}) where {DDT,RDT} =
    joLinearOperator{DDT,RDT}("(-"*A.name*")",A.m,A.n,
        v1->-A.fop(v1),
        v2->-get(A.fop_T)(v2),
        v3->-get(A.fop_A)(v3),
        v4->-get(A.fop_C)(v4),
        v5->-get(A.iop)(v5),
        v6->-get(A.iop_T)(v6),
        v7->-get(A.iop_A)(v7),
        v8->-get(A.iop_C)(v8)
        )
-(A::joMatrix{DDT,RDT}) where {DDT,RDT} =
    joMatrix{DDT,RDT}("(-"*A.name*")",A.m,A.n,
        v1->-A.fop(v1),
        v2->-A.fop_T(v2),
        v3->-A.fop_A(v3),
        v4->-A.fop_C(v4),
        v5->-get(A.iop)(v5),
        v6->-get(A.iop_T)(v6),
        v7->-get(A.iop_A)(v7),
        v8->-get(A.iop_C)(v8)
        )
-(A::joLinearFunction{DDT,RDT}) where {DDT,RDT} =
    joLinearFunction{DDT,RDT}("(-"*A.name*")",A.m,A.n,
        v1->-A.fop(v1),
        v2->-get(A.fop_T)(v2),
        v3->-get(A.fop_A)(v3),
        v4->-get(A.fop_C)(v4),
        A.fMVok,
        v5->-get(A.iop)(v5),
        v6->-get(A.iop_T)(v6),
        v7->-get(A.iop_A)(v7),
        v8->-get(A.iop_C)(v8),
        A.iMVok
        )

# -(jo,jo)
-(A::joAbstractLinearOperator,B::joAbstractLinearOperator) = A+(-B)

# -(jo,mvec)

# -(mvec,jo)

# -(jo,vec)

# -(vec,jo)

# -(jo,num)
-(A::joAbstractLinearOperator,b::Number) = A+(-b)
-(A::joAbstractLinearOperator,b::joNumber) = A+(-b)

# -(num,jo)
-(b::Number,A::joAbstractLinearOperator) = -A+b
-(b::joNumber,A::joAbstractLinearOperator) = -A+b

############################################################
## overloaded Base .*(...jo...)

# .*(num,jo)
Base.Broadcast.broadcasted(::typeof(*),a::Number,A::joAbstractLinearOperator) = a*A
Base.Broadcast.broadcasted(::typeof(*),a::joNumber,A::joAbstractLinearOperator) = a*A

# .*(jo,num)
Base.Broadcast.broadcasted(::typeof(*),A::joAbstractLinearOperator,a::Number) = a*A
Base.Broadcast.broadcasted(::typeof(*),A::joAbstractLinearOperator,a::joNumber) = a*A

############################################################
## overloaded Base .\(...jo...)

############################################################
## overloaded Base .+(...jo...)

# .+(jo,num)
Base.Broadcast.broadcasted(::typeof(+),A::joAbstractLinearOperator,b::Number) = A+b
Base.Broadcast.broadcasted(::typeof(+),A::joAbstractLinearOperator,b::joNumber) = A+b

# .+(num,jo)
Base.Broadcast.broadcasted(::typeof(+),b::Number,A::joAbstractLinearOperator) = A+b
Base.Broadcast.broadcasted(::typeof(+),b::joNumber,A::joAbstractLinearOperator) = A+b

############################################################
## overloaded Base .-(...jo...)

# .-(jo,num)
Base.Broadcast.broadcasted(::typeof(-),A::joAbstractLinearOperator,b::Number) = A+(-b)
Base.Broadcast.broadcasted(::typeof(-),A::joAbstractLinearOperator,b::joNumber) = A+(-b)

# .-(num,jo)
Base.Broadcast.broadcasted(::typeof(-),b::Number,A::joAbstractLinearOperator) = -A+b
Base.Broadcast.broadcasted(::typeof(-),b::joNumber,A::joAbstractLinearOperator) = -A+b

############################################################
## overloaded Base block methods

# hcat(...jo...)
hcat(ops::joAbstractLinearOperator...) = joDict(ops...)

# vcat(...jo...)
vcat(ops::joAbstractLinearOperator...) = joStack(ops...)

# hvcat(...jo...)
hvcat(rows::Tuple{Vararg{Int}}, ops::joAbstractLinearOperator...) = joBlock(collect(rows),ops...)

############################################################
## overloaded LinearAlgebra functions

# mul!(...,jo,...)
mul!(y::LocalVector{RDT},A::joAbstractLinearOperator{DDT,RDT},x::LocalVector{DDT}) where {DDT,RDT} = y[:] = A * x
mul!(y::LocalMatrix{RDT},A::joAbstractLinearOperator{DDT,RDT},x::LocalMatrix{DDT}) where {DDT,RDT} = y[:,:] = A * x

# ldiv!(...,jo,...)
ldiv!(y::LocalVector{DDT},A::joAbstractLinearOperator{DDT,RDT},x::LocalVector{RDT}) where {DDT,RDT} = y[:] = A \ x
ldiv!(y::LocalMatrix{DDT},A::joAbstractLinearOperator{DDT,RDT},x::LocalMatrix{RDT}) where {DDT,RDT} = y[:,:] = A \ x


using Test

@timed_testset "trained_weights/" begin
    include("trained_weights/MNIST.n1.jl")
    include("trained_weights/MNIST.WK17a_linf0.1_authors.jl")
end


function setup_data()
    words = h5read(h5file, "s")
    worddict = Dict(words[i] => i for i=1:length(words))
    chardict = Dict{Char,Int}()

    traindoc = CoNLL.read(".data/wsj_00-18.conll")
    testdoc = CoNLL.read(".data/wsj_22-24.conll")
    info("# sentences of train doc: $(length(traindoc))")
    info("# sentences of test doc: $(length(testdoc))")

    traindata = setup_data(traindoc, worddict, chardict)
    testdata = setup_data(testdoc, worddict, chardict)
    info("# words: $(length(worddict))")
    info("# chars: $(length(chardict))")
    traindata, testdata, worddict, chardict
end

function setup_data(doc::Vector, worddict, chardict)
    data = []
    unkword = worddict["UNKNOWN"]
    root = worddict["PADDING"]
    for sent in doc
        w = Int[root]
        cs = Var[Var([0])]
        h = Int[0]
        for items in sent
            word = items[2]
            word0 = replace(word, r"[0-9]", '0')
            wordid = get(worddict, lowercase(word0), unkword)

            chars = Vector{Char}(word0)
            charids = map(c -> get!(chardict,c,length(chardict)+1), chars)

            head = parse(Int, items[7]) + 1
            push!(w, wordid)
            push!(cs, Var(charids))
            push!(h, head)
        end
        push!(data, (Var(w),cs,Var(h)))
    end
    data
end


function read2(f, block::PlainVariableBlockHeader{T}; fn=_->true, req_pts, batch_size=(req_pts[end][2]-req_pts[end][1])) where T
    if length(block.npts) != length(req_pts)
        ArgumentError("dimension mismatch")
    end

    block_size = prod(block.npts[1:end-1])
    z1, z2 = req_pts[end]
    seek(f, block.base_header.data_location + sizeof(T)*block_size*(z1-1))
    n, m = divrem(z2-z1, batch_size)
    vals = T[]
    idxs = Array{CartesianIndex{length(req_pts)},1}()
    result = Array{T, 1}(undef, block_size*batch_size)
    for i in eachindex(n)
        read_batch!(f, result, block, fn, req_pts, batch_size, block_size, vals, idxs)
    end
    if m != 0
        result = Array{T, 1}(undef, block_size*m)
        read_batch!(f, result, block, fn, req_pts, m, block_size, vals, idxs)
    end

    for (i,v) in enumerate(req_pts)
        idxs .+= Ref(cart_idx(req_pts[i][1], Val(i)))
    end

    vals, idxs
end

function read_batch!(f, result, block, fn, req_pts, batch_size, block_size, vals, idxs)
    read!(f, result)
    result = trim_block(result, block.npts, batch_size, req_pts, Val(length(req_pts)))
    idx = findall(fn, result)
    append!(vals, result[idx])
    append!(idxs, idx)
end

function trim_block(result, npts, batch_size, req_pts, ::Val{1})
    result
end

function trim_block(result, npts, batch_size, req_pts, ::Val{2})
    x1, x2 = req_pts[1]
    result = reshape(result, (npts[1], batch_size))
    result[x1:x2, :]
end

function trim_block(result, npts, batch_size, req_pts, ::Val{3})
    x1, x2 = req_pts[1]
    y1, y2 = req_pts[2]
    result = reshape(result, (npts[1], npts[2], batch_size))
    result[x1:x2, y1:y2, :]
end

cart_idx(pt, ::Val{1}) = CartesianIndex(pt-1, 0, 0)
cart_idx(pt, ::Val{2}) = CartesianIndex(0, pt-1, 0)
cart_idx(pt, ::Val{3}) = CartesianIndex(0, 0, pt-1)


const ThreeAgentPos = SVector{6,Int}

const syma_ego = [[:up,], [:down,], [:left,], [:right,]]
const syma_adv = [[:up, :up], [:down, :up], [:left, :up], [:right, :up],
                [:up, :down], [:down, :down], [:left, :down], [:right, :down],
                [:up, :left], [:down, :left], [:left, :left], [:right, :left],
                [:up, :right], [:down, :right], [:left, :right], [:right, :right]]
const aind_ego = Dict(syma_ego[i] => i for i=1:length(syma_ego))
const aind_adv = Dict(syma_adv[i] => i for i=1:length(syma_adv))
const S2 = ThreeAgentPos
const A2 = Array{Symbol}

# Gridworld with adversary
@with_kw mutable struct DoubleAdversarialGridworldMDP <: MDP{S2, A2}
    size::Tuple{Int, Int} = (10,10)
    rewards::Dict{GWPos, Float64} = Dict()
    walls::Vector{GWPos} = []
    tprob::Float64 = 0.7
    discount::Float64 = 0.95
    agent_gets_action = :ego # :ego or :adversary
    ego_policy = (s, rng::AbstractRNG = Random.GLOBAL_RNG) -> rand(rng, syma_ego)
    adversary_policy = (s, rng::AbstractRNG = Random.GLOBAL_RNG) -> rand(rng, syma_adv)
    failure_penalty = 1.
end

valid_pos(mdp::DoubleAdversarialGridworldMDP, pos::GWPos) = !(pos in mdp.walls || any((pos .> mdp.size) .| (pos .< GWPos(1,1))))

function POMDPs.states(mdp::DoubleAdversarialGridworldMDP)
    lengths = (mdp.size[1], mdp.size[2], mdp.size[1], mdp.size[2], mdp.size[1], mdp.size[2])
    ss = S2[]
    for ijk in CartesianIndices(lengths)
        s = S2(ijk.I...)
        valid_pos(mdp, ego_pos(s)) && valid_pos(mdp, adversary1_pos(s)) && valid_pos(mdp, adversary2_pos(s))  &&  push!(ss, s)
    end
    push!(ss, S2(-1,-1,-1,-1,-1,-1))
    ss
end

function random_valid_pos(mdp::DoubleAdversarialGridworldMDP, rng::AbstractRNG = Random.GLOBAL_RNG, exclude = [], max_trials = 1000)
    trial = 0
    while trial < max_trials
        pos = GWPos(rand(rng, 1:mdp.size[1]), rand(rng, 1:mdp.size[2]))
        if valid_pos(mdp, pos) && !(haskey(mdp.rewards, pos) || pos in exclude)
            return pos
        end
        trial += 1
    end
 end

function POMDPs.initialstate(mdp::DoubleAdversarialGridworldMDP, rng::AbstractRNG = Random.GLOBAL_RNG)
    ego = random_valid_pos(mdp, rng)
    adversary1 = random_valid_pos(mdp, rng, [ego])
    adversary2 = random_valid_pos(mdp, rng, [ego, adversary1])
    Deterministic(S2(ego..., adversary1..., adversary2...))
end

POMDPs.actions(mdp::DoubleAdversarialGridworldMDP) = mdp.agent_gets_action == :ego ? syma_ego : syma_adv
POMDPs.actionindex(mdp::DoubleAdversarialGridworldMDP, a::A2) = mdp.agent_gets_action == :ego ? aind_ego[a] : aind_adv[a]

ego_pos(s::S2) = GWPos(s[1], s[2])
adversary1_pos(s::S2) = GWPos(s[3], s[4])
adversary2_pos(s::S2) = GWPos(s[5], s[6])
agents_overlap(s::S2) = (ego_pos(s) == adversary1_pos(s)) && (ego_pos(s) == adversary2_pos(s))

POMDPs.isterminal(mdp::DoubleAdversarialGridworldMDP, s::S2) = any(s .< 0)
POMDPs.discount(mdp::DoubleAdversarialGridworldMDP) = mdp.discount

# Returns a sample next state and reward
function POMDPs.gen(mdp::DoubleAdversarialGridworldMDP, s::S2, a::A2, rng::AbstractRNG = Random.GLOBAL_RNG)
    if haskey(mdp.rewards, ego_pos(s)) || agents_overlap(s) || isterminal(mdp, s)
        return (sp = S2(-1,-1,-1,-1,-1,-1), r = reward(mdp, s))
    else
        # If this MDP controls the agent then use the adversary policy for the adversary
        # Do the opposite if the adversary is being controlled by the MDP action
        if mdp.agent_gets_action == :ego
            rdir = (rand(rng) < mdp.tprob) ? dir[a[1]] : dir[rand(rng, syma[a[1] .!= syma])]
            new_ego = ego_pos(s) + rdir
            a_adv = mdp.adversary_policy(s, rng)
            new_adv1 = adversary1_pos(s) + dir[a_adv[1]]
            new_adv2 = adversary2_pos(s) + dir[a_adv[2]]
        else
            rdir1 = (rand(rng) < mdp.tprob) ? dir[a[1]] : dir[rand(rng, syma[a[1] .!= syma])]
            rdir2 = (rand(rng) < mdp.tprob) ? dir[a[2]] : dir[rand(rng, syma[a[2] .!= syma])]
            new_ego = ego_pos(s) + dir[mdp.ego_policy(s, rng)[1]]
            new_adv1 = adversary1_pos(s) + rdir1
            new_adv2 = adversary2_pos(s) + rdir2
        end

        # Make sure the moves are in bound and not hitting a wall
        new_ego = valid_pos(mdp, new_ego) ? new_ego : ego_pos(s)
        new_adv1 = valid_pos(mdp, new_adv1) ? new_adv1 : adversary1_pos(s)
        new_adv2 = valid_pos(mdp, new_adv2) ? new_adv2 : adversary2_pos(s)

        return (sp = S2(new_ego..., new_adv1..., new_adv2...), r = reward(mdp, s))
    end
end

# Returns the reward for the provided state
function POMDPs.reward(mdp::DoubleAdversarialGridworldMDP, s::S2)
    isterminal(mdp, s) && return 0
    r = (get(mdp.rewards, ego_pos(s), 0.0) - mdp.failure_penalty*agents_overlap(s)) /  mdp.failure_penalty
    mdp.agent_gets_action == :ego ? r : -r
end

function tocolor(mdp::DoubleAdversarialGridworldMDP, r::Float64)
    maxr = maximum(values(mdp.rewards))
    minr = -maxr
    frac = (r-minr)/(maxr-minr)
    return get(ColorSchemes.redgreensplit, frac)
end

# Renders the mdp
function POMDPModelTools.render(mdp::DoubleAdversarialGridworldMDP, s::S2)
    nx, ny = mdp.size
    cells = []
    for x in 1:nx, y in 1:ny
        pos = GWPos(x,y)
        reward_index = findfirst([pos] .== keys(mdp.rewards))
        wall_index = findfirst([pos] .== mdp.walls)
        ctx = context((x-1)/nx, (ny-y)/ny, 1/nx, 1/ny)
        color = "white"
        if !isnothing(reward_index)
            color = tocolor(mdp, get(mdp.rewards, pos, 0))
        elseif !isnothing(wall_index)
            color = "black"
        end
        cell = compose(ctx, Compose.rectangle(), fill(color))
        push!(cells, cell)
    end
    grid = compose(context(), Compose.stroke("gray"), cells...)
    outline = compose(context(), Compose.rectangle())

    ego = nothing
    x,y = ego_pos(s)
    if all([x,y] .> 0)
        agent_ctx = context((x-1)/nx, (ny-y)/ny, 1/nx, 1/ny)
        ego = compose(agent_ctx, Compose.circle(0.5, 0.5, 0.4), Compose.stroke("black"), fill("blue"))
    end

    adversary1 = nothing
    x,y = adversary1_pos(s)
    if all([x,y] .> 0)
        agent_ctx = context((x-1)/nx, (ny-y)/ny, 1/nx, 1/ny)
        adversary1 = compose(agent_ctx, Compose.circle(0.5, 0.5, 0.4), Compose.stroke("black"), fill("orange"))
    end

    adversary2 = nothing
    x,y = adversary2_pos(s)
    if all([x,y] .> 0)
        agent_ctx = context((x-1)/nx, (ny-y)/ny, 1/nx, 1/ny)
        adversary2 = compose(agent_ctx, Compose.circle(0.5, 0.5, 0.4), Compose.stroke("black"), fill("red"))
    end
    agents_comp = compose(context(), ego, adversary1, adversary2)

    sz = min(w, h)
    return compose(context((w-sz)/2, (h-sz)/2, sz, sz), agents_comp, grid, outline)
end


module KoalaFlux

# new:
export FluxRegressor
export CategoricalEmbedder

# needed in this module:
using LinearAlgebra
using Random
import Koala: Regressor, Classifier, softwarn, BaseType, rms
import Koala: supports_ensembling, Machine, SupervisedMachine
import DataFrames: eltypes, AbstractDataFrame, head
import KoalaTransforms: Transformer, TransformerMachine, ChainTransformer
import KoalaTransforms: MakeCategoricalsIntTransformer, Standardizer
import KoalaTransforms: RegressionTargetTransformer
import KoalaTransforms: ToIntScheme, ToIntTransformer
import Flux

# to be extended (but not explicitly rexported):
import Koala: setup, fit, predict
import Koala: default_transformer_X, default_transformer_y, clean!
import Koala: transform, inverse_transform

# Notes on internals of this module:

# 1. What Flux.jl calls a "model" is not a model in the Koala.jl
# sense. Since we require "models" in the Flux.jl sense to be wrapped
# in `Flux.Chain` objects, we use the word "chain" in place of "Flux
# model" below. Each "chain", then, is a neural network (including
# weights).

# 2. Because a `FluxRegressor` machine allows for categorical input
# features, a FluxRegressor machine includes a chain (neural network)
# *and* a vector of matrices ("tracked" in the sense of Flux.jl)
# encoding the embedding, called "embedding" below. Thus a Koala
# `predictor` consists of a (chain, embedding) tuple.

## HELPERS

A ⊂ B = issubset(A, B)

Base.isempty(chain::Flux.Chain) = isempty(chain.layers)

default_formula(n_classes; threshold=10) =
    min(round(Int, sqrt(n_classes)), threshold)

function default_creator(n_inputs)
    n_hidden = round(Int, sqrt(n_inputs))
    return Flux.Chain(Flux.Dense(n_inputs, n_hidden, Flux.σ),
                      Flux.Dense(n_hidden, 1))
end

function ordinal_categorical_features(X::AbstractDataFrame)

    features = names(X)
    types = eltypes(X)
    ordinal_features = Symbol[]
    categorical_features = Symbol[]
    for j in eachindex(features)
        if types[j] <: AbstractFloat
            push!(ordinal_features, features[j])
        else
            push!(categorical_features, features[j])
        end
    end

    return ordinal_features, categorical_features

end

## MODEL TYPE DEFINITIONS

# type of matrix for embedding a single categorical feature:
EmbType = Flux.TrackedArray{Float64,2,Array{Float64,2}}
WeightsType = Union{Flux.TrackedArray{Float64,2,Array{Float64,2}},
                       Flux.TrackedArray{Float64,1,Array{Float64,1}}}

FluxPredictorType = Tuple{Vector{EmbType},Flux.Chain}

mutable struct FluxRegressor <: Regressor{FluxPredictorType}
    network_creator::Function
    dimension_formula::Function
    learning_rate::Float64
    inertia::Float64
    lambda::Float64
    alpha::Float64
    embedding0::Vector{EmbType}
    chain0::Flux.Chain
    n::Int
end

# keyword constructor:
function FluxRegressor(; network_creator=default_creator,
                       dimension_formula=default_formula,
                       learning_rate=0.001,
                       inertia=0.9, lambda = 1e-5, alpha=0.0,
                       embedding0=EmbType[], chain0=Flux.Chain(), n=20)
    model = FluxRegressor(network_creator, dimension_formula, learning_rate,
                          inertia, lambda, alpha, embedding0, chain0, n)
    softwarn(clean!(model))
    return model
end

function clean!(model::FluxRegressor)
    message = ""
    if model.alpha > 1 || model.alpha < 0
        message = message*"alpha must be in range [0, 1]. Resetting to 0.0. "
        model.alpha = 0.0
    end
    if model.inertia > 1 || model.inertia < 0
        message = message*"inertia must be in range [0, 1]. Resetting to 0.0. "
        model.inertia = 0.0
    end
    return message
end

## TRANSFORMERS

struct FluxInput <: BaseType
    ordinal::Matrix{Float64}
    categorical::Matrix{Int}
end

function Base.getindex(X::KoalaFlux.FluxInput, rows::AbstractVector{Int}, ::Colon)
    ordinal = isempty(X.ordinal) ? Array{Float64}(undef, 0, length(rows)) :
        getindex(X.ordinal, :, rows)
    categorical = isempty(X.categorical) ? Array{Int}(undef, 0, length(rows)) :
        getindex(X.categorical, :, rows)
    return FluxInput(ordinal, categorical)
end

struct FrameToFluxInputTransformerScheme <: BaseType
    make_categoricals_int_transformer_machine::TransformerMachine
    ordinal_features::Vector{Symbol}
    categorical_features::Vector{Symbol}
end

struct FrameToFluxInputTransformer <: Transformer
end

function fit(transformer::FrameToFluxInputTransformer, X::AbstractDataFrame,
             parallel, verbosity)

    make_categoricals_int_tansformer_machine =
        Machine(MakeCategoricalsIntTransformer(sorted=true), X)

    ordinal_features, categorical_features = ordinal_categorical_features(X)

    return FrameToFluxInputTransformerScheme(
        make_categoricals_int_tansformer_machine,
        ordinal_features, categorical_features)

end

function transform(transformer::FrameToFluxInputTransformer,
                   scheme, X::AbstractDataFrame)

    X = transform(scheme.make_categoricals_int_transformer_machine, X)
    ordinal_features, categorical_features = ordinal_categorical_features(X)

    # check X has same features as X used to fit:
    Set(ordinal_features) == Set(scheme.ordinal_features) &&
        Set(categorical_features) == Set(scheme.categorical_features) ||
        throw(DimensionMismatch("Missing or unexpected features encountered."))

    X_ordinal = X[scheme.ordinal_features]
    X_categorical = X[scheme.categorical_features]

    ordinal = isempty(X_ordinal) ? Array{Float64}(undef, size(X, 2), 0) :
        transpose(convert(Array, X_ordinal))
    categorical = isempty(X_categorical) ? Array{Int}(undef, size(X, 2), 0) :
        transpose(convert(Array, X_categorical))

    return FluxInput(ordinal, categorical)

end

default_transformer_X(model::FluxRegressor) = ChainTransformer(FrameToFluxInputTransformer(), Standardizer())
default_transformer_y(model::FluxRegressor) = RegressionTargetTransformer()

## TRAINING AND PREDICTION METHODS

mutable struct FluxCache <: BaseType
    X::FluxInput
    y::Vector{Float64}
    ordinal_features::Vector{Symbol}
    categorical_features::Vector{Symbol}
    make_categoricals_int_transformer_machine::TransformerMachine
    embedding::Vector{EmbType}
    chain::Flux.Chain
    FluxCache(X, y, ordinal_features, categorical_features, machine) =
        new(X, y, ordinal_features, categorical_features, machine)
end

setup(model::FluxRegressor, Xt, yt, scheme_X, parallel, verbosity) =
    FluxCache(Xt, yt, scheme_X[end].ordinal_features, scheme_X[end].categorical_features,
              scheme_X[end].make_categoricals_int_transformer_machine)

# function to concatenate vector of ordinal values with vector of
# embedded categorical values into single vector:
function x(embedding, ordinals, categoricals)
    embedded_categoricals = vcat([embedding[j][:,categoricals[j]]
                              for j in eachindex(embedding)]...)
    return vcat(ordinals, embedded_categoricals)
end

function fit(model::FluxRegressor, cache, add, parallel, verbosity; args...)

    if !add

        # create or load initial embedding matrices:
        if isempty(model.embedding0)
            cache.embedding = Array{EmbType}(undef, length(cache.categorical_features))
        else
            cache.embedding = model.embedding0
        end

        # print embedding details:
        schemes=cache.make_categoricals_int_transformer_machine.scheme.schemes
        dimensions = Int[]
        if verbosity >= 1
            @info "\n"
            @info "Categorical feature embeddings:"
            @info "  feature     \t| embedding dimension"
            @info "------------------|----------------------"
        end
        for j in eachindex(cache.categorical_features)
            n_levels = schemes[j].n_levels
            d = model.dimension_formula(n_levels)
            append!(dimensions, d)
            cache.embedding[j] = Flux.param(randn(d, n_levels))
            verbosity < 1 || @info "  $(cache.categorical_features[j]) \t| $d"
        end

        # create or load a post-embedding network:
        if isempty(model.chain0)
            n_inputs = length(cache.ordinal_features) + sum(dimensions)
            cache.chain = model.network_creator(n_inputs)
        else
            cache.chain = model.chain0
        end
        verbosity < 1 || @info "Flux network architecture: \n  $(cache.chain)"

    end

    Flux.testmode!(cache.chain, false)

    embed_params = cache.embedding
    chain_params = Flux.params(cache.chain)

    ntrain = length(cache.y)
    train = 1:ntrain

    verbosity < 1 || println()

    # Note: In our terminology a "momentum" is a "moving average"
    # estimate of gradient.

    # initialize vectors of momenta:
    embed_momenta = Array{Array{Float64,2}}(undef, length(embed_params))
    chain_momenta = Array{Union{Array{Float64,1},
                                Array{Float64,2}}}(undef, length(chain_params))
    for k in eachindex(embed_params)
        p = embed_params[k]
        embed_momenta[k] = zeros(size(p))
    end
    for k in eachindex(chain_params)
        p = chain_params[k]
        chain_momenta[k] = zeros(size(p))
    end

    η = model.learning_rate
    β = model.inertia

    for eon in 1:model.n
        verbosity < 1 || print("\rTraining eon number: $eon   ")
        scrambled_train = shuffle(train)
        for i in scrambled_train
            ordinals = cache.X.ordinal[:,i]
            categoricals = cache.X.categorical[:,i]
            input_vector = x(embed_params, ordinals, categoricals)
#            @show input_vector
            l = rms(cache.chain(input_vector), cache.y[i])
#            @show cache.y[i]
#            @show chain_params
#            @show cache.chain(input_vector)
#            @show l
            l2_penalty = sum(norm, chain_params)
#            @show l2_penalty
            l1_penalty = sum(x->norm(x, 1), chain_params)
#            @show l1_penalty
            if !isempty(embed_params)
                l2_penalty += sum(norm, embed_params)
                l1_penalty += sum(x->norm(x, 1), embed_params)
            end
            l += model.lambda*(model.alpha*l1_penalty + (1 - model.alpha)*l2_penalty)
            !isnan(l) || @error "A NaN loss encountered."
            # some loss functions are not differentiable at zero, so:
            if abs(l) > eps(Float64)
#                @show l
                Flux.back!(l)
                for k in eachindex(embed_params)
                    p = embed_params[k]
                    embed_momenta[k] =  β*embed_momenta[k] +
                        (1 - β)*Flux.Tracker.grad(p)
                    Flux.Tracker.update!(p, -η*embed_momenta[k])
                end
                for k in eachindex(chain_params)
                    p = chain_params[k]
                    chain_momenta[k] =  β*chain_momenta[k] +
                        (1 - β)*Flux.Tracker.grad(p)
                    Flux.Tracker.update!(p, -η*chain_momenta[k])
                end
            end
        end
    end

    report = Dict{Symbol,Any}()
    report[:embedding] = embed_params
    report[:chain] = cache.chain

    predictor = (embed_params, cache.chain)

    return predictor, report, cache

end

function predict(model::FluxRegressor, predictor::FluxPredictorType,
                 X, parallel, verbosity)

    embedding, chain = predictor
    Flux.testmode!(chain)
    npatterns = size(X.ordinal, 2)

    predictions = Array{Float64}(undef, npatterns)
    for i in 1:npatterns
        ordinals = X.ordinal[:,i]
        categoricals = X.categorical[:,i]
        input_vector = x(embedding, ordinals, categoricals)
        predictions[i] = chain(input_vector).data[1]
    end

    return predictions
end

## FEATURE EMBEDDING TRANSFORMS

# Here we define transforms to embed the categorical features in some
# DataFrame into the multidimensional continuous variables learned by a
# Flux neural network.

mutable struct CategoricalEmbedder <: Transformer
    flux_machine::SupervisedMachine{FluxPredictorType, FluxRegressor}
    features::Vector{Symbol}
end

CategoricalEmbedder(machine::SupervisedMachine{FluxPredictorType, FluxRegressor};
                    features = Symbol[]) = CategoricalEmbedder(machine, features)
CategoricalEmbedder() =
    @error "You must give `CategoricalEmbedder` a `FluxMachine`
         argument to instantiate it."

struct CategoricalEmbedderScheme <: BaseType
    features::Vector{Symbol}
    to_int_transformer::ToIntTransformer
    to_int_scheme_given_feature::Dict{Symbol,ToIntScheme}
    embedding_given_feature::Dict{Symbol,Matrix{Float64}}
end

function fit(transformer::CategoricalEmbedder,
             X::AbstractDataFrame, parallel, verbosity)

    _, categorical = ordinal_categorical_features(X)

    if isempty(transformer.features)
        features = categorical
    else
        features = transformer.features
    end

    flux_categoricals = transformer.flux_machine.scheme_X[end].categorical_features
    Set(features) ⊂ Set(flux_categoricals) ||
        throw(DimensionMismatch("Encountered categorical "*
                                "not appearing in `FluxRegressor` object."))

    to_int_scheme_given_feature = Dict{Symbol,ToIntScheme}()
    embedding_given_feature = Dict{Symbol,Matrix{Float64}}()

    machine = transformer.flux_machine
    scheme = machine.scheme_X[end]
    to_int_schemes =
        scheme.make_categoricals_int_transformer_machine.scheme.schemes
    to_int_transformer =
        scheme.make_categoricals_int_transformer_machine.scheme.to_int_transformer

    for j in eachindex(flux_categoricals)
        if flux_categoricals[j] in features
            ftr = flux_categoricals[j]
            to_int_scheme_given_feature[ftr] =
                to_int_schemes[j]
            embedding_given_feature[ftr] =
                machine.report[:embedding][j].data
        end
    end

    return CategoricalEmbedderScheme(features, to_int_transformer,
                                     to_int_scheme_given_feature,
                                     embedding_given_feature)
end

# Will use the following to get the `kth` component of the emedding
# vector for raw value `x` of feature `ftr`. Will work just as well if
# `x` is replaced with a vector of categorical values, returning a
# Float64 vector.
function transform(scheme::CategoricalEmbedderScheme, ftr::Symbol, k::Int, x)
    x_as_int = transform(scheme.to_int_transformer,
                         scheme.to_int_scheme_given_feature[ftr], x)
    return scheme.embedding_given_feature[ftr][k,x_as_int]
end

function transform(transformer::CategoricalEmbedder, scheme, X::AbstractDataFrame)

    _, categoricals = ordinal_categorical_features(X)
    without_embedding = setdiff(Set(categoricals), Set(scheme.features))
    if !isempty(without_embedding)
        @warn "No embeddings found for the following categorical features: "
        for ftr in without_embedding
            println(ftr)
        end
    end

    Xout = X[:,:] # copy of X, working even if X is subdataframe
    for ftr in intersect(Set(categoricals),Set(scheme.features))
        n_levels = scheme.to_int_scheme_given_feature[ftr].n_levels
        for k in 1:n_levels
            subftr = Symbol(string(ftr,"__",k))
            # in the (rare) case subft is not a new feature label:
            while subftr in names(Xout)
                subftr = Symbol(string(subftr,"_"))
            end
            Xout[subftr] = transform(scheme, ftr, k, X[ftr])
        end
        delete!(Xout, ftr)
    end

    return Xout

end

end # of module