text:
module Continuum

using Bernstein
using Combinatorics
using ComputedFieldTypes
using DifferentialForms
using GrundmannMoeller
using NearestNeighbors
using SimplexPoly
using SparseArrays
using StaticArrays

using ..Funs
using ..Manifolds
using ..SparseOps
using ..ZeroOrOne

################################################################################

export evaluate
"""
Evaluate a function at a point
"""
function evaluate(f::Fun{D,P,R,C,S,T}, x::SVector{C,S}) where {D,P,R,C,S,T}
    @assert C == D
    mfd = f.manifold

    # Find nearest vertex
    if D == 0
        y = first(f.values)
        return Form{D,R,T}((y,))
    end

    simplicesD = get_simplices(mfd, D)
    simplicesD1 = get_simplices(mfd, D - 1)
    lookupD1D = get_lookup(mfd, D - 1, D)
    lookupDD1 = get_lookup(mfd, D, D - 1)
    lookupR = get_lookup(mfd, R, D)
    i, dist = nn(get_simplex_tree(mfd), x)
    i = ID{D}(i)

    found = false
    while !found
        # Simplex vertices
        si = sparse_column_rows(simplicesD, i)
        si = SVector{D + 1}(si[n] for n in 1:(D + 1))
        # Coordinates of simplex vertices
        xs = SVector{D + 1,SVector{C,S}}(get_coords(mfd)[j] for j in si)
        # Calculate barycentric coordinates
        x2λ = cartesian2barycentric_setup(xs)
        λ = cartesian2barycentric(x2λ, x)
        # delta = S(0)
        delta = 10 * eps(S)

        if all(≥(-delta), λ)
            # Found simplex
            @assert P == Pr
            sk = sparse_column_rows(lookupR, i)
            N = DifferentialForms.Forms.binomial(Val(D + 1), Val(R + 1))
            sk = SVector{N}(sk[n] for n in 1:N)
            # Values at `R`-simplices
            ys = f.values[sk]
            # @show D P R x sk ys

            # Evaluate basis functions
            dλ2dx = dbarycentric2dcartesian_setup(Form{D,R}, xs)
            y = zero(Form{D,R,T})
            for n in 1:N
                bx = basis_x(Form{D,R}, x2λ, dλ2dx, n, x)
                # @show bx
                # y += map(b -> ys[n] * b, bx)
                y += ys[n] * bx
            end
            return y
        end

        # Look for neighbouring simplex
        n = argmin(λ)
        n::Int
        @assert 1 ≤ n ≤ D + 1
        # Vertices of face in the direction we need to walk
        si::SVector{D + 1,ID{0}}
        sj = deleteat(si, n)
        sj::SVector{D,ID{0}}
        # All our faces
        sk = sparse_column_rows(lookupD1D, i)
        @assert length(sk) == D + 1
        sk = SVector{D + 1}(sk[n] for n in 1:(D + 1))
        sk::SVector{D + 1,ID{D - 1}}
        # Find face
        ik = findfirst(k -> sparse_column_rows(simplicesD1, k) == sj, sk)
        @assert ik ≢ nothing
        ik::Int
        @assert 1 ≤ ik ≤ D + 1
        k = sk[ik]
        k::ID{D - 1}
        # All the faces neighbouring simplices
        sl = sparse_column_rows(lookupDD1, k)
        @assert length(sl) ≤ 2   # There can be at most 2
        sl = collect(sl)
        sl::Vector{ID{D}}
        @assert any(==(i), sl)   # This must be the current simplex's neighbour
        il = findfirst(≠(i), sl)
        @assert il ≢ nothing     # We ended up at a boundary
        il::Int
        @assert 1 ≤ il ≤ length(sl)
        l = sl[il]
        l::ID{D}
        i = l
    end

    # Dummy return (unreachable)
    return zero(Form{D,R,T})
end

################################################################################

export sample
"""
Sample a function on a manifold

Sample the tangential component at the centre of the `R` simplices.
"""
function sample(::Type{<:Fun{D,P,R,C,S,T}}, f,
                mfd::Manifold{D,C,S}) where {D,P,R,C,S,T}
    @assert P == Pr             # TODO

    # TODO: R > 0 requires a global solve
    @assert R == 0

    simplices = get_simplices(mfd, R)
    coords = get_coords(mfd, R)
    coords0 = get_coords(mfd, 0)
    values = IDVector{R}(Array{T}(undef, length(coords)))
    for i in axes(values, 1)
        xs = coords[i]::SVector{C,S}
        fx = f(xs)::Form{D,R}

        si = sparse_column_rows(simplices, i)
        si = SVector{R + 1}(si[n] for n in 1:(R + 1))
        if C == 0
            t = one(Form{D,R,S})
        else
            xsi = map(x -> Form{C,1,S}(x),
                      coords0[si])::SVector{R + 1,<:Form{C,1,S}}
            ysi = map(x -> x - xsi[1], popfirst(xsi))::SVector{R,<:Form{C,1,S}}
            if isempty(ysi)
                t = one(Form{D,R,S})
            else
                t = ∧(ysi...)::Form{D,R,S}
            end
        end
        val = ((t ⋅ fx) / norm2(t))::Form{D,0}

        values[i] = val[]
    end
    return Fun{D,P,R,C,S,T}(mfd, values)
end

################################################################################

export project
"""
Project a function onto a manifold

    f̃(x) = cⁱ bᵢ(x)

    bⱼ(x) ⋅ f̃(x) = cⁱ bⱼ(x) ⋅ bᵢ(x)
    cⁱ = [bⱼ(x) ⋅ bᵢ(x)]⁻¹ bⱼ(x) ⋅ f̃(x)

    f̃^a(x) = c^i b^a_i(x)

    # tracing over x and over a:
    b_a^j(x) f̃^a(x) = c^i b_a^j(x) b^a_i(x)
    c^i = [b_a^j(x) b^a_i(x)]⁻¹ b_a^j(x) f̃^a(x)
"""
function project(::Type{<:Fun{D,P,R,C,S,T}}, f,
                 mfd::Manifold{D,C,S}) where {D,P,R,C,S,T}
    D::Int
    R::Int
    C::Int
    @assert 0 ≤ R ≤ D ≤ C
    @assert C == D
    N = D + 1
    @assert P == Pr             # TODO

    if D == 0
        value = f(SVector{C,S}())[]
        return Fun{D,P,R,C,S,T}(mfd, IDVector{R}(T[value]))
    end

    order = 4                   # Choice
    scheme = integration_scheme(S, Val(D), Val(order))

    f(zero(SVector{C,S}))::Form{D,R,T}

    values = IDVector{R}(zeros(T, nsimplices(mfd, R)))
    # Loop over all D-simplices
    simplicesD = get_simplices(mfd, D)
    for i in axes(simplicesD, 2)
        si = sparse_column_rows(simplicesD, i)
        @assert length(si) == N
        si = SVector{N}(si[n] for n in 1:N)

        # Coordinates of simplex vertices
        xs = SVector{N,SVector{C,S}}(get_coords(mfd)[n] for n in si)
        x2λ = cartesian2barycentric_setup(xs)
        dλ2dx = dbarycentric2dcartesian_setup(Form{D,R}, xs)

        # Loop over all contained R-simplices
        # TODO: Move this loop inside the integration kernel; return
        # an SVector instead of a scalar
        lookup_RD = get_lookup(mfd, R, D)
        iter = enumerate(sparse_column_rows(lookup_RD, i))
        for (nj, j) in iter
            function kernel(x::SVector{C,S})
                bj = basis_x(Form{D,R}, x2λ, dλ2dx, nj, x)
                return (bj ⋅ f(x))[]
            end
            bf = integrate(kernel, scheme, xs)
            values[j] += bf
        end
    end

    # TODO: memoize this matrix
    # TODO: can we pre-calculate it with barycentric coordinates and
    # then transform it?
    B = basis_products(Val(P), Val(R), mfd)
    @show B
    @show values
    values′ = B \ values
    @show values′

    return Fun{D,P,R,C,S,T}(mfd, values′)
end

################################################################################

function basis_products(::Val{Pr}, ::Val{R},
                        mfd::Manifold{D,C,S}) where {D,R,C,S}
    # Check arguments
    D::Int
    R::Int
    C::Int
    @assert 0 ≤ R ≤ D ≤ C
    @assert C == D
    N = D + 1

    # `order=2` suffices because we only have linear basis functions
    order = 2
    scheme = integration_scheme(S, Val(D), Val(order))

    # Result: sparse matrix
    n = nsimplices(mfd, R)
    B = MakeSparse{S}(n, n)

    # Loop over all D-simplices
    simplicesD = get_simplices(mfd, D)
    for i in axes(simplicesD, 2)
        si = sparse_column_rows(simplicesD, i)
        @assert length(si) == N
        si = SVector{N}(si[n] for n in 1:N)

        # Coordinates of simplex vertices
        xs = SVector{N,SVector{C,S}}(get_coords(mfd)[n] for n in si)
        x2λ = cartesian2barycentric_setup(xs)
        dλ2dx = dbarycentric2dcartesian_setup(Form{D,R}, xs)

        # Double loop over all contained R-simplices
        lookup_RD = get_lookup(mfd, R, D)
        iter = enumerate(sparse_column_rows(lookup_RD, i))
        for (nj, j) in iter
            for (nk, k) in iter
                function kernel(x::SVector{C,S})
                    bj = basis_x(Form{D,R}, x2λ, dλ2dx, nj, x)
                    bk = basis_x(Form{D,R}, x2λ, dλ2dx, nk, x)
                    return (bj ⋅ bk)[]
                end
                b = integrate(kernel, scheme, xs)
                # Note: This inserts contributions for the same basis
                # function multiple times, each time for a different
                # simplex. The sparse matrix constructor will then sum
                # these contributions.
                B[Int(j), Int(k)] = b
            end
        end
    end

    return SparseOp{R,R}(sparse(B))
end

################################################################################

# @fastmath function basis_x(setup, n::Int, x::SVector{C,T}) where {C,T}
#     λ = cartesian2barycentric(setup, x)
#     return basis_λ(n, λ)
# end
#
# # TODO: Remove this, use Bernstein polynomials instead
# @fastmath function basis_λ(n::Int, λ::SVector{N,S}) where {N,S}
#     @assert 1 ≤ n ≤ N
#     return λ[n]
# end

# See: Douglas Arnold, Richard Falk, Ragnar Winther, "Finite element
# exterior calculus, homological techniques, and applications", Acta
# Numerica 15, 1-155 (2006), DOI:10.1017/S0962492906210018, section
# 3.3
#
# P⁻(r) Λ(k): reduced space of polynomial `k`-forms with order up to `r`
# P⁻(r) Λ(k) = P(r-1) Λ(k) ⊕ κ H(r-1) Λ(k+1)

# D     R     [n]
# 2     0     1   x   y
# 2     1     dx   dy   x dy - y dx
# 2     2     dx∧dy

# D     R     [n]
# 3     0     1   x   y   z
# 3     1     dx   dy   dz   x dy - y dx   x dz - z dx   y dz - z dy
# 3     2     dx∧dy   dx∧dz   dy∧dz   x dy dz - y dx dz + z dx dy
# 3     3     dx∧dy∧dz

# Expressed ∈ terms of barycentric coordinates:

# x = xⁱ λᵢ = x¹ λ₁ + x² λ₂ + x³ λ₃
# y = yⁱ λᵢ = y¹ λ₁ + y² λ₂ + y³ λ₃
# 1 = Σᵢ λᵢ = λ₁ + λ₂ + λ₃
# dx = xⁱ dλᵢ = x¹ dλ₁ + x² dλ₂ + x³ dλ₃
# dy = yⁱ dλᵢ = y¹ dλ₁ + y² dλ₂ + y³ dλ₃
# 0  = Σᵢ dλᵢ = dλ₁ + dλ₂ + dλ₃

# D     R     [n]
#
# 1     0     λ₁   λ₂
# 1     1     (λ₁ + λ₂) (- dλ₁ + dλ₂)
#
# 2     0     λ₁   λ₂   λ₃
# 2     1     (λ₁ + λ₂) (- dλ₁ + dλ₂)   (λ₁ + λ₃) (- dλ₁ + dλ₃)
#             (λ₂ + λ₃) (- dλ₂ + dλ₃)
# 2     2     (λ₁ + λ₂ + λ₃) (dλ₁₂ - dλ₁₃ + dλ₂₃)

# To determine signs:
# 1. look at vertices that make up the face
# 2. commute the derivative index to the left; that determines the parity

#
# 3     0     λ₁   λ₂   λ₃   λ₄
# 3     1     (λ₁ + λ₂) (- dλ₁ + dλ₂)   (λ₁ + λ₃) (- dλ₁ + dλ₃)
#             (λ₁ + λ₄) (- dλ₁ + dλ₄)   (λ₂ + λ₃) (- dλ₂ + dλ₃)
#             (λ₂ + λ₄) (- dλ₂ + dλ₄)   (λ₃ + λ₄) (- dλ₃ + dλ₄)
# 3     2     (λ₁ + λ₂ + λ₃) (dλ₁₂ - dλ₁₃ + dλ₂₃)
#             (λ₁ + λ₂ + λ₄) (dλ₁₂ - dλ₁₄ + dλ₂₄)
#             (λ₁ + λ₃ + λ₄) (dλ₁₃ - dλ₁₄ + dλ₃₄)
#             (λ₂ + λ₃ + λ₄) (dλ₂₃ - dλ₂₄ + dλ₃₄)
# 3     3     (λ₁ + λ₂ + λ₃ + λ₄) (- dλ₁₂₃ + dλ₁₂₄ - dλ₁₃₄ + dλ₂₃₄)

# xᵢ = xⁿᵢ λₙ
# dxᵢ = xⁿᵢ dλₙ

# λₙ = Xⁱₙ xᵢ + bₙ
# dλₙ = Xⁱₙ dxᵢ

# dλ₁ + dλ₂                     = dλ₁ - dλ₁                     = 0
# dλ₁₂ + dλ₁₃ + dλ₂₃            = dλ₁₂ - dλ₁₂ + dλ₁₂            = dλ₁₂
# dλ₁₂₃ + dλ₁₂₄ + dλ₁₃₄ + dλ₂₃₄ = dλ₁₂₃ - dλ₁₂₃ + dλ₁₂₃ + dλ₁₂₃ = 2 dλ₁₂₃

