maximum(f, x::AbstractArray{Union{Missing, T},1}; threads = false) where T <: Union{INTEGERS, FLOATS, TimeType} = threads ? hp_maximum(f, x) : stat_maximum(f, x)
maximum(f, x) = Base.maximum(f, x)
maximum(x::AbstractArray{Union{Missing, T},1}; threads = false) where T <: Union{INTEGERS, FLOATS, TimeType}  = threads ? hp_maximum(identity, x) : stat_maximum(identity, x)
maximum(x) = Base.maximum(x)

minimum(f, x::AbstractArray{Union{Missing, T},1}; threads = false) where T <: Union{INTEGERS, FLOATS, TimeType} = threads ? hp_minimum(f, x) : stat_minimum(f, x)
minimum(f, x) = Base.minimum(f, x)
minimum(x::AbstractArray{Union{Missing, T},1}; threads = false) where T <: Union{INTEGERS, FLOATS, TimeType}= threads ? hp_minimum(identity, x) : stat_minimum(identity, x)
minimum(x) = Base.minimum(x)
# not optimised for simd
maximum(f, x::AbstractVector{Union{Missing, T}}) where T <: AbstractString = mapreduce(f, _stat_max_fun, x)
minimum(f, x::AbstractVector{Union{Missing, T}}) where T <: AbstractString = mapreduce(f, _stat_min_fun, x)
maximum(x::AbstractVector{Union{Missing, T}}) where T <: AbstractString = maximum(identity, x)
minimum(x::AbstractVector{Union{Missing, T}}) where T <: AbstractString = minimum(identity, x)

sum(f, x::AbstractArray{Union{Missing, T},1}; threads = false) where T <: Union{INTEGERS, FLOATS} = threads ? hp_sum(f, x) : stat_sum(f, x)
sum(f, x)=Base.sum(f, x)
sum(x::AbstractArray{Union{Missing, T},1}; threads = false) where T <: Union{INTEGERS, FLOATS} = threads ? hp_sum(identity, x) : stat_sum(identity, x)
sum(x) = Base.sum(x)

Statistics.mean(f, x::AbstractArray{Union{T,Missing},1}) where T <: Union{INTEGERS, FLOATS} = stat_mean(f, x)
Statistics.mean(x::AbstractArray{Union{T,Missing},1}) where T <: Union{INTEGERS, FLOATS} = stat_mean(x)

wsum(f, x::AbstractVector, w::AbstractVector) = stat_wsum(f, x, w)
wsum(x::AbstractVector, w::AbstractVector) = stat_wsum(identity, x, w)

wmean(f, x::AbstractVector, w::AbstractVector) = stat_wmean(f, x, w)
wmean(x::AbstractVector, w::AbstractVector) = stat_wmean(identity, x, w)

Statistics.var(f, x::AbstractArray{Union{T,Missing},1}, dof = true) where T <: Union{INTEGERS, FLOATS}= stat_var(f, x, dof)
Statistics.var(x::AbstractArray{Union{T,Missing},1}, dof = true) where T <: Union{INTEGERS, FLOATS}= stat_var(x, dof)

std(f, x::AbstractArray{Union{T,Missing},1}, dof = true) where T <: Union{INTEGERS, FLOATS}= stat_std(f, x, dof)
std(x::AbstractArray{Union{T,Missing},1}, dof = true) where T <: Union{INTEGERS, FLOATS}= stat_std(x, dof)

Statistics.median(x::AbstractArray{Union{T,Missing},1}) where T = stat_median(x)
Statistics.median!(x::AbstractArray{Union{T,Missing},1}) where T = stat_median!(x)

extrema(f, x::AbstractArray{Union{Missing, T},1}; threads = false) where T <: Union{INTEGERS, FLOATS, TimeType} = threads ? (hp_minimum(f, x), hp_maximum(f, x)) : (stat_minimum(f, x), stat_maximum(f, x))
extrema(f, x) = Base.extrema(f, x)
extrema(x::AbstractArray{Union{Missing, T},1}; threads = false) where T  <: Union{INTEGERS, FLOATS, TimeType}= threads ? (hp_minimum(identity, x), hp_maximum(identity, x)) : (stat_minimum(identity, x), stat_maximum(identity, x))
extrema(x) = Base.extrema(x)

# when by is a function the following functions find argmax/min(by.(x))
argmax(x::AbstractArray{Union{Missing, T},1}; by = identity) where T <: Union{INTEGERS, FLOATS, TimeType, AbstractString} = stat_findmax(by, x)[2]
argmax(x) = Base.argmax(x)
argmin(x::AbstractArray{Union{Missing, T},1}; by = identity) where T <: Union{INTEGERS, FLOATS, TimeType, AbstractString} = stat_findmin(by, x)[2]
argmin(x) = Base.argmin(x)

findmax(f, x::AbstractArray{Union{Missing, T},1}) where T <: Union{INTEGERS, FLOATS, TimeType, AbstractString} = stat_findmax(f, x)
findmax(f, x) = Base.findmax(f, x)
findmax(x::AbstractArray{Union{Missing, T},1}) where T <: Union{INTEGERS, FLOATS, TimeType, AbstractString} = stat_findmax(x)
findmax(x) = Base.findmax(x)
findmin(f, x::AbstractArray{Union{Missing, T},1}) where T <: Union{INTEGERS, FLOATS, TimeType, AbstractString} = stat_findmin(f, x)
findmin(f, x) = Base.findmin(f, x)
findmin(x::AbstractArray{Union{Missing, T},1}) where T <: Union{INTEGERS, FLOATS, TimeType, AbstractString} = stat_findmin(x)
findmin(x) = Base.findmin(x)

function cumsum(x::AbstractArray{Union{Missing, T},1}; missings = :ignore) where T <: Union{INTEGERS, FLOATS}
    if missings == :ignore
        stat_cumsum_ignore(x)
    elseif missings == :skip
        stat_cumsum_skip(x)
    else
        throw(ArgumentError("`missings` must be either `:ignore` or `:skip`"))
    end
end
cumsum(x) = Base.cumsum(x)
function cumsum!(outx::AbstractVector, x::AbstractArray{Union{Missing, T},1}; missings = :ignore) where T <: Union{INTEGERS, FLOATS}
    if missings == :ignore
        stat_cumsum!_ignore(outx, x)
    elseif missings == :skip
        stat_cumsum!_skip(outx, x)
    else
        throw(ArgumentError("`missings` must be either `:ignore` or `:skip`"))
    end
end
cumsum!(x,y) = Base.cumsum!(x,y)
function cumprod(x::AbstractArray{Union{Missing, T},1}; missings = :ignore) where T <: Union{INTEGERS, FLOATS}
    if missings == :ignore
        stat_cumprod_ignore(x)
    elseif missings == :skip
        stat_cumprod_skip(x)
    else
        throw(ArgumentError("`missings` must be either `:ignore` or `:skip`"))
    end
end
cumprod(x)=Base.cumprod(x)
function cumprod!(outx::AbstractVector, x::AbstractArray{Union{Missing, T},1}; missings = :ignore) where T <: Union{INTEGERS, FLOATS}
    if missings == :ignore
        stat_cumprod!_ignore(outx, x)
    elseif missings == :skip
        stat_cumprod!_skip(outx, x)
    else
        throw(ArgumentError("`missings` must be either `:ignore` or `:skip`"))
    end
end
cumprod!(x,y) = Base.cumprod!(x,y)
function cummin(x::AbstractArray{Union{Missing, T},1}; missings = :ignore) where T <: Union{INTEGERS, FLOATS, TimeType}
    if missings == :ignore
        stat_cummin_ignore(x)
    elseif missings == :skip
        stat_cummin_skip(x)
    else
        throw(ArgumentError("`missings` must be either `:ignore` or `:skip`"))
    end
end
function cummin!(outx, x::AbstractArray{Union{Missing, T},1}; missings = :ignore) where T <: Union{INTEGERS, FLOATS, TimeType}
    if missings == :ignore
        stat_cummin!_ignore(outx, x)
    elseif missings == :skip
        stat_cummin!_skip(outx, x)
    else
        throw(ArgumentError("`missings` must be either `:ignore` or `:skip`"))
    end
end
function cummax(x::AbstractArray{Union{Missing, T},1}; missings = :ignore) where T <: Union{INTEGERS, FLOATS, TimeType}
    if missings == :ignore
        stat_cummax_ignore(x)
    elseif missings == :skip
        stat_cummax_skip(x)
    else
        throw(ArgumentError("`missings` must be either `:ignore` or `:skip`"))
    end
end
function cummax!(outx, x::AbstractArray{Union{Missing, T},1}; missings = :ignore) where T <: Union{INTEGERS, FLOATS, TimeType}
    if missings == :ignore
        stat_cummax!_ignore(outx, x)
    elseif missings == :skip
        stat_cummax!_skip(outx, x)
    else
        throw(ArgumentError("`missings` must be either `:ignore` or `:skip`"))
    end
end


export TDLearner, DoubleLearner, DifferentialTDLearner, TDλReturnLearner

using .Utils: discount_rewards, discount_rewards_reduced
using Flux.Optimise: Descent
using LinearAlgebra: dot

"""
    TDLearner(approximator::Tapp, γ::Float64, optimizer::Float64; n::Int=0) where {Tapp<:AbstractVApproximator}
    TDLearner(approximator::Tapp, γ::Float64, optimizer::Float64; n::Int=0, method::Symbol=:SARSA) where {Tapp<:AbstractQApproximator}

The `TDLearner`(Temporal Difference Learner) use the latest `n` step experiences to update the `approximator`.
Note that `n` starts with `0`, which means looking forward for the next `n` steps.
`γ` is the discount rate of experience.
`optimizer` is the learning rate.

For [`AbstractVApproximator`](@ref), the only supported update method is `:SRS`, which means
only **S**tates, **R**ewards and next_**S**ates are used to update the `approximator`.

For [`AbstractQApproximator`](@ref), the following methods are supported:

- `:SARS` (aka Q-Learning)
- `:SARSA`
- `:ExpectedSARSA`
"""
mutable struct TDLearner{Tapp<:AbstractApproximator,method,O} <: AbstractLearner
    approximator::Tapp
    γ::Float64
    optimizer::O
    n::Int

    function TDLearner(
        ;
        approximator::Tapp,
        γ = 1.0,
        optimizer = Descent(1.0),
        n = 0,
        method::Symbol = :SARSA,
    ) where {Tapp<:AbstractApproximator}
        if approximator isa AbstractVApproximator
            if method != :SRS
                throw(ArgumentError("method [$method] is unsupported for a value approximator"))
            end
        elseif approximator isa AbstractQApproximator
            if !in(method, [:SARSA, :SARS, :ExpectedSARSA])
                throw(ArgumentError("Supported methods are $supported_methods , your input is $method"))
            end
        else
            throw(ArgumentError("unknown approximator"))
        end
        new{Tapp,method,typeof(optimizer)}(approximator, γ, optimizer, n)
    end
end

"""
    DoubleLearner(;L1, L2)

For now, this only supports [`TDLearner`](@ref) for `L1` and `L2`.
"""
Base.@kwdef struct DoubleLearner{T<:TDLearner} <: AbstractLearner
    L1::T
    L2::T
end

#####
# DifferentialTDLearner
#####

"""
    DifferentialTDLearner(;approximator::A, α::Float64, β::Float64, R̄::Float64 = 0.0, n::Int = 0)
"""
Base.@kwdef mutable struct DifferentialTDLearner{A<:AbstractApproximator} <: AbstractLearner
    approximator::A
    α::Float64
    β::Float64
    R̄::Float64 = 0.0
    n::Int = 0
end

function update!(learner::DifferentialTDLearner, transition)
    states, actions, rewards, terminals, next_states, next_actions = transition
    n, α, β, Q = learner.n, learner.α, learner.β, learner.approximator
    if length(states) ≥ n + 1
        s, a = states[1], actions[1]
        s′, a′ = next_states[end], next_actions[end]
        δ = sum(r -> r - learner.R̄, rewards) + Q(s′, a′) - Q(s, a)
        learner.R̄ += β * δ
        update!(Q, (s, a) => α * δ)
    end
end

extract_transitions(buffer, learner::DoubleLearner) =
    extract_transitions(buffer, learner.L1)

function update!(learner::DoubleLearner, args...)
    rand(Bool) ? update!(learner.L1, learner.L2, args...) :
    update!(learner.L2, learner.L1, args...)
end

(learner::DoubleLearner)(obs::Observation) = learner.L1(obs) .+ learner.L2(obs)

update!(learner::TDLearner{<:AbstractVApproximator,:SRS}, transitions) =
    update!(learner, transitions, nothing)

function update!(learner::TDLearner{<:AbstractVApproximator,:SRS}, transitions, weights)
    states, rewards, terminals, next_states = transitions.states,
        transitions.rewards,
        transitions.terminals,
        transitions.next_states
    n, γ, V, optimizer = learner.n, learner.γ, learner.approximator, learner.optimizer

    if length(terminals) > 0 && terminals[end]
        @views gains = discount_rewards(rewards[max(end - n, 1):end], γ)  # n starts with 0
        cum_weights = isnothing(weights) ? nothing : cumprod!(reverse(weights))
        for (i, G) in enumerate(gains)
            @views s = states[end-length(gains)+i]
            if isnothing(cum_weights)
                update!(V, s => apply!(optimizer, s, G - V(s)))
            else
                update!(V, s => apply!(optimizer, s, cum_weights[i] * (G - V(s))))
            end
        end
    else
        if length(states) ≥ (n + 1)  # n starts with 0
            @views G = discount_rewards_reduced(rewards[end-n:end], γ) +
                       γ^(n + 1) * V(next_states[end])
            @views s = states[end-n]
            w = isnothing(weights) ? 1.0 : reduce(*, weights)
            update!(V, s => apply!(optimizer, s, w * (G - V(s))))
        end
    end
