"""
    struct JSONModel <: MetabolicModel
        json::Dict{String,Any}
    end

A struct used to store the contents of a JSON model, i.e. a model read from a
file ending with `.json`. These model files typically store all the model
parameters in arrays of JSON objects (i.e. Julia dictionaries).

Usually, not all of the fields of the input JSON can be easily represented when
converting to other models, care should be taken to avoid losing information.

Direct work on this precise model type is not very efficient, as the accessor
functions need to repeatedly find the information in the JSON tree. This gets
very slow especially if calling many accessor functions sequentially. To avoid
that, convert to e.g. [`StandardModel`](@ref) as soon as possible.

# Example
````
model = load_json_model("some_model.json")
model.json # see the actual underlying JSON
reactions(model) # see the list of reactions
````
"""
struct JSONModel <: MetabolicModel
    json::Dict{String,Any}
end

# helper access macros, see examples below for usage
macro _json_sectionkey(namesConstant::Symbol, error)
    esc(:(
        begin
            _key = _guesskey(keys(model.json), _constants.keynames.$namesConstant)
            if isnothing(_key)
                $error
            end
            _key
        end
    ))
end

macro _json_section(namesConstant::Symbol, error)
    esc(:(
        begin
            _key = @_json_sectionkey $namesConstant ($error)
            model.json[_key]
        end
    ))
end

macro _json_section_firstid(namesConstant::Symbol, id::Symbol, error)
    return esc(:(
        begin
            _s = @_json_section $namesConstant ($error)
            _s[findfirst(x -> x["id"] == $id, _s)]
        end
    ))
end

function _parse_annotations(x)::Annotations
    Annotations([k => if typeof(v) == String
        [v]
    else
        convert(Vector{String}, v)
    end for (k, v) in x])
end

_parse_notes(x)::Notes = _parse_annotations(x)

"""
    reactions(model::JSONModel)

Extract reaction names (stored as `.id`) from JSON model.
"""
function reactions(model::JSONModel)
    rs = @_json_section rxns throw(
        DomainError(keys(model.json), "JSON model has no reaction keys"),
    )

    return [string(get(rs[i], "id", "rxn$i")) for i in eachindex(rs)]
end

"""
    metabolites(model::JSONModel)

Extract metabolite names (stored as `.id`) from JSON model.
"""
function metabolites(model::JSONModel)
    ms = @_json_section mets throw(
        DomainError(keys(model.json), "JSON model has no metabolite keys"),
    )

    return [string(get(ms[i], "id", "met$i")) for i in eachindex(ms)]
end

"""
    genes(model::JSONModel)

Extract gene names from a JSON model.
"""
function genes(model::JSONModel)
    gs = @_json_section genes return []

    return [string(get(gs[i], "id", "gene$i")) for i in eachindex(gs)]
end

"""
    stoichiometry(model::JSONModel)

Get the stoichiometry. Assuming the information is stored in reaction object
under key `.metabolites`.
"""
function stoichiometry(model::JSONModel)
    rxn_ids = reactions(model)
    met_ids = metabolites(model)

    rs = @_json_section rxns throw(
        DomainError(keys(model.json), "JSON model has no reaction keys"),
    )

    S = SparseArrays.spzeros(length(met_ids), length(rxn_ids))
    for (i, rid) in enumerate(rxn_ids)
        r = rs[findfirst(x -> x["id"] == rid, rs)]
        for (met_id, coeff) in r["metabolites"]
            j = findfirst(==(met_id), met_ids)
            if isnothing(j)
                throw(
                    DomainError(
                        met_id,
                        "Unknown metabolite found in stoichiometry of $(rxn_ids[i])",
                    ),
                )
            end
            S[j, i] = coeff
        end
    end
    return S
end

"""
    bounds(model::JSONModel)

Get the bounds for reactions, assuming the information is stored in
`.lower_bound` and `.upper_bound`.
"""
function bounds(model::JSONModel)
    rs = @_json_section rxns return (
        sparse(fill(-_constants.default_reaction_bound, n_reactions(model))),
        sparse(fill(_constants.default_reaction_bound, n_reactions(model))),
    )
    return (
        sparse([get(rxn, "lower_bound", -_constants.default_reaction_bound) for rxn in rs]),
        sparse([get(rxn, "upper_bound", _constants.default_reaction_bound) for rxn in rs]),
    )
end

"""
    objective(model::JSONModel)

Collect `.objective_coefficient` keys from model reactions.
"""
function objective(model::JSONModel)
    rs = @_json_section rxns return spzeros(n_reactions(model))

    return sparse([float(get(rxn, "objective_coefficient", 0.0)) for rxn in rs])
end

"""
    reaction_gene_associaton(model::JSONModel, rid::String)

Parses the `.gene_reaction_rule` from reactions.
"""
function reaction_gene_association(model::JSONModel, rid::String)
    rxn = @_json_section_firstid rxns rid return nothing
    return _maybemap(_parse_grr, get(rxn, "gene_reaction_rule", nothing))
end

"""
    reaction_subsystem(model::JSONModel, rid::String)

Parses the `.subsystem` out from reactions.
"""
function reaction_subsystem(model::JSONModel, rid::String)
    rxn = @_json_section_firstid rxns rid return nothing
    return get(rxn, "subsystem", nothing)
end

"""
    metabolite_formula(model::JSONModel, mid::String)

Parse and return the metabolite `.formula`
"""
function metabolite_formula(model::JSONModel, mid::String)
    met = @_json_section_firstid mets mid return nothing
    return _maybemap(_parse_formula, get(met, "formula", nothing))
end

"""
    metabolite_charge(model::JSONModel, mid::String)

Return the metabolite `.charge`
"""
function metabolite_charge(model::JSONModel, mid::String)
    met = @_json_section_firstid mets mid return nothing
    return get(met, "charge", 0)
end

"""
    metabolite_compartment(model::JSONModel, mid::String)

Return the metabolite `.compartment`
"""
function metabolite_compartment(model::JSONModel, mid::String)
    met = @_json_section_firstid mets mid return nothing
    return get(met, "compartment", nothing)
end

"""
    gene_annotations(model::JSONModel, gid::String)::Annotations

Gene annotations from the [`JSONModel`](@ref).
"""
function gene_annotations(model::JSONModel, gid::String)::Annotations
    gene = @_json_section_firstid genes gid return Dict()
    _maybemap(_parse_annotations, get(gene, "annotation", nothing))
end

"""
    gene_notes(model::JSONModel, gid::String)::Notes

Gene notes from the [`JSONModel`](@ref).
"""
function gene_notes(model::JSONModel, gid::String)::Notes
    gene = @_json_section_firstid genes gid return Dict()
    _maybemap(_parse_notes, get(gene, "notes", nothing))
end

"""
    reaction_annotations(model::JSONModel, rid::String)::Annotations

Reaction annotations from the [`JSONModel`](@ref).
"""
function reaction_annotations(model::JSONModel, rid::String)::Annotations
    rxn = @_json_section_firstid rxns rid return Dict()
    _maybemap(_parse_annotations, get(rxn, "annotation", nothing))
end

"""
    reaction_notes(model::JSONModel, rid::String)::Notes

Reaction notes from the [`JSONModel`](@ref).
"""
function reaction_notes(model::JSONModel, rid::String)::Notes
    rxn = @_json_section_firstid rxns rid return Dict()
    _maybemap(_parse_notes, get(rxn, "notes", nothing))
end

"""
    metabolite_annotations(model::JSONModel, mid::String)::Annotations

Metabolite annotations from the [`JSONModel`](@ref).
"""
function metabolite_annotations(model::JSONModel, mid::String)::Annotations
    met = @_json_section_firstid mets mid return Dict()
    _maybemap(_parse_annotations, get(met, "annotation", nothing))
end

"""
    metabolite_notes(model::JSONModel, mid::String)::Notes

Metabolite notes from the [`JSONModel`](@ref).
"""
function metabolite_notes(model::JSONModel, mid::String)::Notes
    met = @_json_section_firstid mets mid return Dict()
    _maybemap(_parse_notes, get(met, "notes", nothing))
end

"""
    reaction_stoichiometry(model::JSONModel, rxn_id::String)::Dict{String, Float64}

Return the reaction equation of reaction with id `rxn_id` in model. The reaction
equation maps metabolite ids to their stoichiometric coefficients.
"""
function reaction_stoichiometry(m::JSONModel, rxn_id::String)::Dict{String, Float64}
    ind = findfirst(x -> x["id"] == rxn_id, m.json["reactions"])
    m.json["reactions"][ind]["metabolites"]
end

"""
    Base.convert(::Type{JSONModel}, mm::MetabolicModel)

Convert any [`MetabolicModel`](@ref) to [`JSONModel`](@ref).
"""
function Base.convert(::Type{JSONModel}, mm::MetabolicModel)
    if typeof(mm) == JSONModel
        return mm
    end

    rxn_ids = reactions(mm)
    met_ids = metabolites(mm)
    gene_ids = genes(mm)
    S = stoichiometry(mm)
    lbs, ubs = bounds(mm)
    ocs = objective(mm)

    json_model = JSONModel(Dict{String,Any}())
    json = Dict{String,Any}()
    json["id"] = "model" # default

    #TODO: add notes, names and similar fun stuff when they are available

    json[first(_constants.keynames.genes)] = [
        Dict([
            "id" => gid,
            "annotation" => gene_annotations(mm, gid),
            "notes" => gene_notes(mm, gid),
        ],) for gid in gene_ids
    ]

    json[first(_constants.keynames.mets)] = [
        Dict([
            "id" => mid,
            "formula" => _maybemap(_unparse_formula, metabolite_formula(mm, mid)),
            "charge" => metabolite_charge(mm, mid),
            "compartment" => metabolite_compartment(mm, mid),
            "annotation" => metabolite_annotations(mm, mid),
            "notes" => metabolite_notes(mm, mid),
        ]) for mid in met_ids
    ]

    json[first(_constants.keynames.rxns)] = [
        begin
            res = Dict{String,Any}()
            res["id"] = rid
            res["subsystem"] = reaction_subsystem(mm, rid)
            res["annotation"] = reaction_annotations(mm, rid)
            res["notes"] = reaction_notes(mm, rid)

            grr = reaction_gene_association(mm, rid)
            if !isnothing(grr)
                res["gene_reaction_rule"] = _unparse_grr(String, grr)
            end

            res["lower_bound"] = lbs[ri]
            res["upper_bound"] = ubs[ri]
            res["objective_coefficient"] = ocs[ri]
            I, V = findnz(S[:, ri])
            res["metabolites"] =
                Dict{String,Float64}([met_ids[ii] => vv for (ii, vv) in zip(I, V)])
            res
        end for (ri, rid) in enumerate(rxn_ids)
    ]

    return JSONModel(json)
end


using Test, MultiKDE, Distributions

# Hard-code test cases and expected results
hardcode_data = [[0.75762386,6.0,8.0,],
[0.40709311,2.0,8.0,],
[0.73019489,1.0,7.0,],
[0.9350646,1.0,8.0,],
[0.20622478,7.0,3.0,],
[0.82868765,4.0,9.0,],
[0.83275865,6.0,5.0,],
[0.06151981,6.0,5.0,],
[0.47064985,2.0,0.0,],
[0.81312999,5.0,4.0,],
[0.60976384,5.0,0.0,],
[0.58597128,0.0,9.0,],
[0.97580821,5.0,9.0,],
[0.21986135,2.0,4.0,],
[0.33835384,6.0,6.0,],
[0.52131305,9.0,8.0,],
[0.0651185,6.0,9.0,],
[0.75884609,9.0,3.0,],
[0.59935732,5.0,4.0,],
[0.88377118,6.0,3.0,],]
hardcode_data_pred = [[0.11249392,4.0,6.0,],
[0.28704927,5.0,3.0,],
[0.06442488,1.0,6.0,],
[0.88971842,6.0,9.0,],
[0.58269063,5.0,2.0,],
[0.04070221,3.0,4.0,],
[0.33534097,2.0,3.0,],
[0.25188701,3.0,9.0,],
[0.54533024,5.0,9.0,],
[0.51731902,2.0,2.0,],
[0.40469525,5.0,1.0,],
[0.34114321,8.0,4.0,],
[0.39484506,7.0,9.0,],
[0.45614159,6.0,3.0,],
[0.01384118,2.0,2.0,],
[0.09847895,3.0,5.0,],
[0.95260732,8.0,3.0,],
[0.92190954,2.0,1.0,],
[0.10851023,7.0,7.0,],
[0.41038077,2.0,3.0,],]

