module Ariel

export read_QuestaalXYZ

println("Ariel.jl")
println("Of sulphurous roaring the most mighty Neptune")
println("Seem to besiege and make his bold waves tremble,")

"""
 read_QuestaalXYZ(f::IOStream)

 Read in Questaal .XYZ file format to a trajectory object.

XYZ format looks like (128 H2O):
384 PLAT:   29.73674  0  0  0  29.73674  0  0  0  29.73674
t=103460fs T=285K V=-113.146Ry K.E.+V=-112.1104Ry alat=1; velocities times 1000
O            2.69595    -2.92790    13.17485    0.54262  -0.46064  -0.19684    0.00826   0.03806  -0.06875    1.01293   0.03968   0.25365  -0.09923   0.00000   0.00000   0.00000   0.00000   0.00000
"""
function read_QuestaalXYZ(f::IOStream)
    h=split(readline(f))
    natoms=parse(Int,h[1])
    println("natoms: $natoms")

    PLAT=reshape(map(x->parse(Float64,x), h[3:11] ), 3,3) # lattice parameters
    show(PLAT)

    readline(f) # don't do anything with the comment line, currently

#    traj=[]
#    while !eof(f)
#        append(traj,read_next_frame(f),natoms)
#    end
    A=reshape(readdlm(f, skipstart=2), natoms, 19, 2)
    show(A)
end

"Gobble natoms worth of frame"
function read_next_frame(f::IOStream,natoms)
    println("Oh lordy.")
end

end # module


function update_and_get_max_abs_diff!(  pssrb::PotentialSimulationSetupRB{T, S, DAT, N1, N2},
                                        depletion_handling::Val{depletion_handling_enabled},
                                        only2d::Val{only_2d} = Val{false}(),
                                        is_weighting_potential::Val{_is_weighting_potential} = Val{false}(),
                                        use_nthreads::Int = Base.Threads.nthreads()
                                        )::T where {T, S, DAT, N1, N2, depletion_handling_enabled, only_2d, _is_weighting_potential}
    tmp_potential::Array{T, N2} = copy(pssrb.potential)
    if depletion_handling_enabled
        update!(pssrb, use_nthreads = use_nthreads, depletion_handling = depletion_handling, only2d = only2d, is_weighting_potential = is_weighting_potential)
        slopes::Array{T, N2} = tmp_potential - pssrb.potential
        @inbounds for i in 1:19
            tmp_potential[:] = pssrb.potential[:]
            update!(pssrb, use_nthreads = use_nthreads, depletion_handling = depletion_handling, only2d = only2d, is_weighting_potential = is_weighting_potential)
            slopes += tmp_potential - pssrb.potential
        end
        @inbounds slopes /= 20
        return maximum(abs.(slopes))
    else
        for i in 1:10
            update!(pssrb, use_nthreads = use_nthreads, depletion_handling = depletion_handling, only2d = only2d, is_weighting_potential = is_weighting_potential)
        end
        max_diff::T = maximum(abs.(tmp_potential - pssrb.potential))
        return max_diff
    end
end

function _update_till_convergence!( pssrb::PotentialSimulationSetupRB{T, S, 3},
                                    convergence_limit::T,
                                    ::Type{Array};
                                    n_iterations_between_checks = 500,
                                    depletion_handling::Val{depletion_handling_enabled} = Val{false}(),
                                    only2d::Val{only_2d} = Val{false}(),
                                    is_weighting_potential::Val{_is_weighting_potential} = Val{false}(),
                                    use_nthreads::Int = Base.Threads.nthreads(),
                                    max_n_iterations::Int = -1,
                                    verbose::Bool = true
                                    )::T where {T, S, depletion_handling_enabled, only_2d, _is_weighting_potential}
    n_iterations::Int = 0
    cf::T = Inf
    cfs::Vector{T} = fill(cf, 10)
    cl::T = _is_weighting_potential ? convergence_limit : abs(convergence_limit * pssrb.bias_voltage) # to get relative change in respect to bias voltage
    # To disable automatically ProgressMeters in CI builds:
    is_logging(io) = isa(io, Base.TTY) == false || (get(ENV, "CI", nothing) == "true")
    if verbose prog = ProgressThresh(cl; dt = 0.1, desc = "Convergence: ", output = stderr, enabled = !is_logging(stderr)) end
    while cf > cl
        for i in 1:n_iterations_between_checks
            update!(pssrb, use_nthreads = use_nthreads, depletion_handling = depletion_handling, only2d = only2d, is_weighting_potential = is_weighting_potential)
        end
        cf = update_and_get_max_abs_diff!(pssrb, depletion_handling, only2d, is_weighting_potential, use_nthreads)
        @inbounds cfs[1:end-1] = cfs[2:end]
        @inbounds cfs[end] = cf
        slope::T = abs(mean(diff(cfs)))
        if verbose ProgressMeter.update!(prog, cf) end
        n_iterations += n_iterations_between_checks
        if slope < cl
            # @info "Slope is basically 0 -> Converged: $slope"
            cf = slope
        end
        if max_n_iterations > 0 && n_iterations > max_n_iterations
            # @show n_iterations_between_checks
            @info "Maximum number of iterations reached. (`n_iterations = $(n_iterations)`)"
            break
        end
    end
    if depletion_handling_enabled
        tmp_point_types = pssrb.point_types .& undepleted_bit
        @showprogress "Checking undepleted regions " for i in 1:10
            update!(pssrb, use_nthreads = use_nthreads, depletion_handling = depletion_handling, only2d = only2d, is_weighting_potential = is_weighting_potential)
            @inbounds for i in eachindex(pssrb.point_types)
                if (pssrb.point_types[i] & undepleted_bit == 0) && (tmp_point_types[i] > 0)
                    pssrb.point_types[i] += undepleted_bit
                elseif (pssrb.point_types[i] & undepleted_bit > 0) && (tmp_point_types[i] == 0)
                    tmp_point_types[i] += undepleted_bit
                end
            end
        end
    end

    if verbose ProgressMeter.finish!(prog) end
    return cf
end

# """
#     refine_scalar_potential(p::ScalarPotential{T}, max_diffs::NTuple{3, T}, minimum_distances::NTuple{3, T};
#         only2d::Val{only_2d} = Val(size(p.data, 2)==1)) where {T, only_2d}
#
# Refine any scalar potential `p`.
#
# 1. Extent the grid to be a closed grid in all dimensions.
# 2. Refine the axis of the grid based on `max_diffs` and `minimum_applied_potential`:
#    Insert N new ticks between to existing ticks such that the potential difference between each tick becomes
#    smaller than `max_diff[i]` (i -> dimension) but that the distances between the ticks stays larger than `minimum_distances[i]`.
# 3. Create the new data array for the refined grid and fill it by interpolation of the the initial (coarse) grid.
# """
function refine_scalar_potential(p::ScalarPotential{T}, max_diffs::NTuple{3, T}, minimum_distances::NTuple{3, T};
        only2d::Val{only_2d} = Val(size(p.data, 2)==1)) where {T, only_2d}
    closed_potential = _get_closed_potential(p)
    new_grid = _create_refined_grid(closed_potential, T.(max_diffs), T.(minimum_distances))
    new_data = Array{T, 3}(undef, size(new_grid))
    if only_2d
        int = interpolate_closed_potential(closed_potential, only2d)
        for i3 in axes(new_data, 3)
            x3 = new_grid.axes[3].ticks[i3]
            for i1 in axes(new_data, 1)
                x1 = new_grid.axes[1].ticks[i1]
                new_data[i1, 1, i3] = int(x1, x3)
            end
        end
    else
        int = interpolate_closed_potential(closed_potential, only2d)
        for i3 in axes(new_data, 3)
            x3 = new_grid.axes[3].ticks[i3]
            for i2 in axes(new_data, 2)
                x2 = new_grid.axes[2].ticks[i2]
                for i1 in axes(new_data, 1)
                    x1 = new_grid.axes[1].ticks[i1]
                    new_data[i1, i2, i3] = int(x1, x2, x3)
                end
            end
        end
    end
    return _convert_to_original_potential(p, new_data, new_grid)
end

function interpolate_closed_potential(p::ScalarPotential, ::Val{true}) where {T}
    interpolate((p.grid.axes[1], p.grid.axes[3]), p.data[:,1,:], Gridded(Linear()))
end
function interpolate_closed_potential(p::ScalarPotential, ::Val{false}) where {T}
    interpolate(p.grid.axes, p.data, Gridded(Linear()))
end

_get_closed_ticks(ticks::Vector{T}, int::ClosedInterval{T}) where {T} = ticks
_get_closed_ticks(ticks::Vector{T}, int::Interval{:closed, :open, T}) where {T} = vcat(ticks, int.right)
_get_closed_ticks(ticks::Vector{T}, int::Interval{:open, :closed, T}) where {T} = vcat(int.left, ticks)

_get_closed_values(values::Array{T,3}, dim::Int, int::ClosedInterval{T}) where {T} = values

