# ------- R 2 Julia Conversion ---------------
#function saveModels(d::Dict,fname::String)
#    j = JSON.json([ d])
#    open(fname,"w") do f
#        write(f,j)
#    end
#end

function expandM(di::Dict)
    d2=deepcopy(di)
    if d2[:modelName]==:occ
        d2[:dist] = Poisson()
        d2[:lnk] = LogLink()
        d2[:Buyer_Pos_P1_is1] = true
    end
    if d2[:modelName]==:dolocc
        d2[:dist] = Gamma()
        d2[:lnk] = LogLink()
        d2[:Buyer_Pos_P1_is1] = true
    end
    if d2[:modelName]==:pen
        d2[:dist] = Bernoulli()
        d2[:lnk] = LogitLink()
        d2[:Buyer_Pos_P1_is1] = false
    end
    return d2
end

function readModels(fname::String)
    if isfile(fname)
        d=Dict()
        j = JSON.parsefile(fname)[1]
        #for (key, value) in dict
        for (k, v) in j
            d2=Dict()
            if Symbol(k)==:mocc
                d2[:dist] = Poisson()
                d2[:lnk] = LogLink()
                d2[:Buyer_Pos_P1_is1] = true
            end
            if Symbol(k)==:mdolocc
                d2[:dist] = Gamma()
                d2[:lnk] = LogLink()
                d2[:Buyer_Pos_P1_is1] = true
            end
            if Symbol(k)==:mpen
                d2[:dist] = Bernoulli()
                d2[:lnk] = LogitLink()
                d2[:Buyer_Pos_P1_is1] = false
            end

            if typeof(v) == Array{Any,1}
                d[Symbol(k)] = [Symbol(vx) for vx in v]
            else
                for (k1, v1) in v
                    if typeof(v1) in [ Array{Any,1} ]
                        v2=[Symbol(v1x) for v1x in v1]
                    else
                        v2=Symbol(v1)
                    end
                    d2[Symbol(k1)] = v2
                end
                d[Symbol(k)] = d2
            end
        end
        for (k,v) in d
            if k in [:idolocc, :ipen, :iocc]
                d[k] = expandM(v)
            end
        end
        return d
    else
        return "file not found!"
    end
end

read_dfx(root::String) = readtable(root*"/dfx.csv")
save_dfx(root::String,dfx::DataFrame) = writetable(root*"/dfx.csv", dfx)
read_modelsDict(root::String) = readModels(root*"/modelsDict.json")
function save_modelsDict(root::String, d::Dict)
    j = JSON.json([d])
    open(root*"/modelsDict.json","w") do f
        write(f,j)
    end
end
read_dfd(root::String) = readtable(root*"/dfd.csv",header=true);
save_dfd(root::String,dfd::DataFrame) = writetable(root*"/dfd.csv", dfd)

read_dfr(root::String) = readtable(root*"/scored.csv",header=true);
save_dfr(root::String,dfr::DataFrame) = writetable(root*"/scored.csv", dfr)

# ------- R 2 Julia Conversion - END ---------------

#include("trd.jl")
#function Base.lowercase(df::DataFrame)
function lowercase!(df::DataFrame)
    for (i,v) in enumerate(names(df))   if v != Symbol(lowercase(string(v)))  rename!(df,v,Symbol(lowercase(string(v)))) end end
end
function Base.lowercase(arr::Array{Symbol})    return  Symbol[Symbol(lowercase(string(s))) for s in arr] end
function Base.lowercase(arr::Array{String})    return  String[lowercase(string(s)) for s in arr] end
function Base.lowercase(arr::Array{String})    return  String[lowercase(string(s)) for s in arr] end
function Base.lowercase(s::Symbol) return Symbol(lowercase(string(s))) end
function Base.lowercase(d::OrderedDict)
    vout=OrderedDict()
    for (k, v) in d  vout[lowercase(k)] = typeof(v) in [Bool,Float64,Int64] ? v : lowercase(v) end
    return vout
end

function grep(v::String,vars::Array{Symbol})   #function grep(v::AbstractString) filter(x->ismatch(Regex("$(v)"), string(x)), vars) end
    filter(x->ismatch(Regex("$(v)"), string(x)), vars)
end
function grep1(v::String,vars::Array{Symbol})   #function grep(v::AbstractString) filter(x->ismatch(Regex("$(v)"), string(x)), vars) end
    lst = filter(x->ismatch(Regex("$(v)"), string(x)), vars)
    return length(lst)>0 ? lst[1] : nothing
end

function grep(varr::Array{String},vars::Array{Symbol})
    arr_out=Symbol[]
    for v in varr
        res = filter(x->ismatch(Regex("$(v)"), string(x)), vars)
        arr_out = vcat(arr_out, res)
    end
    return arr_out
end
#function grep(varr::Array{String},vars::Array{Symbol}) return grep(convert(Array{String},varr),vars) end  #grep(["M","pan"],names(df_in))

function df2dict(idf::DataFrame)
    dout=OrderedDict()
    for n in names(idf)
       dout[n] = idf[1,Symbol(n)]
    end
    return dout
end

function df2dict(row::DataFrameRow)
    dout=OrderedDict()
    for n in names(row)
       dout[n] = row[Symbol(n)]
    end
    return dout
end

"""
app.cfg
P2_Competitor = true
pvalue_lvl = 0.20
excludedKeys = key1,key2,key3
exposed_flag_var = new_exposed_flag
sigLevel = 0.2
random_demos =
random_campaigns =
dropvars = goop,snop
scoring_vars =
TotalModelsOnly=false
"""

function loadCFG(d::OrderedDict,fn::String)
    function cvt(v::String, t::DataType)
        if t in [Bool]
            return lowercase(v) =="true" ? true : lowercase(v) =="false" ? false : NA
        elseif t in [nothing]
            return NA
        elseif t in [Array{String,1}]
            #return AbstractString[string(x) split(v,",")]
            return String[string(x) for x in split(replace(v,"\"",""),",")]
        elseif t in [Array{Symbol,1}]
            return Symbol[Symbol(x) for x in split(replace(v,"\"",""),",")]
        elseif t in [Symbol]
            return Symbol(v)
        elseif t in [String]
            return string(v)
        elseif t in [Float64]
            return parse(Float64,v)
        else
            return "__ooops__ TYPE not found!"
        end
    end
    if isfile(fn)
        f = open(fn)
        for ln in eachline(f)
            ln=replace(ln,"\n","")
            sidx=search(ln,"=").start-1
            eidx=search(ln,"=").stop+1
            if sidx>0
                k=strip(ln[1:sidx])
                v=strip(ln[eidx:end])
                if (length(k)>0) & (length(v)>0)
                    println("LINE:: ~"*k*"~"*v*"~")
                    #println("DEF::       ______ ",d[Symbol(k)],"  ----> ",typeof(d[Symbol(k)]) )
                    newk=Symbol(k)
                    dv=get(d,newk,nothing)
                    if dv!=nothing
                        newv = cvt(v,typeof(dv))
                        #println("IN::        ______ ",Symbol(k),"  ----> ",ct, " ....... ", typeof(ct))
                        #println( d[Symbol(k)],"|",v," {",typeof(d[Symbol(k)]),"}  ----->  ",newk,"|",newv,"  {",typeof(newv),"}")
                        d[newk] = newv
                    else
                        println(newk," in app.cfg, but not defaults")
                        #d[newk] = v
                    end
                else
                    #println("CFG ERROR :: invalid length : "*"~"*k*"~"*v*"~",typeof(d[Symbol(k)]))
                end
            else
                if length(ln)>0   println("CFG ERROR :: invalid cfg parameter : ~",ln,"~") end
            end
        end
        close(f)
    else
        println("Using Default configuration")
    end
    return d
end

# Retype DataFrame Columns
# Ensure/Convert DataFrame Column Types
function retypeBools!(dfd::DataFrame,vars::Array{Symbol}=Symbol[])
    if length(vars)==0 vars=names(dfd) end
    for c in vars
        if typeof(dfd[c]) in [DataArray{Any,1}, DataArray{String,1}]  #DataArray{Bool,1}
            println(c," ~ ",typeof(dfd[c]))
            if  typeof(dfd[c]) == DataArray{Any,1}
                if unique(map(x-> typeof(x),dfd[c])) == [Bool]
                    println("Converting Col")
                    dfd[c] = convert(Array{Bool}, dfd[c])
                else
                    println("Type = Any; but not Bool")
                end
            else
                if true in unique(isna(dfd[c]))
                    println("String Col has NA's")
                else
                    ul=convert(Array{String},unique(map(x -> lowercase(x), dfd[c] )))
                    if (length(setdiff(ul,["true","false"]))==0) & (length(intersect(ul,["true","false"]))>0)
                        println("Mapping String Col")
                        dfd[c]=  map(x-> lowercase(x)=="true" ? true :  false , dfd[c])
                        dfd[c] = convert(Array{Bool}, dfd[c])
                    else
                        println("String col not Bool")
                    end
                end
            end
        end
    end
end


# Extracting these closures into types improves type stability a lot.
# Essentially, I have reinveted [zippers](https://wiki.haskell.org/Zipper) ;)

abstract type Context{N, C} end

struct LeftContext{N, C} <: Context{N, C}
    conditions::Dict{Variable{N}, Float64}
    threshold::Float64
    right::DecisionTree{N, C}
    parent::Context{N, C}
end

function (context::LeftContext{N, C})(t::DecisionTree{N, C}) where {N, C}
    context.parent(Decision{N, C}(context.conditions, context.threshold, t, context.right))
end

