text:
using Documenter, DocumenterMarkdown

using Pkg
Pkg.add(PackageSpec(name="PhyloPlots", rev="master"))

using PhyloNetworks
DocMeta.setdocmeta!(PhyloNetworks, :DocTestSetup, :(using PhyloNetworks); recursive=true)
using PhyloPlots # to trigger any precompilation warning outside jldoctests

makedocs(
    sitename = "PhyloNetworks.jl",
    authors = "Claudia Solís-Lemus, Cécile Ané, Paul Bastide and contributors.",
    modules = [PhyloNetworks], # to list methods from PhyloNetworks only, not from Base etc.
    format = Documenter.HTML(prettyurls = get(ENV, "CI", nothing) == "true"), # easier local build
    pages = [
        "Home" => "index.md",
        "Manual" => [
            "Installation" => "man/installation.md",
            "Input Data for SNaQ" => "man/inputdata.md",
            "TICR pipeline" => "man/ticr_howtogetQuartetCFs.md",
            "Network estimation and display" => "man/snaq_plot.md",
            "Network comparison and manipulation" => "man/dist_reroot.md",
            "Candidate Networks" => "man/fixednetworkoptim.md",
            "Extract Expected CFs" => "man/expectedCFs.md",
            "Bootstrap" => "man/bootstrap.md",
            "Multiple Alleles" => "man/multiplealleles.md",
            "Continuous Trait Evolution" => "man/trait_tree.md",
            "Parsimony on networks" => "man/parsimony.md",
            "Discrete Trait Evolution" => "man/fitDiscrete.md",
            "Neighbour Joining" => "man/nj.md",
        ],
        "Library" => [
            "Public" => "lib/public.md",
            "Internals" => "lib/internals.md",
        ]
    ],
)

deploydocs(
    repo = "github.com/crsl4/PhyloNetworks.jl.git",
    push_preview = true,
)


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct FirewallRuleListResult <: SwaggerModel
    value::Any # spec type: Union{ Nothing, Vector{FirewallRule} } # spec name: value
    nextLink::Any # spec type: Union{ Nothing, String } # spec name: nextLink

    function FirewallRuleListResult(;value=nothing, nextLink=nothing)
        o = new()
        validate_property(FirewallRuleListResult, Symbol("value"), value)
        setfield!(o, Symbol("value"), value)
        validate_property(FirewallRuleListResult, Symbol("nextLink"), nextLink)
        setfield!(o, Symbol("nextLink"), nextLink)
        o
    end
end # type FirewallRuleListResult

const _property_map_FirewallRuleListResult = Dict{Symbol,Symbol}(Symbol("value")=>Symbol("value"), Symbol("nextLink")=>Symbol("nextLink"))
const _property_types_FirewallRuleListResult = Dict{Symbol,String}(Symbol("value")=>"Vector{FirewallRule}", Symbol("nextLink")=>"String")
Base.propertynames(::Type{ FirewallRuleListResult }) = collect(keys(_property_map_FirewallRuleListResult))
Swagger.property_type(::Type{ FirewallRuleListResult }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_FirewallRuleListResult[name]))}
Swagger.field_name(::Type{ FirewallRuleListResult }, property_name::Symbol) =  _property_map_FirewallRuleListResult[property_name]

function check_required(o::FirewallRuleListResult)
    true
end

function validate_property(::Type{ FirewallRuleListResult }, name::Symbol, val)
end

ext: jl


text:
# Options
tire_model = "nonlinear"; # "none" || "linear" || "nonlinear"
nonlinear_acceleration_constraints = "on"
vertical_tire_force_constraints = "on"
#integration_scheme =  :backward_euler
integration_scheme =  :trapazoidal
obstacle_exists = "true"
plotting = "false"
steps = 50; # size of t_solve
execution_horizon =  0.5; # was 0.5
if (typeof(execution_horizon) != Float64) error("execution_horizon must be a Float!") end
N = 60;   # preview horizo
x_ref   = 200;
y_ref   = 100;
psi_ref = pi/2; # define goal
if obstacle_exists == "true"
  Q = 3; # number of obstacles
  A = [5,10,2]
  B = [5,10,2]
  s_x = [-2,0,0]
  s_y = [0,0,4]
  X0 = [200,180,200]
  Y0 = [34,75,30]
end


text:
# https://man.openbsd.org/ssh-keygen.1
# https://github.com/JuliaBinaryWrappers/OpenSSH_jll.jl


text:
module DPNPlot

using PyPlot

include("colormaps.jl")
include("complex_colorize.jl")
include("labellines.jl")
include("pilabels.jl")
include("twodmesh.jl")

end # module


text:
using LinearAlgebra
using FermiCG
using Printf
using Test
using Arpack

@testset "fci" begin
    atoms = []
    push!(atoms,Atom(1,"H",[0,0,0]))
    push!(atoms,Atom(2,"H",[0,0,1]))
    push!(atoms,Atom(3,"H",[0,0,2]))
    push!(atoms,Atom(4,"H",[0,0,3]))
    push!(atoms,Atom(5,"H",[0,0,4]))
    push!(atoms,Atom(6,"H",[0,0,5]))
    #basis = "6-31g"
    basis = "sto-3g"
    mol     = Molecule(0,1,atoms,basis)

    mf = FermiCG.pyscf_do_scf(mol)
    nbas = size(mf.mo_coeff)[1]
    ints = FermiCG.pyscf_build_ints(mol,mf.mo_coeff, zeros(nbas,nbas));
    e_fci, d1_fci, d2_fci = FermiCG.pyscf_fci(ints,3,3)
    @printf(" FCI Energy: %12.8f\n", e_fci)

    n_elec_a = 3
    n_elec_b = 3

    norb = size(ints.h1)[1]
    problem = FermiCG.StringCI.FCIProblem(norb, n_elec_a, n_elec_b)

    display(problem)

    @time Hmat = FermiCG.StringCI.build_H_matrix(ints, problem)
    @time e,v = eigs(Hmat, nev = 10, which=:SR)
    e = real(e)
    for ei in e
        @printf(" Energy: %12.8f\n",ei+ints.h0)
    end
    @test isapprox(e[1], e_fci , atol=1e-10)
end

max_stars_repo_name: nmayhall-vt/FermiCG

max_issues_repo_name: nmayhall-vt/FermiCG

max_forks_repo_name: nmayhall-vt/FermiCG


text:
"""
    map_tr!(f, x, X, param)
Return an anonymous function that applies `f` to all elements of `x` and store the result in `X`, with a parameter `param` (which is set to `nothing` for function with no parameter).
"""
function map_tr!(f, x, X, p)
    return ((g, y, Y; param) -> map!(i -> g(i, y; param), Y, 1:length(y)))(f, x, X; param=p)
end
function map_tr!(f, x, X)
    return ((g, y, Y; param) -> map!(i -> g(i, y), Y, 1:length(y)))(f, x, X; param=nothing)
end

"""
    lazy(funcs::Function...)
Generate methods extended to a vector instead of one of its components. A function `f` should have the following signature: `f(i::Int, x::V)`.
"""
function lazy(funcs::Function...)
    for f in Iterators.map(Symbol, funcs)
        eval(
            :(
                function $f(x::V, X; param=nothing) where {V<:AbstractVector}
                    return map_tr!($f, x, X)
                end
            ),
        )
        eval(:($f(x; param=nothing) = $f(x, similar(x); param)))
    end
    return nothing
end

"""
    lazy_param(funcs::Function...)
Generate methods extended to a vector instead of one of its components. A function `f` should have the following signature: `f(i::Int, x::V; param)`.
"""
function lazy_param(funcs::Function...)
    for f in Iterators.map(Symbol, funcs)
        eval(:($f(x::V, X; param) where {V<:AbstractVector} = map_tr!($f, x, X, param)))
        eval(:($f(x; param) = $f(x, similar(x); param)))
    end
    return nothing
end

"""
    as_bitvector(n::Int, max_n::Int = n)
Convert an Int to a BitVector of minimal size (relatively to `max_n`).
"""
function as_bitvector(n::Int, max_n::Int=n)
    nm1 = n - 1
    v = falses(ceil(Int, log2(max_n)))
    i = 0
    @inbounds while !iszero(nm1)
        tz = trailing_zeros(nm1)
        i += (tz + 1)
        v[i] = true
        nm1 >>>= (tz + 1)
    end
    return v
end

"""
    as_int(v::AbstractVector)
Convert a `BitVector` into an `Int`.
"""
function as_int(v::AbstractVector)
    n = 0
    for (i, b) in enumerate(v)
        n += b ? 2^(i - 1) : 0
    end
    return n
end

"""
    reduce_symbols(symbols, sep)
Produce a formatted string that separates the symbols by `sep`. Used internally for `show_composition`.
"""
function reduce_symbols(symbols, sep, parenthesis=true; prefix="")
    str = reduce((x, y) -> "$y$sep$x", map(s -> "$prefix$s", symbols))
    return parenthesis ? "[$str]" : str
end

function incsert!(d::Dictionary, ind)
    return set!(d, ind, isassigned(d, ind) ? d[ind] + 1 : 1)
end

@unroll function tr_in(tr, X, x, param)
    @unroll for i in 1:length(tr)
        tr[i](x, @view(X[:, i]); param)
        # X[:,i] = tr[i](x; param) # NOTE - former version
    end
end


text:
using Revise
using AdFem
using PyCall
using LinearAlgebra
using PyPlot
using SparseArrays
using MAT
using ADCMEKit
np = pyimport("numpy")
using PyPlot
using SpecialFunctions
include("viscosity_accel/viscosity_accel.jl")

ADCME.options.sparse.auto_reorder = false
# matplotlib.use("Agg")
close("all")

# simulation parameter setup
n = 20
NT = 300
ρ = 0.1 # design variable in α-schemes
m = 4n
h = 1 / n
Δt = 3000. / NT
density = 100.

mode = "data"
# mode = "inv"

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

# Dirichlet boundary condition on three sides, and Neumann boundary condition (traction-free) on the top
bdnode = bcnode("left | lower | right", m, n, h)

# viscoelasticity η and shear modulus μ
ηf = (x, y)->begin
    if y <= 0.25
        return 100.
    elseif y <= 0.6
        return 2.
    else
      return 1
    end
end

if mode == "data"
    η = constant(eval_f_on_gauss_pts(ηf, m, n, h))
    # η = ones(4*m*n)
else

  ### 2D inversion
    # η = Variable(eval_f_on_gauss_pts(ηf, m, n, h))

  ### 1D layer inversion
    η_ = [100. * ones(5); ones(15)]
    global v_var = [constant(ones(5)); Variable(2.5ones(n - 5))]
    η = v_var .* η_

    ## invert log η
    # v_var = Variable(log(2.5) * ones(n))
    # η = exp(v_var)

    ## inv 1/η
    # η_ = [10000. * ones(5); ones(15)]
    # global v_var = [constant(ones(5)); Variable(1. / 2.5 * ones(n - 5))]
    # η = 1/v_var .* η_

    # v_var = Variable(2.5 * ones(n))
    # η = v_var

    global η = layer_model(η, m, n, h)

  ### uniform model inversion
  # global v_var = Variable(2.5)*ones(n)
  # η = v_var
  # global η = layer_model(η, m, n, h)

  ## debug
  # global η = placeholder(eval_f_on_gauss_pts(ηf, m, n, h))
end

## DEBUG
# pl = placeholder([1.0])

μ = 0.001 * constant(ones(4m * n))

### uniform model inversion
# μ_ = Variable(2.0)
# μ = μ_ * 0.001 * constant(ones(4m*n))

# linear elasticity matrix
coef = 2μ * η / (η + μ * Δt)
# mapH = c->begin
#     c * diagm(0 => ones(2))
# end
# H = map(mapH, coef)
H = compute_space_varying_tangent_elasticity_matrix(coef, m, n, h)

# H = reshape([coef constant(zeros(4*m*n)) coef constant(zeros(4*m*n))], (4*m*n,2,2))

Δt = Δt * ones(NT)

# generalized mass matrix and stiffness matrix
M = density * constant(compute_fem_mass_matrix1(m, n, h))
K = compute_fem_stiffness_matrix1(H, m, n, h)
C = spzero((m + 1) * (n + 1))

C = 0.1 * M + 0.1 * K

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

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

