module Inference

using ..Core
using ..Core: logZ
using ..Utilities
using DynamicPPL: Metadata, _tail, VarInfo, TypedVarInfo,
    islinked, invlink!, getlogp, tonamedtuple, VarName, getsym, vectorize,
    settrans!, _getvns, getdist, CACHERESET, AbstractSampler,
    Model, Sampler, SampleFromPrior, SampleFromUniform,
    Selector, AbstractSamplerState, DefaultContext, PriorContext,
    LikelihoodContext, MiniBatchContext, set_flag!, unset_flag!, NamedDist, NoDist,
    getspace, inspace
using Distributions, Libtask, Bijectors
using DistributionsAD: VectorOfMultivariate
using LinearAlgebra
using ..Turing: PROGRESS, Turing
using StatsFuns: logsumexp
using Random: GLOBAL_RNG, AbstractRNG, randexp
using DynamicPPL
using AbstractMCMC: AbstractModel, AbstractSampler
using DocStringExtensions: TYPEDEF, TYPEDFIELDS

import AbstractMCMC
import AdvancedHMC; const AHMC = AdvancedHMC
import AdvancedMH; const AMH = AdvancedMH
import ..Core: getchunksize, getADbackend
import DynamicPPL: get_matching_type,
    VarName, _getranges, _getindex, getval, _getvns
import EllipticalSliceSampling
import Random
import MCMCChains
import StatsBase: predict

export  InferenceAlgorithm,
        Hamiltonian,
        GibbsComponent,
        StaticHamiltonian,
        AdaptiveHamiltonian,
        SampleFromUniform,
        SampleFromPrior,
        MH,
        ESS,
        Emcee,
        Gibbs,      # classic sampling
        HMC,
        SGLD,
        SGHMC,
        HMCDA,
        NUTS,       # Hamiltonian-like sampling
        DynamicNUTS,
        IS,
        SMC,
        CSMC,
        PG,
        Prior,
        assume,
        dot_assume,
        observe,
        dot_observe,
        resume,
        predict,
        isgibbscomponent

#######################
# Sampler abstraction #
#######################
abstract type AbstractAdapter end
abstract type InferenceAlgorithm end
abstract type ParticleInference <: InferenceAlgorithm end
abstract type Hamiltonian{AD} <: InferenceAlgorithm end
abstract type StaticHamiltonian{AD} <: Hamiltonian{AD} end
abstract type AdaptiveHamiltonian{AD} <: Hamiltonian{AD} end

getchunksize(::Type{<:Hamiltonian{AD}}) where AD = getchunksize(AD)
getADbackend(::Hamiltonian{AD}) where AD = AD()

# Algorithm for sampling from the prior
struct Prior <: InferenceAlgorithm end

"""
    mh_accept(logp_current::Real, logp_proposal::Real, log_proposal_ratio::Real)

Decide if a proposal ``x'`` with log probability ``\\log p(x') = logp_proposal`` and
log proposal ratio ``\\log k(x', x) - \\log k(x, x') = log_proposal_ratio`` in a
Metropolis-Hastings algorithm with Markov kernel ``k(x_t, x_{t+1})`` and current state
``x`` with log probability ``\\log p(x) = logp_current`` is accepted by evaluating the
Metropolis-Hastings acceptance criterion
```math
\\log U \\leq \\log p(x') - \\log p(x) + \\log k(x', x) - \\log k(x, x')
```
for a uniform random number ``U \\in [0, 1)``.
"""
function mh_accept(logp_current::Real, logp_proposal::Real, log_proposal_ratio::Real)
    # replacing log(rand()) with -randexp() yields test errors
    return log(rand()) + logp_current ≤ logp_proposal + log_proposal_ratio
end

######################
# Default Transition #
######################

struct Transition{T, F<:AbstractFloat}
    θ  :: T
    lp :: F
end

function Transition(spl::Sampler, nt::NamedTuple=NamedTuple())
    theta = merge(tonamedtuple(spl.state.vi), nt)
    lp = getlogp(spl.state.vi)
    return Transition{typeof(theta), typeof(lp)}(theta, lp)
end

function additional_parameters(::Type{<:Transition})
    return [:lp]
end

DynamicPPL.getlogp(t::Transition) = t.lp

##########################################
# Internal variable names for MCMCChains #
##########################################

const TURING_INTERNAL_VARS = (internals = [
    "elapsed",
    "eval_num",
    "lf_eps",
    "lp",
    "weight",
    "le",
    "acceptance_rate",
    "hamiltonian_energy",
    "hamiltonian_energy_error",
    "max_hamiltonian_energy_error",
    "is_accept",
    "log_density",
    "n_steps",
    "numerical_error",
    "step_size",
    "nom_step_size",
    "tree_depth",
    "is_adapt"
],)

#########################################
# Default definitions for the interface #
#########################################

function AbstractMCMC.sample(
    model::AbstractModel,
    alg::InferenceAlgorithm,
    N::Integer;
    kwargs...
)
    return AbstractMCMC.sample(Random.GLOBAL_RNG, model, alg, N; kwargs...)
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    alg::InferenceAlgorithm,
    N::Integer;
    kwargs...
)
    return AbstractMCMC.sample(rng, model, Sampler(alg, model), N; kwargs...)
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    sampler::Sampler{<:InferenceAlgorithm},
    N::Integer;
    chain_type=MCMCChains.Chains,
    resume_from=nothing,
    progress=PROGRESS[],
    kwargs...
)
    if resume_from === nothing
        return AbstractMCMC.mcmcsample(rng, model, sampler, N;
                                       chain_type=chain_type, progress=progress, kwargs...)
    else
        return resume(resume_from, N; chain_type=chain_type, progress=progress, kwargs...)
    end
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    alg::Prior,
    N::Integer;
    chain_type=MCMCChains.Chains,
    resume_from=nothing,
    progress=PROGRESS[],
    kwargs...
)
    if resume_from === nothing
        return AbstractMCMC.mcmcsample(rng, model, SampleFromPrior(), N;
                                       chain_type=chain_type, progress=progress, kwargs...)
    else
        return resume(resume_from, N; chain_type=chain_type, progress=progress, kwargs...)
    end
end

function AbstractMCMC.sample(
    model::AbstractModel,
    alg::InferenceAlgorithm,
    parallel::AbstractMCMC.AbstractMCMCParallel,
    N::Integer,
    n_chains::Integer;
    kwargs...
)
    return AbstractMCMC.sample(Random.GLOBAL_RNG, model, alg, parallel, N, n_chains;
                               kwargs...)
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    alg::InferenceAlgorithm,
    parallel::AbstractMCMC.AbstractMCMCParallel,
    N::Integer,
    n_chains::Integer;
    kwargs...
)
    return AbstractMCMC.sample(rng, model, Sampler(alg, model), parallel, N, n_chains;
                               kwargs...)
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    sampler::Sampler{<:InferenceAlgorithm},
    parallel::AbstractMCMC.AbstractMCMCParallel,
    N::Integer,
    n_chains::Integer;
    chain_type=MCMCChains.Chains,
    progress=PROGRESS[],
    kwargs...
)
    return AbstractMCMC.mcmcsample(rng, model, sampler, parallel, N, n_chains;
                                   chain_type=chain_type, progress=progress, kwargs...)
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    alg::Prior,
    parallel::AbstractMCMC.AbstractMCMCParallel,
    N::Integer,
    n_chains::Integer;
    chain_type=MCMCChains.Chains,
    progress=PROGRESS[],
    kwargs...
)
    return AbstractMCMC.sample(rng, model, SampleFromPrior(), parallel, N, n_chains;
                               chain_type=chain_type, progress=progress, kwargs...)
end

function AbstractMCMC.sample_init!(
    ::AbstractRNG,
    model::AbstractModel,
    spl::Sampler{<:InferenceAlgorithm},
    N::Integer;
    kwargs...
)
    # Resume the sampler.
    set_resume!(spl; kwargs...)

    # Set the parameters to a starting value.
    initialize_parameters!(spl; kwargs...)
end

function initialize_parameters!(
    spl::Sampler;
    init_theta::Union{Nothing,Vector}=nothing,
    verbose::Bool=false,
    kwargs...
)
    islinked(spl.state.vi, spl) && invlink!(spl.state.vi, spl)
    # Get `init_theta`
    if init_theta !== nothing
        verbose && @info "Using passed-in initial variable values" init_theta
        # Convert individual numbers to length 1 vector; `ismissing(v)` is needed as `size(missing)` is undefined`
        init_theta = [ismissing(v) || size(v) == () ? [v] : v for v in init_theta]
        # Flatten `init_theta`
        init_theta_flat = foldl(vcat, map(vec, init_theta))
        # Create a mask to inidicate which values are not missing
        theta_mask = map(x -> !ismissing(x), init_theta_flat)
        # Get all values
        theta = spl.state.vi[spl]
        @assert length(theta) == length(init_theta_flat) "Provided initial value doesn't match the dimension of the model"
        # Update those which are provided (i.e. not missing)
        theta[theta_mask] .= init_theta_flat[theta_mask]
        # Update in `vi`
        spl.state.vi[spl] = theta
    end
end

##########################
# Chain making utilities #
##########################

"""
    getparams(t)

Return a named tuple of parameters.
"""
getparams(t) = t.θ
getparams(t::VarInfo) = tonamedtuple(TypedVarInfo(t))

function _params_to_array(ts::Vector)
    names = Vector{Symbol}()
    # Extract the parameter names and values from each transition.
    dicts = map(ts) do t
        nms, vs = flatten_namedtuple(getparams(t))
        for nm in nms
            if !(nm in names)
                push!(names, nm)
            end
        end
        # Convert the names and values to a single dictionary.
        return Dict(nms[j] => vs[j] for j in 1:length(vs))
    end
    # names = collect(names_set)
    vals = [get(dicts[i], key, missing) for i in eachindex(dicts),
        (j, key) in enumerate(names)]

    return names, vals
end

function flatten_namedtuple(nt::NamedTuple)
    names_vals = mapreduce(vcat, keys(nt)) do k
        v = nt[k]
        if length(v) == 1
            return [(Symbol(k), v)]
        else
            return mapreduce(vcat, zip(v[1], v[2])) do (vnval, vn)
                return collect(FlattenIterator(vn, vnval))
            end
        end
    end
    return [vn[1] for vn in names_vals], [vn[2] for vn in names_vals]
end

function get_transition_extras(ts::AbstractVector{<:VarInfo})
    valmat = reshape([getlogp(t) for t in ts], :, 1)
    return [:lp], valmat
end

function get_transition_extras(ts::AbstractVector)
    # Get the extra field names from the sampler state type.
    # This handles things like :lp or :weight.
    extra_params = additional_parameters(eltype(ts))

    # Get the values of the extra parameters.
    local extra_names
    all_vals = []

    # Iterate through each transition.
    for t in ts
        extra_names = Symbol[]
        vals = []

        # Iterate through each of the additional field names
        # in the struct.
        for p in extra_params
            # Check whether the field contains a NamedTuple,
            # in which case we need to iterate through each
            # key/value pair.
            prop = getproperty(t, p)
            if prop isa NamedTuple
                for (k, v) in pairs(prop)
                    push!(extra_names, Symbol(k))
                    push!(vals, v)
                end
            else
                push!(extra_names, Symbol(p))
                push!(vals, prop)
            end
        end
        push!(all_vals, vals)
    end

    # Convert the vector-of-vectors to a matrix.
    valmat = [all_vals[i][j] for i in 1:length(ts), j in 1:length(all_vals[1])]

    return extra_names, valmat
end

getlogevidence(sampler) = missing
function getlogevidence(sampler::Sampler)
    if isdefined(sampler.state, :average_logevidence)
        return sampler.state.average_logevidence
    elseif isdefined(sampler.state, :final_logevidence)
        return sampler.state.final_logevidence
    else
        return missing
    end
end

# Default MCMCChains.Chains constructor.
# This is type piracy (at least for SampleFromPrior).
function AbstractMCMC.bundle_samples(
    rng::AbstractRNG,
    model::AbstractModel,
    spl::Union{Sampler{<:InferenceAlgorithm},SampleFromPrior},
    N::Integer,
    ts::Vector,
    chain_type::Type{MCMCChains.Chains};
    save_state = false,
    kwargs...
)
    # Convert transitions to array format.
    # Also retrieve the variable names.
    nms, vals = _params_to_array(ts)

    # Get the values of the extra parameters in each transition.
    extra_params, extra_values = get_transition_extras(ts)

    # Extract names & construct param array.
    nms = [nms; extra_params]
    parray = hcat(vals, extra_values)

    # Get the average or final log evidence, if it exists.
    le = getlogevidence(spl)

    # Set up the info tuple.
    if save_state
        info = (range = rng, model = model, spl = spl)
    else
        info = NamedTuple()
    end

    # Conretize the array before giving it to MCMCChains.
    parray = MCMCChains.concretize(parray)

    # Chain construction.
    return MCMCChains.Chains(
        parray,
        nms,
        deepcopy(TURING_INTERNAL_VARS);
        evidence=le,
        info=info,
    ) |> sort
