const Solution = Dict{String, Union{Int, Bool, Set{Int}}}

#TODO add documentation for BeforeRestart
mutable struct Statistics
    infeasibleStatusPerVariable             ::Dict{String, Int}
    numberOfNodes                           ::Int
    numberOfSolutions                       ::Int
    numberOfInfeasibleSolutions             ::Int
    numberOfSolutionsBeforeRestart          ::Int
    numberOfInfeasibleSolutionsBeforeRestart::Int
    numberOfNodesBeforeRestart              ::Int
    AccumulatedRewardBeforeReset            ::Float32
    solutions                               ::Vector{Union{Nothing,Solution}}
    nodevisitedpersolution                  ::Vector{Int}
    objectives                              ::Union{Nothing, Vector{Union{Nothing,Int}}}
    lastPruning                             ::Union{Nothing, Int}
    lastVar                                 ::Union{Nothing, AbstractIntVar} #last var on which we branched
end

mutable struct Limit
    numberOfNodes       ::Union{Int, Nothing}
    numberOfSolutions   ::Union{Int, Nothing}  #the limit can be triggered by set of non-unique solutions
    searchingTime       ::Union{Int, Nothing}
end

"""
    CPModel(trailer::Trailer)
    CPModel()

The structure storing all the informations needed to solve a specific problem and it also stores the solutions.
The CPModel is the center of the solver and evolve during the solving.
The `AbstractStateRepresentation` used by the RL Agent is created from the CPModel.

The CPModel is always created empty and is filled eather by hand by the user (or automatically thanks to written files)
or filled by an `AbstractModelGenerator`.

"""
mutable struct CPModel
    variables               ::Dict{String, AbstractVar}
    branchable              ::Dict{String, Bool}
    constraints             ::Array{Constraint}
    trailer                 ::Trailer
    objective               ::Union{Nothing, AbstractIntVar}
    objectiveBound          ::Union{Nothing, Int}
    statistics              ::Statistics
    limit                   ::Limit
    knownObjective          ::Union{Nothing,Int64}
    adhocInfo               ::Any
    CPModel(trailer) = new(Dict{String, AbstractVar}(), Dict{String, Bool}(), Constraint[], trailer, nothing, nothing, Statistics(Dict{String, Int}(), 0,0, 0, 0, 0, 0, 0, Solution[],Int[], nothing, nothing, nothing), Limit(nothing, nothing, nothing), nothing)
end

CPModel() = CPModel(Trailer())

"""
    addVariable!(model::CPModel, x::AbstractVar; branchable=true)

Add a variable to the model, throwing an error if `x`'s id is already in the model.
The `branchable` argument allows you to tell if we will be able to branch on that variable.
"""
function addVariable!(model::CPModel, x::AbstractVar; branchable=true)
    # Ensure the id is unique
    @assert !haskey(model.variables, x.id) "The id of the variable must be unique"

    @assert !branchable || typeof(x) <: Union{AbstractIntVar, AbstractBoolVar} "You can only branch on Boolean and Integer variables"

    model.statistics.infeasibleStatusPerVariable[id(x)]=0
    model.branchable[x.id] = branchable
    model.variables[x.id] = x
end

"""
    addObjective!(model::CPModel, objective::AbstractVar)

Add an Objective variable to the model. This variable is the variable that needs to be minimized suring the solving.
"""
function addObjective!(model::CPModel, objective::AbstractVar)
    model.objective = objective
    model.statistics.objectives = Int[]  #initialisation of the Array that will contain the score of every solution
end

function addKnownObjective!(model::CPModel, knownObective::Int64)
    model.knownObjective = knownObective
end

function addConstraint!(model::CPModel,constraint::Constraint)
    push!(model.constraints,constraint)
    for var in variablesArray(constraint)
        if haskey(model.branchable, id(var))
            @assert haskey(model.statistics.infeasibleStatusPerVariable, id(var)) "You forget to add the variable $(id(var)) to the model"
            model.statistics.infeasibleStatusPerVariable[id(var)]+=1
        end
    end
end

"""
    function is_branchable(model::CPModel, x::AbstractVar)

Tell if the variable was set as branchable or not.
"""
is_branchable(model::CPModel, x::AbstractVar) = haskey(model.branchable, x.id) && model.branchable[x.id]

"""
    function branchable_variables(model::CPModel)

Return a dict of all branchable variables.
"""
function branchable_variables(model::CPModel)
    to_return = Dict{String, AbstractVar}()
    for x in values(model.variables)
        if is_branchable(model, x)
            to_return[x.id] = x
        end
    end
    to_return
end

"""
    solutionFound(model::CPModel)

Return a boolean, checking whether a solution was found, i.e. every variable is bound.
"""
function solutionFound(model::CPModel)
    for (k, x) in model.variables
        if !isbound(x)
            return false
        end
    end
    return true
end

"""
    triggerFoundSolution!(model::CPModel)

Add the current solution to `model`, and set new constraints for the objective if needed.
"""
function triggerFoundSolution!(model::CPModel)
    @assert solutionFound(model)

    model.statistics.numberOfSolutions += 1
    model.statistics.numberOfSolutionsBeforeRestart += 1

    # Adding the solution
    solution = Solution()
    for (k, x) in model.variables
        solution[k] = assignedValue(x)
    end
    if !(solution in model.statistics.solutions)   #probably not efficient but necessary
        push!(model.statistics.solutions, solution)
        push!(model.statistics.nodevisitedpersolution,model.statistics.numberOfNodes)
        if !isnothing(model.objective)
            @assert !isnothing(model.statistics.objectives)   "did you used SeaPearl.addObjective! to declare your objective function ? "
            push!(model.statistics.objectives, assignedValue(model.objective))
            tightenObjective!(model)
        end
    end
end
"""
    triggerInfeasible!(constraint::Constraint, model::CPModel)

this function increments by one the statistic infeasibleStatusPerVariable for each variable involved in the constraint. infeasibleStatusPerVariable
keeps in track for each variable the number of times the variable was involved in a constraint that led to an infeasible state during a fixpoint. This statistic
is used by the failure-based variable selection heuristic.
"""
function triggerInfeasible!(constraint::Constraint, model::CPModel)
    for var in variablesArray(constraint)
        if haskey(model.branchable, id(var))
            model.statistics.infeasibleStatusPerVariable[id(var)]+=1
        end
    end
    push!(model.statistics.solutions, nothing)
    push!(model.statistics.nodevisitedpersolution,model.statistics.numberOfNodes)

    if !isnothing(model.objective)
        @assert !isnothing(model.statistics.objectives)   "did you used SeaPearl.addObjective! to declare your objective function ? "
        push!(model.statistics.objectives, nothing)
    end
end

"""
    tightenObjective!(model::CPModel)

Set a new constraint to minimize the objective variable.
"""
function tightenObjective!(model::CPModel)
    model.objectiveBound = assignedValue(model.objective)-1
end

"""
    belowLimits(model::CPModel)

Check if `model`' statistics are still under the limits.
"""
belowLimits(model::CPModel) = belowNodeLimit(model) && belowSolutionLimit(model) && belowTimeLimit(model)
belowNodeLimit(model::CPModel) = isnothing(model.limit.numberOfNodes) || model.statistics.numberOfNodes < model.limit.numberOfNodes
belowSolutionLimit(model::CPModel) = isnothing(model.limit.numberOfSolutions) || model.statistics.numberOfSolutions < model.limit.numberOfSolutions
belowTimeLimit(model::CPModel) = isnothing(model.limit.searchingTime) || peektimer() < model.limit.searchingTime
"""
    Base.isempty(model::CPModel)::Bool

Return a boolean describing if the model is empty or not.
"""
function Base.isempty(model::CPModel)::Bool
    (
        isempty(model.variables)
        && isempty(model.constraints)
        && isempty(model.trailer.prior)
        && isempty(model.trailer.current)
        && isnothing(model.objective)
        && isnothing(model.objectiveBound)
        && isempty(model.statistics.solutions)
        && isempty(model.statistics.nodevisitedpersolution)
        && isempty(model.statistics.infeasibleStatusPerVariable)
        && isnothing(model.statistics.objectives)
        && isnothing(model.statistics.lastPruning)
        && isnothing(model.statistics.lastVar)
        && model.statistics.numberOfNodes == 0
        && model.statistics.numberOfSolutions == 0
        && model.statistics.numberOfInfeasibleSolutions == 0
        && model.statistics.numberOfInfeasibleSolutionsBeforeRestart == 0
        && model.statistics.numberOfSolutionsBeforeRestart == 0
        && model.statistics.numberOfNodesBeforeRestart == 0
        && model.statistics.AccumulatedRewardBeforeReset == 0
        && isnothing(model.limit.numberOfNodes)
        && isnothing(model.limit.numberOfSolutions)
        && isnothing(model.limit.searchingTime)
        && isnothing(model.knownObjective)
    )
end

"""
    Base.empty!(model::CPModel)

Empty the CPModel.
"""
function Base.empty!(model::CPModel)
    empty!(model.variables)
    empty!(model.constraints)
    empty!(model.trailer.prior)
    empty!(model.trailer.current)
    model.objective = nothing
    model.objectiveBound = nothing
    empty!(model.statistics.solutions)
    empty!(model.statistics.nodevisitedpersolution)
    empty!(model.statistics.infeasibleStatusPerVariable)
    model.statistics.objectives = nothing
    model.statistics.lastPruning = nothing
    model.statistics.lastVar = nothing
    model.statistics.numberOfNodes = 0
    model.statistics.numberOfSolutions = 0
    model.statistics.numberOfInfeasibleSolutions = 0
    model.statistics.numberOfInfeasibleSolutionsBeforeRestart = 0
    model.statistics.numberOfSolutionsBeforeRestart = 0
    model.statistics.numberOfNodesBeforeRestart = 0
    model.statistics.AccumulatedRewardBeforeReset = 0
    model.limit.numberOfNodes = nothing
    model.limit.numberOfSolutions = nothing
    model.limit.searchingTime = nothing
    model.knownObjective = nothing

    model
end

"""
    reset_model!(model::CPModel)

Reset a given CPModel instance. Make it possible to reuse the same instance instead of having to
delete the old one and create another one. This is used in `launch_experiment!` in order to be able
to use the same CPModel instance to compare different given heuristics.
"""
function reset_model!(model::CPModel)
    restoreInitialState!(model.trailer)
    model.objectiveBound = nothing
    empty!(model.statistics.solutions)
    empty!(model.statistics.nodevisitedpersolution)
    for (key, value) in model.statistics.infeasibleStatusPerVariable
        model.statistics.infeasibleStatusPerVariable[key]=length(getOnDomainChange(model.variables[key]))  #the degree is reset to the initial value : the number of constraints the variable is involved in.
    end
    if !isnothing(model.objective)
        @assert !isnothing(model.statistics.objectives)   "did you used SeaPearl.addObjective! to declare your objective function ?"
        empty!(model.statistics.objectives)
    end
    model.statistics.lastPruning = nothing
    model.statistics.lastVar = nothing
    model.statistics.numberOfNodes = 0
    model.statistics.numberOfSolutions = 0
    model.statistics.numberOfInfeasibleSolutions = 0
    model.statistics.numberOfInfeasibleSolutionsBeforeRestart = 0
    model.statistics.numberOfSolutionsBeforeRestart = 0
    model.statistics.numberOfNodesBeforeRestart = 0
    model.statistics.AccumulatedRewardBeforeReset = 0

end

"""
restart_search!(model::CPModel)

Usefull when dealing with restart based search : ILDS or RBS. Reset to zero usefull statistics on the search that can be used to define
the restart criteria.
"""
function restart_search!(model::CPModel)
    restoreInitialState!(model.trailer)
    model.statistics.lastPruning = nothing
    model.statistics.lastVar = nothing
    model.statistics.numberOfInfeasibleSolutionsBeforeRestart = 0
    model.statistics.numberOfSolutionsBeforeRestart = 0
    model.statistics.numberOfNodesBeforeRestart = 0
end

"""
    domains_cartesian_product(model::CPModel)

Return the cartesian product of the model variables: |D1|x|D2|x ... x|Dn|
Helps providing insights about what is happening during a search.
"""
function domains_cartesian_product(model::CPModel)
    cart_pdt = 1
    for (id, x) in model.variables
        if typeof(x) <: Union{AbstractIntVar, AbstractBoolVar}
            cart_pdt *= length(x.domain)
        end
    end
    return cart_pdt
end

"""
    nb_boundvariables(model::CPModel)

Return the number of variables that have already been assigned to a value.
Helps providing insights about what is happening during a search.
"""
function nb_boundvariables(model::CPModel)
    nb = 0
    for (id, x) in model.variables
        nb += isbound(x) * 1
    end
    return nb
end


import SeisIO: fill_id!, checkbuf!, checkbuf_strict!, checkbuf_8!, fillx_i4!, fillx_i8!, fillx_i16_le!, fillx_i16_be!, fillx_i24_be!, fillx_i32_le!, fillx_i32_be!, fillx_u32_be!, fillx_u32_le!

  # fill_id!(id::Array{UInt8,1},
  # checkbuf!(buf::Array{UInt8,1},
  # checkbuf!(buf::AbstractArray,
  # checkbuf_strict!(buf::AbstractArray,
  # checkbuf_8!(buf::Array{UInt8,1},
printstyled("  low-level read utils\n", color=:light_green)

nx = 4
buf = Array{UInt8, 1}(undef, 4nx)
xl = Array{Float32, 1}(undef, nx)
xb = similar(xl)
yl = rand(UInt32, 4)
yb = bswap.(yl)
buf .= reinterpret(UInt8, yl)
fillx_u32_le!(xl, buf, nx, 0)
fillx_u32_be!(xb, buf, nx, 0)
@test Float32.(yl) == xl
@test Float32.(yb) == xb

