text:
@doc raw"""
    ConjugateGradientOptions <: AbstractGradientOptions

specify options for a conjugate gradient descent algoritm, that solves a
[`GradientProblem`].

# Fields
* `x` – the current iterate, a point on a manifold
* `gradient` – the current gradient
* `δ` – the current descent direction, i.e. also tangent vector
* `β` – the current update coefficient rule, see .
* `coefficient` – a [`DirectionUpdateRule`](@ref) function to determine the new `β`
* `stepsize` – a [`Stepsize`](@ref) function
* `stop` – a [`StoppingCriterion`](@ref)
* `retraction_method` – (`ExponentialRetraction()`) a type of retraction

# See also
[`conjugate_gradient_descent`](@ref), [`GradientProblem`](@ref), [`ArmijoLinesearch`](@ref)
"""
mutable struct ConjugateGradientDescentOptions{
    T,
    TCoeff<:DirectionUpdateRule,
    TStepsize<:Stepsize,
    TStop<:StoppingCriterion,
    TRetr<:AbstractRetractionMethod,
    TVTM<:AbstractVectorTransportMethod,
} <: AbstractGradientOptions
    x::T
    gradient::T
    δ::T
    β::Float64
    coefficient::TCoeff
    stepsize::TStepsize
    stop::TStop
    retraction_method::TRetr
    vector_transport_method::TVTM
    function ConjugateGradientDescentOptions{T}(
        x0::T,
        sC::StoppingCriterion,
        s::Stepsize,
        dC::DirectionUpdateRule,
        retr::AbstractRetractionMethod=ExponentialRetraction(),
        vtr::AbstractVectorTransportMethod=ParallelTransport(),
    ) where {T}
        o = new{T,typeof(dC),typeof(s),typeof(sC),typeof(retr),typeof(vtr)}()
        o.x = x0
        o.stop = sC
        o.retraction_method = retr
        o.stepsize = s
        o.coefficient = dC
        o.vector_transport_method = vtr
        return o
    end
end
function ConjugateGradientDescentOptions(
    x::T,
    sC::StoppingCriterion,
    s::Stepsize,
    dU::DirectionUpdateRule,
    retr::AbstractRetractionMethod=ExponentialRetraction(),
    vtr::AbstractVectorTransportMethod=ParallelTransport(),
) where {T}
    return ConjugateGradientDescentOptions{T}(x, sC, s, dU, retr, vtr)
end

@doc raw"""
    ConjugateDescentCoefficient <: DirectionUpdateRule

Computes an update coefficient for the conjugate gradient method, where
the [`ConjugateGradientDescentOptions`](@ref)` o` include the last iterates
``x_k,ξ_k``, the current iterates ``x_{k+1},ξ_{k+1}`` and the last update
direction ``\delta=\delta_k``, where the last three ones are stored in the
variables with prequel `Old` based on [^Flethcer1987] adapted to manifolds:

```math
β_k =
\frac{ \lVert ξ_{k+1} \rVert_{x_{k+1}}^2 }
{\langle -\delta_k,ξ_k \rangle_{x_k}}.
```

See also [`conjugate_gradient_descent`](@ref)

# Constructor
    ConjugateDescentCoefficient(a::StoreOptionsAction=())

Construct the conjugate descnt coefficient update rule, a new storage is created by default.

[^Flethcer1987]:
    > R. Fletcher, __Practical Methods of Optimization vol. 1: Unconstrained Optimization__
    > John Wiley & Sons, New York, 1987. doi [10.1137/1024028](https://doi.org/10.1137/1024028)
"""
mutable struct ConjugateDescentCoefficient <: DirectionUpdateRule
    storage::StoreOptionsAction
    function ConjugateDescentCoefficient(
        a::StoreOptionsAction=StoreOptionsAction((:x, :gradient))
    )
        return new(a)
    end
end
function (u::ConjugateDescentCoefficient)(
    p::GradientProblem, o::ConjugateGradientDescentOptions, i
)
    if !all(has_storage.(Ref(u.storage), [:x, :gradient]))
        update_storage!(u.storage, o) # if not given store current as old
        return 0.0
    end
    xOld, gradientOld = get_storage.(Ref(u.storage), [:x, :gradient])
    update_storage!(u.storage, o)
    return inner(p.M, o.x, o.gradient, o.gradient) / inner(p.M, xOld, -o.δ, gradientOld)
end

@doc raw"""
    DaiYuanCoefficient <: DirectionUpdateRule

Computes an update coefficient for the conjugate gradient method, where
the [`ConjugateGradientDescentOptions`](@ref)` o` include the last iterates
``x_k,ξ_k``, the current iterates ``x_{k+1},ξ_{k+1}`` and the last update
direction ``\delta=\delta_k``, where the last three ones are stored in the
variables with prequel `Old` based on [^DaiYuan1999]

adapted to manifolds: let ``\nu_k = ξ_{k+1} - P_{x_{k+1}\gets x_k}ξ_k``,
where ``P_{a\gets b}(⋅)`` denotes a vector transport from the tangent space at ``a`` to ``b``.

Then the coefficient reads

````math
β_k =
\frac{ \lVert ξ_{k+1} \rVert_{x_{k+1}}^2 }
{\langle P_{x_{k+1}\gets x_k}\delta_k, \nu_k \rangle_{x_{k+1}}}.
````

See also [`conjugate_gradient_descent`](@ref)

# Constructor
    DaiYuanCoefficient(
        t::AbstractVectorTransportMethod=ParallelTransport(),
        a::StoreOptionsAction=(),
    )

Construct the Dai Yuan coefficient update rule, where the parallel transport is the
default vector transport and a new storage is created by default.

[^DaiYuan1999]:
    > [Y. H. Dai and Y. Yuan, A nonlinear conjugate gradient method with a strong global convergence property,
    > SIAM J. Optim., 10 (1999), pp. 177–182.
    > doi: [10.1137/S1052623497318992](https://doi.org/10.1137/S1052623497318992)
"""
mutable struct DaiYuanCoefficient{TVTM<:AbstractVectorTransportMethod} <:
               DirectionUpdateRule
    transport_method::TVTM
    storage::StoreOptionsAction
    function DaiYuanCoefficient(
        t::AbstractVectorTransportMethod=ParallelTransport(),
        a::StoreOptionsAction=StoreOptionsAction((:x, :gradient, :δ)),
    )
        return new{typeof(t)}(t, a)
    end
end
function (u::DaiYuanCoefficient)(p::GradientProblem, o::ConjugateGradientDescentOptions, i)
    if !all(has_storage.(Ref(u.storage), [:x, :gradient, :δ]))
        update_storage!(u.storage, o) # if not given store current as old
        return 0.0
    end
    xOld, gradientOld, δOld = get_storage.(Ref(u.storage), [:x, :gradient, :δ])
    update_storage!(u.storage, o)

    gradienttr = vector_transport_to(p.M, xOld, gradientOld, o.x, u.transport_method)
    ν = o.gradient - gradienttr #notation y from [HZ06]
    δtr = vector_transport_to(p.M, xOld, δOld, o.x, u.transport_method)
    return inner(p.M, o.x, o.gradient, o.gradient) / inner(p.M, xOld, δtr, ν)
end

@doc raw"""
    FletcherReevesCoefficient <: DirectionUpdateRule

Computes an update coefficient for the conjugate gradient method, where
the [`ConjugateGradientDescentOptions`](@ref)` o` include the last iterates
``x_k,ξ_k``, the current iterates ``x_{k+1},ξ_{k+1}`` and the last update
direction ``\delta=\delta_k``, where the last three ones are stored in the
variables with prequel `Old` based on [^FletcherReeves1964] adapted to manifolds:

````math
β_k =
\frac{\lVert ξ_{k+1}\rVert_{x_{k+1}}^2}{\lVert ξ_{k}\rVert_{x_{k}}^2}.
````

See also [`conjugate_gradient_descent`](@ref)

# Constructor
    FletcherReevesCoefficient(a::StoreOptionsAction=())

Construct the Fletcher Reeves coefficient update rule, a new storage is created by default.

[^FletcherReeves1964]:
    > R. Fletcher and C. Reeves, Function minimization by conjugate gradients,
    > Comput. J., 7 (1964), pp. 149–154.
    > doi: [10.1093/comjnl/7.2.149](http://dx.doi.org/10.1093/comjnl/7.2.149)
"""
mutable struct FletcherReevesCoefficient <: DirectionUpdateRule
    storage::StoreOptionsAction
    function FletcherReevesCoefficient(
        a::StoreOptionsAction=StoreOptionsAction((:x, :gradient))
    )
        return new(a)
    end
end
function (u::FletcherReevesCoefficient)(
    p::GradientProblem, o::ConjugateGradientDescentOptions, i
)
    if !all(has_storage.(Ref(u.storage), [:x, :gradient]))
        update_storage!(u.storage, o) # if not given store current as old
    end
    xOld, gradientOld = get_storage.(Ref(u.storage), [:x, :gradient])
    update_storage!(u.storage, o)
    return inner(p.M, o.x, o.gradient, o.gradient) /
           inner(p.M, xOld, gradientOld, gradientOld)
end

@doc raw"""
    HagerZhangCoefficient <: DirectionUpdateRule

Computes an update coefficient for the conjugate gradient method, where
the [`ConjugateGradientDescentOptions`](@ref)` o` include the last iterates
``x_k,ξ_k``, the current iterates ``x_{k+1},ξ_{k+1}`` and the last update
direction ``\delta=\delta_k``, where the last three ones are stored in the variables with
prequel `Old` based on [^HagerZhang2005]
adapted to manifolds: let ``\nu_k = ξ_{k+1} - P_{x_{k+1}\gets x_k}ξ_k``,
where ``P_{a\gets b}(⋅)`` denotes a vector transport from the tangent space at ``a`` to ``b``.

````math
β_k = \Bigl\langle\nu_k -
\frac{ 2\lVert \nu_k\rVert_{x_{k+1}}^2 }{ \langle P_{x_{k+1}\gets x_k}\delta_k, \nu_k \rangle_{x_{k+1}} }
P_{x_{k+1}\gets x_k}\delta_k,
\frac{ξ_{k+1}}{ \langle P_{x_{k+1}\gets x_k}\delta_k, \nu_k \rangle_{x_{k+1}} }
\Bigr\rangle_{x_{k+1}}.
````

This method includes a numerical stability proposed by those authors.

See also [`conjugate_gradient_descent`](@ref)

# Constructor
    HagerZhangCoefficient(
        t::AbstractVectorTransportMethod=ParallelTransport(),
        a::StoreOptionsAction=(),
    )

Construct the Hager Zhang coefficient update rule, where the parallel transport is the
default vector transport and a new storage is created by default.

[^HagerZhang2005]:
    > [W. W. Hager and H. Zhang, __A new conjugate gradient method with guaranteed descent and an efficient line search__,
    > SIAM J. Optim, (16), pp. 170-192, 2005.
    > doi: [10.1137/030601880](https://doi.org/10.1137/030601880)
"""
mutable struct HagerZhangCoefficient{TVTM<:AbstractVectorTransportMethod} <:
               DirectionUpdateRule
    transport_method::TVTM
    storage::StoreOptionsAction
    function HagerZhangCoefficient(
        t::AbstractVectorTransportMethod=ParallelTransport(),
        a::StoreOptionsAction=StoreOptionsAction((:x, :gradient, :δ)),
    )
        return new{typeof(t)}(t, a)
    end
end
function (u::HagerZhangCoefficient)(
    p::P, o::O, i
) where {P<:GradientProblem,O<:ConjugateGradientDescentOptions}
    if !all(has_storage.(Ref(u.storage), [:x, :gradient, :δ]))
        update_storage!(u.storage, o) # if not given store current as old
        return 0.0
    end
    xOld, gradientOld, δOld = get_storage.(Ref(u.storage), [:x, :gradient, :δ])
    update_storage!(u.storage, o)

    gradienttr = vector_transport_to(p.M, xOld, gradientOld, o.x, u.transport_method)
    ν = o.gradient - gradienttr #notation y from [HZ06]
    δtr = vector_transport_to(p.M, xOld, δOld, o.x, u.transport_method)
    denom = inner(p.M, o.x, δtr, ν)
    νknormsq = inner(p.M, o.x, ν, ν)
    β =
        inner(p.M, o.x, ν, o.gradient) / denom -
        2 * νknormsq * inner(p.M, o.x, δtr, o.gradient) / denom^2
    # Numerical stability from Manopt / Hager-Zhang paper
    ξn = norm(p.M, o.x, o.gradient)
    η = -1 / (ξn * min(0.01, norm(p.M, xOld, gradientOld)))
    return max(β, η)
end