struct RightContext{N, C} <: Context{N, C}
    conditions::Dict{Variable{N}, Float64}
    threshold::Float64
    left::DecisionTree{N, C}
    parent::Context{N, C}
end

function (context::RightContext{N, C})(t::DecisionTree{N, C}) where {N, C}
    context.parent(Decision{N, C}(context.conditions, context.threshold, context.left, t))
end

struct NoContext{N, C} <: Context{N, C} end

function (context::NoContext{N, C})(t::DecisionTree{N, C}) where {N, C}
    t
end

struct Cont{N, C}
    context::Context{N, C}
    chunk::DecisionTree{N, C}
end

function (continuation::Cont{N, C})(k::Function) where {N, C}
    newchunk = k(continuation.chunk)::DecisionTree{N, C}
    tree = continuation.context(newchunk)::DecisionTree{N, C}
    return tree, continuation.chunk
end

function randsplit_impl{N, C}(rng::AbstractRNG, t::Classification{N, C}, n::Int,
                                parent_context::Context{N, C}, continuation::Cont{N, C})
    n += 1
    if rand() ≤ 1/n
        return Cont{N, C}(parent_context, t), n
    else
        return continuation, n
    end
end

function randsplit_impl{N, C}(rng::AbstractRNG, t::Decision{N, C}, n,
                                parent_context::Context{N, C}, continuation::Cont{N, C})
    c1, n = randsplit_impl(rng, t.iftrue, n,
                             LeftContext{N, C}(t.conditions, t.threshold, t.iffalse, parent_context),
                             continuation)
    c2, n = randsplit_impl(rng, t.iffalse, n,
                             RightContext{N, C}(t.conditions, t.threshold, t.iftrue, parent_context),
                             c1)

    n += 1
    if rand() ≤ 1/n
        return Cont{N, C}(parent_context, t), n
    else
        return c2, n
    end
end

"""
    randsplit(action,[ rng,] t::DecisionTree)

Select uniformly at random a point to split `t`, then replace the subtree `s` by `action(s)`.
Returns both the new tree and the replaced subtree.
"""
function randsplit{N, C}(action, rng::AbstractRNG, t::DecisionTree{N, C})
    # Uses reservoir samling of continuations, see: https://stackoverflow.com/a/3272490/1346276.
    default_context = NoContext{N, C}()
    default_continuation = Cont{N, C}(default_context, t)
    cont, _ = randsplit_impl(rng, t, 0, default_context, default_continuation)
    cont(action)
end

randsplit{N, C}(action, t::DecisionTree{N, C}) = randsplit(action, Base.GLOBAL_RNG, t)

"""
    randchild([rng,] t::DecisionTree)

Select uniformly at random a subtree of `t`.
"""
randchild{N, C}(rng::AbstractRNG, t::DecisionTree{N, C}) = randsplit(identity, rng, t)[2]
randchild{N, C}(t::DecisionTree{N, C}) = randchild(Base.GLOBAL_RNG, t)


module CUDNN

using ..CUJulia

if is_windows()
    const libcudnn = Libdl.find_library(["cudnn64_6","cudnn64_5"])
else
    const libcudnn = Libdl.find_library(["libcudnn"])
end
isempty(libcudnn) && throw("CUDNN library cannot be found.")

const version = ccall((:cudnnGetVersion,libcudnn),Cint,())
const major = div(version, 1000)
const minor = div(version - major*1000, 100)

info("CUDNN version: $(version)")
include("../lib/$(CUJulia.major)$(CUJulia.minor)/libcudnn$(major)$(minor).jl")
include("../lib/$(CUJulia.major)$(CUJulia.minor)/libcudnn$(major)$(minor)_types.jl")

function checkstatus(status)
    status == CUDNN_STATUS_SUCCESS && return
    throw(unsafe_string(cudnnGetErrorString(status)))
end

datatype(::Type{Float32}) = CUDNN_DATA_FLOAT
datatype(::Type{Float64}) = CUDNN_DATA_DOUBLE
datatype(::Type{Float16}) = CUDNN_DATA_HALF

const handles = Ptr{Void}[]
function handle(x)
    dev = device(x) + 1
    while dev > length(handles)
        p = Ptr{Void}[0]
        cudnnCreate(p)
        push!(handles, p[1])
    end
    handles[dev]
end
atexit(() -> foreach(cudnnDestroy, handles))

include("cudnn/activation.jl")
#include("batchnorm.jl")
include("cudnn/convolution.jl")
include("cudnn/dropout.jl")
include("cudnn/filter.jl")
include("cudnn/math.jl")
include("cudnn/pooling.jl")
include("cudnn/rnn.jl")
include("cudnn/softmax.jl")
include("cudnn/tensor.jl")

end


using ZipFile
resultspath="/home/ubuntu/data/dat"
resultspathagg="/home/ubuntu/data/agg"

const   agents = Dict{Int64,String}()
const   nodes = Dict{Int64,Array{Int64,1}}()
const   coords = Dict{Int64,Tuple{Float64,Float64}}()
const   routes1 = Dict{Int64,String}()
const   routes2 = Dict{Int64,String}()

const   routes1coords = Dict{Int64,String}()
const   routes2coords = Dict{Int64,String}()

function split_line(line::AbstractString, delim::AbstractString, firstc::Int64)
    res1=String[]
    ix = 1
    while length(res1) < firstc
        ix2 = findnext(delim,line,ix)[1]
        push!(res1,line[ix:(ix2-1)])
        ix = ix2+1
    end
    last_item = line[(findlast(";",line)[1]+1):end]
    (res1,line[ix:end],last_item)
end

function aggregate2!(resultspath,agents,nodes,coords)
    local lineheader_r
    f_count = 0
    for ff in readdir(resultspath)
        global ffn = joinpath(resultspath,ff)

        if isfile(ffn) && endswith(ff,"_nodes.csv.zip") && f_count < 30
            f_count += 1
            println("$f_count reading $ff")
            local line
            r = ZipFile.Reader(ffn);
            f = r.files[1]

            lineheader_s, lineheader_r = split_line(strip(readline(f)),";",4)
            while (!eof(f))
                line = strip(readline(f))
                line_s,line_r,line_last = split_line(line,";",4)
                id = parse(Int64,line_last)

                agents[id] = line_r
                node_id = parse(Int64,line_s[2])
                if !haskey(nodes,node_id)
                    nodes[node_id] = Int64[]
                    coords[node_id] = (parse(Float64,line_s[3]),parse(Float64,line_s[4]))
                end
                push!(nodes[node_id],id)
            end
            close(f)
            close(r)
        end
    end
    return lineheader_r
end

const node_line_header = aggregate2!(resultspath,agents,nodes,coords)

function collect_routes!(routes1,routes2,routes1coords,routes2coords)
    f_count = 0
    for ff in readdir(resultspath)
        ffn = joinpath(resultspath,ff)
        if isfile(ffn) && endswith(ff,"_routes.csv") && f_count < 30
            f_count += 1
            println("$f_count reading $ff")
            f = open(ffn,"r");
            local line
            while (!eof(f))
                line = strip(readline(f))
                if length(line) > 0
                    els = split(line,";")
                    agentid = parse(Int64,els[1])
                    nodes = parse.(Int64,els[3:end])
                    if els[2] == "towork"
                        routes1[agentid] = join(nodes,"#")
                        routes1coords[agentid] = join([join(coords[route], ",") for route in nodes],"#")
                    else
                        routes2[agentid] = join(nodes,"#")
                        routes2coords[agentid] = join([join(coords[route], ",") for route in nodes],"#")
                    end
                end
            end
            close(f)
        end
    end
end

collect_routes!(routes1,routes2,routes1coords,routes2coords)

function export_data(agents,nodes,coords,routes1,routes2,node_line_header)
    fns = open(joinpath(resultspathagg,"Node_Stats_.csv"),"w")
    println(fns,"\"NODE_ID\";\"longitude\";\"latitude\";\"vehicle_count\"")
    routes_tuple = (routes1,routes2)
    coords_tuple = (routes1coords,routes2coords)
    for node in sort!(collect(keys(nodes)))
        println(fns,"$(node);$(coords[node][1]);$(coords[node][2]);$(length(nodes[node]))")
        ftr = open(joinpath(resultspathagg,"Travel_$(node)_$(coords[node][1])_$(coords[node][2]).csv"),"w")
        fn = open(joinpath(resultspathagg,"Routes_n_$(node)_$(coords[node][1])_$(coords[node][2]).csv"),"w")
        fc = open(joinpath(resultspathagg,"Routes_c_$(node)_$(coords[node][1])_$(coords[node][2]).csv"),"w")
        println(ftr,"\"NODE_ID\";\"longitude\";\"latitude\";"*node_line_header)
        for agentid in nodes[node]
            println(ftr,"$(node);$(coords[node][1]);$(coords[node][2]);"*agents[agentid])
            routetype = "towork"
            for i in 1:2
                println(fn,"$agentid;$routetype;$(routes_tuple[i][agentid])")
                # dolozyc do linii pozynizej odpytywanie z coords
                #println(fc,"$agentid;$routetype;"*join([coords[route] for route in routes[agentid]],";"))
                println(fc, "$agentid;$routetype;$(coords_tuple[i][agentid])")
                routetype = "tohome"
            end
        end
        close(ftr)
        close(fn)
        close(fc)
        flush(fns)
    end
    close(fns)
end

export_data(agents,nodes,coords,routes1,routes2,node_line_header)


"""

abstract type FlowChannel{Sensors<:Tuple,Actuators<:Tuple} <: Component end

"""

