# * Atomic System of Equations

"""
    AtomicEquations(atom, equations, integrals)

Structure representing the (e.g. Hartree–Fock) `equations` for `atom`,
along with all `integrals` that are shared between the `equations`.
"""
mutable struct AtomicEquations{T, B<:AbstractQuasiMatrix,
                               O<:AbstractOrbital,
                               A<:Atom{T,B,O},
                               Equations<:AbstractVector{<:AtomicOrbitalEquation}}
    atom::A
    equations::Equations
    integrals::Vector{OrbitalIntegral}
    observables::Dict{Symbol,<:Observable}
end

# Iteration interface
Base.length(hfeqs::AtomicEquations) = length(hfeqs.equations)
Base.iterate(iter::AtomicEquations, args...) = iterate(iter.equations, args...)

"""
    update!(equations::AtomicEquations[, atom::Atom])

Recompute all integrals using the current values for the radial
orbitals (optionally specifying which `atom` from which the orbitals
are taken).
"""
function SCF.update!(equations::AtomicEquations, args...; kwargs...)
    Threads.@threads for integral in equations.integrals
        SCF.update!(integral, args...; kwargs...)
    end
    # Updating the operators is apparently not thread-safe at the
    # moment. Why?
    for eq in equations.equations
        update!(SCF.hamiltonian(eq), args...; kwargs...)
    end
end

"""
    energy_matrix!(H, hfeqs::AtomicEquations[, which=:energy])

Compute the energy matrix by computing the energy observable and
storing it in `H`. Requires that `hfeqs` has the `:energy` and
`:kinetic_energy` [`Observable`](@ref)s registered (this is the
default).
"""
function SCF.energy_matrix!(H::HM, hfeqs::AtomicEquations,
                            which::Symbol=:total_energy) where {HM<:AbstractMatrix}
    observable = hfeqs.observables[which]
    observe!(H, observable)
    H
end

# * Orbital symmetries
"""
    find_symmetries(orbitals)

Group all orbitals according to their symmetries, e.g. ℓ for
`Orbital`s. This is used to determine which off-diagonal Lagrange
multipliers are necessary to maintain orthogonality.
"""
find_symmetries(orbitals::Vector{O}) where {O<:AbstractOrbital} =
    merge!(vcat, [Dict(symmetry(orb) => O[orb])
                  for orb in orbitals]...)

SCF.symmetries(atom::Atom) =
    map(values(find_symmetries(atom.orbitals))) do orbitals
        orbital_index.(Ref(atom), orbitals)
    end

# * Setup orbital equations

function get_operator(::FieldFreeOneBodyHamiltonian, atom::Atom,
                      orbital::aO, source_orbital::bO; kwargs...) where {aO,bO}
    orbital == source_orbital ||
        throw(ArgumentError("FieldFreeOneBodyHamiltonian not pertaining to $(orbital)"))
    AtomicOneBodyHamiltonian(atom, orbital)
end

for (i,HT) in enumerate([:KineticEnergyHamiltonian, :PotentialEnergyHamiltonian])
    @eval begin
        get_operator(::$HT, atom::Atom, orbital::aO, ::bO; kwargs...) where {aO,bO} =
            AtomicOneBodyHamiltonian(one_body_hamiltonian(Tuple, atom, orbital)[$i],
                                     orbital)
    end
end

function get_operator(op::CoulombPotentialMultipole, atom::Atom,
                      orbital::aO, ::bO; kwargs...) where {aO,bO}
    a,b = op.a[1],op.b[1]
    HFPotential(:exchange, op.o.k, a, a, atom, op.o.g; kwargs...)
end

get_operator(top::IdentityOperator, atom::Atom,
             ::aO, source_orbital::bO; kwargs...) where {aO,bO} =
                 SourceTerm(top, source_orbital, view(atom, source_orbital))
SCF.update!(::IdentityOperator; kwargs...) = nothing
SCF.update!(::IdentityOperator, ::Atom; kwargs...) = nothing

function get_operator(op::RadialOperator, atom::Atom, orbital::aO, source_orbital::bO; kwargs...) where {aO, bO}
    if orbital == source_orbital
        # In this case, the operator is diagonal in orbital space,
        # i.e. it maps an orbital onto itself.
        op
    else
        # In this case, the operator maps from source_orbital to
        # orbital.
        SourceTerm(op, source_orbital, view(atom, source_orbital))
    end
end

