##############################
# Observation-based learning #
##############################
"""
Observation-based network learning model type.
"""
mutable struct NetworkLearnerObs{T,U,S,V,
				    R<:Vector{<:AbstractRelationalLearner},
				    C<:AbstractCollectiveInferer,
				    A<:Vector{<:AbstractAdjacency},
				    L<:Union{Nothing, MLLabelUtils.LabelEncoding},
				    OD<:LearnBase.ObsDimension} <: AbstractNetworkLearner
	Ml::T											# local model
	fl_exec::U										# local model execution function
	Mr::S											# relational model
	fr_exec::V										# relational model execution function
	RL::R											# relational learner
	Ci::C											# collective inferer
	Adj::A											# adjacency information
	use_local_data::Bool							# whether to use local data
	target_enc::L									# target encoding
	size_in::Int									# expected input dimensionality
	size_out::Int									# expected output dimensionality
	obsdim::OD										# observation dimension
end

# Printers
Base.show(io::IO, m::NetworkLearnerObs) = begin
	println(io,"NetworkLearner, I/O size $(m.size_in)×$(m.size_out), observation based")
	print(io,"`- local model: "); println(io, m.Ml)
	print(io,"`- relational model: "); println(io, m.Mr)
	print(io,"`- relational learners: "); println(io, m.RL)
	print(io,"`- collective inferer: "); println(io, m.Ci)
	print(io,"`- adjacency: "); println(io, m.Adj)
	println(io,"`- use local data: $(m.use_local_data)");
	println(io,"`- observations are: $(m.obsdim == ObsDim.Constant{1} ? "rows" : "columns")");
	println(io,"`- targets: $(m.target_enc isa Nothing ? "not encoded" : "encoded")");
end

####################
# Training methods #
####################
"""
	fit(::Type{NetworkLearnerObs}, X, y, Adj, fl_train, fl_exec, fr_train, fr_exec [;kwargs])

Training method for the observation-based network learning framework.

# Arguments
  * `X::AbstractMatrix` training data (used by `fl_train`, `fl_exec`; if `use_local_data==true`, it is also used by `fr_train`)
  * `Y::AbstractAray` data targets (used by `fl_train`, `fr_train`)
  * `Adj::Vector{AbstractAdjacency}` a vector containing the observation relational structures (adjacency objects)
  * `fl_train` local model training 'function'; can be anything that supports the call `fl_train((X,y))`
  * `fl_exec` local model prediction 'function'; can be anything that supports the call `fl_exec(Ml,X)` where `Ml = fl_train((X,y))`
  * `fr_train` relational model training `function`; can be anything that suports the call `fr_train((Xr,y))`
  * `fr_exec` relational model prediction `function`; can be anything that suports the call `fr_exec(Mr,Xr)` where `Mr = fr_train((Xr,y))`
  and `Xr` is a dataset of relational variables generated by the relational learner using the results of the local model prediction
  function and adjacency structures.

# Keyword arguments
  * `priors::Vector{Float64}` class priors (if applicable)
  * `learner::Symbol` relational learner (i.e. variable generator); available options `:rn`, `:wrn`, `:bayesrn` and `:cdrn` (default `:wrn`)
  * `inference::Symbol` collective inference method; available options `:rl`, `:ic` and `:gs` (default `:rl`)
  * `normalize::Bool` whether to normalize the relational variables per-observation to the L1 norm (default `true`)
  * `use_local_data::Bool` whether the relational model should use the local data provided (i.e. in `X`) (default `true`)
  * `f_targets::Function` function that extracts targets from estimates generated by the local/relational models
  (default `f_targets = x->MLDataPattern.targets(indmax,x)`)
  * `obsdim::Int` observation dimension (default `2`)
  * `tol::Float64` maximum admissible mean estimate error for collective inference convergence (default `1e-6`)
  * `κ::Float64` relaxation labeling starting constant, used if `learner == :rl` (default `1.0`)
  * `α::Float64` relaxation labeling decay constant, used if `learner == :rl` (default `0.99`)
  * `maxiter::Int` maximum number of iterations for collective inference (default `100`)
  * `bratio::Float64` percentage of iterations i.e. `maxiter` used for Gibbs sampling burn-in (default `0.1`)
"""
function fit(::Type{NetworkLearnerObs}, X::AbstractMatrix, y::AbstractArray, Adj::A where A<:Vector{<:AbstractAdjacency},
		fl_train, fl_exec, fr_train, fr_exec;
		priors::Vector{Float64}=getpriors(y), learner::Symbol=:wvrn, inference::Symbol=:rl,
		normalize::Bool=true, use_local_data::Bool=true, f_targets::Function=x->targets(indmax,x),
		obsdim::Int = 2, tol::Float64=1e-6, κ::Float64=1.0, α::Float64=0.99, maxiter::Int=100, bratio::Float64=0.1)

	# Parse, transform input arguments
	κ = clamp(κ, 1e-6, 1.0)
	α = clamp(α, 1e-6, 1.0-1e-6)
	tol = clamp(tol, 0.0, Inf)
	maxiter = ifelse(maxiter<=0, 1, maxiter)
	bratio = clamp(bratio, 1e-6, 1.0-1e-6)

	@assert all((priors.>=0.0) .& (priors .<=1.0)) "All priors have to be between 0.0 and 1.0."
	@assert obsdim in [1,2] "Observation dimension can have only two values 1 (row-major) or 2 (column-major)."

	# Parse relational learner argument and generate relational learner type
	if learner == :rn
		Rl = SimpleRN
	elseif learner == :wrn
		Rl = WeightedRN
	elseif learner == :cdrn
		Rl = ClassDistributionRN
	elseif learner == :bayesrn
		Rl = BayesRN
	else
		@print_verbose 1 "Unknown relational learner. Defaulting to :wrn."
		Rl = WeightedRN
	end

	# Parse collective inference argument and generate collective inference objects
	if inference == :rl
		Ci = RelaxationLabelingInferer(maxiter, tol, f_targets, κ, α)
	elseif inference == :ic
		Ci = IterativeClassificationInferer(maxiter, tol, f_targets)
	elseif inference == :gs
		Ci = GibbsSamplingInferer(maxiter, tol, f_targets, ceil(Int, maxiter*bratio))
	else
		@print_verbose 1 "Unknown collective inferer. Defaulting to :rl."
		Ci = RelaxationLabelingInferer(maxiter, tol, f_targets, κ, α)
	end

	fit(NetworkLearnerObs, X, y, Adj, Rl, Ci, fl_train, fl_exec, fr_train, fr_exec;
		priors=priors, normalize=normalize, use_local_data=use_local_data,
		obsdim=ObsDim.Constant{obsdim}())
end

"""
Training method for the network learning framework. This method should not be called directly.
"""
function fit(::Type{NetworkLearnerObs}, X::T, y::S, Adj::A, Rl::R, Ci::C, fl_train::U, fl_exec::U2, fr_train::U3, fr_exec::U4;
		priors::Vector{Float64}=getpriors(y), normalize::Bool=true, use_local_data::Bool=true,
		obsdim::LearnBase.ObsDimension=ObsDim.Constant{2}()) where {
			T<:AbstractMatrix,
			S<:AbstractArray,
			A<:Vector{<:AbstractAdjacency},
			R<:Type{<:AbstractRelationalLearner},
			C<:AbstractCollectiveInferer,
			U, U2, U3, U4
		}

	# Step 0: pre-process input arguments and retrieve sizes
	size_in = nvars(X,obsdim)								# number of local variables
	size_out = get_size_out(y)								# number of relational variables / adjacency
	n = nobs(X,obsdim)									# number of observations
	m = length(Adj) * size_out								# number of relational variables

	@assert size_out == length(priors) "Found $c classes, the priors indicate $(length(priors))."

	# Pre-allocate relational variables array
	if use_local_data									# Local observation variable data is used
		Xr = matrix_prealloc(n, size_in+m, obsdim, 0.0)
		_Xr = datasubset(Xr, 1:size_in, oppdim(obsdim))
		_Xr[:] = X
		offset = size_in
	else											# Only relational variables are used
		Xr = matrix_prealloc(n, m, obsdim, 0.0)
		offset = 0
	end

	# Encode targets
	(t_enc, yₑ) = encode_targets(y)

	# Step 1: train and execute local model
	@print_verbose 2 "Training and applying local model ..."
	Dl = (X,yₑ)
	Ml = fl_train(Dl);
	Xl = fl_exec(Ml,X);

	# Step 2: Get relational variables by training and executing the relational learner
	@print_verbose 2 "Calculating relational variables ..."
	RL = [fit(Rl, Aᵢ, Xl, yₑ; obsdim=obsdim, priors=priors, normalize=normalize)
       		for Aᵢ in Adj]									# Train relational learners

	Xrᵢ = matrix_prealloc(n, size_out, obsdim, 0.0)						# Initialize temporary storage
	for (i,(RLᵢ,Aᵢ)) in enumerate(zip(RL,Adj))

		# Apply relational learner
		transform!(Xrᵢ, RLᵢ, Aᵢ, Xl, yₑ)

		# Update relational data output
		_Xr = datasubset(Xr, offset+(i-1)*size_out+1:offset+i*size_out, oppdim(obsdim))
		_Xr[:] = Xrᵢ
	end

	# Step 3 : train relational model
	@print_verbose 2 "Training relational model ..."
	Mr = fr_train((Xr,yₑ))

	# Step 4: remove adjacency data
	Adj_s = AbstractAdjacency[];
	for i in 1:length(Adj)
		push!(Adj_s, strip_adjacency(Adj[i]))
	end

	# Step 5: return network learner
	@print_verbose 2 "Done."
	return NetworkLearnerObs(Ml, fl_exec, Mr, fr_exec, RL, Ci, Adj_s, use_local_data, t_enc, size_in, size_out, obsdim)
