text:
"""
    pmcwf(tspan, psi0, H, J; Ntrajectories, <keyword arguments>)

Integrate parallely `Ntrajectories` MCWF trajectories.
Arguments are passed internally to `QuantumOptics.timeevolution.mcwf`.

# Arguments
* `tspan`: Vector specifying the points of time for which output should
be displayed.
* `psi0`: Initial state vector.
* `H`: Arbitrary Operator specifying the Hamiltonian.
* `J`: Vector containing all jump operators which can be of any arbitrary
operator type.
* `Ntrajectories=1`: Number of MCWF trajectories.
* `parallel_type=:none`: The type of parallelism to employ. The types of
parallelism included are: `:none`, `:threads`, `:split_threads`, `:parfor` and
`:pmap`.
* `progressbar=true`: If `true`, a progression bar is displayed.
* `return_data=true`: If `true`, the solution is returned as a `Tuple`.
* `save_data=true`: If `true`, the solution is saved to disk.
If `return_data=false`, less RAM is used (except for `parallel_type=:threads`).
* `fpath=missing`: savefile path (e.g. `some/valid/path/filename.jld2`).
Directory must pre-exist, the savefile is created.
* `additional_data=missing`: If given a `Dict`, entries are added to the
savefile.
* `seed=rand(UInt)`: seed for each trajectory's random number generator.
* `rates=ones()`: Vector of decay rates.
* `fout`: If given, this function `fout(t, psi)` is called every time an
output should be displayed. ATTENTION: The state `psi` is neither
normalized nor permanent! It is still in use by the ode solve
and therefore must not be changed.
* `Jdagger=dagger.(J)`: Vector containing the hermitian conjugates of the jump
operators. If they are not given they are calculated automatically.
* `display_beforeevent=false`: `fout` is called before every jump.
* `display_afterevent=false`: `fout` is called after every jump.
* `kwargs...`: Further arguments are passed on to the ode solver.

See also: [`timeevolution.mcwf`](@ref)

# Examples
```julia-repl
julia> tspan = collect(t0:dt:t_max);
julia> fb = FockBasis(10); ψ0 = fockstate(fb,0); a = destroy(fb);
julia> H = randoperator(fb); H = H + dagger(H); γ = 1.;
julia> # 100 MCWF trajectories
julia> t, trajs = pmcwf(tspan, ψ0, H, [sqrt(γ)*a];
                        Ntrajectories=100, parallel_type=:pmap);
```
"""
function pmcwf(tspan, psi0::T, H::AbstractOperator{B,B}, J::Vector;
        Ntrajectories=1, parallel_type::Symbol = :none,
        progressbar::Bool = true,
        return_data::Bool = true, save_data::Bool = false,
        fpath::Union{String,Missing}=missing,
        additional_data::Union{Dict{String,T2},Missing}=missing,
        seed=nothing, rates::DecayRates=nothing,
        fout=nothing, Jdagger::Vector=dagger.(J),
        display_beforeevent=false, display_afterevent=false,
        alg=OrdinaryDiffEq.AutoTsit5(OrdinaryDiffEq.Rosenbrock23()),
        kwargs...) where {B<:Basis,T<:Ket{B},T2}
    @assert return_data || save_data "pmcwf outputs nothing"
    save_data && @assert !ismissing(fpath) "ERROR: savefile path is missing"
    save_data && @assert isdir(splitdir(fpath)[1]) "ERROR: accessing "*splitdir(fpath)[1]*": No such directory"
    save_data && @assert !isfile(fpath) "ERROR: "*fpath*" already a file: Choose a free savefile name"

    if parallel_type == :none || Ntrajectories == 1
        return serial_mcwf(tspan,psi0,H,J;Ntrajectories=Ntrajectories,
            progressbar=progressbar,
            return_data=return_data,save_data=save_data,
            fpath=fpath,additional_data=additional_data,
            seed=seed,rates=rates,fout=fout,Jdagger=Jdagger,
            display_beforeevent=display_beforeevent,
            display_afterevent=display_afterevent,
            alg=alg,
            kwargs...);
    elseif parallel_type == :threads
        # TO DO: save_data-only not totally supported. Could use less RAM by
        # writting trajectories directly to disk but would probably require
        # some locking or some parallel process.
        return multithreaded_mcwf(tspan,psi0,H,J;Ntrajectories=Ntrajectories,
            progressbar=progressbar,
            return_data=return_data,save_data=save_data,
            fpath=fpath,additional_data=additional_data,
            seed=seed,rates=rates,fout=fout,Jdagger=Jdagger,
            display_beforeevent=display_beforeevent,
            display_afterevent=display_afterevent,
            alg=alg,
            kwargs...);
    elseif parallel_type == :pmap
        # TO DO: add batch_size as an option
        return pmap_mcwf(tspan,psi0,H,J;Ntrajectories=Ntrajectories,
            progressbar=progressbar,
            return_data=return_data,save_data=save_data,
            fpath=fpath,additional_data=additional_data,
            seed=seed,rates=rates,fout=fout,Jdagger=Jdagger,
            display_beforeevent=display_beforeevent,
            display_afterevent=display_afterevent,
            alg=alg,
            kwargs...);
    elseif parallel_type == :parfor
        return distributed_mcwf(tspan,psi0,H,J;Ntrajectories=Ntrajectories,
            progressbar=progressbar,
            return_data=return_data,save_data=save_data,
            fpath=fpath,additional_data=additional_data,
            seed=seed,rates=rates,fout=fout,Jdagger=Jdagger,
            display_beforeevent=display_beforeevent,
            display_afterevent=display_afterevent,
            alg=alg,
            kwargs...);
    elseif parallel_type == :split_threads
        # TO DO: add batch_size as an option
        return split_threads_mcwf(tspan,psi0,H,J;Ntrajectories=Ntrajectories,
            progressbar=progressbar,
            return_data=return_data,save_data=save_data,
            fpath=fpath,additional_data=additional_data,
            seed=seed,rates=rates,fout=fout,Jdagger=Jdagger,
            display_beforeevent=display_beforeevent,
            display_afterevent=display_afterevent,
            alg=alg,
            kwargs...);
    else
        valptypes = [:none, :threads, :pmap, :parfor, :split_threads];
        error("Invalid parallel type. Type :$parallel_type not available.\n"*
              "Available types are: "*reduce(*,[":$t " for t in valptypes]))
    end
end

function serial_mcwf(tspan, psi0::T, H::AbstractOperator{B,B}, J::Vector;
        Ntrajectories=1, progressbar::Bool = true,
        return_data::Bool = true, save_data::Bool = false,
        fpath::Union{String,Missing}=missing,
        additional_data::Union{Dict{String,T2},Missing}=missing,
        seed=nothing, rates::DecayRates=nothing,
        fout=nothing, Jdagger::Vector=dagger.(J),
        display_beforeevent=false, display_afterevent=false,
        alg=OrdinaryDiffEq.AutoTsit5(OrdinaryDiffEq.Rosenbrock23()),
        kwargs...) where {B<:Basis,T<:Ket{B},T2}
    if save_data
        file = jldopen(fpath, "a+");
        file["t"] = tspan;
        if !ismissing(additional_data)
            for (key, val) in additional_data
                file[key] = val;
            end
        end
    end
    if return_data
        # Pre-allocate an array for holding each MC simulation
        out_type = fout == nothing ? typeof(psi0) : pure_inference(fout, Tuple{eltype(tspan),typeof(psi0)});
        sols::Array{Vector{out_type},1} = fill(Vector{out_type}(),Ntrajectories);
    end
    if progressbar
        # A progress bar is set up to be updated by the master thread
        progress = Progress(Ntrajectories);
        ProgressMeter.update!(progress, 0);
    end
    for i in 1:Ntrajectories
        if isnothing(seed)
            sol = timeevolution.mcwf(tspan,psi0,H,J;
                rates=rates,fout=fout,Jdagger=Jdagger,
                display_beforeevent=display_beforeevent,
                display_afterevent=display_afterevent,
                alg=alg,kwargs...);
        else
            sol = timeevolution.mcwf(tspan,psi0,H,J;
                seed=seed,rates=rates,fout=fout,Jdagger=Jdagger,
                display_beforeevent=display_beforeevent,
                display_afterevent=display_afterevent,
                alg=alg,kwargs...);
        end
        save_data ? file["trajs/"*string(i)] = sol[2] : nothing;
        return_data ? sols[i] = sol[2] : nothing;
        progressbar ? ProgressMeter.next!(progress) : nothing;
    end
    save_data && close(file);
    return return_data ? (tspan, sols) : nothing;
end

function multithreaded_mcwf(tspan, psi0::T, H::AbstractOperator{B,B}, J::Vector;
        Ntrajectories=1, progressbar::Bool = true,
        return_data::Bool = true, save_data::Bool = true,
        fpath::Union{String,Missing}=missing,
        additional_data::Union{Dict{String,T2},Missing}=missing,
        seed=nothing, rates::DecayRates=nothing,
        fout=nothing, Jdagger::Vector=dagger.(J),
        display_beforeevent=false, display_afterevent=false,
        alg=OrdinaryDiffEq.AutoTsit5(OrdinaryDiffEq.Rosenbrock23()),
        kwargs...) where {B<:Basis,T<:Ket{B},T2}

    if progressbar
        # A progress bar is set up to be updated by the master thread
        progress = Progress(Ntrajectories);
        ProgressMeter.update!(progress, 0);
        function update_progressbar(n::Threads.Atomic{Int64})
            if Threads.threadid() == 1  # If first thread: update progress bar
                for i in 1:(n[]+1) ProgressMeter.next!(progress); end
                Threads.atomic_xchg!(n,0);
            else                    # Else: increment the number of pending updates.
                Threads.atomic_add!(n,1);
            end
        end
        nupdates = Threads.Atomic{Int}(0);
    end
    if save_data
        file = jldopen(fpath, "a+");
        file["t"] = tspan;
        if !ismissing(additional_data)
            for (key, val) in additional_data
                file[key] = val;
            end
        end
        return_data = true; # Some duct tape
    end
    if return_data
        # Pre-allocate an array for holding each MC simulation
        out_type = fout == nothing ? typeof(psi0) : pure_inference(fout, Tuple{eltype(tspan),typeof(psi0)});
        sols::Array{Vector{out_type},1} = [Vector{out_type}() for i in 1:Ntrajectories];
    end
    seed = isnothing(seed) ? [rand(UInt) for i in 1:Threads.nthreads()] : [seed for i in 1:Threads.nthreads()];
    # Multi-threaded for-loop over all MC trajectories.
    Threads.@threads for i in 1:Ntrajectories
        sol = timeevolution.mcwf(tspan,psi0,H,J;
                seed=seed[Threads.threadid()],rates=rates,fout=fout,Jdagger=Jdagger,
                display_beforeevent=display_beforeevent,
                display_afterevent=display_afterevent,
                alg=alg, kwargs...);
        #save_data ? file["trajs/"*string(i)] = sol[2] : nothing;
        return_data ? sols[i] = sol[2] : nothing;
        # Updates progress bar if called from the main thread or adds a pending update otherwise
        progressbar ? update_progressbar(nupdates) : nothing;
    end
    progressbar ? update_progressbar(nupdates) : nothing;
    # Sets the progress bar to 100%
    if progressbar && (progress.counter < Ntrajectories) ProgressMeter.update!(progress, Ntrajectories); end;

    # Some additional duct tape
    if save_data
        for i in 1:length(sols)
            file["trajs/"*string(i)] = sols[i];
        end
    end

    save_data && close(file);
    return return_data ? (tspan, sols) : nothing;
end;

function pmap_mcwf(tspan, psi0::T, H::AbstractOperator{B,B}, J::Vector;
        Ntrajectories=1, progressbar::Bool = true,
        return_data::Bool = true, save_data::Bool = true,
        fpath::Union{String,Missing}=missing,
        additional_data::Union{Dict{String,T2},Missing}=missing,
        seed=nothing, rates::DecayRates=nothing,
        fout=nothing, Jdagger::Vector=dagger.(J),
        display_beforeevent=false, display_afterevent=false,
        alg=OrdinaryDiffEq.AutoTsit5(OrdinaryDiffEq.Rosenbrock23()),
        kwargs...) where {B<:Basis,T<:Ket{B},T2}

    # Create a remote channel from where trajectories are read out by the saver
    remch = RemoteChannel(()->Channel{Any}(Inf)); # TO DO: add some finite buffer size

    wp = CachingPool(workers());

    # Create a task fetched by the first available worker that retrieves trajs
    # from the remote channel and writes them to disk. A progress bar is set up
    # as well.
    saver = @async launch_saver(remch; Ntrajectories=Ntrajectories,
        progressbar=progressbar, return_data=return_data, save_data=save_data,
        fpath=fpath, additional_data=additional_data);
        batches = nfolds(1:Ntrajectories,length(wp.workers))
    batches = nfolds(1:Ntrajectories,length(wp.workers))
    # Multi-processed map over all batches of MC trajectories. Maps jobs to
    # workers() from the local process. Jobs consist in computing a batch of
    # MC trajectories and pipe them to the remote channel remch.
    pmap(wp,1:length(wp.workers)) do i
        for j in 1:length(batches[i])
            put!(remch,begin
                if isnothing(seed)
                    timeevolution.mcwf(tspan,psi0,H,J;
                        rates=rates,fout=fout,Jdagger=Jdagger,
                        display_beforeevent=display_beforeevent,
                        display_afterevent=display_afterevent,
                        alg=alg, kwargs...);
                else
                    timeevolution.mcwf(tspan,psi0,H,J;
                        seed=seed,rates=rates,fout=fout,Jdagger=Jdagger,
                        display_beforeevent=display_beforeevent,
                        display_afterevent=display_afterevent,
                        alg=alg, kwargs...);
                end
            end);
        end
    end
    # Once saver has consumed all queued trajectories produced by all workers, an
    # array of MCWF trajs is returned.
    sols = fetch(saver);
    # Clear caching pool
    clear!(wp);

    if return_data
        out_type = fout == nothing ? typeof(psi0) : pure_inference(fout, Tuple{eltype(tspan),typeof(psi0)});
        return (tspan, convert(Array{Vector{out_type},1},sols));
    else
        nothing;
    end
end;

