# export Beta, Bernoulli, Normal

struct Beta{A, B} <: PrimRandVar
  α::A
  β::B
  id::ID
  Beta(α::A, β::B, id = uid()) where {A, B} = new{A, B}(α, β, id)
end
@inline (rv::Beta)(ω::Ω) = apl(rv, ω)
rvtransform(rv::Beta) = betarv
betarv(p::Real, α::Real, β::Real) = quantile(Djl.Beta(α, β), p)
betarv(ω::Ω, α::Real, β::Real) = betarv(rand(ω), α, β)
betarv(ω::Ω, α, β) = (p = rand(ω, anysize(α, β)); betarv.(p, α, β))

betarv(α, β) = Beta(α, β)
betarv(α, β, sz::Dims) = Beta(lift(fill)(α, sz), lift(fill)(β, sz))
const β = betarv

"Bernoulli distribution"
struct Bernoulli{T, A} <: PrimRandVar
  p::A
  id::ID
  Bernoulli{T}(p::A, id = uid()) where {T, A} = new{T, A}(p, id)
end
@inline (rv::Bernoulli)(ω::Ω) = apl(rv, ω)
rvtransform(rv::Bernoulli) = bernoulli
bernoulli(p_::Real, p::Real, T::Type{RT} = Int) where RT =
  RT(quantile(Djl.Bernoulli(p), p_))
bernoulli(ω::Ω, p::Real, T = Int) = bernoulli(rand(ω), p, T)
bernoulli(ω::Ω, p, T = Int) = (p_ = rand(ω, size(p)); bernoulli.(p_, p, T))
ppapl(rv::Bernoulli{T}, ωπ) where T = bernoulli(ωπ, reify(ωπ, params(rv))..., T)

"$(SIGNATURES) Bernoulli distribution with weight `p` and return type `T`"
bernoulli(p, T::Type{RT} = Int) where RT = Bernoulli{RT}(p)
bernoulli(p, sz::Dims, T::Type{RT} = Int) where RT = Bernoulli{RT}(lift(fill)(p, sz))

"Categorical distribution with probability weight vector `p`"
struct Categorical{A} <: PrimRandVar
  p::A    # Probability Vector
  id::ID
  Categorical(p::A, id = uid()) where A = new{A}(p, id)
end
@inline (rv::Categorical)(ω::Ω) = apl(rv, ω)
rvtransform(rv::Categorical) = categorical
categorical(p_::Real, p::Vector) = quantile(Djl.Categorical(p), p_)
categorical(ω::Ω, p::Vector{<:Real}) = categorical(rand(ω), p)
categorical(ω::Ω, p) = categorical(rand(ω, size(p)), p)

categorical(p) = Categorical(p)
categorical(p, sz::Dims) = Categorical(lift(fill)(p, sz))

"Gamma distribution with parameters α, θ"
struct Gamma{A, B} <: PrimRandVar
  α::A
  θ::B
  id::ID
  Gamma(α::A, θ::B, id = uid()) where {A, B} = new{A, B}(α, θ, id)
end
@inline (rv::Gamma)(ω::Ω) = apl(rv, ω)
rvtransform(rv::Gamma) = gammarv
gammarv(p::Real, α::Real, θ::Real) = quantile(Djl.Gamma(α, θ), p)
gammarv(ω::Ω, α::Real, θ::Real) = gammarv(rand(ω), α, θ)
gammarv(ω::Ω, α, θ) = (p = rand(ω, anysize(α, θ)); gammarv.(p, α, θ))

gammarv(α, θ) = Gamma(α, θ)
gammarv(α, θ, sz::Dims) = Gamma(lift(fill)(α, sz), lift(fill)(θ, sz))
const Γ = gammarv

"Gamma distribution with parameters α, θ"
struct InverseGamma{A, B} <: PrimRandVar
  α::A
  θ::B
  id::ID
  InverseGamma(α::A, θ::B, id = uid()) where {A, B} = new{A, B}(α, θ, id)
end
@inline (rv::InverseGamma)(ω::Ω) = apl(rv, ω)
rvtransform(rv::InverseGamma) = invgamma
invgamma(p::Real, α::Real, θ::Real) = quantile(Djl.InverseGamma(α, θ), p)
invgamma(ω::Ω, α::Real, θ::Real) = invgamma(rand(ω), α, θ)
invgamma(ω::Ω, α, θ) = (p = rand(ω, anysize(α, θ)); invgamma.(p, α, θ))

invgamma(α, θ) = InverseGamma(α, θ)
invgamma(α, θ, sz::Dims) = InverseGamma(lift(fill)(α, sz), lift(fill)(θ, sz))
const invΓ = invgamma

"Rademacher distribution"
struct Rademacher{T} <: PrimRandVar
  id::ID
  Rademacher{T}(id = uid()) where {T, A} = new{T}(id)
end
@inline (rv::Rademacher)(ω::Ω) = apl(rv, ω)
rvtransform(::Rademacher) = rademacher
rademacher(p::Real, T::Type{RT} = Int) where RT = bernoulli(p, 0.5, RT) * 2 - 1
rademacher(ω::Ω, T::Type{RT} = Int) where RT = rademacher(rand(ω), RT)

rademacher(p, T::Type{RT} = Int) where RT = Rademacher{RT}()
# FIXME: Nowhere to put size info, since rademacher has no params.  Flaw in design?!

"Poisson distribution with rate parameter `λ`"
struct Poisson{A} <: PrimRandVar
  λ::A
  id::ID
  Poisson(λ::A, id = uid()) where A = new{A}(λ, id)
end
@inline (rv::Poisson)(ω::Ω) = apl(rv, ω)
rvtransform(::Poisson) = poisson
poisson(p::Real, λ::Real) = quantile(Djl.Poisson(λ), p)
poisson(ω::Ω, λ::Real) = poisson(rand(ω), λ)
poisson(ω::Ω, λ) = (p = rand(ω, size(λ)); poisson.(p, λ))

poisson(λ) = Poisson(λ)
poisson(λ, sz::Dims) = Poisson(lift(fill(λ, sz)))

"Normal distribution with parameters μ, σ"
struct Normal{A, B} <: PrimRandVar
  μ::A
  σ::B
  id::ID
  Normal(μ::A, σ::B, id = uid()) where {A, B} = new{A, B}(μ, σ, id)
end

@inline (rv::Normal)(ω::Ω) = apl(rv, ω)
rvtransform(rv::Normal) = normal
normal(p::Real, μ::Real, σ::Real) = quantile(Djl.Normal(μ, σ), p)
normal(ω::Ω, μ::Real, σ::Real) = normal(rand(ω), μ, σ)
normal(ω::Ω, μ, σ) = (p = rand(ω, anysize(μ, σ)); normal.(p, μ, σ))

normal(μ, σ; id::ID = uid()) = Normal(μ, σ, id)
# normal(μ, σ) = Normal(μ, σ)
normal(μ, σ, sz::Dims) = Normal(lift(fill)(μ, sz), lift(fill)(σ, sz))
ciid(x::Normal; id) = Normal(x.μ, x.σ,id)

"Logistic Distribution with mean μ and shape s"
struct Logistic{A, B} <: PrimRandVar
  μ::A
  s::B
  id::ID
  Logistic(μ::A, s::B, id = uid()) where {A, B} = new{A, B}(μ, s, id)
end
@inline (rv::Logistic)(ω::Ω) = apl(rv, ω)
rvtransform(::Logistic) = logistic
logistic(p::Real, μ::Real, s::Real) = μ + s * log(p / (1 - p))
logistic(ω::Ω, μ::Real, s::Real) = logistic(rand(ω), μ, s)
logistic(ω::Ω, μ::Real, s::Real, dims::Dims) = logistic.(rand(ω, dims), μ, s)
logistic(ω::Ω, μ, s) = (p = rand(ω, anysize(μ, s)); logistic.(p, μ, s))

logistic(μ, s) = Logistic(μ, s)
logistic(μ, s, sz::Dims) = Logistic(lift(fill)(μ, sz), lift(fill)(s, sz))

"Exponential Distribution with λ"
struct Exponential{A} <: PrimRandVar
  λ::A
  id::ID
  Exponential(λ::A, id = uid()) where A = new{A}(λ, id)
end
@inline (rv::Exponential)(ω::Ω) = apl(rv, ω)
rvtransform(::Exponential) = exponential
exponential(p::Real, λ::Real) = quantile(Djl.Exponential(λ), p)
exponential(ω::Ω, λ::Real) = exponential(rand(ω), λ)
exponential(ω::Ω, λ) = (p = rand(ω, size(λ)); exponential.(p, λ))

exponential(λ) = Exponential(λ)
exponential(λ, sz::Dims) = Exponential(lift(fill(λ, sz)))

"Kumaraswamy distribution, similar to beta but easier"
struct Kumaraswamy{A, B} <: PrimRandVar
  a::A
  b::B
  id::ID
  Kumaraswamy(a::A, b::B, id = uid()) where {A, B} = new{A, B}(a, b, id)
end
@inline (rv::Kumaraswamy)(ω::Ω) = apl(rv, ω)
rvtransform(::Kumaraswamy) = kumaraswamy
kumaraswamy(p::Real, a::Real, b::Real) = (1 - (1 - p)^(1/b))^(1/a)
kumaraswamy(ω::Ω, a::Real, b::Real) = kumaraswamy(rand(ω), a, b)
kumaraswamy(ω::Ω, a, b) = (p = rand(ω, anysize(a, b)); kumaraswamy.(p, a, b))

kumaraswamy(a, b) = Kumaraswamy(a, b)
kumaraswamy(a, b, sz::Dims) = Kumaraswamy(lift(fill)(a, sz), lift(fill)(b, sz))

"Uniform distribution with parameters a, b"
struct Uniform{A, B} <: PrimRandVar
  a::A
  b::B
  id::ID
  Uniform(a::A, b::B, id = uid()) where {A, B} = new{A, B}(a, b, id)
