const AbstractMatrixBlock{T, M} = SubArray{T,2,M,Tuple{UnitRange{Int},UnitRange{Int}},false}

"""
$(TYPEDEF)

`SegmentedBlockDiagonalMatrix` is an `AbstractMatrix` backed by a parent `AbstractMatrix`, which
additionally stores a sequence of views into the diagonal blocks of the parent matrix. This
type is useful for storing and updating block-diagonal matrices whose block contents
may change but whose overall structure is fixed, such as configuration derivative <-> velocity
jacobians.
"""
struct SegmentedBlockDiagonalMatrix{T, M<:AbstractMatrix{T}} <: AbstractMatrix{T}
    parent::M
    blocks::Vector{AbstractMatrixBlock{T, M}}

    function SegmentedBlockDiagonalMatrix{T, M}(parent::M, block_indices) where {T, M<:AbstractMatrix{T}}
        check_contiguous_block_ranges(parent, block_indices)
        blocks = collect(view(parent, indices...) for indices in block_indices)
        new{T, M}(parent, blocks)
    end
end

SegmentedBlockDiagonalMatrix(parent::M, block_indices) where {T, M<:AbstractMatrix{T}} = SegmentedBlockDiagonalMatrix{T, M}(parent, block_indices)

function SegmentedBlockDiagonalMatrix{T}(initializer, rows::Integer, cols::Integer, block_indices) where T
    parent = Matrix{T}(initializer, rows, cols)
    SegmentedBlockDiagonalMatrix{T}(parent, block_indices)
end

Base.parent(m::SegmentedBlockDiagonalMatrix) = m.parent
Base.size(m::SegmentedBlockDiagonalMatrix) = size(m.parent)
@propagate_inbounds Base.getindex(v::SegmentedBlockDiagonalMatrix, i::Int) = v.parent[i]
@propagate_inbounds Base.setindex!(v::SegmentedBlockDiagonalMatrix, value, i::Int) = v.parent[i] = value
Base.IndexStyle(::Type{<:SegmentedBlockDiagonalMatrix}) = IndexLinear()
blocks(m::SegmentedBlockDiagonalMatrix) = m.blocks

function check_contiguous_block_ranges(parent::AbstractMatrix, block_indices)
    if !_is_contiguous_and_diagonal(parent, block_indices)
        throw(ArgumentError("The `block_indices` should be a vector of index ranges corresponding to non-overlapping contiguous diagonal blocks"))
    end
end

function _is_contiguous_and_diagonal(parent::AbstractMatrix, block_indices)
    expected_starts = first.(axes(parent))
    for inds in block_indices
        if first.(inds) !== expected_starts
            return false
        end
        expected_starts = last.(inds) .+ 1
    end
    if expected_starts !== last.(axes(parent)) .+ 1
        return false
    end
    return true
end

function LinearAlgebra.mul!(C::Matrix, A::Matrix, B::SegmentedBlockDiagonalMatrix)
    # TODO: coordination with BLAS threads
    @boundscheck size(C) == (size(A, 1), size(B, 2)) || throw(DimensionMismatch("Output size mismatch."))
    A′ = Base.unalias(C, A)
    Threads.@threads for block in blocks(B) # allocates 32 bytes (see https://github.com/JuliaLang/julia/issues/29748)
        @inbounds begin
            Acols, Ccols = parentindices(block)
            Aview = uview(A′, :, Acols)
            Cview = uview(C, :, Ccols)
            if block == I
                copyto!(Cview, Aview)
            else
                mul!(Cview, Aview, block)
            end
        end
    end
    return C
end

function LinearAlgebra.mul!(C::Matrix, A::SegmentedBlockDiagonalMatrix, B::Matrix)
    # TODO: coordination with BLAS threads
    @boundscheck size(C) == (size(A, 1), size(B, 2)) || throw(DimensionMismatch("Output size mismatch."))
    B′ = Base.unalias(C, B)
    Threads.@threads for block in blocks(A) # allocates 32 bytes (see https://github.com/JuliaLang/julia/issues/29748)
        @inbounds begin
            Crows, Brows = parentindices(block)
            Bview = uview(B, Brows, :)
            Cview = uview(C, Crows, :)
            if block == I
                copyto!(Cview, Bview)
            else
                mul!(Cview, block, Bview)
            end
        end
    end
    return C
end


using Base.Test
# blob.jl
import WaveletScattering: ScatteredBlob
# node.jl
import WaveletScattering: Node
# path.jl
import WaveletScattering: Path, PathKey

# show
io = IOBuffer()
node = Node([1.0, 2.0], (PathKey(:time)=>1:1:2,))
blob = ScatteredBlob(DataStructures.SortedDict((Path() => node,)))
show(io, blob)
@test String(take!(io)) == "ScatteredBlob(1 node)"
blob = ScatteredBlob(DataStructures.SortedDict(
    (Path((:γ, :time) => 0) => node, Path((:γ, :time) => 1) => node)))
show(io, blob)
@test String(take!(io)) == "ScatteredBlob(2 nodes)"


#
# This file is a part of MolecularGraph.jl
# Licensed under the MIT License http://opensource.org/licenses/MIT
#

export
    PlainDiGraph, plaindigraph

struct PlainDiGraph <: OrderedDiGraph
    outneighbormap::Vector{Dict{Int,Int}}
    inneighbormap::Vector{Dict{Int,Int}}
    edges::Vector{Tuple{Int,Int}}
    cache::Dict{Symbol,Any}
end

"""
    plaindigraph() -> PlainDiGraph

Generate empty `PlainDiGraph`.
"""
plaindigraph() = PlainDiGraph([], [], [], Dict())

"""
    plaindigraph(nodes, edges) -> PlainDiGraph

Generate `DiGraph` that have given nodes and edges represented by the list of
node indices in integer and the list of pairs of node indices, respectively.
"""
function plaindigraph(size::Int, edges)
    outnbrmap = [Dict{Int,Int}() for i in 1:size]
    innbrmap = [Dict{Int,Int}() for i in 1:size]
    edges = collect(edges)
    for (i, (s, t)) in enumerate(edges)
        outnbrmap[s][i] = t
        innbrmap[t][i] = s
    end
    return PlainDiGraph(outnbrmap, innbrmap, edges, Dict{Symbol,Any}())
end

"""
    plaindigraph(graph::OrderedDiGraph) -> PlainDiGraph

Convert an arbitrary `OrderedDiGraph` into a `PlainDiGraph`.
"""
function plaindigraph(graph::OrderedDiGraph)
    newg = plaindigraph()
    for nbr in graph.outneighbormap
        push!(newg.outneighbormap, copy(nbr))
    end
    for nbr in graph.inneighbormap
        push!(newg.inneighbormap, copy(nbr))
    end
    append!(newg.edges, graph.edges)
    return newg
end

"""
    plaindigraph(graph::DirectedGraph) -> PlainDiGraph

Convert an arbitrary `DirectedGraph` into a `PlainDiGraph`.

If the given graph has non-sequential indices (ex. subgraph view), node and
edge indices are sorted in ascending order and are re-indexed.
"""
function plaindigraph(graph::DirectedGraph)
    newg = plaindigraph()
    nkeys = sort(collect(nodeset(graph)))
    ekeys = sort(collect(edgeset(graph)))
    nmap = Dict{Int,Int}()
    for (i, n) in enumerate(nkeys)
        nmap[n] = i
        push!(newg.outneighbormap, Dict())
        push!(newg.inneighbormap, Dict())
    end
    for (i, e) in enumerate(ekeys)
        (olds, oldt) = getedge(graph, e)
        s = nmap[olds]
        t = nmap[oldt]
        push!(newg.edges, (s, t))
        newg.outneighbormap[s][i] = t
        newg.inneighbormap[t][i] = s
    end
    return newg
end

function unlinknodes(graph::PlainDiGraph, nodes)
    # TODO: costful
    subg = nodesubgraph(graph, setdiff(nodeset(graph), nodes))
    return plaindigraph(subg)
end

function unlinkedges(graph::PlainDiGraph, edges)
    # TODO: costful
    subg = SubgraphView(graph, nodeset(graph), setdiff(edgeset(graph), edges))
    return plaindigraph(subg)
end


using Printf # Needed to use printf
using Statistics

include("variables.jl")
include("strings.jl")
include("conditionals.jl")
include("looping.jl")
include("arrays.jl")
include("tuples.jl")
include("dictionaries.jl")
include("sets.jl")
include("functions.jl")
include("anonymous_functions.jl")
include("maths.jl")
include("enums.jl")
include("symbols.jl")
include("structs.jl")
include("abstract_types.jl")
include("exceptions.jl")
include("file_io.jl")
include("macros.jl")
include("map_filter_reduce.jl")

function main()
    variable_examples()
    string_examples()
    conditional_examples()
    looping_examples()
    array_examples()
    tuple_examples()
    dictionary_examples()
    set_examples()
    function_examples()
    anonymous_function_examples()
    maths_examples()
    enum_examples()
    symbol_examples()
    struct_examples()
    abstract_type_examples()
    exception_examples()
    file_io_examples()
    macro_examples()
    map_filter_reduce_examples()
end

main()


__precompile__()
module NoiseWaveform
include("pcfm.jl")
include("optimization.jl")
end


module Causal

using Spec
using DocStringExtensions
using Cassette
using ..Space
import ..Omega
using ..Omega: RandVar, constant, proj
import ..Omega: apl, ppapl, id, params, ciid
using Callbacks: donothing

include("replace.jl")
include("counterfactual.jl")
include("causes.jl")

export replace,
       cf,
       iscausebf
end


##  Conclusion: make my life easier and just define a GraphNode and a TreeNode

print("Hi!")
module tmp
    using GeometryTypes
    struct Node2
        id::Int
        iPrev::Int #parent
        state::Nullable{Point2f0}
    end

    #Node2(id::Int, iPrev::Int) = Node2(id, iPrev, Point2f0())
    Node2(id::Int, iPrev::Int) = Node2(id, iPrev, Point(999,999))
    # Node2(id::Int, iPrev::Int) = Node2(id, iPrev, Point())
#
	# Fails with;
	# WARNING: replacing module tmp
	# ERROR: LoadError: MethodError: Cannot `convert` an object of type Tuple{} to an object of type Float32
	# This may have arisen from a call to the constructor Float32(...),
	# since type constructors fall back to convert methods.

end

    x = tmp.Node2(1,2,Point(2,3))
    y = tmp.Node2(1,2)


"""
Initial query by controller receiving OF message
"""
function do_query!(msg::OFMessage,a::Agent,model)
    # If asset's network does not have any edge, there is no way to transport packets
    ignore = ne(a.ntw_graph) > 0 ? false : true
    path=[]

    if !ignore
        # src (from this sne) to dst
        query = (msg.dpid,msg.data.dst)
        query_time = model.ticks
        #existing precalc paths
        query_paths = a.paths
        #local graph to query
        query_graph = a.ntw_graph

        ####For benchmark#####
        sdir = model.data_dir
        record_benchmark!(sdir,model.run_label,a.id,query_time,query,query_graph,query_paths,model.benchmark)

        path = do_query(query_time,query,query_graph,query_paths)

        # log_info(model.ticks,a.id,"query: $(query) ----path found: $path ----> in precalc paths: $(query_paths)")

        # For no centralised control topo
        if model.ctrl_model != GraphModel(1) # ¬ centralised
            #check path validity as it is unknown how long ago the cache/graph has been updated
            if !isempty(path) && is_invalid_path(path,get_controlled_assets(a.id,model),model)
                path =  []
                clear_cache!(a,model)
            end

            #valid path not found?, need to send queries to nbs
            if isempty(path)
                #should always query nbs even if has already queried with no reply?
                sent_to = query_nbs!(a,msg,query_graph,query,[a.id],model)
                record_pending_query!(a,model.ticks,sent_to,query)
                push!(a.pending,(model.ticks,msg,false)) #false:not reprocess next time
            else
                clear_pending_query!(a,query)
            end
        end
        # update_agent_state!(a,msg,path)
        new_state = get_state(a)
        ## +1 query
        new_state.q_queries += 1.0
        ## +score path
        if !isempty(path) && msg.dpid != msg.data.dst
            push!(new_state.path_scores,(msg.dpid,msg.data.dst,path[2]))
        end
        set_state!(a,new_state)

    end
    return isempty(path) ? [] : last(path)
end