d0 = db
v0 = zeros((m + 1) * (n + 1))
# v0[bcnode("left", m, n, h)] .= 1.0
a0 = vector(idof, M[idof, idof] \ ((-K * db)[idof]), (m + 1) * (n + 1))

function antiplane_visco_αscheme(M::Union{SparseTensor,SparseMatrixCSC},
  K::Union{SparseTensor,SparseMatrixCSC},
  d0::Union{Array{Float64,1},PyObject},
  v0::Union{Array{Float64,1},PyObject},
  a0::Union{Array{Float64,1},PyObject},
  Δt::Array{Float64};
  ρ::Float64 = 1.0)
    nt = length(Δt)

    αm = (2ρ - 1) / (ρ + 1)
    αf = ρ / (1 + ρ)
    γ = 1 / 2 - αm + αf
    β = 0.25 * (1 - αm + αf)^2
    d = length(d0)

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

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

    A = (1 - αm) * M + Kterm
    A =  fem_impose_Dirichlet_boundary_condition1(A, bdnode, m, n, h)[1]

  # for visco_solve
  # ii, jj, vv = find(Kterm)
  # opp = push_matrices(A, Kterm)

    A = factorize(A)

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

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

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

        rhs = - (M * am + C * vf + K * df) + Force
        rhs = scatter_update(rhs, bdnode, zeros(length(bdnode)))

    # @info "visco_solve"
    # op = tf.print(norm(visco_solve(rhs,vv,opp)-A\rhs))
    # rhs = bind(rhs, op)
    # visco_solve(rhs,vv,opp)
    # @info "A \ rhs"
        A \ rhs
    end

    function condition(i, tas...)
        return i <= nt
    end

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

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

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

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

    dM = TensorArray(nt + 1); vM = TensorArray(nt + 1); aM = TensorArray(nt + 1); σM = TensorArray(nt + 1);
    dM = write(dM, 1, d0)
    vM = write(vM, 1, v0)
    aM = write(aM, 1, a0)

    ε0 = eval_strain_on_gauss_pts1(d0, m, n, h)
    σ0 = batch_matmul(H, ε0)
    σM = write(σM, 1, σ0)

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

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

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

function observation(d, v)
    idx = 10:m
    idx_plus = idx .+ 1
    idx_minus = idx .- 1
    idx_t = div(NT, 6):NT + 1
    dobs = d[idx_t, idx]
    vobs = v[idx_t, idx]
    strain_rate_obs = (v[idx_t, idx_plus] - v[idx_t, idx_minus]) / 2h
    return dobs, vobs, strain_rate_obs
end

dobs, vobs, strain_rate_obs = observation(d, v)

sess = Session(); init(sess)

if mode == "data"
    @time disp, vel, strain_rate = run(sess, [dobs, vobs, strain_rate_obs])
    matwrite("data-visc.mat", Dict("vel" => vel, "strain_rate" => strain_rate, "disp" => disp))
    # visulization()
    @info "Data Generated."
end

