text:
function HierarchicalPeriodicHMM_trig(θ_Q::AbstractArray, θ_Y::AbstractArray, T::Int)
    K, D, size_memory = size(θ_Y)
    @argcheck K == size(θ_Q, 1)

    A = zeros(K, K, T)
    [A[k, l, t] = exp(polynomial_trigo(t, θ_Q[k, l, :], T = T)) for k = 1:K, l = 1:K-1, t = 1:T]
    [A[k, K, t] = 1 for k = 1:K, t = 1:T] # last colum is 1/normalization (one could do otherwise)
    normalization_polynomial = [1 + sum(A[k, l, t] for l = 1:K-1) for k = 1:K, t = 1:T]
    [A[k, l, t] /= normalization_polynomial[k, t] for k = 1:K, l = 1:K, t = 1:T]

    p = [1 / (1 + exp(polynomial_trigo(t, θ_Y[k, s, h, :], T = T))) for k = 1:K, t = 1:T, s = 1:D, h = 1:size_memory]

    return HierarchicalPeriodicHMM(A, Bernoulli.(p))
end

function HierarchicalPeriodicHMM_trig(a::AbstractVector, θ_Q::AbstractArray, θ_Y::AbstractArray, T::Int)
    K, D, size_memory = size(θ_Y)
    @assert K == size(θ_Q, 1)

    A = zeros(K, K, T)
    [A[k, l, t] = exp(polynomial_trigo(t, θ_Q[k, l, :], T = T)) for k = 1:K, l = 1:K-1, t = 1:T]
    [A[k, K, t] = 1 for k = 1:K, t = 1:T] # last colum is 1/normalization (one could do otherwise)
    normalization_polynomial = [1 + sum(A[k, l, t] for l = 1:K-1) for k = 1:K, t = 1:T]
    [A[k, l, t] /= normalization_polynomial[k, t] for k = 1:K, l = 1:K, t = 1:T]

    p = [1 / (1 + exp(polynomial_trigo(t, θ_Y[k, s, h, :], T = T))) for k = 1:K, t = 1:T, s = 1:D, h = 1:size_memory]

    return HierarchicalPeriodicHMM(a, A, Bernoulli.(p))
end

function rand(
    rng::AbstractRNG,
    hmm::PeriodicHMM,
    n2t::AbstractArray{Int},
    useless;
    init = rand(rng, Categorical(hmm.a)),
    seq = false
)
    T = size(hmm.B, 2)
    N = length(n2t)

    z = Vector{Int}(undef, N)
    (T >= 1) && (z[1] = init)
    for n = 2:N
        tm1 = n2t[n-1] # periodic t-1
        z[n] = rand(rng, Categorical(hmm.A[z[n-1], :, tm1]))
    end
    y = rand(rng, hmm, n2t, z, useless)
    seq ? (z, y) : y
end

function rand(rng::AbstractRNG, hmm::PeriodicHMM{Univariate}, n2t::AbstractArray{Int}, z::AbstractVector{<:Integer}, useless::String)
    y = Vector(undef, length(z))
    T = size(hmm.B, 2)
    @argcheck length(n2t) == length(z)
    for n in eachindex(z)
        t = n2t[n] # periodic t
        y[n] = rand(rng, hmm.B[z[n], t])
    end
    y
end

function rand(rng::AbstractRNG, hmm::PeriodicHMM{Univariate}, n2t::AbstractArray{Int}, z::AbstractVector{<:Integer}, yini)
    y = Vector(undef, length(z))
    T = size(hmm.B, 2)
    @argcheck length(n2t) == length(z)
    y[1] = yini
    for n in eachindex(z)
        t = n2t[n] # periodic t
        p = yesterday.(params.(hmm.B[z[n], t].v), y[n-1])
        y[n] = rand(rng, Product(Bernoulli.(p)))
    end
    y
end

function rand(
    rng::AbstractRNG,
    hmm::PeriodicHMM{Multivariate},
    n2t::AbstractArray{Int},
    z::AbstractVector{<:Integer},
    useless::String
)
    y = Matrix{eltype(hmm.B[1])}(undef, length(z), size(hmm, 2))
    T = size(hmm.B, 2)
    @argcheck length(n2t) == length(z)
    for n in eachindex(z)
        t = n2t[n] # periodic t
        y[n, :] = rand(rng, hmm.B[z[n], t])
    end
    y
end

function rand(
    rng::AbstractRNG,
    hmm::PeriodicHMM{Multivariate},
    n2t::AbstractArray{Int},
    z::AbstractVector{<:Integer},
    yini
)
    T = size(hmm, 2)
    y = Matrix{eltype(hmm.B[1])}(undef, length(z), T)
    @argcheck length(n2t) == length(z)
    N_ini = size(yini, 2)
    y[1:N_ini, :] = yini
    for n in eachindex(z)[N_ini+1:end]
        t = n2t[n] # periodic t
        p = yesterday.(params.(hmm.B[z[n], t].v), y[n-1, :])
        y[n, :] = rand(rng, Product(Bernoulli.(p)))
    end
    y
end

rand(hmm::PeriodicHMM, n2t::AbstractArray{Int}, useless; kwargs...) = rand(GLOBAL_RNG, hmm, n2t, useless; kwargs...)

rand(hmm::PeriodicHMM, n2t::AbstractArray{Int}, z::AbstractVector{<:Integer}, useless; kwargs...) = rand(GLOBAL_RNG, hmm, n2t, z, useless; kwargs...)

function yesterday(probj::Tuple, yj)
    probj = probj[1]
    probj[day2(yj, 1)] / (probj[day2(yj, 1)] + probj[day2(yj, 0)]) #1-> dd, 2-> wd, 3 -> dw, 4 -> ww
end
day3(b, y, t) = passmissing(Int)(round(1 + 1.4b + 2.4y + 3.4t))
day2(y, t) = (1 + y + 2t)

###

function likelihoods!(L::AbstractMatrix, hmm::PeriodicHMM{Univariate}, observations, n2t::AbstractArray{Int})
    N, K, T = size(observations, 1), size(hmm, 1), size(hmm, 3)
    @argcheck size(LL) == (N, K)
    @inbounds for i in OneTo(K), n in OneTo(T)
        t = n2t[n] # periodic t
        L[n, i] = pdf(hmm.B[i, t], observations[n])
    end
end

function likelihoods!(L::AbstractMatrix, hmm::PeriodicHMM{Multivariate}, observations, n2t::AbstractArray{Int})
    N, K, T = size(observations, 1), size(hmm, 1), size(hmm, 3)
    @argcheck size(LL) == (N, K)
    @inbounds for i in OneTo(K), n in OneTo(T)
        t = n2t[n] # periodic t
        L[n, i] = pdf(hmm.B[i, t], view(observations, n, :))
    end
end

function loglikelihoods!(LL::AbstractMatrix, hmm::PeriodicHMM{Univariate}, observations, n2t::AbstractArray{Int})
    N, K, T = size(observations, 1), size(hmm, 1), size(hmm, 3)
    @argcheck size(LL) == (N, K)
    @inbounds for i in OneTo(K), n in OneTo(N)
        t = n2t[n] # periodic t
        LL[n, i] = logpdf(hmm.B[i, t], observations[n])
    end
end

function loglikelihoods!(LL::AbstractMatrix, hmm::PeriodicHMM{Multivariate}, observations, n2t::AbstractArray{Int})
    N, K, T = size(observations, 1), size(hmm, 1), size(hmm, 3)
    @argcheck size(LL) == (N, K)
    @inbounds for i in OneTo(K), n in OneTo(N)
        t = n2t[n] # periodic t
        LL[n, i] = logpdf(hmm.B[i, t], view(observations, n, :))
    end
end

####

function update_A!(
    A::AbstractArray{T,3} where {T},
    ξ::AbstractArray,
    α::AbstractMatrix,
    β::AbstractMatrix,
    LL::AbstractMatrix,
    n2t::AbstractArray{Int}
)
    @argcheck size(α, 1) == size(β, 1) == size(LL, 1) == size(ξ, 1)
    @argcheck size(α, 2) ==
              size(β, 2) ==
              size(LL, 2) ==
              size(A, 1) ==
              size(A, 2) ==
              size(ξ, 2) ==
              size(ξ, 3)

    N, K = size(LL)
    T = size(A, 3)
    @inbounds for n in OneTo(N - 1)
        t = n2t[n] # periodic t
        m = vec_maximum(view(LL, n + 1, :))
        c = 0.0

        for i in OneTo(K), j in OneTo(K)
            ξ[n, i, j] = α[n, i] * A[i, j, t] * exp(LL[n+1, j] - m) * β[n+1, j]
            c += ξ[n, i, j]
        end

        for i in OneTo(K), j in OneTo(K)
            ξ[n, i, j] /= c
        end
    end

    fill!(A, 0.0)
    ## For periodicHMM only the n observation corresponding to A(t) are used to update A(t)
    n_in_t = [findall(n2t .== t) for t = 1:T] # could probably be speeded up

    @inbounds for t in OneTo(T)
        for i in OneTo(K)
            c = 0.0

            for j in OneTo(K)
                for n in setdiff(n_in_t[t], N)
                    A[i, j, t] += ξ[n, i, j]
                end
                c += A[i, j, t]
            end

            for j in OneTo(K)
                A[i, j, t] /= c
            end
        end
    end
end

# # In-place update of the observations distributions.
function update_B!(B::AbstractMatrix, γ::AbstractMatrix, observations, estimator, n2t::AbstractArray{Int})
    @argcheck size(γ, 1) == size(observations, 1)
    @argcheck size(γ, 2) == size(B, 1)
    @argcheck size(observations, 1) == size(n2t, 1)
    N = size(γ, 1)
    K = size(B, 1)
    T = size(B, 2)
    ## For periodicHMM only the n observation corresponding to A(t) are used to update A(t)
    C = ncategories(B[1].v[1])
    n_in_t = [findall(n2t .== t) for t = 1:T] # could probably be speeded up. For exemple computed outside the function only once
    if typeof(B[1]) in [Product{Discrete,Categorical{Float64,Vector{Float64}},Vector{Categorical{Float64,Vector{Float64}}}}, Categorical{Float64,Vector{Float64}}]
        @inbounds for t in OneTo(T)
            n_t = n_in_t[t]
            for i in OneTo(K)
                if sum(γ[n_t, i]) > 0
                    B[i, t] = fit_mle((B[i, t]), C, permutedims(observations[n_t, :]), γ[n_t, i])
                end
            end
        end
    else
        @inbounds for t in OneTo(T)
            n_t = n_in_t[t]
            for i in OneTo(K)
                if sum(γ[n_t, i]) > 0
                    B[i, t] = estimator((B[i, t]), permutedims(observations[n_t, :]), γ[n_t, i])
                end
            end
        end
    end
end

function fit_mle(hmm::PeriodicHMM, observations, n2t::AbstractArray{Int}; init = :none, kwargs...)
    hmm = copy(hmm)

    if init == :kmeans
        kmeans_init!(hmm, observations, display = get(kwargs, :display, :none))
    end

    history = fit_mle!(hmm, observations, n2t; kwargs...)
    hmm, history
end

function fit_mle!(
    hmm::PeriodicHMM,
    observations,
    n2t::AbstractArray{Int}
    ;
    display = :none,
    maxiter = 100,
    tol = 1e-3,
    robust = false,
    estimator = fit_mle
)
    @argcheck display in [:none, :iter, :final]
    @argcheck maxiter >= 0

    N, K, T = size(observations, 1), size(hmm, 1), size(hmm, 3)
    @argcheck T == size(hmm.B, 2)
    history = EMHistory(false, 0, [])

    # n2t = date # dayofyear_Leap.(date)

    # Allocate memory for in-place updates
    c = zeros(N)
    α = zeros(N, K)
    β = zeros(N, K)
    γ = zeros(N, K)
    ξ = zeros(N, K, K)
    LL = zeros(N, K)

    loglikelihoods!(LL, hmm, observations, n2t)
    robust && replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))

    forwardlog!(α, c, hmm.a, hmm.A, LL, n2t)
    backwardlog!(β, c, hmm.a, hmm.A, LL, n2t)
    posteriors!(γ, α, β)

    logtot = sum(c)
    (display == :iter) && println("Iteration 0: logtot = $logtot")

    for it = 1:maxiter
        update_a!(hmm.a, α, β)
        update_A!(hmm.A, ξ, α, β, LL, n2t)
        update_B!(hmm.B, γ, observations, estimator, n2t)

        # Ensure the "connected-ness" of the states,
        # this prevents case where there is no transitions
        # between two extremely likely observations.
        robust && (hmm.A .+= eps())

        @check isprobvec(hmm.a)
        @check istransmats(hmm.A)

        loglikelihoods!(LL, hmm, observations, n2t)
        robust && replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))

        forwardlog!(α, c, hmm.a, hmm.A, LL, n2t)
        backwardlog!(β, c, hmm.a, hmm.A, LL, n2t)
        posteriors!(γ, α, β)

        logtotp = sum(c)
        (display == :iter) && println("Iteration $it: logtot = $logtotp")

        push!(history.logtots, logtotp)
        history.iterations += 1

        if abs(logtotp - logtot) < tol
            (display in [:iter, :final]) &&
                println("EM converged in $it iterations, logtot = $logtotp")
            history.converged = true
            break
        end

        logtot = logtotp
    end

    if !history.converged
        if display in [:iter, :final]
            println("EM has not converged after $(history.iterations) iterations, logtot = $logtot")
        end
    end

    history
end

