module TestLearningNetworkMachines

using Test
using ..Models
using ..TestUtilities
using MLJBase
using Tables
using StableRNGs
using Serialization
rng = StableRNG(616161)

# A dummy clustering model:
mutable struct DummyClusterer <: Unsupervised
    n::Int
end
DummyClusterer(; n=3) = DummyClusterer(n)
function MLJBase.fit(model::DummyClusterer, verbosity::Int, X)
    Xmatrix = Tables.matrix(X)
    n = min(size(Xmatrix, 2), model.n)
    centres = Xmatrix[1:n, :]
    levels = categorical(1:n)
    report = (centres=centres,)
    fitresult = levels
    return fitresult, nothing, report
end
MLJBase.transform(model::DummyClusterer, fitresult, Xnew) =
    selectcols(Xnew, 1:length(fitresult))
MLJBase.predict(model::DummyClusterer, fitresult, Xnew) =
    [fill(fitresult[1], nrows(Xnew))...]

N = 20
X = (a = rand(N), b = categorical(rand("FM", N)))

@testset "signature helpers" begin
    @test MLJBase._call(NamedTuple()) == NamedTuple()
    a = source(:a)
    b = source(:b)
    W = source(:W)
    yhat = source(:yhat)
    s = (transform=W,
         report=(a=a, b=b),
         predict=yhat)
    @test MLJBase._report_part(s) == (a=a, b=b)
    @test MLJBase._operation_part(s) == (transform=W, predict=yhat)
    @test MLJBase._nodes(s) == (W, yhat, a, b)
    @test MLJBase._operations(s) == (:transform, :predict)
    R = MLJBase._call(MLJBase._report_part(s))
    @test R.a == :a
    @test R.b == :b
end

@testset "wrapping a learning network in a machine" begin

    # unsupervised:
    Xs = source(X)
    W = transform(machine(OneHotEncoder(), Xs), Xs)
    clust = DummyClusterer(n=2)
    m = machine(clust, W)
    yhat = predict(m, W)
    Wout = transform(m, W)
    rnode = source(:stuff)

    # test of `fitted_params(::NamedTuple)':
    fit!(Wout, verbosity=0)

    @test_throws(MLJBase.ERR_BAD_SIGNATURE,
                 machine(Unsupervised(),
                         predict=yhat,
                         fitted_params=rnode))
    @test_throws(MLJBase.ERR_EXPECTED_NODE_IN_SIGNATURE,
                 machine(Unsupervised(),
                         predict=42))
    @test_throws(MLJBase.ERR_EXPECTED_NODE_IN_SIGNATURE,
                 machine(Unsupervised(), Xs;
                         predict=yhat,
                         transform=Wout,
                         report=(some_stuff=42,)))
    mach = machine(Unsupervised(), Xs;
                   predict=yhat,
                   transform=Wout,
                   report=(some_stuff=rnode,))
    @test mach.args == (Xs, )
    @test mach.args[1] == Xs
    fit!(mach, force=true, verbosity=0)
    Θ = mach.fitresult
    @test Θ.predict == yhat
    @test Θ.transform == Wout
    Θ.report.some_stuff == rnode
    @test report(mach).some_stuff == :stuff

    @test report(mach).machines == fitted_params(mach).machines

    # supervised
    y = rand("ab", N) |> categorical;
    ys = source(y)
    mm = machine(ConstantClassifier(), W, ys)
    yhat = predict(mm, W)
    e = @node auc(yhat, ys)

    @test_throws Exception machine(predict=yhat)
    mach = machine(Probabilistic(), Xs, ys;
                   predict=yhat,
                   report=(training_auc=e,))
    @test mach.model isa Probabilistic
    @test_throws ArgumentError machine(Probabilistic(), Xs, ys)
    @test_throws ArgumentError machine(Probabilistic(), Xs, ys;
                                       report=(training_auc=e,))

    # test extra report items coming from `training_auc=e` above
    fit!(mach, verbosity=0)
    err = auc(yhat(), y)
    @test report(mach).training_auc ≈ err

    # supervised - predict_mode
    @test predict_mode(mach, X) == mode.(predict(mach, X))
    predict_mode(mach, rows=1:2) == predict_mode(mach, rows=:)[1:2]

    # evaluate a learning machine
    evaluate!(mach, measure=LogLoss(), verbosity=0)

    # supervised - predict_median, predict_mean
    X1, y1 = make_regression(20)

    Xs = source(X1); ys = source(y1)
    mm = machine(ConstantRegressor(), Xs, ys)
    yhat = predict(mm, Xs)
    mach = fit!(machine(Probabilistic(), Xs, ys; predict=yhat), verbosity=0)
    @test predict_mean(mach, X1) ≈ mean.(predict(mach, X1))
    @test predict_median(mach, X1) ≈ median.(predict(mach, X1))

end

# build a learning network:
x1 = map(n -> mod(n,3), rand(rng, UInt8, 100)) |> categorical;
x2 = randn(rng, 100);
X = (x1=x1, x2=x2);
y = x2.^2;

Xs = source(X)
ys = source(y)
z = log(ys)
stand = UnivariateStandardizer()
standM = machine(stand, z)
u = transform(standM, z)
hot = OneHotEncoder()
hotM = machine(hot, Xs)
W = transform(hotM, Xs)
knn = KNNRegressor()
knnM = machine(knn, W, u)
oak = DecisionTreeRegressor()
oakM = machine(oak, W, u)
uhat = 0.5*(predict(knnM, W) + predict(oakM, W))
zhat = inverse_transform(standM, uhat)
yhat = exp(zhat)
enode = @node mae(ys, yhat)

@testset "replace method for learning network machines, acceleration: $(typeof(accel))" for accel in (CPU1(), CPUThreads())

    fit!(yhat, verbosity=0, acceleration=accel)

    # test nested reporting:
    r = MLJBase.report(yhat)
    d = r.report_given_machine
    ms = machines(yhat)
    @test ms == r.machines |> reverse
    @test all(mach -> report(mach) == d[mach], ms)

    hot2 = deepcopy(hot)
    knn2 = deepcopy(knn)

    # duplicate a learning network machine:
    mach  = machine(Deterministic(), Xs, ys;
                    predict=yhat,
                    report=(mae=enode,))
    mach2 = replace(mach, hot=>hot2, knn=>knn2,
                    ys=>source(ys.data);
                    empty_unspecified_sources=true)
    ss = sources(glb(mach2))
    @test isempty(ss[1])
    mach2 = @test_logs((:warn, r"No replacement"),
                       replace(mach, hot=>hot2, knn=>knn2,
                               ys=>source(ys.data)))
    yhat2 = mach2.fitresult.predict
    fit!(mach, verbosity=0)
    fit!(mach2, verbosity=0)
    @test predict(mach, X) ≈ predict(mach2, X)
    @test report(mach).mae ≈ report(mach2).mae

    @test mach2.args[1]() == Xs()
    @test mach2.args[2]() == ys()

    ## EXTRA TESTS FOR TRAINING SEQUENCE

    # pickout the newly created machines:
    standM2 = machines(yhat2, stand) |> first
    oakM2 = machines(yhat2, oak) |> first
    knnM2 = machines(yhat2, knn) |> first
    hotM2 = machines(yhat2, hot) |> first

    @test_mach_sequence(fit!(yhat2, force=true, acceleration=accel),
                        [(:train, standM2), (:train, hotM2),
                         (:train, knnM2), (:train, oakM2)],
                        [(:train, hotM2), (:train, standM2),
                         (:train, knnM2), (:train, oakM2)],
                        [(:train, standM2), (:train, hotM2),
                         (:train, oakM2), (:train, knnM2)],
                        [(:train, hotM2), (:train, standM2),
                         (:train, oakM2), (:train, knnM2)])

    @test length(MLJBase.machines(yhat)) == length(MLJBase.machines(yhat2))
    @test MLJBase.models(yhat) == MLJBase.models(yhat2)
    @test sources(yhat) == sources(yhat2)
    @test MLJBase.tree(yhat) == MLJBase.tree(yhat2)
    @test yhat() ≈ yhat2()

    # this change should trigger retraining of all machines except the
    # univariate standardizer:
    hot2.drop_last = true
    @test_mach_sequence(fit!(yhat2, acceleration=accel),
                        [(:skip, standM2), (:update, hotM2),
                         (:train, knnM2), (:train, oakM2)],
                        [(:update, hotM2), (:skip, standM2),
                         (:train, knnM2), (:train, oakM2)],
                        [(:skip, standM2), (:update, hotM2),
                         (:train, oakM2), (:train, knnM2)],
                        [(:update, hotM2), (:skip, standM2),
                         (:train, oakM2), (:train, knnM2)])
end

@testset "Test serializable of pipeline" begin
    filename = "pipe_mach.jls"
    X, y = make_regression(100, 1)
    pipe = (X -> coerce(X, :x₁=>Continuous)) |> DecisionTreeRegressor()
    mach = machine(pipe, X, y)
    fit!(mach, verbosity=0)

    # Check serializable function
    smach = MLJBase.serializable(mach)
    @test_logs((:warn, MLJBase.warn_serializable_mach(predict),
                (:warn, MLJBase.warn_serializable_mach(predict)),
                fit!(mach, verbosity=0)))
    TestUtilities.generic_tests(mach, smach)
    @test keys(fitted_params(smach)) == keys(fitted_params(mach))
    @test keys(report(smach)) == keys(report(mach))
    # Check data has been wiped out from models at the first level of composition
    @test length(machines(glb(smach))) == length(machines(glb(mach)))
    for submach in machines(glb(smach))
        TestUtilities.test_data(submach)
    end

    # End to end
    MLJBase.save(filename, mach)
    smach = machine(filename)
    @test predict(smach, X) == predict(mach, X)

    rm(filename)
end

@testset "Test serializable of composite machines" begin
    filename = "stack_mach.jls"
    X, y = make_regression(100, 1)
    model = Stack(
        metalearner = DecisionTreeRegressor(),
        tree1 = DecisionTreeRegressor(min_samples_split=3),
        tree2 = DecisionTreeRegressor(),
        measures=rmse)
    mach = machine(model, X, y)
    fit!(mach, verbosity=0)

    # Check serializable function
    smach = MLJBase.serializable(mach)
    TestUtilities.generic_tests(mach, smach)
    # Check data has been wiped out from models at the first level of composition
    @test length(machines(glb(smach))) == length(machines(glb(mach)))
    for submach in machines(glb(smach))
        TestUtilities.test_data(submach)
    end

    # Testing extra report field : it is a deepcopy
    @test smach.report.cv_report === mach.report.cv_report

    @test smach.fitresult isa MLJBase.CompositeFitresult

    Serialization.serialize(filename, smach)
    smach = Serialization.deserialize(filename)
    MLJBase.restore!(smach)

    @test MLJBase.predict(smach, X) == MLJBase.predict(mach, X)
    @test keys(fitted_params(smach)) == keys(fitted_params(mach))
    @test keys(report(smach)) == keys(report(mach))

    rm(filename)

    # End to end
    MLJBase.save(filename, mach)
    smach = machine(filename)
    @test predict(smach, X) == predict(mach, X)

    rm(filename)
end

@testset "Test serializable of nested composite machines" begin
    filename = "nested_stack_mach.jls"
    X, y = make_regression(100, 1)

    pipe = (X -> coerce(X, :x₁=>Continuous)) |> DecisionTreeRegressor()
    model = Stack(
        metalearner = DecisionTreeRegressor(),
        pipe = pipe)
    mach = machine(model, X, y)
    fit!(mach, verbosity=0)

    MLJBase.save(filename, mach)
    smach = machine(filename)

    @test predict(smach, X) == predict(mach, X)

    # Test data as been erased at the first and second level of composition
    for submach in machines(glb(smach))
        TestUtilities.test_data(submach)
        if submach isa Machine{<:Composite}
            for subsubmach in machines(glb(submach))
                TestUtilities.test_data(subsubmach)
            end
        end
    end

    rm(filename)

end

struct DummyRangeCV
    cv
end

torange(x::UnitRange) = x
torange(x) = minimum(x):maximum(x)

function MLJBase.train_test_pairs(dcv::DummyRangeCV, rows, X, y)
    ttp = MLJBase.train_test_pairs(dcv.cv, rows)
    return [(torange(t),torange(e)) for (t,e) in ttp]
end