end

function extract_transitions(
    buffer::EpisodeTurnBuffer,
    learner::TDLearner{<:AbstractVApproximator,:SRS},
)
    n = learner.n
    if length(buffer) > 0
        @views (
            states = state(buffer)[max(1, end - n - 1):end-1],
            rewards = reward(buffer)[max(1, end - n - 1)+1:end],
            terminals = terminal(buffer)[max(1, end - n - 1)+1:end],
            next_states = state(buffer)[max(1, end - n - 1)+1:end],
        )
    else
        nothing
    end
end

function update!(learner::TDLearner{<:AbstractQApproximator,:SARSA}, transitions)
    states, actions, rewards, terminals, next_states, next_actions = transitions
    n, γ, Q, optimizer = learner.n, learner.γ, learner.approximator, learner.optimizer

    if length(terminals) > 0 && terminals[end]
        @views gains = discount_rewards(rewards[max(end - n, 1):end], γ)  # n starts with 0
        for (i, G) in enumerate(gains)
            @views s, a = states[end-length(gains)+i], actions[end-length(gains)+i]
            update!(Q, (s, a) => apply!(optimizer, (s, a), G - Q(s, a)))
        end
    else
        if length(states) ≥ (n + 1)  # n starts with 0
            @views s, a, s′, a′ = states[end-n],
                actions[end-n],
                next_states[end],
                next_actions[end]
            @views G = discount_rewards_reduced(rewards[end-n:end], γ) +
                       γ^(n + 1) * Q(s′, a′)
            update!(Q, (s, a) => apply!(optimizer, (s, a), G - Q(s, a)))
        end
    end
end

function extract_transitions(
    buffer::EpisodeTurnBuffer,
    learner::Union{TDLearner{<:AbstractQApproximator,:SARSA},DifferentialTDLearner},
)
    n = learner.n
    if length(buffer) > 0
        @views (
            states = state(buffer)[max(1, end - n - 1):end-1],
            actions = action(buffer)[max(1, end - n - 1):end-1],
            rewards = reward(buffer)[max(1, end - n - 1)+1:end],
            terminals = terminal(buffer)[max(1, end - n - 1)+1:end],
            next_states = state(buffer)[max(1, end - n - 1)+1:end],
            next_actions = action(buffer)[max(1, end - n - 1)+1:end],
        )
    else
        nothing
    end
end

function update!(learner::TDLearner{<:AbstractQApproximator,:ExpectedSARSA}, transitions)
    states, actions, rewards, terminals, next_states, probs_of_a′ = transitions
    n, γ, Q, optimizer = learner.n, learner.γ, learner.approximator, learner.optimizer

    if length(terminals) > 0 && terminals[end]
        @views gains = discount_rewards(rewards[max(end - n, 1):end], γ)  # n starts with 0
        for (i, G) in enumerate(gains)
            @views s, a = states[end-length(gains)+i], actions[end-length(gains)+i]
            update!(Q, (s, a) => apply!(optimizer, (s, a), G - Q(s, a)))
        end
    else
        if length(states) ≥ (n + 1)  # n starts with 0
            @views s, a, s′ = states[end-n], actions[end-n], next_states[end]
            @views G = discount_rewards_reduced(rewards[end-n:end], γ) +
                       γ^(n + 1) * dot(Q(s′), probs_of_a′)
            update!(Q, (s, a) => apply!(optimizer, (s, a), G - Q(s, a)))
        end
    end
end

function update!(learner::TDLearner{<:AbstractQApproximator,:SARS}, transitions)
    states, actions, rewards, terminals, next_states = transitions
    n, γ, Q, optimizer = learner.n, learner.γ, learner.approximator, learner.optimizer

    if length(terminals) > 0 && terminals[end]
        @views gains = discount_rewards(rewards[max(end - n, 1):end], γ)  # n starts with 0
        for (i, G) in enumerate(gains)
            @views s, a = states[end-length(gains)+i], actions[end-length(gains)+i]
            update!(Q, (s, a) => apply!(optimizer, (s, a), G - Q(s, a)))
        end
    else
        if length(states) ≥ (n + 1)  # n starts with 0
            @views s, a, s′ = states[end-n], actions[end-n], next_states[end]
            @views G = discount_rewards_reduced(rewards[end-n:end], γ) +
                       γ^(n + 1) * maximum(Q(s′))  # n starts with 0
            update!(Q, (s, a) => apply!(optimizer, (s, a), G - Q(s, a)))
        end
    end
end

function update!(
    learner::T,
    target_learner::T,
    transitions,
) where {T<:TDLearner{<:AbstractQApproximator,:SARS}}
    states, actions, rewards, terminals, next_states = transitions
    n, γ, Q, Qₜ, optimizer = learner.n,
        learner.γ,
        learner.approximator,
        target_learner.approximator,
        learner.optimizer

    if length(terminals) > 0 && terminals[end]
        @views gains = discount_rewards(rewards[max(end - n, 1):end], γ)  # n starts with 0
        for (i, G) in enumerate(gains)
            @views s, a = states[end-length(gains)+i], actions[end-length(gains)+i]
            update!(Q, (s, a) => apply!(optimizer, (s, a), G - Q(s, a)))
        end
    else
        if length(states) ≥ (n + 1)  # n starts with 0
            @views s, a, s′ = states[end-n], actions[end-n], next_states[end]
            @views G = discount_rewards_reduced(rewards[end-n:end], γ) +
                       γ^(n + 1) * Qₜ(s′, argmax(Q(s′)))
            update!(Q, (s, a) => apply!(optimizer, (s, a), G - Q(s, a)))
        end
    end
end

function priority(transition, learner::TDLearner{<:AbstractQApproximator,:SARS})
    s, a, r, d, s′ = transition
    γ, Q, opt = learner.γ, learner.approximator, learner.optimizer
    error = d ? apply!(opt, (s, a), r - Q(s, a)) :
            apply!(opt, (s, a), r + γ^(learner.n + 1) * maximum(Q(s′)) - Q(s, a))
    abs(error)
end

function update!(
    learner::TDLearner{<:AbstractQApproximator,:SARS},
    model::Union{TimeBasedSampleModel,ExperienceBasedSampleModel};
    plan_step = 1,
)
    for _ = 1:plan_step
        # @assert learner.n == 0 "n must be 0 here"
        transitions = extract_transitions(model, learner)
        if !isnothing(transitions)
            update!(learner, transitions)
        end
    end
end

function extract_transitions(
    model::Union{ExperienceBasedSampleModel,TimeBasedSampleModel},
    learner::TDLearner{<:AbstractQApproximator,:SARS},
)
    if length(model.experiences) > 0
        map(Base.vect, sample(model))
    else
        nothing
    end
end

function update!(
    learner::TDLearner{<:AbstractQApproximator,:SARS},
    model::PrioritizedSweepingSampleModel;
    plan_step = 1,
)
    for _ = 1:plan_step
        # @assert learner.n == 0 "n must be 0 here"
        transitions = extract_transitions(model, learner)
        if !isnothing(transitions)
            update!(learner, transitions)
            s, _, _, _, _ = transitions
            s = s[]  # length(s) is assumed to be 1
            for (s̄, ā, r̄, d̄) in model.predecessors[s]
                P = priority((s̄, ā, r̄, d̄, s), learner)
                if P ≥ model.θ
                    model.PQueue[(s̄, ā)] = P
                end
            end
        end
    end
end

function extract_transitions(
    model::PrioritizedSweepingSampleModel,
    learner::TDLearner{<:AbstractQApproximator,:SARS},
)
    if length(model.PQueue) > 0
        map(Base.vect, sample(model))
    else
        nothing
    end
end

function extract_transitions(
    buffer::EpisodeTurnBuffer,
    learner::TDLearner{<:AbstractQApproximator,:SARS},
)
    n = learner.n
    if length(buffer) > 0
        @views (
            states = state(buffer)[max(1, end - n - 1):end-1],
            actions = action(buffer)[max(1, end - n - 1):end-1],
            rewards = reward(buffer)[max(1, end - n - 1)+1:end],
            terminals = terminal(buffer)[max(1, end - n - 1)+1:end],
            next_states = state(buffer)[max(1, end - n - 1)+1:end],
        )
    else
        nothing
    end
end

#####
# TDλReturnLearner
#####
"""
    TDλReturnLearner(;approximator::Tapp, γ::Float64 = 1.0, α::Float64, λ::Float64)
"""
Base.@kwdef struct TDλReturnLearner{Tapp<:AbstractApproximator} <: AbstractLearner
    approximator::Tapp
    γ::Float64 = 1.0
    α::Float64
    λ::Float64
end

function extract_transitions(
    buffer::EpisodeTurnBuffer,
    learner::TDλReturnLearner{<:AbstractVApproximator},
)
    if isfull(buffer)
        @views (
            states = state(buffer)[1:end-1],
            rewards = reward(buffer)[2:end],
            terminals = terminal(buffer)[2:end],
            next_states = state(buffer)[2:end],
        )
    else
        nothing
    end
end

function update!(learner::TDλReturnLearner, transition)
    λ, γ, V, α = learner.λ, learner.γ, learner.approximator, learner.α
    states, rewards, terminals, next_states = transition
    T = length(states)
    for t = 1:T
        G = 0.0
        for n = 1:(T-t)
            G += λ^(n - 1) *
                 (discount_rewards_reduced(@view(rewards[t:t+n-1]), γ) +
                  γ^n * V(next_states[t+n-1]))
        end
        G *= 1 - λ
        G += λ^(T - t) *
             (discount_rewards_reduced(@view(rewards[t:T]), γ) +
              γ^(T - t + 1) * V(next_states[T]))
        sₜ = states[t]
        update!(V, sₜ => α * (G - V(sₜ)))
    end
end


using ReinforcementLearning
using StableRNGs
using Flux
using Flux.Losses
using Nature
# using Serialization
using TensorBoardLogger, Logging
using Infiltrator
using Dates
using Serialization

MAX_STEPS=1_000_000

N_STARTING_PLAYERS = 36
UPDATE_FREQ = 128
WORLD_SIZE = (64, 64)
CLIP = 0.1f0

trial_id = ""

timestamp = Dates.format(now(),"yyyy-mm-dd HH:MM:SS")

function RL.Experiment(
    ::Val{:JuliaRL},
    ::Val{:PPO},
    ::Val{:Nature},
    ::Nothing;
    save_dir = nothing,
    seed = 123,
)

    rng = StableRNG(seed)
    env = NatureEnv(num_starting_players=N_STARTING_PLAYERS,
                    world_size=WORLD_SIZE,
                    episode_len=500,
                    window=3,
                    vocab_size=5,
                    food_generators=[
                        FoodGen([15,15],[60,60]),
                        FoodGen([45,45],[60,60]),
                       ])
    RLBase.reset!(env)

    global trial_id = "$timestamp clip=$CLIP uf=$UPDATE_FREQ maxsteps=$MAX_STEPS ws=$(WORLD_SIZE[1:2]) nf=$(length(env.food_generators))"

    runs = readdir("tensorboard_logs/")
    for (i, run) in enumerate(runs)
        println(i, ". ", run)
    end
    println("n. (NEW RUN)")
    print("Select a run: ")
    trial = readline()
    if trial == "n"
        println("Default name: $trial_id")
        print(  "Enter new name (empty for default) ")
        name = readline()
        if name != ""
            global trial_id = name
        end
        agents = false
        println("Creating new run: $trial_id")
    else
        global trial_id = runs[parse(Int, trial)]
        println("Resuming run: $trial_id")
        agents = deserialize("checkpoints/$trial_id.jls")
    end

    print("Enable TensorBoard Logging? y/n: ")
    if readline() == "y"
        disable_logging(Logging.Debug)
        lg=TBLogger("tensorboard_logs/$trial_id", tb_append)
        global_logger(lg)
    else
        disable_logging(Logging.Info)
        global_logger(ConsoleLogger())
    end

    # We don't build agents at agents = false because
    # we don't want to compile between user inputs
    if agents == false
        agents = build_MultiPPOManager(env, N_STARTING_PLAYERS, UPDATE_FREQ, CLIP)
    end

    stop_condition = StopAfterStep(MAX_STEPS, is_show_progress=!haskey(ENV, "CI"))
    hook = NatureHook(env, trial_id, MAX_STEPS)

    Experiment(agents, env, stop_condition, hook, "# PPO with Nature")
end

ex = E`JuliaRL_PPO_Nature`
run(ex)
serialize("checkpoints/$trial_id.jls", ex.policy)
# step_through_env(ex.env, policy)


export BrickletEnergyMonitorEnergyData
struct BrickletEnergyMonitorEnergyData
    voltage::Integer
    current::Integer
    energy::Integer
    real_power::Integer
    apparent_power::Integer
    reactive_power::Integer
    power_factor::Integer
    frequency::Integer
end

export BrickletEnergyMonitorWaveformLowLevel
struct BrickletEnergyMonitorWaveformLowLevel
    waveform_chunk_offset::Integer
    waveform_chunk_data::Vector{Integer}
end

export BrickletEnergyMonitorTransformerStatus
struct BrickletEnergyMonitorTransformerStatus
    voltage_transformer_connected::Bool
    current_transformer_connected::Bool
end

