function Header(io::IOStream)
    # check magic number
    (read(io, UInt8) == 0x6c && read(io, UInt8) == 0x1b) || throw(ArgumentError("wrong magic number in the input file"))
    # 2 bytes so far
    storage_mode = read(io, UInt8) # 3 bytes
    storage_mode == 0x10 || throw(ArgumentError("Only the standard format supported for now"))
    n_variants = read(io, UInt32) # 7 bytes
    n_samples = read(io, UInt32) # 11 bytes

    if n_samples <= 2 ^ 8
        bytes_per_sample_id = 1
    elseif n_samples <= 2 ^ 16
        bytes_per_sample_id = 2
    elseif n_samples <= 2 ^ 24
        bytes_per_sample_id = 3
    else
        bytes_per_sample_id = 4
    end

    format_num = read(io, UInt8) # 12 bytes

    # figure out bits_per_variant_type and bytes_per_record_length
    format_num_0_3 = format_num & 0x0f

    (bits_per_variant_type, bytes_per_record_length) = variant_type_lengths[format_num_0_3]

    bytes_allele_counts = (format_num & 0x30) >> 4

    # provisional reference
    provisional_reference = (format_num & 0xc0) >> 6

    n_blocks = ceil_int(n_variants, 2 ^ 16) #Int(ceil(n_variants / 2 ^ 16))
    variant_block_offsets = Vector{UInt64}(undef, n_blocks)
    read!(io, variant_block_offsets) # 12 + 8n_blocks
    #variant_block_offsets = reinterpret(UInt64, data[13:(12 + n_blocks * 8)])

    offset = convert(UInt64, 12 + 8n_blocks)

    t_variant_types = nothing
    sectors_variant_types = []

    dtype_variant_length = bytes_to_UInt[bytes_per_record_length]
    t_variant_sizes = nothing
    sectors_variant_lengths = []

    if bytes_allele_counts != 0
        dtype_allele_counts = bytes_to_UInt[bytes_allele_counts]
        t_allele_counts = nothing
        sectors_allele_counts = []
    else
        sectors_allele_counts = nothing
    end
    if provisional_reference == 0x03
        t_provisional_reference_flags = nothing
        sectors_provisional_reference = []
    else
        sectors_provisional_reference = nothing
    end
    for i in 1:n_blocks
        block_size = i < n_blocks ? convert(UInt, 2 ^ 16) : begin
            rem = n_variants % convert(UInt, 2 ^ 16)
            rem > 0 ? rem : convert(UInt, 2 ^ 16)
        end

        # read variant type track
        data_per_byte = 8 ÷ bits_per_variant_type
        size_variant_types = ceil_int(block_size, data_per_byte)
        #convert(UInt, ceil(block_size / (8 ÷ bits_per_variant_type)))
        arr = read(io, size_variant_types) #data[offset + 1 : offset + size_variant_types]
        push!(sectors_variant_types, Ref(arr))
        t_variant_types = typeof(arr)
        offset += size_variant_types

        # read variant record length track
        arr = read(io, block_size * bytes_per_record_length) #data[offset + 1 : offset + block_size * bytes_per_record_length]
        reinterpreted = reinterpret(dtype_variant_length, arr)
        t_variant_sizes = typeof(reinterpreted)
        push!(sectors_variant_lengths, Ref(reinterpreted))
        offset += bytes_per_record_length * block_size

        # read allele counts track
        if sectors_allele_counts !== nothing
            arr = read(io, block_size * bytes_allele_counts)
            #arr = data[offset + 1:offset + block_size * bytes_allele_counts]
            reinterpreted = reinterpret(dtype_allele_counts, arr)
            t_allele_counts = typeof(reinterpreted)
            push!(sectors_allele_counts, Ref(reinterpreted))
            offset += bytes_allele_counts * block_size
        end

        # read provisional reference flags track
        if sectors_provisional_reference !== nothing
            size_pr = ceil_int(block_size, 8)#convert(UInt, ceil(block_size / 8))
            arr = read(io, size_pr)
            #arr = data[offset + 1 : offset + size_pr]
            t_provisional_reference_flags = typeof(arr)
            push!(sectors_provisional_reference, Ref(arr))
            offset += size_pr
        end
    end
    if bits_per_variant_type == 8
        VTT = ScatteredVector{UInt8, t_variant_types}
        variant_types = ScatteredVector{UInt8, t_variant_types}(
            sectors_variant_types, n_variants, n_blocks)
    else
        VTT = ScatteredBitsVector{t_variant_types}
        variant_types = ScatteredBitsVector{t_variant_types}(
            sectors_variant_types, bits_per_variant_type,
            n_variants, n_blocks)
    end
    VLT = ScatteredVector{dtype_variant_length,t_variant_sizes}
    variant_lengths = ScatteredVector{dtype_variant_length,t_variant_sizes}(
        sectors_variant_lengths,
        n_variants, n_blocks)
    if sectors_allele_counts !== nothing
        ACT = ScatteredVector{dtype_allele_counts,t_allele_counts}
        allele_counts = ScatteredVector{dtype_allele_counts,t_allele_counts}(
            sectors_allele_counts,
            n_variants, n_blocks)
    else
        ACT = Nothing
        allele_counts = nothing
    end
    if sectors_provisional_reference !== nothing
        PRFT = ScatteredBitsVector{t_provisional_reference_flags}
        provisional_reference_flags = ScatteredBitsVector{t_provisional_reference_flags}(
            sectors_provisional_reference, 1,
            n_variants, n_blocks)
    else
        PRFT = Nothing
        provisional_reference_flags = nothing
    end

    # Store the most recent non-LD-compressed variant for each LD-compressed variant
    # Useful for random-access of variants
    most_recent_non_ld = Dict{UInt, Variant}()
    buf = nothing
    offset = variant_block_offsets[1]
    for j in 1:n_variants
        vt = variant_types[j]
        vl = variant_lengths[j]
        if !((vt & 0x07 == 0x02) || (vt & 0x07 == 0x03))
            buf = Variant(j, offset, vt, vl)
        else
            most_recent_non_ld[j] = buf
        end
        offset += vl
    end

    Header{VTT, VLT, ACT, PRFT}(storage_mode, n_variants, n_samples,
        bits_per_variant_type, bytes_per_record_length,
        bytes_allele_counts, bytes_per_sample_id, provisional_reference,
        n_blocks, variant_block_offsets,
        variant_types, variant_lengths, allele_counts, provisional_reference_flags,
        most_recent_non_ld)
end

function Header(filename::String)
    io = open(filename)
    Header(io)
end


"""
Abstract type for all kinds of gene product associations
"""
abstract type GeneProductAssociation end

"""
Gene product reference in the association expression
"""
struct GPARef <: GeneProductAssociation
    gene_product::String
end

"""
Boolean binary "and" in the association expression
"""
struct GPAAnd <: GeneProductAssociation
    terms::Vector{GeneProductAssociation}
end

"""
Boolean binary "or" in the association expression
"""
struct GPAOr <: GeneProductAssociation
    terms::Vector{GeneProductAssociation}
end

"""
A simplified representation of MathML-specified math AST
"""
abstract type Math end

"""
A literal value (usually a numeric constant) in mathematical expression
"""
struct MathVal{T} <: Math where {T}
    val::T
end

"""
An identifier (usually a variable name) in mathematical expression
"""
struct MathIdent <: Math
    id::String
end

"""
A constant identified by name (usually something like `pi`, `e` or `true`) in
mathematical expression
"""
struct MathConst <: Math
    id::String
end

"""
A special value representing the current time of the simulation, with a special
name.
"""
struct MathTime <: Math
    id::String
end

"""
Function application ("call by name", no tricks allowed) in mathematical expression
"""
struct MathApply <: Math
    fn::String
    args::Vector{Math}
end

"""
Function definition (aka "lambda") in mathematical expression
"""
struct MathLambda <: Math
    args::Vector{String}
    body::Math
end

"""
SBML Compartment with sizing information.
"""
struct Compartment
    name::Maybe{String}
    constant::Maybe{Bool}
    spatial_dimensions::Maybe{Int}
    size::Maybe{Float64}
    units::Maybe{String}
    notes::Maybe{String}
    annotation::Maybe{String}
    Compartment(na, c, sd, s, u, no = nothing, an = nothing) = new(na, c, sd, s, u, no, an)
end

"""
Reaction with stoichiometry that assigns reactants and products their relative
consumption/production rates (accessible in field `stoichiometry`), lower/upper
bounds (in tuples `lb` and `ub`, with unit names), and objective coefficient
(`oc`). Also may contains `notes` and `annotation`.
"""
struct Reaction
    stoichiometry::Dict{String,Float64}
    lb::Tuple{Float64,String}
    ub::Tuple{Float64,String}
    oc::Float64
    gene_product_association::Maybe{GeneProductAssociation}
    kinetic_math::Maybe{Math}
    reversible::Bool
    notes::Maybe{String}
    annotation::Maybe{String}
    Reaction(s, l, u, o, as, km, r, n = nothing, an = nothing) =
        new(s, l, u, o, as, km, r, n, an)
end

"""
Species metadata -- contains a human-readable `name`, a `compartment`
identifier, `formula`, `charge`, and additional `notes` and `annotation`.
"""
struct Species
    name::Maybe{String}
    compartment::String
    boundary_condition::Maybe{Bool}
    formula::Maybe{String}
    charge::Maybe{Int}
    initial_amount::Maybe{Tuple{Float64,Maybe{String}}}
    initial_concentration::Maybe{Tuple{Float64,Maybe{String}}}
    only_substance_units::Maybe{Bool}
    notes::Maybe{String}
    annotation::Maybe{String}
    Species(na, co, b, f, ch, ia, ic, osu, no = nothing, a = nothing) =
        new(na, co, b, f, ch, ia, ic, osu, no, a)
end

"""
Gene product metadata.
"""
struct GeneProduct
    name::Maybe{String}
    label::Maybe{String}
    notes::Maybe{String}
    annotation::Maybe{String}
    GeneProduct(na, l, no = nothing, a = nothing) = new(na, l, no, a)
end

"""
Custom function definition.
"""
struct FunctionDefinition
    name::Maybe{String}
    body::Maybe{Math}
    notes::Maybe{String}
    annotation::Maybe{String}
    FunctionDefinition(na, b, no = nothing, a = nothing) = new(na, b, no, a)
end

"""
Structure that collects the model-related data. Contains `parameters`, `units`,
`compartments`, `species` and `reactions` and `gene_products`, and additional
`notes` and `annotation` (also present internally in some of the data fields).
The contained dictionaries are indexed by identifiers of the corresponding
objects.
"""
struct Model
    parameters::Dict{String,Float64}
    units::Dict{String,Number}
    compartments::Dict{String,Compartment}
    species::Dict{String,Species}
    reactions::Dict{String,Reaction}
    gene_products::Dict{String,GeneProduct}
    function_definitions::Dict{String,FunctionDefinition}
    notes::Maybe{String}
    annotation::Maybe{String}
    Model(p, u, c, s, r, g, f, n = nothing, a = nothing) = new(p, u, c, s, r, g, f, n, a)