# Expected result of above input
hardcode_expected_result = [
    [4.014556252605015e-07,4.93587849306923e-07,5.017127433753062e-05,4.297840888944418,1.4103686336015966,0.0004239071556898233,0.0002450830674101923,0.001520097546290904,1.6695504130319276e-05,9.585990125997658e-05,0.0024919490954259764,0.024667731573921547,1.2143707817645687e-05,8.952814734339356e-05,3.470292679120249e-09,9.825400944075402e-06,0.00017389979047319838,0.10795426189791625,2.5029651369084925e-05,4.859412084689039,],
    [0.006731618223323871,0.034852108047552656,0.0008987511262355018,0.26494502607810944,0.002093798643830388,0.04134463379287447,0.001734435606041325,0.034237993437472795,0.34814274656951194,0.0014581092899801078,0.002014698187056216,0.002549489235012434,0.03711991361931681,0.32252782622899556,0.0013659916119672653,0.0024460021110439176,0.06131450492717148,0.0012846078414541185,0.000750247341719304,0.0017450935407583167,],
    [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,],
    [0.07012047583564995,0.16682618301528643,0.04595196932831633,0.26088694468743184,0.15452574048888829,0.09336064500373596,0.11310195315610574,0.09485512130103643,0.3361451517102332,0.10902096724118503,0.13788547837516862,0.05885101643100109,0.1329887267194227,0.3387753430435409,0.0787994864890616,0.052985136882964974,0.10988250614008524,0.07668631403187032,0.06941533271304429,0.11576170881780576,],
]
hardcode_dims = [ContinuousDim(), CategoricalDim(10), UnorderedCategoricalDim(10)]
hardcode_bws = [[0.01, 0.1, 1], [1, 0.2, 0.01], [0.02, 1, 1], [0.5, 0.5, 0.5]]

# Arguments for sanity tests
continuous_range = (0, 200)
categorical_level = 10
unordered_extra_level = 10
# zero_val_tolerance = 1e-2
s1_tol = 1
univ_sanity_trial, multi_sanity_trial = 5, 1
num_s2_test = 100
num_obs = 20
dim_types = [ContinuousDim, CategoricalDim, UnorderedCategoricalDim]
bw_s1 = 0.1
bw_s2 = 0.1
diff_delta = 2e-1
multi_test_num_dims = [1, 2, 3, 4, 5]
min_range = 1
# Objects for sanity tests
unordered_objs = Vector{Any}([Vector{Int},  Vector{Float64}, Vector{Array}, Vector{Vector}, Vector{Pair}, Vector{ErrorException}, Vector{Tuple}])
append!(unordered_objs, 1:unordered_extra_level)
unordered_level = length(unordered_objs)

function randint(n_mx)
    ceil(Int,rand()*n_mx)
end

function randint(N, n_mx)
    [ceil(Int,rand()*n_mx) for _ in 1:N]
end

function gen_batch_data(dim_type::DimensionType, n)
    dt = []
    for _ in 1:n
        _dt = gen_data(dim_type)
        if !(_dt isa Real)
            _dt = findall(x->x===_dt, unordered_objs)[1]
        end
        push!(dt, _dt)
    end
    dt
end

function gen_batch_data(dim_types::Vector, n)
    dt = []
    for _ in 1:n
        push!(dt, [gen_data(dim_types[i]) for i in 1:length(dim_types)])
    end
    dt
end

function gen_data(dim_type)
    local _dt
    if dim_type isa CategoricalDim
        _dt = rand(1:categorical_level)
    elseif dim_type isa UnorderedCategoricalDim
        _dt = unordered_objs[rand(1:length(unordered_objs))]
    elseif dim_type isa ContinuousDim
        _dt = rand(Uniform(continuous_range[1], continuous_range[2]))
    else
        error("Nonsupported dim_type")
    end
    _dt
end

# Check if the point is local maximum(gradient=0) using difference differential
function is_maximum(multi_kde::KDEMulti, x; delta=1e-10)
    x1 = [_kde.type isa ContinuousDim ? _x-delta :
                            (_kde.type isa UnorderedCategoricalDim ? _x : _x-delta)
                            for (_kde, _x) in zip(multi_kde.KDEs, x)]
    x2 = [_kde.type isa ContinuousDim ? _x+delta :
                            (_kde.type isa UnorderedCategoricalDim ? _x : _x+delta)
                            for (_kde, _x) in zip(multi_kde.KDEs, x)]
    val, val_x1, val_x2 = MultiKDE.pdf(multi_kde, x), MultiKDE.pdf(multi_kde, x1), MultiKDE.pdf(multi_kde, x2)
    (val-val_x1 >= 0) && (val-val_x2 >= 0)
end

function is_maximum(kde::KDEUniv, x; delta=1e-10)
    x1 = kde.type isa ContinuousDim ? x-delta : (kde.type isa UnorderedCategoricalDim ? x-delta : x-delta)
    x2 = kde.type isa ContinuousDim ? x+delta : (kde.type isa UnorderedCategoricalDim ? x+delta : x+delta)
    val, val_x1, val_x2 = MultiKDE.pdf(kde, x, keep_all=false), MultiKDE.pdf(kde, x1, keep_all=false), MultiKDE.pdf(kde, x2, keep_all=false)
    (val-val_x1 >= 0) && (val-val_x2 >= 0)
end

# Same as above, but calculating the gradient with direction that x points to y
function kde_gradient_to_point(multi_kde::KDEMulti, x, y; delta=1e-10)
    vec = [_kde.type isa ContinuousDim ? _x-_y :
                            (_kde.type isa UnorderedCategoricalDim ? nothing : _x-_y)
                            for (_kde, _x, _y) in zip(multi_kde.KDEs, x, y)]
    x1 = [_vec===nothing ? _x : _x-(_vec*delta) for (_x, _vec) in zip(x, vec)]
    x2 = [_vec===nothing ? _x : _x+(_vec*delta) for (_x, _vec) in zip(x, vec)]
    (MultiKDE.pdf(multi_kde, x2) - MultiKDE.pdf(multi_kde, x1)) / delta
end

function kde_gradient_to_point(kde::KDEUniv, x, y; delta=1e-10)
    vec = kde.type isa ContinuousDim ? x-y : (kde.type isa UnorderedCategoricalDim ? nothing : x-y)
    x1 = vec===nothing ? x : x-(vec*delta)
    x2 = vec===nothing ? x : x+(vec*delta)
    (MultiKDE.pdf(kde, x2, keep_all=false) - MultiKDE.pdf(kde, x1, keep_all=false)) / delta
end

@testset "KDEUniv" begin
    println("Test set for KDEUniv")
    for _ in 1:univ_sanity_trial
        _dim = dim_types[randint(length(dim_types))]
        _dim = _dim==ContinuousDim ? _dim() : (_dim==CategoricalDim ? _dim(categorical_level) : _dim(unordered_level))
        dt = gen_batch_data(_dim, num_obs)
        # Filter observations that are too close
        local dt_s1
        if _dim isa ContinuousDim
            dt_s1 = filter(e->all([abs(e-_dt)>2 for _dt in dt if _dt≠e]), dt)
        else
            dt_s1 = dt
        end
        kde_s1 = KDEUniv(_dim, bw_s1, dt_s1)
        kde_s2 = KDEUniv(_dim, bw_s2, dt)
        # 1. Sanity-check 1: When bandwidth is small enough, the gradient of PDF of observations shoule be close to zero.
        println("Performing sanity check 1...", string(_dim))
        ## Gradient function
        is_maximum_s1 = [is_maximum(kde_s2, _dt, delta=diff_delta) for _dt in kde_s1.data]
        # println("------------------------------")
        # println(sort(kde_s1.data))
        # println([MultiKDE.pdf(kde_s1, _dt, keep_all=false) for _dt in sort(kde_s1.data)])
        # println(is_maximum_s1)
        # println("------------------------------")
        # Debug plotting
        # if sum(is_maximum_s1) < (length(is_maximum_s1)-s1_tol)
        #     x = LinRange(minimum(kde_s1.data), maximum(kde_s1.data), 20000)
        #     y = [MultiKDE.pdf(kde_s1, _x, keep_all=false) for _x in x]
        #     highest = maximum(y)
        #     plot(x, y)
        #     plot!(kde_s1.data, [highest+0.05 for _ in 1:length(y)], seriestype=:scatter, size=(900, 450))
        #     println(sort(kde_s1.data))
        #     println([MultiKDE.pdf(kde_s1, _dt, keep_all=false) for _dt in sort(kde_s1.data)])
        #     println(kde_s1.kernel)
        #     println(sort(kde_s1.data))
        #     println("Plotted")
        #     savefig(string(sum(is_maximum_s1), ".png"))
        # end
        @test sum(is_maximum_s1) >= (length(is_maximum_s1)-s1_tol)
        # 2. Sanity-check 2: Using a hyperbox to include all observations, for every vector out of box that points opposite to box center, the gradient should be negative.
        println("Performing sanity check 2...")
        mn, mx = minimum(kde_s2.data), maximum(kde_s2.data)
        middle = (mx-mn)/2
        range = min(mx-mn, min_range)
        grad_s2 = []
        for i in 1:num_s2_test
            point_out = rand([mn-range, mx+range])
            _grad = kde_gradient_to_point(kde_s2, point_out, middle, delta=diff_delta)
            push!(grad_s2, _grad)
        end
        println(grad_s2)
        @test sum((grad_s2 .≈ 0) .| (grad_s2 .< 0)) == length(grad_s2)
    end
end

@testset "KDEMulti" begin
    println("Test set for KDEMulti")
    # 1. Hard-code cases check
    # result = RealVectorVector()
    result = Vector()
    for bandwidth in hardcode_bws
        multi_kde = KDEMulti(hardcode_dims, bandwidth, hardcode_data)
        _result = [MultiKDE.pdf(multi_kde, _data) for _data in hardcode_data_pred]
        push!(result, _result)
    end
    @test sum(result .≈ hardcode_expected_result) == length(hardcode_expected_result)
    # Variables needed for Sanity-check
    # Try different combinations
    for num_dim in multi_test_num_dims
        for _ in 1:multi_sanity_trial
            println("Number of dimensions = ", num_dim)
            # Prepare Random KDEMulti
            _dims = dim_types[randint(num_dim, length(dim_types))]
            _dims = [_dim==ContinuousDim ? _dim() : (_dim==CategoricalDim ? _dim(categorical_level) : _dim(unordered_level)) for _dim in _dims]
            _bw_s1 = [bw_s1 for _ in 1:num_dim]
            _bw_s2 = [bw_s2 for _ in 1:num_dim]
            dt = gen_batch_data(_dims, num_obs)
            candidates = Dict{Int, Vector}()
            for (i, __dims) in zip(1:length(_dims), _dims)
                if __dims isa UnorderedCategoricalDim
                    candidates[i] = unordered_objs
                end
            end
            local dt_s1
            if any([_dim isa ContinuousDim for _dim in _dims])
                valid_idx = []
                for (_k, _dim) in zip(1:length(_dims), _dims)
                    if _dim isa ContinuousDim
                        _dim_dt = [_dt[_k] for _dt in dt]
                        _valid_idx = filter(_i->all([abs(_dim_dt[_i]-_dim_dt[_j])>2 for _j in 1:length(_dim_dt) if _dim_dt[_i]≠_dim_dt[_j]]), 1:length(_dim_dt))
                        push!(valid_idx, _valid_idx)
                    end
                end
                valid_idx = intersect(valid_idx...)
                dt_s1 = dt[valid_idx]
            else
                dt_s1 = dt
            end
            # println("------------------------------")
            # println(_dims)
            # println(_bw_s1)
            # println(_bw_s2)
            # println(dt_s1)
            # println(candidates)
            # println("------------------------------")
            multi_kde_s1 = KDEMulti(_dims, _bw_s1, dt_s1, candidates)
            multi_kde_s2 = KDEMulti(_dims, _bw_s2, dt, candidates)
            # 2. Sanity-check 1: When bandwidth is small enough, the gradient of PDF of observations shoule be close to zero.
            println("Performing sanity check 1...")
            ## Gradient function
            is_maximum_s1 = [is_maximum(multi_kde_s1, _dt, delta=diff_delta) for _dt in dt_s1]
            println(is_maximum_s1)
            # @test sum(grad_s1 .< zero_val_tolerance) >= (length(grad_s1)-zero_num_tolerance)
            @test sum(is_maximum_s1) >= (length(is_maximum_s1))
            # 3. Sanity-check 2: Using a hyperbox to include all observations, for every vector out of box that points opposite to box center, the gradient should be negative.
            println("Performing sanity check 2...")
            mins = [minimum(kde.data) for kde in multi_kde_s2.KDEs]
            maxs = [maximum(kde.data) for kde in multi_kde_s2.KDEs]
            middle = [_dim isa ContinuousDim ? (_max-_min) / 2 :
                        (_dim isa UnorderedCategoricalDim ? (multi_kde_s2.index_to_unordered[_kde][round(Int, (_max+_min)/2)]) : round(Int, (_max+_min)/2))
                        for (i, _dim, _kde, _min, _max) in zip(1:length(multi_kde_s2.dims), multi_kde_s2.dims, multi_kde_s2.KDEs, mins, maxs)]
            ranges = [min((maxs[i]-mins[i]), min_range) for i in 1:length(mins)]
            grad_s2 = []
            for i in 1:num_s2_test
                point_out = [_dim isa ContinuousDim ? rand([mins[i]-rand(Uniform(0, ranges[i])), maxs[i]+rand(Uniform(0, ranges[i]))]) :
                                (_dim isa UnorderedCategoricalDim ? rand([multi_kde_s2.index_to_unordered[_kde][1], multi_kde_s2.index_to_unordered[_kde][length(multi_kde_s2.index_to_unordered[_kde])]]) :
                                                                            rand([mins[i]-rand(1:ranges[i]), maxs[i]+rand(1:ranges[i])]))
                                for (i, _dim, _kde) in zip(1:length(multi_kde_s2.dims), multi_kde_s2.dims, multi_kde_s2.KDEs)]
                _grad = kde_gradient_to_point(multi_kde_s2, point_out, middle, delta=diff_delta)
                push!(grad_s2, _grad)
            end
            println(grad_s2)
            @test sum((grad_s2 .≈ 0) .| (grad_s2 .< 0)) == length(grad_s2)
        end
    end