# In-place forward pass, where α and c are allocated beforehand.
function forwardlog!(
    α::AbstractMatrix,
    c::AbstractVector,
    a::AbstractVector,
    A::AbstractArray{T,3} where {T},
    LL::AbstractMatrix,
    n2t::AbstractArray{Int}
)
    @argcheck size(α, 1) == size(LL, 1) == size(c, 1)
    @argcheck size(α, 2) == size(LL, 2) == size(a, 1) == size(A, 1) == size(A, 2)

    N, K = size(LL)
    T = size(A, 3)

    fill!(α, 0.0)
    fill!(c, 0.0)

    m = vec_maximum(view(LL, 1, :))

    for j in OneTo(K)
        α[1, j] = a[j] * exp(LL[1, j] - m)
        c[1] += α[1, j]
    end

    for j in OneTo(K)
        α[1, j] /= c[1]
    end

    c[1] = log(c[1]) + m

    for n = 2:N
        tm1 = n2t[n-1] # periodic t-1
        m = vec_maximum(view(LL, n, :))
        #
        for j in OneTo(K)
            for i in OneTo(K)
                α[n, j] += α[n-1, i] * A[i, j, tm1]
            end
            α[n, j] *= exp(LL[n, j] - m)
            c[n] += α[n, j]
        end
        #
        for j in OneTo(K)
            α[n, j] /= c[n]
        end
        #
        c[n] = log(c[n]) + m
    end
end

# In-place backward pass, where β and c are allocated beforehand.
function backwardlog!(
    β::AbstractMatrix,
    c::AbstractVector,
    a::AbstractVector,
    A::AbstractArray{T,3} where {T},
    LL::AbstractMatrix,
    n2t::AbstractArray{Int}
)
    @argcheck size(β, 1) == size(LL, 1) == size(c, 1)
    @argcheck size(β, 2) == size(LL, 2) == size(a, 1) == size(A, 1) == size(A, 2)

    N, K = size(LL)
    T = size(A, 3)
    L = zeros(K)
    (T == 0) && return

    fill!(β, 0.0)
    fill!(c, 0.0)

    for j in OneTo(K)
        β[end, j] = 1.0
    end

    @inbounds for n = N-1:-1:1
        t = n2t[n] # periodic t
        m = vec_maximum(view(LL, n + 1, :))

        for i in OneTo(K)
            L[i] = exp(LL[n+1, i] - m)
        end

        for j in OneTo(K)
            for i in OneTo(K)
                β[n, j] += β[n+1, i] * A[j, i, t] * L[i]
            end
            c[n+1] += β[n, j]
        end

        for j in OneTo(K)
            β[n, j] /= c[n+1]
        end

        c[n+1] = log(c[n+1]) + m
    end

    m = vec_maximum(view(LL, 1, :))

    for j in OneTo(K)
        c[1] += a[j] * exp(LL[1, j] - m) * β[1, j]
    end

    c[1] = log(c[1]) + m
end
## For Array
function forward(a::AbstractVector, A::AbstractArray{T,3} where {T}, LL::AbstractMatrix, n2t::AbstractArray{Int})
    m = Matrix{Float64}(undef, size(LL))
    c = Vector{Float64}(undef, size(LL, 1))
    forwardlog!(m, c, a, A, LL, n2t)
    m, sum(c)
end

function backward(a::AbstractVector, A::AbstractArray{T,3} where {T}, LL::AbstractMatrix, n2t::AbstractArray{Int})
    m = Matrix{Float64}(undef, size(LL))
    c = Vector{Float64}(undef, size(LL, 1))
    backwardlog!(m, c, a, A, LL, n2t)
    m, sum(c)
end

function posteriors(a::AbstractVector, A::AbstractArray{T,3} where {T}, LL::AbstractMatrix, n2t::AbstractArray{Int}; kwargs...)
    α, _ = forward(a, A, LL, n2t; kwargs...)
    β, _ = backward(a, A, LL, n2t; kwargs...)
    posteriors(α, β)
end
## For AbstractHMM + n2t (for inhomogeneous HMM with date or period)

function loglikelihoods(hmm::AbstractHMM, observations, n2t::AbstractArray{Int}; logl = nothing, robust = false)
    (logl !== nothing) && deprecate_kwargs("logl")
    N, K = size(observations, 1), size(hmm, 1)
    LL = Matrix{Float64}(undef, N, K)

    loglikelihoods!(LL, hmm, observations, n2t)
    if robust
        replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))
    end
    LL
end

function loglikelihood(hmm::AbstractHMM, observations, n2t::AbstractArray{Int}; logl = nothing, robust = false)
    (logl !== nothing) && deprecate_kwargs("logl")
    forward(hmm, observations, n2t, robust = robust)[2]
end

function forward(hmm::AbstractHMM, observations, n2t::AbstractArray{Int}; logl = nothing, robust = false)
    (logl !== nothing) && deprecate_kwargs("logl")
    LL = loglikelihoods(hmm, observations, n2t; robust = robust)
    forward(hmm.a, hmm.A, LL, n2t)
end

function backward(hmm::AbstractHMM, observations, n2t::AbstractArray{Int}; logl = nothing, robust = false)
    (logl !== nothing) && deprecate_kwargs("logl")
    LL = loglikelihoods(hmm, observations; robust = robust)
    backward(hmm.a, hmm.A, LL, n2t)
end

function posteriors(hmm::AbstractHMM, observations, n2t::AbstractArray{Int}; logl = nothing, robust = false)
    (logl !== nothing) && deprecate_kwargs("logl")
    LL = loglikelihoods(hmm, observations, n2t; robust = robust)
    posteriors(hmm.a, hmm.A, LL, n2t)
end

## Viterbi
function viterbi(hmm::AbstractHMM, observations, n2t::AbstractArray{Int}; logl = nothing, robust = false)
    (logl !== nothing) && deprecate_kwargs("logl")
    LL = loglikelihoods(hmm, observations, n2t; robust = robust)
    viterbi(hmm.a, hmm.A, LL, n2t)
end

function viterbi(a::AbstractVector, A::AbstractArray, LL::AbstractMatrix, n2t::AbstractArray{Int}; logl = nothing)
    ## < v1.1 compatibility
    (logl !== nothing) && deprecate_kwargs("logl")
    (logl == false) && (LL = log.(LL))
    ## --------------------
    T1 = Matrix{Float64}(undef, size(LL))
    T2 = Matrix{Int}(undef, size(LL))
    z = Vector{Int}(undef, size(LL, 1))
    viterbilog!(T1, T2, z, a, A, LL, n2t::AbstractArray{Int})
    z
end

function viterbi!(
    T1::AbstractMatrix,
    T2::AbstractMatrix,
    z::AbstractVector,
    a::AbstractVector,
    A::AbstractArray{T,3} where {T},
    L::AbstractMatrix,
    n2t::AbstractArray{Int}
)
    N, K = size(L)
    T = size(A, 3)
    (T == 0) && return

    fill!(T1, 0.0)
    fill!(T2, 0)

    c = 0.0

    for i in OneTo(K)
        T1[1, i] = a[i] * L[1, i]
        c += T1[1, i]
    end

    for i in OneTo(K)
        T1[1, i] /= c
    end

    @inbounds for n = 2:N
        tm1 = n2t[n-1] # t-1
        c = 0.0

        for j in OneTo(K)
            # TODO: If there is NaNs in T1 this may
            # stay to 0 (NaN > -Inf == false).
            # Hence it will crash when computing z[t].
            # Maybe we should check for NaNs beforehand ?
            amax = 0
            vmax = -Inf

            for i in OneTo(K)
                v = T1[n-1, i] * A[i, j, tm1]
                if v > vmax
                    amax = i
                    vmax = v
                end
            end

            T1[n, j] = vmax * L[n, j]
            T2[n, j] = amax
            c += T1[n, j]
        end

        for i in OneTo(K)
            T1[n, i] /= c
        end
    end

    z[N] = argmax(T1[N, :])
    for n = N-1:-1:1
        z[n] = T2[n+1, z[n+1]]
    end
end

function viterbilog!(
    T1::AbstractMatrix,
    T2::AbstractMatrix,
    z::AbstractVector,
    a::AbstractVector,
    A::AbstractArray{T,3} where {T},
    LL::AbstractMatrix,
    n2t::AbstractArray{Int}
)
    N, K = size(LL)
    T = size(A, 3)
    (T == 0) && return

    fill!(T1, 0.0)
    fill!(T2, 0)

    al = log.(a)
    Al = log.(A)

    for i in OneTo(K)
        T1[1, i] = al[i] + LL[1, i]
    end
    @inbounds for n = 2:N
        tm1 = n2t[n-1] # t-1
        for j in OneTo(K)
            amax = 0
            vmax = -Inf

            for i in OneTo(K)
                v = T1[n-1, i] + Al[i, j, tm1]
                if v > vmax
                    amax = i
                    vmax = v
                end
            end

            T1[n, j] = vmax + LL[n, j]
            T2[n, j] = amax
        end
    end

    z[N] = argmax(T1[N, :])
    for n = N-1:-1:1
        z[n] = T2[n+1, z[n+1]]
    end
end


text:
using Revise
using AdFem
using PyCall
using LinearAlgebra
using PyPlot
using SparseArrays
using MAT
using ADCMEKit
np = pyimport("numpy")
using PyPlot
using SpecialFunctions

n = 20
NT = 100
ρ = 1.0

m = 5n
h = 1/n
Δt = 1.0/NT

xo = zeros((m+1)*(n+1))
yo = zeros((m+1)*(n+1))
for i = 1:m+1
  for j = 1:n+1
    idx = (j-1)*(m+1)+i
    xo[idx] = (i-1)*h
    yo[idx] = (j-1)*h
  end
end

bdnode = bcnode("left | lower | right", m, n, h)

bdnode = bcnode("all", m, n, h)

μ = 0.5*constant(ones(4*m*n))
η = 1000000. * constant(ones(4*m*n))

coef = 2μ*η/(η + μ*Δt)
mapH = c->begin
  c * diagm(0=>ones(2))
end
H = map(mapH, coef)

Δt = Δt * ones(NT)

M = constant(compute_fem_mass_matrix1(m, n, h))
K = compute_fem_stiffness_matrix1(H, m, n, h)

db = zeros((m+1)*(n+1))
for j = 1:n+1
  idx = (j-1)*(m+1)+1
  if j<=div(n, 4)
    db[idx] = 1.
  else
    db[idx] = (1-(j-div(n, 4))/(3div(n, 4)))
  end
end

# turn the problem into homogeneous Dirichlet problem
homoF = - K * db
idof = ones(Bool, (m+1)*(n+1))
idof[bdnode] .= false
idof = findall(idof)

d0 = zeros((m+1)*(n+1))
v0 = zeros((m+1)*(n+1))
a0 = vector(idof, M[idof, idof]\((-K*db)[idof]), (m+1)*(n+1))

# analytical solution
d0 = @. (5-xo)*xo*(1-yo)*yo
v0 = -d0
a0 = d0

ExtForce = zeros(NT, (m+1)*(n+1))
ts = αscheme_time(Δt)
for i = 1:NT
  t = ts[i]
  f = (x,y)->exp(-t)*(5-x)*x*(1-y)*y - exp(-t)*(-2y*(1-y)-2x*(5-x))
  fval = eval_f_on_gauss_pts(f, m, n, h)
  ExtForce[i,:] = compute_fem_source_term1(fval, m, n, h)
end
ExtForce = constant(ExtForce)

function antiplane_visco_αscheme(M::Union{SparseTensor, SparseMatrixCSC},
  K::Union{SparseTensor, SparseMatrixCSC},
  d0::Union{Array{Float64, 1}, PyObject},
  v0::Union{Array{Float64, 1}, PyObject},
  a0::Union{Array{Float64, 1}, PyObject},
  Δt::Array{Float64};
  ρ::Float64 = 1.0)
  nt = length(Δt)
  αm = (2ρ-1)/(ρ+1)
  αf = ρ/(1+ρ)
  γ = 1/2-αm+αf
  β = 0.25*(1-αm+αf)^2
  d = length(d0)

  M = isa(M, SparseMatrixCSC) ? constant(M) : M
  K = isa(K, SparseMatrixCSC) ? constant(K) : K
  d0, v0, a0, Δt = convert_to_tensor([d0, v0, a0, Δt], [Float64, Float64, Float64, Float64])

  A = (1-αm)*M + (1-αf)*K*β*Δt[1]^2
  A, _ = fem_impose_Dirichlet_boundary_condition1(A, bdnode, m, n, h)

  function equ(dc, vc, ac, dt, εc, σc, i)
    dn = dc + dt*vc + dt^2/2*(1-2β)*ac
    vn = vc + dt*((1-γ)*ac)

    df = (1-αf)*dn + αf*dc
    vf = (1-αf)*vn + αf*vc
    am = αm*ac

    σ_dt = (1-αf)*dt
    Σ = 2* repeat(μ.*η/(η+μ*σ_dt),1,2) .*εc - repeat(η/(η+μ*σ_dt), 1, 2) .* σc
    Force = compute_strain_energy_term1(Σ, m, n, h)

    # rhs = homoF + Force - (M*am + K*df)

    # For wave equation
    rhs = ExtForce[i] - (M*am + K*df)

    rhs = scatter_update(rhs, bdnode, zeros(length(bdnode)))
    A\rhs

  end

  function condition(i, tas...)
    return i<=nt
  end
  function body(i, tas...)
    dc_arr, vc_arr, ac_arr, σc_arr = tas
    dc = read(dc_arr, i)
    vc = read(vc_arr, i)
    ac = read(ac_arr, i)
    σc = read(σc_arr, i)
    εc = eval_strain_on_gauss_pts1(dc, m, n, h)

    op = tf.print(i, norm(σc - εc))
    i = bind(i, op)

    y = equ(dc, vc, ac, Δt[i],εc, σc, i)
    dn = dc + Δt[i]*vc + Δt[i]^2/2*((1-2β)*ac+2β*y)
    vn = vc + Δt[i]*((1-γ)*ac+γ*y)

    εn = eval_strain_on_gauss_pts1(dn, m, n, h)
    σn = 2* repeat(μ.*η/(η+μ*Δt[i]),1,2) .* (εn - εc) + repeat(η/(η+μ*Δt[i]), 1, 2) .* σc

    op = tf.print(i, norm(σn - εn))
    i = bind(i, op)

    i+1, write(dc_arr, i+1, dn), write(vc_arr, i+1, vn), write(ac_arr, i+1, y), write(σc_arr, i+1, σn)
  end

  dM = TensorArray(nt+1); vM = TensorArray(nt+1); aM = TensorArray(nt+1); σM = TensorArray(nt+1);
  dM = write(dM, 1, d0)
  vM = write(vM, 1, v0)
  aM = write(aM, 1, a0)
  # for exact solution
  # ε0 = eval_strain_on_gauss_pts1(d0, m, n, h)
  # σM = write(σM, 1, 2repeat(μ, 1, 2).*ε0)

  σM = write(σM, 1, zeros(4m*n,2))

  i = constant(1, dtype=Int32)
  _, d, v, a = while_loop(condition, body, [i,dM, vM, aM, σM])
  set_shape(stack(d), (nt+1, length(a0))), set_shape(stack(v), (nt+1, length(a0))), set_shape(stack(a), (nt+1, length(a0)))
