"""
    struct HamiltonianMC <: MCMCAlgorithm

Hamiltonian Monte Carlo sampling algorithm.

The gradient of the target density is computed via auto-differentiation.

* Note: The fields of `HamiltonianMC` are still subject to change, and not
yet part of stable public API!*

Constructors:

* ```$(FUNCTIONNAME)(; fields...)```

Fields:

$(TYPEDFIELDS)
"""
@with_kw struct HamiltonianMC <: MCMCAlgorithm
    metric::HMCMetric = DiagEuclideanMetric()
    gradient::Module = ForwardDiff
    integrator::HMCIntegrator = LeapfrogIntegrator()
    proposal::HMCProposal = NUTS()
    adaptor::HMCAdaptor = StanHMCAdaptor()
end

export HamiltonianMC

bat_default(::Type{MCMCSampling}, ::Val{:trafo}, mcalg::HamiltonianMC) = PriorToGaussian()

bat_default(::Type{MCMCSampling}, ::Val{:nsteps}, mcalg::HamiltonianMC, trafo::AbstractDensityTransformTarget, nchains::Integer) = 10^4

get_mcmc_tuning(algorithm::HamiltonianMC) = MCMCNoOpTuning()

# MCMCIterator subtype for HamiltonianMC
mutable struct AHMCIterator{
    AL<:HamiltonianMC,
    D<:AbstractDensity,
    R<:AbstractRNG,
    PR<:RNGPartition,
    SV<:DensitySampleVector,
    I<:AdvancedHMC.AbstractIntegrator,
    P<:AdvancedHMC.AbstractProposal,
    A<:AdvancedHMC.AbstractAdaptor,
    H<:AdvancedHMC.Hamiltonian,
} <: MCMCIterator
    algorithm::AL
    density::D
    rng::R
    rngpart_cycle::PR
    info::MCMCIteratorInfo
    samples::SV
    nsamples::Int64
    stepno::Int64
    hamiltonian::H
    transition::AdvancedHMC.Transition
    integrator::I
    proposal::P
    adaptor::A
end

function AHMCIterator(
    rng::AbstractRNG,
    algorithm::MCMCAlgorithm,
    density::AbstractDensity,
    info::MCMCIteratorInfo,
    x_init::AbstractVector{P},
) where {P<:Real}
    stepno::Int64 = 0

    npar = totalndof(density)

    params_vec = Vector{P}(undef, npar)
    params_vec .= x_init
    !(params_vec in var_bounds(density)) && throw(ArgumentError("Parameter(s) out of bounds"))

    log_posterior_value = eval_logval(density, params_vec, strict = true)

    T = typeof(log_posterior_value)
    W = Float64 # ToDo: Support other sample weight types

    sample_info = AHMCSampleID(info.id, info.cycle, 1, CURRENT_SAMPLE, 0.0, 0, false, 0.0)
    current_sample = DensitySample(params_vec, log_posterior_value, one(W), sample_info, nothing)
    samples = DensitySampleVector{Vector{P},T,W,AHMCSampleID,Nothing}(undef, 0, npar)
    push!(samples, current_sample)

    nsamples::Int64 = 0

    rngpart_cycle = RNGPartition(rng, 0:(typemax(Int16) - 2))

    metric = AHMCMetric(algorithm.metric, npar)
    logval_posterior(v) = eval_logval(density, v)

    hamiltonian = AdvancedHMC.Hamiltonian(metric, logval_posterior, algorithm.gradient)
    hamiltonian, t = AdvancedHMC.sample_init(rng, hamiltonian, params_vec)

    algorithm.integrator.step_size == 0.0 ? algorithm.integrator.step_size = AdvancedHMC.find_good_stepsize(hamiltonian, params_vec) : nothing
    ahmc_integrator = AHMCIntegrator(algorithm.integrator)

    ahmc_proposal = AHMCProposal(algorithm.proposal, ahmc_integrator)
    ahmc_adaptor = AHMCAdaptor(algorithm.adaptor, metric, ahmc_integrator)

    chain = AHMCIterator(
        algorithm,
        density,
        rng,
        rngpart_cycle,
        info,
        samples,
        nsamples,
        stepno,
        hamiltonian,
        t,
        ahmc_integrator,
        ahmc_proposal,
        ahmc_adaptor
    )

    reset_rng_counters!(chain)

    chain
end

function MCMCIterator(
    rng::AbstractRNG,
    algorithm::HamiltonianMC,
    density::AbstractDensity,
    chainid::Integer,
    startpos::AbstractVector{<:Real}
)
    cycle = 0
    tuned = false
    converged = false
    info = MCMCIteratorInfo(chainid, cycle, tuned, converged)
    AHMCIterator(rng, algorithm, density, info, startpos)
end

@inline _current_sample_idx(chain::AHMCIterator) = firstindex(chain.samples)
@inline _proposed_sample_idx(chain::AHMCIterator) = lastindex(chain.samples)

getalgorithm(chain::AHMCIterator) = chain.algorithm

getdensity(chain::AHMCIterator) = chain.density

getrng(chain::AHMCIterator) = chain.rng

mcmc_info(chain::AHMCIterator) = chain.info

nsteps(chain::AHMCIterator) = chain.stepno

nsamples(chain::AHMCIterator) = chain.nsamples

current_sample(chain::AHMCIterator) = chain.samples[_current_sample_idx(chain)]

sample_type(chain::AHMCIterator) = eltype(chain.samples)

function reset_rng_counters!(chain::AHMCIterator)
    set_rng!(chain.rng, chain.rngpart_cycle, chain.info.cycle)
    rngpart_step = RNGPartition(chain.rng, 0:(typemax(Int32) - 2))
    set_rng!(chain.rng, rngpart_step, chain.stepno)
    nothing
end

function samples_available(chain::AHMCIterator)
    i = _current_sample_idx(chain::AHMCIterator)
    chain.samples.info.sampletype[i] == ACCEPTED_SAMPLE
end

function get_samples!(appendable, chain::AHMCIterator, nonzero_weights::Bool)::typeof(appendable)
    if samples_available(chain)
        samples = chain.samples

        for i in eachindex(samples)
            st = samples.info.sampletype[i]
            if (
                (st == ACCEPTED_SAMPLE || st == REJECTED_SAMPLE) &&
                (samples.weight[i] > 0 || !nonzero_weights)
            )
                push!(appendable, samples[i])
            end
        end
    end
    appendable
end

function next_cycle!(chain::AHMCIterator)
    _cleanup_samples(chain)

    chain.info = MCMCIteratorInfo(chain.info, cycle = chain.info.cycle + 1)
    chain.nsamples = 0
    chain.stepno = 0

    reset_rng_counters!(chain)

    resize!(chain.samples, 1)

    i = _proposed_sample_idx(chain)
    @assert chain.samples.info[i].sampletype == CURRENT_SAMPLE
    chain.samples.weight[i] = 1

    t_stat = chain.transition.stat

    chain.samples.info[i] = AHMCSampleID(
        chain.info.id, chain.info.cycle, chain.stepno, CURRENT_SAMPLE,
        t_stat.hamiltonian_energy, t_stat.tree_depth,
        t_stat.numerical_error, t_stat.step_size
    )

    chain
end

function _cleanup_samples(chain::AHMCIterator)
    samples = chain.samples
    current = _current_sample_idx(chain)
    proposed = _proposed_sample_idx(chain)
    if (current != proposed) && samples.info.sampletype[proposed] == CURRENT_SAMPLE
        # Proposal was accepted in the last step
        @assert samples.info.sampletype[current] == ACCEPTED_SAMPLE
        samples.v[current] .= samples.v[proposed]
        samples.logd[current] = samples.logd[proposed]
        samples.weight[current] = samples.weight[proposed]
        samples.info[current] = samples.info[proposed]

        resize!(samples, 1)
    end
end

function mcmc_step!(chain::AHMCIterator)
    _cleanup_samples(chain)

    samples = chain.samples
    algorithm = getalgorithm(chain)

    chain.stepno += 1
    reset_rng_counters!(chain)

    rng = getrng(chain)
    density = getdensity(chain)

    # Grow samples vector by one:
    resize!(samples, size(samples, 1) + 1)
    samples.info[lastindex(samples)] = AHMCSampleID(
        chain.info.id, chain.info.cycle, chain.stepno, PROPOSED_SAMPLE,
        0.0, 0, false, 0.0
    )

    current = _current_sample_idx(chain)
    proposed = _proposed_sample_idx(chain)
    @assert current != proposed

    current_params = samples.v[current]
    proposed_params = samples.v[proposed]

    # Propose new variate:
    samples.weight[proposed] = 0

    ahmc_step!(rng, algorithm, chain, proposed_params, current_params)
    tstat = chain.transition.stat

    current_log_posterior = samples.logd[current]
    T = typeof(current_log_posterior)

    # Evaluate prior and likelihood with proposed variate:
    proposed_log_posterior = eval_logval(density, proposed_params, strict = false)

    samples.logd[proposed] = proposed_log_posterior

    accepted = current_params != proposed_params

    if accepted
        samples.info.sampletype[current] = ACCEPTED_SAMPLE
        samples.info.sampletype[proposed] = CURRENT_SAMPLE
        chain.nsamples += 1

        samples.info.hamiltonian_energy[proposed] = tstat.hamiltonian_energy
        samples.info.tree_depth[proposed] = tstat.tree_depth
        samples.info.divergent[proposed] = tstat.numerical_error
        samples.info.step_size[proposed] = tstat.step_size
    else
        samples.info.sampletype[proposed] = REJECTED_SAMPLE
    end

    delta_w_current, w_proposed = if accepted
        (0, 1)
    else
        (1, 0)
    end

    samples.weight[current] += delta_w_current
    samples.weight[proposed] = w_proposed

    nothing
end

function ahmc_step!(rng, algorithm, chain, proposed_params, current_params)

    chain.transition = AdvancedHMC.step(rng, chain.hamiltonian, chain.proposal, chain.transition.z)

    tstat = AdvancedHMC.stat(chain.transition)

    if typeof(algorithm.adaptor) <: StanHMCAdaptor
        i = chain.adaptor.state.i
        n_adapts = algorithm.adaptor.n_adapts
    else
        i, n_adapts = chain.info.converged ? (3, 2) : (1, 1)
    end

    chain.hamiltonian, chain.proposal, isadapted = AdvancedHMC.adapt!(
        chain.hamiltonian,
        chain.proposal,
        chain.adaptor,
        Int(i),
        Int(n_adapts),
        chain.transition.z.θ,
        tstat.acceptance_rate
    )

    if i == n_adapts
        chain.info = MCMCIteratorInfo(chain.info, tuned = isadapted)
    end

    tstat = merge(tstat, (is_adapt=isadapted,))

    proposed_params[:] = chain.transition.z.θ
    nothing
end


