"""
    symplectic_integrate(x₀, p₀, Λ, U, δUδx, N=50, ϵ=0.1, progress=false)

Do a symplectic integration of the potential energy `U` (with gradient
`δUδx`) starting from point `x₀` with momentum `p₀` and mass matrix
`Λ`. The number of steps is `N` and the step size `ϵ`.

Returns `ΔH, xᵢ, pᵢ` corresponding to change in Hamiltonian, and final
position and momenta. If `history_keys` is specified a history of
requested variables throughout each step is also returned.

"""
function symplectic_integrate(
    x₀::AbstractVector{T}, p₀, Λ, U, δUδx=x->gradient(U,x)[1];
    N=50, ϵ=T(0.1), progress=false, history_keys=nothing
) where {T}

    xᵢ, pᵢ = x₀, p₀
    δUδxᵢ = δUδx(xᵢ)
    H(x,p) = U(x) - p⋅(Λ\p)/2

    history = []

    @showprogress (progress ? 1 : Inf) "Symplectic Integration: " for i=1:N
        xᵢ₊₁    = xᵢ - T(ϵ) * (Λ \ (pᵢ - T(ϵ)/2 * δUδxᵢ))
        δUδxᵢ₊₁ = δUδx(xᵢ₊₁)
        pᵢ₊₁    = pᵢ - T(ϵ)/2 * (δUδxᵢ₊₁ + δUδxᵢ)
        xᵢ, pᵢ, δUδxᵢ = xᵢ₊₁, pᵢ₊₁, δUδxᵢ₊₁

        if !isnothing(history_keys)
            historyᵢ = (;i, x=xᵢ, p=pᵢ, δUδx=δUδxᵢ₊₁, H=(haskey(history_keys,:H) ? H(xᵢ,pᵢ) : nothing))
            push!(history, select(historyᵢ, history_keys))
        end

    end

    ΔH = H(xᵢ,pᵢ) - H(x₀,p₀)

    if isnothing(history)
        return ΔH, xᵢ, pᵢ
    else
        return ΔH, xᵢ, pᵢ, history
    end

end

@doc doc"""
    grid_and_sample(lnP::Function; range::NamedTuple; progress=false, nsamples=1)

Interpolate the log pdf `lnP` with support on `range`, and return  the
integrated log pdf as well `nsamples` samples (drawn via inverse transform
sampling)

`lnP` should either accept a NamedTuple argument and `range` should be a
NamedTuple mapping those same names to `range` objects specifying where to
evaluate `lnP`, e.g.:

```julia
grid_and_sample(nt->-(nt.x^2+nt.y^2)/2, (x=range(-3,3,length=100),y=range(-3,3,length=100)))
```

or `lnP` should accept a single scalar argument and `range` should be directly
the range for this variable:

```julia
grid_and_sample(x->-x^2/2, range(-3,3,length=100))
```

The return value is `(lnP, samples, Px)` where `lnP` is an interpolated/smoothed
log PDF which can be evaluated anywhere within the original range, `Px` are
sampled points of the original PDF, and `samples` is a NamedTuple giving the
Monte-Carlo samples of each of the parameters.

(Note: only 1D sampling is currently implemented, but 2D like in the example
above is planned)
"""
function grid_and_sample(lnP::Function, range::AbstractVector; progress=false, kwargs...)
    lnPs = @showprogress (progress ? 1 : Inf) "Grid Sample: " map(lnP, range)
    grid_and_sample(lnPs, range; progress=progress, kwargs...)
end

function grid_and_sample(lnPs::Vector{<:BatchedReal}, xs::AbstractVector; kwargs...)
    batches = [grid_and_sample(batch_index.(lnPs,i), xs; kwargs...) for i=1:batch_length(lnPs[1])]
    ((batch(getindex.(batches,i)) for i=1:3)...,)
end

function grid_and_sample(lnPs::Vector, xs::AbstractVector; progress=false, nsamples=1, span=0.25, require_convex=false)

    # trim leading/trailing zero-probability regions
    support = findnext(isfinite,lnPs,1):findprev(isfinite,lnPs,length(lnPs))
    xs = xs[support]
    lnPs = lnPs[support]

    if require_convex
        support = longest_run_of_trues(finite_second_derivative(lnPs) .< 0)
        xs = xs[support]
        lnPs = lnPs[support]
    end

    # interpolate PDF
    xmin, xmax = first(xs), last(xs)
    lnPs = lnPs .- maximum(lnPs)
    ilnP = loess(xs, lnPs, span=span)

    # normalize the PDF. note the smoothing is done of the log PDF.
    cdf(x) = quadgk(nan2zero∘exp∘ilnP,xmin,x,rtol=1e-3)[1]
    logA = nan2zero(log(cdf(xmax)))
    lnPs = (ilnP.ys .-= logA)
    ilnP.bs[:,1] .-= logA

    # draw samples via inverse transform sampling
    θsamples = @showprogress (progress ? 1 : Inf) map(1:nsamples) do i
        r = rand()
        if (cdf(xmin)-r)*(cdf(xmax)-r) >= 0
            first(lnPs) > last(lnPs) ? xmin : xmax
        else
            fzero(x->cdf(x)-r, xmin, xmax, xatol=(xmax-xmin)*1e-3)
        end
    end

    (nsamples==1 ? θsamples[1] : θsamples), ilnP, lnPs

end

function grid_and_sample(lnP::Function, range::NamedTuple{S,<:NTuple{1}}; kwargs...) where {S}
    NamedTuple{S}.(Ref.(grid_and_sample(x -> lnP(NamedTuple{S}(x)), first(range); kwargs...)))
end

# allow more convenient evaluation of Loess-interpolated functions
(m::Loess.LoessModel)(x) = Loess.predict(m,x)

@doc doc"""
    sample_joint(ds::DataSet; kwargs...)

Sample the joint posterior, $\mathcal{P}(f,\phi,\theta\,|\,d)$.

Keyword arguments:

* `nsamps_per_chain` — The number of samples per chain.
* `nchains = 1` — Number of chains in parallel.
* `nchunk = 1` — Number of steps between parallel chain communication.
* `nsavemaps = 1` — Number of steps in between saving maps into chain.
* `nburnin_always_accept = 0` — Number of steps at the beginning of
  the chain to always accept HMC steps regardless of integration
  error.
* `nburnin_fixθ = 0` — Number of steps at the beginning of the chain
  before starting to sample `θ`.
* `Nϕ = :qe` — Noise to use in the initial approximation to the
  Hessian. Can give `:qe` to use the quadratic estimate noise.
* `chains = nothing` — `nothing` to start a new chain; the return
  value from a previous call to `sample_joint` to resume those chains;
  `:resume` to resume chains from a file given by `filename`
* `θrange` — Range and density to grid sample parameters as a
  NamedTuple, e.g. `(Aϕ=range(0.7,1.3,length=20),)`.
* `θstart` — Starting values of parameters as a NamedTuple, e.g.
  `(Aϕ=1.2,)`, or nothing to randomly sample from θrange
* `ϕstart` — Starting `ϕ`, either a `Field` object, `:quasi_sample`,
  or `:best_fit`
* `metadata` — Does nothing, but is saved into the chain file
* `nhmc = 1` — Number of HMC passes per `ϕ` Gibbs step.
* `symp_kwargs = fill((N=25, ϵ=0.01), nhmc)` — an array of NamedTupe
  kwargs to pass to [`symplectic_integrate`](@ref). E.g.
  `[(N=50,ϵ=0.1),(N=25,ϵ=0.01)]` would do 50 large steps then 25
  smaller steps per each Gibbs pass. If specified, `nhmc` is ignored.
* `wf_kwargs` — Keyword arguments to pass to [`argmaxf_lnP`](@ref) in
  the Wiener Filter Gibbs step.
* `MAP_kwargs` — Keyword arguments to pass to [`MAP_joint`](@ref) when
  computing the starting point.
"""
function sample_joint(
    ds :: DataSet;
    gibbs_initializers = [
        gibbs_initialize_θ!,
        gibbs_initialize_ϕ!,
        gibbs_initialize_f!
    ],
    gibbs_samplers = [
        gibbs_sample_f!,
        gibbs_mix!,
        gibbs_sample_ϕ!,
        gibbs_unmix!,
        gibbs_postprocess!
    ],
    nsamps_per_chain,
    nchains = nworkers(),
    nfilewrite = 5,
    nsavemaps = 1,
    filename = nothing,
    ϕstart = :prior,
    θstart = :prior,
    θrange = NamedTuple(),
    pmap = (myid() in workers() ? map : pmap),
    pool = default_worker_pool(),
    conjgrad_kwargs = (tol=1e-1, nsteps=500),
    preconditioner = :diag,
    nhmc = 1,
    symp_kwargs = fill((N=25, ϵ=0.01), nhmc),
    MAP_kwargs = (nsteps=40,),
    metadata = nothing,
    progress = false,
    storage = basetype(fieldinfo(ds.d).M),
    grid_and_sample_kwargs = (;),
    kwargs...
)

    # rundat is a Dict with all the args and kwargs minus a few removed ones
    rundat = merge!(foldl(delete!, (:ds, :gibbs_initializers, :gibbs_samplers, :kwargs, :pmap, :pool), init=Base.@locals()), kwargs)
    rundat[:Nbatch] = batch_length(ds.d)

    # dont adapt things passed in kwargs when we adapt the state dict
    _adapt(storage, state) = Dict(k => (haskey(rundat,k) ? v : adapt(storage, v)) for (k,v) in state)

    function filter_for_saving(state, step)
        Dict(k=>v for (k,v) in state if (!(haskey(rundat,k)) && !(k in (:pbar_dict, :timer)) && (step == 1 || (step % nsavemaps) == 0 || !isa(v,Field))))
    end

    # validate arguments
    if !(progress in [false,:summary,:verbose])
        error("`progress` should be one of [false, :summary, :verbose]")
    end
    if (filename!=nothing && splitext(filename)[2]!=".jld2")
        error("Chain filename '$filename' should have '.jld2' extension.")
    end

    # seed
    @everywhere @eval CMBLensing seed!(global_rng_for($storage))

    # distributed dataset object to workers once
    set_distributed_dataset(ds)

    # initialize chains
    if (filename != nothing) && isfile(filename)

        @info "Resuming chain at $filename"
        local chunks_index, prev_chunks
        jldopen(filename,"r") do io
            chunks_index = maximum([parse(Int,k[8:end]) for k in keys(io) if startswith(k,"chunks_")]) + 1
            states = last.(read(io, "chunks_$(chunks_index-1)"))
            merge!.(states, Ref(rundat))
            @unpack step = states[1]
            chain_chunks = map(copy, repeated([], nchains))
        end

    else

        chunks_index = step = 1

        states = pmap(map(copy,repeated(rundat,nchains))) do state
            state = _adapt(storage, state)
            for gibbs_initialize! in gibbs_initializers
                gibbs_initialize!(state, get_distributed_dataset())
            end
            @pack! state = step
            _adapt(Array, state)
        end
        chain_chunks = [Any[filter_for_saving(state, step)] for state in states]

        if filename != nothing
            save(filename, "rundat", cpu(rundat))
        end

    end

    # setup progressbar
    setindex!.(states, copy.(Ref(OrderedDict{String,Any}("step"=>step))), :pbar_dict)
    if progress == :summary
        pbar = Progress(nsamps_per_chain-step, dt=0, desc="Gibbs chain: ")
        ProgressMeter.update!(pbar, showvalues=[("step", step)])
        @everywhere $reset_timer!()
    end

    # start sampling
    for step in (step+1):nsamps_per_chain

        setindex!.(states, step, :step)

        state₁, = states = pmap(pool, states) do state

            state = @⌛ _adapt(storage, state)
            @unpack step, pbar_dict = state

            timing = @⌛ "Gibbs passes" map(gibbs_samplers) do gibbs_sample!
                @elapsed gibbs_sample!(state, get_distributed_dataset())
            end

            state = @⌛ _adapt(Array, state)

            timer = get_defaulttimer()
            @pack! state = timing, timer

            state

        end

        push!.(chain_chunks, filter_for_saving.(states, step))

        if (filename != nothing) && ((step % nfilewrite) == 0)
            @⌛ jldopen(filename,"a+") do io
                wsession = JLDWriteSession()
                write(io, "chunks_$chunks_index", chain_chunks, wsession)
            end
            chunks_index += 1
            empty!.(chain_chunks)
        end

        if progress == :summary
            print("\033[2J")
            next!(pbar, showvalues = [("step",step), ("timing",state₁[:timing]), collect(state₁[:pbar_dict])...])
            merge!(state₁[:timer].inner_timers, get_defaulttimer().inner_timers)
            print(state₁[:timer])
            flush(stdout)
        end

    end

    Chains(chain_chunks)

end

lnP_arg_names(p::Union{Int,Symbol}, ds::DataSet) = lnP_arg_names(Val(p), ds)
lnP_arg_names(::Val{0},    ds::DataSet) = (:f,  :ϕ,  :θ)
lnP_arg_names(::Val{:mix}, ds::DataSet) = (:f°, :ϕ°, :θ)

## initialization

function gibbs_initialize_θ!(state, ds::DataSet)
    @unpack θstart, θrange, nchains, Nbatch = state
    θ = @match θstart begin
        :prior => map(range->batch((first(range) .+ rand(Nbatch) .* (last(range) - first(range)))...), θrange)
        (_::NamedTuple) => θstart
        _ => throw(ArgumentError(θstart))
    end
    lnPθ = map(_->missing, θrange)
    @pack! state = θ, lnPθ
end

function gibbs_initialize_f!(state, ds::DataSet)
    f = missing
    @pack! state = f
end

function gibbs_initialize_ϕ!(state, ds::DataSet)
    @unpack ϕstart, θ, nchains, Nbatch = state
    ϕ = @match ϕstart begin
        :prior     => simulate(ds.Cϕ(θ); Nbatch)
        0          => zero(diag(ds.Cϕ)) * batch(ones(Int,Nbatch)...)
        (_::Field) => ϕstart
        (:quasi_sample|:best_fit) => MAP_joint(
            adapt(storage, ds(θstart)),
            progress = (progress==:verbose ? :summary : false),
            Nϕ = adapt(storage,Nϕ),
            quasi_sample = (ϕstart==:quasi_sample); MAP_kwargs...
        ).ϕ
        _ => throw(ArgumentError(ϕstart))
    end
    @pack! state = ϕ
end

## gibbs passes

@⌛ function gibbs_sample_f!(state, ds::DataSet)
    @unpack f, progress, preconditioner, conjgrad_kwargs = state
    z_other = delete(select(state, lnP_arg_names(0, ds)), :f)
    f = argmaxf_lnP(
        z_other..., ds;
        which = :sample,
        guess = ismissing(f) ? nothing : f,
        preconditioner = preconditioner,
        conjgrad_kwargs = (progress=(progress==:verbose), conjgrad_kwargs...)
    )
    @pack! state = f
end

@⌛ function gibbs_sample_ϕ!(state, ds::DataSet)
    @unpack f°, ϕ°, θ, symp_kwargs, progress, step, nburnin_always_accept = state
    U = ϕ° -> lnP(:mix, f°, ϕ°, θ, ds)
    ϕ°, ΔH, accept = hmc_step(U, ϕ°, mass_matrix_ϕ(θ,ds); symp_kwargs, progress, always_accept=(step<nburnin_always_accept))
    @pack! state = ϕ°, ΔH, accept
end