function _get_closed_values(values::Array{T,3}, dim::Int, int::Interval{:closed, :open, T})::Array{T,3} where {T}
    if dim == 1
        cat(values, values[1:1,:,:], dims=dim)
    elseif dim == 2
        cat(values, values[:,1:1,:], dims=dim)
    else # dim == 3
        cat(values, values[:,:,1:1], dims=dim)
    end
end
function _get_closed_values(values::Array{T,3}, dim::Int, int::Interval{:open, :closed, T})::Array{T,3} where {T}
    if dim == 1
        cat(values[end:end,:,:], values, dims=dim)
    elseif dim == 2
        cat(values[:,end:end,:], values, dims=dim)
    else # dim == 3
        cat(values[:,:,end:end], values, dims=dim)
    end
end

function _get_closed_axis(ax::DiscreteAxis{T, BL, BR}) where {T, BL, BR}
    ticks = _get_closed_ticks(ax.ticks, ax.interval)
    return DiscreteAxis{T, BL, BR, ClosedInterval{T}}(ClosedInterval(ax.interval.left, ax.interval.right), ticks)
end

# """
#     _get_closed_potential(p::ScalarPotential{T,3,CS}) where {T, CS}
#
# Returns an closed Grid & Potential:
# E.g. if one of the axis is {:closed,:open} it will turn this into {:closed,:closed}
# and also extend the `data` field of the potential in the respective dimension and fill
# it with the respective values.
# """
function _get_closed_potential(p::ScalarPotential{T,3,CS}) where {T, CS}
    grid = p.grid
    closed_values = _get_closed_values(p.data, 1, grid.axes[1].interval)
    closed_values = _get_closed_values(closed_values, 2, grid.axes[2].interval)
    closed_values = _get_closed_values(closed_values, 3, grid.axes[3].interval)
    closed_axes = broadcast(i -> _get_closed_axis(grid.axes[i]), (1, 2, 3))
    AT = typeof(closed_axes)
    closed_grid = Grid{T, 3, CS, AT}(closed_axes)
    ScalarPotential(p, closed_values, closed_grid)
end

# """
#     _convert_to_original_potential(::Type{P}, data, grid) where {P, T, CS}
#
# Basically the counterpart to `_get_closed_potential`.
# """
function _convert_to_original_potential(p::ScalarPotential{T,3,CS,ATO}, data::Array{T, 3}, grid::Grid{T, 3, CS, AT}) where {T, CS, AT, ATO}
    ATs = get_axes_type(ATO)
    axs = broadcast(i -> _convert_closed_axis(ATs[i], grid.axes[i]), (1, 2, 3))
    new_data = _reduce_closed_values(ATs, data)
    new_grid = Grid{T,3,CS,ATO}(axs)
    ScalarPotential(p, new_data, new_grid)
end

_convert_closed_axis(::Type{DiscreteAxis{T, BL, BR, I}}, ax::DiscreteAxis{T, BL, BR, I}) where {T, BL, BR, I} = ax

function _convert_closed_axis(::Type{DiscreteAxis{T, BL, BR, Interval{:closed, :open, T}}}, ax::DiscreteAxis{T, BL, BR, ClosedInterval{T}}) where {T, BL, BR}
    DiscreteAxis{T, BL, BR, Interval{:closed, :open, T}}(Interval{:closed, :open, T}(ax.interval.left, ax.interval.right), ax.ticks[1:end-1])
end
function _convert_closed_axis(::Type{DiscreteAxis{T, BL, BR, Interval{:open, :closed, T}}}, ax::DiscreteAxis{T, BL, BR, ClosedInterval{T}}) where {T, BL, BR}
    DiscreteAxis{T, BL, BR, Interval{:open, :closed, T}}(Interval{:open, :closed, T}(ax.interval.left, ax.interval.right), ax.ticks[2:end])
end

function _reduce_closed_values(axTypes::Tuple, a::Array{T, 3}) where {T}
   inds = broadcast(i -> _get_ind_range_of_closed_axis(size(a, i), axTypes[i]), (1, 2, 3))
   a[inds...]
end

_get_ind_range_of_closed_axis(l::Int, ::Type{DiscreteAxis{T, BL, BR, ClosedInterval{T}}}) where {T, BL, BR} = 1:l
_get_ind_range_of_closed_axis(l::Int, ::Type{DiscreteAxis{T, BL, BR, Interval{:closed, :open, T}}}) where {T, BL, BR} = 1:(l-1)
_get_ind_range_of_closed_axis(l::Int, ::Type{DiscreteAxis{T, BL, BR, Interval{:open, :closed, T}}}) where {T, BL, BR} = 2:l

function _create_refined_grid(p::ScalarPotential{T,3}, max_diffs::NTuple{3, T}, minimum_distances::NTuple{3, T}) where {T}
    max_diffs = broadcast(md -> iszero(md) ? Inf : md, max_diffs)
    n_1 = floor.(Int, [maximum(abs.(p.data[i+1,:,:] .- p.data[i,:,:])) for i in 1:size(p.data, 1)-1] ./ max_diffs[1])
    n_2 = floor.(Int, [maximum(abs.(p.data[:,i+1,:] .- p.data[:,i,:])) for i in 1:size(p.data, 2)-1] ./ max_diffs[2])
    n_3 = floor.(Int, [maximum(abs.(p.data[:,:,i+1] .- p.data[:,:,i])) for i in 1:size(p.data, 3)-1] ./ max_diffs[3])
    ns = (n_1, n_2, n_3)
    widths = diff.((p.grid.axes[1].ticks, p.grid.axes[2].ticks, p.grid.axes[3].ticks))
    sub_widths = broadcast(ia -> [widths[ia][i] / (ns[ia][i]+1) for i in eachindex(ns[ia])], (1,2,3))
    for ia in 1:3
        for i in eachindex(ns[ia])
            while sub_widths[ia][i] < minimum_distances[ia] && ns[ia][i] > 0
                ns[ia][i] -= 1
                sub_widths[ia][i] = widths[ia][i] / (ns[ia][i]+1)
            end
        end
    end
    for i in 1:3 # always add an even number of ticks
        if isodd(sum(ns[i]))
            i_max_width = findmax(sub_widths[i])[2]
            ns[i][i_max_width] += 1
            sub_widths[i][i_max_width] = widths[i][i_max_width] / (ns[i][i_max_width]+1)
        end
    end
    new_axes = broadcast(i -> _refine_axis(p.grid.axes[i], ns[i], sub_widths[i]), (1, 2, 3))
    return typeof(p.grid)(new_axes)
end

function _refine_axis(ax::DiscreteAxis{T, <:Any, <:Any, ClosedInterval{T}}, ns::Vector{Int}, sub_widths::Vector{T}) where {T, I}
    @assert length(ns) == length(ax.ticks)-1 # for ClosedInterval axis
    ticks = Vector{T}(undef, length(ax.ticks) + sum(ns))
    i = 1
    for j in eachindex(ns)
        ticks[i] = ax.ticks[j]
        for k in 1:ns[j]
            i += 1
            ticks[i] = ax.ticks[j] + k * sub_widths[j]
        end
        i += 1
    end
    ticks[end] = ax.ticks[end]
    typeof(ax)(ax.interval, ticks)
end

_extend_refinement_limits(rl::Real) = (rl, rl, rl )
_extend_refinement_limits(rl::Tuple{<:Real,<:Real,<:Real}) = rl


module Turing

##############
# Dependency #
########################################################################
# NOTE: when using anything from external packages,                    #
#       let's keep the practice of explictly writing Package.something #
#       to indicate that's not implemented inside Turing.jl            #
########################################################################

using Requires
using Distributions
using ForwardDiff

using LinearAlgebra
using ProgressMeter
using Markdown

#  @init @require Stan="682df890-35be-576f-97d0-3d8c8b33a550" begin
  using Stan
  import Stan: Adapt, Hmc
#  end
import Base: ~, convert, promote_rule, rand, getindex, setindex!
import Distributions: sample
import ForwardDiff: gradient
using Flux: Tracker
import MCMCChain: AbstractChains, Chains

##############################
# Global variables/constants #
##############################

const ADBACKEND = Ref(:reverse_diff)
setadbackend(backend_sym) = begin
  @assert backend_sym == :forward_diff || backend_sym == :reverse_diff
  ADBACKEND[] = backend_sym
end

const ADSAFE = Ref(false)
setadsafe(switch::Bool) = begin
  @info("[Turing]: global ADSAFE is set as $switch")
  ADSAFE[] = switch
end

const CHUNKSIZE = Ref(0) # default chunksize used by AD

setchunksize(chunk_size::Int) = begin
  if ~(CHUNKSIZE[] == chunk_size)
    @info("[Turing]: AD chunk size is set as $chunk_size")
    CHUNKSIZE[] = chunk_size
    global SEEDS = ForwardDiff.construct_seeds(ForwardDiff.Partials{chunk_size,Float64}) # pre-alloced dual parts
  end
end

setchunksize(40)