@testset "utils" begin

    @testset "update_with_cpmodel!()" begin
        trailer = SeaPearl.Trailer()
        model = SeaPearl.CPModel(trailer)

        x = SeaPearl.IntVar(2, 3, "x", trailer)
        y = SeaPearl.IntVar(3, 4, "y", trailer)
        z = SeaPearl.IntVar(4, 5, "z", trailer)
        SeaPearl.addVariable!(model, x)
        SeaPearl.addVariable!(model, y)
        SeaPearl.addVariable!(model, z; branchable=false)
        SeaPearl.addConstraint!(model, SeaPearl.Equal(x, y, trailer))
        SeaPearl.addConstraint!(model, SeaPearl.NotEqual(x, y, trailer))

        lh = SeaPearl.LearnedHeuristic(agent)

        SeaPearl.update_with_cpmodel!(lh, model)

        @test typeof(lh.action_space) == Array{Int64,1}
        @test lh.action_space == [2, 3, 4]
        @test typeof(lh.current_state) == SeaPearl.DefaultStateRepresentation{SeaPearl.DefaultFeaturization, SeaPearl.DefaultTrajectoryState}
        @test lh.reward.value == 0
        @test isa(lh.search_metrics, SeaPearl.SearchMetrics)
    end

    @testset "sync_state!()" begin
        trailer = SeaPearl.Trailer()
        model = SeaPearl.CPModel(trailer)

        x = SeaPearl.IntVar(2, 3, "x", trailer)
        y = SeaPearl.IntVar(2, 3, "y", trailer)
        SeaPearl.addVariable!(model, x)
        SeaPearl.addVariable!(model, y)
        SeaPearl.addConstraint!(model, SeaPearl.Equal(x, y, trailer))
        SeaPearl.addConstraint!(model, SeaPearl.NotEqual(x, y, trailer))

        lh = SeaPearl.LearnedHeuristic(agent)
        SeaPearl.update_with_cpmodel!(lh, model)

        SeaPearl.sync_state!(lh, model, x)

        @test Matrix(LightGraphs.LinAlg.adjacency_matrix(lh.current_state.cplayergraph)) == [0 0 1 1 0 0
                                                        0 0 1 1 0 0
                                                        1 1 0 0 1 1
                                                        1 1 0 0 1 1
                                                        0 0 1 1 0 0
                                                        0 0 1 1 0 0]

        @test lh.current_state.nodeFeatures == Float32[ 1 1 0 0 0 0
                                                            0 0 1 1 0 0
                                                            0 0 0 0 1 1]

        @test lh.current_state.variableIdx == 3

        SeaPearl.assign!(x, 2)
        SeaPearl.sync_state!(lh, model, y)

        @test Matrix(LightGraphs.LinAlg.adjacency_matrix(lh.current_state.cplayergraph)) == [0 0 1 1 0 0
                                                        0 0 1 1 0 0
                                                        1 1 0 0 1 0
                                                        1 1 0 0 1 1
                                                        0 0 1 1 0 0
                                                        0 0 0 1 0 0]

        @test lh.current_state.nodeFeatures == Float32[ 1 1 0 0 0 0
                                                            0 0 1 1 0 0
                                                            0 0 0 0 1 1]

        @test lh.current_state.variableIdx == 4
    end

    @testset "get_observation!()" begin
        trailer = SeaPearl.Trailer()
        model = SeaPearl.CPModel(trailer)

        x = SeaPearl.IntVar(2, 4, "x", trailer)
        y = SeaPearl.IntVar(6, 8, "y", trailer)
        SeaPearl.addVariable!(model, x)
        SeaPearl.addVariable!(model, y)
        SeaPearl.addConstraint!(model, SeaPearl.Equal(x, y, trailer))
        SeaPearl.addConstraint!(model, SeaPearl.NotEqual(x, y, trailer))

        lh = SeaPearl.LearnedHeuristic(agent)
        SeaPearl.update_with_cpmodel!(lh, model)

        obs = SeaPearl.get_observation!(lh, model, x)

        @test obs.reward == 0
        @test obs.terminal == false
        @test obs.actions_index==[1, 2, 3, 4, 5, 6]  #correspond to the index of branchable values [2, 3, 4, 6, 7, 8] for ALL variables
        @test obs.legal_actions == [2, 3, 4]
        @test obs.legal_actions_mask == [true, true, true, false, false, false]

        SeaPearl.remove!(x.domain, 2)

        obs = SeaPearl.get_observation!(lh, model, x)

        @test obs.reward == 0
        @test obs.terminal == false
        @test obs.actions_index==[1, 2, 3, 4, 5, 6]  #correspond to the index of branchable values [2, 3, 4, 6, 7, 8] for ALL variables
        @test obs.legal_actions == [3, 4]
        @test obs.legal_actions_mask == [false, true, true, false, false, false]

        SeaPearl.remove!(y.domain, 7)

        obs = SeaPearl.get_observation!(lh, model, y)

        @test obs.reward == 0
        @test obs.terminal == false
        @test obs.actions_index==[1, 2, 3, 4, 5, 6]  #correspond to the index of branchable values [2, 3, 4, 6, 7, 8] for ALL variables
        @test obs.legal_actions == [6, 8]
        @test obs.legal_actions_mask == [false, false, false, true, false, true]

    end

    @testset "action_to_value(::FixedOutput)" begin
        trailer = SeaPearl.Trailer()
        model = SeaPearl.CPModel(trailer)

        x = SeaPearl.IntVar(2, 3, "x", trailer)
        y = SeaPearl.IntVar(2, 3, "y", trailer)
        SeaPearl.addVariable!(model, x)
        SeaPearl.addVariable!(model, y)

        lh = SeaPearl.LearnedHeuristic{SeaPearl.DefaultStateRepresentation{SeaPearl.DefaultFeaturization,SeaPearl.DefaultTrajectoryState}, SeaPearl.DefaultReward, SeaPearl.FixedOutput}(agent)
        SeaPearl.update_with_cpmodel!(lh, model)

        obs = SeaPearl.get_observation!(lh, model, x)

        @test SeaPearl.action_to_value(lh, 1, obs.state, model) == 2
        @test SeaPearl.action_to_value(lh, 2, obs.state, model) == 3
        @test_throws BoundsError SeaPearl.action_to_value(lh, 3, obs.state, model)
    end

    #TODO repair TSPTW testset
    """
    @testset "action_to_value(::VariableOutput)" begin
        trailer = SeaPearl.Trailer()
        model = SeaPearl.CPModel(trailer)

        x = SeaPearl.IntVar(2, 3, "x", trailer)
        y = SeaPearl.IntVar(2, 3, "y", trailer)
        SeaPearl.addVariable!(model, x)
        SeaPearl.addVariable!(model, y)

        lh = SeaPearl.LearnedHeuristic{SeaPearl.DefaultStateRepresentation{SeaPearl.DefaultFeaturization,SeaPearl.DefaultTrajectoryState}, SeaPearl.DefaultReward, SeaPearl.VariableOutput}(agent)
        SeaPearl.update_with_cpmodel!(lh, model)

        obs = SeaPearl.get_observation!(lh, model, y)
        @test SeaPearl.action_to_value(lh, 1, obs.state, model) == 4
        @test SeaPearl.action_to_value(lh, 2, obs.state, model) == 2
        @test SeaPearl.action_to_value(lh, 3, obs.state, model) == 3
        @test_throws BoundsError SeaPearl.action_to_value(lh, 4, obs.state, model)

        SeaPearl.remove!(y.domain, 3)
        SeaPearl.update_with_cpmodel!(lh, model)

        obs = SeaPearl.get_observation!(lh, model, y)
        @test SeaPearl.action_to_value(lh, 1, obs.state, model) == 4
        @test SeaPearl.action_to_value(lh, 2, obs.state, model) == 2
        @test_throws BoundsError SeaPearl.action_to_value(lh, 3, obs.state, model)

        obs = SeaPearl.get_observation!(lh, model, x)
        @test SeaPearl.action_to_value(lh, 1, obs.state, model) == 2
        @test SeaPearl.action_to_value(lh, 2, obs.state, model) == 3
        @test_throws BoundsError SeaPearl.action_to_value(lh, 4, obs.state, model)
    end
    """

    @testset "branchable_values()" begin
        trailer = SeaPearl.Trailer()
        model = SeaPearl.CPModel(trailer)

        x = SeaPearl.IntVar(2, 3, "x", trailer)
        y = SeaPearl.IntVar(3, 4, "y", trailer)
        z = SeaPearl.IntSetVar(4, 5, "z", trailer)
        SeaPearl.addVariable!(model, x)
        SeaPearl.addVariable!(model, y; branchable=false)
        SeaPearl.addVariable!(model, z; branchable=false)

        @test SeaPearl.branchable_values(model) == [2, 3]
    end
end


@everywhere using ArcadeLearningEnvironment
@everywhere using CGP
@everywhere using Logging
@everywhere using Printf
@everywhere using Random
@everywhere import Images
using ArgParse

@everywhere CGP.Config.init("cfg/atari.yaml")
include("../graphing/graph_utils.jl")

@everywhere function play_atari(c::Chromosome, id::String, seed::Int64;
                    render::Bool=false, folder::String=".", max_frames=18000)
    game = Game(id, seed)

    if get_float(game, "repeat_action_probability") != 0.0
        # We implement action repeating already, so we don't want it enabled in A.L.E. too.
        error("A.L.E. is configured with non-zero repeat_action_probability")
    end

    seed_reset = rand(0:100000)
    Random.seed!(seed)
    reward = 0.0
    frames = 0
    p_action = game.actions[1]
    outputs = zeros(Int64, c.nout)
    while ~game_over(game.ale)
        frames += 1
        output = process(c, get_rgb(game))
        outputs[argmax(output)] += 1
        action = game.actions[argmax(output)]
        reward += act(game.ale, action)
        if rand() < 0.25
            reward += act(game.ale, action) # repeat action here for seeding
        end
        if render
            screen = draw(game)
            filename = string(folder, "/", @sprintf("frame_%06d.png", frames))
            Images.save(filename, screen)
        end
        if frames > max_frames
            @debug(string("Termination due to frame count on ", id))
            break
        end
    end
    # Using println rather than @info here so the output is included
    # with the "From worker N" prefix in the logs
    println("Frames: ", frames, " Reward: ", reward)
    close!(game)
    Random.seed!(seed_reset)
    reward, outputs
end

function get_args()
    s = ArgParseSettings()

    @add_arg_table(
        s,
        "--seed", arg_type = Int, default = 0,
        "--log", arg_type = String, default = "atari.log",
        "--id", arg_type = String, default = "centipede",
        "--ea", arg_type = String, default = "oneplus",
        "--chromosome", arg_type = String, default = "CGPChromo",
        "--frames", arg_type = Int, default = 18000,
        "--render", action = :store_const, constant = true, default = false,
    )

    parse_args(CGP.Config.add_arg_settings!(s))
end

function get_bests(logfile::String)
    bests = []
    for line in readlines(logfile)
        if occursin("C:", line)
            genes = eval(Meta.parse(split(line, "C: ")[2]))
            append!(bests, [genes])
        end
    end
    bests
end

function get_params(args::Dict)
    game = Game(args["id"], args["seed"])
    nin = 3 # r g b
    nout = length(game.actions)
    close!(game)
    nin, nout
end

function render_genes(genes::Array{Float64}, args::Dict;
                    ctype::DataType=CGPChromo, id::Int64=0)
    nin, nout = get_params(args)
    chromo = ctype(genes, nin, nout)
    folder = joinpath("frames", string(args["id"], "_", args["seed"], "_", id))
    mkpath(folder)
    reward, out_counts = play_atari(chromo, args["id"], args["seed"];
                                    render=true, folder=folder,
                                    max_frames=args["frames"])
    println(@sprintf("R: %s %d %d %0.5f %d %d",
                     args["id"], args["seed"], id, reward,
                     sum([n.active for n in chromo.nodes]),
                     length(chromo.nodes)))
    new_genes = deepcopy(chromo.genes)
    for o in 1:nout
        if out_counts[o] == 0.0
            # disable this section of the graph by forcing
            # the output to connect to an input node
            new_genes[nin+o] = 0.00001
        end
    end
    active_outputs = out_counts .> 0
    chromo2 = ctype(new_genes, nin, nout)
    mkpath("graphs")
    filename = string(args["id"], "_", args["seed"], "_", id, ".pdf");
    file = abspath(joinpath("graphs", filename))
    chromo_draw(chromo2, file; active_outputs=active_outputs)
end

if ~isinteractive()
    args = get_args()
    CGP.Config.init(Dict([k=>args[k] for k in setdiff(
        keys(args), ["seed", "log", "id", "ea", "chromosome"])]...))

    Random.seed!(args["seed"])
    global_logger(SimpleLogger(open(args["log"], "a+")))
    nin, nout = get_params(args)
    ea = eval(Meta.parse(args["ea"]))
    ctype = eval(Meta.parse(args["chromosome"]))

    fit = x->play_atari(x, args["id"], args["seed"];
                        max_frames=args["frames"])[1]
    maxfit, best = ea(nin, nout, fit;
                      seed=args["seed"], id=args["id"], ctype=ctype)

    # Note: -maxfit is returned, being the "cost" to be minimized by irace
    @info(@sprintf("Z%0.6f", -maxfit))
    if args["render"]
        render_genes(best, args; ctype=ctype)
    end
end


module CHIP8

using OffsetArrays

include("sprites.jl")           # Defines sprites for all hex digits
include("utils.jl")
include("keyboard.jl")
include("tests.jl")

export Chip8, parseopcode

const memsize = 4096            # 4KiB of memory
const program_start = 0x200 + 1 # Starting address of programs in memory (+1 because of Julia 1 indexing ...)
const disp_columns = 64         # Number of horizontal pixels
const disp_rows = 32            # Number of vertical pixels

"Holds the state of a CHIP8 system"
mutable struct Chip8
  # Registers
  V::AbstractVector{UInt8}      # General purpose registers
  I::UInt16                     # Address register
  DT::UInt8                     # Timer register
  ST::UInt8                     # Sound register
  PC::UInt16                    # Program counter
  # Stack pointer is only needed if the stack is managed manually
  # PC::UInt16                  # Stack pointer

  # Memory
  mem::AbstractVector{UInt8}    # Main memory
  disp::AbstractMatrix{Bool}    # Display buffer

  # Stack
  stack::AbstractVector{UInt16}

  # Keyboard
  keys::AbstractVector{Bool}    # Models which of the 16 keys are pressed at the moment
end

function Chip8()
  ch8 = Chip8(OffsetVector(zeros(UInt8, 16), 0:15),                # General purpose registers
              0, 0, 0,                                             # Address, timer, sound register
              program_start,                                       # Program counter
              OffsetVector(zeros(UInt8, memsize), 0:memsize-1),    # Main memory
              OffsetMatrix(zeros(Bool, 64, 32), 0:63, 0:32),       # Display buffer
              Vector{UInt16}(),                                    # Stack
              OffsetVector(zeros(Bool, 16), 0:15),                 # Keys
              )
  # Write predefined digit sprites to memory (starting at 0x100)
  ch8.mem[0x101:0x101+length(sprites)-1] = sprites

  return ch8
end

include("graphics.jl")

"""
`decode(b::Integer) = reverse(digits(b, base = 2, pad = sizeof(b)))`

Returns the digits of b as an array, where the most significant bit has the lowest index.
Mainly used for decoding lines of sprites saved in memory as bytes.
"""
decode(b::Integer) = reverse(digits(b, base = 2, pad = sizeof(b)))

# Define our instructions

struct Instruction
  asm::String
  opcode::String
  description::String
  f::Function
end

#TODO: Find a good solution for the 1-index problem in Julia ...