function hmc_step(U::Function, x, Λ; symp_kwargs, progress, always_accept)
    local ΔH, accept
    for kwargs in symp_kwargs
        p = simulate(Λ)
        (ΔH, xtest) = symplectic_integrate(
            x, p, Λ, U;
            progress = (progress==:verbose),
            kwargs...
        )
        accept = batch(@. always_accept | (log(rand()) < $unbatch(ΔH)))
        @. x = accept * xtest + (1 - accept) * x
    end
    x, ΔH, accept
end

@⌛ function mass_matrix_ϕ(θ, ds)
    @unpack G, Cϕ, Nϕ = ds(θ)
    pinv(G)^2 * (pinv(Cϕ) + pinv(Nϕ))
end

function gibbs_sample_slice_θ!(k::Symbol)
    @⌛ function gibbs_sample_slice_θ!(state, ds::DataSet)
        @unpack θ, θrange, lnPθ, pbar_dict, progress, grid_and_sample_kwargs = state
        z_other = delete(select(state, lnP_arg_names(:mix, ds)), :θ)
        θₖ, lnPθₖ = grid_and_sample(θₖ -> lnP(:mix, z_other..., @set(θ[k]=θₖ), ds), cpu(θrange[k]); progress=(progress==:verbose), grid_and_sample_kwargs...)
        @set! θ[k] = θₖ
        @set! lnPθ[k] = lnPθₖ
        pbar_dict[string(k)] = string_trunc(θₖ)
        @pack! state = θ, lnPθ
    end
end

@⌛ function gibbs_mix!(state, ds::DataSet)
    @unpack f, ϕ, θ = state
    f°, ϕ° = mix(f, ϕ, θ, ds)
    @pack! state = f°, ϕ°
end

@⌛ function gibbs_unmix!(state, ds::DataSet)
    @unpack f°, ϕ°, θ = state
    f, ϕ = unmix(f°, ϕ°, θ, ds)
    @pack! state = f, ϕ
end

## postprocessing

@⌛ function gibbs_postprocess!(state, ds::DataSet)
    @unpack f, ϕ, θ, pbar_dict, ΔH = state
    lnP = pbar_dict["lnP"] = CMBLensing.lnP(0, select(state, lnP_arg_names(0, ds))..., ds)
    pbar_dict["ΔH"] = ΔH
    f̃ = ds.L(ϕ) * f
    @pack! state = f̃, lnP
end

## util

function once_every(n)
    function (gibbs_sample!)
        function (state, ds)
            if iszero(state[:step] % n)
                gibbs_sample!(state, ds)
            end
        end
    end
end

function start_after_burnin(n)
    function (gibbs_sample!)
        function (state, ds)
            if state[:step] > n
                gibbs_sample!(state, ds)
            end
        end
    end
end


# ---
# title: 1614. Maximum Nesting Depth of the Parentheses
# id: problem1614
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: String
# link: <https://leetcode.com/problems/maximum-nesting-depth-of-the-parentheses/description/>
# hidden: true
# ---
#
# A string is a **valid parentheses string** (denoted **VPS** ) if it meets one
# of the following:
#
#   * It is an empty string `""`, or a single character not equal to `"("` or `")"`,
#   * It can be written as `AB` (`A` concatenated with `B`), where `A` and `B` are **VPS** 's, or
#   * It can be written as `(A)`, where `A` is a **VPS**.
#
# We can similarly define the **nesting depth** `depth(S)` of any VPS `S` as
# follows:
#
#   * `depth("") = 0`
#   * `depth(C) = 0`, where `C` is a string with a single character not equal to `"("` or `")"`.
#   * `depth(A + B) = max(depth(A), depth(B))`, where `A` and `B` are **VPS** 's.
#   * `depth("(" + A + ")") = 1 + depth(A)`, where `A` is a **VPS**.
#
# For example, `""`, `"()()"`, and `"()(()())"` are **VPS** 's (with nesting
# depths 0, 1, and 2), and `")("` and `"(()"` are not **VPS** 's.
#
# Given a **VPS** represented as string `s`, return _the **nesting depth** of
# _`s`.
#
#
#
# **Example 1:**
#
#
#
#     Input: s = "(1+(2*3)+(( _8_ )/4))+1"
#     Output: 3
#     Explanation: Digit 8 is inside of 3 nested parentheses in the string.
#
#
# **Example 2:**
#
#
#
#     Input: s = "(1)+((2))+((( _3_ )))"
#     Output: 3
#
#
# **Example 3:**
#
#
#
#     Input: s = "1+( _2_ *3)/(2-1)"
#     Output: 1
#
#
# **Example 4:**
#
#
#
#     Input: s = " _1_ "
#     Output: 0
#
#
#
#
# **Constraints:**
#
#   * `1 <= s.length <= 100`
#   * `s` consists of digits `0-9` and characters `'+'`, `'-'`, `'*'`, `'/'`, `'('`, and `')'`.
#   * It is guaranteed that parentheses expression `s` is a **VPS**.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


module Utils

import GraphIdx: Const, enumerate_edges, Graph, Weights

"""
    sum2(x)

Sum of squares.

```jldoctest
julia> TreeLas.Utils.sum2([1, 11])
122

```
"""
function sum2(x::Array{F,N})::F where{F<:Real, N}
    local s::F = zero(F)
    for xi in x
        s += xi^2
    end
    s
end

"""
    primal_objective(x, y, graph [, mu])

Compute the objective function, i.e.
```math
\\frac{1}{2} \\sum_{i=1}^n μ_i (x_i - y_i)^2 + \\sum_{(i,j) \\in E} λ_{ij} |x_i - x_j|
```
Hereby the edges (and edge weights ``λ_{ij}``) are obtained via `enumerate_edges(graph)`.
"""
@inline function primal_objective(
    x::Array{F,N},
    y::Array{F,N},
    graph::G,
    mu::F2 = Const(1.0),
)::F where {F, N, G<:Graph, F1<:Weights{F}, F2<:Weights{F}}
    @assert length(y) == length(x)
    local n = length(x)
    local obj::Ref{F} = Ref{F}(0.5 * sum(mu[i] * (x[i] - y[i])^2 for i = 1:n))
    enumerate_edges(graph) do ei::Int, u::Int, v::Int, lam::Float64
        obj[] += lam * abs(x[u] - x[v])
    end
    return obj[]
end

end


## make the surface struct
using transD_GP, Random
surface = transD_GP.Surface(imagegrid;useML=false, σ=noisegrid)
## lengthscale GP options (stationary)
nminlog10λ, nmaxlog10λ = 2, 150
pnorm = 2.
Klog10λ = transD_GP.GP.Mat32()
log10bounds = [log10(20_000) log10(300_000); log10(20_000) log10(300_000)]
δlog10λ = 0.2
# making sure it is not anisotropic ...
λlog10λ = 0.1*abs.([diff([extrema(xall[2,:])...])[1], diff([extrema(xall[2,:])...])[1]])
demean = false # essential for length scale change approximations
sdev_proplog10λ = 0.05*diff(log10bounds, dims=2)[:]
sdev_poslog10λ = 0.05*abs.([diff([extrema(xall[1,:])...])[1], diff([extrema(xall[2,:])...])[1]])
## Initialize a lengthscale model using these options
Random.seed!(12)
optlog10λ = transD_GP.OptionsStat(nmin = nminlog10λ,
                        nmax = nmaxlog10λ,
                        xbounds = [minimum(xall[1,:]) maximum(xall[1,:]);
                                   minimum(xall[2,:]) maximum(xall[2,:])],
                        fbounds = log10bounds,
                        xall = xall,
                        λ = λlog10λ,
                        δ = δlog10λ,
                        needλ²fromlog = true,
                        updatenonstat = true,
                        demean = demean,
                        sdev_prop = sdev_proplog10λ,
                        sdev_pos = sdev_poslog10λ,
                        pnorm = pnorm,
                        quasimultid = false,
                        K = Klog10λ,
                        timesλ = 4.
                        )
## make options for the nonstationary actual properties GP
nmin, nmax = 2, 150
ftrain = surface.d[surface.select]
fbounds = permutedims([extrema(ftrain)...])
sdev_prop = 0.05*diff(fbounds, dims=2)[:]
sdev_pos = 0.05*abs.([diff([extrema(xall[1,:])...])[1], diff([extrema(xall[2,:])...])[1]])
demean_ns = false
sampledc = true
δ = 1.
K = transD_GP.GP.Mat32()
## Initialize model for the nonstationary properties GP
Random.seed!(13)
opt = transD_GP.OptionsNonstat(optlog10λ,
                        nmin = nmin,
                        nmax = nmax,
                        fbounds = fbounds,
                        δ = δ,
                        demean = demean_ns,
                        sampledc = sampledc,
                        sdev_prop = sdev_prop,
                        sdev_pos = sdev_pos,
                        pnorm = pnorm,
                        K = K
                        )


abstract type AbstractSimulatedAnnealing <: AbstractAlgorithm end
abstract type AbstractCooling end

"""
    struct FixedCooling <: AbstractCooling
        temperature::Float64
    end

Using this cooling scheduler, the algorithm accepts a worser solution with probability ``\\alpha^{-\\Delta}``, where
``\\alpha`` is the parameter temperature and ``\\Delta`` is the absolute difference of the fitness functions between the parent and child.

Through `FixedCooling`, you are basicaully using `Metropolis` algorithm.
"""

struct FixedCooling <: AbstractCooling
    temperature::Real
    function FixedCooling(temperature::Real = 0.0)
        new(temperature)
    end
end

function temperature(cooling::FixedCooling, iter::Int)
    return cooling.temperature
end

struct FixedCoolingTag <: AbstractCooling
    temperature::Real
    order::LaTeXString
    function FixedCoolingTag(temperature::Real = 0.0, order::LaTeXString = L"nothing")
        new(temperature, order)
    end
end

function temperature(cooling::FixedCoolingTag, temp::Real)
    return cooling.temperature
end

"""
    struct SimulatedAnnealing <: AbstractSimulatedAnnealing
        cooling::AbstractCooling
        stop::AbstractStop
        name::LaTeXString
        temperature::Float64
    end
"""
struct SimulatedAnnealing <: AbstractSimulatedAnnealing
    cooling::AbstractCooling
    stop::AbstractStop
    name::LaTeXString
    temperature::Real
    mutation::Mutation

    function SimulatedAnnealing(
        cooling::AbstractCooling;
        stop::AbstractStop = FixedBudget(10^11),
        name::LaTeXString = L"Simulated-Annealing",
        temperature::Real = -1.0,
        mutation::Mutation = KBitFlip(1)
    )
        new(cooling, stop, name, temperature, mutation)
    end
end

function optimize(x, setting::SimulatedAnnealing)
    trace = Trace(setting, x)
    x = initial(x)
    n = length(x)

    for iter ∈ 1:niterations(setting.stop)
        α = float(temperature(setting.cooling, iter))
        y = mutation(x, setting.mutation)
        Δ = fitness(y) - fitness(x)

        if Δ ≥ 0
            if Δ > 0
                if isoptimum(y)
                    record(trace, y, iter, isoptimum(y))
                end
            end
            x = y
            if isoptimum(x)
                return trace
            end
        elseif rand() ≤ α^(Δ)
            # record(trace, y, iter, isoptimum(y))
            x = y
        end
    end
    trace
end


@muladd function calc_twopoint_random!(_dW, sqdt, dW)
  @.. _dW = ifelse(sign(dW) > false, sqdt, -sqdt)
  return nothing
end

@muladd function calc_twopoint_random(sqdt, dW)
  return ifelse(sign(dW) > false, sqdt, -sqdt)
end

function calc_threepoint_random!(_dW, sq3dt, quantile, dW_scaled)
  @. _dW = ifelse(abs(dW_scaled) > -quantile, ifelse(dW_scaled < quantile, -sq3dt, sq3dt), zero(sq3dt))
  return nothing
end

@muladd function calc_threepoint_random(sq3dt, quantile, dW_scaled)
  return ifelse(abs(dW_scaled) > -quantile, ifelse(dW_scaled < quantile, -sq3dt, sq3dt), zero(sq3dt))
end

function Ihat2(cache::Union{DRI1ConstantCache,DRI1Cache}, _dW, _dZ, sqdt, k, l)
  # compute elements of I^_(k,l) which is a mxm matrix
  if k<l
    return (_dW[k]*_dW[l]-sqdt*_dZ[k])/2
  elseif l<k
    return (_dW[k]*_dW[l]+sqdt*_dZ[l])/2
  else
    # l == k
    return (_dW[k]^2-sqdt^2)/2
  end
end

function Ihat2(cache::Union{RSConstantCache,RSCache}, _dW, _dZ, sqdt, k, l)
  # compute elements of I^_(k,l) which is a mxm matrix
  if k<l
    return -_dW[k]*_dZ[l]
  elseif l<k
    return _dW[k]*_dZ[l]
  else
    # l == k
    return zero(_dW[k])
  end
end