end


#

# The simulation gets some input from other simulations
# Every x step ?
#
# DataFrame
# Time
# susceptible
# Infected (?)
# Symptomatic
# Hospitalized
# Dead
# Recovered
#
# only SEII2R have compliance because HD are fully-compliant
# E is only knowable from "Nowcasting" or similar things
#

# Every agent has an agent-state from SEIIHDR, but not all are knowable for
# other agents

# From the delta between time stamps recreate the states in our agents.

using Agents

include("../agents/agent.jl")
include("../spaces/network.jl")
include("../dataimport/modelproperties.jl")
include("stepping.jl")
include("evaluation.jl")
include("../dataimport/loadmodel.jl")

# to allow a unique signature of the model configurtion
# the getAgentType function will determine the agent type
# from a configuration file.

mprops = getModelProperties()

function getAgentType()
    return SimplePolicyAgent
end

# same as above
function getSpace()
        props = Dict("type" => "barabasi",
                     "n" => 100,
                     "k" => 2,
                     "n0" => 10
        )
    generateSpace(props)
end

function getSpace()
        props = Dict("type" => "watts",
                     "n" => 100,
                     "k" => 2,
                     "beta" => 0.5
        )
    generateSpace(props)
end

# method for running
function runsimulation(abm)

 run!(abm, agents_step!, 100;
    adata =  a_data_props,
    mdata = m_data_props,
    replicates = 10,
    parallel = true)

end


using CorpusLoaders
using MultiResolutionIterators
using TextAnalysis
using TextAnalysis: CRF, crf_loss
using WordTokenizers
using Embeddings
using Flux
using Flux: onehot, param, onehotbatch, Conv, LSTM, flip, Momentum, reset!, onecold, batchseq
using Tracker
using BSON: @save, @load
using LinearAlgebra
using Base.Iterators: partition
using CorpusLoaders: NERTaggedWord

println("Dependencies Loaded")

device = :gpu

if device == :gpu
    using CuArrays
end

CHAR_EMBED_DIMS = 25
WORD_EMBED_DIMS = 100

train_set = load(CoNLL(), "train") # training set
# test_set = load(CoNLL(), "test") # test set
dev_set = load(CoNLL(), "dev") # dev set

dataset = flatten_levels(train_set, lvls(CoNLL, :document)) |> full_consolidate
dev_dataset = flatten_levels(dev_set, lvls(CoNLL, :document)) |> full_consolidate

get_pos(w::NERTaggedWord) = w.pos

X_train = [CorpusLoaders.word.(sent) for sent in dataset]
Y_train = [get_pos.(sent) for sent in dataset]
@assert length.(X_train) == length.(Y_train)

X_dev = [CorpusLoaders.word.(sent) for sent in dev_dataset]
Y_dev = [get_pos.(sent) for sent in dev_dataset]
@assert length.(X_dev) == length.(Y_dev)

words_vocab = unique(vcat(X_train...))
alphabets = unique(vcat(collect.(words_vocab)...))
labels = unique(vcat(Y_train...))
DENSE_OUT_SIZE = num_labels = length(labels)

#### Word Embeddings
# using these indices to generate the flux one hot input word embeddings vectors.

embtable = load_embeddings(GloVe, 2)
get_word_index = Dict(word => ii for (ii, word) in enumerate(embtable.vocab))
get_word_from_index = Dict(value => key for (key, value) in get_word_index)
W_word_Embed = (embtable.embeddings)

# One Vec for unknown chars
UNK_Word = "<UNK>"
UNK_Word_Idx = length(get_word_index) + 1
embedding_vocab_length = length(get_word_index) + 1
@assert UNK_Word ∉ collect(keys(get_word_index))
get_word_index[UNK_Word] = UNK_Word_Idx
get_word_from_index[UNK_Word_Idx] = UNK_Word
W_word_Embed = hcat(W_word_Embed, rand(WORD_EMBED_DIMS))

# One Vec for PAD chars
PAD_Word = "<pad>"
PAD_Word_Idx = length(get_word_index) + 1
embedding_vocab_length = length(get_word_index) + 1
@assert PAD_Word ∉ collect(keys(get_word_index))
get_word_index[PAD_Word] = PAD_Word_Idx
get_word_from_index[PAD_Word_Idx] = PAD_Word
W_word_Embed = hcat(W_word_Embed, rand(WORD_EMBED_DIMS))

W_word_Embed = Float32.(W_word_Embed)
@assert size(W_word_Embed, 2) == embedding_vocab_length

##### Char Embeddings
UNK_char = '¿'
PAD_char = '█'
@assert UNK_char ∉ alphabets
@assert PAD_char ∉ alphabets
push!(alphabets, UNK_char)
push!(alphabets, PAD_char)
sort!(unique!(alphabets))

# using these indices to generate the flux one hot input Char embeddings vectors.

W_Char_Embed = rand(CHAR_EMBED_DIMS, length(alphabets))
W_Char_Embed = Float32.(W_Char_Embed)
W_Char_Embed = param(W_Char_Embed)
get_char_index = Dict(char => ii for (ii, char) in enumerate(alphabets))
get_char_from_index = Dict(value => key for (key, value) in get_char_index)

println("Embeddings Made")

############# Creating onehoot vectors (useful for embeddings lookup), convenient
BATCH_SIZE = 10
X_train = sort(X_train, by=length, alg=MergeSort)
Y_train = sort(Y_train, by=length, alg=MergeSort) # For stable sorting

onehotword(word) = onehot(get(get_word_index, lowercase(word), get_word_index[UNK_Word]), 1:embedding_vocab_length)
onehotbatchword(words) = onehotbatch([get(get_word_index, lowercase(word), get_word_index[UNK_Word]) for word in words], 1:embedding_vocab_length)
onehotchars(word) = onehotbatch([get(get_char_index, c, get_char_index[UNK_char]) for c in word], 1:length(alphabets))
onehotlabel(label) = onehot(label, labels)
onehotinput(word) = (onehot(get(get_word_index, lowercase(word), get_word_index[UNK_Word]), 1:embedding_vocab_length),
                onehotbatch([get(get_char_index, c, get_char_index[UNK_char]) for c in word], 1:length(alphabets)))

# align(k) = [[j[i]  for j in k] for i in 1:length(k[1])]
# batch_char_input(ohed_chars, p) = [align(padded(b, p)) for b in collect(partition(ohed_chars, BATCH_SIZE))]
# batches(xs, p) = [batchseq(b, p) for b in partition(xs, BATCH_SIZE)]

oh_seq(arr, f) = [f(element) for element in arr]

padded(b, p, m) = [length(elem) < m ? [elem..., repeat([p], m - length(elem))...] : elem for elem in b]
padded(b, p) = padded(b, p, maximum(length.(b)))

align_chars_new(b, p, m) = [[[p, p, elem[i]..., repeat([p], m[i] - length(elem[i]))..., p, p] for i in 1:length(elem)] for elem in b]
align_chars_new(b, p) = align_chars_new(b, p, [maximum([length(j[i]) for j in b])   for i in 1:length(b[1])])

batch_oh_chars(o) = [[char_w_seq[i] for char_w_seq in o] for i in 1:length(o[1])]
gpuit(x) = cu.(x)
batch_new_chars(xs, p) = [gpuit.(batch_oh_chars(oh_seq.(align_chars_new(padded(b, p), '█'), onehotchars))) for b in partition(xs, BATCH_SIZE)]
# X_chars_train = batch_char_input([oh_seq(sentence, onehotchars) for sentence in X_train], onehotchars("█")) # A Bunch Sequences of Array of Chars, done to prevent repeated computations.
X_newchars_train = batch_new_chars([collect.(sentence) for sentence in X_train], collect("█"))

padded_word_batches(b, p, m) = [[length(elem) < i ? p : elem[i] for elem in b] for i in 1:m]
padded_word_batches(b, p) = padded_word_batches(b, p, maximum(length.(b)))
X_words_train = [cu.(oh_seq(padded_word_batches(b, PAD_Word), onehotbatchword)) for b in partition(X_train, BATCH_SIZE)] # A Bunch of Sequences of words, i.e. sentences

X_input_train = [[(k,l) for (k,l) in zip(i,j)] for (i,j) in zip(X_words_train, X_newchars_train)]
Y_oh_train = [cu.(oh_seq.(b, onehotlabel)) for b in partition(Y_train, BATCH_SIZE)]
deleteat!(Y_oh_train, length(Y_oh_train))
deleteat!(X_input_train, length(X_input_train))

X_dev = sort(X_dev, by=length, alg=MergeSort)
Y_dev = sort(Y_dev, by=length, alg=MergeSort) # For stable sorting

X_newchars_dev = batch_new_chars([collect.(sentence) for sentence in X_dev], collect("█"))
X_words_dev = [cu.(oh_seq(padded_word_batches(b, PAD_Word), onehotbatchword)) for b in partition(X_dev, BATCH_SIZE)] # A Bunch of Sequences of words, i.e. sentences
X_input_dev = [[(k,l) for (k,l) in zip(i,j)] for (i,j) in zip(X_words_dev, X_newchars_dev)]
Y_oh_dev = [oh_seq.(b, onehotlabel) for b in partition(Y_dev, BATCH_SIZE)]
# deleteat!(Y_oh_dev, length(Y_oh_dev))
# deleteat!(X_input_dev, length(X_input_dev))

# if device == :gpu
#     X_input_train = [cu.(oh_seq(sentence, onehotinput)) for sentence in X_train]
# else
#     X_input_train = [oh_seq(sentence, onehotinput) for sentence in X_train]
# end
# Y_oh_train = [oh_seq(tags_sequence, onehotlabel) for tags_sequence in Y_train]
#
# if device == :gpu
#     Y_oh_train = cu.(Y_oh_train)
# end
#
# X_input_dev = [oh_seq(sentence, onehotinput) for sentence in X_dev]
# Y_oh_dev = [oh_seq(tags_sequence, onehotlabel) for tags_sequence in Y_dev]
#
# @assert length.(X_train) == length.(X_input_train) == length.(Y_oh_train)

#################### MODEL ######################
# 1. Character Embeddings with CNNs
# 2. Obtain Word Embeddings, concatenate Char & Word Embeddings
# 3. Bi-LSTM
# 4. Softmax or CRF on each unit in sequence

CNN_OUTPUT_SIZE = 30
CONV_WINDOW_LENGTH = 3
LSTM_STATE_SIZE = 200
DROPOUT_RATE_CNN = DROPOUT_INPUT_EMBEDDING = DROPOUT_OUT_LAYER = Float32(0.5)