instructions = [
  Instruction("CLS", "00E0", "Clear Screen",
              function (c)
              c.disp .= false
              end)
  Instruction("RET", "00EE", "Return from subroutine",
              function (c)
              c.PC = pop(c.stack)
              end)
  Instruction("SYS addr", "0nnn", "Call system subroutine",
              function (c, nnn)
              error("Tried to call system subroutine on address $nnn. Not possible in emulation!")
              end)
  Instruction("JP addr", "1nnn", "Jump to address",
              function (c, nnn)
              c.PC = nnn - 2    # -2 because the address is going to be advanced by 2 bytes later
              end)
  Instruction("CALL addr", "2nnn", "Call subroutine at address",
              function (c, nnn)
              push!(c.stack, c.PC)
              c.PC = nnn - 2
              end)
  Instruction("SE Vx, val", "3xkk", "Jump over next instruction if (Vx == val,)",
              function (c, x, kk)
              (c.V[x] == kk) && (c.PC += 2)
              end)
  Instruction("SNE Vx, val", "4xkk", "Jump over next instruction if (Vx != val)",
              function (c, x, kk)
              (c.V[x] != kk) && (c.PC += 2)
              end)
  Instruction("SE Vx, Vy", "5xy0", "Jump over next instruction if (Vx == Vy)",
              function (c, x, y)
              (c.V[x] == c.V[y]) && (c.PC += 2)
              end)
  Instruction("LD Vx, val", "6xkk", "Vx = val",
              function (c, x, kk)
              c.V[x] == kk
              end)
  Instruction("ADD Vx, val", "7xkk", "Vx = Vx + val",
              function (c, x, kk)
              c.V[x] += kk
              end)
  Instruction("LD Vx, Vy", "8xy0", "Vx = Vy",
              function (c, x, y)
              c.V[x] = c.V[y]
              end)
  Instruction("OR Vx, Vy", "8xy1", "Vx = Vx or Vy",
              function (c, x, y)
              c.V[x] = c.V[x] | c.V[y]
              end)
  Instruction("AND Vx, Vy", "8xy2", "Vx = Vx and Vy",
              function (c, x, y)
              c.V[x] = c.V[x] & c.V[y]
              end)
  Instruction("XOR Vx, Vy", "8xy3", "Vx = Vx xor Vy",
              function (c, x, y)
              c.V[x] = c.V[x] ⊻ c.V[y]
              end)
  Instruction("ADD Vx, Vy", "8xy4", "Vx = Vx + Vy",
              function (c, x, y)
              c.V[16] = (UInt(c.V[x]) + UInt(x.V[y])) > 255 # Set carry flag
              c.V[x] += c.V[y]
              end)
  Instruction("SUB Vx, Vy", "8xy5", "Vx = Vx - Vy",
              function (c, x, y)
              c.V[16] = c.V[x] > x.V[y] # Set carry flag
              c.V[x] -= c.V[y]
              end)
  Instruction("SHR Vx {, Vy}", "8xy6", "Vx = Vx >> 1",
              function (c, x, y)
              c.V[16] = c.V[x] & 0x01
              c.V[x] >>=  1
              end)
  Instruction("SUBN Vx, Vy", "8xy7", "Vx = Vy - Vx",
              function (c, x, y)
              c.V[16] = c.V[y] > c.V[x]
              c.V[x] = c.V[y] - c.V[x]
              end)
  Instruction("SHL Vx {, Vy}", "8xyE", "Vx = Vx << 1",
              function (c, x, y)
              c.V[16] = c.V[x] & 0x01 != 0x01
              c.V[x] <<= 1
              end)
  Instruction("SNE Vx, Vy", "9xy0", "Jump over next instruction if (Vx != Vy)",
              function (c, x, y)
              if c.V[x] != c.V[y]
              end
              end)
  Instruction("LD I, addr", "Annn", "I = addr",
              function (c, nnn)
              c.I = nnn
              end)
  Instruction("JP V0, addr", "Bnnn", "Jump to address: V0 + addr",
              function (c, nnn)
              c.I = c.V[1] + nnn
              end)
  Instruction("RND Vx, kk", "Cxkk", "Vx = rand() and kk",
              function (c, x, kk)
              x.V[x] = rand(UInt8) & kk
              # Do we have to set V[16] here?
              end)
  Instruction("DRW Vx, Vy, n", "Dxyn", "Display n rows of sprites beginning at I to the position Vx, Vy",
              function (c, x, y, n)
              c.V[16] = 0       # Will contain 1 if there was a collision
              for line in 1:n
                ypos = (c.V[y] + line) % disp_rows
                spriteline = decode(c.mem[c.I + line - 1]) # Decode one line (one byte) of sprite
                for pixel in eachindex(spriteline)
                  xpos = (c.V[x] + pixel) % disp_columns
                  # Check if there has been a collision while drawing the pixel
                  (spriteline[pixel] == 1) && (c.disp[xpos, ypos] == 1) && (c.V[16] = 1)
                  # Draw pixel of sprite. This is done using xor.
                  # i.e. drawing a white pixel on a white pixel will make it black again
                  c.disp[xpos, ypos] ⊻= spriteline[pixel]
                  end
                end
              end)
  Instruction("SKP Vx", "Ex9E", "Skip next instruction if key Vx is pressed",
              function (c, x)

              end)
  Instruction("SKNP Vx", "ExA1", "Skip next instruction if key Vx is not pressed",
              function (c, x)

              end)
  Instruction("LD Vx, DT", "Fx07", "Vx = DT",
              function (c, x)

              end)
  Instruction("LD Vx, K", "Fx0A", "Wait for key press and put pressed key into Vx",
              function (c, x)

              end)
  Instruction("LD DT, Vx", "Fx15", "DT = Vx",
              function (c, x)

              end)
  Instruction("LD ST, Vx", "Fx18", "ST = Vx",
              function (c, x)

              end)
  Instruction("ADD I, Vx", "Fx1E", "I = I + Vx",
              function (c, x)

              end)
  Instruction("LD F, Vx", "Fx29", "I = address to sprite representing number in Vx",
              function (c, x)

              end)
  Instruction("LD B, Vx", "Fx33", "BCD of Vx is saved to I, I+1, and I+2",
              function (c, x)

              end)
  Instruction("LD [I], Vx", "Fx55", "V0 to Vx are safed to memory beginning at I",
              function (c, x)

              end)
  Instruction("LD Vx, [I]", "Fx65", "V0 to Vx are read from memory beginning at I",
              function (c, x)

              end)
]

"Parses the opcode description and returns mask, pattern to match and mask, type and name of parameters"
function parseopcode(op::String)
  mask = 0x0000
  match = 0x0000
  parameters = Dict{Symbol, Int}()

  for c in op
    mask <<= 4
    match <<= 4
    if isxdigit(c)  # We have a valid hex digit
      mask |= 0xF
      match |= parse(UInt8, c, base = 16)
    else            # We have something else (i.e. a parameter)
      parameters[Symbol(c)] = get(parameters, Symbol(c), 0) + 1
      # Shift mask and match one nibble to the left to set next hex digit
    end
  end

  # Convert number of occurences of parameter into correct type
  parameters = Dict(s => c < 3 ? UInt8 : UInt16 for (s,c) in parameters)

  return mask, match, parameters
end

end # module


@generated function addscalar(v::Vec{W,T}, s::T) where {W, T <: IntegerTypesHW}
    typ = "i$(8sizeof(T))"
    vtyp = "<$W x $typ>"
    instrs = String[]
    push!(instrs, "%ie = insertelement $vtyp zeroinitializer, $typ %1, i32 0")
    push!(instrs, "%v = add $vtyp %0, %ie")
    push!(instrs, "ret $vtyp %v")
    quote
        $(Expr(:meta,:inline))
        Vec($LLVMCALL( $(join(instrs,"\n")), NTuple{$W,Core.VecElement{$T}}, Tuple{NTuple{$W,Core.VecElement{$T}},$T}, data(v), s ))
    end
end
@generated function addscalar(v::Vec{W,T}, s::T) where {W, T <: FloatingTypes}
    typ = LLVM_TYPES[T]
    vtyp = "<$W x $typ>"
    instrs = String[]
    push!(instrs, "%ie = insertelement $vtyp zeroinitializer, $typ %1, i32 0")
    push!(instrs, "%v = fadd nsz arcp contract afn reassoc $vtyp %0, %ie")
    push!(instrs, "ret $vtyp %v")
    quote
        $(Expr(:meta,:inline))
        Vec($LLVMCALL( $(join(instrs,"\n")), NTuple{$W,Core.VecElement{$T}}, Tuple{NTuple{$W,Core.VecElement{$T}},$T}, data(v), s ))
    end
end

@generated function mulscalar(v::Vec{W,T}, s::T) where {W, T <: IntegerTypesHW}
    typ = "i$(8sizeof(T))"
    vtyp = "<$W x $typ>"
    instrs = String[]
    push!(instrs, "%ie = insertelement $vtyp $(llvmconst(W, T, 1)), $typ %1, i32 0")
    push!(instrs, "%v = mul $vtyp %0, %ie")
    push!(instrs, "ret $vtyp %v")
    quote
        $(Expr(:meta,:inline))
        Vec($LLVMCALL( $(join(instrs,"\n")), NTuple{$W,Core.VecElement{$T}}, Tuple{NTuple{$W,Core.VecElement{$T}},$T}, data(v), s ))
    end
end
@generated function mulscalar(v::Vec{W,T}, s::T) where {W, T <: FloatingTypes}
    typ = LLVM_TYPES[T]
    vtyp = "<$W x $typ>"
    instrs = String[]
    push!(instrs, "%ie = insertelement $vtyp $(llvmconst(W, T, 1.0)), $typ %1, i32 0")
    push!(instrs, "%v = fmul nsz arcp contract afn reassoc $vtyp %0, %ie")
    push!(instrs, "ret $vtyp %v")
    quote
        $(Expr(:meta,:inline))
        Vec($LLVMCALL( $(join(instrs,"\n")), NTuple{$W,Core.VecElement{$T}}, Tuple{NTuple{$W,Core.VecElement{$T}},$T}, data(v), s ))
    end
end

function scalar_maxmin(W, ::Type{T}, ismax) where {T}
    comp = if T <: Signed
        ismax ? "icmp sgt" : "icmp slt"
    elseif T <: Unsigned
        ismax ? "icmp ugt" : "icmp ult"
    elseif ismax
        "fcmp ogt"
    else
        "fcmp olt"
    end
    if T <: Integer
        typ = "i$(8sizeof(T))"
        basevalue = llvmconst(W, T, ismax ? typemin(T) : typemax(T))
    else
        opzero = ismax ? -Inf : Inf
        basevalue = llvmconst(W, T, repr(reinterpret(UInt64, opzero)))
        if T === Float64
            typ = "double"
            # basevalue = llvmconst(W, T, repr(reinterpret(UInt64, opzero)))
        elseif T === Float32
            typ = "float"
            # basevalue = llvmconst(W, T, repr(reinterpret(UInt32, Float32(opzero))))
        # elseif T === Float16
            # typ = "half"
            # basevalue = llvmconst(W, T, repr(reinterpret(UInt16, Float16(opzero))))
        else
            throw("T === $T not currently supported.")
        end
    end
    vtyp = "<$W x $typ>"
    instrs = String[]
    push!(instrs, "%ie = insertelement $vtyp $(basevalue), $typ %1, i32 0")
    push!(instrs, "%selection = $comp $vtyp %0, %ie")
    push!(instrs, "%v = select <$W x i1> %selection, $vtyp %0, $vtyp %ie")
    push!(instrs, "ret $vtyp %v")
    instrs
end
@generated function maxscalar(v::Vec{W,T}, s::T) where {W, T <: NativeTypes}
    instrs = scalar_maxmin(W, T, true)
    quote
        $(Expr(:meta,:inline))
        Vec($LLVMCALL( $(join(instrs,"\n")), NTuple{$W,Core.VecElement{$T}}, Tuple{NTuple{$W,Core.VecElement{$T}},$T}, data(v), s ))
    end
end
@generated function minscalar(v::Vec{W,T}, s::T) where {W, T <: NativeTypes}
    instrs = scalar_maxmin(W, T, false)
    quote
        $(Expr(:meta,:inline))
        Vec($LLVMCALL( $(join(instrs,"\n")), NTuple{$W,Core.VecElement{$T}}, Tuple{NTuple{$W,Core.VecElement{$T}},$T}, data(v), s ))
    end
end
for (f,op) ∈ [(:addscalar,:(+)), (:mulscalar,:(*)), (:maxscalar,:max), (:minscalar,:min)]
    @eval begin
        @inline $f(v::VecUnroll, s) = VecUnroll(($f(first(getfield(v, :data)), s), Base.tail(getfield(v, :data))...))
        @inline $f(v::Vec{W,T}, s::NativeTypes) where {W, T <: NativeTypes} = $f(v, vconvert(T, s))
        @inline $f(s::NativeTypes, v::AbstractSIMD{W,T}) where {W,T <: NativeTypes} = $f(v, s)
        @inline $f(a, b) = $op(a, b)
    end
end


colors_tum = Dict{Symbol, RGB{Float64}}(
    :tum_main_blue     => RGB{Float64}(  0/255, 101/255, 189/255),
    :tum_sec_blue1     => RGB{Float64}(  0/255,  82/255, 147/255),
    :tum_sec_blue2     => RGB{Float64}(  0/255,  51/255,  89/255),
    :tum_sec_gray1     => RGB{Float64}( 51/255,  51/255,  51/255),
    :tum_sec_gray2     => RGB{Float64}(128/255, 128/255, 128/255),
    :tum_sec_gray3     => RGB{Float64}(204/255, 204/255, 204/255),
    :tum_accent_gray   => RGB{Float64}(218/255, 215/255, 203/255),
    :tum_accent_orange => RGB{Float64}(227/255, 114/255,  34/255),
    :tum_accent_green  => RGB{Float64}(162/255, 173/255,   0/255),
    :tum_accent_blue1  => RGB{Float64}(152/255, 198/255, 234/255),
    :tum_accent_blue2  => RGB{Float64}(100/255, 160/255, 200/255),
)