# dx1 = x1i dλi
# dx1 ∧ dx2 = x1i dλi ∧ x2j dλj = x1i x2j dλi ∧ dλj
# dx1 ∧ dx2 ∧ dx3 = x1i x2j x3k dλi ∧ dλj ∧ dλk

# x[i] = xs[n][i] λ[n]
# dx[i] = xs[n][i] dλ[n]
#
# dx[i] ∧ dx[j] = xs[n][i] dλ[n] ∧ xs[m][j] dλ[m]
#               = xs[n][i] xs[m][j] dλ[n] ∧ dλ[m]

################################################################################

# WE THINK THESE ARE CORRECT, BUT ARE IN THE WRONG ORDER

# julia> tpc = trimmed_polynomial_complex(Val(2), Int, 1);
#
# julia> whitney(Basis{3,0,Int})
# Basis{3,0,Int64}[
#     ⟦(1 * [0, 0, 1])⟧,
#     ⟦(1 * [0, 1, 0])⟧,
#     ⟦(1 * [1, 0, 0])⟧]
#
# julia> whitney(Basis{3,1,Int})
# Basis{3,1,Int64}[
#     ⟦(), (-1 * [0, 0, 1]), (1 * [0, 1, 0])⟧,
#     ⟦(-1 * [0, 0, 1]), (), (1 * [1, 0, 0])⟧,
#     ⟦(-1 * [0, 1, 0]), (1 * [1, 0, 0]), ()⟧]
#
# julia> whitney(Basis{3,2,Int})
# Basis{3,2,Int64}[
#     ⟦(1 * [0, 0, 1]), (-1 * [0, 1, 0]), (1 * [1, 0, 0])⟧]

# b(D=2, R=0, n=1) = λ1
# b(D=2, R=0, n=2) = λ2
# b(D=2, R=0, n=3) = λ3

# b(D=2, R=1, n=1) = - λ3 dλ2 + λ2 dλ3
# b(D=2, R=1, n=2) = - λ3 dλ1 + λ1 dλ3
# b(D=2, R=1, n=3) = - λ2 dλ1 + λ1 dλ2

# b(D=2, R=2, n=1) = λ3 dλ12 - λ2 dλ13 + λ1 dλ2λ3

################################################################################

# WE THINK THESE ARE CORRECT

# Whitney forms:
#     ϕᵢ = Σₖ (-1)ᵏ⁺¹ λᵢₖ dλᵢ₁ ∧ [dλᵢₖ] ∧ dλᵢₙ

# D = 0:

#     λ₁ = 1
#
#     dλ₁ = 0

#     D=0, R=0:
#         ϕ₁ = λ₁

#     D=0, R=0:
#         ϕ₁ = 1

# D = 1:

#     λ₁ = 1 - x
#     λ₂ = x
#
#     dλ₁ = - dx
#     dλ₂ = dx

#     D=1, R=0:
#         ϕ₁ = λ₁
#         ϕ₂ = λ₂
#     D=1, R=1:
#         ϕ₁₂ = λ₁ dλ₂ - λ₂ dλ₁

#     D=1, R=0:
#         ϕ₁ = 1 - x
#         ϕ₂ = x
#     D=1, R=1:
#         ϕ₁₂ = (1 - x) dx - x (- dx)
#             = dx

# D = 2:

#     λ₁ = 1 - x - y
#     λ₂ = x
#     λ₃ = y
#
#     dλ₁ = - dx - dy
#     dλ₂ = dx
#     dλ₃ = dy

#     D=2, R=0:
#         ϕ₁ = λ₁
#         ϕ₂ = λ₂
#         ϕ₃ = λ₃
#     D=2, R=1:
#         ϕ₁₂ = λ₁ dλ₂ - λ₂ dλ₁
#         ϕ₁₃ = λ₁ dλ₃ - λ₃ dλ₁
#         ϕ₂₃ = λ₂ dλ₃ - λ₃ dλ₂
#     D=2, R=2:
#         ϕ₁₂₃ = λ₁ dλ₂ ∧ dλ₃ - λ₂ dλ₁ ∧ dλ₃ + λ₃ dλ₁ ∧ dλ₂

#     D=2, R=0:
#         ϕ₁ = 1 - x - y
#         ϕ₂ = x
#         ϕ₃ = y
#     D=2, R=1:
#         ϕ₁₂ = (1 - x - y) dx - x (- dx - dy)
#             = - (y - 1) dx + x dy
#         ϕ₁₃ = (1 - x - y) dy - y (- dx - dy)
#             = y dx - (x - 1) dy
#         ϕ₂₃ = x dy - y dx
#             = - y dx + x dy
#     D=2, R=2:
#         ϕ₁₂₃ = (1 - x - y) dx ∧ dy - x (- dx - dy) ∧ dy + y (- dx - dy) ∧ dx
#              = dx ∧ dy

################################################################################

function basis_x(::Type{<:Form{D,R}}, x2λ, dλ2dx, n::Int,
                 x::SVector{C,T}) where {D,R,C,T}
    λ = cartesian2barycentric(x2λ, x)
    bλ = basis_λ(Form{D,R}, n, λ)::Form{D + 1,R,T}
    bx = dbarycentric2dcartesian(dλ2dx, bλ)::Form{D,R,T}
    return bx
end

# @fastmath
function dbarycentric2dcartesian(dλ2dx::SMatrix{Nx,Nλ},
                                 dλ::Form{N,R}) where {Nx,Nλ,N,R}
    N::Int
    @assert N ≥ 1
    D = N - 1
    R::Int
    @assert 0 ≤ R ≤ D
    @assert Nx == length(Form{D,R})
    @assert Nλ == length(Form{D + 1,R})
    return Form{D,R}(dλ2dx * dλ.elts)
end

#TODO @fastmath
function dbarycentric2dcartesian_setup(::Type{<:Form{D,R}},
                                       xs::SVector{N,SVector{C,T}}) where {D,R,
                                                                           N,C,
                                                                           T}
    # Form dimension and rank
    D::Int
    R::Int
    @assert 0 ≤ R ≤ D
    N::Int
    @assert N == D + 1
    C::Int
    @assert C == D
    # Number of barycentric form elements
    Nλ = length(Form{D + 1,R})
    # Number of physical form elements
    Nx = length(Form{D,R})

    # Σₙ λⁿ = 1
    # Σₙ dλⁿ = 0

    # (x,1)ⁱ = (x,1)ₙⁱ λⁿ
    # (dx,0)ⁱ = (x,1)ₙⁱ dλⁿ

    # λⁿ = (x,1)ⁿᵢ (x,1)ⁱ
    # dλⁿ = (x,1)ⁿᵢ (dx,0)ⁱ

    setup = cartesian2barycentric_setup(xs)
    dλ2dx = setup.invA::SMatrix{N,N}

    r = zero(SMatrix{Nx,Nλ,T})
    for nx in 1:Nx, nλ in 1:Nλ
        lstx = DifferentialForms.Forms.lin2lst(Val(D), Val(R), nx)
        lstλ = DifferentialForms.Forms.lin2lst(Val(D + 1), Val(R), nλ)
        @assert length(lstx) == length(lstλ)

        relt = zero(T)
        lstλs = permutations(lstλ)
        for lstλ′ in lstλs
            rterm = one(T)
            _, p = sort_perm(lstλ′)
            rterm *= bitsign(p)
            for (lx, lλ) in zip(lstx, lstλ′)
                rterm *= dλ2dx[lλ, lx]
            end
            relt += rterm
        end

        r = setindex(r, relt, nx, nλ)
    end

    return r
end

#TODO @fastmath
function basis_λ_old(::Type{<:Form{D,R}}, i::Int,
                     λ::SVector{N,T}) where {D,R,C,N,T}
    D::Int
    R::Int
    @assert 0 ≤ R ≤ D
    # Number of barycentric coordinates
    @assert N == D + 1
    # Number of barycentric form elements
    Nλ = length(Form{D + 1,R})
    # Number of basis functions
    I = DifferentialForms.binomial(Val(D + 1), Val(R + 1))
    @assert 1 ≤ i ≤ I

    # Find functional dependency on λᵢ
    inds = DifferentialForms.Forms.lin2lst(Val(D + 1), Val(R + 1), i)
    rval = sum(λ[inds])

    # Find contributing components and their parity
    bits = DifferentialForms.Forms.lin2bit(Val(D + 1), Val(R + 1), i)
    relts = zero(SVector{Nλ,T})
    # TODO: Can this be implemented without this loop?
    for ind in inds
        @assert bits[ind]
        bits2 = Base.setindex(bits, false, ind)
        p = count(bits[1:(ind - 1)])
        s = bitsign(p)
        j = DifferentialForms.Forms.bit2lin(Val(D + 1), Val(R), bits2)
        relts = setindex(relts, relts[j] + s * rval, j)
    end

    return Form{D + 1,R}(relts) / factorial(R)
end

#TODO @fastmath
function basis_λ(::Type{<:Form{D,R}}, i::Int, λ::SVector{N,T}) where {D,R,C,N,T}
    D::Int
    R::Int
    @assert 0 ≤ R ≤ D
    # Number of barycentric coordinates
    @assert N == D + 1
    # Number of barycentric form elements
    Nλ = length(Form{D + 1,R})
    # Number of basis functions
    I = DifferentialForms.binomial(Val(D + 1), Val(R + 1))
    @assert 1 ≤ i ≤ I

    ϕ = whitney(Basis{D + 1,R,Int})::Basis{D + 1,R,Int}
    @assert length(ϕ.forms) == I
    ϕi = ϕ.forms[i]::Form{D + 1,R,Poly{D + 1,Int}}

    belts = map(poly -> evalpoly(poly, λ), ϕi.elts)::SVector{Nλ,T}
    return Form{D + 1,R}(belts)
end

function evalpoly(poly::Poly{D}, x::SVector{D,T}) where {D,T}
    r = zero(T)
    for term in poly.terms
        r += term.coeff * prod(x .^ term.powers)
    end
    return r::T
end

@generated function integration_scheme(::Type{T}, ::Val{D},
                                       ::Val{order}) where {T,D,order}
    return grundmann_moeller(T, Val(D), order + iseven(order))
end

end


text:
module TestInteractionATE

include("helper_fns.jl")

using Test
using TMLE
using MLJ
using Distributions
using Random
using StableRNGs
using Tables
using StatsBase

mutable struct InteractionTransformer <: Static end

function MLJ.transform(a::InteractionTransformer, _, X)
    Xmatrix = MLJ.matrix(X)
    nrows, ncols = size(Xmatrix)
    ninter = Int(ncols*(ncols-1)/2)
    Xinteracts = Matrix{Float64}(undef, nrows, ninter)
    i = 0
    for col₁ in 1:(ncols-1)
        for col₂ in (col₁+1):ncols
            i += 1
            Xinteracts[:, i] = Xmatrix[:, col₁] .* Xmatrix[:, col₂]
        end
    end
    return MLJ.table(hcat(Xmatrix, Xinteracts))
end

LogisticClassifier = @load LogisticClassifier pkg=MLJLinearModels verbosity=0
LinearRegressor = @load LinearRegressor pkg=MLJLinearModels verbosity = 0

cont_interacter = InteractionTransformer |> LinearRegressor
cat_interacter = InteractionTransformer |> LogisticClassifier
Ns = [100, 1000, 10000, 100000]

function binary_target_binary_treatment_pb(rng;n=100)
    μy_fn(W, T₁, T₂) = TMLE.expit(2W[:, 1] .+ 1W[:, 2] .- 2W[:, 3] .- T₁ .+ T₂ .+ 2*T₁ .* T₂)
    # Sampling W: Bernoulli
    W = rand(rng, Bernoulli(0.5), n, 3)

    # Sampling T₁, T₂ from W: Softmax
    θ = rand(rng, 3, 4)
    softmax = exp.(W*θ) ./ sum(exp.(W*θ), dims=2)
    T = [sample(rng, [1, 2, 3, 4], Weights(softmax[i, :])) for i in 1:n]
    T₁ = [t in (1,2) ? true : false for t in T]
    T₂ = [t in (1,3) ? true : false for t in T]

    # Sampling y from T₁, T₂, W: Logistic
    μy = μy_fn(W, T₁, T₂)
    y = [rand(rng, Bernoulli(μy[i])) for i in 1:n]

    # Respect the Tables.jl interface and convert types
    W = MLJ.table(float(W))
    T = (T₁ = categorical(T₁), T₂ = categorical(T₂))
    y = categorical(y)

    # Compute the theoretical ATE
    Wcomb = [1 1 1;
            1 1 0;
            1 0 0;
            1 0 1;
            0 1 0;
            0 0 0;
            0 0 1;
            0 1 1]
    ATE = 0
    for i in 1:8
        w = reshape(Wcomb[i, :], 1, 3)
        temp = μy_fn(w, [1], [1])[1]
        temp += μy_fn(w, [0], [0])[1]
        temp -= μy_fn(w, [1], [0])[1]
        temp -= μy_fn(w, [0], [1])[1]
        ATE += temp*0.5*0.5*0.5
    end
    return T, W, y, ATE
end