export BrickletEnergyMonitorTransformerCalibration
struct BrickletEnergyMonitorTransformerCalibration
    voltage_ratio::Integer
    current_ratio::Integer
    phase_shift::Integer
end

export BrickletEnergyMonitorEnergyDataCallbackConfiguration
struct BrickletEnergyMonitorEnergyDataCallbackConfiguration
    period::Integer
    value_has_to_change::Bool
end

export BrickletEnergyMonitorSPITFPErrorCount
struct BrickletEnergyMonitorSPITFPErrorCount
    error_count_ack_checksum::Integer
    error_count_message_checksum::Integer
    error_count_frame::Integer
    error_count_overflow::Integer
end

export BrickletEnergyMonitorIdentity
struct BrickletEnergyMonitorIdentity
    uid::String
    connected_uid::String
    position::Char
    hardware_version::Vector{Integer}
    firmware_version::Vector{Integer}
    device_identifier::Integer
end

export BrickletEnergyMonitor
"""
Measures Voltage, Current, Energy, Real/Apparent/Reactive Power, Power Factor and Frequency
"""
mutable struct BrickletEnergyMonitor <: TinkerforgeDevice
    ipcon::IPConnection
    deviceInternal::PyObject

    """
    Creates an object with the unique device ID *uid* and adds it to
    the IP Connection *ipcon*.
    """
    function BrickletEnergyMonitor(uid::String, ipcon::IPConnection)
        package = pyimport("tinkerforge.bricklet_energy_monitor")
        deviceInternal = package.BrickletEnergyMonitor(uid, ipcon.ipconInternal)

        return new(ipcon, deviceInternal)
    end
end

export get_energy_data
"""
    $(SIGNATURES)

Returns all of the measurements that are done by the Energy Monitor Bricklet.

* Voltage RMS
* Current RMS
* Energy (integrated over time)
* Real Power
* Apparent Power
* Reactive Power
* Power Factor
* Frequency (AC Frequency of the mains voltage)

The frequency is recalculated every 6 seconds.

All other values are integrated over 10 zero-crossings of the voltage sine wave.
With a standard AC mains voltage frequency of 50Hz this results in a 5 measurements
per second (or an integration time of 200ms per measurement).

If no voltage transformer is connected, the Bricklet will use the current waveform
to calculate the frequency and it will use an integration time of
10 zero-crossings of the current waveform.
"""
function get_energy_data(device::BrickletEnergyMonitor)
    return device.deviceInternal.get_energy_data()
end

export reset_energy
"""
    $(SIGNATURES)

Sets the energy value (see :func:`Get Energy Data`) back to 0Wh.
"""
function reset_energy(device::BrickletEnergyMonitor)
    device.deviceInternal.reset_energy()
end

export get_waveform_low_level
"""
    $(SIGNATURES)

Returns a snapshot of the voltage and current waveform. The values
in the returned array alternate between voltage and current. The data from
one getter call contains 768 data points for voltage and current, which
correspond to about 3 full sine waves.

The voltage is given with a resolution of 100mV and the current is given
with a resolution of 10mA.

This data is meant to be used for a non-realtime graphical representation of
the voltage and current waveforms.
"""
function get_waveform_low_level(device::BrickletEnergyMonitor)
    return device.deviceInternal.get_waveform_low_level()
end

export get_transformer_status
"""
    $(SIGNATURES)

Returns *true* if a voltage/current transformer is connected to the Bricklet.
"""
function get_transformer_status(device::BrickletEnergyMonitor)
    return device.deviceInternal.get_transformer_status()
end

export set_transformer_calibration
"""
    $(SIGNATURES)

Sets the transformer ratio for the voltage and current transformer in 1/100 form.

Example: If your mains voltage is 230V, you use 9V voltage transformer and a
1V:30A current clamp your voltage ratio is 230/9 = 25.56 and your current ratio
is 30/1 = 30.

In this case you have to set the values 2556 and 3000 for voltage ratio and current
ratio.

The calibration is saved in non-volatile memory, you only have to set it once.

Set the phase shift to 0. It is for future use and currently not supported by the Bricklet.
"""
function set_transformer_calibration(device::BrickletEnergyMonitor, voltage_ratio, current_ratio, phase_shift)
    device.deviceInternal.set_transformer_calibration(voltage_ratio, current_ratio, phase_shift)
end

export get_transformer_calibration
"""
    $(SIGNATURES)

Returns the transformer calibration as set by :func:`Set Transformer Calibration`.
"""
function get_transformer_calibration(device::BrickletEnergyMonitor)
    return device.deviceInternal.get_transformer_calibration()
end

export calibrate_offset
"""
    $(SIGNATURES)

Calling this function will start an offset calibration. The offset calibration will
integrate the voltage and current waveform over a longer time period to find the 0
transition point in the sine wave.

The Bricklet comes with a factory-calibrated offset value, you should not have to
call this function.

If you want to re-calibrate the offset we recommend that you connect a load that
has a smooth sinusoidal voltage and current waveform. Alternatively you can also
short both inputs.

The calibration is saved in non-volatile memory, you only have to set it once.
"""
function calibrate_offset(device::BrickletEnergyMonitor)
    device.deviceInternal.calibrate_offset()
end

export set_energy_data_callback_configuration
"""
    $(SIGNATURES)

The period is the period with which the :cb:`Energy Data`
callback is triggered periodically. A value of 0 turns the callback off.

If the `value has to change`-parameter is set to true, the callback is only
triggered after the value has changed. If the value didn't change within the
period, the callback is triggered immediately on change.

If it is set to false, the callback is continuously triggered with the period,
independent of the value.
"""
function set_energy_data_callback_configuration(device::BrickletEnergyMonitor, period, value_has_to_change)
    device.deviceInternal.set_energy_data_callback_configuration(period, value_has_to_change)
end

export get_energy_data_callback_configuration
"""
    $(SIGNATURES)

Returns the callback configuration as set by
:func:`Set Energy Data Callback Configuration`.
"""
function get_energy_data_callback_configuration(device::BrickletEnergyMonitor)
    return device.deviceInternal.get_energy_data_callback_configuration()
end

export get_spitfp_error_count
"""
    $(SIGNATURES)

Returns the error count for the communication between Brick and Bricklet.

The errors are divided into

* ACK checksum errors,
* message checksum errors,
* framing errors and
* overflow errors.

The errors counts are for errors that occur on the Bricklet side. All
Bricks have a similar function that returns the errors on the Brick side.
"""
function get_spitfp_error_count(device::BrickletEnergyMonitor)
    return device.deviceInternal.get_spitfp_error_count()
end

export set_bootloader_mode
"""
    $(SIGNATURES)

Sets the bootloader mode and returns the status after the requested
mode change was instigated.

You can change from bootloader mode to firmware mode and vice versa. A change
from bootloader mode to firmware mode will only take place if the entry function,
device identifier and CRC are present and correct.

This function is used by Brick Viewer during flashing. It should not be
necessary to call it in a normal user program.
"""
function set_bootloader_mode(device::BrickletEnergyMonitor, mode)
    return device.deviceInternal.set_bootloader_mode(mode)
end

export get_bootloader_mode
"""
    $(SIGNATURES)

Returns the current bootloader mode, see :func:`Set Bootloader Mode`.
"""
function get_bootloader_mode(device::BrickletEnergyMonitor)
    return device.deviceInternal.get_bootloader_mode()
end

export set_write_firmware_pointer
"""
    $(SIGNATURES)

Sets the firmware pointer for :func:`Write Firmware`. The pointer has
to be increased by chunks of size 64. The data is written to flash
every 4 chunks (which equals to one page of size 256).

This function is used by Brick Viewer during flashing. It should not be
necessary to call it in a normal user program.
"""
function set_write_firmware_pointer(device::BrickletEnergyMonitor, pointer)
    device.deviceInternal.set_write_firmware_pointer(pointer)
end

export write_firmware
"""
    $(SIGNATURES)

Writes 64 Bytes of firmware at the position as written by
:func:`Set Write Firmware Pointer` before. The firmware is written
to flash every 4 chunks.

You can only write firmware in bootloader mode.

This function is used by Brick Viewer during flashing. It should not be
necessary to call it in a normal user program.
"""
function write_firmware(device::BrickletEnergyMonitor, data)
    return device.deviceInternal.write_firmware(data)
end

export set_status_led_config
"""
    $(SIGNATURES)

Sets the status LED configuration. By default the LED shows
communication traffic between Brick and Bricklet, it flickers once
for every 10 received data packets.

You can also turn the LED permanently on/off or show a heartbeat.

If the Bricklet is in bootloader mode, the LED is will show heartbeat by default.
"""
function set_status_led_config(device::BrickletEnergyMonitor, config)
    device.deviceInternal.set_status_led_config(config)
end

export get_status_led_config
"""
    $(SIGNATURES)

Returns the configuration as set by :func:`Set Status LED Config`
"""
function get_status_led_config(device::BrickletEnergyMonitor)
    return device.deviceInternal.get_status_led_config()
end

export get_chip_temperature
"""
    $(SIGNATURES)

Returns the temperature as measured inside the microcontroller. The
value returned is not the ambient temperature!

The temperature is only proportional to the real temperature and it has bad
accuracy. Practically it is only useful as an indicator for
temperature changes.
"""
function get_chip_temperature(device::BrickletEnergyMonitor)
    return device.deviceInternal.get_chip_temperature()
end

export reset
"""
    $(SIGNATURES)

Calling this function will reset the Bricklet. All configurations
will be lost.

After a reset you have to create new device objects,
calling functions on the existing ones will result in
undefined behavior!
"""
function reset(device::BrickletEnergyMonitor)
    device.deviceInternal.reset()
end

export write_uid
"""
    $(SIGNATURES)

Writes a new UID into flash. If you want to set a new UID
you have to decode the Base58 encoded UID string into an
integer first.

We recommend that you use Brick Viewer to change the UID.
"""
function write_uid(device::BrickletEnergyMonitor, uid)
    device.deviceInternal.write_uid(uid)
end

export read_uid
"""
    $(SIGNATURES)

Returns the current UID as an integer. Encode as
Base58 to get the usual string version.
"""
function read_uid(device::BrickletEnergyMonitor)
    return device.deviceInternal.read_uid()
end

export get_identity
"""
    $(SIGNATURES)

Returns the UID, the UID where the Bricklet is connected to,
the position, the hardware and firmware version as well as the
device identifier.

The position can be 'a', 'b', 'c', 'd', 'e', 'f', 'g' or 'h' (Bricklet Port).
A Bricklet connected to an :ref:`Isolator Bricklet <isolator_bricklet>` is always at
position 'z'.

The device identifier numbers can be found :ref:`here <device_identifier>`.
|device_identifier_constant|
"""
function get_identity(device::BrickletEnergyMonitor)
    return device.deviceInternal.get_identity()
end

export get_waveform
"""
Returns a snapshot of the voltage and current waveform. The values
in the returned array alternate between voltage and current. The data from
one getter call contains 768 data points for voltage and current, which
correspond to about 3 full sine waves.

The voltage is given with a resolution of 100mV and the current is given
with a resolution of 10mA.

This data is meant to be used for a non-realtime graphical representation of
the voltage and current waveforms.
"""
function get_waveform(device::BrickletEnergyMonitor)

    waveform_length = 1536

    lock(device.stream_lock) do
        ret = get_waveform_low_level(device, )

        if ret.waveform_chunk_offset == (1 << 16) - 1 # maximum chunk offset -> stream has no data
            waveform_length = 0
            waveform_out_of_sync = false
            waveform_data = ()
        else
            waveform_out_of_sync = ret.waveform_chunk_offset != 0
        end

        waveform_data = ret.waveform_chunk_data

        while !waveform_out_of_sync && length(waveform_data) < waveform_length
            ret = get_waveform_low_level(device, )
            waveform_out_of_sync = ret.waveform_chunk_offset != length(waveform_data)
            waveform_data += ret.waveform_chunk_data
        end

        if waveform_out_of_sync # discard remaining stream to bring it back in-sync
            while ret.waveform_chunk_offset + 30 < waveform_length
                ret = get_waveform_low_level(device, )
            end

            throw(TinkerforgeStreamOutOfSyncError("Waveform stream is out-of-sync"))
        end
    end

        return waveform_data[:waveform_length]
end

export register_callback
"""
Registers the given *function* with the given *callback_id*.
"""
function register_callback(device::BrickletEnergyMonitor, callback_id, function_)
    if isnothing(function_)
        device.registered_callbacks.pop(callback_id, None)
    else
        device.registered_callbacks[callback_id] = function_
    end
end


using BGEN
using Random
using MLJBase
using MLJLinearModels
using TMLE
using JLD2

genotypesfile = joinpath("data", "genotypes.csv")
confoundersfile = joinpath("data", "confounders.csv")
queryfile = joinpath("config", "query.toml")
continuous_phenotypefile = joinpath("data", "continuous_phenotypes.csv")
binary_phenotypefile = joinpath("data", "binary_phenotypes.csv")
phenotypelist_file = joinpath("data", "phen_list.csv")

function build_query_file(;path=queryfile, threshold=0.9, het="AG")
    queries = Dict(
        "threshold" => threshold,
        "SNPS" => Dict(
            "RSID_10"  => BGEN.datadir("example.8bits.bgen"),
            "RSID_100" => BGEN.datadir("example.8bits.bgen")),
        "QUERY_1" => Dict(
            "RSID_10"  => "GG -> "*het,
            "RSID_100" => "GG -> "*het),
        "QUERY_2" => Dict(
            "RSID_10"  => "GG -> "*het,
            "RSID_100" => "GG -> AA"),
        "QUERY_3" => Dict(
            "RSID_10"  => "TT -> "*het,
            "RSID_100" => "GG -> AA")
    )
    open(path, "w") do io
        TOML.print(io, queries)
    end
