using Pkg
Pkg.activate(normpath(joinpath(@__DIR__, "..")))

using DataFrames
using CSV
using ArgParse
using Logging
using Printf
using ATools
include("../src/ANema.jl")

import ANema: getlorpath

logger = SimpleLogger(stdout, Logging.Warn)
old_logger = global_logger(logger)

function makelors(args)

	rdir   = args["dir"]
	conf   = args["conf"]
	pathf = ANema.pathfinder(rdir, conf)
	println("pathfinder tuple =", pathf)

	n3p = DataFrame(CSV.File(pathf.phipath))
	n3z = DataFrame(CSV.File(pathf.zpath))

	n3mp = ANema.setunits(n3p)
	n3mz = ANema.setunits(n3z)

	println("phi data frame has length of ", nrow(n3mp))
	println("z data frame has length of ", nrow(n3mz))
	println("phi data frame names ", names(n3mp))
	println("z data frame names ", names(n3mz))

	lors_first_true = ANema.dftolor(n3mz, ANema.dtfirst, ANema.postrue)
	ANema.write_lors_hdf5(getlorpath(pathf.lordir,pathf.zlp, "first-true"), lors_first_true)
	lors_minimum_reco = ANema.dftolor(n3mz, ANema.dtminimum, ANema.posreco)
	ANema.write_lors_hdf5(getlorpath(pathf.lordir,pathf.zlp, "minimum-reco"), lors_minimum_reco)
	lors_minimum_recall = ANema.dftolor(n3mz, ANema.dtminimum, ANema.posrecall)
	ANema.write_lors_hdf5(getlorpath(pathf.lordir,pathf.zlp, "minimum-recall"), lors_minimum_recall)
	lors_average_recall = ANema.dftolor(n3mz, ANema.dtaverage, ANema.posrecall)
	ANema.write_lors_hdf5(getlorpath(pathf.lordir,pathf.zlp, "average-recall"), lors_average_recall)

	lors_first_true = ANema.dftolor(n3mp, ANema.dtfirst, ANema.postrue)
	ANema.write_lors_hdf5(getlorpath(pathf.lordir,pathf.philp, "first-true"), lors_first_true)
	lors_minimum_reco = ANema.dftolor(n3mp, ANema.dtminimum, ANema.posreco)
	ANema.write_lors_hdf5(getlorpath(pathf.lordir,pathf.philp, "minimum-reco"), lors_minimum_reco)
	lors_minimum_recall = ANema.dftolor(n3mp, ANema.dtminimum, ANema.posrecall)
	ANema.write_lors_hdf5(getlorpath(pathf.lordir,pathf.philp, "minimum-recall"), lors_minimum_recall)
	lors_average_recall = ANema.dftolor(n3mp, ANema.dtaverage, ANema.posrecall)
	ANema.write_lors_hdf5(getlorpath(pathf.lordir,pathf.philp, "average-recall"), lors_average_recall)

end

function parse_commandline()
    s = ArgParseSettings()

    @add_arg_table s begin
        "--dir", "-d"
            help = "directory with nema data"
            arg_type = String
            default = "../data"
		"--conf", "-c"
            help = "configuration"
            arg_type = String
    end

    return parse_args(s)
end

function main()
	parsed_args = parse_commandline()
	println("Running makelors with arguments", parsed_args)
	makelors(parsed_args)
end

@time main()


#
# Immutable dictionary implemented as a bitmap trie.
#

# Trie nodes.

abstract ImmDictNode{K,V}

immutable ValueNode{K,V} <: ImmDictNode{K,V}
    key::K
    val::V
    key_hc::UInt
end

immutable CollisionNode{K,V} <: ImmDictNode{K,V}
    keys::Vector{K}
    vals::Vector{V}
    key_hc::UInt
end

immutable BitmapNode{K,V} <: ImmDictNode{K,V}
    branches::Vector{ImmDictNode{K,V}}
    bitmap::UInt32

    BitmapNode(branches, bitmap) = new(branches, bitmap)
    BitmapNode() = new(ImmDictNode{K,V}[], 0x00000000)
end

# Lookup in trie nodes.

find{K,V}(node::ValueNode{K,V}, key, key_hc, shift) =
    key_hc == node.key_hc && isequal(key, node.key) ?
        Nullable{V}(node.val) : Nullable{V}()

function find{K,V}(node::CollisionNode{K,V}, key, key_hc, shift)
    if key_hc == node.key_hc
        for j = eachindex(node.keys)
            if isequal(key, node.keys[j])
                return Nullable{V}(node.vals[j])
            end
        end
    end
    return Nullable{V}()
end

mask(hc::UInt, shift::Int) = (hc >>> shift) & 0b11111
bitpos(mask::UInt) = 0x00000001 << mask

function find{K,V}(node::BitmapNode{K,V}, key, key_hc, shift)
    pos = bitpos(mask(key_hc, shift))
    if (node.bitmap & pos) != 0
        j = 1 + count_ones(node.bitmap & (pos-1))
        return find(node.branches[j], key, key_hc, shift+5)::Nullable{V}
    end
    return Nullable{V}()
end

# Iteration.

start{K,V}(node::ValueNode{K,V}) = true
next{K,V}(node::ValueNode{K,V}, flag) = (Pair{K,V}(node.key, node.val), false)
done{K,V}(node::ValueNode{K,V}, flag) = !flag

start{K,V}(node::CollisionNode{K,V}) = 1
next{K,V}(node::CollisionNode{K,V}, k) =
    (Pair{K,V}(node.keys[k], node.vals[k]), k+1)
done{K,V}(node::CollisionNode{K,V}, k) = k > length(node.keys)

function start{K,V}(node::BitmapNode{K,V})
    for j = eachindex(node.branches)
        b = node.branches[j]
        bst = start(b)
        if !done(b, bst)
            return (j, bst)
        end
    end
    return ()
end

function next{K,V}(node::BitmapNode{K,V}, state)
    j, bst = state
    b = node.branches[j]
    pair, bst = next(b, bst)
    if !done(b, bst)
        return (pair, (j, bst))
    end
    while j < length(node.branches)
        j += 1
        b = node.branches[j]
        bst = start(b)
        if !done(b, bst)
            return (pair, (j, bst))
        end
    end
    return (pair, ())
end

done{K,V}(node::BitmapNode{K,V}, state) = state == ()

# Adding an element.

type PutResult{V}
    old::Nullable{V}
    changed::Bool

    PutResult() = new(Nullable{V}(), false)
end

function put{K,V}(node::ValueNode{K,V}, key, val, key_hc, shift, result)
    if key_hc == node.key_hc
        if isequal(key, node.key)
            if isequal(val, node.val)
                return node
            else
                result.old = Nullable{V}(node.val)
                result.changed = true
                return ValueNode{K,V}(key, val, key_hc)
            end
        else
            result.changed = true
            return CollisionNode{K,V}(K[node.key, key], [node.val, val], key_hc)
        end
    else
        branches = ImmDictNode{K,V}[node]
        bitmap = bitpos(mask(node.key_hc, shift))
        return put(BitmapNode{K,V}(branches, bitmap), key, val, key_hc, shift, result)
    end
end

function put{K,V}(node::CollisionNode{K,V}, key, val, key_hc, shift, result)
    if key_hc == node.key_hc
        N = length(node.keys)
        for j = 1:N
            if isequal(key, node.keys[j])
                if isequal(val, node.vals[j])
                    return node
                else
                    result.old = Nullable{V}(node.vals[j])
                    result.changed = true
                    vals = copy(node.vals)
                    vals[j] = val
                    return CollisionNode{K,V}(node.keys, vals, node.key_hc)
                end
            end
        end
        result.changed = true
        keys = Vector{K}(N+1)
        vals = Vector{K}(N+1)
        for j = 1:N
            keys[j] = node.keys[j]
            vals[j] = node.vals[j]
        end
        keys[N+1] = key
        vals[N+1] = val
        return CollisionNode{K,V}(keys, vals, node.key_hc)
    else
        nodes = ImmDictNode{K,V}[node]
        bitmap = bitpos(mask(node.key_hc, shift))
        return put(BitmapNode{K,V}(branches, bitmap), key, val, key_hc, shift, result)
    end
end

function put{K,V}(node::BitmapNode{K,V}, key, val, key_hc, shift, result)
    pos = bitpos(mask(key_hc, shift))
    j = 1 + count_ones(node.bitmap & (pos-1))
    if (node.bitmap & pos) != 0
        j = 1 + count_ones(node.bitmap & (pos-1))
        branch = put(node.branches[j], key, val, key_hc, shift+5, result)
        if result.changed
            branches = copy(node.branches)
            branches[j] = branch
            return BitmapNode{K,V}(branches, node.bitmap)
        else
            return node
        end
    end
    result.changed = true
    branch = ValueNode{K,V}(key, val, key_hc)
    N = length(node.branches)
    branches = Vector{ImmDictNode{K,V}}(N+1)
    for k = 1:j-1
        branches[k] = node.branches[k]
    end
    branches[j] = branch
    for k = j:N
        branches[k+1] = node.branches[k]
    end
    bitmap = node.bitmap | pos
    return BitmapNode{K,V}(branches, bitmap)
end

# Immutable dictionary.

immutable ImmutableDict{K,V} <: Associative{K,V}
    root::ImmDictNode{K,V}
    sz::Int
    hc::UInt

    ImmutableDict() = new(BitmapNode{K,V}(), 0, 0)
    ImmutableDict(root::ImmDictNode{K,V}, sz::Int, hc::UInt) = new(root, sz, hc)
end

ImmutableDict{K,V}(kvs::Pair{K,V}...) =
    let d = ImmutableDict{K,V}()
        for (k, v) in kvs
            d = assoc(d, k, v)
        end
        d
    end

length(idict::ImmutableDict) = idict.sz
isempty(idict::ImmutableDict) = idict.sz == 0

start(idict::ImmutableDict) = start(idict.root)
next(idict::ImmutableDict, state) = next(idict.root, state)
done(idict::ImmutableDict, state) = done(idict.root, state)

get{K,V}(idict::ImmutableDict{K,V}, key) =
    find(idict.root, key, hash(key), 0)::Nullable{V}

get{K,V}(idict::ImmutableDict{K,V}, key, default) =
    let maybe_val = get(idict, key)
        isnull(maybe_val) ? default : get(maybe_val)
    end

getindex{K,V}(idict::ImmutableDict{K,V}, key) =
    let maybe_val = get(idict, key)
        isnull(maybe_val) ? throw(KeyError(key)) : get(maybe_val)
    end

function assoc{K,V}(idict::ImmutableDict{K,V}, key, val)
    result = PutResult{V}()
    new_root = put(idict.root, key, val, hash(key), 0, result)
    if result.changed
        entry_hc = hash(key, hash(val))
        if !isnull(result.old)
            old_entry_hc = hash(key, hash(get(result.old)))
            hc = idict.hc $ entry_hc $ old_entry_hc
            return ImmutableDict{K,V}(new_root, idict.sz, hc)
        else
            sz = idict.sz + 1
            hc = idict.hc $ entry_hc
            return ImmutableDict{K,V}(new_root, sz, hc)
        end
    end
    return idict
end


using DataStructures
using Test
using ChemistryFeaturization
include("../../src/weave_fcns.jl")
using .weave_fcns:
    smiles_atom_features,
    smiles_bond_features,
    chem,
    atom_feat_list,
    bond_feat_bins,
    atom_feat_bins

mol = chem.MolFromSmiles("C1=CC=CC=C1")  # benzene
carbon = get(mol.GetAtoms(), 0)
bond = get(mol.GetBonds(), 0)

@testset "SMILES atom featurization" begin
    @test [k for k in keys(smiles_atom_features(carbon))] == [
        "symbol",
        "degree",
        "implicit_valence",
        "formal_charge",
        "radical_electrons",
        "hybridization",
        "aromaticity",
        "total_H_num",
    ]
    try
        smiles_atom_features(carbon, feature_list = ["invalid_feat"])
    catch e
        @test occursin("is not in supported atomic features list", sprint(showerror, e))
    end
    correct_values = [
        Float32.("C" .== atom_feat_bins["symbol"]),
        Float32.(2 .== atom_feat_bins["degree"]),
        Float32.(1 .== atom_feat_bins["implicit_valence"]),
        Float32(carbon.GetFormalCharge()),
        Float32(carbon.GetNumRadicalElectrons()),
        Float32.(chem.rdchem.HybridizationType.SP2 .== atom_feat_bins["hybridization"]),
        Float32(carbon.GetIsAromatic()),
        Float32.(1 .== atom_feat_bins["total_H_num"]),
    ]
    @test [v for v in values(smiles_atom_features(carbon))] == correct_values
end

@testset "SMILES bond featurization" begin
    @test [k for k in keys(smiles_bond_features(bond))] == ["bond_type", "isConjugated", "isInring"]
    try
        smiles_bond_features(carbon, feature_list = ["invalid_feat"])
    catch e
        @test occursin("is not in supported bond features list", sprint(showerror, e))
    end
    correct_values = [
        Float32.(chem.rdchem.BondType.AROMATIC .== bond_feat_bins["bond_type"]),
        Float32(bond.GetIsConjugated()),
        Float32(bond.IsInRing()),
    ]
    @test [v for v in values(smiles_bond_features(bond))] == correct_values
end