end

# This is type piracy (for SampleFromPrior).
function AbstractMCMC.bundle_samples(
    rng::AbstractRNG,
    model::AbstractModel,
    spl::Union{Sampler{<:InferenceAlgorithm},SampleFromPrior},
    N::Integer,
    ts::Vector,
    chain_type::Type{Vector{NamedTuple}};
    discard_adapt::Bool=true,
    save_state=false,
    kwargs...
)
    nts = Vector{NamedTuple}(undef, N)

    for (i, t) in enumerate(ts)
        params = getparams(t)

        k = collect(keys(params))
        vs = []
        for v in values(params)
            push!(vs, v[1])
        end

        push!(k, :lp)

        nts[i] = NamedTuple{tuple(k...)}(tuple(vs..., getlogp(t)))
    end

    return map(identity, nts)
end

function save(c::MCMCChains.Chains, spl::Sampler, model, vi, samples)
    nt = NamedTuple{(:spl, :model, :vi, :samples)}((spl, model, deepcopy(vi), samples))
    return setinfo(c, merge(nt, c.info))
end

function resume(
    c::MCMCChains.Chains,
    n_iter::Int;
    chain_type=MCMCChains.Chains,
    progress=PROGRESS[],
    kwargs...
)
    @assert !isempty(c.info) "[Turing] cannot resume from a chain without state info"

    # Sample a new chain.
    newchain = AbstractMCMC.mcmcsample(
        c.info[:range],
        c.info[:model],
        c.info[:spl],
        n_iter;
        resume_from=c,
        reuse_spl_n=n_iter,
        chain_type=MCMCChains.Chains,
        progress=progress,
        kwargs...
    )

    # Stick the new samples at the end of the old chain.
    return vcat(c, newchain)
end

function set_resume!(
    s::Sampler;
    resume_from::Union{MCMCChains.Chains, Nothing}=nothing,
    kwargs...
)
    # If we're resuming, grab the sampler info.
    if resume_from !== nothing
        s = resume_from.info[:spl]
    end
end

#########################
# Default sampler state #
#########################

"""
A blank `AbstractSamplerState` that contains only `VarInfo` information.
"""
mutable struct SamplerState{VIType<:VarInfo} <: AbstractSamplerState
    vi :: VIType
end

#######################################
# Concrete algorithm implementations. #
#######################################

include("ess.jl")
include("hmc.jl")
include("mh.jl")
include("is.jl")
include("AdvancedSMC.jl")
include("gibbs.jl")
include("../contrib/inference/sghmc.jl")
include("emcee.jl")

################
# Typing tools #
################

for alg in (:SMC, :PG, :MH, :IS, :ESS, :Gibbs, :Emcee)
    @eval DynamicPPL.getspace(::$alg{space}) where {space} = space
end
for alg in (:HMC, :HMCDA, :NUTS, :SGLD, :SGHMC)
    @eval DynamicPPL.getspace(::$alg{<:Any, space}) where {space} = space
end

floatof(::Type{T}) where {T <: Real} = typeof(one(T)/one(T))
floatof(::Type) = Real # fallback if type inference failed

function get_matching_type(
    spl::AbstractSampler,
    vi,
    ::Type{T},
) where {T}
    return T
end
function get_matching_type(
    spl::AbstractSampler,
    vi,
    ::Type{<:Union{Missing, AbstractFloat}},
)
    return Union{Missing, floatof(eltype(vi, spl))}
end
function get_matching_type(
    spl::AbstractSampler,
    vi,
    ::Type{<:AbstractFloat},
)
    return floatof(eltype(vi, spl))
end
function get_matching_type(spl::AbstractSampler, vi, ::Type{<:Array{T,N}}) where {T,N}
    return Array{get_matching_type(spl, vi, T), N}
end
function get_matching_type(spl::AbstractSampler, vi, ::Type{<:Array{T}}) where T
    return Array{get_matching_type(spl, vi, T)}
end
function get_matching_type(
    spl::Sampler{<:Union{PG, SMC}},
    vi,
    ::Type{TV},
) where {T, N, TV <: Array{T, N}}
    return TArray{T, N}
end

##############
# Utilities  #
##############

DynamicPPL.getspace(spl::Sampler) = getspace(spl.alg)
DynamicPPL.inspace(vn::VarName, spl::Sampler) = inspace(vn, getspace(spl.alg))

"""

    predict([rng::AbstractRNG,] model::Model, chain::MCMCChains.Chains; include_all=false)

Execute `model` conditioned on each sample in `chain`, and return the resulting `Chains`.

If `include_all` is `false`, the returned `Chains` will contain only those variables
sampled/not present in `chain`.

# Details
Internally calls `Turing.Inference.transitions_from_chain` to obtained the samples
and then converts these into a `Chains` object using `AbstractMCMC.bundle_samples`.

# Example
```jldoctest
julia> using Turing; setprogress!(false);
[ Info: [Turing]: progress logging is disabled globally

julia> @model function linear_reg(x, y, σ = 0.1)
           β ~ Normal(0, 1)

           for i ∈ eachindex(y)
               y[i] ~ Normal(β * x[i], σ)
           end
       end;

julia> σ = 0.1; f(x) = 2 * x + 0.1 * randn();

julia> Δ = 0.1; xs_train = 0:Δ:10; ys_train = f.(xs_train);

julia> xs_test = [10 + Δ, 10 + 2 * Δ]; ys_test = f.(xs_test);

julia> m_train = linear_reg(xs_train, ys_train, σ);

julia> chain_lin_reg = sample(m_train, NUTS(100, 0.65), 200);
┌ Info: Found initial step size
└   ϵ = 0.003125

julia> m_test = linear_reg(xs_test, Vector{Union{Missing, Float64}}(undef, length(ys_test)), σ);

julia> predictions = predict(m_test, chain_lin_reg)
Object of type Chains, with data of type 100×2×1 Array{Float64,3}

Iterations        = 1:100
Thinning interval = 1
Chains            = 1
Samples per chain = 100
parameters        = y[1], y[2]

2-element Array{ChainDataFrame,1}

Summary Statistics
  parameters     mean     std  naive_se     mcse       ess   r_hat
  ──────────  ───────  ──────  ────────  ───────  ────────  ──────
        y[1]  20.1974  0.1007    0.0101  missing  101.0711  0.9922
        y[2]  20.3867  0.1062    0.0106  missing  101.4889  0.9903

Quantiles
  parameters     2.5%    25.0%    50.0%    75.0%    97.5%
  ──────────  ───────  ───────  ───────  ───────  ───────
        y[1]  20.0342  20.1188  20.2135  20.2588  20.4188
        y[2]  20.1870  20.3178  20.3839  20.4466  20.5895

julia> ys_pred = vec(mean(Array(group(predictions, :y)); dims = 1));

julia> sum(abs2, ys_test - ys_pred) ≤ 0.1
true
```
"""
function predict(model::Model, chain::MCMCChains.Chains; kwargs...)
    return predict(Random.GLOBAL_RNG, model, chain; kwargs...)
end
function predict(rng::AbstractRNG, model::Model, chain::MCMCChains.Chains; include_all = false)
    spl = DynamicPPL.SampleFromPrior()

    # Sample transitions using `spl` conditioned on values in `chain`
    transitions = [
        transitions_from_chain(rng, model, chain[:, :, chn_idx]; sampler = spl)
        for chn_idx = 1:size(chain, 3)
    ]

    # Let the Turing internals handle everything else for you
    chain_result = reduce(
        MCMCChains.chainscat, [
            AbstractMCMC.bundle_samples(
                rng,
                model,
                spl,
                length(chain),
                transitions[chn_idx],
                MCMCChains.Chains
            ) for chn_idx = 1:size(chain, 3)
        ]
    )

    parameter_names = if include_all
        names(chain_result, :parameters)
    else
        filter(k -> ∉(k, names(chain, :parameters)), names(chain_result, :parameters))
    end

    return chain_result[parameter_names]
end

"""

    transitions_from_chain(
        [rng::AbstractRNG,]
        model::Model,
        chain::MCMCChains.Chains;
        sampler = DynamicPPL.SampleFromPrior()
    )

Execute `model` conditioned on each sample in `chain`, and return resulting transitions.

The returned transitions are represented in a `Vector{<:Turing.Inference.Transition}`.

# Details

In a bit more detail, the process is as follows:
1. For every `sample` in `chain`
   1. For every `variable` in `sample`
      1. Set `variable` in `model` to its value in `sample`
   2. Execute `model` with variables fixed as above, sampling variables NOT present
      in `chain` using `SampleFromPrior`
   3. Return sampled variables and log-joint

# Example
```julia-repl
julia> using Turing

julia> @model function demo()
           m ~ Normal(0, 1)
           x ~ Normal(m, 1)
       end;

julia> m = demo();

julia> chain = Chains(randn(2, 1, 1), ["m"]); # 2 samples of `m`

julia> transitions = Turing.Inference.transitions_from_chain(m, chain);

julia> [Turing.Inference.getlogp(t) for t in transitions] # extract the logjoints
2-element Array{Float64,1}:
 -3.6294991938628374
 -2.5697948166987845

julia> [first(t.θ.x) for t in transitions] # extract samples for `x`
2-element Array{Array{Float64,1},1}:
 [-2.0844148956440796]
 [-1.704630494695469]
```
"""
function transitions_from_chain(
    model::Turing.Model,
    chain::MCMCChains.Chains;
    kwargs...
)
    return transitions_from_chain(Random.GLOBAL_RNG, model, chain; kwargs...)
end
function transitions_from_chain(
    rng::AbstractRNG,
    model::Turing.Model,
    chain::MCMCChains.Chains;
    sampler = DynamicPPL.SampleFromPrior()
)
    vi = Turing.VarInfo(model)

    transitions = map(1:length(chain)) do i
        c = chain[i]
        md = vi.metadata
        for v in keys(md)
            for vn in md[v].vns
                vn_sym = Symbol(vn)

                # Cannot use `vn_sym` to index in the chain
                # so we have to extract the corresponding "linear"
                # indices and use those.
                # `ks` is empty if `vn_sym` not in `c`.
                ks = MCMCChains.namesingroup(c, vn_sym)

                if !isempty(ks)
                    # 1st dimension is of size 1 since `c`
                    # only contains a single sample, and the
                    # last dimension is of size 1 since
                    # we're assuming we're working with a single chain.
                    val = copy(vec(c[ks].value))
                    DynamicPPL.setval!(vi, val, vn)
                    DynamicPPL.settrans!(vi, false, vn)
                else
                    DynamicPPL.set_flag!(vi, vn, "del")
                end
            end
        end
        # Execute `model` on the parameters set in `vi` and sample those with `"del"` flag using `sampler`
        model(rng, vi, sampler)

        # Convert `VarInfo` into `NamedTuple` and save
        theta = DynamicPPL.tonamedtuple(vi)
        lp = Turing.getlogp(vi)
        Transition(theta, lp)
    end

    return transitions
end

end # module


using DifferentiableNAS
using Flux
using Flux: throttle, logitcrossentropy, onecold, onehotbatch, Optimiser
using Zygote: @nograd
using StatsBase: mean
using Parameters
using CUDA
using Distributions
using BSON
using Dates

include("../CIFAR10.jl")
include("../training_utils.jl")

@show beginscript = now()

@show checkpointed = "test/models/darts_6642942/model29.bson"
BSON.@load checkpointed m_cpu

@show args = parse_commandline()

if args["random_seed"] > -1
    Random.seed!(args["random_seed"])
end

m = m_cpu |> gpu

#worth saving these?
optimiser_α = Optimiser(WeightDecay(1e-3),ADAM(3e-4,(0.5,0.999)))
optimiser_w = Optimiser(WeightDecay(3e-4),CosineAnnealing(args["epochs"]),Momentum(0.025, 0.9))

train, val = get_processed_data(args["val_split"], args["batchsize"], args["trainval_fraction"], args["random_seed"])
test = get_test_data(args["test_fraction"], args["random_seed"])

histepoch = historiessml()
histbatch = historiessml()
losses = [0.0, 0.0]

base_folder = prepare_folder("darts", args)

cbepoch = CbAll(CUDA.reclaim, histepoch, save_progress, CUDA.reclaim)
cbbatch = CbAll(CUDA.reclaim, histbatch, CUDA.reclaim)