const PROGRESS = Ref(true)
turnprogress(switch::Bool) = begin
  @info("[Turing]: global PROGRESS is set as $switch")
  PROGRESS[] = switch
end

# Constants for caching
const CACHERESET  = 0b00
const CACHEIDCS   = 0b10
const CACHERANGES = 0b01

#######################
# Sampler abstraction #
#######################

abstract type InferenceAlgorithm end
abstract type Hamiltonian <: InferenceAlgorithm end

"""
    Sampler{T}

Generic interface for implementing inference algorithms.
An implementation of an algorithm should include the following:

1. A type specifying the algorithm and its parameters, derived from InferenceAlgorithm
2. A method of `sample` function that produces results of inference, which is where actual inference happens.

Turing translates models to chunks that call the modelling functions at specified points. The dispatch is based on the value of a `sampler` variable. To include a new inference algorithm implements the requirements mentioned above in a separate file,
then include that file at the end of this one.
"""
mutable struct Sampler{T<:InferenceAlgorithm}
  alg   ::  T
  info  ::  Dict{Symbol, Any}         # sampler infomation
end

include("../deps/deps.jl"); check_deps();
include("helper.jl")
include("transform.jl")
include("core/varinfo.jl")  # core internal variable container
include("trace/trace.jl")   # to run probabilistic programs as tasks

using Turing.Traces
using Turing.VarReplay

###########
# Exports #
###########

# Turing essentials - modelling macros and inference algorithms
export @model, @~, @VarName                   # modelling
export MH, Gibbs                              # classic sampling
export HMC, SGLD, SGHMC, HMCDA, NUTS          # Hamiltonian-like sampling
export IS, SMC, CSMC, PG, PIMH, PMMH, IPMCMC  # particle-based sampling
export sample, setchunksize, resume           # inference
export auto_tune_chunk_size!, setadbackend, setadsafe # helper
export turnprogress  # debugging

# Turing-safe data structures and associated functions
export TArray, tzeros, localcopy, IArray

export @sym_str

export UnivariateGMM2, Flat, FlatPos

##################
# Inference code #
##################

include("core/util.jl")         # utility functions
include("core/compiler.jl")     # compiler
include("core/container.jl")    # particle container
include("core/io.jl")           # I/O
include("samplers/sampler.jl")  # samplers
include("core/ad.jl")           # Automatic Differentiation

end


## Loading data into various data structures ##

struct Require
    versions::VersionSet
    systems::Vector{Symbol}
end

struct Version
    sha1::Union{Nothing, String}
    requires::Dict{String,Require}
end

struct Package
    uuid::UUID
    url::String
    versions::Dict{VersionNumber,Version}
end

Require(versions::VersionSet) = Require(versions, Symbol[])
Require(version::VersionInterval) = Require(VersionSet([version]), Symbol[])
Version(sha1::String) = Version(sha1, Dict{String,Require}())

function load_requires(path::String)
    requires = Dict{String,Require}()
    requires["julia"] = Require(VersionInterval())
    isfile(path) || return requires
    for r in filter!(r->r isa Requirement, Reqs.read(path))
        new = haskey(requires, r.package)
        versions, systems = VersionSet(r.versions.intervals), r.system
        if haskey(requires, r.package)
            versions = versions ∩ requires[r.package].versions
            systems  = systems  ∪ requires[r.package].systems
        end
        requires[r.package] = Require(versions, Symbol.(systems))
    end
    return requires
end

function load_versions(dir::String)
    versions = Dict{VersionNumber,Version}()
    isdir(dir) || return versions
    for ver in readdir(dir)
        path = joinpath(dir, ver)
        sha1 = joinpath(path, "sha1")
        isfile(sha1) || continue
        requires = load_requires(joinpath(path, "requires"))
        versions[VersionNumber(ver)] = Version(readchomp(sha1), requires)
    end
    return versions
end

function load_packages_METADATA(
    dir::AbstractString,
    selected_packages::AbstractVector{<:AbstractString}
)
    pkgs = Dict{String,Package}()
    for pkg in selected_packages
        path = joinpath(dir, pkg)
        url = joinpath(path, "url")
        versions = joinpath(path, "versions")
        isfile(url) || continue
        pkgs[pkg] = Package(
            uuid5(uuid_package, pkg),
            readchomp(url),
            load_versions(versions),
        )
    end
    pkgs
end

@eval julia_versions() = $([VersionNumber(0,m) for m=1:7])
julia_versions(f::Function) = filter(f, julia_versions())
julia_versions(vi::VersionInterval) = julia_versions(v->v in vi)

macro clean(ex) :(x = $(esc(ex)); $(esc(:clean)) &= x; x) end

function check_packages(
    pkgs::AbstractDict{<:AbstractString,Package},
    other_registries::AbstractDict{<:AbstractString, Package}
)
    for (pkg, p) in pkgs, (ver, v) in p.versions
        if ver != thispatch(ver)
            error("$pkg version $ver is illegeal: no prereleases")
        end
        if ver == v"0.0.0"
            error("$pkg version $ver: Version must be greater than 0.0.0")
        end
        for (req, r) in v.requires
            # req == "julia" && continue
            if haskey(other_registries, req)
                req_registry = other_registries
            elseif haskey(pkgs, req)
                req_registry = pkgs
            else
                error("$pkg ver $ver: requirement $req not found")
            end
            available_req_versions = keys(req_registry[req].versions)
            if ! any(w->w in r.versions, available_req_versions)
                error(
"""
$pkg ver $ver: compat range not satisfied
$pkg at $ver requires $req version in $(r.versions)
But known versions for $req are $available_req_versions
"""
                )
            end
        end
    end
    nothing
end


module CommonFunctions

    export nbasis, nnodes, nodes, order, degree

    nbasis() = nothing
    nnodes() = nothing
    nodes()  = nothing
    order()  = nothing
    degree() = nothing

    export evaluate, evaluate!

    evaluate() = nothing
    evaluate!() = nothing

    export derivative, integral

    derivative() = nothing
    integral() = nothing

    export write_to_hdf5

    write_to_hdf5() = nothing

    export reset!, cut_periodic_solution!

    reset!() = nothing
    cut_periodic_solution!() = nothing

    export eachsample, eachtimestep

    eachsample() = nothing
    eachtimestep() = nothing

end


##Pakiety

##Zmienne
#Menu z cenami

#Produkty i ich data waznosci

#Początkowy stan produktów

##Rozkłady
#Ilość samochodów na minutę

##Główna funkcja kroku
#jeden krok = jedna minuta

function OneStep()

    #wylosować z rozkładu samochdów na minutę ile samochdów przyjedzie w tym kroku

    #wylosować ilość osób w kazdym samochodzie, 1 samochod = jedna zmienna?

    # wylosować zamówienia dla kazdego samochodu

    #


#######################################################################
# ParametricHamiltonian
#######################################################################
struct ParametricHamiltonian{P,N,M<:NTuple{N,ElementModifier},D<:NTuple{N,Any},H<:Hamiltonian}
    baseh::H
    h::H
    modifiers::M                   # N modifiers
    ptrdata::D                     # P is an NTuple{N,Vector{Vector{ptrdata}}}, one per harmonic
    allptrs::Vector{Vector{Int}}   # ptrdata may be a nzval ptr, a (ptr,r) or a (ptr, r, dr)
    parameters::NTuple{P,NameType} # allptrs are modified ptrs in each harmonic
end

function Base.show(io::IO, ::MIME"text/plain", pham::ParametricHamiltonian{N}) where {N}
    i = get(io, :indent, "")
    print(io, i, "Parametric")
    show(io, pham.h)
    print(io, i, "\n", "$i  Parameters       : $(parameters(pham))")
end

"""
    parametric(h::Hamiltonian, modifiers::ElementModifier...)

Builds a `ParametricHamiltonian` that can be used to efficiently apply `modifiers` to `h`.
`modifiers` can be any number of `@onsite!(args -> body; kw...)` and `@hopping!(args -> body;
kw...)` transformations, each with a set of parameters `ps` given as keyword arguments of
functions `f = (...; ps...) -> body`. The resulting `ph::ParamtricHamiltonian` can be used
to produced the modified Hamiltonian simply by calling it with those same parameters as
keyword arguments.

Note 1: for sparse `h`, `parametric` only modifies existing onsites and hoppings in `h`,
so be sure to add zero onsites and/or hoppings to `h` if they are originally not present but
you need to apply modifiers to them.

Note 2: `optimize!(h)` is called prior to building the parametric Hamiltonian. This can lead
to extra zero onsites and hoppings being stored in sparse `h`s.

    h |> parametric(modifiers::ElementModifier...)

Function form of `parametric`, equivalent to `parametric(h, modifiers...)`.

# Examples

```jldoctest
julia> ph = LatticePresets.honeycomb() |> hamiltonian(onsite(0) + hopping(1, range = 1/√3)) |>
       unitcell(10) |> parametric(@onsite!((o; μ) -> o - μ))
ParametricHamiltonian{<:Lattice} : Hamiltonian on a 2D Lattice in 2D space
  Bloch harmonics  : 5 (SparseMatrixCSC, sparse)
  Harmonic size    : 200 × 200
  Orbitals         : ((:a,), (:a,))
  Element type     : scalar (Complex{Float64})
  Onsites          : 0
  Hoppings         : 600
  Coordination     : 3.0
  Parameters       : (:μ,)

julia> ph(μ = 2)
Hamiltonian{<:Lattice} : Hamiltonian on a 2D Lattice in 2D space
  Bloch harmonics  : 5 (SparseMatrixCSC, sparse)
  Harmonic size    : 200 × 200
  Orbitals         : ((:a,), (:a,))
  Element type     : scalar (Complex{Float64})
  Onsites          : 200
  Hoppings         : 600
  Coordination     : 3.0
```
# See also
    `@onsite!`, `@hopping!`
"""
function parametric(h::Hamiltonian, ts::ElementModifier...)
    ts´ = resolve(ts, h.lattice)
    optimize!(h)  # to avoid ptrs getting out of sync if optimize! later
    allptrs = Vector{Int}[Int[] for _ in h.harmonics]
    ptrdata = parametric_ptrdata_tuple!(allptrs, h, ts´)
    foreach(sort!, allptrs)
    foreach(unique!, allptrs)
    params = parameters(ts...)
    return ParametricHamiltonian(h, copy(h), ts´, ptrdata, allptrs, params)