# TODO:
# - use keyword argument to enable/disable showing indices?
# - parse the "status" flag for every node and edge? needed for `des_pipes` (design
#   problem); not sure of other use cases
# * for "epanet-data" objects, identify shutoff valves; might want to do that in
#   WaterModels?

"""
Build pygraphviz graph object from a WaterModels network dictionary parsed from an EPANET
file. If a `solution` dict is provided, it should be for the same time as that of the `data`
dict. The `layout` option equates to the layout functions of graphviz (dot, neato, etc.) and
any provided `args` will also be passed to graphviz. Provide `layout=neato` and `args=-n` to
use the provided graph coordinates for the layout.
"""
function build_graph(data::Dict{String, <:Any},
                     solution::Union{Nothing, Dict{String, <:Any}} = nothing;
                     layout::String="dot", args::String="")
    # presumes input is data dict from WaterModels.parse_file() OR WaterModels.parse_epanet()

    # TODO:
    # - check that input dictionary has the expected fields

    G = pgv.AGraph(strict = false, directed = true)

    # Add nodal data to the directed graph object.
    add_nodes!(G, data["node"])
    node_labels!(G, data["demand"])
    node_labels!(G, data["reservoir"])
    node_labels!(G, data["tank"])

    # Add link data to the directed graph object.
    max_pipe_diameter = maximum([x.second["diameter"] for x in data["pipe"]])
    _add_pipe_to_graph!.(Ref(G), values(data["pipe"]), max_pipe_diameter)
    _add_des_pipe_to_graph!.(Ref(G), values(data["des_pipe"]))
    _add_short_pipe_to_graph!.(Ref(G), values(data["short_pipe"]))
    _add_pump_to_graph!.(Ref(G), values(data["pump"]))
    _add_regulator_to_graph!.(Ref(G), values(data["regulator"]))
    _add_valve_to_graph!.(Ref(G), values(data["valve"]))

    # if solution dict provided
    if solution !== nothing
        add_solution!(G, data, solution)
    end

    # set the layout here--can be reused for displaying different solutions
    try
        G.layout(prog=layout, args=(args))
    catch
        G.layout(prog="dot")
        # by putting this warning after, it should display only if `G.draw()` was actually
        # successful using "dot"; if it fails for another reason, that error should be shown
        @warn "$layout $args is not supported; dot without args was used instead"
    end
    # save the layout information for the user to inspect later
    PyCall.set!(G."graph_attr", "layout_args", "$layout $args")
    # get(G.graph_attr, "layout_args") # this is how to inspect

    return G
end # function build_graph

"""
Update a graph for a different time. Updates reservoir elevation (aka head), demands, and
solution information (if provided).
"""
function update_graph!(G::PyCall.PyObject, data::Dict{String,Any},
                       solution::Union{Nothing, Dict{String, <:Any}} = nothing)
    # TODO:
    # - check that reservoir elevation actually changes for a network that has a reservoir
    #   pattern
    # - check that there aren't any other time-dependent information to update

    # update reservoir elevation
    for (key, res) in data["reservoir"]
        nodekey = string(res["node"])
        node = data["node"][nodekey]
        nodeobj = @pycall G."get_node"(nodekey)::PyObject
        # Color according to the node's elevation.
        elmin = parse(Float64, get(G.graph_attr, "elmin", 0.0))
        elmax = parse(Float64, get(G.graph_attr, "elmax", 1.0))
        clr = HSV(get(viridis, node["elevation"], (elmin, elmax)))
        clrstr = string(clr.h / 360) * " " * string(clr.s) * " " * string(clr.v)
        fntclr = clr.v < 0.6 ? "white" : "black"
        PyCall.set!(nodeobj.attr, "fillcolor", clrstr)
        PyCall.set!(nodeobj.attr, "fontcolor", fntclr)
    end

    # update demands
    node_labels!(G, data["demand"])

    # update solution
    if solution !== nothing
        add_solution!(G, data, solution)
    end
end

"""
Add nodes to the pygraphviz graph object, including node attributes for name label,
elevation, and coordinates
"""
function add_nodes!(G::PyCall.PyObject, nodes::Dict{String,Any})
    # Determine minimum and maximum elevations.
    elevation_min = minimum([x["elevation"] for x in values(nodes)])
    elevation_max = maximum([x["elevation"] for x in values(nodes)])

    # Determine minimum and maximum coordinate values.
    if all(map(x -> haskey(x, "coordinates"), values(nodes)))
        x_min = minimum([x["coordinates"][1] for x in values(nodes)])
        x_max = maximum([x["coordinates"][1] for x in values(nodes)])
        y_min = minimum([x["coordinates"][2] for x in values(nodes)])
        y_max = maximum([x["coordinates"][2] for x in values(nodes)])
    else
        x_min, x_max, y_min, y_max = Inf, -Inf, Inf, -Inf
    end

    # scale to use to achieve "good" relative positioning with graphviz output --
    # calculate this systematically from the number of nodes?
    #scale = 20.0 # works with exclamation points as part of `pos` argument
    scale = 2e3 # no `!` but `-n` flag

    # add elmin and elmax as graph attributes for use in creating a colorbar in
    # colorbar (or via write_visualization)
    PyCall.set!(G."graph_attr", "elmin", elevation_min)
    PyCall.set!(G."graph_attr", "elmax", elevation_max)

    # create and populate pygraphviz nodes
    for (key, node) in nodes
        # Get the label to be used for the node.
        label = _get_comp_label(node)

        # Color according to the node's elevation.
        clr = HSV(get(viridis, node["elevation"], (elevation_min, elevation_max)))

        # Convert clr to a string that is in a form suitable for graphviz.
        clrstr = string(clr.h / 360) * " " * string(clr.s) * " " * string(clr.v)

        # Change font color depending on the background color.
        fntclr = clr.v < 0.6 ? "white" : "black"

        if all(map(x -> haskey(x, "coordinates"), values(nodes)))
            # Get scaled version of the coordinates.
            x = scale * (node["coordinates"][1] - x_min) / (x_max - x_min)
            y = scale * (node["coordinates"][2] - y_min) / (y_max - y_min)

            # Exclamation point forces exact positioning using Neato. But probably better to
            # not have it and use "-n" argument instead with `draw` command
            #pos = string(x) * "," * string(y) * "!"
            pos = string(x) * "," * string(y)

            # Add the node to the graph object with specified attributes.
            G.add_node(node["index"], label = label, elevation = node["elevation"],
                pos = pos, style = "filled", fillcolor = clrstr, fontcolor = fntclr)
        else
            # Add the node to the graph object with specified attributes.
            G.add_node(node["index"], label = label, elevation = node["elevation"],
                style = "filled", fillcolor = clrstr, fontcolor = fntclr)
        end
    end
end

"""Add labels for junctions, tanks, and reservoirs."""
function node_labels!(G::PyCall.PyObject, nodes::Dict{String,Any})
    for (key, node) in nodes
        nodekey = node["node"]
        nodeobj = @pycall G."get_node"(nodekey)::PyObject
        name = node["name"]

        # add node index to the label if different from the name
        if name==string(nodekey)
            label = name
        else
            label = name*" ("*string(nodekey)*")"
        end

        # will the node type ever change from the original source_id description? some link
        # types do (pipes -> valves)
        node_type = node["source_id"][1]

        if node_type == "reservoir"
            PyCall.set!(nodeobj.attr, "label", "Rsvr\\n"*label)
            PyCall.set!(nodeobj.attr, "shape", "diamond")
        elseif node_type == "tank"
            PyCall.set!(nodeobj.attr, "label", "Tank\\n"*label)
            PyCall.set!(nodeobj.attr, "shape", "rectangle")
        else
            dem = node["flow_nominal"]

            if dem == 0
                PyCall.set!(nodeobj.attr, "label", label)
            else
                dem = @sprintf("%2.2g", dem)
                PyCall.set!(nodeobj.attr, "label", label*"\\nd: "*dem)
            end
        end
    end
end

function _get_comp_label(comp::Dict{String, <:Any})
    source_name = comp["source_id"][2]

    if string(comp["index"]) == source_name
        return source_name
    else
        return "$(source_name) ($(string(comp["index"])))"
    end
end

function _get_link_arrowhead(link::Dict{String, <:Any})
    return "normal"
end

function _get_link_dir(link::Dict{String, <:Any})
    if link["flow_direction"] == _WM.UNKNOWN
        return "none"
    elseif link["flow_direction"] == _WM.POSITIVE
        return "forward"
    elseif link["flow_direction"] == _WM.NEGATIVE
        return "back"
    end
end

function _add_pipe_to_graph!(graph::PyCall.PyObject, pipe::Dict{String, <:Any}, max_diameter::Float64)
    index, label = string(pipe["index"]), _get_comp_label(pipe)
    penwidth = pipe["diameter"] * inv(max_diameter) * 10.0
    if penwidth > 8.0
        pad = "  "
    elseif penwidth > 4.0
        pad = " "
    else
        pad = ""
    end
    dir, arrowhead = _get_link_dir(pipe), _get_link_arrowhead(pipe)
    length = @sprintf("%.5g m", pipe["length"])
    # save `pad` for use with adding solution labels
    graph.add_edge(pipe["node_fr"], pipe["node_to"], index, dir = dir, headclip = "true",
                   arrowhead = arrowhead, penwidth = penwidth, pad = pad,
                   label = "$(pad)$(label)\\n$(pad)$(length)")
end

function _add_des_pipe_to_graph!(graph::PyCall.PyObject, des_pipe::Dict{String, <:Any})
    _add_pipe_to_graph!(graph, des_pipe, 1.0)
end

function _add_short_pipe_to_graph!(graph::PyCall.PyObject, short_pipe::Dict{String, <:Any})
    index, label = string(short_pipe["index"]), _get_comp_label(short_pipe)
    dir, arrowhead = _get_link_dir(short_pipe), _get_link_arrowhead(short_pipe)
    graph.add_edge(short_pipe["node_fr"], short_pipe["node_to"], index,
        dir = dir, arrowhead = arrowhead, label = "$(label)")
end

function _add_pump_to_graph!(graph::PyCall.PyObject, pump::Dict{String, <:Any})
    index, label = string(pump["index"]), _get_comp_label(pump)
    dir, arrowhead = _get_link_dir(pump), _get_link_arrowhead(pump)
    graph.add_edge(pump["node_fr"], pump["node_to"], index, dir = dir,
        arrowhead = arrowhead, label = "Pmp\\n$(label)", color = "red", style = "bold")
end

function _add_regulator_to_graph!(graph::PyCall.PyObject, regulator::Dict{String, <:Any})
    index, label = string(regulator["index"]), _get_comp_label(regulator)
    dir, arrowhead = _get_link_dir(regulator), _get_link_arrowhead(regulator)
    graph.add_edge(regulator["node_fr"], regulator["node_to"], index, dir = dir,
        arrowhead = arrowhead, label = "Reg\\n$(label)", color = "purple", style = "bold")
end

function _add_valve_to_graph!(graph::PyCall.PyObject, valve::Dict{String, <:Any})
    index, label = string(valve["index"]), _get_comp_label(valve)
    dir, arrowhead = _get_link_dir(valve), _get_link_arrowhead(valve)
    if dir != "none"
        label = "CV\\n$(label)"
    else
        label = "SV\\n$(label)"
    end
    graph.add_edge(valve["node_fr"], valve["node_to"], index, dir = dir,
        arrowhead = arrowhead, label = label, color = "blue", style = "bold")
end

"""
add solution values to the node and link labels
"""
function add_solution!(G::PyCall.PyObject, data::Dict{String,Any},
                       solution::Dict{String,Any})
    # add head to the node labels (could alternatively show pressure)
    for (key,node) in solution["node"]
        head = @sprintf("%2.2g", node["h"])
        nodeobj = @pycall G."get_node"(key)::PyObject
        label = get(nodeobj.attr, "label")
        # remove existing solution string (if it exists)
        idx = findlast("\\n", label)
        if idx != nothing
            idx = idx[1]
            #if contains(label[idx:end], "h:") # requires Julia >=1.5
            if occursin("h:", label[idx:end])
                label = label[1:idx-1]
            end
        end
        PyCall.set!(nodeobj.attr, "label", label*"\\nh: "*head)
    end
    # add flow to the labels for pipes and valves
    # Byron added this set, but not all of these fields exist in the solution object
    #pipesplus = ["pipe", "des_pipe", "pump", "regulator", "short_pipe", "valve"]
    links = ["pump", "pipe", "short_pipe", "valve"]
    for linktype in links
        for (key,pipesol) in solution[linktype]
            flow = _val_string_cut(pipesol["q"], 1e-10)
            link = data[linktype][key]
            # may also need to use `key` if multiple pipes between nodes
            edgeobj = @pycall G."get_edge"(link["node_fr"], link["node_to"])::PyObject
            label = get(edgeobj.attr, "label")
            # remove existing solution string (if it exists)
            idx = findlast("\\n", label)
            if idx != nothing
                idx = idx[1]
                #if contains(label[idx:end], "q:") # requires Julia >=1.5
                if occursin("q:", label[idx:end])
                    label = label[1:idx-1]
                end
            end
            pad = get(edgeobj.attr, "pad")
            PyCall.set!(edgeobj.attr, "label", label*"\\n$(pad)q: "*flow)
        end
    end
end