# Everything else is tested in readers
buf = getfield(BUF, :buf)
x = getfield(BUF, :int32_buf)
nx = 256
os = 3
checkbuf_8!(buf, 65536)
checkbuf_8!(buf, 4*(os + nx))
checkbuf!(x, os + nx)

# test fillx_i16_le!
y = rand(Int16, nx)
copyto!(buf, 1, reinterpret(UInt8, y), 1, 2*nx)
fillx_i16_le!(x, buf, nx, os)
@test x[1+os:nx+os] == y

# test fillx_i32_le!
y = rand(Int32, nx)
copyto!(buf, 1, reinterpret(UInt8, y), 1, 4*nx)
fillx_i32_le!(x, buf, nx, os)
@test x[1+os:nx+os] == y

fillx_i32_be!(x, buf, nx, os)
@test x[1+os:nx+os] == bswap.(y)

resize!(x, 100)
nothing


# Base.null_safe_op is not defined on v0.5
if VERSION < v"0.6.0-dev.0"
    # TODO: Find parsimonious way to support null_safe_op optimizations on 0.5
    null_safe_op(f, xs...) = false
else
    using Base: null_safe_op
end

lift(f, x) = f(x)

abstract MaybeNullable
immutable YesNullable <: MaybeNullable end
immutable NoNullable <: MaybeNullable end
Base.@pure any_nullable(x,y) = ifelse((x<:Nullable) | (y<:Nullable), YesNullable(), NoNullable())
Base.@pure any_nullable(xs) = ifelse(mapreduce(x->x<:Nullable, |, xs.parameters), YesNullable(), NoNullable())

# # back-up in case @pure gives us grief
# any_nullable(x::Nullable, xs...) = true
# any_nullable(x, xs...) = any_nullables(xs...)
# any_nullable(x) = false
# any_nullable() = false

@inline lift(f, x, y) = lift(f, x, y, any_nullable(typeof(x), typeof(y)))
@inline lift(f, xs...) = lift(f, xs, any_nullable(typeof(xs)))
@inline lift(f, x, y, ::NoNullable) = f(x, y)
@inline lift(f, xs, ::NoNullable) = f(xs...)

##############################################################################
##
## Standard lifting semantics
##
## For a function call f(xs...), return null if any x in xs is null;
## otherwise, return f applied to values of xs.
##
##############################################################################

@inline function lift(f, x::Nullable)
    if null_safe_op(f, typeof(x))
        return @compat Nullable(f(x.value), !isnull(x))
    else
        U = Core.Inference.return_type(f, Tuple{eltype(typeof(x))})
        if isnull(x)
            return Nullable{U}()
        else
            return Nullable(f(unsafe_get(x)))
        end
    end
end

@inline function lift(f, x1, x2, ::YesNullable)
    if null_safe_op(f, typeof(x1), typeof(x2))
        return @compat Nullable(
            f(x1.value, x2.value), !(isnull(x1) | isnull(x2)))
    else
        U = Core.Inference.return_type(
            f, Tuple{eltype(typeof(x1)), eltype(typeof(x2))})
        if isnull(x1) | isnull(x2)
            return Nullable{U}()
        else
            return Nullable(f(unsafe_get(x1), unsafe_get(x2)))
        end
    end
end

@inline function lift(f, xs, ::YesNullable)
    if null_safe_op(f, map(typeof, xs)...)
        return @compat Nullable(
            f(map(unsafe_get, xs)...), !(mapreduce(isnull, |, xs)))
    else
        U = Core.Inference.return_type(
            f, Tuple{map(x->eltype(typeof(x)), xs)...})
        if hasnulls(xs)
            return Nullable{U}()
        else
            return Nullable(f(map(unsafe_get, xs)...))
        end
    end
end

##############################################################################
##
## Non-standard lifting semantics
##
##############################################################################

# three-valued logic implementation
@inline function lift(::typeof(&), x, y, ::YesNullable)::Nullable{Bool}
    return ifelse( isnull(x),
        ifelse( isnull(y),
            Nullable{Bool}(),                       # x, y null
            ifelse( unsafe_get(y),
                Nullable{Bool}(),                   # x null, y == true
                Nullable(false)                     # x null, y == false
            )
        ),
        ifelse( isnull(y),
            ifelse( unsafe_get(x),
                Nullable{Bool}(),                   # x == true, y null
                Nullable(false)                     # x == false, y null
            ),
            Nullable(unsafe_get(x) & unsafe_get(y)) # x, y not null
        )
    )
end

# three-valued logic implementation
@inline function lift(::typeof(|), x, y, ::YesNullable)::Nullable{Bool}
    return ifelse( isnull(x),
        ifelse( isnull(y),
            Nullable{Bool}(),                       # x, y null
            ifelse( unsafe_get(y),
                Nullable(true),                     # x null, y == true
                Nullable{Bool}()                    # x null, y == false
            )
        ),
        ifelse( isnull(y),
            ifelse( unsafe_get(x),
                Nullable(true),                     # x == true, y null
                Nullable{Bool}()                    # x == false, y null
            ),
            Nullable(unsafe_get(x) | unsafe_get(y)) # x, y not null
        )
    )
end

# TODO: Decide on semantics for isequal and uncomment the following
# @inline function lift(::typeof(isequal), x, y)
#     return ifelse( isnull(x),
#         ifelse( isnull(y),
#             true,                                   # x, y null
#             false                                   # x null, y not null
#         ),
#         ifelse( isnull(y),
#             false,                                   # x not null, y null
#             isequal(unsafe_get(x), unsafe_get(y))    # x, y not null
#         )
#     )
# end

@inline function lift(::typeof(isless), x, y, ::YesNullable)::Bool
    if null_safe_op(isless, typeof(x), typeof(y))
        return ifelse( isnull(x),
            false,                                      # x null
            ifelse( isnull(y),
                true,                                   # x not null, y null
                isless(unsafe_get(x), unsafe_get(y))    # x, y not null
            )
        )
    else
        return  isnull(x) ? false :
                isnull(y) ? true  : isless(unsafe_get(x), unsafe_get(y))
    end
end

@inline lift(::typeof(isnull), x::Nullable) = isnull(x)
@inline lift(::typeof(isnull), x, ::MaybeNullable) = isnull(x)
@inline lift(::typeof(get), x::Nullable) = get(x)
@inline lift(::typeof(get), x::Nullable, y) = get(x, y)


import Base: show

show(io::IO, ::MIME"text/html", p::MExpr{:Graphics}) =
  print(io, ExportString(p, "SVG"))

show(io::IO, ::MIME"image/png", p::MExpr{:Graphics3D}) =
  print(io, @math ExportString($p, "PNG"))


using DataFrames, FITSIO, ProgressMeter, Dates

import DataFrames.DataFrame, Base.write

export fits2df

DataFrame(hdu::FITSIO.ImageHDU) = error("Can not convert an ImageHDU extension to a DataFrame")
function DataFrame(hdu::FITSIO.HDU)
    fld = join.(split.(FITSIO.colnames(hdu), '.'), '_')
    out = DataFrame()
		@showprogress 0.5 for i in 1:length(fld)
        tmp = read(hdu, (FITSIO.colnames(hdu))[i])
        if ndims(tmp) == 1
            out[!, Symbol(fld[i])] = tmp
        else
            app = DataFrame(Symbol(fld[i])=>Vector{typeof(tmp[1,1])}[])
            for j in 1:(size(tmp))[2]
                push!(app, [tmp[:,j]])
            end
            out = hcat(out, app)
        end
    end
    return out
end

function fits2df(filename::String, hdu=2)
    f = FITS(filename)
    out = DataFrame(f[hdu])
    close(f)
    return out
end

function write(f::FITSIO.FITS, dfr::DataFrame)
    df = deepcopy(dfr)
    data = Array{Any}(undef, 0)
    for name in names(df)
        tt = eltype(df[:, name])
        if isa(tt, Union)
            tt = nonmissingtype(tt)
            if tt == String
                @warn "Using empty string as missing value in field $name"
                df[ismissing.(df[:, name]), name] .= ""
            elseif tt <: AbstractFloat
                @warn "Using NaN as missing value in field $name"
                df[ismissing.(df[:, name]), name] .= NaN
            elseif tt <: Integer
                @warn "Using zero as missing value in field $name"
                df[ismissing.(df[:, name]), name] .= 0
            end
            disallowmissing!(df, name)
        end
        tt = eltype(df[:, name])
        @assert !isa(tt, Union)

        if tt == Symbol
            push!(data, string.(df[:, name]))
        elseif tt == String
            push!(data, string.(df[:, name]))
        elseif tt <: Dates.AbstractTime
            push!(data, string.(df[:, name]))
        else
            push!(data, df[:, name])
        end
    end
    write(f, string.(names(df)), data)
end


"""
    FoldsView(data, train_indices, val_indices, [obsdim])

Description
============

Create a vector-like representation of `data` where each
individual element is partition of `data` in the form of a tuple
of two data subsets (a training- and a validation subset).

The purpose of `FoldsView` is to apply a precomputed sequence of
subset assignment indices to some data container in a convenient
manner. By itself, `FoldsView` is agnostic to any particular
repartitioning strategy (such as k-folds). Instead, the
assignments, `train_indices` and `val_indices`, need to be
precomputed by such a strategy and then passed to `FoldsView`
with a concrete `data` container. The resulting object can then
be queried for its individual folds using `getindex`, or simply
iterated over.

Arguments
==========

- **`data`** : The object describing the dataset. Can be of any
    type as long as it implements [`getobs`](@ref) and
    [`nobs`](@ref) (see Details for more information).

- **`train_indices`** : Vector of integer vectors containing the
    sequence of training assignments. This means that each
    element is a vector of indices that describe each *training*
    subset. The length of this vector must match `val_indices`.

- **`val_indices`** : Vector of integer vectors containing the
    sequence of validation assignments. This means that each
    element is a vector of indices that describe each
    *validation* subset. The length of this vector must match
    `train_indices`.

- **`obsdim`** : Optional. If it makes sense for the type of
    `data`, `obsdim` can be used to specify which dimension of
    `data` denotes the observations. It can be specified in a
    type-stable manner as a positional argument (see
    `?LearnBase.ObsDim`), or more conveniently as a smart keyword
    argument.

Details
========

For `FoldsView` to work on some data structure, the desired type
`MyType` must implement the following interface:

- `LearnBase.getobs(data::MyType, idx, [obsdim::ObsDimension])` :
    Should return the observation(s) indexed by `idx`.
    In what form is up to the user.
    Note that `idx` can be of type `Int` or `AbstractVector`.

- `LearnBase.nobs(data::MyType, [obsdim::ObsDimension])` :
    Should return the total number of observations in `data`

Author(s)
==========

- Christof Stocker (Github: https://github.com/Evizero)
- Tom Breloff (Github: https://github.com/tbreloff)

Examples
=========

    # Load iris data for demonstration purposes
    X, y = MLDataUtils.load_iris()

    # Compute train- and validation-partition indices using kfolds
    train_idx, val_idx = kfolds(nobs(X), 10)

    # These two vectors contain the indices vector for each partitioning
    @assert typeof(train_idx) <: Vector{Vector{Int64}}
    @assert typeof(val_idx)   <: Vector{UnitRange{Int64}}
    @assert length(train_idx) == length(val_idx) == 10

    # Using the repartitioning as an oterator
    for (train_X, val_X) in FoldsView(X, train_idx, val_idx)
        @assert size(train_X) == (4, 135)
        @assert size(val_X) == (4, 15)
    end

    # Calling kfolds with the dataset will create
    # the FoldsView for you automatically.
    # Thus this code is equivalent to above
    for (train_X, val_X) in kfolds(X, 10)
        @assert size(train_X) == (4, 135)
        @assert size(val_X) == (4, 15)
    end

    # leavout is a shortcut for setting k = nobs(X)
    for (train_X, val_X) in leaveout(X)
        @assert size(val_X) == (4, 1)
    end

see also
=========

[`kfolds`](@ref), [`leaveout`](@ref), [`splitobs`](@ref),
[`DataSubset`](@ref)
"""
struct FoldsView{T,D,O,A1<:AbstractArray,A2<:AbstractArray} <: DataView{T,D}
    data::D
    train_indices::A1
    val_indices::A2
    obsdim::O

    function FoldsView{T,D,O,A1,A2}(
            data::D, train_indices::A1, val_indices::A2, obsdim::O) where {T,D,O,A1<:AbstractArray,A2<:AbstractArray}
        n = nobs(data, obsdim)
        (eltype(train_indices) <: AbstractArray{Int}) || throw(ArgumentError("The parameter \"train_indices\" must be an array of integer arrays"))
        (eltype(val_indices)  <: AbstractArray{Int}) || throw(ArgumentError("The parameter \"val_indices\" must be an array of integer arrays"))
        2 <= length(train_indices) <= n || throw(ArgumentError("The amount of train- and validation-indices must be within 2:$n respectively"))
        length(train_indices) == length(val_indices) || throw(DimensionMismatch("The amount of train- and validation-indices must match"))
        new{T,D,O,A1,A2}(data, train_indices, val_indices, obsdim)
    end
end

function FoldsView(data::D, train_indices::A1, val_indices::A2, obsdim::O) where {D,O,A1<:AbstractArray,A2<:AbstractArray}
    n = nobs(data, obsdim)
    # TODO: Move this back into the inner constructor after the
    #       "T = typeof(...)" line below is removed
    (1 <= minimum(minimum.(train_indices)) && maximum(maximum.(train_indices)) <= n) || throw(DimensionMismatch("All training indices must be within 1:$n"))
    (1 <= minimum(minimum.(val_indices))  && maximum(maximum.(val_indices))  <= n) || throw(DimensionMismatch("All validation indices must be within 1:$n"))
    # FIXME: In 0.6 it should be possible to compute just the return
    #        type without executing the function
    T = typeof((datasubset(data, train_indices[1], obsdim), datasubset(data, val_indices[1], obsdim)))
    FoldsView{T,D,O,A1,A2}(data, train_indices, val_indices, obsdim)