end


#*****************************************************************************
# Written by Ritchie Lee, ritchie.lee@sv.cmu.edu
# *****************************************************************************
# Copyright ã 2015, United States Government, as represented by the
# Administrator of the National Aeronautics and Space Administration. All
# rights reserved.  The Reinforcement Learning Encounter Simulator (RLES)
# platform is licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. You
# may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0. Unless required by applicable
# law or agreed to in writing, software distributed under the License is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the specific language
# governing permissions and limitations under the License.
# _____________________________________________________________________________
# Reinforcement Learning Encounter Simulator (RLES) includes the following
# third party software. The SISLES.jl package is licensed under the MIT Expat
# License: Copyright (c) 2014: Youngjun Kim.
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED
# "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
# NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
# PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
# *****************************************************************************

"""
Experiment 3
test combined time series and binary logical operations
"""
module NNGrammarExpt3

export circuit_fgandor, restarts

import Compat.ASCIIString
using TFTools
using Datasets
using TensorFlow
using TensorFlow.CoreTypes
import TensorFlow.API: l2_loss, AdamOptimizer, cast, round_, reshape_,
    reduce_max, reduce_min, pack, minimum_, maximum_, transpose_, reduce_sum,
    moments, gradients
using StatsBase
using RLESUtils, Confusion

using PyCall
@pyimport tensorflow as tf

function restarts(f::Function, N::Int64; kwargs...)
   [f(; kwargs...) for i = 1:N]
end

function circuit_fgandor(;
    dataname::AbstractString="bin_ts_synth",
    labelfield::AbstractString="F_x1_and_x3",
    learning_rate::Float64=0.001,
    max_training_epochs::Int64=500,
    target_cost::Float64=0.001,
    batch_size::Int64=5000,
    embed_hidden_units::Vector{Int64}=Int64[50, 30, 15],
    mux_hidden_units::Vector{Int64}=Int64[50,30],
    display_step::Int64=1,
    b_debug::Bool=false,
    nshow::Int64=20)

    Ds = dataset(dataname) #DFSet
    data_set = TFDataset(Ds, getmeta(Ds)[Symbol(labelfield)])

    # Construct model
    (n_examples, n_steps, n_feats) = size(Ds)
    n_featsflat = n_steps * n_feats

    # inputs
    feats = Placeholder(DT_FLOAT32, [-1, n_steps, n_feats])
    inputs = Tensor(feats)
    feats_flat = reshape_(feats, Tensor([-1, n_featsflat]))

    # common (embedding) layer
    #embed_in = feats_flat
    #embed_blk = ReluStack(embed_in, embed_hidden_units)
    #embed_out = out(embed_blk)

    # mux select input
    muxselect = feats_flat #embed layer not helping, disable for now...

    #softness parameter
    softness = collect(linspace(0.1, 1.0, max_training_epochs))
    softness_pl = Placeholder(DT_FLOAT32, [1])

    #toggle hard or soft output
    harden_pl = Placeholder(DT_BOOL, [1])

    # x mux
    x_muxin = inputs
    x_mux = SoftMux(n_feats, mux_hidden_units, x_muxin, muxselect, Tensor(harden_pl), Tensor(softness_pl))
    x_muxout = out(x_mux)

    # y mux
    y_muxin = inputs
    y_mux = SoftMux(n_feats, mux_hidden_units, y_muxin, muxselect, Tensor(harden_pl), Tensor(softness_pl))
    y_muxout = out(y_mux)

    # logical op block
    ops1_in = (x_muxout, y_muxout)
    ops1_list = [ops1_And, ops1_Or]
    ops1_blk = SoftOpsMux(ops1_in, ops1_list, mux_hidden_units, muxselect, Tensor(harden_pl), Tensor(softness_pl))
    ops1_out = out(ops1_blk)

    # temporal op block
    ops2_in = (ops1_out,)
    ops2_list = [ops2_F, ops2_G]
    ops2_blk = SoftOpsMux(ops2_in, ops2_list, mux_hidden_units, muxselect, Tensor(harden_pl), Tensor(softness_pl))
    ops2_out = out(ops2_blk)

    # outputs
    pred = ops2_out
    labels = Placeholder(DT_FLOAT32, [-1])

    #take nnout of batch, compute moments, take variance component and sum
    x_var = reduce_sum(moments(x_mux.nnout, Tensor(Int32[0]))[2], Tensor(0))
    y_var = reduce_sum(moments(y_mux.nnout, Tensor(Int32[0]))[2], Tensor(0))
    ops1_var = reduce_sum(moments(ops1_blk.softmux.nnout, Tensor(Int32[0]))[2], Tensor(0))
    ops2_var = reduce_sum(moments(ops2_blk.softmux.nnout, Tensor(Int32[0]))[2], Tensor(0))
    sum_var = x_var + y_var + constant(10.0) .* ops1_var + constant(10.0) .* ops2_var

    # Define loss and optimizer
    #cost = l2_loss(pred - labels) # Squared loss
    cost = l2_loss(pred - labels) + constant(50.0) .* sum_var

    #optimizer
    #optimizer = minimize(AdamOptimizer(learning_rate), cost) # Adam Optimizer
    opt = Optimizer(tf.train[:GradientDescentOptimizer](learning_rate))
    gvs = opt.x[:compute_gradients](cost.x)
    capped_gvs = [(tf.nn[:l2_normalize](tf.clip_by_value(grad, -1.0, 1.0), 0), var) for (grad, var) in gvs]
    optimizer = Operation(opt.x[:apply_gradients](capped_gvs))

    #compiled hardselect
    ckt = Circuit([
        x_mux, y_mux,
        ops1_blk,
        ops2_blk,
        ],
        Vector{ASCIIString}[
        recordcolnames(Ds),
        recordcolnames(Ds),
        ["and", "or"],
        ["F", "G"]])

    # Initializing the variables
    init = initialize_all_variables()

    # Rock and roll
    sess = Session()
    run(sess, init)

    # Training cycle
    for epoch in 1:max_training_epochs
        avg_cost = 0.0
        total_batch = div(num_examples(data_set), batch_size)

        # Loop over all batches
        for i in 1:total_batch
            batch_xs, batch_ys = next_batch(data_set, batch_size)
            fd = FeedDict(feats => batch_xs, labels => batch_ys, softness_pl => [softness[epoch]], harden_pl => [false])
            # Fit training using batch data
            run(sess, optimizer, fd)

            #debug
            #@show run(sess, pred - labels, fd)
            #/debug

            # Compute average loss
            batch_average_cost = run(sess, cost, fd)
            avg_cost += batch_average_cost / (total_batch * batch_size)
        end

        # Display logs per epoch step
        if epoch % display_step == 0
            fd = FeedDict(feats => data_set.X[1,:,:], labels => [data_set.Y[1]], softness_pl => [softness[epoch]], harden_pl => [false])
            softsel = softselect_by_example(sess, ckt, fd)
            x_grad = run(sess, Tensor(gradients(cost, x_mux.nnout)), fd)
            y_grad = run(sess, Tensor(gradients(cost, y_mux.nnout)), fd)
            a1_grad = run(sess, Tensor(gradients(cost, ops1_blk.softmux.nnout)), fd)
            t1_grad = run(sess, Tensor(gradients(cost, ops2_blk.softmux.nnout)), fd)
            println("x_sel=", softsel[1][1])
            println("x_grad=",x_grad[1])
            println("y_sel=", softsel[1][2])
            println("y_grad=",y_grad[1])
            println("a1_sel=", softsel[1][3])
            println("a1_grad=",a1_grad[1])
            println("t1_sel=", softsel[1][4])
            println("t1_grad=",t1_grad[1])
            println("Epoch $(epoch)  cost=$(avg_cost)")
            if avg_cost < target_cost
                break;
            end
        end
    end
    println("Optimization Finished")

    # Test model
    correct_prediction = (round_(pred) == labels)
    accuracy = mean(cast(correct_prediction, DT_FLOAT32))

    #reload data_set to recover original order
    data_set = TFDataset(Ds, getmeta(Ds)[Symbol(labelfield)])
    X = data_set.X
    Y = data_set.Y

    #soft metrics
    fd = FeedDict(feats => data_set.X, labels => data_set.Y, softness_pl => [softness[end]], harden_pl => [false])
    Ypred_soft = run(sess, round_(pred), fd)
    acc_soft = run(sess, accuracy, fd)
    stringout = simplestring(sess, ckt, fd; order=[4,1,3,2])
    top5 = topstrings(stringout, 5)
    softsel = softselect_by_example(sess, ckt, fd)
    #grads = run(sess, Tensor([f1_grad, v1_grad, a1_grad, t1_grad]), fd)
    conf = confusion(Ypred_soft.==1.0, Y.==1.0)
    conf_indices = confusion_indices(Ypred_soft.==1.0, Y.==1.0)

    #hard metrics
    fd = FeedDict(feats => data_set.X, labels => data_set.Y, softness_pl => [softness[end]], harden_pl => [true])
    Ypred_hard = run(sess, round_(pred), fd)
    acc_hard = run(sess, accuracy, fd)

    println("Soft Accuracy:", acc_soft)
    println("Hard Accuracy:", acc_hard)
    println(top5)
    d=Dict{ASCIIString,Any}()

    top5, acc_hard
end

function ops2_F(x::Tensor)
   reduce_max(x, Tensor(1))
end

function ops2_G(x::Tensor)
   reduce_min(x, Tensor(1))
end

function ops1_And(x::Tensor, y::Tensor)
    minimum_(x, y) #element-wise
end

function ops1_Or(x::Tensor, y::Tensor)
    maximum_(x, y) #element-wise
end

end #module


"A macro to make calling KNITRO's KN_* C API a little cleaner"
macro kn_ccall(func, args...)
    f = Base.Meta.quot(Symbol("KN_$(func)"))
    args = [esc(a) for a in args]
    quote
        ccall(($f, libknitro), $(args...))
    end
end

macro kn_get_values(function_name, type)
    name_singular = "KN_" * string(function_name)
    name_plural = "KN_" * string(function_name) * "s"
    fname = Symbol(name_plural)
    fnameshort = Symbol(name_singular)
    # Names of C function in knitro.h
    c_fname = Symbol(name_singular)
    c_fnames = Symbol(name_plural)
    c_fnames_all = Symbol(name_plural * "_all")

    n = if occursin("var", name_singular)
        :(KN_get_number_vars(kc))
    elseif occursin("con", name_singular)
        :(KN_get_number_cons(kc))
    end

    quote
        function $(esc(fname))(kc::Model)
            result = zeros($type, $n)
            $c_fnames_all(kc, result)
            return result
        end
        function $(esc(fname))(kc::Model, index::Vector{Cint})
            result = zeros($type, length(index))
            $c_fnames(kc, length(index), index, result)
            return result
        end
        function $(esc(fname))(kc::Model, index::Integer)
            result = zeros($type, 1)
            $c_fname(kc, index, result)
            return result[1]
        end
        $(esc(fnameshort))(kc::Model, index::Integer) = $(esc(fname))(kc, index)
    end
end

macro kn_get_attribute(function_name, type)
    fname = Symbol("KN_" * string(function_name))
    quote
        function $(esc(fname))(m::Model)
            val = zeros($type, 1)
            ret = $fname(m, val)
            return val[1]
        end
    end
end

"Format output returned by KNITRO as proper Julia string."
function _format_output(output::AbstractString)
    # remove trailing whitespace
    res = strip(output)
    # remove special characters
    return res = strip(res, '\0')
end