function binary_target_categorical_treatment_pb(rng;n=100)
    function μy_fn(W, T, Hmach)
        Thot = transform(Hmach, T)
        TMLE.expit(2W[:, 1] .+ 1W[:, 2] .- 2W[:, 3]
                    .- Thot[1] .+ Thot[2] .+ 2Thot[3] .- 3Thot[4]
                    .+ 2*Thot[1].*Thot[2]
                    .+ 1*Thot[1].*Thot[3]
                    .- 4*Thot[1].*Thot[4]
                    .- 3*Thot[2].*Thot[3]
                    .+ 1.5*Thot[2].*Thot[4]
                    .- 2.5*Thot[3].*Thot[4]
                    )
    end
    # Sampling W:
    W = rand(rng, Bernoulli(0.5), n, 3)

    # Sampling T from W:
    # T₁, T₂ will have 3 categories each
    # This is embodied by a 9 dimensional full joint
    θ = rand(rng, 3, 9)
    softmax = exp.(W*θ) ./ sum(exp.(W*θ), dims=2)
    encoding = collect(Iterators.product(["CC", "GG", "CG"], ["TT", "AA", "AT"]))
    T = [sample(rng, encoding, Weights(softmax[i, :])) for i in 1:n]
    T = (T₁=categorical([t[1] for t in T]), T₂=categorical([t[2] for t in T]))

    Hmach = machine(OneHotEncoder(drop_last=true), T)
    fit!(Hmach, verbosity=0)

    # Sampling y from T, W:
    μy = μy_fn(W, T, Hmach)
    y = [rand(rng, Bernoulli(μy[i])) for i in 1:n]

    # Compute the theoretical ATE for the query
    # (CC, AT) against (CG, AA)
    Wcomb = [1 1 1;
            1 1 0;
            1 0 0;
            1 0 1;
            0 1 0;
            0 0 0;
            0 0 1;
            0 1 1]
    ATE = 0
    levels₁ = levels(T.T₁)
    levels₂ = levels(T.T₂)
    for i in 1:8
        w = reshape(Wcomb[i, :], 1, 3)
        temp = μy_fn(w, (T₁=categorical(["CC"], levels=levels₁), T₂=categorical(["AT"], levels=levels₂)), Hmach)[1]
        temp += μy_fn(w, (T₁=categorical(["CG"], levels=levels₁), T₂=categorical(["AA"], levels=levels₂)), Hmach)[1]
        temp -= μy_fn(w, (T₁=categorical(["CC"], levels=levels₁), T₂=categorical(["AA"], levels=levels₂)), Hmach)[1]
        temp -= μy_fn(w, (T₁=categorical(["CG"], levels=levels₁), T₂=categorical(["AT"], levels=levels₂)), Hmach)[1]
        ATE += temp*0.5*0.5*0.5
    end
    return T, MLJ.table(float(W)), categorical(y), ATE
end

function continuous_target_binary_treatment_pb(rng;n=100)
    μy_fn(W, T₁, T₂) = 2W[:, 1] .+ 1W[:, 2] .- 2W[:, 3] .- T₁ .+ T₂ .+ 2*T₁ .* T₂
    # Sampling W: Bernoulli
    W = rand(rng, Bernoulli(0.5), n, 3)

    # Sampling T₁, T₂ from W: Softmax
    θ = rand(rng, 3, 4)
    softmax = exp.(W*θ) ./ sum(exp.(W*θ), dims=2)
    T = [sample(rng, [1, 2, 3, 4], Weights(softmax[i, :])) for i in 1:n]
    T₁ = [t in (1,2) ? true : false for t in T]
    T₂ = [t in (1,3) ? true : false for t in T]

    # Sampling y from T₁, T₂, W: Logistic
    μy = μy_fn(W, T₁, T₂)
    y = μy + rand(rng, Normal(0, 0.1), n)

    # Respect the Tables.jl interface and convert types
    W = MLJ.table(float(W))
    T = (T₁ = categorical(T₁), T₂ = categorical(T₂))

    # Compute the theoretical ATE
    Wcomb = [1 1 1;
            1 1 0;
            1 0 0;
            1 0 1;
            0 1 0;
            0 0 0;
            0 0 1;
            0 1 1]
    ATE = 0
    for i in 1:8
        w = reshape(Wcomb[i, :], 1, 3)
        temp = μy_fn(w, [1], [1])[1]
        temp += μy_fn(w, [0], [0])[1]
        temp -= μy_fn(w, [1], [0])[1]
        temp -= μy_fn(w, [0], [1])[1]
        ATE += temp*0.5*0.5*0.5
    end
    return T, W, y, ATE
end

@testset "Test Double Robustness IATE on binary_target_binary_treatment_pb" begin
    # When Q̅ is misspecified but G is well specified
    query = (T₁=[true, false], T₂=[true, false])
    Q̅ = ConstantClassifier()
    G = FullCategoricalJoint(LogisticClassifier())
    tmle = TMLEstimator(Q̅, G, query)

    abs_mean_rel_errors, abs_vars = asymptotics(
            tmle,
            binary_target_binary_treatment_pb,
            StableRNG(123),
            Ns
            )
    @test all(abs_mean_rel_errors .< [69, 14, 5, 1.5])
    @test all(abs_vars .< [0.05, 0.002, 0.0003, 3.5e-5])

    # When Q̅ is well specified  but G is misspecified
    query = (T₁=[true, false], T₂=[true, false])
    Q̅ = cat_interacter
    G = FullCategoricalJoint(ConstantClassifier())
    tmle = TMLEstimator(Q̅, G, query)

    abs_mean_rel_errors, abs_vars = asymptotics(
            tmle,
            binary_target_binary_treatment_pb,
            StableRNG(123),
            Ns
            )
    @test all(abs_mean_rel_errors .< [53, 13, 4.5, 1.4])
    @test all(abs_vars .< [0.03, 0.002, 0.0003, 2.7e-5])

end

@testset "Test Double Robustness IATE on continuous_target_binary_treatment_pb" begin
    # When Q̅ is misspecified but G is well specified

    # If the order of the variables in the query do no match
    # the order in T it shouldn't matter
    query = (T₂=[true, false], T₁=[true, false])
    Q̅ = MLJ.DeterministicConstantRegressor()
    G = FullCategoricalJoint(LogisticClassifier())
    tmle = TMLEstimator(Q̅, G, query)

    abs_mean_rel_errors, abs_vars = asymptotics(
            tmle,
            continuous_target_binary_treatment_pb,
            StableRNG(123),
            Ns
            )
    @test all(abs_mean_rel_errors .< [18, 1.6, 0.5, 0.2])
    @test all(abs_vars .< [0.08, 0.002, 0.0002, 1.5e-5])

    # When Q̅ is well specified  but G is misspecified
    query = (T₁=[true, false], T₂=[true, false])
    Q̅ = cont_interacter
    G = FullCategoricalJoint(ConstantClassifier())
    tmle = TMLEstimator(Q̅, G, query)

    abs_mean_rel_errors, abs_vars = asymptotics(
            tmle,
            continuous_target_binary_treatment_pb,
            StableRNG(123),
            Ns
            )
    @test all(abs_mean_rel_errors .< [2.3, 0.6, 0.3, 0.06])
    @test all(abs_vars .< [0.003, 0.0003, 2.5e-5, 2e-6])

end

@testset "Test Double Robustness IATE on binary_target_categorical_treatment_pb" begin
    # When Q̅ is misspecified but G is well specified
    query = (T₁=["CC", "CG"], T₂=["AT", "AA"])
    Q̅ = ConstantClassifier()
    G = FullCategoricalJoint(LogisticClassifier())
    tmle = TMLEstimator(Q̅, G, query)

    abs_mean_rel_errors, abs_vars = asymptotics(
            tmle,
            binary_target_categorical_treatment_pb,
            StableRNG(123),
            Ns
            )
    @test all(abs_mean_rel_errors .< [31, 8, 2, 0.6])
    @test all(abs_vars .< [0.03, 0.006, 0.0004, 2.7e-5])

    # When Q̅ is well specified but G is misspecified
    query = (T₁=["CC", "CG"], T₂=["AT", "AA"])
    Q̅ = cat_interacter
    G = FullCategoricalJoint(ConstantClassifier())
    tmle = TMLEstimator(Q̅, G, query)

    abs_mean_rel_errors, abs_vars = asymptotics(
            tmle,
            binary_target_categorical_treatment_pb,
            StableRNG(123),
            Ns
            )
    @test all(abs_mean_rel_errors .< [22, 7, 2.1, 0.7])
    @test all(abs_vars .< [0.04, 0.005, 0.0004, 3.5e-5])
end

end;

true

ext: jl

max_stars_repo_name: olivierlabayle/TMLE

max_issues_repo_name: olivierlabayle/TMLE

max_forks_repo_name: olivierlabayle/TMLE


text:
function sort_matches!(matches::Vector{Int}, indptr::Vector{Int})
    for i in 1:length(indptr)-1
        sort!(view(matches, indptr[i]:indptr[i+1]-1))
    end
end

@testset "Testing deferred_acceptance.jl" begin

    matchings_one_to_one = Array{Dict, 1}()
    matchings_many_to_one = Array{Dict, 1}()
    matchings_many_to_many = Array{Dict, 1}()

    # Test case: one-to-one
    m, n = 4, 3
    # Males' preference orders over females [1, 2, 3] and unmatched 0
    m_prefs = [1 3 2 3;
               2 1 3 1;
               3 2 1 2;
               0 0 0 0]
    # Females' preference orders over males [1, 2, 3, 4] and unmatched 0
    f_prefs = [3 1 3;
               1 2 0;
               2 3 2;
               4 4 1;
               0 0 4]
    # Unique stable matching
    m_matches_m_opt = m_matches_f_opt = [1, 2, 3, 0]
    f_matches_m_opt = f_matches_f_opt = [1, 2, 3]
    d = Dict(
        "m_prefs" => m_prefs,
        "f_prefs" => f_prefs,
        "m_matches_m_opt" => m_matches_m_opt,
        "f_matches_m_opt" => f_matches_m_opt,
        "m_matches_f_opt" => m_matches_f_opt,
        "f_matches_f_opt" => f_matches_f_opt,
    )
    push!(matchings_one_to_one, d)

    # Test case: one-to-one; from Roth and Sotomayor Example 2.9
    m, n = 5, 4
    m_prefs = [
        1, 2, 3, 4, 0,
        4, 2, 3, 1, 0,
        4, 3, 1, 2, 0,
        1, 4, 3, 2, 0,
        1, 2, 4, 0, 3,
    ]
    m_prefs = reshape(m_prefs, n+1, m)
    f_prefs = [
        2, 3, 1, 4, 5, 0,
        3, 1, 2, 4, 5, 0,
        5, 4, 1, 2, 3, 0,
        1, 4, 5, 2, 3, 0,
    ]
    f_prefs = reshape(f_prefs, m+1, n)
    m_matches_m_opt, f_matches_m_opt = [1, 2, 3, 4, 0], [1, 2, 3, 4]
    m_matches_f_opt, f_matches_f_opt = [4, 1, 2, 3, 0], [2, 3, 4, 1]
    d = Dict(
        "m_prefs" => m_prefs,
        "f_prefs" => f_prefs,
        "m_matches_m_opt" => m_matches_m_opt,
        "f_matches_m_opt" => f_matches_m_opt,
        "m_matches_f_opt" => m_matches_f_opt,
        "f_matches_f_opt" => f_matches_f_opt,
    )
    push!(matchings_one_to_one, d)

    # Test case: one-to-one; from Roth and Sotomayor Example 2.17
    m = n = 4
    m_prefs = [
        1, 2, 3, 4, 0,
        2, 1, 4, 3, 0,
        3, 4, 1, 2, 0,
        4, 3, 2, 1, 0,
    ]
    m_prefs = reshape(m_prefs, n+1, m)
    f_prefs = [
        4, 3, 2, 1, 0,
        3, 4, 1, 2, 0,
        2, 1, 4, 3, 0,
        1, 2, 3, 4, 0,
    ]
    f_prefs = reshape(f_prefs, m+1, n)
    m_matches_m_opt = f_matches_m_opt = [1, 2, 3, 4]
    m_matches_f_opt = f_matches_f_opt = [4, 3, 2, 1]
    d = Dict(
        "m_prefs" => m_prefs,
        "f_prefs" => f_prefs,
        "m_matches_m_opt" => m_matches_m_opt,
        "f_matches_m_opt" => f_matches_m_opt,
        "m_matches_f_opt" => m_matches_f_opt,
        "f_matches_f_opt" => f_matches_f_opt,
    )
    push!(matchings_one_to_one, d)

    # Test case: many-to-many; from Roth and Sotomayor Example 5.24
    m_caps = fill(2, m)
    f_caps = fill(2, n)
    m_matches_m_opt = f_matches_m_opt = [1, 2, 1, 2, 3, 4, 3, 4]
    m_matches_f_opt = f_matches_f_opt = [3, 4, 3, 4, 1, 2, 1, 2]
    m_indptr = f_indptr = [1, 3, 5, 7, 9]
    d = Dict(
        "s_prefs" => m_prefs,
        "c_prefs" => f_prefs,
        "s_caps" => m_caps,
        "c_caps" => f_caps,
        "s_matches_s_opt" => m_matches_m_opt,
        "c_matches_s_opt" => f_matches_m_opt,
        "s_matches_c_opt" => m_matches_f_opt,
        "c_matches_c_opt" => f_matches_f_opt,
        "s_indptr" => m_indptr,
        "c_indptr" => f_indptr,
    )
    push!(matchings_many_to_many, d)

    # Test case: many-to-one
    # From Gusfield and Irving (1989, Section 1.6.5)

    m = 11  # Number of students
    n = 5   # Number of colleges

    # Students' preference orders over colleges 1, ..., 5 and unmatched 0
    s_prefs = [3, 1, 5, 4, 0, 2,
               1, 3, 4, 2, 5, 0,
               4, 5, 3, 1, 2, 0,
               3, 4, 1, 5, 0, 2,
               1, 4, 2, 0, 3, 5,
               4, 3, 2, 1, 5, 0,
               2, 5, 1, 3, 0, 4,
               1, 3, 2, 5, 4, 0,
               4, 1, 5, 0, 2, 3,
               3, 1, 5, 2, 4, 0,
               5, 4, 1, 3, 2, 0]
    s_prefs = reshape(s_prefs, n+1, m)

    # Colleges' preference orders over students 1, ..., 11 and unmatched 0
    c_prefs = [3, 7, 9, 11, 5, 4, 10, 8, 6, 1,
               2, 0, 5, 7, 10, 6, 8, 2, 3, 11,
               0, 1, 4, 9, 11, 6, 8, 3, 2, 4,
               7, 1, 10, 0, 5, 9, 10, 1, 2, 11,
               4, 9, 5, 3, 6, 8, 0, 7, 2, 4,
               10, 7, 6, 1, 8, 3, 11, 9, 0, 5]
    c_prefs = reshape(c_prefs, m+1, n)

    # Capacities for colleges
    caps = [4, 1, 3, 2, 1]
    indptr = [1, 5, 6, 9, 11, 12]

    # Optimal stable matchings

    # Student optimal
    s_matches_s_opt = [3, 1, 4, 3, 1, 3, 2, 1, 4, 1, 5]
    c_matches_s_opt = [2, 5, 8, 10, 7, 1, 4, 6, 3, 9, 11]

    # Matching the book claims to be hospital-(or college-)optimal
    # s_matches_c_opt = [4, 4, 3, 1, 1, 3, 2, 3, 1, 5, 1]
    # c_matches_c_opt = [4, 5, 9, 11, 7, 3, 6, 8, 1, 2, 10]

    # Matching actually hospital-(or college-)optimal
    s_matches_c_opt = [4, 5, 3, 1, 1, 3, 2, 3, 1, 4, 1]
    c_matches_c_opt = [4, 5, 9, 11, 7, 3, 6, 8, 1, 10, 2]

    d = Dict(
        "s_prefs" => s_prefs,
        "c_prefs" => c_prefs,
        "caps" => caps,
        "s_matches_s_opt" => s_matches_s_opt,
        "c_matches_s_opt" => c_matches_s_opt,
        "s_matches_c_opt" => s_matches_c_opt,
        "c_matches_c_opt" => c_matches_c_opt,
        "indptr" => indptr,
    )
    push!(matchings_many_to_one, d)

    # Test case: many-to-one; from Roth and Sotomayor Page 16
    m, n = 7, 5
    s_prefs = [
        5, 1, 0, 2, 3, 4,
        2, 5, 1, 0, 3, 4,
        3, 1, 0, 2, 4, 5,
        4, 1, 0, 2, 3, 5,
        1, 2, 0, 3, 4, 5,
        1, 3, 0, 2, 4, 5,
        1, 3, 4, 0, 2, 6,
    ]
    s_prefs = reshape(s_prefs, n+1, m)
    c_prefs = [
        1, 2, 3, 4, 5, 6, 7, 0,
        5, 2, 0, 1, 3, 4, 6, 7,
        6, 7, 3, 0, 1, 2, 4, 5,
        7, 4, 0, 1, 2, 3, 5, 6,
        2, 1, 0, 3, 4, 5, 6, 7,
    ]
    c_prefs = reshape(c_prefs, m+1, n)
    caps = [3, 1, 1, 1, 1]
    indptr = [1, 4, 5, 6, 7,  8]
    s_matches_s_opt = [5, 2, 3, 4, 1, 1, 1]
    c_matches_s_opt = [5, 6, 7, 2, 3, 4, 1]
    s_matches_c_opt = [1, 5, 1, 1, 2, 3, 4]
    c_matches_c_opt = [1, 3, 4, 5, 6, 7, 2]
    d = Dict(
        "s_prefs" => s_prefs,
        "c_prefs" => c_prefs,
        "caps" => caps,
        "s_matches_s_opt" => s_matches_s_opt,
        "c_matches_s_opt" => c_matches_s_opt,
        "s_matches_c_opt" => s_matches_c_opt,
        "c_matches_c_opt" => c_matches_c_opt,
        "indptr" => indptr,
    )
    push!(matchings_many_to_one, d)

    @testset "deferred_acceptance: one-to-one" begin
        for d in matchings_one_to_one
            # Male proposal
            m_matches, f_matches =
                @inferred deferred_acceptance(d["m_prefs"], d["f_prefs"])
            @test m_matches == d["m_matches_m_opt"]
            @test f_matches == d["f_matches_m_opt"]

            # Female proposal
            f_matches, m_matches =
                @inferred deferred_acceptance(d["f_prefs"], d["m_prefs"])
            @test m_matches == d["m_matches_f_opt"]
            @test f_matches == d["f_matches_f_opt"]
        end
    end

    @testset "deferred_acceptance: many-to-one" begin
        for d in matchings_many_to_one
            # Default (Student proposal)
            s_matches, c_matches, indptr =
                @inferred deferred_acceptance(d["s_prefs"], d["c_prefs"], d["caps"])
            sort_matches!(c_matches, indptr)
            @test s_matches == d["s_matches_s_opt"]
            @test c_matches == d["c_matches_s_opt"]
            @test indptr == d["indptr"]

            # Student proposal
            s_matches, c_matches, indptr =
                @inferred deferred_acceptance(d["s_prefs"], d["c_prefs"], d["caps"],
                                    SProposing)
            sort_matches!(c_matches, indptr)
            @test s_matches == d["s_matches_s_opt"]
            @test c_matches == d["c_matches_s_opt"]
            @test indptr == d["indptr"]

            # College proposal
            s_matches, c_matches, indptr =
                @inferred deferred_acceptance(d["s_prefs"], d["c_prefs"], d["caps"],
                                    CProposing)
            sort_matches!(c_matches, indptr)
            @test s_matches == d["s_matches_c_opt"]
            @test c_matches == d["c_matches_c_opt"]
            @test indptr == d["indptr"]
        end
    end

    @testset "deferred_acceptance: many-to-many" begin
        for d in matchings_many_to_many
            # Student proposal
            s_matches, c_matches, s_indptr, c_indptr =
                @inferred deferred_acceptance(d["s_prefs"], d["c_prefs"],
                                    d["s_caps"], d["c_caps"])
            sort_matches!(s_matches, s_indptr)
            sort_matches!(c_matches, c_indptr)
            @test s_matches == d["s_matches_s_opt"]
            @test c_matches == d["c_matches_s_opt"]
            @test s_indptr == d["s_indptr"]
            @test c_indptr == d["c_indptr"]

            # Course proposal
            c_matches, s_matches, c_indptr, s_indptr =
                @inferred deferred_acceptance(d["c_prefs"], d["s_prefs"],
                                    d["c_caps"], d["s_caps"])
            sort_matches!(s_matches, s_indptr)
            sort_matches!(c_matches, c_indptr)
            @test s_matches == d["s_matches_c_opt"]
            @test c_matches == d["c_matches_c_opt"]
            @test s_indptr == d["s_indptr"]
            @test c_indptr == d["c_indptr"]
        end
    end

    # Invalid inputs

    @testset "prefs arrays inconsistent shapes" begin
        prop_prefs = [1 1;
                      2 2]
        resp_prefs = [1 1;
                      2 2]
        @test_throws ArgumentError deferred_acceptance(prop_prefs, resp_prefs)
    end

    @testset "caps inconsistent length" begin
        prop_prefs = [1 1;
                      2 2;
                      0 0]
        resp_prefs = [1 1;
                      2 2;
                      0 0]
        caps = [2, 1, 2]
        @test_throws ArgumentError deferred_acceptance(
            prop_prefs, resp_prefs, caps
        )
    end

