using ForwardDiff, UnPack, LinearAlgebra
using DiffEqBase: DiffCache, get_tmp, dualcache

mutable struct RALΛ{L <: Function, LC}
    Λ::L
    cache::LC
end

function RALΛ(Λ::Function, z::C1, matrix_type::DataType, dims::Tuple{Int, Int}) where {C1 <: AbstractVector{<: Number}}
    cache = matrix_type(undef, 0, 0) # Create empty matrix first, just to check if Λ is in place or not
    if applicable(Λ, cache, z)
        cache = matrix_type(undef, dims)
        Λnew = function _Λ_ip(cache::LCN, z::C1N) where {LCN <: DiffCache, C1N <: AbstractVector{<: Number}}
            Λ(get_tmp(cache, z), z)
            return get_tmp(cache, z)
        end
        return RALΛ(Λnew, dualcache(cache, Val{length(z)}))
    else
        function _Λ_oop(cache::LCN, z::C1N) where {LCN <: Nothing, C1N <: AbstractVector{<: Number}}
            return Λ(z)
        end
        return RALΛ(Λnew, nothing)
    end
end

function RALΛ(Λin::LC, z::C1) where {LC <: AbstractMatrix{<: Number}, C1 <: AbstractVector{<: Number}}
    Λ(cache::LCN, z::C1N) where {LCN <: AbstractMatrix{<: Number}, C1N <: AbstractVector{<: Number}} = cache
    return RALΛ{Function, LC}(Λ, Λin)
end

function (ralλ::RALΛ)(z::C1) where {C1 <: AbstractVector{<: Number}}
    return ralλ.Λ(ralλ.cache, z)
end

mutable struct RALΣ{S <: Function, SC}
    Σ::S
    cache::SC
end

function RALΣ(Σ::Function, z::C1, matrix_type::DataType, dims::Tuple{Int, Int}) where {C1 <: AbstractVector{<: Number}}
    cache = matrix_type(undef, 0, 0)
    if applicable(Σ, cache, z)
        cache = matrix_type(undef, dims)
        Σnew = function _Σ_ip(cache::SCN, z::C1N) where {SCN <: DiffCache, C1N <: AbstractVector{<: Number}}
            du = get_tmp(cache, z)
            Σ(du, z)
            return du
        end
        return RALΣ(Σnew, dualcache(cache, Val{length(z)}))
    else
        Σnew = function _Σ_oop(cache::SCN, z::C1N) where {SCN <: Nothing, C1N <: AbstractVector{<: Number}}
            return Σ(z)
        end
        return RALΣ(Σnew, nothing)
    end
end

function RALΣ(Σin::SC, z::C1) where {SC <: AbstractMatrix{<: Number}, C1 <: AbstractVector{<: Number}}
    Σ(cache::SCN, z::C1N) where {SCN <: AbstractMatrix{<: Number}, C1N <: AbstractVector{<: Number}} = cache
    return RALΣ{Function, SC}(Σ, Σin)
end

function (ralσ::RALΣ)(z::C1) where {C1 <: AbstractVector{<: Number}}
    return ralσ.Σ(ralσ.cache, z)
end

mutable struct RALNonlinearSystem{M <: Function, L <: RALΛ, S <: RALΣ, X <: Function, V <: Function,
                                  VC1 <: AbstractVector{<: Number}, VC2 <: AbstractVector{<: Number}, VC3 <: AbstractVector{<: Number}}
    μ::M         # Functions
    Λ::L         # no type assertion for L b/c it can be Function or Matrix of zeros
    Σ::S         # no type assertion for S b/c it can be Function or constant Matrix
    ξ::X
    𝒱::V
    μ_sss::VC1    # Stochastic steady state values, for caching
    ξ_sss::VC2
    𝒱_sss::VC3
    inplace::NamedTuple{(:μ, :ξ, :𝒱), NTuple{3, Bool}}
end

function RALNonlinearSystem(μ::M, Λ::L, Σ::S, ξ::X, 𝒱::V, μ_sss::VC1, ξ_sss::VC2, 𝒱_sss::VC3,
                            z::C1, y::C1, Ψ::C2, Γ₅::JC5, Γ₆::JC6) where {M <: Function, L <: RALΛ, S <: RALΣ, X <: Function, V <: Function,
                                                                          VC1 <: AbstractVector{<: Number}, VC2 <: AbstractVector{<: Number},
                                                                          VC3 <: AbstractVector{<: Number},
                                                                          C1 <: AbstractVector{<: Number}, C2 <: AbstractMatrix{<: Number},
                                                                          JC5 <: AbstractMatrix{<: Number}, JC6 <: AbstractMatrix{<: Number}}

    inplace = (μ = applicable(μ, μ_sss, z, y), ξ = applicable(ξ, ξ_sss, z, y), 𝒱 = applicable(𝒱, 𝒱_sss, z, Ψ, Γ₅, Γ₆))

    return RALNonlinearSystem{M, L, S, X, V, VC1, VC2, VC3}(μ, Λ, Σ, ξ, 𝒱, μ_sss, ξ_sss, 𝒱_sss, inplace)
end

function update!(m::RALNonlinearSystem, z::C1, y::C1, Ψ::C2,
                 Γ₅::JC5, Γ₆::JC6) where {C1 <: AbstractVector{<: Number}, C2 <: AbstractMatrix{<: Number},
                                          JC5 <: AbstractMatrix{<: Number}, JC6 <: AbstractMatrix{<: Number}}

    if m.inplace[:μ]
        m.μ(m.μ_sss, z, y)
    else
        m.μ_sss .= m.μ(z, y)
    end

    if m.inplace[:ξ]
        m.ξ(m.ξ_sss, z, y)
    else
        m.ξ_sss .= m.ξ(z, y)
    end

    if m.inplace[:𝒱]
        m.𝒱(m.𝒱_sss, z, Ψ, Γ₅, Γ₆)
    else
        m.𝒱_sss .= m.𝒱(z, Ψ, Γ₅, Γ₆)
    end

    m
end

mutable struct RALLinearizedSystem{Mz <: Function, My <: Function, Xz <: Function, Xy <: Function, J <: Function,
                                   JC1 <: AbstractMatrix{<: Number}, JC2 <: AbstractMatrix{<: Number},
                                   JC3 <: AbstractMatrix{<: Number}, JC4 <: AbstractMatrix{<: Number},
                                   JC5 <: AbstractMatrix{<: Number}, JC6 <: AbstractMatrix{<: Number},
                                   JC7 <: AbstractMatrix{<: Number}}
    μz::Mz     # Functions
    μy::My
    ξz::Xz
    ξy::Xy
    J𝒱::J
    Γ₁::JC1    # Jacobians, for caching
    Γ₂::JC2
    Γ₃::JC3
    Γ₄::JC4
    Γ₅::JC5
    Γ₆::JC6
    JV::JC7
    inplace::NamedTuple{(:μz, :μy, :ξz, :ξy, :J𝒱), NTuple{5, Bool}}
end

function RALLinearizedSystem(μz::Mz, μy::My, ξz::Xz, ξy::Xy, J𝒱::J,
                             Γ₁::JC1, Γ₂::JC2, Γ₃::JC3, Γ₄::JC4, Γ₅::JC5, Γ₆::JC6,
                             JV::JC7, z::C1, y::C1, Ψ::C2,
                             μ_sss::VC1, ξ_sss::VC2, 𝒱_sss::VC3) where {Mz <: Function, My <: Function, Xz <: Function,
                                                                        Xy <: Function, J <: Function,
                                                                        JC1 <: AbstractMatrix{<: Number}, JC2 <: AbstractMatrix{<: Number},
                                                                        JC3 <: AbstractMatrix{<: Number}, JC4 <: AbstractMatrix{<: Number},
                                                                        JC5 <: AbstractMatrix{<: Number}, JC6 <: AbstractMatrix{<: Number},
                                                                        JC7 <: AbstractMatrix{<: Number},
                                                                        C1 <: AbstractVector{<: Number}, C2 <: AbstractMatrix{<: Number},
                                                                        VC1 <: AbstractVector{<: Number}, VC2 <: AbstractVector{<: Number},
                                                                        VC3 <: AbstractVector{<: Number},}

    inplace = (μz = applicable(μz, Γ₁, z, y, μ_sss), μy = applicable(μy, Γ₂, z, y, μ_sss), ξz = applicable(ξz, Γ₃, z, y, ξ_sss),
               ξy = applicable(ξy, Γ₄, z, y, ξ_sss), J𝒱 = applicable(J𝒱, JV, z, Ψ, Γ₅, Γ₆, 𝒱_sss))

    return RALLinearizedSystem(μz, μy, ξz, ξy, J𝒱, Γ₁, Γ₂, Γ₃, Γ₄, Γ₅, Γ₆, JV, inplace)
end

function update!(m::RALLinearizedSystem, z::C1, y::C1, Ψ::C2,
                 μ_sss::VC1, ξ_sss::VC2, 𝒱_sss::VC3) where {C1 <: AbstractVector{<: Number}, C2 <: AbstractMatrix{<: Number},
                                                            VC1 <: AbstractVector{<: Number}, VC2 <: AbstractVector{<: Number},
                                                            VC3 <: AbstractVector{<: Number}}

    if m.inplace[:μz]
        m.μz(m.Γ₁, z, y, μ_sss)
    else
        m.μz(m.Γ₁, z, y)
    end

    if m.inplace[:μy]
        m.μy(m.Γ₂, z, y, μ_sss)
    else
        m.μy(m.Γ₂, z, y)
    end

    if m.inplace[:ξz]
        m.ξz(m.Γ₃, z, y, ξ_sss)
    else
        m.ξz(m.Γ₃, z, y)
    end

    if m.inplace[:ξy]
        m.ξy(m.Γ₄, z, y, ξ_sss)
    else
        m.ξy(m.Γ₄, z, y)
    end

    if m.inplace[:J𝒱]
        m.J𝒱(m.JV, z, Ψ, m.Γ₅, m.Γ₆, 𝒱_sss)
    else
        m.J𝒱(m.JV, z, Ψ, m.Γ₅, m.Γ₆)
    end

    m
end

abstract type AbstractRiskAdjustedLinearization end

"""
    RiskAdjustedLinearization(μ, Λ, Σ, ξ, Γ₅, Γ₆, 𝒱, Nz, Ny, Nε)

Creates a first-order perturbation around the stochastic steady state ``(z, y)`` of a discrete-time dynamic model.

(TODO: Move more of the formality to documentation, and make this shorter and concise, w/out explanation of matrix equations)
The affine approximation of the model is
``math
\\begin{aligned}
    \\mathbb{E}[z_{t + 1}] & = \\mu(z, y) + \\Gamma_1(z_t - z) + \\Gamma_2(y_t - y)\\\\
    0                      & = \\xi(z, y) + \\Gamma_3(z_t - z) + \\Gamma_4(y_t - y) + \\Gamma_5 \\mathbb{E}_t z_{t + 1} + \\Gamma_6 \\mathbb{E}_t y_{t + 1} + \\mathscr{V}(z) + J\\mathscr{V}(z)(z_t  - z),
\\end{aligned}
``

where ``\\Gamma_1, \\Gamma_2`` are the Jacobians of ``\\mu`` with respect to ``z_t`` and ``y_t``, respectively;
``\\Gamma_3, \\Gamma_4`` are the Jacobians of ``\\xi`` with respect to ``z_t`` and ``y_t``, respectively;
``\\Gamma_5, \\Gamma_6`` are constant matrices; ``\\mathscr{V}(z)`` is the model's entropy;
``J\\mathscr{V}(z)`` is the Jacobian of the entropy;

and the state variables ``z_t`` and jump variables ``y_t`` follow
``math
\\begin{aligned}
    z_{t + 1} & = z + \\Gamma_1(z_t - z) + \\Gamma_2(y_t - y) + (I_{n_z} - \\Lambda(z_t) \\Psi)^{-1}\\Sigma(z_t)\\varepsilon_{t + 1},\\\\
    y_t       & = y + \\Psi(z_t - z)
\\end{aligned}
``

The unknowns ``(z, y, \\Psi)`` solve the system of equations
``math
\\begin{aligned}
0 & = \\mu(z, y) - z,\\\\
0 & = \\xi(z, y) + \\Gamma_5 z + \\Gamma_6 y + \\mathscr{V}(z),\\\\
0 & = \\Gamma_3 + \\Gamma_4 \\Psi + (\\Gamma_5 + \\Gamma_6 \\Psi)(\\Gamma_1 + \\Gamma_2 \\Psi) + J\\mathscr{V}(z).
\\end{aligned}
``
(TODO: Move the nonlinear model statement to documentation)
The true nonlinear equations defining model are assumed to take the form

``math
\\begin{aligned}
    z_{t + 1} & = \\mu(z_t, y_t) + \\Lambda(z_t)(y_{t + 1} - \\mathbb{E}_t y_{t + 1}) + \\Sigma(z_t) \\varepsilon_{t + 1},\\\\
    0 & = \\log\\mathbb{E}_t[\\exp(\\xi(z_t, y_t) + \\Gamma_5 z_{t + 1} + \\Gamma_6 y_{t + 1})].
\\end{aligned}
``

The vectors ``z_t\\in \\mathbb{R}^{n_z}`` and ``y_t \\in \\mathbb{R}^{n_y}`` are the state and jump variables, respectively.
The first vector equation comprise the model's expectational equations, which are typically
the first-order conditions for the jump variables from agents' optimization problem.
The second vector equation comprise the transition equations of the state variables. The exogenous shocks
``\\varepsilon\\in\\mathbb{R}^{n_\\varepsilon}`` form a martingale difference sequence whose distribution
is described by the differentiable, conditional cumulant generating function (ccgf)

``math
\\begin{aligned}
\\kappa[\\alpha(z_t) \\mid z_t] = \\log\\mathbb{E}_t[\\exp(\\alpha(z_t)' \\varepsilon_{t + 1})],\\quad \text{ for any differentiable map }\\alpha::\\mathbb{R}^{n_z}\\rightarrow\\mathbb{R}^{n_\\varepsilon}.
\\end{aligned}
``

The functions
``math
\\begin{aligned}
\\xi:\\mathbb{R}^{2n_y + 2n_z}\\rightarrow \\mathbb{R}^{n_y},& \\quad \\mu:\\mathbb{R}^{n_y + n_z}\\rightarrow \\mathbb{R}^{n_z},\\\\
\\Lambda::\\mathbb{R}^{n_z} \\rightarrow \\mathbb{R}^{n_z \\times n_y}, & \\quad \\Sigma::\\mathbb{R}^{n_z}\\rightarrow \\mathbb{R}^{n_z\\times n_\\varepsilon}
\\end{aligned}
are differentiable. The first two functions characterize the effects of time ``t`` variables on the expectational and
state transition equations. The function ``\\Lambda`` characterizes heteroskedastic endogenous risk that depends on
innovations in jump variables while the function ``\\Sigma`` characterizes exogenous risk.

Refer to Lopz et al. (2018) "Risk-Adjusted Linearizations of Dynamic Equilibrium Models" for details.
"""
mutable struct RiskAdjustedLinearization{A <: RALNonlinearSystem, B <: RALLinearizedSystem,
                                         C1 <: AbstractVector{<: Number}, C2 <: AbstractMatrix{<: Number}} <: AbstractRiskAdjustedLinearization
    nonlinear::A
    linearization::B
    z::C1        # Coefficients
    y::C1
    Ψ::C2
    Nz::Int      # Dimensions
    Ny::Int
    Nε::Int