end

#####################
# Execution methods #
#####################
"""
Prediction method for the network learning framework.
"""
function predict(model::M, X::T, update::BitVector=trues(nobs(X,model.obsdim))) where {M<:NetworkLearnerObs, T<:AbstractMatrix}
	Xo = matrix_prealloc(nobs(X,model.obsdim), model.size_out, model.obsdim, 0.0)
	predict!(Xo, model, X, update)
	return Xo
end

"""
In-place prediction method for the network learning framework.
"""
function predict!(Xo::S, model::M, X::T, update::BitVector=trues(nobs(X,model.obsdim))) where {M<:NetworkLearnerObs, T<:AbstractMatrix, S<:AbstractMatrix}

	# Step 0: Make initializations and pre-allocations
	obsdim = model.obsdim									# observation dimension
	m = nvars(X,obsdim)									# number of input variables
	n = nobs(X,obsdim)									# number of observations
	l = length(model.Adj)									# number of adjacencies in the model

	@assert model.size_in == m "Expected input dimensionality $(model.size_in), got $m."
	@assert (nvars(Xo,obsdim),nobs(Xo,obsdim)) == (model.size_out, n) "Output dataset size does not match expected size ($(model.size_out) variables×$n observations)."

	if model.use_local_data									# Pre-allocate relational dataset based on local data use
		Xr = matrix_prealloc(n, model.size_out*l+m, obsdim, 0.0)
		_Xr = datasubset(Xr, 1:m, oppdim(obsdim))
		_Xr[:] = X									#  - initialize (local) dimensions
		offset = m
	else											# Skip local dataset dimensions
		Xr = matrix_prealloc(n, model.size_out*l, obsdim, 0.0)
		offset = 0
	end

	# Step 1: Apply local model, initialize output
	@print_verbose 2 "Executing local model ..."
	Xl = model.fl_exec(model.Ml, getobs(datasubset(X, update, obsdim)))
	@assert nvars(Xo,obsdim) == nvars(Xl,obsdim) "Local model outputs $(nvars(Xl,obsdim)) estimates/obs, the output indicates $(nvars(Xo,obsdim))."
	_Xo = datasubset(Xo, update, obsdim);
	_Xo[:] = Xl

	# Step 2: Apply collective inference
	@print_verbose 2 "Collective inference ..."
	transform!(Xo, model.Ci, model.obsdim, model.Mr, model.fr_exec, model.RL, model.Adj, offset, Xr, update)

	# Step 3: Return output estimates
	@print_verbose 2 "Done."
	return Xo
end

# It may be necessary to add adjacency information to the model, regarding the test data
"""
	add_adjacency!(model, Av)

Function that adds or, creates and then adds, adjacency objects contained in a vector `Av` to a network
learning `model`. The method is used in 'ouf-of-graph' learning i.e. the training data is not used in
the predictions for new samples. In such circumstances, the adjacency structures - graphs, matrices etc.
pertinent to the new observations have to be created either using the same functions used in training or,
separately supplied, so that relational variables can be created.
"""
function add_adjacency!(model::M, Av::Vector{T}) where {M<:NetworkLearnerObs, T<:AbstractAdjacency}
	@assert length(Av) == length(model.Adj) "New adjacency vector must have a length of $(length(model.Adj))."
	model.Adj = Av
end

function add_adjacency!(model::M, Av::Vector{T}) where {M<:NetworkLearnerObs, T}
	@assert length(Av) == length(model.Adj) "Adjacency data vector must have a length of $(length(model.Adj))."
	model.Adj = adjacency.(Av)
end


# cubic is default

struct Linear end
struct Cubic end

struct CompletePolnomial{order} end
struct Smolyak end

# incomplete type inference:
DecisionRule(exo_grid, endo_grid, i::Int64) = DecisionRule(exo_grid, endo_grid, Val{i})
DecisionRule(exo_grid, endo_grid, i::Int64, interp_type) = DecisionRule(exo_grid, endo_grid, Val{i}, interp_type)

# cubic interpolation
function DecisionRule(exo_grid::EmptyGrid, endo_grid::CartesianGrid, ::Type{Val{nx}}, ::Type{Cubic}) where nx
    CubicDR(exo_grid, endo_grid, Val{nx})
end

function DecisionRule(exo_grid::UnstructuredGrid, endo_grid::CartesianGrid, ::Type{Val{nx}}, ::Type{Cubic}) where nx
    CubicDR(exo_grid, endo_grid, Val{nx})
end

function DecisionRule(exo_grid::CartesianGrid, endo_grid::CartesianGrid, ::Type{Val{nx}}, ::Type{Cubic}) where nx
    CubicDR(exo_grid, endo_grid, Val{nx})
end

# Smolyak interpolation
function DecisionRule(
        exo_grid::S, endo_grid::SmolyakGrid, ::Type{Val{nx}}, ::Type{Smolyak}
    ) where S <: Union{EmptyGrid,<:UnstructuredGrid} where nx
    SmolyakDR(exo_grid, endo_grid, Val{nx})
end

# Complete polynomials
function DecisionRule(
        exo_grid::S, endo_grid::Grid, ::Type{Val{nx}}, ::Type{CompletePolnomial{order}}
    ) where S <: Union{EmptyGrid,<:UnstructuredGrid} where nx where order
    CompletePolyDR(exo_grid, endo_grid, Val{nx}, order)
end

# default to order 3
function DecisionRule(
        exo_grid::S, endo_grid::Grid, ::Type{Val{nx}}, ::Type{<:CompletePolnomial}
    ) where S <: Union{EmptyGrid,<:UnstructuredGrid} where nx
    CompletePolyDR(exo_grid, endo_grid, Val{nx})
end

# generic constructors

function DecisionRule(exo_grid::Grid, endo_grid::Grid, vals::Vector{ListOfPoints{n_x}}) where n_x
    dr = DecisionRule(exo_grid, endo_grid, Val{n_x}, Cubic)
    set_values!(dr, vals)
    dr
end

function DecisionRule(exo_grid::Grid, endo_grid::Grid, vals::Vector{ListOfPoints{n_x}}, tt) where n_x
    dr = DecisionRule(exo_grid, endo_grid, Val{n_x}, tt)
    set_values!(dr, vals)
    dr
end

# #####
# ##### Cached Decision Rules (do we really need them ?)
# #####

mutable struct CachedDecisionRule{T,S}
    dr::T
    process::S
end

CachedDecisionRule(process::AbstractDiscretizedProcess, grid::Grid, n_x::Int) =
    CachedDecisionRule(DecisionRule(process.grid, grid, n_x), process)

CachedDecisionRule(process::AbstractDiscretizedProcess, grid::Grid, values) =
    CachedDecisionRule(DecisionRule(process.grid, grid, values), process)

set_values!(cdr::CachedDecisionRule, v) = set_values!(cdr.dr, v)

# defaults

(cdr::CachedDecisionRule)(v::Union{<:Point,<:ListOfPoints}) = evaluate(cdr.dr,v)
(cdr::CachedDecisionRule)(m::Union{<:Point,<:ListOfPoints}, s::Union{<:Point,<:ListOfPoints}) = evaluate(cdr.dr, m, s)