function distributed_mcwf(tspan, psi0::T, H::AbstractOperator{B,B}, J::Vector;
        Ntrajectories=1, progressbar::Bool = true,
        return_data::Bool = true, save_data::Bool = true,
        fpath::Union{String,Missing}=missing,
        additional_data::Union{Dict{String,T2},Missing}=missing,
        seed=nothing, rates::DecayRates=nothing,
        fout=nothing, Jdagger::Vector=dagger.(J),
        display_beforeevent=false, display_afterevent=false,
        alg=OrdinaryDiffEq.AutoTsit5(OrdinaryDiffEq.Rosenbrock23()),
        kwargs...) where {B<:Basis,T<:Ket{B},T2}

    # Create a remote channel from where trajectories are read out by the saver
    remch = RemoteChannel(()->Channel{Any}(Inf)); # TO DO: add some finite buffer size

    # Create a task fetched by the first available worker that retrieves trajs
    # from the remote channel and writes them to disk. A progress bar is set up
    # as well.
    saver = @async launch_saver(remch; Ntrajectories=Ntrajectories,
        progressbar=progressbar, return_data=return_data, save_data=save_data,
        fpath=fpath, additional_data=additional_data);
    # Multi-processed for-loop over all MC trajectories. @distributed feeds
    # workers() asynchronously with jobs from the local process and is fetched.
    # Jobs consist in computing a trajectory and pipe it to the remote channel remch.
    @sync @distributed for i in 1:Ntrajectories
        put!(remch,begin
                        if isnothing(seed)
                            timeevolution.mcwf(tspan,psi0,H,J;
                            rates=rates,fout=fout,Jdagger=Jdagger,
                            display_beforeevent=display_beforeevent,
                            display_afterevent=display_afterevent,
                            alg=alg, kwargs...);
                        else
                            timeevolution.mcwf(tspan,psi0,H,J;
                            seed=seed,rates=rates,fout=fout,Jdagger=Jdagger,
                            display_beforeevent=display_beforeevent,
                            display_afterevent=display_afterevent,
                            alg=alg, kwargs...);
                        end
                    end);
    end
    # Once saver has consumed all queued trajectories produced by all workers, an
    # array of MCWF trajs is returned.
    sols = fetch(saver);

    if return_data
        out_type = fout == nothing ? typeof(psi0) : pure_inference(fout, Tuple{eltype(tspan),typeof(psi0)});
        return (tspan, convert(Array{Vector{out_type},1},sols));
    else
        nothing;
    end
end;

function split_threads_mcwf(tspan, psi0::T, H::AbstractOperator{B,B}, J::Vector;
        Ntrajectories=1, progressbar::Bool = true,
        return_data::Bool = true, save_data::Bool = true,
        fpath::Union{String,Missing}=missing,
        additional_data::Union{Dict{String,T2},Missing}=missing,
        seed=nothing, rates::DecayRates=nothing,
        fout=nothing, Jdagger::Vector=dagger.(J),
        display_beforeevent=false, display_afterevent=false,
        alg=OrdinaryDiffEq.AutoTsit5(OrdinaryDiffEq.Rosenbrock23()),
        kwargs...) where {B<:Basis,T<:Ket{B},T2}

    # Create a remote channel from where trajectories are read out by the saver
    remch = RemoteChannel(()->Channel{Any}(Inf)); # TO DO: add some finite buffer size

    # Create a task fetched by the first available worker that retrieves trajs
    # from the remote channel and writes them to disk. A progress bar is set up
    # as well.
    saver = @async launch_saver(remch; Ntrajectories=Ntrajectories,
        progressbar=progressbar, return_data=return_data, save_data=save_data,
        fpath=fpath, additional_data=additional_data);
    wp = CachingPool(workers());

    batches = nfolds(1:Ntrajectories,length(wp.workers))
    # Multi-processed map over all batches of MC trajectories. Maps jobs to
    # workers() from the local process. Jobs consist in computing a batch of
    # MC trajectories and pipe them to the remote channel remch.
    pmap(wp,1:length(wp.workers)) do i
        sol_batch = Array{Any,1}(undef, length(batches[i]))
        Threads.@threads for j in 1:length(batches[i])
            if isnothing(seed)
                sol_batch[j] = timeevolution.mcwf(tspan,psi0,H,J;
                    rates=rates,fout=fout,Jdagger=Jdagger,
                    display_beforeevent=display_beforeevent,
                    display_afterevent=display_afterevent,
                    alg=alg, kwargs...);
            else
                sol_batch[j] = timeevolution.mcwf(tspan,psi0,H,J;
                    seed=seed,rates=rates,fout=fout,Jdagger=Jdagger,
                    display_beforeevent=display_beforeevent,
                    display_afterevent=display_afterevent,
                    alg=alg, kwargs...);
            end
        end
        for sol in sol_batch
            put!(remch,sol);
        end
    end
    # Once saver has consumed all queued trajectories produced by all workers, an
    # array of MCWF trajs is returned.
    sols = fetch(saver);
    # Clear caching pool
    clear!(wp);

    if return_data
        out_type = fout == nothing ? typeof(psi0) : pure_inference(fout, Tuple{eltype(tspan),typeof(psi0)});
        return (tspan, convert(Array{Vector{out_type},1},sols));
    else
        nothing;
    end
end;

function launch_saver(readout_ch::RemoteChannel{Channel{T1}};
        Ntrajectories=1, progressbar::Bool = true, return_data::Bool = true,
        save_data::Bool = true, fpath::Union{String,Missing}=missing,
        additional_data::Union{Dict{String,T2},Missing}=missing) where {T1, T2}
    if progressbar
        # Set up a progress bar
        progress = Progress(Ntrajectories);
        ProgressMeter.update!(progress, 0);
    end
    if save_data
        file = jldopen(fpath, "a+");
        println("Saving data to ",fpath)
        if !ismissing(additional_data)
            for (key, val) in additional_data
                file[key] = val;
            end
        end
    end
    if return_data
        sols::Array{Any,1} = Array{Any,1}(undef,Ntrajectories);
    end

    for i in 1:Ntrajectories
        # Retrieve a queued traj
        sol = take!(readout_ch);
        if save_data && (i == 1) file["t"] = sol[1]; end
        save_data ? file["trajs/" * string(i)] = sol[2] : nothing;
        return_data ? sols[i] = sol[2] : nothing;
        progressbar ? ProgressMeter.next!(progress) : nothing;
    end
    # Set progress bar to 100%
    if progressbar && (progress.counter < Ntrajectories) ProgressMeter.update!(progress, Ntrajectories); end;

    save_data && close(file);
    return return_data ? sols : nothing;
end;

function nfolds(arr,n::Integer)
    foldsize = cld(length(arr),n)
    lastfoldsize = length(arr) - (n-1)*foldsize
    return [[arr[foldsize*(i-1)+1:foldsize*i] for i in 1:n-1]; [arr[end-lastfoldsize+1:end]]]
end;


text:
using DynamicHMCModels

Random.seed!(1)

ProjDir = @__DIR__
cd(ProjDir)

include("simulatePoisson.jl")

struct PoissonProb
   y::Array{Int64,1}
   x::Array{Float64,1}
   idx::Array{Int64,1}
   N::Int64
   Ns::Int64
end

function (problem::PoissonProb)(θ)
   @unpack y, x, idx, N, Ns = problem   # extract the data
   @unpack a0,a1,a0s,a0_sig = θ
   LL = 0.0
   LL += logpdf(Cauchy(0, 1),a0_sig)
   LL += sum(logpdf.(Normal(0,a0_sig),a0s))
   LL += logpdf.(Normal(0, 10),a0)
   LL += logpdf.(Normal(0, 1),a1)
   for i in 1:N
      λ = exp(a0 + a0s[idx[i]] + a1*x[i])
      LL += logpdf(Poisson(λ),y[i])
   end
   return LL
end

y, x, idx, N, Ns = simulatePoisson(;Nd=1,Ns=10,a0=1.0,a1=.5,a0_sig=.3)
p = PoissonProb(y,x,idx,N,Ns)
p((a0=0.0,a1=0.0,a0s=fill(0.0,Ns),a0_sig=.3))

# Write a function to return properly dimensioned transformation.

problem_transformation(p::PoissonProb) =
 as( (a0 = asℝ, a1 = asℝ, a0s = as(Array, Ns), a0_sig = asℝ₊) )

# Use Flux for the gradient.

P = TransformedLogDensity(problem_transformation(p), p)
#∇P = LogDensityRejectErrors(ADgradient(:ForwardDiff, P))

import Zygote
∇P = ADgradient(:Zygote, P)

# FSample from the posterior.
chain, NUTS_tuned = NUTS_init_tune_mcmc(∇P, 1000);
# Undo the transformation to obtain the posterior from the chain.
posterior = TransformVariables.transform.(Ref(problem_transformation(p)), get_position.(chain));

chns = nptochain(posterior,NUTS_tuned)

ext: jl


text:
####
#   normalform (reduction to Zumbroich basis)

"""
    normalform(α::Cyclotomic)
Reduces `α` to the Zumbroich basis.
"""
normalform(α::Cyclotomic) = normalform!(deepcopy(α))

"""
    normalform!(α::Cyclotomic[, tmp=dense(α)]; basis_forbidden=...)
Reduces `α` to the Zumbroich basis in-place. Note that unless `α` has dense
storage the operation will need at least one dense `tmp::Cyclotomic`.
"""
function normalform!(
    α::Cyclotomic{T},
    tmp = dense(α);
    basis_forbidden = zumbroich_viacomplement(conductor(α)),
) where {T}
    # @debug "normalizing:" conductor(α) coeffs(α)
    isnormalized(α, first(basis_forbidden)) && return α

    copyto!(coeffs(tmp), coeffs(α))
    normalform!(tmp, basis_forbidden = basis_forbidden)
    copyto!(coeffs(α), coeffs(tmp))

    return α
end

function normalform!(
    α::Cyclotomic{T,<:DenseVector};
    basis_forbidden = zumbroich_viacomplement(conductor(α)),
) where {T}

    basis, forbidden = basis_forbidden
    isnormalized(α, basis) && return α
    for fb in forbidden
        for exp in exponents(α)
            exp in basis && continue
            _replace_exponent!(α, exp, fb)
        end
    end
    return α
end

function _replace_exponent!(α::Cyclotomic, exp::Integer, fb)
    val = α[exp]
    p, q, forbidden_exps = fb

    # @debug p forbidden_exps
    if exp % q ∈ forbidden_exps
        # @debug "removing $(α[exp])*ζ^$exp from" α
        m = conductor(α) ÷ p
        # @debug " " collect(α[exp .+ 1:m:m*(p-1)])

        # TODO: this doesn't work
        # α[exp .+ 1:m:m*(p-1)] .-= val

        for i = 1:p-1
            α[exp+i*m] -= val
            # exp + i*m is not guaranteed to be allowed, but it's larger than exp and hence will be dealt later
        end
        α[exp] = 0
        # @debug "after removal:" α
    end
    return α
end

"""
    embed(α::Cyclotomic, m::Integer)
Embed `α` into the `m`-th cyclotomic field. `m` can be either larger, or smaller
than the conductor of `α`, however either `conductor(α)` must divide `m`, or
the other way around.
"""
function embed(α::Cyclotomic, m::Integer)
    conductor(α) == m && return deepcopy(α)
    if conductor(α) > m
        @assert conductor(α) % m == 0 "Embeding of ℚ(ζ$(subscriptify(m))) ↪ ℚ(ζ$(subscriptify(conductor(α)))) is not possible."
        return reduced_embedding(α, m)
    else
        k, _r = divrem(m, conductor(α))
        @assert _r == 0 "Embeding of ℚ(ζ$(subscriptify(conductor(α)))) ↪ ℚ(ζ$(subscriptify(m))) is not possible."

        res = zero!(similar(α, valtype(α), m))

        @inbounds for e in exponents(α)
            res[e*k] = α[e]
        end
        return res
    end
end

function _tmp_for_reduced_embedding(α::Cyclotomic{T}) where {T}

    all(iszero, exponents(α)) && return Cyclotomic{T,Vector{T}}(1, [α[0]])

    k = gcd(conductor(α), exponents(α)...)
    @debug "gcd(exponents(α)) = $k" collect(exponents(α))

    tmp = if k > 1
        d = div(conductor(α), k)
        @debug "Reducing the embedding ring to $T(ζ$(subscriptify(d)))"
        tmp = Cyclotomic{T,Vector{T}}(d, zeros(T, d))
    else
        n = conductor(α)
        tmp = Cyclotomic{T,Vector{T}}(n, zeros(T, n))
    end
    # now tmp has dense storage

    if k > 1
        # the trivial reduction E(n)^e → E(n÷k)^(e÷k)
        @debug "Performing trivial reduction from ℚ(ζ$(subscriptify(conductor(α)))) → ℚ(ζ$(subscriptify(conductor(tmp))))"
        @inbounds for (e, c) in exps_coeffs(α)
            tmp[div(e, k)] = c
        end
    else
        @debug "No trivial reduction is possible"
        copyto!(tmp.coeffs, coeffs(α))
    end
    return tmp
end

