text:
module Pipeline

using Statistics
using Random
using JLD2
import Dates

using ..DistributionUtils
using ..ReferenceModels
using ..ReferenceStats
using ..TurbulenceConvectionUtils
using ..NetCDFIO
using ..HelperFuncs
using ..KalmanProcessUtils

# Import EKP modules
using EnsembleKalmanProcesses
using EnsembleKalmanProcesses.ParameterDistributions
import EnsembleKalmanProcesses: update_ensemble!

export init_calibration, ek_update, versioned_model_eval, restart_calibration

"""
    init_calibration(job_id::String, config::Dict{Any, Any})

Initializes a calibration process given a configuration, and a pipeline mode.

Inputs:

 - job_id :: Unique job identifier for sbatch communication.
 - config :: User-defined configuration dictionary.
 - mode :: Whether the calibration process is parallelized through HPC resources or using Julia's pmap.
"""
function init_calibration(config::Dict{Any, Any}; mode::String = "hpc", job_id::String = "12345", config_path = nothing)
    @info "Initialize calibration on $(Dates.now())"
    @assert mode in ["hpc", "pmap"]

    ref_config = config["reference"]
    y_ref_type = ref_config["y_reference_type"]
    batch_size = get_entry(ref_config, "batch_size", nothing)
    kwargs_ref_model = get_ref_model_kwargs(ref_config)

    reg_config = config["regularization"]
    kwargs_ref_stats = get_ref_stats_kwargs(ref_config, reg_config)
    l2_reg = get_entry(reg_config, "l2_reg", nothing)

    out_config = config["output"]
    overwrite_scm_file = get_entry(out_config, "overwrite_scm_file", false)
    outdir_root = get_entry(out_config, "outdir_root", pwd())

    proc_config = config["process"]
    N_ens = proc_config["N_ens"]
    N_iter = proc_config["N_iter"]
    algo_name = get_entry(proc_config, "algorithm", "Inversion")

    Δt_scheduler = get_entry(proc_config, "Δt", 1.0)
    Δt = get_Δt(Δt_scheduler, 1)

    augmented = get_entry(proc_config, "augmented", false)
    failure_handler = get_entry(proc_config, "failure_handler", "high_loss")

    params = config["prior"]["constraints"]
    unc_σ = get_entry(config["prior"], "unconstrained_σ", 1.0)
    prior_μ = get_entry(config["prior"], "prior_mean", nothing)

    namelist_args = get_entry(config["scm"], "namelist_args", nothing)

    val_config = get(config, "validation", nothing)

    # Dimensionality
    n_param = sum(map(length, collect(values(params))))
    if algo_name == "Unscented"
        @assert N_ens == 2 * n_param + 1 "Number of ensemble members must be 2p + 1 in Unscented Kalman Inversion."
    end

    # Minibatch mode
    if !isnothing(batch_size)
        @info "Training using mini-batches."
        ref_model_batch = ReferenceModelBatch(kwargs_ref_model)
        global_ref_models = deepcopy(ref_model_batch.ref_models)
        ref_models, batch_indices = get_minibatch!(ref_model_batch, batch_size)
        ref_stats = ReferenceStatistics(ref_models; kwargs_ref_stats...)
        ref_model_batch = reshuffle_on_epoch_end(ref_model_batch)
        # Generate global reference statistics for diagnostics
        io_ref_stats = ReferenceStatistics(global_ref_models; kwargs_ref_stats...)
        io_ref_models = global_ref_models
    else
        ref_models = construct_reference_models(kwargs_ref_model)
        ref_stats = ReferenceStatistics(ref_models; kwargs_ref_stats...)
        io_ref_stats = ref_stats
        io_ref_models = ref_models
        batch_indices = nothing
    end

    outdir_path = create_output_dir(
        ref_stats,
        outdir_root,
        algo_name,
        Δt,
        n_param,
        N_ens,
        N_iter,
        batch_size,
        config_path,
        y_ref_type,
    )

    if !isnothing(prior_μ)
        @assert collect(keys(params)) == collect(keys(prior_μ))
    else
        prior_μ = nothing
    end
    priors = construct_priors(params, outdir_path = outdir_path, unconstrained_σ = unc_σ, prior_mean = prior_μ)
    # parameters are sampled in unconstrained space
    if algo_name in ["Inversion", "Sampler", "SparseInversion"]
        if algo_name == "Inversion"
            algo = Inversion()
        elseif algo_name == "SparseInversion"
            γ = get_entry(proc_config, "l1_norm_limit", 1 / eps(Float64))
            prune_below = get_entry(proc_config, "prune_below", 0.0)
            sparse_params = get_entry(proc_config, "sparse_params", true)
            sparse_idx = get_sparse_indices(sparse_params, n_param)
            convex_opt_reg = get_entry(proc_config, "convex_opt_reg", 0.0)
            algo = SparseInversion(γ, prune_below, sparse_idx, convex_opt_reg)
        elseif algo_name == "Sampler"
            algo = Sampler(vcat(mean(priors)...), cov(priors))
        end
        initial_params = construct_initial_ensemble(priors, N_ens, rng_seed = rand(1:1000))
        if augmented
            ekobj = generate_tekp(
                ref_stats,
                priors,
                algo,
                initial_params,
                outdir_path = outdir_path,
                l2_reg = l2_reg,
                failure_handler = failure_handler,
            )
        else
            ekobj = generate_ekp(
                ref_stats,
                algo,
                initial_params,
                outdir_path = outdir_path,
                failure_handler = failure_handler,
            )
        end
    elseif algo_name == "Unscented"
        algo = Unscented(vcat(mean(priors)...), cov(priors), α_reg = 1.0, update_freq = 1)
        if augmented
            ekobj = generate_tekp(
                ref_stats,
                priors,
                algo,
                outdir_path = outdir_path,
                l2_reg = l2_reg,
                failure_handler = failure_handler,
            )
        else
            ekobj = generate_ekp(ref_stats, algo, outdir_path = outdir_path, failure_handler = failure_handler)
        end
    end

    params_cons_i = transform_unconstrained_to_constrained(priors, get_u_final(ekobj))
    params = [c[:] for c in eachcol(params_cons_i)]
    mod_evaluators = [ModelEvaluator(param, get_name(priors), ref_models, ref_stats) for param in params]
    versions = generate_scm_input(mod_evaluators, outdir_path, batch_indices)
    # Store version identifiers for this ensemble in a common file
    write_versions(versions, 1, outdir_path = outdir_path)
    # Store ReferenceModelBatch
    if !isnothing(batch_size)
        write_ref_model_batch(ref_model_batch, outdir_path = outdir_path)
    end

    # Initialize validation
    val_ref_models, val_ref_stats =
        isnothing(val_config) ? repeat([nothing], 2) :
        init_validation(
            val_config,
            reg_config,
            ekobj,
            priors,
            versions,
            outdir_path,
            overwrite = overwrite_scm_file,
            namelist_args = namelist_args,
        )

    # Diagnostics IO
    init_diagnostics(config, outdir_path, io_ref_models, io_ref_stats, ekobj, priors, val_ref_models, val_ref_stats)

    if mode == "hpc"
        open("$(job_id).txt", "w") do io
            write(io, "$(outdir_path)\n")
        end
    end
    return outdir_path
end

function get_ref_model_kwargs(ref_config::Dict{Any, Any})
    n_cases = length(ref_config["case_name"])
    Σ_dir = expand_dict_entry(ref_config, "Σ_dir", n_cases)
    Σ_t_start = expand_dict_entry(ref_config, "Σ_t_start", n_cases)
    Σ_t_end = expand_dict_entry(ref_config, "Σ_t_end", n_cases)
    n_obs = expand_dict_entry(ref_config, "n_obs", n_cases)
    scm_path_kwargs = if haskey(ref_config, "scm_dir")
        Dict(:scm_dir => ref_config["scm_dir"])
    elseif all(haskey.(Ref(ref_config), ["scm_parent_dir", "scm_suffix"]))
        Dict(:scm_parent_dir => ref_config["scm_parent_dir"], :scm_suffix => ref_config["scm_suffix"])
    else
        throw(ArgumentError("config must either contain `scm_dir` or both `scm_parent_dir` and `scm_suffix`."))
    end

    rm_kwargs = Dict(
        :y_names => ref_config["y_names"],
        # Reference path specification
        :y_dir => ref_config["y_dir"],
        :Σ_dir => Σ_dir,
        # Case name
        :case_name => ref_config["case_name"],
        # Define observation window (s)
        :t_start => ref_config["t_start"],
        :t_end => ref_config["t_end"],
        :Σ_t_start => Σ_t_start,
        :Σ_t_end => Σ_t_end,
        :n_obs => n_obs,
    )
    merge!(rm_kwargs, scm_path_kwargs)
    n_RM = length(rm_kwargs[:case_name])
    for (k, v) in pairs(rm_kwargs)
        @assert length(v) == n_RM "Entry `$k` in the reference config file has length $(length(v)). Should have length $n_RM."
    end
    return rm_kwargs
end

function get_ref_stats_kwargs(ref_config::Dict{Any, Any}, reg_config::Dict{Any, Any})
    y_ref_type = ref_config["y_reference_type"]
    Σ_ref_type = get_entry(ref_config, "Σ_reference_type", y_ref_type)
    perform_PCA = get_entry(reg_config, "perform_PCA", true)
    variance_loss = get_entry(reg_config, "variance_loss", 1.0e-2)
    normalize = get_entry(reg_config, "normalize", true)
    tikhonov_mode = get_entry(reg_config, "tikhonov_mode", "relative")
    tikhonov_noise = get_entry(reg_config, "tikhonov_noise", 1.0e-6)
    dim_scaling = get_entry(reg_config, "dim_scaling", true)
    return Dict(
        :perform_PCA => perform_PCA,
        :normalize => normalize,
        :variance_loss => variance_loss,
        :tikhonov_noise => tikhonov_noise,
        :tikhonov_mode => tikhonov_mode,
        :dim_scaling => dim_scaling,
        :y_type => y_ref_type,
        :Σ_type => Σ_ref_type,
    )
end

"Create the calibration output directory and copy the config file into it"
function create_output_dir(
    ref_stats::ReferenceStatistics,
    outdir_root::String,
    algo_name::String,
    Δt::FT,
    n_param::IT,
    N_ens::IT,
    N_iter::IT,
    batch_size::Union{IT, Nothing},
    config_path::Union{String, Nothing},
    y_ref_type,
) where {FT <: Real, IT <: Integer}
    # Output path
    d = isnothing(batch_size) ? "d$(pca_length(ref_stats))" : "mb"
    now = Dates.format(Dates.now(), "yyyy-mm-dd_HH-MM")
    suffix = randstring(3)  # ensure output folder is unique
    outdir_path = joinpath(
        outdir_root,
        "results_$(algo_name)_dt_$(Δt)_p$(n_param)_e$(N_ens)_i$(N_iter)_$(d)_$(typeof(y_ref_type))_$(now)_$(suffix)",
    )
    @info "Name of outdir path for this EKP is: $outdir_path"
    mkpath(outdir_path)
    if !isnothing(config_path)
        cp(config_path, joinpath(outdir_path, "config.jl"))
    end
    return outdir_path
end

"Initialize the validation process."
function init_validation(
    val_config::Dict{Any, Any},
    reg_config::Dict{Any, Any},
    ekp::EnsembleKalmanProcess,
    priors::ParameterDistribution,
    versions::Vector{IT},
    outdir_path::String;
    overwrite::Bool = false,
    namelist_args = nothing,
) where {FT <: Real, IT <: Integer}

    kwargs_ref_model = get_ref_model_kwargs(val_config)
    kwargs_ref_stats = get_ref_stats_kwargs(val_config, reg_config)
    batch_size = get_entry(val_config, "batch_size", nothing)

    if !isnothing(batch_size)
        @info "Validation using mini-batches."
        ref_model_batch = ReferenceModelBatch(kwargs_ref_model)
        ref_models, batch_indices = get_minibatch!(ref_model_batch, batch_size)
        ref_model_batch = reshuffle_on_epoch_end(ref_model_batch)
        write_val_ref_model_batch(ref_model_batch, outdir_path = outdir_path)
    else
        ref_models = construct_reference_models(kwargs_ref_model)
        batch_indices = nothing
    end
    ref_stats = ReferenceStatistics(ref_models; kwargs_ref_stats...)
    params_cons_i = transform_unconstrained_to_constrained(priors, get_u_final(ekp))
    params = [c[:] for c in eachcol(params_cons_i)]
    mod_evaluators = [ModelEvaluator(param, get_name(priors), ref_models, ref_stats) for param in params]
    [
        jldsave(scm_val_init_path(outdir_path, version); model_evaluator, version, batch_indices) for
        (model_evaluator, version) in zip(mod_evaluators, versions)
    ]
    return ref_models, ref_stats
end