"""
    Query by neighbour control agent after receiving AGMessage
"""
function do_query!(msg::AGMessage,a::Agent,model)
    # log_info(model.ticks,a.id,"query msg is: $msg")

        # visited control ag
        trace = deepcopy(msg.body[:trace])
        push!(trace,a.id)
        query = msg.body[:query]

        if get_state(a).up
            # Aggregated information from other agents
            ## join graph received
            msg_ntw_g = create_subgraph(msg.body[:ntw_edgel],msg.body[:ntw_equiv],:eid)
            jg = join_subgraphs(a.ntw_graph,msg_ntw_g)

            # Update Knowledge Base
            # update local graph, however the problem it can grow too much
            # hence there is a periodic clear_cache to restore to original
            if model.clear_cache_graph_freq != 0
                a.ntw_graph = jg
            end

            query_time = model.ticks
            query_paths = a.paths
            query_graph = jg
            # path = Dict()

            #path for benchmark recording
            sdir = model.data_dir * model.run_label
            record_benchmark!(sdir,model.run_label,a.id,query_time,query,query_graph,query_paths,model.benchmark)

            #query

            path = do_query(query_time,query,query_graph,query_paths)

            if isempty(path)
                if !haskey(a.previous_queries,query)
                    sent_to = query_nbs!(a,msg,jg,query,trace,model)
                    record_pending_query!(a,model.ticks,sent_to,query)
                end
            else
                do_match!(path,msg,a,model)
                clear_pending_query!(a,query)
            end

            new_state = get_state(a)
            new_state.q_queries += 1.0
            set_state!(a,new_state)

        #TODO consider case where multiple paths are found
        else
            fw_query_nbs!(a,msg,trace,model)
        end

end

"""
Query local calculated paths and local graph
"""
function do_query(time::Int64,query::Tuple{Int64,Int64},lg::MetaGraph,paths::Dict{Tuple{Int64,Int64},Array{Tuple{Int64,Float64,Float64,Array{Int64}}}})
    path = []
    cp_paths = []
    lg_paths = []

    #query pre-calculated (cache) paths
    cp_paths = haskey(paths,query) ? paths[query] : []

    #query graph path regardless of cache, in case there is another
    # TODO: Do this only if cache path is too old
    path_state = query_paths(lg,query)

    #reshape paths found local graph
    for lg_path in path_state.paths
        default_confidence = 0.5
        push!(lg_paths,(time,default_confidence,last(path_state.dists),last(path_state.paths)))
    end

    #assumes query_paths is sorted by tick,score
    lg_ps = !isempty(lg_paths) ? first(lg_paths)[3] : -1.0
    cp_ps = !isempty(cp_paths) ? first(cp_paths)[3] : -1.0

    path = @match (lg_ps,cp_ps) begin
        #no path
        (-1.0,-1.0)  => []
        #path in local graph
        (a::Float64,-1.0), if a > 0 end => first(lg_paths)
        #path in cache
        (-1.0,b::Float64), if b > 0 end => first(cp_paths)
        #path in cache paths and local graph
        (_::Float64,_::Float64) => first(cp_paths)[3] < first(lg_paths)[3] ? first(cp_paths) : first(lg_paths) ##lower is better
    end
    return path
end

"""
Prepare query msg and send to neighbors
"""

function query_nbs!(a::Agent,msg::CTLMessage,query_graph::MetaGraph,query::Tuple{Int64,Int64},trace::Array{Int64},model::ABM)::Array{Int64}
    ntw_edgel = [ e for e in edges(query_graph) if src(e) <  dst(e) ]
    ntw_equiv = [(v,query_graph[v,:eid]) for v in vertices(query_graph)]
    # when length == 1, msg is the original OFMessage otherwise is AGMessage
    of_mid = length(trace) == 1 ? msg.id : msg.body[:of_mid]
    body = Dict(:query=>query,:trace=>trace,:ntw_edgel => ntw_edgel, :ntw_equiv=>ntw_equiv, :of_mid=>of_mid)
    msg_template = AGMessage(-1,model.ticks,a.id,-1,QUERY_PATH,body)
    return send_to_nbs!(msg_template,a,model)
end

"""
    Forward query to neighbors when control agent is down
"""

function fw_query_nbs!(a::Agent,msg::AGMessage,trace::Array{Int64},model::ABM)
    of_mid = msg.body[:of_mid]
    nbody = Dict(:query=>msg.body[:query],:trace=>trace,:ntw_edgel => msg.body[:ntw_edgel], :ntw_equiv=>msg.body[:ntw_equiv], :of_mid=>of_mid)
    msg_template = AGMessage(-1,model.ticks,a.id,-1,QUERY_PATH,nbody)
    send_to_nbs!(msg_template,a,model)
end

"""
records when control ag send a query to other control ags
"""

function record_pending_query!(a::Agent,time::Int64,sent_to::Array{Int64},query::Tuple{Int64,Int64})
    #track query sent, if not previously tracked
    if !haskey(a.previous_queries,query)
        a.previous_queries[query] = (time,sent_to)
    end
end

"""
clear when a solution is found/received to a query
"""

function clear_pending_query!(a::Agent,query::Tuple{Int64,Int64})
    #no need to clear pending?
    new_previous_queries = Dict()
    #path found, need to remove potential pending queries
    for k in keys(a.previous_queries)
        if k != query #exclude current query as path has been found
            new_previous_queries[k] = a.previous_queries[k]
        end
    end
    a.previous_queries = new_previous_queries
end

"""
clear queries requested by sne to controller that have expired, so the same query can be requested after some time model.ofmsg_reattempt
"""
function clear_pending_query!(a::SimNE,model::ABM)
    #no need to clear pending?
    new_requested_ctl = Dict()
    #path found, need to remove potential pending queries
    for k in keys(a.requested_ctl)
        if model.ticks - a.requested_ctl[k] < model.ofmsg_reattempt #
            new_requested_ctl[k] = a.requested_ctl[k]
        end
    end
    a.requested_ctl = new_requested_ctl
end

"""
When a solution is received for a query it marks the pending msg
so it can be reprocessed in the next tick
"""
function mark_reprocess_of_msg!(a::Agent,msg::AGMessage)
    #reprocess of msg right after, to do local query with new path found
    new_pending = []
    for p in a.pending
        if p[2].id == msg.body[:of_mid]
            push!(new_pending,(p[1],p[2],true))#reprocess next tick
        else
            push!(new_pending,p)
        end
    end
    a.pending = new_pending
end


#####
##### Benchmark.jl
##### Script to run a new benchmark on a previous training session
##### (requires training with option `save_intermediate=true`)
#####

using AlphaZero

include("games.jl")
GAME = get(ENV, "GAME", "connect-four")
SelectedGame = GAME_MODULE[GAME]
using .SelectedGame: Game, Training

SESSION_DIR = "sessions/$GAME"

# Examples: run a network-only benchmark

netonly = Benchmark.NetworkOnly(τ=0.5, use_gpu=true)
alphazero = Benchmark.Full(Training.arena.mcts)
baselines = [Training.mcts_baseline, Training.minmax_baseline]

make_duel(player, baseline) =
  Benchmark.Duel(
    player,
    baseline,
    num_games=200,
    reset_every=40,
    flip_probability=0.5,
    color_policy=CONTENDER_WHITE)

netonly_benchmark   = [make_duel(netonly, b) for b in baselines]
alphazero_benchmark = [make_duel(alphazero, b) for b in baselines]

benchmark = alphazero_benchmark
label = "full"

AlphaZero.UserInterface.run_new_benchmark(
  Game, Training.Network{Game}, SESSION_DIR,
  label, benchmark,
  params=Training.params, itcmax=nothing)


"""
    get_action_timers() : Nothing

Test performance of current interpolation routines.
"""
function get_action_timers() :: Nothing

    # init test dummys (taking action_su2 as an example)
    list = get_mesh(rand(), 1.0, 30, 0.4)
    m    = Mesh(31, 31, 31, 31, list, list, list, list, list, list)
    p1   = get_param(rand(), list)
    p2   = get_param(rand(), list)
    p3   = get_param(rand(), list)
    l    = get_lattice("square", 10, verbose = false)
    r    = get_reduced_lattice("heisenberg", [[1.0]], l, verbose = false)
    a    = get_action_empty("su2", r, m); init_action!(l, r, a)
    temp = zeros(Float64, length(r.sites), 1, 1)

    # fill self energy with random values
    a.Σ .= rand(31)

    # fill s channel with random values
    a.Γ[1].ch_s.q3 .= rand(length(r.sites), 31, 31, 31)

    # init timer
    to = TimerOutput()

    # time self energy interpolation and extrapolation
    val   = rand()
    outer = m.σ[end] + val

    @timeit to "=> interpolation / extrapolation Σ" begin
        for rep in 1 : 10
            @timeit to "-> interpolation" get_Σ(val, m, a)
            @timeit to "-> extrapolation" get_Σ(outer, m, a)
        end
    end

    # deref channel and generate temp view
    ch    = a.Γ[1].ch_s
    vtemp = view(temp, :, 1, 1)

    # time q3 interpolations
    @timeit to "=> interpolation q3 kernel" begin
        # time sequential routine
        for rep in 1 : 10
            @timeit to "-> sequential" begin
                for i in 1 : size(temp, 1)
                    temp[i, 1, 1] = get_q3(i, p1, p2, p3, ch)
                end
            end
        end

        # time vectorized routine
        for rep in 1 : 10
            @timeit to "-> vectorized" begin
                get_q3_avx!(r, p1, p2, p3, ch, vtemp, false, 1.0)
            end
        end
    end

    # time q2_2 interpolations
    @timeit to "=> interpolation q2_2 kernel" begin
        # time sequential routine
        for rep in 1 : 10
            @timeit to "-> sequential" begin
                for i in 1 : size(temp, 1)
                    temp[i, 1, 1] = get_q2_2(i, p1, p3, ch)
                end
            end
        end

        # time vectorized routine
        for rep in 1 : 10
            @timeit to "-> vectorized" begin
                get_q2_2_avx!(r, p1, p3, ch, vtemp, false, 1.0)
            end
        end
    end

    # time q2_1 interpolations
    @timeit to "=> interpolation q2_1 kernel" begin
        # time sequential routine
        for rep in 1 : 10
            @timeit to "-> sequential" begin
                for i in 1 : size(temp, 1)
                    temp[i, 1, 1] = get_q2_1(i, p1, p2, ch)
                end
            end
        end

        # time vectorized routine
        for rep in 1 : 10
            @timeit to "-> vectorized" begin
                get_q2_1_avx!(r, p1, p2, ch, vtemp, false, 1.0)
            end
        end
    end

    # time q1 interpolations
    @timeit to "=> interpolation q1 kernel" begin
        # time sequential routine
        for rep in 1 : 10
            @timeit to "-> sequential" begin
                for i in 1 : size(temp, 1)
                    temp[i, 1, 1] = get_q1(i, p1, ch)
                end
            end
        end

        # time vectorized routine
        for rep in 1 : 10
            @timeit to "-> vectorized" begin
                get_q1_avx!(r, p1, ch, vtemp, false, 1.0)
            end
        end
    end

    show(to)

    return nothing
end


module Layered

using StaticArrays

using Reexport
@reexport using Colors
using Match
import Cairo

export Transform, Layer, Shape, upward_transform, solve!, →, Angle, rad, deg, needed_attributes

include("types.jl")
include("angles.jl")
include("geometry.jl")
include("gradient.jl")
include("attributes.jl")
include("transform.jl")
include("layer.jl")
include("shape.jl")
include("canvas.jl")
include("drawing.jl")
include("video.jl")
include("helpers.jl")

Base.Broadcast.broadcastable(g::GeometricObject) = Ref(g)

function add_shape!(shape, layer::Layer, args...)

end