"""
    reduced_embedding(α::Cyclotomic{T,V}[, m::Integer=1])
Return the reduced embedding of `α` into `m`-th cyclotomic field.

When `m=1`, the embedding into possibly smallest degree is returned.
"""
function reduced_embedding(α::Cyclotomic{T,V}, m::Integer = 1) where {T,V}

    tmp = _tmp_for_reduced_embedding(normalform!(α))

    if conductor(tmp) == 1
        res = Cyclotomic{T,V}(conductor(tmp), coeffs(tmp))
        return res
    end

    basis, forbidden = Cyclotomics.zumbroich_viacomplement(conductor(tmp))

    normalform!(tmp, basis_forbidden = (basis, forbidden))

    phi_nc = length(basis)
    nz = count(!iszero, coeffs(tmp))

    if all(p == q for (p, q, _) in forbidden) && # conductor(tmp) is square-free
       phi_nc == nz && # tmp is supported on φ(n) elements
       _all_equal(tmp, exponents(tmp)) # all nz coeffs of tmp are equal

        val = if iseven(length(forbidden))
            tmp[first(exponents(tmp))]
        else
            -tmp[first(exponents(tmp))]
        end
        @debug "Cyclotomic is real:" α = val
        res = Cyclotomic{T,V}(m, similar(coeffs(α), m))
        zero!(res)
        res[0] = val
        return res
    end

    for (p, q, fb_res) in forbidden
        p == 2 && continue
        n = conductor(tmp)
        nz = count(!iszero, coeffs(tmp))

        p == q || continue
        @debug "$(p)² does NOT divide conductor $(conductor(tmp))" p q

        m % p != 0 || continue
        @debug "prime $p does not divide $m so embedding is possible"

        # there are n/p residue classes, each contains p-1 elts, and tmp is constant on each
        nnz, _r = divrem(nz, p - 1)
        @debug "the number of non-zero coeffs $nz ≡ $_r mod $(p-1)"
        _r == 0 || continue

        # Check that coeffs(tmp) are constant on each congruence classes.
        # To illustrate by example write:
        #
        # 1 = - ζ₅ - ζ₅² - ζ₅³ - ζ₅⁴
        #
        # and consider
        #
        # c*ζ₉^e = c*ζ₉^e * 1 = c*ζ₉^e * (- ζ₅ - ζ₅² - ζ₅³ - ζ₅⁴) =
        # - c*ζ₄₅^(5e+1*9) - c*ζ₄₅^(5e+2*9) - c*ζ₄₅^(5e+3*9) - c*ζ₄₅^(5e+4*9)
        #
        # If tmp ∈ ℚ(ζ_n) can be reduced to ℚ(ζ_{n÷p})
        # (n = 45, p = 5 in the example), we need to check that all coefficients
        #
        # tmp[p*e + i*n÷p] are equal (for 1 ≤ i ≤ p-1)
        #
        # where exponents 0 ≤ e ≤ n-1 represent congruence classes modulo n÷p

        n_p = n ÷ p

        equal_on_classes = true
        for pe = 0:p:(n-1)
            equal_on_classes || break
            if !_all_equal(tmp, pe .+ (n_p .* (1:p-1)))
                equal_on_classes = false
                @debug "skipping reduction over prime $p as cyclotomic is not constant over class $(pe + (n_p)) (mod $(n_p)) " (
                    pe + n_p,
                    tmp[pe+n_p],
                ) (pe .+ (n_p .* (2:p-1)), tmp[pe.+(n_p.*(2:p-1))])
            end
        end
        equal_on_classes || continue # to the next prime

        @debug "cyclotomic is constant over residue classes (mod $(n_p)), reducing"

        for pe = 0:p:(n-1)
            # replace sum_i(tmp[p*e + i*n÷p]) (1 ≤ i ≤ p-1)
            # inverse to normalform! over p
            @debug "replacing $pe-th power by $(pe + n_p)-th:" tmp[pe] -tmp[pe+n_p]
            tmp[pe] = -tmp[pe+n_p] # tmp[pe + 0*n_p] = -tmp[pe + 1*n_p]
            tmp[pe.+(n_p.*(1:p-1))] .= zero(T)
        end

        for i = 1:n_p
            tmp[i] = tmp[i*p]
            tmp[i*p] = zero(T)
        end

        tmp = Cyclotomic{T,Vector{T}}(n_p, resize!(tmp.coeffs, n_p))
    end

    res = Cyclotomic{T,V}(conductor(tmp), coeffs(tmp))

    return res
end


text:
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.6.1
#     language: julia
#     name: julia-1.6
# ---

# %% [markdown]
# # Getting even faster by reducing allocations and by multi-threading
#
# * Gen Kuroki
# * 2021-06-28

# %%
VERSION

# %% [markdown]
# ## Original version
#
# * https://twitter.com/elbersb/status/1409454534666207232
# * https://elbersb.com/public/posts/interaction_simulation/

# %%
using Distributions
using DataFrames

const std_normal = Normal(0, 1)

function gen_data()
    ## total time periods in the the panel = 500
    tt = 500

    # x1 and x2 covariates
    x1_A = 1 .+ rand(std_normal, tt)
    x1_B = 1/4 .+ rand(std_normal, tt)
    x2_A = 1 .+ x1_A .+ rand(std_normal, tt)
    x2_B = 1 .+ x1_B .+ rand(std_normal, tt)

    # outcomes (notice different slope coefs for x2_A and x2_B)
    y_A = x1_A .+ 1*x2_A + rand(std_normal, tt)
    y_B = x1_B .+ 2*x2_B + rand(std_normal, tt)

    # combine
    DataFrame(
        id = vcat(fill(0, length(x1_A)), fill(1, length(x1_B))),
        x1 = vcat(x1_A, x1_B),
        x2 = vcat(x2_A, x2_B),
        x1_dmean = vcat(x1_A .- mean(x1_A), x1_B .- mean(x1_B)),
        x2_dmean = vcat(x2_A .- mean(x2_A), x2_B .- mean(x2_B)),
        y = vcat(y_A, y_B))
end

# %%
using GLM

function coefs_lm_formula(data)
    mod_level = lm(@formula(y ~ id + x1 * x2), data)
    mod_dmean = lm(@formula(y ~ id + x1_dmean * x2_dmean), data)
    (coef(mod_level)[end], coef(mod_dmean)[end])
end

# example
data = gen_data()
coefs_lm_formula(data)

# %%
function run_simulations(nsim)
    sims = zeros(nsim, 2);
    for i in 1:nsim
        data = gen_data()
        sims[i, :] .= coefs_lm_formula(data)
    end
    sims
end

using BenchmarkTools
n = 20000
@btime run_simulations($n);

# %%
using Plots
sims = run_simulations(n)
histogram(sims, label = ["level" "dmean"])

# %%
function coefs_lm_formula(data)
    constant = fill(1, nrow(data))
    X = Float64[constant data.id data.x1 data.x2 data.x1 .* data.x2]
    mod_level = fit(LinearModel, X, data.y)
    X[:, 5] .= data.x1_dmean .* data.x2_dmean
    mod_dmean = fit(LinearModel, X, data.y)
    (coef(mod_level)[end], coef(mod_dmean)[end])
end

@btime run_simulations($n);

