text:
export *,+, basis_matrix, ambient_space, base_ring, base_field, root_lattice,
kernel_lattice, invariant_lattice
# scope & verbose scope: :Lattice

basis_matrix(L::ZLat) = L.basis_matrix

ambient_space(L::ZLat) = L.space

base_ring(L::ZLat) = FlintZZ

base_field(L::ZLat) = base_ring(gram_matrix(ambient_space(L)))

################################################################################
#
#  Creation
#
################################################################################

@doc Markdown.doc"""
    Zlattice([B::MatElem]; gram) -> ZLat

Return the Z-lattice with basis matrix $B$ inside the quadratic space with
Gram matrix `gram`.

If the keyword `gram` is not specified, the Gram matrix is the identity matrix.
If $B$ is not specified, the basis matrix is the identity matrix.

# Examples

```jldoctest
julia> L = Zlattice(matrix(QQ, 2, 2, [1//2, 0, 0, 2]));

julia> gram_matrix(L) == matrix(QQ, 2, 2, [1//4, 0, 0, 4])
true

julia> L = Zlattice(gram = matrix(ZZ, [2 -1; -1 2]));

julia> gram_matrix(L) == matrix(ZZ, [2 -1; -1 2])
true
```
"""
function Zlattice(B::fmpq_mat; gram = identity_matrix(FlintQQ, ncols(B)))
  @req issymmetric(gram) "Gram matrix must be symmetric"
  V = quadratic_space(FlintQQ, gram)
  return lattice(V, B)
end

function Zlattice(B::fmpz_mat; gram = identity_matrix(FlintQQ, ncols(B)))
  @req issymmetric(gram) "Gram matrix must be symmetric"
  V = quadratic_space(FlintQQ, gram)
  return lattice(V, B)
end

function Zlattice(;gram)
  @req issymmetric(gram) "Gram matrix must be symmetric"
  n = nrows(gram)
  return lattice(quadratic_space(FlintQQ, gram), identity_matrix(FlintQQ, n))
end

@doc Markdown.doc"""
    lattice(V::QuadSpace, B::MatElem) -> ZLat

Return the Z-lattice with basis matrix $B$ inside the quadratic space $V$.
"""
function lattice(V::QuadSpace{FlintRationalField, fmpq_mat}, B::MatElem; isbasis::Bool = true, check::Bool = true)
  @req dim(V) == ncols(B) "Ambient space and the matrix B have incompatible dimension"
  local Gc
  try
    Gc = change_base_ring(FlintQQ, B)
  catch e
    throw(error("Cannot convert entries of the matrix to the rationals"))
  end
  if typeof(Gc) !== fmpq_mat
    throw(error("Cannot convert entries of the matrix to the rationals"))
  end

  # We need to produce a basis matrix

  if !isbasis
    BB = fmpq_mat(hnf(FakeFmpqMat(B), :upper_right))
    i = nrows(BB)
    while i > 0 && iszero_row(BB, i)
      i = i - 1
    end
    return ZLat(V, BB[1:i, :])
  else
    return ZLat(V, Gc)
  end
end

function lattice_in_same_ambient_space(L::ZLat, B::MatElem)
  V = L.space
  return lattice(V,B)
end

function rescale(G::ZLat, r)
  B = basis_matrix(G)
  gram_space = gram_matrix(G.space)
  Vr = quadratic_space(QQ, r*gram_space)
  return lattice(Vr, B)
end

################################################################################
#
#  Gram matrix
#
################################################################################

@doc Markdown.doc"""
    gram_matrix(L::ZLat) -> fmpq_mat

Return the gram matrix of $L$.

# Examples
```jldoctest
julia> L = Zlattice(matrix(ZZ, [2 0; -1 2]));

julia> gram_matrix(L)
[ 4   -2]
[-2    5]
```
"""
function gram_matrix(L::ZLat)
  if isdefined(L, :gram_matrix)
    return L.gram_matrix
  end
  b = basis_matrix(L)
  G = b * gram_matrix(ambient_space(L)) * transpose(b)
  L.gram_matrix = G
  return G
end

gram_matrix_of_rational_span(L::ZLat) = gram_matrix(L)

################################################################################
#
#  Rational span
#
################################################################################

@doc Markdown.doc"""
    rational_span(L::ZLat) -> QuadSpace

Return the rational span of $L$, which is the quadratic space with Gram matrix
equal to `gram_matrix(L)`.

# Examples
```jldoctest
julia> L = Zlattice(matrix(ZZ, [2 0; -1 2]));

julia> rational_span(L)
Quadratic space over
Rational Field
with Gram matrix
[4 -2; -2 5]
```
"""
function rational_span(L::ZLat)
  if isdefined(L, :rational_span)
    return L.rational_span
  else
    G = gram_matrix(L)
    V = quadratic_space(FlintQQ, G)
    L.rational_span = V
    return V
  end
end

################################################################################
#
#  Orthogonal sum
#
################################################################################

@doc Markdown.doc"""
    orthogonal_sum(L1::ZLat, L2::ZLat)

Return the orthogonal direct sum of the lattices `L1` and `L2`.
"""
function orthogonal_sum(L1::ZLat, L2::ZLat)
  V1 = ambient_space(L1)
  V2 = ambient_space(L2)
  B1 = basis_matrix(L1)
  B2 = basis_matrix(L2)
  B = diagonal_matrix(B1, B2)
  V, f1, f2 = orthogonal_sum(V1, V2)
  return lattice(V, B), f1, f2
end

@doc Markdown.doc"""
    orthogonal_submodule(L::ZLat, S::ZLat) -> ZLat

Return the largest submodule of `L` orthogonal to `S`.
"""
function orthogonal_submodule(L::ZLat, S::ZLat)
  @assert ambient_space(L)==ambient_space(S) "L and S must have the same ambient space"
  B = basis_matrix(L)
  C = basis_matrix(S)
  V = ambient_space(L)
  G = gram_matrix(V)
  M = B * G * transpose(C)
  _, K = left_kernel(M)
  K = change_base_ring(ZZ, K*denominator(K))
  Ks = saturate(K)
  return lattice(V, Ks*B)
end
################################################################################
#
#  String I/O
#
################################################################################

function Base.show(io::IO, L::ZLat)
  print(io, "Quadratic lattice of rank ", rank(L),
            " and degree ", degree(L), " over the rationals")
end

################################################################################
#
#  Automorphism groups
#
################################################################################

# This is an internal function, which sets
# L.automorphism_group_generators
# L.automorphism_group_order
function assert_has_automorphisms(L::ZLat; redo::Bool = false,
                                           try_small::Bool = true)

  if !redo && isdefined(L, :automorphism_group_generators)
    return nothing
  end

  if rank(L) == 0
    L.automorphism_group_generators = fmpz_mat[identity_matrix(ZZ, 0)]
    L.automorphism_group_order = one(fmpz)
    return nothing
  end

  V = ambient_space(L)
  GL = gram_matrix(L)
  d = denominator(GL)
  res = fmpz_mat[change_base_ring(FlintZZ, d * GL)]
  # So the first one is either positive definite or negative definite
  # Make it positive definite. This does not change the automorphisms.
  if res[1][1, 1] < 0
    res[1] = -res[1]
  end
  Glll, T = lll_gram_with_transform(res[1])
  Ttr = transpose(T)
  res_orig = copy(res)
  res[1] = Glll

  bm = basis_matrix(L)

  # Make the Gram matrix small

  C = ZLatAutoCtx(res)
  if try_small
    fl, Csmall = try_init_small(C)
    if fl
      auto(Csmall)
      _gens, order = _get_generators(Csmall)
      gens = fmpz_mat[matrix(FlintZZ, g) for g in _gens]
    else
      init(C)
      auto(C)
      gens, order = _get_generators(C)
    end
  else
    init(C)
    auto(C)
    gens, order = _get_generators(C)
  end

  # Now translate back
  Tinv = inv(T)
  for i in 1:length(gens)
    gens[i] = Tinv * gens[i] * T
  end

  # Now gens are with respect to the basis of L
  @hassert :Lattice 1 all(change_base_ring(FlintQQ, gens[i]) * GL *
                          transpose(change_base_ring(FlintQQ, gens[i])) == GL
                                                        for i in 1:length(gens))

  L.automorphism_group_generators = gens
  L.automorphism_group_order = order

  return nothing
end

# documented in ../Lattices.jl

function automorphism_group_generators(L::ZLat; ambient_representation::Bool = true)

  @req rank(L) == 0 || isdefinite(L) "The lattice must be definite"
  assert_has_automorphisms(L)

  gens = L.automorphism_group_generators
  if !ambient_representation
    return fmpq_mat[ change_base_ring(FlintQQ, g) for g in gens]
  else
    # Now translate to get the automorphisms with respect to basis_matrix(L)
    bm = basis_matrix(L)
    gens = L.automorphism_group_generators
    V = ambient_space(L)
    if rank(L) == rank(V)
      bminv = inv(bm)
      res = fmpq_mat[bminv * change_base_ring(FlintQQ, g) * bm for g in gens]
    else
      # Extend trivially to the orthogonal complement of the rational span
      !isregular(V) &&
        throw(error(
          """Can compute ambient representation only if ambient space is
             regular"""))
      C = orthogonal_complement(V, basis_matrix(L))
      C = vcat(basis_matrix(L), C)
      Cinv = inv(C)
      D = identity_matrix(FlintQQ, rank(V) - rank(L))
      res = fmpq_mat[Cinv * diagonal_matrix(change_base_ring(FlintQQ, g), D) * C for g in gens]
    end
    @hassert :Lattice 1 all(g * gram_matrix(V) * transpose(g) == gram_matrix(V)
                            for g in res)
    return res
  end
end

# documented in ../Lattices.jl

function automorphism_group_order(L::ZLat)
  assert_has_automorphisms(L)
  return L.automorphism_group_order
end

################################################################################
#
#  Isometry
#
################################################################################

# documented in ../Lattices.jl

function isisometric(L::ZLat, M::ZLat; ambient_representation::Bool = true)
  @req isdefinite(L) && isdefinite(M) "The lattices must be definite"

  if rank(L) != rank(M)
    return false, zero_matrix(FlintQQ, 0, 0)
  end

  if rank(L) == 0
    return true, identity_matrix(FlintQQ, 0, 0)
  end

  i = sign(gram_matrix(L)[1,1])
  j = sign(gram_matrix(M)[1,1])
  @req i==j "The lattices must have the same signatures"

  if i < 0
    L = rescale(L,-1)
    M = rescale(M,-1)
  end

  GL = gram_matrix(L)
  dL = denominator(GL)
  GLint = change_base_ring(FlintZZ, dL * GL)
  cL = content(GLint)
  GLint = divexact(GLint, cL)

  GM = gram_matrix(M)
  dM = denominator(GM)
  GMint = change_base_ring(FlintZZ, dM * GM)
  cM = content(GMint)
  GMint = divexact(GMint, cM)

  # GLint, GMint are integral, primitive scalings of GL and GM
  # If they are isometric, then the scalars must be identitcal.
  if dL//cL != dM//cM
    return false, zero_matrix(FlintQQ, 0, 0)
  end

  # Now compute LLL reduces gram matrices

  GLlll, TL = lll_gram_with_transform(GLint)
  @hassert :Lattice 1 TL * change_base_ring(FlintZZ, GL) * transpose(TL) * dL == GLlll *cL
  GMlll, TM = lll_gram_with_transform(GMint)
  @hassert :Lattice 1 TM * change_base_ring(FlintZZ, GM) * transpose(TM) * dM == GMlll *cM

  # Setup for Plesken--Souvignier

  G1 = fmpz_mat[GLlll]
  G2 = fmpz_mat[GMlll]

  fl, CLsmall, CMsmall = _try_iso_setup_small(G1, G2)
  if fl
    b, _T = isometry(CLsmall, CMsmall)
    T = matrix(FlintZZ, _T)
  else
    CL, CM = _iso_setup(fmpz_mat[GLlll], fmpz_mat[GMlll])
    b, T = isometry(CL, CM)
  end

  if b
    T = change_base_ring(FlintQQ, inv(TL)*T*TM)
    if !ambient_representation
      @hassert :Lattice 1 T * gram_matrix(M) * transpose(T) == gram_matrix(L)
      return true, T
    else
      V = ambient_space(L)
      W = ambient_space(M)
      if rank(L) == rank(V)
        T = inv(basis_matrix(L)) * T * basis_matrix(M)
      else
        (!isregular(V) || !isregular(W)) &&
          throw(error(
            """Can compute ambient representation only if ambient space is
               regular"""))
          (rank(V) != rank(W)) &&
          throw(error(
            """Can compute ambient representation only if ambient spaces
            have the same dimension."""))

        CV = orthogonal_complement(V, basis_matrix(L))
        CV = vcat(basis_matrix(L), CV)
        CW = orthogonal_complement(W, basis_matrix(M))
        CW = vcat(basis_matrix(M), CW)
        D = identity_matrix(FlintQQ, rank(V) - rank(L))
        T = inv(CV) * diagonal_matrix(T, D) * CW
      end
      @hassert :Lattice 1 T * gram_matrix(ambient_space(M))  * transpose(T) ==
                  gram_matrix(ambient_space(L))
      return true, T
    end
  else
    return false, zero_matrix(FlintQQ, 0, 0)
  end