@doc raw"""
    HeestenesStiefelCoefficient <: DirectionUpdateRule

Computes an update coefficient for the conjugate gradient method, where
the [`ConjugateGradientDescentOptions`](@ref)` o` include the last iterates
``x_k,ξ_k``, the current iterates ``x_{k+1},ξ_{k+1}`` and the last update
direction ``\delta=\delta_k``, where the last three ones are stored in the
variables with prequel `Old` based on [^HeestensStiefel1952]

adapted to manifolds as follows: let ``\nu_k = ξ_{k+1} - P_{x_{k+1}\gets x_k}ξ_k``.
Then the update reads

````math
β_k = \frac{\langle ξ_{k+1}, \nu_k \rangle_{x_{k+1}} }
    { \langle P_{x_{k+1}\gets x_k} \delta_k, \nu_k\rangle_{x_{k+1}} },
````
where ``P_{a\gets b}(⋅)`` denotes a vector transport from the tangent space at ``a`` to ``b``.

# Constructor
    HeestenesStiefelCoefficient(
        t::AbstractVectorTransportMethod=ParallelTransport(),
        a::StoreOptionsAction=()
    )

Construct the Heestens Stiefel coefficient update rule, where the parallel transport is the
default vector transport and a new storage is created by default.

See also [`conjugate_gradient_descent`](@ref)

[^HeestensStiefel1952]:
    > M.R. Hestenes, E.L. Stiefel, Methods of conjugate gradients for solving linear systems,
    > J. Research Nat. Bur. Standards, 49 (1952), pp. 409–436.
    > doi: [10.6028/jres.049.044](http://dx.doi.org/10.6028/jres.049.044)
"""
mutable struct HeestenesStiefelCoefficient{TVTM<:AbstractVectorTransportMethod} <:
               DirectionUpdateRule
    transport_method::TVTM
    storage::StoreOptionsAction
    function HeestenesStiefelCoefficient(
        transport_method::AbstractVectorTransportMethod=ParallelTransport(),
        storage_action::StoreOptionsAction=StoreOptionsAction((:x, :gradient, :δ)),
    )
        return new{typeof(transport_method)}(transport_method, storage_action)
    end
end
function (u::HeestenesStiefelCoefficient)(
    p::GradientProblem, o::ConjugateGradientDescentOptions, i
)
    if !all(has_storage.(Ref(u.storage), [:x, :gradient, :δ]))
        update_storage!(u.storage, o) # if not given store current as old
        return 0.0
    end
    xOld, gradientOld, δOld = get_storage.(Ref(u.storage), [:x, :gradient, :δ])
    update_storage!(u.storage, o)
    gradienttr = vector_transport_to(p.M, xOld, gradientOld, o.x, u.transport_method)
    δtr = vector_transport_to(p.M, xOld, δOld, o.x, u.transport_method)
    ν = o.gradient - gradienttr #notation from [HZ06]
    β = inner(p.M, o.x, o.gradient, ν) / inner(p.M, o.x, δtr, ν)
    return max(0, β)
end

@doc raw"""
    LiuStoreyCoefficient <: DirectionUpdateRule

Computes an update coefficient for the conjugate gradient method, where
the [`ConjugateGradientDescentOptions`](@ref)` o` include the last iterates
``x_k,ξ_k``, the current iterates ``x_{k+1},ξ_{k+1}`` and the last update
direction ``\delta=\delta_k``, where the last three ones are stored in the
variables with prequel `Old` based on [^LuiStorey1991]
adapted to manifolds: let ``\nu_k = ξ_{k+1} - P_{x_{k+1}\gets x_k}ξ_k``,
where ``P_{a\gets b}(⋅)`` denotes a vector transport from the tangent space at ``a`` to ``b``.

Then the coefficient reads

````math
β_k = -
\frac{ \langle ξ_{k+1},\nu_k \rangle_{x_{k+1}} }
{\langle \delta_k,ξ_k \rangle_{x_k}}.
````

See also [`conjugate_gradient_descent`](@ref)

# Constructor
    LiuStoreyCoefficient(
        t::AbstractVectorTransportMethod=ParallelTransport(),
        a::StoreOptionsAction=()
    )

Construct the Lui Storey coefficient update rule, where the parallel transport is the
default vector transport and a new storage is created by default.

[^LuiStorey1991]:
    > [Y. Liu and C. Storey, Efficient generalized conjugate gradient algorithms, Part 1: Theory
    > J. Optim. Theory Appl., 69 (1991), pp. 129–137.
    > doi: [10.1007/BF00940464](https://doi.org/10.1007/BF00940464)
"""
mutable struct LiuStoreyCoefficient{TVTM<:AbstractVectorTransportMethod} <:
               DirectionUpdateRule
    transport_method::TVTM
    storage::StoreOptionsAction
    function LiuStoreyCoefficient(
        t::AbstractVectorTransportMethod=ParallelTransport(),
        a::StoreOptionsAction=StoreOptionsAction((:x, :gradient, :δ)),
    )
        return new{typeof(t)}(t, a)
    end
end
function (u::LiuStoreyCoefficient)(
    p::GradientProblem, o::ConjugateGradientDescentOptions, i
)
    if !all(has_storage.(Ref(u.storage), [:x, :gradient, :δ]))
        update_storage!(u.storage, o) # if not given store current as old
    end
    xOld, gradientOld, δOld = get_storage.(Ref(u.storage), [:x, :gradient, :δ])
    update_storage!(u.storage, o)
    gradienttr = vector_transport_to(p.M, xOld, gradientOld, o.x, u.transport_method)
    ν = o.gradient - gradienttr # notation y from [HZ06]
    return inner(p.M, o.x, o.gradient, ν) / inner(p.M, xOld, -δOld, gradientOld)
end

@doc raw"""
    PolakRibiereCoefficient <: DirectionUpdateRule

Computes an update coefficient for the conjugate gradient method, where
the [`ConjugateGradientDescentOptions`](@ref)` o` include the last iterates
``x_k,ξ_k``, the current iterates ``x_{k+1},ξ_{k+1}`` and the last update
direction ``\delta=\delta_k``, where the last three ones are stored in the
variables with prequel `Old` based on [^PolakRibiere1969][^Polyak1969]

adapted to manifolds: let ``\nu_k = ξ_{k+1} - P_{x_{k+1}\gets x_k}ξ_k``,
where ``P_{a\gets b}(⋅)`` denotes a vector transport from the tangent space at ``a`` to ``b``.

Then the update reads
````math
β_k =
\frac{ \langle ξ_{k+1}, \nu_k \rangle_{x_{k+1}} }
{\lVert ξ_k \rVert_{x_k}^2 }.
````

# Constructor
    PolakRibiereCoefficient(
        t::AbstractVectorTransportMethod=ParallelTransport(),
        a::StoreOptionsAction=()
    )

Construct the PolakRibiere coefficient update rule, where the parallel transport is the
default vector transport and a new storage is created by default.

See also [`conjugate_gradient_descent`](@ref)

[^PolakRibiere1969]:
    > E. Polak, G. Ribiere, Note sur la convergence de méthodes de directions conjuguées
    > ESAIM: Mathematical Modelling and Numerical Analysis - Modélisation Mathématique et Analyse Numérique, Tome 3 (1969) no. R1, p. 35-43,
    > url: [http://www.numdam.org/item/?id=M2AN_1969__3_1_35_0](http://www.numdam.org/item/?id=M2AN_1969__3_1_35_0)

[^Polyak1969]:
    > B. T. Polyak, The conjugate gradient method in extreme problems,
    > USSR Comp. Math. Math. Phys., 9 (1969), pp. 94–112.
    > doi: [10.1016/0041-5553(69)90035-4](https://doi.org/10.1016/0041-5553(69)90035-4)
"""
mutable struct PolakRibiereCoefficient{TVTM<:AbstractVectorTransportMethod} <:
               DirectionUpdateRule
    transport_method::TVTM
    storage::StoreOptionsAction
    function PolakRibiereCoefficient(
        t::AbstractVectorTransportMethod=ParallelTransport(),
        a::StoreOptionsAction=StoreOptionsAction((:x, :gradient)),
    )
        return new{typeof(t)}(t, a)
    end
end
function (u::PolakRibiereCoefficient)(
    p::GradientProblem, o::ConjugateGradientDescentOptions, i
)
    if !all(has_storage.(Ref(u.storage), [:x, :gradient]))
        update_storage!(u.storage, o) # if not given store current as old
    end
    xOld, gradientOld = get_storage.(Ref(u.storage), [:x, :gradient])
    update_storage!(u.storage, o)

    gradienttr = vector_transport_to(p.M, xOld, gradientOld, o.x, u.transport_method)
    ν = o.gradient - gradienttr
    β = inner(p.M, o.x, o.gradient, ν) / inner(p.M, xOld, gradientOld, gradientOld)
    return max(0, β)
end

@doc raw"""
    SteepestDirectionUpdateRule <: DirectionUpdateRule

The simplest rule to update is to have no influence of the last direction and
hence return an update ``β = 0`` for all [`ConjugateGradientDescentOptions`](@ref)` o`

See also [`conjugate_gradient_descent`](@ref)
"""
struct SteepestDirectionUpdateRule <: DirectionUpdateRule end
function (u::SteepestDirectionUpdateRule)(
    ::GradientProblem, ::ConjugateGradientDescentOptions, i
)
    return 0.0
end


text:
using Stheno
using Optim
using Zygote: gradient

using GPARatScale

export create_optim_gp_post, create_optim_gpar_post

"""
Script including utilities to optimize the hyperparameters of your GP.
"""

"""
Compute the optimised posterior for a EQ kenel GP.

input_locations: array of length L
outputs: array of length L

# Examples:
'''julia-repl
julia> create_optim_gp_post([0.5, 0.2], [-2, 2])

In this example we have the mappings:
0.5 -> -2
0.2 -> 2
"""
function create_optim_gp_post(
    input_locations,
    outputs;
    kernel_structure::Kernel = EQ(),
    # Initial log vals for hyperparameters which will be optimized
    i_log_l=nothing, i_log_process_var=nothing, i_log_noise_sigma=nothing,
    debug::Bool = false,
)
    # Helper function used to compute negative log marignal lik for params
    function nlml(params)
        l, process_var, noise_sigma = unpack_gp(params)
        curr_kernel = process_var^2 * stretch(kernel_structure, 1/l)
        # curr_kernel = Stheno.kernel(kernel_structure; l = l, s = process_var^2)
        f = GP(curr_kernel, GPC())
        result = -logpdf(f(input_locations, noise_sigma^2), outputs)

        return result
    end
    # Optimize the parameters
    params = parse_initial_gp_params(i_log_l, i_log_process_var, i_log_noise_sigma)

    if debug
        i_l, i_var, i_noise = unpack_gp(params)
        println("Generating GP with initial parameters:")
        println("\tl=$(i_l); var=$(i_var); noise=$(i_noise)")
    end
    results = Optim.optimize(nlml, x->gradient(nlml, x)[1], params, BFGS(); inplace=false)
    opt_l, opt_process_var, opt_noise_sigma = unpack_gp(results.minimizer)
    if debug
        println("Finished optimizing parameters:")
        println("\tOptimum L: $(opt_l) ")
        println("\tOptimum Process Variance: $(opt_process_var)")
        println("\tOptimum noise: $(opt_noise_sigma)")
        println()
    end
    gp_kernel = kernel(kernel_structure, l = opt_l, s = opt_process_var^2)
    gp = GP(gp_kernel, GPC())
    gp_post = gp | (gp(input_locations, opt_noise_sigma^2) ← outputs)

    return gp_post
end

"""
Compute the optimised posterior for a EQ kenel GPAR.

time_input: array of length L
prev_outputs: array of arrays of length L
outputs: array of length L

# Examples:
'''julia-repl
julia> create_optim_gpar_post([1, 2], [[0.5, 0.2], [4.2, 4.5]], [-2, 2])

In this example we have the mappings:
(1, 0.5, 4.2) -> -2
(2, 0.2, 4.5) -> 2
"""
function create_optim_gpar_post(
    input_locations,
    outputs;
    time_kernel::Kernel = EQ(),
    out_kernel::Kernel = EQ(),
    # Initial log param values which will be put in the optimization
    i_log_time_l=nothing, i_log_time_var=nothing, i_log_out_l=nothing,
    i_log_out_var=nothing, i_log_noise_sigma=nothing,
    multi_input::Bool = true,  # False if it is the first gpar kernel
    debug::Bool = false,
)
    if !multi_input  # same case as usual gp
        return create_optim_gp_post(
            input_locations,
            outputs,
            kernel_structure = time_kernel;
            i_log_l=i_log_time_l, i_log_process_var=i_log_time_var,
            i_log_noise_sigma=i_log_noise_sigma,
            debug = debug,
        )
    end
    # Otherwise create an optimised GPAR
    input_length = length(input_locations)
    input_locations = to_ColVecs(input_locations)  # transform into ColVecs
    # Helper function that strings together the GPAR kernel
    function create_gpar_kernel(time_l, time_var, out_l, out_var)
        time_mask = get_time_mask(input_length)
        masked_time_kernel = stretch(time_kernel, time_mask)
        scaled_time_kernel =
            kernel(masked_time_kernel, l = time_l, s = time_var^2)

        out_mask = get_output_mask(input_length)
        masked_out_kernel = stretch(out_kernel, out_mask)
        scaled_out_kernel = kernel(masked_out_kernel, l = out_l, s = out_var^2)

        final_kernel = scaled_time_kernel + scaled_out_kernel
        return final_kernel
    end

    # Helper function used to compute negative log marignal lik for params
    function nlml(params)
        time_l, time_var, out_l, out_var, noise_sigma = unpack_gpar(params)
        kernel = create_gpar_kernel(time_l, time_var, out_l, out_var)
        f = GP(kernel, GPC())

        return -logpdf(f(input_locations, noise_sigma^2), outputs)
    end
    # Optimize the parameters
    params = parse_initial_gpar_params(
        i_log_time_l, i_log_time_var, i_log_out_l,
        i_log_out_var, i_log_noise_sigma)
    if debug
        i_time_l, i_time_var, i_out_l, i_out_var, i_noise_sigma = unpack_gpar(params)
        println("Generating GPAR with initial parameters:")
        println("\ti_time_l=$(i_time_l); i_time_var=$(i_time_var); i_out_l=$(i_out_l); i_out_var=$(i_out_var); i_noise_sigma=$(i_noise_sigma)")
    end
    results = Optim.optimize(nlml, params, NelderMead())

    opt_time_l, opt_time_var, opt_out_l, opt_out_var, opt_noise_sigma =
        unpack_gpar(results.minimizer)
    if debug
        println("Finished optimizing parameters:")
        println("\tOptimum time L: $(opt_time_l) ")
        println("\tOptimum time var: $(opt_time_var)")
        println("\tOptimum outputs l: $(opt_out_l)")
        println("\tOptimum outputs var: $(opt_out_var)")
        println("\tOptimum Noise std: $(opt_noise_sigma)")
        println()
    end
    gpar_kernel =
        create_gpar_kernel(opt_time_l, opt_time_var, opt_out_l, opt_out_var)
    gpar = GP(gpar_kernel, GPC())
    gpar_post = gpar | (gpar(input_locations, opt_noise_sigma^2) ← outputs)

    return gpar_post