# channels always sense the reversal of the currents they actuate
FlowChannel(T) = FlowChannel{Tuple{Reversal{T}},Tuple{T}}
FlowChannel(S, A) = FlowChannel{Tuple{S,Reversal{A}},Tuple{A}}
FlowChannel() = FlowChannel{Tuple{Nothing},Tuple{Nothing}}

sensed(::FlowChannel{S,A}) where {S,A} = typeflatten(S)
actuated(::FlowChannel{S,A}) where {S,A} = typeflatten(A)
sensed(::FlowChannel{Tuple{Nothing},A}) where {A} = Vector{DataType}()
actuated(::FlowChannel{S,Tuple{Nothing}}) where {S} = Vector{DataType}()

function typeflatten(s::DataType)
    map(fieldtypes(s)) do el
        el <: Tuple && return typeflatten(el)
        return (el,)
    end |> Iterators.flatten |> collect |> unique!
end

voltage(el) = (Voltage,)
ionic(x) = x <: Ion

currents(i::FlowChannel) =
    map(filter(ionic, actuated(i))) do thing
        Current{thing}
    end

sensedvars(i::FlowChannel) =
    map(sensed(i)) do thing
        Symbol(thing |> shorthand_name)
    end

mrna(i::FlowChannel) =
    map(filter(ionic, actuated(i))) do thing
        mRNA{thing}
    end

reversals(i::FlowChannel) =
    map(filter(ionic, actuated(i))) do thing
        Reversal{thing}
    end

conductances(i::FlowChannel) =
    map(filter(ionic, actuated(i))) do thing
        Conductance{thing}
    end

sensed_ions(i::FlowChannel) = filter(ionic, sensed(i))

instantiate_hooks(pre::Compartment, to::Component, args::Function...) =
    map(args) do fun
        map(fun(to)) do thing
            _name = shorthand_name(thing)
            if has_dynamics(pre, thing)
                return @variables $_name(t)
            else
                return @parameters $_name
            end
        end |> Iterators.flatten |> collect
    end |> Iterators.flatten |> collect

instantiate_variables(c::Component, args::Function...) =
    map(args) do fun
        map(fun(c)) do thing
            _name = shorthand_name(thing)
            return @variables $_name(t)
        end |> Iterators.flatten |> collect
    end |> Iterators.flatten |> collect

instantiate_parameters(c::Component, args::Function...) =
    map(args) do fun
        map(fun(c)) do thing
            _name = shorthand_name(thing)
            return @parameters $_name
        end |> Iterators.flatten |> collect
    end |> Iterators.flatten |> collect

instantiate_variables(v::Vector{Symbol}) =
    map(v) do f
        @variables $f(t)
    end |> Iterators.flatten |> collect

instantiate_variables(s::Symbol) = @variables $s(t)

instantiate_parameters(v::Vector{Symbol}) =
    map(v) do f
        @variables $f
    end |> Iterators.flatten |> collect

instantiate_parameters(s::Symbol) = @parameters $f

get_actuator(c::FlowChannel, sys::ODESystem, v::Type{Voltage}) =
    sum(currents(c)) do I
        ModelingToolkit.getvar(sys, shorthand_name(I); namespace=true)
    end

function get_actuator(c::FlowChannel, sys::ODESystem, f::DataType)
    indx = findfirst(x -> x == f, actuated(c))
    isnothing(indx) && return Num(0.0)
    return ModelingToolkit.getvar(sys, shorthand_name.(currents(c))[indx]; namespace=true)
end

function get_sensor(c::FlowChannel, sys::ODESystem, f::DataType)
    if f ∈ sensed(c)
        return ModelingToolkit.getvar(sys, shorthand_name(f); namespace=true)
    end
end
get_sensor(c::FlowChannel, sys::ODESystem, f::Type{Voltage}) =
    ModelingToolkit.getvar(sys, shorthand_name(f); namespace=true)


using ThreeBodyTB
using JLD2

fillist = readdir()

for f in fillist
    if occursin("jld2",f)
        if occursin("2body" ,f )
            dim = 2
        elseif occursin("3body", f)
            dim = 3
        else
            println("err $f")
        end

        t = @load f
        savename=f[1:end-5]
        c = ThreeBodyTB.CalcTB.make_coefs(dat.names, dim, datH=dat.datH, datS=dat.datS, min_dist = dat.min_dist, maxmin_val_train=dat.maxmin_val_train, dist_frontier=dat.dist_frontier)
        ThreeBodyTB.CalcTB.write_coefs("$savename.xml", c)

    end
end


using Associations, AssociationsGUI, Base.Dates
using Base.Test

@testset "util" begin

    @test AssociationsGUI.second2hms(Second(59)) == Dict(Hour => 0, Minute => 0, Second => 59)
    @test AssociationsGUI.second2hms(Second(60 + 1)) == Dict(Hour => 0, Minute => 1, Second => 1)
    @test AssociationsGUI.second2hms(Second(60*60 + 1)) == Dict(Hour => 1, Minute => 0, Second => 1)
    @test AssociationsGUI.second2hms(Second(60*60*60)) == Dict(Hour => 60, Minute => 0, Second => 0)

    @test length(AssociationsGUI.shorten("1234567", 3)) == 7
    @test length(AssociationsGUI.shorten("1234567", 2)) == 5

    @test AssociationsGUI.findshortfile("aaa", Dict("aaa" => "b", "c" => "aaa")) == "c"
    @test_throws ErrorException AssociationsGUI.findshortfile("aaa", Dict("aaa" => "b", "c" => "aa")) == "c"

    t = (1977, 06, 01, 12, 0, 1)
    d1 = AssociationsGUI.validargs(:a, t...)
    @test isnull(d1)

    t = (1977, 02, 31, 12, 0, 1)
    d2 = AssociationsGUI.validargs(:a, t...)
    @test !isnull(d2)

end

#=@testset "GUIs" begin

    @test_nowarn AssociationsGUI.poi_gui(POI("a", Point("1",0), Point("2",0), "", ""), ["a", "b"], Dict(string(i) => string(i) for i in 1:2), "a")

end=#


# load apm library
include("../apm.jl")

# Integrate model and return solution
z = apm_solve("demo",7)


import Luxor
import CairoMakie: Makie

@testset "plot_structure" begin
    w = "((((.....))))."
    s = "GGCGAAUACCGCCU"
    @test plot_structure(w) isa Luxor.Drawing
    @test plot_structure(w; sequence=s) isa Luxor.Drawing
    for layout in [:simple, :naview, :circular, :turtle, :puzzler]
        @test plot_structure(w; sequence=s, layout_type=layout) isa Luxor.Drawing
        @test plot_structure(w; sequence=s, layout_type=layout,
                             base_colors=rand(length(w))) isa Luxor.Drawing
    end
end

@testset "plot_structure_makie" begin
    w = "((((.....))))."
    s = "GGCGAAUACCGCCU"
    @test ViennaRNA.plot_structure_makie(w) isa Makie.Figure
end


export forw_model

function forw_model(v,Q,model,opts)
    Null = Array{Float64}(0)
    PDEfunc!(:forw_model,vec(v),vec(Q),Null,Null,model,opts)
end


const!(builder, x::Union{AbstractArray,XScalar,Bool}) = xlaclient.ops.Constant(builder, x)
const!(builder, xs::Tuple) = build!(builder, XTuple(), const!.((builder,), xs)...)

struct Lambda
  vars::Vector{Any}
  func::IR
end

for op in :[Neg, Sign, Floor, Ceil, Round, Exp, Log, Expm1, Log1p, Tanh,
            Sin, Cos, Lgamma, Digamma, Erf, Erfc, ErfInv, Sqrt, Rsqrt, Not].args
  @eval begin
    struct $op end
    build!(builder, ::$op, x) = xlaclient.ops.$op(x)
  end
end

for op in :[Atan2, Pow, And, Or, Xor, Add, Sub, Mul, SafeMul, Div, Rem, Dot,
            Max, Min, ShiftLeft, ShiftRightArithmetic, ShiftRightLogical,
            Gt, Ge, Lt, Le].args
  @eval begin
    struct $op end
    build!(builder, ::$op, x, y) = xlaclient.ops.$op(x, y)
  end
end

struct Rev
  dims
end

build!(builder, op::Rev, x) =
  xlaclient.ops.Rev(x, op.dims .- 1)

struct Conv
  strides
  padding
  lhs_dilation
  rhs_dilation
end

function build!(builder, op::Conv, x, w)
  dims = xlaclient.make_convolution_dimension_numbers(("WHCN", "WHIO", "WHCN"), 2)
  xlaclient.ops.ConvGeneralDilated(x, w, op.strides, op.padding, op.lhs_dilation, op.rhs_dilation, dims)
end

struct DynamicSlice
  size
end

function build!(builder, slice::DynamicSlice, x, start...)
  one = const!(builder, 1)
  start = map(i -> xlaclient.ops.Sub(i, one), start)
  xlaclient.ops.DynamicSlice(x, start, slice.size)
end

struct BroadcastInDim
  size
  map
end

build!(builder, bc::BroadcastInDim, x) = xlaclient.ops.BroadcastInDim(x, bc.size, bc.map.-1)

alldims(builder, xs) = [0:ndims(shapeof(builder, xs))-1;]

struct Map end

function build!(builder, ::Map, f, args...)
  settypes!(builder, f, args..., with = eltype)
  xlaclient.ops.Map(builder, args, build(f), alldims(builder, args[1]))
end

struct Reduce
  dims
end

function build!(builder, op::Reduce, f, xs, init)
  settypes!(builder, f, xs, xs, with = eltype)
  dims = op.dims == (:) ? alldims(builder, xs) : sort([op.dims...].-1)
  xlaclient.ops.Reduce(builder, [xs], [init], build(f), dims)