end

FoldsView(data, train_indices::AbstractArray, val_indices::AbstractArray; obsdim = default_obsdim(data)) =
    FoldsView(data, train_indices, val_indices, convert(LearnBase.ObsDimension,obsdim))

function FoldsView(data::T, train_indices::AbstractArray, val_indices::AbstractArray, obsdim) where T<:DataView
    @assert obsdim == data.obsdim
    @warn string("Trying to nest a ", T.name, " into an FoldsView, which is not supported. Returning FoldsView(parent(_)) instead")
    FoldsView(parent(data), train_indices, val_indices, obsdim)
end

# compare if both FoldsViews describe the same folds of the same data
# we don't care how the indices are stored, just that they match
# in their order and values
function Base.:(==)(fv1::FoldsView,fv2::FoldsView)
    fv1.data == fv2.data &&
        all(all(i1==i2 for (i1,i2) in zip(I1,I2)) for (I1,I2) in zip(fv1.train_indices,fv2.train_indices)) &&
        all(all(i1==i2 for (i1,i2) in zip(I1,I2)) for (I1,I2) in zip(fv1.val_indices,fv2.val_indices)) &&
        fv1.obsdim == fv2.obsdim
end

nobs(iter::FoldsView) = nobs(iter.data, iter.obsdim)
Base.parent(iter::FoldsView) = iter.data
Base.length(iter::FoldsView) = length(iter.train_indices)

function Base.getindex(iter::FoldsView, i::Int)
    (datasubset(iter.data, iter.train_indices[i], iter.obsdim),
     datasubset(iter.data, iter.val_indices[i], iter.obsdim))
end

function Base.getindex(iter::FoldsView, i::AbstractVector)
    FoldsView(iter.data, iter.train_indices[i], iter.val_indices[i], iter.obsdim)
end

# compatibility with nested functions
default_obsdim(iter::FoldsView) = iter.obsdim

function Base.summary(A::FoldsView)
    string(length(A), "-fold ", typeof(A).name,
           " of ", nobs(A), " observations")
end

_datastr(data) = summary(data)
_datastr(data::Tuple) = string("(", join(map(_datastr, data), ", "), ")")
function Base.show(io::IO, ::MIME"text/plain", A::FoldsView)
    println(io, summary(A), ":")
    println(io, "  data: ", _datastr(A.data))
    # TODO: consider the case where fold sizes differ
    ntrain = length(first(A.train_indices))
    println(io, "  training: ", ntrain, " observations/fold")
    nval = length(first(A.val_indices))
    println(io, "  validation: ", nval, " observations/fold")
    print(io, "  obsdim: ", obsdim_string(A.obsdim))
end

"""
    kfolds(n::Integer, [k = 5]) -> Tuple

Compute the train/validation assignments for `k` repartitions of
`n` observations, and return them in the form of two vectors. The
first vector contains the index-vectors for the training subsets,
and the second vector the index-vectors for the validation subsets
respectively. A general rule of thumb is to use either `k = 5` or
`k = 10`. The following code snippet generates the indices
assignments for `k = 5`

```julia
julia> train_idx, val_idx = kfolds(10, 5);
```

Each observation is assigned to the validation subset once (and
only once). Thus, a union over all validation index-vectors
reproduces the full range `1:n`. Note that there is no random
assignment of observations to subsets, which means that adjacent
observations are likely to be part of the same validation subset.

```julia
julia> train_idx
5-element Array{Array{Int64,1},1}:
 [3,4,5,6,7,8,9,10]
 [1,2,5,6,7,8,9,10]
 [1,2,3,4,7,8,9,10]
 [1,2,3,4,5,6,9,10]
 [1,2,3,4,5,6,7,8]

julia> val_idx
5-element Array{UnitRange{Int64},1}:
 1:2
 3:4
 5:6
 7:8
 9:10
```
"""
function kfolds(n::Integer, k::Integer = 5)
    2 <= k <= n || throw(ArgumentError("n must be positive and k must to be within 2:$(max(2,n))"))
    # Compute the size of each fold. This is important because
    # in general the number of total observations might not be
    # divideable by k. In such cases it is custom that the remaining
    # observations are divided among the folds. Thus some folds
    # have one more observation than others.
    sizes = fill(floor(Int, n/k), k)
    for i = 1:(n % k)
        sizes[i] = sizes[i] + 1
    end
    # Compute start offset for each fold
    offsets = cumsum(sizes) .- sizes .+ 1
    # Compute the validation indices using the offsets and sizes
    val_indices = map((o,s)->(o:o+s-1), offsets, sizes)
    # The train indices are then the indicies not in validation
    train_indices = map(idx->setdiff(1:n,idx), val_indices)
    # We return a tuple of arrays
    train_indices, val_indices
end

"""
    kfolds(data, [k = 5], [obsdim]) -> FoldsView

Repartition a `data` container `k` times using a `k` folds
strategy and return the sequence of folds as a lazy
[`FoldsView`](@ref). The resulting `FoldsView` can then be
indexed into or iterated over. Either way, only data subsets are
created. That means that no actual data is copied until
[`getobs`](@ref) is invoked.

Conceptually, a k-folds repartitioning strategy divides the given
`data` into `k` roughly equal-sized parts. Each part will serve
as validation set once, while the remaining parts are used for
training. This results in `k` different partitions of `data`.

In the case that the size of the dataset is not dividable by the
specified `k`, the remaining observations will be evenly
distributed among the parts.

```julia
for (x_train, x_val) in kfolds(X, k = 10)
    # code called 10 times
    # nobs(x_val) may differ up to ±1 over iterations
end
```

Multiple variables are supported (e.g. for labeled data)

```julia
for ((x_train, y_train), val) in kfolds((X, Y), k = 10)
    # ...
end
```

By default the folds are created using static splits. Use
[`shuffleobs`](@ref) to randomly assign observations to the
folds.

```julia
for (x_train, x_val) in kfolds(shuffleobs(X), k = 10)
    # ...
end
```

see [`FoldsView`](@ref) for more info, or [`leaveout`](@ref) for
a related function.
"""
function kfolds(data, k::Integer, obsdim)
    n = nobs(data, obsdim)
    train_indices, val_indices = kfolds(n, k)
    FoldsView(data, train_indices, val_indices, obsdim)
end

kfolds(data, k::Integer; obsdim = default_obsdim(data)) =
    kfolds(data, k, convert(LearnBase.ObsDimension,obsdim))

kfolds(data; k = 5, obsdim = default_obsdim(data)) =
    kfolds(data, k, convert(LearnBase.ObsDimension,obsdim))

kfolds(data, obsdim::Union{Tuple,ObsDimension}) =
    kfolds(data, 5, obsdim)

"""
    leaveout(n::Integer, [size = 1]) -> Tuple

Compute the train/validation assignments for `k ≈ n/size`
repartitions of `n` observations, and return them in the form of
two vectors. The first vector contains the index-vectors for the
training subsets, and the second vector the index-vectors for the
validation subsets respectively. Each validation subset will have
either `size` or `size+1` observations assigned to it. The
following code snippet generates the index-vectors for `size = 2`.

```julia
julia> train_idx, val_idx = leaveout(10, 2);
```

Each observation is assigned to the validation subset once (and
only once). Thus, a union over all validation index-vectors
reproduces the full range `1:n`. Note that there is no random
assignment of observations to subsets, which means that adjacent
observations are likely to be part of the same validation subset.

```julia
julia> train_idx
5-element Array{Array{Int64,1},1}:
 [3,4,5,6,7,8,9,10]
 [1,2,5,6,7,8,9,10]
 [1,2,3,4,7,8,9,10]
 [1,2,3,4,5,6,9,10]
 [1,2,3,4,5,6,7,8]

julia> val_idx
5-element Array{UnitRange{Int64},1}:
 1:2
 3:4
 5:6
 7:8
 9:10
```
"""
function leaveout(n::Integer, size::Integer = 1)
    1 <= size <= floor(n/2) || throw(ArgumentError("size must to be within 1:$(floor(Int,n/2))"))
    k = floor(Int, n / size)
    kfolds(n, k)
end

"""
    leaveout(data, [size = 1], [obsdim]) -> FoldsView

Repartition a `data` container using a k-fold strategy, where `k`
is chosen in such a way, that each validation subset of the
resulting folds contains roughly `size` observations. Defaults to
`size = 1`, which is also known as "leave-one-out" partitioning.

The resulting sequence of folds is returned as a lazy
[`FoldsView`](@ref), which can be index into or iterated over.
Either way, only data subsets are created. That means no actual
data is copied until [`getobs`](@ref) is invoked.

```julia
for (train, val) in leaveout(X, size = 2)
    # if nobs(X) is dividable by 2,
    # then nobs(val) will be 2 for each iteraton,
    # otherwise it may be 3 for the first few iterations.
end
```

see [`FoldsView`](@ref) for more info, or [`kfolds`](@ref) for a
related function.
"""
function leaveout(data, size, obsdim)
    n = nobs(data, obsdim)
    train_indices, val_indices = leaveout(n, size)
    FoldsView(data, train_indices, val_indices, obsdim)
end

leaveout(data, size::Integer; obsdim = default_obsdim(data)) =
    leaveout(data, size, convert(LearnBase.ObsDimension,obsdim))

leaveout(data; size = 1, obsdim = default_obsdim(data)) =
    leaveout(data, size, convert(LearnBase.ObsDimension,obsdim))

leaveout(data, obsdim::Union{Tuple,ObsDimension}) =
    leaveout(data, 1, obsdim)


abstract type AbstractExistence end

struct DynamicExistence <: AbstractExistence
    existence::Vector{Bool}
    sample_freq::typeof(1.0Hz)
    time::typeof(1.0s)
end

function DynamicExistence(existence, sample_freq; time = 0.0s)
    DynamicExistence(existence, sample_freq, time)
end

"""
$(SIGNATURES)

Simulates the static existence of either a satellite signal or a interference signal. Type:
Boolean (true if signal exists).
"""
function propagate(existence::Bool, Δt, rng)
    existence
end

function get_existence(existence::Bool)
    existence
end

"""
$(SIGNATURES)

Simulates the dynamic existence of either a satellite signal or a interference signal. Type:
Boolean (true if signal exists).
If time index exceeds data length, last available value is returned
"""
function propagate(existence::DynamicExistence, Δt, rng)
    next_time = existence.time + Δt
    DynamicExistence(existence.existence, existence.sample_freq, next_time)
end

function get_existence(existence::DynamicExistence)
    get_sampled_value(existence.time, existence.sample_freq, existence.existence)
end


"""
    AbstractSampler

Abstract supertype for sampler objects.
"""
abstract type AbstractSampler{S <: AbstractScenario} end
"""
    sample(sampler::AbstractSampler{S})

Sample a scenario of type `S` using `sampler`.
"""
function sample(sampler::AbstractSampler)
    return sampler()
end
function Base.show(io::IO, sampler::AbstractSampler{S}) where S <: AbstractScenario
    print(io, "$(S.name.name) sampler")
    return io
end
function Base.show(io::IO, sampler::AbstractSampler{S}) where S <: Scenario
    print(io, "Scenario sampler")
    return io
end
"""
    sample(sampler::AbstractSampler{S}, π::AbstractSampler)

Sample a scenario of type `S` using `sampler` and set the probability of the sampled scenario to `π`.
"""
function sample(sampler::AbstractSampler, π::AbstractFloat)
    scenario = sampler()
    set_probability!(scenario, π)
    return scenario
end
"""
    Sampler

General purpose sampler object that samples Scenario.

See also: [`Scenario`](@ref), [`@sampler`](@ref)
"""
struct Sampler <: AbstractSampler{Scenario}
    sampler::Function
end

function (sampler::Sampler)()
    return sampler.sampler()
end


using Test
@isdefined(TestHelpers) || include("../../TestHelpers.jl")

TestHelpers.@timed_testset "generated_weights/" begin
    include("generated_weights/mfc+mfc+softmax.jl")
    include("generated_weights/conv+fc+softmax.jl")
end


using MPI, Test
include(joinpath("..", "..", "testhelpers.jl"))

@testset "DGmethods" begin
    runmpi(joinpath(@__DIR__, "vars_test.jl"))
    runmpi(joinpath(@__DIR__, "integral_test.jl"))
    runmpi(joinpath(@__DIR__, "integral_test_sphere.jl"))
    runmpi(joinpath(@__DIR__, "Euler/isentropicvortex.jl"))
    runmpi(joinpath(@__DIR__, "Euler/isentropicvortex_imex.jl"))
    runmpi(joinpath(@__DIR__, "Euler/isentropicvortex_multirate.jl"))
    runmpi(joinpath(
        @__DIR__,
        "advection_diffusion/pseudo1D_advection_diffusion.jl",
    ))
    runmpi(joinpath(
        @__DIR__,
        "advection_diffusion/advection_diffusion_model_1dimex_bgmres.jl",
    ))
    runmpi(joinpath(@__DIR__, "compressible_Navier_Stokes/ref_state.jl"))
    runmpi(joinpath(@__DIR__, "horizontal_integral_test.jl"))
    runmpi(joinpath(@__DIR__, "courant.jl"), ntasks = 2)
end


#--------------------------------------------------------------------
# Spacetime Discretization methods in Julia
# Soham 01-2018
#--------------------------------------------------------------------

fbase = fdistribute(x->sin(x), y->sin(y), (x,y)->0, 2, 2, 2)
dbase = distribute(x->sin(x), y->sin(y), (x,y)->0, 2, 2, 2)

@test fbase[[1,1]].value ≈ dbase[[1,1]].value
@test fbase[[1,2]].value ≈ dbase[[1,2]].value
@test fbase[[2,1]].value ≈ dbase[[2,1]].value
@test fbase[[2,2]].value ≈ dbase[[2,2]].value