end

function build_ate_query_file(;path=queryfile, threshold=0.9)
    queries = Dict(
        "threshold" => threshold,
        "SNPS" => Dict("RSID_10"  => BGEN.datadir("example.8bits.bgen")),
        "QUERY_1" => Dict("RSID_10"  => "GG -> AG"),
        "QUERY_2" => Dict("RSID_10"  => "GG -> AA")
    )
    open(path, "w") do io
        TOML.print(io, queries)
    end
end

function clean_hdf5estimates_files(prefix)
    rm(string(prefix, "_batch_1_Real.hdf5"))
    rm(string(prefix, "_batch_1_Bool.hdf5"))
    rm(string(prefix, "_batch_2_Bool.hdf5"))
end

function build_results_files(grm_ids, prefix)
    rng = Xoshiro(0)
    n = size(grm_ids, 1)
    T = (t₁=categorical(rand(rng, [0, 1], n)),)
    W = (w₁=rand(rng, n), w₂=rand(rng, n))
    height = rand(rng, n)
    bmi = 2convert(Vector{Float64}, T.t₁) + 0.2rand(rng, n)
    cancer = categorical(rand(rng, [0,1], n))
    query_1 = Query(case=(t₁=1,), control=(t₁=0,), name="QUERY_1")
    query_2 = Query(case=(t₁=0,), control=(t₁=1,), name="QUERY_2")

    # Continuous single batch
    cont_path = string(prefix, "_batch_1_Real.hdf5")
    tmle_reg = TMLEstimator(
            LinearRegressor(),
            LogisticClassifier(),
            query_1,
            query_2
    )
    TMLE.fit(tmle_reg, T, W, (height=height, bmi=bmi);
        verbosity=0,
        cache=false,
        callbacks=[TMLEEpistasis.JLD2Saver(cont_path, true)])
    jldopen(cont_path, "a") do io
            io["SAMPLE_IDS"] = Dict(
                "height" => string.(grm_ids.SAMPLE_ID),
                "bmi" => string.(grm_ids.SAMPLE_ID)
            )
    end
    # Binary single batch
    bin_path = string(prefix, "_batch_1_Bool.hdf5")
    tmle_bin = TMLEstimator(
        LogisticClassifier(),
        LogisticClassifier(),
        query_1,
        query_2)
    TMLE.fit(tmle_bin, T, W, (cancer=cancer,);
        verbosity=0,
        cache=false,
        callbacks=[TMLEEpistasis.JLD2Saver(bin_path, true)])
    jldopen(bin_path, "a") do io
        io["SAMPLE_IDS"] = Dict(
            "cancer" => string.(grm_ids.SAMPLE_ID)
        )
    end

    # Empty file that may be output by the tmle procedure
    jldopen(string(prefix, "_batch_2_Bool.hdf5"), "a") do io
        JLD2.Group(io, "TMLEREPORTS")
    end
end

function test_queries(queries, expected_queries)
    for (i, query) in enumerate(queries)
        @test query.case == expected_queries[i].case
        @test query.control == expected_queries[i].control
        @test query.name == expected_queries[i].name
    end
end

function _test_wiped_data(mach)
    @test mach.data == ()
    @test :data ∉ mach.cache
    @test mach.args == ()
end

function test_data_has_been_removed(tmle_mach)
    _test_wiped_data(tmle_mach)

    for submach in tmle_mach.report.machines
        _test_wiped_data(submach)
    end

    for submach in report(tmle_mach).G.machines
        _test_wiped_data(submach)
    end

    for submach in report(tmle_mach).Q̅.machines
        _test_wiped_data(submach)
    end
end


"""
    UnravelledLayout

Parameters are

- `x` (a function to get the value on x)
- `y` (a function to get the value on y)

Both of these functions *must* accept a unipartite network as input, and return
a dictionary with species and a single numerical value as output. Note that `x`
and/or `y` can be `λ`s.
"""
struct UnravelledLayout{TX<:Function,TY<:Function}
    x::TX
    y::TY
end

"""
    UnravelledLayout(;x=omnivory, y=fractional_trophic_level)

By default, the two functions used are omnivory and fractional trophic level,
but this can be changed.
"""
UnravelledLayout(;x=omnivory, y=fractional_trophic_level) = UnravelledLayout(x, y)

"""
    position!(LA::LT, L::Dict{K,NodePosition}, N::T) where {LT <: UnravelledLayout, T <: AbstractEcologicalNetwork} where {K}

Position species according to the function defined in the `UnravelledLayout`.
"""
function position!(LA::LT, L::Dict{K,NodePosition}, N::T) where {LT <: UnravelledLayout, T <: AbstractEcologicalNetwork} where {K}
  X = LA.x(N)
  Y = LA.y(N)
  for (i, s) in enumerate(species(N))
    L[s].x = X[s]
    L[s].y = Y[s]
  end
end

"""
    NestedBipartiteLayout

Parameters are

- `align` (whether the two levels should be centered together)
- `relative` (whether the two levels should occupy a length equal to their relative richness)
- `spread` (the distance between the two)

Note that to see the effect of `spread`, you may have to use `aspectratio=1`; if
not, the spacing between levels will be determined by the dimensions of the
plot.
"""
struct NestedBipartiteLayout
    align::Bool
    relative::Bool
    spread::Float64
end

"""
    NestedBipartiteLayout()

By default, a `NestedBipartiteLayout` is aligned, centered, and with a spread of
`1.0`.
"""
NestedBipartiteLayout() = NestedBipartiteLayout(true, true, 1.0)
NestedBipartiteLayout(spread::Float64) = NestedBipartiteLayout(true, true, spread)

"""
    position!(LA::NestedBipartiteLayout, L::Dict{K,NodePosition}, N::T) where {T <: AbstractBipartiteNetwork} where {K}

Rank species according to their degree.
"""
function position!(LA::NestedBipartiteLayout, L::Dict{K,NodePosition}, N::T) where {T <: AbstractBipartiteNetwork} where {K}
    r_top = ordinalrank(collect(values(degree(N; dims=1))); rev=true).-1
    r_bot = ordinalrank(collect(values(degree(N; dims=2))); rev=true).-1

    r_top = r_top./maximum(r_top)
    r_bot = r_bot./maximum(r_bot)
    if LA.relative
        if richness(N; dims=2)>richness(N; dims=1)
            r_bot = r_bot .* (richness(N; dims=2)./richness(N; dims=1))
        end
        if richness(N; dims=1)>richness(N; dims=2)
            r_top = r_top .* (richness(N; dims=1)./richness(N; dims=2))
        end
    end

    if LA.align
        r_bot = r_bot .+ (0.5 - (maximum(r_bot)-minimum(r_bot))/2.0)
        r_top = r_top .+ (0.5 - (maximum(r_top)-minimum(r_top))/2.0)
    end

    d_bot = Dict(zip(keys(degree(N; dims=2)), r_bot))
    d_top = Dict(zip(keys(degree(N; dims=1)), r_top))

    d_all = merge(d_bot, d_top)

    for s in species(N)
        L[s].x = d_all[s]
        L[s].y = LA.spread * L[s].y
    end
end


# %j is number of days in year
time_step = Dict('s' => (Dates.Millisecond, 3), 'S' => (Dates.Second, 2),
                 'M' => (Dates.Minute, 2), 'H' => (Dates.Hour, 2),
                 'd' => (Dates.Day, 2), 'j' => (Dates.Day, 3),
                 'U' => (Dates.Month, 1), 'u' => (Dates.Month, 1), 'm' => (Dates.Month, 2),
                 'Y' => (Dates.Year, 1), 'y' => (Dates.Year, 2))
date_steps = ['s', 'S', 'M', 'H',
              'd', 'j', 'U', 'u', 'm', 'Y', 'y']
# ^^ Not using keys(time_step) because dict is not ordered. #FIXME change to ordict

function smaller_step{T <: String}(pattern::T)
    # list of date steps in increasing order
    # Loop over all the possibles step and
    # return the smallest one.
    for step in date_steps
        if step in pattern
            return step
        end
    end
    error("Date pattern $(pattern) not encountered")
end

function extract_date{T <: String}(text::T)
    # Look up for %X patterns to convert to date
    dates = collect(matchall(r"%.", text)) |> join
    replace(dates, "%", "")
end

function check_steps{T <: String}(text::T)
    # Check whether all %pattern are dates
    all(x->x in date_steps, text)
end

"""
    dates_directories(urlpattern)

Extracts the dates patterns from an url pattern
# Example
```jldoctest
julia> dates_directories("http://this.is/my/%Y/%m/%d/ofBirth.txt")
"Ymd"
```
"""
function dates_directories{T <: String}(urlpattern::T)
    directories = collect(matchall(r"%.*/", urlpattern)) |> join
    extract_date(directories)
end

"""
    create_urls(urlpattern, timeStart, timeEnd)

Generates urls directories where to look for a particular time
range.
# Example
```jldoctest
julia> dt0, dt1 = DateTime(2017, 2, 3), DateTime(2017, 2, 5)
julia> create_urls("http://this.is/my/%Y/%m/%d/ofBirth.txt", dt0, dt1)
3-element Array{String, 1}:
 "http://this.is/my/2017/02/03/"
 "http://this.is/my/2017/02/04/"
 "http://this.is/my/2017/02/05/"
```
"""
function create_urls{T <: String, DT <: Base.Dates.TimeType}(urlpattern::T, timeStart::DT, timeEnd::DT)
    #NOTE Should the steprange preceed over urlpattern?

    # Find directory step range
    dates = dates_directories(urlpattern)
    if check_steps(dates)
        step = smaller_step(dates)
    end

    # Only extract the directory path
    urlpattern_dir = matchall(r".*/", urlpattern)[1]
    function convert_url(tstep)
        # Multiple replace at once, when %j using the custom format.
        # Need to take the second element of each match (x[2]) because "%" is the first one
        replace(urlpattern_dir, Regex("%[$(join(keys(time_step)))]"),
                x -> x == "%j" ? dec(Dates.dayofyear(tstep), 3) : Dates.format(tstep, "$(x[2])" ^ time_step[x[2]][2]))
    end

    dt = timeStart:time_step[step][1](1):timeEnd
    map(convert_url, dt)

end


export
    Simulation,
    simulate

"""
    Simulation

Describe a PDSampling: information about the initial point, the time of the
simulation, the function to sample from an IPP, etc.
"""
struct Simulation
    x0::Vector{Real}       # Starting point
    v0::Vector{Real}       # Starting velocity
    T::Real                # Simulation time
    nextevent::Function    # Appropriate simulation for first arrival time
    gll::Function          # Gradient of Log Lik (potentially CV)
    nextboundary::Function # Where/When is the next boundary hit
    lambdaref::Real        # Refreshment rate
    algname::String        # BPS, ZZ, GBPS
    # derived
    dim::Int               # dimensionality
    # optional named arguments
    mass::Matrix{Real}     # mass matrix (preconditioner)
    blocksize::Int         # increment the storage by blocks
    maxsimtime::Real       # max. simulation time (s)
    maxsegments::Int       # max. num. of segments
    maxgradeval::Int       # max. num. grad. evals
    refresh!::Function     # refreshment function (TODO: examples)

    # constructor
    function Simulation(x0, v0, T, nextevent, gradloglik, nextboundary,
                lambdaref=1.0, algname="BPS";
                mass=diagm(0=>ones(0)), blocksize=1000, maxsimtime=4e3,
                maxsegments=1_000_000, maxgradeval=100_000_000,
                refresh! = refresh_global! )
        # check none of the default named arguments went through
        @assert !(x0 == :undefined || v0 == :undefined || T == :undefined ||
                  nextevent == :undefined || gradloglik == :undefined || nextboundary == :undefined) "Essential arguments undefined"
        # basic check to see that things are reasonable
        @assert length(x0) == length(v0) > 0 "Inconsistent arguments"
        @assert T > 0.0 "Simulation time must be positive"
        @assert lambdaref >= 0.0 "Refreshment rate must be >= 0"
        ALGNAME = uppercase(algname)
        @assert (ALGNAME ∈ ["BPS", "ZZ","GBPS"]) "Unknown algorithm <$algname>"
        new( x0, v0, T,
             nextevent, gradloglik, nextboundary,
             lambdaref, ALGNAME, length(x0),
             mass, blocksize, maxsimtime,
             maxsegments, maxgradeval, refresh! )
    end
end
# Constructor with named arguments
function Simulation(;
            x0 = :undefined,
            v0 = :undefined,
            T = :undefined,
            nextevent = :undefined,
            gradloglik = :undefined,
            nextboundary = :undefined,
            lambdaref = 1.0,
            algname = "BPS",
            mass = diagm(0=>ones(0)),
            blocksize = 1000,
            maxsimtime = 4e3,
            maxsegments = Int(1e6),
            maxgradeval = Int(1e8),
            refresh! = refresh_global! )
    # calling the unnamed constructor
    Simulation( x0, v0, T,
                nextevent,
                gradloglik,
                nextboundary,
                lambdaref,
                algname;
                mass = mass,
                blocksize = blocksize,
                maxsimtime = maxsimtime,
                maxsegments = maxsegments,
                maxgradeval = maxgradeval,
                refresh! = refresh! )
end