"""
    update_validation(
        val_config::Dict{Any, Any},
        reg_config::Dict{Any, Any},
        ekp_old::EnsembleKalmanProcess,
        priors::ParameterDistribution,
        versions::Vector{String},
        outdir_path::String,
        iteration::IT
        )

Updates the validation diagnostics and writes to file the validation ModelEvaluators
for the next calibration step.

Inputs:

 - val_config    :: Validation model configuration.
 - reg_config    :: Regularization configuration.
 - ekp_old       :: EnsembleKalmanProcess updated using the past forward model evaluations.
 - priors        :: The priors over parameter space.
 - versions      :: String versions identifying the forward model evaluations.
 - outdir_path   :: Output path directory.
"""
function update_validation(
    val_config::Dict{Any, Any},
    reg_config::Dict{Any, Any},
    ekp_old::EnsembleKalmanProcess,
    priors::ParameterDistribution,
    versions::Vector{String},
    outdir_path::String,
    iteration::IT,
) where {IT <: Integer}

    batch_size = get_entry(val_config, "batch_size", nothing)

    if !isnothing(batch_size)
        ref_model_batch = load(joinpath(outdir_path, "val_ref_model_batch.jld2"))["ref_model_batch"]
        ref_models, batch_indices = get_minibatch!(ref_model_batch, batch_size)
        ref_model_batch = reshuffle_on_epoch_end(ref_model_batch)
        kwargs_ref_stats = get_ref_stats_kwargs(val_config, reg_config)
        ref_stats = ReferenceStatistics(ref_models; kwargs_ref_stats...)
        rm(joinpath(outdir_path, "val_ref_model_batch.jld2"))
        write_val_ref_model_batch(ref_model_batch, outdir_path = outdir_path)
    else
        mod_evaluator = load(scm_val_output_path(outdir_path, versions[1]))["model_evaluator"]
        ref_models = mod_evaluator.ref_models
        ref_stats = mod_evaluator.ref_stats
        batch_indices = nothing
    end
    params_cons_i = transform_unconstrained_to_constrained(priors, get_u_final(ekp_old))
    params = [c[:] for c in eachcol(params_cons_i)]
    mod_evaluators = [ModelEvaluator(param, get_name(priors), ref_models, ref_stats) for param in params]
    # Save new ModelEvaluators using the new versions
    versions = readlines(joinpath(outdir_path, "versions_$(iteration + 1).txt"))
    [
        jldsave(scm_val_init_path(outdir_path, version); model_evaluator, version, batch_indices) for
        (model_evaluator, version) in zip(mod_evaluators, versions)
    ]
    return
end

"""
    ek_update(
        ekobj::EnsembleKalmanProcess,
        priors::ParameterDistribution,
        iteration::Int64,
        config::Dict{Any, Any},
        versions::Vector{String},
        outdir_path::String,
    )

Updates an EnsembleKalmanProcess using forward model evaluations stored
in output files defined by their `versions`, and generates the parameters
for the next ensemble for forward model evaluations. The updated EnsembleKalmanProcess
and new ModelEvaluators are both written to file.

Inputs:

 - ekobj         :: EnsembleKalmanProcess to be updated.
 - priors        :: Priors over parameters, used for unconstrained-constrained mappings.
 - iteration     :: Current iteration of the calibration process.
 - config        :: Process configuration dictionary.
 - versions      :: String versions identifying the forward model evaluations.
 - outdir_path   :: Output path directory.
"""
function ek_update(
    ekobj::EnsembleKalmanProcess,
    priors::ParameterDistribution,
    iteration::Int64,
    config::Dict{Any, Any},
    versions::Vector{String},
    outdir_path::String,
)
    @info "Update the Ensemble Kalman Process at iteration $iteration. $(Dates.now())"
    @info "Output is stored in: $outdir_path"
    # Get config
    proc_config = config["process"]
    N_iter = proc_config["N_iter"]
    algo_name = get_entry(proc_config, "algorithm", "Inversion")

    Δt_scheduler = get_entry(proc_config, "Δt", 1.0)
    Δt = get_Δt(Δt_scheduler, iteration)

    deterministic_forward_map = get_entry(proc_config, "noisy_obs", false)
    augmented = get_entry(proc_config, "augmented", false)

    ref_config = config["reference"]
    batch_size = get_entry(ref_config, "batch_size", nothing)

    val_config = get(config, "validation", nothing)

    reg_config = config["regularization"]
    l2_reg = get_entry(reg_config, "l2_reg", nothing)

    scm_args = load(scm_output_path(outdir_path, versions[1]))
    mod_evaluator = scm_args["model_evaluator"]
    batch_indices = scm_args["batch_indices"]
    ref_stats = mod_evaluator.ref_stats
    ref_models = mod_evaluator.ref_models

    # Advance EKP
    if augmented
        g, g_full = get_ensemble_g_eval_aug(outdir_path, versions, priors, l2_reg)
    else
        g, g_full = get_ensemble_g_eval(outdir_path, versions)
    end

    if isa(ekobj.process, Inversion) || isa(ekobj.process, SparseInversion)
        update_ensemble!(ekobj, g, Δt_new = Δt, deterministic_forward_map = deterministic_forward_map)
    elseif isa(ekobj.process, Unscented)
        update_ensemble!(ekobj, g, Δt_new = Δt)
    else
        Δt ≈ 1.0 ? nothing : @warn "Ensemble Kalman Sampler does not accept a custom Δt."
        update_ensemble!(ekobj, g)
    end

    # Diagnostics IO
    if !isnothing(val_config)
        update_val_diagnostics(outdir_path, ekobj, priors, versions, augmented, l2_reg)
    end
    update_diagnostics(outdir_path, ekobj, priors, ref_stats, g_full, versions, batch_indices)

    if iteration < N_iter
        # Prepare updated EKP and ReferenceModelBatch if minibatching.
        if !isnothing(batch_size)
            ref_model_batch = load(joinpath(outdir_path, "ref_model_batch.jld2"))["ref_model_batch"]
            ekp, ref_models, ref_stats, ref_model_batch, batch_indices =
                update_minibatch_inverse_problem(ref_model_batch, ekobj, priors, batch_size, outdir_path, config)
            rm(joinpath(outdir_path, "ref_model_batch.jld2"))
            write_ref_model_batch(ref_model_batch, outdir_path = outdir_path)
        else
            ekp = ekobj
        end

        # Write to file new EKP and ModelEvaluators
        jldsave(ekobj_path(outdir_path, iteration + 1); ekp)
        write_model_evaluators(ekp, priors, ref_models, ref_stats, outdir_path, iteration, batch_indices)

        # Update validation ModelEvaluators
        if !isnothing(val_config)
            reg_config = config["regularization"]
            update_validation(val_config, reg_config, ekobj, priors, versions, outdir_path, iteration)
        end
    end
    # Clean up
    for version in versions
        rm(scm_output_path(outdir_path, version))
        !isnothing(val_config) ? rm(scm_val_output_path(outdir_path, version)) : nothing
    end
    return
end

"""
    restart_calibration(
        ekobj::EnsembleKalmanProcess,
        priors::ParameterDistribution,
        last_iteration::Int64,
        config::Dict{Any, Any},
        outdir_path::String,
    )

Restarts a calibration process from an EnsembleKalmanProcess, the parameter priors and the
calibration process config file. If batching, it requires access to the last ReferenceModelBatch,
stored in the results directory of the previous calibration, `outdir_path`.

Writes to file the ModelEvaluators necessary to continue the calibration process.

Inputs:

 - ekobj          :: EnsembleKalmanProcess to be updated.
 - priors         :: Priors over parameters, used for unconstrained-constrained mappings.
 - last_iteration :: Last iteration of the calibration process to be restarted.
 - config         :: Configuration dictionary.
 - outdir_path    :: Output path directory of the calibration process to be restarted.
 - mode :: Whether the calibration process is parallelized through HPC resources or using Julia's pmap.
 - job_id :: Unique job identifier for sbatch communication.
"""
function restart_calibration(
    ekobj::EnsembleKalmanProcess,
    priors::ParameterDistribution,
    last_iteration::Int64,
    config::Dict{Any, Any},
    outdir_path::String;
    mode::String = "hpc",
    job_id::String = "12345",
)
    @info "Restart calibration on $(Dates.now())"
    # Get config
    ref_config = config["reference"]
    batch_size = get_entry(ref_config, "batch_size", nothing)
    kwargs_ref_model = get_ref_model_kwargs(ref_config)

    reg_config = config["regularization"]
    kwargs_ref_stats = get_ref_stats_kwargs(ref_config, reg_config)

    out_config = config["output"]
    overwrite_scm_file = get_entry(out_config, "overwrite_scm_file", false)

    namelist_args = get_entry(config["scm"], "namelist_args", nothing)

    val_config = get(config, "validation", nothing)

    # Prepare updated EKP and ReferenceModelBatch if minibatching.
    if !isnothing(batch_size)
        ref_model_batch = load(joinpath(outdir_path, "ref_model_batch.jld2"))["ref_model_batch"]
        global_ref_models = deepcopy(ref_model_batch.ref_models)
        ekp, ref_models, ref_stats, ref_model_batch, batch_indices =
            update_minibatch_inverse_problem(ref_model_batch, ekobj, priors, batch_size, outdir_path, config)
        rm(joinpath(outdir_path, "ref_model_batch.jld2"))
        write_ref_model_batch(ref_model_batch, outdir_path = outdir_path)
    else
        ekp = ekobj
        ref_models = construct_reference_models(kwargs_ref_model)
        ref_stats = ReferenceStatistics(ref_models; kwargs_ref_stats...)
        batch_indices = nothing
    end

    # Write to file new EKP and ModelEvaluators
    jldsave(ekobj_path(outdir_path, last_iteration + 1); ekp)
    write_model_evaluators(ekp, priors, ref_models, ref_stats, outdir_path, last_iteration, batch_indices)

    # Restart validation
    if !isnothing(val_config)
        restart_validation(
            val_config,
            reg_config,
            ekobj,
            priors,
            outdir_path,
            last_iteration,
            overwrite = overwrite_scm_file,
            namelist_args = namelist_args,
        )
    end
    if mode == "hpc"
        open("$(job_id).txt", "w") do io
            write(io, "$(outdir_path)\n")
        end
    end
    return outdir_path
end

"""
    restart_validation(
        val_config::Dict{Any, Any},
        reg_config::Dict{Any, Any},
        ekp_old::EnsembleKalmanProcess,
        priors::ParameterDistribution,
        outdir_path::String,
        last_iteration::IT;
        overwrite::Bool = false,
        namelist_args = nothing,
    ) where {IT <: Integer}

Restarts a validation process by writing to file the validation ModelEvaluators necessary to continue
the validation process. If batching, it requires access to the last validation ReferenceModelBatch,
stored in the results directory of the previous calibration process, `outdir_path`.

Inputs:

 - val_config    :: Validation model configuration.
 - reg_config    :: Regularization configuration.
 - ekp_old       :: EnsembleKalmanProcess updated using the past forward model evaluations.
 - priors        :: The priors over parameter space.
 - outdir_path   :: Output path directory.
"""
function restart_validation(
    val_config::Dict{Any, Any},
    reg_config::Dict{Any, Any},
    ekp_old::EnsembleKalmanProcess,
    priors::ParameterDistribution,
    outdir_path::String,
    last_iteration::IT;
    overwrite::Bool = false,
    namelist_args = nothing,
) where {IT <: Integer}

    kwargs_ref_model = get_ref_model_kwargs(val_config)
    kwargs_ref_stats = get_ref_stats_kwargs(val_config, reg_config)
    batch_size = get_entry(val_config, "batch_size", nothing)

    if !isnothing(batch_size)
        ref_model_batch = load(joinpath(outdir_path, "val_ref_model_batch.jld2"))["ref_model_batch"]
        ref_models, batch_indices = get_minibatch!(ref_model_batch, batch_size)
        ref_model_batch = reshuffle_on_epoch_end(ref_model_batch)
        rm(joinpath(outdir_path, "val_ref_model_batch.jld2"))
        write_val_ref_model_batch(ref_model_batch, outdir_path = outdir_path)
    else
        ref_models = construct_reference_models(kwargs_ref_model)
        batch_indices = nothing
    end
    ref_stats = ReferenceStatistics(ref_models; kwargs_ref_stats...)

    params_cons_i = transform_unconstrained_to_constrained(priors, get_u_final(ekp_old))
    params = [c[:] for c in eachcol(params_cons_i)]
    mod_evaluators = [ModelEvaluator(param, get_name(priors), ref_models, ref_stats) for param in params]
    # Save new ModelEvaluators using the new versions
    versions = readlines(joinpath(outdir_path, "versions_$(last_iteration + 1).txt"))
    [
        jldsave(scm_val_init_path(outdir_path, version); model_evaluator, version, batch_indices) for
        (model_evaluator, version) in zip(mod_evaluators, versions)
    ]
    return
end

"""
    get_ensemble_g_eval(outdir_path::String, versions::Vector{String}; validation::Bool = false)

Recovers forward model evaluations from the particle ensemble stored in jld2 files,
after which the files are deleted from disk.

Inputs:
 - `outdir_path`  :: Path to output directory.
 - `versions`     :: Version identifiers of the files containing forward model evaluations.
 - `validation`   :: Whether the function should return training or validation forward model evaluations.
Outputs:
 - `g`            :: Forward model evaluations in the reduced space of the inverse problem.
 - `g_full`       :: Forward model evaluations in the original physical space.
"""
function get_ensemble_g_eval(outdir_path::String, versions::Vector{String}; validation::Bool = false)
    # Find train/validation path
    scm_path(x) = validation ? scm_val_output_path(outdir_path, x) : scm_output_path(outdir_path, x)
    # Get array sizes with first file
    scm_outputs = load(scm_path(first(versions)))
    d = length(scm_outputs["g_scm_pca"])
    d_full = length(scm_outputs["g_scm"])
    N_ens = length(versions)
    g = zeros(d, N_ens)
    g_full = zeros(d_full, N_ens)
    for (ens_index, version) in enumerate(versions)
        scm_outputs = load(scm_path(version))
        g[:, ens_index] = scm_outputs["g_scm_pca"]
        g_full[:, ens_index] = scm_outputs["g_scm"]
    end
    return g, g_full
end