"""
    chamber(dim=100.0)

Construct a high vacuum chamber to act as the outer-most Region.
"""
chamber(dim = 100.0) = #
    Region(
        RectangularShape([-dim, -dim, -dim], 2.0 * [dim, dim, dim]),
        parse(Material, "H", density = 5.0e-11),
        nothing,
        "Chamber",
        0,
    )

"""
    gun(::Type{T}, energy::Float64, width::Float64=1.0e-7, initial::Position=Position(0.0,0.0,-10.0), direction=Position(0.0,0.0,1.0)::T where {T <: Particle}

A helper to construct the initial `Particle` in a randomized Gaussian distribution.
"""
function gun(
    ::Type{T},
    energy::Float64,
    width::Float64 = 1.0e-7,
    initial::Position = Position(0.0, 0.0, -10.0),
    direction = Position(0.0, 0.0, 1.0),
)::T where {T<:Particle}
    r, th = sqrt(-2.0 * log(rand())) * width, 2.0π * rand()
    st = initial .+ Position(r * cos(th), r * sin(th), 0.0)
    return T(st .- direction, st, energy)
end

"""
    bulk(mat::Material)

Construct a bulk homogeneous sample at the origin.
"""
function bulk(mat::Material)
    c = chamber()
    Region(RectangularShape([-0.1, -0.1, 0.0], [0.2, 0.2, 0.2]), mat, c, "Bulk homogeneous")
    return c
end

"""
    particle(mat::Material, radius::Float64, substrate::Union{Nothing,Material} = nothing)

Construct a spherical particle at the origin with an optional bulk substrate.
"""
function particle(
    mat::Material,
    radius::Float64,
    substrate::Union{Nothing,Material} = nothing,
)
    c = chamber()
    Region(SphericalShape([0.0, 0.0, -radius], radius), mat, c, "Particle")
    if !isnothing(substrate)
        Region(
            RectangularShape([-0.1, -0.1, 0.0], [0.2, 0.2, 0.2]),
            substrate,
            c,
            "Substrate",
        )
    end
    return c
end

"""
    coated_particle(mat::Material, radius::Float64, coating::Material, thickness::Float64, substrate::Union{Nothing,Material} = nothing)

Construct a coated particle on an optional substrate.

This model is a good example of how more complex models are constructed.  Use `dump(..)` to display the structure.  You'll notice that
the chamber serves as the root `Region`.  The `coating` and `substrate` are in the chamber (children of the chamber `Region`).  The
`particle` is a child of the `coating` `Region` because the particle is fully enclosed by the coating.  An electron inside of the coating will
enter the particle and leave the coating as soon as it enters the volume representing the particle.  The electron only appears to be
within the child-most region at any point in space.  So a typical trajectory might start in the chamber, enter the coating, traverse
the thickness of the coating and then enter the particle.  It may eventually leave the particle and reenter the coating, exit the coating
and reenter the chamber and finally come to rest in the substrate.
"""
function coated_particle(
    mat::Material,
    radius::Float64,
    coating::Material,
    thickness::Float64,
    substrate::Union{Nothing,Material} = nothing,
)
    c = chamber()
    cr = Region(
        SphericalShape([0.0, 0.0, -(radius + thickness)], radius + thickness),
        coating,
        c,
        "Coating",
    )
    Region(SphericalShape([0.0, 0.0, -(radius + thickness)], radius), mat, cr, "Particle")
    if !isnothing(substrate)
        Region(
            RectangularShape([-0.1, -0.1, 0.0], [0.2, 0.2, 0.2]),
            substrate,
            c,
            "Substrate",
        )
    end
    return c
end

"""
    thin_film(prs::Pair{Material, Float64}...; substrate::Union{Nothing,Material} = nothing)

Construct sample consisting of one or more thin films on an optional bulk substrate.
"""
function thin_film(
    prs::Pair{Material,Float64}...;
    substrate::Union{Nothing,Material} = nothing,
)
    c = chamber()
    t = 0.0
    for (i, (mat, thick)) in enumerate(prs)
        Region(RectangularShape([-0.1, -0.1, t], [0.2, 0.2, thick]), mat, c, "Layer[$i]")
        t += thick
    end
    if !isnothing(substrate)
        Region(
            RectangularShape([-0.1, -0.1, t], [0.2, 0.2, 0.2 - t]),
            substrate,
            c,
            "Substrate",
        )
    end
    return c
end

"""
    allmaterials(reg::Region)

Generates a list of all the unique materials in a sample Region.
"""
function allmaterials(reg::Region)
    res = [reg.material]
    for ch in reg.children
        append!(res, allmaterials(ch))
    end
    return unique(res)
end

"""
    colorize(reg::Region)::Dict{Material, Color}

Generate a `Dict{Material, Color}` for all the `Materials` in the specified `Region`.
Designed for distinctive but not necessarily attractive colors.
"""
function colorize(reg::Region)::Dict{Material,Color}
    mats = allmaterials(reg)
    colors = distinguishable_colors(
        length(mats) + 2,
        [RGB(0.12, 0.14, 0.47), RGB(0.42, 0.63, 0.85), RGB(0.9, 0.9, 0.9)],
        transform = deuteranopic,
    )[3:end]
    return Dict(mat => col for (mat, col) in zip(mats, colors))
end

"""
    sample(ch::Region)

Extracts the sample portion from within the chamber.
"""
sample(ch::Region)::Vector{Region} = ch.children

function Base.dump(reg::Region, indent = 0)
    print("$(" → "^indent)$reg\n")
    for ch in reg.children
        dump(ch, indent + 1)
    end
end


#=
Here I show the exact correspondence
my spiking implementation and code by
Litwin-Kumar / Auguste Schulz

Here I just compare weight matrix,definition of assemblies, input function

=#

# include inbuilt modules
using Distributions
using SparseArrays
using Dates
using LinearAlgebra
using Random
using Plots ; theme(:dark)
using Random ; Random.seed!(0)
using ProgressMeter
using SpikingRNNs; const global S = SpikingRNNs
using SparseArrays

# --------------------- include functions  ------------------------------------------------
# include runs the respective julia code, i.e. defined functions are then in the workspace

#function gennextsequencepretraining(stim::Array{Float64,2}, firstimg, novelimg; Nimg = 4, Ntrain = 20, Nreps = 20, Nseq = 5, Nblocks = 1, stimstart = 1000, lenstim = 1000, lenpause = 3000, strength = 8)
function genstimparadigmpretraining(stim::Array{Float64,2}; Nimg = 4, Nass = 20, Ntrain = 20, stimstart = 1000, lenstim = 1000, lenpause = 3000, strength = 8)
	# generate a stimulus sequence
	# where each assembly is
	# Nimg Number of images per sequence
	# Nreps Number of seqeunce repititions in one block
	# Nseq Number of sequences
	# Nblocks Number of block repititions

	if Ntrain != 0
		Nstim = Nass * Ntrain
		tempstim = zeros(Nstim, 4)

		#tempstim[1,1] = firstimg
		# repeat sequence several times, i.e. assemblie numbers 123412341234...
		#tempstim[:,1] = repeat(firstimg:(firstimg+Nimg-1), outer = Nreps)
		# ensure that images are shuffled but never the same image after the other
		images = 1:Nass
		assemblysequence = shuffle(copy(images))
		for rep = 2:Ntrain
			shuffleimg = shuffle(images)
			#println("begin $(shuffleimg[1])")
			while shuffleimg[1] == assemblysequence[end]
				println("shuffle again")
				shuffleimg = shuffle(images)
				#println("new begin $(shuffleimg[1])")
			end
		assemblysequence = vcat(assemblysequence, shuffleimg)
		end
		tempstim[:,1] .= assemblysequence
		if stim[end,3] == 0
			tempstim[1,2] = stimstart
		else
			tempstim[1,2] = stim[end,3]+lenpause
		end
		tempstim[1,3] = tempstim[1,2] + lenstim
		tempstim[1,4] = strength

		for i = 2:Nstim#Nimg*Nreps #number in image times
			# # lenstim ms stimulation and lenpause ms wait time
			tempstim[i,2] = tempstim[i-1,2] + lenstim + lenpause
			tempstim[i,3] = tempstim[i,2]+lenstim
			tempstim[i,4] = strength
		end
		if stim[end,3] == 0
			stim = tempstim
		else
			stim = vcat(stim, tempstim)
		end
		#tempstim = nothing
	end
	return stim
end # function genstim

function gennextsequenceNonovel(stim::Array{Float64,2}, firstimg, novelimg; Nimg = 4, Nreps = 20, Nseq = 5, Nblocks = 1, stimstart = 1000, lenstim = 1000, lenpause = 3000, strength = 8)
	# generate a stimulus sequence
	# where each assembly is
	# Nimg Number of images per sequence
	# Nreps Number of seqeunce repititions in one block
	# Nseq Number of sequences
	# Nblocks Number of block repititions
	Nstim = Nimg*Nreps
	tempstim = zeros(Nstim, 4)
    #tempstim[1,1] = firstimg
	tempstim[:,1] = repeat(firstimg:(firstimg+Nimg-1), outer = Nreps)

	if stim[end,3] == 0
		tempstim[1,2] = stimstart
	else
    	tempstim[1,2] = stim[end,3]+lenpause
	end
    tempstim[1,3] = tempstim[1,2] + lenstim
    tempstim[1,4] = strength

    for i = 2:Nimg*Nreps #number in image times
        # lenstim ms stimulation and lenpause ms wait time
        tempstim[i,2] = tempstim[i-1,2] + lenstim + lenpause
        tempstim[i,3] = tempstim[i,2]+lenstim
        tempstim[i,4] = strength
    end
	#tempstim[end-Nimg,1] = novelimg
	if stim[end,3] == 0
		stim = tempstim
	else
		stim = vcat(stim, tempstim)
	end
	#tempstim = nothing
    return stim
end # function genstim

function genstimparadigmNonovel(stimulus; Nimg = 4, Nreps = 20, Nseq = 5, Nblocks = 1, stimstart = 1000, lenstim = 1000, lenpause = 3000, strength = 8)

	firstimages = collect(1:Nimg:Nimg*Nseq)
	novelimg = firstimages .+ Nimg*Nseq
	novelrep = collect(Nimg*Nseq:Nimg*Nseq+Nimg).+1

	# for img in firstimages
	# 	stimulus = gennextsequence!(stimulus, img)
	# end
	blockonset = []
	storefirstimages = copy(firstimages)
	for b = 1:Nblocks
		if b == 1
			for img in firstimages
				append!(blockonset, size(stimulus,1) + 1)
				stimulus = gennextsequenceNonovel(stimulus, img, img + Nseq * Nimg, Nimg = Nimg, Nreps = Nreps, Nseq = Nseq, Nblocks = Nblocks, stimstart = stimstart, lenstim = lenstim, lenpause = lenpause, strength = strength)
			end
		else
			shuffleimg = shuffle(firstimages)
			while shuffleimg[1] == storefirstimages[end]
				shuffleimg = shuffle(firstimages)
			end
			storefirstimages = vcat(storefirstimages, shuffleimg)
			for img in shuffleimg # since we shuffle ensure that it is still the correct novel image
				append!(blockonset, size(stimulus,1) + 1)
				stimulus = gennextsequenceNonovel(stimulus, img, img + Nseq * Nimg,Nimg = Nimg, Nreps = Nreps, Nseq = Nseq, Nblocks = Nblocks, stimstart = stimstart, lenstim = lenstim, lenpause = lenpause, strength = strength)
			end
		end
	end
	# ensure mapping from novel img to right value despite shuffling
	for i = 1:length(novelimg)
		stimulus[stimulus[:,1] .== novelimg[i],1] .= novelrep[i]
	end

	# ensure that blockonset 1 when we start with an emptz array
	if blockonset[1] == 2
		blockonset[1] = 1
	end
	return stimulus, convert(Array{Int64,1}, blockonset)

end

function genstimparadigmnovelcont(stimulus; Nimg = 4, Nreps = 20, Nseq = 5, Nblocks = 1, stimstart = 1000, lenstim = 1000, lenpause = 3000, strength = 8)
  """	generate the stimulation paradigm ensure that for each block there is a different novel image and
  that two consecutive blocks are never the same

  generate seq. violation stimulation paradigm
  ABCABCABC....ABN1ABC
  DEFDEFDEF....DEN2DEF
  XYZXYZXYZ...XYN3XYZ
  DEFDEFDEF....DEN4DEF
  ABCABCABC....ABN1ABC
  """

    firstimages = collect(1:Nimg:Nimg*Nseq)
    novelimg = Nimg*Nseq + 1 # start counting from the first assembly after the core assemblies
    novelrep = collect(Nimg*Nseq:Nimg*Nseq+Nimg).+1

    storefirstimages = copy(firstimages)
    blockonset = []

    for b = 1:Nblocks
      #println("block $(b) -------------------")
      if b == 1
        for img in firstimages
          append!(blockonset, size(stimulus,1) + 1)
          stimulus = gennextsequencenovel(stimulus, img, novelimg, Nimg = Nimg, Nreps = Nreps, Nseq = Nseq, Nblocks = Nblocks, stimstart = stimstart, lenstim = lenstim, lenpause = lenpause, strength = strength)
          novelimg += 1
        end
      else
        # ensure that the first images of two consecutive blocks differ
        shuffleimg = shuffle(firstimages)
        #println("begin $(shuffleimg[1])")
        while shuffleimg[1] == storefirstimages[end] && Nseq > 1 # added Nseq larger 1 cause otgerwise always the same
          println("shuffle again")
          shuffleimg = shuffle(firstimages)
          #println("new begin $(shuffleimg[1])")
        end
        storefirstimages = vcat(storefirstimages, shuffleimg)
        #println(storefirstimages)
        for img in shuffleimg  # since we shuffle ensure that it is still the correct novel image
          append!(blockonset, size(stimulus,1) + 1)
          stimulus = gennextsequencenovel(stimulus, img, novelimg ,Nimg = Nimg, Nreps = Nreps, Nseq = Nseq, Nblocks = Nblocks, stimstart = stimstart, lenstim = lenstim, lenpause = lenpause, strength = strength)
          novelimg += 1
        end
      end
    end

    # ensure that blockonset 1 when we start with an emptz array
    if blockonset[1] == 2
      blockonset[1] = 1
    end
    return stimulus, convert(Array{Int64,1}, blockonset)

  end