@testset "Test serialized filesize does not increase with datasize" begin
    # At the moment it is necessary to resort to a custom resampling strategy,
    # for this test. This is because partial functions working on nodes,
    # such as `selectrows`in learning networks store data.
    # A basic CV would store vectors which would grow in size as the dataset grows.

    dcv = DummyRangeCV(CV(nfolds=3))
    model = Stack(
        metalearner = FooBarRegressor(lambda=1.),
        resampling = dcv,
        model_1 = DeterministicConstantRegressor(),
        model_2=ConstantRegressor())

    filesizes = []
    for n in [100, 500, 1000]
        filename = "serialized_temp_$n.jls"
        X, y = make_regression(n, 1)
        mach = machine(model, X, y)
        fit!(mach, verbosity=0)
        MLJBase.save(filename, mach)
        push!(filesizes, filesize(filename))
        rm(filename)
    end
    @test all(x==filesizes[1] for x in filesizes)
    # What if no serializable procedure had happened
    filename = "full_of_data.jls"
    X, y = make_regression(1000, 1)
    mach = machine(model, X, y)
    fit!(mach, verbosity=0)
    serialize(filename, mach)
    @test filesize(filename) > filesizes[1]

    @test_logs (:warn, MLJBase.warn_bad_deserialization(mach.state)) machine(filename)

    rm(filename)
end

mutable struct DummyComposite <: DeterministicComposite
    stand1
    stand2
end

@testset "issue 377" begin
    stand = Standardizer()
    model = DummyComposite(stand, stand)

    Xs = source()
    mach1 = machine(model.stand1, Xs)
    X1 = transform(mach1, Xs)
    mach2 = machine(model.stand2, X1)
    X2 = transform(mach2, X1)

    mach = machine(Unsupervised(), Xs; transform=X2)
    @test_logs((:error, r"The hyper"),
               @test_throws(ArgumentError,
                            MLJBase.network_model_names(model, mach)))
end

end

true


#=
Fundamental result of the entire algorithm:
for each point in the masked region of the target image,
the coordinates (in the corpus frame)
of the best matching patch.

When the algorithm is "with replacement"
the target image is also mutated as the algorithm proceeds.

It is not the same thing to execute without replacement,
and only then mutate (set the color of) the target image
with the color from the best match in the corpus.
Since replacing as you go along gives an ever changing comparand for searches.

But you could also say that the mutated target is a result.
Thats the most used result, say for the "retouching" application.

After the algorithm,
1) hasValue is all true
2) and mapFromTargetToCorpusPoints is valid CartesianIndex in the synth region
and uninitialized elsewhere.
=#

#=
The arrays are the same shape as the target.

The arrays are ordinary.
The mapFromTargetToCorpusPoints is sparse:
most values not initialized, only values in the synth region ever get values.
But we don't use an implementation that uses less storage when sparse.

The frame of the arrays is the same as the frame of the target image.

TODO use a subarray with the same frame, but less storage???
i.e. an array that encloses the synthesized region of the target
=#

# Depends on image.jl because methods take instance of MaskedImage
include("../image.jl")  # MaskedImage

# OLD array of CartesianIndex
# MapElementType =
# MapElementType = Int64

struct SynthResult{DimensionCount}
    #=
    Both the array, and its elements, have same DimensionCount.
    !!! The elements of the array are CartesianIndex.
    The eltype() is (CartesianIndex{DimensionCount})

    NOT AbstractArray.
    Resynthesizer chooses an implementation.
    =#

    # WAS
    mapFromTargetToCorpusPoints::Array{CartesianIndex{DimensionCount}, DimensionCount}

    # NEW array of linear index i.e. Int64
    # mapFromTargetToCorpusPoints::Array{MapElementType, DimensionCount}

    hasValue::BitArray{DimensionCount}
end

#=
Return the initial SynthResult

No points in the synth region have been synthesized.
The coordinates of the best match for points in the synth region is unknown.

In original C, prepare_target_sources()
=#
function initialSynthResult(
        targetImage::MaskedImage{ValueType, DimensionCount}
        ) where {ValueType, DimensionCount}
    @debug "SynthResult initializer called"

    #=
    Uninitialized array of CartesianIndex.
    Whose elements have same concrete type as the indices of targetImage
    e.g. CartesianIndex{N}

    TODO Really the indices of the corpusImage.
    For more generality, allow targetImage and corpusImage to have differing dimensions.
    =#
    # elementType = CartesianIndex{ndims(targetImage.image)}
    elementType = CartesianIndex{DimensionCount}

    # TODO this is doing one allocation for each element of the image
    # Why isn't it just one huge allocation?
    mapFromTargetToCorpusPoints = similar(targetImage.image, elementType) # {DimensionCount})

    # TODO understand what an undef is
    # mapFromTargetToCorpusPoints = Array{CartesianIndex{DimensionCount}}(undef, size(targetImage.image))
    #mapFromTargetToCorpusPoints = Array{MapElementType}(undef, size(targetImage.image))

    # boolean array false only in the synth region (the selected region give by the mask)
    hasValue = trues(size(targetImage.image))
    hasValue[targetImage.mask] .= false
    @debug "Initial hasValue map" hasValue

    result = SynthResult(mapFromTargetToCorpusPoints, hasValue)
    @debug "typeof SynthResult" typeof(result)
    # Useless  to assert isconcretetype(result) since all structs are concrete
    return result
end

#=
Setter method of the SynthResult object.
=#
function setBetterSearchResult(
        synthResult::SynthResult{DimensionCount},
        targetIndex::CartesianIndex{DimensionCount},
        corpusIndex::CartesianIndex{DimensionCount}
        )::Nothing where {DimensionCount}
    synthResult.hasValue[targetIndex] = true
    synthResult.mapFromTargetToCorpusPoints[targetIndex] = corpusIndex
    return nothing
end

#=
Does the point have a value?
I.E. has been synthesized already?
Getter method of the SynthResult object.

All points in target outside the synth region have a value
(from initial image, never mutated.)
Otherwise only some points in the synth region of the target have a value
(have been synthesized i.e. mutated as search progresses and improves.)
This is only pertinent for the first pass, since the algorithm starts
by nullify or blacken the values of point in the synth region.
After the first pass, they have a value, albeit a crude value.
=#
function doesPointHaveValue(synthResult, patchPoint )::Bool
    # assert patchPoint and synthResult in same frame, the target
    return synthResult.hasValue[patchPoint]
end

#=
Is the targetPoint in the synth region (not the context)
AND has been synthesized (has a best match in the corpus.)

In original C code: has_source_neighbor()

!!! This is not the original implementation.
The original relies on sourceOf field of Neighbor,
i.e. on caching SynthResult.mapFromTargetToCorpusPoints[targetIndex] in the Neighbor.
Which probably gives better cpu cache coherence.
Since the implementation below:
1) accesses the original mask
2) accesses an inverted copy of the mask (SynthResult.hasValue)
both of which are large and both might not fit in the cpu cache.
=#
function isInSynthAndWasSynthesized(
        synthResult,
        targetMaskedImage, patchPoint   # framed point ::PointInMaskedImage)
        )::Bool
        # assert patchPoint is in target frame

    #=
    Order of the terms is important for performance:
    hasValue is  general, for any target point, and is usually true when target is much larger than synth region.
    =#
    return ( isPointSelectedInMaskedImage(targetMaskedImage, patchPoint)  # is in synth region
            && doesPointHaveValue(synthResult, patchPoint)
            )
end


"""
    ArcEager()

Arc-Eager transition system for dependency parsing.

# Transitions

| Transition  | Definition                                    |
|:----------- |:--------------------------------------------- |
| LeftArc(l)  | (σ\\|s, b\\|β, A) → (σ, b\\|β, A ∪ (b, l, s)) |
| RightArc(l) | (σ\\|s, b\\|β, A) → (σ, b\\|β, A ∪ (b, l, s)) |
| Reduce      | (σ\\|s, β,  A) → (σ, β,   A)                  |
| Shift	      | (σ,  b\\|β, A) → (σ\\|b, β, A)                |

# Preconditions

| Transition  | Condition                        |
|:----------- |:-------------------------------- |
| LeftArc(l)  | ¬[s = 0], ¬∃k∃l'[(k, l', i) ϵ A] |
| RightArc(l) | ¬∃k∃l'[(k, l', j) ϵ A]           |
| Reduce      | ∃k∃l[(k, l, i) ϵ A]              |

# References

[Nivre 2003](http://stp.lingfil.uu.se/~nivre/docs/iwpt03.pdf), [Nivre 2008](https://www.aclweb.org/anthology/J08-4003.pdf).
"""
struct ArcEager <: AbstractTransitionSystem end

initconfig(::ArcEager, graph::DependencyTree) =
    ArcEagerConfig(graph)
initconfig(::ArcEager, deptype, words) =
    ArcEagerConfig{deptype}(words)

transition_space(::ArcEager, labels=[]) =
    isempty(labels) ? [LeftArc(), RightArc(), Reduce(), Shift()] :
    [LeftArc.(labels)..., RightArc.(labels)..., Reduce(), Shift()]

projective_only(::ArcEager) = true

struct ArcEagerConfig{T} <: AbstractParserConfiguration{T}
    c::StackBufferConfiguration{T}
end

@stackbufconfig ArcEagerConfig

leftarc(cfg::ArcEagerConfig, args...; kwargs...) =
    ArcEagerConfig(leftarc_reduce(cfg.c, args...; kwargs...))

rightarc(cfg::ArcEagerConfig, args...; kwargs...) =
    ArcEagerConfig(rightarc_shift(cfg.c, args...; kwargs...))

reduce(cfg::ArcEagerConfig) = ArcEagerConfig(reduce(cfg.c))

shift(cfg::ArcEagerConfig) = ArcEagerConfig(shift(cfg.c))

isfinal(cfg::ArcEagerConfig) = all(a -> head(a) >= 0, cfg.c.A)
hashead(cfg::ArcEagerConfig, k) = head(token(cfg, k)) != -1

"""
    static_oracle(cfg, tree, transition)

Default static oracle function for arc-eager dependency parsing.

See [Goldberg & Nivre 2012](https://www.aclweb.org/anthology/C12-1059.pdf).
(Also called Arc-Eager-Reduce in [Qi & Manning 2017](https://nlp.stanford.edu/pubs/qi2017arcswift.pdf)).
"""
function static_oracle(cfg::ArcEagerConfig, gold_tree, arc)
    l = i -> arc(token(gold_tree, i))
    gold_arc = (a, b) -> has_arc(gold_tree, a, b)
    if stacklength(cfg) >= 1
        s = last(stack(cfg))
        if bufferlength(cfg) >= 1
            b = first(buffer(cfg))
            if gold_arc(b, s)
                return LeftArc(l(s)...)
            elseif gold_arc(s, b)
                return RightArc(l(b)...)
            end
        end
        if all(k -> k > 0 && hashead(cfg, k), [s ; dependents(gold_tree, s)])
            return Reduce()
        end
    end
    return Shift()
end

"""
    static_oracle_prefer_shift(cfg, gold_tree, arc=untyped)

Static oracle for arc-eager dependency parsing. Similar to the
"regular" static oracle, but always Shift when ambiguity is present.

See [Qi & Manning 2017](https://nlp.stanford.edu/pubs/qi2017arcswift.pdf).
"""
function static_oracle_prefer_shift(cfg::ArcEagerConfig, gold_tree, arc=untyped)
    l = i -> arc(token(gold_tree, i))
    gold_arc = (a, b) -> has_arc(gold_tree, a, b)
    (σ, s), (b, β) = popstack(cfg), shiftbuffer(cfg)
    if gold_arc(b, s)
        return LeftArc(l(s)...)
    elseif gold_arc(s, b)
        return RightArc(l(b)...)
    end
    must_reduce = false
    for k in stack(cfg)
        if gold_arc(k, b) || gold_arc(b, k)
            must_reduce = true
            break
        elseif head(token(cfg, k)) < 0
            break
        end
    end
    has_right_children = any(k -> s in rightdeps(gold_tree, k), buffer(cfg))
    if !must_reduce || s > 0 && !hashead(cfg, s) || has_right_children
        return Shift()
    else
        return Reduce()
    end
end

"""
    dynamic_oracle(t, cfg::ArgEagerConfig, tree)

Dynamic oracle function for arc-eager parsing.

For details, see [Goldberg & Nivre 2012](https://aclweb.org/anthology/C12-1059).
"""
dynamic_oracle(t, cfg::ArcEagerConfig, tree) = cost(t, cfg, tree) == 0

# see figure 2 in goldberg & nivre 2012 "a dynamic oracle..."
possible_transitions(cfg::ArcEagerConfig, graph::DependencyTree, arc=untyped) =
    possible_transitions(cfg, arc)

function possible_transitions(cfg::ArcEagerConfig, arc=untyped)
    s, b = last(stack(cfg)), first(buffer(cfg))
    l = i -> arc(token(cfg, i))
    transitions = [LeftArc(l(s)...), RightArc(l(b)...), Reduce(), Shift()]
    return filter(t -> is_possible(t, cfg), transitions)
end

function cost(t::LeftArc, cfg::ArcEagerConfig, gold)
    # left arc cost: num of arcs (k,l',s), (s,l',k) s.t. k ϵ β
    σ, s = popstack(cfg)
    b, β = shiftbuffer(cfg)
    if has_arc(gold, b, s)
        0
    else
        count(k -> has_arc(gold, k, s) || has_arc(gold, s, k), β)
    end
