text:
"""
Support for generic maximal orders over any PID

  final result:
    integral_closure(R, F)
    where R is "any" PID and F a finite extension of some quotient field of R

  R needs to support
   - euclidean (hnf, pseudo_inv, gcd, lcm, mod, div, divrem)
   - factorisation
   - a useful ResidueField (need to know characteristic and finiteness)
   - integral_split, numerator, denominator
     given a in Frac(R), decompose into num, den
     (all Localisations of Z have QQ as quotient field,
     Q[x], Z[x] and Localisation(Q(x), degree) use Q(t))
   - isdomain_type

Seems to work for
-  R = ZZ, F = AnticNumberField
-  R = Loc{fmpz}, F = AnticNumberField

-  R = k[x], F = FunctionField (for k = QQ, F_q)
-  R = Localization(k(x), degree), F = FunctionField
-  R = Z[x], F = FunctionField/ QQ(t)
"""
module GenericRound2

using Hecke, Markdown
import AbstractAlgebra, Nemo
import Base: +, -, *, gcd, lcm, divrem, div, rem, mod, ^, ==
export integral_closure, extension_field
import AbstractAlgebra: expressify

#TODO: type parametrisation....
mutable struct Order <: AbstractAlgebra.Ring
  F::AbstractAlgebra.Field
  R::AbstractAlgebra.Ring
  trans::MatElem #both matrices are over coefficient_ring(F)
  itrans::MatElem

  function Order(R::AbstractAlgebra.Ring, F::AbstractAlgebra.Field, empty::Bool = false; check::Bool = true)
    #empty allows to create an Order that is none:
    # Z[x]/3x+1 is no order. This will be "fixed" by using any_order, but
    #the intial shell needs to be empty (illegal)
    r = new()
    r.F = F
    r.R = R
    empty && return r
    Qt = base_field(F)
    d = reduce(lcm, map(x->denominator(x, R), coefficients(defining_polynomial(F))))
    f = map_coefficients(x->numerator(Qt(d)*x, R), defining_polynomial(F))
    if !ismonic(f) #need Lenstra Order
      d = degree(F)
      M = zero_matrix(Qt, d, d)
      M[1, 1] = one(Qt)
      for i in 2:d
        for j in i:-1:1
          M[i, j] = coeff(f, d - (i - j))
        end
      end
      O = Order(r, M, one(Qt), check = check)
      return O
    end
    return r
  end

  function Order(O::Order, T::MatElem, d::RingElem; check::Bool = true)
    F = base_field(O.F)
    T = map_entries(F, T)
    T = divexact(T, base_ring(T)(d))
    Ti = inv(T)
    r = Order(O.R, O.F, true)
    if isdefined(O, :trans)
      r.trans = T*O.trans
      r.itrans = O.itrans*Ti
    else
      r.trans = T
      r.itrans = Ti
    end
    check && map(representation_matrix, basis(r))
    return r
  end
end

function Base.show(io::IO, O::Order)
  print(io, AbstractAlgebra.obj_to_string(O, context = io))
end

function expressify(O::Order; context = nothing)
  return Expr(:sequence, Expr(:text, "generic order in "),
                expressify(O.F, context = context),
                Expr(:text, " over "),
                expressify(O.R, context = context))
end

Hecke.base_ring(O::Order) = O.R
Hecke.coefficient_ring(O::Order) = O.R

mutable struct OrderElem <: RingElem
  parent :: Order
  data :: FieldElem
  coord :: Vector{RingElem}

  function OrderElem(O::Order, f::FieldElem, check::Bool = false)
    @assert parent(f) == O.F
    if check && !isone(integral_split(f, O)[2])
      error("element not in order")
    end
    r = new()
    r.parent = O
    r.data = f
    return r
  end
  function OrderElem(O::Order, f::fmpz)
    return OrderElem(O, O.F(f))
  end
end

function Base.show(io::IO, a::OrderElem)
  print(io, AbstractAlgebra.obj_to_string(a.data, context = io))
end

function expressify(a::OrderElem; context = nothing)
  return expressify(base_ring(R), context = context)
end

Nemo.elem_type(::Order) = OrderElem
Nemo.parent_type(::OrderElem) = Order
Nemo.parent_type(::Type{OrderElem}) = Order
Nemo.isdomain_type(::Type{OrderElem}) = true

Base.parent(a::OrderElem) = a.parent

(R::Order)(a::FieldElem, check::Bool = true) = OrderElem(R, a, check)
(R::Order)(a::fmpz) = OrderElem(R, a)
(R::Order)(a::Integer) = OrderElem(R, fmpz(a))
(R::Order)(a::OrderElem, check::Bool = true) = OrderElem(R, a.data, check)
(R::Order)() = R(0)

function Base.in(a::OrderElem, O::Order)
  return isone(integral_split(a.data, O)[2])
end

function Base.in(a::FieldElem, O::Order)
  return isone(integral_split(a.data, O)[2])
end

Nemo.iszero(a::OrderElem) = iszero(a.data)
Nemo.isone(a::OrderElem) = isone(a.data)

Nemo.zero(R::Order) = R(0)
Nemo.one(R::Order) = R(1)
Nemo.canonical_unit(a::OrderElem) = OrderElem(parent(a), fmpz(1))

Base.deepcopy_internal(a::OrderElem, dict::IdDict) = OrderElem(parent(a), Base.deepcopy_internal(a.data, dict))

check_parent(a::OrderElem, b::OrderElem) = parent(a) == parent(b) || error("Incompatible orders")

+(a::OrderElem, b::OrderElem) = check_parent(a, b) && OrderElem(parent(a), a.data + b.data)
-(a::OrderElem, b::OrderElem) = check_parent(a, b) && OrderElem(parent(a), a.data - b.data)
-(a::OrderElem) = OrderElem(parent(a), -a.data)
*(a::OrderElem, b::OrderElem) = check_parent(a, b) && OrderElem(parent(a), a.data * b.data)

==(a::OrderElem, b::OrderElem) = check_parent(a, b) && parent(a) == parent(b) && a.data == b.data

function Hecke.mul!(a::OrderElem, b::OrderElem, c::OrderElem)
  a.data = Hecke.mul!(a.data, b.data, c.data)
  return a
end

function Hecke.add!(a::OrderElem, b::OrderElem, c::OrderElem)
  a.data = Hecke.add!(a.data, b.data, c.data)
  return a
end

function Hecke.addeq!(a::OrderElem, b::OrderElem)
  a.data = Hecke.addeq!(a.data, b.data)
  return a
end

#TODO: move elsewhere?
function Hecke.lcm(a::Vector{<:RingElem})
  if length(a) == 0
    error("don't know the ring")
  end
  return reduce(lcm, a)
end
#TODO: move elsewhere?
Nemo.ngens(R::MPolyRing) = Nemo.nvars(R)

function Hecke.tr(a::OrderElem)
  return parent(a).R(trace(a.data))
end

function Hecke.coordinates(a::FieldElem, O::Order)
  c = coordinates(a)
  if isdefined(O, :itrans)
    d = matrix(c)'*O.itrans
  else
    d = matrix(c)'
  end
  return d
end

function Hecke.coordinates(a::OrderElem)
  return coordinates(a.data, parent(a))
end

function Hecke.coordinates(a::Generic.FunctionFieldElem)
  return [coeff(a, i) for i=0:degree(parent(a))-1]
end

@doc Markdown.doc"""
    integral_split(a::FieldElem, O::Ring)

For an element of the quotient field of some ring $O$, decompose
$a$ as
an element $n$ in $O$ and some denominator $d$, either in $O$ or the
coefficient ring of $O$.

# EXAMPLES
```julia
julia> integral_split(1//3, ZZ)
(1, 3)

julia> k, a = quadratic_field(5);

julia> zk = maximal_order(k);

julia> integral_split(1//a, zk)
(5, [-1 2])
```
"""
function Hecke.integral_split(a::Generic.FunctionFieldElem, O::Order)
  d = integral_split(coordinates(a, O), base_ring(O))[2]
  return O(base_ring(parent(a))(d)*a, false), d
end

function Hecke.integral_split(a::nf_elem, O::Order)
  d = integral_split(coordinates(a, O), base_ring(O))[2]
  return O(d.data*a, false), d #evil, but no legal way found
end

Hecke.degree(O::Order) = degree(O.F)

function Hecke.basis(O::Order)
  if isdefined(O, :itrans)
    return [O(O.F(vec(collect(O.trans[i,:])))) for i=1:degree(O)]
  else
    return map(O, basis(O.F))
  end
end

