text:
module TurbulenceConvectionUtils

import Logging
using JLD2
using JSON
using Random
using DocStringExtensions

using ..ReferenceModels
import ..ReferenceModels: NameList

using ..ReferenceStats
# EKP modules
using EnsembleKalmanProcesses.ParameterDistributions
import EnsembleKalmanProcesses: construct_initial_ensemble
using TurbulenceConvection

using ..HelperFuncs

export ModelEvaluator
export run_SCM, run_SCM_handler, run_reference_SCM
export generate_scm_input, get_gcm_les_uuid, eval_single_ref_model
export save_full_ensemble_data
export precondition

"""
    ModelEvaluator

A structure containing the information required to perform
a forward model evaluation at a given parameter vector.

# Fields

$(TYPEDFIELDS)

# Constructors

$(METHODLIST)
"""
Base.@kwdef struct ModelEvaluator{FT <: Real}
    "Parameter vector in constrained (i.e. physical) space"
    param_cons::Vector{FT}
    "Parameter names associated with parameter vector"
    param_names::Vector{String}
    "Vector of reference models"
    ref_models::Vector{ReferenceModel}
    "Reference statistics for the inverse problem"
    ref_stats::ReferenceStatistics

    function ModelEvaluator(
        param_cons::Vector{FT},
        param_names::Vector{String},
        RM::Vector{ReferenceModel},
        RS::ReferenceStatistics,
    ) where {FT <: Real}
        return new{FT}(param_cons, param_names, RM, RS)
    end
end

"""
    run_SCM(
        u::Vector{FT},
        u_names::Vector{String},
        RM::Vector{ReferenceModel},
        RS::ReferenceStatistics;
        error_check::Bool = false,
        namelist_args = nothing,
        failure_handler = "high_loss",
    ) where {FT <: Real}
    run_SCM(
        ME::ModelEvaluator;
        error_check::Bool = false,
        namelist_args = nothing,
        failure_handler = "high_loss",
    ) where {FT <: Real}

Run the single-column model (SCM) using a set of parameters
and return the value of outputs defined in y_names, possibly
after normalization and projection onto lower dimensional
space using PCA.

The function also outputs a boolean diagnosing whether the
requested SCM simulation failed.

Inputs:

 - `u`               :: Values of parameters to be used in simulations.
 - `u_names`         :: SCM names for parameters `u`.
 - `RM`              :: Vector of `ReferenceModel`s
 - `RS`              :: reference statistics for simulation
 - `error_check`     :: Returns as an additional argument whether the SCM call errored.
 - `namelist_args`   :: Additional arguments passed to the TurbulenceConvection namelist.
 - `failure_handler` :: Method used to handle failed simulations.

Outputs:

 - `sim_dirs`    :: Vector of simulation output directories
 - `g_scm`       :: Vector of model evaluations concatenated for all flow configurations.
 - `g_scm_pca`   :: Projection of `g_scm` onto principal subspace spanned by eigenvectors.
 - `model_error` :: Whether the simulation errored with the requested configuration.
"""
function run_SCM(
    u::Vector{FT},
    u_names::Vector{String},
    RM::Vector{ReferenceModel},
    RS::ReferenceStatistics;
    error_check::Bool = false,
    namelist_args = nothing,
    failure_handler = "high_loss",
) where {FT <: Real}

    mkpath(joinpath(pwd(), "tmp"))
    result_arr = map(x -> eval_single_ref_model(x..., RS, u, u_names, namelist_args), enumerate(RM))
    # Unpack
    sim_dirs, g_scm, g_scm_pca, sim_errors = (getindex.(result_arr, i) for i in 1:4)
    g_scm = vcat(g_scm...)
    g_scm_pca = vcat(g_scm_pca...)

    model_error = any(sim_errors)
    if model_error && failure_handler == "high_loss"
        g_scm .= 1e5
        g_scm_pca .= 1e5
    end

    @info "Length of g_scm (full, pca): ($(length(g_scm)), $(length(g_scm_pca)))"
    if error_check
        return sim_dirs, g_scm, g_scm_pca, model_error
    else
        return sim_dirs, g_scm, g_scm_pca
    end
end
function run_SCM(
    ME::ModelEvaluator;
    error_check::Bool = false,
    namelist_args = nothing,
    failure_handler = "high_loss",
) where {FT <: Real}
    return run_SCM(
        ME.param_cons,
        ME.param_names,
        ME.ref_models,
        ME.ref_stats,
        error_check = error_check,
        namelist_args = namelist_args,
        failure_handler = failure_handler,
    )
end

"""
    eval_single_ref_model(
        m_index::IT,
        m::ReferenceModel,
        RS::ReferenceStatistics,
        u::Vector{FT},
        u_names::Vector{String},
        namelist_args = nothing,
    ) where {FT <: Real, IT <: Int}

Runs the single-column model (SCM) under a single configuration
(i.e., ReferenceModel) using a set of parameters u, and returns
the forward model evaluation in both the original and the latent
PCA space.

Inputs:

 - `m_index`       :: The index of the ReferenceModel within the overarching
                    ref_models vector used to construct the ReferenceStatistics.
 - `m`             :: A ReferenceModel.
 - `RS`            :: reference statistics for simulation
 - `u`             :: Values of parameters to be used in simulations.
 - `u_names`       :: SCM names for parameters `u`.
 - `namelist_args` :: Additional arguments passed to the TurbulenceConvection namelist.

Outputs:

 - `sim_dir`     ::  Simulation output directory.
 - `g_scm`       :: Forward model evaluation in original output space.
 - `g_scm_pca`   :: Projection of `g_scm` onto principal subspace spanned by eigenvectors.
 - `model_error` :: Whether the simulation errored with the requested configuration.
"""
function eval_single_ref_model(
    m_index::IT,
    m::ReferenceModel,
    RS::ReferenceStatistics,
    u::Vector{FT},
    u_names::Vector{String},
    namelist_args = nothing,
) where {FT <: Real, IT <: Int}
    # create temporary directory to store SCM data in
    tmpdir = mktempdir(joinpath(pwd(), "tmp"))
    # run TurbulenceConvection.jl. Get output directory for simulation data
    sim_dir, model_error = run_SCM_handler(m, tmpdir, u, u_names, namelist_args)
    filename = get_stats_path(sim_dir)
    if model_error
        g_scm = fill(NaN, length(get_z_obs(m)) * length(m.y_names))
    else
        g_scm = get_profile(m, filename, z_scm = get_z_obs(m))
        g_scm = normalize_profile(g_scm, length(m.y_names), RS.norm_vec[m_index])
    end
    # perform PCA reduction
    g_scm_pca = RS.pca_vec[m_index]' * g_scm
    return sim_dir, g_scm, g_scm_pca, model_error
end

"""
    run_reference_SCM(m::ReferenceModel; overwrite::Bool = false, run_single_timestep = true)

Run the single-column model (SCM) for a reference model object
using default parameters, and write the output to file.

Inputs:
 - `m`                    :: A `ReferenceModel`.
 - `overwrite`            :: if true, run TC.jl and overwrite existing simulation files.
 - `run_single_timestep`  :: if true, run only one time step.
"""
function run_reference_SCM(m::ReferenceModel; overwrite::Bool = false, run_single_timestep = true)
    output_dir = scm_dir(m)
    if ~isdir(joinpath(output_dir, "stats")) | overwrite
        namelist = get_scm_namelist(m)

        default_t_max = namelist["time_stepping"]["t_max"]
        default_adapt_dt = namelist["time_stepping"]["adapt_dt"]
        if run_single_timestep
            # Run only 1 timestep -- since we don't need output data, only simulation config
            namelist["time_stepping"]["adapt_dt"] = false
            namelist["time_stepping"]["t_max"] = namelist["time_stepping"]["dt_min"]
        end
        # run TurbulenceConvection.jl
        logger = Logging.ConsoleLogger(stderr, Logging.Warn)
        _, ret_code = Logging.with_logger(logger) do
            main1d(namelist; time_run = false)
        end
        if ret_code ≠ :success
            @warn "Default TurbulenceConvection.jl simulation $(basename(m.y_dir)) failed."
        end
        if run_single_timestep
            # reset t_max to default and overwrite stored namelist file
            namelist["time_stepping"]["t_max"] = default_t_max
            namelist["time_stepping"]["adapt_dt"] = default_adapt_dt
            open(namelist_directory(output_dir, m), "w") do io
                JSON.print(io, namelist, 4)
            end
        end
    end
end

"""
    run_SCM_handler(
        m::ReferenceModel,
        tmpdir::String,
        u::Array{FT, 1},
        u_names::Array{String, 1},
        namelist_args = nothing,
    ) where {FT<:AbstractFloat}

Run a case using a set of parameters `u_names` with values `u`,
and return directory pointing to where data is stored for simulation run.

Inputs:

 - m             :: Reference model
 - tmpdir        :: Temporary directory to store simulation results in
 - u             :: Values of parameters to be used in simulations.
 - u_names       :: SCM names for parameters `u`.
 - namelist_args :: Additional arguments passed to the TurbulenceConvection namelist.

Outputs:

 - output_dir   :: directory containing output data from the SCM run.
 - model_error   :: Boolean specifying whether the simulation failed.
"""
function run_SCM_handler(
    m::ReferenceModel,
    tmpdir::String,
    u::Vector{FT},
    u_names::Vector{String},
    namelist_args = nothing,
) where {FT <: AbstractFloat}

    # fetch namelist
    namelist = get_scm_namelist(m)
    # output subset of variables needed for calibration
    namelist["stats_io"]["calibrate_io"] = true

    # run TurbulenceConvection.jl with modified parameters
    return run_SCM_handler(
        m.case_name,
        tmpdir;
        u = u,
        u_names = u_names,
        namelist = namelist,
        namelist_args = namelist_args,
        uuid = basename(tmpdir), # set random uuid
        les = get(namelist["meta"], "lesfile", nothing),
    )
end

"""
    run_SCM_handler(
        case_name::String,
        out_dir::String;
        u::Vector{FT},
        u_names::Vector{String},
        namelist::Union{Dict, Nothing} = nothing,
        namelist_args::Union{Tuple, Nothing} = nothing,
        uuid::String = "01",
        les::Union{NamedTuple, String} = nothing,
    )

Run a TurbulenceConvection.jl case and return directory pointing to where data is stored for simulation run.

Inputs:
 - case_name     :: case name
 - out_dir       :: Directory to store simulation results in.
 Optional Inputs:
 - u             :: Values of parameters to be used in simulations.
 - u_names       :: SCM names for parameters `u`.
 - namelist      :: namelist to use for simulation.
 - namelist_args :: Additional arguments passed to the TurbulenceConvection namelist.
 - uuid          :: uuid of SCM run
 - les           :: path to LES stats file, or NamedTuple with keywords {forcing_model, month, experiment, cfsite_number} needed to specify path.
 Outputs:
 - output_dir   :: directory containing output data from the SCM run.
 - model_error   :: Boolean specifying whether the simulation failed.
"""
function run_SCM_handler(
    case_name::String,
    out_dir::String;
    u::Vector{FT},
    u_names::Vector{String},
    namelist::Union{Dict, Nothing} = nothing,
    namelist_args::Union{Vector, Nothing} = nothing,
    uuid::String = "01",
    les::Union{NamedTuple, String, Nothing} = nothing,
) where {FT <: AbstractFloat}
    model_error = false

    # fetch default namelist if not provided
    if isnothing(namelist)
        namelist = NameList.default_namelist(case_name)
    end

    namelist["meta"]["uuid"] = uuid
    # set output dir to `out_dir`
    namelist["output"]["output_root"] = out_dir

    u_names, u = create_parameter_vectors(u_names, u)
    # Set optional namelist args
    if !isnothing(namelist_args)
        for namelist_arg in namelist_args
            change_entry!(namelist, namelist_arg)
        end
    end

    # update learnable parameter values
    @assert length(u_names) == length(u)
    for (pName, pVal) in zip(u_names, u)
        if haskey(namelist["turbulence"]["EDMF_PrognosticTKE"], pName)
            namelist["turbulence"]["EDMF_PrognosticTKE"][pName] = pVal
        elseif haskey(namelist["microphysics"], pName)
            namelist["microphysics"][pName] = pVal
        elseif haskey(namelist["time_stepping"], pName)
            namelist["time_stepping"][pName] = pVal
        else
            throw(
                ArgumentError("Parameter $pName cannot be calibrated. Consider adding namelist dictionary if needed."),
            )
        end
    end

    if case_name == "LES_driven_SCM"
        if isnothing(les)
            error("les path or keywords required for LES_driven_SCM case!")
        elseif isa(les, NamedTuple)
            les = get_stats_path(
                get_cfsite_les_dir(
                    les.cfsite_number;
                    forcing_model = les.forcing_model,
                    month = les.month,
                    experiment = les.experiment,
                ),
            )
        end
        namelist["meta"]["lesfile"] = les
    end

    # run TurbulenceConvection.jl with modified parameters
    logger = Logging.ConsoleLogger(stderr, Logging.Warn)
    _, ret_code = Logging.with_logger(logger) do
        main1d(namelist; time_run = false)
    end
    if ret_code ≠ :success
        model_error = true
        message = ["TurbulenceConvection.jl simulation $out_dir failed with parameters: \n"]
        append!(message, ["$param_name = $param_value \n" for (param_name, param_value) in zip(u_names, u)])
        @warn join(message)
    end

    return data_directory(out_dir, case_name, uuid), model_error