end

d, v, a = antiplane_visco_αscheme(M, K, d0, v0, a0, Δt, ρ=ρ)

# function solver(A, rhs)
#   A, _ = fem_impose_Dirichlet_boundary_condition1(A, bdnode, m, n, h)
#   rhs = scatter_update(rhs, bdnode, zeros(length(bdnode)))
#   A\rhs
# end
# d, v, a = αscheme(M, spzero((m+1)*(n+1), (m+1)*(n+1)), K, ExtForce, d0, v0, a0, Δt; solve = solver, ρ=ρ)

sess = Session(); init(sess)
d_, v_, a_ = run(sess, [d, v, a])

# pcolormesh(reshape(d_[2,:]+db, m+1, n+1))
# colorbar()

close("all")
for (k,tid) in enumerate(LinRange{Int64}(1, NT+1, 5))
  t = (tid-1)*Δt[1]
  dd = @. exp(-t)*(5-xo)*xo*(1-yo)*yo
  plot(dd[div(n,2)*(m+1) .+ (1:m+1)], "C$k-", label="$tid")
  plot(d_[tid, :][div(n,2)*(m+1) .+ (1:m+1)],"C$k--o", label="$tid", markersize=3)
end
legend()

# close("all")
# for tid in LinRange{Int64}(1, NT+1, 5)
#   plot((a_[tid, :])[1:m+1], label="$tid")
# end
# legend()

ext: jl

max_stars_repo_name: ADCMEMarket/ADCMEImages

max_issues_repo_name: ADCMEMarket/ADCMEImages

max_forks_repo_name: ADCMEMarket/ADCMEImages


text:
"""Attempting to implement bitonic sort"""

using CUDA

# sort kernel
bisort!(shared) = begin
  k = UInt16(2)
  NUM = length(cushared)
  while (k <= NUM)
    j = div(k, 2)
    while j >= 1
      tid = UInt(((blockIdx().x - 1) * blockDim().x + threadIdx().x)-1)
      ixj = tid⊻UInt(j)
      if ixj > tid
        if (tid & k) == 0
          if shared[tid+1] > shared[ixj+1]
            tmp = shared[ixj+1]
            shared[ixj+1] = shared[tid+1]
            shared[tid+1] = tmp
          end
        else
          if shared[tid+1] < shared[ixj+1]
            tmp = shared[ixj+1]
            shared[ixj+1] = shared[tid+1]
            shared[tid+1] = tmp
          end
        end
      end
    end
    k *= 2
  end
  return
end

bitonicsort!(cushared, NUM) = begin
  nblocks = ceil(Int, NUM/256)
  @cuda threads = 256 blocks = 1024 bisort!(cushared)
end

using SortingAlgorithms, BenchmarkTools
shared = rand(Float32, 2^26)
cpusort = @belapsed sort!($shared, alg=RadixSort) #0.788

shared = rand(Float32, 2^26)

measure_gpu_sort(shared) = begin
  res = Float64[]
  for i = 1:3
    cushared = cu(shared)
    # sorted_shared = sort(shared, alg=RadixSort)
    # println("exp false;")
    # println("got $(collect(cushared) |> issorted)")
    t = Base.@elapsed begin
      bitonicsort!(cushared, length(shared))
      CUDA.synchronize()
    end
    # xx = collect(cushared)
    # println("exp true;")
    # println("got $(xx |> issorted); max error: $(1_000_000_000maximum(xx .- sorted_shared))")
    # push!(res, t)
  end

  res
end

@time measure_gpu_sort(shared)
# 6.457073086
# 2.774147852
# 2.771599214
# 2.770980271
# 2.778133025
# 2.769555927
# 2.799603755
# 2.774497496
# 2.790657341
# 2.790034242


text:
# BlockBandedMatrix with block range indexes is also block-banded
const SubBlockSkylineMatrix{T,LL,UU,R1,R2} =
    SubArray{T,2,BlockSkylineMatrix{T,LL,UU},<:Tuple{BlockSlice{<:R1},BlockSlice{<:R2}}}

function add_bandwidths(A::AbstractBlockBandedMatrix,B::AbstractBlockBandedMatrix)
    Al,Au = colblockbandwidths(A)
    Bl,Bu = colblockbandwidths(B)

    l = Vector(Bl)
    u = Vector(Bu)

    for (v,Av) in [(l,Al),(u,Au)]
        n = length(v)
        for i = 1:n
            sel = max(i-Bu[i],1):min(i+Bl[i],length(Av))
            isempty(sel) && continue
            v[i] += maximum(Av[sel])
        end
    end

    l,u
end

function add_bandwidths(A::BlockBandedMatrix,B::BlockBandedMatrix)
    l,u = blockbandwidths(A) .+ blockbandwidths(B)
    Fill(l,blocksize(B,2)), Fill(u,blocksize(B,2))
end

function similar(M::MulAdd{<:AbstractBlockBandedLayout,<:AbstractBlockBandedLayout}, ::Type{T}) where T
    A,B = M.A, M.B

    if !blockisequal(axes(A,2), axes(B,1))
        # diagonal matrices can be converted
        if isdiag(B) && size(A,2) == size(B,1) == size(B,2)
            B = BlockBandedMatrix(B.data, BlockSizes((Acols,Acols)), 0, 0, 0, 0)
        elseif isdiag(A) && size(A,2) == size(B,1) == size(A,1)
            A = BlockBandedMatrix(A.data, BlockSizes((Brows,Brows)), 0, 0, 0, 0)
        else
            throw(DimensionMismatch("*"))
        end
    end
    n,m = size(A,1), size(B,2)
    l,u = add_bandwidths(A,B)
    BlockSkylineMatrix{T}(undef, (axes(A,1),axes(B,2)), (l,u))
end

function similar(M::MulAdd{<:AbstractBandedBlockBandedLayout,<:AbstractBandedBlockBandedLayout}, ::Type{T}) where T
    A,B = M.A, M.B

    if !blockisequal(axes(A,2), axes(B,1))
        # diagonal matrices can be converted
        if isdiag(B) && size(A,2) == size(B,1) == size(B,2)
            # TODO: fix
            B = BandedBlockBandedMatrix(B.data, BlockSizes((Acols,Acols)), 0, 0, 0, 0)
        elseif isdiag(A) && size(A,2) == size(B,1) == size(A,1)
            A = BandedBlockBandedMatrix(A.data, BlockSizes((Brows,Brows)), 0, 0, 0, 0)
        else
            throw(DimensionMismatch("*"))
        end
    end
    n,m = size(A,1), size(B,2)
    Al,Au = blockbandwidths(A)
    Bl,Bu = blockbandwidths(B)
    Aλ,Aμ = subblockbandwidths(A)
    Bλ,Bμ = subblockbandwidths(B)

    BandedBlockBandedMatrix{T}(undef, (axes(A,1),axes(B,2)), (Al+Bl, Au+Bu), (Aλ+Bλ, Aμ+Bμ))
end

similar(M::MulAdd{<:DiagonalLayout,<:AbstractBlockBandedLayout}, ::Type{T}) where T = similar(M.B,T)
similar(M::MulAdd{<:AbstractBlockBandedLayout,<:DiagonalLayout}, ::Type{T}) where T = similar(M.A,T)

function blockbandwidths(V::SubBlockSkylineMatrix{<:Any,LL,UU,<:BlockRange1,<:BlockRange1}) where {LL,UU}
    A = parent(V)

    KR = parentindices(V)[1].block.indices[1]
    JR = parentindices(V)[2].block.indices[1]

    shift =  first(KR) - first(JR)
    l,u = blockbandwidths(A)
    l-shift, u+shift
end

####
# BlockIndexRange subblocks
####
sublayout(::AbstractBlockBandedLayout, ::Type{<:Tuple{<:BlockSlices, <:BlockSlices}}) = BlockBandedLayout()

sublayout(::BlockBandedColumnMajor, ::Type{<:Tuple{<:BlockSlices, <:BlockSlices}}) = BlockBandedColumnMajor()
sublayout(::BlockBandedColumnMajor, ::Type{<:Tuple{<:BlockSlices, <:BlockSlice{Block1}}}) = ColumnMajor()
sublayout(::BlockBandedColumnMajor, ::Type{<:Tuple{<:BlockSlices, <:BlockSlice{<:BlockIndexRange1}}}) = ColumnMajor()
sublayout(::BlockBandedColumnMajor, ::Type{<:Tuple{<:BlockSlice{<:BlockIndexRange1}, <:BlockSlice{<:BlockIndexRange1}}}) = ColumnMajor()

isblockbanded(V::SubArray{<:Any,2,<:Any,<:Tuple{<:BlockSlices, <:BlockSlices}}) =
    isblockbanded(parent(V))

sub_materialize(::AbstractBlockBandedLayout, V, _) = BlockBandedMatrix(V)
sub_materialize(::AbstractBlockBandedLayout, V, ::Tuple{<:BlockedUnitRange,<:BlockedUnitRange}) = BlockBandedMatrix(V)
sub_materialize(::BlockLayout{<:AbstractBandedLayout}, V, _) = BlockBandedMatrix(V)
sub_materialize(::BlockLayout{<:AbstractBandedLayout}, V, ::Tuple{<:BlockedUnitRange,<:BlockedUnitRange}) = BlockBandedMatrix(V)

strides(V::SubBlockSkylineMatrix{<:Any,LL,UU,<:Union{BlockRange1,Block1},Block1}) where {LL,UU} =
    (1,parent(V).block_sizes.block_strides[Int(parentindices(V)[2].block)])

function unsafe_convert(::Type{Ptr{T}}, V::SubBlockSkylineMatrix{T,LL,UU,<:Union{BlockRange1,Block1},Block1}) where {T,LL,UU}
    A = parent(V)
    JR = parentindices(V)[2]
    KR = parentindices(V)[1].block
    J = parentindices(V)[2].block
    p = unsafe_convert(Ptr{T}, view(A, first(KR), J))
end

strides(V::SubBlockSkylineMatrix{<:Any,LL,UU,<:BlockRange1,<:BlockIndexRange1}) where {LL,UU} =
    (1,parent(V).block_sizes.block_strides[Int(Block(parentindices(V)[2]))])

function unsafe_convert(::Type{Ptr{T}}, V::SubBlockSkylineMatrix{T,LL,UU,<:BlockRange1,<:BlockIndexRange1}) where {T,LL,UU}
    A = parent(V)
    JR = parentindices(V)[2]
    K = first(parentindices(V)[1].block)
    J = Block(JR)
    K ∈ blockcolsupport(A, J) || throw(ArgumentError("Pointer is only defined when inside blockcolsupport"))
    p = unsafe_convert(Ptr{T}, view(A, K, J))
    p + sizeof(T)*(JR.block.indices[1][1]-1)*stride(V,2)
end

function unsafe_convert(::Type{Ptr{T}}, V::SubBlockSkylineMatrix{T,LL,UU,BlockIndexRange1,BlockIndexRange1}) where {T,LL,UU}
    A = parent(V)
    JR = parentindices(V)[2]
    K = parentindices(V)[1].block.block
    kr = parentindices(V)[1].block.indices[1]
    J = parentindices(V)[2].block.block
    jr = parentindices(V)[2].block.indices[1]
    p = unsafe_convert(Ptr{T}, view(A, K, J))
    p + sizeof(T)*(kr[1]-1 + (jr[1]-1)*stride(V,2))
end

strides(V::SubBlockSkylineMatrix{T,LL,UU,BlockIndexRange1,BlockIndexRange1}) where {T,LL,UU} =
    (1,parent(V).block_sizes.block_strides[Int(parentindices(V)[2].block.block)])

MemoryLayout(V::SubBlockSkylineMatrix{T,LL,UU,BlockIndexRange1,BlockIndexRange1}) where {T,LL,UU} = ColumnMajor()


text:
"""
    parse_loglevel(input_level::AbstractString)::LogLevel
    parse_loglevel(input_level::Integer)::LogLevel

Parse a log level from either an integer or string.

# Examples
    parse_loglevel("info") -> Logging.Info
    parse_loglevel(0) -> LogLevel(0) (== Logging.Info)
"""
function parse_loglevel(input_level::AbstractString)::LogLevel
    level = lowercase(input_level)
    if level == "debug"
        return Logging.Debug
    elseif level == "info"
        return Logging.Info
    elseif level == "warn"
        return Logging.Warn
    elseif level == "error"
        return Logging.Error
    else
        error("loglevel $input_level not recognized")
    end
end

parse_loglevel(input_level::Integer) = LogLevel(input_level)