function Ihat2(cache::Union{PL1WMConstantCache,PL1WMCache}, _dW, _dZ, sqdt, k, l)
  # compute elements of I^_(k,l) which is a mxm matrix
  if k<l
    return -_dZ[Int(1+1//2*(l-3)*l+k)]
  elseif l<k
    return _dZ[Int(1+1//2*(k-3)*k+l)]
  else
    # l == k
    return -sqdt^2
  end
end

function Ihat2(cache::Union{NONConstantCache,NONCache}, _dW, _dZ, sqdt, k, l)
  # compute elements of I^_(k,l) which is a mxm matrix
  if k<l
    return _dZ[k]
  elseif l<k
    return _dW[l]
  else
    # l == k
    return _dW[k]
  end
end


using RCall
using SigmaRidgeRegression
using LinearAlgebra
using StatsBase
using Random
using Test
using MLJ
using Plots

Random.seed!(1)
n = 200
p = 160
X = randn(n, p)
Xtable = MLJ.table(X);
βs = randn(p)./sqrt(p)
Y = X*βs .+ randn(n)
groups = GroupedFeatures([60;60;40])

_scale = true
_center = true
_param_min_ratio = 0.9
seagull_glasso = TunedSeagull(groups=groups, scale=_scale, center=_center,
                      param_min_ratio=_param_min_ratio)
seagull_mach = machine(seagull_glasso, X, Y)
fit!(seagull_mach)

glasso = GroupLassoRegressor(groups=groups,λ=0.2, scale=_scale, center=_center)
glasso_machine = machine(glasso, Xtable, Y)
fit!(glasso_machine)
λ_max =	SigmaRidgeRegression._default_hyperparameter_maximum(glasso, glasso_machine)

@test seagull_mach.report.λ_max ≈ λ_max atol=0.02

λs_opt = seagull_mach.report.best_λs

multiridge = MultiGroupRidgeRegressor(groups=groups, λs=λs_opt, scale=_scale, center=_center)
multiridge_mach = machine(multiridge, X, Y)
fit!(multiridge_mach)

@test multiridge_mach.fitresult.coef ≈ seagull_mach.fitresult.coef atol=1e-3

group_summary(groups, multiridge_mach.fitresult.coef, norm)
group_summary(groups, seagull_mach.fitresult.coef, norm)

@test MLJ.predict(seagull_mach, Xtable) ≈ MLJ.predict(multiridge_mach, Xtable) atol=0.0001


function equilibrate(A::SparseMatrixCSC{T}; max_iters=10, max_val::T=2., is_symm::Bool=false,
                     blk_sizes::Vector{Int64}=ones(Int64, size(A,1))) where T <: AbstractFloat
    m, n = size(A)
    At = copy(A)
    D1 = ones(T, m)
    D2 = ones(T, n)
    Dr = ones(T, m)
    Dc = ones(T, n)
    @assert all(isfinite.(A.nzval))
    for i in 1:max_iters
        abs_At = abs.(At)
        max_dim1 = maximum(abs_At, dims=1)
        max_dim2 = maximum_by_block(abs_At, blk_sizes)
        if max(maximum(max_dim1), maximum(max_dim2)) < max_val
            break
        end
        Dc = vec(sqrt.(max_dim1))
        Dr = vec(sqrt.(max_dim2))
        bound_below!(Dr, 1e-10)
        bound_below!(Dc, 1e-10)
        D1 = D1 ./ Dr
        D2 = D2 ./ Dc
        At = Diagonal(D1) * A * Diagonal(D2)
    end
    @assert all(isfinite.(D1))
    @assert all(isfinite.(D2))
    return At, D1, D2
end

function bound_below!(x::Vector{T}, eps::T) where T <: AbstractFloat
    @inbounds for i in 1:length(x)
        x[i] = max(x[i], eps)
    end
end

function maximum_by_block(A::AbstractMatrix{T}, blk_sizes::Vector{Int}) where T <: Real
    ptrs = [1, cumsum(blk_sizes)]
    offset = 0
    blk_maxes = Vector{T}(undef, size(A, 1))
    for sz in blk_sizes
        idx = offset+1:offset+sz
        max_blk = maximum(A[idx, :])
        blk_maxes[idx] .= max_blk
        offset += sz
    end
    return blk_maxes
end

struct ProblemEquilibration{T <: AbstractFloat}
    x_scale::Vector{T}
    y_scale::Vector{T}
    z_scale::Vector{T}
    obj_scale::T
    const_scale::T
end

function equilibrate(prob::Problem{T}; max_iters=5) where T
    n, m, p = get_prob_sizes(prob)
    blk_sizes = [ones(Int64, m+1); elem_cone_dims(prob.C)]
    M = [prob.c' spzeros(1,1); prob.A prob.b; prob.G prob.h]
    Mt, D1, D2 = equilibrate(M, max_iters=max_iters, blk_sizes=blk_sizes)

    x_scale = D2[1:end-1]
    const_scale = D2[end]
    y_scale = D1[2:get_num_constrs(prob)+1]
    z_scale = D1[get_num_constrs(prob)+2:end]
    obj_scale = D1[1]

    At = copy(prob.A)
    Gt = copy(prob.G)
    ct = copy(prob.c)
    bt = copy(prob.b)
    ht = copy(prob.h)

    At = Diagonal(y_scale) * prob.A * Diagonal(x_scale)
    Gt = Diagonal(z_scale) * prob.G * Diagonal(x_scale)
    ct = Diagonal(x_scale) * prob.c * obj_scale
    bt = Diagonal(y_scale) * prob.b * const_scale
    ht = Diagonal(z_scale) * prob.h * const_scale

    new_prob = Problem(At, Gt, ct, bt, ht, prob.C)
    equil = ProblemEquilibration(x_scale, y_scale, z_scale, obj_scale, const_scale)
    return new_prob, equil
end

function unscale!(sol::Solution{T}, equil::ProblemEquilibration{T}) where T <: AbstractFloat
    sol.x = sol.x .* equil.x_scale ./ equil.const_scale
    sol.y = sol.y .* equil.y_scale ./ equil.obj_scale
    sol.z = sol.z .* equil.z_scale ./ equil.obj_scale
    sol.s = sol.s ./ equil.z_scale ./ equil.const_scale
    sol.optval = sol.optval / (equil.obj_scale * equil.const_scale)
end


#===============================================================================

    Various random walkers, including:
        ✔ uni- and multivariate Uniform random walk
        ✔ multivariate Gaussian random walk
        ✔ mixture of two Gaussian random walks
    All of the above allow for incorporating a restriction on the positivity
    of any chosen subset of parameter's coordinates.

===============================================================================#

"""
    RandomWalk <: TransitionKernel

Supertype for all random walkers.
"""
abstract type RandomWalk <: TransitionKernel end

#===============================================================================
                            Uniform random walk
===============================================================================#

@doc raw"""
    mutable struct UniformRandomWalk{T,S} <: RandomWalk
        ϵ::T
        pos::S
    end

A uniform random walker that moves all **unrestricted** coordinates according to
the additive rule
```math
x_i+U,
    \quad\mbox{with}
        \quad U\sim\texttt{Unif}([-ϵ_i,ϵ_i]),
            \qquad i\in\texttt{indices},
```
and moves all coordinates **restricted to be positive** according to the rule
```math
x_i\exp(U),
    \quad\mbox{with}
        \quad U\sim\texttt{Unif}([-ϵ_i,ϵ_i]),
            \qquad i\in\texttt{indices}.
```
"""
mutable struct UniformRandomWalk{T,S} <: RandomWalk
    ϵ::T
    pos::S

    function UniformRandomWalk(ϵ, pos=custom_zero(ϵ, Bool))
        @assert all(ϵ .> 0.0)
        new{typeof(ϵ),typeof(pos)}(ϵ, pos)
    end
end

Base.eltype(::UniformRandomWalk{T}) where T = eltype(T)
Base.length(rw::UniformRandomWalk) = length(rw.ϵ)

"""
    Base.rand(rw::UniformRandomWalk, θ)

Sample a new state for a random walker `rw` that is currently in a state `θ`.
"""
Base.rand(rw::UniformRandomWalk, θ) = rand(Random.GLOBAL_RNG, rw, θ)

function Base.rand(
        rng::Random.AbstractRNG,
        rw::UniformRandomWalk,
        θ
    )
    _rand(u) = rand(rng, u)
    U = _rand.(Uniform.(-rw.ϵ, rw.ϵ))
    θ.* ( exp.(U).*rw.pos .+ 1.0.*.!rw.pos ) .+ U.*.!rw.pos
end

function Random.rand!(rng::Random.AbstractRNG, rw::UniformRandomWalk, θ, θ°)
    θ° .= rand(rng, rw, θ)
end

Random.rand!(rw::UniformRandomWalk, θ, θ°) = rand!(Random.GLOBAL_RNG, rw, θ, θ°)

"""
    Distributions.logpdf(rw::UniformRandomWalk, θ, θ°)

Evaluate the log-probability density function for a move of a random walker
from θ to θ°.
"""
function Distributions.logpdf(rw::UniformRandomWalk, θ, θ°)
    mapreduce(
        i->( rw.pos[i] ? -log(2.0*rw.ϵ[i])-log(θ°[i]) : 0.0 ),
        +,
        1:length(rw.ϵ)
    )
end

#===============================================================================
                            Gaussian random walk
===============================================================================#
@doc raw"""
    mutable struct GaussianRandomWalk{T} <: RandomWalk
        Σ::Symmetric{T,Array{T,2}}
        pos::Vector{Bool}
    end

A Gaussian random walker that first transforms all coordinates restricted to be
positive to a log-scale via
```math
θ_i\leftarrow\log(θ_i),\quad i\in\{i: θ_i\texttt{_restricted_to_be_positive}\},
```
then moves according to
```math
θ°∼N(θ, Σ),\quad \forall i,
```
and finally, transforms the restricted coordinates back to the original scale
via
```math
θ°_i\leftarrow\exp(θ°_i),\quad i\in\{i: θ_i\texttt{_restricted_to_be_positive}\}.
```
Restrictions are specified in the vector `pos` (and no restrictions correspond
to all entries in `pos` being `false`).
"""
mutable struct GaussianRandomWalk{T} <: RandomWalk
    Σ::Symmetric{T,Array{T,2}}
    pos::Vector{Bool}

    function GaussianRandomWalk(Σ, pos=nothing)
        @assert size(Σ, 1) == size(Σ, 2)
        Σ = reshape(Σ, (size(Σ, 1), size(Σ, 2)))
        T = eltype(Σ)
        pos = (pos === nothing) ? fill(false, size(Σ, 1)) : pos
        new{T}(Symmetric(Σ), pos)
    end
end

Base.eltype(::GaussianRandomWalk{T}) where T = T
Base.length(rw::GaussianRandomWalk) = size(rw.Σ, 1)
remove_constraints!(rw::GaussianRandomWalk, θ) = (θ[rw.pos] .= log.(θ[rw.pos]))
reimpose_constraints!(rw::GaussianRandomWalk, θ) = (
    θ[rw.pos] .= exp.(θ[rw.pos])
)

Base.rand(rw::GaussianRandomWalk, θ) = rand(Random.GLOBAL_RNG, rw, θ)

function Base.rand(rng::Random.AbstractRNG, rw::GaussianRandomWalk, θ)
    remove_constraints!(rw, θ)
    θ° = rand(MvNormal(θ, rw.Σ))
    reimpose_constraints!(rw, θ°)
    reimpose_constraints!(rw, θ)
    θ°
end

function Random.rand!(rw::GaussianRandomWalk, θ, θ°)
    rand!(Random.GLOBAL_RNG, rw, θ, θ°)
end

function Random.rand!(rng::Random.AbstractRNG, rw::GaussianRandomWalk, θ, θ°)
    θ° .= rand(rng, rw, θ)
end

_logjacobian(rw::GaussianRandomWalk, θ) = -sum(log.(θ[rw.pos]))

function Distributions.logpdf(rw::GaussianRandomWalk, θ, θ°)
    logJ = _logjacobian(rw, θ°)
    remove_constraints!(rw, θ)
    remove_constraints!(rw, θ°)
    lpdf = logpdf(MvNormal(θ, rw.Σ), θ°) + logJ
    reimpose_constraints!(rw, θ)
    reimpose_constraints!(rw, θ°)
    lpdf
end

#===============================================================================
            A mixture of two multidimensional Gaussian random walks
===============================================================================#
@doc raw"""
    mutable struct GaussianRandomWalkMix{T} <: RandomWalk
        gsn_A::GaussianRandomWalk{T}
        gsn_B::GaussianRandomWalk{T}
        λ::Float64
    end

A mixture of two Gaussian random walkers. It performs updates according to
```math
X_{i+1} = X_i + λZ_A + (1-λ)Z_B,
```
where
```math
Z_A ∼ N(0,Σ_A),\quad\mbox{and}\quad Z_B ∼ N(0, Σ_B),
```
(with $X_i$ being already appropriately log-transformed if needed).
"""
mutable struct GaussianRandomWalkMix{T} <: RandomWalk
    gsn_A::GaussianRandomWalk{T}
    gsn_B::GaussianRandomWalk{T}
    λ::Float64

    function GaussianRandomWalkMix(Σ_A, Σ_B, λ=0.5, pos=nothing)
        @assert 0.0 <= λ <= 1.0
        @assert size(Σ_A) == size(Σ_B)
        gsn_A = GaussianRandomWalk(Σ_A, pos)
        gsn_B = GaussianRandomWalk(Σ_B, pos)
        T = eltype(gsn_A)
        new{T}(gsn_A, gsn_B, λ)
    end
end

Base.eltype(::GaussianRandomWalkMix{T}) where T = T
Base.length(rw::GaussianRandomWalkMix) = length(rw.gsn_A)

Base.rand(rw::GaussianRandomWalkMix, θ) = rand(Random.GLOBAL_RNG, rw, θ)

function Base.rand(rng::Random.AbstractRNG, rw::GaussianRandomWalkMix, θ)
    rw_i = pick_kernel(rng, rw)
    rand(rng, rw_i, θ)
end

Random.rand!(rw::GaussianRandomWalkMix, θ, θ°) = rand!(Random.GLOBAL_RNG, rw, θ, θ°)

function Random.rand!(rng::Random.AbstractRNG, rw::GaussianRandomWalkMix, θ, θ°)
    rw_i = pick_kernel(rng, rw)
    rand!(rng, rw_i, θ, θ°)
end

function pick_kernel(rng::Random.AbstractRNG, rw::GaussianRandomWalkMix)
    rand(rng, Bernoulli(rw.λ)) ? rw.gsn_B : rw.gsn_A
end

function Distributions.logpdf(rw::GaussianRandomWalkMix, θ, θᵒ)
    log( (1-rw.λ)*exp(logpdf(rw.gsn_A, θ, θᵒ))
          + rw.λ *exp(logpdf(rw.gsn_B, θ, θᵒ)) )
end


@testset "Test integrated function Nx = 1" begin

Nx = 1
Ne = 50

atol = 1e-6

X = randn(Nx, Ne)

Blist = [ProHermiteBasis(8); PhyHermiteBasis(8); CstProHermiteBasis(8); CstPhyHermiteBasis(8); CstLinProHermiteBasis(8); CstLinPhyHermiteBasis(8)]
for b in Blist
    B = MultiBasis(b, Nx)

    idx = reshape([0; 1; 2; 3; 4], (5, 1))
    truncidx = idx[1:2:end,:]
    Nψ = 5

    coeff = randn(Nψ)

    f = ExpandedFunction(B, idx, coeff)
    R = IntegratedFunction(f)

    # Test evaluate
    ψt = zeros(Ne)

    for i=1:Ne
        x = X[:,i]
        ψt[i] = R.f(vcat(x[1:end-1], 0.0)) +
                quadgk(t->R.g(ForwardDiff.gradient(y->R.f(y), vcat(x[1:end-1],t))[end]), 0, x[end])[1]
    end

    ψ = evaluate(R, X)

    @test norm(ψ - ψt)<atol

    # Test grad_x
    dψ = grad_x(R, X)

    function Gradient(x)
        return ForwardDiff.gradient(y->begin
                               y[end] = 0.0
                               return  R.f(y)
            end, x) + quadgk(t->ForwardDiff.gradient(
                      z->R.g(ForwardDiff.gradient(y->R.f(y), vcat(z[1:end-1],t))[end]),
                x), 0.0, x[end])[1] + vcat(zeros(size(x,1)-1), R.g(ForwardDiff.gradient(y->R.f(y), x)[end]))
    end

    for i=1:Ne
        xi = X[:,i]
        dψt = Gradient(xi)
        @test norm(dψ[i,:] - dψt)<1e-8
    end

    # Test hess_x
    H = hess_x(R, X)
    Ht = zeros(Nx, Nx)
    function Hessian!(H, x)
        fill!(H, 0.0)

        H[1:Nx-1,1:Nx-1] .= ForwardDiff.hessian(y->begin
                            y[Nx] = 0.0
                            return R.f(y) end, x)[1:Nx-1,1:Nx-1]

        # Hessian of the integral term
        H[1:Nx-1,1:Nx-1] .+= (quadgk(t -> ForwardDiff.hessian(
        z-> R.g(ForwardDiff.gradient(y->R.f(y), vcat(z[1:end-1], t))[end]), x),
                0.0, x[end])[1])[1:Nx-1,1:Nx-1]

        # H[Nx,:] and H[:,Nx]
        H[:, Nx] .= ForwardDiff.gradient(z->R.g(ForwardDiff.gradient(y->R.f(y), z)[end]), x)
        H[Nx, :] .= H[:, Nx]
        return H
    end

    @inbounds for i=1:Ne
        @test norm(H[i,:,:] - Hessian!(Ht, X[:,i]))<atol
    end

    # Test grad_xd
    gdψ =  grad_xd(R, X)

    gdψt = zeros(Ne)

    for i=1:Ne
        gdψt[i] = R.g(ForwardDiff.gradient(R.f, X[:,i])[end])
    end

    @test norm(gdψ - gdψt) <atol

    # Test grad_coeff_grad_xd
    dψ_xd_dc = grad_coeff_grad_xd(R, X)
    dψ_xd_dct = zeros(Ne, Nψ)

    ∂kf(y) = ForwardDiff.gradient(R.f, y)[end]

    for j=1:Nψ
        fj = MultiFunction(R.f.B, f.idx[j,:])
        ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]
        for i=1:Ne
            dψ_xd_dct[i,j] = ∂kfj(X[:,i])*grad_x(R.g, ∂kf(X[:,i]))
        end
    end

    @test norm(dψ_xd_dc - dψ_xd_dct)<atol

    # Test hess_coeff_grad_xd
    d2ψ_xd_dc = hess_coeff_grad_xd(R, X)
    d2ψ_xd_dct = zeros(Ne, Nψ, Nψ)

    for j=1:Nψ
        fj = MultiFunction(R.f.B, f.idx[j,:])
        ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]
        for l=1:Nψ
        fl = MultiFunction(R.f.B, f.idx[l,:])
        ∂kfl(y) = ForwardDiff.gradient(fl, y)[end]

            for i=1:Ne
                d2ψ_xd_dct[i,j,l] = 0.0*grad_x(R.g, ∂kf(X[:,i])) +
                    hess_x(R.g, ∂kf(X[:,i]))*∂kfj(X[:,i])*∂kfl(X[:,i])
            end
        end
    end

    @test norm(d2ψ_xd_dc - d2ψ_xd_dct)<atol

    # Test integrate_xd
    intψt = zeros(Ne)

    for i=1:Ne
        xi = X[:,i]
        intψt[i] = quadgk(t->R.g(ForwardDiff.gradient(y->R.f(y), vcat(xi[1:end-1],t))[end]), 0, xi[end])[1]
    end

    intψ = integrate_xd(R, X)

    @test norm(intψt - intψ)<atol

    # Test grad_coeff_integrate_xd

    dcintψt = zeros(Ne, Nψ)
    xi = zeros(Nx)

    for j=1:Nψ
        fj = MultiFunction(R.f.B, f.idx[j,:])
        ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]
        for i=1:Ne
            xi .= X[:,i]
            dcintψt[i,j] = quadgk(t->∂kfj(vcat(xi[1:end-1],t))*grad_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]
        end
    end

    dcintψ = grad_coeff_integrate_xd(R, X)
    @test norm(dcintψ - dcintψt)<atol

    # Test hess_coeff_integrate_xd
    d2cintψt = zeros(Ne, Nψ, Nψ)
    xi = zeros(Nx)
    for j=1:Nψ
        fj = MultiFunction(R.f.B, f.idx[j,:])
        ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]

        for l=1:Nψ

            fl = MultiFunction(R.f.B, f.idx[l,:])
            ∂kfl(y) = ForwardDiff.gradient(fl, y)[end]

            for i=1:Ne
                xi .= X[:,i]
    #                 @show quadgk(t->∂kfj(vcat(xi[1:end-1],t))*∂kfl(vcat(xi[1:end-1],t))*hess_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]
                    d2cintψt[i,j,l] = quadgk(t->∂kfj(vcat(xi[1:end-1],t))*∂kfl(vcat(xi[1:end-1],t))*hess_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]
            end
        end
    end

    d2cintψ = hess_coeff_integrate_xd(R, X)

    @test norm(d2cintψt - d2cintψ)<atol

    # Test grad_coeff
    dcRt = zeros(Ne, Nψ)

    X0 = deepcopy(X)
    X0[end,:] .= zeros(Ne)
    # ∂_c f(x_{1:k-1},0)
    dcRt += evaluate_basis(R.f, X0)
    xi = zeros(Nx)

    for j=1:Nψ
        fj = MultiFunction(R.f.B, f.idx[j,:])
        ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]

        for i=1:Ne
            xi .= X[:,i]
            dcRt[i,j] += quadgk(t->∂kfj(vcat(xi[1:end-1],t))*grad_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]

        end
    end

    dcR = grad_coeff(R, X)

    @test norm(dcRt - dcR)<atol

    # Test hess_coeff
    d2cRt = zeros(Ne, Nψ, Nψ)

    xi = zeros(Nx)

    for j=1:Nψ
        fj = MultiFunction(R.f.B, f.idx[j,:])
        ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]

        for l=1:Nψ

            fl = MultiFunction(R.f.B, f.idx[l,:])
            ∂kfl(y) = ForwardDiff.gradient(fl, y)[end]

            for i=1:Ne
                xi .= X[:,i]
    #                 @show quadgk(t->∂kfj(vcat(xi[1:end-1],t))*∂kfl(vcat(xi[1:end-1],t))*hess_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]
                    d2cRt[i,j,l] = quadgk(t->∂kfj(vcat(xi[1:end-1],t))*∂kfl(vcat(xi[1:end-1],t))*hess_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]
            end
        end
    end

    d2cR = hess_coeff(R, X)

    @test norm(d2cRt - d2cR)<atol

    end