# 1. Dropout before conv, Max poll layer, PADDING on both sides, hyperparams = window size, output vector size.
W_Char_Embed = W_Char_Embed |> gpu
conv1 = Conv((CHAR_EMBED_DIMS, CONV_WINDOW_LENGTH), 1=>CNN_OUTPUT_SIZE, pad=(0,2)) |> gpu
dropout1 = Dropout(DROPOUT_RATE_CNN) |> gpu

slice_embed(x) = W_Char_Embed * x
char_features = Chain(x -> slice_embed.(x),
                      x -> cat(x..., dims=4),
                      dropout1,
                      conv1,
                      x -> maximum(x, dims=2),
                      x -> reshape(x, CNN_OUTPUT_SIZE, BATCH_SIZE))
# 2. Input embeddings

W_word_Embed = param(W_word_Embed) |> gpu # For trainable
get_word_embedding(w) = W_word_Embed * w |> gpu # works coz - onehot

# Dropout before LSTM
dropout_embed = Dropout(DROPOUT_INPUT_EMBEDDING) |> gpu
input_embeddings((w, cs)) = dropout_embed(vcat(get_word_embedding(w), char_features(cs))) |> gpu

# 3. Bi-LSTM

if device == :gpu
    forward_lstm = LSTM(CNN_OUTPUT_SIZE + WORD_EMBED_DIMS, LSTM_STATE_SIZE) |> gpu
    backward = LSTM(CNN_OUTPUT_SIZE + WORD_EMBED_DIMS, LSTM_STATE_SIZE) |> gpu
else
    forward_lstm = LSTM(CNN_OUTPUT_SIZE + WORD_EMBED_DIMS, LSTM_STATE_SIZE)
    backward = LSTM(CNN_OUTPUT_SIZE + WORD_EMBED_DIMS, LSTM_STATE_SIZE)
end
backward_lstm(x) = reverse(backward.(reverse(x)))

bilstm_layer(x) = vcat.(forward_lstm.(x), backward_lstm(x))

# 4. Softmax / CRF Layer
# Dropout after LSTM

if device == :gpu
    dropout2 = Dropout(DROPOUT_OUT_LAYER) |> gpu
    d_out = Dense(LSTM_STATE_SIZE * 2, DENSE_OUT_SIZE + 2) |> gpu
else
    dropout2 = Dropout(DROPOUT_OUT_LAYER)
    d_out = Dense(LSTM_STATE_SIZE * 2, DENSE_OUT_SIZE + 2)
end

m = Chain(x -> input_embeddings.(x),
                bilstm_layer,
                x -> dropout2.(x),
                x -> d_out.(x))

if device == :gpu
    m = gpu(m)
end
# dropout.(bilstm_layer(input_embeddings.(w_cs))) |> gpu

using TextAnalysis: crf_loss, CRF
Flux.@treelike TextAnalysis.CRF

c = TextAnalysis.CRF(num_labels)

init_α = Float32.(fill(-10000, (c.n + 2, 1)))
init_α[c.n + 1] = 0
if device == :gpu
    init_α = cu(init_α)
    c = c |> gpu
end

geti(x, i) = x[:, i]
loss(x, y) = sum([crf_loss(c, [geti(x[j], i) for j in 1:length(y[i])], y[i], init_α) for i in 1:BATCH_SIZE])

η = 0.01 # learning rate
β = 0.05 # rate decay
ρ = 0.9 # momentum

opt = Flux.Optimiser(ExpDecay(β), Momentum(η, ρ))
data = zip(X_input_train, Y_oh_train)

println("Model Built")