end

struct ReduceWindow
  window
  stride
  base_dilation
  window_dilation
  padding
end

ReduceWindow(window) = ReduceWindow(window, window, one.(window), one.(window), map(_ -> (0, 0), window))

function build!(builder, op::ReduceWindow, f, xs, init)
  xlaclient.ops.ReduceWindowWithGeneralPadding(
    xs, init, build(f), op.window, op.stride,
    op.base_dilation, op.window_dilation, op.padding)
end

struct SelectAndScatter
  window
  stride
  padding
end

SelectAndScatter(window) = SelectAndScatter(window, window, map(_ -> (0, 0), window))

function build!(builder, op::SelectAndScatter, select, scatter, xs, source, init)
  xlaclient.ops.SelectAndScatterWithGeneralPadding(
    xs, build(select), op.window, op.stride, op.padding, source, init, build(scatter))
end

struct XTuple end

build!(builder, ::XTuple, xs...) = xlaclient.ops.Tuple(builder, xs)

struct GetTupleElement
  idx::Int
end

build!(builder, op::GetTupleElement, x) = xlaclient.ops.GetTupleElement(x, op.idx)

struct ConvertElementType
  to::Type{<:XScalar}
end

build!(builder, op::ConvertElementType, x) =
  xlaclient.ops.ConvertElementType(x, primitivetype(op.to))

struct Reshape
  dims
  size
end

build!(builder, op::Reshape, x) =
  xlaclient.ops.Reshape(x, op.dims.-1, op.size)

struct Conditional end

function build!(builder, ::Conditional, pred,
                true_operand, true_computation,
                false_operand, false_computation)
  settypes!(builder, true_computation, true_operand)
  settypes!(builder, false_computation, false_operand)
  xlaclient.ops.Conditional(pred, true_operand, build(true_computation),
                            false_operand, build(false_computation))
end

struct While end

function build!(builder, ::While, condition, body, init)
  settypes!(builder, condition, init)
  settypes!(builder, body, init)
  xlaclient.ops.While(build(condition), build(body), init)
end


function get_halfcheetah(;
    timestep=0.01,
    gravity=[0.0; 0.0; -9.81],
    friction_coefficient=0.4,
    contact_feet=true,
    contact_body=true,
    limits=true,
    spring=[240.0, 180.0, 120.0, 180.0, 120.0, 60.0],
    damper=0.0,
    parse_damper=true,
    joint_limits=[[-0.52, -0.785, -0.400, -1.0, -1.20, -0.5],
                  [ 1.05,  0.785,  0.785,  0.7,  0.87,  0.5]],
    T=Float64)

    path = joinpath(@__DIR__, "../deps/halfcheetah.urdf")
    mech = Mechanism(path; floating=false, T,
        gravity,
        timestep,
        parse_damper)

    # Adding springs and dampers
    set_springs!(mech.joints, spring)
    set_dampers!(mech.joints, damper)

    # joint limits
    joints = deepcopy(mech.joints)

    if limits
        bthigh = get_joint(mech, :bthigh)
        joints[bthigh.id] = add_limits(mech, bthigh,
            rot_limits=[SVector{1}(joint_limits[1][1]), SVector{1}(joint_limits[2][1])])

        bshin = get_joint(mech, :bshin)
        joints[bshin.id] = add_limits(mech, bshin,
            rot_limits=[SVector{1}(joint_limits[1][2]), SVector{1}(joint_limits[2][2])])

        bfoot = get_joint(mech, :bfoot)
        joints[bfoot.id] = add_limits(mech, bfoot,
            rot_limits=[SVector{1}(joint_limits[1][3]), SVector{1}(joint_limits[2][3])])

        fthigh = get_joint(mech, :fthigh)
        joints[fthigh.id] = add_limits(mech, fthigh,
            rot_limits=[SVector{1}(joint_limits[1][4]), SVector{1}(joint_limits[2][4])])

        fshin = get_joint(mech, :fshin)
        joints[fshin.id] = add_limits(mech, fshin,
            rot_limits=[SVector{1}(joint_limits[1][5]), SVector{1}(joint_limits[2][5])])

        ffoot = get_joint(mech, :ffoot)
        joints[ffoot.id] = add_limits(mech, ffoot,
            rot_limits=[SVector{1}(joint_limits[1][6]), SVector{1}(joint_limits[2][6])])

        mech = Mechanism(Origin{T}(), [mech.bodies...], [joints...];
            gravity,
            timestep)
    end

    if contact_feet
        origin = Origin{T}()
        bodies = mech.bodies
        joints = mech.joints

        normal = [0.0; 0.0; 1.0]
        names = contact_body ? getfield.(mech.bodies, :name) : [:ffoot, :bfoot]
        models = []
        for name in names
            body = get_body(mech, name)
            if name == :torso # need special case for torso
                # torso
                pf = [+0.5 * body.shape.shape[1].rh[2]; 0.0; 0.0]
                pb = [-0.5 * body.shape.shape[1].rh[2]; 0.0; 0.0]
                o = body.shape.shape[1].rh[1]
                push!(models, contact_constraint(body, normal;
                    friction_coefficient,
                    contact_origin=pf,
                    contact_radius=o))
                push!(models, contact_constraint(body, normal;
                    friction_coefficient,
                    contact_origin=pb,
                    contact_radius=o))

                # head
                pf = [+0.5 * body.shape.shape[1].rh[2] + 0.214; 0.0; 0.1935]
                o = body.shape.shape[2].rh[1]
                push!(models, contact_constraint(body, normal;
                    friction_coefficient,
                    contact_origin=pf,
                    contact_radius=o))
            else
                p = [0;0; -0.5 * body.shape.rh[2]]
                o = body.shape.rh[1]
                push!(models, contact_constraint(body, normal;
                    friction_coefficient,
                    contact_origin=p,
                    contact_radius=o))
            end
        end

        set_minimal_coordinates!(mech, get_joint(mech, :floating_joint), [0.576509, 0.0, 0.02792])

        mech = Mechanism(origin, bodies, joints, [models...];
            gravity,
            timestep)
    end
    return mech
end

function initialize_halfcheetah!(mechanism::Mechanism{T};
    body_position=[0.0, 0.0],
    body_orientation=0.0) where T

    set_minimal_coordinates!(mechanism,
                 get_joint(mechanism, :floating_joint),
                 [body_position[2] + 0.576509, -body_position[1], -body_orientation + 0.02792])
    for joint in mechanism.joints
        (joint.name != :floating_joint) && set_minimal_coordinates!(mechanism, joint, zeros(input_dimension(joint)))
    end
    zero_velocity!(mechanism)
end


using LaserTag
using Plots
using POMDPToolbox
using Reel
using ProgressMeter
using PmapProgressMeter
using ParticleFilters
using ContinuousPOMDPTreeSearchExperiments
using Plots
using QMDP
using JLD
# using DESPOT
using BasicPOMCP
using ARDESPOT

file_contents = readstring(@__FILE__())

@everywhere begin
# begin
    using POMDPs
    using POMDPToolbox
    using ContinuousPOMDPTreeSearchExperiments
    using DiscreteValueIteration
    using ParticleFilters
    using POMCPOW
    using LaserTag
    using QMDP
    import DESPOT
    using BasicPOMCP
    using ARDESPOT

    N = 1000
    K = 5000
    T_max = 10.0
    P = typeof(gen_lasertag(rng=MersenneTwister(18)))

    solvers = Dict{String, Union{Policy, Solver}}(

        "qmdp" => QMDPSolver(max_iterations=1000),

        "ardespot" => begin
            DESPOTSolver(lambda=0.0,
                         max_trials=1_000_000,
                         T_max=T_max,
                         K=K,
                         bounds=LaserBounds{P}(),
                         default_action=nogap_tag,
                         rng=MersenneTwister(4))
        end,

        "despot" => begin
            DESPOT.DESPOTSolver{LTState,
                                Int,
                                DMeas,
                                LaserBounds{P},
                                DESPOT.MersenneStreamArray}(
                bounds=LaserBounds{P}(),
                time_per_move=T_max,
                max_trials=1_000_000,
                n_particles=K,
                random_streams=DESPOT.MersenneStreamArray(MersenneTwister(4)),
                next_state=LTState([-1,-1],[-1,-1], false)
            )
        end
    )
end

@show K
@show N
@show T_max

rdict = Dict{String, Any}()
for (k,sol) in solvers
    prog = Progress(N, desc="Simulating...")
    @show k
    rewards = pmap(prog, 1:N) do i
    # rewards = map(1:N) do i
        pomdp = gen_lasertag(rng=MersenneTwister(i+600_000))
        if isa(sol,Solver)
            p = solve(deepcopy(sol), pomdp)
        else
            p = sol
        end
        hr = HistoryRecorder(max_steps=100, rng=MersenneTwister(i))
        up_rng = MersenneTwister(i+100_000)
        up = ObsAdaptiveParticleFilter(deepcopy(pomdp), LowVarianceResampler(100_000), 0.05, up_rng)
        hist = simulate(hr, pomdp, p, up)
        discounted_reward(hist)
    end
    @show mean(rewards)
    @show std(rewards)/sqrt(N)
    rdict[k] = rewards
end

filename = Pkg.dir("ContinuousPOMDPTreeSearchExperiments", "data", "despot_comparison_run_$(Dates.format(now(), "E_d_u_HH_MM")).jld")
println("saving to $filename...")
@save(filename, rdict, file_contents)
println("done.")