# Interpolating rules
(cdr::CachedDecisionRule{<:AbstractDecisionRule{T,<:Grid}, <:AbstractDiscretizedProcess})(i::Int, s::Union{<:Point,<:ListOfPoints}) where T<:Union{CartesianGrid, RandomGrid} = evaluate(cdr.dr,node(Point, cdr.process, i), s)

(cdr::CachedDecisionRule{<:AbstractDecisionRule{T,<:Grid}, <:AbstractDiscretizedProcess})(i::Int, j::Int, s::Union{<:Point,<:ListOfPoints}) where T<:Union{CartesianGrid, RandomGrid} = evaluate(cdr.dr,inode(Point, cdr.process, i, j), s)

(cdr::CachedDecisionRule{<:AbstractDecisionRule{<:EmptyGrid,<:Grid}, <:AbstractDiscretizedProcess})(i::Int, j::Int, s::Union{<:Point,<:ListOfPoints}) = evaluate(cdr.dr, s)

(cdr::CachedDecisionRule{<:AbstractDecisionRule{<:EmptyGrid,<:Grid}, <:AbstractDiscretizedProcess})(i::Int, s::Union{<:Point,<:ListOfPoints}) = evaluate(cdr.dr, s)

# maybe keep only the i,j,s calls.

(cdr::CachedDecisionRule{<:AbstractDecisionRule{T,<:Grid}, <:DiscreteMarkovProcess})(i::Int, s::Union{<:Point,<:ListOfPoints}) where T<:Union{CartesianGrid, RandomGrid} = evaluate(cdr.dr,node(Point, cdr.process, i), s)
(cdr::CachedDecisionRule{<:AbstractDecisionRule{T,<:Grid}, <:DiscreteMarkovProcess})(i::Int, j::Int, s::Union{<:Point,<:ListOfPoints}) where T<:Union{CartesianGrid, RandomGrid} = evaluate(cdr.dr,inode(Point, cdr.process, i, j), s)

(cdr::CachedDecisionRule{<:AbstractDecisionRule{<:UnstructuredGrid,<:Grid}, <:DiscreteMarkovProcess})(i::Int, s::Union{<:Point,<:ListOfPoints}) = evaluate(cdr.dr, i, s)
(cdr::CachedDecisionRule{<:AbstractDecisionRule{<:UnstructuredGrid,<:Grid}, <:DiscreteMarkovProcess})(i::Int, j::Int, s::Union{<:Point,<:ListOfPoints}) = evaluate(cdr.dr, j, s)

#
# # We don't do compat for CachedDR cuz it's a nigthmare.
#
# (cdr::CachedDecisionRule{<:AbstractDecisionRule{<:UnstructuredGrid,<:Grid}, <:DiscreteMarkovProcess})(i::Int, s::Union{Vector{Float64},Matrix{Float64}}) = evaluate(cdr.dr, i, s)
# (cdr::CachedDecisionRule{<:AbstractDecisionRule{<:UnstructuredGrid,<:Grid}, <:DiscreteMarkovProcess})(i::Int, j::Int, s::Union{Vector{Float64},Matrix{Float64}}) = evaluate(cdr.dr, j, s)
#
#
# (cdr::CachedDecisionRule)(v::Union{Vector{Float64},Matrix{Float64}}) = evaluate(cdr.dr,v)
# (cdr::CachedDecisionRule)(m::Union{Vector{Float64},Matrix{Float64}}, s::Union{Vector{Float64},Matrix{Float64}}) = evaluate(cdr.dr, m, s)
# (cdr::CachedDecisionRule)(i::Int, s::Union{Vector{Float64},Matrix{Float64}}) = evaluate(cdr.dr,node(cdr.process, i), s)
# (cdr::CachedDecisionRule)(i::Int, j::Int, s::Union{Vector{Float64},Matrix{Float64}}) = evaluate(cdr.dr,inode(cdr.process, i, j), s)


# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    CookieCutter(master, others)

A cookie-cutter simulation solver.

## Parameters

* `master` - Master simulation solver (a.k.a. facies solver)
* `others` - A list of pairs mapping categories to solvers

## Examples

Simulate lithology facies with image quilting and fill property
with direct Gaussian simulation:

```julia
julia> fsolver  = ImgQuilt(:facies => (TI=Strebelle, template=(30,30,1)))
julia> psolver₀ = DirectGaussSim(:property => (variogram=SphericalVariogram(range=10.),))
julia> psolver₁ = DirectGaussSim(:property => (variogram=SphericalVariogram(range=20.),))
julia> solver   = CookieCutter(fsolver, [0 => psolver₀, 1 => psolver₁])
```
"""
struct CookieCutter <: AbstractSimulationSolver
  master::AbstractSimulationSolver
  others::Vector{Pair{Number,AbstractSimulationSolver}}
end

function solve(problem::SimulationProblem, solver::CookieCutter)
  # retrieve problem info
  pdata   = data(problem)
  pdomain = domain(problem)
  pvars   = variables(problem)
  preals  = nreals(problem)
  pmaps   = datamap(problem)

  # master variable
  mvars = collect(keys(solver.master.params))
  @assert length(mvars) == 1 "one single variable must be specified in master solver"
  mvar = mvars[1]
  @assert mvar ∈ keys(pvars) "invalid variable in master solver"
  mtype = pvars[mvar]

  # other variables
  ovars = Tuple(var => V for (var, V) in pvars if var ≠ mvar)
  @assert length(ovars) > 0 "cookie-cutter requires problem with more than one target variable"

  # copy mappings for master variable
  mmapper = CopyMapper(collect(values(pmaps[mvar])), collect(keys(pmaps[mvar])))

  # define master problem
  if hasdata(problem)
    mproblem = SimulationProblem(pdata, pdomain, mvar, preals, mapper=mmapper)
  else
    mproblem = SimulationProblem(pdomain, mvar => mtype, preals, mapper=mmapper)
  end

  # solve master problem
  msolution = solve(mproblem, solver.master)

  # pre-allocate memory for result
  realizations = msolution.realizations
  for (var, V) in ovars
    reals = map(1:preals) do i
      # find inactive locations in master variable
      inactive = findall(isnan, realizations[mvar][i])

      # copy inactive locations to other variable
      real = Vector{V}(undef, npoints(pdomain))
      view(real, inactive) .= V(NaN)

      real
    end
    realizations[var] = reals
  end

  # find mappings for all other variables
  omaps = merge([pmaps[var] for var in first.(ovars)]...)

  # solve other problems
  for (i, mreal) in enumerate(realizations[mvar])
    for (mval, osolver) in solver.others
      # lookup indices with given master value
      inds = findall(isequal(mval), mreal)

      if !isempty(inds)
        # define subdomain
        odomain = view(pdomain, inds)

        # find mappings for subdomain
        datlocs = Vector{Int}()
        domlocs = Vector{Int}()
        for (j, ind) in enumerate(inds)
          if ind ∈ keys(omaps)
            push!(datlocs, omaps[ind])
            push!(domlocs, j)
          end
        end

        # copy hard data if needed
        omapper = CopyMapper(datlocs, domlocs)

        # define subproblem
        if hasdata(problem)
          oproblem = SimulationProblem(pdata, odomain, first.(ovars), 1, mapper=omapper)
        else
          oproblem = SimulationProblem(odomain, ovars, 1, mapper=omapper)
        end

        # solve subproblem
        osolution = solve(oproblem, osolver)

        # save result and continue
        for (var, V) in ovars
          view(realizations[var][i], inds) .= osolution.realizations[var][1]
        end
      end
    end
  end

  SimulationSolution(pdomain, realizations)
end

# ------------
# IO methods
# ------------
function Base.show(io::IO, solver::CookieCutter)
  print(io, "CookieCutter")
end

function Base.show(io::IO, ::MIME"text/plain", solver::CookieCutter)
  mvar = collect(keys(solver.master.params))[1]
  println(io, solver)
  println(io, "  └─", mvar, " ⇨ ", solver.master)
  for (val, osolver) in solver.others
    println(io, "    └─", val, " ⇨ ", osolver)
  end
end


# Default implementations of `sample`.
const PROGRESS = Ref(true)

"""
    setprogress!(progress::Bool)

Enable progress logging globally if `progress` is `true`, and disable it otherwise.
"""
function setprogress!(progress::Bool)
    @info "progress logging is $(progress ? "enabled" : "disabled") globally"
    PROGRESS[] = progress
    return progress
end

function StatsBase.sample(
    model::AbstractModel,
    sampler::AbstractSampler,
    arg;
    kwargs...
)
    return StatsBase.sample(Random.GLOBAL_RNG, model, sampler, arg; kwargs...)
end

"""
    sample([rng, ]model, sampler, N; kwargs...)