function try_outs()
    Flux.testmode!(dropout1)
    Flux.testmode!(dropout2)
    Flux.testmode!(dropout_embed)

    confusion_matrix = zeros(Int, (num_labels, num_labels))
    for (x, y) in zip(X_input_dev, Y_oh_dev)
        reset!(forward_lstm)
        reset!(backward)
        x_seq = m(x)
        sum([crf_loss(c, [geti(x_seq[j], i) for j in 1:length(y[i])], y[i], init_α) for i in 1:BATCH_SIZE])

        ohs = [TextAnalysis.viterbi_decode(cpu(c), cpu.([geti(x_seq[j], i) for j in 1:length(y[i])]), cpu(init_α)) for i in 1:BATCH_SIZE]

        for i in 1:BATCH_SIZE
            for d in zip(ohs[i], y[i])
                confusion_matrix[d[1].ix, d[2].ix] += 1
            end
        end
    end

    s1 = sum(confusion_matrix, dims = 2)
    s2 = sum(confusion_matrix, dims = 1)
    dg = diag(confusion_matrix)

    a = sum([s1[i] == 0 ? 0 : dg[i]/s1[i] for i in eachindex(dg)]) / num_labels
    b = sum(dg ./ s2') / num_labels

    println("Precision and recall are:", a, " ", b)
    println("F1 is:", (2 * a * b) / (a + b))

    Flux.testmode!(dropout1, false)
    Flux.testmode!(dropout2, false)
    Flux.testmode!(dropout_embed, false)
end

function save_weights(epoch, conv1, W_word_Embed, W_Char_Embed, forward_lstm, backward, d_out, c)
    conv_cpu = conv1 |> cpu
    W_word_cpu = W_word_Embed |> cpu
    W_char_cpu = W_Char_Embed |> cpu
    forward_lstm_cpu = forward_lstm |> cpu
    backward_lstm_cpu = backward |> cpu
    crf_cpu = c |> cpu
    d_cpu = d_out |> cpu

    @save "./weights/conv_cpu$(epoch).bson" conv_cpu
    @save "./weights/W_word_cpu$(epoch).bson" W_word_cpu
    @save "./weights/W_char_cpu$(epoch).bson" W_char_cpu
    @save "./weights/forward_lstm$(epoch).bson" forward_lstm_cpu
    @save "./weights/backward_lstm$(epoch).bson" backward_lstm_cpu
    @save "./weights/d_cpu$(epoch).bson" d_cpu
    @save "./weights/crf$(epoch).bson" crf_cpu
end

UPPER_BOUND = Float32(5.0)

# Gradient Clipping
grad_clipping(g) = min(g, UPPER_BOUND)

# Backward
function back!(p::Flux.Params, l, opt)
    l = Tracker.hook(x -> grad_clipping(x), l)
    grads = Tracker.gradient(() -> l, p)
    Tracker.update!(opt, p, grads)
    return
end

function train(EPOCHS)
    reset!(forward_lstm)
    reset!(backward)

    ps = params(params(conv1)..., params(W_word_Embed)..., params(W_Char_Embed)..., params(forward_lstm, backward)..., params(d_out)..., params(c)...)
    for epoch in 1:EPOCHS
        if epoch % 2 == 1
            save_weights(epoch, conv1, W_word_Embed, W_Char_Embed, forward_lstm, backward, d_out, c)
        end
        reset!(forward_lstm)
        reset!(backward)
        try_outs()
        println("----------------------- Starting with epoch : $(epoch) ----------------------")

        for d in data
            reset!(forward_lstm)
            reset!(backward)

            l = loss(m(d[1]), d[2])
            back!(ps, l, opt)
        end
        reset!(forward_lstm)
        reset!(backward)
    end
end

train(100)
reset!(forward_lstm)
reset!(backward)
# save_weights(51, conv_cpu, W_word_Embed, W_Char_Embed, forward_lstm, backward, d_out, c)
try_outs()


#Type definition for the aerodynamic database, including a function to
# add new contributions per contribution. (TBD: Including a remove function)
export AbstractDataBase

export AeroDataBase

abstract type AbstractDataBase end

mutable struct AeroDataBase <: AbstractDataBase
  Variable::String
  ListOfInputs::Array{String}
  ListOfContributions::Array{Any}
  ListOfInputValues::Array{Any}
  ListOfInputOrder::Array{Any}
  AddContribution::Function

  function AeroDataBase(varkey::String)
    this = new()
    this.Variable = varkey
    this.ListOfInputs = []
    this.ListOfContributions = []
    this.ListOfInputValues = []
    this.ListOfInputOrder = []

    #Add method to add contributions to this specific database
    this.AddContribution = function (Contr::Array{Float64},varargs::Tuple)
      #check if correct number of dimensions have been assigned
      nvarargs = length(varargs)
      ndims(Contr) != nvarargs ? error("Incorrect number of inputs") : nothing

      #Add Contr to ListOfContributions
      this.ListOfContributions = [this.ListOfContributions..., Contr]
      NewInput = Vector(nvarargs)
      NewInputOrder = Vector(nvarargs)

      for i=1:length(varargs)
        if IsContainedInStringArray(this.ListOfInputs,varargs[i][1])
          NewInput[i] = varargs[i][2:end]
          NewInputOrder[i] = IdInStringArray(this.ListOfInputs,varargs[i][1])
        else
          this.ListOfInputs = [this.ListOfInputs...,varargs[i][1]]
          NewInput[i] = varargs[i][2:end]
          NewInputOrder[i] = length(this.ListOfInputs)
        end
      end
      this.ListOfInputValues = push!(this.ListOfInputValues, NewInput)
      this.ListOfInputOrder = push!(this.ListOfInputOrder, NewInputOrder)
      return nothing
    end #this.AddContribution
    return this
  end #Constructor
end #typedef

#Example
#=
Cm1 = rand(4,7)
Cm2 = rand(5,7)
x = ["alpha", 2,3,4,5]
y = ["beta",1,5,2,7,9,4,9]
z = ["sigma",1,2,3,4,5]
Cmdb = AeroDataBase("Cm")
Cmdb.AddContribution(Cm1,(x,y))
Cmdb.AddContribution(Cm2,(z,y))
=#


using consonance

push!(LOAD_PATH, "/home/ian/Grad/Research/Disks/DiskJockey/")
push!(LOAD_PATH, "/n/home07/iczekala/DiskJockey/")
push!(LOAD_PATH, "/pool/scout0/DiskJockey/")

using ArgParse

s = ArgParseSettings()
@add_arg_table s begin
    # "--opt1"
    # help = "an option with an argument"
    "--run_index", "-r"
    help = "Output run index"
    arg_type = Int
end

parsed_args = parse_args(ARGS, s)
run_index = parsed_args["run_index"]

using constants
using image

import PyPlot.plt
using LaTeXStrings

function spec_plot(img::image.SkyImage, id::Int)

    fig = plt.figure()
    ax = fig[:add_subplot](111)

    spec = imToSpec(img)

    # CO 2-1 rest frame
    lam0 = cc/230.538e9 * 1e4 # [microns]

    # convert wavelengths to velocities
    vels = c_kms * (img.lams .- lam0)/lam0 # [km/s]

    ax[:plot](vels, spec[:,2], ls="steps-mid")
    # ax[:plot](vels, reverse(spec[:,2]), ls="steps-mid")
    ax[:set_ylabel](L"$f_\nu$ [Jy]")
    ax[:set_xlabel](L"$v$ [km/s]")

    # Annotate the parameters
    p = pars[id]
    mass = p.M_star
    r_c = p.r_c
    incl = p.incl
    label = L"$M_\ast$: " * @sprintf("%.2f", mass) * L" $M_\odot$   $r_c$: " * @sprintf("%.0f", r_c) * L" AU   $i$: " * @sprintf("%.0f", incl) * L"${}^\circ$"

    fig[:text](0.25, 0.1, label)

    plt.savefig(outdir * @sprintf("spec%04d.png", id), dpi=150)
    plt.close("all")

end

# How many frames per process?
start = run_index * nframes_per_proc + 1

ids = Int[i for i=start:(start + nframes_per_proc)]

for id in ids
    if id <= nframes
        fname = outdir * @sprintf("gimage%04d.out", id)
        im = imread(fname)
        skim = imToSky(im, 73.)
        spec_plot(skim, id)
    else
        break
    end
end


"""
```
forecast(m, system, z0; enforce_zlb = false, shocks = Matrix{S}(0,0))

forecast(system, z0, shocks; enforce_zlb = false)
```

### Inputs

- `m::AbstractModel`: model object. Only needed for the method in which `shocks`
  are not provided.
- `system::System{S}`: state-space system matrices
- `kal::Kalman{S}` or `z0::Vector{S}`: result of running the Kalman filter or
  state vector in the final historical period (aka initial forecast period)

where `S<:AbstractFloat`.

### Keyword Arguments

- `cond_type::Symbol`: one of `:none`, `:semi`, or `:full`, used to determine
  how many periods to forecast ahead. If `cond_type in [:semi, :full]`, the
  forecast horizon is reduced by the number of periods of conditional
  data. Defaults to `:none`.
- `enforce_zlb::Bool`: whether to enforce the zero lower bound. Defaults to
  `false`.
- `shocks::Matrix{S}`: matrix of size `nshocks` x `shock_horizon` of shock
  innovations under which to forecast. If `shock_horizon > horizon`, the extra
  periods of shocks will be ignored; if `shock_horizon < horizon`, zeros will be
  filled in for the shocks hitting the remaining forecasted periods.
- `draw_shocks::Bool`: if `isempty(shocks)`, indicates whether to draw shocks
  according to:

  1. If `forecast_tdist_shocks(m)`, draw `horizons` many shocks from a
     `Distributions.TDist(forecast_tdist_df_val(m))`
  2. Otherwise, draw `horizons` many shocks from a
     `DegenerateMvNormal(zeros(nshocks), sqrt(system[:QQ]))`

  or to set `shocks` to a `nshocks` x `horizon` matrix of zeros. Defaults to
  `false`. If `shocks` is provided as a keyword argument, this flag has no
  effect.

### Outputs

- `states::Matrix{S}`: matrix of size `nstates` x `horizon` of forecasted states
- `obs::Matrix{S}`: matrix of size `nobs` x `horizon` of forecasted observables
- `pseudo::Matrix{S}`: matrix of size `npseudo` x `horizon` of forecasted
  pseudo-observables
- `shocks::Matrix{S}`: matrix of size `nshocks` x `horizon` of shock innovations
"""
function forecast{S<:AbstractFloat}(m::AbstractModel, system::System{S},
    z0::Vector{S}; cond_type::Symbol = :none, enforce_zlb::Bool = false,
    shocks::Matrix{S} = Matrix{S}(0, 0), draw_shocks::Bool = false)

    # Numbers of things
    nshocks = n_shocks_exogenous(m)
    horizon = forecast_horizons(m; cond_type = cond_type)

    if isempty(shocks)
        # Populate shocks matrix
        if draw_shocks
            μ = zeros(S, nshocks)
            σ = sqrt.(system[:QQ])
            dist = if forecast_tdist_shocks(m)
                # Use t-distributed shocks
                ν = forecast_tdist_df_val(m)
                DegenerateDiagMvTDist(μ, σ, ν)
            else
                # Use normally distributed shocks
                DegenerateMvNormal(μ, σ)
            end

            shocks = rand(dist, horizon)

            # Forecast without anticipated shocks
            if n_anticipated_shocks(m) > 0
                ind_ant1 = m.exogenous_shocks[:rm_shl1]
                ind_antn = m.exogenous_shocks[Symbol("rm_shl$(n_anticipated_shocks(m))")]
                ant_shock_inds = ind_ant1:ind_antn
                shocks[ant_shock_inds, :] = 0
            end
        else
            shocks = zeros(S, nshocks, horizon)
        end
    else
        # Adjust size of shocks matrix, padding with zeros or cutting off
        # periods of shocks if necessary
        shock_horizon = size(shocks, 2)
        if shock_horizon <= horizon
            shocks0 = zeros(nshocks, horizon - shock_horizon)
            shocks = hcat(shocks, shocks0)
        else
            shocks = shocks[:, 1:horizon]
        end
    end

    # Populate shocks matrix under alternative policy, if
    # user has specified a function to do so
    alt_policy = alternative_policy(m)
    if alt_policy.solve != identity &&
        alt_policy.forecast_init != identity

        shocks, z0 = alt_policy.forecast_init(m, shocks, z0, cond_type = cond_type)
    end

    # Get variables necessary to enforce the zero lower bound in the forecast
    ind_r = m.observables[:obs_nominalrate]
    ind_r_sh = m.exogenous_shocks[:rm_sh]
    zlb_value = forecast_zlb_value(m)

    forecast(system, z0, shocks; enforce_zlb = enforce_zlb,
        ind_r = ind_r, ind_r_sh = ind_r_sh, zlb_value = zlb_value)
end

function forecast{S<:AbstractFloat}(system::System{S}, z0::Vector{S},
    shocks::Matrix{S}; enforce_zlb::Bool = false, ind_r::Int = -1,
    ind_r_sh::Int = -1, zlb_value::S = 0.13/4)

    # Unpack system
    T, R, C = system[:TTT], system[:RRR], system[:CCC]
    Q, Z, D = system[:QQ], system[:ZZ], system[:DD]
    Z_pseudo, D_pseudo = system[:ZZ_pseudo], system[:DD_pseudo]

    # Setup
    nshocks = size(R, 2)
    nstates = size(T, 2)
    nobs    = size(Z, 1)
    npseudo = size(Z_pseudo, 1)
    horizon = size(shocks, 2)

    # Define our iteration function
    function iterate(z_t1, ϵ_t)
        z_t = C + T*z_t1 + R*ϵ_t

        # Change monetary policy shock to account for 0.13 interest rate bound
        if enforce_zlb
            interest_rate_forecast = getindex(D + Z*z_t, ind_r)
            if interest_rate_forecast < zlb_value
                # Solve for interest rate shock causing interest rate forecast to be exactly ZLB
                ϵ_t[ind_r_sh] = 0.
                z_t = C + T*z_t1 + R*ϵ_t
                ϵ_t[ind_r_sh] = getindex((zlb_value - D[ind_r] - Z[ind_r, :]'*z_t) / (Z[ind_r, :]' * R[:, ind_r_sh]), 1)

                # Forecast again with new shocks
                z_t = C + T*z_t1 + R*ϵ_t

                # Confirm procedure worked
                interest_rate_forecast = getindex(D + Z*z_t, ind_r)
                @assert interest_rate_forecast >= zlb_value - 0.01 "interest_rate_forecast = $interest_rate_forecast must be >= zlb_value - 0.01 = $(zlb_value - 0.01)"
            end
        end
        return z_t, ϵ_t
    end

    # Iterate state space forward
    states = zeros(S, nstates, horizon)
    states[:, 1], shocks[:, 1] = iterate(z0, shocks[:, 1])
    for t in 2:horizon
        states[:, t], shocks[:, t] = iterate(states[:, t-1], shocks[:, t])
    end

    # Apply measurement and pseudo-measurement equations
    obs    = D .+ Z*states
    pseudo = D_pseudo .+ Z_pseudo * states

    # Return forecasts
    return states, obs, pseudo, shocks
end


using Zygote

export PrioritizedDQNLearner

"""
    PrioritizedDQNLearner(;kwargs...)

See paper: [Prioritized Experience Replay](https://arxiv.org/abs/1511.05952)

# Keywords

- `approximator`::[`AbstractQApproximator`](@ref): used to get Q-values of a state.
- `target_approximator`::[`AbstractQApproximator`](@ref): similar to `approximator`, but used to estimate the target (the next state).
- `loss_fun`: the loss function.
- `γ::Float32=0.99f0`: discount rate.
- `batch_size::Int=32`
- `update_horizon::Int=1`: length of update ('n' in n-step update).
- `min_replay_history::Int=32`: number of transitions that should be experienced before updating the `approximator`.
- `update_freq::Int=4`: the frequency of updating the `approximator`.
- `target_update_freq::Int=100`: the frequency of syncing `target_approximator`.
- `stack_size::Union{Int, Nothing}=4`: use the recent `stack_size` frames to form a stacked state.
- `default_priority::Float64=100.`: the default priority for newly added transitions.
"""
Base.@kwdef mutable struct PrioritizedDQNLearner{Tq<:AbstractQApproximator, Tt<:AbstractQApproximator,Tf, Ts<:Union{Int, Nothing}} <: AbstractLearner
    approximator::Tq
    target_approximator::Tt
    loss_fun::Tf
    stack_size::Ts = 4
    γ::Float32 = 0.99f0
    batch_size::Int = 32
    update_horizon::Int = 1
    min_replay_history::Int = 32
    update_freq::Int = 1
    target_update_freq::Int = 100
    update_step::Int = 0
    default_priority::Float64 = 100.0

    function PrioritizedDQNLearner(approximator::Tq, target_approximator::Tt, loss_fun::Tf, stack_size::Ts, args...) where {Tq, Tt, Tf, Ts}
        copyto!(approximator, target_approximator)  # force sync
        new{Tq,Tt, Tf, Ts}(approximator, target_approximator, loss_fun, stack_size, args...)
    end
end

function update!(learner::PrioritizedDQNLearner{<:NeuralNetworkQ}, batch)
    learner.update_step += 1
    learner.update_step % learner.update_freq == 0 || return nothing

    Q, Qₜ, γ, loss_fun, update_horizon, batch_size = learner.approximator,
        learner.target_approximator,
        learner.γ,
        learner.loss_fun,
        learner.update_horizon,
        learner.batch_size
    states, rewards, terminals, next_states = map(x->to_device(Q, x), (batch.states, batch.rewards, batch.terminals, batch.next_states))
    actions = CartesianIndex.(batch.actions, 1:batch_size)
    priorities = Vector{Float32}()

    gs = gradient(Q) do
        q = batch_estimate(Q, states)[actions]
        q′ = dropdims(maximum(batch_estimate(Qₜ, next_states); dims = 1), dims = 1)
        G = rewards .+ γ^update_horizon .* (1 .- terminals) .* q′

        batch_losses = loss_fun(G, q)
        priorities = (Zygote.dropgrad(batch_losses) .+ 1f-10)
        mean(batch_losses)
    end

    update!(Q, gs)

    if learner.update_step % learner.target_update_freq == 0
        copyto!(Qₜ, Q)
    end

    priorities |> to_host
end

function extract_transitions(
    buffer::CircularTurnBuffer{PRTSA},
    learner::PrioritizedDQNLearner,
)
    if length(buffer) > learner.min_replay_history
        inds = sample_indices(buffer, learner.batch_size, learner.update_horizon, learner.stack_size)
        inds, extract_SARTS(buffer, inds, learner.γ, learner.update_horizon, learner.stack_size)
    else
        nothing
    end
end


using Molly
using DelimitedFiles

data_dir = normpath(dirname(pathof(Molly)), "..", "data")
ff_dir = joinpath(data_dir, "force_fields")
openmm_dir = joinpath(data_dir, "openmm_6mrr")

ff = OpenMMForceField(joinpath.(ff_dir, ["ff99SBildn.xml", "tip3p_standard.xml", "his.xml"])...)

atoms, atoms_data, specific_inter_lists, general_inters, neighbor_finder, coords, box_size = setupsystem(
    joinpath(data_dir, "6mrr_equil.pdb"), ff)

n_steps = 500
timestep = 0.0005u"ps"
temp = 300.0u"K"
velocities = SVector{3}.(eachrow(readdlm(joinpath(openmm_dir, "velocities_300K.txt"))))u"nm * ps^-1"

s = Simulation(
    simulator=VelocityVerlet(),
    atoms=atoms,
    atoms_data=atoms_data,
    specific_inter_lists=specific_inter_lists,
    general_inters=general_inters,
    coords=coords,
    velocities=velocities,
    temperature=temp,
    box_size=box_size,
    neighbor_finder=neighbor_finder,
    thermostat=AndersenThermostat(1.0u"ps"),
    timestep=timestep,
    n_steps=n_steps,
)

parallel = false
simulate!(s, 5; parallel=parallel)
@time simulate!(s, n_steps; parallel=parallel)

parallel = true
simulate!(s, 5; parallel=parallel)
@time simulate!(s, n_steps; parallel=parallel)


export @surf_expr, @surf_toplevel, @surf_type
function surf_type(s::Symbol)
    Surf.TVar(s)
end

function surf_type(ex)
    @match ex begin
        QuoteNode(a::Symbol) => Surf.TSym(a)
        :(NewType($(a::Symbol))) => Surf.TNew(a)
        :(implicit[$t]) => Surf.TImplicit(surf_type(t))
        :(Fn[$a, $b]) => Surf.TArrow(surf_type(a), surf_type(b))
        :($a.?($label)) => Surf.TQuery(string(label), surf_type(a))
        :($a.$(b::Symbol)) => Surf.TApp(Surf.TVar(:field), Surf.TSym(b))
        :($f[$a]) => Surf.TApp(surf_type(f), surf_type(a))
        Expr(:tuple, args...) =>
            Surf.TTuple(Surf.TyExpr[surf_type(a) for a in args])
        Expr(:->, a, Expr(:block, ::LineNumberNode, b)) =>
            Surf.TArrow(surf_type(a), surf_type(b))
        :($t where {$(tvars...)}) =>
            Surf.TForall(Symbol[tvars...], surf_type(t))
    end
end

function surf_expr(n::Symbol)
    Surf.EVar(n)
end

function surf_expr(i::Union{Integer, AbstractString, AbstractChar, AbstractFloat})
    Surf.EVal(i)
end

function surf_expr(exp::Expr)
    @match exp begin
    :(@julia $(::LineNumberNode) $jlex) => Surf.EExt(jlex)
    :(@extern $(::LineNumberNode) $jlex) => Surf.EExt(jlex)
    :($a.?($label)) => Surf.EQuery(string(label), surf_expr(a))
    :($a.$(b::Symbol)) => Surf.EApp(Surf.EVar(:op_Getter), Surf.EField(surf_expr(a), b))
    :($f($(args...))) =>
        length(args) === 1 ?
        Surf.EApp(surf_expr(f), surf_expr(args[1])) :
        Surf.EApp(
            surf_expr(f),
            Surf.ETup(Surf.Expr[surf_expr(arg) for arg in args]))
    Expr(:tuple, elts...) =>
        Surf.ETup(Surf.Expr[surf_expr(arg) for arg in elts])

    Expr(:->, s::Symbol, exp && Expr(:block, ln::LineNumberNode, _...)) =>
        Surf.ELoc(
            ln,
            Surf.EFun(s, surf_expr(exp)))

    Expr(:if, cond, arm1, arm2) =>
        Surf.EInt(surf_expr(cond), surf_expr(arm1), surf_expr(arm2))

    Expr(:let, :($_ = $_) && bind, body) =>
        surf_expr(Expr(:let, Expr(:block, bind), body))

    Expr(:let, Expr(:block, decls...), body) =>
        Surf.ELet(surf_decls(decls), surf_expr(body))

    Expr(:block, suite...) =>
        begin
            if isempty(suite)
                Surf.ETup(Surf.Expr[])
            else
                ret = surf_expr(suite[end])
                foldr((@view suite[1:end-1]), init=ret) do new, last
                    @match new begin
                        ln::LineNumberNode =>
                            @match last begin
                                Surf.ELet(vec, _) =>
                                    begin
                                        pushfirst!(vec, DLoc(ln))
                                        last
                                    end
                                _ => Surf.ELoc(ln, last)
                            end
                        :($_ :: $_) || :($_ = $_) || :($_.?($_)) =>
                            @match last begin
                                Surf.ELet(vec, _) =>
                                    begin
                                        pushfirst!(vec, surf_decl(new))
                                        last
                                    end
                                _ => Surf.ELet([surf_decl(new)], last)
                            end
                        _ => Surf.ELet([Surf.DBind(:_, surf_expr(new))], last)
                    end
                end
            end
        end
    end
end

surf_decl(decl::Any) = @match decl begin
    :($(x::Symbol).?($label)) => Surf.DQuery(string(label), x)
    :($(x::Symbol) :: $t) => Surf.DAnn(x, surf_type(t))
    :($(x::Symbol) = $v) => Surf.DBind(x, surf_expr(v))
    ln::LineNumberNode => Surf.DLoc(ln)
end

function surf_decls(decls::AbstractVector)
    Surf.Decl[surf_decl(decl) for decl in decls]
end

macro surf_toplevel(ex)
    @match ex begin
        Expr(:block, stmts...) => surf_decls(stmts)
        _ => error("syntax error")
    end
end

function surf_module(ex)
    @match ex begin
        Expr(:module, true, name, stmts...) => (name, surf_decls(stmts))
        _ => error("syntax error")
    end
end

macro surf_expr(ex)
    surf_expr(ex)
end

macro surf_type(ex)
    surf_type(ex)
end


"""
The gene-reaction rules (grr) are written as
[[gene1 and gene2...] or [gene3 and...] ...]
so for a reaction to be available it is sufficient that one group
is available, but inside a group all of the genes need to be available
"""

@testset "knockout_single_gene" begin
    m = StandardModel()
    add_metabolite!(m, Metabolite("A"))
    add_metabolite!(m, Metabolite("B"))
    add_gene!(m, Gene("g1"))
    add_gene!(m, Gene("g2"))
    add_reaction!(
        m,
        Reaction("v1", metabolites = Dict("A" => -1.0, "B" => 1.0), grr = [["g1"]]),
    )
    add_reaction!(
        m,
        Reaction("v2", metabolites = Dict("A" => -1.0, "B" => 1.0), grr = [["g1", "g2"]]),
    )
    add_reaction!(
        m,
        Reaction("v3", metabolites = Dict("A" => -1.0, "B" => 1.0), grr = [["g1"], ["g2"]]),
    )
    add_reaction!(
        m,
        Reaction(
            "v4",
            metabolites = Dict("A" => -1.0, "B" => 1.0),
            grr = [["g1", "g2"], ["g2"]],
        ),
    )

    remove_gene!(m, "g1", knockout_reactions = true)

    @test length(m.reactions) == 2
    @test !haskey(m.reactions, "v1")
    @test !haskey(m.reactions, "v2")
end

@testset "knockout_multiple_genes" begin
    m = StandardModel()
    add_metabolite!(m, Metabolite("A"))
    add_metabolite!(m, Metabolite("B"))
    add_gene!(m, Gene("g1"))
    add_gene!(m, Gene("g2"))
    add_gene!(m, Gene("g3"))
    add_reaction!(
        m,
        Reaction(
            "v1",
            metabolites = Dict("A" => -1.0, "B" => 1.0),
            grr = [["g1"], ["g2"], ["g3"]],
        ),
    )
    add_reaction!(
        m,
        Reaction("v2", metabolites = Dict("A" => 1.0, "B" => -1.0), grr = [["g1"], ["g3"]]),
    )

    remove_genes!(m, ["g1", "g3"], knockout_reactions = true)

    @test haskey(m.reactions, "v1")
    @test !haskey(m.reactions, "v2")
end


export Q_update!

"""
update the Q matrix in SparseMatrixCSC or CDS format. Q=sum_{i=1}^p AtA[i]rho[i]
if any of the rho[i] change
"""
function Q_update!(
                  Q            ::Union{Array{TF,2},SparseMatrixCSC{TF,Int64},joAbstractLinearOperator{TF, TF}},
                  AtA          ::Union{Vector{Array{TF, 2}}, Vector{SparseMatrixCSC{TF, TI}}, Vector{joAbstractLinearOperator{TF, TF}}, Vector{Union{Array{TF, 2}, SparseMatrixCSC{TF, TI}, joAbstractLinearOperator{TF, TF}}}},
                  set_Prop,
                  rho          ::Vector{TF},
                  ind_updated  ::Vector{TI},
                  log_PARSDMM,
                  i            ::Integer,
                  Q_offsets=[] ::Vector{TI}
                  ) where {TF<:Real,TI<:Integer}

  if isempty(ind_updated) == false
    if typeof(Q) <: joAbstractLinearOperator #some of the AtA are JOLI operators
      #JOLI operators cannot be defined recursively for many iterations for many operators
      #because this causes memory issues and slow performance -> re-assemble Q from scratch
      # for ii=1:length(ind_updated)
      #   if typeof(AtA[ind_updated[ii]]) <: joAbstractLinearOperator
      #     Q = Q + ( AtA[ind_updated[ii]] )*( rho[ind_updated[ii]]-log_PARSDMM.rho[i,ind_updated[ii]] );
      #   else
      #     Q = Q + ( joMatrix(AtA[ind_updated[ii]]) )*( rho[ind_updated[ii]]-log_PARSDMM.rho[i,ind_updated[ii]] );
      #   end
      # end
      if typeof(AtA[1]) <: joAbstractLinearOperator
        Q = rho[1]*AtA[1]
      else
        Q = rho[1]*joMatrix(AtA[1])
      end
      for i=2:length(AtA)
        if typeof(AtA[i]) <: joAbstractLinearOperator
          Q = Q + rho[i]*AtA[i]
        else
          Q = Q + rho[i]*joMatrix(AtA[i])
        end
      end
    elseif typeof(Q) == SparseMatrixCSC{TF,Int64} #all AtA are sparse arrays
      for ii=1:length(ind_updated)
          Q = Q + ( AtA[ind_updated[ii]] )*( rho[ind_updated[ii]]-log_PARSDMM.rho[i,ind_updated[ii]] );
      end
    else #all AtA are in CDS format
      for ii=1:length(ind_updated)
        CDS_scaled_add!(Q,AtA[ind_updated[ii]],Q_offsets,set_Prop.AtA_offsets[ind_updated[ii]],rho[ind_updated[ii]] .- log_PARSDMM.rho[i,ind_updated[ii]] )
      end
    end
  end

return Q
end


function bellmanford(G::DictGraph, s::Int)

end


using Distributed # addprocs, @everywhere
using BenchmarkTools # @btime
using UnicodePlots # barplot on terminal
using Statistics # mean
using Base.Threads
using PiSMC

# checks if we're running Julia in multithread mode
PiSMC.Multithread.check_julia_num_threads()

qty_workers = 8
addprocs(qty_workers)

@everywhere begin
    using Pkg
    Pkg.activate(joinpath(@__DIR__, ".."))

    using PiSMC

    # Warmup
    PiSMC.run_naive(2)
    PiSMC.run_singlethread(2)
    PiSMC.Multithread.run_multithread(2)
    PiSMC.run_distributed(2)
end

@info("Started $(length(workers())) workers")
@info("Threads available: $(nthreads())")

num_sims = 100_000_000

let
    naive_result = PiSMC.run_naive(num_sims)
    singlethread_result = PiSMC.run_singlethread(num_sims)
    multithread_result = PiSMC.Multithread.run_multithread(num_sims)
    distributed_result = PiSMC.run_distributed(num_sims)

    @info("Naive result: $naive_result")
    @info("Singlethread result: $singlethread_result")
    @info("Multithread result: $multithread_result")
    @info("Distributed result: $distributed_result")
end

@info("Benchmarking naive...")
naive_bench = @benchmark PiSMC.run_naive(num_sims)
@info("Benchmarking singlethread...")
singlethread_bench = @benchmark PiSMC.run_singlethread(num_sims)
@info("Benchmarking multithread...")
multithread_bench = @benchmark PiSMC.Multithread.run_multithread(num_sims)
@info("Benchmarking distributed...")
distributed_bench = @benchmark PiSMC.run_distributed(num_sims)

plt = barplot( [ "naive", "singlethread", "multithread ($(nthreads()) threads)", "Distributed ($qty_workers workers)"],
              [ round(mean(v.times)/1E9, digits=4) for v in [ naive_bench, singlethread_bench, multithread_bench, distributed_bench ]],
        title="Benchmark (seconds)")

show(plt)


"""
Mutable Struct: Contains the created scale factors from code to physical units
"""
mutable struct ScalesType
# exported

   # length
   Mpc::Float64
   kpc::Float64
   pc::Float64
   mpc::Float64
   ly::Float64
   Au::Float64
   km::Float64
   m::Float64
   cm::Float64
   mm::Float64
   μm::Float64

   # volume
   Mpc3::Float64
   kpc3::Float64
   pc3::Float64
   mpc3::Float64
   ly3::Float64
   Au3::Float64
   km3::Float64
   m3::Float64
   cm3::Float64
   mm3::Float64
   μm3::Float64

   # density
   Msol_pc3::Float64
   g_cm3::Float64

   # surface
   Msol_pc2::Float64
   g_cm2::Float64

   # time
   Gyr::Float64
   Myr::Float64
   yr::Float64
   s::Float64
   ms::Float64

   # mass
   Msol::Float64
   Mearth::Float64
   Mjupiter::Float64
   g::Float64

   # speed
   km_s::Float64
   m_s::Float64
   cm_s::Float64

   nH::Float64
   erg::Float64
   g_cms2::Float64

   T_mu::Float64
   Ba::Float64
   ScalesType() = new()
end

"""
Mutable Struct: Contains the physical constants in cgs units
"""
 mutable struct PhysicalUnitsType
# exported
    # in cgs units
     Au::Float64#cm: Astronomical unit
     Mpc::Float64 #cm: Parsec
     kpc::Float64 #cm: Parsec
     pc::Float64 #cm: Parsec
     mpc::Float64 #cm: MilliParsec
     ly::Float64 #cm: Light year
     Msol::Float64 #g: Solar mass
     Mearth::Float64 #g: Earth mass
     Mjupiter::Float64 #g: Jupiter mass
     Rsol::Float64 #cm: Solar radius
    # Lsol = #erg s-2: Solar luminosity
    # Mearth = #g: Earh mass

     me::Float64 #g: electron mass
     mp::Float64 #g: proton mass
     mn::Float64 #g: neutron mass
     mH::Float64 #g: hydrogen mass
     amu::Float64 #g: atomic mass unit
     NA::Float64 # Avagadro's number
     c::Float64 #cm s-1: speed of light in a vacuum
    # h = #erg s: Planck constant
    # hbar = #erg s
     G::Float64 # cm3 g-1 g-2 Gravitational constant
     kB::Float64 #erg k-1 Boltzmann constant
     Gyr::Float64 #sec: defined as 365.25 days
     Myr::Float64 #sec: defined as 365.25 days
     yr::Float64 #sec: defined as 365.25 days
     PhysicalUnitsType() = new()
end

mutable struct FileNamesType
    output::String
    info::String
    amr::String
    hydro::String
    hydro_descriptor::String
    gravity::String
    particles::String
    part_descriptor::String
    rt::String
    rt_descriptor::String
    rt_descriptor_v0::String
    clumps::String
    timer::String
    header::String
    namelist::String
    compilation::String
    makefile::String
    patchfile::String
    FileNamesType() = new()
end

"""
Mutable Struct: Contains the collected information about grid
"""
mutable struct GridInfoType
# exported
    ngridmax::Int
    nstep_coarse::Int
    nx::Int
    ny::Int
    nz::Int
    nlevelmax::Int
    nboundary::Int
    ngrid_current::Int
    bound_key::Array{Float64,1}
    cpu_read::Array{Bool,1}
    GridInfoType() = new()
end

"""
Mutable Struct: Contains the collected information about particles
"""
mutable struct PartInfoType
# exported
    eta_sn::Float64
    age_sn::Float64
    f_w::Float64
    Npart::Int
    Ndm::Int
    Nstars::Int
    Nsinks::Int
    Ncloud::Int
    Ndebris::Int
    Nother::Int
    Nundefined::Int
    other_tracer1::Int
    debris_tracer::Int
    cloud_tracer::Int
    star_tracer::Int
    other_tracer2::Int
    gas_tracer::Int
    PartInfoType() = new()
end

"""
Mutable Struct: Contains the collected information about the compilation of RAMSES
"""
mutable struct CompilationInfoType
# exported
    compile_date::String
    patch_dir::String
    remote_repo::String
    local_branch::String
    last_commit::String
    CompilationInfoType() = new()
end

"""
Mutable Struct: Contains the collected information about the descriptors
"""
mutable struct DescriptorType
# exported
    hversion::Int
    hydro::Array{Symbol,1}
    htypes::Array{String,1}
    usehydro::Bool
    hydrofile::Bool

    pversion::Int
    particles::Array{Symbol,1}
    ptypes::Array{String,1}
    useparticles::Bool
    particlesfile::Bool

    gravity::Array{Symbol,1}
    usegravity::Bool
    gravityfile::Bool

    rtversion::Int
    rt::Dict{Any,Any}
    rtPhotonGroups::Dict{Any,Any}
    usert::Bool
    rtfile::Bool

    clumps::Array{Symbol,1}
    useclumps::Bool
    clumpsfile::Bool

    sinks::Array{Symbol,1}
    usesinks::Bool
    sinksfile::Bool

    DescriptorType() = new()
end

mutable struct FilesContentType
    makefile::Array{String,1}
    timerfile::Array{String,1}
    patchfile::Array{String,1}
    FilesContentType() = new()
end

"""
Mutable Struct: Collected information about the selected simulation output
"""
mutable struct InfoType
    # exported
    output::Real
    path::String
    fnames::FileNamesType
    simcode::String
    mtime::DateTime
    ctime::DateTime
    ncpu::Int
    ndim::Int
    levelmin::Int
    levelmax::Int
    boxlen::Float64
    time::Float64
    aexp::Float64
    H0::Float64
    omega_m::Float64
    omega_l::Float64
    omega_k::Float64
    omega_b::Float64
    unit_l::Float64
    unit_d::Float64
    unit_m::Float64
    unit_v::Float64
    unit_t::Float64
    gamma::Float64

    hydro::Bool
    nvarh::Int   # number of hydro variables
    nvarp::Int   # number of particle variables
    nvarrt::Int  # number of rt variables

    variable_list::Array{Symbol,1}  # hydro variable list
    gravity_variable_list::Array{Symbol,1}
    particles_variable_list::Array{Symbol,1}
    rt_variable_list::Array{Symbol,1}
    clumps_variable_list::Array{Symbol,1}
    sinks_variable_list::Array{Symbol,1}
    descriptor::DescriptorType

    amr::Bool
    gravity::Bool
    particles::Bool
    rt::Bool
    clumps::Bool
    sinks::Bool

    namelist::Bool
    namelist_content::Dict{Any,Any}
    headerfile::Bool
    makefile::Bool
    files_content::FilesContentType
    timerfile::Bool
    compilationfile::Bool
    patchfile::Bool
    Narraysize::Int
    scale::ScalesType
    grid_info::GridInfoType
    part_info::PartInfoType
    compilation::CompilationInfoType
    constants::PhysicalUnitsType
    #overview::simulation_overview
    #cpu_overview::cpu_overview_type
    #boxcenter::Array{Float64,1}
    InfoType() = new()
end

mutable struct LevelType
  imin::Int
  imax::Int
  jmin::Int
  jmax::Int
  kmin::Int
  kmax::Int
end

"""
Abstract Supertype of all the different dataset types
> HydroPartType <: ContainMassDataSetType <: DataSetType
"""
abstract type DataSetType end # exported

"""
Abstract Supertype of all datasets that contain mass variables
> HydroPartType <: ContainMassDataSetType <: DataSetType
"""
abstract type ContainMassDataSetType <: DataSetType end # exported

"""
Abstract Supertype of data-sets that contain hydro and particle data
> HydroPartType <: ContainMassDataSetType <: DataSetType
"""
abstract type HydroPartType <: ContainMassDataSetType end # exported

"""
Mutable Struct: Contains hydro data and information about the selected simulation
> HydroDataType <: HydroPartType
"""
mutable struct HydroDataType <: HydroPartType
    # exported
    data::JuliaDB.AbstractIndexedTable
    info::InfoType
    lmin::Int
    lmax::Int
    boxlen::Float64
    ranges::Array{Float64,1}
    selected_hydrovars::Array{Int,1}
    used_descriptors::Dict{Any,Any}
    smallr::Float64
    smallc::Float64
    scale::ScalesType
    HydroDataType() = new()
end

# exported
# mutable struct GravDataType <: DataSetType
#     data
#     info::InfoType
#     lmin::Int
#     lmax::Int
#     boxlen::Float64
#     ranges::Array{Float64,1}
#     selected_gravvars::Array{Int,1}
#     scale::ScalesType
#     GravDataType() = new()
# end

"""
Mutable Struct: Contains particle data and information about the selected simulation
> PartDataType <: HydroPartType
"""
mutable struct PartDataType <: HydroPartType
#exported
    data::JuliaDB.AbstractIndexedTable
    info::InfoType
    lmin::Int
    lmax::Int
    boxlen::Float64
    ranges::Array{Float64,1}
    selected_partvars::Array{Symbol,1}
    used_descriptors::Dict{Any,Any}
    scale::ScalesType
    PartDataType() = new()
end

"""
Mutable Struct: Contains clump data and information about the selected simulation
> ClumpDataType <: ContainMassDataSetType
"""
mutable struct ClumpDataType <: ContainMassDataSetType
# exported
    data
    info::InfoType
    boxlen::Float64
    ranges::Array{Float64,1}
    selected_clumpvars::Array{Symbol,1}
    used_descriptors::Dict{Any,Any}
    scale::ScalesType
    ClumpDataType() = new()
end

"""
Union Type: Mask-array that is of type Bool or BitArray
MaskType = Union{Array{Bool,1},BitArray{1}}
"""
MaskType = Union{Array{Bool,1},BitArray{1}} # exported

MaskArrayType = Union{ Array{Array{Bool,1},1}, Array{BitArray{1},1} }
MaskArrayAbstractType = Union{ MaskArrayType, Array{AbstractArray{Bool,1},1} } # used for the combined center_of_mass function
#HydroPartType = Union{HydroDataType, PartDataType}

"""
Mutable Struct: Contains the output statistics returned by wstat
"""
mutable struct WStatType
    mean::Float64
    median::Float64
    std::Float64
    skewness::Float64
    kurtosis::Float64
    min::Float64
    max::Float64
end

"""
Abstract Supertype of all the different dataset type maps
HydroMapsType <: DataMapsType
PartMapsType <: DataMapsType
"""
abstract type DataMapsType end # exported

"""
Mutable Struct: Contains the maps/units returned by the hydro-projection information about the selected simulation
"""
mutable struct HydroMapsType <: DataMapsType
    maps::DataStructures.SortedDict{Any,Any,Base.Order.ForwardOrdering}
    maps_unit::DataStructures.SortedDict{Any,Any,Base.Order.ForwardOrdering}
    maps_lmax::DataStructures.SortedDict{Any,Any,Base.Order.ForwardOrdering}
    maps_mode::DataStructures.SortedDict{Any,Any,Base.Order.ForwardOrdering}
    lmax_projected::Real
    lmin::Int
    lmax::Int
    ranges::Array{Float64,1}
    extent::Array{Float64,1}
    cextent::Array{Float64,1}
    ratio::Float64
    boxlen::Float64
    smallr::Float64
    smallc::Float64
    scale::ScalesType
    info::InfoType
end

"""
Mutable Struct: Contains the maps/units returned by the particles-projection information about the selected simulation
"""
mutable struct PartMapsType <: DataMapsType
    maps::DataStructures.SortedDict{Any,Any,Base.Order.ForwardOrdering}
    maps_unit::DataStructures.SortedDict{Any,Any,Base.Order.ForwardOrdering}
    maps_lmax::DataStructures.SortedDict{Any,Any,Base.Order.ForwardOrdering}
    maps_mode::DataStructures.SortedDict{Any,Any,Base.Order.ForwardOrdering}
    lmax_projected::Real
    lmin::Int
    lmax::Int
    ref_time::Real
    ranges::Array{Float64,1}
    extent::Array{Float64,1}
    cextent::Array{Float64,1}
    ratio::Float64
    boxlen::Float64
    scale::ScalesType
    info::InfoType
end

"""
Mutable Struct: Contains the 2D histogram returned by the function: histogram2 and information about the selected simulation
"""
mutable struct Histogram2DMapType
    map::Array{Float64,2}
    closed::Symbol
    weight::Tuple{Symbol,Symbol}
    nbins::Array{Int,1}
    xrange::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}
    yrange::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}
    extent::Array{Float64,1}
    ratio::Float64
    scale::ScalesType
    info::InfoType