end

ext: jl

max_stars_repo_name: willtebbutt/GPAR-at-scale

max_issues_repo_name: willtebbutt/GPAR-at-scale

max_forks_repo_name: willtebbutt/GPAR-at-scale


text:
module StarBattle

using Parameters
using SimpleGraphs
using SimpleGraphAlgorithms

include("graph.jl")

export solve, graph_solve
export adj_matrix

end


text:
using StippleCharts
using Test

@testset "StippleCharts.jl" begin
    # Write your own tests here.
end


text:
############################################################################################
"""
$(SIGNATURES)
Flatten array x. Not exported.

# Examples
```julia
```

"""
@inline function flatten_array(mat::AbstractArray{R}) where {R<:Real}
    return vec(mat)
end

"""
$(SIGNATURES)
Fill array with elements of vec. Not exported.

# Examples
```julia
```

"""
@inline function fill_array!(
    buffer::AbstractArray{T}, vec::Union{F,AbstractArray{F}}
) where {T<:Real,F<:Real}
    @inbounds for iter in eachindex(vec)
        buffer[iter] = vec[iter]
    end
    return buffer
end
function ChainRulesCore.rrule(
    ::typeof(fill_array!), mat::AbstractMatrix{R}, v::Union{T,AbstractVector{T}}
) where {R<:Real,T<:Real}
    # forward pass: Fill Matrix with Vector elements
    L = fill_array!(mat, v)
    # backward pass: Fill Vector with Matrix elements
    function pullback_Idx(Δ)
        return ChainRulesCore.NoTangent(),
        ChainRulesCore.unthunk(Δ),
        flatten_array(ChainRulesCore.unthunk(Δ))
    end
    return L, pullback_Idx
end

############################################################################################
"""
$(SIGNATURES)
Flatten vector x to k-1 dimensions. Not exported.

# Examples
```julia
```

"""
function flatten_Simplex(x::AbstractVector{R}) where {R<:Real}
    return x[1:(end - 1)]
end

"""
$(SIGNATURES)
Expand vector of k-1 dimensions back to k dimensions. Not exported.

# Examples
```julia
```

"""
function Simplex_from_flatten(x_vec::Union{R,AbstractVector{R}}) where {R<:Real}
    buffer = zeros(eltype(x_vec), length(x_vec) + 1)
    @inbounds for iter in eachindex(x_vec)
        buffer[iter] = x_vec[iter]
    end
    buffer[end] = 1.0 - sum(x_vec)
    return buffer
end
function ChainRulesCore.rrule(
    ::typeof(Simplex_from_flatten), v::Union{R,AbstractVector{R}}
) where {R<:Real}
    # forward pass: From k-1 to k dimensions
    L = Simplex_from_flatten(v)
    # backward pass: From k to k-1 dimensions
    function pullback_Simplex(Δ)
        return ChainRulesCore.NoTangent(), flatten_Simplex(ChainRulesCore.unthunk(Δ))
    end
    return L, pullback_Simplex
end

"""
$(SIGNATURES)
Inplace version of Simplex_from_flatten. Not exported.

# Examples
```julia
```

"""
function Simplex_from_flatten!(
    buffer::AbstractVector{R}, x_vec::Union{T,AbstractVector{T}}
) where {R<:Real,T<:Real}
    @inbounds for iter in eachindex(x_vec)
        buffer[iter] = x_vec[iter]
    end
    buffer[end] = 1.0 - sum(x_vec)
    return buffer
end
function ChainRulesCore.rrule(
    ::typeof(Simplex_from_flatten!), p::AbstractVector{R}, v::Union{T,AbstractVector{T}}
) where {R<:Real,T<:Real}
    # forward pass: From k-1 to k dimensions
    L = Simplex_from_flatten!(p, v)
    # backward pass: From k to k-1 dimensions
    function pullback_Simplex(Δ)
        return ChainRulesCore.NoTangent(),
        ChainRulesCore.unthunk(Δ),
        flatten_Simplex(ChainRulesCore.unthunk(Δ))
    end
    return L, pullback_Simplex
end

############################################################################################
"""
$(SIGNATURES)
Flatten matrix to vector. Not exported.

# Examples
```julia
```

"""
@inline function flatten_Symmetric(mat::AbstractMatrix{R}, idx::BitMatrix) where {R<:Real}
    return vec(mat[idx])
end

"""
$(SIGNATURES)
Expand vector back to (symmetric) Matrix. Not exported.

# Examples
```julia
```

"""
@inline function Symmetric_from_flatten(
    x_vec::Union{R,AbstractVector{R}}, idx::BitMatrix
) where {R<:Real}
    ArgCheck.@argcheck size(idx, 1) == size(idx, 2) "'idx' has to be a square matrix."
    mat = zeros(eltype(x_vec), size(idx))
    counter = 0
    @inbounds @simd for Ncol in axes(mat, 2)
        for Nrow in axes(mat, 1)
            if idx[Nrow, Ncol]
                counter += 1
                mat[Nrow, Ncol] = mat[Ncol, Nrow] = x_vec[counter]
            end
        end
    end
    return mat
end
function ChainRulesCore.rrule(
    ::typeof(Symmetric_from_flatten), v::Union{R,AbstractVector{R}}, idx::BitMatrix
) where {R<:Real}
    # forward pass: Fill Matrix with Vector elements
    L = Symmetric_from_flatten(v, idx)
    # backward pass: Fill Vector with Matrix elements
    function pullback_Idx(Δ)
        return ChainRulesCore.NoTangent(),
        flatten_Symmetric(ChainRulesCore.unthunk(Δ), idx),
        ChainRulesCore.NoTangent()
    end
    return L, pullback_Idx
end

"""
$(SIGNATURES)
Inplace version of Symmetric_from_flatten. Not exported.

# Examples
```julia
```

"""
@inline function Symmetric_from_flatten!(
    mat::AbstractMatrix{T}, x_vec::Union{R,AbstractVector{R}}, idx::BitMatrix
) where {T<:Real,R<:Real}
    ArgCheck.@argcheck size(mat) == size(idx) "Dimension of matrix and entries do not match."
    counter = 0
    @inbounds @simd for Ncol in axes(mat, 2)
        for Nrow in axes(mat, 1)
            if idx[Nrow, Ncol]
                counter += 1
                mat[Nrow, Ncol] = mat[Ncol, Nrow] = x_vec[counter]
            end
        end
    end
    return mat
end
function ChainRulesCore.rrule(
    ::typeof(Symmetric_from_flatten!),
    mat::AbstractMatrix{T},
    v::Union{R,AbstractVector{R}},
    idx::BitMatrix,
) where {T<:Real,R<:Real}
    # forward pass: Fill Matrix with Vector elements
    L = Symmetric_from_flatten!(mat, v, idx)
    # backward pass: Fill Vector with Matrix elements
    function pullback_Idx(Δ)
        return ChainRulesCore.NoTangent(),
        ChainRulesCore.unthunk(Δ), flatten_Symmetric(ChainRulesCore.unthunk(Δ), idx),
        ChainRulesCore.NoTangent()
    end
    return L, pullback_Idx
end

############################################################################################
"""
$(SIGNATURES)
Assign subset of elements to track in Matrix mat. Not exported.

# Examples
```julia
```

"""
@inline function tag(
    mat::AbstractMatrix{R}, upper::Bool=true, diag::Bool=true
) where {R<:Real}
    ArgCheck.@argcheck size(mat, 1) == size(mat, 2) "Dimension of matrix and entries do not match."
    idx =
        upper ? LinearAlgebra.triu(trues(size(mat))) : LinearAlgebra.tril(trues(size(mat)))
    if !diag
        @inbounds for iter in axes(idx, 1)
            idx[iter, iter] = 0.0
        end
    end
    return idx
end

############################################################################################
# Export


text:
using SpecialFunctions
using OffsetArrays

export overlap, kinetic, nuclear_attraction

"""
overlap(a::PGBF,b::PGBF)

Compute the overlap between primitive Gaussian basis functions `a` and `b`.
"""
function overlap(a::PGBF,b::PGBF)
    return a.norm*b.norm*overlap(a.expn,a.xyz,a.I,a.J,a.K,
    b.expn,b.xyz,b.I,b.J,b.K)
end

"""
overlap(a::CGBF,b::CGBF)

Compute the overlap between contracted Gaussian basis functions `a` and `b`.
"""
overlap(a::CGBF,b::CGBF) = contract(overlap,a,b)

"""
overlap(aexpn,ax,ay,az,aI,aJ,aK,bexpn,bx,by,bz,bI,bJ,bK)

Compute the overlap between primitive Gaussian basis functions
defined by centers `ax,ay,az`, `bx,by,bz`, powers `aI,aJ,aK`
`bI,bJ,bK`, and exponents `aexpn,bexpn`.
"""
function overlap(aexpn,ax,ay,az,aI,aJ,aK,bexpn,bx,by,bz,bI,bJ,bK)
    return overlap(aexpn,[ax,ay,az],aI,aJ,aK,bexpn,[bx,by,bz],bI,bJ,bK)
end

"""
overlap(aexpn,axyz,aI,aJ,aK,bexpn,bxyz,bI,bJ,bK)

Compute the overlap between primitive Gaussian basis functions
defined by centers `axyz`, `bxyz`, powers `aI,aJ,aK`
`bI,bJ,bK`, and exponents `aexpn,bexpn`.
"""
function overlap(aexpn,axyz,aI,aJ,aK,bexpn,bxyz,bI,bJ,bK)
    gamma = aexpn+bexpn
    pxyz = gaussian_product_center(aexpn,axyz,bexpn,bxyz)
    pa = pxyz-axyz
    pb = pxyz-bxyz
    rab2 = dist2(axyz-bxyz)
    pre = (pi/gamma)^1.5*exp(-aexpn*bexpn*rab2/gamma)
    wx = overlap1d(aI,bI,pa[1],pb[1],gamma)
    wy = overlap1d(aJ,bJ,pa[2],pb[2],gamma)
    wz = overlap1d(aK,bK,pa[3],pb[3],gamma)
    return pre*wx*wy*wz
end

"""
gaussian_product_center(aexpn,ax,ay,az,bexpn,bx,by,bz)

Compute the Gaussian function center defined by the
product of Gaussian functions at  `ax,ay,az`, `bx,by,bz`,
and exponents `aexpn,bexpn`.
"""
gaussian_product_center(aexpn,ax,ay,az,bexpn,bx,by,bz) = gaussian_product_center(aexpn,[ax,ay,az],bexpn,[bx,by,bz])

"""
gaussian_product_center(a::PGBF,b::PGBF)

Compute the Gaussian function center defined by the
product of Gaussian functions `a` and `b`.
"""
gaussian_product_center(a::PGBF,b::PGBF) = gaussian_product_center(a.expn,a.xyz,b.expn,b.xyz)

"""
gaussian_product_center(aexpn,axyz,bexpn,bxyz)

Compute the Gaussian function center defined by the
product of Gaussian functions at  `axyz`, `bxyz`,
and exponents `aexpn,bexpn`.
"""
function gaussian_product_center(aexpn,axyz,bexpn,bxyz)
    ab = aexpn+bexpn
    a = aexpn/ab
    b = bexpn/ab
    return a*axyz+b*bxyz
end

function overlap1d(la,lb,ax,bx,gamma)
    total = 0
    for i in 0:div(la+lb,2)
        total += binomial_prefactor(2i,la,lb,ax,bx)*factorial2(2i-1)/(2gamma)^i
    end
    return total
end

binomial_prefactor(s,ia,ib,xpa,xpb) = sum(binomial_kernel(ia,s-t,xpa)*binomial_kernel(ib,t,xpb) for t in 0:s if (s-ia) <= t <= ib)
binomial_kernel(i,t,x) = binomial(i,t)x^(i-t)

"""
kinetic(a::PGBF,b::PGBF)

Compute the kinetic energy between primitive Gaussian basis functions `a` and `b`.
"""
function kinetic(a::PGBF,b::PGBF)
    return a.norm*b.norm*kinetic(a.expn,a.xyz,a.I,a.J,a.K,
                                b.expn,b.xyz,b.I,b.J,b.K)