end

"""
    create_parameter_vectors(u_names::Vector{String}, u::Vector{FT}) where {FT <: AbstractFloat}

Given vector of parameter names and corresponding values, combine any vector components
into single parameter vectors for input into SCM.

Inputs:

 - `u_names` :: SCM names for parameters `u`, which may contain vector components.
 - `u` :: Values of parameters to be used in simulations, which may contain vector components.

Outputs:

 -  `u_names_out` :: SCM names for parameters `u`.
 -  `u_out` :: Values of parameters to be used in simulations.
"""
function create_parameter_vectors(u_names::Vector{String}, u::Vector{FT}) where {FT <: AbstractFloat}

    u_names_out = String[]
    u_out = []

    vector_param_inds = occursin.(r"{?}", u_names)
    pv_name_elem = rsplit.(u_names[vector_param_inds], "_", limit = 2) # get param name and element index
    u_vec_names, uvi = ~isempty(pv_name_elem) ? eachrow(reduce(hcat, pv_name_elem)) : ([], [])  # "transpose" `pv_name_elem`
    u_vec_inds = @. parse(Int64, only(split(uvi, (('{', '}'),), keepempty = false)))  # get `i` from "{i}" as Int64

    # collect scalar parameters
    scalar_param_inds = .!vector_param_inds
    append!(u_names_out, u_names[scalar_param_inds])
    append!(u_out, u[scalar_param_inds])

    # collect vector parameters
    for u_name in unique(u_vec_names)
        u_name_inds = u_vec_names .== u_name
        u_vals = u[vector_param_inds][u_name_inds]
        u_vals_sort_inds = u_vec_inds[u_name_inds]
        permute!(u_vals, u_vals_sort_inds)
        push!(u_names_out, u_name)
        push!(u_out, u_vals)
    end

    return u_names_out, u_out
end

"""
    generate_scm_input(
        model_evaluators::Vector{ModelEvaluator{FT}},
        outdir_path::String = pwd(),
    ) where {FT <: AbstractFloat}

Writes to file a set of ModelEvaluator used to initialize SCM
evaluations at different parameter vectors, as well as their
assigned numerical version.
"""
function generate_scm_input(
    model_evaluators::Vector{ModelEvaluator{FT}},
    outdir_path::String = pwd(),
    batch_indices::Union{Vector{Int}, Nothing} = nothing,
) where {FT <: AbstractFloat}
    # Generate versions conditioned on being unique within the batch.
    used_versions = Vector{Int}()
    for model_evaluator in model_evaluators
        version = rand(11111:99999)
        while version in used_versions
            version = rand(11111:99999)
        end
        jldsave(scm_init_path(outdir_path, version); model_evaluator, version, batch_indices)
        push!(used_versions, version)
    end
    return used_versions
end

"""
    get_gcm_les_uuid(
        cfsite_number::Integer;
        forcing_model::String,
        month::Integer,
        experiment::String,)

Generate unique and self-describing uuid given information about a GCM-driven LES
simulation from [Shen2022](@cite).
"""
function get_gcm_les_uuid(
    cfsite_number::Integer;
    forcing_model::String = "HadGEM2-A",
    month::Integer = 7,
    experiment::String = "amip",
)
    cfsite_number = string(cfsite_number)
    month = string(month, pad = 2)
    return join([cfsite_number, forcing_model, month, experiment], '_')
end

""" Save full EDMF data from every ensemble"""
function save_full_ensemble_data(save_path, sim_dirs_arr, ref_models)
    # get a simulation directory `.../Output.SimName.UUID`, and corresponding parameter name
    for (ens_i, sim_dirs) in enumerate(sim_dirs_arr)  # each ensemble returns a list of simulation directories
        ens_i_path = joinpath(save_path, "ens_$ens_i")
        mkpath(ens_i_path)
        for (ref_model, sim_dir) in zip(ref_models, sim_dirs)
            scm_name = ref_model.case_name
            # Copy simulation data to output directory
            dirname = splitpath(sim_dir)[end]
            @assert dirname[1:7] == "Output."  # sanity check
            # Stats file
            tmp_data_path = joinpath(sim_dir, "stats/Stats.$scm_name.nc")
            save_data_path = joinpath(ens_i_path, "Stats.$scm_name.$ens_i.nc")
            cp(tmp_data_path, save_data_path)
            # namefile
            tmp_namefile_path = namelist_directory(sim_dir, scm_name)
            save_namefile_path = namelist_directory(ens_i_path, scm_name)
            cp(tmp_namefile_path, save_namefile_path)
        end
    end
end

"""
    precondition(
        param::Vector{FT},
        priors,
        ref_models::Vector{ReferenceModel},
        ref_stats::ReferenceStatistics,
        namelist_args = nothing;
        counter::Integer = 0,
        max_counter::Integer = 10,
    ) where {FT <: Real}

Substitute parameter vector `param` by a parameter vector drawn
from the same prior, conditioned on the forward model being stable.

Inputs:

 - `param`      :: A parameter vector that may possibly result in unstable
    forward model evaluations (in unconstrained space).
 - `priors`      :: Priors from which the parameters were drawn.
 - `ref_models`  :: Vector of ReferenceModels to check stability for.
 - `ref_stats`   :: ReferenceStatistics of the ReferenceModels.
 - `namelist_args` :: Arguments passed to the TC.jl namelist.
 - `counter` :: Accumulator tracking number of recursive calls to preconditioner.
 - `max_counter` :: Maximum number of recursive calls to the preconditioner.

Outputs:

 - `new_param` :: A new parameter vector drawn from the prior, conditioned on simulations being stable (in unconstrained space).
"""
function precondition(
    param::Vector{FT},
    priors::ParameterDistribution,
    ref_models::Vector{ReferenceModel},
    ref_stats::ReferenceStatistics,
    namelist_args = nothing;
    counter::Integer = 0,
    max_counter::Integer = 10,
) where {FT <: Real}
    param_names = priors.name
    # Wrapper around SCM
    g_(u::Array{Float64, 1}) =
        run_SCM(u, param_names, ref_models, ref_stats, error_check = true, namelist_args = namelist_args)

    param_cons = deepcopy(transform_unconstrained_to_constrained(priors, param))
    _, _, _, model_error = g_(param_cons)
    if model_error && counter < max_counter
        message = ["Unstable parameter vector found: \n"]
        append!(message, ["$param_name = $param \n" for (param_name, param) in zip(param_names, param_cons)])
        @warn join(message)
        @warn "Sampling new parameter vector from prior..."
        return precondition(
            vec(construct_initial_ensemble(priors, 1)),
            priors,
            ref_models,
            ref_stats,
            namelist_args,
            counter = counter + 1,
            max_counter = max_counter,
        )
    elseif model_error
        @error "Number of recursive calls to preconditioner exceeded $(max_counter). Returning last failed parameter."
        return param
    else
        @info "Preconditioning finished."
        return param
    end
end

"""
    precondition(ME::ModelEvaluator, priors)

Substitute the parameter vector of a ModelEvaluator by another
one drawn from the given `priors`, conditioned on the forward
model being stable.

Inputs:
 - ME          :: A ModelEvaluator.
 - priors      :: Priors from which the parameters were drawn.
Outputs:
 - A preconditioned ModelEvaluator.
"""
function precondition(ME::ModelEvaluator, priors; namelist_args = nothing)
    # Precondition in unconstrained space
    u_orig = transform_constrained_to_unconstrained(priors, ME.param_cons)
    u = precondition(u_orig, priors, ME.ref_models, ME.ref_stats, namelist_args)
    # Transform back to constrained space
    param_cons = transform_unconstrained_to_constrained(priors, u)
    return ModelEvaluator(param_cons, ME.param_names, ME.ref_models, ME.ref_stats)
end

end # module


text:
# Implementation of d-dimensional Boomerang and Bouncy particle samplers (the two
# most known not-factorised PDMC)
using LinearAlgebra

grad_correct!(y, x, F::Union{BouncyParticle, ZigZag}) = y
function grad_correct!(y, x, F::FactBoomerang)
    @. y -= x - F.μ
    y