end

parametric(ts::ElementModifier...) = h -> parametric(h, ts...)

parametric_ptrdata_tuple!(allptrs, h, ts´::Tuple) = _parametric_ptrdata!(allptrs, h, ts´...)
_parametric_ptrdata!(allptrs, h, t, ts...) =
    (parametric_ptrdata!(allptrs, h, t), _parametric_ptrdata!(allptrs, h, ts...)...)
_parametric_ptrdata!(allptrs, h) = ()

function parametric_ptrdata!(allptrs, h::Hamiltonian{LA,L,M,<:AbstractSparseMatrix}, t::ElementModifier) where {LA,L,M}
    harmonic_ptrdata = empty_ptrdata(h, t)
    lat = h.lattice
    selector = t.selector
    for (har, ptrdata, allptrs_har) in zip(h.harmonics, harmonic_ptrdata, allptrs)
        matrix = har.h
        dn = har.dn
        rows = rowvals(matrix)
        for col in 1:size(matrix, 2), ptr in nzrange(matrix, col)
            row = rows[ptr]
            selected = selector(lat, (row, col), (dn, zero(dn)))
            if selected
                push!(ptrdata, ptrdatum(t, lat, ptr, (row, col), dn))
                push!(allptrs_har, ptr)
            end
        end
    end
    return harmonic_ptrdata
end

# Uniform case, one vector of nzval ptr per harmonic
empty_ptrdata(h, t::UniformModifier)  = Vector{Int}[Int[] for _ in h.harmonics]

# Non-uniform case, one vector of (ptr, r, dr) per harmonic
function empty_ptrdata(h, t::OnsiteModifier)
    S = positiontype(h.lattice)
    return Vector{Tuple{Int,S}}[Tuple{Int,S}[] for _ in h.harmonics]
end

function empty_ptrdata(h, t::HoppingModifier)
    S = positiontype(h.lattice)
    return Vector{Tuple{Int,S,S}}[Tuple{Int,S,S}[] for _ in h.harmonics]
end

# Uniform case
ptrdatum(t::UniformModifier, lat, ptr, (row, col), dn) = ptr

# Non-uniform case
ptrdatum(t::OnsiteModifier, lat, ptr, (row, col), dn) = (ptr, sites(lat)[col])

function ptrdatum(t::HoppingModifier, lat, ptr, (row, col), dn)
    r, dr = _rdr(sites(lat)[col], sites(lat)[row] + bravais(lat) * dn)
    return (ptr, r, dr)
end

function (ph::ParametricHamiltonian)(; kw...)
    checkconsistency(ph, false) # only weak check for performance
    reset_harmonic!.(ph.h.harmonics, ph.baseh.harmonics, ph.allptrs)
    applymodifier_ptrdata!.(Ref(ph.h), ph.modifiers, ph.ptrdata, Ref(values(kw)))
    return ph.h
end

function reset_harmonic!(har, basehar, ptrs)
    nz = nonzeros(har.h)
    onz = nonzeros(basehar.h)
    @simd for ptr in ptrs
        nz[ptr] = onz[ptr]
    end
    return har
end

function applymodifier_ptrdata!(h, modifier, ptrdata, kw)
    for (har, hardata) in zip(h.harmonics, ptrdata)
        nz = nonzeros(har.h)
        @simd for data in hardata  # @simd is valid because ptrs are not repeated
            ptr = first(data)
            args = modifier_args(nz, data)
            val = modifier(args...; kw...)
            nz[ptr] = val
        end
    end
    return h
end

# A negative ptr corresponds to a forced-hermitian element
modifier_args(nz, ptr::Int) = (nz[ptr],)
modifier_args(nz, (ptr, r)::Tuple{Int,SVector}) = (nz[ptr], r)
modifier_args(nz, (ptr, r, dr)::Tuple{Int,SVector,SVector}) = (nz[ptr], r, dr)

function checkconsistency(ph::ParametricHamiltonian, fullcheck = true)
    isconsistent = true
    length(ph.baseh.harmonics) == length(ph.h.harmonics) || (isconsitent = false)
    if fullcheck && isconsistent
        for (basehar, har) in zip(ph.baseh.harmonics, ph.h.harmonics)
            length(nonzeros(har.h)) == length(nonzeros(basehar.h)) || (isconsistent = false; break)
            rowvals(har.h) == rowvals(basehar.h) || (isconsistent = false; break)
            getcolptr(har.h) == getcolptr(basehar.h) || (isconsistent = false; break)
        end
    end
    isconsistent ||
        throw(error("ParametricHamiltonian is not internally consistent, it may have been modified after creation"))
    return nothing
end

# ParametricHamiltonian's are already optimized upon creation
optimize!(ph::ParametricHamiltonian) = ph

"""
    parameters(ph::ParametricHamiltonian)

Return the names of the parameter that `ph` depends on
"""
parameters(ph::ParametricHamiltonian) = ph.parameters

matrixtype(ph::ParametricHamiltonian) = matrixtype(parent(ph))

blockeltype(ph::ParametricHamiltonian) = blockeltype(parent(ph))

bravais(ph::ParametricHamiltonian) = bravais(ph.hamiltonian.lattice)

Base.parent(ph::ParametricHamiltonian) = ph.h

Base.copy(ph::ParametricHamiltonian) =
    ParametricHamiltonian(copy(ph.baseh), copy(ph.h), ph.modifiers, copy(h.ptrdata), copy(h.allptrs))

Base.size(ph::ParametricHamiltonian, n...) = size(ph.h, n...)

Base.eltype(ph::ParametricHamiltonian) = eltype(ph.h)

#######################################################################
# bloch! for parametric
#######################################################################

bloch(ph::ParametricHamiltonian, args...) = bloch!(similarmatrix(ph), ph, args...)

bloch!(matrix, ph::ParametricHamiltonian, pϕs = (), axis = 0) =
    bloch!(matrix, h_phases(ph, toSVector(pϕs))..., axis)

@inline function h_phases(ph::ParametricHamiltonian, pϕs)
    pnames = parameters(ph)
    ps, ϕs = extract_parameters_phases(pnames, pϕs)
    h = ph(; ps...)
    return (h, ϕs)
end

extract_parameters_phases(pnames::NTuple{N,NameType}, ϕs::SVector{M}) where {N,M} =
    (NamedTuple{pnames}(ntuple(i->ϕs[i], Val(N))), ntuple(i->ϕs[i+N], Val(M-N)))


export Experiment

using Dates
using ReinforcementLearningBase
using ReinforcementLearningCore
using .ReinforcementLearningEnvironments
using Flux
using BSON
using TensorBoardLogger
using Logging

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:BasicDQN},
    ::Val{:CartPole},
    ::Nothing;
    save_dir = nothing,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyymmddHHMMSS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_BasicDQN_CartPole_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    env = CartPoleEnv(; T = Float32, seed = 11)
    ns, na = length(rand(get_observation_space(env))), length(get_action_space(env))
    agent = Agent(
        policy = QBasedPolicy(
            learner = BasicDQNLearner(
                approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 128, relu; initW = seed_glorot_uniform(seed = 17)),
                        Dense(128, 128, relu; initW = seed_glorot_uniform(seed = 23)),
                        Dense(128, na; initW = seed_glorot_uniform(seed = 39)),
                    ) |> cpu,
                    optimizer = ADAM(),
                ),
                batch_size = 32,
                min_replay_history = 100,
                loss_func = huber_loss,
                seed = 22,
            ),
            explorer = EpsilonGreedyExplorer(
                kind = :exp,
                ϵ_stable = 0.01,
                decay_steps = 500,
                seed = 33,
            ),
        ),
        trajectory = CircularCompactSARTSATrajectory(
            capacity = 1000,
            state_type = Float32,
            state_size = (ns,),
        ),
    )

    stop_condition = StopAfterStep(10000)

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        DoEveryNStep() do t, agent, env, obs
            with_logger(lg) do
                @info "training" loss = agent.policy.learner.loss
            end
        end,
        DoEveryNEpisode() do t, agent, env, obs
            with_logger(lg) do
                @info "training" reward = total_reward_per_episode.rewards[end]
            end
        end,
        DoEveryNStep(10000) do t, agent, env, obs
            RLCore.save(save_dir, agent)
            BSON.@save joinpath(save_dir, "stats.bson") total_reward_per_episode time_per_step
        end,
    )

    description = """
    This experiment uses three dense layers to approximate the Q value.
    The testing environment is CartPoleEnv.

    Agent and statistic info will be saved to: `$save_dir`
    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`
    To load the agent and statistic info:
    ```
    agent = RLCore.load("$save_dir", Agent)
    BSON.@load joinpath("$save_dir", "stats.bson") total_reward_per_episode time_per_step
    ```
    """

    Experiment(agent, env, stop_condition, hook, description)