end

"""
kinetic(aexpn,ax,ay,az,aI,aJ,aK,bexpn,bx,by,bz,bI,bJ,bK)

Compute the kinetic energy between primitive Gaussian basis functions
defined by centers `ax,ay,az`, `bx,by,bz`, powers `aI,aJ,aK`
`bI,bJ,bK`, and exponents `aexpn,bexpn`.
"""
function kinetic(aexpn,ax,ay,az,aI,aJ,aK,bexpn,bx,by,bz,bI,bJ,bK)
    return kinetic(aexpn,[ax,ay,az],aI,aJ,aK,bexpn,[bx,by,bz],bI,bJ,bK)
end

"""
kinetic(aexpn,axyz,aI,aJ,aK,bexpn,bxyz,bI,bJ,bK)

Compute the kinetic energy between primitive Gaussian basis functions
defined by centers `ax,yz`, `bxyz`, powers `aI,aJ,aK`
`bI,bJ,bK`, and exponents `aexpn,bexpn`.
"""
function kinetic(aexpn,axyz,aI,aJ,aK,bexpn,bxyz,bI,bJ,bK)
    overlap0 = overlap(aexpn,axyz,aI,aJ,aK,bexpn,bxyz,bI,bJ,bK)
    overlapx1 = overlap(aexpn,axyz,aI,aJ,aK,bexpn,bxyz,bI+2,bJ,bK)
    overlapy1 = overlap(aexpn,axyz,aI,aJ,aK,bexpn,bxyz,bI,bJ+2,bK)
    overlapz1 = overlap(aexpn,axyz,aI,aJ,aK,bexpn,bxyz,bI,bJ,bK+2)
    overlapx2 = overlap(aexpn,axyz,aI,aJ,aK,bexpn,bxyz,bI-2,bJ,bK)
    overlapy2 = overlap(aexpn,axyz,aI,aJ,aK,bexpn,bxyz,bI,bJ-2,bK)
    overlapz2 = overlap(aexpn,axyz,aI,aJ,aK,bexpn,bxyz,bI,bJ,bK-2)
    term0 = bexpn*(2*(bI+bJ+bK)+3)*overlap0
    term1 = -2*(bexpn^2)*(overlapx1+overlapy1+overlapz1)
    term2 = -0.5*(bI*(bI-1)*overlapx2+bJ*(bJ-1)*overlapy2+bK*(bK-1)*overlapz2)
    return term0+term1+term2
end

"""
kinetic(a::CGBF,b::CGBF)

Compute the kinetic energy between contracted Gaussian basis functions `a` and `b`.
"""
kinetic(a::CGBF,b::CGBF) = contract(kinetic,a,b)

function Aterm(i,r,u,l1,l2,ax,bx,cx,gamma)
    term1 = (-1)^i*binomial_prefactor(i,l1,l2,ax,bx)
    term2 = (-1)^u*factorial(i)*cx^(i-2r-2u)
    term3 = (1/4/gamma)^(r+u)/factorial(r)/factorial(u)/factorial(i-2r-2u)
    return term1*term2*term3
end

function Aarray(l1,l2,a,b,c,g)
    Imax = l1+l2+1
    A = OffsetArray(zeros(Float64,Imax),0:(Imax-1))
    for i in 0:(Imax-1)
        for r in 0:div(i,2)
            for u in 0:div(i-2r,2)
                I = i-2r-u
                A[I] += Aterm(i,r,u,l1,l2,a,b,c,g)
            end
        end
    end
    return A
end

"""
nuclear_attraction(a::PGBF,b::PGBF,cxyz)

Compute the nuclear attraction energy between primitive Gaussian basis functions `a` and `b`
and center `cxyz`.
"""
function nuclear_attraction(a::PGBF,b::PGBF,cxyz)
    return a.norm*b.norm*nuclear_attraction(a.expn,a.xyz,a.I,a.J,a.K,
                                            b.expn,b.xyz,b.I,b.J,b.K,cxyz)
end

"""
nuclear_attraction(a::PGBF,b::PGBF,c::Atom)

Compute the nuclear attraction energy between primitive Gaussian basis functions `a` and `b`
and atom `c`.
"""
nuclear_attraction(a::PGBF,b::PGBF,c::Atom) = c.atno*nuclear_attraction(a,b,c.xyz)

"""
nuclear_attraction(a::PGBF,b::PGBF,m::Vector{Atom})

Compute the sum of nuclear attraction energy between primitive Gaussian basis
functions `a` and `b` and the vector of atoms in `m`.
"""
nuclear_attraction(a::PGBF,b::PGBF,m::Vector{Atom}) = sum([nuclear_attraction(a,b,c) for c in m])

"""
nuclear_attraction(aexpn,ax,ay,az,aI,aJ,aK,bexpn,bx,by,bz,bI,bJ,bK,cx,cy,cz)

Compute the nuclear attraction energy between primitive Gaussian basis functions
defined by `ax,ay,az`, `bx,by,bz`, powers `aI,aJ,aK`, `bI,bJ,bK`, and exponents `aexpn,bexpn`
and center `cxyz`.
"""
function nuclear_attraction(aexpn,ax,ay,az,aI,aJ,aK,bexpn,bx,by,bz,bI,bJ,bK,cx,cy,cz)
    return nuclear_attraction(aexpn,[ax,ay,az],aI,aJ,aK,bexpn,[bx,by,bz],bI,bJ,bK,[cx,cy,cz])
end

"""
nuclear_attraction(aexpn,ax,ay,az,aI,aJ,aK,bexpn,bx,by,bz,bI,bJ,bK,cx,cy,cz)

Compute the nuclear attraction energy between primitive Gaussian basis functions
defined by `ax,ay,az`, `bx,by,bz`, powers `aI,aJ,aK`, `bI,bJ,bK`, and exponents `aexpn,bexpn`
and center `cxyz`.
"""
function nuclear_attraction(aexpn,axyz,aI,aJ,aK,bexpn,bxyz,bI,bJ,bK,cxyz)
    pxyz = gaussian_product_center(aexpn,axyz,bexpn,bxyz)
    gamma = aexpn+bexpn
    pa = pxyz-axyz
    pb = pxyz-bxyz
    pc = pxyz-cxyz
    rab2 = dist2(axyz-bxyz)
    rcp2 = dist2(pc)
    Ax = Aarray(aI,bI,pa[1],pb[1],pc[1],gamma)
    Ay = Aarray(aJ,bJ,pa[2],pb[2],pc[2],gamma)
    Az = Aarray(aK,bK,pa[3],pb[3],pc[3],gamma)
    total = 0
    for I in 0:(aI+bI)
        for J in 0:(aJ+bJ)
            for K in 0:(aK+bK)
                total += Ax[I]*Ay[J]*Az[K]*Fgamma(I+J+K,rcp2*gamma)
            end
        end
    end
    return -2pi*exp(-aexpn*bexpn*rab2/gamma)*total/gamma
end

"""
nuclear_attraction(a::CGBF,b::CGBF,cxyz)

Compute the nuclear attraction energy between contracted Gaussian
basis functions `a` and `b` and center `cxyz`.
"""
function nuclear_attraction(a::CGBF,b::CGBF,cxyz)
    na(a,b) = nuclear_attraction(a,b,cxyz)
    contract(na,a,b)
end

"""
nuclear_attraction(a::CGBF,b::CGBF,c::Atom)

Compute the nuclear attraction energy between contracted Gaussian
basis functions `a` and `b` and atom `c`.
"""
function nuclear_attraction(a::CGBF,b::CGBF,c::Atom)
    na(a,b) = nuclear_attraction(a,b,c)
    contract(na,a,b)
end

"""
nuclear_attraction(a::CGBF,b::CGBF,m::Vector{Atom})

Compute the sum of nuclear attraction energy between contracted Gaussian
basis functions `a` and `b` and vector of atoms in `m`.
"""
function nuclear_attraction(a::CGBF,b::CGBF,m::Vector{Atom})
    na(a,b) = nuclear_attraction(a,b,m)
    contract(na,a,b)
end

"""
all_1e_ints(bfs::Vector{CGBF},mol::Vector{Atom})

Return the overlap, the kinetic energy, and the nuclear attraction
integrals between vector of contracted functions `bfs` and
atoms `mol`.
"""
function all_1e_ints(bfs::Vector{CGBF},mol::Vector{Atom})
    n = length(bfs)
    S = Array{Float64}(n,n)
    T = Array{Float64}(n,n)
    V = Array{Float64}(n,n)
    for (i,j) in pairs(n)
        a,b = bfs[i],bfs[j]
        S[i,j] = S[j,i] = overlap(a,b)
        T[i,j] = T[j,i] = kinetic(a,b)
        V[i,j] = V[j,i] = nuclear_attraction(a,b,mol)
    end
    return S,T,V
end


text:
module TimeStepping

using DataStructures
using JuMP
using Mosek
using MosekTools
using Ipopt
using LinearAlgebra
using Revise
import MathOptInterface

include("moreau.jl")
include("integrator.jl")

export Moreau, step, set_state
export Integrator, integrate

end


text:
"""
    iterative!(rho0, H, J, [method!], args...; [log=false], kwargs...) -> rho[, log]

Compute the steady state density matrix of a Hamiltonian and a set of jump
operators by solving `L rho = 0` via an iterative method provided as argument.

# Arguments
* `rho0`: Initial density matrix.
* `H`: Non-lazy operator or arbitrary matrix specifying the Hamiltonian.
* `J`: Vector containing all jump operators which can be non-lazy operators types or matrices.
* `method!::Function`: The iterative method to be used. Defaults to `IterativeSolvers.bicgstabl!` or `IterativeSolvers.idrs!` depending on arguments' types.
* `rates=nothing`: Vector or matrix specifying the coefficients (decay rates) for the jump operators. If nothing is specified all rates are assumed to be 1.
* `Jdagger=dagger.(J)`: Vector containing the hermitian conjugates of the jump operators. If they are not given they are calculated automatically.
* `args...`: Further arguments are passed on to the iterative solver.
* `kwargs...`: Further keyword arguments are passed on to the iterative solver.

See also: [`iterative`](@ref)
"""
function iterative!(rho0::AbstractOperator{B,B}, H::AbstractOperator{B,B}, J::Vector, method!::Union{Function,Missing}=missing, args...;
                    rates::DecayRates=nothing, Jdagger::TJd=nothing, kwargs...) where {B<:Basis,TJd<:Union{Vector,Nothing}}
    _check_jump_ops(B,rates,J,Jdagger)

    sol = if Jdagger===nothing
        iterative!(rho0.data,H.data,map(x->x.data,J),method!,args...;rates=rates,Jdagger=nothing,kwargs...)
    else
        T = typeof(first(J).data)
        iterative!(rho0.data,H.data,map(x->x.data,J),method!,args...;rates=rates,Jdagger=map(x->T(x.data),Jdagger),kwargs...)
    end
    if typeof(sol) <: Tuple
        rho = deepcopy(H)
        rho.data .= sol[1]
        return rho, sol[2]
    else
        rho = deepcopy(H)
        rho.data .= sol
        return rho
    end
end

"""
    iterative(H, J, [method!], args...; [log=false], kwargs...) -> rho[, log]

Compute the steady state density matrix of a Hamiltonian and a set of jump
operators by solving `L rho = 0` via an iterative method provided as argument.

# Arguments
* `H`: Non-lazy operator or arbitrary matrix specifying the Hamiltonian.
* `J`: Vector containing all jump operators which can be non-lazy operators types or matrices.
* `method!::Function`: The iterative method to be used. Defaults to `IterativeSolvers.bicgstabl!` or `IterativeSolvers.idrs!` depending on arguments' types.
* `rates=nothing`: Vector or matrix specifying the coefficients (decay rates) for the jump operators. If nothing is specified all rates are assumed to be 1.
* `Jdagger=dagger.(J)`: Vector containing the hermitian conjugates of the jump operators. If they are not given they are calculated automatically.
* `args...`: Further arguments are passed on to the iterative solver.
* `kwargs...`: Further keyword arguments are passed on to the iterative solver.

See also: [`iterative`](@ref)
"""
function iterative(H::AbstractOperator{B,B}, J::Vector, method!::Union{Function,Missing}=missing, args...;
                   rates::DecayRates=nothing, Jdagger::TJd=nothing, kwargs...) where {B<:Basis,TJd<:Union{Vector,Nothing}}
    _check_jump_ops(B,rates,J,Jdagger)

    sol = if Jdagger===nothing
        iterative(H.data,map(x->x.data,J),method!,args...;rates=rates,Jdagger=nothing,kwargs...)
    else
        iterative(H.data,map(x->x.data,J),method!,args...;rates=rates,Jdagger=map(x->x.data,Jdagger),kwargs...)
    end
    if typeof(sol) <: Tuple
        rho = deepcopy(H)
        rho.data .= sol[1]
        return rho, sol[2]
    else
        rho = deepcopy(H)
        rho.data .= sol
        return rho
    end
end