end
@inline (rv::Uniform)(ω::Ω) = apl(rv, ω)
rvtransform(rv::Uniform) = uniform
uniform(p::Real, a::Real, b::Real) = quantile(Djl.Uniform(a, b), p)
uniform(ω::Ω, a::Real, b::Real) = uniform(rand(ω), a, b)
uniform(ω::Ω, a, b) = (p = rand(ω, anysize(a, b)); uniform.(p, a, b))

uniform(a, b) = Uniform(a, b)
uniform(a, b, sz::Dims) = Uniform(lift(fill)(a, sz), lift(fill)(b, sz))

"Categorical distribution with probability weight vector `p`"
struct UniformChoice{A} <: PrimRandVar
  vals::A    # Probability Vector
  id::ID
  UniformChoice(vals::A, id = uid()) where A = new{A}(vals, id)
end
@inline (rv::UniformChoice)(ω::Ω) = apl(rv, ω)
rvtransform(rv::UniformChoice) = uniform
# zt: FIXME abstractvector?
uniform(ω::Ω, vals::Union{Vector, UnitRange}) = rand(ω, vals)
uniform(ω::Ω, vals::Union{Vector, UnitRange}, dims::Dims) = rand(ω, vals, dims)
uniform(ω::Ω, vals::Union{Vector, UnitRange}, n::Integer) = rand(ω, vals, n)

uniform(vals) = UniformChoice(vals)
uniform(vals, sz::Dims) = UniformChoice(lift(fill)(vals, sz))

"Discrete uniform distribution over unit range `range`"
uniform(range::MaybeRV{UnitRange}) = UniformChoice(range) #FIXME

"Discrete uniform distribution over array"
uniform(arr::MaybeRV{Array}) = UniformChoice(arr)