end

@testset "Test integrated function Nx = 2" begin

Nx = 2
Ne = 50

atol = 1e-6

X = randn(Nx, Ne)

Blist = [ProHermiteBasis(8); PhyHermiteBasis(8); CstProHermiteBasis(8); CstPhyHermiteBasis(8); CstLinProHermiteBasis(8); CstLinPhyHermiteBasis(8)]
for b in Blist
    B = MultiBasis(b, Nx)

    idx = [0 0; 0 1; 1 0; 1 1; 1 2]
    truncidx = idx[1:2:end,:]
    Nψ = 5

    coeff = randn(Nψ)

    f = ExpandedFunction(B, idx, coeff)
    R = IntegratedFunction(f)

    # Test evaluate
    ψt = zeros(Ne)

    for i=1:Ne
        x = X[:,i]
        ψt[i] = R.f(vcat(x[1:end-1], 0.0)) +
                quadgk(t->R.g(ForwardDiff.gradient(y->R.f(y), vcat(x[1:end-1],t))[end]), 0, x[end])[1]
    end

    ψ = evaluate(R, X)

    @test norm(ψ - ψt)<atol

    # Test grad_x
    dψ = grad_x(R, X)

    function Gradient(x)
        return ForwardDiff.gradient(y->begin
                               y[end] = 0.0
                               return  R.f(y)
            end, x) + quadgk(t->ForwardDiff.gradient(
                      z->R.g(ForwardDiff.gradient(y->R.f(y), vcat(z[1:end-1],t))[end]),
                x), 0.0, x[end])[1] + vcat(zeros(size(x,1)-1), R.g(ForwardDiff.gradient(y->R.f(y), x)[end]))
    end

    for i=1:Ne
        xi = X[:,i]
        dψt = Gradient(xi)
        @test norm(dψ[i,:] - dψt)<1e-8
    end

    # Test hess_x
    H = hess_x(R, X)
    Ht = zeros(Nx, Nx)
    function Hessian!(H, x)
        fill!(H, 0.0)

        H[1:Nx-1,1:Nx-1] .= ForwardDiff.hessian(y->begin
                            y[Nx] = 0.0
                            return R.f(y) end, x)[1:Nx-1,1:Nx-1]

        # Hessian of the integral term
        H[1:Nx-1,1:Nx-1] .+= (quadgk(t -> ForwardDiff.hessian(
        z-> R.g(ForwardDiff.gradient(y->R.f(y), vcat(z[1:end-1], t))[end]), x),
                0.0, x[end])[1])[1:Nx-1,1:Nx-1]

        # H[Nx,:] and H[:,Nx]
        H[:, Nx] .= ForwardDiff.gradient(z->R.g(ForwardDiff.gradient(y->R.f(y), z)[end]), x)
        H[Nx, :] .= H[:, Nx]
        return H
    end

    @inbounds for i=1:Ne
        @test norm(H[i,:,:] - Hessian!(Ht, X[:,i]))<atol
    end

    # Test grad_xd
    gdψ =  grad_xd(R, X)

    gdψt = zeros(Ne)

    for i=1:Ne
        gdψt[i] = R.g(ForwardDiff.gradient(R.f, X[:,i])[end])
    end

    @test norm(gdψ - gdψt) <atol

    # Test grad_coeff_grad_xd
    dψ_xd_dc = grad_coeff_grad_xd(R, X)
    dψ_xd_dct = zeros(Ne, Nψ)

    ∂kf(y) = ForwardDiff.gradient(R.f, y)[end]

    for j=1:Nψ
        fj = MultiFunction(R.f.B, f.idx[j,:])
        ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]
        for i=1:Ne
            dψ_xd_dct[i,j] = ∂kfj(X[:,i])*grad_x(R.g, ∂kf(X[:,i]))
        end
    end

    @test norm(dψ_xd_dc - dψ_xd_dct)<atol

    # Test hess_coeff_grad_xd
    d2ψ_xd_dc = hess_coeff_grad_xd(R, X)
    d2ψ_xd_dct = zeros(Ne, Nψ, Nψ)

    for j=1:Nψ
        fj = MultiFunction(R.f.B, f.idx[j,:])
        ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]
        for l=1:Nψ
        fl = MultiFunction(R.f.B, f.idx[l,:])
        ∂kfl(y) = ForwardDiff.gradient(fl, y)[end]

            for i=1:Ne
                d2ψ_xd_dct[i,j,l] = 0.0*grad_x(R.g, ∂kf(X[:,i])) +
                    hess_x(R.g, ∂kf(X[:,i]))*∂kfj(X[:,i])*∂kfl(X[:,i])
            end
        end
    end

    @test norm(d2ψ_xd_dc - d2ψ_xd_dct)<atol

    # Test integrate_xd
    intψt = zeros(Ne)

    for i=1:Ne
        xi = X[:,i]
        intψt[i] = quadgk(t->R.g(ForwardDiff.gradient(y->R.f(y), vcat(xi[1:end-1],t))[end]), 0, xi[end])[1]
    end

    intψ = integrate_xd(R, X)

    @test norm(intψt - intψ)<atol

    # Test grad_coeff_integrate_xd

    dcintψt = zeros(Ne, Nψ)
    xi = zeros(Nx)

    for j=1:Nψ
        fj = MultiFunction(R.f.B, f.idx[j,:])
        ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]
        for i=1:Ne
            xi .= X[:,i]
            dcintψt[i,j] = quadgk(t->∂kfj(vcat(xi[1:end-1],t))*grad_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]
        end
    end

    dcintψ = grad_coeff_integrate_xd(R, X)
    @test norm(dcintψ - dcintψt)<atol

    # Test hess_coeff_integrate_xd
    d2cintψt = zeros(Ne, Nψ, Nψ)
    xi = zeros(Nx)
    for j=1:Nψ
        fj = MultiFunction(R.f.B, f.idx[j,:])
        ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]

        for l=1:Nψ

            fl = MultiFunction(R.f.B, f.idx[l,:])
            ∂kfl(y) = ForwardDiff.gradient(fl, y)[end]

            for i=1:Ne
                xi .= X[:,i]
    #                 @show quadgk(t->∂kfj(vcat(xi[1:end-1],t))*∂kfl(vcat(xi[1:end-1],t))*hess_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]
                    d2cintψt[i,j,l] = quadgk(t->∂kfj(vcat(xi[1:end-1],t))*∂kfl(vcat(xi[1:end-1],t))*hess_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]
            end
        end
    end

    d2cintψ = hess_coeff_integrate_xd(R, X)

    @test norm(d2cintψt - d2cintψ)<atol

    # Test grad_coeff
    dcRt = zeros(Ne, Nψ)

    X0 = deepcopy(X)
    X0[end,:] .= zeros(Ne)
    # ∂_c f(x_{1:k-1},0)
    dcRt += evaluate_basis(R.f, X0)
    xi = zeros(Nx)

    for j=1:Nψ
        fj = MultiFunction(R.f.B, f.idx[j,:])
        ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]

        for i=1:Ne
            xi .= X[:,i]
            dcRt[i,j] += quadgk(t->∂kfj(vcat(xi[1:end-1],t))*grad_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]

        end
    end

    dcR = grad_coeff(R, X)

    @test norm(dcRt - dcR)<atol

    # Test hess_coeff
    d2cRt = zeros(Ne, Nψ, Nψ)

    xi = zeros(Nx)

    for j=1:Nψ
        fj = MultiFunction(R.f.B, f.idx[j,:])
        ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]

        for l=1:Nψ

            fl = MultiFunction(R.f.B, f.idx[l,:])
            ∂kfl(y) = ForwardDiff.gradient(fl, y)[end]

            for i=1:Ne
                xi .= X[:,i]
    #                 @show quadgk(t->∂kfj(vcat(xi[1:end-1],t))*∂kfl(vcat(xi[1:end-1],t))*hess_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]
                    d2cRt[i,j,l] = quadgk(t->∂kfj(vcat(xi[1:end-1],t))*∂kfl(vcat(xi[1:end-1],t))*hess_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]
            end
        end
    end

    d2cR = hess_coeff(R, X)

    @test norm(d2cRt - d2cR)<atol

    end
end