cb = (vs, iter, loss)->begin
    if mod(iter, 10) == 0 || iter < 10
        # clf()
        # plot(vs[1])
        x_tmp, y_tmp, z_tmp = visualize_scalar_on_gauss_points(vs[2], m, n, h)
        clf()
        pcolormesh(x_tmp, y_tmp, z_tmp', vmax=3, vmin = 0, rasterized=true)
        colorbar(shrink=0.2)
        axis("scaled")
        xlabel("x")
        ylabel("y")
        gca().invert_yaxis()
        title("Iter = $iter")
        savefig("figures2/inv_$(lpad(iter,5,"0")).png", bbox_size="tight")
        matwrite("results2/inv_$(lpad(iter,5,"0")).mat", Dict("var" => vs[1], "eta" => vs[2]))
    end
    printstyled("[#iter $iter] eta = $(vs[1])\n", color = :green)
end

if mode != "data"
    data = matread("data-visc.mat")
    global disp, vel, strain_rate =  data["disp"], data["vel"], data["strain_rate"]
    # global loss = 1e10 * sum((vel - vobs)^2)
    global loss = 1e10 * sum((disp - dobs)^2)
    @info run(sess, loss)
    global loss_ = BFGS!(sess, loss * 1e10, vars = [v_var, η], callback = cb)
end

## DEBUG
# η_ = [10000. * ones(5); ones(15)]
# v_var = [constant(ones(5)); Variable(2.5ones(n - 5))]
# η0 = v_var .* η_
# η0 = layer_model(η0, m, n, h)
# sess = Session();init(sess)
# @show run(sess, loss)
# η0 = run(sess, η0)
# lineview(sess, η, loss, eval_f_on_gauss_pts(ηf, m, n, h), η0)
# gradview(sess, η, loss, η0)
# gradview(sess, pl, loss, [1.0])

# @time run(sess, loss)

## plot model

x, y, z = visualize_scalar_on_gauss_points(run(sess, η), m, n, h)
figure()
pcolormesh(x, y, z', vmax=3, vmin = 0, rasterized=true)
colorbar(shrink=0.2)
axis("scaled")
xlabel("x")
ylabel("y")
gca().invert_yaxis()
if mode == "data"
  title(L"""Viscosity Model
  (top elastic layer: $η=100$)""")
  savefig("viscosity-model.png", bbox_size="tight")
else
  title("Inverted Model")
  savefig("viscosity-inversion.png", bbox_size="tight")
end

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

figure()
pl, = plot([], [], "o-", markersize = 3)
t = title("0")
xi = (0:m) * h
xlim(-h, (m + 1) * h)
xlabel("Distance")
ylim(-0.1, 1.1)
ylabel("Displacement")
tight_layout()
function update(i)
    pl.set_data(xi[:], d_[i,1:m + 1])
    t.set_text("time = $(i * Δt[1])")
end
p = animate(update, [NT÷6:10:NT+1;])
# p = animate(update, [NT÷2:5:NT+1])
# saveanim(p, "displacement.gif")

# figure()
# pl, = plot([], [], "o-", markersize = 3)
# t = title("time = 0")
# xi = (0:m)*h
# xlim(-h, (m+1)*h)
# xlabel("Distance")
# ylim(-0.0001, 0.0005)
# ylabel("Velocity")
# tight_layout()
# function update(i)
#   pl.set_data(xi[:], v_[i,1:m+1])
#   t.set_text("time = $(i*Δt[1])")
# end
# p = animate(update, [NT÷6:10:NT+1;])
# saveanim(p, "velocity.gif")

# figure()
# pl, = plot([], [], "o-", markersize = 3)
# t = title("time = 0")
# xi = (0:m)*h
# xlim(-h, (m+1)*h)
# xlabel("Distance")
# ylim(-0.001, 0.002)
# ylabel("Strain Rate")
# tight_layout()
# function update(i)
#   pl.set_data(xi[1:end-1], (v_[i,2:m+1]-v_[i,1:m])/h)
#   t.set_text("time = $(i*Δt[1])")
# end
# p = animate(update, [NT÷6:10:NT+1;])
# saveanim(p, "strain_rate.gif")

max_stars_repo_name: ADCMEMarket/ADCMEImages

max_issues_repo_name: ADCMEMarket/ADCMEImages

max_forks_repo_name: ADCMEMarket/ADCMEImages


text:
##########################
# FunctionCell Interface #
##########################
"""
	whiten([;kwargs])

Constructs an untrained cell that when piped data inside, calculates a transform `W`
of the input data `X` so that `Wᵀcov(X)W=I`.

# Keyword arguments (same as in `MultivariateStats`)
  * `regcoef::Real` is the ratio of variaces preserved in the principal subspace
  * `mean` is the mean vector, can be `nothing`, `0` or precomputed mean vector (default `nothing`)

Read the `MultivariateStats.jl` documentation for more information.
"""
whiten(;kwargs...) = FunctionCell(whiten, (), ModelProperties(), kwtitle("Whitening", kwargs); kwargs...)

############################
# DataCell/Array Interface #
############################
"""
	whiten(x, [;kwargs])

Trains a function cell that when receiving data into will whiten the input
the observations based on the projection matrix W extracted from `x`.
"""
# Training
whiten(x::T where T<:CellData; kwargs...) = whiten(getx!(x); kwargs...)
whiten(x::T where T<:AbstractVector; kwargs...) = whiten(mat(x, LearnBase.ObsDim.Constant{2}()); kwargs...)
whiten(x::T where T<:AbstractMatrix; kwargs...) = begin

	whdata = fit(MultivariateStats.Whitening, getobs(x); kwargs...)

	# Build model properties
	modelprops = ModelProperties(nvars(x), nvars(x))

	# Returned trained cell
	FunctionCell(whiten, Model(whdata, modelprops), kwtitle("Whitening",kwargs))
end

# Execution
whiten(x::T where T<:CellData, model::Model{<:MultivariateStats.Whitening}) =
	datacell(whiten(getx!(x), model), gety(x))
whiten(x::T where T<:AbstractVector, model::Model{<:MultivariateStats.Whitening}) =
	whiten(mat(x, LearnBase.ObsDim.Constant{2}()), model)
whiten(x::T where T<:AbstractMatrix, model::Model{<:MultivariateStats.Whitening}) =
	MultivariateStats.transform(model.data, getobs(x))


text:
using Oceananigans.Operators: Δzᵃᵃᶜ, Δzᵃᵃᶠ
using Oceananigans.AbstractOperations: flip
using Oceananigans.Solvers: BatchedTridiagonalSolver, solve!

#####
##### Batched Tridiagonal solver for implicit diffusion
#####

"""
    z_viscosity(closure, args...)

Returns the "vertical" (z-direction) viscosity associated with `closure`.
"""
function z_viscosity end

"""
    z_diffusivity(closure, ::Val{tracer_index}, args...) where tracer_index

Returns the "vertical" (z-direction) diffusivity associated with `closure` and `tracer_index`.
"""
function z_diffusivity end

implicit_step!(field, ::Nothing, args...; kwargs...) = NoneEvent()
implicit_diffusion_solver(::ExplicitTimeDiscretization, args...; kwargs...) = nothing

#####
##### Solver kernel functions for tracers / horizontal velocities and for vertical velocities
##### Note: "ivd" stands for implicit vertical diffusion.
#####

@inline κ_Δz²(i, j, kᶜ, kᶠ, grid, κ) = κ / Δzᵃᵃᶜ(i, j, kᶜ, grid) / Δzᵃᵃᶠ(i, j, kᶠ, grid)

instantiate(X) = X()

# Tracers and horizontal velocities at cell centers in z

@inline function ivd_upper_diagonal(i, j, k, grid, LX, LY, ::Center, clock, Δt, κ⁻⁻ᶠ, κ)
    κᵏ⁺¹ = κ⁻⁻ᶠ(i, j, k+1, grid, clock, κ)

    return ifelse(k > grid.Nz-1,
                  zero(eltype(grid)),
                  - Δt * κ_Δz²(i, j, k, k+1, grid, κᵏ⁺¹))
end

@inline function ivd_lower_diagonal(i, j, k, grid, LX, LY, ::Center, clock, Δt, κ⁻⁻ᶠ, κ)
    k′ = k + 1 # Shift to adjust for Tridiagonal indexing convenction
    κᵏ = κ⁻⁻ᶠ(i, j, k′, grid, clock, κ)

    return ifelse(k < 1,
                  zero(eltype(grid)),
                  - Δt * κ_Δz²(i, j, k′, k′, grid, κᵏ))
end

# Vertical velocity kernel functions (at cell interfaces in z)
#
# Note: these coefficients are specific to vertically-bounded grids (and so is
# the BatchedTridiagonalSolver).

@inline function ivd_upper_diagonal(i, j, k, grid, LX, LY, ::Face, clock, Δt, νᶜᶜᶜ, ν)
    νᵏ = νᶜᶜᶜ(i, j, k, grid, clock, ν)

    return ifelse(k < 1, # should this be k < 2? #should this be grid.Nz - 1?
                  zero(eltype(grid)),
                  - Δt * κ_Δz²(i, j, k, k, grid, νᵏ))
end

@inline function ivd_lower_diagonal(i, j, k, grid, LX, LY, ::Face, clock, Δt, νᶜᶜᶜ, ν)
    k′ = k + 1 # Shift to adjust for Tridiagonal indexing convenction
    νᵏ⁻¹ = νᶜᶜᶜ(i, j, k′-1, grid, clock, ν)
    return ifelse(k < 1,
                  zero(eltype(grid)),
                  - Δt * κ_Δz²(i, j, k′, k′-1, grid, νᵏ⁻¹))
end

### Diagonal terms

@inline ivd_diagonal(i, j, k, grid, LX, LY, LZ, clock, Δt, interp_κ, κ) =
    one(eltype(grid)) - ivd_upper_diagonal(i, j, k, grid, LX, LY, LZ, clock, Δt, interp_κ, κ) -
                        ivd_lower_diagonal(i, j, k-1, grid, LX, LY, LZ, clock, Δt, interp_κ, κ)

#####
##### Solver constructor
#####

"""
    implicit_diffusion_solver(::VerticallyImplicitTimeDiscretization, grid)

Build tridiagonal solvers for the elliptic equations

```math
(1 - Δt ∂z κ ∂z) cⁿ⁺¹ = c★
```

and

```math
(1 - Δt ∂z ν ∂z) wⁿ⁺¹ = w★
```

where `cⁿ⁺¹` and `c★` live at cell `Center`s in the vertical,
and `wⁿ⁺¹` and `w★` lives at cell `Face`s in the vertical.
"""
function implicit_diffusion_solver(::VerticallyImplicitTimeDiscretization, grid)

    topo = topology(grid)

    topo[3] == Periodic && error("VerticallyImplicitTimeDiscretization can only be specified on " *
                                 "grids that are Bounded in the z-direction.")

    z_solver = BatchedTridiagonalSolver(grid;
                                        lower_diagonal = ivd_lower_diagonal,
                                        diagonal = ivd_diagonal,
                                        upper_diagonal = ivd_upper_diagonal)

    return z_solver
end

#####
##### Implicit step functions
#####

is_c_location(loc) = loc === (Center, Center, Center)
is_u_location(loc) = loc === (Face, Center, Center)
is_v_location(loc) = loc === (Center, Face, Center)
is_w_location(loc) = loc === (Center, Center, Face)

"""
    implicit_step!(field, solver::BatchedTridiagonalSolver, clock, Δt,
                   closure, tracer_index, args...; dependencies = Event

Initialize the right hand side array `solver.batched_tridiagonal_solver.f`, and then solve the
tridiagonal system for vertically-implicit diffusion, passing the arguments
`clock, Δt, κ⁻⁻ᶠ, κ` into the coefficient functions that return coefficients of the
lower diagonal, diagonal, and upper diagonal of the resulting tridiagonal system.

`args...` are passed into `z_diffusivity` and `z_viscosity` appropriately for the purpose of retrieving
the diffusivities / viscosities associated with `closure`.
"""
function implicit_step!(field::AbstractField{LX, LY, LZ},
                        implicit_solver::BatchedTridiagonalSolver,
                        clock,
                        Δt,
                        closure,
                        tracer_index,
                        args...;
                        dependencies) where {LX, LY, LZ}

    location = (LX, LY, LZ)

    if is_c_location(location)

        locate_coeff = κᶜᶜᶠ
        coeff = z_diffusivity(closure, Val(tracer_index), args...)

    elseif is_u_location(location)

        locate_coeff = νᶠᶜᶠ
        coeff = z_viscosity(closure, args...)

    elseif is_v_location(location)

        locate_coeff = νᶜᶠᶠ
        coeff = z_viscosity(closure, args...)

    elseif is_w_location(location)

        locate_coeff = νᶜᶜᶜ
        coeff = z_viscosity(closure, args...)

    else
        error("Cannot take an implicit_step! for a field at $location")
    end

    return solve!(field, implicit_solver, field, instantiate.(location)...,
                  clock, Δt, locate_coeff, coeff; dependencies = dependencies)
end


text:
function Integrators.initialize!(int::AbstractIntegratorHSPARK, sol::AtomicSolutionPDAE)
    sol.t̄ = sol.t - timestep(int)

    equation(int, :v̄)(sol.t, sol.q, sol.p, sol.v)
    equation(int, :f̄)(sol.t, sol.q, sol.p, sol.f)

    initialize!(int.iguess, sol.t, sol.q, sol.p, sol.v, sol.f,
                            sol.t̄, sol.q̄, sol.p̄, sol.v̄, sol.f̄)
end

function initial_guess!(int::AbstractIntegratorHSPARK{DT}, sol::AtomicSolutionPDAE{DT},
                        cache::IntegratorCacheSPARK{DT}=int.caches[DT]) where {DT}
    for i in eachstage(int)
        evaluate!(int.iguess, sol.q̄, sol.p̄, sol.v̄, sol.f̄,
                              sol.q, sol.p, sol.v, sol.f,
                              cache.q̃, cache.p̃, cache.ṽ, cache.f̃,
                              tableau(int).q.c[i], tableau(int).p.c[i])

        for k in eachdim(int)
            int.solver.x[2*(ndims(int)*(i-1)+k-1)+1] = (cache.q̃[k] - sol.q[k])/timestep(int)
            int.solver.x[2*(ndims(int)*(i-1)+k-1)+2] = (cache.p̃[k] - sol.p[k])/timestep(int)
        end
    end

    for i in 1:pstages(int)
        evaluate!(int.iguess, sol.q̄, sol.p̄, sol.v̄, sol.f̄,
                              sol.q, sol.p, sol.v, sol.f,
                              cache.q̃, cache.p̃, cache.ṽ, cache.f̃,
                              tableau(int).q̃.c[i], tableau(int).p̃.c[i])

        for k in eachdim(int)
            int.solver.x[2*ndims(int)*nstages(int)+3*(ndims(int)*(i-1)+k-1)+1] = (cache.q̃[k] - sol.q[k])/timestep(int)
            int.solver.x[2*ndims(int)*nstages(int)+3*(ndims(int)*(i-1)+k-1)+2] = (cache.p̃[k] - sol.p[k])/timestep(int)
            int.solver.x[2*ndims(int)*nstages(int)+3*(ndims(int)*(i-1)+k-1)+3] = 0
        end
    end

    if isdefined(tableau(int), :λ) && tableau(int).λ.c[1] == 0
        for k in eachdim(int)
            int.solver.x[2*ndims(int)*nstages(int)+3*(k-1)+3] = sol.λ[k]
        end
    end
end


text:
using LinearAlgebra
using KrylovKit

export AbstractLattice, SquareLattice
abstract type AbstractLattice end
struct SquareLattice <: AbstractLattice end

export BCVUMPSRuntime, SquareBCVUMPSRuntime

# NOTE: should be renamed to more explicit names
"""
    BCVUMPSRuntime{LT}

a struct to hold the tensors during the `bcvumps` algorithm, each is a `Ni` x `Nj` Matrix, containing
- `d × d × d × d'` `M[i,j]` tensor
- `D × d' × D` `AL[i,j]` tensor
- `D × D`     `C[i,j]` tensor
- `D × d' × D` `AR[i,j]` tensor
- `D × d' × D` `FL[i,j]` tensor
- `D × d' × D` `FR[i,j]` tensor
and `LT` is a AbstractLattice to define the lattice type.
"""
struct BCVUMPSRuntime{LT,T,N,AT <: AbstractArray{<:AbstractArray,2},CT,ET}
    M::AT
    AL::ET
    C::CT
    AR::ET
    FL::ET
    FR::ET
    function BCVUMPSRuntime{LT}(M::AT, AL::ET, C::CT, AR::ET, FL::ET, FR::ET) where {LT <: AbstractLattice,AT <: AbstractArray{<:AbstractArray,2}, CT <: AbstractArray{<:AbstractArray,2}, ET <: AbstractArray{<:AbstractArray,2}}
        T, N = eltype(M[1,1]), ndims(M[1,1])
        new{LT,T,N,AT,CT,ET}(M, AL, C, AR, FL, FR)
    end
end

const SquareBCVUMPSRuntime{T,AT} = BCVUMPSRuntime{SquareLattice,T,4,AT}
function SquareBCVUMPSRuntime(M::AT, AL, C, AR, FL, FR) where {AT <: AbstractArray{<:AbstractArray,2}}
    ndims(M[1,1]) == 4 || throw(DimensionMismatch("M dimensions error, should be `4`, got $(ndims(M[1,1]))."))
    BCVUMPSRuntime{SquareLattice}(M, AL, C, AR, FL, FR)
end

@doc raw"
    SquareBCVUMPSRuntime(M::AbstractArray{T,4}, env::Val, χ::Int)

create a `SquareBCVUMPSRuntime` with M-tensor `M`. The `NixNj` `AL,C,AR,FL,FR`
tensors are initialized according to `env`. If `env = Val(:random)`,
the `A[i,j]` is initialized as a random `D×d×D` tensor,and `AL[i,j],C[i,j],AR[i,j]` are the corresponding
canonical form. `FL,FR` is the left and right environment.

# example

```jldoctest; setup = :(using BCVUMPS)
julia> Ni, Nj = 2, 2;

julia> M = Array{Array{Float64,3},2}(undef, Ni, Nj);

julia> for j = 1:Nj, i = 1:Ni
           M[i,j] = rand(2,2,2,2)
       end

julia> rt = SquareBCVUMPSRuntime(M, Val(:random), 4);

julia> size(rt.AL) == (2,2)
true

julia> size(rt.AL[1,1]) == (4,2,4)
true
```
"
function SquareBCVUMPSRuntime(M::AbstractArray{<:AbstractArray,2}, env, D::Int; verbose=false)
    return SquareBCVUMPSRuntime(M, _initializect_square(M, env, D; verbose=verbose)...)
end

function _initializect_square(M::AbstractArray{<:AbstractArray,2}, env::Val{:random}, D::Int; verbose=false)
    A = initialA(M, D)
    AL, L = leftorth(A)
    R, AR = rightorth(AL)
    _, FL = leftenv(AL, M)
    _, FR = rightenv(AR, M)
    C = LRtoC(L,R)
    Ni, Nj = size(M)
    verbose && print("random initial bcvumps $(Ni)×$(Nj) environment-> ")
    AL, C, AR, FL, FR
end

function _initializect_square(M::AbstractArray{<:AbstractArray,2}, chkp_file::String, D::Int; verbose=false)
    env = load(chkp_file)["env"]
    Ni, Nj = size(M)
    verbose && print("bcvumps $(Ni)×$(Nj) environment load from $(chkp_file) -> ")
    env.AL, env.C, env.AR, env.FL, env.FR
end

function bcvumps(rt::BCVUMPSRuntime; tol::Real, maxiter::Int, verbose=false)
    # initialize
    olderror = Inf

    stopfun = StopFunction(olderror, -1, tol, maxiter)
    rt, err = fixedpoint(res -> bcvumpstep(res...), (rt, olderror), stopfun)
    verbose && println("bcvumps done@step: $(stopfun.counter), error=$(err)")
    return rt
end

function bcvumpstep(rt::BCVUMPSRuntime, err)
    M, AL, C, AR, FL, FR = rt.M, rt.AL, rt.C, rt.AR, rt.FL, rt.FR
    AL, C, AR = ACCtoALAR(AL, C, AR, M, FL, FR)
    _, FL = leftenv(AL, M, FL)
    _, FR = rightenv(AR, M, FR)
    err = error(AL, C, FL, M, FR)
    return SquareBCVUMPSRuntime(M, AL, C, AR, FL, FR), err
end


text:
module MusicTransformer

include("transformer.jl")
include("positional_encoding.jl")
include("mask.jl")

export Attention, Transformer, Encoder, Decoder,
        get_positional_encoding,
        mask

end # module


text:
module TimeoutRequest

import ..Layer, ..request
using ..ConnectionPool
import ..@debug, ..DEBUG_LEVEL

"""
    request(TimeoutLayer, ::IO, ::Request, body) -> HTTP.Response

Close `IO` if no data has been received for `timeout` seconds.
"""
abstract type TimeoutLayer{Next <: Layer} <: Layer end
export TimeoutLayer

function request(::Type{TimeoutLayer{Next}}, io::IO, req, body;
                 readtimeout::Int=60, kw...) where Next

    wait_for_timeout = Ref{Bool}(true)

    @async while wait_for_timeout[]
        if isreadable(io) && inactiveseconds(io) > readtimeout
            close(io)
            @debug 1 "💥  Read inactive > $(readtimeout)s: $io"
            break
        end
        sleep(8 + rand() * 4)
    end

    try
        return request(Next, io, req, body; kw...)
    finally
        wait_for_timeout[] = false
    end
end

end # module TimeoutRequest


text:
# 1) given a Vector{T}, T<:Number, construct an object of type S using the same underlying data
# e.g.
#  array = [1.0,2.0,3.0], S = Tuple{Complex{Float64},Float64} => (1.0 + 2.0im, 3.0)

# 2) for an Array{T,N}, and some 1 <= D <= N,
#   reconstruct S from a slice of array#

# The main difference between StructArrays and what we want is that we want 1 underlying (rather than # fields)
# there is no heterogeneity in the struct (can store recursive representations) of field types

# get_offset(array, S, offset) => S(...), next_offset
import Base: @propagate_inbounds

"""
    basetype(S...)

Compute the "base" floating point type of one or more types `S`. This will throw
an error if there is no unique type.
"""
basetype(::Type{FT}) where {FT <: AbstractFloat} = FT
basetype(::Type{NamedTuple{names, T}}) where {names, T} = basetype(T)
function basetype(::Type{S}) where {S}
    isprimitivetype(S) && error("$S is not a floating point type")
    basetype(ntuple(i -> fieldtype(S, i), fieldcount(S))...)
end
function basetype(::Type{S1}, Sx...) where {S1}
    if sizeof(S1) == 0
        return basetype(Sx...)
    end
    FT1 = basetype(S1)
    FT2 = basetype(Sx...)
    FT1 !== FT2 && error("Inconsistent basetypes $FT1 and $FT2")
    return FT1
end

replace_basetype(::Type{S}, ::Type{FT}) where {S <: AbstractFloat, FT} = FT
function replace_basetype(::Type{S}, ::Type{FT}) where {S <: Tuple, FT}
    Tuple{ntuple(i -> replace_basetype(fieldtype(S, i), FT), fieldcount(S))...}
end
function replace_basetype(
    ::Type{NamedTuple{names, T}},
    ::Type{FT},
) where {names, T, FT}
    NamedTuple{names, replace_basetype(T, FT)}
end

"""
    parent_array_type(::Type{<:AbstractArray})

Returns the parent array type underlying the `SubArray` wrapper type
"""
parent_array_type(::Type{A}) where {A <: AbstractArray{FT}} where {FT} =
    Array{FT}
# TODO: extract interface to overload for backends into separate file

"""
    StructArrays.bypass_constructor(T, args)

Create an instance of type `T` from a tuple of field values `args`, bypassing
possible internal constructors. `T` should be a concrete type.
"""
@generated function bypass_constructor(::Type{T}, args) where {T}
    vars = ntuple(_ -> gensym(), fieldcount(T))
    assign = [
        :($var::$(fieldtype(T, i)) = getfield(args, $i)) for
        (i, var) in enumerate(vars)
    ]
    construct = Expr(:new, :T, vars...)
    Expr(:block, assign..., construct)
end

"""
    fieldtypeoffset(T,S,i)

Similar to `fieldoffset(S,i)`, but gives result in multiples of `sizeof(T)` instead of bytes.
"""
fieldtypeoffset(::Type{T}, ::Type{S}, i) where {T, S} =
    Int(div(fieldoffset(S, i), sizeof(T)))

"""
    typesize(T,S)

Similar to `sizeof(S)`, but gives the result in multiples of `sizeof(T)`.
"""
function typesize(::Type{T}, ::Type{S}) where {T, S}
    isbitstype(T) || error("$T is not isbitstype")
    isbitstype(S) || error("$S is not isbitstype")
    div(sizeof(S), sizeof(T))
end

# TODO: this assumes that the field struct zero type is the same as the backing
# zero'd out memory, which should be true in all "real world" cases
# but is something that should be revisited
@inline function _mzero!(out::MArray{S, T, N, L}, FT) where {S, T, N, L}
    TdivFT = DataLayouts.typesize(FT, T)
    Base.GC.@preserve out begin
        @inbounds for i in 1:(L * TdivFT)
            Base.unsafe_store!(
                Base.unsafe_convert(Ptr{FT}, Base.pointer_from_objref(out)),
                zero(FT),
                i,
            )
        end
    end
    return out
end

"""
    get_struct(array, S[, offset=0])

Construct an object of type `S` from the values of `array`, optionally offset by `offset` from the start of the array.
"""
function get_struct(array::AbstractArray{T}, ::Type{S}, offset) where {T, S}
    if @generated
        tup = :(())
        for i in 1:fieldcount(S)
            push!(
                tup.args,
                :(get_struct(
                    array,
                    fieldtype(S, $i),
                    offset + fieldtypeoffset(T, S, $i),
                )),
            )
        end
        :(bypass_constructor(S, $tup))
    else
        args = ntuple(fieldcount(S)) do i
            get_struct(array, fieldtype(S, i), offset + fieldtypeoffset(T, S, i))
        end
        return bypass_constructor(S, args)
    end
end
@propagate_inbounds function get_struct(
    array::AbstractArray{S},
    ::Type{S},
    offset,
) where {S}
    return array[offset + 1]
end

@propagate_inbounds get_struct(
    array::AbstractArray{T},
    ::Type{S},
) where {T, S} = get_struct(array, S, 0)

function set_struct!(array::AbstractArray{T}, val::S, offset) where {T, S}
    if @generated
        errorstring = "Expected type $T, got type $S"
        ex = quote
            # TODO: need to figure out a better way to handle the case where we require conversion
            # e.g. if T = Dual or Double64

            if isprimitivetype(S)
                # error if we dont hit triangular dispatch method (defined below):
                # set_struct!(::AbstractArray{S}, ::S) where {S}
                error($errorstring)
            end
            # TODO: we get a segfault here when trying to pass propogate_inbounds
            # with a generated function ctx (in the quote block or attached to the generated function)
            # passing it in the quoted expr for args seems to work :( @propagate_inbounds set_struct! ...)
            # https://github.com/JuliaArrays/StaticArrays.jl/blob/52fc10278667dd5fa82ded1edcfd5f7fedfae1c4/src/indexing.jl#L16-L38
            # Base.@_propagate_inbounds_meta
        end
        for i in 1:fieldcount(S)
            push!(
                ex.args,
                :(set_struct!(
                    array,
                    getfield(val, $i),
                    offset + fieldtypeoffset(T, S, $i),
                )),
            )
        end
        push!(ex.args, :(return nothing))
        ex
    else
        if isprimitivetype(S)
            return error("Expected type $T, got type $S")
        end
        for i in 1:fieldcount(S)
            set_struct!(
                array,
                getfield(val, i),
                offset + fieldtypeoffset(T, S, i),
            )
        end
        return nothing
    end
end

@propagate_inbounds function set_struct!(
    array::AbstractArray{S},
    val::S,
    offset,
) where {S}
    array[offset + 1] = val
end

@propagate_inbounds function set_struct!(array, val)
    set_struct!(array, val, 0)
end


text:
using DataFrames
using CSV
using Plots
#%%
run(`./harmonic`)
data_v = open("harmonic_verlet.dat")
data_v4 = open("harmonic_verlet4.dat")
data_rk = open("harmonic_rk.dat")
data_rk4 = open("harmonic_rk4.dat")

df_v=CSV.read(data_v, header=false, delim=" ")
df_v4=CSV.read(data_v4, header=false, delim=" ")
df_rk=CSV.read(data_rk, header=false, delim=" ")
df_rk4=CSV.read(data_rk4, header=false, delim=" ")

t_v=df_v[1]
x_v=df_v[2]
v_v=df_v[3]

t_v4=df_v4[1]
x_v4=df_v4[2]
v_v4=df_v4[3]

t_rk=df_rk[1]
x_rk=df_rk[2]
v_rk=df_rk[3]

t_rk4=df_rk4[1]
x_rk4=df_rk4[2]
v_rk4=df_rk4[3]
#%%

K_v=1//2*v_v.^2
U_v=1//2*x_v.^2
En_v=K_v+U_v

K_v4=1//2*v_v4.^2
U_v4=1//2*x_v4.^2
En_v4=K_v4+U_v4

K_rk=1//2*v_rk.^2
U_rk=1//2*x_rk.^2
En_rk=K_rk+U_rk

K_rk4=1//2*v_rk4.^2
U_rk4=1//2*x_rk4.^2
En_rk4=K_rk4+U_rk4
#%%
plot(t_v, En_v, label="Energy Verlet", dpi=600)
plot!(t_v4, En_v4, label="Energy Verlet 4", dpi=600)
plot!(t_rk, En_rk, label="Energy Runge-Kutta", dpi=600)
plot!(t_rk4, En_rk4, label="Energy Runge-Kutta 4", dpi=600, line=:dash, color="black")
plot!(legend=:topleft)
xlabel!("step")
ylabel!("Energy")
savefig("harmonic4.png")

#%%
plot(t_v4, En_v4, label="Energy Verlet 4", dpi=600)
plot!(t_rk4, En_rk4, label="Energy Runge-Kutta 4", dpi=600, line=:dash)
plot!(legend=:bottomleft)
xlabel!("step")
ylabel!("Energy")
savefig("harmonic2.png")

pi/4

max_stars_repo_name: aurelio-amerio/corso-cpp


text:
#=
Author:= Shreeda Bhat
License:= MIT
UTF-8
=#

######################################Array Types##########################################
println("######################################Array Types##########################################")
println(Array{Int64}(undef,3))

#***************# 1-D Array #***************#
println("#***************# 1-D Array #***************#")
println(Array{Int64,1}(undef,3))

#***************# 2-D Array #***************#
println("#***************# 2-D Array #***************#")
println(Array{Int64,2}(undef,3,3))

#***************# Array with String #***************#

println("#***************# Array with String #***************#")
println(Array{String}(undef,3)) # you get undef, should we should to an variable

#***************# Array with any Types #***************#

println("#***************# Array with any Types #***************#")
println(Array{any}(undef,3)) # Same here

#***************#Let's Greet #***********************#
println("#***************#Let's Greet #***********************#")
greeting = "Hello, Thanks for checking my practice"
println(greeting)

#***************# What's the variable keyword/name avaiable in Julia #***********************#
println("#***************# What's the variable keyword/name avaiable in Julia #***********************#")
#= variable in julia can be anything letter, word and number, including special character. Even Mandarin,Chinese can also be a variable name, even unicode
Even underscore(_) is allowed but not a standard suggested by Julia Developer

=#
println("#***************# Let's try Mandarin variable#***********************#")
的 = "hello, world"
println(的)

dimension_array = ['a','b','c','d']
println(dimension_array)


text:
#=
Implement a URL shortener with the following methods:

    shorten(url), which shortens the url into a six-character alphanumeric string, such as zLg6wl.
    restore(short), which expands the shortened string into the original url. If no such shortened string exists, return null.

Hint: What if we enter the same URL twice?
=#
using Random
global url_list = Dict{String, String}()

function shorten(url::String)
    if !in(url, values(url_list))
        shortened_url = randstring(6)
        while in(shortened_url, keys(url_list))
            shortened_url = randstring(6)
        end
        url_list[shortened_url] = url
    else
        shortened_url = collect(keys(url_list))[collect(values(url_list)) .== url][1]
    end
    return shortened_url
end

function restore(shortened_url)
    if in(shortened_url, keys(url_list))
        return url_list[shortened_url]
    else
        return nothing
    end
end

max_stars_repo_name: DominiqueCaron/daily-coding-problem

max_issues_repo_name: DominiqueCaron/daily-coding-problem

max_forks_repo_name: DominiqueCaron/daily-coding-problem


text:
abstract type SearchStrategy end
draw(s::SearchStrategy) = error("draw not defined by $(s)")

function Base.iterate(strategy::SearchStrategy, state=nothing)
    state = state === nothing ? 0 : state
    #TODO: Remove hardcoded count
    state >= PARAM_TEST_SETTINGS.samples && return nothing
    return (draw(strategy), state + 1)
end

Base.eltype(::SearchStrategy) = Any


text:
using Bonsai, FilePaths
using FilePathsBase: /
using AbstractTrees: print_tree

const app = App()
const static = Path(@__DIR__)

app.get("/") do stream
    Bonsai.write(stream, static / "index.html")
end

app.get("/ws") do stream
    ws = Bonsai.ws_upgrade(stream)
    try
        while !eof(ws)
            data = readavailable(ws)
            s = String(data)
            write(ws, s)
        end
    catch e
        @error e
    finally
        close(ws)
    end
end

start(app, port=9999, verbose=true)
stop(app)
wait(router)

max_stars_repo_name: onetonfoot/Tree

max_issues_repo_name: onetonfoot/Tree

max_forks_repo_name: onetonfoot/Tree


text:
##########################################################################
# Copyright 2017 Samuel Ridler.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
##########################################################################

# Dynamic Maximum Expected Coverage Location Problem (DMEXCLP)
# From paper: "An efficient heuristic for real-time ambulance redeployment"

"""
	function initDmexclp!(sim::Simulation;
		busyFraction::Float = 0.5,
		demandWeights::Dict{Priority,Float} = Dict([p => 1.0 for p in priorities]))
Initialise data for the Dynamic Maximum Expected Coverage Location Problem (DMEXCLP).

# Keyword arguments
- `busyFraction` is the fraction of time that ambulances are busy; should be within [0,1] though this is not enforced
- `demandWeights` is the weight of each demand priority on the objective function

Mutates: `sim.moveUpData.dmexclpData`
"""
function initDmexclp!(sim::Simulation;
	busyFraction::Float = 0.5,
	demandWeights::Dict{Priority,Float} = Dict([p => 1.0 for p in priorities]))

	# initialise demand and demand coverage data if not already initialised
	sim.demand.initialised || initDemand!(sim)
	sim.demandCoverage.initialised || initDemandCoverage!(sim)

	# shorthand
	@unpack numAmbs, numStations = sim

	dcd = sim.moveUpData.dmexclpData # shorthand
	dcd.busyFraction = busyFraction
	dcd.demandWeights = demandWeights

	# calculate cover benefit values, for single demand
	dcd.marginalBenefit = (busyFraction.^[0:numAmbs-1;])*(1-busyFraction)

	# values that will be calculated when needed
	dcd.stationNumFreeAmbs = zeros(Int, numStations)
	dcd.stationMarginalCoverages = zeros(Float, numStations) # stationMarginalCoverages[i] gives extra coverage provided from placing newly freed ambulance at station i
	# dcd.pointSetsCoverCounts = [zeros(Int, length(pointsCoverageMode.pointSets)) for pointsCoverageMode in sim.demandCoverage.pointsCoverageModes] # pointSetsCoverCounts[i][j] = number of free ambulances covering node set j, for demand.pointsCoverageModes i
end

function dmexclpMoveUp(sim::Simulation, newlyFreedAmb::Ambulance)
	@assert(sim.moveUpData.useMoveUp)
	@assert(newlyFreedAmb.status == ambFreeAfterCall)
	@assert(sim.demand.initialised && sim.demandCoverage.initialised)

	# shorthand names:
	dcd = sim.moveUpData.dmexclpData
	@unpack ambulances, stations, numStations = sim

	# calculate the number of free ambulances at (or travelling to) each station
	dcd.stationNumFreeAmbs[:] .= 0
	for (i,ambulance) in enumerate(ambulances)
		# do not count newly freed ambulance, it has not been assigned a station
		if isAmbMovable(ambulance) && i != newlyFreedAmb.index
			dcd.stationNumFreeAmbs[ambulance.stationIndex] += 1
		end
	end

	# ignoring newly freed amb, count number of ambulances covering each point set
	demandsPointSetsCoverCounts = calcPointSetsCoverCounts!(sim, sim.time, dcd.stationNumFreeAmbs)

	# find station allocation for newly freed ambulance that gives greatest
	# increase in expected demand coverage
	dcd.stationMarginalCoverages[:] .= 0.0
	for demandPriority in priorities
		if !haskey(dcd.demandWeights, demandPriority) || dcd.demandWeights[demandPriority] == 0
			continue
		end
		pointSetsCoverCounts = demandsPointSetsCoverCounts[demandPriority]
		pointsCoverageMode = getPointsCoverageMode!(sim, demandPriority, sim.time)
		demandMode = getDemandMode!(sim.demand, demandPriority, sim.time)
		pointSetsDemands = sim.demandCoverage.pointSetsDemands[pointsCoverageMode.index, demandMode.rasterIndex]
		# to do: if marginal coverage has already been calculated for the combination of pointsCoverageMode and rasterIndex (but for a different demand priority), the marginal coverage value should be reused (accounting for differences in old and new values of demandMode.rasterMultiplier and dcd.demandWeights[demandPriority]) to save on computation.
		for i = 1:length(pointsCoverageMode.pointSets)
			pointSetDemand = pointSetsDemands[i] * demandMode.rasterMultiplier
			pointSetMarginalCoverage = pointSetDemand * dcd.marginalBenefit[pointSetsCoverCounts[i]+1] * dcd.demandWeights[demandPriority]
			for j in pointsCoverageMode.stationSets[i]
				dcd.stationMarginalCoverages[j] += pointSetMarginalCoverage
			end
		end
	end
	(bestMarginalCoverage, bestStationIndex) = findmax(dcd.stationMarginalCoverages)

	return [newlyFreedAmb], [stations[bestStationIndex]]
end


text:
porosity(image :: AbstractArray{Bool}) =
    count(iszero, image) / length(image)

function corrfn(image  :: AbstractArray{Bool},
                cutoff :: Integer)
    return Directional.s2(image, false; periodic = true, len = cutoff) |> mean
end

function likelihood_point(x  :: AbstractFloat,
                          xs :: AbstractVector{<:AbstractFloat})
    σ = var(xs)
    μ = mean(xs)

    return -log(sqrt(2π * σ)) - (x - μ)^2/(2σ)
end

function likelihood(cf, size,
                    porosity   :: AbstractFloat,
                    dist       :: Distribution;
                    ncloud     :: Integer = 100)
    cutoff = length(cf)
    cloud = mapreduce(vcat, 1:ncloud) do _
        corrfn(generate_image(size, porosity, dist), cutoff)'
    end

    cloud_points = (cloud[:, n] for n in 1:cutoff)
    res = sum(likelihood_point(x, xs) for (x, xs) in zip(cf, cloud_points))
    return isnan(res) ? -Inf : res
end

function likelihood(image  :: AbstractArray{Bool},
                    dist   :: Distribution;
                    ncloud :: Integer = 100,
                    cutoff :: Integer = typemax(Int))
    likelihood(corrfn(image, cutoff), size(image), porosity(image), dist;
               ncloud = ncloud)
end


text:
export Tensor,
       tensor,
       inds,
       ind,
       store

"""
Tensor{StoreT,IndsT}

A plain old tensor (with order independent
interface and no assumption of labels)
"""
struct Tensor{ElT,
              N,
              StoreT<:TensorStorage,
              IndsT} <: AbstractArray{ElT,N}
  store::StoreT
  inds::IndsT
  """
  Tensor{ElT,N,StoreT,IndsT}(inds, store::StorageType)

  Internal constructor for creating a Tensor from the
  storage and indices.

  The Tensor is a view of the tensor storage.

  For normal usage, use the Tensor(store::TensorStorage, inds)
  and tensor(store::TensorStorage, inds) constructors.
  """
  function Tensor{ElT,N,StoreT,IndsT}(inds,
                                      store::TensorStorage) where {ElT,
                                                                   N,
                                                                   StoreT,
                                                                   IndsT}
    new{ElT,N,StoreT,IndsT}(store, inds)
  end
end

"""
tensor(store::TensorStorage, inds)

Construct a Tensor from a tensor storage and indices.
The Tensor is a view of the tensor storage.

See also Tensor(store::TensorStorage, inds)
which makes a copy of the storage.
"""
function tensor(store::StoreT,
                inds::IndsT) where {StoreT<:TensorStorage,
                                    IndsT}
  return Tensor{eltype(store),length(inds),StoreT,IndsT}(inds, store)
end

"""
Tensor(store::TensorStorage, inds)

Construct a Tensor from a tensor storage and indices.
The storage data is copied into the Tensor.

See also tensor(store::TensorStorage, inds)
which makes a view of the storage.
"""
Tensor(store::TensorStorage, inds) = tensor(copy(store), inds)

store(T::Tensor) = T.store

data(T::Tensor) = data(store(T))

storetype(::Tensor{ElT,N,StoreT}) where {ElT,N,StoreT} = StoreT

inds(T::Tensor) = T.inds

ind(T::Tensor, j::Integer) = inds(T)[j]

Base.eachindex(T::Tensor) = CartesianIndices(dims(inds(T)))

Base.eltype(::Tensor{ElT}) where {ElT} = ElT

Base.strides(T::Tensor) = strides(inds(T))

#
# Generic Tensor functions
#

# The size is obtained from the indices
dims(T::Tensor) = dims(inds(T))
dim(T::Tensor) = dim(inds(T))
dim(T::Tensor,i::Int) = dim(inds(T),i)
mindim(T::Tensor) = mindim(inds(T))
diaglength(T::Tensor) = mindim(T)
Base.size(T::Tensor) = dims(T)
Base.size(T::Tensor,i::Int) = dim(T,i)

# Needed for passing Tensor{T,2} to BLAS/LAPACK
# TODO: maybe this should only be for DenseTensor?
function Base.unsafe_convert(::Type{Ptr{ElT}},
                             T::Tensor{ElT}) where {ElT}
  return Base.unsafe_convert(Ptr{ElT},store(T))
end

Base.copy(T::Tensor) = tensor(copy(store(T)), inds(T))

Base.copyto!(R::Tensor,T::Tensor) = (copyto!(store(R), store(T)); R)

Base.complex(T::Tensor) = tensor(complex(store(T)), inds(T))

#
# Necessary to overload since the generic fallbacks are
# slow
#

LinearAlgebra.norm(T::Tensor) = norm(store(T))

Base.conj(T::Tensor;
          kwargs...) = tensor(conj(store(T); kwargs...), inds(T))

Random.randn!(T::Tensor) = (randn!(store(T)); T)

LinearAlgebra.rmul!(T::Tensor,α::Number) = (rmul!(store(T),α); T)
scale!(T::Tensor,α::Number) = rmul!(store(T),α)

Base.fill!(T::Tensor,α::Number) = (fill!(store(T),α); T)

#function Base.similar(::Type{<:Tensor{ElT,N,StoreT}},dims) where {ElT,N,StoreT}
#  return tensor(similar(StoreT,dim(dims)),dims)
#end

# TODO: make sure these are implemented correctly
#Base.similar(T::Type{<:Tensor},::Type{S}) where {S} = tensor(similar(store(T),S),inds(T))
#Base.similar(T::Type{<:Tensor},::Type{S},dims) where {S} = tensor(similar(store(T),S),dims)

Base.similar(T::Tensor) = tensor(similar(store(T)), inds(T))

# TODO: for BlockSparse, this needs to include the offsets
# TODO: for Diag, the storage is not just the total dimension
#Base.similar(T::Tensor,dims) = _similar_from_dims(T,dims)

# To handle method ambiguity with AbstractArray
#Base.similar(T::Tensor,dims::Dims) = _similar_from_dims(T,dims)

Base.similar(T::Tensor,
             ::Type{S}) where {S} = tensor(similar(store(T),S),
                                           inds(T))

Base.similar(T::Tensor,
             ::Type{S},
             dims) where {S<:Number} = _similar_from_dims(T, S, dims)

# To handle method ambiguity with AbstractArray
Base.similar(T::Tensor,
             ::Type{S},
             dims::Dims) where {S<:Number} = _similar_from_dims(T, S, dims)

_similar_from_dims(T::Tensor,
                   dims) = tensor(similar(store(T), dim(dims)), dims)

function _similar_from_dims(T::Tensor,::Type{S},dims) where {S<:Number}
  return tensor(similar(store(T),S,dim(dims)),dims)
end

function Base.convert(::Type{<:Tensor{<:Number,N,StoreR,Inds}},
                      T::Tensor{<:Number,N,<:Any,Inds}) where {N,Inds,StoreR}
  return tensor(convert(StoreR,store(T)),inds(T))
end

function Base.zeros(::Type{<:Tensor{ElT,N,StoreT}},
                    inds) where {ElT,N,StoreT}
  return tensor(zeros(StoreT,dim(inds)),inds)
end

# This is to fix a method ambiguity with a Base array function
function Base.zeros(::Type{<:Tensor{ElT,N,StoreT}},
                    inds::Dims{N}) where {ElT,N,StoreT}
  return tensor(zeros(StoreT,dim(inds)),inds)
end

function Base.promote_rule(::Type{<:Tensor{ElT1,N1,StoreT1,IndsT1}},
                           ::Type{<:Tensor{ElT2,N2,StoreT2,IndsT2}}) where {ElT1,ElT2,
                                                                            N1,N2,
                                                                            StoreT1,StoreT2,
                                                                            IndsT1,IndsT2}
  StoreR = promote_type(StoreT1,StoreT2)
  ElR = eltype(StoreR)
  return Tensor{ElR,N3,StoreR,IndsR} where {N3,IndsR}
end

function Base.promote_rule(::Type{<:Tensor{ElT1,N,StoreT1,Inds}},
                           ::Type{<:Tensor{ElT2,N,StoreT2,Inds}}) where {ElT1,ElT2,N,
                                                                         StoreT1,StoreT2,Inds}
  StoreR = promote_type(StoreT1,StoreT2)
  ElR = eltype(StoreR)
  return Tensor{ElR,N,StoreR,Inds}
end

# Convert the tensor type to the closest dense
# type
function dense(::Type{<:Tensor{ElT,NT,StoreT,IndsT}}) where {ElT,NT,StoreT,IndsT}
  return Tensor{ElT,NT,dense(StoreT),dense(IndsT)}
end

dense(T::Tensor) = tensor(dense(store(T)), inds(T))

function similar_type(::Type{<:Tensor{ElT,<:Any,StoreT,<:Any}},
                      ::Type{IndsR}) where {ElT,StoreT,IndsR}
  return Tensor{ElT,length(IndsR),StoreT,IndsR}
end

function similar_type(::Type{<:Tensor{ElT,<:Any,StoreT,<:Any}},
                      ::Type{IndsR}) where {ElT,StoreT,IndsR<:NTuple{NR}} where {NR}
  return Tensor{ElT,NR,StoreT,IndsR}
end

# Convert to Array, avoiding copying if possible
array(T::Tensor) = array(dense(T))
matrix(T::Tensor{<:Number,2}) = array(T)
vector(T::Tensor{<:Number,1}) = array(T)

#
# Helper functions for BlockSparse-type storage
#

"""
nzblocks(T::Tensor)

Return a vector of the non-zero blocks of the BlockSparseTensor.
"""
nzblocks(T::Tensor) = nzblocks(store(T))

blockoffsets(T::Tensor) = blockoffsets(store(T))
nnzblocks(T::Tensor) = nnzblocks(store(T))
nnz(T::Tensor) = nnz(store(T))
nblocks(T::Tensor) = nblocks(inds(T))
blockdims(T::Tensor,block) = blockdims(inds(T),block)
blockdim(T::Tensor,block) = blockdim(inds(T),block)

"""
offset(T::Tensor, block::Block)

Get the linear offset in the data storage for the specified block.
If the specified block is not non-zero structurally, return nothing.

offset(T::Tensor,pos::Int)

Get the offset of the block at position pos
in the block-offsets list.
"""
offset(T::Tensor, block) = offset(store(T), block)

nzblock(T::Tensor, n::Int) = nzblock(store(T), n)

"""
blockdim(T::Tensor,pos::Int)

Get the block dimension of the block at position pos
in the block-offset list.
"""
blockdim(T::Tensor, pos::Int) = blockdim(store(T), pos)

findblock(T::Tensor,
          block;
          sorted=true) = findblock(store(T), block; sorted=sorted)

"""
isblocknz(T::Tensor,
          block::Block)

Check if the specified block is non-zero
"""
isblocknz(T::Tensor,block) = isblocknz(store(T),block)

function blockstart(T::Tensor{<:Number,N},block) where {N}
  start_index = @MVector ones(Int,N)
  for j in 1:N
    ind_j = ind(T,j)
    for block_j in 1:block[j]-1
      start_index[j] += blockdim(ind_j,block_j)
    end
  end
  return Tuple(start_index)
end

function blockend(T::Tensor{<:Number,N},
                  block) where {N}
  end_index = @MVector zeros(Int,N)
  for j in 1:N
    ind_j = ind(T,j)
    for block_j in 1:block[j]
      end_index[j] += blockdim(ind_j,block_j)
    end
  end
  return Tuple(end_index)
end

"""
blockview(T::Tensor,block::Block)

Given a specified block, return a Dense/Diag Tensor that is a view to the data
in that block.
"""
function blockview(T::Tensor,block; sorted=true)
  pos = findblock(T,block; sorted=sorted)
  return blockview(T,pos)
end

#
# Some generic getindex and setindex! functionality
#

"""
getdiagindex

Get the specified value on the diagonal
"""
function getdiagindex(T::Tensor{<:Number,N},ind::Int) where {N}
  return getindex(T,CartesianIndex(ntuple(_->ind,Val(N))))
end

"""
setdiagindex!

Set the specified value on the diagonal
"""
function setdiagindex!(T::Tensor{<:Number,N},val,ind::Int) where {N}
  setindex!(T,val,CartesianIndex(ntuple(_->ind,Val(N))))
  return T
end

#
# Some generic contraction functionality
#

function zero_contraction_output(T1::TensorT1,
                                 T2::TensorT2,
                                 indsR::IndsR) where {TensorT1<:Tensor,
                                                      TensorT2<:Tensor,
                                                      IndsR}
  return zeros(contraction_output_type(TensorT1,TensorT2,IndsR),indsR)
end

#
# Broadcasting
#

Base.BroadcastStyle(::Type{T}) where {T<:Tensor} = Broadcast.ArrayStyle{T}()

function Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{T}},
                      ::Type{<:Any}) where {T<:Tensor}
  A = find_tensor(bc)
  return similar(A)