end
function grad_correct!(y, x, F::Boomerang)
    y .-= (F.L'\(F.L\(x - F.μ)))
    y
end
λ(∇ϕx, θ, F::Union{BouncyParticle, Boomerang}) = pos(dot(∇ϕx, θ))

#=
function refresh!(rng, θ, F::BouncyParticle)
    ρ̄ = sqrt(1-F.ρ^2)
    U = F.U
    θ .= F.ρ*θ + ρ̄*(U\randn(rng, length(θ)))
    θ
end
=#

# Here use sparsity as the factorised samplers
function ab(x, θ, C::GlobalBound, ∇ϕx, v, B::BouncyParticle)
    (C.c + θ'*(B.Γ*(x-B.μ)), θ'*(B.Γ*θ), Inf)
end
function ab(x, θ, C::LocalBound, ∇ϕx, v, B::BouncyParticle)
    (C.c + dot(θ, ∇ϕx), v, 2.0/C.c/norm(θ, Inf))
end

function ab(x, θ, C::GlobalBound, ∇ϕx, v, B::Boomerang)
    (sqrt(normsq(θ) + normsq((x - B.μ)))*C.c, 0.0, Inf)
end
ab(x, θ, c, Flow) =  ab(x, θ, GlobalBound(c), nothing, nothing, Flow)

function event(t, x, θ, Z::Union{BouncyParticle,Boomerang})
    t, copy(x), copy(θ), nothing
end

function next_time(t, abc, z = rand(rng))
    Δt = poisson_time(abc[1], abc[2], z)
    if Δt > abc[3]
        return t + abc[3], true
    else
        return t + Δt, false
    end
end

function pdmp_inner!(rng, ∇ϕ!, ∇ϕx, t, x, θ, c::Bound, abc, (t′, renew), τref, v, (acc, num),
     Flow::Union{BouncyParticle, Boomerang}, args...; subsample=false, factor=1.5, adapt=false)
    while true
        if τref < t′
            t, x, θ = move_forward!(τref - t, t, x, θ, Flow)
            θ = refresh!(rng, θ, Flow)
            ∇ϕx, v = ∇ϕ!(∇ϕx, t, x, θ, args...)
            ∇ϕx = grad_correct!(∇ϕx, x, Flow)
            l = λ(∇ϕx, θ, Flow)
            τref = t + waiting_time_ref(rng, Flow)
            abc = ab(x, θ, c, ∇ϕx, v, Flow)
            t′, renew = next_time(t, abc, rand(rng))
            return t, x, θ, (acc, num), c, abc, (t′, renew), τref, v
        elseif renew
            τ = t′ - t
            t, x, θ = move_forward!(τ, t, x, θ, Flow)
            ∇ϕx, v = ∇ϕ!(∇ϕx, t, x, θ, args...)
            ∇ϕx = grad_correct!(∇ϕx, x, Flow)
            abc = ab(x, θ, c, ∇ϕx, v, Flow)
            t′, renew = next_time(t, abc, rand(rng))
        else
            τ = t′ - t
            t, x, θ = move_forward!(τ, t, x, θ, Flow)
            ∇ϕx, v = ∇ϕ!(∇ϕx, t, x, θ, args...)
            ∇ϕx = grad_correct!(∇ϕx, x, Flow)
            l, lb = λ(∇ϕx, θ, Flow), pos(abc[1] + abc[2]*τ)
            num += 1
            if rand(rng)*lb <= l
                acc += 1
                if l > lb
                    !adapt && error("Tuning parameter `c` too small.")
                    c *= factor
                end
                θ = reflect!(∇ϕx, x, θ, Flow)
                ∇ϕx, v = ∇ϕ!(∇ϕx, t, x, θ, args...)
                abc = ab(x, θ, c, ∇ϕx, v, Flow)
                t′, renew = next_time(t, abc, rand(rng))
                !subsample && return t, x, θ, (acc, num), c, abc, (t′, renew), τref, v
            else
                abc = ab(x, θ, c, ∇ϕx, v, Flow)
                t′, renew = next_time(t, abc, rand(rng))
            end
        end
    end
end

addnothing(u::Tuple) = u
addnothing(u) = u, nothing

"""
    pdmp(∇ϕ!, t0, x0, θ0, T, c::Bound, Flow::Union{BouncyParticle, Boomerang}; adapt=false, factor=2.0)

Run a Bouncy particle sampler (`BouncyParticle`) or `Boomerang` sampler from time,
location and velocity `t0, x0, θ0` until time `T`. `∇ϕ!(y, x)` writes the gradient
of the potential (neg. log density) into y.
`c` is a tuning parameter for the upper bound of the Poisson rate.
If `adapt = false`, `c = c*factor` is tried, otherwise an error is thrown.

It returns a `PDMPTrace` (see [`Trace`](@ref)) object `Ξ`, which can be collected
into pairs `t => x` of times and locations and discretized with `discretize`.
Also returns the `num`ber of total and `acc`epted Poisson events and updated bounds
`c` (in case of `adapt==true` the bounds are multiplied by `factor` if they turn
out to be too small.)
"""
function pdmp(∇ϕ!, t0, x0, θ0, T, c::Bound, Flow::Union{BouncyParticle, Boomerang}, args...; adapt=false, subsample=false, progress=false, progress_stops = 20, islocal = false, seed=Seed(), factor=2.0)
    t, x, θ, ∇ϕx = t0, copy(x0), copy(θ0), copy(θ0)
    rng = Rng(seed)
    Ξ = Trace(t0, x0, θ0, Flow)
    τref = waiting_time_ref(rng, Flow)
    ∇ϕx, v = ∇ϕ!(∇ϕx, t, x, θ, args...)
    ∇ϕx = grad_correct!(∇ϕx, x, Flow)
    num = acc = 0
    #l = 0.0
    abc = ab(x, θ, c, ∇ϕx, v, Flow)
    if progress
        prg = Progress(progress_stops, 1)
    else
        prg = missing
    end
    stops = ismissing(prg) ? 0 : max(prg.n - 1, 0) # allow one stop for cleanup
    tstop = T/stops

    t′, renew = next_time(t, abc, rand(rng))
    while t < T
        t, x, θ, (acc, num), c, abc, (t′, renew), τref, v = pdmp_inner!(rng, ∇ϕ!, ∇ϕx, t, x, θ, c, abc, (t′, renew), τref, v, (acc, num), Flow, args...; subsample=subsample, factor=factor, adapt=adapt)
        push!(Ξ, event(t, x, θ, Flow))

        if t > tstop
            tstop += T/stops
            next!(prg)
        end
    end
    ismissing(prg) || ProgressMeter.finish!(prg)
    return Ξ, (t, x, θ), (acc, num), c
end

wrap(f) = wrap_(f,  methods(f)...)
wrap_(f, args...) = f
@inline wrap_(f, m) = m.nargs <= 4 ? Wrapper(f) : f
struct Wrapper{F}
    f::F
end
(F::Wrapper)(y, t, x, θ, args...) = F.f(y, x, args...), nothing

pdmp(∇ϕ!, t0, x0, θ0, T, c, Flow::Union{BouncyParticle, Boomerang}, args...; nargs...) =
pdmp(Wrapper(∇ϕ!), t0, x0, θ0, T, GlobalBound(c), Flow, args...; nargs...)

ext: jl


text:
module Lattices

using LinearAlgebra

using ..UnitCells: UnitCell, monkhorst_pack_mesh, calc_site_pos!

export Lattice
export loc_to_cell, loc_to_site, site_to_site
export translationally_equivalent_sets
export calc_neighbor_table, sorted_neighbor_table_perm!
export site_to_site_vec!, site_to_site_vec

"""
Represents a finite lattice.
"""
struct Lattice{T}

    "Describes unit cell."
    unit_cell::UnitCell{T}

    "Number of unit cells in direction of first lattice vector."
    L1::Int

    "Number of unit cells in direction of second lattice vector."
    L2::Int

    "Number of unit cells in direction of third lattice vector."
    L3::Int

    "Dimensions of lattice in vector form: [L1,L2,L3]"
    dims::Vector{Int}

    "Number of sites in lattice."
    nsites::Int

    "Number of unit cells in lattice."
    ncells::Int

    "Location of each unit cell in lattice."
    cell_loc::Matrix{Int}

    "Records orbital type of each site in lattice."
    site_to_orbit::Vector{Int}

    "Records which unit cell each site live in."
    site_to_cell::Vector{Int}

    #######################
    ## INNER CONSTRUCTOR ##
    #######################
    """
        Lattice(unit_cell::UnitCell,L1::Int,L2::Int,L3::Int)::Lattice

    Constructor for Lattice type.
    """
    function Lattice(unit_cell::UnitCell{T},L1::Int,L2::Int,L3::Int)::Lattice where {T<:AbstractFloat}

        # making sure dimensions of lattice are valid
        @assert (L1>=1 && L2>=1 && L3>=1)

        # dimensions of lattice
        dims = [L1,L2,L3]

        # number of unit cell in lattice
        ncells = L1*L2*L3

        # orbitals/sites per unit cell
        norbits = unit_cell.norbits

        # calculating number of sites in lattice
        nsites = ncells*norbits

        # allocating array to contain unit cell locations
        cell_loc = Matrix{Int}(undef,3,ncells)

        # allcoating array that maps site to orbit in lattice
        site_to_orbit = Vector{Int}(undef,nsites)

        # allocating array that maps site to unit cell in lattice
        site_to_cell = Vector{Int}(undef,nsites)

        # filling in allocated arrays
        site = 1 # keeps track to site number
        cell = 1 # keeps track of unit cell number

        # iterating over unit cells in lattice
        for l3=0:L3-1
            for l2=0:L2-1
                for l1=0:L1-1
                    # recording location of unit cell in lattice
                    cell_loc[1,cell] = l1
                    cell_loc[2,cell] = l2
                    cell_loc[3,cell] = l3
                    # iterating over orbitals in unit cell
                    for orbit=1:norbits
                        # recording the orbit and unit cell associated with site in lattice
                        site_to_orbit[site] = orbit
                        site_to_cell[site] = cell
                        site += 1
                    end
                    cell += 1
                end
            end
        end

        return new{T}(unit_cell,L1,L2,L3,dims,nsites,ncells,cell_loc,site_to_orbit,site_to_cell)
    end

end

########################
## OUTER CONSTRUCTORS ##
########################

function Lattice(unit_cell::UnitCell,L1::Int,L2::Int)::Lattice

    @assert unit_cell.ndim==2
    return Lattice(unit_cell,L1,L2,1)
end

function Lattice(unit_cell::UnitCell,L1::Int)::Lattice

    if unit_cell.ndim==1
        L2 = 1
        L3 = 1
    elseif unit_cell.ndim==2
        L2 = L1
        L3 = 1
    elseif unit_cell.ndim==3
        L2 = L1
        L3 = L1
    end
    return Lattice(unit_cell,L1,L2,L3)
end

############################################
## DEFINING METHODS THAT USE LATTICE TYPE ##
############################################

"""
Given a location of a cell in a lattice, return the corresponding cell.
"""
function loc_to_cell(lattice::Lattice,loc::AbstractVector{Int})::Int

    return loc_to_cell(lattice,loc[1],loc[2],loc[3])
end

function loc_to_cell(lattice::Lattice,l1::Int,l2::Int,l3::Int)::Int

    l1p, l2p, l3p = _pbc!(l1,l2,l3,lattice)
    return l1p + l2p*lattice.L1 + l3p*lattice.L1*lattice.L2 + 1
end

"""
Given the location of a site in the lattice, return the correpsonding site.
"""
function loc_to_site(lattice::Lattice, orbit::Int, loc::AbstractVector{Int})::Int

    return loc_to_site(lattice, orbit, loc[1], loc[2], loc[3])
end

function loc_to_site(lattice::Lattice, orbit::Int, l1::Int, l2::Int=0, l3::Int=0)::Int

    norbits = lattice.unit_cell.norbits::Int
    return norbits * ( loc_to_cell(lattice,l1,l2,l3) - 1 ) + orbit
end

"""
    site_to_site(lattice::Lattice,isite::Int,displacement::AbstractVector{Int},orbit::Int)::Int

Get the new site following a displacement from an initial site.
"""
function site_to_site(lattice::Lattice,isite::Int,displacement::AbstractVector{Int},orbit::Int)::Int

    @assert length(displacement)==3
    # get unit cell
    cell = lattice.site_to_cell[isite]
    # getting location of the unit cell that the initial site lives in
    l1 = lattice.cell_loc[ 1 , cell ]
    l2 = lattice.cell_loc[ 2 , cell ]
    l3 = lattice.cell_loc[ 3 , cell ]
    # displacing the location of the unit cell
    l1 += displacement[1]
    l2 += displacement[2]
    l3 += displacement[3]
    # calculating the final site location after the displacement
    return loc_to_site(lattice, orbit, l1, l2, l3)
end

"""
Constructs translationally equivalent sets of sites in lattice.
The translationally equivalent sets are stored in 7-dimensional array
of size ( 2 x numorbits x L1 x L2 x L3 x norbits x norbits ) where
numorbits is the number of sites of a given orbital type in lattice
i.e. numorbits=nsites/norbits.
"""
function translationally_equivalent_sets(lattice::Lattice)::Array{Int,7}

    # getting info about lattice
    L1      = lattice.L1
    L2      = lattice.L2
    L3      = lattice.L3
    nsites  = lattice.nsites
    norbits = lattice.unit_cell.norbits

    # number of orbitals of a given type in lattice.
    # this is also equal to the number of translationally equivalent pairs of sites
    # in a given translationally equivlent set.
    numorbits = div(nsites,norbits)

    # declaring 7 dimensional array to contain translationally equivalent pairs of sites
    sets = zeros(Int, 2, numorbits, L1, L2, L3, norbits, norbits)

    # counter for tracking numbers of paired sites in translationally equivalent set
    setcount = 0

    # stores second sites in a pairs of sites
    site2 = 0

    # to store displacement in unit cells
    displacement = zeros(Int,3)

    # iterating over all possible combinations of orbitals
    for orbit1 in 1:norbits
        for orbit2 in 1:norbits
            # iterating over all possible unit cell displacements
            for l3 in 0:L3-1
                for l2 in 0:L2-1
                    for l1 in 0:L1-1
                        # reseting counter for tracking numbers of paired sites in set
                        setcount = 0
                        # iterating over sites in lattice of orbital type orbit1
                        for site1 in orbit1:norbits:nsites
                            # incrementing size of set count
                            setcount += 1
                            # setting displacement in unit cells
                            displacement[1] = l1
                            displacement[2] = l2
                            displacement[3] = l3
                            # getting site2 based on displacement vector away from site1
                            site2 = site_to_site(lattice,site1,displacement,orbit2)
                            # recording pair of sites
                            sets[1,setcount,l1+1,l2+1,l3+1,orbit2,orbit1] = site1
                            sets[2,setcount,l1+1,l2+1,l3+1,orbit2,orbit1] = site2
                        end
                    end
                end
            end
        end
    end

    return sets
end

"""
    calc_neighbor_table(lattice::Lattice,orbit1::Int,orbit2::Int,displacement::AbstractVector{Int})::Array{Int,2}

Construct the neighbor table for a certain type of displacement in the lattice.
"""
function calc_neighbor_table(lattice::Lattice,orbit1::Int,orbit2::Int,displacement::AbstractVector{Int},remove_duplicates::Bool=true)::Matrix{Int}

    # number of orbits/sites per unit cell
    norbits = lattice.unit_cell.norbits::Int

    # ensuring valid rule is specified for defining neighbor relations
    @assert length(displacement)==3
    @assert (1<=orbit1<=norbits && 1<=orbit2<=norbits)

    # sites in the lattice
    nsites = lattice.nsites

    # total number pairs of neighbors
    N = div(nsites,norbits)

    # allocating neighbor table array
    neighbor_table = zeros(Int,2,N)

    # keeps track of the number of neighbor relations that have been calcualted.
    neighbor_count = 1

    # only iterates over those sites that are an orbit of type orbit1
    for isite in orbit1:norbits:nsites
        # get final site
        fsite = site_to_site(lattice,isite,displacement,orbit2)
        # record neighbor_table
        neighbor_table[1,neighbor_count] = isite
        neighbor_table[2,neighbor_count] = fsite
        neighbor_count += 1
    end

    # remove duplicate neighbor pairs
    if remove_duplicates
        keep = ones(Bool,N)
        for i in 1:N-1
            if keep[i]
                isite = neighbor_table[1,i]
                fsite = neighbor_table[2,i]
                for j in i+1:N
                    isite′ = neighbor_table[1,j]
                    fsite′ = neighbor_table[2,j]
                    if (isite==isite′ && fsite==fsite′)||(isite==fsite′ && fsite==isite′)
                        keep[j] = false
                    end
                end
            end
        end
        neighbor_table = neighbor_table[:,keep]
    end

    return neighbor_table
end

"""
Returns the permutation that sorts the neighbor table so that the first row is in strictly ascending order,
and for fixed values in the first row, the second row is also in ascending order.
"""
function sorted_neighbor_table_perm!(neighbor_table::Matrix{Int})::Vector{Int}

    @assert size(neighbor_table,1)==2

    # make sure smaller number is always in first column of neighbor table
    for i in 1:size(neighbor_table,2)
        c1 = neighbor_table[1,i]
        c2 = neighbor_table[2,i]
        if c1 > c2
            neighbor_table[1,i] = c2
            neighbor_table[2,i] = c1
        end
    end

    vals = maximum(neighbor_table)*neighbor_table[1,:] + neighbor_table[2,:]
    perm = sortperm(vals)
    return perm
end

"""
    site_to_site_vec!(vector::AbstractVector{Float64},lattice::Lattice,unit_cell::UnitCell,site1::Int,site2::Int)

Calculates the displacement vector between two sites in the lattice accounting for periodic boundary conditions.
"""
function site_to_site_vec!(vector::AbstractVector{T},lattice::Lattice,site1::Int,site2::Int) where {T<:AbstractFloat}

    # iterating over each lattice vector direction
    delta = 0 # shift in unit cells
    for direction in 1:3
        # displacement in unit cells
        delta = _cell_displacement(lattice,site1,site2,direction)
        # updating displacement vector
        @. vector += delta * unit_cell.lvecs[:,direction]
    end
    # accounting for basis vector positions of intial and final orbitals
    @. vector += lattice.unit_cell.bvecs[:,lattice.site_to_orbit[site1]] - lattice.unit_cell.bvecs[:,lattice.site_to_orbit[site2]]
    return nothing
end

function site_to_site_vec(lattice::Lattice,site1::Int,site2::Int)::Vector{T} where {T<:AbstractFloat}

    vector = zeros(T,3)
    site_to_site_vec!(vector,lattice,site1,site2)
    return vector
end

############################################################
## PRIVATE FUNCTIONS NOT TO BE CALLED OUTSIDE THIS MODULE ##
############################################################

"""
Applies periodic boundary conditions.
"""
function _pbc!(loc::AbstractVector{Int},lattice::Lattice)

    @. loc = mod(loc,lattice.dims)
    return nothing
end

function _pbc!(l1::Int,l2::Int,l3::Int,lattice::Lattice)::Tuple{Int,Int,Int}

    l1new = mod(l1,lattice.L1)
    l2new = mod(l2,lattice.L2)
    l3new = mod(l3,lattice.L3)

    return l1new, l2new, l3new
end

"""
    _cell_displacement(lattice::Lattice,site1::Int,site2::Int,direction::Int)::Int

Calculates a displacement in unit cells between two sites in the lattice
in the 'direction' of a specified lattice vector, accounting for peridic
boundary conditions.
"""
function _cell_displacement(lattice::Lattice,site1::Int,site2::Int,direction::Int)::Int

    @assert 1<=direction<=3
    # width of lattice in unit cells in direction of specified lattice vector
    L = lattice.dims[direction]
    # half-width of lattice in unit cells in direction of specified lattice vector
    Lhalf = div(L,2)
    # unit cell that each site lives in
    cell1 = lattice.site_to_cell[site1]
    cell2 = lattice.site_to_cell[site2]
    # displacement in unit cells
    delta = lattice.cell_loc[direction,cell2] - lattice.cell_loc[direction,cell1]
    # accounting for periodic boundary conditions
    if delta > Lhalf
        delta -= L
    elseif delta < -Lhalf
        delta += L
    end

    return delta
end

end

max_stars_repo_name: cohensbw/Langevin

max_issues_repo_name: cohensbw/ElPhDynamics

max_forks_repo_name: cohensbw/ElPhDynamics


text:
immutable SArray{Size, T, N, L} <: StaticArray{T, N}
    data::NTuple{L,T}

    function SArray(x::NTuple{L,T})
        check_sarray_parameters(Val{Size}, T, Val{N}, Val{L})
        new(x)
    end

    function SArray(x::NTuple{L})
        check_sarray_parameters(Val{Size}, T, Val{N}, Val{L})
        new(convert_ntuple(T, x))
    end
end

@generated function check_sarray_parameters{Size,T,N,L}(::Type{Val{Size}}, ::Type{T}, ::Type{Val{N}}, ::Type{Val{L}})
    if !(isa(Size, Tuple{Vararg{Int}}))
        error("SArray parameter Size must be a tuple of Ints (e.g. `SArray{(3,3)}`)")
    end

    if L != prod(Size) || L < 0 || minimum(Size) < 0 || length(Size) != N
        error("Size mismatch in SArray parameters. Got size $Size, dimension $N and length $L.")
    end

    return nothing
end

@generated function (::Type{SArray{Size,T,N}}){Size,T,N}(x)
    return quote
        $(Expr(:meta, :inline))
        SArray{Size,T,N,$(prod(Size))}(x)
    end
end

@generated function (::Type{SArray{Size,T}}){Size,T}(x)
    return quote
        $(Expr(:meta, :inline))
        SArray{Size,T,$(length(Size)),$(prod(Size))}(x)
    end
end

@generated function (::Type{SArray{Size}}){Size}(x)
    return quote
        $(Expr(:meta, :inline))
        SArray{Size,$(promote_tuple_eltype(x)),$(length(Size)),$(prod(Size))}(x)
    end
end

@inline SArray(a::StaticArray) = SArray{size(typeof(a))}(Tuple(a))

# Some more advanced constructor-like functions
@inline eye{Size}(::Type{SArray{Size}}) = eye(SArray{Size,Float64})
@inline zeros{Size}(::Type{SArray{Size}}) = zeros(SArray{Size,Float64})
@inline ones{Size}(::Type{SArray{Size}}) = ones(SArray{Size,Float64})

####################
## SArray methods ##
####################

@pure size{Size}(::Union{SArray{Size},Type{SArray{Size}}}) = Size
@pure size{Size,T}(::Type{SArray{Size,T}}) = Size
@pure size{Size,T,N}(::Type{SArray{Size,T,N}}) = Size
@pure size{Size,T,N,L}(::Type{SArray{Size,T,N,L}}) = Size

function getindex(v::SArray, i::Integer)
    Base.@_inline_meta
    v.data[i]
end

@inline Tuple(v::SArray) = v.data

@inline function Base.unsafe_convert{Size,T}(::Type{Ptr{T}}, a::SArray{Size,T})
    Base.unsafe_convert(Ptr{T}, Base.data_pointer_from_objref(a))
end

macro SArray(ex)
    if !isa(ex, Expr)
        error("Bad input for @SArray")
    end

    if ex.head == :vect  # vector
        return esc(Expr(:call, SArray{(length(ex.args),)}, Expr(:tuple, ex.args...)))
    elseif ex.head == :ref # typed, vector
        return esc(Expr(:call, Expr(:curly, :SArray, ((length(ex.args)-1),), ex.args[1]), Expr(:tuple, ex.args[2:end]...)))
    elseif ex.head == :hcat # 1 x n
        s1 = 1
        s2 = length(ex.args)
        return esc(Expr(:call, SArray{(s1, s2)}, Expr(:tuple, ex.args...)))
    elseif ex.head == :typed_hcat # typed, 1 x n
        s1 = 1
        s2 = length(ex.args) - 1
        return esc(Expr(:call, Expr(:curly, :SArray, (s1, s2), ex.args[1]), Expr(:tuple, ex.args[2:end]...)))
    elseif ex.head == :vcat
        if isa(ex.args[1], Expr) && ex.args[1].head == :row # n x m
            # Validate
            s1 = length(ex.args)
            s2s = map(i -> ((isa(ex.args[i], Expr) && ex.args[i].head == :row) ? length(ex.args[i].args) : 1), 1:s1)
            s2 = minimum(s2s)
            if maximum(s2s) != s2
                error("Rows must be of matching lengths")
            end

            exprs = [ex.args[i].args[j] for i = 1:s1, j = 1:s2]
            return esc(Expr(:call, SArray{(s1, s2)}, Expr(:tuple, exprs...)))
        else # n x 1
            return esc(Expr(:call, SArray{(length(ex.args), 1)}, Expr(:tuple, ex.args...)))
        end
    elseif ex.head == :typed_vcat
        if isa(ex.args[2], Expr) && ex.args[2].head == :row # typed, n x m
            # Validate
            s1 = length(ex.args) - 1
            s2s = map(i -> ((isa(ex.args[i+1], Expr) && ex.args[i+1].head == :row) ? length(ex.args[i+1].args) : 1), 1:s1)
            s2 = minimum(s2s)
            if maximum(s2s) != s2
                error("Rows must be of matching lengths")
            end

            exprs = [ex.args[i+1].args[j] for i = 1:s1, j = 1:s2]
            return esc(Expr(:call, Expr(:curly, :SArray, (s1, s2), ex.args[1]), Expr(:tuple, exprs...)))
        else # typed, n x 1
            return esc(Expr(:call, Expr(:curly, :SArray, (length(ex.args)-1, 1), ex.args[1]), Expr(:tuple, ex.args[2:end]...)))
        end
    elseif isa(ex, Expr) && ex.head == :comprehension
        if length(ex.args) != 1 || !isa(ex.args[1], Expr) || ex.args[1].head != :generator
            error("Expected generator in comprehension, e.g. [f(i,j) for i = 1:3, j = 1:3]")
        end
        ex = ex.args[1]
        n_rng = length(ex.args) - 1
        rng_args = [ex.args[i+1].args[1] for i = 1:n_rng]
        rngs = [eval(current_module(), ex.args[i+1].args[2]) for i = 1:n_rng]
        rng_lengths = map(length, rngs)

        f = gensym()
        f_expr = :($f = ($(Expr(:tuple, rng_args...)) -> $(ex.args[1])))

        # TODO figure out a generic way of doing this...
        if n_rng == 1
            exprs = [:($f($j1)) for j1 in rngs[1]]
        elseif n_rng == 2
            exprs = [:($f($j1, $j2)) for j1 in rngs[1], j2 in rngs[2]]
        elseif n_rng == 3
            exprs = [:($f($j1, $j2, $j3)) for j1 in rngs[1], j2 in rngs[2], j3 in rngs[3]]
        elseif n_rng == 4
            exprs = [:($f($j1, $j2, $j3, $j4)) for j1 in rngs[1], j2 in rngs[2], j3 in rngs[3], j4 in rngs[4]]
        elseif n_rng == 5
            exprs = [:($f($j1, $j2, $j3, $j4, $j5)) for j1 in rngs[1], j2 in rngs[2], j3 in rngs[3], j4 in rngs[4], j5 in rngs[5]]
        elseif n_rng == 6
            exprs = [:($f($j1, $j2, $j3, $j4, $j5, $j6)) for j1 in rngs[1], j2 in rngs[2], j3 in rngs[3], j4 in rngs[4], j5 in rngs[5], j6 in rngs[6]]
        elseif n_rng == 7
            exprs = [:($f($j1, $j2, $j3, $j4, $j5, $j6, $j7)) for j1 in rngs[1], j2 in rngs[2], j3 in rngs[3], j4 in rngs[4], j5 in rngs[5], j6 in rngs[6], j7 in rngs[7]]
        elseif n_rng == 8
            exprs = [:($f($j1, $j2, $j3, $j4, $j5, $j6, $j7, $j8)) for j1 in rngs[1], j2 in rngs[2], j3 in rngs[3], j4 in rngs[4], j5 in rngs[5], j6 in rngs[6], j7 in rngs[7], j8 in rngs[8]]
        else
            error("@SArray only supports up to 8-dimensional comprehensions")
        end

        return quote
            $(Expr(:meta, :inline))
            $(esc(f_expr))
            $(esc(Expr(:call, Expr(:curly, :SArray, (rng_lengths...)), Expr(:tuple, exprs...))))
        end
    elseif isa(ex, Expr) && ex.head == :typed_comprehension
        if length(ex.args) != 2 || !isa(ex.args[2], Expr) || ex.args[2].head != :generator
            error("Expected generator in typed comprehension, e.g. Float64[f(i,j) for i = 1:3, j = 1:3]")
        end
        T = ex.args[1]
        ex = ex.args[2]
        n_rng = length(ex.args) - 1
        rng_args = [ex.args[i+1].args[1] for i = 1:n_rng]
        rngs = [eval(current_module(), ex.args[i+1].args[2]) for i = 1:n_rng]
        rng_lengths = map(length, rngs)

        f = gensym()
        f_expr = :($f = ($(Expr(:tuple, rng_args...)) -> $(ex.args[1])))

        # TODO figure out a generic way of doing this...
        if n_rng == 1
            exprs = [:($f($j1)) for j1 in rngs[1]]
        elseif n_rng == 2
            exprs = [:($f($j1, $j2)) for j1 in rngs[1], j2 in rngs[2]]
        elseif n_rng == 3
            exprs = [:($f($j1, $j2, $j3)) for j1 in rngs[1], j2 in rngs[2], j3 in rngs[3]]
        elseif n_rng == 4
            exprs = [:($f($j1, $j2, $j3, $j4)) for j1 in rngs[1], j2 in rngs[2], j3 in rngs[3], j4 in rngs[4]]
        elseif n_rng == 5
            exprs = [:($f($j1, $j2, $j3, $j4, $j5)) for j1 in rngs[1], j2 in rngs[2], j3 in rngs[3], j4 in rngs[4], j5 in rngs[5]]
        elseif n_rng == 6
            exprs = [:($f($j1, $j2, $j3, $j4, $j5, $j6)) for j1 in rngs[1], j2 in rngs[2], j3 in rngs[3], j4 in rngs[4], j5 in rngs[5], j6 in rngs[6]]
        elseif n_rng == 7
            exprs = [:($f($j1, $j2, $j3, $j4, $j5, $j6, $j7)) for j1 in rngs[1], j2 in rngs[2], j3 in rngs[3], j4 in rngs[4], j5 in rngs[5], j6 in rngs[6], j7 in rngs[7]]
        elseif n_rng == 8
            exprs = [:($f($j1, $j2, $j3, $j4, $j5, $j6, $j7, $j8)) for j1 in rngs[1], j2 in rngs[2], j3 in rngs[3], j4 in rngs[4], j5 in rngs[5], j6 in rngs[6], j7 in rngs[7], j8 in rngs[8]]
        else
            error("@SArray only supports up to 8-dimensional comprehensions")
        end

        return quote
            $(Expr(:meta, :inline))
            $(esc(f_expr))
            $(esc(Expr(:call, Expr(:curly, :SArray, (rng_lengths...), T), Expr(:tuple, exprs...))))
        end
    elseif isa(ex, Expr) && ex.head == :call
        if ex.args[1] == :zeros || ex.args[1] == :ones || ex.args[1] == :rand || ex.args[1] == :randn
            if length(ex.args) == 1
                error("@SArray got bad expression: $(ex.args[1])()")
            else
                return quote
                    $(Expr(:meta, :inline))
                    if isa($(esc(ex.args[2])), DataType)
                        $(ex.args[1])($(esc(Expr(:curly, SArray, Expr(:tuple, ex.args[3:end]...), ex.args[2]))))
                    else
                        $(ex.args[1])($(esc(Expr(:curly, SArray, Expr(:tuple, ex.args[2:end]...)))))
                    end
                end
            end
        elseif ex.args[1] == :fill
            if length(ex.args) == 1
                error("@SArray got bad expression: $(ex.args[1])()")
            elseif length(ex.args) == 2
                error("@SArray got bad expression: $(ex.args[1])($(ex.args[2]))")
            else
                return quote
                    $(Expr(:meta, :inline))
                    $(esc(ex.args[1]))($(esc(ex.args[2])), SArray{$(esc(Expr(:tuple, ex.args[3:end]...)))})
                end
            end
        elseif ex.args[1] == :eye
            if length(ex.args) == 2
                return quote
                    $(Expr(:meta, :inline))
                    eye(SArray{($(esc(ex.args[2])), $(esc(ex.args[2])))})
                end
            elseif length(ex.args) == 3
                # We need a branch, depending if the first argument is a type or a size.
                return quote
                    $(Expr(:meta, :inline))
                    if isa($(esc(ex.args[2])), DataType)
                        eye(SArray{($(esc(ex.args[3])), $(esc(ex.args[3]))), $(esc(ex.args[2]))})
                    else
                        eye(SArray{($(esc(ex.args[2])), $(esc(ex.args[3])))})
                    end
                end
            elseif length(ex.args) == 4
                return quote
                    $(Expr(:meta, :inline))
                    eye(SArray{($(esc(ex.args[3])), $(esc(ex.args[4]))), $(esc(ex.args[2]))})
                end
            else
                error("Bad eye() expression for @SArray")
            end
        else
            error("@SArray only supports the zeros(), ones(), rand(), randn() and eye() functions.")
        end
    else
        error("Bad input for @SArray")
    end
end


text:
@testset "Constraints" begin
    using ONNXNaiveNASflux.NaiveNASflux
    using ONNXNaiveNASflux: SizePseudoTransparent, MeasureNout
    using NaiveNASflux: actdim

    dv(name, invertex, outsize) = dv(name, invertex, nout(invertex), outsize)
    dv(name, invertex, insize, outsize) = fluxvertex(name, Dense(insize, outsize), invertex)
    cv(name, invertex, outsize) = cv(name, invertex, nout(invertex), outsize)
    cv(name, invertex, insize, outsize) =fluxvertex(name, Conv((3,3), insize => outsize, pad=(1,1)), invertex)

    @testset "Reshape" begin
        using ONNXNaiveNASflux: Reshape

        rv(name, invertex, dims) = absorbvertex(name, MeasureNout(Reshape(dims)), invertex; traitdecoration=SizePseudoTransparent)

        @testset "Reshape 2D -> 4D variable batch" begin
            v0 = inputvertex("in", 3)
            v1 = dv("v1", v0, 4)
            v2 = rv("v2", v1, (5,2,3,Colon()))
            v3 = cv("v3", v2, 4)

            @test actdim(v2) == [actdim(layer(v2))] == [3]
            @test layer(v2) isa Reshape

            g = CompGraph(v0, v3)

            @test size(g(ones(Float32, 3, 15))) == (5,2,4,2)

            @test_logs (:warn, r"Could not change nout") Δnout!(v1, 3)

            @test nout.(vertices(g)) == [3, 6, 3, 4]
            @test size(g(ones(Float32, 3, 15))) == (5,2,4,3)
        end

        @testset "Reshape 2D -> 4D variable shape" begin
            v0 = inputvertex("in", 3)
            v1 = dv("v1", v0, 4)
            v2 = rv("v2", v1, (5, Colon(),3, 2))
            v3 = cv("v3", v2, 4)

            @test actdim(v2) == [actdim(layer(v2))] == [3]
            @test layer(v2) isa Reshape

            g = CompGraph(v0, v3)

            @test size(g(ones(Float32, 3, 15))) == (5,2,4,2)

            @test_logs (:warn, r"Could not change nout") Δnout!(v1, 3)

            @test nout.(vertices(g)) == [3, 6, 3, 4]
            @test size(g(ones(Float32, 3, 15))) == (5,3,4,2)
        end

        @testset "Reshape 2D -> 4D variable nout" begin
            v0 = inputvertex("in", 3)
            v1 = dv("v1", v0, 4)
            v2 = rv("v2", v1, (5, 3, Colon(), 2))
            v3 = cv("v3", v2, 2, 4)

            @test actdim(v2) == [actdim(layer(v2))] == [3]
            @test layer(v2) isa Reshape

            g = CompGraph(v0, v3)

            @test size(g(ones(Float32, 3, 15))) == (5,3,4,2)

            # note that this works only because MeasureNout got the size when we evaluated the graph absorbvertex
            # In other cases one might need try_infer_sizes!
            @test_logs (:warn, r"Could not change nout") Δnout!(v1, 3)

            @test nout.(vertices(g)) == [3, 6, 3, 4]
            @test size(g(ones(Float32, 3, 15))) == (5,3,4,2)
        end

        @testset "Reshape 2D -> 4D variable nout 0-dim" begin
            v0 = inputvertex("in", 3)
            v1 = dv("v1", v0, 20)
            v2 = rv("v2", v1, (2, 5, Colon(), 0))
            v3 = cv("v3", v2, 2, 4)

            @test actdim(v2) == [actdim(layer(v2))] == [3]
            @test layer(v2) isa Reshape

            g = CompGraph(v0, v3)

            @test size(g(ones(Float32, 3, 2))) == (2,5,4,2)

            # note that this works only because MeasureNout got the size when we evaluated the graph absorbvertex
            # In other cases one might need try_infer_sizes!
            @test_logs (:warn, r"Could not change nout") Δnout!(v1, 3)

            @test nout.(vertices(g)) == [3, 30, 3, 4]
            @test size(g(ones(Float32, 3, 2))) == (2,5,4,2)
        end

        @testset "Reshape 2D -> 4D all fixed" begin
            v0 = inputvertex("in", 3)
            v1 = dv("v1", v0, 4)
            v2 = rv("v2", v1, (5, 2, 3, 2))
            v3 = cv("v3", v2, 4)

            @test actdim(v2) == [actdim(layer(v2))] == [3]
            @test layer(v2) isa Reshape

            g = CompGraph(v0, v3)

            @test size(g(ones(Float32, 3, 15))) == (5,2,4,2)

            @test @test_logs((:warn, r"Could not change nout"), match_mode=:any, Δnout!(v1, 3)) == false

            @test nout.(vertices(g)) == [3, 4, 3, 4]
            @test size(g(ones(Float32, 3, 15))) == (5,2,4,2)
        end

        @testset "Reshape 2D -> 4D all fixed 0-dim" begin
            v0 = inputvertex("in", 3)
            v1 = dv("v1", v0, 4)
            v2 = rv("v2", v1, (5, 2, 0, 2))
            v3 = cv("v3", v2, 4, 4)

            @test actdim(v2) == [actdim(layer(v2))] == [3]
            @test layer(v2) isa Reshape

            g = CompGraph(v0, v3)

            @test size(g(ones(Float32, 3, 20))) == (5,2,4,2)

            @test_logs (:warn, r"Could not change nout") Δnout!(v1, 3)

            @test nout.(vertices(g)) == [3, 6, 6, 4]
            @test size(g(ones(Float32, 3, 20))) == (5,2,4,2)
        end

        @testset "Reshape 4D -> 2D variable batch" begin
            v0 = inputvertex("in", 3)
            v1 = cv("v1", v0, 4)
            v2 = rv("v2", v1, (5 ,Colon()))
            v3 = dv("v3", v2, 4)

            @test actdim(v2) == [actdim(layer(v2))] == [1]
            @test layer(v2) isa Reshape

            g = CompGraph(v0, v3)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 16)

            @test_logs (:warn, r"Could not change nout") Δnout!(v1, 3)

            @test nout.(vertices(g)) == [3, 8, 10, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 16)
        end

        @testset "Reshape 4D -> 2D variable nout" begin
            v0 = inputvertex("in", 3)
            v1 = cv("v1", v0, 4)
            v2 = rv("v2", v1, (Colon(), 16))
            v3 = dv("v3", v2, 5, 4)

            @test actdim(v2) == [actdim(layer(v2))] == [1]
            @test layer(v2) isa Reshape

            g = CompGraph(v0, v3)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 16)

            @test_logs (:warn, r"Could not change nout") Δnout!(v1, 3)

            @test nout.(vertices(g)) == [3, 16, 20, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 16)
        end

        @testset "Reshape 4D -> 2D variable nout 0-dim" begin
            v0 = inputvertex("in", 3)
            v1 = cv("v1", v0, 4)
            v2 = rv("v2", v1, (Colon(), 0))
            v3 = dv("v3", v2, 40, 4)

            @test actdim(v2) == [actdim(layer(v2))] == [1]
            @test layer(v2) isa Reshape

            g = CompGraph(v0, v3)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 2)

            Δnout!(v1, -2)

            @test nout.(vertices(g)) == [3, 2, 20, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 2)
        end
    end

    @testset "Flatten" begin
        import ONNXNaiveNASflux: Flatten
        fv(name, invertex, dim) = absorbvertex(name, MeasureNout(Flatten(dim)), invertex; traitdecoration=SizePseudoTransparent)

        function tg(outsize, dim)
            v0 = inputvertex("in", 3)
            v1 = cv("v1", v0, 4)
            v2 = fv("v2", v1, dim)
            v3 = dv("v3", v2, outsize, 4)

            @test actdim(v2) == [actdim(layer(v2))] == [1]
            @test layer(v2) isa Flatten

            g = CompGraph(v0, v3)
        end

        @testset "Flatten 4D dim 0" begin
            g = tg(80, 4)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 1)

            Δnout!(vertices(g)[2], -3)

            @test nout.(vertices(g)) == [3, 1, 20, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 1)
        end

        @testset "Flatten 4D dim 1" begin
            g = tg(5, 1)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 16)

            Δnout!(vertices(g)[2], 3)

            @test nout.(vertices(g)) == [3, 7, 5, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 28)
        end

        @testset "Flatten 4D dim 2" begin
            g = tg(10, 2)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 8)

            Δnout!(vertices(g)[2], -1)

            @test nout.(vertices(g)) == [3, 3, 10, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 6)
        end

        @testset "Flatten 4D dim 3" begin
            g = tg(40, 3)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 2)

            Δnout!(vertices(g)[2], -2)

            @test nout.(vertices(g)) == [3, 2, 20, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 2)
        end

        @testset "Flatten 4D dim 4" begin
            g = tg(80, 4)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 1)

            Δnout!(vertices(g)[2], -3)

            @test nout.(vertices(g)) == [3, 1, 20, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 1)
        end

        @testset "Flatten 4D dim -4" begin
            g = tg(1, -4)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 80)

            Δnout!(vertices(g)[2], -2)

            @test nout.(vertices(g)) == [3, 2, 1, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 40)
        end

        @testset "Flatten 4D dim -3" begin
            g = tg(5, -3)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 16)

            Δnout!(vertices(g)[2], 3)

            @test nout.(vertices(g)) == [3, 7, 5, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 28)
        end

        @testset "Flatten 4D dim -2" begin
            g = tg(10, -2)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 8)

            Δnout!(vertices(g)[2], -1)

            @test nout.(vertices(g)) == [3, 3, 10, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 6)
        end

        @testset "Flatten 4D dim -1" begin
            g = tg(40, -1)

            @test size(g(ones(Float32, 5,2,3,2))) == (4, 2)

            Δnout!(vertices(g)[2], -2)

            @test nout.(vertices(g)) == [3, 2, 20, 4]
            @test size(g(ones(Float32, 5,2,3,2))) == (4, 2)
        end
    end