"""
    module PICD
---
The module that does `p`ivoted `i`ncomplete `C`holesky `d`ecomposition(PICD).
"""
module PICD
using Mmap, Statistics, Octavian, LinearAlgebra

include("genotype.jl")
include("picd.jl")

function picd_H()
end

end


module Objects
  using CG;
  export houseOfSantaClaus;
  function houseOfSantaClaus()
    return Object(Vec4f(-1, -1, 0, 1),
    Vec4f(1, -1, 0, 1),
    Vec4f(-1, 1, 0, 1),
    Vec4f(0, 2, 0, 1),
    Vec4f(1, 1, 0, 1),
    Vec4f(1, 1, 0, 1),
    Vec4f(-1, -1, 0, 1),
    Vec4f(-1, 1, 0, 1),
    Vec4f(1, 1, 0, 1),
    Vec4f(1, -1, 0, 1));
  end
end


module Auto2D

using AutomotiveDrivingModels
using AutoViz
using PDMats
using NGSIM
using ForwardNets
import Reel

export gen_simparams, reset, tick, reward, observe, step, isdone, action_space_bounds, observation_space_bounds, render
export SimParams, reel_drive, GaussianMLPDriver, load_gru_driver

##################################
# Gaussian MLP Driver

type GaussianMLPDriver{A<:DriveAction, F<:Real, G<:Real, E<:AbstractFeatureExtractor, M<:MvNormal} <: DriverModel{A, IntegratedContinuous}
    net::ForwardNet
    rec::SceneRecord
    pass::ForwardPass
    input::Vector{F}
    output::Vector{G}
    extractor::E
    mvnormal::M
    context::IntegratedContinuous
end

_get_Σ_type{Σ,μ}(mvnormal::MvNormal{Σ,μ}) = Σ
function GaussianMLPDriver{A <: DriveAction}(::Type{A}, net::ForwardNet, extractor::AbstractFeatureExtractor, context::IntegratedContinuous;
    input::Symbol = :input,
    output::Symbol = :output,
    Σ::Union{Real, Vector{Float64}, Matrix{Float64},  Distributions.AbstractPDMat} = 0.1,
    rec::SceneRecord = SceneRecord(2, context.Δt),
    )

    pass = calc_forwardpass(net, [input], [output])
    input_vec = net[input].tensor
    output = net[output].tensor
    mvnormal = MvNormal(Array(Float64, 2), Σ)
    GaussianMLPDriver{A, eltype(input_vec), eltype(output), typeof(extractor), typeof(mvnormal)}(net, rec, pass, input_vec, output, extractor, mvnormal, context)
end

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

include("../validation/load_train_test_split.jl")
include("../pull_traces/multifeatureset.jl")
include("../validation/RootDir.jl")

function get_train_segments(trajdatas::Dict{Int, Trajdata}, nsteps::Int)

    assignment = load_assignment()
    evaldata = load_evaldata()

    # pull all segments that are for training
    all_train_segments = evaldata.segments[find(assignment .== FOLD_TRAIN)]

    # only keep segments that are long enough and are not within 500 frames of start or end
    train_segments = TrajdataSegment[]
    for seg in all_train_segments
        if !haskey(trajdatas, seg.trajdata_index) # only keep valid segments
            continue
        end
        trajdata = trajdatas[seg.trajdata_index]
        frame_lo = max(seg.frame_lo, 500)
        frame_hi = min(seg.frame_hi, nframes(trajdata) - 1000)
        if frame_hi - frame_lo > nsteps # if it is long enough
            push!(train_segments, TrajdataSegment(seg.trajdata_index, seg.egoid, frame_lo, frame_hi))
        end
    end

    train_segments
end

##################################
# SimParams
include("../validation/load_policy.jl")

type SimState
    frame::Int
    start_frame::Int
    egoid::Int
    trajdata_index::Int

    scene::Scene
    rec::SceneRecord

    playback_reactive_active_vehicle_ids::Set{Int}
    playback_reactive_speeds::Dict{Int, Float64}
    hidden_dict::Dict{Int, Vector{Float64}}

    function SimState(context::IntegratedContinuous, rec_size::Int)
        retval = new()
        retval.scene = Scene()
        retval.rec = SceneRecord(rec_size, context.Δt)
        retval.playback_reactive_active_vehicle_ids = Set{Int}()
        retval.playback_reactive_speeds = Dict{Int, Float64}()
        retval.hidden_dict = Dict{Int, Vector{Float64}}()
        retval
    end
end

type SimParams
    col_weight::Float64 # reward for collision
    off_weight::Float64 # reward for being offroad
    rev_weight::Float64 # reward for reversal (negative speeds)
    jrk_weight::Float64 # reward scale for square jerk
    acc_weight::Float64 # reward scale for square acceleration
    cen_weight::Float64 # reward scale for lane center offset
    ome_weight::Float64 # reward scale for turn acceleration (omega)
    use_debug_reward::Bool # false by default

    context::IntegratedContinuous
    prime_history::Int
    ego_action_type::DataType

    safety_model::DriverModel{LatLonAccel, IntegratedContinuous}

    use_playback_reactive::Bool
    playback_reactive_model::DriverModel{LatLonAccel, IntegratedContinuous}
    playback_reactive_threshold_brake::Float64
    playback_reactive_scene_buffer::Scene

    model_all::Bool
    driver_model::GaussianMLPDriver

    trajdatas::Dict{Int, Trajdata}
    segments::Vector{TrajdataSegment}

    nsteps::Int
    step_counter::Int
    simstates::Vector{SimState}
    features::Vector{Float64}
    extractor::MultiFeatureExtractor
end

function SimParams(trajdatas::Dict{Int, Trajdata}, segments::Vector{TrajdataSegment},
    col_weight::Float64,
    off_weight::Float64,
    rev_weight::Float64,
    jrk_weight::Float64,
    acc_weight::Float64,
    cen_weight::Float64,
    ome_weight::Float64,
    use_debug_reward::Bool,
    use_playback_reactive::Bool,
    model_all::Bool,
    playback_reactive_threshold_brake::Float64, # [m/s²]
    nsimstates::Int,
    prime_history::Int,
    nsteps::Int,
    ego_action_type::DataType,
    extractor::MultiFeatureExtractor,
    context = IntegratedContinuous(NGSIM_TIMESTEP,1),
    )
    println("Begin SimParams")
    simstates = Array(SimState, nsimstates)
    for i in 1 : length(simstates)
        simstates[i] = SimState(context, prime_history+1)
    end
    println("loop end")
    features = Array(Float64, length(extractor))

    safety_policy = LatLonSeparableDriver(
        context,
        ProportionalLaneTracker(σ=0.1, kp=3.0, kd=2.0),
        IntelligentDriverModel(σ=0.1, k_spd=1.0, T=0.5, s_min=1.0, a_max=3.0, d_cmf=2.5),
        )
    println("safe policy end")
    playback_reactive_model = LatLonSeparableDriver(
        context,
        ProportionalLaneTracker(σ=0.1, kp=3.0, kd=2.0),
        IntelligentDriverModel(σ=0.1, k_spd=1.0, T=0.5, s_min=1.0, a_max=3.0, d_cmf=2.5),
        )
    playback_reactive_scene_buffer = Scene()
    println("playback reactive model end")
    filepath = joinpath(ROOT_FILEPATH,"julia","validation","models","gail_gru.h5")
    iteration = 413
    driver_model = load_gru_driver(filepath, iteration)
    println("SimParams initial process end")
    SimParams(
        col_weight, off_weight, rev_weight, jrk_weight, acc_weight, cen_weight, ome_weight, use_debug_reward,
        context, prime_history, ego_action_type, safety_policy,
        use_playback_reactive, playback_reactive_model,
        playback_reactive_threshold_brake, playback_reactive_scene_buffer, model_all, driver_model,
        trajdatas, segments, nsteps, 0, simstates, features, extractor)
end
function gen_simparams(trajdata_indeces::Vector,
    col_weight::Float64,
    off_weight::Float64,
    rev_weight::Float64,
    jrk_weight::Float64,
    acc_weight::Float64,
    cen_weight::Float64,
    ome_weight::Float64,
    use_debug_reward::Bool,
    use_playback_reactive::Bool,
    model_all::Bool,
    playback_reactive_threshold_brake::Float64,
    nsimstates::Int,
    prime_history::Int,
    nsteps::Int,
    ego_action_type::DataType,
    extractor::MultiFeatureExtractor,
    )

    println("loading trajdatas: ", trajdata_indeces); tic()
    trajdatas = Dict{Int, Trajdata}()
    println("loading done")
    println(load_trajdata(1))
    for i in trajdata_indeces
        trajdatas[i] = load_trajdata(i)
    end
    toc()

    println("loading training segments"); tic()
    segments = get_train_segments(trajdatas, nsteps)
    toc()

    println("HIT THIS THING")
    SimParams(trajdatas, segments, col_weight, off_weight, rev_weight, jrk_weight, acc_weight, cen_weight, ome_weight,
              use_debug_reward, use_playback_reactive, model_all, playback_reactive_threshold_brake,
              nsimstates, prime_history, nsteps, ego_action_type, extractor)