Return `N` samples from the `model` with the Markov chain Monte Carlo `sampler`.
"""
function StatsBase.sample(
    rng::Random.AbstractRNG,
    model::AbstractModel,
    sampler::AbstractSampler,
    N::Integer;
    kwargs...
)
    return mcmcsample(rng, model, sampler, N; kwargs...)
end

"""
    sample([rng, ]model, sampler, isdone; kwargs...)

Sample from the `model` with the Markov chain Monte Carlo `sampler` until a
convergence criterion `isdone` returns `true`, and return the samples.

The function `isdone` has the signature
```julia
isdone(rng, model, sampler, samples, iteration; kwargs...)
```
and should return `true` when sampling should end, and `false` otherwise.
"""
function StatsBase.sample(
    rng::Random.AbstractRNG,
    model::AbstractModel,
    sampler::AbstractSampler,
    isdone;
    kwargs...
)
    return mcmcsample(rng, model, sampler, isdone; kwargs...)
end

function StatsBase.sample(
    model::AbstractModel,
    sampler::AbstractSampler,
    parallel::AbstractMCMCEnsemble,
    N::Integer,
    nchains::Integer;
    kwargs...
)
    return StatsBase.sample(Random.GLOBAL_RNG, model, sampler, parallel, N, nchains;
                            kwargs...)
end

"""
    sample([rng, ]model, sampler, parallel, N, nchains; kwargs...)