end

function cost(t::RightArc, cfg::ArcEagerConfig, gold)
    # right arc cost: num of gold arcs (k,l',b), s.t. k ϵ σ or k ϵ β,
    #                 plus num of gold arcs (b,l',k) s.t. k ϵ σ
    σ, s = popstack(cfg)
    b, β = shiftbuffer(cfg)
    if has_dependency(gold, s, b)
        0
    else
        count(k -> has_arc(gold, k, b), [σ ; β]) + count(k -> has_arc(gold, b, k), σ)
    end
end

function cost(t::Reduce, cfg::ArcEagerConfig, gold)
    # num of gold arcs (s,l',k) s.t. k ϵ b|β
    σ, s = popstack(cfg)
    count(k -> has_arc(gold, s, k), buffer(cfg))
end

function cost(t::Shift, cfg::ArcEagerConfig, gold)
    # num of gold arcs (k,l',b), (b,l',k) s.t. k ϵ s|σ
    b, β = shiftbuffer(cfg)
    count(k -> has_arc(gold, k, b) || has_arc(gold, b, k), stack(cfg))
end

function is_possible(::LeftArc, cfg::ArcEagerConfig)
    s = last(stack(cfg))
    return s != 0 && !hashead(token(cfg, s))
end

is_possible(::RightArc, cfg::ArcEagerConfig) =
    !hashead(token(cfg, first(buffer(cfg))))

is_possible(::Reduce, cfg::ArcEagerConfig) =
    hashead(token(cfg, last(stack(cfg))))

is_possible(::Shift, cfg::ArcEagerConfig) = true

==(cfg1::ArcEagerConfig, cfg2::ArcEagerConfig) = cfg1.c == cfg2.c

Base.getindex(cfg::ArcEagerConfig, i) = arc(cfg, i)


"""
    annotations(strings::Vector{String}, positions::Vector{Point})

Plots an array of texts at each position in `positions`.

## Attributes
$(ATTRIBUTES)
"""
@recipe(Annotations, text, position) do scene
    default_theme(scene, Text)
end

function convert_arguments(::Type{<: Annotations},
                           strings::AbstractVector{<: AbstractString},
                           text_positions::AbstractVector{<: Point{N}}) where N
    return (map(strings, text_positions) do str, pos
        (String(str), Point{N, Float32}(pos))
    end,)
end

function plot!(plot::Annotations)
    # annotations are not necessary anymore with the different text behavior
    text!(plot, plot[1]; plot.attributes...)
    plot
end


# --------------------------------------------------------------------------- #
# Labels -- a list of cell position coordinates

struct Label
  data :: Vector{Tuple{Int, Int}}
end

function Base.convert(::Type{Label}, mat::Matrix{Int})
  @assert (size(mat, 1) == 2)
  return Label([(mat[:,i]...) for i in 1:size(mat, 2)])
end

function Base.convert(::Type{Matrix{Int}}, lab::Label)
  return hcat([[coords...] for coords in lab.data]...)
end

Label() = Label(Tuple{Int, Int}[])
Label(mat::Matrix{Int}) = convert(Label, mat)

# --------------------------------------------------------------------------- #
# See tuples of images and labels as "labeled Images"

const LabeledGreyscaleImage = Tuple{GreyscaleImage, Label}
const LabeledRGBImage = Tuple{RGBImage, Label}
const LabeledImage = Union{LabeledGreyscaleImage, LabeledRGBImage}

# --------------------------------------------------------------------------- #
# Proximity maps
# Proxymaps are "washed out" 2d arrays with peaks at given label positions

function proxymap(width, height, lbls :: Vector{Label};
                  kernelsize :: Int = 7,
                  stddev :: Real = kernelsize/4,
                  peakheight :: Real = 100)

  # Create proximity maps from labels
  # Each point in the label gets converted to a gaussian activation

  n = length(lbls)

  # Pixel map
  prmap = zeros(Float32, width, height, 1, n)
  for i in 1:n
    for (x, y) in lbls[i]
      prmap[y, x, 1, i] = 1.
    end
  end

  # Convolve the pixel-map with a suitable kernel
  c = ceil(Int, kernelsize/2)
  kernel = Float32[ exp(- ((i - c)^2 + (j - c)^2) / (2stddev^2))
                    for i in 1:kernelsize, j in 1:kernelsize ]
  kernel = reshape(kernel, size(kernel)..., 1, 1)

  # Scale the output such that deviations from real labels are
  # punished more than deviations from the background
  pad = floor(Int, kernelsize/2)
  prmap = reshape(conv4(kernel, prmap, padding=pad), width, height, n)

  return peakheight * prmap
end

function proxymap(width, height, lbl :: Label; kwargs...)
  return reshape(proxymap(width, height, [lbl]; kwargs...), width, height)
end

# --------------------------------------------------------------------------- #
# Adjacency
# Calculates a list of smallest distances between the spots in two labels

function adjacency(dlbl, tlbl)
  ddata, tdata = convert.(Matrix{Int}, (dlbl, tlbl))
  if isempty(ddata) && isempty(tdata)
    return 0.
  elseif isempty(ddata) || isempty(tdata)
    return Inf
  end

  n, m = length(tlbl), length(dlbl)
  vals = zeros(n + m)
  for i in 1:n
    dists = sqrt.(sum(abs2, ddata .- tdata[:,i], 1))[1,:]
    vals[i] = minimum(dists)
  end
  for i in 1:m
    dists = sqrt.(sum(abs2, ddata[:,i] .- tdata, 1))[1,:]
    vals[n+i] = minimum(dists)
  end
  return vals
end

function adjacency(model, img :: Image, lbl :: Label; kwargs...)
  return adjacency(label(model, img; kwargs...), lbl)
end

# --------------------------------------------------------------------------- #
# Convenience functions

Base.size(l :: Label, args...; kwargs...) = size(l.data, args..., kwargs...)
Base.length(l :: Label, args...; kwargs...) = length(l.data, args..., kwargs...)
Base.getindex(l :: Label, k :: Integer) = l.data[k]
Base.setindex!(l :: Label, v, k :: Integer) = setindex!(l.data, (v...), k)
Base.endof(l :: Label) = length(l)

Base.start(l::Label) = start(l.data)
Base.next(l::Label, state) = next(l.data, state)
Base.done(l::Label, state) = done(l.data, state)
Base.eltype(::Type{Label}) = Tuple{Int, Int}

Base.push!(l :: Label, coords) = push!(l.data, coords)
Base.pop!(l :: Label) = pop!(l.data)
Base.deleteat!(l :: Label, idx) = deleteat!(l.data, idx)
Base.join(l1 :: Label, l2 :: Label) = Label([l1.data; l2.data])

Base.isempty(l :: Label) = isempty(l.data)

function crop(lbl :: Label, x :: Int, y :: Int, w :: Int, h :: Int)
  o = (x-1, y-1)
  data = [z .- o for z in lbl.data if x <= z[1] < x+w && y <= z[2] < y+h]
  return Label(data)
end

crop(lmg :: LabeledImage, x, y, w, h) = crop.(lmg, x, y, w, h)

function lblscale(lbl :: Label, s :: Real)
  return Label([round.(Int, s .* c) for c in lbl])
end


using LasIO
using LazIO
using FileIO
using DataFrames
using ProgressMeter
using StaticArrays
using BenchmarkTools
using PointCloudRasterizers
using NearestNeighbors #use it for spatial search

#get the pulse
function classify(points)
    prev_point = points[1]
    pulse =  Dict{Int,String}() #dictionary with indexes and tags
    n = 0
    for (i,p) in enumerate(points)
        point = p
        inten = intensity(p)
        rn = return_number(p)
        nr = number_of_returns(p)
        dt = gps_time(point) - gps_time(prev_point)
        #rn = return number, nr = number_of_returns
        if (rn == n + 1) && (nr > 1)
            if rn == 1
                pulse[i] = "firstpoint"
            elseif rn < nr && dt < 1e-7
                pulse[i] = "nextpoint"
            elseif rn == nr && dt < 1e-7
                pulse[i] = "lastpoint"
            end
            if nr == rn
                n = 0
            end
        elseif nr == 1 && rn == 1
            pulse[i] = "firstpoint"
            n = 0
        else
            pulse[i] ="unclassified"
            n = 0
        end
        n+=1
        prev_point = p
    end
    pulse
end

#write point cloud
function write_output(writefiles,ds, pulse)
    LazIO.write(writefiles, ds.header) do io
        for (i,p) in enumerate(ds) #I need to interate only to dataset, since it is Lazpoints
            allkeys = keys(pulse)
            if i in allkeys
                if pulse[i] == "unclassified"
                    p.classification = UInt(31)
                elseif pulse[i] == "firstpoint"
                    p.classification = UInt(3)
                elseif pulse[i] == "nextpoint"
                    p.classification = UInt(4)
                elseif pulse[i] == "lastpoint"
                    p.classification = UInt(5)
                end
                LazIO.writepoint(io,p)
            end
        end
    end
end

##MAIN##
#working with 4 different datasets
workdir = dirname(@__FILE__)
# dataname_1 = "NL3_clipped_water_cropheight_pointformat3_classification_sorted"
# dataname_2 = "3_splitted_sorted_pointformat3"
# dataname_3 = "12_splitted_sorted_pointformat3_clipZ"
dataname_4 = "14_splitted_sorted_pointformat3_clipZ"

const filenn_NL1 = joinpath(workdir,dataname_4 * ".laz")
const filenn_NL1_las = File{format"LAZ_"}(filenn_NL1)

#read specific LAS file
header, points = LazIO.load(filenn_NL1_las) #point cloud as LAS format
n = length(points) #length

#output files
outputname = "_classified"
writefiles = joinpath(workdir, dataname_4 * outputname * ".laz")

#read dataset
ds = LazIO.open(filenn_NL1)#dataset

#pulse
classified = classify(points)

#write output
write_output(writefiles,ds,classified)


###
### grow_beg!
###
# can't call grow_beg! on this
# function grow_beg!(x::DynamicAxis{T}, n::Unsigned) where {T} end

function grow_beg!(x, n::Integer)
    n < 0 && throw(ArgumentError("new length must be ≥ 0"))
    return unsafe_grow_beg!(x, n)
end
unsafe_grow_beg!(x::Vector, n) = Base._growbeg!(x, n)
unsafe_grow_beg!(x::MutableRange, n) = setfield!(x, :parent, unsafe_grow_beg(parent(x), n))

###
### grow_end!
###
""" grow_end! """
function grow_end!(x, n::Integer)
    n < 0 && throw(ArgumentError("new length must be ≥ 0"))
    unsafe_grow_end!(x, n)
    return x
end
unsafe_grow_end!(x::Vector, n) = Base._growend!(x, n)
unsafe_grow_end!(x::MutableRange, n) = setfield!(x, :parent, unsafe_grow_end(parent(x), n))

""" grow_end """
function grow_end(x, n::Integer)
    n < 0 && throw(ArgumentError("n must be positive; got $n"))
    return unsafe_grow_end(x, n)
end
function unsafe_grow_end(x::AbstractRange, n::Integer)
    if known_step(x) === 0
        s = step(x)
        return first(x):s:(last(x) + (s * n))
    else
        return first(x):(last(x) + n)
    end
end

###
### grow_to!
###
function grow_to!(x, n::Integer)
    len = length(x)
    if len < n
        unsafe_grow_to!(x, n)
        return x
    elseif len == n
        return x
    else
        throw(ArgumentError("new length must be ≥ than length of collection, got length $(length(x))."))
    end
end

unsafe_grow_to!(x::AbstractRange, n::Integer) = unsafe_grow_end!(x, n - length(x))

function grow_to(x, n::Integer)
    len = length(x)
    if len <= n
        return unsafe_grow_to(x, n)
    else
        throw(ArgumentError("new length must be ≥ than length of collection, got length $(length(x))."))
    end
end
unsafe_grow_to(x::AbstractRange, n::Integer) = unsafe_grow_end(x, n - length(x))

###
### shrink_beg!
###
function shrink_beg!(x, n::Integer)
    n > length(x) && throw(ArgumentError("new length cannot be < 0"))
    return unsafe_shrink_beg!(x, n)
end
unsafe_shrink_beg!(x::Vector, n) = Base._deletebeg!(x, n)
unsafe_shrink_beg!(x::MutableRange, n) = setfield!(x, :parent, unsafe_shrink_beg(parent(x), n))

###
### shrink_end!
###
function shrink_end!(x, n::Integer)
    n < 0 && throw(ArgumentError("new length must be ≥ 0"))
    return unsafe_shrink_end!(x, n)
end
unsafe_shrink_end!(x::Vector, n) = Base._deleteend!(x, n)
unsafe_shrink_end!(x::MutableRange, n) = setfield!(x, :parent, unsafe_shrink_end(parent(x), n))