"Return the current KNITRO version."
function get_release()
    len = 15
    out = zeros(Cchar, len)
    KN_get_release(len, out)
    return String(strip(String(convert(Vector{UInt8}, out)), '\0'))
end

"Wrapper for KNITRO KN_context."
mutable struct Env
    ptr_env::Ptr{Cvoid}

    function Env()
        ptrptr_env = Ref{Ptr{Cvoid}}()
        res = KN_new(ptrptr_env)
        if res != 0
            error("Fail to retrieve a valid KNITRO KN_context. Error $res")
        end
        return new(ptrptr_env[])
    end

    Env(ptr::Ptr{Cvoid}) = new(ptr)
end

Base.unsafe_convert(ptr::Type{Ptr{Cvoid}}, env::Env) = env.ptr_env::Ptr{Cvoid}
is_valid(env::Env) = env.ptr_env != C_NULL

"""
Free all memory and release any Knitro license acquired by calling KN_new.
"""
function free_env(env::Env)
    if env.ptr_env != C_NULL
        ptrptr_env = Ref{Ptr{Cvoid}}(env.ptr_env)
        KN_free(ptrptr_env)
        env.ptr_env = C_NULL
    end
    return
end

"""
Structure specifying the callback context.

Each evaluation callbacks (for objective, gradient or hessian)
is attached to a unique callback context.

"""
mutable struct CallbackContext
    context::Ptr{Cvoid}
    n::Int
    m::Int
    # Add a dictionnary to store user params.
    userparams::Any

    # Oracle's callbacks are context dependent, so store
    # them inside dedicated CallbackContext.
    eval_f::Function
    eval_g::Function
    eval_h::Function
    eval_rsd::Function
    eval_jac_rsd::Function

    function CallbackContext(ptr_cb::Ptr{Cvoid})
        return new(ptr_cb, 0, 0, nothing)
    end
end

Base.unsafe_convert(ptr::Type{Ptr{Cvoid}}, cb::CallbackContext) = cb.context::Ptr{Cvoid}

mutable struct Model
    # KNITRO context environment.
    env::Env
    # Keep reference to callbacks for garbage collector.
    callbacks::Vector{CallbackContext}
    # Some structures for userParams
    puts_user::Any
    multistart_user::Any
    mip_user::Any
    newpoint_user::Any

    # Solution values.
    # Optimization status. Equal to 1 if problem is unsolved.
    status::Cint
    obj_val::Cdouble
    x::Vector{Cdouble}
    mult::Vector{Cdouble}

    # Special callbacks (undefined by default).
    # (this functions do not depend on callback environments)
    ms_process::Function
    newpt_callback::Function
    mip_callback::Function
    user_callback::Function
    ms_initpt_callback::Function
    puts_callback::Function

    # Constructor.
    function Model()
        model = new(
            Env(),
            CallbackContext[],
            nothing,
            nothing,
            nothing,
            nothing,
            1,
            Inf,
            Cdouble[],
            Cdouble[],
        )
        # Add a destructor to properly delete model.
        finalizer(KN_free, model)
        return model
    end
    # Instantiate a new Knitro instance in current environment `env`.
    function Model(env::Env)
        return new(
            env,
            CallbackContext[],
            nothing,
            nothing,
            nothing,
            nothing,
            1,
            Inf,
            Cdouble[],
            Cdouble[],
        )
    end
end

Base.unsafe_convert(ptr::Type{Ptr{Cvoid}}, kn::Model) = kn.env.ptr_env::Ptr{Cvoid}

"Free solver object."
KN_free(m::Model) = free_env(m.env)

"Create solver object."
KN_new() = Model()

is_valid(m::Model) = is_valid(m.env)
has_callbacks(m::Model) = !isempty(m.callbacks)
register_callback(model::Model, cb::CallbackContext) = push!(model.callbacks, cb)

function Base.show(io::IO, m::Model)
    if is_valid(m)
        println(io, "$(get_release())")
        println(io, "-----------------------")
        println(io, "Problem Characteristics")
        println(io, "-----------------------")
        println(io, "Objective goal:  Minimize")
        println(io, "Objective type:  $(KN_get_obj_type(m))")
        println(
            io,
            "Number of variables:                             $(KN_get_number_vars(m))",
        )
        println(
            io,
            "Number of constraints:                           $(KN_get_number_cons(m))",
        )
        println(
            io,
            "Number of nonzeros in Jacobian:                  $(KN_get_jacobian_nnz(m))",
        )
        println(
            io,
            "Number of nonzeros in Hessian:                   $(KN_get_hessian_nnz(m))",
        )

    else
        println(io, "KNITRO Problem: NULL")
    end
end

#=
    LM license manager
=#
"""
Type declaration for the Artelys License Manager context object.
Applications must not modify any part of the context.
"""
mutable struct LMcontext
    ptr_lmcontext::Ptr{Cvoid}
    # Keep a pointer to instantiated models in order to free
    # memory properly.
    linked_models::Vector{Model}

    function LMcontext()
        ptrref = Ref{Ptr{Cvoid}}()
        res = KN_checkout_license(ptrref)
        if res != 0
            error("KNITRO: Error checkout license")
        end
        lm = new(ptrref[], Model[])
        finalizer(KN_release_license, lm)
        return lm
    end
end
Base.unsafe_convert(ptr::Type{Ptr{Cvoid}}, lm::LMcontext) = lm.ptr_lmcontext::Ptr{Cvoid}

function Env(lm::LMcontext)
    ptrptr_env = Ref{Ptr{Cvoid}}()
    res = KN_new_lm(lm, ptrptr_env)
    if res != 0
        error("Fail to retrieve a valid KNITRO KN_context. Error $res")
    end
    return Env(ptrptr_env[])
end

function attach!(lm::LMcontext, model::Model)
    push!(lm.linked_models, model)
    return
end

# create Model with license manager
function Model(lm::LMcontext)
    model = Model(Env(lm))
    attach!(lm, model)
    return model
end
KN_new_lm(lm::LMcontext) = Model(lm)

function KN_release_license(lm::LMcontext)
    # First, ensure that all linked models are properly freed
    # before releasing license manager!
    KN_free.(lm.linked_models)

    if lm.ptr_lmcontext != C_NULL
        refptr = Ref{Ptr{Cvoid}}(lm.ptr_lmcontext)
        KN_release_license(refptr)
        lm.ptr_lmcontext = C_NULL
    end
end

#=
    VARIABLES
=#
function KN_add_var(m::Model)
    index = Cint[0]
    KN_add_var(m, index)
    return index[1]
end
function KN_add_vars(m::Model, nvars::Int)
    indexes = zeros(Cint, nvars)
    KN_add_vars(m, nvars, indexes)
    return indexes
end

@kn_get_values get_var_lobnd Cdouble
@kn_get_values get_var_upbnd Cdouble
@kn_get_values get_var_eqbnd Cdouble
@kn_get_values get_var_fxbnd Cdouble
@kn_get_values get_var_primal_value Cdouble
@kn_get_values get_var_dual_value Cdouble
@kn_get_values get_var_type Cint

#=
    OBJECTIVE
=#
function KN_add_obj_linear_struct(
    m::Model,
    objIndices::Vector{Cint},
    objCoefs::Vector{Cdouble},
)
    nnz = length(objIndices)
    @assert nnz == length(objCoefs)
    return KN_add_obj_linear_struct(m, nnz, objIndices, objCoefs)
end
function KN_add_obj_linear_struct(m::Model, objindex::Int, objCoefs::Cdouble)
    return KN_add_obj_linear_struct(m, Cint[objindex], [objCoefs])
end

function KN_add_obj_quadratic_struct(
    m::Model,
    indexVars1::Vector{Cint},
    indexVars2::Vector{Cint},
    coefs::Vector{Cdouble},
)
    nnz = length(indexVars1)
    @assert nnz == length(indexVars2) == length(coefs)
    return KN_add_obj_quadratic_struct(m, nnz, indexVars1, indexVars2, coefs)
end

#=
    CONSTRAINTS
=#
function KN_add_cons(m::Model, ncons::Integer)
    indexes = zeros(Cint, ncons)
    KN_add_cons(m, ncons, indexes)
    return indexes
end
function KN_add_con(m::Model)
    index = Cint[0]
    KN_add_con(m, index)
    return index[1]
end

@kn_get_values get_con_lobnd Cdouble
@kn_get_values get_con_upbnd Cdouble
@kn_get_values get_con_eqbnd Cdouble
@kn_get_values get_con_dual_value Cdouble
@kn_get_values get_con_value Cdouble

function KN_add_con_linear_struct(
    m::Model,
    index_cons::Vector{Cint},
    index_vars::Vector{Cint},
    coefs::Vector{Cdouble},
)
    @assert length(index_cons) == length(index_vars) == length(coefs)
    nnz = length(index_cons)
    return KN_add_con_linear_struct(m, nnz, index_cons, index_vars, coefs)
end
function KN_add_con_linear_struct(
    m::Model,
    index_con::Integer,
    index_vars::Vector{Cint},
    coefs::Vector{Cdouble},
)
    @assert length(index_vars) == length(coefs)
    nnz = length(index_vars)
    return KN_add_con_linear_struct_one(m, nnz, index_con, index_vars, coefs)
end
function KN_add_con_linear_struct(
    m::Model,
    index_con::Integer,
    index_var::Integer,
    coef::Cdouble,
)
    nnz = 1
    return KN_add_con_linear_struct_one(m, nnz, index_con, [index_var], [coef])
end

function KN_add_con_quadratic_struct(
    m::Model,
    index_cons::Vector{Cint},
    index_vars1::Vector{Cint},
    index_vars2::Vector{Cint},
    coefs::Vector{Cdouble},
)
    @assert length(index_cons) ==
            length(index_vars1) ==
            length(index_vars2) ==
            length(coefs)
    nnz = length(index_cons)
    return KN_add_con_quadratic_struct(m, nnz, index_cons, index_vars1, index_vars2, coefs)
end
function KN_add_con_quadratic_struct(
    m::Model,
    index_con::Integer,
    index_vars1::Vector{Cint},
    index_vars2::Vector{Cint},
    coefs::Vector{Cdouble},
)
    @assert length(index_vars1) == length(index_vars2) == length(coefs)
    nnz = length(index_vars1)
    return KN_add_con_quadratic_struct_one(
        m,
        nnz,
        index_con,
        index_vars1,
        index_vars2,
        coefs,
    )
end
function KN_add_con_quadratic_struct(
    m::Model,
    index_con::Integer,
    index_var1::Integer,
    index_var2::Integer,
    coef::Cdouble,
)
    nnz = 1
    return KN_add_con_quadratic_struct_one(
        m,
        nnz,
        index_con,
        Cint[index_var1],
        Cint[index_var2],
        [coef],
    )
end

function KN_set_compcons(
    m::Model,
    ccTypes::Vector{Cint},
    indexComps1::Vector{Cint},
    indexComps2::Vector{Cint},
)
    # get number of constraints
    nnc = length(ccTypes)
    @assert nnc == length(indexComps1) == length(indexComps2)
    return KN_set_compcons(m, nnc, ccTypes, indexComps1, indexComps2)
end

#=
    RESIDUALS
=#
function KN_add_rsds(m::Model, ncons::Integer)
    indexes = zeros(Cint, ncons)
    KN_add_rsds(m, ncons, indexes)
    return indexes
end
function KN_add_rsd(m::Model)
    index = Cint[0]
    KN_add_rsd(m, index)
    return index
end

function KN_add_rsd_linear_struct(
    m::Model,
    indexRsds::Vector{Cint},
    indexVars::Vector{Cint},
    coefs::Vector{Cdouble},
)
    nnz = length(indexRsds)
    @assert nnz == length(indexVars) == length(coefs)
    return KN_add_rsd_linear_struct(m, nnz, indexRsds, indexVars, coefs)
end
function KN_add_rsd_linear_struct(
    m::Model,
    indexRsd::Integer,
    indexVar::Vector{Cint},
    coefs::Vector{Cdouble},
)
    nnz = length(indexVar)
    @assert nnz == length(coefs)
    return KN_add_rsd_linear_struct_one(m, nnz, indexRsd, indexVar, coefs)
end

#=
    SOLVE
=#
function KN_solve(m::Model)
    # Check sanity. If model has Julia callbacks, we need to ensure
    # that Knitro is not multithreaded. Otherwise, the code will segfault
    # as we have trouble calling Julia code from multithreaded C
    # code. See issue #93 on https://github.com/jump-dev/KNITRO.jl.
    if has_callbacks(m)
        if KNITRO_VERSION >= v"13.0"
            KN_set_param(m, KN_PARAM_MS_NUMTHREADS, 1)
            KN_set_param(m, KN_PARAM_NUMTHREADS, 1)
            KN_set_param(m, KN_PARAM_MIP_NUMTHREADS, 1)
        else
            KN_set_param(m, "par_numthreads", 1)
            KN_set_param(m, "par_msnumthreads", 1)
        end
    end
    # For KN_solve, we do not return an error if ret is different of 0.
    m.status = KN_solve(m.env)
    return m.status