for geom in geoms

    lowerc = Symbol(lowercase(String(geom)))
    lowerc_plural = Symbol(lowercase(String(geom)) * "s")
    lowerc_mutating = Symbol(lowercase(String(geom)) * "!")
    lowerc_plural_mutating = Symbol(lowercase(String(geom)) * "s!")

    @eval begin
        $geom(g::$geom) = $geom((getfield(g, k) for k ∈ fieldnames($geom))...)

        """
        $($lowerc)(args...)

        Creates a shape containing a `GeometricObject` of type `$($geom)`,
        passing any trailing arguments to the constructor `$($geom)()`."""
        $lowerc(args...) = Shape($geom, args...)

        """
        $($lowerc)(f::Function, args...)

        Creates a shape containing a `GeometricObject` of type `$($geom)`,
        storing any trailing arguments as dependencies, later to be passed as arguments
        to the given `Function` `f` that should return a `GeometricObject` of type `$($geom)`
        and will be evaluated when `solve!` is called on the shape during
        the drawing process."""
        $lowerc(f::Function, args...) = Shape(f, $geom, args...)

        """
        $($lowerc_mutating)(layer::Layer, args...)

        Creates a shape containing a `GeometricObject` of type `$($geom)`,
        passing any trailing arguments to the constructor `$($geom)()`. Then
        appends the created shape to the given `Layer` `layer`."""
        function $lowerc_mutating(layer::Layer, args...)
            x = $lowerc(args...)
            push!(layer, x)
            x
        end

        """
        $($lowerc_mutating)(f::Function, layer::Layer, args...)

        Creates a shape containing a `GeometricObject` of type `$($geom)`,
        storing any trailing arguments as dependencies, later to be passed as arguments
        to the given `Function` `f` that should return a `GeometricObject` of type `$($geom)`
        and will be evaluated when `solve!` is called on the shape during
        the drawing process. This function then appends the created shape to the given `Layer` `layer`."""
        function $lowerc_mutating(f::Function, layer::Layer, args...)
            x = $lowerc(f, args...)
            push!(layer, x)
            x
        end

        """
        $($lowerc_plural)(args...)

        Creates `Shapes` containing `GeometricObject`s of type `$($geom)`,
        passing any trailing arguments to the constructor `$($geom).()`."""
        $lowerc_plural(args...) = Shapes($geom, args...)

        """
        $($lowerc_plural)(f::Function, args...)

        Creates shapes containing `GeometricObject`s of type `$($geom)`,
        storing any trailing arguments as dependencies, later to be passed as arguments
        to the given `Function` `f` that should return `GeometricObject`s of type `$($geom)`
        and will be evaluated when `solve!` is called on the shapes during
        the drawing process."""
        $lowerc_plural(f::Function, args...) = Shapes(f, $geom, args...)

        """
        $($lowerc_plural_mutating)(layer::Layer, args...)

        Creates `Shapes` containing `GeometricObject`s of type `$($geom)`,
        passing any trailing arguments to the constructor `$($geom).()`. Then
        appends the created `Shapes` to the given `Layer` `layer`."""
        function $lowerc_plural_mutating(layer::Layer, args...)
            x = $lowerc_plural(args...)
            push!(layer, x)
            x
        end

        """
        $($lowerc_plural_mutating)(f::Function, layer::Layer, args...)

        Creates multiple shapes containing `GeometricObject`s of type `$($geom)`,
        storing any trailing arguments as dependencies, later to be passed as arguments
        to the given `Function` `f` that should return the `GeometricObject`s of type `$($geom)`
        and will be evaluated when `solve!` is called on the shapes during
        the drawing process. This function then appends the created shapes to the given `Layer` `layer`."""
        function $lowerc_plural_mutating(f::Function, layer::Layer, args...)
            x = $lowerc_plural(f, args...)
            push!(layer, x)
            x
        end

        export $geom, $lowerc, $lowerc_mutating, $lowerc_plural, $lowerc_plural_mutating

    end
end

end # module


function evaluate!{F, LD}(output::AbstractArray{F}, mlp::MultilayerPerceptron{F, LD}, input::AbstractArray, layerstorage::Storage{MultilayerPerceptron{F, LD}})
  l = length(LD)
  #run a batch evaluate WITH using buffers for intermediate values.
  evaluate!(layer(layerstorage, Val{2}), mlp.layers[1], input)
  for idx = 2:(l - 2)
    evaluate!(layer(layerstorage, Val{idx + 1}), mlp.layers[idx], layer(layerstorage, Val{idx}))
  end
  evaluate!(output, mlp.layers[end], layer(layerstorage, Val{l - 1}))
end


using JudiLing # our package
using CSV # read csv files
using DataFrames # parse data into dataframes

mkpath(joinpath(@__DIR__, "data"))
download("https://osf.io/2ejfu/download", joinpath(@__DIR__, "data", "latin.csv"))
# load latin file
latin = DataFrame(CSV.File(joinpath(@__DIR__, "data", "latin.csv")))
display(latin)

# create C matrixes for training datasets
cue_obj = JudiLing.make_cue_matrix(
    latin,
    grams = 3,
    target_col = :Word,
    tokenized = false,
    keep_sep = false,
)

# retrieve dim of C
# we set the S matrixes as the same dimensions
n_features = size(cue_obj.C, 2)
S = JudiLing.make_S_matrix(
    latin,
    ["Lexeme"],
    ["Person", "Number", "Tense", "Voice", "Mood"],
    ncol = n_features,
)

# we use cholesky function to calculate mapping G from S to C
G = JudiLing.make_transform_matrix(S, cue_obj.C)

# we calculate Chat matrixes by multiplying S and G
Chat = S * G
@show JudiLing.eval_SC(Chat, cue_obj.C)

# we calculate F as we did for G
F = JudiLing.make_transform_matrix(cue_obj.C, S)

# we calculate Shat as we did for Chat
Shat = cue_obj.C * F
@show JudiLing.eval_SC(Shat, S)

# here we only use a adjacency matrix as we got it from training dataset
A = cue_obj.A

# we calculate how many timestep we need for learn_paths and huo function
max_t = JudiLing.cal_max_timestep(latin, :Word)

# we calculate learn_paths and build_paths function
res_learn, gpi_learn = JudiLing.learn_paths(
    latin,
    latin,
    cue_obj.C,
    S,
    F,
    Chat,
    A,
    cue_obj.i2f,
    cue_obj.f2i, # api changed in 0.3.1
    check_gold_path = true,
    gold_ind = cue_obj.gold_ind,
    Shat_val = Shat,
    max_t = max_t,
    max_can = 10,
    grams = 3,
    threshold = 0.05,
    tokenized = false,
    keep_sep = false,
    target_col = :Word,
    verbose = true,
)

acc_learn = JudiLing.eval_acc(res_learn, cue_obj.gold_ind, verbose = false)

println("Acc for learn: $acc_learn")

res_build = JudiLing.build_paths(
    latin,
    cue_obj.C,
    S,
    F,
    Chat,
    A,
    cue_obj.i2f,
    cue_obj.gold_ind,
    max_t = max_t,
    n_neighbors = 3,
    verbose = true,
)

acc_build = JudiLing.eval_acc(res_build, cue_obj.gold_ind, verbose = false)

println("Acc for build: $acc_build")

# you can save results into csv files or dfs
JudiLing.write2csv(
    res_learn,
    latin,
    cue_obj,
    cue_obj,
    "latin_learn_res.csv",
    grams = 3,
    tokenized = false,
    sep_token = nothing,
    start_end_token = "#",
    output_sep_token = "",
    path_sep_token = ":",
    target_col = :Word,
    root_dir = @__DIR__,
    output_dir = "latin_out",
)

df_learn = JudiLing.write2df(
    res_learn,
    latin,
    cue_obj,
    cue_obj,
    grams = 3,
    tokenized = false,
    sep_token = nothing,
    start_end_token = "#",
    output_sep_token = "",
    path_sep_token = ":",
    target_col = :Word,
)

JudiLing.write2csv(
    res_build,
    latin,
    cue_obj,
    cue_obj,
    "latin_build_res.csv",
    grams = 3,
    tokenized = false,
    sep_token = nothing,
    start_end_token = "#",
    output_sep_token = "",
    path_sep_token = ":",
    target_col = :Word,
    root_dir = @__DIR__,
    output_dir = "latin_out",
)

df_build = JudiLing.write2df(
    res_build,
    latin,
    cue_obj,
    cue_obj,
    grams = 3,
    tokenized = false,
    sep_token = nothing,
    start_end_token = "#",
    output_sep_token = "",
    path_sep_token = ":",
    target_col = :Word,
)

display(df_learn)
display(df_build)

# cross-validation
download("https://osf.io/2ejfu/download", joinpath(@__DIR__, "data", "latin_train.csv"))
download("https://osf.io/bm7y6/download", joinpath(@__DIR__, "data", "latin_val.csv"))

latin_train =
    DataFrame(CSV.File(joinpath(@__DIR__, "data", "latin_train.csv")))
latin_val =
    DataFrame(CSV.File(joinpath(@__DIR__, "data", "latin_val.csv")))

# create C matrices for both training and validation datasets
cue_obj_train, cue_obj_val = JudiLing.make_cue_matrix(
    latin_train,
    latin_val,
    grams = 3,
    target_col = :Word,
    tokenized = false,
    keep_sep = false,
)

# create S matrices
n_features = size(cue_obj_train.C, 2)
S_train, S_val = JudiLing.make_S_matrix(
    latin_train,
    latin_val,
    ["Lexeme"],
    ["Person", "Number", "Tense", "Voice", "Mood"],
    ncol = n_features,
)

# here we learning mapping only from training dataset
G_train = JudiLing.make_transform_matrix(S_train, cue_obj_train.C)
F_train = JudiLing.make_transform_matrix(cue_obj_train.C, S_train)

# we predict S and C for both training and validation datasets
Chat_train = S_train * G_train
Chat_val = S_val * G_train
Shat_train = cue_obj_train.C * F_train
Shat_val = cue_obj_val.C * F_train

# we evaluate them
@show JudiLing.eval_SC(Chat_train, cue_obj_train.C)
@show JudiLing.eval_SC(Chat_val, cue_obj_val.C)
@show JudiLing.eval_SC(Shat_train, S_train)
@show JudiLing.eval_SC(Shat_val, S_val)

# we can use build path and learn path
A = cue_obj_train.A
max_t = JudiLing.cal_max_timestep(latin_train, latin_val, :Word)

res_learn_train, gpi_learn_train = JudiLing.learn_paths(
    latin_train,
    latin_train,
    cue_obj_train.C,
    S_train,
    F_train,
    Chat_train,
    A,
    cue_obj_train.i2f,
    cue_obj_train.f2i, # api changed in 0.3.1
    gold_ind = cue_obj_train.gold_ind,
    Shat_val = Shat_train,
    check_gold_path = true,
    max_t = max_t,
    max_can = 10,
    grams = 3,
    threshold = 0.05,
    tokenized = false,
    sep_token = "_",
    keep_sep = false,
    target_col = :Word,
    issparse = :dense,
    verbose = true,
)

res_learn_val, gpi_learn_val = JudiLing.learn_paths(
    latin_train,
    latin_val,
    cue_obj_train.C,
    S_val,
    F_train,
    Chat_val,
    A,
    cue_obj_train.i2f,
    cue_obj_train.f2i, # api changed in 0.3.1
    gold_ind = cue_obj_val.gold_ind,
    Shat_val = Shat_val,
    check_gold_path = true,
    max_t = max_t,
    max_can = 10,
    grams = 3,
    threshold = 0.05,
    is_tolerant = true,
    tolerance = -0.1,
    max_tolerance = 2,
    tokenized = false,
    sep_token = "-",
    keep_sep = false,
    target_col = :Word,
    issparse = :dense,
    verbose = true,
)

acc_learn_train =
    JudiLing.eval_acc(res_learn_train, cue_obj_train.gold_ind, verbose = false)
acc_learn_val = JudiLing.eval_acc(res_learn_val, cue_obj_val.gold_ind, verbose = false)

res_build_train = JudiLing.build_paths(
    latin_train,
    cue_obj_train.C,
    S_train,
    F_train,
    Chat_train,
    A,
    cue_obj_train.i2f,
    cue_obj_train.gold_ind,
    max_t = max_t,
    n_neighbors = 3,
    verbose = true,
)

res_build_val = JudiLing.build_paths(
    latin_val,
    cue_obj_train.C,
    S_val,
    F_train,
    Chat_val,
    A,
    cue_obj_train.i2f,
    cue_obj_train.gold_ind,
    max_t = max_t,
    n_neighbors = 20,
    verbose = true,
)

acc_build_train =
    JudiLing.eval_acc(res_build_train, cue_obj_train.gold_ind, verbose = false)
acc_build_val = JudiLing.eval_acc(res_build_val, cue_obj_val.gold_ind, verbose = false)

@show acc_learn_train
@show acc_learn_val
@show acc_build_train
@show acc_build_val

# Once you are done, you may want to clean up the workspace
rm(joinpath(@__DIR__, "data"), force = true, recursive = true)
rm(joinpath(@__DIR__, "latin_out"), force = true, recursive = true)


module Decode
import ..Utils
import ..Frames
import ..Signals
import ..Messages

"""
"""
function decode(signal::Signals.NamedSignal{T},
                can_frame::Frames.CANFrame)::Union{Nothing,T} where {T}
    try
        sig = Signals.signal(signal)
        return decode(sig, can_frame)
    catch
        return Signals.default(signal)
    end
end

"""
"""
function decode(signal::Signals.UnnamedSignal{T}, can_frame::Frames.CANFrame,
                default::D)::Union{T,D} where {T,D}
    try
        return decode(signal, can_frame)
    catch
        return default
    end
end

"""
"""
function decode(signal::Signals.Bit, can_frame::Frames.CANFrame)::Bool
    start = Signals.start(signal)

    if start >= 8*Frames.dlc(can_frame)
        throw(DomainError(start, "CANFrame does not have data at bit position"))
    else
        mask = Utils.mask(UInt64, 1, start)
        value = Utils.from_bytes(UInt64, Frames.data(can_frame))
        if mask & value != 0
            return true
        else
            return false
        end
    end
end