end
# TODO
# 1.UPDATE THE PRINTING, maybe just write out "risk-adjusted linearization with dimensions ()"
#
# 2. Test update! functions for the various blocks as well as access functions for RiskAdjustedLinearization
#
# 3. Check inplace inference is correct, check construction of each block plus main block
#=
TODO: Finish this once the final struct is completed
# A series of lower level constructors
function RiskAdjustedLinearization(μ::M, Λ::L, Σ::S, ξ::X, 𝒱::V, μz::Mz, μy::My, ξz::Xz, ξy::Xy, J𝒱::J,
                                   μ_sss::AbstractVector{T}, ξ_sss::AbstractVector{T}, 𝒱_sss::AbstractVector{T},
                                   Γ₁::AbstractMatrix{T}, Γ₂::AbstractMatrix{T}, Γ₃::AbstractMatrix{T}
                                   Γ₄::AbstractMatrix{T}, Γ₅::AbstractMatrix{T}, Γ₆::AbstractMatrix{T},
                                   JV::AbstractMatrix{T}, z::AbstractVector{T}, y::AbstractVector{T}, Ψ::AbstractMatrix{T},
                                   Nε::Int = -1) where {T <: Number, M <: Function, L,
                                                        S, X <: Function, V <: Function,
                                                        Mz <: Function, My <: Function, Xz <: Function,
                                                        Xy <: Function, J <: Function}

    Nz = length(z)
    Ny = length(y)
    if Nε < 0
        Nε = size(Σ(z), 2)
    end

    return RiskAdjustedLinearization{T, M, L, S, X, V, J}(μ, Λ, Σ, ξ, 𝒱, J𝒱, μ_sss, ξ_sss, 𝒱_sss,
                                                          Γ₁, Γ₂, Γ₃, Γ₄, Γ₅, Γ₆,
                                                          JV, z, y, Ψ, Nz, Ny, Nε)
end

function RiskAdjustedLinearization(μ::M, Λ::L, Σ::S, ξ::X, 𝒱::V, μz::Mz, μy::My, ξz::Xz, ξy::Xy, J𝒱::J,
                                   Γ₁::AbstractMatrix{T}, Γ₂::AbstractMatrix{T}, Γ₃::AbstractMatrix{T}
                                   Γ₄::AbstractMatrix{T}, Γ₅::AbstractMatrix{T}, Γ₆::AbstractMatrix{T},
                                   JV::AbstractMatrix{T}, z::AbstractVector{T}, y::AbstractVector{T}, Ψ::AbstractMatrix{T},
                                   Nε::Int = -1) where {T <: Number, M <: Function, L,
                                                        S, X <: Function, V <: Function,
                                                        Mz <: Function, My <: Function, Xz <: Function,
                                                        Xy <: Function, J <: Function}
    Nz = length(z)
    Ny = length(y)
    if Nε < 0
        Nε = size(Σ(z), 2)
    end

    # Cache stochastic steady state vectors
    μ_sss, ξ_sss, 𝒱_sss = _cache_sss_vectors(z, y)

    return RiskAdjustedLinearization{T, M, L, S, X, V, J}(μ, Λ, Σ, ξ, 𝒱, J𝒱, μ_sss, ξ_sss, 𝒱_sss,
                                                          Γ₁, Γ₂, Γ₃, Γ₄, Γ₅, Γ₆,
                                                          JV, z, y, Ψ, Nz, Ny, Nε)
end

function RiskAdjustedLinearization(μ::M, Λ::L, Σ::S, ξ::X, 𝒱::V, μz::Mz, μy::My, ξz::Xz, ξy::Xy, J𝒱::J,
                                   z::AbstractVector{T}, y::AbstractVector{T}, Ψ::AbstractMatrix{T},
                                   Nε::Int = -1) where {T <: Number, M <: Function, L,
                                                        S, X <: Function, V <: Function,
                                                        Mz <: Function, My <: Function, Xz <: Function,
                                                        Xy <: Function, J <: Function}
    # Get dimensions
    Nz = length(z)
    Ny = length(y)
    if Nε < 0
        Nε = size(Σ(z), 2)
    end

    # Cache stochastic steady state vectors
    μ_sss, ξ_sss, 𝒱_sss = _cache_sss_vectors(z, y)

    # Cache stochastic steady state Jacobians
    Γ₁, Γ₂, Γ₃, Γ₄, Γ₅, Γ₆, JV = _cache_jacobians(Ψ, Nz, Ny)

    return RiskAdjustedLinearization{T, M, L, S, X, V, J}(μ, Λ, Σ, ξ, 𝒱, J𝒱, μ_sss, ξ_sss, 𝒱_sss,
                                                          Γ₁, Γ₂, Γ₃, Γ₄, Γ₅, Γ₆,
                                                          JV, z, y, Ψ, Nz, Ny, Nε)
end
=#
function RiskAdjustedLinearization(nonlinear::A, linearization::B, z::C1, y::C1, Ψ::C2,
                                   Nz::Int, Ny::Int, Nε::Int;
                                   check_inputs::Bool = true) where {A <: RALNonlinearSystem, B <: RALLinearizedSystem,
                                                                     C1 <: AbstractVector{<: Number}, C2 <: AbstractMatrix{<: Number}}

    # Make sure inputs are well-formed
    if check_inputs
        _check_inputs(nonlinear, linearization, z, y, Ψ)
    end

    return RiskAdjustedLinearization{A, B, C1, C2}(nonlinear, linearization, z, y, Ψ, Nz, Ny, Nε)
end

# Constructor that uses ForwardDiff to calculate Jacobian functions
# NOTE THAT here we pass in the ccgf, rather than 𝒱
function RiskAdjustedLinearization(μ::M, Λ::L, Σ::S, ξ::X, Γ₅::JC5, Γ₆::JC6, ccgf::CF,
                                   z::AbstractVector{T}, y::AbstractVector{T}, Ψ::AbstractMatrix{T},
                                   Nz::Int, Ny::Int, Nε::Int; sss_vector_type::DataType = Vector{T},
                                   jacobian_type::DataType = Matrix{T}) where {T <: Number, M <: Function, L <: RALΛ, S <: RALΣ,
                                                                               X <: Function,
                                                                               JC5 <: AbstractMatrix{<: Number},
                                                                               JC6 <: AbstractMatrix{<: Number},
                                                                               CF <: Function}

    # Cache stochastic steady state vectors
    μ_sss, ξ_sss, 𝒱_sss = _cache_sss_vectors(z, y)

    # Cache stochastic steady state Jacobians
    Γ₁, Γ₂, Γ₃, Γ₄, JV = _cache_jacobians(Ψ, Nz, Ny, jacobian_type)

    # Use cached Jacobians to create Jacobian functions for μ, ξ
    if applicable(μ, z, y) # Check if μ is in place or not
        μz = (F, z, y) -> ForwardDiff.jacobian!(F, x -> μ(x, y), z) # not in place
        μy = (F, z, y) -> ForwardDiff.jacobian!(F, x -> μ(z, x), y)
    else # in place
        μz = (F, z, y, μ_sss) -> ForwardDiff.jacobian!(F, (G, x) -> μ(G, x, y), μ_sss, z)
        μy = (F, z, y, μ_sss) -> ForwardDiff.jacobian!(F, (G, x) -> μ(G, z, x), μ_sss, y)
    end

    if applicable(ξ, z, y) # Check if ξ is in place or not
        ξz = (F, z, y) -> ForwardDiff.jacobian!(F, x -> ξ(x, y), z) # not in place
        ξy = (F, z, y) -> ForwardDiff.jacobian!(F, x -> ξ(z, x), y)
    else # in place
        ξz = (F, z, y, ξ_sss) -> ForwardDiff.jacobian!(F, (G, x) -> ξ(G, x, y), ξ_sss, z)
        ξy = (F, z, y, ξ_sss) -> ForwardDiff.jacobian!(F, (G, x) -> ξ(G, z, x), ξ_sss, y)
    end

    # Create 𝒱 and its Jacobian J𝒱
    if applicable(ccgf, Γ₅, z) # Check if ccgf is in place or not
        𝒱 = function _𝒱(F, z, Ψ, Γ₅, Γ₆)
            F .= ccgf((Γ₅ + Γ₆ * Ψ) * ((I - Λ(z) * Ψ) \ Σ(z)), z)
        end
    else # in place
        𝒱 = (F, z, Ψ, Γ₅, Γ₆) -> ccgf(F, (Γ₅ + Γ₆ * Ψ) * ((I - Λ(z) * Ψ) \ Σ(z)), z)
    end
    J𝒱 = function _J𝒱(F, z, Ψ, Γ₅, Γ₆, 𝒱_sss)
        ForwardDiff.jacobian!(F, (G, x) -> 𝒱(G, x, Ψ, Γ₅, Γ₆), 𝒱_sss, z)
    end

    # Form underlying RAL blocks
    nonlinear_system  = RALNonlinearSystem(μ, Λ, Σ, ξ, 𝒱, μ_sss, ξ_sss, 𝒱_sss, z, y, Ψ, Γ₅, Γ₆)
    linearized_system = RALLinearizedSystem(μz, μy, ξz, ξy, J𝒱, Γ₁, Γ₂, Γ₃, Γ₄, Γ₅, Γ₆, JV, z, y, Ψ, μ_sss, ξ_sss, 𝒱_sss)

    return RiskAdjustedLinearization(nonlinear_system, linearized_system, z, y, Ψ, Nz, Ny, Nε)
end

function RiskAdjustedLinearization(μ::M, Λ::L, Σ::S, ξ::X, Γ₅::JC5, Γ₆::JC6, ccgf::CF,
                                   z::AbstractVector{T}, y::AbstractVector{T}, Ψ::AbstractMatrix{T},
                                   Nε::Int; sss_vector_type::DataType = Vector{T}, sss_matrix_type::DataType = Matrix{T},
                                   jacobian_type::DataType = Matrix{T}) where {T <: Number, M <: Function, L <: Function, S <: Function,
                                                                               X <: Function,
                                                                               JC5 <: AbstractMatrix{<: Number},
                                                                               JC6 <: AbstractMatrix{<: Number},
                                                                               CF <: Function}
    # Get dimensions
    Nz = length(z)
    Ny = length(y)
    if Nε < 0
        error("Nε cannot be negative")
    end

    # Create wrappers enabling caching for Λ and Σ
    Λ = RALΛ(Λ, z, sss_matrix_type, (Nz, Ny))
    Σ = RALΣ(Σ, z, sss_matrix_type, (Nz, Nε))

    return RiskAdjustedLinearization(μ, Λ, Σ, ξ, Γ₅, Γ₆, ccgf, z, y, Ψ, Nz, Ny, Nε, sss_vector_type = sss_vector_type,
                                     jacobian_type = jacobian_type)
end

function RiskAdjustedLinearization(μ::M, Λ::L, Σ::S, ξ::X, Γ₅::JC5, Γ₆::JC6, ccgf::CF,
                                   z::AbstractVector{T}, y::AbstractVector{T}, Ψ::AbstractMatrix{T},
                                   Nε::Int = -1; sss_vector_type::DataType = Vector{T}, sss_matrix_type::DataType = Matrix{T},
                                   jacobian_type::DataType = Matrix{T}) where {T <: Number, M <: Function, L <: AbstractMatrix{<: Number}, S <: Function,
                                                                               X <: Function,
                                                                               JC5 <: AbstractMatrix{<: Number},
                                                                               JC6 <: AbstractMatrix{<: Number},
                                                                               CF <: Function}
    # Get dimensions
    Nz = length(z)
    Ny = length(y)
    if Nε < 0
        error("Nε cannot be negative")
    end

    # Create wrappers enabling caching for Λ and Σ
    Λ = RALΛ(Λ, z)
    Σ = RALΣ(Σ, z, sss_matrix_type, (Nz, Nε))

    return RiskAdjustedLinearization(μ, Λ, Σ, ξ, Γ₅, Γ₆, ccgf, z, y, Ψ, Nz, Ny, Nε, sss_vector_type = sss_vector_type,
                                     jacobian_type = jacobian_type)
end

function RiskAdjustedLinearization(μ::M, Λ::L, Σ::S, ξ::X, Γ₅::JC5, Γ₆::JC6, ccgf::CF,
                                   z::AbstractVector{T}, y::AbstractVector{T}, Ψ::AbstractMatrix{T},
                                   Nε::Int = -1; sss_vector_type::DataType = Vector{T}, sss_matrix_type::DataType = Matrix{T},
                                   jacobian_type::DataType = Matrix{T}) where {T <: Number, M <: Function, L <: Function, S <: AbstractMatrix{<: Number},
                                                                               X <: Function,
                                                                               JC5 <: AbstractMatrix{<: Number},
                                                                               JC6 <: AbstractMatrix{<: Number},
                                                                               CF <: Function}
    # Get dimensions
    Nz = length(z)
    Ny = length(y)
    if Nε < 0
        error("Nε cannot be negative")
    end

    # Create wrappers enabling caching for Λ and Σ
    Λ = RALΛ(Λ, z, sss_matrix_type, (Nz, Ny))
    Σ = RALΣ(Σ, z)

    return RiskAdjustedLinearization(μ, Λ, Σ, ξ, Γ₅, Γ₆, ccgf, z, y, Ψ, Nz, Ny, Nε, sss_vector_type = sss_vector_type,
                                     jacobian_type = jacobian_type)
end

function RiskAdjustedLinearization(μ::M, Λ::L, Σ::S, ξ::X, Γ₅::JC5, Γ₆::JC6, ccgf::CF,
                                   z::AbstractVector{T}, y::AbstractVector{T}, Ψ::AbstractMatrix{T},
                                   Nε::Int = -1; sss_vector_type::DataType = Vector{T}, sss_matrix_type::DataType = Matrix{T},
                                   jacobian_type::DataType = Matrix{T}) where {T <: Number, M <: Function,
                                                                               L <: AbstractMatrix{<: Number}, S <: AbstractMatrix{<: Number},
                                                                               X <: Function,
                                                                               JC5 <: AbstractMatrix{<: Number},
                                                                               JC6 <: AbstractMatrix{<: Number},
                                                                               CF <: Function}
    # Get dimensions
    Nz = length(z)
    Ny = length(y)
    if Nε < 0
        error("Nε cannot be negative")
    end

    # Create wrappers enabling caching for Λ and Σ
    Λ = RALΛ(Λ, z)
    Σ = RALΣ(Σ, z)

    return RiskAdjustedLinearization(μ, Λ, Σ, ξ, Γ₅, Γ₆, ccgf, z, y, Ψ, Nz, Ny, Nε, sss_vector_type = sss_vector_type,
                                     jacobian_type = jacobian_type)