end

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:DQN},
    ::Val{:CartPole},
    ::Nothing;
    save_dir = nothing,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyymmddHHMMSS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_DQN_CartPole_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    env = CartPoleEnv(; T = Float32, seed = 11)
    ns, na = length(rand(get_observation_space(env))), length(get_action_space(env))

    agent = Agent(
        policy = QBasedPolicy(
            learner = DQNLearner(
                approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 128, relu; initW = seed_glorot_uniform(seed = 17)),
                        Dense(128, 128, relu; initW = seed_glorot_uniform(seed = 23)),
                        Dense(128, na; initW = seed_glorot_uniform(seed = 39)),
                    ) |> cpu,
                    optimizer = ADAM(),
                ),
                target_approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 128, relu; initW = seed_glorot_uniform(seed = 17)),
                        Dense(128, 128, relu; initW = seed_glorot_uniform(seed = 23)),
                        Dense(128, na; initW = seed_glorot_uniform(seed = 39)),
                    ) |> cpu,
                    optimizer = ADAM(),
                ),
                loss_func = huber_loss,
                stack_size = nothing,
                batch_size = 32,
                update_horizon = 1,
                min_replay_history = 100,
                update_freq = 1,
                target_update_freq = 100,
                seed = 22,
            ),
            explorer = EpsilonGreedyExplorer(
                kind = :exp,
                ϵ_stable = 0.01,
                decay_steps = 500,
                seed = 33,
            ),
        ),
        trajectory = CircularCompactSARTSATrajectory(
            capacity = 1000,
            state_type = Float32,
            state_size = (ns,),
        ),
    )

    stop_condition = StopAfterStep(10_000)

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        DoEveryNStep() do t, agent, env, obs
            if agent.policy.learner.update_step % agent.policy.learner.update_freq == 0
                with_logger(lg) do
                    @info "training" loss = agent.policy.learner.loss
                end
            end
        end,
        DoEveryNStep(10000) do t, agent, env, obs
            RLCore.save(save_dir, agent)
            BSON.@save joinpath(save_dir, "stats.bson") total_reward_per_episode time_per_step
        end,
    )

    description = """
    This experiment uses the `DQNLearner` method with three dense layers to approximate the Q value.
    The testing environment is CartPoleEnv.

    Agent and statistic info will be saved to: `$save_dir`
    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`
    To load the agent and statistic info:

    ```
    agent = RLCore.load("$save_dir", Agent)
    BSON.@load joinpath("$save_dir", "stats.bson") total_reward_per_episode time_per_step
    ```
    """

    Experiment(agent, env, stop_condition, hook, description)
end

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:PrioritizedDQN},
    ::Val{:CartPole},
    ::Nothing;
    save_dir = nothing,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyymmddHHMMSS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_PrioritizedDQN_CartPole_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    env = CartPoleEnv(; T = Float32, seed = 11)
    ns, na = length(rand(get_observation_space(env))), length(get_action_space(env))

    agent = Agent(
        policy = QBasedPolicy(
            learner = PrioritizedDQNLearner(
                approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 128, relu; initW = seed_glorot_uniform(seed = 17)),
                        Dense(128, 128, relu; initW = seed_glorot_uniform(seed = 23)),
                        Dense(128, na; initW = seed_glorot_uniform(seed = 39)),
                    ) |> cpu,
                    optimizer = ADAM(),
                ),
                target_approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 128, relu; initW = seed_glorot_uniform(seed = 17)),
                        Dense(128, 128, relu; initW = seed_glorot_uniform(seed = 23)),
                        Dense(128, na; initW = seed_glorot_uniform(seed = 39)),
                    ) |> cpu,
                    optimizer = ADAM(),
                ),
                loss_func = huber_loss_unreduced,
                stack_size = nothing,
                batch_size = 32,
                update_horizon = 1,
                min_replay_history = 100,
                update_freq = 1,
                target_update_freq = 100,
                seed = 22,
            ),
            explorer = EpsilonGreedyExplorer(
                kind = :exp,
                ϵ_stable = 0.01,
                decay_steps = 500,
                seed = 33,
            ),
        ),
        trajectory = CircularCompactPSARTSATrajectory(
            capacity = 1000,
            state_type = Float32,
            state_size = (ns,),
        ),
    )

    stop_condition = StopAfterStep(10_000)

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        DoEveryNStep() do t, agent, env, obs
            if agent.policy.learner.update_step % agent.policy.learner.update_freq == 0
                with_logger(lg) do
                    @info "training" loss = agent.policy.learner.loss
                end
            end
        end,
        DoEveryNStep(10000) do t, agent, env, obs
            RLCore.save(save_dir, agent)
            BSON.@save joinpath(save_dir, "stats.bson") total_reward_per_episode time_per_step
        end,
    )

    description = """
    This experiment uses the `PrioritizedDQNLearner` method with three dense layers to approximate the Q value.
    The testing environment is CartPoleEnv.

    Agent and statistic info will be saved to: `$save_dir`
    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`
    To load the agent and statistic info:

    ```
    agent = RLCore.load("$save_dir", Agent)
    BSON.@load joinpath("$save_dir", "stats.bson") total_reward_per_episode time_per_step
    ```
    """

    Experiment(agent, env, stop_condition, hook, description)
end

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:Rainbow},
    ::Val{:CartPole},
    ::Nothing;
    save_dir = nothing,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyymmddHHMMSS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_Rainbow_CartPole_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    env = CartPoleEnv(; T = Float32, seed = 11)
    ns, na = length(rand(get_observation_space(env))), length(get_action_space(env))

    n_atoms = 51
    agent = Agent(
        policy = QBasedPolicy(
            learner = RainbowLearner(
                approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 128, relu; initW = seed_glorot_uniform(seed = 17)),
                        Dense(128, 128, relu; initW = seed_glorot_uniform(seed = 23)),
                        Dense(128, na * n_atoms; initW = seed_glorot_uniform(seed = 39)),
                    ) |> cpu,
                    optimizer = ADAM(0.0005),
                ),
                target_approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 128, relu; initW = seed_glorot_uniform(seed = 17)),
                        Dense(128, 128, relu; initW = seed_glorot_uniform(seed = 23)),
                        Dense(128, na * n_atoms; initW = seed_glorot_uniform(seed = 39)),
                    ) |> cpu,
                    optimizer = ADAM(0.0005),
                ),
                n_actions = na,
                n_atoms = n_atoms,
                Vₘₐₓ = 200.0f0,
                Vₘᵢₙ = 0.0f0,
                update_freq = 1,
                γ = 0.99f0,
                update_horizon = 1,
                batch_size = 32,
                stack_size = nothing,
                min_replay_history = 100,
                loss_func = logitcrossentropy_unreduced,
                target_update_freq = 100,
                seed = 22,
            ),
            explorer = EpsilonGreedyExplorer(
                kind = :exp,
                ϵ_stable = 0.01,
                decay_steps = 500,
                seed = 33,
            ),
        ),
        trajectory = CircularCompactPSARTSATrajectory(
            capacity = 1000,
            state_type = Float32,
            state_size = (ns,),
        ),
    )

    stop_condition = StopAfterStep(10_000)

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        DoEveryNStep() do t, agent, env, obs
            if agent.policy.learner.update_step % agent.policy.learner.update_freq == 0
                with_logger(lg) do
                    @info "training" loss = agent.policy.learner.loss
                end
            end
        end,
        DoEveryNStep(10000) do t, agent, env, obs
            RLCore.save(save_dir, agent)
            BSON.@save joinpath(save_dir, "stats.bson") total_reward_per_episode time_per_step
        end,
    )

    description = """
    This experiment uses the `RainbowLearner` method with three dense layers to approximate the distributed Q value.
    The testing environment is CartPoleEnv.

    Agent and statistic info will be saved to: `$save_dir`
    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`
    To load the agent and statistic info:

    ```
    agent = RLCore.load("$save_dir", Agent)
    BSON.@load joinpath("$save_dir", "stats.bson") total_reward_per_episode time_per_step
    ```
    """

    Experiment(agent, env, stop_condition, hook, description)