"""
"""
function decode(signal::Signals.Unsigned{T}, can_frame::Frames.CANFrame) where {T}
    start = Signals.start(signal)
    length = Signals.length(signal)
    factor = Signals.factor(signal)
    offset = Signals.offset(signal)
    byte_order = Signals.byte_order(signal)

    if byte_order == :little_endian
        end_bit = start + length - 1
        if end_bit >= 8 * Frames.dlc(can_frame)
            throw(DomainError(end_bit, "The bit($end_bit) cannot be selected"))
        end

        value = Utils.from_bytes(UInt64, Frames.data(can_frame))
        value = value >> start
    elseif byte_order == :big_endian
        start_bit_in_byte = start % 8
        start_byte = div(start, 8)

        start = 8*start_byte + (7 - start_bit_in_byte)
        new_shift = Int64(8*Frames.dlc(can_frame)) - Int64(start) - Int64(length)

        if new_shift < 0
            throw(DomainError(new_shift, "The bits cannot be selected"))
        end

        value = Utils.from_bytes(UInt64, reverse(Frames.data(can_frame)))
        value = value >> new_shift
    else
        throw(DomainError(byte_order, "Byte order not supported"))
    end

    value = value & Utils.mask(UInt64, length)
    result = T(value) * factor + offset
    return result
end

function decode(signal::Signals.Signed{T}, can_frame::Frames.CANFrame) where {T}
    start = Signals.start(signal)
    length = Signals.length(signal)
    factor = Signals.factor(signal)
    offset = Signals.offset(signal)
    byte_order = Signals.byte_order(signal)

    if byte_order == :little_endian
        end_bit = start + length - 1
        if end_bit >= 8 * Frames.dlc(can_frame)
            throw(DomainError(end_bit, "The bit($end_bit) cannot be selected"))
        end

        value = Utils.from_bytes(Int64, Frames.data(can_frame))
        value = value >> start
    elseif byte_order == :big_endian
        start_bit_in_byte = start % 8
        start_byte = div(start, 8)

        start = 8*start_byte + (7 - start_bit_in_byte)
        new_shift = Int64(8*Frames.dlc(can_frame)) - Int64(start) - Int64(length)

        if new_shift < 0
            throw(DomainError(new_shift, "The bits cannot be selected"))
        end

        value = Utils.from_bytes(Int64, reverse(Frames.data(can_frame)))
        value = value >> new_shift
    else
        throw(DomainError(byte_order, "Byte order not supported"))
    end

    value = value & Utils.mask(Int64, length)
    # sign extend value if most-significant bit is 1
    if (value >> (length - 1)) & 0x01 != 0
        value = value + ~Utils.mask(Int64, length)
    end

    result = T(value) * factor + offset
    return result
end

"""
"""
function decode(signal::Signals.FloatSignal{T}, can_frame::Frames.CANFrame) where {T}
    start = Signals.start(signal)
    length = Signals.length(signal)
    factor = Signals.factor(signal)
    offset = Signals.offset(signal)
    byte_order = Signals.byte_order(signal)

    if byte_order == :little_endian
        end_bit = start + length - 1
        if end_bit >= 8 * Frames.dlc(can_frame)
            throw(DomainError(end_bit, "The bit($end_bit) cannot be selected"))
        end

        value = Utils.from_bytes(UInt64, Frames.data(can_frame))
        value = value >> start
    elseif byte_order == :big_endian
        start_bit_in_byte = start % 8
        start_byte = div(start, 8)

        start = 8*start_byte + (7 - start_bit_in_byte)
        new_shift = Int64(8*Frames.dlc(can_frame)) - Int64(start) - Int64(length)

        if new_shift < 0
            throw(DomainError(new_shift, "The bits cannot be selected"))
        end

        value = Utils.from_bytes(UInt64, reverse(Frames.data(can_frame)))
        value = value >> new_shift
    else
        throw(DomainError(byte_order, "Byte order not supported"))
    end

    value = value & Utils.mask(UInt64, length)
    result = Utils.from_bytes(T, Utils.to_bytes(value))
    result = factor * result + offset
    return result
end

"""
"""
function decode(signal::Signals.Raw, can_frame::Frames.CANFrame)::UInt64
    start = Signals.start(signal)
    length = Signals.length(signal)
    byte_order = Signals.byte_order(signal)

    if byte_order == :little_endian
        end_bit = start + length - 1
        if end_bit >= 8 * Frames.dlc(can_frame)
            throw(DomainError(end_bit, "The bit($end_bit) cannot be selected"))
        end

        value = Utils.from_bytes(UInt64, Frames.data(can_frame))
        value = value >> start
    elseif byte_order == :big_endian
        start_bit_in_byte = start % 8
        start_byte = div(start, 8)

        start = 8*start_byte + (7 - start_bit_in_byte)
        new_shift::Int16 = Int64(8*Frames.dlc(can_frame)) - Int64(start) - Int64(length)

        if new_shift < 0
            throw(DomainError(new_shift, "The bits cannot be selected"))
        end

        value = Utils.from_bytes(UInt64, reverse(Frames.data(can_frame)))
        value = value >> new_shift
    else
        throw(DomainError(byte_order, "Byte order not supported"))
    end

    result = value & Utils.mask(UInt64, length)
    return UInt64(result)
end

function decode(message::Messages.Message, can_frame::Frames.CANFrame)
    d = Dict()
    for (key, signal) in message
        value = decode(signal, can_frame)
        d[key] = value
    end

    return d
end

export decode
end


using LinearAlgebra

using ModelingToolkit: isdifferenceeq, has_continuous_events, generate_rootfinding_callback, generate_difference_cb, merge_cb

const MAX_INLINE_NLSOLVE_SIZE = 8

function torn_system_with_nlsolve_jacobian_sparsity(state, var_eq_matching, var_sccs, nlsolve_scc_idxs, eqs_idxs, states_idxs)
    fullvars = state.fullvars
    graph = state.structure.graph

    # The sparsity pattern of `nlsolve(f, u, p)` w.r.t `p` is difficult to
    # determine in general. Consider the "simplest" case, a linear system. We
    # have
    #                   A u = p.
    # Clearly, the sparsity of `u` depends on the sparsity of both `p` and `A`
    # in a non-trivial way. However, in the generic case, `u` is dense even when
    # `A` and `p` are sparse. For instance
    #
    # ```julia
    # julia> using Random, SparseArrays
    #
    # julia> A = sprand(MersenneTwister(1234), 100, 100, 0.1);
    #
    # julia> p = sprand(MersenneTwister(12345), 100, 0.05);
    #
    # julia> count(x->abs(x) < 1e-5, A \ Vector(p))
    # 0
    # ```
    #
    # Let 𝑇 be the set of tearing variables and 𝑉 be the set of all *states* in
    # the residual equations. In the following code, we are going to assume the
    # connection between 𝑇 (the `u` in from above) and 𝑉 ∖ 𝑇 (the `p` in from
    # above) has full incidence.
    #
    # Note that as we are reducing algebraic equations numerically, it could be
    # the case that a later partition (a BLT block) contains tearing variables
    # from other partitions.
    #
    # We know that partitions are BLT ordered. Hence, the tearing variables in
    # each partition is unique, and all states in a partition must be
    # either differential variables or algebraic tearing variables that are
    # from previous partitions. Hence, we can build the dependency chain as we
    # traverse the partitions.

    var_rename = ones(Int64, ndsts(graph))
    nlsolve_vars = Int[]
    for i in nlsolve_scc_idxs, c in var_sccs[i]
        append!(nlsolve_vars, c)
        for v in c
            var_rename[v] = 0
        end
    end
    masked_cumsum!(var_rename)

    dig = DiCMOBiGraph{true}(graph, var_eq_matching)

    fused_var_deps = map(1:ndsts(graph)) do v
        BitSet(v′ for v′ in neighborhood(dig, v, Inf; dir=:in) if var_rename[v′] != 0)
    end

    for scc in var_sccs[nlsolve_scc_idxs]
        if length(scc) >= 2
            deps = fused_var_deps[scc[1]]
            for c in 2:length(scc)
                union!(deps, fused_var_deps[c])
                fused_var_deps[c] = deps
            end
        end
    end

    var2idx = Dict{Int,Int}(v => i for (i, v) in enumerate(states_idxs))
    eqs2idx = Dict{Int,Int}(v => i for (i, v) in enumerate(eqs_idxs))
    nlsolve_vars_set = BitSet(nlsolve_vars)

    I = Int[]; J = Int[]
    s = state.structure
    for ieq in 𝑠vertices(graph)
        nieq = get(eqs2idx, ieq, 0)
        nieq == 0 && continue
        for ivar in 𝑠neighbors(graph, ieq)
            isdervar(s, ivar) && continue
            if var_rename[ivar] != 0
                push!(I, nieq)
                push!(J, var2idx[ivar])
            else
                for dvar in fused_var_deps[ivar]
                    isdervar(s, dvar) && continue
                    niv = get(var2idx, dvar, 0)
                    niv == 0 && continue
                    push!(I, nieq)
                    push!(J, niv)
                end
            end
        end
    end
    sparse(I, J, true)
end

function gen_nlsolve!(is_not_prepended_assignment, eqs, vars, u0map::AbstractDict, assignments, (deps, invdeps), var2assignment; checkbounds=true)
    isempty(vars) && throw(ArgumentError("vars may not be empty"))
    length(eqs) == length(vars) || throw(ArgumentError("vars must be of the same length as the number of equations to find the roots of"))
    rhss = map(x->x.rhs, eqs)
    # We use `vars` instead of `graph` to capture parameters, too.
    paramset = ModelingToolkit.vars(r for r in rhss)

    # Compute necessary assignments for the nlsolve expr
    init_assignments = [var2assignment[p] for p in paramset if haskey(var2assignment, p)]
    if isempty(init_assignments)
        needed_assignments_idxs = Int[]
        needed_assignments = similar(assignments, 0)
    else
        tmp = [init_assignments]
        # `deps[init_assignments]` gives the dependency of `init_assignments`
        while true
            next_assignments = reduce(vcat, deps[init_assignments])
            isempty(next_assignments) && break
            init_assignments = next_assignments
            push!(tmp, init_assignments)
        end
        needed_assignments_idxs = reduce(vcat, unique(reverse(tmp)))
        needed_assignments = assignments[needed_assignments_idxs]
    end

    # Compute `params`. They are like enclosed variables
    rhsvars = [ModelingToolkit.vars(r.rhs) for r in needed_assignments]
    vars_set = Set(vars)
    outer_set = BitSet()
    inner_set = BitSet()
    for (i, vs) in enumerate(rhsvars)
        j = needed_assignments_idxs[i]
        if isdisjoint(vars_set, vs)
            push!(outer_set, j)
        else
            push!(inner_set, j)
        end
    end
    init_refine = BitSet()
    for i in inner_set
        union!(init_refine, invdeps[i])
    end
    intersect!(init_refine, outer_set)
    setdiff!(outer_set, init_refine)
    union!(inner_set, init_refine)

    next_refine = BitSet()
    while true
        for i in init_refine
            id = invdeps[i]
            isempty(id) && break
            union!(next_refine, id)
        end
        intersect!(next_refine, outer_set)
        isempty(next_refine) && break
        setdiff!(outer_set, next_refine)
        union!(inner_set, next_refine)

        init_refine, next_refine = next_refine, init_refine
        empty!(next_refine)
    end
    global2local = Dict(j=>i for (i, j) in enumerate(needed_assignments_idxs))
    inner_idxs = [global2local[i] for i in collect(inner_set)]
    outer_idxs = [global2local[i] for i in collect(outer_set)]
    extravars = reduce(union!, rhsvars[inner_idxs], init=Set())
    union!(paramset, extravars)
    setdiff!(paramset, vars)
    setdiff!(paramset, [needed_assignments[i].lhs for i in inner_idxs])
    union!(paramset, [needed_assignments[i].lhs for i in outer_idxs])
    params = collect(paramset)

    # splatting to tighten the type
    u0 = []
    for v in vars
        v in keys(u0map) || (push!(u0, 1e-3); continue)
        u = substitute(v, u0map)
        for i in 1:length(u0map)
            u = substitute(u, u0map)
            u isa Number && (push!(u0, u); break)
        end
        u isa Number || error("$v doesn't have a default.")
    end
    u0 = [u0...]
    # specialize on the scalar case
    isscalar = length(u0) == 1
    u0 = isscalar ? u0[1] : SVector(u0...)

    fname = gensym("fun")
    # f is the function to find roots on
    f = Func(
        [
         DestructuredArgs(vars, inbounds=!checkbounds)
         DestructuredArgs(params, inbounds=!checkbounds)
        ],
        [],
        Let(
            needed_assignments[inner_idxs],
            isscalar ? rhss[1] : MakeArray(rhss, SVector),
            false
           )
    ) |> SymbolicUtils.Code.toexpr

    # solver call contains code to call the root-finding solver on the function f
    solver_call = LiteralExpr(quote
            $numerical_nlsolve(
                               $fname,
                               # initial guess
                               $u0,
                               # "captured variables"
                               ($(params...),)
                              )
        end)

    preassignments = []
    for i in outer_idxs
        ii = needed_assignments_idxs[i]
        is_not_prepended_assignment[ii] || continue
        is_not_prepended_assignment[ii] = false
        push!(preassignments, assignments[ii])
    end

    nlsolve_expr = Assignment[
                    preassignments
                    fname ← @RuntimeGeneratedFunction(f)
                    DestructuredArgs(vars, inbounds=!checkbounds) ← solver_call
                   ]

    nlsolve_expr