end

#=
    GETTERS
=#
function KN_get_solution(m::Model)
    # we first check that the model is well defined to avoid segfault
    @assert m.env != C_NULL
    nx = KN_get_number_vars(m)
    nc = KN_get_number_cons(m)

    x = zeros(Cdouble, nx)
    lambda = zeros(Cdouble, nx + nc)
    status = Cint[0]
    obj = Cdouble[0.0]

    KN_get_solution(m, status, obj, x, lambda)

    # Keep solution in cache.
    m.status = status[1]
    m.x = x
    m.mult = lambda
    m.obj_val = obj[1]
    return status[1], obj[1], x, lambda
end

# some wrapper functions for MOI
function get_status(m::Model)
    @assert m.env != C_NULL
    if m.status != 1
        return m.status
    end
    status = Cint[0]
    obj = Cdouble[0.0]
    KN_get_solution(m, status, obj, C_NULL, C_NULL)
    # Keep status in cache.
    m.status = status[1]
    return status[1]
end

function get_objective(m::Model)
    @assert m.env != C_NULL
    if isfinite(m.obj_val)
        return m.obj_val
    end
    status = Cint[0]
    obj = Cdouble[0.0]
    KN_get_solution(m, status, obj, C_NULL, C_NULL)
    # Keep objective value in cache.
    m.obj_val = obj[1]
    return obj[1]
end

function get_solution(m::Model)
    # We first check that the model is well defined to avoid segfault.
    @assert m.env != C_NULL
    if !isempty(m.x)
        return m.x
    end
    nx = KN_get_number_vars(m)
    x = zeros(Cdouble, nx)
    status = Cint[0]
    obj = Cdouble[0.0]
    KN_get_solution(m, status, obj, x, C_NULL)
    # Keep solution in cache.
    m.x = x
    return x
end
get_solution(m::Model, ix::Int) = isempty(m.x) ? get_solution(m)[ix] : m.x[ix]

function get_dual(m::Model)
    # we first check that the model is well defined to avoid segfault
    @assert m.env != C_NULL
    if !isempty(m.mult)
        return m.mult
    end
    nx = KN_get_number_vars(m)
    nc = KN_get_number_cons(m)
    lambda = zeros(Cdouble, nx + nc)
    status = Cint[0]
    obj = Cdouble[0.0]
    KN_get_solution(m, status, obj, C_NULL, lambda)
    # Keep multipliers in cache.
    m.mult = lambda
    return lambda
end
get_dual(m::Model, ix::Int) = isempty(m.mult) ? get_dual(m)[ix] : m.mult[ix]

function KN_get_objgrad_values(m::Model)
    nnz = KN_get_objgrad_nnz(m)
    indexVars = zeros(Cint, nnz)
    objGrad = zeros(Cdouble, nnz)
    KN_get_objgrad_values(m, indexVars, objGrad)
    return indexVars, objGrad
end

function KN_get_jacobian_values(m::Model)
    nnz = KN_get_jacobian_nnz(m)
    jacvars = zeros(Cint, nnz)
    jaccons = zeros(Cint, nnz)
    jaccoef = zeros(Cdouble, nnz)
    KN_get_jacobian_values(m, jacvars, jaccons, jaccoef)
    return jacvars, jaccons, jaccoef
end

function KN_get_rsd_jacobian_values(m::Model)
    nnz = KN_get_rsd_jacobian_nnz(m)
    jacvars = zeros(Cint, nnz)
    jaccons = zeros(Cint, nnz)
    jaccoef = zeros(Cdouble, nnz)
    KN_get_rsd_jacobian_values(m, jacvars, jaccons, jaccoef)
    return jacvars, jaccons, jaccoef
end

function KN_get_hessian_values(m::Model)
    nnz = KN_get_hessian_nnz(m)
    indexVars1 = zeros(Cint, nnz)
    indexVars2 = zeros(Cint, nnz)
    hess = zeros(Cdouble, nnz)
    KN_get_hessian_values(m, indexVars1, indexVars2, hess)
    return indexVars1, indexVars2, hess
end

function KN_get_var_viols(kc::Model, index::Vector{Cint})
    bndInfeas = zeros(Cint, length(index))
    intInfeas = zeros(Cint, length(index))
    viols = zeros(Cdouble, length(index))
    KN_get_var_viols(kc, length(index), index, bndInfeas, intInfeas, viols)
    return bndInfeas, intInfeas, viols
end

function KN_get_con_viols(kc::Model, index::Vector{Cint})
    infeas = zeros(Cint, length(index))
    viols = zeros(Cdouble, length(index))
    KN_get_con_viols(kc, length(index), index, infeas, viols)
    return infeas, viols
end

function KN_get_presolve_error(m::Model)
    @assert m.env != C_NULL
    component, index, error = Cint[0], Cint[0], Cint[0]
    viol = Cdouble[0.0]
    KN_get_presolve_error(m, component, index, error, viol)
    return convert(Bool, component[1]),
    convert(Int64, index[1]),
    convert(Int64, error[1]),
    convert(Float64, viol[1])
end

@kn_get_attribute get_number_vars Cint
@kn_get_attribute get_number_cons Cint
@kn_get_attribute get_obj_value Cdouble
@kn_get_attribute get_obj_type Cint

@kn_get_attribute get_number_iters Cint
@kn_get_attribute get_number_cg_iters Cint
@kn_get_attribute get_abs_feas_error Cdouble
@kn_get_attribute get_rel_feas_error Cdouble
@kn_get_attribute get_abs_opt_error Cdouble
@kn_get_attribute get_rel_opt_error Cdouble
@kn_get_attribute get_objgrad_nnz Cint
@kn_get_attribute get_jacobian_nnz KNLONG
@kn_get_attribute get_rsd_jacobian_nnz KNLONG
@kn_get_attribute get_hessian_nnz KNLONG
@kn_get_attribute get_solve_time_cpu Cdouble
@kn_get_attribute get_solve_time_real Cdouble

@kn_get_attribute get_mip_number_nodes Cint
@kn_get_attribute get_mip_number_solves Cint
@kn_get_attribute get_mip_abs_gap Cdouble
@kn_get_attribute get_mip_rel_gap Cdouble
@kn_get_attribute get_mip_incumbent_obj Cdouble
@kn_get_attribute get_mip_relaxation_bnd Cdouble
@kn_get_attribute get_mip_lastnode_obj Cdouble

#=
PARAMS
=#
function KN_set_param(m::Model, id::Integer, value::Integer)
    return KN_set_int_param(m, id, value)
end
function KN_set_param(m::Model, param::AbstractString, value::Integer)
    return KN_set_int_param_by_name(m, param, value)
end
function KN_set_param(m::Model, id::Integer, value::Cdouble)
    return KN_set_double_param(m, id, value)
end
function KN_set_param(m::Model, param::AbstractString, value::Cdouble)
    return KN_set_double_param_by_name(m, param, value)
end
function KN_set_param(m::Model, id::Integer, value::AbstractString)
    return KN_set_char_param(m, id, value)
end
function KN_set_param(m::Model, param::AbstractString, value::AbstractString)
    return KN_set_char_param_by_name(m, param, value)
end

function KN_get_int_param(m::Model, id::Integer)
    res = Cint[0]
    KN_get_int_param(m, id, res)
    return res[1]
end
function KN_get_int_param(m::Model, param::AbstractString)
    res = Cint[0]
    KN_get_int_param_by_name(m, param, res)
    return res[1]
end

function KN_get_double_param(m::Model, id::Integer)
    res = Cdouble[0.0]
    KN_get_double_param(m, id, res)
    return res[1]
end
function KN_get_double_param(m::Model, param::AbstractString)
    res = Cdouble[0.0]
    KN_get_double_param_by_name(m, param, res)
    return res[1]
end

function KN_get_param_name(m::Model, id::Integer)
    output_size = 128
    res = " "^output_size
    KN_get_param_name(m, id, res, output_size)
    return _format_output(res)
end

function KN_get_param_doc(m::Model, id::Integer)
    output_size = 128
    res = " "^output_size
    KN_get_param_doc(m, id, res, output_size)
    return _format_output(res)
end

function KN_get_param_type(m::Model, id::Integer)
    res = Cint[0]
    KN_get_param_type(m, id, res)
    return res[1]
end

function KN_get_num_param_values(m::Model, id::Integer)
    res = Cint[0]
    KN_get_num_param_values(m, id, res)
    return res[1]
end

function KN_get_param_value_doc(m::Model, id::Integer, value_id::Integer)
    output_size = 128
    res = " "^output_size
    KN_get_param_value_doc(m, id, value_id, res, output_size)
    return _format_output(res)
end

function KN_get_param_id(m::Model, name::AbstractString)
    res = Cint[0]
    KN_get_param_id(m, name, res)
    return res[1]
end

#=
    NAMES
=#
function KN_get_var_names(m::Model, max_length=1024)
    return String[
        KN_get_var_names(m, Cint(id - 1), max_length) for id in 1:KN_get_number_vars(m)
    ]
end
function KN_get_var_names(m::Model, index::Vector{Cint}, max_length=1024)
    return String[KN_get_var_names(m, id, max_length) for id in index]
end
function KN_get_var_names(m::Model, index::Cint, max_length=1024)
    rawname = zeros(Cchar, max_length)
    ret = KN_get_var_name(m, index, max_length, rawname)
    return String(strip(String(convert(Vector{UInt8}, rawname)), '\0'))
end

function KN_get_con_names(m::Model, max_length=1024)
    return String[
        KN_get_con_names(m, Cint(id - 1), max_length) for id in 1:KN_get_number_cons(m)
    ]
end
function KN_get_con_names(m::Model, index::Vector{Cint}, max_length=1024)
    return String[KN_get_con_names(m, id, max_length) for id in index]
end
function KN_get_con_names(m::Model, index::Cint, max_length=1024)
    rawname = zeros(Cchar, max_length)
    ret = KN_get_con_name(m, index, max_length, rawname)
    return String(strip(String(convert(Vector{UInt8}, rawname)), '\0'))
end


module CellMapApply

using Gridap
using Gridap.Helpers
using Gridap.CachedArrays
using Gridap.CachedValues
using Gridap.CellNumberApply: _getvalues
using Gridap.MapApply: _compute_S, _compute_M, setinputs!
using Gridap.MapApply: MapFromKernel
using Gridap.Kernels: _compute_N, _compute_T

import Gridap: evaluate
import Gridap.MapApply: _stype, _m
import Gridap: apply
import Base: iterate
import Base: length
import Base: size
import Base: getindex

function apply(k::ArrayKernel,m::CellMap,v::Vararg{<:CellValue})
  CellMapFromKernel(k,m,v...)
end

function apply(k::ArrayKernel,m::CellValue,v::CellMap)
  CellMapFromKernel(k,m,v)
end

function apply(k::ArrayKernel,m::CellMap,v::CellMap)
  CellMapFromKernel(k,m,v)
end

function apply(k::ArrayKernel,m::IndexCellMap,v::Vararg{<:IndexCellValue})
  IndexCellMapFromKernel(k,m,v...)
end

function apply(k::ArrayKernel,m::IndexCellValue,v::IndexCellMap)
  IndexCellMapFromKernel(k,m,v)
end

function apply(k::ArrayKernel,m::IndexCellMap,v::IndexCellMap)
  IndexCellMapFromKernel(k,m,v)
end

struct CellMapFromKernel{S,M,T,N,R,K,V} <: IterCellMap{S,M,T,N,R}
  kernel::K
  cellvalues::V
end

function CellMapFromKernel(k::ArrayKernel,v::Vararg{<:CellValue})
  S = _compute_S(v)
  M = _compute_M(v)
  T = _compute_T(k,v)
  N = _compute_N(k,v)
  R = _compute_R(k,v)
  K = typeof(k)
  V = typeof(v)
  CellMapFromKernel{S,M,T,N,R,K,V}(k,v)
end

function length(self::CellMapFromKernel)
  vi, = self.cellvalues
  length(vi)
end

@inline function iterate(self::CellMapFromKernel{T,N}) where {T,N}
  zipped = zip(self.cellvalues...)
  znext = iterate(zipped)
  if znext === nothing; return nothing end
  a, zstate = znext
  v = MapFromKernel(self.kernel,a...)
  _iterate(self,znext,zipped,v)
end

@inline function iterate(self::CellMapFromKernel,state)
  zipped, zstate, v = state
  znext = iterate(zipped,zstate)
  _iterate(self,znext,zipped,v)
end