iterative(rho0::AbstractOperator{B,B}, H::AbstractOperator{B,B}, J::Vector, method!::Union{Function,Missing}=missing, args...; rates::DecayRates=nothing, Jdagger::TJd=nothing, kwargs...) where {B<:Basis,TJd<:Union{Vector,Nothing}} = iterative!(copy(rho0),H,J,method!,args...;rates=rates,Jdagger=Jdagger,kwargs...)
iterative(psi0::Ket{B}, H::AbstractOperator{B,B}, J::Vector, method!::Union{Function,Missing}=missing, args...; rates::DecayRates=nothing, Jdagger::TJd=nothing, kwargs...) where {B<:Basis,TJd<:Union{Vector,Nothing}} = iterative!(dm(psi0),H,J,method!,args...;rates=rates,Jdagger=Jdagger,kwargs...)

function iterative!(rho0::AbstractMatrix{T1}, H::AbstractMatrix{T2}, J::Vector{TJ}, method!::Union{Function,Missing}=missing, args...;
                   rates::DecayRates=nothing, Jdagger::VJd=nothing, kwargs...) where {T1<:Number,T2<:Number,TJ<:AbstractMatrix,TJd<:AbstractMatrix,VJd<:Union{Vector{TJd},Nothing}}
    _check_jump_ops(rates,J,Jdagger)
    if ismissing(method!)
        if isblascompatible(rho0) && isblascompatible(H) && all(isblascompatible.(J))
            return steadystate_iterative!(rho0,H,rates,J,Jdagger,bicgstabl!,args...;kwargs...)
        else
            return steadystate_iterative!(rho0,H,rates,J,Jdagger,idrs!,args...;kwargs...)
        end
    else
        @assert eltype(J) <: AbstractMatrix{T} where {T<:Number} "The jump operators must be matrices of numbers."
        return steadystate_iterative!(rho0,H,rates,J,Jdagger,method!,args...;kwargs...)
    end
end

function iterative(H::AbstractMatrix{T}, J::Vector, method!::Union{Function,Missing}=missing, args...; rates::DecayRates=nothing, Jdagger::TJd=nothing, kwargs...) where {T<:Number,TJd<:Union{Vector,Nothing}}
    rho0 = similar(H)
    rho0 .= zero(T)
    rho0[1,1] = one(T)
    return iterative!(rho0,H,J,method!,args...;rates=rates,Jdagger=Jdagger,kwargs...)
end

function iterative(rho0::AbstractMatrix{T1}, H::AbstractMatrix{T2}, J::Vector{TJ}, method!::Union{Function,Missing}=missing, args...;
                   rates::DecayRates=nothing, Jdagger::TJd=nothing, kwargs...) where {T1<:Number,T2<:Number,TJ<:AbstractMatrix,TJd<:Union{Vector{TJ},Nothing}}
    return iterative!(copy(rho0),H,J,method!,args...;rates=rates,Jdagger=Jdagger,kwargs...)
end

function _check_jump_ops(B::DataType, rates::DecayRates, J::Vector, Jdagger::Union{Vector,Nothing})
    @assert all(typeof(j) <: AbstractOperator{B,B} for j in J) "Jump operators have incompatible bases."
    if Jdagger!==nothing
        @assert all(typeof(j) <: AbstractOperator{B,B} for j in Jdagger) "Adjoint jump operators have incompatible bases."
    end
end

function _check_jump_ops(rates::DecayRates, J::Vector, Jdagger::Union{Vector,Nothing})
    @assert all(typeof(j) <: AbstractMatrix for j in J) "Jump operators must all be of the same type."
    if Jdagger!==nothing
        @assert all(typeof(j)==typeof(first(J)) for j in Jdagger) "Adjoint jump operators must all be of the same type."
    end
    if rates!==nothing
        if typeof(rates) <: Vector
            @assert length(rates) == length(J) "The number of rates must match that of jump operators."
            if Jdagger!==nothing
                @assert length(J) == length(Jdagger) "The number of jump operators must match that of adjoint jump operators."
            end
        else # rates <: Matrix
            @assert size(rates,1) == length(J) "The number of jump operators must be compatible with the rates."
            if Jdagger!==nothing
                @assert size(rates,2) == length(Jdagger) "The number of adjoint jump operators must be compatible with the rates."
            end
        end
    else # rates = nothing
        if Jdagger!==nothing
            @assert length(J) == length(Jdagger) "The number of jump operators must match that of adjoint jump operators."
        end
    end
end


text:
using GeometricIntegrators.Integrators
using GeometricIntegrators.Tableaus
using GeometricIntegrators.TestProblems.HarmonicOscillatorProblem
using GeometricIntegrators.Utils
using Test

using GeometricIntegrators.TestProblems.HarmonicOscillatorProblem: Δt, nt, refx

sode = harmonic_oscillator_sode()

sint = Integrator(sode, getTableauLieA(), Δt)
ssol = integrate(sint, nt)

@test rel_err(ssol.q, refx) < 5E-2

sint = Integrator(sode, getTableauLieB(), Δt)
ssol = integrate(sint, nt)

@test rel_err(ssol.q, refx) < 5E-2

sint = Integrator(sode, getTableauStrang(), Δt)
ssol = integrate(sint, nt)

@test rel_err(ssol.q, refx) < 1E-3

sint = Integrator(sode, getTableauMcLachlan2(), Δt)
ssol = integrate(sint, nt)

@test rel_err(ssol.q, refx) < 1E-4

sint = Integrator(sode, getTableauMcLachlan4(), Δt)
ssol = integrate(sint, nt)

@test rel_err(ssol.q, refx) < 5E-4

sint = Integrator(sode, getTableauTripleJump(), Δt)
ssol = integrate(sint, nt)

@test rel_err(ssol.q, refx) < 5E-6

sint = Integrator(sode, getTableauSuzukiFractal(), Δt)
ssol = integrate(sint, nt)

@test rel_err(ssol.q, refx) < 5E-7


text:
###########################
# FunctionCell Interface  #
###########################
"""
	filterdomain!(opts::Dict)

Returns a cell that processes data piped to it according to the domains specified in `opts`.
The format of `opts` is `Dict(idx=>domain)` where:

  * `idx` can be an `Int`, `Vector{Int}` or `UnitRange` and specifies the variable indices
  * `domain` can be a vector or admissible values, matrix specifying intervals with the
first column specifying the lower bounds and second column the upper bounds, a single value
representing the only admissible value or function in which case the function is applied
to each element of the variable.

Filtered values are replaced by `NaN`,`""` and `:()` depending on the type of the data being
filtered.

# Examples
```
julia> a=[1,2,3,4,5.0]
5-element Array{Float64,1}:
 1.0
 2.0
 3.0
 4.0
 5.0

julia> a |> filterdomain!(Dict(1=>[1.0 2; 4 6.0]))
5-element Array{Float64,1}:
 1.0
 NaN
 NaN
 4.0
 5.0

julia> a = rand(["a","b","c"],3,3)
3×3 Array{String,2}:
 "c"  "b"  "a"
 "b"  "a"  "c"
 "a"  "c"  "c"

julia> a |> filterdomain!(Dict(1=>["a","b"], 2=>"a", 3=>x->x*"X"))
3×3 Array{String,2}:
 ""    "b"   "a"
 ""    "a"   ""
 "aX"  "cX"  "cX"
```
"""
filterdomain!(opts::T where T<:Dict) = FunctionCell(filterdomain!, (opts,), "Domain filter!")

############################
# DataCell/Array Interface #
############################
"""
	filterdomain!(data, opts::Dict)

Processes `data` according to the domains specified in `opts`.
"""
filterdomain!(x::T where T<:CellData, opts::U where U<:Dict) = datacell(filterdomain!(strip(x),opts))

filterdomain!(x::Tuple{T,S} where T<:AbstractArray where S<:AbstractArray, opts::U where U<:Dict) = (filterdomain!(x[1],opts),x[2])

filterdomain!(x::T where T<:AbstractArray, opts::U where U<:Dict) = begin

	# Construct out-of-bounds values
	_obvals_(x::Number) = NaN
	_obvals_(::AbstractFloat) = NaN
	_obvals_(::AbstractString) = ""
	_obvals_(::Symbol) = :()

	# Construct methods for each domain processing option
	_filter_(v::Function) = v
	_filter_(v::AbstractVector) = (x)->(x in v) ? x : _obvals_(x)::typeof(x)
 	_filter_(v::AbstractMatrix) =
		(x)->begin
			any( (x >= v[i,1])&(x < v[i,2]) for i in 1:size(v,1) ) && return x
			_obvals_(x)::typeof(x)
		end
	_filter_(v::T) where T<:DataElement = (x)->(x == v) ? x : _obvals_(x)

	# Apply the filters to the each variable
	for (idx, method) in opts
		variable = _variable_(x,idx) 			# Assign temp variable

		variable[:] = _filter_(method).(variable)	# Apply domain filter
	end

	return x
end


text:
# delegation of samplers

for fname in ["categorical.jl",
              "binomial.jl",
              "poisson.jl",
              "exponential.jl",
              "gamma.jl",
              "multinomial.jl",
              "vonmises.jl",
              "vonmisesfisher.jl"]

    include(joinpath("samplers", fname))
end


text:
include("common.jl")

facts("Optimize single objective problems in 5, 10, 30 and 100 dimensions with sNES") do
  simple_problems = ["Sphere", "Schwefel2.22", "Schwefel2.22"]
  for(problem in simple_problems)
    context(problem) do
      p = BlackBoxOptim.Problems.examples[problem]

      @fact fitness_for_opt(p, 5, 20,  1e2, separable_nes) < 0.01 => true
      @fact fitness_for_opt(p, 5, 20,  5e2, xnes) < 0.01 => true

      @fact fitness_for_opt(p, 10, 20, 3e2, separable_nes) < 0.01 => true
      @fact fitness_for_opt(p, 10, 20, 1e3, xnes) < 0.01 => true

      @fact fitness_for_opt(p, 30, 25, 1e3, separable_nes) < 0.01 => true
      @fact fitness_for_opt(p, 30, 25, 4e3, xnes) < 0.01 => true

      @fact fitness_for_opt(p, 100, 25, 5e3, separable_nes) < 0.01 => true
      # Cannot run xnes in 100 dimensions since it scales badly.
    end
  end

  context("Schwefel1.2") do
    problem = "Schwefel1.2"
    p = BlackBoxOptim.Problems.examples[problem]

    @fact fitness_for_opt(p, 30, 50, 4e3, separable_nes) < 10.0 => true
    @fact fitness_for_opt(p, 30, 50, 8e3, xnes) < 10.0 => true
  end

  context("Rosenbrock") do
    problem = "Rosenbrock"
    p = BlackBoxOptim.Problems.examples[problem]

    @fact fitness_for_opt(p, 30, 40, 2e4, separable_nes) < 100.0 => true
    @fact fitness_for_opt(p, 30, 40, 5e4, xnes) < 100.0 => true

    @fact fitness_for_opt(p, 50, 40, 3e4, separable_nes) < 100.0 => true
  end

end


text:
#=

@author : Spencer Lyon <spencer.lyon@nyu.edu>

@date: 2014-06-27

References
----------

Simple port of the file quantecon.models.jv

http://quant-econ.net/jl/jv.html
=#

using Distributions
using QuantEcon

# NOTE: only brute-force approach is available in bellman operator. Waiting on a simple constrained optimizer to be written in pure Julia

"""
A Jovanovic-type model of employment with on-the-job search.

The value function is given by

\\[V(x) = \\max_{\\phi, s} w(x, \\phi, s)\\]

for

    w(x, phi, s) := x(1 - phi - s) + beta (1 - pi(s)) V(G(x, phi)) +
                    beta pi(s) E V[ max(G(x, phi), U)

where

* `x`: : human capital
* `s` : search effort
* `phi` : investment in human capital
* `pi(s)` : probability of new offer given search level s
* `x(1 - \phi - s)` : wage
* `G(x, \phi)` : new human capital when current job retained
* `U` : Random variable with distribution F -- new draw of human capital

##### Fields

- `A::Real` : Parameter in human capital transition function
- `alpha::Real` : Parameter in human capital transition function
- `bet::Real` : Discount factor in (0, 1)
- `x_grid::Vector` : Grid for potential levels of x
- `G::Function` : Transition `function` for human captial
- `pi_func::Function` : `function` mapping search effort to the probability of getting a new job offer
- `F::UnivariateDistribution` : A univariate distribution from which the value of new job offers is drawn
- `quad_nodes::Vector` : Quadrature nodes for integrating over phi
- `quad_weights::Vector` : Quadrature weights for integrating over phi
- `epsilon::Float64` : A small number, used in optimization routine

"""
type JvWorker
    A::Real
    alpha::Real
    bet::Real
    x_grid::Vector{Float64}
    G::Function
    pi_func::Function
    F::UnivariateDistribution
    quad_nodes::Vector
    quad_weights::Vector
    epsilon::Float64
end