Sample `nchains` Monte Carlo Markov chains from the `model` with the `sampler` in parallel
using the `parallel` algorithm, and combine them into a single chain.
"""
function StatsBase.sample(
    rng::Random.AbstractRNG,
    model::AbstractModel,
    sampler::AbstractSampler,
    parallel::AbstractMCMCEnsemble,
    N::Integer,
    nchains::Integer;
    kwargs...
)
    return mcmcsample(rng, model, sampler, parallel, N, nchains; kwargs...)
end

# Default implementations of regular and parallel sampling.

function mcmcsample(
    rng::Random.AbstractRNG,
    model::AbstractModel,
    sampler::AbstractSampler,
    N::Integer;
    progress = PROGRESS[],
    progressname = "Sampling",
    callback = nothing,
    discard_initial = 0,
    thinning = 1,
    chain_type::Type=Any,
    kwargs...
)
    # Check the number of requested samples.
    N > 0 || error("the number of samples must be ≥ 1")
    Ntotal = thinning * (N - 1) + discard_initial + 1

    # Start the timer
    start = time()
    local state

    @ifwithprogresslogger progress name=progressname begin
        # Determine threshold values for progress logging
        # (one update per 0.5% of progress)
        if progress
            threshold = Ntotal ÷ 200
            next_update = threshold
        end

        # Obtain the initial sample and state.
        sample, state = step(rng, model, sampler; kwargs...)

        # Discard initial samples.
        for i in 1:(discard_initial - 1)
            # Update the progress bar.
            if progress && i >= next_update
                ProgressLogging.@logprogress i/Ntotal
                next_update = i + threshold
            end

            # Obtain the next sample and state.
            sample, state = step(rng, model, sampler, state; kwargs...)
        end

        # Run callback.
        callback === nothing || callback(rng, model, sampler, sample, state, 1; kwargs...)

        # Save the sample.
        samples = AbstractMCMC.samples(sample, model, sampler, N; kwargs...)
        samples = save!!(samples, sample, 1, model, sampler, N; kwargs...)

        # Update the progress bar.
        itotal = 1 + discard_initial
        if progress && itotal >= next_update
            ProgressLogging.@logprogress itotal / Ntotal
            next_update = itotal + threshold
        end

        # Step through the sampler.
        for i in 2:N
            # Discard thinned samples.
            for _ in 1:(thinning - 1)
                # Obtain the next sample and state.
                sample, state = step(rng, model, sampler, state; kwargs...)

                # Update progress bar.
                if progress && (itotal += 1) >= next_update
                    ProgressLogging.@logprogress itotal / Ntotal
                    next_update = itotal + threshold
                end
            end

            # Obtain the next sample and state.
            sample, state = step(rng, model, sampler, state; kwargs...)

            # Run callback.
            callback === nothing || callback(rng, model, sampler, sample, state, i; kwargs...)

            # Save the sample.
            samples = save!!(samples, sample, i, model, sampler, N; kwargs...)

            # Update the progress bar.
            if progress && (itotal += 1) >= next_update
                ProgressLogging.@logprogress itotal / Ntotal
                next_update = itotal + threshold
            end
        end
    end

    # Get the sample stop time.
    stop = time()
    duration = stop - start
    stats = SamplingStats(start, stop, duration)

    return bundle_samples(
        samples,
        model,
        sampler,
        state,
        chain_type;
        stats=stats,
        discard_initial=discard_initial,
        thinning=thinning,
        kwargs...
    )
end

function mcmcsample(
    rng::Random.AbstractRNG,
    model::AbstractModel,
    samplers::Vector{<:AbstractSampler},
    swap_every::Integer,
    N::Integer;
    progress = PROGRESS[],
    progressname = "Replica Exchange: Sampling ($(length(samplers)) replicas)",
    callback = nothing,
    discard_initial = 0,
    thinning = 1,
    chain_type::Type=Any,
    kwargs...
)
    """
    Replica exchange
    """
    println("Replica Exchange AbstractMCMC.mcmcsample")
    function swap_β(samplers::Vector{<:AbstractSampler}, states, start::Integer)
        L = length(samplers) - 1
        for sampler_id in start:2:L
            println("Attempting swap b/w $sampler_id, $(sampler_id+1)")
            T = typeof(states[sampler_id].z.ℓπ)
            for (name, typ) in zip(fieldnames(T), T.types)
                println("type of the fieldname $name is $typ")
            end
            println(states[sampler_id].z.ℓπ.value)
            logα = - (samplers[sampler_id].alg.β - samplers[sampler_id + 1].alg.β) * (states[sampler_id].z.ℓπ.value - states[sampler_id+1].z.ℓπ.value)
            if log(1-Random.rand(rng)) ≤ logα
                println("Successful swap b/w $sampler_id, $(sampler_id+1)")
                @set samplers[sampler_id].alg.β, samplers[sampler_id + 1].alg.β = samplers[sampler_id + 1].alg.β, samplers[sampler_id].alg.β
            else
                println("Failed swap b/w $sampler_id, $(sampler_id+1)")
            end
        end
    end
    # Check the number of requested samples.
    N > 0 || error("the number of samples must be ≥ 1")
    Ntotal = thinning * (N - 1) + discard_initial + 1

    # Start the timer
    start = time()
    local states, state, itotal
    states, samples_per_beta = [], Dict()
    βs = [sampler.alg.β for sampler in samplers]
    @ifwithprogresslogger progress name=progressname begin
        # Determine threshold values for progress logging
        # (one update per 0.5% of progress)
        if progress
            threshold = Ntotal ÷ 200
            next_update = threshold
        end

        for (sampler_id, sampler) in enumerate(samplers)
            # Obtain the initial sample and state.
            sample, state = step(rng, model, sampler; kwargs...)

            push!(states, state)

            # Discard initial samples.
            for i in 1:(discard_initial - 1)
                # Update the progress bar.
                if progress && i >= next_update && sampler_id == length(samplers)
                    ProgressLogging.@logprogress i/Ntotal
                    next_update = i + threshold
                end

                # Obtain the next sample and state.
                sample, states[sampler_id] = step(rng, model, sampler, states[sampler_id]; kwargs...)
            end

            # Run callback.
            callback === nothing || callback(rng, model, sampler, sample, states[sampler_id], 1; kwargs...)

            # Save the sample.
            samples = AbstractMCMC.samples(sample, model, sampler, N; kwargs...)
            samples = save!!(samples, sample, 1, model, sampler, N; kwargs...)
            samples_per_beta[sampler.alg.β] = samples
            # Update the progress bar if it is the last replica
            itotal = 1 + discard_initial
            println(itotal)
            if progress && itotal >= next_update && sampler_id == length(samplers)
                ProgressLogging.@logprogress itotal / Ntotal
                next_update = itotal + threshold
            end
        end
        println(itotal)
        # Step through the sampler.
        for i in 2:N
            # # Discard thinned samples.
            # for _ in 1:(thinning - 1)
            #     # Obtain the next sample and state.
            #     sample, state = step(rng, model, sampler, state; kwargs...)

            #     # Update progress bar.
            #     if progress && (itotal += 1) >= next_update
            #         ProgressLogging.@logprogress itotal / Ntotal
            #         next_update = itotal + threshold
            #     end
            # end

            if i % swap_every == 0
                if i % (2*swap_every) == 0
                    swap_β(samplers, states, 2) # swap even indices
                else
                    swap_β(samplers, states, 1) # swap odd indices
                end
            end
            for (sampler_id, sampler) in enumerate(samplers)
                # Load state of replica 'sampler_id'
                state = states[sampler_id]

                # Obtain the next sample and state for the replica
                sample, state = step(rng, model, sampler, state; kwargs...)

                # Run callback for the replica
                callback === nothing || callback(rng, model, sampler, sample, state, i; kwargs...)

                # Save state of replica 'sampler_id'
                states[sampler_id] = state

                # Save the sample for the replica
                samples_per_beta[sampler.alg.β] = save!!(samples_per_beta[sampler.alg.β], sample, i, model, sampler, N; kwargs...)

                # Update the progress bar if it is the last replica
                if progress && (itotal += 1) >= next_update && sampler_id == length(samplers)
                    ProgressLogging.@logprogress itotal / Ntotal
                    next_update = itotal + threshold
                end
            end
        end
    end

    # Get the sample stop time.
    stop = time()
    duration = stop - start
    stats = SamplingStats(start, stop, duration)

    # println(samples_per_beta)

    bundled_samples = Dict()
    for β in βs
        for (sampler_id, sampler) in enumerate(samplers)
            if sampler.alg.β == β
                state = states[sampler_id]
                samples = samples_per_beta[β]
                bundled_samples[β] =  bundle_samples(
                    samples,
                    model,
                    sampler,
                    state,
                    chain_type;
                    stats=stats,
                    discard_initial=discard_initial,
                    thinning=thinning,
                    kwargs...
                )
            end
        end
    end
    return bundled_samples
end

function mcmcsample(
    rng::Random.AbstractRNG,
    model::AbstractModel,
    sampler::AbstractSampler,
    isdone;
    chain_type::Type=Any,
    progress = PROGRESS[],
    progressname = "Convergence sampling",
    callback = nothing,
    discard_initial = 0,
    thinning = 1,
    kwargs...
)

    # Start the timer
    start = time()
    local state

    @ifwithprogresslogger progress name=progressname begin
        # Obtain the initial sample and state.
        sample, state = step(rng, model, sampler; kwargs...)

        # Discard initial samples.
        for _ in 2:discard_initial
            # Obtain the next sample and state.
            sample, state = step(rng, model, sampler, state; kwargs...)
        end

        # Run callback.
        callback === nothing || callback(rng, model, sampler, sample, state, 1; kwargs...)

        # Save the sample.
        samples = AbstractMCMC.samples(sample, model, sampler; kwargs...)
        samples = save!!(samples, sample, 1, model, sampler; kwargs...)

        # Step through the sampler until stopping.
        i = 2

        while !isdone(rng, model, sampler, samples, state, i; progress=progress, kwargs...)
            # Discard thinned samples.
            for _ in 1:(thinning - 1)
                # Obtain the next sample and state.
                sample, state = step(rng, model, sampler, state; kwargs...)
            end

            # Obtain the next sample and state.
            sample, state = step(rng, model, sampler, state; kwargs...)

            # Run callback.
            callback === nothing || callback(rng, model, sampler, sample, state, i; kwargs...)

            # Save the sample.
            samples = save!!(samples, sample, i, model, sampler; kwargs...)

            # Increment iteration counter.
            i += 1
        end
    end

    # Get the sample stop time.
    stop = time()
    duration = stop - start
    stats = SamplingStats(start, stop, duration)

    # Wrap the samples up.
    return bundle_samples(
        samples,
        model,
        sampler,
        state,
        chain_type;
        stats=stats,
        discard_initial=discard_initial,
        thinning=thinning,
        kwargs...
    )
end

function mcmcsample(
    rng::Random.AbstractRNG,
    model::AbstractModel,
    sampler::AbstractSampler,
    ::MCMCThreads,
    N::Integer,
    nchains::Integer;
    progress = PROGRESS[],
    progressname = "Sampling ($(min(nchains, Threads.nthreads())) threads)",
    kwargs...
)
    # Check if actually multiple threads are used.
    if Threads.nthreads() == 1
        @warn "Only a single thread available: MCMC chains are not sampled in parallel"
    end

    # Check if the number of chains is larger than the number of samples
    if nchains > N
        @warn "Number of chains ($nchains) is greater than number of samples per chain ($N)"
    end

    # Copy the random number generator, model, and sample for each thread
    # NOTE: As of May 17, 2020, this relies on Julia's thread scheduling functionality
    #       that distributes a for loop into equal-sized blocks and allocates them
    #       to each thread. If this changes, we may need to rethink things here.
    interval = 1:min(nchains, Threads.nthreads())
    rngs = [deepcopy(rng) for _ in interval]
    models = [deepcopy(model) for _ in interval]
    samplers = [deepcopy(sampler) for _ in interval]

    # Create a seed for each chain using the provided random number generator.
    seeds = rand(rng, UInt, nchains)

    # Set up a chains vector.
    chains = Vector{Any}(undef, nchains)

    @ifwithprogresslogger progress name=progressname begin
        # Create a channel for progress logging.
        if progress
            channel = Channel{Bool}(length(interval))
        end

        Distributed.@sync begin
            if progress
                # Update the progress bar.
                Distributed.@async begin
                    # Determine threshold values for progress logging
                    # (one update per 0.5% of progress)
                    threshold = nchains ÷ 200
                    nextprogresschains = threshold

                    progresschains = 0
                    while take!(channel)
                        progresschains += 1
                        if progresschains >= nextprogresschains
                            ProgressLogging.@logprogress progresschains/nchains
                            nextprogresschains = progresschains + threshold
                        end
                    end
                end
            end

            Distributed.@async begin
                try
                    Threads.@threads for i in 1:nchains
                        # Obtain the ID of the current thread.
                        id = Threads.threadid()

                        # Seed the thread-specific random number generator with the pre-made seed.
                        subrng = rngs[id]
                        Random.seed!(subrng, seeds[i])

                        # Sample a chain and save it to the vector.
                        chains[i] = StatsBase.sample(subrng, models[id], samplers[id], N;
                                                     progress = false, kwargs...)

                        # Update the progress bar.
                        progress && put!(channel, true)
                    end
                finally
                    # Stop updating the progress bar.
                    progress && put!(channel, false)
                end
            end
        end
    end

    # Concatenate the chains together.
    return chainsstack(tighten_eltype(chains))
end

function mcmcsample(
    rng::Random.AbstractRNG,
    model::AbstractModel,
    sampler::AbstractSampler,
    ::MCMCDistributed,
    N::Integer,
    nchains::Integer;
    progress = PROGRESS[],
    progressname = "Sampling ($(Distributed.nworkers()) processes)",
    kwargs...
)
    # Check if actually multiple processes are used.
    if Distributed.nworkers() == 1
        @warn "Only a single process available: MCMC chains are not sampled in parallel"
    end

    # Check if the number of chains is larger than the number of samples
    if nchains > N
        @warn "Number of chains ($nchains) is greater than number of samples per chain ($N)"
    end

    # Create a seed for each chain using the provided random number generator.
    seeds = rand(rng, UInt, nchains)

    # Set up worker pool.
    pool = Distributed.CachingPool(Distributed.workers())

    local chains
    @ifwithprogresslogger progress name=progressname begin
        # Create a channel for progress logging.
        if progress
            channel = Distributed.RemoteChannel(() -> Channel{Bool}(Distributed.nworkers()))
        end

        Distributed.@sync begin
            if progress
                # Update the progress bar.
                Distributed.@async begin
                    # Determine threshold values for progress logging
                    # (one update per 0.5% of progress)
                    threshold = nchains ÷ 200
                    nextprogresschains = threshold

                    progresschains = 0
                    while take!(channel)
                        progresschains += 1
                        if progresschains >= nextprogresschains
                            ProgressLogging.@logprogress progresschains/nchains
                            nextprogresschains = progresschains + threshold
                        end
                    end
                end
            end

            Distributed.@async begin
                try
                    chains = Distributed.pmap(pool, seeds) do seed
                        # Seed a new random number generator with the pre-made seed.
                        Random.seed!(rng, seed)

                        # Sample a chain.
                        chain = StatsBase.sample(rng, model, sampler, N;
                                                 progress = false, kwargs...)

                        # Update the progress bar.
                        progress && put!(channel, true)

                        # Return the new chain.
                        return chain
                    end
                finally
                    # Stop updating the progress bar.
                    progress && put!(channel, false)
                end
            end
        end
    end

    # Concatenate the chains together.
    return chainsstack(tighten_eltype(chains))
end

function mcmcsample(
    rng::Random.AbstractRNG,
    model::AbstractModel,
    sampler::AbstractSampler,
    ::MCMCSerial,
    N::Integer,
    nchains::Integer;
    progressname = "Sampling",
    kwargs...
)
    # Check if the number of chains is larger than the number of samples
    if nchains > N
        @warn "Number of chains ($nchains) is greater than number of samples per chain ($N)"
    end

    # Sample the chains.
    chains = map(
        i -> StatsBase.sample(rng, model, sampler, N; progressname = string(progressname, " (Chain ", i, " of ", nchains, ")"),
        kwargs...),
        1:nchains
    )

    # Concatenate the chains together.
    return chainsstack(tighten_eltype(chains))
end

tighten_eltype(x) = x
tighten_eltype(x::Vector{Any}) = map(identity, x)


using JuLIP.Chemistry: atomic_number
using JuLIP: Atoms
using StaticArrays: SVector

abstract type MSitePotential <: AbstractCalculator end

abstract type MPairPotential <: MSitePotential end

# Experimental Multi-component codes

# ----------------------------------------------------------------------
#    Managing a list of species
# ----------------------------------------------------------------------

abstract type AbstractZList end

"""
`ZList` and `SZList{NZ}` : simple data structures that store a list
of species and convert between atomic numbers and the index in the list.
Can be constructed via
* `ZList(zors)` : where `zors` is an Integer  or `Symbol` (single species)
* `ZList(zs1, zs2, ..., zsn)`
* `ZList([sz1, zs2, ..., zsn])`
* All of these take a kwarg `static = {true, false}`; if `true`, then `ZList`
will return a `SZList{NZ}` for (possibly) faster access.
"""
struct ZList <: AbstractZList
   list::Vector{Int16}
end

Base.length(zlist::AbstractZList) = length(zlist.list)

struct SZList{N} <: AbstractZList
   list::SVector{N, Int16}
end

ZList(zlist::AbstractVector{<: Integer}; static = false) = (
   static ? SZList(SVector( (Int16.(sort(zlist)))... ))
          :  ZList( convert(Vector{Int16}, sort(zlist)) ))

ZList(s::Symbol; kwargs...) =
      ZList( [ atomic_number(s) ]; kwargs... )

ZList(S::AbstractVector{Symbol}; kwargs...) =
      ZList( atomic_number.(S); kwargs... )

ZList(args...; kwargs) =
      ZList( [args...]; kwargs...)

i2z(Zs::AbstractZList, i::Integer) = Zs.list[i]

function z2i(Zs::AbstractZList, z::Integer)
   for j = 1:length(Zs.list)
      if Zs.list[j] == z
         return j
      end
   end
   error("z = $z not found in ZList $(Zs.list)")
end

i2z(V, i::Integer) = i2z(V.zlist, i)
z2i(V, z::Integer) = z2i(V.zlist, z)

Dict(zlist::ZList) = Dict("__id__" => "JuLIP_ZList",
                          "list" => zlist.list)
Base.convert(::Val{:JuLIP_ZList}, D::Dict) = ZList(D)
ZList(D::Dict) = ZList(D["list"])

Dict(zlist::SZList) = Dict("__id__" => "JuLIP_SZList",
                          "list" => zlist.list)
Base.convert(::Val{:JuLIP_SZList}, D::Dict) = SZList(D)
SZList(D::Dict) = SZList(SVector(Int16.(D["list"])...))

# ----------------------------------------------------------------------
#       Basic calculator extensions for multi-species
# ----------------------------------------------------------------------

"""
`neigsz!(tmp, nlist::PairList, at::Atoms, i::Integer) -> j, R Z`