end

function build_torn_function(
        sys;
        expression=false,
        jacobian_sparsity=true,
        checkbounds=false,
        max_inlining_size=nothing,
        kw...
    )

    max_inlining_size = something(max_inlining_size, MAX_INLINE_NLSOLVE_SIZE)
    rhss = []
    eqs = equations(sys)
    eqs_idxs = Int[]
    for (i, eq) in enumerate(eqs)
        isdiffeq(eq) || continue
        push!(eqs_idxs, i)
        push!(rhss, eq.rhs)
    end

    state = get_or_construct_tearing_state(sys)
    fullvars = state.fullvars
    var_eq_matching, var_sccs = algebraic_variables_scc(state)
    condensed_graph = MatchedCondensationGraph(
        DiCMOBiGraph{true}(complete(state.structure.graph), complete(var_eq_matching)), var_sccs)
    toporder = topological_sort_by_dfs(condensed_graph)
    var_sccs = var_sccs[toporder]

    states_idxs = collect(diffvars_range(state.structure))
    mass_matrix_diag = ones(length(states_idxs))

    assignments, deps, sol_states = tearing_assignments(sys)
    invdeps = map(_->BitSet(), deps)
    for (i, d) in enumerate(deps)
        for a in d
            push!(invdeps[a], i)
        end
    end
    var2assignment = Dict{Any,Int}(eq.lhs => i for (i, eq) in enumerate(assignments))
    is_not_prepended_assignment = trues(length(assignments))

    torn_expr = Assignment[]

    defs = defaults(sys)
    nlsolve_scc_idxs = Int[]

    needs_extending = false
    @views for (i, scc) in enumerate(var_sccs)
        torn_vars_idxs = Int[var for var in scc if var_eq_matching[var] !== unassigned]
        torn_eqs_idxs = [var_eq_matching[var] for var in torn_vars_idxs]
        isempty(torn_eqs_idxs) && continue
        if length(torn_eqs_idxs) <= max_inlining_size
            nlsolve_expr = gen_nlsolve!(is_not_prepended_assignment, eqs[torn_eqs_idxs], fullvars[torn_vars_idxs], defs, assignments, (deps, invdeps), var2assignment, checkbounds=checkbounds)
            append!(torn_expr, nlsolve_expr)
            push!(nlsolve_scc_idxs, i)
        else
            needs_extending = true
            append!(eqs_idxs, torn_eqs_idxs)
            append!(rhss, map(x->x.rhs, eqs[torn_eqs_idxs]))
            append!(states_idxs, torn_vars_idxs)
            append!(mass_matrix_diag, zeros(length(torn_eqs_idxs)))
        end
    end

    mass_matrix = needs_extending ? Diagonal(mass_matrix_diag) : I

    out = Sym{Any}(gensym("out"))
    funbody = SetArray(
        !checkbounds,
        out,
        rhss
    )

    states = fullvars[states_idxs]
    syms = map(Symbol, states_idxs)

    pre = get_postprocess_fbody(sys)

    expr = SymbolicUtils.Code.toexpr(
        Func(
             [
              out
              DestructuredArgs(states, inbounds=!checkbounds)
              DestructuredArgs(parameters(sys), inbounds=!checkbounds)
              independent_variables(sys)
             ],
             [],
             pre(Let(
                     [torn_expr; assignments[is_not_prepended_assignment]],
                     funbody,
                     false
                ))
            ),
        sol_states
    )
    if expression
        expr, states
    else
        observedfun = let state = state, dict=Dict(), assignments=assignments, deps=(deps, invdeps), sol_states=sol_states, var2assignment=var2assignment
            function generated_observed(obsvar, u, p, t)
                obs = get!(dict, value(obsvar)) do
                    build_observed_function(state, obsvar, var_eq_matching, var_sccs,
                                            assignments, deps, sol_states, var2assignment,
                                            checkbounds=checkbounds,
                                           )
                end
                obs(u, p, t)
            end
        end

        ODEFunction{true}(
                          @RuntimeGeneratedFunction(expr),
                          sparsity = jacobian_sparsity ? torn_system_with_nlsolve_jacobian_sparsity(state, var_eq_matching, var_sccs, nlsolve_scc_idxs, eqs_idxs, states_idxs) : nothing,
                          syms = syms,
                          observed = observedfun,
                          mass_matrix = mass_matrix,
                         ), states
    end
end

"""
    find_solve_sequence(sccs, vars)

given a set of `vars`, find the groups of equations we need to solve for
to obtain the solution to `vars`
"""
function find_solve_sequence(sccs, vars)
    subset = filter(i -> !isdisjoint(sccs[i], vars), 1:length(sccs))
    isempty(subset) && return []
    vars′ = mapreduce(i->sccs[i], union, subset)
    if vars′ == vars
        return subset
    else
        return find_solve_sequence(sccs, vars′)
    end
end

function build_observed_function(
        state, ts, var_eq_matching, var_sccs,
        assignments,
        deps,
        sol_states,
        var2assignment;
        expression=false,
        output_type=Array,
        checkbounds=true,
    )

    is_not_prepended_assignment = trues(length(assignments))
    if (isscalar = !(ts isa AbstractVector))
        ts = [ts]
    end
    ts = Symbolics.scalarize.(value.(ts))

    vars = Set()
    sys = state.sys
    foreach(Base.Fix1(vars!, vars), ts)
    ivs = independent_variables(sys)
    dep_vars = collect(setdiff(vars, ivs))

    fullvars = state.fullvars
    s = state.structure
    graph = s.graph
    diffvars = map(i->fullvars[i], diffvars_range(s))
    algvars = map(i->fullvars[i], algvars_range(s))

    required_algvars = Set(intersect(algvars, vars))
    obs = observed(sys)
    observed_idx = Dict(map(x->x.lhs, obs) .=> 1:length(obs))
    # FIXME: this is a rather rough estimate of dependencies.
    maxidx = 0
    sts = Set(states(sys))
    for (i, s) in enumerate(dep_vars)
        idx = get(observed_idx, s, nothing)
        if idx === nothing
            if !(s in sts)
                throw(ArgumentError("$s is either an observed nor a state variable."))
            end
            continue
        end
        idx > maxidx && (maxidx = idx)
    end
    vs = Set()
    for idx in 1:maxidx
        vars!(vs, obs[idx].rhs)
        union!(required_algvars, intersect(algvars, vs))
        empty!(vs)
    end

    varidxs = findall(x->x in required_algvars, fullvars)
    subset = find_solve_sequence(var_sccs, varidxs)
    if !isempty(subset)
        eqs = equations(sys)

        nested_torn_vars_idxs = []
        for iscc in subset
            torn_vars_idxs = Int[var for var in var_sccs[iscc] if var_eq_matching[var] !== unassigned]
            isempty(torn_vars_idxs) || push!(nested_torn_vars_idxs, torn_vars_idxs)
        end
        torn_eqs = [[eqs[var_eq_matching[i]] for i in idxs] for idxs in nested_torn_vars_idxs]
        torn_vars = [fullvars[idxs] for idxs in nested_torn_vars_idxs]
        u0map = defaults(sys)
        assignments = copy(assignments)
        solves = map(zip(torn_eqs, torn_vars)) do (eqs, vars)
            gen_nlsolve!(is_not_prepended_assignment, eqs, vars,
                         u0map, assignments, deps, var2assignment; checkbounds=checkbounds)
        end
    else
        solves = []
    end

    subs = []
    for sym in vars
        eqidx = get(observed_idx, sym, nothing)
        eqidx === nothing && continue
        push!(subs, sym ← obs[eqidx].rhs)
    end
    pre = get_postprocess_fbody(sys)

    ex = Code.toexpr(Func(
        [
         DestructuredArgs(diffvars, inbounds=!checkbounds)
         DestructuredArgs(parameters(sys), inbounds=!checkbounds)
         independent_variables(sys)
        ],
        [],
        pre(Let(
            [
             collect(Iterators.flatten(solves))
             map(eq -> eq.lhs←eq.rhs, obs[1:maxidx])
             subs
             assignments[is_not_prepended_assignment]
            ],
            isscalar ? ts[1] : MakeArray(ts, output_type),
            false
           ))
       ), sol_states)

    expression ? ex : @RuntimeGeneratedFunction(ex)
end

struct ODAEProblem{iip}
end

ODAEProblem(args...; kw...) = ODAEProblem{true}(args...; kw...)

"""
    ODAEProblem{iip}(sys, u0map, tspan, parammap = DiffEqBase.NullParameters(); kw...)

This constructor acts similar to the one for [`ODEProblem`](@ref) with the following changes:
`ODESystem`s can sometimes be further reduced if `structural_simplify` has
already been applied to them. This is done this constructor.
In these cases, the constructor uses the knowledge of the strongly connected
components calculated during the process of simplification as the basis for
building pre-simplified nonlinear systems in the implicit solving.
In summary: these problems are structurally modified, but could be
more efficient and more stable. Note, the returned object is still of type [`ODEProblem`](@ref).
"""
function ODAEProblem{iip}(
                          sys,
                          u0map,
                          tspan,
                          parammap=DiffEqBase.NullParameters();
                          callback = nothing,
                          kwargs...
                         ) where {iip}
    fun, dvs = build_torn_function(sys; kwargs...)
    ps = parameters(sys)
    defs = defaults(sys)

    u0 = ModelingToolkit.varmap_to_vars(u0map, dvs; defaults=defs)
    p = ModelingToolkit.varmap_to_vars(parammap, ps; defaults=defs)

    has_difference = any(isdifferenceeq, equations(sys))
    if has_continuous_events(sys)
        event_cb = generate_rootfinding_callback(sys; kwargs...)
    else
        event_cb = nothing
    end
    difference_cb = has_difference ? generate_difference_cb(sys; kwargs...) : nothing
    cb = merge_cb(event_cb, difference_cb)
    cb = merge_cb(cb, callback)

    if cb === nothing
        ODEProblem{iip}(fun, u0, tspan, p; kwargs...)
    else
        ODEProblem{iip}(fun, u0, tspan, p; callback=cb, kwargs...)
    end
end


function drop!(cur_player::AbstractPlayer, cur_piece::AbstractPiece, cur_grid::AbstractGrid, cur_direction::Integer, is_main_piece::Bool, is_step::Bool)
  can_move(cur_player) || return false

  cur_blocks = map(
    cur_block -> calc_block_coords(cur_block),
    cur_piece.blocks
  )

  did_move = false

  new_direction = 0

  for tmp_direction in sign(cur_direction):sign(cur_direction):cur_direction

    for (cur_index, cur_block) in enumerate(cur_blocks)
      cur_blocks[cur_index] = (
        first(cur_block) + sign(cur_direction),
        last(cur_block) + 0
      )
    end

    check(cur_blocks, cur_grid) || break

    did_move = true

    new_direction = tmp_direction

  end

  did_move || return false

  cur_piece.i += new_direction

  if is_main_piece && !is_step
    drop_count = abs(new_direction)

    if drop_count == 1
      score!(cur_player, "soft", -1)
    else
      score!(cur_player, "hard", drop_count)
    end
  end

  if is_main_piece
    move!(cur_player, false)

    cur_player.clock.drag =
      Nullable{UUIDs.UUID}()
  end

  true
end

function drop!(cur_player::AbstractPlayer, cur_piece::Nullable{AbstractPiece}, cur_grid::AbstractGrid, cur_direction::Integer, is_main_piece::Bool, is_step::Bool)
  false
end


#################################### Branch Variables ##################################################
get_variable_binary(_, ::Type{<:PSY.DCBranch}, ::AbstractDCLineFormulation) = false

get_variable_multiplier(::FlowActivePowerVariable, ::Type{<:PSY.DCBranch}, _) = NaN

get_variable_multiplier(
    ::FlowActivePowerFromToVariable,
    ::Type{<:PSY.DCBranch},
    ::AbstractDCLineFormulation,
) = -1.0
get_variable_multiplier(
    ::FlowActivePowerToFromVariable,
    ::Type{<:PSY.DCBranch},
    ::AbstractDCLineFormulation,
) = 1.0

get_variable_lower_bound(::FlowActivePowerVariable, d::PSY.DCBranch, ::HVDCUnbounded) =
    nothing

get_variable_upper_bound(::FlowActivePowerVariable, d::PSY.DCBranch, ::HVDCUnbounded) =
    nothing

get_variable_lower_bound(
    ::FlowActivePowerVariable,
    d::PSY.DCBranch,
    ::AbstractDCLineFormulation,
) = max(PSY.get_active_power_limits_from(d).min, PSY.get_active_power_limits_to(d).min)

get_variable_upper_bound(
    ::FlowActivePowerVariable,
    d::PSY.DCBranch,
    ::AbstractDCLineFormulation,
) = min(PSY.get_active_power_limits_from(d).max, PSY.get_active_power_limits_to(d).max)