end


text:
using CUDA
import Zygote
using Base: issingletontype
using .NeuralNetwork

const use_cuda = Ref{Union{Nothing,Bool}}(nothing)

function check_use_cuda()
  if use_cuda[] === nothing
    use_cuda[] = CUDA.functional()
    if use_cuda[] && !CUDA.has_cudnn()
      @warn "CUDA.jl found cuda, but did not find libcudnn. Some functionality will not be available."
    end
    if !(use_cuda[])
      @info """The GPU function is being called but the GPU is not accessible.
               Defaulting back to the CPU. (No action is required if you want to run on the CPU).""" maxlog=1
    end
  end
end
Zygote.@nograd check_use_cuda

@generated new_struct(DT, args...) = Expr(:call, nameof(DT), :(args...))
new_struct(t::T, args...) where {T<:Tuple} = ntuple(i->args[i], length(args))

isleaf(t::T) where {T}  = issingletontype(T), !isstructtype(T), false
isleaf(t::AbstractArray{T}) where {T} = false, isprimitivetype(T), true

adapt(x) = CUDA.cu(x)
adapt(x::Zygote.FillArrays.AbstractFill) = CUDA.cu(collect(x))
adapt(x::Zygote.OneElement) = CUDA.cu(collect(x))

cpu(x::AbstractArray) = Array(x)

for processor in (:gpu, :cpu)
    @eval begin
        $(Symbol("array_$(processor)_obj"))(objs::AbstractArray{T}) where {T} = collect([$(processor)(objs[i]) for i in 1 : length(objs)])
    end

    @eval begin
        function $(processor)(obj::T) where {T}
            is_singleton, is_leaf, is_array = isleaf(obj)
            is_leaf && return $((processor==:gpu) ? :adapt : :cpu)(obj)
            if is_array
                return $(Symbol("array_$(processor)_obj"))(obj)
            elseif is_singleton
                return obj
            else
                fields = fieldnames(T)
                return new_struct(obj, [$(processor)(getfield(obj, f)) for f in fields]...)
            end
        end
    end
end

"""
    gpu(model)
Transform the model so that it can be trained on the GPU. When called in an environment without a GPU, it does nothing and returns the original model.
!!! note
    This function is included in the HorseML module and can only be used with `using HorseML`.

# Example
```jldoctest gpu
julia> model = NetWork(Dense(10=>5, relu), Dense(5=>1, tanh)) |> gpu
Layer1 : Dense(IO:10 => 5, σ:relu)
Layer2 : Dense(IO:5 => 1, σ:tanh)

julia> model[1].w |> typeof
CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}
```
"""
function gpu(model::NetWork)
    check_use_cuda()
    if !use_cuda[]
        @warn "your computer doesn't have a GPU, or couldn't recognized. So a GPU isn't used."
        return model
    end
    N = NetWork()
    n = length(model.net)
    for i in 1 : n
        add_layer!(N, gpu(model[i]))
    end
    return N
end

"""
    cpu(model)
Put the model trained on the GPU back on the CPU.
!!! note
    This function is included in the HorseML module and can only be used with `using HorseML`.

# Example
```jldoctest
julia> model_on_cpu = model |> cpu #This model is made with description of gpu function
Layer1 : Dense(IO:10 => 5, σ:relu)
Layer2 : Dense(IO:5 => 1, σ:tanh)

julia> model_on_cpu[1].w |> typeof
Matrix{Float32} (alias for Array{Float32, 2})
```
"""
function cpu(model::NetWork)
    N = NetWork()
    n = length(model.net)
    for i in 1 : n
        add_layer!(N, cpu(model[i]))
    end
    return N
end


text:
abstract type FSubType end
struct FSubVar <: FSubType
	name::Symbol
end
struct FSubTop <: FSubType end
struct FSubUni <: FSubType
	binding::Symbol
	ub::FSubType
	body::FSubType
end

subst(var::Symbol, fort::FSubType, int::FSubTop) = int
subst(var::Symbol, fort::FSubType, int::FSubVar) = if int.name==var fort else int end
function subst(var::Symbol, fort::FSubType, int::FSubUni)
	if int.binding == var
		return int
	end
	return FSubUni(int.binding, subst(var, fort, int.ub), subst(var, fort, int.body))
end

fsub(a::FSubType, b::FSubTop, env::Dict{Symbol, FSubType})=true
fsub(a::FSubVar, b::FSubTop, env::Dict{Symbol, FSubType})=true
function fsub(a::FSubVar, b::FSubType, env::Dict{Symbol, FSubType})
	if b isa FSubVar && a.name == b.name
		return true
	end
	return fsub(env[a.name], b, env)
end
function fsub(a::FSubUni, b::FSubUni, env::Dict{Symbol, FSubType})
	if !fsub(b.ub, a.ub, env)
		return false
	end
	nsvar = gensym(a.binding)
	nvar = FSubVar(nsvar)
	na = subst(a.binding, nvar, a.body)
	nb = subst(b.binding, nvar, b.body)
	env[nsvar] = b.ub
	return fsub(na, nb, env)
end
function fsub(a::FSubType, b::FSubType, env::Dict{Symbol, FSubType})
	println("Failed at $env |- $a <: $b")
	return false
end

function enc(v::FSubVar, eenv::Dict{Symbol,TypeVar})
	return eenv[v.name] # from Ref{...} ; check that not invariant
end
function enc(v::FSubTop, eenv::Dict{Symbol, TypeVar})
	return Union{}
end
function enc(v::FSubUni, eenv::Dict{Symbol, TypeVar})
	nvn = TypeVar(gensym(v.binding))
	lb = enc(v.ub, eenv)
	nvn.lb = lb
	eenv[v.binding] = nvn
	return UnionAll(nvn, Tuple{Ref{nvn}, enc(v.body, eenv)})
end
function enc(v::FSubType)
	return enc(v, Dict{Symbol, TypeVar}())
end

function esub(a::FSubType,b::FSubType)
	return enc(b) <: enc(a)
end
function enc(env::Dict{Symbol,FSubType})
	return Dict(k => begin nvn = TypeVar(gensym(k)); nvn.lb = enc(v); nvn end for (k,v) in env)
end

function esub(a::FSubType, b::FSubType, env::Dict{Symbol,FSubType})
	tenv = enc(env)
	A = enc(b, tenv)
	B = enc(a, tenv)
	lhs = foldl((t,v) -> UnionAll(v,t), values(tenv); init=Tuple{Ref{A}, Ref{V} where V>:B})
	rhs = (Tuple{Ref{T}, Ref{V} where V>:T} where T)
	return lhs <: rhs
end

fsnot(a::FSubType) = begin isym=gensym(:notv); FSubUni(isym, a, FSubVar(isym)) end
theta = FSubUni(:a, FSubTop(), fsnot(FSubUni(:b, FSubVar(:a), fsnot(FSubVar(:b)))))
ge = Dict{Symbol, FSubType}(:a0=>theta)
gl = FSubVar(:a0)
gr = FSubUni(:a1, FSubVar(:a0), fsnot(FSubVar(:a1)))

esub(gl, gr, ge)

max_stars_repo_name: prl-julia/julia-sub

max_issues_repo_name: prl-julia/julia-sub

max_forks_repo_name: prl-julia/julia-sub


text:
module TokenizerTest

using MiniLoggers
using MiniLoggers: Color, Token, colorparse, iscolored, tokenize
using ReTest

comptokens(t1, t2) = all(x == y for (x, y) in zip(t1, t2))

@testset "iscolored" begin
    @test iscolored(Color(:green, false))
    @test iscolored(Color(:green, true))
    @test iscolored(Color(-1, true))
    @test !iscolored(Color(-1, false))
end

@testset "colorparse correct" begin
    i, col = colorparse("xxx:red")
    @test i == 3
    @test col == Color(:red, false)

    i, col = colorparse("xxx:red:bold")
    @test i == 3
    @test col == Color(:red, true)

    i, col = colorparse("xxx:bold")
    @test i == 3
    @test col == Color(:normal, true)

    i, col = colorparse("xxx:5")
    @test i == 3
    @test col == Color(5, false)

    i, col = colorparse("xxx:5:bold")
    @test i == 3
    @test col == Color(5, true)