end

function _cache_jacobians(Ψ::AbstractMatrix{T}, Nz::Int, Ny::Int, mat_type::DataType) where {T <: Number}

    Γ₁ = mat_type(undef, Nz, Nz)
    Γ₂ = mat_type(undef, Nz, Ny)
    Γ₃ = similar(Ψ)
    Γ₄ = mat_type(undef, Ny, Ny)
    JV = similar(Ψ)

    return Γ₁, Γ₂, Γ₃, Γ₄, JV
end

function _cache_sss_vectors(z::AbstractVector{T}, y::AbstractVector{T}) where {T <: Number, L, S}

    μ_sss = similar(z)
    ξ_sss = similar(y)
    𝒱_sss = similar(y)

   return μ_sss, ξ_sss, 𝒱_sss
end

function _check_inputs(nonlinear::A, linearization::B, z::C1, y::C1, Ψ::C2) where {A <: RALNonlinearSystem, B <: RALLinearizedSystem,
                                                                                C1 <: AbstractVector{<: Number}, C2 <: AbstractMatrix{<: Number}}

    # Get contents of nonlinear and linearization blocks
    @unpack μ, ξ, 𝒱, μ_sss, ξ_sss, 𝒱_sss = nonlinear
    @unpack μz, μy, ξz, ξy, J𝒱, Γ₁, Γ₂, Γ₃, Γ₄, Γ₅, Γ₆, JV = linearization

    @assert applicable(μ, z, y) ||
        applicable(μ, μ_sss, z, y) "The function μ must take either the form " *
        "μ(z, y) or the in-place equivalent μ(F, z, y)"

    @assert applicable(ξ, z, y) ||
        applicable(ξ, ξ_sss, z, y) "The function μ must take either the form " *
        "ξ(z, y) or the in-place equivalent ξ(F, z, y)"

    @assert applicable(𝒱, z, Ψ, Γ₅, Γ₆) ||
        applicable(𝒱, y, z, Ψ, Γ₅, Γ₆) "The function 𝒱 must take either the form " *
        "𝒱(z, Ψ, Γ₅, Γ₆) or the in-place equivalent 𝒱(F, z, Ψ, Γ₅, Γ₆)"

    @assert applicable(μz, Γ₁, z, y) ||
        applicable(μz, Γ₁, z, y, μ_sss) "The function μz must take either the form " *
        "μz(F, z, y) or μz(F, z, y, μ_sss)"

    @assert applicable(μy, Γ₂, z, y) ||
        applicable(μy, Γ₂, z, y, μ_sss) "The function μy must take either the form " *
        "μy(F, z, y) or μy(F, z, y, μ_sss)"

    @assert applicable(ξz, Γ₃, z, y) ||
        applicable(ξz, Γ₃, z, y, ξ_sss) "The function ξz must take either the form " *
        "ξz(F, z, y) or ξz(F, z, y, ξ_sss)"

    @assert applicable(ξy, Γ₄, z, y) ||
        applicable(ξy, Γ₄, z, y, ξ_sss) "The function ξy must take either the form " *
        "ξy(F, z, y) or ξy(F, z, y, ξ_sss)"

    @assert applicable(J𝒱, z, Ψ, Γ₅, Γ₆) ||
        applicable(J𝒱, JV, z, Ψ, Γ₅, Γ₆, 𝒱_sss) "The function J𝒱 must take either the form " *
        "J𝒱(F, z, Ψ, Γ₅, Γ₆) or J𝒱(F, z, Ψ, Γ₅, Γ₆, 𝒱_sss)"
end

## Methods for using RiskAdjustedLinearization
@inline Γ₁(m::RiskAdjustedLinearization) = m.linearization.Γ₁
@inline Γ₂(m::RiskAdjustedLinearization) = m.linearization.Γ₂
@inline Γ₃(m::RiskAdjustedLinearization) = m.linearization.Γ₃
@inline Γ₄(m::RiskAdjustedLinearization) = m.linearization.Γ₄
@inline Γ₅(m::RiskAdjustedLinearization) = m.linearization.Γ₅
@inline Γ₆(m::RiskAdjustedLinearization) = m.linearization.Γ₆
@inline JV(m::RiskAdjustedLinearization) = m.linearization.JV
@inline getvalues(m::RiskAdjustedLinearization) = (m.z, m.y, m.Ψ)
@inline getvecvalues(m::RiskAdjustedLinearization) = vcat(m.z, m.y, vec(m.Ψ))
@inline nonlinear_system(m::RiskAdjustedLinearization) = m.nonlinear
@inline linearized_system(m::RiskAdjustedLinearization) = m.linearization

function update!(m::RiskAdjustedLinearization, z::C1, y::C1, Ψ::C2;
                 update_cache::Bool = true) where {C1 <: AbstractVector{<: Number}, C2 <: AbstractMatrix{<: Number}}

    # Update values of the affine approximation
    m.z .= z
    m.y .= y
    m.Ψ .= Ψ

    # Update the cached vectors and Jacobians
    if update_cache
        update!(nonlinear_system(m), m.z, m.y, m.Ψ, Γ₅(m), Γ₆(m))

        update!(linearized_system(m), m.z, m.y, m.Ψ, m.nonlinear.μ_sss, m.nonlinear.ξ_sss, m.nonlinear.𝒱_sss)
    end

    m
end

function Base.show(io::IO, m::RiskAdjustedLinearization)
    @printf io "Risk-Adjusted Linearization of a Dynamic Economic Model\n"
    @printf io "No. of state variables:              %i\n" m.Nz
    @printf io "No. of jump variables:               %i\n" m.Ny
    @printf io "No. of exogenous shocks:             %i\n" m.Nε
end


module NBody

# include("librationpoints.jl")

end


using Test
using LoopVectorization

using LinearAlgebra
import InteractiveUtils

InteractiveUtils.versioninfo(stdout; verbose = true)

const START_TIME = time()

function clenshaw(x, coeff)
    len_c = length(coeff)
    tmp = zero(x)
    ret = zero(x)
    for i in len_c:-1:2
        ret     = muladd(x,2tmp,coeff[i]-ret)
        ret,tmp = tmp,ret
    end
    ret = muladd(x,tmp,coeff[1]-ret)
    return ret
end

"""
Causes `check_args` to fail.
"""
struct FallbackArrayWrapper{T,N} <: AbstractArray{T,N}
    data::Array{T,N}
end
Base.size(A::FallbackArrayWrapper) = size(A.data)
Base.@propagate_inbounds Base.getindex(A::FallbackArrayWrapper, i::Vararg{Int, N}) where {N} = getindex(A.data, i...)
Base.@propagate_inbounds Base.setindex!(A::FallbackArrayWrapper, v, i::Vararg{Int, N}) where {N} = setindex!(A.data, v, i...)
Base.IndexStyle(::Type{<:FallbackArrayWrapper}) = IndexLinear()

@show LoopVectorization.REGISTER_COUNT

const RUN_SLOW_TESTS = LoopVectorization.REGISTER_COUNT ≤ 16 || !parse(Bool, get(ENV, "GITHUB_ACTIONS", "false"))
@show RUN_SLOW_TESTS

@time @testset "LoopVectorization.jl" begin

    @test isempty(detect_unbound_args(LoopVectorization))

    @time include("printmethods.jl")

    @time include("can_avx.jl")

    @time include("fallback.jl")

    @time include("utils.jl")

    @time include("arraywrappers.jl")

    @time include("check_empty.jl")

    @time include("offsetarrays.jl")

    @time include("tensors.jl")

    @time include("map.jl")

    @time include("filter.jl")

    @time include("mapreduce.jl")

    @time include("ifelsemasks.jl")

    @time include("dot.jl")

    @time include("special.jl")

    @time include("gemv.jl")

    @time include("miscellaneous.jl")

    @time include("copy.jl")

    @time include("broadcast.jl")

    @time include("gemm.jl")
end

const ELAPSED_MINUTES = (time() - START_TIME)/60
@test ELAPSED_MINUTES < 180


module DensityConsistency

using ExtractMacro, LinearAlgebra, SparseArrays, Random
export DCState, density_consistency

include("types.jl")
include("closures.jl")
include("dc_main.jl")

end # end module


module Endowment

if  pwd() ∉ LOAD_PATH
  push!(LOAD_PATH,pwd())
end

if "$(pwd())\\data" ∉ LOAD_PATH
  push!(LOAD_PATH,"$(pwd())\\data")
end

#Use this to print
#=
using Weave
codeName = "train"
weave(Pkg.dir("$(pwd())\\$(codeName).jl"),
  informat="script",
  out_path = "$(pwd())\\output\\$(codeName)_Appendix.html",
  doctype = "md2html")
=#
##############0.7/1.0 compatability hacks
#NOTE: Delete these upon 0.7

#############Dependencies####################
using DataFrames, Distributions, CSV, GZip, ForwardDiff,
  Gadfly, NLopt, Measures, Formatting, HTTP, TexTables,
  CodecZlib, StatsBase, Distributions, DataStructures

#using CT

##################Exports########################
export initializePar, ###Acquire NCCSData
cleanup,
downloadNCCSCSVs,
preProcessData,
validateSingleTypeCSV,
readSingleTypeCSV,
readMetaIndex,
formMetaData,
MetaTuple,
MetaDictionary,
NCCSFile,
NCCS_PATH,

MBool,
MInt,
MFloat64,
MSymbol,
MString

#############################Shared Constants###########
const OUTPUT_PATH = pwd() * "\\output"
const WORKING_PATH = pwd() * "\\working"
const DATA_PATH = pwd() * "\\data"
const NCCS_PATH = DATA_PATH * "\\NCCS"
const DOCUMENTATION_PATH = DATA_PATH * "\\documentation"

const HEADER_NAME = "header.tex"
const FOOTER_NAME = "footer.tex"
const METADATA_NAME = "NCCSMetadata"
const META_INDEX_NAME = "NCCSMetaIndex"
const META_PATH = DOCUMENTATION_PATH * "\\raw meta"
const SUB_TABLES_PATH = META_PATH * "\\subtables"
const FILES_INDEX_NAME = "NCCSFiles"
const KEEPER_NAME = "Keeper List"
const Nothing = Void #NOTE: 0.7 compatabiltiy hack, delete on ugrade
const VAR_TYPES = Dict(:Numeric=>Float64, :Character=>String)

const MBool = Union{Bool, Missing}
const MInt = Union{Int, Missing}
const MFloat64 = Union{Float64, Missing}
const MSymbol = Union{Symbol, Missing}
const MString = Union{String, Missing}

#default compressor algorithms
I_CSV_STREAMER(p::String) = GzipDecompressorStream(open(p))
I_JLS_STREAMER(p::String) = GzipDecompressorStream(open(p))
O_JLS_STREAMER(p::String) = GzipCompressorStream(open(p, "w"))

include("EndowmentCode\\EndowmentPar.jl") #holds supporting functions for parallelization
include("EndowmentCode\\MetaTuples.jl") #holds methods related to field-level metadata
include("EndowmentCode\\MetaDictionaries.jl") #holds methods related to file level metadata
include("EndowmentCode\\Metadata.jl") #holds all other metadata information
include("EndowmentCode\\DownloadNCCSCSVs.jl") #holds methods related to downloading the data
include("EndowmentCode\\NCCSFiles.jl") #holds methods related to the data files
include("EndowmentCode\\EndowmentData.jl") #holds high-level scripts and other miscellaneous

end


### A Pluto.jl notebook ###
# v0.19.0

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local iv = try Base.loaded_modules[Base.PkgId(Base.UUID("6e696c72-6542-2067-7265-42206c756150"), "AbstractPlutoDingetjes")].Bonds.initial_value catch; b -> missing; end
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : iv(el)
        el
    end
end

# ╔═╡ aa854f02-46bf-4a2b-b1e9-5fb52b99925c
begin
    import Pkg
    Pkg.activate(Base.current_project())
    Pkg.instantiate()
end

# ╔═╡ 8b4426b8-588b-49ad-83c5-a79ed698d704
using PlutoUI

# ╔═╡ d47c1cf4-4844-4c66-8a21-b62ad7740dcf
using JLD2

# ╔═╡ 3c07c33b-7f91-4e2d-a27a-09a8924f328c
using DataFrames

# ╔═╡ b59dd2c0-4c3e-4c3d-a4d4-f8da9dcbe286
using CSV

# ╔═╡ 76e2e659-c708-4a08-8263-def6c3dfa930
using Dates

# ╔═╡ 41b0f95e-d806-4167-8199-7392717782f9
using OffsetArrays

# ╔═╡ 95f2bd0b-7d8b-43b9-b2f9-cfa6628ec485
using LinearAlgebra

# ╔═╡ 356cc4c2-e565-4847-9286-cfc1f835654d
using StatsBase

# ╔═╡ 730340bb-8176-4282-8559-8d60b1d28d57
using Statistics

# ╔═╡ f07ab63f-dfb4-493c-a40c-6b3363a858c3
using Plots

# ╔═╡ 342cbcff-5a05-4199-b5b5-6cc0fa70c202
using TemplateMatching

# ╔═╡ 7ddcab3c-a100-44be-baa3-af28cc08b2d7
using Interpolations

# ╔═╡ 341b9066-dcb3-4d99-b985-82e24c11f28e
using Optim

# ╔═╡ c2f22858-6d3c-4f92-a05b-659deae566f6
md"## Reading Continuous Data"

# ╔═╡ a86b6c8a-9b90-4388-8208-d257ec6951d2
md"""Path of the JLD2 file containing continuous data: $(@bind datapath TextField(default="../data/2021-01-12_20-25-30/data.jld2"))"""

# ╔═╡ f3f17c65-dab6-46da-83a9-87aad1181524
data = load(datapath, "data")

# ╔═╡ 8271f7e6-d8ac-472a-9538-75d0afacf42e
samplefreq = 2 # MHz

# ╔═╡ 57b35604-727b-4eac-bcbd-ea302e4c79a2
md"## Reading events catalogue"

# ╔═╡ b62d7318-3a6d-4442-b6cc-3166b87dff8a
md"""Path of the CSV catalogue $(@bind cataloguepath TextField(default="../data/2021-01-12_20-25-30/cat.csv"))"""

# ╔═╡ f3139f19-e386-4bb7-8aca-fa54062d4093
begin
	import Base.+, Base.show

	+(x::Dates.UTInstant{Microsecond}, d::Microsecond) = Dates.UTInstant(x.periods + d)

	function show(io::IO, ::MIME"text/plain", x::Dates.UTInstant{Microsecond})
		date = DateTime(Dates.UTM(round(Int, 1e-3 * x.periods.value)))
		show(io, MIME"text/plain"(), date)
	end
end