function gennextsequencenovel(stim::Array{Float64,2}, firstimg, novelimg; Nimg = 4, Nreps = 20, Nseq = 5, Nblocks = 1, stimstart = 1000, lenstim = 1000, lenpause = 3000, strength = 8)
	"""generate a stimulus sequence with novel stim
	Nimg Number of images per sequence
	Nreps Number of seqeunce repititions in one block
	Nseq Number of sequences
	Nblocks Number of block repititions"""
	Nstim = Nimg*Nreps
	tempstim = zeros(Nstim, 4)
    #tempstim[1,1] = firstimg
	# repeat sequence several times, i.e. assemblie numbers 123412341234...
	tempstim[:,1] = repeat(firstimg:(firstimg+Nimg-1), outer = Nreps)
	if stim[end,3] == 0
		tempstim[1,2] = stimstart
	else
    	tempstim[1,2] = stim[end,3]+lenpause
	end
    tempstim[1,3] = tempstim[1,2] + lenstim
    tempstim[1,4] = strength

    for i = 2:Nimg*Nreps #number in image times
        # # lenstim ms stimulation and lenpause ms wait time
        tempstim[i,2] = tempstim[i-1,2] + lenstim + lenpause
        tempstim[i,3] = tempstim[i,2]+lenstim
        tempstim[i,4] = strength
    end
	tempstim[end-Nimg,1] = novelimg
	if stim[end,3] == 0
		stim = tempstim
	else
		stim = vcat(stim, tempstim)
	end
	#tempstim = nothing
    return stim
end # function genstim

function genstimparadigmssa(stimulus; Nimg = 4, Nreps = 20, Nseq = 5, Nblocks = 1, stimstart = 1000, lenstim = 1000, lenpause = 3000, strength = 8)
    """
    generate SSA stimulation paradigm
    A A A A A A A B A A A ... Block 1
    B B B B B 				 Block 2 ...
    Nreps 20 times
    Blocks
    """
    Nimg = 2
    Nreps = 5
    Nseq = 5
    Nblocks = 10
    stimulus, blockonset = genstimparadigmnovelcont(stimulus, Nimg = 2, Nreps = 5, Nseq = 5, Nblocks = 10, stimstart = stimstart, lenstim = 100, lenpause = 300, strength = strength )

    #last_onset = blockonset[1]
    freq = 1
    dev = 2
    for bb =1:length(blockonset)
      #idxfreq = collect(blockonset[bb]:(blockonset[bb]+Nimg*Nreps))
      freq_start = blockonset[bb]
      freq_end = blockonset[bb]+Nimg*Nreps
      dev_end = blockonset[bb]+Nimg*Nreps
      dev_start = blockonset[bb]+Nreps
      if bb == length(blockonset)
        stimulus[freq_start:end,1] .= freq
        #println(idxfreq)
        #idxdev = collect(blockonset[bb]+Nreps:Nreps:(blockonset[bb]+Nimg*Nreps))
        stimulus[dev_start:Nreps:end,1] .= dev
        #printl
      else
        stimulus[freq_start:freq_end,1] .= freq
        #println(idxfreq)
        #idxdev = collect(blockonset[bb]+Nreps:Nreps:(blockonset[bb]+Nimg*Nreps))
        stimulus[dev_start:Nreps:dev_end,1] .= dev
        #println(idxdev)
        # for the nect block switch deviant and frequent
      end
      temp =freq
      freq = dev
      dev = temp
    end

    return stimulus, blockonset
end

function initweights((Ne, Ni, p, Jee0, Jei0, Jie, Jii))
	""" initialise the connectivity weight matrix based on initial
	E-to-E Jee0 initial weight e to e plastic
	I-to-E Jei0 initial weight i to e plastic
	E-to-I Jie0 constant weight e to i not plastic
	I-to-I Jii constant weight i to i not plastic"""
	Ncells = Ne+Ni

	# initialise weight matrix
	# w[i,j] is the weight from pre- i to postsynaptic j neuron
	weights = zeros(Float64,Ncells,Ncells)
	weights[1:Ne,1:Ne] .= Jee0
	weights[1:Ne,(1+Ne):Ncells] .= Jie
	weights[(1+Ne):Ncells,1:Ne] .= Jei0
	weights[(1+Ne):Ncells,(1+Ne):Ncells] .= Jii
	# set diagonal elements to 0
	# for cc = 1:Ncells
	# 	weights[cc,cc] = 0
	# end
	weights[diagind(weights)] .= 0.0
	# ensure that the connection probability is only p
	weights = weights.*(rand(Ncells,Ncells) .< p)
	return weights
end

function weightpars(;Ne = 4000, Ni = 1000, p = 0.2 )
	"""Ne, Ni number of excitatory, inhibitory neurons
	p initial connection probability"""
	Jee0 = 2.86 #initial weight e to e plastic
	Jei0 = 48.7 #initial weight i to e plastic
	Jie = 1.27 #constant weight e to i not plastic
	Jii = 16.2 #constant weight i to i not plastic
	return Ne,Ni,p,Jee0,Jei0,Jie,Jii
end

function initassemblymembers(;Nassemblies = 20, pmember = .05, Nmembersmax = 300, Ne = 4000)
	"""Nassemblies number of assemblies
	pmember probability of belonging to any assembly
	Nmembersmax maximum number of neurons in a population (to set size of matrix)

	set up excitatory assemblies"""
	#seed = 1
	#Random.seed!(seed) # ensure same Assemblies when repeating
	assemblymembers = ones(Int,Nassemblies,Nmembersmax)*(-1)
	for pop = 1:Nassemblies
		members = findall(rand(Ne) .< pmember)
		assemblymembers[pop,1:length(members)] = members
	end
	#println(assemblymembers)
	return assemblymembers
end

function initinhibassemblymembers(;Nassemblies = 20, pmember = .15, Nmembersmax = 200, Ne = 4000, Ni = 1000)
	"""Nassemblies number of assemblies
	pmember probability of belonging to any assembly
	Nmembersmax maximum number of neurons in a population (to set size of matrix)

	set up inhibitory assemblies
	higher connection probability"""
	#seed = 1
	#Random.seed!(seed) # ensure same Assemblies when repeating
	inhibmembers = ones(Int,Nassemblies,Nmembersmax)*(-1)
	for pop = 1:Nassemblies
		members = findall(rand(Ni) .< pmember) .+ Ne
		inhibmembers[pop,1:length(members)] = members
	end
	#println(assemblymembers)
	return inhibmembers
end

#include("../simulation/runsimulation_inhibtuning.jl")
#include("../simulation/sequencefunctions.jl")
#include("../simulation/helperfunctions.jl")
#include("../evaluation/evaluationfunctions.jl")

# --------------------- initialise simulation --------------------------------------------

# Define number of excitatory and inhibitory neurons
const Ne = 4000
const Ni = 1000

Ncells = Ne + Ni

weights_au = initweights(weightpars(Ne = Ne, Ni = Ni))
weights_compare = permutedims(weights_au)
heatmap(weights_au)

pweights=0.2
Jee = 2.86 #initial weight e to e plastic
weights_ee = S.make_sparse_weights(Ne,Ne,pweights,Jee)
Jei = 48.7 #initial weight i to e plastic
weights_ei = S.make_sparse_weights(Ne,Ni,pweights,Jei)
Jie = 1.27 #constant weight e to i not plastic
weights_ie = S.make_sparse_weights(Ni,Ne,pweights,Jie)
Jii = 16.2 #constant weight i to i not plastic
weights_ii = S.make_sparse_weights(Ni,Ni,pweights,Jii)

weights_all = Matrix(vcat(hcat(weights_ee,weights_ei), hcat(weights_ie,weights_ii)))

heatmap(weights_all)
heatmap(weights_compare)

##
# Nimg 3, Nreps 20, Nseq 10, Nblocks 1, lenstim 300, strength 12, Ntrain 5, adjustfactor 1.0, adaptive neurons false, inhibfactor 0.1
ARGS = ["3", "20", "10", "1","300", "12", "5", "10", "0", "10"] # unique sequences
# ARGS = ["3", "20", "5", "10","300", "12", "5", "10", "0", "10"] # repeated sequences

# stimulus parameters
Nimg = parse(Int64, ARGS[1]) # number of stimuli per sequence
Nreps = parse(Int64, ARGS[2]) # number of repetitions per sequence cycle
Nseq = parse(Int64, ARGS[3]) # number of sequences
Nblocks = parse(Int64, ARGS[4]) # number of sequence block repetitions
stimstart = 1_000  # 4000 # start time of stimulation in ms
lenstimpre = parse(Int64, ARGS[5]) # duration of assembly stimulation per image in ms
lenpausepre = 0 # duration of no stimulation between two images in ms
strengthpre = parse(Int64, ARGS[6]) # strength of the stimulation in kHz added to baseline input
lenstim = lenstimpre # duration of assembly stimulation per image in ms
lenpause = lenpausepre # duration of no stimulation between two images in ms
strength = strengthpre # strength of the stimulation in kHz added to baseline input
Ntrain = parse(Int64, ARGS[7]) # number of pretraining iterations
Nass = (Nimg) * Nseq +  Nseq * Nblocks # total number of assemblies Nimg * Nseq + all novelty assemblies
##

p_ase = 0.05
p_asi = 0.15
Nassemblies = 40

assemblymembers = initassemblymembers(Nassemblies = Nass,Ne = Ne)
# inhibitory tuning
inhibassemblies = initinhibassemblymembers(Nassemblies = Nass, Ne = Ne, Ni = Ni)

assemblies_e = map(1:Nassemblies) do _
  sort(findall(rand(Ne) .< p_ase))
end
assemblies_i = map(1:Nassemblies) do _
  sort(findall(rand(Ni) .< p_asi))
end

map(r->count(r.>0),eachrow(assemblymembers))
length.(assemblies_e)

map(r->count(r.>0),eachrow(inhibassemblies))
length.(assemblies_i)

##

stimulus = zeros(1, 4)	# initialisation of the stimulus

# ---------------------- generate the stimulus --------------------------------

# specify aspects of the stimulation paradigm
withnovelty 	= true 	# novelty response
pretrainig 		= true	# include pretraining phase

# ----------- standard setting when the following three false -----------------
shuffledstim 	= false	# shuffle the stimulus randomly
lastshuffled 	= false	# shuffle the last two stimuli
reducednovelty 	= false	# reduce the stimulation strength of the novel stimulus

if withnovelty
	if pretrainig
		# --------- main stim generation function ---------------
		stimulus = genstimparadigmpretraining(stimulus, Nass = Nass, Ntrain = Ntrain, stimstart = stimstart, lenstim = lenstimpre, lenpause = lenpausepre, strength = strengthpre)
		lenpretrain = size(stimulus,1)
	end

	if shuffledstim
		if pretrainig
			stimulus, blockonset = genstimparadigmnovelshuffledpretrain(stimulus, lenpretrain, Nimg = Nimg, Nreps = Nreps, Nseq = Nseq, Nblocks = Nblocks, stimstart = stimstart, lenstim = lenstim, lenpause = lenpause, strength = strength)
			tag = "repeatedsequencesshuffled.h5"

		else
			stimulus, blockonset = genstimparadigmnovelshuffled(stimulus, Nimg = Nimg, Nreps = Nreps, Nseq = Nseq, Nblocks = Nblocks, stimstart = stimstart, lenstim = lenstim, lenpause = lenpause, strength = strength )
			tag = "repeatedsequencesshuffled.h5"
		end
	else # if shuffled stim

		# --------- main stim generation function ---------------
		stimulus, blockonset = genstimparadigmnovelcont(stimulus, Nimg = Nimg, Nreps = Nreps, Nseq = Nseq, Nblocks = Nblocks, stimstart = stimstart, lenstim = lenstim, lenpause = lenpause, strength = strength )
		tag = "repeatedsequences.h5"

	end
else # if with novelty
	# even if no novelty assemblies are shown still train with all novelty assemblies
	if pretrainig
		#println(stimulus)
		stimulus = genstimparadigmpretraining(stimulus, Nass = Nass, Ntrain = Ntrain, stimstart = stimstart, lenstim = lenstimpre, lenpause = lenpausepre, strength = strengthpre)
		#println(stimulus)
		lenpretrain = size(stimulus,1)
		tag = "repeatedsequences_imprinting.h5"

	end
	tag = "repeatedsequencesnoNovelty.h5"
	stimulus, blockonset = genstimparadigmNonovel(stimulus, Nimg = Nimg, Nreps = Nreps, Nseq = Nseq, Nblocks = Nblocks, stimstart = stimstart, lenstim = lenstim, lenpause = lenpause, strength = strength )
end

function inputfun(t::R,i::Integer,low::R,
    stimulus_mat::Matrix{Float64},
    assembly_v::Vector{Vector{Int64}};scale_factor::Float64=1.0) where R<:Real
  # catch the first or last
  (!(stimulus_mat[1,2] <= t <= stimulus_mat[end,3])) && (return low)
  stimstarts = view(stimulus_mat,:,2)
  # find row number
  therow = searchsortedfirst(stimstarts,t) - 1
  # account for stim duration
  (t >= stimulus_mat[therow,3]) && (return low)
  # ok, the assembly is active!
  as_v = assembly_v[Int64(stimulus_mat[therow,1])]
  if _is_in_assembly(i,as_v)
    return low + scale_factor*stimulus_mat[therow,4] # add to stimulus
  else
    return low
  end
  return low # in case the compiler is dumb