end

"`A = find_tensor(As)` returns the first Tensor among the arguments."
find_tensor(bc::Base.Broadcast.Broadcasted) = find_tensor(bc.args)
find_tensor(args::Tuple) = find_tensor(find_tensor(args[1]), Base.tail(args))
find_tensor(x) = x
find_tensor(a::Tensor, rest) = a
find_tensor(::Any, rest) = find_tensor(rest)

function Base.summary(io::IO,
                      T::Tensor)
  println(io,typeof(inds(T)))
  for (dim,ind) in enumerate(inds(T))
    println(io,"Dim $dim: ",ind)
  end
  println(io,typeof(store(T)))
  println(io," ",Base.dims2string(dims(T)))
end

#
# Printing
#

print_tensor(io::IO,T::Tensor) = Base.print_array(io,T)
print_tensor(io::IO,T::Tensor{<:Number,1}) = Base.print_array(io,reshape(T,(dim(T),1)))


text:
#=
  Original Author: Anastasia Yendiki

  Copyright © 2022 The General Hospital Corporation (Boston, MA) "MGH"

  Terms and conditions for use, reproduction, distribution and contribution
  are found in the 'FreeSurfer Software License Agreement' contained
  in the file 'LICENSE' found in the FreeSurfer distribution, and here:

  https://surfer.nmr.mgh.harvard.edu/fswiki/FreeSurferSoftwareLicense

  Reporting: freesurfer@nmr.mgh.harvard.edu