end

"""
Mutable Struct: Contains the existing simulation snapshots in a folder and a list of the empty output-folders
"""
mutable struct CheckOutputNumberType
    outputs::Array{Int,1}
    missing::Array{Int,1}
end


import REPL

using REPL.TerminalMenus
using CSV
using DataFrames, Query

export rebalance

function list_csv(path::String)
  files = [file for file in readdir(path) if isfile(abspath(joinpath(path, file)))]
  files = [m.match for m in match.(r".*\.csv", files) if m != nothing]
  return (path * "/").*(files)
end

function choose(prompt, options)
  menu = RadioMenu(options)
  choice = request(prompt, menu)
  return options[choice]
end

function request_csv(path::String, prompt::String)
  options = list_csv(path)
  return choose(prompt, options)
end

function open_fidelity_csv(path::String=homedir() * "/Downloads")
  filename = request_csv(path, "Open Fidelity CSV:")
  df = DataFrame(CSV.File(filename; footerskip=6, normalizenames=true, stringtype=String))
  return df[:,2:8]
end

function open_strategy_csv(path::String=homedir() * "/Documents")
  filename = request_csv(path, "Open Targets CSV:")
  return DataFrame(CSV.File(filename; normalizenames=true, stringtype=String))
end

function by_account(accounts::DataFrame, name::String)
  @from a in accounts begin
    @where a.Account_Name == name
    @select a
    @collect DataFrame
  end