for epoch in 1:args["epochs"]
    @show epoch
    display(Dates.format(convert(DateTime,now()-beginscript), "HH:MM:SS"))
    @time DARTStrain1st!(loss, m, train, val, optimiser_α, optimiser_w, losses, epoch; cbepoch = cbepoch, cbbatch = cbbatch)
end
display(("done", Dates.format(convert(DateTime,now()-beginscript), "HH:MM:SS")))


# This file is a part of BinaryCommitteeMachineFBP.jl. License is MIT: http://github.com/carlobaldassi/BinaryCommitteeMachineFBP.jl/LICENCE.md

module BinaryCommitteeMachineFBP

export focusingBP, MagT64, MagP64,
       read_messages, write_messages,
       FocusingProtocol, StandardReinforcement, Scoping, PseudoReinforcement, FreeScoping

using StatsFuns
using GZip
using ExtractMacro
using SpecialFunctions
using Printf, LinearAlgebra, Random
import Base: copyto!

include("Magnetizations.jl")
using .Magnetizations

include("Util.jl")
using .Util

struct Messages{F<:Mag64}
    M::Int
    N::Int
    K::Int

    # 2 layers:
    #   0  --> external fields
    #   1  --> theta (perceptron) nodes
    #   2  --> theta (perceptron) nodes
    #
    # notation:
    #   m* --> total fields (includes external fields)
    #   u* --> messages directed down (node->variable)
    #   U* --> messages directed up (node->variable)
    #
    # variable names (quantity):
    #   w  --> weights (NxK)
    #   τ1 --> outputs from the first layer of perceptrons (KxM)
    #          also inputs to the second layer
    #   τ2 --> outputs from the second layer of perceptrons (M)
    #
    #                 DEPTH
    ux::MagVec2{F}    # 0
    mw::MagVec2{F}    # 0+1
    mτ1::MagVec2{F}   # 1+2
    uw::MagVec3{F}    # 1
    Uτ1::MagVec2{F}   # 1
    mτ2::MagVec{F}    # 2+
    uτ1::MagVec2{F}   # 2

    function Messages{F}(M::Int, N::Int, K::Int, ux::MagVec2{F}, mw::MagVec2{F}, mτ1::MagVec2{F}, uw::MagVec3{F},
                         Uτ1::MagVec2{F}, mτ2::MagVec{F}, uτ1::MagVec2{F}; check::Bool = true) where {F<:Mag64}
        if check
            checkdims(ux, K, N)
            checkdims(mw, K, N)
            checkdims(mτ1, M, K)
            checkdims(uw, M, K, N)
            checkdims(Uτ1, M, K)
            checkdims(mτ2, M)
            checkdims(uτ1, M, K)
        end
        new(M, N, K, ux, mw, mτ1, uw, Uτ1, mτ2, uτ1)
    end

end

function Messages(::Type{F}, M::Integer, N::Integer, K::Integer, x::Real) where {F<:Mag64}
    ux = [mflatp(F, N) for k = 1:K]
    mw = [mflatp(F, N) for k = 1:K]
    mτ1 = [mflatp(F, K) for a = 1:M]
    uw = [MagVec{F}[map(F, x*(2*rand(N) .- 1)) for k = 1:K] for a = 1:M]
    Uτ1 = [mflatp(F, K) for a = 1:M]
    mτ2 = mflatp(F, M)
    uτ1 = [mflatp(F, K) for a = 1:M]

    for k = 1:K, i = 1:N, a = 1:M
        mw[k][i] = mw[k][i] ⊗ uw[a][k][i]
    end
    for a = 1:M, k = 1:K
        mτ1[a][k] = mτ1[a][k] ⊗ Uτ1[a][k] ⊗ uτ1[a][k]
    end

    return Messages{F}(M, N, K, ux, mw, mτ1, uw, Uτ1, mτ2, uτ1, check=false)
end

Messages(::Type{F}, messages::Messages{F}) where {F<:Mag64} = messages
function Messages(::Type{F}, messages::Messages) where {F<:Mag64}
    @extract messages : M N K ux mw mτ1 uw Uτ1 mτ2 uτ1
    return Messages{F}(M, N, K,
                       chgeltype(ux, F), chgeltype(mw, F), chgeltype(mτ1, F), chgeltype(uw, F),
                       chgeltype(Uτ1, F), chgeltype(mτ2, F), chgeltype(uτ1, F),
                       check=false)
end

function read_messages(io::IO, ::Type{F}) where {F<:Mag64}
    l = split(readline(io))
    length(l) == 2 && l[1] == "fmt:" || error("invalid messages file")
    fmt = Val{Symbol(l[2])}
    l = split(readline(io))
    length(l) == 4 && l[1] == "N,K,M:" || error("invalid messgaes file")
    N, K, M = parse(Int, l[2]), parse(Int, l[3]), parse(Int, l[4])

    ux = [mflatp(F, N) for k = 1:K]
    mw = [mflatp(F, N) for k = 1:K]
    mτ1 = [mflatp(F, K) for a = 1:M]
    uw = [[mflatp(F, N) for k = 1:K] for a = 1:M]
    Uτ1 = [mflatp(F, K) for a = 1:M]
    mτ2 = mflatp(F, M)
    uτ1 = [mflatp(F, K) for a = 1:M]

    expected_lines = K + M + M*K + M + 1 + M + K
    for (i,l) in enumerate(eachline(io))
        i > expected_lines && (strip(l) == "END" || error("invalid messages file"); break)
        @readmagvec(l, fmt, ux, mw, mτ1, uw, Uτ1, mτ2, uτ1)
    end
    eof(io) || error("invalid messages file")
    return Messages{F}(M, N, K, ux, mw, mτ1, uw, Uτ1, mτ2, uτ1, check=false)
end

"""
    read_messages(filename, mag_type)

Reads messages from a file. `mag_type` is the internal storage format used in the resulting `Messages` object,
it can be either `MagT64` (uses tanhs, accurate but slower) or `MagP64` (plain format, faster but inaccurate).

The file format is the one produced by [`write_messages`](@ref).
"""
read_messages(filename::AbstractString, ::Type{F}) where {F<:Mag64} = gzopen(io->read_messages(io, F), filename, "r")

"""
    write_messages(filename, messages)

Writes messages to a file. The messages can be read back with [`read_messages`](@ref). Note that the output
is a plain text file compressed with gzip.
"""
function write_messages(filename::AbstractString, messages::Messages)
    gzopen(filename, "w") do f
        write_messages(f, messages)
    end
end

function write_messages(io::IO, messages::Messages{F}) where {F<:Mag64}
    @extract messages : N K M ux mw mτ1 uw Uτ1 mτ2 uτ1

    println(io, "fmt: ", magformat(F))
    println(io, "N,K,M: $N $K $M")
    @dumpmagvecs(io, ux, mw, mτ1, uw, Uτ1, mτ2, uτ1)
    println(io, "END")
end

Base.eltype(messages::Messages{F}) where {F<:Mag64} = F

function copyto!(dest::Messages{F}, src::Messages{F}) where {F<:Mag64}
    dest.N == src.N || throw(ArgumentError("incompatible arguments: dest.N=$(dest.N) src.N=$(src.N)"))
    dest.K == src.K || throw(ArgumentError("incompatible arguments: dest.K=$(dest.K) src.K=$(src.K)"))
    dest.M == src.M || throw(ArgumentError("incompatible arguments: dest.M=$(dest.M) src.M=$(src.M)"))
    for k = 1:dest.K
        copyto!(dest.ux[k], src.ux[k])
        copyto!(dest.mw[k], src.mw[k])
    end
    for a = 1:dest.M, k = 1:dest.K
        copyto!(dest.uw[a][k], src.uw[a][k])
    end
    for a = 1:dest.M
        copyto!(dest.mτ1[a], src.mτ1[a])
        copyto!(dest.Uτ1[a], src.Uτ1[a])
        copyto!(dest.uτ1[a], src.uτ1[a])
    end
    copyto!(dest.mτ2, src.mτ2)
    return dest
end

function set_outfields!(messages::Messages{F}, output::Vector, β::Float64) where {F<:Mag64}
    @extract messages : N K M mτ2
    @assert length(output) == M
    t = tanh(β / 2)
    for a = 1:M
        mτ2[a] = forcedmag(F, output[a] * t) # forced avoids clamping
    end
end

print_mags(messages::Messages) = print_mags(stdout, messages)
function print_mags(io::IO, messages::Messages)
    @extract messages : N K mw
    for k = 1:K, i = 1:N
        @printf(io, "%i %i %.15f\n", k, i, Float64(mw[k][i]))
    end
end

mutable struct Params{F<:Mag64}
    damping::Float64
    ϵ::Float64
    β::Float64
    max_iters::Int
    accuracy1::Symbol
    accuracy2::Symbol
    r::Float64
    pol::F
    dγ::Float64
    quiet::Bool
end

struct Patterns
    M::Int
    X::Vec2
    output::IVec
    function Patterns(X::AbstractVector, output::AbstractVector)
        M = length(X)
        length(output) == M || throw(ArgumentError("incompatible lengths of inputs and outputs: $M vs $(length(output))"))
        all(ξ->all(ξi->abs(ξi) == 1, ξ), X) || throw(ArgumentError("inputs must be ∈ {-1,1}"))
        all(o->abs(o) == 1, output) || throw(ArgumentError("outputs must be ∈ {-1,1}"))
        new(M, X, output)
    end
end
Patterns(Xo::Tuple{Vec2,Vec}) = Patterns(Xo...)

@doc """
    Patterns(inputs::AbstractVector, outputs::AbstractVector)

Construct a `Patterns` object. `inputs` and `outputs` must have the same length.
`outputs` entries must be ∈ {-1,1}. `intputs` entries must be Vectors in which
each element is ∈ {-1,1}, and all the vectors must have the same length.
""" Patterns(inputs::AbstractVector, outputs::AbstractVector)

"""
    Patterns(NM::Tuple{Integer,Integer}; teacher::Union{Bool,Int,Vector{Vector{Float64}}} = false)

Constructs random (unbiasad, unifom i.i.d.) binary patterns. The `NM` argument is a Tuple with
two items, the size of the inputs (`N`) and the number of patterns (`M`).
The keyword argument `teacher` controls how the outputs are generated:
* If `false`, they are random i.i.d. (the default)
* If `true`, a teacher unit with a single hidden layer (a perceptron) is generated randomly
  and it's used to compute the outputs
* If an `Int` is given, it's the number of hidden units of the teacher, which is generated
  randomly and used to compute the outputs
* If a `Vector{Vector{Float64}}` is given, it represents the weights of the teacher, with one
  `Vector{Float64}` for each hidden unit.
"""
function Patterns(NM::Tuple{Integer,Integer}; teacher::Union{Bool,Int,Vec2} = false)
    N, M = NM
    X = [rand(-1.0:2.0:1.0, N) for a = 1:M]
    if teacher == false
        output = ones(M)
    else
        if !isa(teacher, Vec2)
            K::Int = isa(teacher, Bool) ? 1 : teacher
            tw = [rand(-1.0:2.0:1.0, N) for k = 1:K]
        else
            K = length(teacher)
            K > 0 || throw(ArgumentError("empty teacher"))
            tw = teacher
            all(w -> length(w) == N, tw) || throw(ArgumentError("invalid teacher length, expected $N, given: $(sort!(unique(collect(w->length(w) for w in tw))))"))
            all(w -> all(x -> abs(x) == 1, w), tw) || throw(ArgumentError("invalid teacher, entries must be all ±1"))
        end
        r0 = Array{Float64}(undef, K)
        output = Float64[transf1!(r0, tw, ξ) for ξ in X]
        @assert all(o->abs(o) == 1, output)
    end
    Patterns(X, output)
end
Patterns(patterns::Patterns) = deepcopy(patterns)

"""
    Patterns(patternsfile::AbstractString)

Read patterns from a file. It only reads the inputs, one pattern per line, entries separated by
whitespace. All lines must have the same length. All outputs are assumed to be `1`. The file may
optionally be gzipped.
"""
function Patterns(patternsfile::AbstractString)
    X = Vec[]
    N = 0
    M = 0
    gzopen(patternsfile) do f
        M = 0
        for l in eachline(f)
            push!(X, map(x->parse(Float64, x), split(l)))
            M += 1
        end
    end
    o = Int[1.0 for a = 1:M]

    return Patterns(X, o)
end

function computeσ²(w::Vec)
    σ² = 0.0
    @inbounds @simd for wi in w
        # σ² += (1 - wi^2)
        σ² += (1 - wi*wi)
    end
    return σ²
end

# function computeσ²(w::Vec, ξ::Vec)
#     σ² = 0.0
#     @inbounds for (wi,ξi) in zip(w,ξ)
#         σ² += (1 - wi^2) * ξi^2
#     end
#     return σ²
# end