@inline function _iterate(self::CellMapFromKernel,znext,zipped,v)
  if znext === nothing; return nothing end
  a, zstate = znext
  setinputs!(v,a...)
  state = (zipped,zstate,v)
  (v, state)
end

function evaluate(
  m::CellMapFromKernel{S,M,T,N},a::CellArray{<:S,M}) where {S,M,T,N}
  v = [ _eval(mi,a) for mi in m.cellvalues ]
  apply(m.kernel,v...)
end

_eval(m::CellMap,a) = evaluate(m,a)

_eval(m::CellArray,a) = m

_eval(m::CellNumber,a) = m

function _compute_R(k,v)
  m = tuple([eltype(vi) for vi in v]...)
  S = _compute_S(m)
  M = _compute_M(m)
  T = _compute_T(k,m)
  N = _compute_N(k,m)
  K = typeof(k)
  V = Tuple{m...}
  C = Tuple{[ _cache_type(mi) for mi in m]...}
  MapFromKernel{S,M,T,N,K,V,C}
end

_cache_type(::Type{<:Map{S,M,T,N}}) where {S,M,T,N} = CachedArray{T,N,Array{T,N}}

_cache_type(::Type{<:AbstractArray{T,N}}) where {T,N} = CachedArray{T,N,Array{T,N}}

_cache_type(::Type{T}) where T<:Number = Vector{T}

_stype(v::CellMap{S}) where S = S

_stype(v::CellArray) = nothing

_stype(v::CellNumber) = nothing

_m(v::CellMap{S,M}) where {S,M} = M

_m(v::CellArray) = nothing

_m(v::CellNumber) = nothing

_stype(v::Type{<:Map{S}}) where S = S

_stype(v::Type{<:AbstractArray}) = nothing

_stype(v::Type{<:Number}) = nothing

_m(v::Type{<:Map{S,M}}) where {S,M} = M

_m(v::Type{<:AbstractArray}) = nothing

_m(v::Type{<:Number}) = nothing

struct IndexCellMapFromKernel{S,M,T,N,R<:Map{S,M,T,N},K,V,F} <: IndexCellValue{R,1}
  kernel::K
  cellvalues::V
  cache::CachedValue{F}
end

function IndexCellMapFromKernel(k::ArrayKernel,v::Vararg{<:CellValue})
  S = _compute_S(v)
  M = _compute_M(v)
  T = _compute_T(k,v)
  N = _compute_N(k,v)
  R = _compute_R(k,v)
  K = typeof(k)
  V = typeof(v)
  F = Union{Nothing,R}
  cache = CachedValue{F}(nothing)
  IndexCellMapFromKernel{S,M,T,N,R,K,V,F}(k,v,cache)
end

function length(self::IndexCellMapFromKernel)
  vi, = self.cellvalues
  length(vi)
end

size(self::IndexCellMapFromKernel) = (length(self),)

function getindex(self::IndexCellMapFromKernel,i::Integer)
  vals = _getvalues(i,self.cellvalues...)
  if self.cache.value === nothing
    self.cache.value = MapFromKernel(self.kernel,vals...)
  else
  setinputs!(self.cache.value,vals...)
  end
  self.cache.value
end

function evaluate(
  m::IndexCellMapFromKernel{S,M,T,N},a::IndexCellArray{<:S,M}) where {S,M,T,N}
  v = [ _eval(mi,a) for mi in m.cellvalues ]
  apply(m.kernel,v...)
end

function evaluate(
  m::IndexCellMapFromKernel{S,M,T,N},a::CellArray{<:S,M}) where {S,M,T,N}
  v = [ _eval(mi,a) for mi in m.cellvalues ]
  apply(m.kernel,v...)
end

end # module


#=
    Helper functions
=#
"""
    get_column(model::LinQuadOptimizer, index::MOI.VariableIndex)

Return the column of the variable `index` in `model`.
"""
function get_column(model::LinQuadOptimizer, index::VarInd)
    return model.variable_mapping[index]
end

"""
    get_column(model::LinQuadOptimizer, variable::MOI.SingleVariable)

Return the column of the variable inside a `MOI.SingleVariable` function in
`model`.
"""
function get_column(model::LinQuadOptimizer, variable::SinVar)
    return get_column(model, variable.variable)
end

#=
    Get variable names
=#
MOI.supports(::LinQuadOptimizer, ::MOI.VariableName, ::Type{VarInd}) = true
function MOI.get(model::LinQuadOptimizer, ::MOI.VariableName, index::VarInd)
    return get(model.variable_to_name, index, "")
end

function MOI.set(model::LinQuadOptimizer,
                 ::MOI.VariableName, index::VarInd, name::String)
    if name != ""
        model.variable_to_name[index] = name
    else
        # delete! does nothing if no index exists.
        delete!(model.variable_to_name, index)
    end
    model.name_to_variable = nothing
    return
end

#=
    Get variable by name. Rebuild the lazy reverse-lookup if required.
=#

function _rebuild_name_to_variable(model::LinQuadOptimizer)
    model.name_to_variable = Dict{String, MOI.VariableIndex}()
    for (var, v_name) in model.variable_to_name
        v_name == "" && continue
        if haskey(model.name_to_variable, v_name)
            duplicate_var = model.name_to_variable[v_name]
            model.name_to_variable = nothing
            error("Variable $(var) and $(duplicate_var) have the same ",
                  "name: $(v_name)")
        else
            model.name_to_variable[v_name] = var
        end
    end
    return
end

function MOI.get(model::LinQuadOptimizer, ::Type{MOI.VariableIndex}, name::String)
    if model.name_to_variable === nothing
        _rebuild_name_to_variable(model)
    end
    return get(model.name_to_variable, name, nothing)
end

#=
    Get number of variables
=#

function MOI.get(model::LinQuadOptimizer, ::MOI.NumberOfVariables)
    # TODO(odow): work out why this is not the same as
    # length(model.variable_references)
    return get_number_variables(model)
end

#=
    List of Variable References
=#

function MOI.get(model::LinQuadOptimizer, ::MOI.ListOfVariableIndices)
    return model.variable_references
end

#=
    Add a single variable
=#

function MOI.add_variable(model::LinQuadOptimizer)
    add_variables!(model, 1)
    # assumes we add columns linearly
    model.last_variable_reference += 1
    index = VarInd(model.last_variable_reference)
    model.variable_mapping[index] = MOI.get(model, MOI.NumberOfVariables())
    push!(model.variable_references, index)
    push!(model.variable_primal_solution, NaN)
    push!(model.variable_dual_solution, NaN)
    model.variable_type[index] = CONTINUOUS
    return index
end

#=
    Add multiple variables
=#

function MOI.add_variables(model::LinQuadOptimizer, number_to_add::Int)
    previous_vars = MOI.get(model, MOI.NumberOfVariables())
    add_variables!(model, number_to_add)
    variable_indices = VarInd[]
    sizehint!(variable_indices, number_to_add)
    for i in 1:number_to_add
        # assumes we add columns linearly
        model.last_variable_reference += 1
        index = VarInd(model.last_variable_reference)
        push!(variable_indices, index)
        # m.last_variable_reference might not be the column if we have
        # deleted variables.
        model.variable_mapping[index] = previous_vars + i
        push!(model.variable_references, index)
        push!(model.variable_primal_solution, NaN)
        push!(model.variable_dual_solution, NaN)
        model.variable_type[index] = CONTINUOUS
    end
    return variable_indices
end

#=
    Check if reference is valid
=#

function MOI.is_valid(model::LinQuadOptimizer, index::VarInd)
    if haskey(model.variable_mapping, index)
        column = get_column(model, index)
        if 1 <= column <= MOI.get(model, MOI.NumberOfVariables())
            return true
        end
    end
    return false
end

#=
    Delete a variable
=#

function MOI.delete(model::LinQuadOptimizer, index::VarInd)
    __assert_valid__(model, index)
    column = get_column(model, index)
    delete_variables!(model, column, column)
    # delete from problem
    deleteat!(model.variable_references, column)
    delete!(model.variable_type, index)
    deleteat!(model.variable_primal_solution, column)
    deleteat!(model.variable_dual_solution, column)
    deleteref!(model.variable_mapping, column, index)
    # delete name if not ""
    if haskey(model.variable_to_name, index)
        model.name_to_variable = nothing
        delete!(model.variable_to_name, index)
    end
    # Delete any bounds (deleting from a dict without the key does nothing).
    constraint_map = cmap(model)
    delete_from_dictionary_by_value(constraint_map.upper_bound, index)
    delete_from_dictionary_by_value(constraint_map.lower_bound, index)
    delete_from_dictionary_by_value(constraint_map.fixed_bound, index)
    delete_from_dictionary_by_value(constraint_map.interval_bound, index)
    return
end

"""
    delete_from_dictionary_by_value(dict::Dict{K, V}, index::V)

Delete the entries of a dictionary where the value of a `key=>value` pair is
equal to `index`.

This is used to remove bounds when deleting a variable.
"""
function delete_from_dictionary_by_value(dict::Dict{K, V}, index::V) where {K, V}
    for (key, value) in dict
        if value == index
            delete!(dict, key)
        end
    end
end

#=
    MIP starts
=#
function MOI.supports(
        ::LinQuadOptimizer, ::MOI.VariablePrimalStart, ::Type{VarInd})
    return false
end

function MOI.set(model::LinQuadOptimizer, ::MOI.VariablePrimalStart,
                  indices::Vector{VarInd}, values::Vector{Float64})
    if !MOI.supports(model, MOI.VariablePrimalStart(), MOI.VariableIndex)
        throw(MOI.UnsupportedAttribute(MOI.VariablePrimalStart()))
    end
    add_mip_starts!(model, get_column.(Ref(model), indices), values)
end

function MOI.set(model::LinQuadOptimizer, ::MOI.VariablePrimalStart,
                  index::VarInd, value::Float64)
    MOI.set(model, MOI.VariablePrimalStart(), [index], [value])
end


# support functions

# _tablesz and hashindex are defined in Base, but are not exported,
# so they are redefined here.
_tablesz(x::Integer) = x < 16 ? 16 : one(x)<<((sizeof(x)<<3)-leading_zeros(x-1))
hashindex(key, sz) = (reinterpret(Int,(hash(key))) & (sz-1)) + 1

function not_iterator_of_pairs(kv)
    return any(x->isempty(methodswith(typeof(kv), x, true)),
               [start, next, done]) ||
           any(x->!isa(x, Union{Tuple,Pair}), kv)
end


# Machine
# =======

struct Node
    state::Int
    edges::Vector{Tuple{Edge,Node}}
end

function Node(state::Int)
    return Node(state, Tuple{Edge,Node}[])
end

function Base.show(io::IO, node::Node)
    print(io, summary(node), "(<state=$(node.state),#edges=$(length(node.edges))>)")
end

function findedge(s::Node, b::UInt8)
    for (e, t) in s.edges
        if b ∈ e.labels
            return (e, t)
        end
    end
    error("$(b) ∈ label not found")
end

struct Machine
    start::Node
    states::UnitRange{Int}
    start_state::Int
    final_states::Set{Int}
    eof_actions::Dict{Int,ActionList}
end

function Base.show(io::IO, machine::Machine)
    print(io, summary(machine), "(<states=", machine.states, ",start_state=", machine.start_state, ",final_states=", machine.final_states, ">)")
end

function compile(re::RegExp.RE; optimize::Bool=true)
    dfa = nfa2dfa(remove_dead_nodes(re2nfa(re)))
    if optimize
        dfa = remove_dead_nodes(reduce_nodes(dfa))
    end
    validate(dfa)
    return dfa2machine(dfa)
end

function dfa2machine(dfa::DFA)
    newnodes = Dict{DFANode,Node}()
    new(s) = get!(() -> Node(length(newnodes) + 1), newnodes, s)
    final_states = Set{Int}()
    eof_actions = Dict{Int,ActionList}()
    for s in traverse(dfa.start)
        s′ = new(s)
        if s.final
            push!(final_states, s′.state)
            eof_actions[s′.state] = s.eof_actions
        end
        for (e, t) in s.edges
            push!(s′.edges, (e, new(t)))
        end
    end
    start = new(dfa.start)
    @assert start.state == 1
    return Machine(start, 1:length(newnodes), 1, final_states, eof_actions)
end

function execute(machine::Machine, data::Vector{UInt8})
    s = machine.start
    cs = s.state
    actions = Symbol[]
    for d in data
        try
            e, s = findedge(s, d)
            cs = s.state
            append!(actions, action_names(e.actions))
        catch ex
            if !isa(ex, ErrorException)
                rethrow()
            end
            cs = -cs
        end
    end
    if cs ∈ machine.final_states && haskey(machine.eof_actions, s.state)
        append!(actions, action_names(machine.eof_actions[s.state]))
    end
    if cs > 0
        if cs ∈ machine.final_states
            cs = 0
        else
            cs = -cs
        end
    end
    return cs, actions