end

################################################################################
#
#  Is sublattice?
#
################################################################################

function issublattice(M::ZLat, N::ZLat)
  if ambient_space(M) != ambient_space(N)
    return false
  end

  hassol, _rels = can_solve_with_solution(basis_matrix(M), basis_matrix(N), side=:left)

  if !hassol || !isone(denominator(_rels))
    return false
  end

  return true
end

@doc Markdown.doc"""
    issublattice_with_relations(M::ZLat, N::ZLat) -> Bool, fmpq_mat

Returns whether $N$ is a sublattice of $M$. In this case, the second return
value is a matrix $B$ such that $B B_M = B_N$, where $B_M$ and $B_N$ are the
basis matrices of $M$ and $N$ respectively.
"""
function issublattice_with_relations(M::ZLat, N::ZLat)
   if ambient_space(M) != ambient_space(N)
     return false, basis_matrix(M)
   end

   hassol, _rels = can_solve_with_solution(basis_matrix(M), basis_matrix(N), side=:left)

   if !hassol || !isone(denominator(_rels))
     return false, basis_matrix(M)
   end

   return true, _rels
 end

################################################################################
#
#  Root lattice
#
################################################################################

@doc Markdown.doc"""
    root_lattice(R::Symbol, n::Int)

Return the root lattice of type `R` given by ``:A``, ``:D`` or ``:E`` with parameter `n`.
"""
function root_lattice(R::Symbol, n::Int)
  if R === :A
    return Zlattice(gram = _root_lattice_A(n))
  elseif R === :E
    return Zlattice(gram = _root_lattice_E(n))
  elseif R === :D
    return Zlattice(gram = _root_lattice_D(n))
  else
    error("Type (:$R) must be :A, :D or :E")
  end
end

function _root_lattice_A(n::Int)
  n < 0 && error("Parameter ($n) for root lattice of type :A must be positive")
  z = zero_matrix(FlintQQ, n, n)
  for i in 1:n
    z[i, i] = 2
    if i > 1
      z[i, i - 1] = -1
    end
    if i < n
      z[i, i + 1] = -1
    end
  end
  return z
end

function _root_lattice_D(n::Int)
  n < 2 && error("Parameter ($n) for root lattices of type :D must be greater or equal to 2")
  if n == 2
    G = matrix(ZZ, [2 0 ;0 2])
  elseif n == 3
    return _root_lattice_A(n)
  else
    G = zero_matrix(ZZ, n, n)
    G[1,3] = G[3,1] = -1
    for i in 1:n
      G[i,i] = 2
      if 2 <= i <= n-1
        G[i,i+1] = G[i+1,i] = -1
      end
    end
  end
  return G
end

function _root_lattice_E(n::Int)
  n in [6,7,8] || error("Parameter ($n) for lattice of type :E must be 6, 7 or 8")
  if n == 6
    G = [2 -1 0 0 0 0;
        -1 2 -1 0 0 0;
        0 -1 2 -1 0 -1;
        0 0 -1 2 -1 0;
        0 0 0 -1 2 0;
        0 0 -1 0 0 2]
  elseif n == 7
    G = [2 -1 0 0 0 0 0;
        -1 2 -1 0 0 0 0;
        0 -1 2 -1 0 0 -1;
        0 0 -1 2 -1 0 0;
        0 0 0 -1 2 -1 0;
        0 0 0 0 -1 2 0;
        0 0 -1 0 0 0 2]
  else
    G = [2 -1 0 0 0 0 0 0;
        -1 2 -1 0 0 0 0 0;
        0 -1 2 -1 0 0 0 -1;
        0 0 -1 2 -1 0 0 0;
        0 0 0 -1 2 -1 0 0;
        0 0 0 0 -1 2 -1 0;
        0 0 0 0 0 -1 2 0;
        0 0 -1 0 0 0 0 2]
  end
  return matrix(ZZ, G)
end

################################################################################
#
#  Dual lattice
#
################################################################################

# documented in ../Lattices.jl

function dual(L::ZLat)
  G = gram_matrix(L)
  new_bmat = inv(G)*basis_matrix(L)
  return lattice(ambient_space(L), new_bmat)
end

################################################################################
#
#  Scale
#
################################################################################

function scale(L::ZLat)
  if isdefined(L, :scale)
    return L.scale
  end
  G = gram_matrix(L)
  s = zero(fmpq)
  for i in 1:nrows(G)
    for j in 1:i
      s = gcd(s, G[i, j])
    end
  end
  L.scale = s
  return s
end

################################################################################
#
#  Norm
#
################################################################################

function norm(L::ZLat)
  if isdefined(L, :norm)
    return L.norm
  end
  n = 2 * scale(L)
  G = gram_matrix(L)
  for i in 1:nrows(G)
    n = gcd(n, G[i, i])
  end
  L.norm = n
  return n
end

################################################################################
#
#  Eveness
#
################################################################################

iseven(L::ZLat) = isintegral(L) && iseven(numerator(norm(L)))

################################################################################
#
#  Discriminant
#
################################################################################

function discriminant(L::ZLat)
  d = det(gram_matrix(L))
  return d
end

################################################################################
#
#  Determinant
#
################################################################################

function det(L::ZLat)
  return det(gram_matrix(L))
end

################################################################################
#
#  Rank
#
################################################################################

function rank(L::ZLat)
  return nrows(basis_matrix(L))
end

################################################################################
#
#  Local basis matrix
#
################################################################################

# so that abstract lattice functions also work with Z-lattices

local_basis_matrix(L::ZLat, p) = basis_matrix(L)

################################################################################
#
#  Intersection
#
################################################################################

function intersect(M::ZLat, N::ZLat)
  @req ambient_space(M) === ambient_space(N) "Lattices must have same ambient space"
  BM = basis_matrix(M)
  BN = basis_matrix(N)
  dM = denominator(BM)
  dN = denominator(BN)
  d = lcm(dM, dN)
  BMint = change_base_ring(FlintZZ, d * BM)
  BNint = change_base_ring(FlintZZ, d * BN)
  H = vcat(BMint, BNint)
  k, K = left_kernel(H)
  BI = divexact(change_base_ring(FlintQQ, hnf(view(K, 1:k, 1:nrows(BM)) * BMint)), d)
  return lattice(ambient_space(M), BI)
end

################################################################################
#
#  Sum
#
################################################################################

function +(M::ZLat, N::ZLat)
  @req ambient_space(M) === ambient_space(N) "Lattices must have same ambient space"
  BM = basis_matrix(M)
  BN = basis_matrix(N)
  B = fmpq_mat(hnf(FakeFmpqMat(vcat(BM, BN))))
  i = 1
  while iszero_row(B, i)
    i += 1
  end
  return lattice(ambient_space(M), B[i:end, 1:ncols(B)])
end

################################################################################
#
#  Local isometry
#
################################################################################

# TODO: This is slow. We need a dedicated implementation

function islocally_isometric(L::ZLat, M::ZLat, p::Int)
  return islocally_isometric(L, M, fmpz(p))
end

function islocally_isometric(L::ZLat, M::ZLat, p::fmpz)
  LL = _to_number_field_lattice(L)
  K = base_field(LL)
  MM = _to_number_field_lattice(L, K = K)
  OK = maximal_order(K)
  P = prime_decomposition(OK, p)[1][1]
  return islocally_isometric(LL, MM, P)
end

################################################################################
#
#  Conversion between ZLat and QuadLat
#
################################################################################

function _to_number_field_lattice(L::ZLat, K, V)
  LL = lattice(V, change_base_ring(K, basis_matrix(L)))
  return LL
end

function _to_number_field_lattice(L::ZLat;
                                  K::AnticNumberField = rationals_as_number_field()[1],
                                  V::QuadSpace = quadratic_space(K, gram_matrix(ambient_space(L))))
  return _to_number_field_lattice(L, K, V)
end

function _to_ZLat(L::QuadLat, K, V)
  pm = pseudo_matrix(L)
  cm = coefficient_ideals(pm)
  pmm = matrix(pm)
  bm = zero_matrix(FlintQQ, rank(L), dim(V))
  for i in 1:nrows(pm)
    a = norm(cm[i])
    for j in 1:ncols(pm)
      bm[i, j] = a * FlintQQ(pmm[i, j])
    end
  end
  return lattice(V, bm)
end

function _to_ZLat(L::QuadLat;
                  K::FlintRationalField = FlintQQ,
                  V::QuadSpace = quadratic_space(K, map_entries(FlintQQ, gram_matrix(ambient_space(L)))))
  return _to_ZLat(L, K, V)
end

################################################################################
#
#  Mass
#
################################################################################

mass(L::ZLat) = mass(_to_number_field_lattice(L))

################################################################################
#
#  Genus representatives
#
################################################################################

function genus_representatives(L::ZLat)
  LL = _to_number_field_lattice(L)
  K = base_field(L)
  G = genus_representatives(LL)
  res = ZLat[]
  for N in G
    push!(res, _to_ZLat(N, K = K))
  end
  return res
end

################################################################################
#
#  Maximal integral lattice
#
################################################################################

function maximal_integral_lattice(L::ZLat)
  LL = _to_number_field_lattice(L)
  M = maximal_integral_lattice(LL)
  return _to_ZLat(M, V = ambient_space(L))
end

################################################################################
#
#  Scalar multiplication
#
################################################################################

@doc Markdown.doc"""
    *(a, L::ZLat) -> ZLat

Returns the lattice $aM$ inside the ambient space of $M$.
"""
function Base.:(*)(a, L::ZLat)
  @assert has_ambient_space(L)
  B = a*L.basis_matrix
  return lattice_in_same_ambient_space(L, B)
end

function Base.:(*)(L::ZLat, a)
  return a * L
end

################################################################################
#
#  Equality
#
################################################################################

@doc Markdown.doc"""
Return ``true`` if both lattices have the same ambient quadratic space
and the same underlying module.
"""
function Base.:(==)(L1::ZLat, L2::ZLat)
  V1 = ambient_space(L1)
  V2 = ambient_space(L2)
  if V1 != V2
    return false
  end
  B1 = basis_matrix(L1)
  B2 = basis_matrix(L2)
  return hnf(FakeFmpqMat(B1)) == hnf(FakeFmpqMat(B2))
end

@doc Markdown.doc"""
    local_modification(M::ZLat, L::ZLat, p)

Return a local modification of `M` that matches `L` at `p`.

INPUT:

- ``M`` -- a `\ZZ_p`-maximal lattice
- ``L`` -- the a lattice
            isomorphic to `M` over `\QQ_p`
- ``p`` -- a prime number

OUTPUT:

an integral lattice `M'` in the ambient space of `M` such that `M` and `M'` are locally equal at all
completions except at `p` where `M'` is locally isometric to the lattice `L`.
"""
function local_modification(M::ZLat, L::ZLat, p)
  # notation
  d = denominator(inv(gram_matrix(L)))
  level = valuation(d,p)
  d = p^(level+1) # +1 since scale(M) <= 1/2 ZZ

  @req isisometric(L.space, M.space, p) "quadratic spaces must be locally isometric at m"
  L_max = maximal_integral_lattice(L)

  # invert the gerstein operations
  GLm, U = padic_normal_form(gram_matrix(L_max), p, prec=level+3)
  B1 = inv(U*basis_matrix(L_max))

  GM, UM = padic_normal_form(gram_matrix(M), p, prec=level+3)
  # assert GLm == GM at least modulo p^prec
  B2 = B1 * UM * basis_matrix(M)
  Lp = lattice(M.space, B2)

  # the local modification
  S = intersect(Lp, M) + d * M
  # confirm result
  @hassert :Lattice 2 genus(S, p)==genus(L, p)
  return S
end

################################################################################
#
#  Kernel lattice
#
################################################################################

@doc Markdown.doc"""
    kernel_lattice(L::ZLat, f::MatElem;
                   ambient_representation::Bool = true) -> ZLat

Given a $\mathbf{Z}$-lattice $L$ and a matrix $f$ inducing an endomorphism of
$L$, return $\ker(f)$ is a sublattice of $L$.

If `ambient_representation` is `true` (the default), the endomorphism is
represented with respect to the ambient space of $L$. Otherwise, the
endomorphism is represented with respect to the basis of $L$.
"""
function kernel_lattice(L::ZLat, f::MatElem; ambient_representation::Bool = true)
  bL = basis_matrix(L)
  if ambient_representation
    if !issquare(bL)
      fl, finL = can_solve_with_solution(bL, bL * f, side = :left)
      @req fl "f must preserve the lattice L"
    else
      finL = bL * f * inv(bL)
    end
  else
    finL = f
  end
  k, K = left_kernel(change_base_ring(ZZ, finL))
  return lattice(ambient_space(L), K * basis_matrix(L))
