text:
# We first define a "Flat" projection where the metric in the
# resulting cartesian space is assumed to be "flat" / "euclidean" and
# convolution operators are diagonalized by the 2D FFT. This is also
# known as the "flat-sky approximation". However, we don't actually
# define any functionality for it, this is mostly a placeholder for
# future extensions, and also provides backwards compatibility with
# previous use of FlatMap, etc.. in CMBLensing.jl. The only such
# projection we have implemented is ProjLambert, and everything in
# this file is actually defined on LambertFields, not FlatFields.
# However, with some thought, most of what is in here could be
# generalized to other fields.
abstract type FlatProj <: CartesianProj end
make_field_aliases("Flat", FlatProj)
default_proj(::Type{F}) where {F<:BaseField{<:Any,<:FlatProj}} = ProjLambert

# A Lambert azimuthal equal-area projection
#
# The `rotator` field is an argument passed to `healpy.Rotator`
# specifying a rotation which rotates the north pole to the
# center of the desired field.
#
struct ProjLambert{T, V<:AbstractVector{T}, M<:AbstractMatrix{T}} <: FlatProj
    # these must be the same to broadcast together
    Ny        :: Int
    Nx        :: Int
    θpix      :: Float64
    rotator   :: NTuple{3,Float64}
    # these can be different and still broadcast (including different types)
    storage
    Δx        :: T
    Ωpix      :: T
    nyquist   :: T
    Δℓx       :: T
    Δℓy       :: T
    ℓy        :: V
    ℓx        :: V
    ℓmag      :: M
    sin2ϕ     :: M
    cos2ϕ     :: M
end

ProjLambert(;Ny, Nx, θpix=1, rotator=(0,90,0), T=Float32, storage=Array) =
    ProjLambert(Ny, Nx, Float64(θpix), Float64.(rotator), real_type(T), storage)