"""
    simulate(sim)

Launch a PD simulation defined in the sim variable. Return the corresponding
Path and a dictionary of indicators (clocktime, ...).
"""
function simulate(sim::Simulation)
    # keep track of how much time we've been going for
    start = time()
    # dimensionality
    d = sim.dim
    # initial states
    x, v = copy(sim.x0), copy(sim.v0)
    # time counter, and segment counter
    t, i = 0.0, 1
    # counters for the number of effective loops and
    # for the number of evaluations of the gradient
    # this will be higher than the number of segments i
    lcnt, gradeval = 0, 0

    # storing by blocks of blocksize nodes at the time
    blocksize = sim.blocksize

    # storing xs as a single column for efficient resizing
    xs, ts = zeros(d*blocksize), zeros(blocksize)
    # store initial point
    xs[1:d] = x
    # mass matrix?
    mass = copy(sim.mass) # store it here as we may want to adapt it

    # check if nextevent takes 2 or 3 parameters
    nevtakes2 = (length(methods(sim.nextevent).ms[1].sig.parameters)-1) == 2

    # compute current reference bounce time
    lambdaref = sim.lambdaref
    tauref = (lambdaref>0.0) ? randexp() / lambdaref : Inf
    # Compute time to next boundary + normal
    (taubd, normalbd) = sim.nextboundary(x, v)

    # keep track of how many refresh events
    nrefresh = 0
    # keep track of how many boundary hits
    nboundary = 0
    # keep track of how many standard bounces
    nbounce = 0

    while (t < sim.T) && (gradeval < sim.maxgradeval)
        # increment the counter to keep track of the number of effective loops
        lcnt += 1
        # simulate first arrival from IPP
        bounce = nevtakes2 ? sim.nextevent(x, v) : sim.nextevent(x, v, tauref)
        # find next event (unconstrained case taubd=NaN (ignored))
        tau = min(bounce.tau, taubd, tauref)
        # standard bounce
        if tau == bounce.tau
            # there will be an evaluation of the gradient
            gradeval += 1
            # updating position/time
            t += tau
            x += tau*v
            # exploiting the memoryless property
            tauref -= tau
            # ---- BOUNCE ----
            g = sim.gll(x)
            if bounce.dobounce(g, v) # e.g.: thinning, acc/rej
                # if accept
                nbounce += 1
                if sim.algname == "BPS"
                    # if a mass matrix is provided
                    if length(mass) > 0
                        v = reflect_bps!(g, v, mass)
                    # standard BPS bounce
                    else
                        v = reflect_bps!(g,v)
                    end
                elseif sim.algname == "GBPS"
                    v = reflect_gbps(g, v)
                elseif sim.algname == "ZZ"
                    v = reflect_zz!(bounce.flipindex, v)
                end
            else
                # move closer to the boundary/refreshment time
                taubd -= tau
                # we don't need to record when rejecting
                continue
            end
        # hard bounce against boundary
        elseif tau == taubd
            #
            nboundary += 1
            # Record point epsilon from boundary for numerical stability
            x += (tau - 1e-10) * v
            t += tau
            # exploiting the memoryless property
            tauref -= tau
            # ---- BOUNCE (boundary) ----
            if sim.algname ∈ ["BPS", "GBPS"]
                # Specular reflection (possibly with mass matrix)
                if length(mass) > 0
                    v = reflect_bps!(normalbd, v, mass)
                else
                    v = reflect_bps!(normalbd,v)
                end
            elseif sim.algname == "ZZ"
                v = reflect_zz!(find((v.*normalbd).<0.0), v)
            end
        # random refreshment
        else
            #= to be in this part, lambdaref should be greater than 0.0
            because if lambdaref=0 then tauref=Inf. There may be a weird
            corner case in which an infinity filters through which we would
            skip =#
            if !isinf(tau)
                #
                nrefresh += 1
                #
                x += tau*v
                t += tau
                # ---- REFRESH ----
                v = sim.refresh!(v)
                # update tauref
                tauref = randexp()/lambdaref
            end
        end
        # check when the next boundary hit will occur
        (taubd, normalbd) = sim.nextboundary(x, v)
        # increment the counter for the number of segments
        i += 1

        # Increase storage on a per-need basis.
        if mod(i,blocksize)==0
            resize!(xs, length(xs) + d * blocksize)
            resize!(ts, length(ts) + blocksize)
        end

        # Storing path times
        ts[i] = t
         # storing columns vertically, a vector is simpler/cheaper to resize
        xs[((i-1) * d + 1):(i * d)] = x

        # Safety checks to break long loops every 100 iterations
        if mod(lcnt, 100) == 0
            if  (time() - start) > sim.maxsimtime
                println("Max simulation time reached. Stopping")
                break
            end
            if  i > sim.maxsegments
                println("Too many segments generated. Stopping")
                break
            end
        end
    end # End of while loop

    details = Dict(
        "clocktime" => time()-start,
        "ngradeval" => gradeval,
        "nloops"    => lcnt,
        "nsegments" => i,
        "nbounce"   => nbounce,
        "nboundary" => nboundary,
        "nrefresh"  => nrefresh
    )

    (Path(reshape(xs[1:(i*d)], (d,i)), ts[1:i]), details)
end


module Virtuals

using Base.Broadcast: Broadcasted, Extruded
using LinearAlgebra

using Swizzles.ArrayifiedArrays
using Swizzles.Properties

export Virtual, VirtualArray, virtualize, virtual

struct Virtual{T}
    ex
end

function virtual(root, Data::Type)
    ins = instance(Data)
    if ins !== nothing
        return something(ins)
    else
        return Virtual{Data}(root)
    end
end

Base.ndims(::Virtual{T}) where {T} = ndims(T)
Base.ndims(::Virtual{<:Tuple}) = 1
Base.ndims(::Type{Virtual{T}}) where {T} = ndims(T)
Base.ndims(::Type{<:Virtual{<:Tuple}}) = 1
Base.axes(::Virtual{T}) where {T} = ndims(T) == 0 ? () : error() #FIXME return virtual axes!
Base.Broadcast.broadcastable(x::Virtual{<:Union{AbstractArray,Number,Ref,Tuple,Broadcasted}}) = x
Base.eltype(::Virtual{T}) where {T} = eltype(T)

struct VirtualArray{T, N, Data<:AbstractArray{T, N}} <: AbstractArray{T, N}
    ex
end

function virtualtuple(root, Data::Type{<:Tuple})
    data = map(((i, Datum),) -> virtualize(:($root[$i]), Datum), enumerate(Data.parameters))
    return (data...,)
end

function virtualize(root, Data::Type)
    ins = instance(Data)
    if ins !== nothing
        return something(ins)
    else
        return Virtual{Data}(root)
    end
end

function virtualize(root, Data::AbstractArray{T, N}) where {T, N}
    return VirtualArray{T, N, Data}(root)
end

function virtualize(root, ::Type{<:Adjoint{T, Arg}}) where {T, Arg}
    arg = virtualize(:($root.parent), Arg)
    return Adjoint{T, typeof(arg)}(arg)
end

function virtualize(root, ::Type{<:Transpose{T, Arg}}) where {T, Arg}
    arg = virtualize(:($root.parent), Arg)
    return Transpose{T, typeof(arg)}(arg)
end

function virtualize(root, ::Type{<:ArrayifiedArray{T, N, Arg}}) where {T, N, Arg}
    arg = virtualize(:($root.arg), Arg)
    return ArrayifiedArray{T, N, typeof(arg)}(arg)
end

function virtualize(root, Data::Type{<:AbstractArray{T, N}}) where {T, N}
    try
        return Data.instance
    catch Error
        return VirtualArray{T, N, Data}(root)
    end
end

function virtualize(root, ::Type{<:Broadcasted{Style, Axes, F, Args}}) where {Style, Axes, F, Args}
    f = virtual(:($root.f), F)
    args = (map(((i, Arg),) -> virtualize(:($root.args[$i]), Arg), enumerate(Args.parameters))...,)
    axes = (map(((i, Axis),) -> virtualize(:($root.axes[$i]), Axes), enumerate(Axes.parameters))...,)
    return Broadcasted{Style}(f, args, axes)
end

function virtualize(root, ::Type{<:Extruded{Arg, Keeps, Defaults}}) where {Arg, Keeps, Defaults}
    arg = virtualize(:($root.arg), Arg)
    args = (map(((i, K),) -> virtual(:($root.keeps[$i]), K), enumerate(Keeps.parameters))...,)
    axes = (map(((i, D),) -> virtual(:($root.defaults[$i]), D), enumerate(Defaults.parameters))...,)
    return Extruded(arg, keeps, defaults)
end

end


using Word2Vec
using LinearAlgebra

const stopWords = Set{String}([
    "a",
	"an",
	"and",
	"as",
	"based",
	"by",
	"for",
	"from",
	"in",
	"on",
	"of",
    "over",
	"that",
	"the",
	"this",
	"to",
	"via",
	"with",
	"without",
	"is",
	"are",
	"be",
	"you",
	"we",
	"they",
	"it",
	"your",
	"our",
	"their",
	"its",
	"what",
	"when",
	"where",
	"how",
	"do",
	"use",
	"here",
	"there",
	"using",
    "than",
	"、",
	",",
	"，",
	":",
	"：",
	".",
	"。",
	"‧",
	"!",
	"！",
	"?",
	"？",
	";",
	"；",
	"(",
	"（",
	")",
	"）",
	"'",
	"‘",
	"’",
	"\"",
	"「",
	"」",
	"“",
	"”",
	"`",
	"…",
    "</s>",
])

function getNonStopWords(wordVectors::WordVectors)::Tuple{Vector{String}, Dict{String,Int}}
    nonStopWords = Vector{String}()
    nonStopWordIDs = Dict{String,Int}()
    for (wordID, word) in enumerate(wordVectors.vocab)
        if !(word in stopWords)
            push!(nonStopWords, word)
            nonStopWordIDs[word] = wordID
        end
    end
    (nonStopWords, nonStopWordIDs)
end

function normalizeVectors(wordVectors::WordVectors, nonStopWords::Vector{String})::Array{Float64,2}
    numDims, numWords = size(wordVectors)
    numVectors = length(nonStopWords)
    result = Array{Float64,2}(undef, (numDims, numVectors))
    for i in 1:numVectors
        vector = get_vector(wordVectors, nonStopWords[i])
        w = 1.0 / norm(vector)
        result[:,i] = vector * w
    end
    result
end

struct IDPair
    i::Int
    j::Int
end

mutable struct Biterm
    pair::IDPair
    topicID::Int
    gpuFlag::Bool
end

struct NBTMWEParameters
    numTopics::Int
    maxIter::Int
    simThres::Float64
    alpha::Float64
    beta::Float64
    phrases::Vector{String}
    phraseIDs::Dict{String,Int}
    docs::Vector{Vector{Int}}
end

mutable struct NBTMWE
    parameters::NBTMWEParameters
    biterms::Vector{Biterm}
    similarPhraseIDs::Dict{IDPair, Vector{Int}}
    numBitermsOfTopic::Vector{Int}
    numRelatedOfTopic::Vector{Int}
    phraseFreqsOfTopic::Array{Float64,2}
end

function saveNBTMWE(fileName::String, btm::NBTMWE)
    file = open(fileName, "w")
    param = btm.parameters
    numPhrases = length(param.phrases)
    numDocs = length(param.docs)
    numBiterms = length(btm.biterms)
    numSimilar = length(btm.similarPhraseIDs)
    println(file, "$(param.numTopics) $(param.maxIter) $(param.simThres) $(param.alpha) $(param.beta) $numPhrases $numDocs $numBiterms $numSimilar")
    for phraseID in 1:numPhrases
        println(file, param.phrases[phraseID])
    end
    for docID in 1:numDocs
        doc = param.docs[docID]
        docSize = length(doc)
        for i in 1:docSize
            if i > 1
                print(file, " ")
            end
            print(file, doc[i])
        end
        print(file, "\n")
    end
    for bitermID in 1:numBiterms
        biterm = btm.biterms[bitermID]
        i = biterm.pair.i
        j = biterm.pair.j
        topicID = biterm.topicID
        gpuFlag = biterm.gpuFlag
        println(file, "$i $j $topicID $gpuFlag")
    end
    for (pair, relatedIDs) in btm.similarPhraseIDs
        print(file, "$(pair.i) $(pair.j)")
        for id in relatedIDs
            print(file, " $id")
        end
        print(file, "\n")
    end
    for topicID in 1:param.numTopics
        if topicID > 1
            print(file, " ")
        end
        print(file, btm.numBitermsOfTopic[topicID])
    end
    print(file, "\n")
    for topicID in 1:param.numTopics
        if topicID > 1
            print(file, " ")
        end
        print(file, btm.numRelatedOfTopic[topicID])
    end
    print(file, "\n")
    for topicID in 1:param.numTopics
        for phraseID in 1:numPhrases
            if phraseID > 1
                print(file, " ")
            end
            print(file, btm.phraseFreqsOfTopic[topicID,phraseID])
        end
        print(file, "\n")
    end
    close(file)
end