end


text:
############################################################################################
"""
$(TYPEDEF)
Diagnostics fields for NUTS sampler.

# Fields
$(TYPEDFIELDS)
"""
struct DiagnosticsNUTS{T<:AbstractFloat} <: MCMCKernelDiagnostics
    "Log density (negative energy)."
    ℓH::T
    "Depth of the tree."
    depth::Int64
    "Reason for termination. See [`InvalidTree`](@ref) and [`REACHED_MAX_DEPTH`](@ref)."
    termination::InvalidTree
    "Acceptance rate statistic."
    acceptance_rate::T
    "Discretization size"
    ϵ::T
    "Number of leapfrog steps evaluated."
    steps::Int64
    "Directions for tree doubling (useful for debugging)."
    directions::Directions
    function DiagnosticsNUTS(
        ℓH::T,
        depth::Int64,
        termination::InvalidTree,
        acceptance_rate::T,
        ϵ::T,
        steps::Int64,
        directions::Directions,
    ) where {T<:AbstractFloat}
        @argcheck ϵ > 0.0 "Discretization size has to be positive"
        return new{T}(ℓH, depth, termination, acceptance_rate, ϵ, steps, directions)
    end
end

"""
$(SIGNATURES)
Summarize results of Vector of diagnostics.

# Examples
```julia
```

"""
function results(
    diagnostics::AbstractVector{DiagnosticsNUTS{F}}, Ndigits::Integer, quantiles::Vector{T}
) where {F<:AbstractFloat,T<:Real}
    println(
        "NUTS sampler had ",
        sum(is_divergent(diagnostics[iter].termination) for iter in eachindex(diagnostics)),
        " divergent transitions. Average (final) stepsize of ",
        round.(
            mean(diagnostics[iter].ϵ for iter in eachindex(diagnostics)); digits=Ndigits
        ),
        " (",
        round.(diagnostics[end].ϵ; digits=Ndigits),
        ") with ",
        round.(
            mean(diagnostics[iter].steps for iter in eachindex(diagnostics)); digits=Ndigits
        ),
        " (",
        diagnostics[end].steps,
        ") number of steps and depth of ",
        round.(
            mean(diagnostics[iter].depth for iter in eachindex(diagnostics)); digits=Ndigits
        ),
        " (",
        diagnostics[end].depth,
        ").",
    )
    return nothing