end

function inputfun_upper(t::R,i::Integer,low::R,stimmat::Matrix{R},asv::Vector;
   scale_factor::Float64=1.0) where R<:Real
  return max(inputfun(t,i,low,stimmat,asv;scale_factor=scale_factor),
              inputfun(t+2dt,i,low,stimmat,asv;scale_factor=scale_factor))
end

function _is_in_assembly(neu_idx::Integer,as_vec::Vector{<:Integer})
  if as_vec[1] <= neu_idx <= as_vec[end]
    idx = searchsortedfirst(as_vec,neu_idx)
    return as_vec[idx] == neu_idx
  else
    return false
  end
end
const assembly_vectors =
    map(eachrow(hcat(assemblymembers,inhibassemblies))) do r
    sort(filter(>(0),r))
end
as_ordering = let Ntot=Ncells
  S.order_by_pattern_idxs(assembly_vectors,Ntot)
end

const stimulus_sec = copy(stimulus)
stimulus_sec[:,[2,3]] ./= 1E3
stimulus_sec[:,4] .*= 1E3

_ = let low_e = 2E3, low_i = 3E3,
  _scal = 0.1
  ratefun = function (t,i)
    if i<=Ne
      inputfun(t,i,low_e,stimulus_sec,assembly_vectors)
    else
      inputfun(t,i,low_i,stimulus_sec,assembly_vectors,scale_factor=_scal)
    end
  end
  times = range(0,100;length=200)
  nshow = Ncells
  neushow = as_ordering[1:nshow]
  vals = hcat([ ratefun.(t,neushow) for t in times]...)
  # vals_toplot = vals[as_ordering,:][1:nshow,:]
  heatmap(times,1:nshow,vals)
end

# now mine!

##
Ttot = 60.0
lowrate_e =2E3
highrate_e =12E3
lowrate_i =3E3
highrate_i = 12E3
scale_fact_i = 0.1
Δt_as = 0.3
t_as_delay = 1.0
Δt_blank = 0.0

inputfun_e = S.PatternPresentation(Δt_as,Ttot,lowrate_e,lowrate_e+highrate_e,Ne,
    assemblies_e;t_pattern_delay=t_as_delay,Δt_pattern_blank=Δt_blank)
inputfun_i = S.PatternPresentation(Δt_as,Ttot,lowrate_i,lowrate_i+scale_fact_i*highrate_i,
  Ni, assemblies_i;t_pattern_delay=t_as_delay,Δt_pattern_blank=Δt_blank)
# force same sequence
inputfun_i.sequence .= inputfun_e.sequence

##

as_ordering_mine = let
  assemblies_joint=map(ei->sort(vcat(ei[1],Ne.+ei[2])),zip(assemblies_e,assemblies_i))
  S.order_by_pattern_idxs(assemblies_joint,Ne+Ni)
end

##

_ = let
  ratefun = function (t,i)
    if i<=Ne
      inputfun_e(t,i)
    else
      inputfun_i(t,i-Ne)
    end
  end
  times = range(0,100;length=200)
  nshow = Ncells
  neushow = as_ordering_mine[1:nshow]
  vals = hcat([ ratefun.(t,neushow) for t in times]...)
  # vals_toplot = vals[as_ordering,:][1:nshow,:]
  heatmap(times,1:nshow,vals)
end


module Simulations

    using ProgressMeter

    using ..Equations
    using ..Integrators
    using ..Solutions

    export Simulation, run!

    include("simulations/simulation.jl")

end


# definition of some distribution types

## Distribution & Make

abstract type Distribution{T} end

Base.eltype(::Type{<:Distribution{T}}) where {T} = T

struct Make{T,X<:Tuple, XX<:Tuple} <: Distribution{T}
    # XX is X with Type{A} replaced by Nothing (otherwise, X (e.g. Tuple{Type{Int}}) can't have instances)
    x::XX
end

# @inline necessary for one @inferred test on arrays
@inline function Base.getindex(m::Make{T,X}, i::Integer) where {T,X}
    i = Int(i)
    fieldtype(X, i) <: Type ?
        fieldtype(X, i).parameters[1] :
        m.x[i]
end

@inline Base.getindex(m::Make, idxs::AbstractVector{<:Integer}) =
    ntuple(i->m[idxs[i]], length(idxs))
# seems faster than `Tuple(m[i] for i in idxs)`

Base.lastindex(m::Make) = lastindex(m.x)

@generated function Make{T}(X...) where T
    XX = Tuple{(x <: Type ? Nothing : x for x in X)...}
    y = [x <: Type ? nothing : :(X[$i]) for (i, x) in enumerate(X)]
    :(Make{T,Tuple{$X...},$XX}(tuple($(y...))))
end

Make0{T} = Make{T,Tuple{}}
Make1{T} = Make{T,Tuple{X}}         where X
Make2{T} = Make{T,Tuple{X,Y}}       where {X, Y}
Make3{T} = Make{T,Tuple{X,Y,Z}}     where {X, Y, Z}
Make4{T} = Make{T,Tuple{X,Y,Z,U}}   where {X, Y, Z, U}
Make5{T} = Make{T,Tuple{X,Y,Z,U,V}} where {X, Y, Z, U, V}

Make0{T}()              where {T} = Make{T}()
Make1{T}(x)             where {T} = Make{T}(x)
Make2{T}(x, y)          where {T} = Make{T}(x, y)
Make3{T}(x, y, z)       where {T} = Make{T}(x, y, z)
Make4{T}(x, y, z, u)    where {T} = Make{T}(x, y, z, u)
Make5{T}(x, y, z, u, v) where {T} = Make{T}(x, y, z, u, v)

# default maketype & make & Make(...)

# Make(...) is not meant to be specialized, i.e. Make(a, b, c) always create a Make3,
# and is equal to the *default* make(...)
# (it's a fall-back for client code which can help break recursivity)
# TODO: add tests for Make(...)

maketype(::Type{T}, x...) where {T} = T

Make(::Type{T}, x...) where {T} = Make{maketype(T, x...)}(x...)
make(::Type{T}, x...) where {T} = Make{maketype(T, x...)}(x...)

# make(x) is defined in sampling.jl, and is a special case wrapping already valid
# distributions (explicit or implicit)
Make(x1, x2, xs...) = Make{maketype(x1, x2, xs...)}(x1, x2, xs...)
make(x1, x2, xs...) = Make{maketype(x1, x2, xs...)}(x1, x2, xs...)

find_deduced_type(::Type{T}, ::X,     ) where {T,X} = deduce_type(T, gentype(X))
find_deduced_type(::Type{T}, ::Type{X}) where {T,X} = deduce_type(T, X)

find_deduced_type(::Type{T}, ::X,       ::Y)       where {T,X,Y} = deduce_type(T, gentype(X), gentype(Y))
find_deduced_type(::Type{T}, ::Type{X}, ::Y)       where {T,X,Y} = deduce_type(T, X,          gentype(Y))
find_deduced_type(::Type{T}, ::X,       ::Type{Y}) where {T,X,Y} = deduce_type(T, gentype(X), Y)
find_deduced_type(::Type{T}, ::Type{X}, ::Type{Y}) where {T,X,Y} = deduce_type(T, X,          Y)

deduce_type(::Type{T}, ::Type{X}, ::Type{Y}) where {T,X,Y} = _deduce_type(T, Val(isconcretetype(T)), X, Y)
deduce_type(::Type{T}, ::Type{X}) where {T,X} = _deduce_type(T, Val(isconcretetype(T)), X)

_deduce_type(::Type{T}, ::Val{true},  ::Type{X}, ::Type{Y}) where {T,X,Y} = T
_deduce_type(::Type{T}, ::Val{false}, ::Type{X}, ::Type{Y}) where {T,X,Y} = deduce_type(T{X}, Y)

_deduce_type(::Type{T}, ::Val{true},  ::Type{X}) where {T,X} = T
_deduce_type(::Type{T}, ::Val{false}, ::Type{X}) where {T,X} = T{X}

# show(::Make)

function Base.show(io::IO, m::Make{T}) where {T}
    M = typeof(m)
    P = M.parameters[2].parameters
    t = ntuple(length(m.x)) do i
        P[i] isa Type{<:Type} ? P[i].parameters[1] : m.x[i]
    end
    Base.show_type_name(io, M.name)
    print(io, "{", T, "}", t)
end

## Uniform

abstract type Uniform{T} <: Distribution{T} end

struct UniformType{T} <: Uniform{T} end

Uniform(::Type{T}) where {T} = UniformType{T}()

Base.getindex(::UniformType{T}) where {T} = T

struct UniformWrap{T,E} <: Uniform{E}
    val::T
end

Uniform(x::T) where {T} = UniformWrap{T,gentype(T)}(x)

Base.getindex(x::UniformWrap) = x.val

## Normal & Exponential

abstract type Normal{T} <: Distribution{T} end

struct Normal01{T} <: Normal{T} end

struct Normalμσ{T} <: Normal{T}
    μ::T
    σ::T
end

const NormalTypes = Union{AbstractFloat,Complex{<:AbstractFloat}}

Normal(::Type{T}=Float64) where {T<:NormalTypes} = Normal01{T}()
Normal(μ::T, σ::T) where {T<:NormalTypes} = Normalμσ(μ, σ)
Normal(μ::T, σ::T) where {T<:Real} = Normalμσ(AbstractFloat(μ), AbstractFloat(σ))
Normal(μ, σ) = Normal(promote(μ, σ)...)

abstract type Exponential{T} <: Distribution{T} end

struct Exponential1{T} <: Exponential{T} end

struct Exponentialθ{T} <: Exponential{T}
    θ::T
end

Exponential(::Type{T}=Float64) where {T<:AbstractFloat} = Exponential1{T}()
Exponential(θ::T) where {T<:AbstractFloat} = Exponentialθ(θ)
Exponential(θ::Real) = Exponentialθ(AbstractFloat(θ))

## floats

abstract type FloatInterval{T<:AbstractFloat} <: Uniform{T} end

abstract type CloseOpen{ T<:AbstractFloat} <: FloatInterval{T} end
abstract type OpenClose{ T<:AbstractFloat} <: FloatInterval{T} end
abstract type CloseClose{T<:AbstractFloat} <: FloatInterval{T} end
abstract type OpenOpen{  T<:AbstractFloat} <: FloatInterval{T} end

struct CloseOpen12{T<:AbstractFloat} <: CloseOpen{T} end # interval [1,2)

struct CloseOpen01{ T<:AbstractFloat} <: CloseOpen{T}  end # interval [0,1)
struct OpenClose01{ T<:AbstractFloat} <: OpenClose{T}  end # interval (0,1]
struct CloseClose01{T<:AbstractFloat} <: CloseClose{T} end # interval [0,1]
struct OpenOpen01{  T<:AbstractFloat} <: OpenOpen{T}   end # interval (0,1)

struct CloseOpenAB{T<:AbstractFloat} <: CloseOpen{T} # interval [a,b)
    a::T
    b::T

    CloseOpenAB{T}(a::T, b::T) where {T} = (check_interval(a, b); new{T}(a, b))
end

struct OpenCloseAB{T<:AbstractFloat} <: OpenClose{T} # interval (a,b]
    a::T
    b::T

    OpenCloseAB{T}(a::T, b::T) where {T} = (check_interval(a, b); new{T}(a, b))
end

struct CloseCloseAB{T<:AbstractFloat} <: CloseClose{T} # interval [a,b]
    a::T
    b::T

    CloseCloseAB{T}(a::T, b::T) where {T} = (check_interval(a, b); new{T}(a, b))
end

struct OpenOpenAB{T<:AbstractFloat} <: OpenOpen{T} # interval (a,b)
    a::T
    b::T

    OpenOpenAB{T}(a::T, b::T) where {T} = (check_interval(a, b); new{T}(a, b))
end

check_interval(a, b) = a >= b && throw(ArgumentError("invalid interval specification"))

const FloatInterval_64 = FloatInterval{Float64}
const CloseOpen01_64   = CloseOpen01{Float64}
const CloseOpen12_64   = CloseOpen12{Float64}

CloseOpen01(::Type{T}=Float64) where {T<:AbstractFloat} = CloseOpen01{T}()
CloseOpen12(::Type{T}=Float64) where {T<:AbstractFloat} = CloseOpen12{T}()

CloseOpen(::Type{T}=Float64) where {T<:AbstractFloat} = CloseOpen01{T}()
CloseOpen(a::T, b::T) where {T<:AbstractFloat} = CloseOpenAB{T}(a, b)

OpenClose(::Type{T}=Float64) where {T<:AbstractFloat} = OpenClose01{T}()
OpenClose(a::T, b::T) where {T<:AbstractFloat} = OpenCloseAB{T}(a, b)

CloseClose(::Type{T}=Float64) where {T<:AbstractFloat} = CloseClose01{T}()
CloseClose(a::T, b::T) where {T<:AbstractFloat} = CloseCloseAB{T}(a, b)

OpenOpen(::Type{T}=Float64) where {T<:AbstractFloat} = OpenOpen01{T}()
OpenOpen(a::T, b::T) where {T<:AbstractFloat} = OpenOpenAB{T}(a, b)

# convenience functions

CloseOpen(a, b) = CloseOpen(promote(a, b)...)
CloseOpen(a::T, b::T) where {T} = CloseOpen(AbstractFloat(a), AbstractFloat(b))

OpenClose(a, b) = OpenClose(promote(a, b)...)
OpenClose(a::T, b::T) where {T} = OpenClose(AbstractFloat(a), AbstractFloat(b))

CloseClose(a, b) = CloseClose(promote(a, b)...)
CloseClose(a::T, b::T) where {T} = CloseClose(AbstractFloat(a), AbstractFloat(b))

OpenOpen(a, b) = OpenOpen(promote(a, b)...)
OpenOpen(a::T, b::T) where {T} = OpenOpen(AbstractFloat(a), AbstractFloat(b))