end

function get_account(accounts::DataFrame)
  account_names = @from a in accounts begin
    @group a by a.Account_Name into n
    @select {Account_Name=key(n)}
    @collect DataFrame
  end
  account_name = choose("Choose account:", account_names[:, 1])
  return by_account(accounts, account_name)
end

function target_symbol(taxable::Bool, target)
  if taxable
    return target.Symbol_Tax
  else
    return target.Symbol_Adv
  end
end

function make_target(strategy::DataFrame)
  taxable = ("Taxable" == choose("Account Type:", ["Taxable", "Tax-Advantaged"]))
  levels = ["Conservative", "Moderate", "Balanced", "Growth/Income", "Growth", "Aggressive"]
  target = choose("Target:", levels)
  sort!(strategy, :Target, rev=true)
  # starts at 30%
  equity_percent = (findfirst(t -> t == target, levels) + 2) * 10
  equities = @from t in strategy begin
    @where !Bool(t.Income)
    @select {Symbol=target_symbol(taxable, t), t.Target}
    @collect DataFrame
  end
  equities.Target = allocate(equity_percent, equities.Target)
  income = @from t in strategy begin
    @where Bool(t.Income)
    @select {Symbol=target_symbol(taxable, t), t.Target}
    @collect DataFrame
  end
  income_percent = 100 - equity_percent
  income.Target = allocate(income_percent, income.Target)
  return sort(vcat(equities, income), :Target, rev=true)