=#

module FreeSurfer

include("mri.jl")
include("trk.jl")
include("show.jl")
#include("view.jl")
include("util.jl")
include("dti.jl")
include("odf.jl")
include("gqi.jl")
include("rusd.jl")

function __init__()

  println("FREESURFER_HOME: " * (haskey(ENV, "FREESURFER_HOME") ?
                                 ENV["FREESURFER_HOME"] : "not defined"))
end

end # module

max_stars_repo_name: freesurfer/jupysurf

max_issues_repo_name: freesurfer/jupysurf

max_forks_repo_name: freesurfer/jupysurf


text:
"""
    FourierBasis([::Type,] num_inputs::Int, dorder::Int, iorder::Int [, full::Bool=false])

Creates a struct to generate fourier features up to a given order. Both coupled fourier features, e.g.,
``\\cos(π(3x₁ + 2x₂))`` and uncoupled features, ``[\\cos(πx₁), \\cos(π2x₁), …]`` can be generated with this basis function.
The dorder parameter controls order for the coupling features. The number of coupled features generated, ``(\\text{dorder}+1)^{\\text{num\\_inputs}}``,
 grows exponentially with dorder, so it is not reccomended for use with high deminsional vectors.
The iorder parameter controls the order of the independent features generated. The full parameter determines if both
``\\sin`` and ``\\cos`` features are generated, if false only cos features are generated.

See also: [`FourierBasisBuffer`](@ref)

# Examples
```jldoctest
julia> f = FourierBasis(2, 1, 2);

julia> x = [0.0, 0.5];

julia> feats = f(x)
6-element Array{Float64,1}:
  1.0
  6.123233995736766e-17
  1.0
  6.123233995736766e-17
  1.0
 -1.0

```
"""
struct FourierBasis{T,TF} <: Any where {T,TF}
    C::T

    function FourierBasis(::Type{T}, num_inputs::Int, dorder::Int, iorder::Int, both::Bool=false) where {T}
        C = make_Cmat(T, num_inputs, dorder, iorder)
        new{typeof(C),both}(C)
    end

    function FourierBasis(num_inputs::Int, dorder::Int, iorder::Int, both::Bool=false) where {T}
        return FourierBasis(Float64, num_inputs, dorder, iorder, both)
    end