@testset "Test integrated function Nx = 3" begin

    Nx = 3
    Ne = 50

    X = randn(Nx, Ne)

    atol = 1e-6

    Blist = [ProHermiteBasis(8); PhyHermiteBasis(8); CstProHermiteBasis(8); CstPhyHermiteBasis(8); CstLinProHermiteBasis(8); CstLinPhyHermiteBasis(8)]
    for b in Blist
        B = MultiBasis(b, Nx)

        # idx = [0 0; 0 1; 1 0; 1 1; 1 2]
        idx = [0 0 0 ;0  0 1; 0 1 0; 1 0 0;1 1 0; 0 1 1; 1 0 1; 1 1 1; 1 2 0; 2 1 0; 2 2 2]

        truncidx = idx[1:2:end,:]
        Nψ = 11

        coeff = randn(Nψ)

        f = ExpandedFunction(B, idx, coeff)
        R = IntegratedFunction(f)

        # Test grad_x
        dψ = grad_x(R, X)

        function Gradient(x)
            return ForwardDiff.gradient(y->begin
                                   y[end] = 0.0
                                   return  R.f(y)
                end, x) + quadgk(t->ForwardDiff.gradient(
                          z->R.g(ForwardDiff.gradient(y->R.f(y), vcat(z[1:end-1],t))[end]),
                    x), 0.0, x[end])[1] + vcat(zeros(size(x,1)-1), R.g(ForwardDiff.gradient(y->R.f(y), x)[end]))
        end

        for i=1:Ne
            xi = X[:,i]
            dψt = Gradient(xi)
            @test norm(dψ[i,:] - dψt)<atol
        end

        # Test grad_xd
        gdψ =  grad_xd(R, X)

        gdψt = zeros(Ne)

        for i=1:Ne
            gdψt[i] = R.g(ForwardDiff.gradient(R.f, X[:,i])[end])
        end

        @test norm(gdψ - gdψt) <atol

        # Test grad_coeff_grad_xd
        dψ_xd_dc = grad_coeff_grad_xd(R, X)
        dψ_xd_dct = zeros(Ne, Nψ)

        ∂kf(y) = ForwardDiff.gradient(R.f, y)[end]

        for j=1:Nψ
            fj = MultiFunction(R.f.B, f.idx[j,:])
            ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]
            for i=1:Ne
                dψ_xd_dct[i,j] = ∂kfj(X[:,i])*grad_x(R.g, ∂kf(X[:,i]))
            end
        end

        @test norm(dψ_xd_dc - dψ_xd_dct)<atol

        # Test hess_coeff_grad_xd
        d2ψ_xd_dc = hess_coeff_grad_xd(R, X)
        d2ψ_xd_dct = zeros(Ne, Nψ, Nψ)

        for j=1:Nψ
            fj = MultiFunction(R.f.B, f.idx[j,:])
            ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]
            for l=1:Nψ
            fl = MultiFunction(R.f.B, f.idx[l,:])
            ∂kfl(y) = ForwardDiff.gradient(fl, y)[end]

                for i=1:Ne
                    d2ψ_xd_dct[i,j,l] = 0.0*grad_x(R.g, ∂kf(X[:,i])) +
                        hess_x(R.g, ∂kf(X[:,i]))*∂kfj(X[:,i])*∂kfl(X[:,i])
                end
            end
        end

        @test norm(d2ψ_xd_dc - d2ψ_xd_dct)<atol

        # Test integrate_xd
        intψt = zeros(Ne)

        for i=1:Ne
            xi = X[:,i]
            intψt[i] = quadgk(t->R.g(ForwardDiff.gradient(y->R.f(y), vcat(xi[1:end-1],t))[end]), 0, xi[end])[1]
        end

        intψ = integrate_xd(R, X)

        @test norm(intψt - intψ)<atol

        # Test grad_coeff_integrate_xd

        dcintψt = zeros(Ne, Nψ)
        xi = zeros(Nx)
        for j=1:Nψ
            fj = MultiFunction(R.f.B, f.idx[j,:])
            ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]
            for i=1:Ne
                xi .= X[:,i]
                dcintψt[i,j] = quadgk(t->∂kfj(vcat(xi[1:end-1],t))*grad_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]
            end
        end

        dcintψ = grad_coeff_integrate_xd(R, X)
        @test norm(dcintψ - dcintψt)<atol

        # Test hess_coeff_integrate_xd
        d2cintψt = zeros(Ne, Nψ, Nψ)
        xi = zeros(Nx)
        for j=1:Nψ
            fj = MultiFunction(R.f.B, f.idx[j,:])
            ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]

            for l=1:Nψ

                fl = MultiFunction(R.f.B, f.idx[l,:])
                ∂kfl(y) = ForwardDiff.gradient(fl, y)[end]

                for i=1:Ne
                    xi .= X[:,i]
        #                 @show quadgk(t->∂kfj(vcat(xi[1:end-1],t))*∂kfl(vcat(xi[1:end-1],t))*hess_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]
                        d2cintψt[i,j,l] = quadgk(t->∂kfj(vcat(xi[1:end-1],t))*∂kfl(vcat(xi[1:end-1],t))*hess_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]
                end
            end
        end

        d2cintψ = hess_coeff_integrate_xd(R, X)

        @test norm(d2cintψt - d2cintψ)<atol

        # Test grad_coeff
        dcRt = zeros(Ne, Nψ)

        X0 = deepcopy(X)
        X0[end,:] .= zeros(Ne)
        # ∂_c f(x_{1:k-1},0)
        dcRt += evaluate_basis(R.f, X0)
        xi = zeros(Nx)

        for j=1:Nψ
            fj = MultiFunction(R.f.B, f.idx[j,:])
            ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]
            for i=1:Ne
                xi .= X[:,i]
                dcRt[i,j] += quadgk(t->∂kfj(vcat(xi[1:end-1],t))*grad_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]

            end
        end

        dcR = grad_coeff(R, X)

        @test norm(dcRt - dcR)<atol

        # Test hess_coeff
        d2cRt = zeros(Ne, Nψ, Nψ)

        xi = zeros(Nx)

        for j=1:Nψ
            fj = MultiFunction(R.f.B, f.idx[j,:])
            ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]

            for l=1:Nψ

                fl = MultiFunction(R.f.B, f.idx[l,:])
                ∂kfl(y) = ForwardDiff.gradient(fl, y)[end]

                for i=1:Ne
                    xi .= X[:,i]
        #                 @show quadgk(t->∂kfj(vcat(xi[1:end-1],t))*∂kfl(vcat(xi[1:end-1],t))*hess_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]
                    d2cRt[i,j,l] = quadgk(t->∂kfj(vcat(xi[1:end-1],t))*∂kfl(vcat(xi[1:end-1],t))*hess_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]
                end
            end
        end

        d2cR = hess_coeff(R, X)

        @test norm(d2cRt - d2cR)<atol
    end
end

@testset "Test integrated function Nx = 4" begin

    Nx = 4
    Ne = 50

    X = randn(Nx, Ne)

    atol = 5e-6

    Blist = [ProHermiteBasis(8); PhyHermiteBasis(8); CstProHermiteBasis(8); CstPhyHermiteBasis(8); CstLinProHermiteBasis(8); CstLinPhyHermiteBasis(8)]
    for b in Blist
        B = MultiBasis(b, Nx)

        # idx = [0 0; 0 1; 1 0; 1 1; 1 2]
        idx = [0 0 0 0; 0 2 0 1; 0 2 3 0; 2 0 2 1; 0 0 1 2; 1 2 0 2;3 2 2 2]

        truncidx = idx[1:2:end,:]
        Nψ = 7

        coeff = randn(Nψ)
        f = ExpandedFunction(B, idx, coeff)
        R = IntegratedFunction(f)

        # Test grad_x
        dψ = grad_x(R, X)

        function Gradient(x)
            return ForwardDiff.gradient(y->begin
                                   y[end] = 0.0
                                   return  R.f(y)
                end, x) + quadgk(t->ForwardDiff.gradient(
                          z->R.g(ForwardDiff.gradient(y->R.f(y), vcat(z[1:end-1],t))[end]),
                    x), 0.0, x[end])[1] + vcat(zeros(size(x,1)-1), R.g(ForwardDiff.gradient(y->R.f(y), x)[end]))
        end

        for i=1:Ne
            xi = X[:,i]
            dψt = Gradient(xi)
            @test norm(dψ[i,:] - dψt)<atol
        end

        # Test hess_x
        H = hess_x(R, X)
        Ht = zeros(Nx, Nx)
        function Hessian!(H, x)
            fill!(H, 0.0)

            H[1:Nx-1,1:Nx-1] .= ForwardDiff.hessian(y->begin
                                y[Nx] = 0.0
                                return R.f(y) end, x)[1:Nx-1,1:Nx-1]

            # Hessian of the integral term
            H[1:Nx-1,1:Nx-1] .+= (quadgk(t -> ForwardDiff.hessian(
            z-> R.g(ForwardDiff.gradient(y->R.f(y), vcat(z[1:end-1], t))[end]), x),
                    0.0, x[end])[1])[1:Nx-1,1:Nx-1]

            # H[Nx,:] and H[:,Nx]
            H[:, Nx] .= ForwardDiff.gradient(z->R.g(ForwardDiff.gradient(y->R.f(y), z)[end]), x)
            H[Nx, :] .= H[:, Nx]
            return H
        end

        @inbounds for i=1:Ne
            @test norm(H[i,:,:] - Hessian!(Ht, X[:,i]))<atol
        end

        # Test grad_xd
        gdψ =  grad_xd(R, X)

        gdψt = zeros(Ne)

        for i=1:Ne
            gdψt[i] = R.g(ForwardDiff.gradient(R.f, X[:,i])[end])
        end

        @test norm(gdψ - gdψt) <atol

        # Test grad_coeff_grad_xd
        dψ_xd_dc = grad_coeff_grad_xd(R, X)
        dψ_xd_dct = zeros(Ne, Nψ)

        ∂kf(y) = ForwardDiff.gradient(R.f, y)[end]

        for j=1:Nψ
            fj = MultiFunction(R.f.B, f.idx[j,:])
            ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]
            for i=1:Ne
                dψ_xd_dct[i,j] = ∂kfj(X[:,i])*grad_x(R.g, ∂kf(X[:,i]))
            end
        end

        @test norm(dψ_xd_dc - dψ_xd_dct)<atol

        # Test hess_coeff_grad_xd
        d2ψ_xd_dc = hess_coeff_grad_xd(R, X)
        d2ψ_xd_dct = zeros(Ne, Nψ, Nψ)

        for j=1:Nψ
            fj = MultiFunction(R.f.B, f.idx[j,:])
            ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]
            for l=1:Nψ
            fl = MultiFunction(R.f.B, f.idx[l,:])
            ∂kfl(y) = ForwardDiff.gradient(fl, y)[end]

                for i=1:Ne
                    d2ψ_xd_dct[i,j,l] = 0.0*grad_x(R.g, ∂kf(X[:,i])) +
                        hess_x(R.g, ∂kf(X[:,i]))*∂kfj(X[:,i])*∂kfl(X[:,i])
                end
            end
        end

        @test norm(d2ψ_xd_dc - d2ψ_xd_dct)<atol

        # Test integrate_xd
        intψt = zeros(Ne)

        for i=1:Ne
            xi = X[:,i]
            intψt[i] = quadgk(t->R.g(ForwardDiff.gradient(y->R.f(y), vcat(xi[1:end-1],t))[end]), 0, xi[end])[1]
        end

        intψ = integrate_xd(R, X)

        @test norm(intψt - intψ)<atol

        # Test grad_coeff_integrate_xd

        dcintψt = zeros(Ne, Nψ)
        xi = zeros(Nx)
        for j=1:Nψ
            fj = MultiFunction(R.f.B, f.idx[j,:])
            ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]
            for i=1:Ne
                xi .= X[:,i]
                dcintψt[i,j] = quadgk(t->∂kfj(vcat(xi[1:end-1],t))*grad_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]
            end
        end

        dcintψ = grad_coeff_integrate_xd(R, X)
        @test norm(dcintψ - dcintψt)<atol

        # Test hess_coeff_integrate_xd
        d2cintψt = zeros(Ne, Nψ, Nψ)
        xi = zeros(Nx)
        for j=1:Nψ
            fj = MultiFunction(R.f.B, f.idx[j,:])
            ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]
            for l=1:Nψ
                fl = MultiFunction(R.f.B, f.idx[l,:])
                ∂kfl(y) = ForwardDiff.gradient(fl, y)[end]
                for i=1:Ne
                    xi .= X[:,i]
                    d2cintψt[i,j,l] = quadgk(t->∂kfj(vcat(xi[1:end-1],t))*∂kfl(vcat(xi[1:end-1],t))*hess_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]
                end
            end
        end

        d2cintψ = hess_coeff_integrate_xd(R, X)

        @test norm(d2cintψt - d2cintψ)<atol

        # Test grad_coeff
        dcRt = zeros(Ne, Nψ)

        X0 = deepcopy(X)
        X0[end,:] .= zeros(Ne)
        # ∂_c f(x_{1:k-1},0)
        dcRt += evaluate_basis(R.f, X0)
        xi = zeros(Nx)

        for j=1:Nψ
            fj = MultiFunction(R.f.B, f.idx[j,:])
            ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]
            for i=1:Ne
                xi .= X[:,i]
                dcRt[i,j] += quadgk(t->∂kfj(vcat(xi[1:end-1],t))*grad_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]

            end
        end

        dcR = grad_coeff(R, X)

        @test norm(dcRt - dcR)<atol

        # Test hess_coeff
        d2cRt = zeros(Ne, Nψ, Nψ)

        xi = zeros(Nx)

        for j=1:Nψ
            fj = MultiFunction(R.f.B, f.idx[j,:])
            ∂kfj(y) = ForwardDiff.gradient(fj, y)[end]

            for l=1:Nψ

                fl = MultiFunction(R.f.B, f.idx[l,:])
                ∂kfl(y) = ForwardDiff.gradient(fl, y)[end]

                for i=1:Ne
                    xi .= X[:,i]
        #                 @show quadgk(t->∂kfj(vcat(xi[1:end-1],t))*∂kfl(vcat(xi[1:end-1],t))*hess_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]
                        d2cRt[i,j,l] = quadgk(t->∂kfj(vcat(xi[1:end-1],t))*∂kfl(vcat(xi[1:end-1],t))*hess_x(R.g, ∂kf(vcat(xi[1:end-1],t))),0,xi[end])[1]
                end
            end
        end

        d2cR = hess_coeff(R, X)

        @test norm(d2cRt - d2cR)<atol
    end
end


import CompactBasisFunctions: Lagrange
import LinearAlgebra

@doc raw"""
The s-stage Lobatto nodes are defined as the roots of the following polynomial of degree $s$:
```math
\frac{d^{s-2}}{dx^{s-2}} \big( (x - x^2)^{s-1} \big) .
```
"""
function get_lobatto_nodes(::Type{T}, s) where {T}
    if s == 1
        throw(ErrorException("Lobatto nodes for one stage are not defined."))
    end

    D(k) = Polynomials.derivative(Polynomial(T[0, 1, -1])^(k-1), k-2)
    c = sort(T.(Polynomials.roots(D(s))))
    c[begin] = 0; c[end] = 1; c
end

get_lobatto_nodes(s) = get_lobatto_nodes(BigFloat, s)

@doc raw"""
The Lobatto weights can be explicitly computed by the formula
```math
b_j = \frac{1}{s (s-1) P_{s-1}(2 c_j - 1)^2} \qquad j = 1 , \, ... , \, s ,
```
where $P_k$ is the $k$th Legendre polynomial, given by
```math
P_k (x) = \frac{1}{k! 2^k} \big( \frac{d^k}{dx^k} (x^2 - 1)^k \big) .
```
"""
function get_lobatto_weights(::Type{T}, s) where {T}
    if s == 1
        throw(ErrorException("Lobatto weights for one stage are not defined."))
    end

    P(k,x) = Polynomials.derivative(Polynomial(T[-1, 0, 1])^k, k)(x) / LinearAlgebra.factorial(k) / 2^k
    c = get_lobatto_nodes(T,s)
    b = [ 1 / ( s*(s-1) * P(s-1, 2c[i] - 1)^2 ) for i in 1:s ]
end

get_lobatto_weights(s) = get_lobatto_weights(BigFloat, s)