requires a temporary storage array `tmp` with fields
`tmp.R, tmp.Z`.
"""
function neigsz!(tmp, nlist::PairList, at::Atoms, i::Integer)
   j, R = neigs!(tmp.R, nlist, i)
   Z = tmp.Z
   for n = 1:length(j)
      Z[n] = at.Z[j[n]]
   end
   return j, R, (@view Z[1:length(j)])
end

function neigsz(nlist::PairList, at::Atoms, i::Integer)
   j, R = NeighbourLists.neigs(nlist, i)
   return j, R, at.Z[j]
end

alloc_temp(V::MSitePotential, at::AbstractAtoms) =
      alloc_temp(V, maxneigs(neighbourlist(at, cutoff(V))))

alloc_temp(V::MSitePotential, N::Integer) =
      ( R = zeros(JVecF, N),
        Z = zeros(Int16, N), )

alloc_temp_d(V::MSitePotential, at::AbstractAtoms) =
      alloc_temp_d(V, maxneigs(neighbourlist(at, cutoff(V))))

alloc_temp_d(V::MSitePotential, N::Integer) =
      (dV = zeros(JVecF, N),
        R = zeros(JVecF, N),
        Z = zeros(Int16, N), )

energy(V::MSitePotential, at::AbstractAtoms; kwargs...) =
      energy!(alloc_temp(V, at), V, at; kwargs...)

virial(V::MSitePotential, at::AbstractAtoms; kwargs...) =
      virial!(alloc_temp_d(V, at), V, at; kwargs...)

forces(V::MSitePotential, at::AbstractAtoms{T}; kwargs...) where {T} =
      forces!(zeros(JVec{T}, length(at)), alloc_temp_d(V, at), V, at; kwargs...)

function energy!(tmp, calc::MSitePotential, at::Atoms{T};
                 domain=1:length(at)) where {T}
   E = zero(T)
   nlist = neighbourlist(at, cutoff(calc))
   for i in domain
      j, R, Z = neigsz!(tmp, nlist, at, i)
      E += evaluate!(tmp, calc, R, Z, Int16(at.Z[i]))
   end
   return E
end

function forces!(frc, tmp, calc::MSitePotential, at::Atoms{T};
                 domain=1:length(at), reset=true) where {T}
   if reset; fill!(frc, zero(JVec{T})); end
   nlist = neighbourlist(at, cutoff(calc))
   for i in domain
      j, R, Z = neigsz!(tmp, nlist, at, i)
      if length(j) > 0
         evaluate_d!(tmp.dV, tmp, calc, R, Z, Int16(at.Z[i]))
         for a = 1:length(j)
            frc[j[a]] -= tmp.dV[a]
            frc[i]    += tmp.dV[a]
         end
      end
   end
   return frc
end

function virial!(tmp, calc::MSitePotential, at::Atoms{T};
                 domain=1:length(at)) where {T}
   nlist = neighbourlist(at, cutoff(calc))
   vir = zero(JMat{T})
   for i in domain
      j, R, Z = neigsz!(tmp, nlist, at, i)
      if length(j) > 0
         evaluate_d!(tmp.dV, tmp, calc, R, Z, Int16(at.Z[i]))
         vir += site_virial(tmp.dV, R)
      end
   end
   return vir
end

evaluate(V::MSitePotential, R, Z, z) =
      evaluate!(alloc_temp(V, length(R)), V, R, Z, z)

evaluate_d(V::MSitePotential, R::AbstractVector{JVec{T}}, Z, z) where {T} =
      evaluate_d!(zeros(JVec{T}, length(R)),
                  alloc_temp_d(V, length(R)),
                  V, R, Z, z)

site_energy(V::MSitePotential, at::AbstractAtoms, i0::Integer) =
      energy(V, at; domain = (i0,))

site_energy_d(V::MSitePotential, at::AbstractAtoms, i0::Integer) =
      rmul!(forces(V, at; domain = (i0,)), -one(eltype(at)))

# -------- Dispatch for MPairPotentials ------------------------------------

evaluate!(tmp, V::MPairPotential,
          R::AbstractVector{<: JVec{T}}, Z::AbstractVector, z0) where {T} =
   length(R) == 0 ? zero(T) : T(0.5) * sum( evaluate!(tmp, V, norm(R[i]), Z[i], z0)
                                            for i = 1:length(R) )

function evaluate_d!(dV, tmp, V::MPairPotential, R, Z, z0)
   dV = tmp.dV
   for i = 1:length(R)
      r = norm(R[i])
      dV[i] = ((evaluate_d!(tmp, V, r, Z[i], z0)/r) * R[i]) / 2
   end
   return dV
end

evaluate(V::MPairPotential, r, z1, z0) =
      evaluate!(alloc_temp(V, 1), V, r, z1, z0)

evaluate_d(V::MPairPotential, r, z1, z0) =
      evaluate_d!(alloc_temp_d(V, 1), V, r, z1, z0)

# ------------------ Preconditioners for MPairPotentials ----------------------

function precon!(hEs, tmp, V::MPairPotential, R::AbstractVector{<: JVec},
                 Z::AbstractVector, z0, innerstab=T(0.0))
   n = length(R)
   for i = 1:n
      hEs[i,i] = precon!(tmp, V, norm(R[i]), R[i], Z[i], z0, innerstab)
   end
   return hEs
end

# # an FF preconditioner for pair potentials
# function precon!(tmp, V::PairPotential, r::T, R::JVec{T}, innerstab=T(0.1)
#                  ) where {T <: Number}
#    r = norm(R)
#    dV = evaluate_d!(tmp, V, r)
#    ddV = evaluate_dd!(tmp, V, r)
#    R̂ = R/r
#    return (1-innerstab) * (abs(ddV) * R̂ * R̂' + abs(dV / r) * (I - R̂ * R̂')) +
#              innerstab  * (abs(ddV) + abs(dV / r)) * I
# end

# struct MPairPotential <: MSitePotential
#    Z2V::Dict{Tuple{Int16,Int16}, WrappedPairPotential}
# end
#
# MPairPotential() =
#       MPairPotential(Dict{(Int16, Int16), WrappedPairPotential}())
#
# function MPairPotential(args...)
#    Vm = MPairPotential()
#    for p in args
#       @assert p isa Pair
#       z, V = _convert_multi_pp(sym::Symbol, p)
#       Vm.Z2V[z] = V
#    end
#    return Vm
# end
#
# F64fun(V::F64fun) = V
# _convert_multi_pp(sym::Tuple{Symbol, Symbol}, V) =
#       _convert_multi_pp(atomic_number.(sym), V)
# _convert_multi_pp(z::Tuple{<:Integer,<:Integer}, V) =
#       Int16.(z), V
#
# cutoff(V::MPairPotential) =
#    maximum( cutoff(V) for (key, V) in V.Z2V )


"""
    Code for convenience functions for calculating CCFs for each chunk of an AbstractChunkListTimeseries (see RvSpectMLBase)