"Print a log message to a single line, for Delft-FEWS"
function format_message(io::IO, args)::Nothing
    kwargs = IOBuffer()
    for p in args.kwargs
        print(kwargs, string(" |", p[1], " = ", p[2], "|"))
    end
    println(
        io,
        now(),
        " | ",
        args._module,
        " | [",
        args.level,
        "] ",
        args.message,
        String(take!(kwargs)),
    )
    return nothing
end

"Initialize a logger, which is different if `fews_run` is set in the Config."
function init_logger(config::Config; silent = false)::Tuple{TeeLogger,IOStream}
    loglevel = parse_loglevel(get(config, "loglevel", "info"))
    path_log = output_path(config, get(config, "path_log", "log.txt"))
    log_handle = open(path_log, "w")
    fews_run = get(config, "fews_run", false)::Bool

    file_logger = if fews_run
        # Format the log message to be printed on a single line.
        MinLevelLogger(FormatLogger(format_message, log_handle), loglevel)
    else
        # use ConsoleLogger instead of FileLogger to be able to print full stacktraces there
        # https://github.com/JuliaLogging/LoggingExtras.jl/issues/46#issuecomment-803716480
        ConsoleLogger(
            IOContext(log_handle, :compact => false, :limit => false, :color => false),
            loglevel,
            show_limited = false,
        )
    end

    terminal_logger = if silent
        NullLogger()
    else
        # avoid debug information overflowing the terminal, -1 is the level of ProgressLogging
        # avoid double stacktraces by filtering the @error catch in Wflow.run
        EarlyFilteredLogger(
            log -> log.id !== :wflow_run,
            MinLevelLogger(TerminalLogger(), LogLevel(-1)),
        )
    end

    logger = TeeLogger(
        # avoid progresslogger and NCDatasets debug messages in the file
        EarlyFilteredLogger(
            log -> log.group !== :ProgressLogging && log._module != NCDatasets,
            file_logger,
        ),
        terminal_logger,
    )
    with_logger(logger) do
        @debug "Logger initialized." loglevel silent fews_run path_log
    end
    return logger, log_handle
end


text:
"""
# LoggingTestSets

`LoggingTestSet` is an [`AbstractTestSet`](https://docs.julialang.org/en/v1/stdlib/Test/#Creating-Custom-AbstractTestSet-Types)
that logs test results using `@info` and `@error` from the
[Logging][https://docs.julialang.org/en/v1/stdlib/Logging/] module.

```julia
julia> using LoggingTestSets
julia> using Test
julia> @testset LoggingTestSet "Some tests" begin
           @test 1 == 1
           @test 1 == 2
       end

┌ Error: Some tests: Test Failed at REPL[6]:3
│   Expression: 1 == 2
│    Evaluated: 1 == 2
└ @ LoggingTestSets LoggingTestSets/src/LoggingTestSets.jl:49
┌ Info: Test Summary: Some tests
│   passes = 1
│   fails = 1
│   errors = 0
└   broken = 0
Test Summary: | Pass  Fail  Total
Some tests    |    1     1      2
┌ Error: Some tests did not pass: 1 passed, 1 failed, 0 errored, 0 broken.
└ @ LoggingTestSets LoggingTestSets/src/LoggingTestSets.jl:69
ERROR: Some tests did not pass: 1 passed, 1 failed, 0 errored, 0 broken.
```

`TestFileLogger` is an [`AbstractLogger`](https://docs.julialang.org/en/v1/stdlib/Logging/#AbstractLogger-interface)
that writes log messages to a file.

```julia
julia> using LoggingTestSets
julia> using Logging
julia> using LoggingExtras
julia> tee = TeeLogger(global_logger(), TestFileLogger("test.log"))

julia> global_logger(tee)
```
"""
module LoggingTestSets
export LoggingTestSet
export TestFileLogger
export ColumnFileLogger
export ColumnConsoleLogger

using Test
using Test: DefaultTestSet, AbstractTestSet, Fail, Error, scrub_backtrace

using Logging
using LoggingExtras
using DataStructures
using Crayons
using Dates: now, Time

# Logging test results.

struct LoggingTestSet <: AbstractTestSet
    ts::DefaultTestSet
    function LoggingTestSet(name; kw...)
        @info "Test Set: $name"
        new(DefaultTestSet(name; kw...))
    end
end

function Test.record(ts::LoggingTestSet, t::Union{Fail, Error})

    io = IOBuffer()

    begin # copied from: stdlib/Test/src/Test.jl https://git.io/JqZQk
        print(io, ts.ts.description, ": ")
        # don't print for interrupted tests
        if !(t isa Error) || t.test_type !== :test_interrupted
            print(io, t)
            if t isa Error # if not gets printed in the show method
                Base.show_backtrace(io, scrub_backtrace(backtrace()))
            end
            println(io)
        end
    end

    @error String(take!(io))

    push!(ts.ts.results, t)
end

Test.record(ts::LoggingTestSet, args...) = Test.record(ts.ts, args...)

function Test.finish(ts::LoggingTestSet)

    # Copied from: stdlib/Test/src/Test.jl https://git.io/JqZ70
    np, nf, ne, nb, ncp, ncf, nce, ncb = Test.get_test_counts(ts.ts)
    passes = np + ncp
    fails  = nf + ncf
    errors = ne + nce
    broken = nb + ncb

    if fails + errors + broken == 0
        @info "Test Set: $(ts.ts.description) -- All tests passed."
    else
        @info "Test Summary: $(ts.ts.description)" passes fails errors broken
    end

    try
        Test.finish(ts.ts)
    catch err
        @error err
        rethrow(err)
    end
end

# Logging Failure Context.

struct TestFailContextLogger{T <: AbstractLogger} <: AbstractLogger
    logger::T
    buffer::CircularBuffer{Tuple}
    TestFailContextLogger(l::T) where T = new{T}(l, buffer::CircularBuffer{Tuple}(5))
end

function Logging.handle_message(logger::TestFailContextLogger,
                                level, message, _module, group, id, file, line;
                                kwargs...)
    if !comp_shouldlog(logger.logger, level, _module, group, id)
        return
    end

    args = (level, message, _module, group, id, file, line),
    if contains(message, "Test Failed")
        while !isempty(logger.buffer)
            a, k = popfirst!(logger.buffer)
            handle_message(logger.logger, a...; k...)
        end
        handle_message(logger.logger, args...; kwargs...)
    else
        push!(logger.buffer, (args, kwargs))
    end
end

Logging.shouldlog(logger::TestFailContextLogger, args...) = true
Logging.min_enabled_level(logger::TestFailContextLogger) = Logging.BelowMinLevel
Logging.catch_exceptions(logger::TestFailContextLogger) = catch_exceptions(logger.logger)

# Logging to file.

struct TestFileLogger <: AbstractLogger
    io::IO
    TestFileLogger(filename) = new(open(filename, append=true))
end

Logging.shouldlog(::TestFileLogger, args...) = true
Logging.min_enabled_level(::TestFileLogger) = Logging.BelowMinLevel
Logging.catch_exceptions(::TestFileLogger) = false

function Logging.handle_message(l::TestFileLogger,
                                level, message, _module, group, id, file, line;
                                kwargs...)

    # Copied from: base/logging.jl https://git.io/JqZ5Q
    msglines = split(chomp(string(message)), '\n')

    # Copied from: ConsoleLogger.jl https://git.io/JqZdY
    valbuf=IOBuffer()
    for (key,val) in pairs(kwargs)
        Logging.showvalue(valbuf, val)
        vallines = split(String(take!(valbuf)), '\n')
        if length(vallines) == 1
            push!(msglines, "$key = $(vallines[1])")
        else
            push!(msglines, "$key =")
            append!(msglines, vallines)
        end
    end

    # Copied from: base/logging.jl https://git.io/JqZ5Q
    println(l.io, "┌ ", string(level), " ", now(),  ": ", msglines[1])
    for i in 2:length(msglines)
        println(l.io, "│ ", msglines[i])
    end
    println(l.io, "└ @ ", _module, " ", basename(file), ":", line)
    flush(l.io)
    nothing
end

# Logging to file with columns.

const default_width = 160

struct ColumnFileLogger <: AbstractLogger
    io::IO
    width::Int
    ColumnFileLogger(io::IO, width) = new(io, width)
    ColumnFileLogger(filename; width=default_width) =
        new(open(filename, append=true), width)
end

function ColumnConsoleLogger()
    display_height, display_width = displaysize(stdout)
    return ColumnFileLogger(stdout, display_width)
end

Logging.shouldlog(::ColumnFileLogger, args...) = true
Logging.min_enabled_level(::ColumnFileLogger) = Logging.BelowMinLevel
Logging.catch_exceptions(::ColumnFileLogger) = false

function Logging.handle_message(l::ColumnFileLogger, args...; kwargs...)
    write(l.io, column_format_log(l.width, args...; kwargs...))
    flush(l.io)
    nothing
end

function column_format_log(width,
                           level, message, _module, group, id, file, line;
                           kwargs...)

    # Split message into lines.
    # Copied from: base/logging.jl https://git.io/JqZ5Q
    msglines = split(chomp(string(message)), '\n')

    # Append lines for extra log args.
    # Copied from: ConsoleLogger.jl https://git.io/JqZdY
    valbuf=IOBuffer()
    for (key,val) in kwargs
        Logging.showvalue(valbuf, val)
        vallines = split(String(take!(valbuf)), '\n')
        if length(vallines) == 1
            push!(msglines, "$key = $(vallines[1])")
        else
            push!(msglines, "$key =")
            append!(msglines, vallines)
        end
    end

    # Prepare prefix, file:line, and stuffix.
    prefix = length(msglines) > 1 ? "┌ " : "[ "
    fileline = string(basename(file), ":", rpad(string(line), 4))
    if group != _module
        fileline = "$group:$fileline"
    end
    suffix = length(msglines) == 1 ? fileline : ""

    # Calculate column widths and padding.
    widths = [2, 5, 1, # level
                12, 1, # time
                20, 3, # module
                textwidth(msglines[1]),
                textwidth(suffix)]
    pad = max(0, width - sum(widths))

    # Colours.
    bg = module_background(string(_module))
    lc = level_color(level)
    dim = dim_color

    buffer = IOBuffer()

    # Print the first line:
    #
    # ┌ Level 10:11:05.339  Module │ Log Message [ ----- pad ----- ] file:line
    print(buffer,
          "\n",
          bg,
          lc,  prefix, lpad(string(level), widths[2]),  inv(lc),  " ",
          dim,         rpad(Time(now()),   widths[4]),  inv(dim), " ",
                       lpad(_module,       widths[6]),            " │ ",
                       msglines[1],        repeat(" ", pad),
          dim,         suffix,                          inv(dim),
          inv(bg)
    )

    # Print the other lines.
    #
    # │ [ -------- pad1 -------- ] │ msglines[2] [ ---------- pad2 ----------]
    # │ [ -------- pad1 -------- ] │ msglines[3...] [ -------- pad2 ---------]
    # └------------------------------------------------------------- file:line
    if length(msglines) > 1
        # Copied from: base/logging.jl https://git.io/JqZ5Q
        for i in 2:length(msglines)
            pad1 = sum(widths[2:6])
            pad2 = max(0, width - pad1 - 5 - textwidth(msglines[i]))
            print(buffer, "\n",
                          bg,
                          lc, "│ ", inv(lc),
                          repeat(" ", pad1), " │ ",
                          msglines[i],
                          repeat(" ", pad2),
                          inv(bg))
        end
        pad = max(0, width - length(fileline) - 2)
        print(buffer, "\n",
                      bg,
                      lc, "└", repeat("-", pad), inv(lc), " ",
                      dim, fileline, inv(dim),
                      inv(bg))
    end

    return String(take!(buffer))
end

# Colours.

const dim_color = crayon"fg:246"

level_color(level) = level < Logging.Info  ? crayon"blue" :
                     level < Logging.Warn  ? crayon"cyan" :
                     level < Logging.Error ? crayon"yellow" :
                                             crayon"light_red"

const module_backgrounds = Dict{String, Crayon}(
    "LoggingTestSets" => crayon"bg:16"
)

function module_background(m)
    bg = get(module_backgrounds, m, nothing)
    if bg != nothing
        return bg
    end
    bg = popfirst!(background_colors)
    module_backgrounds[m] = bg
    return bg
end

const background_colors =
    Iterators.Stateful(
    Iterators.Cycle([
    #crayon"bg:16",
    #crayon"bg:18",
    crayon"bg:233",
    crayon"bg:234",
    crayon"bg:235",
    crayon"bg:236",
    crayon"bg:237",
    crayon"bg:238",
    crayon"bg:239",
    crayon"bg:240",
    crayon"bg:241",
    crayon"bg:242",
    crayon"bg:17",
    crayon"bg:18",
    crayon"bg:19",
    crayon"bg:20",
    crayon"bg:21",
    crayon"bg:22",
    crayon"bg:23",
    crayon"bg:24",
    crayon"bg:25",
    crayon"bg:58",
    crayon"bg:59",
    crayon"bg:60",]))

# Documentation.

readme() = Docs.doc(@__MODULE__)

end # module


text:
using IntervalTrees
using DataFrames

# FeatureTable:
#   Description can be any metadata e.g. it could note what each column means, or
#   describe what organism this is about, or any other key - value.
#

type FeatureTable
  description::Dict{ASCIIString,Any}
  dataframe::DataFrame
end

#  Feature:  An interval, point, possibly an area??
#    represented by a single row in a dataframe - it might have a start and a stop a seqid a name and so on.
type Feature
   feature::DataFrame
   function Feature(start,stop,region_id="REGION1"; args...)
       return DataFrame(region_id=region_id,start=start,stop=stop)
       # add in other columns
   end
end