###
### shrink_to!(x::AbstractRange, n::Integer
###
function shrink_to!(x, n::Integer)
    len = length(x)
    if len > n
        return unsafe_shrink_to!(x, n)
    elseif len == n
        return x
    else
        throw(ArgumentError("new length must be ≤ than length of collection, got length $(length(x))."))
    end
end

unsafe_shrink_to!(x, n) = unsafe_shrink_end!(x, length(x) - n)

function grow_beg(x, n::Integer)
    n < 0 && throw(ArgumentError("n must be positive; got $n"))
    return unsafe_grow_beg(x, n)
end
function unsafe_grow_beg(x::AbstractRange, n::Integer)
    if known_step(x) === 0
        s = step(x)
        return (first(x) - (s * n)):s:last(x)
    else
        return (first(x) - n):last(x)
    end
end
function shrink_beg(x, n::Integer)
    n < 0 && throw(ArgumentError("n must be positive; got $n"))
    return unsafe_shrink_beg(x, n)
end
function unsafe_shrink_beg(x::AbstractRange, n::Integer)
    if known_step(x) === 0
        s = step(x)
        return (first(x) + (s * n)):s:last(x)
    else
        return (first(x) + n):last(x)
    end
end

function shrink_end(x, n::Integer)
    n < 0 && throw(ArgumentError("n must be positive; got $n"))
    return unsafe_shrink_end(x, n)
end

function unsafe_shrink_end(x::AbstractRange, n::Integer)
    if known_step(x) === 0
        s = step(x)
        return first(x):s:(last(x) - (s * n))
    else
        return first(x):(last(x) - n)
    end
end


for fn_fits in fns
    local file = FITS(fn_fits)
    local l1 = read(file[2],"wavelength")
    local f1 = read(file[2],"spectrum")
    local v1 = read(file[2],"uncertainty")
    fn_jld2 = replace(fn_fits,"fits"=>"jld2")
    println(fn_jld2)
    @save fn_jld2 λ=l1 flux=f1 var=v1
end

@time begin
    for fn_fits in fns
    fn_jld2 = replace(fn_fits,"fits"=>"jld2")
    local f = jldopen(fn_jld2, "r")
    local l1 = f["λ"]
    local f1 = f["flux"]
    local v1 = f["var"]
    println("sum(l) = ", sum(view(l1[:,10:20],.!isnan.(l1[:,10:20]))),
    " sum(f) = ", sum(view(f1[:,10:20],.!isnan.(f1[:,10:20]))),
    " sum(v) = ", sum(view(v1[:,10:20],.!isnan.(v1[:,10:20]))) )
end
end

@time begin for fn_fits in fns
    local file = FITS(fn_fits)
    local l1 = read(file[2],"wavelength")
    local f1 = read(file[2],"spectrum")
    local v1 = read(file[2],"uncertainty")
    println("sum(l) = ", sum(view(l1[:,10:20],.!isnan.(l1[:,10:20]))),
    " sum(f) = ", sum(view(f1[:,10:20],.!isnan.(f1[:,10:20]))),
    " sum(v) = ", sum(view(v1[:,10:20],.!isnan.(v1[:,10:20]))) )
end
end

@time begin
    for fn_fits in fns
    fn_jld2 = replace(fn_fits,"fits"=>"jld2")
    local f = h5open(fn_jld2, "r")
    local l1 = readmmap(f["λ"])
    local f1 = readmmap(f["flux"])
    local v1 = readmmap(f["var"])
    println("sum(l) = ", sum(view(l1[:,10:20],.!isnan.(l1[:,10:20]))),
    " sum(f) = ", sum(view(f1[:,10:20],.!isnan.(f1[:,10:20]))),
    " sum(v) = ", sum(view(v1[:,10:20],.!isnan.(v1[:,10:20]))) )
end
end


using DeepForest, Flux
using StatsBase: mode, mean
using Test

@testset "Leaf precompute" begin
    l = DeepForest.Leaf()
    x = Float32.([1.0 2.0; 3.0 4.0])
    y = [1, 1]
    DeepForest.precompute!(l, x, y)
    @test l.pred == 1
    y = [2, 2]
    DeepForest.precompute!(l, x, y)
    @test l.pred == 2
end

@testset "Leaf predict" begin
    l = DeepForest.Leaf()
    x = Float32.([1.0 2.0; 3.0 4.0])
    y = [1, 1]
    DeepForest.precompute!(l, x, y)
    @test DeepForest.predict(l, x) == [1, 1]
end

@testset "Node init" begin
    feat_sub = Dict(
        1 => [1, 2]
    )

    n = Node(feat_sub, 5, 1, 1)
    @test n.subset == [1, 2]
    @test typeof(n.left) == DeepForest.Leaf

    feat_sub = Dict(
        1 => [1, 2],
        2 => [2, 3],
        3 => [1, 3]
    )

    n = Node(feat_sub, 5, 2, 1)
    @test n.subset == [1, 2]
    @test n.left.subset == [2, 3]
    @test n.right.subset == [1, 3]
end

@testset "Node parameters" begin
    feat_sub = Dict(
        1 => [1, 2],
        2 => [2, 3],
        3 => [1, 3]
    )

    n = Node(feat_sub, 5, 2, 1)
    @test params(n) != Flux.Params([])
end

@testset "Node precompute" begin
    feat_sub = Dict(
        1 => [1, 2],
        2 => [2, 3],
        3 => [1, 3]
    )

    n = Node(feat_sub, 5, 2, 1)
    x = rand(Float32, 3, 4)
    y = [1, 2, 2, 2]
    DeepForest.precompute!(n, x, y)
    @test n.best.first != 1 || n.best.second != 1
end

@testset "Node predict" begin
    feat_sub = Dict(
        1 => [1, 2],
        2 => [2, 3],
        3 => [1, 3]
    )

    n = Node(feat_sub, 5, 2, 1)
    x = rand(Float32, 3, 4)
    y = [1, 2, 2, 2]
    DeepForest.precompute!(n, x, y)
    @test length(predict(n, x)) == 4
    @test length(unique(predict(n, x))) <= 2
end

@testset "Node loss" begin
    feat_sub = Dict(
        1 => [1, 2],
        2 => [2, 3],
        3 => [1, 3]
    )

    n = Node(feat_sub, 5, 2, 1)
    x = rand(Float32, 3, 4)
    y = [1, 2, 2, 2]
    DeepForest.precompute!(n, x, y)

    @test loss!(n, x, y) != 0

    gs = gradient(() -> loss!(n, x, y), params(n))
    @test gs != IdDict()
end

@testset "Node training" begin
    x = rand(Float32, 2, 1000);
    x[1, :] .*= 2*π;
    x[1, :] .-= π;
    x[2, :] .*= 3;
    x[2, :] .-= 1.5;
    y = Int64.(x[2, :] .< sin.(x[1, :])) .+ 1;

    features = Dict(
        7 => [1, 2],
        6 => [1, 2],
        5 => [1, 2],
        4 => [1, 2],
        3 => [1, 2],
        2 => [1, 2],
        1 => [1, 2]
    )

    model = Node(features, 10, 2, 1)

    tree_train!(1000, model, x, y)
    @test mean(predict(model, x) .== y) > .9
end

@testset "Node importances" begin
    x = rand(Float32, 2, 1000);
    x[1, :] .*= 2*π;
    x[1, :] .-= π;
    x[2, :] .*= 3;
    x[2, :] .-= 1.5;
    y = Int64.(x[2, :] .< sin.(x[1, :])) .+ 1;

    features = Dict(
        7 => [1, 2],
        6 => [1, 2],
        5 => [1, 2],
        4 => [1, 2],
        3 => [1, 2],
        2 => [1, 2],
        1 => [1, 2]
    )

    model = Node(features, 10, 2, 1)

    tree_train!(1000, model, x, y)
    imp = importance(model, x)
    @test imp[2] >= imp[1]
end


using QMTK
using Compat.Test

@testset "Sites" begin
    include("Sites.jl")
end

@testset "SubSites" begin
    include("SubSites.jl")
end

@testset "Tags" begin
    include("Tag/Tags.jl")
end

@testset "Ket" begin
    include("Ket.jl")
end

@testset "Space" begin
    include("Space/Space.jl")
end


#####
# some common components for Prioritized Experience Replay based methods
#####

const PERLearners = Union{PrioritizedDQNLearner,RainbowLearner,IQNLearner}

function RLBase.update!(learner::Union{DQNLearner, QRDQNLearner,REMDQNLearner,PERLearners}, t::AbstractTrajectory)
    length(t[:terminal]) - learner.sampler.n <= learner.min_replay_history && return

    learner.update_step += 1

    if learner.update_step % learner.target_update_freq == 0
        copyto!(learner.target_approximator, learner.approximator)
    end

    learner.update_step % learner.update_freq == 0 || return

    inds, batch = sample(learner.rng, t, learner.sampler)

    if t isa PrioritizedTrajectory
        priorities = update!(learner, batch)
        t[:priority][inds] .= priorities
    else
        update!(learner, batch)
    end
end

function RLBase.update!(
    trajectory::PrioritizedTrajectory,
    p::QBasedPolicy{<:PERLearners},
    env::AbstractEnv,
    ::PostActStage,
)
    push!(trajectory[:reward], reward(env))
    push!(trajectory[:terminal], is_terminated(env))
    push!(trajectory[:priority], p.learner.default_priority)
end

"""
    DuelingNetwork(;base, val, adv)

Dueling network automatically produces separate estimates of the state value function network and advantage function network. The expected output size of val is 1, and adv is the size of the action space.
"""
struct DuelingNetwork{B,V,A}
    base::B
    val::V
    adv::A
end

Flux.@functor DuelingNetwork

function (m::DuelingNetwork)(state)
    x = m.base(state)
    val = m.val(x)
    return val .+ m.adv(x) .- mean(m.adv(x), dims=1)
end


ops = Dict{Symbol, Any}()
has_lstm = false

ops[:InputLayer] = function(a)
    return nothing
end

ops[:Input] = function(a)
    return vcall(:.+, a, 0)
end

#data_format = "channels_last"
ops[:Conv] = function(a)
    activation = a.fields["activation"]
    if activation =="linear"
        activation = identity
    end
    if !haskey(weight[a.fields["name"]] ,a.fields["name"])
        dummy_name = a.fields["name"]*"_1"
        weight[a.fields["name"]][a.fields["name"]] = weight[a.fields["name"]][dummy_name]
    end
    w = weight[a.fields["name"]][a.fields["name"]]["kernel:0"]
    kernel_weight = permutedims(w, (3,4,2,1))
    #Flip the kernel
    num_filters = size(kernel_weight)[4]
    num_layers = size(kernel_weight)[3]
    kernel_weight = kernel_weight[end:-1:1, end:-1:1,:,:]
    #kernel_weight = reshape(kernel_weight, (2,2,1,1))
    if !haskey(weight[a.fields["name"]][a.fields["name"]], "bias:0")
        weight[a.fields["name"]][a.fields["name"]]["bias:0"] = [0]
    end
    kernel_bias = weight[a.fields["name"]][a.fields["name"]]["bias:0"]
    strides = (a.fields["strides"]...)
    if a.fields["padding"] == "valid"
        pads = (0,0)
    elseif a.fields["padding"] == "same"
        pads = (Int64.((a.fields["kernel_size"] .-1)./2)...)
    end
    f = (x,) -> begin
        x = permutedims(x, (3,2,4,1))
        temp = Conv(activation, kernel_weight, kernel_bias, strides, pads, (1,1))(x)
        return temp
    end
    return f
end

ops[:Conv1D] = function(a)
    activation = a.fields["activation"]
    if activation == "linear"
        activation = identity
    elseif activation == "relu"
        activation = relu
    end
    if !haskey(weight[a.fields["name"]] ,a.fields["name"])
        dummy_name = a.fields["name"]*"_1"
        weight[a.fields["name"]][a.fields["name"]] = weight[a.fields["name"]][dummy_name]
    end
    kernel_weight = weight[a.fields["name"]][a.fields["name"]]["kernel:0"]
    kernel_weight = permutedims(kernel_weight, (3,2,1))
    kernel_weight = kernel_weight[end:-1:1, end:-1:1, :]
    s =size(kernel_weight)
    new_size =(s[1], s[2], s[3], 1)
    kernel_weight = reshape(kernel_weight, new_size)
    if !haskey(weight[a.fields["name"]][a.fields["name"]], "bias:0")
        weight[a.fields["name"]][a.fields["name"]]["bias:0"] = [0]
    end
    kernel_bias = weight[a.fields["name"]][a.fields["name"]]["bias:0"]
    strides = (a.fields["strides"]...)[1]
    if a.fields["padding"] == "valid"
        pads = (0,0)
    elseif a.fields["padding"] == "same"
        pads = (Int64.((a.fields["kernel_size"] .-1)./2)...)
    end
    dilation = a.fields["dilation_rate"][1]
    f = (x,) -> begin
        if ndims(x) == 2
            n_shape = (1, size(x)[1], size(x)[2])
            x = reshape(x, n_shape)
        elseif ndims(x) == 3
            x = permutedims(x, (3,1,2))
        end
        x = permutedims(x, (2,3,1))
        s = size(x)
        new_size = (s[1], s[2], s[3], 1)
        x = reshape(x, new_size)
        kernel_weight = permutedims(kernel_weight, (1,2,4,3))
        res = Conv(activation, kernel_weight, kernel_bias, (strides, strides), pads, (dilation, dilation))(x)
        return permutedims(res[:,:,:,1], (2,1,3))
    end
    return f