"""
    get_ensemble_g_eval_aug(
        outdir_path::String,
        versions::Vector{String},
        priors::ParameterDistribution,
        l2_reg::Union{Dict{String, Vector{R}}, R};
        validation::Bool = false,
    ) where {R}

Recovers forward model evaluations from the particle ensemble stored in jld2 files,
and augments the projected output state with the input parameters in unconstrained form
to enable regularization.

Inputs:
 - `outdir_path`  :: Path to output directory.
 - `versions`     :: Version identifiers of the files containing forward model evaluations.
 - `priors`       :: Parameter priors, used to transform between constrained and unconstrained spaces.
 - `l2_reg`      :: The config entry specifying l2 regularization.
 - `validation`   :: Whether the function should return training or validation forward model evaluations.
Outputs:
 - `g_aug`        :: Forward model evaluations in the reduced space of the inverse problem, augmented with
                    the unconstrained input parameters.
 - `g_full`       :: Forward model evaluations in the original physical space.
"""
function get_ensemble_g_eval_aug(
    outdir_path::String,
    versions::Vector{String},
    priors::ParameterDistribution,
    l2_reg::Union{Dict{String, Vector{R}}, R};
    validation::Bool = false,
) where {R}
    # Find train/validation path
    scm_path(x) = validation ? scm_val_output_path(outdir_path, x) : scm_output_path(outdir_path, x)
    # Get array sizes with first file
    scm_outputs = load(scm_path(first(versions)))

    aug_indices =
        isa(l2_reg, Dict) ? get_regularized_indices(l2_reg) : 1:length(scm_outputs["model_evaluator"].param_cons)

    # Set dimensionality
    d = length(scm_outputs["g_scm_pca"])
    d_aug = d + length(aug_indices)
    d_full = length(scm_outputs["g_scm"])
    N_ens = length(versions)
    g_aug = zeros(d_aug, N_ens)
    g_full = zeros(d_full, N_ens)

    for (ens_index, version) in enumerate(versions)
        scm_outputs = load(scm_path(version))
        g_aug[1:d, ens_index] = scm_outputs["g_scm_pca"]
        g_full[:, ens_index] = scm_outputs["g_scm"]
        θ = transform_constrained_to_unconstrained(priors, scm_outputs["model_evaluator"].param_cons)
        g_aug[(d + 1):d_aug, ens_index] = θ[aug_indices]
    end
    return g_aug, g_full
end

"""
   versioned_model_eval(version::Union{String, Int}, outdir_path::String, mode::String, config::Dict{Any, Any})

Performs or omits a model evaluation given the parsed mode and provided config,
 and writes to file the model output.

Inputs:
 - version       :: The version associated with the ModelEvaluator to be used.
 - outdir_path   :: The path to the results directory of the calibration process.
 - mode          :: Whether the ModelEvaluator is used for training or validation.
 - config        :: The general configuration dictionary.
"""
function versioned_model_eval(version::Union{String, Int}, outdir_path::String, mode::String, config::Dict{Any, Any})
    @assert mode in ["train", "validation"]
    # Omits validation if unsolicited
    if mode == "validation" && isnothing(get(config, "validation", nothing))
        return
    elseif mode == "validation"
        input_path = scm_val_init_path(outdir_path, version)
        output_path = scm_val_output_path(outdir_path, version)
    else
        input_path = scm_init_path(outdir_path, version)
        output_path = scm_output_path(outdir_path, version)
    end
    # Load inputs
    scm_args = load(input_path)
    namelist_args = get_entry(config["scm"], "namelist_args", nothing)
    failure_handler = get_entry(config["process"], "failure_handler", "high_loss")
    # Check consistent failure method for given algorithm
    @assert failure_handler == "sample_succ_gauss" ? config["process"]["algorithm"] != "Sampler" : true
    model_evaluator = scm_args["model_evaluator"]
    batch_indices = scm_args["batch_indices"]
    # Eval
    sim_dirs, g_scm, g_scm_pca =
        run_SCM(model_evaluator, namelist_args = namelist_args, failure_handler = failure_handler)
    # Store output and delete input
    jldsave(output_path; sim_dirs, g_scm, g_scm_pca, model_evaluator, version, batch_indices)
    rm(input_path)
end

"""
    update_minibatch_inverse_problem(
        rm_batch::ReferenceModelBatch,
        ekp_old::EnsembleKalmanProcess,
        priors::ParameterDistribution,
        batch_size::Integer,
        outdir_path::String,
        config::Dict{Any, Any},
    )

Returns the EnsembleKalmanProcess and ReferenceStatistics consistent with the
new ReferenceModel minibatch, and updates the evaluation order of the ReferenceModelBatch.

Inputs:
 - rm_batch    :: The global ReferenceModelBatch with the current model evaluation order.
 - ekp_old     :: The EnsembleKalmanProcess from the previous minibatch evaluation.
 - batch_size  :: The batch size of the current minibatch.
 - outdir_path :: The output directory.
 - config      :: The configuration dictionary.
Outputs:
 - ekp             :: The EnsembleKalmanProcess for the current minibatch.
 - ref_models      :: The current minibatch of ReferenceModels.
 - ref_stats       :: The ReferenceStatistics consistent with the current minibatch.
 - ref_model_batch :: The global ReferenceModelBatch with the updated model evaluation order.
"""
function update_minibatch_inverse_problem(
    rm_batch::ReferenceModelBatch,
    ekp_old::EnsembleKalmanProcess,
    priors::ParameterDistribution,
    batch_size::Integer,
    outdir_path::String,
    config::Dict{Any, Any},
)
    # Construct new reference minibatch, new ref_stats, and new ekp
    ref_model_batch = deepcopy(rm_batch)
    ref_models, batch_indices = get_minibatch!(ref_model_batch, batch_size)
    ref_model_batch = reshuffle_on_epoch_end(ref_model_batch)

    ref_config = config["reference"]
    reg_config = config["regularization"]
    proc_config = config["process"]

    augmented = get_entry(proc_config, "augmented", false)
    failure_handler = get_entry(proc_config, "failure_handler", "high_loss")
    l2_reg = get_entry(reg_config, "l2_reg", nothing)
    kwargs_ref_stats = get_ref_stats_kwargs(ref_config, reg_config)
    ref_stats = ReferenceStatistics(ref_models; kwargs_ref_stats...)
    process = ekp_old.process

    if isa(process, Unscented)
        # Reconstruct UKI using regularization toward the prior
        algo = Unscented(
            process.u_mean[end],
            process.uu_cov[end],
            α_reg = 1.0,
            update_freq = 1,
            prior_mean = vcat(mean(priors)...),
        )
        if augmented
            ekp = generate_tekp(
                ref_stats,
                priors,
                algo,
                outdir_path = outdir_path,
                l2_reg = l2_reg,
                failure_handler = failure_handler,
            )
        else
            ekp = generate_ekp(ref_stats, algo, outdir_path = outdir_path, failure_handler = failure_handler)
        end
    else
        if augmented
            ekp = generate_tekp(
                ref_stats,
                priors,
                process,
                get_u_final(ekp_old),
                outdir_path = outdir_path,
                l2_reg = l2_reg,
                failure_handler = failure_handler,
            )
        else
            ekp = generate_ekp(
                ref_stats,
                process,
                get_u_final(ekp_old),
                failure_handler = failure_handler,
                outdir_path = outdir_path,
            )
        end
    end
    return ekp, ref_models, ref_stats, ref_model_batch, batch_indices
end

"""
    write_model_evaluators(
        ekp::EnsembleKalmanProcess,
        priors::ParameterDistribution,
        ref_models::Vector{ReferenceModel},
        ref_stats::ReferenceStatistics,
        outdir_path::String,
        iteration::Int,
    )

Creates and writes to file the ModelEvaluators for the current particle ensemble.

Inputs:
 - `ekp`         :: The EnsembleKalmanProcess with the current ensemble of parameter values.
 - `priors`      :: The parameter priors.
 - `ref_models`  :: The ReferenceModels defining the new model evaluations.
 - `ref_stats`   :: The ReferenceStatistics corresponding to passed `ref_models`.
 - `outdir_path` :: The output directory.
 - `iteration`   :: The current process iteration.
"""
function write_model_evaluators(
    ekp::EnsembleKalmanProcess,
    priors::ParameterDistribution,
    ref_models::Vector{ReferenceModel},
    ref_stats::ReferenceStatistics,
    outdir_path::String,
    iteration::Int,
    batch_indices::Union{Vector{Int}, Nothing} = nothing,
)
    params_cons_i = transform_unconstrained_to_constrained(priors, get_u_final(ekp))
    params = [c[:] for c in eachcol(params_cons_i)]
    mod_evaluators = [ModelEvaluator(param, get_name(priors), ref_models, ref_stats) for param in params]
    versions = generate_scm_input(mod_evaluators, outdir_path, batch_indices)
    # Store version identifiers for this ensemble in a common file
    write_versions(versions, iteration + 1, outdir_path = outdir_path)
    return
end

"""
    init_diagnostics(
        config::Dict{Any, Any},
        outdir_path::String,
        ref_stats::ReferenceStatistics,
        ekp::EnsembleKalmanProcess,
        priors::ParameterDistribution,
    )

Creates a diagnostics netcdf file.

Inputs:

 - config :: User-defined configuration dictionary.
 - outdir_path :: Path of results directory.
 - ref_stats :: ReferenceStatistics.
 - ekp :: Initial EnsembleKalmanProcess, containing parameter information, but no forward model evaluations.
 - priors:: Prior distributions of the parameters.
"""
function init_diagnostics(
    config::Dict{Any, Any},
    outdir_path::String,
    ref_models::Vector{ReferenceModel},
    ref_stats::ReferenceStatistics,
    ekp::EnsembleKalmanProcess,
    priors::ParameterDistribution,
    val_ref_models::Union{Vector{ReferenceModel}, Nothing} = nothing,
    val_ref_stats::Union{ReferenceStatistics, Nothing} = nothing,
)
    write_full_stats = get_entry(config["reference"], "write_full_stats", true)
    N_ens = size(get_u_final(ekp), 2)
    diags = NetCDFIO_Diags(config, outdir_path, ref_stats, N_ens, priors, val_ref_stats)
    # Write prior and reference diagnostics
    io_prior(diags, priors)
    io_reference(diags, ref_stats, ref_models, write_full_stats)
    # Add diags, write first state diags
    init_iteration_io(diags)
    init_metrics(diags)
    init_ensemble_diags(diags, ekp, priors)
    init_particle_diags(diags, ekp, priors)
    if !isnothing(val_ref_models)
        write_full_stats = get_entry(config["validation"], "write_full_stats", true)
        init_val_diagnostics(diags, val_ref_stats, val_ref_models, write_full_stats)
    end
end

"""
    update_diagnostics(outdir_path::String, ekp::EnsembleKalmanProcess, priors::ParameterDistribution)

Appends diagnostics of the current iteration evaluations (i.e., forward model output metrics)
and the next iteration state (i.e., parameters and parameter metrics) to a diagnostics netcdf file.

Inputs:

 - outdir_path :: Path of results directory.
 - ekp :: Current EnsembleKalmanProcess.
 - priors:: Prior distributions of the parameters.
 - ref_stats :: ReferenceStatistics.
 - g_full :: The forward model evaluation in primitive space.
 - versions :: Version identifiers of the forward model evaluations at the current iteration.
 - val_config :: The validation configuration, if given.
 - batch_indices :: The indices of the ReferenceModels used in the current batch.
"""
function update_diagnostics(
    outdir_path::String,
    ekp::EnsembleKalmanProcess,
    priors::ParameterDistribution,
    ref_stats::ReferenceStatistics,
    g_full::Array{FT, 2},
    versions::Union{Vector{Int}, Vector{String}},
    batch_indices::Union{Vector{Int}, Nothing} = nothing,
) where {FT <: Real}

    mse_full = compute_mse(g_full, ref_stats.y_full)
    diags = NetCDFIO_Diags(joinpath(outdir_path, "Diagnostics.nc"))

    io_diagnostics(diags, ekp, priors, mse_full, g_full, batch_indices)
end

function update_val_diagnostics(
    outdir_path::String,
    ekp::EnsembleKalmanProcess,
    priors::ParameterDistribution,
    versions::Union{Vector{Int}, Vector{String}},
    augmented::Bool,
    l2_reg,
) where {FT <: Real}
    scm_args = load(scm_val_output_path(outdir_path, versions[1]))
    mod_evaluator = scm_args["model_evaluator"]
    val_batch_indices = scm_args["batch_indices"]
    val_ref_stats = mod_evaluator.ref_stats
    if augmented
        g, g_full = get_ensemble_g_eval_aug(outdir_path, versions, priors, l2_reg, validation = true)
    else
        g, g_full = get_ensemble_g_eval(outdir_path, versions, validation = true)
    end
    # Compute diagnostics
    mse_full = compute_mse(g_full, val_ref_stats.y_full)
    diags = NetCDFIO_Diags(joinpath(outdir_path, "Diagnostics.nc"))

    io_val_diagnostics(diags, ekp, mse_full, g, g_full, val_ref_stats, val_batch_indices)
end

end # module


text:
include(joinpath(@__DIR__, "..", "CoordTrans", "CoordTrans.jl"))

using .CoordTrans

if length(ARGS) == 3
    in_filename = ARGS[1]
    out_filename = ARGS[2]
    wgt_filename = ARGS[3]
else
    println("ARGS does not equal to 3. Maybe parameters have been set already. Keep going...")
end

CoordTrans.convertFile(
    in_filename,
    out_filename,
    wgt_filename,
    varnames=("T_ML",);
    xdim = "Nx",
    ydim = "Ny",
    tdim = "time",
)

ext: jl

max_stars_repo_name: meteorologytoday/CESM-diagnostic

max_issues_repo_name: meteorologytoday/CESM-diagnostic

max_forks_repo_name: meteorologytoday/CESM-diagnostic


text:
using SDE
using Cubature
import SDE.b, SDE.sigma, SDE.a

srand(9)
if !isdefined(:MvTest)

type MvTest <: MvPro
    B
    beta
    Sigma
    A
    d::Int
end
end