# ╔═╡ 9de113e1-36ce-4251-9fa5-5baa27941a85
function DateTimeMicrosecond(y, m, d, h, mi, s, us)
	rata = us + 1_000_000 * (s + 60mi + 3600h + 86400 * Dates.totaldays(y, m, d))
	Dates.UTInstant{Microsecond}(Microsecond(rata))
end

# ╔═╡ 96cf6e3b-1c2b-49f6-979d-0fce28a55a06
starttime = DateTimeMicrosecond(2021, 01, 12, 20, 25, 30, 0)

# ╔═╡ 5d5f3ca7-260d-4dbb-bde1-59f2bd58c0c9
begin
	df = CSV.read(cataloguepath, DataFrame)
	sec = floor.(Int, df.Second)
	usec = round.(Int, 1e6 * (df.Second - sec))
	templates = DataFrame()
	templates.datetime = DateTimeMicrosecond.(df.Year, df.Month, df.Day, df.Hour, df.Minute, sec, usec)
	templates.sample = map(x -> round(Int, x.value * samplefreq), templates.datetime .- starttime)
	templates.north = df.North
	templates.east = df.East
	templates.up = df.Up
	templates
end

# ╔═╡ 6ac6cce3-09ae-421a-b502-b5c8807bf555
md"## Reading Sensors Positions"

# ╔═╡ be090485-7917-464a-8e8c-aec4a2c009d3
begin
	sensors_positions = CSV.read("../data/2021-01-12_20-25-30/passive-xyz.csv", DataFrame, header=[:north, :east, :up])
	sensors_positions .*= 100
	sensors_positions.sensor = axes(sensors_positions, 1) .- 1
	filter!(r -> r.sensor in keys(data), sensors_positions)
	sensors_positions
end

# ╔═╡ 6ab4a8da-4613-4004-807f-c36876285ebd
md" ## Preparing a single template"

# ╔═╡ fe42e7e6-2c56-4f11-aae6-ec06ad7a219f
md"""template number: $(@bind template_num Select(1:nrow(templates), default=116))"""

# ╔═╡ 298010b3-c34c-40ba-9e1b-b79f7378f1d4
begin
	displacement = DataFrame()
	for s in [:north, :east, :up]
		displacement[!, s] = sensors_positions[!, s] .- templates[template_num, s]
	end
	displacement
end

# ╔═╡ a1fc5e84-48df-45ea-b6a1-0073ed9f8340
distances = hypot.(displacement[!, :north], displacement[!, :east], displacement[!, :up]) # in cm

# ╔═╡ c7f0582a-8997-4641-b607-0b95a7f73a6c
v_p = 0.67345 # in cm/us

# ╔═╡ 074983c0-2747-4b09-a0c0-4d4adba24c0b
offsets = @. round(Int, samplefreq * (distances / v_p))

# ╔═╡ cd5f912a-fafb-4947-bd70-318151c2c49c
shifts = templates[template_num, :sample] .+ offsets

# ╔═╡ 033473d3-6958-4e7b-8c95-a9e15d882118
md"""
Template window pre (in samples): $(@bind t_pre Slider(0:10:1000, default=100, show_value=true))

Template window post (in samples): $(@bind t_post Slider(0:10:1000, default=500, show_value=true))
"""

# ╔═╡ be80dd80-d39d-4338-bb76-c644ea9d970c
cuts = [shift - t_pre: shift + t_post for shift in shifts]

# ╔═╡ aa23d53f-8333-4c94-8932-fac99ce9196d
template_data = [view(series, cut) for (series, cut) in zip(values(data), cuts)]

# ╔═╡ 4ef908f4-d8dc-4159-97ad-fc8d4621a652
let
	template_aligned = OffsetVector.(template_data, shifts)
	xlim = minimum(firstindex(series, 1) for series in template_aligned), maximum(lastindex(series, 1) for series in template_aligned)
	plots = []
	for (nch, series) = zip(keys(data), template_aligned)
		ylim = minimum(series), maximum(series)
		push!(plots,
			  plot(axes(series, 1), series;
				   xlim,
				   ylim,
			       title="Channel $nch",
				   titlealign=:left,
				   showaxis=false,
				   ticks=false,
				   label=nothing))
	end
	numchannels = length(template_data)
	plot(plots..., layout=(numchannels, 1), size=(800, 100numchannels), dpi=300)
end

# ╔═╡ 368f7fc1-0afc-439f-905a-76af04e8ca8d
md" ## Computing the cross-correlation"

# ╔═╡ 0d122aef-1767-4389-b043-97d07ecef1fe
md"""
tolerance $(@bind tolerance Slider(0:10, default=5, show_value=true))
"""

# ╔═╡ a7a589fc-4dcc-4137-a423-b8d2e7d41a0e
signal = correlatetemplate(collect(values(data)), template_data, offsets, tolerance)

# ╔═╡ 91924a12-6ee4-4d1b-8c49-9b54142cf733
let
	template_ref = templates[template_num, :sample] - t_pre
	window = t_pre + t_post - 1
	a = max(firstindex(signal, 1), template_ref - div(window, 2))
	b = min(lastindex(signal, 1), template_ref + div(window, 2))
	plot(a:b, signal[a:b], label=nothing, dpi=300)
end

# ╔═╡ abb8b867-39c2-40f2-9d7a-f40bacf8f550
md" ## Finding peaks"

# ╔═╡ 7d1f1903-8342-4715-8e88-3e3d063ea5db
md"""
correlation threshold $(@bind threshold Slider(0.0:0.05:1, default=0.4, show_value=true))

distance (in unit of template length) $(@bind reldistance Slider(0:0.5:5, default=2, show_value=true))
"""

# ╔═╡ 666c8d5f-c1c5-4cb4-9364-bf997b18a683
peaks, heights = findpeaks(signal, threshold, reldistance * (t_pre + t_post - 1))

# ╔═╡ 96deaa3d-fc40-466f-bf1b-d8d3dd58dcd4
md"""
Peak number $(@bind peak_num Slider(axes(peaks, 1), default=1, show_value=true)) / $(length(peaks))
"""

# ╔═╡ cc0c3e0a-4208-4339-bea1-75396b2778bd
let
	plt = plot(ylims=(0.0, 1.0))
	scale = round(Int, reldistance * (t_post + t_pre - 1))
	a = max(firstindex(signal, 1), peaks[peak_num] - scale)
	b = min(lastindex(signal, 1), peaks[peak_num] + scale)
	plot!(plt, a:b, signal[a:b], label=nothing)
	plot!(plt, peaks[a .< peaks .< b], heights[a.< peaks .< b],
		seriestype=:scatter, label=nothing, dpi=300)
	plt
end

# ╔═╡ b6848d59-e97e-4662-a2ee-78612694d6d8
md"## Relocalization"

# ╔═╡ 89b036db-5c47-4a17-8466-9062e8d5d452
toas = estimatetoa.(collect(values(data)), template_data, peaks[peak_num] .+ offsets, tolerance)

# ╔═╡ 6d6d27c9-3b1a-4646-8420-61a775f7842f
let
	peak = peaks[peak_num]
	peak_start = peak - div(t_pre, 2)
	peak_stop = peak + (t_pre + t_post - 1) + div(t_post, 2)
	data_start = minimum(series -> firstindex(series, 1), values(data))
	data_stop = maximum(series -> lastindex(series, 1), values(data))
	xlim = max(data_start, peak_start), min(data_stop, peak_stop)

	plots = []
	for (n, channel_num) = enumerate(keys(data))
		series_start = 	max(firstindex(data[channel_num], 1), peak_start)
		series_stop = min(lastindex(data[channel_num], 1), peak_stop)
		data_series = OffsetVector(view(data[channel_num], series_start:series_stop), series_start:series_stop)

		toa_f, toa_i = modf(toas[n][1])
		template_interpolated = subsampleshift(template_data[n], toa_f)
		template_aligned = OffsetVector(template_interpolated, Int(toa_i))

		data_inf, data_sup = extrema(data_series)
		data_center = mean(data_series)
		template_inf, template_sup = extrema(template_interpolated)
		template_center = mean(template_aligned)
		β = (data_sup - data_inf) / (template_sup - template_inf)
		push!(plots,
			  plot([axes(data_series, 1), axes(template_aligned, 1)],
				   [data_series, data_center .+ β .* (template_aligned .- template_center)];
				   xlim,
				   ylim=(data_inf, data_sup),
				   title="Channel $channel_num",
				   titlealign=:left,
				   showaxis=false,
				   ticks=false,
				   label=nothing))
	end
	numchannels = length(template_data)
	plot(plots..., layout=(numchannels, 1), size=(800, 100numchannels), dpi=300)
end

# ╔═╡ 2911437d-4775-4a1b-819e-adf6e91acd61
measurements = let correlation_threshold = 0.5
	measurements = copy(sensors_positions)
	measurements.toas = [(sample + t_pre) / samplefreq for (sample, _) in toas]
	measurements.cc = [cc for (_, cc) in toas]
	filter(row -> row.cc > correlation_threshold, measurements)
end

# ╔═╡ 2e6b7b35-9512-406d-9a97-f1f94bea8732
if nrow(measurements) > 4
	x0 = Vector(templates[template_num, [:north, :east, :up]])
	t0 = (peaks[peak_num] + t_pre) / samplefreq
	sensors_readings = eachrow(Matrix(measurements[!, [:north, :east, :up, :toas]]))
	candidate = locate(sensors_readings, v_p, [x0; t0])
end

# ╔═╡ 32ab23a5-60b6-4f07-8c0f-cfa2e0a9d7e1
candidate.minimizer

# ╔═╡ 91f9aee4-6f80-4b02-ab09-1ace3759ea22
let dx = 1, dy = 1, npx = 100, npy = 100
	y0, x0, z0, t0 = candidate.minimizer
	xrange = range(x0 - dx, x0 + dx, npx)
	yrange = range(y0 - dy, y0 + dy, npy)
	sensors_readings = eachrow(Matrix(measurements[!, [:north, :east, :up, :toas]]))
	residues = reshape(
		[residue_rms([y, x, z0, t0], sensors_readings, v_p)
			for x = xrange, y = yrange],
		npx, npy)
	contour(xrange, yrange, residues, fill=:true, c=:heat)
	y1, x1 = templates[template_num, [:north, :east]]
	scatter!([x1], [y1], markershape=:circle, label="Template")
	scatter!([x0], [y0], markershape=:x, label="Event")
	plot!(title="X-Y Section")
end

# ╔═╡ 5c40564b-bb44-47f4-a98e-e09f161e2d68
let dx = 1, dz = 1, npx = 100, npz = 100
	y0, x0, z0, t0 = candidate.minimizer
	xrange = range(x0 - dx, x0 + dx, npx)
	zrange = range(z0 - dz, z0 + dz, npz)
	sensors_readings = eachrow(Matrix(measurements[!, [:north, :east, :up, :toas]]))
	residues = reshape(
		[TemplateMatching.residue_rms([y0, x, z, t0], sensors_readings, v_p)
			for x = xrange, z = zrange],
		npx, npz)
	contour(xrange, zrange, residues, fill=:true, c=:heat)
	z1, x1 = templates[template_num, [:up, :east]]
	scatter!([x1], [z1], markershape=:circle, label="Template")
	scatter!([x0], [z0], markershape=:x, label="Event")
	plot!(title="X-Z Section")
end

# ╔═╡ 102e4534-6be4-4436-be36-69726a393253
md" ## Statistics of the detections "

# ╔═╡ 33bc3279-c65f-47b5-acf4-60853c2bdb91
begin
	catalogue = DataFrame()
	samples = peaks .+ t_pre
	catalogue.datetimes = @. starttime + Microsecond(samples)
	catalogue.sample = samples
	catalogue.correlation = heights
	catalogue.rel_mag = [magnitude(values(data), template_data, peak .+ offsets) for peak in peaks]
	catalogue
end

# ╔═╡ 471f21ba-23d3-42f4-8d62-12fe2f36124f
histogram(collect(signal), label=nothing)

# ╔═╡ 6e2e9172-4cf5-45f3-a239-12c09b6ca3da
summarystats(collect(signal))

# ╔═╡ feebcd5f-91c4-4a70-a064-e642e3172614
histogram(catalogue[!, :correlation], label=nothing)

# ╔═╡ 476101af-0108-4628-9ac1-f233a456a869
summarystats(catalogue[!, :correlation])

# ╔═╡ 3be237b0-9f19-4e9a-aa13-c6355771b8e8
histogram(catalogue[!, :rel_mag], label=nothing)

# ╔═╡ Cell order:
# ╠═aa854f02-46bf-4a2b-b1e9-5fb52b99925c
# ╟─c2f22858-6d3c-4f92-a05b-659deae566f6
# ╠═8b4426b8-588b-49ad-83c5-a79ed698d704
# ╠═d47c1cf4-4844-4c66-8a21-b62ad7740dcf
# ╟─a86b6c8a-9b90-4388-8208-d257ec6951d2
# ╠═f3f17c65-dab6-46da-83a9-87aad1181524
# ╠═8271f7e6-d8ac-472a-9538-75d0afacf42e
# ╟─57b35604-727b-4eac-bcbd-ea302e4c79a2
# ╠═3c07c33b-7f91-4e2d-a27a-09a8924f328c
# ╠═b59dd2c0-4c3e-4c3d-a4d4-f8da9dcbe286
# ╠═76e2e659-c708-4a08-8263-def6c3dfa930
# ╟─b62d7318-3a6d-4442-b6cc-3166b87dff8a
# ╠═9de113e1-36ce-4251-9fa5-5baa27941a85
# ╠═96cf6e3b-1c2b-49f6-979d-0fce28a55a06
# ╠═f3139f19-e386-4bb7-8aca-fa54062d4093
# ╠═5d5f3ca7-260d-4dbb-bde1-59f2bd58c0c9
# ╟─6ac6cce3-09ae-421a-b502-b5c8807bf555
# ╠═be090485-7917-464a-8e8c-aec4a2c009d3
# ╟─6ab4a8da-4613-4004-807f-c36876285ebd
# ╠═41b0f95e-d806-4167-8199-7392717782f9
# ╠═95f2bd0b-7d8b-43b9-b2f9-cfa6628ec485
# ╠═356cc4c2-e565-4847-9286-cfc1f835654d
# ╠═730340bb-8176-4282-8559-8d60b1d28d57
# ╠═f07ab63f-dfb4-493c-a40c-6b3363a858c3
# ╟─fe42e7e6-2c56-4f11-aae6-ec06ad7a219f
# ╠═298010b3-c34c-40ba-9e1b-b79f7378f1d4
# ╠═a1fc5e84-48df-45ea-b6a1-0073ed9f8340
# ╠═c7f0582a-8997-4641-b607-0b95a7f73a6c
# ╠═074983c0-2747-4b09-a0c0-4d4adba24c0b
# ╠═cd5f912a-fafb-4947-bd70-318151c2c49c
# ╟─033473d3-6958-4e7b-8c95-a9e15d882118
# ╠═be80dd80-d39d-4338-bb76-c644ea9d970c
# ╠═aa23d53f-8333-4c94-8932-fac99ce9196d
# ╟─4ef908f4-d8dc-4159-97ad-fc8d4621a652
# ╟─368f7fc1-0afc-439f-905a-76af04e8ca8d
# ╠═342cbcff-5a05-4199-b5b5-6cc0fa70c202
# ╟─0d122aef-1767-4389-b043-97d07ecef1fe
# ╠═a7a589fc-4dcc-4137-a423-b8d2e7d41a0e
# ╟─91924a12-6ee4-4d1b-8c49-9b54142cf733
# ╟─abb8b867-39c2-40f2-9d7a-f40bacf8f550
# ╟─7d1f1903-8342-4715-8e88-3e3d063ea5db
# ╠═666c8d5f-c1c5-4cb4-9364-bf997b18a683
# ╟─96deaa3d-fc40-466f-bf1b-d8d3dd58dcd4
# ╟─cc0c3e0a-4208-4339-bea1-75396b2778bd
# ╟─b6848d59-e97e-4662-a2ee-78612694d6d8
# ╠═7ddcab3c-a100-44be-baa3-af28cc08b2d7
# ╠═341b9066-dcb3-4d99-b985-82e24c11f28e
# ╠═89b036db-5c47-4a17-8466-9062e8d5d452
# ╠═6d6d27c9-3b1a-4646-8420-61a775f7842f
# ╠═2911437d-4775-4a1b-819e-adf6e91acd61
# ╠═2e6b7b35-9512-406d-9a97-f1f94bea8732
# ╠═32ab23a5-60b6-4f07-8c0f-cfa2e0a9d7e1
# ╟─91f9aee4-6f80-4b02-ab09-1ace3759ea22
# ╟─5c40564b-bb44-47f4-a98e-e09f161e2d68
# ╟─102e4534-6be4-4436-be36-69726a393253
# ╠═33bc3279-c65f-47b5-acf4-60853c2bdb91
# ╠═471f21ba-23d3-42f4-8d62-12fe2f36124f
# ╠═6e2e9172-4cf5-45f3-a239-12c09b6ca3da
# ╠═feebcd5f-91c4-4a70-a064-e642e3172614
# ╠═476101af-0108-4628-9ac1-f233a456a869
# ╠═3be237b0-9f19-4e9a-aa13-c6355771b8e8