end

ops[:Activation] = function(a)
    if a.fields["activation"] == "linear"
        return relu
    end
end

ops[:Concatenate] = function(a)
    return :cat
end

ops[:Dropout] = function(a)
    f = (x,) -> begin
        return x
    end
    return f
end

ops[:MaxPool] = function(a)
    f = (x,) -> begin
    x = permutedims(x, (3,2,4,1))
    return maxpool(x, (a.fields["pool_size"]...), pad=(0,0), stride=(a.fields["strides"]...))
end
    return f
end

ops[:MaxPooling1D] = function(a)
    if a.fields["padding"] == "valid"
        pad = (0,0)
    end
    pool_size = a.fields["pool_size"][1]
    stride = a.fields["strides"][1]

    f = (x,) -> begin
    if (size(x)[2] - pool_size + stride)/ stride % 1 == 0
        fin_size_middle = Int((size(x)[2] - pool_size + stride)/ stride)
    else
        t = size(x)[2] / pool_size
        fin_size_middle = Int(t - (t%1))
    end
    fin_size = (size(x)[1], fin_size_middle, size(x)[3])
    temp = []
    for i=1:size(x)[3]
        push!(temp, maxpool(reshape(x[1,:,i], (size(x)[2],1,1,1)), (pool_size,1), pad=pad, stride=(stride,1)))
    end
    res = temp[1]
    for i=2:size(x)[3]
        res = hcat(res, temp[i])
    end
    res = reshape(res, fin_size)
    println(size(res))
    return permutedims(res, (2,3,1))
    end
    return f
end

ops[:Flatten] = function(a)

    f = (x,) -> begin
    l = prod(size(x))
    if ndims(x) == 4
        return reshape(permutedims(x, (3,1,2,4)), (l,1))
    else
        x = permutedims(x, reverse(range(1, ndims(x))))
        return reshape(x, (l,1))
    end
    end
    return f
end

ops[:BatchNormalization] = function(a)
    epsilon = a.fields["epsilon"]
    momentum = a.fields["momentum"]
    return x -> BatchNorm(size(x)[3], ϵ=epsilon, momentum=momentum)(x)
end

ops[:Dense] = function(a)
    name = a.fields["name"]
    if !haskey(weight[name], name)
        weight[name][name] = weight[name][name*"_1"]
    end
    weight_kernel = weight[name][name]["kernel:0"]
    bias = weight[name][name]["bias:0"]
    if !haskey(a.fields, "activation")
       return Dense(weight_kernel, bias)
    elseif a.fields["activation"] == "linear" || a.fields["activation"] == "relu"
        a.fields["activation"] = "relu"
        return Dense(weight_kernel, bias), Symbol(a.fields["activation"])
    elseif a.fields["activation"] == "sigmoid"
        return Dense(weight_kernel, bias), x ->(sigmoid.(x))
    elseif a.fields["activation"] == "softmax"
        return Dense(weight_kernel, bias), softmax
    end
end

ops[:ZeroPadding2D] = function(a)
    pads = (a.fields["padding"][1]...)
    return x -> meanpool(x, (1,1), pad=pads, stride=(1,1))
end

ops[:AveragePooling1D] = function(a)
    if a.fields["padding"] == "valid"
        pad = (0,0)
    end
    pool_size = a.fields["pool_size"][1]
    stride = a.fields["strides"][1]

    f = (x,) -> begin
    fin_size_middle = size(x)[2] / 2
    fin_size = (size(x)[1], fin_size_middle, size(x)[3])
    temp = []
    for i=1:size(x)[3]
        push!(temp, meanpool(reshape(x[1,:,i], (size(x)[2],1,1,1)), (pool_size,1), pad=pad, stride=(stride,1)))
    end
    res = temp[1]
    for i=2:size(x)[3]
        res = hcat(res, temp[i])
    end
    return reshape(res, (size(x)[1],Int(size(x)[2]/2),size(x)[3]))
    end
    return f
end

ops[:AveragePooling2D] = function(a)
    pool_size = (a.fields["pool_size"]...)
    strides = (a.fields["strides"]...)
    if a.fields["padding"] == "valid"
        pads = (0,0)
    elseif a.fields["padding"] == "same"
        pads = (Int64.((a.fields["pool_size"] .-1)./2)...)
    end
    return x -> meanpool(x, pool_size, pad=pads, stride=strides)
end

ops[:GlobalAveragePooling2D] = function(a)
    return x-> mean(x, (1,2))
end

ops[:Add] = function(a)
    return :+
end

ops[:Reshape] = function(a)
    return (x -> reshape(x, (a.fields["target_shape"]...)))
end

ops[:relu] = function(a)
    return x -> broadcast(relu, x)
end

ops[:tanh] = function(a)
    return tanh
end

ops[:sigmoid] = function(a)
    return sigmoid
end

ops[:elu] = function(a)
    return elu
end

ops[:softplus] = function(a)
    return softplus
end

ops[:softmax] = function(a)
    return relu
end

# Embeddings

ops[:Embedding] = function(a)
    name = a.fields["name"]
    embedding_matrix = weight[name][name]["embeddings:0"]
    output_dim = a.fields["output_dim"]
    input_dim = a.fields["input_dim"]
    f = (x,) -> begin
        if ndims(x) == 1 || ndims(x) == 2
        temp = embedding_matrix[:, Int64(x[1])+1]
        for i=2:length(x)
            temp = hcat(temp, embedding_matrix[:, Int64(x[i])+1])
        end
            return permutedims(temp,reverse(range(1, ndims(temp))))
        end
    end
    return f
end

ops[:LSTM] = function(a)
    name = a.fields["name"]
    lstm_weight = weight[name][name]
    lstm_recurrent_kernel = lstm_weight["recurrent_kernel:0"]
    lstm_bias = lstm_weight["bias:0"]
    lstm_kernel = lstm_weight["kernel:0"]
    vec_size = size(lstm_recurrent_kernel)[2]
    avg_ = 178 # Add general case here
    limit = sqrt(3/avg_)
    kernel_init = linspace(-1*limit, limit, 256)

    if a.fields["return_sequences"] == false
        f = (x,)-> begin
            res= 0
            model = LSTM(Flux.LSTMCell(lstm_kernel, lstm_recurrent_kernel, lstm_bias, zeros(vec_size), zeros(vec_size)))
            x = permutedims(x, reverse(range(1, ndims(x))))
            for i=1:size(x)[2]
                res = model(x[:, i])
            end
            return res
        end
    else
        f = (x,)-> begin
            res= 0
            model = LSTM(Flux.LSTMCell(lstm_kernel, lstm_recurrent_kernel, lstm_bias, zeros(vec_size), zeros(vec_size)))
            ans = model(x[1])
            for i=2:length(x)
                res = model(x[i])
                ans = hcat(ans, res)
            end
            return ans[:, end]
        end
    end
    return f
end


module ADjson

using Statistics
using GZip
using JSON
using TensorCast
using Dates
using TimeZones
using ADerrors

"""
    load_json(fname::String)

Load data from a json.gz file

# Currently supported datatypes
- `Obs`
- `List`

# Arguments
- `fname::string`: File name. Suffix '.json.gz' is automatically added if not specified explicitly.

# Notes
- ADerrors always assumes that replica are labeled as r0, r1, etc. Custom replica names get lost in the conversion.
- Irregular Monte Carlo chains are currently not supported as not all cases pyerrors can output can be safely converted to ADerrors
"""
function load_json(fname::String)

    if !endswith(fname, ".json.gz")
        fname *= ".json.gz"
    end

    if !(isfile(fname) || islink(fname))
        error("File '" * fname * "' does not exist.")
    end

    df = GZip.open(fname, "r") do io
        JSON.parse(io)
    end

    println("Data has been written using ", df["program"])
    println("Format version ", df["version"])
    println("Written by ", df["who"], " on ", df["date"], " on host ", df["host"])
    if haskey(df, "description")
        println("\nDescription: ", df["description"])
    end

    res_list = Vector{Any}()

    for entry in df["obsdata"]
        if !(entry["type"] in ["Obs", "List", "Array", "Corr"])
            error("Type '" * entry["type"] * "' is not implemented." )
        end

        res = Vector{uwreal}()
        out_length = prod([parse(Int, o) for o in split(entry["layout"], ", ")])

        if haskey(entry, "data")
            for element in entry["data"]
                conc_deltas = [Float64[] for b in 1:out_length]
                int_cnfg_numbers = Int[]
                rep_lengths = Int[]

                for rep in element["replica"]
                    my_arr = rep["deltas"]
                    @cast deltas[i][j] := my_arr[j][i]
                    cnfg_numbers = convert(Array{Int,1}, deltas[1])

                    if length(element["replica"]) > 1
                        for ch in 1:length(cnfg_numbers)
                            if ch != cnfg_numbers[ch]
                                error("Irregular Monte Carlo chains for multiple replica cannot be safely read into ADerrors.")
                            end
                        end
                    end

                    append!(int_cnfg_numbers, cnfg_numbers)
                    append!(rep_lengths, length(cnfg_numbers))
                    for i in 1:out_length
                        append!(conc_deltas[i], Vector{Float64}(deltas[i + 1]) .- mean(deltas[i + 1]))
                    end
                end

                if length(res) == 0
                    for i in 1:out_length
                        if length(element["replica"]) > 1
                            push!(res, uwreal(Vector{Float64}(conc_deltas[i]), element["id"], rep_lengths))
                        else
                            push!(res, uwreal(Vector{Float64}(conc_deltas[i]), element["id"], int_cnfg_numbers, int_cnfg_numbers[end]))
                        end
                    end
                else
                    for i in 1:out_length
                        if length(element["replica"]) > 1
                            res[i] += uwreal(Vector{Float64}(conc_deltas[i]), element["id"], rep_lengths)
                        else
                            res[i] += uwreal(Vector{Float64}(conc_deltas[i]), element["id"], int_cnfg_numbers, sum(rep_lengths))
                        end
                    end
                end
            end
        end

        if haskey(entry, "cdata")
            for element in entry["cdata"]
                if element["layout"] != "1, 1"
                   error("ADerror does not support non-scalar constants at the moment.")
                end

                if length(res) == 0
                    for i in 1:out_length
                        push!(res, uwreal([0.0, sqrt(element["grad"][1][i] * element["cov"][1] * element["grad"][1][i])], element["id"]))
                    end
                else
                    for i in 1:out_length
                        res[i] += uwreal([0.0, sqrt(element["grad"][1][i] * element["cov"][1] * element["grad"][1][i])], element["id"])
                    end
                end
            end
        end

        for i in 1:out_length
            res[i] += entry["value"][i]
        end

        if length(res) == 1
            push!(res_list, res[1])
        else
            push!(res_list, reshape(res, Tuple(Int(x) for x in [parse(Int, o) for o in split(entry["layout"], ", ")])))
        end
    end

    if length(res_list) == 1
        return res_list[1]
    else
        return res_list
    end
end