function _val_string_cut(val::Real, cut::Real)
    if val < cut
        return "0"
    else
        return @sprintf("%2.2g", val)
    end
end


abstract type AbstractScenarioProblems{S <: AbstractScenario} end

abstract type AbstractBlockStructure{N} <: AbstractStochasticStructure{N} end

# MOI #
# ========================== #
function MOI.is_valid(structure::AbstractBlockStructure, index::MOI.VariableIndex, stage::Integer, scenario_index::Integer)
    stage == 1 && error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    return MOI.is_valid(scenarioproblems(structure, stage), index, scenario_index)
end

function MOI.delete(structure::AbstractBlockStructure{N}, index::MOI.VariableIndex, stage::Integer, scenario_index::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    MOI.delete(scenarioproblems(structure, stage), index, scenario_index)
    return nothing
end
function MOI.delete(structure::AbstractBlockStructure{N}, indices::Vector{MOI.VariableIndex}, stage::Integer, scenario_index::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    MOI.delete(scenarioproblems(structure, stage), indices, scenario_index)
    return nothing
end

# JuMP #
# ========================== #
function JuMP.objective_function_type(structure::AbstractBlockStructure{N}, stage::Integer, scenario_index::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    return objective_function_type(scenarioproblems(structure, stage), scenario_index)
end

function JuMP.objective_function(structure::AbstractBlockStructure{N},
                                 proxy::JuMP.Model,
                                 stage::Integer,
                                 scenario_index::Integer,
                                 FunType::Type{<:AbstractJuMPScalar}) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    objective::FunType = objective_function(scenarioproblems(structure, stage), proxy, scenario_index, FunType)
    return objective
end

function JuMP._moi_optimizer_index(structure::AbstractBlockStructure, index::VI, scenario_index::Integer)
    return JuMP._moi_optimizer_index(scenarioproblems(structure), index, scenario_index)
end
function JuMP._moi_optimizer_index(structure::AbstractBlockStructure, ci::CI, scenario_index::Integer)
    return JuMP._moi_optimizer_index(scenarioproblems(structure), ci, scenario_index)
end

function DecisionRef(proxy::JuMP.Model, structure::AbstractBlockStructure, index::VI, scenario_index::Integer)
    return DecisionRef(proxy, index)
end

# Getters #
# ========================== #
function scenarioproblems(structure::AbstractBlockStructure{N}, stage::Integer) where N
    1 < stage <= N || error("Stage $stage not in range 2 to $N.")
    stage == 1 && error("Stage 1 does not have scenario problems.")
    N == 2 && (stage == 2 || error("Stage $stage not available in two-stage model."))
    return structure.scenarioproblems[stage-1]
end
function scenarioproblems(structure::AbstractBlockStructure{2})
    return scenarioproblems(structure, 2)
end
function scenario_types(structure::AbstractBlockStructure{N}) where N
    return ntuple(Val{N-1}()) do s
        scenario_type(scenarioproblems(structure), s + 1)
    end
end
function decision(structure::AbstractBlockStructure{N}, index::MOI.VariableIndex, stage::Integer, scenario_index::Integer) where N
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    return decision(scenarioproblems(structure, stage), index, scenario_index)
end
function scenario(structure::AbstractBlockStructure, stage::Integer, scenario_index::Integer)
    scenario(scenarioproblems(structure, stage), scenario_index)
end
function scenario(structure::AbstractBlockStructure{2}, scenario_index::Integer)
    scenario(scenarioproblems(structure, 2), scenario_index)
end
function scenarios(structure::AbstractBlockStructure, stage::Integer)
    scenarios(scenarioproblems(structure, stage))
end
function scenarios(structure::AbstractBlockStructure{2})
    scenarios(scenarioproblems(structure, 2))
end
function expected(structure::AbstractBlockStructure, stage::Integer)
    return expected(scenarioproblems(structure, stage))
end
function expected(structure::AbstractBlockStructure{2})
    return expected(scenarioproblems(structure, 2))
end
function scenario_type(structure::AbstractBlockStructure, stage::Integer)
    return scenario_type(scenarioproblems(structure, stage))
end
function scenario_type(structure::AbstractBlockStructure{2})
    return scenario_type(scenarioproblems(structure, 2))
end
function probability(structure::AbstractBlockStructure, stage::Integer, scenario_index::Integer)
    return probability(scenarioproblems(structure, stage), scenario_index)
end
function probability(structure::AbstractBlockStructure{2}, scenario_index::Integer)
    return probability(scenarioproblems(structure, 2), scenario_index)
end
function stage_probability(structure::AbstractBlockStructure, stage::Integer)
    return probability(scenarioproblems(structure, stage))
end
function stage_probability(structure::AbstractBlockStructure{2})
    return probability(scenarioproblems(structure, 2))
end
function subproblem(structure::AbstractBlockStructure, stage::Integer, scenario_index::Integer)
    return subproblem(scenarioproblems(structure, stage), scenario_index)
end
function subproblem(structure::AbstractBlockStructure{2}, scenario_index::Integer)
    return subproblem(scenarioproblems(structure, 2), scenario_index)
end
function subproblems(structure::AbstractBlockStructure, stage::Integer)
    return subproblems(scenarioproblems(structure, stage))
end
function subproblems(structure::AbstractBlockStructure{2})
    return subproblems(scenarioproblems(structure, 2))
end
function num_subproblems(structure::AbstractBlockStructure, stage::Integer)
    return num_subproblems(scenarioproblems(structure, stage))
end
function num_subproblems(structure::AbstractBlockStructure{2})
    return num_subproblems(scenarioproblems(structure, 2))
end
function num_scenarios(structure::AbstractBlockStructure, stage::Integer)
    return num_scenarios(scenarioproblems(structure, stage))
end
function num_scenarios(structure::AbstractBlockStructure{2})
    return num_scenarios(scenarioproblems(structure, 2))
end
deferred(structure::AbstractBlockStructure{N}) where N = deferred(structure, Val(N))
deferred(structure::AbstractBlockStructure, ::Val{1}) = deferred_first_stage(structure)
function deferred(structure::AbstractBlockStructure, ::Val{N}) where N
    return deferred_stage(structure, N) || deferred(structure, Val(N-1))
end
deferred_first_stage(structure::AbstractBlockStructure) = false
function deferred_stage(structure::AbstractBlockStructure{N}, stage::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage == 1 && return deferred_first_stage(structure)
    num_subproblems(structure, stage) < num_scenarios(structure, stage)
end
function distributed(structure::AbstractBlockStructure{N}, stage::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage == 1 && return false
    return distributed(scenarioproblems(structure, stage))
end

# ========================== #

# Setters
# ========================== #
function update_known_decisions!(structure::AbstractBlockStructure, change::DecisionModification, stage::Integer, scenario_index::Integer)
    update_known_decisions!(scenarioproblems(structure, stage), change, scenario_index)
    return nothing
end
function add_scenario!(structure::AbstractBlockStructure, stage::Integer, scenario::AbstractScenario)
    add_scenario!(scenarioproblems(structure, stage), scenario)
    return nothing
end
function add_worker_scenario!(structure::AbstractBlockStructure, stage::Integer, scenario::AbstractScenario, w::Integer)
    add_scenario!(scenario(structure, stage), scenario, w)
    return nothing
end
function add_scenario!(scenariogenerator::Function, structure::AbstractBlockStructure, stage::Integer)
    add_scenario!(scenariogenerator, scenarioproblems(structure, stage))
    return nothing
end
function add_worker_scenario!(scenariogenerator::Function, structure::AbstractBlockStructure, stage::Integer, w::Integer)
    add_scenario!(scenariogenerator, scenarioproblems(structure, stage), w)
    return nothing
end
function add_scenarios!(structure::AbstractBlockStructure, stage::Integer, scenarios::Vector{<:AbstractScenario})
    add_scenarios!(scenarioproblems(structure, stage), scenarios)
    return nothing
end
function add_worker_scenarios!(structure::AbstractBlockStructure, stage::Integer, scenarios::Vector{<:AbstractScenario}, w::Integer)
    add_scenarios!(scenarioproblems(structure, stage), scenarios, w)
    return nothing
end
function add_scenarios!(scenariogenerator::Function, structure::AbstractBlockStructure, stage::Integer, n::Integer)
    add_scenarios!(scenariogenerator, scenarioproblems(structure, stage), n)
    return nothing
end
function add_worker_scenarios!(scenariogenerator::Function, structure::AbstractBlockStructure, stage::Integer, n::Integer, w::Integer)
    add_scenarios!(scenariogenerator, scenarioproblems(structure, stage), n, w)
    return nothing
end
function sample!(structure::AbstractBlockStructure, stage::Integer, sampler::AbstractSampler, n::Integer)
    sample!(scenarioproblems(structure, stage), sampler, n)
    return nothing
end
# ========================== #

# Includes
# ========================== #
include("scenarioproblems.jl")
include("vertical.jl")
include("horizontal.jl")


# Copyright (c) 2020-2021, Bank of Canada
# All rights reserved.

using Base.Broadcast

# custom broadcasting style for TSeries
struct MVTSeriesStyle{F<:Frequency} <: Broadcast.BroadcastStyle end
@inline frequencyof(::Type{<:MVTSeriesStyle{F}}) where {F<:Frequency} = F

@inline Base.Broadcast.BroadcastStyle(::Type{<:MVTSeries{F}}) where {F<:Frequency} = MVTSeriesStyle{F}()
@inline Base.Broadcast.BroadcastStyle(S1::MVTSeriesStyle, S2::MVTSeriesStyle) = mixed_freq_error(S1, S2)
@inline Base.Broadcast.BroadcastStyle(S::MVTSeriesStyle{F}, ::MVTSeriesStyle{F}) where {F<:Frequency} = S

# mixing with TSeries - we treat is as a vector column
@inline Base.Broadcast.BroadcastStyle(S1::MVTSeriesStyle, S2::TSeriesStyle) = mixed_freq_error(S1, S2)
@inline Base.Broadcast.BroadcastStyle(S::MVTSeriesStyle{F}, ::TSeriesStyle{F}) where {F<:Frequency} = S

# mixing with other things - we take over as long as we have 2 dimensons or less
@inline Base.Broadcast.BroadcastStyle(S::MVTSeriesStyle, AS::Broadcast.AbstractArrayStyle{N}) where {N} = N <= 2 ? S : AS

function Base.similar(bc::Broadcast.Broadcasted{<:MVTSeriesStyle}, ::Type{ElType}) where {ElType}
    if !isa(bc.axes, Nothing) && length(bc.axes) == 2
        similar(Vector, ElType, bc.axes)
    else
        throw(DimensionMismatch("Cannot create a MVTSeries with axes $(bc.axes)."))
    end
end

# disable extrusion
@inline Base.Broadcast.extrude(x::MVTSeries) = x
@inline Base.Broadcast.extrude(x::TSeries) = x

const MVTSeriesIndexType = Tuple{UnitRange{<:MIT},NTuple{N,Symbol}} where {N}

@inline mvts_get_index(x::Number, p::MIT, c::Symbol) = x
@inline mvts_get_index(x::MVTSeries, p::MIT, c::Symbol) = x[p, c]
@inline mvts_get_index(x::TSeries, p::MIT, c::Symbol) = x[p]
@inline mvts_get_index(x::Base.Broadcast.Broadcasted, p::MIT, c::Symbol) = x[p, c]

function Base.Broadcast.getindex(bc::Base.Broadcast.Broadcasted, p::MIT, c::Symbol)
    args = (mvts_get_index(arg, p, c) for arg in bc.args)
    return bc.f(args...)
end

function Base.copyto!(dest::MVTSeries, bc::Broadcast.Broadcasted{Nothing})
    bc′ = Base.Broadcast.preprocess(dest, bc)
    bcrng, bcvars = axes(bc)
    xrng, xvars = axes(dest)
    for (p, c) in Iterators.product(intersect(bcrng, xrng), intersect(bcvars, xvars))
        @inbounds dest[p, c] = bc′[p, c]
    end
    return dest
end

function do_instantiate(bc)
    shape = mvts_combine_axes(bc.args...)
    if bc.axes !== nothing
        shape = mvts_broadcast_shape(bc.axes, shape)
    end
    mvts_instantiate(bc, shape)
end

@inline function Base.Broadcast.instantiate(bc::Base.Broadcast.Broadcasted{S,A}) where {N,S<:Base.Broadcast.BroadcastStyle,A<:Tuple{UnitRange{<:MIT},NTuple{N,Symbol}}}
    do_instantiate(bc)
end

@inline function Base.Broadcast.instantiate(bc::Base.Broadcast.Broadcasted{S,A}) where {N,S<:Base.Broadcast.AbstractArrayStyle{0},A<:Tuple{UnitRange{<:MIT},NTuple{N,Symbol}}}
    do_instantiate(bc)
end

@inline function Base.Broadcast.instantiate(bc::Base.Broadcast.Broadcasted{S,Nothing}) where {S<:MVTSeriesStyle}
    do_instantiate(bc)
end

# combine_axes works recursively
@inline mvts_combine_axes(A) = axes(A)
@inline mvts_combine_axes(A, B) = mvts_broadcast_shape(axes(A), axes(B))
@inline mvts_combine_axes(A, B...) = mvts_broadcast_shape(axes(A), mvts_combine_axes(B...))

@inline mvts_broadcast_shape(shape::Tuple) = shape
@inline mvts_broadcast_shape(::Tuple{}, shape::Tuple{}) = shape
@inline mvts_broadcast_shape(::Tuple{}, shape::Tuple) = shape
@inline mvts_broadcast_shape(shape::Tuple, ::Tuple{}) = shape
@inline mvts_broadcast_shape(shape1::Tuple{A}, shape2::Tuple) where {A} = mvts_broadcast_shape(shape2, shape1)
function mvts_broadcast_shape(shape1::Tuple, shape2::Tuple)
    if length(shape1) > 2 || length(shape2) > 2
        throw(ArgumentError("broadcasting MVTSeries with ndims > 2."))
    end
    return (mit_common_axes(shape1[1], shape2[1]),
        length(shape2) == 1 ? shape1[2] : sym_common_axes(shape1[2], shape2[2]),
    )
end

@inline mit_common_axes(a::UnitRange{<:MIT}, b::Base.OneTo) =
    length(b) == 1 || length(b) == length(a) ? a :
    throw(DimensionMismatch("Cannot broadcast with $a and $b."))

# mit_common_axes is the same as in TSeries
@inline sym_common_axes(a::NTuple{N1,Symbol}, b::NTuple{N2,Symbol}) where {N1,N2} = a == b ? a : tuple(intersect(a, b)...)
@inline sym_common_axes(a::NTuple{N,Symbol}, b::Any) where {N} = length(a) == length(b) && first(b) == 1 ? a : throw(DimensionMismatch("Cannot broadcast with $(a) and $b."))
@inline sym_common_axes(a::Any, b::NTuple{N,Symbol}) where {N} = sym_common_axes(b, a)

function mvts_instantiate(bc::Base.Broadcast.Broadcasted{S}, shape) where {S<:Base.Broadcast.BroadcastStyle}
    args = map(bc.args) do arg
        mvts_check_axes(shape, arg)
    end
    return Base.Broadcast.Broadcasted{S}(bc.f, args, shape)
end

struct MVTSBroadcasted{Axes,BC}
    singleton::Tuple{Bool,Bool}
    shape::Axes
    bc::BC
    function MVTSBroadcasted(shape, bc)
        bcshape = axes(bc)
        singleton = length(bcshape) == 0 ? (true, true) :
                    length(bcshape) == 1 ? (length(bcshape[1]) == 1, true) :
                    (length(bcshape[1]) == 1, length(bcshape[2]) == 1)
        # bc1 = Base.Broadcast.instantiate(bc)
        new{typeof(shape),typeof(bc)}(singleton, shape, bc)
    end
end

@inline Base.eltype(x::MVTSBroadcasted) = Base.Broadcast.combine_eltypes(x.bc.f, x.bc.args)

function mvts_get_index(x::MVTSBroadcasted, p::MIT, c::Symbol)
    ip = x.singleton[1] ? 1 : Int(p - first(x.shape[1]) + 1)
    ic = x.singleton[2] ? 1 : indexin([c], collect(x.shape[2]))[1]
    ind = CartesianIndex(ip, ic)
    x.bc[ind]
end

# nested broadcasts are processed recursively
@inline mvts_check_axes(shape::MVTSeriesIndexType, bc::Base.Broadcast.Broadcasted{<:MVTSeriesStyle}) = mvts_instantiate(bc, shape)
@inline mvts_check_axes(shape::MVTSeriesIndexType, bc::Base.Broadcast.Broadcasted{<:TSeriesStyle}) = mvts_instantiate(bc, shape)
@inline mvts_check_axes(shape::MVTSeriesIndexType, bc::Base.Broadcast.Broadcasted) = MVTSBroadcasted(shape, bc)

# MVTSs are left alone. If axes are wrong it'll error when indexing later.
@inline mvts_check_axes(shape::MVTSeriesIndexType, x::MVTSeries) = x
# TSeries are also left alone
@inline mvts_check_axes(shape::MVTSeriesIndexType, x::TSeries) = x
# Leave numbers alone too
@inline mvts_check_axes(shape::MVTSeriesIndexType, x::Number) = x
# For Vector, we wrap it in a TSeries
@inline mvts_check_axes(shape::MVTSeriesIndexType, x::AbstractVector) = TSeries(shape[1], x)
# For Matrix, we wrap them in an MVTSeries with the same dimensions
@inline mvts_check_axes(shape::MVTSeriesIndexType, x::AbstractMatrix) = MVTSeries(first(shape[1]), shape[2], x)

function Base.axes(bc::Base.Broadcast.Broadcasted{<:MVTSeriesStyle})
    bc.axes === nothing ? mvts_combine_axes(bc.args...) : bc.axes
end

function Base.axes(bc::Base.Broadcast.Broadcasted{<:MVTSeriesStyle}, d::Integer)
    1 <= d <= 2 ? axes(bc)[d] : Base.OneTo(1)
end


"Causal Inductive Synthesis Corpus"
module CISC

f(x) = abs(x)

end


function sparse_from_coo(::Type{<:SparseMatrixCSR{Bi}}, I,J,V,m,n) where Bi
  sparsecsr(Val(Bi),I,J,V,m,n)
end

@inline function is_entry_stored(::Type{<:SparseMatrixCSR},i,j)
  true
end

function finalize_coo!(::Type{<:SparseMatrixCSR},I,J,V,m,n)
  nothing
end

function nz_index(A::SparseMatrixCSR,i0,i1)
  if !(1 <= i0 <= size(A, 1) && 1 <= i1 <= size(A, 2)); throw(BoundsError()); end
  o = getoffset(A)
  Bi = getBi(A)
  r1 = Int(A.rowptr[i0]+o)
  r2 = Int(A.rowptr[i0+1]-Bi)
  (r1 > r2) && return -1
  i1o = i1-o
  k = searchsortedfirst(colvals(A), i1o, r1, r2, Base.Order.Forward)
  ((k > r2) || (colvals(A)[k] != i1o)) ? -1 : k
end

@inline function push_coo!(::Type{<:SparseMatrixCSR},I,J,V,i,j,v)
 push!(I,i)
 push!(J,j)
 push!(V,v)
 nothing
end


using Base.Test
using AutoEnvs

function test_basics()
    # ctor
    filepath = Pkg.dir("NGSIM", "data", "trajdata_i80_trajectories-0400-0415.txt")
    n_veh = 2
    params = Dict(
        "trajectory_filepaths"=>[filepath],
        "H"=>50,
        "n_veh"=>n_veh
    )
    env = MultiagentNGSIMEnv(params)

    # reset, step
    x = reset(env)
    @test size(x) == (2, 66)
    a = zeros(n_veh, 2)
    nx, r, terminal, infos = step(env, a)
    @test x != nx
    @test terminal == [false, false]
    nnx, r, terminal, infos = step(env, a)
    @test nx != nnx

    # obs spec
    shape, spacetype, infos = observation_space_spec(env)
    @test spacetype == "Box"
    @test in("high", keys(infos))
    @test in("low", keys(infos))

    # does accel reflect applied value
    fns = obs_names(env)
    acc_idx = [i for (i,n) in enumerate(fns) if "accel" == n][1]
    tur_idx = [i for (i,n) in enumerate(fns) if "turn_rate_global" == n][1]
    a = zeros(n_veh, 2)
    a[:, 2] = 1
    for _ in 1:10
        nx, _, _, _ = step(env, a)
    end

    @test all(abs.(nx[:, acc_idx]) .<= 1e-1)
    @test all(abs.(nx[:, tur_idx] - 1) .<= 1e-1)

    # test infos
    reset(env)
    _, _, _, infos = step(env, ones(n_veh, 2))
    @test all(infos["rmse_pos"] .!= 0.)

    # test multiple episodes
    reset(env, [true for _ in 1:n_veh])
    for _ in 1:50
        nx, _, terminal, _ = step(env, a)
    end
    @test all(terminal)
    _, _, terminal, _ = step(env, a)
    @test all(.!terminal)
    for _ in 1:49
        _, _, terminal, _ = step(env, a)
    end
    @test all(terminal)

end

function test_render()
    srand(2)
    n_veh = 50
    filepath = Pkg.dir("NGSIM", "data", "trajdata_i101_trajectories-0805am-0820am.txt")
    params = Dict(
        "trajectory_filepaths"=>[filepath],
        "H"=>50,
        "primesteps"=>50,
        "n_veh"=>n_veh,
        "remove_ngsim_veh"=>true
    )
    env = MultiagentNGSIMEnv(params)

    reset(env)
    a = ones(n_veh, 2)
    for _ in 1:50
        render(env)
        x, r, terminal, _ = step(env, a)
    end
end

@time test_basics()
# manual tests
# @time test_render()


"Generic state difference represented as additions, deletions, and assignments."
struct GenericDiff <: Diff
    add::Vector{Term} # List of additions
    del::Vector{Term} # List of deletions
    ops::Dict{Term,Term} # Dictionary of assignment operations
end

GenericDiff() = GenericDiff(Term[], Term[], Dict{Term,Any}())

function combine!(d1::GenericDiff, d2::GenericDiff)
    append!(d1.add, d2.add)
    append!(d1.del, d2.del)
    merge!(d1.ops, d2.ops)
    return d1
end

function as_term(diff::GenericDiff)
    add = diff.add
    del = [Compound(:not, [t]) for t in diff.del]
    ops = [Compound(:assign, [t, v]) for (t, v) in diff.ops]
    return Compound(:and, [add; del; ops])
end

is_redundant(diff::GenericDiff) =
    issetequal(diff.add, diff.del) && all(k == v for (k, v) in diff.ops)

Base.empty(diff::GenericDiff) = GenericDiff()

"Conditional state difference, represented as paired conditions and sub-diffs."
struct ConditionalDiff{D <: Diff} <: Diff
    conds::Vector{Vector{Term}}
    diffs::Vector{D}
end

ConditionalDiff{D}() where {D} = ConditionalDiff{D}([], Vector{D}())

function combine!(d1::ConditionalDiff{D}, d2::D) where {D}
    for (i, cs) in enumerate(d1.conds)
        if isempty(cs)
            combine!(d1.diffs[i], d2)
            return d1
        end
    end
    push!(d1.conds, Term[])
    push!(d1.diffs, d2)
    return d1
end

function combine!(d1::ConditionalDiff{D}, d2::ConditionalDiff{D}) where {D}
    append!(d1.conds, d2.conds)
    append!(d1.diffs, d2.diffs)
    return d1
end

function as_term(diff::ConditionalDiff)
    branches = map(zip(diff.conds, diff.diffs)) do (cs, d)
        effect = as_term(d)
        isempty(cs) && return effect
        cond = length(cs) == 1 ? cs[1] : Compound(:and, cs)
        return Compound(:when, [cond, effect])
    end
    return Compound(:and, branches)
end

is_redundant(diff::ConditionalDiff) = all(is_redundant.(diff.diffs))

Base.empty(diff::ConditionalDiff{D}) where {D} = ConditionalDiff{D}()


using CamelCup

game = CamelCup.Game(6)
p, st = CamelCup.setup!(game)
while !game.over
    p, st = CamelCup.taketurn!(game, p[1].action)
    show(p)
end
p
st
p, st = CamelCup.taketurn!(game, CamelCup.TakeBet(CamelCup.Green))
p, st = CamelCup.taketurn!(game, CamelCup.Oasis(5, true))
p, st = CamelCup.taketurn!(game, CamelCup.Roll())

p, st = CamelCup.taketurn!(game, CamelCup.FinalBet(CamelCup.White, true))

p, st = CamelCup.taketurn!(game, CamelCup.Roll())
p, st = CamelCup.taketurn!(game, CamelCup.Roll())
p, st = CamelCup.taketurn!(game, CamelCup.Roll())
CamelCup.taketurn!(game, CamelCup.Roll())


#############################################################
# Data Layers
#############################################################
include("layers/hdf5-data.jl")

#############################################################
# General Computation Layers
#############################################################
include("layers/inner-product.jl")
include("layers/convolution.jl")
include("layers/pooling.jl")
include("layers/softmax.jl")
include("layers/power.jl")
include("layers/split.jl")
include("layers/element-wise.jl")
include("layers/channel-pooling.jl")
include("layers/dropout.jl")
include("layers/argmax.jl")
include("layers/crop.jl")
include("layers/concat.jl")
include("layers/tied-inner-product.jl")
include("layers/index2onehot.jl")
include("layers/random-normal.jl")

#############################################################
# Loss Layers
#############################################################
include("layers/square-loss.jl")
include("layers/hinge-loss.jl")
include("layers/multinomial-logistic-loss.jl")
include("layers/softmax-loss.jl")
include("layers/softlabel-softmax-loss.jl")
include("layers/wasserstein-loss.jl")
include("layers/binary-cross-entropy-loss.jl")
include("layers/gaussian-kl-loss.jl")

#############################################################
# Statistics Layers
#############################################################
include("layers/accuracy.jl")
include("layers/binary-accuracy.jl")


using MonteCarloMeasurements: Particles

@testset "from_namedtuple" begin
    rng = MersenneTwister(42)

    nchains, ndraws = 4, 10
    sizes = (x=(), y=(2,), z=(3, 4))
    dims = Dict("y" => ["yx"], "z" => ["zx", "zy"])
    coords = Dict("yx" => ["y1", "y2"], "zx" => 1:3, "zy" => 1:4)

    nts = [
        "NamedTuple" =>
            (; (k => randn(rng, nchains, ndraws, v...) for (k, v) in pairs(sizes))...),
        "Vector{NamedTuple}" => [
            (; (k => randn(rng, ndraws, v...) for (k, v) in pairs(sizes))...) for
            _ in 1:nchains
        ],
        "Matrix{NamedTuple}" => [
            (; (k => randn(rng, v...) for (k, v) in pairs(sizes))...) for
            _ in 1:nchains, _ in 1:ndraws
        ],
        "Vector{Vector{NamedTuple}}" => [
            [(; (k => randn(rng, v...) for (k, v) in pairs(sizes))...) for _ in 1:ndraws] for _ in 1:nchains
        ],
        "Vector{NamedTuple} particles" => [
            (; (k => Particles(randn(rng, ndraws, v...)) for (k, v) in pairs(sizes))...) for _ in 1:nchains
        ],
    ]

    @testset "posterior::$(type)" for (type, nt) in nts
        idata1 = from_namedtuple(nt; dims=dims, coords=coords, library="MyLib")
        test_namedtuple_data(
            idata1,
            :posterior,
            keys(sizes),
            nchains,
            ndraws;
            library="MyLib",
            coords=coords,
            dims=dims,
        )

        idata2 = convert_to_inference_data(nt; dims=dims, coords=coords, library="MyLib")
        test_namedtuple_data(
            idata2,
            :posterior,
            keys(sizes),
            nchains,
            ndraws;
            library="MyLib",
            coords=coords,
            dims=dims,
        )
    end

    @testset "$(group)" for group in [
        :posterior_predictive, :sample_stats, :predictions, :log_likelihood
    ]
        @testset "::$(type)" for (type, nt) in nts
            idata1 = convert_to_inference_data(
                nt; (group => nt,)..., dims=dims, coords=coords, library="MyLib"
            )
            test_namedtuple_data(
                idata1,
                :posterior,
                keys(sizes),
                nchains,
                ndraws;
                library="MyLib",
                coords=coords,
                dims=dims,
            )
            test_namedtuple_data(
                idata1,
                group,
                keys(sizes),
                nchains,
                ndraws;
                library="MyLib",
                coords=coords,
                dims=dims,
            )

            idata2 = convert_to_inference_data(
                nt; (group => keys(sizes),)..., dims=dims, coords=coords, library="MyLib"
            )
            test_namedtuple_data(
                idata2,
                group,
                keys(sizes),
                nchains,
                ndraws;
                library="MyLib",
                coords=coords,
                dims=dims,
            )
        end
    end

    @testset "$(group)" for group in [:prior_predictive, :sample_stats_prior]
        @testset "::$(type)" for (type, nt) in nts
            idata1 = from_namedtuple(;
                prior=nt, (group => nt,)..., dims=dims, coords=coords, library="MyLib"
            )
            test_namedtuple_data(
                idata1,
                :prior,
                keys(sizes),
                nchains,
                ndraws;
                library="MyLib",
                coords=coords,
                dims=dims,
            )
            test_namedtuple_data(
                idata1,
                group,
                keys(sizes),
                nchains,
                ndraws;
                library="MyLib",
                coords=coords,
                dims=dims,
            )

            idata2 = from_namedtuple(;
                prior=nt,
                (group => keys(sizes),)...,
                dims=dims,
                coords=coords,
                library="MyLib",
            )
            test_namedtuple_data(
                idata2,
                group,
                keys(sizes),
                nchains,
                ndraws;
                library="MyLib",
                coords=coords,
                dims=dims,
            )
        end
    end

    @testset "$(group)" for group in
                            [:observed_data, :constant_data, :predictions_constant_data]
        _, nt = nts[1]
        idata = from_namedtuple(
            nt;
            (group => Dict("w" => [1.0, 2.0]),)...,
            dims=Dict("w" => ["wx"]),
            coords=Dict("wx" => 1:2),
            library="MyLib",
        )
        @test idata isa InferenceData
        @test group in ArviZ.groupnames(idata)
        ds = getproperty(idata, group)
        sizes = dimsizes(ds)
        @test length(sizes) == 1
        @test "w" in keys(vardict(ds))
        @test dimdict(ds)["w"] == ("wx",)
        @test "inference_library" in keys(attributes(ds))
        @test attributes(ds)["inference_library"] == "MyLib"

        # ensure that dims are matched to named tuple keys
        # https://github.com/arviz-devs/ArviZ.jl/issues/96
        idata = from_namedtuple(
            nt;
            (group => (w=[1.0, 2.0],),)...,
            dims=Dict("w" => ["wx"]),
            coords=Dict("wx" => 1:2),
            library="MyLib",
        )
        @test idata isa InferenceData
        @test group in ArviZ.groupnames(idata)
        ds = getproperty(idata, group)
        sizes = dimsizes(ds)
        @test length(sizes) == 1
        @test "w" in keys(vardict(ds))
        @test dimdict(ds)["w"] == ("wx",)
        @test "inference_library" in keys(attributes(ds))
        @test attributes(ds)["inference_library"] == "MyLib"
    end
end


using ArgParse
using DrWatson
@quickactivate
using CSV
using BSON
using Random
using FileIO
using DataFrames
using PrettyTables
using PrettyTables.Crayons

using GenerativeAD.Evaluation: _prefix_symbol, PAT_METRICS, aggregate_stats_mean_max
using GenerativeAD.Evaluation: rank_table, print_rank_table

s = ArgParseSettings()
@add_arg_table! s begin
	"filename"
		arg_type = String
		default = "evaluation/images_eval.bson"
		help = "Location of cached DataFrame."
	"-c", "--criterion-metric"
		arg_type = String
		default = "val_auc"
		help = "Criterion to sort the models."
	"-r", "--rank-metric"
		arg_type = String
		default = ""
		help = "Metric to rank models."
	"-b", "--backend"
		arg_type = String
		default = "txt"
		help = "Backend for PrettyTable print. Either of [txt (default), tex, html] is allowed."
	"-o", "--output-prefix"
		arg_type = String
		default = "evaluation/images_eval"
		help = "Output prefix for storing results."
	"-v", "--verbose"
		action = :store_true
		help = "Print all results instead of storing to files."
	"-p", "--proportional"
		action = :store_true
		help = "Overloads criterion and uses pat@x% with incresing x. Prints only ranks."
	"--best-params"
		action = :store_true
		help = "Stores CSV files for each model's best parameters."
end

function main(args)
	f = datadir(args["filename"])
	df = load(f)[:df]
	@info "Loaded $(nrow(df)) rows from $f"

	if args["proportional"]
		ranks = []
		if args["rank-metric"] != ""
			for criterion in _prefix_symbol.("val", PAT_METRICS)
				df_agg = aggregate_stats_mean_max(df, criterion)
				push!(ranks, rank_table(df_agg, args["rank-metric"])[end:end, :])
			end
		else # pat/pat scenario if no rank-metric is provided
			for (criterion, metric) in zip(_prefix_symbol.("val", PAT_METRICS), _prefix_symbol.("tst", PAT_METRICS))
				df_agg = aggregate_stats_mean_max(df, criterion)
				push!(ranks, rank_table(df_agg, metric)[end:end, :])
			end
		end

		df_ranks = vcat(ranks...)
		df_ranks[:, :dataset] .= PAT_METRICS
		hl_best_rank = Highlighter(
					f = (data, i, j) -> (data[i,j] == minimum(df_ranks[i, 2:end])),
					crayon = crayon"green bold")

		rank_metric = (args["rank-metric"] != "") ? args["rank-metric"] : "tst_pat_x"

		@info "Best models chosen by validation pat@x% with incresing x"
		@info "Ranking by $(rank_metric)"

		if ~args["verbose"]
			target_filename = datadir(args["output-prefix"],"ranks_pat_$(rank_metric).txt")
			(~isdir(dirname(target_filename))) && mkdir(dirname(target_filename))
			open(target_filename, "w") do io
				pretty_table(
					io,
					df_ranks,
					formatters = ft_round(2),
					highlighters = (hl_best_rank),
				)
			end
		else
			pretty_table(
					df_ranks,
					formatters = ft_round(2),
					highlighters = (hl_best_rank),
				)
		end
	else
		df_agg = aggregate_stats_mean_max(df, Symbol(args["criterion-metric"]))
		rt = rank_table(df_agg, args["rank-metric"])

		@info "Best models chosen by $(args["criterion-metric"])"
		@info "Ranking by $(args["rank-metric"])"
		if ~args["verbose"]
			target_filename = datadir(args["output-prefix"],"$(args["criterion-metric"])_$(args["rank-metric"]).$(args["backend"])")
			(~isdir(dirname(target_filename))) && mkdir(dirname(target_filename))
			open(target_filename, "w") do io
				print_rank_table(io, rt; backend=Symbol(args["backend"]))
			end
		else
			print_rank_table(rt; backend=Symbol(args["backend"]))
		end

		if args["best-params"]
			target_dir = datadir(args["output-prefix"], "params_"*args["criterion-metric"])
			@info "Storing best parameters as separate CSVs in $(target_dir)"
			(~isdir(target_dir)) && mkdir(target_dir)
			all_models = unique(df_agg[:modelname])
			for m in all_models
				magg = filter(x -> (x.modelname == m), df_agg)
				CSV.write(joinpath(target_dir, "$(m).csv"), magg)
			end
		end
	end
end

main(parse_args(ARGS, s))


__precompile__()

module Agtor

using Parameters
using DataFrames, Dates, YAML, CSV
using Distributed

import OrderedCollections: LittleDict
import DataStructures: OrderedDict, DefaultDict

include("AgBase/Units.jl")
include("AgBase/Component.jl")
include("AgBase/Parameter.jl")
include("AgBase/io.jl")
include("AgBase/properties.jl")
include("AgBase/ResultsHandler.jl")
include("Infrastructure/Infrastructure.jl")
include("Infrastructure/Irrigation.jl")
include("Infrastructure/Pump.jl")
include("Climate.jl")
include("WaterSource.jl")
include("Crop.jl")
include("Field.jl")
include("Zone.jl")
include("Manager.jl")
include("RigidManager.jl")
include("Basin.jl")
include("Model.jl")

AgUnion = Union{Int64, Float64, Agtor.AgParameter}

const component_sep = "└──"
const component_name_sep = "-"
const component_field_sep = "~"

export @def, AgUnion

export load_yaml, load_spec, load_climate, create
export generate_agparams, create, reset!
export AgComponent, AgParameter, Infrastructure, Irrigation, Pump
export WaterSource, Crop, FarmField, CropField, FarmZone
export Manager, BaseManager, RigidManager
export Climate, Basin, run_timestep!, subtotal_costs, total_costs, update_available_water!
export collect_results, min_max, extract_values, param_info, extract_spec, add_prefix!
export set_params!, extract_agparams, collect_agparams!, collect_agparams, update_model!
export water_used_by_source
export get_data_for_timestep, aggregate_field_logs, run_model
export run_scenarios!, save_results!, save_state!
export collate_results, collate_results!, component_relationship
export select_results, scenario_stats, results_df

end # module


export
    Event,
    EventList,
    AllEventList,
    getevent,
    getlastevent,
    pushevent!,
    samplelocalpath,
    quadpathpoly,
    pathmean,
    tmax

"""
    AllowedVarType

Syntactic sugar to make it explicit what variable types are handled in the
local BPS (i.e, the type of the variables on each node of the factor graph).
See also, Event, EventList.
"""
const AllowedVarType = Union{Real, Vector{<:Real}}

"""
    Event{T <: AllowedVarType}

Encapsulation of a local event: a tuple made out of (position, velocity, time)
for a node in the local BPS algorithm.
See also: `EventList`.
"""
struct Event{T <: AllowedVarType}
    x::T
    v::T
    t::Real
end

"""
    EventList{T <: AllowedVarType}

Storage for all events that happened for a variable in the local BPS.
Each variable is stored explicitly (as opposed to storing a list of Event)
so that the times can be traversed efficiently. The xs, vs must all have the
same type within the list.
See also: `Event`, `AllEventList`.
"""
mutable struct EventList{T <: AllowedVarType}
    xs::Vector{T}     # coordinates
    vs::Vector{T}     # associated velocities
    ts::Vector{Real}  # associated times
end

function EventList(T::Type)
    @assert T <: AllowedVarType
    EventList(Vector{T}(undef, 0), Vector{T}(undef, 0), Vector{Real}(undef, 0))
end

"""
    AllEventList

Storage for all the `EventList` objects corresponding to a factor graph in a
local BPS setting. The element `evl` can be accessed by the variable index so
`a.evl[k]` will return the eventlist associated with variable `k`. Note that
the different `EventList` may have different types (in particular, different
dimensions). See also: `Event`, `EventList`.
"""
mutable struct AllEventList
    evl::Vector{EventList}
    types::Vector{DataType}
    function AllEventList(types::Vector{DataType})
        new([EventList(types[i]) for i ∈ 1:length(types)], types)
    end
end
AllEventList(T::DataType, nvars::Int) = AllEventList([T for i ∈ 1:nvars])

tmax(aev::AllEventList) = maximum(aev.evl[i].ts[end] for i ∈ 1:length(aev.evl))

"""
    getevent(evl, evidx)

Retrieve a specfic event (at index `evidx`) in an `EventList` object `evl` and
return it as an `Event` object.
See also: `Event`, `EventList`.
"""
function getevent(evl::EventList, evidx::Int)
    Event(evl.xs[evidx], evl.vs[evidx], evl.ts[evidx])
end

"""
    getlastevent(evl)

Retrieve the last `Event` in an `EventList` object `evl`.
See also: `getevent`, `EventList`.
"""
function getlastevent(evl::EventList)
    Event(evl.xs[end], evl.vs[end], evl.ts[end])
end

"""
    pushevent!(evl, ev)

Append an `Event` object `ev` to an `EventList` object `evl`. The `Event`
object is unpacked and pushed to the vectors encapsulated by `evl`.
See also: `EventList`.
"""
function pushevent!(evl::EventList, ev::Event)
    push!(evl.xs, ev.x)
    push!(evl.vs, ev.v)
    push!(evl.ts, ev.t)
    evl
end

"""
    getlocalsegment(evl, j)

Retrieve the `j`th event in an `EventList` object `evl` and return as a tuple
that event and the next event. If there is no next event (end of event list),
an "infinity event" is returned (in order to allow extrapolation).
See also: `samplelocalpath`.
"""
function getlocalsegment(evl::EventList, j::Int)
    if j == length(evl.ts)
        return (getevent(evl,j), Event(NaN, NaN, Inf))
    end
    (getevent(evl,j), getevent(evl,j+1))
end

"""
    samplelocalpath(evl, t)

Sample x_k(t) given the corresponding `EventList` object `evl`. The times `t`
must be a list of sorted times (or a single positive time).
See also: `getlocalsegment`.
"""
function samplelocalpath(evl::EventList, t::AllowedTimeType)
    @assert issorted(t) "list of times should be sorted"
    ti  = t[1]
    @assert ti >= 0.0 "the first time should be greater than 0.0"
    # find the starting local segment
    j = searchsortedlast(evl.ts, ti)
    # event and next event (describes the "segment")
    eva, evb = getlocalsegment(evl,j)
    # allocation of samples
    samples = Vector{typeof(eva.x)}(undef, length(t))
    # recursively work with local segments
    i=1
    while i<= length(t)
        ti = t[i]
        if ti <= evb.t
            # stay in the current segment
            samples[i] = eva.x + (ti - eva.t) * eva.v
            i += 1
        else
            # get the next segment
            j += searchsortedlast(evl.ts[j+1:end], ti)
            eva, evb = getlocalsegment(evl, j)
        end
    end
    samples
end
samplelocalpath(evl::EventList, t::Real) = samplelocalpath(evl, [t])[1]

"""
    quadpathpoly(evl, pol, T)

Integrate a polynomial on the elements along a path for a single node.
So if we are sampling in Rd and have phi(X)=Poly(xk), then we can integrate the
polynomial along the path. Note that we can't do the same thing if phi(X) mixes
2 or more nodes. See also syntactic sugar to perform the same operation on all
nodes.
"""
function quadpathpoly(evl::EventList, pol::Poly, T::Real)
    dim  = length(evl.xs[1])
    res  = zeros(dim)
    nseg = searchsortedlast(evl.ts,T)
    for segidx = 1:nseg-1
        # for each segment in the path
        eva, evb = getlocalsegment(evl, segidx)
        # segment time, initial point and velocity
        tau = evb.t - eva.t
        xa  = eva.x
        v   = (evb.x - xa) / tau
        # integrate along that segment
        for d = 1:dim
            res[d] += polyint(polyval(pol,Poly([xa[d],v[d]])))(tau)
        end
    end
    # last segment
    eva = getevent(evl, nseg)
    # characteristics
    tau = T - eva.t
    for d = 1:dim
        res[d] += polyint(polyval(pol,Poly([eva.x[d],eva.v[d]])))(tau)
    end
    return res / T
end
function quadpathpoly(aev::AllEventList, pol::Poly, T::Real)
    res = Vector{Vector{<:Real}}(undef, length(aev.evl))
    for (d, evl) in enumerate(aev.evl)
        res[d] = quadpathpoly(evl, pol, T)
    end
    return res
 end
pathmean(evl::EventList, T::Real) = quadpathpoly(evl, Poly([0.0,1.0]), T)
pathmean(aev::AllEventList, T::Real) = quadpathpoly(aev, Poly([0.0,1.0]), T)
pathmean(aev::AllEventList) = pathmean(aev, tmax(aev))


module LifeValData

using XLSX, DataFrames, JuliaDB, DelimitedFiles

using LifeValStructures

export Initialise, Read_Data

# set some global constants that need to be worked into functions at some point
function Initialise(filePath, fileName)
    main = DataFrame(XLSX.readtable(fileName, "MAIN")...)
    runsettings = DataFrame(XLSX.readtable(fileName, main.RS[1])...)
    spcode = DataFrame(XLSX.readtable(fileName, main.SP[1])...)

        #writes these as lastrun SETTINGS
    XLSX.writetable(filePath * "lastrun.xlsx", MAIN =( DataFrames._columns(main), DataFrames._names(main) ),
        TEMP1 =( DataFrames._columns(runsettings), DataFrames._names(runsettings) ),
        TEMP2 =( DataFrames._columns(spcode), DataFrames._names(spcode) ), overwrite=true)

    XLSX.openxlsx(filePath * "lastrun.xlsx", mode="rw") do xf
        sheet = xf[2]
        XLSX.rename!(sheet, main.RS[1])
        sheet = xf[3]
        XLSX.rename!(sheet, main.SP[1])
    end

    #read bases - this should happen before reading policies for e.g. impact of unit pricing on policies
    qxname = Vector{String}()
    wxname = Vector{String}()
    vtname = Vector{String}()
    expname = Vector{String}()
    chname = Vector{String}()
    fname = Vector{String}()
    fbname = Vector{String}()

    #qxLookup = unique(spcode, :qxbasis)
    for runtype in eachrow(runsettings)
        qxn = filePath * runtype.qxname
        wxn = filePath * runtype.wxname
        vtn = filePath * runtype.vtname
        expn = filePath * runtype.expname
        chn = filePath * runtype.chname
        fn = filePath * runtype.fname
        fbn = runtype.fbasis

        push!(qxname, qxn)
        push!(wxname, wxn)
        push!(vtname, vtn)
        push!(expname, expn)
        push!(chname, chn)
        push!(fname, fn)
        push!(fbname, fbn)
    end

    names = DataFrame([qxname, wxname, vtname, expname, chname, fname, fbname])
    rename!(names, [:qxname,:wxname,:vtname,:expname,:chname,:fname,:fbname])

    outpn = main.outv[1]

    return [runsettings, spcode, names, outpn]
end

function Read_Data(runsettings, spcode, names, filePath)
    #read policyData - []
    policyData = Vector{IndexedTable}(undef, nrow(runsettings))
    fB = DataFrame(XLSX.readtable(names.fname[1], names.fbname[1])...) #temporary, assume all f has same #s
    fData = Array{LifeValStructures.fTables}(undef, nrow(runsettings), nrow(fB))
    for runtype in eachrow(runsettings)
        policyData[runtype.runno] = load(filePath * runtype.policyfile)

        #fund data
        if runtype.fbasis != ""
            fB = DataFrame(XLSX.readtable(names.fname[runtype.runno], runtype.fbasis)...)
            for fbItem in eachrow(fB)
                fData[runtype.runno, fbItem.ID]=LifeValStructures.fTables(runtype.runno, fbItem.ID, DataFrame(fbItem))
            end
        end
    end

    #find all (unique) spcodes
    spcodes = []
    for runno in eachrow(spcode)
        spcod = runno.spc
        #if spcod !in spcodes
            push!(spcodes, spcod)
        #end
    end
    spcodes = unique(spcodes)

    qxData = Array{LifeValStructures.qxTables}(undef, nrow(runsettings), 100)
    qxpData = Array{LifeValStructures.qxTables}(undef, nrow(runsettings), 100)
    wxData = Array{LifeValStructures.wxTables}(undef, nrow(runsettings), 100)
    wxpData = Array{LifeValStructures.wxTables}(undef, nrow(runsettings), 100)
    vtData = Array{LifeValStructures.vtTables}(undef, nrow(runsettings), 100)
    vtpData = Array{LifeValStructures.vtTables}(undef, nrow(runsettings), 100)
    expData = Array{LifeValStructures.expTables}(undef, nrow(runsettings), 100)
    exppData = Array{LifeValStructures.expTables}(undef, nrow(runsettings), 100)
    chData = Array{LifeValStructures.chTables}(undef, nrow(runsettings), 100)

    for runno in eachrow(spcode)
        runtemp = runno.runno
        spcod = runno.spc

        qxB = DataFrame(XLSX.readtable(names.qxname[runtemp], runno.qxbasis)...)
        qxData[runtemp,spcod]=LifeValStructures.qxTables(runtemp,spcod,qxB)

        if runno.qxbasisp != ""
            qxpB = DataFrame(XLSX.readtable(names.qxname[runtemp], runno.qxbasisp)...)
            qxpData[runtemp,spcod]=LifeValStructures.qxTables(runtemp,spcod,qxpB)
        end

        wxB = DataFrame(XLSX.readtable(names.wxname[runtemp], runno.wxbasis)...)
        wxData[runtemp,spcod]=LifeValStructures.wxTables(runtemp,spcod,wxB)

        if runno.wxbasisp != ""
            wxpB = DataFrame(XLSX.readtable(names.wxname[runtemp], runno.wxbasisp)...)
            wxpData[runtemp,spcod]=LifeValStructures.wxTables(runtemp,spcod,wxpB)
        end

        vtB = DataFrame(XLSX.readtable(names.vtname[runtemp], runno.vtbasis)...)
        vtData[runtemp,spcod]=LifeValStructures.vtTables(runtemp,spcod,vtB)

        if runno.vtbasisp != ""
            vtpB = DataFrame(XLSX.readtable(names.vtname[runtemp], runno.vtbasisp)...)
            vtpData[runtemp,spcod]=LifeValStructures.vtTables(runtemp,spcod,vtpB)
        end

        expB = DataFrame(XLSX.readtable(names.expname[runtemp], runno.ebasis)...)
        expItem = filter(row -> row.SPCODE == runno.spc, expB) #don't need every spcode on spcode level
        expData[runtemp,spcod]=LifeValStructures.expTables(runtemp,spcod,expItem)

        if runno.ebasisp != ""
            exppB = DataFrame(XLSX.readtable(names.expname[runtemp], runno.ebasisp)...)
            exppItem = filter(row -> row.SPCODE == runno.spc, exppB) #don't need every spcode on spcode level
            exppData[runtemp,spcod]=LifeValStructures.expTables(runtemp,spcod,exppItem)
        end

        if runno.chbasis != ""
            chB = DataFrame(XLSX.readtable(names.chname[runtemp], runno.chbasis)...)
            chData[runtemp,spcod]=LifeValStructures.chTables(runtemp,spcod,chB)
        end
    end

    basisData = [qxData, qxpData, wxData, wxpData, vtData, vtpData, expData, exppData, chData]
    #names!(basisData, [:qxData,:qxpData,:wxData,:wxpData,:vtData,:vtpData,:expData,:exppData,:chData,:fData])

    return [policyData, basisData, fData]
end

end


# Query.jl is an embarrassment when it comes to group by Performance

# using StatsBase, RCall, Query, DataFrames, DataFramesMeta, Distributions, PooledArrays, IndexedTables
#
# function query_benches(df::DataFrame)
#
#     # timings
#     ti = Dict()
#
#     ti[:sum1] = @elapsed @from i in df begin
#                              @group i by i.id1 into g
#                              @select {r=sum(g..v1)}
#                              @collect DataFrame
#                          end
#     ti[:sum2] = @elapsed @from i in df begin
#                              @group i by i.id1 into g
#                              @select {r=sum(g..v1)}
#                              @collect DataFrame
#                          end
#     ti[:sum3] = @elapsed @from i in df begin
#                              @group i by (i.id1,i.id2) into g
#                              @select {r=sum(g..v1)}
#                              @collect DataFrame
#                          end
#     ti[:sum4] = @elapsed @from i in df begin
#                              @group i by (i.id1,i.id2) into g
#                              @select {r=sum(g..v1)}
#                              @collect DataFrame
#                          end
#     ti[:sum_mean1] = @elapsed @from i in df begin
#                              @group i by i.id3 into g
#                              @select {s=sum(g..v1),m=mean(g..v3)}
#                              @collect DataFrame
#                          end
#     ti[:sum_mean2] = @elapsed @from i in df begin
#                              @group i by i.id3 into g
#                              @select {s=sum(g..v1),m=mean(g..v3)}
#                              @collect DataFrame
#                          end
#     ti[:mean7_9_by_id4_1] = @elapsed @from i in df begin
#                              @group i by i.id4 into g
#                              @select {m7=mean(g..v1),m8=mean(g..v2),m9=mean(g..v3)}
#                              @collect DataFrame
#                          end
#     ti[:mean7_9_by_id4_2] = @elapsed @from i in df begin
#                              @group i by i.id4 into g
#                              @select {m7=mean(g..v1),m8=mean(g..v2),m9=mean(g..v3)}
#                              @collect DataFrame
#                          end
#     ti[:sum7_9_by_id6_1] = @elapsed @from i in df begin
#                              @group i by i.id6 into g
#                              @select {m7=mean(g..v1),m8=mean(g..v2),m9=mean(g..v3)}
#                              @collect DataFrame
#                          end
#     ti[:sum7_9_by_id6_2] = @elapsed @from i in df begin
#                              @group i by i.id6 into g
#                              @select {m7=mean(g..v1),m8=mean(g..v2),m9=mean(g..v3)}
#                              @collect DataFrame
#                          end
#     return ti
# end
#
# function query_benches(df::IndexedTable)
#
#     # timings
#     ti = Dict()
#
#     ti[:sum1] = @elapsed @from i in df begin
#                              @group i by i.id1 into g
#                              @select {r=sum(g..v1)}
#                              @collect DataFrame
#                          end
#     ti[:sum2] = @elapsed @from i in df begin
#                              @group i by i.id1 into g
#                              @select {r=sum(g..v1)}
#                              @collect DataFrame
#                          end
#     ti[:sum3] = @elapsed @from i in df begin
#                              @group i by (i.id1,i.id2) into g
#                              @select {r=sum(g..v1)}
#                              @collect DataFrame
#                          end
#     ti[:sum4] = @elapsed @from i in df begin
#                              @group i by (i.id1,i.id2) into g
#                              @select {r=sum(g..v1)}
#                              @collect DataFrame
#                          end
#     ti[:sum_mean1] = @elapsed @from i in df begin
#                              @group i by i.id3 into g
#                              @select {s=sum(g..v1),m=mean(g..v3)}
#                              @collect DataFrame
#                          end
#     ti[:sum_mean2] = @elapsed @from i in df begin
#                              @group i by i.id3 into g
#                              @select {s=sum(g..v1),m=mean(g..v3)}
#                              @collect DataFrame
#                          end
#     ti[:mean7_9_by_id4_1] = @elapsed @from i in df begin
#                              @group i by i.id4 into g
#                              @select {m7=mean(g..v1),m8=mean(g..v2),m9=mean(g..v3)}
#                              @collect DataFrame
#                          end
#     ti[:mean7_9_by_id4_2] = @elapsed @from i in df begin
#                              @group i by i.id4 into g
#                              @select {m7=mean(g..v1),m8=mean(g..v2),m9=mean(g..v3)}
#                              @collect DataFrame
#                          end
#     ti[:sum7_9_by_id6_1] = @elapsed @from i in df begin
#                              @group i by i.id6 into g
#                              @select {m7=mean(g..v1),m8=mean(g..v2),m9=mean(g..v3)}
#                              @collect DataFrame
#                          end
#     ti[:sum7_9_by_id6_2] = @elapsed @from i in df begin
#                              @group i by i.id6 into g
#                              @select {m7=mean(g..v1),m8=mean(g..v2),m9=mean(g..v3)}
#                              @collect DataFrame
#                          end
#     return ti
# end
#
#
# function run_query_benches(N=1_000_000;K=100)
#     # get small data for JIT warmup
#     d_ = createSynDataFrame(10,3)
#     # dt_ = createDataTable(10,3)
#
#
#     # warm up julia benchmarks
#     query_benches(d_);
#     # benches(dt_);
#     benches(di_);
#     DfMeta_benches(d_);
#
#     # get real data
#     d = createDataFrame(N,K)
#     # measure DataFrames
#     query_df = benches(d)
#     meta = DfMeta_benches(d)
#     d = 0
#     gc()
#     # measure DataTables
#     # dt = createDataTable(N,K)
#     # query_dt = benches(dt)
#     # dt = 0
#     # gc()
#     # measure IndexedTables
#     di = createIndexedTable(N,K)
#     query_di = benches(di)
#     di = 0
#     gc()
#
#
#     # get R time
#     R = R_bench(N,K)
#
#     # get
#     k = collect(keys(query_df))
#     #out = DataFrame(bench = k,Query_DF = [query_df[kk] for kk in k],Query_DT= [query_dt[kk] for kk in k],Query_idxT= [query_di[kk] for kk in k],DataFramesMeta=[meta[kk] for kk in k], Rdatatable=[R[kk] for kk in k])
#     out = DataFrame(bench = k,Query_DF = [query_df[kk] for kk in k] ,Query_idxT= [query_di[kk] for kk in k],DataFramesMeta=[meta[kk] for kk in k], Rdatatable=[R[kk] for kk in k])
#     sort!(out,cols=:bench)
#     rel = deepcopy(out)
#     #rel = @transform(rel,Query_DF = :Query_DF./:Rdatatable, Query_DT = :Query_DT./:Rdatatable, Query_idxT = :Query_idxT./:Rdatatable, DataFramesMeta = :DataFramesMeta ./ :Rdatatable, Rdatatable = 1.0)
#     rel = @transform(rel,Query_DF = :Query_DF./:Rdatatable, Query_idxT = :Query_idxT./:Rdatatable, DataFramesMeta = :DataFramesMeta ./ :Rdatatable, Rdatatable = 1.0)
#     return (out,rel)
# end
#
# function run_all()
#     d=Dict()
#     for n in [1_000_000, 2_000_000]
#         d[n] = run_benches(n)
#         gc()
#         info("results with $n rows:")
#         info("-- absolute time in seconds:")
#         println(d[n][1])
#         info("-- time relative to R data.table:")
#         println(d[n][2])
#         println()
#         println()
#     end
#     return d
# end


println("Starting")

using MetaGen
using JSON
import YAML
using Pipe: @pipe
using Random

config_path = ARGS[1]
config = YAML.load_file(config_path)
output_dir = "results/$(config["experiment_name"])_" * ENV["SLURM_JOB_ID"]
mkdir(output_dir)

include("useful_functions.jl")
dict = []
#for i = 0:config["batches_upto"]
	#to_add =  @pipe "$(config["input_file_dir"])$(i)_data_labelled.json" |> open |> read |> String |> JSON.parse
to_add =  @pipe "$(config["input_file_dir"])data_labelled.json" |> open |> read |> String |> JSON.parse
append!(dict, to_add)
#end
#dict = @pipe "../../scratch_work_07_16_21/0_data_labelled.json" |> open |> read |> String |> JSON.parse
#dict = @pipe "../../scratch_work_07_16_21/0_data_labelled.json" |> open |> read |> String |> JSON.parse

#Random.seed!(15)
#try to make objects_observed::Array{Array{Array{Array{Detection2D}}}} of observed objects.
#outer array is for scenes, then frames, the receptive fields, then last is an array of detections

################################################################################
num_videos = config["num_videos"]
num_frames = config["num_frames"]
threshold = config["threshold"]

params = Video_Params(n_possible_objects = 5)
top_n = 5

receptive_fields = make_receptive_fields(params)
objects_observed, camera_trajectories = make_observations_office(dict, receptive_fields, num_videos, num_frames, threshold, top_n)

num_particles = config["num_particles"]
mcmc_steps_outer = config["mcmc_steps_outer"]
mcmc_steps_inner = config["mcmc_steps_inner"]

################################################################################
#=
#Set up the output file
online_outfile = output_dir * "/online_output.csv"
online_file = open(online_outfile, "w")
file_header(online_file)

################################################################################
#Online MetaGen
println("start pf for online")

#@profilehtml unfold_particle_filter(false, num_particles, objects_observed, camera_trajectories, params, file)
traces, inferred_world_states, avg_v = unfold_particle_filter(nothing,
	num_particles, mcmc_steps_outer, mcmc_steps_inner, objects_observed,
	camera_trajectories, params, online_file)
close(online_file)

println("done with pf for online")

## Commented out for testing

################################################################################
#Retrospective MetaGen
println("start retrospective")

#Set up the output file
retro_outfile = output_dir * "/retrospective_output.csv"
retro_file = open(retro_outfile, "w")
file_header(retro_file)

unfold_particle_filter(avg_v, num_particles, mcmc_steps_outer, mcmc_steps_inner,
	objects_observed, camera_trajectories, params, retro_file)
close(retro_file)

println("done with pf for retrospective")

################################################################################
#run Lesioned MetaGen
=#

shuffle_type = config["shuffle_type"]

num_videos_train = convert(Int64, num_videos/2)

if shuffle_type==0
	order = collect(1:num_videos_train)
elseif shuffle_type==1
	Random.seed!(1)
	order = shuffle(1:num_videos_train)
else shuffle_type==2
	Random.seed!(2)
	order = shuffle(1:num_videos_train)
end

training_objects_observed = objects_observed[order, :]
training_camera_trajectories = camera_trajectories[order, :]
#training set and test set
input_objects_observed = vcat(objects_observed[order, :], objects_observed[(num_videos_train+1):num_videos, :])
input_camera_trajectories = vcat(camera_trajectories[order, :], camera_trajectories[(num_videos_train+1):num_videos, :])

#Set up the output file
lesioned_V_file = open(output_dir * "/lesioned_V.csv", "w")
file_header_V(lesioned_V_file, params)
lesioned_ws_file = open(output_dir * "/lesioned_ws.csv", "w")
file_header_ws(lesioned_ws_file, params, num_particles)

v = zeros(length(params.possible_objects), 2)
v[:,1] .= 1.0
v[:,2] .= 0.5
unfold_particle_filter(v, num_particles, mcmc_steps_outer, mcmc_steps_inner,
	input_objects_observed, input_camera_trajectories, params, lesioned_V_file,
	lesioned_ws_file)
close(lesioned_V_file)
close(lesioned_ws_file)

println("done with pf for lesioned metagen")

#=

################################################################################
#for writing an output file for a demo using MetaGen

###### add to dictionary
out = write_to_dict(dict, camera_trajectories, inferred_world_states, num_videos, num_frames)

#open("../../scratch_work_07_16_21/output_tiny_set_detections.json","w") do f
open(output_dir * "/output.json","w") do f
	JSON.print(f,out)
end

println("finished writing json")
=#


function LUDecomposition(arg0::RealMatrix)
    return LUDecomposition((RealMatrix,), arg0)
end

function LUDecomposition(arg0::RealMatrix, arg1::jdouble)
    return LUDecomposition((RealMatrix, jdouble), arg0, arg1)
end

function get_determinant(obj::LUDecomposition)
    return jcall(obj, "getDeterminant", jdouble, ())
end

function get_l(obj::LUDecomposition)
    return jcall(obj, "getL", RealMatrix, ())
end

function get_p(obj::LUDecomposition)
    return jcall(obj, "getP", RealMatrix, ())
end

function get_pivot(obj::LUDecomposition)
    return jcall(obj, "getPivot", Vector{jint}, ())
end

function get_solver(obj::LUDecomposition)
    return jcall(obj, "getSolver", DecompositionSolver, ())
end

function get_u(obj::LUDecomposition)
    return jcall(obj, "getU", RealMatrix, ())
end


import MathOptInterface
import MathOptInterface: Utilities

const MOI = MathOptInterface
const MOIU = MOI.Utilities

# indices
const VI = MOI.VariableIndex
const CI = MOI.ConstraintIndex

# function aliases
const SAF = MOI.ScalarAffineFunction{Float64}
const SQF = MOI.ScalarQuadraticFunction{Float64}

# set aliases
const Bounds{T} = Union{
    MOI.EqualTo{T},
    MOI.GreaterThan{T},
    MOI.LessThan{T},
    MOI.Interval{T}
}

mutable struct Optimizer <: MOI.AbstractOptimizer
    inner::BaronModel
    nlp_block_data::Union{Nothing, MOI.NLPBlockData}
end

Optimizer(; options...) = Optimizer(BaronModel(; options...), nothing)

MOI.get(model::Optimizer, ::MOI.SolverName) = "BARON"

# empty
function MOI.is_empty(model::Optimizer)
    BARON.is_empty(model.inner) && model.nlp_block_data === nothing
end

function MOI.empty!(model::Optimizer)
    model.inner = BaronModel(; ((Symbol(key), val) for (key, val) in model.inner.options)...)
    model.nlp_block_data = nothing
    return
end

# copy
# MOIU.supports_default_copy_to(::Optimizer, copy_names::Bool) = !copy_names
MOI.supports_incremental_interface(::Optimizer) = true
function MOI.copy_to(dest::Optimizer, src::MOI.ModelLike; kws...)
    return MOIU.default_copy_to(dest, src; kws...)
end

# optimize
function MOI.optimize!(model::Optimizer)
    if !IS_SOLVER_SET
        error(("""BARON.jl was not built correctly.
                 Set the environment variable `BARON_EXEC` and run `using Pkg; Pkg.build("BARON")`."""))
    end
    write_bar_file(model.inner)
    if model.inner.print_input_file
        println("\nBARON input file: $(model.inner.problem_file_name)\n")
        println(read(model.inner.problem_file_name, String))
    end
    try
        run(`$baron_exec $(model.inner.problem_file_name)`)
    catch e
        println("$e")
        println(read(model.inner.problem_file_name, String))
        error("failed to call BARON exec $baron_exec")
    end
    read_results(model.inner)
end

# RawOptimizerAttribute
MOI.supports(::Optimizer, ::MOI.RawOptimizerAttribute) = true
function MOI.set(model::Optimizer, param::MOI.RawOptimizerAttribute, value)
    model.inner.options[param.name] = value
    return
end
function MOI.get(model::Optimizer, param::MOI.RawOptimizerAttribute)
    return get(model.inner.options, param.name) do
        throw(ErrorException("Requested parameter $(param.name) is not set."))
    end
end

# TimeLimitSec
MOI.supports(::Optimizer, ::MOI.TimeLimitSec) = true
function MOI.set(model::Optimizer, ::MOI.TimeLimitSec, val::Real)
    model.inner.options["MaxTime"] = Float64(val)
    return
end

# BARON's default time limit is 1000 seconds.
function MOI.get(model::Optimizer, ::MOI.TimeLimitSec)
    return get(model.inner.options, "MaxTime", 1000.0)
end

struct PrintInputFile <: MOI.AbstractOptimizerAttribute end
function MOI.set(model::Optimizer, ::PrintInputFile, val::Bool)
    model.inner.print_input_file = val
    return
end
function MOI.get(model::Optimizer, ::PrintInputFile)
    model.inner.print_input_file
end

include(joinpath("moi", "util.jl"))
include(joinpath("moi", "variables.jl"))
include(joinpath("moi", "constraints.jl"))
include(joinpath("moi", "objective.jl"))
include(joinpath("moi", "nlp.jl"))
include(joinpath("moi", "results.jl"))


for (functor, impl) in ((ElementWiseFunctors.Add, :(CuVec.add!)),
                        (ElementWiseFunctors.Subtract, :(CuVec.sub!)),
                        (ElementWiseFunctors.Multiply, :(CuVec.mul!)),
                        (ElementWiseFunctors.Divide, :(CuVec.div!)))
  @eval begin
    function forward(backend::GPUBackend, state::ElementWiseLayerState{$functor},
        inputs::Vector{Blob})

      input1 = inputs[1]
      input2 = inputs[2]
      output = state.blobs[1]
      data_type = eltype(input1)
      len = length(input1)

      CuBLAS.copy(backend.cublas_ctx, data_type, len, input1.ptr.p, 1, output.ptr.p, 1)
      $impl(backend, output, input2)
    end
  end
end

function backward(backend::GPUBackend, state::ElementWiseLayerState{ElementWiseFunctors.Subtract},
    inputs::Vector{Blob}, diffs::Vector{Blob})

  if !isa(diffs[1], NullBlob)
    copy!(diffs[1], state.blobs_diff[1])
  end
  if !isa(diffs[2], NullBlob)
    copy!(diffs[2], state.blobs_diff[1])
    CuBLAS.scal(backend.cublas_ctx, length(diffs[2]), convert(eltype(diffs[2]),-1),
        diffs[2].ptr, 1)
  end
end
function backward(backend::GPUBackend, state::ElementWiseLayerState{ElementWiseFunctors.Multiply},
    inputs::Vector{Blob}, diffs::Vector{Blob})

  for i = 1:length(diffs)
    if !isa(diffs[i], NullBlob)
      copy!(diffs[i], state.blobs_diff[1])
      CuVec.mul!(backend, diffs[i], inputs[i%2 + 1])
    end
  end
end
function backward(backend::GPUBackend, state::ElementWiseLayerState{ElementWiseFunctors.Divide},
    inputs::Vector{Blob}, diffs::Vector{Blob})

  data_type = eltype(inputs[1])
  len = length(inputs[1])

  if !isa(diffs[1], NullBlob)
    copy!(diffs[1], state.blobs_diff[1])
    CuVec.div!(backend, diffs[1], inputs[2])
  end
  if !isa(diffs[2], NullBlob)
    copy!(diffs[2], state.blobs_diff[1])
    CuVec.mul!(backend, diffs[2], state.blobs[1])
    CuVec.div!(backend, diffs[2], inputs[2])
    CuBLAS.scal(backend.cublas_ctx, len, convert(data_type,-1), diffs[2].ptr, 1)
  end
end


import Pkg
Pkg.rm(["Grids", "GridGames"])
Pkg.develop(["Grids", "GridGames"])

using CUDA

# Make sure no errors or
# warnings with CUDA
CUDA.versioninfo()

using Base.Threads: @spawn
using MuZero
using GridGames
cd(@__DIR__)

# You need at least 5
# threads.
Threads.nthreads()

# User defined state and game implementations
include("game.jl")

# Play a few sample games to make sure the
# game dynamics are correct. Need to copy and
# paste "play" function into REPL, if using VSCode.
#
# play(Tac)

# Initialize MuZero configuration
include("muzero_config.jl")
config = getconfig(Tac)
config.num_unroll_steps = 5
config.num_actors = 1

# Initialize the policy tests
include("inference_tests.jl")
rprv_tests = makerrvptests(Tac)
rprv_tests.period = 1000
pv_tests = makepolicyvaluetests(Tac)
pv_tests.period = 1000
#valuefn_debug = getvaluefn(Tac)

# Configure Learner
config.num_actors
config.num_unroll_steps
learner = Learner(
    config;
    uuid="b234fc40-a07b-4418-9fbc-adb0278dd4e1",
    dir="~/models/MuZeroModels.jl/tictactoe/",
    tests=[pv_tests, rprv_tests], #
    schedule=tacschedule!,#cycleschedule(min_lr=1e-6, max_lr=9e-3),#
    #valuefn_debug=valuefn_debug
)

xyz = randrrvptest(learner)
xyz.state |> display
xyz.state_next |> display

getpolicyvaluetests(learner)[300] |> svg

l = Learner(Tac)
Learner(Tac; uuid="20B4993D-D1B8-41B7-ABB2-9BC8C70EA5F9")
Learner(Tac; dir="~/models/MuZeroModels.jl/tictactoe/")
Learner(
    Tac;
    uuid="8bcecce1-c870-4e34-8ae9-df6433d033de",
    dir="./models/MuZeroModels.jl/tictactoe/"
)

Learner(config)
Learner(config; uuid="20B4993D-D1B8-41B7-ABB2-9BC8C70EA5F9")
Learner(config; dir="~/models/MuZeroModels.jl/tictactoe/")
Learner(
    config;
    uuid="8bcecce1-c870-4e34-8ae9-df6433d033de",
    dir="./models/MuZeroModels.jl/tictactoe/"
)

Learner(Tac, config)
Learner(Tac, config; uuid="20B4993D-D1B8-41B7-ABB2-9BC8C70EA5F9")
Learner(Tac, config; dir="~/models/MuZeroModels.jl/tictactoe/")
Learner(
    Tac,
    config;
    uuid="8bcecce1-c870-4e34-8ae9-df6433d033de",
    dir="./models/MuZeroModels.jl/tictactoe/"
)

getpvtestresult(learner, 300) |> svg
acc, fails, value_mse = runpvtests(learner)
reward_mse, value_mse, acc, fails = runrrvptests(learner)
getpolicy(learner, getstate(Tac([2,8,6]))) |> display
getvalue(learner, getstate(Tac([2,8,6])), 1) |> display

# Let's make sure our policy test cases look reasonable
randpolicyvaluetest(learner) |> svg
randpolicytest(learner) |> svg
randpolicytest(learner) |> svg
randpolicytest(learner) |> svg

showrrvptestresults(learner)

# The predictions from this untrained model should be random
# and mostly incorrect.
showpolicyvaluetestresults(learner)
randpolicyvaluetestresult(learner) |> svg

# Self-play 10 sample games to ensure proper game implementation.
selfplay!(learner; limit=10)

# Show a random self-played game
g = randgame(learner)
g |> display
g |> svg
randtrainingsample(learner, g) |> display

# Get a random batch of training samples
randtrainingbatch(learner) |> display
randtrainingbatch(learner)[1] |> svg

# If game play seems to be working properly, then start indefinite self play on multiple threads.
r1 = @spawn play!(learner)

# Let's make sure threading is working propery and the games are being recorded
length(learner.replay_buffers)
getavggamestep(learner)
getgamestepquality(learner)
get_steps(learner)
r2 = @spawn fit!(learner, steps=1100)

learner.is_playing = true
res3 = @spawn selfplay!(learner)
learner.is_playing = false

updatecallback!(
    learner,
    LearnerCallback(
        "SaveGamePlayStats",
        savegameplaystats!,
        period=1000,
        event=:afterstep
    )
)

learner.is_training = false
updateschedule!(learner, tacschedule!)
updateschedule!(learner, newcycleschedule(learner, min_lr=1e-5, max_lr=3e-5, cycle_steps=10_000))

# Train three steps to make sure there are no errors. This 1st step
# may take a several minutes if your machine has no usable
# Nvidia GPU.
@show @time fit!(learner, steps=3)

# Let's ensure the initial loss was stored, which means the model is training
showlosses(learner)

# View training progress
showresults(learner)

# Train network for a few more steps
fit!(learner, steps=20)

# View training progress. We should see steady decrease
# in the loss and a steady increase in the accuracy.
showresults(learner)

# Let's train indefinitely on a separate thread and periodically check the
# progress.
res = @spawn fit!(learner; steps=1000000)

fit!(learner; steps=3)

# function losses_svg()
#     function Base.show(io::IO, m::MIME"image/svg+xml", g::GridSVG)
#         draw(g)
#         print(io, read(output_file(g), String))
#     end
# end

# # Visualize any game from batch using the encoded state key
# df = get_df(batch)
# state = TicTacToeState(df."State Key"[1])
# state |> show

# # Or, use the following convenience function to get the state
# # @todo Stopped here!
# get_state(batch, 1) |> show

# Learning rate finder
# @todo implement

# Open up "losses.png" and "accuracy.png" in separate windows
# to view training loss and testing accuracy graphs in real-time.

# Alternatively, print the "losses" and/or "accuracy" arrays
# in the REPL to view progress.
losses
accuracy

# If losses is empty and/or not increasing, see status.txt
# to check if there was an error during training.
`cat ./status.txt` |> run

# Also, you can check "fetch(result)" to see if there was an exception.
# If no exception prints (and the call to fetch hangs), then the training is still
# starting, and you can interrupt the call to fetch
fetch(result)

# Export Flux model
MuZero.export(learner)

@show @time selfplay!(learner; limit=3)
@show @time selfplay!(learner; limit=3)

using BenchmarkTools
@btime selfplay!(learner; limit=3)

@btime selfplay!(learner; limit=10)

@spawn selfplay!(learner)
@spawn selfplay!(learner)
@spawn selfplay!(learner)

length(learner.replay_buffer)

@show @time fit!(learner; steps=3)
@spawn fit!(learner)

learner.stats.losses |> length
learner |> gettraindf |> display
learner |> gettraindf |> mean1000 |> display
learner |> gettraindfall |> mean1000 |> display
stoptraining!(learner)
stopplaying!(learner)
clearreplaybuffer!(learner)
length(learner.replay_buffer)

learner |> show_pv_tests
learner |> show_reward_tests

#@learnertest 4854 |> svg
#@learnersample 33

selfplay!(learner; limit=10)

learner |> show_game

# still a fence post problem here
s = learner |> randtrainingsample
learner |> randtrainingsample |> svg

learner |> randtrainingbatch

using Base.Threads: @spawn
for i = 1:num_actors(learner)
    @spawn selfplay!(learner)
end

length(learner.replay_buffer)

@show get_steps(learner), learner.stats.losses, learner.stats.accuracies, learner.times

@time fit!(learner, steps=1)

learner |> gettraindfall |> display

@show learner.stats.losses, learner.stats.accuracies, learner.times

get_steps(learner), get_accuracy(learner)

# STOPPED HERE

tor = "/home/ubuntu/models/MuZeroModels.jl/tictactoe/eb14fc40-a07b-4418-9fbc-adb0278dd4e1/replay_buffers/e4ffabbe-1638-4580-8cc6-b57c8e7c6dbd/buffer.bin"
tor2 = "/home/ubuntu/models/MuZeroModels.jl/tictactoe/eb14fc40-a07b-4418-9fbc-adb0278dd4e1/replay_buffers/d5281c62-455a-4bd4-a145-0ff86597d058/buffer.bin"
reanalyzebuffer(learner, tor2)

# DEBUGGING PV FAILS
getpolicyvaluetestresults(learner)[1935:1935]
fit!(learner, steps=100) ; learner.stats.policy_test_fails

pvdf = getdf(getpolicyvaluetestresults(learner))
pvdf[learner.stats.policy_test_fails, :]

getpolicyvaluetestresults(learner)[1935:1935]
pt1935 = getpolicyvaluetests(learner)[1935]
getpolicy(learner, pt1935.state)
getvalue(learner, pt1935.state, 5)
getvalue(learner, pt1935.state, 2)
getvalue(learner, pt1935.state, 3)
getvalue(learner, pt1935.state, 7)

pt1154 = getpolicyvaluetests(learner)[1154]
pt1154 |> svg
getpolicy(learner, pt2438.state)

pvdf[pvdf."P/F" .== "Fail", :]

using Plots
plot(getlosses(learner))

# LR finder

writable = learner.replay_buffers.d[learner.replay_buffers.writable_uuid]
g1 = writable[end]
g1.net_step
g1
g1 |> svg
g1s5 = get_state(g1, 5)
g1s5 |> typeof

g2 = writable[end-1]

using Profile
@time play_game(learner)

@profile play_game(learner)

open("/tmp/prof.txt", "w") do s
    Profile.print(IOContext(s))
end

learner.play.representation_network[1].weight
learner.play.representation_network[2]
learner.play.representation_network[2].μ
learner.play.representation_network[2].μ
learner.play.representation_network[2].active
learner.train.representation_network[2].active
learner.test.representation_network[2].active