end
function gen_simparams_from_trajdatas(trajdata_filepaths::Vector, roadway_filepaths::Vector,
    col_weight::Float64,
    off_weight::Float64,
    rev_weight::Float64,
    jrk_weight::Float64,
    acc_weight::Float64,
    cen_weight::Float64,
    ome_weight::Float64,
    use_debug_reward::Bool,
    use_playback_reactive::Bool,
    model_all::Bool,
    playback_reactive_threshold_brake::Float64,
    nsimstates::Int,
    prime_history::Int,
    nsteps::Int,
    ego_action_type::DataType,
    extractor::MultiFeatureExtractor,
    )

    println("loading trajdatas"); tic()
    roadways = Dict{String, Roadway}()
    trajdatas = Dict{Int, Trajdata}()
    for (i,filepath) in enumerate(trajdata_filepaths)
        roadwayfp = roadway_filepaths[i]
        if !haskey(roadways, roadwayfp)
            roadways[roadwayfp] = open(io->read(io, Roadway), roadwayfp, "r")
        end
        trajdatas[i] = open(io->read(io, Trajdata, roadways[roadwayfp]), filepath, "r")
    end
    toc()

    println("generating training segments"); tic()
    segments = TrajdataSegment[]
    for (trajdata_index,trajdata) in trajdatas
        append!(segments, pull_continuous_segments(trajdata, trajdata_index))
    end
    toc()

    SimParams(trajdatas, segments, col_weight, off_weight, rev_weight, jrk_weight, acc_weight, cen_weight, ome_weight,
              use_debug_reward, use_playback_reactive, model_all, playback_reactive_threshold_brake,
              nsimstates, prime_history, nsteps, ego_action_type, extractor)
end
function gen_simparams(batch_size::Int, args::Dict)

    col_weight = get(args, "col_weight", -2.0)
    off_weight = get(args, "off_weight", -0.75)
    rev_weight = get(args, "rev_weight", -0.50)
    jrk_weight = get(args, "jrk_weight", -0.050)
    acc_weight = get(args, "acc_weight", -0.050)
    cen_weight = get(args, "cen_weight", -0.050)
    ome_weight = get(args, "ome_weight", -0.500)

    use_debug_reward = get(args, "use_debug_reward", false)
    use_playback_reactive = convert(Bool, get(args, "use_playback_reactive", false))
    playback_reactive_threshold_brake = get(args, "playback_reactive_threshold_brake", -2.0) # [m/s²]
    prime_history = get(args, "prime_history", 2)
    nsteps = get(args, "nsteps", 100)

    model_all = get(args, "model_all", false)

    extract_core = get(args, "extract_core", true)
    extract_temporal = get(args, "extract_temporal", true)
    extract_well_behaved = get(args, "extract_well_behaved", true)
    extract_neighbor_features = get(args, "extract_neighbor_features", false)
    extract_carlidar_rangerate = get(args, "extract_carlidar_rangerate", true)
    carlidar_nbeams = get(args, "carlidar_nbeams", 0)
    roadlidar_nbeams = get(args, "roadlidar_nbeams", 0)
    roadlidar_nlanes = get(args, "roadlidar_nlanes", 2)
    carlidar_max_range = get(args, "carlidar_max_range", 100.0) # [m]
    roadlidar_max_range = get(args, "roadlidar_max_range", 100.0) # [m]

    extractor = MultiFeatureExtractor(
        extract_core,
        extract_temporal,
        extract_well_behaved,
        extract_neighbor_features,
        extract_carlidar_rangerate,
        carlidar_nbeams,
        roadlidar_nbeams,
        roadlidar_nlanes,
        carlidar_max_range=carlidar_max_range,
        roadlidar_max_range=roadlidar_max_range,
        )

    ego_action_type = AccelTurnrate
    T = get(args, "action_type", "AccelTurnrate")
    if T == "AccelTurnrate"
        ego_action_type = AccelTurnrate
    elseif T == "LatLonAccel"
        ego_action_type = LatLonAccel
    end

    if haskey(args, "trajdata_filepaths")
        trajdata_filepaths = args["trajdata_filepaths"]
        roadway_filepaths = args["roadway_filepaths"]
        gen_simparams_from_trajdatas(trajdata_filepaths, roadway_filepaths,
            col_weight, off_weight, rev_weight, jrk_weight, acc_weight, cen_weight, ome_weight,
            use_debug_reward, use_playback_reactive, model_all, playback_reactive_threshold_brake,
            batch_size, prime_history, nsteps, ego_action_type, extractor)
    else
        trajdata_indeces = get(args, "trajdata_indeces", [1,2,3,4,5,6])
        gen_simparams(trajdata_indeces,
            col_weight, off_weight, rev_weight, jrk_weight, acc_weight, cen_weight, ome_weight,
            use_debug_reward, use_playback_reactive, model_all, playback_reactive_threshold_brake,
            batch_size, prime_history, nsteps, ego_action_type, extractor)
    end
end

###########################################
AutomotiveDrivingModels.get_name(::GaussianMLPDriver) = "GaussianMLPDriver"
AutomotiveDrivingModels.action_context(model::GaussianMLPDriver) = model.context

function AutomotiveDrivingModels.reset_hidden_state!(model::GaussianMLPDriver)
    empty!(model.rec)
    model
end

function AutomotiveDrivingModels.observe!{A,F,G,E,P}(
                                            model::GaussianMLPDriver{A,F,G,E,P},
                                            simparams::SimParams,
                                            scene::Scene,
                                            roadway::Roadway,
                                            egoid::Int)

    update!(model.rec, scene)
    vehicle_index = get_index_of_first_vehicle_with_id(scene, egoid)
    o = pull_features!(simparams.extractor, simparams.features, model.rec, roadway, vehicle_index)
    model.net[:hidden_0].input = (o - model.extractor.feature_means)./model.extractor.feature_std
    forward!(model.pass)
    copy!(model.mvnormal.μ, model.output[1:2])

    model
end
Base.rand{A,F,G,E,P}(model::GaussianMLPDriver{A,F,G,E,P}) = convert(A, rand(model.mvnormal))
Distributions.pdf{A,F,G,E,P}(model::GaussianMLPDriver{A,F,G,E,P}, a::A) = pdf(model.mvnormal, convert(Vector{Float64}, a))
Distributions.logpdf{A,F,G,E,P}(model::GaussianMLPDriver{A,F,G,E,P}, a::A) = logpdf(model.mvnormal, convert(Vector{Float64}, a))
############################################

Base.show(io::IO, simparams::SimParams) = print(io, "SimParams")

function restart!(simstate::SimState, simparams::SimParams)

    # pick a random segment
    local train_seg_index
    for i in 1 : 100
        train_seg_index = rand(1:length(simparams.segments))
        seg = simparams.segments[train_seg_index]
        candidate_frame_lo = seg.frame_lo
        candidate_frame_hi = seg.frame_hi - simparams.nsteps - simparams.prime_history
        if candidate_frame_hi > candidate_frame_lo
        break
    end
    if i > 95
        assert(false)
        end
    end

    seg = simparams.segments[train_seg_index]
    simstate.egoid = seg.egoid
    simstate.trajdata_index = seg.trajdata_index

    # pull the trajdata
    trajdata = simparams.trajdatas[simstate.trajdata_index]

    # pick a random sub-trajectory
    candidate_frame_lo = seg.frame_lo
    candidate_frame_hi = seg.frame_hi - simparams.nsteps - simparams.prime_history
    simstate.frame = rand(candidate_frame_lo:candidate_frame_hi)

    # clear rec and make first observations
    empty!(simstate.rec)
    for i in 1 : simparams.prime_history + 1
        get!(simstate.scene, trajdata, simstate.frame)
        update!(simstate.rec, simstate.scene)
        simstate.frame += 1
    end
    simstate.frame -= 1
    simstate.start_frame = simstate.frame

    # empty playback reactive
    empty!(simstate.playback_reactive_active_vehicle_ids)

    # simparams.extractor.road_lidar_culling = simparams.roadway_cullers[simstate.trajdata_index]

    simstate
end

function Base.reset(simparams::SimParams)
    for state in simparams.simstates
        restart!(state, simparams)
    end
    simparams.step_counter = 0

    simparams
end