Author: Eric Ford
Created: August 2020
"""

"""  `calc_order_ccf_chunklist_timeseries( chunklist_timeseries, ccf_plan )`
Convenience function to compute separate CCFs for each chunk (potentially an order or view around one or two lines) of each spectrum in a timeseries.
    CCF is evaluated using line list and mask_shape provided by the ccf plan for each chunk.
Uses multiple threads if avaliable.
# Inputs:
- chunklist_timeseries
# Optional Arguments:
- ccf_plan (BasicCCFPlan())
# Return:
A 3-d array containing the CCF at each (velocity, chunk, spectrum)
Note that the ccf_plan provided is used as a template for creating a custom ccf_plan for each chunk that
    only includes lines that reliably appear in that order for all spectra in the chunklist_timeseries.
"""
function calc_order_ccf_chunklist_timeseries(clt::AbstractChunkListTimeseries,
                plan::PlanT = BasicCCFPlan(); verbose::Bool = false ,
                Δfwhm::AbstractVector{T} = zeros(0)
                ) where { PlanT<:AbstractCCFPlan, T<:Real  }

  #@assert issorted( plan.line_list.λ )
  #nvs = length(calc_ccf_v_grid(plan))
  nvs = calc_length_ccf_v_grid(plan)
  norders = length(clt.chunk_list[1].data)
  nobs =  length(clt.chunk_list)
  order_ccfs = zeros(nvs, norders, nobs)
  num_lines = length(plan.line_list.λ)
  plan_for_chunk = Vector{BasicCCFPlan}(undef,num_chunks(clt))
  for chid in 1:num_chunks(clt)
      order = clt[1].order[chid]
      if typeof(plan.line_list) <: BasicLineList2D
          start_order_idx = searchsortedfirst(plan.line_list.order,order)
          stop_order_idx = searchsortedlast(view(plan.line_list.order,start_order_idx:num_lines),order) + start_order_idx-1
      else
          start_order_idx = 1
          stop_order_idx = num_lines
      end
      if (1<=start_order_idx<=stop_order_idx<=num_lines)
          # find the maximum lower wavelength for the chunk, and the minumum upper wavelength, over all observations
          λmin = maximum(map(obsid->first(clt.chunk_list[obsid].data[chid].λ), 1:length(clt) ))
          λmax = minimum(map(obsid->last( clt.chunk_list[obsid].data[chid].λ), 1:length(clt) ))
          # extend λmin/λmax by the velocity range over which we don't want the mask to change
          λmin  = λmin / (calc_doppler_factor(plan.v_center)*calc_doppler_factor(-plan.v_range_no_mask_change))
          λmax  = λmax / (calc_doppler_factor(plan.v_center)*calc_doppler_factor(plan.v_range_no_mask_change))
          # extend λmin/λmax by the mask width
          upper_edge_of_mask_for_line_at_λmin = λ_max(plan.mask_shape,λmin)
          lower_edge_of_mask_for_line_at_λmax = λ_min(plan.mask_shape,λmax)
          # find the first and last mask entries to use in each chunk
          start_line_idx = searchsortedfirst(view(plan.line_list.λ,start_order_idx:stop_order_idx),upper_edge_of_mask_for_line_at_λmin) + start_order_idx-1
          #stop_line_idx  = num_lines+1 - searchsortedfirst(view(plan.line_list.λ,num_lines:-1:1),lower_edge_of_mask_for_line_at_λmax,rev=true)
          stop_line_idx  = start_order_idx-1 + searchsortedlast(view(plan.line_list.λ,start_order_idx:stop_order_idx),lower_edge_of_mask_for_line_at_λmax)
          if verbose
              flush(stdout)
              println("extrema(plan.line_list.λ) = ",extrema(plan.line_list.λ) )
              if (1 <= start_line_idx <= length(plan.line_list.λ)) && (1 <= stop_line_idx <= length(plan.line_list.λ))
                  println("start_order_idx = ", start_order_idx, "  stop_order_idx = ", stop_order_idx)
                  println("start_line_idx = ", start_line_idx, " λ= ", plan.line_list.λ[start_line_idx]) #, " order= ", plan.line_list.order[start_line_idx])
                  println("stop_line_idx = ", stop_line_idx, " λ= ", plan.line_list.λ[stop_line_idx]) #, " order= ", plan.line_list.order[stop_line_idx])
              end
              println("upper_edge_of_mask_for_line_at_λmin = ", upper_edge_of_mask_for_line_at_λmin, " Lower_edge_of_mask_for_line_at_λmax = ", lower_edge_of_mask_for_line_at_λmax)
          end

          if (1 <= start_line_idx <= length(plan.line_list.λ)) && (1 <= stop_line_idx <= length(plan.line_list.λ))
              if verbose && typeof(plan.line_list) <: BasicLineList2D
                  println("start_line_idx = ", start_line_idx, " λ= ", plan.line_list.λ[start_line_idx], " order= ", plan.line_list.order[start_line_idx])
                  println("stop_line_idx = ", stop_line_idx, " λ= ", plan.line_list.λ[stop_line_idx], " order= ", plan.line_list.order[stop_line_idx])
                  println(" Using ", length(start_line_idx:stop_line_idx), " lines for chunk ", chid, " order ", order)
              end
              if ! issorted( view(plan.line_list.λ, start_line_idx:stop_line_idx ) )
                  println("# View is not sorted.  Why?")
                  for i in start_line_idx:stop_line_idx
                      println("i= ", i, " λ= ", plan.line_list.λ[i], " order= ", plan.line_list.order[i])
                  end
              end
              @assert issorted( view(plan.line_list.λ, start_line_idx:stop_line_idx ) )
              # line_list_for_chunk = BasicLineList2D(view(plan.line_list.λ,start_line_idx:stop_line_idx), view(plan.line_list.weight,start_line_idx:stop_line_idx), fill(order,length(start_line_idx:stop_line_idx)) )
              line_list_for_chunk = BasicLineList(view(plan.line_list.λ,start_line_idx:stop_line_idx), view(plan.line_list.weight,start_line_idx:stop_line_idx)  )
          else  # No lines in this chunk!
              line_list_for_chunk = EmptyBasicLineList()
          end
      else
        line_list_for_chunk = EmptyBasicLineList()
      end

      #create a plan for this chunk that only includes the mask entries we want for this chunk
      plan_for_chunk[chid] = BasicCCFPlan( line_list=line_list_for_chunk, midpoint=plan.v_center, step=plan.v_step, max=plan.v_max, mask_shape=plan.mask_shape )
  end
  #@threaded mapreduce(obsid->calc_ccf_chunklist(clt.chunk_list[obsid], plan_for_chunk),hcat, 1:length(clt) )
  #Threads.@threads for i in 1:nobs
    #  order_ccfs[:,:,i] .= calc_order_ccfs_chunklist(clt.chunk_list[i], plan)
  #end
  Threads.@threads for i in 1:nobs
      this_Δfwhm = length(Δfwhm) == nobs ? Δfwhm[i] : 0.0
      order_ccfs[:,:,i] .= calc_order_ccfs_chunklist(clt.chunk_list[i], plan_for_chunk, Δfwhm=this_Δfwhm, assume_sorted=true )
  end

  return order_ccfs
end

"""  `calc_order_ccf_and_var_chunklist_timeseries( chunklist_timeseries, ccf_plan )`
Convenience function to compute separate CCFs for each chunk (potentially an order or view around one or two lines) of each spectrum in a timeseries.
    CCF is evaluated using line list and mask_shape provided by the ccf plan for each chunk.