"""
    dump_to_json(p, fname::String, description="", indent=1)

Dump data to a json.gz file

# Currently supported datatypes
- `Obs`

# Arguments
- `p::uwreal` or `p::Vector{uwreal}`: Data to be written to the file
- `fname::String`: File name. Suffix '.json.gz' is automatically added if not specified explicitly.
- `description::String`: Optional description to be added to the file
- `indent::Int`: Indent of the json structure. Use `ident=nothing` for smallest file size but less human readability.
"""
function dump_to_json(data, fname::String, description="", indent=1)
    jsonstring = Dict{String, Any}()
    jsonstring["program"] = "ADjson 1.1"
    jsonstring["who"] = ENV["USER"]
    jsonstring["host"] = gethostname() * ", " * Sys.MACHINE
    jsonstring["date"] = Dates.format(now(localzone()), "Y-mm-dd HH:MM:SS zzzz")
    jsonstring["version"] = "1.1"
    jsonstring["obsdata"] = []
    if !isempty(description)
        jsonstring["description"] = description
    end

    ws = ADerrors.wsg

    vec_data = Vector{uwreal}()
    if data isa uwreal
        push!(vec_data, data)
    elseif data isa Vector{uwreal}
        append!(vec_data, data)
    else
        error("Unkown data type.")
    end

    for (index, p) in enumerate(vec_data)
        push!(jsonstring["obsdata"], Dict{String, Any}())
        jsonstring["obsdata"][index]["layout"] = "1"
        jsonstring["obsdata"][index]["type"] = "Obs"
        jsonstring["obsdata"][index]["value"] = [p.mean]

        i_data = 0
        i_cdata = 0
        for i in 1:convert(Int32, ADerrors.unique_ids!(p, ws))
            # Collect and rename replica names
            my_rep_names = Vector{String}()
            for el in ADerrors.get_repnames_from_id(p.ids[i], ws)
                loc = findlast("_", el)
                push!(my_rep_names, el[1:Vector{Int}(loc)[1] - 1] * "|" * el[Vector{Int}(loc)[1] + 1:end])
            end

            # Get deltas
            dt = zeros(Float64, ws.fluc[ws.map_ids[p.ids[i]]].nd)
            for j in 1:length(p.prop)
                if (p.prop[j] && (ws.map_nob[j] == p.ids[i]))
                    dt .= dt .+ p.der[j] .* ws.fluc[j].delta
                end
            end

            idc = convert(Vector{Int32}, ADerrors.get_repidc_from_id(p.ids[i], ws))
            if length(idc) == 2
                # Write covobs data
                i_cdata += 1
                if !haskey(jsonstring["obsdata"][index], "cdata")
                    jsonstring["obsdata"][index]["cdata"] = []
                end
                push!(jsonstring["obsdata"][index]["cdata"], Dict{String, Any}())
                jsonstring["obsdata"][index]["cdata"][i_cdata]["id"] = ADerrors.get_name_from_id(p.ids[i], ws)
                jsonstring["obsdata"][index]["cdata"][i_cdata]["layout"] = "1, 1"
                jsonstring["obsdata"][index]["cdata"][i_cdata]["cov"] = [dt[1] ^ 2]
                jsonstring["obsdata"][index]["cdata"][i_cdata]["grad"] = [[1.0]]
            else
                # Write Monte Carlo data
                i_data += 1
                if !haskey(jsonstring["obsdata"][index], "data")
                    jsonstring["obsdata"][index]["data"] = []
                end
                push!(jsonstring["obsdata"][index]["data"], Dict{String, Any}())
                jsonstring["obsdata"][index]["data"][i_data]["id"] = ADerrors.get_name_from_id(p.ids[i], ws)
                jsonstring["obsdata"][index]["data"][i_data]["replica"] = []
                rep_indices = prepend!(cumsum(convert(Vector{Int32}, ws.fluc[ws.map_ids[p.ids[i]]].ivrep)), 0)
                for j in 1:length(my_rep_names)
                    for delta_entry in dt[rep_indices[j] + 1:rep_indices[j + 1]]
                        if isapprox(delta_entry, 0.0; atol=1e-14)
                            error("Irregular Monte Carlo chain cannot be safely written to json.gz file.")
                        end
                    end
                    tmp = []
                    push!(jsonstring["obsdata"][index]["data"][i_data]["replica"], Dict{String, Any}())
                    jsonstring["obsdata"][index]["data"][i_data]["replica"][j]["name"] = my_rep_names[j]
                    push!(tmp, idc[rep_indices[j] + 1:rep_indices[j + 1]])
                    push!(tmp, dt[rep_indices[j] + 1:rep_indices[j + 1]])
                    @cast deltas[i][j] := tmp[j][i]
                    jsonstring["obsdata"][index]["data"][i_data]["replica"][j]["deltas"] = deltas
                end
            end
        end
    end
    if !endswith(fname, ".json.gz")
        fname *= ".json.gz"
    end

    GZip.open(fname, "w") do f
        JSON.print(f, jsonstring, indent)
    end
end

export load_json, dump_to_json

end # module


module SomaticEvolution

using Distributions
using Statistics
using Random
using StatsBase
using DataFrames
using RecipesBase
using GLM
using LaTeXStrings
using Printf
using CSV
using Plots
using JSON
using DelimitedFiles
using AbstractTrees

export Cell,
MultiSimulation,
SimulationTracker,
SimulationInput,
MultilevelInput,
BranchingMoranInput,
BranchingInput,
MoranInput,
MultilevelInput,
Simulation,
SampledData,
PlotVAF,
PlotInverseVAF,
VAFResult,

multilevel_simulation,
multilevel_simulation_fast,
multiplesimulations,
run1simulation,
getallelefreq,
sampledallelefreq,
fitinverse,
gethist,
plotvaf,
plotvaf!,
plotinversevaf,
plotinversevaf!,
getstats,
saveinput,
branchingprocess,
moranprocess,
get_simulation,
pairwise_fixed_differences,
pairwise_fixed_differences_matrix,
pairwise_fixed_differences_statistics,
shared_fixed_mutations,
newmoduletimes,
cellpopulationsize,
average_mutations,
average_mutations_per_module,
mutations_per_cell,
mutation_ids_by_cell,
clonal_mutation_ids,
clonal_mutations,
subclonefreq,
getVAFresult,
age,
saveinput,
loadinput,
run_multilevel_from_file

include("definetypes.jl")
include("simulations.jl")
include("multilevel.jl")
include("process.jl")
include("sampling.jl")
include("analyse.jl")
include("multirun.jl")
include("plot.jl")
include("util.jl")
include("statistics.jl")
include("io.jl")
include("trees.jl")

end


import CVRPSEP_jll
const LIB_CVRPSEP = CVRPSEP_jll.libcvrpsep

# /* Definition of constants for identification of constraint types. */
@enum ConstraintType begin
    CMGR_CT_MIN_ROUTES        = 101
    CMGR_CT_NODE_DEGREE       = 102  # /* Degree = 2 for each customer. */
    CMGR_CT_CAP               = 103  # /* Capacity constraint. */
    CMGR_CT_GENCAP            = 104  # /* Generalized capacity constraint. */
    # CMGR_CT_FCI               = 104  # /* For public version. */ # duplicate
    CMGR_CT_TWOMATCHING       = 105
    CMGR_CT_COMB              = 106
    CMGR_CT_STR_COMB          = 107  # /* Strengthened comb. */
    CMGR_CT_HYPOTOUR          = 108
    CMGR_CT_EXT_HYPOTOUR      = 109
    CMGR_CT_MSTAR             = 110  # /* Homogeneous multistar. */
    CMGR_CT_WMSTAR            = 111  # /* Weak multistar. */
    CMGR_CT_DJCUT             = 112  # /* Disjunctive cut. */
    CMGR_CT_GOMORY            = 113  # /* By variable numbers */
    CMGR_CT_TWOEDGES_HYPOTOUR = 114  # /* 2EH inequality */
    CMGR_BT_CLIQUE_DOWN       = 201  # /* x(S:S) <= RHS */
    CMGR_BT_CLIQUE_UP         = 202  # /* x(S:S) >= RHS */
    CMGR_BT_STAR_DOWN         = 301  # /* x(i:F) <=RHS */
    CMGR_BT_STAR_UP           = 302  # /* x(i:F) >=RHS */
    CMGR_CT_SLB               = 401  # /* x(F) >= RHS. Simple lower bound */
end

mutable struct CnstrRecord
    Ctype       :: Cint
    Key         :: Cint
    IntListSize :: Cint
    IntList     :: Ptr{Cint}
    ExtListSize :: Cint
    ExtList     :: Ptr{Cint}
    CListSize   :: Cint
    CList       :: Ptr{Cint}
    CoeffList   :: Ptr{Cdouble}
    A           :: Cint
    B           :: Cint
    L           :: Cint
    RHS         :: Cdouble
    BranchLevel :: Cint
    GlobalNr    :: Cint
end

CnstrPointer = Ptr{CnstrRecord}
CnstrPointerList = Ptr{CnstrPointer}

mutable struct CnstrMgrRecord
    CPL         :: CnstrPointerList
    Dim         :: Cint
    Size        :: Cint
end

CnstrMgrPointer = Ptr{CnstrMgrRecord}

# test code

Dim = 100
old_cmp_ref = Ref{CnstrMgrPointer}()
my_cmp_ref = Ref{CnstrMgrPointer}()

ccall(
    (:CMGR_CreateCMgr, LIB_CVRPSEP),
    Cvoid,
    (Ptr{CnstrMgrPointer}, Cint),
    my_cmp_ref, Dim
)
ccall(
    (:CMGR_CreateCMgr, LIB_CVRPSEP),
    Cvoid,
    (Ptr{CnstrMgrPointer}, Cint),
    old_cmp_ref, Dim
)

my_cmp = my_cmp_ref[]
old_cmp = old_cmp_ref[]

@show my_cmp_ref
@show my_cmp

################################
n_customers = 36
demand = [16, 18, 1, 13, 8, 23, 7, 27, 1, 3, 6, 24, 19, 2, 5, 16, 7, 4, 22, 7, 23, 16, 2, 2, 9, 2, 12, 1, 9, 23, 6, 19, 7, 7, 20, 20]
capacity = 100
edge_tail = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 17, 17, 17, 17, 17, 17, 17, 17, 17, 17, 17, 17, 17, 17, 17, 17, 17, 17, 17, 17, 18, 18, 18, 18, 18, 18, 18, 18, 18, 18, 18, 18, 18, 18, 18, 18, 18, 18, 18, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21, 22, 22, 22, 22, 22, 22, 22, 22, 22, 22, 22, 22, 22, 22, 22, 23, 23, 23, 23, 23, 23, 23, 23, 23, 23, 23, 23, 23, 23, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 25, 25, 25, 25, 25, 25, 25, 25, 25, 25, 25, 25, 26, 26, 26, 26, 26, 26, 26, 26, 26, 26, 26, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 28, 28, 28, 28, 28, 28, 28, 28, 28, 29, 29, 29, 29, 29, 29, 29, 29, 30, 30, 30, 30, 30, 30, 30, 31, 31, 31, 31, 31, 31, 32, 32, 32, 32, 32, 33, 33, 33, 33, 34, 34, 34, 35, 35, 36]
edge_head = [2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 29, 30, 31, 32, 33, 34, 35, 36, 37, 30, 31, 32, 33, 34, 35, 36, 37, 31, 32, 33, 34, 35, 36, 37, 32, 33, 34, 35, 36, 37, 33, 34, 35, 36, 37, 34, 35, 36, 37, 35, 36, 37, 36, 37, 37]
edge_x = [1.0, 0.0, 2.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 2.0, 2.0, 0.0, 0.0, 0.0, 0.0, 2.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, 1.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.5, 0.0, 0.0, 0.0, 0.0, 0.5, 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.0, 0.0, 1.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, 1.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, 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.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.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, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.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, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.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, 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, 1.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, 1.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, 1.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, 0.0, 0.0, 0.0, 1.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, 1.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, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.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.5, 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.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.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, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.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, 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.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, 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.5, 0.0, 0.0, 0.0, 0.0, 0.5, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 1.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, 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.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, 1.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, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.5, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.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, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.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, 1.0, 0.0]
n_edges = length(edge_head)
max_n_cuts = 10
integrality_tolerance = 1e-6
integer_and_feasible = Ref{Cchar}(-9)
max_violation = Ref{Cdouble}(-9.9)

ccall(
    (:CAPSEP_SeparateCapCuts, LIB_CVRPSEP),
    Cvoid,
    (Cint, Ptr{Cint}, Cint, Cint, Ptr{Cint}, Ptr{Cint}, Ptr{Cdouble}, CnstrMgrPointer, Cint, Cdouble, Ref{Cchar}, Ref{Cdouble}, CnstrMgrPointer),
    n_customers, demand, capacity, n_edges, edge_tail, edge_head, edge_x, old_cmp, max_n_cuts, integrality_tolerance, integer_and_feasible, max_violation, my_cmp
)
@show integer_and_feasible[]
@show max_violation[]

@show old_cmr = unsafe_load(old_cmp)
@show my_cmr = unsafe_load(my_cmp)

@show cpl = my_cmr.CPL

jl_cpl = unsafe_wrap(Array, cpl, my_cmr.Size)

for i in 1:length(jl_cpl)
    cr = unsafe_load(jl_cpl[i])
    # @show cr
    S = unsafe_wrap(Vector{Cint}, cr.IntList, cr.IntListSize)[2:end]
    @show S
    # @show unsafe_wrap(Vector{Cint}, cr.ExtList, cr.ExtListSize)
    # @show unsafe_wrap(Vector{Cint}, cr.CList, cr.CListSize)
    # @show unsafe_wrap(Vector{Cdouble}, cr.CoeffList, cr.CListSize)
    @show cr.RHS
end


function diagonal_elements(n)

  reps = parse(Int64, ENV["LAMP_REPS"])

  A = randn(n,n)
  B = randn(n,n)

  Benchmarker.add_data(csv, "diag_mult_nai", Benchmarker.measure(reps, implicit_diagonal_mult, A, B))
  Benchmarker.add_data(csv, "diag_add_nai", Benchmarker.measure(reps, implicit_diagonal, A, B))
  Benchmarker.add_data(csv, "diag_add_rec", Benchmarker.measure(reps, explicit_diagonal, A, B))

  @test isapprox(implicit_diagonal(A, B)[2], explicit_diagonal(A, B)[2])