end

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:IQN},
    ::Val{:CartPole},
    ::Nothing;
    save_dir = nothing,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyymmddHHMMSS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_IQN_CartPole_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    env = CartPoleEnv(; T = Float32, seed = 111)
    ns, na = length(rand(get_observation_space(env))), length(get_action_space(env))

    init = seed_glorot_uniform(seed = 17)
    Nₑₘ = 16
    n_hidden = 64
    κ = 1.0f0

    nn_creator() =
        ImplicitQuantileNet(
            ψ = Dense(ns, n_hidden, relu; initW = init),
            ϕ = Dense(Nₑₘ, n_hidden, relu; initW = init),
            header = Dense(n_hidden, na; initW = init),
        ) |> cpu

    agent = Agent(
        policy = QBasedPolicy(
            learner = IQNLearner(
                approximator = NeuralNetworkApproximator(
                    model = nn_creator(),
                    optimizer = ADAM(0.001),
                ),
                target_approximator = NeuralNetworkApproximator(model = nn_creator()),
                κ = κ,
                N = 8,
                N′ = 8,
                Nₑₘ = Nₑₘ,
                K = 32,
                γ = 0.99f0,
                stack_size = nothing,
                batch_size = 32,
                update_horizon = 1,
                min_replay_history = 100,
                update_freq = 1,
                target_update_freq = 100,
                default_priority = 1.0f2,
                seed = 123,
                device_seed = 321,
            ),
            explorer = EpsilonGreedyExplorer(
                kind = :exp,
                ϵ_stable = 0.01,
                decay_steps = 500,
                seed = 33,
            ),
        ),
        trajectory = CircularCompactPSARTSATrajectory(
            capacity = 1000,
            state_type = Float32,
            state_size = (ns,),
        ),
    )

    stop_condition = StopAfterStep(10_000)

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        DoEveryNStep() do t, agent, env, obs
            if agent.policy.learner.update_step % agent.policy.learner.update_freq == 0
                with_logger(lg) do
                    @info "training" loss = agent.policy.learner.loss
                end
            end
        end,
        DoEveryNStep(10000) do t, agent, env, obs
            RLCore.save(save_dir, agent)
            BSON.@save joinpath(save_dir, "stats.bson") total_reward_per_episode time_per_step
        end,
    )

    description = """
    This experiment uses the `IQNLearner` method with a `ImplicitQuantileNet`.
    The testing environment is CartPoleEnv.

    Agent and statistic info will be saved to: `$save_dir`
    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`
    To load the agent and statistic info:

    ```
    agent = RLCore.load("$save_dir", Agent)
    BSON.@load joinpath("$save_dir", "stats.bson") total_reward_per_episode time_per_step
    ```
    """

    Experiment(agent, env, stop_condition, hook, description)
end

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:A2C},
    ::Val{:CartPole},
    ::Nothing;
    save_dir = nothing,
)
    N_ENV = 16
    UPDATE_FREQ = 10
    env = MultiThreadEnv([CartPoleEnv(; T = Float32, seed = i) for i in 1:N_ENV])
    ns, na = length(rand(get_observation_space(env[1]))), length(get_action_space(env[1]))
    RLBase.reset!(env, is_force = true)
    agent = Agent(
        policy = QBasedPolicy(
            learner = A2CLearner(
                approximator = ActorCritic(
                    actor = Chain(
                        Dense(ns, 256, relu; initW = seed_glorot_uniform(seed = 17)),
                        Dense(256, na; initW = seed_glorot_uniform(seed = 23)),
                    ),
                    critic = Chain(
                        Dense(ns, 256, relu; initW = seed_glorot_uniform(seed = 29)),
                        Dense(256, 1; initW = seed_glorot_uniform(seed = 29)),
                    ),
                    optimizer = ADAM(1e-3),
                ) |> cpu,
                γ = 0.99f0,
                actor_loss_weight = 1.0f0,
                critic_loss_weight = 0.5f0,
                entropy_loss_weight = 0.001f0,
            ),
            explorer = BatchExplorer(GumbelSoftmaxExplorer()),#= seed = nothing =#
        ),
        trajectory = CircularCompactSARTSATrajectory(;
            capacity = UPDATE_FREQ,
            state_type = Float32,
            state_size = (ns, N_ENV),
            action_type = Int,
            action_size = (N_ENV,),
            reward_type = Float32,
            reward_size = (N_ENV,),
            terminal_type = Bool,
            terminal_size = (N_ENV,),
        ),
    )

    hook = TotalBatchRewardPerEpisode(N_ENV)
    stop_condition = StopAfterStep(100000)
    Experiment(agent, env, stop_condition, hook, "# A2C with CartPole")
end

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:A2CGAE},
    ::Val{:CartPole},
    ::Nothing;
    save_dir = nothing,
)
    N_ENV = 16
    UPDATE_FREQ = 10
    env = MultiThreadEnv([CartPoleEnv(; T = Float32, seed = i) for i in 1:N_ENV])
    ns, na = length(rand(get_observation_space(env[1]))), length(get_action_space(env[1]))
    RLBase.reset!(env, is_force = true)
    agent = Agent(
        policy = QBasedPolicy(
            learner = A2CGAELearner(
                approximator = ActorCritic(
                    actor = NeuralNetworkApproximator(
                        model = Chain(
                            Dense(ns, 256, relu; initW = seed_glorot_uniform(seed = 17)),
                            Dense(256, na; initW = seed_glorot_uniform(seed = 23)),
                        ),
                        optimizer = ADAM(1e-3),
                    ),
                    critic = NeuralNetworkApproximator(
                        model = Chain(
                            Dense(ns, 256, relu; initW = seed_glorot_uniform(seed = 29)),
                            Dense(256, 1; initW = seed_glorot_uniform(seed = 29)),
                        ),
                        optimizer = ADAM(1e-3),
                    ),
                ) |> cpu,
                γ = 0.99f0,
                λ = 0.97f0,
                actor_loss_weight = 1.0f0,
                critic_loss_weight = 0.5f0,
                entropy_loss_weight = 0.001f0,
            ),
            explorer = BatchExplorer(GumbelSoftmaxExplorer(;)), #= seed = nothing =#
        ),
        trajectory = CircularCompactSARTSATrajectory(;
            capacity = UPDATE_FREQ,
            state_type = Float32,
            state_size = (ns, N_ENV),
            action_type = Int,
            action_size = (N_ENV,),
            reward_type = Float32,
            reward_size = (N_ENV,),
            terminal_type = Bool,
            terminal_size = (N_ENV,),
        ),
    )
    Experiment(
        agent,
        env,
        StopAfterStep(100000),
        TotalBatchRewardPerEpisode(N_ENV),
        "# A2CGAE with CartPole",
    )
end

function RLCore.Experiment(::Val{:JuliaRL}, ::Val{:DDPG}, ::Val{:Pendulum}, ::Nothing;)
    inner_env = PendulumEnv(T = Float32, seed = 9231)
    action_space = get_action_space(inner_env)
    low = action_space.low
    high = action_space.high
    ns = length(rand(get_observation_space(inner_env)))

    env = WrappedEnv(;
        env = inner_env,
        postprocessor = ((x,),) -> (low + (x + 1) * 0.5 * (high - low),), # rescale [-1, 1] -> (low, high)
    )

    init = seed_glorot_uniform(seed = 17)

    create_actor() = Chain(
        Dense(ns, 256, relu; initW = init),
        Dense(256, 256, relu; initW = init),
        Dense(256, 1, tanh; initW = init),
    )

    create_critic() = Chain(
        Dense(ns + 1, 256, relu; initW = init),
        Dense(256, 256, relu; initW = init),
        Dense(256, 1; initW = init),
    )

    agent = Agent(
        policy = DDPGPolicy(
            behavior_actor = NeuralNetworkApproximator(
                model = create_actor(),
                optimizer = ADAM(),
            ),
            behavior_critic = NeuralNetworkApproximator(
                model = create_critic(),
                optimizer = ADAM(),
            ),
            target_actor = NeuralNetworkApproximator(
                model = create_actor(),
                optimizer = ADAM(),
            ),
            target_critic = NeuralNetworkApproximator(
                model = create_critic(),
                optimizer = ADAM(),
            ),
            γ = 0.99f0,
            ρ = 0.995f0,
            batch_size = 64,
            start_steps = 1000,
            start_policy = RandomPolicy(ContinuousSpace(-1.0, 1.0); seed = 923),
            update_after = 1000,
            update_every = 1,
            act_limit = 1.0,
            act_noise = 0.1,
            seed = 131,
        ),
        trajectory = CircularCompactSARTSATrajectory(
            capacity = 10000,
            state_type = Float32,
            state_size = (ns,),
            action_type = Float32,
        ),
    )

    description = """
    # Play Pendulum with DDPG
    """

    Experiment(agent, env, StopAfterStep(10000), TotalRewardPerEpisode(), description)