function step_forward!(simstate::SimState, simparams::SimParams, action_ego::Vector{Float64})

    trajdata = simparams.trajdatas[simstate.trajdata_index]
    veh_ego = get_vehicle(simstate.scene, simstate.egoid)

    safety_policy = simparams.safety_model

    act_dim = length(action_ego)
    if act_dim == 3
        if action_ego[end] > 0
            #set_desired_speed!(safety_policy, simstate.playback_reactive_speeds[simstate.egoid])
            AutomotiveDrivingModels.observe!(safety_policy,simstate.scene,trajdata.roadway,simstate.egoid)
            action_ego = rand(safety_policy)::LatLonAccel # is this only returning acceleration?
            #action_ego = convert(simparams.ego_action_type, action_ego)

        else
            action_ego = action_ego[1:end-1]
            action_ego = convert(simparams.ego_action_type, action_ego)

        end
    else
        action_ego = convert(simparams.ego_action_type, action_ego)
    end

    # propagate the ego vehicle
    trajdata = simparams.trajdatas[simstate.trajdata_index]
    veh_ego = get_vehicle(simstate.scene, simstate.egoid)
    #action_ego = convert(simparams.ego_action_type, action_ego)

    ego_state = propagate(veh_ego, action_ego, simparams.context, trajdata.roadway)

    simstate.frame += 1
    if simparams.model_all
        model = simparams.driver_model

        empty!(simparams.playback_reactive_scene_buffer)

        tdframe = trajdata.frames[simstate.frame]
        for i in tdframe.lo : tdframe.hi
            s = trajdata.states[i]
            if s.id != simstate.egoid

                veh_index = get_index_of_first_vehicle_with_id(simstate.scene, s.id)
                if veh_index != 0 # iscarinframe
                    if Symbol("gru") in fieldnames(model.net)
                        if !(s.id in simstate.hidden_dict)
                            model.net[:gru].h_prev = zeros(length(model.net[:gru].h_prev))
                        else
                            model.net[:gru].h_prev = simstate.hidden_dict[s.id]
                        end
                    end

                    AutomotiveDrivingModels.observe!(model, simparams, simstate.scene, trajdata.roadway, s.id)
                    action = rand(model)::AccelTurnrate
                    a = clamp(action.a, -5.0, 3.0)
                    ω = clamp(action.ω, -0.01, 0.01)

                    # Propagate scene
                    veh = simstate.scene[veh_index]
                    veh_state = propagate(veh, AccelTurnrate(a, ω), simparams.context, trajdata.roadway)
                    veh2 = Vehicle(veh_state, veh.def)
                    push!(simparams.playback_reactive_scene_buffer, veh2)

                    if Symbol("gru") in fieldnames(model.net)
                        simstate.hidden_dict[s.id] = model.net[:gru].h
                    end
                end
            end
        end

        # move simparams.playback_reactive_scene_buffer over to simparams.scene
        # and place new ego pos
        ego_def = veh_ego.def
        copy!(simstate.scene, simparams.playback_reactive_scene_buffer)
        push!(simstate.scene, Vehicle(ego_state, ego_def))
    elseif simparams.use_playback_reactive

        model = simparams.playback_reactive_model

        empty!(simparams.playback_reactive_scene_buffer)

        tdframe = trajdata.frames[simstate.frame]
        for i in tdframe.lo : tdframe.hi
            s = trajdata.states[i]
            if s.id != simstate.egoid

                use_playback = true
                veh_index = get_index_of_first_vehicle_with_id(simstate.scene, s.id)
                ego_index = get_index_of_first_vehicle_with_id(simstate.scene, simstate.egoid)
                if veh_index != 0 # iscarinframe

                    if !in(s.id, simstate.playback_reactive_active_vehicle_ids)
                        simstate.playback_reactive_speeds[s.id] = simstate.scene[veh_index].state.v
                    end

                    set_desired_speed!(model, simstate.playback_reactive_speeds[s.id])
                    AutomotiveDrivingModels.observe!(model, simstate.scene, trajdata.roadway, s.id)
                    action = rand(model)::LatLonAccel

                    if s.id in simstate.playback_reactive_active_vehicle_ids ||
                       (action.a_lon < simparams.playback_reactive_threshold_brake &&
                        simstate.scene[veh_index].state.posF.s < simstate.scene[ego_index].state.posF.s)

                        # use IDM
                        use_playback = false
                        push!(simstate.playback_reactive_active_vehicle_ids, s.id)

                        veh = simstate.scene[veh_index]
                        if veh.state.v < 0.0
                            action = LatLonAccel(action.a_lat, max(action.a_lon, 0.0))
                        end
                        veh_state = propagate(veh, action, simparams.context, trajdata.roadway)
                        veh2 = Vehicle(veh_state, veh.def)
                        push!(simparams.playback_reactive_scene_buffer, veh2)
                    end
                end

                if use_playback
                    # if it is a new car not in previous frame or it doesn't need IDM
                    veh = Vehicle(s.state, get_vehicledef(trajdata, s.id))
                    push!(simparams.playback_reactive_scene_buffer, veh)
                end
            end
        end

        # move simparams.playback_reactive_scene_buffer over to simparams.scene
        # and place new ego pos
        ego_def = veh_ego.def
        copy!(simstate.scene, simparams.playback_reactive_scene_buffer)
        push!(simstate.scene, Vehicle(ego_state, ego_def))
    else
        # pull new frame from trajdata
        get!(simstate.scene, trajdata, simstate.frame)

        # move in propagated ego vehicle
        veh_index = get_index_of_first_vehicle_with_id(simstate.scene, simstate.egoid)
        simstate.scene[veh_index].state = ego_state
    end

    # update record
    update!(simstate.rec, simstate.scene)

    simstate
end

function tick(simparams::SimParams, u::Vector{Float64}, batch_index::Int=1)
    step_forward!(simparams.simstates[batch_index], simparams, u)
    simparams
end
function step_forward(simparams::SimParams, U::Matrix{Float64})
    # note: actions are batch_size × action_size
    for (i,s) in enumerate(simparams.simstates)
        step_forward!(s, simparams, U[i,:])
    end
    simparams
end

function reward(simstate::SimState, simparams::SimParams)

    veh_index = get_index_of_first_vehicle_with_id(simstate.scene, simstate.egoid)
    veh_ego = simstate.scene[veh_index]

    reward = 0.0
    if !simparams.use_debug_reward
        trajdata = simparams.trajdatas[simstate.trajdata_index]
        acc_ego = convert(Float64, get(ACC, simstate.rec, trajdata.roadway, veh_index))
        jrk_ego = convert(Float64, get(JERK, simstate.rec, trajdata.roadway, veh_index))
        ome_ego = convert(Float64, get(ANGULARRATEG, simstate.rec, trajdata.roadway, veh_index))
        cen_ego = veh_ego.state.posF.t

        d_ml = convert(Float64, get(MARKERDIST_LEFT, simstate.rec, trajdata.roadway, veh_index))
        d_mr = convert(Float64, get(MARKERDIST_RIGHT, simstate.rec, trajdata.roadway, veh_index))

        reward += simparams.col_weight * get_first_collision(simstate.scene, veh_index).is_colliding
        reward += simparams.off_weight * (d_ml < -1.0 || d_mr < -1.0)
        reward += simparams.rev_weight * (veh_ego.state.v < 0.0)
        reward += simparams.jrk_weight * jrk_ego*jrk_ego
        reward += simparams.acc_weight * acc_ego*acc_ego
        reward += simparams.cen_weight * cen_ego*cen_ego
        reward += simparams.ome_weight * ome_ego*ome_ego
    else
        reward -= (veh_ego.state.v^2)
        reward -= ((veh_ego.state.posG.θ - π/2)^2)
    end

    reward
end
function reward(simparams::SimParams, u::Vector{Float64}, batch_index::Int = 1)
    reward(simparams.simstates[batch_index], simparams)
end

function observe(simparams::SimParams, batch_index::Int=1)
    simstate = simparams.simstates[batch_index]
    trajdata = simparams.trajdatas[simstate.trajdata_index]
    veh_index = get_index_of_first_vehicle_with_id(simstate.scene, simstate.egoid)
    pull_features!(simparams.extractor, simparams.features, simstate.rec, trajdata.roadway, veh_index)
end

isdone(simparams::SimParams) = simparams.step_counter ≥ simparams.nsteps

function Base.step(simparams::SimParams, u::Vector{Float64}, batch_index::Int=1)

    r = reward(simparams, u, batch_index)
    tick(simparams, u, batch_index)
    features = observe(simparams, batch_index)
    simparams.step_counter += 1
    # done = isdone(simparams)
    simstate = simparams.simstates[batch_index]

    # End if collision or reverse or off-road
    veh_index = get_index_of_first_vehicle_with_id(simstate.scene, simstate.egoid)
    trajdata = simparams.trajdatas[simstate.trajdata_index]
    d_ml = convert(Float64, get(MARKERDIST_LEFT, simstate.rec, trajdata.roadway, veh_index))
    d_mr = convert(Float64, get(MARKERDIST_RIGHT, simstate.rec, trajdata.roadway, veh_index))
    done = false
    done = done || get_first_collision(simstate.scene, veh_index).is_colliding
    done = done || (d_ml < -1.0 || d_mr < -1.0)
    done = done || (simstate.scene[veh_index].state.v < 0.0)

    (features, r, done)
end
function Base.step(simparams::SimParams, U::Matrix{Float64})

    batch_size = length(simparams.states)
    feature_mat = Array(Float64, batch_size, obssize(simparams))
    rewards = Array(Float64, batch_size)
    dones = Array(Float64, batch_size)

    step_counter = simparams.step_counter
    for batch_index in 1 : batch_size
        simparams.step_counter = step_counter
        features, reward, done = step(simparams, U[batch_index, :], batch_index)
        feature_mat[batch_index, :] = features
        rewards[batch_index] = reward
        dones[batch_index] = done
    end
    simparams.step_counter = step_counter + 1

    (feature_mat, rewards, dones)
end

action_space_bounds(simparams::SimParams) = ([-5.0, -1.0], [3.0, 1.0])
observation_space_bounds(simparams::SimParams) = (fill(-Inf, length(simparams.extractor)), fill(Inf, length(simparams.extractor)))

##################################
# Visualization

function AutoViz.render(simparams::SimParams, batch_index::Int=1)

    simstate = simparams.simstates[batch_index]
    trajdata = simparams.trajdatas[simstate.trajdata_index]

    carcolors = Dict{Int,Colorant}()
    carcolors[simstate.egoid] = COLOR_CAR_EGO

    for id in simstate.playback_reactive_active_vehicle_ids
        carcolors[id] = colorant"0xF9971FFF" # orange
    end

    render(simstate.scene, trajdata.roadway,
          cam=CarFollowCamera(simstate.egoid, 5.0), car_colors=carcolors)
end

function AutoViz.render(simparams::SimParams, image::Matrix{UInt32}, batch_index::Int=1)

    simstate = simparams.simstates[batch_index]
    trajdata = simparams.trajdatas[simstate.trajdata_index]

    c = CairoImageSurface(image, Cairo.FORMAT_ARGB32, flipxy=false)
    ctx = CairoContext(c)

    carcolors = Dict{Int,Colorant}()
    carcolors[simstate.egoid] = COLOR_CAR_EGO

    for id in simstate.playback_reactive_active_vehicle_ids
        carcolors[id] = colorant"0xF9971FFF" # orange
    end

    #render(ctx, simstate.scene, trajdata.roadway,
    #      cam=CarFollowCamera(simstate.egoid, 5.0), car_colors=carcolors)

    render(ctx, simstate.scene, trajdata.roadway, cam=CarFollowCamera(simstate.egoid, 5.0))

    image