end

function implicit_diagonal_mult(A, B)
  result = zeros(1, size(A,1))
  start::Float64 = 0.0
  finish::Float64 = 0.0
  Benchmarker.cachescrub()
  Benchmarker.gcscrub()
  start = time_ns()

  result = diag(A*B)

  finish = time_ns()
  return ((finish-start)*1e-9, result)
end

function implicit_diagonal(A, B)
  result = zeros(1, size(A,1))
  start::Float64 = 0.0
  finish::Float64 = 0.0
  Benchmarker.cachescrub()
  Benchmarker.gcscrub()
  start = time_ns()

  result = diag(A+B)

  finish = time_ns()
  return ((finish-start)*1e-9, result)
end

function explicit_diagonal(A, B)
  result = zeros(1, size(A,1))
  start::Float64 = 0.0
  finish::Float64 = 0.0
  Benchmarker.cachescrub()
  Benchmarker.gcscrub()
  start = time_ns()

  result = diag(A) + diag(B)

  finish = time_ns()
  return ((finish-start)*1e-9, result)
end


export  alphanumeric_names

# ----------
# types

# ----------
# functions
function learn{S<:DSL_LearningAlgorithm}( mat::Matrix{Int}, alg::Type{S} )
	dset = matrix_to_dataset(mat)
	net  = Network()
	worked = learn!( net, dset, alg )
	(net, worked)
end
function learn{S<:DSL_LearnParams}( mat::Matrix{Int}, params::S )
	dset = matrix_to_dataset(mat)
	net  = Network()
	worked = learn!( net, dset, params )
	(net, worked)
end
function learn!{S<:DSL_LearningAlgorithm}( net::Network, mat::Matrix{Int}, alg::Type{S} )
	dset = matrix_to_dataset(mat)
	learn!( net, dset, alg )
end
function learn!{S<:DSL_LearnParams}( net::Network, mat::Matrix{Int}, params::S )
	dset = matrix_to_dataset(mat)
	learn!( net, dset, params )
end

function matrix_to_dataset{R <: Real, S<:AbstractString}( mat::Matrix{R}, names::Vector{S}=alphanumeric_names(size(mat,2)) )
	# converts a Matrix{Int} to a Dataset by writing it to a tempfile and loading it back in
	# each column in mat is one variable's data

	ftempname = tempname()
	writetable(ftempname, mat, names)
	dset = Dataset()
	if !read_file(dset, ftempname)
		error("Could not read data file $ftempname")
	end
	rm(ftempname)
	dset
end

function writetable{R <: Real, S<:AbstractString}(filename::AbstractString, mat::Matrix{R}, names::Vector{S}; separator::Char='\t')
	n,m = size(mat)
	f = open(filename, "w")

	for name in names[1:end-1]
		@printf(f, "%s%c", name, separator)
	end
	@printf(f, "%s\n", names[end])

	if isa(R, Integer)
		for i in 1 : n
			for j in 1 : m-1
				@printf(f, "%d%c", mat[i,j], separator)
			end
			@printf(f, "%d\n", mat[i,m])
		end
	else
		for i in 1 : n
			for j in 1 : m-1
				@printf(f, "%f%c", mat[i,j], separator)
			end
			@printf(f, "%f\n", mat[i,m])
		end
	end

	close(f)
end

function increment_alphanumeric_string(n,i=1)
	if i > length(n)
		n = "a"^i
	elseif n[i] != 'z'
		n = n[1:i-1]*string(n[i]+1)*n[i+1:end]
	else
		n = increment_alphanumeric_string(n[1:i-1]*"a"*n[i+1:end], i+1)
	end
	n
end
function alphanumeric_names(n::Int)
	# generate the first n lexographic names
	# a,b,c,...,z,aa,ab,...
	@assert(n > 0)
	names = Array(AbstractString, n)
	names[1] = "a"
	for i in 2 : n
		names[i] = increment_alphanumeric_string(names[i-1])
	end
	names
end


using Printf
using ArgParse
using HTTP
using JSON2
using Jtac

# -------- List of supported games ------------------------------------------- #

const GAMES = Dict(
    "tictactoe" => [:TicTacToe]
  , "mnkgame"   => [:MNKGame, Int, Int, Int]
  , "metatac"   => [:MetaTac]
  , "morris"    => [:Morris]
)

# -------- En-/decoding visual / renderable content -------------------------- #

include("rencoder.jl")

# -------- En-/decoding game states ------------------------------------------ #

struct GameState
  typ :: String
  params :: Vector{String}
  game :: Game
end

function GameState(game :: Game)

  typ = string(typeof(game))
  ff = findfirst('{', typ)
  fl = findlast('}', typ)

  if !isnothing(ff) && !isnothing(fl)
    params = split(typ[ff+1:fl-1], ",")
    typ = lowercase(typ[1:ff-1])
  else
    typ = lowercase(typ)
    params = String[]
  end

  GameState(typ, params, game)
end

function JSON2.read(io :: IO, :: Type{GameState})

  fields = JSON2.read(io)
  typ = fields.typ

  gametype = GAMES[typ][1]
  paramtypes = GAMES[typ][2:end]
  params = map(parse, paramtypes, fields.params)

  cexpr = isempty(params) ? gametype : Expr(:curly, gametype, params...)
  game = eval(cexpr)()

  for field in fieldnames(typeof(game))
    value = convert(getfield(game, field) |> typeof, getfield(fields.game, field))
    setfield!(game, field, value)
  end

  GameState(game)

end

# -------- Requests and Responses -------------------------------------------- #

struct TurnRequest
  game :: GameState
  action :: Int
end

TurnRequest(; game, action) = TurnRequest(game, action)
JSON2.@format TurnRequest keywordargs

struct ApplyRequest
  game :: GameState
  model :: String
  power :: Int
  temperature :: Float32
  exploration :: Float32
end

function ApplyRequest(; game, model, power, temperature, exploration)
  ApplyRequest(game, model, power, temperature, exploration)
end

JSON2.@format ApplyRequest keywordargs

struct ApplyResponse
  value :: Float32
  policy :: Vector{Float32}
  features :: Dict{String,Vector{Float32}}
end

# -------- Auxiliary Functions ----------------------------------------------- #

wrap(f, args...) = HTTP.Response(200, JSON2.write(f(args...)))

function load_models(dir, gpu)

  # Get all files/directories in dir
  files = readdir(dir)

  # Extract the files with '.jtm' ending and load them as async (gpu) models
  models = map(filter(x->length(x) > 4 && x[end-3:end] == ".jtm", files)) do fname
    name = fname[1:end - 4] # strip '.jtm'
    path = joinpath(dir, name)
    model = gpu ? load_model(path) |> to_gpu : load_model(path)
    name => Async(model)
  end

  Dict{String,Model}( "rollout" => RolloutModel()
                    , "random" => RandomModel()
                    , models... )

end

# -------- Handling GET Requests --------------------------------------------- #

function GET_index(req::HTTP.Request, index)
  if index == ""
    content ="Welcome the the Jtac AI server. No index file was provided"
  else
    content = read(index)
  end
  HTTP.Response(200, content)
end

function GET_reload(req::HTTP.Request, dir, gpu, models)

  # Remove previous models
  empty!(models)

  # Add new ones
  # TODO: We should be more efficient here and only look for newly added
  # / deleted models
  merge!(models, load_models(dir, gpu))

  "Successfully loaded $(length(models)) models."

end

function GET_models(req::HTTP.Request, models)
  map(collect(models)) do (name, model)
    Dict( "name"     => name, "game"     => string(gametype(model))
        , "params"   => sum([0; length.(Knet.params(model))])
        , "features" => Jtac.feature_name.(Jtac.features(model)))
  end
end

function GET_games(req::HTTP.Request)
  map(keys(GAMES), values(GAMES)) do k, v
    Dict("typ" => k, "params" => v[2:end])
  end
end

function GET_create(req::HTTP.Request)
  typ = HTTP.URIs.splitpath(req.target)[3]
  params = HTTP.URIs.splitpath(req.target)[4:end]

  gametype = GAMES[typ][1]
  paramtypes = GAMES[typ][2:end]
  params = map(parse, paramtypes, params)

  cexpr = isempty(params) ? gametype : Expr(:curly, gametype, params...)
  GameState(eval(cexpr)())
end

function GET_dummyimage(req::HTTP.Request)
  game = TicTacToe()
  for i in rand(1:5) random_turn!(game) end
  Image(game, rand(Float32, 9), 0.5f0) |> Typed
end

function POST_turn(req::HTTP.Request)
  req = JSON2.read(IOBuffer(HTTP.payload(req)), TurnRequest)
  apply_action!(req.game.game, req.action)
  GameState(req.game.game)
end

function POST_apply(req::HTTP.Request, models)

  # Parse target and json and sent via POST
  paths = HTTP.URIs.splitpath(req.target)

  req = JSON2.read(IOBuffer(HTTP.payload(req)), ApplyRequest)

  # Get model, game, power
  model  = models[req.model]
  game   = req.game.game
  power  = req.power

  # A value power <= 1 means that we pick the IntuitionPlayer
  if power <= 1
    player = IntuitionPlayer(model, temperature = req.temperature)
  else
    player = MCTSPlayer( model
                       , power = power
                       , temperature = req.temperature
                       , exploration = req.exploration )
  end

  # TODO: Get (improved) value of state from player!
  v, _, f = apply_features(base_model(model), game) # get value and features
  p = think(player, game)                # get improved policy

  # Bring features in nicer format as dictionary
  feats = Jtac.features(base_model(model))
  fdict = Dict{String,Vector{Float32}}(
            Jtac.feature_name(feat) => Jtac.feature_conv(feat, f[idx])
            for (feat, idx) in zip(feats, Jtac.feature_indices(feats, typeof(game))) )

  # If a visual response is desired, return the corresponding image data
  if length(paths) == 3 && paths[3] == "visual"

    Image(game, p, v) |> Typed

  # Else, send value/policy/feature data directly
  else

    ApplyResponse(v, p, fdict)

  end

end

# -------- Main entry function ----------------------------------------------- #

function aiserver( ip = "127.0.0.1"
                 , port = "4242"
                 , dir = "."
                 , gpu = "false"
                 , index = "" )

  # Load models from model_dir
  models = load_models(dir, gpu)

  # Create a router that maps addresses to function calls
  router = HTTP.Router()

  # Register addresses
  HTTP.@register router "GET"  "/" (r->GET_index(r, index))

  HTTP.@register router "GET"  "/api/reload" (r->wrap(GET_reload, r, dir, gpu, models))
  HTTP.@register router "GET"  "/api/models" (r->wrap(GET_models, r, models))
  HTTP.@register router "GET"  "/api/games"  (r->wrap(GET_games, r))
  HTTP.@register router "GET"  "/api/create" (r->wrap(GET_create, r))

  HTTP.@register router "GET"  "/api/dummyimage" (r->wrap(GET_dummyimage, r))

  HTTP.@register router "POST" "/api/turn"   (r->wrap(POST_turn, r))
  HTTP.@register router "POST" "/api/apply" (r->wrap(POST_apply, r, models))

  # Start the server
  println("Server listens on $ip:$port.")

  try

    HTTP.serve(router, ip, port)

  catch err

    println(err)

  end

end

# -------- Parsing arguments and calling entry point ------------------------- #

s = ArgParseSettings()
@add_arg_table s begin
    "--ip"
    help = "IP address that the server listens to"
    default = "127.0.0.1"
  "--port"
    help = "port that the server listens to"
    arg_type = Int
    default = 4242
  "--dir"
    help = "directory path that is searched for '.jtm' model files"
    default = "."
  "--gpu"
    help = "whether the GPU will be used if CUDA is supported"
    action = :store_true
  "--index"
    help = "statically served index page"
    default = ""
end

args = parse_args(s)

aiserver(args["ip"], args["port"], args["dir"], args["gpu"], args["index"])


__precompile__()

module TikzGraphs

export plot, Layouts
import LightGraphs: DiGraph, Graph, vertices, edges, src, dst

preamble = read(joinpath(dirname(@__FILE__), "..", "src", "preamble.tex"), String)

const AbstractGraph = Union{Graph, DiGraph}

using TikzPictures

module Layouts
    export Layered, Spring, SpringElectrical, SimpleNecklace

    abstract type Layout end

    struct Layered <: Layout
        sib_dist
        lev_dist
        Layered(;sib_dist=-1,lev_dist=-1) = new(sib_dist,lev_dist)
    end

    struct Spring <: Layout
        randomSeed
        dist
        Spring(;randomSeed=42,dist=-1) = new(randomSeed,dist)
    end

    struct SpringElectrical <: Layout
        randomSeed
        charge
        dist
        SpringElectrical(;randomSeed=42,charge=1.,dist=-1) = new(randomSeed,charge,dist)
    end

    struct SimpleNecklace <: Layout
    end