end

############################################################################################
"""
$(SIGNATURES)
Show relevant diagnostic results.

# Examples
```julia
```

"""
function generate_showvalues(diagnostics::D) where {D<:DiagnosticsNUTS}
    return function showvalues()
        return (:NUTS, "diagnostics"),
        (:ϵ, diagnostics.ϵ), (:steps, diagnostics.steps),
        (:ℓH, diagnostics.ℓH)
    end
end

############################################################################################
# Export
export DiagnosticsNUTS, generate_showvalues


text:
#=
This file is auto-generated. Do not edit.
=#
"""
    mutable struct StaticDemand <: WaterDemand
        name::String
        available::Bool
        junction::Junction
        base_demand::Float64
        pattern_name::String
        _forecasts::InfrastructureSystems.Forecasts
        internal::InfrastructureSystemsInternal
    end

Temporal demand at a junction that is fixed.

# Arguments
- `name::String`
- `available::Bool`
- `junction::Junction`
- `base_demand::Float64`: 'base' demand in m^3/second
- `pattern_name::String`: name of forecast pattern array
- `_forecasts::InfrastructureSystems.Forecasts`
- `internal::InfrastructureSystemsInternal`: internal reference, do not modify
"""
mutable struct StaticDemand <: WaterDemand
    name::String
    available::Bool
    junction::Junction
    "'base' demand in m^3/second"
    base_demand::Float64
    "name of forecast pattern array"
    pattern_name::String
    _forecasts::InfrastructureSystems.Forecasts
    "internal reference, do not modify"
    internal::InfrastructureSystemsInternal