function (O::Order)(c::Vector)
  if isdefined(O, :itrans)
    return O(O.F(vec(collect(matrix(map(base_ring(O.trans), c))'*O.trans))))
  else
    return O(O.F(c))
  end
end

Base.:^(a::OrderElem, n::Integer) = parent(a)(a.data^n)

function Hecke.representation_matrix(a::OrderElem)
  O = parent(a)
  b = basis(O)
  m = zero_matrix(base_ring(O), degree(O), degree(O))
  for i=1:degree(O)
    c = coordinates(b[i]*a)
    for j=1:degree(O)
      m[i,j] = numerator(c[j], base_ring(O))
      @assert isone(denominator(c[j], base_ring(O)))
    end
  end
  return m
end

function Hecke.charpoly(a::OrderElem)
  return charpoly(representation_matrix(a))
end

function Hecke.minpoly(a::OrderElem)
  return minpoly(representation_matrix(a))
end

######################################################################
#
# pure function field
#
######################################################################
struct FFElemCoeffs{T <: RingElement}
   f::T
end

function Nemo.coefficients(f::Generic.FunctionFieldElem)
   return FFElemCoeffs(f)
end

function Base.iterate(PC::FFElemCoeffs, st::Int = -1)
   st += 1
   if st >= degree(parent(PC.f))
       return nothing
   else
       return coeff(PC.f, st), st
   end
end
Base.length(PC::FFElemCoeffs) = degree(parent(PC.f))

#####################################################################
#
# towards round2:
#   p-radical
#     via powers in perfect field case, finite char.
#     via powers in F_p(t) (non-perfect)
#     via trace if char 0 or char > degree
#   ring of multipliers
#

function Hecke.mod(a::OrderElem, p::RingElem)
  O = parent(a)
  R = parent(p)
  S = base_field(O.F)

  if isdefined(O, :itrans)
    a = map(x->S(R(x) % p), coordinates(a))
    b = a*O.trans
    return O(O.F(vec(collect(b'))))
  else
    mu = elem_type(O.R)[O.R(x) % p for x = coefficients(a.data)]
    return O(O.F(mu))
  end
end

function Hecke.powermod(a::OrderElem, n::fmpz, p::RingElem)
  c = parent(a)(1)
  for i = Hecke.BitsMod.bits(n)
    c *= c
    if i
      c *= a
    end
    c = Hecke.mod(c, p)
  end
  return c
end

# we don't have ideals, so radical is given via a matrix where
# rows are an S-basis
#in pos. char: O/p -> O/p : x-> x^(p^l) has the radical as kernel, perfect field
function radical_basis_power(O::Order, p::RingElem)
  t = ResidueField(parent(p), p)
  if isa(t, Tuple)
    F, mF = t
  else
    F = t
    mF = MapFromFunc(x->F(x), y->lift(y), parent(p), F)
  end
#  @assert characteristic(F) == 0 || (isfinite(F) && characteristic(F) > degree(O))
  q = characteristic(F)
  @assert q > 0
  while q < degree(O)
    q *= characteristic(F)
  end

  b = basis(O)
  m = zero_matrix(F, degree(O), degree(O))
  for i=1:degree(O)
    c = coordinates(powermod(b[i], q, p))
    for j=1:degree(O)
      m[j,i] = mF(O.R(c[j]))
    end
  end
  k, B = kernel(m)

  M2 = B[:, 1:k]'
  M2 = map_entries(x->preimage(mF, x), M2)
  M3 = Hecke.hnf_modular(M2, p, true)
  return M3 #[O(vec(collect((M3[i, :])))) for i=1:degree(O)]
end

#in char 0 and small char: rad = {x : Tr(xO) in pR} perfect field
function radical_basis_trace(O::Order, p::RingElem)
  T = trace_matrix(O)

  t = ResidueField(parent(p), p)
  if isa(t, Tuple)
    R, mR = t
  else
    R = t
    mR = MapFromFunc(x->R(x), y->lift(y), parent(p), R)
  end

  TT = map_entries(mR, T)
  k, B = kernel(TT)
  M2 = map_entries(x->preimage(mR, x), B[:, 1:k])'
  M3 = Hecke.hnf_modular(M2, p, true)
  return M3 #[O(vec(collect((M3[i, :])))) for i=1:degree(O)]
end

#pos. char, non-perfect (residue) field
function radical_basis_power_non_perfect(O::Order, p::RingElem)
  t = ResidueField(parent(p), p)
  if isa(t, Tuple)
    F, mF = t
  else
    F = t
    mF = MapFromFunc(x->F(x), y->lift(y), parent(p), F)
  end
  @assert isa(F, Generic.RationalFunctionField) && characteristic(F) != 0
  q = characteristic(F)
  @assert q > 0
  while q < degree(O)
    q *= characteristic(F)
  end
#=
  rad is still kernel of O/pO -> O/pO x -> x^(p^l), but
  this map is F_p linear, but not F-linear where F is the residue field.
  We need lin. comb. where the coefficients are all p^l-th powers, so we
  think in terms of a field extension
  F = F_p(t)/F_p(s) for s = t^(p^l)
  we want the kernel over F_p(s), not F_p(t)
=#

  q = Int(q)
  b = basis(O)
  dd = denominator(F(1))
  mm = zero_matrix(F, degree(O), degree(O))
  m = zero_matrix(F, q*degree(O), degree(O))
  for i=1:degree(O)
    c = coordinates(powermod(b[i], fmpz(q), p))
    for j=1:degree(O)
      d = mF(O.R(c[j]))
      d2 = denominator(d)
      l = lcm(dd, d2)
      d *= l
      if l != dd
        mm *= divexact(l, dd)
      end
      dd = l
      mm[i,j] = d
    end
  end

  for i=1:degree(O)
    for j=1:degree(O)
      d = numerator(mm[i,j])
      for k=0:degree(d)
        iszero(coeff(d, k)) && continue
        m[(j-1)*q+(k%q)+1,i] += gen(parent(d))^div(k, q)*coeff(d, k)
      end
    end
  end
  k, B = kernel(m)

  M2 = B[:, 1:k]'
  M2 = map_entries(x->preimage(mF, x), M2)
  M3 = Hecke.hnf_modular(M2, p, true)
  return return M3 #[O(vec(collect((M3[i, :])))) for i=1:degree(O)]
end

function Hecke.representation_matrix(a::Generic.FunctionFieldElem)
  F = parent(a)
  g = gen(F)
  m = zero_matrix(base_field(F), degree(F), degree(F))
  b = a
  for i=1:degree(F)
    for j=1:degree(F)
      m[i,j] = coeff(b, j-1)
    end
    if i<degree(F)
      b *= g
    end
  end
  return m
end

"""
    hnf_modular(M::MatElem{T}, d::T, isprime::Bool = false)

Return the `hnf` of `vcat(M, identity_matrix(parent(d), ncols(M)))`
if `isprime` is set, then instead of an `hnf` internally a `rref` over the
residue field modulo `d` is used.
"""
function Hecke.hnf_modular(M::MatElem{T}, d::T, isprime::Bool = false) where {T}
  if isprime
    x = ResidueField(parent(d), d)
    if isa(x, Tuple)
      R, mR = x
    else
      R = x
      mR = MapFromFunc(x->R(x), x->lift(x), parent(d), R)
    end
    r, h = rref(map_entries(mR, M))
    H = map_entries(x->preimage(mR, x), h[1:r, :])
  else
    R, mR = ResidueRing(parent(d), d)
    H = map_entries(x->preimage(mR, x), hnf(map_entries(mR, M)))
  end
  H = vcat(H, d*identity_matrix(parent(d), ncols(M)))
  H = hnf(H)
  @assert iszero(H[ncols(M)+1:end, :])
  return H[1:ncols(M), :]
end

function Base.divrem(a::fmpz_mod, b::fmpz_mod)
  R = parent(a)
  r = rem(a, b)
  return divexact(a-r, b), r
end

function Base.div(a::fmpz_mod, b::fmpz_mod)
  R = parent(a)
  r = rem(a, b)
  return divexact(a-r, b)
end

function Base.rem(a::fmpz_mod, b::fmpz_mod)
  R = parent(a)
  r = R(rem(lift(a), gcd(modulus(R), lift(b))))
  return r
end

function ring_of_multipliers(O::Order, I::MatElem{T}, p::T, isprime::Bool = false) where {T}
  #TODO: modular big hnf, peu-a-peu, not all in one
  @vprint :NfOrd 2 "ring of multipliers module $p (isprime: $isprime) of ideal with basis matrix $I\n"
  II, d = pseudo_inv(I)
  @assert II*I == d

  m = hcat([divexact(representation_matrix(O(vec(collect(I[i, :]))))*II, d) for i=1:nrows(I)]...)
  m = m'
  if isprime
    x = ResidueField(parent(p), p)
    if isa(x, Tuple)
      R, mR = x
    else
      R = x
      mR = MapFromFunc(x->R(x), x->lift(x), parent(p), R)
    end
    ref = x->rref(x)[2]
  else
    R, mR = ResidueRing(parent(p), p)
#    R = parent(p)
#    mR = MapFromFunc(x->x, x->x, R, R)
    ref = hnf
  end
  m = map_entries(mR, m)

  n = degree(O)
  mm = ref(m[1:n, 1:n])
  for i=2:n
    mm = vcat(mm, ref(m[(i-1)*n+1:i*n, 1:n]))
    mm = ref(mm)
    @assert iszero(mm[n+1:end, :])
    mm = mm[1:n, 1:n]
  end
#  H = hnf(map_entries(x->preimage(mR, x), mm))
  H = hnf_modular(map_entries(x->preimage(mR, x), mm), p, isprime)

  @vtime :NfOrd 2 Hi, d = pseudo_inv(H)

  O = Order(O, Hi', d, check = false)
  return O
end

function ring_of_multipliers(O::Order, I::MatElem)
  #TODO: modular big hnf, peu-a-peu, not all in one
  @vprint :NfOrd 2 "ring of multipliers of ideal with basis matrix $I\n"
  II, d = pseudo_inv(I)
  @assert II*I == d

  m = hcat([divexact(representation_matrix(O(vec(collect(I[i, :]))))*II, d) for i=1:nrows(I)]...)
  m = m'
  n = degree(O)
  mm = hnf(m[1:n, 1:n])
  for i=2:n
    mm = vcat(mm, hnf(m[(i-1)*n+1:i*n, 1:n]))
    mm = hnf(mm)[1:n, 1:n]
  end
  H = mm

  @vtime :NfOrd 2 Hi, d = pseudo_inv(H)

  O = Order(O, Hi', d, check = false)
  return O
end

function Hecke.trace_matrix(O::Order)
  return trace_matrix(basis(O))
end

function Hecke.trace_matrix(b::Vector{OrderElem})
  O = parent(b[1])
  m = zero_matrix(O.R, length(b), length(b))
  for i=1:length(b)
    m[i,i] = trace(b[i]*b[i])
    for j=i+1:length(b)
      m[i,j] = m[j,i] = trace(b[i]*b[j])
    end
  end
  return m
end

function Hecke.trace_matrix(b::Vector{OrderElem}, c::Vector{OrderElem}, exp::fmpz = fmpz(1))
  O = parent(b[1])
  m = zero_matrix(O.R, length(b), length(c))
  for i=1:length(b)
    for j=1:length(c)
      m[i,j] = trace((b[i]*c[j])^exp)
    end
  end
  return m
end

function Hecke.pmaximal_overorder(O::Order, p::RingElem, isprime::Bool = false)
  @vprint :NfOrd 1 "computing a $p-maximal orderorder\n"

  t = ResidueField(parent(p), p)

  if isa(t, Tuple)
    R, mR = t
  else
    R = t
    mR = MapFromFunc(x->R(x), y->lift(y), parent(p), R)
  end
#  @assert characteristic(F) == 0 || (isfinite(F) && characteristic(F) > degree(O))
  if characteristic(R) == 0 || characteristic(R) > degree(O)
    @vprint :NfOrd 1 "using trace-radical for $p\n"
    rad = radical_basis_trace
  elseif isa(R, Generic.RationalFunctionField)
    @vprint :NfOrd 1 "non-perfect case for radical for $p\n"
    rad = radical_basis_power_non_perfect
  else
    @vprint :NfOrd 1 "using radical-by-power for $p\n"
    rad = radical_basis_power
  end
  while true #TODO: check the discriminant to maybe skip the last iteration
    I = rad(O, p)
    S = ring_of_multipliers(O, I, p, isprime)
    if discriminant(O) == discriminant(S)
      return O
    end
    O = S
  end
end

function Hecke.function_field(f::PolyElem{<:Generic.Rat}, s::String = "_a"; check::Bool = true, cached::Bool = false)
  return FunctionField(f, s, cached = cached)
end

function Hecke.function_field(f::PolyElem{<:Generic.Rat}, s::Symbol; check::Bool = true, cached::Bool = false)
  return FunctionField(f, s, cached = cached)
end

function extension_field(f::PolyElem{<:Generic.Rat}, s::String = "_a"; check::Bool = true, cached::Bool = false)
  return FunctionField(f, s, cached = cached)
end

function extension_field(f::PolyElem{<:Generic.Rat}, s::Symbol; check::Bool = true, cached::Bool = false)
  return FunctionField(f, s, cached = cached)
end

@doc Markdown.doc"""
    integral_closure(R, F)

Computes the integral closure of the ring $R$ in the field $F$. $F$ needs
to be a finite extension over the fraction field of $R$. The algorithm
uses a variant of the Round-2 method.

Currently supported are

- $R$ the integers and $F$ an (absolute simple) number field. Here the result is an number
  field order.
- $R$ a localisation of the integers and $F$ an (absolute simple) number. Here the result is
  a generic order.
- $R$ a univariate polynomial ring over a field $k$ and $F$ a function field (finite extension of $k(t)$.
- $R$ the degree-localisation of a univariate polynomial ring over a field $k$ amd $F$ a finite extension of $k(t)$
- $R$ a univariate polynomial ring over the integers and $F$ a finite extension of $Q(t)$ for the rational field $Q$.

# EXAMPLES
```julia
julia> k, a = quadratic_field(12);

julia> integral_closure(ZZ, k)

Maximal order of Real quadratic field defined by x^2 - 12
with basis nf_elem[1, 1//2*sqrt(12)]
```
"""
function integral_closure(::FlintIntegerRing, F::AnticNumberField)
  return Hecke.maximal_order(F)
end

function integral_closure(S::Loc{fmpz}, F::AnticNumberField)
  return _integral_closure(S, F)
end

function integral_closure(S::PolyRing{T}, F::Generic.FunctionField{T}) where {T}
  return _integral_closure(S, F)
end

function integral_closure(S::KInftyRing{T}, F::Generic.FunctionField{T}) where {T}
  return _integral_closure(S, F)
end

function _integral_closure(S::AbstractAlgebra.Ring, F::AbstractAlgebra.Ring)
  O = Order(S, F)
  return Hecke.maximal_order(O)
end

function Hecke.maximal_order(O::Order)
  @vprint :NfOrd 1 "starting maximal order...\n"
  S = base_ring(O)
  d = discriminant(O)
  @vprint :NfOrd 2 "factoring the discriminant...\n"
  @vtime :NfOrd 2 ld = factor(d)
  local Op
  first = true
  for (p,k) = ld.fac
    if k<2
      continue
    end
    OO = pmaximal_overorder(O, p, true)
    if !isdefined(OO, :trans)
      continue
    end
    T = integral_split(OO.trans, S)
    isone(T[2]) && continue
    if first
      Op = T
      first = false
    else
      Op = (Hecke._hnf(vcat(Op[1]*T[2], T[1]*Op[2]), :lowerleft)[degree(O)+1:end, :], T[2]*Op[2])
    end
  end
  if first
    return O
  else
    return Order(O, Op[1], Op[2])
  end
end

function Hecke.discriminant(O::Order)
  d = discriminant(O.F)
  if isdefined(O, :trans)
    d *= det(O.trans)^2
  end
  return O.R(d)
end

function Hecke.basis(O::Order, F::Generic.FunctionField)
  return map(F, basis(O))
end

function (K::AnticNumberField)(b::OrderElem)
  O = parent(b)
  @assert O.F == K
  return b.data
end

function Hecke.basis(O::Order, F::AnticNumberField)
  return map(F, basis(O))
end

function Hecke.basis(F::Generic.FunctionField)
  a = gen(F)
  bas = [a^0, a]
  while length(bas) < degree(F)
    push!(bas, bas[end]*a)
  end
  return bas
end

#=
function (R::PolyRing{T})(a::Generic.Rat{T}) where {T}
  @assert isone(denominator(a))
  return R(numerator(a))
end
=#
function (R::FmpqPolyRing)(a::Generic.Rat{fmpq})
  @assert isone(denominator(a))
  return R(numerator(a))
end

function Hecke.ResidueField(R::FmpqPolyRing, p::fmpq_poly)
  K, _ = number_field(p)
  return K, MapFromFunc(x->K(x), y->R(y), R, K)
end

function (F::Generic.FunctionField{T})(p::PolyElem{<:AbstractAlgebra.Generic.Rat{T}}) where {T <: FieldElem}
  @assert parent(p) == parent(F.pol)
  @assert degree(p) < degree(F) # the reduction is not implemented
  R = parent(gen(F).num)
  S = coefficient_ring(R)
  d = reduce(lcm, map(denominator, coefficients(p)), init = one(S))
  return F(map_coefficients(S, d*p, parent = R), d)
end

function (F::Generic.FunctionField)(c::Vector{<:RingElem})
  return F(parent(F.pol)(map(base_field(F), c)))
end

(F::Generic.FunctionField)(a::OrderElem) = a.data

function Hecke.charpoly(a::Generic.FunctionFieldElem)
  return charpoly(representation_matrix(a))
end

function Hecke.minpoly(a::Generic.FunctionFieldElem)
  return minpoly(representation_matrix(a))
end

function Hecke.discriminant(F::Generic.FunctionField)
  return discriminant(defining_polynomial(F))
end

#######################################################################
#
# support for ZZ
#
#######################################################################
Hecke.denominator(a::fmpq, ::FlintIntegerRing) = denominator(a)
Hecke.numerator(a::fmpq, ::FlintIntegerRing) = numerator(a)
Hecke.integral_split(a::fmpq, ::FlintIntegerRing) = (numerator(a), denominator(a))
Hecke.integral_split(a::Rational, R::FlintIntegerRing) = integral_split(fmpq(a), R)

#######################################################################
#
# support for Loc{fmpz}
#
#######################################################################
function Hecke.integral_split(a::fmpq, R::Loc{fmpz})
  d = denominator(a)
  p = R.prime
  q,w = Hecke.ppio(d, p)
  if R.comp # den can ONLY use prime
    return R(numerator(a)//q), R(w)
  else
    return R(numerator(a)//w), R(q)
  end
end
Hecke.denominator(a::fmpq, R::Loc{fmpz}) = integral_split(a, R)[2]
Hecke.numerator(a::fmpq, R::Loc{fmpz}) = integral_split(a, R)[1]
(::FlintRationalField)(a::LocElem{fmpz}) = data(a)

function Hecke.factor(a::LocElem{fmpz})
  c = canonical_unit(a)
  b = a*inv(c)
  L = parent(a)
  @assert isone(denominator(data(b)))
  lf = factor(numerator(data(b)))
  return Fac(c, Dict(L(p)=>v for (p,v) = lf.fac))
end

function Hecke.ResidueField(R::Loc{fmpz}, p::LocElem{fmpz})
  pp = numerator(data(p))
  @assert isprime(pp) && isone(denominator(p))
  F = GF(pp)
  return F, MapFromFunc(x->F(data(x)), y->R(lift(y)), R, F)
end

Hecke.isdomain_type(::Type{LocElem{fmpz}}) = true

#######################################################################
#
# support for Rat{T}
#
#######################################################################
# Rat{T}, KInftyRing{T}

Base.denominator(x::AbstractAlgebra.Generic.Rat{T}, R::KInftyRing{T}) where {T} = Hecke.integral_split(x, R)[2]
Base.numerator(x::AbstractAlgebra.Generic.Rat{T}, R::KInftyRing{T}) where {T} = Hecke.integral_split(x, R)[1]

function Hecke.integral_split(x::AbstractAlgebra.Generic.Rat{T}, R::KInftyRing{T}) where {T}
  if iszero(x)
    return zero(R), one(R)
  end
  a = degree(numerator(x))
  b = degree(denominator(x))
  if b >= a
    return R(x), one(R)
  end
  t = gen(parent(x))
  return R(x*t^(b-a)), R(t^(b-a))
end

(R::Generic.RationalFunctionField{T})(x::KInftyElem{T}) where {T <: FieldElem} = x.d

base_ring_type(::Type{AbstractAlgebra.Generic.PolyRing{T}}) where {T} = parent_type(T)

base_ring_type(::Type{AcbPolyRing}) = AcbField

base_ring_type(::Type{ArbPolyRing}) = ArbField

base_ring_type(::Type{FmpqPolyRing}) = FlintRationalField

base_ring_type(::Type{FmpzModPolyRing}) = Nemo.FmpzModRing

base_ring_type(::Type{FmpzPolyRing}) = FlintIntegerRing

base_ring_type(::Type{FqDefaultPolyRing}) = FqDefaultFiniteField

base_ring_type(::Type{FqNmodPolyRing}) = FqNmodFiniteField

base_ring_type(::Type{FqPolyRing}) = FqFiniteField

base_ring_type(::Type{GFPFmpzPolyRing}) = Nemo.GaloisFmpzField

base_ring_type(::Type{GFPPolyRing}) = Nemo.GaloisField

base_ring_type(::Type{NmodPolyRing}) = Nemo.NmodRing

function (::PolyRing{T})(x::AbstractAlgebra.Generic.Rat{T}) where {T}
  @assert isone(denominator(x))
  return numerator(x)
end

# Rat{T}, PolyRing{T}
function Hecke.numerator(a::Generic.Rat{T}, S::PolyRing{T}) where {T}
  return numerator(a)
end

function Hecke.denominator(a::Generic.Rat{T}, S::PolyRing{T}) where {T}
  return denominator(a)
end

function Hecke.integral_split(a::Generic.Rat{T}, S::PolyRing{T}) where {T}
  return numerator(a), denominator(a)
end

function Hecke.factor(a::Generic.Rat{T}, R::S) where {T, S<:PolyRing{T}}
  @assert parent(numerator(a)) == R
  f1 = factor(numerator(a))
  f2 = factor(denominator(a))
  for (p,e) = f2.fac
    @assert !haskey(f1.fac, p)
    f1.fac[p] = -e
  end
  f1.unit = divexact(f1.unit, f2.unit)
  return f1
end

########################################################################
#
# Matrices
#
########################################################################
function Hecke.integral_split(M::MatElem{<:AbstractAlgebra.FieldElem}, S::Generic.Ring)
  m = zero_matrix(S, nrows(M), ncols(M))
  den = one(S)
  for i=1:nrows(M)
    for j=1:ncols(M)
      n, d = integral_split(M[i,j], S)
      q, r = divrem(den, d)
      if iszero(r)
        m[i,j] = n*q
      else
        t = lcm(den, d)
        m *= divexact(t, den)
        m[i,j] = n*divexact(t, d)
        den = t
      end
    end
  end
  return m, den
end

end  # ModuleRound2

using .GenericRound2

export integral_closure


text:
function PutCallParity(CallPutFlag, OptionPrice, F, X, T, r)

  # returns for a known options price and the corresponding underlying and Strike
  # the parity price of the other option gender

  if CallPutFlag == "c"
    p = OptionPrice - (F - X) * exp(-r * T)
  elseif CallPutFlag == "p"
    c = OptionPrice + (F - X) * exp(-r * T)
  end

end

ext: jl


text:
module MultilineStrings

export @m_str, indent, multiline

const DEFAULT_STYLE = :folded
const DEFAULT_CHOMP = :strip

"""
    multiline(str; style=$(repr(DEFAULT_STYLE)), chomp=$(repr(DEFAULT_CHOMP))) -> AbstractString

Manipulate a multiline string according to the provided `style` and `chomp`. Works similarly
to YAML multiline strings (also known as block scalars).

# Arguments
- `str::AbstractString`: The multiline string to be processed

# Keywords
- `style::Symbol`: Replace newlines with spaces (`:folded`) or keep newlines (`:literal`)
- `chomp::Symbol`: No newlines at the end (`:strip`), single newline at the end (`:clip`),
  or keep all newlines from the end (`:keep`)
"""
function multiline(str::AbstractString; style=DEFAULT_STYLE, chomp=DEFAULT_CHOMP)
    multiline(str, style, chomp)
end

"""
    multiline(str, indicators) -> AbstractString

Manipulate a multiline string according to the provided style and chomp encoded in the
`indicators` string.

# Arguments
- `str::AbstractString`: The multiline string to be processed
- `indicators::AbstractString`: A terse string representing the style and chomp. Indicators
  can be either in letter-form or in YAML-form:

    - "fs" / ">-": folded and strip
    - "fc" / ">": folded and clip
    - "fk" / ">+": folded and keep
    - "ls" / "|-": literal and strip
    - "lc" / "|": literal and clip
    - "lk" / "|+": literal and keep
"""
function multiline(str::AbstractString, indicators::AbstractString)
    style, chomp = _process_indicators(indicators)
    return multiline(str, style, chomp)
end

function multiline(str::AbstractString, style::Symbol, chomp::Symbol)
    if style === :folded
        str = replace(str, r"(?<=\S)\n(?=\S)" => " ")
        str = replace(str, r"(?<=\n)\n(?=\S)" => "")
    elseif style === :literal
        # no-op
    else
        throw(ArgumentError("Unknown style indicator: $style"))
    end

    if chomp === :clip
        if endswith(str, '\n')
            str = rstrip(str, '\n') * '\n'
        end
    elseif chomp === :strip
        str = rstrip(str, '\n')
    elseif chomp === :keep
        # no-op
    else
        throw(ArgumentError("Unknown chomping indicator: $chomp"))
    end

    return str
end

"""
    @m_str -> String

Manipulate a multiline string literal according to a style and chomp indicator (provided
after the ending quote):

- Style indicator:
    - `f` replace newlines with spaces (folded, default)
    - `l` keep newlines (literal)
- Chomp indicator:
    - `s` no newlines at the end (strip, default)
    - `c` single newline at the end (clip)
    - `k` keep all newlines from the end (keep)

If both a style and chomp indicator is provided the style indicator must be specified first.

Note string interpolation is still respected any newlines added from interpolation will be
also be processed.

# Examples

```jldoctest; setup = :(using MultilineStrings)
julia> m\"\"\"
       A string written
       over multiple lines
       \"\"\"
"A string written over multiple lines"
```
"""
macro m_str(str::AbstractString, indicators::AbstractString="")
    style, chomp = _process_indicators(indicators)
    parsed = interpolate(str)

    # When no string interpolation needs to take place we can just process the multiline
    # string during parse time. If string interpolation needs to take place we'll evaluate
    # the multiline string at runtime so that we can process after interpolation has taken
    # place.
    result = if parsed isa String
        multiline(unescape_string(parsed), style, chomp)
    else
        Expr(:call, :(MultilineStrings.multiline), parsed, QuoteNode(style), QuoteNode(chomp))
    end

    return esc(result)
end

function interpolate(str::AbstractString)
    components = []
    start = 1
    lastind = lastindex(str)

    state = iterate(str)
    while state !== nothing
        c, i = state

        if c == '$'
            ending = prevind(str, i, 2)
            start <= ending && push!(components, SubString(str, start, ending))

            expr, i = Meta.parse(str, i; greedy=false)
            push!(components, expr)
            start = i
        end

        state = iterate(str, i)
    end

    # When interpolation isn't used we can just return the original string
    start == 1 && return str

    ending = lastind
    start <= ending && push!(components, SubString(str, start, ending))

    return Expr(:string, components...)
end

function _process_indicators(indicators::AbstractString)
    indicators_len = length(indicators)
    indicators_len > 2 && throw(ArgumentError("Too many indicators provided"))

    # Note: Using '\0` to indicate undefined
    yaml_chomp = false
    style_char, chomp_char = if indicators_len == 2
        indicators
    elseif indicators_len == 1
        ind = indicators[1]
        if ind in ('f', 'l')
            ind, '\0'
        elseif ind in ('>', '|')
            yaml_chomp = true
            ind, '\0'
        else
            '\0', ind
        end
    else
        '\0', '\0'
    end

    if style_char != '\0' && chomp_char != '\0' && isletter(style_char) ⊻ isletter(chomp_char)
        throw(ArgumentError("Can't mix YAML style block indicators with letter indicators"))
    end

    style = if style_char == 'f' || style_char == '>'
        :folded
    elseif style_char == 'l' || style_char == '|'
        :literal
    elseif style_char == '\0'
        DEFAULT_STYLE
    else
        throw(ArgumentError("Unknown style indicator: $(repr(style_char))"))
    end

    chomp = if chomp_char == 'c' || yaml_chomp
        :clip
    elseif chomp_char == 's' || chomp_char == '-'
        :strip
    elseif chomp_char == 'k' || chomp_char == '+'
        :keep
    elseif chomp_char == '\0'
        DEFAULT_CHOMP
    else
        throw(ArgumentError("Unknown chomping indicator: $(repr(chomp_char))"))
    end

    return style, chomp
end

include("indent.jl")

end


text:
#=
Julia can natively call C and Fortran functions, and with the help of packages it can also access code from other languages, such as Python and R.

Julia has a "no boilerplate" philosophy for C and Fortran code.
There's no "glue" code, code generation of compilation.
To exemplify this, let's consider two files, one in Fortran and one in C:
```
! File lu_fact.f90
subroutine lu_factorization (nrow, ncol, A)

  implicit none

  integer, intent(in) :: nrow, ncol
  real, intent(inout), dimension(nrow, ncol) :: A

  integer :: i, j, k

  do j = 1,nrow
    do i = j+1,nrow
      A(i,j) = A(i,j) / A(j,j)
      do k = j+1,ncol
        A(i,k) = A(i,k) - A(j,k) * A(i,j)
      enddo
    enddo
  enddo

end
```
The code above implements the [LU factorization](https://en.wikipedia.org/wiki/LU_decomposition), without pivoting or checks.
The input matrix is also used as output.
The factorization is used to solve linear systems.

The code below implements the [forward and back substitutions](https://en.wikipedia.org/wiki/Triangular_matrix#Forward_and_back_substitution), using the resulting matrix from above as input.
The input vector `x` corresponds to the right-hand side vector on input, and as the solution of the linear system on output.
```
// File lu_solve.c
#include <stdio.h>

void lu_solve(long long n, float *LU, float *x)
{
  int i, j;

  for (i = 0; i < n; i++)
    for (j = 0; j < i; j++)
      x[i] -= LU[i + j * n] * x[j];

  for (i = n - 1; i >= 0; i--)
  {
    for (j = i + 1; j < n; j++)
      x[i] -= LU[i + j * n] * x[j];
    x[i] /= LU[i + i * n];
  }
}
```

The file `lu_fact.f90` is in Fortran 90 and `lu_solve.c` is in C, as the extension implies.
To compile these file to use with Julia, we have to prepare them into a shared library.
To that end, we'll compile them using the GCC/GFortran compilers with the `-fPIC -shared` flags, and then link them with `ld`.

The following commands were tested on linux:

```
gfortran -fPIC -shared lu_fact.f90 -o lu_fact.o
gcc -fPIC -shared lu_solve.c -o lu_solve.o
ld -shared $PWD/lu_fact.o $PWD/lu_solve.o -o lu.so
```

The `$PWD` variable is necessary because the shared library should know where the compiled packages are.
There are other ways of doing this.

Now we can use the library `lu.so` on Julia.
Here's an example:
=#
using Libdl, LinearAlgebra

path = "assets/interoperability"
lib = joinpath(path, "lu.so")

n = 10
L = tril(rand(Float32, n, n), -1) + I
U = triu(rand(Float32, n, n))
A = L * U
b = A * ones(Float32, n)

LU = copy(A)
x = copy(b)

ccall(
  (:lu_factorization_, lib),
  Cvoid,
  (Ref{Int}, Ref{Int}, Ref{Cfloat}),
  n, n, LU,
)

ccall(
  (:lu_solve, lib),
  Cvoid,
  (Int, Ref{Cfloat}, Ref{Cfloat}),
  n, LU, x,
)

println("‖LU - (L + U - I)‖ = ", norm(LU - L - U + I))
println("‖x - 1‖ = ", norm(x .- 1))
println("‖A * x - b‖ = ", norm(A * x - b))

#=
The first part of the code defines a "good" matrix `A` and `b` such that the solution to the linear system $Ax = b$ is `x = ones(n)`.
The `LU = copy(A)` and `x = copy(b)` are so that we can use `A` and `b` for testing later.
One important part of the code is that since we used the default variables for Fortran, we used `float` for C, and we have to use an appropriate type for Julia as well.
`Float32` refers to a 32 bits real number, i.e. `float`.
`Cfloat` is just an alias for `Float32`.

The `ccall` signature is simple:
- The first argument is the function or a tuple with the function and it's library;
- The second argument is the return type;
- The third argument is a tuple with the types of every argument;
- The other arguments are the arguments to the C/Fortran function.

Noteworthy:
- Fortran with `gfortran` added a `_` to the function name;
- `Ref{T}` indicates that a reference will be passed, allowing the argument to be modified (`Ptr{T}` also exists, to indicate returned pointers by C);
- Fortran always passes variables by reference.

## FFTW

For a more substantial example, we can look at the [FFTW Library](https://www.fftw.org) for the discrete Fourier transform.
The compilation of this example takes some time and the installation may be less straightforward than in my system, so be warned.

- First, download the library. I'm using v3.3.10. Untar it too.
- Then, enter the folder and issue a `./configure --enable-shared`.
- Now, enter `make`.

This will produce a shared library `.libs/libfftw3.so`.

Now we can write a function to compute the FFT of a given vector using this lib.

We'll be create a simple version of the following C code, available in the [online docs of FFTW](https://www.fftw.org/fftw3_doc/Complex-One_002dDimensional-DFTs.html#Complex-One_002dDimensional-DFTs).
```
#include <fftw3.h>
...
{
    fftw_complex *in, *out;
    fftw_plan p;
    ...
    in = (fftw_complex*) fftw_malloc(sizeof(fftw_complex) * N);
    out = (fftw_complex*) fftw_malloc(sizeof(fftw_complex) * N);
    p = fftw_plan_dft_1d(N, in, out, FFTW_FORWARD, FFTW_ESTIMATE);
    ...
    fftw_execute(p); /* repeat as needed */
    ...
    fftw_destroy_plan(p);
    fftw_free(in); fftw_free(out);
}
```
=#
function myfft(x::Vector{ComplexF64})
  y = copy(x)
  FFTW_FORWARD = Int32(-1)
  FFTW_ESTIMATE = Int32(64)

  plan = ccall(
    (:fftw_plan_dft_1d, "assets/interoperability/fftw-3.3.10/.libs/libfftw3.so"),
    Ptr{Cvoid},
    (Int32, Ref{ComplexF64}, Ref{ComplexF64}, Int32, UInt32),
    length(x), x, y, FFTW_FORWARD, FFTW_ESTIMATE
  )
  ccall(
    (:fftw_execute, "assets/interoperability/fftw-3.3.10/.libs/libfftw3.so"),
    Cvoid,
    (Ptr{Cvoid},),
    plan
  )
  ccall(
    (:fftw_destroy_plan, "assets/interoperability/fftw-3.3.10/.libs/libfftw3.so"),
    Cvoid,
    (Ptr{Cvoid},),
    plan
  )
  return y
end

x = rand(ComplexF64, 8)
y = myfft(x)

#=
We can implement the a basic DFT to compare:
=#
function mydft(x)
  n = length(x)
  return [
    sum(x[i+1] * exp(-im * 2π * i * k / n) for i = 0:n-1)
    for k = 0:n-1
  ]
end
mydft(x)

#=
You may have noticed:
- Memory is handled by Julia in this case, so we don't use `malloc`.
- The C struct `fftw_complex` simply stores two `double`s. In Julia we would need a similar struct, and `ComplexF64` is exactly that, so it's extra smooth.
- We can use `Ptr{Cvoid}` to avoid dealing with the return type of `fftw_plan_dft_1d`, which is an opaque pointer.

### FFTW.jl (and where are the .so files?)

The ability to easily integrate C/Fortran functions is also one of the foundations of Julia.
In fact, the example above is adapted from a [2011 commit](https://github.com/JuliaLang/julia/commit/021afbfbac5ea941dc5402af2a9dea6d91fc7267) almost unmodified.
The FFTW package is now wrapped in the package `FFTW.jl`. You can add and use it easily:
=#
using FFTW
y = fft(x)
#=
If you install `FFTW.jl`, you will notice that you don't need to install the C library `FFTW` manually, i.e., the Julia package automatically does that for you.
Julia has a built-in [`Artifacts` system](https://pkgdocs.julialang.org/dev/artifacts/) and a commonly used package `BinaryBuilder.jl` that work in conjunction to allow binaries to be pre-built and distributed.
The details are complicated, but the user doesn't have to worry about them.
=#

#=
## Python

Julia can also can Python and R code, using `PyCall` and `RCall`, respectively.
We'll focus on Python for this example.

We're gonna use `PyCall`. Like `IJulia`,
=#
using Plots, PyCall, Random
pyplot()

Random.seed!(123)
n = 2^8
X = (2 * rand(n, 2) .- 1) * 4
y = [
  X[i,2]^2 * 16 - X[i,1]^2 * 16 ≤ 2 * randn() ? 1 : -1 for i = 1:n
]
y = [
  (X[i,1] - 1)^2 + (4 - randn()) * (X[i,2] - X[i,1]^2)^2 ≤ 10 ? 1 : -1 for i = 1:n
]
plot(leg=false)
idx = findall(y .== -1)
scatter!(X[idx,1], X[idx,2], m=(4,:red,:square))
idx = findall(y .== 1)
scatter!(X[idx,1], X[idx,2], m=(4,:blue,:circle))
png(joinpath(@OUTPUT, "pycall-1")) # hide

#=
\fig{pycall-1}
=#

svm = pyimport_conda("sklearn.svm", "scikit-learn")
clf = svm.SVC(C=1e-2, gamma=5.0, probability=true)
clf.fit(X, y)
x1g = range(extrema(X[:,1])..., length=100)
x2g = range(extrema(X[:,2])..., length=100)
Z = [
  clf.predict_proba([x1i x2j;])[2] for x2j in x2g, x1i in x1g
]
contourf(x1g, x2g, Z, c=cgrad([:pink,:magenta,:lightblue]), levels=50)
idx = findall(y .== -1)
scatter!(X[idx,1], X[idx,2], m=(4,:red,:square), lab="", opacity=0.5)
idx = findall(y .== 1)
scatter!(X[idx,1], X[idx,2], m=(4,:blue,:circle), lab="", opacity=0.5)
png(joinpath(@OUTPUT, "pycall-2")) # hide

#=
\fig{pycall-2}
=#

max_stars_repo_name: chunggnuhc/julia-tutorial-at-nlesc

max_issues_repo_name: chunggnuhc/julia-tutorial-at-nlesc

max_forks_repo_name: chunggnuhc/julia-tutorial-at-nlesc


text:
module MRIQuant

	include("Constants.jl")
	include("Relaxation.jl")
	include("ExtendedPhaseGraphs.jl")

end


text:
using PQEcolaPoint
using Dates

# Simulate a PUT credit spread -

# create a SHORT put contract -
short_put_contract = PutContractModel()
short_put_contract.ticker = "XYZ"
short_put_contract.expiration_date = Date(2022, 04, 14)
short_put_contract.strike_price = 315.0
short_put_contract.premium = 5.25
short_put_contract.number_of_contracts = 1
short_put_contract.direction = -1

# create a LONG put contract -
long_put_contract = PutContractModel()
long_put_contract.ticker = "XYZ"
long_put_contract.expiration_date = Date(2022, 04, 14)
long_put_contract.strike_price = 310.0
long_put_contract.premium = 4.31
long_put_contract.number_of_contracts = 1
long_put_contract.direction = 1

# build model -
put_credit_spread_contract_array = Array{AbstractAssetModel,1}()
push!(put_credit_spread_contract_array, short_put_contract)
push!(put_credit_spread_contract_array, long_put_contract)

# setup the underlying -
underlying_range = range(305.0, stop = 320.0, length = 1000) |> collect

# compute the table -
short_credit_spread_table = expiration(put_credit_spread_contract_array, underlying_range)


text:
# generate LLVM wrappers

using LLVM_full_jll

using Clang.Generators
import Clang

cd(@__DIR__)

@add_def off_t

function main()
    options = load_options(joinpath(@__DIR__, "wrap_llvmextra.toml"))

    includedir = LLVM_full_jll.llvm_config() do config
        readchomp(`$config --includedir`)
    end
    cppflags = LLVM_full_jll.llvm_config() do config
        split(readchomp(`$config --cppflags`))
    end

    args = get_default_args("x86_64-linux-gnu")
    push!(args, "-Isystem$includedir")
    append!(args, cppflags)

    header_files = detect_headers(realpath(joinpath(@__DIR__, "..", "deps", "LLVMExtra", "include")), args)

    ctx = create_context(header_files, args, options)

    build!(ctx, BUILDSTAGE_NO_PRINTING)

    # custom rewriter
    function rewrite!(dag::ExprDAG)
        replace!(get_nodes(dag)) do node
            filename = normpath(Clang.get_filename(node.cursor))
            if !contains(filename, "LLVMExtra")
                return ExprNode(node.id, Generators.Skip(), node.cursor, Expr[], node.adj)
            end
            return node
        end
    end

    rewrite!(ctx.dag)

    # print
    build!(ctx, BUILDSTAGE_PRINTING_ONLY)

end

isinteractive() || main()


text:
using LinearAlgebra
import Base: *
import LinearAlgebra: mul!

mutable struct LBFGSOperator{M,R,I,T}
    currmem::I
    curridx::I
    s::T
    y::T
    s_M::Vector{T}
    y_M::Vector{T}
    ys_M::Vector{R}
    alphas::Vector{R}
    H::R
end

function LBFGSOperator{M}(x::T) where {M,T}
    R = real(eltype(x))
    s_M = [zero(x) for i = 1:M]
    y_M = [zero(x) for i = 1:M]
    s = zero(x)
    y = zero(x)
    ys_M = zeros(R, M)
    alphas = zeros(R, M)
    LBFGSOperator{M,R,typeof(0),T}(0, 0, s, y, s_M, y_M, ys_M, alphas, one(R))
end

LBFGSOperator(M, x) = LBFGSOperator{M}(x)

function update!(L::LBFGSOperator{M}, s, y) where M
    L.s .= s
    L.y .= y
    ys = real(dot(L.s, L.y))
    if ys > 0
        L.curridx += 1
        if L.curridx > M
            L.curridx = 1
        end
        L.currmem += 1
        if L.currmem > M
            L.currmem = M
        end
        L.ys_M[L.curridx] = ys
        copyto!(L.s_M[L.curridx], L.s)
        copyto!(L.y_M[L.curridx], L.y)
        yty = real(dot(L.y, L.y))
        L.H = ys / yty
    end
    return L
end

function reset!(L::LBFGSOperator{M,R,I}) where {M,R,I}
    L.currmem, L.curridx = zero(I), zero(I)
    L.H = one(R)
end

function (*)(L::LBFGSOperator, v)
    w = similar(v)
    mul!(w, L, v)
end

# Two-loop recursion

function mul!(d, L::LBFGSOperator, v)
    d .= v
    idx = loop1!(d, L)
    d .*= L.H
    loop2!(d, idx, L)
    return d
end

function loop1!(d, L::LBFGSOperator{M}) where M
    idx = L.curridx
    for i = 1:L.currmem
        L.alphas[idx] = real(dot(L.s_M[idx], d)) / L.ys_M[idx]
        d .-= L.alphas[idx] .* L.y_M[idx]
        idx -= 1
        if idx == 0
            idx = M
        end
    end
    return idx
end

function loop2!(d, idx, L::LBFGSOperator{M}) where M
    for i = 1:L.currmem
        idx += 1
        if idx > M
            idx = 1
        end
        beta = real(dot(L.y_M[idx], d)) / L.ys_M[idx]
        d .+= (L.alphas[idx] - beta) .* L.s_M[idx]
    end
    return d
end

struct LBFGS{M} end

LBFGS(M) = LBFGS{M}()

acceleration_style(::Type{<:LBFGS}) = QuasiNewtonStyle()

function initialize(::LBFGS{M}, x) where M
    return LBFGSOperator{M}(x)
end


text:
@symbol_func function r_le(cur_reactor::AbstractReactor, cur_rho::AbstractSymbol)
  cur_T = T_rho(cur_reactor, cur_rho)

  cur_r = 1.07e-4

  cur_r *= sqrt(cur_T)

  cur_r /= cur_reactor.B_0

  cur_r
end


text:
module ArchGDAL

    using Compat
    import Compat.String
    import GDAL, GeoInterface
    import DataStreams: Data

    include("utils.jl")
    include("types.jl")
    include("driver.jl")
    include("gcp.jl")
    include("spatialref.jl")
    include("dataset.jl")
    include("raster/rasterband.jl")
    include("raster/rasterio.jl")
    include("raster/rasterattributetable.jl")
    include("raster/colortable.jl")
    include("ogr/geometry.jl")
    include("ogr/feature.jl")
    include("ogr/featurelayer.jl")
    include("ogr/featuredefn.jl")
    include("ogr/fielddefn.jl")
    include("ogr/styletable.jl")
    include("context.jl")
    include("base/iterators.jl")
    include("base/display.jl")
    include("datastreams.jl")
    include("geointerface.jl")

end # module


text:
# some convenience functions
function structname(equ)
    if occursin('.', equ)
        return equ[findlast(isequal('.'), equ)+1:end]
    else
        return equ
    end
end

teststring(equ) = structname(string(equ))
teststring(equ, pert) = structname(string(equ)) * " + " * structname(string(pert))

function test_equilibrium(equ_mod, t, x)
    @test equ_mod.X(x...) == equ_mod.X(t,x)
    @test equ_mod.Y(x...) == equ_mod.Y(t,x)
    @test equ_mod.Z(x...) == equ_mod.Z(t,x)

    @test equ_mod.J(x...) == equ_mod.J(t,x)
    @test equ_mod.B(x...) == equ_mod.B(t,x)
    @test equ_mod.φ(x...) == equ_mod.φ(t,x)

    @test equ_mod.A₁(x...) == equ_mod.A₁(t,x)
    @test equ_mod.A₂(x...) == equ_mod.A₂(t,x)
    @test equ_mod.A₃(x...) == equ_mod.A₃(t,x)

    @test equ_mod.B₁(x...) == equ_mod.B₁(t,x)
    @test equ_mod.B₂(x...) == equ_mod.B₂(t,x)
    @test equ_mod.B₃(x...) == equ_mod.B₃(t,x)

    @test equ_mod.b₁(x...) == equ_mod.b₁(t,x)
    @test equ_mod.b₂(x...) == equ_mod.b₂(t,x)
    @test equ_mod.b₃(x...) == equ_mod.b₃(t,x)

    @test equ_mod.E₁(x...) == equ_mod.E₁(t,x)
    @test equ_mod.E₂(x...) == equ_mod.E₂(t,x)
    @test equ_mod.E₃(x...) == equ_mod.E₃(t,x)

    @test equ_mod.A¹(x...) == equ_mod.A¹(t,x)
    @test equ_mod.A²(x...) == equ_mod.A²(t,x)
    @test equ_mod.A³(x...) == equ_mod.A³(t,x)

    @test equ_mod.B¹(x...) == equ_mod.B¹(t,x)
    @test equ_mod.B²(x...) == equ_mod.B²(t,x)
    @test equ_mod.B³(x...) == equ_mod.B³(t,x)

    @test equ_mod.b¹(x...) == equ_mod.b¹(t,x)
    @test equ_mod.b²(x...) == equ_mod.b²(t,x)
    @test equ_mod.b³(x...) == equ_mod.b³(t,x)

    @test equ_mod.E¹(x...) == equ_mod.E¹(t,x)
    @test equ_mod.E²(x...) == equ_mod.E²(t,x)
    @test equ_mod.E³(x...) == equ_mod.E³(t,x)

    @test equ_mod.dA₁dx₁(x...) == equ_mod.dA₁dx₁(t,x)
    @test equ_mod.dA₁dx₂(x...) == equ_mod.dA₁dx₂(t,x)
    @test equ_mod.dA₁dx₃(x...) == equ_mod.dA₁dx₃(t,x)

    @test equ_mod.dA₂dx₁(x...) == equ_mod.dA₂dx₁(t,x)
    @test equ_mod.dA₂dx₂(x...) == equ_mod.dA₂dx₂(t,x)
    @test equ_mod.dA₂dx₃(x...) == equ_mod.dA₂dx₃(t,x)

    @test equ_mod.dA₃dx₁(x...) == equ_mod.dA₃dx₁(t,x)
    @test equ_mod.dA₃dx₂(x...) == equ_mod.dA₃dx₂(t,x)
    @test equ_mod.dA₃dx₃(x...) == equ_mod.dA₃dx₃(t,x)
end

# modules that will hold the generated equilibrium code
module AxisymmetricTokamakCartesianEquilibrium end
module AxisymmetricTokamakCylindricalEquilibrium end
module AxisymmetricTokamakToroidalEquilibrium end
module SymmetricQuadraticEquilibrium end
module ThetaPinchEquilibrium end
module ABCEquilibrium end
module SolovevEquilibrium end
module SolovevXpointEquilibrium end
module SolovevQuadraticEquilibrium end

module SymmetricQuadraticEquilibriumEzCosZPerturbation end
module ThetaPinchEquilibriumEzCosZPerturbation end

# equilibrium list (equilibrium, parameters, periodicity, module)
eqs = (
    (AxisymmetricTokamakCartesian,      (2., 3., 2.),   [0., 0., 0.],   AxisymmetricTokamakCartesianEquilibrium),
    (AxisymmetricTokamakCylindrical,    (2., 3., 2.),   [0., 0., 2π],   AxisymmetricTokamakCylindricalEquilibrium),
    (AxisymmetricTokamakToroidal,       (2., 3., 2.),   [0., 2π, 2π],   AxisymmetricTokamakToroidalEquilibrium),
    (SymmetricQuadratic,                (1.),           [0., 0., 0.],   SymmetricQuadraticEquilibrium),
    (ThetaPinch,                        (1.),           [0., 0., 0.],   ThetaPinchEquilibrium),
    (ABC,                               (1., 0.5, 0.5), [0., 0., 0.],   ABCEquilibrium),
    (Solovev,           (6.2, 5.3, 0.32, 1.8, 0.45, -0.155),                [0., 0., 2π],   SolovevEquilibrium),
    (SolovevXpoint,     (6.2, 5.3, 0.32, 1.8, 0.45, -0.155, 0.88, -0.60),   [0., 0., 2π],   SolovevXpointEquilibrium),
    (SolovevQuadratic,  (6.2, 5.3, 1., 1.),                                 [0., 0., 2π],   SolovevQuadraticEquilibrium),
)

# perturbation list (equilibrium, parameters, perturbation, parameters, module)
perts = (
    (SymmetricQuadratic,    (1.),   EzCosZ,    (2.),   SymmetricQuadraticEquilibriumEzCosZPerturbation),
    (ThetaPinch,            (1.),   EzCosZ,    (2.),   ThetaPinchEquilibriumEzCosZPerturbation),
)

# testing parameters
t = 1.
x = [1., 1., 1.]

# test equilibria
@testset "$(rpad(teststring(equ[1]),60))" for equ in eqs begin
        equ_obj = equ[1](equ[2]...)
        load_equilibrium(equ_obj, target_module=equ[4])
        test_equilibrium(equ[4], t, x)
        @test periodicity(x, equ_obj) == equ[3]
    end
end
println()

# test perturbations
@testset "$(rpad(teststring(equ[1], equ[3]),60))" for equ in perts begin
        equ_obj = equ[1](equ[2]...)
        prt_obj = equ[3](equ[4]...)
        load_equilibrium(equ_obj, prt_obj, target_module=equ[5])
        test_equilibrium(equ[5], t, x)
    end
end
println()

# test correctness of some of the magnetic fields

function test_axisymmetric_tokamak_cartesian_equilibrium(equ_mod, t=0., x=[1.5, 0.0, 0.5])
    @test equ_mod.B¹(t,x) == equ_mod.B₁(t,x)
    @test equ_mod.B²(t,x) == equ_mod.B₂(t,x)
    @test equ_mod.B³(t,x) == equ_mod.B₃(t,x)

    @test equ_mod.B₁(t,x) == - equ_mod.B₀ / equ_mod.q₀ * (equ_mod.q₀ * equ_mod.R₀ * equ_mod.Y(t,x) + equ_mod.X(t,x) * equ_mod.Z(t,x) ) / equ_mod.R(t,x)^2
    @test equ_mod.B₂(t,x) == + equ_mod.B₀ / equ_mod.q₀ * (equ_mod.q₀ * equ_mod.R₀ * equ_mod.X(t,x) - equ_mod.Y(t,x) * equ_mod.Z(t,x) ) / equ_mod.R(t,x)^2
    @test equ_mod.B₃(t,x) == + equ_mod.B₀ / equ_mod.q₀ * (equ_mod.R(t,x) - equ_mod.R₀) / equ_mod.R(t,x)
end

function test_axisymmetric_tokamak_cylindrical_equilibrium(equ_mod, t=0., x=[1.5, 0.5, π/5])
    @test equ_mod.B¹(t,x) == - equ_mod.B₀ / equ_mod.q₀ * equ_mod.Z(t,x) / equ_mod.R(t,x)
    @test equ_mod.B²(t,x) == + equ_mod.B₀ / equ_mod.q₀ * (equ_mod.R(t,x) - equ_mod.R₀) / equ_mod.R(t,x)
    @test equ_mod.B³(t,x) == - equ_mod.B₀ * equ_mod.R₀ / equ_mod.R(t,x)^2

    @test equ_mod.B₁(t,x) == - equ_mod.B₀ / equ_mod.q₀ * equ_mod.Z(t,x) / equ_mod.R(t,x)
    @test equ_mod.B₂(t,x) == + equ_mod.B₀ / equ_mod.q₀ * (equ_mod.R(t,x) - equ_mod.R₀) / equ_mod.R(t,x)
    @test equ_mod.B₃(t,x) == - equ_mod.B₀ * equ_mod.R₀
end

function test_axisymmetric_tokamak_toroidal_equilibrium(equ_mod, t=0., x=[0.5, π/10, π/5])
    @test equ_mod.B¹(t,x) == 0
    @test equ_mod.B²(t,x) == equ_mod.B₀ / equ_mod.q₀ / equ_mod.R(t,x)
    @test equ_mod.B³(t,x) == equ_mod.B₀ * equ_mod.R₀ / equ_mod.R(t,x)^2

    @test equ_mod.B₁(t,x) == 0
    @test equ_mod.B₂(t,x) == equ_mod.B₀ / equ_mod.q₀ * equ_mod.r(t,x)^2 / equ_mod.R(t,x)
    @test equ_mod.B₃(t,x) == equ_mod.B₀ * equ_mod.R₀
end

function test_symmetric_quadratic_equilibrium(equ_mod, t=0., x=[1.0, 0.5, 0.5])
    @test equ_mod.B¹(t,x) == equ_mod.B₁(t,x)
    @test equ_mod.B²(t,x) == equ_mod.B₂(t,x)
    @test equ_mod.B³(t,x) == equ_mod.B₃(t,x)

    @test equ_mod.B₁(t,x) == 0
    @test equ_mod.B₂(t,x) == 0
    @test equ_mod.B₃(t,x) == equ_mod.B₀ * (1 + equ_mod.X(t,x)^2 + equ_mod.Y(t,x)^2)

    @test equ_mod.B(t,x)  == equ_mod.B₀ * (1 + equ_mod.X(t,x)^2 + equ_mod.Y(t,x)^2)

    @test equ_mod.b¹(t,x) == 0
    @test equ_mod.b²(t,x) == 0
    @test equ_mod.b³(t,x) == 1

    @test equ_mod.b₁(t,x) == 0
    @test equ_mod.b₂(t,x) == 0
    @test equ_mod.b₃(t,x) == 1
end

function test_theta_pinch_equilibrium(equ_mod, t=0., x=[1.0, 0.5, 0.5])
    @test equ_mod.B¹(t,x) == equ_mod.B₁(t,x)
    @test equ_mod.B²(t,x) == equ_mod.B₂(t,x)
    @test equ_mod.B³(t,x) == equ_mod.B₃(t,x)

    @test equ_mod.B₁(t,x) == 0
    @test equ_mod.B₂(t,x) == 0
    @test equ_mod.B₃(t,x) == equ_mod.B₀

    @test equ_mod.B(t,x)  == equ_mod.B₀

    @test equ_mod.b¹(t,x) == 0
    @test equ_mod.b²(t,x) == 0
    @test equ_mod.b³(t,x) == 1

    @test equ_mod.b₁(t,x) == 0
    @test equ_mod.b₂(t,x) == 0
    @test equ_mod.b₃(t,x) == 1
end

function test_abc_equilibrium(equ_mod, t=0., x=[1.0, 0.5, 0.5])
    @test equ_mod.B¹(t,x) == equ_mod.B₁(t,x)
    @test equ_mod.B²(t,x) == equ_mod.B₂(t,x)
    @test equ_mod.B³(t,x) == equ_mod.B₃(t,x)

    @test equ_mod.B₁(t,x) == equ_mod.A₁(t,x)
    @test equ_mod.B₂(t,x) == equ_mod.A₂(t,x)
    @test equ_mod.B₃(t,x) == equ_mod.A₃(t,x)

    @test equ_mod.B₁(t,x) == equ_mod.a * sin(equ_mod.Z(t,x)) + equ_mod.c * cos(equ_mod.Y(t,x))
    @test equ_mod.B₂(t,x) == equ_mod.b * sin(equ_mod.X(t,x)) + equ_mod.a * cos(equ_mod.Z(t,x))
    @test equ_mod.B₃(t,x) == equ_mod.c * sin(equ_mod.Y(t,x)) + equ_mod.b * cos(equ_mod.X(t,x))
end

@testset "$(rpad("Magnetic Fields",60))" begin
    test_axisymmetric_tokamak_cartesian_equilibrium(AxisymmetricTokamakCartesianEquilibrium)
    test_axisymmetric_tokamak_cylindrical_equilibrium(AxisymmetricTokamakCylindricalEquilibrium)
    test_axisymmetric_tokamak_toroidal_equilibrium(AxisymmetricTokamakToroidalEquilibrium)
    test_symmetric_quadratic_equilibrium(SymmetricQuadraticEquilibrium)
    test_theta_pinch_equilibrium(ThetaPinchEquilibrium)
    test_abc_equilibrium(ABCEquilibrium)
end
println()


text:
struct Vertex
    name::Symbol
    inputs::Vector{Symbol}
    outputs::Vector{Symbol}
end

overwrites(vertex::Vertex) = intersect(vertex.inputs, vertex.outputs)

get_vertex_names(vetrices::AbstractArray{Vertex}) = map(vertex -> vertex.name, vetrices)

function simpledigraph(vertices::AbstractArray{Vertex})
    N = length(vertices)
    g = SimpleDiGraph(N)
    for i in 1:N, j in 1:N
        vi, vj = vertices[i], vertices[j]
        isdisjoint(vi.outputs, vj.inputs) || add_edge!(g, i, j)
    end
    return g
end

# nested case
function simpledigraph(nested_vertices::AbstractArray{<:AbstractArray{Vertex}})

    vertices = reduce(vcat, nested_vertices)
    groups = mapreduce(append!, enumerate(nested_vertices), init=Int[]) do (idx, vertices)
        return fill(idx, length(vertices))
    end

    N = length(vertices)
    g = SimpleDiGraph(N)
    for j in 1:N
        inputs = copy(vertices[j].inputs)
        ii = searchsortedlast(groups, groups[j])
        for i in ii:-1:1
            outputs = vertices[i].outputs
            isdisjoint(outputs, inputs) && continue
            i == j && continue
            setdiff!(inputs, outputs)
            add_edge!(g, i, j)
        end
    end

    return g
end


text:
#== # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#   Tests related to the position of the Moon.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
#
#   [1] Vallado, D. A (2013). Fundamentals of Astrodynamics and Applications.
#       Microcosm Press, Hawthorn, CA, USA.
#   [2] Meeus, J (1998). Astronomical algorithms. Willmann-Bell, Inc, Richmond,
#       VA.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # ==#

# File: ./src/sun/moon_position.jl
# ================================

# Function moon_position_i
# ------------------------

################################################################################
#                                 Test Results
################################################################################
#
# Example 47.a: Calculate the geocentric longitude, latitude, and equatorial
#               horizontal parallax of the Moon for 1992 April 12, at 0h TD.
#
#   In the ecliptic plane w.r.t. the mean equator of date:
#
#       λ = 133°.162655 (longitude)
#       β = -3°.229126 (latitude)
#       Δ = 368409.7 km
#
################################################################################

@testset "Function moon_position_i (Meeus)" begin
    JD_TDB = date_to_jd(1992, 04, 12, 0, 0, 0)

    r_moon_mod = moon_position_i(JD_TDB)

    # We need to compute the mean ecliptic to obtain the variables in the
    # example.
    T_TDB = (JD_TDB - JD_J2000)/36525.0
    ϵ = @evalpoly(T_TDB, 23.439_291, -0.013_004_2, -1.64e-7, +5.04e-7)
    ϵ = mod2pi(deg2rad(ϵ))

    r_moon_ecl = angle_to_dcm(ϵ, :X)*r_moon_mod

    λ = atand(r_moon_ecl[2], r_moon_ecl[1])
    β = atand(r_moon_ecl[3], √(r_moon_ecl[1]^2 + r_moon_ecl[2]^2))
    Δ = norm(r_moon_ecl)/1000

    @test λ ≈ 133.162655 atol = 1e-6
    @test β ≈ -3.229126  atol = 1e-6
    @test Δ ≈ 368409.7   atol = 1e-1
end

################################################################################
#                                 Test Results
################################################################################
#
# Example 5-3: Finding the Moon position vector [1, p. 289]
#
#   In April 28, 1994, 00:00 UTC, one gets:
#
#             | -134_240.626 |
#   r_J2000 = | -311_571.590 | km
#             | -126_693.785 |
#
#   Note: it assumes that UTC ≈ UT1 ≈ TDB.
#
################################################################################

@testset "Function moon_position_i (Vallado)" begin
    JD_UTC = date_to_jd(1994, 4, 28, 0, 0, 0)
    JD_TBD = JD_UTC

    r_J2000 = moon_position_i(JD_TBD, Val(:Vallado))

    @test r_J2000[1] / 1000 ≈ -134_240.626 atol = 1e-3
    @test r_J2000[2] / 1000 ≈ -311_571.590 atol = 1e-3
    @test r_J2000[3] / 1000 ≈ -126_693.785 atol = 1e-3
end


text:
function _pluck(cur_array::AbstractArray, remove_item=false)
  cur_length = length(cur_array)
  ( iszero(cur_length) ) && throw(BoundsError())

  cur_index = rand(1:cur_length)
  cur_value = cur_array[cur_index]

  remove_item && deleteat!(cur_array, cur_index)

  cur_value
end

function pluck(cur_array::AbstractArray)
  _pluck(cur_array, false)
end

function pluck!(cur_array::AbstractArray)
  _pluck(cur_array, true)
end

function pluck(cur_iterator::Union{Base.KeyIterator, Base.ValueIterator})
  pluck(collect(cur_iterator))
end


text:
using PkgSRA, Test, Random
Random.seed!(13)

###
### First, Lorenz
###
include("../examples/example_lorenz.jl")
dat = Array(solve_lorenz_system())
true_grad = core_dyn_true(dat)

# Numerical derivative: the main function to be tested
numerical_grad = numerical_derivative(dat, ts)

# Basic
@test !isempty(numerical_grad)
@test size(numerical_grad) == size(true_grad)

# Accuracy
@test isapprox(numerical_grad, true_grad, rtol=1e-2)
@test !isapprox(numerical_grad, true_grad, rtol=1e-3)

###
### Second, SIR
###
include("../examples/example_sir.jl")
dat = Array(solve_sir_system())
true_grad = core_dyn_true(dat)

# Numerical derivative: the main function to be tested
numerical_grad = numerical_derivative(dat, ts)

# Basic
@test !isempty(numerical_grad)
@test size(numerical_grad) == size(true_grad)

# Accuracy
@test isapprox(numerical_grad, true_grad, rtol=1e-2)
@test !isapprox(numerical_grad, true_grad, rtol=1e-3)

# Check why the above doesn't work...
plot(numerical_grad[1,:], label="numerical");
    plot!(true_grad[1,:], label="true")

dt = ts[2] - ts[1]
bad_grad = (dat[:,2:end] .- dat[:,1:end-1]) / dt
plot(ts, numerical_grad[1,:], label="numerical");
    plot!(ts, bad_grad[1,:], label="one-step")
    plot!(ts, true_grad[1,:], label="true")


text:
"""
    point(xs...)

Create point of default type (double precision coordinates)
"""
point(xs...) = point(Float64, xs...)

"""
    point(type, xs...)

Create point of default type and supplied precision for the coordinates
"""
@generated function point(T::Type,xs...)
    D = length(xs)
    xp = :(SVector{$D,T}())
    for d in 1:D
        push!(xp.args, :(xs[$d]))
    end
    return xp
end

"""
    index(i1, i2, ...) -> ids

Create a tuple of vertex indices.
"""
index(is...) = SVector{length(is),Int}(is...)

"""
  euclidian_basis(type, dim)

Returns the origin and default unit vectors for Euclidian space of dimension dim
"""
function euclidianbasis(dim, T=Float64::Type)
  P = SVector{dim,T}
  #id = eye(dim)
  id = Matrix{T}(I,dim,dim)
  r = P[ P(id[:,i]...) for i in 1:dim ]
  z = P(zeros(T,dim)...)
  return pushfirst!(r, z)
end


text:
#########
# broadcasting
#########

for op in (:+, :-)
    @eval broadcasted(::DefaultArrayStyle{N}, ::typeof($op), r1::AbstractFill{T,N}, r2::AbstractFill{V,N}) where {T,V,N} =
            $op(r1, r2)
end

function broadcasted(::DefaultArrayStyle{N}, ::typeof(*), a::Zeros{T,N}, b::Zeros{V,N}) where {T,V,N}
    axes(a) ≠ axes(b) && throw(DimensionMismatch("dimensions must match."))
    Zeros{promote_type(T,V)}(axes(a))
end

function broadcasted(::DefaultArrayStyle{N}, ::typeof(*), a::AbstractArray{T,N}, b::Zeros{V,N}) where {T,V,N}
    axes(a) ≠ axes(b) && throw(DimensionMismatch("dimensions must match."))
    Zeros{promote_type(T,V)}(axes(a))
end

function broadcasted(::DefaultArrayStyle{N}, ::typeof(*), a::Zeros{T,N}, b::AbstractArray{V,N}) where {T,V,N}
    axes(a) ≠ axes(b) && throw(DimensionMismatch("dimensions must match."))
    Zeros{promote_type(T,V)}(axes(a))
end

broadcasted(::DefaultArrayStyle{N}, op, r::AbstractFill{T,N}) where {T,N} = Fill(op(getindex_value(r)), size(r))
broadcasted(::DefaultArrayStyle{N}, op, r::AbstractFill{T,N}, x::Number) where {T,N} = Fill(op(getindex_value(r),x), size(r))
broadcasted(::DefaultArrayStyle{N}, op, x::Number, r::AbstractFill{T,N}) where {T,N} = Fill(op(x, getindex_value(r)), size(r))
function broadcasted(::DefaultArrayStyle{N}, op, r1::AbstractFill{T,N}, r2::AbstractFill{V,N}) where {T,V,N}
    size(r1) ≠ size(r2) && throw(DimensionMismatch("dimensions must match."))
    Fill(op(getindex_value(r1),getindex_value(r2)), size(r1))
end

for op in (:*, :/, :\)
    @eval function broadcasted(::DefaultArrayStyle{N}, ::typeof($op), r1::Ones{T,N}, r2::Ones{V,N}) where {T,V,N}
        size(r1) ≠ size(r2) && throw(DimensionMismatch("dimensions must match."))
        Ones{promote_type(T,V)}(size(r1))
    end
end


text:
using MPI
using MPIArrays
using Test

MPI.Init()
const comm = MPI.COMM_WORLD
const rank = MPI.Comm_rank(comm)
const nb_procs = MPI.Comm_size(comm)

do_serial = (rank == 0)

include(joinpath(@__DIR__, "matmul_setup.jl"))

function timed_mul(A::MPIArray,x::MPIArray)
    times = zeros(nsamples)
    for i in 1:length(times)
        if rank == 0
            val, times[i], bytes, gctime, memallocs = @timed begin
                b = A*x
                sync(b)
            end
        else
            b = A*x
            sync(b)
        end
        free(b)
    end
    return extrema(times[2:end])
end

@testset "MatMul" begin

A1 = MPIArray{Float64}(comm, (nb_procs,1), N, N)
A2 = MPIArray{Float64}(comm, (1,nb_procs), N, N)
x = MPIArray{Float64}(comm, (nb_procs,), N)

if rank == 0
    putblock!(As,A1[:,:])
    putblock!(As,A2[:,:])
    putblock!(xs,x[:])
end

pres1 = A1*x
pres2 = A2*x

if rank == 0
    @test all(getblock(pres1[:]) .≈ ref)
    @test all(getblock(pres2[:]) .≈ ref)
end

sync(pres2)

global (mintime, maxtime) = timed_mul(A1,x)
rank == 0 && write_timings("out-mpi.txt", 1, mintime, maxtime)
global (mintime, maxtime) = timed_mul(A2,x)
rank == 0 && write_timings("out-mpi.txt", 2, mintime, maxtime)

sync(pres2)

end

MPI.Finalize()


text:
# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    DirectionPartitioner(direction; tol=1e-6)

A method for partitioning spatial objects along a given `direction`
with bandwidth tolerance `tol`.
"""
struct DirectionPartitioner{T,N} <: AbstractSpatialPredicatePartitioner
  direction::SVector{N,T}
  tol::Float64

  function DirectionPartitioner{T,N}(direction, tol) where {N,T}
    new(normalize(direction), tol)
  end
end

DirectionPartitioner(direction::SVector{N,T}; tol=1e-6) where {T,N} =
  DirectionPartitioner{T,N}(direction, tol)

DirectionPartitioner(direction::NTuple{N,T}; tol=1e-6) where {T,N} =
  DirectionPartitioner(SVector(direction), tol=tol)

(p::DirectionPartitioner)(x, y) = begin
  δ = x - y
  d = p.direction
  norm(δ - (δ⋅d)*d) < p.tol
end


text:
##################################################################################
# This file is part of StateSpaceEcon.jl
# BSD 3-Clause License
# Copyright (c) 2020, Bank of Canada
# All rights reserved.
##################################################################################

"""
    SteadyStateSolver

A module that is part of StateSpaceEcon package.
Contains methods for finding a steady state of a model.
"""
module SteadyStateSolver

using SparseArrays
using LinearAlgebra
using ForwardDiff
using DiffResults
using Printf

using ModelBaseEcon

include("steadystate/1dsolvers.jl")
include("steadystate/presolve.jl")
include("steadystate/initial.jl")
include("steadystate/global.jl")
include("steadystate/solverdata.jl")
include("steadystate/nr.jl")
include("steadystate/lm.jl")
include("steadystate/diagnose.jl")
include("steadystate/sssolve.jl")

end # module

using .SteadyStateSolver
export clear_sstate!
export initial_sstate!
export check_sstate
export diagnose_sstate
export sssolve!


text:
using FractalTools
using Makie
using GeometryBasics

# Construct a mesh
pts = [
    [0., 0.],
    [1., 0.],
    [1., 1.]
]
msh = trimesh(pts...)
trange = 1 : length(faces(msh))

# Plot mesh
fig = Figure()
ax, plt = mesh(fig[1, 1], msh, color=last.(msh.position))
wireframe!(ax, msh)
cb = Colorbar(fig[1, 2], plt, width=20)
button1 = Button(fig[2, 1], tellwidth=false, label="Increase")
button2 = Button(fig[2, 2], tellwidth=false, label="Decrease")
label = Label(fig[3, :], tellwidth=false)

# Plot a triangle
coords = coordinates(msh)
trigs = faces(msh)
idx = Node(1)
x = @lift(getindex.(coords[trigs[$idx]], 1) |> collect)
y = @lift(getindex.(coords[trigs[$idx]], 2) |> collect)
scatter!(ax, x, y)

on(button1.clicks) do n
    idx[] += 1
    label.text = "idx = $(idx[]), x = $(x[]), y = $(y[])"
end

on(button2.clicks) do n
    idx[] -= 1
    label.text = "idx = $(idx[]), x = $(x[]), y = $(y[])"
end

# Bind keybord interacton
scene = fig.scene
on(scene.events.keyboardbuttons) do button
    if ispressed(button, Keyboard.up)
        idx[] += 1
        label.text = "idx = $(idx[]), x = $(x[]), y = $(y[])"
    end
    if ispressed(button, Keyboard.down)
        idx[] -= 1
        label.text = "idx = $(idx[]), x = $(x[]), y = $(y[])"
    end
end

display(fig)


text:
VERSION >= v"0.4.0-dev+6521" && __precompile__()

module TextWrap

using Compat
import Compat: String

export
    wrap,
    print_wrapped,
    println_wrapped

function _expand_tabs(text::AbstractString, i0::Int)
    out_buf = IOBuffer()
    i = i0 % 8
    for c in text
        if c == '\t'
            while i < 8
                print(out_buf, ' ')
                i += 1
            end
            i = 0
        elseif c == '\n' || c == '\r'
            print(out_buf, c)
            i = 0
        else
            print(out_buf, c)
            i = (i + 1) % 8
        end
    end
    return takebuf_string(out_buf)
end

function _check_width(width::Integer)
    if width <= 0
        error("invalid width $width (must be > 0)")
    end
    return true
end
function _check_indent(indent::Integer, width::Integer)
    if indent < 0 || indent >= width
        error("invalid intent $indent (must be an integer between 0 and width-1, or an AbstractString)")
    end
    return true
end
function _check_indent(indent::AbstractString, width::Integer)
    if length(indent) >= width
        error("invalid intent (must be shorter than width-1)")
    end
end

function _put_chunks(chunk::AbstractString, out_str,
                    cln, cll, bol, soh,
                    width, initial_indent, subsequent_indent,
                    break_on_hyphens, break_long_words)

    # This function just performs breaks-on-hyphens and passes
    # individual chunks to _put_chunk

    _hyphen_re = r"^(\w+-)*?\w+[^0-9\W]+-(?=\w+[^0-9\W])"

    while break_on_hyphens
        m = match(_hyphen_re, chunk)
        if m === nothing
            break
        end
        c = m.match
        cln, cll, bol, lcise = _put_chunk(c, out_str,
                    cln, cll, bol, soh,
                    width, initial_indent, subsequent_indent,
                    break_long_words)
        soh = ""
        chunk = chunk[m.offset+endof(c):end]
    end

    cln, cll, bol, lcise = _put_chunk(chunk, out_str,
                cln, cll, bol, soh,
                width, initial_indent, subsequent_indent,
                break_long_words)
    return cln, cll, bol, lcise
end

function _put_chunk(chunk::AbstractString, out_str,
                    cln, cll, bol, soh,
                    width, initial_indent, subsequent_indent,
                    break_long_words)

    # Writes a chunk to out_str, based on the current position
    # as encoded in (cln, cll, bol) = (current_line_number,
    # current_line_length, beginning_of_line), and returns the
    # updated position (plus a flag to signal that an end-of-sentence
    # was detected).
    # The argument soh (=space_on_hold) is the spacing which should
    # go in front of chunk, and it may or may not be printed.
    # The rest are options.

    liindent = length(initial_indent)
    lsindent = length(subsequent_indent)
    lchunk = length(chunk)
    lsoh = length(soh)

    if cll + lsoh > width
        soh = ""
        lsoh = 0
        print(out_str, "\n")
        cln += 1
        cll = 0
        bol = true
    end

    if bol
        if cln == 1
            indent = initial_indent
            lindent = liindent
        else
            indent = subsequent_indent
            lindent = lsindent
        end

        print(out_str, indent)
        cll += lindent

        soh = ""
        lsoh = 0
    end

    # is there enough room for the chunk? or is this the
    # beginning of the text and we cannot break words?
    if cll + lsoh + lchunk <= width ||
            (cln == 1 && bol && !break_long_words)
        print(out_str, soh, chunk)
        cll += lchunk + lsoh
        bol = false
    # does the chunk fit into next line? or are we
    # forced to put it there?
    elseif lchunk <= width - lsindent || !break_long_words
        print(out_str, bol?"":"\n", subsequent_indent, chunk)
        cll = lsindent + lchunk
        cln += 1
        bol = false
    # break it until it fits
    else
        while cll + lsoh + lchunk > width
            if width - cll - lsoh > 0
                print(out_str, soh, chunk[1:chr2ind(chunk, width-cll-lsoh)], "\n",
                        subsequent_indent)
                chunk = chunk[chr2ind(chunk, width-cll-lsoh+1):end]
                lchunk = length(chunk)
            else
                print(out_str, "\n", subsequent_indent)
            end
            cll = lsindent
            cln += 1
            soh = ""
            lsoh = 0
        end
        print(out_str, chunk)
        cll += lchunk
        bol = false
    end

    # detect end-of-sentences
    _sentence_end_re = r"\w([\.\!\?…]|\.\.\.)[\"\'´„]?\Z"
    lcise = ismatch(_sentence_end_re, chunk)

    return cln, cll, bol, lcise
end

@compat function wrap(text::AbstractString;
                      width::Int = 70,
                      initial_indent::Union{Integer,AbstractString} = "",
                      subsequent_indent::Union{Integer,AbstractString} = "",
                      expand_tabs::Bool = true,
                      replace_whitespace::Bool = true,
                      fix_sentence_endings::Bool = false,
                      break_long_words::Bool = true,
                      break_on_hyphens::Bool = true)

    # Reformat the single paragraph in 'text' so it fits in lines of
    # no more than 'opts.width' columns, and return an AbstractString.

    # Sanity checks
    _check_width(width)
    _check_indent(initial_indent, width)
    _check_indent(subsequent_indent, width)

    if isa(initial_indent, Integer)
        initial_indent = " "^initial_indent
    end
    if isa(subsequent_indent, Integer)
        subsequent_indent = " "^subsequent_indent
    end

    # State variables initialization
    cln = 1 # current line number
    cll = 0 # current line length
    bol = true # beginning of line
    lcise = false # last chunk is sentence ending
    soh = "" # space on hold

    # Preliminary conversion to UTF8
    text = @compat String(text)

    # We iterate over the text, looking for whitespace
    # where to split.
    i = start(text)
    l = endof(text)
    out_str = IOBuffer()

    wsrng = search(text, r"\s+", i)
    j = first(wsrng)
    k = last(wsrng) + 1
    while 0 < j <= l
        if i < k
            if i < j
                # This is non-whitespace. We write it out according
                # to the current cursor position and the leading space.
                chunk = text[i:j-1]

                cln, cll, bol, lcise = _put_chunks(chunk, out_str,
                            cln, cll, bol, soh,
                            width, initial_indent, subsequent_indent,
                            break_on_hyphens, break_long_words)
            end
            i = k
        end

        # This is whitespace. We mangle it (expand tabs, fix
        # sentence endings, replace it with single spaces) and
        # then we keep it on hold.
        soh = text[j:k-1]
        if expand_tabs && ismatch(r"\t", soh)
            soh = _expand_tabs(soh, cll)
        end
        if fix_sentence_endings && lcise && soh == " "
            soh = "  "
        end
        if replace_whitespace
            soh = " "^length(soh)
        end

        # Continue the search

        if k <= j; k = nextind(text,j) end
        wsrng = search(text, r"\s+", k)
        j = first(wsrng)
        k = last(wsrng) + 1
    end
    if !done(text,i)
        # Some non-whitespace is left at the end.
        chunk = text[i:end]
        cln, cll, bol = _put_chunks(chunk, out_str,
                    cln, cll, bol, soh,
                    width, initial_indent, subsequent_indent,
                    break_on_hyphens, break_long_words)
    end
    return takebuf_string(out_str)
end

# print functions signature:
#   first arg: IO
#   last arg: Options
#   inbetween: anything printable
#
#   all arguments are optional
#
function _print_wrapped(newline::Bool, args...; kwargs...)
    if !isempty(args) && isa(args[1], IO)
        io = args[1]
        args = args[2:end]
    else
        io = STDOUT
    end

    if !isempty(args)
        ws = wrap(string(args...); kwargs...)
    else
        ws = ""
    end

    if newline
        println(io, ws)
    else
        print(io, ws)
    end
end
print_wrapped(args...; kwargs...) = _print_wrapped(false, args...; kwargs...)
println_wrapped(args...; kwargs...) = _print_wrapped(true, args...; kwargs...)

end # module TextWrap


text:
struct Sample
	β
	#TODO: adaptive sampling
end

function sample_params(b, T)
	β = [b for t = 1:T]
	return Sample(β)
end

struct SampleDynamics <: Constraints
	n
	ineq
	w0
end

function sample_dynamics_constraints(prob, N, M)
	T = prob.nom.T

	# sample mean
	n = prob.mean.model.n * T

	# resample
	for t = 1:T-1
		n += prob.mean.model.n * N
	end

	ineq = (1:0)
	w0 = [zeros(p.model.n) for p in prob.sample]

	return SampleDynamics(n, ineq, w0)
end

function constraints!(c, Z, con::SampleDynamics,
	prob::DPOProblems, idx::DPOIndices,
	N, D,
	dist,
	sample)

	T = prob.nom.T
	h = prob.nom.h
	M = N + D

	con_shift = 0

	# sample mean
	for t = 1:T

		# mean
		μ = view(Z, idx.mean[prob.mean.idx.x[t]])

		# samples
		x = [view(Z, idx.sample[i][prob.sample[i].idx.x[t]]) for i = 1:N]

		# sample mean
		c[con_shift .+ (1:prob.mean.model.n)] = μ - sample_mean(x)
		con_shift += prob.mean.model.n
	end

	# resample
	n_resample = N * prob.mean.model.n

	for t = 1:T-1
		xt = view(Z, idx.xt[t])
		ut = view(Z, idx.ut[t])
		μ = view(Z, idx.mean[prob.mean.idx.x[t]])
		ν = view(Z, idx.mean[prob.mean.idx.u[t]])

		# samples
		xt⁺ = view(Z, idx.xt[t + 1])

		c[con_shift .+ (1:n_resample)] = sample_dynamics(prob.mean.model,
			xt, ut, μ, ν, dist.w, h, t, sample.β)[1] - xt⁺
		con_shift += n_resample
	end

    return nothing
end

function constraints_jacobian!(∇c, Z, con::SampleDynamics,
	prob::DPOProblems, idx::DPOIndices,
	N, D,
	dist,
	sample)

	T = prob.nom.T
	h = prob.nom.h
	M = N + D

	con_shift = 0
	jac_shift = 0

	# sample mean
	for t = 1:T
		# samples
		x_vec = vcat([view(Z, idx.sample[i][prob.sample[i].idx.x[t]]) for i = 1:N]...)
		sample_mean_vec(y) = sample_mean([view(y, (i - 1) * prob.mean.model.n
			.+ (1:prob.mean.model.n)) for i = 1:N])

		r_idx = con_shift .+ (1:prob.mean.model.n)

		c_idx = idx.mean[prob.mean.idx.x[t]]
		len = length(r_idx) * length(c_idx)
		∇c[jac_shift .+ (1:len)] = vec(Diagonal(ones(prob.mean.model.n)))
		jac_shift += len

		c_idx = vcat([idx.sample[i][prob.sample[i].idx.x[t]] for i = 1:N]...)
		len = length(r_idx) * length(c_idx)
		∇c[jac_shift .+ (1:len)] = vec(-1.0 * ForwardDiff.jacobian(sample_mean_vec, x_vec))
		jac_shift += len

		con_shift += prob.mean.model.n
	end

	# resample
	n_resample = N * prob.mean.model.n
	for t = 1:T-1
		xt = view(Z, idx.xt[t])
		ut = view(Z, idx.ut[t])
		μ = view(Z, idx.mean[prob.mean.idx.x[t]])
		ν = view(Z, idx.mean[prob.mean.idx.u[t]])

		# samples
		# xt⁺ = view(Z, idx.xt[t + 1])

		a1, a2, a3, a4 = sample_dynamics_jacobian(model, xt, ut, μ, ν,
			dist.w, h, t, sample.β)

		r_idx = con_shift .+ (1:n_resample)

		c_idx = idx.xt[t]
		len = length(r_idx) * length(c_idx)
		∇c[jac_shift .+ (1:len)] = vec(a1)
		jac_shift += len

		c_idx = idx.ut[t]
		len = length(r_idx) * length(c_idx)
		∇c[jac_shift .+ (1:len)] = vec(a2)
		jac_shift += len

		c_idx = idx.mean[prob.mean.idx.x[t]]
		len = length(r_idx) * length(c_idx)
		∇c[jac_shift .+ (1:len)] = vec(a3)
		jac_shift += len

		c_idx = idx.mean[prob.mean.idx.u[t]]
		len = length(r_idx) * length(c_idx)
		∇c[jac_shift .+ (1:len)] = vec(a4)
		jac_shift += len

		c_idx = idx.xt[t + 1]
		len = length(r_idx) * length(c_idx)
		∇c[jac_shift .+ (1:len)] = vec(-1.0 * Diagonal(ones(n_resample)))
		jac_shift += len
		con_shift += n_resample
	end

	return nothing
end

function constraints_sparsity(con::SampleDynamics,
	prob::DPOProblems, idx::DPOIndices,
	N, D;
	shift_row = 0, shift_col = 0)

	row = []
    col = []

	T = prob.nom.T
	h = prob.nom.h
	M = N + D

	con_shift = 0

	# sample mean
	for t = 1:T
		# samples
		r_idx = shift_row + con_shift .+ (1:prob.mean.model.n)

		c_idx = shift_col .+ idx.mean[prob.mean.idx.x[t]]
		row_col!(row, col, r_idx, c_idx)

		c_idx = shift_col .+ vcat([idx.sample[i][prob.sample[i].idx.x[t]] for i = 1:N]...)
		row_col!(row, col, r_idx, c_idx)

		con_shift += prob.mean.model.n
	end

	# resample
	n_resample = N * prob.mean.model.n
	for t = 1:T-1
		r_idx = shift_row + con_shift .+ (1:n_resample)

		c_idx = shift_col .+ idx.xt[t]
		row_col!(row, col, r_idx, c_idx)

		c_idx = shift_col .+ idx.ut[t]
		row_col!(row, col, r_idx, c_idx)

		c_idx = shift_col .+ idx.mean[prob.mean.idx.x[t]]
		row_col!(row, col, r_idx, c_idx)

		c_idx = shift_col .+ idx.mean[prob.mean.idx.u[t]]
		row_col!(row, col, r_idx, c_idx)

		c_idx = shift_col .+ idx.xt[t + 1]
		row_col!(row, col, r_idx, c_idx)

		con_shift += n_resample
	end

    return collect(zip(row, col))
end

function sample_dynamics(model, xt, ut, μ, ν, w, h, t, β)
	N = 2 * model.n
	D = 2 * model.d
	M = N + D

	w0 = zeros(model.d)

	# propagate samples
	s = zeros(model.n * M)

	for j = 1:M
		if j <= N
			xi = view(xt, (j - 1) * model.n .+ (1:model.n))
			ui = view(ut, (j - 1) * model.m .+ (1:model.m))
			s[(j - 1) * model.n .+ (1:model.n)] = propagate_dynamics(model,
				xi, ui, w0, h, t)
		else
			k = j - N
			s[(j - 1) * model.n .+ (1:model.n)] = propagate_dynamics(model,
				μ, ν,
				β[t] * w[t][k], h, t)
		end
	end

	# resample
	xt⁺ = resample_vec(s, model.n, M, β[t + 1])

	return xt⁺, s
end

function sample_dynamics_jacobian(model, xt, ut, μ, ν, w, h, t, β)
	N = 2 * model.n
	D = 2 * model.d
	M = N + D

	w0 = zeros(model.d)

	dx⁺dxt = zeros(model.n * M, model.n * N)
	dx⁺dut = zeros(model.n * M, model.m * N)
	dx⁺dμ = zeros(model.n * M, model.n)
	dx⁺dν = zeros(model.n * M, model.m)

	dsdxt = zeros(model.n * M, model.n * N)
	dsdut = zeros(model.n * M, model.m * N)
	dsdμ = zeros(model.n * M, model.n)
	dsdν = zeros(model.n * M, model.m)

	xt⁺, s = sample_dynamics(model, xt, ut, μ, ν, w, h, t, β)
	r(y) = resample_vec(y, model.n, M, β[t + 1])
	dx⁺ds = real.(FiniteDiff.finite_difference_jacobian(r, s))

	# s = []
	# A = []
	# B = []

	for j = 1:M
		if j <= N
			xi = view(xt, (j - 1) * model.n .+ (1:model.n))
			ui = view(ut, (j - 1) * model.m .+ (1:model.m))
			_, _A, _B = propagate_dynamics_jacobian(model, xi, ui,
				w0, h, t)
			# push!(s, _s)
			# push!(A, _A)
			# push!(B, _B)

			dsdxt[(j - 1) * model.n .+ (1:model.n),
				(j - 1) * model.n .+ (1:model.n)] = _A
			dsdut[(j - 1) * model.n .+ (1:model.n),
				(j - 1) * model.m .+ (1:model.m)] = _B
		else
			k = j - N
			_, _A, _B = propagate_dynamics_jacobian(model, μ, ν,
				β[t] * w[t][k], h, t)

			# push!(s, _s)
			# push!(A, _A)
			# push!(B, _B)

			dsdμ[(j - 1) * model.n .+ (1:model.n), :] = _A
			dsdν[(j - 1) * model.n .+ (1:model.n), :] = _B
		end
	end

	dx⁺dxt = dx⁺ds * dsdxt
	dx⁺dut = dx⁺ds * dsdut
	dx⁺dμ = dx⁺ds * dsdμ
	dx⁺dν = dx⁺ds * dsdν

	return dx⁺dxt, dx⁺dut, dx⁺dμ, dx⁺dν
end

max_stars_repo_name: zacmanchester/DirectMotionPlanning

max_issues_repo_name: zacmanchester/DirectMotionPlanning

max_forks_repo_name: zacmanchester/DirectMotionPlanning


text:
@testset "Selection" begin

    # Helper methods
    nt(name) = t -> NamedTrait(t, name)
    tf(name) = t -> nt(name)(t)
    iv(size, name="in") = inputvertex(name, size)
    av(in, outsize, name) = absorbvertex(MatMul(nout(in), outsize), outsize, in, traitdecoration=tf(name))
    tv(in, name) = invariantvertex(identity, in, traitdecoration=tf(name))

    cc(ins...; name) = conc(ins...; dims=2, traitdecoration=tf(name))
    nc(name) = traitconf(nt(name))

    select_outputs_and_change(v, values) = select_outputs_and_change(NoutExact(), v, values)
    function select_outputs_and_change(s, v, values)
        execute, selected = select_outputs(s, v, values)
        if execute
            Δnout(v, selected)
        end
    end

    @testset "SelectDirection" begin

        mutable struct TestProbe <: AbstractSelectionStrategy
            vs
            function TestProbe(v)
                tp = new(nothing)
                Δoutputs(SelectDirection(tp), v, v -> error("Shall not be called!"))
                return tp
            end
        end
        NaiveNASlib.Δoutputs(s::TestProbe, vs::AbstractVector{<:AbstractVertex}, vfun::Function) = s.vs = Set(vs)

        v1 = av(iv(3), 5, "v1")
        v2 = av(v1, 4, "v2")
        v3 = av(v2, 3, "v3")

        tp = TestProbe(v2)
        @test tp.vs == nothing

        Δnin(v2, -1)
        tp = TestProbe(v2)
        @test tp.vs == Set([v2, v1])

        Δnout(v2, 1)
        tp = TestProbe(v2)
        @test tp.vs == Set([v1, v2, v3])

        Δnin(v2, 1)
        tp = TestProbe(v2)
        @test tp.vs == Set([v2, v3])

    end

    @testset "ApplyAfter" begin
        v1 = av(iv(3), 5, "v1")
        v2 = av(v1, 4, "v2")
        v3 = av(v2, 3, "v3")

        res = []
        apply(v) = push!(res, v)

        Δnout(v2, -1)
        @test Δoutputs(ApplyAfter(apply, OutSelectExact()), v2, v -> 1:nout_org(v))

        @test res == [v2, v3]
    end

    @testset "Absorb 2 Absorb" begin
        inpt = iv(3)
        v1 = av(inpt, 5, "v1")
        v2 = av(v1, 4, "v2")

        g = CompGraph(inpt, v2)

        Δnout(v1, -2)
        @test Δoutputs(v1, v -> 1:nout_org(v))

        @test out_inds(op(v1)) == in_inds(op(v2))[] == [3,4,5]
        apply_mutation(g)

        @test size(g(ones(1, 3))) == (1, nout(v2))

        Δnout(v1, 3)
        @test Δoutputs(Output(), v1, v->1:nout_org(v))

        @test out_inds(op(v1)) == in_inds(op(v2))[] == [1,2,3,-1,-1,-1]
        apply_mutation(g)

        @test size(g(ones(1, 3))) == (1, nout(v2))
    end

    @testset "Absorb 2 Absorb revert" begin
        inpt = iv(3)
        v1 = av(inpt, 5, "v1")
        v2 = av(v1, 4, "v2")

        g = CompGraph(inpt, v2)

        Δnout(v1, -2)
        @test !Δoutputs(NoutRevert(), v1, v->1:nout_org(v))

        @test out_inds(op(v1)) == in_inds(op(v2))[] == [1,2,3,4,5]
        apply_mutation(g)

        @test size(g(ones(1, 3))) == (1, nout(v2))

        Δnout(v1, +3)

        @test !Δoutputs(NoutRevert(), v1, v-> 1:nout_org(v))

        @test out_inds(op(v1)) == in_inds(op(v2))[] == [1,2,3,4,5]
        apply_mutation(g)

        @test size(g(ones(1, 3))) == (1, nout(v2))
    end

    @testset "Absorb 2 Absorb fail" begin
        inpt = iv(3)
        v1 = av(inpt, 5, "v1")
        v2 = av(v1, 4, "v2")

        Δnout(v1, -2)
        @test_throws ErrorException Δoutputs(SelectionFail(), v1, v -> 1:nout_org(v))
    end

    @testset "SizeStack duplicate" begin
        inpt = iv(3)
        v1 = av(inpt, 7, "v1")
        v2 = av(inpt, 4, "v2")
        v3 = cc(v2, v1, name="v3")
        v4 = cc(v3, v2, name="v4")

        g = CompGraph(inpt, v4)
        @test size(g(ones(1, 3))) == (1, nout(v4))

        @test minΔnoutfactor(v4) == 2
        Δnout(v4, -4)

        @test nout(v1) == 5
        @test nout(v2) == 3

        @test Δoutputs(v4, v -> 1:nout_org(v))
        apply_mutation(g)

        @test nout(v1) == 5
        @test nout(v2) == 3

        @test size(g(ones(1, 3))) == (1, nout(v4))

        Δnout(v4, 6)

        @test nout(v1) == 7
        @test nout(v2) == 5

        @test Δoutputs(Output(), v4, v->1:nout_org(v))
        apply_mutation(g)

        @test nout(v1) == 7
        @test nout(v2) == 5

        @test size(g(ones(1, 3))) == (1, nout(v4))
    end

    @testset "SizeInvariant duplicate" begin
        inpt = iv(3)
        v1 = av(inpt, 7, "v1")
        v2 = av(inpt, 4, "v2")
        v3 = av(inpt, 3, "v3")

        v4 = cc(v1, v2, name="v4")
        v5 = cc(v2, v3, v2, name="v5")

        v6 = nc("v6") >> v4 + v5

        g = CompGraph(inpt, v6)
        @test size(g(ones(1, 3))) == (1, nout(v6))

        @test minΔnoutfactor(v6) == 2
        Δnout(v6, -4)

        @test nout(v1) == 5
        @test nout(v2) == 2
        @test nout(v3) == 3

        @test Δoutputs(v6, v->1:nout_org(v))
        apply_mutation(g)

        @test nout(v1) == 5
        @test nout(v2) == 2
        @test nout(v3) == 3

        @test size(g(ones(1, 3))) == (1, nout(v6))

        Δnout(v6, 6)

        @test nout(v1) == 9
        @test nout(v2) == 4
        @test nout(v3) == 5

        @test Δoutputs(Output(), v6, v->1:nout_org(v))
        apply_mutation(g)

        @test nout(v1) == 9
        @test nout(v2) == 4
        @test nout(v3) == 5

        @test size(g(ones(1, 3))) == (1, nout(v6))
    end

    @testset "SizeStack one immutable" begin
        inpt = iv(3)
        v1 = av(inpt, 5, "v1")
        v2 = cc(inpt, v1, name="v2")

        g = CompGraph(inpt, v2)
        @test size(g(ones(1, 3))) == (1, nout(v2))

        Δnout(v1, -3)

        @test nin(v2) == [nout(inpt), nout(v1)] == [3, 2]
        @test nout(v2) == 5

        # "Tempt" optimizer to not select inputs from inpt
        @test Δoutputs(OutSelect{NaiveNASlib.Relaxed}(SelectionFail()), v2, v -> v == v2 ? (-nout(inpt):nout_org(v1)-1) : 1:nout_org(v))
        apply_mutation(g)

        @test nin(v2) == [nout(inpt), nout(v1)] == [3, 2]
        @test nout(v2) == 5

        @test size(g(ones(1, 3))) == (1, nout(v2))
    end

    @testset "SizeInvariant exact infeasible" begin
        inpt = iv(3)
        v1 = av(inpt, 10, "v1")
        v2 = av(inpt, 5, "v2")
        v3 = av(inpt, 10, "v3")
        v4 = av(inpt, 5, "v4")

        v5 = cc(v1, v2, v3, name="v5")
        v6 = cc(v2, v1, v2, v4, name="v6")

        v7 = nc("v7") >> v5 + v6

        g = CompGraph(inpt, v7)
        @test size(g(ones(1, 3))) == (1, nout(v7))

        @test minΔnoutfactor(v7) == 2
        Δnout(v7, -6)

        @test nout(v1) == 7
        @test nout(v2) == 4
        @test nout(v3) == 8
        @test nout(v4) == 4

        @test @test_logs (:warn, "Selection for vertex v7 failed! Relaxing size constraint...")  match_mode=:any Δoutputs(v7, v->-1:nout_org(v)-2)
        apply_mutation(g)

        @test nout(v1) == 6
        @test nout(v2) == 3
        @test nout(v3) == 8
        @test nout(v4) == 5

        @test size(g(ones(1, 3))) == (1, nout(v7))

        Δnout(v7, 4)

        @test nout(v1) == 8
        @test nout(v2) == 4
        @test nout(v3) == 9
        @test nout(v4) == 5

        # Works on the first try this time around
        @test Δoutputs(Output(), v7, v->1:nout_org(v))
        apply_mutation(g)

        @test nout(v1) == 8
        @test nout(v2) == 4
        @test nout(v3) == 9
        @test nout(v4) == 5

        @test size(g(ones(1, 3))) == (1, nout(v7))
    end

    @testset "SizeInvariant increase exact infeasible" begin
        inpt = iv(3)
        v1 = av(inpt, 3, "v1")
        v2 = av(inpt, 4, "v2")
        v3 = av(inpt, 4, "v3")
        v4 = av(inpt, 5, "v4")

        v5 = cc(v1, v2, v4, name="v5")
        v6 = cc(v3, v4, v1, name="v6")

        v7 = nc("v7") >> v5 + v6

        g = CompGraph(inpt, v7)
        @test size(g(ones(1, 3))) == (1, nout(v7))

        @test minΔnoutfactor(v7) == 1
        Δnout(v7, 6)

        @test nout(v1) == 4
        @test nout(v2) == 6
        @test nout(v3) == 6
        @test nout(v4) == 8

        @test @test_logs (:warn, "Selection for vertex v7 failed! Relaxing size constraint...")  match_mode=:any Δoutputs(Output(), v7, v -> 1:nout_org(v))
        apply_mutation(g)

        # Sizes can't change when increasing, even if problem is relaxed :(
        @test nout(v1) == 4
        @test nout(v2) == 6
        @test nout(v3) == 6
        @test nout(v4) == 8

        @test size(g(ones(1, 3))) == (1, nout(v7))
    end

    @testset "SizeStack increase decrease" begin
        inpt = iv(3)
        v1 = av(inpt, 3, "v1")
        v2 = av(inpt, 10, "v2")
        v3 = cc(v1,v2, name="v3")
        v4 = av(v3, 4, "v4")

        g = CompGraph(inpt, v3)
        @test size(g(ones(1, 3))) == (1, nout(v3))

        Δnout(v3, -5)
        Δnout(v1, 3)

        @test nin(v3) == [nout(v1), nout(v2)] == [5, 6]
        @test nout(v3) == sum(nin(v3)) == 11

        @test Δoutputs(Output(), v3, v -> 1:nout_org(v))

        @test in_inds(op(v3)) == [out_inds(op(v1)), out_inds(op(v2))] == [[1,2,3,-1,-1],[5,6,7,8,9,10]]
        @test [out_inds(op(v3))] == in_inds(op(v4)) == [[1,2,3,-1,-1,8,9,10,11,12,13]]

        apply_mutation(g)
        @test size(g(ones(1, 3))) == (1, nout(v3))

    end

    @testset "Constrained by remote subtree" begin
        inpt = iv(3)
        v0 = av(inpt, 10, "v0")

        # Main branch, the one we want to change
        v1 = cc(v0, v0, name="v1")

        # Path to subtree
        # First hop
        v2 = av(inpt, 10, "v2")
        v3 = nc("v3") >> v2 + v0

        # Second hop (v4a and v4b should not be touhed)
        v4a = av(inpt, 3, "v4a")
        v4b = av(inpt, 5, "v4b")
        v5 = cc(v4a, v3, v4b, name="v5")

        # Subtree
        v6 = av(inpt, 4, "v6")
        v7 = av(inpt, nout(v5) - nout(v6), "v7")
        v8 = cc(v6,v7,name="v8")

        # Aaaand connect it to the path
        v9 = nc("v9") >> v8 + v5
        v10 = av(v9, 5, "v10")

        g = CompGraph(inpt, [v1, v9])
        @test size.(g(ones(1,3))) == ((1, nout(v1)), (1, nout(v9)))

        Δnout(v2, -1)

        @test nout(v6) == 4
        @test nout(v7) == 13
        @test nout(v0) == 9

        @test Δoutputs(v1, v -> 1:nout_org(v))
        apply_mutation(g)

        @test nout(v6) == 4
        @test nout(v7) == 13
        @test nout(v0) == 9

        @test size.(g(ones(1,3))) == ((1, nout(v1)), (1, nout(v9)))

        Δnout(v2, 5)

        @test nout(v6) == 4
        @test nout(v7) == 16
        @test nout(v0) == 14

        @test @test_logs (:warn, "Selection for vertex v1 failed! Relaxing size constraint...") Δoutputs(Output(), v1, v->1:nout_org(v))
        apply_mutation(g)

        @test nout(v6) == 4
        @test nout(v7) == 16
        @test nout(v0) == 14

        @test size.(g(ones(1,3))) == ((1, nout(v1)), (1, nout(v9)))
    end

    @testset "Increase after vertex removal" begin
        inpt = iv(3)
        v1 = av(inpt, 2, "v1")
        v2 = tv(v1, "v2")
        v3 = av(v2, 4, "v3")
        v4 = av(v3, 3, "v4")

        g = CompGraph(inpt, v4)
        @test size(g(ones(Float32, 1,3))) == (1, nout(v4))

        remove!(v3, RemoveStrategy())

        @test Δoutputs(v2, v -> 1:nout_org(v))

        @test in_inds(op(v4))[] == [1,2,-1,-1] # TODO: Should be [1,2,3,4]
        @test out_inds(op(v1)) == out_inds(op(v2)) == [1, 2, -1, -1]

        apply_mutation(g)

        @test size(g(ones(Float32, 1,3))) == (1, nout(v4))
    end

    @testset "Decrease after vertex removal" begin
        inpt = iv(3)
        v1 = av(inpt, 2, "v1")
        v2 = tv(v1, "v2")
        v3 = av(v2, 4, "v3")
        v4 = av(v3, 3, "v4")

        g = CompGraph(inpt, v4)
        @test size(g(ones(Float32, 1,3))) == (1, nout(v4))

        remove!(v3, RemoveStrategy(DecreaseBigger()))

        # What happened now is that nin(v4) got decreased from 4 to 2. We now need to select which inputs to keep
        # However, there is absolutely no need at all to select anything from v2 and before as they have not changed.
        @test Δoutputs(Output(), v2, v -> 1:nout_org(v))

        @test in_inds(op(v4))[] == out_inds(op(v1)) == out_inds(op(v2)) == out_inds(op(v3)) == [1, 2]

        apply_mutation(g)

        @test size(g(ones(Float32, 1,3))) == (1, nout(v4))
    end

    @testset "Increase after vertex removal SizeStack" begin
        inpt = iv(3)
        v1 = av(inpt, 3, "v1")
        v2 = av(inpt, 4, "v2")
        v3 = av(v1, 5, "v3")
        v4 = cc(v2, v3, name="v4")

        g = CompGraph(inpt, v4)
        @test size(g(ones(Float32, 1,3))) == (1, nout(v4))

        remove!(v3, RemoveStrategy())

        @test Δoutputs(Output(), v2, v -> 1:nout_org(v))

        @test in_inds(op(v4)) == [out_inds(op(v2)), out_inds(op(v1))] == [[1,2,3,4], [1,2,3,-1,-1]]
        @test out_inds(op(v4)) == 1:9

        apply_mutation(g)

        @test size(g(ones(Float32, 1,3))) == (1, nout(v4))
    end

    @testset "Decrease after vertex removal SizeStack" begin
        inpt = iv(3)
        v1 = av(inpt, 3, "v1")
        v2 = av(inpt, 4, "v2")
        v3 = av(v1, 5, "v3")
        v4 = cc(v2, v3, name="v4")

        g = CompGraph(inpt, v4)
        @test size(g(ones(Float32, 1,3))) == (1, nout(v4))

        remove!(v3, RemoveStrategy(DecreaseBigger()))

        @test Δoutputs(v2, v -> 1:nout_org(v))

        @test in_inds(op(v4)) == [out_inds(op(v2)), out_inds(op(v1))] == [[1,2,3,4], [1,2,3]]
        @test out_inds(op(v4)) == [1,2,3,4,7,8,9]

        apply_mutation(g)

        @test size(g(ones(Float32, 1,3))) == (1, nout(v4))
    end

    @testset "SelectOutputs after increse due to vertex removal" begin
        inpt = iv(3)
        v1 = av(inpt, 2, "v1")
        v2 = av(v1, 3, "v2")
        v3 = av(v2, 5, "v3")
        v4 = av(v3, 3, "v4")

        g = CompGraph(inpt, v4)
        @test size(g(ones(Float32, 1,3))) == (1, nout(v4))

        remove!(v3, RemoveStrategy(SelectOutputs(select=SelectDirection(), valuefun = v -> 1:nout_org(v))))

        @test out_inds(op(v2)) == [1,2,3,-1,-1]
        @test in_inds(op(v4))[] == [1,2,3,4,5]

        apply_mutation(g)

        @test size(g(ones(Float32, 1,3))) == (1, nout(v4))
    end

    @testset "SelectOutputs after decrease due to vertex removal" begin
        inpt = iv(3)
        v1 = av(inpt, 2, "v1")
        v2 = av(v1, 3, "v2")
        v3 = av(v2, 5, "v3")
        v4 = av(v3, 3, "v4")

        g = CompGraph(inpt, v4)
        @test size(g(ones(Float32, 1,3))) == (1, nout(v4))

        remove!(v3, RemoveStrategy(SelectOutputs(align=DecreaseBigger(), valuefun= v -> 1:nout_org(v))))

        @test in_inds(op(v4))[] == out_inds(op(v3)) == [3,4,5]

        apply_mutation(g)

        @test size(g(ones(Float32, 1,3))) == (1, nout(v4))
    end

    @testset "ApplyMutation vertex removal SizeStack" begin
        inpt = iv(3)
        v1 = av(inpt, 3, "v1")
        v2 = av(inpt, 4, "v2")
        v3 = av(v1, 5, "v3")
        v4 = cc(v2, v3, name="v4")

        g = CompGraph(inpt, v4)
        @test size(g(ones(Float32, 1,3))) == (1, nout(v4))

        remove!(v3, RemoveStrategy(ApplyMutation()))
        @test nout(v1) == 5

        @test size(g(ones(Float32, 1,3))) == (1, nout(v4))
    end

    @testset "TruncateInIndsToValid" begin
        inpt = iv(3)
        v1 = av(inpt, 3, "v1")
        v2 = tv(v1, "v2")
        v3 = av(v2, 2, "v3")

        g = CompGraph(inpt, v2)
        @test size(g(ones(Float32, 1,3))) == (1, nout(v2))

        vnew = av(inpt, 7, "vnew")
        create_edge!(vnew, v2;strategy=PostAlignJuMP())
        @test nin_org(v2) == [3, 0] != [nout(v1), nout(vnew)]

        Δoutputs(TruncateInIndsToValid(), g, v -> 1:nout_org(v))

        @test out_inds(op(vnew)) == [4,5,6,7]
        @test in_inds(op(v2)) == [[1,2,3,-1], [1,2,3,-1]]
        @test out_inds(op(v2)) == [1,2,3,-1]
    end

    @testset "CompConstraint" begin

        struct CompConstraint end
        function NaiveNASlib.compconstraint!(::AbstractJuMPSelectionStrategy, ::CompConstraint, data)
            var = data.outselectvars[data.vertex];
            JuMP.@constraint(data.model, var[[1, 3]] .== 0)
        end

        inpt = iv(3)
        v1 = absorbvertex(CompConstraint(), 4, inpt, traitdecoration = tf("v1"))
        v2 = av(v1, 5, "v2")

        Δnout(v1, 3)
        @test @test_logs (:warn, "Selection for vertex v1 failed! Relaxing size constraint...") Δoutputs(Output(), v1, v->nout_org(v):-1:1)

        @test out_inds(op(v1)) == in_inds(op(v2))[] == [2, 4, -1, -1, -1, -1, -1]
    end
end