end

using .Layouts

plot(g, layout::Layouts.Layout, labels::Vector{T}=map(string, vertices(g)); args...) where {T<:AbstractString} = plot(g; layout=layout, labels=labels, args...)
plot(g, labels::Vector{T}; args...) where {T<:AbstractString} = plot(g; layout=Layered(), labels=labels, args...)

function edgeHelper(o::IOBuffer, a, b, edge_labels, edge_styles, edge_style)
    print(o, " [$(edge_style),")
    if haskey(edge_labels, (a,b))
        print(o, "edge label={$(edge_labels[(a,b)])},")
    end
    if haskey(edge_styles, (a,b))
        print(o, "$(edge_styles[(a,b)]),")
    end
    print(o, "] ")
end

function nodeHelper(o::IOBuffer, v, labels, node_styles, node_style)
    print(o, "$v/\"$(labels[v])\" [$(node_style)")
    if haskey(node_styles, v)
        print(o, ",$(node_styles[v])")
    end
    println(o, "],")
end

# helper function for edge type
edge_str(g::DiGraph) = "->"
edge_str(g::Graph) = "--"

function plot(g::AbstractGraph; layout::Layouts.Layout = Layered(), labels::Vector{T}=map(string, vertices(g)), edge_labels::Dict = Dict(), node_styles::Dict = Dict(), node_style="", edge_styles::Dict = Dict(), edge_style="", options="",prepend_preamble::String="") where T<:AbstractString
    o = IOBuffer()
    println(o, "\\graph [$(layoutname(layout)), $(options_str(layout))] {")
    for v in vertices(g)
        nodeHelper(o, v, labels, node_styles, node_style)
    end
    println(o, ";")
    for e in edges(g)
        a = src(e)
        b = dst(e)
        print(o, "$a $(edge_str(g))")
        edgeHelper(o, a, b, edge_labels, edge_styles, edge_style)
        println(o, "$b;")
    end
    println(o, "};")
    mypreamble = prepend_preamble * preamble * "\n\\usegdlibrary{$(libraryname(layout))}"
    TikzPicture(String(take!(o)), preamble=mypreamble, options=options)
end

for (_layout, _libraryname, _layoutname) in [
    (:Layered, "layered", "layered layout"),
    (:Spring, "force", "spring layout"),
    (:SpringElectrical, "force", "spring electrical layout"),
    (:SimpleNecklace, "circular", "simple necklace layout")
    ]
    @eval libraryname(p::$(_layout)) = $_libraryname
    @eval layoutname(p::$(_layout)) = $_layoutname
end

options_str(p::Layouts.Layout) = ""
function options_str(p::Layouts.Layered)
    sib_str = ""
    lev_str = ""
    if p.sib_dist > 0
        sib_str="sibling distance=$(p.sib_dist)mm,"
    end
    if p.lev_dist > 0
        lev_str = "level distance=$(p.lev_dist)mm,"
    end
    return sib_str*lev_str
end
function options_str(p::Spring)
    if p.dist == -1
        return "random seed = $(p.randomSeed),"
    else
        return "random seed = $(p.randomSeed), node distance=$(p.dist),"
    end
end
function options_str(p::SpringElectrical)
    if p.dist == -1
        return "random seed = $(p.randomSeed), electric charge=$(p.charge),"
    else
        return "random seed = $(p.randomSeed), electric charge=$(p.charge), node distance=$(p.dist),"
    end
end

end # module


using Test
using MPI

MPI.Init()

function scatterv_array(A, counts::Vector{Cint}, root)
    comm = MPI.COMM_WORLD
    T = eltype(A)
    if MPI.Comm_rank(comm) == root
        B = copy(A)
    else
        B = Array{T}(undef,1)
    end
    C = MPI.Scatterv(B, counts, root, comm)
end

comm = MPI.COMM_WORLD
size = MPI.Comm_size(comm)
rank = MPI.Comm_rank(comm)
root = 0
isroot = (MPI.Comm_rank(comm) == root)
# Defining this to make ones work for Char
Base.one(::Type{Char}) = '\01'

for typ in Base.uniontypes(MPI.MPIDatatype)

    # Allocating
    A = ones(typ, 3 * div(size,2) + mod(size,2))
    counts = Cint[ mod(i,2) + 1 for i in 0:(size-1)]
    B = scatterv_array(A, counts, root)
    @test B == ones(typ, mod(rank,2) + 1)

    # Non Allocating
    A = ones(typ, 3 * div(size,2) + mod(size,2))
    counts = Cint[ mod(i,2) + 1 for i in 0:(size-1)]
    B = Array{typ}(undef, counts[MPI.Comm_rank(comm) + 1])
    MPI.Scatterv!(A, B, counts, root, comm)
    @test B == ones(typ, mod(rank,2) + 1)

    # IN_PLACE
    counts = Cint[ mod(i,2) + 1 for i in 0:(size-1)]
    A = ones(typ, 3 * div(size,2) + mod(size,2))
    if isroot
        B = deepcopy(A)
    else
        B = Array{typ}(undef, counts[MPI.Comm_rank(comm) + 1])
    end
    MPI.Scatterv_in_place!(B, counts, root, comm)
    isroot  && @test B == A
    !isroot && @test B == ones(typ, mod(rank,2) + 1)
end

MPI.Finalize()
@test MPI.Finalized()


mutable struct Cons{T}
    car::T
    cdr::Any
end

struct Nil{T} end

const List{T} = Union{Cons{T},Nil{T}}

listof(::Type{T}) where {T} = Nil{T}()
listof(::Type{T}, x::T, xs::T...) where {T} = Cons{T}(x, listof(T, xs...))
Base.empty(::List{T}) where {T} = Nil{T}()

car(x::Cons) = x.car
cdr(x::Cons{T}) where {T} = x.cdr::Union{Cons{T},Nil{T}}

Base.eltype(::Type{List{T}}) where {T} = T
Base.IteratorSize(::Type{<:List}) = Base.SizeUnknown()
Base.iterate(::Nil) = nothing
Base.iterate(lst::Cons{T}, x = lst) where {T} =
    if x isa Cons{T}
        car(x), cdr(x)
    else
        nothing
    end

mutable struct LockedLinkedList{T}
    list::List{T}
    lock::ReentrantLock
    isempty::Threads.Atomic{Bool}
end

LockedLinkedList{T}() where {T} =
    LockedLinkedList{T}(Nil{T}(), ReentrantLock(), Threads.Atomic{Bool}(true))

setlist(l::LockedLinkedList{T}, list::List{T}) where {T} = @set l.list = list

function locked(f, l::LockedLinkedList)
    lock(l.lock) do
        f(l.list)
    end
end

function Base.empty!(l::LockedLinkedList)
    lock(l.lock) do
        l.list = empty(l.list)
        l.isempty[] = true
    end
end

Base.isempty(l::LockedLinkedList) = l.isempty[]

function setcdr!(lst::Cons{T}, cdr::List{T}) where {T}
    lst.cdr = cdr
    return lst
end

function setcdr!(l::LockedLinkedList{T}, cdr::List{T}) where {T}
    lock(l.lock) do
        if l.list isa Nil
            l.list = cdr
            if !(cdr isa Nil)
                l.isempty[] = false
            end
        else
            setcdr!(l.list, cdr)
            l.isempty[] = false
        end
    end
    return
end

function trypopfirst!(l::LockedLinkedList)
    # isempty(l) && return nothing  # TODO: is it OK?
    lock(l.lock) do
        list = l.list
        if list isa Nil
            return nothing
        else
            l.list = newlist = cdr(list)
            if newlist isa Nil
                l.isempty[] = true
            end
            return Some(car(list))
        end
    end
end


module TestVar

using Modia
using Modia.Measurements
using Modia.StaticArrays

struct nnnVar
	var::NamedTuple
	Var(; kwargs...) = new((; kwargs...) )
end

struct Model
	model::NamedTuple
	Model(; kwargs...) = new((; kwargs...) )
end

struct Map
	map::NamedTuple
	Map(; kwargs...) = new((; kwargs...) )
end

Var(;kwargs...) = (;kwargs...)
Var(value; kwargs...) = (; value=value, kwargs...)

#Var(args...; kwargs...) = begin println("1"); (;args..., kwargs...) end
#Var(value, args...; kwargs...) = typeof(value) <: NamedTuple ? begin println("2", value, args, kwargs); (;value..., args..., kwargs...) end : begin println("3"); (;value = value, args..., kwargs...) end
xar(value, args...; kwargs...) = begin
	@show value args kwargs merge(value, args[1], args[2], (;kwargs...))
	end

var(value, args...; kwargs...) = typeof(value) <: NamedTuple ?
	begin println("2", value, args, kwargs); value end : #merge(value, args, kwargs ) end :
	begin println("3"); (;value = value, args..., kwargs...) end
#Var(value, args...; kwargs...) = (;value = value, args..., kwargs...)
#Var(value::Union{Number, String, Expr, AbstractArray}, args...; kwargs...) = (;value = value, args..., kwargs...)

#NotPair = Union{Number, String, Expr, AbstractArray} # ..., ....}
#Var(value::NotPair, args...; kwargs...) = (;value = value, args..., kwargs...)

Model(; kwargs...) = (; kwargs...)

constant = Var(constant = true)
parameter = Var(parameter = true)
input = Var(input = true)
output = Var(output = true)
potential = Var(potential = true)
flow = Var(flow = true)

range(min, max) = Var(min=min, max=max)

recursiveMerge(x, ::Nothing) = x
recursiveMerge(x, y) = y
recursiveMerge(x::Expr, y::Expr) = begin dump(x); dump(y); Expr(x.head, x.args..., y.args...) end
recursiveMerge(x::Expr, y::Tuple) = begin x = copy(x); xargs = x.args; xargs[y[2]] = y[3]; Expr(x.head, xargs...) end

function recursiveMerge(nt1::NamedTuple, nt2::NamedTuple)
    all_keys = union(keys(nt1), keys(nt2))
    gen = Base.Generator(all_keys) do key
        v1 = get(nt1, key, nothing)
        v2 = get(nt2, key, nothing)
        key => recursiveMerge(v1, v2)
    end
    return Var(; gen...)
end

Base.:|(m::Model, n::Model) =  recursiveMerge(m.model, n.model)
Base.:|(m::Model, n::Map) = recursiveMerge(m.model, n.map)
Base.:|(m, n::Map) = recursiveMerge(m, n.map)
Base.:|(m, n) = if !(typeof(n) <: NamedTuple); recursiveMerge(m, (; value=n)) else recursiveMerge(n, (value=m,)) end

#Base.:|(m::Model, n::Model) =  (:MergeModel, m.model, n.model)
#Base.:|(m::Model, n::Map) =  (:MergeMap, m.model, n.map)
#Base.:|(m, n::Map) =  (:MergeMap, m, n.map)
#Base.:|(m::Var, n::Var) =  (:MergeModel, m.var, n.var)

# ---------------------------------------------------------------------

#v = Var(potential, nominal=10)
#@show v

v1 = input | output | flow | Var(min=0, nominal=10)
@show v1

v2 = parameter | 5
@show v2

v3 = 10u"m/s" | parameter
@show v3

v4 = parameter | 5.5 | Var(min=0)
@show v4

v5 = parameter | Var(5, min=0)
@show v5

Pin = Model( v = potential | Var(nominal=10), i = flow )
@show Pin

v6 = Var(2.0)
@show v6

v7 = Var(2.0 ± 0.1)
@show v7

v8 = Var(SVector(1.0,2.0,3.0))
@show v8

v9 = v4 | Var(parameter=false)
@show v9

v10 = v4 | Var(parameter=nothing)
@show v10

v11 = v4 | range(0,10)
@show v11

end


module UtilTest

using Test
using Random
using MeshCat: mergesorted!

@testset "mergesorted!" begin
    @testset "basic" begin
        rng = MersenneTwister(1)
        for i = 1 : 1000
            a = sort(rand(1 : 100, rand(0 : 20)))
            b = sort(rand(1 : 100, rand(0 : 20)))
            result = similar(a, length(a) + length(b))
            mergesorted!(result, a, b)
            @test issorted(result)
            @test a ⊆ result
            @test b ⊆ result
        end
    end

    @testset "by" begin
        rng = MersenneTwister(1)
        for i = 1 : 1000
            a = sort([rand(1 : 100) => rand(1 : 100) for _ in rand(0 : 20)]; dims=1, by=first)
            b = sort([rand(1 : 100) => rand(1 : 100) for _ in rand(0 : 20)]; dims=1, by=first)
            result = similar(a, length(a) + length(b))
            mergesorted!(result, a, b; by=first)
            @test issorted(result; by=first)
            @test a ⊆ result
            @test b ⊆ result
        end
    end
end

end