"""
Constructor with default values for `JvWorker`

##### Arguments

 - `A::Real(1.4)` : Parameter in human capital transition function
 - `alpha::Real(0.6)` : Parameter in human capital transition function
 - `bet::Real(0.96)` : Discount factor in (0, 1)
 - `grid_size::Int(50)` : Number of points in discrete grid for `x`
 - `epsilon::Float(1e-4)` : A small number, used in optimization routine

##### Notes

There is also a version of this function that accepts keyword arguments for
each parameter

"""
function JvWorker(A=1.4, alpha=0.6, bet=0.96, grid_size=50, epsilon=1e-4)
    G(x, phi) = A .* (x .* phi).^alpha
    pi_func = sqrt
    F = Beta(2, 2)

    # integration bounds
    a, b, = quantile(F, 0.005), quantile(F, 0.995)

    # quadrature nodes/weights
    nodes, weights = qnwlege(21, a, b)

    # Set up grid over the state space for DP
    # Max of grid is the max of a large quantile value for F and the
    # fixed point y = G(y, 1).
    grid_max = max(A^(1.0 / (1.0 - alpha)), quantile(F, 1 - epsilon))

    # range for linspace(epsilon, grid_max, grid_size). Needed for
    # CoordInterpGrid below
    x_grid = collect(linspace(epsilon, grid_max, grid_size))

    JvWorker(A, alpha, bet, x_grid, G, pi_func, F, nodes, weights, epsilon)
end

# make kwarg version
JvWorker(;A=1.4, alpha=0.6, bet=0.96, grid_size=50, epsilon=1e-4) = JvWorker(A, alpha, bet, grid_size, epsilon)

"""
Apply the Bellman operator for a given model and initial value, returning only the value function

##### Arguments

- `jv::JvWorker` : Instance of `JvWorker`
- `V::Vector`: Current guess for the value function
- `new_V::Vector` : Storage for updated value function

##### Returns

None, `new_V` is updated in place with the value function.

##### Notes

Currently, only the brute-force approach is available. We are waiting on a simple constrained optimizer to be written in pure Julia

"""
function bellman_operator!(jv::JvWorker, V::Vector, new_V::Vector)

    # simplify notation
    G, pi_func, F, bet, epsilon = jv.G, jv.pi_func, jv.F, jv.bet, jv.epsilon
    nodes, weights = jv.quad_nodes, jv.quad_weights

    # prepare interpoland of value function
    Vf = LinInterp(jv.x_grid, V)

    # instantiate the linesearch variables
    max_val = -1.0
    cur_val = 0.0
    max_s = 1.0
    max_phi = 1.0
    search_grid = linspace(epsilon, 1.0, 15)

    for (i, x) in enumerate(jv.x_grid)

        function w(z)
            s, phi = z
            function h(u)
              out = similar(u)
              for j in 1:length(u)
                out[j] = Vf(max(G(x, phi), u[j])) * pdf(F, u[j])
              end
              out
            end
            integral = do_quad(h, nodes, weights)
            q = pi_func(s) * integral + (1.0 - pi_func(s)) * Vf(G(x, phi))

            return - x * (1.0 - phi - s) - bet * q
        end

        for s in search_grid
            for phi in search_grid
                if s + phi <= 1.0
                    cur_val = -w((s, phi))
                else
                    cur_val = -1.0
                end
                if cur_val > max_val
                    max_val, max_s, max_phi = cur_val, s, phi
                end
            end
        end

        new_V[i] = max_val
    end
end

"""
Apply the Bellman operator for a given model and initial value, returning policies

##### Arguments

- `jv::JvWorker` : Instance of `JvWorker`
- `V::Vector`: Current guess for the value function
- `out::Tuple{Vector, Vector}` : Storage for the two policy rules

##### Returns

None, `out` is updated in place with the two policy functions.

##### Notes

Currently, only the brute-force approach is available. We are waiting on a simple constrained optimizer to be written in pure Julia

"""
function bellman_operator!(jv::JvWorker, V::Vector, out::Tuple{Vector, Vector})

    # simplify notation
    G, pi_func, F, bet, epsilon = jv.G, jv.pi_func, jv.F, jv.bet, jv.epsilon
    nodes, weights = jv.quad_nodes, jv.quad_weights

    # prepare interpoland of value function
    Vf = LinInterp(jv.x_grid, V)

    # instantiate variables
    s_policy, phi_policy = out[1], out[2]

    # instantiate the linesearch variables
    max_val = -1.0
    cur_val = 0.0
    max_s = 1.0
    max_phi = 1.0
    search_grid = linspace(epsilon, 1.0, 15)

    for (i, x) in enumerate(jv.x_grid)

        function w(z)
            s, phi = z
            function h(u)
              out = similar(u)
              for j in 1:length(u)
                out[j] = Vf(max(G(x, phi), u[j])) * pdf(F, u[j])
              end
              out
            end
            integral = do_quad(h, nodes, weights)
            q = pi_func(s) * integral + (1.0 - pi_func(s)) * Vf(G(x, phi))

            return - x * (1.0 - phi - s) - bet * q
        end

        for s in search_grid
            for phi in search_grid
                if s + phi <= 1.0
                    cur_val = -w((s, phi))
                else
                    cur_val = -1.0
                end
                if cur_val > max_val
                    max_val, max_s, max_phi = cur_val, s, phi
                end
            end
        end

      s_policy[i], phi_policy[i] = max_s, max_phi
  end
end

function bellman_operator(jv::JvWorker, V::Vector; ret_policies=false)
    if ret_policies
        out = (similar(V), similar(V))
    else
        out = similar(V)
    end
    bellman_operator!(jv, V, out)
    return out
end


text:
using Distributions
export MarkovChain, DirichletRows
export draw_sample, draw_sample!
export get_transition_matrix

struct MarkovChain
    P::Matrix{Float64}
    P_c::Matrix{Float64}
    mu::Categorical
end
MarkovChain(P::Matrix{Float64}, x::Int) = MarkovChain(
    P,
    cumsum(P, dims=2),
    Categorical([i==x ? 1.0 : 0.0 for i in 1:size(P,1)])
)

MarkovChain(P::Matrix{Float64}, μ::Categorical) = MarkovChain(
    P,
    cumsum(P, dims=2),
    μ
)

MarkovChain(P::Matrix{Float64}) = MarkovChain(
    P,
    cumsum(P, dims=2),
    Categorical([1/size(P,1) for i in 1:size(P,1)])
)

function rand_categorical(μ_cusum::AbstractArray{Float64})
    β = rand()
    for i in 1:length(μ_cusum)
        if β < μ_cusum[i]
            if i == 1
                return i
            else
                return i
            end
        else
            continue
        end
    end
end

function rand_categorical_with_prob(μ_cusum::AbstractArray{Float64})
    β = rand()
    for i in 1:length(μ_cusum)
        if β < μ_cusum[i]
            if i == 1
                return i, μ_cusum[i]
            else
                return i, μ_cusum[i] - μ_cusum[i-1]
            end
        else
            continue
        end
    end
end

function draw_sample!(
    out::Vector{Int},
    model::MarkovChain,
    n::Int)

    # Sample initial value
    curr = rand(model.mu)
    out[1] = curr
    for i in 2:n
        @views curr = rand_categorical(model.P_c[curr,:])
        out[i] = curr
    end
end

function draw_sample(
    model::MarkovChain,
    n::Int)

    out = zeros(Int, n)
    draw_sample!(out, model, n)
    return out
end

struct DirichletRows
    A::Matrix{Float64}
end

function draw_sample(model::DirichletRows)
    P = [rand(Gamma(a)) for a in model.A]
    P ./= sum(P, dims=2)
    return P
end

function draw_sample!(
    out::Vector{Matrix{Float64}},
    model::DirichletRows
    )

    for P in out
        for j in 1:size(P,2)
            for i in 1:size(P,1)
                P[i,j] = rand(Gamma(model.A[i,j]))
            end
        end
        P ./= sum(P, dims=2)
    end
end

function draw_sample(
    model::DirichletRows,
    n::Int
    )

    out = [similar(model.A) for i in 1:n]
    draw_sample!(out,model)
    return out
end

function get_transition_matrix(
    data::Vector{Int},
    V::Int
    )

    T = zeros(Int, V, V)
    @inbounds for i in Iterators.rest(eachindex(data),1)
        T[data[i-1], data[i]] += 1
    end
    return T
end

function get_transition_matrix(
    data::Vector{Int}
    )

    V = maximum(data)
    return get_transition_matrix(data, V)
end


text:
import ONNX: from_nnlib, from_onnx

@testset "Conversions" begin

    for n=1:4
        x = rand([5 for i=1:n]...)
        xo = from_nnlib(x)
        # test conversion forward and back
        @test from_onnx(xo) == x

        # test elements in x[k1, k2, ..., kn] are equal to xo[kn, ..., k2, k1]
        idx = ntuple(i -> rand(1:size(x, i)), n)
        @test x[idx...] == xo[reverse(idx)...]
    end
end


text:
# This file is a part of GERDADeepLearning.jl, licensed under the MIT License (MIT).

using HDF5, Compat, LsqFit
import Base: filter, length, getindex, haskey, copy, string, print, deepcopy

export EventCollection
@compat abstract type EventCollection
end

type NoSuchEventException <: Exception
end

export EventLibrary
type EventLibrary <: EventCollection
  # EventLibraries can be lazily initialized.
  # Possible initializers: load from file, filter by label
  initialization_function::Union{Function,Void}

  waveforms::Array{Float32, 2}
  labels::Dict{Symbol,Vector} # same length as waveforms
  prop::Dict{Symbol,Any}

  EventLibrary(init::Function) = new(init, zeros(Float32, 0,0), Dict{Symbol,Vector}(), Dict{Symbol,Any}())

  EventLibrary(waveforms::Matrix{Float64}) = EventLibrary(convert(Matrix{Float32}, waveforms))

  EventLibrary(waveforms::Matrix{Float32}) = new(nothing, waveforms, Dict{Symbol,Vector}(), Dict{Symbol,Any}())
end

export DLData
type DLData <: EventCollection
  entries :: Vector{EventLibrary}
  dir :: Union{AbstractString,Void}

  function DLData(entries::Vector{EventLibrary})
    for e in entries
      @assert e != nothing
    end
    new(entries, nothing)
  end
end

export no_events
no_events() = DLData(EventLibrary[])

function Base.cat(libs::Vector{EventLibrary})
    if length(libs) == 1
        return libs[1]
    end
    return DLData(libs)
end
Base.cat(libs::EventLibrary...) = cat([lib for lib in libs])

function _set_shallow(lib::EventLibrary, from::EventLibrary)
  lib.initialization_function = from.initialization_function
  lib.waveforms = from.waveforms
  lib.labels = from.labels
  lib.prop = from.prop
end

is_initialized(lib::EventLibrary) = lib.initialization_function == nothing

export initialize
function initialize(lib::EventLibrary)
  if lib.initialization_function != nothing
    lib.initialization_function(lib)
    lib.initialization_function = nothing
  end
  return lib
end

export dispose
function dispose(lib::EventLibrary)
  lib.waveforms = zeros(Float32, 0, 0)
  empty!(lib.labels)
end
dispose(data::DLData) = for lib in data dispose(lib) end

function initialize(data::DLData)
  for lib in data
    initialize(lib)
  end
  return data
end

Base.start(data::DLData) = 1
Base.done(data::DLData, state) = length(data.entries) == state-1
Base.next(data::DLData, state) = data.entries[state], state+1
Base.length(data::DLData) = length(data.entries)

Base.start(lib::EventLibrary) = false
Base.done(lib::EventLibrary, state) = state
Base.next(lib::EventLibrary, state) = lib, true
Base.length(lib::EventLibrary) = 1

export set_property!
set_property!(lib::EventLibrary, key::Symbol, value) = lib.prop[key] = value
set_property!(data::DLData, key::Symbol, value) = [set_property!(lib, key, value) for lib in data.entries]

export flatten
function flatten(data::DLData)
  if length(data.entries) == 0
    return data
  end
  if length(data.entries) == 1
    return data.entries[1]
  end
  return cat_events(data.entries...)
end
flatten(lib::EventLibrary) = lib

export waveforms
waveforms(lib::EventLibrary) = initialize(lib).waveforms
waveforms(data::DLData) = waveforms(flatten(data))

export sampling_rate
sampling_rate(lib::EventLibrary) = lib[:sampling_rate]
sampling_rate(data::DLData) = sampling_rate(data.entries[1])

export set_sampling_rate!
set_sampling_rate!(events::EventCollection, rate::Real) = set_property!(events, :sampling_rate, rate)

export sampling_period
sampling_period(lib::EventLibrary) = 1/sampling_rate(lib)
sampling_period(data::DLData) = sampling_period(data.entries[1])

export sample_size
sample_size(lib::EventLibrary) = size(initialize(lib).waveforms, 1)
sample_size(data::DLData) = sample_size(data.entries[1])

export sample_times
sample_times(events::EventCollection) = linspace(0, (sample_size(events)-1)*sampling_period(events), sample_size(events))

export filter

filter(lib::EventCollection, key::Symbol, value::Union{AbstractString,Number}) = filter(lib, key, x -> x==value)
filter{T<:Union{AbstractString,Number}}(lib::EventCollection, key::Symbol, values::Vector{T}) = filter(lib, key, x -> x in values)

function filter(lib::EventLibrary, predicate_key::Symbol, predicate::Function)
    if haskey(lib.prop, predicate_key)
        propval = lib.prop[predicate_key]
        if !predicate(propval)
            return nothing
        else
            return lib
        end
    end

  if is_initialized(lib)
    indices = find(predicate, lib.labels[predicate_key])
    return lib[indices]
  else
    result = copy(lib) # shallow copy
    result.initialization_function = lib2 -> _set_shallow(lib2, filter(initialize(lib), predicate_key, predicate))
    delete!(result.prop, :eventcount)
    return result
  end
end