computeσ(σ²::Float64) = √(2σ²)
computeσ(w::Vec) = √(2computeσ²(w))
# computeσ(w::Vec, ξ::Vec) = √(2computeσ²(w, ξ))

# transf0(w::Vec, ξ::Vec) = transf0(w, computeσ(w, ξ), ξ)
transf0(w::Vec, ξ::Vec) = transf0(w, computeσ(w), ξ) # NOTE: assuming ξ = ±1 here!!!
transf0(w::Vec, σ::Float64, ξ::Vec) = erf((ξ ⋅ w) / σ)

function transf0!(r::Vec, w1::Vec2, ξ::Vec)
    for (i,w) in enumerate(w1)
        r[i] = transf0(w, ξ)
    end
    return r
end

gauss(x, dσ²) = ℯ^(-x^2 / dσ²) / √(π * dσ²)
gauss(x) = ℯ^(-x^2 / 2) / √(2π)

function subfield!(h::MagVec, m::MagVec, u::MagVec)
    @inbounds @simd for i = 1:length(m)
        h[i] = m[i] ⊘ u[i]
    end
end

function addfield!(m::MagVec, h::MagVec, u::MagVec)
    @inbounds @simd for i = 1:length(m)
        m[i] = h[i] ⊗ u[i]
    end
end

let hsT = Dict{Int,MagVec{MagT64}}(), hsP = Dict{Int,MagVec{MagP64}}(), vhs = Dict{Int,Vec}(),
    vHs = Dict{Int,Vec}(), leftCs = Dict{Int,Vec2}(), rightCs = Dict{Int,Vec2}()

    geth(::Type{MagT64}, N::Int) = Base.get(()->Array{MagT64}(undef, N), hsT, N)
    geth(::Type{MagP64}, N::Int) = Base.get(()->Array{MagP64}(undef, N), hsP, N, )

    global theta_node_update_approx!
    function theta_node_update_approx!(m::MagVec{F}, M::F, ξ::Vec, u::MagVec{F}, U::F, params::Params) where {F<:Mag64}
        @extract params : λ=damping

        N = length(m)
        h::MagVec{F} = geth(F, N)
        vh = Base.get(()->Array{Float64}(undef, N), vhs, N)

        subfield!(h, m, u)
        H = M ⊘ U

        @inbounds for i = 1:N
            vh[i] = h[i]
        end

        vH = Float64(H)
        # σ² = computeσ²(vh, ξ)
        σ² = computeσ²(vh) # assume ξ = ±1

        μ = vh ⋅ ξ

        dσ² = 2σ²
        newU = merf(F, μ / √dσ²)

        maxdiff = abs(U - newU)
        U = damp(newU, U, λ)
        newM = H ⊗ U
        M = newM

        g = gauss(μ, dσ²)

        p0 = 2vH * g / (1 + vH * U)

        pμ = p0 * (p0 + μ / σ²)

        pσ = p0 * (1 - μ / σ² - μ * p0) / dσ²

        @inbounds for i = 1:N
            ξi = ξ[i]
            hi = vh[i]
            newu = convert(F, clamp(ξi * (p0 + ξi * (hi * pμ + ξi * (1-hi^2) * pσ)), -1+3e-16, 1-3e-16)) # use mag-functions?
            d = conv_diff(newu, u[i])
            maxdiff = ifelse(maxdiff < d, d, maxdiff)
            u[i] = damp(newu, u[i], λ)
        end
        addfield!(m, h, u)

        return maxdiff, U, M
    end

    global theta_node_update_accurate!
    function theta_node_update_accurate!(m::MagVec{F}, M::F, ξ::Vec, u::MagVec{F}, U::F, params::Params) where {F<:Mag64}
        @extract params : λ=damping

        N = length(m)
        h::MagVec{F} = geth(F, N)
        vh = Base.get(()->Array{Float64}(undef, N), vhs, N)

        subfield!(h, m, u)
        H = M ⊘ U

        @inbounds for i = 1:N
            vh[i] = h[i]
        end

        # σ² = computeσ²(vh, ξ)
        σ² = computeσ²(vh) # assume ξ = ±1
        μ = vh ⋅ ξ

        dσ² = 2σ²
        newU = merf(F, μ / √dσ²)

        maxdiff = 0.0
        U = damp(newU, U, λ)
        M = H ⊗ U

        @inbounds for i = 1:N
            ξi = ξ[i]
            hi = vh[i]
            μ̄ = μ - ξi * hi
            σ̄² = σ² - (1-hi^2) # * ξi^2 # assume ξ = ±1
            sdσ̄² = √(2σ̄²)
            m₊ = (μ̄ + ξi) / sdσ̄²
            m₋ = (μ̄ - ξi) / sdσ̄²
            newu = erfmix(H, m₊, m₋)
            d = abs(newu - u[i])
            maxdiff = ifelse(maxdiff < d, d, maxdiff)
            u[i] = damp(newu, u[i], λ)
        end
        addfield!(m, h, u)
        return maxdiff, U, M
    end

    global theta_node_update_exact!
    function theta_node_update_exact!(m::MagVec{F}, M::F, ξ::Vec, u::MagVec{F}, U::F, params::Params) where {F<:Mag64}
        @extract params : λ=damping

        N = length(m)
        h::MagVec{F} = geth(F, N)
        vh = Base.get(()->Array{Float64}(undef, N), vhs, N)
        leftC = Base.get(()->[zeros(i+1) for i = 1:N], leftCs, N)
        rightC = Base.get(()->[zeros((N-i+1)+1) for i = 1:N], rightCs, N)

        subfield!(h, m, u)
        H = M ⊘ U

        @inbounds for i = 1:N
            vh[i] = h[i]
        end

        leftC[1][1] = (1-ξ[1]*vh[1])/2
        leftC[1][2] = (1+ξ[1]*vh[1])/2
        for i = 2:N
            lC0 = leftC[i-1]
            lC = leftC[i]
            hi = ξ[i] * vh[i]
            lC[1] = lC0[1] * (1-hi)/2
            for j = 2:i
                lC[j] = lC0[j-1] * (1+hi)/2 + lC0[j] * (1-hi)/2
            end
            lC[end] = lC0[end] * (1+hi)/2
        end

        rightC[end][1] = (1-ξ[end]*vh[end])/2
        rightC[end][2] = (1+ξ[end]*vh[end])/2
        for i = (N-1):-1:1
            rC0 = rightC[i+1]
            rC = rightC[i]
            hi = ξ[i] * vh[i]
            rC[1] = rC0[1] * (1-hi)/2
            for j = 2:(N-i+1)
                rC[j] = rC0[j-1] * (1+hi)/2 + rC0[j] * (1-hi)/2
            end
            rC[end] = rC0[end] * (1+hi)/2
        end

        @assert maximum(abs.(leftC[end] .- rightC[1])) ≤ 1e-10 (leftC[end], rightC[1])

        @assert isodd(N)
        z = (N+1) ÷ 2
        pm = sum(rightC[1][1:z])
        pp = sum(rightC[1][(z+1):end])

        newU = Mag64(F, pp, pm)

        @assert isfinite(newU)

        maxdiff = 0.0
        U = damp(newU, U, λ)
        newM = H ⊗ U
        M = newM

        @assert isfinite(newM) (H, U)

        u1 = ones(1)

        @inbounds for i = 1:N
            ξi = ξ[i]
            @assert ξi^2 == 1

            lC = i > 1 ? leftC[i-1] : u1
            rC = i < N ? rightC[i+1] : u1

            pm = 0.0
            pz = 0.0
            pp = 0.0
            for j = 1:N
                p = 0.0
                for k = max(1,j+i-N):min(j,i)
                    p += lC[k] * rC[j-k+1]
                end
                if j < z
                    pm += p
                elseif j == z
                    pz = p
                else
                    pp += p
                end
            end

            mp = convert(F, clamp(pp + ξi * pz - pm, -1.0, 1.0))
            mm = convert(F, clamp(pp - ξi * pz - pm, -1.0, 1.0))
            newu = exactmix(H, mp, mm)

            d = abs(newu - u[i])
            maxdiff = ifelse(maxdiff < d, d, maxdiff)
            u[i] = damp(newu, u[i], λ)
            @assert isfinite(u[i]) (u[i],)
        end
        addfield!(m, h, u)
        return maxdiff, U, M
    end

    global free_energy_theta
    function free_energy_theta(m::MagVec{F}, M::F, ξ::Vec, u::MagVec{F}, U::F) where {F<:Mag64}
        N = length(m)
        h::MagVec{F} = geth(F, N)
        vh = Base.get(()->Array{Float64}(undef, N), vhs, N)

        f = 0.0

        subfield!(h, m, u)
        H = M ⊘ U

        @inbounds for i = 1:N
            vh[i] = h[i]
        end

        # σ = computeσ(vh, ξ)
        σ = computeσ(vh) # assume ξ = ±1
        μ = vh ⋅ ξ

        b = merf(F, μ / σ)

        f -= log1pxy(H, b)
        @assert isfinite(f)

        for i = 1:N
            f += log1pxy(h[i], u[i])
        end
        return f
    end

    global free_energy_theta_exact
    function free_energy_theta_exact(m::MagVec{F}, M::F, ξ::Vec, u::MagVec{F}, U::F) where {F<:Mag64}

        N = length(m)
        h::MagVec{F} = geth(F, N)
        vh = Base.get(()->Array{Float64}(undef, N), vhs, N)
        leftC = Base.get(()->[zeros(i+1) for i = 1:N], leftCs, N)
        rightC = Base.get(()->[zeros((N-i+1)+1) for i = 1:N], rightCs, N)

        f = 0.0

        subfield!(h, m, u)
        H = M ⊘ U

        @inbounds for i = 1:N
            vh[i] = h[i]
        end
        vH = Float64(H)

        leftC[1][1] = (1-ξ[1]*vh[1])/2
        leftC[1][2] = (1+ξ[1]*vh[1])/2
        for i = 2:N
            lC0 = leftC[i-1]
            lC = leftC[i]
            hi = ξ[i] * vh[i]
            lC[1] = lC0[1] * (1-hi)/2
            for j = 2:i
                lC[j] = lC0[j-1] * (1+hi)/2 + lC0[j] * (1-hi)/2
            end
            lC[end] = lC0[end] * (1+hi)/2
        end

        rightC[end][1] = (1-ξ[end]*vh[end])/2
        rightC[end][2] = (1+ξ[end]*vh[end])/2
        for i = (N-1):-1:1
            rC0 = rightC[i+1]
            rC = rightC[i]
            hi = ξ[i] * vh[i]
            rC[1] = rC0[1] * (1-hi)/2
            for j = 2:(N-i+1)
                rC[j] = rC0[j-1] * (1+hi)/2 + rC0[j] * (1-hi)/2
            end
            rC[end] = rC0[end] * (1+hi)/2
        end

        @assert maximum(abs.(leftC[end] .- rightC[1])) ≤ 1e-10 (leftC[end], rightC[1])

        @assert isodd(N)
        z = (N+1) ÷ 2
        pm = sum(rightC[1][1:z])
        pp = sum(rightC[1][(z+1):end])

        b = Mag64(F, pp, pm)
        f -= log1pxy(H, b)
        @assert isfinite(f)

        for i = 1:N
            f += log1pxy(h[i], u[i])
        end
        return f
    end

end

function entro_node_update(m::F, u::F, params::Params{F}) where {F<:Mag64}
    @extract params : λ=damping r pol

    h = m ⊘ u
    if r == 0 || iszero(pol)
        newu = zero(F)
    elseif r == Inf
        newu = ifelse(iszero(h), zero(F), copysign(pol, h))
    else
        newu::F = ((h * pol) ↑ r) * pol
    end

    diff = abs(newu - u)
    newu = damp(newu, u, λ)
    newm = h ⊗ newu

    return diff, newu, newm
end

function iterate!(messages::Messages{F}, patterns::Patterns, params::Params) where {F<:Mag64}
    @extract messages : N M K ux mw mτ1 uw Uτ1 mτ2 uτ1
    @extract patterns : X output
    @extract params   : accuracy1 accuracy2
    maxdiff = 0.0
    tnu1! = accuracy1 == :exact ? theta_node_update_exact! :
            accuracy1 == :accurate ? theta_node_update_accurate! :
            accuracy1 == :none ? theta_node_update_approx! :
            error("accuracy must be one of :exact, :accurate, :none (was given $accuracy)")
    tnu2! = accuracy2 == :exact ? theta_node_update_exact! :
            accuracy2 == :accurate ? theta_node_update_accurate! :
            accuracy2 == :none ? theta_node_update_approx! :
            error("accuracy must be one of :exact, :accurate, :none (was given $accuracy)")
    for a = randperm(M + N*K)
        if a ≤ M
            ξ = X[a]
            for k = 1:K
                diff, Uτ1[a][k], mτ1[a][k] = tnu1!(mw[k], mτ1[a][k], ξ, uw[a][k], Uτ1[a][k], params)
                maxdiff = max(maxdiff, diff)
            end
            diff, _, mτ2[a] = tnu2!(mτ1[a], mτ2[a], ones(K), uτ1[a], zero(F), params)
            maxdiff = max(maxdiff, diff)
        else
            (params.r == 0 || iszero(params.pol)) && continue
            j = a - M
            k = (j-1) ÷ N + 1
            i = (j-1) % N + 1

            diff, ux[k][i], mw[k][i] = entro_node_update(mw[k][i], ux[k][i], params)
            maxdiff = max(diff, maxdiff)
        end
    end
    return maxdiff