# Features are features of something, this something define as "Region".
#
#  A region is any range of something, probably tied to a sequence often, such as
#  a gene (although that isn't tied necessarily to one set of coordinates)
#  or an entire chromosome
#
#  The Features type indexes Features by Region (i.e. a region_id such as chr).
#

#
# Features is for features that are indexed by a sequence_id
# SF.features["seq1"] or SF.features["chr1"]

type Features
  description::Dict{ASCIIString,Any}
  features::Dict{Any,Any}

  function Features(FT::FeatureTable; start_field = :start, stop_field = :stop, region_id_field = :chr)
     region_features_dict = Dict()
     DF = FT.dataframe
     add_dataframe_to_interval_tree(DF,start_field=start_field,stop_field=stop_field,region_id_field=region_id_field)
     desc = FT.description
     if ! haskey(desc, "organism")
        desc["organism"] = ""  #add in organism default if it doesnt exist
     end
     new( desc, region_features_dict )
  end

  function Features(DF::DataFrame; start_field = :start, stop_field = :stop, region_id_field = :chr, regions=false )
    desc = Dict{ASCIIString,Any}()
    if regions
        region_features_dict = add_dataframe_to_interval_tree(DF,start_field=:notanything,stop_field=:notanything,region_id_field=region_id_field)
        add_dataframe_to_description(DF,description=desc, region_id_field=region_id_field)
        desc["organism"]=""
     else
       region_features_dict = add_dataframe_to_interval_tree(DF,start_field=start_field,stop_field=stop_field,region_id_field=region_id_field)
       desc["organism"]=""
     end
     new( desc, region_features_dict )
  end

  # Empty Feautures
  function Features()
      desc = Dict{ASCIIString,Any}()
      region_features_dict=Dict()
      new( desc, region_features_dict )
  end
end

function add_dataframe_to_description(DF::DataFrame; description=Dict{ASCIIString,Any}(),region_id_field = :chr)
    nrows = size(DF,1)
        for row = 1:nrows
          feature = DF[row, :]
          if( haskey(feature, region_id_field ))
             region_id  =  feature[region_id_field][1]
          else
             region_id = "REGION1"
          end
          if ! haskey( description, region_id )
              description[region_id] = feature
          else
              error("Regions must unique - region_id $region_id occurs at least twice")
          end
     end
    return description
end

function add_dataframe_to_interval_tree(DF::DataFrame; region_features_dict=Dict(), start_field = :start, stop_field = :stop, region_id_field = :chr )
        #region_features_dict = Dict()
        nrows = size(DF,1)
        (start_nothing_count,stop_nothing_count) = (0,0)
        for row = 1:nrows
          feature = DF[row, :]

          if(haskey( feature,start_field ) )
             start = feature[start_field][1]
          else
             start_nothing_count +=1
             start = nothing
          end

          if(haskey(feature, stop_field ))
              stop  = feature[stop_field][1]
          else
             stop_nothing_count +=1
             stop = nothing
          end

          if( haskey(feature, :name) )
            feature_id = feature[:name][1]
          else
            feature_id = "FEATURE$row"
          end

          if( haskey(feature, region_id_field ))
             region_id  =  feature[region_id_field][1]
          else
             region_id = "REGION1"
          end
          insertintervals!(region_features_dict, region_id, start, stop,feature )
     end
     if stop_nothing_count > 0 || start_nothing_count > 0
         warn( "Start pos set to nothing: $start_nothing_count. Stop pos set to nothing $stop_nothing_count")
     end
     return region_features_dict
end

# takes one set of features and adds them to an already existing
# container of features
function insertfeatures!( target::Features, features::Features; start_field = :start, stop_field = :stop, region_id_field = :chr )
end

function add_dataframe_to_features!(target::Features, DF::DataFrame; start_field = :start, stop_field = :stop, region_id_field = :chr)
    add_dataframe_to_interval_tree( DF, region_features_dict = target.features, start_field=start_field, stop_field=stop_field,region_id_field=region_id_field )
end

function insertintervals!( region_features_dict, region_id, start, stop, value::DataFrame)
    if haskey(region_features_dict, region_id)
        interval_tree = region_features_dict[region_id]
        if start != nothing && stop != nothing
             interval_tree[(start,stop)] = value
        end
        #interval_tree[(start,stop)] = value
        #println("inserted $start $stop")
    else
        try
           interval_tree = IntervalTree{Int32,DataFrame}()
           if start != nothing && stop != nothing
             interval_tree[(start,stop)] = value
           end
           region_features_dict[region_id]  = interval_tree
        catch e
           println("Exception chr: $region_id, start: $start, stop: $stop, value $value ")
           println( join( [  "Exception, ", typeof(region_id), " ", typeof(start), " ", typeof(stop) ], "\t" ) )
           error(e)
        end
    end
end

## --- IN PROGRESS ----##
# A more memory efficient version of Features
# Don't use dataframes as the value
# Interval Tree positions should be Int32
# See how that goes
type FeaturesSlim
  description::Dict{ASCIIString,Any}
  features::Dict{Any,Any}
  function FeaturesSlim()
      desc = Dict{ASCIIString,Any}()
      region_features_dict=Dict()
      new( desc, region_features_dict )
  end
end

function insertfeatures!(target::FeaturesSlim, region_id,start,stop,value)
   region_features_dict = target.features
   if haskey(region_features_dict, region_id)
        interval_tree = region_features_dict[region_id]
        if start != nothing && stop != nothing
             interval_tree[(start,stop)] = value
        end
    else
        try
           interval_tree = IntervalTree{Int32,Int64}()
           if start != nothing && stop != nothing
             interval_tree[(start,stop)] = value
           end
           region_features_dict[region_id]  = interval_tree
        catch e
           println("Exception chr: $region_id, start: $start, stop: $stop, value $value ")
           println( join( [  "Exception, ", typeof(region_id), " ", typeof(start), " ", typeof(stop) ], "\t" ) )
           error(e)
        end
    end
end

# Further thoughts
# ----------------
# A feature is a region of some size
# and it could have multiple regions
# there might also be feature_type - i.e. gene? gene_non_coding? - it gets all ontological then.
#
# Are Regions just Features or vice versa
# Is it necessary to make the distinction


text:
using Test
using OMEinsum
using OMEinsum: StaticEinCode, DynamicEinCode
using Zygote

@doc raw"
    bpcheck(f, args...; η = 1e-5, verbose=false)
returns a `Bool` indicating whether Zygote calculates the gradient of `f(args...) -> scalar`
correctly using the relation `f(x - ηg) ≈ f(x) - η|g|²` with a relative tolerance
of 1e-2 and an absolute tolerance of 1e-8.
If `verbose=true`, print `f(x) - f(x - ηg)`and `η|g|²`.

# example

```jldoctest; setup = :(using OMEinsum)
julia> using OMEinsum: bpcheck

julia> a, b = rand(2,2), rand(2,2);

julia> bpcheck(sum ∘ ein\"ij,jk -> ik\", a, b)
true
```
"
function bpcheck(f, args...; η = 1e-5, verbose = false)
    g = gradient(f, args...)
    all(==(nothing), g) && error()
    dy_ref = 0
    for x in g
        x === nothing && continue
        x isa Tuple && (dy_ref += η * mapreduce(y -> y == nothing ? 0 : sum(abs2,y), +, x))
        x isa AbstractArray && (dy_ref += η * sum(abs2,x))
    end
    dy = f(args...) - f([gi === nothing ? arg : arg .- η .* gi for (arg, gi) in zip(args,g)]...)

    verbose && @show dy
    verbose && @show dy_ref

    isapprox(dy, dy_ref, rtol=1e-2, atol=1e-8)
end

@testset "einsum bp" begin
    for T in (Float64, ComplexF64)
        @testset "$T" begin
            # matrix and vector multiplication
            a,b,c = rand(T,2,2), rand(T,2,2), rand(T,2,2)
            v = rand(T,2)
            t = randn(2,2,2,2)
            @test bpcheck( (a,b,c) -> einsum(EinCode(((1,2),(2,3),(3,4)), (1,4)), (a,b,c)) |> abs ∘ sum ,a,b,c)
            @test bpcheck( (a,b,c) -> einsum(EinCode(((1,2),(2,3),(3,4)), (4,1)), (a,b,c)) |> abs ∘ sum ,a,b,c)
            @test bpcheck((a,v) -> einsum(EinCode(((1,2),(2,)), (1,)), (a,v)) |> abs ∘ sum , a, v)
            @test bpcheck((a,v) -> einsum(StaticEinCode{((1,2),(2,)), (1,)}(), (a,v)) |> abs ∘ sum , a, v)
            @test bpcheck((a,t,v) -> ein"(ab,abcd),c->ad"(a,t,v) |> abs ∘ sum ,a,T.(t),v)

            # contract to 0-dim array
            @test bpcheck((a,b) -> einsum(EinCode(((1,2),(1,2)), ()), (a,b)) |> abs ∘ sum , a,b)

            # trace
            @test bpcheck(a -> einsum(EinCode(((1,1),), ()), (a,)) |> abs ∘ sum, a)
            aa = rand(T,2,4,4,2)
            @test bpcheck(aa -> einsum(EinCode(((1,2,2,1),), ()), (aa,)) |> abs ∘ sum, aa)

            # partial trace
            @test bpcheck(aa -> einsum(EinCode(((1,2,2,3),), (1,3)), (aa,)) |> abs ∘ sum, aa)

            # diag
            @test bpcheck(aa -> einsum(EinCode(((1,2,2,3),), (1,2,3)), (aa,)) |> abs ∘ sum, aa)

            # permutation
            @test bpcheck(a -> einsum(EinCode(((1,2),), (2,1)), (a,)) |> abs ∘ sum, a)
            @test bpcheck(t -> einsum(EinCode(((1,2,3,4),),(2,3,1,4)), (t,)) |> abs ∘ sum, t)

            # tensor contraction
            @test bpcheck((t,a) -> einsum(EinCode(((1,2,3,4), (2,3)), (1,4)), (t,a)) |> abs ∘ sum, t,a)
            @test bpcheck((t,a) -> einsum(EinCode(((4,3,2,1), (2,3)), (1,4)), (t,a)) |> abs ∘ sum, t,a)

            # star-contraction
            @test bpcheck((a,b,c) -> einsum(EinCode(((1,2),(1,3),(1,4)), (2,3,4)), (a,b,c)) |> abs ∘ sum, a,b,c)

            # star and contract
            @test bpcheck((a,b,c) -> einsum(EinCode(((1,2),(1,2),(1,3)), (3,)), (a,b,c)) |> abs ∘ sum, a,b,c)

            # index-sum
            a3 = rand(T,2,2,2)
            @test bpcheck(a -> einsum(EinCode(((1,2,3),),(1,2)),(a,)) |> abs ∘ sum, a3)

            # Hadamard product
            @test bpcheck((a,b) -> einsum(EinCode(((1,2),(1,2)), (1,2)), (a,b)) |> abs ∘ sum, a, b)

            # Outer
            @test bpcheck((a,b) -> einsum(EinCode(((1,2),(3,4)),(1,2,3,4)),(a,b)) |> abs ∘ sum, a, b)
        end
    end
end

@testset "gradient type check" begin
    array_match(x, y) = typeof(x) == typeof(y) && size(x) == size(y)
    a = randn(3,3)
    b = randn(3,3)
    @test array_match(gradient(a->einsum(EinCode(((1,2), (2,1)), ()), (a, b))[] |> abs, a)[1], a)
    b = randn(ComplexF64,3,3)
    @test array_match(gradient(a->einsum(EinCode(((1,2), (2,1)), ()), (a, b))[] |> abs, a)[1], a)
    a = randn(ComplexF64,3,3)
    @test array_match(gradient(a->einsum(EinCode(((1,2), (2,3)), ()), (a, b))[] |> abs, a)[1], a)
    b = randn(3,3)
    @test array_match(gradient(a->einsum(EinCode(((1,2), (2,3)), ()), (a, b))[] |> abs, a)[1], a)
end

@testset "string-specification" begin
    a,b,c = rand(2,2), rand(2,2), rand(2,2)
    v = rand(2)
    @test bpcheck((a,b,c) -> einsum(ein"ij,jk,kl -> il", (a,b,c)) |> abs ∘ sum ,a,b,c)
    @test bpcheck((a,b,c) -> einsum(ein"ij,jk,kl -> li", (a,b,c)) |> abs ∘ sum ,a,b,c)
    @test bpcheck((a,v) -> einsum(ein"ij,j -> i", (a,v)) |> abs ∘ sum , a, v)
end

@testset "sequence specification" begin
    a, b, c = rand(2,2), rand(2,2), rand(2,2)
    @test all(gradient(sum ∘ ein"ij,jk,kl -> il", a, b, c) .≈
          gradient(sum ∘ ein"(ij,(jk,kl)) -> il", a, b, c))
end


text:
const DDMMatrix{T} = DDMArray{T,2}

"""
    CartesianDDMMatrix{T,C,A}

Objects of `CartesianDDMMatrix` behave like a normal matrix, except
that the storage is not stored continuously but by DDM subdomain.

For now, only Dirichlet matrices are considered.

"""
struct CartesianDDMMatrix{T,C,A} <: DDMMatrix{T}
    context::C
    parent::A
end

parent(x::CartesianDDMMatrix) = getproperty(x, :parent)
getcontext(x::CartesianDDMMatrix) = getproperty(x, :context)

function size(x::CartesianDDMMatrix)
    (; context, parent) = x
    (; dims) = context
    ntuple(i -> mapreduce(getdof, *, dims), 2)
end