#! format: on

function get_default_time_series_names(
    ::Type{U},
    ::Type{V},
) where {U <: PSY.DCBranch, V <: AbstractDCLineFormulation}
    return Dict{Type{<:TimeSeriesParameter}, String}()
end

function get_default_attributes(
    ::Type{U},
    ::Type{V},
) where {U <: PSY.DCBranch, V <: AbstractDCLineFormulation}
    return Dict{String, Any}()
end

get_initial_conditions_device_model(
    ::OperationModel,
    ::DeviceModel{T, <:AbstractDCLineFormulation},
) where {T <: PSY.HVDCLine} = DeviceModel(T, HVDCDispatch)

#################################### Rate Limits Constraints ##################################################
function branch_rate_bounds!(
    container::OptimizationContainer,
    devices::IS.FlattenIteratorWrapper{B},
    ::DeviceModel{B, HVDCLossless},
    ::Type{<:PM.AbstractPowerModel},
) where {B <: PSY.HVDCLine}
    var = get_variable(container, FlowActivePowerVariable(), B)
    for d in devices
        name = PSY.get_name(d)
        for t in get_time_steps(container)
            _var = var[name, t]
            max_val = max(
                PSY.get_active_power_limits_to(d).max,
                PSY.get_active_power_limits_from(d).max,
            )
            min_val = min(
                PSY.get_active_power_limits_to(d).min,
                PSY.get_active_power_limits_from(d).min,
            )
            JuMP.set_upper_bound(_var, max_val)
            JuMP.set_lower_bound(_var, min_val)
        end
    end
    return
end

function branch_rate_bounds!(
    container::OptimizationContainer,
    devices::IS.FlattenIteratorWrapper{B},
    ::DeviceModel{B, <:AbstractDCLineFormulation},
    ::Type{<:PM.AbstractPowerModel},
) where {B <: PSY.HVDCLine}
    vars = [
        get_variable(container, FlowActivePowerVariable(), B),
        get_variable(container, HVDCTotalPowerDeliveredVariable(), B),
    ]
    for d in devices
        name = PSY.get_name(d)
        for t in get_time_steps(container)
            JuMP.set_upper_bound(vars[1][name, t], PSY.get_active_power_limits_from(d).max)
            JuMP.set_lower_bound(vars[1][name, t], PSY.get_active_power_limits_from(d).min)
            JuMP.set_upper_bound(vars[2][name, t], PSY.get_active_power_limits_to(d).max)
            JuMP.set_lower_bound(vars[2][name, t], PSY.get_reactive_power_limits_to(d).min)
        end
    end
    return
end

function add_constraints!(
    container::OptimizationContainer,
    cons_type::Type{FlowRateConstraint},
    devices::IS.FlattenIteratorWrapper{B},
    ::DeviceModel{B, HVDCLossless},
    ::Type{<:PM.AbstractDCPModel},
) where {B <: PSY.DCBranch}
    var = get_variable(container, FlowActivePowerVariable(), B)
    time_steps = get_time_steps(container)
    names = [PSY.get_name(d) for d in devices]
    constraint = add_constraints_container!(container, cons_type(), B, names, time_steps)
    for t in time_steps, d in devices
        min_rate = max(
            PSY.get_active_power_limits_from(d).min,
            PSY.get_active_power_limits_to(d).min,
        )
        max_rate = min(
            PSY.get_active_power_limits_from(d).max,
            PSY.get_active_power_limits_to(d).max,
        )
        constraint[PSY.get_name(d), t] = JuMP.@constraint(
            container.JuMPmodel,
            min_rate <= var[PSY.get_name(d), t] <= max_rate
        )
    end
    return
end

add_constraints!(
    ::OptimizationContainer,
    ::Type{<:Union{FlowRateConstraintFromTo, FlowRateConstraintToFrom, FlowRateConstraint}},
    ::IS.FlattenIteratorWrapper{<:PSY.DCBranch},
    ::DeviceModel{<:PSY.DCBranch, HVDCUnbounded},
    ::Type{<:PM.AbstractPowerModel},
) = nothing

function add_constraints!(
    container::OptimizationContainer,
    cons_type::Type{
        <:Union{FlowRateConstraintFromTo, FlowRateConstraintToFrom, FlowRateConstraint},
    },
    devices::IS.FlattenIteratorWrapper{B},
    ::DeviceModel{B, <:AbstractDCLineFormulation},
    ::Type{<:PM.AbstractPowerModel},
) where {B <: PSY.DCBranch}
    time_steps = get_time_steps(container)
    names = [PSY.get_name(d) for d in devices]

    var = get_variable(container, FlowActivePowerVariable(), B)
    constraint = add_constraints_container!(container, cons_type(), B, names, time_steps)
    for t in time_steps, d in devices
        min_rate = max(
            PSY.get_active_power_limits_from(d).min,
            PSY.get_active_power_limits_to(d).min,
        )
        max_rate = min(
            PSY.get_active_power_limits_from(d).max,
            PSY.get_active_power_limits_to(d).max,
        )
        constraint[PSY.get_name(d), t] = JuMP.@constraint(
            container.JuMPmodel,
            min_rate <= var[PSY.get_name(d), t] <= max_rate
        )
    end
    return
end

function add_constraints!(
    container::OptimizationContainer,
    cons_type::Type{HVDCPowerBalance},
    devices::IS.FlattenIteratorWrapper{B},
    ::DeviceModel{B, <:AbstractDCLineFormulation},
    ::Type{<:PM.AbstractPowerModel},
) where {B <: PSY.DCBranch}
    time_steps = get_time_steps(container)
    names = [PSY.get_name(d) for d in devices]

    delivered_power_var = get_variable(container, HVDCTotalPowerDeliveredVariable(), B)
    flow_var = get_variable(container, FlowActivePowerVariable(), B)
    constraint = add_constraints_container!(container, cons_type(), B, names, time_steps)
    for t in get_time_steps(container), d in devices
        constraint[PSY.get_name(d), t] = JuMP.@constraint(
            container.JuMPmodel,
            delivered_power_var[PSY.get_name(d), t] ==
            -PSY.get_loss(d).l1 * flow_var[PSY.get_name(d), t] - PSY.get_loss(d).l0,
        )
    end
    return
end


using PyCall,DataFrames,RCall
unshift!(PyVector(pyimport("sys")["path"]), "")
@pyimport ncdata
len = length
filename = "edwards11.nc"

data = ncdata.get(filename)
specs = names!(DataFrame(data["spec"]),[Symbol(i)for i in data["sc"]])
rates = names!(DataFrame(data["rate"]),[Symbol(i)for i in data["rc"]])
rates = rates[:,[Symbol(i) for i in filter(x->!ismatch(r"EMISS|DUMMY",x),data["rc"])]]
specs = specs[:,[Symbol(i) for i in filter(x->!ismatch(r"EMISS|DUMMY",x),data["sc"])]]
rates = rates[7:len(rates)]
rates = rates[:,[Symbol(i) for i in filter(x->sum(rates[x])>0,names(rates))]]
joint = "\n"*(join(names(rates),"\n"))*"\n"

products= [split(i[4:len(i)],"+") for i in matchall(r"-->([A-z0-9+]*)",joint)]
reactants = [split(i[2:len(i)-1],"+") for i in matchall(r"\n([A-z0-9+]{1,60})([-->]{0,1})",joint)]

node_set = Set([ i[2:len(i)-1] for i in matchall(r"[\n+>-]{1}[A-z0-9]+[\n+>-]{1}",joint) ])
nodes = [i for i in node_set]
dict = Dict( [[i,x] for (x,i) in enumerate(nodes)])

flux = []
edges = []
for i in 1:length(reactants)
  rcol=[]
  #println(i)
  for j in reactants[i]
    coeff = match( r"([\d\s\.]*)(\D[\d\D]*)", j)
    dummy = specs[Symbol(coeff[2])]
    try
      push!(rcol,parse(Float32,coeff[1]))
    catch
      push!(rcol,1.0f0)
    end

    prod = 1
    for k in rcol
      prod*=k
    end

    push!(flux, prod*rates[i])
    # make edges array
    for l in products[i]
        push!(edges,[match( r"([\d\s\.]*)(\D[\d\D]*)", j)[2],match( r"([\d\s\.]*)(\D[\d\D]*)", l)[2],i])
    end

  end
end

normalise(x) = x/norm(x)

function sorteqn(x,delim)
    r,p = split(string(x),delim)
    r = join(sort(split(r,'+')),'+')
    p = join(sort(split(p,'+')),'+')
    return join([r,p],delim)
end

rateeqn = [Symbol(sorteqn(i,"-->")) for i in names(rates)]

grouprecorded = [dict[Symbol(sorteqn(i,"-->"))] for i in names(rates)]

################READ CATEGORIES##################
dict = Dict()
open("edgecat.json", "r") do f
    global dict
    dict=JSON.parse(readstring(f))  # parse and transform data
end

dict = Dict(Symbol( sorteqn( replace(key,"=","-->") , "-->" )) => value for (key, value) in dict)

reactiontypes= [try dict[i]; catch err; "missing" end  for i in rateeqn]

filter(i->ismatch(r".*NC7H16.*", i), [string(i) for i in keys(dict)])

smiles =DataFrame(readtable("smiles_mined.csv"))

t = 325

links = filter(i -> flux[i][t]>0 , 1:len(flux))
tflux = [log10(flux[i][t]) for i in links]
weight =  tflux - minimum(tflux)
weight = 1 - weight/maximum(weight)

#1+normalise(tflux)

newflux = [flux[i][t] for i in 1:len(flux)]
counter = 0
for i in links
  newflux[i] = weight[counter+=1]
end

edge = filter(i -> newflux[i[3]]>0 , edges)
source = [i[1] for i in edge]
target=[i[2] for i in edge]
weighted = [newflux[i[3]]+0.0001 for i in edge]
grouping  = [reactiontypes[i[3]] for i in edge]

JSONdata =  [(i[1],i[2],newflux[i[3]],reactiontypes[i[3]]) for i  in edge]


import Knet.Ops20: elu, relu, selu, sigm, eluback, reluback, seluback, sigmback
import Base.Broadcast: broadcasted
import Knet
using Knet.KnetArrays: KnetArray, DevArray, Bcasted
using Knet.LibKnet8: @knet8
using CUDA: CuArray, CuPtr
using AutoGrad: AutoGrad, @primitive

# Specialize tanh gradient for DevArrays here. The others have been declared in Ops20 as generic gradients.
tanhback(dyi::T,yi::T) where {T<:Number} = dyi*(T(1)-yi*yi)
@primitive tanh(x::DevArray),dy,y tanhback.(dy,y)
@primitive tanhback(dy,y),ddx  ddx.*(1 .- y.*y)  ddx.*(-2 .* dy.*y)

for (R,P) in ((KnetArray,Ptr), (CuArray,CuPtr)), T in (Float32,Float64); S = sizeof(T) * 8
    for f in ("elu","relu","selu","sigm")
        J, F = Symbol(f), "$(f)_$S"; M = which(@__MODULE__,J)
        @eval begin
            function broadcasted(::typeof($J),x::$R{$T})
                y = similar(x)
                @knet8($F,(Cint,$P{$T},$P{$T}),length(y),x,y)
                return y
            end
            # Bcasted methods -- only needed for KnetArray
            ($M).$J(x::Bcasted{<:$R{$T}}) = broadcasted($J, x.value) |> Bcasted
            broadcasted(::typeof($J),x::Bcasted{<:$R{$T}}) = broadcasted($J, x.value) |> Bcasted
        end
    end
    for f in ("eluback","reluback","seluback","sigmback","tanhback")
        J, F = Symbol(f), "$(f)_$(S)_11"; M = which(@__MODULE__,J)
        @eval begin
            function broadcasted(::typeof($J),x::$R{$T},y::$R{$T})
                z = similar(x)
                @knet8($F,(Cint,$P{$T},$P{$T},$P{$T}),length(z),x,y,z)
                return z
            end
            # Bcasted methods -- only needed for KnetArray
            ($M).$J(x::Bcasted{<:$R{$T}}, y::Bcasted{<:$R{$T}}) = broadcasted($J, x.value, y.value) |> Bcasted
            ($M).$J(x::$R{$T}, y::Bcasted{<:$R{$T}}) = broadcasted($J, x, y.value) |> Bcasted
            ($M).$J(x::Bcasted{<:$R{$T}}, y::$R{$T}) = broadcasted($J, x.value, y) |> Bcasted
            broadcasted(::typeof($J),x::Bcasted{<:$R{$T}}, y::Bcasted{<:$R{$T}}) = broadcasted($J, x.value, y.value) |> Bcasted
            broadcasted(::typeof($J),x::$R{$T}, y::Bcasted{<:$R{$T}}) = broadcasted($J, x, y.value) |> Bcasted
            broadcasted(::typeof($J),x::Bcasted{<:$R{$T}}, y::$R{$T}) = broadcasted($J, x.value, y) |> Bcasted
        end
    end