end

@testset "colorparse edge" begin
    i, col = colorparse(":red")
    @test i == 0
    @test col == Color(:red, false)

    i, col = colorparse(":red:bold")
    @test i == 0
    @test col == Color(:red, true)

    i, col = colorparse(":green:red:bold")
    @test i == 6
    @test col == Color(:red, true)

    i, col = colorparse(":3000:bold")
    @test i == 5
    @test col == Color(:normal, true)

    i, col = colorparse(":bold:3000")
    @test i == 10
    @test col == Color(-1, false)
end

@testset "basic token" begin
    @test Token("world:green:bold") == Token("world", Color(:green, true))
    @test Token(":green:bold") == Token("", Color(:green, true))
    @test Token("hello") == Token("hello", Color(-1, false))
end

@testset "basic tokenizer" begin
    res1 = tokenize("hello:red")
    res2 = [Token("hello", Color(:red, false))]
    @test comptokens(res1, res2)

    res1 = tokenize("{hello}")
    res2 = [Token("hello", Color(-1, false))]
    @test comptokens(res1, res2)

    res1 = tokenize("{hello} world:red")
    res2 = [Token("hello", Color(:red, false)), Token(" world", Color(:red, false))]
    @test comptokens(res1, res2)

    res1 = tokenize("and {hello:green:bold} world:red")
    res2 = [Token("and ", Color(:red, false)),
            Token("hello", Color(:green, true)),
            Token(" world", Color(:red, false))]
    @test comptokens(res1, res2)

    res1 = tokenize("{and {hello:green:bold} world:red}")
    res2 = [Token("and ", Color(:red, false)),
            Token("hello", Color(:green, true)),
            Token(" world", Color(:red, false))]
    @test comptokens(res1, res2)

    res1 = tokenize("and {hello:green:bold} world:")
    res2 = [Token("and ", Color(-1, false)),
            Token("hello", Color(:green, true)),
            Token(" world:", Color(-1, false))]
    @test comptokens(res1, res2)

    res1 = tokenize("{hello}{world}")
    res2 = [Token("hello"), Token("world")]
    @test comptokens(res1, res2)
end

end # module


text:
LC = LocalCongruence

"""
	cellCongruenceSM(
		cop::Lar.ChainOp,
		lo_cls::Array{Array{Int,1},1},
		lo_sign::Array{Array{Int8,1},1};
		imp = false,
		d = 0
	)::Tuple{ Lar.ChainOp,  Array{Array{Int,1},1},  Array{Array{Int8,1},1} }

Evaluates the Cell Congruence for a Cochain `cop with classes `lo_cls.

The function determines the new Cochain Operator built from `cop` where
the lower order cells are merged according to `lo_cls` map.
`lo_sign` specifies whether a cell must be considered in reverse order.

If optional paramenter `imp` is set to `true` then FP imprecisions
are taken into account in the sense that lower order cells may have collided.
The parameter `d` represent then the order of the cell (that also is the
least number of lower order cells a current order cell is made of).

The method returns:
 - the new Cochain Operator
 - a map that, for every new cell, identifies the old cells it is made of
 - a map that, for every new cell, specify if old cells have changed ordering.
"""
function cellCongruenceSM(cop, lo_cls, lo_sign; imp = false, d = 0)

	# Collide columns
	copCols = []
	for i = 1 : length(lo_cls)
		col = sum([
			cop[:, lo_cls[i][j]] .* lo_sign[i][j]
			for j = 1 : length(lo_cls[i])
		])
#		if imp
#			#TODO remove zeros
#			if length(col.nzind) > d  push!(copCols, col)  end
#		end
		push!(copCols, col)
	end

	# Retrieve Matrix Form and extract rows
	cop = hcat(copCols...)
	rows = [cop[row, :] for row = 1 : cop.m]

	# Adjustinfg signs such that the first value always is `-1`
	sign = ones(Int8, cop.m)
	for row = 1 : cop.m  if rows[row].nzval[1] > 0
		rows[row] = -rows[row]
		sign[row] = -1
	end  end

	# Sort rows in order to collide them quickly
	rows_ord = sortperm(rows)
	nrows = unique(rows[rows_ord])
	ho_cls = [Array{Int,1}() for i = 1 : length(nrows)]
	nidx = 1

	# Collide cells with same structure
	for cidx = 1 : cop.m
		if rows[rows_ord[cidx]] != nrows[nidx]  nidx = nidx + 1;  end
		push!(ho_cls[nidx], rows_ord[cidx])
	end

	# Reshaping Sign and Cochain
	ho_sign = [[sign[el] for el in cell] for cell in ho_cls]
#	cop = convert(Lar.ChainOp, hcat(nrows...)')

	return hcat(nrows...)', ho_cls, ho_sign
end

"""
	chainCongruenceSM(
		G::Lar.Points, T::Array{Lar.ChainOp};
		imp=false, ϵ=1e-6
	)::Tuple{Lar.Points, Array{Lar.ChainOp}}

Perform the Geometry `G` congruence and coherently reshape the Topology `T`
"""
function chainCongruenceSM(G, T; imp=false, ϵ=1e-6)

	# Perform the Geometry Congruence
	G, cls = LC.vertCongruence(G; ϵ=ϵ)
	# Build default sign
	sign = [ones(Int8, length(cl)) for cl in cls]
	# Update the Topology coherently
	Tn = Array{Lar.ChainOp, 1}(undef, length(T))
	for d = 1 : length(T)
		Tn[d], cls, sign = LC.cellCongruenceSM(T[d], cls, sign; imp=imp, d=d)
	end
	# Euler Characteristic
	# @show size(G,2) - size(T[1],1) + size(T[2],1)
	return G, Tn
end


text:
function nx_max(S::GphysData)
  N = 0
  window_lengths = Array{Int64,1}(undef,0)
  t = Array{Int64,2}(undef,0,0)
  for i = 1:S.n
    if S.fs[i] > 0.0
      t = getfield(S, :t)[i]
      window_lengths = diff(t[:,1])
      window_lengths[end] += 1
      N = max(N, maximum(window_lengths))
    end
  end
  return N
end

function nx_max(S::GphysData, C::Union{Integer, UnitRange, Array{Int64,1}})
  N = 0
  window_lengths = Array{Int64,1}(undef,0)
  t = Array{Int64,2}(undef,0,0)
  for i = 1:S.n
    if i in C
      if S.fs[i] > 0.0
        t = getfield(S, :t)[i]
        window_lengths = diff(t[:,1])
        window_lengths[end] += 1
        N = max(N, maximum(window_lengths))
      end
    end
  end
  return N
end

function nx_max(C::GphysChannel)
  @assert C.fs > 0.0
  window_lengths = diff(C.t[:,1])
  window_lengths[end] += 1
  return maximum(window_lengths)
end


text:
module QUIP

using ASE
using JuLIP
using PyCall

quippy = pyimport("quippy")

export soap

function soap(at, cutoff; n_max=nothing, l_max=nothing, sigma=0.5, normalise=true)

all_unique_nums = unique(collect(at.Z))

#Setting SOAP Descriptors
desc_definition = "soap" *
                  " cutoff=" * string(cutoff) *
                  " l_max=" * string(l_max) *
                  " n_max=" * string(n_max) *
                  " atom_sigma=" * string(sigma) *
                  " normalise=" * (normalise ? 'T' : 'F') *
                  " n_Z=" * string(length(all_unique_nums)) *
                  " Z={" * join([string(nm) for nm in all_unique_nums], ", ") * "}" *
                  " n_species=" * string(length(all_unique_nums)) *
                  " species_Z={" * join([string(nm) for nm in all_unique_nums], " ") * "}"

soap_desc = quippy.descriptors.Descriptor(desc_definition)
atom_struct = ASEAtoms(at)
# Calculate descriptor
soap_rtn_desc = soap_desc.calc(atom_struct.po)
#has_data = soap_rtn_desc["has_data"]
return soap_rtn_desc["data"][:,1:end-1]
end

end # module


text:
using DSGE
using Base.Test

# Test `hessizero` in context of Rosenbrock function
function rosenbrock(x::Vector)
    a = 100
    b = 1
    return a*(x[2]-x[1]^2.0)^2.0 + b*(1-x[1])^2.0
end

function rosenbrock_hessian(x::Vector)
    H = zeros(eltype(x), 2, 2)
    H[1,1] = 2.0 - 400.0 * x[2] + 1200.0 * x[1]^2.0
    H[1,2] = -400.0 * x[1]
    H[2,1] = -400.0 * x[1]
    H[2,2] = 200.0
    return H
end

# At the min, ensure no negatives in diagonal
x0 = [1.0, 1.0]
hessian_expected = rosenbrock_hessian(x0)
hessian, _ = DSGE.hessizero(rosenbrock, x0; check_neg_diag=true)
@testset "Check valid hessian at the minimum" begin
    @test_matrix_approx_eq hessian_expected hessian
end

# Not at the min (indeed, we are at max), ensure throws error
x1 = [1.0, 1.0]
rosenbrock_neg(x) = -rosenbrock(x)
@testset "Throw error for hessian calculation not at the min" begin
    @test_throws Exception hessian, _ = DSGE.hessizero(rosenbrock_neg, x1; check_neg_diag=true)
end

# Not at the min, check matches closed form
x1 = Vector[[0.5, 1.5], [-1.0, -1.0]]
@testset "Check closed-form of the hessian" begin
    for x in x1
        hessian_expected = rosenbrock_hessian(x)
        hessian, _ = DSGE.hessizero(rosenbrock, x)
        @test @test_matrix_approx_eq hessian_expected hessian
    end
end


text:
using NahaJuliaLib
using Logging
using VectorLogging

@testset "test trace" begin
    logger = VectorLogger()

    @trace trace_f1 f1(a, b::Int, c) = a + b + c
    trace_f1 = true
    with_logger(logger) do
        f1(1, 2, 3)
    end
    @test length(logger.log) == 2
    @test logger.log[1].message == "Trace enter"
    @test logger.log[1].keys isa AbstractDict
    # @test logger.log[1].keys[:fcall] == "(f1)(1, 2, 3)"    # FAILS
    @test logger.log[2].message == "Trace exit"
    @test logger.log[2].keys[:result] == 6
    @test logger.log[1].keys[:id] == 1
    @test logger.log[1].keys[:id] == logger.log[2].keys[:id]

    @trace trace_f2 f2(a; b=0, c=0) = a + b + c
    trace_f2 = true
    with_logger(logger) do
        f2(1; b=2, c=3)
    end
    @test length(logger.log) == 4
    @test logger.log[3].message == "Trace enter"
    # @test logger.log[1].keys[:fcall] == "(f2)(1; b = 2, c = 3)"    # FAILS
    @test logger.log[4].message == "Trace exit"
    @test logger.log[4].keys[:result] == 6
    @test length(logger.log) == 4
    @test logger.log[3].keys[:id] == 2
    @test logger.log[3].keys[:id] == logger.log[4].keys[:id]
end

trace_hanoi = false

@trace(trace_hanoi,
       function hanoi(from, to, other, count)
           if count == 0
               return nothing
           else
               hanoi(from, other, to, count - 1)
               println("move 1 from $from to $to")
               hanoi(other, to, from, count - 1)
               return (from, to)   # arbitrary result to show
           end
       end
       )

const SHOWN_HANOI_TRACE = """(hanoi)(a, b, c, 4) => (:a, :b)
  (hanoi)(a, c, b, 3) => (:a, :c)
    (hanoi)(a, b, c, 2) => (:a, :b)
      (hanoi)(a, c, b, 1) => (:a, :c)
        (hanoi)(a, b, c, 0) => nothing
        (hanoi)(b, c, a, 0) => nothing
      (hanoi)(c, b, a, 1) => (:c, :b)
        (hanoi)(c, a, b, 0) => nothing
        (hanoi)(a, b, c, 0) => nothing
    (hanoi)(b, c, a, 2) => (:b, :c)
      (hanoi)(b, a, c, 1) => (:b, :a)
        (hanoi)(b, c, a, 0) => nothing
        (hanoi)(c, a, b, 0) => nothing
      (hanoi)(a, c, b, 1) => (:a, :c)
        (hanoi)(a, b, c, 0) => nothing
        (hanoi)(b, c, a, 0) => nothing
  (hanoi)(c, b, a, 3) => (:c, :b)
    (hanoi)(c, a, b, 2) => (:c, :a)
      (hanoi)(c, b, a, 1) => (:c, :b)
        (hanoi)(c, a, b, 0) => nothing
        (hanoi)(a, b, c, 0) => nothing
      (hanoi)(b, a, c, 1) => (:b, :a)
        (hanoi)(b, c, a, 0) => nothing
        (hanoi)(c, a, b, 0) => nothing
    (hanoi)(a, b, c, 2) => (:a, :b)
      (hanoi)(a, c, b, 1) => (:a, :c)
        (hanoi)(a, b, c, 0) => nothing
        (hanoi)(b, c, a, 0) => nothing
      (hanoi)(c, b, a, 1) => (:c, :b)
        (hanoi)(c, a, b, 0) => nothing
        (hanoi)(a, b, c, 0) => nothing
"""

@testset "test trace analasys" begin
    logger = VectorLogger()
    old = trace_hanoi
    try
        global trace_hanoi = true
        with_logger(logger) do
            hanoi(:a, :b, :c, 4)
        end
        traces = analyze_traces(logger)
        @test length(traces) == 1
        out = IOBuffer()
        show_trace(out, traces[1])
        @test String(take!(out)) == SHOWN_HANOI_TRACE
    finally
        global trace_hanoi = old
    end
end


text:
using aoc2019.computer: load_program, io, run

input = joinpath(@__DIR__, "input")
p = load_program(input)

stdin, stdout = io()
put!(stdin, 1)
run(copy(p), stdin, stdout)
close(stdout)
signals = [signal for signal in stdout]
p1 = last(signals)
@assert p1 == 6745903

stdin, stdout = io()
put!(stdin, 5)
run(copy(p), stdin, stdout)
close(stdout)
signals = [signal for signal in stdout]
p2 = last(signals)
@assert p2 == 9168267