"""
    reproject(source::GeoFormat, target::GeoFormat, dim::Dimension, val)

`reproject` uses ArchGDAL.reproject, but implemented for a reprojecting
a value array of values, a single dimension at a time.
"""
reproject(source, target, dim::Dimension, val) = val
function reproject(source::GeoFormat, target::GeoFormat, dim::XDim, val::Number)
    AG.reproject((val, zero(val)), source, target; order=:trad)[1]
end
function reproject(source::GeoFormat, target::GeoFormat, dim::YDim, val::Number)
    AG.reproject((zero(val), val), source, target; order=:trad)[2]
end
function reproject(source::GeoFormat, target::GeoFormat, ::XDim, vals::AbstractArray)
    rep = AG.reproject(map(v -> (v, zero(v)), vals), source, target; order=:trad)
    map(r -> r[1], rep)
end
function reproject(source::GeoFormat, target::GeoFormat, dim::YDim, vals::AbstractArray)
    rep = AG.reproject(map(v -> (zero(v), v), vals), source, target; order=:trad)
    map(r -> r[2], rep)
end
function reproject(source::GeoFormat, target::GeoFormat, ::XDim, vals::Tuple)
    reps = AG.reproject([(v, zero(v)) for v in vals], source, target; order=:trad)
    Tuple(r[1] for r in reps)
end
function reproject(source::GeoFormat, target::GeoFormat, dim::YDim, vals::Tuple)
    reps = AG.reproject([(zero(v), v) for v in vals], source, target; order=:trad)
    Tuple(r[2] for r in reps)
end


abstract type AbstractContext <: AbstractElement end

struct Style <: AbstractElement
    context::Union{AbstractContext, Nothing}
    value::NamedTuple
end
Style(s::NamedTuple=NamedTuple()) = Style(nothing, s)
Style(c::AbstractContext) = Style(c, NamedTuple())
Style(s::Style) = s

function Base.show(io::IO, s::Style)
    print(io, "Style with entries ")
    show(io, keys(s.value))
end

style(args...; kwargs...) = Style(namedtuple(args...; kwargs...))

function Base.:*(s1::Union{Style, AbstractContext}, s2::Union{Style, AbstractContext})
    merge(Style(s1), Style(s2))
end

function Base.merge(s1::Style, s2::Style)
    c1, c2 = s1.context, s2.context
    c = c2 === nothing ? c1 : c2
    if c !== c1 && !isempty(s1.value)
        @warn "Changing context on a non empty style"
    end
    return _merge(c, s1, s2)
end
Base.pairs(s::Style) = _pairs(s.context, s)

# interface and fallbacks

# Must return a Style
_merge(c, s1::Style, s2::Style) = Style(c, merge(s1.value, s2.value))

# Must return a vector of NamedTuple => Style pairs
_pairs(c, s::Style) = [NamedTuple() => s]

## Dims context

struct DimsSelector{N} <: AbstractContext
    dims::NTuple{N, Int}
end
dims(args...) = DimsSelector(args)

Broadcast.broadcastable(d::DimsSelector) = fill(d)

Base.isless(s1::DimsSelector, s2::DimsSelector) = isless(s1.dims, s2.dims)

function aos(d::NamedTuple{names}) where names
    v = broadcast((args...) -> NamedTuple{names}(args), d...)
    return v isa NamedTuple ? fill(v) : v
end

function adjust(val::DimsSelector, c, nts)
    is = collect(val.dims)
    ax, cf = axes(nts)[is], Tuple(c)[is]
    return LinearIndices(ax)[cf...]
end

# make it a pair list
function _pairs(::DimsSelector{0}, s::Style)
    nts = aos(s.value)
    ps = map(CartesianIndices(nts)) do c
        nt = nts[c]
        dskeys = filter(keys(nt)) do key
            nt[key] isa DimsSelector
        end
        dsvalues = map(dskeys) do key
            val = nt[key]
            return adjust(val, c, nts)
        end
        ds = (; zip(dskeys, dsvalues)...)
        nds = Base.structdiff(nt, ds)
        return ds => Style(nds)
    end
    return vec(ps)
end

function _pairs(d::DimsSelector, s::Style)
    value = map(s.value) do v
        v isa AbstractArray ? mapslices(x -> [x], v; dims=d.dims) : v
    end
    return pairs(Style(dims(), value))
end

# data context: integers and symbols are columns

struct DataContext{NT, I<:AbstractVector{Int}} <: AbstractContext
    data::ColumnDict
    pkeys::NT
    perm::I
end

function DataContext(data)
    col = getcolumn(data, 1)
    DataContext(data, NamedTuple(), axes(col, 1))
end

data(x) = DataContext(ColumnDict(x))

iscategorical(v) = !(eltype(v) <: Number)

iswrappedcategorical(::Any) = false
iswrappedcategorical(v::AbstractArray{<:Any, 0}) = iscategorical(v[])

# unwrap categorical vectors from the 0-dim array that contains them
function unwrap_categorical(values)
    iter = (k => v[] for (k, v) in pairs(values) if iswrappedcategorical(v))
    return (; iter...)
end

function refine_perm(perm, pc, n)
    if n == length(pc)
        perm
    elseif n == 0
        sortperm(StructArray(pc))
    else
        x, y = pc[n], pc[n+1]
        refine_perm!(copy(perm), pc, n, x, y, 1, length(x))
    end
end

optimize_cols(v::NamedDimsArray) = refs(parent(v))
optimize_cols(v::Union{Tuple, NamedTuple}) = map(optimize_cols, v)

function _merge(c::DataContext, s1::Style, s2::Style)
    data, pkeys, perm = c.data, c.pkeys, c.perm
    nt = map(val -> extract_columns(data, val), s2.value)
    newpkeys = unwrap_categorical(keyword(nt))
    @assert isempty(intersect(keys(pkeys), keys(newpkeys)))
    allpkeys = merge(pkeys, newpkeys)
    # unwrap from NamedDimsArray to perform the sorting
    allperm = refine_perm(perm, optimize_cols(allpkeys), length(pkeys))
    ctx = DataContext(data, allpkeys, allperm)
    return Style(ctx, merge(s1.value, Base.structdiff(nt, newpkeys)))
end

function _pairs(c::DataContext, s::Style)
    data, pkeys, perm = c.data, c.pkeys, c.perm
    isempty(pkeys) && return pairs(Style(dims(), s.value))
    # unwrap for sorting computations
    itr = GroupPerm(StructArray(optimize_cols(pkeys)), perm)
    sa = StructArray(pkeys)
    nestedpairs = map(itr) do idxs
        i1 = perm[first(idxs)]
        # keep value categorical and preserve name by taking a mini slice (document this!)
        # consider using these levels directly in discrete scales
        k = map(col -> col[i1:i1], pkeys)
        cols = map(s.value) do val
            extract_views(val, perm[idxs])
        end
        [merge(k, p) => v for (p, v) in pairs(Style(dims(), cols))]
    end
    return reduce(vcat, nestedpairs)
end

extract_views(cols, idxs) = map(t -> view(t, idxs), cols)
extract_views(cols::DimsSelector, idxs) = cols

function extract_column(t, (nm, f)::Pair)
    colname = nm isa Symbol ? nm : columnnames(t)[nm]
    vals = f(getcolumn(t, nm))
    p = iscategorical(vals) ? categorical(vals) : vals
    return NamedDimsArray{(colname,)}(p)
end
extract_column(t, nm::Union{Symbol, Int}) = extract_column(t, nm => identity)

extract_columns(t, val::DimsSelector) = val
extract_columns(t, val::Union{Tuple, AbstractArray}) = map(v -> extract_column(t, v), val)
extract_columns(t, val) = fill(extract_column(t, val))


# --------------------------------------------------------------------------
# ACE.jl and SHIPs.jl: Julia implementation of the Atomic Cluster Expansion
# Copyright (c) 2019 Christoph Ortner <christophortner0@gmail.com>
# Licensed under ASL - see ASL.md for terms and conditions.
# --------------------------------------------------------------------------

using JuLIP, ACE, ZipFile, JSON

#---

species_list = [ [:Si], [:Al, :Ti] ] #, [:C, :O, :H] ]
degrees = [ 15, 14, 13, 12, 11, 12 ]
wLs = [ 1.5, 1.4, 1.3, 1.6, 1.55, 1.62 ]

zipname = dirname(pathof(ACE))[1:end-3] * "test/models/v08_compat.zip"
zipdir = ZipFile.Writer(zipname)

@info("Exporting several test basis sets and test potentials.")
for species in species_list, N = 1:length(degrees)
   local D, fptr
   @info("   ... species=$(species), N=$(N)");
   r0 = minimum(rnn.(species))
   _params = Dict("species" => species, "maxdeg" => degrees[N], "wL" => wLs[N],
                  "r0" => r0, "rcut" => 2.5 * r0, "N" => N,
                  "degreetype" => "SparsePSHDegree")
   basis = ACE.Testing.test_basis(_params)

   # file label
   flabel = "_$(N)_" * prod(string.(species)) * ".json"

   # write basis and basis tests
   D = write_dict(basis)
   D["_params"] = _params
   D["_tests"] = ACE.Testing.createtests(basis, 3; tests = ["E"])
   fptr = ZipFile.addfile(zipdir, "rpibasis" * flabel; method = ZipFile.Deflate)
   write(fptr, JSON.json(D))

   # write potential and potential tests
   V = ACE.Random.randcombine(basis; diff=1)
   D = write_dict(V)
   D["_tests"] = ACE.Testing.createtests(V, 3)
   fptr = ZipFile.addfile(zipdir, "rpipot" * flabel; method = ZipFile.Deflate)
   write(fptr, JSON.json(D))
end

close(zipdir)


export
  # ITensors
  maxlinkdim,
  siteinds,
  state,

  # circuits/circuits.jl
  randomlayer,
  gatelayer,
  randomcircuit,
  qft,
  ghz,
  dag,

  # circuits/trottersuzuki.jl
  trottercircuit,

  # circuits/gates.jl
  # Methods
  gate,
  # Macros
  @GateName_str,
  randomparams,

  # circuits/getsamples.jl
  # Methods
  getsamples,
  fullbases,
  randombases,
  fullpreparations,
  randompreparations,

  # noise.jl
  insertnoise,

  # productstates.jl
  qubits,
  qudits,
  productstate,
  productoperator,

  # circuits/runcircuit.jl
  # Methods
  buildcircuit,
  runcircuit,
  choimatrix,

  # optimizers.jl
  optimizer,

  # lpdo.jl
  LPDO,
  normalize!,
  logtr,
  tr,

  # randomstates.jl
  # Methods
  randomstate,
  randomprocess,

  # tomography/statetomography.jl
  #           /processtomography.jl
  tomography,

  # measurements.jl
  measure,
  entanglemententropy,

  # distances.jl
  fidelity,
  fidelity_bound,
  frobenius_distance,

  # observer.jl
  Observer,
  save,
  results,

  # inputoutput.jl
  writesamples,
  readsamples,

  # utils.jl
  split_dataset,
  nqubits


JSON.lower(hf::HavingSpec) = non_nothing_dict(hf)

"""
    EqualTo(aggregation::String, value::Real)

aggregation = value filter.
"""
struct EqualTo <: HavingSpec
    type::String
    aggregation::String
    value::Real
    EqualTo(aggregation, value) = new("equalTo", aggregation, value)
end

"""
    GreaterThan(aggregation::String, value::Real)

aggregation >= value filter.
"""
struct GreaterThan <: HavingSpec
    type::String
    aggregation::String
    value::Real
    GreaterThan(aggregation, value) = new("greaterThan", aggregation, value)
end

"""
    LessThan(aggregation::String, value::Real)

aggregation <= value filter.
"""
struct LessThan <: HavingSpec
    type::String
    aggregation::String
    value::Real
    LessThan(aggregation, value) = new("lessThan", aggregation, value)
end

"""
    DimSelector(dimension::String, value)

dimension = value filter.
"""
struct DimSelector <: HavingSpec
    type::String
    dimension::String
    value
    DimSelector(dimension, value) = new("dimSelector", dimension, value)
end

"""
    AndH(havingSpecs::Vector{<:HavingSpec})

Match all `havingSpecs`.
"""
struct AndH <: HavingSpec
    type::String
    havingSpecs::Vector{<:HavingSpec}
    AndH(havingSpecs) = new("and", havingSpecs)
end

"""
    OrH(havingSpecs::Vector{<:HavingSpec})

Match at leats one of the `havingSpecs`.
"""
struct OrH <: HavingSpec
    type::String
    havingSpecs::Vector{<:HavingSpec}
    OrH(havingSpecs) = new("or", havingSpecs)