end


"""
NNMPCdriver
Driver that randomly changes lanes and speeds.

# Constructors
BafflingDriver(timestep::Float64;mlon::LaneFollowingDriver=IntelligentDriverModel(), mlat::LateralDriverModel=ProportionalLaneTracker(), mlane::LaneChangeModel=RandLaneChanger(timestep),rec::SceneRecord = SceneRecord(1, timestep))

# Fields
- `rec::SceneRecord` A record that will hold the resulting simulation results
- `mlon::LaneFollowingDriver = IntelligentDriverModel()` Longitudinal driving model
- `mlat::LateralDriverModel = ProportionalLaneTracker()` Lateral driving model
- `mlane::LaneChangeModel =RandLaneChanger` Lane change model (randomly)
"""

mutable struct NNMPCdriver <: DriverModel{AccelSteeringAngle}
    rec::SceneRecord
    n_ticks::Integer
    mlane::LaneChangeModel
    δ_max::Float64
    δ_min::Float64
    a_max::Float64
    a_min::Float64
    N_sim::Float64
    T::Float64
    ΔT::Float64
    T_obs_given::Float64
    T_obs::Float64
    T_pred::Float64
    Δa_max::Float64
    Δδ_max::Float64
    λ_v::Float64
    λ_div::Float64
    λ_δ::Float64
    λ_a::Float64
    λ_Δδ::Float64
    λ_Δa::Float64
    v_des::Float64
    a::Float64
    δ::Float64
    height::Float64
    width::Float64
    models::Dict{Int, DriverModel}
    thred_safety::Float64
    isDebugMode::Bool
    pred_model_type::AbstractString  # {"sgan", "perfect"}
    predictor::Any
    sgan_path::AbstractString
    sgan_model_path::AbstractString
    lane_change_action::LaneChangeChoice
    temp_frame_count::Integer
    veh_shape_type::AbstractString

    function NNMPCdriver(
        timestep::Float64;
        n_ticks::Integer = 5,
        mlane::LaneChangeModel= BafflingLaneChanger(timestep),
        rec::SceneRecord = SceneRecord(n_ticks, timestep),
        δ_max::Float64 = 0.6,
        δ_min::Float64 = -0.4,
        a_max::Float64 = 3.5,
        a_min::Float64 = -4.0,
        N_sim::Float64 = 200.0,
        T::Float64 = 0.2,
        ΔT::Float64 = timestep,
        T_obs::Float64 = 1.0,
        T_obs_given::Float64 = T_obs,
        T_pred::Float64 = 1.0,
        Δa_max::Float64 = 2.0,
        Δδ_max::Float64 = 0.4,
        λ_v::Float64 = 10.0,
        λ_div::Float64 = 10e2+0.0,
        λ_δ::Float64 = 1.0,
        λ_a::Float64 = 1.0,
        λ_Δδ::Float64 = 3.0,
        λ_Δa::Float64 = 1.0,
        height::Float64 = 4.0,
        width::Float64 = 1.8,
        v_des::Float64 = 30.0,
        thred_safety::Float64 = 0.3,
        isDebugMode::Bool = false,
        pred_model_type::AbstractString = "sgan",
        sgan_path::AbstractString = "/home/sbae/NNMPC.jl/python/sgan",
        sgan_model_path::AbstractString = "/home/sbae/NNMPC.jl/python/sgan/models/sgan-models/eth_8_model.pt",
        lane_change_action::LaneChangeChoice = LaneChangeChoice(DIR_RIGHT),
        veh_shape_type::AbstractString = "circle")

        retval = new()

        pushfirst!(PyVector(pyimport("sys")."path"),sgan_path)
        sganPredictor = pyimport("sgan.predictor")
        predictor = sganPredictor.Predictor(sgan_model_path)

        retval.rec = rec
        retval.mlane = mlane
        retval.δ_max = δ_max
        retval.δ_min = δ_min
        retval.a_max = a_max
        retval.a_min = a_min
        retval.N_sim = trunc(Int,N_sim)
        retval.T = T
        retval.ΔT = ΔT
        retval.T_obs = T_obs
        retval.T_obs_given = T_obs_given
        retval.T_pred = T_pred
        retval.Δa_max = Δa_max
        retval.Δδ_max = Δδ_max
        retval.λ_v = λ_v
        retval.λ_div = λ_div
        retval.λ_δ = λ_δ
        retval.λ_a = λ_a
        retval.λ_Δδ = λ_Δδ
        retval.λ_Δa = λ_Δa
        retval.v_des = v_des
        retval.a = a_min
        retval.δ = 0
        retval.height = height
        retval.width = width
        retval.models = Dict{Int64,DriverModel}()
        retval.thred_safety = thred_safety
        retval.isDebugMode = isDebugMode
        retval.predictor = predictor
        retval.pred_model_type = pred_model_type
        retval.lane_change_action = lane_change_action
        retval.temp_frame_count = 0
        retval.veh_shape_type = veh_shape_type

        retval
    end
end

get_name(::NNMPCdriver) = "NNMPCdriver"
function AutomotiveDrivingModels.set_desired_speed!(model::NNMPCdriver, v_des::Float64)
    model.v_des = v_des
    model
end

function set_other_models!(model::NNMPCdriver, models::Dict{Int, DriverModel})
    model.models = models
    model
end

"""
    compute_longitudinal_acceleration(model::NNMPCdriver, scene::Scene, roadway::Roadway, egoid::Int)

"""
function compute_longitudinal_acceleration(model::NNMPCdriver, scene::Scene, roadway::Roadway, ind_ego::Int)
    # local variables -- from IDM
    k_spd = 1.0 # proportional constant for speed tracking when in freeflow [s⁻¹]
    δ = 4.0 # acceleration exponent [-]
    T = 1.5 # desired time headway [s]
    # v_des = 29.0 # desired speed [m/s]
    s_min = 5.0 # minimum acceptable gap [m]
    a_max = 3.0 # maximum acceleration ability [m/s²]
    d_cmf = 2.0 # comfortable deceleration [m/s²] (positive)
    d_max = 9.0 # maximum deceleration [m/s²] (positive)
    ΔT = 0.2 # timestep to simulate [s]
    v_offset = 5.0 # offset from the desired speed [m/s]

    fore = get_neighbor_fore_along_lane(scene, ind_ego, roadway, VehicleTargetPointFront(), VehicleTargetPointRear(), VehicleTargetPointFront())
    v_ego = scene[ind_ego].state.v
    if fore.ind != nothing
        headway, v_oth = fore.Δs, scene[fore.ind].state.v
    else
        headway, v_oth = NaN, NaN
    end
    if !isnan(v_oth)
        Δv = v_oth - v_ego
        s_des = s_min + v_ego*T - v_ego*Δv / (2*sqrt(model.a_max*abs(model.a_min)))
        v_ratio = model.v_des > 0.0 ? (v_ego/model.v_des) : 1.0
        a = model.a_max * (1.0 - v_ratio^δ - (s_des/headway)^2)
    else
        # no lead vehicle, just drive to match desired speed
        Δv = model.v_des - v_ego
        a = Δv*k_spd # predicted accel to match target speed
    end

    if v_ego + model.ΔT * a < 0
        a = max(model.a_min, -v_ego/model.ΔT)
    end

    return a
end

"""
    observe!(model::NNMPCdriver, scene::Scene, roadway::Roadway, egoid::Int)

(1) randomly generates a sequence of control inputs (acceleration and steering angle)
(2) evaluates Social GAN to predict motions of neighboring vehicles within a range
(3) checks the safety constraint at each time step
(4) evalautes the cumulative cost over the receding time horizon T
(5) finds the optimal pair of control inputs
"""
function AutomotiveDrivingModels.observe!(model::NNMPCdriver, scene::Scene, roadway::Roadway, egoid::Int)
    # δ_max,δ_min,a_max,a_min,N_sim,T,ΔT,T_obs,T_pred,Δa_max,Δδ_max=0.6,-0.4,3.0,-3.0,1000,4,0.5,2,2,5,0.6
    ind_ego = findfirst(egoid, scene)

    update!(model.rec, scene)
    observe!(model.mlane, scene, roadway, egoid) # receive lane change action from higher level decision maker
    # global y_ref = DEFAULT_LANE_WIDTH * lane_change_action.dir
    # global y_ref = -1*DEFAULT_LANE_WIDTH

    # check if we have enough dataset to use SGAN
    isSganDataReady = model.rec.nframes*model.ΔT >= model.T_obs_given
    model.isDebugMode ? println("isSganDataReady: ", isSganDataReady) : nothing

    # if sgan data is not ready, then it just keeps the lane
    if !isSganDataReady
        @warn "Not enough records to run SGAN. IDM is running."
        model.a = compute_longitudinal_acceleration(model, scene, roadway, ind_ego)
        model.δ = 0
        return model
    end

    # ============== TEMP: test================
    # if model.lane_change_action.dir == DIR_MIDDLE
    #     if model.temp_frame_count == 20
    #         model.temp_frame_count = 0
    #
    #         left_lane_exists = convert(Float64, get(N_LANE_LEFT, model.rec, roadway, ind_ego)) > 0
    #         right_lane_exists = convert(Float64, get(N_LANE_RIGHT, model.rec, roadway, ind_ego)) > 0
    #
    #         if left_lane_exists
    #             model.lane_change_action = LaneChangeChoice(DIR_LEFT)
    #         elseif right_lane_exists
    #             model.lane_change_action = LaneChangeChoice(DIR_RIGHT)
    #         end
    #     else
    #         model.temp_frame_count += 1
    #     end
    #
    # end
    # =========================================

    # handle NaN lane offset. later, this must be addressed in the higher level
    lane_offset = get_lane_offset(model.lane_change_action, model.rec, roadway, ind_ego)
    if isnan(lane_offset)
        @warn "NaN lane offset. Fixed to keeping lane command."
        model.lane_change_action = LaneChangeChoice(DIR_MIDDLE)
        model.temp_frame_count = 0
    end

    lane_offset = get_lane_offset(model.lane_change_action, model.rec, roadway, ind_ego)
    if model.lane_change_action.dir == DIR_MIDDLE
        println("keeping lane")
    elseif model.lane_change_action.dir == DIR_LEFT
        println("turning to left lane")
    else
        println("turning to right lane")
    end

    # local variables
    N_receding = trunc(Int, round(model.T/model.ΔT))
    N_sgan_obs = trunc(Int, round(model.T_obs/model.ΔT))
    N_sgan_pred = trunc(Int, round(model.T_pred/model.ΔT))

    # generate sequences of control inputs
    # println("generates random control sequences")
    # @time begin
        a_seq, δ_seq = gen_rand_inputs(model)
    # end

    cost_min = Inf
    n_opt = -1            # index of the optimal sequence of the controls
    model.a = model.a_min # maximum braking by default
    model.δ = 0           # zero steering by default
    # initialize the vehicle object for the other vehicles
    ind_near_vehs = get_indices_near_vehs(ind_ego, scene, 15.0)
    println("number of vehicles in the range: ", length(ind_near_vehs))

    println("find optimal contorls")
    @time begin
        for n = 1 : model.N_sim
            model.isDebugMode ? println("n: ",n) : nothing

            # initialize the vehicle object
            ego_veh = scene[ind_ego] # Entity{VehicleStateBuffer,BicycleModel,Int}

            # initialize the record scene queue
            rec′ = SceneRecord(model.rec.nframes + N_receding, model.ΔT)
            for k = 1 : N_sgan_obs
                update!(rec′,model.rec[-N_sgan_obs + k])
            end

            # initialize lane change choice
            # lane_change_action′= LaneChangeChoice(lane_change_action.dir)
            lane_offset′ = lane_offset

            # randomly sample the sequence of control
            a_ = a_seq[:,rand(1:size(a_seq,2))]
            δ_ = δ_seq[:,rand(1:size(δ_seq,2))]

            # initialize the scene with the recent one
            scene′= Scene()
            for i = 1 : length(scene)
                push!(scene′,scene[i])
            end

            # evaluate the control sequences
            cost_n = 0
            for ℓ = 1 : N_receding
                model.isDebugMode ? println("ℓ: ", ℓ) : nothing
                # (1) predict other drivers' motion (SGAN) (for the next step)
                if length(ind_near_vehs) >= 1
                    if model.pred_model_type == "perfect"
                        scene′ = propagate_other_vehs(Any, model, ind_ego, ind_near_vehs, scene′, roadway, rec′)
                    elseif model.pred_model_type == "sgan"
                        scene′ = predict(model, ind_ego, ind_near_vehs, scene′, roadway, rec′)
                    else
                        error("ERROR: prediction model type error. It must be either {perfect} (benchmark) or {sgan}")
                    end
                end

                # (2) update the vehicle states with a,δ
                a,δ = a_[ℓ],δ_[ℓ]
                model.isDebugMode ? println("try a = ",a, ", δ = ", δ) : nothing
                state′ = propagate(ego_veh, AccelSteeringAngle(a,δ), roadway, model.ΔT)
                ego_veh = Entity(state′, ego_veh.def, ego_veh.id)
                scene′[ind_ego] = ego_veh
                update!(rec′,scene′)

                # (3) check the feasibility (break if violated)
                # println("check safety")
                # @time begin
                if length(ind_near_vehs) >= 1
                    if !isSafe(model, ind_ego, ind_near_vehs, scene′)
                        cost_n = Inf
                        break
                    end
                end
                # end

                # (4) compute the cumulative costs
                a_prev = ℓ == 1 ? 0 : a_[ℓ-1]
                δ_prev = ℓ == 1 ? 0 : δ_[ℓ-1]
                # if rec′[-1][ind_ego].state.posF.t * rec′[0][ind_ego].state.posF.t < 0 # lane changed
                #     lane_change_action′ = LaneChangeChoice(DIR_MIDDLE)
                # end
                lane_offset′ += ego_veh.state.v*sin(ego_veh.state.posF.ϕ)*model.ΔT
                cost_n += ( model.λ_v*(ego_veh.state.v-model.v_des)^2
                        + model.λ_div*(abs(lane_offset′)+1)^2
                        + model.λ_δ*(δ)^2
                        + model.λ_a*(a)^2
                        + model.λ_Δδ*(δ-δ_prev)^2
                        + model.λ_Δa*(a-a_prev)^2)
                        # + model.λ_div*(ego_veh.state.posG.y-y_ref)^2
            end

            if cost_n < cost_min
                # update with the optimal control sequence
                cost_min = cost_n
                model.a = a_[1]
                model.δ = δ_[1]
                n_opt = n
            end
        end
    end

    model.isDebugMode ? println("optimal controls (n = ", n_opt,") : a = ",model.a, ", δ = ", model.δ) : nothing
    if n_opt == -1
        @warn "No feasible solution found. Maximum braking applied."
    end

    # prevent the vehicle going backward
    if scene[ind_ego].state.v + model.ΔT*model.a < 0
        model.a = max(-scene[ind_ego].state.v/model.ΔT, model.a_min)
    end
    println("lane_offset: ",get_lane_offset(model.lane_change_action, model.rec, roadway, ind_ego))
    println("a: ", model.a)
    println("δ: ", model.δ)

    # check if the ego vehicle changes the lane with the next action
    state′ = propagate(scene[ind_ego], AccelSteeringAngle(model.a,model.δ), roadway, model.ΔT)
    if abs(lane_offset + state′.v*sin(state′.posF.ϕ)*model.ΔT) <= DEFAULT_LANE_WIDTH/2
        model.lane_change_action = LaneChangeChoice(DIR_MIDDLE)
    end

    model