end

################################################################################
#
#  Invariant lattice
#
################################################################################

@doc Markdown.doc"""
    invariant_lattice(L::ZLat, G::Vector{MatElem};
                      ambient_representation::Bool = true) -> ZLat

Given a $\mathbf{Z}$-lattice $L$ and a list of matrices $G$ inducing
endomorphisms of $L$, return the lattice $L^G$, consisting of elements fixed by
$G$.

If `ambient_representation` is `true` (the default), the endomorphism is
represented with respect to the ambient space of $L$. Otherwise, the
endomorphism is represented with respect to the basis of $L$.
"""
function invariant_lattice(L::ZLat, G::Vector{<:MatElem};
                           ambient_representation::Bool = true)
  if length(G) == 0
    return L
  end

  M = kernel_lattice(L, G[1] - 1,
                     ambient_representation = ambient_representation)
  for i in 2:length(G)
    N = kernel_lattice(L, G[i] - 1,
                       ambient_representation = ambient_representation)
    M = intersect(M, N)
  end
  return M
end

function invariant_lattice(L::ZLat, G::MatElem;
                           ambient_representation::Bool = true)
  return kernel_lattice(L, G - 1, ambient_representation = ambient_representation)
end

################################################################################
#
#  Membership check
#
################################################################################

@doc Markdown.doc"""
    Base.in(v::Vector, L::ZLat) -> Bool

  Check if the vector `v` lies in the lattice `L` or not.
"""
function Base.in(v::Vector, L::ZLat)
  @assert length(v) == degree(L) "The vector should have the same length as the degree of the lattice."
  V = matrix(QQ, 1, length(v), v)
  return V in L
end

@doc Markdown.doc"""
    Base.in(v::fmpq_mat, L::ZLat) -> Bool

  Check if the row span of `v` lies in the lattice `L` or not.
"""
function Base.in(v::fmpq_mat, L::ZLat)
  @assert ncols(v)==degree(L) "The vector should have the same length as the degree of the lattice."
  B = basis_matrix(L)
  fl, w = can_solve_with_solution(B, v, side=:left)
  return fl && isone(denominator(w))
end


text:
using BinDeps

@BinDeps.setup

ENV["JULIA_ROOT"] = abspath(JULIA_HOME, "../../")

libportaudio = library_dependency("libportaudio")
libsndfile = library_dependency("libsndfile")

# TODO: add other providers with correct names
provides(AptGet, {"portaudio19-dev" => libportaudio})
provides(AptGet, {"libsndfile1-dev" => libsndfile})
provides(Pacman, {"portaudio" => libportaudio})
provides(Pacman, {"libsndfile" => libsndfile})

@osx_only begin
    using Homebrew
    provides(Homebrew.HB, {"portaudio" => libportaudio})
    provides(Homebrew.HB, {"libsndfile" => libsndfile})
end

@BinDeps.install [:libportaudio => :libportaudio,
                  :libsndfile => :libsndfile]

ext: jl


text:
# Regular functions for evaluating Chebyshev polynominals

function chebyshev_evaluate(weights::AbstractArray{T,N},x::AbstractArray{R,1},order::Array{S,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,R<:Number,N,S<:Integer}

  poly = Array{Array{R,2},1}(undef,N)
  @inbounds for i = 1:N
    poly[i] = chebyshev_polynomial(order[i],normalize_node(x[i],domain[:,i]))
  end

  yhat = zero(T)
  @inbounds for i in CartesianIndices(weights)
    poly_product = poly[1][i[1]]
    @inbounds for j = 2:N
      poly_product *= poly[j][i[j]]
    end
    yhat += weights[i]*poly_product
  end

  return yhat

end

function chebyshev_evaluate(weights::AbstractArray{T,N},x::AbstractArray{R,1},order::S,domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,R<:Number,N,S<:Integer}

  poly = Array{Array{R,2},1}(undef,N)
  @inbounds for i = 1:N
    poly[i] = chebyshev_polynomial(order,normalize_node(x[i],domain[:,i]))
  end

  yhat = zero(T)
  @inbounds for i in CartesianIndices(weights)
    if sum(Tuple(i)) <= order+N
      poly_product = poly[1][i[1]]
      @inbounds for j = 2:N
        poly_product *= poly[j][i[j]]
      end
      yhat += weights[i]*poly_product
    end
  end

  return yhat

end

function chebyshev_evaluate(cheb_poly::ChebPoly,x::AbstractArray{R,1}) where {R<:Number}

  yhat = chebyshev_evaluate(cheb_poly.weights,x,cheb_poly.order,cheb_poly.domain)

  return yhat

end

function cheb_interp(cheb::ChebInterpRoots,x::AbstractArray{R,1}) where {R<:Number}

  weights = chebyshev_weights_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)
  yhat    = chebyshev_evaluate(weights,x,cheb.order,cheb.domain)

  return yhat

end

function cheb_interp(cheb::ChebInterpExtrema,x::AbstractArray{R,1}) where {R<:Number}

  weights = chebyshev_weights_extrema_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)
  yhat    = chebyshev_evaluate(weights,x,cheb.order,cheb.domain)

  return yhat

end

function cheb_interp(cheb::ChebInterpExtended,x::AbstractArray{R,1}) where {R<:Number}

  weights = chebyshev_weights_extended_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)
  yhat    = chebyshev_evaluate(weights,x,cheb.order,cheb.domain)

  return yhat

end

function cheb_interp(cheb::ChebInterpVertesi,x::AbstractArray{R,1}) where {R<:Number}

  weights = chebyshev_weights_vertesi_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)
  yhat    = chebyshev_evaluate(weights,x,cheb.order,cheb.domain)

  return yhat

end

function chebyshev_evaluate(weights::AbstractArray{T,N},order::Array{S,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T <: AbstractFloat,N,S <: Integer}

  function chebeval(x::AbstractArray{R,1}) where {R <: Number}

    return chebyshev_evaluate(weights,x,order,domain)

  end

  return chebeval

end

function chebyshev_evaluate(weights::AbstractArray{T,N},order::S,domain=[ones(T,1,N);-ones(T,1,N)]) where {T <: AbstractFloat,N,S <: Integer}

  function chebeval(x::AbstractArray{R,1}) where {R <: Number}

    return chebyshev_evaluate(weights,x,order,domain)

  end

  return chebeval

end

function chebyshev_evaluate(cheb_poly::ChebPoly)

  function chebeval(x::AbstractArray{R,1}) where {R <: Number}

    return chebyshev_evaluate(cheb_poly,x)

  end

  return chebeval

end

function cheb_interp(cheb::ChebInterpRoots)

  weights = chebyshev_weights_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)

  function chebeval(x::AbstractArray{R,1}) where {R <: Number}

    return chebyshev_evaluate(weights,x,cheb.order,cheb.domain)

  end

  return chebeval

end

function cheb_interp(cheb::ChebInterpExtrema)

  weights = chebyshev_weights_extrema_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)

  function chebeval(x::AbstractArray{R,1}) where {R <: Number}

    return chebyshev_evaluate(weights,x,cheb.order,cheb.domain)

  end

  return chebeval

end

function cheb_interp(cheb::ChebInterpExtended)

  weights = chebyshev_weights_extended_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)

  function chebeval(x::AbstractArray{R,1}) where {R <: Number}

    return chebyshev_evaluate(weights,x,cheb.order,cheb.domain)

  end

  return chebeval

end

function cheb_interp(cheb::ChebInterpVertesi)

  weights = chebyshev_weights_vertesi_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)

  function chebeval(x::AbstractArray{R,1}) where {R <: Number}

    return chebyshev_evaluate(weights,x,cheb.order,cheb.domain)

  end

  return chebeval

end

max_stars_repo_name: RJDennis/ChebyshevApprox

max_issues_repo_name: RJDennis/ChebyshevApprox

max_forks_repo_name: RJDennis/ChebyshevApprox


text:
using DirectGaussianSimulation
using GeoStatsBase
using Variography
using Plots, VisualRegressionTests
using Test, Pkg, Random

# workaround GR warnings
ENV["GKSwstype"] = "100"

# environment settings
islinux = Sys.islinux()
istravis = "TRAVIS" ∈ keys(ENV)
datadir = joinpath(@__DIR__,"data")
visualtests = !istravis || (istravis && islinux)
if !istravis
  Pkg.add("Gtk")
  using Gtk
end

@testset "DirectGaussianSimulation.jl" begin
  @testset "Conditional simulation" begin
    sdata = georef((z=[0.,1.,0.,1.,0.],), [0. 25. 50. 75. 100.])
    sdomain = RegularGrid(100)
    problem = SimulationProblem(sdata, sdomain, :z, 2)
    solver = DirectGaussSim(:z => (variogram=SphericalVariogram(range=10.),))

    Random.seed!(2018)
    solution = solve(problem, solver)

    if visualtests
      @plottest plot(solution,size=(600,400),layout=(2,1)) joinpath(datadir,"CondSimSol.png") !istravis
    end
  end

  @testset "Unconditional simulation" begin
    sdomain = RegularGrid(100)
    problem = SimulationProblem(sdomain, :z => Float64, 2)
    solver = DirectGaussSim(:z => (variogram=SphericalVariogram(range=10.),))

    Random.seed!(2018)
    solution = solve(problem, solver)

    if visualtests
      @plottest plot(solution,size=(600,400),layout=(2,1)) joinpath(datadir,"UncondSimSol.png") !istravis
    end
  end

  @testset "Cosimulation" begin
    sdata = georef((z=[0.,1.,0.,1.,0.], y=[0.,1.,0.,1.,0.]), [0. 25. 50. 75. 100.])
    sdomain = RegularGrid(500)
    problem = SimulationProblem(sdomain, (:z=>Float64,:y=>Float64), 1)
    solver = DirectGaussSim(:z => (variogram=SphericalVariogram(range=10.),),
                            :y => (variogram=GaussianVariogram(range=10.),),
                            (:z,:y) => (correlation=0.95,))

    Random.seed!(2020)
    solution = solve(problem, solver)

    if visualtests
      @plottest plot(solution,size=(600,400),layout=(2,1)) joinpath(datadir,"CoSimSol.png") !istravis
    end
  end
end


text:
## Check functions work on real data

fname = joinpath(dirname(@__FILE__), "../data", "test_Hz19.5-testing.bdf")

s = read_SSR(fname)
println(s)

s.modulationrate = assr_frequency(40)

s = rereference(s, "Cz")

s = kalman_filter(s)

println(s.processing["statistics"])

## Check values using simulated data

model = Private.acoustic_model(assr_frequency(40), 8192.0, 1e-10, 5.0, 0.01, [1.0, 1.0])

new_data = Array(Float64, 8192*10, 6)
for n in 1:size(s.data, 2)
    x, new_data[:, n] = TimeModels.simulate(model, 8192*10)
end
s.data = new_data

s = kalman_filter(s)
println(s.processing["statistics"])

@test_approx_eq_eps mean(s.processing["statistics"][:SignalAmplitude][7:end]) sqrt(2) 0.2

#
# Test setting nans
#

fname = joinpath(dirname(@__FILE__), "..", "data", "test_Hz19.5-testing.bdf")
a = read_SSR(fname)
a = keep_channel!(a, "Cz")
@test any(isnan(a.data)) == false
a = set_nans(a)
@test any(isnan(a.data)) == true
@test (sum(isnan(a.data)) / length(a.data)) > 0.1
@test (sum(isnan(a.data)) / length(a.data)) < 0.2
# Running it again shouldnt change anything
a = set_nans(a)
@test any(isnan(a.data)) == true
@test (sum(isnan(a.data)) / length(a.data)) > 0.1
@test (sum(isnan(a.data)) / length(a.data)) < 0.2
# Up the limits
a = set_nans(a, perc = 0.3)
@test (sum(isnan(a.data)) / length(a.data)) > 0.3
@test (sum(isnan(a.data)) / length(a.data)) < 0.5

#
# Artifact model
#

# Generate a third of second of data with 300 pulses per second
p = Private.triangle_artifact(0:1/81920:0.33, 300)
Plots.unicodeplots()
Plots.plot(p)


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

name = "alsa"
# It's actually v1.2.5.1
version = v"1.2.5"