function get_operator(M::AbstractMatrix, atom::Atom, a::aO, b::bO; kwargs...) where {aO, bO}
    R = radial_basis(atom)
    get_operator(R*M*R', atom, a, b; kwargs...)
end

# RadialOperators (which are built from matrices), are independent of
# the atom.
SCF.update!(::RadialOperator; kwargs...) = nothing
SCF.update!(::RadialOperator, ::Atom; kwargs...) = nothing

# function get_operator(top::Projector, ::Atom, orbital::aO, source_orbital::bO) where {aO,bO}
#     orbital == source_orbital || return 0
#     SourceTerm(IdentityOperator{1}(), source_orbital,
#                top.ϕs[findfirst(isequal(source_orbital), top.orbitals)])
# end

get_operator(op::QO, atom::Atom, ::aO, ::bO; kwargs...) where {QO,aO,bO} =
    throw(ArgumentError("Unsupported operator $op"))

"""
    generate_atomic_orbital_equations(atom::Atom, eqs::MCEquationSystem,
                                      integrals, integral_map)

For each variationally derived orbital equation in `eqs`, generate the
corresponding [`AtomicOrbitalEquation`](@ref).
"""
function generate_atomic_orbital_equations(atom::Atom{T,B,O}, eqs::MCEquationSystem,
                                           integrals::Vector,
                                           integral_map::Dict{Any,Int},
                                           symmetries::Dict;
                                           verbosity=0,
                                           kwargs...) where {T,B,O}
    map(eqs.equations) do equation
        orbital = equation.orbital
        terms = Vector{OrbitalHamiltonianTerm{O,O,T}}()

        for (integral,equation_terms) in equation.terms
            if integral > 0
                operator = get_integral(integrals, integral_map, eqs.integrals[integral])
                # We first add all terms with operators diagonal in
                # orbital space.
                pushterms!(terms, operator, filter(t -> t.source_orbital == orbital, equation_terms),
                           integrals, integral_map, eqs.integrals)
                # We then add those terms are that are off-diagonal in
                # orbital space.
                for t in filter(t -> t.source_orbital ≠ orbital, equation_terms)
                    pushterms!(terms, SourceTerm(operator, t.source_orbital, view(atom, t.source_orbital)), [t],
                               integrals, integral_map, eqs.integrals)
                end
            else
                for t in equation_terms
                    operator = get_operator(t.operator, atom, orbital, t.source_orbital; kwargs...)
                    iszero(operator) && continue

                    pushterms!(terms, operator, [t],
                               integrals, integral_map, eqs.integrals)
                end
            end
        end

        # Find all other orbitals of the same symmetry as the current
        # one. These will be used to create a projector, that projects
        # out their components.
        #
        # TODO: Think of what the non-orthogonalities due to
        # `overlaps` imply for the Lagrange multipliers/projectors.
        symmetry_orbitals = filter(!isequal(orbital), symmetries[symmetry(orbital)])
        verbosity > 1 && println("Symmetry: ", symmetry_orbitals)

        AtomicOrbitalEquation(atom, equation, orbital, terms, symmetry_orbitals)
    end
end

atomic_hamiltonian(::Atom{T,B,O,TC,CV,P}) where {T,B,O,TC,CV,P<:AbstractPotential} =
    FieldFreeOneBodyHamiltonian() + CoulombInteraction()

"""
    diff(atom; H=atomic_hamiltonian(atom), overlaps=[], selector=outsidecoremodel, verbosity=0)

Differentiate the energy expression of the Hamiltonian `H` associated
with the `atom`'s configurations(s) with respect to the atomic
orbitals to derive the Hartree–Fock equations for the orbitals.

By default, the Hamiltonian
`H=FieldFreeOneBodyHamiltonian()+CoulombInteraction()`.

Non-orthogonality between orbitals can be specified by providing
`OrbitalOverlap`s between these pairs. Only those electrons not
modelled by `atom.potential` of each configuration are considered for
generating the energy expression, this can be changed by choosing
another value for `selector`.
"""
function Base.diff(atom::Atom,
                   energy_expression::EnergyExpression;
                   H::QuantumOperator=atomic_hamiltonian(atom),
                   overlaps::Vector{<:OrbitalOverlap}=OrbitalOverlap[],
                   selector::Function=cfg -> outsidecoremodel(cfg, atom.potential),
                   configurations = selector.(atom.configurations),
                   orbitals = unique_orbitals(configurations),
                   symmetries = find_symmetries(orbitals),
                   observables::Dict{Symbol,Tuple{<:QuantumOperator,Bool}} =
                   Dict{Symbol,Tuple{<:QuantumOperator,Bool}}(
                       :total_energy => (H,false),
                       :double_counted_energy => (H,true),
                       :kinetic_energy => (KineticEnergyHamiltonian(),false),
                   ),
                   modify_eoms!::Function = eqs -> nothing,
                   modify_integrals!::Function = (eqs,integrals,integral_map) -> nothing,
                   modify_equations!::Function = hfeqs -> nothing,
                   verbosity=0)
    eqs = diff(energy_expression, conj.(orbitals))
    modify_eoms!(eqs)

    if verbosity > 0
        println("Energy expression:")
        display(energy_expression)
        println()

        println("Hartree–Fock equations:")
        display(eqs)
        println()

        println("Symmetries:")
        display(symmetries)
        println()

        if verbosity > 1
            println("Common integrals:")
            display(eqs.integrals)
            println()
        end
    end

    integrals = Vector{OrbitalIntegral}()
    integral_map = Dict{Any,Int}()
    poisson_cache = Dict{Int,CoulombIntegrals.PoissonCache}()
    append_common_integrals!(integrals, integral_map,
                             atom, eqs.integrals,
                             poisson_cache=poisson_cache)
    modify_integrals!(eqs, integrals, integral_map)

    hfeqs = generate_atomic_orbital_equations(atom, eqs,
                                              integrals, integral_map,
                                              symmetries; verbosity=verbosity,
                                              poisson_cache=poisson_cache)
    modify_equations!(hfeqs)

    observables = map(collect(pairs(observables))) do (k,(operator,double_counted))
        verbosity > 2 && println("Observable: $k ($operator)")
        k => Observable(operator, atom, overlaps,
                        integrals, integral_map,
                        symmetries, selector;
                        double_counted=double_counted,
                        poisson_cache=poisson_cache,
                        verbosity=verbosity)
    end |> Dict{Symbol,Observable}

    AtomicEquations(atom, hfeqs, integrals, observables)
end

Base.Matrix(H::QuantumOperator, atom::Atom;
            overlaps::Vector{<:OrbitalOverlap}=OrbitalOverlap[],
            selector::Function=cfg -> outsidecoremodel(cfg, atom.potential),
            configurations = selector.(atom.configurations),
            kwargs...) =
    Matrix(H, configurations, overlaps)

function Base.diff(atom::Atom; H::QuantumOperator=atomic_hamiltonian(atom),
                   kwargs...)
    energy_expression = Matrix(H, atom; kwargs...)
    diff(atom, energy_expression; H=H, kwargs...)
end

# * Overlap matrix
function SCF.overlap_matrix(atom::Atom)
    R = radial_basis(atom)
    R'R
end


using QuantumControlBase.QuantumPropagators: propstep!, write_to_storage!, get_from_storage!
using QuantumControlBase: evalcontrols!, resetgradvec!
using QuantumControlBase.Functionals: grad_J_T_via_chi!, make_gradient, make_chi
using QuantumControlBase.ConditionalThreads: @threadsif
using LinearAlgebra
using Printf

"""Optimize a control problem using GRAPE.

```julia
result = optimize_grape(problem)
```

optimizes the given
control [`problem`](@ref QuantumControlBase.ControlProblem),
returning a [`GrapeResult`](@ref).

!!! note

    It is recommended to call [`optimize`](@ref QuantumControlBase.optimize)
    with `method=:GRAPE` instead of calling `optimize_grape` directly.

Keyword arguments that control the optimization are taken from the keyword
arguments used in the instantiation of `problem`.

# Required problem keyword arguments

* `J_T`: A function `J_T(ϕ, objectives; τ=τ)` that evaluates the final time
  functional from a vector `ϕ` of forward-propagated states and
  `problem.objectives`. For all `objectives` that define a `target_state`, the
  element `τₖ` of the vector `τ` will contain the overlap of the state `ϕₖ`
  with the `target_state` of the `k`'th objective, or `NaN` otherwise.

# Optional problem keyword arguments

* `gradient_via`: A flag indicating how the gradient of `J_T` should be
   calculated. One of `:tau`, `:chi`. If all objectives in `problem` define a
   `target_function`, the default is `gradient_via=:tau`. This understands the
   functional `J_T` as a function of the complex overlaps between the
   propagated states and the target states, and evaluates the gradient via the
   chain rule. For functionals that cannot be expressed in terms of the overlap
   of propagated and target states, and/or if not all objectives define a
   target state, `gradient_via=:chi` because the default. In this case, the
   functional `J_T` is understood as a function of the forward-propagated
   states directly, and the full gradient is again calculated by the chain
   rule. See [`make_gradient`](@ref) for details.
* `gradient`:  A function to evaluate the gradient of `J_T`. By default, it is
  constructed via [`make_gradient`](@ref). If given manually, it must meet the
  interface described by [`make_gradient`](@ref). Most importantly, it must be
  consistent with the chosen `gradient_via`.
* `chi`: If `gradient_via=:chi`, a function that constructs the χ-states for
   the backward propagation, see [`make_gradient`](@ref) for details. By
   default, it is constructed via [`make_chi`](@ref). If given manually, it
   must meet the same interface described in [`make_chi`](@ref).
* `force_zygote=false`: Whether to force the use of automatic differentiation
  in [`make_gradient`](@ref) and [`make_chi`](@ref). This disables analytic
  gradients. The only reason to do this is for testing/benchmarking analytic vs
  automatic gradients.
* `update_hook`: Not implemented
* `info_hook`: A function that receives the same arguments as `update_hook`, in
  order to write information about the current iteration to the screen or to a
  file. The default `info_hook` prints a table with convergence information to
  the screen. Runs after `update_hook`. The `info_hook` function may return a
  tuple, which is stored in the list of `records` inside the
  [`GrapeResult`](@ref) object.
* `check_convergence`: A function to check whether convergence has been
  reached. Receives a [`GrapeResult`](@ref) object `result`, and should set
  `result.converged` to `true` and `result.message` to an appropriate string in
  case of convergence. Multiple convergence checks can be performed by chaining
  functions with `∘`. The convergence check is performed after any calls to
  `update_hook` and `info_hook`.
* `x_tol`: Parameter for Optim.jl
* `f_tol`: Parameter for Optim.jl
* `g_tol`: Parameter for Optim.jl
* `show_trace`: Parameter for Optim.jl
* `extended_trace`:  Parameter for Optim.jl
* `show_every`: Parameter for Optim.jl
* `allow_f_increases`: Parameter for Optim.jl
* `optimizer`: An optional Optim.jl optimizer (`Optim.AbstractOptimizer`
  instance). If not given, an [L-BFGS-B](https://github.com/Gnimuc/LBFGSB.jl)
  optimizer will be used.
* `prop_method`/`fw_prop_method`/`bw_prop_method`: The propagation method to
  use for each objective, see below.
* `prop_method`/`fw_prop_method`/`grad_prop_method`: The propagation method to
  use for the extended gradient vector for each objective, see below.
* `verbose=false`: If `true`, print information during initialization

The propagation method for the forward propagation of each objective is
determined by the first available item of the following:

* a `fw_prop_method` keyword argument
* a `prop_method` keyword argument
* a property `fw_prop_method` of the objective
* a property `prop_method` of the objective
* the value `:auto`

The propagation method for the backward propagation is determined similarly,
but with `bw_prop_method` instead of `fw_prop_method`. The propagation method
for the forward propagation of the extended gradient vector for each objective
is determined from `grad_prop_method`, `fw_prop_method`, `prop_method` in order
of precedence.
"""
function optimize_grape(problem)
    update_hook! = get(problem.kwargs, :update_hook, (args...) -> nothing)
    # TODO: implement update_hook
    # TODO: streamline the interface for info_hook
    # TODO: check if x_tol, f_tol, g_tol are used necessary / used correctly
    info_hook = get(problem.kwargs, :info_hook, print_table)
    check_convergence! = get(problem.kwargs, :check_convergence, res -> res)
    verbose = get(problem.kwargs, :verbose, false)

    wrk = GrapeWrk(problem; verbose)

    χ = wrk.bw_states
    Ψ = wrk.fw_states
    χ̃ = wrk.bw_grad_states
    Ψ₀ = [obj.initial_state for obj ∈ wrk.objectives]
    Ψtgt = Union{eltype(Ψ₀),Nothing}[
        (hasproperty(obj, :target_state) ? obj.target_state : nothing) for
        obj ∈ wrk.objectives
    ]

    if any(isnothing, Ψtgt)
        gradient_via = get(wrk.kwargs, :gradient_via, :chi)
        if gradient_via == :tau
            error("Evaluating gradients via τ requires target_states for all objectives")
        end
    else
        gradient_via = get(wrk.kwargs, :gradient_via, :tau)
    end
    allowed_gradient_via = (:tau, :chi)  # TODO: :tau_matrix
    J_T_func = wrk.kwargs[:J_T]
    force_zygote = get(wrk.kwargs, :get_zygote, false)
    default_gradfunc! =
        make_gradient(J_T_func, wrk.objectives; via=gradient_via, force_zygote)
    gradfunc! = get(wrk.kwargs, :gradient, default_gradfunc!)
    if gradient_via == :tau
        # chi! is not used
    elseif gradient_via == :chi
        default_chi! = make_chi(J_T_func, wrk.objectives; force_zygote)
        chi! = get(wrk.kwargs, :chi, default_chi!)
        if gradfunc! ≢ grad_J_T_via_chi!
            @warn "gradient_via=:chi requires gradient=grad_J_T_via_chi! Ignoring passed gradient"
            gradfunc! = grad_J_T_via_chi!
        end
    else
        throw(
            ArgumentError(
                "gradient_via $(repr(gradient_via)) is not one of $(repr(allowed_gradient_via))"
            )
        )
    end

    τ = wrk.result.tau_vals
    ∇τ = wrk.tau_grads
    N_T = length(wrk.result.tlist) - 1
    N = length(wrk.objectives)
    L = length(wrk.controls)
    Φ = wrk.fw_storage

    # Calculate the functional only; optionally store.
    # Side-effects: set Ψ, τ, wrk.result.f_calls, wrk.fg_count
    function f(F, G, pulsevals; storage=nothing, count_call=true)
        @assert !isnothing(F)
        @assert isnothing(G)
        if count_call
            wrk.result.f_calls += 1
            wrk.fg_count[2] += 1
        end
        @threadsif wrk.use_threads for k = 1:N
            local Φₖ = isnothing(storage) ? nothing : storage[k]
            copyto!(Ψ[k], Ψ₀[k])
            (Φₖ !== nothing) && write_to_storage!(Φₖ, 1, Ψ₀[k])
            for n = 1:N_T  # `n` is the index for the time interval
                local (G, dt) = _fw_gen(pulsevals, k, n, wrk)
                propstep!(Ψ[k], G, dt, wrk.fw_prop_wrk[k])
                (Φₖ !== nothing) && write_to_storage!(Φₖ, n + 1, Ψ[k])
            end
            τ[k] = isnothing(Ψtgt[k]) ? NaN : (Ψtgt[k] ⋅ Ψ[k])
        end
        return J_T_func(Ψ, wrk.objectives; τ=τ)
    end

    # Calculate the functional and the gradient G ≡ ∇J_T
    function fg!(F, G, pulsevals)

        if isnothing(G)  # functional only
            return f(F, G, pulsevals)
        end
        wrk.result.fg_calls += 1
        wrk.fg_count[1] += 1

        # forward propagation and storage of states
        J_T_val = f(wrk.result.J_T, nothing, pulsevals; storage=Φ, count_call=false)

        # backward propagation of combined χ-state and gradient
        if gradient_via == :tau
            for k = 1:N
                copyto!(χ[k], Ψtgt[k])
            end
        elseif gradient_via == :chi
            chi!(χ, Ψ, wrk.objectives)  # Ψ were set in forward prop
        end
        @threadsif wrk.use_threads for k = 1:N
            resetgradvec!(χ̃[k], χ[k])
            for n = N_T:-1:1  # N_T is the number of time slices
                local (G̃, dt) = _bw_gradgen(pulsevals, k, n, wrk)
                propstep!(χ̃[k], G̃, dt, wrk.grad_prop_wrk[k])
                get_from_storage!(Ψ[k], Φ[k], n)
                for l = 1:L
                    ∇τ[k][l, n] = χ̃[k].grad_states[l] ⋅ Ψ[k]
                end
                resetgradvec!(χ̃[k])
            end
        end

        gradfunc!(G, τ, ∇τ)
        return J_T_val

    end

    optimizer = get_optimizer(wrk)
    res = run_optimizer(optimizer, wrk, fg!, info_hook, check_convergence!)
    finalize_result!(wrk, res)
    return wrk.result

end

function get_optimizer(wrk)
    n = length(wrk.pulsevals)
    m = 10 # TODO: kwarg for number of limited memory corrections
    optimizer = get(wrk.kwargs, :optimizer, LBFGSB.L_BFGS_B(n, m))
    return optimizer
end

# The dynamical generator for the forward propagation
function _fw_gen(pulse_vals, k, n, wrk)
    vals_dict = wrk.vals_dict[k]
    t = wrk.result.tlist
    L = length(wrk.controls)
    ϵ = reshape(pulse_vals, L, :)
    for (l, control) in enumerate(wrk.controls)
        vals_dict[control] = ϵ[l, n]
    end
    dt = t[n+1] - t[n]
    evalcontrols!(wrk.G[k], wrk.objectives[k].generator, vals_dict)
    return wrk.G[k], dt
end

# The dynamical generator for the backward gradient propagation
function _bw_gradgen(pulse_vals, k, n, wrk)
    vals_dict = wrk.vals_dict[k]
    t = wrk.result.tlist
    L = length(wrk.controls)
    ϵ = reshape(pulse_vals, L, :)
    for (l, control) in enumerate(wrk.controls)
        vals_dict[control] = ϵ[l, n]
    end
    dt = t[n+1] - t[n]
    evalcontrols!(wrk.gradG[k], wrk.TDgradG[k], vals_dict)
    return wrk.gradG[k], -dt
end


"""
Surprise-REINFORCE
------------------

REINFORCE with surprise-modulated learning rate

See REINFORCE.jl for some more implementational info
See learner/REINFORCEmodulatedlearner.jl for the core implementation.

Note 1
------
With particle filtering we need to run multiple times with different seeds.
As final LL we report the MEAN across seeds.
For signals/results we keep track of the ones of all seeds
and in the end report the one with the MAXIMUM LL.

Note 2
------
get_all_puzzles(), get_all_episodes_in_puzzle(), get_data_by_puzzle_episode()
are in rl_data_toolsforalgos.jl
"""
function getLL_REINFORCE_modulated(data::Array{Int,2},
                                    alpha::Float64,
                                    surpriseprobability::Float64,
                                    stochasticity::Float64,
                                    gamma::Float64,
                                    temp::Float64,
                                    nparticles::Union{Int, Float64};
                                    particlefilterseeds = [1,2,3],
                                    do_debug_print = false)
    nparticles = converttoInt(nparticles)
    data_stats, nr_states, nr_actions, nr_samples = get_data_stats_and_dimensions(data)
    if data_stats.nr_subjects != 1
        error("Invalid argument: getLL_REINFORCE_modulated supports only one subject")
    end
    if nr_actions!=2; error("REINFORCE supports two actions only"); end
    all_puzzles = get_all_puzzles(data)

    action_selection_probs_seeds = Array{Array{Float64, 1}, 1}(undef, length(particlefilterseeds))
    [action_selection_probs_seeds[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]
    LL_seeds = zeros(Float64, length(particlefilterseeds))

    theta_seeds = Array{Any, 1}(undef, length(particlefilterseeds))
    Tsas_seeds = Array{Any, 1}(undef, length(particlefilterseeds))
    #signals:
    Surprise_seeds =Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [Surprise_seeds[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]
    γ_surprise_seeds =Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [γ_surprise_seeds[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]
    SPE_seeds =Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [SPE_seeds[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]

    PolicyParametersDiffChosenVsUnchosen_record_seed =Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [PolicyParametersDiffChosenVsUnchosen_record_seed[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]
    PolicyDiffChosenVsUnchosen_record_seed = Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [PolicyDiffChosenVsUnchosen_record_seed[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]

    finaltheta_record = 0
    for i_seed in 1:length(particlefilterseeds)
        action_counter = 0
        Testimate = []
        learner = []
        for i_puzzle in all_puzzles
            # Puzzle = independent block of the experiment.
            # The code was made for general purposes, in this experiment we
            # have only ONE puzzle.
            Testimate = TParticleFilterJump(nr_states = nr_states,
                                    nr_actions = nr_actions,
                                    nparticles = nparticles,
                                    surpriseprobability = surpriseprobability,
                                    stochasticity = stochasticity,
                                    seed = particlefilterseeds[i_seed])

            learner = REINFORCEModulated(nr_states = nr_states,
                            nr_actions = nr_actions,
                            α = alpha,
                            γ = gamma,
                            do_debug_print = do_debug_print)

            all_epi_in_puzzle = get_all_episodes_in_puzzle(data, i_puzzle)
            for i_episode in all_epi_in_puzzle
                episode_data = get_data_by_puzzle_episode(data, i_puzzle, i_episode)
                path_length = size(episode_data)[1]  # = T
                if do_debug_print
                    println("i_puzzle:$i_puzzle | i_episode:$i_episode | path_length:$path_length")
                end
                reward = 0. # we have a single reward at goal state. If the episode was
                # not completed, this code still correctly works (updates are 0)
                for t in 1:path_length
                    action_counter += 1
                    state = episode_data[t, Int(c_state)]
                    action = episode_data[t, Int(c_action)]
                    reward = episode_data[t, Int(c_reward)]
                    next_state = episode_data[t, Int(c_nextState)]
                    done = reward>0 ? true : false

                    p_actions = getSoftmaxProbs([learner.θ[state,1],learner.θ[state,2]], temp, [0., 0.])

                    LL_seeds[i_seed] += log(p_actions[action])
                    action_selection_probs_seeds[i_seed][action_counter] = copy(p_actions[action])

                    # Particle Filtering model learner
                    updatet!(Testimate, state, action, next_state, done)
                    γ_surprise_seeds[i_seed][action_counter] = copy(Testimate.γ_surprise[1])
                    Surprise_seeds[i_seed][action_counter] = copy(Testimate.Sgm[1])
                    Tsas_seeds[i_seed]= deepcopy(Testimate.Ps1a0s0)
                    SPE_seeds[i_seed][action_counter] = copy(Testimate.SPE[1])

                    PolicyParametersDiffChosenVsUnchosen_record_seed[i_seed][action_counter] = learner.θ[state, action] - learner.θ[state, mod(action,2)+1]
                    PolicyDiffChosenVsUnchosen_record_seed[i_seed][action_counter] = p_actions[action] - p_actions[mod(action,2)+1]

                    # REINFORCE modulated
                    update!(learner, state, action, reward, p_actions[action], (1. - Testimate.γ_surprise[1]))

                    if do_debug_print
                        println("t:$t | s:$state | a:$action | R:$reward | s':$next_state | p(s,a):$(round(p_action,digits=2))")
                    end
                end # path_length (in single episode)
          end # episode
      end # puzzle
      theta_seeds[i_seed] = copy(learner.θ)
  end # seed

  LL = get_meanLL_acrossseeds(LL_seeds)

  action_selection_probs_cat = hcat(action_selection_probs_seeds...)
  action_selection_probs = mean(action_selection_probs_cat, dims=2)[:]
  trial_by_trial_p_sa = (action_selection_probs)

  # For signals: return the ones corresponding to the Max LL across seeds
  LL_max, seed_idx = findmax(LL_seeds)
  trial_by_trial_p_sa_max = (action_selection_probs_seeds[seed_idx])

  SPE_record = copy(SPE_seeds[seed_idx])

  algo = "REINFORCE_modulated"
  param_v = [alpha, surpriseprobability, stochasticity, gamma, temp, nparticles]
  param_str = "alpha=$(round(alpha,digits=4)),
              surpriseprobability=$(round(surpriseprobability,digits=4)),
              stochasticity=$(round(stochasticity,digits=4)),
              gamma=$(round(gamma,digits=4)),
              temp=$(round(temp,digits=4)),
              nparticles=$nparticles"

  results_dict = Dict{String,Any}("final_Theta" => theta_seeds[seed_idx],
                                "final_Tsas" => Tsas_seeds[seed_idx],
                                "Surprise" => Surprise_seeds[seed_idx],
                                "gamma_surprise" => γ_surprise_seeds[seed_idx],
                                "nparticles" => nparticles,
                                "particlefilterseeds" => particlefilterseeds,
                                "trial_by_trial_p_sa_max" => trial_by_trial_p_sa_max,
                                "PolicyParametersDiffChosenVsUnchosen" => PolicyParametersDiffChosenVsUnchosen_record_seed[seed_idx],
                                "PolicyDiffChosenVsUnchosen" => PolicyDiffChosenVsUnchosen_record_seed[seed_idx],)

  signals = Step_by_step_signals(
    data_stats.all_subject_ids[1],
    [], # no RPE
    SPE_record, # no SPE
    trial_by_trial_p_sa,
    LL,
    data_stats,
    algo,
    param_v,
    param_str,
    99, #biased_action_idx, dummy value
    results_dict)
  return LL, signals

end


using Test
using ScottishTaxBenefitModel
using .ModelHousehold
using .Results
using .Intermediate: MTIntermediate, make_intermediate
using .ScottishBenefits
using .STBIncomes
using .ExampleHelpers

sys = get_system( scotland=true )

@testset "Scottish Child Payment" begin

    for (hht,hh) in get_all_examples()
        bus = get_benefit_units( hh )
        intermed = make_intermediate(
            hh,
            sys.hours_limits,
            sys.age_limits,
            sys.child_limits )
        hres = init_household_result( hh )
        hhead = get_head( hh )
        hpid = hhead.pid
        for buno in eachindex(bus)
            bures = hres.bus[buno]
            bu = bus[buno]
            bint = intermed.buint[buno]
            ncs = num_children( bu )
            bhead = get_head( bus[buno ])
            bspouse = get_spouse( bus[buno ])
            if ncs > 0
                println( "on hhld $(hht)")
                # no qualifying benefit
                ages = fill( 6, ncs ) # 1 year too old
                set_childrens_ages!( hh, ages... )
                calc_scottish_child_payment!(
                    bures,
                    bu,
                    bint,
                    sys.scottish_child_payment
                )
                aggregate!( bures )
                @test ! has_any( bures, SCOTTISH_CHILD_PAYMENT )
                ti = rand( sys.scottish_child_payment.qualifying_benefits )
                bures.pers[bhead.pid].income[ti] = 1.0
                calc_scottish_child_payment!(
                    bures,
                    bu,
                    bint,
                    sys.scottish_child_payment
                )
                aggregate!( bures )
                # still should be zero since all over age
                @test ! has_any( bures, SCOTTISH_CHILD_PAYMENT )
                # all qualify on age
                ages = fill( 3, ncs )
                set_childrens_ages!( hh, ages... )
                calc_scottish_child_payment!(
                    bures,
                    bu,
                    bint,
                    sys.scottish_child_payment
                )
                aggregate!( bures )
                @test bures.income[SCOTTISH_CHILD_PAYMENT] == ncs*10
                # 1 qualifying
                set_childrens_ages!( hh, 3,8,9,10,12 )
                calc_scottish_child_payment!(
                    bures,
                    bu,
                    bint,
                    sys.scottish_child_payment
                )
                aggregate!( bures )
                @test bures.income[SCOTTISH_CHILD_PAYMENT] == 10
            end # with children
        end
    end

end


using PyCall

@pyimport theano.sandbox.cuda as cuda

const long = pybuiltin("long")

function cudandarray(blob::CuTensorBlob)
  # Create a Theano CudaNdArray wrapper for the GPU blob data
  if eltype(blob) == Float64
    error("Only Float32 is supported for Theano GPU code")
  end
  ptr = long(blob.ptr.p)
  shape = reverse(blob.shape)
  stride = 1
  strides = Int[]
  for sz in blob.shape
    push!(strides, stride)
    stride *= sz
  end
  return cuda.from_gpu_pointer(ptr, blob.shape, strides, blob)
end

function forward(backend::GPUBackend, state::TheanoLossLayerState, inputs::Vector{Blob})
  data_type = eltype(inputs[1])
  n = length(inputs[1])
  # map Mocha input to Theano
  ins = map(cudandarray, inputs)

  # TODO factor out common code cpu/gpu

  # apply the forward function to get the loss (it's returned as a 0-d array)
  loss = state.forward(ins...)[]

  state.loss += loss*state.layer.weight

  # accumulate statistics
  state.loss_accum *= state.n_accum
  state.loss_accum += state.loss * n
  state.loss_accum /= state.n_accum + n

  state.n_accum += n
end

function backward(backend::GPUBackend, state::TheanoLossLayerState,
                  inputs::Vector{Blob}, diffs::Vector{Blob})
  ins = map(cudandarray, inputs)

  for (g, diff) in zip(state.gradients, diffs)
    if isa(diff, CuTensorBlob)
      # TODO Get theano to output to diffs directly without copying

      copy!(diff, g(ins...))
    end
  end
end


#=

This file contains all types that are used in the Absorption module:

- `HitranTable` stores HITRAN transition line parameters
- `AbstractBroadeningFunctions` specify the type of line-shape broadening to perform
- `AbstractComplexErrorFunctions` specify the error function to use for Voigt calculations
- `AbstractCrossSectionModel` specify whether to calculate the line-shape from scratch or
  using an interpolator.

=#

"""
    struct HitranTable{FT}

A struct, which provides all HITRAN line parameters needed to compute
absorption cross sections

See https://hitran.org/docs/definitions-and-units/ for details

# Fields
$(DocStringExtensions.FIELDS)
"""
Base.@kwdef struct HitranTable{FT<:AbstractFloat}
    "The molecular species identification (ID) number"
    mol::Array{Int,1}
    "The isotopologue ID number"
    iso::Array{Int,1}
    "The wavenumber of the spectral line transition (cm-1) in vacuum"
    νᵢ::Array{FT,1}
    "The spectral line intensity (cm−1/(molecule·cm−2)) at Tref=296K"
    Sᵢ::Array{FT,1}
    "The Einstein-A coefficient (s-1) of a transition"
    Aᵢ::Array{FT,1}
    "The air-broadened half width at half maximum (HWHM) (cm−1/atm) at Tref=296K and reference pressure pref=1atm"
    γ_air::Array{FT,1}
    "The self-broadened half width at half maximum (HWHM) (cm−1/atm) at Tref=296K and reference pressure pref=1atm"
    γ_self::Array{FT,1}
    "The lower-state energy of the transition (cm-1)"
    E″::Array{FT,1}
    "The coefficient of the temperature dependence of the air-broadened half width"
    n_air::Array{FT,1}
    "The pressure shift (cm−1/atm) at Tref=296K and pref=1atm of the line position with respect to the vacuum transition wavenumber νij"
    δ_air::Array{FT,1}
    "The electronic and vibrational quantum numbers and labels for the upper state of a transition"
    global_upper_quanta::Array{String,1}
    "The electronic and vibrational quantum numbers and labels for the lower state of a transition"
    global_lower_quanta::Array{String,1}
    "Rotational, hyperfine and other quantum numbers and labels for the upper state of a transition"
    local_upper_quanta::Array{String,1}
    "Rotational, hyperfine and other quantum numbers and labels for the lower state of a transition"
    local_lower_quanta::Array{String,1}
    "Ordered list of indices corresponding to uncertainty estimates of transition parameters"
    ierr::Array{String,1}
    "Ordered list of reference identifiers for transition parameters"
    iref::Array{String,1}
    "A flag indicating the presence of additional data and code relating to line-mixing"
    line_mixing_flag::Array{String,1}
    "The upper state degeneracy"
    g′::Array{FT,1}
    "The lower state degeneracy"
    g″::Array{FT,1}
end

"""
    struct AbscoTable{FT}

A struct, which provides all ABSCO tabulated cross section parameters

# Fields
$(DocStringExtensions.FIELDS)
"""
Base.@kwdef struct AbscoTable{FT<:AbstractFloat}
    "The molecular species identification (ID) number"
    mol::Int
    "The isotopologue ID number"
    iso::Int
    "The wavenumber array (cm⁻¹)"
    ν::Array{Float64,1}
    "The 4D cross section table"
    σ::Array{FT,4}
    "The pressure coordinates (hPa)"
    p::Array{FT,1}
    "The temperature coordinates (K)"
    T::Array{FT,2}
end

#=

Types of broadening functions (currently Doppler, Lorentz, and Voigt)

=#

"""
    type AbstractBroadeningFunction
Abstract Broadening Function type for generic line broadening function
"""
abstract type AbstractBroadeningFunction end

"Doppler line broadening"
struct Doppler <: AbstractBroadeningFunction end

"Lorentz line broadening"
struct Lorentz <: AbstractBroadeningFunction end

"Voigt line broadening"
struct Voigt <: AbstractBroadeningFunction end

#=

Types of complex error functions for use in Voigt computations

=#

"""
    type AbstractComplexErrorFunction
Abstract Complex Error Function type for generic complex error functions
"""
abstract type AbstractComplexErrorFunction end

"Humlicek only formulation for Complex Error Function"
struct HumlicekErrorFunction <: AbstractComplexErrorFunction end

"Mix of Humlicek and Weidemann (N=32) Error Function, suggested for Voigt function"
struct HumlicekWeidemann32VoigtErrorFunction <: AbstractComplexErrorFunction end

"Mix of Humlicek and Weidemann (N=32) Error Function, suggested for Speed Dependent Voigt function"
struct HumlicekWeidemann32SDErrorFunction <: AbstractComplexErrorFunction end

"Humlicek with a single rational approximation."
struct CPF12ErrorFunction <: AbstractComplexErrorFunction end

"Mix of Humplicek and erfc Special Function for Voigt"
struct ErfcHumliErrorFunctionVoigt  <: AbstractComplexErrorFunction end

"Mix of Humplicek and erfc Special Function for Voigt"
struct ErfcHumliErrorFunctionSD  <: AbstractComplexErrorFunction end

"erfc Special Function for Voigt"
struct ErfcErrorFunction  <: AbstractComplexErrorFunction end

#=

Types of models that can be used to calculate an absorption cross-section:

- HitranModel
- InterpolationModel

=#

"""
    type AbstractCrossSectionModel
Abstract Cross Section Model type for generic cross section calculations
Currently either a HitranModel or an InterpolationModel
"""
abstract type AbstractCrossSectionModel end

"""
    struct HitranModel{FT}

A struct which provides all model parameters needed for cross-section
calculations using HITRAN data

# Fields
$(DocStringExtensions.FIELDS)
"""
Base.@kwdef struct HitranModel <: AbstractCrossSectionModel

    "Struct with hitran data"
    hitran::HitranTable
    "Broadening function (Doppler/Lorentz/Voigt)"
    broadening::AbstractBroadeningFunction
    "Wing cutoff [cm-1]"
    wing_cutoff::Real
    "VMR of gas itself [0-1]"
    vmr::Union{Real, Vector}
    "Complex Error Function to Use"
    CEF::AbstractComplexErrorFunction
    "Computer `Architecture` on which `Model` is run"
    architecture::AbstractArchitecture
end

"""
    struct InterpolationModel{FT}

A struct which provides all model parameters needed for cross-section
calculations using an Interpolator

# Fields
$(DocStringExtensions.FIELDS)
"""
Base.@kwdef struct InterpolationModel <: AbstractCrossSectionModel

    "The interpolator"
    itp

    # Everything below is metadata associated with the interpolator

    "The molecular species identification (ID) number"
    mol::Int
    "The isotopologue ID number"
    iso::Int
    "Wavelength grids"
    ν_grid::AbstractRange{<:Real}
    "Wavelength grids"
    p_grid::AbstractRange{<:Real}
    "Wavelength grids"
    t_grid::AbstractRange{<:Real}

end

# Types of Errors that may be thrown

struct HitranEmptyError <: Exception end
Base.showerror(io::IO, e::HitranEmptyError) = print(io, e, "No HITRAN records match the parameters")


# To use the solution, the input data is required

mutable struct Action
    id::String
    location_id::String
    operation_type::Int # 0- Depot, 1 - Pickup, 2 - Delivery, 3 - otherOperation # should be picked quantity
    request_id::String # if any request is linked to a request the operation
    scheduled_start_time::Float64 # post-compute
end

mutable struct Route
    id::String
    vehicle_set_id::String
    sequence::Vector{Action}
    end_status::Int # 0 returnToStartDepot, 1 returnToOtherDepot, 2 ongoing # post-compute
    # path::OSRMpath     # TODO: add OSRM path here
end

mutable struct RvrpSolution
    id::String
    instance_id::String
    cost::Float64 # post-compute
    routes::Vector{Route}
    unassigned_request_ids::Vector{String} # among semi-mandatory or optional request # post-compute
end


module tst_cov
    using Test
    using CloudCovErr
    using FITSIO
    using PyCall

    test_dir = dirname(@__FILE__)

    @testset "DECam" begin
        ref = test_dir*"/data/decapsi/c4d_170119_085651_ood_r_v1.I.fits.fz"
        @test ref == CloudCovErr.inject_rename(test_dir*"/data/decaps/c4d_170119_085651_ood_r_v1.fits.fz")

        ref_im, d_im = CloudCovErr.read_decam(test_dir*"/data/decaps/c4d_","170119_085651","r","v1","S6";corrects7=true)
        @test d_im[1,1] == 1
        @test abs(ref_im[1]-416.4855651855) < 1e-7

        decam_dir = dirname(pwd())*"/src/decam_dir"

        py"""
        import os
        os.environ['DECAM_DIR'] = $decam_dir
        """
        ref_im, d_im = CloudCovErr.read_decam(test_dir*"/data/decaps/c4d_","170119_085651","r","v1","S7";corrects7=true)
        @test d_im[1,1] == 1
        @test abs(ref_im[1]-375.392761230468) < 1e-7

        out = CloudCovErr.read_crowdsource(test_dir*"/data/decaps/","170119_085651","r","v1","S6")
        @test length(out) == 9
        wcol = out[end-1]
        w = out[end]

        out = CloudCovErr.load_psfmodel_cs(test_dir*"/data/decaps/","170119_085651","r","v1","S6")
        @test size(out(1,1,11)) == (11,11)

        if isfile(test_dir*"/data/decaps/cer/c4d_170119_085651_ooi_r_v1.cat.cer.fits")
            rm(test_dir*"/data/decaps/cer/c4d_170119_085651_ooi_r_v1.cat.cer.fits")
        end
        CloudCovErr.save_fxn(wcol,w,test_dir*"/data/decaps/","170119_085651","r","v1","S6")
        f = FITS(test_dir*"/data/decaps/cer/c4d_170119_085651_ooi_r_v1.cat.cer.fits")
        @test length(read(f["S6_CAT"],"x")) == 14572
        close(f)
        if isfile(test_dir*"/data/decaps/cer/c4d_170119_085651_ooi_r_v1.cat.cer.fits")
            rm(test_dir*"/data/decaps/cer/c4d_170119_085651_ooi_r_v1.cat.cer.fits")
        end

        f = FITS(test_dir*"/data/decaps/cat/c4d_170119_085651_ooi_r_v1.cat.fits")
        @test CloudCovErr.get_catnames(f) == ["S6","S7"]

        # run one real example end to end
        @test_nowarn CloudCovErr.proc_all(test_dir*"/data/decaps/c4d_","170119_085651","r","v1",test_dir*"/data/decaps/",ccdlist=["S6"],resume=true,corrects7=true,thr=20,outthr=20000,Np=33,widx=129,tilex=8)
        @test_nowarn CloudCovErr.proc_all(test_dir*"/data/decaps/c4d_","170119_085651","r","v1",test_dir*"/data/decaps/",ccdlist=["S6"],resume=true,corrects7=true,thr=20,outthr=20000,Np=33,widx=129,tilex=8)
        #based = test_dir*"/data/decaps/c4d_"
        #catd = test_dir*"/data/decaps/"
        #decap2f = test_dir*"/decaps2.jl"
        #@test_nowarn run(`julia $decap2f --ccdlist S6 --cS7 -r --tilex8 $based 170119_085651 r v1 $catd`)
    end
end


module TestPolyModule

using Test

using JuMP
using PolyJuMP

include("testpolymodule.jl")

function test_PolyModule()
    m = Model()
    # Triggers the creation of polydata
    @test PolyJuMP.getpolydata(m).nonnegpoly_default === nothing
    @test_throws ErrorException PolyJuMP.getdefault(m, PolyJuMP.NonNegPoly)
    @test PolyJuMP.getpolydata(m).posdefpolymatrix_default === nothing
    @test_throws ErrorException PolyJuMP.getdefault(m, PolyJuMP.PosDefPolyMatrix)
    setpolymodule!(m, DummyPolyModule)
    @test PolyJuMP.getdefault(m, PolyJuMP.NonNegPoly) == DummyPolyModule.DummyNonNeg
    @test PolyJuMP.getdefault(m, PolyJuMP.PosDefPolyMatrix) == DummyPolyModule.DummyPosDefMatrix
    PolyJuMP.setdefault!(m, PolyJuMP.PosDefPolyMatrix, DummyPolyModule.DummyNonNeg)
    @test PolyJuMP.getdefault(m, PolyJuMP.PosDefPolyMatrix) == DummyPolyModule.DummyNonNeg
end

function runtests()
    for name in names(@__MODULE__; all = true)
        if startswith("$name", "test_")
            @testset "$(name)" begin
                getfield(@__MODULE__, name)()
            end
        end
    end
end

end

TestPolyModule.runtests()


#
Random.seed!(123)
using OrdinaryDiffEq, DiffEqOperators, LinearAlgebra

"""
GOALS

1. have defaultlinsolve use predefined ldiv! if available on that type
    - allow for `Diagonal`, `LinearAlgebra.I`

2. add misc shit in each Krylov iteration call:
    https://github.com/SciML/DiffEqBase.jl/issues/366

3. Make SEM.jl operators AbstractDiffEqOperators
4. Make SEM fields subsets of AbstractVector

"""

n = 10
x = Array(range(start=0,stop=1,length=n))
A = 0.01*Tridiagonal(-ones(n-1),2ones(n),-ones(n-1))
rn = (du, u, p, t) -> begin
    mul!(du, A, u)
end
u0 = rand(n)
prob = ODEProblem(ODEFunction(rn, jac_prototype=JacVecOperator{Float64}(rn, u0; autodiff=false)), u0, (0, 10.))
#=============================================#
# DiffEqBase/linear_nonlinear.jl
#=============================================#
"""
 OrdinaryDiffEq/test/interface/linear_nonlinear_tests.jl
"""

""" default """
#sol = solve(prob, QNDF(autodiff=false)) # linsolve=DefaultLinSolve()
#@show sol.retcode
""" GMRES """
#sol = solve(prob, QNDF(autodiff=false, linsolve=LinSolveGMRES()))
#@show sol.retcode
""" CG """
#sol = solve(prob, QNDF(autodiff=false, linsolve=LinSolveCG()))
#@show sol.retcode
#=============================================#
# LinearSolve.jl
#=============================================#

#=============================================#
"""
 OrdinaryDiffEq/test/interface/mass_matrix_tests.jl

 how to provide inverse to mass_matrix. where is mass_matrix being inverted???
 how to set preconditioner,
"""

#---------------------------------------------#
return


using Distributed

addprocs(2)

@everywhere using MonteCarlo, Distributed

@everywhere function simulate(mu)
    # Connect other workers
    pool = filter(x -> x != myid(), workers())
    println(myid(), " => ", pool)
    connect(pool, wait = true)

    println("Creating model")
    model = HubbardModelAttractive(4, 2, U = 1.0, mu = mu)

    println("Creating scheduler")
    # scheduler = SimpleScheduler(
    #     LocalSweep(3), GlobalFlip(),
    #     LocalSweep(3), GlobalShuffle(),
    #     LocalSweep(3), ReplicaExchange(pool[1]),
    #     LocalSweep(3), GlobalFlip(),
    #     LocalSweep(3), GlobalShuffle(),
    #     LocalSweep(3), ReplicaPull()
    # )
    scheduler = AdaptiveScheduler(
        (
            LocalSweep(3), Adaptive(),
            LocalSweep(3), ReplicaPull(),
            LocalSweep(3), Adaptive(),
            LocalSweep(3), ReplicaExchange(pool[1])
        ),
        (GlobalFlip(), GlobalShuffle())
    )

    println("Creating dqmc")
    dqmc = DQMC(
        model,
        beta=5.0, recorder=Discarder, scheduler = scheduler,
        thermalization=10000, sweeps=10000
    )

    dqmc[:G]    = greens_measurement(dqmc, model)
    dqmc[:E]    = total_energy(dqmc, model)
    dqmc[:Occs] = occupation(dqmc, model)
    dqmc[:CDC]  = charge_density_correlation(dqmc, model)
    dqmc[:Mx]   = magnetization(dqmc, model, :x)
    dqmc[:My]   = magnetization(dqmc, model, :y)
    dqmc[:Mz]   = magnetization(dqmc, model, :z)
    dqmc[:SDCx] = spin_density_correlation(dqmc, model, :x)
    dqmc[:SDCy] = spin_density_correlation(dqmc, model, :y)
    dqmc[:SDCz] = spin_density_correlation(dqmc, model, :z)
    dqmc[:PC]   = pairing_correlation(dqmc, model, K=4)

    println("Running Simulation")
    run!(dqmc, verbose=!true)

    dqmc
end

begin
    @everywhere MonteCarlo.enable_benchmarks()
    @time dqmcs = pmap(simulate, [1.0, 0.0])
    @everywhere MonteCarlo.reset_timer!()
    @time dqmcs = pmap(simulate, [1.0, 0.0])
    MonteCarlo.show_statistics(stdout, dqmcs[1].scheduler)
    MonteCarlo.show_statistics(stdout, dqmcs[2].scheduler)
end

TOs = pmap(_ -> MonteCarlo.TimerOutputs.DEFAULT_TIMER, workers());


"""
    Model
Abstract type representing a model.

Child types of this abstract type are recognized as models by JITrench.

For more information, refer to the "Neural Networks" section of the documentation (here).
"""
abstract type Model end

"""
    plot_model(model::Model, inputs...,  to_file="model.png")
plot model by graphviz.
Note that JITrench builds the graph at the time of calculation, so some kind of input is required.
"""
function plot_model(model::Model, inputs...; to_file="model.png")
    y = model(inputs...)
    plot_graph(y, to_file=to_file)
end

#TODO: better impl
function parameters(model::Model)
    params = []
    for layer in layers(model)
        for param in parameters(layer)
            push!(params, param)
        end
    end
    return params
end

"""
    cleargrads!(model::Model; skip_uninit=false)

clear information about grad of the layers which is returned from `layers`.
if skip_uninit is true, parameters which is uninit(i.e nothing) is skipped.
"""
function cleargrads!(model::T, layers::Function; skip_uninit=false) where T <: Model
    for layer in layers(model)
        for param in parameters(layer)
            if param isa Nothing
                if skip_uninit
                    continue
                else
                    throw(DomainError("It is not possible to operate to clear uninit parameters."))
                end
            else
                param.grad = nothing
            end
        end
    end
end


"""
    SerialExecution

Functor object for using serial execution in a progressive-hedging algorithm. Create by supplying a [`Serial`](@ref) object through `execution` in the `ProgressiveHedgingSolver` factory function and then pass to a `StochasticPrograms.jl` model.

"""
struct SerialExecution{T <: AbstractFloat,
                       A <: AbstractVector,
                       S <: LQSolver,
                       PT <: PenaltyTerm} <: AbstractExecution
    subproblems::Vector{SubProblem{T,A,S,PT}}

    function SerialExecution(::Type{T}, ::Type{A}, ::Type{S}, ::Type{PT}) where {T <: AbstractFloat, A <: AbstractVector, S <: LQSolver, PT <: PenaltyTerm}
        return new{T,A,S,PT}(Vector{SubProblem{T,A,S,PT}}())
    end
end

function initialize_subproblems!(ph::AbstractProgressiveHedgingSolver, subsolver::QPSolver, penaltyterm::PenaltyTerm, execution::SerialExecution)
    for i = 1:ph.nscenarios
        push!(execution.subproblems,SubProblem(WS(ph.stochasticprogram, scenario(ph.stochasticprogram,i); solver = subsolver),
                                               i,
                                               probability(ph.stochasticprogram,i),
                                               decision_length(ph.stochasticprogram),
                                               subsolver,
                                               copy(penaltyterm)))
    end
    update_iterate!(ph)
    return ph
end

function resolve_subproblems!(ph::AbstractProgressiveHedgingSolver, execution::SerialExecution{T,A}) where {T <: AbstractFloat, A <: AbstractVector}
    Qs = A(undef, length(execution.subproblems))
    # Update subproblems
    reformulate_subproblems!(execution.subproblems, ph.ξ, penalty(ph))
    # Solve sub problems
    for (i, subproblem) ∈ enumerate(execution.subproblems)
        Qs[i] = subproblem(ph.ξ)
    end
    # Return current objective value
    return sum(Qs)
end

function update_iterate!(ph::AbstractProgressiveHedgingSolver, execution::SerialExecution)
    # Update the estimate
    ξ_prev = copy(ph.ξ)
    ph.ξ[:] = sum([subproblem.π*subproblem.x for subproblem in execution.subproblems])
    # Update δ₁
    ph.data.δ₁ = norm(ph.ξ-ξ_prev, 2)^2
    return nothing
end

function update_subproblems!(ph::AbstractProgressiveHedgingSolver, execution::SerialExecution)
    # Update dual prices
    update_subproblems!(execution.subproblems, ph.ξ, penalty(ph))
    return nothing
end

function update_dual_gap!(ph::AbstractProgressiveHedgingSolver, execution::SerialExecution)
    # Update δ₂
    ph.data.δ₂ = sum([s.π*norm(s.x-ph.ξ,2)^2 for s in execution.subproblems])
    return nothing
end

function calculate_objective_value(ph::AbstractProgressiveHedgingSolver, execution::SerialExecution)
    return sum([get_objective_value(s) for s in execution.subproblems])
end

function fill_first_stage!(ph::AbstractProgressiveHedgingSolver, stochasticprogram::StochasticProgram, nrows::Integer, ncols::Integer, execution::SerialExecution)
    subproblem = execution.subproblems[1]
    μ = try
        MPB.getreducedcosts(subproblem.solver.lqmodel)[1:ncols]
    catch
        fill(NaN, ncols)
    end
    StochasticPrograms.set_first_stage_redcosts!(stochasticprogram, μ)
    λ = try
        MPB.getconstrduals(subproblem.solver.lqmodel)[1:nrows]
    catch
        fill(NaN, nrows)
    end
    StochasticPrograms.set_first_stage_duals!(stochasticprogram, λ)
    return nothing
end

function fill_submodels!(ph::AbstractProgressiveHedgingSolver, scenarioproblems::StochasticPrograms.ScenarioProblems, nrows::Integer, ncols::Integer, execution::SerialExecution)
    for (i, submodel) in enumerate(scenarioproblems.problems)
        fill_submodel!(submodel, execution.subproblems[i], nrows, ncols)
    end
    return nothing
end

function fill_submodels!(ph::AbstractProgressiveHedgingSolver, scenarioproblems::StochasticPrograms.DScenarioProblems, nrows::Integer, ncols::Integer, execution::SerialExecution)
    j = 0
    @sync begin
        for w in workers()
            n = remotecall_fetch((sp)->length(fetch(sp).problems), w, scenarioproblems[w-1])
            for i in 1:n
                k = i+j
                @async remotecall_fetch((sp,i,x,μ,λ,nrows,ncols) -> fill_submodel!(fetch(sp).problems[i],x,μ,λ,nrows,ncols),
                                        w,
                                        scenarioproblems[w-1],
                                        i,
                                        get_solution(execution.subproblems[k])...,
                                        nrows,
                                        ncols)
            end
            j += n
        end
    end
    return nothing
end
# API
# ------------------------------------------------------------
function (execution::Serial)(::Type{T}, ::Type{A}, ::Type{S}, ::Type{PT}) where {T <: AbstractFloat, A <: AbstractVector, S <: LQSolver, PT <: PenaltyTerm}
    return SerialExecution(T,A,S,PT)
end

function str(::Serial)
    return "serial execution"
end


function score!(cur_player::AbstractPlayer, cur_label::AbstractString, cur_value::Int)

  # --------------------
  #  verify score comes
  #   from valid label
  # --------------------

  in(cur_label, allowed_labels) ||
    error("Invalid label in score call: $(cur_label).")

  # ------------------
  #  handle soft drop
  # ------------------

  if cur_label == "soft"
    cur_player.score += 1
    cur_player.stash += 1

    return 1
  end

  # ---------------------
  #  calculate the score
  # ---------------------

  cur_dict = OrderedDict(
    "level" => cur_player.level,
    "lines" => cur_player.lines,
    "combo" => cur_player.combo,
    "stash" => cur_player.stash,
    "glues" => cur_player.glues,
    "total" => cur_player.score,
    "chain" => cur_player.chain,
    "piece" => cur_player.piece.name
  )

  cur_function = getfield(
    Tetris,
    Symbol("_score_$(cur_label)")
  )

  cur_score = cur_function(cur_value, cur_dict)

  # -------------------
  #  log current score
  # -------------------

  cur_round = cur_player.round

  if cur_round.is_keeping_score
    cur_position, cur_round.log_count =
      [0, 1] .+ cur_round.log_count

    cur_log = Log(
      cur_position,
      cur_label,
      cur_value,
      cur_score,
      values(cur_dict)...
    )

    push!(cur_round.logs, cur_log)
  end

  # ---------------------
  #  update player score
  # ---------------------

  if cur_label == "glue" || cur_label == "hold"
    cur_player.stash = 0
  else
    cur_player.score += cur_score
  end

  if is_repl
    cur_string = []
    push!(cur_string, crayon_dict["grey"])
    push!(cur_string, "\x1b[$(14);$(11)H")
    push!(cur_string, lpad(cur_player.score, 8, "0"))
    push!(cur_string, inv(crayon_dict["grey"]))
    push!(cur_string, "\x1b[u")
    print(cur_string...)
  end

  return cur_score

end

function _score_clear(cleared_count::Int, cur_dict::OrderedDict)
  if cleared_count == 4 && cur_dict["chain"] == "tetris"
    cur_score = 1200
  else
    cur_score = score_array[cleared_count]
  end

  cur_score *= cur_dict["level"]

  cur_score
end

function _score_combo(cleared_count::Int, cur_dict::OrderedDict)
  cur_score = ( cleared_count == 1 ) ? 20 : 50

  cur_score *= cur_dict["combo"]

  cur_score *= cur_dict["level"]

  cur_score
end

function _score_hard(drop_count::Int, cur_dict::OrderedDict)
  cur_score = drop_count

  cur_score *= 2

  cur_score
end

function _score_glue(stub_value::Int, cur_dict::OrderedDict)
  cur_score = cur_dict["stash"]

  cur_score
end

function _score_hold(stub_value::Int, cur_dict::OrderedDict)
  cur_score = cur_dict["stash"]

  cur_score
end


#################### ModelChains ####################

#################### Constructors ####################

function ModelChains(c::Chains, m::Model)
  ModelChains(c.value, c.range, c.names, c.chains, m)
end

#################### Conversions ####################

Base.convert(::Type{Chains}, mc::ModelChains) =
  Chains(mc.value, mc.range, mc.names, mc.chains)

#################### Indexing ####################

function Base.getindex(mc::ModelChains, window, names, chains)
  c = getindex(convert(Chains, mc), window, names2inds(mc, names), chains)
  ModelChains(c, mc.model)
end

names2inds(mc::ModelChains, nodekey::Symbol) = names2inds(mc, [nodekey])

function names2inds(mc::ModelChains, nodekeys::Vector{Symbol})
  inds = Int[]
  missing = Symbol[]
  for key in nodekeys
    keyinds = indexin(names(mc.model, all=key), mc.names)
    0 in keyinds ? push!(missing, key) : append!(inds, keyinds)
  end
  if !isempty(missing)
    throw(ArgumentError(string(
      "chain values are missing for nodes : ",
      join(map(string, missing), ", ")
    )))
  end
  inds
end

function Base.keys(mc::ModelChains, ntype::Symbol, at...)
  values = Symbol[]
  m = mc.model
  nodekeys = ntype == :dependent ?
    keys(m, :dependent) :
    intersect(keys(m, ntype, at...), keys(m, :dependent))
  for key in nodekeys
    all(name -> name in mc.names, names(m, all=key)) && push!(values, key)
  end
  values
end

#################### Auxilliary Functions ####################

function link(c::ModelChains)
  cc = copy(c.value)
  inds_queue = 1:length(c.names)
  for key in intersect(keys(c.model, :monitor), keys(c.model, :stochastic))
    node = c.model[key]
    inds = findall(in(names(node)), c.names)
    if !isempty(inds)
      f(x) = unlist(node, relist(node, x), true)
      cc[:, inds, :] = mapslices(f, cc[:, inds, :], 2)
      inds_queue = setdiff(inds_queue, inds)
    end
  end
  for j in inds_queue
    x = cc[:, j, :]
    if minimum(x) > 0.0
      cc[:, j, :] = maximum(x) < 1.0 ? StatsFuns.logit.(x) : log.(x)
    end
  end
  cc
end


"""
    FieldModule

Module for abstract fields.
"""
module FieldModule

using FinEtools.FTypesModule: FInt, FFlt, FCplxFlt, FFltVec, FIntVec, FFltMat, FIntMat, FMat, FVec, FDataDict
import Base.copyto!

"""
    Abstract field.

Expected  attributes:
  + values::FMat{T} # Array of degree of freedom parameters,  indexed by node
  + dofnums::FIntMat # Array of degree of freedom numbers, indexed by node
  + is_fixed::Matrix{Bool} # Array of Boolean flags, indexed by node
  + fixed_values::FMat{T} # Array of fixed values, indexed by node
  + nfreedofs::FInt # Total number of free degrees of freedom
See the macro `add_Field_fields()`.
"""
abstract type Field end

macro add_Field_fields()
    return esc(:(
    values::Array{T, 2};
    dofnums::Array{FInt, 2};
    is_fixed::Matrix{Bool};
    fixed_values::Array{T, 2};;
    nfreedofs::FInt;
    )
    )
end

"""
    ndofs{F<:Field}(self::T)

Dimension of the degree of freedom parameters (i. e.  how many degrees of
freedom per entity).
"""
ndofs(self::Field)  = size(self.values, 2)

"""
    nents{F<:Field}(self::F)

Number of nodes associated with the field.
"""
nents(self::Field)  = size(self.values, 1)

"""
    copyto!(DEST::F,  SRC::F) where {F<:Field}

Copy data from one field to another.
"""
function copyto!(DEST::F,  SRC::F) where {F<:Field}
    copyto!(DEST.values, SRC.values)
    copyto!(DEST.dofnums, SRC.dofnums)
    copyto!(DEST.is_fixed, SRC.is_fixed)
    copyto!(DEST.fixed_values, SRC.fixed_values)
    DEST.nfreedofs = SRC.nfreedofs
    return  DEST
end

"""
    wipe!(self::Field)

Wipe all the data from the field.
"""
function wipe!(self::Field)
    Zer = zero(eltype(self.fixed_values[1]))
    self.nfreedofs = 0
    fill!(self.dofnums, 0)
    fill!(self.is_fixed, false)
    fill!(self.values, Zer)
    fill!(self.fixed_values, Zer)
    return self
end

"""
    gathersysvec{F<:Field}(self::F)

Gather values from the field for the whole system vector.
"""
function gathersysvec(self::F) where {F<:Field}
    nents,dim = size(self.values)
    vec = zeros(typeof(self.values[1,1]),self.nfreedofs,1)
    for i = 1:nents
        for j = 1:dim
            en = self.dofnums[i,j]
            if (en > 0) && (en <= self.nfreedofs)
                vec[en] = self.values[i,j]
            end
        end
    end
    return vec
end

"""
    gathersysvec!(self::F, vec::FVec{T}) where {F<:Field, T}

Gather values from the field for the whole system vector.
"""
function gathersysvec!(self::F, vec::FVec{T}) where {F<:Field, T}
    nents,dim = size(self.values)
    @assert length(vec) == self.nfreedofs
    for i = 1:nents
        for j = 1:dim
            en = self.dofnums[i,j]
            if (en > 0) && (en <= self.nfreedofs)
                vec[en] = self.values[i,j]
            end
        end
    end
    return vec
end

"""
    gathervalues_asvec!(self::Field, dest::AbstractArray{T, 1},
        conn::CC) where {CC, T}

Gather values from the field into a vector.

The order is: for each node  in the connectivity, copy into the buffer all the
degrees of freedom,  then the next node and so on.

`dest` = destination buffer: overwritten  inside,  must be preallocated
in the correct size
"""
function gathervalues_asvec!(self::Field, dest::AbstractArray{T, 1}, conn::CC) where {CC, T}
    en::FInt = 1;
    for i = 1:length(conn)
        for j = 1:size(self.values,2)
            dest[en] = self.values[conn[i],j];
            en = en + 1;
        end
    end
    return dest
end

"""
    gathervalues_asmat!(self::Field, dest::AbstractArray{T, 2},
        conn::CC) where {CC, T}

Gather values from the field into a two-dimensional array.

The order is: for each node  in the connectivity, copy into the corresponding
row of the buffer all the degrees of freedom,  then the next node into the next
row and so on.

`dest` = destination buffer: overwritten  inside,  must be preallocated
in the correct size
"""
function gathervalues_asmat!(self::Field, dest::AbstractArray{T, 2},    conn::CC) where {CC, T}
    for i = 1:length(conn)
        for j = 1:size(self.values,2)
            dest[i, j] = self.values[conn[i], j];
        end
    end
    return dest
end

"""
    gatherfixedvalues_asvec!(self::Field, dest::AbstractArray{T, 1},
        conn::CC) where {CC, T}

Gather FIXED values from the field into a vector.

The order is: for each node  in the connectivity, copy into the buffer all the
fixed degrees of freedom,  then the next node and so on. If a degree of freedom
is NOT fixed, the corresponding entry is  set to zero.

`dest` = destination buffer: overwritten  inside,  must be preallocated
in the correct size
"""
function gatherfixedvalues_asvec!(self::Field, dest::AbstractArray{T, 1},    conn::CC) where {CC, T}
    en::FInt = 1;
    for i = 1:length(conn)
        for j = 1:size(self.fixed_values,2)
            if self.is_fixed[conn[i],j] # free degree of freedom
                dest[en] = self.fixed_values[conn[i], j];
            else
                dest[en] = 0.0
            end
            en = en + 1;
        end
    end
    return dest
end

"""
    gatherfixedvalues_asmat!(self::Field, dest::AbstractArray{T, 2},
        conn::CC) where {CC, T}

Gather FIXED values from the field into a two-dimensional array.

The order is: for each node  in the connectivity, copy into the corresponding
row of the buffer all the degrees of freedom,  then the next node into the next
row and so on.  If a degree of freedom
is NOT fixed, the corresponding entry is  set to zero.

`dest` = destination buffer: overwritten  inside,  must be preallocated
in the correct size
"""
function gatherfixedvalues_asmat!(self::Field, dest::AbstractArray{T, 2},    conn::CC) where {CC, T}
    for i = 1:length(conn)
        for j = 1:size(self.fixed_values,2)
            if self.is_fixed[conn[i],j] # fixed degree of freedom
                dest[i, j] = self.fixed_values[conn[i], j];
            else
                dest[i, j] = 0.0
            end
        end
    end
    return dest
end

function anyfixedvaluenz(self::Field, conn::CC) where {CC}
    for i = 1:length(conn)
        for j = 1:size(self.fixed_values,2)
            if self.is_fixed[conn[i],j] # free degree of freedom
                if  abs(self.fixed_values[conn[i], j]) > 0.0
                    return true
                end
            end
        end
    end
    return false
end

"""
    gatherdofnums!(self::Field, dest::A, conn::CC) where {A, CC}

Gather dofnums from the field.
"""
function gatherdofnums!(self::Field, dest::A, conn::CC) where {A, CC}
    en::FInt = 1;
    for i = 1:length(conn)
        for j = 1:size(self.dofnums,2)
            dest[en] = self.dofnums[conn[i],j];
            en = en+1;
        end
    end
    return dest
end

"""
    numberdofs!(self::Field)

Number the degrees of freedom.

The free components in the field are numbered consecutively. No effort is
made to optimize the numbering in any way. If you'd like to optimize the
numbering of the degrees of freedom, use the above form that sets the
permutation of the degrees of freedom, or the permutation of the nodes.
"""
function numberdofs!(self::Field)
    fixed_dofnum::FInt = 0
    nents,dim = size(self.values)
    self.nfreedofs::FInt =0
    for i=1:nents
        for j=1:dim
            if !self.is_fixed[i,j] # free degree of freedom
                self.nfreedofs = self.nfreedofs + 1
                self.dofnums[i,j] = self.nfreedofs
            else # fixed degree of freedom: no equation
                self.dofnums[i,j] = fixed_dofnum
            end
        end
    end
    return  self
end

"""
    setebc!(self::Field, fenids::FIntVec, is_fixed::Bool, comp::FInt,
      val::FVec{T}) where {T<:Number}

Set the EBCs (essential boundary conditions).

`fenids`         - array of N node identifiers
`is_fixed` = scaler Boolean: are the degrees of freedom being fixed (true)
             or released (false),
`comp` = integer, which  degree of freedom (component),
`val` = array of N values of type T

Note:  Any call to setebc!() potentially changes the current assignment
which degrees of freedom are free and which are fixed
and therefore is presumed to invalidate the
current degree-of-freedom numbering. In such a case this method sets
`nfreedofs = 0`; and  `dofnums=0`.
"""
function setebc!(self::Field, fenids::FIntVec, is_fixed::Bool, comp::FInt,    val::FVec{T}) where {T<:Number}
    @assert comp <= size(self.values,2) "Requested  nonexistent  degree of freedom"
    @assert maximum(fenids) <= size(self.values,1) "Requested nonexistent node"
    @assert size(fenids) == size(val) "Arrays of mismatched sizes"
    for  j = 1:length(fenids)
        self.is_fixed[fenids[j],comp] = is_fixed;
        if self.is_fixed[fenids[j],comp]
            self.fixed_values[fenids[j],comp] = val[j];
        else
            self.fixed_values[fenids[j],comp] = zero(T)
        end
    end
    self.nfreedofs = 0
    fill!(self.dofnums, 0)
    return  self
end

"""
    setebc!(self::Field, fenids::FIntVec, is_fixed::Bool, comp::FInt,
      val::T) where {T<:Number}

Set the EBCs (essential boundary conditions).

`fenids`         - array of N node identifiers
`is_fixed` = scaler Boolean: are the degrees of freedom being fixed (true)
             or released (false),
`comp` = integer, which  degree of freedom (component),
`val` = scalar of type T

Note:  Any call to setebc!() potentially changes the current assignment
which degrees of freedom are free and which are fixed
and therefore is presumed to invalidate the
current degree-of-freedom numbering. In such a case this method sets
`nfreedofs = 0`; and  `dofnums=0`.
"""
function setebc!(self::Field, fenids::FIntVec, is_fixed::Bool, comp::FInt,    val::T) where {T<:Number}
    @assert (comp >= 1 && comp <= size(self.values,2)) "Requested  nonexistent  degree of freedom"
    @assert maximum(fenids) <= size(self.values,1) "Requested nonexistent node"
    @assert minimum(fenids) >= 1 "Requested nonexistent node"
    for  j = 1:length(fenids)
        self.is_fixed[fenids[j],comp] = is_fixed;
        if self.is_fixed[fenids[j],comp]
            self.fixed_values[fenids[j],comp] = val;
        else
            self.fixed_values[fenids[j],comp] = zero(T)
        end
    end
    self.nfreedofs = 0
    fill!(self.dofnums, 0)
    return  self
end

"""
    setebc!(self::Field, fenids::FIntVec, comp::FInt,
      val::FVec{T}) where {T<:Number}

Set the EBCs (essential boundary conditions).

`fenids`         - array of N node identifiers
`comp` = integer, which  degree of freedom (component),
`val` = array of N values of type T

Note:  Any call to setebc!() potentially changes the current assignment
which degrees of freedom are free and which are fixed
and therefore is presumed to invalidate the
current degree-of-freedom numbering. In such a case this method sets
`nfreedofs = 0`; and  `dofnums=0`.
"""
function setebc!(self::Field, fenids::FIntVec, comp::FInt, val::FVec{T}) where {T<:Number}
    return setebc!(self, fenids, true, comp, val)
end

"""
    setebc!(self::Field, fenids::FIntVec, comp::FInt;
      val::T=0.0) where {T<:Number}

Set the EBCs (essential boundary conditions).

`fenids`         - array of N node identifiers
`comp` = integer, which  degree of freedom (component),
`val` = scalar of type T

Note:  Any call to setebc!() potentially changes the current assignment
which degrees of freedom are free and which are fixed
and therefore is presumed to invalidate the
current degree-of-freedom numbering. In such a case this method sets
`nfreedofs = 0`; and  `dofnums=0`.
"""
function setebc!(self::Field, fenids::FIntVec, comp::FInt; val::T=0.0) where {T<:Number}
    return setebc!(self, fenids, true, comp, val)
end

"""
    setebc!(self::Field, fenids::FIntVec, is_fixed::Bool, comp::FInt;
      val::T=0.0) where {T<:Number}

Set the EBCs (essential boundary conditions).

`fenids`         - array of N node identifiers
`comp` = integer, which  degree of freedom (component),
`val` = scalar of type T

Note:  Any call to setebc!() potentially changes the current assignment
which degrees of freedom are free and which are fixed
and therefore is presumed to invalidate the
current degree-of-freedom numbering. In such a case this method sets
`nfreedofs = 0`; and  `dofnums=0`.
"""
function setebc!(self::Field, fenids::FIntVec, is_fixed::Bool, comp::FInt;    val::T=0.0) where {T<:Number}
    j = comp
    @assert (j >= 1) && (j <= ndofs(self))
    setebc!(self, fenids, is_fixed, j, val)
    return self
end

"""
    setebc!(self::Field, fenids::FIntVec)

Set the EBCs (essential boundary conditions).

Suppress all degrees of freedom at the given nodes.

`fenids`         - array of N node identifiers

Note:  Any call to setebc!() potentially changes the current assignment
which degrees of freedom are free and which are fixed
and therefore is presumed to invalidate the
current degree-of-freedom numbering. In such a case this method sets
`nfreedofs = 0`; and  `dofnums=0`.
"""
function setebc!(self::Field, fenids::FIntVec)
    Zer = zero(eltype(self.fixed_values[1]))
    for comp = 1:size(self.values, 2)
        setebc!(self, fenids, true, comp, Zer)
    end
    return self
end

"""
    setebc!(self::Field, fenid::FInt)

Set the EBCs (essential boundary conditions).

Suppress all degrees of freedom at the given node.

`fenid`         - One integer as a node identifier

Note:  Any call to setebc!() potentially changes the current assignment
which degrees of freedom are free and which are fixed
and therefore is presumed to invalidate the
current degree-of-freedom numbering. In such a case this method sets
`nfreedofs = 0`; and  `dofnums=0`.
"""
function setebc!(self::Field, fenid::FInt)
    Zer = zero(eltype(self.fixed_values[1]))
    for comp = 1:size(self.values, 2)
        setebc!(self, [fenid], true, comp, Zer)
    end
    return self
end

"""
    setebc!(self::Field)

Set the EBCs (essential boundary conditions).

All essential boundary conditions are CLEARED.

Note:  Any call to setebc!() potentially changes the current assignment
which degrees of freedom are free and which are fixed
and therefore is presumed to invalidate the
current degree-of-freedom numbering. In such a case this method sets
`nfreedofs = 0`; and  `dofnums=0`.
"""
function setebc!(self::Field)
    self.nfreedofs = 0
    fill!(self.dofnums, 0)
    fill!(self.is_fixed, false)
    fill!(self.fixed_values, zero(eltype(self.fixed_values[1])))
    return  self
end

"""
    applyebc!(self::Field)

Apply EBCs (essential boundary conditions).
"""
function applyebc!(self::Field)
    nents,dim = size(self.values);
    for i = 1:nents
        for j = 1:dim
            if self.is_fixed[i,j]
                self.values[i,j] = self.fixed_values[i,j];
            end
        end
    end
    return  self
end

"""
    scattersysvec!(self::Field, vec::FVec{T}) where {T<:Number}

Scatter values to the field from a system vector.
"""
function scattersysvec!(self::Field, vec::FVec{T}) where {T<:Number}
    nents,dim = size(self.values);
    for i = 1:nents
        for j = 1:dim
            dn = self.dofnums[i,j];
            if (dn > 0) && (dn <= self.nfreedofs)
                self.values[i,j] = vec[dn];
            end
        end
    end
    return  self
end

end


############################################################################################
# Types to distinguish between Gibbs and Metropolis updating
"Referencing struct: no updating via reference index"
struct Marginal <: ParticleReferencing end
"Referencing struct: last element in particle history always reference index"
struct Conditional <: ParticleReferencing end
"Referencing struct: last element in particle history weighted via reference index"
struct Ancestral <: ParticleReferencing end #PG with Ancestor Sampling

############################################################################################
"""
$(SIGNATURES)
Assign reference, depending on whether reference is tracked in Model.

# Examples
```julia
```

"""
function get_reference(referencing, model::ModelWrapper, data::D, tagged::Tagged) where {D}
    ## Check if only 1 parameter is tracked that has equal number of samples as data
    ArgCheck.@argcheck length(tagged.parameter) == 1 "Only 1 parameter can be tracked for Particle Filter."
    referenceᵖ = getindex(model.val, keys(tagged.parameter)[begin])
    return referenceᵖ
end
function get_reference(referencing, objective::Objective)
    return get_reference(referencing, objective.model, objective.data, objective.tagged)
end

############################################################################################
"""
$(SIGNATURES)
Inplace calculate ancestor weights.

# Examples
```julia
```

"""
function ancestralweight!(
    weights::BaytesCore.ParameterWeights,
    valₜ::Union{P,AbstractArray{P}},
    kernel::ParticleKernel,
    val::AbstractArray,
    iter::Integer,
) where {P}
    #!NOTE: buffer first overwritten in ℓtransition! and then assigned again on LHS
    weights.buffer .=
        weights.ℓweights .+ ℓtransition!(weights.buffer, valₜ, kernel, val, iter)
    weights.buffer .-= logsumexp(weights.buffer)
    return nothing
end

############################################################################################
"""
$(SIGNATURES)
Sample ancestor reference for particle history.

# Examples
```julia
```

"""
function sample_ancestor(
    _rng::Random.AbstractRNG,
    weights::BaytesCore.ParameterWeights,
    valₜ::Union{P,AbstractArray{P}},
    kernel::ParticleKernel,
    val::AbstractArray,
    iter::Integer,
) where {P}
    #!NOTE: buffer will be first overwritten on right side to inplace calculate new weights, and then assigned inplace on left side
    ancestralweight!(weights, valₜ, kernel, val, iter)
    #!NOTE Need to exp() particles.weights.buffer due to ancestralweights! calculation in log space
    weights.buffer .= exp.(weights.buffer)
    return BaytesCore.randcat(_rng, weights.buffer)
end

############################################################################################
"""
$(SIGNATURES)
Resampling function for particles, dispatched on ParticleReferencing subtype.

# Examples
```julia
```

"""
function ancestors!(
    _rng::Random.AbstractRNG,
    referencing::ParticleReferencing,
    resampling::R,
    ancestor::AbstractArray, #referencing::Marginal
    iter::Integer,
    Nparticles::Integer,
    weights::AbstractVector{F},
) where {R<:BaytesCore.ResamplingMethod,F<:AbstractFloat}
    #!NOTE: Weights.buffer already exp.(weightsₙ) during BaytesCore.issmaller( BaytesCore.computeESS(particles.weights), X) in previous step
    resample!(
        _rng,
        resampling,
        ancestor,
        iter,
        weights,
        isa(referencing, Marginal) ? Nparticles : Nparticles - 1,
    )
    return nothing
end

############################################################################################
"""
$(SIGNATURES)
Set last particle at current iteration.

# Examples
```julia
```

"""
function set_reference!(
    referencing::Marginal,
    val::AbstractVector{P},
    reference::AbstractVector{P},
    iter::Integer,
) where {P}
    #!NOTE: Nothing to be done in Marginal case
    return nothing
end
function set_reference!(
    referencing::R, val::AbstractVector{P}, reference::AbstractVector{P}, iter::Integer
) where {R<:Union{Conditional,Ancestral},P}
    #!NOTE: Reference particle at t is the same for PG and PGAS
    val[end] = reference[iter]
    return nothing
end
function set_reference!(
    referencing, val::AbstractMatrix{P}, reference::AbstractVector{P}, iter::Integer
) where {P}
    return set_reference!(referencing, view(val, :, iter), reference, iter)
end

############################################################################################
"""
$(SIGNATURES)
Compute new reference index from current reference index.

# Examples
```julia
```

"""
function get_reference!(
    _rng::Random.AbstractRNG,
    referencing::Marginal,
    ancestor::AbstractVector{I},
    particles::F,
    tune::ParticleFilterTune,
    reference::Union{P,AbstractArray{P}}
) where {I<:Integer, F<:AbstractParticles, P}
    return nothing
end

function get_reference!(
    _rng::Random.AbstractRNG,
    referencing::Conditional,
    ancestor::AbstractVector{I},
    particles::F,
    tune::ParticleFilterTune,
    reference::Union{P,AbstractArray{P}}
) where {I<:Integer, F<:AbstractParticles, P}
    ancestor[end] = length(ancestor)
    return nothing
end

function get_reference!(
    _rng::Random.AbstractRNG,
    referencing::Ancestral,
    ancestor::AbstractVector{I},
    particles::F,
    tune::ParticleFilterTune,
    reference::Union{P,AbstractArray{P}}
) where {I<:Integer, F<:AbstractParticles, P}
    referenceₜ = BaytesCore.grab(reference, tune.iter.current, tune.config.particle)
    #!NOTE: Weights.buffer already exp.(weightsₙ) during BaytesCore.issmaller( BaytesCore.computeESS(particles.weights), X) in previous step
    ancestor[end] = sample_ancestor(
        _rng,
        particles.weights,
        referenceₜ,
        particles.kernel,
        particles.val,
        tune.iter.current
    )
    return nothing
end

#!NOTE: If buffer is a matrix, reference will be assigned at iter-1 within function, but input here is iter as we grab data for forward looking ancestral step.
function get_reference!(
    _rng::Random.AbstractRNG,
    referencing::R,
    ancestor::AbstractMatrix{I},
    particles::F,
    tune::ParticleFilterTune,
    reference::Union{P,AbstractArray{P}}
) where {R<:ParticleReferencing, I<:Integer, F<:AbstractParticles, P}
    return get_reference!(
        _rng,
        referencing,
        view(ancestor, :, tune.iter.current - 1),
        particles,
        tune,
        reference
    )
end
function get_reference!(
    _rng::Random.AbstractRNG,
    particles::F,
    tune::ParticleFilterTune,
    reference::Union{P,AbstractArray{P}}
) where {F<:AbstractParticles, P}
    return get_reference!(_rng, tune.referencing, particles.ancestor, particles, tune, reference)
end

############################################################################################
#export
export
    Marginal,
    Conditional,
    Ancestral,
    get_reference,
    get_reference!,
    set_reference!,
    sample_ancestor


function AESCTR(plaintext, iv::Array{UInt8, 1}, key::AbstractAESKey, cache::AbstractAESCache)
    len = length(plaintext)
    ciphertext = similar(Array{UInt8, 1}, len)
    AESCTR!(ciphertext, plaintext, iv, key, cache)
end

function AESCTR!(ciphertext, plaintext, iv::Array{UInt8, 1}, key::AbstractAESKey, cache::AbstractAESCache)
    len = length(plaintext)
    pad = 16 - (len % 16)
    iters = Int((len + pad) / 16)
    counter = cache.modecache
    counter .= iv
    ciphertextblock = AESBlock(ciphertext)
    plaintextblock = AESBlock(plaintext)
    for i in 1:(iters-1)
        AESEncryptBlock!(ciphertextblock, counter, key, cache)
        @. ciphertextblock ⊻= plaintextblock
        increment!(counter)
        increment!(ciphertextblock)
        increment!(plaintextblock)
    end
    # for the last iteration, we dont need counter no more
    # we can use the inplace property of AES
    AESEncryptBlock!(counter, counter, key, cache)
    offset = 16(iters-1)
    for i in 1:(16-pad)
        ciphertext[offset+i] = UInt8(plaintext[offset+i]) ⊻ counter[i]
    end
    ciphertext
end

# Encryption is same as decryption in CTR mode for a give Key and IV
const AESCTR_D = AESCTR
const AESCTR_D! = AESCTR!


# random functions that dispatch either to CURAND or GPUArrays' generic RNG

using Random

export rand_logn!, rand_poisson!

# native RNG

"""
    CUDA.RNG()

A random number generator using `rand()` in a device kernel.

!!! warning

    This generator is experimental, and not used by default yet. Currently, it does not
    pass SmallCrush (but neither does the current GPUArrays.jl-based generator). If you
    can, use a CURAND-based generator for now.

See also: [SharedTauswortheGenerator](@ref)
"""
struct RNG <: AbstractRNG
    state::CuVector{UInt32}

    function RNG(seed)
        @assert length(seed) == 32
        new(seed)
    end
end

RNG() = RNG(rand(UInt32, 32))

function Random.seed!(rng::RNG, seed::AbstractVector{UInt32})
    @assert length(seed) == 32
    copyto!(rng.state, seed)
end

function Random.seed!(rng::RNG)
    Random.rand!(rng.state)
    return
end

function Random.rand!(rng::RNG, A::AnyCuArray)
    function kernel(A::AbstractArray{T}, state::AbstractVector{UInt32}) where {T}
        device_rng = Random.default_rng()

        # initialize the state
        tid = threadIdx().x
        if tid <= 32
            # we know for sure the seed will contain 32 values
            @inbounds Random.seed!(device_rng, state)
        end

        sync_threads()

        # grid-stride loop
        offset = (blockIdx().x - 1) * blockDim().x
        while offset < length(A)
            # generating random numbers synchronizes threads, so needs to happen uniformly.
            val = Random.rand(device_rng, T)

            i = tid + offset
            if i <= length(A)
                @inbounds A[i] = val
            end

            offset += (blockDim().x - 1) * gridDim().x
        end

        sync_threads()

        # save the device rng state of the first block (other blocks are derived from it)
        # so that subsequent launches generate different random numbers
        if blockIdx().x == 1 && tid <= 32
            @inbounds state[tid] = device_rng.state[tid]
        end

        return
    end

    kernel = @cuda launch=false name="rand!" kernel(A, rng.state)
    config = launch_configuration(kernel.fun; max_threads=64)
    threads = max(32, min(config.threads, length(A)))
    blocks = min(config.blocks, cld(length(A), threads))
    kernel(A, rng.state; threads=threads, blocks=blocks)

    # XXX: updating the state from within the kernel is racey,
    #      so we should have a RNG per stream

    A
end

# TODO: `randn!`; cannot reuse from Base or RandomNumbers, as those do scalar indexing

# generic functionality

function Random.rand!(rng::Union{RNG,CURAND.RNG,GPUArrays.RNG}, A::AbstractArray{T}) where {T}
    B = CuArray{T}(undef, size(A))
    Random.rand!(rng, B)
    copyto!(A, B)
end

function Random.rand(rng::Union{RNG,CURAND.RNG,GPUArrays.RNG}, T::Type)
    assertscalar("scalar rand")
    A = CuArray{T}(undef, 1)
    Random.rand!(rng, A)
    A[]
end

# RNG-less interface

# the interface is split in two levels:
# - functions that extend the Random standard library, and take an RNG as first argument,
#   will only ever dispatch to CURAND and as a result are limited in the types they support.
# - functions that take an array will dispatch to either CURAND or GPUArrays
# - non-exported functions are provided for constructing GPU arrays from only an eltype

const curand_rng = CURAND.default_rng
gpuarrays_rng() = GPUArrays.default_rng(CuArray)

function seed!(seed=Base.rand(UInt64))
    Random.seed!(curand_rng(), seed)
    Random.seed!(gpuarrays_rng(), seed)
end

# CURAND in-place
Random.rand!(A::CURAND.UniformArray) = Random.rand!(curand_rng(), A)
Random.randn!(A::CURAND.NormalArray; kwargs...) = Random.randn!(curand_rng(), A; kwargs...)
rand_logn!(A::CURAND.LognormalArray; kwargs...) = CURAND.rand_logn!(curand_rng(), A; kwargs...)
rand_poisson!(A::CURAND.PoissonArray; kwargs...) = CURAND.rand_poisson!(curand_rng(), A; kwargs...)

# CURAND out-of-place
rand(T::CURAND.UniformType, dims::Dims) = Random.rand(curand_rng(), T, dims)
randn(T::CURAND.NormalType, dims::Dims; kwargs...) = Random.randn(curand_rng(), T, dims; kwargs...)
rand_logn(T::CURAND.LognormalType, dims::Dims; kwargs...) = CURAND.rand_logn(curand_rng(), T, dims; kwargs...)
rand_poisson(T::CURAND.PoissonType, dims::Dims; kwargs...) = CURAND.rand_poisson(curand_rng(), T, dims; kwargs...)

# support all dimension specifications
rand(T::CURAND.UniformType, dim1::Integer, dims::Integer...) =
    Random.rand(curand_rng(), T, Dims((dim1, dims...)))
randn(T::CURAND.NormalType, dim1::Integer, dims::Integer...; kwargs...) =
    Random.randn(curand_rng(), T, Dims((dim1, dims...)); kwargs...)
rand_logn(T::CURAND.LognormalType, dim1::Integer, dims::Integer...; kwargs...) =
    CURAND.rand_logn(curand_rng(), T, Dims((dim1, dims...)); kwargs...)
rand_poisson(T::CURAND.PoissonType, dim1::Integer, dims::Integer...; kwargs...) =
    CURAND.rand_poisson(curand_rng(), T, Dims((dim1, dims...)); kwargs...)

# GPUArrays in-place
Random.rand!(A::AnyCuArray) = Random.rand!(gpuarrays_rng(), A)
Random.randn!(A::AnyCuArray; kwargs...) =
    error("CUDA.jl does not support generating normally-distributed random numbers of type $(eltype(A))")
# FIXME: GPUArrays.jl has a randn! nowadays, but it doesn't work with e.g. Cuint
rand_logn!(A::AnyCuArray; kwargs...) =
    error("CUDA.jl does not support generating lognormally-distributed random numbers of type $(eltype(A))")
rand_poisson!(A::AnyCuArray; kwargs...) =
    error("CUDA.jl does not support generating Poisson-distributed random numbers of type $(eltype(A))")

# GPUArrays out-of-place
rand(T::Type, dims::Dims) = Random.rand!(CuArray{T}(undef, dims...))
randn(T::Type, dims::Dims; kwargs...) = Random.randn!(CuArray{T}(undef, dims...); kwargs...)
rand_logn(T::Type, dims::Dims; kwargs...) = rand_logn!(CuArray{T}(undef, dims...); kwargs...)
rand_poisson(T::Type, dims::Dims; kwargs...) = rand_poisson!(CuArray{T}(undef, dims...); kwargs...)

# support all dimension specifications
rand(T::Type, dim1::Integer, dims::Integer...) =
    Random.rand!(CuArray{T}(undef, dim1, dims...))
randn(T::Type, dim1::Integer, dims::Integer...; kwargs...) =
    Random.randn!(CuArray{T}(undef, dim1, dims...); kwargs...)
rand_logn(T::Type, dim1::Integer, dims::Integer...; kwargs...) =
    rand_logn!(CuArray{T}(undef, dim1, dims...); kwargs...)
rand_poisson(T::Type, dim1::Integer, dims::Integer...; kwargs...) =
    rand_poisson!(CuArray{T}(undef, dim1, dims...); kwargs...)

# untyped out-of-place
rand(dim1::Integer, dims::Integer...) =
    Random.rand(curand_rng(), Dims((dim1, dims...)))
randn(dim1::Integer, dims::Integer...; kwargs...) =
    Random.randn(curand_rng(), Dims((dim1, dims...)); kwargs...)
rand_logn(dim1::Integer, dims::Integer...; kwargs...) =
    CURAND.rand_logn(curand_rng(), Dims((dim1, dims...)); kwargs...)
rand_poisson(dim1::Integer, dims::Integer...; kwargs...) =
    CURAND.rand_poisson(curand_rng(), Dims((dim1, dims...)); kwargs...)


@testset "SimulatedAnnealing" begin
    img = generate_surface(SynthSphere(50), 0.5, [0,0,1])
    p,q = synthetic_gradient(SynthSphere(40))
    p,q = SimulatedAnnealing(pin = p, qin = q, max_iter = 100)(img)
    @test 1 == 1
end


import HoTRG_lemon.SpinModule:
	getTrotterparameter,
	getTrotteriteration,
	getTrotterlayers,
	getBeta,
	setTrotterparameter!,
	setTrotterIteration!

abstract QuantumLattice{T} <: Lattice{T}

function buildQuantumLattice(latticeinfo::LatticeInfo, spinmodel::SpinModel)
	if isClassical(spinmodel)
		error("wrong lattice")
	end
	if isTwoDimension(latticeinfo) & isSquareLattice(latticeinfo)
		return Quantum2dSquareLattice(spinmodel)
	elseif isTwoDimension(latticeinfo) & isFractalLattice(latticeinfo)
		return Quantum2dFractalLattice(spinmodel)
	else
		error("wrong latticeinfo")
	end
end

#---
# functions inherit from QuantumSpinModel
get_functions = [
				:getTrotterparameter,
				:getTrotteriteration,
				:getTrotterlayers,
				:getBeta
				]
for f in get_functions
	eval(quote
	  function $(f){T}(lattice::QuantumLattice{T})
	  	$(f)(lattice.spinmodel)
		end
	end)
end

set_functions = [
				:setTrotterparameter!,
				:setTrotterIteration!
				]

for f in set_functions
	eval(quote

		  function $(f){T}(lattice::QuantumLattice{T}, parameter)
		  $(f)(lattice.spinmodel, parameter)
			initialize!(lattice)
		end

	end)
end


module Direct

using jInv.Mesh
using jInv.Utils
using jInv.InverseSolve
using ShapeReconstructionPaLS.Utils
using ShapeReconstructionPaLS.ParamLevelSet
using MAT
using SparseArrays
using Distributed
import jInv.ForwardShare.getData
import jInv.ForwardShare.getSensTMatVec
import jInv.ForwardShare.getSensMatVec
import jInv.ForwardShare.ForwardProbType

## The rotation planes are:
# first rotation (theta) is for planes x1 and x2
# second rotation (phi) is for planes x1 and x3
# Sampling is done by summing planes of x1 and x2 for all values of x3.

export DirectParam, getDirectParam
mutable struct DirectParam <: ForwardProbType
    Mesh      				:: RegularMesh
	ndipsAll				:: Int64
	workerSubIdxs
    #theta_phi_rad			:: Array{Float64,2}
	#b						:: Array{Float64,2}
	samplingBinning			:: Int64
	method					:: String
	Jacobian				:: SparseMatrixCSC{Float32,Int32}
	#S						:: SparseMatrixCSC{Float32,Int32}
end

function getDirectParamInternal(Mesh::RegularMesh,ndipsAll::Int64,workerSubIdxs,samplingBinning::Int64,method)
	return DirectParam(Mesh, ndipsAll, workerSubIdxs, samplingBinning, method,spzeros(Float32,Int32,0,0));
end

function getDirectParam(Mesh::RegularMesh,samplingBinning::Int64,numWorkers::Int64,method = MATBased)
	## This function does use the parallel mechanism of jInv (i.e., returns a RemoteChannel), even if numWorkers=1.
	if numWorkers > nworkers()
		numWorkers = nworkers();
	end
	#SourcesSubInd = Array{Array{Int64,1}}(numWorkers);
	SourcesSubInd = Array{Array{Int64}}(undef,numWorkers);
	ActualWorkers = workers();
	if numWorkers < nworkers()
		ActualWorkers = ActualWorkers[1:numWorkers];
	end
	#pFor   = Array{RemoteChannel}(numWorkers)
	pFor   = Array{RemoteChannel}(undef,numWorkers);
	i = 1; nextidx() = (idx=i; i+=1; idx)
	idx=i;
	ndips  = 1;
	# send out jobs
	@sync begin
		for p = ActualWorkers
			@async begin
				while true
					idx = nextidx();
					if idx > numWorkers
						break
					end
					I_p = getIndicesOfKthWorker(numWorkers,idx,ndips);
					# println("Sending ",collect(I_p)," to worker ", p);
					# find src and rec on mesh
					#SEE IF CORRECT!!!!!!!!!!!!!!!!
					tmp = initRemoteChannel(getDirectParamInternal,p, Mesh,ndips,I_p,samplingBinning,method);
					pFor[1]  = initRemoteChannel(getDirectParamInternal,p, Mesh,ndips,I_p,samplingBinning,method);
					wait(pFor[1]);

				end
			end
		end
	end
	return pFor # Array of Remote Refs
end

import jInv.Utils.clear!
function clear!(pFor::DirectParam)
	pFor.SampleMat = spzeros(0);
	clear!(pFor.Mesh);
	return pFor;
end

#include("generateDipMatrix.jl");
#include("prepareSyntheticDipData.jl");
include("getData.jl");
include("sensitivityFuncs.jl");
end


import Base: copy, +

#*********************************#
#       COST FUNCTION CLASS       #
#*********************************#

"""
Abstract type that represents a scalar-valued function that accepts a state and control
at a single knot point.
"""
abstract type CostFunction end

import RobotDynamics: diffmethod, ForwardAD, FiniteDifference
diffmethod(::CostFunction) = ForwardAD()

# Automatic differentiation methods for generic costs
@inline function gradient!(E, cost::CostFunction, z::AbstractKnotPoint, cache=ExpansionCache(cost))
    _gradient!(diffmethod(cost), E, cost, z, cache)
end

@inline function hessian!(E, cost::CostFunction,  z::AbstractKnotPoint, cache=ExpansionCache(cost))
    _hessian!(diffmethod(cost), E, cost, z, cache)
end

# ForwardDiff methods
function _gradient!(::ForwardAD, E, cost::CostFunction, z::AbstractKnotPoint, cache=nothing)
    if is_terminal(z)
        costfun_term(x) = stage_cost(cost, x)
        ForwardDiff.gradient!(E.x, costfun_term, state(z))
    else
        ix,iu = z._x, z._u
        costfun(z) = stage_cost(cost, z[ix], z[iu])
        ForwardDiff.gradient!(E.grad, costfun, z.z)
    end
    return false
end

function _hessian!(::ForwardAD, E, cost::CostFunction, z::AbstractKnotPoint, cache=nothing)
    if is_terminal(z)
        costfun_term(x) = stage_cost(cost, x)
        ForwardDiff.hessian!(E.xx, costfun_term, state(z))
    else
        ix,iu = z._x, z._u
        costfun(z) = stage_cost(cost, z[ix], z[iu])
        ForwardDiff.hessian!(E.hess, costfun, z.z)
    end
    return false
end

# FiniteDiff methods
function _gradient!(::FiniteDifference, E, cost::CostFunction, z::AbstractKnotPoint,
        cache=ExpansionCache(cost))
    if is_terminal(z)
        cache = cache[3]
        costfun_term(x) = stage_cost(cost, x)
        cache.c3 .= state(z)
        FiniteDiff.finite_difference_gradient!(E.x.data, costfun_term, cache.c3, cache)
    else
        cache = cache[1]
        ix,iu = z._x, z._u
        costfun(z) = stage_cost(cost, z[ix], z[iu])
        cache.c3 .= z.z
        FiniteDiff.finite_difference_gradient!(E.grad, costfun, cache.c3, cache)
    end
    return false
end

function _hessian!(::FiniteDifference, E, cost::CostFunction, z::AbstractKnotPoint,
        cache=ExpansionCache(cost))
    if is_terminal(z)
        cache = cache[4]
        costfun_term(x) = stage_cost(cost, x)
        x = state(z)
        cache.xmm .= x
        cache.xmp .= x
        cache.xpm .= x
        cache.xpp .= x
        FiniteDiff.finite_difference_hessian!(E.xx.data, costfun_term, cache.xmm, cache)
    else
        cache = cache[2]
        ix,iu = z._x, z._u
        costfun(z) = stage_cost(cost, z[ix], z[iu])
        cache.xmm .= z.z
        cache.xmp .= z.z
        cache.xpm .= z.z
        cache.xpp .= z.z
        FiniteDiff.finite_difference_hessian!(E.hess, costfun, cache.xmm, cache)
    end
    return false
end

"""
    ExpansionCache(cost)

Allocate the cache needed to evaluate the gradient and Hessian of the cost function.
Returns a 4-element vector with the following elements: `[grad, hess, grad_term, hess_term]`,
which are the caches need for the gradient and Hessian for the state and terminal cost
functions.
"""
@inline ExpansionCache(cost::CostFunction) =
    ExpansionCache(diffmethod(cost), state_dim(cost), control_dim(cost))
@inline ExpansionCache(::ForwardAD, n::Int, m::Int) = (nothing,nothing,nothing,nothing)
function ExpansionCache(::FiniteDifference, n::Int, m::Int)
    grad_cache0 = FiniteDiff.GradientCache(zeros(n), zeros(n), Val(:forward))
    hess_cache0 = FiniteDiff.HessianCache(zeros(n))
    grad_cache  = FiniteDiff.GradientCache(zeros(n+m), zeros(n+m), Val(:forward))
    hess_cache  = FiniteDiff.HessianCache(zeros(n+m))
    return (grad_cache, hess_cache, grad_cache0, hess_cache0)
end