export VMPOPolicy, VMPOTrajectory

using LinearAlgebra: ⋅  # dot product
using Zygote: ignore, dropgrad

## trajectory

const VMPOTrajectory = CircularArraySARTTrajectory

## policy definition

"""
    VMPOPolicy(;kwargs)

V-MPO, an on-policy adaptation of Maximum a Posteriori Policy Optimization (MPO)
that performs policy iteration based on a learned state-value function.

# Keyword arguments

- `approximator`: an [`ActorCritic`](@ref) based on [`NeuralNetworkApproximator`](@ref)
- `update_freq`: update policy every n timesteps
- `γ = 0.99f0`: discount factor
- `ϵ_η = 0.02f0`: temperature η hyperparameter
- `ϵ_α = 0.1f0`: Lagrange multiplier α (discrete) hyperparameter
- `ϵ_αμ = 0.005f0`: Lagrange multiplier α_mu (continuous) hyperparameter
- `ϵ_ασ = 0.00005f0`: Lagrange multiplier α_σ (continuous) hyperparameter
- `n_epochs = 8`: update policy for n epochs
- `dist = Categorical`: `Categorical` - discrete, `Normal` - continuous
- `rng = Random.GLOBAL_RNG`

By default, `dist` is set to `Categorical`, which means it will only work
on environments of discrete actions. To work with environments of continuous
actions `dist` should be set to `Normal` and the `actor` in the `approximator`
should be a `GaussianNetwork`. This algorithm only supports one-dimensional
action space for now.

# Ref paper

[V-MPO: On-Policy Maximum a Posteriori Policy Optimization for Discrete and Continuous Control](https://arxiv.org/abs/1909.12238)
"""
mutable struct VMPOPolicy{A<:ActorCritic,D,R} <: AbstractPolicy
    approximator::A
    γ::Float32
    ϵ_η::Float32
    ϵ_α::Float32
    ϵ_αμ::Float32
    ϵ_ασ::Float32
    η::Vector{Float32}  # 1-element array for temperature η
    α::Vector{Float32}  # 3-element array for Lagrange multipliers: [α, αμ, ασ]
    n_epochs::Int
    update_freq::Int
    update_step::Int
    rng::R
end

function VMPOPolicy(;
    approximator,
    update_freq,
    γ=0.99f0,
    ϵ_η=0.02f0,
    ϵ_α=0.1f0,
    ϵ_αμ=0.005f0,
    ϵ_ασ=0.00005f0,
    n_epochs=8,
    dist=Categorical,
    rng=Random.GLOBAL_RNG
)
    VMPOPolicy{typeof(approximator),dist,typeof(rng)}(
        approximator,
        γ,
        ϵ_η,
        ϵ_α,
        ϵ_αμ,
        ϵ_ασ,
        [1.0f0],                # [η]
        [1.0f0, 1.0f0, 1.0f0],  # [α, αμ, ασ]
        n_epochs,
        update_freq,
        0,
        rng,
    )
end

## policy environment interactions

# discrete action

function RLBase.prob(policy::VMPOPolicy{<:ActorCritic,Categorical}, env::AbstractEnv)
    s = send_to_device(device(policy.approximator), state(env))
    p = policy.approximator.actor(s) |> softmax |> send_to_host
    Categorical(p; check_args=false)
end

function (agent::Agent{<:VMPOPolicy{<:ActorCritic,Categorical}})(env::AbstractEnv)
    dist = prob(agent.policy, env)
    rand(agent.policy.rng, dist)
end

function (policy::VMPOPolicy{<:ActorCritic,Categorical})(
    state::AbstractArray,
    action::AbstractArray,
)
    p = policy.approximator.actor(state) |> softmax
    idx = ignore() do
        CartesianIndex.(action, 1:length(action))
    end
    action_log_prob = log.(p)[idx]
    p, action_log_prob
end

# continuous action

function RLBase.prob(
    policy::VMPOPolicy{<:ActorCritic{<:GaussianNetwork},Normal},
    env::AbstractEnv,
)
    s = send_to_device(device(policy.approximator), state(env))
    μ, logσ = policy.approximator.actor(agent.policy.rng, s)
    Normal(μ, exp(logσ))
end

function (agent::Agent{<:VMPOPolicy{<:ActorCritic{<:GaussianNetwork},Normal}})(
    env::AbstractEnv,
)
    s = send_to_device(device(agent.policy.approximator), state(env))
    # the action is an output of GaussianNetwork which is normalised by tanh(),
    # we increase its stability by limiting it to [-1 + eps, 1 - eps]
    a = agent.policy.approximator.actor(agent.policy.rng, s, is_sampling=true)
    m = one(eltype(a)) - eps(eltype(a))
    clamp.(a, -m, m) |> send_to_host |> first
end

function (policy::VMPOPolicy{<:ActorCritic{<:GaussianNetwork},Normal})(
    state::AbstractArray,
    action::AbstractArray,
)
    μ, logσ = policy.approximator.actor(policy.rng, state)
    action = atanh.(action)
    μ, logσ, normlogpdf(μ, exp.(logσ), reshape(action, size(μ)...))
end

## update policy

function RLBase.update!(p::VMPOPolicy, t::VMPOTrajectory, env::AbstractEnv, ::PreActStage)
    # in the first update, only state & action are inserted into trajectory
    length(t) == 0 && return

    p.update_step += 1
    if p.update_step % p.update_freq == 0
        _update!(p, t)
    end
end

function _update!(p::VMPOPolicy, t::VMPOTrajectory)
    AC = p.approximator
    D = device(AC)
    s = send_to_device(D, t[:state][:, 1:end-1])  # drop the last extra state
    a = send_to_device(D, t[:action][1:end-1])  # drop the last extra action
    is_discrete = isa(p, VMPOPolicy{<:ActorCritic,Categorical})
    π_old = is_discrete ? AC.actor(s) |> softmax : nothing
    μ_old, σ_old = is_discrete ? (nothing, nothing) : begin
        μ, logσ = AC.actor(s)
        μ, exp.(logσ)
    end

    # normalise rewards based on each trajectory
    inds = vcat(0, findall(t[:terminal]), length(t[:terminal]))
    rewards =
        map(
            Flux.normalise,
            view.(Ref(t[:reward]), (:).(inds[1:end-1] .+ 1, inds[2:end])),
        ) |>
        Iterators.flatten |>
        collect

    # calculate discounted rewards
    rewards = discount_rewards(rewards, p.γ, terminal=t[:terminal])
    rewards = send_to_device(D, rewards)

    # calculate advantages and sample top half of it
    advantages = rewards .- AC.critic(s)[1, :]
    top_k_idx = sortperm(advantages |> send_to_host, rev=true)[1:length(advantages)÷2]
    top_k_advs = advantages[top_k_idx]

    for epoch in 1:p.n_epochs
        ps = Flux.params(AC, p.η, p.α)
        gs = gradient(ps) do
            η, α, αμ, ασ = p.η[1], p.α...
            π, logπ = is_discrete ? p(s, a) : (nothing, nothing)
            μ, σ, logπ =
                is_discrete ? (nothing, nothing, logπ) : begin
                    μ, logσ, logπ = p(s, a)
                    μ, exp.(logσ), logπ
                end

            # critic loss
            v = AC.critic(s)[1, :]
            Lᵥ = Flux.Losses.mse(v, rewards) / 2

            # actor loss
            ψ = softmax(top_k_advs / dropgrad(η))
            Lπ = -ψ ⋅ logπ[top_k_idx]

            # loss for temperature η
            Lη = η * p.ϵ_η + η * log(mean(exp.(top_k_advs / η)))

            # loss for Lagrange multiplier α
            Lα = if is_discrete
                KL = Flux.kldivergence(π_old, π, agg=identity)
                mean(α * (p.ϵ_α .- dropgrad(KL)) + dropgrad(α) * KL)
            else
                KLμ = 0.5f0 * (μ .- μ_old) .^ 2 ./ (σ .^ 2)
                Lαμ = mean(αμ * (p.ϵ_αμ .- dropgrad(KLμ)) + dropgrad(αμ) * KLμ)
                KLσ = 0.5f0 * ((σ_old ./ σ) .^ 2 .- 1 .+ 2 * log.(σ ./ σ_old))
                Lασ = mean(ασ * (p.ϵ_ασ .- dropgrad(KLσ)) + dropgrad(ασ) * KLσ)
                Lαμ + Lασ
            end

            loss = Lᵥ + Lπ + Lη + Lα
        end

        # skip if gradient contains NaN or Inf
        if any(x -> !isnothing(x) && any(y -> isnan(y) || isinf(y), x), gs)
            break
        end

        Flux.Optimise.update!(AC.optimizer, ps, gs)

        p.η = clamp.(p.η, 1e-8, p.η)
        p.α = clamp.(p.α, 1e-8, p.α)
    end
end