end

function execute(machine::Machine, data::String)
    return execute(machine, convert(Vector{UInt8}, codeunits(data)))
end


export LazyDiscretePost,
       ApproximatingDiscretePost

import LazySets.use_precise_ρ

"""
    LazyDiscretePost <: DiscretePost

Textbook implementation of a discrete post operator, but with lazy intersections.

### Fields

- `options` -- an `Options` structure that holds the algorithm-specific options

### Algorithm

The algorithm is based on [Flowpipe-Guard Intersection for Reachability
Computations with Support Functions](http://spaceex.imag.fr/sites/default/files/frehser_adhs2012.pdf).
"""
struct LazyDiscretePost <: DiscretePost
    options::Options

    function LazyDiscretePost(𝑂::Options)
        𝑂copy = copy(𝑂)
        # TODO: pass 𝑂 directly?
        check_aliases_and_add_default_value!(𝑂.dict, 𝑂copy.dict, [:check_invariant_intersection], false)
        check_aliases_and_add_default_value!(𝑂.dict, 𝑂copy.dict, [:overapproximation], Hyperrectangle)
        check_aliases_and_add_default_value!(𝑂.dict, 𝑂copy.dict, [:lazy_R⋂I], false)
        check_aliases_and_add_default_value!(𝑂.dict, 𝑂copy.dict, [:lazy_R⋂G], true)
        check_aliases_and_add_default_value!(𝑂.dict, 𝑂copy.dict, [:lazy_A⌜R⋂G⌟], "invertible")
        check_aliases_and_add_default_value!(𝑂.dict, 𝑂copy.dict, [:lazy_A⌜R⋂G⌟⋂I], true)
        check_aliases_and_add_default_value!(𝑂.dict, 𝑂copy.dict, [:combine_invariant_guard], 𝑂copy[:lazy_R⋂I])

        if 𝑂copy[:combine_invariant_guard] && !𝑂copy[:lazy_R⋂I]
            throw(ArgumentError("option :combine_invariant_guard only makes " *
                                "sense in combination with option :lazy_R⋂I"))
        end

        return new(𝑂copy)
    end
end

# convenience constructor from pairs of symbols
LazyDiscretePost(𝑂::Pair{Symbol,<:Any}...) = LazyDiscretePost(Options(Dict{Symbol,Any}(𝑂)))

# default options for the LazyDiscretePost discrete post operator
LazyDiscretePost() = LazyDiscretePost(Options())

"""
    ApproximatingDiscretePost()

Textbook implementation of a discrete post operator, but with lazy intersections
followed by an overapproximation. This is a particular case of the
`LazyDiscretePost`.
"""
function ApproximatingDiscretePost()
    return LazyDiscretePost(:check_invariant_intersection=>false,
                            :overapproximation=>Hyperrectangle,
                            :lazy_R⋂I=>false,
                            :lazy_R⋂G=>false,
                            :lazy_A⌜R⋂G⌟⋂I=>false)
end

function ApproximatingDiscretePost(𝑂::Options)
    𝑂_default = Options(:lazy_R⋂I=>false,
                        :lazy_R⋂G=>false,
                        :lazy_A⌜R⋂G⌟⋂I=>false)
    merge!(𝑂_default, 𝑂)
    LazyDiscretePost(𝑂_default)
end

init(𝒫::LazyDiscretePost, 𝒮::AbstractSystem, 𝑂::Options) = init!(𝒫, 𝒮, copy(𝑂))

# TODO: use 𝑂 only?
function init!(𝒫::LazyDiscretePost, 𝒮::AbstractSystem, 𝑂::Options)
    𝑂[:n] = statedim(𝒮, 1)

    # solver-specific options (adds default values for unspecified options)
    𝑂out = validate_solver_options_and_add_default_values!(𝑂)

    return 𝑂out
end

function tube⋂inv(𝒫::LazyDiscretePost,
                  reach_tube::Vector{<:AbstractReachSet{<:LazySet{N}}},
                  invariant,
                  start_interval
                 ) where {N}

    dirs = 𝒫.options[:overapproximation]

    Rsets = Vector{AbstractReachSet{<:LazySet{N}}}()
    @inbounds for reach_set in reach_tube
        R⋂I = Intersection(set(reach_set), invariant)
        if 𝒫.options[:check_invariant_intersection] && isempty(R⋂I)
            break
        end
        if !𝒫.options[:lazy_R⋂I]
            R⋂I = overapproximate(R⋂I, dirs)
        end
        push!(Rsets,
              substitute(reach_set, set=R⋂I,
                         time_start=time_start(reach_set) + start_interval[1],
                         time_end=time_end(reach_set) + start_interval[2]))
    end

    return Rsets
end

function post(𝒫::LazyDiscretePost,
              HS::HybridSystem,
              waiting_list::Vector{Tuple{Int, <:AbstractReachSet{LazySet{N}}, Int}},
              passed_list,
              source_loc_id,
              tube⋂inv,
              jumps,
              options
             ) where {N}
    jumps += 1
    oa = 𝒫.options[:overapproximation]
    source_invariant = HS.modes[source_loc_id].X
    inv_isa_Hrep, inv_isa_H_polytope = get_Hrep_info(source_invariant)

    for trans in out_transitions(HS, source_loc_id)
        info("Considering transition: $trans")
        target_loc_id = target(HS, trans)
        target_loc = HS.modes[target(HS, trans)]
        target_invariant = target_loc.X
        constrained_map = resetmap(HS, trans)
        guard = stateset(constrained_map)

        guard_isa_Hrep, guard_isa_H_polytope = get_Hrep_info(guard)
        combine_constraints = 𝒫.options[:combine_invariant_guard] &&
                              inv_isa_Hrep && guard_isa_Hrep
        if combine_constraints # combine the constraints of invariant and guard
            T = inv_isa_H_polytope || guard_isa_H_polytope ? HPolytope : HPolyhedron
            # TODO: remove redundant constraints => use intersection(..)
            invariant_guard = T([constraints_list(source_invariant);
                                 constraints_list(guard)])
        end

        # perform jumps
        post_jump = Vector{ReachSet{LazySet{N}}}()
        sizehint!(post_jump, length(tube⋂inv))
        for reach_set in tube⋂inv
            # check intersection with guard
            if combine_constraints
                R⋂G = Intersection(set(reach_set).X, invariant_guard)
            else
                R⋂G = Intersection(set(reach_set), guard)
            end
            if isempty(R⋂G)
                continue
            end
            if !𝒫.options[:lazy_R⋂G]
                R⋂G = overapproximate(R⋂G, oa)
            end

            # apply assignment
            A⌜R⋂G⌟ = apply_assignment(𝒫, constrained_map, R⋂G)
            if 𝒫.options[:lazy_A⌜R⋂G⌟] == "always" ||
                    (𝒫.options[:lazy_A⌜R⋂G⌟] == "invertible" &&
                     !isinvertible(constrained_map))
                A⌜R⋂G⌟ = overapproximate(A⌜R⋂G⌟, oa)
            end

            # intersect with target invariant
            A⌜R⋂G⌟⋂I = Intersection(target_invariant, A⌜R⋂G⌟)
            if isempty(A⌜R⋂G⌟⋂I)
                continue
            end
            if !𝒫.options[:lazy_A⌜R⋂G⌟⋂I]
                A⌜R⋂G⌟⋂I = overapproximate(A⌜R⋂G⌟⋂I, oa)
            end

            # store result
            push!(post_jump, ReachSet{LazySet{N}}(A⌜R⋂G⌟⋂I,
                                                  time_start(reach_set),
                                                  time_end(reach_set)))
        end

        postprocess(𝒫, HS, post_jump, options, waiting_list, passed_list,
            target_loc_id, jumps)
    end
end

function get_Hrep_info(set::LazySet)
    return (false, false)
end

function get_Hrep_info(set::AbstractPolytope)
    return (true, true)
end

function get_Hrep_info(set::AbstractPolyhedron)
    return (true, false)
end

# --- line-search policies ---

# usually do not use line search
function use_precise_ρ(𝒫::LazyDiscretePost,
                       cap::Intersection{N})::Bool where N<:Real
    return false
end

# use line search for the outermost level, which is a LinearMap
function use_precise_ρ(𝒫::LazyDiscretePost,
                       cap::Intersection{N, <:LinearMap{N}}
                       )::Bool where N<:Real
    return true
end


module TestOpenml

using Test
using MLJOpenML
using ScientificTypes

for M in [OpenML, MLJOpenML]
    @test length(M.list_tags()) > 200
    table = M.list_datasets(
        tag = "OpenML100",
        filter = "number_instances/100..1000/number_features/1..10")
    @test issubset([:id, :name, :status], schema(table).names)
    M.describe_dataset(61)
    table = M.load(61)
    @test :sepallength in schema(table).names
end

end
true


using Test
using Dates
using Pretend
using DataFrames

using Discord:
    Client, DiscordChannel, Message, User

using HoJBot:
    IgHolding, IgPortfolio, IgUserError, PrettyView, SimpleView,
    current_date, discord_channel, discord_reply, discord_upload_file,
    recent_date_range, retrieve, retrieve_users, upload_file,
    ig_affirm_non_player, ig_affirm_player,
    ig_buy, ig_cash_entry, ig_chart, ig_count_shares, ig_execute,
    ig_file_path, ig_get_real_time_quote, ig_grouped_holdings, ig_hey,
    ig_historical_prices, ig_holdings_data_frame, ig_is_player,
    ig_load_all_portfolios, ig_load_portfolio, ig_mark_to_market!,
    ig_mark_to_market_portfolio, ig_ranking_table,
    ig_reformat_view!, ig_remove_game, ig_save_portfolio,
    ig_sell, ig_start_new_game, ig_value_all_portfolios, ig_view_table

Pretend.activate()

const USER_ID = UInt64(0)
const USER_ID2 = UInt64(1)
const USER_ID3 = UInt64(2)
const USER_NAME = "Joe"

function test_ig_cases(f, name)
    @testset "$name" begin
        try
            @test_nowarn ig_start_new_game(USER_ID)
            @test_nowarn ig_start_new_game(USER_ID2)
            f()
        finally
            @test_nowarn ig_remove_game(USER_ID)
            @test_nowarn ig_remove_game(USER_ID2)
        end
    end
end

function clean_ig_game_files()
    for id in (USER_ID, USER_ID2, USER_ID3)
        path = ig_file_path(id)
        isfile(path) && rm(path)
    end
end

# Mock functions
mock_ig_get_quote_100(symbol::AbstractString) = 100.0
mock_currrent_date() = Date(2021, 1, 11)
mock_reply(c::Client, m::Message, s::AbstractString) = nothing
mock_upload_file(c::Client, ch::DiscordChannel, filename::AbstractString; kwargs...) = nothing
mock_channel(c::Client, id::UInt64) = DiscordChannel(; id = 0x00, type = 0x00)

function mock_retrieve_users(c::Client, ids::Vector{UInt64})
    return Dict(id => User(; id=id, username="$id") for id in ids)
end

# Mock objects
mocked_client() = Client("hey")