end

function RLCore.Experiment(::Val{:JuliaRL}, ::Val{:PPO}, ::Val{:CartPole}, ::Nothing;)
    N_ENV = 8
    UPDATE_FREQ = 16
    env = MultiThreadEnv([CartPoleEnv(; T = Float32, seed = i) for i in 1:N_ENV])
    ns, na = length(rand(get_observation_space(env[1]))), length(get_action_space(env[1]))
    RLBase.reset!(env, is_force = true)
    agent = Agent(
        policy = QBasedPolicy(
            learner = PPOLearner(
                approximator = ActorCritic(
                    actor = NeuralNetworkApproximator(
                        model = Chain(
                            Dense(ns, 256, relu; initW = seed_glorot_uniform(seed = 17)),
                            Dense(256, na; initW = seed_glorot_uniform(seed = 23)),
                        ),
                        optimizer = ADAM(1e-3),
                    ),
                    critic = NeuralNetworkApproximator(
                        model = Chain(
                            Dense(ns, 256, relu; initW = seed_glorot_uniform(seed = 29)),
                            Dense(256, 1; initW = seed_glorot_uniform(seed = 29)),
                        ),
                        optimizer = ADAM(1e-3),
                    ),
                ) |> cpu,
                γ = 0.99f0,
                λ = 0.95f0,
                clip_range = 0.1f0,
                max_grad_norm = 0.5f0,
                n_epochs = 4,
                n_microbatches = 4,
                actor_loss_weight = 1.0f0,
                critic_loss_weight = 0.5f0,
                entropy_loss_weight = 0.001f0,
            ),
            explorer = BatchExplorer(GumbelSoftmaxExplorer(; seed = 1)),
        ),
        trajectory = PPOTrajectory(;
            capacity = 32,
            state_type = Float32,
            state_size = (ns, N_ENV),
            action_type = Int,
            action_size = (N_ENV,),
            action_log_prob_type = Float32,
            action_log_prob_size = (N_ENV,),
            reward_type = Float32,
            reward_size = (N_ENV,),
            terminal_type = Bool,
            terminal_size = (N_ENV,),
        ),
    )
    Experiment(
        agent,
        env,
        StopAfterStep(100000),
        TotalBatchRewardPerEpisode(N_ENV),
        "# PPO with CartPole",
    )
end

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:BasicDQN},
    ::Val{:MountainCar},
    ::Nothing;
    save_dir = nothing,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyymmddHHMMSS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_BasicDQN_MountainCar_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    env = MountainCarEnv(; T = Float32, max_steps = 5000, seed = 11)
    ns, na = length(rand(get_observation_space(env))), length(get_action_space(env))
    agent = Agent(
        policy = QBasedPolicy(
            learner = BasicDQNLearner(
                approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 64, relu; initW = seed_glorot_uniform(seed = 17)),
                        Dense(64, 64, relu; initW = seed_glorot_uniform(seed = 23)),
                        Dense(64, na; initW = seed_glorot_uniform(seed = 39)),
                    ) |> cpu,
                    optimizer = ADAM(),
                ),
                batch_size = 32,
                min_replay_history = 100,
                loss_func = huber_loss,
                seed = 22,
            ),
            explorer = EpsilonGreedyExplorer(
                kind = :exp,
                ϵ_stable = 0.01,
                decay_steps = 500,
                seed = 33,
            ),
        ),
        trajectory = CircularCompactSARTSATrajectory(
            capacity = 50000,
            state_type = Float32,
            state_size = (ns,),
        ),
    )

    stop_condition = StopAfterStep(40000)

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        DoEveryNStep() do t, agent, env, obs
            with_logger(lg) do
                @info "training" loss = agent.policy.learner.loss
            end
        end,
        DoEveryNEpisode() do t, agent, env, obs
            with_logger(lg) do
                @info "training" reward = total_reward_per_episode.rewards[end]
            end
        end,
        DoEveryNStep(10000) do t, agent, env, obs
            RLCore.save(save_dir, agent)
            BSON.@save joinpath(save_dir, "stats.bson") total_reward_per_episode time_per_step
        end,
    )

    description = """
    This experiment uses three dense layers to approximate the Q value.
    The testing environment is MountainCarEnv.

    Agent and statistic info will be saved to: `$save_dir`
    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`
    To load the agent and statistic info:
    ```
    agent = RLCore.load("$save_dir", Agent)
    BSON.@load joinpath("$save_dir", "stats.bson") total_reward_per_episode time_per_step
    ```
    """

    Experiment(agent, env, stop_condition, hook, description)
end

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:DQN},
    ::Val{:MountainCar},
    ::Nothing;
    save_dir = nothing,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyymmddHHMMSS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_DQN_MountainCar_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    env = MountainCarEnv(; T = Float32, max_steps = 5000, seed = 11)
    ns, na = length(rand(get_observation_space(env))), length(get_action_space(env))

    agent = Agent(
        policy = QBasedPolicy(
            learner = DQNLearner(
                approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 64, relu; initW = seed_glorot_uniform(seed = 17)),
                        Dense(64, 64, relu; initW = seed_glorot_uniform(seed = 23)),
                        Dense(64, na; initW = seed_glorot_uniform(seed = 39)),
                    ) |> cpu,
                    optimizer = ADAM(),
                ),
                target_approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 64, relu; initW = seed_glorot_uniform(seed = 17)),
                        Dense(64, 64, relu; initW = seed_glorot_uniform(seed = 23)),
                        Dense(64, na; initW = seed_glorot_uniform(seed = 39)),
                    ) |> cpu,
                    optimizer = ADAM(),
                ),
                loss_func = huber_loss,
                stack_size = nothing,
                batch_size = 32,
                update_horizon = 1,
                min_replay_history = 100,
                update_freq = 1,
                target_update_freq = 100,
                seed = 22,
            ),
            explorer = EpsilonGreedyExplorer(
                kind = :exp,
                ϵ_stable = 0.01,
                decay_steps = 500,
                seed = 33,
            ),
        ),
        trajectory = CircularCompactSARTSATrajectory(
            capacity = 50000,
            state_type = Float32,
            state_size = (ns,),
        ),
    )

    stop_condition = StopAfterStep(40_000)

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        DoEveryNStep() do t, agent, env, obs
            if agent.policy.learner.update_step % agent.policy.learner.update_freq == 0
                with_logger(lg) do
                    @info "training" loss = agent.policy.learner.loss
                end
            end
        end,
        DoEveryNStep(10000) do t, agent, env, obs
            RLCore.save(save_dir, agent)
            BSON.@save joinpath(save_dir, "stats.bson") total_reward_per_episode time_per_step
        end,
    )

    description = """
    This experiment uses the `DQNLearner` method with three dense layers to approximate the Q value.
    The testing environment is MountainCarEnv.

    Agent and statistic info will be saved to: `$save_dir`
    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`
    To load the agent and statistic info:

    ```
    agent = RLCore.load("$save_dir", Agent)
    BSON.@load joinpath("$save_dir", "stats.bson") total_reward_per_episode time_per_step
    ```
    """

    Experiment(agent, env, stop_condition, hook, description)
end


include("imports.jl")

path_to_output = ("/scratchdata/alui2/cytof/results/repfam/test-sim-6-7-2/" *
                  "maxtemp10-ntempts20-degree2-N500/output.bson")

output = BSON.load(path_to_output)

# 1. For a fixed state (theta), compute_marg_loglike for different temperatures.
tempers = output[:tempers]
# tempers = CytofRepFAM.MCMC.WSPT.gentempers(2, 20, degree=1)
lls = [CytofRepFAM.Model.compute_marg_loglike(output[:lastState].theta,
                                              output[:c].constants,
                                              output[:d].data,
                                              tau) for tau in tempers];
plt.plot(tempers, lls, marker="o")
plt.savefig("tmp.pdf", bbox_inches="tight")
plt.close()

# plt.plot(log.(tempers), lls - [ll[end] for ll in output[:lls]], marker="o")
# plt.savefig("tmp.pdf", bbox_inches="tight")
# plt.close()


module BounWrite
export setboun, timelimit
include("PathFile.jl");
##############################

"""

    setboun(times, heads, location, FEHMprefix)

Sets boundary conditions in the `FEHMprefix.boun` file for FEHM simulations.

Arguments:
- `times`: vector of years at which pumping rate from well is changed
- `hears`: pumping rate.
- `location`: location (FEHM code) of pumping well.
- `FEHMprefix` : FEHM files prefix in `./contaminantion-model/

"""

function setboun(times, heads, location, FEHMprefix)