end

function unshift(trades::DataFrame)
  return trades[1, :], trades[2:nrow(trades), :]
end

function is_core(symbol)
  return endswith(symbol, "**") || symbol == "Pending Activity"
end

function Base.get(a::String, b)
  return a
end

function sum_dtoi(c::Vector{String})
  return sum(dtoi.(c))
end

function sum_dtoi(c::Vector{Union{Missing, String}})
  result = 0
  for d in c
    result += dtoi(get(d, "\$0.00"))
  end
  return result
end

function get_core(account::DataFrame)
  core = @from asset in account begin
    @select {asset.Symbol, Current_Value=get(asset.Current_Value, asset.Last_Price_Change)}
    @collect DataFrame
  end
  return filter(asset -> is_core(get(asset.Symbol, "?")), core)
end

function get_non_core(account::DataFrame)
  nonCore = @from asset in account begin
    @select {asset.Symbol, asset.Description, asset.Quantity, asset.Current_Value}
    @collect DataFrame
  end
  return filter(asset -> !is_core(get(asset.Symbol, "?")), nonCore)
end

function get_holdings(target::DataFrame, nonCore::DataFrame)
  return @from asset in leftjoin(target, nonCore, on = :Symbol) begin
    @orderby descending(asset.Target), descending(dtoi(get(asset.Current_Value, "\$0.00")))
    @select {asset.Symbol, asset.Description, Quantity=get(asset.Quantity, 0.0), Current_Value=get(asset.Current_Value, "\$0.00"), asset.Target}
    @collect DataFrame
  end
end

function get_exiting(target::DataFrame, nonCore::DataFrame)
  exits = @from asset in nonCore begin
    @orderby descending(dtoi(get(asset.Current_Value, "\$0.00")))
    @select {asset.Symbol, asset.Description, asset.Quantity, asset.Current_Value, Target=0, Trade_Type="SELL", Trade_Value=itod(-dtoi(get(asset.Current_Value, "\$0.00"))), Drift=NaN, DriftPct=NaN}
    @collect DataFrame
  end
  return filter(e -> !(e.Symbol in target.Symbol), exits)
end

function generate_trades(account::DataFrame, target::DataFrame, deposit::Int=0)
  # core cash and pending
  core = get_core(account)
  println("\n---- HOLDINGS ----")
  if nrow(core) > 0
    printframe(core)
  end
  # assets other than core position
  nonCore = get_non_core(account)
  # account totals
  core_value = sum_dtoi(core.Current_Value)
  non_core_value = sum_dtoi(nonCore.Current_Value)
  cash = core_value + deposit
  accountTotal = core_value + non_core_value
  newTotal = non_core_value + cash
  # assets not in target (sell them)
  exiting = get_exiting(target, nonCore)
  # target assets (trade them)
  holdings = get_holdings(target, nonCore)
  # calculate drift from target
  holdings[!, :Drift] = map(h -> round(((dtoi(h.Current_Value)*100/accountTotal) - h.Target)*10, RoundNearestTiesAway)/10, eachrow(holdings))
  holdings[!, :DriftPct] = map(h -> round(h.Drift*100/h.Target*10)/10, eachrow(holdings))
  printframe(holdings)
  println("CASH | ", itod(core_value))
  # trade value in dollars
  tradeAmount = allocate(newTotal, holdings.Target[:,1]) - map(h -> dtoi(h), holdings.Current_Value[:,1])
  tradeType = map(t -> t < 0 ? "SELL" : "BUY", tradeAmount)
  holdings[!, :Trade_Type] = tradeType
  holdings[!, :Trade_Value] = itod.(tradeAmount)
  # combine all tradable holdings
  trades = vcat(holdings, exiting)
  # create exchange trades, sort
  first, second = optimize_trades(trades, cash)
  println("\n---- TRADES ----")
  settle = print_trades(first, core_value)
  if nrow(second) > 0
    settle = print_trades(second, settle)
  end