"""
    sqrt_kalman_filter(model::StateSpaceModel{Typ}; tol::Typ = Typ(1e-5)) where Typ

Square-root Kalman filter with big Kappa initialization.
"""
function sqrt_kalman_filter(model::StateSpaceModel{Typ}; tol::Typ = Typ(1e-5)) where Typ

    time_invariant = model.mode == "time-invariant"

    # Load dimensions
    n, p, m, r = size(model)

    # Load system
    y = model.y
    Z, T, R = ztr(model)

    ensure_pos_sym!(model.H)
    ensure_pos_sym!(model.Q)

    sqrtH = cholesky(model.H).L
    sqrtQ = cholesky(model.Q).L

    # Predictive state and its sqrt-covariance
    a     = Matrix{Typ}(undef, n+1, m)
    sqrtP = Array{Typ, 3}(undef, m, m, n+1)

    # Innovation and its sqrt-covariance
    v     = Matrix{Typ}(undef, n, p)
    sqrtF = Array{Typ, 3}(undef, p, p, n)

    # Kalman gain
    K     = Array{Typ, 3}(undef, m, p, n)

    # Auxiliary matrices
    U2star = Array{Typ, 3}(undef, m, p, n)

    # Steady state initialization
    steadystate  = false
    tsteady      = n+1

    # Initial state: big Kappa initialization
    a[1, :]        = zeros(m, 1)
    sqrtP[:, :, 1]    = 1e3 .* Matrix(I, m, m)

    # Pre-allocating for performance
    zeros_pr = zeros(p, r)
    zeros_mp = zeros(m, p)
    range1   = (p + 1):(p + m)
    range2   = 1:p
    sqrtH_zeros_pr  = [sqrtH zeros_pr]
    zeros_mp_RsqrtQ = [zeros_mp R*sqrtQ]

    !isempty(model.missing_observations) && error("Treatment of missing values not implemented for SquareRootFilter, please use KalmanFilter.")

    # Square-root Kalman filter
    for t = 1:n
        v[t, :] = y[t, :] - Z[:, :, t]*a[t, :] - model.d[t, :]
        if steadystate
            sqrtF[:, :, t] = sqrtF[:, :, t-1]
            K[:, :, t] = K[:, :, t-1]
            a[t+1, :] = T * a[t, :] + K[:, :, t] * v[t, :]
            sqrtP[:, :, t+1] = sqrtP[:, :, t]
        else
            # Manipulation of auxiliary matrices
            U         = SMatrix{p + m, m + p + r}([Z[:, :, t]*sqrtP[:, :, t] sqrtH_zeros_pr;
                                                   T*sqrtP[:, :, t]         zeros_mp_RsqrtQ])
            G         = qr(U').Q
            Ustar     = U*G
            U2star[:, :, t] = Ustar[range1, range2]
            sqrtF[:, :, t]  = Ustar[range2, range2]

            # Kalman gain and predictive state update
            K[:, :, t]       = U2star[:, :, t]*inv(sqrtF[:, :, t])
            a[t+1, :]        = T*a[t, :] + K[:, :, t]*v[t, :] + model.c[t, :]
            sqrtP[:, :, t+1] = Ustar[range1, range1]

            # Checking if steady state was attained
            if time_invariant && check_steady_state(sqrtP, t, tol)
                steadystate = true
                tsteady     = t
            end
        end
    end

    # Return the auxiliary filter structre
    return SquareRootFilter(a, v, sqrtP, sqrtF, steadystate, tsteady, K)
end

"""
    filtered_state(model::StateSpaceModel{Typ}, sqrt_filter::SquareRootFilter{Typ}) where Typ

Obtain the filtered state estimates and their covariance matrices.
"""
function filtered_state(model::StateSpaceModel{Typ}, sqrt_filter::SquareRootFilter{Typ}) where Typ

    # Load dimensions data
    n, p, m, r = size(model)

    # Load system data
    Z, T, R = ztr(model)

    # Load filter data
    a = sqrt_filter.a
    v = sqrt_filter.v
    F = gram_in_time(sqrt_filter.sqrtF)
    P = gram_in_time(sqrt_filter.sqrtP)

    # Filtered state and its covariance
    att = Matrix{Typ}(undef, n, m)
    Ptt = Array{Typ, 3}(undef, m, m, n)

    for t = 1:n
        PZF = P[:, :, t] * Z[:, :, t]' * inv(F[:, :, t])
        att[t, :]    = a[t, :] + PZF * v[t, :] + model.c[t, :]
        Ptt[:, :, t] = P[:, :, t] - PZF * Z[:, :, t] * P[:, :, t]
        ensure_pos_sym!(Ptt, t)
    end

    return att, Ptt

end

"""
    sqrt_smoother(model::StateSpaceModel, sqrt_filter::SquareRootFilter) where Typ

Square-root smoother for state space model.
"""
function sqrt_smoother(model::StateSpaceModel{Typ}, sqrt_filter::SquareRootFilter{Typ}) where Typ

    # Load dimensions data
    n, p, m, r = size(model)

    # Load system data
    Z, T, R = ztr(model)

    # Load filter data
    a           = sqrt_filter.a
    v           = sqrt_filter.v
    tsteady     = sqrt_filter.tsteady
    sqrtF       = sqrt_filter.sqrtF
    sqrtP       = sqrt_filter.sqrtP
    K           = sqrt_filter.K

    # Smoothed state and its covariance
    alpha = Matrix{Typ}(undef, n, m)
    V     = Array{Typ, 3}(undef, m, m, n)
    L     = Array{Typ, 3}(undef, m, m, n)
    r     = Matrix{Typ}(undef, n, m)
    sqrtN = Array{Typ, 3}(undef, m, m, n)

    # Initialization
    sqrtN[:, :, end]  = zeros(m, m)
    r[end, :]      = zeros(m, 1)
    sqrtPsteady = sqrtP[:, :, end]
    sqrtFsteady = sqrtF[:, :, end]

    # Iterating backwards
    for t = n:-1:tsteady
        Psteady = gram(sqrtP[:, :, end])
        Fsteady = gram(sqrtF[:, :, end])
        N_t_1   = gram(sqrtN[:, :, t-1])
        L[:, :, t]   = T - K[:, :, end] * Z[:, :, t]
        r[t-1, :] = Z[:, :, t]' * inv(Fsteady) * v[t, :] + L[:, :, t]' * r[t, :]

        # QR decomposition of auxiliary matrix Nstar
        Nstar        = [Z[:, :, t]' * inv(sqrtF[:, :, end]) L[:, :, t]' * sqrtN[:, :, t]]
        G            = qr(Nstar').Q
        NstarG       = Nstar * G
        sqrtN[:, :, t-1]   = NstarG[1:m, 1:m]

        # Smoothed state and its covariance
        alpha[t, :] = a[t, :] + Psteady * r[t-1, :]
        V[:, :, t]  = Psteady - Psteady * N_t_1 * Psteady
    end

    for t = tsteady-1:-1:2
        L[:, :, t]   = T - K[:, :, t] * Z[:, :, t]
        r[t-1, :] = Z[:, :, t]' * inv(sqrtF[:, :, t] * sqrtF[:, :, t]') * v[t, :] + L[:, :, t]'*r[t, :]
        Nstar  = [Z[:, :, t]' * inv(sqrtF[:, :, t]) L[:, :, t]' * sqrtN[:, :, t]]

        # QR decomposition of auxiliary matrix Nstar
        NstarG     = Nstar*qr(Nstar').Q
        sqrtN[:, :, t-1] = NstarG[1:m, 1:m]

        # Smoothed state and its covariance
        P = gram(sqrtP[:, :, t])
        N = gram(sqrtN[:, :, t-1])
        alpha[t, :] = a[t, :] + P * r[t-1, :]
        V[:, :, t]  = P - (P * N * P)
    end

    L[:, :, 1]   = T - K[:, :, 1] * Z[:, :, 1]
    r_0    = Z[:, :, 1]' * inv(gram(sqrtF[:, :, 1])) * v[1, :] + L[:, :, 1]' * r[1, :]
    Nstar  = [Z[:, :, 1]' * inv(sqrtF[:, :, 1]) L[:, :, 1]' * sqrtN[:, :, 1]]
    G      = qr(Nstar').Q
    NstarG = Nstar*G

    sqrtN_0  = NstarG[1:m, 1:m]
    P_1 = gram(sqrtP[:, :, 1])
    alpha[1, :] = a[1, :] + P_1 * r_0
    V[:, :, 1]  = P_1 - (P_1 * gram(sqrtN_0) * P_1)

    # Return the Square Root kalman filter smoothed state
    return Smoother(alpha, V)
end

function get_log_likelihood_params(model::StateSpaceModel, filter_type::Type{SquareRootFilter{T}}) where T

    # Obtain innovation v and its variance F
    sqrt_filter = sqrt_kalman_filter(model)
    # Return v and F
    return sqrt_filter.v, gram_in_time(sqrt_filter.sqrtF)
end

function kfas(model::StateSpaceModel{T}, filter_type::Type{SquareRootFilter{T}}) where T
    # Do the SquareRootFilter
    filter_output  = sqrt_kalman_filter(model)
    smoothed_state = sqrt_smoother(model, filter_output)
    att, Ptt       = filtered_state(model, filter_output)

    return FilterOutput(filter_output.a, att, filter_output.v,
                         gram_in_time(filter_output.sqrtP), Ptt, gram_in_time(filter_output.sqrtF),
                         filter_output.steadystate, filter_output.tsteady) ,
           SmoothedState(smoothed_state.alpha, smoothed_state.V)
end


import Base: +, -, *, /, ==
import Polynomials: polyder, polyint

struct PPoly{T}
    pis::Vector{Pair{Interval{T},Poly{T}}}
    function PPoly(pis::Vector{Pair{Interval{T},Poly{T}}}) where T
        for (j,i) in enumerate(first.(pis))
            for i′ in first.(pis[j+1:end])
                g = i ∩ i′
                !isempty(g) && error("Overlapping intervals: $(i) ∩ $(i′) = $(g) ≠ ∅")
            end
        end
        sort!(pis, by = ip -> ip[1].a)
        filter!(ip -> !isempty(ip[1]) && ip[2] != zero(T), pis)
        new{T}(pis)
    end
end

PPoly() = PPoly{Int64}(Vector{Pair{Interval{Int64},Poly{Int64}}}())

# Convenience constructor
PPoly(pis::Vector{Pair{Tuple{T,T},Vector{T}}}) where T =
    PPoly([Interval(i...)=>Poly(p)
           for (i,p) in pis])

# Constructor promoting all arguments to a common type
function PPoly(pis::Vector)
    isempty(pis) && PPoly()
    T = typeof(pis[1][1][1])
    for (i,p) in pis
        T = promote_type(T, typeof(i[1]), typeof(i[2]), eltype(p))
    end
    PPoly([Interval(T.(i)...)=>Poly(T.(p))
           for (i,p) in pis])
end

function show(io::IO, pp::PPoly{T}) where T
    write(io, "Piecewise polynomial{$(T)}")
    for (i,p) in pp.pis
        write(io, "\n $(p) on $(i)")
    end
end

# Evaluation
function (pp::PPoly)(x::Number)
    S = 0
    for (i,p) ∈ pp.pis
        if x ∈ i
            S += polyval(p, x)
        end
    end
    S
end

# Arithmetic
function +(p₁::PPoly{T}, p₂::PPoly{T}) where T
    pis = Vector{Pair{Interval{T},Poly{T}}}()
    intersections = Matrix{Bool}(undef, length(p₁.pis), length(p₂.pis))
    isub = Vector{Vector{Interval{T}}}()
    i′sub = Vector{Vector{Interval{T}}}()
    for (k,(i,p)) ∈ enumerate(p₁.pis)
        push!(isub, [i])
        for (l,(i′,p′)) ∈ enumerate(p₂.pis)
            ii′ = i ∩ i′
            intersections[k,l] = !isempty(ii′)
        end
    end
    for (i′,p′) ∈ p₂.pis
        push!(i′sub, [i′])
    end

    irows = reduce(|, intersections, dims=2)
    icols = reduce(|, intersections, dims=1)
    for (k,(i,p)) ∈ zip(eachindex(irows), p₁.pis)
        !irows[k] && continue
        for (l,(i′,p′)) ∈ zip(eachindex(icols), p₂.pis)
            !icols[l] && continue
            ii′ = i ∩ i′
            push!(pis, ii′ => p+p′)
            isub[k] = collapse!(Vector{Interval{T}}(vcat([(is \ ii′)[1] for is ∈ isub[k]]...)))
            i′sub[l] = collapse!(Vector{Interval{T}}(vcat([(is \ ii′)[1] for is ∈ i′sub[l]]...)))
        end
    end
    for (is,p) ∈ zip(isub,last.(p₁.pis))
        for i ∈ is
            push!(pis, i => p)
        end
    end
    for (i′s,p′) ∈ zip(i′sub,last.(p₂.pis))
        for i′ ∈ i′s
            push!(pis, i′ => p′)
        end
    end
    PPoly(pis)
end

-(p::PPoly) = PPoly([i=>-p for (i,p) in p.pis])
-(p₁::PPoly{T}, p₂::PPoly{T}) where T = p₁ + (-p₂)

==(p₁::PPoly{T}, p₂::PPoly{T}) where T = isempty((p₁-p₂).pis)

*(pp::PPoly, n::Number) = PPoly([i=>p*n for (i,p) in pp.pis])
*(n::Number, pp::PPoly) = pp*n
/(pp::PPoly, n::Number) = PPoly([i=>p/n for (i,p) in pp.pis])

*(pp::PPoly, p′::Poly) = PPoly([i=>p*p′ for (i,p) in pp.pis])
*(p′::Poly,pp::PPoly) = pp*p′

function *(p₁::PPoly{T}, p₂::PPoly{T}) where T
    pis = Vector{Pair{Interval{T},Poly{T}}}()
    intersections = Matrix{Bool}(undef, length(p₁.pis), length(p₂.pis))
    isub = Vector{Vector{Interval{T}}}()
    i′sub = Vector{Vector{Interval{T}}}()
    for (k,(i,p)) ∈ enumerate(p₁.pis)
        push!(isub, [i])
        for (l,(i′,p′)) ∈ enumerate(p₂.pis)
            ii′ = i ∩ i′
            intersections[k,l] = !isempty(ii′)
        end
    end
    for (i′,p′) ∈ p₂.pis
        push!(i′sub, [i′])
    end

    irows = reduce(|, intersections, dims=2)
    icols = reduce(|, intersections, dims=1)
    for (k,(i,p)) ∈ zip(eachindex(irows), p₁.pis)
        !irows[k] && continue
        for (l,(i′,p′)) ∈ zip(eachindex(icols), p₂.pis)
            !icols[l] && continue
            ii′ = i ∩ i′
            push!(pis, ii′ => p*p′)
        end
    end
    PPoly(pis)
end

# /(pp::PPoly, p′::Poly) = PPoly([i=>p/p′ for (i,p) in pp.pis])

polyder(pp::PPoly, k=1) = PPoly([i => polyder(p,k) for (i,p) in pp.pis])
polyint(pp::PPoly) = PPoly([i=>polyint(p) for (i,p) in pp.pis])

function polyint(pp::PPoly{T}, a::Number, b::Number) where T
    ab = Interval(a, b)
    S = 0
    for (i,p) ∈ pp.pis
        if !(isempty(i ∩ ab))
            S += polyint(p, max(i.a,a), min(i.b,b))
        end
    end
    S
end

export PPoly, polyder, polyint


# This script actually solves the RBCCampbellCochrane model with a risk-adjusted linearization
# and times the methods, if desired
using BenchmarkTools, RiskAdjustedLinearizations, Test
include("rbc_cc.jl")

# Settings: what do you want to do?
time_methods        = false
numerical_algorithm = :relaxation
autodiff            = false
n_strips            = 3

# Set up
autodiff_method = autodiff ? :forward : :central
m_rbc_cc = RBCCampbellCochraneHabits()
m = rbc_cc(m_rbc_cc, n_strips)
z0 = copy(m.z)
y0 = copy(m.y)
Ψ0 = copy(m.Ψ)

# Solve!
solve!(m; algorithm = numerical_algorithm, autodiff = autodiff_method)

if n_strips == 0
    if numerical_algorithm == :relaxation
        sssout = JLD2.jldopen(joinpath(dirname(@__FILE__), "..", "..", "test", "reference", "rbccc_sss_iterative_output.jld2"), "r")
    elseif numerical_algorithm == :homotopy
        sssout = JLD2.jldopen(joinpath(dirname(@__FILE__), "..", "..", "test", "reference", "rbccc_sss_homotopy_output.jld2"), "r")
    end
elseif n_strips == 3
    if numerical_algorithm == :relaxation
        sssout = JLD2.jldopen(joinpath(dirname(@__FILE__), "..", "..", "test", "reference", "rbccc_sss_iterative_N3_output.jld2"), "r")
    elseif numerical_algorithm == :homotopy
        sssout = JLD2.jldopen(joinpath(dirname(@__FILE__), "..", "..", "test", "reference", "rbccc_sss_homotopy_N3_output.jld2"), "r")
    end
end

if n_strips in [0, 3]
    @test isapprox(sssout["z_rss"], m.z, atol=1e-4)
    @test isapprox(sssout["y_rss"], m.y, atol=2e-4)
    @test isapprox(sssout["Psi_rss"], m.Ψ, atol=1e-4)
end

if time_methods
    println("Deterministic steady state")
    @btime begin
        solve!(m, z0, y0; algorithm = :deterministic, autodiff = autodiff_method, verbose = :none)
    end

    # Use deterministic steady state as guesses
    solve!(m, z0, y0; algorithm = :deterministic, autodiff = autodiff_method, verbose = :none)
    zdet = copy(m.z)
    ydet = copy(m.y)
    Ψdet = copy(m.Ψ)

    println("Relaxation method")
    @btime begin # called the "iterative" method in the original paper
        solve!(m, zdet, ydet, Ψdet; algorithm = :relaxation, autodiff = autodiff_method, verbose = :none)
    end

    println("Relaxation method with automatic differentiation")
    @btime begin # called the "iterative" method in the original paper
        solve!(m, zdet, ydet, Ψdet; algorithm = :relaxation, autodiff = :forward, verbose = :none)
    end

    println("Relaxation method with Anderson acceleration")
    @btime begin # called the "iterative" method in the original paper
        solve!(m, zdet, ydet, Ψdet; algorithm = :relaxation, use_anderson = true, m = 3, autodiff = autodiff_method, verbose = :none)
    end

    println("Homotopy method")
    @btime begin # called the "continuation" method in the original paper, but is called homotopy in the original code
        solve!(m, zdet, ydet, Ψdet; algorithm = :homotopy, autodiff = autodiff_method, verbose = :none)
    end
end


using PQEcolaPoint
using Dates

# what is the contract?
short_put_contract = PutContractModel()
short_put_contract.ticker = "XYZ"
short_put_contract.expiration_date = Date(2022, 04, 14)
short_put_contract.strike_price = 65.0
short_put_contract.premium = 0.0
short_put_contract.number_of_contracts = 1
short_put_contract.direction = 1

# other version -
# origin_point = encode(JIT_lattice; S = 70.12, DTE = 53.0, IV = 0.4926, K = 65.0)
val_γ = γ(short_put_contract; Sₒ=70.12, number_of_levels=80, σ=0.5129, T=(53.0 / 365), μ=0.0047)


module straight_duct_examples
using FinEtools
using FinEtoolsAcoustics
using UnicodePlots

function straight_duct_H8_example()
    t0  =  time()

    rho = 1.21*phun("kg/m^3");# mass density
    c  = 343.0*phun("m/s");# sound speed
    bulk =  c^2*rho;
    omega =  54.5901*phun("rev/s")
    vn0 =  -1.0*phun("m/s")
    Lx = 10.0*phun("m");# length of the box, millimeters
    Ly = 1.0*phun("m"); # length of the box, millimeters
    n = 20;#number of elements along the length

    println("""

    Straight duct with anechoic termination.
    Example from Boundary element acoustics: Fundamentals and computer codes, TW Wu, page 44.
    Both real and imaginary components of the pressure should have amplitude of
    rho*c = $(rho*c).

    Hexahedral mesh.
    """)

    fens,fes  =  H8block(Lx,Ly,Ly,n,2,2); # Mesh
    bfes  =  meshboundary(fes)
    L0 = selectelem(fens,bfes,facing = true, direction = [-1.0 0.0 0.0])
    L10 = selectelem(fens,bfes,facing = true, direction = [+1.0 0.0 0.0])
    nLx = selectnode(fens,box = [0.0 Lx  0.0 0.0 0.0 0.0], inflate = Lx/1.0e5)

    geom  =  NodalField(fens.xyz)
    P  =  NodalField(zeros(ComplexF64,size(fens.xyz,1),1))

    numberdofs!(P)

    material = MatAcoustFluid(bulk,rho)
    femm  =  FEMMAcoust(IntegDomain(fes, GaussRule(3, 2)), material)

    S  =  acousticstiffness(femm, geom, P);
    C  =  acousticmass(femm, geom, P);

    E10femm  =  FEMMAcoustSurf(IntegDomain(subset(bfes,L10),GaussRule(2, 2)), material)
    D  =  acousticABC(E10femm, geom, P);

    E0femm  =  FEMMBase(IntegDomain(subset(bfes,L0), GaussRule(2,  2)))
    fi  =  ForceIntensity(-1.0im*omega*rho*vn0);
    F  =  distribloads(E0femm, geom, P, fi, 2);

    p = (-omega^2*S +omega*1.0im*D + C)\F
    scattersysvec!(P, p[:])

    println("Pressure amplitude bounds: ")
    println("  real $(minimum(real(P.values)))/$(maximum(real(P.values)))")
    println("  imag $(minimum(imag(P.values)))/$(maximum(imag(P.values)))")

    println("Total time elapsed  =  ",time() - t0,"s")

    File  =   "straight_duct.vtk"
    scalars = real(P.values);
    vtkexportmesh(File, connasarray(fes), geom.values, FinEtools.MeshExportModule.H8;
    scalars = [("Pressure", scalars)])
    @async run(`"paraview.exe" $File`)

    # @pyimport matplotlib.pyplot as plt
    # plt.style[:use]("seaborn-whitegrid")
    # fig = plt.figure()
    # ax = plt.axes()
    # ix = sortperm(geom.values[nLx,1])
    # ax[:plot](geom.values[nLx,1][ix], real(P.values)[nLx][ix], color = :blue, label = "real")
    # ax[:plot](geom.values[nLx,1][ix], imag(P.values)[nLx][ix], color = :red, label  =  "imag")
    # ax[:set_xlabel]("x coordinate")
    # ax[:set_ylabel]("Pressure")
    # plt.show()
    # set(axis="normal", plotstyle="linespoints", linewidth=3, pointsize = 2, color = "black", xlabel = "x coordinate", ylabel = "Pressure", grid="on", title = "")
    ix = sortperm(geom.values[nLx,1])
    # f = figure()
    # plot(geom.values[nLx,1][ix], real(P.values)[nLx][ix], legend = "real", marker = "ecircle")
    # plot!(geom.values[nLx,1][ix], imag(P.values)[nLx][ix], legend = "imag", marker = "edmd")
    # figure(f)
    plt = lineplot(geom.values[nLx,1][ix], real(P.values)[nLx][ix], canvas = DotCanvas, title = "Steady-state pressure", name = "real", xlabel = "x", ylabel = "P")
    plt = lineplot!(plt, geom.values[nLx,1][ix], imag(P.values)[nLx][ix], name = "imag")
    display(plt)

    true

end # straight_duct_H8_example

function straight_duct_T10_example()
    t0  =  time()

    rho = 1.21*phun("kg/m^3");# mass density
    c  = 343.0*phun("m/s");# sound speed
    bulk =  c^2*rho;
    omega =  54.5901*phun("rev/s")
    vn0 =  -1.0*phun("m/s")
    Lx = 10.0*phun("m");# length of the box, millimeters
    Ly = 1.0*phun("m"); # length of the box, millimeters
    n = 20;#number of elements along the length

    println("""

    Straight duct with anechoic termination.
    Example from Boundary element acoustics: Fundamentals and computer codes, TW Wu, page 44.
    Both real and imaginary components of the pressure should have amplitude of
    rho*c = $(rho*c).

    Quadratic tetrahedral mesh.
    """)

    fens,fes  =  T10block(Lx,Ly,Ly,n,2,2); # Mesh
    bfes  =  meshboundary(fes)
    L0 = selectelem(fens,bfes,facing = true, direction = [-1.0 0.0 0.0])
    L10 = selectelem(fens,bfes,facing = true, direction = [+1.0 0.0 0.0])
    nLx = selectnode(fens,box = [0.0 Lx  0.0 0.0 0.0 0.0], inflate = Lx/1.0e5)

    geom  =  NodalField(fens.xyz)
    P  =  NodalField(zeros(ComplexF64,size(fens.xyz,1),1))

    numberdofs!(P)

    material = MatAcoustFluid(bulk,rho)
    femm  =  FEMMAcoust(IntegDomain(fes, TetRule(5)), material)

    S  =  acousticstiffness(femm, geom, P);
    C  =  acousticmass(femm, geom, P);

    E10femm  =  FEMMAcoustSurf(IntegDomain(subset(bfes,L10),TriRule(6)), material)
    D  =  acousticABC(E10femm, geom, P);

    E0femm  =  FEMMBase(IntegDomain(subset(bfes,L0), TriRule(6)))
    fi  =  ForceIntensity(-1.0im*omega*rho*vn0);
    F  =  distribloads(E0femm, geom, P, fi, 2);

    p = (-omega^2*S +omega*1.0im*D + C)\F
    scattersysvec!(P, p[:])

    println("Pressure amplitude bounds: ")
    println("  real $(minimum(real(P.values)))/$(maximum(real(P.values)))")
    println("  imag $(minimum(imag(P.values)))/$(maximum(imag(P.values)))")

    println("Total time elapsed  =  ",time() - t0,"s")

    File  =   "straight_duct_T10.vtk"
    scalars = real(P.values);
    vtkexportmesh(File, connasarray(fes), geom.values, FinEtools.MeshExportModule.T10;
    scalars = [("Pressure", scalars)])
    @async run(`"paraview.exe" $File`)

    # @pyimport matplotlib.pyplot as plt
    # plt.style[:use]("seaborn-whitegrid")
    # fig = plt.figure()
    # ax = plt.axes()
    # ix = sortperm(geom.values[nLx,1])
    # ax[:plot](geom.values[nLx,1][ix], real(P.values)[nLx][ix], color = :blue, label = "real")
    # ax[:plot](geom.values[nLx,1][ix], imag(P.values)[nLx][ix], color = :red, label  =  "imag")
    # ax[:set_xlabel]("x coordinate")
    # ax[:set_ylabel]("Pressure")
    # plt.show()
    # set(axis="normal", plotstyle="linespoints", linewidth=3, pointsize = 2, color = "black", xlabel = "x coordinate", ylabel = "Pressure", grid="on", title = "")
    ix = sortperm(geom.values[nLx,1])
    # f = figure()
    # plot(geom.values[nLx,1][ix], real(P.values)[nLx][ix], legend = "real", marker = "ecircle")
    # plot!(geom.values[nLx,1][ix], imag(P.values)[nLx][ix], legend = "imag", marker = "edmd")
    # figure(f)
    plt = lineplot(geom.values[nLx,1][ix], real(P.values)[nLx][ix], canvas = DotCanvas, title = "Steady-state pressure", name = "real", xlabel = "x", ylabel = "P")
    plt = lineplot!(plt, geom.values[nLx,1][ix], imag(P.values)[nLx][ix], name = "imag")
    display(plt)

    true

end # straight_duct_T10_example

function straight_duct_T4_example()
    t0  =  time()

    rho = 1.21*phun("kg/m^3");# mass density
    c  = 343.0*phun("m/s");# sound speed
    bulk =  c^2*rho;
    omega =  54.5901*phun("rev/s")
    vn0 =  -1.0*phun("m/s")
    Lx = 10.0*phun("m");# length of the box, millimeters
    Ly = 1.0*phun("m"); # length of the box, millimeters
    n = 20;#number of elements along the length

    println("""

    Straight duct with anechoic termination.
    Example from Boundary element acoustics: Fundamentals and computer codes, TW Wu, page 44.
    Both real and imaginary components of the pressure should have amplitude of
    rho*c = $(rho*c).

    Tetrahedral mesh.
    """)

    fens,fes  =  T4block(Lx,Ly,Ly,n,2,2); # Mesh
    bfes  =  meshboundary(fes)
    L0 = selectelem(fens,bfes,facing = true, direction = [-1.0 0.0 0.0])
    L10 = selectelem(fens,bfes,facing = true, direction = [+1.0 0.0 0.0])
    nLx = selectnode(fens,box = [0.0 Lx  0.0 0.0 0.0 0.0], inflate = Lx/1.0e5)

    geom  =  NodalField(fens.xyz)
    P  =  NodalField(zeros(ComplexF64,size(fens.xyz,1),1))

    numberdofs!(P)

    material = MatAcoustFluid(bulk,rho)
    femm  =  FEMMAcoust(IntegDomain(fes, SimplexRule(3, 1)), material)

    S  =  acousticstiffness(femm, geom, P);
    C  =  acousticmass(femm, geom, P);

    E10femm  =  FEMMAcoustSurf(IntegDomain(subset(bfes,L10),TriRule(3)), material)
    D  =  acousticABC(E10femm, geom, P);

    E0femm  =  FEMMBase(IntegDomain(subset(bfes,L0), TriRule(3)))
    fi  =  ForceIntensity(-1.0im*omega*rho*vn0);
    F  =  distribloads(E0femm, geom, P, fi, 2);

    p = (-omega^2*S +omega*1.0im*D + C)\F
    scattersysvec!(P, p[:])

    println("Pressure amplitude bounds: ")
    println("  real $(minimum(real(P.values)))/$(maximum(real(P.values)))")
    println("  imag $(minimum(imag(P.values)))/$(maximum(imag(P.values)))")

    println("Total time elapsed  =  ",time() - t0,"s")

    File  =   "straight_duct_T4.vtk"
    scalars = real(P.values);
    vtkexportmesh(File, connasarray(fes), geom.values, FinEtools.MeshExportModule.T4;
    scalars = [("Pressure", scalars)])
    @async run(`"paraview.exe" $File`)

    # @pyimport matplotlib.pyplot as plt
    # plt.style[:use]("seaborn-whitegrid")
    # fig = plt.figure()
    # ax = plt.axes()
    # ix = sortperm(geom.values[nLx,1])
    # ax[:plot](geom.values[nLx,1][ix], real(P.values)[nLx][ix], color = :blue, label = "real")
    # ax[:plot](geom.values[nLx,1][ix], imag(P.values)[nLx][ix], color = :red, label  =  "imag")
    # ax[:set_xlabel]("x coordinate")
    # ax[:set_ylabel]("Pressure")
    # plt.show()
    # set(axis="normal", plotstyle="linespoints", linewidth=3, pointsize = 2, color = "black", xlabel = "x coordinate", ylabel = "Pressure", grid="on", title = "")
    ix = sortperm(geom.values[nLx,1])
    # f = figure()
    # plot(geom.values[nLx,1][ix], real(P.values)[nLx][ix], legend = "real", marker = "ecircle")
    # plot!(geom.values[nLx,1][ix], imag(P.values)[nLx][ix], legend = "imag", marker = "edmd")
    # figure(f)
    plt = lineplot(geom.values[nLx,1][ix], real(P.values)[nLx][ix], canvas = DotCanvas, title = "Steady-state pressure", name = "real", xlabel = "x", ylabel = "P")
    plt = lineplot!(plt, geom.values[nLx,1][ix], imag(P.values)[nLx][ix], name = "imag")
    display(plt)

    true

end # straight_duct_T4_example

function allrun()
    println("#####################################################")
    println("# straight_duct_H8_example ")
    straight_duct_H8_example()
    println("#####################################################")
    println("# straight_duct_T10_example ")
    straight_duct_T10_example()
    println("#####################################################")
    println("# straight_duct_T4_example ")
    straight_duct_T4_example()
end # function allrun

end # module straight_duct_examples


using StanBase

ProjDir = @__DIR__

bernoulli = "
data {
  int<lower=1> N;
  int<lower=0,upper=1> y[N];
  real empty[0];
}
parameters {
  real<lower=0,upper=1> theta;
}
model {
  theta ~ beta(1,1);
  y ~ bernoulli(theta);
}
";

sm = HelpModel( "bernoulli", bernoulli;
  method = StanBase.Help(:sample))

# Other options for method are:
# :optimize, :diagnose or :variational
# see ?CmdStanModels

rc = stan_help(sm; n_chains=1)

if success(rc)
	run(`cat $(sm.log_file[1])`)
end


VERSION >= v"0.4.0-dev+6521" && __precompile__()

module MDCT
export mdct, imdct

import Base.FFTW.fftwNumber, Base.FFTW.r2r, Base.FFTW.r2r!, Base.FFTW.REDFT11

fftwsimilar{T<:fftwNumber}(X::AbstractArray{T}, sz) = Array(T, sz...)
fftwsimilar{T<:Real}(X::AbstractArray{T}, sz) = Array(Float64, sz...)
fftwsimilar{T<:Complex}(X::AbstractArray{T}, sz) = Array(Complex128, sz...)

# The following two routines compute the MDCT and IMDCT via
# a type-IV DCT (FFTW's REDFT11 r2r transform).  For a review
# of this relationship, see the notes I posted on Wikipedia:
#     http://en.wikipedia.org/wiki/Modified_discrete_cosine_transform

function mdct{T<:Number}(X::AbstractVector{T})
    sz = length(X)
    if isodd(sz)
        throw(ArgumentError("mdct requires an even-length vector"))
    end
    N = div(sz, 2);
    Y = fftwsimilar(X, N)
    if isodd(N)
        throw(ArgumentError("mdct requires a multiple-of-4 vector length"))
        # FIXME: handle odd case via DCT-III?
    else
        N2 = div(N,2)
        for i = 1:N2
            Y[i] = -0.5 * (X[(3*N2+1)-i] + X[3*N2+i])
            Y[N2+i] = -0.5 * (X[(2*N2+1)-i] - X[i])
        end
        r2r!(Y, REDFT11)
    end
    return Y
end

function imdct{T<:Number}(X::StridedVector{T})
    N = length(X)
    Z = fftwsimilar(X, 2*N)
    if isodd(N)
        throw(ArgumentError("imdct requires an even vector length"))
        # FIXME: handle odd case via DCT-II?
    else
        Y = r2r(X, REDFT11)
        N2 = div(N,2)
        s = 0.5 / N
        for i = 1:N2
            Z[N+1-i] = -(Z[i] = Y[N2+i]*s)
            Z[N+N2+1-i] = (Z[N+N2+i] = -Y[i]*s)
        end
    end
    return Z
end

imdct{T<:Number}(X::AbstractVector{T}) = imdct(copy!(fftwsimilar(X, size(X)),
                                                     X))

end # MDCT


using ApproxFun, Base.Test

f=Fun(exp,Jacobi(2.,.5))
@test_approx_eq f(.1) exp(.1)

f=Fun(x->cos(100x),Jacobi(2.124,.5),500)
@test_approx_eq f(.1) cos(100*.1)

sp=Jacobi(2.124,.5)
f=Fun(exp,sp)
sp2=Jacobi(2.124,1.5)
f2=Fun(exp,sp2)
sp3=Jacobi(3.124,1.5)
f3=Fun(exp,sp3)
sp4=Jacobi(4.124,2.5)
f4=Fun(exp,sp4)
@test norm((Fun(f,sp2)-f2).coefficients)<10eps()
@test norm((Fun(f,sp3)-f3).coefficients)<10eps()
@test norm((Fun(f,sp4)-f4).coefficients)<20eps()

m=20
f=Fun(x->((1-x)/2).^m.*exp(x),JacobiWeight(0.,m,Jacobi(2m+1,0.)))
@test abs(f(.1)-(x->((1-x)/2).^m.*exp(x))(.1))<10eps()

m=10
f=Fun(x->besselj(m,m*(1-x)),JacobiWeight(0.,m,Jacobi(2m+1,0.)))
@test_approx_eq f(0.) besselj(m,m)

## Conversion

@test norm(Fun(Fun(exp),Jacobi(-.5,-.5))-Fun(exp,Jacobi(-.5,-.5))) < 100eps()

x=Fun(identity)
ri=0.5./(1-x)
@test_approx_eq ((1-x)./2.*Fun(exp,JacobiWeight(0.,0.,Jacobi(1.,0.))))(.1) (1-.1)./2*exp(.1)

@test_approx_eq ((1-x)./2.*Fun(exp,JacobiWeight(0.,0.,Jacobi(1.,0.))))(.1) (1-.1)./2*exp(.1)

@test_approx_eq (ri.*Fun(exp,JacobiWeight(0.,0.,Jacobi(1.,0.))))(.1) .5/(1-.1)*exp(.1)

## Derivative

S=JacobiWeight(0.,0.,Jacobi(1.,0.,Interval(1.,0.)))
D=Derivative(S)
f=Fun(exp,domainspace(D))
@test (D*f-f).coefficients|>norm < eps(100000.)
@test (f'-f).coefficients|>norm < eps(100000.)
@test (D^2*f-f).coefficients|>norm < eps(100000000.)
@test (D*(D*f)-f).coefficients|>norm < eps(100000000.)

### Jacobi multiplication

x=Fun(identity,Jacobi(0.,0.))
f=Fun(exp,Jacobi(0.,0.))

@test_approx_eq (x*f)(.1) .1exp(.1)

x=Fun(identity,Jacobi(0.123,12.324))
f=Fun(exp,Jacobi(0.,0.))

@test_approx_eq (x*f)(.1) .1exp(.1)

x=Fun(identity,Jacobi(0.123,12.324))
f=Fun(exp,Jacobi(0.213,0.590))

@test_approx_eq (x*f)(.1) .1exp(.1)

g=Fun(cos,Jacobi(0.123,12.324))
f=Fun(exp,Jacobi(0.213,0.590))

@test_approx_eq (g*f)(.1) cos(.1)*exp(.1)

## Jacobi integrate and sum

@test_approx_eq sum(Fun(exp,Legendre([0,2]))) sum(Fun(exp,[0,2]))

a=Arc(0.,.1,0.,π/2)
g=Fun(exp,Legendre(a))

@test_approx_eq sum(g) sum(Fun(exp,a))

## Test special derivative

x=Fun()
f=exp(x)*sqrt(1-x^2)
D=Derivative(WeightedJacobi(.5,.5))
g=(D*Fun(f,domainspace(D)))
@test_approx_eq f'(0.1) g(0.1)

## Test implementation of conversion between Chebyshev and Jacobi spaces using FastTransforms

f = Fun(x->cospi(1000x))
g = Fun(f,Legendre())
h = Fun(g,Chebyshev())
@test norm(f.coefficients-h.coefficients,Inf) < 10eps()
h = Fun(h,Legendre())
@test norm(g.coefficients-h.coefficients,Inf) < 1000eps()

## Legendre conversions

@test norm(Fun(exp,Ultraspherical(1//2))-Fun(exp,Jacobi(0,0))) < 100eps()

C=Conversion(Jacobi(0,0),Chebyshev())
@test norm(C*Fun(exp,Jacobi(0,0))  - Fun(exp)) < 100eps()

C=Conversion(Ultraspherical(1//2),Chebyshev())
@test norm(C*Fun(exp,Ultraspherical(1//2))  - Fun(exp)) < 100eps()

C=Conversion(Chebyshev(),Ultraspherical(1//2))
@test norm(C*Fun(exp)-Fun(exp,Legendre())) < 100eps()

C=Conversion(Chebyshev(),Jacobi(0,0))
@test norm(C*Fun(exp)  - Fun(exp,Jacobi(0,0))) < 100eps()

C=Conversion(Chebyshev(),Jacobi(1,1))
@test norm(C*Fun(exp) - Fun(exp,Jacobi(1,1))) < 100eps()

C=Conversion(Ultraspherical(1//2),Ultraspherical(1))
@test norm(C*Fun(exp,Ultraspherical(1//2))-Fun(exp,Ultraspherical(1))) < 100eps()

C=Conversion(Jacobi(0,0),Ultraspherical(1))
@test norm(C*Fun(exp,Jacobi(0,0))-Fun(exp,Ultraspherical(1))) < 100eps()

C=Conversion(Ultraspherical(1),Jacobi(0,0))
@test norm(C*Fun(exp,Ultraspherical(1))-Fun(exp,Jacobi(0,0))) < 100eps()


# Fourier method
# First published by Milstein, 1994

struct Milstein <: AbstractIteratedIntegralAlgorithm end

convorder(::Milstein) = 1//2
errcoeff(m, h, ::Milstein, ::MaxL2) = h/(√2*π)
norv(m, n, ::Milstein) = 2*m*n+m

"""
    levyarea(W, n, alg::Milstein)

Simulates an approximation of the iterated Itô-integrals ``\\int_0^1W_s\\otimes dW_s``
of the given ``m``-dimensional increment of a Wiener process with step size 1.
The parameter ``n`` specifies the number of terms in the approximation and thus determines the accuracy.
This is an efficient implementation of the algorithm proposed in [Milstein, 1994](@ref milstein1994).
It is based on a Fourier expansion of the Wiener process.
The algorithm needs approximately ``m^2+2\\cdot m\\cdot n`` Float's.
The time complexity is ``\\mathcal{O}(m^2\\cdot n)``.
"""
function levyarea(W::AbstractVector{T}, n::Integer, alg::Milstein) where {T}
    rng = default_rng()
    m = length(W)
    X = randn(rng, T, n, m) # allocates m*n Floats
    Y = randn(rng, T, m, n) # allocates m*n Floats
    Y .= (Y .- √(T(2)).*W) ./ (1:n)'
    A = Y * X # allocates m*m Floats
    # Add simple rest approximation (a₀)
    # M = randn(rng, T, m) # allocates m Floats
    M = randn!(rng, view(Y, :, 1))
    a = T(√(2*trigamma(n+1)))
    A .+= a .* W .* M'
    # Antisymmetrize
    # G = inv(2*T(π)).*(A .- A') # allocates m*m Floats
    for i = 1:m
        @inbounds A[i,i] = zero(T)
        for j = (i+1):m
            @inbounds A[i,j] = (A[i,j] - A[j,i])/(2*T(π))
            @inbounds A[j,i] = -A[i,j]
        end
    end
    return A
end


module ChangePoint

include("../src/Jaynes.jl")
using .Jaynes
using Gen

# ------------ Change point model ------------ #

jfunc = @jaynes function change_point_model(N::Int)
    n ~ DiscreteUniform(N)
    λ₁ ~ Gamma(2, 2)
    λ₂ ~ Gamma(3, 2)
    x = [i >= n ? {:x => i} ~ Poisson(λ₂) : {:x => i} ~ Poisson(λ₁) for i in 1 : N]
    x
end

# Sample trace.
tr = simulate(jfunc, (50, ))
display(tr)

# ------------ Gibbs ------------ #

@kern function gibbs_kernel(trace)
    trace ~ mh(trace, select(:n))
    trace ~ mh(trace, select(:λ₁))
    trace ~ mh(trace, select(:λ₂))
end

# ------------ Inference ------------ #

infer = () -> begin
    obs = constrain([i <= 30 ? (:x => i, ) => rand(Poisson(3)) : (:x => i, ) => rand(Poisson(7)) for i in 1 : 50])

    # Sample trace.
    tr, w = generate(jfunc, (50, ), obs)

    trs = Trace[]
    for i in 1 : 5000
        tr, _ = gibbs_kernel(tr)
        i % 50 == 0 && push!(trs, tr)
    end

    # Posterior est.
    λ₁_est, λ₂_est, n_est = zip(map(trs) do tr
                                    (tr[:λ₁], tr[:λ₂], tr[:n])
                                end...)
    len = length(λ₁_est)
    println("est: $((sum(λ₁_est) / len,
                     sum(λ₂_est) / len,
                     sum(n_est) / len))")

end
infer()
@time infer()

end # module


# Fisher metric
struct FisherMetric{N} <: AbstractMetric
    I::Function
    dom::Function
end

# Constructors
FisherMetricManifold(N::Int,I::Function,dom::Function=p->true) = MetricManifold(
    Euclidean(N),
    FisherMetric{N}(I,dom)
)
FisherMetricManifold(d,dom::Function=p->true) = MetricManifold(
    Euclidean(length(fieldnames(d))),
    FisherMetric{length(fieldnames(d))}(p -> fim(eval(Expr(:call,d,p...))),dom)
)

# Methods
function local_metric(M::MetricManifold{ℝ,<:AbstractManifold,<:FisherMetric},p,B::InducedBasis)
    metric(M).I(p)
end
function check_point(M::MetricManifold{ℝ,<:AbstractManifold,<:FisherMetric{N}}, p; kwargs...) where {N}
    (size(p)) == (N,) || return DomainError(size(p),"The size of $p is not $((N,)).")
    metric(M).dom(p) || return DomainError(p,"p is not a valid point in the metric domain")
    return nothing
end
function log!(M::MetricManifold{ℝ,<:AbstractManifold,<:FisherMetric},X,p,q)
    # Use shooting
    X .= optimize(X -> norm(q - geodesic(M,p,X)(1.0)), q - p * norm(exp(M,p,X) - p) / norm(q - p)).minimizer
end
function inner(M::MetricManifold{ℝ,<:AbstractManifold,<:FisherMetric},p,X,Y)
    X' * metric(M).I(p) * Y
end

# Default basis (to make things easy)
function default_basis(M,p)
    induced_basis(M,Manifolds.RetractionAtlas(),p,TangentSpace)
end


@testset "approximators" begin
    @testset "TabularQ" begin
        A = TabularQ([1 2; 3 4])
        @test A(1) == [1., 2.]
        @test A(2) == [3., 4.]
        @test A(2, 2) == 4
        @test A(2,Val(:dist)) == [3., 4.]
        @test A(2,Val(:max)) == 4.
        @test A(2,Val(:argmax)) == 2

        update!(A, 2, 2, -1.)
        @test A(2) == [3., 3.]
        @test A(2, 2) == 3

        update!(A, 1, [-1., -1.])
        @test A(1) == [0., 1.]
        @test A(1, 1) == 0
        @test A(1, 2) == 1
    end

    @testset "TabularV" begin
        V = TabularV([1., 2., 3., 2., 1.])
        @test V(1) == 1.
        @test V(3) == 3.

        update!(V, 3, -1.)
        @test V(3) == 2.
    end

    @testset "AggregationV" begin
        table = [-1., 1.]
        V = AggregationV(table, s -> s < 0 ? 1 : 2)

        @test V(-5) == -1.
        @test V(5) == 1.

        update!(V, -1, 0.5)
        update!(V, 1, -0.5)

        @test V(-5) == -.5
        @test V(5) == 0.5
    end

    @testset "FourierV" begin
        V = FourierV([1, 1, 1])
        @test V(1) ≈ 1.0
        @test V(2) ≈ 3.0

        update!(V, 1, 0.5)
        @test V(1) ≈ 2.5
        @test V(2) ≈ 3.5
    end

    @testset "LinearV" begin
        features = [1 0 1; 0 1 1]
        weights = [1., 2., 3.]
        V = LinearV(features, weights)

        @test V(1) ≈ 1. + 3.
        @test V(2) ≈ 2. + 3.

        update!(V, 1, 0.5)
        @test V(1) ≈ 1.5 + 3.5
        @test V(2) ≈ 2. + 3.5
    end

    @testset "PolynomialV" begin
        V = PolynomialV([1., 2., 3.])

        @test V(2) ≈ 1 + 2. * 2  + 3. * 2^2

        update!(V, 2, 1.)
        @test V(2) ≈ 2 + 4. * 2 + 7. * 2^2
    end

    @testset "tilings" begin
        init_tiling = Tiling((0:5:50, 0:10:100))
        tilings = [init_tiling, init_tiling - [1, 2]]
        @testset "TilingsV" begin
            V = TilingsV(tilings)
            update!(V, [0, 0], 1.)
            @test V([0,0]) == 2.
            @test V([4,8]) == 1.
            @test V([5,10]) == 0.
        end
        @testset "TilingsQ" begin
            Q = TilingsQ(tilings, 2)
            update!(Q, [0, 0], 1, 1.)

            @test Q([0, 0]) ≈ [2., 0.]
            @test Q([0, 0], 1) ≈ 2.
            @test Q([0, 0], Val(:max)) ≈ 2.
            @test Q([0, 0], Val(:argmax)) == 1
        end
    end
end


using Distributions, Statistics

"""
    KernelScenarios(data::Union{Array{Int64,2},Array{Float64,2}},kernelDistribution = Logistic)

    This function takes a matrix of data in 2 dimension (NxT) where N is the number of trajectories and T is the number of stages.
    It uses function closures and above is a function.
    To get results a trajectory from above, `KernelScenarios(data,kernelDistribution)()` gives the result.
    User specifies both data and the distribution of the kernel he/she wants to use. The default kernel is the Logistic kernel.
"""

function KernelScenarios(data::Union{Array{Int64,2},Array{Float64,2}},kernelDistribution = Logistic; Markovian::Bool = true)
    # We use the function closure here because we want to use this function in the stochastic approximation process.
    # That is, we wwant to follow the specifications of the TreeApproximation and Latticeapproximation inputs.
    # The way we have created the stochastic approximation functions is that it takes a function which does not take any inputs.
    # And so the outer function takes the data and the distribution of the kernel while the inner function takes no inputs.
    # To obtain the result then, call "KernelScenarios(data,Logistic)()".
    # Notice we have an input of whether you are creating Markovian trajectories or not.
    # This is important for generating either scenario trees or scenario lattices.
    # The default for Markovian is true hence the trajectory created is Markovian which is used for scenario lattices.
    function closure()
        N,T = size(data)                                                  # Dimensions of the data
        d = 1
        w = fill(1.0,N)                                                     # initialize the weights
        ξ = Array{Float64,1}(undef,T)                                     # trajectory to be created
        for t = 1:T
            w = w / sum(w)                                                # normalized weights
            Nt = sum(w)^2 / sum(w.^2)                                     # effective sample size
            σt = sqrt(sum(w .* ((data[:,t] .- sum(w .* data[:,t])).^2)))  # effective standard devistion
            ht = σt * Nt^(-1/(d + 4)) + eps()
            # Composition method comes in here
            u = rand(Uniform(0,1))
            _,jstar = findmin(cumsum(w) .< u * sum(w))
            # The cumulative sum of weights leads to a high probability of picking a data path near an observation (Sequin && Pichler 2017)
            ξ[t] = rand(kernelDistribution(data[jstar,t],ht))
            # The new generated value is according to the distribution of density of the current stage and dependent on the history of all the data paths.
            # All the rows in the column are associated with a certain weight as follows:
            # Each weight is a product of the kernels of the data at that point.
            if t<T
                for j = 1:N
                    # The choice of the kernel does not have any important effect on density estimation (Jones (1990)).
                    if Markovian
                        # Markovian is used for scenario lattices
                        w[j] = pdf(kernelDistribution(data[j,t],ht),ξ[t])
                    else
                        # Non-Markovian is used for scenario lattices
                        w[j] *= pdf(kernelDistribution(data[j,t],ht),ξ[t])
                    end
                end
            end
        end
        return ξ                         # The length of ξ is equal to the number stages = T
    end
end


ra = 89.014303
dec = 13.924912
ω = 37.59

α = deg2rad(ra)
δ = deg2rad(dec)
d = 1000.0 / ω

pmra = 372.72
pmdec = -483.69
radial_velocity = 0.37

η = deg2rad(58.5986320306)
αGC = deg2rad(266.4051)
δGC = - deg2rad(-28.936175)#the sign in the rotation matrix as per documentation seems to be inverted; cannot understand the root cause of the problem
dGC = 8300
xGC = [1, 0, 0]
z = 27
θ = asin(z / dGC)

r_icrs = d * [cos(α) * cos(δ), sin(α) * cos(δ), sin(δ)]

R₁ = Matrix(undef, 3, 3)
R₁[1,1] = cos(δGC)
R₁[1,2] = 0
R₁[1,3] = - sin(δGC)
R₁[2,:] = [0 1 0]
R₁[3,1] = sin(δGC)
R₁[3,2] = 0
R₁[3,3] = cos(δGC)
println("R₁:", R₁)

R₂ = Matrix(undef, 3, 3)
R₂[1,1] = cos(αGC)
R₂[1,2] = sin(αGC)
R₂[1,3] = 0
R₂[2,1] = - sin(αGC)
R₂[2,2] = cos(αGC)
R₂[2,3] = 0
R₂[3,:] = [0 0 1]
println("R₂:", R₂)

R₃ = Matrix(undef, 3, 3)
R₃[1,:] = [1 0 0]
R₃[2,1] = 0
R₃[2,2] = cos(η)
R₃[2,3] = sin(η)
R₃[3,1] = 0
R₃[3,2] = - sin(η)
R₃[3,3] = cos(η)
println("R₃:", R₃)

R = R₃ * R₁ * R₂
println("R:", R)

R = Matrix(undef, 3, 3)
R[1,1] = cos(αGC) * cos(δGC)
R[1,2] = cos(δGC) * sin(αGC)
R[1,3] = - sin(δGC)
R[2,1] = cos(αGC) * sin(δGC) * sin(η) - sin(αGC) * cos(η)
R[2,2] = sin(αGC) * sin(δGC) * sin(η) + cos(αGC) * cos(η)
R[2,3] = cos(δGC) * sin(η)
R[3,1] = cos(αGC) * sin(δGC) * cos(η) + sin(αGC) * sin(η)
R[3,2] = sin(αGC) * sin(δGC) * cos(η) - cos(αGC) * sin(η)
R[3,3] = cos(δGC) * cos(η)
println("R:", R)

H = Matrix(undef, 3, 3)
H[1,1] = cos(θ)
H[1,2] = 0
H[1,3] = sin(θ)
H[2,1] = 0
H[2,2] = 1
H[2,3] = 0
H[3,1] = - sin(θ)
H[3,2] = 0
H[3,3] = cos(θ)
println("H:", H)

r = R * r_icrs - dGC * xGC
rGC = H * r

println("distance:", d)
println("r_icrs:", r_icrs)
println("r:", r)
println("rGC:", rGC)


using SpecialFunctions; # need to evaluate the Hankel function (Green's function)

function matricesAlpert(k,prt,order)

    (nskip, chi, weights) = Alpert_params(order)
    # nskip = a
    ninterp = order + 3;

    τ = prt.τ
    ν = prt.ν
    s = prt.s
    ds= prt.ds

    N = prt.N
    Δs = prt.Δs

    x1 = prt.x[:,1];     x2 = prt.x[:,2]
    dx1 = prt.dx[:,1];   dx2 = prt.dx[:,2]

    S  = zeros(N,N)+im*zeros(N,N)
    K  = zeros(N,N)+im*zeros(N,N)
    Kp = zeros(N,N)+im*zeros(N,N)
    T  = zeros(N,N)+im*zeros(N,N)

    for i=1:N

        for j=1:N

            r = sqrt((x1[i]-x1[j])^2+(x2[i]-x2[j])^2);

            H0 = hankelh1(0,k*r); H1 = hankelh1(1,k*r)

            # if i!=j

                S2  = im/4*H0*τ[j]*Δs*ds[j]
                K2  = im*k/4*H1/r*((x1[i]-x1[j])*dx2[j]-(x2[i]-x2[j])*dx1[j])*Δs*ds[j]
                Kp2 = im*k/4*H1/r*((x1[j]-x1[i])*dx2[i]-(x2[j]-x2[i])*dx1[i])*(τ[j]/τ[i])*Δs*ds[j]

                T2 = ((x1[i]-x1[j])*dx1[j]+(x2[i]-x2[j])*dx2[j])*((x1[i]-x1[j])*dx1[i]+(x2[i]-x2[j])*dx2[i])/r^2
                T2 = im/2*T2*(k^2*H0-2*k*H1/r)+im*k/2*(dx1[i]*dx1[j]+dx2[i]*dx2[j])*H1/r
                T2 = -0.5*T2*Δs*ds[j]/τ[j]

            # end

            S[i,j]  = S2
            K[i,j]  = K2
            Kp[i,j] = Kp2
            T[i,j]  = T2

        end

        # @. S[i, mod(i+(-nskip+1:nskip-1)-1, N)+1] = 0

        S[i, mod.(i.+(-nskip+1:nskip-1).-1, N).+1] .= 0
        K[i, mod.(i.+(-nskip+1:nskip-1).-1, N).+1] .= 0
        Kp[i,mod.(i.+(-nskip+1:nskip-1).-1, N).+1] .= 0
        T[i, mod.(i.+(-nskip+1:nskip-1).-1, N).+1] .= 0

    end

    for p = 1:size(chi, 1)

        if prt.PD == 0
            t = Δs*collect(0:N-1)
            w = mod.(t .+ Δs*chi[p], 2*pi)
            dw = ones(size(w))
        else
            t = Δs*collect(1:N)
            (w,dw) = polySigmoid(mod.(t .+ Δs*chi[p], 2*pi),prt.PD)
        end

        (aux_x, aux_dx) = prt.param(w)
        # @show aux_x
        # return

        aux_x1  = aux_x[:,1]
        aux_x2  = aux_x[:,2]

        aux_dx1 = aux_dx[:,1]
        aux_dx2 = aux_dx[:,2]

        tau_aux = sqrt.(aux_dx1.^2+aux_dx2.^2)

        aux_n1 = aux_dx2./tau_aux
        aux_n2 =-aux_dx1./tau_aux

        aux_r = sqrt.((x1-aux_x1).^2+(x2-aux_x2).^2)

        aux_H0 = hankelh1.(0,k*aux_r)
        aux_H1 = hankelh1.(1,k*aux_r)

        ind = floor.(chi[p] - ninterp/2 + 1) .+ (0 : ninterp - 1)

        L   = lagrangeInterp(ind', chi[p])

        kerS = im/4*aux_H0.*tau_aux.*dw
        kerK = im*k/4*aux_H1./aux_r.*((x1-aux_x1).*aux_dx2-(x2-aux_x2).*aux_dx1).*dw
        kerKp= im*k/4*aux_H1./aux_r.*((aux_x1-x1).*ν[:,1]+(aux_x2-x2).*ν[:,2]).*(tau_aux.*dw)

        kerT = ((x1-aux_x1).*aux_dx1+(x2-aux_x2).*aux_dx2).*((x1-aux_x1).*dx1+(x2-aux_x2).*dx2)./aux_r.^2
        kerT = im/2*kerT.*(k^2*aux_H0-2*k*aux_H1./aux_r)+
                im*k/2*(dx1.*aux_dx1+dx2.*aux_dx2)./aux_r.*aux_H1
        kerT = -0.5*kerT.*dw./tau_aux

        for i0=1:N
           j0  = i0 .+ ind
           j0  = floor.(Int,mod.(j0.-1, N) .+ 1)
           S[i0,j0] .= S[i0,j0]  + Δs*weights[p]*kerS[i0]*L[:,1]
           K[i0,j0] .= K[i0,j0]  + Δs*weights[p]*kerK[i0]*L[:,1]
           Kp[i0,j0].= Kp[i0,j0] + Δs*weights[p]*kerKp[i0]*L[:,1]
           T[i0,j0] .= T[i0,j0]  + Δs*weights[p]*kerT[i0]*L[:,1]
        end

    end

    T = k^2*S.*(ν*ν') + T

    mat = (S=S,K=K,Kp=Kp,T=T);

end

########################################################################

function  Alpert_params(order)

    if order == 2
        a = 1
        chi_p = [.591549430918953e-01]
        w_p = [5.000000000000000e-01]

    elseif order == 6
        a = 3
        chi_p= [4.004884194926570e-03;
                7.745655373336686e-02;
                3.972849993523248e-01;
                1.075673352915104e+00;
                2.003796927111872e+00]

        w_p = [ 1.671879691147102e-02;
                1.636958371447360e-01;
                4.981856569770637e-01;
                8.372266245578912e+00;
                9.841730844088381e+00]

    elseif order == 10
        a = 6
        chi_p = [1.175089381227308e-03;
             1.877034129831289e-02;
             9.686468391426860e-02;
             3.004818668002884e-01;
             6.901331557173356e-01;
             1.293695738083659e+00;
             2.090187729798780e+00;
             3.016719313149212e+00;
             4.001369747872486e+00;
             5.000025661793423e+00]

        w_p = [4.560746882084207e-03;
           3.810606322384757e-02;
           1.293864997289512e-01;
           2.884360381408835e-01;
           4.958111914344961e-01;
           7.077154600594529e-01;
           8.741924365285083e-01;
           9.661361986515218e-01;
           9.957887866078700e-01;
           9.998665787423845e-01]

    elseif order == 16
        a = 10;
        chi_p = [8.371529832014113e-04
                  1.239382725542637e-02
                  6.009290785739468e-02
                  1.805991249601928e-01
                  4.142832599028031e-01
                  7.964747731112430e-01
                  1.348993882467059e+00
                  2.073471660264395e+00
                  2.947904939031494e+00
                  3.928129252248612e+00
                  4.957203086563112e+00
                  5.986360113977494e+00
                  6.997957704791519e+00
                  7.999888757524622e+00
                  8.999998754306120e+00]

        w_p = [3.190919086626234e-03
               2.423621380426338e-02
               7.740135521653088e-02
               1.704889420286369e-01
               3.029123478511309e-01
               4.652220834914617e-01
               6.401489637096768e-01
               8.051212946181061e-01
               9.362411945698647e-01
               1.014359775369075e+00
               1.035167721053657e+00
               1.020308624984610e+00
               1.004798397441514e+00
               1.000395017352309e+00
               1.000007149422537e+00]

    end

    chi_p = [chi_p' -chi_p[end:-1:1]']';
    w_p = [w_p' w_p[end:-1:1]']';

    return (a, chi_p, w_p)

end

##############################################################################

function lagrangeInterp(xx, t)
    #  L is a row vector calculating the coefficients of lagrange interpolation
    #  at t xx are data
    n = length(xx) # here xx = index = [-2, -1, 0, 1, 2]
    denom = xx'*ones(1, n) - ones(n, 1)*xx
    numer = t*ones(n, n) - ones(n, 1)*xx
    for i = 1:n
      denom[i, i] = 1.0
      numer[i, i] = 1.0
    end
    temp = numer./denom
    L = prod(temp,dims=2)
    return L
end


"""
    latitudes(layer::T) where {T <: SimpleSDMLayer}

Returns an iterator with the latitudes of the SDM layer passed as its argument.
This returns the latitude at the center of each cell in the grid.
"""
latitudes(layer::T) where {T <: SimpleSDMLayer} = LinRange(layer.bottom+stride(layer, 2), layer.top-stride(layer, 2), size(layer,1))

"""
    longitudes(layer::T) where {T <: SimpleSDMLayer}

Returns an iterator with the longitudes of the SDM layer passed as its argument.
This returns the longitudes at the center of each cell in the grid.
"""
longitudes(layer::T) where {T <: SimpleSDMLayer} = LinRange(layer.left+stride(layer, 1), layer.right-stride(layer, 1), size(layer,2))

"""
    _layers_are_compatible(l1::X, l2::Y) where {X <: SimpleSDMLayer, Y <: SimpleSDMLayer}

    Internal function to verify if layers are compatible, i.e. have the same
    size and bounding coordinates.

"""
function _layers_are_compatible(l1::X, l2::Y) where {X <: SimpleSDMLayer, Y <: SimpleSDMLayer}
    size(l1) == size(l2) || throw(ArgumentError("The layers have different sizes"))
    l1.top == l2.top || throw(ArgumentError("The layers have different top coordinates"))
    l1.left == l2.left || throw(ArgumentError("The layers have different left coordinates"))
    l1.bottom == l2.bottom || throw(ArgumentError("The layers have different bottom coordinates"))
    l1.right == l2.right || throw(ArgumentError("The layers have different right coordinates"))
end

function _layers_are_compatible(layers::Array{T}) where {T <: SimpleSDMLayer}
    all(layer -> _layers_are_compatible(layer, layers[1]), layers)
end


# Create example data
x = randn(100_000)
y = randn(100_000)
# Draw an hexagonal plot in the bigger bottom-right region
subplot(3, 3, (5, 9))
hexbin(x, y, colorbar = false)
# Draw marginal histograms
subplot(3, 3, (2, 3))
histogram(x)
subplot(3, 3, (4, 7))
histogram(y, horizontal = true, xflip = true)
# Draw a shade plot in the smaller top-left region
subplot(3, 3, 1)
shade(x, y)


facts("KDTrees.knn") do

context("KDTrees.knn") do

# 8 node rectangle
data = [0.0 0.0 0.0 0.5 0.5 1.0 1.0 1.0;
        0.0 0.5 1.0 0.0 1.0 0.0 0.5 1.0]
tree = KDTree(data)

idxs, dists = knn(tree, [0.8, 0.8], 1)
@fact idxs[1] --> 8 # Should be closest to top right corner
@fact sqrt(0.2^2 + 0.2^2) --> roughly(dists[1])

idxs, dists = knn(tree, [0.1, 0.8], 3)
@fact idxs --> [3, 2, 5]

idxs, dists = knn(tree, [1//10, 8//10], 3)
@fact idxs --> [3, 2, 5]

@fact_throws knn(tree, [0.1, 0.8], 10) # k > n_points

@fact_throws knn(tree, [0.1], 10) # n_dim != trees dim

end  #context

end # facts