end

"""
    FourierBasisBuffer(ϕ::FourierBasis)

Creates preallocated buffers for a fourier basis to use to avoid making allocations on every call of the basis.

See also: [`FourierBasis`](@ref)

# Examples
```jldoctest
julia> f = FourierBasis(2, 1, 2);

julia> buff = FourierBasisBuffer(f);

julia> x = [0.0, 0.5];

julia> feats = f(buff, x)
6-element Array{Float64,1}:
  1.0
  6.123233995736766e-17
  1.0
  6.123233995736766e-17
  1.0
 -1.0
"""
struct FourierBasisBuffer{T} <: Any
    y::T
    z::T

    function FourierBasisBuffer(ϕ::FourierBasis{TC,TB}) where {TC, TB}
        T = eltype(ϕ.C)
        n = size(ϕ.C, 2)
        y = zeros(T, n)
        if TB == true
            z = zeros(T, n*2)
        else
            z = zeros(T, n)
        end
        return new{typeof(y)}(y, z)
    end
end

function increment_counter!(counter::Vector{Int}, maxDigit::Int)
    for i in length(counter):-1:1
        counter[i] += 1
        if (counter[i] > maxDigit)
            counter[i] = 0
        else
            break
        end
    end
end

function make_Cmat(T::Type, num_inputs::Int, dorder::Int, iorder::Int)
    iTerms = iorder * num_inputs
    dTerms = (dorder+1)^num_inputs
    oTerms = min(iorder, dorder) * num_inputs
    num_feats = iTerms + dTerms - oTerms
    C = zeros(T, (num_inputs, num_feats))
    counter = zeros(Int, num_inputs)
    termCount::Int = 1

    while termCount <= dTerms
        for i in 1:num_inputs
            C[i, termCount] = counter[i]
        end
        increment_counter!(counter, dorder)
        termCount += 1
    end
    for i in 1:num_inputs
        for j in (dorder+1):(iorder)
            C[i, termCount] = j
            termCount += 1
        end
    end
    C .*= π

    return C
end