end

"""
    NotH(havingSpec::HavingSpec)

Do not match the `havingSpec`.
"""
struct NotH <: HavingSpec
    type::String
    havingSpec::HavingSpec
    NotH(havingSpec) = new("not", havingSpec)
end


function go_around(r, side)
    counter_horisontal = 0
    counter_vertical = 0
    while(isborder(r, side))
        move!(r, turn_side(side)) # turn_side(side)
        counter_vertical = counter_vertical + 1
    end
    move!(r, side)
    while(isborder(r, invert_side(turn_side(side)))) # invert_side(turn_side(side))
        move!(r, side)
        counter_horisontal = counter_horisontal + 1
    end
    move!(r, invert_side(turn_side(side)))
    c = counter_vertical
    while(c > 1)
        move!(r, invert_side(turn_side(side))) # invert_side(turn_side(side))
        c = c - 1
    end
    # counter_horisontal - в побочную сторону
    # counter_vertical - в главную сторону
    return (counter_horisontal, counter_vertical)
end
function turn_side(side)
    if side == Nord
        return West
    end
    if side == Sud
        return West
    end
    if side == Ost
        return Sud
    end
    if side == West
        return Sud
    end
end
function invert_side(side)
    if side == Nord
        return Sud
    end
    if side == Sud
        return Nord
    end
    if side == Ost
        return West
    end
    if side == West
        return Ost
    end
end
function gotobegin(r)
    a = 0
    b = 0
    while(!((isborder(r, West)) & (isborder(r, Sud))))
        a = a + gotoborder(r, West)
        b = b + gotoborder(r, Sud)
    end
    return(a, b)
end
function gotoborder(r, side)
    count = 0
    while(!isborder(r, side))
        move!(r, side)
        count = count + 1
    end
    return count
end
function measure_field(r)
    y = gotoborder(r, Nord)
    gotoborder(r, Sud)
    x = gotoborder(r, Ost)
    gotoborder(r, West)
    return (x,y)
end
function push_and_mark(r, side, counter)
    b = 0
    while(counter>0)
        putmarker!(r)
        if(isborder(r, side))
            b,count = go_around(r, side)
            counter = counter - count - 1
        else
            move!(r, side)
            counter = counter - 1
        end
    end
end
function push_hor(r, side, counter)
    b = 0
    while(counter>0)
        if(isborder(r, side))
            count,b = go_around(r, side)
            counter = counter - count - 1
        else
            move!(r, side)
            counter = counter - 1
        end
    end
end
function goto_coord(r, side, coord)
    while(coord > 0)
        coord = coord - 1
        move!(r, side)
    end
end
function main(r)
    c_hor,c_vert = gotobegin(r)
    rt_x,rt_y=c_hor,c_vert
    x_dim,y_dim = measure_field(r)
    side = Nord
    move!(r, Ost)
    d = 1
    obstacle_count = 0
    while(x_dim>0)
        d = d + 1
        x_dim = x_dim - 1
        c = y_dim
        while(c>0)
            if(isborder(r, side))
                count,b = go_around(r, side)
                c = c - count - 1
                if((b == 1) & (count*b != 0))
                    obstacle_count = obstacle_count + 1
                end
                # print(count, " ", b, "\n")
            else
                move!(r, side)
                c = c - 1
            end
        end
        gotobegin(r)
        if(x_dim != 0)
            goto_coord(r, Ost, d)
        end
    end
    obstacle_count = obstacle_count
    # print(obstacle_count)
    return obstacle_count
end


"""
    add_nodes!(g, dep::Pair; id=nextid(g))

Recursively add nodes via the dependency chain specified by `dep`.
If any intermediarry node doesn't exist, it is created.
A new path is created starting from the first node to the last one, but
if there is an existing part, it is continued (see [`nextid`](@ref)).
"""
function add_nodes!(g, dep::Pair; id=nextid(g, dep))
    if dep[2] isa Pair
        dest = add_nodes!(g, dep[2], id=id)
    else
        dest = dep[2]
        if id == get_prop(g)
            # We are on a new path
            set_prop!(g, id+1)
        end
    end
    add_path!(g, dep[1], dest, id=id)

    return dep[1]
end

"""
    add_path!(g, source, dest; id=maxid(g))

Create a path between the source node and the destination one.
If the nodes do not exist, they are created.
"""
function add_path!(g, source, dest; id=maxid(g))
    sv = haskey(g.index, source) ? g[source] : nv(g) + 1
    sv > nv(g) && add_vertex!(g, source)
    dv = haskey(g.index, dest) ? g[dest] : nv(g) + 1
    dv > nv(g) && add_vertex!(g, dest)
    if has_edge(g, sv, dv)
        set_prop!(g, sv, dv, id)
    else
        add_edge!(g, sv, dv, id)
    end
end

function endof(dep)
    if dep[2] isa Pair
        return endof(dep[2])
    else
        return dep[2]
    end
end

"""
    add_bulk!(g, dep, vals)

Add the multiple values (`vals`) of the things identified by the keys of `vals`,
with the dependency chain given by `dep`. The values of `vals` are assumed to
be _equal length_ vectors. Each added node will correspond to an element of the vectors.
Note: The dependency chain must contain all relevant information for
identifying the values.

The function returns the ids of the paths corresponding to the added values.
"""
function add_bulk!(g, dep, vals)
    dep_end = endof(dep)
    # nextid(g, dep) is expensive
    # We can campute all ids in advance with only one call (for the first)
    ids = range(nextid(g, dep), length=length(values(vals[1])), step=1)
    for i in eachindex(values(vals[1]))
        # created the path up to the nodes to be added
        add_nodes!(g, dep, id=ids[i])
        val = (v[i] for v in vals)
        # add the values
        add_nodes!(g, dep_end=>NamedTuple{keys(vals)}(val), id=ids[i])
    end
    return ids
end


#===============================================================================

                                Cloning

                This is quite inefficient, but has a nice syntax

===============================================================================#

"""
    clone(P::T, θ::AbstractDict) where T <: DiffusionProcess

Simplified cloning of diffusion law `P`. Substitute relevant parameters with
new values. `θ` must be a dict corresponding to parameters returned after a call
to `var_parameters`.
"""
function clone(P::T, θ::AbstractDict) where T <: DiffusionProcess
    cp = const_parameters(P)
    remove_curly(T)(;cp..., θ...)
end

#===============================================================================

                        in-place parameter setting

            Efficient way of re-parameterizing the diffusion law

===============================================================================#
"""
    set_parameters!(P::DiffusionProcess, θ, entries)

Set parameters of a diffusion law `P` in-place. `entries` should be a collection
of pairs Pair{Int64,Symbol} that list the relevant entries in `θ` for
reparameterization, together with the corresponding parameter names.
"""
function set_parameters!(P::DiffusionProcess, θ, entries)
    for (i,s) in entries
        setfield!(P, s, θ[i])
    end
end

function set_parameters!(P::DiffusionProcess, P°::DiffusionProcess, entries)
    for e in entries
        setfield!(P°, e, getfield(P, e))
    end
end

function set_parameters!(P::T, θ::Dict) where T <: DiffusionProcess
    for k in keys(θ)
        hasfield(T, k) && setfield!(P, k, θ[k])
    end
end

function same_entries(P::DiffusionProcess, P°::DiffusionProcess, entries)
    for e in entries
        getfield(P, e) == getfield(P°, e) || return false
    end
    true
end

function has_any(P::T, entries) where T <: DiffusionProcess
    for (i,s) in entries
        hasfield(T, s) && return true
    end
    false
end


"""
The **pcl_octree** library provides efficient methods for creating a
hierarchical tree data structure from point cloud data.

http://docs.pointclouds.org/trunk/group__octree.html

## Exports

$(EXPORTS)
"""
module PCLOctree

using DocStringExtensions
using LibPCL
using PCLCommon
using Cxx

const libpcl_octree = LibPCL.find_library_e("libpcl_octree")
try
    Libdl.dlopen(libpcl_octree, Libdl.RTLD_GLOBAL)
catch e
    warn("You might need to set DYLD_LIBRARY_PATH to load dependencies proeprty.")
    rethrow(e)
end

end # module


#= Read and write plain-text edgelists

This source code borrows the general parsing scheme from `GraphIO.jl`. =#

const DELIM_REGEX = raw"\s,;|"

function read_edgelist_ids!(io::IO, g::Graph, delim::String, start_id::Integer)
    re = Regex(raw"(\d+)[" * delim * raw"]+(\d+)")
    Δid = 1 - start_id
    for line in eachline(io)
        line = strip(line)
        if length(line) > 0 && line[1] != '#'
            m = match(re, line).captures
            m == nothing && return line
            src = parse(Int, m[1]) + Δid
            tgt = parse(Int, m[2]) + Δid
            while nodecount(g) < max(src, tgt)
                addnode!(g)
            end
            addedge!(g, src, tgt)
        end
    end
end

function read_edgelist_strings!(io::IO, g::Graph, delim::String)
    re = Regex(raw"(\w+)[" * delim * raw"]+(\w+)")
    addnodeattr!(g, :id, String[])
    id_map = Dict{String,Int}()
    create = (x) -> id_map[x] = addnode!(g, id=String(x))
    for line in eachline(io)
        line = strip(line)
        if length(line) > 0 && line[1] != '#'
            m = match(re, line).captures
            m == nothing && return line
            src, tgt = m
            haskey(id_map, src) || create(src)
            haskey(id_map, tgt) || create(tgt)
            addedge!(g, id_map[src], id_map[tgt])
        end
    end
end

"""
    read_edgelist(io::IO[; delim, kvargs...])

Read graph from a plain text file where each line specifies a single edge.
Lines should have the format `source<delims...>target`.
Set the possible delimiters as a string `delim`.

Keyword argument `mode` specifies how to treat the parsed node names:
* `:auto` — try both of the following two options (default),
* `:int` — try to parse numerical ids with the offset `start_id`,
* `:str` — treat node ids as strings (slower and the ordering of resulting
  node ids can be arbitrary).

If `:int` mode is chosen, it is assumed that node ids in the file start
with `start_id`, which is zero by default.
If `:str` mode is chosen, the resulting graph has a `String` attribute `:id`
associated with its nodes.
"""
function read_edgelist(io::IO; delim=DELIM_REGEX, mode=:auto, start_id=0, kwargs...)
    if mode == :int || mode == :auto
        pos = position(io)
        g = Graph(; kwargs...)
        res = read_edgelist_ids!(io, g, delim, start_id)
        if res == nothing
            return g
        elseif mode == :int
            error("can't parse line\"$res\"")
        end
        seek(io, pos)
    end
    if mode == :str || mode == :auto
        g = Graph(; kwargs...)
        res = read_edgelist_strings!(io, g, delim)
        if res == nothing
            return g
        else
            error("can't parse line\"$res\"")
        end
    end
end

"""
    write_edgelist(io::IO, g::Graph[; delim, ids])

Write the list of graph's edges, one per line.
If `ids` is not customized, default node ids are used.
"""
function write_edgelist(io::IO, g::Graph; delim='\t', ids=nodes(g) - 1)
    for e = edges(g)
        println(io, repr(ids[e.source]), delim, repr(ids[e.target]))
    end
end


import HierarchicalUtils: NodeType, LeafNode, InnerNode, noderepr, children

NodeType(::Type{<:Union{Missing, ArrayNode, ArrayModel, LazyNode}}) = LeafNode()
NodeType(::Type{<:AbstractNode}) = InnerNode()
NodeType(::Type{<:AbstractMillModel}) = InnerNode()
NodeType(::Type{<:LazyModel}) = InnerNode()

noderepr(::T) where T <: Union{AbstractNode, AbstractMillModel} = "$(Base.nameof(T))"
noderepr(::Missing) = "∅"
noderepr(n::ArrayNode) = "ArrayNode$(size(n.data))"
noderepr(n::ArrayModel) = "ArrayModel($(n.m))"
noderepr(n::BagNode) = if ismissing(n.data) || nobs(n.data) == 0
    "BagNode with $(length(n.bags)) empty bag(s)"
else
    "BagNode with $(length(n.bags)) bag(s)"
end
noderepr(n::BagModel) = "BagModel ↦ $(repr("text/plain", n.a)) ↦ $(repr("text/plain", n.bm))"
noderepr(n::WeightedBagNode) = "WeightedNode with $(length(n.bags)) bag(s) and weights Σw = $(sum(n.weights))"
noderepr(n::AbstractProductNode) = "ProductNode"
noderepr(n::ProductModel) = "ProductModel ↦ $(noderepr(n.m))"
noderepr(n::MissingNode) = "Missing"
noderepr(n::MissingModel) = "Missing"
noderepr(n::LazyNode{N,D}) where {N,D} = (data_len = length(n.data); "$N $data_len item$(data_len > 1 ? "s" : "")")
noderepr(n::LazyNode{N,D}) where {N,D<:Nothing} = "$(N) ∅"
noderepr(n::LazyModel{Name}) where {Name} = "Lazy$(Name)Model"

children(n::AbstractBagNode) = (n.data,)
children(n::BagModel) = (n.im,)
children(n::ProductNode) = n.data
children(n::ProductModel) = n.ms
children(n::MissingNode) = (n.data,)
children(n::MissingModel) = (n.data,)
children(n::LazyModel) = (n.m,)


function strip_maybe(T)
    T === Missing && return Union{}
    !isa(T, Union) && return T
    return Union{strip_maybe(T.a), strip_maybe(T.b)}
end

optional_eltype(::Type{Union{T, Void}}) where {T} = T
optional_eltype(_) = Union{}

macro from_json(T, renamings)
    syms = Symbol[]
    body = Expr[]
    for (N,fT) in zip(fieldnames(T), T.types)
        sym = gensym()
        push!(syms, sym)
        name = string(N)
        if haskey(renamings, N)
            val = renamings[N]
            if isa(val, String)
                name = val
            else
                @assert isa(val, Expr)
                push!(body, :($sym = $val))
                continue
            end
        end
        sfT = strip_maybe(fT)
        oeT = optional_eltype(sfT)
        if oeT !== Union{}
            sfT = oeT
        end
        if sfT <: QboRef
            name *= "Ref"
            op = :($sym = from_json($sfT, data[$(name)]))
        elseif sfT <: Vector
            op = :($sym = map(x->from_json($(eltype(sfT)), x),data[$(name)]))
        else
            op = :($sym = from_json($sfT, data[$(name)]))
        end
        if oeT !== Union{}
            push!(body, quote
                if haskey(data, $(name))
                    $op
                else
                    $sym = nothing
                end
            end)
        else
            push!(body, op)
        end
    end
    esc(quote
        $(body...)
        $T($(syms...))
    end)
end