@testset "inference.jl" begin
    # Only test threading if 1.3+.
    if VERSION > v"1.2"
        @testset "threaded sampling" begin
            # Test that chains with the same seed will sample identically.
            @testset "rng" begin
                model = gdemo_default

                # multithreaded sampling with PG causes segfaults on Julia 1.5.4
                # https://github.com/TuringLang/Turing.jl/issues/1571
                samplers = @static if VERSION <= v"1.5.3" || VERSION >= v"1.6.0"
                    (
                        HMC(0.1, 7),
                        PG(10),
                        IS(),
                        MH(),
                        Gibbs(PG(3, :s), HMC(0.4, 8, :m)),
                        Gibbs(HMC(0.1, 5, :s), ESS(:m)),
                    )
                else
                    (
                        HMC(0.1, 7),
                        IS(),
                        MH(),
                        Gibbs(HMC(0.1, 5, :s), ESS(:m)),
                    )
                end
                for sampler in samplers
                    Random.seed!(5)
                    chain1 = sample(model, sampler, MCMCThreads(), 1000, 4)

                    Random.seed!(5)
                    chain2 = sample(model, sampler, MCMCThreads(), 1000, 4)

                    @test chain1.value == chain2.value
                end
            end

            # Smoke test for default sample call.
            Random.seed!(100)
            chain = sample(gdemo_default, HMC(0.1, 7), MCMCThreads(), 1000, 4)
            check_gdemo(chain)

            # run sampler: progress logging should be disabled and
            # it should return a Chains object
            sampler = Sampler(HMC(0.1, 7), gdemo_default)
            chains = sample(gdemo_default, sampler, MCMCThreads(), 1000, 4)
            @test chains isa MCMCChains.Chains
        end
    end
    @testset "chain save/resume" begin
        Random.seed!(1234)

        alg1 = HMCDA(1000, 0.65, 0.15)
        alg2 = PG(20)
        alg3 = Gibbs(PG(30, :s), HMC(0.2, 4, :m))

        chn1 = sample(gdemo_default, alg1, 5000; save_state=true)
        check_gdemo(chn1)

        chn1_resumed = Turing.Inference.resume(chn1, 1000)
        check_gdemo(chn1_resumed)

        chn1_contd = sample(gdemo_default, alg1, 5000; resume_from=chn1)
        check_gdemo(chn1_contd)

        chn1_contd2 = sample(gdemo_default, alg1, 5000; resume_from=chn1, reuse_spl_n=1000)
        check_gdemo(chn1_contd2)

        chn2 = sample(gdemo_default, alg2, 1000; save_state=true)
        check_gdemo(chn2)

        chn2_contd = sample(gdemo_default, alg2, 1000; resume_from=chn2)
        check_gdemo(chn2_contd)

        chn3 = sample(gdemo_default, alg3, 5000; save_state=true)
        check_gdemo(chn3)

        chn3_contd = sample(gdemo_default, alg3, 1000; resume_from=chn3)
        check_gdemo(chn3_contd)
    end
    @testset "Contexts" begin
        # Test LikelihoodContext
        @model function testmodel1(x)
            a ~ Beta()
            lp1 = getlogp(__varinfo__)
            x[1] ~ Bernoulli(a)
            global loglike = getlogp(__varinfo__) - lp1
        end
        model = testmodel1([1.0])
        varinfo = Turing.VarInfo(model)
        model(varinfo, Turing.SampleFromPrior(), Turing.LikelihoodContext())
        @test getlogp(varinfo) == loglike

        # Test MiniBatchContext
        @model function testmodel2(x)
            a ~ Beta()
            x[1] ~ Bernoulli(a)
        end
        model = testmodel2([1.0])
        varinfo1 = Turing.VarInfo(model)
        varinfo2 = deepcopy(varinfo1)
        model(varinfo1, Turing.SampleFromPrior(), Turing.LikelihoodContext())
        model(varinfo2, Turing.SampleFromPrior(), Turing.MiniBatchContext(Turing.LikelihoodContext(), 10))
        @test isapprox(getlogp(varinfo2) / getlogp(varinfo1), 10)
    end
    @testset "Prior" begin
        N = 5000

        # Note that all chains contain 3 values per sample: 2 variables + log probability
        Random.seed!(100)
        chains = sample(gdemo_d(), Prior(), N)
        @test chains isa MCMCChains.Chains
        @test size(chains) == (N, 3, 1)
        @test mean(chains, :s) ≈ 3 atol=0.1
        @test mean(chains, :m) ≈ 0 atol=0.1

        Random.seed!(100)
        chains = sample(gdemo_d(), Prior(), MCMCThreads(), N, 4)
        @test chains isa MCMCChains.Chains
        @test size(chains) == (N, 3, 4)
        @test mean(chains, :s) ≈ 3 atol=0.1
        @test mean(chains, :m) ≈ 0 atol=0.1

        Random.seed!(100)
        chains = sample(gdemo_d(), Prior(), N; chain_type = Vector{NamedTuple})
        @test chains isa Vector{<:NamedTuple}
        @test length(chains) == N
        @test all(length(x) == 3 for x in chains)
        @test all(haskey(x, :lp) for x in chains)
        @test mean(x[:s][1] for x in chains) ≈ 3 atol=0.1
        @test mean(x[:m][1] for x in chains) ≈ 0 atol=0.1
    end

    @testset "chain ordering" begin
        for alg in (Prior(), Emcee(10, 2.0))
            chain_sorted = sample(gdemo_default, alg, 1, sort_chain=true)
            @test names(MCMCChains.get_sections(chain_sorted, :parameters)) == [:m, :s]

            chain_unsorted = sample(gdemo_default, alg, 1, sort_chain=false)
            @test names(MCMCChains.get_sections(chain_unsorted, :parameters)) == [:s, :m]
        end
    end

    @testset "chain iteration numbers" begin
        for alg in (Prior(), Emcee(10, 2.0))
            chain = sample(gdemo_default, alg, 10)
            @test range(chain) == 1:10

            chain = sample(gdemo_default, alg, 10; discard_initial=5, thinning=2)
            @test range(chain) == range(6; step=2, length=10)
        end
    end

    # Copy-paste from integration tests in DynamicPPL.
    @testset "assume" begin
        @model function test_assume()
            x ~ Bernoulli(1)
            y ~ Bernoulli(x / 2)
            return x, y
        end

        smc = SMC()
        pg = PG(10)

        res1 = sample(test_assume(), smc, 1000)
        res2 = sample(test_assume(), pg, 1000)

        check_numerical(res1, [:y], [0.5]; atol=0.1)
        check_numerical(res2, [:y], [0.5]; atol=0.1)

        # Check that all xs are 1.
        @test all(isone, res1[:x])
        @test all(isone, res2[:x])
    end
    @testset "beta binomial" begin
        prior = Beta(2, 2)
        obs = [0, 1, 0, 1, 1, 1, 1, 1, 1, 1]
        exact = Beta(prior.α + sum(obs), prior.β + length(obs) - sum(obs))
        meanp = exact.α / (exact.α + exact.β)

        @model function testbb(obs)
            p ~ Beta(2, 2)
            x ~ Bernoulli(p)
            for i in 1:length(obs)
                obs[i] ~ Bernoulli(p)
            end
            return p, x
        end

        smc = SMC()
        pg = PG(10)
        gibbs = Gibbs(HMC(0.2, 3, :p), PG(10, :x))

        chn_s = sample(testbb(obs), smc, 1000)
        chn_p = sample(testbb(obs), pg, 2000)
        chn_g = sample(testbb(obs), gibbs, 1500)

        check_numerical(chn_s, [:p], [meanp]; atol=0.05)
        check_numerical(chn_p, [:x], [meanp]; atol=0.1)
        check_numerical(chn_g, [:x], [meanp]; atol=0.1)
    end
    @testset "forbid global" begin
        xs = [1.5 2.0]
        # xx = 1

        @model function fggibbstest(xs)
            s ~ InverseGamma(2, 3)
            m ~ Normal(0, sqrt(s))
            # xx ~ Normal(m, sqrt(s)) # this is illegal

            for i in 1:length(xs)
                xs[i] ~ Normal(m, sqrt(s))
                # for xx in xs
                # xx ~ Normal(m, sqrt(s))
            end
            return s, m
        end

        gibbs = Gibbs(PG(10, :s), HMC(0.4, 8, :m))
        chain = sample(fggibbstest(xs), gibbs, 2)
    end
    @testset "new grammar" begin
        x = Float64[1 2]

        @model function gauss(x)
            priors = TArray{Float64}(2)
            priors[1] ~ InverseGamma(2, 3)         # s
            priors[2] ~ Normal(0, sqrt(priors[1])) # m
            for i in 1:length(x)
                x[i] ~ Normal(priors[2], sqrt(priors[1]))
            end
            return priors
        end

        chain = sample(gauss(x), PG(10), 10)
        chain = sample(gauss(x), SMC(), 10)

        @model function gauss2(::Type{TV}=Vector{Float64}; x) where {TV}
            priors = TV(undef, 2)
            priors[1] ~ InverseGamma(2, 3)         # s
            priors[2] ~ Normal(0, sqrt(priors[1])) # m
            for i in 1:length(x)
                x[i] ~ Normal(priors[2], sqrt(priors[1]))
            end
            return priors
        end

        chain = sample(gauss2(; x=x), PG(10), 10)
        chain = sample(gauss2(; x=x), SMC(), 10)

        chain = sample(gauss2(Vector{Float64}; x=x), PG(10), 10)
        chain = sample(gauss2(Vector{Float64}; x=x), SMC(), 10)
    end
    @testset "new interface" begin
        obs = [0, 1, 0, 1, 1, 1, 1, 1, 1, 1]

        @model function newinterface(obs)
            p ~ Beta(2, 2)
            for i in 1:length(obs)
                obs[i] ~ Bernoulli(p)
            end
            return p
        end

        chain = sample(
            newinterface(obs), HMC{Turing.ForwardDiffAD{2}}(0.75, 3, :p, :x), 100
        )
    end
    @testset "no return" begin
        Random.seed!(5)
        @model function noreturn(x)
            s ~ InverseGamma(2, 3)
            m ~ Normal(0, sqrt(s))
            for i in 1:length(x)
                x[i] ~ Normal(m, sqrt(s))
            end
        end

        chain = sample(noreturn([1.5 2.0]), HMC(0.1, 10), 4000)
        check_numerical(chain, [:s, :m], [49 / 24, 7 / 6])
    end
    @testset "observe" begin
        Random.seed!(5)
        @model function test()
            z ~ Normal(0, 1)
            x ~ Bernoulli(1)
            1 ~ Bernoulli(x / 2)
            0 ~ Bernoulli(x / 2)
            return x
        end

        is = IS()
        smc = SMC()
        pg = PG(10)

        res_is = sample(test(), is, 10000)
        res_smc = sample(test(), smc, 1000)
        res_pg = sample(test(), pg, 100)

        @test all(isone, res_is[:x])
        @test res_is.logevidence ≈ 2 * log(0.5)

        @test all(isone, res_smc[:x])
        @test res_smc.logevidence ≈ 2 * log(0.5)

        @test all(isone, res_pg[:x])
    end
    @testset "sample" begin
        alg = Gibbs(HMC(0.2, 3, :m), PG(10, :s))
        chn = sample(gdemo_default, alg, 1000)
    end
    @testset "vectorization @." begin
        # https://github.com/FluxML/Tracker.jl/issues/119
        if Turing.Core.ADBackend() !== Turing.Core.TrackerAD
            @model function vdemo1(x)
                s ~ InverseGamma(2, 3)
                m ~ Normal(0, sqrt(s))
                @. x ~ Normal(m, sqrt(s))
                return s, m
            end

            alg = HMC(0.01, 5)
            x = randn(100)
            res = sample(vdemo1(x), alg, 250)

            @model function vdemo1b(x)
                s ~ InverseGamma(2, 3)
                m ~ Normal(0, sqrt(s))
                @. x ~ Normal(m, $(sqrt(s)))
                return s, m
            end

            res = sample(vdemo1b(x), alg, 250)

            @model function vdemo2(x)
                μ ~ MvNormal(zeros(size(x, 1)), I)
                @. x ~ $(MvNormal(μ, I))
            end

            D = 2
            alg = HMC(0.01, 5)
            res = sample(vdemo2(randn(D, 100)), alg, 250)

            # Vector assumptions
            N = 10
            setchunksize(N)
            alg = HMC(0.2, 4)

            @model function vdemo3()
                x = Vector{Real}(undef, N)
                for i in 1:N
                    x[i] ~ Normal(0, sqrt(4))
                end
            end

            t_loop = @elapsed res = sample(vdemo3(), alg, 1000)

            # Test for vectorize UnivariateDistribution
            @model function vdemo4()
                x = Vector{Real}(undef, N)
                @. x ~ Normal(0, 2)
            end

            t_vec = @elapsed res = sample(vdemo4(), alg, 1000)

            @model vdemo5() = x ~ MvNormal(zeros(N), 4 * I)

            t_mv = @elapsed res = sample(vdemo5(), alg, 1000)

            println("Time for")
            println("  Loop : ", t_loop)
            println("  Vec  : ", t_vec)
            println("  Mv   : ", t_mv)

            # Transformed test
            @model function vdemo6()
                x = Vector{Real}(undef, N)
                @. x ~ InverseGamma(2, 3)
            end

            sample(vdemo6(), alg, 1000)

            N = 3
            @model function vdemo7()
                x = Array{Real}(undef, N, N)
                @. x ~ [InverseGamma(2, 3) for i in 1:N]
            end

            sample(vdemo7(), alg, 1000)
        end
    end
    @testset "vectorization .~" begin
        @model function vdemo1(x)
            s ~ InverseGamma(2, 3)
            m ~ Normal(0, sqrt(s))
            x .~ Normal(m, sqrt(s))
            return s, m
        end

        alg = HMC(0.01, 5)
        x = randn(100)
        res = sample(vdemo1(x), alg, 250)

        @model function vdemo2(x)
            μ ~ MvNormal(zeros(size(x, 1)), I)
            return x .~ MvNormal(μ, I)
        end

        D = 2
        alg = HMC(0.01, 5)
        res = sample(vdemo2(randn(D, 100)), alg, 250)

        # Vector assumptions
        N = 10
        setchunksize(N)
        alg = HMC(0.2, 4)

        @model function vdemo3()
            x = Vector{Real}(undef, N)
            for i in 1:N
                x[i] ~ Normal(0, sqrt(4))
            end
        end

        t_loop = @elapsed res = sample(vdemo3(), alg, 1000)

        # Test for vectorize UnivariateDistribution
        @model function vdemo4()
            x = Vector{Real}(undef, N)
            return x .~ Normal(0, 2)
        end

        t_vec = @elapsed res = sample(vdemo4(), alg, 1000)

        @model vdemo5() = x ~ MvNormal(zeros(N), 4 * I)

        t_mv = @elapsed res = sample(vdemo5(), alg, 1000)

        println("Time for")
        println("  Loop : ", t_loop)
        println("  Vec  : ", t_vec)
        println("  Mv   : ", t_mv)

        # Transformed test
        @model function vdemo6()
            x = Vector{Real}(undef, N)
            return x .~ InverseGamma(2, 3)
        end

        sample(vdemo6(), alg, 1000)

        @model function vdemo7()
            x = Array{Real}(undef, N, N)
            return x .~ [InverseGamma(2, 3) for i in 1:N]
        end

        sample(vdemo7(), alg, 1000)
    end
    @testset "Type parameters" begin
        N = 10
        setchunksize(N)
        alg = HMC(0.01, 5)
        x = randn(1000)
        @model function vdemo1(::Type{T}=Float64) where {T}
            x = Vector{T}(undef, N)
            for i in 1:N
                x[i] ~ Normal(0, sqrt(4))
            end
        end

        t_loop = @elapsed res = sample(vdemo1(), alg, 250)
        t_loop = @elapsed res = sample(vdemo1(Float64), alg, 250)

        vdemo1kw(; T) = vdemo1(T)
        t_loop = @elapsed res = sample(vdemo1kw(; T=Float64), alg, 250)

        @model function vdemo2(::Type{T}=Float64) where {T<:Real}
            x = Vector{T}(undef, N)
            @. x ~ Normal(0, 2)
        end

        t_vec = @elapsed res = sample(vdemo2(), alg, 250)
        t_vec = @elapsed res = sample(vdemo2(Float64), alg, 250)

        vdemo2kw(; T) = vdemo2(T)
        t_vec = @elapsed res = sample(vdemo2kw(; T=Float64), alg, 250)

        @model function vdemo3(::Type{TV}=Vector{Float64}) where {TV<:AbstractVector}
            x = TV(undef, N)
            @. x ~ InverseGamma(2, 3)
        end

        sample(vdemo3(), alg, 250)
        sample(vdemo3(Vector{Float64}), alg, 250)

        vdemo3kw(; T) = vdemo3(T)
        sample(vdemo3kw(; T=Vector{Float64}), alg, 250)
    end
end


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

using JuliaFEM
using LightXML