end

function reel_drive(
    gif_filename::AbstractString,
    actions::Matrix{Float64}, # columns are actions
    simparams::SimParams
    )

    frames = Reel.Frames(MIME("image/png"), fps=framerate)

    action = [NaN,NaN]
    push!(frames, render(simparams))
    for frame_index in 1:size(actions,2)

        action[1] = actions[frame_index,1]
        action[2] = actions[frame_index,2]

        step_forward(simparams, action)

        push!(frames, render(simparams))
    end

    Reel.write(gif_filename, frames) # Write to a gif file
end

end # module


include("CommonFuncs.jl")
threshold = 1e7 # σY ≈ 1e8

idx = 0
H_function = spd_Chol_Orth

if length(ARGS)==1
    global idx = parse(Int64, ARGS[1])
elseif length(ARGS)==2
    global idx = parse(Int64, ARGS[1])
    global tid = parse(Int64, ARGS[2])
else
    @warn("No ARGS provided")
end

@show H_function
if H_function==spd_zero_to_H || H_function==spd_Chol_Orth
    global nout = 4
end

if nntype=="ae_scaled"
    global nout = 3
end

if idx == 0
    global config=[20,20,20,nout]
elseif idx == 1
    global config=[20, 20, 20, 20, nout]
elseif idx == 2
    global config=[20, 20, 20, 20, 20, nout]
elseif idx == 3
    global config=[20, 20, 20, 20, 20, 20,nout]
elseif idx == 5
    global config=[nout]
end
printstyled("idx = $idx, config=$config, H_function=$H_function\n", color=:green)

function nn(ε, ε0, σ0) # ε, ε0, σ0 450x3
    local y, z
    global H0
    if nntype=="linear"
        y = ε*H0*stress_scale
        y
    elseif nntype=="ae_scaled"
        x = ε/strain_scale
        if isa(x, Array)
            x = constant(x)
        end
        # y = ae(x, [50,50,50,50,50,50,50,3], nntype)*stress_scale
        # y = ae(x, [50,50,50,50,50,50,50,3], nntype)*stress_scale
        y = ae(x, config, nntype)*stress_scale
    elseif nntype=="stiffmat"
        x = ε/strain_scale
        x = constant(x)

        σ0 = constant(σ0)

        y = ae(x, config, nntype)

        if H_function==spd_H
            z = spd_H(y, H0)
        else
            z = H_function(y)
        end

        σnn = squeeze(tf.matmul(z, tf.reshape((ε-ε0)/strain_scale, (-1,3,1))))

        ####################################
        # σH = (ε-ε0)/strain_scale * H0
        # s = σ0[:,1]^2-σ0[:,1]*σ0[:,2]+σ0[:,2]^2+3*σ0[:,3]^2

        # i = sigmoid(1000*(s-threshold)/1e9)
        # i = [i i i]
        # out = σnn .* i + σH .* (1-i)  + σ0/stress_scale
        ###############################################

        out = σnn + σ0/stress_scale

        out*stress_scale
    elseif nntype=="piecewise"
        x = [ε/strain_scale ε0/strain_scale σ0/stress_scale]
        x = constant(x)
        ε = constant(ε)
        ε0 = constant(ε0)
        σ0 = constant(σ0)

        y = ae(x, config, nntype)

        if H_function==spd_H
            z = spd_H(y, H0)
        else
            z = H_function(y)
        end

	    # op = tf.print(z)
	    # z = bind(z, op)
        # z = sym_H(y)

        σnn = squeeze(tf.matmul(z, tf.reshape((ε-ε0)/strain_scale, (-1,3,1))))
        σH = (ε-ε0)/strain_scale * H0
        s = σ0[:,1]^2-σ0[:,1]*σ0[:,2]+σ0[:,2]^2+3*σ0[:,3]^2

        i = sigmoid(1000*(s-threshold)/1e9)
        i = [i i i]
        out = σnn .* i + σH .* (1-i)  + σ0/stress_scale
        out*stress_scale
    elseif nntype=="doublenn"
        x = [ε/strain_scale ε0/strain_scale σ0/stress_scale]
        x = constant(x)
        ε = constant(ε)
        ε0 = constant(ε0)
        σ0 = constant(σ0)

        y = ae(x, config, nntype)
        i = sigmoid(ae(x, [20,20,1], nntype*"i"))
        # op = tf.print("call spd_H")
        if H_function==spd_H
            z = spd_H(y, H0)
        else
            z = H_function(y)
        end
        # z = bind(z, op)
        # z = sym_H(y)

        σnn = squeeze(tf.matmul(z, tf.reshape((ε-ε0)/strain_scale, (-1,3,1))))
        σH = (ε-ε0)/strain_scale * H0
        s = σ0[:,1]^2-σ0[:,1]*σ0[:,2]+σ0[:,2]^2+3*σ0[:,3]^2

        # i = sigmoid(1000*(s-threshold)/1e9)
        i = [i i i]
        out = σnn .* i + σH .* (1-i)  + σ0/stress_scale
        out*stress_scale
    else
        error("$nntype does not exist")
    end
end

function sigmoid_(z)

    return 1.0 / (1.0 + exp(-z))

end

function nn_helper(ε, ε0, σ0)
    local y1
    if nntype=="linear"
        x = reshape(reshape(ε,1,3)*H0,3,1)
    elseif nntype=="ae_scaled"
        x = reshape(ε/strain_scale,1, 3)
        reshape(nnae_scaled(x)*stress_scale,3,1)
    elseif nntype=="stiffmat"
        ε = ε/strain_scale
        ε0 = ε0/strain_scale
        σ0 = σ0/stress_scale
        x = reshape(ε, 1, 3)
        # y1 = (reshape(ε, 1, 3) - reshape(ε0, 1, 3))*sym_H(nnpiecewise(x))
        if H_function==spd_H
            y1 = (reshape(ε, 1, 3) - reshape(ε0, 1, 3))*spd_H(nnstiffmat(x), H0)
        else
            y1 = (reshape(ε, 1, 3) - reshape(ε0, 1, 3))*H_function(nnstiffmat(x))
        end
        y1 = reshape(y1, 3, 1)
        ################################################
        # y2 = reshape((reshape(ε, 1, 3) - reshape(ε0, 1, 3))*H0, 3,1)
        # s = σ0[1]^2-σ0[1]*σ0[2]+σ0[2]^2+3*σ0[3]^2
        # i = sigmoid_(1000*(s*stress_scale^2 - threshold)/1e9)
        # out = y1 * i + y2 * (1-i)  + reshape(σ0, 3, 1)
        #####################################################

        out = y1  + reshape(σ0, 3, 1)

        out*stress_scale
    elseif nntype=="piecewise"
        ε = ε/strain_scale
        ε0 = ε0/strain_scale
        σ0 = σ0/stress_scale
        x = reshape([ε;ε0;σ0],1, 9)
        # y1 = (reshape(ε, 1, 3) - reshape(ε0, 1, 3))*sym_H(nnpiecewise(x))
        if H_function==spd_H
            y1 = (reshape(ε, 1, 3) - reshape(ε0, 1, 3))*spd_H(nnpiecewise(x), H0)
        else
            y1 = (reshape(ε, 1, 3) - reshape(ε0, 1, 3))*H_function(nnpiecewise(x))
        end
        y1 = reshape(y1, 3, 1)
        y2 = reshape((reshape(ε, 1, 3) - reshape(ε0, 1, 3))*H0, 3,1)
        s = σ0[1]^2-σ0[1]*σ0[2]+σ0[2]^2+3*σ0[3]^2
        i = sigmoid_(1000*(s*stress_scale^2 - threshold)/1e9)
        out = y1 * i + y2 * (1-i)  + reshape(σ0, 3, 1)
        out*stress_scale
    elseif nntype=="doublenn"
        ε = ε/strain_scale
        ε0 = ε0/strain_scale
        σ0 = σ0/stress_scale
        x = reshape([ε;ε0;σ0],1, 9)
        # y1 = (reshape(ε, 1, 3) - reshape(ε0, 1, 3))*sym_H(nnpiecewise(x))
        if H_function==spd_H
            y1 = (reshape(ε, 1, 3) - reshape(ε0, 1, 3))*spd_H(nndoublenn(x), H0)
        else
            y1 = (reshape(ε, 1, 3) - reshape(ε0, 1, 3))*H_function(nndoublenn(x))
        end
        y1 = reshape(y1, 3, 1)
        y2 = reshape((reshape(ε, 1, 3) - reshape(ε0, 1, 3))*H0, 3,1)
        s = σ0[1]^2-σ0[1]*σ0[2]+σ0[2]^2+3*σ0[3]^2
        # i = sigmoid_(1000*(s*stress_scale^2 - threshold)/1e9)
        i = sigmoid_(nndoublenni(x)[1,1])
        out = y1 * i + y2 * (1-i)  + reshape(σ0, 3, 1)
        out*stress_scale
    elseif nntype=="stress"
        ε = ε/strain_scale
        ε0 = ε0/strain_scale
        σ0 = σ0/stress_scale
        x = reshape([ε;ε0;σ0],1, 9)
        out = nnstress(x)
        reshape(out*stress_scale,3,1)
    else
        error("$nntype does not exist")
    end
end

function post_nn(ε, ε0, σ0, Δt)
    f = x -> nn_helper(x, ε0, σ0)
    df = ForwardDiff.jacobian(f, ε)
    return f(ε), df
end