## Bernoulli

struct Bernoulli{T<:Number} <: Distribution{T}
    p::Float64

    Bernoulli{T}(p::Real) where {T} = let pf = Float64(p)
        0.0 <= pf <= 1.0 ? new(pf) :
            throw(DomainError(p, "Bernoulli: parameter p must satisfy 0.0 <= p <= 1.0"))
    end
end

Bernoulli(p::Real=0.5) = Bernoulli(Int, p)
Bernoulli(::Type{T}, p::Real=0.5) where {T<:Number} = Bernoulli{T}(p)

## Categorical

struct Categorical{T<:Number} <: Distribution{T}
    cdf::Vector{Float64}

    function Categorical{T}(weigths) where T
        if !isa(weigths, AbstractArray)
            # necessary for accumulate
            # TODO: will not be necessary anymore in Julia 1.5
            weigths = collect(weigths)
        end
        weigths = vec(weigths)

        isempty(weigths) &&
            throw(ArgumentError("Categorical requires at least one category"))

        s = Float64(sum(weigths))
        cdf = accumulate(weigths; init=0.0) do x, y
            x + Float64(y) / s
        end
        @assert isapprox(cdf[end], 1.0) # really?
        cdf[end] = 1.0 # to be sure the algo terminates
        new{T}(cdf)
    end
end

Categorical(weigths) = Categorical{Int}(weigths)

Categorical(n::Number) =
    Categorical{typeof(n)}(Iterators.repeated(1.0 / Float64(n), Int(n)))


@kwdef mutable struct SDMHandler <: AbstractHandler
    id = 0
end

@register(SDMHandler, SDM, [Number] => [DBit])

function (handler::SDMHandler)(netlist::Netlist,
                               inputs::Vector{Variable},
                               outputs::Vector{Variable})
    outstring = """
        $stdcomment
        // BEGIN sdm$(handler.id)
        sdm #(
                .BIT_WIDTH(BIT_WIDTH),
                .INT_WIDTH(INT_WIDTH)
            ) sdm$(handler.id) (
                .CLK(CLK),
                .nRST(nRST),
                .x($(inputs[1].name)),
                .y($(outputs[1].name))
            );
        // END sdm$(handler.id)
        \n"""

    handler.id += 1

    return outstring
end


using TcpInstruments
using Plots
scope = initialize(AgilentDSOX4034A,  "10.1.30.32")
data = get_data(scope, 1)
plot(data)


module HyperHessians

using SIMD: SIMD, Vec, vstore
using CommonSubexpressions: cse

include("rules.jl")

# This is currently "artifically" restricted to "square" HyperDuals where
# ϵ1 and ϵ2 have the same length.
struct HyperDual{N, T} <: Real
    v::T
    ϵ1::Vec{N,T} # Vec{M,T}
    ϵ2::Vec{N,T}
    ϵ12::NTuple{N, Vec{N,T}} # NTuple{M, Vec{N,T}}
end
HyperDual(v::T, ϵ1::Vec{N,T}, ϵ2::Vec{N,T}) where {N, T} = HyperDual(v, ϵ1, ϵ2, ntuple(i -> zero(Vec{N,T}), Val(N)))
HyperDual{N}(v::T) where {N, T} = HyperDual(v, zero(Vec{N,T}),zero(Vec{N,T}))
HyperDual{N,T}(v) where {N, T} = HyperDual{N}(T(v))

function HyperDual(v::T1, ϵ1::Vec{N,T2}, ϵ2::Vec{N,T2}, ϵ12::NTuple{N, Vec{N,T2}}) where {N, T1, T2}
    T = promote_type(T1, T2)
    HyperDual(T(v), convert(Vec{N,T}, ϵ1), convert(Vec{N,T}, ϵ2), convert.(Vec{N,T}, ϵ12))
end

Base.promote_rule(::Type{HyperDual{N,T1}}, ::Type{HyperDual{N,T2}}) where {N,T1,T2} = HyperDual{N, promote_type(T1, T2)}
Base.convert(::Type{HyperDual{N,T1}}, h::HyperDual{N,T2}) where {N,T1,T2} =
    HyperDual{N,T1}(T1(h.v), convert(Vec{N,T1}, h.ϵ1), convert(Vec{N,T1}, h.ϵ2), convert.(Vec{N,T1}, h.ϵ12))

# Make this look nicer
function Base.show(io::IO, h::HyperDual)
    print(io, h.v, " + ", Tuple(h.ϵ1), "ϵ1", " + ", Tuple(h.ϵ2), "ϵ2", " + ",  map(Tuple, h.ϵ12), "ϵ12")
end

Base.one(::Type{HyperDual{N, T}}) where {N, T} = HyperDual{N}(one(T))
Base.zero(::Type{HyperDual{N, T}}) where {N, T} = HyperDual{N}(zero(T))
Base.one(::HyperDual{N, T}) where {N, T} = one(HyperDual{N,T})
Base.zero(::HyperDual{N, T}) where {N, T} = zero(HyperDual{N,T})
Base.float(h::HyperDual{N,T}) where {N,T} = convert(HyperDual{N,float(T)}, h)

# Unary
@inline Base.:(-)(h::HyperDual{N}) where {N} = HyperDual(-h.v, -h.ϵ1, -h.ϵ2, ntuple(i -> -h.ϵ12[i], Val(N)))
@inline Base.:(+)(h::HyperDual) = h

# Binary
for f in (:+, :-)
    @eval begin
        @inline Base.$f(h1::HyperDual{N,T}, h2::HyperDual{N,T}) where {N,T} =
            HyperDual($f(h1.v, h2.v), $f(h1.ϵ1, h2.ϵ1), $f(h1.ϵ2, h2.ϵ2), ntuple(i -> $f(h1.ϵ12[i], h2.ϵ12[i]), Val(N)))
        @inline Base.$f(h1::HyperDual{N,T1}, h2::HyperDual{N,T2}) where {N,T1,T2} = $f(promote(h1, h2)...)
        @inline Base.$f(h::HyperDual{N}, r::Real) where {N} =
            HyperDual($f(h.v, r), h.ϵ1, h.ϵ2, h.ϵ12)
        @inline Base.$f(r::Real, h::HyperDual{N}) where {N} =
            HyperDual($f(r, h.v), h.ϵ1, h.ϵ2, h.ϵ12)
    end
end

for f in (:*, :/)
    @eval @inline Base.$f(h::HyperDual{N}, r::Real) where {N} =
        HyperDual($f(h.v, r), $f(h.ϵ1, r), $f(h.ϵ2, r), ntuple(i -> $f(h.ϵ12[i], r), Val(N)))
    @eval @inline Base.$f(r::Real, h::HyperDual{N}) where {N} =
        HyperDual($f(r, h.v), $f(r, h.ϵ1), $f(r, h.ϵ2), ntuple(i -> $f(r, h.ϵ12[i]), Val(N)))
end

@inline Base.:(/)(h1::HyperDual{N,T}, h2::HyperDual{N,T}) where {N,T} = h1 * inv(h2)

# Get's rid of a bunch of boundserror throwing in the LLVM IR, not sure it matters for runtime performance...
unsafe_getindex(v::SIMD.Vec, i::SIMD.IntegerTypes) = SIMD.Intrinsics.extractelement(v.data, i-1)
@inline ⊗(v1::Vec{N}, v2::Vec{N}) where {N} = ntuple(i -> unsafe_getindex(v1, i) * v2, Val(N))

@inline Base.:(*)(h1::HyperDual{N,T1}, h2::HyperDual{N,T2}) where {N,T1,T2} = *(promote(h1, h2)...)
@inline function Base.:(*)(h1::HyperDual{N,T}, h2::HyperDual{N,T}) where {N,T}
    r = h1.v * h2.v
    ϵ1 = muladd(h1.v, h2.ϵ1, h1.ϵ1 * h2.v)
    ϵ2 = muladd(h1.v, h2.ϵ2, h1.ϵ2 * h2.v)
    ϵ12_1 = h1.ϵ1 ⊗ h2.ϵ2
    ϵ12_2 = h2.ϵ1 ⊗ h1.ϵ2
    ϵ12 = ntuple(i -> muladd(h1.v, h2.ϵ12[i], muladd(h1.ϵ12[i], h2.v, ϵ12_1[i] + ϵ12_2[i])), Val(N))
    return HyperDual(r, ϵ1, ϵ2, ϵ12)
end
@inline Base.literal_pow(::typeof(^), x::HyperDual, ::Val{0}) = one(typeof(x))
@inline Base.literal_pow(::typeof(^), x::HyperDual, ::Val{1}) = x
@inline Base.literal_pow(::typeof(^), x::HyperDual, ::Val{2}) = x*x
@inline Base.literal_pow(::typeof(^), x::HyperDual, ::Val{3}) = x*x*x

for (f, f′, f′′) in DIFF_RULES
    expr = quote
        # Verify that the cse still works properly when changing this.
        v = $f(x)
        f′ = $f′
        f′′ = $f′′
        x23 = (f′′ * h.ϵ1) ⊗ h.ϵ2
        return HyperDual(v, h.ϵ1 * f′, h.ϵ2 * f′, ntuple(i -> h.ϵ12[i]*f′ + x23[i], Val(N)))
    end
    cse_expr = cse(expr; warn=false)
    @eval @inline function Base.$f(h::HyperDual{N,T}) where {N,T}
        x = h.v
        $cse_expr
    end
end

@inline function Base.sin(h::HyperDual{N}) where {N}
    s, c = sincos(h.v)
    f′, f′′ = c, -s
    x23 = (f′′ * h.ϵ1) ⊗ h.ϵ2
    return HyperDual(s, h.ϵ1 * f′, h.ϵ2 * f′, ntuple(i -> h.ϵ12[i]*f′ + x23[i], Val(N)))
end

@inline function Base.cos(h::HyperDual{N}) where {N}
    s, c = sincos(h.v)
    f′,  f′′ = -s, -c
    x23 = (f′′ * h.ϵ1) ⊗ h.ϵ2
    return HyperDual(c, h.ϵ1 * f′,  h.ϵ2 * f′, ntuple(i -> h.ϵ12[i]*f′ + x23[i], Val(N)))
end

##########
# Chunks #
##########

const DEFAULT_CHUNK_THRESHOLD = 8

struct Chunk{N} end

const CHUNKS = [Chunk{i}() for i in 1:DEFAULT_CHUNK_THRESHOLD]

function Chunk(input_length::Integer)
    N = pickchunksize(input_length)
    0 < N <= DEFAULT_CHUNK_THRESHOLD && return CHUNKS[N] # TODO: Check if this matters
    return Chunk{N}()
end

Chunk(x::AbstractArray) = Chunk(length(x))

# TODO: MOre testing needed here.
function pickchunksize(input_length)
    input_length == 1 && return 1
    input_length == 2 && return 2
    3 <= input_length <= 4 && return 4 # 3 seems to be very bad for some reason
    return 8
end

(chunksize(::Chunk{N})::Int) where {N} = N

##################
# Hessian Config #
##################

abstract type AbstractHessianConfig end

struct HessianConfig{D <: AbstractVector{<:HyperDual}, S} <: AbstractHessianConfig
    duals::D
    seeds::S
end
(chunksize(cfg::HessianConfig)::Int) = length(cfg.seeds)

function HessianConfig(x::AbstractVector{T}, chunk=Chunk(x)::Chunk) where {T}
    N = chunksize(chunk)
    @assert 0 < N
    duals = similar(x, HyperDual{N,T}) # not Vector
    seeds = collect(construct_seeds(NTuple{N,T}))
    return HessianConfig(duals, seeds)
end

struct HessianConfigThreaded{D <: AbstractVector{<:HyperDual}, S} <: AbstractHessianConfig
    duals::Vector{D}
    seeds::S
end

function HessianConfigThreaded(x::AbstractVector{T}, chunk=Chunk(x)::Chunk) where {T}
    N = chunksize(chunk)
    @assert 0 < N
    duals = [similar(x, HyperDual{N,T}) for i in 1:Threads.nthreads()] # not Vector
    seeds = collect(construct_seeds(NTuple{N,T}))
    return HessianConfigThreaded(duals, seeds)
end

(chunksize(cfg::HessianConfigThreaded)::Int) = length(cfg.seeds)
HessianConfig(cfg::HessianConfigThreaded) = HessianConfig(first(cfg.duals), cfg.seeds)

###########
# Seeding #
###########

@generated function single_seed(::Type{NTuple{N,T}}, ::Val{i}) where {N,T,i}
    ex = Expr(:tuple, [ifelse(i === j, :(one(T)), :(zero(T))) for j in 1:N]...)
    return :(Vec($(ex)))
end

@generated construct_seeds(::Type{NTuple{N,T}}) where {N,T} =
    Expr(:tuple, [:(single_seed(NTuple{N,T}, Val{$i}())) for i in 1:N]...)

seed_epsilon_1(d::HyperDual{N,T}, ϵ1) where {N,T} = HyperDual{N,T}(d.v, ϵ1, d.ϵ2, d.ϵ12)
seed_epsilon_2(d::HyperDual{N,T}, ϵ2) where {N,T} = HyperDual{N,T}(d.v, d.ϵ1, ϵ2, d.ϵ12)

function seed!(d::AbstractVector{<:HyperDual{N}}, x, seeds, block_i, block_j) where {N}
    d .= HyperDual{N}.(x)
    index_i = (block_i-1) * N + 1
    index_j = (block_j-1) * N + 1
    range_i = index_i : min(length(x), (index_i + N -1))
    range_j = index_j : min(length(x), (index_j + N -1))
    chunks_i = length(range_i)
    chunks_j = length(range_j)

    d[range_i] .= seed_epsilon_1.(view(d, range_i), view(seeds, 1:chunks_i))
    d[range_j] .= seed_epsilon_2.(view(d, range_j), view(seeds, 1:chunks_j))
    return d