b(t,x, P::MvTest) =   exp(-0.2*t)*P.B*x + P.beta
sigma(t,x, P::MvTest) = exp(-0.1*(t))*P.Sigma
a(t,x,  P::MvTest) = exp(-0.2*(t))*P.A

const d = 2

const Q95 = sqrt(2.)*erfinv(0.95)
function mc(M)
  m = mean(M)
  ste = std(M)/sqrt(length(M))
  round([m, Q95*ste], int(2-log(10,Q95*ste)))
end

@Test.test 5.0 == soft(tofs(5., 1., 10.,), 1.,10.)

const T =  1.7
const tmax = T
const tmin = 0.

const B0 = [ -2.21225   -0.49590;   0.631753  -1.34795]
const beta0 = [0.5,-0.5]
const sigma0 = .1* [1.5*[1, 1]  1.5*[1,-1]]
ph0(t,s ) =  5.*exp(-0.2*s)-5.*exp(-0.2*t)

begin

    MvTest() = MvTest(B0, beta0,sigma0, sigma0*sigma0', 2)

    u = [1.,0.]
    v = [0.5, -0.2]

    #taken slightly off
    B = 1.2exp(-0.2*T)*B0

    N = 100

    P = MvTest()
    Pt = MvLinPro(B, beta0, sigma(T, v, P))

    tt = linspace(0., T, N)
    dt = (tt[N]-tt[1])/(N-1)

    ss = linspace(0.,T, N)
    ds = (ss[N]-ss[1])/(N-1)

    ttofss = map(s -> tofs(s, 0, T),ss)

    M = 500

    Yll = zeros(M)
    Y2ll = zeros(M)
    Ull = zeros(M)
    Ull2 = zeros(M)

    Yt0 = zeros(2, M)
    Y2t0 = zeros(2, M)
    YofUt0 = zeros(2, M)
    Us0 = zeros(2, M)

    nt = int(floor(5N/10))

    t0 =  dt*(nt-1)
    W = sample(tt, Wiener(d))
    @time for i in 1:M
	    resample!(W, Wiener(d))
	    Y = guidedeuler(u, W, T, v, Pt, P)
	    ll = llikeliXcirc(Y, Pt, P)
	    Yt0[:,i] = Y.yy[:,nt]
	    Yll[i] = ll
    end

    println("X° at ", round(t0,3), "($nt): ", "[", repr(mc(Yt0[1,:])), ", ",repr(mc(Yt0[2,:])), "]"	)

    x1 = zeros(2, M)
    s0 = soft(t0, 0, T)
    ns = 1 + int(floor(s0/ds))
    s0 = ds*(ns-1)

    u0 = uofx(0.0, u,  T, v, Pt)
    W = sample(ss, Wiener(d))
    U = eulerU(u0, W, 0., T, v, Pt, P)
    srand(9)
    @time for i in 1:M
 	    resample!(W, Wiener(d))
	    eulerU!(U, u0, W, 0., T, v, Pt, P)
	    ll =  llikeliU(U, 0., T, v, Pt, P)
	    Us0[:,i] =  U.yy[:,ns]
	    YofUt0[:,i] =  xofu(s0, U.yy[:,ns],  T, v, Pt)
	    Ull[i] = ll
    end
    srand(9)
    #Profile.init(10^7, 0.0003)
    #@profile  begin
    @time begin
    Phim = [ Phims(ss[i], T, Pt.B) for i in 1:N-1] #precompute Phim
    U2 = eulerU(u0, W, 0., T, v, Pt, P, Phim)
    for i in 1:M
 	    resample!(W, Wiener(d))
	    eulerU!(U2, u0, W, 0., T, v, Pt, P, Phim)
	    ll2 = llikeliU(U2, 0., T, v, Pt, P, Phim)
	    Us0[:,i] =  U2.yy[:,ns]
	    xofu(s0, U2.yy[:,ns],  T, v, Pt, Phim[ns])
	    Ull2[i] = ll2
    end
    end

    println("diff normal/precompiled: ", norm(Ull .- Ull2))

    println("X°(U) at ", round(tofs(s0,0., T),3), "($ns): ",  "[", repr(mc(YofUt0[1,:])), ",", repr(mc(YofUt0[2,:])), "]")

    W = sample(ttofss, Wiener(d))
    @time for i in 1:M
	    resample!(W, Wiener(d))
	    Y2 = guidedeuler(u, W, T, v, Pt, P)
	    ll = llikeliXcirc(Y2, Pt, P)
	    Y2t0[:,i] = Y2.yy[:,ns]
	    Y2ll[i] = ll
    end

    println("X°° at ", round(ttofss[ns],3), "($ns): ", "[", repr(mc(Y2t0[1,:])), ", ", repr(mc(Y2t0[2,:])), "]"	)

    pt = exp(lp(0., u, T, v, Pt))
	p = exp(SDE.varlp(0., u, T, v, ph0, B0, beta0, s->a(s, u, P)))
	println("p ", round(p,5), " ", repr(mc(exp(Yll)*pt)), repr(mc(exp(Ull)*pt)) , repr(mc(exp(Y2ll)*pt)))
	println("~p ", round(pt,5))
	w1 = mc(pt/p*exp(Yll))
	w2 = mc(pt/p*exp(Ull))
	w2b = mc(pt/p*exp(Ull2))
	w3 = mc(pt/p*exp(Y2ll))
	println(" Σ ~p/p*psi(X°) = ", repr(w1), " ≈ 1\n", " Σ ~p/p*psi(U) = ", repr(w2), repr(w2b), " ≈ 1\n", " Σ ~p/p*psi(X°°)= ", repr(w3), " ≈ 1")
	println("lmax ", maximum(exp(Yll)), " ", maximum(exp(Ull)), " ", maximum(exp(Y2ll)))

    Test.@test abs(w1[1] - 1.) < 1.96w1[2]
    Test.@test abs(w2[1] - 1.) < 1.96w2[2]
    Test.@test abs(w3[1] - 1.) < 1.96w3[2]

end


text:
# Test setup
#Beware of size of tests and time they take.
tic()
#@testsetup will run everything after the begin and end and report
#Optional tic() toc() to observe large performance regressions.
@testset "Abstract Hierarchies" begin
    include("testtypes.jl")
end

toc()


text:
module Utils
export chain, isiterable, @ifsomething, iterate_named

# include("unionall_implementationdetails.jl")  # TODO do we need this?

chain(itr...) = Iterators.flatten(itr)

isiterable(value) = isiterable(typeof(value))
isiterable(type::Type) = Base.isiterable(type)

"""
    @ifsomething expr

If `expr` evaluates to `nothing`, equivalent to `return nothing`, otherwise the macro
evaluates to the value of `expr`. Not exported, useful for implementing iterators.
```jldoctest
julia> using TypeClasses.Utils

julia> @ifsomething iterate(1:2)
(1, 1)
julia> let elt, state = @ifsomething iterate(1:2, 2); println("not reached"); end
```
"""
macro ifsomething(ex)
  quote
    result = $(esc(ex))
    result === nothing && return nothing
    result
  end
end

"""
    result = iterate_named(iterable)
    result = iterate_named(iterable, state)

Exactly like `iterate`, with the addition that you can use `result.value` in addition to 'result[1]'
and `result.state` for `result[2]`. It will return a named tuple respectively if not nothing, hence also the name.

Can be used instead of `iterate`.

```jldoctest
julia> using TypeClasses.Utils

julia> iterate_named(["one", "two"])
(value = "one", state = 2)
```
"""
function iterate_named(iter, state...)
  value, state′ = @ifsomething iterate(iter, state...)
  (value=value, state=state′)
end

end # module


text:
import JavaCall

function jsprit_java_unit_tests()

    setup_jvm_tests()
    jsprit_build_and_add_locations_tests()
    jsprit_build_vehicle_type_tests()
    jsprit_add_vehicles_tests()
    jsprit_add_jobs_tests()
    jsprit_create_input_tests()
    jsprit_create_algorithm_tests()
    jsprit_search_best_solution_tests()
    jsprit_transform_solution_tests()

    rvrp_jsprit_solve_tests()
    rvrp_jsprit_supported_features_tests()

end

function setup_jvm_tests()
    RVRP.setup_jvm()
    ghcoord = JavaCall.@jimport com.graphhopper.jsprit.core.util.Coordinate
    c = ghcoord((JavaCall.jdouble, JavaCall.jdouble,), 13.0, 124.0)
    x = JavaCall.jcall(c, "getX", JavaCall.jdouble, ())
    @test x == 13.0
end

function jsprit_build_and_add_locations_tests()
    data = RVRP.parse_cvrplib(dirname(@__FILE__)* "/../../data/cvrplib/P/P\\P-n16-k8.vrp")
    computed_data = RVRP.build_computed_data(data)
    jsp_types = RVRP.JspritJavaTypes()
    vrp_builder = jsp_types.JVrpBuilder(())
    locs = RVRP.jsprit_build_and_add_locations(data.locations,
                                                      vrp_builder, jsp_types)
    @test eltype(locs) == jsp_types.JLocation
end

function jsprit_build_vehicle_type_tests()
    data = RVRP.parse_cvrplib(dirname(@__FILE__)* "/../../data/cvrplib/P/P\\P-n16-k8.vrp")
    computed_data = RVRP.build_computed_data(data)
    jsp_types = RVRP.JspritJavaTypes()
    vrp_builder = jsp_types.JVrpBuilder(())
    jsvt = RVRP.jsprit_build_vehicle_type(data.vehicle_categories,
                                                 data.vehicle_sets[1],
                                                 computed_data, jsp_types)
    @test typeof(jsvt) == jsp_types.JVType
    @test JavaCall.jcall(jsvt, "getTypeId", JavaCall.JString, ()) == data.vehicle_sets[1].id
end

function jsprit_add_vehicles_tests()
    data = RVRP.parse_cvrplib(dirname(@__FILE__)* "/../../data/cvrplib/P/P\\P-n16-k8.vrp")
    computed_data = RVRP.build_computed_data(data)
    jsp_types = RVRP.JspritJavaTypes()
    vrp_builder = jsp_types.JVrpBuilder(())
    jsp_locs = RVRP.jsprit_build_and_add_locations(data.locations,
                                                          vrp_builder, jsp_types)
    RVRP.jsprit_add_vehicles(
        data.vehicle_categories, data.vehicle_sets, data.locations,
        data.location_groups, computed_data, vrp_builder, jsp_locs, jsp_types
    )
end

function jsprit_add_jobs_tests()
    data = RVRP.parse_cvrplib(dirname(@__FILE__)* "/../../data/cvrplib/P/P\\P-n16-k8.vrp")
    computed_data = RVRP.build_computed_data(data)
    jsp_types = RVRP.JspritJavaTypes()
    vrp_builder = jsp_types.JVrpBuilder(())
    jsp_locs = RVRP.jsprit_build_and_add_locations(data.locations,
                                                          vrp_builder, jsp_types)
    RVRP.jsprit_add_jobs(
        data.requests, data.locations, data.location_groups,
        data.product_specification_classes, computed_data, vrp_builder, jsp_locs,
        jsp_types
    )
end

function jsprit_create_input_tests()
    data = RVRP.parse_cvrplib(dirname(@__FILE__)* "/../../data/cvrplib/P/P\\P-n16-k8.vrp")
    computed_data = RVRP.build_computed_data(data)
    jsp_types = RVRP.JspritJavaTypes()
    jsp_vrp = RVRP.jsprit_create_input(data, computed_data,
                                                   jsp_types)
    @test typeof(jsp_vrp) == jsp_types.JVrp
    jsp_jobs = JavaCall.jcall(jsp_vrp, "getJobs", jsp_types.JMap, ())
    n_jobs = JavaCall.jcall(jsp_jobs, "size", JavaCall.jint, ())
    @test n_jobs == 15
end

function jsprit_create_algorithm_tests()
    solver = RVRP.JspritSolver()
    solver.params["iterations"] = "1000"
    data = RVRP.parse_cvrplib(dirname(@__FILE__) *
                              "/../../data/cvrplib/P/P\\P-n16-k8.vrp")
    computed_data = RVRP.build_computed_data(data)
    jsp_types = RVRP.JspritJavaTypes()
    jsp_vrp = RVRP.jsprit_create_input(data, computed_data, jsp_types)
    jsp_algo = RVRP.jsprit_create_algorithm(jsp_vrp, jsp_types, solver)
    @test typeof(jsp_algo) == jsp_types.JAlgorithm
end

function jsprit_search_best_solution_tests()
    solver = RVRP.JspritSolver()
    data = RVRP.parse_cvrplib(dirname(@__FILE__)*
                              "/../../data/cvrplib/P/P\\P-n16-k8.vrp")
    computed_data = RVRP.build_computed_data(data)
    jsp_types = RVRP.JspritJavaTypes()
    jsp_vrp = RVRP.jsprit_create_input(data, computed_data, jsp_types)
    jsp_algo = RVRP.jsprit_create_algorithm(jsp_vrp, jsp_types, solver)
    jsp_sol = RVRP.jsprit_search_best_solution(jsp_algo, jsp_types)
    cost = JavaCall.jcall(jsp_sol, "getCost", JavaCall.jdouble, ())
    @test typeof(jsp_sol) == jsp_types.JSolution
    @test cost == 450
end

function jsprit_transform_solution_tests()
    solver = RVRP.JspritSolver()
    data = RVRP.parse_cvrplib(dirname(@__FILE__)* "/../../data/cvrplib/P/P\\P-n16-k8.vrp")
    computed_data = RVRP.build_computed_data(data)
    jsp_types = RVRP.JspritJavaTypes()
    jsp_vrp = RVRP.jsprit_create_input(data, computed_data, jsp_types)
    jsp_algo = RVRP.jsprit_create_algorithm(jsp_vrp, jsp_types, solver)
    jsp_sol = RVRP.jsprit_search_best_solution(jsp_algo, jsp_types)
    rvrp_sol = RVRP.jsprit_transform_solution(data, solver, computed_data,
                                       jsp_sol, jsp_types)
    @test typeof(rvrp_sol) == RVRP.RvrpSolution
    @test rvrp_sol.cost == 450
    @test length(rvrp_sol.routes) == 8
    for r in rvrp_sol.routes
        @test r.sequence[1].location_id == r.sequence[end].location_id
    end
    @test length(rvrp_sol.unassigned_request_ids) == 0

    data.requests[1].product_quantity_range = RVRP.Range(100.0, 100.0)
    computed_data = RVRP.build_computed_data(data)
    jsp_vrp = RVRP.jsprit_create_input(data, computed_data, jsp_types)
    jsp_algo = RVRP.jsprit_create_algorithm(jsp_vrp, jsp_types, solver)
    jsp_sol = RVRP.jsprit_search_best_solution(jsp_algo, jsp_types)
    rvrp_sol = RVRP.jsprit_transform_solution(data, solver, computed_data,
                                       jsp_sol, jsp_types)
    @test rvrp_sol.unassigned_request_ids == ["req_1"]

end

function rvrp_jsprit_solve_tests()
    data = RVRP.parse_cvrplib(dirname(@__FILE__)* "/../../data/cvrplib/P/P\\P-n16-k8.vrp")
    computed_data = RVRP.build_computed_data(data)
    jsp_types = RVRP.JspritJavaTypes()
    rvrp_sol = RVRP.solve(data, RVRP.JspritSolver())
    @test typeof(rvrp_sol) == RVRP.RvrpSolution
    @test rvrp_sol.cost == 450
    @test length(rvrp_sol.routes) == 8
    for r in rvrp_sol.routes
        @test r.sequence[1].location_id == r.sequence[end].location_id
    end
    @test length(rvrp_sol.unassigned_request_ids) == 0

    # Empty instance
    data = RVRP.RvrpInstance()
    rvrp_sol = RVRP.solve(data, RVRP.JspritSolver())
    @test rvrp_sol.cost == 0.0
    @test rvrp_sol.routes == RVRP.Route[]
    @test rvrp_sol.unassigned_request_ids == String[]
end

function rvrp_jsprit_supported_features_tests()
    features_list = RVRP.supported_features(RVRP.JspritSolver)
    example_features = BitSet([RVRP.HAS_PICKUPONLY_REQUESTS,
                                     RVRP.HAS_DELIVERYONLY_REQUESTS])
    @test (issubset(example_features, features_list[1]))
end


text:
include("./single_bh_superradiance_param.jl")

s = [ p5 ]

for x in s
    solve_wave_equation_2D(x)
end


text:
using LaserTag
using POMDPToolbox
using ContinuousPOMDPTreeSearchExperiments
using POMDPs
using ProfileView
using ParticleFilters
using ARDESPOT

p = gen_lasertag(rng=MersenneTwister(4))

solver = DESPOTSolver(T_max=Inf,
             lambda=0.01,
             max_trials=1,
             bounds=LaserBounds{typeof(p)}(),
             rng=MersenneTwister(4))

planner = solve(solver, p)

b0 = initial_state_distribution(p)
s0 = rand(MersenneTwister(5), b0)
println(LaserTagVis(p, s=s0))

tree = ARDESPOT.build_despot(planner, b0)

println(TreeView(tree, 1, 3))

#=
@time action(policy, initial_state_distribution(p))

Profile.clear()
@profile action(policy, initial_state_distribution(p))
ProfileView.view()
=#

# hr = HistoryRecorder(max_steps=5, show_progress=true)
# filter = SIRParticleFilter(p, 100_000)
#
# simulate(hr, p, policy, filter)


text:
function approximation_term(
    i::MinimumBudget{SpanningTreeInstance{T, Maximise}, U},
    ::LagrangianRefinementAlgorithm,
) where {T, U}
    return maximum(values(i.instance.rewards))
end
function approximation_ratio(
    ::MinimumBudget{SpanningTreeInstance{T, Maximise}, U},
    ::LagrangianRefinementAlgorithm,
) where {T, U}
    return NaN
end

function solve(
    i::MinimumBudget{SpanningTreeInstance{T, Maximise}, U},
    ::LagrangianRefinementAlgorithm;
    ε::Float64=1.0e-3,
    ζ⁻::Float64=0.2,
    ζ⁺::Float64=5.0,
    stalling⁻::Float64=1.0e-5,
) where {T, U}
    # Approximately solve the following problem:
    #     \max_{x spanning tree} rewards x  s.t.  weights x >= budget
    # This algorithm provides an additive approximation to this problem. If x* is the optimum solution and x~ the one
    # returned by this algorithm,
    #     weights x* >= budget   and   weights x~ >= budget                 (the returned solution is feasible)
    #     rewards x~ >= rewards x* - \max{e edge} reward[e]                 (additive approximation)

    # Check assumptions.
    if ζ⁻ >= 1.0
        error(
            "ζ⁻ must be strictly less than 1.0: the dual multiplier λ is multiplied by ζ⁻ to reach an infeasible solution by less penalising the budget constraint.",
        )
    end
    if ζ⁺ <= 1.0
        error(
            "ζ⁺ must be strictly greater than 1.0: the dual multiplier λ is multiplied by ζ⁺ to reach a feasible solution by penalising more the budget constraint.",
        )
    end

    # Ensure the problem is feasible by only considering the budget constraint.
    feasible_rewards = Dict{Edge{T}, Float64}(
        e => i.weights[e] for e in keys(i.instance.rewards)
    )
    feasible_instance = SpanningTreeInstance(i.instance.graph, feasible_rewards)
    feasible_solution = solve(feasible_instance, PrimAlgorithm())

    if _budgeted_spanning_tree_compute_value(
        feasible_instance,
        feasible_solution.variables,
    ) < i.min_budget
        # By maximising the left-hand side of the budget constraint, impossible to reach the target budget. No solution!
        return SimpleBudgetedSpanningTreeSolution(i)
    end

    # Solve the Lagrangian relaxation to optimality.
    lagrangian = solve(i, LagrangianAlgorithm(), ε=ε)
    λ0, v0, st0 = lagrangian.λ, lagrangian.value, lagrangian.variables
    λmax = lagrangian.λmax
    b0 = _budgeted_spanning_tree_compute_weight(i, st0) # Budget consumption of this first solution.

    # If already respecting the budget constraint exactly, done!
    if b0 == i.min_budget
        return SimpleBudgetedSpanningTreeSolution(i, st0)
    end

    # Find two solutions: one above the budget x⁺ (i.e. respecting the constraint), the other not x⁻.
    x⁺, x⁻ = nothing, nothing

    λi = λ0
    if b0 > i.min_budget
        x⁺ = st0
        @assert _budgeted_spanning_tree_compute_weight(i, x⁺) > i.min_budget

        stalling = false
        while true
            # Penalise less the constraint: it should no more be satisfied.
            λi *= ζ⁻
            _, sti = _st_prim_budgeted_lagrangian(i, λi)
            if _budgeted_spanning_tree_compute_weight(i, sti) < i.min_budget
                x⁻ = sti
                break
            end

            # Is the process stalling?
            if λi <= stalling⁻
                stalling = true
                break
            end
        end

        # Specific handling of stallings.
        if stalling # First test: don't penalise the constraint at all.
            _, sti = _st_prim_budgeted_lagrangian(i, 0.0)
            new_budget = _budgeted_spanning_tree_compute_weight(i, sti)
            if new_budget < i.min_budget
                x⁻ = sti
                stalling = false
            end
        end

        if stalling # Second test: minimise the left-hand side of the budget constraint, in hope of finding a feasible solution.
            # This process is highly similar to the computation of feasible_solution, but with a reverse objective function.
            infeasible_rewards = Dict{Edge{T}, Float64}(
                e => -i.weights[e] for e in keys(i.weights)
            )
            infeasible_solution =
                solve(
                    SpanningTreeInstance(i.instance.graph, infeasible_rewards),
                    PrimAlgorithm(),
                ).variables

            if _budgeted_spanning_tree_compute_weight(i, infeasible_solution) <
               i.min_budget
                x⁻ = infeasible_solution
                stalling = false
            end
        end

        if stalling # Third: decide there is no solution strictly below the budget. No refinement is possible.
            # As x⁺ is feasible, return it.
            return SimpleBudgetedSpanningTreeSolution(i, x⁺)
        end
    else
        x⁻ = st0
        @assert _budgeted_spanning_tree_compute_weight(i, x⁻) < i.min_budget

        while true
            # Penalise more the constraint: it should become satisfied at some point.
            λi *= ζ⁺
            _, sti = _st_prim_budgeted_lagrangian(i, λi)
            if _budgeted_spanning_tree_compute_weight(i, sti) >= i.min_budget
                x⁺ = sti
                break
            end

            # Is the process stalling? If so, reuse feasible_solution, which is guaranteed to be feasible.
            if λi >= λmax
                x⁺ = feasible_solution.variables
                break
            end
        end
    end

    # Normalise the solutions: the input graph is undirected, the direction of the edges is not important.
    # In case one solution has the edge v -> w and the other one w -> v, make them equal. s
    sort_edge(e::Edge{T}) where {T} = (src(e) < dst(e)) ? e : reverse(e)
    x⁺ = [sort_edge(e) for e in x⁺]
    x⁻ = [sort_edge(e) for e in x⁻]

    # Iterative refinement. Stop as soon as there is a difference of at most one edge between the two solutions.
    while Kombinator._solution_symmetric_difference_size(x⁺, x⁻) > 2
        # Enforce the loop invariant.
        @assert x⁺ !== nothing
        @assert x⁻ !== nothing
        @assert _budgeted_spanning_tree_compute_weight(i, x⁺) >= i.min_budget # Feasible.
        @assert _budgeted_spanning_tree_compute_weight(i, x⁻) < i.min_budget # Infeasible.

        # Switch elements from one solution to another.
        only_in_x⁺, only_in_x⁻ = _solution_symmetric_difference(x⁺, x⁻)
        e1 = first(only_in_x⁺)
        e2 = first(only_in_x⁻)

        # Create the new solution (don't erase x⁺ nor x⁻: only one of the two will be forgotten, the other will be kept).
        new_x = copy(x⁺)
        filter!(e -> e != e1, new_x)
        push!(new_x, e2)

        # Replace one of the two solutions, depending on whether this solution is feasible (x⁺) or not (x⁻).
        if _budgeted_spanning_tree_compute_weight(i, new_x) >= i.min_budget
            x⁺ = new_x
        else
            x⁻ = new_x
        end
    end

    # Done!
    return SimpleBudgetedSpanningTreeSolution(i, x⁺)
end


text:
import Base64

@inline function _is_raw_ssh_private_key(content::AbstractString)
    result = ( occursin("-", content) ) && ( occursin(" ", content) ) &&
                                           ( occursin("BEGIN ", content) ) &&
                                           ( occursin("END ", content) ) &&
                                           ( occursin(" PRIVATE KEY", content) )

    return result
end

@inline function _decode_ssh_private_key(content::AbstractString)
    if _is_raw_ssh_private_key(content)
        @info("This is a raw SSH private key.")
        return content
    end
    @info("This doesn't look like a raw SSH private key. I will assume that it is a Base64-encoded SSH private key. I will now try to Base64-decode it.")
    decoded_result_array = Base64.base64decode(content)
    decoded_result_string = String(decoded_result_array)
    @info("The Base64-decoding completed successfully.")
    if !_is_raw_ssh_private_key(decoded_result_string)
        throw(BadSSHPrivateKeyError("Private key is not in the correct format. It must be an OpenSSH private key. You should use the OpenSSH `ssh-keygen` command to generate the private key."))
    end
    @info("This was a Base64-encoded SSH private key. I Base64-decoded it, and now I have the raw SSH private key.")
    return decoded_result_string
end

@inline function compathelper_priv_is_defined(env)
    COMPATHELPER_PRIV_is_defined = ( haskey(env, "COMPATHELPER_PRIV") ) && ( isa(env["COMPATHELPER_PRIV"], AbstractString) ) &&
                                                                           ( length(strip(env["COMPATHELPER_PRIV"])) > 0 ) &&
                                                                           ( !occursin("false", strip(env["COMPATHELPER_PRIV"])) )
    return COMPATHELPER_PRIV_is_defined
end


text:
function circle!(r)
    for side in [Sud,Ost,Nord,West]
        while isborder(r,side)==0 && ismarker(r)==0
            putmarker!(r)
            move!(r,side)
         end
    end
end

function homenowallmind!(r)
    x=0
    y=0
    while isborder(r,Nord)==0
        move!(r,Nord)
        x=x+1
    end
    while isborder(r,West)==0
        move!(r,West)
        y=y+1
    end
    return x,y
end
function homenowall!(r,x,y)
    for i in 1:x
        move!(r,Sud)
    end
    for i in 1:y
        move!(r,Ost)
    end
end

function circleinborder!(r)
   x,y= homenowallmind!(r)
    circle!(r)
    homenowall!(r,x,y)
end

circleinborder!(r)


text:
using Plots
using HDF5
using Pkg
using Statistics

data=h5open("test/mnist.hdf5","r")
xtrain=transpose(read(data,"x_train"))
ytrain=read(data,"t_train")
xvalid=transpose(read(data,"x_valid"))
yvalid=read(data,"t_valid")

n = size(xtrain,1)
xtrain = [ones(n,1) xtrain]
d = size(xtrain,2)

# Choose network structure and randomly initialize weights
include("NeuralNet.jl")
nHidden = [3,3,4,3]
# nHidden = [5,5,5,5]
nParams = NeuralNet_nParams(d,nHidden)
w = randn(nParams,1)

# Train with stochastic gradient
include("sgdfunc.jl")
maxIter = 10000

#----Uncomment below for hyperparameter search:----#

# aTest = [0.1, 0.01, 0.001, 0.0001, 0.00001]
# BTest = [1, 5, 10, 20, 100] #200 500 1000
#
# bestvalid = Inf
# bestw = []
# bestparams = [0,0]
# for B in BTest
#     for a in aTest
#         global bestvalid, bestw, bestparams
#         wb, validerr = VanillaSGD(NeuralNet_backprop, NeuralNet_predict, maxIter, nHidden,
#                                     nParams, xtrain, ytrain, xvalid, yvalid, a, B)
#         if (validerr[end]<bestvalid)
#             bestvalid=validerr[end]
#             bestw=wb
#             bestparams = [a,B]
#         end
#     end
# end

#best hyperparams found to be [0.0001, 20]

wbSGDB, validSGDB = SGDBabysitter(NeuralNet_backprop, NeuralNet_predict,
maxIter, nHidden, nParams, xtrain, ytrain, xvalid, yvalid)
savefig("SGDB_MNIST.pdf")

wbVan, validVan = VanillaSGD(NeuralNet_backprop, NeuralNet_predict, maxIter, nHidden, nParams, xtrain, ytrain, xvalid, yvalid, bestparams[1], convert(Int64, bestparams[2]))
savefig("Vanilla_MNIST.pdf")


text:
export RandomStartPolicy

Base.@kwdef mutable struct RandomStartPolicy{P,R<:RandomPolicy} <: AbstractPolicy
    policy::P
    random_policy::R
    num_rand_start::Int
end

function (p::RandomStartPolicy)(env)
    p.num_rand_start -= 1
    if p.num_rand_start < 0
        p.policy(env)
    else
        p.random_policy(env)
    end
end

function RLBase.update!(
    p::RandomStartPolicy,
    t::AbstractTrajectory,
    e::AbstractEnv,
    s::AbstractStage,
)
    update!(p.policy, t, e, s)
end

for f in (:prob, :priority)
    @eval function RLBase.$f(p::RandomStartPolicy, args...)
        if p.num_rand_start < 0
            $f(p.policy, args...)
        else
            $f(p.random_policy, args...)
        end
    end
end


text:
a = 3
b = 2
c = 1

fs = 100
T = 1/fs
f1 = 5
f2 = 10
f3 = 30

t = T:T:fs

x1 = a*sin.(2*pi*f1*T*t)+b*sin.(2*pi*f2*T*t)+c*sin.(2*pi*f3*T*t);

@testset "Peak prominence" begin
    @testset "Reciprocity" begin
        maxpks = argmaxima(x1)
        minpks = argminima(-x1)
        _, maxprom = peakproms(maxpks, x1)
        _, minprom = peakproms(minpks, -x1)
        @test maxprom == minprom

        _, maxprom = peakproms(maxpks, x1; strict=false)
        _, minprom = peakproms(minpks, -x1; strict=false)
        @test maxprom == minprom
    end

    @testset "Prominence values" begin
        x2 = sin.(1e-5:1e-5:9*pi)
        i, p = peakproms(argmaxima(x2), x2)

        @test p[[1,5]] ≈ [1., 1.] atol=1e-4
        @test p[[2,3,4]] ≈ [2., 2., 2.] atol=1e-4

        p2 = [1,0,2,0,1]
        @test last(peakproms(argmaxima(p2), p2)) == [2]
        @test last(peakproms(argmaxima(p2; strict=false), p2; strict=false)) == [1,2,1]

        # Prominence should be the same regardless of window size
        p1 = [0,0,3,1,2,0,4,0,0,5]
        @test last(peakproms(argmaxima(p1, 1), p1)) == [3,1,4]
        @test last(peakproms(argmaxima(p1, 2), p1)) == [3,4]

        # A peaks of the same height count as an intersection for reference intervals
        p4 = [0,4,2,4,3,4,0]
        @test last(peakproms(argmaxima(p4), p4)) == [2,1,1]
        @test last(peakproms(argmaxima(p4[2:end-1]), p4[2:end-1])) == [1]
        @test last(peakproms(argmaxima(p4[2:end-1]; strict=false), p4[2:end-1]; strict=false)) == [2,1,1]

        # The presence of a missing/NaN in either bounding interval poisons the prominence
        m4 = [missing; p4; missing]
        n4 = [NaN; p4; NaN]
        @test isequal(last(peakproms(argmaxima(m4), m4)), [missing,1,missing])
        @test isequal(last(peakproms(argmaxima(n4), n4)), [NaN,1.,NaN])
        @test last(peakproms(argmaxima(m4), m4; strict=false)) == [2,1,1]
        @test last(peakproms(argmaxima(n4), n4; strict=false)) == [2.,1.,1.]

        m5 = [missing, 1, missing]
        n5 = [NaN, 1, NaN]
        @test last(peakproms(argmaxima(m5; strict=false), m5; strict=false)) == [0]
        @test last(peakproms(argmaxima(n5; strict=false), n5; strict=false)) == [0]

        p5 = [-1,6,3,4,2,4,2,5,-2,0]
        @test last(peakproms(argmaxima(p5, 3; strict=false), p5; strict=false)) == [7,3]
        @test last(peakproms(argmaxima(reverse(p5), 3; strict=false), reverse(p5); strict=false)) == [3,7]

    end

    @testset "Min/max prominence" begin
        sint = sin.(T:T:6pi)
        maxs = argmaxima(sint)
        @test length(first(peakproms(maxs, sint; minprom=1.5))) == 2
        @test length(first(peakproms(maxs, sint; maxprom=1.5))) == 1

        @test_throws ArgumentError peakproms([1,2,3], sint; maxprom=0.1, minprom=1)
    end

    # issue #4
    let i, p
        i, p = peakproms(Int[], zeros(10))
        @test isempty(i)
        @test isempty(p)
        i, p = peakproms([1], zeros(10); strict=false)
        @test p == [0.0]
    end

    issue91ddaa9 = [1,2,3,2,3,1]
    @test all(!iszero,
        last(peakproms(argminima(issue91ddaa9; strict=false), issue91ddaa9; strict=false)))

    @test_deprecated peakprom(Minima(), x1)
    @test_deprecated peakprom(Maxima(), x1)

    # Deprecated keyword "strictbounds"
    @test_deprecated r"renamed" peakproms([2], [0,1,0]; strictbounds=true)
end


text:
endof(s::MutableASCIIString) = length(s.data)
getindex(s::MutableASCIIString, i::Int) = (x=s.data[i]; x < 0x80 ? char(x) : '\ufffd')

getindex(s::MutableASCIIString, r::Vector) = ASCIIString(getindex(s.data,r))
getindex(s::MutableASCIIString, r::UnitRange{Int}) = ASCIIString(getindex(s.data,r))
getindex(s::MutableASCIIString, indx::AbstractVector{Int}) = ASCIIString(s.data[indx])

setindex!(s::MutableASCIIString, x, i0::Real) = (s.data[i0] = x)
setindex!(s::MutableASCIIString, r::ASCIIString, I::UnitRange) = (s.data[I] = r.data)
setindex!(s::MutableASCIIString, c::Char, I::UnitRange) = (s.data[I] = uint8(c))

search(s::MutableASCIIString, c::Char, i::Integer) = c < 0x80 ? search(s.data,uint8(c),i) : 0
rsearch(s::MutableASCIIString, c::Char, i::Integer) = c < 0x80 ? rsearch(s.data,uint8(c),i) : 0

string(c::MutableASCIIString) = c
function string(c::MutableASCIIString...)
    n = 0
    for s in c
        n += length(s.data)
    end
    v = Array(Uint8,n)
    o = 1
    for s in c
        ls = length(s.data)
        unsafe_copy!(v, o, s.data, 1, ls)
        o += ls
    end
    MutableASCIIString(v)
end

function ucfirst!(s::MutableASCIIString)
    if length(s) > 0 && 'a' <= s[1] <= 'z'
        s.data[1] -= 32
    end
    nothing
end
function lcfirst!(s::MutableASCIIString)
    if length(s) > 0 && 'A' <= s[1] <= 'Z'
        s.data[1] += 32
    end
    nothing
end

function uppercase!(s::MutableASCIIString)
    d = s.data
    for i = 1:length(d)
        if 'a' <= d[i] <= 'z'
            d[i] -= 32
        end
    end
end

function lowercase!(s::MutableASCIIString)
    d = s.data
    const l::Int = length(d)
    for i = 1:l
        if 'A' <= d[i] <= 'Z'
            d[i] += 32
        end
    end
end

reverse!(s::MutableASCIIString) = reverse!(s.data)

function _splice!{T<:Integer}(a::Vector, r::UnitRange{T}, ins::AbstractArray=Base._default_splice, m::Int=length(ins))
    #m = length(ins)
    if m == 0
        deleteat!(a, r)
        return
    end

    n = length(a)
    f = first(r)
    l = last(r)
    d = length(r)

    if m < d
        delta = d - m
        if f-1 < n-l
            Base._deleteat_beg!(a, f, delta)
        else
            Base._deleteat_end!(a, l-delta+1, delta)
        end
    elseif m > d
        delta = m - d
        if f-1 < n-l
            Base._growat_beg!(a, f, delta)
        else
            Base._growat_end!(a, l+1, delta)
        end
    end

    for k = 1:m
        a[f+k-1] = ins[k]
    end
end

function replace!{T}(str::MutableASCIIString, pattern::T, repl::Union(Function,ASCIIString,Char), limit::Integer=0)
    n = 1
    es = endof(str)
    r = search(str,pattern,1)
    j, k = first(r), last(r)
    while j != 0
        rv::Union(ASCIIString,Char) = isa(repl, Function) ? repl(SubString(str, j, k)) : repl
        if isa(rv, ASCIIString)
            _splice!(str.data, r, rv.data)
            lrv = length(rv)
            lr = length(r)
            if lrv != lr
                k = j+lrv
                es = es - lr + lrv
            end
        else
            str[r] = rv
        end
        ((k+1) > es) && break
        r = search(str, pattern, k+1)
        j, k = first(r), last(r)
        n == limit && break
        n += 1
    end
end

function replace!(str::MutableASCIIString, re::Regex, repl::Union(Function,ASCIIString,Char), limit::Integer=0)
    regex = Base.compile(re).regex
    extra = re.extra
    n = length(str.data)
    offset = int32(0)
    opts = re.options & Base.PCRE.EXECUTE_MASK
    opts_nonempty = opts | Base.PCRE.ANCHORED | Base.PCRE.NOTEMPTY_ATSTART
    prevempty = false
    ovec = Array(Int32, 3)
    n_repl = 1
    while true
        result = ccall((:pcre_exec, :libpcre), Int32,
                       (Ptr{Void}, Ptr{Void}, Ptr{Uint8}, Int32,
                       Int32, Int32, Ptr{Int32}, Int32),
                       regex, extra, str.data, n,
                       offset, prevempty ? opts_nonempty : opts, ovec, 3)

        if result < 0
            if prevempty && offset < n
                offset = int32(nextind(str, offset + 1) - 1)
                prevempty = false
                continue
            else
                break
            end
        end

        rv::Union(ASCIIString,Char) = isa(repl, Function) ? repl(SubString(str, ovec[1]+1, ovec[2])) : repl
        r = (ovec[1]+1):(ovec[2])
        if isa(rv, ASCIIString)
            _splice!(str.data, r, rv.data)

            lrv = length(rv)
            lr = length(r)
            if lrv != lr
                n = n - lr + lrv
            end

            prevempty = offset == ovec[2]
            offset = ovec[1]+1+lrv
        else
            str[r] = rv
            prevempty = offset == ovec[2]
            offset = ovec[2]
        end
        n_repl == limit && break
        n_repl += 1
    end
    nothing
end

function map!(f, s::MutableASCIIString)
    d = s.data
    for i = 1:length(d)
        d[i] = convert(Uint8, f(d[i]))
    end
end


text:
export fit_statespace, fit_statespace_gd, fit_statespace_constrained, fit_statespace_gd!
using ProximalOperators, SparseArrays, LinearAlgebra, DiffResults

function fit_statespace_gd(d::AbstractIdData,λ; initializer::Symbol=:kalman, extend=false,kwargs...)
    @assert hasinput(d) "The identification data must have inputs for this method to apply"
    model = SimpleLTVModel(d, extend=extend)
    y,A     = matrices(model,d)
    T = length(d)
    n = nstates(d)
    m = ninputs(d)
    T      -= 1
    if initializer != :kalman
        k = A\y
        k = repmat(k',T,1)
        k .+= 0.00001randn(size(k))
        model = statevec2model(k,n,m,false)
    else
        R1 = 0.1*eye(n^2+n*m)
        R2 = 10eye(n)
        P0 = 10000R1
        model = KalmanModel(d,R1,R2,P0, extend=false)
    end
    fit_statespace_gd!(model,d,λ; extend=extend, kwargs...)
end

function fit_statespace_gd!(model::AbstractModel,d::AbstractIdData,λ; normType = 1, D = 1, step=0.001, iters=10000, lasso=0, decay_rate=0.999, momentum=0.9, print_period = 100, reduction=0, extend=true, opt=LBFGS(m=100), kwargs...)
    if reduction > 0
        decay_rate = decayfun(iters, reduction)
    end
    T = length(d)
    n = nstates(d)
    m = ninputs(d)
    x = state(d)
    k = model2statevec(model)
    y, A     = matrices(model,d)
    nparams = size(A,2)
    T      -= 1
    diff_fun = D == 2 ? x-> diff(diff(x,dims=1),dims=1) : x-> diff(x,dims=1)
    function lossfun(k2)
        loss    = 0.
        for i = 1:T
            ii = (i-1)*n+1
            ii2 = ii+n-1
            @inbounds @views loss += mean((y[ii:ii2,:] - A[ii:ii2,:]*k2[i,:]).^2)
        end
        NK = length(k2)
        if normType == 1
            loss += λ^2/NK*sum(sqrt,sum(diff_fun(k2./std(k2,dims=1)).^2, dims=2) )
        else
            loss += λ/NK*sum(diff_fun(k2).^2)
        end
        if lasso > 0
            loss += lasso^2*sum(abs, k2)
        end
        loss
    end

    # inputs      = (k,)
    # loss_tape   = GradientTape(lossfun, inputs)
    # results     = similar.(inputs)
    # all_results = DiffResults.GradientResult.(results)
    #
    # mom         = zeros(size(k))
    # @show bestcost    = lossfun(k)
    # costs       = Inf*ones(iters+1); costs[1] = bestcost
    # steps       = zeros(iters)
    # # opt         = RMSpropOptimizer(k, step, 0.8, momentum)
    # opt         = ADAMOptimizer(k; α = step,  β1 = 0.9, β2 = 0.999, ɛ = 1e-8)
    # for iter = 1:iters
    #     steps[iter] = step
    #     gradient!(all_results, loss_tape, inputs)
    #     costs[iter] = all_results[1].value
    #     # mom .= step.*all_results[1].derivs[1] .+ momentum.*mom
    #     # k .-= mom
    #     opt(all_results[1].derivs[1], iter)
    #     # opt(all_results.derivs, iter)
    #     if iter % print_period == 0
    #         println("Iteration: ", iter, " cost: ", round(costs[iter],digits=6), " stepsize: ", step)
    #     end
    #     step *=  decay_rate
    #     opt.α = step
    # end

    inputs      = k
    loss_tape   = GradientTape(lossfun, inputs)
    results     = similar(inputs)
    all_results = DiffResults.GradientResult(results)
    gradfun = (G,inputs) -> gradient!(G, loss_tape, inputs)
    costs = Optim.optimize(lossfun, gradfun, k, opt, Optim.Options(store_trace=true, show_trace=true, show_every=10, iterations=10000, allow_f_increases=false, time_limit=100, x_tol=0, f_tol=0, g_tol=1e-8, f_calls_limit=iters, g_calls_limit=0))

    k = costs.minimizer

    At,Bt = ABfromk(k,n,m,T)
    SimpleLTVModel(At,Bt,extend), costs
end

function fit_statespace_constrained(d::AbstractIdData,changepoints::AbstractVector; extend=true)
    T             = length(d)
    n             = nstates(d)
    m             = ninputs(d)
    y,A           = matrices(SimpleLTVModel(d), d)
    nc            = length(changepoints)
    changepointse = [1; changepoints; T-1]
    Ai            = zeros(n,n,nc+1)
    Bi            = zeros(n,m,nc+1)
    k = map(1:nc+1) do i
        ii = (changepointse[i]-1)*n+1
        ii2 = changepointse[i+1]*n-1
        inds = ii:ii2
        k = (A[inds,:]\y[inds])'
    end
    At,Bt = segments2full(k,changepoints,n,m,T)
    return SimpleLTVModel(At,Bt,extend)
end

"""
model = fit_admm(modeltype::Type, d::AbstractIdData, λ; initializer::Symbol=:kalman)

Fit a model by solving `minimize ||y-ŷ||² + λ²||Dₓ k||` where `x` is 1 or 2
using linearized ADMM

See README for usage example

# Keyword arguments
iters      = 10000
D          = 1 # Order of differentiation
tol        = 1e-5
printerval = 100
zeroinit   = false
cb         = nothing # Callback function `model -> cb(model)`
γ          = 0.05 # Regularization parameter
μ          = γ/4/D^2/(ridge == 0 ? 1 : 2), # 32 is the biggest possible ||A||₂² # ADMM parameter
ridge      = 0 # `ridge > 0` Add some L2 regularization (`||k||`)
"""
function fit_admm(model::AbstractModel, d::AbstractIdData,λ; initializer::Symbol=:kalman, zeroinit = false, kwargs...)
    y,A  = matrices(model, d)
    T = length(d)
    n = nstates(d)
    m = ninputs(d)
    FT = eltype(y)
    T -= 1
    if !zeroinit
        if initializer != :kalman
            k = A\y
            k = repmat(k',T,1)
            model = statevec2model(model,k,n,m,false)
        else
            R1 = FT.(0.1*eye(n^2+n*m))
            R2 = FT.(10eye(n))
            P0 = 10000R1
            model = KalmanModel(model,d,R1,R2,P0, extend=false)
        end
    end
    fit_admm!(model, d,λ; zeroinit=zeroinit, kwargs...)
end

function fit_admm!(model::AbstractModel,d::AbstractIdData,λ;
    iters      = 10000,
    D          = 1,
    tol        = Float64(1e-5),
    printerval = 100,
    zeroinit   = false,
    cb         = nothing,
    γ          = Float64(0.05),
    ridge      = 0,
    μ          = Float64(γ/4/D^2/(ridge == 0 ? 1 : 2)), # 32 is the biggest possible ||A||₂²
    kwargs...)

    T = length(d)
    n = nstates(d)
    m = ninputs(d)
    @show γ, ridge
    k       = LTVModels.model2statevec(model)' |> copy
    y, Φ    = matrices(model,d)
    nparams = size(Φ,2)
    y       = reshape(y,n,:)
    FT      = eltype(y)
    T      -= 1
    NK      = length(k)
    x       = !zeroinit*copy(k[:])
    A       = sparse(FT(1.0)*I,NK,NK)
    if D == 1
        normA2 = ridge > 0 ? 2*3.91 : 3.91
        z       = !zeroinit*diff(k,dims=2)[:]
        for i = 1:NK-nparams
            A[i, i+nparams] = -1
        end
        A       = A[1:end-nparams,:]
    elseif D == 2
        normA2 = ridge > 0 ? 2*15.1 : 15.1
        z       = !zeroinit*diff(diff(k,dims=2),dims=2)[:]
        for i = 1:NK-nparams
            A[i, i+nparams] = -2
        end
        for i = 1:NK-2nparams
            A[i, i+2nparams] = 1
        end
        A       = A[1:end-2nparams,:]
    end
    @assert 0 ≤ μ ≤ γ/normA2 "μ should be ≤ γ/$normA2"

    proxf = prox_ls(model, y, Φ)

    gs = ntuple(t->NormL2(λ), T-D)
    indsg = ntuple(t->((t-1)*nparams+1:t*nparams, ) ,T-D)

    ## To add extra penalty
    if ridge > 0
        Q     = sparse(FT(ridge)*I(nparams))
        q     = zeros(FT,nparams)
        gs2   = fill(Quadratic(Q,q), T-D)
        gs    = (gs...,gs2...)
        indsg = (indsg..., indsg...)
        A     = [A; A]
        z     = [z;z]
    end
    ##

    proxg = SlicedSeparableSum(gs, indsg)
    Ax        = A*x
    u         = zeros(FT,size(z))
    Axz       = Ax .- z
    Axzu      = similar(u)
    proxf_arg = similar(x)
    proxg_arg = similar(u)
    x,z = _fit_admm_inner(Axzu, Axz,u,μ,γ,A,x,proxf_arg,proxf,z,Ax,proxg,proxg_arg,n,m,T,tol,cb,printerval,iters)
    k = reshape(x,nparams,T)' |> copy
    model = LTVModels.statevec2model(typeof(model),k,n,m,true)
    # plot(flatten(model))
    model
end

function _fit_admm_inner(Axzu, Axz,u,μ,γ,A,x,proxf_arg,proxf,z,Ax,proxg,proxg_arg,n,m,T,tol,cb,printerval,iters)
    for i = 1:iters

        Axzu .= Axz.+u
        proxf_arg .= -(μ/γ) .* A'Axzu .+ x # proxf_arg .= x - (μ/γ)*A'*Axzu
        # proxf_arg .+= x

        prox!(x, proxf, proxf_arg, μ)
        Ax .= A*x # Ax       .= A*x
        proxg_arg .= Ax .+ u
        prox!(z, proxg, proxg_arg, γ)
        Axz .= Ax .- z
        u  .+= Axz

        nAxz = norm(Axz)
        if i % printerval == 0
            @printf("%d ||Ax-z||₂ %.6f\n", i,  nAxz)
            if cb !== nothing
                cb(reshape(x,:,T)' |> copy)
            end
        end
        if nAxz < tol
            @info("||Ax-z||₂ ≤ tol")
            break
        end
    end
    x,z
end

function prox_ls(model::SimpleLTVModel, y, Φ)
    nparams = size(Φ,2)
    n,T = size(y)
    fs = ntuple(T) do t
        ii  = (t-1)*n+1
        ii2 = ii+n-1
        a   = Φ[ii:ii2,:]
        Q   = Matrix(a'a)
        q   = -a'y[:,t]
        ProximalOperators.QuadraticIterative(2Q,2q)
    end
    indsf = ntuple(t->((t-1)*nparams+1:t*nparams, ), T)
    proxf = SlicedSeparableSum(fs, indsf)
end

function prox_ls(model::LTVAutoRegressive, y, Φ)
    LeastSquares(Φ,y)
end


text:
# Euclidean distance from a set

immutable DistL2{R <: Real} <: ProximableFunction
  ind::IndicatorConvex
  lambda::R
  function DistL2(ind::IndicatorConvex, lambda::R)
    if lambda < 0
      error("parameter λ must be nonnegative")
    else
      new(ind, lambda)
    end
  end
end

DistL2{R <: Real}(ind::IndicatorConvex, lambda::R=1.0) = DistL2{R}(ind, lambda)

function (f::DistL2){T <: RealOrComplex}(x::AbstractArray{T})
  p, = prox(f.ind, x)
  return f.lambda*vecnorm(x-p)
end

function prox!{T <: RealOrComplex}(f::DistL2, x::AbstractArray{T}, y::AbstractArray{T}, gamma::Real=1.0)
  p, = prox(f.ind, x)
  d = vecnorm(x-p)
  gamlam = (gamma*f.lambda)
  if gamlam < d
    gamlamd = gamlam/d
    for k in eachindex(p)
      y[k] = (1-gamlamd)*x[k] + gamlamd*p[k]
    end
    return f.lambda*(d-gamlam)
  end
  y[:] = p
  return 0.0
end

fun_name(f::DistL2) = "Euclidean distance from a convex set"
fun_dom(f::DistL2) = fun_dom(f.ind)
fun_expr(f::DistL2) = "x ↦ λ inf { ||x-y|| : y ∈ S} "
fun_params(f::DistL2) = string("λ = $(f.lambda), S = ", typeof(f.ind))

function prox_naive{T <: RealOrComplex}(f::DistL2, x::AbstractArray{T}, gamma::Real=1.0)
  p, = prox(f.ind, x)
  d = vecnorm(x-p)
  gamlam = gamma*f.lambda
  if d > gamlam
    return x + gamlam/d*(p-x), f.lambda*(d-gamlam)
  end
  return p, 0.0
end


text:
using  DataStructures
using  Unitful

t  = range(0.0, stop=10.0, length=100)
t2 = range(0.0, stop=10.0, length=110)

result = OrderedDict{String,Any}()

result["time"] = t*u"s"
result["phi"]  = sin.(t)u"rad"
result["phi2"] = 0.5 * sin.(t)u"rad"
result["w"]    = cos.(t)u"rad/s"
result["w2"]   = 0.6 * cos.(t)u"rad/s"
result["r"]    = [[0.4 * cos(t[i]),
                   0.5 * sin(t[i]),
                   0.3 * cos(t[i])] for i in eachindex(t)]*u"m"

result["nothingSignal"]   = nothing
result["emptySignal"]     = Float64[]
result["wrongSizeSignal"] = sin.(t2)u"rad"

println("... Next plots should give warnings:")
plot(result, ("phi", "r", "signalNotDefined"), heading="Plot with warning 1" , figure=1)
plot(result, ("signalNotDefined",
              "nothingSignal",
              "emptySignal",
              "wrongSizeSignal"),
              heading="Plot with warning 2" , figure=2)

println("\n... result info:")
printResultInfo(result)


text:
module SMC

export HiddenMarkovModel
export num_states
export num_observations
export hmm_simulate
export hmm_log_marginal_likelihood
export hmm_posterior_sample
export hmm_log_joint_probability
export HMMPriorSMCScheme
export HMMConditionalSMCScheme
export render_hmm!
export render_hmm_states!
export render_hmm_posterior_marginals!
export render_hmm_observations!

export logsumexp
export SMCScheme
export smc
export conditional_smc

include("mathutils.jl")
include("generic_smc.jl")
include("hmm.jl")
include("hmm_render.jl")

end


text:
function covidPlot3(cases,death,recov,pop,name)
p1 = plot()
hspan!(p1,[0,1e2], color  = :red, alpha   = 0.6, labels = false);
hspan!(p1,[0,1e1], color  = :white, alpha = 0.1, labels = false);
hspan!(p1,[0,1e0], color  = :white, alpha = 0.1, labels = false);
hspan!(p1,[0,1e-1], color = :white, alpha = 0.1, labels = false);
hspan!(p1,[0,1e-2], color = :white, alpha = 0.1, labels = false);
hspan!(p1,[0,1e-3], color = :white, alpha = 0.1, labels = false);
hspan!(p1,[0,1e-4], color = :white, alpha = 0.1, labels = false);
hspan!(p1,[0,1e-5], color = :white, alpha = 0.1, labels = false);
hspan!(p1,[0,1e-6], color = :white, alpha = 0.1, labels = false);
scatter!(1:size(cases)[2],
      100*cases[1,:]./pop,lab=name*" COVID Cases",
      scale=:log10,
      yticks = [1e-6,1e-5,1e-4,1e-3,1e-2,1e-1,1,1e1,1e2],
      yaxis="Percent of Country", xaxis="Days since Jan 22",
      marker =:o,
      color =:red)
scatter!(1:size(cases)[2],
      100*death[1,:]./pop,lab=name*" COVID Deaths",
      scale=:log10,
      yticks = [1e-6,1e-5,1e-4,1e-3,1e-2,1e-1,1,1e1,1e2],
      yaxis="Percent of Country", xaxis="Days since Jan 22",
      marker =:x,
      color =:black)
scatter!(1:size(cases)[2],
      100*recov[1,:]./pop,lab=name*" COVID Recoveries",
      scale=:log10,
      yticks = [1e-6,1e-5,1e-4,1e-3,1e-2,1e-1,1,1e1,1e2],
      yaxis="Percent of Country", xaxis="Days since Jan 22",
      marker =:^,
      color =:blue)
 plot!(xscale=:identity,xticks=0:10:100)
 labelCases = cases[1,end];
 annotate!(size(cases)[2]+5, 200*cases[1,end]./pop, ha="left",va="bottom","$labelCases")
 xlims!((0,100))
 ylims!((1e-6,100))
 png(p1,"CountryFigures/$name.png")
end


text:
###
### **acb_elliptic.h** -- elliptic integrals and functions of complex variables
###

### Complete elliptic integrals
arbcall"void acb_elliptic_k(acb_t res, const acb_t m, slong prec)"
arbcall"void acb_elliptic_k_jet(acb_ptr res, const acb_t m, slong len, slong prec)"
arbcall"void _acb_elliptic_k_series(acb_ptr res, acb_srcptr m, slong mlen, slong len, slong prec)"
arbcall"void acb_elliptic_k_series(acb_poly_t res, const acb_poly_t m, slong len, slong prec)"
arbcall"void acb_elliptic_e(acb_t res, const acb_t m, slong prec)"
arbcall"void acb_elliptic_pi(acb_t res, const acb_t n, const acb_t m, slong prec)"

### Legendre incomplete elliptic integrals
arbcall"void acb_elliptic_f(acb_t res, const acb_t phi, const acb_t m, int pi, slong prec)"
arbcall"void acb_elliptic_e_inc(acb_t res, const acb_t phi, const acb_t m, int pi, slong prec)"
arbcall"void acb_elliptic_pi_inc(acb_t res, const acb_t n, const acb_t phi, const acb_t m, int pi, slong prec)"

### Carlson symmetric elliptic integrals
arbcall"void acb_elliptic_rf(acb_t res, const acb_t x, const acb_t y, const acb_t z, int flags, slong prec)"
arbcall"void acb_elliptic_rg(acb_t res, const acb_t x, const acb_t y, const acb_t z, int flags, slong prec)"
arbcall"void acb_elliptic_rj(acb_t res, const acb_t x, const acb_t y, const acb_t z, const acb_t p, int flags, slong prec)"
arbcall"void acb_elliptic_rj_carlson(acb_t res, const acb_t x, const acb_t y, const acb_t z, const acb_t p, int flags, slong prec)"
arbcall"void acb_elliptic_rj_integration(acb_t res, const acb_t x, const acb_t y, const acb_t z, const acb_t p, int flags, slong prec)"
arbcall"void acb_elliptic_rc1(acb_t res, const acb_t x, slong prec)"

### Weierstrass elliptic functions
arbcall"void acb_elliptic_p(acb_t res, const acb_t z, const acb_t tau, slong prec)"
arbcall"void acb_elliptic_p_jet(acb_ptr res, const acb_t z, const acb_t tau, slong len, slong prec)"
arbcall"void _acb_elliptic_p_series(acb_ptr res, acb_srcptr z, slong zlen, const acb_t tau, slong len, slong prec)"
arbcall"void acb_elliptic_p_series(acb_poly_t res, const acb_poly_t z, const acb_t tau, slong len, slong prec)"
arbcall"void acb_elliptic_invariants(acb_t g2, acb_t g3, const acb_t tau, slong prec)"
arbcall"void acb_elliptic_roots(acb_t e1, acb_t e2, acb_t e3, const acb_t tau, slong prec)"
arbcall"void acb_elliptic_inv_p(acb_t res, const acb_t z, const acb_t tau, slong prec)"
arbcall"void acb_elliptic_zeta(acb_t res, const acb_t z, const acb_t tau, slong prec)"
arbcall"void acb_elliptic_sigma(acb_t res, const acb_t z, const acb_t tau, slong prec)"


text:
# Function to check parameters

# Input:
# θ:        parameters
# model:    Model specifications

# Output:
# Admissible parameters or not?!

"""
    parametercheck(θ, model)
Check if parameters are admissable.

* `θ`: Parameters (vector or parameter)
* `model`: Model specification

# Example
```julia-repl
model = Model(pastObs = 1)
parametercheck([10, 0.3], model)
```

Parameters are admissible if they yield strictly positive conditional means and
further fulfill stationarity properties.
"""
function parametercheck(θ::parameter, model::INGARCHModel)
    logl = model.link == "Log"
    lin = !logl
    nb = model.distr == "NegativeBinomial"

    β0 = θ.β0
    α = θ.α
    β = θ.β
    αβ = [α; β]
    η = θ.η
    ϕ = θ.ϕ
    ω = θ.ω

    p = length(α)
    q = length(β)

    if !(0 <= ω <= 1) return false end
    if any(ϕ .<= 0) return false end

    r = length(model.external)

    if r > 0
        if lin
            if !all(η .>= 0)
                return false
            end
        end
    end

    if logl
        if p + q > 0
            if any(abs.(αβ) .>= 1) | (abs(sum(αβ)) >= 1)
                return false
            end
        end
    else
        if β0 <= 0
            return false
        end

        if p + q > 0
            if any(αβ .< 0) | (sum(αβ) > 1)
                return false
            end
        end
    end

    return true
end

function parametercheck(θ::parameter, model::INARCHModel)
    logl = model.link == "Log"
    lin = !logl
    nb = model.distr == "NegativeBinomial"

    β0 = θ.β0
    α = θ.α
    η = θ.η
    ϕ = θ.ϕ
    ω = θ.ω

    p = length(α)

    if !(0 <= ω <= 1) return false end
    if any(ϕ .<= 0) return false end

    r = length(model.external)

    if r > 0
        if lin
            if !all(η .>= 0)
                return false
            end
        end
    end

    if logl
        if p > 0
            if any(abs.(α) .>= 1) | (abs(sum(α)) >= 1)
                return false
            end
        end

    else
        if β0 <= 0
            return false
        end

        if p > 0
            if any(α .< 0) | (sum(α) > 1)
                return false
            end
        end
    end

    return true
end

function parametercheck(θ::parameter, model::IIDModel)
    logl = model.link == "Log"
    lin = !logl
    nb = model.distr == "NegativeBinomial"

    β0 = θ.β0
    η = θ.η
    ϕ = θ.ϕ
    ω = θ.ω

    if !(0 <= ω <= 1) return false end
    if any(ϕ .<= 0) return false end

    r = length(model.external)

    if r > 0
        if lin
            if !all(η .>= 0)
                return false
            end
        end
    end

    if !logl
        if β0 <= 0
            return false
        end
    end

    return true
end

function parametercheck(θ::parameter, model::INARMAModel)
    logl1 = model.link[1] == "Log"
    lin1 = !logl1

    logl2 = model.link[2] == "Log"
    lin2 = !logl2

    nb1 = model.distr[1] == "NegativeBinomial"
    nb2 = model.distr[2] == "NegativeBinomial"

    p = length(model.pastObs)
    q = length(model.pastMean)
    r = length(model.external)

    β0 = θ.β0
    α = θ.α
    β = θ.β
    η = θ.η
    ϕ = θ.ϕ
    ω = θ.ω

    if lin1 & (β0 <= 0)
        return false
    end

    if p > 0
        if any(α .< 0)
            return false
        end

        if sum(α) >= 1
            return false
        end
    end

    if q > 0
        if !all(0 .<= β .< 1)
            return false
        end
    end

    if r > 0
        if lin1
            if !all(η[findall(.!model.external)] .>= 0)
                return false
            end
        end

        if lin2
            if !all(η[findall(model.external)] .>= 0)
                return false
            end
        end
    end

    if nb1 | nb2
        if any(ϕ .<= 0)
            return false
        end
    end

    if !(0 <= ω <= 1)
        return false
    end

    return true
end

function parametercheck(θ::parameter, model::INARModel)
    logl1 = model.link[1] == "Log"
    lin1 = !logl1

    logl2 = model.link[2] == "Log"
    lin2 = !logl2

    nb1 = model.distr[1] == "NegativeBinomial"
    nb2 = model.distr[2] == "NegativeBinomial"

    p = length(model.pastObs)
    r = length(model.external)

    β0 = θ.β0
    α = θ.α
    η = θ.η
    ϕ = θ.ϕ
    ω = θ.ω

    if lin1 & (β0 <= 0)
        return false
    end

    if p > 0
        if any(α .< 0)
            return false
        end

        if sum(α) >= 1
            return false
        end
    end

    if r > 0
        if lin1
            if !all(η[findall(.!model.external)] .>= 0)
                return false
            end
        end

        if lin2
            if !all(η[findall(model.external)] .>= 0)
                return false
            end
        end
    end

    if nb1 | nb2
        if any(ϕ .<= 0)
            return false
        end
    end

    if !(0 <= ω <= 1)
        return false
    end

    return true
end

function parametercheck(θ::parameter, model::INMAModel)
    logl1 = model.link[1] == "Log"
    lin1 = !logl1

    logl2 = model.link[2] == "Log"
    lin2 = !logl2

    nb1 = model.distr[1] == "NegativeBinomial"
    nb2 = model.distr[2] == "NegativeBinomial"

    q = length(model.pastMean)
    r = length(model.external)

    β0 = θ.β0
    β = θ.β
    η = θ.η
    ϕ = θ.ϕ
    ω = θ.ω

    if lin1 & (β0 <= 0)
        return false
    end

    if q > 0
        if !all(0 .<= β .< 1)
            return false
        end
    end

    if r > 0
        if lin1
            if !all(η[findall(.!model.external)] .>= 0)
                return false
            end
        end

        if lin2
            if !all(η[findall(model.external)] .>= 0)
                return false
            end
        end
    end

    if nb1 | nb2
        if any(ϕ .<= 0)
            return false
        end
    end

    if !(0 <= ω <= 1)
        return false
    end

    return true
end

function parametercheck(θ::Array{T, 1} where T<: AbstractFloat, model::T where T<:CountModel)
    parametercheck(θ2par(θ, model), model)
end


text:
using DelimitedFiles
using StatsBase

function to_decimal(binary_array)
    to_dec = collect(range(size(binary_array,1)-1,0,size(binary_array,1)))
    Int(sum(binary_array .*  2 .^to_dec))
end

function answer1()
    input = readdlm("input", String)

    data = split.(input,"")
    data = map(x->parse.(Int,x), data)
    data = hcat(data...)
    data
    mode(test)
    data = [data[i,:] for i in 1:size(data,1)]
    data = mode.(data)
    γrate = to_decimal(data)
    ϵrate = to_decimal(-1 .* (data .- 1))
    γrate, ϵrate, γrate * ϵrate
end
@show answer1()

using Statistics
function answer2()
    input = readdlm("input", String)
    input_in_use = copy(input)
    ninput = map(x->parse.(Int,x), input)

    i=0
    while size(input_in_use,1) > 1
        i += 1
        data_in_use = split.(input_in_use,"")
        data_in_use = map(x->parse.(Int,x), data_in_use)
        data_in_use = hcat(data_in_use...)
        data_in_use = data_in_use[i,:]
        # @show data_in_use
        # @show median(data_in_use)
        ninput = map(x->parse.(Int,x), input_in_use)
        if median(data_in_use) > 0
            ii = permutedims(hcat(split.(input_in_use, "")...))[: , i] .== "1"
        else
            ii = permutedims(hcat(split.(input_in_use, "")...))[: , i] .== "0"
        end
        # @show ii
        input_in_use = input_in_use[ii]
    end
    O2_rating = to_decimal(map(x->parse.(Int,x), split(input_in_use[1],"")))

    input_in_use = copy(input)
    ninput = map(x->parse.(Int,x), input)
    i=0
    while size(input_in_use,1) > 1
        i += 1
        data_in_use = split.(input_in_use,"")
        data_in_use = map(x->parse.(Int,x), data_in_use)
        data_in_use = hcat(data_in_use...)
        data_in_use = data_in_use[i,:]
        # @show data_in_use
        # @show median(data_in_use)
        ninput = map(x->parse.(Int,x), input_in_use)
        if median(data_in_use) > 0
            # ii = ninput .>= 10^(size(data,1)-i)
            ii = permutedims(hcat(split.(input_in_use, "")...))[: , i] .== "0"
        else
            ii = permutedims(hcat(split.(input_in_use, "")...))[: , i] .== "1"
        end
        # @show ii
        input_in_use = input_in_use[ii]
    end
    CO2_rating = to_decimal(map(x->parse.(Int,x), split(input_in_use[1],"")))

    O2_rating, CO2_rating, O2_rating * CO2_rating
end
test = answer2()
permutedims(hcat(split.(test, "")...))[: , 2] .== "1"


text:
M_example = M(model,u₀,u₀_type,tspan,[],[10],15,Rodas5)

rates_example = ks_γϵ
sol = solve_prob(M_example,ks_γϵ;tᵢ=[])

ŷ = sum(sol,dims=1)
mean_t = mean_lifetime_linear(sol)

L = ŷ / mean_t

pyplot()
# plot(sol)
plot(sol.t,ŷ',linewidth=5,label="P")
plot!(sol.t,L',linewidth=5,label="pdf")
xlabel!("t")
ylabel!("distribution")