print("-----------------------------------------------------------------------\n")
print("sunny with a chance of asteroids -- part one\n    final value : $p1\n")
print("sunny with a chance of asteroids -- part two\n    final value : $p2\n")
print("-----------------------------------------------------------------------\n")


text:
# Chemfiles.jl, a modern library for chemistry file reading and writing
# Copyright (C) Guillaume Fraux and contributors -- BSD license

module Chemfiles
    module lib
        const ROOT = dirname(@__FILE__)
        const depsfile = normpath(joinpath(ROOT, "..", "deps", "deps.jl"))
        if isfile(depsfile)
            include(depsfile)
        else
            error("Chemfiles is not installed (the '$depsfile' file is missing).\nPlease run Pkg.build(\"Chemfiles\")")
        end
        include("generated/types.jl")
        include("generated/cdef.jl")
    end

    include("misc.jl")
    include("utils.jl")

    export Trajectory, Topology, Atom, UnitCell, Frame, Selection, Residue

    function version()
        unsafe_string(lib.chfl_version())
    end

    """
    A ``Property`` is a generic container for various forms of metadata
    stored for other structures.
    """
    struct Property
        __handle :: CxxPointer{lib.CHFL_PROPERTY}
    end

    """
    An ``Atom`` is a particle in the current ``Frame``.

    An atom stores the following atomic properties:
      - atom name
      - atom type
      - atom mass
      - atom charge

    The atom name is usually an unique identifier ("H1", "C_a") while the atom
    type will be shared among all particles of the same type: "H", "Ow", "CH3".
    """
    struct Atom
        __handle :: CxxPointer{lib.CHFL_ATOM}
    end

    """
    A ``Residue`` is a group of atoms belonging to the same logical unit. They
    can be small molecules, amino-acids in a protein, monomers in polymers,
    *etc.*
    """
    struct Residue
        __handle :: CxxPointer{lib.CHFL_RESIDUE}
    end

    """
    A ``Topology`` describes the organisation of the particles in the system:
    what their names are, how they are bonded together, *etc.* A ``Topology``
    is a list of ``Atoms`` in the system, together with the list of bonds
    between the atoms.
    """
    struct Topology
        __handle :: CxxPointer{lib.CHFL_TOPOLOGY}
    end

    """
    A ``UnitCell`` describes the bounding box of a system. It is represented by
    three base vectors of lengths ``a``, ``b`` and ``c``; and the angles
    between these vectors are ``alpha``, ``beta`` and ``gamma``.
    """
    struct UnitCell
        __handle :: CxxPointer{lib.CHFL_CELL}
    end

    """
    A ``Frame`` holds data for one step of a simulation. As not all formats
    provide all the types of information, some fields may be initialized to a
    default value. A ``Frame`` may contain the following data:

    - Positions for all the atoms in the system;
    - Velocities for all the atoms in the system;
    - The ``Topology`` of the system;
    - The ``UnitCell`` of the system.
    """
    struct Frame
        __handle :: CxxPointer{lib.CHFL_FRAME}
    end

    """
    A ``Selection`` is used to select a group of atoms. Examples of selections are
    "name H" and "(x < 45 and name O) or name C". See the `full documentation
    <http://chemfiles.org/chemfiles/latest/selections.html>`_ for more
    information about the selection language.
    """
    struct Selection
        __handle :: CxxPointer{lib.CHFL_SELECTION}
    end

    """
    A ``Trajectory`` represents a simulation file on the hard drive. It can read
    or write one or many ``Frame``s to this file. The file format can be
    automatically determined from the extention, or manually specified.
    Writing to a ``Trajectory`` is buffered, which means that one needs to
    ``close()`` the trajectory and flush the buffer before being able to read the
    file again.
    """
    struct Trajectory
        __handle :: CxxPointer{lib.CHFL_TRAJECTORY}
    end

    include("Property.jl")
    include("Atom.jl")
    include("Residue.jl")
    include("Topology.jl")
    include("UnitCell.jl")
    include("Frame.jl")
    include("Selection.jl")
    include("Trajectory.jl")

    function __init__()
        if !startswith(version(), "0.9")
            error(
                """Chemfiles.jl requires the 0.9 version of libchemfiles,
                but $(version()) is installed. You can try to run
                Pkg.build(\"Chemfiles\") to update libchemfiles."""
            )
        end
        set_warning_callback(__default_warning_callback)
    end
end


text:
module BioArgParse

using Bio: Seq
import ArgParse.parse_item

"Parse a string passed in on the command line into a biological sequence."
function parse_item{A<:Alphabet}(::Type{BioSequence{A}}, x::AbstractString)
    return BioSequence{A}(x)
end

"Parse a string passed in on the command line into a nucleotide."
function parse_item{N<:Nucleotide}(::Type{N}, x::AbstractString)
    return convert(N, x[1])
end

"Parse a string passed in on the command line into a vector of nucleotides."
function parse_item{N<:Nucleotide}(::Type{Vector{N}}, x::AbstractString)
    return [convert(N, char) for char in x]
end

end


text:
function bin(df::AbstractDataFrame, @nospecialize(f::FormulaTerm), n::Integer = 20;
            weights::Union{Symbol, Nothing} = nothing)
    df = partial_out(df, _shift(f); weights = weights, align = false, add_mean = true)[1]
    cols = names(df)
    df.__cut = cut(df[!, end], n)
    df = groupby(df, :__cut, sort = true)
    out = combine(df, cols .=> mean .=> cols; keepkeys = false)
    return out
end

function bin(df::GroupedDataFrame, @nospecialize(f::FormulaTerm), n::Integer = 20;
            weights::Union{Symbol, Nothing} = nothing)
    combine(d -> bin(d, f, n; weights = weights), df; ungroup = false)
end

function bin(df, @nospecialize(f::FormulaTerm), n::Integer = 20;
            weights::Union{Symbol, Nothing} = nothing)
    bin(DataFrame(df), f, n; weights = weights)
end

# simplified version of CategoricalArrays' cut
function cut(x::AbstractArray, ngroups::Integer)
    xnm = eltype(x) >: Missing ? skipmissing(x) : x
    cut(x, Statistics.quantile(xnm, (1:ngroups-1)/ngroups))
end

function cut(x::AbstractArray, breaks::AbstractVector)
    xnm = eltype(x) >: Missing ? skipmissing(x) : x
    min_x, max_x = extrema(xnm)
    if first(breaks) > min_x
        breaks = [min_x; breaks]
    end
    if last(breaks) < max_x
        breaks = [breaks; max_x]
    end
    refs = Array{UInt32}(undef, size(x))
    fill_refs!(refs, x, breaks)
end

function fill_refs!(refs::AbstractArray, X::AbstractArray, breaks::AbstractVector)
    upper = last(breaks)
    @inbounds for i in eachindex(X)
        x = X[i]
        if ismissing(x)
            refs[i] = 0
        elseif x == upper
            refs[i] = length(breaks)-1
        else
            refs[i] = searchsortedlast(breaks, x)
        end
    end
    return refs
end

#transform lhs ~ x + rhs to lhs + x ~ rhs
function _shift(@nospecialize(formula::FormulaTerm))
    lhs = formula.lhs
    rhs = formula.rhs
    if !(lhs isa Tuple)
        lhs = (lhs,)
    end
    if !(rhs isa Tuple)
        rhs = (rhs,)
    end
    i = findfirst(x -> x isa Term, rhs)
    FormulaTerm(tuple(lhs..., rhs[i]), Tuple(term for term in rhs if term != rhs[i]))
end


text:
import SpecialFunctions
include("../examples/FPDE_problem.jl")

Γ = SpecialFunctions.gamma
α = 1.6
F(z) = z^(-α)
N = 14

t = range(0.1, 5.0, length=201)
y = Float64[ Bromich_integral(tval, F, N) for tval in t ]

err = maximum(abs.( y - t.^(α-1) / Γ(α) ))


text:
# NEEDED MODULES -- This can take time at the first launch
# Code main be needed to run twice to plot correctly (shift-enter)
using DifferentialEquations
using Plots
using LaTeXStrings
# Load the backend (engine) for plotting based on javascript
plotlyjs()

function weighted_leastsquare(O,R,E)
    denom  = sum(O.^2)
    weights = 1. / (E.^2)
    return sqrt(sum(weights.*(O-R).^2)/denom)
end

function leastsquare(O,R)
    denom  = sum(O.^2)
    return sqrt(sum((O-R).^2)/denom)
end

# Saturation function to positive values
# (concentrations are positive function)
# This is just used as a safety for the parameter search phase
@inline function sat(x)
    if x < 0.0
        return 0.0
    else
        x
    end
end

# Gaussian model of plasma Glucose appearance rate
function GlcGauss(t)
    σ = 35.0 # width
    μ = 30.0 # peak time
    dose = 4040 # total glucose dose (area under curve)
    return dose*1/(σ*sqrt(2*pi))*exp((-(t-μ)^2/(2*σ^2)))
end

# Parametric ODEs of te Glucose response model
# t : time
# u : variables
# p : parameters
# Here I call the guassian model of glucose input GlcGauss(t)
# And I saturate the variables to the positive cone
function f_param!(du,u,p,t)
 du[1] = -p[1]*(sat(u[1])-p[2])- p[10]*sat(u[2])*sat(u[1])+GlcGauss(t)
 du[2] = -p[3]*sat(u[2])+p[4]*(sat(u[3])-p[5])
 du[3] = -p[6]*sat(u[3])+p[7]*(sat(u[1])-p[8])*t + p[9]*u[4]
 du[4] = -p[9]*u[4]
end

# Initial conditions at 26j
u0 = [76.0,0.0,34.0,5950]
tspan = (0.0,120.0)

# #  A Parameter set from the MINMOD model as a reference
# #       p1  Gb    p2    p3     Ib  n   γ    h     |added parameters => |     p4    r
# p0 = [0.01,100.0,0.01,0.00001,8.0,0.1,0.01,100.0,                           0.07,1.0]

# Fitted parameter set to the control group at 26 Days ===> err: 0.00315
#         p1    Gb    p2     p3     Ib     n      γ      h     p4     r
 p0 = [0.01, 100.0, 0.61, 0.0245, 05.0, 09.44, 0.0110, 79.0, 0.0215,  1.0]
###################################################
## Simulation of the Control experiment at 26Days
prob_param = ODEProblem(f_param!,u0,tspan,p0)
sctrl = solve(prob_param, alg_hints = [:stiff])
## Experimental results
texp = [0,10,20,30,45,60,90,120]
glcexp = [75.4,133.4,166.5,184.7,186.5,164.8,126.4,111.7]
uppererror = [3.2,6.7,5.9,5.6,4.1,3.6,3.3,3.4]
lowererror = [3.2,6.7,5.9,5.6,4.1,3.6,3.3,3.4]
errs = [lowererror;uppererror]
res_tot = [sctrl(ti) for ti in texp]
res = [ri[1] for ri in res_tot]
## Computation of the weighted least square error and least square
wls_err_Ctrl = weighted_leastsquare(glcexp,res,uppererror)
print("CONTROL wls = ", wls_err_Ctrl, "\n")
ls_err_Ctrl = leastsquare(glcexp,res)
## Generation of the figure and first plots
plot(sctrl,vars=[(0,1)],label="Glucose (control) simulations")
scatter!(texp,glcexp,yerr = errs,marker = stroke(2, :blue),
       xlabel="time (min)", ylabel= "Concentration: mg/dL",
    label="Control exp results", legend=:top, markerstrokecolor = :auto)

# Parameters sets for the different hypothesis for the group CD1 at 26 Days
########################################################################
# Control parameter set  ===>   err: 0.00972
#        p1    Gb     p2     p3    Ib     n       γ      h     p4      r
# p0 = [0.01, 100.0, 0.61, 0.0245, 05.0, 09.44, 0.0110, 79.0, 0.0215,  1.0]
# hypothese 1.1: sensibility to interstitial insuline (r) ===> err : 0.00411
# p0 = [0.01, 100.0, 0.61, 0.0245, 05.0, 09.44, 0.0110, 79.0, 0.0215,  0.91]
# hypothese 1.2: disponibility of insuline for X(t) (p3) ==> err: 0.004105
# p0 = [0.01, 100.0, 0.61, 0.0224, 05.0, 09.44, 0.0110, 79.0, 0.0215,  1.0]
# hypothese 2: degradation  insuline (n) ==> err : 0.00395
 p0 = [0.01, 100.0, 0.61, 0.0245, 05.0, 9.93, 0.0110, 79.0, 0.0215,  1.0]
# hypothese 3: interstitial insuline degradation (p2) ==> err : 0.00398
# p0 = [0.01, 100.0, 0.67, 0.0245, 05.0, 09.44, 0.0110, 79.0, 0.0215,  1.0]
# hypothesis 4: γ (Insuline production acceleration) decreased  ==> err =  0.00762
# p0 = [0.01, 100.0, 0.61, 0.0245, 05.0, 09.44, 0.0100, 79.0, 0.0215,  1.0]
# hypothesis 5:  change p4 (stored insuline release rate)  ==> err = 0.00638
# p0 = [0.01, 100.0, 0.61, 0.0245, 05.0, 09.44, 0.0110, 79.0, 0.0200,  1.0]
# hypothese 6: reponse Glucose threshold (h)  ==> err:  0.00819
# p0 = [0.01, 100.0, 0.61, 0.0245, 05.0, 09.44, 0.0110, 85.0, 0.0215,  1.0]
########################################################################
# Simulation CD1 26 days
prob_param = ODEProblem(f_param!,u0,tspan,p0)
scd1 = solve(prob_param, alg_hints = [:stiff])
# Experimental results
glcexpCd1 = [75.7,133.3,184.5,196.8,188.6,167.4,131.2,110.9]
uppererror = [1.97,5.34,4.96,5.21,2.52,4.47,4.30,2.64]
lowererror = [1.97,5.34,4.96,5.21,2.52,4.47,4.30,2.64]
errs = [lowererror;uppererror]
res_tot = [scd1(ti) for ti in texp]
res = [ri[1] for ri in res_tot]
# Error computation
wls_err_Cd1 = weighted_leastsquare(glcexpCd1,res,uppererror)
print("CD1 wls = ", wls_err_Cd1, "\n")
ls_err_Cd1 = leastsquare(glcexpCd1,res)
# Plotting
plot!(scd1,vars=[(0,1)],label="Glucose(Cd1) simulations")
scatter!(texp,glcexpCd1,yerr = errs,marker = stroke(2, :green),
      xlabel="time (min)", ylabel= "Concentration: mg/dL",
      label="Cd1 exp results", legend=:top, markerstrokecolor = :auto)