function newNBTMWE(fileName::String, nonStopWords::Vector{String}, vectors::Array{Float64,2}, numTopics::Int, maxIter::Int, simThres::Float64, alpha::Float64, beta::Float64)::NBTMWE
    phrases = nonStopWords
    phraseIDs = Dict{String,Int}()
    for (phraseID, phrase) in enumerate(phrases)
        phraseIDs[phrase] = phraseID
    end

    docs = Vector{Vector{Int}}()
    file = open(fileName)
    for line in readlines(file)
        phrasesAtLine = split(line)
        doc = Vector{Int}()
        for phrase in phrasesAtLine
            if haskey(phraseIDs, phrase)
                push!(doc, phraseIDs[phrase])
            end
        end
        if length(doc) > 0
            push!(docs, doc)
        end
    end

    numPhrases = length(phrases)
    numBitermsOfTopic = fill(0, numTopics)
    numRelatedOfTopic = fill(0, numTopics)
    phraseFreqsOfTopic = fill(0.0, (numTopics, numPhrases))
    biterms = Vector{Biterm}()
    for doc in docs
        docSize = length(doc)
        for i in 1:docSize
            for j in 1:docSize
                if i != j
                    topicID = rand(1:numTopics)
                    push!(biterms, Biterm(IDPair(doc[i], doc[j]), topicID, false))
                    numBitermsOfTopic[topicID] += 1
                    phraseFreqsOfTopic[topicID, doc[i]] += 1.0
                    phraseFreqsOfTopic[topicID, doc[j]] += 1.0
                end
            end
        end
    end

    similarPhraseIDs = Dict{IDPair, Vector{Int}}()
    numBiterms = length(biterms)
    println("building similar phrases list...")
    simMat = 0.5 .+ 0.5 .* (vectors' * vectors)
    for (bitermID, biterm) in enumerate(biterms)
        if haskey(similarPhraseIDs, biterm.pair)
            continue
        end
        relatedIDs = Vector{Int}()
        for phraseID in 1:numPhrases
            if phraseID == biterm.pair.i || phraseID == biterm.pair.j
                continue
            end
            sim1 = simMat[biterm.pair.i, phraseID]
            sim2 = simMat[biterm.pair.j, phraseID]
            if sim1 >= simThres && sim2 >= simThres
                push!(relatedIDs, phraseID)
            end
        end
        similarPhraseIDs[biterm.pair] = relatedIDs
        if bitermID % 10000 == 0
            println("similar phrase list $bitermID of $numBiterms")
        end
    end
    println("finish building similar phrases list")

    NBTMWE(NBTMWEParameters(numTopics, maxIter, simThres, alpha, beta, phrases, phraseIDs, docs), biterms, similarPhraseIDs, numBitermsOfTopic, numRelatedOfTopic, phraseFreqsOfTopic)
end

function loadNBTMWE(fileName::String)::NBTMWE
    file = open(fileName)
    firstLine = readline(file)
    fieldsOfFirstLine = split(firstLine)
    @assert length(fieldsOfFirstLine) == 9
    numTopics = parse(Int, fieldsOfFirstLine[1])
    maxIter = parse(Int, fieldsOfFirstLine[2])
    simThres = parse(Float64, fieldsOfFirstLine[3])
    alpha = parse(Float64, fieldsOfFirstLine[4])
    beta = parse(Float64, fieldsOfFirstLine[5])
    numPhrases = parse(Int, fieldsOfFirstLine[6])
    numDocs = parse(Int, fieldsOfFirstLine[7])
    numBiterms = parse(Int, fieldsOfFirstLine[8])
    numSimilar = parse(Int, fieldsOfFirstLine[9])
    phrases = Vector{String}()
    phraseIDs = Dict{String,Int}()
    for phraseID in 1:numPhrases
        phrase = readline(file)
        push!(phrases, phrase)
        phraseIDs[phrase] = phraseID
    end
    docs = Vector{Vector{Int}}()
    for docID in 1:numDocs
        doc = Vector{Int}()
        lineOfDoc = readline(file)
        fieldsOfDoc = split(lineOfDoc)
        docSize = length(fieldsOfDoc)
        for i in 1:docSize
            phraseID = parse(Int, fieldsOfDoc[i])
            push!(doc, phraseID)
        end
    end
    biterms = Vector{Biterm}(undef, numBiterms)
    for bitermID in 1:numBiterms
        lineOfBiterm = readline(file)
        fieldsOfBiterm = split(lineOfBiterm)
        @assert length(fieldsOfBiterm) == 4
        i = parse(Int, fieldsOfBiterm[1])
        j = parse(Int, fieldsOfBiterm[2])
        topicID = parse(Int, fieldsOfBiterm[3])
        gpuFlag = parse(Bool, fieldsOfBiterm[4])
        biterm = Biterm(IDPair(i, j), topicID, gpuFlag)
        biterms[bitermID] = biterm
    end
    similarPhraseIDs = Dict{IDPair, Vector{Int}}()
    for iSimilar in 1:numSimilar
        lineOfSimilar = readline(file)
        fieldsOfSimilar = split(lineOfSimilar)
        @assert length(fieldsOfSimilar) >= 2
        i = parse(Int, fieldsOfSimilar[1])
        j = parse(Int, fieldsOfSimilar[2])
        relatedIDs = Vector{Int}()
        for k in 3:length(fieldsOfSimilar)
            id = parse(Int, fieldsOfSimilar[k])
            push!(relatedIDs, id)
        end
        similarPhraseIDs[IDPair(i,j)] = relatedIDs
    end
    numBitermsOfTopic = Vector{Int}(undef, numTopics)
    lineOfNumBiterms = readline(file)
    fieldsOfNumBiterms = split(lineOfNumBiterms)
    @assert length(fieldsOfNumBiterms) == numTopics
    for topicID in 1:numTopics
        numBitermsOfTopic[topicID] = parse(Int, fieldsOfNumBiterms[topicID])
    end
    numRelatedOfTopic = Vector{Int}(undef, numTopics)
    lineOfNumRelated = readline(file)
    fieldsOfNumRelated = split(lineOfNumRelated)
    @assert length(fieldsOfNumRelated) == numTopics
    for topicID in 1:numTopics
        numRelatedOfTopic[topicID] = parse(Int, fieldsOfNumRelated[topicID])
    end
    phraseFreqsOfTopic = Array{Float64,2}(undef, (numTopics, numPhrases))
    for topicID in 1:numTopics
        lineOfTopic = readline(file)
        fieldsOfTopic = split(lineOfTopic)
        @assert length(fieldsOfTopic) == numPhrases
        for phraseID in 1:numPhrases
            freq = parse(Float64, fieldsOfTopic[phraseID])
            phraseFreqsOfTopic[topicID, phraseID] = freq
        end
    end
    close(file)

    NBTMWE(NBTMWEParameters(numTopics, maxIter, simThres, alpha, beta, phrases, phraseIDs, docs), biterms, phraseFreqsOfTopic, numBitermsOfTopic, numRelatedOfTopic, phraseFreqsOfTopic)
end

function computeProbTopicBiterm(btm::NBTMWE, biterm::Biterm)::Vector{Float64}
    numTopics = btm.parameters.numTopics
    numPhrases = length(btm.parameters.phrases)
    numBiterms = length(btm.biterms)
    numTotalRelated = sum(btm.numRelatedOfTopic)
    numBitermRelated = if biterm.gpuFlag
        length(btm.similarPhraseIDs[biterm.pair])
    else
        0
    end
    alpha = btm.parameters.alpha
    beta = btm.parameters.beta
    prob = Vector{Float64}(undef, numTopics)
    for topicID in 1:numTopics
        wRelated = 0.5 * btm.parameters.simThres * (numTotalRelated - numBitermRelated)
        if topicID != biterm.topicID
            w = 1.0 / (2 * btm.numBitermsOfTopic[topicID] + numPhrases * beta)
            prob1 = (btm.phraseFreqsOfTopic[topicID, biterm.pair.i] + beta) * w
            prob2 = (btm.phraseFreqsOfTopic[topicID, biterm.pair.j] + beta) * w
            wTopicRelated = 0.5 * btm.parameters.simThres * btm.numRelatedOfTopic[topicID]
            probTopic = (btm.numBitermsOfTopic[topicID] + wTopicRelated + alpha) / (numBiterms - 1 + wRelated + numTopics * alpha)
            prob[topicID] = prob1 * prob2 * probTopic
        else
            w = 1.0 / (2 * btm.numBitermsOfTopic[topicID] - 2 + numPhrases * btm.parameters.beta)
            prob1 = (btm.phraseFreqsOfTopic[topicID, biterm.pair.i] - 1 + btm.parameters.beta) * w
            prob2 = (btm.phraseFreqsOfTopic[topicID, biterm.pair.j] - 1 + btm.parameters.beta) * w
            wTopicRelated = 0.5 * btm.parameters.simThres * (btm.numRelatedOfTopic[topicID] - numBitermRelated)
            probTopic = (btm.numBitermsOfTopic[topicID] - 1 + wTopicRelated + alpha) / (numBiterms - 1 + wRelated + numTopics * alpha)
            prob[topicID] = prob1 * prob2 * probTopic
        end
    end
    c = 1.0 / sum(prob)
    prob .*= c
    prob
end

function sampleProb(prob::Vector{Float64})::Int
    x = rand()
    partialSum = 0.0
    for i in 1:length(prob)
        partialSum += prob[i]
        if partialSum >= x
            return i
        end
    end
    length(prob)
end

function sampleGPUFlag(prob::Vector{Float64}, topicID::Int)::Bool
    maxProb = max(prob...)
    ratio = prob[topicID] / maxProb
    x = rand()
    ratio > x
end

function trainNBTMWE!(btm::NBTMWE, fileName::String)
    numBiterms = length(btm.biterms)
    numCPUs = length(Sys.cpu_info())
    Segments = Vector{Tuple{Int,Int}}
    TopFlag = Tuple{Int, Bool}
    for iter in 1:btm.parameters.maxIter
        jobs = Channel{Tuple{Int,Int}}(numCPUs)
        jobOutputs = Channel{Tuple{Segments,Vector{Vector{TopFlag}}}}(numCPUs)

        function makeJobs(batchSize::Int)
            for i in 1:batchSize:numBiterms
                put!(jobs, (i, min(i + batchSize - 1, numBiterms)))
            end
        end

        function runJob(iCPU::Int)
            mySegments = Segments()
            myOutputs = Vector{Vector{TopFlag}}()
            for job in jobs
                push!(mySegments, job)
                jobOutput = Vector{TopFlag}(undef, job[2] - job[1] + 1)
                for bitermID in job[1]:job[2]
                    prob = computeProbTopicBiterm(btm, btm.biterms[bitermID])
                    newTopID = sampleProb(prob)
                    newFlag = sampleGPUFlag(prob, newTopID)
                    jobOutput[bitermID - job[1] + 1] = (newTopID,newFlag)
                end
                push!(myOutputs, jobOutput)
            end
            put!(jobOutputs, (mySegments, myOutputs))
        end

        bind(jobs, @async makeJobs(100))
        for iCPU in 1:numCPUs
            Threads.@spawn runJob(iCPU)
        end

        prevPhraseFreqsOfTopic = copy(btm.phraseFreqsOfTopic)
        simThres = btm.parameters.simThres
        for iCPU in 1:numCPUs
            segmentedOutputs = take!(jobOutputs)
            for (segID, segment) in enumerate(segmentedOutputs[1])
                segOutput = segmentedOutputs[2][segID]
                for bitermID in segment[1]:segment[2]
                    biterm = btm.biterms[bitermID]
                    topicFlag = segOutput[bitermID - segment[1] + 1]
                    newTopicID = topicFlag[1]
                    newFlag = topicFlag[2]
                    if newTopicID == biterm.topicID && newFlag == biterm.gpuFlag
                        continue
                    end
                    btm.numBitermsOfTopic[biterm.topicID] -= 1
                    btm.phraseFreqsOfTopic[biterm.topicID, biterm.pair.i] -= 1
                    btm.phraseFreqsOfTopic[biterm.topicID, biterm.pair.j] -= 1
                    if biterm.gpuFlag
                        relatedIDs = btm.similarPhraseIDs[biterm.pair]
                        btm.numRelatedOfTopic[biterm.topicID] -= length(relatedIDs)
                        for id in relatedIDs
                            btm.phraseFreqsOfTopic[biterm.topicID, id] -= simThres
                        end
                    end
                    btm.numBitermsOfTopic[newTopicID] += 1
                    btm.phraseFreqsOfTopic[newTopicID, biterm.pair.i] += 1
                    btm.phraseFreqsOfTopic[newTopicID, biterm.pair.j] += 1
                    btm.biterms[bitermID].topicID = newTopicID
                    if newFlag
                        relatedIDs = btm.similarPhraseIDs[biterm.pair]
                        btm.numRelatedOfTopic[newTopicID] += length(relatedIDs)
                        for id in relatedIDs
                            btm.phraseFreqsOfTopic[newTopicID, id] += simThres
                        end
                    end
                end
            end
        end
        deltaFreq = sum(abs.(btm.phraseFreqsOfTopic - prevPhraseFreqsOfTopic))
        println("iter $iter: deltaFreq = $deltaFreq")

        if iter % 100 == 0
            saveNBTMWE(fileName, btm)
            println("iter $iter: output saved")
        end
    end
    if maxIter % 100 != 0
        saveNBTMWE(fileName, btm)
    end
    println("final output saved")
end

numTopics = 100
maxIter = 300
simThres = 0.9
wordVectors = wordvectors("DualAspectTopicRepresentation/workspace/NIPS.w2v.model.txt")
nonStopWords, nonStopWordIDs = getNonStopWords(wordVectors)
vectors = normalizeVectors(wordVectors, nonStopWords)
btm = newNBTMWE("DualAspectTopicRepresentation/workspace/NIPS.w2v.training.txt", nonStopWords, vectors, numTopics, maxIter, simThres, 2.0, 0.005)
#saveBTM("DualAspectTopicRepresentation/workspace/NIPS.btm.txt", btm)
trainNBTMWE!(btm, "DualAspectTopicRepresentation/workspace/NIPS.nbtmwe.txt")


export vlplot_axes

function asvlplot_scale(scale)
  if scale in [:linear,:log,:sqrt,:time,:utc,:sequential]
    string(scale)
  else
    error("Unsupported scale type $scale")
  end
end

function vlplot_axes(data,args...;colors="reds",width=300,height=300,kwds...)
  df, axes = asplotable(data,args...;kwds...)
  if any(x -> x isa QualitativePlotAxis,axes)
    error("VegaLite interface does not yet support qualitative axes.")
  end

  vlplot_axes_(df,axes,names(df)[2:end]...;colors=colors,
               width=width,height=height)
end

function vlplot_axes_(df,axes,x;colors,width,height)
  df |>
    @vlplot(:line, width=width, height=height,
            x={field=x,typ="quantitative", bin={step=axes[1].step},
               scale={typ=asvlplot_scale(axes[1].scale)}},
            y={field=:value,typ="quantitative"})
end

function vlplot_axes_(df,axes,x,y;colors,width,height)
  df |>
  @vlplot(:rect, width=width, height=height,
          x={field=x,typ="quantitative", bin={step=axes[1].step},
             scale={typ=asvlplot_scale(axes[1].scale)}},
          y={field=y,typ="quantitative", bin={step=axes[2].step},
             scale={typ=asvlplot_scale(axes[2].scale)}},
          color={field=:value, aggregate="mean", typ="quantitative"},
          config={view={stroke="transparent"},
                  scale={bandPaddingInner=0, bandPaddingOuter=0},
                  range={heatmap={scheme=colors}}})
end

function vlplot_axes_(df,axes,x,y,z;colors,width,height)
  df |>
    @vlplot(:rect, width=width, height=height,
            x={field=x,typ="quantitative", bin={step=axes[1].step},
               scale={typ=asvlplot_scale(axes[1].scale)}},
            y={field=y,typ="quantitative", bin={step=axes[2].step},
               scale={typ=asvlplot_scale(axes[2].scale)}},
            color={field=:value, aggregate="mean", typ="quantitative"},
            column={field=z},
            config={view={stroke="transparent"},
                    scale={bandPaddingInner=0, bandPaddingOuter=0},
                    range={heatmap={scheme=colors}}})
end

function vlplot_axes_(df,axes,x,y,z,w;colors,width,height)
  df |>
    @vlplot(:rect, width=width, height=height,
            x={field=x,typ="quantitative", bin={step=axes[1].step},
               scale={typ=asvlplot_scale(axes[1].scale)}},
            y={field=y,typ="quantitative", bin={step=axes[2].step},
               scale={typ=asvlplot_scale(axes[2].scale)}},
            color={field=:value, aggregate="mean", typ="quantitative"},
            column={field=z},row={field=w},
            config={view={stroke="transparent"},
                    scale={bandPaddingInner=0, bandPaddingOuter=0},
                    range={heatmap={scheme=colors}}})
end

function vlplot_axes_(df,axes,args...;kwds...)
  error("Plotting data with $(length(axes)) dims along
        $(length(args)) axes is not supported.")
end


import BioSimulator: execute_jump!, get_site, initialize_datastructs!, get_neighbor_class, get_ptype, coordinates

@testset "Execute" begin
  @testset "Predator-Prey Model" begin
    ##### reaction definition #####
    intermediate = @def_reactions begin
      A + B --> A + A, k1 # predation / predator reproduction
      B + 0 --> 0 + B, k2 # prey movement
      A + 0 --> 0 + A, k3 # predator movement
      B + 0 --> B + B, k4 # prey reproduction
      A --> 0, k5         # predator death
      B --> 0, k6         # prey death
    end k1 k2 k2 k3 k4 k5 k6

    ##### assign parameter values #####
    params = [2.0, 1.0, 1.0, 2.0, 1.0, 1.0]

    ##### generate the full reaction list #####
    model = @enumerate_with_sclass intermediate VonNeumann() 2 params

    ##### initial configuration #####
    initial_sites = [(0,1), (1,0), (1,1), (1,2), (1,3), (2,0), (2,1), (2,2), (3,0), (3,1)]
    initial_sites = [initial_sites[j][i] for i in 1:2, j in 1:length(initial_sites)]
    ptypes        = [1, 1, 1, 2, 2, 1, 2, 2, 2, 2]

    lattice = Lattice(initial_sites, ptypes)

    initialize_datastructs!(lattice, model)
    enum = model.enumeration

    # Test Reaction:
    execute_jump!(lattice, model, 9)

    # aggregate coordinates
    test_coord = [tuple(coordinates(x)...) for x in lattice.site]

    # check new open sites
    @test (2, -1) in test_coord

    # check type changes
    @test get_ptype(get_site(lattice, (2,  0))) == 1
    @test get_ptype(get_site(lattice, (2, -1))) == 2

    # check neighbor class changes
    # (2, 1): 13 -> 9
    # (1, 0): 4  -> 2
    # (3, 0): 5  -> 3
    # (2, 0): 9  -> 13
    # (2, -1): 2 -> 1
    points_to_check = [(2,1), (1,0), (3,0), (2,0), (2,-1)]
    list = [get_site(lattice, x) for x in points_to_check]

    @test get_neighbor_class(list[1]) == 9
    @test get_neighbor_class(list[2]) == 2
    @test get_neighbor_class(list[3]) == 3
    @test get_neighbor_class(list[4]) == 13
    @test get_neighbor_class(list[5]) == 1

    # check sample class changes
    # (2, 1): [14] -> [14, 5]
    # (1, 0): [10, 13] -> [13, 11]
    # (3, 0): [6, 14] -> [14, 7]
    # (2, 0): [2,9,13] -> []
    # (2, -1): [] -> [12, 13]
    particle_sample_classes = [
      (2, 1)  => [5, 14],
      (1, 0)  => [11, 13],
      (3, 0)  => [7, 14],
      (2, 0)  => Int[],
      (2, -1) => [12, 13]
      ]

    for (coord, expected_sample_classes) in particle_sample_classes
      x = get_site(lattice, coord)

      l = get_ptype(x)
      k = get_neighbor_class(x)

      test_sample_classes = sort(get(enum.pair_to_classes, (l,k), Int[]))

      @test test_sample_classes == expected_sample_classes
    end

    # check invariants
    @test issorted(lattice.site)
    @test issorted(lattice.coord_order, by = coordinates)
  end
end


using Test, YaoBlocks, YaoArrayRegister, LuxurySparse, CacheServers

@testset "constructor" begin
    @test CacheFragment(X) isa CacheFragment{XGate,UInt8,Any}
    @test CacheFragment{XGate,Int,PermMatrix{ComplexF64,Int}}(X) isa
          CacheFragment{XGate,Int,PermMatrix{ComplexF64,Int}}
    @test CacheFragment{XGate,Int}(X) isa CacheFragment{XGate,Int,Any}
end

@testset "CacheServer API" begin
    frag = CacheFragment(X)
    @test update!(frag, mat(X)) === frag
    @test pull(frag) == mat(X)
    clear!(frag)
    @test_throws KeyError pull(frag)
end

test_server = DefaultServer{AbstractBlock,CacheFragment}()

@testset "constructor" begin
    c = CachedBlock(test_server, X, 2)
    @test c isa CachedBlock{DefaultServer{AbstractBlock,CacheFragment},XGate}

    blk = kron(4, 2 => Rx(0.3))
    @test first(chsubblocks(c, blk) |> subblocks) == blk
end

@testset "methods" begin
    g = CachedBlock(test_server, X, 3)
    @test_throws KeyError pull(g)

    update_cache(ComplexF64, g)
    @test pull(g) ≈ mat(X)

    clear!(g)
    @test_throws KeyError pull(g)

    @test mat(g) ≈ mat(X)

    clear!(g)

    @test state(apply!(arrayreg(bit"1"), g)) ≈ state(arrayreg(bit"0"))
    @test pull(g) ≈ mat(X)

    clear!(g)
    @test state(YaoBlocks._apply!(arrayreg(bit"1"), g, 2)) ≈ state(arrayreg(bit"0"))
    @test_throws KeyError pull(g)
end

@testset "direct inherited methods" begin
    g = chain([X, Z, Y, I2])
    g = CachedBlock(test_server, g, 2)

    @test g[1] isa XGate
    @test g[3] isa YGate

    g[4] = Z
    @test g[4] isa ZGate

    g = chain(X, Y)
    g = CachedBlock(test_server, g, 2)

    @test g[1] isa XGate
    @test g[2] isa YGate

    #@test iterate(g) == iterate(g.content)
    #@test length(g) == length(g.content)

    @test subblocks(g) == (g.content,)

    gg = chain(g, g)
    cgg = CachedBlock(test_server, gg, 2)
    @test cgg isa CachedBlock
    @test content(cgg)[1] == g
end

@testset "matrix-chain cache" begin
    A = matblock(rand(ComplexF64, 2, 2))
    B = matblock(rand(ComplexF64, 2, 2))
    C = cache(chain(A, B))

    update_cache(ComplexF64, C)
    @test pull(C) ≈ mat(C)
end


#####
##### Interface to Pascal Pons' Connect 4 Solver
##### https://github.com/PascalPons/connect4
#####

# Problem: no Connect4 module. We can change this:

module Solver

import ..Game, ..history, ..WHITE, ..NUM_CELLS
import AlphaZero: GI, GameInterface, Benchmark, AbstractPlayer, think

const DEFAULT_SOLVER_DIR = joinpath(@__DIR__, "solver", "connect4")

struct Player <: AbstractPlayer{Game}
  process :: Base.Process
  function Player(;
      solver_dir=DEFAULT_SOLVER_DIR,
      solver_name="c4solver",
      disable_stderr=true)
    cmd = Cmd(`./$solver_name`, dir=solver_dir)
    if disable_stderr
      cmd = pipeline(cmd, stderr=devnull)
    end
    p = open(cmd, "r+")
    return new(p)
  end
end

# Solver protocol
# - Standard input: one position per line
# - Standard output: space separated
#     position, score, number of explored node, computation time in μs

struct SolverOutput
  score :: Int
  num_explored_nodes :: Int64
  time :: Int64 # in μs
end

history_string(game) = reduce(*, map(string, history(game)))

function query_solver(p::Player, g)
  hstr = history_string(g)
  println(p.process, hstr)
  l = readline(p.process)
  args = map(split(l)[2:end]) do x
    parse(Int64, x)
  end
  return SolverOutput(args...)
end

function remaining_stones(game, player)
  @assert !isnothing(game.history)
  n = length(game.history)
  p = n ÷ 2
  (n % 2 == 1 && player == WHITE) && (p += 1)
  return NUM_CELLS ÷ 2 - p
end

function value(player, game)
  z = GI.white_reward(game)
  if isnothing(z)
    return query_solver(player, game).score
  elseif iszero(z)
    return 0
  else
    v = remaining_stones(game, game.winner) + 1
    z < 0 && (v = -v)
    GI.white_playing(game) || (v = -v)
    return v
  end
end

function qvalue(player, game, action)
  @assert isnothing(GI.white_reward(game))
  wp = GI.white_playing(game)
  game = copy(game)
  GI.play!(game, action)
  pswitch = wp != GI.white_playing(game)
  nextv = value(player, game)
  return pswitch ? - nextv : nextv
end

function think(p::Player, g, turn=nothing)
  as = GI.available_actions(g)
  qs = [qvalue(p, g, a) for a in as]
  maxq = maximum(qs)
  opt = findall(>=(maxq), qs)
  π = zeros(length(as))
  π[opt] .= 1 / length(opt)
  return as, π
end

Benchmark.PerfectPlayer(::Type{Game}) = Player

end


# Plotting setup

# This file uses RecipesBase.jl to define plotting behaviour and layout without
# actually loading any plotting backend. To plot a planetary structure, use the
# Plots package and then just call the plot command on a planetary structure. At
# the moment, only planets with atmospheric layers are handled this way - the
# old plotting behaviour is stored in plots-old.jl until it is all transitioned.

using RecipesBase
using LaTeXStrings

# Plotting planetary structures

@recipe function plot(p::PlanetStructure{WithAtmosphere})
    M = vec(mass(p))
    r = vec(radius(p))
    P = vec(pressure(p))
    T = vec(temperature(p))
    τ = vec(opticaldepth(p))

    legend := false
    markershape := :circle
    markersize := 1
    linewidth := 2

    # T-P profile
    # xlims := (300, Inf)
    # ylims := (1e-3, 1e3)
    # yflip := true
    # yscale := :log10
    # (T/K, P/bar)

    # τ-P profile
    xguide := L"Optical depth $τ$"
    yguide := "Pressure P / bar"
    ylims := (1e-5, 1e2)
    xscale := :log10
    yscale := :log10
    (τ, P/bar)

    # # all profiles
    # layout := 4
    # xguide := L"Radius r / R$_⊕$"
    # yguide := hcat(L"Mass m / M$_⊕$",
    #                "Pressure P / bar",
    #                "Temperature T / K",
    #                L"Optical depth $τ$")
    # ylims := [nothing]
    # yscale := [:identity :log10 :identity :log10]
    #
    # x = r/R_earth
    # ys = hcat(M/M_earth, P/1bar, T/K, τ)
    #
    # (x, ys)
end


#####
##### Hyperparameters
#####
Network = ResNet

netparams = ResNetHP(
  num_filters=128,
  num_blocks=15,
  conv_kernel_size=(3, 3),
  num_policy_head_filters=256,
  num_value_head_filters=256,
  batch_norm_momentum=0.1)

self_play = SelfPlayParams(
  num_games=1000,
  reset_mcts_every=1,
  mcts=MctsParams(
    use_gpu=true,
    num_workers=32,
    num_iters_per_turn=500,
    cpuct=3.0,
    prior_temperature=0.7,
    temperature=PLSchedule([0, 12], [1.0, 0.1]),
    dirichlet_noise_ϵ=0.25,
    dirichlet_noise_α=0.75))

arena = ArenaParams(
  num_games=400,
  reset_mcts_every=1,
  flip_probability=0.1,
  update_threshold=1.0,
  play_swiss_style=true,
  mcts=MctsParams(
    self_play.mcts,
    temperature=ConstSchedule(0.2),
    dirichlet_noise_ϵ=0.05))

learning = LearningParams(
  use_gpu=true,
  use_position_averaging=true,
  samples_weighing_policy=LOG_WEIGHT,
  batch_size=64,
  loss_computation_batch_size=256,
  optimiser=Adam(lr=1e-3),
  l2_regularization=1e-4,
  nonvalidity_penalty=1.,
  min_checkpoints_per_epoch=1,
  max_batches_per_checkpoint=8000,
  num_checkpoints=1)

params = Params(
  arena=arena,
  self_play=self_play,
  learning=learning,
  num_iters=100,
  ternary_rewards=false,
  use_symmetries=true,
  memory_analysis=MemAnalysisParams(
    num_game_stages=4),
  mem_buffer_size=PLSchedule(
  [      0,        100],
  [200_000,    750_000]))

#####
##### Benchmarks
#####

mcts_baseline =
  Benchmark.MctsRollouts(
    MctsParams(
      arena.mcts,
      num_iters_per_turn=1000,
      num_workers=1,
      cpuct=1.))

minmax_baseline = Benchmark.MinMaxTS(depth=5, amplify_rewards=true, τ=0.2)

solver_baseline = Benchmark.Solver(ϵ=0)

players = [
  #Benchmark.Full(arena.mcts),
  Benchmark.Full(arena.mcts),
  Benchmark.Full(arena.mcts),
  Benchmark.NetworkOnly(τ=0.5, use_gpu=true)]

baselines = [
  #solver_baseline,
  mcts_baseline,
  minmax_baseline,
  mcts_baseline]

make_duel(player, baseline) =
  Benchmark.Duel(
    player,
    baseline,
    num_games=200,
    flip_probability=0.1,
    color_policy=ALTERNATE_COLORS,
    play_swiss_style=true)

benchmark = [make_duel(p, b) for (p, b) in zip(players, baselines)]


function OnBoardAntennaRangeModifier(arg0::Vector3D)
    return OnBoardAntennaRangeModifier((Vector3D,), arg0)
end

function get_parameters_drivers(obj::OnBoardAntennaRangeModifier)
    return jcall(obj, "getParametersDrivers", List, ())
end

function modify(obj::OnBoardAntennaRangeModifier, arg0::EstimatedMeasurement)
    return jcall(obj, "modify", void, (EstimatedMeasurement,), arg0)
end


@inline function get_interpolation(wpot::Interpolations.Extrapolation{T, 3}, pt::CartesianPoint{T}, ::Type{Cartesian})::T where {T <: SSDFloat}
    return wpot(pt.x, pt.y, pt.z)::T
end
@inline function get_interpolation(wpot::Interpolations.Extrapolation{T, 3}, pt::CartesianPoint{T}, ::Type{Cylindrical})::T where {T <: SSDFloat}
    pt_cyl::CylindricalPoint{T} = CylindricalPoint(pt)
    return wpot(pt_cyl.r, pt_cyl.φ, pt_cyl.z)::T
end
@inline function get_interpolation(wpot::Interpolations.Extrapolation{T, 3}, pt::CylindricalPoint{T}, ::Type{Cartesian})::T where {T <: SSDFloat}
    pt_car::CartesianPoint{T} = CartesianPoint(pt)
    return wpot(pt_car.x, pt_car.y, pt_car.z)::T
end
@inline function get_interpolation(wpot::Interpolations.Extrapolation{T, 3}, pt::CylindricalPoint{T}, ::Type{Cylindrical})::T where {T <: SSDFloat}
    return wpot(pt.r, pt.φ, pt.z)::T
end

function add_signal!(signal::AbstractVector{T}, timestamps::AbstractVector{T}, path::Vector{CartesianPoint{T}}, pathtimestamps::AbstractVector{T}, charge::T,
                        wpot::Interpolations.Extrapolation{T, 3}, S::CoordinateSystemType)::Nothing where {T <: SSDFloat}
    tmp_signal = Vector{T}(undef, length(pathtimestamps))
    @inbounds for i in eachindex(tmp_signal)
        tmp_signal[i] = get_interpolation(wpot, path[i], S)::T * charge
    end
    itp = interpolate( (pathtimestamps,), tmp_signal, Gridded(Linear()))
    t_max::T = last(pathtimestamps)
    i_max::Int = searchsortedlast(timestamps, t_max)
    signal[1:i_max] += itp(timestamps[1:i_max])
    if length(signal) > i_max
        signal[i_max+1:end] .+= tmp_signal[end]
    end
    nothing
end

function add_signal!(signal::AbstractVector{T}, timestamps::AbstractVector{TT}, path::EHDriftPath{T, TT}, charge::T, wpot::Interpolations.Extrapolation{T, 3}, S::CoordinateSystemType)::Nothing where {T <: SSDFloat, TT}
    add_signal!(signal, timestamps, path.e_path, path.timestamps_e, -charge, wpot, S) # electrons induce negative charge
    add_signal!(signal, timestamps, path.h_path, path.timestamps_h,  charge, wpot, S)
    nothing
end

function add_signal!(signal::AbstractVector{T}, timestamps::AbstractVector{<:RealQuantity}, paths::Vector{<:EHDriftPath{T}}, charges::Vector{T}, wpot::Interpolations.Extrapolation{T, 3}, S::CoordinateSystemType)::Nothing where {T <: SSDFloat}
    for ipath in eachindex(paths)
        add_signal!(signal, timestamps, paths[ipath], charges[ipath], wpot, S)
    end
    nothing
end


module AES

	using StaticArrays, Random

	abstract type AbstractSymmetricKey end
	abstract type AbstractCipher end
	abstract type AbstractCipherCache end

	abstract type AbstractAESKey <: AbstractSymmetricKey end
	abstract type AbstractAESCache <: AbstractCipherCache end

	include("constants.jl")
	include("types.jl")
	include("block_encryption.jl")
	include("block_decryption.jl")
	include("modes/cbc.jl")
	include("modes/ctr.jl")
	include("modes/ecb.jl")
	include("encrypt.jl")
	include("decrypt.jl")

	export AESCipher
	export AES128Key, AES192Key, AES256Key
	export AESCache
	export encrypt, decrypt

end # module


const braille_signs = ['⠁' '⠂' '⠄' '⡀';
                       '⠈' '⠐' '⠠' '⢀']

"""
The type of canvas with the highest resolution
for Unicode-based plotting.
It uses the Unicode characters for
the Braille symbols to represent individual pixel.
This effectively turns every character into 8 pixels
that can individually be manipulated using binary operations.
"""
type BrailleCanvas <: Canvas
    grid::Array{Char,2}
    colors::Array{UInt8,2}
    pixel_width::Int
    pixel_height::Int
    origin_x::Float64
    origin_y::Float64
    width::Float64
    height::Float64
end

@inline pixel_width(c::BrailleCanvas) = c.pixel_width
@inline pixel_height(c::BrailleCanvas) = c.pixel_height
@inline origin_x(c::BrailleCanvas) = c.origin_x
@inline origin_y(c::BrailleCanvas) = c.origin_y
@inline width(c::BrailleCanvas) = c.width
@inline height(c::BrailleCanvas) = c.height
@inline nrows(c::BrailleCanvas) = size(c.grid, 2)
@inline ncols(c::BrailleCanvas) = size(c.grid, 1)

@inline x_pixel_per_char(::Type{BrailleCanvas}) = 2
@inline y_pixel_per_char(::Type{BrailleCanvas}) = 4

function BrailleCanvas(char_width::Int, char_height::Int;
                       origin_x::Real = 0.,
                       origin_y::Real = 0.,
                       width::Real  = 1.,
                       height::Real = 1.)
    char_width = max(char_width, 5)
    char_height = max(char_height, 2)
    pixel_width = char_width * x_pixel_per_char(BrailleCanvas)
    pixel_height = char_height * y_pixel_per_char(BrailleCanvas)
    width > 0 || throw(ArgumentError("Width has to be positive"))
    height > 0 || throw(ArgumentError("Height has to be positive"))
    grid = fill(Char(0x2800), char_width, char_height)
    colors = fill(0x00, char_width, char_height)
    BrailleCanvas(grid, colors,
                  pixel_width, pixel_height,
                  Float64(origin_x), Float64(origin_y),
                  Float64(width), Float64(height))
end

function pixel!(c::BrailleCanvas, pixel_x::Int, pixel_y::Int, color::Symbol)
    0 <= pixel_x <= c.pixel_width || return c
    0 <= pixel_y <= c.pixel_height || return c
    pixel_x = pixel_x < c.pixel_width ? pixel_x : pixel_x - 1
    pixel_y = pixel_y < c.pixel_height ? pixel_y : pixel_y - 1
    cw, ch = size(c.grid)
    tmp = pixel_x / c.pixel_width * cw
    char_x = floor(Int, tmp) + 1
    char_x_off = (pixel_x % 2) + 1
    if char_x < round(Int, tmp, RoundNearestTiesUp) + 1 && char_x_off == 1
        char_x = char_x + 1
    end
    char_y = floor(Int, pixel_y / c.pixel_height * ch) + 1
    char_y_off = (pixel_y % 4) + 1
    c.grid[char_x,char_y] = Char(UInt64(c.grid[char_x,char_y]) | UInt64(braille_signs[char_x_off, char_y_off]))
    c.colors[char_x,char_y] = c.colors[char_x,char_y] | color_encode[color]
    c
end

function printrow(io::IO, c::BrailleCanvas, row::Int)
    nunrows = nrows(c)
    0 < row <= nunrows || throw(ArgumentError("Argument row out of bounds: $row"))
    y = row
    for x in 1:ncols(c)
        print_color(c.colors[x,y], io, c.grid[x,y])
    end
end


@with_kw mutable struct Axis
    prefix::String # x, y, x2, y2, z
    # ---
    ticks    ::Ticks     = Ticks()     # ticks of the axis
    textstyle::TextStyle = TextStyle() # parent textstyle of axis
    # ---
    title ::Option{Title}   = ∅ # title of the axis
    base  ::Option{Float64} = ∅ # scale font and ticks
    lwidth::Option{Float64} = ∅ # width of the axis spine
    min   ::Option{Float64} = ∅ # minimum span of the axis
    max   ::Option{Float64} = ∅ # maximum span of the axis
    # -- toggle-able
    off   ::Bool = false # if true, axis is not shown
    log   ::Bool = false # log scale
end
Axis(p::String; o...) = Axis(prefix=p; o...)

abstract type Axes{B <: Backend} end

@with_kw mutable struct Axes2D{B} <: Axes{B}
    parent::String # id of the parent figure
    # --
    xaxis ::Axis = Axis("x")
    x2axis::Axis = Axis("x2")
    yaxis ::Axis = Axis("y")
    y2axis::Axis = Axis("y2")
    # ---
    drawings::Vector{Drawing2D} = Vector{Drawing2D}()
    objects ::Vector{Object2D}  = Vector{Object2D}()
    # ---
    title ::Option{Title}  = ∅
    size  ::Option{T2F}    = ∅ # (width cm, height cm)
    legend::Option{Legend} = ∅
    origin::Option{T2F}    = ∅ # related to layout
    # -- toggle-able
    math::Bool = false # axis crossing (0, 0)
    off::Bool = false
    # --
    scale::String = "auto"
end

@with_kw mutable struct Axes3D{B} <: Axes{B}
    parent::String
    # --
    xaxis::Axis = Axis("x"; min=0, max=1) # NOTE color = ticks, not spine if box
    yaxis::Axis = Axis("y"; min=0, max=1) # if nobox, then spine
    zaxis::Axis = Axis("z"; min=0, max=1)
    # ---
    drawings::Vector{Drawing3D} = Vector{Drawing3D}()
    objects ::Vector{Object3D}  = Vector{Object3D}()
    # ---
    title::Option{Title} = ∅
    size ::T2F           = (20.,10.) # box size
    # cube
    nocube   ::Bool      = false # XXX if true, then xaxis have an expressed linestyle
    cubedims ::T3F       = (20.,20.,10.) # cube sides x,y,z
    linestyle::LineStyle = LineStyle() # XXX only lstyle, color, see set
    # XXX legend::Legend ...
    origin::Option{T2F} = ∅
    # rotation and view
    rotate::Option{T2F} = ∅
    # XXX view x y p
    off::Bool = false # do not show
end

parent(a::Axes) = Figure(a.parent; _noreset=true)

function Base.show(io::IO, ::MIME"text/plain", a::Axes2D{GLE})
    s = "GPlot.Axes2D{GLE}" *
        "\n\t"*rpad("Title:", 15) * (isdef(a.title) ? "\"$(a.title.text)\"" : "none") *
        "\n\t"*rpad("N. drawings:", 15) * "$(length(a.drawings))" *
        "\n\t"*rpad("N. objects:", 15) * "$(length(a.objects))" *
        "\n\t"*rpad("Math mode:", 15) * "$(a.math)" *
        "\n\t"*rpad("Layout origin:", 15) * (isdef(a.origin) ? "$(round.(a.origin, digits=1))" : "auto")
    write(io, s)
    return nothing
end