end

function StaticDemand(name, available, junction, base_demand, pattern_name, _forecasts=InfrastructureSystems.Forecasts(), )
    StaticDemand(name, available, junction, base_demand, pattern_name, _forecasts, InfrastructureSystemsInternal(), )
end

function StaticDemand(; name, available, junction, base_demand, pattern_name, _forecasts=InfrastructureSystems.Forecasts(), )
    StaticDemand(name, available, junction, base_demand, pattern_name, _forecasts, )
end

# Constructor for demo purposes; non-functional.
function StaticDemand(::Nothing)
    StaticDemand(;
        name="init",
        available=true,
        junction=Junction(nothing),
        base_demand=0.0,
        pattern_name="init",
        _forecasts=InfrastructureSystems.Forecasts(),
    )
end

"""Get StaticDemand name."""
get_name(value::StaticDemand) = value.name
"""Get StaticDemand available."""
get_available(value::StaticDemand) = value.available
"""Get StaticDemand junction."""
get_junction(value::StaticDemand) = value.junction
"""Get StaticDemand base_demand."""
get_base_demand(value::StaticDemand) = value.base_demand
"""Get StaticDemand pattern_name."""
get_pattern_name(value::StaticDemand) = value.pattern_name
"""Get StaticDemand _forecasts."""
get__forecasts(value::StaticDemand) = value._forecasts
"""Get StaticDemand internal."""
get_internal(value::StaticDemand) = value.internal


text:
t =Template(;
    user="YayeIrene",
    authors="Irene Ndindabahizi",
    plugins=[
        License(; name="MIT"),
        Codecov(),
        TravisCI(),
        Coveralls(),
        AppVeyor(),
    ],
)

generate("ErrorBudget", t)


text:
using StatsBase
export dist_to_mode_sample, mean_inner_dim_sample, outer_dim_sample
export pred_missing, get_prediction, get_truth, was_correct, get_pred_accuracy

function pred_missing(
    S::InteractionSequence{Int},
    ind::Tuple{Int,Int},
    mcmc_post_out::SisPosteriorMcmcOutput
    )

    posterior = mcmc_post_out.posterior
    d = posterior.dist
    Sₓ = deepcopy(S)
    # Maps a mode to the V (num vertices) different distances to value S with different value in ind
    dists_to_vals = Dict{InteractionSequence{Int},Vector{Real}}()
    V = length(posterior.V)
    # @show V
    dist_vec = zeros(V)
    μ_tmp = zeros(V)
    μ = zeros(V)

    for (mode, γ) in zip(mcmc_post_out.S_sample, mcmc_post_out.γ_sample)
        if mode ∉ keys(dists_to_vals)
            for x in 1:V
                Sₓ[ind[1]][ind[2]] = x
                dist_vec[x] = d(Sₓ, mode)
            end
            dists_to_vals[mode] = dist_vec
        end
        map!(x -> exp(- γ * x ), μ_tmp, dists_to_vals[mode])
        μ_tmp /= sum(μ_tmp)
        μ += μ_tmp
    end

    μ /= sum(μ)

    return SingleMissingPredictive(S, ind, μ)
end

function pred_missing(
    S::InteractionSequence{Int},
    ind::Tuple{Int,Int},
    model::SIS
    )

    d, γ = (model.dist, model.γ)
    μ = zeros(length(model.V))
    Sₓ = deepcopy(S)
    i,j = ind
    for x in model.V
        Sₓ[i][j] = x
        μ[x] = exp(-γ * d(Sₓ, model.mode))
    end
    μ /= sum(μ)
    return SingleMissingPredictive(S, ind, μ)
end

function get_prediction(
    predictive::SingleMissingPredictive
    )
    max_prob = maximum(predictive.p)  # MAP
    vals = findall(predictive.p .== max_prob) # Vertices with max MAP
    pred = rand(vals) # Choose randomly from said vertices
    H = entropy(predictive.p) # Evaluate entropy
    return pred, H
end

function get_truth(
    predictive::SingleMissingPredictive
    )
    i,j = predictive.ind
    return predictive.S[i][j]
end

was_correct(predictive::SingleMissingPredictive) = (get_prediction(predictive)[1] == get_truth(predictive))

function get_pred_accuracy(
    predictives::Vector{SingleMissingPredictive}
    )
    return sum(was_correct.(predictives))/length(predictives)
end

function draw_sample!(
    out::Vector{InteractionSequence{Int}},
    mcmc::T,
    predictive::PosteriorPredictive
    ) where {T<:Union{SisMcmcSampler, SimMcmcSampler}}

    posterior_mcmc = predictive.posterior
    S_sample = posterior_mcmc.S_sample
    γ_sample = posterior_mcmc.γ_sample
    posterior = posterior_mcmc.posterior
    d = posterior.dist
    V = posterior.V
    K_I, K_O = (posterior.K_inner, posterior.K_outer)
    n_samples = length(S_sample)

    for i in eachindex(out)
        ind = rand(1:n_samples)
        model = predictive.model_type(S_sample[ind],γ_sample[ind], d, V, K_I, K_O)
        draw_sample!(view(out,i:i), mcmc, model)
    end
end

function draw_sample!(
    out::Vector{InteractionSequenceSample{Int}}, # Note difference here with prev function
    mcmc::T,
    predictive::PosteriorPredictive
    ) where {T<:Union{SisMcmcSampler, SimMcmcSampler}}

    # Aliases for info in given objects
    posterior_mcmc = predictive.posterior
    S_sample = posterior_mcmc.S_sample
    γ_sample = posterior_mcmc.γ_sample
    posterior = posterior_mcmc.posterior
    d = posterior.dist
    V = posterior.V
    K_I, K_O = (posterior.K_inner, posterior.K_outer)
    n_samples = length(S_sample)

    # For each length n vector in out we
    # (i) sample parameters from posterior
    # (ii) store n samples from model at parameters in said vector
    for i in eachindex(out)
        ind = rand(1:n_samples)
        model = predictive.model_type(S_sample[ind],γ_sample[ind], d, V, K_I, K_O)
        draw_sample!(out[i], mcmc, model)
    end
end

function draw_sample(
    mcmc::T,
    predictive::PosteriorPredictive;
    n_samples::Int=500,  # Number of draws from the posterior
    n_reps::Int=100  # Number of draws from predictive at sampled parameters
    ) where {T<:Union{SisMcmcSampler, SimMcmcSampler}}

    if n_reps == 1
        out = InteractionSequenceSample{Int}(undef, n_samples)
    else
        out = [InteractionSequenceSample{Int}(undef, n_samples) for i in 1:n_reps]
    end

    draw_sample!(out, mcmc, predictive)

    return out
end

# Distance to mode
# ----------------