# %%
using LinearAlgebra
fastfit(X, y) = cholesky!(Symmetric(X' * X)) \ (X' * y)

function coefs_lm_formula(data)
    constant = fill(1, nrow(data))
    X = Float64[constant data.id data.x1 data.x2 data.x1 .* data.x2]
    mod_level = fastfit(X, data.y)
    X[:, 5] .= data.x1_dmean .* data.x2_dmean
    mod_dmean = fastfit(X, data.y)
    (mod_level[end], mod_dmean[end])
end

@btime run_simulations($n);

# %% [markdown]
# ## Optimized version: Reducing allocations
#
# Ref. [Performance Tips: Pre-allocating outputs](https://docs.julialang.org/en/v1/manual/performance-tips/#Pre-allocating-outputs)

# %%
using DataFrames
using Distributions
using LinearAlgebra
using Random: rand!
using BenchmarkTools
using Plots

const std_normal = Normal(0, 1)

function undef_data(tt)
    id = vcat(fill(0, tt), fill(1, tt))
    x1 = similar(id, Float64)
    x2 = similar(x1)
    x1_dmean = similar(x1)
    x2_dmean = similar(x2)
    y = similar(x1)
    DataFrame(id = id, x1 = x1, x2 = x2, x1_dmean = x1_dmean, x2_dmean = x2_dmean, y = y)
end

function gen_data!(data, tmp)
    tt = length(tmp) # Assume 2tt = nrow(data)
    A, B = 1:tt, tt+1:2tt

    # Warning: dataframe is type-unstable
    x1::Vector{Float64} = data.x1
    x2::Vector{Float64} = data.x2
    x1_dmean::Vector{Float64} = data.x1_dmean
    x2_dmean::Vector{Float64} = data.x2_dmean
    y::Vector{Float64} = data.y

    @views x1[A] .= 1   .+ rand!(std_normal, tmp)
    @views x1[B] .= 1/4 .+ rand!(std_normal, tmp)
    @views x2[A] .= 1   .+ x1[A] .+ rand!(std_normal, tmp)
    @views x2[B] .= 1   .+ x1[B] .+ rand!(std_normal, tmp)

    # outcomes (notice different slope coefs for x2[A] and x2[B])
    @views y[A] .= x1[A] .+ 1 .* x2[A] .+ rand!(std_normal, tmp)
    @views y[B] .= x1[B] .+ 2 .* x2[B] .+ rand!(std_normal, tmp)

    @views x1_dmean[A] .= x1[A] .- mean(x1[A])
    @views x1_dmean[B] .= x1[B] .- mean(x1[B])
    @views x2_dmean[A] .= x2[A] .- mean(x2[A])
    @views x2_dmean[B] .= x2[B] .- mean(x2[B])

    data
end

function fastfit!(mod, X, y, XtX, Xty)
    mul!(XtX, X', X)
    mul!(Xty, X', y)
    ldiv!(mod, cholesky!(Symmetric(XtX)), Xty)
end

function coefs_lm_formula!(result, data, X, mod_level, mod_dmean, XtX, Xty)
    # Warning: dataframe is type-unstable.
    id::Vector{Int} = data.id
    x1::Vector{Float64} = data.x1
    x2::Vector{Float64} = data.x2
    x1_dmean::Vector{Float64} = data.x1_dmean
    x2_dmean::Vector{Float64} = data.x2_dmean
    y::Vector{Float64} = data.y

    X[:, 1] .= 1
    X[:, 2] .= id
    X[:, 3] .= x1
    X[:, 4] .= x2
    X[:, 5] .= x1 .* x2
    fastfit!(mod_level, X, y, XtX, Xty)

    X[:, 5] .= x1_dmean .* x2_dmean
    fastfit!(mod_dmean, X, y, XtX, Xty)

    result .= (mod_level[end], mod_dmean[end])
end

function run_simulations_optimized(nsim, tt=500)
    sims = Matrix{Float64}(undef, nsim, 2)
    data = undef_data(tt)
    tmp = Vector{Float64}(undef, tt)
    X = Matrix{Float64}(undef, 2tt, 5)
    mod_level = Vector{Float64}(undef, 5)
    mod_dmean = similar(mod_level)
    XtX = Matrix{Float64}(undef, 5, 5)
    Xty = Vector{Float64}(undef, 5)
    for i in 1:nsim
        gen_data!(data, tmp)
        @views coefs_lm_formula!(sims[i, :], data, X, mod_level, mod_dmean, XtX, Xty)
    end
    sims
end

@show n = 20000
@time sims = run_simulations_optimized(n)
histogram(sims, label = ["level" "dmean"])

# %%
tt = 500
data = undef_data(tt)
tmp = Vector{Float64}(undef, tt)
X = Matrix{Float64}(undef, 2tt, 5)
mod_level = Vector{Float64}(undef, 5)
mod_dmean = similar(mod_level)
XtX = Matrix{Float64}(undef, 5, 5)
Xty = Vector{Float64}(undef, 5)
sim = zeros(2)
n = 20000

println("Optimized version:")
@btime gen_data!($data, $tmp)
@btime coefs_lm_formula!($sim, $data, $X, $mod_level, $mod_dmean, $XtX, $Xty)
@btime run_simulations_optimized($n);

# %%
println("Original version:")
@btime gen_data()
@btime coefs_lm_formula($data)
@btime run_simulations($n);

# %% [markdown]
# ## Multi-thread original version

# %%
function run_simulations_multithread(nsim)
    sims = zeros(nsim, 2);
    Threads.@threads for i in 1:nsim
        data = gen_data()
        sims[i, :] .= coefs_lm_formula(data)
    end
    sims
end

@show Threads.nthreads()
@time sims = run_simulations_multithread(n)
histogram(sims, label = ["level" "dmean"])

# %%
println("Multi-thread original version: ")
@show Threads.nthreads()
@btime run_simulations_multithread($n);

# %% [markdown]
# ## Multi-thread optimized version

# %%
# pkg> dev https://github.com/genkuroki/MyUtils.jl
using MyUtils: @my_threads

function run_simulations_multithread_optimized(nsim, tt=500)
    sims = Matrix{Float64}(undef, nsim, 2)
    @my_threads begin
        data = undef_data(tt)
        tmp = Vector{Float64}(undef, tt)
        X = Matrix{Float64}(undef, 2tt, 5)
        mod_level = Vector{Float64}(undef, 5)
        mod_dmean = similar(mod_level)
        XtX = Matrix{Float64}(undef, 5, 5)
        Xty = Vector{Float64}(undef, 5)
    end for i in 1:nsim
        gen_data!(data, tmp)
        @views coefs_lm_formula!(sims[i, :], data, X, mod_level, mod_dmean, XtX, Xty)
    end begin
    end
    sims
end

@show Threads.nthreads()
@time sims = run_simulations_multithread_optimized(n)
histogram(sims, label = ["level" "dmean"])

# %%
println("Multi-thread optimized version: ")
@show Threads.nthreads()
@btime run_simulations_multithread_optimized($n);

# %% [markdown]
# ## non-DataFrame version
#
# https://discourse.julialang.org/t/simulation-of-regression-coefficients/63683/2

# %%
using Distributions
using LinearAlgebra
using BenchmarkTools
using Plots
using Base.Threads

const std_normal = Normal(0, 1)
fastfit(X, y) = cholesky!(Symmetric(X' * X)) \ (X' * y)

function gen_data2!(X, y, x1_dmean, x2_dmean)
    ## total time periods in the the panel = 500
    tt = 500

    # x1 and x2 covariates
    x1_A = 1 .+ rand(std_normal, tt)
    x1_B = 1/4 .+ rand(std_normal, tt)
    x2_A = 1 .+ x1_A .+ rand(std_normal, tt)
    x2_B = 1 .+ x1_B .+ rand(std_normal, tt)

    # outcomes (notice different slope coefs for x2_A and x2_B)
    y_A = x1_A .+ 1*x2_A + rand(std_normal, tt)
    y_B = x1_B .+ 2*x2_B + rand(std_normal, tt)
    x1 = vcat(x1_A, x1_B)
    x2 = vcat(x2_A, x2_B)

    X[:, 2] .= vcat(fill(0, length(x1_A)), fill(1, length(x1_B)))
    X[:, 3] .= x1
    X[:, 4] .= x2
    X[:, 5] .= (x1 .* x2)

    y .= vcat(y_A, y_B)

    x1_dmean .= vcat(x1_A .- mean(x1_A), x1_B .- mean(x1_B))
    x2_dmean .= vcat(x2_A .- mean(x2_A), x2_B .- mean(x2_B))
end

function coefs_lm_formula4(X, y, x1_dmean, x2_dmean)
    mod_level = fastfit(X, y)
    @inbounds X[:, 5] .= x1_dmean .* x2_dmean
    mod_dmean = fastfit(X, y)
    (mod_level[end], mod_dmean[end])
end

function run_simulations5(nsim)
    sims = zeros(nsim, 2);
    X = [zeros(Float64, 1000, 5) for _ in 1:nthreads()]
    y = [zeros(Float64, 1000) for _ in 1:nthreads()]
    x1_dmean = [zeros(Float64, 1000) for _ in 1:nthreads()]
    x2_dmean = [zeros(Float64, 1000) for _ in 1:nthreads()]

    # set constant
    for i in 1:nthreads()
        X[i][:, 1] .= 1
    end

    @threads for i in 1:nsim
        gen_data2!(X[threadid()], y[threadid()], x1_dmean[threadid()], x2_dmean[threadid()])
        @inbounds sims[i, 1], sims[i, 2] = coefs_lm_formula4(X[threadid()], y[threadid()], x1_dmean[threadid()], x2_dmean[threadid()])
    end
    sims
end

# %%
n = 20000
@show Threads.nthreads()
@time sims = run_simulations5(n)
histogram(sims, label = ["level" "dmean"])

# %%
println("Multi-thread non-dataframe version (run_simulations5): ")
@show Threads.nthreads()
@btime run_simulations5($n);

# %% [markdown]
# ## Optimized non-DataFrame version

# %%
using Distributions
using LinearAlgebra
using Random: rand!
using BenchmarkTools
using Plots

const std_normal = Normal(0, 1)

# use NamedTuple indtead of DataFrame
function undef_data_namedtuple(tt)
    id = vcat(fill(0, tt), fill(1, tt))
    x1 = similar(id, Float64)
    x2 = similar(x1)
    x1_dmean = similar(x1)
    x2_dmean = similar(x2)
    y = similar(x1)
    (; id, x1, x2, x1_dmean, x2_dmean, y)
end

function gen_data!(data, tmp)
    tt = length(tmp) # Assume 2tt = nrow(data)
    A, B = 1:tt, tt+1:2tt

    # Warning: dataframe is type-unstable
    x1::Vector{Float64} = data.x1
    x2::Vector{Float64} = data.x2
    x1_dmean::Vector{Float64} = data.x1_dmean
    x2_dmean::Vector{Float64} = data.x2_dmean
    y::Vector{Float64} = data.y

    @views x1[A] .= 1   .+ rand!(std_normal, tmp)
    @views x1[B] .= 1/4 .+ rand!(std_normal, tmp)
    @views x2[A] .= 1   .+ x1[A] .+ rand!(std_normal, tmp)
    @views x2[B] .= 1   .+ x1[B] .+ rand!(std_normal, tmp)

    # outcomes (notice different slope coefs for x2[A] and x2[B])
    @views y[A] .= x1[A] .+ 1 .* x2[A] .+ rand!(std_normal, tmp)
    @views y[B] .= x1[B] .+ 2 .* x2[B] .+ rand!(std_normal, tmp)

    @views x1_dmean[A] .= x1[A] .- mean(x1[A])
    @views x1_dmean[B] .= x1[B] .- mean(x1[B])
    @views x2_dmean[A] .= x2[A] .- mean(x2[A])
    @views x2_dmean[B] .= x2[B] .- mean(x2[B])

    data
end

function fastfit!(mod, X, y, XtX, Xty)
    mul!(XtX, X', X)
    mul!(Xty, X', y)
    ldiv!(mod, cholesky!(Symmetric(XtX)), Xty)
end

function coefs_lm_formula!(result, data, X, mod_level, mod_dmean, XtX, Xty)
    # Warning: dataframe is type-unstable.
    id::Vector{Int} = data.id
    x1::Vector{Float64} = data.x1
    x2::Vector{Float64} = data.x2
    x1_dmean::Vector{Float64} = data.x1_dmean
    x2_dmean::Vector{Float64} = data.x2_dmean
    y::Vector{Float64} = data.y

    X[:, 1] .= 1
    X[:, 2] .= id
    X[:, 3] .= x1
    X[:, 4] .= x2
    X[:, 5] .= x1 .* x2
    fastfit!(mod_level, X, y, XtX, Xty)

    X[:, 5] .= x1_dmean .* x2_dmean
    fastfit!(mod_dmean, X, y, XtX, Xty)

    result .= (mod_level[end], mod_dmean[end])
end

function run_simulations_optimized_nondataframe(nsim, tt=500)
    sims = Matrix{Float64}(undef, nsim, 2)
    data = undef_data_namedtuple(tt) # use NamedTuple indtead of DataFrame
    tmp = Vector{Float64}(undef, tt)
    X = Matrix{Float64}(undef, 2tt, 5)
    mod_level = Vector{Float64}(undef, 5)
    mod_dmean = similar(mod_level)
    XtX = Matrix{Float64}(undef, 5, 5)
    Xty = Vector{Float64}(undef, 5)
    for i in 1:nsim
        gen_data!(data, tmp)
        @views coefs_lm_formula!(sims[i, :], data, X, mod_level, mod_dmean, XtX, Xty)
    end
    sims
end

@show n = 20000
@time sims = run_simulations_optimized_nondataframe(n)
histogram(sims, label = ["level" "dmean"])

# %%
println("Optimized non-dataframe version (single thread): ")
@btime run_simulations_optimized_nondataframe($n);

# %% [markdown]
# ## Multi-thread optimized non-DataFrame version

# %%
# pkg> dev https://github.com/genkuroki/MyUtils.jl
using MyUtils: @my_threads

function run_simulations_multithread_optimized_nondataframe(nsim, tt=500)
    sims = Matrix{Float64}(undef, nsim, 2)
    @my_threads begin
        data = undef_data_namedtuple(tt)
        tmp = Vector{Float64}(undef, tt)
        X = Matrix{Float64}(undef, 2tt, 5)
        mod_level = Vector{Float64}(undef, 5)
        mod_dmean = similar(mod_level)
        XtX = Matrix{Float64}(undef, 5, 5)
        Xty = Vector{Float64}(undef, 5)
    end for i in 1:nsim
        gen_data!(data, tmp)
        @views coefs_lm_formula!(sims[i, :], data, X, mod_level, mod_dmean, XtX, Xty)
    end begin
    end
    sims
end

@show Threads.nthreads()
@time sims = run_simulations_multithread_optimized_nondataframe(n)
histogram(sims, label = ["level" "dmean"])

# %%
println("Multi-thread optimized non-dataframe version: ")
@show Threads.nthreads()
@btime run_simulations_multithread_optimized_nondataframe($n);

# %% [markdown]
# ## Summary

# %%
n = 20000

@show VERSION
@show n
println()

println("Original version (with DataFrame, single thread):")
@btime run_simulations($n)
println()

println("Optimized version (with DataFrame, single thread):")
@btime run_simulations_optimized($n)
println()

println("Multi-thread optimized version (with DataFrame): ")
@show Threads.nthreads()
@btime run_simulations_multithread_optimized($n)
println()

println("Multi-thread non-dataframe version (run_simulations5): ")
@show Threads.nthreads()
@btime run_simulations5($n)
println()

println("Optimized non-dataframe version (single thread): ")
@btime run_simulations_optimized_nondataframe($n)
println()

println("Multi-thread optimized non-dataframe version: ")
@show Threads.nthreads()
@btime run_simulations_multithread_optimized_nondataframe($n)
println()

# %%
versioninfo()

# %%

max_stars_repo_path: 0008/Getting even faster by reducing allocations and by multi-threading.jl

max_stars_repo_name: genkuroki/public

max_issues_repo_name: genkuroki/public

max_forks_repo_name: genkuroki/public


text:
module Parallelism

using Base.Threads
using Distributed
using LinearAlgebra
using Memento

export robust_pmap, tmap, tmap_with_warmup

const MODULE = @__MODULE__
const LOGGER = getlogger(MODULE)

include("robust_pmap.jl")
include("tmap.jl")

end # module


text:
module MatrixPencils

const BlasFloat = Union{Float64,Float32,ComplexF64,ComplexF32}
const BlasReal = Union{Float64,Float32}
const BlasComplex = Union{ComplexF64,ComplexF32}

using LinearAlgebra
using Polynomials

include("lapackutil2.jl")
import .LapackUtil2: larfg!, larfgl!, larf!, gghrd!, hgeqz!, tgexc!, tgsen!, lanv2, lag2, safemin, tgsyl!

export preduceBF, klf, klf_left, klf_leftinf, klf_right, klf_rightinf, klf_rlsplit,
       klf_right!, klf_right_refine!, klf_right_refineut!, klf_right_refineinf!,
       klf_left!, klf_left_refine!, klf_left_refineut!, klf_left_refineinf!
export prank, pkstruct, peigvals, pzeros, KRInfo
export isregular, isunimodular, fisplit, _svdlikeAE!, sfisplit
export sreduceBF, sklf, gsklf, sklf_right, sklf_left, sklf_right!, sklf_right2!, sklf_left!, sklf_rightfin!, sklf_rightfin2!, sklf_leftfin!
export sprank, spkstruct, speigvals, spzeros
export lsminreal, lsminreal2, lsequal, lseval, lps2ls
export lpsminreal, lpsequal, lpseval
export poldivrem, polgcdvw, pollcm, conv, poldiv, convmtx, gcdvwupd, qrsolve!, poldeg, poldeg1
export poly2pm, pm2poly, pmdeg, pmreverse, pmeval, pmdivrem
export pm2lpCF1, pm2lpCF2, pm2lps, pm2ls, spm2ls, spm2lps, ls2pm, lps2pm
export pmkstruct, pmeigvals, pmzeros, pmzeros1, pmzeros2, pmroots, pmpoles, pmpoles1, pmpoles2, pmrank, ispmregular, ispmunimodular
export rmeval, rm2lspm, rm2ls, ls2rm, rm2lps, lps2rm
export lpmfd2ls, rpmfd2ls, lpmfd2lps, rpmfd2lps, pminv2ls, pminv2lps
export rmkstruct, rmzeros, rmzeros1, rmpoles, rmpoles1, rmrank
export saloc, salocd, salocinf, salocinfd, ordeigvals, isqtriu, eigselect1, eigselect2, saloc2
export fihess, fischur, fischursep, sfischursep, fiblkdiag, gsblkdiag, ssblkdiag
import LinearAlgebra: copy_oftype

include("klftools.jl")
include("regtools.jl")
include("klfapps.jl")
include("sklftools.jl")
include("sklfapps.jl")
include("lstools.jl")
include("pmtools.jl")
include("pmapps.jl")
include("poltools.jl")
include("rmtools.jl")
include("rmapps.jl")
include("lputil.jl")
include("slputil.jl")
include("gsfstab.jl")
include("gsep.jl")

end


text:
reshapeVec2Mat(vec::Vector, rows::Int) = reshape(vec, rows, round(Int,length(vec)/rows))
# function reshapeVec2Mat(vec::Vector, rows::Int)
#   M = reshape(vec, rows, round(Int,length(vec)/rows))
#   return ndims(M) < 2 ? (M')' : M
# end

# get vertex from factor graph according to label symbol "x1"
getVert(fgl::FactorGraph, lbl::Symbol; api::DataLayerAPI=dlapi, nt::Symbol=:var) = api.getvertex(fgl, lbl, nt=nt)
getVert(fgl::FactorGraph, id::Int; api::DataLayerAPI=dlapi) = api.getvertex(fgl, id)

# TODO -- upgrade to dedicated memory location in Graphs.jl
# see JuliaArchive/Graphs.jl#233
getData(v::Graphs.ExVertex) = v.attributes["data"]
# Convenience functions
getData(fgl::FactorGraph, lbl::Symbol; api::DataLayerAPI=dlapi) = getData(getVert(fgl, lbl, api=api))
getData(fgl::FactorGraph, id::Int; api::DataLayerAPI=dlapi) = getData(getVert(fgl, id, api=api))

function setData!(v::Graphs.ExVertex, data)
  v.attributes["data"] = data
  nothing
end

function getVal(v::Graphs.ExVertex)
  return getData(v).val
end
function getVal(v::Graphs.ExVertex, idx::Int)
  return getData(v).val[:,idx]
end

# Convenience function to get values for given variable label
function getVal(fgl::FactorGraph, lbl::Symbol; api::DataLayerAPI=dlapi)
  #getVal(dlapi.getvertex(fgl, lbl))
  getVal(getVert(fgl, lbl, api=api))
end
function getVal(fgl::FactorGraph, exvertid::Int; api::DataLayerAPI=dlapi)
  # getVal(dlapi.getvertex(fgl, exvertid))
  getVal(getVert(fgl, exvertid, api=api))
end

function getfnctype(vertl::Graphs.ExVertex)
  data = getData(vertl)
  if typeof(data).name.name == :VariableNodeData
    return VariableNodeData
  end
  data.fnc.usrfnc!
end

function getfnctype(fgl::FactorGraph, exvertid::Int; api::DataLayerAPI=dlapi)
  #
  # data = getData(fgl, exvertid, api=api)
  # data.fnc.usrfnc!
  getfnctype(getVert(fgl, exvertid, api=api))
end

# setVal! assumes you will update values to database separate, this used for local graph mods only
function setVal!(v::Graphs.ExVertex, val::Array{Float64,2})
  v.attributes["data"].val = val
  nothing
end
function getBWVal(v::Graphs.ExVertex)
  return getData(v).bw
end
function setBW!(v::Graphs.ExVertex, bw::Array{Float64,2})
  v.attributes["data"].bw = bw
  nothing
end
function setVal!(v::Graphs.ExVertex, val::Array{Float64,2}, bw::Array{Float64,2})
  setVal!(v,val)
  setBW!(v,bw)
  nothing
end
function setVal!(v::Graphs.ExVertex, val::Array{Float64,2}, bw::Vector{Float64})
  setVal!(v,val,reshape(bw,length(bw),1)) #(bw')')
  nothing
end
function setValKDE!(v::Graphs.ExVertex, val::Array{Float64,2})
  p = kde!(val)
  setVal!(v,val,getBW(p)[:,1]) # TODO -- this can be little faster
  nothing
end
function setValKDE!(v::Graphs.ExVertex, em::EasyMessage)
  setVal!(v, em.pts, em.bws ) # getBW(p)[:,1]
  nothing
end
function setValKDE!(v::Graphs.ExVertex, p::BallTreeDensity)
  pts = getPoints(p)
  setVal!(v, pts, getBW(p)[:,1]) # BUG ...al!(., val, . ) ## TODO -- this can be little faster
  nothing
end
setVal!(v::Graphs.ExVertex, em::EasyMessage) = setValKDE!(v, em)
setVal!(v::Graphs.ExVertex, p::BallTreeDensity) = setValKDE!(v, p)

function kde!(em::EasyMessage)
  return kde!(em.pts,em.bws)
end

# TODO -- there should be a better way, without retrieving full vertex
getOutNeighbors(fgl::FactorGraph, v::ExVertex; api::DataLayerAPI=dlapi, needdata::Bool=false, ready::Int=1,backendset::Int=1 ) = api.outneighbors(fgl, v, needdata=needdata, ready=ready, backendset=backendset )
getOutNeighbors(fgl::FactorGraph, vertid::Int; api::DataLayerAPI=dlapi, needdata::Bool=false, ready::Int=1,backendset::Int=1 ) = api.outneighbors(fgl, api.getvertex(fgl,vertid), needdata=needdata, ready=ready, backendset=backendset )

function updateFullVert!(fgl::FactorGraph, exvert::ExVertex;
            api::DataLayerAPI=IncrementalInference.dlapi,
            updateMAPest::Bool=false  )
  #
  warn("use of updateFullVert! should be clarified for local or remote operations.")
  api.updatevertex!(fgl, exvert, updateMAPest=updateMAPest)
end

function setDefaultNodeData!(v::Graphs.ExVertex, initval::Array{Float64,2},
                             stdev::Array{Float64,2}, dodims::Int, N::Int, dims::Int;
                             gt=nothing, initialized::Bool=true,
                             softtype=nothing)
  # TODO review and refactor this function, exists as legacy from pre-v0.3.0
  # this should be the only function allocating memory for the node points (unless number of points are changed)
  data = nothing
  if initialized
    if size(initval,2) < N && size(initval, 1) == dims
      warn("setDefaultNodeData! -- deprecated use of stdev.")
      p = kde!(initval,diag(stdev));
      pN = resample(p,N)
    elseif size(initval,2) < N && size(initval, 1) != dims
      println("Node value memory allocated but not initialized")
      pN = kde!(randn(dims, N));
    else
      pN = kde!(initval)
    end
    # dims = size(initval,1) # rows indicate dimensions
    sp = round.(Int,linspace(dodims,dodims+dims-1,dims))
    gbw = getBW(pN)[:,1]
    gbw2 = Array{Float64}(length(gbw),1)
    gbw2[:,1] = gbw[:]
    pNpts = getPoints(pN)
    data = VariableNodeData(initval, stdev, pNpts,
                            gbw2, Int[], sp,
                            dims, false, 0, Int[], gt, softtype, true) #initialized
  else
      sp = round.(Int,linspace(dodims,dodims+dims-1,dims))
      data = VariableNodeData(initval, stdev, zeros(dims, N),
                              zeros(dims,1), Int[], sp,
                              dims, false, 0, Int[], gt, softtype, false) #initialized
  end
  #
  v.attributes["data"] = data

  nothing
end

function addNewVarVertInGraph!(fgl::FactorGraph, vert::Graphs.ExVertex, id::Int, lbl::Symbol, ready::Int)
  vert.attributes = Graphs.AttributeDict() #fg.v[fg.id]
  vert.attributes["label"] = string(lbl) #fg.v[fg.id]
  fgl.IDs[lbl] = id

  # used for cloudgraph solving
  vert.attributes["ready"] = ready
  vert.attributes["backendset"] = 0

  # fgl.g.vertices[id] = vert # will be inserted during addvertex! call
  # fgl.v[id] = vert #  -- this is likely not required, but is used in subgraph methods
  nothing
end

# must set either dims or initval for proper initialization
# Add node to graph, given graph struct, labal, init values,
# std dev [TODO -- generalize], particle size and ready flag for concurrency
function addNode!(fg::FactorGraph,
      lbl::Symbol,
      initval::Array{Float64}=zeros(1,1),
      stdev::Array{Float64}=ones(1,1); # this is bad and should be removed TODO
      N::Int=100,
      ready::Int=1,
      labels::Vector{T}=String[],
      api::DataLayerAPI=dlapi,
      uid::Int=-1,
      dims::Int=-1  ) where {T <: AbstractString}
  #
  warn("this addNode! will be deprecated, please use FactorGraph01.jl:addNode!(fg::FactorGraph, lbl::Symbol, softtype::Type{T}).")
  currid = fg.id+1
  if uid==-1
    fg.id=currid
  else
    currid = uid
  end
  dims = dims != -1 ? dims : size(initval,1)

  lblstr = string(lbl)
  vert = ExVertex(currid,lblstr)
  addNewVarVertInGraph!(fg, vert, currid, lbl, ready)
  # dlapi.setupvertgraph!(fg, vert, currid, lbl) #fg.v[currid]
  dodims = fg.dimID+1
  # TODO -- vert should not loose information here
  setDefaultNodeData!(vert, initval, stdev, dodims, N, dims) #fg.v[currid]

  vnlbls = deepcopy(labels)
  push!(vnlbls, fg.sessionname)
  # addvert!(fg, vert, api=api)
  api.addvertex!(fg, vert, labels=vnlbls) #fg.g ##vertr =

  fg.dimID+=dims # rows indicate dimensions, move to last dimension
  push!(fg.nodeIDs, currid)

  return vert #fg.v[fg.id]
end

"""
$(SIGNATURES)

Add a node (variable) to a graph. Use this over the other dispatches.
"""
function addNode!(fg::FactorGraph,
      lbl::Symbol,
      softtype::Type{<:InferenceVariable};
      N::Int=100,
      autoinit=true,  # does init need to be separate from ready? TODO
      ready::Int=1,
      labels::Vector{<:AbstractString}=String[],
      api::DataLayerAPI=dlapi,
      uid::Int=-1,
      dims::Int=-1  ) # where {T , S }
  #
  currid = fg.id+1
  if uid==-1
    fg.id=currid
  else
    currid = uid
  end
  st = softtype()
  dims = dims != -1 ? dims : st.dims

  lblstr = string(lbl)
  vert = ExVertex(currid,lblstr)
  addNewVarVertInGraph!(fg, vert, currid, lbl, ready)
  # dlapi.setupvertgraph!(fg, vert, currid, lbl) #fg.v[currid]
  dodims = fg.dimID+1
  setDefaultNodeData!(vert, zeros(dims,N), zeros(0,0), dodims, N, dims, initialized=!autoinit, softtype=st) #fg.v[currid]

  vnlbls = union(string.(labels), st.labels)
  push!(vnlbls, fg.sessionname)
  # addvert!(fg, vert, api=api)
  api.addvertex!(fg, vert, labels=vnlbls) #fg.g ##vertr =

  fg.dimID+=dims # rows indicate dimensions, move to last dimension
  push!(fg.nodeIDs, currid)

  vert
end

# rethink abstraction, maybe closer to CloudGraph use case a better solution
# function addEdge!(g::FGG,n1,n2)
#   edge = dlapi.makeedge(g, n1, n2)
#   dlapi.addedge!(g, edge)
#   # edge = Graphs.make_edge(g, n1, n2)
#   # Graphs.add_edge!(g, edge)
# end

function getVal(vA::Array{Graphs.ExVertex,1})
  len = length(vA)
  vals = Array{Array{Float64,2},1}()
  cols = Array{Int,1}()
  push!(cols,0)
  rows = Array{Int,1}()
  for v in vA
      push!(vals, getVal(v))
      c = size(vals[end],2)
      r = size(vals[end],1)
      push!(cols, floor(Int,c))
      push!(rows, floor(Int,r))
  end
  cols = cumsum(cols)
  sc = cols[end]
  rw = floor(Int,rows[1])
  val = Array{Float64,2}(rw, sc)
  for i in 1:(len-1)
      val[:,(cols[i]+1):cols[i+1]] = vals[i]
  end
  val[:,(cols[len]+1):cols[len+1]] = vals[len] # and the last one
  return val
end

function registerCallback!(fgl::FactorGraph, fnc::Function)
  m = Symbol(typeof(fnc).name.module)
  fgl.registeredModuleFunctions[m] = fnc
  nothing
end

# idea for registering, will likely be removed again
# function createregistercallback!(fnc::Function)
#   fgl = initfg()
#   m = Symbol(typeof(fnc).name.module)
#   fgl.registeredModuleFunctions[m] = fnc
#   nothing
# end
#
function prepareparamsarray!(ARR::Array{Array{Float64,2},1},Xi::Vector{Graphs.ExVertex}, N::Int, solvefor::Int)
  LEN = Int[]
  maxlen = N
  count = 0
  sfidx = 0
  for xi in Xi
    push!(ARR, getVal(xi))
    len = size(ARR[end], 2)
    push!(LEN, len)
    if len > maxlen
      maxlen = len
    end
    count += 1
    if xi.index == solvefor
      sfidx = count #xi.index
    end
  end
  SAMP=LEN.<maxlen
  for i in 1:count
    if SAMP[i]
      ARR[i] = KernelDensityEstimate.sample(getKDE(Xi[i]), maxlen)[1]
    end
  end
  # we are generating a proposal distribution, not direct replacement for existing
  if sfidx > 0 ARR[sfidx] = deepcopy(ARR[sfidx]) end
  return maxlen, sfidx
end

function prepgenericwrapper{T <: FunctorInferenceType}(
      Xi::Vector{Graphs.ExVertex},
      usrfnc::T,
      samplefnc::Function )
  #
  ARR = Array{Array{Float64,2},1}()
  maxlen, sfidx = prepareparamsarray!(ARR, Xi, 0, 0)
  # test if specific zDim or partial constraint used
  fldnms = fieldnames(usrfnc)
  # sum(fldnms .== :zDim) >= 1
  return GenericWrapParam{T}(usrfnc, ARR, 1, 1, (zeros(0,1),), samplefnc, sum(fldnms .== :zDim) >= 1, sum(fldnms .== :partial) >= 1)
end

function setDefaultFactorNode!{T <: Union{FunctorInferenceType, InferenceType}}(
      fgl::FactorGraph,
      vert::Graphs.ExVertex,
      Xi::Vector{Graphs.ExVertex},
      usrfnc::T  )
  #
  ftyp = typeof(usrfnc) # maybe this can be T
  m = Symbol(ftyp.name.module)
  # samplefnc2 = fgl.registeredModuleFunctions[m]
  gwpf = prepgenericwrapper(Xi, usrfnc, getSample) #samplefnc2)

  data = FunctionNodeData{GenericWrapParam{T}}(Int[], false, false, Int[], m, gwpf)
  vert.attributes["data"] = data

  nothing
end

function addNewFncVertInGraph!(fgl::FactorGraph, vert::Graphs.ExVertex, id::Int, lbl::Symbol, ready::Int)
  vert.attributes = Graphs.AttributeDict() #fg.v[fg.id]
  vert.attributes["label"] = lbl #fg.v[fg.id]
  # fgl.f[id] = vert #  -- not sure if this is required, using fg.g.vertices
  fgl.fIDs[lbl] = id # fg.id

  # used for cloudgraph solving
  vert.attributes["ready"] = ready
  vert.attributes["backendset"] = 0

  # for graphviz drawing
  vert.attributes["shape"] = "point"
  vert.attributes["width"] = 0.2
  nothing
end
addNewFncVertInGraph!{T <: AbstractString}(fgl::FactorGraph, vert::Graphs.ExVertex, id::Int, lbl::T, ready::Int) =
    addNewFncVertInGraph!(fgl,vert, id, Symbol(lbl), ready)

function isInitialized(vert::Graphs.ExVertex)::Bool
  return getData(vert).initialized
end
function isInitialized(fgl::FactorGraph, vsym::Symbol)::Bool
  # TODO, make cloudgraphs work and make faster by avoiding all the getVerts
  isInitialized(getVert(fgl, vsym))
end

"""
    doautoinit!(fg, Xi[,api=dlapi])

initialize destination variable nodes based on this factor in factor graph, fg, generally called
during addFactor!. Destination factor is first (singletons) or second (dim 2 pairwise) variable vertex in Xi.
"""
function doautoinit!(fgl::FactorGraph, Xi::Vector{Graphs.ExVertex}; api::DataLayerAPI=dlapi, singles::Bool=true)
  # Mighty inefficient function, since we only need very select fields nearby from a few neighboring nodes
  # do double depth search for variable nodes
  # TODO this should maybe stay localapi only...
  for xi in Xi
    if !isInitialized(xi)
       # @show "doautoinit!", size(getVal(xi))
      vsym = Symbol(xi.label)
      neinodes = ls(fgl, vsym)
      if (length(neinodes) > 1 || singles) # && !isInitialized(xi)
        # Which of the factors can be used for initialization
        useinitfct = Symbol[]
        # println("Consider all pairwise factors connected to $vsym...")
        for xifct in neinodes #potntlfcts
          xfneivarnodes = lsf(fgl, xifct)
          for vsym2 in xfneivarnodes
            # println("find all variables that are initialized for $vsym2")
            vert2 = getVert(fgl, vsym2)
            if (isInitialized(vert2) && sum(useinitfct .== xifct) == 0 ) || length(xfneivarnodes) == 1      # OR singleton  TODO get faster version of isInitialized for database version
              # println("adding $xifct to init factors list")
              push!(useinitfct, xifct)
            end
          end
        end
        # println("Consider all singleton (unary) factors to $vsym...")

        # calculate the predicted belief over $vsym
        pts = predictbelief(fgl, vsym, useinitfct, api=api)
        # println("doautoinit! Past predictbelief...")
        setValKDE!(xi, pts)
        getData(xi).initialized = true
        # println("doautoinit! just before update vertex...")
        api.updatevertex!(fgl, xi, updateMAPest=false)
      end
    end
  end

  # len = length(Xi)
  # # pts = Array{Float64,2}()
  # if length(Xi) == 1
  #   pts = evalFactor2(fgl, fc, Xi[1].index)
  #   setValKDE!(Xi[1], pts)
  #   api.updatevertex!(fgl, Xi[1], updateMAPest=false)
  # elseif len == 2
  #   pts = evalFactor2(fgl, fc, Xi[2].index)
  #   setValKDE!(Xi[2], pts)
  #   api.updatevertex!(fgl, Xi[2], updateMAPest=false)
  # else
  #   # consider specifying an init order in the constraint type
  #   # also consider taking product between all incoming densities which have been inited
  #   error("don't know how to autoinit with pairwise dimension > 2")
  # end
  # println("doautoinit! Done in autoinit!")
  nothing
end

function ensureAllInitialized!(fgl::FactorGraph; api::DataLayerAPI=dlapi)
  xx, xl = ls(fgl)
  allvarnodes = union(xx, xl)
  for vsym in allvarnodes
    if !isInitialized(fgl, vsym)
      println("$vsym is not initialized, and will do so now...")
      doautoinit!(fgl, Graphs.ExVertex[getVert(fgl, vsym, api=api);], api=api, singles=true)
    end
  end
  nothing
end

function assembleFactorName(fgl::FactorGraph, Xi::Vector{Graphs.ExVertex}; maxparallel::Int=50)
  namestring = ""
  for vert in Xi #f.Xi
    namestring = string(namestring,vert.attributes["label"])
  end
  for i in 1:maxparallel
    tempnm = string(namestring, "f$i")
    if !haskey(fgl.fIDs, Symbol(tempnm))
      namestring = tempnm
      break
    end
    i != maxparallel ? nothing : error("Cannot currently add more than $(maxparallel) factors in parallel, please open an issue if this is too restrictive.")
  end
  return namestring
end

function addFactor!(fgl::FactorGraph,
      Xi::Vector{Graphs.ExVertex},
      usrfnc::R;
      ready::Int=1,
      api::DataLayerAPI=dlapi,
      labels::Vector{T}=String[],
      uid::Int=-1,
      autoinit::Bool=true) where
        {R <: Union{FunctorInferenceType, InferenceType},
         T <: AbstractString}
  #
  currid = fgl.id+1
  if uid==-1
    fgl.id=currid
  else
    currid = uid
  end
  namestring = assembleFactorName(fgl, Xi)
  # fgl.id+=1
  newvert = ExVertex(currid,namestring)
  addNewFncVertInGraph!(fgl, newvert, currid, namestring, ready)
  setDefaultFactorNode!(fgl, newvert, Xi, deepcopy(usrfnc))
  push!(fgl.factorIDs,currid)

  for vert in Xi
    push!(newvert.attributes["data"].fncargvID, vert.index)
  end

  fnlbls = deepcopy(labels)
  push!(fnlbls, "FACTOR")
  push!(fnlbls, fgl.sessionname)
  # TODO -- multiple accesses to DB with this method, must refactor!
  newvert = api.addvertex!(fgl, newvert, labels=fnlbls)  # used to be two be three lines up ##fgl.g
  for vert in Xi
    api.makeaddedge!(fgl, vert, newvert)
  end

  # TODO change this operation to update a conditioning variable
  if autoinit
    doautoinit!(fgl, Xi, api=api, singles=false)
  end

  return newvert
end

function addFactor!(
      fgl::FactorGraph,
      xisyms::Vector{Symbol},
      usrfnc::R;
      ready::Int=1,
      api::DataLayerAPI=dlapi,
      labels::Vector{T}=String[],
      uid::Int=-1,
      autoinit::Bool=true ) where
        {R <: Union{FunctorInferenceType, InferenceType},
         T <: AbstractString}
  #
  verts = Vector{Graphs.ExVertex}()
  for xi in xisyms
      push!( verts, api.getvertex(fgl,xi) )
  end
  addFactor!(fgl, verts, usrfnc, ready=ready, api=api, labels=labels, uid=uid, autoinit=autoinit)
end

function prtslperr(s)
  println(s)
  sleep(0.1)
  error(s)
end

# for computing the Bayes Net-----------------------------------------------------
function getEliminationOrder(fg::FactorGraph; ordering::Symbol=:qr)
    s = fg.nodeIDs
    lens = length(s)
    sf = fg.factorIDs
    lensf = length(sf)
    adjm, dictpermu = adjacency_matrix(fg.g,returnpermutation=true)
    permuteds = Vector{Int}(lens)
    permutedsf = Vector{Int}(lensf)
    for j in 1:length(dictpermu)
      semap = 0
      for i in 1:lens
        if dictpermu[j] == s[i]
          permuteds[i] = j#dictpermu[j]
          semap += 1
          if semap >= 2  break; end
        end
      end
      for i in 1:lensf
        if dictpermu[j] == sf[i]
          permutedsf[i] = j#dictpermu[j]
          semap += 1
          if semap >= 2  break; end
        end
      end
    end

      A=convert(Array{Int},adjm[permutedsf,permuteds]) # TODO -- order seems brittle
      p = Int[]
      if ordering==:chol
        p = cholfact(A'A,:U,Val{true})[:p] #,pivot=true
      elseif ordering==:qr
        q,r,p = qr(A,Val{true})
      else
        prtslperr("getEliminationOrder -- cannot do the requested ordering $(ordering)")
      end

    # we need the IDs associated with the Graphs.jl and our Type fg
    return  dictpermu[permuteds[p]] # fg.nodeIDs[p]
end

# lets create all the vertices first and then deal with the elimination variables thereafter
function addBayesNetVerts!(fg::FactorGraph, elimOrder::Array{Int,1})
  for p in elimOrder
    vert = getVert(fg, p, api=localapi)
    @show vert.label, getData(vert).BayesNetVertID
    if getData(vert).BayesNetVertID == 0
      fg.bnid+=1
      vert.attributes["data"].BayesNetVertID = p
      localapi.updatevertex!(fg, vert)
    else
      println("addBayesNetVerts -- something is very wrong, should not have a Bayes net vertex")
    end
  end
end

function addConditional!(fg::FactorGraph, vertID::Int, lbl, Si)
  bnv = getVert(fg, vertID, api=localapi) #fg.v[vertID]
  bnvd = getData(bnv) # bnv.attributes["data"]
  bnvd.separator = Si
  for s in Si
    push!(bnvd.BayesNetOutVertIDs, s)
  end
  localapi.updatevertex!(fg, bnv)
  nothing
end

function addChainRuleMarginal!(fg::FactorGraph, Si)
  lbls = String[]
  genmarg = GenericMarginal()
  Xi = Graphs.ExVertex[]
  for s in Si
    push!(Xi, getVert(fg, s, api=localapi))
  end
  println("adding marginal to")
  for x in Xi @show x.index end
  addFactor!(fg, Xi, genmarg, api=localapi)
  nothing
end

# TODO -- Cannot have any CloudGraph calls at this level, must refactor
function rmVarFromMarg(fgl::FactorGraph, fromvert::Graphs.ExVertex, gm::Array{Graphs.ExVertex,1})
  for m in gm
    # get all out edges
    # get neighbors
    for n in localapi.outneighbors(fgl, m)
      if n.index == fromvert.index
        alleids = m.attributes["data"].edgeIDs
        i = 0
        for id in alleids
          i+=1
          edge = localapi.getedge(fgl, id)
          if edge != nothing # hack to avoid dictionary use case
            if edge.SourceVertex.exVertexId == m.index || edge.DestVertex.exVertexId == m.index
              warn("removing edge $(edge.neo4jEdgeId), between $(m.index) and $(n.index)")
              localapi.deleteedge!(fgl, edge)
              m.attributes["data"].edgeIDs = alleids[[collect(1:(i-1));collect((i+1):length(alleids))]]
              localapi.updatevertex!(fgl, m)
            end
          end
        end
      end
    end
    # if 0 edges, delete the marginal
    if length(localapi.outneighbors(fgl, m)) <= 1
      warn("removing vertex id=$(m.index)")
      localapi.deletevertex!(fgl,m)
    end
  end
  nothing
end

function buildBayesNet!(fg::FactorGraph, p::Array{Int,1})
    addBayesNetVerts!(fg, p)
    for v in p
      println()
      println("Eliminating $(v)")
      println("===============")
      println()
      # which variable are we eliminating

      # all factors adjacent to this variable
      fi = Int[]
      Si = Int[]
      gm = ExVertex[]
      # TODO -- optimize outneighbor calls like this
      vert = localapi.getvertex(fg,v)
      for fct in localapi.outneighbors(fg, vert)
        if (getData(fct).eliminated != true)
          push!(fi, fct.index)
          for sepNode in localapi.outneighbors(fg, fct)
            if sepNode.index != v && length(findin(sepNode.index,Si)) == 0
              push!(Si,sepNode.index)
            end
          end
          getData(fct).eliminated = true #fct.attributes["data"].eliminated = true
          localapi.updatevertex!(fg, fct) # TODO -- this might be a premature statement
        end

        if typeof(getData(fct).fnc) == GenericMarginal
          push!(gm, fct)
        end
      end

      if v != p[end]
        addConditional!(fg, v, "", Si)
        # not yet inserting the new prior p(Si) back into the factor graph
      end

      tuv = localapi.getvertex(fg, v) # TODO -- This may well through away valuable data
      tuv.attributes["data"].eliminated = true # fg.v[v].
      localapi.updatevertex!(fg, tuv)

      # TODO -- remove links from current vertex to any marginals
      rmVarFromMarg(fg, vert, gm)

      #add marginal on remaining variables... ? f(xyz) = f(x | yz) f(yz)
      # new function between all Si
      addChainRuleMarginal!(fg, Si)

    end
    nothing
end

# some plotting functions on the factor graph
function stackVertXY(fg::FactorGraph, lbl::String)
    v = dlapi.getvertex(fg,lbl)
    vals = getVal(v)
    X=vec(vals[1,:])
    Y=vec(vals[2,:])
    return X,Y
end

function getKDE(v::Graphs.ExVertex)
  return kde!(getVal(v), getBWVal(v)[:,1])
end

function getVertKDE(v::Graphs.ExVertex)
  return getKDE(v)
end

function getVertKDE(fgl::FactorGraph, id::Int; api::DataLayerAPI=dlapi)
  v = api.getvertex(fgl,id)
  return getKDE(v)
end
function getVertKDE(fgl::FactorGraph, lbl::Symbol; api::DataLayerAPI=dlapi)
  v = api.getvertex(fgl,lbl)
  return getKDE(v)
end

function drawCopyFG(fgl::FactorGraph)
  fgd = deepcopy(fgl)
  for (sym,vid) in fgd.IDs
    delete!(fgd.g.vertices[vid].attributes,"data")
    !haskey(fgd.g.vertices[vid].attributes,"frtend") ? nothing : delete!(fgd.g.vertices[vid].attributes,"frtend")
  end
  for (sym,vid) in fgd.fIDs
    delete!(fgd.g.vertices[vid].attributes,"data")
    !haskey(fgd.g.vertices[vid].attributes,"frtend") ? nothing : delete!(fgd.g.vertices[vid].attributes,"frtend")
  end
  return fgd
end

# function writeGraphPdf(fgl::FactorGraph)
#   fgd = drawCopyFG(fgl)
#   println("Writing factor graph file")
#   fid = open("fg.dot","w+")
#   write(fid,Graphs.to_dot(fgd.g))
#   close(fid)
#   run(`dot fg.dot -Tpdf -o fg.pdf`)
#   nothing
# end

function expandEdgeListNeigh!(fgl::FactorGraph,
                              vertdict::Dict{Int,Graphs.ExVertex},
                              edgedict::Dict{Int,Graphs.Edge{Graphs.ExVertex}})
  #asfd
  for vert in vertdict
    for newedge in out_edges(vert[2],fgl.g)
      if !haskey(edgedict, newedge.index)
        edgedict[newedge.index] = newedge
      end
    end
  end

  nothing
end

# dictionary of unique vertices from edgelist
function expandVertexList!(fgl::FactorGraph,
  edgedict::Dict{Int,Graphs.Edge{Graphs.ExVertex}},
  vertdict::Dict{Int,Graphs.ExVertex})

  # go through all source and target nodes
  for edge in edgedict
    if !haskey(vertdict, edge[2].source.index)
      vertdict[edge[2].source.index] = edge[2].source
    end
    if !haskey(vertdict, edge[2].target.index)
      vertdict[edge[2].target.index] = edge[2].target
    end
  end
  nothing
end

function edgelist2edgedict(edgelist::Array{Graphs.Edge{Graphs.ExVertex},1})
  edgedict = Dict{Int,Graphs.Edge{Graphs.ExVertex}}()
  for edge in edgelist
    edgedict[edge.index] = edge
  end
  return edgedict
end

# TODO -- convert to use add_vertex! instead, since edges type must be made also
function addVerticesSubgraph(fgl::FactorGraph,
    fgseg::FactorGraph,
    vertdict::Dict{Int,Graphs.ExVertex})

    for vert in vertdict
      fgseg.g.vertices[vert[1]] = vert[2]
      if haskey(fgl.v,vert[1])
        fgseg.g.vertices[vert[1]] = vert[2]
        fgseg.IDs[Symbol(vert[2].label)] = vert[1]

        # add edges going in opposite direction
        elr = Graphs.out_edges(vert[2], fgl.g)
        len = length(elr)
        keeprm = trues(len)
        j = 0
        for i in 1:len
          if !haskey(vertdict, elr[i].target.index) # a function node in set, so keep ref
            keeprm[i] = false
            j+=1
          end
        end
        if j < len
          elridx = elr[1].source.index
          fgseg.g.inclist[elridx] = elr[keeprm]
        end
      elseif haskey(fgl.f, vert[1])
        fgseg.f[vert[1]] = vert[2] # adding element to subgraph
        fgseg.fIDs[Symbol(vert[2].label)] = vert[1]
        # get edges associated with function nodes and push edges onto incidence list
        el = Graphs.out_edges(vert[2], fgl.g)
        elidx = el[1].source.index
        fgseg.g.inclist[elidx] = el # okay because treating function nodes only
        fgseg.g.nedges += length(el)
      else
        error("Unknown type factor graph vertex type, something is wrong")
      end
    end
    nothing
end

# NOTICE, nodeIDs and factorIDs are not brough over by this method yet
# must sort out for incremental updates
function genSubgraph(fgl::FactorGraph,
    vertdict::Dict{Int,Graphs.ExVertex})
    # edgedict::Dict{Int,Graphs.Edge{Graphs.ExVertex}},

  fgseg = FactorGraph() # new handle for just a segment of the graph
  fgseg.g = Graphs.inclist(Graphs.ExVertex,is_directed=false)

  fgseg.v = Dict{Int,Graphs.ExVertex}()
  fgseg.f = Dict{Int,Graphs.ExVertex}()
  fgseg.IDs = Dict{AbstractString,Int}()
  fgseg.fIDs = Dict{AbstractString,Int}()

  # TODO -- convert to use empty constructor since Graphs.incdict now works
  fgseg.g.vertices = Array{Graphs.ExVertex,1}(length(fgl.g.vertices))
  fgseg.g.inclist = Array{Array{Graphs.Edge{Graphs.ExVertex},1},1}(length(fgl.g.inclist))

  addVerticesSubgraph(fgl, fgseg, vertdict)

  fgseg.id = fgl.id
  fgseg.bnid = fgl.bnid
  fgseg.dimID = fgl.dimID

  return fgseg
end

function getShortestPathNeighbors(fgl::FactorGraph;
    from::Graphs.ExVertex=nothing,
    to::Graphs.ExVertex=nothing,
    neighbors::Int=0 )

  edgelist = shortest_path(fgl.g, ones(num_edges(fgl.g)), from, to)
  vertdict = Dict{Int,Graphs.ExVertex}()
  edgedict = edgelist2edgedict(edgelist)
  expandVertexList!(fgl, edgedict, vertdict) # grow verts
  for i in 1:neighbors
    expandEdgeListNeigh!(fgl, vertdict, edgedict) # grow edges
    expandVertexList!(fgl, edgedict, vertdict) # grow verts
  end
  return vertdict
end

function subgraphShortestPath(fgl::FactorGraph;
    from::Graphs.ExVertex=nothing,
    to::Graphs.ExVertex=nothing,
    neighbors::Int=0  )

  vertdict = getShortestPathNeighbors(fgl, from=from, to=to, neighbors=neighbors)
  return genSubgraph(fgl, vertdict)
end

# explore all shortest paths combinations in verts, add neighbors and reference subgraph
function subgraphFromVerts(fgl::FactorGraph,
    verts::Dict{Int,Graphs.ExVertex};
    neighbors::Int=0  )

    allverts = Dict{Int,Graphs.ExVertex}()
    allkeys = collect(keys(verts))
    len = length(allkeys)
    # union all shortest path combinations in a vertdict
    for i in 1:len, j in (i+1):len
      from = verts[allkeys[i]]
      to = verts[allkeys[j]]
      vertdict = getShortestPathNeighbors(fgl, from=from, to=to, neighbors=neighbors)
      for vert in vertdict
        if !haskey(allverts, vert[1])
          allverts[vert[1]] = vert[2]
        end
      end
    end

  return genSubgraph(fgl, allverts)
end

# explore all shortest paths combinations in verts, add neighbors and reference subgraph
# Using unique index into graph data structure
function subgraphFromVerts(fgl::FactorGraph,
    verts::Array{Int,1};
    neighbors::Int=0  )

  vertdict = Dict{Int,Graphs.ExVertex}()
  for vert in verts
    vertdict[vert] = fgl.g.vertices[vert]
  end

  return subgraphFromVerts(fgl,vertdict,neighbors=neighbors)
end

# explore all shortest paths combinations in verts, add neighbors and reference subgraph
# Using unique index into graph data structure
function subgraphFromVerts(fgl::FactorGraph,
    verts::Array{String,1};
    neighbors::Int=0  )

  vertdict = Dict{Int,Graphs.ExVertex}()
  for vert in verts
    id = -1
    if haskey(fgl.IDs, vert)
      id = fgl.IDs[Symbol(vert)]
    else
      error("FactorGraph01 only supports variable node subgraph search at this point")
    end
    vertdict[id] = fgl.g.vertices[id]
  end

  return subgraphFromVerts(fgl,vertdict,neighbors=neighbors)
end


text:
### More advanced data structures

struct FaceBasedDS
    faces::AbstractArray{Int,2}
    neighs::AbstractArray{Int,2} ### For now keeping simple
    vfaces::AbstractArray{Int,1}
end

function FaceBasedDS(faces::Array{Int,2})

    vfaces = Array{Int}(undef,maximum(faces))
    neighs = Array{Int}(undef,size(faces)...)

    for vi in 1:maximum(faces)
        vfaces[vi] = find_triangle_vertex(vi,faces)
    end

    for ti in 1:size(faces,2)
        v1, v2, v3 = faces[:,ti]
        t1 = find_other_triangle_edge(v2,v3,ti,faces)
        t2 = find_other_triangle_edge(v1,v3,ti,faces)
        t3 = find_other_triangle_edge(v1,v2,ti,faces)
        neighs[:,ti] = [t1,t2,t3]
    end

    return FaceBasedDS(faces,neighs,vfaces)
end

"""
Iterator for getting all triangles around vertex
"""
struct FaceVRingFB
    v::Int
    faces::Array{Int,2}
    neighs::Array{Int,2} ### For now keeping simple
    vfaces::Array{Int,1}
end

FaceVRing(v::Int,ds::FaceBasedDS) = FaceVRingFB(v,ds.faces,ds.neighs,ds.vfaces)

function start(iter::FaceVRingFB)
    vface = iter.vfaces[iter.v]
    i0 = 1
    return (i0,vface)
end

function done(iter::FaceVRingFB,state::Tuple{Int,Int})
    i, face = state
    i0, vface = start(iter)
    if !(i==i0) & (face==vface)
        return true
    else
        return false
    end
end

function next(iter::FaceVRingFB,state::Tuple{Int,Int})

    i, tri = state
    v = iter.v
    face = iter.faces[:,tri]
    neighbours = iter.neighs[:,tri]

    index = face.==v
    w = index[[1,2,3]]
    cw = index[[3,1,2]]

    nexttri, = neighbours[cw]

    if nexttri==-1
        error("The surface is not closed")
    end

    return tri, (i+1,nexttri)
end

Base.iterate(iter::FaceVRingFB) = next(iter,start(iter))
function Base.iterate(iter::FaceVRingFB,ti)
    if done(iter,ti)
        return nothing
    else
        return next(iter,ti)
    end
end

###

struct DoubleVertexVRingFB
    v::Int
    faces::Array{Int,2}
    neighs::Array{Int,2} ### For now keeping simple
    vfaces::Array{Int,1}
end

DoubleVertexVRing(v::Int,ds::FaceBasedDS) = DoubleVertexVRingFB(v,ds.faces,ds.neighs,ds.vfaces)

function start(iter::DoubleVertexVRingFB)
    vface = iter.vfaces[iter.v]
    i0 = 1
    return (i0,vface)
end

function done(iter::DoubleVertexVRingFB,state::Tuple{Int,Int})
    i, face = state
    i0, vface = start(iter)
    if !(i==i0) & (face==vface)
        return true
    else
        return false
    end
end

function next(iter::DoubleVertexVRingFB,state::Tuple{Int,Int})

    i, tri = state
    v = iter.v
    face = iter.faces[:,tri]
    neighbours = iter.neighs[:,tri]

    index = face.==v
    w = index[[1,2,3]]
    cw = index[[3,1,2]]
    ccw = index[[2,3,1]]

    nexttri, = neighbours[cw]

    if nexttri==-1
        error("The surface is not closed")
    end

    return (face[cw]...,face[ccw]...), (i+1,nexttri)
end

Base.iterate(iter::DoubleVertexVRingFB) = next(iter,start(iter))
function Base.iterate(iter::DoubleVertexVRingFB,ti)
    if done(iter,ti)
        return nothing
    else
        return next(iter,ti)
    end
end

"""
Now the magic of making VertexIterator
"""
struct VertexVRingFB
    v::Int
    faces::Array{Int,2}
    neighs::Array{Int,2} ### For now keeping simple
    vface::Int
    # function VertexRingFBDS(v,faces,neighs,vfaces)
    #     new(v,faces,neighs,vfaces[v])
    # end
end

VertexVRing(v::Int,ds::FaceBasedDS) = VertexVRingFB(v,ds.faces,ds.neighs,ds.vfaces[v])

function start(iter::VertexVRingFB)
    vface = iter.vface
    i0 = 1
    return (i0,vface)
end

function done(iter::VertexVRingFB,state::Tuple{Int,Int})
    i, face = state
    if !(i==1) & (face==iter.vface)
        return true
    else
        return false
    end
end

function next(iter::VertexVRingFB,state::Tuple{Int,Int})

    i, tri = state
    v = iter.v
    face = iter.faces[:,tri]
    neighbours = iter.neighs[:,tri]

    index = face.==v
    w = index[[1,2,3]]
    cw = index[[3,1,2]]

    nexttri, = neighbours[cw]

    if nexttri==-1
        error("The surface is not closed")
    end

    ### Code for extracting vertex from face tri

    face = iter.faces[:,tri]
    cw = (face.==v)[[3,1,2]]
    vi, = face[cw]

    return vi, (i+1,nexttri)
end

Base.iterate(iter::VertexVRingFB) = next(iter,start(iter))
function Base.iterate(iter::VertexVRingFB,ti)
    if done(iter,ti)
        return nothing
    else
        return next(iter,ti)
    end
end

### Using conectivity for improving performance

struct ConnectivityDS
    connectivity
end

function ConnectivityTable(faces,valence)
    vmax = maximum(faces)
    connectivity = zeros(Int,valence,vmax)
    for i in 1:vmax
        v1 = Array{Int}(undef,0)
        v2 = Array{Int}(undef,0)
        for (v1i,v2i) in DoubleVertexVRing(i,faces)
            push!(v1,v1i)
            push!(v2,v2i)
        end

        vj = v2[1]
        connectivity[1,i] = vj
        for j in 2:size(v1,1)
            w = (v1.==vj)
            vj, = v2[w]
            connectivity[j,i] = vj
        end
    end
    return connectivity
end

ConnectivityDS(faces,valence) = ConnectivityDS(ConnectivityTable(faces,valence))

struct VertexVRingCon{T<:Integer}
    v::T
    connectivity::Array{T,2}
end

VertexVRing(v::Int,ds::ConnectivityDS) = VertexVRingCon(v,ds.connectivity)

function start(iter::VertexVRingCon)
    return 1
end

function done(iter::VertexVRingCon,i::T) where {T<:Integer}
    return (i>size(iter.connectivity,1)) || (iter.connectivity[i,iter.v]==0) ? true : false
end

function next(iter::VertexVRingCon,i::T) where {T<:Integer}
    return iter.connectivity[i,iter.v], i+1
end

Base.iterate(iter::VertexVRingCon) = next(iter,start(iter))
function Base.iterate(iter::VertexVRingCon,ti)
    if done(iter,ti)
        return nothing
    else
        return next(iter,ti)
    end
end

struct DoubleVertexVRingCon{T<:Integer}
    v::T
    connectivity::Array{T,2}
end

DoubleVertexVRing(v::Int,ds::ConnectivityDS) = DoubleVertexVRingCon(v,ds.connectivity)

function start(iter::DoubleVertexVRingCon)
    return 1
end

function done(iter::DoubleVertexVRingCon,i::T) where {T<:Integer}
    return (i>size(iter.connectivity,1)) || (iter.connectivity[i,iter.v]==0) ? true : false
end

function next(iter::DoubleVertexVRingCon,i::T) where {T<:Integer}
    v1 = iter.connectivity[i,iter.v]
    v2 = done(iter,i+1) ? iter.connectivity[1,iter.v] : iter.connectivity[i+1,iter.v] # (i==size(iter.connectivity,1)) || (iter.connectivity[i+1,iter.v]==0)
    return (v1,v2), i+1
end

Base.iterate(iter::DoubleVertexVRingCon) = next(iter,start(iter))
function Base.iterate(iter::DoubleVertexVRingCon,ti)
    if done(iter,ti)
        return nothing
    else
        return next(iter,ti)
    end
end


text:
using ArgParse,CairoMakie

CairoMakie.activate!(type="pdf")

function parse_argument()

end

function binominal_prob_desity(n,p,x)
    return binomial(n,x)*p^x*(1-p)^(n-x)
end

function main()
    graph_name = "binominal_prob_mass"
    xs_max = 50
    xs = 0:xs_max
    fig = Figure(resolution = (800, 600))

    axis = Axis(fig[1,1], title = "binominal probability mass function")

    graph_params = [Dict("prob"=>0.1, "color"=>:gray0,  "marker"=>:circle,  "label"=>"p=0.1"),
                    Dict("prob"=>0.5, "color"=>:gray30, "marker"=>:diamond, "label"=>"p=0.5"),
                    Dict("prob"=>0.7, "color"=>:gray50, "marker"=>:rect,    "label"=>"p=0.7")]

    for params in graph_params
        p_s = [ binominal_prob_desity(xs_max,params["prob"],x) for x in xs ]
        stem!(fig[1,1], xs, p_s, color = params["color"], marker = params["marker"], label=params["label"])
    end

    axislegend(axis)
    save("$graph_name.pdf", fig)

end

main()

max_stars_repo_name: seisyuu-hantatsushi/julia

max_issues_repo_name: seisyuu-hantatsushi/julia

max_forks_repo_name: seisyuu-hantatsushi/julia


text:
"""
Return default marking value based on `PNTD`. Has meaning of empty, as in `zero`.
"""
function default_term end
default_term(::PNTD) where {PNTD <: PnmlType} = zero(Integer) #!
default_term(::PNTD) where {PNTD <: AbstractContinuousCore} = zero(Float64) #!
default_term(::PNTD) where {PNTD <: AbstractHLCore} = Term() #!

"""
$(TYPEDEF)
$(TYPEDFIELDS)

Part of the many-sorted algebra attached to nodes on a Petri Net Graph.

 ast variants:
  - variable
  - operator
"""
struct Term{T<:AbstractDict}  <: AbstractTerm #TODO make mutable?
  tag::Symbol
  dict::T
  #TODO xml
end

Term() = Term(:empty, PnmlDict())

Term(p::Pair{Symbol,PnmlDict}; kw...) = Term(p.first, p.second)

convert(::Type{Maybe{Term}}, pdict::PnmlDict) = Term(pdict)

tag(t::Term) = t.tag
dict(t::Term) = t.dict
#xml(t::Term) = t.xml

"""
Evaluate a term by returning the ':value' in `dict`.
Assumes that `dict` is an `AbstractDictionary`.
"""
(t::Term)() = get(t.dict, :value, 1)


text:
module CellValuesGallery

using Gridap
using Gridap.Helpers
using Gridap.CachedSubVectors
using Gridap.CachedArrays

using Base: @propagate_inbounds, @pure

export CellValueFromArray
export CellArrayFromArrayOfArrays
export CellVectorFromDataAndPtrs
export CellVectorFromDataAndStride
export CellVectorFromLocalToGlobal
export CellVectorFromLocalToGlobalPosAndNeg
export CellVectorByComposition
import Base: iterate
import Base: length
import Base: eltype
import Base: size
import Base: getindex
import Base: IndexStyle
import Gridap: compress

struct CellValueFromArray{T,N,V<:AbstractArray{T,N}} <: IndexCellValue{T,N}
  v::V
end

@propagate_inbounds function getindex(
  self::CellValueFromArray{T,N},I::Vararg{Int,N}) where {T,N}
  @inbounds self.v[I...]
end

size(self::CellValueFromArray) = size(self.v)

@pure IndexStyle(::Type{CellValueFromArray{T,N,V}}) where {T,N,V} = IndexStyle(V)

struct CellVectorFromDataAndPtrs{T,V,P} <: IndexCellArray{T,1,CachedSubVector{T,V},1}
  data::V
  ptrs::P
  cv::CachedSubVector{T,V}
end

function CellVectorFromDataAndPtrs(data::AbstractArray{T,1},ptrs::AbstractArray{Int,1}) where T
  V = typeof(data)
  P = typeof(ptrs)
  cv = CachedSubVector(data,0,0)
  CellVectorFromDataAndPtrs{T,V,P}(data,ptrs,cv)
end

@propagate_inbounds function getindex(self::CellVectorFromDataAndPtrs,cell::Int)
  pini = self.ptrs[cell]
  pend = self.ptrs[cell+1]-1
  locate!(self.cv,pini,pend)
  self.cv
end

size(self::CellVectorFromDataAndPtrs) = (length(self.ptrs)-1,)

IndexStyle(::Type{CellVectorFromDataAndPtrs{T,V,P}}) where {T,V,P} = IndexLinear()

function compress(cv::CellVectorFromDataAndPtrs)
  (cv.data, cv.ptrs)
end

struct CellVectorFromDataAndStride{T,V} <: IndexCellArray{T,1,CachedSubVector{T,V},1}
  data::V
  stride::Int
  cv::CachedSubVector{T,V}
end

function CellVectorFromDataAndStride(data::AbstractArray{T,1},stride::Int) where T
  V = typeof(data)
  cv = CachedSubVector(data,0,0)
  CellVectorFromDataAndStride{T,V}(data,stride,cv)
end

@propagate_inbounds function getindex(self::CellVectorFromDataAndStride,cell::Int)
  pini = self.stride*(cell - 1) + 1
  pend = self.stride*cell
  locate!(self.cv,pini,pend)
  self.cv
end

size(self::CellVectorFromDataAndStride) = (ceil(Int,length(self.data)/self.stride),)

IndexStyle(::Type{CellVectorFromDataAndStride{T,V}}) where {T,V} = IndexLinear()

struct CellVectorFromLocalToGlobal{
  T,D,L<:IndexCellArray{Int,1},V<:IndexCellValue{T}} <: IndexCellArray{T,1,CachedArray{T,1,Array{T,1}},D}
  lid_to_gid::L
  gid_to_val::V
  cv::CachedArray{T,1,Array{T,1}}
end

function CellVectorFromLocalToGlobal(
  lid_to_gid::IndexCellArray{Int,1,A,D},
  gid_to_val::IndexCellValue{T}) where {T,A,D}
  L = typeof(lid_to_gid)
  V = typeof(gid_to_val)
  cv = CachedArray(T,1)
  CellVectorFromLocalToGlobal{T,D,L,V}(lid_to_gid,gid_to_val,cv)
end

function CellVectorFromLocalToGlobal(
  lid_to_gid::IndexCellArray{Int,1,A,D},
  gid_to_val::AbstractArray{T}) where {T,A,D}
  _gid_to_val = CellValueFromArray(gid_to_val)
  CellVectorFromLocalToGlobal(lid_to_gid,_gid_to_val)
end

@propagate_inbounds function getindex(self::CellVectorFromLocalToGlobal,cell::Vararg{Int,N} where N)
  lid_to_gid = self.lid_to_gid[cell...]
  setsize!(self.cv,(length(lid_to_gid),))
  for (lid,gid) in enumerate(lid_to_gid)
    self.cv[lid] = self.gid_to_val[gid]
  end
  self.cv
end

size(self::CellVectorFromLocalToGlobal) = size(self.lid_to_gid)

IndexStyle(::Type{CellVectorFromLocalToGlobal{T,D,L,V}}) where {T,D,L,V} = IndexStyle(L)

struct CellVectorFromLocalToGlobalPosAndNeg{
  T,D,L<:IndexCellArray{Int,1},V<:IndexCellValue{T},W<:IndexCellValue{T}} <: IndexCellArray{T,1,CachedArray{T,1,Array{T,1}},D}
  lid_to_gid::L
  gid_to_val_pos::V
  gid_to_val_neg::W
  cv::CachedArray{T,1,Array{T,1}}
end

function CellVectorFromLocalToGlobalPosAndNeg(
  lid_to_gid::IndexCellArray{Int,1,A,D},
  gid_to_val_pos::IndexCellValue{T},
  gid_to_val_neg::IndexCellValue{T}) where {T,A,D}
  L = typeof(lid_to_gid)
  V = typeof(gid_to_val_pos)
  W = typeof(gid_to_val_neg)
  cv = CachedArray(T,1)
  CellVectorFromLocalToGlobalPosAndNeg{T,D,L,V,W}(lid_to_gid,gid_to_val_pos,gid_to_val_neg,cv)
end

function CellVectorFromLocalToGlobalPosAndNeg(
  lid_to_gid::IndexCellArray{Int,1,A,D},
  gid_to_val_pos::AbstractArray{T},
  gid_to_val_neg::AbstractArray{T}) where {T,A,D}
  _gid_to_val_pos = CellValueFromArray(gid_to_val_pos)
  _gid_to_val_neg = CellValueFromArray(gid_to_val_neg)
  CellVectorFromLocalToGlobalPosAndNeg(lid_to_gid,_gid_to_val_pos,_gid_to_val_neg)
end

IndexStyle(::Type{CellVectorFromLocalToGlobalPosAndNeg{T,D,L,V,W}}) where {T,D,L,V,W} = IndexStyle(L)

@propagate_inbounds function getindex(
  self::CellVectorFromLocalToGlobalPosAndNeg,cell::Vararg{Int,N} where N)
  lid_to_gid = self.lid_to_gid[cell...]
  setsize!(self.cv,(length(lid_to_gid),))
  for (lid,gid) in enumerate(lid_to_gid)
    if gid > 0
      self.cv[lid] = self.gid_to_val_pos[gid]
    elseif gid < 0
      self.cv[lid] = self.gid_to_val_neg[-gid]
    else
      @unreachable
    end
  end
  self.cv
end

size(self::CellVectorFromLocalToGlobalPosAndNeg) = size(self.lid_to_gid)

struct CellVectorByComposition{
  T,L<:IndexCellVector{Int},V<:IndexCellVector{T}} <: IndexCellVector{T,CachedVector{T,Vector{T}},1}

  cell_to_x::L
  x_to_vals::V
  cv::CachedVector{T,Vector{T}}
end

function CellVectorByComposition(
  cell_to_x::IndexCellVector{Int}, x_to_vals::IndexCellVector{T}) where T
  L = typeof(cell_to_x)
  V = typeof(x_to_vals)
  cv = CachedArray(T,1)
  CellVectorByComposition{T,L,V}(cell_to_x, x_to_vals, cv)
end

@propagate_inbounds function getindex(self::CellVectorByComposition,cell::Int)
  cell_to_x = self.cell_to_x[cell]
  l = 0
  for x in cell_to_x
    l += length(self.x_to_vals[x])
  end
  setsize!(self.cv,(l,))
  l = 1
  for x in cell_to_x
    for val in self.x_to_vals[x]
      for iv in val
        self.cv[l] = iv
        l += 1
      end
    end
  end
  self.cv
end

size(self::CellVectorByComposition) = (length(self.cell_to_x),)

IndexStyle(::Type{CellVectorByComposition{T,L,V}}) where {T,L,V} = IndexLinear()

end # module


text:
using Oceananigans.BoundaryConditions: fill_halo_regions!
using Oceananigans.ImmersedBoundaries: mask_immersed_field!

import Oceananigans.TimeSteppers: update_state!

"""
    update_state!(model::ShallowWaterModel)

Fill halo regions for `model.solution` and `model.tracers`.
"""
function update_state!(model::ShallowWaterModel)

    # Mask immersed fields
    masking_events = Tuple(mask_immersed_field!(field) for field in merge(model.solution, model.tracers))

    wait(device(model.architecture), MultiEvent(masking_events))

    # Fill halos for velocities and tracers
    fill_halo_regions!(merge(model.solution, model.tracers),
                       model.architecture,
                       model.clock,
                       fields(model))

    return nothing
end


text:
#--------------------------------------------------------------------
# Spacetime Discretization methods in Julia
# Soham 08-2018
# Define operations for ND spaces
#--------------------------------------------------------------------

function Base. reshape(u::Field{S}) where {S}
    return reshape(u.value, (prod(size(u.space))))
end

function Base. reshape(A::Operator{S}) where {S}
    return reshape(A.value, (prod(size(A.space)), prod(size(A.space))))
end

function Base. reshape(space::S, u::AbstractArray{T,1})::Field{S} where {T,S}
    return Field(space, reshape(u, size(space)))
end

function Base. reshape(space::S, A::AbstractArray{T,2})::Operator{S} where {T,S}
    return Operator(space, reshape(A, (size(space)..., size(space)...)))
end

function Base. *(A::Operator{S}, u::Field{S})::Field{S} where {S}
    @assert range(A.space) == range(u.space)
    return reshape(A.space, reshape(A)*reshape(u))
end

function Base. *(A::Operator{S}, B::Operator{S})::Operator{S} where {S}
    @assert range(A.space) == range(B.space)
    return reshape(A.space, reshape(A)*reshape(B))
end

function Base. +(A::Operator{S}, B::Operator{S})::Operator{S} where {S}
    @assert range(A.space) == range(B.space)
    return reshape(A.space, reshape(A) + reshape(B))
end

function Base. -(A::Operator{S}, B::Operator{S})::Operator{S} where {S}
    @assert range(A.space) == range(B.space)
    return reshape(A.space, reshape(A) - reshape(B))
end

function Base. *(u::Field{S}, B::Operator{S})::Operator{S} where {S}
    @assert range(u.space) == range(B.space)
    return reshape(u.space, Diagonal(reshape(u))*reshape(B))
end

function Base. *(a::Number, B::Operator{S})::Operator{S} where {S}
    return reshape(B.space, eltype(reshape(B))(a).*reshape(B))
end

function Base. isapprox(u::Field{S}, v::Field{S})::Bool where {S}
    @assert range(u.space) == range(v.space)
    return u.value ≈ v.value
end

function Base. *(u::Field{S}, v::Field{S})::Field{S} where {S}
    @assert range(u.space) == range(v.space)
    return reshape(u.space, reshape(u).*reshape(v))
end

function Base. /(u::Field{S}, v::Field{S})::Field{S} where {S}
    @assert range(u.space) == range(v.space)
    return reshape(u.space, reshape(u)./reshape(v))
end

function Base. +(u::Field{S}, v::Field{S})::Field{S} where {S}
    @assert range(u.space) == range(v.space)
    return reshape(u.space, reshape(u) + reshape(v))
end

function Base. -(u::Field{S}, v::Field{S})::Field{S} where {S}
    @assert range(u.space) == range(v.space)
    return reshape(u.space, reshape(u) - reshape(v))
end

function Base. ^(u::Field{S}, a::Number)::Field{S} where {S}
    return reshape(u.space, reshape(u).^a)
end

function Base. abs(u::Field{S})::Field{S} where {S}
    return reshape(u.space, abs.(reshape(u)))
end

function Base. maximum(u::Field{S})::Number where {S}
    return maximum(reshape(u))
end

function Base. minimum(u::Field{S})::Number where {S}
    return minimum(reshape(u))
end
function Base. *(a::Number, u::Field{S})::Field{S} where {S}
    return reshape(u.space, a.*reshape(u))
end

function Base. /(a::Number, u::Field{S})::Field{S} where {S}
    return reshape(u.space, a./reshape(u))
end

function Base. display(A::Operator{S}) where {S}
    display(reshape(A))
    println()
end

function Base. display(u::Field{S}) where {S}
    display(u.value)
    println()
end

function Base. sum(u::Field{S})::Number  where {S}
    return sum(reshape(u))
end

function Base. -(u::Field{S})::Field{S} where {S}
    return reshape(u.space, -reshape(u))
end

function Base. +(a::Number, u::Field{S})::Field{S} where {S}
    return reshape(u.space, a .+ reshape(u))
end

function Base. /(u::Field{S}, a::Number)::Field{S} where {S}
    return reshape(u.space, reshape(u)./a)
end

function Base. log(u::Field{S})::Field{S} where {S}
    return reshape(u.space, log.(reshape(u)))
end

function Base. log10(u::Field{S})::Field{S} where {S}
    return reshape(u.space, log10.(reshape(u)))
end

function Base. exp(u::Field{S})::Field{S} where {S}
    return reshape(u.space, exp.(reshape(u)))
end

function Base. sin(u::Field{S})::Field{S} where {S}
    return reshape(u.space, sin.(reshape(u)))
end

function Base. cos(u::Field{S})::Field{S} where {S}
    return reshape(u.space, cos.(reshape(u)))
end

function Base. sqrt(u::Field{S})::Field{S} where {S}
    return reshape(u.space, sqrt.(reshape(u)))
end

function Base. transpose(u::Field{S})::Field{S} where {S}
    return Field(u.space, transpose(u.value))
end

function LinearAlgebra. cond(A::Operator{S}) where {S}
    return cond(reshape(A))
end

function LinearAlgebra. eigvals(A::Operator{S}) where {S}
    return sort(abs.(eigvals(reshape(A))))
end

function LinearAlgebra. norm(u::Field{S})::Number where {S}
    W = integral(u.space)
    return sqrt(sum(W*(u*u)))
end

function value(u::Field{S, D, T})::Array{T, D} where {S, D, T}
    return u.value
end

function space(u::Field{S, D, T})::Array{T, D} where {S, D, T}
    return u.space
end

function linsolve(A::Operator{S}, u::Field{S})::Field{S} where {S}
    @assert range(A.space) == range(u.space)
    return reshape(A.space, reshape(A)\reshape(u))
end

function  ⊕(A::Operator{S}, B::Operator{S})::Operator{S} where {S}
    @assert range(A.space) == range(B.space)
    ID = identity(A.space)
    return (ID - B)*A + B
end

function  ⊕(A::Operator{S}, B::Operator{S}, C::Operator{S})::Operator{S} where {S}
    @assert range(A.space) == range(B.space)
    ID = identity(A.space)
    return (ID - B)*A + B*C
end

function  ⊕(u::Field{S}, v::Field{S})::Field{S} where {S}
    @assert range(u.space) == range(v.space)
    I = identity(u.space)
    B = incomingboundary(u.space)
    return (I - B)*v + B*u
end

function Base. isnan(u::Field{S})::Bool where {S}
    return any(isnan.(u.value))
end


text:
# Code derived from metabolicEP (https://github.com/anna-pa-m/Metabolic-EP)

function inplaceinverse!(dest::AbstractArray, source::AbstractArray)
    copyto!(dest, source)
    inv!(cholesky!(Hermitian(dest)))
end

Φ(x) = 0.5*(1.0+erf(x/sqrt(2.0)))
ϕ(x) = exp(-x.^2/2.0)/sqrt(2π)
ϕ(x, μ, σ) = inv(σ*sqrt(2π))*exp(-((x - μ)^2)/(2σ^2))