"""
    getindex(x::CartesianDDMMatrix, i)

Decomposed matrix along rows.
Block-diagonal matrix.

"""
function getindex(x::CartesianDDMMatrix, i::Int, j::Int)
    @boundscheck checkbounds(x, i, j)

    (; context, parent) = x

    (; dims) = context
    gi = getindex(CartesianIndices(getdof.(dims)), i)

    nocontext = removeoverlap(context)
    noglb, _ = ranges(nocontext)

    index = findfirst(noglb) do el
        in(gi, CartesianIndices(el))
    end

    glb, _, lcl = ranges(context)
    domain = CartesianIndices(glb[index])

    li = findfirst(domain) do li
        li == gi
    end

    ii = LinearIndices(CartesianIndices(lcl[index]))[li]

    gj = getindex(CartesianIndices(getdof.(dims)), j)

    lj = findfirst(domain) do lj
        lj == gj
    end

    # something fishy when over = 1
    isnothing(lj) && return zero(eltype(x))

    jj = LinearIndices(CartesianIndices(lcl[index]))[lj]

    parent[index][ii, jj]
end

function decompose(context::CartesianDDMContext, a::AbstractMatrix)
    (; dims) = context

    glb, _ = ranges(context)

    cindices = CartesianIndices(getdof.(dims))
    lindices = LinearIndices(cindices)

    parent = map(glb) do el
        domain = reshape(lindices[CartesianIndices(el)], :)
        a[domain, domain]
    end

    CartesianDDMMatrix{eltype(a),typeof(context),typeof(parent)}(context, parent)
end

function (*)(A::CartesianDDMMatrix, x::CartesianDDMVector)
    context = A.context
    @assert context == x.context

    parent = A.parent .* x.parent

    T = eltype(eltype(parent))

    CartesianDDMVector{T,typeof(context),typeof(parent)}(context, parent)
end


text:
export Impenetrable, Penetrable, NoSlip, FreeSlip, KinematicStress

abstract type VelocityBC end
abstract type VelocityDragBC end

"""
    Impenetrable(drag::VelocityDragBC) :: VelocityBC

Defines an impenetrable wall model for velocity. This implies:
  - no flow in the direction normal to the boundary, and
  - flow parallel to the boundary is subject to the `drag` condition.
"""
struct Impenetrable{D <: VelocityDragBC} <: VelocityBC
    drag::D
end

"""
    Penetrable(drag::VelocityDragBC) :: VelocityBC

Defines an penetrable wall model for velocity. This implies:
  - no constraint on flow in the direction normal to the boundary, and
  - flow parallel to the boundary is subject to the `drag` condition.
"""
struct Penetrable{D <: VelocityDragBC} <: VelocityBC
    drag::D
end

"""
    NoSlip() :: VelocityDragBC

Zero velocity at the boundary.
"""
struct NoSlip <: VelocityDragBC end

"""
    ocean_velocity_boundary_state!(::NumericalFluxFirstOrder, ::Impenetrable{NoSlip}, ::HBModel)

apply no slip boundary condition for velocity
sets reflective ghost point
"""

function ocean_velocity_boundary_state!(
    nf::NumericalFluxFirstOrder,
    bc_velocity::Impenetrable{NoSlip},
    ocean,
    Q⁺,
    A⁺,
    n⁻,
    Q⁻,
    A⁻,
    t,
)
    Q⁺.u = -Q⁻.u
    A⁺.w = -A⁻.w

    return nothing
end

"""
    ocean_velocity_boundary_state!(::NumericalFluxGradient, ::Impenetrable{NoSlip}, ::HBModel)

apply no slip boundary condition for velocity
set numerical flux to zero for u
"""
function ocean_velocity_boundary_state!(
    nf::NumericalFluxGradient,
    bc_velocity::Impenetrable{NoSlip},
    ocean,
    Q⁺,
    A⁺,
    n⁻,
    Q⁻,
    A⁻,
    t,
)
    FT = eltype(Q⁺)
    Q⁺.u = SVector(-zero(FT), -zero(FT))
    A⁺.w = -zero(FT)

    return nothing
end

"""
    ocean_velocity_boundary_state!(::NumericalFluxSecondOrder, ::Impenetrable{NoSlip}, ::HBModel)

apply no slip boundary condition for velocity
sets ghost point to have no numerical flux on the boundary for u
"""
@inline function ocean_velocity_boundary_state!(
    nf::NumericalFluxSecondOrder,
    bc_velocity::Impenetrable{NoSlip},
    ocean,
    Q⁺,
    D⁺,
    A⁺,
    n⁻,
    Q⁻,
    D⁻,
    A⁻,
    t,
)
    Q⁺.u = -Q⁻.u
    A⁺.w = -A⁻.w
    D⁺.ν∇u = D⁻.ν∇u

    return nothing
end

"""
    FreeSlip() :: VelocityDragBC

No surface drag on velocity parallel to the boundary.
"""
struct FreeSlip <: VelocityDragBC end

"""
    ocean_velocity_boundary_state!(::NumericalFluxFirstOrder, ::Impenetrable{FreeSlip}, ::HBModel)

apply free slip boundary condition for velocity
sets reflective ghost point
"""
function ocean_velocity_boundary_state!(
    nf::NumericalFluxFirstOrder,
    bc_velocity::Impenetrable{FreeSlip},
    ocean,
    Q⁺,
    A⁺,
    n⁻,
    Q⁻,
    A⁻,
    t,
)
    v⁻ = @SVector [Q⁻.u[1], Q⁻.u[2], A⁻.w]
    v⁺ = v⁻ - 2 * n⁻ ⋅ v⁻ .* SVector(n⁻)
    Q⁺.u = @SVector [v⁺[1], v⁺[2]]
    A⁺.w = v⁺[3]

    return nothing
end

"""
    ocean_velocity_boundary_state!(::NumericalFluxGradient, ::Impenetrable{FreeSlip}, ::HBModel)

apply free slip boundary condition for velocity
sets non-reflective ghost point
"""
function ocean_velocity_boundary_state!(
    nf::NumericalFluxGradient,
    bc_velocity::Impenetrable{FreeSlip},
    ocean,
    Q⁺,
    A⁺,
    n⁻,
    Q⁻,
    A⁻,
    t,
)
    v⁻ = @SVector [Q⁻.u[1], Q⁻.u[2], A⁻.w]
    v⁺ = v⁻ - n⁻ ⋅ v⁻ .* SVector(n⁻)
    Q⁺.u = @SVector [v⁺[1], v⁺[2]]
    A⁺.w = v⁺[3]

    return nothing
end

"""
    ocean_velocity_normal_boundary_flux_second_order!(::NumericalFluxSecondOrder, ::Impenetrable{FreeSlip}, ::HBModel)

apply free slip boundary condition for velocity
apply zero numerical flux in the normal direction
"""
function ocean_velocity_boundary_state!(
    nf::NumericalFluxSecondOrder,
    bc_velocity::Impenetrable{FreeSlip},
    ocean,
    Q⁺,
    D⁺,
    A⁺,
    n⁻,
    Q⁻,
    D⁻,
    A⁻,
    t,
)
    Q⁺.u = Q⁻.u
    A⁺.w = A⁻.w
    D⁺.ν∇u = n⁻ * (@SVector [-0, -0])'

    return nothing
end

"""
    ocean_velocity_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::Penetrable{FreeSlip}, ::HBModel)

apply free slip boundary condition for velocity
sets non-reflective ghost point
"""
function ocean_velocity_boundary_state!(
    nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},
    bc_velocity::Penetrable{FreeSlip},
    ocean,
    args...,
)
    return nothing
end

"""
    ocean_velocity_boundary_state!(::NumericalFluxSecondOrder, ::Penetrable{FreeSlip}, ::HBModel)

apply free slip boundary condition for velocity
sets non-reflective ghost point
"""
function ocean_velocity_boundary_state!(
    nf::NumericalFluxSecondOrder,
    bc_velocity::Penetrable{FreeSlip},
    ocean,
    args...,
)
    Q⁺.u = Q⁻.u
    A⁺.w = A⁻.w
    D⁺.ν∇u = n⁻ * (@SVector [-0, -0])'

    return nothing
end

"""
    KinematicStress(stress) :: VelocityDragBC

Applies the specified kinematic stress on velocity normal to the boundary.
Prescribe the net inward kinematic stress across the boundary by `stress`,
a function with signature `stress(problem, state, aux, t)`, returning the flux (in m²/s²).
"""
struct KinematicStress <: VelocityDragBC end

"""
    ocean_velocity_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::Impenetrable{KinematicStress}, ::HBModel)

apply kinematic stress boundary condition for velocity
applies free slip conditions for first-order and gradient fluxes
"""
function ocean_velocity_boundary_state!(
    nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},
    bc_velocity::Impenetrable{KinematicStress},
    ocean,
    args...,
)
    return ocean_velocity_boundary_state!(
        nf,
        Impenetrable(FreeSlip()),
        ocean,
        args...,
    )
end

"""
    ocean_velocity_boundary_state!(::NumericalFluxSecondOrder, ::Impenetrable{KinematicStress}, ::HBModel)

apply kinematic stress boundary condition for velocity
sets ghost point to have specified flux on the boundary for ν∇u
"""
@inline function ocean_velocity_boundary_state!(
    nf::NumericalFluxSecondOrder,
    bc_velocity::Impenetrable{KinematicStress},
    ocean,
    Q⁺,
    D⁺,
    A⁺,
    n⁻,
    Q⁻,
    D⁻,
    A⁻,
    t,
)
    Q⁺.u = Q⁻.u
    D⁺.ν∇u = n⁻ * kinematic_stress(ocean.problem, A⁻.y, ocean.ρₒ)'

    return nothing
end

"""
    ocean_velocity_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::Penetrable{KinematicStress}, ::HBModel)

apply kinematic stress boundary condition for velocity
applies free slip conditions for first-order and gradient fluxes
"""
function ocean_velocity_boundary_state!(
    nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},
    bc_velocity::Penetrable{KinematicStress},
    ocean,
    args...,
)
    return ocean_velocity_boundary_state!(
        nf,
        Penetrable(FreeSlip()),
        ocean,
        args...,
    )
end

"""
    ocean_velocity_boundary_state!(::NumericalFluxSecondOrder, ::Penetrable{KinematicStress}, ::HBModel)

apply kinematic stress boundary condition for velocity
sets ghost point to have specified flux on the boundary for ν∇u
"""
@inline function ocean_velocity_boundary_state!(
    nf::NumericalFluxSecondOrder,
    bc_velocity::Penetrable{KinematicStress},
    ocean,
    Q⁺,
    D⁺,
    A⁺,
    n⁻,
    Q⁻,
    D⁻,
    A⁻,
    t,
)
    Q⁺.u = Q⁻.u
    D⁺.ν∇u = n⁻ * kinematic_stress(ocean.problem, A⁻.y, ocean.ρₒ)'

    return nothing
end

max_stars_repo_name: sandreza/CLIMA

max_issues_repo_name: sandreza/CLIMA

max_forks_repo_name: sandreza/CLIMA


text:
"""
```
imge = imcorner(img; [method], [border], [blockSize], [k])
```
Performs corner detection, using either the Harris method or the Shi-Tomasi
method. `method` is `'harris'` (default) or `'shi-tomasi'`. `'border'` is the border
mode used in computing the gradient (default is `'replicate'`). `'blockSize'` is
the size of the box filter kernel (default of `3`). `k` is the Harris free parameter
(default of `0.04`). It is only used when method is set to `'harris'`.

"""
function imcorner(img::AbstractArray; method::AbstractString="harris", border::AbstractString="replicate", blockSize::Int=3, k::Float64=0.04)

    (grad_x,grad_y) = imgradients(img, "sobel",border);

    cov_xx = grad_x .* grad_x;
    cov_xy = grad_x .* grad_y;
    cov_yy = grad_y .* grad_y;

    box_filter_kernel = (1/(blockSize*blockSize)) * ones(blockSize,blockSize);

    cov_xx = imfilter(cov_xx,box_filter_kernel);
    cov_xy = imfilter(cov_xy,box_filter_kernel);
    cov_yy = imfilter(cov_yy,box_filter_kernel);

    if method == "harris"
        corner = (cov_xx.*cov_yy - cov_xy.*cov_xy - k*(cov_xx + cov_yy).*(cov_xx + cov_yy));
    elseif method == "shi-tomasi"
        cov_yy = 0.5*cov_yy;
        cov_xx = 0.5*cov_xx;
        corner = ((cov_xx + cov_yy) - sqrt((cov_xx - cov_yy).*(cov_xx - cov_yy) + cov_xy.*cov_xy));
    else
        error("Unknown corner method: $method")
    end

    return corner
end


text:
# compiler driver and main interface

# (::CompilerContext)
const compile_hook = Ref{Union{Nothing,Function}}(nothing)