# Parameters sets for the different hypothesis for the group CD2 at 26 Days
########################################################################
# Control parameter set  ===>  err: 0.00600
#        p1    Gb     p2     p3    Ib     n       γ      h     p4      r
# p0 = [0.01, 100.0, 0.61, 0.0245, 05.0, 09.44, 0.0110, 79.0, 0.0215,  1.0]
# hypothese 1.1: sensibility to interstitial insuline (r) ==> err: 0.00444
# p0 = [0.01, 100.0, 0.61, 0.0245, 05.0, 09.44, 0.0110, 79.0, 0.0215,  0.93]
# hypothese 1.2: disponibility of insuline for X(t) (p3) ==> err : 0.00435
# p0 = [0.01, 100.0, 0.61, 0.0230, 05.0, 09.44, 0.0110, 79.0, 0.0215,  1.0]
# hypothese 2: degradation  insuline (n) ==> err : 0.00442
# p0 = [0.01, 100.0, 0.61, 0.0245, 05.0, 09.75, 0.0110, 79.0, 0.0215,  1.0]
# hypothese 3: interstitial insuline degradation (p2) ==> err: 0.00434
# p0 = [0.01, 100.0, 0.65, 0.0245, 05.0, 09.44, 0.0110, 79.0, 0.0215,  1.0]
# hypothesis 4: γ (Insuline production acceleration) decreased  ==> err =  0.00584
# p0 = [0.01, 100.0, 0.61, 0.0245, 05.0, 09.44, 0.0107, 79.0, 0.0215,  1.0]
# hypothesis 5:  change p4 (stored insuline release rate)  ==> err = 0.00388
 p0 = [0.01, 100.0, 0.61, 0.0245, 05.0, 09.44, 0.0110, 79.0, 0.0201,  1.0]
# hypothese 6: reponse Glucose threshold (h)  ==> err: 0.0058
# p0 = [0.01, 100.0, 0.61, 0.0245, 05.0, 09.44, 0.0110, 81.0, 0.0215,  1.0]
##########################################################################
# Simulation Group CD2 26 days
prob_param = ODEProblem(f_param!,u0,tspan,p0)
scd2 = solve(prob_param, alg_hints = [:stiff])
# Experimental results
glcexpCd2 = [77.9,132,182.8,195,183,162.3,125.5,112.2]
uppererror = [1.6,7.1,7.7,7.3,5.7,3.8,3.9,4.2]
lowererror = [1.6,7.1,7.7,7.3,5.7,3.8,3.9,4.2]
errs = [lowererror;uppererror]
res_tot = [scd2(ti) for ti in texp]
res = [ri[1] for ri in res_tot]
# Error Computation
wls_err_Cd2 = weighted_leastsquare(glcexpCd2,res,uppererror)
print("CD2 wls = ", wls_err_Cd2, "\n")
ls_err_Cd2 = leastsquare(glcexpCd2,res)
# Plotting
plot!(scd2,vars=[(0,1)],label="Glucose (Cd2) simulations")
cur_colors = get_color_palette(:auto, plot_color(:white), 17)
scatter!(texp,glcexpCd2,yerr = errs,marker = stroke(2, cur_colors[5]),
       xlabel="time (min)", ylabel= "Concentration: mg/dL",
       label="Cd2 exp results", legend=:top, markerstrokecolor = :auto)

# Additionnal plot options
plot!(size=(1280,960),legendfontsize=15, xtickfont=font(13), ytickfont=font(13), guidefont=font(18))
# savefig("OGTT_26.pdf")


text:
#   Unit "spectra" of the FourierAnalysis Package for julia language
#   v 0.0.1 - last update 5th of September 2019
#
#   MIT License
#   Copyright (c) 2019, Marco Congedo, CNRS, Grenobe, France:
#   https://sites.google.com/site/marcocongedo/home

# ? CONTENTS :
#   This unit implements Welch power spectral estimates using FFTW

#############################################################################
# Generic constructors of Spectra objects

# Enable construction giving only `y`, `sr`, `wl`, `DC` and `taper` argument.
# `flabels` is generated automatically, `smoothing` is set to `noSmoother`
# and `func` is set to `identity`
Spectra(y, sr, wl, DC, taper) =
    Spectra(y, sr, wl, DC, taper, fdf(sr, wl; DC=DC), identity, noSmoother)
#
# As above, but setting by default also `DC` to false
Spectra(y, sr, wl, taper) =
    Spectra(y, sr, wl, false, taper, fdf(sr, wl; DC=false), identity, noSmoother)
#
#
# Get the spectra from a CrossSpectra object and construct a Spectra object
# The `func` function can be applied element-wise to the spectra (see `spectra`).
Spectra(𝙎::CrossSpectra; func::Function=identity) =
    Spectra(func.(Real.([𝙎.y[i][j, j] for i=1:length(𝙎.y), j=1:size(𝙎.y[1], 2)])), 𝙎.sr, 𝙎.wl, 𝙎.DC, 𝙎.taper, 𝙎.flabels, func, 𝙎.smoothing)
#
# Get the spectra from a CrossSpectraVector and construct a SpectraVector
# The `func` function can be applied element-wise to the spectra (see `spectra`).
Spectra(𝙎::CrossSpectraVector; func::Function=identity) =
    SpectraVector([Spectra(S; func=func) for S in 𝙎])

#############################################################################
"""
```
(1)
function spectra( X     :: Union{Vector{T}, Matrix{T}},
                  sr    :: Int,
                  wl    :: Int;
            tapering  :: Union{Taper, TaperKind} = harris4,
            planner   :: Planner                 = getplanner,
            slide     :: Int                     = 0,
            DC        :: Bool                    = false,
            func      :: Function                = identity, # equal to x->x
            smoothing :: Smoother                = noSmoother,
            ⏩       :: Bool                    = true) where T<:Real

(2)
function spectra( 𝐗   :: Vector{Matrix{T}},
            < same argument sr, ..., ⏩ of method (1) > where T<:Real
```

(1)

Construct a [Spectra](@ref) object from real univariate or
multivariate data. Given sampling rate `sr` and epoch length `wl`,
compute the Welch power spectra of a vector (univariate) or of a data matrix
(multivariate) `X` of dimension ``t``x``n``, where ``t`` is the
number of samples (rows) and ``n`` is the number of time-series (columns).

The spectra are hold in the `.y` field of the created object.
If `X` is a vector, `.y` is a vector, whereas if `X` is a matrix,
`.y` is a matrix holding in its columns the spectra of the
signals in the columns of `X`. The size of the spectra depends on
the `DC` optional keyword argument (see below),
as reported in the documentation of the [Spectra](@ref) structure.

**Optional Keyword Arguments**:

`tapering`: this is a tapering object of type [Taper](@ref) or a tapering
window kind of type [TaperKind](@ref).
By default the *harris4* tapering window is used.
If no tapering is sought, pass as argument `tapering=rectangular`.
This same syntax is the most convenient way to specify all simple
tapering window, e.g., `tapering=hann`, `tapering=hamming`, etc.
For *discrete prolate spheroidal sequences (dpss)* multi-tapering,
use instead the [`slepians`](@ref) constructor, e.g.,
pass as argument something like ```tapering=slepians(sr, wl, 2)```.

`planner`: this is an instance of the [`Planner`](@ref) object, holding the
forward FFTW plan used to compute the FFTs.
By default the planner is computed by this method,
but it can be passed as an
argumet here if it is pre-computed. This is interesting if this
function is to be invoked repeatedly.

`slide`: this is the number of samples the windows slide on (Welch method).
By default the number of samples is chosen to allow 50% overlap.

`DC`: if true, the spectrum/a of the DC level is returned
in the first row of `y` (see the fields of the [Spectra](@ref)
object), otherwise (default) the rows in `y` start with the first positive
discrete frequency, that is, ``sr/wl`` Hz.

`func`: this is a function that operate element-wise to transfrom the power
spectra before returning them, including anonymous functions.
Common choices are:
- `func=sqrt` return the amplitude spectra,
- `func=log` return the log-power spectra,
- `func=dB` return the power spectra in deciBels (see [`dB`](@ref)),
By default no function is applied and the power spectra are returned.
If smoothing has been requested (see below), it is applied after the function.

`smoothing`: it applies a smoothing function of type [Smoother](@ref)
to the spectra across adjacent frequencies. By default no smoothing is applied.

`⏩`: if true (default), the method run in multi-threaded mode
across the series in `X` if the number of series is at least twice
the number of threads Julia is instructed to use. See [Threads](@ref).

(2)

Construct a [SpectraVector](@ref) object from a vector of real
univariate (vectors) or multivariate data (matrices).
Compute the spectra as per method (1)
for all ``k`` data vectors or data matrices in `𝐗`.

If `𝐗` holds data matrices, the ``k`` matrices in `𝐗` must have the same
number of columns (i.e., the same number of time series),
but may have any number of (at least `wl`) rows (samples).
All other arguments have
the same meaning as in method (1), with the following difference:

`⏩`: if true (default), the method run in multi-threaded mode across the
``k`` data matrices if ``k`` is at least twice the number of threads Julia
is instructed to use, otherwise this method attempts to run each spectra
estimation in multi-threaded mode across series as per method (1).
See [Threads](@ref).

If a `planner` is not explicitly passed as an argument,
the FFTW plan is computed once and applied for all spectral estimations.

**See**: [Spectra](@ref), [`plot`](@ref)

**See also**: [`crossSpectra`](@ref), [`coherence`](@ref), [goertzel.jl](@ref).

**Examples**:
```
using FourierAnalysis, Plots

###################################################################

# (1)

# Check that correct amplitude spectra is returned at all discrete
# Fourier Frequency (using a rectangular taper).
# Sinusoids are created at all frequencies with amplitude 10 at the
# first frequency and then incrementing by 10 units along frequencies.
# NB when t is even, correct amplitude for the last frequency is obtained
# only if the sinusoidal has a particular phase.

sr, t, wl= 16, 32, 16
V=Matrix{Float64}(undef, t, wl)
for i=1:wl V[:, i]=sinusoidal(10*i, b2f(i, sr, t), sr, t, π/6) end

# using FFTW.jl only
P=plan_rfft(V, 1)*(2/t)
Σ=abs.(P*V)
bar(Σ[brange(t, DC=true), :], labels="")

# using FourierAnalysis.jl
Σ2=spectra(V, sr, t; tapering=rectangular, func=√, DC=true)
bar(Σ2.y[brange(t, DC=true), :], labels="")

#############################################################################

### Check amplitude spectra on long signals obtained by welch methods
# one sinusoidal is at an exact discrete Fourier Frequency and the other not
# Rectangular window
sr, t, f, a = 128, 128, 10, 0.5
v=sinusoidal(a, f, sr, t*16)+sinusoidal(a, f*3.5+0.5, sr, t*16)+randn(t*16)
Σ=spectra(v, sr, t; tapering=rectangular, func=√)
bar(Σ.y, labels="raw")

# harris4 window (default)
sr, t, f, a = 128, 128, 10, 0.5
v=sinusoidal(a, f, sr, t*16)+sinusoidal(a, f*3.5+0.5, sr, t*16)+randn(t*16)
Σ=spectra(v, sr, t; func=√)
bar(Σ.y, labels="raw")

#smooth spectra
Σ2=smooth(blackmanSmoother, Σ)
bar!(Σ2.y, labels="smoothed")

#############################################################################

function generateSomeData(sr::Int, t::Int; noise::Real=1.)
    # four sinusoids of length t samples and sr sampling rate
    # peak amplitude: 0.7, 0.6, 0.5, 0.4
    # frequency:        5,   7,  13,  27
    # phase:            0, π/4, π/2,   π
    v1=sinusoidal(0.7, 5,  sr, t, 0)
    v2=sinusoidal(0.6, 7,  sr, t, π/4)
    v3=sinusoidal(0.5, 13, sr, t, π/2)
    v4=sinusoidal(0.4, 27, sr, t, π)
    return hcat(v1, v2, v3, v4) + (randn(t, 4)*noise)
end

sr, wl, t = 128, 512, 8192
X=generateSomeData(sr, t) # multivariate data matrix 8192x4

# compute spectra
S=spectra(X, sr, wl)

# check the spectrum of first series
S.y[1, :]

# plot spectra using FourierAnalysis
plot(S; maxf=32)
plot(S; maxf=32, space=2)

# use a custom simple taperig window
S=spectra(X, sr, wl; tapering=riesz)

# use Slepian's multi-tapering
S=spectra(X, sr, wl; tapering=slepians(sr, wl, 1.5))

# construct with smoothing
S=spectra(X, sr, wl; tapering=slepians(sr, wl, 1.5), smoothing=hannSmoother)

# compute Amplitude Spectra instead
S=spectra(X, sr, wl; tapering=slepians(sr, wl, 1.5), func=√)

# plot Aplitude spectra
plot(S; maxf=32, space=2, ylabel="Amplitude")

# smooth the spectra a-posteriori
S=smooth(blackmanSmoother, S)

# extract spectra in range (8Hz to 12Hz)
e=extract(S, (8, 12))

# extract spectra in range (8Hz to 12Hz) for series 1 and 2
e=extract(S, (8, 12))[:, 1:2]

# extract the spectra at 10Hz only (1 bin per series)
e=extract(S, 10)

# average spectra in the 8Hz-12Hz range
bar(mean(S, (8.0, 12.0)))

# average across series of the average spectra in the 8Hz-12Hz range
mean(mean(S, (8.0, 12.0)))

# average spectrum across all frequencies for each series
bar(mean(S, :))

# average spectra in equally-spaced 2Hz-band-pass regions for all series
plot(bands(S, 2))

# average spectra in equally-spaced 2Hz-band-pass regions for series 1 and 2
plot(bands(S, 2)[:, 1:2])

# (2)

# generate 3 multivariate data matrices 8192x4
X=[generateSomeData(sr, t) for i=1:3]

# Now the call to the spectra function will generate 3 Spectra objects
S=spectra(X, sr, wl)
plot(S[1])
plot(S[2])
plot(S[3])

# when you want to compute the spectra of many data matrices you may want
# to do it using a fast FFTW plan (wait up to 10s for computing the plan)
plan=Planner(plan_exhaustive, 10.0, wl)
S=spectra(X, sr, wl; planner=plan)

# average spectra in range (8Hz-12Hz) for all series of all objects
M=mean(S, (8, 12))

# plot the average spectrum across all series for the three objects
# using Julia standard mean function
plot(mean(S[1].y[:, i] for i=1:size(S[1].y, 2)))
plot!(mean(S[2].y[:, i] for i=1:size(S[2].y, 2)))
plot!(mean(S[3].y[:, i] for i=1:size(S[3].y, 2)))

# average spectra in range (4Hz-32.5Hz) across objects for the 4 series
plot(mean(mean(S, (4, 32.5))))

# extract spectra in range (8Hz-12Hz) for all series and all subjects
extract(S, (8, 12))

# if you enter en illegal range, nothing will be done and you will get
# an error in the REPL explaining what is wrong in your argument
extract(S, (0, 12))
mean(S, (1, 128))

# extract 4Hz-band-pass average spectra for all series and all objects
bands(S, 4)

# Apply smoothing in the spectra computations
S=spectra(X, sr, t; smoothing=blackmanSmoother)
plot(S[1])
plot(S[2])
plot(S[3])

# plot spectra in in 1Hz band-pass regions for all series in S[1]
plot(bands(S[1], 1))

# use slepian multi-tapering
S=spectra(X, sr, wl; tapering=slepians(sr, wl, 1.))
plot(S[1])

# average spectra across objects
plot(mean(s.y for s ∈ S))

```
"""
function spectra( X   :: Union{Vector{T}, Matrix{T}},
                  sr  :: Int,
                  wl  :: Int;
            tapering  :: Union{Taper, TaperKind} = harris4,
            planner   :: Planner                 = getplanner,
            slide     :: Int                     = 0,
            DC        :: Bool                    = false,
            func      :: Function                = identity, # equal to x->x
            smoothing :: Smoother                = noSmoother,
            ⏩       :: Bool                    = true) where T<:Real

    # function to set all parameters
    n, 𝚙, 𝜏, s, e, 𝜎, η, wl½, wl½_, cT = _paramSP!(X, tapering, wl, planner, slide, DC)

    # allocate memory and initialize the spectra :  𝚂 is a (wl½_ x n)-Matrix
    𝚂 = zeros(T, wl½_, n)

    # add to 𝚂 the power spectra of a vector or of the column vectors of matrix `X`,
    # computed on the epoch(s) of duration `t` samples starting at sample `at`.
    # ARGUMENTS:
    # `η` is a normalization Vector to be applied to the FFT coefficients, computed
    #    by function `_fftNormalization`. It is applied element-wise to the FFT coefficient vector(s).
    # `n` (Int) is the number of columns in `X`; it must be 1 if `X` is a vector.
    # `𝜏` is a tapering window (Vector), `rectangular` if no tapering is applied (rectangular taper)
    # `𝚙` is the FFTW plan performing the FFT.
    # `DC` (bool) if true the DC level(s) is(are) computed and stored in the first element (row) of 𝚂
    _spectra!(at, 𝚗) =
        𝜏.kind == rectangular ? ( 𝚂[:, 𝚗] += @. abs2( η*$((𝚙*X[at:at+wl-1, 𝚗])[1+!DC:end]) ) ) :
        ( 𝚂[:, 𝚗] += mean(abs2.( η.*(𝚙*(𝜏.y[:, i].*X[at:at+wl-1, 𝚗]))[1+!DC:end] ) for i=1:𝜏.n) )

    # function cumulating the avarege spectra
    sp!(𝚎, 𝚗) = _spectra!((𝚎*s)+1, 𝚗)

    # average spectra across sliding windows, threaded over `n` if `n` is large
    _thread(⏩, n) ? (for 𝚎=0:e-1 @threads for 𝚗=1:n sp!(𝚎, 𝚗) end end) : (for 𝚗=1:n, 𝚎=0:e-1 sp!(𝚎, 𝚗) end)

    # return spectra as a Vector if input data is a vector, as a matrix otherwise
    return smooth(smoothing, Spectra(func.(X isa Vector ? 𝚂[:] : 𝚂), sr, wl, DC, taperinfo(𝜏), fdf(sr, wl; DC=DC), func, smoothing))