end

# Hessian
@noinline check_scalar(x) =
    x isa Number || throw(error("expected a scalar to be returned from function passed to `hessian`"))

function extract_hessian!(H::AbstractMatrix, v::HyperDual)
    Base.require_one_based_indexing(H)
    @inbounds for i in 1:size(H, 2)
        H[:, i] .= Tuple(v.ϵ12[i])
    end
    return H
end

function symmetrize!(H::AbstractMatrix)
    Base.require_one_based_indexing(H)
    for i in 1:size(H, 1)
        for j in i:size(H,2)
            H[j,i] = H[i,j]
        end
    end
    return H
end

function extract_hessian!(H::AbstractMatrix, v::HyperDual{N}, block_i::Int, block_j::Int) where {N}
    index_i = (block_i-1) * N + 1
    index_j = (block_j-1) * N + 1
    range_i = index_i : (index_i + N -1)
    range_j = index_j : (index_j + N -1)

    for (I, i) in enumerate(range_i)
        for (J, j) in enumerate(range_j)
            if checkbounds(Bool, H, i, j)
                H[i, j] = v.ϵ12[I][J]
            end
        end
    end
    return H
end

function extract_gradient!(G::AbstractVector, v::HyperDual, block_i::Int)
    index_i = (block_i-1) * N + 1
    range_i = index_i : (index_i + N -1)
    for (I, i) in enumerate(range_i)
        G[i] .= v.ϵ1[I]
    end
    return G
end

###############
# Hessian API #
###############

# Scalar
function hessian(f, x::Real)
    dual = HyperDual(x, Vec(one(x)), Vec(one(x)))
    v = f(dual)
    check_scalar(v)
    return @inbounds v.ϵ12[1][1]
end

hessian(f, x::AbstractVector)                             = hessian!(similar(x, axes(x,1), axes(x,1)), f, x, HessianConfig(x))
hessian(f, x::AbstractVector, cfg::AbstractHessianConfig) = hessian!(similar(x, axes(x,1), axes(x,1)), f, x, cfg)

function hessian!(H::AbstractMatrix, f, x::AbstractVector{T}, cfg::AbstractHessianConfig) where {T,N}
    if chunksize(cfg) == length(x)
        if cfg isa HessianConfigThreaded
            cfg = HessianConfig(cfg)
        end
        return hessian_vector!(H, f, x, cfg)
    else
        return hessian_chunk!(H, f, x, cfg)
    end
end

function hessian_vector!(H::AbstractMatrix, f, x::AbstractVector, cfg::HessianConfig)
    @assert size(H,1) == size(H,2) == length(x)
    cfg.duals .= HyperDual.(x, cfg.seeds, cfg.seeds)
    v = f(cfg.duals)
    check_scalar(v)
    return extract_hessian!(H, v)
end

function hessian_chunk!(H::AbstractMatrix, f, x::AbstractVector{T}, cfg::HessianConfig) where {T,N}
    @assert size(H,1) == size(H,2) == length(x)
    n_chunks = ceil(Int, length(x) / chunksize(cfg))
    for i in 1:n_chunks
        for j in i:n_chunks
            seed!(cfg.duals, x, cfg.seeds, i, j)
            v = f(cfg.duals)
            check_scalar(v)
            extract_hessian!(H, v, i, j)
        end
    end
    symmetrize!(H)
    return H
end

# Threaded chunk mode
function linear_to_cartesian(n::Int, k::Int)
    k -= 1
    n += 1
    i = n - 2 - floor(sqrt(-8*k + 4*n*(n-1)-7)/2 - 0.5)
    j = k + i + 1 - n*(n-1)/2 + (n-i)*((n-i)-1)/2
    return round(Int, i+1), round(Int, j)
end

function hessian_chunk!(H::AbstractMatrix, f, x::AbstractVector{T}, cfg::HessianConfigThreaded) where {T,N}
    @assert size(H,1) == size(H,2) == length(x)
    n_chunks = ceil(Int, length(x) / chunksize(cfg))
    Threads.@threads for k in 1:((n_chunks*(n_chunks+1)) ÷ 2)
        i, j = linear_to_cartesian(n_chunks, k)
        duals = cfg.duals[Threads.threadid()]
        seed!(duals, x, cfg.seeds, i, j)
        v = f(duals)
        check_scalar(v)
        extract_hessian!(H, v, i, j)
    end
    symmetrize!(H)
    return H
end

end # module


function checkArgumentOrder(mdp::MDP)

  statedim = length(mdp.statemap)
  actiondim = length(mdp.actionmap)

  ntransitionargs = length(mdp.transition.argnames)
  nrewardargs = length(mdp.reward.argnames)

  sadim = statedim + actiondim
  saspdim = sadim + statedim

  if sadim != nrewardargs
    @error(string(
      "the number of reward function input arguments must be the same as ",
      "the sum of the number of state and action variables"))
  elseif sadim != ntransitionargs && saspdim != ntransitionargs
    @error(string(
      "the number of transition function input arguments must be the same as ",
      "the sum of the number of state, action, and next state variables or ",
      "the sum of the number of state and action variables"))
  end

  for iarg in 1:nrewardargs
    if mdp.reward.argnames[iarg] != mdp.transition.argnames[iarg]
      @error(string(
        "transition and reward function state and action variable input arguments",
        "must be consistent in both naming and order"))
    end
  end

  if length(mdp.transition.argnames) > length(mdp.reward.argnames)
    for iarg in 1:nrewardargs - actiondim
      if mdp.transition.argnames[iarg] != mdp.transition.argnames[iarg + sadim]
        @error(string(
          "transition type T(s,a,s')'s state s and next state s' variable ",
          "input arguments must be consistent in both naming and order"))
      end
    end
  end

end

function checkTransition(mdp::MDP)
  args = randargs(mdp)
  transitionval = mdp.transition.fn(args...)

  if isa(transitionval, Real) || isa(transitionval, Vector)
    checkTransition(mdp, args, transitionval)
  else
    @error(string(
      "transition function provided is not a correctly defined ",
      "T(s,a,s') or T(s,a) type function, check the return type\n",
      "return type: ", typeof(transitionval)))
  end
end

function checkTransition(mdp::MDP, args::Vector, transitionval::Real)
  if transitionval < 0 || transitionval > 1
    @warn(string(  # warn not error because we might have sampled a non-existent state
      "transition function provided is of type T(s,a,s'), ",
      "but the value returned from a random state is not ",
      "a valid probability value bounded to [0,1]\n",
      "argument names: ", mdp.transition.argnames, "\n",
      "random state: ", args, "\n",
      "return value: ", transitionval))
  end
end

function checkTransition(mdp::MDP, args::Vector, transitionval::Vector)
  nargs = length(mdp.transition.argnames)
  sumprob = 0.0

  for returnval in transitionval  # |returnval| is a (state, prob) pair
    state = returnval[1]
    prob = returnval[2]
    sumprob += prob

    if length(state) + length(prob) != nargs || !isvalid(mdp, state)
      @error(string(
        "transition function provided is of type T(s,a), ",
        "but one of the states returned from a random state is ",
        "either not bounded by its range or not in the set of values\n",
        "argument names: ", mdp.transition.argnames, "\n",
        "random state: ", args, "\n",
        "return value: ", state, "\n",
        "probability: ", prob))
    elseif prob < 0 || prob > 1
      @error(string(
        "transition function provided is of type T(s,a), ",
        "but one of the probabilities returned from a random state is ",
        "a valid probability value bounded to [0,1]\n",
        "argument names: ", mdp.transition.argnames, "\n",
        "random state: ", args, "\n",
        "return value: ", state, "\n",
        "probability: ", prob))
    end
  end

  if sumprob != 1
    @warn(string(  # warn not error because we might have sampled a non-existent state
      "transition function provided is of type T(s,a), ",
      "but the sum of transition probabilities returned from a random state ",
      "does not sum to 1\n",
      "argument names: ", mdp.transition.argnames, "\n",
      "random state: ", args, "\n",
      "return value: ", returnval, "\n",
      "sum of transition probabilities: ", sumprob))
  end
end

function randargs(mdp::MDP)
  nargs = length(mdp.transition.argnames)
  args = Array{Any,1}(undef, nargs)
  for iarg in 1:nargs
    argname = mdp.transition.argnames[iarg]
    if haskey(mdp.statemap, argname)
      args[iarg] = lazySample(mdp.statemap[argname])
    elseif haskey(mdp.actionmap, argname)
      args[iarg] = lazySample(mdp.actionmap[argname])
    else
      @error(string("variable ", lazyvar.varname, " has not been defined"))
    end
  end
  args
end

function lazySample(lazyvar::LazyVar)
  if isa(lazyvar, RangeVar)
    return lazyvar.minval + rand() * (lazyvar.maxval - lazyvar.minval)
  elseif isa(lazyvar, ValuesVar)
    return lazyvar.values[rand(1:length(lazyvar.values))]
  else
    @error(string("variable", lazyvar.varname, " is not a valid subtype of LazyVar"))
  end
end

function isvalid(mdp::MDP, state::Vector)
  for iarg in length(mdp.statemap)
    argname = mdp.transition.argnames[iarg]
    if isa(mdp.statemap[argname], RangeVar) &&
        (state[iarg] < mdp.statemap[argname].minval ||
        state[iarg] > mdp.statemap[argname].maxval)
      return false
    elseif isa(mdp.statemap[argname], ValuesVar) &&
        !(state[iarg] in mdp.statemap[argname].values)
      return false
    end
  end
  return true
end


using AutoPreallocation
using BenchmarkTools
using Test

# functions to test on
f_ones() = ones(64)
f_matmul() = ones(32,64) * ones(64, 2)

f_matmul_noprealloc() = ones(32,64) * @no_prealloc(ones(64, 2))

@testset "ones example" begin
    @assert (@ballocated f_ones()) === 624
    val, record = record_allocations(f_ones)
    @test (@ballocated avoid_allocations($record, f_ones)) <= 64
end

@testset "matmul example" begin
    @assert (@ballocated f_matmul()) === 18_304
    val, record = record_allocations(f_matmul)
    # NOTE: (@Roger-luo) not sure why this is 256 on my machine
    @test (@ballocated avoid_allocations($record, f_matmul)) <= 352

    _, pf = preallocate(f_matmul)
    @test (@ballocated $pf()) <= 352
end

@testset "noprealloc example" begin
    @assert (@ballocated f_matmul_noprealloc()) === 18_304
    val, record = record_allocations(f_matmul_noprealloc)
    @test length(record.allocations) == 2
    @test record.initial_sizes == [(32, 64), (32, 2)]
    @test (@ballocated avoid_allocations($record, f_matmul_noprealloc)) <= 1520
end

@testset "check thread-safe" begin
    f(x, y) = x * y
    n = Threads.nthreads()
    results = Vector{Any}(undef, n)
    As = [rand(4, 4) for _ in 1:n]
    Bs = [rand(4, 4) for _ in 1:n]
    _, pf = preallocate(f, As[1], Bs[1])
    Threads.@threads for k in 1:n
        results[k] = pf(As[k], Bs[k])
    end

    # if it's not thread-safe, the result
    # will be modified by other multiplications
    # since they will share the same memory
    for k in 1:n
        @test results[k] ≈ f(As[k], Bs[k])
    end
end

@testset "resizing operations" begin
    function push_pop_test(a)
        x = [0, 0, 0]
        @test length(x) == 3
        push!(x, a)
        push!(x, 10a)
        @test length(x) == 5
        pop!(x)
        @test length(x) == 4
    end

    _, p_push_pop_test = preallocate(push_pop_test, 1)
    p_push_pop_test(2)  # check works on second call
    p_push_pop_test(3)  # check works on third call
    p_push_pop_test(4)  # check works on fourth call
end


# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder

# Collection of sources required to build MPCBuilder
name = "ARB"
version = v"2.16.0"
sources = [
    "https://github.com/fredrik-johansson/arb/archive/$version.tar.gz" =>
    "77464be4d34a511bb004457f862fec857ff934b0ed58d56d6f52d76ebadd4daf",
]

# Bash recipe for building across all platforms

script = """cd arb-$version
"""*raw"""
UNAME=`uname`
if [[ ${UNAME} == MSYS_NT-6.3 ]]; then
    ln -sf $prefix/bin/libflint-13.dll $prefix/lib/libflint.dll.a
fi
./configure --prefix=$prefix --enable-shared --disable-static --with-gmp=$prefix --with-mpfr=$prefix --with-flint=$prefix
make -j
make install
if [[ ${UNAME} == MSYS_NT-6.3 ]]; then
    rm -f $prefix/lib/libflint.dll.a
fi
"""
# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = [
    BinaryProvider.Linux(:aarch64, :glibc, :blank_abi),
    BinaryProvider.Windows(:i686, :blank_libc, :blank_abi),
    BinaryProvider.Linux(:armv7l, :glibc, :eabihf),
    BinaryProvider.Windows(:x86_64, :blank_libc, :blank_abi),
    BinaryProvider.Linux(:x86_64, :glibc, :blank_abi),
    BinaryProvider.MacOS(:x86_64, :blank_libc, :blank_abi),
    BinaryProvider.Linux(:i686, :glibc, :blank_abi)
]

# The products that we will ensure are always built
products(prefix) = [
    LibraryProduct(prefix, "libarb", :libarb)
]

# Dependencies that must be installed before this package can be built
dependencies = [
    "https://github.com/JuliaMath/GMPBuilder/releases/download/v6.1.2-2/build_GMP.v6.1.2.jl"
    "https://github.com/JuliaMath/MPFRBuilder/releases/download/v4.0.1/build.jl"
    "https://github.com/isuruf/FLINTBuilder/releases/download/v2.5.3%2Bdev0.ea9f102a0b/build_FLINT.v2.5.3+dev0.ea9f102a0b.jl"
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies)