text:
export AffineConstantFlow, AffineHalfFlow, SlowMAF, MAF, IAF, ActNorm,
       Invertible1x1Conv, NormalizingFlow, NormalizingFlowModel, NeuralCouplingFlow,
       autoregressive_network, FlowOp, LinearFlow, TanhFlow, AffineScalarFlow,
       InvertFlow, AffineFlow

abstract type FlowOp end

#------------------------------------------------------------------------------------------
# Users can define their own FlowOp in the Main
forward(fo::T, x::Union{Array{<:Real}, PyObject}) where T<:FlowOp  = Main.forward(fo, x)
backward(fo::T, z::Union{Array{<:Real}, PyObject}) where T<:FlowOp = Main.backward(fo, z)

#------------------------------------------------------------------------------------------
mutable struct LinearFlow <: FlowOp
    dim::Int64
    indim::Int64
    outdim::Int64
    A::PyObject
    b::PyObject
end

@doc raw"""
    LinearFlow(indim::Int64, outdim::Int64, A::Union{PyObject, Array{<:Real,2}, Missing}=missing,
        b::Union{PyObject, Array{<:Real,1}, Missing}=missing, name::Union{String, Missing}=missing)
    LinearFlow(A::Union{PyObject, Array{<:Real,2}},
        b::Union{PyObject, Array{<:Real,1}, Missing} = missing; name::Union{String, Missing}=missing)

Creates a linear flow operator,

$$x = Az + b$$

The operator is not necessarily invertible. $A$ has the size `outdim×indim`, $b$ is a length `outdim` vector.

# Example 1: Reconstructing a Gaussian random variable using invertible A
```julia
using ADCME
using Distributions

d = MvNormal([2.0;3.0],[2.0 1.0;1.0 2.0])
flows = [LinearFlow(2, 2)]
prior = ADCME.MultivariateNormalDiag(loc = zeros(2))
model = NormalizingFlowModel(prior, flows)
x = rand(d, 10000)'|>Array
zs, prior_logpdf, logdet = model(x)
log_pdf = prior_logpdf + logdet
loss = -sum(log_pdf)
sess = Session(); init(sess)
BFGS!(sess, loss)
run(sess, flows[1].b) # approximately [2.0 3.0]
run(sess, flows[1].A * flows[1].A') # approximate [2.0 1.0;1.0 2.0]
```

# Example 2: Reconstructing a Gaussian random variable using non-invertible A
```julia
using ADCME
using Distributions

d = Normal()
A = reshape([2. 1.], 2, 1)
b = Variable(ones(2))
flows = [LinearFlow(A, b)]
prior = ADCME.MultivariateNormalDiag(loc = zeros(1))
model = NormalizingFlowModel(prior, flows)

x = rand(d, 10000, 1)
x = x * [2.0 1.0] .+ [1.0 3.0]

x = [2.0 2.0]
zs, prior_logpdf, logdet = model(x)
log_pdf = prior_logpdf + logdet
loss = -sum(log_pdf)
sess = Session(); init(sess)
BFGS!(sess, loss)
run(sess, b[1]*2 + b[2]) # approximate 6.0
```
"""
function LinearFlow(indim::Int64, outdim::Int64, A::Union{PyObject, Array{<:Real,2}, Missing}=missing,
        b::Union{PyObject, Array{<:Real,1}, Missing}=missing; name::Union{String, Missing}=missing)
    if ismissing(name)
        name = randstring(10)
    end
    if ismissing(A)
        if indim==outdim
            A = diagm(0=>ones(indim))
            A = get_variable(A, name=name*"_A")
        else
            A = get_variable(Float64, shape=[outdim,indim], name=name*"_A")
        end
    end

    if ismissing(b)
        b = get_variable(zeros(1,outdim), name=name*"_b")
    end

    A = constant(A)
    b = constant(b)
    b = reshape(b, (1,-1))

    if size(A)!=(outdim,indim)
        error("Expected A shape: $((outdim,indim)), but got $(size(A))")
    end
    if size(b)!=(1,outdim)
        error("Expected b shape: $((1,outdim)), but got $(size(b))")
    end

    LinearFlow(outdim, indim, outdim, A, b)
end

function LinearFlow(dim::Int64)
    LinearFlow(dim,dim)
end

function LinearFlow(A::Union{PyObject, Array{<:Real,2}},
    b::Union{PyObject, Array{<:Real}, Missing} = missing; name::Union{String, Missing}=missing)
    outdim, indim = size(A)
    if ismissing(b)
        b = zeros(outdim)
    end
    A = constant(A)
    b = constant(b)
    LinearFlow(outdim, indim, outdim, A, b)
end

function forward(fo::LinearFlow, x::Union{Array{<:Real}, PyObject})
    if fo.indim==fo.outdim
        return solve_batch(fo.A, x - fo.b), -log(abs(det(fo.A))) * ones(size(x, 1))
    else
        return solve_batch(fo.A, x - fo.b), -1/2*log(abs(det(fo.A'*fo.A))) * ones(size(x, 1))
    end
end

function backward(fo::LinearFlow, z::Union{Array{<:Real}, PyObject})
    d = nothing
    if fo.indim==fo.outdim
        d = log(abs(det(fo.A)))
    end
    z * fo.A' + fo.b, d
end

#------------------------------------------------------------------------------------------
mutable struct AffineConstantFlow <: FlowOp
    dim::Int64
    s::Union{Missing,PyObject}
    t::Union{Missing,PyObject}
end

function AffineConstantFlow(dim::Int64, name::Union{String, Missing}=missing; scale::Bool=true, shift::Bool=true)
    if ismissing(name)
        name = randstring(10)
    end
    s = missing
    t = missing
    if scale
        s = get_variable(Float64, shape=[1,dim], name=name*"_s")
    end
    if shift
        t = get_variable(Float64, shape=[1,dim], name=name*"_t")
    end
    AffineConstantFlow(dim, s, t)
end

function forward(fo::AffineConstantFlow, x::Union{Array{<:Real}, PyObject})
    s = ismissing(fo.s) ? zeros_like(x) : fo.s
    t = ismissing(fo.t) ? zeros_like(x) : fo.t
    z = x .* exp(s) + t
    log_det = sum(s, dims=2)
    return z, log_det
end

function backward(fo::AffineConstantFlow, z::Union{Array{<:Real}, PyObject})
    s = ismissing(fo.s) ? zeros_like(z) : fo.s
    t = ismissing(fo.t) ? zeros_like(z) : fo.t
    x = (z-t) .* exp(-s)
    log_det = sum(-s, dims=2)
    return x, log_det
end

#------------------------------------------------------------------------------------------
mutable struct ActNorm <: FlowOp
    dim::Int64
    fo::AffineConstantFlow
    initialized::PyObject
end

function ActNorm(dim::Int64, name::Union{String, Missing}=missing)
    if ismissing(name)
        name = "ActNorm_"*randstring(10)
    end
    initialized = get_variable(1, name = name)
    ActNorm(dim, AffineConstantFlow(dim), initialized)
end

function forward(actnorm::ActNorm, x::Union{Array{<:Real}, PyObject})
    x = constant(x)
    actnorm.fo.s, actnorm.fo.t = tf.cond(tf.equal(actnorm.initialized,constant(0)),
        ()->(actnorm.fo.s, actnorm.fo.t), #actnorm.fo.s, actnorm.fo.t),
        ()->begin
            setzero = assign(actnorm.initialized, 0)
            x0 = copy(x)
            s0 = reshape(-log(std(x0, dims=1)), 1, :)
            t0 = reshape(mean(-x0 .* exp(s0), dims=1), 1, :)
            op = assign([actnorm.fo.s, actnorm.fo.t], [s0, t0])
            p = tf.print("ActNorm: initializing s and t...")
            op[1] = bind(op[1], setzero)
            op[1] = bind(op[1], p)
            op[1], op[2]
        end)
    forward(actnorm.fo, x)
end

function backward(actnorm::ActNorm, z::Union{Array{<:Real}, PyObject})
    backward(actnorm.fo, z)
end

#------------------------------------------------------------------------------------------
mutable struct SlowMAF <: FlowOp
    dim::Int64
    layers::Array{Function}
    order::Array{Int64}
    p::PyObject
end

function SlowMAF(dim::Int64, parity::Bool, nns::Array)
    local order
    @assert length(nns)==dim - 1
    if parity
        order = Array(1:dim)
    else
        order = reverse(Array(1:dim))
    end
    p = Variable(zeros(2))
    SlowMAF(dim, nns, order, p)
end

function forward(fo::SlowMAF, x::Union{Array{<:Real}, PyObject})
    z = tf.zeros_like(x)
    log_det = zeros(size(x,1))
    for i = 1:fo.dim
        if i==1
            st = repeat(fo.p', size(x,1), 1)
        else
            st = fo.layers[i-1](x[:,1:i-1])
            @assert size(st, 2)==2
        end
        s, t = st[:,1], st[:,2]
        z = scatter_update(z, :, fo.order[i],  x[:, i]*exp(s) + t)
        log_det += s
    end
    return z, log_det
end

function backward(fo::SlowMAF, z::Union{Array{<:Real}, PyObject})
    x = tf.zeros_like(z)
    log_det = zeros(size(z,1))
    for i = 1:fo.dim
        if i==1
            st = repeat(fo.p', size(x,1), 1)
        else
            st = fo.layers[i-1](x[:,1:i-1])
        end
        s, t = st[:,1], st[:,2]
        x = scatter_update(x, :, i, (z[:, fo.order[i]] - t) * exp(-s))
        log_det += -s
    end
    return x, log_det
end

#------------------------------------------------------------------------------------------
mutable struct MAF <: FlowOp
    dim::Int64
    net::Function
    parity::Bool
end

function MAF(dim::Int64, parity::Bool, config::Array{Int64}; name::Union{String, Missing} = missing,
    activation::Union{Nothing,String} = "relu", kwargs...)
    push!(config, 2)
    kwargs = jlargs(kwargs)
    if ismissing(name)
        name = "MAF_"*randstring(10)
    end
    net = x->autoregressive_network(x, config, name; activation = activation)
    MAF(dim, net, parity)
end

function forward(fo::MAF, x::Union{Array{<:Real}, PyObject})
    st = fo.net(x)
    # @info st, split(st, 1, dims=2)
    s, t = squeeze.(split(st, 2, dims=3))
    z = x .* exp(s) + t
    if fo.parity
        z = reverse(z, dims=2)
    end
    log_det = sum(s, dims=2)
    return z, log_det
end

function backward(fo::MAF, z::Union{Array{<:Real}, PyObject})
    x = zeros_like(z)
    log_det = zeros(size(z, 1))
    if fo.parity
        z = reverse(z, dims=2)
    end
    for i = 1:fo.dim
        st = fo.net(copy(x))
        s, t = squeeze.(split(st, 2, dims=3))
        x = scatter_update(x, :, i, (z[:, i] - t[:, i]) * exp(-s[:, i]))
        log_det += -s[:,i]
    end
    return x, log_det
end

"""
    autoregressive_network(x::Union{Array{Float64}, PyObject}, config::Array{<:Integer},
    scope::String="default"; activation::String=nothing, kwargs...)

Creates an masked autoencoder for distribution estimation.
"""
function autoregressive_network(x::Union{Array{Float64}, PyObject}, config::Array{<:Integer},
    scope::String="default"; activation::Union{Nothing,String}="tanh", kwargs...)
    local y
    x = constant(x)
    params = config[end]
    hidden_units = PyVector(config[1:end-1])
    event_shape = [size(x, 2)]
    if haskey(STORAGE, scope*"/autoregressive_network/")
        made = STORAGE[scope*"/autoregressive_network/"]
        y = made(x)
    else
        made = ADCME.tfp.bijectors.AutoregressiveNetwork(params=params, hidden_units = hidden_units,
            event_shape = event_shape, activation = activation, kwargs...)
        STORAGE[scope*"/autoregressive_network/"] = made
        y = made(x)
    end
    return y
end

#------------------------------------------------------------------------------------------
mutable struct IAF <: FlowOp
    maf::MAF
end

function IAF(dim::Int64, parity::Bool, config::Array{Int64}; name::Union{String, Missing} = missing,
    activation::Union{Nothing,String} = "relu", kwargs...)
    MAF(dim, parity, config; name=name, activation=activation, kwargs...)
end

function forward(fo::IAF, x::Union{Array{<:Real}, PyObject})
    backward(fo.maf, x)
end

function backward(fo::IAF, z::Union{Array{<:Real}, PyObject})
    forward(fo.maf, z)
end

#------------------------------------------------------------------------------------------
mutable struct Invertible1x1Conv <: FlowOp
    dim::Int64
    P::Array{Float64}
    L::PyObject
    S::PyObject
    U::PyObject
end

function Invertible1x1Conv(dim::Int64, name::Union{Missing,String}=missing)
    if ismissing(name)
        name = randstring(10)
    end
    Q = lu(randn(dim,dim))
    P, L, U = Q.P, Q.L, Q.U
    L = get_variable(L, name=name*"_L")
    S = get_variable(diag(U), name=name*"_S")
    U = get_variable(triu(U, 1), name=name*"_U")
    Invertible1x1Conv(dim, P, L, S, U)
end

function _assemble_W(fo::Invertible1x1Conv)
    L = tril(fo.L, -1) + diagm(0=>ones(fo.dim))
    U = triu(fo.U, 1)
    W = fo.P * L * (U + diagm(fo.S))
    return W
end

function forward(fo::Invertible1x1Conv, x::Union{Array{<:Real}, PyObject})
    W = _assemble_W(fo)
    z = x*W
    log_det = sum(log(abs(fo.S)))
    return z, log_det
end

function backward(fo::Invertible1x1Conv, z::Union{Array{<:Real}, PyObject})
    W = _assemble_W(fo)
    Winv = inv(W)
    x = z*Winv
    log_det = -sum(log(abs(fo.S)))
    return x, log_det
end

#------------------------------------------------------------------------------------------
mutable struct AffineHalfFlow <: FlowOp
    dim::Int64
    parity::Bool
    s_cond::Function
    t_cond::Function
end

"""
    AffineHalfFlow(dim::Int64, parity::Bool, s_cond::Union{Function, Missing} = missing, t_cond::Union{Function, Missing} = missing)

Creates an `AffineHalfFlow` operator.
"""
function AffineHalfFlow(dim::Int64, parity::Bool, s_cond::Union{Function, Missing} = missing, t_cond::Union{Function, Missing} = missing)
    if ismissing(s_cond)
        s_cond = x->constant(zeros(size(x,1), (dim+1)÷2))
    end
    if ismissing(t_cond)
        t_cond = x->constant(zeros(size(x,1), dim÷2))
    end
    AffineHalfFlow(dim, parity, s_cond, t_cond)
end

function forward(fo::AffineHalfFlow, x::Union{Array{<:Real}, PyObject})
    x = constant(x)
    x0 = x[:,1:2:end]
    x1 = x[:,2:2:end]
    if fo.parity
        x0, x1 = x1, x0
    end
    s = fo.s_cond(x0)
    t = fo.t_cond(x0)
    z0 = x0
    z1 = exp(s) .* x1 + t
    if fo.parity
        z0, z1 = z1, z0
    end
    z = [z0 z1]
    log_det = sum(s, dims=2)
    return z, log_det
end

function backward(fo::AffineHalfFlow, z::Union{Array{<:Real}, PyObject})
    z = constant(z)
    z0 = z[:,1:2:end]
    z1 = z[:,2:2:end]
    if fo.parity
        z0, z1 = z1, z0
    end
    s = fo.s_cond(z0)
    t = fo.t_cond(z0)
    x0 = z0
    x1 = (z1-t).*exp(-s)
    if fo.parity
        x0, x1 = x1, x0
    end
    x = [x0 x1]
    log_det = sum(-s, dims=2)
    return x, log_det
end

#------------------------------------------------------------------------------------------
mutable struct NeuralCouplingFlow <: FlowOp
    dim::Int64
    K::Int64
    B::Int64
    f1::Function
    f2::Function
end

function NeuralCouplingFlow(dim::Int64, f1::Function, f2::Function, K::Int64=8, B::Int64=3)
    @assert mod(dim,2)==0
    NeuralCouplingFlow(dim, K, B, f1, f2)
end

function forward(fo::NeuralCouplingFlow, x::Union{Array{<:Real}, PyObject})
    x = constant(x)
    dim, K, B, f1, f2 = fo.dim, fo.K, fo.B, fo.f1, fo.f2
    RQS = tfp.bijectors.RationalQuadraticSpline

    log_det = constant(zeros(size(x,1)))
    lower, upper = x[:,1:dim÷2], x[:,dim÷2+1:end]

    f1_out = f1(lower)
    @assert size(f1_out,2)==(3K-1)*size(lower,2)
    out = reshape(f1_out, (-1, size(lower,2), 3K-1))
    W, H, D = split(out, [K, K, K-1], dims=3)

    W, H = softmax(W, dims=3), softmax(H, dims=3)
    W, H = 2B*W, 2B*H
    D = softplus(D)

    rqs = RQS(W, H, D, range_min=-B)
    upper, ld = rqs.forward(upper), rqs.forward_log_det_jacobian(upper, 1)
    log_det += ld

    f2_out = f2(upper)
    @assert size(f2_out,2)==(3K-1)*size(upper,2)
    out = reshape(f2_out, (-1, size(upper,2), 3K -1))
    W, H, D = split(out, [K, K, K-1], dims=3)
    W, H = softmax(W, dims=3), softmax(H, dims=3)
    W, H = 2B*W, 2B*H
    D = softplus(D)

    rqs = RQS(W, H, D, range_min=-B)
    lower, ld = rqs.forward(lower), rqs.forward_log_det_jacobian(lower, 1)
    log_det += ld

    return [lower upper], log_det
end

function backward(fo::NeuralCouplingFlow, x::Union{Array{<:Real}, PyObject})
    x = constant(x)
    dim, K, B, f1, f2 = fo.dim, fo.K, fo.B, fo.f1, fo.f2
    RQS = tfp.bijectors.RationalQuadraticSpline

    log_det = constant(zeros(size(x,1)))
    lower, upper = x[:,1:dim÷2], x[:,dim÷2+1:end]

    f2_out = f2(upper)
    @assert size(f2_out,2)==(3K-1)*size(upper,2)
    out = reshape(f2_out, (-1, size(upper,2), 3K-1))
    W, H, D = split(out, [K, K, K-1], dims=3)
    W, H = softmax(W, dims=3), softmax(H, dims=3)
    W, H = 2B*W, 2B*H
    D = softplus(D)

    rqs = RQS(W, H, D, range_min=-B)
    lower, ld = rqs.inverse(lower), rqs.inverse_log_det_jacobian(lower, 1)
    log_det += ld

    f1_out = f1(lower)
    @assert size(f1_out,2)==(3K-1)*size(lower,2)
    out = reshape(f1_out, (-1, size(lower,2), 3K -1))
    W, H, D = split(out, [K, K, K-1], dims=3)
    W, H = softmax(W, dims=3), softmax(H, dims=3)
    W, H = 2B*W, 2B*H
    D = softplus(D)
    rqs = RQS(W, H, D, range_min=-B)
    upper, ld = rqs.inverse(upper), rqs.inverse_log_det_jacobian(upper, 1)
    log_det += ld

    return [lower upper], log_det
end

#------------------------------------------------------------------------------------------
mutable struct TanhFlow <: FlowOp
    dim::Int64
    o::PyObject
end

function TanhFlow(dim::Int64)
    TanhFlow(dim, tfp.bijectors.Tanh())
end

function forward(fo::TanhFlow, x::Union{Array{<:Real}, PyObject})
    fo.o.forward(x), fo.o.forward_log_det_jacobian(x, 1)
end

function backward(fo::TanhFlow, z::Union{Array{<:Real}, PyObject})
    fo.o.inverse(z), fo.o.inverse_log_det_jacobian(z, 1)
end

#------------------------------------------------------------------------------------------
mutable struct AffineScalarFlow <: FlowOp
    dim::Int64
    o::PyObject
end

function AffineScalarFlow(dim::Int64, shift::Float64 = 0.0, scale::Float64 = 1.0)
    AffineScalarFlow(dim, tfp.bijectors.AffineScalar(shift=shift, scale=constant(scale)))
end

function forward(fo::AffineScalarFlow, x::Union{Array{<:Real}, PyObject})
    fo.o.forward(x), fo.o.forward_log_det_jacobian(x, 1)
end

function backward(fo::AffineScalarFlow, z::Union{Array{<:Real}, PyObject})
    fo.o.inverse(z), fo.o.inverse_log_det_jacobian(z, 1)
end

#------------------------------------------------------------------------------------------
mutable struct InvertFlow <: FlowOp
    dim::Int64
    o::FlowOp
end

"""
    InvertFlow(fo::FlowOp)

Creates a flow operator that is the invert of `fo`.
# Example
```julia
inv_tanh_flow = InvertFlow(TanhFlow(2))
```
"""
function InvertFlow(fo::FlowOp)
    InvertFlow(fo.dim, fo)
end
forward(fo::InvertFlow, x::Union{Array{<:Real}, PyObject}) = backward(fo.o, x)
backward(fo::InvertFlow, z::Union{Array{<:Real}, PyObject}) = forward(fo.o, z)

#------------------------------------------------------------------------------------------
@doc raw"""
    AffineFlow(args...;kwargs...)

Creates a linear flow operator,

$$z = Ax + b$$

See [`LinearFlow`](@ref)
"""
AffineFlow(args...;kwargs...) = InvertFlow(LinearFlow(args...;kwargs...))

#------------------------------------------------------------------------------------------
mutable struct NormalizingFlow
    flows::Array
    NormalizingFlow(flows::Array) = new(flows)
end

function forward(fo::NormalizingFlow, x::Union{PyObject, Array{<:Real,2}})
    x = constant(x)
    m = size(x,1)
    log_det = constant(zeros(m))
    zs = PyObject[x]
    for flow in fo.flows
        x, ld = forward(flow, x)
        log_det += ld
        push!(zs, x)
    end
    return zs, log_det
end

function backward(fo::NormalizingFlow, z::Union{PyObject, Array{<:Real,2}})
    z = constant(z)
    m = size(z,1)
    log_det = constant(zeros(m))
    xs = [z]
    for flow in reverse(fo.flows)
        z, ld = backward(flow, z)
        if isnothing(ld) || isnothing(log_det)
            log_det = nothing
        else
            log_det += ld
        end
        push!(xs, z)
    end
    return xs, log_det
end

#------------------------------------------------------------------------------------------
mutable struct NormalizingFlowModel
    prior::ADCMEDistribution
    flow::NormalizingFlow
end

function Base.:show(io::IO, model::NormalizingFlowModel)
    typevec = typeof.(model.flow.flows)
    println("( $(string(typeof(model.prior))[7:end]) )")
    println("\t↓")
    for i = length(typevec):-1:1
        println("$(typevec[i]) (dim = $(model.flow.flows[i].dim))")
        println("\t↓")
    end
    println("(    Data    )")
end

function NormalizingFlowModel(prior::T, flows::Array{<:FlowOp}) where T<:ADCMEDistribution
    NormalizingFlowModel(prior, NormalizingFlow(flows))
end

function forward(nf::NormalizingFlowModel, x::Union{PyObject, Array{<:Real,2}})
    x = constant(x)
    zs, log_det = forward(nf.flow, x)
    prior_logprob = logpdf(nf.prior, zs[end])
    return zs, prior_logprob, log_det
end

function backward(nf::NormalizingFlowModel, z::Union{PyObject, Array{<:Real,2}})
    z = constant(z)
    xs, log_det = backward(nf.flow, z)
    return xs, log_det
end

function Base.:rand(nf::NormalizingFlowModel, num_samples::Int64)
    z = rand(nf.prior, num_samples)
    xs, _ = backward(nf.flow, z)
    return xs
end

#------------------------------------------------------------------------------------------
(o::NormalizingFlowModel)(x::Union{PyObject, Array{<:Real,2}}) = forward(o, x)

#------------------------------------------------------------------------------------------
# From TensorFlow Probability
for (op1, op2) in [(:Abs, :AbsoluteValue), (:Exp, :Exp), (:Log, :Log),
        (:MaskedAutoregressiveFlow, :MaskedAutoregressiveFlow), (:Permute,:Permute),
        (:PowerTransform, :PowerTransform), (:Sigmoid,:Sigmoid), (:Softplus,:Softplus),
        (:Square, :Square)]
    @eval begin
        export $op1
        mutable struct $op1 <: FlowOp
            dim::Int64
            fo::PyObject
        end
        function $op1(dim::Int64, args...;kwargs...)
            fo = tfp.bijectors.$op2(args...;kwargs...)
            $op1(dim, fo)
        end
        forward(o::$op1, x::Union{PyObject, Array{<:Real,2}}) = (o.fo.forward(x), o.fo.forward_log_det_jacobian(x, 1))
        backward(o::$op1, z::Union{PyObject, Array{<:Real,2}}) = (o.fo.inverse(z), o.fo.inverse_log_det_jacobian(z, 1))
    end
end


text:
using Weave

weave("../examples/FIR_design.jmd",
     informat="markdown", out_path = "build/examples", doctype = "pandoc")

weave("../examples/FIR_design.jmd",
      informat="markdown", out_path = "build/examples", doctype = "md2html")

weave("../examples/FIR_design.jmd",
        informat="markdown", out_path = "build/examples", doctype = "md2pdf")

cp("../examples/FIR_design.jmd",
    "build/examples/FIR_design.jmd", force = true)

cp("build/examples/FIR_design.md",
      "build/examples/FIR_design.txt", force = true)

weave("../examples/FIR_design_plots.jl", out_path = "build/examples")
weave("../examples/FIR_design_plots.jl", doctype = "md2pdf", out_path = "build/examples")
cp("../examples/FIR_design_plots.jl",
    "build/examples/FIR_design_plots.jl", force = true)

ext: jl


text:
using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation

using Plasm
View = Plasm.view

square = ([[0; 0] [0; 1] [1; 0] [1; 1]], [[1, 2, 3, 4]],
[[1,2], [1,3], [2,4], [3,4]])
V,FV,EV  =  square
model  =  V,([[1],[2],[3],[4]],EV,FV)
table = Lar.apply(Lar.t(-0.5,-0.5), square)
chair = Lar.Struct([Lar.t(0.75,0),Lar.s(0.35,0.35),table])
structo = Lar.Struct([Lar.t(2,1),table,repeat([Lar.r(pi/2),chair],outer = 4)...])
structo1 = Lar.Struct(repeat([structo,Lar.t(0,2.5)],outer = 10));
structo2 = Lar.Struct(repeat([structo1,Lar.t(3,0)],outer = 10));

scene = Lar.evalStruct(structo2);
View(scene)
W,FW,EW = Lar.struct2lar(structo2);
View(Plasm.lar2hpc(W,EW))
assembly = Lar.Struct([Lar.sphere()(), Lar.t(3,0,-1), Lar.cylinder()()])
View(assembly)
View(Lar.struct2lar(assembly))

cube = Lar.apply( Lar.t(-.5,-.5,0), Lar.cuboid([1,1,1]))
tableTop = Lar.Struct([ Lar.t(0,0,.85), Lar.s(1,1,.05), cube ])
tableLeg = Lar.Struct([ Lar.t(-.475,-.475,0), Lar.s(.1,.1,.89), cube ])
tablelegs = Lar.Struct( repeat([ tableLeg, Lar.r(0,0,pi/2) ],outer=4) )
table = Lar.Struct([ tableTop, tablelegs ])
table = Lar.struct2lar(table)
View(table)

cylndr = Lar.rod(.06, .5, 2*pi)([8,1])
chairTop = Lar.Struct([ Lar.t(0,0,0.5), Lar.s(0.5,0.5,0.04), cube ])
chairLeg = Lar.Struct([ Lar.t(-.22,-.22,0), Lar.s(.5,.5,1), Lar.r(0,0,pi/8), cylndr ])
chairlegs = Lar.Struct( repeat([ chairLeg, Lar.r(0,0,pi/2) ],outer=4) );
chair = Lar.Struct([ chairTop, chairlegs ]);
chair = Lar.struct2lar(chair)
View(chair)

theChair = Lar.Struct([ Lar.t(-.8,0,0), chair ])
fourChairs = Lar.Struct( repeat([Lar.r(0,0,pi/2), theChair],outer=4) );
fourSit = Lar.Struct([fourChairs,table]);
View(fourSit)
singleRow=Lar.Struct(repeat([fourSit,Lar.t(0,2.5,0)],outer=10));
View(singleRow)
refectory=Lar.Struct(repeat([singleRow,Lar.t(3,0,0)],outer=10));
View(refectory)


text:
#=
void acb_elliptic_p(acb_t res, const acb_t z, const acb_t tau, slong prec)¶
Computes Weierstrass’s elliptic function

void acb_elliptic_inv_p(acb_t res, const acb_t z, const acb_t tau, slong prec)¶
Computes the inverse of the Weierstrass elliptic function

void acb_elliptic_zeta(acb_t res, const acb_t z, const acb_t tau, slong prec)¶
Computes the Weierstrass zeta function

void acb_elliptic_sigma(acb_t res, const acb_t z, const acb_t tau, slong prec)¶
Computes the Weierstrass sigma function
=#

"""
    weierstrass_p(z, tau)

- weierstrass_p(z, tau) == weierstrass_p(z+1, tau) == weierstrass_p(z+tau, tau)
"""
function weierstrass_p(z::ArbComplex{P}, tau::ArbComplex{P}) where {P}
    result = ArbComplex{P}()
    flag = 0
    ccall(@libarb(acb_elliptic_p), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint),
          result, z, tau, P)
    return result
end

"""
    weierstrass_invp(z, tau)

- weierstrass_p(weierstrass_invp(z, tau), tau) == z
"""
function weierstrass_invp(z::ArbComplex{P}, tau::ArbComplex{P}) where {P}
    result = ArbComplex{P}()
    flag = 0
    ccall(@libarb(acb_elliptic_inv_p), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint),
          result, z, tau, P)
    return result
end

"""
    weierstrass_zeta(z, tau)

- weierstrass_zeta(z+1, tau) == weierstrass_zeta(z, tau) + weierstrass_zeta(1/2, tau)
- weierstrass_zeta(z+tau, tau) == weierstrass_zeta(z, tau) + weierstrass_zeta(tau/2, tau)
"""
function weierstrass_zeta(z::ArbComplex{P}, tau::ArbComplex{P}) where {P}
    result = ArbComplex{P}()
    flag = 0
    ccall(@libarb(acb_elliptic_zeta), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint),
          result, z, tau, P)
    return result
end

"""
    weierstrass_sigma(z, tau)

- weierstrass_sigma(z+1, tau) == −exp(2*(z+τ/2)*weierstrass_zeta(τ/2,τ)) * weierstrass_sigma(z, tau)
"""
function weierstrass_sigma(z::ArbComplex{P}, tau::ArbComplex{P}) where {P}
    result = ArbComplex{P}()
    flag = 0
    ccall(@libarb(acb_elliptic_sigma), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint),
          result, z, tau, P)
    return result
end

weierstrass_p(z::ArbReal{P}, tau::ArbReal{P}) where {P} =
    weierstrass_p(ArbComplex{P}(z), ArbComplex{P}(tau))
weierstrass_p(z::ArbFloat{P}, tau::ArbFloat{P}) where {P} =
    weierstrass_p(ArbComplex{P}(z), ArbComplex{P}(tau))

weierstrass_pinv(z::ArbReal{P}, tau::ArbReal{P}) where {P} =
    weierstrass_pinv(ArbComplex{P}(z), ArbComplex{P}(tau))
weierstrass_pinv(z::ArbFloat{P}, tau::ArbFloat{P}) where {P} =
    weierstrass_pinv(ArbComplex{P}(z), ArbComplex{P}(tau))

weierstrass_zeta(z::ArbReal{P}, tau::ArbReal{P}) where {P} =
    weierstrass_zeta(ArbComplex{P}(z), ArbComplex{P}(tau))
weierstrass_zeta(z::ArbFloat{P}, tau::ArbFloat{P}) where {P} =
    weierstrass_zeta(ArbComplex{P}(z), ArbComplex{P}(tau))

weierstrass_sigma(z::ArbReal{P}, tau::ArbReal{P}) where {P} =
    weierstrass_sigma(ArbComplex{P}(z), ArbComplex{P}(tau))
weierstrass_sigma(z::ArbFloat{P}, tau::ArbFloat{P}) where {P} =
    weierstrass_sigma(ArbComplex{P}(z), ArbComplex{P}(tau))


text:
@testset "tref" begin
    # Test atomic values.
    @testset "atomic" begin
        function f()
            t = TRef(1)
            t[] = 0
            for _ in 1:6
                produce(t[])
                t[]
                t[] += 1
            end
        end

        ctask = TapedTask(f)

        consume(ctask)
        consume(ctask)

        a = copy(ctask)
        consume(a)
        consume(a)

        @test consume(ctask) == 2
        @test consume(a) == 4
    end

    # Test dictionary functionality.
    @testset "dictionary" begin
        function f()
            t = TRef(Dict("A" => 1, 5 => "B"))
            t["A"] = 0
            for _ in 1:6
                produce(t["A"])
                t["A"]
                t["A"] += 1
            end
        end

        ctask = TapedTask(f)

        consume(ctask)
        consume(ctask)

        a = copy(ctask)
        consume(a)
        consume(a)

        @test consume(ctask) == 2
        @test consume(a) == 4
    end

    @testset "array" begin
        # Create a TRef storing a matrix.
        x = TRef([1 2 3; 4 5 6])
        x[1, 3] = 900
        @test x[1,3] == 900

        # TRef holding an array.
        y = TRef([1,2,3])
        y[2] = 19
        @test y[2] == 19
    end
end

max_stars_repo_name: TuringLang/Libtask

max_issues_repo_name: TuringLang/Libtask

max_forks_repo_name: TuringLang/Libtask


text:
# using DataFrames

function get_fields
end

function get_records
end

function build_result
end

# interface for `Dict`s
get_fields(df :: Dict{Symbol, Vector{T} where T}) = collect(keys(df))
get_records(df :: Dict{Symbol, Vector{T} where T}) = zip((df[key] for key in keys(df))...)
function build_result(::Type{Dict{Symbol, Vector{T} where T}}, fields, typs, source :: Base.Generator)
    res = Tuple(typ[] for typ in typs)
    for each in source
        push!.(res, each)
    end
    Dict(zip(fields, res))
end

# interface for DataFrames

# get_fields(df :: DataFrame) = collect(names(df))
# get_records(df :: DataFrame) = DataFrames.columns(df)
# function build_result(::Type{DataFrame}, fields, typs, source :: Base.Generator)
#     res = Tuple(typ[] for typ in typs)
#     for each in source
#         push!.(res, each)
#     end
#     DataFrame(collect(res), fields)
# end


text:
@doc raw"""
    PiecewisePolynomialKernel(; degree::Int=0, dim::Int)
    PiecewisePolynomialKernel{degree}(dim::Int)

Piecewise polynomial kernel of degree `degree` for inputs of dimension `dim` with support in
the unit ball.

# Definition

For inputs ``x, x' \in \mathbb{R}^d`` of dimension ``d``, the piecewise polynomial kernel
of degree ``v \in \{0,1,2,3\}`` is defined as
```math
k(x, x'; v) = \max(1 - \|x - x'\|, 0)^{\alpha(v,d)} f_{v,d}(\|x - x'\|),
```
where ``\alpha(v, d) = \lfloor \frac{d}{2}\rfloor + 2v + 1`` and ``f_{v,d}`` are
polynomials of degree ``v`` given by
```math
\begin{aligned}
f_{0,d}(r) &= 1, \\
f_{1,d}(r) &= 1 + (j + 1) r, \\
f_{2,d}(r) &= 1 + (j + 2) r + \big((j^2 + 4j + 3) / 3\big) r^2, \\
f_{3,d}(r) &= 1 + (j + 3) r + \big((6 j^2 + 36j + 45) / 15\big) r^2 + \big((j^3 + 9 j^2 + 23j + 15) / 15\big) r^3,
\end{aligned}
```
where ``j = \lfloor \frac{d}{2}\rfloor + v + 1``.

The kernel is ``2v`` times continuously differentiable and the corresponding Gaussian
process is hence ``v`` times mean-square differentiable.
"""
struct PiecewisePolynomialKernel{D,C<:Tuple} <: SimpleKernel
    alpha::Int
    coeffs::C

    function PiecewisePolynomialKernel{D}(dim::Int) where {D}
        dim > 0 || error("number of dimensions has to be positive")
        j = div(dim, 2) + D + 1
        alpha = j + D
        coeffs = piecewise_polynomial_coefficients(Val(D), j)
        return new{D,typeof(coeffs)}(alpha, coeffs)
    end
end

# TODO: remove `maha` keyword argument in next breaking release
function PiecewisePolynomialKernel(; v::Int=-1, degree::Int=0, maha=nothing, dim::Int=-1)
    if v != -1
        Base.depwarn(
            "keyword argument `v` is deprecated, use `degree` instead",
            :PiecewisePolynomialKernel,
        )
        degree = v
    end

    if maha !== nothing
        Base.depwarn(
            "keyword argument `maha` is deprecated, use a `LinearTransform` instead",
            :PiecewisePolynomialKernel,
        )
        dim = size(maha, 1)
        return transform(
            PiecewisePolynomialKernel{degree}(dim), LinearTransform(cholesky(maha).U)
        )
    else
        return PiecewisePolynomialKernel{degree}(dim)
    end
end

piecewise_polynomial_coefficients(::Val{0}, ::Int) = (1,)
piecewise_polynomial_coefficients(::Val{1}, j::Int) = (1, j + 1)
piecewise_polynomial_coefficients(::Val{2}, j::Int) = (1, j + 2, (j^2 + 4 * j)//3 + 1)
function piecewise_polynomial_coefficients(::Val{3}, j::Int)
    return (1, j + 3, (2 * j^2 + 12 * j)//5 + 3, (j^3 + 9 * j^2 + 23 * j)//15 + 1)
end
function piecewise_polynomial_coefficients(::Val{D}, ::Int) where {D}
    return error("invalid degree $D, only 0, 1, 2, or 3 are supported")
end

kappa(κ::PiecewisePolynomialKernel, r) = max(1 - r, 0)^κ.alpha * evalpoly(r, κ.coeffs)

metric(::PiecewisePolynomialKernel) = Euclidean()

function Base.show(io::IO, κ::PiecewisePolynomialKernel{D}) where {D}
    return print(
        io,
        "Piecewise Polynomial Kernel (degree = ",
        D,
        ", ⌊dim/2⌋ = ",
        κ.alpha - 1 - 2 * D,
        ")",
    )
end


text:
function phpdf(ph::GPH{Tv}, t::TimeT; eps=Tv(1.0e-8), ufact=Tv(1.01)) where {TimeT,Tv}
    _phcalc(t, ph.alpha, ph.T, ph.tau, eps, ufact)
end

function phccdf(ph::GPH{Tv}, t::TimeT; eps=Tv(1.0e-8), ufact=Tv(1.01)) where {TimeT,Tv}
    _phcalc(t, ph.alpha, ph.T, ones(Tv, ph.dim), eps, ufact)
end

function phcdf(ph::GPH{Tv}, t::TimeT; eps=Tv(1.0e-8), ufact=Tv(1.01)) where {TimeT,Tv}
    alpha = zeros(Tv, ph.dim+1)
    tau = zeros(Tv, ph.dim+1)
    for i = eachindex(ph.alpha)
        alpha[i] = ph.alpha[i]
    end
    tau[ph.dim+1] = Tv(1)
    T = sparse(BlockCOO(2, 2, [(1, 1, ph.T), (1, 2, reshape(ph.tau, ph.dim, 1)), (2, 2, zeros(1,1))]))
    _phcalc(t, alpha, T, tau, eps, ufact)
end

function phpdf(ph::CF1{Tv}, t::TimeT, ::Type{MatT} = SparseMatrixCSC; eps=Tv(1.0e-8), ufact=Tv(1.01)) where {TimeT,Tv,MatT}
    phpdf(GPH(ph, MatT), t, eps=eps, ufact=ufact)
end

function phcdf(ph::CF1{Tv}, t::TimeT, ::Type{MatT} = SparseMatrixCSC; eps=Tv(1.0e-8), ufact=Tv(1.01)) where {TimeT,Tv,MatT}
    phcdf(GPH(ph, MatT), t, eps=eps, ufact=ufact)
end

function phccdf(ph::CF1{Tv}, t::TimeT, ::Type{MatT} = SparseMatrixCSC; eps=Tv(1.0e-8), ufact=Tv(1.01)) where {TimeT,Tv,MatT}
    phccdf(GPH(ph, MatT), t, eps=eps, ufact=ufact)
end

function _phcalc(t::Tv, alpha::Vector{Tv}, T::AbstractMatrix{Tv}, tau::Vector{Tv}, eps::Tv, ufact::Tv) where {Tv}
    P, qv = unif(T, ufact)
    right = rightbound(qv*t, eps)
    weight, poi = poipmf(qv*t, right, left=0)
    tmpv = zero(alpha)
    unifstep!(:T, P, poi, (0, right), weight, copy(alpha), tmpv)
    return @dot(tmpv, tau)
end

function _phcalc(t::AbstractVector{Tv}, alpha::Vector{Tv}, T::AbstractMatrix{Tv}, tau::Vector{Tv}, eps::Tv, ufact::Tv) where {Tv}
    P, qv = unif(T, ufact)
    perm = sortperm(t)
    dt, maxt = itime(t[perm])
    right = rightbound(qv*maxt, eps)
    poi = Vector{Tv}(undef, right + 1)
    m = length(t)
    result = Vector{Tv}(undef, m)
    vf = alpha
    tmpv = similar(alpha)
    for k = eachindex(dt)
        right = rightbound(qv*dt[k], eps)
        weight = poipmf!(qv*dt[k], poi, left=0, right=right)
        tmpv .= Tv(0)
        unifstep!(:T, P, poi, (0, right), weight, copy(vf), tmpv)
        vf .= tmpv
        result[perm[k]] = @dot(vf, tau)
    end
    return result
end


text:
using Distributions, Plots, Turing, DataFrames, MCMCChains, StatsPlots

mutable struct QLearningAgent
    q::Array{Real, 1}
    α::Real
    β::Real
    function QLearningAgent(α::Real, β::Real, K::Int64)
        q = Array{Real, 1}(zeros(K))
        new(q, α, β)
    end
end

function update(agent::QLearningAgent, action::Real, reward::Real)
    agent.q[action] += agent.α * (reward - agent.q[action])
end

function policy(agent::QLearningAgent)
    qmax = maximum(agent.q)
    qexp = exp.(agent.β .* (agent.q .- qmax))
    return qexp ./ sum(qexp)
end

struct Bandit
    p::Array{Real, 1}
end

function step(env::Bandit, action::Int64)
    k = length(env.p)
    rewards = env.p .>= rand(k)
    return rewards[action]
end

function step(agent::QLearningAgent, env::Bandit)
    p = policy(agent)
    action = rand(Categorical(p))
    reward = step(env, action)
    update(agent, action, reward)
    d = DataFrame(action = action, reward = reward)
    return hcat(d, DataFrame(agent.q', :auto))
end

function run(agent::QLearningAgent, env::Bandit, trial::Int64)
    return reduce(vcat, map(_ -> step(agent, env), 1:trial))
end

bandit = Bandit([0.2, 0.8])
k = length(bandit.p)
agent = QLearningAgent(0.05, 2., k)
result = run(agent, bandit, 500)

plot(result.x1, label="Q1")
plot!(result.x2, label="Q2")
scatter!(result.action .- 1,
         markersize=1,
         markershape=:vline,
         label="Response")

@model QLearningModel(actions::Array{Real, 1}, rewards::Array{Real, 1}, k::Int64) = begin
    T = length(actions)
    α ~ Beta(1, 1)
    β ~ Gamma(1, 100)
    agent = QLearningAgent(α, β, k)
    for t in 1:T
        action = actions[t]
        reward = rewards[t]
        p = policy(agent)
        actions[t] ~ Categorical(p)
        update(agent, action, reward)
    end
end

actions = Array{Real, 1}(result.action)
rewards = Array{Real, 1}(result.reward)
chains = sample(QLearningModel(actions, rewards, 2), NUTS(), 1000)
plot(chains)

chains = sample(QLearningModel(actions, rewards, 2), NUTS(), MCMCThreads(), 1000, 4)
plot(chains)


text:
module TestCandidatePool

import Wordle
import Test: @test, @test_throws

candidates = Wordle.candidate_pool(true)

@test eltype(candidates) === Wordle.Candidate
@test length(candidates) === 2315
@test String(candidates[1]) === "cigar"
@test String(candidates[2315]) === "shave"

candidates = Wordle.candidate_pool(false)

@test eltype(candidates) === Wordle.Candidate
@test length(candidates) > 2315

end # module


text:
using StaticArrays

abstract type AbstractSplitting end

export HamiltonianSplitting

"""
    HamiltonianSplitting( maxwell_solver,
                          kernel_smoother_0, kernel_smoother_1,
                          particle_group, e_dofs, b_dofs)

Hamiltonian splitting type for Vlasov-Maxwell

- Integral over the spline function on each interval (order p+1)
- Integral over the spline function on each interval (order p)
- `e_dofs` describing the two components of the electric field
- `b_dofs` describing the magnetic field
- `j_dofs` for kernel representation of current density.
"""
struct HamiltonianSplitting
    maxwell_solver::AbstractMaxwellSolver
    kernel_smoother_0::ParticleMeshCoupling1D
    kernel_smoother_1::ParticleMeshCoupling1D
    particle_group::ParticleGroup

    spline_degree::Int
    Lx::Float64
    x_min::Float64
    delta_x::Float64

    cell_integrals_0::SVector
    cell_integrals_1::SVector

    e_dofs::Array{Array{Float64,1}}
    b_dofs::Array{Float64,1}
    j_dofs::Array{Array{Float64,1}}

    buffer::Array{Array{Float64,1}}
    chunks::Iterators.PartitionIterator

    function HamiltonianSplitting(
        maxwell_solver, kernel_smoother_0, kernel_smoother_1, particle_group, e_dofs, b_dofs
    )

        # Check that n_dofs is the same for both kernel smoothers.
        @assert kernel_smoother_0.n_dofs == kernel_smoother_1.n_dofs

        j_dofs = [zeros(Float64, kernel_smoother_0.n_dofs) for i in 1:2]

        x_min = maxwell_solver.xmin
        Lx = maxwell_solver.Lx
        spline_degree = 3
        delta_x = Lx / kernel_smoother_1.n_dofs

        cell_integrals_1 = SVector{3}([0.5, 2.0, 0.5] ./ 3.0)
        cell_integrals_0 = SVector{4}([1.0, 11.0, 11.0, 1.0] ./ 24.0)

        np = particle_group.n_particles
        n_jobs = nthreads()

        @assert np % n_jobs == 0

        buffer = [zeros(Float64, kernel_smoother_0.n_dofs) for i in 1:n_jobs]

        chunks = Iterators.partition(1:np, np ÷ n_jobs)

        return new(
            maxwell_solver,
            kernel_smoother_0,
            kernel_smoother_1,
            particle_group,
            spline_degree,
            Lx,
            x_min,
            delta_x,
            cell_integrals_0,
            cell_integrals_1,
            e_dofs,
            b_dofs,
            j_dofs,
            buffer,
            chunks,
        )
    end
end

export strang_splitting!

"""
    strang_splitting( h, dt, number_steps)

Strang splitting
- time splitting object
- time step
- number of time steps
"""
function strang_splitting!(h::HamiltonianSplitting, dt::Float64, number_steps::Int)
    for i_step in 1:number_steps
        operatorHB(h, 0.5dt)
        operatorHE(h, 0.5dt)
        operatorHp2(h, 0.5dt)
        operatorHp1(h, 1.0dt)
        operatorHp2(h, 0.5dt)
        operatorHE(h, 0.5dt)
        operatorHB(h, 0.5dt)
    end
end

export lie_splitting!
"""
    lie_splitting( h, dt, number_steps)

Lie splitting
"""
function lie_splitting!(h::HamiltonianSplitting, dt::Float64, number_steps::Int)
    for i_step in 1:number_steps
        operatorHE(dt)
        operatorHB(dt)
        operatorHp1(dt)
        operatorHp2(dt)
    end
end

export lie_splitting_back!
"""
    lie_splitting_back(h, dt, number_steps)

Lie splitting (oposite ordering)
"""
function lie_splitting_back!(h::HamiltonianSplitting, dt::Float64, number_steps::Int)
    for i_step in 1:number_steps
        operatorHp2(dt)
        operatorHp1(dt)
        operatorHB(dt)
        operatorHE(dt)
    end
end

"""

    operatorHp1(h, dt)

```math
\\begin{aligned}
\\partial_t f + v_1 \\partial_{x_1} f = 0 & \\rightarrow  X_{new} = X_{old} + dt V_1 \\\\
V_{new},2 = V_{old},2 + \\int_0 h V_{old},1 B_{old} & \\\\
\\partial_t E_1 = - \\int v_1 f(t,x_1, v) dv & \\rightarrow E_{1,new} = E_{1,old} -
\\int \\int v_1 f(t,x_1+s v_1,v) dv ds  & \\\\
\\partial_t E_2 = 0 & \\rightarrow E_{2,new} = E_{2,old} \\\\
\\partial_t B = 0 & \\rightarrow B_{new} = B_{old}
\\end{aligned}
```

Here we have to accumulate j and integrate over the time interval.
At each ``k=1,...,n_{grid}``, we have for ``s \\in [0,dt]:
j_k(s) =  \\sum_{i=1,..,N_p} q_i N((x_k+sv_{1,k}-x_i)/h) v_k``,
where ``h`` is the grid spacing and ``N`` the normalized B-spline
 In order to accumulate the integrated ``j``, we normalize the values of
``x`` to the grid spacing,
    calling them ``y``, we have
```math
 j_k(s) = \\sum_{i=1,..,N_p} q_i N(y_k+\\frac{s}{h} v_{1,k}-y_i) v_k.
```
 Now, we want the integral

```math
\\int_{0..dt} j_k(s) d s = \\sum_{i=1}^{N_p} q_i v_k
\\int_{0}{dt} N(y_k+\\frac{s}{h} v_{1,k}-y_i) ds
=  \\sum_{i=1}^{N_p} q_i v_k  \\int_{0}^{dt}  N(y_k + w v_{1,k}-y_i) dw
```

For each particle compute the index of the first DoF on the grid it contributes
to and its position (normalized to cell size one). Note: `j_dofs`
does not hold the values for `j` itself but for the integrated `j`.

Then update particle position:
``X_{new} = X_{old} + dt * V``
"""
function operatorHp1(h::HamiltonianSplitting, dt::Float64)
    np = h.particle_group.n_particles
    n_cells = h.kernel_smoother_0.n_dofs

    fill!(h.j_dofs[1], 0.0)
    fill!(h.j_dofs[2], 0.0)

    @sync for i_chunk in h.chunks
        @spawn begin
            fill!(h.buffer[threadid()], 0.0)
            for i_part in i_chunk

                # Read out particle position and velocity
                x_old = h.particle_group.array[1, i_part]
                v1_old = h.particle_group.array[2, i_part]
                v2_old = h.particle_group.array[3, i_part]

                # Then update particle position:  X_new = X_old + dt * V
                x_new = x_old + dt * v1_old

                # Get charge for accumulation of j
                wi = h.particle_group.array[4, i_part]
                wi = wi * h.particle_group.charge
                wi = wi * h.particle_group.common_weight

                qoverm = h.particle_group.q_over_m

                v2_new = add_current_update_v!(
                    h.buffer[threadid()],
                    h.kernel_smoother_1,
                    x_old,
                    x_new,
                    wi,
                    qoverm,
                    h.b_dofs,
                    v2_old,
                )

                x_new = mod(x_new, h.Lx)

                h.particle_group.array[1, i_part] = x_new
                h.particle_group.array[3, i_part] = v2_new
            end
        end
    end

    h.j_dofs[1] .= reduce(+, h.buffer)

    @sync for i_chunk in h.chunks
        @spawn begin
            fill!(h.buffer[threadid()], 0.0)
            for i_part in i_chunk

                # Read out particle position and charge
                x = h.particle_group.array[1, i_part]
                w = h.particle_group.array[4, i_part]
                w = w * h.particle_group.charge
                w = w * h.particle_group.common_weight

                # Accumulate rho for Poisson diagnostics
                add_charge!(h.buffer[threadid()], h.kernel_smoother_0, x, w)
            end
        end
    end

    h.j_dofs[2] .= reduce(+, h.buffer)

    # Update the electric field.
    return compute_e_from_j!(h.e_dofs[1], h.maxwell_solver, h.j_dofs[1], 1)
end

"""
    operatorHp2(h, dt)

Push Hp2: Equations to solve are

```math
\\begin{aligned} X_{new}  =  X_{old} & \\\\
V_{new,1} = V_{old,1} + \\int_0 h V_{old,2} B_{old} & \\\\
\\partial_t E_1 = 0 & \\rightarrow  E_{1,new} = E_{1,old}  \\\\
\\partial_t E_2 = - \\int v_2 f(t,x_1, v) dv & \\rightarrow
E_{2,new} = E_{2,old} - \\int \\int v_2 f(t,x_1 + s v_1,v) dv ds\\\\
\\partial_t B = 0 & => & B_{new} = B_{old} & \\\\
\\end{aligned}
```
"""
function operatorHp2(h::HamiltonianSplitting, dt::Float64)
    n_cells = h.kernel_smoother_0.n_dofs

    fill!(h.j_dofs[1], 0.0)
    fill!(h.j_dofs[2], 0.0)

    qm = h.particle_group.q_over_m

    np = h.particle_group.n_particles

    # Update v_1

    @sync for i_chunk in h.chunks
        @spawn begin
            fill!(h.buffer[threadid()], 0.0)
            for i_part in i_chunk

                # Evaluate b at particle position (splines of order p)
                x1 = h.particle_group.array[1, i_part]
                v1 = h.particle_group.array[2, i_part]
                v2 = h.particle_group.array[3, i_part]

                b = evaluate(h.kernel_smoother_1, x1, h.b_dofs)
                v1 = v1 + dt * qm * v2 * b

                h.particle_group.array[2, i_part] = v1

                # Scale vi by weight to combine both factors
                #for accumulation of integral over j
                w = h.particle_group.array[4, i_part]
                w = w * h.particle_group.charge
                w = w * h.particle_group.common_weight
                w = w * v2

                add_charge!(h.buffer[threadid()], h.kernel_smoother_0, x1, w)
            end
        end
    end

    h.j_dofs[2] .= reduce(+, h.buffer)

    # Update the electric field. Also, we still need to scale with 1/Lx

    h.j_dofs[2] .= h.j_dofs[2] .* dt

    return compute_e_from_j!(h.e_dofs[2], h.maxwell_solver, h.j_dofs[2], 2)
end

"""
    operatorHE(h, dt)

Push ``H_E``: Equations to be solved
```math
\\begin{aligned}
\\partial_t f + E_1 \\partial_{v_1} f + E_2 \\partial_{v_2} f = 0 &\\rightarrow& V_{new} = V_{old} + dt * E \\\\
\\partial_t E_1 = 0 &\\rightarrow& E_{1,new} = E_{1,old} \\\\
\\partial_t E_2 = 0 &\\rightarrow& E_{2,new} = E_{2,old} \\\\
\\partial_t B + \\partial_{x_1} E_2 = 0 &\\rightarrow& B_{new} = B_{old} - dt \\partial_{x_1} E_2
\\end{aligned}
```
"""
function operatorHE(h::HamiltonianSplitting, dt::Float64)
    qm = h.particle_group.q_over_m

    np = h.particle_group.n_particles

    # V_new = V_old + dt * E

    @sync for i_chunk in h.chunks
        @spawn begin
            for i_part in i_chunk
                v_old1 = h.particle_group.array[2, i_part]
                v_old2 = h.particle_group.array[3, i_part]

                # Evaluate efields at particle position

                xi = h.particle_group.array[1, i_part]

                e1 = evaluate(h.kernel_smoother_1, xi, h.e_dofs[1])
                e2 = evaluate(h.kernel_smoother_0, xi, h.e_dofs[2])

                v_new1 = v_old1 + dt * qm * e1
                v_new2 = v_old2 + dt * qm * e2

                h.particle_group.array[2, i_part] = v_new1
                h.particle_group.array[3, i_part] = v_new2
            end
        end
    end

    # Update bfield
    return compute_b_from_e!(h.b_dofs, h.maxwell_solver, dt, h.e_dofs[2])
end

@doc raw"""
    operatorHB(h, dt)

Push ``H_B``: Equations to be solved ``V_{new} = V_{old}``

```math
\begin{aligned}
\partial_t E_1 = 0 & \rightarrow & E_{1,new} = E_{1,old} \\
\partial_t E_2 = - \partial_{x_1} B & \rightarrow & E_{2,new} = E_{2,old}-dt*\partial_{x_1} B \\
\partial_t B = 0 & \rightarrow & B_{new} = B_{old} \\
\end{aligned}
```
"""
function operatorHB(h::HamiltonianSplitting, dt::Float64)
    return compute_e_from_b!(h.e_dofs[2], h.maxwell_solver, dt, h.b_dofs)
end


text:
include("preparephi.jl")
include("coefexp_ab.jl")
"""
    PrepareTwoScalesPureAB(nb_t, t_max, order, par_u0::PrepareU0)

Immutable structure, to share calculations, needed for the twoscale function

# Arguments :
- `nb_t::Int` : number of time slices
- `t_max`: end of the time
- `order` : order for compute the coefficients
- `par_u0::PrepareU0` : prepared initial data

# Keywords :
- `p_coef::Union{CoefExpAB,Missing}=missing` : precomputed coefficients of AB method
- `verbose=100` : trace level

# Fields :
- nb_t : number of time slices
- t_max : end of the time
- order : order for compute the coefficients
- parphi : prepared parameters for phi (from par_u0)
- par_u0 : prepared initial data
- p_coef : computed coefficients
- exptau : exp( -im*dt*'\tau'/epsilon) for all '\tau' values
- exptau_inv : inverse of exptau
- verbose : trace level

"""
struct PrepareTwoScalesPureAB
    nb_t::Any
    t_max::Any
    dt::Any
    order::Any
    parphi::PreparePhi
    par_u0::PrepareU0
    p_coef::CoefExpAB
    exptau::Any
    exptau_inv::Any
    verbose::Any
    function PrepareTwoScalesPureAB(
        nb_t,
        t_max,
        order,
        par_u0::PrepareU0;
        p_coef::Union{CoefExpAB,Missing} = missing,
        verbose = 100,
    )
        parphi = par_u0.parphi
        T = typeof(parphi.epsilon)
        dt = T(t_max - parphi.t_0) / nb_t
        p_coef = if ismissing(p_coef)
            CoefExpAB(order, parphi.epsilon, parphi.n_tau, dt)
        else
            p_coef
        end
        exptau = collect(transpose(exp.(-im * dt / parphi.epsilon * parphi.tau_list)))
        exptau_inv = collect(transpose(exp.(im * dt / parphi.epsilon * parphi.tau_list)))
        return new(
            nb_t,
            t_max,
            dt,
            order,
            parphi,
            par_u0,
            p_coef,
            exptau,
            exptau_inv,
            verbose,
        )
    end
end

function _calculfft(parphi::PreparePhi, u_caret)
    f = filtredfct(parphi, real(ifftgen(parphi.par_fft, u_caret)))
    return fftgen(parphi.par_fft, f)
end

function _calcul_ab(par::PrepareTwoScalesPureAB, ord, fftfct, fft_u, dec, sens)
    resfft = fft_u[dec-sens] .* ((sens == 1) ? par.exptau : par.exptau_inv)
    tab_coef = (sens == 1) ? par.p_coef.tab_coef : par.p_coef.tab_coef_neg
    for k = 1:ord
        indice = dec - sens * k
        resfft .+= transpose(tab_coef[:, k, ord]) .* fftfct[indice]
    end
    fft_u[dec] = resfft
    if par.verbose >= 200 && isexactsol(par.parphi)
        t = par.parphi.t_0 + (dec - par.order) * par.dt
        u_exact = getexactsol(par.parphi, par.par_u0.up0[1:(end-1)], t)
        u_computed = _getresult(fft_u[dec], t - par.parphi.t_0, par.parphi)
        err = Float64(norm(u_exact - u_computed, Inf))
        i = dec - par.order
        traceln(200, "i=$i err=$err", verbose = par.verbose)
    end
    fftfct[dec] = _calculfft(par.parphi, resfft)
end

# function _init_ab(par::PrepareTwoScalesPureAB, fftfct, fft_u)
# #    println("_init_ab order=$(par.order)")
#     if par.order != 1
#         for new_ord=2:par.order
#             _calcul_ab(par, new_ord-1, fftfct, fft_u, par.order+new_ord-1, 1)
#             for k = 1:new_ord-1
#                 _calcul_ab(par, new_ord, fftfct, fft_u, par.order-k, -1)
#             end
#             for k = 1:new_ord-1
#                 _calcul_ab(par, new_ord, fftfct, fft_u, par.order+k, 1)
#             end
#         end
#     end
# end
function _init_ab(par::PrepareTwoScalesPureAB, fftfct, fft_u)
    #    println("_init_ab order=$(par.order)")
    if par.order != 1
        for new_ord = 2:par.order
            for k = 1:new_ord-1
                _calcul_ab(par, new_ord - 1, fftfct, fft_u, par.order - k, -1)
            end
            for k = 1:new_ord-1
                _calcul_ab(par, new_ord, fftfct, fft_u, par.order + k, 1)
            end
        end
    end
end

function _tr_ab(par::PrepareTwoScalesPureAB, fftfct, u_chap)
    resu_c = par.exptau .* u_chap
    bound = par.order - 1
    for k = 0:bound
        resu_c .+= transpose(par.p_coef.tab_coef[:, k+1, par.order]) .* fftfct[end-k]
    end
    resfft = _calculfft(par.parphi, resu_c)
    return resfft, resu_c
end
function traceln(refniv, str; verbose = 100)::Nothing
    if verbose >= refniv
        println(str)
    end
end
function trace(refniv, str; verbose = 100)::Nothing
    if verbose >= refniv
        print(str)
    end
end

# for this function only, t is the time from the beginning
function _getresult(u_chap, t, par::PreparePhi)
    # matlab : u1=real(sum(fft(ut)/Ntau.*exp(1i*Ltau*T/ep)));
    u1 = real(u_chap * exp.(1im * par.tau_list * t / par.epsilon)) / par.n_tau
    res = exp(t / par.epsilon * par.sparse_Ap) * u1
    return res[1:(end-1)]
end
function _getresult(tab_u_chap, t, par::PreparePhi, t_begin, t_max, order)
    nb = size(tab_u_chap, 1) - 1
    dt = (t_max - t_begin) / nb
    t_ex = (t - t_begin) / dt
    t_int = convert(Int64, floor(t_ex))
    t_int_begin = t_int - div(order, 2)
    if t_int_begin < 0
        t_int_begin = 0
    end
    if t_int_begin > nb - order
        t_int_begin = nb - order
    end
    t_ex -= t_int_begin
    t1 = t_int_begin + 1
    N = (typeof(par.epsilon) == BigFloat) ? BigInt : Int64
    u_chap = interpolate(tab_u_chap[t1:(t1+order)], order, t_ex, N)
    return _getresult(u_chap, t - t_begin, par)
end
function _getresult(tab_t, tab_u_chap, t, par::PreparePhi, t_begin, t_max, order)
    nb = size(tab_u_chap, 1) - 1
    dt = (t_max - t_begin) / nb
    t_ex = (t - t_begin) / dt
    t_int = convert(Int64, floor(t_ex))
    t_int_begin = t_int - div(order, 2)
    if t_int_begin < 0
        t_int_begin = 0
    end
    if t_int_begin > nb - order
        t_int_begin = nb - order
    end
    t1 = t_int_begin + 1
    u_chap = interpolate(tab_t[t1:(t1+order)], tab_u_chap[t1:(t1+order)], order, t)
    return _getresult(u_chap, t - t_begin, par)
end
"""
    twoscales_pure_ab(par::PrepareTwoScalesPureAB; only_end::Bool=false, diff_fft::Bool=false, res_fft::Bool=false, verbose::Integer=100)

compute the data to get solution of the differential equation

# Arguments :
- `par::PrepareTwoScalesPureAB` : contains all the parameters and prepared data

# Keywords :
- `only_end=false` : if true return only the result for t_end
- `diff_fft::Bool=false` : if true return data about diff
- `res_fft::Bool=false` : if true return u_caret data indispensable for interpolation
- `verbose::Integer`: level off traces (0 means no output)

"""
function twoscales_pure_ab(
    par::PrepareTwoScalesPureAB;
    only_end::Bool = false,
    diff_fft::Bool = false,
    res_fft::Bool = false,
    verbose::Integer = 100,
)
    levelref = 90
    T = typeof(par.parphi.epsilon)
    fftfct = Vector{Array{Complex{T},2}}(undef, 2par.order - 1)
    fft_u = Vector{Array{Complex{T},2}}(undef, 2par.order - 1)
    res_u = par.par_u0.ut0
    up0 = par.par_u0.up0
    fftfct[par.order] = fftgen(par.parphi.par_fft, filtredfct(par.parphi, res_u))
    fft_u[par.order] = fftgen(par.parphi.par_fft, res_u)
    result = zeros(typeof(par.parphi.epsilon), par.parphi.size_vect - 1, par.nb_t + 1)
    result[:, 1] = up0[1:(end-1)]
    result_fft = undef
    _init_ab(par, fftfct, fft_u)
    memfft = fftfct[par.order:(2par.order-1)]
    if res_fft
        result_fft = Vector{Array{Complex{T},2}}(undef, par.nb_t + 1)
        res_u_chap = Vector{Array{Complex{T},2}}(undef, par.nb_t + 1)
        for i = 1:min(par.order, par.nb_t + 1)
            result_fft[i] = memfft[i]
        end
        for i = 1:min(par.order, par.nb_t + 1)
            res_u_chap[i] = fft_u[i+par.order-1]
        end
    end
    if !only_end
        for i = 2:min(par.order, par.nb_t + 1)
            result[:, i] = _getresult(fft_u[par.order+i-1], (i - 1) * par.dt, par.parphi)
        end
    end
    tabdifffft = undef
    tabdifffft_2 = undef
    if diff_fft
        tabdifffft = zeros(BigFloat, par.nb_t)
        tabdifffft_2 = zeros(BigFloat, par.nb_t)
        for i = 1:(par.order-1)
            tabdifffft[i] = norm(memfft[i] - memfft[i+1], Inf)
            tabdifffft_2[i] = norm(memfft[i] - memfft[i+1])
        end
    end
    # ring permutation where the beginning becomes the end and the rest is shifted by one
    permut = collect(Iterators.flatten((2:par.order, 1:1)))
    ut0_fft = fft_u[par.order-1+min(par.order, par.nb_t + 1)]
    traceln(levelref, "", verbose = verbose)
    norm_delta_fft = 0
    norm_delta_fft_2 = 0
    nbnan = 0
    borne_nm = 0
    c_mult = 1.1
    for i = par.order:par.nb_t
        resfft, ut0_fft = _tr_ab(par, memfft, ut0_fft)
        if res_fft
            result_fft[i] = resfft
            #            res_u_chap[i+par.order] = ut0_fft
            res_u_chap[i+1] = ut0_fft
        end
        if !only_end
            result[:, i+1] = _getresult(ut0_fft, i * par.dt, par.parphi)
        end
        if diff_fft
            nm = norm(resfft - memfft[end], Inf)
            nm_2 = norm(resfft - memfft[end])
            norm_delta_fft = max(nm, norm_delta_fft)
            norm_delta_fft_2 = max(nm_2, norm_delta_fft_2)
            tabdifffft[i] = nm
            tabdifffft_2[i] = nm_2
        end
        memfft = memfft[permut]
        memfft[end] = resfft
        if i % 100 == 0
            trace(levelref, "x", verbose = verbose)
            #            GC.gc()
        end
        if i % 10000 == 0 || i == par.nb_t
            traceln(levelref, " $i/$(par.nb_t)", verbose = verbose)
        end
    end

    #    println("norm diff fft = $norm_delta_fft")

    ret = only_end ? _getresult(ut0_fft, par.t_max - par.parphi.t_0, par.parphi) : result

    if res_fft
        if diff_fft
            return ret,
            result_fft,
            res_u_chap,
            tabdifffft,
            tabdifffft_2,
            norm_delta_fft,
            norm_delta_fft_2
        else
            return ret, result_fft, res_u_chap
        end
    else
        if diff_fft
            return ret, tabdifffft, tabdifffft_2, norm_delta_fft, norm_delta_fft_2
        else
            return ret
        end
    end
end


text:
"""
Reject epochs based on the maximum peak to peak voltage within an epoch across all channels

#### Arguments

* `epochs`: Array containing the epoch data in the format samples x epochs x channels
* `retain_percentage`: The percentage of epochs to retain
* `rejection_method`: Method to be used for epoch rejection (peak2peak)

#### Returns

* An array with a reduced amount of entries in the epochs dimension
"""
function epoch_rejection(
    epochs::Array{T,3},
    retain_percentage::AbstractFloat;
    rejection_method::Function = Neuroimaging.peak2peak,
) where {T<:Number}

    if (0 > retain_percentage) || (1 < retain_percentage)
        @warn("Non valid percentage value for retaining epochs $(retain_percentage)")
    end

    @info(
        "Rejected $(round.(Int, (1 - retain_percentage) * 100))% of epochs based on $(string(rejection_method))"
    )

    # Epoch value should be a value or score per epoch where a lower value is better
    # The lowest `retain_percentage` amount of epoch values will be kept
    epoch_values = rejection_method(epochs)

    cut_off_value = sort(epoch_values)[floor(Int, length(epoch_values) * retain_percentage)]
    epochs = epochs[:, epoch_values.<=cut_off_value, :]
end

"""
Find the peak to peak value for each epoch to be returned to epoch_rejection()
"""
function peak2peak(epochs)

    epochsNum = size(epochs)[2]

    peakvalues = Array{AbstractFloat}(undef, epochsNum)
    for epoch = 1:epochsNum
        peakvalues[epoch] =
            abs.(maximum(epochs[:, epoch, :]) - minimum(epochs[:, epoch, :]))
    end

    return peakvalues
end

"""
Reject channels with too great a variance.

Rejection can be based on a threshold or dynamicly chosen based on the variation of all channels.

#### Arguments

* `signals`: Array of data in format samples x channels
* `threshold_abs`: Absolute threshold to remove channels with variance above this value
* `threshold_std`: Reject channels with a variance more than n times the std of all channels

#### Returns

An array indicating the channels to be kept
"""
function channel_rejection(
    sigs::Array{T},
    threshold_abs::Number,
    threshold_var::Number,
) where {T<:Number}

    @debug(
        "Rejecting channels for signal of $(size(sigs,2)) chanels and $(size(sigs,1)) samples"
    )

    variances = var(sigs, dims = 1)        # Determine the variance of each channel
    valid_nonzero = variances .!= 0    # The reference channel will have a variance of 0 so ignore it

    # Reject channels above the threshold
    valid_threshold_abs = variances .< threshold_abs
    @debug("Static rejection threshold: $(threshold_abs)")

    # Reject channels outside median + n * std
    variances_median = median(variances[valid_nonzero])    # Use the median as usually not normal
    variances_std = std(variances[valid_nonzero])       # And ignore the reference channel
    valid_threshold_var = variances .< (variances_median + threshold_var * variances_std)
    @debug(
        "Dynamic rejection threshold: $(variances_median + threshold_var * variances_std)"
    )

    valid_nonzero .& valid_threshold_abs .& valid_threshold_var   # Merge all methods
end


text:
# The sequence of triangle numbers is generated by adding the natural numbers.
# So the 7th triangle number would be 1 + 2 + 3 + 4 + 5 + 6 + 7 = 28. The
# first ten terms would be:
#
# 1, 3, 6, 10, 15, 21, 28, 36, 45, 55, ...
#
# Let us list the factors of the first seven triangle numbers:
#
#  1: 1
#  3: 1,3
#  6: 1,2,3,6
# 10: 1,2,5,10
# 15: 1,3,5,15
# 21: 1,3,7,21
# 28: 1,2,4,7,14,28
# We can see that 28 is the first triangle number to have over five divisors.
#
# What is the value of the first triangle number to have over five hundred
# divisors?

using ProjectEulerSolutions

# Cycle through triangular numbers, starting with a initial value using a high
# value starting point and incrementing the triangular number each iteration.
function p012solution_increment(numdivisors::Integer=1)::Integer

    i = 10000
    # Uses the sum of series equation, allowing it to start at a larger number
    trinumber::Integer = i * (i - 1) / 2
    while length(factors(trinumber)) + 1 <= numdivisors
        trinumber += i
        i += 1
    end

    return trinumber
end

# Cycle through triangular numbers, setting the triangular number at each step
# using the sum of series equation.
function p012solution_sum(numdivisors::Integer=1)::Integer

    i = 10000
    trinumber::Integer = 1
    while length(factors(trinumber)) + 1 <= numdivisors
        # Uses the sum of series equation, allowing it to start at a larger number
        trinumber = i * (i + 1) / 2
        i += 1
    end

    return trinumber
end

p012 = Problems.Problem(Dict("Increment" => p012solution_increment,
                             "Sum" => p012solution_sum))

Problems.benchmark(p012, 500)

max_stars_repo_name: gnujosh/julia-euler-project

max_issues_repo_name: gnujosh/julia-euler-project

max_forks_repo_name: gnujosh/julia-euler-project


text:
using DrWatson
@quickactivate :Swalbe
using Plots, CUDA, DataFrames, BSON
# CUDA.device!(1)

# Fluid dynamics we need for the experiment
"""
    measure_substratewave

Simulates a changing wettability and the response from the film.
"""
function measure_dropcoalesence(
    sys::Swalbe.SysConst,
    device::String;
    R₀ = 1.0,
    dist = 1,
    dump = 1000,
    θₛ=ones(sys.Lx, sys.Ly),
    fluid=zeros(sys.Tmax÷dump, sys.Lx*sys.Ly),
    Svelx=zeros(sys.Tmax÷dump, sys.Lx*sys.Ly),
    Svely=zeros(sys.Tmax÷dump, sys.Lx*sys.Ly),
    verbos=true,
    T=Float64
)
    println("Simulating two coalesing droplets")
    fout, ftemp, feq, height, velx, vely, vsq, pressure, dgrad, Fx, Fy, slipx, slipy, h∇px, h∇py = Swalbe.Sys(sys, device, false, T)
    if device == "CPU"
        for i in 1:sys.Lx, j in 1:sys.Ly
            height[i,j] = h₀ + ϵ * sin(2π*wave_x*(i-1)/sys.Lx) * sin(2π*wave_y*(j-1)/sys.Ly)
        end
    elseif device == "GPU"
        h = zeros(size(height))
        for i in 1:sys.Lx, j in 1:sys.Ly
            h[i,j] = h₀ + ϵ * sin(2π*wave_x*(i-1)/sys.Lx) * sin(2π*wave_y*(j-1)/sys.Ly)
        end
        # theta = CUDA.zeros(Float64, sys.Lx, sys.Ly)
        height = CUDA.adapt(CuArray, h)
    end
    Swalbe.equilibrium!(fout, height, velx, vely, vsq)
    ftemp .= fout
    for t in 1:sys.Tmax
        if t % sys.tdump == 0
            mass = 0.0
            mass = sum(height)
            if verbos
                println("Time step $t mass is $(round(mass, digits=3))")
            end
        end

        Swalbe.filmpressure!(pressure, height, dgrad, sys.γ, θₛ, sys.n, sys.m, sys.hmin, sys.hcrit)
        Swalbe.∇f!(h∇px, h∇py, pressure, dgrad, height)
        Swalbe.slippage!(slipx, slipy, height, velx, vely, sys.δ, sys.μ)
        # Forces are the pressure gradient and the slippage due to substrate liquid boundary conditions
        Fx .= h∇px .+ slipx
        Fy .= h∇py .+ slipy
        # New equilibrium
        Swalbe.equilibrium!(feq, height, velx, vely, vsq)
        Swalbe.BGKandStream!(fout, feq, ftemp, -Fx, -Fy)
        # New moments
        Swalbe.moments!(height, velx, vely, fout)
        # Measurements, in this case only snapshots of simulational arrays
        Swalbe.snapshot!(fluid, height, t, dumping = dump)
        Swalbe.snapshot!(Svelx, velx, t, dumping = dump)
        Swalbe.snapshot!(Svelx, vely, t, dumping = dump)
        Swalbe.snapshot!(theta, θₛ, t, dumping = dump)
        move_substrate!(slipx, θₛ, t, sub_speed, direction=dire)
    end
    return fluid, Svelx, Svely, theta
    CUDA.reclaim()
end

"""
    twodroplets(height, R_drop1, R_drop2, θ, center_drop1, center_drop2)

Initial condition to have two droplets next to each other.
"""
function twodroplets!(height, R_drop1, R_drop2, θ_drop1, θ_drop2, center_drop1, center_drop2)
    # A lot to do here!
    lx, ly = size(height)
    h1 = zeros(lx, ly)
    h2 = zeros(lx, ly)
    for i in 1:lx, j in 1:ly
        circ1 = sqrt((i-center_drop1[1])^2 + (j-center_drop1[2])^2)
        circ2 = sqrt((i-center_drop2[1])^2 + (j-center_drop2[2])^2)
        if circ1 <= R_drop1
            h1[i,j] = (cos(asin(circ1/R_drop1)) - cospi(θ_drop1)) * R_drop1
        elseif circ2 <= R_drop2
            h1[i,j] += (cos(asin(circ2/R_drop2)) - cospi(θ_drop2)) * R_drop2
        else
            h1[i,j] = 0.05
        end
    end

    return h1, h2
end

sys = Swalbe.SysConst(Lx=512, Ly=512, γ=0.01, δ=1.0, n=3, m=2, hmin=0.07, Tmax=10000, tdump=1000)
df_fluid = DataFrame()
df_velx = DataFrame()
df_vely = DataFrame()
df_sub = DataFrame()
θₚ = fill(1/9, sys.Lx,sys.Ly)
# Make a cuarray with the substrate pattern
θ_in = CUDA.adapt(CuArray, θₚ)
# Actual simulation
fluid, velx, vely, substrate = measure_substratewave(sys, "GPU", "no", sub_speed=speed, θₛ=θ_in, dire=direction, dump=sys.tdump)
println("Writing measurements to dataframes")
# Filling the dataframes
for t in 1:sys.Tmax÷sys.tdump
    df_fluid[!, Symbol("h_$(t*sys.tdump)")] = fluid[t,:]
    df_velx[!, Symbol("vx_$(t*sys.tdump)")] = velx[t,:]
    df_vely[!, Symbol("vy_$(t*sys.tdump)")] = vely[t,:]
    df_sub[!, Symbol("theta_$(t*sys.tdump)")] = substrate[t,:]
end
folder = "Drop_bridge"
println("Saving dataframe subdirection $direction subvel $speed and sines $waves to disk")
bson("data/$(folder)/height_direc_$(direction)_sp_$(speed)_sines_$(waves).bson", Dict(:f => df_fluid))
bson("data/$(folder)/velx_direc_$(direction)_sp_$(speed)_sines_$(waves).bson", Dict(:v => df_velx))
bson("data/$(folder)/vely_direc_$(direction)_sp_$(speed)_sines_$(waves).bson", Dict(:u => df_vely))
bson("data/$(folder)/theta_direc_$(direction)_sp_$(speed)_sines_$(waves).bson", Dict(:t => df_sub))
CUDA.reclaim()

println("Script done, let's have a look at the data :)")


text:
# By default, Julia/LLVM does not use fused multiply-add operations (FMAs).
# Since these FMAs can increase the performance of many numerical algorithms,
# we need to opt-in explicitly.
# See https://ranocha.de/blog/Optimizing_EC_Trixi for further details.
@muladd begin

abstract type AbstractIndicator end

function create_cache(typ::Type{IndicatorType}, semi) where {IndicatorType<:AbstractIndicator}
  create_cache(typ, mesh_equations_solver_cache(semi)...)
end

function get_element_variables!(element_variables, indicator::AbstractIndicator, ::VolumeIntegralShockCapturingHG)
  element_variables[:indicator_shock_capturing] = indicator.cache.alpha
  return nothing
end

"""
    IndicatorHennemannGassner

Indicator used for shock-capturing or AMR used by
- Hennemann, Gassner (2020)
  "A provably entropy stable subcell shock capturing approach for high order split form DG"
  [arXiv: 2008.12044](https://arxiv.org/abs/2008.12044)
"""
struct IndicatorHennemannGassner{RealT<:Real, Variable, Cache} <: AbstractIndicator
  alpha_max::RealT
  alpha_min::RealT
  alpha_smooth::Bool
  variable::Variable
  cache::Cache
end

# this method is used when the indicator is constructed as for shock-capturing volume integrals
function IndicatorHennemannGassner(equations::AbstractEquations, basis;
                                   alpha_max=0.5,
                                   alpha_min=0.001,
                                   alpha_smooth=true,
                                   variable)
  alpha_max, alpha_min = promote(alpha_max, alpha_min)
  cache = create_cache(IndicatorHennemannGassner, equations, basis)
  IndicatorHennemannGassner{typeof(alpha_max), typeof(variable), typeof(cache)}(
    alpha_max, alpha_min, alpha_smooth, variable, cache)
end

# this method is used when the indicator is constructed as for AMR
function IndicatorHennemannGassner(semi::AbstractSemidiscretization;
                                   alpha_max=0.5,
                                   alpha_min=0.001,
                                   alpha_smooth=true,
                                   variable)
  alpha_max, alpha_min = promote(alpha_max, alpha_min)
  cache = create_cache(IndicatorHennemannGassner, semi)
  IndicatorHennemannGassner{typeof(alpha_max), typeof(variable), typeof(cache)}(
    alpha_max, alpha_min, alpha_smooth, variable, cache)
end

function Base.show(io::IO, indicator::IndicatorHennemannGassner)
  @nospecialize indicator # reduce precompilation time

  print(io, "IndicatorHennemannGassner(")
  print(io, indicator.variable)
  print(io, ", alpha_max=", indicator.alpha_max)
  print(io, ", alpha_min=", indicator.alpha_min)
  print(io, ", alpha_smooth=", indicator.alpha_smooth)
  print(io, ")")
end

function Base.show(io::IO, ::MIME"text/plain", indicator::IndicatorHennemannGassner)
  @nospecialize indicator # reduce precompilation time

  if get(io, :compact, false)
    show(io, indicator)
  else
    setup = [
             "indicator variable" => indicator.variable,
             "max. α" => indicator.alpha_max,
             "min. α" => indicator.alpha_min,
             "smooth α" => (indicator.alpha_smooth ? "yes" : "no"),
            ]
    summary_box(io, "IndicatorHennemannGassner", setup)
  end
end

"""
    IndicatorLöhner (equivalent to IndicatorLoehner)

AMR indicator adapted from a FEM indicator by Löhner (1987), also used in the
FLASH code as standard AMR indicator.
The indicator estimates a weighted second derivative of a specified variable locally.
- Löhner (1987)
  "An adaptive finite element scheme for transient problems in CFD"
  [doi: 10.1016/0045-7825(87)90098-3](https://doi.org/10.1016/0045-7825(87)90098-3)
- http://flash.uchicago.edu/site/flashcode/user_support/flash4_ug_4p62/node59.html#SECTION05163100000000000000
"""
struct IndicatorLöhner{RealT<:Real, Variable, Cache} <: AbstractIndicator
  f_wave::RealT # TODO: Taal documentation
  variable::Variable
  cache::Cache
end

# this method is used when the indicator is constructed as for shock-capturing volume integrals
function IndicatorLöhner(equations::AbstractEquations, basis;
                         f_wave=0.2, variable)
  cache = create_cache(IndicatorLöhner, equations, basis)
  IndicatorLöhner{typeof(f_wave), typeof(variable), typeof(cache)}(f_wave, variable, cache)
end

# this method is used when the indicator is constructed as for AMR
function IndicatorLöhner(semi::AbstractSemidiscretization;
                         f_wave=0.2, variable)
  cache = create_cache(IndicatorLöhner, semi)
  IndicatorLöhner{typeof(f_wave), typeof(variable), typeof(cache)}(f_wave, variable, cache)
end

function Base.show(io::IO, indicator::IndicatorLöhner)
  @nospecialize indicator # reduce precompilation time

  print(io, "IndicatorLöhner(")
  print(io, "f_wave=", indicator.f_wave, ", variable=", indicator.variable, ")")
end

function Base.show(io::IO, ::MIME"text/plain", indicator::IndicatorLöhner)
  @nospecialize indicator # reduce precompilation time

  if get(io, :compact, false)
    show(io, indicator)
  else
    setup = [
             "indicator variable" => indicator.variable,
             "f_wave" => indicator.f_wave,
            ]
    summary_box(io, "IndicatorLöhner", setup)
  end
end

const IndicatorLoehner = IndicatorLöhner

# dirty Löhner estimate, direction by direction, assuming constant nodes
@inline function local_löhner_estimate(um::Real, u0::Real, up::Real, löhner::IndicatorLöhner)
  num = abs(up - 2 * u0 + um)
  den = abs(up - u0) + abs(u0-um) + löhner.f_wave * (abs(up) + 2 * abs(u0) + abs(um))
  return num / den
end

struct IndicatorMax{Variable, Cache<:NamedTuple} <: AbstractIndicator
  variable::Variable
  cache::Cache
end

# this method is used when the indicator is constructed as for shock-capturing volume integrals
function IndicatorMax(equations::AbstractEquations, basis;
                      variable)
  cache = create_cache(IndicatorMax, equations, basis)
  IndicatorMax{typeof(variable), typeof(cache)}(variable, cache)
end

# this method is used when the indicator is constructed as for AMR
function IndicatorMax(semi::AbstractSemidiscretization;
                      variable)
  cache = create_cache(IndicatorMax, semi)
  return IndicatorMax{typeof(variable), typeof(cache)}(variable, cache)
end

function Base.show(io::IO, indicator::IndicatorMax)
  @nospecialize indicator # reduce precompilation time

  print(io, "IndicatorMax(")
  print(io, "variable=", indicator.variable, ")")
end

function Base.show(io::IO, ::MIME"text/plain", indicator::IndicatorMax)
  @nospecialize indicator # reduce precompilation time

  if get(io, :compact, false)
    show(io, indicator)
  else
    setup = [
             "indicator variable" => indicator.variable,
            ]
    summary_box(io, "IndicatorMax", setup)
  end
end

"""
    IndicatorNeuralNetwork

Artificial neural network based indicator used for shock-capturing or AMR.
Depending on the indicator_type, different input values and corresponding trained networks are used.

`indicator_type = NeuralNetworkPerssonPeraire()`
- Input: The energies in lower modes as well as nnodes(dg).

`indicator_type = NeuralNetworkRayHesthaven()`
- 1d Input: Cell average of the cell and its neighboring cells as well as the interface values.
- 2d Input: Linear modal values of the cell and its neighboring cells.

- Ray, Hesthaven (2018)
  "An artificial neural network as a troubled-cell indicator"
  [doi:10.1016/j.jcp.2018.04.029](https://doi.org/10.1016/j.jcp.2018.04.029)
- Ray, Hesthaven (2019)
  "Detecting troubled-cells on two-dimensional unstructured grids using a neural network"
  [doi:10.1016/j.jcp.2019.07.043](https://doi.org/10.1016/j.jcp.2019.07.043)

`indicator_type = CNN (Only in 2d)`
- Based on convolutional neural network.
- 2d Input: Interpolation of the nodal values of the `indicator.variable` to the 4x4 LGL nodes.

If `alpha_continuous == true` the continuous network output for troubled cells (`alpha > 0.5`) is considered.
If the cells are good (`alpha < 0.5`), `alpha` is set to `0`.
If `alpha_continuous == false`, the blending factor is set to `alpha = 0` for good cells and
`alpha = 1` for troubled cells.

!!! warning "Experimental implementation"
    This is an experimental feature and may change in future releases.

"""
struct IndicatorNeuralNetwork{IndicatorType, RealT<:Real, Variable, Chain, Cache} <: AbstractIndicator
  indicator_type::IndicatorType
  alpha_max::RealT
  alpha_min::RealT
  alpha_smooth::Bool
  alpha_continuous::Bool
  alpha_amr::Bool
  variable::Variable
  network::Chain
  cache::Cache
end

# this method is used when the indicator is constructed as for shock-capturing volume integrals
function IndicatorNeuralNetwork(equations::AbstractEquations, basis;
                                indicator_type,
                                alpha_max=0.5,
                                alpha_min=0.001,
                                alpha_smooth=true,
                                alpha_continuous=true,
                                alpha_amr=false,
                                variable,
                                network)
  alpha_max, alpha_min = promote(alpha_max, alpha_min)
  IndicatorType = typeof(indicator_type)
  cache = create_cache(IndicatorNeuralNetwork{IndicatorType}, equations, basis)
  IndicatorNeuralNetwork{IndicatorType, typeof(alpha_max), typeof(variable), typeof(network), typeof(cache)}(
      indicator_type, alpha_max, alpha_min, alpha_smooth, alpha_continuous, alpha_amr, variable,
      network, cache)
end

# this method is used when the indicator is constructed as for AMR
function IndicatorNeuralNetwork(semi::AbstractSemidiscretization;
                                indicator_type,
                                alpha_max=0.5,
                                alpha_min=0.001,
                                alpha_smooth=true,
                                alpha_continuous=true,
                                alpha_amr=true,
                                variable,
                                network)
  alpha_max, alpha_min = promote(alpha_max, alpha_min)
  IndicatorType = typeof(indicator_type)
  cache = create_cache(IndicatorNeuralNetwork{IndicatorType}, semi)
  IndicatorNeuralNetwork{IndicatorType, typeof(alpha_max), typeof(variable), typeof(network), typeof(cache)}(
      indicator_type, alpha_max, alpha_min, alpha_smooth, alpha_continuous, alpha_amr, variable,
      network, cache)
end

function Base.show(io::IO, indicator::IndicatorNeuralNetwork)
  @nospecialize indicator # reduce precompilation time

  print(io, "IndicatorNeuralNetwork(")
  print(io, indicator.indicator_type)
  print(io, ", alpha_max=", indicator.alpha_max)
  print(io, ", alpha_min=", indicator.alpha_min)
  print(io, ", alpha_smooth=", indicator.alpha_smooth)
  print(io, ", alpha_continuous=", indicator.alpha_continuous)
  print(io, indicator.variable)
  print(io, ")")
end

function Base.show(io::IO, ::MIME"text/plain", indicator::IndicatorNeuralNetwork)
  @nospecialize indicator # reduce precompilation time

  if get(io, :compact, false)
    show(io, indicator)
  else
    setup = [
             "indicator type" => indicator.indicator_type,
             "max. α" => indicator.alpha_max,
             "min. α" => indicator.alpha_min,
             "smooth α" => (indicator.alpha_smooth ? "yes" : "no"),
             "continuous α" => (indicator.alpha_continuous ? "yes" : "no"),
             "indicator variable" => indicator.variable,
            ]
    summary_box(io, "IndicatorNeuralNetwork", setup)
  end
end

# Convert probability for troubled cell to indicator value for shockcapturing/AMR
@inline function probability_to_indicator(probability_troubled_cell, alpha_continuous, alpha_amr,
                                          alpha_min, alpha_max)
  # Initialize indicator to zero
  alpha_element = zero(probability_troubled_cell)

  if alpha_continuous && !alpha_amr
    # Set good cells to 0 and troubled cells to continuous value of the network prediction
    if probability_troubled_cell > 0.5
      alpha_element = probability_troubled_cell
    else
      alpha_element = zero(probability_troubled_cell)
    end

    # Take care of the case close to pure FV
    if alpha_element > 1 - alpha_min
      alpha_element = one(alpha_element)
    end

    # Scale the probability for a troubled cell (in [0,1]) to the maximum allowed alpha
    alpha_element *= alpha_max
  elseif !alpha_continuous && !alpha_amr
    # Set good cells to 0 and troubled cells to 1
    if probability_troubled_cell > 0.5
      alpha_element = alpha_max
    else
      alpha_element = zero(alpha_max)
    end
  elseif alpha_amr
    # The entire continuous output of the neural network is used for AMR
    alpha_element  = probability_troubled_cell

    # Scale the probability for a troubled cell (in [0,1]) to the maximum allowed alpha
    alpha_element *= alpha_max
  end

  return alpha_element
end

"""
    NeuralNetworkPerssonPeraire

Indicator type for creating an `IndicatorNeuralNetwork` indicator.

!!! warning "Experimental implementation"
    This is an experimental feature and may change in future releases.

See also: [`IndicatorNeuralNetwork`](@ref)
"""
struct NeuralNetworkPerssonPeraire end

"""
    NeuralNetworkRayHesthaven

Indicator type for creating an `IndicatorNeuralNetwork` indicator.

!!! warning "Experimental implementation"
    This is an experimental feature and may change in future releases.

See also: [`IndicatorNeuralNetwork`](@ref)
"""
struct NeuralNetworkRayHesthaven end

"""
    NeuralNetworkCNN

Indicator type for creating an `IndicatorNeuralNetwork` indicator.

!!! warning "Experimental implementation"
    This is an experimental feature and may change in future releases.

See also: [`IndicatorNeuralNetwork`](@ref)
"""
struct NeuralNetworkCNN end

end # @muladd


text:
@testset "Operators under mixed barrier conditions" begin
    regular_grids = [0:4, -1:6]
    irregular_grids = [[-1.0; 0.0; 2.0; 5.0; 9.0], [0.2; 0.3; 0.5; 0.9; 1.6; 2.5]]
        for x̄ in [regular_grids; irregular_grids],
            (L_bc_generator, L_generator) in [(L₁₋bc, L₁₋), (L₁₊bc, L₁₊), (L₂bc, L₂)],
            (ξ_lb, ξ_ub) in [(-2, 1), (-3, -2), (3, -4), (4, 3)],
            directions in [(:forward, :backward), (:backward, :forward),
                            (:backward, :backward), (:forward, :forward)]
            # Setup
            # RHS function
            f(x) = x^2

            # grid setup
            # x̄ = 0:4
            x = interiornodes(x̄)
            M = length(x)
            Δ_0p = x̄[2] - x̄[1]
            Δ_Mp = x̄[end] - x̄[end-1]

            # boundary condition
            bc = (Mixed(ξ = ξ_lb, direction = directions[1]),
                Mixed(ξ = ξ_ub, direction = directions[2]))

            # corresponding boundary condition matrix
            B_lb = (directions[1] != :backward) ? [-1+ξ_lb*Δ_0p; 1; zeros(M)] : [-1; 1+ξ_lb*Δ_0p; zeros(M)]
            B_ub = (directions[2] != :forward) ? [zeros(M); -1; 1+ξ_ub*Δ_Mp] : [zeros(M); -1+ξ_ub*Δ_Mp; 1]

            B = transpose([B_lb B_ub])
            b = [0; 0]

            # operator on the interior, with bc applied
            L_bc = L_bc_generator(x̄, bc)
            # operator on the exterior
            L = L_generator(x̄)

            # solution for interior v from L_bc
            v_bc = L_bc \ f.(x)
            # solution for extended v from L by stacking up boundary condition matrices
            v̄ = [L; B] \ [f.(x); b]

            # test if interior solutions are identical
            @test v_bc ≈ v̄[2:end-1]
            # test if extrapolated solutions are identical
            @test extrapolatetoboundary(v_bc, x̄, bc) ≈ v̄
        end
end

@testset "Accuracy test for absorbing boundary conditions" begin
    using LinearAlgebra
    regular_grids = [1:20, -5:-1]
    irregular_grids = [[-1.0; 0.0; 2.0; 5.0; 9.0], [0.2; 0.3; 0.5; 0.9; 1.6; 2.5]]
    for x̄ in [regular_grids; irregular_grids],
        (L_bc_generator, L_generator, L_affine_generator) in [(L₁₋bc, L₁₋, L₁₋affine), (L₁₊bc, L₁₊, L₁₊affine), (L₂bc, L₂, L₂affine)],
        loc in 0:3
        # Setup
        # RHS function
        f(x) = x^2

        # grid setup
        # x̄ = 0:4
        x = interiornodes(x̄)
        M = length(x)
        Δ_0p = x̄[2] - x̄[1]
        Δ_Mp = x̄[end] - x̄[end-1]

        # boundary condition
        bc = (Absorbing(loc), Absorbing())

        # corresponding boundary condition matrix
        B = transpose([[1; zeros(M+1)] [zeros(M+1); 1]])
        b = [0; 0]

        # operator on the interior, with bc applied
        L_bc = L_bc_generator(x̄, bc)
        # operator on the exterior
        L = L_generator(x̄)
        # apply interior boundary conditions (if needed)
        RHS = f.(x)
        if (bc[1].loc > 0)
            L[1:bc[1].loc,:] = zeros(bc[1].loc,M+2)
            L[1:bc[1].loc,2:(1+bc[1].loc)] = Diagonal(ones(bc[1].loc))
            RHS[1:bc[1].loc] .= zero(eltype(L))
        end

        # solution for interior v from L_bc
        v_bc = L_bc \ (f.(x) + L_affine_generator(x̄, f.(x), bc))
        # solution for extended v from L by stacking up boundary condition matrices
        v̄ = [L; B] \ [RHS; b]

        # test if extrapolated solutions are identical
        @test extrapolatetoboundary(v_bc, x̄, bc) ≈ v̄

    end
end

@testset "Accuracy test for nonhomogeneous absorbing boundary conditions" begin
    regular_grids = [0:4, -1:6]
    irregular_grids = [[-1.0; 0.0; 2.0; 5.0; 9.0], [0.2; 0.3; 0.5; 0.9; 1.6; 2.5]]
    for x̄ in [regular_grids; irregular_grids],
        (L_bc_generator, L_generator, L_affine_generator) in [(L₁₋bc, L₁₋, L₁₋affine), (L₁₊bc, L₁₊, L₁₊affine), (L₂bc, L₂, L₂affine)],
        (S_lower, S_upper) in [(0.0, 1.0), (1.0, 0.0), (1.5, -1.5), (1.2, 1.3), (1.0, 1.0)]
        # Setup
        # RHS function
        f(x) = x^2

        # grid setup
        # x̄ = 0:4
        x = interiornodes(x̄)
        M = length(x)
        Δ_0p = x̄[2] - x̄[1]
        Δ_Mp = x̄[end] - x̄[end-1]

        # boundary condition
        bc = (NonhomogeneousAbsorbing(S_lower), NonhomogeneousAbsorbing(S_upper))

        # corresponding boundary condition matrix
        B = transpose([[1; zeros(M+1)] [zeros(M+1); 1]])
        b = [S_lower; S_upper]

        # operator on the interior, with bc applied
        L_bc = L_bc_generator(x̄, bc)
        # operator on the exterior
        L = L_generator(x̄)

        # solution for interior v from L_bc
        v_bc = L_bc \ (f.(x) + L_affine_generator(x̄, f.(x), bc))
        # solution for extended v from L by stacking up boundary condition matrices
        v̄ = [L; B] \ [f.(x); b]

        # test if extrapolated solutions are identical
        @test extrapolatetoboundary(v_bc, x̄, bc) ≈ v̄

    end
end


text:
# pre-defined factors

@deffactor EntropyFactor [x] H(x)

@deffactor LogNormalFactor [x, μ, τ] begin
    -(1/2) * ((E(τ) * ( V(x) + V(μ) + (E(x) - E(μ))^2 ) + log(2π) - Elog(τ)))
end

@deffactor LogGammaCanonFactor [x, α, β] begin
    (E(α) - 1) * Elog(x) - E(x) * E(β) + E(α) * Elog(β) - Eloggamma(α)
end

@deffactor LogGammaFactor [x, α, θ] begin
    (E(α) - 1) * Elog(x) - E(x)/E(θ) - E(α) * Elog(θ) - Eloggamma(α)
end

@deffactor LogMvNormalCanonFactor [x, μ, Λ] begin
    δ = E(x) - E(μ)
    EΛ = E(Λ)
    -(1/2) * (trace(EΛ * (C(x) .+ C(μ) .+ δ * δ')) + length(x) * log(2π) - Elogdet(Λ))
end

@deffactor LogMvNormalDiagCanonFactor [x, μ, τ] begin
    δ = E(x) - E(μ)
    Eτ = E(τ)
    -(1/2) * (sum(Eτ .* (V(x) .+ V(μ) .+ δ.^2 )) + length(x) * log(2π) - Elogdet(τ))
end

@deffactor LogDirichletFactor [x, α] begin
    dot(E(α) - 1, Elog(x)) - ElogB(α)
end

@deffactor LogWishartFactor [X ν V] begin
    EX = E(X)
    p = size(EX, 1)
    n = E(ν)
    0.5 * (n - p - 1) * Elogdet(X) - 0.5 * n * Elogdet(V) - 0.5 * trace(Einv(V) * EX) - 0.5 * n * p * log(2) - Elogmvgamma(p, ν/2)
end

################### Dealing with HMMs in factors #####################
@deffactor LogMarkovChainFactor [z, π0, A] begin
    dot(E(z)[:, 1], Elog(π0)) + sum(C(z) .* Elog(A))
end

@deffactor LogMarkovMatrixFactor [x, A] begin
    sum((E(A) - 1) .* Elog(x) .- ElogB(A))
end


text:
using Pkg; Pkg.add("Documenter")
using Documenter
using SeisProcessing

makedocs(
    sitename = "SeisProcessing.jl",
    modules = [SeisProcessing],
    format = Documenter.HTML(prettyurls = get(ENV, "CI", nothing)== "true"),
    pages = [
          "Home" => "index.md",
          "Library" => Any[
                       "Public" => "lib/public.md",
			],
		],
)

# Documenter can also automatically deploy documentation to gh-pages.
# See "Hosting Documentation" and deploydocs() in the Documenter manual
# for more information.
deploydocs(
    repo = "github.com/SeismicJulia/SeisProcessing.jl.git",
    target = "build",
    deps = nothing,
    make = nothing
)


text:
using MDP
using Printf

# Only needed for rle function
using StatsBase

const L = 50
const M = 30

#alpha = [0.75 0.75]
alpha = [0.75 0.75]

const P = [ alpha[1]		1-alpha[1]
	         1 - alpha[2]  alpha[2]  ]

p1 = p2 = 0.3

const Q1 = [ 1 - p1   p1
              0        1 ]

const Q2 = [ 1 - p2   p2
              0        1 ]

const xi = [ P^n for n = 1:M+1]

function probabilityVector(Q, n)
    if n < L+2
        [1 0]*Q^n
    else # L+2 component corresponds to z_inf
        [0;1]
    end
end

z1 = [ probabilityVector(Q1, n)  for n = 1:L+2]
z2 = [ probabilityVector(Q2, n)  for n = 1:L+2]

r = 1.0
c = 0.4

function bellmanUpdate!(v_next, g_next, v; discount=1.0)
    # Assume that v is L+2 x L+2 x 2 x M+1

    # v_next(k, ell, s, m)
    # k   = last time state of user 1 was observed
    # ell = last time state of user 2 was observed
    #  s  = last observed state of channel {Idle, Busy}
    #  m  = last time the state of channel was observed
    for m = 1:M+1, s = 1:2, ell = 1:L+2, k = 1:L+2
        next_m = min(m+1, M+1)
        next_ell = (ell < L+1) ? ell+1 : ell
        next_k = (k < L+1) ? k+1 : k

        q00 = discount * v[next_k, next_ell, s, next_m]

        q10 = ( z1[k][2] * xi[m][s,1] * r - z1[k][2] * c
              + discount * ( z1[k][1] * v[1, next_ell, s, next_m]
                           + z1[k][2] * xi[m][s,1] * v[1, next_ell, 1, 1]
                           + z1[k][2] * xi[m][s,2] * v[L+2, next_ell, 2, 1]
                           )
              )

        q01 = ( z2[ell][2] * xi[m][s,1] * r - z2[ell][2] * c
              + discount * ( z2[ell][1] * v[next_k, 1, s, next_m]
                           + z2[ell][2] * xi[m][s,1] * v[next_k, 1, 1, 1]
                           + z2[ell][2] * xi[m][s,2] * v[next_k, L+2, 2, 1]
                           )
              )

        q11 = ( ( z1[k][2] * z2[ell][1] + z1[k][1] * z2[ell][2] ) * xi[m][s,1] * r
              - ( z1[k][2] + z2[ell][2] ) * c
              + discount * ( z1[k][1] * z2[ell][1] * v[1, 1, s, next_m]
                           + (z1[k][2] * z2[ell][1] + z1[k][1] * z2[ell][2]) * xi[m][s,1] * v[1, 1, 1, 1]
                           + z1[k][2] * z2[ell][2] * xi[m][s,1] * v[L+2, L+2, 1, 1]
                           + z1[k][2] * z2[ell][1] * xi[m][s,2] * v[L+2, 1, 2, 1]
                           + z1[k][1] * z2[ell][2] * xi[m][s,2] * v[1, L+2, 2, 1]
                           + z1[k][2] * z2[ell][2] * xi[m][s,2] * v[L+2, L+2, 2, 1]
                           )
              )

        next = [q00, q10, q01, q11]
        v_next[k, ell, s, m],  g_next[k, ell, s, m] = findmax(next)
    end

    @printf("Debug: v[1,1,1,1] = %0.8f\n", v_next[1,1,1,1])
    return (v_next, g_next)
end

model = DynamicModel(bellmanUpdate!; objective=:Max)
v_initial = zeros(L+2, L+2, 2, M+1)

@time (v,g) = valueIteration(model, v_initial; discount=0.9)
@printf("Value = %0.4f\n", v[1,1,1,1])

for s = 1:2
    for m = 1:5
        bottom  = vec(g[:,1,s,m])
        top     = vec(g[:,L+2,s,m])

        left    = vec(g[1,:,s,m])
        right   = vec(g[L+2,:,s,m])

        @printf(" ========= s = %d m = %d =============\n", s, m)
        println(rle(bottom))
        println(rle(top))
        println(rle(left))
        println(rle(right))
    end
end


text:
using Documenter, NASAEarthObservations

makedocs(;
    modules=[NASAEarthObservations],
    format=Documenter.HTML(),
    pages=[
        "Home" => "index.md",
    ],
    repo="https://github.com/asinghvi17/NASAEarthObservations.jl/blob/{commit}{path}#L{line}",
    sitename="NASAEarthObservations.jl",
    authors="Anshul Singhvi",
    assets=String[],
)


text:
using JuAFEM
using Tensors

# Stiffness using normal matrices
function ke_element_mat!(Ke, X::Vector{Vec{dim, T}}, fe_values::CellScalarValues{dim}, Ee, B, DB, BDB) where {T, dim}
    n_basefuncs = getnbasefunctions(fe_values)
    @assert length(X) == n_basefuncs

    reinit!(fe_values, X)
    for q_point in 1:getnquadpoints(fe_values)
        for i in 1:n_basefuncs
            dNdx = shape_gradient(fe_values, q_point, i)[1]
            dNdy = shape_gradient(fe_values, q_point, i)[2]
            dNdz = shape_gradient(fe_values, q_point, i)[3]

            B[1, i * 3-2] = dNdx
            B[2, i * 3-1] = dNdy
            B[3, i * 3-0] = dNdz
            B[4, 3 * i-1] = dNdz
            B[4, 3 * i-0] = dNdy
            B[5, 3 * i-2] = dNdz
            B[5, 3 * i-0] = dNdx
            B[6, 3 * i-2] = dNdy
            B[6, 3 * i-1] = dNdx
        end

        A_mul_B!(DB, Ee, B)
        At_mul_B!(BDB, B, DB)
        scale!(BDB, getdetJdV(fe_values, q_point))
        for p in 1:size(Ke,1)
            for q in 1:size(Ke,2)
                Ke[p, q] += BDB[p, q]
            end
        end
    end

    return Ke
end;

# Stiffness using scalar values
function ke_element!(Ke, X::Vector{Vec{dim, T}}, fe_values::CellScalarValues{dim}, C) where {T,dim}
    n_basefuncs = getnbasefunctions(fe_values)
    @assert length(X) == n_basefuncs
    reinit!(fe_values, X)
    @inbounds for q_point in 1:getnquadpoints(fe_values)
        for a in 1:n_basefuncs
            for b in 1:n_basefuncs
                ∇ϕa = shape_gradient(fe_values, q_point, a)
                ∇ϕb = shape_gradient(fe_values, q_point, b)
                Ke_e = dotdot(∇ϕa, C, ∇ϕb) * getdetJdV(fe_values, q_point)
                for d1 in 1:dim, d2 in 1:dim
                    Ke[dim*(a-1) + d1, dim*(b-1) + d2] += Ke_e[d1,d2]
                end
            end
        end
    end
    return Ke
end;

# Stiffness using vector values
function ke_element2!(Ke, X::Vector{Vec{dim, T}}, fe_values::CellVectorValues{dim}, C) where {T,dim}
    n_basefuncs = getnbasefunctions(fe_values)
    @assert length(X) * dim == n_basefuncs
    reinit!(fe_values, X)
    ɛ = [zero(SymmetricTensor{2, dim, T}) for i in 1:n_basefuncs]
    @inbounds for q_point in 1:getnquadpoints(fe_values)
        for i in 1:n_basefuncs
            ɛ[i] = symmetric(shape_gradient(fe_values, q_point, i))
        end
        dΩ = getdetJdV(fe_values, q_point)
        for i in 1:n_basefuncs
            ɛC = ɛ[i] ⊡ C
            for j in 1:n_basefuncs
                Ke[i, j] += (ɛC ⊡ ɛ[j]) * dΩ
            end
        end
    end
    return Ke
end;

E = 200e9
ν = 0.3
λ = E*ν / ((1 + ν) * (1 - 2ν))
μ = E / (2(1 + ν))
δ(i,j) = i == j ? 1.0 : 0.0
g(i,j,k,l) = λ*δ(i,j)*δ(k,l) + μ*(δ(i,k)*δ(j,l) + δ(i,l)*δ(j,k))

C = SymmetricTensor{4, 3}(g)

M = λ/ν * (1 - ν)

Cmat = [ M      λ      λ    0.0    0.0   0.0;
         λ      M      λ    0.0    0.0   0.0;
         λ      λ      M    0.0    0.0   0.0;
        0.0    0.0    0.0    μ     0.0   0.0;
        0.0    0.0    0.0   0.0     μ    0.0;
        0.0    0.0    0.0   0.0    0.0    μ]

interpolation = Lagrange{3, RefCube, 1}()
quad_rule = QuadratureRule{3, RefCube}(1)
values = CellScalarValues(quad_rule, interpolation);
vector_values = CellVectorValues(quad_rule, interpolation);

# Generate some coordinates
x = [-1.0 -1.0 -1.0;
      1.0 -1.0 -1.0;
      1.0  1.0 -1.0;
     -1.0  1.0 -1.0;
     -1.0 -1.0  1.0;
      1.0 -1.0  1.0;
      1.0  1.0  1.0;
     -1.0  1.0  1.0;]
x = x .+ 0.05 * rand()
x_vec = reinterpret(Vec{3, Float64}, x, (8,));

n_basefunctions = getnbasefunctions(vector_values)
Ke = zeros(n_basefunctions, n_basefunctions)
Ke2 = copy(Ke)
Ke3 = copy(Ke)

B   =  zeros(6, n_basefunctions)
DB  =  zeros(6, n_basefunctions)
BDB =  zeros(n_basefunctions, n_basefunctions);

fill!(Ke, 0)
fill!(Ke2, 0)
fill!(Ke3, 0)
ke_element!(Ke2, x_vec, values, C)
ke_element2!(Ke3, x_vec, vector_values, C);
ke_element_mat!(Ke, x_vec, values, Cmat, B, DB, BDB);

using Test
@test norm(Ke - Ke2) / norm(Ke) < 1e-14
@test norm(Ke - Ke3) / norm(Ke) < 1e-14

println("Stiffness successful")

Ke


text:
#! /usr/bin/env julia

println("Start test 'tests/patched.jl'")

include("../src/DCellC.jl")
using DCellC

image, label = synthesize(512, 512, (500, 1000))
model = Multiscale3(GreyscaleImage, bn = true, seed = 42)

dens = density(model, image)

println("Test patched.jl: Calculating density (unpatched) successful")

dens2 = density_patched(model, image, patchsize=200)

println("Test patched.jl: Calculating density (patched) successful")

error = mean(abs.(dens - dens2))

println("Test patched.jl: Mean difference between patched and unpatched density: $error")

println("Test patched.jl completed")

max_stars_repo_name: tscode/DCellC

max_issues_repo_name: tscode/DCellC


text:
parent(img::AbstractImage) = parent(data(img))
first_index(A::Array) = 1
first_index(A::SubArray) = A.first_index

function iterate_spatial(img::AbstractArray)
    sz = size(img)
    s = strides(img)
    colorsz = 1
    colorstride = 0
    timesz = 1
    timestride = 0
    cd = colordim(img)
    if cd != 0
        colorsz = sz[cd]
        colorstride = s[cd]
    end
    td = timedim(img)
    if td != 0
        timesz = sz[td]
        timestride = s[td]
    end
    cspatial = setdiff(1:ndims(img), [cd, td])
    first_index(data(img)), sz[cspatial], s[cspatial], colorsz, colorstride, timesz, timestride
end


text:
# medotdo da bisseccao
function met_bis(a,b,f,tol=1.0e-2)
    if a>=b
        error("entre com a < b")
    elseif !(f(a)*f(b) < 0)
        error("entre com a,b tal que f(a)*f(b) < 0")
    end
    itermax = 100
    iter= 0
    xk= (a+b)/2
    while !(abs(f(xk))< tol) && iter < itermax
        if f(xk)*f(a) < 0
            b = xk
        elseif f(xk)*f(b) < 0
            a = xk
        end
        xk=(a+b)/2
        iter += 1 #iter = iter +1
    end
    println("$iter  é o numero de interações")
    if abs(f(xk)) < tol
        println("$xk é zero de f")
    end
    return xk
end

# Método de Newton
function newton(xk , f , df , tol = 1.0e-5)
    itermax = 100;
    iter= 0;
    while !(abs(f(xk))< tol) && iter < itermax
        xk = xk - (f(xk)/df(xk))
        iter += 1
        println(iter ,":", f(xk))
    end
    println("$iter  é o numero de interações")
    if abs(f(xk)) < tol
        println("$xk é zero de f")
    end
    return xk
end

#Método de Newton Secante
function newton_secante(xk, xk1, f, tol = 1.0e-5)
    itermax = 100;
    iter= 0;
    while !(abs(f(xk))< tol) && iter < itermax
        mk = (f(xk)-f(xk1))/(xk-xk1)
        xk1=xk
        xk = xk - (f(xk)/mk)
        iter += 1 #iter = iter +1
        println(iter ,":", f(xk))
    end
    println("$iter  é o numero de interações")
    if abs(f(xk)) < tol
        println("$xk é zero de f")
    end
    return xk
end

max_stars_repo_path: 002 - Raizes_de_Polinomios.jl

max_stars_repo_name: LudwigFonseca/Julia-Language-Examples

max_issues_repo_path: 002 - Raizes_de_Polinomios.jl

max_issues_repo_name: LudwigFonseca/Julia-Language-Examples

max_forks_repo_path: 002 - Raizes_de_Polinomios.jl

max_forks_repo_name: LudwigFonseca/Julia-Language-Examples