@memoize function ProjLambert(Ny, Nx, θpix, rotator, ::Type{T}, storage) where {T}

    Δx           = T(deg2rad(θpix/60))
    Δℓx          = T(2π/(Nx*Δx))
    Δℓy          = T(2π/(Ny*Δx))
    nyquist      = T(2π/(2Δx))
    Ωpix         = T(Δx^2)
    ℓy           = adapt(storage, (ifftshift(-Ny÷2:(Ny-1)÷2) .* Δℓy)[1:Ny÷2+1])
    ℓx           = adapt(storage, (ifftshift(-Nx÷2:(Nx-1)÷2) .* Δℓx))
    ℓmag         = @. sqrt(ℓx'^2 + ℓy^2)
    ϕ            = @. angle(ℓx' + im*ℓy)
    sin2ϕ, cos2ϕ = @. sin(2ϕ), cos(2ϕ)
    if iseven(Ny)
        sin2ϕ[end, end:-1:(Nx÷2+2)] .= sin2ϕ[end, 2:Nx÷2]
    end

    ProjLambert(Ny,Nx,θpix,rotator,storage,Δx,Ωpix,nyquist,Δℓx,Δℓy,ℓy,ℓx,ℓmag,sin2ϕ,cos2ϕ)

end

# make LambertMap, LambertFourier, etc... type aliases
make_field_aliases("Lambert", ProjLambert)

# for printing
typealias_def(::Type{F}) where {B,M<:ProjLambert,T,A,F<:LambertField{B,M,T,A}} = "Lambert$(typealias(B)){$(typealias(A))}"
function Base.summary(io::IO, f::LambertField)
    @unpack Nx,Ny,Nbatch,θpix = f
    print(io, "$(length(f))-element $Ny×$Nx$(Nbatch==1 ? "" : "(×$Nbatch)")-pixel $(θpix)′-resolution ")
    Base.showarg(io, f, true)
end

### promotion

# used in broadcasting to decide the resulting metadata when
# broadcasting over two fields
function promote_metadata_strict(metadata₁::ProjLambert{T₁}, metadata₂::ProjLambert{T₂} ) where {T₁,T₂}

    if (
        metadata₁.θpix === metadata₂.θpix &&
        metadata₁.Ny   === metadata₂.Ny   &&
        metadata₁.Nx   === metadata₂.Nx
    )

        # always returning the "wider" metadata even if T₁==T₂ helps
        # inference and is optimized away anyway
        promote_type(T₁,T₂) == T₁ ? metadata₁ : metadata₂

    else
        error("""Can't broadcast two fields with the following differing metadata:
        1: $(select(fields(metadata₁),(:θpix,:Ny,:Nx)))
        2: $(select(fields(metadata₂),(:θpix,:Ny,:Nx)))
        """)
    end

end

# used in non-broadcasted algebra to decide the resulting metadata
# when performing some operation across two fields. this is free to do
# more generic promotion than promote_metadata_strict (although this
# is currently not used, but in the future could include promoting
# resolution, etc...). the result should be a common metadata which we
# can convert both fields to then do a succesful broadcast
promote_metadata_generic(metadata₁::ProjLambert, metadata₂::ProjLambert) =
    promote_metadata_strict(metadata₁, metadata₂)

### preprocessing
# defines how ImplicitFields and BatchedReals behave when broadcasted
# with ProjLambert fields. these can return arrays, but can also
# return `Broadcasted` objects which are spliced into the final
# broadcast, thus avoiding allocating any temporary arrays.

function preprocess((_,proj)::Tuple{<:Any,<:ProjLambert{T,V}}, r::Real) where {T,V}
    r isa BatchedReal ? adapt(V, reshape(r.vals, 1, 1, 1, :)) : r
end
# need custom adjoint here bc Δ can come back batched from the
# backward pass even though r was not batched on the forward pass
@adjoint function preprocess(m::Tuple{<:Any,<:ProjLambert{T,V}}, r::Real) where {T,V}
    preprocess(m, r), Δ -> (nothing, Δ isa AbstractArray ? batch(real.(Δ[:])) : Δ)
end

function preprocess((_,proj)::Tuple{BaseFieldStyle{S,B},<:ProjLambert}, ∇d::∇diag) where {S,B}

    (B <: Union{Fourier,QUFourier,IQUFourier}) ||
        error("Can't broadcast ∇[$(∇d.coord)] as a $(typealias(B)), its not diagonal in this basis.")

    # turn both into 2D matrices so this function is type-stable
    # (reshape doesnt actually make a copy here, so this doesn't
    # impact performance)
    if ∇d.coord == 1
        broadcasted(*, ∇d.prefactor * im, reshape(proj.ℓx, 1, :))
    else
        broadcasted(*, ∇d.prefactor * im, reshape(proj.ℓy, :, 1))
    end
end

function preprocess((_,proj)::Tuple{BaseFieldStyle{S,B},<:ProjLambert}, ::∇²diag) where {S,B}

    (B <: Union{Fourier,<:Basis2Prod{<:Any,Fourier},<:Basis3Prod{<:Any,<:Any,Fourier}}) ||
        error("Can't broadcast a BandPass as a $(typealias(B)), its not diagonal in this basis.")

    broadcasted(+, broadcasted(^, proj.ℓx', 2), broadcasted(^, proj.ℓy, 2))
end

function preprocess((_,proj)::Tuple{<:Any,<:ProjLambert}, bp::BandPass)
    Cℓ_to_2D(bp.Wℓ, proj)
end

function Cℓ_to_2D(Cℓ, proj::ProjLambert{T}) where {T}
    Complex{T}.(nan2zero.(Cℓ.(proj.ℓmag)))
end

### adapting

# dont adapt the fields in proj, instead re-call into the memoized
# ProjLambert so we always get back the singleton ProjLambert object
# for the given set of parameters (helps reduce memory usage and
# speed-up subsequent broadcasts which would otherwise not hit the
# "===" branch of the "promote_*" methods)
function adapt_structure(storage, proj::ProjLambert{T}) where {T}
    @unpack Ny, Nx, θpix = proj
    T′ = eltype(storage)
    ProjLambert(;Ny, Nx, θpix, T=(T′==Any ? T : real(T′)), storage)
end
adapt_structure(::Nothing, proj::ProjLambert{T}) where {T} = proj

@doc doc"""
    pixwin(θpix, ℓ)

Returns the pixel window function for square flat-sky pixels of width `θpix` (in
arcmin) evaluated at some `ℓ`s. This is the scaling of k-modes, the scaling of
the power spectrum will be pixwin^2.
"""
pixwin(θpix, ℓ) = @. sinc(ℓ*deg2rad(θpix/60)/2π)

### serialization
# makes it so the arrays in ProjLambert objects aren't actually
# serialized, instead just (Ny, Nx, θpix, rotator, T) are stored, and
# deserializing just recreates the ProjLambert object, possibly from
# the memoized cache if it already exists.

function _serialization_key(proj::ProjLambert{T}) where {T}
    @unpack Ny, Nx, θpix, rotator, storage = proj
    (;Ny, Nx, θpix, rotator, T, storage)
end

# Julia serialization
function Serialization.serialize(s::AbstractSerializer, proj::ProjLambert)
    @unpack Ny, Nx, θpix, rotator = proj
    Serialization.writetag(s.io, Serialization.OBJECT_TAG)
    Serialization.serialize(s, ProjLambert)
    Serialization.serialize(s, _serialization_key(proj))
end
function Serialization.deserialize(s::AbstractSerializer, ::Type{ProjLambert})
    ProjLambert(; Serialization.deserialize(s)...)
end

# JLD2 serialization
# (always deserialize as Array)
function JLD2.writeas(::Type{<:ProjLambert})
    Tuple{Val{ProjLambert},NamedTuple{(:Ny,:Nx,:θpix,:rotator,:T),Tuple{Int,Int,Float64,NTuple{3,Float64},DataType}}}
end
function JLD2.wconvert(::Type{<:Tuple{Val{ProjLambert},NamedTuple}}, proj::ProjLambert)
    (Val(ProjLambert), delete(_serialization_key(proj), :storage))
end
function JLD2.rconvert(::Type{<:ProjLambert}, (_,s)::Tuple{Val{ProjLambert},NamedTuple})
    ProjLambert(; storage=Array, s...)
end

### basis conversion
## spin-0
Fourier(f::LambertMap) = LambertFourier(m_rfft(f.arr, (1,2)), f.metadata)
Fourier(f′::LambertFourier, f::LambertMap) = (m_rfft!(f′.arr, f.arr, (1,2)); f′)
Map(f::LambertFourier) = LambertMap(m_irfft(f.arr, f.Ny, (1,2)), f.metadata)
Map(f′::LambertMap, f::LambertFourier) = (m_irfft!(f′.arr, f.arr, (1,2)); f′)

## spin-2
QUFourier(f::LambertQUMap) = LambertQUFourier(m_rfft(f.arr, (1,2)), f.metadata)
QUFourier(f::LambertEBMap) = f |> EBFourier |> QUFourier
QUFourier(f::LambertEBFourier) = begin
    @unpack El, Bl, sin2ϕ, cos2ϕ = fieldinfo(f)
    Ql = @. - El * cos2ϕ + Bl * sin2ϕ
    Ul = @. - El * sin2ϕ - Bl * cos2ϕ
    LambertQUFourier(Ql, Ul, f.metadata)
end

QUMap(f::LambertQUFourier) = LambertQUMap(m_irfft(f.arr, f.Ny, (1,2)), f.metadata)
QUMap(f::LambertEBMap)      = f |> EBFourier |> QUFourier |> QUMap
QUMap(f::LambertEBFourier)  = f |> QUFourier |> QUMap

EBFourier(f::LambertEBMap) = LambertEBFourier(m_rfft(f.arr, (1,2)), f.metadata)
EBFourier(f::LambertQUMap) = f |> QUFourier |> EBFourier
EBFourier(f::LambertQUFourier) = begin
    @unpack Ql, Ul, sin2ϕ, cos2ϕ = fieldinfo(f)
    El = @. - Ql * cos2ϕ - Ul * sin2ϕ
    Bl = @.   Ql * sin2ϕ - Ul * cos2ϕ
    LambertEBFourier(El, Bl, f.metadata)
end

EBMap(f::LambertEBFourier) = LambertEBMap(m_irfft(f.arr, f.Ny, (1,2)), f.metadata)
EBMap(f::LambertQUMap)     = f |> QUFourier |> EBFourier |> EBMap
EBMap(f::LambertQUFourier) = f |> EBFourier |> EBMap

# in-place
QUMap(f′::LambertQUMap, f::LambertQUFourier) = (m_irfft!(f′.arr, f.arr, (1,2)); f′)
QUFourier(f′::LambertQUFourier, f::LambertQUMap) = (m_rfft!(f′.arr, f.arr, (1,2)); f′)

## spin-(0,2)
IQUFourier(f::LambertIQUMap) = LambertIQUFourier(m_rfft(f.arr, (1,2)), f.metadata)
IQUFourier(f::LambertIEBMap) = f |> IEBFourier |> IQUFourier
IQUFourier(f::LambertIEBFourier) = LambertIQUFourier(f.I, QUFourier(f.P))

IQUMap(f::LambertIQUFourier) = LambertIQUMap(m_irfft(f.arr, f.Ny, (1,2)), f.metadata)
IQUMap(f::LambertIEBMap)      = f |> IEBFourier |> IQUFourier |> IQUMap
IQUMap(f::LambertIEBFourier)  = f |> IQUFourier |> IQUMap

IEBFourier(f::LambertIEBMap) = LambertIEBFourier(m_rfft(f.arr, (1,2)), f.metadata)
IEBFourier(f::LambertIQUMap) = f |> IQUFourier |> IEBFourier
IEBFourier(f::LambertIQUFourier) = LambertIEBFourier(f.I, EBFourier(f.P))

IEBMap(f::LambertIEBFourier) = LambertIEBMap(m_irfft(f.arr, f.Ny, (1,2)), f.metadata)
IEBMap(f::LambertIQUMap)     = f |> IQUFourier |> IEBFourier |> IEBMap
IEBMap(f::LambertIQUFourier) = f |> IEBFourier |> IEBMap

# in-place
IQUMap(f′::LambertIQUMap, f::LambertIQUFourier) = (m_irfft!(f′.arr, f.arr, (1,2)); f′)
IQUFourier(f′::LambertIQUFourier, f::LambertIQUMap) = (m_rfft!(f′.arr, f.arr, (1,2)); f′)

## for diagonal operator
function getindex(D::DiagOp{<:LambertEBFourier}, k::Symbol)
    @unpack El, Bl, metadata = diag(D)
    @unpack sin2ϕ, cos2ϕ = fieldinfo(diag(D))
    f = @match k begin
        (:QQ)        => LambertFourier((@. Bl*sin2ϕ^2 + El*cos2ϕ^2),   metadata)
        (:QU || :UQ) => LambertFourier((@. (El - Bl) * sin2ϕ * cos2ϕ), metadata)
        (:UU)        => LambertFourier((@. Bl*cos2ϕ^2 + El*sin2ϕ^2),   metadata)
        _            => getproperty(D.diag, k)
    end
    Diagonal(f)
end

### dot products
# do in Map space (the LenseBasis, Ł) for simplicity
function dot(a::LambertField, b::LambertField)
    z = Ł(a) .* Ł(b)
    batch(sum_dropdims(z.arr, dims=nonbatch_dims(z)))
end

### logdets

function logdet(L::Diagonal{<:Union{Real,Complex},<:LambertField{B}}) where {B<:Union{Fourier,Basis2Prod{<:Any,Fourier},Basis3Prod{<:Any,<:Any,Fourier}}}
    # half the Fourier plane needs to be counted twice since the real
    # FFT only stores half of it
    @unpack Ny, arr = L.diag
    λ = adapt(typeof(arr), rfft_degeneracy_fac(Ny))
    # note: since our maps are required to be real, the logdet of any
    # operator which preserves this property is also guaranteed to be
    # real, hence the `real` and `abs` below are valid
    batch(real.(sum_dropdims(nan2zero.(log.(abs.(arr)) .* λ), dims=nonbatch_dims(L.diag))))
end

function logdet(L::Diagonal{<:Real,<:LambertField{B}}) where {B<:Union{Map,Basis2Prod{<:Any,Map},Basis3Prod{<:Any,<:Any,Map}}}
    batch(
        sum_dropdims(log.(abs.(L.diag.arr)), dims=nonbatch_dims(L.diag))
        .+ dropdims(log.(prod(sign.(L.diag.arr), dims=nonbatch_dims(L.diag))), dims=nonbatch_dims(L.diag))
    )
end

### traces

function tr(L::Diagonal{<:Union{Real,Complex},<:LambertField{B}}) where {B<:Union{Fourier,Basis2Prod{<:Any,Fourier},Basis3Prod{<:Any,<:Any,Fourier}}}
    @unpack Ny, Nx, arr = L.diag
    λ = adapt(typeof(arr), rfft_degeneracy_fac(Ny))
    # the `real` is ok bc the imaginary parts of the half-plane which
    # is stored would cancel with those from the other half-plane
    batch(real.(sum_dropdims(arr .* λ, dims=nonbatch_dims(L.diag))))
end

function tr(L::Diagonal{<:Real,<:LambertField{B}}) where {B<:Union{Map,Basis2Prod{<:Any,Map},Basis3Prod{<:Any,<:Any,Map}}}
    batch(sum_dropdims(L.diag.arr, dims=nonbatch_dims(L.diag)))
end

### simulation
_white_noise(ξ::LambertField, rng::AbstractRNG) =
    (randn!(rng, similar(ξ.arr, real(eltype(ξ)), ξ.Ny, size(ξ.arr)[2:end]...)), ξ.metadata)
white_noise(ξ::LambertS0,  rng::AbstractRNG) = LambertMap(_white_noise(ξ,rng)...)
white_noise(ξ::LambertS2,  rng::AbstractRNG) = LambertEBMap(_white_noise(ξ,rng)...)
white_noise(ξ::LambertS02, rng::AbstractRNG) = LambertIEBMap(_white_noise(ξ,rng)...)

### creating covariance operators
# fixed covariances
Cℓ_to_Cov(pol::Symbol, args...; kwargs...) = Cℓ_to_Cov(Val(pol), args...; kwargs...)
function Cℓ_to_Cov(::Val{:I}, proj::ProjLambert, Cℓ::InterpolatedCℓs; units=proj.Ωpix)
    Diagonal(LambertFourier(Cℓ_to_2D(Cℓ,proj), proj) / units)
end
function Cℓ_to_Cov(::Val{:P}, proj::ProjLambert, CℓEE::InterpolatedCℓs, CℓBB::InterpolatedCℓs; units=proj.Ωpix)
    Diagonal(LambertEBFourier(Cℓ_to_2D(CℓEE,proj), Cℓ_to_2D(CℓBB,proj), proj) / units)
end
function Cℓ_to_Cov(::Val{:IP}, proj::ProjLambert, CℓTT, CℓEE, CℓBB, CℓTE; kwargs...)
    ΣTT, ΣEE, ΣBB, ΣTE = [Cℓ_to_Cov(:I,proj,Cℓ; kwargs...) for Cℓ in (CℓTT,CℓEE,CℓBB,CℓTE)]
    BlockDiagIEB(@SMatrix([ΣTT ΣTE; ΣTE ΣEE]), ΣBB)
end
# ParamDependentOp covariances scaled by amplitudes in different ℓ-bins
function Cℓ_to_Cov(::Val{:I}, proj::ProjLambert{T}, (Cℓ, ℓedges, θname)::Tuple; kwargs...) where {T}
    # we need an @eval here since we want to dynamically select a
    # keyword argument name, θname. the @eval happens into Main rather
    # than CMBLensing as a workaround for
    # https://discourse.julialang.org/t/closure-not-shipping-to-remote-workers-except-from-main/38831
    C₀ = diag(Cℓ_to_Cov(:I, proj, Cℓ; kwargs...))
    @eval Main let ℓedges=$((T.(ℓedges))...,), C₀=$C₀
        ParamDependentOp(function (;$θname=ones($T,length(ℓedges)-1),_...)
            _A = $preprocess.(Ref((nothing,C₀.metadata)), $T.($ensure1d($θname)))
            Diagonal(LambertFourier($bandpower_rescale!(ℓedges, C₀.ℓmag, copy(C₀.arr), _A...), C₀.metadata))
        end)
    end
end
function Cℓ_to_Cov(::Val{:P}, proj::ProjLambert{T}, (CℓEE, ℓedges, θname)::Tuple, CℓBB::InterpolatedCℓs; kwargs...) where {T}
    C₀ = diag(Cℓ_to_Cov(:P, proj, CℓEE, CℓBB; kwargs...))
    @eval Main let ℓedges=$((T.(ℓedges))...,), C₀=$C₀
        ParamDependentOp(function (;$θname=ones($T,length(ℓedges)-1),_...)
            _E = $preprocess.(Ref((nothing,C₀.metadata)),      $T.($ensure1d($θname)))
            _B = $preprocess.(Ref((nothing,C₀.metadata)), one.($T.($ensure1d($θname))))
            Diagonal(LambertEBFourier($bandpower_rescale!(ℓedges, C₀.ℓmag, copy(C₀.El), _E...), C₀.Bl .* _B[1], C₀.metadata))
        end)
    end
end
# this is written weird because the stuff inside the broadcast! needs
# to work as a GPU kernel
function bandpower_rescale!(ℓedges, ℓ, Cℓ, A...)
    length(A)==length(ℓedges)-1 || error("Expected $(length(ℓedges)-1) bandpower parameters, got $(length(A)).")
    eltype(A[1]) <: Real || error("Bandpower parameters must be real numbers.")
    if length(A)>30
        # if more than 30 bandpowers, we need to chunk the rescaling
        # because of a maximum argument limit of CUDA kernels
        for p in partition(1:length(A), 30)
            bandpower_rescale!(ℓedges[p.start:(p.stop+1)], ℓ, Cℓ, A[p]...)
        end
    else
        broadcast!(Cℓ, ℓ, Cℓ, A...) do ℓ, Cℓ, A...
            for i=1:length(ℓedges)-1
                (ℓedges[i] < ℓ < ℓedges[i+1]) && return A[i] * Cℓ
            end
            return Cℓ
        end
    end
    Cℓ
end
# cant reliably get Zygote's gradients to work through these
# broadcasts, which on GPU use ForwardDiff, so write the adjoint by
# hand for now. likely more performant, in any case.
@adjoint function bandpower_rescale!(ℓedges, ℓ, Cℓ, A...)
    function back(Δ)
        Ā = map(1:length(A)) do i
            sum(
                real,
                broadcast(Δ, ℓ, Cℓ) do Δ, ℓ, Cℓ
                    (ℓedges[i] < ℓ < ℓedges[i+1]) ? Cℓ*Δ : zero(Cℓ)
                end,
                dims = ndims(Δ)==4 ? (1,2) : (:)
            )
        end
        (nothing, nothing, nothing, Ā...)
    end
    bandpower_rescale!(ℓedges, ℓ, Cℓ, A...), back
end
function cov_to_Cℓ(C::DiagOp{<:LambertS0}; kwargs...)
    @unpack Nx, Ny, Δx = diag(C)
    α = Nx*Ny/Δx^2
    get_Cℓ(sqrt.(diag(C)); kwargs...)*sqrt(α)
end

### spin adjoints
function *(a::SpinAdjoint{F}, b::F) where {B<:Union{Map,Basis2Prod{<:Any,Map},Basis3Prod{<:Any,<:Any,Map}},F<:LambertField{B}}
    r = sum(a.f.arr .* b.arr, dims=3)
    LambertMap(dropdims(r, dims=(ndims(r)==3 ? 3 : ())), get_metadata_strict(a, b))
end
function mul!(dst::LambertMap, a::SpinAdjoint{F}, b::F) where {F<:LambertField{<:Union{Map,Basis2Prod{<:Any,Map},Basis3Prod{<:Any,<:Any,Map}}}}
    dst.arr .= reshape(sum(a.f.arr .* b.arr, dims=3), size(dst.arr))
    dst
end

### batching

batch_length(f::LambertField) = f.Nbatch

"""
    batch(fs::LambertField...)
    batch(fs::Vector{<:LambertField})

Concatenate one of more LambertFields along the "batch" dimension
(dimension 4 of the underlying array). For the inverse operation, see
[`unbatch`](@ref).
"""
batch(fs::LambertField{B}...) where {B} = batch(collect(fs))
batch(fs::AbstractVector{<:LambertField{B}}) where {B} =
    LambertField{B}(cat(getfield.(fs,:arr)..., dims=Val(4)), only(unique(getfield.(fs,:metadata))))

"""
    unbatch(f::LambertField)

Return an Array of LambertFields corresponding to each batch index. For
the inverse operation, see [`batch`](@ref).
"""
unbatch(f::LambertField{B}) where {B} = [f[!,i] for i=1:batch_length(f)]

batch_index(f::LambertField{B}, I) where {B<:Union{Map,Fourier}} = LambertField{B}(f.arr[:,:,1,I], f.metadata)
batch_index(f::LambertField{B}, I) where {B} = LambertField{B}(f.arr[:,:,:,I], f.metadata)

###

make_mask(f::LambertField; kwargs...) = make_mask((f.Ny,f.Nx), f.θpix; kwargs...)

### power spectra

function get_Cℓ(f₁::LambertS0, f₂::LambertS0=f₁; Δℓ=50, ℓedges=0:Δℓ:16000, Cℓfid=ℓ->1, err_estimate=false)
    @unpack Nx, Ny, Δx, ℓmag = fieldinfo(f₁)
    ℓmag = unfold(ℓmag, Ny)
    α = Nx*Ny/Δx^2

    # faster to excise unused parts:
    ℓmask = (ℓmag .> minimum(ℓedges)) .&  (ℓmag .< maximum(ℓedges))
    L = ℓmag[ℓmask]
    CLobs = 1/α .* real.(dot.(
        adapt(Array,f₁)[:Il, full_plane=true][ℓmask],
        adapt(Array,f₂)[:Il, full_plane=true][ℓmask]
    ))
    w = @. nan2zero((2*Cℓfid(L)^2/(2L+1))^-1)

    sum_in_ℓbins(x) = Float64.(fit(Histogram, L, Weights(x), ℓedges).weights)

    local A, Cℓ, ℓ, N, Cℓ²
    Threads.@sync begin
        Threads.@spawn A  = sum_in_ℓbins(w)
        Threads.@spawn Cℓ = sum_in_ℓbins(w .* CLobs)
        Threads.@spawn ℓ  = sum_in_ℓbins(w .* L)
        if err_estimate
            Threads.@spawn N   = sum_in_ℓbins(one.(w)) / 2
            Threads.@spawn Cℓ² = sum_in_ℓbins(w .* CLobs.^2)
        end
    end

    if err_estimate
        σℓ  = sqrt.((Cℓ² ./ A .- Cℓ.^2) ./ N)
        InterpolatedCℓs(ℓ./A,  Cℓ./A .± σℓ)
    else
        InterpolatedCℓs(ℓ./A,  Cℓ./A)
    end
end

function get_Cℓ(f1::LambertS2, f2::LambertS2=f1; which=(:EE,:BB), kwargs...)
    Cℓ = (;[Symbol(x1*x2) => get_Cℓ(getindex(f1,Symbol(x1)),getindex(f2,Symbol(x2)); kwargs...) for (x1,x2) in split.(string.(ensure1d(which)),"")]...)
    which isa Symbol ? Cℓ[1] : Cℓ
end

function get_Cℓ(f1::LambertS02, f2::LambertS02=f1; which=(:II,:EE,:BB,:IE,:IB,:EB), kwargs...)
    Cℓ = (;[Symbol(x1*x2) => get_Cℓ(getindex(f1,Symbol(x1)),getindex(f2,Symbol(x2)); kwargs...) for (x1,x2) in split.(string.(ensure1d(which)),"")]...)
    which isa Symbol ? Cℓ[1] : Cℓ
end

"""
    ud_grade(f::Field, θnew, mode=:map, deconv_pixwin=true, anti_aliasing=true)

Up- or down-grades field `f` to new resolution `θnew` (only in integer steps).
Two modes are available specified by the `mode` argument:

* `:map`     — Up/downgrade by replicating/averaging pixels in map-space
* `:fourier` — Up/downgrade by extending/truncating the Fourier grid

For `:map` mode, two additional options are possible. If `deconv_pixwin` is
true, deconvolves the pixel window function from the downgraded map so the
spectrum of the new and old maps are the same. If `anti_aliasing` is true,
filters out frequencies above Nyquist prior to down-sampling.

"""
function ud_grade(
    f :: LambertField{B},
    θnew;
    mode = :map,
    deconv_pixwin = (mode==:map),
    anti_aliasing = (mode==:map)
) where {B}

    N,θ,T = (f.Ny, f.Nx), f.θpix, f.T
    θnew==θ && return f
    (mode in [:map,:fourier]) || throw(ArgumentError("Available modes: [:map,:fourier]"))

    fac = θnew > θ ? θnew÷θ : θ÷θnew
    (round(Int, fac) ≈ fac) || throw(ArgumentError("Can only ud_grade in integer steps"))
    fac = round(Int, fac)
    Ny_new, Nx_new = @. round(Int, N * θ ÷ θnew)
    proj = ProjLambert(;Ny=Ny_new, Nx=Nx_new, θpix=θnew, T=real(T), f.storage)
    @unpack Δx,ℓx,ℓy,Nx,Ny,nyquist = proj

    PWF = deconv_pixwin ? Diagonal(LambertFourier((@. T((pixwin(θnew,ℓy)*pixwin(θnew,ℓx)')/(pixwin(θ,ℓy)*pixwin(θ,ℓx)'))), proj)) : I

    if θnew > θ
        # downgrade
        if anti_aliasing
            f = Diagonal(LambertFourier(ifelse.((abs.(f.ℓy) .>= nyquist) .| (abs.(f.ℓx') .>= nyquist), 0, 1), f.metadata)) * f
        end
        if mode == :map
            fnew = LambertField{Map(B)}(dropdims(mean(reshape(Map(f).arr, fac, Ny, fac, Nx, size.(Ref(f.arr),nonbatch_dims(f)[3:end])...), dims=(1,3)), dims=(1,3)), proj)
        else
            fnew = LambertField{Fourier(B)}(Fourier(f).arr[1:(Ny_new÷2+1), [1:(isodd(Nx_new) ? Nx_new÷2+1 : Nx_new÷2); (end-Nx_new÷2+1):end], repeated(:, length(nonbatch_dims(f))-2)...], proj)
        end
        if deconv_pixwin
            fnew = Diagonal(LambertFourier((@. T((pixwin(θnew,ℓy)*pixwin(θnew,ℓx)')/(pixwin(θ,ℓy)*pixwin(θ,ℓx)'))), proj)) \ fnew
        end
    else
        error("Not implemented")
        # # upgrade
        # @assert fieldinfo(f).Nside isa Int "Upgrading resolution only implemented for maps where `Nside isa Int`"
        # if mode==:map
        #     fnew = LambertMap{Pnew}(permutedims(hvcat(N,(x->fill(x,(fac,fac))).(f[:Ix])...)))
        #     deconv_pixwin ? LambertFourier{Pnew}(fnew[:Il] .* PWF' .* PWF[1:Nnew÷2+1]) : fnew
        # else
        #     fnew = LambertFourier{Pnew}(zeros(Nnew÷2+1,Nnew))
        #     setindex!.(Ref(fnew.Il), f[:Il], 1:(N÷2+1), [findfirst(fieldinfo(fnew).k .≈ fieldinfo(f).k[i]) for i=1:N]')
        #     deconv_pixwin ? fnew * fac^2 : fnew
        # end

    end
    return fnew
end


text:
struct AdversarialDebiasing{M<:MLJBase.Model}<: Deterministic
    # grp::Symbol
    classifier::M
    # alpha::Float64
end

function AdversarialDebiasing(; classifier::MLJBase.Model = nothing)
    model = AdversarialDebiasing(classifier)
    message = MLJBase.clean!(model)
    isempty(message)||@warn message
    return model
end

function MLJBase.clean!(model::AdversarialDebiasing)
    warning = ""
    model.classifier!=nothing || (warning *= "No classifier specified in model\n")
    target_scitype(model) <: AbstractVector{<:Finite} || (warning *= "Only Binary Classifiers are supported\n")
    # (model.alpha>=0 && model.alpha<=1) || (warning*="alpha should be between 0 and 1 (inclusive)\n")
    return warning
end

function classifier_model(features_dim, classifier_num_hidden_units, keep_prob)
	layer1 = Dense(features_dim,classifier_num_hidden_units; bias=true)
	layer2 = Dense(classifier_num_hidden_units, 1, bias = true, σ)
	model = Chain(layer1, LayerNorm(classifier_num_hidden_units, relu), layer2)
	return model
end

function adversary_model()
	layer1 = Dense(3,1;bias=true)
	return layer1
end

loss(ŷ,y) = Flux.logitbinarycrossentropy(ŷ, y; agg = mean)

classifier_opt = Flux.Optimise.Optimiser(Flux.Optimise.ExpDecay(0.001, 0.1, 1000, 1e-4), Flux.Optimise.ADAM())
classifier_opt2 = Flux.Optimise.Optimiser(Flux.Optimise.ExpDecay(0.001, 0.1, 1000, 1e-4), Flux.Optimise.ADAM())

normalise(temp) = temp/(norm(temp)+ 1e-8)

function MMI.fit(model::AdversarialDebiasing, verbosity::Int,X,y)
    # 	@show "Hey i am being called"
        df_n1 = DataFrames.DataFrame()
        for name in names(X)
            if typeof(X[:,name])==CategoricalArray{String,1,UInt32,String,CategoricalValue{String,UInt32},Union{}}
                df_n1[:,name] = coerce(X[:,name], Count, verbosity=0)
            else
                df_n1[:,name] = X[:,name]
            end
        end

        @show "this was called"

        train_loader = Flux.Data.DataLoader((Matrix(df_n1)', rand(Bool,6907)*1), batchsize=100, shuffle=true)

        # a3 = @pipeline ContinuousEncoder  NeuralNetworkClassifier
        # @show model.classifier
        # mch = machine(model.classifier, X, y)
        # @show mch
        # fit!(mch)
        classifier_net = classifier_model(8,5,1)
        adversary_net = adversary_model()
        classifier_vars = Flux.params(classifier_net)
        adversary_vars = Flux.params(adversary_net)
        for (x,y) in train_loader
            classifier_grad = gradient(() -> loss(σ.(classifier_net(x)),y), classifier_vars)
            adversary_grad = gradient(() -> loss(σ.(adversary_net(cat(σ.(2*classifier_net(x)), σ.(2*classifier_net(x)).*y',σ.(2*classifier_net(x)).*(1 .- y)';dims=1))),x[4,:]), classifier_vars)
            for var in classifier_vars
                unit_adversary_grad = normalise(adversary_grad[var])
                classifier_grad[var] = classifier_grad[var]-sum(classifier_grad[var].*unit_adversary_grad)*unit_adversary_grad
                classifier_grad[var] = classifier_grad[var] - 0.01*adversary_grad[var]
            end
            adversary_grad2 = gradient(() -> loss(σ.(adversary_net(cat(σ.(2*classifier_net(x)), σ.(2*classifier_net(x)).*y',σ.(2*classifier_net(x)).*(1 .- y)';dims=1))), x[4,:]), adversary_vars)
            Flux.Optimise.update!(classifier_opt, classifier_vars, classifier_grad)
            Flux.Optimise.update!(adversary_opt, adversary_vars, adversary_grad2)
            # append!(losses, loss(σ.(classifier_net(x)),y))
            # append!(losses_adversary, loss(σ.(adversary_net(cat(σ.(2*classifier_net(x)), σ.(2*classifier_net(x)).*y',σ.(2*classifier_net(x)).*(1 .- y)';dims=1))), rand(1,100)))
        end
        @show Flux.params(classifier_net)
        fitresult =  Flux.params(classifier_net)
        return fitresult
end

function MMI.predict(model::AdversarialDebiasing, fitresult, X)
	df_n1 = DataFrames.DataFrame()
	for name in names(X)
		if typeof(X[:,name])==CategoricalArray{String,1,UInt32,String,CategoricalValue{String,UInt32},Union{}}
			df_n1[:,name] = coerce(X[:,name], Count, verbosity=0)
		else
			df_n1[:,name] = X[:,name]
		end
	end

	@show size(fitresult)

	@show "this was called"
	classifier_net2 = classifier_model(8,5,1)
	@show Flux.params(classifier_net2)
	Flux.loadparams!(classifier_net2, fitted_params(mch2))
	ŷ = classifier_net2(Matrix(df_n1)')
	ŷ=round.(ŷ)
	return ŷ
end

MMI.input_scitype(::Type{<:AdversarialDebiasing{M}}) where M = input_scitype(M)
MMI.target_scitype(::Type{<:AdversarialDebiasing{M}}) where M = AbstractVector{<:Finite{2}}

ext: jl


text:
let
        tot = 0
        totR = 0
        open("data.txt") do file
            for line in eachline(file)
                    tokens = split(line, "x")
                    ints = sort(map(x -> parse(Int, x), tokens))
                    a = 3*ints[1]*ints[2] + 2*ints[2]*ints[3] + 2*ints[1]*ints[3]
                    r = 2*ints[1] + 2*ints[2] + ints[1]*ints[2]*ints[3]
                tot += a
                totR += r
            end
        end
        println(tot)
        println(totR)
end


text:
using UnitfulConventionalMoles, Unitful, Test

@testset "conversion to grams" begin
    @test uconvert(u"g", 2u"molH") == 2.016u"g"
end

max_stars_repo_name: rafaqz/UnitfulConventionalMoles

max_issues_repo_name: rafaqz/UnitfulConventionalMoles

max_forks_repo_name: rafaqz/UnitfulConventionalMoles


text:
struct SimpleStiffSystem
    λ::Float64
    de::Function
    function SimpleStiffSystem(de::Function)
        return new(1.0, de)
    end
end

function stiff_de!(xx::Vector{T}, x::Vector{T}, s::SimpleStiffSystem, t::Float64=0.0) where {T}
    xx .= -s.λ * x
    return nothing
end

my_stiff_struct = SimpleStiffSystem(stiff_de!)
x0 = zeros(Float64, 4) .+ 1.0
n_rule_max = 3

@testset "basic_test" begin
    for NC = 1:10  # check different chunk sizes
        rr = makeRadauIntegrator(my_stiff_struct, x0, 1.0e-16, n_rule_max, NC)
        for k_rule = 1:n_rule_max  # chick different rules
            rr.rule.s = 3
            t_final = 0.2
            update_h!(rr, t_final)
            h, x_final = solveRadau(rr, x0)
            if k_rule != 1
                x_ana = exp(-t_final)
                @test x_ana ≈ x_final[1]
            end
        end
    end
end


text:
# function to convert elements in mongo_keys into bigData CloudGraphs container

using Caesar, CloudGraphs, Neo4j
using LibBSON, Mongo
using JSON

using IncrementalInference

using PyCall

# lcmtpath = joinpath(dirname(@__FILE__),"python")
# unshift!(PyVector(pyimport("sys")["path"]),lcmtpath)
# include(joinpath(dirname(@__FILE__),"blandauthremote.jl"))
# collection = "bindata"

@pyimport numpy as np
@pyimport cv2 as opencv

# include("VisualizationUtilities.jl") # @pyimport getimages as gi

## Uncomment out for command line operation
# cloudGraph, addrdict = standardcloudgraphsetup()
# session = addrdict["session"]
# mongoaddress = addrdict["mongo addr"]

# interactive operation
session = "SESSROX"
Nparticles = 100
include(joinpath(dirname(@__FILE__),"blandauthremote.jl"))

# also connect to mongo separately
client = pymongo.MongoClient(mongoaddress)
db = client[:CloudGraphs]

fg = initfg(sessionname=session, cloudgraph=cloudGraph)

# totally reset to frontend built state in DB
# resetentireremotesession(conn,session)

# updatenewverts!(fg, N=N)
# --
# sortedvd = getnewvertdict(fg.cg.neo4j.connection, fg.sessionname)
# populatenewvariablenodes!(fg, sortedvd, N=N)

fullLocalGraphCopy!(fg)

IDs = getPoseExVertexNeoIDs(fg.cg.neo4j.connection, session=session, reqbackendset=false);

xx,ll = ls(fg)
LD = Array{Array{Float64,1},1}()
C = Vector{AbstractString}()

lbl = xx[2]

mongk = getmongokeys(fg, lbl, IDs)

img = fetchmongoimg(db[collection], mongk["depthframe_image"])
# mongo_keydepth = bson.ObjectId(mongk["depthframe_image"])
# dpim, ims = gi.fastdepthimg(db[collection], mongo_keydepth)

# ===============================================================

# get the mongokeys from Neo

# ===============================================================
using PyCall

lcmtpath = joinpath(dirname(@__FILE__),"python")
unshift!(PyVector(pyimport("sys")["path"]),lcmtpath)
include(joinpath(dirname(@__FILE__),"blandauthremote.jl"))

@show mongoaddress

@pyimport bson
@pyimport pymongo

@pyimport numpy as np
@pyimport cv2 as opencv

@pyimport getimages as gi

collection = "bindata"

client = pymongo.MongoClient(mongoaddress)
db = client[:CloudGraphs]

key = "58af67265d7625647859fa73"
@show gi.test(key,1)

# @show gi.getrgbimg(mongoaddress, collection, key)

mongo_key = bson.ObjectId(key)

im = gi.fastrgbimg(db[collection], mongo_key)
@show size(im),typeof(im)
# opencv.imshow("image", im)

keydepth = "58b07be05d7625699151a01d"
mongo_keydepth = bson.ObjectId(keydepth)

dpim, ims = gi.fastdepthimg(db[collection], mongo_keydepth)

key = "58b07bdf5d7625699151a01c"
mongo_key = bson.ObjectId(key)

im, ims = gi.fastrgbimg(db[collection], mongo_key)

@show size(im)

results = (db[collection])[:find]( Dict{String, Any}("_id" => mongo_key) )

for rec in results
  @show rec["description"]
  # im = rec["val"]
  val = "np.fromstring(str($(rec)))"
  pycall()
end

# =============================================================

mongoId = BSONOID("58af67265d7625647859fa73")

im = []

val = find(fg.cg.mongo.cgBindataCollection, query("_id" => mongoId))
elem = first(val)

val = elem["val"]

@show typeof(val)

# unsafe_load(val,1) # doesn't work

LibBSON.libbson

for doc in find(fg.cg.mongo.cgBindataCollection, query("_id" => mongoId))
  im = doc["val"]
end

numNodes = count(fg.cg.mongo.cgBindataCollection, ("_id" => mongoId));
if numNodes != 1
  error("The query for $(mongoId) returned $(numNodes) values, expected a 1-to-1 mapping!");
end

results = first(find(fg.cg.mongo.cgBindataCollection, ("_id" => eq(mongoId))) )

data = CloudGraphs._stringToUint8Array(results["val"]);

@show length(results)

o = next(results, nothing)

for (k,v) in results
  @show k
end

dr = dict(results)

id = slam.fg.IDs[sp2[1]]
cgid = slam.fg.cgIDs[id]
cv = CloudGraphs.get_vertex(slam.fg.cg, cgid)
# for f in sp2[2:end]
fid = open(sp2[2],"r")
imageData = readbytes(fid) # imageData::Vector{UInt8}
close(fid)
bdei = CloudGraphs.BigDataElement("keyframe-image", imageData)
push!(cv.bigData.dataElements, bdei);

fid = open(sp2[3],"r")
imageData = readbytes(fid) # imageData::Vector{UInt8}
close(fid)
bdei = CloudGraphs.BigDataElement("depthframe-image", imageData)
push!(cv.bigData.dataElements, bdei);

CloudGraphs.save_BigData!(slam.fg.cg, cv)

cursor = find(cloudGraph.mongo.cgBindataCollection, query())
for o in cursor
    # println( ", oid: ", o["_id"])
    println("id: ", o["_id"]);
end

#


text:
using DataFrames, DiscriminantAnalysis;

iris_df = readtable("iris.csv");
pool!(iris_df, [:Species]);  # Ensure species is made a pooled data vector

X = convert(Array{Float64}, iris_df[[:PetalLength, :PetalWidth, :SepalLength, :SepalWidth]]);
y = iris_df[:Species].refs;  # Class indices

#== Fitting the LDA model ==#

model1 = lda(X, y)
y_pred1 = classify(model1, X);
accuracy1 = sum(y_pred1 .== y)/length(y)

# The following illustrates column-major ordering with gamma-regularization of 0.1 and non-uniform
# class priors:
model2 = lda(X', y, order = Val{:col}, gamma = 0.1, priors = [2/5; 2/5; 1/5])
y_pred2 = classify(model2, X');
accuracy2 = sum(y_pred2 .== y)/length(y)

#== Fitting the QDA model ==#

model3 = qda(X, y, lambda = 0.1, gamma = 0.1, priors = [1/3; 1/3; 1/3])
y_pred3 = classify(model3, X);
accuracy3 = sum(y_pred3 .== y)/length(y)

#== Fitting the CDA model ==#

model4 = cda(X, y, gamma = Nullable{Float64}(), priors = [1/3; 1/3; 1/3])
y_pred4 = classify(model4, X);
accuracy4 = sum(y_pred4 .== y)/length(y)


text:
# exp(-t)*(2-x)*x*(1-y)
using Revise
using ADCME
using AdFem
using PyPlot

n = 20
m = 2n
NT = 200
ρ = 0.0
Δt = 1/NT
h = 1/n
x = zeros((m+1)*(n+1))
y = zeros((m+1)*(n+1))
for i = 1:m+1
    for j = 1:n+1
        idx = (j-1)*(m+1)+i
        x[idx] = (i-1)*h
        y[idx] = (j-1)*h
    end
end
bd = bcnode("all", m, n, h)

u1 = (x,y,t)->exp(-t)*(0.5-x)^2*(2-y)^2
u2 = (x,y,t)->exp(-t)*(0.5-x^2)^2*(2-sin(y))^2

ts = Δt * ones(NT)
dt = αscheme_time(ts, ρ = ρ )
F = zeros(NT, 2(m+1)*(n+1))
for i = 1:NT
    t = dt[i]
    f1 = (x,y)->-9.87654320987654*x*(0.5 - x^2)*(2 - sin(y))*exp(-t)*cos(y) + (0.5 - x)^2*(2 - y)^2*exp(-t) - 0.740740740740741*(0.5 - x)^2*exp(-t) - 3.20987654320988*(2 - y)^2*exp(-t)
  	f2 = (x,y)->-2.96296296296296*x^2*(2 - sin(y))^2*exp(-t) + (0.5 - x^2)^2*(2 - sin(y))^2*exp(-t) - 3.20987654320988*(0.5 - x^2)^2*(2 - sin(y))*exp(-t)*sin(y) - 3.20987654320988*(0.5 - x^2)^2*exp(-t)*cos(y)^2 + 1.48148148148148*(0.5 - x^2)*(2 - sin(y))^2*exp(-t) - (0.740740740740741*x - 0.37037037037037)*(2*y - 4)*exp(-t) - (2*x - 1.0)*(1.72839506172839*y - 3.45679012345679)*exp(-t)
    fval1 = eval_f_on_gauss_pts(f1, m, n, h)
  	fval2 = eval_f_on_gauss_pts(f2, m, n, h)
    F[i,:] = compute_fem_source_term(fval1, fval2, m, n, h)
end

abd = zeros(NT, (m+1)*(n+1)*2)
for i = 1:NT
    t = dt[i]
    abd[i,:] = [(@. u1(x, y, t)); (@. u2(x, y, t))]
end
abd = constant(abd[:, [bd; bd .+ (m+1)*(n+1)]])

E = 1.0
ν = 0.35
H = E/(1+ν)/(1-2ν)*[
  1-ν ν 0
  ν 1-ν 0
  0 0 (1-2ν)/2
]
M = constant(compute_fem_mass_matrix(m, n, h))
K = constant(compute_fem_stiffness_matrix(H, m, n, h))

a0 = [(@. u1(x, y, 0.0)); (@. u2(x, y, 0.0))]
u0 = -[(@. u1(x, y, 0.0)); (@. u2(x, y, 0.0))]
d0 = [(@. u1(x, y, 0.0)); (@. u2(x, y, 0.0))]

function solver(A, rhs, i)
    A, Abd = fem_impose_Dirichlet_boundary_condition_experimental(A, bd, m, n, h)
    rhs = rhs - Abd * abd[i]
    rhs = scatter_update(rhs, [bd; bd .+ (m+1)*(n+1)], abd[i])
    return A\rhs
end
d, u, a = αscheme(M, spzero(2(m+1)*(n+1)), K, F, d0, u0, a0, ts; extsolve=solver, ρ = ρ  )

sess = Session()
d_, u_, a_ = run(sess, [d, u, a])

function plot_traj(idx)
    figure(figsize=(12,3))
    subplot(131)
    plot((0:NT)*Δt, u1.(x[idx], y[idx],(0:NT)*Δt), "b-", label="x-Acceleration")
    plot((0:NT)*Δt, a_[:,idx], "y--", markersize=2)
  	plot((0:NT)*Δt, u2.(x[idx], y[idx],(0:NT)*Δt), "r-", label="y-Acceleration")
    plot((0:NT)*Δt, a_[:,idx+(m+1)*(n+1)], "c--", markersize=2)
    legend()
    xlabel("Time")
    ylabel("Value")

    subplot(132)
    plot((0:NT)*Δt, u1.(x[idx], y[idx],(0:NT)*Δt), "b-", label="x-Displacement")
    plot((0:NT)*Δt, d_[:,idx], "y--", markersize=2)
  	plot((0:NT)*Δt, u2.(x[idx], y[idx],(0:NT)*Δt), "r-", label="y-Displacement")
    plot((0:NT)*Δt, d_[:,idx+(m+1)*(n+1)], "c--", markersize=2)
    legend()
    xlabel("Time")
    ylabel("Value")

    subplot(133)
    plot((0:NT)*Δt, -u1.(x[idx], y[idx],(0:NT)*Δt), "b-", label="x-Velocity")
    plot((0:NT)*Δt, u_[:,idx], "y--", markersize=2)
  	plot((0:NT)*Δt, -u2.(x[idx], y[idx],(0:NT)*Δt), "r-", label="y-Velocity")
    plot((0:NT)*Δt, u_[:,idx+(m+1)*(n+1)], "c--", markersize=2)
    legend()
    xlabel("Time")
    ylabel("Value")

    tight_layout()
end

idx2 = (n÷3)*(m+1) + m÷2
plot_traj(idx2)

max_stars_repo_name: ADCMEMarket/ADCMEImages

max_issues_repo_name: ADCMEMarket/ADCMEImages

max_forks_repo_name: ADCMEMarket/ADCMEImages


text:
# Reducing allocations as much as possible.
# Included after the non-mutating versions, so that we can
# fall back to a less efficient version where needed

function cross3!(c,a,b)
    a1 = @view a[:,1]
    a2 = @view a[:,2]
    a3 = @view a[:,3]

    b1 = @view b[:,1]
    b2 = @view b[:,2]
    b3 = @view b[:,3]

    c[:,1] .= @. a2*b3 - a3*b2
    c[:,2] .= @. a3*b1 - a1*b3
    c[:,3] .= @. a1*b2 - a2*b1
end

function dot3!(c,a,b)
    @inbounds for i = 1:length(c)
        c[i] = a[i,1]*b[i,1] + a[i,2]*b[i,2] + a[i,3]*b[i,3]
    end
end

function norm3!(a)
    @inbounds for i = 1:size(a,1)
        tmp = sqrt(a[i,1]*a[i,1] + a[i,2]*a[i,2] + a[i,3]*a[i,3])
        a[i,1] /= tmp
        a[i,2] /= tmp
        a[i,3] /= tmp
    end
end

function norm3!(out, a)
    @inbounds for i = 1:length(out)
        out[i] = sqrt(a[i,1]*a[i,1] + a[i,2]*a[i,2] + a[i,3]*a[i,3])
    end
end

function edges!(e, x)
    nv = size(x, 1)
    e .= view(x, 2:nv, :) .- view(x, 1:nv-1, :)
end

function tangents!(t, x)
    edges!(t,x)
    norm3!(t)
end

# ----------------------------------------------------------------------------------
# Parallel Transport (Vectorized)
# ----------------------------------------------------------------------------------
#
# Uses the Rodriguez rotation formula. There is a singularity in
# this parametrization of rotation matrices, for t1 and t2 almost
# aligned. This singularity is handled in ptransport_inner_coefficient,
# by switching to a small-θ expansion.
#
# The conditional does not affect the "true" value of the function,
# it simply switches to a more stable method of computing the same
# quantity. Consequently, it is hidden from reverse-mode AD by defining
# a custom pullback for the function 'ptransport_inner_coefficient'

const PTRANSPORT_TOL2 = PTRANSPORT_TOL^2
function ptransport_inner_coefficient!(out, s2, c)
    out .= 1/2 .+ (1 .- c)./4 .+ (1 .- c).^2 .*(1 .+ s2./4)./8
    @inbounds for i = 1:length(out)
        if s2[i] > PTRANSPORT_TOL2 #ind[i]
            out[i] = (1 - c[i])/s2[i]
        end
    end
end

# Cache needs to be ne x 6
function ptransport!(v2, v1, t1, t2, cache)
    X = view(cache, :, 1:3)
    s2 = view(cache, :, 4)
    c = view(cache, :, 5)
    ptic = view(cache, :, 6)

    cross3!(X, t1, t2)
    dot3!(s2, X, X)
    dot3!(c, t1, t2)
    ptransport_inner_coefficient!(ptic,s2,c)

    dXv = view(cache, :, 4) # Overwrites s2
    dot3!(dXv,X,v1)
    cross3!(v2,X,v1)
    v2 .+= c.*v1 .+ ptic.*dXv.*X
end

# ------------------------------------------------------------------------------
# Rotate about Unit Vector (Vectorized)
# ------------------------------------------------------------------------------
# Rodriguez rotation formula, specialized to the case where the vector
# and the axis are unit length and orthogonal.

# cache is ne x 5
function rotate_orthogonal_unit!(v, t, θ, cache)
    tmp = view(cache, :, 1:3)
    c = view(cache, :, 4:4)
    s = view(cache, :, 5:5)

    c .= cos.(θ)
    s .= sin.(θ)
    cross3!(tmp, t, v)
    v .= c.*v .+ s.*tmp
end

# Update kinematics given displacement and twist
function rod_update!(x1, d1, x0, d0, Δx, Δθ, cache)
    t0 = view(cache, :, 1:3)
    t1 = view(cache, :, 4:6)
    cache_pt = view(cache, :, 7:12)
    cache_rt = view(cache, :, 7:11)

    tangents!(t0, x0)
    x1 .= x0 .+ Δx
    tangents!(t1, x1)

    ptransport!(d1, d0, t0, t1, cache_pt)
    rotate_orthogonal_unit!(d1, t1, Δθ, cache_pt)
end

allocate_cache(f::typeof(rod_update!), T::Type, ns) = zeros(T, ns+1, 12)

function full_kinematics!(l, κ, τ, x, d1, caches)

    # Unpack and divy up caches
    cache1, cache2 = caches
    t = view(cache1, :, 1:3)
    d2 = view(cache1, :, 4:6)
    k = view(cache2, :, 1:3)
    tmp1 = view(cache2, :, 4:4)
    tmp3 = view(cache2, :, 4:6)
    pt_cache = view(cache2, :, 1:6) # overrides k, tmp
    d1trans = view(cache2, :, 7:9)
    cosτ = view(cache2, :, 1)
    sinτ = view(cache2, :, 2)

    # Lengths
    edges!(t, x)
    norm3!(l, t)
    t ./= l
    cross3!(d2, t, d1)

    # Break tangents and directors into sets
    nv = size(x,1)
    t1 = view(t, 1:nv-2, :)
    t2 = view(t, 2:nv-1, :)
    d1l = view(d1, 1:nv-2, :)
    d1r = view(d1, 2:nv-1, :)
    d2l = view(d2, 1:nv-2, :)
    d2r = view(d2, 2:nv-1, :)

    # Curvature
    cross3!(k, t1, t2)
    dot3!(tmp1, t1, t2)
    k .*= 2.0./(1.0 .+ tmp1)
    tmp3 .= k.*(d2l .+ d2r)./2
    sum!(view(κ, :, 1:1), tmp3)
    tmp3 .= -1.0.*k.*(d1l .+ d1r)./2
    sum!(view(κ, :, 2:2), tmp3)

    # # Twists (could move in part of ptransport to reduce intermediates)
    ptransport!(d1trans, d1l, t1, t2, pt_cache)
    dot3!(cosτ, d1r, d1trans)
    dot3!(sinτ, d2r, d1trans)
    τ .= -1.0.*atan.(sinτ, cosτ)
end

allocate_cache(f::typeof(full_kinematics!), T::Type, ns) =
    (zeros(T, ns+1, 6), zeros(T, ns, 9))

# Elastic Energy
function elastic_energy(l0,κ0,τ0, l,κ,τ, k,B,β, cache)
    ne = length(l)
    nv = ne - 1
    cache_s = view(cache, 1:ne)
    cache_s .= k./l0.*(l .- l0).^2
    Es = 1/2*sum(cache_s)

    # Reference voronoi lengths
    vl0 = view(cache, 1:nv, 1)
    vl0[1] = l0[1] + l0[2]/2
    vl0[2:nv-1] .= (view(l0, 2:ne-2) .+ view(l0, 3:ne-1))./2
    vl0[nv] = l0[ne-1]/2 + l0[ne]

    cache_t = view(cache, 1:nv, 2)
    cache_t .= β./vl0.*(τ .- τ0).^2
    Et = 1/2*sum(cache_t)

    cache_b = view(cache, 1:nv, 2:3)
    cache_b .= B./vl0.*(κ .- κ0).^2
    Eb = 1/2*sum(cache_b)

    Es + Et + Eb
end

allocate_cache(f::typeof(elastic_energy), T::Type, ns) = zeros(T, ns+1, 3)

# cache is at least ne
function gravitational_energy(x, m, g, cache)
    nv = size(x,1)
    tmp = view(cache, 1:nv-1)
    tmp .= (view(x, 1:nv-1, 3) .+ view(x, 2:nv, 3)).*m
    return sum(tmp)*g/2
end

# ------------------------------------------------------------------------------
# Overload for convenience
# ------------------------------------------------------------------------------
function edges!(e, r::basic_rod)
    edges!(e, r.x)
end

function tangents!(e, r::basic_rod)
    tangents!(e, r.x)
end

function rod_update!(new::basic_rod, old::basic_rod, Δr::rod_delta, cache)
    rod_update!(new.x, new.d, old.x, old.d, Δr.Δx, Δr.Δθ, cache)
end

function full_kinematics!(s::rod_strains, r::basic_rod, cache)
    full_kinematics!(s.l, s.κ, s.τ, r.x, r.d, cache)
end

function elastic_energy(ref_strains::rod_strains, strains::rod_strains,
                        rod_props::rod_properties, cache)
    l0, κ0, τ0 = ref_strains.l, ref_strains.κ, ref_strains.τ
    l, κ, τ = strains.l, strains.κ, strains.τ
    k, B, β = rod_props.k, rod_props.B, rod_props.β
    elastic_energy(l0,κ0,τ0, l,κ,τ, k,B,β, cache)
end

# ------------------------------------------------------------------------------
# Overload for convenience
# ------------------------------------------------------------------------------

# immutable DiffCache{T<:AbstractArray, S<:AbstractArray}
#     du::T
#     dual_du::S
# end
# #
# function DiffCache{chunk_size}(T, size, ::Type{Val{chunk_size}})
#     DiffCache(zeros(T, size...), zeros(ForwardDiff.Dual{nothing,T,chunk_size}, size...))
# end
#
# DiffCache(u::AbstractArray) = DiffCache(eltype(u),size(u),Val{ForwardDiff.pickchunksize(length(u))})
# DiffCache(u::AbstractArray, size) = DiffCache(eltype(u),size,Val{ForwardDiff.pickchunksize(length(u))})
# get_tmp{T<:ForwardDiff.Dual}(dc::DiffCache, ::Type{T}) = dc.dual_du
# get_tmp(dc::DiffCache, T) = dc.du


text:
using Distributions
using SpecialFunctions
using StatsBase
using DelimitedFiles

function higibbs(Y, T, nu0 = 1, s20 = 100.0, eta0 = 1.0, t20 = 100.0, mu0 = 50.0, g20 = 25.0)
    m = size(unique(Y[:,1]), 1)
    # starting value
    ybar = ones(m)
    sv = ones(m)
    n = ones(Int, m)
    for j = 1:m
        yj = Y[ [Y[i,1] == j for i = 1:end], 2]
        ybar[j] = mean(yj)
        sv[j] = var(yj)
        n[j] = size(yj, 1)
    end
    theta = copy(ybar)
    sigma2 = mean(sv)
    mu = mean(theta)
    tau2 = var(theta)

    THETA = ones(T, m)
    # sigma2 mu tau2
    SMT = ones(T, 3)

    for t = 1:T
        # sample theta
        for j = 1:m
            vartheta = 1 / (n[j] / sigma2 + 1 / tau2)
            meantheta = vartheta * ( n[j] * ybar[j] / sigma2 + mu / tau2)
            rnorm = Normal(meantheta, sqrt(vartheta))
            theta[j] = rand(rnorm, 1)[1]
        end

        # sample new sigma2
        nun = nu0 + sum(n)
        ss = nu0 * s20
        for j = 1:m
            ss += sum((Y[ [Y[i, 1] == j for i = 1:end], 2] .- theta[j]).^2)
        end
        rgamma = Gamma(nun/2, 2/ss)
        sigma2 = 1 / rand(rgamma, 1)[1]

        # sample mu
        varmu = 1 / (m / tau2 + 1 / g20)
        meanmu = varmu * (m * mean(theta) / tau2 + mu0 / g20)
        rnorm = Normal(meanmu, sqrt(varmu))
        mu = rand(rnorm, 1)[1]

        # sample 1/tau2
        shapetau = (eta0 + m) / 2
        ratetau = ( eta0 * t20 + sum((theta .- mu).^2) ) / 2
        rgamma = Gamma(shapetau, 1/ratetau)
        tau2 = 1 / rand(rgamma, 1)[1]

        # store results
        THETA[t, :] .= theta
        SMT[t, :] .= [sigma2, mu, tau2]
    end
    return THETA, SMT
end

# run
Y = readdlm("math-score-Y.csv")
THETA, SMT = higibbs(Y, 50);
#using PyPlot
#plt[:hist](SMT[:,2])
#show()

max_stars_repo_name: szcf-weiya/MCMC

max_issues_repo_name: szcf-weiya/MCMC

max_forks_repo_name: szcf-weiya/MCMC


text:
include("2.1 (c+d x)^m (a+b (F^(g (e+f x)))^n)^p.jl")
include("2.2 (c+d x)^m (F^(g (e+f x)))^n (a+b (F^(g (e+f x)))^n)^p.jl")
include("2.3 Miscellaneous exponentials.jl")

max_stars_repo_path: src/rules/2 Exponentials/.jl

max_issues_repo_path: src/rules/2 Exponentials/.jl

max_forks_repo_path: src/rules/2 Exponentials/.jl


text:
# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""
    AbstractScenario

Abstract supertype for scenario objects.
"""
abstract type AbstractScenario end
"""
    Probability

A type-safe wrapper for `Float64` used to represent probability of a scenario occuring.
"""
mutable struct Probability
    π::Float64
end
"""
    probability(scenario::AbstractScenario)

Return the probability of `scenario` occuring.

Is always defined for scenarios created through @scenario. Other user defined scenario types must implement this method to generate a proper probability. The default behaviour is to assume that `scenario` has a `probability` field of type [`Probability`](@ref)

See also: [`Probability`](@ref)
"""
probability(scenario::AbstractScenario)::Float64 = scenario.probability.π
"""
    probability(scenarios::Vector{<:AbstractScenario})

Return the probability of that any scenario in the collection `scenarios` occurs.
"""
probability(scenarios::Vector{<:AbstractScenario}) = sum(probability.(scenarios))
"""
    set_probability!(scenario::AbstractScenario, probability::AbstractFloat)

Set the probability of `scenario` occuring.

Is always defined for scenarios created through @scenario. Other user defined scenario types must implement this method.
"""
function set_probability!(scenario::AbstractScenario, π::AbstractFloat)
    scenario.probability.π = π
end
function Base.zero(::Type{S}) where S <: AbstractScenario
    error("zero not implemented for scenario type: ", S)
end
function Base.show(io::IO, scenario::S) where S <: AbstractScenario
    print(io, "$(S.name.name) with probability $(probability(scenario))")
    scenariotext(io, scenario)
    return io
end
"""
    scenariotext(io::IO, scenario::AbstractScenario)

Custom printout called when printing `scenario`.
"""
function scenariotext(io::IO, scenario::AbstractScenario)
    return io
end
function Base.getindex(ξ::AbstractScenario, field::Symbol)
    return getfield(ξ, field)
end
"""
    ExpectedScenario{S <: AbstractScenario}

Wrapper type around an `AbstractScenario`. Should for convenience be used as the result of a call to `expected`.

See also [`expected`](@ref)
"""
struct ExpectedScenario{S <: AbstractScenario} <: AbstractScenario
    scenario::S

    function ExpectedScenario(scenario::AbstractScenario)
        return new{typeof(scenario)}(scenario)
    end
end
function Base.show(io::IO, scenario::ExpectedScenario{S}) where S <: AbstractScenario
    print(io, "Expected scenario of type $(S.name.name)")
    scenariotext(io, scenario.scenario)
    return io
end
"""
    expected(scenarios::Vector{<:AbstractScenario})

Return the expected scenario out of the collection `scenarios` in an [`ExpectedScenario`](@ref) wrapper.

This is defined through classical expectation: sum([probability(s)*s for s in scenarios]), and is always defined for scenarios created through @scenario, if the requested fields support it.

Otherwise, user-defined scenario types must implement this method for full functionality.

See also [`ExpectedScenario`](@ref)
"""
function expected(scenarios::Vector{S}) where S <: AbstractScenario
    isempty(scenarios) && return ExpectedScenario(zero(S))
    return reduce(expected, scenarios; init = ExpectedScenario(zero(S)))
end
function expected(ξ₁::ExpectedScenario{S}, ξ₂::S) where S <: AbstractScenario
    set_probability!(ξ₁.scenario, 1.0)
    return expected(ξ₁.scenario, ξ₂)
end
function expected(ξ₁::S, ξ₂::ExpectedScenario{S}) where S <: AbstractScenario
    set_probability!(ξ₂.scenario, 1.0)
    return expected(ξ₁, ξ₂.scenario)
end
function expected(ξ₁::ExpectedScenario{S}, ξ₂::ExpectedScenario{S}) where S <: AbstractScenario
    set_probability!(ξ₁.scenario, 1.0)
    set_probability!(ξ₂.scenario, 1.0)
    return expected(ξ₁.scenario, ξ₂.scenario)
end

"""
    Scenario

Conveniece type that adheres to the [`AbstractScenario`](@ref) abstraction. Useful when uncertain parameters are defined using [`@uncertain`](@ref) and instances are created using [`@scenario`](@ref).
"""
struct Scenario{T} <: AbstractScenario
    probability::Probability
    data::T

    function Scenario(data::T; probability::AbstractFloat = 1.0) where T
        return new{T}(Probability(probability), data)
    end

    function Scenario(; probability::AbstractFloat = 1.0, kw...)
        data = values(kw)
        NT = typeof(data)
        return new{NT}(Probability(probability), data)
    end
end

function Base.zero(::Type{Scenario{NT}}) where NT <: NamedTuple
    return Scenario(NamedTuple{Tuple(fieldnames(NT))}(zero.(NT.types)); probability = 1.0)
end
function Base.zero(::Type{Scenario{D}}) where {T, N, D <: Array{T,N}}
    return Scenario(Array{T,N}(undef, ntuple(Val{N}()) do i 0 end); probability = 1.0)
end
function Base.zero(::Type{Scenario{D}}) where {T, N, Ax <: NTuple, D <: DenseAxisArray{T,N,Ax}}
    axs = ntuple(Val{N}()) do i Ax.types[i]() end
    return Scenario(DenseAxisArray(Array{T,N}(undef, ntuple(Val{N}()) do i 0 end), axs...); probability = 1.0)
end
function Base.zero(::Type{Scenario{D}}) where {T, N, K, D <: SparseAxisArray{T,N,K}}
    return Scenario(Dict{K,T}(); probability = 1.0)
end
function scenariotext(io::IO, scenario::Scenario{NT}) where NT <: NamedTuple
    for (k,v) in pairs(scenario.data)
        print(io, "\n  $k: $v")
    end
    return io
end
function scenariotext(io::IO, scenario::Scenario)
    print(io, " and underlying data:\n\n")
    print(io, scenario.data)
    return io
end

Scenarios{S <: AbstractScenario} = Vector{S}
ScenarioTypes{N} = NTuple{N, Union{DataType, UnionAll}}

function expected(ξ₁::Scenario{NT}, ξ₂::Scenario{NT}) where NT <: NamedTuple
    keys(ξ₁.data) == keys(ξ₂.data) || error("Inconsistent scenarios. $(keys(ξ₁)) and $(keys(ξ₂)) do not match.")
    data = map(zip(values(ξ₁.data), values(ξ₂.data))) do (x, y)
        probability(ξ₁) * x + probability(ξ₂) * y
    end
    expected = Scenario(NamedTuple{Tuple(keys(ξ₁.data))}(data))
    return ExpectedScenario(expected)
end

function expected(ξ₁::Scenario{D}, ξ₂::Scenario{D}) where D <: Array
    if isempty(ξ₁.data)
        weighted = Scenario(probability(ξ₂) * ξ₂.data)
        return ExpectedScenario(weighted)
    end
    if isempty(ξ₂.data)
        weighted = Scenario(probability(ξ₁) * ξ₁.data)
        return ExpectedScenario(weighted)
    end
    size(ξ₁.data) == size(ξ₂.data) || error("Inconsistent scenarios. $(size(ξ₁.data)) and $(size(ξ₂.data)) do not match.")
    expected = Scenario(probability(ξ₁) * ξ₁.data + probability(ξ₂) * ξ₂.data)
    return ExpectedScenario(expected)
end

function expected(ξ₁::Scenario{D}, ξ₂::Scenario{D}) where D <: DenseAxisArray
    if isempty(ξ₁.data.data)
        weighted = Scenario(DenseAxisArray(probability(ξ₂) * ξ₂.data.data, axes(ξ₂.data)...))
        return ExpectedScenario(weighted)
    end
    if isempty(ξ₂.data.data)
        weighted = Scenario(DenseAxisArray(probability(ξ₁) * ξ₁.data.data, axes(ξ₁.data)...))
        return ExpectedScenario(weighted)
    end
    axes(ξ₁.data) == axes(ξ₂.data) || error("Inconsistent scenarios. $(axes(ξ₁.data)) and $(axes(ξ₂.data)) do not match.")
    size(ξ₁.data) == size(ξ₂.data) || error("Inconsistent scenarios. $(size(ξ₁.data)) and $(size(ξ₂.data)) do not match.")
    data = DenseAxisArray(probability(ξ₁) * ξ₁.data.data + probability(ξ₂) * ξ₂.data.data, axes(ξ₁.data)...)
    expected = Scenario(data)
    return ExpectedScenario(expected)
end

function expected(ξ₁::Scenario{D}, ξ₂::Scenario{D}) where D <: SparseAxisArray
    if isempty(ξ₁.data.data)
        weighted = Scenario(Dict([key => probability(ξ₂) * ξ₂.data.data[key] for key in keys(ξ₂.data.data)]))
        return ExpectedScenario(weighted)
    end
    if isempty(ξ₂.data.data)
        weighted = Scenario(Dict([key => probability(ξ₁) * ξ₁.data.data[key] for key in keys(ξ₁.data.data)]))
        return ExpectedScenario(weighted)
    end
    keys(ξ₁.data.data) == keys(ξ₂.data.data) || error("Inconsistent scenarios. $(keys(ξ₁.data.data)) and $(keys(ξ₂.data.data)) do not match.")
    expected = Scenario(Dict([key => probability(ξ₁) * ξ₁.data.data[key] + probability(ξ₂) * ξ₂.data.data[key] for key in keys(ξ₁.data.data)]))
    return ExpectedScenario(expected)
end


text:
using BenchmarkTools, Printf
using StaticKernels

_bench_i = 1
if !@isdefined _bench_times
    _bench_times = Float64[]
end

amean(v) = mapreduce(x->x, +, v) / length(v)
gmean(v) = mapreduce(x->x^(1/length(v)), *, v)
hmean(v) = length(v) / mapreduce(x->1/x, +, v)

function showreg(trial, tol=0.05)
    global _bench_times, _bench_i

    while length(_bench_times) < _bench_i
        push!(_bench_times, 0)
    end

    old = _bench_times[_bench_i]
    new = minimum(trial.times)

    rel = (new - old) / old

    str = abs(rel) < tol ? "→" : new < old ? "↓" : "↑"

    @printf("  [%3d]: %10.0f %s %10.0f (%+1.3f)\n", _bench_i, round(old), str, round(new), rel)

    _bench_times[_bench_i] = new
    _bench_i += 1

    return nothing
end

macro mybench(x)
    b = esc(:(@benchmark($x, seconds=0.3)))
    return :(showreg($b))
end

# Benchmarks

function laplace(n)
    a = rand(n, n)
    k = Kernel{(-1:1,-1:1)}(@inline function f(w) @inbounds w[0,-1] + w[-1,0] - 4*w[0,0] + w[1,0] + w[0,1] end)
    b = similar(a, size(k, a))

    @mybench map!($k, $b, $a)
end
println("laplace")
foreach(laplace, (1000, 3000))

a = extend(rand(1000, 1000), StaticKernels.ExtensionConstant(0))
@inline kf(w) = @inbounds w[0,-1] + w[-1,0] - 4*w[0,0] + w[1,0] + w[0,1]
k = Kernel{(-1:1,-1:1)}(kf)
b = similar(a, size(k, a))

@mybench map!($k, $b, $a)


text:
type FEProblem
    name::ASCIIString
    nodes::Vector{FENode2}
    bcs::Vector{DirichletBC}
    loads::Vector{NodeLoad}
    sections::Vector{FESection}
    node_doftype_bc::Dict{Tuple{Int, DofType}, DirichletBC}
    doftype_eqs::Dict{DofType, Vector{Int}}
    dof_vals::DofVals
    n_eqs::Int
    n_fixed::Int
end

#=
function FEProblem(name::ASCIIString, geomesh)
    nodes::Vector{FENode2}(0)
    bcs::Vector{DirichletBC}(0)
    loads::Vector{NodalLoad}(0)
    sections::Vector{FESection}(0)
    node_doftype_bc::Dict{(Int, DofType), DirichletBC}()
    dof_vals::Vector{Float64}(0)
    n_eqs=0
    n_fixed=0
end
=#

push!(fp::FEProblem, bc::DirichletBC) = push!(fp.bcs, bc)
push!(fp::FEProblem, load::NodeLoad) = push!(fp.loads, load)

function FEProblem(name::ASCIIString, nodes::Vector{FENode2}, bcs,
                    loads, sections)
    node_doftype_bc = Dict{Int, Vector{DofType}}()
    doftype_eqs = Dict{DofType, Int}()
    dof_vals = DofVals()
    FEProblem(name, nodes, bcs, loads, sections, node_doftype_bc, doftype_eqs, dof_vals, 0, 0)
end

function create_feproblem(name, geomesh, element_regions, material_regions,
                          bcs, loads=NodeLoad[])

    gps = Dict{DataType, Vector{GaussPoint2}} ()
    interps = Dict{DataType, AbstractInterpolator} ()
    storage = Dict{DataType, AbstractElemStorage} ()
    elem_types = Array(DataType, 0)

    for element_region in element_regions
        elem_type = element_region.element_type
        interps[elem_type] = createinterp(elem_type)
        gps[elem_type] = creategps(elem_type)
        storage[elem_type] = createstorage(elem_type)
    end

    nodes = Array(FENode2, 0)
    for node in geomesh.nodes
        push!(nodes, FENode2(node.n, node.coords))
    end

    sections = Array(FESection, 0)
    for matregion in material_regions
        material = matregion.material
        matstat = create_matstat(typeof(material))
        for eleregion in element_regions
            ele_type = eleregion.element_type
            common = intersect(matregion.elements, eleregion.elements)
            common = collect(common)
            sort!(common) # TODO: Remove?
            gps_ele = gps[ele_type]
            elem_storage = storage[ele_type]
            interp = interps[ele_type]

            section = FESection(material, ele_type, typeof(matstat))

            for ele_id in common
                vertices = geomesh.elements[ele_id].vertices
                element = ele_type(vertices, ele_id, interp,
                                   elem_storage, gps_ele, matstat)
                push!(section, element)

            end
            push!(sections, section)
        end
    end
    fe = FEProblem(name, nodes, bcs, loads, sections)
    createdofs(fe)
    return fe
end

# Type stability sanitized
function set_dof_types_section!{T<:FESection}(section::T,
                                       node_doftypes::Dict{Int, Vector{DofType}})
    for element in section.elements
        for (i, v) in enumerate(element.vertices)
            dof_types = doftypes(element, i)
            node_doftypes[v] = dof_types
        end
    end
end

function createdofs(fp::FEProblem)

    node_doftypes = Dict{Int, Vector{DofType}}()
    # Create a dictionary between a node_id to
    # what dof types are in that node.
    eq_n = 1
    for section in fp.sections
        set_dof_types_section!(section, node_doftypes)
    end

    # Create a dictionary between a tuple of node_id
    # and dof type to the BC for that tuple.
    for bc in fp.bcs
        node_ids = bc.node_set.node_ids
        for node_id in node_ids
            node_number = fp.nodes[node_id].n
            for doftype in bc.dof_types
                fp.node_doftype_bc[(node_id, doftype)] = bc
            end
        end
    end

    # Create the dofs
    eq_n = 0
    pres_n = 0
    id = 0
    dofs = Array(Dof, 0)

    # TODO: Optimize
    for node in fp.nodes
        for doftype in node_doftypes[node.n]
            if haskey(fp.node_doftype_bc, (node.n, doftype))
                bc = fp.node_doftype_bc[(node.n, doftype)]
                pres_n += 1
                push!(node.dofs, Dof(pres_n, doftype, false))
            else
                eq_n += 1
                push!(node.dofs, Dof(eq_n, doftype, true))
                try
                    fp.doftype_eqs[doftype]
                catch KeyError
                    fp.doftype_eqs[doftype] = Vector{Int}(0)
                end
                push!(fp.doftype_eqs[doftype], eq_n)
            end
        end
    end
    #TODO: Make these have same name
    resize!(fp.dof_vals.free_dof_values, eq_n)
    fill!(fp.dof_vals.free_dof_values, 0.0)
    resize!(fp.dof_vals.presc_dof_values, pres_n)
    fill!(fp.dof_vals.presc_dof_values, 0.0)
    fp.n_eqs = eq_n
    fp.n_fixed = pres_n
end

function extload(fp::FEProblem, t::Number=0.0)
    f = zeros(fp.n_eqs)
    for load in fp.loads
        for node_id in load.node_set.node_ids
            for dof in fp.nodes[node_id].dofs
                if dof.dof_type in load.dof_types
                     extload(f, load, dof, fp.nodes[node_id], t)
                end
            end
        end
    end
    return f
end

@inline function extload{T}(f::Vector{Float64}, load::NodeLoad{T},
                            dof::Dof, node::FENode2, t::Number)
    f[dof.eq_n] += evaluate(load, node, t)
end

function assembleK!(K::SparseMatrixCSC, fp::FEProblem, colptrs::Vector{Int}, dofvals::DofVals)
    fill!(K.nzval, 0.0)
    z = 1
    for section in fp.sections
        z = assembleK!(K, colptrs, z, section, fp.nodes, dofvals)
    end
    return K
end

function assembleK!(K::SparseMatrixCSC, colptrs::Vector{Int}, z::Int,
                    section::FESection, nodes::Vector{FENode2}, dofvals::DofVals)
    mat = section.material
    for element in section.elements
        Ke = stiffness(element, nodes, mat, dofvals)
        dof1_n = 0
        for vertex1 in element.vertices
            for dof1 in nodes[vertex1].dofs
                dof1_n += 1
                if dof1.active
                    dof2_n = 0
                    for vertex2 in element.vertices
                        for dof2 in nodes[vertex2].dofs
                            dof2_n += 1
                            if dof2.active
                                K.nzval[colptrs[z]] += Ke[dof1_n, dof2_n]
                                z+=1
                            end
                        end
                    end
                end
            end
        end
    end
    return z
end

# This is nicer but the activedofs iterators are too slow right now
#=
function assembleK!(K::SparseMatrixCSC, colptrs::Vector{Int}, z::Int,
                    section::FESection, nodes::Vector{FENode2})
    mat = section.material
    for element in section.elements
        Ke = stiffness(element, nodes, mat)
        for (dof1, i) in activedofs(element, nodes)
            for (dof2, j) in activedofs(element, nodes)
                v = K[dof1.eq_n, dof2.eq_n]
                K.nzval[colptrs[z]] += Ke[i, j]
                z += 1
            end
        end
    end
    return z
end
=#

function create_sparse_structure(fp::FEProblem)
    dof_rows = Int[]
    dof_cols = Int[]
    for section in fp.sections
        create_sparse_structure(section, fp.nodes, dof_rows, dof_cols)
    end
    # Using ones until we get a sparse structure initializer in base
    return Base.sparse(dof_rows, dof_cols, ones(length(dof_rows)), fp.n_eqs, fp.n_eqs)
end

function create_sparse_structure(section::FESection, nodes::Vector{FENode2},
                                dof_rows::Vector{Int}, dof_cols::Vector{Int})
    for element in section.elements
                dof1_n = 0
        for vertex1 in element.vertices
            for dof1 in nodes[vertex1].dofs
                dof1_n += 1
                if dof1.active
                    dof2_n = 0
                    for vertex2 in element.vertices
                        for dof2 in nodes[vertex2].dofs
                            dof2_n += 1
                            if dof2.active
                                push!(dof_rows, dof1.eq_n)
                                push!(dof_cols, dof2.eq_n)
                             end
                        end
                    end
                end
            end
        end
    end
end

function assembleK(fp::FEProblem)
    dof_rows = Array(Int, 0)
    dof_cols = Array(Int, 0)
    k_values = Array(Float64, 0)
    for section in fp.sections
        assemble_K_section(section, fp.nodes, dof_rows, dof_cols, k_values)
    end

    return Base.sparse(dof_rows, dof_cols, k_values, fp.n_eqs, fp.n_eqs)
end

function assemble_K_section{T<:FESection}(section::T, nodes::Vector{FENode2},
                                          dof_rows::Vector{Int}, dof_cols::Vector{Int},
                                          k_values::Vector{Float64}, dofvals::DofVals)
    mat = section.material
    for element in section.elements
        Ke = stiffness(element, nodes, mat, dofvals)
        dof1_n = 0
        for vertex1 in element.vertices
            for dof1 in nodes[vertex1].dofs
                dof1_n += 1
                if dof1.active
                    dof2_n = 0
                    for vertex2 in element.vertices
                        for dof2 in nodes[vertex2].dofs
                            dof2_n += 1
                            if dof2.active
                                push!(dof_rows, dof1.eq_n)
                                push!(dof_cols, dof2.eq_n)
                                push!(k_values, Ke[dof1_n, dof2_n])
                            end
                        end
                    end
                end
            end
        end
    end
end

function assemble_intf(fp::FEProblem)
    fint = zeros(fp.n_eqs)
    for section in fp.sections
        assemble_intf_section(section, fint, fp.nodes, fp.dof_vals)
    end
    return fint
end

function assemble_intf_section{T<:FESection}(section::T,
                                            int_forces::Vector{Float64},
                                            nodes::Vector{FENode2},
                                            dof_vals::DofVals)
    mat = section.material
    for element in section.elements
        finte = intf(element, mat, nodes, dof_vals)
        i = 1
        for vertex in element.vertices
            for dof in nodes[vertex].dofs
                if dof.active
                    int_forces[dof.eq_n] += finte[i]
                end
                i+=1
            end
        end

    end
end

function updatedofs!(fp::FEProblem, du::Vector{Float64})
    @devec fp.dof_vals.free_dof_values[:] += du
end

function updatebcs!(fp::FEProblem, t::Number=0.0)
    for node in fp.nodes
        for dof in node.dofs
            if !dof.active
                bc = fp.node_doftype_bc[(node.n, dof.dof_type)]
                updatebc!(bc, fp.dof_vals, dof.eq_n, node, t)
            end
        end
    end
end

@inline function updatebc!{f}(bc::DirichletBC{f}, dof_vals::DofVals, eq_n:: Int, node::FENode2, t::Number)
    dof_vals.presc_dof_values[eq_n] = evaluate(bc, node, t)
end

function update_feproblem(fp::FEProblem)
    for section in fp.sections
        for element in section.elements
            for i in 1:length(element.matstats)
                element.matstats[i] = copy(element.temp_matstats[i])
            end
        end
    end
end


text:
using DataStructures: OrderedDict

function k8s_create(manifest::IO)
    open(`$(kubectl()) apply --force -f -`, "w", stdout) do p
        write(p.in, read(manifest))
    end
end

function get_job(name::AbstractString; jsonpath=nothing)
    output = if jsonpath !== nothing
        "jsonpath=$jsonpath"
    else
        "json"
    end

    kubectl_cmd = `$(kubectl()) get job/$name -o $output`
    err = IOBuffer()
    out = read(pipeline(ignorestatus(kubectl_cmd), stderr=err), String)

    err.size > 0 && throw(KubeError(err))
    return output == "json" ? JSON.parse(out; dicttype=OrderedDict) : out
end

function wait_job(job_name; condition=!isempty, timeout=60)
    timedwait(timeout; pollint=10) do
        condition(get_job(job_name, jsonpath="{.status..type}"))
    end
end

pod_exists(pod_name) = success(`$(kubectl()) get pod/$pod_name`)

# Will fail if called and the job is in state "Waiting"
pod_logs(pod_name) = read(ignorestatus(`$(kubectl()) logs $pod_name`), String)

# https://kubernetes.io/docs/concepts/workloads/pods/pod-lifecycle/#pod-phase
pod_phase(pod_name) = read(`$(kubectl()) get pod/$pod_name -o 'jsonpath={.status.phase}'`, String)

function pod_names(labels::Pair...; sort_by=nothing)
    selectors = Dict{String,String}(labels)
    selector = join(map(p -> join(p, '='), collect(pairs(selectors))), ',')

    # Adapted from: https://kubernetes.io/docs/concepts/workloads/controllers/job/#running-an-example-job
    jsonpath = "{range .items[*]}{.metadata.name}{\"\\n\"}{end}"
    sort_by_opt = sort_by !== nothing ? `--sort-by=$sort_by` : ``
    kubectl_cmd = `$(kubectl()) get pods -l $selector -o=jsonpath=$jsonpath $sort_by_opt`
    output = readchomp(kubectl_cmd)
    return !isempty(output) ? split(output, '\n') : String[]
end

# Use the double-quoted flow scalar style to allow us to have a YAML string which includes
# newlines without being aware of YAML indentation (block styles)
#
# The double-quoted style allows us to use escape sequences via `\` but requires us to
# escape uses of `\` and `"`. It so happens that `escape_string` follows the same rules
escape_yaml_string(str::AbstractString) = escape_string(str)

randsuffix(len=5) = randstring(['a':'z'; '0':'9'], len)

function report(job_name, pods::Pair...)
    kubectl_cmd = `$(kubectl()) describe job/$job_name`
    @info "Describe job:\n" * read(ignorestatus(kubectl_cmd), String)

    # Note: A job doesn't contain a direct reference to the pod it starts so
    # re-using the job name can result in us identifying the wrong manager pod.
    kubectl_cmd = `$(kubectl()) get pods -L job-name=$job_name`
    @info "List pods for job $job_name:\n" * read(ignorestatus(kubectl_cmd), String)

    for (title, pod_name) in pods
        if pod_exists(pod_name)
            kubectl_cmd = `$(kubectl()) describe pod $pod_name`
            @info "Describe $title pod:\n" * read(kubectl_cmd, String)
        else
            @info "$(titlecase(title)) pod \"$pod_name\" not found"
        end
    end

    for (title, pod_name) in pods
        if pod_exists(pod_name)
            @info "Logs for $title ($pod_name):\n" * pod_logs(pod_name)
        else
            @info "No logs for $title ($pod_name)"
        end
    end
end


text:
module ReadVTK

using Base64: base64decode
using Downloads: download

using CodecZlib: ZlibDecompressor
using LightXML: LightXML, XMLDocument, XMLElement, parse_string, attribute, has_attribute,
                child_elements, free, content, find_element

export VTKFile, get_points, get_cells, VTKData, get_point_data, get_cell_data, VTKDataArray,
       get_data, VTKCells, get_example_file

"""
    VTKFile

Hold all relevant information about a VTK XML file that has been read in.

# Fields
- `filename`: original path to the VTK file that has been read in
- `xml_file`: object that represents the XML file
- `file_type`: currently only `UnstructuredGrid` is supported
- `version`: currently only v1.0 is supported
- `byte_order`: can be `LittleEndian` or `BigEndian` and must currently be the same as the system's
- `compressor`: can be empty (no compression) or `vtkZLibDataCompressor`
- `appended_data`: in case of appended data (see XML documentation), the data is stored here for
                   convenient retrievel (otherwise it is empty)
- `n_points`: number of points in the VTK file
- `n_cells`: number of cells in the VTK file`
"""
mutable struct VTKFile
  filename::String
  xml_file::XMLDocument
  file_type::String
  version::VersionNumber
  byte_order::String
  compressor::String
  appended_data::Vector{UInt8}
  n_points::Int
  n_cells::Int

  # Create inner constructor to add finalizer (this requires VTKFile to be a *mutable* struct)
  function VTKFile(filename, xml_file, file_type, version, byte_order, compressor, appended_data,
                   n_points, n_cells)
    # Create new object
    vtk_file = new(filename, xml_file, file_type, version, byte_order, compressor, appended_data,
                   n_points, n_cells)

    # Create finalizer that releases memory for VTK file
    f(vtk_file) = free(vtk_file.xml_file)

    # Register finalizer
    finalizer(f, vtk_file)
  end
end

# Header type is hardcoded and corresponds to VTK XML version 1.0
header_type(::VTKFile) = UInt64

# Return true if data is compressed (= XML attribute `compressor` is non-empty in VTK file)
is_compressed(vtk_file::VTKFile) = !isempty(vtk_file.compressor)

"""
    VTKFile(filename)

Read in and parse the VTK XML file specified by its `filename`.
"""
function VTKFile(filename)
  # Read in file into memory as a string
  raw_file_contents = read(filename, String)

  # Check if file begins with string that indicates this is a VTK file but *not* in XML format
  if startswith(raw_file_contents, "# vtk DataFile Version")
    error("bad VTK file format (found legacy format; only VTK XML files are currently supported)")
  end

  # Check if data is stored in "appended" mode
  marker = findfirst("<AppendedData encoding=\"raw\">", raw_file_contents)
  if isnothing(marker)
    # `appended_data` remains empty and the LightXML library may parse the entire string
    appended_data = UInt8[]
    xml_file_contents = raw_file_contents
  else
    # Since data is stored in appended mode, the VTK file is technically not valid XML. Therefore we
    # extract the appended data and store it as a byte stream in `appended_data`. We then clean up
    # the string (i.e., remove appended data) and pass a valid XML file string to the LightXML
    # library
    offset_begin = first(findnext("_", raw_file_contents, last(marker))) + 1
    offset_end = first(findnext("</AppendedData>", raw_file_contents, offset_begin)) - 1
    appended_data = Vector{UInt8}(rstrip(raw_file_contents[offset_begin:offset_end]))
    xml_file_contents = raw_file_contents[1:offset_begin-1] * "\n  </AppendedData>\n</VTKFile>"
  end

  # Open file and ensure that it is a valid VTK file
  xml_file = parse_string(xml_file_contents)
  root = LightXML.root(xml_file)
  @assert LightXML.name(root) == "VTKFile"

  # Extract attributes (use `required=true` to fail fast & hard in case of unexpected content)
  file_type = attribute(root, "type", required=true)
  version = VersionNumber(attribute(root, "version", required=true))
  byte_order = attribute(root, "byte_order", required=true)
  header_type = attribute(root, "header_type", required=true)
  if has_attribute(root, "compressor")
    compressor = attribute(root, "compressor", required=true)
  else
    compressor = ""
  end

  # Ensure matching file type
  @assert file_type == "UnstructuredGrid"

  # Ensure correct version
  @assert version == v"1.0"

  # Ensure matching byte order
  is_little_endian = ENDIAN_BOM == 0x04030201
  @assert ((byte_order == "LittleEndian" && is_little_endian) ||
           (byte_order == "BigEndian" && !is_little_endian))

  # Ensure matching header type
  @assert header_type == "UInt64"

  # Ensure supported compression type
  @assert in(compressor, ("", "vtkZLibDataCompressor"))

  # Extract piece
  piece = root[file_type][1]["Piece"][1]
  n_points = parse(Int, attribute(piece, "NumberOfPoints", required=true))
  n_cells = parse(Int, attribute(piece, "NumberOfCells", required=true))

  # Create and return VTKFile
  VTKFile(filename, xml_file, file_type, version, byte_order, compressor, appended_data, n_points,
          n_cells)
end

# Show basic information on REPL
function Base.show(io::IO, vtk_file::VTKFile)
  print(io, "VTKFile(",
            "\"", vtk_file.filename, "\", ",
            "<XMLDocument>", ", ",
            "\"", vtk_file.file_type, "\", ",
            "\"", vtk_file.version, "\", ",
            "\"", vtk_file.byte_order, "\", ",
            "\"", vtk_file.compressor, "\", ",
            "<appended_data>", ", ",
            vtk_file.n_points, ", ",
            vtk_file.n_cells, ")")
end

# Return `Piece` XML element that contains all VTK data
piece(vtk_file::VTKFile) = LightXML.root(vtk_file.xml_file)[vtk_file.file_type][1]["Piece"][1]

"""
    VTKData

Convenience type to hold information about available data arrays for cells or points.

Supports a collectible/dictionary-like syntax, e.g., `keys(vtk_data)` to show available data arrays
or `vtk_data["varname"]` to retrieve the `VTKDataArray` for variable `varname`.
"""
struct VTKData
  names::Vector{String}
  data_arrays::Vector{XMLElement}
  vtk_file::VTKFile
end

# Reduce REPL noise by defining `show`
Base.show(io::IO, data::VTKData) = print(io, "VTKData()")

# Retrieve a data section (should be `CellData` or `PointData`) from the VTK file
function get_data_section(vtk_file, section)
  names = String[]
  data_arrays = XMLElement[]

  # Iterate over XML elemens in the section
  for xml_element in child_elements(piece(vtk_file)[section][1])
    # We do not know how to handle anything other than `DataArray`s
    @assert LightXML.name(xml_element) == "DataArray"

    # Store the name and the XML element for each found data array
    push!(names, attribute(xml_element, "Name", required=true))
    push!(data_arrays, xml_element)
  end

  VTKData(names, data_arrays, vtk_file)
end

"""
    get_cell_data(vtk_file)

Retrieve a lightweight `VTKData` object with the cell data of the given VTK file.

See also: [`VTKData`](@ref), [`get_point_data`](@ref)
"""
get_cell_data(vtk_file) = get_data_section(vtk_file, "CellData")

"""
    get_point_data(vtk_file)

Retrieve a lightweight `VTKData` object with the point data of the given VTK file.

See also: [`VTKData`](@ref), [`get_cell_data`](@ref)
"""
get_point_data(vtk_file) = get_data_section(vtk_file, "PointData")

# Auxiliary methods for conveniently using `VTKData` objects like a dictionary/collectible
Base.firstindex(data::VTKData) = first(data.names)
Base.lastindex(data::VTKData) = last(data.names)
Base.length(data::VTKData) = length(data.names)
Base.size(data::VTKData) = (length(data),)
Base.keys(data::VTKData) = tuple(data.names...)

function Base.iterate(data::VTKData, state=1)
  if state > length(data)
    return nothing
  else
    return (data.names[state] => data[data.names[state]], state + 1)
  end
end

function Base.getindex(data::VTKData, name)
  index = findfirst(isequal(name), data.names)

  if isnothing(index)
    throw(KeyError(name))
  end

  return VTKDataArray(data.data_arrays[index], data.vtk_file)
end

Base.eltype(::VTKData) = Pair{String, VTKDataArray}

"""
    VTKDataArray{T, N, Format}

Hold information about a single VTK data array (cell data or point data).

The data type is encoded as `T`, `N` represents the size of the second dimension in case of
multi-dimensonal arrays (or `1` for a vector), and `Format` encodes in which format the data is
stored in the XML file.

The actual data can be retrieved by calling `get_data` on the `VTKDataArray` object.

See also: [`get_data`](@ref)
"""
struct VTKDataArray{T, N, Format}
  name::String
  offset::Int
  data_array::XMLElement
  vtk_file::VTKFile
end

# Auxiliary types for type stability
struct FormatBinary end
struct FormatAppended end
struct FormatAscii end

# Types taken from https://vtk.org/doc/release/7.0/html/vtkType_8h_source.html
function string_to_data_type(s)
  if s == "Int8"
    Int8
  elseif s == "Int16"
    Int16
  elseif s == "Int32"
    Int32
  elseif s == "Int64"
    Int64
  elseif s == "UInt8"
    UInt8
  elseif s == "UInt16"
    UInt16
  elseif s == "UInt32"
    UInt32
  elseif s == "UInt64"
    UInt64
  elseif s == "Float32"
    Float32
  elseif s == "Float64"
    Float64
  else
    error("unknown data type: ", s)
  end
end

"""
    VTKDataArray(xml_element, vtk_file)

Create a lightweight container for a given `xml_element` and the corresponding `vtk_file`.

The `xml_element` must be of type `DataArray` and the `vtk_file` parameter is required for
retrieving the actual data. You can pass a `VTKDataArray` object to `get_data` to retrieve the actual data.

See also: [`get_data`](@ref)
"""
function VTKDataArray(xml_element, vtk_file)
  # Ensure the correct type of the XML element
  @assert LightXML.name(xml_element) == "DataArray"

  # Extract information about the underlying data
  data_type = string_to_data_type(attribute(xml_element, "type", required=true))
  name = attribute(xml_element, "Name", required=true)
  n_components = parse(Int, attribute(xml_element, "NumberOfComponents", required=true))
  format_string = attribute(xml_element, "format", required=true)

  # An offset is only used when the format is `appended`
  if has_attribute(xml_element, "offset")
    @assert format_string == "appended"
    offset = parse(Int, attribute(xml_element, "offset", required=true))
  else
    offset = -1
  end

  # Convert string to type for type stability later on
  if format_string == "binary"
    format = FormatBinary
  elseif format_string == "appended"
    format = FormatAppended
  elseif format_string == "ascii"
    format = FormatAscii
  else
    error("unknown data array format: ", format_string)
  end

  VTKDataArray{data_type, n_components, format}(name, offset, xml_element, vtk_file)
end

# Reduce REPL noise by defining `show`
Base.show(io::IO, data_array::VTKDataArray) = print(io, "VTKDataArray(\"", data_array.name, "\")")

# Return true if data is compressed (= XML attribute `compressor` is non-empty in VTK file)
is_compressed(data_array::VTKDataArray) = is_compressed(data_array.vtk_file)

"""
    get_data(data_array::VTKDataArray)

Retrieve actual data from a `VTKDataArray` as a one- or two-dimensional array-like container.

Note: This function is not type stable but could be - help wanted!
"""
function get_data(data_array::VTKDataArray) end

# Retrieve actual data for XML data array (version for storage format "binary")
function get_data(data_array::VTKDataArray{T,N,<:FormatBinary}) where {T,N}
  # To retrieve the raw data,
  # * first get the content of the corresponding XML data array
  # * then remove leading/trailing whitespace
  # * finally decode from Base64 to binary representation
  raw = base64decode(strip(content(data_array.data_array)))

  if is_compressed(data_array)
    # If data is stored compressed, the first four integers of type `header_type` are the header and
    # must be discarded
    first = 1 + 4 * sizeof(header_type(data_array.vtk_file))
    last = length(raw)

    # Pass data through ZLib decompressor
    uncompressed = transcode(ZlibDecompressor, raw[first:last])
  else
    # If data is stored uncompressed, the first integer of type `header_type` is the header and must
    # be discarded
    first = 1 + sizeof(header_type(data_array.vtk_file))
    last = length(raw)
    uncompressed = view(raw, first:last)
  end

  # Convert from binary representation to actual type
  v = reinterpret(T, uncompressed)

  if N == 1
    # If it is a one dimensional array, just return it as is
    return v
  else
    # Otherwise, reshape into two-dimensional array
    return reshape(v, N, :)
  end
end

# Retrieve actual data for XML data array (version for storage format "appended")
function get_data(data_array::VTKDataArray{T,N,<:FormatAppended}) where {T,N}
  raw = data_array.vtk_file.appended_data
  HeaderType = header_type(data_array.vtk_file)

  if is_compressed(data_array)
    # If data is stored compressed, the first four integers of type `header_type` are the header and
    # the fourth value contains the number of bytes to read
    first = data_array.offset + 1
    last = data_array.offset + 4 * sizeof(HeaderType)
    header = Int.(reinterpret(HeaderType, raw[first:last]))
    n_bytes = header[4]

    first = data_array.offset + 4 * sizeof(HeaderType) + 1
    last = first + n_bytes - 1
    uncompressed = transcode(ZlibDecompressor, raw[first:last])
  else
    # If data is stored uncompressed, the first integer of type `header_type` is the header and
    # contains the number of bytes to read
    first = data_array.offset + 1
    last = data_array.offset + sizeof(HeaderType)
    header = Int.(reinterpret(HeaderType, raw[first:last]))
    n_bytes = header[1]

    first = data_array.offset + sizeof(HeaderType) + 1
    last = first + n_bytes - 1
    uncompressed = view(raw, first:last)
  end

  v = reinterpret(T, uncompressed)
  if N == 1
    # If it is a one dimensional array, just return it as is
    return v
  else
    # Otherwise, reshape into two-dimensional array
    return reshape(v, N, :)
  end
end

"""
    get_points(vtk_file)

Retrieve VTK points as a two-dimensional array-like container.

The points are stored in dimension × points format. Note that in VTK, points are always stored
three-dimensional, even for 1D or 2D files.

See also: [`get_cells`](@ref)
"""
function get_points(vtk_file)
  # Retrieve `Points` section
  points = find_element(piece(vtk_file), "Points")
  @assert !isnothing(points)

  # Get the first `DataArray` XML element (there should be only one)
  xml_data_array = find_element(points, "DataArray")
  @assert !isnothing(xml_data_array)

  # Create data array and return actual data
  data_array = VTKDataArray(xml_data_array, vtk_file)

  get_data(data_array)
end

"""
    VTKCells{Connectivity, Offsets, Types}

Store the `connectivity`, `offsets`, and `types` information from the VTK file as one-dimensional
array-like containers. See the
[VTK file format documentation](http://vtk.org/VTK/img/file-formats.pdf) for information on how to
connect the `connectivity` and `offset` arrays to the actual geometric points.

You may use `length` to determine the number of cells from a `VTKCells` object.

See also: [`get_points`](@ref)
"""
struct VTKCells{Connectivity, Offsets, Types}
  connectivity::Connectivity
  offsets::Offsets
  types::Types
end

"""
    get_cells(vtk_file)

Retrieve VTK cell information as an object of type `VTKCells`.

See also: [`VTKCells`](@ref)
"""
function get_cells(vtk_file)
  # Retrieve `Cells` section
  cells = find_element(piece(vtk_file), "Cells")
  @assert !isnothing(cells)

  # Iterate over available XML data arrays and store corresponding data in `VTKDataArray`s
  connectivity = offsets = types = nothing
  for xml_element in cells["DataArray"]
    a = attribute(xml_element, "Name", required=true)
    if a == "connectivity"
      connectivity = VTKDataArray(xml_element, vtk_file)
    elseif a == "offsets"
      offsets = VTKDataArray(xml_element, vtk_file)
    elseif a == "types"
      types = VTKDataArray(xml_element, vtk_file)
    else
      error("unknown cell information: ", a)
    end
  end

  # Ensure that nothing is missing
  @assert !isnothing(connectivity)
  @assert !isnothing(offsets)
  @assert !isnothing(types)

  # Create VTKCells container
  VTKCells(get_data(connectivity), get_data(offsets), get_data(types))
end

# Convenience functions for working with `VTKCells` container
Base.length(vtk_cells::VTKCells) = length(vtk_cells.types)
Base.size(vtk_cells::VTKCells) = (length(vtk_cells),)

"""
    get_example_file(filename; head="main", output_directory=".", force=false)

Retrieve an example file from the
[`ReadVTK_examples` repository](https://github.com/trixi-framework/ReadVTK_examples)
at commit/branch `head` and store it in the `output_directory`. If the file already
exists locally, do not download the file again unless `force` is true. Return the local path to the
downloaded file.
"""
function get_example_file(filename; head="main", output_directory=".", force=false)
  filepath = joinpath(output_directory, filename)
  if !isfile(filepath) || force
    url = ("https://github.com/trixi-framework/ReadVTK_examples/raw/"
           * head
           * "/examples/"
           * filename)
    download(url, filepath)
  end

  return filepath
end

end # module


text:
module LoadStore

# Use LLVM to generate the low level code for these aligned load and store operations.
#
# There are two main pain points to defining these ourselves:
#
# 1. Normal `unsafe_load` and `unsafe_store!` operations don't align based on the data size.
# 2. The normal pointer load/store instructions don't include TBAA data, which means LLVM
#    gets confused when generating code involving loads from these pointers and normal
#    arrays, which results in things not getting vectorized.
#
# This gets around both those issues using `LLVM.jl` to get the Julia `arraybuf` TBAA nodes
# and annotate the `load` and `store` instructions with TBAA access tags derived from the
# `arraybuf` struct.
#
# This should be legal as long as we make sure the types of our arrays are `isbitstype`s.
# If they aren't ... we're in BIG trouble.

using LLVM
using LLVM.Interop

# Type-based Alias Analysis for pointer accesses.
function arraybuf_tbaa(ctx)
    # Get the array_buf tbaa node.
    # Look in `codegen.cpp` for the name of this node.
    tbaa_root = MDNode([MDString("jtbaa"; ctx)]; ctx)
    tbaa_struct = MDNode([
        MDString("jtbaa_arraybuf"; ctx),
        tbaa_root,
        LLVM.ConstantInt(0; ctx)
    ]; ctx)

    tbaa_tag = MDNode([
        tbaa_struct,
        tbaa_struct,
        # Offset - must set to 0
        LLVM.ConstantInt(0; ctx),
        # Constant - set to 0
        LLVM.ConstantInt(0; ctx),
    ]; ctx)

    return tbaa_tag
end

@generated function unsafe_custom_load(ptr::Ptr{T}, i::Integer=1) where {T}
    Context() do ctx
        # Convert the element type
        eltyp = convert(LLVMType, T; ctx)

        # LLVM Int and Pointer types
        T_int = LLVM.IntType(sizeof(Int)*8; ctx)
        T_ptr = LLVM.PointerType(eltyp)

        # Create a function
        paramtyps = [T_int, T_int]
        llvmf, _ = create_function(eltyp, paramtyps)

        # Generate IR
        Builder(ctx) do builder
            entry = BasicBlock(llvmf, "entry"; ctx)
            position!(builder, entry)

            ptr = inttoptr!(builder, parameters(llvmf)[1], T_ptr)
            ptr = gep!(builder, ptr, [parameters(llvmf)[2]])

            val = load!(builder, ptr)

            # Assign alignment
            alignment!(val, sizeof(T))
            # Add tbaa data to match julia's `arraybuf` tbaa.
            metadata(val)[LLVM.MD_tbaa] = arraybuf_tbaa(ctx)

            ret!(builder, val)
        end

        return call_function(llvmf, T, Tuple{Ptr{T}, Int}, :(ptr), :(Int(i-1)))
    end
end

@generated function unsafe_custom_store!(ptr::Ptr{T}, v, i::Integer=1) where {T}
    Context() do ctx
        # Convert the element type
        eltyp = convert(LLVMType, T; ctx)

        # LLVM Int and Pointer Types
        T_int = LLVM.IntType(sizeof(Int) * 8; ctx)
        T_ptr = LLVM.PointerType(eltyp)

        # Create a function
        paramtyps = [T_int, eltyp, T_int]
        llvmf, _ = create_function(LLVM.VoidType(ctx), paramtyps)

        # Generate IR
        Builder(ctx) do builder
            entry = BasicBlock(llvmf, "entry"; ctx)
            position!(builder, entry)

            ptr = inttoptr!(builder, parameters(llvmf)[1], T_ptr)
            ptr = gep!(builder, ptr, [parameters(llvmf)[3]])

            # Emit store instruction
            val = parameters(llvmf)[2]
            st = store!(builder, val, ptr)

            # Add alignment and tbaa information
            alignment!(st, sizeof(T))
            metadata(st)[LLVM.MD_tbaa] = arraybuf_tbaa(ctx)

            ret!(builder)
        end

        return call_function(
            llvmf,
            Cvoid,
            Tuple{Ptr{T}, T, Int},
            :(ptr),
            :(convert(T, v)),
            :(Int(i-1)),
        )
    end
end

#####
##### For testing purposes.
#####

mutable struct UnsafeArray{T,N} <: DenseArray{T,N}
    ptr::Ptr{T}
    dims::NTuple{N,Int}
end

Base.pointer(x::UnsafeArray) = x.ptr
Base.size(x::UnsafeArray) = x.dims
Base.IndexStyle(::Type{<:UnsafeArray}) = Base.IndexLinear()

Base.getindex(x::UnsafeArray, i::Int) = unsafe_custom_load(pointer(x), i)
Base.setindex!(x::UnsafeArray, v, i::Int) = unsafe_custom_store!(pointer(x), v, i)

end


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule Pixman_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("Pixman")
JLLWrappers.@generate_main_file("Pixman", UUID("30392449-352a-5448-841d-b1acce4e97dc"))
end  # module Pixman_jll


text:
# Copyright 2016, 2017, 2018 Martin Holters
# See accompanying license file.

using ACME

# Model of "Der Super Over", see
# http://diy.musikding.de/wp-content/uploads/2013/06/superoverschalt.pdf
# for schematics.
# Note: R19 and LED D5 have been omitted from the model.

function superover(::Type{Circuit}; drive=nothing, tone=nothing, level=nothing, sym::Bool=false)
    circ = @circuit begin
        # power supply
        j3 = voltagesource(9), [+] ⟷ vcc, [-] ⟷ gnd # 9V power supply
        d4 = diode(is=12e-9, η=2), [-] ⟷ vcc, [+] ⟷ gnd # 1N4001
        c11 = capacitor(100e-6), [1] ⟷ vcc, [2] ⟷ gnd
        r17 = resistor(33e3), [1] ⟷ vcc, [2] ⟷ vb
        r18 = resistor(33e3), [1] ⟷ vb, [2] ⟷ gnd
        c12 = capacitor(47e-6), [1] ⟷ vb, [2] ⟷ gnd

        # input stage
        j1 = voltagesource(), [-] ⟷ gnd # input
        r1 = resistor(2.2e6), [1] ⟷ j1[+], [2] ⟷ gnd
        c1 = capacitor(47e-9), [1] ⟷ j1[+]
        r2 = resistor(10e3), [1] ⟷ c1[2]
        r3 = resistor(470e3), [1] ⟷ r2[2], [2] ⟷ vb
        q1 = bjt(:npn, is=80e-15, βf=500, βr=10), [base] ⟷ r2[2], [collector] ⟷ vcc
        r4 = resistor(10e3), [1] ⟷ q1[emitter], [2] ⟷ gnd
        c2 = capacitor(18e-9), [1] ⟷ q1[emitter]
        r5 = resistor(100e3), [1] ⟷ c2[2], [2] ⟷ vb

        # distortion stage
        ic1a = opamp(), ["in+"] ⟷ c2[2], ["out-"] ⟷ gnd
        d1 = diode(is=4e-9, η=2), [-] ⟷ ic1a["out+"], [+] ⟷ ic1a["in-"] # 1N914
        d2 = diode(is=3e-9, η=2), [-] ⟷ ic1a["in-"] # 1N914
        d3 = diode(is=5e-9, η=2), [+] ⟷ ic1a["out+"], [-] ⟷ d2[+] # 1N914
        p1 = potentiometer(1e6, (drive == nothing ? () : (drive,))...), [2] ⟷ [3] ⟷ ic1a["out+"]
        r6 = resistor(33e3), [1] ⟷ ic1a["in-"], [2] ⟷ p1[1]
        c4 = capacitor(47e-9), [1] ⟷ ic1a["in-"]
        r7 = resistor(4.7e3), [1] ⟷ c4[2], [2] ⟷ vb

        # tone control stage
        r8 = resistor(10e3), [1] ⟷ ic1a["out+"]
        ic1b = opamp(), ["in+"] ⟷ r8[2], ["out-"] ⟷ gnd
        c5 = capacitor(18e-9), [1] ⟷ ic1b["in+"], [2] ⟷ gnd
        r10 = resistor(10e3), [1] ⟷ ic1b["out+"], [2] ⟷ ic1b["in-"]
        c7 = capacitor(10e-9), [1] ⟷ ic1b["out+"], [2] ⟷ ic1b["in-"]
        p2 = potentiometer(20e3, (tone == nothing ? () : (tone,))...), [1] ⟷ ic1b["in+"], [3] ⟷ ic1b["in-"]
        c6 = capacitor(27e-9), [1] ⟷ p2[2]
        r11 = resistor(470), [1] ⟷ c6[2], [2] ⟷ gnd

        # output stage
        c8 = capacitor(1e-3), [1] ⟷ ic1b["out+"]
        r12 = resistor(4.7e3), [1] ⟷ c8[2]
        p3 = potentiometer(10e3, (level == nothing ? () : (level,))...), [1] ⟷ vb, [3] ⟷ r12[2]
        r20 = resistor(22e3), [1] ⟷ p3[2]
        c9 = capacitor(47e-9), [1] ⟷ r20[2]
        r13 = resistor(1e6), [1] ⟷ c9[2], [2] ⟷ vb
        q2 = bjt(:npn, is=80e-15, βf=500, βr=10), [base] ⟷ c9[2], [collector] ⟷ vcc
        r14 = resistor(10e3), [1] ⟷ q2[emitter], [2] ⟷ gnd
        r15 = resistor(1e3), [1] ⟷ q2[emitter]
        c10 = capacitor(1e-6), [1] ⟷ r15[2]
        r16 = resistor(100e3), [1] ⟷ c10[2], [2] ⟷ gnd
        j2 = voltageprobe(), [+] ⟷ c10[2], [-] ⟷ gnd # output
    end

    if sym
        connect!(circ, (:d3, :-), (:d3, :+))
    end

    return circ
end

superover(::Type{DiscreteModel}=DiscreteModel;drive=nothing, tone=nothing, level=nothing, sym::Bool=false, fs=44100, solver=nothing) =
    solver === nothing ?
        DiscreteModel(superover(Circuit, drive=drive, tone=tone, level=level, sym=sym), 1//fs) :
        DiscreteModel(superover(Circuit, drive=drive, tone=tone, level=level, sym=sym), 1//fs, solver)