function filter(data::DLData, predicate_key::Symbol, predicate::Function)
  if hasproperty(data, predicate_key)
    # Filter list of datasets
    filtered_indices = find(lib -> predicate(lib.prop[predicate_key]), data.entries)
    if length(filtered_indices) == 1
        return data.entries[filtered_indices[1]]
    end
    return DLData(data.entries[filtered_indices])
  else
    # Filter individual datasets
    return DLData([filter(lib, predicate_key, predicate) for lib in data])
  end
end

Base.filter!(lib::EventCollection, key::Symbol, value::Union{AbstractString,Number}) = filter!(lib, key, x -> x==value)
Base.filter!{T<:Union{AbstractString,Number}}(lib::EventCollection, key::Symbol, values::Vector{T}) = filter!(lib, key, x -> x in values)

function Base.filter!(lib::EventLibrary, predicate_key::Symbol, predicate::Function)
    # First check whether it's a property
  if haskey(lib.prop, predicate_key)
        propval = lib.prop[predicate_key]
        if !predicate(propval)
            lib.waveforms = zeros(Float32, 0, 0)
            lib.initialization_function = nothing
            for label in keys(lib.labels)
              lib.labels[label] = zeros(Float32, 0)
            end
        end
        return lib
    end

    # Else it has to be a label
  if is_initialized(lib)
    indices = find(predicate, lib.labels[predicate_key])
    lib.waveforms = lib.waveforms[:,indices]
    for label in keys(lib.labels)
      lib.labels[label] = lib.labels[label][indices]
    end
  else
    prev_initialization = lib.initialization_function
    lib.initialization_function = lib2 -> begin
      lib2.initialization_function = nothing
      prev_initialization(lib2)
      filter!(lib2, predicate_key, predicate)
    end
  end
  delete!(lib.prop, :eventcount)
  return lib
end

function Base.filter!(data::DLData, predicate_key::Symbol, predicate::Function)
  if hasproperty(data, predicate_key)
    # Filter list of datasets
    filtered_indices = find(lib -> predicate(lib.prop[predicate_key]), data.entries)
    data.entries = data.entries[filtered_indices]
    return data
  else
    # Filter individual datasets
    for lib in data
      filter!(lib, predicate_key, predicate)
    end
  end
end

function filter(lib::EventLibrary, predicate_keys::Vector{Symbol}, predicate::Function)
    if is_initialized(lib)
        indices = Int64[]
        arg_arrays = [lib.labels[predicate_key] for predicate_key in predicate_keys]
        for i in 1:eventcount(lib)
            args = [arg_array[i] for arg_array in arg_arrays]
            if predicate(args...)
                push!(indices, i)
            end
        end
        return lib[indices]
    else
        result = copy(lib) # shallow copy
        result.initialization_function = lib2 -> _set_shallow(lib2, filter(initialize(lib), predicate_keys, predicate))
        delete!(result.prop, :eventcount)
        return result
    end
end

function filter(data::DLData, predicate_keys::Vector{Symbol}, predicate::Function)
    return cat([filter(lib, predicate_keys, predicate) for lib in data])
end

export filter_by_proxy
function filter_by_proxy(lib::EventLibrary, predicate_lib::EventLibrary, predicate_key, predicate)
  @assert eventcount(lib) == eventcount(predicate_lib)
  indices = find(predicate, predicate_lib[predicate_key])
  return lib[indices]
end
function filter_by_proxy(data::DLData, predicate_data::DLData, predicate_key, predicate)
  return DLData([filter_by_proxy(data.entries[i], predicate_data.entries[i], predicate_key, predicate) for i in 1:length(data.entries)])
end

export getindex
function getindex(lib::EventLibrary, key::Symbol)
    if key == :wf || key == :waveforms || key == :waveform
        return waveforms(lib)
    end
    if key == :idx || key == :index || key == :indices
        return 1:eventcount(lib)
    end
    if haskey(lib.prop, key)
        return lib.prop[key]
    end
    if key == :name
        return "Unnamed"
    end
    return initialize(lib).labels[key]
end
function getindex(data::DLData, key::Symbol)
    if length(data) == 0
        return []
    end
    if haskey(data.entries[1].prop, key)
        return [lib[key] for lib in data]
    end
    return flatten(data)[key]
end

function getindex(events::EventLibrary, key::AbstractString)
  return getindex(events, Symbol(key))
end

function getindex(events::EventLibrary, selection::Union{Range, Integer, Array{Int64}, Array{Int32}})
  initialize(events)
  result = EventLibrary(events.waveforms[:,selection])
  for (key, value) in events.labels
    result.labels[key] = value[selection]
  end
  for (key,value) in events.prop
    result.prop[key] = value
  end
  delete!(result.prop, :eventcount)
  return result
end
getindex(data::DLData, selection::Union{UnitRange,Integer, Vector{Int64}, Vector{Int32}}) = flatten(data)[selection]

function getindex{T<:Union{AbstractString,Number}}(events::EventCollection, predicates::Pair{Symbol,T}...)
  for predicate in predicates
    events = filter(events, predicate[1], predicate[2])
  end
  return events
end

export eventcount
function eventcount(events::EventLibrary)
    if size(events.waveforms, 2) > 0
        return size(events.waveforms, 2)
    elseif haskey(events, :eventcount)
        return events[:eventcount]
    else
        return size(initialize(events).waveforms, 2)
    end
end
function eventcount(data::DLData)
    if length(data.entries) == 0
        return 0
    end
    return sum([eventcount(lib) for lib in data.entries])
end

Base.haskey(events::EventLibrary, key::Symbol) = haskey(events.labels, key) || haskey(events.prop, key)
Base.haskey(data::DLData, key::Symbol) = length(data) > 0 && haskey(data.entries[1])

export hasproperty
hasproperty(events::EventLibrary, key::Symbol) = haskey(events.prop, key)
hasproperty(data::DLData, key::Symbol) = length(data) > 0 && hasproperty(data.entries[1], key)

export copy
function copy(events::EventLibrary)
  result = EventLibrary(events.waveforms)
  for (key, value) in events.labels
    result.labels[key] = value
  end
  for (key,value) in events.prop
    result.prop[key] = value
  end
  return result
end

function copy(data::DLData)
  return DLData([copy(lib) for lib in data])
end

export deepcopy
function deepcopy(events::EventLibrary)
  result = EventLibrary(copy(events.waveforms))
  result.initialization_function = events.initialization_function
  for (key, value) in events.labels
    result.labels[key] = copy(value)
  end
  for (key,value) in events.prop
    result.prop[key] = deepcopy(value)
  end
  return result
end

function Base.string(events::EventLibrary)
  if is_initialized(events)
    return "$(events[:name]) ($(eventcount(events)) events)"
  else
    return "$(events[:name]) (not yet initialized)"
  end
end
Base.println(lib::EventLibrary) = println(string(lib))
Base.string(data::DLData) = "DLData ($(length(data.entries)) subsets, $(eventcount(data)) events)"
Base.show(data::DLData) = println(data)
Base.show(io::IO, data::DLData) = println(io, string(data))
Base.print(io::IOBuffer, data::DLData) = println(io, string(data))
Base.print(data::DLData) = print(string(data))
function Base.println(data::DLData)
  println(string(data))
  for lib in data
    println(string(lib))
  end
end

Base.summary(events::EventCollection) = string(events)

export eventinfo
function eventinfo(events::EventCollection, index::Integer)
    lib, i = unravel_index(events, index)
    println("Event properties:")
    for (key, arr) in lib.labels
        println("  $key: $(arr[i])")
    end
    println("General properties:")
    for (key, val) in lib.prop
        println("  $key: $val")
    end
end

export setname!
function setname!(events::EventLibrary, name::String)
  events.prop[:name] = name
end

function setname!(data::DLData, names::Vector{String})
  for i in 1:length(data.entries)
    setname!(data.entries[i], names[i])
  end
end

function unravel_index(events::EventCollection, index::Integer)
    total = 0
    for lib in events
        if index > total+eventcount(lib)
            total += eventcount(lib)
        else
            return lib, index-total
        end
    end
    return throw(NoSuchEventException())
end

# export convert
# function convert(target::Type{String}, events::EventLibrary)
#   if target == String
#     return string(events)
#   end
#   throw(MethodError(convert, target, events))
# end

export get_classifiers
function get_classifiers(events::EventLibrary)
  if haskey(events.prop, :classifier_names)
    return events.prop[:classifier_names]
  else
    return String[]
  end
end

export push_classifier!
function push_classifier!(events::EventLibrary, clname::AbstractString)
  if haskey(events.prop, :classifier_names)
    push!(events.prop[:classifier_names], clname)
  else
    events.prop[:classifier_names] = [clname]
  end
end

export cat_events
function cat_events(libs::EventLibrary...)
  @assert length(libs) > 0
  nonempty = find(lib->eventcount(lib)>0, libs)

  result = EventLibrary(identity)
  result.initialization_function = nothing

  # Cat waveforms
  if length(nonempty) > 0
    result.waveforms = hcat([waveforms(lib) for lib in libs[nonempty]]...) # this initializes the entries
    # Cat labels
    for key in keys(libs[1].labels)
      result.labels[key] = vcat([lib.labels[key] for lib in libs[nonempty]]...)
    end
  else
    result.waveforms = libs[1].waveforms
    result.labels = copy(libs[1]).labels
  end

  # Adopt shared property values
  for key in keys(libs[1].prop)
    values = [lib.prop[key] for lib in libs]
    if !any(x->x!=values[1], values)
      result.prop[key] = values[1]
    end
  end
    delete!(result.prop, :eventcount)
  return result
end

function Base.sort(lib::EventLibrary, labelkey::Symbol)
    p = sortperm(lib[labelkey])
    return lib[p]
end
Base.sort(data::DLData, labelkey::Symbol) = sort(flatten(data), labelkey)

function Base.split(data::DLData, datasets::Dict{AbstractString,Vector{AbstractFloat}})
    result = DLData(EventLibrary[])
    for lib in copy(data.entries)
      split_result = split(lib, datasets)
      for (set_name, lib) in split_result
        push!(result.entries, lib)
      end
    end
    return result
end

# function split_uninitialized(lib::EventLibrary, fractions::Dict{AbstractString,Vector{AbstractFloat}})
# end

function Base.split(lib::EventLibrary, fractions::Dict{AbstractString,Vector{AbstractFloat}})
  keylist_ids = lib[:keylist] # 1-based
  if eventcount(lib) > 0
    kl_count = Int(maximum(keylist_ids))
  else
    kl_count = 0
  end

  dsets = Dict{AbstractString,Vector{EventLibrary}}()
  for (dset_name, dset_sizes) in fractions
    dsets[dset_name] = EventLibrary[]
  end

  # for each keylist split separately
  for keylist_id in 1:kl_count
    keylist_indices = find(x->x==keylist_id, keylist_ids)
    N = length(keylist_indices)
    index_perm = randperm(MersenneTwister(0), N)

    depleted_fraction = 0.0
    # split shuffled indices into data sets
    for (dset_name, dset_sizes) in fractions
      dset_frac = dset_sizes[keylist_id]
      start_i = Int(round(depleted_fraction*N)) + 1
      end_i = Int(round((depleted_fraction+dset_frac)*N))
      dset_indices = keylist_indices[index_perm[start_i : end_i]]
      push!(dsets[dset_name], lib[dset_indices])
      depleted_fraction += dset_frac
    end
    if !(depleted_fraction ≈ 1)
      info("Some events were not assigned to any data set during split.")
    end
  end

  # Combine different keylists
  result = Dict{AbstractString,EventLibrary}()
  for dset_name in keys(fractions)
    if eventcount(lib) > 0
      result[dset_name] = cat_events(dsets[dset_name]...)
    else
      result[dset_name] = copy(lib)
    end
    result[dset_name].prop[:set] = dset_name
    setname!(result[dset_name], result[dset_name][:name]*"_"*dset_name)
  end
  return result
end

export label_energy_peaks!
function label_energy_peaks!(events::EventLibrary; label_key=:SSE, peaks0=[1620.7], peaks1=[1592.5], half_window=2.0)
  labels = [_get_label(events[:E][i], peaks0, peaks1, half_window) for i in 1:eventcount(events)]
  events.labels[label_key] = labels
  return events
end
function label_energy_peaks!(data::DLData; label_key=:SSE, peaks0=[1620.7], peaks1=[1592.5], half_window=2.0)
    for lib in data
        label_energy_peaks!(lib; label_key=label_key, peaks0=peaks0, peaks1=peaks1, half_window=half_window)
    end
end

function _get_label(energy, peaks0, peaks1, half_window)
  for peak0 in peaks0
    if abs(energy-peak0) <= half_window
      return 0
    end
  end
  for peak1 in peaks1
    if(abs(energy-peak1)) <= half_window
      return 1
    end
  end
  return -1
end

export equalize_counts_by_label
function equalize_counts_by_label(events::EventLibrary, label_key=:SSE)
  if !haskey(events, label_key)
    label_energy_peaks!(events, label_key)
  end
  labels = events.labels[label_key]
  i_SSE = find(x -> x==1, labels)
  i_MSE = find(x -> x==0, labels)
  count = min(length(i_SSE), length(i_MSE))
  info("Equalizing $(events[:name]) to $count counts. SSE: $(length(i_SSE)), MSE: $(length(i_MSE))")

  # trim and shuffle
  used_indices = [i_SSE[1:count];i_MSE[1:count]]
  shuffled_indices = used_indices[randperm(length(used_indices))]

  return events[shuffled_indices], shuffled_indices