end

function converge!(messages::Messages, patterns::Patterns, params::Params)
    @extract params : ϵ max_iters λ₀=damping quiet

    λ = λ₀
    ok = false
    strl = 0
    t = @elapsed for it = 1:max_iters
        diff = iterate!(messages, patterns, params)

        if !quiet
            str = "[it=$it Δ=$diff λ=$λ]"
            print("\r", " "^strl, "\r", str)
            strl = length(str)
            #println(str)
            flush(stdout)
            strl = length(str)
        end
        if diff < ϵ
            ok = true
            quiet || println("\nok")
            break
        end
    end
    if !quiet
        ok || println("\nfailed")
        println("elapsed time = $t seconds")
    end
    return ok
end

transf1(w::Vec) = sign(sum(w))

transf1!(r0::Vec, ws::Vec2, ξ::Vec) = transf1(transf0!(r0, ws, ξ))

function test!(r0::Vec, ws::Vec2, ξ::Vec, out::Int)
    o = transf1!(r0, ws, ξ)
    return o != out
end

let r0d = Dict{Int,Vec}()
    global test
    function test(ws::Vec2, ξs::Vec2, output::IVec)
        N = length(ws)
        r0 = get!(r0d, N) do; Vec(undef, N); end
        sum([test!(r0, ws, ξ, out) for (ξ,out) in zip(ξs, output)])
    end
end

function test(messages::Messages, patterns::Patterns)
    @extract messages : N K mw
    @extract patterns : X output
    ws = [Float64[mw[k][i] for i = 1:N] for k = 1:K]
    return test(ws, X, output)
end

function nonbayes_test(messages::Messages, patterns::Patterns)
    @extract messages : N K mw
    @extract patterns : X output
    ws = [Float64[sign0(mw[k][i]) for i = 1:N] for k = 1:K]
    return test(ws, X, output)
end

function free_energy(messages::Messages{F}, patterns::Patterns, params::Params{F}) where {F<:Mag64}
    @extract messages : M N K ux mw mτ1 uw Uτ1 mτ2 uτ1
    @extract patterns : X output
    @extract params   : r pol

    f = 0.0

    for a = 1:M
        ξ = X[a]
        for k = 1:K
            f += free_energy_theta(mw[k], mτ1[a][k], ξ, uw[a][k], Uτ1[a][k])
        end
        f += free_energy_theta_exact(mτ1[a], mτ2[a], ones(K), uτ1[a], zero(F))
    end

    for k = 1:K, i = 1:N
        f -= logZ(ux[k][i], F[uw[a][k][i] for a = 1:M])

        f -= logtwo / 2
        f += log1pxy(pol, -pol) / 2
        hkix = mw[k][i] ⊘ ux[k][i]
        f += log1pxy(hkix, ux[k][i])
        hpol = hkix * pol
        mx = hpol ↑ (r + 1)
        f += mcrossentropy(mx, hpol)
    end

    return f / (N * K)
end

function compute_S(messages::Messages{F}, params::Params{F}) where {F<:Mag64}
    @extract messages : N K ux mw
    @extract params   : r pol
    S = 0.0
    for k = 1:K, i = 1:N
        hkix = mw[k][i] ⊘ ux[k][i]
        hxki = (hkix * pol) ↑ r
        hh = hkix * hxki
        S += Float64(hh ⊗ pol)
    end
    return S / (N * K)
end

function compute_q̃(messages::Messages{F}, params::Params{F}) where {F<:Mag64}
    @extract messages : N K ux mw
    @extract params   : r pol
    q̃ = 0.0
    for k = 1:K, i = 1:N
        hkix = mw[k][i] ⊘ ux[k][i]
        mx = (hkix * pol) ↑ (r + 1)
        q̃ += mx^2
    end
    return q̃ / (N * K)
end

function compute_q(messages::Messages)
    @extract messages : N K mw
    q = 0.0
    for k = 1:K, i = 1:N
        q += Float64(mw[k][i])^2
    end
    return q / (N * K)
end

function mags_symmetry(messages::Messages)
    @extract messages : N K mw
    overlaps = Matrix(1.0I, K, K)
    qs = zeros(K)
    for k1 = 1:K
        z = 0.0
        for i = 1:N
            z += Float64(mw[k1][i])^2
        end
        qs[k1] = √z
    end
    for k1 = 1:K, k2 = k1+1:K
        s = 0.0
        for i = 1:N
            s += Float64(mw[k1][i]) * Float64(mw[k2][i])
        end
        s /= qs[k1] * qs[k2]
        overlaps[k1,k2] = s
        overlaps[k2,k1] = s
    end
    return overlaps, qs / N
end

"""
    FocusingProtocol

Abstract type representing a protocol for the focusing procedure, i.e. a way to produce
successive values for the quantities `γ`, `y` and `β`. Currently, however, only `β=Inf`
is supported. To be provided as an argument to [`focusingBP`](@ref).

Available protocols are: [`StandardReinforcement`](@ref), [`Scoping`](@ref), [`PseudoReinforcement`](@ref) and
[`FreeScoping`](@ref).
"""
abstract type FocusingProtocol end

struct StandardReinforcement <: FocusingProtocol
    r::AbstractRange{Float64}
    StandardReinforcement(r::AbstractRange{T}) where {T<:Real} = new(r)
end

@doc """
    StandardReinforcement(r::AbstractRange) <: FocusingProtocol

Standard reinforcement protocol, returns `γ=Inf` and `y=1/(1-x)`, where `x` is taken from the given range `r`.
""" StandardReinforcement(r::AbstractRange)

"""
    StandardReinforcement(dr::Float64) <: FocusingProtocol

Shorthand for [`StandardReinforcement`](@ref)`(0:dr:(1-dr))`.
"""
StandardReinforcement(dr::Float64) = StandardReinforcement(0.0:dr:(1-dr))

function Base.iterate(s::StandardReinforcement, i = 1)
    n = iterate(s.r, i)
    n ≡ nothing && return nothing
    return (Inf, 1/(1-n[1]), Inf), n[2]
end

"""
    Scoping(γr::AbstractRange, y) <: FocusingProtocol

Focusing protocol with fixed `y` and a varying `γ` taken from the given `γr` range.
"""
struct Scoping <: FocusingProtocol
    γr::AbstractRange{Float64}
    y::Float64
    Scoping(γr::AbstractRange, y) = new(γr, y)
end

function Base.iterate(s::Scoping, i = 1)
    n = iterate(s.γr, i)
    n ≡ nothing && return nothing
    return (n[1], s.y, Inf), n[2]
end

struct PseudoReinforcement <: FocusingProtocol
    r::Vector{Float64}
    x::Float64
    PseudoReinforcement(r::AbstractRange{T}...; x::Real=0.5) where {T<:Real} = new(vcat(map(collect, r)...), x)
end
@doc """
    PseudoReinforcement(r::AbstractRange...; x=0.5) <: FocusingProtocol

A focusing protocol in which both `γ` and `y` are progressively increased, according to
the formulas

```julia
γ = atanh(ρ^x)
y = 1+ρ^(1-2x)/(1-ρ)
```

where `ρ` is taken from the given range(s) `r`. With `x=0`, this is basically the same as
[`StandardReinforcement`](@ref).
""" PseudoReinforcement(r::AbstractRange...)

"""
    PseudoReinforcement(dr::Float64; x=0.5) <: FocusingProtocol

Shorthand for [`PseudoReinforcement`](@ref)`(0:dr:(1-dr); x=x)`.
""" PseudoReinforcement(dr::Float64; x::Real=0.5) = PseudoReinforcement(0.0:dr:(1-dr), x=x)

function Base.iterate(s::PseudoReinforcement, i = 1)
    n = iterate(s.r, i)
    n ≡ nothing && return nothing
    x = s.x
    ρ = n[1]
    # some special cases just to avoid possible 0^0
    if x == 0.5
        return (atanh(√ρ), (2-ρ)/(1-ρ), Inf), n[2]
    elseif x == 0
        return (Inf, 1/(1-ρ), Inf), n[2]
    else
        return (atanh(ρ^x), 1+ρ^(1-2x)/(1-ρ), Inf), n[2]
    end
end

"""
    FreeScoping(list::Vector{NTuple{2,Float64}}) <: FocusingProtocol

A focusing protocol which just returns the values of `(γ,y)` from the given `list`.

Example:

```julia
FreeScoping([(1/(1-x), (2-x)/(1-x)) for x = 0:0.01:0.99])
```
"""
struct FreeScoping <: FocusingProtocol
    list::Vector{NTuple{3,Float64}}
    FreeScoping(list::Vector{NTuple{3,Float64}}) = new(list)
end
FreeScoping(list::Vector{NTuple{2,Float64}}) = FreeScoping(NTuple{3,Float64}[(γ,y,Inf) for (γ,y) in list])

Base.iterate(s::FreeScoping, i = 1) = iterate(s.list, i)