"""
    compile(dev::CuDevice, f, tt; kwargs...)

Compile a function `f` invoked with types `tt` for device `dev`, returning the compiled
function module respectively of type `CuFuction` and `CuModule`.

For a list of supported keyword arguments, refer to the documentation of
[`cufunction`](@ref).
"""
function compile(dev::CuDevice, @nospecialize(f), @nospecialize(tt);
                 strip_ir_metadata::Bool=false, kwargs...)
    CUDAnative.configured || error("CUDAnative.jl has not been configured; cannot JIT code.")
    isa(f, Core.Function) || throw(ArgumentError("Kernel argument to `compile` should be a function."))

    ctx = CompilerContext(f, tt, supported_capability(dev), #=kernel=# true; kwargs...)

    if compile_hook[] != nothing
        global globalUnique
        previous_globalUnique = globalUnique
        compile_hook[](ctx)
        globalUnique = previous_globalUnique
    end

    (module_asm, module_entry) = compile(ctx; strip_ir_metadata=strip_ir_metadata)

    # enable debug options based on Julia's debug setting
    jit_options = Dict{CUDAdrv.CUjit_option,Any}()
    if Base.JLOptions().debug_level == 1
        jit_options[CUDAdrv.GENERATE_LINE_INFO] = true
    elseif Base.JLOptions().debug_level >= 2
        jit_options[CUDAdrv.GENERATE_DEBUG_INFO] = true
    end
    cuda_mod = CuModule(module_asm, jit_options)
    cuda_fun = CuFunction(cuda_mod, module_entry)

    return cuda_fun, cuda_mod
end

# Compile a function to PTX, returning the assembly and an entry point.
# FIXME: this pipeline should be partially reusable from eg. code_llvm
#        also, does the kernel argument belong in the compiler context?
function compile(ctx::CompilerContext; strip_ir_metadata::Bool=false)
    ## high-level code generation (Julia AST)

    @debug "(Re)compiling function" ctx

    check_method(ctx)

    ## low-level code generation (LLVM IR)

    mod, entry = irgen(ctx)

    need_library(lib) = any(f -> isdeclaration(f) &&
                                 intrinsic_id(f) == 0 &&
                                 haskey(functions(lib), LLVM.name(f)),
                            functions(mod))

    libdevice = load_libdevice(ctx.cap)
    if need_library(libdevice)
        link_libdevice!(ctx, mod, libdevice)
    end

    # optimize the IR
    entry = optimize!(ctx, mod, entry)

    runtime = load_runtime(ctx.cap)
    if need_library(runtime)
        link_library!(ctx, mod, runtime)
    end

    check_invocation(ctx, entry)

    # check generated IR
    check_ir(ctx, mod)
    verify(mod)

    if strip_ir_metadata
        strip_debuginfo!(mod)
    end

    ## machine code generation (PTX assembly)

    module_asm = mcgen(ctx, mod, entry)

    return module_asm, LLVM.name(entry)
end

max_stars_repo_name: Cyvadra/CUDAnative

max_issues_repo_name: Cyvadra/CUDAnative

max_forks_repo_name: Cyvadra/CUDAnative


text:
using DrWatson
@quickactivate "ScoreDrivenExponentialRandomGraphs"
DrWatson.greet()

max_stars_repo_name: domenicodigangi/SDERGM

max_issues_repo_name: domenicodigangi/SDERGM

max_forks_repo_name: domenicodigangi/SDERGM


text:
module M4Test

using M4
using Test

@testset "M4.jl" begin
    # Write your own tests here.
end


text:
"""
    getPhenotype(gramm::Grammar, genotype::GEGenotype)

Obtains the phenotype as expresion in Julia from a grammatical tree.
"""
function getPhenotype(gramm::Grammar, genotype::GEGenotype)

    if !genotype._valid
        return quote Inf end
    end
    tree = genotype._tree
    stack = [tree]
    result = Array{String}(undef, 0)

    while !isempty(stack)
        node = popfirst!(stack)
        if isLeaf(node)
            push!(result, getRepresentation(gramm, node))
        else
            children = getChildren(node)
            nChildren = length(children)
            for i=1:nChildren
                pushfirst!(stack, children[i])
            end
        end
    end
    if length(result) == 1
        result = Meta.parse(result[1])
        return quote $result end
    else
        aux = Meta.parse(join(result))
        if !(typeof(aux) <: Expr)
            return quote $aux end
        else
            return aux
        end
    end
end

"""
    compareFunctions(genotype::GEGenotype,
                    gpExperimentInfo::GPExperimentInfo, objectives::Array,
                    values...)

This GE-exclusive fitness function compares the objective values given with the
ones obtained by the individual's evaluation using the values for the variables.
The fitness value that this function calculates is the MSE between the results of
the genotype for the given variable values and the objective values.

`Self-provided Arguments` are provided by the library, so only `User Arguments` must be provided.

# Self-provided Arguments
- `genotype::GPGenotype`: genotype of the individual that is going to be evaluated.
- `gpExperimentInfo::GPExperimentInfo`: information about the GP experiment.

# User Arguments
- `objectives::Array`: objective values that are going to be compared.
- `values`: values of the variables. Must be in the same order as the variables
    are in the GP experiment information (see documentation of [`GPExperimentInfo`](@ref)
    and [`getVariables`](@ref)).

# Returns
The MSE between the objective values and the results of evaluating the genotype
with the given variables values.

# Examples
```jdoctests
julia> objectives = [0, 1, 4, 9, 16, 25] #Objective values
6-element Array{Int64,1}:
[...]

julia> xValues = [0, 1, 2, 3, 4 , 5] #Values for variable x
6-element Array{Int64,1}:
[...]

julia> setEvaluator([FitnessFunction(compareFunctions, objectives, xValues, weight=-1)])
Evaluator(...)
```

See also: [`setEvaluator`](@ref), [`FitnessFunction`](@ref)
"""
function compareFunctions(genotype::GEGenotype,
                          gpExperimentInfo::GEInfo, objectives::Array,
                          values...)

    vars = getVariables(gpExperimentInfo)
    nVars = length(vars)

    vars = [Meta.parse(x) for x in vars]
    vars = (vars...,)
    phenotype = getPhenotype(gpExperimentInfo._grammar, genotype)

    values = [i for i in values]
    nValues = 1
    y = 0
    y_est = 0
    acc = 0
    failures = 0

    if nVars == 0

        y = objectives[1]

        try
            y_est = evalPhenotype(phenotype)

            if isnan(y_est) || abs(y_est) == Inf
                failures += 1
            else
                acc += (abs(y - y_est))^2
            end
        catch
            failures += 1
        end

    else
        nValues = length(values[1])
        varValues = Array{Number}(undef, nVars)
        for i=1:nValues

            for j=1:nVars
                varValues[j] = values[j][i]
            end
            varsdict = NamedTuple{vars}(varValues)

            y = objectives[i]

            try
                y_est = evaluate(phenotype, varsdict)

                if isnan(y_est) || abs(y_est) == Inf
                    failures += 1
                else
                    acc += (abs(y - y_est))^2
                end
            catch e

                failures += 1
            end
        end
    end

    if failures > (nValues * 0.15)
        fitness = Inf
    else
        mse = acc/(nValues - failures)
        fitness = (acc + failures * mse * 3) / nValues
    end

    return fitness
end # function

max_stars_repo_name: SergioGmezM/GeneticJulia

max_issues_repo_name: SergioGmezM/GeneticJulia

max_forks_repo_name: SergioGmezM/GeneticJulia


text:
doc"""
    Levy(μ, σ)

The *Lévy distribution* with location `μ` and scale `σ` has probability density function

$f(x; \mu, \sigma) = \sqrt{\frac{\sigma}{2 \pi (x - \mu)^3}}
\exp \left( - \frac{\sigma}{2 (x - \mu)} \right), \quad x > \mu$

```julia
Levy()         # Levy distribution with zero location and unit scale, i.e. Levy(0.0, 1.0)
Levy(u)        # Levy distribution with location u and unit scale, i.e. Levy(u, 1.0)
Levy(u, c)     # Levy distribution with location u ans scale c

params(d)      # Get the parameters, i.e. (u, c)
location(d)    # Get the location parameter, i.e. u
```

External links

* [Lévy distribution on Wikipedia](http://en.wikipedia.org/wiki/Lévy_distribution)
"""
immutable Levy <: ContinuousUnivariateDistribution
    μ::Float64
    σ::Float64

    Levy(μ::Real, σ::Real) = (@check_args(Levy, σ > zero(σ)); new(μ, σ))
    Levy(μ::Real) = new(μ, 1.0)
    Levy() = new(0.0, 1.0)
end

@distr_support Levy d.μ Inf

#### Parameters

location(d::Levy) = d.μ
params(d::Levy) = (d.μ, d.σ)

#### Statistics

mean(d::Levy) = Inf
var(d::Levy) = Inf
skewness(d::Levy) = NaN
kurtosis(d::Levy) = NaN

mode(d::Levy) = d.σ / 3.0 + d.μ

entropy(d::Levy) = (1.0 - 3.0 * digamma(1.0) + log(16π * d.σ^2)) / 2.0

median(d::Levy) = d.μ + d.σ / 0.4549364231195728  # 0.454... = (2.0 * erfcinv(0.5)^2)

#### Evaluation

function pdf(d::Levy, x::Float64)
    μ, σ = params(d)
    z = x - μ
    (sqrt(σ) / sqrt2π) * exp((-σ) / (2.0 * z)) / z^1.5
end

function logpdf(d::Levy, x::Float64)
    μ, σ = params(d)
    z = x - μ
    0.5 * (log(σ) - log2π - σ / z - 3.0 * log(z))
end

cdf(d::Levy, x::Float64) = erfc(sqrt(d.σ / (2.0 * (x - d.μ))))
ccdf(d::Levy, x::Float64) = erf(sqrt(d.σ / (2.0 * (x - d.μ))))

quantile(d::Levy, p::Float64) = d.μ + d.σ / (2.0 * erfcinv(p)^2)
cquantile(d::Levy, p::Float64) = d.μ + d.σ / (2.0 * erfinv(p)^2)

mgf(d::Levy, t::Real) = t == zero(t) ? 1.0 : NaN

function cf(d::Levy, t::Real)
    μ, σ = params(d)
    exp(im * μ * t - sqrt(-2.0 * im * σ * t))
end

#### Sampling

rand(d::Levy) = d.μ + d.σ / randn()^2


text:
using RobotDynamics
using Rotations
using StaticArrays, LinearAlgebra
using BenchmarkTools

# Define the model struct to inherit from `RigidBody{R}`
struct Satellite{R,T} <: RigidBody{R}
    mass::T
    J::Diagonal{T,SVector{3,T}}
end
RobotDynamics.control_dim(::Satellite) = 6

# Define some simple "getter" methods that are required to evaluate the dynamics
RobotDynamics.mass(model::Satellite) = model.mass
RobotDynamics.inertia(model::Satellite) = model.J

# Define the 3D forces at the center of mass, in the world frame
function RobotDynamics.forces(model::Satellite, x::StaticVector, u::StaticVector)
    q = orientation(model, x)
    F = @SVector [u[1], u[2], u[3]]
    q*F  # world frame
end

# Define the 3D moments at the center of mass, in the body frame
function RobotDynamics.moments(model::Satellite, x::StaticVector, u::StaticVector)
    return @SVector [u[4], u[5], u[6]]  # body frame
end

# Build model
T = Float64
R = UnitQuaternion{T}
mass = 1.0
J = Diagonal(@SVector ones(3))
model = Satellite{R,T}(mass, J)

# Initialization
x,u = rand(model)
z = KnotPoint(x,u,0.1)
∇f = RobotDynamics.DynamicsJacobian(model)

# Continuous dynamics
dynamics(model, x, u)
jacobian!(∇f, model, z)

# Performance improvements for continuous dynamics Jacobian
b1 = @benchmark jacobian!($∇f, $model, $z)

function RobotDynamics.wrench_jacobian!(F, model::Satellite, z)
    x = state(z)
    u = control(z)
    q = orientation(model, x)
    ir, iq, iv, iω, iu = RobotDynamics.gen_inds(model)
    iF = SA[1,2,3]
    iM = SA[4,5,6]
    F[iF, iq] .= Rotations.∇rotate(q, u[iF])
    F[iF, iu[iF]] .= RotMatrix(q)
    for i = 1:3
        F[iM[i], iu[i+3]] = 1
    end
    return F
end
b2 = @benchmark jacobian!($∇f, $model, $z)

function RobotDynamics.wrench_sparsity(::Satellite)
    SA[false true  false false true;
       false false false false true]
end
b3 = @benchmark jacobian!($∇f, $model, $z)
println("Analytical Wrench Jacobian:  ", judge(minimum(b2), minimum(b1)))
println("w/ Wrench Jacobian Sparsity: ", judge(minimum(b3), minimum(b1)))

# Discrete dynamics
discrete_dynamics(RK2, model, z)
discrete_jacobian!(RK2, ∇f, model, z)


text:
"""
    FENodeToFEMapModule

Module to construct a map from finite element nodes to the finite elements.
"""
module FENodeToFEMapModule

using ..FTypesModule: FInt, FFlt, FCplxFlt, FFltVec, FIntVec, FFltMat, FIntMat, FMat, FVec, FDataDict

"""
Map from finite element nodes to the finite elements connecting
them.

For each  node referenced in the connectivity of
the finite element set on input, the numbers of the individual
finite elements that reference that node is stored in an array in
the array `map`.
        Example:
```
fes.conn= [7,6,5;
            4,1,3;
            3,7,5];
The map reads
    map[1] = [2];
    map[2] = [];#  note that node number 2 is not referenced by the connectivity
    map[3] = [2,3];
    map[4] = [2];
    map[5] = [1,3];
    map[6] = [1];
    map[7] = [1,3];
```
The individual elements from the connectivity that reference node number 5 are
1 and 3, so that `fes.conn(map[5],:) `includes all the nodes that are
connected to node 5 (including node 5 itself).
"""
struct FENodeToFEMap
	"""
	Map as a vector of vectors.
	"""
    map::Array{Vector{FInt},1}
end

"""
    FENodeToFEMap(conn::Vector{NTuple{N, IT}}, nmax::FInt) where {N, IT<:Integer}

Map from finite element nodes to the finite elements connecting them.

- `conns` = connectivities as a vector of tuples
- `nmax` = largest possible node number
"""
function FENodeToFEMap(conn::Vector{NTuple{N, IT}}, nmax::FInt) where {N, IT<:Integer}
    map = FIntVec[]; sizehint!(map, nmax)
    for i = 1:nmax
        push!(map, [])  # initially empty arrays
    end
    for j = 1:length(conn)
        for i = 1:length(conn[j])
            ni = conn[j][i];
            push!(map[ni],j)
        end
    end
    return FENodeToFEMap(map)
end

"""
    FENodeToFEMap(conns::FIntMat, nmax::FInt)

Map from finite element nodes to the finite elements connecting them.

- `conns` = integer array of the connectivities
- `nmax` = largest possible node number
"""
function FENodeToFEMap(conns::FIntMat, nmax::FInt)
    map = FIntVec[]; sizehint!(map, nmax)
    for i = 1:nmax
        push!(map, [])  # initially empty arrays
    end
    for j = 1:size(conns,1)
        for i = 1:size(conns,2)
            ni = conns[j,i];
            push!(map[ni],j)
        end
    end
    return FENodeToFEMap(map)
end

end


text:
# include("./utils.jl")

abstract type Smoother end
mutable struct GaussianProcessSmoother <: Smoother
    smoother_function::Function
    kernel
    mean_function
    GaussianProcessSmoother(x, y, z) = GaussianProcessSmoother(x, y, z)
    GaussianProcessSmoother() = new()
    function GaussianProcessSmoother(mean_function, kernel)
        gp = new()
        gp.smoother_function = GP
        gp.mean_function = mean_function
        gp.kernel = kernel
        return gp
    end
end

function fit(smoother::GaussianProcessSmoother, x, y)
    model = smoother.smoother_function(
                        x,
                        y,
                        smoother.mean_function,
                        smoother.kernel
                        )
    return model
end

function predict(model::GPBase, smoothed_t)
    smoothed_x, _ = predict_y(model, smoothed_t)
    return smoothed_x
end

function smoothing(x::AbstractVector, smoother::Smoother;step=0.01)
    t = 1.0:length(x)
    standardized_x = standardize(x)
    model = fit(smoother, t, standardized_x)
    lower_bound, upper_bound = minimum(t), maximum(t)
    smoothed_t = range(lower_bound, stop=upper_bound, step=step)
    n_points = length(smoothed_t)
    smoothed_x = predict(model, smoothed_t)
    smoothed_x = reverse_standardization(smoothed_x, x)
    return smoothed_x
end
function smoothing(x::AbstractVector;step=0.01)
    smoother =  GaussianProcessSmoother(MeanZero(), SE(1.0, 0.0))
    smoothed_x = smoothing(x, smoother, step=step)
    return smoothed_x
end

function smoothing(x::Matrix{T}, smoothers::Vector{<:Smoother};step=0.01) where {T <: Number}
    n_row, n_col = size(x, 1), size(x, 2)
    @assert n_col == length(smoothers)
    t = 1.0:n_row
    lower_bound, upper_bound = minimum(t), maximum(t)
    smoothed_t = range(lower_bound, stop=upper_bound, step=step)
    n_points = length(smoothed_t)
    new_x = zeros(eltype(x), n_points, n_col)

    for i in 1:n_col
        smoothed_x = smoothing(x[:,i], smoothers[i], step=step)
        new_x[:, i] = smoothed_x
    end
    return new_x
end
function smoothing(x::Matrix{T}, smoother::Smoother;step=0.01) where {T <: Number}
    smoother_x = smoother
    smoother_y = deepcopy(smoother)
    smoother_z = deepcopy(smoother)
    smoothers = [smoother_x,smoother_y,smoother_z]
    return smoothing(x, smoothers;step=step)
end
function smoothing(x::Matrix{T};step=0.01) where {T <: Number}
    n_row, n_col = size(x, 1), size(x, 2)
    t = 1.0:n_row
    lower_bound, upper_bound = minimum(t), maximum(t)
    smoothed_t = range(lower_bound, stop=upper_bound, step=step)
    n_points = length(smoothed_t)
    new_x = zeros(eltype(x), n_points, n_col)
    for i in 1:n_col
        smoothed_x = smoothing(x[:,i], step=step)
        new_x[:, i] = smoothed_x
    end
    return new_x
end


text:
# ---
# title: 1409. Queries on a Permutation With Key
# id: problem1409
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Array
# link: <https://leetcode.com/problems/queries-on-a-permutation-with-key/description/>
# hidden: true
# ---
#
# Given the array `queries` of positive integers between `1` and `m`, you have
# to process all `queries[i]` (from `i=0` to `i=queries.length-1`) according to
# the following rules:
#
#   * In the beginning, you have the permutation `P=[1,2,3,...,m]`.
#   * For the current `i`, find the position of `queries[i]` in the permutation `P` ( **indexing from 0** ) and then move this at the beginning of the permutation `P.` Notice that the position of `queries[i]` in `P` is the result for `queries[i]`.
#
# Return an array containing the result for the given `queries`.
#
#
#
# **Example 1:**
#
#
#
#     Input: queries = [3,1,2,1], m = 5
#     Output: [2,1,2,1]
#     Explanation: The queries are processed as follow:
#     For i=0: queries[i]=3, P=[1,2,3,4,5], position of 3 in P is **2** , then we move 3 to the beginning of P resulting in P=[3,1,2,4,5].
#     For i=1: queries[i]=1, P=[3,1,2,4,5], position of 1 in P is **1** , then we move 1 to the beginning of P resulting in P=[1,3,2,4,5].
#     For i=2: queries[i]=2, P=[1,3,2,4,5], position of 2 in P is **2** , then we move 2 to the beginning of P resulting in P=[2,1,3,4,5].
#     For i=3: queries[i]=1, P=[2,1,3,4,5], position of 1 in P is **1** , then we move 1 to the beginning of P resulting in P=[1,2,3,4,5].
#     Therefore, the array containing the result is [2,1,2,1].
#
#
# **Example 2:**
#
#
#
#     Input: queries = [4,1,2,2], m = 4
#     Output: [3,1,2,0]
#
#
# **Example 3:**
#
#
#
#     Input: queries = [7,5,5,8,3], m = 8
#     Output: [6,5,0,7,5]
#
#
#
#
# **Constraints:**
#
#   * `1 <= m <= 10^3`
#   * `1 <= queries.length <= m`
#   * `1 <= queries[i] <= m`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
module Time

export @time2

macro time2(expr)
    eexpr = esc(expr)
    quote
        @time $eexpr
        @time $eexpr
    end
end

end


text:
# Control the move generation used by the Annealing scheme.
"""
API
---
* [`generate_move`](@ref)
* [`distancelimit`](@ref)
* [`initialize!`](@ref)
* [`update!`](@ref)

Implementations
---------------
* [`CachedMoveGenerator`](@ref)
"""
abstract type MoveGenerator end

# Document the interface for an MoveGenerator.
"""
    generate_move(sa_struct, move_generator, node_idx)

Generate a valid for move for the node with index `node_idx` in `sa_struct`.
If `isflat(sa_struct) == true`, return `CartesianIndex{D}` where
`D = dimension(sa_struct)`. Otherwise, return `Location{D}`.
"""
function generate_move end

"""
    distancelimit(move_generator, sa_struct)

Return the maximum move distance of `move_generator` for `sa_struct`.
"""
function distancelimit end

"""
    initialize!(move_generator, sa_struct, [limit = distancelimit(move_generator, sa_struct)])

Initialize the state of `move_generator` based on `sa_struct`. Common operations
include establishing an initial move distance limit or caching a list of
possible moves.

If an initial limit is not supplied, it defaults to the maximum limit of
the move generator for the given architecture.
"""
function initialize!(move_generator, sa_struct)
    return initialize!(move_generator, sa_struct, distancelimit(move_generator, sa_struct))
end

"""
    update!(move_generator, sa_struct, limit)

Update the move generator to a new move distance limit.
"""
function update! end

# Top level function - serves as entry point for the methods listed above.
@propagate_inbounds function generate_move!(
    sa_struct::SAStruct, move_generator::MoveGenerator, move_cookie
)

    # Pick a random address
    node_idx = rand(1:length(sa_struct.nodes))

    new_location = generate_move(sa_struct, move_generator, node_idx)
    return move_with_undo(sa_struct, move_cookie, node_idx, new_location)
end

################################################################################
# Cached Move Generator
################################################################################

# The idea for this data structure is that the Locations or Addresses within
# some distance D of each Location/Address is precomputed. When the distance
# limit changes, this structure will have to be modified - hence why it's
# mutable.

"""
Look-up table for moves for a single node class starting at some base address.

The main invariant of this a MoveLUT ``L`` is with base address ``\\alpha`` is

```math
\\text{distance}\\left( L_\\text{targets}[i] - \\alpha \\right) \\leq \\delta
    \\, \\forall i \\in 1 \\ldots L_\\text{idx}
```

where ``\\text{distance}`` is the distance between to addresses in the SAStruct
and `\\delta` is the current move distance limit.

Thus, to generate a random move within ``\\delta`` of ``\\alpha``, we must
need to perform the operation
```
L.targets[rand(1:L.idx)]
```
assuming `L.idx` has been configured correctly.

To aid in the configuration of `L.idx`, the field `indices` is constructed
such that for a move limit ``\\delta``, `L.idx = L.indices[δ]`.
"""
mutable struct MoveLUT{T}
    """
    Vector of destination addresses from the base address. Sorted in
    increasing order of distance from the base addresses according to the
    distance metric of the parent `SAStruct`.
    """
    targets::Vector{T}

    """
    The index of the last entry in `targets` that is within the current move
    distance limit of the base address.
    """
    idx::Int

    """
    Cached `idx` for various move distance limits.
    """
    indices::Vector{Int}
end

# Get a random element of distance "d"
function Base.rand(m::MoveLUT)
    @inbounds target = m.targets[rand(1:(m.idx))]
    return target
end

"""
This move generator precomputes all of the valid moves for each class at all
addresses, and references this cached database to generator moves.

Standard classes are used to index into the first level of `moves`. The inner
dictionary is a mapping from a base address to a [`MoveLUT`](@ref) for that
address.
"""
mutable struct CachedMoveGenerator{T} <: MoveGenerator
    # Outer vector: class index
    # inner dict: starting Location
    moves::Vector{Dict{T,MoveLUT{T}}}

    CachedMoveGenerator{T}() where {T} = new{T}(Dict{T,MoveLUT{T}}[])
end
function CachedMoveGenerator(sa_struct::SAStruct{A,U,D}) where {A,U,D}
    return CachedMoveGenerator{location_type(sa_struct.maptable)}()
end

# Configure the distance limit to be the maximum value in the distance table
# to be sure that initially, a task may be moved anywhere on the proecessor
# array.
distancelimit(::CachedMoveGenerator, sa_struct) = maxdistance(sa_struct, sa_struct.distance)

function initialize!(move_generator::CachedMoveGenerator, sa_struct::SAStruct, limit)
    maptable = sa_struct.maptable
    num_classes = length(maptable.mask)

    moves = map(1:num_classes) do class
        # Get all locations that can be occupied by this class.
        class_locations = getlocations(maptable, class)

        # Create a dictionary of MoveLUTs for this class.
        return Dict(
            map(class_locations) do source
                # Define a function that returns the distance of a location from the current
                # source address.
                dist(x) = getdistance(sa_struct.distance, source, x)

                # Sort destinations for this address by distance.
                dests = sort(
                    class_locations;
                    # Comparison of distance. If distances are equal, sort by
                    # CartesianIndex for easier testing.
                    lt = (x, y) ->
                        dist(x) < dist(y) ||
                            (dist(x) == dist(y) && location(x) < location(y)),
                )

                # Create an look up vector.
                # Find the index of the last entry in `dests` whose distance from the source is
                # within the limit `i`.
                θ = [findlast(x -> dist(x) <= i, dests) for i in 1:limit]

                # Will throw if any entry in θ is `Nothing` - which is a good thing
                return source => MoveLUT(dests, length(dests), θ)
            end,
        )
    end

    move_generator.moves = moves
    return nothing
end

# Modify the index in each MoveLUT to correspond to the new limit.
function update!(move_generator::CachedMoveGenerator, sa_struct::SAStruct, limit)
    for dict in move_generator.moves, lut in values(dict)
        lut.idx = lut.indices[min(length(lut.indices), limit)]
    end
    return nothing
end

@propagate_inbounds function generate_move(
    sa_struct::SAStruct, move_generator::CachedMoveGenerator, node_idx
)
    node = sa_struct.nodes[node_idx]

    # Unpack current location and class of the node.
    this_location = location(node)
    class = getclass(node)

    # Use the class and current location to get a `MoveLUT` from the move generator
    move_lut = move_generator.moves[class][this_location]

    # Return a random address from the move_lut.
    return rand(move_lut)
end


text:
using Blosc
blosc_compress(s, level = 9, shuffle = false, compressor = "blosclz") = begin
  Blosc.set_compressor(compressor)
  compress(s; level = level, shuffle = shuffle)
end

blosclzlength(s) = length(blosc_compress(s, 9, false, "blosclz"))
lz4length(s) = length(blosc_compress(s, 9, false, "lz4"))
lz4hclength(s) = length(blosc_compress(s, 9, false, "lz4hc"))
snappylength(s) = length(blosc_compress(s, 9, false, "snappy"))
blosczliblength(s) = length(blosc_compress(s, 9, false, "zlib"))
qlz4length(s) = length(blosc_compress(s, 1, false, "lz4"))
qsnappylength(s) = length(blosc_compress(s, 1, false, "snappy"))
qblosczliblength(s) = length(blosc_compress(s, 1, false, "zlib"))

function find_min_len_for_compression(compressLenFunc, seedStr = "a", maxReps = 1000)
  prevclen = currclen = compressLenFunc(seedStr)
  for i in 1:maxReps
    prevclen = currclen
    str = seedStr ^ i
    currclen = compressLenFunc(str)
    #print("(", length(str), ", ", currclen, "), ")
    if currclen < prevclen
      return i, length(str), currclen, length(str)+length(seedStr), compressLenFunc(seedStr^(i+1))
    end
  end
  return -1, str
end

find_min_len_for_compression(blosclzlength, "arne")
find_min_len_for_compression(lz4length, "arne")
find_min_len_for_compression(lz4hclength, "arne")

# Seems that there is a min length around 128 bytes. Below that no compression happens.