macro to_json(T, renamings)
    body = Expr[:(ret = Dict{String, Any}())]
    for (N,fT) in zip(fieldnames(T), T.types)
        name = string(N)
        if haskey(renamings, N)
            val = renamings[N]
            if isa(val, String)
                name = val
            else
                @assert isa(val, Expr)
                push!(body, val)
                continue
            end
        end
        sfT = strip_maybe(fT)
        oeT = optional_eltype(sfT)
        if oeT !== Union{}
            sfT = oeT
        end
        if sfT <: QboRef
            name *= "Ref"
            op = :(ret[$name] = to_json($sfT, getfield(data, $(quot(N)))))
        elseif sfT <: Vector
            op = :(ret[$name] = map(to_json,getfield(data, $(quot(N)))))
        else
            op = :(ret[$name] = to_json(getfield(data, $(quot(N)))))
        end
        push!(body, quote
            if !ismissing(getfield(data, $(quot(N)))) && getfield(data, $(quot(N))) !== nothing
                $op
            end
        end)
    end
    a = esc(quote
        $(body...)
        ret
    end)
end

from_json(::Type{Int}, s::String) = parse(Int, s)
from_json(::Type{String}, s::String) = s
from_json(::Type{Date}, s::String) = Date(s)
from_json(::Type{Decimal}, d::Decimal) = d
from_json(::Type{Bool}, b::Bool) = b

to_json(s::String) = s
to_json(d::Date) = repr(d)
to_json(d::Decimal) = d


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

@recipe function f(data::Data, vars=nothing)
  # valid variables
  validvars = Tables.schema(values(data)).names

  # plot all variables by default
  isnothing(vars) && (vars = validvars)
  @assert vars ⊆ validvars "invalid variable names"

  # underlying domain where variables live
  dom = domain(data)

  # shared plot specs
  layout --> length(vars)

  for (i, var) in enumerate(vars)
    # data for variable
    vals = data[var]

    # handle categorical values
    v = eltype(vals) <: CategoricalValue ? levelcode.(vals) : copy(vals)

    # handle missing values
    v = replace(v, missing => NaN)

    @series begin
      subplot := i
      title --> string(var)
      legend --> false
      dom, v
    end
  end
end


push!(LOAD_PATH, "./")

using Printf
using Statistics
using stiffness

const S = StiffnessData()
const x = Float64[0., 1., 0.]
const y = Float64[0., 0., 1.]
m = zeros(21)

nb_runs = 200

times = zeros(nb_runs)

for irun in 1:nb_runs
    times[irun] = @elapsed stiffness.op!(S, x, y, m)
end

@printf("%.3f µs\n", median(times) * 1e6)


function random(numbits)
  numbytes = fld(numbits, 8)
  if numbits % 8 != 0
    warn("The requested number of random bits is not divisible by 8. Rounding to the nearest multiple of 8.")
    numbytes += 1
  end

  ret = zeros(Uint8, numbytes)
  ccall((:RAND_bytes, "libcrypto"), Void, (Ptr{Uint8}, Int), ret, numbytes)

  return ret
end


include("../src/Hash.jl")
include("testing.jl")

fname = tempname()
test_strings = ["foo", "foobar", "baz", "this is a test", ""]
function check_hash(algorithm, fn)
    print("checking " * algorithm * " sums ")
    for value in test_strings
        f = open(fname, "w")
        write(f, value)
        close(f)

        command_line = strip(split(readall(`openssl dgst -$algorithm $fname`), " ")[2])
        library_with_passed_vaue = Hash.hex(fn(value))

        @assert_equal(command_line, library_with_passed_vaue)

        f = open(fname, "r")
        library_from_file = Hash.hex(fn(f))
        close(f)

        @assert_equal(command_line, library_from_file)
    end
    println()
end

check_hash("md4", Hash.md4)
check_hash("md5", Hash.md5)
check_hash("mdc2", Hash.mdc2)
check_hash("sha1", Hash.sha1)
check_hash("sha224", Hash.sha224)
check_hash("sha256", Hash.sha256)
check_hash("sha384", Hash.sha384)
check_hash("sha512", Hash.sha512)
println()


using Test
using LogicCircuits

module TestNodes

    using Test
    using ..LogicCircuits

    mutable struct TestINode <: Dag
        id::Int
        children::Vector{Dag}
        data
        counter::UInt32
        TestINode(i,c) = new(i,c,nothing,false)
    end

    mutable struct TestLNode <: Dag
        id::Int
        data
        counter::UInt32
        TestLNode(i) = new(i,nothing,false)
    end

    LogicCircuits.NodeType(::Type{<:TestINode}) = Inner()
    LogicCircuits.NodeType(::Type{<:TestLNode}) = Leaf()
    LogicCircuits.children(n::TestINode) = n.children

    @testset "Graphs utils for TestNodes" begin

        l1 = TestLNode(1)
        l2 = TestLNode(2)

        @test !has_children(l1)
        @test num_children(l1) == 0
        @test isleaf(l1)
        @test !isinner(l1)

        i1 = TestINode(1,[l1])
        i2 = TestINode(2,[l2])
        i12 = TestINode(3,[l1,l2])

        @test has_children(i1)
        @test has_children(i12)

        @test num_children(i1) == 1
        @test num_children(i12) == 2

        j1 = TestINode(1,[i1,i12])
        j2 = TestINode(2,[i2])
        j12 = TestINode(3,[i1,i2])

        r = TestINode(5,[j1,j2,j12])

        @test has_children(r)
        @test num_children(r) == 3

        reset_counter(r,5)
        @test r.counter == 5
        @test l1.counter == 5
        @test i12.counter == 5

        reset_counter(r)
        @test r.counter == 0
        @test l1.counter == 0
        @test i12.counter == 0

        foreach(r) do n
            n.id += 1
        end
        @test l1.id == 2
        @test l2.id == 3
        @test i12.id == 4
        @test j2.id == 3
        @test r.id == 6
        @test r.counter == 0
        @test l1.counter == 0
        @test i12.counter == 0

        foreach(r, l -> l.id += 1, i -> i.id -= 1)
        @test l1.id == 2+1
        @test l2.id == 3+1
        @test i12.id == 4-1
        @test j2.id == 3-1
        @test r.id == 6-1
        @test r.counter == 0
        @test l1.counter == 0
        @test i12.counter == 0

        @test filter(n -> iseven(n.id), r) == [l2,i2,j2]

        lastvisited = nothing
        foreach(n -> lastvisited=n,r)
        @test lastvisited === r

        lastvisited = nothing
        foreach_down(n -> lastvisited=n,r)
        @test isleaf(lastvisited)

        @test num_nodes(r) == 9
        @test num_edges(r) == 12

        @test isempty(inodes(l1))
        @test leafnodes(l1) == [l1]

        @test issetequal(inodes(r), [i1,i2,i12,j1,j2,j12,r])
        @test issetequal(innernodes(r), [i1,i2,i12,j1,j2,j12,r])
        @test issetequal(leafnodes(r), [l1,l2])

        @test tree_num_edges(r) == 14 # unverified

        @test linearize(r)[end] == r
        @test linearize(r)[1] == l1 || linearize(r)[1] == l2
        @test linearize(l2) == [l2]
        @test length(linearize(i12)) == 3

        @test eltype(linearize(r)) == Dag
        @test eltype(linearize(l1)) == TestLNode
        @test eltype(linearize(r, Any)) == Any

        @test left_most_descendent(r) == l1
        @test right_most_descendent(r) == l2

    end

end

@testset "Node stats tests" begin

    lc = load_logic_circuit(zoo_lc_file("little_4var.circuit"));

    lstats = leaf_stats(lc);
    istats = inode_stats(lc);
    nstats = node_stats(lc);

    @test !(PlainTrueNode in keys(lstats));
    @test !(PlainFalseNode in keys(lstats));
    @test lstats[PlainLiteralNode] == 8;

    @test istats[(Plain⋀Node, 2)] == 9;
    @test istats[(Plain⋁Node, 1)] == 10;
    @test istats[(Plain⋁Node, 4)] == 2;

    for t in keys(lstats)
        @test lstats[t] == nstats[t]
    end
    for t in keys(istats)
        @test istats[t] == nstats[t]
    end

    @test num_nodes(lc) == 29
    @test num_variables(lc) == 4;
    @test num_edges(lc) == 36;

end


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

"""
    fit(estimator, data)

Fit estimator to geospatial `data` and return a fitted estimator.
"""
function fit end

"""
    predict(estimator, var, uₒ)

Predict variable `var` with the `estimator` at point or geometry `uₒ`.
"""
function predict end

"""
    status(fitted)

Return the status of the `fitted` estimator. (e.g. the
factorization of the linear system was successful)
"""
function status end


module Mathematica

# This package is based on MathLink
import MathLink
import MathLink: WExpr, WSymbol
import IterTools: takewhile

export weval, @importsymbol

include("types.jl")
include("operators.jl")
include("show.jl")
include("eval.jl")

end # module


using Overseer: Indices, @indices_in, @reverse_indices_in

@testset "Construction, collect" begin
    data_in = (1,5,100)
    s = Indices(data_in)
    data_out = collect(s)
    @test all(map(d->in(d,data_out), data_in))
    @test length(data_out) == length(data_in)
end

@testset "eltype, empty" begin
    @test eltype(Indices()) == Int
    @test eltype(typeof(Indices())) == Int
    @test isequal(empty(Indices([1,2,3])), Indices())
end

@testset "Core Functionality" begin
    s = Indices([1,2,10,20,200,300,1000,10000,10002])
    @test last(s) == 10002
    @test first(s) == 1
    @test length(s) == 9
    @test pop!(s) == 10002
    @test length(s) == 8
    @test popfirst!(s) == 1
    @test length(s) == 7
    @test !in(1,s)
    @test !in(10002,s)
    @test in(10000,s)
    @test_throws ArgumentError first(Indices())
    @test_throws ArgumentError last(Indices())
    t = copy(s)
    s = Indices()
    push!(s, 1, 2, 100)
    @test 1 in s
    @test !(3 in s)
    @test 2 in s
    @test 100 in s
    @test !(101 in s)
    @test !(1000 in s)
    @test first(s) == 1
    @test last(s) == 100
    @test s == Indices([1, 2, 100])
    push!(s, 1000)
    @test [i for i in s] == [1, 2, 100, 1000]
    @test pop!(s) == 1000
    @test s == Indices([1, 2, 100])
    push!(s, 5000)
    push!(s, 2000)
    pop!(s, 5000)
    @test s.reverse[end] === Overseer.NULL_INT_PAGE
    b = 1:1000
    s = Indices(b)
    @test collect(s) == collect(b)
    @test length(s) == length(b)
    @test pop!(s, 100) == 100
    @test_throws BoundsError pop!(s, 100)
    @test pop!(s, 100, 1) == 1
    @test pop!(s, 99, 1) == 99
    @test !in(500000, s)
    @test !in(99, s)
end

@testset "setdiff / symdiff" begin
    @test setdiff(Indices([1, 2, 3, 4]), Indices([2, 4, 5, 6])) == Indices([1, 3])
end

@testset "setdiff!" begin
    s2 = Indices([1, 2, 3, 4])
    setdiff!(s2, Indices([2, 4, 5, 6]))

    @test s2 == Indices([1, 3])
end

@testset "issue #7851" begin
    @test_throws DomainError Indices(-1)
    @test !(-1 in Indices(1:10))
end
@testset "Copy, copy!, empty" begin
    s1 = Indices([1,2,3])
    s2 = empty(s1)
    push!(s2, 10000)
    @test !in(10000, s1)
    copy!(s2, s1)
    @test !in(10000, s2)
    push!(s2, 10000)
    @test !in(10000, s1)
    s3 = copy(s2)
    push!(s3, 1000)
    @test !in(1000, s2)
    pop!(s3, 1000)
    pop!(s2, 10000)
    @test in(10000, s3)
    pop!(s3, 10000)
    @test s3 == s2 == s1
    @test collect(s3) == collect(s2) == [1,2,3]

end

@testset "Push, union" begin
    # Push, union
    s1 = Indices()
    @test_throws DomainError push!(s1, -1)
    push!(s1, 1, 10, 100, 1000)
    @test collect(s1) == [1, 10, 100, 1000]
    push!(s1, 606)
    @test collect(s1) == [1, 10, 100, 1000, 606]
    s2 = Indices()
    @test s2 === union!(s2, s1)
    s3 = Indices([1, 10, 100])
    union!(s3, [1, 606, 1000])
    s4 = union(Indices([1, 100, 1000]), Indices([10, 100, 606]))
    @test s1 == s2 == s3 == s4
end

@testset "pop!, delete!" begin
    s = Indices(1:2:10)
    @test pop!(s, 1) == 1
    @test !(1 in s)
    @test_throws BoundsError pop!(s, 1)
    @test_throws ArgumentError pop!(s, -1)
    @test_throws ArgumentError pop!(s, -1, 1)
    @test pop!(s, 1, 0) == 0
    is = copy(s.packed)
    for i in is; pop!(s, i); end
    @test isempty(s)
    push!(s, 1:2:10...)
    @test pop!(s) == 9
    @test pop!(s) == 7
    @test popfirst!(s) == 1
    @test collect(s) == [5, 3]
    empty!(s)
    @test isempty(s)
end

@testset "Intersect" begin
    @test isempty(intersect(Indices()))
    @test isempty(intersect(Indices(1:10), Indices()))
    @test isempty(intersect(Indices(), Indices(1:10)))

    @test intersect(Indices([1,2,3])) == Indices([1,2,3])

    @test intersect(Indices(1:7), Indices(3:10)) ==
          intersect(Indices(3:10), Indices(1:7)) == Indices(3:7)

    @test intersect!(Indices(1:10), Indices(1:4), 1:5, [1,2,10]) == Indices(1:2)
end

@testset "Setdiff" begin
    s1 = Indices(1:100)
    setdiff!(s1, Indices(1:2:100))
    s2 = setdiff(Indices(1:100), Indices(1:2:100))
    @test s1 == s2 == Indices(2:2:100)

    s1 = Indices(1:10)
    s2 = Indices([1:2; 6:100])
    @test setdiff(s1, s2) == setdiff(s1, [1:2; 6:100]) == Indices(3:5)
end

@testset "Subsets, equality" begin
    @test Indices(2:2:10) < Indices(1:10)
    @test !(Indices(2:2:10) < Indices(2:2:10))
    @test Indices(2:2:10) <= Indices(2:10)
    @test Indices(2:2:10) <= Indices(2:2:10)
end

@testset "indices_in" begin
    a = Indices(2:2:10)
    b = Indices(10:3:20)
    c = [Indices(3:10)]

    d = Indices(1)
    t = 0
    it = @indices_in(((a && c[1]) || d) && !b)
    for e in it
        t+=e
    end
    @test t == 4+6+8+1
    @test length(it) == 4
    t = 0
    it = @reverse_indices_in(((a && c[1]) || d) && !b)
    for e in it
        t+=e
    end
    @test t == 4+6+8+1
    @test length(it) == 4
    @test collect(it) == [1, 8, 6, 4]
end


#################################################################
#
#    Main file of MCMC.jl module
#
#################################################################