"""
    focusingBP(N, K, patternspec; keywords...)

Run the Focusing Belief Propagation algorithm on a fully-connected committee machine with binary weights.
`N` is the input (first layer) size, `K` the number of the hidden units (second layer size), and
`patternspec` specifies how to build the patterns for the training set. Note that with the defult
settings `K` must be odd (see notes for the `accuracy1` and `accuracy2` arguments below).

Possible values of `patternspec` are:

* a `Float64` number: this is interpreted as the `α` parameter, and `M = α*N*K` random ±1 patterns are generated.
* a `Tuple` with `Vector{Vector{Float64}}` and a `Vector{Float64}`: these are the inputs and associated desired outputs.
* a string: the patterns are read from a file (one input pattern per line, entries separated by whitespace, outputs are
            assumed to be all 1); the file can be gzipped.
* a [`Patterns`](@ref) object (which could be the output of a previous run of the function).

*Note*: all inputs and outputs must be ∈ {-1,1}.

The keyword arguments are:

* `max_iters` (default = `1000`): maximum number of BP iterations per step. If convergence is not achieved in this many iterations,
                                  the algorithm proceeds to the next step.
* `max_steps` (default = `typemax(Int)`): maximum number of focusing steps.
* `seed` (default = `1`): random seed.
* `damping` (default = `0`): BP damping parameter (between `0` and `1`; `0` = no damping).
* `quiet` (default = `false`): whether to print on screen
* `accuracy1` (default = `:accurate`): accuracy of the messages computation at the hidden units level. Possible values are:
                                       `:accurate` (Gaussian approximation, good for large `N`, works in `O(N)` time),
                                       `:exact` (no approximation, uses convolutions, good for small `N`, works in `O(N³)` time, requires
                                       `N` to be odd),
                                       `:none` (a TAP-like approximation, fast but never very good, should probably be removed or done
                                       properly...).
* `accuracy2` (default = `:exact`): accuracy of the messages computation at the output node level.
                                    See `accuracy1` (and think of `K` instead of `N`).
* `randfact` (default = `0.01`): random factor used in the initialization of the messages. Must be between `0` and `1`. Large values are
                                 not a good idea.
* `fprotocol` (default = `StandardReinforcement(1e-2)`): focusing protocol specification. See [`FocusingProtocol`](@ref).
* `ϵ` (default = `1e-3`): convergence criterion: BP is assumed to have converged when the difference between messages in two successive
                          iterations is smaller than this value. Reduce it (e.g. to 1e-6) for more precise results, while increasing
                          `max_iters`.
* `messfmt` (default = `:tanh`): internal storage format for messages: it can be either `:tanh` or `:plain`; `:tanh` is much more
                                 precise but slower.
* `initmess` (default = `nothing`): how to initialize the messages. If `nothing`, they are initialized randomly; If a string is given,
                                    they are read from a file (see also [`read_messages`](@ref) and [`write_messages`](@ref)). If a
                                    `Messages` object is given (e.g. one returned from an earlier run of the function, or from
                                    [`read_messages`](@ref)) it is used (and overwritten).
* `outatzero` (default = `true`): if `true`, the algorithm exits as soon as a solution to the learning problem is found, without waiting
                                  for the focusing protocol to terminate.
* `writeoutfile` (default = `:auto`): whether to write results on an output file. Can be `:never`, `:always` or `:auto`. The latter means
                                      that a file is written when `outatzero` is set to `false` and only when BP converges. It can make sense
                                      setting this to `:always` even when `outfile` is `nothing` to force the computation of the local
                                      entropy and other thermodynamic quantities.
* `outfile` (default = `nothing`): the output file name. `nothing` means no output file is written. An empty string means using a default
                                   file name of the form: `"results_BPCR_N\$(N)_K\$(K)_M\$(M)_s\$(seed).txt"`.
* `outmessfiletmpl` (default = `nothing`): template file name for writing the messages at each focusing step. The file name should include a
                                           substring `"%gamma%"` which will be substituted with the value of the `γ` parameter at each step.
                                           If `nothing`, it is not used. If empty, the default will be used:
                                           `"messages_BPCR_N\$(N)_K\$(K)_M\$(M)_g%gamma%_s\$(seed).txt.gz"`.
                                           *Note*: this can produce a lot of fairly large files!.

The function returns three objects: the number of training errors, the messages and the patterns. The last two can be used as inputs to
successive runs of the algorithms, as the `initmessages` keyword argument and the `patternspec` argument, respectively.

Example of a run which solves a problem with `N * K = 1605` synapses with `K = 5` at `α = 0.3`:
```
julia> errs, messages, patterns = focusingBP(321, 5, 0.3, randfact=0.1, seed=135, max_iters=1, damping=0.5);
```
"""
function focusingBP(N::Integer, K::Integer,
                    initpatt::Union{AbstractString,Tuple{Vec2,Vec},Real,Patterns};

                    max_iters::Integer = 1000,
                    max_steps::Integer = typemax(Int),
                    seed::Integer = 1,
                    damping::Real = 0.0,
                    quiet::Bool = false,
                    accuracy1::Symbol = :accurate,
                    accuracy2::Symbol = :exact,
                    randfact::Real = 0.01,
                    fprotocol::FocusingProtocol = StandardReinforcement(1e-2),
                    ϵ::Real = 1e-3,
                    messfmt::Symbol = :tanh,
                    initmess::Union{Messages,Nothing,AbstractString} = nothing,
                    outatzero::Bool = true,
                    writeoutfile::Symbol = :auto, # note: ∈ [:auto, :always, :never]; auto => !outatzero && converged
                    outfile::Union{AbstractString,Nothing} = nothing, # note: "" => default, nothing => no output
                    outmessfiletmpl::Union{AbstractString,Nothing} = nothing) # note: same as outfile

    Random.seed!(seed)

    N > 0 || throw(ArgumentError("N must be positive; given: $N"))
    K > 0 || throw(ArgumentError("K must be positive; given: $K"))

    writeoutfile ∈ [:auto, :always, :never] || error("invalide writeoutfile, expected one of :auto, :always, :never; given: $writeoutfile")
    max_iters ≥ 0 || throw(ArgumentError("max_iters must be non-negative; given: $max_iters"))
    max_steps ≥ 0 || throw(ArgumentError("max_steps must be non-negative; given: $max_steps"))
    0 ≤ damping < 1 || throw(ArgumentError("damping must be ∈ [0,1); given: $damping"))
    0 ≤ randfact ≤ 1 || throw(ArgumentError("randfact must be ∈ [0,1]; given: $randfact"))
    accuracy1 ∈ [:exact, :accurate, :none] || error("accuracy1 must be one of :exact, :accurate, :none; given: $accuracy1")
    accuracy2 ∈ [:exact, :accurate, :none] || error("accuracy2 must be one of :exact, :accurate, :none; given: $accuracy2")

    accuracy1 == :exact && iseven(N) && throw(ArgumentError("when accuracy1==:exact N must be odd, given: $N"))
    accuracy2 == :exact && iseven(K) && throw(ArgumentError("when accuracy2==:exact K must be odd, given: $K"))

    messfmt ∈ [:tanh, :plain] || throw(ArgumentError("invalid messfmt, should be :tanh or :plain; given: $messfmt"))
    F = messfmt == :tanh ? MagT64 : MagP64

    if isa(initpatt, Real)
        initpatt ≥ 0 || throw(ArgumentError("invalide negative initpatt; given: $initpatt"))
        initpatt = (N, round(Int, K * N * initpatt))
    end

    patterns = Patterns(initpatt)

    M = patterns.M

    messages::Messages = initmess ≡ nothing ? Messages(F, M, N, K, randfact) :
                         isa(initmess, AbstractString) ? read_messages(initmess, F) :
                         initmess

    messages.N == N || throw(ArgumentError("wrong messages size, expected N=$N; given: $(messages.N)"))
    messages.K == K || throw(ArgumentError("wrong messages size, expected K=$K; given: $(messages.K)"))
    messages.M == M || throw(ArgumentError("wrong messages size, expected M=$M; given: $(messages.M)"))
    F = eltype(messages)

    params = Params{F}(damping, ϵ, NaN, max_iters, accuracy1, accuracy2, 0.0, 0.0, 0.0, quiet)

    outfile == "" && (outfile = "results_BPCR_N$(N)_K$(K)_M$(M)_s$(seed).txt")
    outmessfiletmpl == "" && (outmessfiletmpl = "messages_BPCR_N$(N)_K$(K)_M$(M)_g%gamma%_s$(seed).txt.gz")
    lockfile = "bpcomm.lock"
    if outfile ≢ nothing && writeoutfile ∈ [:always, :auto]
        !quiet && println("writing outfile $outfile")
        exclusive(lockfile) do
            !isfile(outfile) && open(outfile, "w") do f
                println(f, "#1=pol 2=y 3=β 4=S 5=q 6=q̃ 7=βF 8=𝓢ᵢₙₜ 9=Ẽ")
            end
        end
    end

    ok = true
    if initmess ≢ nothing
        errs = nonbayes_test(messages, patterns)
        !quiet && println("initial errors = $errs")

        outatzero && errs == 0 && return errs, messages, patterns
    end
    !quiet && K > 1 && (println("mags overlaps="); display(mags_symmetry(messages)[1]); println())

    for (it, (γ,y,β)) in enumerate(fprotocol)
        it > max_steps && break
        isfinite(β) && error("finite β not yet supported; given: $β")
        pol = mtanh(F, γ)
        params.pol = pol
        params.r = y - 1
        params.β = β
        set_outfields!(messages, patterns.output, params.β)
        ok = converge!(messages, patterns, params)
        !quiet && K > 1 && (println("mags overlaps="); display(mags_symmetry(messages)[1]); println())
        errs = nonbayes_test(messages, patterns)

        if writeoutfile == :always || (writeoutfile == :auto && !outatzero)
            S = compute_S(messages, params)
            q = compute_q(messages)
            q̃ = compute_q̃(messages, params)
            βF = free_energy(messages, patterns, params)
            Σint = -βF - γ * S

            !quiet && println("it=$it pol=$pol y=$y β=$β (ok=$ok) S=$S βF=$βF Σᵢ=$Σint q=$q q̃=$q̃ Ẽ=$errs\n")
            (ok || writeoutfile == :always) && outfile ≢ nothing && exclusive(lockfile) do
                open(outfile, "a") do f
                    println(f, "$pol $y $β $S $q $q̃ $βF $Σint $errs")
                end
            end
            if outmessfiletmpl ≢ nothing
                outmessfile = replace(outmessfiletmpl, "%gamma%" => γ)
                write_messages(outmessfile, messages)
            end
        else
            !quiet && println("it=$it pol=$pol y=$y β=$β (ok=$ok) Ẽ=$errs\n")
            errs == 0 && return 0, messages, patterns
        end
    end
    return errs, messages, patterns
end

end # module


function show_volume(vdata::Array{Float64, 3}; cmap::Matrix{Int64}=[(i==1)*(j-1) + (i==3)*(256-j) for i in 1:3, j in 1:256], default_cmap="viridis", default_renderstyle="iso", clim1=0, clim2=1, isothreshold=0.15)
    HTML{String}(
"""<!DOCTYPE html>

	<title>three.js webgl - volume rendering example</title>
	<meta charset="utf-8">
	<meta name="viewport" content="width=device-width, user-scalable=no, minimum-scale=1.0, maximum-scale=1.0">
	<link type="text/css" rel="stylesheet" href="main.css">

""")
end


# Bruker I/O functions

"""     read_bruker_binary(fname)
Bruker FID, and processed frequency-domain data files (fid, 1r, 1i) contain a flat
binary representation of their data points, each of which is an Int32."""
function read_bruker_binary(fname)
    reinterpret(Int32, open(read, fname))
end

function parse_float_list(m)
    lines = split(m)
    # The first line of a match is "(0..n)"; we don't need this.
    [parse(Float64, s) for s in split(m)[2:end]]
end

function parse_or(::Type{T}, s, default::T) where T
    v = tryparse(T, s)
    v === nothing ? default : v
end

filters = [ ( Set(["SW", "O1", "SFO1", "SF", "BF1"]),
              s -> parse(Float64, s) ),
            ( Set(["TD", "NS", "DS", "SI"]),
              s -> parse(Int, s) ),
            ( Set(["D", "P"]),
              parse_float_list),
            ( Set(["PULPROG"]),
              x -> strip(x)[2:end-1] ),
            ( Set(["AUTOPOS"]),
              x -> parse_or(Int, strip(x)[2:end-1], 0) )
]

function read_params(file)
    contents = read(file, String)
    matches = eachmatch(r"##\$?(.*?)=\s?([^#]+)"s, contents)
    res = Dict(m.captures[1] => parse_param(m.captures[1], m.captures[2]) for m in matches)

    # a few little tweaks
    if "O1" in keys(res) && !("O1P" in keys(res))
        res["O1P"] = res["O1"] / res["SFO1"]
    end
    res
end

function read_intrng(file)
    lines = try
        readlines(file)
    catch
        return []
    end
    if length(lines) > 2 && strip(lines[1])[1] == 'A'
        [tuple(map(s -> parse(Float64, s), split(line)[1:2])...) for line in lines[3:end]]
    elseif length(lines) > 1 && strip(lines[1])[1] == 'P'
        [tuple(map(s -> parse(Float64, s), split(line))...) for line in lines[2:end]]
    else
        []
    end
end

function parse_param(param, val)
    for (names, fun) in filters
        if param in names
            return fun(val)
        end
    end
    return strip(val)
end

function ProcessedSpectrum(path :: AbstractString, procno :: Int)
    re_ft = float(read_bruker_binary(joinpath(path, "1r")))
    im_ft = float(read_bruker_binary(joinpath(path, "1i")))
    params = read_params(joinpath(path, "proc"))
    title = read(joinpath(path, "title"), String)
    intrng = read_intrng(joinpath(path, "intrng"))
    ProcessedSpectrum(re_ft, im_ft, params, intrng, procno, title)
end

ProcessedSpectrum(path::AbstractString) = ProcessedSpectrum(path, parse(Int, basename(path)))

Spectrum(path :: AbstractString, procnos :: AbstractArray{Int}, default_proc :: Int) = begin
    fid = float(read_bruker_binary(joinpath(path, "fid")))
    acqu = read_params(joinpath(path, "acqu"))
    name = basename(dirname(path))
    expno = parse(Int, basename(path))
    procs = Dict()
    for procno in procnos
        proc_path = joinpath(path, "pdata", string(procno))
        procs[procno] = ProcessedSpectrum(proc_path, procno)
    end
    Spectrum(fid, acqu, procs, default_proc, name, expno)
end

Spectrum(path :: AbstractString, procno :: Int) = Spectrum(path, [procno], procno)

Spectrum(path :: AbstractString, procnos :: AbstractArray{Int}) = Spectrum(path, procnos, minimum(procnos))

Spectrum(path :: AbstractString) = begin
    procnos = [parse(Int,n) for n in readdir(joinpath(path, "pdata"))]
    Spectrum(path, procnos)
end


# License for this file: MIT (expat)
# Copyright 2017-2021, DLR Institute of System Dynamics and Control
#
# This file is part of module
#   Modia3D.Composition (Modia3D/Composition/_module.jl)
#
import Modia3D.Frames