"""
```julia
get_lobatto_nullvector(::Type, s; normalize=false)
get_lobatto_nullvector(s; kwargs...)
```

Computes the nullvector of the matrix containing the derivatives of the
Lagrange basis on the `s` Lobatto nodes evaluated on these nodes.
"""
function get_lobatto_nullvector(::Type{T}, s; normalize=false) where {T}
    if s == 1
        throw(ErrorException("Lobatto nullvector for one stage is not defined."))
    end

    q = get_lobatto_nodes(s)
    l = Lagrange(q)
    v = [l'[x, j] for x in q, j in eachindex(l)]
    w = LinearAlgebra.nullspace(v')[:,begin]
    normalize ? T.(LinearAlgebra.normalize(w) .* sign(w[begin])) : T.(w)
end

get_lobatto_nullvector(s; kwargs...) = get_lobatto_nullvector(Float64, s; kwargs...)

@doc raw"""
The Lobatto IIIA coefficients are implicitly given by the so-called simplifying assumption $C(s)$:
```math
\sum \limits_{j=1}^{s} a_{ij} c_{j}^{k-1} = \frac{c_i^k}{k}  \qquad i = 1 , \, ... , \, s , \; k = 1 , \, ... , \, s .
```
"""
function get_lobatto_a_coefficients(::Type{T}, s) where {T}
    if s == 1
        throw(ErrorException("Lobatto IIIA coefficients for one stage are not defined."))
    end
    solve_simplifying_assumption_c(get_lobatto_nodes(T,s))
end

@doc raw"""
The Lobatto IIIB coefficients are implicitly given by the so-called simplifying assumption $D(s)$:
```math
\sum \limits_{i=1}^{s} b_i c_{i}^{k-1} a_{ij} = \frac{b_j}{k} ( 1 - c_j^k)  \qquad j = 1 , \, ... , \, s , \; k = 1 , \, ... , \, s .
```
"""
function get_lobatto_b_coefficients(::Type{T}, s) where {T}
    if s == 1
        throw(ErrorException("Lobatto IIIB coefficients for one stage are not defined."))
    end

    solve_simplifying_assumption_d(get_lobatto_weights(T,s), get_lobatto_nodes(T,s))
end

@doc raw"""
The Lobatto IIIC coefficients are determined by setting $a_{i,1} = b_1$ and
solving the so-called simplifying assumption $C(s-1)$, given by
```math
\sum \limits_{j=1}^{s} a_{ij} c_{j}^{k-1} = \frac{c_i^k}{k}  \qquad i = 1 , \, ... , \, s , \; k = 1 , \, ... , \, s-1 ,
```
for $a_{i,j}$ with $i = 1, ..., s$ and $j = 2, ..., s$.
"""
function get_lobatto_c_coefficients(::Type{T}, s) where {T}
    if s == 1
        throw(ErrorException("Lobatto IIIC coefficients for one stage are not defined."))
    end

    b = get_lobatto_weights(T,s)
    c = get_lobatto_nodes(T,s)
    M = [ c[j]^(k-1) for k in 1:s-1, j in 2:s ]

    row(i) = begin
        r = [ c[i]^k / T(k) - c[1]^(k-1) * b[1] for k in 1:s-1 ]
        M \ r
    end

    hcat(b[1] * ones(T,s), vcat([transpose(row(i)) for i in 1:s]...))
end

@doc raw"""
The Lobatto IIIC̄ coefficients are determined by setting $a_{i,s} = 0$ and
solving the so-called simplifying assumption $C(s-1)$, given by
```math
\sum \limits_{j=1}^{s} a_{ij} c_{j}^{k-1} = \frac{c_i^k}{k}  \qquad i = 1 , \, ... , \, s , \; k = 1 , \, ... , \, s-1 ,
```
for $a_{i,j}$ with $i = 1, ..., s$ and $j = 1, ..., s-1$.
"""
function get_lobatto_c̄_coefficients(::Type{T}, s) where {T}
    if s == 1
        throw(ErrorException("Lobatto IIIC̄ coefficients for one stage are not defined."))
    end

    c = get_lobatto_nodes(T,s)
    M = [ c[j]^(k-1) for k in 1:s-1, j in 1:s-1 ]

    row(i) = begin
        r = [ c[i]^k / T(k) for k in 1:s-1 ]
        M \ r
    end

    hcat(vcat([transpose(row(i)) for i in 1:s]...), zeros(T,s))
end

get_lobatto_d_coefficients(::Type{T}, s) where {T} = (get_lobatto_c_coefficients(T,s) .+ get_lobatto_c̄_coefficients(T,s)) ./ 2

get_lobatto_e_coefficients(::Type{T}, s) where {T} = (get_lobatto_a_coefficients(T,s) .+ get_lobatto_b_coefficients(T,s)) ./ 2

function get_lobatto_f_coefficients(::Type{T}, s) where {T}
    if s == 1
        throw(ErrorException("Lobatto IIIF coefficients for one stage are not defined."))
    end

    c = get_lobatto_nodes(T,s)
    M = [ 1 / T(k + j - 1) for k in 1:s, j in 1:s ]
    r = [ 1 / T(s) / T(s + k) for k in 1:s ]
    α = M \ r

    Vₛ = [ c[i]^(j-1) for i in 1:s, j in 1:s ]
    Aₛ = zeros(T, s, s)
    for i in 2:s
       Aₛ[i,i-1] = 1 / T(i-1)
    end
    Aₛ[:,s] = α

    Vₛ * Aₛ * inv(Vₛ)
end

function get_lobatto_g_coefficients(::Type{T}, s) where {T}
    a = get_lobatto_f_coefficients(T,s)
    b = get_lobatto_weights(T,s)
    ā = get_symplectic_conjugate_coefficients(a,b)
    return (a .+ ā) ./ 2
end

get_lobatto_a_coefficients(s) = get_lobatto_a_coefficients(BigFloat, s)
get_lobatto_b_coefficients(s) = get_lobatto_b_coefficients(BigFloat, s)
get_lobatto_c_coefficients(s) = get_lobatto_c_coefficients(BigFloat, s)
get_lobatto_c̄_coefficients(s) = get_lobatto_c̄_coefficients(BigFloat, s)
get_lobatto_d_coefficients(s) = get_lobatto_d_coefficients(BigFloat, s)
get_lobatto_e_coefficients(s) = get_lobatto_e_coefficients(BigFloat, s)
get_lobatto_f_coefficients(s) = get_lobatto_f_coefficients(BigFloat, s)
get_lobatto_g_coefficients(s) = get_lobatto_g_coefficients(BigFloat, s)

"""
Lobatto IIIA tableau with s stages

```julia
TableauLobattoIIIA(::Type{T}, s)
TableauLobattoIIIA(s) = TableauLobattoIIIA(Float64, s)
```

The constructor takes the number of stages `s` and optionally the element type `T` of the tableau.

References:

    Byron Leonard Ehle
    On Padé approximations to the exponential function and a-stable methods for the numerical solution of initial value problems.
    Research Report CSRR 2010, Dept. AACS, University of Waterloo, 1969.

    Laurent O. Jay.
    Lobatto Methods.
    In: Engquist B. (eds). Encyclopedia of Applied and Computational Mathematics. Springer, Berlin, Heidelberg. 2015.
    doi: 10.1007/978-3-540-70529-1_123

"""
function TableauLobattoIIIA(::Type{T}, s) where {T}
    Tableau{T}(Symbol("LobattoIIIA($s)"), 2s-2, get_lobatto_a_coefficients(s), get_lobatto_weights(s), get_lobatto_nodes(s))
end

"""
Lobatto IIIB tableau with s stages

```julia
TableauLobattoIIIB(::Type{T}, s)
TableauLobattoIIIB(s) = TableauLobattoIIIB(Float64, s)
```

The constructor takes the number of stages `s` and optionally the element type `T` of the tableau.

References:

    Byron Leonard Ehle.
    On Padé approximations to the exponential function and a-stable methods for the numerical solution of initial value problems.
    Research Report CSRR 2010, Dept. AACS, University of Waterloo, 1969.

    Laurent O. Jay.
    Lobatto Methods.
    In: Engquist B. (eds). Encyclopedia of Applied and Computational Mathematics. Springer, Berlin, Heidelberg. 2015.
    doi: 10.1007/978-3-540-70529-1_123.

"""
function TableauLobattoIIIB(::Type{T}, s) where {T}
    Tableau{T}(Symbol("LobattoIIIB($s)"), 2s-2, get_lobatto_b_coefficients(s), get_lobatto_weights(s), get_lobatto_nodes(s))
end

"""
Lobatto IIIC tableau with s stages

```julia
TableauLobattoIIIC(::Type{T}, s)
TableauLobattoIIIC(s) = TableauLobattoIIIC(Float64, s)
```

The constructor takes the number of stages `s` and optionally the element type `T` of the tableau.

References:

    F. H. Chipman.
    A-stable Runge-Kutta processes.
    BIT, Volume 11, Pages 384-388, 1971.
    doi: 10.1007/BF01939406.

    Laurent O. Jay.
    Lobatto Methods.
    In: Engquist B. (eds). Encyclopedia of Applied and Computational Mathematics. Springer, Berlin, Heidelberg. 2015.
    doi: 10.1007/978-3-540-70529-1_123.

"""
function TableauLobattoIIIC(::Type{T}, s) where {T}
    Tableau{T}(Symbol("LobattoIIIC($s)"), 2s-2, get_lobatto_c_coefficients(s), get_lobatto_weights(s), get_lobatto_nodes(s))
end

"""
Lobatto IIIC̄ tableau with s stages

```julia
TableauLobattoIIIC̄(::Type{T}, s)
TableauLobattoIIIC̄(s) = TableauLobattoIIIC̄(Float64, s)
```

The constructor takes the number of stages `s` and optionally the element type `T` of the tableau.

References:

    John C. Butcher.
    Integration processes based on Radau quadrature formulas
    Mathematics of Computation, Volume 18, Pages 233-244, 1964.
    doi: 10.1090/S0025-5718-1964-0165693-1.

    Laurent O. Jay.
    Lobatto Methods.
    In: Engquist B. (eds). Encyclopedia of Applied and Computational Mathematics. Springer, Berlin, Heidelberg. 2015.
    doi: 10.1007/978-3-540-70529-1_123.

"""
function TableauLobattoIIIC̄(::Type{T}, s) where {T}
    Tableau{T}(Symbol("LobattoIIIC̄($s)"), 2s-2, get_lobatto_c̄_coefficients(s), get_lobatto_weights(s), get_lobatto_nodes(s))
end

"""
Lobatto IIID tableau with s stages

```julia
TableauLobattoIIID(::Type{T}, s)
TableauLobattoIIID(s) = TableauLobattoIIID(Float64, s)
```

The constructor takes the number of stages `s` and optionally the element type `T` of the tableau.

References:

    R.P.K. Chan.
    On symmetric Runge-Kutta methods of high order.
    Computing, Volume 45, Pages 301-309, 1990.
    doi: 10.1007/BF02238798

    Laurent O. Jay.
    Lobatto Methods.
    In: Engquist B. (eds). Encyclopedia of Applied and Computational Mathematics. Springer, Berlin, Heidelberg. 2015.
    doi: 10.1007/978-3-540-70529-1_123.

"""
function TableauLobattoIIID(::Type{T}, s) where {T}
    Tableau{T}(Symbol("LobattoIIID($s)"), 2s-2, get_lobatto_d_coefficients(s), get_lobatto_weights(s), get_lobatto_nodes(s))
end

"""
Lobatto IIIE tableau with s stages

```julia
TableauLobattoIIIE(::Type{T}, s)
TableauLobattoIIIE(s) = TableauLobattoIIIE(Float64, s)
```

The constructor takes the number of stages `s` and optionally the element type `T` of the tableau.

References:

    R.P.K. Chan.
    On symmetric Runge-Kutta methods of high order.
    Computing, Volume 45, Pages 301-309, 1990.
    doi: 10.1007/BF02238798

    Laurent O. Jay.
    Lobatto Methods.
    In: Engquist B. (eds). Encyclopedia of Applied and Computational Mathematics. Springer, Berlin, Heidelberg. 2015.
    doi: 10.1007/978-3-540-70529-1_123.

"""
function TableauLobattoIIIE(::Type{T}, s) where {T}
    Tableau{T}(Symbol("LobattoIIIE($s)"), 2s-2, get_lobatto_e_coefficients(s), get_lobatto_weights(s), get_lobatto_nodes(s))
end

"""
Lobatto IIIF tableau with s stages

```julia
TableauLobattoIIIF(::Type{T}, s)
TableauLobattoIIIF(s) = TableauLobattoIIIF(Float64, s)
```

The constructor takes the number of stages `s` and optionally the element type `T` of the tableau.

References:

    Wang Fangzong and Liao Xiaobing.
    A Class of Lobatto Methods of Order 2s.
    Journal of Applied Mathematics, Volume 46, Pages 6-10, 2016.

"""
function TableauLobattoIIIF(::Type{T}, s) where {T}
    Tableau{T}(Symbol("LobattoIIIF($s)"), 2s,   get_lobatto_f_coefficients(s), get_lobatto_weights(s), get_lobatto_nodes(s))
end

@doc raw"""
Lobatto IIIG tableau with s stages

```julia
TableauLobattoIIIG(::Type{T}, s)
TableauLobattoIIIG(s) = TableauLobattoIIIG(Float64, s)
```

The constructor takes the number of stages `s` and optionally the element type `T` of the tableau.

Symplectizied algorithm for [`TableauLobattoIIIF`](@ref)

Coefficients are taken as $a^G = \frac{1}{2} ( a^F + \bar{a}^F )$ where the coefficients $\bar{a}^F$ are computed such that
the symplecticity conditions $b_{i} \bar{a}_{i,j} + \bar{b}_{j} a_{j,i} = b_{i} \bar{b}_{j}$ and $b_{i} = \bar{b}_i$ hold
for all $1 \le i,j \le s$.
"""
function TableauLobattoIIIG(::Type{T}, s) where {T}
    Tableau{T}(Symbol("LobattoIIIG($s)"), 2s,   get_lobatto_g_coefficients(s), get_lobatto_weights(s), get_lobatto_nodes(s))
end

TableauLobattoIIIA(s) = TableauLobattoIIIA(Float64, s)
TableauLobattoIIIB(s) = TableauLobattoIIIB(Float64, s)
TableauLobattoIIIC(s) = TableauLobattoIIIC(Float64, s)
TableauLobattoIIIC̄(s) = TableauLobattoIIIC̄(Float64, s)
TableauLobattoIIID(s) = TableauLobattoIIID(Float64, s)
TableauLobattoIIIE(s) = TableauLobattoIIIE(Float64, s)
TableauLobattoIIIF(s) = TableauLobattoIIIF(Float64, s)
TableauLobattoIIIG(s) = TableauLobattoIIIG(Float64, s)


function infilt!(I, b, P)
  I  .= P.α.*(b.+P.q*P.f)./(b.+P.q)
end

function infilt(b::Float64, P)
  return P.α*(b+P.q*P.f)/(b+P.q)
end

function infilt(b, P)
  return P.α.*(b.+P.q*P.f)./(b.+P.q)
end

function rhs_h!(ht, h, p, t)
     In, P = p
@.   ht = P.p - In*h + P.dh*$laplacian(h.*h,P.dx)
     nothing
end

function steadyh!(h, In, P)
    probh=SteadyStateProblem(rhs_h!, h, (In, P))
    hh=solve(probh, DynamicSS(Tsit5(), tspan=0.2*P.dx*P.dx/P.dw), save_everystep=false, save_start = false);
    h.=hh.u
    nothing
end

function rhs_stat!(ut, u, p, t)

    global iint

    P, fg, σ, h, bint, wint = p
    b = @view u[:,:,1]
    w = @view u[:,:,2]
    bt = @view ut[:,:,1]
    wt = @view ut[:,:,2]

    I = infilt(b, P)

    # Compute steady h solution
    steadyh!(h, I, P);

    # Compute integrals only every nint steps
    if(mod(iint,P.nint)==0)
        α = initapprox(b, P.η, σ)
        bint .= approxintb(w, fg, α)
        wint .= approxintw(b, fg, α)
        iintb = 0
    end
    iint += 1

@.  bt = P.ν*b*(1-b)*bint - b  + P.db*$laplacian(b, P.dx)
@.  wt = I*h - P.ν*w/(1+P.ρ*b) - P.γ*w*wint + P.dw*$laplacian(w, P.dx)

    nothing
end


module LinearClassifiers

export AbstractLinearClassifier, CentroidLinearClassifier, classify

"""
Supertype of all linear classifiers.

Each linear classifier should have the fields:
- `classify::Function`
"""
abstract type AbstractLinearClassifier end

struct CentroidLinearClassifier <: AbstractLinearClassifier
    classify::Function
    # `coeffs` and `offset` are parameters for the linear decision boundary, for
    # a function of the form:
    # f(x) = (coeffs \dot x) + offset
    # where the decision boundary is defined by f(x) = 0
    coeffs::AbstractVector{<:Number}
    offset::Number
end

"""
    classify(point, model)

Classifies `point` using the `classify` function of `model`.
"""
classify(point::AbstractVector{<:Number},
         model::AbstractLinearClassifier) = model.classify(point)

end # module


@testset "Lasso small (h. split, $T)" for T in [Float32, Float64, ComplexF32, ComplexF64]

    using ProximalOperators
    using ProximalAlgorithms
    using LinearAlgebra
    using AbstractOperators: MatrixOp
    using RecursiveArrayTools: ArrayPartition

    A1 = T[  1.0  -2.0   3.0;
             2.0  -1.0   0.0;
            -1.0   0.0   4.0;
            -1.0  -1.0  -1.0]
    A2 = T[ -4.0  5.0;
            -1.0  3.0;
            -3.0  2.0;
             1.0  3.0]
    A = hcat(MatrixOp(A1), MatrixOp(A2))
    b = T[1.0, 2.0, 3.0, 4.0]

    m, n1 = size(A1)
    _, n2 = size(A2)
    n = n1 + n2

    R = real(T)

    lam = R(0.1)*norm(A'*b, Inf)
    @test typeof(lam) == R

    f = Translate(SqrNormL2(R(1)), -b)
    g = SeparableSum(NormL1(lam), NormL1(lam))

    x_star = ArrayPartition(
        T[-3.877278911564627e-01, 0, 0],
        T[2.174149659863943e-02, 6.168435374149660e-01]
    )

    TOL = R(1e-4)

    @testset "ForwardBackward" begin

        ## Nonfast/Nonadaptive

        x0 = ArrayPartition(zeros(T, n1), zeros(T, n2))
        solver = ProximalAlgorithms.ForwardBackward{R}(tol=TOL)
        x, it = solver(x0, f=f, A=A, g=g, L=opnorm([A1 A2])^2)
        @test eltype(x) == T
        @test norm(x - x_star, Inf) <= TOL
        @test it < 150

        # Nonfast/Adaptive

        x0 = ArrayPartition(zeros(T, n1), zeros(T, n2))
        solver = ProximalAlgorithms.ForwardBackward{R}(tol=TOL, adaptive=true)
        x, it = solver(x0, f=f, A=A, g=g)
        @test eltype(x) == T
        @test norm(x - x_star, Inf) <= TOL
        @test it < 300

        # Fast/Nonadaptive

        x0 = ArrayPartition(zeros(T, n1), zeros(T, n2))
        solver = ProximalAlgorithms.ForwardBackward{R}(tol=TOL, fast=true)
        x, it = solver(x0, f=f, A=A, g=g, L=opnorm([A1 A2])^2)
        @test eltype(x) == T
        @test norm(x - x_star, Inf) <= TOL
        @test it < 100

        # Fast/Adaptive

        x0 = ArrayPartition(zeros(T, n1), zeros(T, n2))
        solver = ProximalAlgorithms.ForwardBackward{R}(tol=TOL, adaptive=true, fast=true)
        x, it = solver(x0, f=f, A=A, g=g)
        @test eltype(x) == T
        @test norm(x - x_star, Inf) <= TOL
        @test it < 200
    end

    @testset "ZeroFPR" begin

        # ZeroFPR/Nonadaptive

        x0 = ArrayPartition(zeros(T, n1), zeros(T, n2))
        solver = ProximalAlgorithms.ZeroFPR{R}(tol=TOL)
        x, it = solver(x0, f=f, A=A, g=g, L=opnorm([A1 A2])^2)
        @test eltype(x) == T
        @test norm(x - x_star, Inf) <= TOL
        @test it < 20

        # ZeroFPR/Adaptive

        x0 = ArrayPartition(zeros(T, n1), zeros(T, n2))
        solver = ProximalAlgorithms.ZeroFPR{R}(adaptive=true, tol=TOL)
        x, it = solver(x0, f=f, A=A, g=g)
        @test eltype(x) == T
        @test norm(x - x_star, Inf) <= TOL
        @test it < 20

    end

    @testset "PANOC" begin

        # PANOC/Nonadaptive

        x0 = ArrayPartition(zeros(T, n1), zeros(T, n2))
        solver = ProximalAlgorithms.PANOC{R}(tol=TOL)
        x, it = solver(x0, f=f, A=A, g=g, L=opnorm([A1 A2])^2)
        @test eltype(x) == T
        @test norm(x - x_star, Inf) <= TOL
        @test it < 20

        ## PANOC/Adaptive

        x0 = ArrayPartition(zeros(T, n1), zeros(T, n2))
        solver = ProximalAlgorithms.PANOC{R}(adaptive=true, tol=TOL)
        x, it = solver(x0, f=f, A=A, g=g)
        @test eltype(x) == T
        @test norm(x - x_star, Inf) <= TOL
        @test it < 20

    end

end


export Newton

"""
A globalized Newton algorithm with Line Search
Inputs:
    - An AbstractNLPModel (our problem)
    - An NLPStopping (stopping fonctions/criterion of our problem)
    - (opt) hessian_rep, different way to implement the Hessian. The default
        way is a LinearOperator
Outputs:
    - An NLPAtX, so all the information of the last iterate
    - A boolean true if we reahced an optimal solution, false otherwise
"""
function Newton(nlp             :: AbstractNLPModel,
                 nlp_stop       :: NLPStopping;
                 linesearch     :: Function = armijo_ls,
                 Nwtdirection   :: Function = NwtdirectionLDLT,
                 hessian_rep    :: Function = hessian_dense,
                 kwargs...)

    nlp_at_x = nlp_stop.current_state

    n = nlp.meta.nvar

    xt = copy(nlp.meta.x0)
	T = eltype(xt)
    ∇ft = Array{Float64}(undef, n)

    f = obj(nlp, xt)
    ∇f = grad(nlp, xt)
	nlp_stop.meta.optimality0 = norm(copy(∇f))
    OK = update_and_start!(nlp_stop, x = xt, fx = f, gx = ∇f)#, g0 = ∇f)
    # ∇fNorm = BLAS.nrm2(n, nlp_at_x.gx, 1)
    ∇fNorm = norm(nlp_at_x.gx)
    !OK && Stopping.update!(nlp_at_x, Hx = hessian_rep(nlp, xt))

    iter = 0

    @info log_header([:iter, :f, :nrm_g, :α], [Int64, T, T, T])
	@info log_row(Any[iter, nlp_at_x.fx, ∇fNorm, 0.0])
    β = 0.0
    # d = zeros(nlp_at_x.gx)
    d = zero(nlp_at_x.gx)

    h = LineSearch.LineModel(nlp, xt, d)

    while !OK
        d = Nwtdirection(nlp_at_x.Hx, nlp_at_x.gx)
        # slope = BLAS.dot(n, d, 1, nlp_at_x.gx, 1)
        slope = d' * nlp_at_x.gx

        h = LineSearch.redirect!(h, xt, d)

        ls_at_t = LSAtT(0.0, h₀ = nlp_at_x.fx, g₀ = slope)
        stop_ls = LS_Stopping(h, (x, y) -> armijo(x, y, τ₀ = 1e-09), ls_at_t)
        ls_at_t, good_step_size = linesearch(h, stop_ls; kwargs...)
        # good_step_size || (nlp_stop.meta.stalled_linesearch = true)

        xt = nlp_at_x.x + ls_at_t.x * d
        ft = obj(nlp, xt); ∇ft = grad(nlp, xt)

        # BLAS.blascopy!(n, nlp_at_x.x, 1, xt, 1)
        # BLAS.axpy!(n, ls_at_t.x, d, 1, xt, 1) #BLAS.axpy!(n, t, d, 1, xt, 1)
        ∇ft = grad!(nlp, xt, ∇ft)

        # Move on.
        OK = update_and_stop!(nlp_stop, x = xt, fx = ft, gx = ∇ft, Hx = hessian_rep(nlp, xt))

        # norm(∇f) bug: https://github.com/JuliaLang/julia/issues/11788
        # ∇fNorm = BLAS.nrm2(n, nlp_at_x.gx, 1)
        ∇fNorm = norm(nlp_at_x.gx)
        iter = iter + 1
		@show nlp_at_x.x[1]
		@show nlp_at_x.x[2]
		@info log_row(Any[iter, nlp_at_x.fx, ∇fNorm, ls_at_t.x])
    end

    return nlp_stop, nlp_stop.meta.optimal
end


"""
LinSolvers (C) 2020, Martin Gurtner
Solvers of linear systems of equations tailored to the needs of
DistCtrl4DistMan.jl package
"""
module LinSolvers

using LinearAlgebra

export kernel_solve, kernel_solve!, LinSystemData

struct LinSystemData{T<:Real, I<:Int}
    n::I;
    m::I;
    A::Matrix{T};
    b::Vector{T};
    R::Matrix{T};
    x_sol::Vector{T};

    LinSystemData{T}(n::I, m::I) where {I<:Int, T<:Real} = new{T, I}(
            n, m,
            zeros(T, m+n,m),                   # A
            zeros(T, m),                       # b
            zeros(T, m,m),    # R
            zeros(T, n)
    );

    LinSystemData(n::I, m::I, x_sol::Vector{T}) where {I<:Int, T<:Real} = new{T, I}(
        n, m,
        zeros(T, m+n,m),  # A
        zeros(T, m),      # b
        zeros(T, m,m),    # R
        x_sol
    );
end # struct

"""
    update!(lsd::LinSystemData, Dᵀ::Matrix, λ<:Real, b::Vector)

Update 'lsd' structure storing data for solving matrix equation (DᵀD + 1/λI)x=Dᵀb.
"""
function update!(lsd::LinSystemData{T}, Dᵀ::Matrix{T}, λ::T, b::Vector{T}) where T<:Real
    @assert (size(Dᵀ, 1) == lsd.n) && (size(Dᵀ, 2) == lsd.m) "LinSystemData does not have compatible dimensions with matrix Dᵀ."
    @assert length(b) == lsd.m "LinSystemData does not have compatible dimensions with vector b."

    @inbounds for i in 1:lsd.n, j in 1:lsd.m
        lsd.A[i,j] = Dᵀ[i,j];
    end

    diag_val = 1/sqrt(λ);
    @inbounds for i in 1:lsd.m, j in 1:lsd.m
        lsd.A[i+lsd.n,j] = i==j ? diag_val : 0;
    end

    copyto!(lsd.b, b);

    nothing
end

"""
    kernel_solve(Dᵀ::Matrix, λ<:Real, b::Vector)

Solve a matrix equation in form (DᵀD + 1/λ)x=Dᵀb.

# Examples
```julia-repl
julia> n, m = 10, 3;
julia> for i in 1:10
           x_sol = kernel_solve(rand(n,m), 1.0, rand(m));
           print("Solution is \$x_sol\\n");
       end
Solution is ...
```
"""
function kernel_solve(Dt::Matrix{T}, λ::T, b::Vector{T}) where T<:Real
    n, m = size(Dt);
    A = zeros(m+n, m);
    b_tmp = zeros(m);
    x_sol = zeros(n);
    R = zeros(m, m);

    A[1:n,1:m] = Dt;
    @simd for i in (n+1):(n+m)
        A[i,i-n] = 1/sqrt(λ);
    end

    qr_r!(A, R);

    copyto!(b_tmp, b);
    forward_substitution!(R, b_tmp)
    backward_substitution!(R, b_tmp)

    mul!(x_sol, Dt, b_tmp);
    return x_sol;
end

"""
    kernel_solve!(lsd::LinSystemData, Dᵀ::Matrix, λ<:Real, b::Vector)

Solve a matrix equation in form (DᵀD + 1/λ)x=Dᵀb using the pre-allocated data structure 'lsd'.

# Example
```julia-repl
julia> n, m = 10, 3;
julia> lsd = LinSystemData{Float64}(n, m);
julia> for i in 1:10
           kernel_solve!(lsd, rand(n,m), 1.0, rand(m));
           print("Solution is \$(lsd.x_sol)\\n");
       end
Solution is ...
```
"""
function kernel_solve!(lsd::LinSystemData{T}, Dt::Matrix{T}, λ::T, b::Vector{T}) where T<:Real
    update!(lsd, Dt, λ, b);

    qr_r!(lsd.A, lsd.R);

    forward_substitution!(lsd.R, lsd.b);
    backward_substitution!(lsd.R, lsd.b);

    mul!(lsd.x_sol, Dt, lsd.b);

    nothing
end

"""
    kernel_solve_prealloc!(Dt::Matrix{T}, λ::T, Fi::Vector{T}, x_sol::Vector{T},
            A::Matrix{T},
            b_tmp::Vector{T},
            R::Matrix{T}) where T<:Real

Solve a matrix equation in form (DᵀD + 1/λ)x=Dᵀb with preallocated vectors and matrices provided.

# Examples
```julia-repl
julia> n, m = 10, 3;
julia> for i in 1:10
           x_sol = kernel_solve(rand(n,m), 1.0, rand(m));
           print("Solution is \$x_sol\\n");
       end
Solution is ...
```
"""
function kernel_solve_prealloc!(Dt::Matrix{T}, λ::T, b::Vector{T}, x_sol::Vector{T},
        A::Matrix{T},
        b_tmp::Vector{T},
        R::Matrix{T}) where T<:Real
    n, m = size(Dt);

    @inbounds for i in 1:n, j in 1:m
        A[i,j] = Dt[i,j];
    end

    diag_val = 1/sqrt(λ);
    @inbounds for i in 1:m, j in 1:m
        A[i+n,j] = i==j ? diag_val : 0;
    end

    qr_r!(A, R);

    copyto!(b_tmp, b);
    forward_substitution!(R, b_tmp)
    backward_substitution!(R, b_tmp)

    mul!(x_sol, Dt, b_tmp);

    nothing
end

# Forward substitution
function forward_substitution!(R::Matrix, b::Vector)
    n = size(R,1);
    @assert n==length(b) "R and b must have compatible dimensions!"

    @fastmath @inbounds for i in 1:n
        @simd for j in 1:(i-1)
            b[i] -= b[j]*R[j,i];
        end
        b[i] /= R[i,i];
    end

    nothing
end

# Backward substitution
function backward_substitution!(R::Matrix, b::Vector)
    n = size(R,1);
    @assert n==length(b) "R and b must have compatible dimensions!"

    @fastmath @inbounds for i in n:-1:1
        @simd for j in (i+1):n
            b[i] -= b[j]*R[i,j];
        end
        b[i] /= R[i,i];
    end

    nothing
end

function qr_r!(A::Matrix, R::Matrix)
    m, n = size(A);
    @assert n==size(R,1) "Dimensions of A and R must be compatible!"

    @fastmath @inbounds for i in 1:n
        R[i,i] = sqrt(colTimesCol(A, i, i));
        for j in range(i+1, stop=n)
            R[i,j] = colTimesCol(A, i , j)/R[i,i];
            α = R[i,j]/R[i,i];
            for k in 1:m
                A[k,j] -= A[k,i]*α;
            end
        end
    end

    nothing
end

function colTimesCol(A::Matrix, i::Int, j::Int)
    res = 0.0;
    n = size(A,1);
    @inbounds @simd for k in 1:n
        res += A[k,i]*A[k,j];
    end

    return res
end

end


time_start = time()

using PowerModels
using Ipopt
using JuMP

pkg_load_time = time() - time_start

time_start = time()

file_name = "../data/pglib_opf_case5_pjm.m"
#file_name = "../data/pglib_opf_case118_ieee.m"

data = PowerModels.parse_file(file_name)
PowerModels.standardize_cost_terms!(data, order=2)
PowerModels.calc_thermal_limits!(data)
ref = PowerModels.build_ref(data)[:it][:pm][:nw][0]

data_load_time = time() - time_start

time_start = time()

model = Model(Ipopt.Optimizer)
#set_optimizer_attribute(model, "print_level", 0)

@variable(model, va[i in keys(ref[:bus])])
@variable(model, -ref[:bus][i]["vmax"] <= vr[i in keys(ref[:bus])] <= ref[:bus][i]["vmax"], start=1.0)
@variable(model, -ref[:bus][i]["vmax"] <= vi[i in keys(ref[:bus])] <= ref[:bus][i]["vmax"], start=0.0)

@variable(model, ref[:gen][i]["pmin"] <= pg[i in keys(ref[:gen])] <= ref[:gen][i]["pmax"])
@variable(model, ref[:gen][i]["qmin"] <= qg[i in keys(ref[:gen])] <= ref[:gen][i]["qmax"])

@variable(model, -ref[:branch][l]["rate_a"] <= p[(l,i,j) in ref[:arcs]] <= ref[:branch][l]["rate_a"])
@variable(model, -ref[:branch][l]["rate_a"] <= q[(l,i,j) in ref[:arcs]] <= ref[:branch][l]["rate_a"])

@objective(model, Min, sum(gen["cost"][1]*pg[i]^2 + gen["cost"][2]*pg[i] + gen["cost"][3] for (i,gen) in ref[:gen]))

for (i,bus) in ref[:ref_buses]
    @constraint(model, vi[i] == 0)
end

for (i,bus) in ref[:bus]
    bus_loads = [ref[:load][l] for l in ref[:bus_loads][i]]
    bus_shunts = [ref[:shunt][s] for s in ref[:bus_shunts][i]]

    @constraint(model, bus["vmin"]^2 <= vr[i]^2 + vi[i]^2)
    @constraint(model, bus["vmax"]^2 >= vr[i]^2 + vi[i]^2)

    @constraint(model,
        sum(p[a] for a in ref[:bus_arcs][i]) ==
        sum(pg[g] for g in ref[:bus_gens][i]) -
        sum(load["pd"] for load in bus_loads) -
        sum(shunt["gs"] for shunt in bus_shunts)*(vr[i]^2 + vi[i]^2)
    )

    @constraint(model,
        sum(q[a] for a in ref[:bus_arcs][i]) ==
        sum(qg[g] for g in ref[:bus_gens][i]) -
        sum(load["qd"] for load in bus_loads) +
        sum(shunt["bs"] for shunt in bus_shunts)*(vr[i]^2 + vi[i]^2)
    )

end

# Branch power flow physics and limit constraints
for (i,branch) in ref[:branch]
    f_idx = (i, branch["f_bus"], branch["t_bus"])
    t_idx = (i, branch["t_bus"], branch["f_bus"])

    p_fr = p[f_idx]                     # p_fr is a reference to the optimization variable p[f_idx]
    q_fr = q[f_idx]                     # q_fr is a reference to the optimization variable q[f_idx]
    p_to = p[t_idx]                     # p_to is a reference to the optimization variable p[t_idx]
    q_to = q[t_idx]                     # q_to is a reference to the optimization variable q[t_idx]

    vr_fr = vr[branch["f_bus"]]         # vm_fr is a reference to the optimization variable vm on the from side of the branch
    vr_to = vr[branch["t_bus"]]         # vm_to is a reference to the optimization variable vm on the to side of the branch
    vi_fr = vi[branch["f_bus"]]         # va_fr is a reference to the optimization variable va on the from side of the branch
    vi_to = vi[branch["t_bus"]]         # va_fr is a reference to the optimization variable va on the to side of the branch

    g, b = PowerModels.calc_branch_y(branch)
    tr, ti = PowerModels.calc_branch_t(branch)
    tm = tr^2 + ti^2
    g_fr = branch["g_fr"]
    b_fr = branch["b_fr"]
    g_to = branch["g_to"]
    b_to = branch["b_to"]

    # From side of the branch flow
    @constraint(model, p_fr ==  (g+g_fr)/tm*(vr_fr^2 + vi_fr^2) + (-g*tr+b*ti)/tm*(vr_fr*vr_to + vi_fr*vi_to) + (-b*tr-g*ti)/tm*(vi_fr*vr_to - vr_fr*vi_to) )
    @constraint(model, q_fr == -(b+b_fr)/tm*(vr_fr^2 + vi_fr^2) - (-b*tr-g*ti)/tm*(vr_fr*vr_to + vi_fr*vi_to) + (-g*tr+b*ti)/tm*(vi_fr*vr_to - vr_fr*vi_to) )

    # To side of the branch flow
    @constraint(model, p_to ==  (g+g_to)*(vr_to^2 + vi_to^2) + (-g*tr-b*ti)/tm*(vr_fr*vr_to + vi_fr*vi_to) + (-b*tr+g*ti)/tm*(-(vi_fr*vr_to - vr_fr*vi_to)) )
    @constraint(model, q_to == -(b+b_to)*(vr_to^2 + vi_to^2) - (-b*tr+g*ti)/tm*(vr_fr*vr_to + vi_fr*vi_to) + (-g*tr-b*ti)/tm*(-(vi_fr*vr_to - vr_fr*vi_to)) )

    # Voltage angle difference limit
    @constraint(model, (vi_fr*vr_to - vr_fr*vi_to) <= tan(branch["angmax"])*(vr_fr*vr_to + vi_fr*vi_to))
    @constraint(model, (vi_fr*vr_to - vr_fr*vi_to) >= tan(branch["angmin"])*(vr_fr*vr_to + vi_fr*vi_to))

    # Apparent power limit, from side and to side
    @constraint(model, p_fr^2 + q_fr^2 <= branch["rate_a"]^2)
    @constraint(model, p_to^2 + q_to^2 <= branch["rate_a"]^2)
end

model_build_time = time() - time_start

#println("model building done")
#println(model)

time_start = time()

optimize!(model)
cost = objective_value(model)

solve_time = time() - time_start

println("")
println("\033[1mSummary\033[0m")
println("   case......: $(file_name)")
println("   cost......: $(round(Int, cost))")
println("   pkg time..: $(pkg_load_time)")
println("   data time.: $(data_load_time)")
println("   build time: $(model_build_time)")
println("   solve time: $(solve_time)")
println("")


push!(LOAD_PATH,"simulator")
using EVQueues, Plots

lambda=120.0;
mu=1.0;
gamma=0.5;

C=60.0;

Tfinal=1000.0;

sim = ev_llf(lambda,mu,gamma,Tfinal,C,snapshots=[Tfinal])
compute_statistics!(sim)

#sim2 = ev_edffixed(lambda,mu,gamma,Tfinal,C,snapshots=[Tfinal])
#compute_statistics!(sim2)


### A Pluto.jl notebook ###
# v0.14.4

using Markdown
using InteractiveUtils

# ╔═╡ 09d8b052-ac4a-11eb-09db-b3560d91e761
using Omega

# ╔═╡ 7c492773-17bf-4dd1-9b45-2ce200441bc2
using Distributions

# ╔═╡ bdd33342-30da-4f57-ba56-ec71009ee082
md"# Causal VAE"

# ╔═╡ 300a118e-cd19-4515-a103-c9623554da8d
z = 1 ~ Normal(0, 1)

# ╔═╡ 4f788865-83dc-44fc-bd60-93e129cffd28
f(θ, z) = (-θ + z, θ + z)

# ╔═╡ 2a926551-e019-4385-b60b-55682344e11f
θs = rand(5)

# ╔═╡ 9f4f2521-03ef-4baa-8588-44555c938a6e

# ╔═╡ dbcca834-f021-4c18-beaa-f6afcd172e99
function f(ω)
	z = 1 ~ Normal(0, 1)
	μ_X, σ_X = f(θs[1], z)
	x = ~ Normal(μ_X, σ_X)

	# t = 0
	μ_y_0, σ_y_0 = h(θs[2], z)
	y_0 ~ normal(mu_y_0, sigma_y_0)

	# = 1
	μ_y_1, σ_y_1 = h(θs[3], z)
	y_1 = 3 ~ Normal(μ_y_1, σ_y_1)
	return y_1 - y_0
end

# ╔═╡ 3a4faf28-565e-41b4-b278-bbb5b32c48ad
μ_X, σ_X = f(θs[1], z)

# ╔═╡ 7b60e041-3c72-44b8-ab88-4d1af04d2a8a

# ╔═╡ Cell order:
# ╟─bdd33342-30da-4f57-ba56-ec71009ee082
# ╠═09d8b052-ac4a-11eb-09db-b3560d91e761
# ╠═7c492773-17bf-4dd1-9b45-2ce200441bc2
# ╠═300a118e-cd19-4515-a103-c9623554da8d
# ╠═4f788865-83dc-44fc-bd60-93e129cffd28
# ╠═2a926551-e019-4385-b60b-55682344e11f
# ╠═3a4faf28-565e-41b4-b278-bbb5b32c48ad
# ╠═9f4f2521-03ef-4baa-8588-44555c938a6e
# ╠═dbcca834-f021-4c18-beaa-f6afcd172e99
# ╠═7b60e041-3c72-44b8-ab88-4d1af04d2a8a


g = 9.81
#Insira o valor de L
L = 1000 #Em metros
T = 2pi*sqrt(L/g)


AcousticCircleParticle{T} = Particle{T,2,Acoustic{T,2},Sphere{T,2}}

"""
    t_matrix(Particle{T,2,Acoustic{T,2},Sphere{T,2}}, Acoustic{T,2}, ω, order)

The T-matrix for a 2D circlular acoustic particle in a 2D acoustic medium.
"""
function t_matrix(p::Particle{T,2,Acoustic{T,2},Sphere{T,2}}, outer_medium::Acoustic{T,2}, ω::T, basis_order::Integer)::Diagonal{Complex{T}} where T <: AbstractFloat

    # Check for material properties that don't make sense or haven't been implemented
    check_material(p, outer_medium)

    "Returns a ratio used in multiple scattering which reflects the material properties of the particles"
    function Zn(m::Integer)::Complex{T}
        m = T(abs(m))
        ak = outer_radius(p)*ω/outer_medium.c

        # set the scattering strength and type
        if isinf(p.medium.c) || isinf(p.medium.ρ)
            numer = diffbesselj(m, ak)
            denom = diffhankelh1(m, ak)
        elseif iszero(outer_medium.ρ)
            numer = diffbesselj(m, ak)
            denom = diffhankelh1(m, ak)
        elseif iszero(p.medium.ρ) || iszero(p.medium.c)
            numer = besselj(m, ak)
            denom = hankelh1(m, ak)
        else
            q = impedance(p.medium)/impedance(outer_medium) # Impedance ratio
            γ = outer_medium.c/p.medium.c #speed ratio
            numer = q * diffbesselj(m, ak) * besselj(m, γ * ak) - besselj(m, ak)*diffbesselj(m, γ * ak)
            denom = q * diffhankelh1(m, ak) * besselj(m, γ * ak) - hankelh1(m, ak)*diffbesselj(m, γ * ak)
        end

        return numer / denom
    end

    # Get Zns for positive m
    Zns = map(Zn,0:basis_order)

    return - Diagonal(vcat(reverse(Zns), Zns[2:end]))
end


########################################################################
# Mapped atoms.
########################################################################

"Mapped atom."
struct MappedAtom{T1<:Int64, T2<:AbstractLinearOp, T3<:AbstractAtom} <: AbstractAtom
    n::T1
    map::T2
    atom::T3
    function MappedAtom(map::AbstractLinearOp, atom::AbstractAtom)
        size(map, 2) == length(atom) || error("size(M) incompatible with atom")
        n = size(map, 1)
        new{Int64, AbstractLinearOp, AbstractAtom}(n, map, atom)
    end
end

Base.length(a::MappedAtom) = a.n
Base.vec(a::MappedAtom) = a.map * a.atom
Base.:(*)(M::AbstractLinearOp, a::MappedAtom) = M*vec(a)
LinearAlgebra.mul!(Ma::AbstractVector{Float64}, M::LinearOp, a::MappedAtom) = mul!(Ma, M, vec(a))

########################################################################
# Mapped atomic set.
########################################################################

"""
    MappedAtomicSet(n::Integer, A::AbstractAtomicSet, M::LinearMap)

Linear transformation another atomic set, including scaling.
"""
struct MappedAtomicSet{T1<:Int64, T2<:AbstractLinearOp, T3<:AbstractAtomicSet} <: AbstractAtomicSet
    n::T1
    map::T2
    set::T3
    function MappedAtomicSet(map::AbstractLinearOp, set::AbstractAtomicSet)
        size(map, 2) == length(set) || error("size(M) incompatible with atomic set")
        n = size(map, 1)
        return new{Int64, AbstractLinearOp, AbstractAtomicSet}(n, map, set)
    end
end

"""
    *(M, A::AbstractAtomicSet)

Apply the linear operator `M` on the atomic set `A`. Results in
another atomic set (MA) with the defining property

    MA = { Ma | a ∈ A }.
"""
function Base.:(*)(M::AbstractLinearOp, A::AbstractAtomicSet)
    return MappedAtomicSet(M, A)
end

function Base.:(*)(λ::Float64, A::AbstractAtomicSet)
    return MappedAtomicSet( LinearMaps.UniformScalingMap(λ, length(A)), A)
end

"""
    support(A::MappedAtomicSet, z::Vector)

Support value of a gauge with respect to a scaled atomic set.
"""
support(A::MappedAtomicSet, z::Vector{Float64}) = support(A.set, A.map'*z)

"""
    expose(A::MappedAtomicSet, z::Vector)

Obtain an atom in the face of the face of the scaled atomic set exposed by
the vector `z`.
"""
function expose!(A::MappedAtomicSet, z::Vector{Float64}, a::MappedAtom)
    expose!(A.set, A.map'*z, a.atom)
    return nothing
end

function expose(A::MappedAtomicSet, z::Vector{Float64})
    atom = expose(A.set, A.map'*z)
    return MappedAtom(A.map, atom)
end

Base.length(A::MappedAtomicSet) = A.n
atom_name(A::MappedAtomicSet) = "Map of $(atom_name(A.set))"
atom_description(A::MappedAtomicSet) = "M ⋅ $(atom_description(A.set))"
atom_parameters(A::MappedAtomicSet) = "n = $(length(A))"

########################################################################
# Face of the Mapped Atomic Set
########################################################################

struct MappedFace{T1<:Int64, T2<:AbstractLinearOp, T3<:AbstractFace} <: AbstractFace
    n::T1
    k::T1
    map::T2
    fa::T3
    function MappedFace(map::AbstractLinearOp, fa::AbstractFace)
        size(map, 2) == length(fa) || error("size(M) incompatible with face")
        n = size(map, 1)
        k = rank(fa)
        new{Int64, AbstractLinearOp, AbstractFace}(n, k, map, fa)
    end
end

"""
    face(A, z)

Return a face of the atomic set `A` exposed by the vector `z`.
"""
function face!(A::MappedAtomicSet, z::Vector{Float64}, F::MappedFace)
    face!(A.set, A.map'*z, F.fa)
    return nothing
end

function face(A::MappedAtomicSet, z::Vector{Float64})
    fa = face(A.set, A.map'*z)
    return MappedFace(A.map, fa)
end

"""
Given a face and a set of weights, reveal the corresponding
point on the face.
"""
Base.:(*)(F::MappedFace, c::Vector) = (F.map*F.fa)*c
Base.:(*)(M::AbstractLinearOp, F::MappedFace) = M*F.map*F.fa
Base.:(*)(λ::Real, F::MappedFace) = λ*F.map*F.fa
Base.length(F::MappedFace) = F.n
rank(F::MappedFace) = F.k
vec(F::MappedFace) = F.map*F.fa
face_name(F::MappedFace) = "Mapped Atomic Set"
face_parameters(F::MappedFace) = "rank = $(rank(F)); n = $(length(F))"