# element codes: http://www.paraview.org/pipermail/paraview/2013-July/028859.html
# > from  ./VTK/ThirdParty/xdmf2/vtkxdmf2/libsrc/XdmfTopology.h
# >
# > // Topologies
# > #define XDMF_NOTOPOLOGY     0x0
# > #define XDMF_POLYVERTEX     0x1
# > #define XDMF_POLYLINE       0x2
# > #define XDMF_POLYGON        0x3
# > #define XDMF_TRI            0x4
# > #define XDMF_QUAD           0x5
# > #define XDMF_TET            0x6
# > #define XDMF_PYRAMID        0x7
# > #define XDMF_WEDGE          0x8
# > #define XDMF_HEX            0x9
# > #define XDMF_EDGE_3         0x0022
# > #define XDMF_TRI_6          0x0024
# > #define XDMF_QUAD_8         0x0025
# > #define XDMF_QUAD_9         0x0023
# > #define XDMF_TET_10         0x0026
# > #define XDMF_PYRAMID_13     0x0027
# > #define XDMF_WEDGE_15       0x0028
# > #define XDMF_WEDGE_18       0x0029
# > #define XDMF_HEX_20         0x0030
# > #define XDMF_HEX_24         0x0031
# > #define XDMF_HEX_27         0x0032
# > #define XDMF_MIXED          0x0070
# > #define XDMF_2DSMESH        0x0100
# > #define XDMF_2DRECTMESH     0x0101
# > #define XDMF_2DCORECTMESH   0x0102
# > #define XDMF_3DSMESH        0x1100
# > #define XDMF_3DRECTMESH     0x1101
# > #define XDMF_3DCORECTMESH   0x1102

get_xdmf_element_code(element::Element{Poi1})  = 0x0001
get_xdmf_element_code(element::Element{Seg2})  = 0x0002
get_xdmf_element_code(element::Element{Seg3})  = 0x0003
get_xdmf_element_code(element::Element{Tri3})  = 0x0004
get_xdmf_element_code(element::Element{Quad4}) = 0x0005
get_xdmf_element_code(element::Element{Tet4})  = 0x0006
get_xdmf_element_code(element::Element{Hex8})  = 0x0009
get_xdmf_element_code(element::Element{Tet10}) = 0x0026

type XDMF
    dimension :: Int
    use_hdf :: Bool
    xdoc :: XMLDocument
    domain :: XMLElement
    temporal_collection :: XMLElement
    current_grid
    permutation :: Vector{Int}
end

function XDMF()
    xdoc = XMLDocument()
    xroot = create_root(xdoc, "Xdmf")
    set_attribute(xroot, "xmlns:xi", "http://www.w3.org/2001/XInclude")
    set_attribute(xroot, "Version", "2.1")
    domain = new_child(xroot, "Domain")
    temporal_collection = new_child(domain, "Grid")
    set_attribute(temporal_collection, "CollectionType", "Temporal")
    set_attribute(temporal_collection, "GridType", "Collection")
    set_attribute(temporal_collection, "Name", "Collection")
    return XDMF(3, false, xdoc, domain, temporal_collection, Union{}, [])
end

function xdmf_new_result!(xdmf::XDMF, elements::Vector, time)
    grid = new_child(xdmf.temporal_collection, "Grid")
    set_attribute(grid, "Name", "Grid")
    time_ = new_child(grid, "Time")
    set_attribute(time_, "Value", time)
    xdmf.current_grid = grid

    # 1. calculate permutation
    nids = Set()
    X = Dict{Int64, Vector{Float64}}()
    for element in elements
        conn = get_connectivity(element)
        push!(nids, conn...)
        X_el = element["geometry"](time)
        for (i, c) in enumerate(conn)
            X[c] = X_el[i]
        end
    end
    xdmf.permutation = sort(collect(nids))
    iperm = Dict{Int64, Int64}()
    for (i, j) in enumerate(xdmf.permutation)
        iperm[j] = i
    end

    # 2. write nodes
    geometry = new_child(grid, "Geometry")
    set_attribute(geometry, "Type", xdmf.dimension == 3 ? "XYZ" : "XY")
    dataitem = new_child(geometry, "DataItem")
    set_attribute(dataitem, "DataType", "Float")
    set_attribute(dataitem, "Format", "XML")
    #set_attribute(dataitem, "Precision", 8)
    s = []
    ndim = 0
    for i in xdmf.permutation
        ndim += length(X[i])
        push!(s, join(round(X[i], 5), " "))
    end
    set_attribute(dataitem, "Dimensions", ndim)
    add_text(dataitem, "\n"*join(s, "\n")*"\n")

    # 3. write elements
    topology = new_child(grid, "Topology")
    set_attribute(topology, "TopologyType", "Mixed")
    set_attribute(topology, "NumberOfElements", length(elements))
    dataitem = new_child(topology, "DataItem")
    set_attribute(dataitem, "Format", "XML")
    set_attribute(dataitem, "DataType", "Int")
#   set_attribute(dataitem, "Precision", 8)
    s = []
    eldim = 0
    for element in elements
        eltype = get_xdmf_element_code(element)
        # note: id numbers start from 0 in Xdmf
        conn = [iperm[j] for j in get_connectivity(element)] - 1
        data = [eltype; conn]
        eldim += length(data)
        push!(s, join(data, " "))
    end
    set_attribute(dataitem, "Dimensions", eldim)
    add_text(dataitem, "\n"*join(s, "\n")*"\n")
end

function xdmf_save_field!(xdmf, elements::Vector, time, field_name; field_type="Scalar", debug=false)
    f = Dict()
    field_dim = 0
    for element in elements
        haskey(element, field_name) || continue
        g = element[field_name](time)
        conn = get_connectivity(element)
        for (i, c) in enumerate(conn)
            gi = g[i]
            if (field_type == "Vector") && (length(gi) < 3)
                # paraview goes crazy if 2d model with 2d displacement vector
                gi = [gi; 0.0]
            end
            if field_dim == 0
                field_dim = length(gi)
            end
            field_dim == length(gi) || error("several dimensions in field, dim = $field_dim.")
            f[c] = gi
        end
    end

    if length(f) == 0
        warn("xdmf_save_field!(): field $field_name was not found from set of elements")
        return
    end

    attribute = new_child(xdmf.current_grid, "Attribute")
    set_attribute(attribute, "Center", "Node")
    set_attribute(attribute, "Name", ucfirst(field_name))
    set_attribute(attribute, "Type", field_type)
    dataitem = new_child(attribute, "DataItem")
    set_attribute(dataitem, "DataType", "Float")
    set_attribute(dataitem, "Format", "XML")
    #set_attribute(dataitem, "Precision", 8)
    debug && info("field dim = $field_dim")
    debug && info(f)
    s = []
    dim = 0
    for i in xdmf.permutation
        gi = zeros(field_dim)
        if haskey(f, i)
            gi = f[i]
        end
        push!(s, join(round(gi, 5), " "))
        dim += length(gi)
    end
    set_attribute(dataitem, "Dimensions", dim)
    add_text(dataitem, "\n"*join(s, "\n")*"\n")
end

function xdmf_save_field!(xdmf, problem::Problem, time, field_name; field_type="Scalar")
    xdmf_save_field!(xdmf, problem.elements, time, field_name; field_type=field_type)
end

function xdmf_new_result!(xdmf, problem::Problem, time)
    xdmf_new_result!(xdmf, problem.elements, time)
end

function xdmf_save!(xdmf, filename)
    save_file(xdmf.xdoc, filename)
end


import .QueueUnion: Range, Queues
import Printf: @printf

const EPS = 1e-10
const DEBUG = false

@deprecate(
    clip_front(elements, pqs, i, slope, offset, t),
    clip(elements, Ref(pqs, i), +slope, +offset - t, Val(true))
)

@deprecate(
    clip_back(elements, pqs, i, slope, offset, t),
    clip(elements, Ref(pqs, i), -slope, -offset + t, Val(false))
)

@deprecate(
    clip_front(qs::Queues{Event}, i::I, slope::F, offset::F, t::F) where {F,I},
    clip(qs, i, slope, offset - t, Val(true))
)

@deprecate(
    clip_back(qs::Queues{Event}, i::I, slope::F, offset::F, t::F) where {F,I},
    clip(qs, i, -slope, -offset + t, Val(false))
)

clip(qs::Queues{Event}, i::I, slope::F, offset::F, v::Val{D}) where {F,I,D} =
    clip(qs.events, Ref(qs.pq, i), slope, offset, v)

function clip(
    elements::Vector{Event},
    pq::Ref{Range},
    slope::F,
    offset::F,
    ::Val{forward},
    ::Val{check},
)::F where {F,forward,check}
    local start::Int = pq[].start
    local stop::Int = pq[].stop
    local e::Event = elements[forward ? start : stop]::Event
    @static DEBUG && local dir::String = forward ? "f" : "b"
    @static DEBUG && @printf("clip_%s: (%+g, %+.2f) [%d] %s\n",
                             dir, slope, offset, check, string(elements[start:stop]))
    while start <= stop && slope * e.x + offset < 0
        offset += intercept(e)
        slope += e.slope
        @static DEBUG && @printf(" lip_%s: (%+g, %+.2f)\n", dir, slope, offset)
        e = elements[forward ? start += 1 : stop -= 1]
    end
    local x::F = if check && abs(slope) <= EPS
        forward ? -Inf : +Inf
    else
        let x::F = -offset/slope
            elements[forward ? start -= 1 : stop += 1] = Event(x, slope)
            x
        end
    end
    @static DEBUG && @printf("  ip_%s:  --> %g\n", dir, x)
    pq[] = Range(start, stop)
    return x
end


"""
    cas_parse(str)

given a CAS number string, returns a tuple of 3 Int32 containing the CAS numbers, it performs no validation on the data

"""
function cas_parse(str)
    a1,a2,a3 = split(str,'-')
    n1 = parse(Int32,a1)
    n2 = parse(Int16,a2)
    n3 = parse(Int16,a3)
    return (n1,n2,n3)
end

function unique_idxs_sorted(vec)
    n = length(vec)
    res = zeros(Int,n)
    k = 1
    res[1] = k
    @inbounds for i = 2:n
        if !isequal(vec[i],vec[i-1])
            k+=1
            res[k] = i
        end
    end
    resize!(res,k)
    return res
end
"""
    load_db!(dbtype::Symbol)
    downloads, processes and stores a database corresponding to the one with the same key stored in DATA_INFO

"""

function load_db!(dbtype::Symbol)
    data = DATA_INFO[dbtype]
    if !isfile(data.textdb)
        @info ":" * string(dbtype) * " database file not found, downloading from " * data.url
        url  = data.url
        fname = data.textdb
        path = Downloads.download(url,fname)
        @info ":" * string(dbtype) *  " database file downloaded."
    end

    path = data.textdb
    if !isfile(data.db)
        @info ":" * string(dbtype)  * " arrow file not generated, processing..."
        arrow_db,arrow_synonym_db,arrow_sort_db =parse_and_write_db!(dbtype)
    else
        arrow_db = Arrow.Table(data.db)
        arrow_synonym_db = Arrow.Table(data.symsdb)
        arrow_sort_db = Arrow.Table(data.sorteddb)

    end
    DATA_DB[dbtype] = (arrow_db,arrow_synonym_db,arrow_sort_db)
    return arrow_db,arrow_synonym_db,arrow_sort_db
end

function parse_and_write_db!(dbtype::Symbol)
    data = DATA_INFO[dbtype]
    path = data.textdb
    i = 0
    for _ in eachline(path)
        i +=1
    end

    pubchemid = zeros(Int64,i)
    CAS = Vector{Tuple{Int32,Int16,Int16}}(undef,i)
    formula = Vector{String}(undef,i)
    MW = Vector{Float64}(undef,i)
    smiles = Vector{String}(undef,i)
    InChI = Vector{String}(undef,i)
    InChI_key = Vector{String}(undef,i)
    iupac_name = Vector{String}(undef,i)
    common_name = Vector{String}(undef,i)
    _synonyms  = Vector{Vector{String}}(undef,i)

    i = 0
    for line in eachline(path)
        i += 1
        strs = line |> z->rstrip(z) |> z->split(z,'\t',limit=10)
        pubchemid[i] = parse(Int64,strs[1])
        CAS[i] = cas_parse(strs[2])
        formula[i] = strs[3]
        MW[i] = parse(Float64,strs[4])
        smiles[i] = strs[5]
        InChI[i] = strs[6]
        InChI_key[i] = strs[7]
        iupac_name[i] = strs[8]
        common_name[i] = strs[9]
        if length(strs) >= 10 #not any synonyms
            sym_i = split(strs[10],('\t',';'))
            push!(sym_i,strs[8])
            _synonyms[i]  = sym_i
        else
            _synonyms[i] = String[strs[8]]
        end
    end

    syms_i = mapreduce(length,+,_synonyms)
    synonyms_list = Vector{String}(undef,syms_i)
    synonyms_index = Vector{Int}(undef,syms_i)

    #for some reason,some empty strings are generated as synonyms.
    #those are eliminated here.

    k = 0
    for (ii,sym_vec) in pairs(_synonyms)

        for (jj,sym) in pairs(sym_vec)
            if !isempty(sym)
                k+=1
                synonyms_list[k] = sym
                synonyms_index[k] = ii
            end
        end
    end

    resize!(synonyms_list,k)
    resize!(synonyms_index,k)

    pubchemid_sort =sortperm(pubchemid)
    CAS_sort = sortperm(CAS)
    formula_sort =sortperm(formula)
    MW_sort =sortperm(MW)
    smiles_sort =sortperm(smiles)
    InChI_sort =sortperm(InChI)
    InChI_key_sort =sortperm(InChI_key)
    iupac_name_sort =sortperm(iupac_name)
    common_name_sort =sortperm(common_name)
    synonyms_sort = sortperm(synonyms_list)

    list = synonyms_list[synonyms_sort]
    index = synonyms_index[synonyms_sort]
    #there is the posibility of repeated elements.
    #this is not present in CalebBell/chemicals because it uses a dict, so equal keys
    #store the same value
    list_unique_idx = unique_idxs_sorted(list)
    list = list[list_unique_idx]
    index = index[list_unique_idx]

    db = (;pubchemid, CAS, formula, MW, smiles, InChI, InChI_key, iupac_name, common_name)
    synonym_db = (;list,index)
    sort_db = (;pubchemid_sort, CAS_sort, formula_sort, MW_sort, smiles_sort, InChI_sort, InChI_key_sort, iupac_name_sort, common_name_sort)

    Arrow.write(data.db,db)
    Arrow.write(data.symsdb,synonym_db)
    Arrow.write(data.sorteddb,sort_db)

    arrow_db = Arrow.Table(data.db)
    arrow_synonym_db = Arrow.Table(data.symsdb)
    arrow_sort_db = Arrow.Table(data.sorteddb)
    return arrow_db,arrow_synonym_db,arrow_sort_db