"""
    joint = Revolute(;obj1, obj2, path="", axis=3, phi=0, w=0, canCollide=false)

Return a Revolute `joint` that rotates `obj1::`[`Object3D`](@ref) into
`obj2::`[`Object3D`](@ref) along the axis `axis` of `obj1` (`axis = 1,2,3,-1,-2,-3`).
The initial start angle is `phi` and the initial angular velocity
is `w`. Negative `axis` values describe axes in negative axis directions, i.e.
the signs of the values of `phi` and `w` are reversed. If `canCollide=false`, no
collision detection will occur between `obj1` and `obj2` (and `Object3D`s that are
directly or indirectly rigidly fixed to `obj1` or `obj2`).

It is currently not supported that a `Revolute` joint *closes a kinematic loop*.
"""
mutable struct Revolute <: Modia3D.AbstractJoint
    path::String

    obj1::Object3D
    obj2::Object3D

    posAxis::Int      # = 1,2,3
    posMovement::Bool # = true, if movement along posAxis, otherwise in negative posAxis
    ndof::Int
    canCollide::Bool  # = false, if no collision between obj1 and obj2

    phi::Float64
    w::Float64
    a::Float64
    tau::Float64
    residue::Float64

    function Revolute(; obj1::Object3D,
                        obj2::Object3D,
                        path::String="",
                        axis::Int = 3,
                        phi::Real = 0.0,
                        w::Real   = 0.0,
                        canCollide::Bool = false)

        (parent,obj,cutJoint) = attach(obj1, obj2)
        if cutJoint
            error("\nError from Revolute joint connecting ", Modia3D.fullName(obj1), " with ", Modia3D.fullName(obj2), ":\n",
                "    This joint is a cut-joint which is currently not supported.!")
        end

        if !(1 <= abs(axis) <= 3)
            error("\nError from Revolute joint connecting ", Modia3D.fullName(obj1), " with ", Modia3D.fullName(obj2), ":\n",
                "    axis = $axis, but must be 1,2,3 or -1,-2-3.!")
        end

        if !(typeof(phi) <: AbstractFloat) || !(typeof(w) <: AbstractFloat)
            error("\nError from Revolute joint connecting ", Modia3D.fullName(obj1), " with ", Modia3D.fullName(obj2), ":\n",
                "    phi and/or w are not <: AbstractFloat!")
        end

        axis = obj === obj2 ? axis : -axis
        phi = Modia3D.convertAndStripUnit(Float64, u"rad"  , phi)
        w   = Modia3D.convertAndStripUnit(Float64, u"rad/s", w)

        posAxis     = abs(axis)
        posMovement = axis > 0

        obj.joint = new(path, parent, obj, posAxis, posMovement, 1, canCollide, phi, w, 0.0, 0.0, 0.0)
        obj.jointKind  = RevoluteKind
        obj.jointIndex = 0
        obj.ndof       = 1
        obj.canCollide = canCollide
        obj.r_rel      = Modia3D.ZeroVector3D
        obj.R_rel      = Frames.rotAxis(posAxis, posMovement, phi)

        parent.hasChildJoint = true
        return obj.joint
    end
end


###############################################################################
#
# Method to fit soil albedo
#
###############################################################################
"""
    abstract type AbstractAlbedoFitting

Hierachy of AbstractAlbedoFitting
- [`FourBandsFittingCurve`](@ref)
- [`FourBandsFittingHybrid`](@ref)
- [`FourBandsFittingPoint`](@ref)
- [`TwoBandsFittingCurve`](@ref)
- [`TwoBandsFittingHybrid`](@ref)
- [`TwoBandsFittingPoint`](@ref)
"""
abstract type AbstractAlbedoFitting end

"""
$(TYPEDEF)

Method to use all four GSV PCA bands to fit two flat lines
"""
struct FourBandsFittingCurve <: AbstractAlbedoFitting end

"""
$(TYPEDEF)

Method to use all four GSV PCA bands to fit a mean and a flat curve
"""
struct FourBandsFittingHybrid <: AbstractAlbedoFitting end

"""
$(TYPEDEF)

Method to use all four GSV PCA bands to fit 2 mean values
"""
struct FourBandsFittingPoint <: AbstractAlbedoFitting end

"""
$(TYPEDEF)

Method to use 2 GSV PCA bands to fit two flat lines
"""
struct TwoBandsFittingCurve <: AbstractAlbedoFitting end

"""
$(TYPEDEF)

Method to use 2 GSV PCA bands to fit a mean and a flat curve
"""
struct TwoBandsFittingHybrid <: AbstractAlbedoFitting end

"""
$(TYPEDEF)

Method to use 2 GSV PCA bands to fit 2 mean values
"""
struct TwoBandsFittingPoint <: AbstractAlbedoFitting end

###############################################################################
#
# Calculate soil albedo based on color class and soil moisture
#
###############################################################################
"""
    soil_albedos(color::Int, swc::FT) where {FT<:AbstractFloat}

Calculate soil broad band albedos for PAR and NIR, given
- `color` Soil color class, 1~20
- `swc` Soil volumetric water content
"""
function soil_albedos(color::Int, swc::FT) where {FT<:AbstractFloat}
    # soil class must be from 1 to 20
    @assert 1 <= color <=20;

    # calculate albedos for PAR and NIR
    _delta   = max(0, FT(0.11) - FT(0.4) * swc);
    _alb_par::FT = max(SOIL_BNDS[color,1], SOIL_BNDS[color,3] + _delta);
    _alb_nir::FT = max(SOIL_BNDS[color,2], SOIL_BNDS[color,4] + _delta);

    return _alb_par,_alb_nir
end

function soil_albedos(
            color::Int,
            rswc::FT,
            rel::Bool
) where {FT<:AbstractFloat}
    # soil class must be from 1 to 20
    @assert 1 <= color <=20;

    # calculate albedos for PAR and NIR
    _alb_par::FT = SOIL_BNDS[color,1] * (1 - rswc) + rswc * SOIL_BNDS[color,3];
    _alb_nir::FT = SOIL_BNDS[color,2] * (1 - rswc) + rswc * SOIL_BNDS[color,4];

    return _alb_par,_alb_nir
end

"""
    fit_soil_mat!(
                soil::SoilOpticals{FT},
                wls::WaveLengths{FT},
                swc::FT,
                method::AbstractAlbedoFitting;
                clm::Bool = false
    ) where {FT<:AbstractFloat}

Fit soil albedo bands parameters, given
- `soil` [`SoilOpticals`](@ref) type structure
- `wls` [`Wavelengths`](@ref) type structure
- `swc` Soil volumetric water content
- `method` [`AbstractAlbedoFitting`](@ref) type fitting method
- `clm` If true, use CLM method, else use new method
"""
function fit_soil_mat!(
            soil::SoilOpticals{FT},
            wls::WaveLengths{FT},
            swc::FT,
            method::AbstractAlbedoFitting;
            clm::Bool = false
) where {FT<:AbstractFloat}
    @unpack iWLF = wls;

    # fit the curve only if the values mismatch
    if clm
        _ref_par,_ref_nir = soil_albedos(soil.color, swc);
    else
        _ref_par,_ref_nir = soil_albedos(soil.color, swc, true);
    end;
    if soil.ρ_PAR !== _ref_par || soil.ρ_NIR !== _ref_nir
        fit_soil_mat!(soil, wls, _ref_par, _ref_nir, method);
    end

    # update the soil information
    soil.ρ_SW_SIF .= view(soil.ρ_SW, iWLF);
    soil.ε_SW .= 1 .- soil.ρ_SW;

    return nothing
end

"""
    fit_soil_mat!(
                soil::SoilOpticals{FT},
                wls::WaveLengths{FT},
                ref_PAR::FT,
                ref_NIR::FT,
                method::FourBandsFittingCurve
    ) where {FT<:AbstractFloat}

Fit soil albedo bands parameters, given
- `soil` [`SoilOpticals`](@ref) type structure
- `wls` [`Wavelengths`](@ref) type structure
- `ref_PAR` Target albedo for PAR
- `ref_NIR` Target albedo for NIR
- `method` [`FourBandsFittingCurve`](@ref) type fitting method
"""
function fit_soil_mat!(
            soil::SoilOpticals{FT},
            wls::WaveLengths{FT},
            ref_PAR::FT,
            ref_NIR::FT,
            method::FourBandsFittingCurve
) where {FT<:AbstractFloat}
    # update soil PAR and NIR albedo
    soil.ρ_PAR = ref_PAR;
    soil.ρ_NIR = ref_NIR;

    # solve for weights using pinv
    soil.ρ_SW_raw[32:end] .= ref_NIR;
    soil.ρ_SW_raw[1:31] .= ref_PAR;
    soil.SW_vec_4 .= pinv(soil.SW_mat_raw_4) * soil.ρ_SW_raw;

    # update vectors in soil
    mul!(soil.ρ_SW, soil.SW_mat_4, soil.SW_vec_4);

    return nothing
end

"""
    fit_soil_mat!(
                soil::SoilOpticals{FT},
                wls::WaveLengths{FT},
                ref_PAR::FT,
                ref_NIR::FT,
                method::FourBandsFittingHybrid
    ) where {FT<:AbstractFloat}

Fit soil albedo bands parameters, given
- `soil` [`SoilOpticals`](@ref) type structure
- `wls` [`Wavelengths`](@ref) type structure
- `ref_PAR` Target albedo for PAR
- `ref_NIR` Target albedo for NIR
- `method` [`FourBandsFittingHybrid`](@ref) type fitting method
"""
function fit_soil_mat!(
            soil::SoilOpticals{FT},
            wls::WaveLengths{FT},
            ref_PAR::FT,
            ref_NIR::FT,
            method::FourBandsFittingHybrid
) where {FT<:AbstractFloat}
    # update soil PAR and NIR albedo
    soil.ρ_PAR = ref_PAR;
    soil.ρ_NIR = ref_NIR;

    # solve for weights using pinv
    soil.ρ_SW_raw[32:end] .= ref_NIR;
    soil.ρ_SW_raw[1:31] .= ref_PAR;
    soil.SW_vec_4 .= pinv(soil.SW_mat_raw_4) * soil.ρ_SW_raw;

    # solve for weights
    @inline _fit(x::Vector{FT}) where {FT<:AbstractFloat} = (
        mul!(soil.ρ_SW_raw, soil.SW_mat_raw_4, x);
        _diff = ( mean(soil.ρ_SW_raw[1:31]) - ref_PAR ) ^ 2 +
                mean( abs.(soil.ρ_SW_raw[32:end] .- ref_NIR) ) ^ 2;
        return -_diff
    );

    _ms = ReduceStepMethodND{FT}(x_mins = FT[-2,-2,-2,-2],
                                 x_maxs = FT[2,2,2,2],
                                 x_inis = soil.SW_vec_4,
                                 Δ_inis = FT[0.1,0.1,0.1,0.1]);
    _tol = SolutionToleranceND{FT}(FT[0.001,0.001,0.001,0.001], 50);
    _sol = find_peak(_fit, _ms, _tol);
    soil.SW_vec_4 .= _sol;

    # update vectors in soil
    mul!(soil.ρ_SW, soil.SW_mat_4, soil.SW_vec_4);

    return nothing
end

"""
    fit_soil_mat!(
                soil::SoilOpticals{FT},
                wls::WaveLengths{FT},
                ref_PAR::FT,
                ref_NIR::FT,
                method::FourBandsFittingPoint
    ) where {FT<:AbstractFloat}

Fit soil albedo bands parameters, given
- `soil` [`SoilOpticals`](@ref) type structure
- `wls` [`Wavelengths`](@ref) type structure
- `ref_PAR` Target albedo for PAR
- `ref_NIR` Target albedo for NIR
- `method` [`FourBandsFittingPoint`](@ref) type fitting method
"""
function fit_soil_mat!(
            soil::SoilOpticals{FT},
            wls::WaveLengths{FT},
            ref_PAR::FT,
            ref_NIR::FT,
            method::FourBandsFittingPoint
) where {FT<:AbstractFloat}
    # update soil PAR and NIR albedo
    soil.ρ_PAR = ref_PAR;
    soil.ρ_NIR = ref_NIR;

    # solve for weights using pinv
    soil.ρ_SW_raw[32:end] .= ref_NIR;
    soil.ρ_SW_raw[1:31] .= ref_PAR;
    soil.SW_vec_4 .= pinv(soil.SW_mat_raw_4) * soil.ρ_SW_raw;

    # solve for weights
    @inline _fit(x::Vector{FT}) where {FT<:AbstractFloat} = (
        mul!(soil.ρ_SW_raw, soil.SW_mat_raw_4, x);
        _diff = ( mean(soil.ρ_SW_raw[1:31]) - ref_PAR ) ^ 2 +
                ( mean(soil.ρ_SW_raw[32:end]) - ref_NIR ) ^ 2;
        return -_diff
    );

    _ms = ReduceStepMethodND{FT}(x_mins = FT[-2,-2,-2,-2],
                                 x_maxs = FT[2,2,2,2],
                                 x_inis = soil.SW_vec_4,
                                 Δ_inis = FT[0.1,0.1,0.1,0.1]);
    _tol = SolutionToleranceND{FT}(FT[0.001,0.001,0.001,0.001], 50);
    _sol = find_peak(_fit, _ms, _tol);
    soil.SW_vec_4 .= _sol;

    # update vectors in soil
    mul!(soil.ρ_SW, soil.SW_mat_4, soil.SW_vec_4);

    return nothing