end

export put_label!
function put_label!(events::EventLibrary, label_name::Symbol, label_data::Vector)
    if length(label_data) != eventcount(events)
        info("Provided label array $label_name has wrong length $(length(label_data)) ($(eventcount(events)) needed)")
        return
    end
    events.labels[label_name] = label_data
end

function put_label!(data::DLData, label_name::Symbol, label_data::Vector)
    @assert eventcount(data) == length(label_data)
    start_index = 1
    for lib in data
        put_label!(lib, label_name, label_data[start_index : start_index+eventcount(lib)-1])
        start_index += eventcount(lib)
    end
end

function Base.keys(lib::EventLibrary)
    return vcat(collect(keys(lib.labels)), collect(keys(lib.prop)))
end
function Base.keys(data::DLData)
    if length(data.entries) == 0
        return Symbol[]
    else
        return keys(data.entries[1])
    end
end

export detectors
function detectors(data::DLData)
    return unique([lib[:detector_name] for lib in data])
end
detectors(lib::EventLibrary) = [lib[:detector_name]]

export SSE_at
function SSE_at(lib::EventLibrary, energy::Real)
    if startswith(lib[:detector_name], "GD")
        SSE_indices = find(AoE->(AoE>0.98)&&(AoE<1.02), lib[:AoE])
    else
        SSE_indices = find(c->c==0, lib[:ANN_mse_class])
    end
    sorted_indices = sortperm(abs.(lib[:E].-energy))
    for idx in sorted_indices
        if idx in SSE_indices
            return idx
        end
    end
    throw(NoSuchEventException())
end
SSE_at(data::DLData, energy::Real) = SSE_at(flatten(data), energy)

export MSE_at
function MSE_at(lib::EventLibrary, energy::Real; AoE_cut=0.7)
    if startswith(lib[:detector_name], "GD")
        MSE_indices = find(AoE->(AoE<AoE_cut)&&(AoE!=0), lib[:AoE])
    else
        MSE_indices = find(c->c==1, lib[:ANN_mse_class])
    end
    sorted_indices = sortperm(abs.(lib[:E].-energy))
    for idx in sorted_indices
        if idx in MSE_indices
            return idx
        end
    end
    throw(NoSuchEventException())
end
MSE_at(data::DLData, energy::Real) = MSE_at(flatten(data), energy)

export any_at
function any_at(lib::EventCollection, energy::Real)
    return findmin(abs.(lib[:E].-energy))[2]
end

export equal_event_count_edges
function equal_event_count_edges(events::EventCollection, label::Symbol; events_per_bin::Real=2000)
    labels = sort(events[label])
    total_event_count = eventcount(events)
    equal_count_edges = []
    i = 0
    while (i*events_per_bin) < total_event_count
        push!(equal_count_edges, labels[i*events_per_bin+1])
        i += 1
    end
    return equal_count_edges
end

export lookup_property
function lookup_property(source::EventLibrary, sourceindex::Integer, target::EventCollection, propkey::Symbol; sameattrs=[:E, :E1, :E2, :E3, :E4, :filenum])
  sourceattr = [source[attr][sourceindex] for attr in sameattrs]
    for lib in target
        targetattrs = hcat([lib[attr] for attr in sameattrs]...)
        for i in 1:eventcount(lib)
          if targetattrs[i,:] == sourceattr
                if isa(lib[propkey], Array)
                    dim = length(size(lib[propkey]))
                    if dim == 0
                        return lib[propkey]
                    elseif dim == 1
                        return lib[propkey][i]
                    else
                        return lib[i:i][propkey]
                    end
                else
                    return lib[propkey]
                end
            end
        end
    end
    return nothing
end

export lookup_event
function lookup_event(source::EventLibrary, sourceindex::Integer, target::EventCollection; sameattrs=[:E, :E1, :E2, :E3, :E4, :filenum])
  sourceattr = [source[attr][sourceindex] for attr in sameattrs]
    for lib in target
        targetattrs = hcat([lib[attr] for attr in sameattrs]...)
        for i in 1:eventcount(lib)
          if targetattrs[i,:] == sourceattr
                return lib[i:i]
            end
        end
    end
    return nothing
end

export normalize_AoE!
function normalize_AoE!(events::EventCollection, peak_energy=1592.5, peak_window=2)
  dep_events = filter(events, :E, E->(E>=peak_energy-peak_window)&&(E<peak_energy-peak_window))
  events[:AoE] ./= mean(dep_events[:AoE])
end

export calculate_AoE!
function calculate_AoE!(events::EventCollection; denoise=false, correction=nothing, mse_cutoff=0.7)
    orig_events = events
    if denoise
        events = deepcopy(events)
        denoise_waveforms!(events)
    end
    waveforms = events[:wf]
    energies = events[:E]
    AoE = [maximum(waveforms[:,i]) for i in 1:eventcount(events)]
    AoE /= median(AoE)

    if correction == "exponential"
      exp_curve(E,p) = p[1] * exp.(-p[2]*E) + p[3]
      fit_result = curve_fit(exp_curve, Array{Float64}(energies[1:1000:end]), Array{Float64}(AoE[1:1000:end]), Float64[5, 1/1000, 0.5]).param
      println("A/E Best fit: $fit_result")
      for i in 1:length(energies)
        E = energies[i]
        AoE[i] = AoE[i] - exp_curve(E, fit_result)
      end
      AoE += 1-median(AoE)
    end
    put_label!(orig_events, :AoE, AoE)
    AoE_class = [(AoE[i] < mse_cutoff) ? 1 : 0 for i in 1:length(AoE)]
    put_label!(orig_events, :ANN_mse_class, AoE_class)
end

export calculate_SingleSeg!
function calculate_SingleSeg!(events::EventCollection; key=:SSeg, dE=10, segcount_key=:SegCount)
  core = events[:E]
  E1 = events[:E1]
  E2 = events[:E2]
  E3 = events[:E3]
  E4 = events[:E4]
  SSeg = zeros(Float32, length(core))
  SegCount = zeros(Float32, length(core))
  for i in 1:length(core)
        if (E1[i] > core[i]-dE) && (E1[i] < core[i]+dE)
          SSeg[i] = 1
        elseif (E2[i] > core[i]-dE) && (E2[i] < core[i]+dE)
          SSeg[i] = 2
        elseif (E3[i] > core[i]-dE) && (E3[i] < core[i]+dE)
          SSeg[i] = 3
        elseif (E4[i] > core[i]-dE) && (E4[i] < core[i]+dE)
          SSeg[i] = 4
        end

        n = 0
        if E1[i] > dE n+=1 end
        if E2[i] > dE n+=1 end
        if E3[i] > dE n+=1 end
        if E4[i] > dE n+=1 end
        SegCount[i] = n

    end
    put_label!(events, key, SSeg)
    put_label!(events, segcount_key, SegCount)
end

export fit_SSE_band
function fit_SSE_band(events::EventCollection)
    hist1D = fit(Histogram{Float64}, events[:AoE], linspace(0,2,400), closed=:left)

    model(x, p) = p[1]./sqrt(2*pi.*abs.(p[3])).*exp.(-((x.-p[2]).^2 ./(2 .*p[3].^2)))

    maxidx = findmax(hist1D.weights)[2]
    cutoff = [(i < maxidx*0.96)?0:1 for i in 1:length(hist1D.edges[1])]

    weights = sqrt.(1.+vcat(0, hist1D.weights)) .* cutoff
    fit_result = curve_fit(model, hist1D.edges[1], vcat(0, hist1D.weights), weights, [1000*2*pi, 1.1, 0.1])
    covar = estimate_covar(fit_result)

    return fit_result.param[2], fit_result.param[3], x->model(x,fit_result.param), sqrt(covar[2,2]), sqrt(covar[3,3])
end

export calculate_normalized_AoE!
function calculate_normalized_AoE!(events::EventCollection; energy_slices = [1005., 1025., 1045.,
      1115., 1135., 1155., 1175., 1195., 1205., 1225., 1245.,
      1305., 1325., 1345., 1365., 1385., 1405., 1425., 1445., 1465., 1485.,
      1535., 1555.,
      1830., 1850., 1870., 1890., 1910., 1930., 1950., 1970., 1990., 2010., 2030., 2050., 2070.,
      2140., 2160., 2180., 2200., 2220., 2240., 2260., 2280.], log=false)
  means = Float64[]
  mean_errs = Float64[]
  sigmas = Float64[]
  sigma_errs = Float64[]
  curves = []

  for slice in energy_slices
    log && info("Fitting slice $slice keV")
      slice_events = filter(events, :E, E->(E>slice-10)&&(E<slice+10))
      gmean, gsigma, curve, mean_err, sigma_err = fit_SSE_band(slice_events)
      push!(means, gmean)
      push!(sigmas, gsigma)
      push!(curves, curve)
      push!(mean_errs, mean_err)
      push!(sigma_errs, sigma_err)
  end

  mean_model(x,p) = p[1] * x + p[2]
  line_fit = curve_fit(mean_model, energy_slices, means, 1./mean_errs.^2, [0, 1.1])
  log && info("A/E Slope: $(line_fit.param[1])")

  sigma_model(x,p) = p[1]
  sigma_fit = curve_fit(sigma_model, energy_slices, sigmas, 1./sigma_errs.^2, [0.1])
  mean_sigma = sigma_fit.param[1]
  log && info("Sigma fit to $(sigma_fit.param[1])")

  normalize_AoE(E, AoE) = (AoE/mean_model(E, line_fit.param) - 1) / mean_sigma

  put_label!(events, :AoE_norm, normalize_AoE.(events[:E], events[:AoE]))
  return normalize_AoE
end

export peak_FWHM
function peak_FWHM(events::EventCollection, peak_center::Float64; half_window::Float64=10.0, bin_width::Float64=0.5)
    start_energy = peak_center-half_window-bin_width
    end_energy = peak_center+half_window+bin_width
    nbins = Int64(round((end_energy-start_energy)/bin_width+2))
    bin_edges = linspace(start_energy, end_energy, nbins)
    energy_hist = fit(Histogram{Float64}, events[:E], bin_edges, closed=:left)
    energy_axis = collect(bin_edges[2:end] + bin_edges[1:end-1]) ./ 2

    mean_count = mean(energy_hist.weights)

    f(e, p) = p[1] + p[2]*(e-peak_center) + p[3] * exp.(-0.5 .* ((e-peak_center) ./ p[4]).^2)
    fit_result = curve_fit(f, energy_axis, energy_hist.weights, [mean_count*0.2, 0.0, mean_count*0.8, 5.0])
    return fit_result.param[4] * 2.35482
end


text:
const IdealGasConst = 8.314472

function R() return IdealGasConst end


text:
function select_tobereduced!(
                basis::Basis, tobereduced::Basis,
                matrix::MacaulayMatrix,
                symbol_ht::MonomialHashtable, ht::MonomialHashtable)

    # prepare to load all elems from tobereduced
    # into low rows
    reinitialize_matrix!(matrix, 2^6)
    resize!(matrix.lowrows, tobereduced.ntotal)

    # TODO
    etmp = zeros(UInt16, ht.explen)

    for i in 1:tobereduced.ntotal
        matrix.nrows += 1
        gen = tobereduced.gens[i]
        h = UInt32(0)
        matrix.lowrows[matrix.nrows] = multiplied_poly_to_matrix_row!(symbol_ht, ht, h, etmp, gen)
        matrix.low2coef[matrix.nrows] = i
    end

    basis.ntotal
    basis.nlead = basis.ndone = basis.ntotal
    basis.isred .= 0
    for i in 1:basis.nlead
        basis.nonred[i] = i
        basis.lead[i] = ht.hashdata[basis.gens[i][1]].divmask
    end
end

function normal_form_f4!(
            ring::PolyRing,
            basis::Basis,
            ht::MonomialHashtable,
            tobereduced::Basis)

    matrix = initialize_matrix(ring)
    symbol_ht = initialize_secondary_hash_table(ht)

    select_tobereduced!(basis, tobereduced, matrix, symbol_ht, ht)

    symbolic_preprocessing!(basis, matrix, ht, symbol_ht)

    convert_hashes_to_columns!(matrix, symbol_ht)

    sort_matrix_rows_decreasing!(matrix) # for pivots,  AB part
    exact_sparse_linear_algebra_nf!(matrix, tobereduced, basis)

    convert_nf_rows_to_basis_elements!(matrix, tobereduced, ht, symbol_ht)

    tobereduced
end

function normal_form_f4(
                ring::PolyRing,
                basisexps::Vector{Vector{Vector{UInt16}}},
                basiscoeffs::Vector{Vector{UInt64}},
                tobereducedexps::Vector{Vector{Vector{UInt16}}},
                tobereducedcfs::Vector{Vector{UInt64}},
                rng::Rng,
                tablesize::Int=2^16) where {Rng}

    basis, ht = initialize_structures(ring, basisexps,
                                        basiscoeffs, rng, tablesize)

    tobereduced, ht = initialize_structures(ring, tobereducedexps,
                                        tobereducedcfs, rng, tablesize,
                                        ht)
    tobereduced = normal_form_f4!(ring, basis, ht, tobereduced)

    export_basis_data(tobereduced, ht)
end