function dist_to_mode_sample!(
    out::Vector{Float64},
    mcmc::T,
    predictive::PosteriorPredictive
    ) where {T<:Union{SisMcmcSampler, SimMcmcSampler}}

    posterior_mcmc = predictive.posterior
    S_sample = posterior_mcmc.S_sample
    γ_sample = posterior_mcmc.γ_sample
    posterior = posterior_mcmc.posterior
    d = posterior.dist
    V = posterior.V
    K_I, K_O = (posterior.K_inner, posterior.K_outer)
    n_samples = length(S_sample)
    sample_store = [[Int[]]]
    for i in eachindex(out)
        ind = rand(1:n_samples)
        model = predictive.model_type(S_sample[ind],γ_sample[ind], d, V, K_I, K_O)
        draw_sample!(sample_store, mcmc, model)
        out[i] = d(sample_store[1], model.mode)
    end
end

function dist_to_mode_sample!(
    out::Vector{Vector{Float64}},
    mcmc::T,
    predictive::PosteriorPredictive
    ) where {T<:Union{SisMcmcSampler, SimMcmcSampler}}

    posterior_mcmc = predictive.posterior
    S_sample = posterior_mcmc.S_sample
    γ_sample = posterior_mcmc.γ_sample
    posterior = posterior_mcmc.posterior
    d = posterior.dist
    V = posterior.V
    K_I, K_O = (posterior.K_inner, posterior.K_outer)
    n_samples = length(S_sample)
    sample_store = [[Int[]] for i in 1:length(out[1])]
    for i in eachindex(out)
        ind = rand(1:n_samples)
        model = predictive.model_type(S_sample[ind],γ_sample[ind], d, V, K_I, K_O)
        draw_sample!(sample_store, mcmc, model)
        out[i] = map(x -> d(x, model.mode), sample_store)
    end
end

function dist_to_mode_sample(
    mcmc::T,
    predictive::PosteriorPredictive;
    n_samples::Int=500,  # Number of draws from the posterior
    n_reps::Int=100  # Number of draws from predictive at sampled parameters
    ) where {T<:Union{SisMcmcSampler, SimMcmcSampler}}

    if n_reps == 1
        out = zeros(n_samples)
    else
        out = [zeros(n_samples) for i in 1:n_reps]
    end

    dist_to_mode_sample!(out, mcmc, predictive)

    return out
end

function mean_inner_dim_sample!(
    out::Vector{Float64},
    mcmc::T,
    predictive::PosteriorPredictive
    ) where {T<:Union{SisMcmcSampler, SimMcmcSampler}}

    posterior_mcmc = predictive.posterior
    S_sample = posterior_mcmc.S_sample
    γ_sample = posterior_mcmc.γ_sample
    posterior = posterior_mcmc.posterior
    d = posterior.dist
    V = posterior.V
    K_I, K_O = (posterior.K_inner, posterior.K_outer)
    n_samples = length(S_sample)
    sample_store = [[Int[]]]
    for i in eachindex(out)
        ind = rand(1:n_samples)
        model = predictive.model_type(S_sample[ind],γ_sample[ind], d, V, K_I, K_O)
        draw_sample!(sample_store, mcmc, model)
        out[i] = mean(lenth.(sample_store[1]))
    end
end

function outer_dim_sample!(
    out::Vector{Int},
    mcmc::T,
    predictive::PosteriorPredictive
    ) where {T<:Union{SisMcmcSampler, SimMcmcSampler}}

    posterior_mcmc = predictive.posterior
    S_sample = posterior_mcmc.S_sample
    γ_sample = posterior_mcmc.γ_sample
    posterior = posterior_mcmc.posterior
    d = posterior.dist
    V = posterior.V
    K_I, K_O = (posterior.K_inner, posterior.K_outer)
    n_samples = length(S_sample)
    sample_store = [[Int[]]]

    for i in eachindex(out)
        ind = rand(1:n_samples)
        model = predictive.model_type(S_sample[ind],γ_sample[ind], d, V, K_I, K_O)
        draw_sample!(sample_store, mcmc, model)
        out[i] = length(sample_store[1])
    end
end

function mean_inner_dim_sample!(
    out::Vector{Vector{Float64}},
    mcmc::T,
    predictive::PosteriorPredictive
    ) where {T<:Union{SisMcmcSampler, SimMcmcSampler}}

    posterior_mcmc = predictive.posterior
    S_sample = posterior_mcmc.S_sample
    γ_sample = posterior_mcmc.γ_sample
    posterior = posterior_mcmc.posterior
    d = posterior.dist
    V = posterior.V
    K_I, K_O = (posterior.K_inner, posterior.K_outer)
    n_samples = length(S_sample)
    sample_store = [[Int[]] for i in 1:length(out[1])]
    for i in eachindex(out)
        ind = rand(1:n_samples)
        model = predictive.model_type(S_sample[ind],γ_sample[ind], d, V, K_I, K_O)
        draw_sample!(sample_store, mcmc, model)
        out[i] = map(x -> mean(length.(x)), sample_store)
    end
end

function outer_dim_sample!(
    out::Vector{Vector{Int}},
    mcmc::T,
    predictive::PosteriorPredictive
    ) where {T<:Union{SisMcmcSampler, SimMcmcSampler}}

    posterior_mcmc = predictive.posterior
    S_sample = posterior_mcmc.S_sample
    γ_sample = posterior_mcmc.γ_sample
    posterior = posterior_mcmc.posterior
    d = posterior.dist
    V = posterior.V
    K_I, K_O = (posterior.K_inner, posterior.K_outer)
    n_samples = length(S_sample)
    sample_store = [[Int[]] for i in 1:length(out[1])]
    for i in eachindex(out)
        ind = rand(1:n_samples)
        model = predictive.model_type(S_sample[ind],γ_sample[ind], d, V, K_I, K_O)
        draw_sample!(sample_store, mcmc, model)
        out[i] = map(x -> length(x), sample_store)
    end
end

function mean_inner_dim_sample(
    mcmc::T,
    predictive::PosteriorPredictive;
    n_samples::Int=500,  # Number of draws from the posterior
    n_reps::Int=100  # Number of draws from predictive at sampled parameters
    ) where {T<:Union{SisMcmcSampler, SimMcmcSampler}}

    if n_reps == 1
        out = zeros(n_samples)
    else
        out = [zeros(n_samples) for i in 1:n_reps]
    end

    mean_inner_dim_sample!(out, mcmc, predictive)

    return out
end

function outer_dim_sample(
    mcmc::T,
    predictive::PosteriorPredictive;
    n_samples::Int=500,  # Number of draws from the posterior
    n_reps::Int=100  # Number of draws from predictive at sampled parameters
    ) where {T<:Union{SisMcmcSampler, SimMcmcSampler}}

    if n_reps == 1
        out = zeros(Int, n_samples)
    else
        out = [zeros(Int, n_samples) for i in 1:n_reps]
    end

    outer_dim_sample!(out, mcmc, predictive)

    return out
end


text:
using PathWeightSampling: GaussianSystem
using Distributions
using Test
using LinearAlgebra

system = GaussianSystem(delta_t=0.05, duration=1.0)
initial = generate_configuration(system)

n_dim = length(initial) ÷ 2
c_xx = system.joint.Σ[n_dim + 1:end, n_dim + 1:end]
@test issymmetric(c_xx)

marginal = MvNormal(c_xx)
val1 = logpdf(marginal, initial[n_dim + 1:end]) # this is the analytically correct value of the log marginal

# we now test for every implemented estimation algorithm whether they yield results within 1% of the correct value

algorithms = [
    DirectMCEstimate(2^16),
    AnnealingEstimate(10, 50, 2000),
    TIEstimate(2^10, 10, 2^17)
]

for algorithm in algorithms
    local result = PathWeightSampling.simulate(algorithm, initial, system, scale=0.6)
    val2 = PathWeightSampling.log_marginal(result)
    @show algorithm (val2 / val1 - 1)
    if :acceptance in fieldnames(typeof(result))
        @show mean(result.acceptance)
    end
    @test isapprox(val1, val2, rtol=1e-2)
end


text:
# export table

# Base.show
function Base.show(io::IO, vertical::Vertical)
  print(io, vertical.data)
end

function Base.show(io::IO, dash::Dash)
  print(io, repeat(dash.data, dash.repeat))
end

function Base.show(io::IO, connector::Connector)
  print(io, connector.data)
end

function Base.show(io::IO, mill::Mill)
  print(io, decking(mill))
end

function Base.show(io::IO, linear::Linear)
  firstcell = linear[2]
  assert(isa(firstcell, Cell))
  height = firstcell.height
  @inbounds for i=1:height
    for j=1:length(linear)
      cv = linear[j]
      if isa(cv, Vertical)
        print(io, cv)
      else
        cel = cv
        rows,cols = arraysize(cel.data)
        if rows >= i
          for n=1:cols
            print(io, cel.data[i,n])
          end
        else
          print(io, repeat(" ", cel.width))
        end
      end
    end
    println(io)
  end
end

function Base.show(io::IO, horizon::Horizontal)
  for j=1:length(horizon)
    hr = horizon[j]
    print(io, hr)
  end
end

function Base.show(io::IO, plates::PlateVector)
  len = length(plates)
  @inbounds for i=1:len
    plate = plates[i]
    print(io, plate)
    if isa(plate, Horizontal) && len > i
      println(io)
    end
  end
end

# Base.isless
Base.isless(::String, ::Symbol) = false

# arraysize
arraysize(a::AbstractArray) = size(a[:,:])

# precell
precell(::Void) = PreCell([], 0, 0)

function precell(el::String)
  if contains(el, "\n")
    a = map(split(el, "\n")) do x
       width = strwidth(x)
       string(x, repeat(" ", width-strwidth(x)))
    end
    m,n = arraysize(a)
    PreCell(a, maximum(map(length, a)), m)
  else
    width = strwidth(el)
    PreCell([el], width, 1)
  end
end

function precell(el::AbstractArray)
  rows,cols = arraysize(el)
  if 0==rows
    PreCell([""], 0, 0)
  else
    prelines = similar(el, String)
    widths = zeros(Int, rows, cols)
    for i=1:rows
      for j=1:cols
        s = replace(string(el[i,j]), "\n", " ")
        widths[i,j] = length(s)
        prelines[i,j] = s
      end
    end
    maxwidths = zeros(Int, cols)
    for j=1:cols
      maxwidths[j] = maximum(widths[:,j])
    end

    lines = AbstractArray{String}([])
    for i=1:rows
      line = AbstractArray{String}([])
      for (j,x) in enumerate(prelines[i,:])
        push!(line, lpad(x, maxwidths[j]))
      end
      push!(lines, join(line, " "))
    end
    PreCell(lines, sum(maxwidths) + (cols-1), rows)
  end
end

function precell(el::Any)
  PreCell([el], length(repr(el)), 1)
end

# postcell
function postcell(precel::PreCell, width::Int, height::Int, margin::Margin)
  data = precel.data
  rows,cols = arraysize(data)
  A = Array{String}(rows, cols)
  @inbounds for i=1:rows
    for j=1:cols
      el = data[i,j]
      x = isa(el, String) ? el : repr(el)
      prep = repeat(" ", margin.leftside + width - strwidth(x))
      A[i,j] = string(prep, x, repeat(" ", margin.rightside))
    end
  end
  Cell(A, width+margin.leftside+margin.rightside, height)
end

function horizontal(maxwidths::Vector{Int}, cols::Int, margin::Margin; dash="-")
  horizon = Horizontal{Union{Dash,Connector}}([])
  push!(horizon, Connector())
  for j=1:cols
    width = maxwidths[j]
    push!(horizon, Dash(dash, width + margin.leftside + margin.rightside))
    push!(horizon, Connector())
  end
  horizon
end

function vertical(maxheight::Int)
  Vertical(maxheight)
end

function marginal(option::Dict)
  Margin(1, 1) # fixed
end