@testset "Investment Game" begin

    # ensure everything is clean before we start
    clean_ig_game_files()

    test_ig_cases("Life cycle") do
        @test ig_is_player(USER_ID) == true
        @test ig_hey(USER_NAME, "wat") isa AbstractString
        @test_nowarn ig_affirm_player(USER_ID)
        @test_throws IgUserError ig_affirm_non_player(USER_ID)
    end

    test_ig_cases("Load/save") do
        # load
        pf = @test_nowarn ig_load_portfolio(USER_ID)
        @test pf.cash > 0
        @test isempty(pf.holdings)

        # save with holdings
        push!(pf.holdings, IgHolding("AAPL", 1, today(), 100))
        @test_nowarn ig_save_portfolio(USER_ID, pf)
        pf = @test_nowarn ig_load_portfolio(USER_ID)
        @test length(pf.holdings) == 1

        @test_nowarn ig_save_portfolio(USER_ID2, pf)
        pfs = ig_load_all_portfolios()
        @test length(pfs) >= 2   # two test pf's, plus whatever is there for others
        @test haskey(pfs, USER_ID)
        @test haskey(pfs, USER_ID2)
    end

    test_ig_cases("Buy/sell") do
        apply(ig_get_real_time_quote => mock_ig_get_quote_100) do
            # Buy something, portfolio should be updated
            executed_price = @test_nowarn ig_buy(USER_ID, "AAPL", 50)
            @test executed_price == 100.0

            # Load it again and check
            pf = ig_load_portfolio(USER_ID)
            @test length(pf.holdings) == 1
            @test pf.holdings[1].symbol == "AAPL"
            @test pf.holdings[1].shares == 50
            @test pf.cash == 1_000_000 - 50 * 100.0

            # Sell more than what you have
            @test_throws IgUserError ig_sell(USER_ID, "AAPL", 70)

            # Partial sell
            executed_price = @test_nowarn ig_sell(USER_ID, "AAPL", 30)
            @test executed_price == 100.0

            # What's left?
            pf = ig_load_portfolio(USER_ID)
            @test length(pf.holdings) == 1
            @test pf.holdings[1].symbol == "AAPL"
            @test pf.holdings[1].shares == 20

            # Sell completely
            ig_sell(USER_ID, "AAPL", 20)
            pf = ig_load_portfolio(USER_ID)
            @test length(pf.holdings) == 0
            @test pf.cash == 1_000_000

            # buy 3 lots relief
            ig_buy(USER_ID, "IBM", 50)
            ig_buy(USER_ID, "IBM", 50)
            ig_buy(USER_ID, "IBM", 50)
            ig_sell(USER_ID, "IBM", 80)
            pf = ig_load_portfolio(USER_ID)
            @test pf.cash == 1_000_000 - 70 * 100.0
            @test length(pf.holdings) == 2
            @test sort([h.shares for h in pf.holdings]) == [20, 50]
        end
    end

    test_ig_cases("Views") do
        # empty portfolio with pre-filled cash only
        df = @test_nowarn ig_mark_to_market_portfolio(USER_ID)
        @test nrow(df) == 1
        @test df.symbol[1] == "CASH:USD"

        # setup
        ig_buy(USER_ID, "AAPL", 100, 120)
        ig_buy(USER_ID, "IBM", 50, 130)
        ig_buy(USER_ID, "IBM", 50, 150)
        pf = ig_load_portfolio(USER_ID)

        # counting
        @test ig_count_shares(pf, "AAPL") == 100
        @test ig_count_shares(pf, "IBM") == 100

        # convert to holdings data frame
        df = @test_nowarn ig_holdings_data_frame(pf)
        @test df isa AbstractDataFrame
        @test sort(propertynames(df)) == sort([:symbol, :shares, :purchase_price, :purchase_date])

        # grouped with average pricing
        gdf = ig_grouped_holdings(df)
        @test hasproperty(gdf, :symbol) == true
        @test hasproperty(gdf, :shares) == true
        @test hasproperty(gdf, :purchase_price) == true
        @test gdf[gdf.symbol .== "IBM", :].purchase_price[1] == 140

        # mark to market
        df2 = copy(df)
        apply(ig_get_real_time_quote => mock_ig_get_quote_100) do
            @test_nowarn ig_mark_to_market!(df2)
            @test hasproperty(df2, :current_price)
            @test hasproperty(df2, :market_value)
            @test unique(df2.current_price) == [100]

            df3 = @test_nowarn ig_mark_to_market_portfolio(USER_ID)
            @test df3.current_price .* df3.shares == df3.market_value

            # compatible column names
            @test Set(propertynames(ig_cash_entry(pf))) == Set(propertynames(df3))

            # translation to string
            @test_nowarn ig_view_table(PrettyView(), df2)

            # SimpleView requires specific format
            @test_nowarn ig_reformat_view!(df3)
            @test_nowarn ig_view_table(SimpleView(), df3)
        end
    end

    test_ig_cases("Ranking") do
        ig_buy(USER_ID, "AAPL", 100, 60)
        ig_buy(USER_ID, "IBM", 100, 60)
        ig_buy(USER_ID2, "AAPL", 100, 120)  # bought at higher price i.e. less cash remaining
        ig_buy(USER_ID2, "IBM", 100, 120)
        apply(
            ig_get_real_time_quote => mock_ig_get_quote_100,
            retrieve_users => mock_retrieve_users,
        ) do
            valuations = @test_nowarn ig_value_all_portfolios()
            # ranking order: highest valuation is at the top
            @test valuations[1].id == USER_ID
            @test valuations[2].id == USER_ID2

            client = mocked_client()
            rt = @test_nowarn ig_ranking_table(client)
            @test rt[1, :player] == "$USER_ID"
            @test rt[2, :player] == "$USER_ID2"
        end
    end

    test_ig_cases("Executors") do
        c = Client("test")
        m = Message(; id = 0x00, channel_id = 0x00)
        u = User(; id = USER_ID3, username = USER_NAME)
        apply(
            discord_channel => mock_channel,
            discord_reply => mock_reply,
            discord_upload_file => mock_upload_file,
            ig_get_real_time_quote => mock_ig_get_quote_100,
            retrieve_users => mock_retrieve_users,
        ) do
            @test_nowarn ig_execute(c, m, u, Val(Symbol("start-game")), [])
            @test_nowarn ig_execute(c, m, u, Val(Symbol("abandon-game")), [])
            @test_nowarn ig_execute(c, m, u, Val(Symbol("abandon-game-really")), [])

            @test_nowarn ig_execute(c, m, u, Val(Symbol("start-game")), [])
            @test_nowarn ig_execute(c, m, u, Val(:view), [])
            @test_nowarn ig_execute(c, m, u, Val(:view), ["simple"])
            @test_nowarn ig_execute(c, m, u, Val(:buy), ["100", "aapl"])
            @test_nowarn ig_execute(c, m, u, Val(:sell), ["100", "aapl"])
            @test_nowarn ig_execute(c, m, u, Val(Symbol("abandon-game-really")), [])

            @test_nowarn ig_execute(c, m, u, Val(:quote), ["aapl"])
            @test_nowarn ig_execute(c, m, u, Val(:chart), ["aapl"])
            @test_nowarn ig_execute(c, m, u, Val(:chart), ["aapl", "100d"])

            @test_nowarn ig_execute(c, m, u, Val(:rank), [])
            @test_nowarn ig_execute(c, m, u, Val(:rank), ["10"])
        end
    end

    @testset "Misc" begin
        apply(current_date => mock_currrent_date) do
            r = @test_nowarn recent_date_range(Day(10))
            @test r isa Tuple{Date,Date}
            @test r[1] == Date(2021, 1, 1)
            @test r[2] == Date(2021, 1, 11)
        end
    end

    @testset "Charting" begin
        dates = collect(Date(2021, 1, 1):Day(1):Date(2021,12,31))
        values = collect(1:365)
        @test_nowarn ig_chart("AAPL", dates, values)
    end

    @testset "Integration tests" begin
        from_date, to_date = Date(2021, 1, 1), Date(2021, 1, 31)
        @test_nowarn ig_historical_prices("AAPL", from_date, to_date)
        @test_throws IgUserError ig_historical_prices("BADSYMBOL", from_date, to_date)
    end

    clean_ig_game_files()
end


# An abstract type for fixed-point acceleration methods
# Fixed point problem x = g(x) with residual f(x) = x - g(x)
export AbstractAccelerator
"""
    AbstractAccelerator

Abstract supertype for acceleration objects that can be used to speed up a fixed-point iterations g = g(x) of a nonexpansive operator `g`.
They must implement the following methods to communicate with the fixed-point algorithm:
  - update!(aa::AbstractAccelerator, g, x, num_iter) #stores the fixed-point iterates
  - accelerate!(g::AbstractVector, x, aa::AbstractAccelerator, num_iter) #recombines past iterates to determine an accelerated point and overwrites `g`
  - restart!(aa::AbstractAccelerator, args...; kwargs...) # algorithm tells the accelerator to restart

The algorithm has to be able to query the following information:
  - was_successful(aa::AbstractAccelerator) --> Bool #indicate whether accelerate! was succesful at the last iteration

The following method is optional:
  - log!(aa::AbstractAccelerator, args...; kwargs...) # algorithm tells accelerator to log certain information for debugging
"""
abstract type AbstractAccelerator end

# update!(aa::AbstractAccelerator, g, x, iter) = nothing
# accelerate!(g, x, aa::AbstractAccelerator, iter) = nothing
# was_successful(aa::AbstractAccelerator) = false
# restart!(aa::AbstractAccelerator) = nothing

log!(aa::AbstractAccelerator, args...; kwargs...) = false


module Models

import StatsBase: fit, predict

export Model, Template
export fit, predict, submodels, estimate_type, output_type, predict_input_type
export EstimateTrait, PointEstimate, DistributionEstimate
export OutputTrait, SingleOutput, MultiOutput
export PredictInputTrait, PointPredictInput, PointOrDistributionPredictInput

"""
   Template

A Template is an untrained [`Model`](@ref) that can be [`fit`](@ref) to data.
Defined as well are the traits:
- [`output_type`](@ref): [`SingleOutput`](@ref) or [`MultiOutput`](@ref).
- [`estimate_type`](@ref): [`PointEstimate`](@ref) or [`DistributionEstimate`](@ref).
"""
abstract type Template end

"""
   Model

A Model is a trained [`Template`](@ref) with which one can [`predict`](@ref) on inputs.
Defined as well are the traits:
- [`output_type`](@ref): [`SingleOutput`](@ref) or [`MultiOutput`](@ref).
- [`estimate_type`](@ref): [`PointEstimate`](@ref) or [`DistributionEstimate`](@ref).
"""
abstract type Model end

"""
   fit(::Template, output::AbstractMatrix, input::AbstractMatrix, [weights]) -> Model

Fit the [`Template`](@ref) to the `output` and `input` data and return a trained
[`Model`](@ref).
Convention is that `weights` defaults to `StatsBase.uweights(Float32, size(outputs, 2))`
"""
function fit end

"""
    predict(model::Model, inputs::AbstractMatrix)
    predict(model::Model, inputs::AbstractVector{<:AbstractVector})

Predict targets for the provided the collection of `inputs` and [`Model`](@ref).

A [`Model`](@ref) subtype for which the `predict_input_type(model)` is
[`PointPredictInput`](@ref) will only need to implement a `predict` function that operates
on an `AbstractMatrix` of inputs.

If the `estimate_type(model)` is [`PointEstimate`](@ref) then this function should return
another `AbstractMatrix` in which each column contains the prediction for a single input.

If the `estimate_type(model)` is [`DistributionEstimate`](@ref) then this function should
return a `AbstractVector{<:Distribution}`.
"""
function predict end

function predict(model::Model, inputs::AbstractVector{<:AbstractVector})
   return predict(model, reduce(hcat, inputs))
end

"""
    submodels(::Union{Template, Model})

Return all submodels within a multistage model/template.

Submodels are models within a model that have their own inputs (which may or may not be
combined with outputs of _earlier_ submodels, before actually being passed as input to the submodel).
Such multistage models take a tuple of inputs (which may be nested if the submodel itself
has submodels).
The order of submodels returned by `submodels` is as per the order of the inputs in the
tuple.

For single-stage models, (i.e. ones that simply take a matrix as input), this returns an
empty tuple.
Wrapper models which do not expose their inner models to seperate inputs, including ones
that only wrap a single model, should **not** define `submodels` as they are
(from the outside API perspective) single-stage models.
"""
submodels(::Union{Template, Model}) = ()

include("traits.jl")
include("test_utils.jl")

end # module


_productions = Vector{Vector{Int64}}()
_width = nothing
_height = nothing

function serializeMoveSet(moves::Vector{Move})
  returnString = ""
  for move in moves
    returnString *= string(move.loc.x) * " " * string(move.loc.y) * " " * string(move.direction) * " "
  end
  returnString
end

function deserializeMapSize(inputString::String)
  splitString = reverse(split(inputString, " "))

  global _width, _height
  _width = parse(pop!(splitString))
  _height = parse(pop!(splitString))
end

function deserializeProductions(inputString::String)
  splitString = reverse(split(inputString, " "))

  for a in 1:_height
    row = Vector{Int64}()
    for b in 1:_width
      push!(row, parse(pop!(splitString)))
    end
    push!(_productions, row)
  end
end

function deserializeMap(inputString::String)
  splitString = reverse(split(inputString, " "))

  m = GameMap(_width, _height)

  y = 1
  x = 1
  counter = 0
  owner = 0
  while ~(y > m.height)
    counter = parse(pop!(splitString))
    owner = parse(pop!(splitString))
    for a in 0:counter-1
      m.contents[y][x].owner = owner
      x += 1
      if x > m.width
        x = 1
        y += 1
      end
    end
  end

  for a in 1:_height
    for b in 1:_width
      m.contents[a][b].strength = parse(pop!(splitString))
      m.contents[a][b].production = _productions[a][b]
    end
  end

  m
end

function sendString(toBeSent::String)
  toBeSent *= "\n"

  write(STDOUT, toBeSent)
  flush(STDOUT)
end

function getString()
  rstrip(readline(), '\n')
end

function getInit()
  playerTag = parse(getString())
  deserializeMapSize(getString())
  deserializeProductions(getString())
  m = deserializeMap(getString())

  (playerTag, m)
end

function sendInit(name::String)
  sendString(name)
end

function getFrame()
  return deserializeMap(getString())
end

function sendFrame(moves::Vector{Move})
  sendString(serializeMoveSet(moves))
end