ff = open(string(fehmdir, FEHMprefix, ".boun"),"w+");
write(ff,"boun
model 3 w21
tran
time\n");
write(ff,string(length(times)));
write(ff,"\n");
times = times./8;
for i = 1:length(times)
	tt = times[i];
	write(ff,"$tt \n");
end
write(ff,"year \ndsw \n");
for i = 1:length(heads)
	hh = heads[i];
	write(ff, "$hh \n");
end
write(ff,"end\n");
write(ff, "-");
write(ff, string(location));
write(ff, " 0 0 1\n ");
write(ff," ");
close(ff);

end

##################################################

"""

    timelimit(days, location, FEHMprefix)

Sets simulation horizing in the `FEHMprefix.data` FEHM file.

Argument:

- `days` : number of days limit
- `location` : location (FEHM code) of pumping well.
- `FEHMprefix` : FEHM files prefix in `./contaminantion-model/

"""

function timelimit(days, location, FEHMprefix)

gg = open(string(fehmdir,FEHMprefix,".data"),"w+");
write(gg,"** Cr model
cont
avs 1 1e20
geo
liquid
concentration
head
velocity
material
endavs
head 1000
nobr
bous
1
airw
3
30.0 0.1
zone
file
gs01_outside.zone
flow
-3 0 0 1798.613472817     1 1e10
-5 0 0 1772.054959807     1 1e10

pres
1 0 0 1776 1 2

zone
110
4.971500000E+05 501750          501750          4.971500000E+05 4.971500000E+05 501750          501750          4.971500000E+05
5.373500000E+05 5.373500000E+05 5.406500000E+05 5.406500000E+05 5.373500000E+05 5.373500000E+05 5.406500000E+05 5.406500000E+05
-1 -1 -1 -1 1 1 1 1

flow
-110 0 0 0.000717864089902394     1 0

zonn
file
hole_coord-20120702-recharge-600p.zonn
flow
-9420 0 0 -0.0003413616876478    1 0
-9620 0 0 -0.00007255554857619   1 0

hyco
1 0 0 -2.51795 -2.51795 -2.51795

rock
1 0 0 2500. 1010. 1.0E-02

rlp
1 0. 0. 1. 1. 0. 1.

1 0 0 1

ppor
-1
1 0 0 0.000215968

stea
shea 1e-6
stim 1e20
smul 1e10
end
time
365.25 $days 10000 001 2006 1 0
0.0  -1.2 1.0 365.25

ctrl
-5 1e-6 1 350 bcgs
1 0 0 2

1.0 0 1.0
10 1000 1e-8 1e20
0 2
sol
1 -1
iter
1e-5 1e-5 0.1 1e-4 1.1
13 0 0 5 1200.
node
1
5471
zone
file
gs01_outside.zone
zonn
file
hole_coord-20140523-screen-change-100.zonn
flxz
1\n")
write(gg,"$location")
write(gg,"\nboun
file\n")
write(gg,FEHMprefix)
write(gg,".boun
zonn
file\n")
write(gg,FEHMprefix)
write(gg,".source.zonn
trac
file\n")
write(gg,FEHMprefix)
write(gg,".trac
stop");
close(gg);
end
end


export PackedFluxModelsDistribution

Base.@kwdef mutable struct PackedFluxModelsDistribution <: PackedSamplableBelief
  # standardized _type field
  _type::String
  # shape of the input data
  inputDim::Vector{Int}
  # shape of the output data
  outputDim::Vector{Int}
  # actual Flux models (Base64 encoded binary)
  mimeTypeModel::String
  models::Vector{String}
  # the data used for prediction, must be <: AbstractArray
  mimeTypeData::String
  data::String
  # shuffle model predictions relative to particle index at each sampling
  shuffle::Bool
  # false for default serialization with model info, set true for separate storage of models
  serializeHollow::Bool
    # TODO remove requirement and standardize sampler API
    # specialSampler::Symbol
  # TODO, only use ._type.  Legacy, field name usage to direct the IIF serialization towards JSON method
  PackedSamplableTypeJSON::String
end


function runfig6()
    p = Params(Ln = 2000, W = 2000, Ls = 20, Ws = 20, scale = 40, λ = 5, α = 0,
    Δ = 1, d = 0, τ = 1)

    # PANELS a-b
    p = reconstruct(p, EZ = SA[0, 2, 0], μN = 0.6)
    presets_fig = Fig4_presets(0.000, 0π/180, 0, 8, true)
    data = ldosonlattice_averaged_sc(p, presets_fig)
    savepsi("fig6psi", p, data[1])
    savepsi("fig6psi", p, data[2])

    p = reconstruct(p, EZ = SA[0, 4, 0], μN = 1.192)
    presets_fig = Fig4_presets(0, 0π/180, 0, 4, true)
    data = ldosonlattice_averaged_sc(p, presets_fig)
    savepsi("fig6psi", p, data[1])

   # PANELS c-d
    p = reconstruct(p, EZ = SA[0, 2, 0], μN = 0.6)
    presets_fig = Fig4_presets(0.002, 2π/180, 4, 8, true)
    data = ldosonlattice_averaged_sc(p, presets_fig)
    savepsi("fig6psi", p, data[1])
    savepsi("fig6psi", p, data[2])

    p = reconstruct(p, EZ = SA[0, 4, 0], μN = 1.192)
    presets_fig = Fig4_presets(0.002, 2π/180, 4, 4, true)
    data = ldosonlattice_averaged_sc(p, presets_fig)
    savepsi("fig6psi", p, data[1])

    # PANELS e-f
    sp = splittingvsrotation(reconstruct(p, Ln = 400, EZ = SA[0, 2, 0], μN = 0.6),
        0:.25:10, true, ϕ0 = pi)
    savespectrum("fig6", p, sp[1], sp[2])

    sp = splittingvsrotation(reconstruct(p, Ln = 400, EZ = SA[0, 4, 0], μN = 1.192),
        0:.25:10, true, ϕ0 = pi)
    savespectrum("fig6", p, sp[1], sp[2])

    # Study vs spectrum vs disorder (not included in the manuscript)
    # sp = splittingvsdisorder(reconstruct(p, Ln = 800, W = 2000, EZ = SA[0, 2, 0]),
    # 0:0.02:0.5, true, ϕ0 = pi)
    # savespectrum("fig6testdisorder2", p, sp[1], sp[2])

    # sp = splittingvsdisorder(reconstruct(p, Ln = 800, W = 2000, EZ = SA[0, 5, 0]),
    #0:0.02:0.5, true, ϕ0 = pi)
    # savespectrum("fig6testdisorder5", p, sp[1], sp[2])
end


using JuLIP: Atoms, AbstractCalculator
export EnergyMixing, prepare_qmmm!

import JuLIP: energy, forces

# ------------------------------------------------------------------

function get_domains(at::Atoms)
   Iqm = get_data(at, "qmmm_e_Iqm")
   # Iqmbuf = get_data(at, "qmmm_e_Iqmbuf")
   Imm = setdiff(1:length(at), Iqm)
   return Iqm, Imm, nothing
end

struct EnergyMixing{T} <: AbstractCalculator
   atqm::Atoms{T}
   QM::AbstractCalculator
   MM::AbstractCalculator
end

function energy(calc::EnergyMixing{T}, at::Atoms{T}) where {T}
   Iqm, Imm, Iqmbuf = get_domains(at)
   # set_positions!(calc.atqm, positions(at)[Iqmbuf])
   Eqm = energy(calc.QM, at; domain = Iqm)
   Emm = energy(calc.MM, at; domain = Imm)
   return Eqm + Emm
end

function forces(calc::EnergyMixing{T}, at::Atoms{T}) where {T}
   Iqm, Imm, Iqmbuf = get_domains(at)
   # set_positions!(calc.atqm, positions(at)[Iqmbuf])
   Fqm = forces(calc.QM, at; domain = Iqm)
   Fmm = forces(calc.MM, at; domain = Imm)
   return Fqm + Fmm
end

# ------------------------------------------------------------------

struct ForceMixing{T}
   atqm::Atoms{T}
   QM::AbstractCalculator
   MM::AbstractCalculator
end

function forces(calc::ForceMixing{T}, at::Atoms{T}) where {T}
   Iqm, Imm, Iqmbuf = get_domains(at)
   Fqm = forces(calc.QM, at; domain = Iqm)
   Fmm = forces(calc.MM, at; domain = Imm)
   return Fqm + Fmm
end

# ------------------------------------------------------------------

prepare_qmmm!(at; kwargs...) = prepare_qmmm!(at, EnergyMixing; kwargs...)

function prepare_qmmm!(at::Atoms, ::Type{EnergyMixing};
                       Vqm = nothing, Vmm = nothing,
                       Iqm = nothing)
   # perform some sanity checks
   Iqm = collect(Iqm)
   @assert Iqm isa Vector{<: Integer}
   @assert Vqm isa AbstractCalculator
   @assert Vmm isa AbstractCalculator
   @assert sort(unique(Iqm)) == sort(Iqm)
   @assert minimum(Iqm) > 0
   @assert maximum(Iqm) <= length(at)

   # create an atoms object for the QM calculator
   # => TODO: add buffer, carve out region, etc...

   # store the QM region in `at`
   at["qmmm_e_Iqm"] = Iqm
   # store the new calculator
   set_calculator!(at, EnergyMixing(at, Vqm, Vmm))

   return at
end