function decking(mill::Mill)
  board = mill.board
  option = mill.option

  # prepare precells
  preplates = PlateVector{Union{Linear}}([])
  rows,cols = arraysize(board)
  widths = zeros(Int, rows+1, cols+1)
  heights = zeros(Int, rows+1, cols+1)

  headlinear = Linear{Union{AbstractCell}}([])
  zerocolname = nothing
  if haskey(option, :zerocolname)
    zerocolname = option[:zerocolname]
  end
  cel = precell(zerocolname)
  widths[1,1] = cel.width
  heights[1,1] = cel.height
  push!(headlinear, cel)
  precols = Vector{PreCell}(cols)
  for j=1:cols
    precols[j] = precell(j)
  end
  if haskey(option, :colnames)
    @inbounds for (j,name) in enumerate(option[:colnames])
      if cols >= j
        precols[j] = precell(name)
      end
    end
  end
  for j=1:cols
    cel = precols[j]
    widths[1,j+1] = cel.width
    heights[1,j+1] = cel.height
    push!(headlinear, cel)
  end
  push!(preplates, headlinear)

  prerows = Vector{PreCell}(rows)
  for i=1:rows
    prerows[i] = precell(i)
  end
  if haskey(option, :rownames)
    @inbounds for (i,name) in enumerate(option[:rownames])
      if rows >= i
        prerows[i] = precell(name)
      end
    end
  end
  for i=1:rows
    linear = Linear{Union{AbstractCell}}([])
    rownamecell = prerows[i]
    widths[1+i,1] = rownamecell.width
    heights[1+i,1] = rownamecell.height
    push!(linear, rownamecell)
    for j=1:cols
      cel = precell(board[i,j])
      widths[1+i,1+j] = cel.width
      heights[1+i,1+j] = cel.height
      push!(linear, cel)
    end
    push!(preplates, linear)
  end

  maxwidths = zeros(Int, cols+1)
  for j=1:cols+1
    maxwidths[j] = maximum(widths[:,j])
  end
  maxheights = zeros(Int, rows+1)
  for i=1:rows+1
    maxheights[i] = maximum(heights[i,:])
  end
  margin = marginal(option)
  if 0==rows
    cols = 1
  else
    rows += 1
    cols += 1
  end

  # decking
  plates = PlateVector{Union{Linear,Horizontal}}([])
  push!(plates, horizontal(maxwidths, cols, margin, dash="="))
  @inbounds for i=1:rows
    preplate = preplates[i]
    linear = Linear{Union{AbstractCell}}([])
    push!(linear, vertical(maxheights[i]))
    maxheight = maxheights[i]
    for j=1:cols
      precel = preplate[j]
      cel = postcell(precel, maxwidths[j], maxheight, margin)
      if 1==j && cel.height > 1
        prep = [repeat(" ", cel.width) for x in 1:(cel.height-length(cel.data))]
        cel.data = vcat(prep, cel.data)
      end
      push!(linear, cel)
      push!(linear, vertical(maxheight))
    end
    push!(plates, linear)
    push!(plates, horizontal(maxwidths, cols, margin, dash=1==i ? "=" : "-"))
  end
  plates
end

# table
table{T}(board::T) = table(board, Dict())
table{T}(board::T, option::Pair...) = table(board, Dict(option))
table(board::Any, option::Dict) = table([board], option)
table(board::AbstractArray, option::Dict) = Mill(board, option)
function table(board::Tuple, option::Dict)
  if 0==length(board)
    Mill([], option)
  else
    Mill(rotl90(collect(board)[:,:]), option)
  end
end
function table(board::Dict, option::Dict)
  if 0==length(board)
    Mill([], option)
  else
    opt = option
    if !haskey(option, :colnames)
      opt[:zerocolname] = "key"
      opt[:colnames] = ["value"]
    end
    opt[:rownames] = sort(collect(keys(board)))
    valuez = map(opt[:rownames]) do key
      board[key]
    end
    table(valuez, opt)
  end
end
table{T<:Base.Enums.Enum}(board::T, option::Dict) = table(typeof(board), option)
function table{T<:Base.Enums.Enum}(board::Type{T}, option::Dict)
  opt = option
  enums = instances(board)
  opt[:rownames] = map(string, enums)
  if !haskey(option, :colnames)
    opt[:zerocolname] = string(board)
    opt[:colnames] = ["value"]
  end
  valuez = collect(map(enums) do typename
    string(Int(typename))
  end)
  table(valuez, opt)
end


text:
using CSyntax.CStatic
using Test

function cstatic_basic()
    @cstatic x=0 begin
        x += 1
    end
end

function cstatic_scope()
    w = false
    x1, y1 = @cstatic x=30 y=40 begin
        x += 1
        y -= 1
        w = true
    end
    x = 10
    x2 = @cstatic x=3.1415 begin
       x *= 2
    end
    return x, w, x1, y1, x2
end

@testset "CStatic" begin
    for i = 1:10
        @test cstatic_basic() == i
    end
    cstatic_scope()
    cstatic_scope()
    @test cstatic_scope() == (10, true, 33, 37, 25.132)
end


text:
using PyPlot
using CGS
nperiod = 100
p = logspace(log10(2),log10(5000/365.25),nperiod)
lnR = log(4*REARTH/RJUPITER)
qdet = zeros(nperiod)
for i=1:nperiod
  qdet[i] = minimum([maximum([(-0.13*log(p[i])+0.95)*(-0.20*lnR+0.90),0]),1])/(1+exp(-(0.7*log(p[i])+3.06)*(lnR-(-0.07*log(p[i])-0.91))))
end
fduty = 0.97
qwin = 1. - (1.-fduty).^(4./p)
qgeom = 0.005./p.^(2/3)
#loglog(p,qdet)
#loglog(p,qdet.*qwin)
semilogx(p*365.25,qdet.*qwin.*qgeom.*1e4)
plot([200,200,340,340,590,590],[0,15,15,12,12,0])
plot([590,590,1000,1000,1700,1700,3000,3000,5000,5000],[0,1,1,6,6,5,5,3,3,0])

max_stars_repo_name: dfm/single-transits

max_issues_repo_name: dfm/single-transits

max_forks_repo_name: dfm/single-transits


text:
Pkg.update()

pkgs = ["Images", "GLFW"]
for pkg in pkgs
	Pkg.add(pkg)
end

dev_pkgs = [
	"https://github.com/SimonDanisch/GLAbstraction.jl.git"
	"https://github.com/SimonDanisch/ModernGL.jl.git"
	"https://github.com/SimonDanisch/GLWindow.jl.git"
	]

for pkg in dev_pkgs
	Pkg.clone(pkg)
end

Pkg.build()
Pkg.update()


text:
using SeisAcoustic

# the path to segy file
work_dir = "/home/wgao1/australia_OBN";

# work_dir = "/Users/wenlei/Desktop/OBC_2D";
path_sgy = joinpath(work_dir, "hydro.segy");

# read text file header
text_header = read_text_header(path_sgy);
print_text_header(text_header);

# read file header
file_header = read_file_header(path_sgy);

# number of traces
tmp = filesize(path_sgy);
num_traces = convert(Int64, (tmp - 3600) / (240+sizeof(Float32)*file_header.ns))

# extract all the traces' header from segy file
traces_header = extract_traces_header(path_sgy; print_interval=100000);

# print the value of the fields of trace header
fields = [:sx, :sy, :gx, :gy]
coor = extract_field_trace_header(traces_header, fields; istart=1, iend=length(traces_header));

fields = [:ep]
ep = extract_field_trace_header(traces_header, fields; istart=1, iend=length(traces_header));

# extract the coordinate
function sort_CSG(coor::Matrix{Real})

    while

end

for i = 1 : 500
    sx = traces_header[i].sx
    sy = traces_header[i].sy
    gx = traces_header[i].gx
    gy = traces_header[i].gy
    println("sx: $sx, sy: $sy, gx: $gx, gy:$gy")
    # x = traces_header[i].cdp
    # println("$x")
end

(sx, sy, gx, gy) = create_acquisition_geometry(traces_header);

path_tab = joinpath(work_dir, "lookup_table.txt");
create_lookup_table(path_tab, path_sgy);

work_dir = "/Users/wenlei/Desktop/OBC_2D";
path_sgy = joinpath(work_dir, "hydro.segy");
path_tab = joinpath(work_dir, "lookup_table.txt");
(fh_size, ns, ntrace, swap_bytes, data_format, shot_idx) = read_lookup_table(path_tab);

i = 16
(trace_header, d) = read_one_shot(path_sgy, path_tab, i);
figure(); imshow(d, vmax=10, vmin=-10, cmap="seismic", aspect=0.2);

i = 1000
(shot_x, shot_y) = extract_shot_gather(sx[i], sy[i], traces_header);
figure(figsize=(8,8)); scatter(gx, gy, s=5, c="b"); scatter(sx, sy, s=5, c="r");
scatter(shot_x, shot_y, s=10, c="g"); scatter(sx[i], sy[i], s=80, c="r");
xlim(351689, 366488); ylim(7780276, 7794926);

function extract_shot_gather(sx::T, sy::T, traces_header::Vector{TraceHeader}; print_interval=100000) where {T<:Int32}

    num = length(traces_header)
    rx  = Vector{Int32}(undef, 0)
    ry  = Vector{Int32}(undef, 0)

    for i = 1 : num
        if sx == traces_header[i].sx && sy == traces_header[i].sy
           push!(rx, traces_header[i].gx)
           push!(ry, traces_header[i].gy)
        end

        if (i % print_interval) == 0
           println("$i")
        end
    end
    return rx, ry
end

# # extract unique source and receiver's location
# (sx, sy, gx, gy) = extract_sensors_location(traces_header; print_interval=100000);

# temporarily save the unique position of sources and receivers
# path_sco = joinpath(work_dir, "unique_source_location.bin");
# fp = open(path_sco, "w")
# write(fp, convert(Int32, length(sx)))
# write(fp, sx, sy); close(fp);
path_sco = "/Users/wenlei/Desktop/unique_source_location.bin"
fp = open(path_sco, "r")
num_s = read(fp, Int32)
sx = Vector{Int32}(undef, num_s); read!(fp, sx);
sy = Vector{Int32}(undef, num_s); read!(fp, sy);

path_rco = "/Users/wenlei/Desktop/unique_receiver_location.bin"
fp = open(path_rco, "r")
num_r = read(fp, Int32)
gx = Vector{Int32}(undef, num_r); read!(fp, gx);
gy = Vector{Int32}(undef, num_r); read!(fp, gy);

# create the acquisition geometry for each shot
(src_x, src_y, rec_x, rec_y) = create_acquisition_geometry(traces_header);

figure(figsize=(8,8)); scatter(gx, gy, s=5, c="b"); scatter(sx, sy, s=5, c="r");
i = 20000; scatter(rec_x[i], rec_y[i], s=10, c="g"); scatter(src_x[i], src_y[i], s=80, c="r");
xlim(351689, 366488); ylim(7780276, 7794926);

(usx1, usy1) = remove_repeat(sx, sy);
(usx2, usy2) = remove_repeat(src_x, src_y);

figure(figsize=(8,8)); scatter(usx1[1:100], usy1[1:100], s=5, c="b");
xlim(351689, 366488); ylim(7780276, 7794926);

function min_distance(sx::Vector{T}, sy::Vector{T}) where {T<:Real}

    r = Inf32
    num = length(sx)

    for i = 1 : length(sx)

    end

end

function remove_repeat(sx::Vector{T}, sy::Vector{T}; tol=0.1) where {T<:Real}

    num = length(sx)
    usx = [sx[1]]
    usy = [sy[1]]

    res = zeros(eltype(sx), 2)
    dis = Inf32

    for i = 2 : num

        # source location
        dis = Inf32
        for j = 1 : length(usx)
            res[1] = usx[j] - sx[i]
            res[2] = usy[j] - sx[i]
            tmp = norm(res)
            if tmp < dis
               dis = tmp
            end
        end
        if dis > tol
           push!(usx, sx[i])
           push!(usy, sy[i])
        end

        println("$i")
    end

    return usx, usy
end

path_rco = joinpath(work_dir, "unique_receiver_location.bin");
fp = open(path_rco, "w")
write(fp, convert(Int32, length(gx)))
write(fp, gx, gy); close(fp);

# create the unique source location
sh = RegularSampleHeader(n1=length(sx), n2=2, label1="coordinate", label2="dimension",
     title="unique source coordinate", data_format=Int32);
src_loc = hcat(sx, sy)

path_sloc = joinpath(work_dir, "unique_source_location.rsd");
write_RSdata(path_sloc, sh, hcat(sx, sy))

#create the unique receiver location
gh = RegularSampleHeader(n1=length(gx), n2=2, label1="coordinate", label2="dimension",
     title="unique receiver coordinate", data_format=Float32);
rec_loc = hcat(gx, gy)

# i = 50;
# shot = data[:, trace_start[i]:trace_end[i]]
# SeisPlot(shot, pclip=90);
#
# path = "/Users/wenlei/Desktop/for_Wenlei/data_Cynthia.sgy"
# (file_header, trace_header, data) = read_segy_file(path);
# (f, a) = amplitude_spectra(d[:,300], 0.004, 2048; fhigh_end=100);
#
# path = "/Users/wenlei/Desktop/for_Wenlei/source_wavelet.sgy"
# (file_header1, trace_header1, wavelet) = read_segy_file(path);
#
# path = "/Users/wenlei/Desktop/for_Wenlei/velocity.sgy"
# (file_header2, trace_header2, velocity) = read_segy_file(path);
#
# SeisPlot(wavelet, pclip=99);
# SeisPlot(velocity, vmin=minimum(velocity), vmax=maximum(velocity), cmap="rainbow", wbox=10, hbox=4); colorbar();
#
# (trace_start, trace_end, source_x, source_y, receiver_x, receiver_y) = create_acquisition_geometry(trace_header)