function (ϕ::FourierBasis{T,false})(x) where {T}
    y = zeros(size(ϕ.C, 2))
    mul!(y, ϕ.C', x)
    @. y = cos(y)
    return y
end

function (ϕ::FourierBasis{T,false})(buff::FourierBasisBuffer, x) where {T}
    y = buff.y
    mul!(y, ϕ.C', x)
    @. y = cos(y)
    return y
end

function (ϕ::FourierBasis{T,true})(x) where {T}
    n = size(ϕ.C, 2)
    y = zeros(n)
    z = zeros(n*2)
    mul!(y, ϕ.C', x)
    @. z[1:n] = cos(y)
    @. z[n+1:end] = sin(y)
    return z
end

function (ϕ::FourierBasis{T,true})(buff::FourierBasisBuffer, x) where {T}
    y, z = buff.y, buff.z
    n = length(y)
    mul!(y, ϕ.C', x)
    @. z[1:n] = cos(y)
    @. z[n+1:end] = sin(y)
    return z
end

"""
    length(ϕ::FourierBasis)

Returns the number of feautes produced by the Fourier basis.
"""
function length(ϕ::FourierBasis{T,false}) where {T}
    return size(ϕ.C, 2)
end

function length(ϕ::FourierBasis{T,true}) where {T}
    return size(ϕ.C, 2)*2
end


text:
module ImageIO

using FileIO: File, DataFormat, Stream, stream

## PNGs
const PNGFiles_LOADED = Ref(false)

function load(f::File{DataFormat{:PNG}}; kwargs...)
    if !PNGFiles_LOADED[]
        @eval ImageIO import PNGFiles
        PNGFiles_LOADED[] = isdefined(ImageIO, :PNGFiles)
    end
    return Base.invokelatest(PNGFiles.load, f.filename, kwargs...)
end
function load(s::Stream{DataFormat{:PNG}}; kwargs...)
    if !PNGFiles_LOADED[]
        @eval ImageIO import PNGFiles
        PNGFiles_LOADED[] = isdefined(ImageIO, :PNGFiles)
    end
    return Base.invokelatest(PNGFiles.load, stream(s), kwargs...)
end
function load(s::IO; kwargs...)
    if !PNGFiles_LOADED[]
        @eval ImageIO import PNGFiles
        PNGFiles_LOADED[] = isdefined(ImageIO, :PNGFiles)
    end
    return Base.invokelatest(PNGFiles.load, s, kwargs...)
end

function save(f::File{DataFormat{:PNG}}, image::S; kwargs...) where {T, S<:Union{AbstractMatrix, AbstractArray{T,3}}}
    if !PNGFiles_LOADED[]
        @eval ImageIO import PNGFiles
        PNGFiles_LOADED[] = isdefined(ImageIO, :PNGFiles)
    end
    return Base.invokelatest(PNGFiles.save, f.filename, image, kwargs...)
end
function save(s::Stream{DataFormat{:PNG}}, image::S; kwargs...) where {T, S<:Union{AbstractMatrix, AbstractArray{T,3}}}
    if !PNGFiles_LOADED[]
        @eval ImageIO import PNGFiles
        PNGFiles_LOADED[] = isdefined(ImageIO, :PNGFiles)
    end
    return Base.invokelatest(PNGFiles.save, stream(s), image, kwargs...)
end

end # module


text:
using StackCollections
using Test

@testset "DigitSet" begin
include("digitset.jl")
end

@testset "StackSet" begin
include("stackset.jl")
end

@testset "StackVector" begin
include("stackvector.jl")
end

@testset "OneHotVector" begin
include("onehotvector.jl")
end

@testset "Cross-type" begin
    @test hash(DigitSet([1, 6, 2])) != hash(StackSet([1, 6, 2]))
    @test hash(DigitSet([1, 6, 2])) == hash(DigitSet([1, 6, 2]))

    @test hash(StackVector([false, false, true])) != hash(OneHotVector([false, false, true]))
end


text:
# Formulation K2: (if regul==:classic, adds additional Regularization parmeters -ρ (top left) and δ (bottom right))
# [-Q_X - D      sqrt(X1X2)A' ] [̃x] = rhs
# [ A sqrt(X1x2)         0    ] [y]
# where Q_X = sqrt(X1X2) Q sqrt(X1X2) and D = s_l X2 + s_u X1
# and  Δ x = sqrt(X1 X2) Δ ̃x
export K2_5LDLParams

"""
Type to use the K2.5 formulation with a LDLᵀ factorization, using the package
[`LDLFactorizations.jl`](https://github.com/JuliaSmoothOptimizers/LDLFactorizations.jl).
The outer constructor

    sp = K2_5LDLParams(; regul :: Symbol = :classic)

creates a [`RipQP.SolverParams`](@ref) that should be used to create a [`RipQP.InputConfig`](@ref).
`regul = :dynamic` uses a dynamic regularization (the regularization is only added if the LDLᵀ factorization
encounters a pivot that has a small magnitude).
`regul = :none` uses no regularization (not recommended).
"""
struct K2_5LDLParams <: SolverParams
    regul :: Symbol
end

function K2_5LDLParams(; regul :: Symbol = :classic)
    regul == :classic || regul == :dynamic || regul == :none || error("regul should be :classic or :dynamic or :none")
    return K2_5LDLParams(regul)
end

mutable struct PreallocatedData_K2_5{T<:Real} <: PreallocatedData{T}
    D                :: Vector{T}                                        # temporary top-left diagonal
    regu             :: Regularization{T}
    diag_Q           :: SparseVector{T,Int} # Q diagonal
    K                :: SparseMatrixCSC{T,Int} # augmented matrix
    K_fact           :: LDLFactorizations.LDLFactorization{T,Int,Int,Int} # factorized matrix
    fact_fail        :: Bool # true if factorization failed
    diagind_K        :: Vector{Int} # diagonal indices of J
    K_scaled         :: Bool # true if K is scaled with X1X2
end

function PreallocatedData(sp :: K2_5LDLParams, fd :: QM_FloatData{T}, id :: QM_IntData,
                          iconf :: InputConfig{Tconf}) where {T<:Real, Tconf<:Real}
    # init Regularization values
    if iconf.mode == :mono
        regu = Regularization(T(sqrt(eps())*1e5), T(sqrt(eps())*1e5), 1e-5*sqrt(eps(T)), 1e0*sqrt(eps(T)), sp.regul)
        D = -T(1.0e0)/2 .* ones(T, id.nvar)
    else
        regu = Regularization(T(sqrt(eps())*1e5), T(sqrt(eps())*1e5), T(sqrt(eps(T))*1e0), T(sqrt(eps(T))*1e0), sp.regul)
        D = -T(1.0e-2) .* ones(T, id.nvar)
    end
    diag_Q = get_diag_Q(fd.Q.colptr, fd.Q.rowval, fd.Q.nzval, id.nvar)
    K = create_K2(id, D, fd.Q, fd.AT, diag_Q, regu)

    diagind_K = get_diag_sparseCSC(K.colptr, id.ncon+id.nvar)
    K_fact = ldl_analyze(Symmetric(K, :U))
    if regu.regul == :dynamic
        Amax = @views norm(K.nzval[diagind_K], Inf)
        regu.ρ, regu.δ = -T(eps(T)^(3/4)), T(eps(T)^(0.45))
        K_fact.r1, K_fact.r2 = regu.ρ, regu.δ
        K_fact.tol = Amax*T(eps(T))
        K_fact.n_d = id.nvar
    elseif regu.regul == :none
        regu.ρ, regu.δ = zero(T), zero(T)
    end
    K_fact = ldl_factorize!(Symmetric(K, :U), K_fact)
    K_fact.__factorized = true

    return PreallocatedData_K2_5(D,
                                 regu,
                                 diag_Q, #diag_Q
                                 K, #K
                                 K_fact, #K_fact
                                 false,
                                 diagind_K, #diagind_K
                                 false,
                                 )
end

function convertpad(::Type{<:PreallocatedData{T}}, pad :: PreallocatedData_K2_5{T_old}, T0 :: DataType) where {T<:Real, T_old<:Real}
    pad =  PreallocatedData_K2_5(convert(Array{T}, pad.D),
                                 convert(Regularization{T}, pad.regu),
                                 convert(SparseVector{T,Int}, pad.diag_Q),
                                 convert(SparseMatrixCSC{T,Int}, pad.K),
                                 convertldl(T, pad.K_fact),
                                 pad.fact_fail,
                                 pad.diagind_K,
                                 pad.K_scaled
                                 )

    if pad.regu.regul == :classic
        if T == Float64 && T0 == Float64
            pad.regu.ρ_min, pad.regu.δ_min = T(sqrt(eps())*1e-5), T(sqrt(eps())*1e0)
        else
            pad.regu.ρ_min, pad.regu.δ_min = T(sqrt(eps(T))*1e1), T(sqrt(eps(T))*1e1)
        end
        pad.regu.ρ /= 10
        pad.regu.δ /= 10
    elseif pad.regu.regul == :dynamic
        pad.regu.ρ, pad.regu.δ = -T(eps(T)^(3/4)), T(eps(T)^(0.45))
        pad.K_fact.r1, pad.K_fact.r2 = pad.regu.ρ, pad.regu.δ
    end

    return pad
end

# solver LDLFactorization
function solver!(pad :: PreallocatedData_K2_5{T}, dda :: DescentDirectionAllocs{T}, pt :: Point{T}, itd :: IterData{T},
                 fd :: Abstract_QM_FloatData{T}, id :: QM_IntData, res :: Residuals{T}, cnts :: Counters, T0 :: DataType,
                 step :: Symbol) where {T<:Real}

    if step == :aff
        dda.Δxy_aff[1:id.nvar] .*= pad.D
        LDLFactorizations.ldiv!(pad.K_fact, dda.Δxy_aff)
        dda.Δxy_aff[1:id.nvar] .*= pad.D
    else
        if pad.K_scaled
            itd.Δxy[1:id.nvar] .*= pad.D
            LDLFactorizations.ldiv!(pad.K_fact, itd.Δxy)
            itd.Δxy[1:id.nvar] .*= pad.D
        else
            LDLFactorizations.ldiv!(pad.K_fact, itd.Δxy)
        end
    end

    if step == :cc || step == :IPF  # update regularization and restore K. Cannot be done in update_pad since x-lvar and uvar-x will change.
        out = 0
        if pad.regu.regul == :classic # update ρ and δ values, check K diag magnitude
            out = update_regu_diagK2_5!(pad.regu, pad.D, itd.pdd, itd.l_pdd, itd.mean_pdd, cnts, T, T0)
        end

        # restore J for next iteration
        if pad.K_scaled
            pad.D .= one(T)
            pad.D[id.ilow] ./= sqrt.(itd.x_m_lvar)
            pad.D[id.iupp] ./= sqrt.(itd.uvar_m_x)
            lrmultilply_J!(pad.K.colptr, pad.K.rowval, pad.K.nzval, pad.D, id.nvar)
            pad.K_scaled = false
        end
        out == 1 && return out
    end
    return 0
end

function update_pad!(pad :: PreallocatedData_K2_5{T}, dda :: DescentDirectionAllocs{T}, pt :: Point{T}, itd :: IterData{T},
                     fd :: Abstract_QM_FloatData{T}, id :: QM_IntData, res :: Residuals{T}, cnts :: Counters, T0 :: DataType) where {T<:Real}

    pad.K_scaled = false
    out = factorize_K2_5!(pad.K, pad.K_fact, pad.D, pad.diag_Q, pad.diagind_K, pad.regu,
                          pt.s_l, pt.s_u, itd.x_m_lvar, itd.uvar_m_x, id.ilow, id.iupp,
                          id.ncon, id.nvar, cnts, itd.qp, T, T0)
    out == 1 && return out
    pad.K_scaled = true

    return out
end

function lrmultilply_J!(J_colptr, J_rowval, J_nzval, v, nvar)
    T = eltype(v)
    @inbounds @simd for i=1:nvar
        for idx_row=J_colptr[i]: J_colptr[i+1]-1
                J_nzval[idx_row] *= v[i] * v[J_rowval[idx_row]]
        end
    end

    n = length(J_colptr)
    @inbounds @simd for i=nvar+1:n-1
        for idx_row= J_colptr[i]: J_colptr[i+1] - 1
            if J_rowval[idx_row] <= nvar
                J_nzval[idx_row] *= v[J_rowval[idx_row]] # multiply row i by v[i]
            end
        end
    end
end

# function lrmultiply_J2!(J, v)
#     lmul!(v, J)
#     rmul!(J, v)
# end

# iteration functions for the K2.5 system
function factorize_K2_5!(K, K_fact, D, diag_Q , diagind_K, regu, s_l, s_u, x_m_lvar, uvar_m_x,
                         ilow, iupp, ncon, nvar, cnts, qp, T, T0)

    if regu.regul == :classic
        D .= -regu.ρ
        K.nzval[view(diagind_K, nvar+1:ncon+nvar)] .= regu.δ
    else
        D .= zero(T)
    end
    D[ilow] .-= s_l ./ x_m_lvar
    D[iupp] .-= s_u ./ uvar_m_x
    D[diag_Q.nzind] .-= diag_Q.nzval
    K.nzval[view(diagind_K,1:nvar)] = D

    D .= one(T)
    D[ilow] .*= sqrt.(x_m_lvar)
    D[iupp] .*= sqrt.(uvar_m_x)
    lrmultilply_J!(K.colptr, K.rowval, K.nzval, D, nvar)

    if regu.regul == :dynamic
        # Amax = @views norm(K.nzval[diagind_K], Inf)
        Amax = minimum(D)
        if Amax < sqrt(eps(T)) && cnts.c_pdd < 8
            if T == Float32
                # restore J for next iteration
                D .= one(T)
                D[ilow] ./= sqrt.(x_m_lvar)
                D[iupp] ./= sqrt.(uvar_m_x)
                lrmultilply_J!(K.colptr, K.rowval, K.nzval, D, nvar)
                return one(Int) # update to Float64
            elseif qp || cnts.c_pdd < 4
                cnts.c_pdd += 1
                regu.δ /= 10
                K_fact.r2 = max(sqrt(Amax), regu.δ)
                # regu.ρ /= 10
            end
        end
        K_fact.tol = min(Amax, T(eps(T)))
        K_fact = ldl_factorize!(Symmetric(K, :U), K_fact)

    elseif regu.regul == :classic
        ldl_factorize!(Symmetric(K, :U), K_fact)
        while !factorized(K_fact)
            out = update_regu_trycatch!(regu, cnts, T, T0)
            if out == 1
                # restore J for next iteration
                D .= one(T)
                D[ilow] ./= sqrt.(x_m_lvar)
                D[iupp] ./= sqrt.(uvar_m_x)
                lrmultilply_J!(K.colptr, K.rowval, K.nzval, D, nvar)
                return out
            end
            cnts.c_catch += 1
            cnts.c_catch >= 4 && return 1
            D .= -regu.ρ
            D[ilow] .-= s_l ./ x_m_lvar
            D[iupp] .-= s_u ./ uvar_m_x
            D[diag_Q.nzind] .-= diag_Q.nzval
            K.nzval[view(diagind_K,1:nvar)] = D

            D .= one(T)
            D[ilow] .*= sqrt.(x_m_lvar)
            D[iupp] .*= sqrt.(uvar_m_x)
            K.nzval[view(diagind_K,1:nvar)] .*= D.^2
            K.nzval[view(diagind_K, nvar+1:ncon+nvar)] .= regu.δ
            ldl_factorize!(Symmetric(K, :U), K_fact)
        end

    else # no Regularization
        K_fact = ldl_factorize!(Symmetric(K, :U), K_fact)
    end

    return 0 # factorization succeeded
end


text:
@composed type LineAttributes
    Model::Mat4f0
    Color::Vec4f0
    Thickness::Float32
end

@composed type LineSegments
    Vertices::AbstractArray{<: AbstractLineVertex}
    <: LineAttributes
end

immutable Vert2Geom
    position::Vec4f0
    color::Vec4f0
    thickness::Float32
end

immutable Geom2Fragment
    thickness::Float32
    color::Vec4f0
    uv::Vec2f0
end

function screen_space(vertex::Vec4f0, scene)
    return (vertex[Vec(1, 2)] / vertex[4]) .* scene.resolution
end

function vert_linesegments(vertex, scene, uniforms)
    m = scene.projectionview * uniforms.model
    geomout = Vert2Geom(
        m * to_vec4(getposition(vertex)),
        getcolor(vertex, uniforms),
        getthickness(vertex, uniforms)
    )
    return geomout
end

function emit_line_vertex(emit!, geom_in, scene, uniforms, position::Vec2f0, uv::Vec2f0, index)
    inpos = geom_in[index].position
    x = (position ./ scene.resolution) * inpos[4]
    outpos = Vec4f0(x[1], x[2], inpos[3], inpos[4])
    fragout = Geom2Fragment(
        geom_in[index].thickness + 2f0,
        geom_in[index].color,
        uv,
    )
    emit!(outpos, fragout)
    return
end

function geom_linesegments(emit!, vertex_out, scene, uniforms)
    # get the four vertices passed to the shader:
    p0 = screen_space(vertex_out[1].position, scene) # start of previous segment
    p1 = screen_space(vertex_out[2].position, scene) # end of previous segment, start of current segment

    thickness0 = vertex_out[1].thickness
    thickness1 = vertex_out[2].thickness
    thickness_aa0 = thickness0 + 2f0
    thickness_aa1 = thickness1 + 2f0

    # determine the direction of each of the 3 segments (previous, current, next)
    vun0 = p1 - p0
    v0 = normalize(vun0)
    # determine the normal of each of the 3 segments (previous, current, next)
    n0 = Vec2f0(-v0[2], v0[1])
    l = norm(p1 - p0)

    uv0 = thickness_aa0 / thickness0
    uv1 = thickness_aa1 / thickness_aa1
    emit_line_vertex(emit!, vertex_out, scene, uniforms, p0 + thickness_aa0 * n0, Vec2f0(0f0, -uv0), 1)
    emit_line_vertex(emit!, vertex_out, scene, uniforms, p0 - thickness_aa0 * n0, Vec2f0(0f0, uv0), 1)
    emit_line_vertex(emit!, vertex_out, scene, uniforms, p1 + thickness_aa1 * n0, Vec2f0(l, -uv1), 2)
    emit_line_vertex(emit!, vertex_out, scene, uniforms, p1 - thickness_aa1 * n0, Vec2f0(l, uv1), 2)
    return
end

function frag_linesegments(geom_out, scene, uniforms)
    uv = geom_out.uv; color = geom_out.color
    xy = Vec2f0(0.5f0, uv[2])
    alpha = aastep(0f0, xy[1])
    alpha2 = aastep(-1f0, 1f0, xy[2])
    outcolor = Vec4f0(color[1], color[2], color[3], color[4] * alpha * alpha2)
    (outcolor, )
end

broadcast_conv(x::Vector) = x
broadcast_conv(x::Number) = x
broadcast_conv(x::StaticVector) = StaticArrays.Scalar(x)

function default(::Type{Vertices}, p::Partial{LineSegments})
    positions = get!(p, Position)::Vector{<: StaticVector}
    #(isa(positions, Vector) && eltype(positions) <: SVector) || error("Must be vector of static vectors")
    color = get!(p, Color)::Union{<: StaticVector, Vector{<: StaticVector,}}
    thickness = get!(p, Thickness)::Union{<: AbstractFloat, Vector{<: AbstractFloat,}}
    args = broadcast_conv.((positions, thickness, color))
    # TODO sort args into uniforms (if scalar) or vertex array (if array)
    # A bit tricky, since that means we need to figure out vertex & uniform type in a semi
    # dynamic way
    LineVertex.(args...)
end

function Drawable(window::AbstractWindow, primitive::LineSegments)
    verts = primitive[Vertices]
    vbo = reinterpret(NTuple{2, eltype(verts)}, verts)
    uniforms = LineAttributes(primitive)
    args = (window[Scene], uniforms)
    rasterizer(
        window,
        vbo, args,
        vert_linesegments, frag_linesegments;
        geometryshader = geom_linesegments,
        primitive_in = :lines
    )
end


text:
"""
Atomic solution for an PDAE.

### Parameters

* `DT`: data type
* `TT`: time step type
* `AT`: array type
* `IT`: internal variable types

### Fields

* `t`: time of current time step
* `t̄`: time of previous time step
* `q`: current solution of q
* `q̄`: previous solution of q
* `q̃`: compensated summation error of q
* `p`: current solution of p
* `p̄`: previous solution of p
* `p̃`: compensated summation error of p
* `λ`: current solution of λ
* `λ̄`: previous solution of λ
* `v`: vector field of q
* `v̄`: vector field of q̄
* `f`: vector field of p
* `f̄`: vector field of p̄
* `u`: projective vector field of q
* `ū`: projective vector field of q̄
* `g`: projective vector field of p
* `ḡ`: projective vector field of p̄
* `internal`: internal variables of the integrator (e.g., internal stages of a Runge-Kutta methods or solver output)

### Constructors

```julia
AtomicSolutionPDAE{DT,TT,AT,IT}(nd, nm, internal::IT)
AtomicSolutionPDAE{DT,TT,AT,IT}(t::TT, q::AT, p::AT, λ::AT, internal::IT)
AtomicSolutionPDAE(DT, TT, AT, nd, nm, internal::IT=NamedTuple())
AtomicSolutionPDAE(t::TT, q::AT, p::AT, λ::AT, internal::IT=NamedTuple())
```

* `nd`: dimension of the state vector
* `nm`: dimension of the constraint

"""
mutable struct AtomicSolutionPDAE{DT <: Number, TT <: Real, AT <: AbstractArray{DT}, IT <: NamedTuple} <: AtomicSolution{DT,TT,AT}
    t::TT
    t̄::TT

    q::AT
    q̄::AT
    q̃::AT

    p::AT
    p̄::AT
    p̃::AT

    λ::AT
    λ̄::AT

    v::AT
    v̄::AT
    f::AT
    f̄::AT

    u::AT
    ū::AT
    g::AT
    ḡ::AT

    internal::IT

    function AtomicSolutionPDAE{DT,TT,AT,IT}(nd, nm, internal::IT) where {DT,TT,AT,IT}
        new(zero(TT), zero(TT),
            AT(zeros(DT, nd)), AT(zeros(DT, nd)), AT(zeros(DT, nd)),
            AT(zeros(DT, nd)), AT(zeros(DT, nd)), AT(zeros(DT, nd)),
            AT(zeros(DT, nm)), AT(zeros(DT, nm)),
            AT(zeros(DT, nd)), AT(zeros(DT, nd)), AT(zeros(DT, nd)), AT(zeros(DT, nd)),
            AT(zeros(DT, nd)), AT(zeros(DT, nd)), AT(zeros(DT, nd)), AT(zeros(DT, nd)),
            internal)
    end

    function AtomicSolutionPDAE{DT,TT,AT,IT}(t::TT, q::AT, p::AT, λ::AT, internal::IT) where {DT,TT,AT,IT}
        new(zero(t), zero(t),
            zero(q), zero(q), zero(q),
            zero(p), zero(p), zero(p),
            zero(λ), zero(λ),
            zero(q), zero(q), zero(p), zero(p),
            zero(λ), zero(λ), zero(λ), zero(λ),
            internal)
    end
end

AtomicSolutionPDAE(DT, TT, AT, nd, nm, internal::IT=NamedTuple()) where {IT} = AtomicSolutionPDAE{DT,TT,AT,IT}(nd, nm, internal)
AtomicSolutionPDAE(t::TT, q::AT, p::AT, λ::AT, internal::IT=NamedTuple()) where {DT, TT, AT <: AbstractArray{DT}, IT} = AtomicSolutionPDAE{DT,TT,AT,IT}(t, q, p, λ, internal)

function set_initial_conditions!(asol::AtomicSolutionPDAE, equ::AbstractEquationPDAE, i::Int=1)
    @assert i ≥ nsamples(equ)
    t, q, p, λ = initial_conditions(equ)
    asol.t  = t
    asol.q .= q[i]
    asol.p .= p[i]
    asol.λ .= λ[i]
    asol.v .= 0
    asol.f .= 0
end

function set_solution!(asol::AtomicSolutionPDAE, sol)
    t, q, p, λ = sol
    asol.t  = t
    asol.q .= q
    asol.p .= p
    asol.λ .= λ
    asol.v .= 0
    asol.f .= 0
end

function get_solution(asol::AtomicSolutionPDAE)
    (asol.t, asol.q, asol.p, asol.λ)
end

function GeometricBase.reset!(asol::AtomicSolutionPDAE, Δt)
    asol.t̄  = asol.t
    asol.q̄ .= asol.q
    asol.p̄ .= asol.p
    asol.λ̄.= asol.λ
    asol.v̄ .= asol.v
    asol.f̄ .= asol.f
    asol.ū .= asol.u
    asol.ḡ .= asol.g
    asol.t += Δt
end

function update!(asol::AtomicSolutionPDAE{DT}, y::Vector{DT}, z::Vector{DT}, λ::Vector{DT}) where {DT}
    for k in eachindex(y,z)
        update!(asol, y[k], z[k], k)
    end
    for k in eachindex(λ)
        asol.λ[k] = λ[k]
    end
end

function update!(asol::AtomicSolutionPDAE{DT}, y::DT, z::DT, k::Union{Int,CartesianIndex}) where {DT}
    asol.q[k], asol.q̃[k] = compensated_summation(y, asol.q[k], asol.q̃[k])
    asol.p[k], asol.p̃[k] = compensated_summation(z, asol.p[k], asol.p̃[k])
end