end

"""
    load_data!(key::Symbol;url=nothing,file=nothing)

generates and adds to the global DATA_INFO dict a new database. download and process this database with ´load_db(key)´
"""
function load_data!(key::Symbol;url=nothing,file=nothing)
    if url == file == nothing
        throw(ArgumentError("a file or a url must be provided."))
    elseif (url !== nothing) & (file !== nothing)
        throw(ArgumentError("a file or a url must be provided."))
    elseif (url !== nothing)
        fname = joinpath(download_cache, "pubchem_" * string(key))
        farrow = fname * ".arrow"
        fsyms = fname * "_synonyms.arrow"
        sorted = fname * "_sorted.arrow"
        data = (url=url,textdb=fname,db=farrow,symsdb=fsyms,sorteddb=sorted)

    else #file
        fname = file
        url = fname
        farrow = fname * ".arrow"
        fsyms = fname * "_synonyms.arrow"
        sorted = fname * "_sorted.arrow"
        data = (url=url,textdb=fname,db=farrow,symsdb=fsyms,sorteddb=sorted)
    end
    DATA_INFO[key] = data
end

#3 * x^0.7 - 2 * x + 1
#a,b,c,a
#3,2,1,0.7
#=
ax = bx-c
c = (b-a)x
=#


struct Prop2DAcoTTIDenQ_DEO2_FDTD
    p::Ptr{Cvoid}
end

function Prop2DAcoTTIDenQ_DEO2_FDTD(;
        nz=0,
        nx=0,
        nsponge=60,
        nbz=512,
        nbx=8,
        dz=5.0,
        dx=5.0,
        dt=0.001,
        nthreads=Sys.CPU_THREADS,
        freesurface=true,
        freqQ=5.0,
        qMin=0.1,
        qInterior=100.0)
    nz,nx,nsponge,nbz,nbx,nthreads = map(x->round(Int,x), (nz,nx,nsponge,nbz,nbx,nthreads))
    dz,dx,dt = map(x->Float32(x), (dz,dx,dt))

    @assert nx > 0
    @assert nz > 0
    @assert nsponge > 0
    @assert nthreads > 0
    @assert nbx > 0
    @assert nbz > 0

    fs = freesurface ? 1 : 0

    p = ccall((:Prop2DAcoTTIDenQ_DEO2_FDTD_alloc, libprop2DAcoTTIDenQ_DEO2_FDTD),
        Ptr{Cvoid},
        (Clong, Clong,    Clong, Clong, Clong,   Cfloat, Cfloat, Cfloat, Clong, Clong),
         fs,    nthreads, nx,    nz,    nsponge, dx,     dz,     dt,     nbx,   nbz)

    ccall((:Prop2DAcoTTIDenQ_DEO2_FDTD_SetupDtOmegaInvQ, libprop2DAcoTTIDenQ_DEO2_FDTD),
        Cvoid,
        (Ptr{Cvoid}, Cfloat, Cfloat, Cfloat),
         p,          freqQ,  qMin,   qInterior)

    Prop2DAcoTTIDenQ_DEO2_FDTD(p)
end

free(prop::Prop2DAcoTTIDenQ_DEO2_FDTD) = ccall((:Prop2DAcoTTIDenQ_DEO2_FDTD_free, libprop2DAcoTTIDenQ_DEO2_FDTD), Cvoid, (Ptr{Cvoid},), prop.p)

function size(prop::Prop2DAcoTTIDenQ_DEO2_FDTD)
    nz = ccall((:Prop2DAcoTTIDenQ_DEO2_FDTD_getNz, libprop2DAcoTTIDenQ_DEO2_FDTD), (Clong), (Ptr{Cvoid},), prop.p)
    nx = ccall((:Prop2DAcoTTIDenQ_DEO2_FDTD_getNx, libprop2DAcoTTIDenQ_DEO2_FDTD), (Clong), (Ptr{Cvoid},), prop.p)
    (nz,nx)
end

for _f in (:V, :Eps, :Eta, :B, :F, :PSpace, :MSpace, :PCur, :POld, :MCur, :MOld, :SinTheta, :CosTheta, :DtOmegaInvQ)
    symf = "Prop2DAcoTTIDenQ_DEO2_FDTD_get" * string(_f)
    @eval $(_f)(prop::Prop2DAcoTTIDenQ_DEO2_FDTD) = unsafe_wrap(Array, ccall(($symf, libprop2DAcoTTIDenQ_DEO2_FDTD), Ptr{Float32}, (Ptr{Cvoid},), prop.p), size(prop), own=false)
end

propagateforward!(prop::Prop2DAcoTTIDenQ_DEO2_FDTD) = ccall((:Prop2DAcoTTIDenQ_DEO2_FDTD_TimeStep, libprop2DAcoTTIDenQ_DEO2_FDTD), Cvoid, (Ptr{Cvoid},), prop.p)
propagateadjoint!(prop::Prop2DAcoTTIDenQ_DEO2_FDTD) = propagateforward!(prop) # self-adjoint

scale_spatial_derivatives!(prop::Prop2DAcoTTIDenQ_DEO2_FDTD) =
    ccall((:Prop2DAcoTTIDenQ_DEO2_FDTD_ScaleSpatialDerivatives, libprop2DAcoTTIDenQ_DEO2_FDTD), Cvoid, (Ptr{Cvoid},), prop.p)

abstract type Prop2DAcoTTIDenQ_DEO2_FDTD_Model end
struct Prop2DAcoTTIDenQ_DEO2_FDTD_Model_VEA <: Prop2DAcoTTIDenQ_DEO2_FDTD_Model end
struct Prop2DAcoTTIDenQ_DEO2_FDTD_Model_V <: Prop2DAcoTTIDenQ_DEO2_FDTD_Model end

# v,ϵ and η in model-space
function forwardBornInjection!(prop::Prop2DAcoTTIDenQ_DEO2_FDTD, modeltype::Prop2DAcoTTIDenQ_DEO2_FDTD_Model_VEA, dmodel, wavefields)
    ccall((:Prop2DAcoTTIDenQ_DEO2_FDTD_ForwardBornInjection_VEA, libprop2DAcoTTIDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat}, Ptr{Cfloat}, Ptr{Cfloat},        Ptr{Cfloat},        Ptr{Cfloat},          Ptr{Cfloat}),
         prop.p,     dmodel["v"], dmodel["ϵ"], dmodel["η"], wavefields["pold"], wavefields["mold"], wavefields["pspace"], wavefields["mspace"])
end

function adjointBornAccumulation!(prop::Prop2DAcoTTIDenQ_DEO2_FDTD, modeltype::Prop2DAcoTTIDenQ_DEO2_FDTD_Model_VEA,
        imagingcondition::ImagingConditionStandard, dmodel, wavefields)
    ccall((:Prop2DAcoTTIDenQ_DEO2_FDTD_AdjointBornAccumulation_VEA, libprop2DAcoTTIDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat}, Ptr{Cfloat}, Ptr{Cfloat},        Ptr{Cfloat},        Ptr{Cfloat},          Ptr{Cfloat}),
         prop.p,     dmodel["v"], dmodel["ϵ"], dmodel["η"], wavefields["pold"], wavefields["mold"], wavefields["pspace"], wavefields["mspace"])
end

# v in model-space
function forwardBornInjection!(prop::Prop2DAcoTTIDenQ_DEO2_FDTD, modeltype::Prop2DAcoTTIDenQ_DEO2_FDTD_Model_V, dmodel, wavefields)
    ccall((:Prop2DAcoTTIDenQ_DEO2_FDTD_ForwardBornInjection_V, libprop2DAcoTTIDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat},          Ptr{Cfloat}),
         prop.p,     dmodel["v"], wavefields["pspace"], wavefields["mspace"])
end

function adjointBornAccumulation!(prop::Prop2DAcoTTIDenQ_DEO2_FDTD, modeltype::Prop2DAcoTTIDenQ_DEO2_FDTD_Model_V,
        imagingcondition::ImagingConditionStandard, dmodel, wavefields)
    ccall((:Prop2DAcoTTIDenQ_DEO2_FDTD_AdjointBornAccumulation_V, libprop2DAcoTTIDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat},          Ptr{Cfloat}),
         prop.p,     dmodel["v"], wavefields["pspace"], wavefields["mspace"])
end

function adjointBornAccumulation!(prop::Prop2DAcoTTIDenQ_DEO2_FDTD, modeltype::Prop2DAcoTTIDenQ_DEO2_FDTD_Model_V,
        imagingcondition::ImagingConditionWaveFieldSeparationFWI, dmodel, wavefields)
    ccall((:Prop2DAcoTTIDenQ_DEO2_FDTD_AdjointBornAccumulation_wavefieldsep_V, libprop2DAcoTTIDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat},          Ptr{Cfloat},          Clong),
         prop.p,     dmodel["v"], wavefields["pspace"], wavefields["mspace"], 1)
end

function adjointBornAccumulation!(prop::Prop2DAcoTTIDenQ_DEO2_FDTD, modeltype::Prop2DAcoTTIDenQ_DEO2_FDTD_Model_V,
        imagingcondition::ImagingConditionWaveFieldSeparationRTM, dmodel, wavefields)
    ccall((:Prop2DAcoTTIDenQ_DEO2_FDTD_AdjointBornAccumulation_wavefieldsep_V, libprop2DAcoTTIDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat},          Ptr{Cfloat},          Clong),
         prop.p,     dmodel["v"], wavefields["pspace"], wavefields["mspace"], 0)
end

function show(io::IO, prop::Prop2DAcoTTIDenQ_DEO2_FDTD)
    nx = ccall((:Prop2DAcoTTIDenQ_DEO2_FDTD_getNx, libprop2DAcoTTIDenQ_DEO2_FDTD), (Clong), (Ptr{Cvoid},), prop.p)
    nz = ccall((:Prop2DAcoTTIDenQ_DEO2_FDTD_getNz, libprop2DAcoTTIDenQ_DEO2_FDTD), (Clong), (Ptr{Cvoid},), prop.p)
    write(io, "Prop2DAcoTTIDenQ_DEO2_FDTD -- nx,nz; $nx,$nz")
end


module GridapGmsh

using Libdl
using Gridap
using Gridap.CellValuesGallery
using StaticArrays
using UnstructuredGrids.Kernels

export GmshDiscreteModel
export GmshCLagrangianFESpace
export GmshDLagrangianFESpace

using Gridap.DiscreteModels: DiscreteModelFromData

deps_jl = joinpath(@__DIR__, "..", "deps", "deps.jl")
if !isfile(deps_jl)
  s = """
  Package GridapGmsh not installed properly.
  Run Pkg.build(\"GridapGmsh\"), restart Julia and try again
  """
  error(s)
end

include(deps_jl)

include(gmsh_jl)

# Hack taken from MPI.jl
function __init__()
    @static if Sys.isunix()
        Libdl.dlopen(gmsh.lib, Libdl.RTLD_LAZY | Libdl.RTLD_GLOBAL)
    end
end

include("GmshDiscreteModels.jl")

include("LagrangianFESpaces.jl")

end # module


module MocosSimLauncher

using ArgParse
using Base.Threads
using DataFrames
using FileIO
using JLD2
using JSON
using ProgressMeter
using Random

import MocosSim
import MocosSim: ContactKind, NoContact, contactkind, time

const OptTimePoint = Union{Missing, MocosSim.TimePoint}
optreal2float32(optreal::Union{Missing,T} where T<:Real) = ismissing(optreal) ? NaN32 : Float32(optreal)

include("cmd_parsing.jl")
include("callback.jl")
include("load_params.jl")
include("outputs.jl")

export launch

function launch()
  @info "Stated" nthreads()
  if nthreads() == 1
    @warn "using single thread, set more threads by passing --threads agrument to julia or setting JULIA_NUM_THREADS environment variable"
  end

  cmd_args = parse_commandline()
  @info "Parsed args" cmd_args
  json = JSON.parsefile(cmd_args["JSON"])

  max_num_infected = json["stop_simulation_threshold"] |> Int
  num_trajectories = json["num_trajectories"] |> Int
  num_initial_infected = json["initial_conditions"]["cardinalities"]["infectious"] |> Int
  params_seed = get(json, "params_seed", 0)

  @info "loading population and setting up parameters" params_seed
  rng = MersenneTwister(params_seed)
  params = read_params(json, rng)
  num_individuals =  MocosSim.numindividuals(params)

  states = [MocosSim.SimState(num_individuals) for _ in 1:nthreads()]
  callbacks = [DetectionCallback(num_individuals, max_num_infected) for _ in 1:nthreads()]

  outputs = make_outputs(cmd_args, num_trajectories)
  foreach(o->beforetrajectories(o, params), outputs)

  @info "starting simulation" num_trajectories
  writelock = ReentrantLock()
  progress = ProgressMeter.Progress(num_trajectories)
  GC.gc()
  @threads for trajectory_id in 1:num_trajectories
    state = states[threadid()]
    MocosSim.reset!(state, trajectory_id)
    MocosSim.initialfeed!(state, num_initial_infected)

    callback = callbacks[threadid()]
    reset!(callback)
    try
      MocosSim.simulate!(state, params, callback)
      foreach(o->pushtrajectory!(o, trajectory_id, writelock, state, params, callback), outputs)
    catch err
      println(stderr, "Failed on thread ", threadid(), " iteration ", trajectory_id, " failed: ", err)
      foreach(x -> println(stderr, x), stacktrace(catch_backtrace()))
    end

    ProgressMeter.next!(progress) # is thread-safe
  end
  foreach(o->aftertrajectories(o, params), outputs)