Uses multiple threads if avaliable.
# Inputs:
- chunklist_timeseries
# Optional Arguments:
- ccf_plan (BasicCCFPlan())
# Return:
A 3-d array containing the CCF at each (velocity, chunk, spectrum)
Note that the ccf_plan provided is used as a template for creating a custom ccf_plan for each chunk that
    only includes lines that reliably appear in that order for all spectra in the chunklist_timeseries.
"""
function calc_order_ccf_and_var_chunklist_timeseries(clt::AbstractChunkListTimeseries,
                plan::PlanT = BasicCCFPlan(); verbose::Bool = false,
                ccf_var_scale::Real = 1.0, Δfwhm::AbstractVector{T} = zeros(0)
                 ) where {
                    PlanT<:AbstractCCFPlan, T<:Real }
  #@assert issorted( plan.line_list.λ )
  #nvs = length(calc_ccf_v_grid(plan))
  nvs = calc_length_ccf_v_grid(plan)
  norders = length(clt.chunk_list[1].data)
  nobs =  length(clt.chunk_list)
  order_ccfs = zeros(nvs, norders, nobs)
  order_ccf_vars = zeros(nvs, norders, nobs)
  num_lines = length(plan.line_list.λ)
  plan_for_chunk = Vector{BasicCCFPlan}(undef,num_chunks(clt))
  for chid in 1:num_chunks(clt)
      order = clt[1].order[chid]
      if typeof(plan.line_list) <: BasicLineList2D
          start_order_idx = searchsortedfirst(plan.line_list.order,order)
          stop_order_idx = searchsortedlast(view(plan.line_list.order,start_order_idx:num_lines),order) + start_order_idx-1
      else
          start_order_idx = 1
          stop_order_idx = num_lines
      end
      if (1<=start_order_idx<=stop_order_idx<=num_lines)
          # find the maximum lower wavelength for the chunk, and the minumum upper wavelength, over all observations
          λmin = maximum(map(obsid->first(clt.chunk_list[obsid].data[chid].λ), 1:length(clt) ))
          λmax = minimum(map(obsid->last( clt.chunk_list[obsid].data[chid].λ), 1:length(clt) ))
          # extend λmin/λmax by the velocity range over which we don't want the mask to change
          λmin  = λmin / (calc_doppler_factor(plan.v_center)*calc_doppler_factor(-plan.v_range_no_mask_change))
          λmax  = λmax / (calc_doppler_factor(plan.v_center)*calc_doppler_factor(plan.v_range_no_mask_change))
          # extend λmin/λmax by the mask width
          upper_edge_of_mask_for_line_at_λmin = λ_max(plan.mask_shape,λmin)
          lower_edge_of_mask_for_line_at_λmax = λ_min(plan.mask_shape,λmax)
          # find the first and last mask entries to use in each chunk
          start_line_idx = searchsortedfirst(view(plan.line_list.λ,start_order_idx:stop_order_idx),upper_edge_of_mask_for_line_at_λmin) + start_order_idx-1
          stop_line_idx  = start_order_idx-1 + searchsortedlast(view(plan.line_list.λ,start_order_idx:stop_order_idx),lower_edge_of_mask_for_line_at_λmax)
          if verbose
              flush(stdout)
              println("extrema(plan.line_list.λ) = ",extrema(plan.line_list.λ) )
              if (1 <= start_line_idx <= length(plan.line_list.λ)) && (1 <= stop_line_idx <= length(plan.line_list.λ))
                  println("start_order_idx = ", start_order_idx, "  stop_order_idx = ", stop_order_idx)
                  println("start_line_idx = ", start_line_idx, " λ= ", plan.line_list.λ[start_line_idx]) #, " order= ", plan.line_list.order[start_line_idx])
                  println("stop_line_idx = ", stop_line_idx, " λ= ", plan.line_list.λ[stop_line_idx]) #, " order= ", plan.line_list.order[stop_line_idx])
              end
              println("upper_edge_of_mask_for_line_at_λmin = ", upper_edge_of_mask_for_line_at_λmin, " Lower_edge_of_mask_for_line_at_λmax = ", lower_edge_of_mask_for_line_at_λmax)
          end

          if (1 <= start_line_idx <= length(plan.line_list.λ)) && (1 <= stop_line_idx <= length(plan.line_list.λ))
              if verbose && typeof(plan.line_list) <: BasicLineList2D
                  println("start_line_idx = ", start_line_idx, " λ= ", plan.line_list.λ[start_line_idx], " order= ", plan.line_list.order[start_line_idx])
                  println("stop_line_idx = ", stop_line_idx, " λ= ", plan.line_list.λ[stop_line_idx], " order= ", plan.line_list.order[stop_line_idx])
                  println(" Using ", length(start_line_idx:stop_line_idx), " lines for chunk ", chid, " order ", order)
              end
              if ! issorted( view(plan.line_list.λ, start_line_idx:stop_line_idx ) )
                  println("# View is not sorted.  Why?")
                  for i in start_line_idx:stop_line_idx
                      println("i= ", i, " λ= ", plan.line_list.λ[i], " order= ", plan.line_list.order[i])
                  end
              end
              @assert issorted( view(plan.line_list.λ, start_line_idx:stop_line_idx ) )
              # line_list_for_chunk = BasicLineList(view(plan.line_list.λ,start_line_idx:stop_line_idx), view(plan.line_list.weight,start_line_idx:stop_line_idx) )
              line_list_for_chunk = BasicLineList(view(plan.line_list.λ,start_line_idx:stop_line_idx), view(plan.line_list.weight,start_line_idx:stop_line_idx)  )
          else  # No lines in this chunk!
              line_list_for_chunk = EmptyBasicLineList()
          end
      else
        line_list_for_chunk = EmptyBasicLineList()
      end
      #create a plan for this chunk that only includes the mask entries we want for this chunk
      plan_for_chunk[chid] = BasicCCFPlan( line_list=line_list_for_chunk, midpoint=plan.v_center, step=plan.v_step, max=plan.v_max, mask_shape=plan.mask_shape )
  end
  #@threaded mapreduce(obsid->calc_ccf_chunklist(clt.chunk_list[obsid], plan_for_chunk),hcat, 1:length(clt) )
  #Threads.@threads for i in 1:nobs
    #  order_ccfs[:,:,i] .= calc_order_ccfs_chunklist(clt.chunk_list[i], plan)
  #end
  Threads.@threads for i in 1:nobs
      this_Δfwhm = length(Δfwhm) == nobs ? Δfwhm[i] : 0.0
      for j in 1:num_chunks(clt)
          ( ccf_tmp, ccf_var_tmp ) = calc_ccf_and_var_chunk(clt.chunk_list[i][j], plan_for_chunk[j],
                                            ccf_var_scale=ccf_var_scale, Δfwhm=this_Δfwhm, assume_sorted=true )
          order_ccfs[:,j,i] .= ccf_tmp
          order_ccf_vars[:,j,i] .= ccf_var_tmp
      end
  end

  return (ccf=order_ccfs, ccf_var=order_ccf_vars)
end