# Collection of sources required to complete build
sources = [
    ArchiveSource("https://www.alsa-project.org/files/pub/lib/alsa-lib-1.2.5.1.tar.bz2",
                  "628421d950cecaf234de3f899d520c0a6923313c964ad751ffac081df331438e"),
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir/alsa-lib*/
./configure --prefix=${prefix} --build=${MACHTYPE} --host=${target}
make
make install
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = filter(Sys.islinux, supported_platforms(;experimental=true))

# The products that we will ensure are always built
products = [
    LibraryProduct("libasound", :libasound),
    LibraryProduct("libatopology", :libatopology),
]

# Dependencies that must be installed before this package can be built
dependencies = [
]

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

max_stars_repo_name: sharanry/Yggdrasil

max_issues_repo_name: sharanry/Yggdrasil

max_forks_repo_name: sharanry/Yggdrasil


text:
using Base.Test
import WaveletScattering: Log1P

x = [0.0, 1.0, 2.0, 3.0]
data_in = expm1.(x)
data_out = zeros(data_in)
ρ = Log1P(1.0)
Base.map!(ρ, data_out, data_in)
@test_approx_eq x data_out


text:
# This file is a part of TypeDBClient.  License is MIT: https://github.com/Humans-of-Julia/TypeDBClient.jl/blob/main/LICENSE

mutable struct  BidirectionalStream
    resCollector::ResponseCollector
    dispatcher::Dispatcher
    is_open::Threads.Atomic{Bool}
    input_channel::Channel{Proto.Transaction_Client}
    output_channel::Channel{Proto.Transaction_Server}
    status::Task
    error_break_time::Real
end

function BidirectionalStream(input_channel::Channel{Proto.Transaction_Client},
                             output_channel::Channel{Proto.Transaction_Server},
                             status::Task;
                             error_break_time::Real)

    dispatcher = Dispatcher(input_channel)
    res_collector = ResponseCollector(output_channel)

    return BidirectionalStream(res_collector, dispatcher, Threads.Atomic{Bool}(true),input_channel, output_channel, status, error_break_time)
end

function single_request(bidirect_stream::BidirectionalStream, request::Proto.ProtoType)
    return single_request(bidirect_stream, request, true)
end

"""
function single_request(bidirect_stream::BidirectionalStream, request::T, batch::Bool) where {T<: Proto.ProtoType}
    This function process one single request and give back the results from the server to the calling functions. It is an intern
    function.
"""
function single_request(bidirect_stream::BidirectionalStream, request::Proto.Transaction_Req, batch::Bool)
    answer = _process_request(bidirect_stream, request, batch)
    return isempty(answer) ? nothing : answer[1]
end

"""
function stream_request(bidirect_stream::BidirectionalStream, request::T, batch::Bool) where {T<: Proto.ProtoType}
    Here we let the user decide what to do with the result channel. The returned value is the pure result_channel.
"""
function stream_request(bidirect_stream::BidirectionalStream, request::Proto.Transaction_Req, batch::Bool)
    answer = _process_request(bidirect_stream, request, batch)
    return answer
end

# ToDo: Make the request timout determined by user at session level as a default and in the transaction level
function _process_request(bidirect_stream::BidirectionalStream, request::Proto.Transaction_Req, batch::Bool)

    res_channel = _open_result_channel(bidirect_stream, request, batch)
    # start a task to collect the results asynchronusly
    result_task = Threads.@spawn collect_result(res_channel, bidirect_stream, request.req_id)

    # until solving the absent possibility to detect grpc errors in the gRPCClient a pure time
    # dependent solutionresult = nothing
    answer = Vector{Transaction_Res_All}()
    contr = Controller(true, bidirect_stream.error_break_time)
    @async sleeper(contr)
    while contr.running
        yield()
        if istaskdone(result_task)
            answer = fetch(result_task)
            break
        end
    end

    if !istaskdone(result_task)
        close(bidirect_stream.input_channel)
        failure_stat = fetch(bidirect_stream.status)
        close(bidirect_stream)
        try
            gRPCCheck(failure_stat)
        catch ex
            throw(TypeDBClientException(ex))
        end
    end

    return answer
end

function _open_result_channel(bidirect_stream::BidirectionalStream, request::Proto.Transaction_Req, batch::Bool)
    # This function serves as the opener for the result channel for both single- and stream request

    # put the needed request_id on the request
    request.req_id = bytes(uuid4())
    # get the channel which stores the result of the request
    res_channel = push!(bidirect_stream.resCollector, request.req_id)

    # direct the request to the right channel direct or batched processing
    if batch
        put!(bidirect_stream.dispatcher.dispatch_channel, request)
    else
        put!(bidirect_stream.dispatcher.direct_dispatch_channel, request)
    end
    return res_channel
end

"""
function collect_result(res_channel::Channel{T}) where {T<:ProtoProtoType}
    The function will be called for each single request. She works until
    the whole result set will be collected.
"""
function collect_result(res_channel::Channel{Transaction_Res_All}, bidirect_stream::BidirectionalStream, request_id::Bytes)
    answers = Vector{Transaction_Res_All}()
    for tmp_result in res_channel
        req_push, loop_break = _is_stream_respart_done(tmp_result, bidirect_stream)
        req_push && push!(answers, tmp_result)
        if loop_break
            delete!(bidirect_stream.resCollector, request_id)
            break
        end
    end
    return answers
end

"""
function _is_stream_respart_done(req_result::Proto.ProtoType)
    This function decides how to treat the result. It returns whether it should push the
    request to the answers and if it should break the retrieving loop.
"""
function _is_stream_respart_done(req_result::Transaction_Res_All, bidirect_stream::BidirectionalStream)
    kind_of_content = Proto.which_oneof(req_result, :res)
    request_content = getproperty(req_result, kind_of_content)
    type_of_result = typeof(request_content)
    req_push = false
    loop_break = false

    if type_of_result == Proto.Transaction_Stream_ResPart
        if request_content.state == Proto.Transaction_Stream_State[:DONE]
            req_push = false
            loop_break = true
        else
            req_push = false
            loop_break = false
            # if the Stream_State is :CONTINUE put a Stream requ on the dispatch_channel
            put!(bidirect_stream.dispatcher.direct_dispatch_channel,
                  TransactionRequestBuilder.stream_req(req_result.req_id))
        end
    elseif typeof(req_result) == Proto.Transaction_Res
            req_push = true
            loop_break = true
            @debug "answer to req_id: $(req_result.req_id)"
    elseif typeof(req_result) == Proto.Transaction_ResPart
        req_push = true
        loop_break = false
        @debug "answer to req_id: $(req_result.req_id)"
    else
        throw(TypeDBClientException("Not known result type"))
    end

    return req_push, loop_break
end

function close(stream::BidirectionalStream)
    safe_close(stream.dispatcher)
    task_to_close = @async begin
        while true
            yield()
            isempty(stream.resCollector.collectors) && break
            for (k,_) in stream.resCollector.collectors
                delete!(stream.resCollector.collectors, k)
            end
        end
    end
    wait(task_to_close)
    safe_close(stream.resCollector)
    return true
end


text:
# MultinomialSeries.jl -- Part of 'MultinomialSeries',
# a package for The Julia Programming Language
# Copyright (C) 2020-2021 Markus J. Weber
# License: See accompanying file 'LICENSE.md'

"""
    module MultinomialSeries

Compute multinomial coefficients and natively iterate over multinomial expansions in Julia.

Exported functions are:

    multinomial(m,n)

    eachmultinomial(m,n)

"""
module MultinomialSeries

export multinomial, eachmultinomial

import Base: eltype, iterate, length,
             IteratorSize, IteratorEltype

using StaticArrays: MVector

struct MultinomialIterator{T,N,P} end

eachmultinomial(dims, order) = MultinomialIterator{Int,dims,order}()
eachmultinomial(::Type{T}, dims, order) where T<:Integer = MultinomialIterator{T,dims,order}()

multinomial(k::NTuple{N,T}) where {N,T} = T(factorial(sum(k)) ÷ prod(factorial.(k)))

# See https://en.wikipedia.org/wiki/Multinomial_theorem?section=5#Number_of_multinomial_coefficients
# here: P ≡ m, N ≡ m
Base.length(::MultinomialIterator{T,N,P}) where {T,N,P} = binomial(P+N-1, N-1)

Base.eltype(::MultinomialIterator{T,N,P}) where {T,N,P} = NTuple{N,T}

Base.IteratorSize(::MultinomialIterator) = Base.HasLength()
Base.IteratorEltype(::MultinomialIterator) = Base.HasEltype()

# iterate impementation note:
# State is a MVector of k₁...kₘ (m=N), plus a pointer to the current active iteration element in k.
# The element data type returned is a NTuple{N,T}, which is simply pulled out from the MVector.
function Base.iterate(
        ::MultinomialIterator{T,N,P}
    ) where {T,N,P}

    local k::NTuple{N,T} = (T(P), (zero(T) for n in 1:N-1)...)
    return ( k,                      # element
            (MVector{N,T}(k), 1) )   # state

end

function Base.iterate(
        ::MultinomialIterator{T,N,P},
        state
    ) where {T,N,P}

    k, ptr = state

    while @inbounds k[ptr] == 0
        ptr -= 1
        P == 0 && ptr == 0 && @goto finish  # only relevant if order == 0
    end

    if ptr == N  # ==length(k)
        while true
            ptr == 1 && @goto finish        # done !
            ptr -= 1
            @inbounds k[ptr] == 0 || break
        end
        @inbounds balls = sum(k[ptr+1:end])
        @inbounds k[ptr+1:end] .= 0         # reset everything after current position to zero
        @inbounds k[ptr+1] += balls
    end

    if @inbounds k[ptr] > 0
        @inbounds k[ptr] -= 1
        @inbounds k[ptr += 1] += 1
    end

    return (k.data, (k, ptr))

    @label finish
    return nothing

end  # function iterate

end  # module MultinomialSeries

# end of file


text:
"""
Non-relativistic configurations and CSF definition lists, for generating inputs
to `rcsfgenerate`.
"""
module Configurations

import ..GRASP: angularmomentum, nelectrons, maxelectrons, nexcitations
import ..GRASP: SPECTROSCOPIC_NAMES, CSF, specname
using AtomicLevels: κ2ℓ, κ2j

struct CSFOrbital
    n :: Int
    l :: Int

    function CSFOrbital(n, l)
        @assert n >= 1 && l >= 0
        @assert l < n
        new(n, l)
    end
end

import Base: ==
==(a::CSFOrbital, b::CSFOrbital) = (a.n == b.n) && (a.l == b.l)

Base.isless(a::CSFOrbital, b::CSFOrbital) = (a.n < b.n) || (a.n == b.n && a.l < b.l)

function Base.show(io::IO, co::CSFOrbital)
    write(io, "CSFOrbital(\"")
    print(io, co)
    write(io, "\")")
end

function Base.print(io::IO, co::CSFOrbital)
    write(io, string(co.n), specname(co.l));
    nothing
end

maxelectrons(co::CSFOrbital) = 4*co.l + 2

struct CSFDefinition
    orbitals :: Vector{CSFOrbital}
    "Number of allowed excitations."
    nelectrons :: Vector{Int}
    "Number of allowed excitations."
    nexcitations :: Vector{Int}
end
CSFDefinition() = CSFDefinition([], [], [])
function CSFDefinition(csfdef::CSFDefinition)
    CSFDefinition(
        copy(csfdef.orbitals),
        copy(csfdef.nelectrons),
        copy(csfdef.nexcitations)
    )
end

function Base.push!(cdef::CSFDefinition, orbital::CSFOrbital, nelec, nexc)
    # TODO: Sanity checks!
    push!(cdef.orbitals, orbital)
    push!(cdef.nelectrons, nelec)
    push!(cdef.nexcitations, nexc)
    cdef
end

Base.length(cdef::CSFDefinition) = length(cdef.orbitals)

function Base.iterate(cdef::CSFDefinition, s = 1)
    @assert length(cdef.orbitals) == length(cdef.nelectrons) == length(cdef.nexcitations) # sanity check
    s > length(cdef.orbitals) && return nothing
    item = (cdef.orbitals[s], cdef.nelectrons[s], cdef.nexcitations[s])
    return item, s + 1
end

nelectrons(cdef::CSFDefinition) = sum(cdef.nelectrons)

function nexcitations(csf_from::CSFDefinition, csf_to::CSFDefinition)
    @assert nelectrons(csf_from) == nelectrons(csf_to)
    let from_orbs = [(orb.n, orb.l) for orb in csf_from.orbitals],
        to_orbs = [(orb.n, orb.l) for orb in csf_to.orbitals]
        @assert length(from_orbs) == length(unique(from_orbs))
        @assert length(to_orbs) == length(unique(to_orbs))
    end

    to_unchecked = fill(true, length(csf_to.orbitals))

    nexcitations = 0
    nremoved = 0 # for sanity checking only
    for (i, orb) in enumerate(csf_from.orbitals)
        # Let's count all the electrons that have been excited into some of
        # the occupied orbitals of csf_from.
        toids = findall(o -> o.n == orb.n && o.l == orb.l, csf_to.orbitals)
        @assert length(toids) < 2
        if length(toids) == 0
            nremoved += csf_from.nelectrons[i]
        end
        if length(toids) == 1
            toid = first(toids)
            @assert to_unchecked[toid]
            to_unchecked[toid] = false
            diff = csf_to.nelectrons[toid] - csf_from.nelectrons[i]
            # If diff < 0, electrons have been excited from this orbital
            (diff < 0) && (nremoved += abs(diff)) # for sanity checking
            nexcitations += (diff > 0) ? diff : 0
        end
    end
    for (i, orb) in enumerate(csf_to.orbitals)
        to_unchecked[i] || continue
        nexcitations += csf_to.nelectrons[i]
    end
    @assert nexcitations == nremoved # sanity check
    return nexcitations
end

function Base.print(io::IO, cdef::CSFDefinition)
    for (orb, nelec, nexc) in cdef
        nexc_str = if nexc == nelec
            "*"
        elseif nexc == 0
            "i"
        else
            string(nelec - nexc)
        end
        print(io, orb)
        write(io, "(", string(nelec), ",", nexc_str, ")")
    end
    nothing
end

import Base: *
"""
For two [`CSFDefinition`](@ref) objects, merge the CSF definitions lists together, or thrown
an error if the lists have overlapping orbitals.
"""
function *(a::CSFDefinition, b::CSFDefinition)
    if !isempty(intersect(a.orbitals, b.orbitals))
        isect = join(intersect(a.orbitals, b.orbitals), ",")
        error("""
        CSFDefinitions contain overlapping orbitals ($(isect)).
          1: $(a)
          2: $(b)
        """)
    end

    CSFDefinition(
        vcat(a.orbitals, b.orbitals),
        vcat(a.nelectrons, b.nelectrons),
        vcat(a.nexcitations, b.nexcitations)
    )
end

mutable struct CSFDefinitionList
    cdefs :: Vector{CSFDefinition}
end

function Base.print(io::IO, cdl::CSFDefinitionList)
    first = true
    for cdef in cdl.cdefs
        first ? (first = false) : println(io)
        print(io, cdef)
    end
    nothing
end

"""
    csfdefinition(csf::CSF)

Reduces a `CSF` down to a `CSFDefinition`, by merging the orbitals with same `l` but a
different `j` value.
"""
function csfdefinition(csf::CSF)
    nrelos = Vector{Tuple{Int,Int}}()
    nelecs = Dict{Tuple{Int,Int}, Int}()
    for (orb, nelec, orbcoupling, csfcoupling) in csf
        nl = orb.n, κ2ℓ(orb.κ)
        if !(nl in keys(nelecs))
            push!(nrelos, nl)
            nelecs[nl] = 0
        end
        nelecs[nl] += nelec
    end

    cd = CSFDefinition()
    for nl in nrelos
        push!(cd, CSFOrbital(nl...), nelecs[nl], 0)
    end
    return cd
end

const NUMBERS = [Char(48+i) for i=0:9]

function Base.parse(::Type{CSFDefinition}, str)
     # Find all the opening and closing parens first.
    lparens, rparens = Int[], Int[]
    for (i, c) in enumerate(str)
        if c == '('
            (length(lparens) == length(rparens)) || throw(ArgumentError("Unexpected '(' at $i in '$str'"))
            push!(lparens, nextind(str, 0, i))
        elseif c == ')'
            (length(lparens) - length(rparens) == 1) || throw(ArgumentError("Unexpected ')' at $i in '$str'"))
            push!(rparens, nextind(str, 0, i))
        end
    end
    (length(lparens) == length(rparens)) || throw(ArgumentError("Missing final '(' in '$str'"))

    cd = CSFDefinition()
    for i = 1:length(lparens)
        ss_idx = (i == 1) ? 1 : nextind(str, 0, length(str, 1, rparens[i-1]) + 1)
        ss_end = nextind(str, 0, length(str, 1, lparens[i]) - 1)
        ss_nl = str[ss_idx:ss_end]

        ss_idx = nextind(str, 0, length(str, 1, lparens[i]) + 1)
        ss_end = nextind(str, 0, length(str, 1, rparens[i]) - 1)
        ss_el = str[ss_idx:ss_end]

        # parse the n and l value
        ss_nl = strip(ss_nl)

        last_number = 0
        for (i, c) in enumerate(ss_nl)
            c in NUMBERS || break
            last_number = i
        end
        last_number == 0 && throw(ArgumentError("Missing principal value in `$str`"))

        ss_n = ss_nl[1:nextind(ss_nl, 0, last_number)]
        n = parse(Int, ss_n)

        ss_l = ss_nl[nextind(ss_nl, 0, last_number+1):end]
        l = parse_l(ss_l)
        l >= 0 || throw(ArgumentError("Bad spectroscopic orbital `$ss_l` in `$str`"))

        # parse the electron / excitation counts
        ss_el_split = split(ss_el, ',')
        (1 <= length(ss_el_split) <= 2) || throw(ArgumentError("Bad number of terms in ss_el: $ss_el"))
        nelec = parse(Int, ss_el_split[1])
        nexc = if length(ss_el_split) == 1 || ss_el_split[2] == "i"
            0
        elseif ss_el_split[2] == "*"
            nelec
        else
            parse(Int, ss_el_split[2])
        end

        push!(cd, CSFOrbital(n, l), nelec, nexc)
    end

    return cd
end

function parse_l(lstring)
    idx = findfirst(ls -> ls == strip(lstring), SPECTROSCOPIC_NAMES)
    if idx === nothing
        throw(ArgumentError("Unable to parse `$(lstring)`"))
    end
    return idx - 1
end

function parse_j(s)
    if endswith(s,"-")
        @assert parse_l(s[1:end-1]) != 0
        parse_l(s[1:end-1])
    else
        - parse_l(s) - 1
    end
end

end


text:
#=
This file is auto-generated. Do not edit.
=#
"""
    mutable struct FullMachine <: Machine
        R::Float64
        R_f::Float64
        R_1d::Float64
        R_1q::Float64
        L_d::Float64
        L_q::Float64
        L_ad::Float64
        L_aq::Float64
        L_f1d::Float64
        L_ff::Float64
        L_1d::Float64
        L_1q::Float64
        ext::Dict{String, Any}
        inv_d_fluxlink::Array{Float64,2}
        inv_q_fluxlink::Array{Float64,2}
        states::Vector{Symbol}
        n_states::Int64
        internal::InfrastructureSystemsInternal
    end

Parameter of a full order flux stator-rotor model without zero sequence flux in the stator.
 The derivative of stator fluxes (ψd and ψq) is NOT neglected. Only one q-axis damping circuit is considered. All parameters are in machine per unit.
 Refer to Chapter 3 of Power System Stability and Control by P. Kundur or Chapter 11 of Power System Dynamics: Stability and Control, by J. Machowski, J. Bialek and J. Bumby, for more details.
 Note that the models are somewhat different (but equivalent) due to the different Park Transformation used in both books.

# Arguments
- `R::Float64`: Resistance after EMF in machine per unit, validation range: (0, nothing)
- `R_f::Float64`: Field rotor winding resistance in per unit, validation range: (0, nothing)
- `R_1d::Float64`:  Damping rotor winding resistance on d-axis in per unit. This value is denoted as RD in Machowski., validation range: (0, nothing)
- `R_1q::Float64`: Damping rotor winding resistance on q-axis in per unit. This value is denoted as RQ in Machowski., validation range: (0, nothing)
- `L_d::Float64`: Inductance of fictitious damping that represent the effect of the three-phase stator winding in the d-axis of the rotor, in per unit. This value is denoted as L_ad + L_l in Kundur (and Ld in Machowski)., validation range: (0, nothing)
- `L_q::Float64`: Inductance of fictitious damping that represent the effect of the three-phase stator winding in the q-axis of the rotor, in per unit. This value is denoted as L_aq + L_l in Kundur., validation range: (0, nothing)
- `L_ad::Float64`: Mutual inductance between stator winding and rotor field (and damping) winding inductance on d-axis, in per unit, validation range: (0, nothing)
- `L_aq::Float64`: Mutual inductance between stator winding and rotor damping winding inductance on q-axis, in per unit, validation range: (0, nothing)
- `L_f1d::Float64`: Mutual inductance between rotor field winding and rotor damping winding inductance on d-axis, in per unit, validation range: (0, nothing)
- `L_ff::Float64`: Field rotor winding inductance, in per unit, validation range: (0, nothing)
- `L_1d::Float64`: Inductance of the d-axis rotor damping circuit, in per unit, validation range: (0, nothing)
- `L_1q::Float64`: Inductance of the q-axis rotor damping circuit, in per unit, validation range: (0, nothing)
- `ext::Dict{String, Any}`
- `inv_d_fluxlink::Array{Float64,2}`: Equations 3.127, 3.130, 3.131 From Kundur
- `inv_q_fluxlink::Array{Float64,2}`: Equations 3.128, 3.132 From Kundur
- `states::Vector{Symbol}`
- `n_states::Int64`
- `internal::InfrastructureSystemsInternal`: power system internal reference, do not modify
"""
mutable struct FullMachine <: Machine
    "Resistance after EMF in machine per unit"
    R::Float64
    "Field rotor winding resistance in per unit"
    R_f::Float64
    " Damping rotor winding resistance on d-axis in per unit. This value is denoted as RD in Machowski."
    R_1d::Float64
    "Damping rotor winding resistance on q-axis in per unit. This value is denoted as RQ in Machowski."
    R_1q::Float64
    "Inductance of fictitious damping that represent the effect of the three-phase stator winding in the d-axis of the rotor, in per unit. This value is denoted as L_ad + L_l in Kundur (and Ld in Machowski)."
    L_d::Float64
    "Inductance of fictitious damping that represent the effect of the three-phase stator winding in the q-axis of the rotor, in per unit. This value is denoted as L_aq + L_l in Kundur."
    L_q::Float64
    "Mutual inductance between stator winding and rotor field (and damping) winding inductance on d-axis, in per unit"
    L_ad::Float64
    "Mutual inductance between stator winding and rotor damping winding inductance on q-axis, in per unit"
    L_aq::Float64
    "Mutual inductance between rotor field winding and rotor damping winding inductance on d-axis, in per unit"
    L_f1d::Float64
    "Field rotor winding inductance, in per unit"
    L_ff::Float64
    "Inductance of the d-axis rotor damping circuit, in per unit"
    L_1d::Float64
    "Inductance of the q-axis rotor damping circuit, in per unit"
    L_1q::Float64
    ext::Dict{String, Any}
    "Equations 3.127, 3.130, 3.131 From Kundur"
    inv_d_fluxlink::Array{Float64,2}
    "Equations 3.128, 3.132 From Kundur"
    inv_q_fluxlink::Array{Float64,2}
    states::Vector{Symbol}
    n_states::Int64
    "power system internal reference, do not modify"
    internal::InfrastructureSystemsInternal
end

function FullMachine(R, R_f, R_1d, R_1q, L_d, L_q, L_ad, L_aq, L_f1d, L_ff, L_1d, L_1q, ext=Dict{String, Any}(), )
    FullMachine(R, R_f, R_1d, R_1q, L_d, L_q, L_ad, L_aq, L_f1d, L_ff, L_1d, L_1q, ext, inv([[-L_d L_ad L_ad]; [-L_ad L_ff L_f1d]; [-L_ad L_f1d L_1d]]), inv([[-L_q L_aq]; [-L_aq L_1q]]), [:ψd, :ψq, :ψf, :ψ1d, :ψ1q], 5, InfrastructureSystemsInternal(), )
end

function FullMachine(; R, R_f, R_1d, R_1q, L_d, L_q, L_ad, L_aq, L_f1d, L_ff, L_1d, L_1q, ext=Dict{String, Any}(), )
    FullMachine(R, R_f, R_1d, R_1q, L_d, L_q, L_ad, L_aq, L_f1d, L_ff, L_1d, L_1q, ext, )
end

# Constructor for demo purposes; non-functional.
function FullMachine(::Nothing)
    FullMachine(;
        R=0,
        R_f=0,
        R_1d=0,
        R_1q=0,
        L_d=1.0,
        L_q=1.0,
        L_ad=2.0,
        L_aq=2.0,
        L_f1d=1.0,
        L_ff=2.0,
        L_1d=1.0,
        L_1q=1.0,
        ext=Dict{String, Any}(),
    )
end

"""Get FullMachine R."""
get_R(value::FullMachine) = value.R
"""Get FullMachine R_f."""
get_R_f(value::FullMachine) = value.R_f
"""Get FullMachine R_1d."""
get_R_1d(value::FullMachine) = value.R_1d
"""Get FullMachine R_1q."""
get_R_1q(value::FullMachine) = value.R_1q
"""Get FullMachine L_d."""
get_L_d(value::FullMachine) = value.L_d
"""Get FullMachine L_q."""
get_L_q(value::FullMachine) = value.L_q
"""Get FullMachine L_ad."""
get_L_ad(value::FullMachine) = value.L_ad
"""Get FullMachine L_aq."""
get_L_aq(value::FullMachine) = value.L_aq
"""Get FullMachine L_f1d."""
get_L_f1d(value::FullMachine) = value.L_f1d
"""Get FullMachine L_ff."""
get_L_ff(value::FullMachine) = value.L_ff
"""Get FullMachine L_1d."""
get_L_1d(value::FullMachine) = value.L_1d
"""Get FullMachine L_1q."""
get_L_1q(value::FullMachine) = value.L_1q
"""Get FullMachine ext."""
get_ext(value::FullMachine) = value.ext
"""Get FullMachine inv_d_fluxlink."""
get_inv_d_fluxlink(value::FullMachine) = value.inv_d_fluxlink
"""Get FullMachine inv_q_fluxlink."""
get_inv_q_fluxlink(value::FullMachine) = value.inv_q_fluxlink
"""Get FullMachine states."""
get_states(value::FullMachine) = value.states
"""Get FullMachine n_states."""
get_n_states(value::FullMachine) = value.n_states
"""Get FullMachine internal."""
get_internal(value::FullMachine) = value.internal

"""Set FullMachine R."""
set_R!(value::FullMachine, val::Float64) = value.R = val
"""Set FullMachine R_f."""
set_R_f!(value::FullMachine, val::Float64) = value.R_f = val
"""Set FullMachine R_1d."""
set_R_1d!(value::FullMachine, val::Float64) = value.R_1d = val
"""Set FullMachine R_1q."""
set_R_1q!(value::FullMachine, val::Float64) = value.R_1q = val
"""Set FullMachine L_d."""
set_L_d!(value::FullMachine, val::Float64) = value.L_d = val
"""Set FullMachine L_q."""
set_L_q!(value::FullMachine, val::Float64) = value.L_q = val
"""Set FullMachine L_ad."""
set_L_ad!(value::FullMachine, val::Float64) = value.L_ad = val
"""Set FullMachine L_aq."""
set_L_aq!(value::FullMachine, val::Float64) = value.L_aq = val
"""Set FullMachine L_f1d."""
set_L_f1d!(value::FullMachine, val::Float64) = value.L_f1d = val
"""Set FullMachine L_ff."""
set_L_ff!(value::FullMachine, val::Float64) = value.L_ff = val
"""Set FullMachine L_1d."""
set_L_1d!(value::FullMachine, val::Float64) = value.L_1d = val
"""Set FullMachine L_1q."""
set_L_1q!(value::FullMachine, val::Float64) = value.L_1q = val
"""Set FullMachine ext."""
set_ext!(value::FullMachine, val::Dict{String, Any}) = value.ext = val
"""Set FullMachine inv_d_fluxlink."""
set_inv_d_fluxlink!(value::FullMachine, val::Array{Float64,2}) = value.inv_d_fluxlink = val
"""Set FullMachine inv_q_fluxlink."""
set_inv_q_fluxlink!(value::FullMachine, val::Array{Float64,2}) = value.inv_q_fluxlink = val
"""Set FullMachine states."""
set_states!(value::FullMachine, val::Vector{Symbol}) = value.states = val
"""Set FullMachine n_states."""
set_n_states!(value::FullMachine, val::Int64) = value.n_states = val
"""Set FullMachine internal."""
set_internal!(value::FullMachine, val::InfrastructureSystemsInternal) = value.internal = val


text:
"""
DuckDB data chunk
"""
mutable struct DataChunk
    handle::duckdb_data_chunk

    function DataChunk(handle::duckdb_data_chunk, destroy::Bool)
        result = new(handle)
        if destroy
            finalizer(_destroy_data_chunk, result)
        end
        return result
    end
end

function GetColumnCount(chunk::DataChunk)
    return duckdb_data_chunk_get_column_count(chunk.handle)
end

function GetSize(chunk::DataChunk)
    return duckdb_data_chunk_get_size(chunk.handle)
end

function SetSize(chunk::DataChunk, size::Int64)
    return duckdb_data_chunk_set_size(chunk.handle, size)
end

function GetVector(chunk::DataChunk, col_idx::Int64)::Vec
    if col_idx < 1 || col_idx > GetColumnCount(chunk)
        throw(
            InvalidInputException(
                string(
                    "GetArray column index ",
                    col_idx,
                    " out of range, expected value between 1 and ",
                    GetColumnCount(chunk)
                )
            )
        )
    end
    return Vec(duckdb_data_chunk_get_vector(chunk.handle, col_idx))
end

function GetArray(chunk::DataChunk, col_idx::Int64, ::Type{T})::Vector{T} where {T}
    return GetArray(GetVector(chunk, col_idx), T)
end

function GetValidity(chunk::DataChunk, col_idx::Int64)::ValidityMask
    return GetValidity(GetVector(chunk, col_idx))
end

function AllValid(chunk::DataChunk, col_idx::Int64)
    return AllValid(GetVector(chunk, col_idx))
end

# this is only required when we own the data chunk
function _destroy_data_chunk(chunk::DataChunk)
    if chunk.handle != C_NULL
        duckdb_destroy_data_chunk(chunk.handle)
    end
    return chunk.handle = C_NULL
end


text:
export
    run_simulation,
    simulate_encounters,
    simulate_encounter

#=
The main script that simulates encounters.
=#

#=
Description:
This method simulates a single encounter. Abstractly, this
entails initializing an encounter, simulating
each uav in the environment taking actions, and then
simulating the trajectories resulting from these actions.

Parameters:
- environment: the environment in which to simulate the encounter

Return Value:
- encounter: a single encounter object

=#
function simulate_encounter(environment::Environment)
    # select a random start state
    state = start_encounter(environment)

    # create an encounter to track information
    encounter = Encounter()
    update!(encounter, state)

    # simulate the actual encounter
    for idx in 1:environment.max_steps
        actions = get_uav_actions!(environment, state, encounter)
        state = step!(environment, state, actions, encounter)
    end

    return encounter
end

#=
Description:
This method simulates a variable number of encounters using the
environment provided to it.

Parameters:
- environment: the environment in which to simulate encounters
- num_encounters: the number of encounters to simulate

Return Value:
- encounters: a list of encounter objects which contain all the
                information needed to visualize each encounter
=#
function simulate_encounters(environment::Environment, num_encounters::Int64)

    # simulate num_encounters, collecting encounter information
    encounters = Array(Encounter, num_encounters)
    for idx in 1:num_encounters
        encounters[idx] = simulate_encounter(environment)
    end
    return encounters
end

#=
Description:
Main method that performs setup and then calls simulate_encounters
to actually simulate encounters.
=#
function run_simulation(;num_encounters = 1)
    # get the environment to simulate encounters in
    environment = build_environment()

    # simulate the encounters
    encounters = simulate_encounters(environment, num_encounters)

end


text:
@testset "ruleset_loading.jl" begin
    @testset "on_new_rule" begin
        frule_history = []
        rrule_history = []
        on_new_rule(frule) do sig
            op = sig.parameters[1]
            push!(frule_history, op)
        end
        on_new_rule(rrule) do sig
            op = sig.parameters[1]
            push!(rrule_history, op)
        end

        @testset "new rules hit the hooks" begin
            # Now define some rules
            @scalar_rule x + y (1, 1)
            @scalar_rule x - y (1, -1)
            refresh_rules()

            @test Set(frule_history[end-1:end]) == Set((typeof(+), typeof(-)))
            @test Set(rrule_history[end-1:end]) == Set((typeof(+), typeof(-)))
        end

        @testset "# Make sure nothing happens anymore once we clear the hooks" begin
            ChainRulesCore.clear_new_rule_hooks!(frule)
            ChainRulesCore.clear_new_rule_hooks!(rrule)

            old_frule_history = copy(frule_history)
            old_rrule_history = copy(rrule_history)

            @scalar_rule sin(x) cos(x)
            refresh_rules()

            @test old_rrule_history == rrule_history
            @test old_frule_history == frule_history
        end

    end

    @testset "_primal_sig" begin
        _primal_sig = ChainRulesCore._primal_sig
        @testset "frule" begin
            @test isequal(  # DataType without shared type but with constraint
                _primal_sig(frule, Tuple{typeof(frule), Any, typeof(*), Int, Vector{Int}}),
                Tuple{typeof(*), Int, Vector{Int}}
            )
            @test isequal(  # UnionAall without shared type but with constraint
                _primal_sig(frule, Tuple{typeof(frule), Any, typeof(*),  T, Int} where T<:Real),
                Tuple{typeof(*), T, Int} where T<:Real
            )
            @test isequal(  # UnionAall with share type
                _primal_sig(frule, Tuple{typeof(frule), Any, typeof(*), T, Vector{T}} where T),
                Tuple{typeof(*), T, Vector{T}} where T
            )
        end

        @testset "rrule" begin
            @test isequal(  # DataType without shared type but with constraint
                _primal_sig(rrule, Tuple{typeof(rrule), typeof(*), Int, Vector{Int}}),
                Tuple{typeof(*), Int, Vector{Int}}
            )
            @test isequal(  # UnionAall without shared type but with constraint
                _primal_sig(rrule, Tuple{typeof(rrule), typeof(*),  T, Int} where T<:Real),
                Tuple{typeof(*), T, Int} where T<:Real
            )
            @test isequal(  # UnionAall with share type
                _primal_sig(rrule, Tuple{typeof(rrule), typeof(*), T, Vector{T}} where T),
                Tuple{typeof(*), T, Vector{T}} where T
            )
        end
    end
end


text:
import Base: +, -, *, /, length, iterate, zero, convert, Vector, Array
"""
This correspond to coordinate of atom.
This contain 3 fields like below.

- x <: Real
- y <: Real
- z <: Real

Some operators and functions were overloaded to make Coordinate objects broadcastable.
"""
mutable struct Coordinate{RealT <: Real}
    # This class is for make vector operation broadcastable.
    x::RealT
    y::RealT
    z::RealT

    function Coordinate(vector::VectorT
        ) where VectorT <: AbstractArray{RealT} where RealT <: Real

        if length(vector) != 3
            ArgumentError("Argument for Coordinate constructor must be array of length 3")
        end
        new{RealT}(vector[1], vector[2], vector[3])
    end

    function Coordinate(x::RealT, y::RealT, z::RealT) where RealT <: Real
        new{RealT}(x, y, z)
    end
end

function +(arg1::Coordinate{T}, arg2::Coordinate{U}) where T <: Real where U <: Real
    Coordinate(arg1.x + arg2.x, arg1.y + arg2.y, arg1.z + arg2.z)
end

function -(arg1::Coordinate{T}, arg2::Coordinate{U}) where T <: Real where U <: Real
    Coordinate(arg1.x - arg2.x, arg1.y - arg2.y, arg1.z - arg2.z)
end

function *(arg1::Coordinate{T}, arg2::U) where T <: Real where U <: Real
    Coordinate(arg1.x * arg2, arg1.y * arg2, arg1.z * arg2)
end

function *(arg1::U, arg2::Coordinate{T}) where T <: Real where U <: Real
    Coordinate(arg1 * arg2.x, arg1 * arg2.y, arg1 * arg2.z)
end

function *(arg1::Coordinate{T}, arg2::Coordinate{U}) where T <: Real where U <: Real
    arg1.x * arg2.x + arg1.y * arg2.y + arg1.z * arg2.z
end

function /(arg1::Coordinate{T}, arg2::U) where T <: Real where U <: Real
    Coordinate(arg1.x / arg2, arg1.y / arg2, arg1.z / arg2)
end

# for operator broadcast
function length(arg::Coordinate{T}) where T <: Real
    1
end

function iterate(arg::Coordinate{T}) where T <: Real
    (arg, nothing)
end

function iterate(arg::Coordinate{T}, nothing) where T <: Real
    nothing
end

function zero(arg::Coordinate{T}) where T <: Real
    Coordinate(T(0.0), T(0.0), T(0.0))
end

function convert(::Type{Coordinate{RealT}}, coord::Coordinate{<:Real}) where RealT <: Real
    return Coordinate(RealT(coord.x), RealT(coord.y), RealT(coord.z))
end

function Vector(arg::Coordinate{RealT}) where RealT <: Real
    [arg.x, arg.y, arg.z]
end

function Array(arg::Coordinate{RealT}) where RealT <: Real
    [arg.x, arg.y, arg.z]
end

function norm(arg::Coordinate{T}) where T <: Real
    sqrt(abs2(arg.x) + abs2(arg.y) + abs2(arg.z))
end

function distance(first_coord::Coordinate{T}, second_coord::Coordinate{T}) where T <: Real
    norm(first_coord - second_coord)
end


text:
module Centrality

using ..AbstractStaticGraph, ..AbstractEvolvingGraph

import ..num_nodes, ..nodes, ..node_index, ..timestamps, ..unique_timestamps, ..sparse_adjacency_matrix, ..adjacency_matrix

export katz, betweenness, closeness

include("katz.jl")
include("betweenness.jl")
include("closeness.jl")

end


text:
println("# HDP TESTS...")

# run demo code at first
println(" * try demo...")

# define size of bar image
img_size = 5

# define amount of noise
noise_level = 0.01

# define probabilities of generating a particular number of bars
num_per_mixture = [0 ones(1,3)] ./ 3

# define number of groups (J)
num_group = 40

# define number of data items drawn from each group
num_data = 50;

# generate data
(X, bars) = generateBarsDataset(img_size, noise_level, num_per_mixture, num_group, num_data);

# Dimensionality of the data
D = img_size * img_size

# We assume a Multinomial Distribution with a Dirichlet Prior as base distribution
H = MultinomialDirichlet(D, 1.0)

# train a Hierarchical Dirichlet Process Mixture Model
# guessing 10 shared Distributions
models = train(HDP(H), Gibbs(), RandomInitialisation(k = 10), X);

# access the log likelihood values of all iteration
LLH = map(model -> model.energy, models)

# access the weights of all iterations
W = map(model -> model.weights, models)

# access the distributions of all iterations
D = map(model -> model.distributions, models)

println(" * test randnumtable()")
weights = ones(10, 10)
table = ones(Int, 10, 10)

@test weights == BNP.randnumtable(weights, table)

println("# FINISHED")


text:
@testset "1207.unique-number-of-occurrences.jl" begin
    @test unique_number_of_occurrences([1,2,2,1,1,3])
    @test !(unique_number_of_occurrences([1,2]))
    @test unique_number_of_occurrences([-3,0,1,-3,1,1,1,-3,10,0])
end


text:
module PiecewiseAffineTransforms

export PAWarpParams,
       pa_warp_params,
       pa_warp,
       delaunayindexes,
       ## probably these fucntions should not reside here,
       ## but I couldn't find better place for them
       inpolygon,
       fillpoly!,
       poly2mask

include("core.jl")

end


text:
# Julia Mad Libs game adapted from: https://github.com/logankilpatrick/10-Julia-Projects-for-Beginners
# I added some small randomness to it by providing 2 other sentences.
# Randomness comes from the random seed based on the user inputs.
using Random
function play_mad_libs()
	print("Enter a verb (action): ")
	verb1 = readline()
	print("Enter an adjective (descriptive word): ")
	adj1 = readline()
	print("Enter a noun (person, place, or thing): ")
	noun1 = readline()
	print("Enter another noun (person, place, or thing): ")
	noun2 = readline()
	print("Enter a catchphrase (something like 'hands up!'): ")
	phrase1 = readline()
	+(a::String,b::String,c::String,x::String,y::String) = string(a,b,c,x,y)
	i = length(verb1+adj1+noun1+noun2+phrase1)
	Random.seed!(i)
	j = rand(1:3)
	if j == 1
		base_sentence = "John $verb1 down the street one night, playing with his $adj1 $noun1. When all of a
		sudden, a $noun2 jumped out at him and said '$phrase1'"
		elseif j == 2
		base_sentence = "Cooter $verb1 in the woods on a $adj1 day with his $noun1. On the hike, they ran into
		$noun2. Cooter then yelled $phrase1 at the bear and ran with his $noun1."
		else
		base_sentence = "A $adj1 $noun1 $verb1 into the lake. Then, a $noun2 came upon the lake and yelled $phrase1"
		end
	print("\n\n", base_sentence)
	end
play_mad_libs()

max_stars_repo_path: Mad libs with Randomness.jl

max_stars_repo_name: slatteew/Julia

max_issues_repo_path: Mad libs with Randomness.jl

max_issues_repo_name: slatteew/Julia

max_forks_repo_path: Mad libs with Randomness.jl

max_forks_repo_name: slatteew/Julia


text:
@doc raw"""
    GCNConv(in => out, σ=identity; bias=true, init=glorot_uniform, add_self_loops=true)

Graph convolutional layer from paper [Semi-supervised Classification with Graph Convolutional Networks](https://arxiv.org/abs/1609.02907).

Performs the operation
```math
\mathbf{x}'_i = \sum_{j\in N(i)} \frac{1}{c_{ij}} W \mathbf{x}_j
```
where ``c_{ij} = \sqrt{|N(i)||N(j)|}``.

The input to the layer is a node feature array `X`
of size `(num_features, num_nodes)`.

# Arguments

- `in`: Number of input features.
- `out`: Number of output features.
- `σ`: Activation function.
- `bias`: Add learnable bias.
- `init`: Weights' initializer.
- `add_self_loops`: Add self loops to the graph before performing the convolution.
"""
struct GCNConv{A<:AbstractMatrix, B, F} <: GNNLayer
    weight::A
    bias::B
    σ::F
    add_self_loops::Bool
end

@functor GCNConv

function GCNConv(ch::Pair{Int,Int}, σ=identity;
                 init=glorot_uniform, bias::Bool=true,
                 add_self_loops=true)
    in, out = ch
    W = init(out, in)
    b = bias ? Flux.create_bias(W, true, out) : false
    GCNConv(W, b, σ, add_self_loops)
end

## Matrix operations are more performant,
## but cannot compute the normalized adjacency of sparse cuda matrices yet,
## therefore fallback to message passing framework on gpu for the time being

function (l::GCNConv)(g::GNNGraph, x::AbstractMatrix{T}) where T
    Ã = normalized_adjacency(g, T; dir=:out, l.add_self_loops)
    l.σ.(l.weight * x * Ã .+ l.bias)
end

compute_message(l::GCNConv, xi, xj, eij) = xj
update_node(l::GCNConv, m, x) = m

function (l::GCNConv)(g::GNNGraph, x::CuMatrix{T}) where T
    if l.add_self_loops
        g = add_self_loops(g)
    end
    c = 1 ./ sqrt.(degree(g, T, dir=:in))
    x = x .* c'
    x, _ = propagate(l, g, +, x)
    x = x .* c'
    return l.σ.(l.weight * x .+ l.bias)
end

function Base.show(io::IO, l::GCNConv)
    out, in = size(l.weight)
    print(io, "GCNConv($in => $out")
    l.σ == identity || print(io, ", ", l.σ)
    print(io, ")")
end

@doc raw"""
    ChebConv(in => out, k; bias=true, init=glorot_uniform)

Chebyshev spectral graph convolutional layer from
paper [Convolutional Neural Networks on Graphs with Fast Localized Spectral Filtering](https://arxiv.org/abs/1606.09375).

Implements

```math
X' = \sum^{K-1}_{k=0}  W^{(k)} Z^{(k)}
```

where ``Z^{(k)}`` is the ``k``-th term of Chebyshev polynomials, and can be calculated by the following recursive form:

```math
Z^{(0)} = X \\
Z^{(1)} = \hat{L} X \\
Z^{(k)} = 2 \hat{L} Z^{(k-1)} - Z^{(k-2)}
```

with ``\hat{L}`` the [`scaled_laplacian`](@ref).

# Arguments

- `in`: The dimension of input features.
- `out`: The dimension of output features.
- `k`: The order of Chebyshev polynomial.
- `bias`: Add learnable bias.
- `init`: Weights' initializer.
"""
struct ChebConv{A<:AbstractArray{<:Number,3}, B} <: GNNLayer
    weight::A
    bias::B
    k::Int
end

function ChebConv(ch::Pair{Int,Int}, k::Int;
                  init=glorot_uniform, bias::Bool=true)
    in, out = ch
    W = init(out, in, k)
    b = bias ? Flux.create_bias(W, true, out) : false
    ChebConv(W, b, k)
end

@functor ChebConv

function (c::ChebConv)(g::GNNGraph, X::AbstractMatrix{T}) where T
    check_num_nodes(g, X)
    @assert size(X, 1) == size(c.weight, 2) "Input feature size must match input channel size."

    L̃ = scaled_laplacian(g, eltype(X))

    Z_prev = X
    Z = X * L̃
    Y = view(c.weight,:,:,1) * Z_prev
    Y += view(c.weight,:,:,2) * Z
    for k = 3:c.k
        Z, Z_prev = 2*Z*L̃ - Z_prev, Z
        Y += view(c.weight,:,:,k) * Z
    end
    return Y .+ c.bias
end

function Base.show(io::IO, l::ChebConv)
    out, in, k = size(l.weight)
    print(io, "ChebConv(", in, " => ", out)
    print(io, ", k=", k)
    print(io, ")")
end

@doc raw"""
    GraphConv(in => out, σ=identity; aggr=+, bias=true, init=glorot_uniform)

Graph convolution layer from Reference: [Weisfeiler and Leman Go Neural: Higher-order Graph Neural Networks](https://arxiv.org/abs/1810.02244).

Performs:
```math
\mathbf{x}_i' = W_1 \mathbf{x}_i + \square_{j \in \mathcal{N}(i)} W_2 \mathbf{x}_j
```

where the aggregation type is selected by `aggr`.

# Arguments

- `in`: The dimension of input features.
- `out`: The dimension of output features.
- `σ`: Activation function.
- `aggr`: Aggregation operator for the incoming messages (e.g. `+`, `*`, `max`, `min`, and `mean`).
- `bias`: Add learnable bias.
- `init`: Weights' initializer.
"""
struct GraphConv{A<:AbstractMatrix, B} <: GNNLayer
    weight1::A
    weight2::A
    bias::B
    σ
    aggr
end

@functor GraphConv

function GraphConv(ch::Pair{Int,Int}, σ=identity; aggr=+,
                   init=glorot_uniform, bias::Bool=true)
    in, out = ch
    W1 = init(out, in)
    W2 = init(out, in)
    b = bias ? Flux.create_bias(W1, true, out) : false
    GraphConv(W1, W2, b, σ, aggr)
end

compute_message(l::GraphConv, x_i, x_j, e_ij) =  x_j
update_node(l::GraphConv, m, x) = l.σ.(l.weight1 * x .+ l.weight2 * m .+ l.bias)

function (l::GraphConv)(g::GNNGraph, x::AbstractMatrix)
    check_num_nodes(g, x)
    x, _ = propagate(l, g, l.aggr, x)
    x
end

function Base.show(io::IO, l::GraphConv)
    in_channel = size(l.weight1, ndims(l.weight1))
    out_channel = size(l.weight1, ndims(l.weight1)-1)
    print(io, "GraphConv(", in_channel, " => ", out_channel)
    l.σ == identity || print(io, ", ", l.σ)
    print(io, ", aggr=", l.aggr)
    print(io, ")")
end

@doc raw"""
    GATConv(in => out, σ=identity;
            heads=1,
            concat=true,
            init=glorot_uniform
            bias=true,
            negative_slope=0.2f0)

Graph attentional layer from the paper [Graph Attention Networks](https://arxiv.org/abs/1710.10903).

Implements the operation
```math
\mathbf{x}_i' = \sum_{j \in N(i)} \alpha_{ij} W \mathbf{x}_j
```
where the attention coefficients ``\alpha_{ij}`` are given by
```math
\alpha_{ij} = \frac{1}{z_i} \exp(LeakyReLU(\mathbf{a}^T [W \mathbf{x}_i \,\|\, W \mathbf{x}_j]))
```
with ``z_i`` a normalization factor.

# Arguments

- `in`: The dimension of input features.
- `out`: The dimension of output features.
- `bias`: Learn the additive bias if true.
- `heads`: Number attention heads.
- `concat`: Concatenate layer output or not. If not, layer output is averaged over the heads.
- `negative_slope`: The parameter of LeakyReLU.
"""
struct GATConv{T, A<:AbstractMatrix, B} <: GNNLayer
    weight::A
    bias::B
    a::A
    σ
    negative_slope::T
    channel::Pair{Int, Int}
    heads::Int
    concat::Bool
end

@functor GATConv
Flux.trainable(l::GATConv) = (l.weight, l.bias, l.a)

function GATConv(ch::Pair{Int,Int}, σ=identity;
                 heads::Int=1, concat::Bool=true, negative_slope=0.2,
                 init=glorot_uniform, bias::Bool=true)
    in, out = ch
    W = init(out*heads, in)
    if concat
        b = bias ? Flux.create_bias(W, true, out*heads) : false
    else
        b = bias ? Flux.create_bias(W, true, out) : false
    end
    a = init(2*out, heads)
    negative_slope = convert(eltype(W), negative_slope)
    GATConv(W, b, a, σ, negative_slope, ch, heads, concat)
end

function compute_message(l::GATConv, Wxi, Wxj)
    aWW = sum(l.a .* vcat(Wxi, Wxj), dims=1)   # 1 × nheads × nedges
    α = exp.(leakyrelu.(aWW, l.negative_slope))
    return (α = α, m = α .* Wxj)
end

update_node(l::GATConv, d̄, x) = d̄.m ./ d̄.α

function (l::GATConv)(g::GNNGraph, x::AbstractMatrix)
    check_num_nodes(g, x)
    g = add_self_loops(g)
    chin, chout = l.channel
    heads = l.heads

    Wx = l.weight * x
    Wx = reshape(Wx, chout, heads, :)                   # chout × nheads × nnodes

    x, _ = propagate(l, g, +, Wx)                 ## chout × nheads × nnodes

    if !l.concat
        x = mean(x, dims=2)
    end
    x = reshape(x, :, size(x, 3))  # return a matrix
    x = l.σ.(x .+ l.bias)

    return x
end

function Base.show(io::IO, l::GATConv)
    out_channel, in_channel = size(l.weight)
    print(io, "GATConv(", in_channel, "=>", out_channel ÷ l.heads)
    print(io, ", LeakyReLU(λ=", l.negative_slope)
    print(io, "))")
end

@doc raw"""
    GatedGraphConv(out, num_layers; aggr=+, init=glorot_uniform)

Gated graph convolution layer from [Gated Graph Sequence Neural Networks](https://arxiv.org/abs/1511.05493).

Implements the recursion
```math
\mathbf{h}^{(0)}_i = \mathbf{x}_i || \mathbf{0} \\
\mathbf{h}^{(l)}_i = GRU(\mathbf{h}^{(l-1)}_i, \square_{j \in N(i)} W \mathbf{h}^{(l-1)}_j)
```

 where ``\mathbf{h}^{(l)}_i`` denotes the ``l``-th hidden variables passing through GRU. The dimension of input ``\mathbf{x}_i`` needs to be less or equal to `out`.

# Arguments

- `out`: The dimension of output features.
- `num_layers`: The number of gated recurrent unit.
- `aggr`: Aggregation operator for the incoming messages (e.g. `+`, `*`, `max`, `min`, and `mean`).
- `init`: Weight initialization function.
"""
struct GatedGraphConv{A<:AbstractArray{<:Number,3}, R} <: GNNLayer
    weight::A
    gru::R
    out_ch::Int
    num_layers::Int
    aggr
end

@functor GatedGraphConv

function GatedGraphConv(out_ch::Int, num_layers::Int;
                        aggr=+, init=glorot_uniform)
    w = init(out_ch, out_ch, num_layers)
    gru = GRUCell(out_ch, out_ch)
    GatedGraphConv(w, gru, out_ch, num_layers, aggr)
end

compute_message(l::GatedGraphConv, x_i, x_j, e_ij) = x_j
update_node(l::GatedGraphConv, m, x) = m

# remove after https://github.com/JuliaDiff/ChainRules.jl/pull/521
@non_differentiable fill!(x...)

function (l::GatedGraphConv)(g::GNNGraph, H::AbstractMatrix{S}) where {S<:Real}
    check_num_nodes(g, H)
    m, n = size(H)
    @assert (m <= l.out_ch) "number of input features must less or equals to output features."
    if m < l.out_ch
        Hpad = similar(H, S, l.out_ch - m, n)
        H = vcat(H, fill!(Hpad, 0))
    end
    for i = 1:l.num_layers
        M = view(l.weight, :, :, i) * H
        M, _ = propagate(l, g, l.aggr, M)
        H, _ = l.gru(H, M)
    end
    H
end

function Base.show(io::IO, l::GatedGraphConv)
    print(io, "GatedGraphConv(($(l.out_ch) => $(l.out_ch))^$(l.num_layers)")
    print(io, ", aggr=", l.aggr)
    print(io, ")")
end

@doc raw"""
    EdgeConv(nn; aggr=max)

Edge convolutional layer from paper [Dynamic Graph CNN for Learning on Point Clouds](https://arxiv.org/abs/1801.07829).

Performs the operation
```math
\mathbf{x}_i' = \square_{j \in N(i)} nn(\mathbf{x}_i || \mathbf{x}_j - \mathbf{x}_i)
```

where `nn` generally denotes a learnable function, e.g. a linear layer or a multi-layer perceptron.

# Arguments

- `nn`: A (possibly learnable) function acting on edge features.
- `aggr`: Aggregation operator for the incoming messages (e.g. `+`, `*`, `max`, `min`, and `mean`).
"""
struct EdgeConv <: GNNLayer
    nn
    aggr
end

@functor EdgeConv

EdgeConv(nn; aggr=max) = EdgeConv(nn, aggr)

compute_message(l::EdgeConv, x_i, x_j, e_ij) = l.nn(vcat(x_i, x_j .- x_i))

update_node(l::EdgeConv, m, x) = m

function (l::EdgeConv)(g::GNNGraph, X::AbstractMatrix)
    check_num_nodes(g, X)
    X, _ = propagate(l, g, l.aggr, X)
    X
end

function Base.show(io::IO, l::EdgeConv)
    print(io, "EdgeConv(", l.nn)
    print(io, ", aggr=", l.aggr)
    print(io, ")")
end

@doc raw"""
    GINConv(f, ϵ; aggr=+)

Graph Isomorphism convolutional layer from paper [How Powerful are Graph Neural Networks?](https://arxiv.org/pdf/1810.00826.pdf)

```math
\mathbf{x}_i' = f_\Theta\left((1 + \epsilon) \mathbf{x}_i + \sum_{j \in N(i)} \mathbf{x}_j \right)
```
where ``f_\Theta`` typically denotes a learnable function, e.g. a linear layer or a multi-layer perceptron.

# Arguments

- `f`: A (possibly learnable) function acting on node features.
- `ϵ`: Weighting factor.
"""
struct GINConv{R<:Real} <: GNNLayer
    nn
    ϵ::R
    aggr
end

@functor GINConv
Flux.trainable(l::GINConv) = (l.nn,)

GINConv(nn, ϵ; aggr=+) = GINConv(nn, ϵ, aggr)

compute_message(l::GINConv, x_i, x_j, e_ij) = x_j
update_node(l::GINConv, m, x) = l.nn((1 + ofeltype(x, l.ϵ)) * x + m)

function (l::GINConv)(g::GNNGraph, X::AbstractMatrix)
    check_num_nodes(g, X)
    X, _ = propagate(l, g, l.aggr, X)
    X
end

function Base.show(io::IO, l::GINConv)
    print(io, "GINConv($(l.nn)")
    print(io, ", $(l.ϵ)")
    print(io, ")")
end

@doc raw"""
    NNConv(in => out, f, σ=identity; aggr=+, bias=true, init=glorot_uniform)

The continuous kernel-based convolutional operator from the
[Neural Message Passing for Quantum Chemistry](https://arxiv.org/abs/1704.01212) paper.
This convolution is also known as the edge-conditioned convolution from the
[Dynamic Edge-Conditioned Filters in Convolutional Neural Networks on Graphs](https://arxiv.org/abs/1704.02901) paper.

Performs the operation

```math
\mathbf{x}_i' = W \mathbf{x}_i + \square_{j \in N(i)} f_\Theta(\mathbf{e}_{j\to i})\,\mathbf{x}_j
```

where ``f_\Theta``  denotes a learnable function (e.g. a linear layer or a multi-layer perceptron).
Given an input of batched edge features `e` of size `(num_edge_features, num_edges)`,
the function `f` will return an batched matrices array whose size is `(out, in, num_edges)`.
For convenience, also functions returning a single `(out*in, num_edges)` matrix are allowed.

# Arguments

- `in`: The dimension of input features.
- `out`: The dimension of output features.
- `f`: A (possibly learnable) function acting on edge features.
- `aggr`: Aggregation operator for the incoming messages (e.g. `+`, `*`, `max`, `min`, and `mean`).
- `σ`: Activation function.
- `bias`: Add learnable bias.
- `init`: Weights' initializer.
"""
struct NNConv <: GNNLayer
    weight
    bias
    nn
    σ
    aggr
end

@functor NNConv

function NNConv(ch::Pair{Int,Int}, nn, σ=identity; aggr=+, bias=true, init=glorot_uniform)
    in, out = ch
    W = init(out, in)
    b = bias ? Flux.create_bias(W, true, out) : false
    return NNConv(W, b, nn, σ, aggr)
end

function compute_message(l::NNConv, x_i, x_j, e_ij)
    nin, nedges = size(x_i)
    W = reshape(l.nn(e_ij), (:, nin, nedges))
    x_j = reshape(x_j, (nin, 1, nedges)) # needed by batched_mul
    m = NNlib.batched_mul(W, x_j)
    return reshape(m, :, nedges)
end

function update_node(l::NNConv, m, x)
    l.σ.(l.weight*x .+ m .+ l.bias)
end

function (l::NNConv)(g::GNNGraph, x::AbstractMatrix, e)
    check_num_nodes(g, x)
    x, _ = propagate(l, g, l.aggr, x, e)
    return x
end

(l::NNConv)(g::GNNGraph) = GNNGraph(g, ndata=l(g, node_features(g), edge_features(g)))

function Base.show(io::IO, l::NNConv)
    out, in = size(l.weight)
    print(io, "NNConv( $in => $out")
    print(io, ", aggr=", l.aggr)
    print(io, ")")
end

@doc raw"""
    SAGEConv(in => out, σ=identity; aggr=mean, bias=true, init=glorot_uniform)

GraphSAGE convolution layer from paper [Inductive Representation Learning on Large Graphs](https://arxiv.org/pdf/1706.02216.pdf).

Performs:
```math
\mathbf{x}_i' = W [\mathbf{x}_i \,\|\, \square_{j \in \mathcal{N}(i)} \mathbf{x}_j]
```

where the aggregation type is selected by `aggr`.

# Arguments

- `in`: The dimension of input features.
- `out`: The dimension of output features.
- `σ`: Activation function.
- `aggr`: Aggregation operator for the incoming messages (e.g. `+`, `*`, `max`, `min`, and `mean`).
- `bias`: Add learnable bias.
- `init`: Weights' initializer.
"""
struct SAGEConv{A<:AbstractMatrix, B} <: GNNLayer
    weight::A
    bias::B
    σ
    aggr
end

@functor SAGEConv

function SAGEConv(ch::Pair{Int,Int}, σ=identity; aggr=mean,
                   init=glorot_uniform, bias::Bool=true)
    in, out = ch
    W = init(out, 2*in)
    b = bias ? Flux.create_bias(W, true, out) : false
    SAGEConv(W, b, σ, aggr)
end

compute_message(l::SAGEConv, x_i, x_j, e_ij) =  x_j
update_node(l::SAGEConv, m, x) = l.σ.(l.weight * vcat(x, m) .+ l.bias)

function (l::SAGEConv)(g::GNNGraph, x::AbstractMatrix)
    check_num_nodes(g, x)
    x, _ = propagate(l, g, l.aggr, x)
    x
end

function Base.show(io::IO, l::SAGEConv)
    out_channel, in_channel = size(l.weight)
    print(io, "SAGEConv(", in_channel ÷ 2, " => ", out_channel)
    l.σ == identity || print(io, ", ", l.σ)
    print(io, ", aggr=", l.aggr)
    print(io, ")")
end

@doc raw"""
    ResGatedGraphConv(in => out, act=identity; init=glorot_uniform, bias=true)

The residual gated graph convolutional operator from the [Residual Gated Graph ConvNets]((https://arxiv.org/abs/1711.07553)) paper.

The layer's forward pass is given by

```math
\mathbf{x}_i' = act\big(U\mathbf{x}_i + \sum_{j \in N(i)} \eta_{ij} V \mathbf{x}_j\big),
```
where the edge gates ``\eta_{ij}`` are given by

```math
\eta_{ij} = sigmoid(A\mathbf{x}_i + B\mathbf{x}_j).
```

# Arguments

- `in`: The dimension of input features.
- `out`: The dimension of output features.
- `act`: Activation function.
- `init`: Weight matrices' initializing function.
- `bias`: Learn an additive bias if true.
"""
struct ResGatedGraphConv <: GNNLayer
    A
    B
    U
    V
    bias
    σ
end

@functor ResGatedGraphConv

function ResGatedGraphConv(ch::Pair{Int,Int}, σ=identity;
                      init=glorot_uniform, bias::Bool=true)
    in, out = ch
    A = init(out, in)
    B = init(out, in)
    U = init(out, in)
    V = init(out, in)
    b = bias ? Flux.create_bias(A, true, out) : false
    return ResGatedGraphConv(A, B, U, V, b, σ)
end

function compute_message(l::ResGatedGraphConv, di, dj)
    η = sigmoid.(di.Ax .+ dj.Bx)
    return η .* dj.Vx
end

update_node(l::ResGatedGraphConv, m, x) = m

function (l::ResGatedGraphConv)(g::GNNGraph, x::AbstractMatrix)
    check_num_nodes(g, x)

    Ax = l.A * x
    Bx = l.B * x
    Vx = l.V * x

    m, _ = propagate(l, g, +, (; Ax, Bx, Vx))

    return l.σ.(l.U*x .+ m .+ l.bias)
end

function Base.show(io::IO, l::ResGatedGraphConv)
    out_channel, in_channel = size(l.A)
    print(io, "ResGatedGraphConv(", in_channel, "=>", out_channel)
    l.σ == identity || print(io, ", ", l.σ)
    print(io, ")")
end