end

function spectra( 𝐗   :: Vector{Matrix{T}},
                  sr  :: Int,
                  wl  :: Int;
            tapering  ::  Union{Taper, TaperKind} = harris4,
            planner   ::  Planner                =  getplanner,
            slide     ::  Int                    =  0,
            DC        ::  Bool                   =  false,
            func      ::  Function               =  identity, # equal to x->x
            smoothing ::  Smoother               =  noSmoother,
            ⏩       ::  Bool                   =  true) where T<:Real

    # set some parameters. k=# input data matrices, n=# time series
    k, n, wl½, thr = length(𝐗), size(𝐗[1], 2), (wl>>1), nthreads()

    # get the FFTW plan (for one time-series) if not provided
    planner ≠ getplanner ? plan=planner : plan = _Planner(plan_rfft(zeros(T, wl)))

    # get the tapering window if not provided
    tapering isa Taper ? 𝜏=tapering : 𝜏=taper(tapering, wl)

    # allocate memory for spectra : 𝗦 is a k-Vector of Spectra objects
    𝗦 = SpectraVector(undef, k)

    # function to compute spectra for the ith input data matrix `𝐗[i]`
    sp!(i, ⏩) = 𝗦[i]=spectra(𝐗[i],
                              sr,
                              wl;
                        tapering  = 𝜏,
                        planner   = plan,
                        slide     = slide,
                        DC        = DC,
                        func      = func,
                        smoothing = smoothing,
                        ⏩       = ⏩)

    # threads over k if k is large, otherwise let spectra function decide if to thread over n
    _thread(⏩, k) ? (@threads for i=1:k sp!(i, false) end) : (for i=1:k sp!(i, ⏩) end)

    return 𝗦
end

# Internal function: create a vector for normalizing element-wise the
# FFT coefficients for spectra estimations.
# With this normalization the `spectra` function
# returns amplitude spectra values equal to the peak-amplitude
# when taking as input a synusoid at any phase and any Fourier discrete
# frequency (with a rectangular tapering window).
# NB when wl is even, correct amplitude for the last frequency is obtained
# only if the sinusoidal has a phase=π/6.
#
# ARGUMENTS:
# `wl` = FFT epoch length (Int)
# `wl½` = t÷2 (Int)
# `wl⁻¹` = inv(t) (Float64)
# `two_wl⁻¹` = 2*inv(t)
# `DC` (Bool); is false the output is of size t÷2, otherwise it is of size t÷2+1
_fftNormalization(wl           :: Int,
                  wl½          :: Int,
                  wl⁻¹         :: Real,
                  two_wl⁻¹     :: Real,
                  DC           :: Bool) =
    DC ? [wl⁻¹; [two_wl⁻¹ for i=1:wl½]] : [two_wl⁻¹ for i=1:wl½]

# Internal function: prepare all parameters that are needed for Welch analysis
# in spectral, cross-spectral and coherence estimates
#
# ARGUMENTS:
# `X` is the input data vector or matrix (n time-series)
# `taper` is a a tapering window ( of enum type `TaperTD`)
# `inrange` is the range (UnitRange{Int}, e.g., 128:1280-1) of the vectors
#      or of rows of `X` on which the analysis is to be performed
# `wl` (Int) is the epoch length for FFT
# `plan` a FFTW.rFFTWPlan forward plan for computing the FFT
# `planflag` flags for the computation of the plan by FFYT. See doc of `spectra`
# `plantime` (Int) is the maximum time allowed (in seconds) to optimize the FFTW plan
# `slide` (Int) is the number of samples the windows slides on (Welch method)
# `⏩` if true multi-threaded computations are requested
#
# RETURN
# `n` (Int) 1 if `X` is a Vector, the number of columns of `X` otherwise
# `𝚙` return the forward FFTW plan. It is computed if argument `plan`=`computeplan`
# `𝜏` (Vector or matrix) the computed tapering window(s)
#       or 0 if no tapering is requested (rectangular taper)
# `n𝜏` the number of tapers. 1 for all tapers (𝜏 is a Vector), but for slepians (𝜏 is a matrix)
# `slide` (Int) checked to be valid
# `e` (Int) the number of sliding windows for the validated `inrange`
# `𝜎=√e` (Real) the squre toot of e, used for in-line computation of the avarage
#       across sliding windows as a normalization factor for the FFT coefficients
# `from` the first sample in the validated `inrange`
# `wl½` `wl`÷2 computed by the right aritmetic bit-wise operation `wl`>>1
# `wl½_`, the length of FFT vectors:  wl½+1 if DC=true, wl½ is DC=false
# `cT` the complex type corresponding to type `T` in function definition (e.g., ComplexF64 if T=Float64)
function _paramSP!( X        :: AbstractArray{T},
                    tapering :: Union{Taper, TaperKind},
                    wl       :: Int,
                    planner  :: Planner,
                    slide    :: Int,
                    DC       :: Bool) where T<:Real

    tₐₗₗ, wl½, wl⁻¹, n = size(X, 1), (wl>>1), inv(wl), size(X, 2)
    wl½_ = wl½ + DC
    tₐₗₗ < wl && @error 📌*", the number of samples in input matrix is smaller than the desired window length."
    planner ≠ getplanner ? 𝚙=planner.p :
        𝚙=plan_rfft(zeros(T, wl), flags=plan_estimate, timelimit=-1.0)
    tapering isa Taper ? 𝜏=tapering : 𝜏=taper(tapering, wl) # tapering window
    if slide <= 0 || slide>tₐₗₗ-wl slide = wl½ end # check the number of sliding samples
    e = (tₐₗₗ-wl)÷slide+1 # number of overlapping epochs
    𝜎 = T(√e)
    η = _fftNormalization(wl, wl½, wl⁻¹/𝜎, T(2/wl)/𝜎, DC)
    cT = typeof(Complex(T(0)))

    return n, 𝚙, 𝜏, slide, e, 𝜎, η, wl½, wl½_, cT
end

# ++++++++++++++++++++  Show override  +++++++++++++++++++ # (REPL output)
function Base.show(io::IO, ::MIME{Symbol("text/plain")}, S::Spectra)
    r=size(S.y, 1)
    vec=S.y isa Vector
    vec ? c=1 : c=size(S.y, 2)
    l=length(S.flabels)
    println(io, titleFont, "▥ Spectra type; $r freq. x $c series")
    #println(io, "□  □    □      □        □           □", defaultFont)
    println(io, separatorFont, "⭒  ⭒    ⭒      ⭒        ⭒           ⭒", defaultFont)
    println(io, "sampling rate   (.sr): $(S.sr)")
    println(io, "epoch length    (.wl): $(S.wl)")
    println(io, "DC level        (.DC): $(S.DC)")
    println(io, "taper kind(.tapering): $(S.taper)")
    println(io, "freq. lab. (.flabels): $(l)-", typeof(S.flabels))
    vec ? println(io, "data             (.y): $(r)-", typeof(S.y)) :
          println(io, "data             (.y): $(r)x$(c)-", typeof(S.y))
    # println(io, S.y)
    println(io, "function      (.func): ", string(S.func))
    println(io, "smoother (.smoothing): ", string(S.smoothing))
    r≠l && @warn "number of frequency labels does not match the spectra data size" l r
end

function Base.show(io::IO, ::MIME{Symbol("text/plain")}, 𝐒::SpectraVector)
    println(io, titleFont, "▥ ⋯ ▥ SpectraVector Type")
    println(io, separatorFont, "⭒  ⭒    ⭒      ⭒        ⭒           ⭒", defaultFont)
    println(io, "$(length(𝐒))-element Vector{Spectra}")
end
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ #


text:
using Jute

runtests_dir = Jute.get_runtests_dir()
test_files = [joinpath(runtests_dir, "runtests_benchmark_testcases.jl")]

test_include_only = (ARGS[1] == "test_include_only")

run_options = Jute.build_run_options(options=Dict(:verbosity => 0))

t = time_ns()
Jute.include_test_files!(test_files)
tcs = task_local_storage(Jute.TESTCASE_ACCUM_ID)
if test_include_only
    exitcode = 0
else
    exitcode = Jute.runtests_internal(run_options, tcs)
end
println((time_ns() - t) / 1e9)
exit(exitcode)


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder, Pkg

name = "Yosys"
version = v"0.1.0"

# Collection of sources required to complete build
sources = [
    GitSource("https://github.com/YosysHQ/yosys.git", "2e421feb0ea526468493ab008f3c72beb12c2bc6")
]

dependencies = [
    Dependency("boost_jll"; compat="=1.76.0"), # max gcc7
    Dependency("Readline_jll"),
    Dependency("Tcl_jll"),
    Dependency("Zlib_jll"; compat="~1.2.11"),
    Dependency("Libffi_jll"; compat="~3.2.2"),
]

# Bash recipe for building across all platforms
script = raw"""
cd yosys
if [[ "${target}" == *-apple-* ]] || [[ "${target}" == *-freebsd* ]]; then
    CONFIG=clang
else
    CONFIG=gcc
fi
make CONFIG=${CONFIG} PREFIX=${prefix} -j${nproc}
make install PREFIX=${prefix}
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = filter!(p -> Sys.isapple(p) || Sys.islinux(p), supported_platforms())
platforms = expand_cxxstring_abis(platforms)
# For some reasons, building for CXX03 string ABI doesn't actually work, skip it
filter!(x -> cxxstring_abi(x) != "cxx03", platforms)

# The products that we will ensure are always built
products = Product[
	ExecutableProduct("yosys", :yosys)
	ExecutableProduct("yosys-config", :yosys_config)
	ExecutableProduct("yosys-filterlib", :yosys_filterlib)
	ExecutableProduct("yosys-smtbmc", :yosys_smtbmc)
	ExecutableProduct("yosys-abc", :yosys_abc)
]

# Build the tarballs, and possibly a `build.jl` as well.
# gcc7 constraint from boost
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies; julia_compat="1.6", preferred_gcc_version=v"7")

max_stars_repo_name: wookay/Yggdrasil

max_issues_repo_name: wookay/Yggdrasil

max_forks_repo_name: wookay/Yggdrasil