end

"""
    get_indices_near_vehs(model::NNMPCdriver, scene::Scene; range::Float64 = 5 [m])

find indices of near vehicles around the ego vehicle
"""
function get_indices_near_vehs(ind_ego::Int, scene::Scene, range::Float64)
    ego_veh = scene[ind_ego]
    x = ego_veh.state.posG.x
    y = ego_veh.state.posG.y
    ind_near_vehs = []
    for i in 1 : length(scene)
        if i == ind_ego
            continue
        end
        xᵢ = scene[i].state.posG.x
        yᵢ = scene[i].state.posG.y
        if sqrt((x-xᵢ)^2+(y-yᵢ)^2) <= range
            append!(ind_near_vehs, i)
        end
    end
    return ind_near_vehs
end

"""
    gen_rand_inputs(model::NNMPCdriver)

returns randomly generated sequence of a and δ, both of which satisfies the
maximum rate constraints (jerk, steering rate)
"""
function gen_rand_inputs(model::NNMPCdriver)
    n_row = trunc(Int,model.T/model.ΔT)
    n_col = trunc(Int,model.N_sim)

    # generate random sequences of action pair (a, δ)
    a_seq_raw = rand(Random.MersenneTwister(),model.a_min:0.01:model.a_max, (n_row,n_col)) #0.002 sec
    δ_seq_raw = rand(Random.MersenneTwister(),model.δ_min:0.01:model.δ_max, (n_row,n_col))

    # filter out the infeasible jerks
    a_seq_check = [zeros(1,n_col);a_seq_raw]
    a_diff = abs.(reshape([diff(reshape(a_seq_check, length(a_seq_check)));0],n_row+1,n_col)[1:end-1,1:end])
    bitarr = sum(a_diff .< model.Δa_max,dims=1) .== n_row
    a_seq = []
    for i in 1:length(bitarr)
        if bitarr[i] == true
            append!(a_seq, a_seq_raw[1:end,i])
        end
    end
    a_seq = reshape(a_seq, n_row, trunc(Int,length(a_seq)/n_row))

    # filter out the infeasible steering rate
    δ_seq_check = [zeros(1,n_col);δ_seq_raw]
    δ_diff = abs.(reshape([diff(reshape(δ_seq_check, length(δ_seq_check)));0],n_row+1,n_col)[1:end-1,1:end])
    bitarr = sum(δ_diff .< model.Δδ_max,dims=1) .== n_row
    δ_seq = []
    for i in 1:length(bitarr)
        if bitarr[i] == true
            append!(δ_seq, δ_seq_raw[1:end,i])
        end
    end
    δ_seq = reshape(δ_seq, n_row, trunc(Int,length(δ_seq)/n_row))

    return a_seq, δ_seq
end

"""
    function propagate(model::NNMPCdriver, ind_near_vehs::Vector{Any}, rec::SceneRecord)
returns positions of the neighboring vehicles, based on their Driver Model. This is equivalent to the SGAN that is perfectly trained.
"""
function propagate_other_vehs(::Type{A}, model::NNMPCdriver, ind_ego::Int, ind_near_vehs::Vector{Any}, scene::Scene, roadway::Roadway, rec::SceneRecord) where {A}
    models = Dict{Int, DriverModel}()
    for i in 1 : length(model.models)
        if i == ind_ego
            models[i] = BafflingDriver(rec.timestep)
        else
            models[i] = model.models[i]
        end
    end
    actions = Array{A}(undef, length(scene))
    get_actions!(actions, scene, roadway, models)
    tick!(scene, roadway, actions, rec.timestep)

    return scene
end

"""
    predict(model::NNMPCdriver, ind_ego::Integer, ind_near_vehs::Vector{Any}, scene::Scene, roadway::Roadway, rec::SceneRecord)
predicts motions of all vehicles within a range, using Social GAN.
"""
function predict(model::NNMPCdriver, ind_ego::Integer, ind_near_vehs::Vector{Any}, scene::Scene, roadway::Roadway, rec::SceneRecord)
    # arrange input for SGAN
    # example:
    # obs_traj = [  [ 1,  1.000e+00,  8.460e+00,  3.590e+00],
    #           [ 1,  2.000e+00,  1.364e+01,  5.800e+00],
    #           [ 2,  1.000e+00,  9.570e+00,  3.790e+00],
    #           [ 2,  2.000e+00,  1.364e+01,  5.800e+00],
    #           [ 3,  1.000e+00,  1.067e+01,  3.990e+00],
    #           [ 3,  2.000e+00,  1.364e+01,  5.800e+00],
    #           [ 4,  1.000e+00,  1.173e+01,  4.320e+00],
    #           [ 4,  2.000e+00,  1.209e+01,  5.750e+00],
    #           [ 5,  1.000e+00,  1.281e+01,  4.610e+00],
    #           [ 5,  2.000e+00,  1.137e+01,  5.800e+00],
    #           [ 6,  1.000e+00,  1.281e+01,  4.610e+00],
    #           [ 6,  2.000e+00,  1.031e+01,  5.970e+00],
    #           [ 7,  1.000e+00,  1.194e+01,  6.770e+00],
    #           [ 7,  2.000e+00,  9.570e+00,  6.240e+00],
    #           [ 8,  1.000e+00,  1.103e+01,  6.840e+00],
    #           [ 8,  2.000e+00,  8.730e+00,  6.340e+00]]

    N_obs = trunc(Int, round(model.T_obs/model.ΔT))
    N_vehs = length(ind_near_vehs)

    # bulid the input matrix for sgan
    obs_traj = []
    for seq in 1 : N_obs # sequence
        for i in 1 : N_vehs + 1 # vehicles
            ind = i == 1 ? ind_ego : ind_near_vehs[i-1]
            veh = rec[0-(N_obs-1)+(seq-1)][ind]
            l = [seq, i, veh.state.posG.x, veh.state.posG.y]
            obs_traj = [obs_traj;[l]]
        end
    end

    if model.isDebugMode
        println("abs_traj:")
        for i in 1:length(obs_traj)
              println(obs_traj[i])
        end
    end

    # call SGAN and return the predicted positions
    next_pred_traj = model.predictor.predict(obs_traj)
    if model.isDebugMode
        println("pred_traj:")
        for i in 1:size(next_pred_traj,1)
            ind = i == 1 ? ind_ego : ind_near_vehs[i-1]
            if i == 1
                println("ego vehicle: ", next_pred_traj[i,:])
            else
                println("vehicle ",ind, ": ", next_pred_traj[i,:])
            end
        end
    end

    # update the scene with the predicted positions
    for i in 1 : N_vehs + 1
        ind = i == 1 ? ind_ego : ind_near_vehs[i-1]
        veh = scene[ind]

        x′ = next_pred_traj[i,1]; x_diff = x′-veh.state.posG.x
        y′ = next_pred_traj[i,2]; y_diff = y′-veh.state.posG.y
        θ′ = atan(y_diff/x_diff)
        v′ = sqrt(x_diff^2+y_diff^2)/model.ΔT

        posG = VecSE2(x′, y′, θ′)
        state′  = VehicleState(posG, roadway, v′)

        scene[ind] = Entity(state′, veh.def, veh.id)
    end

    return scene
end

"""
    isSafe(ind_ego::Int, ind_others::Array{Int}, scene::Scene)
check collision with any of adjacent vehicles
"""
function isSafe(model::NNMPCdriver, ind_ego::Int, ind_others::Vector{Any}, scene::Scene)
    ego_veh = scene[ind_ego]
    ϵ = model.thred_safety
    # println("ind_ego: ", ind_ego)
    for i in ind_others
        dist = min_distance(model, ego_veh, scene[i])
        model.isDebugMode ? println("distance to vehicle ", i, ": ", dist) : nothing
        if dist > ϵ
            continue
        else
            model.isDebugMode ? println("   COLISION DETECTED. actions discarded") : nothing
            return false
        end
    end
    return true
end

"""
    min_distance(ego::Entity{VehicleState, D, Int},
                other::Entity{VehicleState, D, Int}) where {D}

returns distance between ego vehicle and an adjacent vehicle
"""
function min_distance(model::NNMPCdriver,
                    ego::Entity{VehicleState, D, Int},
                    other::Entity{VehicleState, D, Int}) where {D}
    # ego: ego vehicle object
    # other: other vehicle object
    x,y,θ = ego.state.posG.x, ego.state.posG.y, ego.state.posG.θ
    xᵢ,yᵢ,θᵢ = other.state.posG.x, other.state.posG.y, other.state.posG.θ
    h = model.height/2
    w = model.width/2

    if model.veh_shape_type == "ellipsoid"
        prob = Model(with_optimizer(Ipopt.Optimizer, print_level=0))
        @variable(prob, x′); @variable(prob, y′)
        @variable(prob, xᵢ′); @variable(prob, yᵢ′)
        @objective(prob, MOI.MIN_SENSE, (x′-xᵢ′)^2+(y′-yᵢ′)^2)
        @constraint(prob, ((x′-x)*cos(θ) + (y′-y)*sin(θ))^2/(h^2) + ((x′-x)*sin(θ) + (y′-y)*cos(θ))^2/(w^2) <= 1.0)
        @constraint(prob, ((xᵢ′-xᵢ)*cos(θᵢ) + (yᵢ′-yᵢ)*sin(θᵢ))^2/(h^2) + ((xᵢ′-xᵢ)*sin(θᵢ) + (yᵢ′-yᵢ)*cos(θᵢ))^2/(w^2) <= 1.0)
        optimize!(prob)
        return sqrt(max(objective_value(prob), 0))
    elseif model.veh_shape_type == "circle"
        r = w
        min_dis = Inf
        for i in [-1,0,1]
            for j in [-1,0,1]
                dis = sqrt(
                          ( (x + i*r*cos(θ)) - (xᵢ+ j*r*cos(θᵢ)) )^2
                         +( (y + i*r*sin(θ)) - (yᵢ+ j*r*sin(θᵢ)) )^2
                         )     - 2*r
                min_dis = min(dis, min_dis)
            end
        end
        return max(min_dis, 0)
    else
        error("ERROR: vehicle shape type error. It must be either circle or ellipsoid")
    end
end

Base.rand(model::NNMPCdriver) = AccelSteeringAngle(model.a,model.δ)