end

end


# *********************************************************************************
# REopt, Copyright (c) 2019-2020, Alliance for Sustainable Energy, LLC.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
#
# Redistributions of source code must retain the above copyright notice, this list
# of conditions and the following disclaimer.
#
# Redistributions in binary form must reproduce the above copyright notice, this
# list of conditions and the following disclaimer in the documentation and/or other
# materials provided with the distribution.
#
# Neither the name of the copyright holder nor the names of its contributors may be
# used to endorse or promote products derived from this software without specific
# prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
# IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
# OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
# OF THE POSSIBILITY OF SUCH DAMAGE.
# *********************************************************************************
"""
    Storage

Currently only electric storage is modeled but thermal storage will be added soon.

TODO update this doc string and perhaps the way we define arbitrary storage.

```julia
Base.@kwdef struct ElecStorage <: AbstractStorage
    min_kw::Float64 = 0.0
    max_kw::Float64 = 1.0e4
    min_kwh::Float64 = 0.0
    max_kwh::Float64 = 1.0e6
    internal_efficiency_pct::Float64 = 0.975
    inverter_efficiency_pct::Float64 = 0.96
    rectifier_efficiency_pct::Float64 = 0.96
    soc_min_pct::Float64 = 0.2
    soc_init_pct::Float64 = 0.5
    can_grid_charge::Bool = true
    installed_cost_per_kw::Float64 = 840.0
    installed_cost_per_kwh::Float64 = 420.0
    replace_cost_per_kw::Float64 = 410.0
    replace_cost_per_kwh::Float64 = 200.0
    inverter_replacement_year::Int = 10
    battery_replacement_year::Int = 10
    macrs_option_years::Int = 7
    macrs_bonus_pct::Float64 = 1.0
    macrs_itc_reduction::Float64 = 0.5
    total_itc_pct::Float64 = 0.0
    total_rebate_per_kw::Float64 = 0.0
    total_rebate_per_kwh::Float64 = 0.0
end
```
"""
Base.@kwdef struct ElecStorage <: AbstractStorage
    min_kw::Float64 = 0.0
    max_kw::Float64 = 1.0e4
    min_kwh::Float64 = 0.0
    max_kwh::Float64 = 1.0e6
    internal_efficiency_pct::Float64 = 0.975
    inverter_efficiency_pct::Float64 = 0.96
    rectifier_efficiency_pct::Float64 = 0.96
    soc_min_pct::Float64 = 0.2
    soc_init_pct::Float64 = 0.5
    can_grid_charge::Bool = true
    installed_cost_per_kw::Float64 = 840.0
    installed_cost_per_kwh::Float64 = 420.0
    replace_cost_per_kw::Float64 = 410.0
    replace_cost_per_kwh::Float64 = 200.0
    inverter_replacement_year::Int = 10
    battery_replacement_year::Int = 10
    macrs_option_years::Int = 7
    macrs_bonus_pct::Float64 = 1.0
    macrs_itc_reduction::Float64 = 0.5
    total_itc_pct::Float64 = 0.0
    total_rebate_per_kw::Float64 = 0.0
    total_rebate_per_kwh::Float64 = 0.0
end

struct Storage <: AbstractStorage
    types::Array{Symbol,1}
    raw_inputs::Dict{Symbol, AbstractStorage}
    min_kw::Dict{Symbol, Float64}
    max_kw::Dict{Symbol, Float64}
    min_kwh::Dict{Symbol, Float64}
    max_kwh::Dict{Symbol, Float64}
    charge_efficiency::Dict{Symbol, Float64}
    discharge_efficiency::Dict{Symbol, Float64}
    soc_min_pct::Dict{Symbol, Float64}
    soc_init_pct::Dict{Symbol, Float64}
    installed_cost_per_kw::Dict{Symbol, Float64}
    installed_cost_per_kwh::Dict{Symbol, Float64}
    can_grid_charge::Array{Symbol,1}
    grid_charge_efficiency::Dict{Symbol, Float64}
end

"""
    # function Storage(d::Dict{Symbol,Dict}, f::Financial)

Construct Storage struct from Dict with keys for each storage type (eg. :elec) and values with
    input dicts for each storage type. Note that the REopt inputs are indexed on the storage type.
"""
function Storage(d::Dict, f::Financial)  # nested dict
    types = Symbol[]
    can_grid_charge = Symbol[]
    raw_vals = Dict(zip(fieldnames(Storage), [Any[] for _ in range(1, stop=fieldcount(Storage))]))
    # raw_vals has keys = fieldnames(Storage) and values = arrays of values for each type in types

    for (storage_type, input_dict) in d

        push!(types, storage_type)
        struct_name = string(titlecase(string(storage_type)) * "Storage")  # eg. ElecStorage
        storage_instance = eval(Meta.parse(struct_name * "(;$input_dict...)"))

        if storage_instance.can_grid_charge
            push!(can_grid_charge, storage_type)
        end
        fill_storage_vals!(raw_vals, storage_instance, f)
    end

    storage_args = Dict(
        :types => types,
        :can_grid_charge => can_grid_charge,
    )
    d2 = Dict()  # Julia won't let me use storage_args: "unable to check bounds for indices of type Symbol"
    for k in keys(raw_vals)
        d2[k] = Dict(zip(types, raw_vals[k]))
    end

    grid_charge_efficiency = Dict(:elec => convert(Float64, d2[:charge_efficiency][:elec]))

    return Storage(
        storage_args[:types],
        d2[:raw_inputs],
        d2[:min_kw],
        d2[:max_kw],
        d2[:min_kwh],
        d2[:max_kwh],
        d2[:charge_efficiency],
        d2[:discharge_efficiency],
        d2[:soc_min_pct],
        d2[:soc_init_pct],
        d2[:installed_cost_per_kw],
        d2[:installed_cost_per_kwh],
        storage_args[:can_grid_charge],
        grid_charge_efficiency
    )
    # TODO expand for smart thermostat
end

"""
    fill_storage_vals!(d::Dict, s::AbstractStorage, f::Financial)

Fill `d`'s values using `s` and `f`. The dict `d` must have keys equivalent to the fieldnames(Storage)
"""
function fill_storage_vals!(d::Dict, s::AbstractStorage, f::Financial)
    push!(d[:raw_inputs], s)
    push!(d[:min_kw], s.min_kw)
    push!(d[:max_kw], s.max_kw)
    push!(d[:min_kwh], s.min_kwh)
    push!(d[:max_kwh], s.max_kwh)
    push!(d[:soc_min_pct], s.soc_min_pct)
    push!(d[:soc_init_pct], s.soc_init_pct)

    push!(d[:charge_efficiency], s.rectifier_efficiency_pct * s.internal_efficiency_pct^0.5)
    push!(d[:discharge_efficiency], s.inverter_efficiency_pct * s.internal_efficiency_pct^0.5)

    push!(d[:installed_cost_per_kw], effective_cost(;
        itc_basis=s.installed_cost_per_kw,
        replacement_cost=s.replace_cost_per_kw,
        replacement_year=s.inverter_replacement_year,
        discount_rate=f.owner_discount_pct,
        tax_rate=f.owner_tax_pct,
        itc=s.total_itc_pct,
        macrs_schedule = s.macrs_option_years == 7 ? f.macrs_seven_year : f.macrs_five_year,
        macrs_bonus_pct=s.macrs_bonus_pct,
        macrs_itc_reduction = s.macrs_itc_reduction,
        rebate_per_kw = s.total_rebate_per_kw
    ))
    push!(d[:installed_cost_per_kwh], effective_cost(;
        itc_basis=s.installed_cost_per_kwh,
        replacement_cost=s.replace_cost_per_kwh,
        replacement_year=s.inverter_replacement_year,
        discount_rate=f.owner_discount_pct,
        tax_rate=f.owner_tax_pct,
        itc=s.total_itc_pct,
        macrs_schedule=s.macrs_option_years == 7 ? f.macrs_seven_year : f.macrs_five_year,
        macrs_bonus_pct=s.macrs_bonus_pct,
        macrs_itc_reduction = s.macrs_itc_reduction
    ))
    d[:installed_cost_per_kwh][end] -= s.total_rebate_per_kwh
end


module GPLinearODEMaker

using LinearAlgebra
using SpecialFunctions
using Random
using Distributed
using SharedArrays
using SymEngine
using IterativeSolvers
using Statistics

import Base.ndims

include("general_functions.jl")
include("GLO_functions.jl")
include("gp_functions.jl")
include("prior_functions.jl")
include("kernel_base_functions.jl")
include("kernel_creation_functions.jl")
include("diagnostic_functions.jl")

end # module


using Revise
using Flipping
using Guilia
using MixedModels
using Distributions
##
Directory = "/home/beatriz/mainen.flipping.5ht@gmail.com/Flipping/run_task_2/Test"
Directory = "/home/beatriz/mainen.flipping.5ht@gmail.com/Flipping/run_task_2/Dev_raw_data"
pokes, streaks, DataIndex = create_exp_dataframes(Directory)
###
Directory = "/home/beatriz/mainen.flipping.5ht@gmail.com/Flipping/run_task_2/Dev_raw_data/Results/"
file = joinpath(Directory,"streaksResults.csv")
streaks = Guilia.carica(file)
exclude_mice = ["CN21", "CN22","CN23","CN24","CN41","CN42","CN43"]# CNO GROUP WITHOUT CNO THE FIRST DAy
streaks = @filter streaks !in(:MouseID, exclude_mice)
Ncs = ["NC"* string(i) for i in collect(1:14)]
Ccs = ["CC"* string(i) for i in [3,4,5,9,10,14,15,16]]
Rcs = ["RC"* string(i) for i in [1,2,3,4]]
Nbs = ["NB"* string(i) for i in [21,22,23,24,25,41,42,43,44,45]]
Bns = ["BN"* string(i) for i in [21,22,23,24,41,42,43,44]]
Its = ["IT"* string(i) for i in [21,22,23,24,41,42,43,44]]
age_exp = vcat(Nbs,Bns,Its)
youngs = ["BN21", "BN22", "BN41", "BN42", "IT21", "IT22","IT41", "IT42", "NB21", "NB22","NB23","NB41","NB42"]
dev_below15_day1 = ["BN21","BN23","BN41","BN42","BN44","IT23","NB25","NB43","NB45"]
proj_below15_day1 = ["CC9","NC10","NC12","NC14","RC3","RC8"]
below_thrs = vcat(dev_below15_day1,proj_below15_day1)
cnos_animals = vcat(Ncs,Ccs,Rcs)
wt =  append!(["NC"* string(i) for i in [1,2,3,4,9,10,11,12]],age_exp)
"BN21" in youngs
streaks = @apply streaks begin
    @transform {Age = :MouseID in youngs ? "Young" : "Adult"}
    @transform {Exp_type = :MouseID in age_exp ? "Development" : "Projections"}
    @transform {Treatment = :MouseID in cnos ? "CNO" : "SAL"}
    @transform {Gen = :MouseID in wt ? "WT" : "HET"}
    @transform {Low_performance = :MouseID in below_thrs}
    @transform {Combo = :Gen*"_"*:Treatment}
end

union(column(streaks,:MouseID))
union(column(streaks,:Age))
filetosave = joinpath(Directory,"Full_info_streaks.csv")
CSV.write(filetosave,streaks)

cno_only = @apply streaks begin
    @filter !in(:MouseID,age_exp)
    @filter :Exp_Day == 1
end

check = @filter cno_only :Treatment == "CNO"
filetosave = joinpath(Directory,"All_dreadds_streaks.csv")
CSV.write(filetosave,streaks)
#######
data = launch()
t = data[][:data][]
df = DataFrame(t)

categorical!(df,:Treatment)
categorical!(df,:Gen)
##
function Likelyhood_Ratio_test(simple,full)
    degrees = dof(full) - dof(simple)
    ccdf(Distributions.Chisq(degrees), deviance(simple) - deviance(full))
end

function AIC_test(candidate, simpler)
    exp((aic(candidate) - aic(simpler))/2)
end
function AICc(model)
    aic(model) + ((2*(dof(model)^2) + 2*dof(model))/(nobs(model) - dof(model) - 1))
end
function AICc_test(candidate, simpler)
    exp((AICc(candidate) - AICc(simpler))/2)
end
##
No_interaction = fit(LinearMixedModel, @formula( AfterLast ~ 1 + Treatment + Gen + (1|MouseID)), df)
Interaction = fit(LinearMixedModel, @formula( AfterLast ~ 1 + Treatment + Gen + Treatment * Gen + (1|MouseID)), df)

AIC_test(protocol_effect,zero_effects)
AICc_test(protocol_effect,zero_effects)
Likelyhood_Ratio_test(zero_effects,protocol_effect)