end

function optimize_trades(trades::DataFrame, cash::Int)
  # split trades into buys and sells
  sells = @from trade in trades begin
    @where trade.Trade_Type == "SELL"
    @orderby dtoi(trade.Trade_Value)
    @select {trade.Symbol, trade.Description, trade.Quantity, trade.Current_Value, trade.Trade_Type, trade.Trade_Value, Trade_For="CASH", Trade_Shares=NaN}
    @collect DataFrame
  end
  buys = @from trade in trades begin
    @where trade.Trade_Type == "BUY"
    @orderby descending(dtoi(trade.Trade_Value))
    @select {trade.Symbol, trade.Description, trade.Quantity, trade.Current_Value, trade.Trade_Type, trade.Trade_Value, Trade_For="", Trade_Shares=NaN}
    @collect DataFrame
  end
  # find mutual fund exchanges
  trades = similar(sells, 0)
  while true
    if nrow(sells) == 0
      # out of sells, append remaining buys
      append!(trades, buys)
      break
    end
    sell, sells = unshift(sells)
    if !endswith(get(sell.Symbol, "?"), "X")
      # not exchangable sell, push it
      push!(trades, sell)
      continue
    end
    if nrow(buys) == 0
      # out of buys, append remaining sells
      push!(trades, sell)
      append!(trades, sells)
      break
    end
    buy, buys = unshift(buys)
    if !endswith(get(buy.Symbol, "?"), "X")
      # not exchangable buy, push it
      push!(trades, buy)
      push!(sells, sell)
      sells = @from trade in sells begin
        @orderby dtoi(trade.Trade_Value)
        @select trade
        @collect DataFrame
      end
      continue
    end
    # exchange mutual funds
    if abs(dtoi(sell.Trade_Value)) >= dtoi(buy.Trade_Value)
      # sell larger than buy - reduce sell and turn buy into exchange
      sell.Trade_Value = itod(dtoi(sell.Trade_Value) + dtoi(buy.Trade_Value))
      buy.Trade_For = buy.Symbol
      buy.Symbol = sell.Symbol
      buy.Description = sell.Description
      buy.Quantity = sell.Quantity
      buy.Current_Value = sell.Current_Value
      buy.Trade_Type = "EXCHANGE"
      buy.Trade_Value = itod(-dtoi(buy.Trade_Value))
      push!(sells, sell)
      push!(trades, buy)
      sells = @from trade in sells begin
        @orderby dtoi(trade.Trade_Value)
        @select trade
        @collect DataFrame
      end
    else
      # buy larger than sell, reduce buy and turn sell into exchange
      buy.Trade_Value = itod(dtoi(sell.Trade_Value) + dtoi(buy.Trade_Value))
      sell.Trade_Type = "EXCHANGE"
      sell.Trade_For = buy.Symbol
      push!(buys, buy)
      push!(trades, sell)
      buys = @from trade in buys begin
        @orderby descending(dtoi(trade.Trade_Value))
        @select trade
        @collect DataFrame
      end
    end
  end
  # filter-out zero-value trades
  trades = @from trade in trades begin
    @where dtoi(trade.Trade_Value) != 0
    @orderby trade.Symbol
    @select trade
    @collect DataFrame
  end
  # calculate number of shares for sells and exchanges
  symbol = ""
  remaining_quantity = 0
  remaining_value = 0
  for trade in eachrow(trades)
    if trade.Trade_Type == "BUY"
      # skip trade-quantity for buys
      trade.Trade_Shares = NaN
      trade.Trade_For = ""
      continue
    end
    if trade.Symbol != symbol
      # reset on next symbol
      symbol = trade.Symbol
      remaining_quantity = qtoi(trade.Quantity)
      remaining_value = dtoi(trade.Current_Value)
    end
    # value for this sell/exchange, and remainder
    trade_value = abs(dtoi(trade.Trade_Value))
    remaining_value -= trade_value
    if (remaining_value == 0)
      trade_quantity = remaining_quantity
    else
      # allocate shares in thousandths
      trade_quantity, remaining_quantity = allocate(remaining_quantity, [trade_value, remaining_value])
    end
    trade.Trade_Shares = itoq(trade_quantity)
  end
  # split trades into sells-first, and buys-second
  first = @from trade in trades begin
    @where trade.Trade_Type != "BUY"
    @orderby endswith(get(trade.Symbol, "?"), "XX"), descending(trade.Trade_Type), trade.Symbol, descending(abs(dtoi(trade.Trade_Value)))
    @select trade
    @collect DataFrame
  end
  second = @from trade in trades begin
    @where trade.Trade_Type == "BUY"
    # least-expensive buys first
    @orderby dtoi(trade.Trade_Value)
    @select trade
    @collect DataFrame
  end
  # non-mutual fund sales - for funding trades
  non_mutual = @from trade in first begin
    @where !endswith(get(trade.Symbol, "?"), "X")
    @select trade
    @collect DataFrame
  end
  # amount of funding for first trades
  available = max(0, cash) - sum_dtoi(non_mutual.Trade_Value)
  # move buys from second to first, if we can afford them
  while true
    if nrow(second) == 0
      break
    end
    # least-expensive buy...
    buy, second = unshift(second)
    if dtoi(buy.Trade_Value) <= available
      available -= dtoi(buy.Trade_Value)
      push!(first, buy)
    else
      push!(second, buy)
      break
    end
  end
  # final sort
  first = @from t in first begin
    @orderby endswith(get(t.Trade_For, "?"), "XX"), endswith(get(t.Symbol, "?"), "XX"), endswith(get(t.Symbol, "?"), "X"), descending(t.Trade_Type), t.Symbol, descending(abs(dtoi(t.Trade_Value)))
    @select {t.Symbol, t.Trade_Type, t.Trade_Value, t.Trade_Shares, t.Trade_For}
    @collect DataFrame
  end
  second = @from t in second begin
    @orderby endswith(get(t.Symbol, "?"), "XX"), endswith(get(t.Symbol, "?"), "X"), descending(t.Trade_Type), t.Symbol, descending(abs(dtoi(t.Trade_Value)))
    @select {t.Symbol, t.Trade_Type, t.Trade_Value, t.Trade_Shares, t.Trade_For}
    @collect DataFrame
  end
  return first, second
end

function printframe(df::DataFrame)
  println(replace(sprint(show, df, context=:compact=>false), r".*DataFrame" => ""))
end

function print_trades(trades::DataFrame, core::Int)
  printframe(trades)
  nonex = @from trade in trades begin
    @where trade.Trade_Type != "EXCHANGE"
    @select trade
    @collect DataFrame
  end
  settle = max(0, core) - sum_dtoi(nonex.Trade_Value)
  println("CASH | ", itod(settle))
  return settle
end

function rebalance(deposit::Int=0)
  accounts = open_fidelity_csv()
  strategy = open_strategy_csv()
  println()
  account = get_account(accounts)
  target = make_target(strategy)
  generate_trades(account, target, deposit)
end


using Merlin
using HDF5

const wordembeds_file = "wordembeds_nyt100.h5"

function train()
    words = h5read(wordembeds_file, "s")
    worddict = Dict(words[i] => i for i=1:length(words))
    chardict = Dict{Char,Int}(' ' => 1)
    tagdict = Dict{String,Int}()
    #traindoc = UD_English.traindata()
    #testdoc = UD_English.testdata()
    train_w, train_c, train_t = readdata(".data/wsj_00-18.conll", worddict, chardict, tagdict)
    test_w, test_c, test_t = readdata(".data/wsj_22-24.conll", worddict, chardict, tagdict)
    info("# Training sentences:\t$(length(train_w))")
    info("# Testing sentences:\t$(length(test_w))")
    info("# Words:\t$(length(worddict))")
    info("# Chars:\t$(length(chardict))")
    info("# Tags:\t$(length(tagdict))")

    wordembeds = h5read(wordembeds_file, "v")
    T = eltype(wordembeds)
    charembeds = rand(T, 10, length(chardict))
    charembeds = charembeds * T(0.002) - T(0.001)
    nn = Model(wordembeds, charembeds, length(tagdict))
    opt = SGD(momentum=0.95, nesterov=true)

    batchsize = 8
    train_data = makebatch(batchsize, train_w, train_c, train_t)
    train_data = collect(zip(train_data...))
    test_data = makebatch(100, test_w, test_c, test_t)
    test_data = collect(zip(test_data...))

    for epoch = 1:10
        println("Epoch: $epoch")
        opt.rate = 0.001 / batchsize / epoch
        #opt.rate = 0.0075 / (1 + 0.05*(epoch-1))

        loss = 0
        for i in randperm(length(train_data))
            w, c, t = train_data[i]
            y = nn(w, c)
            y = softmax_crossentropy(t, y)
            loss += sum(y.data)
            gradient!(y)
            update!(nn.fw, opt)
            update!(nn.fc, opt)
            update!(nn.fs, opt)
        end
        loss /= batchsize * length(train_data)
        # loss = minimize!(train_f, opt, train_data)
        println("Loss: $loss")

        # test
        ys = Int[]
        zs = Int[]
        for (w,c,t) in test_data
            append!(ys, t.data)
            y = nn(w, c)
            z = vec(argmax(y.data,1))
            append!(zs, z)
        end
        length(ys) == length(zs) || throw("Length mismatch.")

        acc = mean(i -> ys[i] == zs[i] ? 1.0 : 0.0, 1:length(ys))
        acc = round(acc, 5)
        println("Test acc.: $acc")
        println()
    end
    # Merlin.save("postagger.h5", nn)
end

include("data.jl")
include("model.jl")
train()


#species_noneq = []
#species_noneq = ["H2"]
species_noneq = ["H2", "H+"]
const N_neq = length(species_noneq)
const NONEQ = N_neq > 0

const grRec = true

const file_path = "/ptmp/huchiayu/snapshots/tallbox/SFSNPI_N3e6_gS10H250dS40H250_soft0p3_from200Myr_SFcut_Z0p1_Lsh200"
const Zp = 0.1
const ShieldingLength = 0.2
include("treecol.jl")

#const Zp = parse(T, ARGS[2])

@show file_path
@show Zp

#snaps = collect(740:-10:150)
snaps = collect(550:10:1000)
for i in snaps
    abund_all, ga, X, NH, NH2, NCO, tree = solve_chem_all_particles(i, file_path, Zp);
end
0