end

"""
    fit_soil_mat!(
                soil::SoilOpticals{FT},
                wls::WaveLengths{FT},
                ref_PAR::FT,
                ref_NIR::FT,
                method::TwoBandsFittingCurve
    ) where {FT<:AbstractFloat}

Fit soil albedo bands parameters, given
- `soil` [`SoilOpticals`](@ref) type structure
- `wls` [`Wavelengths`](@ref) type structure
- `ref_PAR` Target albedo for PAR
- `ref_NIR` Target albedo for NIR
- `method` [`TwoBandsFittingCurve`](@ref) type fitting method
"""
function fit_soil_mat!(
            soil::SoilOpticals{FT},
            wls::WaveLengths{FT},
            ref_PAR::FT,
            ref_NIR::FT,
            method::TwoBandsFittingCurve
) where {FT<:AbstractFloat}
    # update soil PAR and NIR albedo
    soil.ρ_PAR = ref_PAR;
    soil.ρ_NIR = ref_NIR;

    # solve for weights using pinv
    soil.ρ_SW_raw[32:end] .= ref_NIR;
    soil.ρ_SW_raw[1:31] .= ref_PAR;
    soil.SW_vec_2 .= pinv(soil.SW_mat_raw_2) * soil.ρ_SW_raw;

    # update vectors in soil
    mul!(soil.ρ_SW, soil.SW_mat_2, soil.SW_vec_2);

    return nothing
end

"""
    fit_soil_mat!(
                soil::SoilOpticals{FT},
                wls::WaveLengths{FT},
                ref_PAR::FT,
                ref_NIR::FT,
                method::TwoBandsFittingHybrid
    ) where {FT<:AbstractFloat}

Fit soil albedo bands parameters, given
- `soil` [`SoilOpticals`](@ref) type structure
- `wls` [`Wavelengths`](@ref) type structure
- `ref_PAR` Target albedo for PAR
- `ref_NIR` Target albedo for NIR
- `method` [`TwoBandsFittingHybrid`](@ref) type fitting method
"""
function fit_soil_mat!(
            soil::SoilOpticals{FT},
            wls::WaveLengths{FT},
            ref_PAR::FT,
            ref_NIR::FT,
            method::TwoBandsFittingHybrid
) where {FT<:AbstractFloat}
    # update soil PAR and NIR albedo
    soil.ρ_PAR = ref_PAR;
    soil.ρ_NIR = ref_NIR;

    # solve for weights using pinv
    soil.ρ_SW_raw[32:end] .= ref_NIR;
    soil.ρ_SW_raw[1:31] .= ref_PAR;
    soil.SW_vec_2 .= pinv(soil.SW_mat_raw_2) * soil.ρ_SW_raw;

    # solve for weights
    @inline _fit(x::Vector{FT}) where {FT<:AbstractFloat} = (
        mul!(soil.ρ_SW_raw, soil.SW_mat_raw_2, x);
        _diff = ( mean(soil.ρ_SW_raw[1:31]) - ref_PAR ) ^ 2 +
                mean( abs.(soil.ρ_SW_raw[32:end] .- ref_NIR) ) ^ 2;
        return -_diff
    );

    _ms = ReduceStepMethodND{FT}(x_mins = FT[-2,-2],
                                 x_maxs = FT[2,2],
                                 x_inis = soil.SW_vec_2,
                                 Δ_inis = FT[0.1,0.1]);
    _tol = SolutionToleranceND{FT}(FT[0.001,0.001], 50);
    _sol = find_peak(_fit, _ms, _tol);
    soil.SW_vec_2 .= _sol;

    # update vectors in soil
    mul!(soil.ρ_SW, soil.SW_mat_2, soil.SW_vec_2);

    return nothing
end

"""
    fit_soil_mat!(
                soil::SoilOpticals{FT},
                wls::WaveLengths{FT},
                ref_PAR::FT,
                ref_NIR::FT,
                method::TwoBandsFittingPoint
    ) where {FT<:AbstractFloat}

Fit soil albedo bands parameters, given
- `soil` [`SoilOpticals`](@ref) type structure
- `wls` [`Wavelengths`](@ref) type structure
- `ref_PAR` Target albedo for PAR
- `ref_NIR` Target albedo for NIR
- `method` [`TwoBandsFittingPoint`](@ref) type fitting method
"""
function fit_soil_mat!(
            soil::SoilOpticals{FT},
            wls::WaveLengths{FT},
            ref_PAR::FT,
            ref_NIR::FT,
            method::TwoBandsFittingPoint
) where {FT<:AbstractFloat}
    @unpack dry_NIR, dry_PAR, wet_NIR, wet_PAR = soil;

    # update soil PAR and NIR albedo
    soil.ρ_PAR = ref_PAR;
    soil.ρ_NIR = ref_NIR;

    # solve for weights
    soil.SW_vec_2[1] = (ref_PAR * wet_NIR - ref_NIR * wet_PAR) /
                       (wet_NIR * dry_PAR - wet_PAR * dry_NIR);
    soil.SW_vec_2[2] = (ref_PAR * dry_NIR - ref_NIR * dry_PAR) /
                       (dry_NIR * wet_PAR - dry_PAR * wet_NIR);

    # update vectors in soil
    mul!(soil.ρ_SW, soil.SW_mat_2, soil.SW_vec_2);

    return nothing
end


function simurecwf(
    model::Union{elmod2d,acmod2d},
    source::Array{MTsource},
    rec::receiver,
    RecDataPath::String,
    Slices::Union{UnitRange{Int64},StepRange{Int64,Int64},Array{Int64,1},Int64},
    WFDataPath::String,
    OptCpnt::Array{String},
    )

    nT = model.medium.nT

    #============  write both Recordings and Wavefields ===========#
    fid_wf = open(RecDataPath, "a+")
    fid_record = open(WFDataPath, "a+")
    # wIB(fid_record,rec,model.medium.nT,OptCpnt)
    # wIB(fid_wf,model,Slices,OptCpnt)
    for it = 1 : nT
        for isn in 1:length(sou)
            addmt!(model.wf, sou[isn], it)
        end
          run!(model)
          wDB(fid_record,fid_wf,model.wf,rec.BDnloc,
            model.nwf,OptCpnt,it,Slices)
    end
    close(fid_wf)
end


module POMCPDPW

using POMDPs
using BasicPOMCP
using ParticleFilters
using Parameters
using MCTS
using CPUTime
using D3Trees
using Colors
using Random
using Printf
using POMDPPolicies

using BasicPOMCP: convert_estimator

import Base: insert!
import POMDPs: action, solve, mean, rand, updater, currentobs, history
import POMDPModelTools: action_info

import MCTS: n_children, next_action, isroot, node_tag, tooltip_tag

export
    POMCPOWSolver,
    POMCPOWPlanner,
    POMCPOWTree,
    POWNodeBelief,
    POWTreeObsNode,
    CategoricalVector,
    FORollout,
    FOValue,

    MaxUCB,
    MaxQ,
    MaxTries,

    init_N,
    init_V,

    n_children,
    belief,
    sr_belief,
    isroot,

    POMCPOWVisualizer,
    blink

const init_V = init_Q

include("categorical_vector.jl")
include("beliefs.jl")

include("tree.jl")
include("criteria.jl")

"""
    POMCPOWSolver

Partially observable Monte Carlo planning with observation widening.

Fields:

- `eps::Float64`:
    Rollouts and tree expansion will stop when discount^depth is less than this.
    default: `0.01`
- `max_depth::Int`:
    Rollouts and tree expension will stop when this depth is reached.
    default: `10`
- `criterion::Any`:
    Criterion to decide which action to take at each node. e.g. `MaxUCB(c)`, `MaxQ`, or `MaxTries`
    default: `MaxUCB(1.0)`
- `final_criterion::Any`:
    Criterion for choosing the action to take after the tree is constructed.
    default: `MaxQ()`
- `tree_queries::Int`:
    Number of iterations during each action() call.
    default: `1000`
- `max_time::Float64`:
    Time limit for planning at each steps (seconds).
    default: `Inf`
- `rng::AbstractRNG`:
    Random number generator.
    default: Base.GLOBAL_RNG
- `node_sr_belief_updater::Updater`:
    Updater for state-reward distribution at the nodes.
    default: `POWNodeFilter()`
- `estimate_value::Any`: (rollout policy can be specified by setting this to RolloutEstimator(policy))
    Function, object, or number used to estimate the value at the leaf nodes.
    If this is a function `f`, `f(pomdp, s, h::BeliefNode, steps)` will be called to estimate the value.
    If this is an object `o`, `estimate_value(o, pomdp, s, h::BeliefNode, steps)` will be called.
    If this is a number, the value will be set to that number
    default: `RolloutEstimator(RandomSolver(rng))`
- `enable_action_pw::Bool`:
    Controls whether progressive widening is done on actions; if `false`, the entire action space is used.
    default: `true`
- `check_repeat_obs::Bool`:
    Check if an observation was sampled multiple times. This has some dictionary maintenance overhead, but prevents multiple nodes with the same observation from being created. If the observation space is discrete, this should probably be used, but can be turned off for speed.
    default: `true`
- `check_repeat_act::Bool`:
    Check if an action was sampled multiple times. This has some dictionary maintenance overhead, but prevents multiple nodes with the same action from being created. If the action space is discrete, this should probably be used, but can be turned off for speed.
    default: `true`
- `tree_in_info::Bool`:
    If `true`, return the tree in the info dict when action_info is called, this can use a lot of memory if histories are being saved.
    default: `false`
- `k_action::Float64`, `alpha_action::Float64`, `k_observation::Float64`, `alpha_observation::Float64`:
        These constants control the double progressive widening. A new observation
        or action will be added if the number of children is less than or equal to kN^alpha.
        defaults: k: `10`, alpha: `0.5`
- `init_V::Any`:
    Function, object, or number used to set the initial V(h,a) value at a new node.
    If this is a function `f`, `f(pomdp, h, a)` will be called to set the value.
    If this is an object `o`, `init_V(o, pomdp, h, a)` will be called.
    If this is a number, V will be set to that number
    default: `0.0`
- `init_N::Any`:
    Function, object, or number used to set the initial N(s,a) value at a new node.
    If this is a function `f`, `f(pomdp, h, a)` will be called to set the value.
    If this is an object `o`, `init_N(o, pomdp, h, a)` will be called.
    If this is a number, N will be set to that number
    default: `0`
- `next_action::Any`
    Function or object used to choose the next action to be considered for progressive widening.
    The next action is determined based on the POMDP, the belief, `b`, and the current `BeliefNode`, `h`.
    If this is a function `f`, `f(pomdp, b, h)` will be called to set the value.
    If this is an object `o`, `next_action(o, pomdp, b, h)` will be called.
    default: `RandomActionGenerator(rng)`
- `default_action::Any`:
    Function, action, or Policy used to determine the action if POMCP fails with exception `ex`.
    If this is a Function `f`, `f(pomdp, belief, ex)` will be called.
    If this is a Policy `p`, `action(p, belief)` will be called.
    If it is an object `a`, `default_action(a, pomdp, belief, ex)` will be called, and
    if this method is not implemented, `a` will be returned directly.

"""
@with_kw mutable struct POMCPOWSolver{RNG<:AbstractRNG} <: AbstractPOMCPSolver
    eps::Float64                = 0.01
    max_depth::Int              = typemax(Int)
    criterion                   = MaxUCB(1.0)
    final_criterion             = MaxQ()
    tree_queries::Int           = 1000
    max_time::Float64           = Inf
    rng::RNG                    = Random.GLOBAL_RNG
    node_sr_belief_updater      = POWNodeFilter()

    estimate_value::Any         = RolloutEstimator(RandomSolver(rng))

    enable_action_pw::Bool      = true
    check_repeat_obs::Bool      = true
    check_repeat_act::Bool      = true
    tree_in_info::Bool          = false

    alpha_observation::Float64  = 0.5
    k_observation::Float64      = 10.0
    alpha_action::Float64       = 0.5
    k_action::Float64           = 10.0
    init_V::Any                 = 0.0
    init_N::Any                 = 0
    next_action::Any            = RandomActionGenerator(rng)
    default_action::Any         = ExceptionRethrow()
end

# unweighted ParticleCollections don't get anything pushed to them
function push_weighted!(::ParticleCollection, sp) end

include("planner2.jl")
include("solver2.jl")

function solve(solver::POMCPOWSolver, problem::POMDP)
    return POMCPOWPlanner(solver, problem)
end

include("updater.jl")
include("visualization.jl")

end # module