module MCMC

using Distributions
using StatsBase

import Base.*, Base.show
export show, *
export MCMCTask, MCMCChain, MCMCLikModel, MCMCSampler, MCMCTuner, MCMCRunner

# Abstract types
abstract Model
abstract MCMCModel <: Model
abstract MCMCSampler
abstract MCMCTuner
abstract MCMCRunner

typealias FunctionOrNothing Union(Function, Nothing)
typealias MatrixF64OrNothing Union(Matrix{Float64}, Nothing)
typealias F64OrVectorF64 Union(Float64, Vector{Float64})

###########  model DSL parsing and translation  ##########
include("parsers/modelparser.jl")

###########  Models  ##########
include("modellers/models.jl")
include("modellers/MCMCLikModel.jl")

### MCMCTask type, generated by combining a MCMCModel, a MCMCSampler and a MCMCRunner
type MCMCTask
  task::Task
  model::MCMCModel
  sampler::MCMCSampler
  runner::MCMCRunner
end
reset(t::MCMCTask, x) = t.task.storage[:reset](x)

#############  samplers  ########################
include("samplers/samplers.jl")  # Common definitions for samplers
include("samplers/ARS.jl")    # Acceptance Rejection sampler
include("samplers/IMH.jl")    # Independent Metropolis-Hastings sampler
include("samplers/RWM.jl")    # Random-walk Metropolis sampler
include("samplers/RAM.jl")    # Robust adaptive Metropolis sampler
include("samplers/MALA.jl")   # Metropolis adjusted Langevin algorithm sampler
include("samplers/HMC.jl")    # Hamiltonian Monte-Carlo sampler
include("samplers/HMCDA.jl")    # Adaptive Hamiltonian Monte-Carlo sampler with dual averaging
include("samplers/NUTS.jl")   # No U-Turn Hamiltonian Monte-Carlo sampler
include("samplers/SMMALA.jl") # Simplified manifold Metropolis adjusted Langevin algorithm sampler
include("samplers/PMALA.jl")  # Position-dependent Metropolis adjusted Langevin algorithm sampler
include("samplers/RMHMC.jl")  # Riemannian manifold Hamiltonian Monte Carlo sampler
include("samplers/SliceSampler.jl")  # Slice sampler

### MCMCChain, the result of running a MCMCTask
type MCMCChain
  range::Range{Int}
  samples::Matrix{Float64}
  logtargets::Vector{Float64}
  gradients::MatrixF64OrNothing
  diagnostics::Dict
  task::Union(MCMCTask, Array{MCMCTask})
  runTime::Float64

  function MCMCChain(r::Range{Int},
                     s::Matrix{Float64},
                     l::Vector{Float64},
                     g::MatrixF64OrNothing,
                     d::Dict,
                     t::Union(MCMCTask, Array{MCMCTask}),
                     rt::Float64)
    @assert size(s,1) == size(l,1) "samples and logtargets do not have the same size"
    if g != nothing
      @assert size(s) == size(g) "samples and gradients must have the same number of rows and columns"
    end
    new(r, s, l, g, d, t, rt)
  end
end

MCMCChain(r::Range{Int}, s::Matrix{Float64}, l::Vector{Float64},
          d::Dict, t::Union(MCMCTask, Array{MCMCTask}), rt::Float64) =
	MCMCChain(r, s, l, nothing,      d, t,    rt)

MCMCChain(r::Range{Int}, s::Matrix{Float64}, l::Vector{Float64},
          t::Union(MCMCTask, Array{MCMCTask}), rt::Float64) =
	MCMCChain(r, s, l, nothing, Dict(), t,    rt)

MCMCChain(r::Range{Int}, s::Matrix{Float64}, l::Vector{Float64},
          d::Dict, t::Union(MCMCTask, Array{MCMCTask})) =
	MCMCChain(r, s, l, nothing,      d, t,   NaN)

MCMCChain(r::Range{Int}, s::Matrix{Float64}, l::Vector{Float64},
          t::Union(MCMCTask, Array{MCMCTask})) =
	MCMCChain(r, s, l, nothing, Dict(), t,   NaN)

function show(io::IO, res::MCMCChain)
  nsamples, npars = size(res.samples)
  println(io, "$npars parameters, $nsamples samples (per parameter), $(round(res.runTime, 1)) sec.")
end

#  Definition of * as a shortcut operator for (model, sampler, runner) combination
*{M<:MCMCModel, S<:MCMCSampler, R<:MCMCRunner}(m::M, s::S, r::R) = spinTask(m, s, r)
*{M<:MCMCModel, S<:MCMCSampler, R<:MCMCRunner}(m::Array{M}, s::S, r::R) = map((me) -> spinTask(me, s, r), m)
*{M<:MCMCModel, S<:MCMCSampler, R<:MCMCRunner}(m::M, s::Array{S}, r::R) = map((se) -> spinTask(m, se, r), s)
*{M<:MCMCModel, S<:MCMCSampler, R<:MCMCRunner}(m::M, s::S, r::Array{R}) = map((re) -> spinTask(m, s, re), r)
*{M<:MCMCModel, S<:MCMCSampler, R<:MCMCRunner}(m::Array{M}, s::Array{S}, r::R) =
  map((me, se) -> spinTask(me, se, r), m, s)
*{M<:MCMCModel, S<:MCMCSampler, R<:MCMCRunner}(m::Array{M}, s::S, r::Array{R}) =
  map((me, re) -> spinTask(me, s, re), m, r)
*{M<:MCMCModel, S<:MCMCSampler, R<:MCMCRunner}(m::M, s::Array{S}, r::Array{R}) =
  map((se, re) -> spinTask(m, se, re), s, r)
*{M<:MCMCModel, S<:MCMCSampler, R<:MCMCRunner}(m::Array{M}, s::Array{S}, r::Array{R}) =
  map((me, se, re) -> spinTask(me, se, re), m, s, r)

#############  runners    ########################
include("runners/runners.jl")
include("runners/SerialMC.jl") # Ordinary serial MCMC runner
include("runners/SerialTempMC.jl") # Serial Tempering Monte-Carlo runner
include("runners/SeqMC.jl") # Sequential Monte-Carlo runner

#############  MCMC output analysis and diagnostics    ########################
include("stats/mean.jl") # MCMC mean estimators
include("stats/var.jl") # MCMC variance estimators
include("stats/ess.jl") # Effective sample size and integrated autocorrelation time functions
include("stats/summary.jl") # Summary statistics for MCMCChain
include("stats/zv.jl")  # ZV-MCMC estimators
end


using TableTraits
using NamedTuples
import IterableTables

@testset "table tests" begin
  let x = K_Table(a=[1, 2])
    @test x == K_Table(Any[[1,2]], [:a])
    @test xt(kpointer(x)) == XT
    @test ncol(x) == 1
    @test nrow(x) == 2
    @test Q.colnames(x) == [:a]
    @test DataFrames.index(x) == DataFrames.Index(Dict(:a=>1), Symbol[:a])
    @test x[1] == [1, 2]
    @test x[2,1] == 2
    @test (x′ = K(r1(kpointer(x))); x == x′)
    @test begin
      t = K_Table(a=[0, 1, 2], b=[0., 10., 20.])
      show_to_string(t) == """
      3×2 Q.K_Table
      │ Row │ a │ b    │
      ├─────┼───┼──────┤
      │ 1   │ 0 │ 0.0  │
      │ 2   │ 1 │ 10.0 │
      │ 3   │ 2 │ 20.0 │"""
    end
    @test (x = DataFrame(a=1:3, b=[10., 20, 30]); x == K(x));
    @test (x = K_Table(a=Int[]); names(names!(x, [:b])) == [:b])
    @test (x = K_Table(a=Int[]); x[:a] == Int[])
  end
  @test (x = K_Table(@NT(a::Int, b::Float64), 3); size(x) == (3, 2))
  @test begin
    t = K_Table(a=[1])
    Q.coldata(t)[1] = push!(t[1], 2)
    Q.coldata(t)[1] = push!(t[1], 3)
    t[1] == [1, 2, 3]
  end
end

@testset "table iterator tests" begin
  let t = K_Table(a=[0, 1, 2], b=[0., 10., 20.]), ti = getiterator(t)
    @test isiterable(t)
    @test isiterabletable(t)
    @test length(ti) == 3
    @test eltype(ti) == typeof(@NT(a=0, b=0.))
    @test start(ti) == 1
    @test collect(ti) == [
      @NT(a=0, b=0.),
      @NT(a=1, b=10.),
      @NT(a=2, b=20.),
    ]
  end
  @test K_Table([@NT(a=1), @NT(a=2)]) == DataFrame(a=[1,2])
  @test K_Table(@NT(a=i) for i in 1:3 if i != 2) == DataFrame(a=[1,3])
end


"""
    module Simulate_Billiards_OneBall

Simulates a billiard ball on a (non-standard) square billiard table
that is hitting the cushing several times. This model is used to
test the collision with the cushing.
"""
module Simulate_Billiards_OneBall

using  Modia3D
import Modia3D.ModiaMath

vmatBalls = Material(color="Red"      , transparency=0.0)
vmatTable = Material(color="DarkGreen", transparency=0.4)

diameter = 0.06
radius   = diameter/2

TableX = 1.24
TableY = 1.12
TableZ = 0.03
LyBox  = 0.03
hz     = 7*radius/5
hzz    = hz/3

@assembly Cushion(world) begin
  cushionPart11 = Solid(SolidBox(LyBox, TableY, hz) , "BilliardCushion", vmatTable)
  cushionPart12 = Solid(SolidBox(1.5*LyBox, TableY, hzz,rsmall=0.0) , "BilliardCushion", vmatTable; contactMaterial="BilliardCushion")
  cushionPart21 = Solid(SolidBox(TableX, LyBox , hz) , "BilliardCushion", vmatTable)
  cushionPart22 = Solid(SolidBox(TableX, 1.5*LyBox, hzz,rsmall=0.0) , "BilliardCushion", vmatTable; contactMaterial="BilliardCushion")

  cushion11 = Object3D(world    , cushionPart11, fixed=true, r=[TableX/2 - LyBox/2 , 0.0                 , hz/2])
  cushion12 = Object3D(cushion11, cushionPart12, fixed=true, r=[-LyBox/4           , 0.0                 , hz/2+hzz/2])
  cushion31 = Object3D(world    , cushionPart11, fixed=true, r=[-TableX/2 + LyBox/2, 0.0                 , hz/2])
  cushion32 = Object3D(cushion31, cushionPart12, fixed=true, r=[LyBox/4            , 0.0                 , hz/2+hzz/2])
  cushion21 = Object3D(world    , cushionPart21, fixed=true, r=[0.0                , TableY/2 - LyBox/2  , hz/2])
  cushion22 = Object3D(cushion21, cushionPart22, fixed=true, r=[0.0                , -LyBox/4            , hz/2+hzz/2])
  cushion41 = Object3D(world    , cushionPart21, fixed=true, r=[0.0                , -TableY/2 + LyBox/2 , hz/2])
  cushion42 = Object3D(cushion21, cushionPart22, fixed=true, r=[0.0                , LyBox/4             , hz/2+hzz/2])
end

distance_balls = sqrt(3)/2*diameter
dist           = 0.0001
solidBall      = Solid(SolidSphere(diameter), "BilliardBall", vmatBalls ; contactMaterial="BilliardBall")
solidTable     = Solid(SolidBox(TableX, TableY, TableZ), "BilliardTable", vmatTable; contactMaterial="BilliardTable")

@assembly Billiards_OneBall() begin
  world   = Object3D()
  cushion = Cushion(world)
  table   = Object3D(world, solidTable, fixed=true , r=[0.0, 0.0, -TableZ/2])
  ball    = Object3D(world, solidBall , fixed=false, r=[0.0, 0.0, diameter/2], v_start=[6.0, 3.0, 0.0], visualizeFrame=true)
end

gravField         = UniformGravityField(g=9.81, n=[0,0,-1])
billiards_OneBall = Billiards_OneBall(sceneOptions=SceneOptions(gravityField=gravField, visualizeFrames=false,
                                                                defaultFrameLength=0.1, enableContactDetection=true))

# Modia3D.visualizeAssembly!( billiards_OneBall )
model = SimulationModel( billiards_OneBall, maxNumberOfSteps=1000 )   # default maxNumberOfSteps gives an error
# ModiaMath.print_ModelVariables(model)
result = ModiaMath.simulate!(model; stopTime=0.5, tolerance=1e-8, log=false)

ModiaMath.plot(result, [("ball.r[1]", "ball.r[2]"),
                        ("ball.v[1]", "ball.v[2]"),
                        ("ball.w[1]", "ball.w[2]")])

println("... success of examples/collisions/Simulate_Billiards_OneBall!")
end


# This file contains
# a type for input data

"""
Module containing the type to interact with input data.
"""
module Input

# Export
export InputStruct

# Load extras
using ..Extras

# Arrays with non-standard indexing
using OffsetArrays

# Precisions and formats of numbers
using ..Prec

"""
    InputStruct()

Instantiate this type to interact with input data.

# Data
- `N::`[`IT`](@ref)`=0`: sample size;
- `Δt::`[`RT`](@ref)`=0`: sample step;
- `q::`[`RT`](@ref)`=0`: significance level;
- `t::Vector{`[`RT`](@ref)`}=[]`: time array;
- `x::Vector{`[`RT`](@ref)`}=[]`: values array.

# Methods
- [`read!`](@ref)`(file::AbstractString)`: read input data from a file;
- [`write`](@ref)`(file::AbstractString)`: write input data to a file;
- [`example`](@ref)`(file::AbstractString)`: generate an example of the input/output file;
- [`reset!`](@ref)`()`: reset an instance to default values.

# Note
Data can be also read when calling an instance like so:
```jldoctest; output = false
using Scats
s = Scats.API()
file, _ = mktemp()
s.Input.example(file)
s.Input(file)

# output

```
"""
mutable struct InputStruct

    # Data
    N::IT         # Sample size
    Δt::RT        # Sample step
    q::RT         # significance level
    t::Vector{RT} # Time array
    x::Vector{RT} # Values array

    # Methods
    read!::Function   # Read input data from a file
    write::Function   # Write input data to a file
    example::Function # Generate an example of the input/output file
    reset!::Function  # Reset an instance to default values

    # Construct an object of this type
    function InputStruct()
        this = new(0, 0, 0, [], [])
        this.read! = function (file::AbstractString) read!(this, file) end
        this.write = function (file::AbstractString) write(this, file) end
        this.example = example
        this.reset! = function () reset!(this) end
        return this
    end

    # Read the data when calling an instance
    function (input::InputStruct)(file::AbstractString)
        input.read!(file)
    end

end

# Sources
include("input_exceptions.jl") # Exceptions
include("input_read.jl")       # Read input data from a file
include("input_write.jl")      # Write input data to a file
include("input_example.jl")    # Generate an example of the input/output file
include("input_reset.jl")      # Reset an instance to default values

end