text:
export overorders, isbass, isgorenstein

################################################################################
#
#  Defines minimal overorder
#
################################################################################
#H must be in lower left hnf form
function iszero_mod_hnf!(a::fmpz_mat, H::fmpz_mat)
  j = ncols(H)
  for i = min(nrows(H), ncols(H)):-1:1
    if iszero(a[i])
      continue
    end
    while j >= 1 && iszero(H[i, j])
      j -= 1
    end
    if iszero(j)
      return iszero(a)
    end
    q, r = divrem(a[1, j], H[i, j])
    if r != 0
      return false
    end
    a[1, j] = r
    for k=1:(j-1)
      a[1, k] = a[1, k] - q * H[i, k]
    end
  end
  return iszero(a)
end

function defines_minimal_overorder(B::Vector{nf_elem}, l::Vector{nf_elem})
  M = basis_mat(B)
  hnf!(M)
  x = M.den * l[1]^2
  if denominator(x) != 1
    return false, M
  end
  m = zero_matrix(FlintZZ, 1, ncols(M))
  for i = 1:ncols(M)
    m[1, i] = numerator(coeff(x, i - 1))
  end
  fl = iszero_mod_hnf!(m, M.num)
  return fl, M
end

################################################################################
#
#  Quotients of two orders
#
################################################################################

# For convenience, there is a quotient constructor for an extension of orders.
# The quotient will be represented as an abelian group.
mutable struct GrpAbFinGenToNfOrdQuoNfOrd <: Map{GrpAbFinGen, NfOrd, HeckeMap, GrpAbFinGenToNfOrdQuoNfOrd}
  domain::GrpAbFinGen
  codomain::NfOrd
  bottom::NfOrd
  offset::Int
  top_snf_basis::Vector{nf_elem}
  top_snf_basis_in_order::Vector{NfOrdElem}
  bottom_snf_basis::Vector{nf_elem}
  top_basis_mat_inv::FakeFmpqMat

  function GrpAbFinGenToNfOrdQuoNfOrd(M::NfOrd, O::NfOrd)
    z = new()
    d = degree(O)
    K = nf(O)
    B = zero_matrix(FlintZZ, d, d)
    for i in 1:d
      v = coordinates(M(elem_in_nf(basis(O)[i])))
      for j in 1:d
        B[i, j] = v[j]
      end
    end
    S, U, V = snf_with_transform(B, true, true)
    Vinv = inv(V)
    basis_O = basis(O)
    basis_M = basis(M)
    new_basis_O = Vector{nf_elem}(undef, d)
    new_basis_M = Vector{nf_elem}(undef, d)
    for i in 1:d
      new_basis_O[i] = elem_in_nf(sum(U[i, j] * basis_O[j] for j in 1:d))
    end

    for i in 1:d
      new_basis_M[i] = elem_in_nf(sum(Vinv[i, j] * basis_M[j] for j in 1:d))
    end

    offset = 1
    while offset <= d && isone(S[offset, offset])
      offset += 1
    end

    offset -= 1

    # offset is the first index - 1, where the diagonal of S is not one.

    A = DiagonalGroup(fmpz[S[i, i] for i in (offset + 1):d])

    z.offset = offset
    z.domain = A
    z.codomain = M
    z.bottom = O
    z.top_snf_basis = new_basis_M
    z.bottom_snf_basis = new_basis_O
    z.top_basis_mat_inv = inv(basis_mat(new_basis_M))
    z.top_snf_basis_in_order = map(M, z.top_snf_basis)

    return z
  end
end

domain(f::GrpAbFinGenToNfOrdQuoNfOrd) = f.domain

codomain(f::GrpAbFinGenToNfOrdQuoNfOrd) = f.codomain

function image(f::GrpAbFinGenToNfOrdQuoNfOrd, x::GrpAbFinGenElem)
  t = zero(codomain(f))
  z = deepcopy(f.top_snf_basis_in_order[1 + f.offset])
  mul!(z, x.coeff[1], z)
  for i in 2:length(x.coeff)
    mul!(t, x.coeff[i], f.top_snf_basis_in_order[i + f.offset])
    add!(z, z, t)
  end
  return z
end

function preimage(f::GrpAbFinGenToNfOrdQuoNfOrd, x::NfOrdElem)
  v = elem_in_nf(x)
  t = f.codomain.tcontain
  elem_to_mat_row!(t.num, 1, t.den, v)
  t = mul!(t, t, f.top_basis_mat_inv)
  @assert isone(t.den)
  return domain(f)(sub(t, 1:1, (1 + f.offset):ncols(t)).num)
end

function quo(M::NfOrd, O::NfOrd)
  f = GrpAbFinGenToNfOrdQuoNfOrd(M, O)
  return domain(f), f
end

################################################################################
#
#  Inducing an endomorphism of the quotient
#
################################################################################

# Let f : A -> M/O be a quotient of orders and g: M -> M an endomorphism of
# O-modules (so g(x) is in O for all x in O).
# This functions computes an endomorphism h: A -> A such that
#
# A <- M/O
# |     |
# v     v
# A <- M/O
#
# commutes. The horizontal maps is the section of f.

function induce(f::GrpAbFinGenToNfOrdQuoNfOrd, g)
  G = domain(f)
  imgs = Vector{GrpAbFinGenElem}(undef, ngens(G))
  d = degree(codomain(f))
  m = zero_matrix(FlintZZ, d, d)
  for i in 1:ngens(G)
    imgs[i] = f\(g(f(G[i])))
  end
  return hom(G, G, imgs)
end

################################################################################
#
#  Minimal overorders
#
################################################################################

# Compute the minimal overorders of O in M by computing the orders generated
# by minimal O-submodules of M.
function _minimal_overorders_meataxe(O::NfOrd, M::NfOrd)

  orders = NfOrd[]
  A, mA = quo(M, O)
  if order(A) == 1
    return orders
  end
  B = mA.bottom_snf_basis
  @assert isone(B[1])
  autos = Vector{GrpAbFinGenMap}(undef, degree(O) - 1)
  # We skip the first basis element, since it acts trivially
  for i in 1:(degree(O) - 1)
    autos[i] = induce(mA, x -> M(elem_in_nf(B[i + 1]))*x)
  end

  d = degree(O)
  K = nf(O)

  potential_basis = Vector{nf_elem}(undef, d)
  subs = stable_subgroups(A, autos, minimal = true, op = (G, z) -> sub(G, z, false))

  for i in 1:mA.offset
    potential_basis[i] = mA.bottom_snf_basis[i]
  end

  offset = mA.offset

  subs = stable_subgroups(A, autos, minimal = true, op = (G, z) -> sub(G, z, false))

  for s in subs
    T = image(s[2], false)
    G = domain(T[2])
    for i in 1:(d - offset)
      v = T[2](G[i]).coeff
      if iszero(v)
        potential_basis[i + offset] = mA.bottom_snf_basis[i + offset]
      else
        @assert ncols(v) == d - offset
        potential_basis[i + offset] = sum(v[1, j] * mA.top_snf_basis[j + offset] for j in 1:(d - offset))
      end
    end
    #L = _Order(K, potential_basis)
    # A minimal overorder must be a minimal submodule
    fl, bL = defines_order(K, potential_basis)
    if !fl
      continue
    end
    bL = hnf(bL)
    #bL = basis_mat(L, copy = false)
    if any(x -> basis_mat(x, copy = false) == bL, orders)
      continue
    else
      push!(orders, Order(K, bL, check = false, cached = false))
    end
  end
  return orders
end

function _minimal_poverorders(O::NfOrd, P::NfOrdIdl, excess = Int[0], use_powering::Bool = false)
  M = ring_of_multipliers(P)
  p = minimum(P)
  f = valuation(norm(P), p)
  if p == 2
    return _minimal_poverorders2(O, P, excess)
  end
  M = ring_of_multipliers(P)
  A, mA = quo(M, O)
  orders = NfOrd[]
  if order(A) == 1
    return orders
  end
  if order(A) == p^f
    O1 = Order(nf(O), hnf(basis_mat(M)), check = false, cached = false)
    push!(orders, O1)
    return orders
  end
  B = mA.bottom_snf_basis
  d = degree(O)
  K = nf(O)
  @assert isone(B[1])
  if use_powering
    autos = Vector{GrpAbFinGenMap}(undef, d)
    autos[d] = induce(mA, x -> x^p)
  else
    autos = Vector{GrpAbFinGenMap}(undef, d - 1)
  end
  # We skip the first basis element, since it acts trivially
  for i in 1:(d - 1)
    autos[i] = induce(mA, x -> M(elem_in_nf(B[i + 1]))*x)
  end

  d = degree(O)
  K = nf(O)

  potential_basis = Vector{nf_elem}(undef, d)

  filter!( x -> !iszero(x.map), autos)
  if iszero(length(autos))
    subs = subgroups(A, subtype = [Int(p)^f], fun = (G, z) -> sub(G, z, false))
  else
    subs = stable_subgroups(A, autos, minimal = true, op = (G, z) -> sub(G, z, false))
  end

  for i in 1:mA.offset
    potential_basis[i] = mA.bottom_snf_basis[i]
  end

  offset = mA.offset

  lQ = prime_ideals_over(M, P)
  rel_fs = fmpz[divexact(valuation(norm(Q), p), f) for Q in lQ]
  fac = factor(lcm(rel_fs))
  for (q, _) in fac
    if q == 2
      continue
    else
      subs = subgroups(A, subtype = [Int(p) for j in 1:Int(f*(Int(q) - 1))], fun = (G, z) -> sub(G, z, false))
      for s in subs
        T = image(s[2], false)
        G = domain(T[2])
        new_element = 0
        for i in 1:(d - offset)
          v = T[2](G[i]).coeff
          if iszero(v)
            potential_basis[i + offset] = mA.bottom_snf_basis[i + offset]
          else
            @assert ncols(v) == d - offset
            potential_basis[i + offset] = sum(v[1, j] * mA.top_snf_basis[j + offset] for j in 1:(d - offset))
            new_element = i + offset
          end
        end
        @assert new_element != 0
        b, bL = defines_order(K, potential_basis)
        if !b
          excess[] = excess[] + 1
          continue
        end
        L = Order(K, bL, check = false, cached = false)
        lQL = prime_ideals_over(L, P)
        if length(lQL) == 1 && norm(lQL[1]) == norm(P)^q
          push!(orders, L)
        end
      end
    end
  end

  if !(2 in fac) && length(lQ) == 1 && norm(P)^rel_fs[1] == order(A)
    return orders
  end

  for s in subs
    T = image(s[2], false)
    G = domain(T[2])
    new_element = 0
    for i in 1:(d - offset)
      v = T[2](G[i]).coeff
      if iszero(v)
        potential_basis[i + offset] = mA.bottom_snf_basis[i + offset]
      else
        @assert ncols(v) == d - offset
        potential_basis[i + offset] = sum(v[1, j] * mA.top_snf_basis[j + offset] for j in 1:(d - offset))
        new_element = i + offset
      end
    end
    @assert new_element != 0
    b, bL = defines_minimal_overorder(potential_basis, nf_elem[potential_basis[new_element]])
    if !b
      excess[] = excess[] + 1
      continue
    end
    L = Order(K, bL, check = false, cached = false)
    push!(orders, L)
  end

  return orders
end

function _minimal_poverorders2(O::NfOrd, P::NfOrdIdl, excess = Int[])
  M = ring_of_multipliers(P)
  A, mA = quo(M, O)
  orders = NfOrd[]
  if order(A) == 1
    return orders
  end
  B = mA.bottom_snf_basis
  d = degree(O)
  K = nf(O)
  if norm(P) == order(A)
    O1 = Order(K, hnf(basis_mat(M, copy = false)), check = false, cached = false)
    push!(orders, O1)
    return orders
  end

  f = valuation(norm(P), 2)
  @assert isone(B[1])
  autos = Vector{GrpAbFinGenMap}(undef, d)
  # We skip the first basis element, since it acts trivially
  for i in 1:(d - 1)
    autos[i] = induce(mA, x -> M(elem_in_nf(B[i + 1]))*x)
  end
  autos[d] = induce(mA, x -> x^2)

  filter!( x -> !iszero(x.map), autos)
  if iszero(length(autos))
    subs = subgroups(A, subtype = [2], fun = (G, z) -> sub(G, z, false))
  else
    R = GF(2)
    W = MatrixSpace(R, ngens(A), ngens(A), false)
    V = ModAlgAss([W(l.map) for l in autos])

    subm = minimal_submodules(V, f)
    subs = (sub(A, lift(x), false) for x in subm)
  end

  potential_basis = Vector{nf_elem}(undef, d)

  offset = mA.offset
  for i in 1:offset
    potential_basis[i] = mA.bottom_snf_basis[i]
  end

  for s in subs
    T = image(s[2], false)
    G = domain(T[2])
    new_element = 0
    for i in 1:(d - offset)
      v = T[2](G[i]).coeff
      if iszero(v)
        potential_basis[i + offset] = mA.bottom_snf_basis[i + offset]
      else
        @assert ncols(v) == d - offset
        potential_basis[i + offset] = sum(v[1, j] * mA.top_snf_basis[j + offset] for j in 1:(d - offset))
        new_element = i + offset
      end
    end
    @assert new_element != 0
    bL = hnf!(basis_mat(potential_basis))
    L = Order(K, bL, check = false, cached = false)
    push!(orders, L)
  end
  return orders
end

################################################################################
#
#  Primary overorders
#
################################################################################

function new_poverorders(O::NfOrd, p::fmpz)
  lP = prime_ideals_over(O, p)
  res = typeof(O)[O]
  K = nf(O)
  tz = zero_matrix(FlintZZ, 2 * degree(O), degree(O))
  for P in lP
    nres = typeof(O)[]
    tP = @elapsed Pprim = pprimary_overorders(O, P)
    @show length(Pprim), length(res)
    trec = @elapsed for R in Pprim
      #@show R
      for S in res
        #@show S
        #@show sum_as_Z_modules(R, S)
        push!(nres, sum_as_Z_modules(R, S, tz), triangular = true)
      end
    end
    println("time for recombination: $trec")
    res = nres
  end
  return res
end

function bass_overorders(O::NfOrd)
  orders = typeof(O)[]
  M = maximal_order(O)
  fac = factor(FlintZZ(div(gen_index(M), gen_index(O))))
  for (p, ) in fac
    lp = prime_ideals_over(O, p)
    for P in lp
      Pprim = pprimary_overorders(O, P)
      bassP = Vector{typeof(O)}()
      for A in Pprim
        add = true
        pabove = prime_ideals_over(A, P)
        for Q in pabove
          if !isbass(A, Q)
            add = false
            break
          end
        end
        if add
          push!(bassP, A)
        end
      end
      if isempty(orders)
        append!(orders, bassP)
      else
        orders1 = Vector{typeof(O)}(undef, length(orders) * length(bassP))
        z = zero_matrix(FlintZZ, 2*degree(O), degree(O))
        kk = 1
        for O1 in orders
          @time for O2 in bassP
            orders1[kk] = sum_as_Z_modules(O1, O2, z, triangular = true)
            kk += 1
          end
        end
        orders = orders1
      end
    end
  end

  if length(orders) == 0
    return typeof(O)[O]
  end
  return orders
end

function gorenstein_overorders(O::NfOrd)
  orders = typeof(O)[]
  M = maximal_order(O)
  fac = factor(FlintZZ(div(gen_index(M), gen_index(O))))
  for (p, ) in fac
    lp = prime_ideals_over(O, p)
    for P in lp
      Pprim = pprimary_overorders(O, P)
      gorP = Vector{typeof(O)}()
      for A in Pprim
        add = true
        pabove = prime_ideals_over(A, P)
        for Q in pabove
          if !isgorenstein(A, Q)
            add = false
            break
          end
        end
        if add
          push!(gorP, A)
        end
      end
      if isempty(orders)
        append!(orders, gorP)
      else
        orders1 = Vector{typeof(O)}(undef, length(orders) * length(gorP))
        z = zero_matrix(FlintZZ, 2*degree(O), degree(O))
        kk = 1
        for O1 in orders
          @time for O2 in gorP
            orders1[kk] = sum_as_Z_modules(O1, O2, z, triangular = true)
            kk += 1
          end
        end
        orders = orders1
      end
    end
  end

  if length(orders) == 0
    return typeof(O)[O]
  end
  return orders
end

function new_overorders(O::NfOrd)
  orders = typeof(O)[]
  M = maximal_order(O)
  for (p, ) in factor(FlintZZ(div(gen_index(M), gen_index(O))))
    print("Time for $p: ")
    tp = @elapsed new_p = new_poverorders(O, p)
    if isempty(orders)
      append!(orders, new_p)
    else
      orders1 = Vector{typeof(O)}(undef, length(orders) * length(new_p))
      sizehint!(orders1, length(orders) * length(new_p))
      z = zero_matrix(FlintZZ, 2*degree(O), degree(O))
      @show length(orders), length(new_p)
      kk = 1
      for O1 in orders
        @time for O2 in new_p
          orders1[kk] = sum_as_Z_modules(O1, O2, z, triangular = true)
          kk += 1
        end
      end
      orders = orders1
    end
    println(tp)
    #@show p, length(orders[end])
  end

  if length(orders) == 0
    return typeof(O)[O]
  end
  return orders
end

function pprimary_overorders(O::NfOrd, P::NfOrdIdl)
  to_enlarge = typeof(O)[O]
  current = Dict{fmpq, Dict{FakeFmpqMat, typeof(O)}}()
  excess = Int[0]
  while length(to_enlarge) > 0
    #@show length(to_enlarge)
    #if length(current) > 0
    #  @show sum([length(x) for x in values(current)])
    #end
    N = popfirst!(to_enlarge)
    lQ = prime_ideals_over(N, P)
    for Q in lQ
      #new = append!(new, _minimal_overorders_meataxe(N, ring_of_multipliers(Q)))
      if length(lQ) == 1 && isbass(N, Q)
        new = new_pprimary_overorders_bass(N, Q)
        for S in new
          #@show ishnf(basis_mat(S).num, :lowerleft)
          #t += @elapsed H = hnf(basis_mat(S, copy = false))
          H = basis_mat(S, copy = false)
          ind = prod(H.num[i, i] for i in 1:degree(O))//(H.den)^degree(O)
          if haskey(current, ind)
            c = current[ind]
            if !haskey(c, H)
              c[H] = S
            end
          else
            c = Dict{FakeFmpqMat, typeof(O)}()
            current[ind] = c
            c[H] = S
          end
        end
      else
        new = _minimal_poverorders(N, Q, excess)
        for S in new
          #@show ishnf(basis_mat(S).num, :lowerleft)
          #t += @elapsed H = hnf(basis_mat(S, copy = false))
          H = basis_mat(S, copy = false)
          #@assert ishnf(H.num, :lowerleft)
          ind = prod(H.num[i, i] for i in 1:degree(O))//(H.den)^degree(O)
          if haskey(current, ind)
            c = current[ind]
            if !haskey(c, H)
              c[H] = S
              push!(to_enlarge, S)
            end
          else
            c = Dict{FakeFmpqMat, typeof(O)}()
            current[ind] = c
            c[H] = S
            push!(to_enlarge, S)
          end
        end
      end
    end
  end
  push!(to_enlarge, O)
  for d in values(current)
    for e in values(d)
      push!(to_enlarge, e)
    end
  end
  println("excess: $(excess[])")
  return to_enlarge
end

function pprimary_overorders_old(O::NfOrd, P::NfOrdIdl)
  #if isbass(O, P)
  #  return pprimary_overorders_bass(O, P)
  #end

  if index(E) != index(O)
    minimaloverorders = _minimal_poverorders(O, P)#_minimal_overorders_meataxe(O, E)
  else
    return typeof(O)[O]
  end

  res = typeof(O)[O]

  for R in minimaloverorders
    primes = prime_ideals_over(R, P)

    if length(primes) == 1
      rres = pprimary_overorders(R, primes[1])
    else
      rres1 = pprimary_overorders(R, primes[1])
      rres2 = pprimary_overorders(R, primes[2])
      rres = typeof(rres1)()
      for R1 in rres1
        for R2 in rres2
          push!(rres, R1 + R2)
        end
      end
    end
    nres = typeof(res)()
    for C in res
      for T in rres
        push!(nres, C + T)
      end
    end
    append!(res, nres)
    res = unique(OO -> basis_mat(OO, copy = false), res)
  end
  @show length(res)
  return res
end

function overorders_naive(O::NfOrd, M::NfOrd = maximal_order(nf(O)))
  d = degree(O)
  K = nf(O)
  B = zero_matrix(FlintZZ, d, d)
  for i in 1:d
    v = coordinates(M(elem_in_nf(basis(O)[i])))
    for j in 1:d
      B[i, j] = v[j]
    end
  end
  S, U, V = snf_with_transform(B, true, true)
  Vinv = inv(V)
  basis_O = basis(O)
  basis_M = basis(M)
  new_basis_O = Vector{nf_elem}(undef, d)
  new_basis_M = Vector{nf_elem}(undef, d)
  for i in 1:d
    new_basis_O[i] = elem_in_nf(sum(U[i, j] * basis_O[j] for j in 1:d))
  end

  for i in 1:d
    new_basis_M[i] = elem_in_nf(sum(Vinv[i, j] * basis_M[j] for j in 1:d))
  end

  A = AbelianGroup(S)

  if order(A) == 1
    return typeof(O)[O]
  end

  # Don't put things in the subgroup lattice
  subs = subgroups(A, fun = (G, z) -> sub(G, z, false))
  potential_basis = Vector{nf_elem}(undef, d)
  oorders = typeof(O)[]
  println("#subgroups: $(length(collect(subs)))")
  for (i, s) in enumerate(subs)
    #@show i
    T = image(s[2], false)
    G = domain(T[2])
    for i in 1:d
      v = T[2](G[i]).coeff
      if iszero(v)
        potential_basis[i] = new_basis_O[i]
      else
        potential_basis[i] = sum(v[1, j] * new_basis_M[j] for j in 1:d)
      end
    end
    b, bmat = defines_order(K, potential_basis)
    if b
      push!(oorders, Order(K, hnf!(bmat)))
    end
  end
  return oorders
end

function poverorders_meataxe(O::NfOrd, p::Integer)
  return poverorders_meataxe(O, fmpz(p))
end

function minimal_poverorders_naive(O::NfOrd, p::fmpz)
  orders = poverorders_meataxe(O, p)
  res = Vector{typeof(O)}()
  for S in orders
    if length(overorders_meataxe(O, S)) == 2
      push!(res, S)
    end
  end
  return res
end

function poverorders_from_multipliers(O, p::fmpz)
  M = MaximalOrder(O)
  lP = prime_ideals_over(O, p)
  orders = typeof(O)[]
  for P in lP
    E = ring_of_multipliers(P)
    if index(E) != index(O)
      append!(orders, _overorders_meataxe(O, E))
    end
  end
  if length(orders) == 0
    push!(orders, O)
  end
  return orders
end

function poverorders(O, p::fmpz)
  to_enlarge = typeof(O)[O]
  current = Dict{fmpq, Dict{FakeFmpqMat, typeof(O)}}()
  while length(to_enlarge) > 0
    N = pop!(to_enlarge)
    new = poverorders_from_multipliers(N, p)
    for S in new
      H = hnf(basis_mat(S, copy = false))
      ind = prod(H.num[i, i] for i in 1:degree(O))//H.den
      if haskey(current, ind)
        c = current[ind]
        if haskey(c, H)
          continue
        else
          c[H] = S
          push!(to_enlarge, S)
        end
      else
        c = Dict{FakeFmpqMat, typeof(O)}()
        current[ind] = c
        c[H] = S
        push!(to_enlarge, S)
      end
    end
  end
  for d in values(current)
    for e in values(d)
      push!(to_enlarge, e)
    end
  end
  return to_enlarge
end

function _overorders_meataxe(O::NfOrd, M::NfOrd)
  K = nf(O)
  d = degree(O)
  B = zero_matrix(FlintZZ, d, d)
  orders = Vector{typeof(O)}()
  for i in 1:d
    v = coordinates(M(elem_in_nf(basis(O)[i])))
    for j in 1:d
      B[i, j] = v[j]
    end
  end
  S::fmpz_mat, U::fmpz_mat, V::fmpz_mat = snf_with_transform(B, true, true)
  Vinv = inv(V)
  basis_O = basis(O)
  basis_M = basis(M)
  new_basis_O = Vector{nf_elem}(undef, d)
  new_basis_M = Vector{nf_elem}(undef, d)
  for i in 1:d
    new_basis_O[i] = elem_in_nf(sum(U[i, j] * basis_O[j] for j in 1:d))
  end

  for i in 1:d
    new_basis_M[i] = elem_in_nf(sum(Vinv[i, j] * basis_M[j] for j in 1:d))
  end

  new_basis_mat_M_inv = inv(basis_mat(new_basis_M))

  autos = GrpAbFinGenMap[]

  A = DiagonalGroup(fmpz[ S[i, i] for i in 1:d])

  for i in 1:d
    b = new_basis_O[i]
    m = zero_matrix(FlintZZ, d, d)
    for j in 1:d
      v = elem_in_nf(M(b* new_basis_M[j]))
      t = O.tcontain
      elem_to_mat_row!(t.num, 1, t.den, v)
      t = mul!(t, t, new_basis_mat_M_inv)
      # I need the representation with respect to new_basis_M
      for k in 1:d
        m[j, k] = t.num[1, k]
      end
    end
    push!(autos, hom(A, A, m))
  end

  potential_basis = Vector{nf_elem}(undef, d)

  subs = stable_subgroups(A, autos, op = (G, z) -> sub(G, z, false))
  for s in subs
    T = image(s[2], false)
    G = domain(T[2])
    for i in 1:d
      v = T[2](G[i]).coeff
      if iszero(v)
        potential_basis[i] = new_basis_O[i]
      else
        potential_basis[i] = sum(v[1, j] * new_basis_M[j] for j in 1:d)
      end
    end
    b, bmat = defines_order(K, deepcopy(potential_basis))
    if b
      push!(orders, Order(K, hnf(bmat)))
    end
  end
  return orders
end

function poverorders_meataxe(O::NfOrd, p::fmpz, N::NfOrd = pmaximal_overorder(O, p))
  K = nf(O)
  d = degree(O)
  M = intersect(N, pmaximal_overorder(O, p))
  A, mA = quo(M, O)
  orders = NfOrd[O]
  if order(A) == 1
    return orders
  end
  B = mA.bottom_snf_basis
  d = degree(O)
  K = nf(O)
  @assert isone(B[1])
  autos = Vector{GrpAbFinGenMap}(undef, d - 1)
  # We skip the first basis element, since it acts trivially
  for i in 1:(d - 1)
    autos[i] = induce(mA, x -> M(elem_in_nf(B[i + 1]))*x)
  end

  d = degree(O)
  K = nf(O)

  potential_basis = Vector{nf_elem}(undef, d)
  orders = NfOrd[]

  for i in 1:mA.offset
    potential_basis[i] = mA.bottom_snf_basis[i]
  end

  offset = mA.offset

  #B = zero_matrix(FlintZZ, d, d)
  #orders = Vector{typeof(O)}()
  #for i in 1:d
  #  v = coordinates(M(elem_in_nf(basis(O)[i])))
  #  for j in 1:d
  #    B[i, j] = v[j]
  #  end
  #end
  #S::fmpz_mat, U::fmpz_mat, V::fmpz_mat = snf_with_transform(B, true, true)
  #Vinv = inv(V)
  #basis_O = basis(O)
  #basis_M = basis(M)
  #new_basis_O = Vector{nf_elem}(undef, d)
  #new_basis_M = Vector{nf_elem}(undef, d)
  #for i in 1:d
  #  new_basis_O[i] = elem_in_nf(sum(U[i, j] * basis_O[j] for j in 1:d))
  #end

  #for i in 1:d
  #  new_basis_M[i] = elem_in_nf(sum(Vinv[i, j] * basis_M[j] for j in 1:d))
  #end

  #new_basis_mat_M_inv = inv(basis_mat(new_basis_M))

  #autos = GrpAbFinGenMap[]

  #A = DiagonalGroup(fmpz[ S[i, i] for i in 1:d])

  #for i in 1:d
  #  b = new_basis_O[i]
  #  m = zero_matrix(FlintZZ, d, d)
  #  for j in 1:d
  #    v = elem_in_nf(M(b* new_basis_M[j]))
  #    t = O.tcontain
  #    elem_to_mat_row!(t.num, 1, t.den, v)
  #    t = mul!(t, t, new_basis_mat_M_inv)
  #    # I need the representation with respect to new_basis_M
  #    for k in 1:d
  #      m[j, k] = t.num[1, k]
  #    end
  #  end
  #  push!(autos, hom(A, A, m))
  #end
  #
  #potential_basis = Vector{nf_elem}(undef, d)

  excess = 0

  subs = stable_subgroups(A, autos, op = (G, z) -> sub(G, z, false))

  for s in subs
    T = image(s[2], false)
    G = domain(T[2])
    new_element = 0
    for i in 1:(d - offset)
      v = T[2](G[i]).coeff
      if iszero(v)
        potential_basis[i + offset] = mA.bottom_snf_basis[i + offset]
      else
        @assert ncols(v) == d - offset
        potential_basis[i + offset] = sum(v[1, j] * mA.top_snf_basis[j + offset] for j in 1:(d - offset))
        new_element = i + offset
      end
    end
    b, bmat = defines_order(K, deepcopy(potential_basis))
    if b
      push!(orders, Order(K, bmat))
    else
      excess += 1
    end
  end
  @show excess, length(orders)
  return orders
end

function overorders_meataxe(O::NfOrd, M::NfOrd = maximal_order(O))
  orders = Vector{typeof(O)}[]

  for (p, ) in factor(div(index(M), index(O)))
    push!(orders, poverorders_meataxe(O, p, M))
  end

  if length(orders) == 0
    return typeof(O)[O]
  end

  res = Vector{typeof(O)}(undef, prod(length(orders[i]) for i in 1:length(orders)))

  if length(orders) == 1
    return orders[1]
  else
    I = Iterators.product(orders...)
    for (j, i) in enumerate(I)
      res[j] = sum(i)
    end
    return res
  end
end

@doc Markdown.doc"""
    overorders(O::NfOrd)

Returns all overorders of $\mathcal O$.
"""
function overorders(O::NfOrd)
  return overorders_meataxe(O)
end

function new_pprimary_overorders_bass(O::NfOrd, P::NfOrdIdl)
  res = NfOrd[]
  O1 = ring_of_multipliers(P)
  if gen_index(O1) == gen_index(O)
    return res
  end
  K = nf(O)
  O1n = Order(K, hnf(basis_mat(O1, copy = false)), check = false, cached = false)
  push!(res, O1n)
  primes = prime_ideals_over(O1, P)
  while length(primes) == 1
    O2 = ring_of_multipliers(primes[1])
    if index(O2) == index(O1)
      return res
    end
    O2n = Order(K, hnf(basis_mat(O2, copy = false)), check = false, cached = false)
    push!(res, O2n)
    O1 = O2
    primes = prime_ideals_over(O1, primes[1])
  end
  @assert length(primes) == 2
  #There was branching.
  res1 = NfOrd[]
  O3 = O1
  O2 = ring_of_multipliers(primes[1])
  while index(O3) != index(O2)
    O2n = Order(K, hnf(basis_mat(O2, copy = false)), check = false, cached = false)
    push!(res1, O2n)
    O3 = O2
    P = prime_ideals_over(O2, primes[1])[1]
    O2 = ring_of_multipliers(P)
  end

  res2 = NfOrd[]
  O3 = O1
  O2 = ring_of_multipliers(primes[2])
  while index(O3) != index(O2)
    O2n = Order(K, hnf(basis_mat(O2, copy = false)), check = false, cached = false)
    push!(res2, O2)
    O3 = O2
    P = prime_ideals_over(O2, primes[2])[1]
    O2 = ring_of_multipliers(P)
  end
  append!(res, res1)
  append!(res, res2)
  for N in res1
    for M in res2
      push!(res, N + M)
    end
  end
  return res

end

function pprimary_overorders_bass(O::NfOrd, P::NfOrdIdl; branching::Bool = true)
  res = NfOrd[]
  R = ring_of_multipliers(P)
  if index(R) == index(O)
    return typeof(O)[O]
  end
  lQ = prime_ideals_over(R, P)
  if !branching
    res = typeof(O)[O]
    for Q in lQ
      if intersect(Q, O) == P
        return append!(res, pprimary_overorders_bass(R, Q, branching = false))
      end
    end
  end

  res = typeof(O)[O]

  primes = NfOrdIdl[]

  k = 0
  for Q in lQ
    if k == 2
      break
    end
    if intersect(Q, O) == P
      push!(primes, Q)
      k = k + 1
    end
  end

  if k == 1
    append!(res, pprimary_overorders_bass(R, primes[1], branching = false))
    return res
  else
    @assert k == 2
    S1 = pprimary_overorders_bass(R, primes[1], branching = false)
    S2 = pprimary_overorders_bass(R, primes[2], branching = false)
    for T1 in S1
      for T2 in S2
        push!(res, T1 + T2)
      end
    end
    return res
  end
end

################################################################################
#
#  IsBass function
#
################################################################################

function isbass(O::NfOrd, P::NfOrdIdl)
  M = maximal_order(O)
  Q = extend(P, M)
  p = minimum(P)
  resfield_dim = valuation(norm(P), p)
  ext_dim = valuation(norm(Q), p)
  @assert mod(ext_dim, resfield_dim) == 0
  return div(ext_dim, resfield_dim) <= 2
end

function isbass(O::NfOrd, p::fmpz)
  M = maximal_order(O)
  p_critical_primes = Set{ideal_type(O)}()
  lp = prime_decomposition(M, p)
  for (P, e) in lp
    push!(p_critical_primes, contract(P, O))
  end
  for Q in p_critical_primes
    resfield_dim = valuation(norm(Q), p)
    ext_dim = valuation(norm(Q * M), p)
    @assert mod(ext_dim, resfield_dim) == 0
    if div(ext_dim, resfield_dim) > 2
      return false
    end
  end
  return true
end

function isbass(O::NfOrd)
  f = minimum(conductor(O))
  M = maximal_order(nf(O))
  for (p, ) in factor(f)
    if !isbass(O, p)
      return false
    end
  end
  return true
end

################################################################################
#
#  IsGorenstein function
#
################################################################################

function isgorenstein(O)
  codiff = codifferent(O)
  R = simplify(simplify(colon(1*O, codiff.num) * codiff) * codiff.den)
  return isone(norm(R))
end

function isgorenstein(O::NfOrd, p::Int)
  return isgorenstein(O, fmpz(p))
end

function isgorenstein(O::NfOrd, p::fmpz)
  codiff = codifferent(O)
  R = simplify(simplify(colon(1*O, codiff.num) * codiff) * codiff.den)
  v = valuation(norm(R), p)
  return v == 0
end

function isgorenstein(O::NfOrd, P::NfOrdIdl)
  J = colon(ideal(O, 1), P)
  return isone(norm(P)*det(basis_mat(J)))
end

# This is very slow!
function intersect(x::NfOrd, y::NfOrd)
  d = degree(x)
  g = lcm(denominator(basis_mat(x)), denominator(basis_mat(y)))
  H = vcat(divexact(g * basis_mat(x).num, basis_mat(x).den), divexact(g * basis_mat(y).num, basis_mat(y).den))
  K = left_kernel(H)[2]
  return Order(nf(x), FakeFmpqMat(_hnf(sub(K, 1:d, 1:d)*divexact(g * basis_mat(x).num, basis_mat(x).den), :lowerleft), g))
end

################################################################################
#
#  Goursat
#
################################################################################

function ideals_with_norm(O::NfOrd, p::fmpz, n::Int)
  pn = p^n
  pInt = Int(p)
  K = nf(O)
  d = degree(O)
  ideals = []
  B = basis(O)

  autos = GrpAbFinGenMap[]

  A = DiagonalGroup(fmpz[pn for i in 1:d])

  for i in 1:d
    m = representation_matrix(B[i])
    push!(autos, hom(A, A, m))
  end

  potential_basis = Vector{elem_type(O)}(d)
  ideals = Vector{Tuple{Vector{Int}, ideal_type(O)}}()

  for par in AllParts(n)
    ppar = Vector(par)

    subs = stable_subgroups(A, autos, quotype=[pInt^y for y in ppar], op = (G, z) -> sub(G, z, false))

    for s in subs
      new_basis_mat = zero_matrix(FlintZZ, d, d)
      T = image(s[2], false)
      G = domain(T[2])
      for i in 1:d
        v = T[2](G[i]).coeff
        if iszero(v)
          new_basis_mat[i, i] = pn
        else
          for j in 1:d
            new_basis_mat[i, j] = v[1, j]
          end
        end
      end
      push!(ideals, ([pInt^y for y in ppar], ideal(O, new_basis_mat)))
    end
  end
  return ideals
end

function index(R::NfOrd, S::NfOrd)
  r = gen_index(R)
  s = gen_index(S)
  i = r^-1 * s
  @assert isinteger(i)
  return FlintZZ(i)
end

function poverorders_goursat(O1::NfOrd, O2::NfOrd, p::fmpz)
  l1 = poverorders(O1, p)
  l2 = poverorders(O2, p)
  data_from_l2 = Dict{Vector{Int}, Vector{Tuple{typeof(O1), ideal_type(O1)}}}()
  d = degree(O2)
  for O in l2
    i = index(O2, O)
    e, _ = ispower(i)
    for k in 1:e
      ideals = ideals_with_norm(O, p, k)
      for (typ, I) in ideals
        if haskey(data_from_l2, typ)
          push!(data_from_l2[typ], (O, I))
        else
          data_from_l2[typ] = [(O, I)]
        end
      end
    end
  end

  return data_from_l2
end

function abelian_group(Q::NfOrdQuoRing)
  A = AbelianGroup(Q.basis_mat)
  S, mS = snf(A)
  B = basis(Q.base_ring, copy = false)
  f = a -> begin aa = mS(a); Q(sum(aa.coeff[i] * B[i] for i in 1:degree(Q.base_ring))) end
  g = b -> mS\A(coordinates(b.elem))
  return S, f, g
end

function isisomorphic(Q1::NfOrdQuoRing, Q2::NfOrdQuoRing)
  Q1_A, Q1_mA, Q1_mA_inv = abelian_group(Q1)
  Q2_A, Q2_mA, Q2_mA_inv = abelian_group(Q2)

  if Q1_A.snf != Q2_A.snf
    return false
  end

  Q1_U, Q1_mU = multiplicative_group(Q1)
  Q2_U, Q2_mU = multiplicative_group(Q2)

  if Q1_U.snf != Q2_U.snf
    return false
  end

  Q1_gens = [ Q1_mA(g) for g in gens(Q1_A) ]

  orders = [ order(g) for g in gens(Q1_A) ]

  l = length(Q1_A.snf)

  elements_with_correct_order = Dict{fmpz, Vector{GrpAbFinGenElem}}()

  for g in Q2_A
    o = order(g)
    if o in orders
      #elem_g = Q2_mA(g)
      if !haskey(elements_with_correct_order, o)
        elements_with_correct_order[o] = GrpAbFinGenElem[g]
      else
        push!(elements_with_correct_order[o], g)
      end
    end
  end

  isos = []

  genQ1A = gens(Q1_A)

  O1 = base_ring(Q1)
  O2 = base_ring(Q2)

  basis_O1 = basis(O1)

  for poss in Iterators.product([ elements_with_correct_order[o] for o in orders]...)
    #@show poss
    h = hom(Q1_A, collect(poss))
    if !isbijective(h)
      continue
    end
    if h(Q1_mA_inv(one(Q1))) != Q2_mA_inv(one(Q2))
      continue
    end
    multiplicative = true
    for i in 1:l
      for j in i:l
        @assert h(genQ1A[i]) == poss[i]
        @assert h(genQ1A[j]) == poss[j]
        if h(Q1_mA_inv(Q1_mA(genQ1A[i]) * Q1_mA(genQ1A[j]))) != Q2_mA_inv(Q2_mA(poss[i]) * Q2_mA(poss[j]))
          multiplicative = false
          break
        end
      end
      if !multiplicative
        break
      end
    end
    if multiplicative
      M = Array{fmpz}(degree(O1), degree(O2))
      for i in 1:degree(O1)
        v = coordinates(Q2_mA(h(Q1_mA_inv(Q1(basis_O1[i])))).elem)
        for j in 1:degree(O2)
          M[i, j] = v[j]
        end
      end
      push!(isos, M)
    end
  end
  return isos
end


text:
using CondaBinDeps, Compat, Compat.Test, BinDeps
import CondaBinDeps.Conda

# manager for test environment
testenv = :CondaBinDeps_test
TestManager = CondaBinDeps.EnvManager{testenv}

if "libpng" in Conda._installed_packages(testenv)
    Conda.rm("libpng", testenv)
end

# force installation by Conda
empty!(BinDeps.defaults)
push!(BinDeps.defaults, BinDeps.PackageManager)

@BinDeps.setup
libpng = library_dependency("libpng", aliases = ["libpng16"])
provides(TestManager, "libpng", libpng)
@BinDeps.install Dict(:libpng => :libpng_lib)

@test "libpng" in Conda._installed_packages(testenv)

deps = joinpath(@__DIR__, "deps.jl")
@test isfile(deps)
include(deps)
@test isfile(libpng_lib)

ext: jl


text:
export Headers, append, get, set, unset, delete

import Base

struct Headers
    data::Dict{AbstractString, AbstractString}
    keymap::Dict{AbstractString, AbstractString}
end # Headers

function Headers()
    Headers(Dict{AbstractString, AbstractString}(), Dict{AbstractString, AbstractString}())
end

function Headers(entries::Pair{K,V}...) where {K <: AbstractString, V <: AbstractString}
    function reducer(headers, entry)
        append(headers, first(entry), last(entry))
    end
    reduce(reducer, entries; init=Headers())
end

function Base.getindex(headers::Headers, name::AbstractString)
    if !haskey(headers, name)
        throw(BoundsError(headers, name))
    end

    get(headers, name, "")
end

function Base.haskey(headers::Headers, name::AbstractString)
    lcname = lowercase(name)
    haskey(headers.keymap, lcname)
end

function Base.get(headers::Headers, name::AbstractString, default=nothing)
    lcname = lowercase(name)
    if !haskey(headers.keymap, lcname)
        default
    else
        realname = get(headers.keymap, lcname, "")
        get(headers.data, realname, default)
    end
end

function set(headers::Headers, name::AbstractString, value::AbstractString)
    lcname = lowercase(name)
    newdata = merge!(Dict{AbstractString, AbstractString}(), headers.data)
    newkeymap = merge!(Dict{AbstractString, AbstractString}(), headers.keymap)

    if haskey(newkeymap, lcname)
        oldname = get(newkeymap, lcname, "")
        Base.delete!(newdata, oldname)
        Base.delete!(newkeymap, lcname)
    end

    Base.setindex!(newkeymap, name, lcname)
    Base.setindex!(newdata, value, name)

    Headers(newdata, newkeymap)
end

set(headers::Headers, name::AbstractString, value::Nothing) = unset(headers, name)

function unset(headers::Headers, name::AbstractString)
    lcname = lowercase(name)
    newdata = merge!(Dict{AbstractString, AbstractString}(), headers.data)
    newkeymap = merge!(Dict{AbstractString, AbstractString}(), headers.keymap)

    if haskey(newkeymap, lcname)
        oldname = get(newkeymap, lcname, "")
        Base.delete!(newdata, oldname)
    end

    Base.delete!(newkeymap, lcname)
    Base.delete!(newdata, name)

    Headers(newdata, newkeymap)
end

function append(headers::Headers, name::AbstractString, value::AbstractString)
    lcname = lowercase(name)
    if haskey(headers.keymap, lcname)
        name = Base.get(headers.keymap, lcname, name)
        value = string(Base.get(headers.data, name, ""), ", ", value)
    end

    set(headers, name, value)
end

function delete(headers::Headers, name::AbstractString)
    lcname = lowercase(name)
    if haskey(headers.keymap, lcname)
        newdata = merge!(Dict{AbstractString, AbstractString}(), headers.data)
        newkeymap = merge!(Dict{AbstractString, AbstractString}(), headers.keymap)
        oldname = get(newkeymap, lcname, "")
        Base.delete!(newdata, oldname)
        Base.delete!(newkeymap, lcname)
        Headers(newdata, newkeymap)
    else
        headers
    end
end

function Base.show(io::IO, headers::Headers)
    comma = false
    print(io, "Headers(")
    for (name, value) in headers.data
        if comma
            print(io, ", ")
        end
        show(io, name)
        print(io, " => ")
        show(io, value)

        comma = true
    end
    print(io, ")")
end

Base.iterate(headers::Headers) = Base.iterate(headers.data)
Base.iterate(headers::Headers, state) = Base.iterate(headers.data, state)
# Base.start(headers::Headers) = Base.start(headers.data)
# Base.done(headers::Headers, state) = Base.done(headers.data, state)
# Base.next(headers::Headers, state) = Base.next(headers.data, state)
Base.isempty(headers::Headers) = Base.isempty(headers.data)
Base.length(headers::Headers) = Base.length(headers.data)
Base.eltype(headers::Headers) = Base.eltype(headers.data)
Base.keytype(headers::Headers) = Base.keytype(headers.data)
Base.valtype(headers::Headers) = Base.valtype(headers.data)
Base.reduce(fn, headers::Headers; init=nothing) = Base.reduce(fn, headers.data; init=init)
Base.keys(headers::Headers) = Base.keys(headers.data)
Base.values(headers::Headers) = Base.values(headers.data)
function Base.merge(headers::Headers, others::Headers...)
    if isempty(others)
        headers
    else
        (head, tail) = (others[1], others[2:end])
        for (k,v) in head
            headers = set(headers, k, v)
        end

        Base.merge(headers, tail)
    end
end


text:
using FFTW
using Plots

n = 128
L = 30
dx = L/n
t = -L/2:dx:L/2-dx
f(x) = cos(x)*exp(-x^2/25) #Function
df(x) = -sin(x)*exp(-x^2/25) - 2*x*f(x)/25 #Derivative

# Approximate derivative using finite Differences...
dfFD = zeros(n)
for k = 1:n-1
    dfFD[k] = (f(t[k+1])-f(t[k]))/dx
end
dfFD[n] = dfFD[n-1];

# Derivative using FFT (spectral derivative)
fhat = fft(f.(t))
kappa = (2pi/L)*(-n/2:n/2-1)
kappa = fftshift(kappa) # Re-order fft frequencies
dfhat = im*kappa.*fhat
dfFFT = real(ifft(dfhat))

# Plotting commands
plot(t, f)
plot!(t, df)
plot!(t, dfFD)
plot!(t, dfFFT)


text:
mutable struct TestMDP <: MDP{Bool, Bool} end
mutable struct TestPOMDP <: POMDP{Bool, Bool, Bool} end

@test actions(TestMDP()) == (true, false)
@test actions(TestPOMDP()) == (true, false)

a = [1,2,3]
@test iterator(a) == a
@test iterator((1,2,3)) == (1,2,3)

@test states(TestMDP()) == (true, false)
@test states(TestPOMDP()) == (true, false)

@test observations(TestPOMDP()) == (true, false)
@test n_observations(TestPOMDP()) == 2

@test state_index(TestMDP(), 1) == 1
@test action_index(TestMDP(), 2) == 2
@test obs_index(TestMDP(), 3) == 3


text:
using Makie
using AbstractPlotting: textslider

"""
    example by @pbouffard from JuliaPlots/Makie.jl#307
    https://github.com/pbouffard/miniature-garbanzo/
"""

struct Joint
    angle::Float32
    axis::Vec3f0
    offset::Vec3f0
end
function AbstractPlotting.plot!(p::Plot(Joint))
    triad = linesegments!(
        p, [Point3f0(0) => unit(Point3f0, i) for i in 1:3],
        color = [:red, :green, :blue],
        linewidth = 3
    ).plots[end]
    on(p[1]) do joint
        translate!(triad, joint.offset)
        rotate!(triad, joint.axis, joint.angle)
        return
    end
    return
end

function Joint(j::Joint; offset::Point3f0 = (0,0,0), axis=(0, 1, 0), angle=0)
    jnew = Joint(j.scene)
    translate!(jnew.scene, j.offset)
    linesegments!(jnew.scene, [Point3f0(0) => offset], linewidth=4, color=:magenta, show_axis=false)
    jnew.axis = axis
    jnew.offset = offset
    setangle!(jnew, angle)
    return jnew
end

scene = Scene(camera = cam3d!)

joints = Node([
    Joint(0, (0,0,1), (0, 0, 1)),
    Joint(-pi/4, (0,1,0), (3,0,0)),
    Joint(pi/2, (0,1,0), (3,0,0)),
    Joint(-pi/4, (0,1,0), (3,0,0)),
    Joint(0, (0,1,0), (3,0,0)),
    Joint(0, (0,1,0), (3,0,0)),
])

[]
sliders = map(1:length(joints)) do i
    slider, val = textslider(-180.0:1.0:180.0, "Joint $(i)", start=rad2deg(joints[i].angle))
    on(val) do x
        setangle!(joints[i], deg2rad(x))
    end
    slider
end

# Add sphere to end effector:
mesh!(joints[end].scene, Sphere(Point3f0(0.5, 0, 0), 0.25f0), color=:cyan, show_axis=false)
center!(s)


text:
# Generate Small example, d=800, n=1601, 2n=3202
example='L'
K,M=GenerateExamples(example)

# Create DampedSystem structure
DS=DampedSystem(M,K,[1.0 0;0 1])

# Compute changes of basis
Ω,Φ,Ξ=ChangeOfBasis(DS)

# Eigenvalue decomposition of a CSymDPR1 matrix.
# l is position of damper, ρ is viscosity
T=ComplexF64
l=50
ρ=50.0
y=Φ[l,:]
A=CSymDPR1(Ξ,y,map(T,ρ))
println("Time of single eigenvalue decomposition of a CSymDPR1 matrix.")
println("Second timing is relevant\n")
@time Λ,S=eigen(A)
@time Λ,S=eigen(A)

# Computation of trace(X) for a single damper in position l with viscosity ρ, and
# the number of damped frequencies s
s=27

function f(ρ::Number)
    A=CSymDPR1(Ξ,Φ[l,:],ComplexF64(ρ))
    Λ,S=eigen(A)
    Tr=traceX(S,s)
    return real(Tr)
end
println("\nTime of trace computation for a single damper.")
println("Second timing is relevant.\n")
@time Tr=f(50.0)
@time Tr=f(50.0)

#  MULTIPLE DAMPERS WITH DIFFERENT VISCOSITIES
# Computation of trace(X) for three dampers in positions l with viscosities ρ, and
# the number of damped frequencies s

# Positions of dampers and viscosities
l=[50,950,220]
ρ=[100.0,200.0,300.0]

# Auxiliary arrays
n₁=800
n=2*n₁+1
U=Matrix{T}(undef,2n,length(l))
Uvec=Matrix{T}(undef,2n,length(l))

function f(ρ::Vector)
    Uvec[:,1]=Φ[l[1],:]
    U[:,1]=Φ[l[1],:]
    A=CSymDPR1(Ξ,Φ[l[1],:],ComplexF64(ρ[1]))
    Λ,S=eigen(A)
    y=similar(Λ)
    # Loop
    for j=2:length(l)
        U[:,j]=Φ[l[j],:]
        if j==2
           U[:,j]-=Φ[l[j+1]+n₁,:]
        end
        lmult!(y,S,U[:,j])
        Uvec[:,j]=y
        A=CSymDPR1(Λ,y,ComplexF64(ρ[j]))
        Λ,S₁=eigen(A)
        # Multiplication of linked Cauchy-like matrices
        S=S*S₁
    end
    Tr=traceX(S,s)
    println("ρ = ", ρ," trace = ", Tr)
    return real(Tr)
end

println("\nTime of trace computation for three dampers.")
println("Second timing is relevant.\n")
@time Tr=f(ρ)
@time Tr=f(ρ)

# Optimization. Given starting ρ, we compute the optmial viscosities
ρ=[100.0,100,100]
println("\nTime for optimization of viscosities for three dampers.\n")
@time result=optimize(f,ρ,method=ConjugateGradient(eta=0.01),
    x_tol=1e-2,f_tol=1e-6)

println(result)
println("\nOptimal viscosities: ",result.minimizer)
println("\nMinimal trace: ", result.minimum)


text:
import ReverseDiffExpressionsBase: RESERVED_INCREMENT_SEED_RESERVED!

using PaddedMatrices: UninitializedVector

struct ∂DiagLowerTri∂Diag{M,T,L <: AbstractLowerTriangularMatrix{M,T}}
    data::L
end
struct ∂DiagLowerTri∂LowerTri{M,T,V <: AbstractFixedSizeVector{M,T}}
    data::LinearAlgebra.Diagonal{T,V}
end
@inline LinearAlgebra.adjoint(x::∂DiagLowerTri∂Diag) = x
@inline LinearAlgebra.adjoint(x::∂DiagLowerTri∂LowerTri) = x

@inline function RESERVED_INCREMENT_SEED_RESERVED!(
    seedout::AbstractFixedSizeVector,
    seedin::AbstractLowerTriangularMatrix,
    jac::∂DiagLowerTri∂Diag
)
    row_sum_prod_add!(seedout, jac.data, seedin); nothing
end
@inline function RESERVED_INCREMENT_SEED_RESERVED!(
    seedout::UninitializedVector,
    seedin::AbstractLowerTriangularMatrix,
    jac::∂DiagLowerTri∂Diag
)
    row_sum_prod!(seedout, jac.data, seedin); nothing
end
@inline function RESERVED_INCREMENT_SEED_RESERVED!(
    # seedout::Union{<:AbstractMutableFixedSizeVector,<:AbstractLowerTriangularMatrix},
    seedout::AbstractLowerTriangularMatrix,
    jac::∂DiagLowerTri∂LowerTri,
    seedin::AbstractLowerTriangularMatrix
)
    muladd!(seedout, jac.data, seedin); nothing
end
@inline function RESERVED_INCREMENT_SEED_RESERVED!(
    # seedout::Union{<:UnitializedVector,<:UninitializedLowerTriangularMatrix},
    seedout::UninitializedLowerTriangularMatrix,
    jac::∂DiagLowerTri∂LowerTri,
    seedin::AbstractLowerTriangularMatrix
)
    mul!(seedout, jac.data, seedin); nothing
end
# @inline function RESERVED_MULTIPLY_SEED_RESERVED(
    # sp::StackPointer,
    # seedin::AbstractLowerTriangularMatrix,
    # jac::∂DiagLowerTri∂LowerTri
# )
    # *(sp, jac.data, seedin)
# end
@generated function RESERVED_INCREMENT_SEED_RESERVED(
    A::AbstractMutableDiagMatrix{M,T},
    D::LinearAlgebra.Diagonal{T,<:AbstractFixedSizeVector{M,T}}
) where {M,T}
    quote
        $(Expr(:meta,:inline))
        LoopVectorization.@vvectorize $T for m ∈ 1:$M
            A[m] = D[m] + A[m]
        end
    end
end
@generated function RESERVED_INCREMENT_SEED_RESERVED!(
    B::AbstractLowerTriangularMatrix{P,T,L},
    A::AbstractLowerTriangularMatrix{P,T,L}
) where {P,T,L}
    quote
        $(Expr(:meta,:inline))
        @vvectorize $T for l in 1:$L
            B[l] += A[l]
        end
    end
end


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

module CrossfilterCharts

using DataFrames

export
	dc,

	DCOut,
	DCChart,

	BaseChart,
	ColorChart,
	CoordinateGridChart,
	PieChart,
	BarChart,
	LineChart,

	rowchart,
	barchart,
	piechart,
	linechart,
	bubblechart,
	datacountwidget,
	datatablewidget,

	get_group_by_name,
	get_groups_by_col,
	get_dim_by_col,
	get_charts,

	add_chart!,
	add_widget!,
	add_group!,
	add_datacountwidget!,
	add_datatablewidget!,
	add_bubblechart!,

	remove_group!,
	remove_chart!,

	quick_add!,
	clear_charts!,
	clear_widgets!,
	set_elastic_height!,
	randomize_ids!,

	reduce_count,
	reduce_sum,
	reduce_master,
	can_infer_chart,
	infer_dimension,
	infer_group,
	infer_chart,

	infer_dimensions!,
	infer_groups!

include("dimensions.jl")
include("groups.jl")
include("charts.jl")
include("dcout.jl")
include("io.jl")

end # module


text:
@doc raw"""
This formulation of the duration constraints adds over the start times looking backwards.

# LaTeX

* Minimum up-time constraint:

If ``t \leq d_{min}^{up} - d_{init}^{up}`` and ``d_{init}^{up} > 0``

`` 1 + \sum_{i=t-d_{min}^{up} + 1}^t x_i^{start} - x_t^{on} \leq 0 ``

for i in the set of time steps. Otherwise:

`` \sum_{i=t-d_{min}^{up} + 1}^t x_i^{start} - x_t^{on} \leq 0 ``

for i in the set of time steps.

* Minimum down-time constraint:

If ``t \leq d_{min}^{down} - d_{init}^{down}`` and ``d_{init}^{down} > 0``

`` 1 + \sum_{i=t-d_{min}^{down} + 1}^t x_i^{stop} + x_t^{on} \leq 1 ``

for i in the set of time steps. Otherwise:

`` \sum_{i=t-d_{min}^{down} + 1}^t x_i^{stop} + x_t^{on} \leq 1 ``

for i in the set of time steps.

# Arguments
* psi_container::PSIContainer : the psi_container model built in PowerSimulations
* duration_data::Vector{UpDown} : gives how many time steps variable needs to be up or down
* initial_duration::Matrix{InitialCondition} : gives initial conditions for up (column 1) and down (column 2)
* cons_name::Symbol : name of the constraint
* var_names::Tuple{Symbol, Symbol, Symbol}) : names of the variables
- : var_names[1] : varon
- : var_names[2] : varstart
- : var_names[3] : varstop
"""
function device_duration_retrospective!(
    psi_container::PSIContainer,
    duration_data::Vector{UpDown},
    initial_duration::Matrix{InitialCondition},
    cons_name::Symbol,
    var_names::Tuple{Symbol, Symbol, Symbol},
)
    time_steps = model_time_steps(psi_container)

    varon = get_variable(psi_container, var_names[1])
    varstart = get_variable(psi_container, var_names[2])
    varstop = get_variable(psi_container, var_names[3])

    name_up = middle_rename(cons_name, PSI_NAME_DELIMITER, "up")
    name_down = middle_rename(cons_name, PSI_NAME_DELIMITER, "dn")

    set_names = [device_name(ic) for ic in initial_duration[:, 1]]
    con_up = add_cons_container!(psi_container, name_up, set_names, time_steps)
    con_down = add_cons_container!(psi_container, name_down, set_names, time_steps)

    for t in time_steps
        for (ix, ic) in enumerate(initial_duration[:, 1])
            name = device_name(ic)
            # Minimum Up-time Constraint
            lhs_on = JuMP.GenericAffExpr{Float64, _variable_type(psi_container)}(0)
            for i in (t - duration_data[ix].up + 1):t
                if i in time_steps
                    JuMP.add_to_expression!(lhs_on, varstart[name, i])
                end
            end
            if t <= max(0, duration_data[ix].up - ic.value) && ic.value > 0
                JuMP.add_to_expression!(lhs_on, 1)
            end
            con_up[name, t] =
                JuMP.@constraint(psi_container.JuMPmodel, lhs_on - varon[name, t] <= 0.0)
        end

        for (ix, ic) in enumerate(initial_duration[:, 2])
            name = device_name(ic)
            # Minimum Down-time Constraint
            lhs_off = JuMP.GenericAffExpr{Float64, _variable_type(psi_container)}(0)
            for i in (t - duration_data[ix].down + 1):t
                if i in time_steps
                    JuMP.add_to_expression!(lhs_off, varstop[name, i])
                end
            end
            if t <= max(0, duration_data[ix].down - ic.value) && ic.value > 0
                JuMP.add_to_expression!(lhs_off, 1)
            end
            con_down[name, t] =
                JuMP.@constraint(psi_container.JuMPmodel, lhs_off + varon[name, t] <= 1.0)
        end
    end
    return
end
@doc raw"""
This formulation of the duration constraints looks ahead in the time frame of the model.

# LaTeX

* Minimum up-time constraint:

If ``t \leq d_{min}^{up}``

`` d_{min}^{down}x_t^{stop} - \sum_{i=t-d_{min}^{up} + 1}^t x_i^{on} - x_{init}^{up} \leq 0 ``

for i in the set of time steps. Otherwise:

`` d_{min}^{down}x_t^{stop} - \sum_{i=t-d_{min}^{up} + 1}^t x_i^{on} \leq 0 ``

for i in the set of time steps.

* Minimum down-time constraint:

If ``t \leq d_{min}^{down}``

`` d_{min}^{up}x_t^{start} - \sum_{i=t-d_{min^{down} + 1}^t (1 - x_i^{on}) - x_{init}^{down} \leq 0 ``

for i in the set of time steps. Otherwise:

`` d_{min}^{up}x_t^{start} - \sum_{i=t-d_{min^{down} + 1}^t (1 - x_i^{on}) \leq 0 ``

for i in the set of time steps.

# Arguments
* psi_container::PSIContainer : the psi_container model built in PowerSimulations
* duration_data::Vector{UpDown} : gives how many time steps variable needs to be up or down
* initial_duration::Matrix{InitialCondition} : gives initial conditions for up (column 1) and down (column 2)
* cons_name::Symbol : name of the constraint
* var_names::Tuple{Symbol, Symbol, Symbol}) : names of the variables
- : var_names[1] : varon
- : var_names[2] : varstart
- : var_names[3] : varstop
"""
function device_duration_look_ahead!(
    psi_container::PSIContainer,
    duration_data::Vector{UpDown},
    initial_duration::Matrix{InitialCondition},
    cons_name::Symbol,
    var_names::Tuple{Symbol, Symbol, Symbol},
)
    time_steps = model_time_steps(psi_container)
    varon = get_variable(psi_container, var_names[1])
    varstart = get_variable(psi_container, var_names[2])
    varstop = get_variable(psi_container, var_names[3])

    name_up = middle_rename(cons_name, PSI_NAME_DELIMITER, "up")
    name_down = middle_rename(cons_name, PSI_NAME_DELIMITER, "dn")

    set_names = [device_name(ic) for ic in initial_duration[:, 1]]
    con_up = add_cons_container!(psi_container, name_up, set_names, time_steps)
    con_down = add_cons_container!(psi_container, name_down, set_names, time_steps)

    for t in time_steps
        for (ix, ic) in enumerate(initial_duration[:, 1])
            name = device_name(ic)
            # Minimum Up-time Constraint
            lhs_on = JuMP.GenericAffExpr{Float64, _variable_type(psi_container)}(0)
            for i in (t - duration_data[ix].up + 1):t
                if i in time_steps
                    JuMP.add_to_expression!(lhs_on, varon[name, i])
                end
            end
            if t <= duration_data[ix].up
                lhs_on += ic.value
            end
            con_up[name, t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                varstop[name, t] * duration_data[ix].up - lhs_on <= 0.0
            )
        end

        for (ix, ic) in enumerate(initial_duration[:, 2])
            name = device_name(ic)
            # Minimum Down-time Constraint
            lhs_off = JuMP.GenericAffExpr{Float64, _variable_type(psi_container)}(0)
            for i in (t - duration_data[ix].down + 1):t
                if i in time_steps
                    JuMP.add_to_expression!(lhs_off, (1 - varon[name, i]))
                end
            end
            if t <= duration_data[ix].down
                lhs_off += ic.value
            end
            con_down[name, t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                varstart[name, t] * duration_data[ix].down - lhs_off <= 0.0
            )
        end
    end

    return
end

@doc raw"""
This formulation of the duration constraints considers parameters.

# LaTeX

* Minimum up-time constraint:

If ``t \leq d_{min}^{up}``

`` d_{min}^{down}x_t^{stop} - \sum_{i=t-d_{min}^{up} + 1}^t x_i^{on} - x_{init}^{up} \leq 0 ``

for i in the set of time steps. Otherwise:

`` \sum_{i=t-d_{min}^{up} + 1}^t x_i^{start} - x_t^{on} \leq 0 ``

for i in the set of time steps.

* Minimum down-time constraint:

If ``t \leq d_{min}^{down}``

`` d_{min}^{up}x_t^{start} - \sum_{i=t-d_{min^{down} + 1}^t (1 - x_i^{on}) - x_{init}^{down} \leq 0 ``

for i in the set of time steps. Otherwise:

`` \sum_{i=t-d_{min}^{down} + 1}^t x_i^{stop} + x_t^{on} \leq 1 ``

for i in the set of time steps.

# Arguments
* psi_container::PSIContainer : the psi_container model built in PowerSimulations
* duration_data::Vector{UpDown} : gives how many time steps variable needs to be up or down
* initial_duration_on::Vector{InitialCondition} : gives initial number of time steps variable is up
* initial_duration_off::Vector{InitialCondition} : gives initial number of time steps variable is down
* cons_name::Symbol : name of the constraint
* var_names::Tuple{Symbol, Symbol, Symbol}) : names of the variables
- : var_names[1] : varon
- : var_names[2] : varstart
- : var_names[3] : varstop
"""
function device_duration_parameters!(
    psi_container::PSIContainer,
    duration_data::Vector{UpDown},
    initial_duration::Matrix{InitialCondition},
    cons_name::Symbol,
    var_names::Tuple{Symbol, Symbol, Symbol},
)
    time_steps = model_time_steps(psi_container)

    varon = get_variable(psi_container, var_names[1])
    varstart = get_variable(psi_container, var_names[2])
    varstop = get_variable(psi_container, var_names[3])

    name_up = middle_rename(cons_name, PSI_NAME_DELIMITER, "up")
    name_down = middle_rename(cons_name, PSI_NAME_DELIMITER, "dn")

    set_names = [device_name(ic) for ic in initial_duration[:, 1]]
    con_up = add_cons_container!(psi_container, name_up, set_names, time_steps)
    con_down = add_cons_container!(psi_container, name_down, set_names, time_steps)

    for t in time_steps
        for (ix, ic) in enumerate(initial_duration[:, 1])
            name = device_name(ic)
            @assert typeof(ic.value) == PJ.ParameterRef
            # Minimum Up-time Constraint
            lhs_on = JuMP.GenericAffExpr{Float64, _variable_type(psi_container)}(0)
            for i in (t - duration_data[ix].up + 1):t
                if t <= duration_data[ix].up
                    if in(i, time_steps)
                        JuMP.add_to_expression!(lhs_on, varon[name, i])
                    end
                else
                    JuMP.add_to_expression!(lhs_on, varstart[name, i])
                end
            end
            if t <= duration_data[ix].up
                lhs_on += ic.value
                con_up[name, t] = JuMP.@constraint(
                    psi_container.JuMPmodel,
                    varstop[name, t] * duration_data[ix].up - lhs_on <= 0.0
                )
            else
                con_up[name, t] = JuMP.@constraint(
                    psi_container.JuMPmodel,
                    lhs_on - varon[name, t] <= 0.0
                )
            end
        end

        for (ix, ic) in enumerate(initial_duration[:, 2])
            @assert typeof(ic.value) == PJ.ParameterRef
            name = device_name(ic)
            # Minimum Down-time Constraint
            lhs_off = JuMP.GenericAffExpr{Float64, _variable_type(psi_container)}(0)
            for i in (t - duration_data[ix].down + 1):t
                if t <= duration_data[ix].down
                    if in(i, time_steps)
                        JuMP.add_to_expression!(lhs_off, (1 - varon[name, i]))
                    end
                else
                    JuMP.add_to_expression!(lhs_off, varstop[name, i])
                end
            end
            if t <= duration_data[ix].down
                lhs_off += ic.value
                con_down[name, t] = JuMP.@constraint(
                    psi_container.JuMPmodel,
                    varstart[name, t] * duration_data[ix].down - lhs_off <= 0.0
                )
            else
                con_down[name, t] = JuMP.@constraint(
                    psi_container.JuMPmodel,
                    lhs_off + varon[name, t] <= 1.0
                )
            end
        end
    end
    return
end

@doc raw"""
This formulation of the duration constraints adds over the start times looking backwards.

# LaTeX

* Minimum up-time constraint:

`` \sum_{i=t-min(d_{min}^{up}, T)+ 1}^t x_i^{start} - x_t^{on} \leq 0 ``

for i in the set of time steps.

* Minimum down-time constraint:

`` \sum_{i=t-min(d_{min}^{down}, T) + 1}^t x_i^{stop} + x_t^{on} \leq 1 ``

for i in the set of time steps.

# Arguments
* psi_container::PSIContainer : the psi_container model built in PowerSimulations
* duration_data::Vector{UpDown} : gives how many time steps variable needs to be up or down
* initial_duration::Matrix{InitialCondition} : gives initial conditions for up (column 1) and down (column 2)
* cons_name::Symbol : name of the constraint
* var_names::Tuple{Symbol, Symbol, Symbol}) : names of the variables
- : var_names[1] : varon
- : var_names[2] : varstart
- : var_names[3] : varstop
"""
function device_duration_compact_retrospective!(
    psi_container::PSIContainer,
    duration_data::Vector{UpDown},
    initial_duration::Matrix{InitialCondition},
    cons_name::Symbol,
    var_names::Tuple{Symbol, Symbol, Symbol},
)
    time_steps = model_time_steps(psi_container)

    varon = get_variable(psi_container, var_names[1])
    varstart = get_variable(psi_container, var_names[2])
    varstop = get_variable(psi_container, var_names[3])

    name_up = middle_rename(cons_name, PSI_NAME_DELIMITER, "up")
    name_down = middle_rename(cons_name, PSI_NAME_DELIMITER, "dn")

    set_names = [device_name(ic) for ic in initial_duration[:, 1]]
    con_up = add_cons_container!(psi_container, name_up, set_names, time_steps)
    con_down = add_cons_container!(psi_container, name_down, set_names, time_steps)
    total_time_steps = length(time_steps)
    for t in time_steps
        for (ix, ic) in enumerate(initial_duration[:, 1])
            name = device_name(ic)
            # Minimum Up-time Constraint
            lhs_on = JuMP.GenericAffExpr{Float64, _variable_type(psi_container)}(0)
            if t in min(duration_data[ix].up, total_time_steps):total_time_steps
                for i in (t - duration_data[ix].up + 1):t
                    if i in time_steps
                        JuMP.add_to_expression!(lhs_on, varstart[name, i])
                    end
                end
            elseif t <= max(0, duration_data[ix].up - ic.value) && ic.value > 0
                JuMP.add_to_expression!(lhs_on, 1)
            else
                continue
            end
            con_up[name, t] =
                JuMP.@constraint(psi_container.JuMPmodel, lhs_on - varon[name, t] <= 0.0)
        end

        for (ix, ic) in enumerate(initial_duration[:, 2])
            name = device_name(ic)
            # Minimum Down-time Constraint
            lhs_off = JuMP.GenericAffExpr{Float64, _variable_type(psi_container)}(0)
            if t in min(duration_data[ix].down, total_time_steps):total_time_steps
                for i in (t - duration_data[ix].down + 1):t
                    if i in time_steps
                        JuMP.add_to_expression!(lhs_off, varstop[name, i])
                    end
                end
            elseif t <= max(0, duration_data[ix].down - ic.value) && ic.value > 0
                JuMP.add_to_expression!(lhs_off, 1)
            else
                continue
            end
            con_down[name, t] =
                JuMP.@constraint(psi_container.JuMPmodel, lhs_off + varon[name, t] <= 1.0)
        end
    end
    return
end


text:
# Copyright (c) 2022 Code Komali
#
# This software is released under the MIT License.
# https://opensource.org/licenses/MIT

module Utils

using Agents
using InteractiveDynamics
using GLMakie
using LinearAlgebra

flatten(x) = reduce(vcat, x)

reduce_y(line, val) = (line[1] .- (0, val), line[2] .- (0, val))

increase_y(line,val)= (line[1] .+ (0, val), line[2] .+ (0, val))

reduce_x(line, val) = (line[1] .- (val, 0), line[2] .- (val, 0))

increase_x(line, val) = (line[1] .+ (val, 0), line[2] .+ (val, 0))

toVectorPos(tuple::NTuple{2,Float64})= [ round(x, digits=2) for x in tuple ]

function toTuplePos(vector::Vector{Float64})
    length(vector)==2 || return error("Position vector has more than 2 dims")
    return NTuple{2, Float64}(round(i, digits=2) for i in vector)
end

orientation(startPos, endPos) = endPos .- startPos |> toVectorPos |> normalize |> toTuplePos

end


text:
# # Gross-Pitaevskii equation in one dimension
# In this example we will use DFTK to solve
# the Gross-Pitaevskii equation, and use this opportunity to explore a few internals.

# ## The model
# The [Gross-Pitaevskii equation](https://en.wikipedia.org/wiki/Gross%E2%80%93Pitaevskii_equation) (GPE)
# is a simple non-linear equation used to model bosonic systems
# in a mean-field approach. Denoting by ``ψ`` the effective one-particle bosonic
# wave function, the time-independent GPE reads in atomic units:
# ```math
#     H ψ = \left(-\frac12 Δ + V + 2 C |ψ|^2\right) ψ = μ ψ \qquad \|ψ\|_{L^2} = 1
# ```
# where ``C`` provides the strength of the boson-boson coupling.
# It's in particular a favorite model of applied mathematicians because it
# has a structure simpler than but similar to that of DFT, and displays
# interesting behavior (especially in higher dimensions with magnetic fields, see
# [Gross-Pitaevskii equation with magnetism](@ref)).

# We wish to model this equation in 1D using DFTK.
# First we set up the lattice. For a 1D case we supply two zero lattice vectors,
a = 10
lattice = a .* [[1 0 0.]; [0 0 0]; [0 0 0]];

# which is special cased in DFTK to support 1D models.
#
# For the potential term `V` we just pick a harmonic
# potential. The real-space grid is in ``[0,1)``
# in fractional coordinates( see
# [Lattices and lattice vectors](@ref conventions-lattice)),
# therefore:
pot(x) = (x - a/2)^2;

# We setup each energy term in sequence: kinetic, potential and nonlinear term.
# For the non-linearity we use the `PowerNonlinearity(C, α)` term of DFTK.
# This object introduces an energy term ``C ∫ ρ(r)^α dr``
# to the total energy functional, thus a potential term ``α C ρ^{α-1}``.
# In our case we thus need the parameters
C = 1.0
α = 2;

# ... and with this build the model
using DFTK
using LinearAlgebra

n_electrons = 1  # Increase this for fun
terms = [Kinetic(),
         ExternalFromReal(r -> pot(r[1])),
         PowerNonlinearity(C, α),
]
model = Model(lattice; n_electrons=n_electrons, terms=terms,
              spin_polarization=:spinless);  # use "spinless electrons"

# We discretize using a moderate Ecut (For 1D values up to `5000` are completely fine)
# and run a direct minimization algorithm:
Ecut = 500
basis = PlaneWaveBasis(model, Ecut)
scfres = direct_minimization(basis, tol=1e-8) # This is a constrained preconditioned LBFGS
scfres.energies

# ## Internals
# We use the opportunity to explore some of DFTK internals.
#
# Extract the converged density and the obtained wave function:
ρ = real(scfres.ρ.real)[:, 1, 1]  # converged density
ψ_fourier = scfres.ψ[1][:, 1];    # first kpoint, all G components, first eigenvector

# Transform the wave function to real space and fix the phase:
ψ = G_to_r(basis, basis.kpoints[1], ψ_fourier)[:, 1, 1]
ψ /= (ψ[div(end, 2)] / abs(ψ[div(end, 2)]));

# Check whether ``ψ`` is normalised:
x = a * vec(first.(DFTK.r_vectors(basis)))
N = length(x)
dx = a / N  # real-space grid spacing
@assert sum(abs2.(ψ)) * dx ≈ 1.0

# The density is simply built from ψ:
norm(scfres.ρ.real - abs2.(ψ))

# We summarize the ground state in a nice plot:
using Plots

p = plot(x, real.(ψ), label="real(ψ)")
plot!(p, x, imag.(ψ), label="imag(ψ)")
plot!(p, x, ρ, label="ρ")

# The `energy_hamiltonian` function can be used to get the energy and
# effective Hamiltonian (derivative of the energy with respect to the density matrix)
# of a particular state (ψ, occupation).
# The density ρ associated to this state is precomputed
# and passed to the routine as an optimization.
E, ham = energy_hamiltonian(basis, scfres.ψ, scfres.occupation; ρ=scfres.ρ)
@assert sum(values(E)) == sum(values(scfres.energies))

# Now the Hamiltonian contains all the blocks corresponding to kpoints. Here, we just have one kpoint:
H = ham.blocks[1];

# `H` can be used as a linear operator (efficiently using FFTs), or converted to a dense matrix:
ψ11 = scfres.ψ[1][:, 1] # first kpoint, first eigenvector
Hmat = Array(H) # This is now just a plain Julia matrix,
##                which we can compute and store in this simple 1D example
@assert norm(Hmat * ψ11 - H * ψ11) < 1e-10

# Let's check that ψ11 is indeed an eigenstate:
norm(H * ψ11 - dot(ψ11, H * ψ11) * ψ11)

# Build a finite-differences version of the GPE operator ``H``, as a sanity check:
A = Array(Tridiagonal(-ones(N - 1), 2ones(N), -ones(N - 1)))
A[1, end] = A[end, 1] = -1
K = A / dx^2 / 2
V = Diagonal(pot.(x) + C .* α .* (ρ.^(α-1)))
H_findiff = K + V;
maximum(abs.(H_findiff*ψ - (dot(ψ, H_findiff*ψ) / dot(ψ, ψ)) * ψ))


text:
"""
    apply_textstyle!(g, s, parent_font; addset)

Internal function to apply the textstyle `s` in a GLE context.
"""
function apply_textstyle!(
            g::GS,
            s::TextStyle,
            parent_font::String = "";
            #
            addset::Bool = false
        )::Nothing

    if !isdef(s.font) && !isempty(parent_font)
        s.font = parent_font
    end
    isanydef(s) || return

    addset && "\nset" |> g
    add(g, s, :font, :hei, :color)
    return
end

"""
    apply_linestyle!(g, s; nosmooth, addset)

Internal function to apply the linestyle `s` in a GLE context.

## Kwargs

    * nosmooth: while for a line plot the smooth option is allowed, for the
                line property of e.g. a tick, it doesn't make sense.
    * addset: whether to add a set.
"""
function apply_linestyle!(
            g::GS,
            s::LineStyle;
            #
            nosmooth::Bool = false,
            addset::Bool   = false
        )::Nothing

    isanydef(s) || return

    addset && "\nset" |> g
    add(g, s, :lstyle, :lwidth, :color)
    if !nosmooth
        add(g, s, :smooth)
    end
    return
end

"""
    apply_markerstyle!(g, s; mcol, mscale)

Internal function to apply the markerstyle `s` in a GLE context.
"""
function apply_markerstyle!(
            g::GS,
            s::MarkerStyle;
            #
            mcol::Bool      = false,
            mscale::Float64 = 1.0
        )::Nothing

    isanydef(s) || return
    isdef(s.marker) && s.marker == "none" && return

    if !mcol
        add(g, s, :marker, :msize, :color)
    else
        "marker $(str(s))" |> g
        "msize"            |> g
        if isdef(s.msize)
            ms = "$(fl2str(s.msize / mscale))"
        else
            ms = "$(0.4 / mscale)"
        end
        ms |> g
    end
    return
end

"""
    apply_barstyle!(g, s)

Internal function to apply the barstyle `s` in a GLE context.
"""
function apply_barstyle!(
            g::GS,
            s::BarStyle
        )::Nothing

    isanydef(s) || return
    add(g, s, :color, :fill)
    return
end

"""
    apply_barstyles_nostack!(g, v)

Internal function to apply the Vector of barstyles `v` in a GLE context where
the bars are not stacked.
"""
function apply_barstyles_nostack!(
            g::GS,
            v::Vector{BarStyle}
        )::Nothing
    # assumption that if one is defined, all are defined (this is checked
    # with the set_properties!)
    isanydef(v[1]) || return
    if isdef(v[1].color)
        cv = svec2str((col2str(s.color) for s ∈ v))
        "color $cv" |> g
    end
    if isdef(v[1].fill)
        cv = svec2str((col2str(s.fill) for s ∈ v))
        "fill $cv"  |> g
    end
    return
end

"""
    apply_boxplotstyle!(g, s)

Internal function to apply the BoxplotStyle `s` in a GLE context.
Note that this is the context of a subroutine so we just pass raw
arguments, see also `gle_sub/boxplot`.
"""
function apply_boxplotstyle!(
            g::GS,
            s::BoxplotStyle,
            f::Figure
        )::Nothing

    join(
        vec2str.(
            # BOX
            s.bwidth,          # box width
            s.wwidth,          # whiskers width
            s.lstyle,          # box line style
            s.blstyle.lwidth,  # box line width
            s.blstyle.color,   # box line colour
            # MEDIAN LINE
            s.mlstyle.lstyle,   # median line style
            s.mlstyle.lwidth,   # median line width
            s.mlstyle.color,    # median line color
            # MEAN MARKER
            Int(s.mshow),                       # show mean marker
            s.mmstyle.marker,                   # marker symbol
            s.mmstyle.msize / f.textstyle.hei,  # scaled marker size (≈ scatter msize)
            s.mmstyle.color
            ),
        " "
    ) |> g
    return
end


text:
# --------------------------------------------------------------------------- #
# Prepare layout for batches of images, weights, states and prmaps
# For suitable arraytypes, this will transfer the data into the graphics
# memory

function packimage(img :: Image; at = Array{Float32})
 pack = reshape(imgdata(img), imgsize(img)..., imgchannels(img), 1)
 return convert(at, pack)
end

function packimage(imgs :: Vector{<: Image}; at = Array{Float32})
  pack = cat(4, map(packimage, imgs)...)
  return convert(at, pack)
end

function packweights(w :: Vector{Any}; at = Array{Float32})
  return Any[ convert(at, x) for x in w ]
end

function packbnmoments(m; at = Array{Float32})
  mean = (m.mean != nothing) ? convert(at, m.mean) : nothing
  var  = (m.var  != nothing) ? convert(at, m.var)  : nothing
  return bnmoments(momentum = m.momentum, mean = mean, var = var)
end

function packstate(s :: Vector{Any}; at = Array{Float32})
  return Any[ packbnmoments(x, at = at) for x in s ]
end

function packproxymap(prmap; at = Array{Float32})
  if length(size(prmap)) < 3
    return convert(at, reshape(prmap, size(prmap)..., 1))
  else
    return convert(at, prmap)
  end
end

unpackweights(w) = packweights(w, at = Array{Float32})
unpackstate(s) = packstate(s, at = Array{Float32})

# --------------------------------------------------------------------------- #
# Check consistency between Array types

# TODO: This should be possible without relying on explicit
# knowledge of KnetArrays
arraytype(a :: Array{Float32}) = Array{Float32}
arraytype(a :: KnetArray{Float32}) = KnetArray{Float32}

# --------------------------------------------------------------------------- #
# Loss function

function loss(w, s, imgdata, prmap,
              mt :: Type{<: Model}; at = nothing, kwargs...)

  # Loss that can be applied to raw data without conversion
  # if `at` is not specified. This can be used to efficiently
  # calculate the loss for already gpu-packed data. If `at`
  # is specified, conversion might take place.

  if at != nothing
    w       = packweights(w, at = at)
    s       = packstate(s,  at = at)
    prmap   = packproxymap(prmap, at = at)
    imgdata = packimage(imgdata, at = at)
  end

  return sum(abs2, density(w, s, imgdata, mt) .- prmap) / length(prmap)
end

function loss(w, s, imgdata, lbls :: Vector{Label},
              mt :: Type{<: Model}; at = nothing, kwargs...)

  # This function will build proximity maps from the given labels.
  # If the argument `at` is not specified, these proximity maps
  # inherit their arraytype from `imgdata`

  _at = (at != nothing) ? at : arraytype(imgdata)

  imgdata = packimage(imgdata, at = _at)
  prmap   = proxymap(imgsize(imgdata)..., lbls; kwargs...)
  prmap   = packproxymap(prmap)
  return loss(w, s, imgdata, prmap, mt; at = at, kwargs...)
end

function loss{I <: Image}(w, s, img :: I, lbl :: Label,
                          mt :: Type{<: Model{I}}; at = nothing, kwargs...)

  # Loss calculated for an image
  # Creates proximity maps and then calculates the loss

  _at = (at != nothing) ? at : arraytype(imgdata(img))

  x = packimage(img, at = _at)
  return loss(w, s, x, [lbl], mt; at = at, kwargs...)
end

function loss(m :: Model, img :: GreyscaleImage, lbl :: Label)
  return loss(weights(m), state(m), img, lbl, typeof(m))
end

# --------------------------------------------------------------------------- #
# Loss gradient

_gradloss = gradloss(loss)  # Calculate both the gradient and loss value
_grad     = grad(loss)

lossgrad(args...; kwargs...) = _gradloss(args...; kwargs...)[2:-1:1]

function lossgrad{I <: Image}(m :: Model{I}, img :: I, lbl :: Label; kwargs...)
  return lossgrad(weights(m), state(m), img, lbl, typeof(m); kwargs...)
end

# --------------------------------------------------------------------------- #
# Divide dataset in batches

function makebatches(imgs, lbls, batchsize; shuffle = false)
  n = length(imgs)
  idx = shuffle ? Base.shuffle(collect(1:n)) : collect(1:n)
  sel = collect(1:batchsize:n)
  m = length(sel)

  idxs = [ idx[sel[i] : (i < m ? sel[i+1] - 1 : n)] for i in 1:m ]
  return [ (imgs[range], lbls[range]) for range in idxs ]
end

function packbatches(batches; at = Array{Float32}, kwargs...)
  return map(batches) do batch
    imgs   = packimage(batch[1], at = at)
    prmaps = proxymap(imgsize(imgs)..., batch[2]; kwargs...)
    prmaps = packproxymap(prmaps, at = at)
    return (imgs, prmaps)
  end
end

# --------------------------------------------------------------------------- #
# High level training function

"""
    train!(model, imgs, lbls; kwargs...)

    train!(model, lmgs; kwargs...)

Train a model `model` using the images `imgs` and labels `lbls`. Alternatively, an iterable `lmgs` of labeled images can be provided.

# Arguments
- `testset`: an iterable of labeled images to be used for testing the network.
- `epochs = 10`: the number of epochs (iterations of the training images)
  used for training.
- `batchsize = 1`: the size of a batch used for one update step of the model.
- `shuffle = true`: whether to shuffle the training images for each epoche
- `log = true`: whether to print information about the model's training
  performance.
- `logpath`: redirect logging to this file.
- `modelpath`: file to store the trained model in
- `record = true`: record the training performance in a `Dict` and return it.
- `patchsize`: providing an integer for this option causes that the
  effective training images are patches of this size, generated from the
  input image during each epoche.
- `patchmode`: if `patchmode < 1`, create patches by cutting ordered
  patches from the given input images. If `patchmode > 0`, create
  `patchmode` patches by random selection.
- `at = Array{Float32}`: the arraytyped used for training and evaluating.
  Currently, this is restricted to `Array{Float32}` for cpu computation or
  `KnetArray{Float32}` for gpu computation.
- `kernelsize = 7`: size of the kernel used to construct proximity maps.
- `peakheight = 100`: height of the kernels used to construct proximity
  maps.
- `opt = Adam`: Knet optimizer to be used. All other keyword arguments are
  passed to the constructor of the optimizer.
"""

function train!(model :: Model, imgs, lbls;
                epochs = 10,
                opt = Adam,
                log = true,
                logpath = nothing,
                record = true,
                batchsize = 1,
                patchsize = nothing,
                patchmode = 0,
                shuffle = false,
                modelpath = nothing,
                testset = nothing,
                kernelsize = 7,
                peakheight = 100.,
                imageop = Id(),
                at = nothing,
                kwargs...)

  # Sanity checks

  @assert (length(imgs) == length(lbls))
  if patchsize == nothing
    @assert all(x -> x == imgsize(imgs[1]), imgsize.(imgs))
  end

  # Pack the model weights and state.
  # In practice this means to transfer them to
  # gpu memory if a gpu based arratype is given.

  if at == nothing
    at = gpu() >= 0 ? KnetArray{Float32} : Array{Float32}
  end

  w = packweights(weights(model), at = at)
  s = packstate(state(model), at = at)

  # Initialize the optimizers and conduct some
  # general preparations

  optim = optimizers(w, opt; kwargs...)
  test  = (testset != nothing) && (log || record)

  if patchsize == nothing
    n = ceil(Int, length(imgs) / batchsize)
  elseif patchmode > 0
    n = ceil(Int, length(imgs) * patchmode / batchsize)
  else
    # WARNING: This will not work if margins or manual offsets are given to
    # `ordered_patches`
    m1 = floor(imgsize(imgs[1])[1] / patchsize)
    m2 = floor(imgsize(imgs[1])[2] / patchsize)
    n = ceil(Int, length(imgs) * m1 * m2 / batchsize)
  end

  if record
    rec = Dict{Symbol, Array{Float64}}(
               :loss_train    => zeros(n, epochs),
               :count_train   => zeros(n, epochs),
               :meanadj_train => zeros(n, epochs),
               :maxadj_train  => zeros(n, epochs))
    if test
      rec[:loss_test]    = zeros(epochs)
      rec[:count_test]   = zeros(epochs)
      rec[:meanadj_test] = zeros(epochs)
      rec[:maxadj_test]  = zeros(epochs)
    end
  end

  # Bring the testset in the correct format

  if test
    if typeof(testset) <: NTuple{2}
      timgs, tlbls = testset
    else
      timgs = [ t[1] for t in testset ]
      tlbls = [ t[2] for t in testset ]
    end
    tbatches = makebatches(timgs, tlbls, batchsize, shuffle = false)
    tpacks   = packbatches(tbatches, at = at,
                           peakheight = peakheight,
                           kernelsize = kernelsize)
  end

  # Print overview

  if log
    @printf "# CellCC training progress\n"
    @printf "#\n"
    @printf "# model      %s\n" typeof(model)
    @printf "# date       %s\n" now()
    @printf "# images     %d\n" length(imgs)
    @printf "# imagesize  %s\n" imgsize(imgs[1])
    @printf "# patchsize  %s\n" patchsize
    @printf "# patchmode  %s\n" patchmode
    @printf "# batches    %d\n" n
    @printf "# batchsize  %d\n" batchsize
    @printf "# epochs     %d\n" epochs
    @printf "# testset    %s\n" test ? length(timgs) : "-"
    @printf "# shuffle    %s\n" shuffle
    @printf "# record     %s\n" record
    @printf "# optimizer  %s\n" optim[1]
    @printf "# modelpath  %s\n" modelpath == nothing ? "-" : modelpath
    @printf "# logpath    %s\n" logpath == nothing ? "-" : logpath
  end

  # Start the training loop

  time = @elapsed for i in 1:epochs

    # Generate batches used for training. These batches
    # are also 'packed', meaning that they are brought in
    # a dense format consisting of image and proximity
    # map data arrays. Depending on `at`, this might
    # transfer the batch-data to gpu memory.

    # TODO: make the batch-acquicision process more dynamic!

    if patchsize != nothing && patchmode > 0
      patches = random_patches.(imgs, lbls, patchmode,
                                imageop = imageop,
                                size = (patchsize, patchsize))

    elseif patchsize != nothing && patchmode <= 0
      patches = ordered_patches.(imgs, lbls,
                                 imageop = imageop,
                                 size = (patchsize, patchsize))

    else
      patches = apply.(imageop, imgs, lbls)
    end

    timgs, tlbls = vcat(first.(patches)...), vcat(second.(patches)...)

    batches = makebatches(timgs, tlbls, batchsize, shuffle = shuffle)
    packs   = packbatches(batches, at = at,
                          peakheight = peakheight,
                          kernelsize = kernelsize)

    if log
      @printf "\n"
      @printf "# Epoche  Batch   Loss  Count  MeanAdj   MaxAdj \n"
      @printf "# ----------------------------------------------\n"
    end

    # Training epoche

    for j in 1:n

      # Make sure that everything is printed

      flush(STDOUT)
      flush(STDERR)

      # Extract image data and proximity maps

      data, prmaps = packs[j]
      lbl = batches[j][2]

      l, lg = lossgrad(w, s, data, prmaps, typeof(model))
      update!(w, lg, optim)

      if log || record

        # Let the model predict labels

        dlbl = label(w, s, data, typeof(model))

        # Evaluate the quality of the prediction

        c      = mean(length.(dlbl) ./ length.(lbl))
        adj    = adjacency.(dlbl, lbl)
        mm, mx = mean(mean, adj), mean(maximum, adj)

        if log
          @printf("  %6d  %5d  %5.1f  %5.3f  %7.2f  %7.2f\n",
                  i, j, l, c, mm, mx)
        end

        if record
          rec[:loss_train][j, i]    = l
          rec[:count_train][j, i]   = c
          rec[:meanadj_train][j, i] = mm
          rec[:maxadj_train][j, i]  = mx
        end
      end
    end

    # Evaluate performance on the test set

    if test

      l = c = mm = mx = 0.

      for j in 1:length(tpacks)
        data, prmaps = tpacks[j]
        timg, tlbl = tbatches[j]

        dlbl = label(w, s, data, typeof(model))
        adj  = adjacency.(dlbl, tlbl)

        l  += loss(w, s, data, prmaps, typeof(model)) * length(timg)
        c  += sum(length.(dlbl) ./ length.(tlbl))
        mm += sum(mean, adj)
        mx += sum(maximum, adj)
      end

      l  /= length(timgs)
      c  /= length(timgs)
      mm /= length(timgs)
      mx /= length(timgs)

      if log
          @printf "\n"
          @printf("  %6d      0  %5.1f  %5.3f  %7.2f  %7.2f\n",
                   i, l, c, mm, mx)
      end

      if record
        rec[:loss_test][i]    = l
        rec[:count_test][i]   = c
        rec[:meanadj_test][i] = mm
        rec[:maxadj_test][i]  = mx
      end
    end
  end

  # Update the model weights and its state

  weights(model)[:] = unpackweights(w)
  state(model)[:]   = unpackstate(s)

  if log
    @printf "\n"
    @printf "# End of training process after %.1f seconds\n" time
  end

  if logpath != nothing
    # TODO
    @printf "# Warning: Logpath not yet implemented\n"
    @printf "# Did not save a version of this output in file %s\n" logpath
  end

  if modelpath != nothing
    modelsave(modelpath, model)
  end

  return record ? rec : nothing
end

function train!(model :: Model, data; kwargs...)
  imgs = [ d[1] for d in data ]
  lbls = [ d[2] for d in data ]
  train!(model, imgs, lbls; kwargs...)
end

function train!(model :: Model, gen :: Function; samples = 10, kwargs...)
  data = [gen() for i in 1:samples]
  train!(model, data; kwargs...)
end

# --------------------------------------------------------------------------- #
# Function for training with a single image frame
# Intended for testing purposes

function train!(model :: Model, img :: Image, lbl :: Label;
                epochs = 10, record = true, opt = Adam,
                log = true, logpath = nothing, modelpath = nothing, kwargs...)

  optim = optimizers(weights(model), opt; kwargs...)

  if record
    rec = Dict{Symbol, Vector{Float64}}(
               :loss    => zeros(epochs),
               :count   => zeros(epochs),
               :meanadj => zeros(epochs),
               :maxadj  => zeros(epochs))
  end

  if log
    @printf "# CellCC training progress (single image mode)\n"
    @printf "#\n"
    @printf "# model      %s\n" typeof(model)
    @printf "# date       %s\n" now()
    @printf "# epochs     %d\n" epochs
    @printf "# record     %s\n" record
    @printf "# optimizer  %s\n" optim[1]
    @printf "# modelpath  %s\n" modelpath == nothing ? "-" : modelpath
    @printf "# logpath    %s\n" logpath == nothing ? "-" : logpath
    @printf "\n"
    @printf "# Epoche   Loss  Count  MeanAdj   MaxAdj\n"
    @printf "# --------------------------------------\n"
  end

  time = @elapsed for i in 1:epochs
    lg, l = lossgrad(model, img, lbl)
    update!(weights(model), lg, optim)

    if log || record

      dens = density(model, img)
      dlbl = label(dens)

      c      = length.(dlbl) ./ length.(lbl)
      adj    = adjacency(dlbl, lbl)
      mm, mx = mean(adj), maximum(adj)

      if log
        @printf("  %6d  %5.1f  %5.3f  %7.2f  %7.2f\n",
                 i, l, c, mm, mx)
      end

      if record
        rec[:loss][i]    = l
        rec[:count][i]   = c
        rec[:meanadj][i] = mm
        rec[:maxadj][i]  = mx
      end
    end
  end

  @printf "\n"
  @printf "# End of training process after %.1f seconds\n" time

  if logpath != nothing
    # TODO
    @printf "# Warning: Logpath not yet implemented\n"
    @printf "# Did not save a version of this output in file %s\n" logpath
  end

  if modelpath != nothing
    modelsave(model, modelpath)
    @printf "# Model saved at %s\n" modelpath
  end

  return record ? rec : nothing
end

max_stars_repo_name: tscode/DCellC

max_issues_repo_name: tscode/DCellC


text:
# `topological_sort`
let
    # simple graph
    G = DiGraph(4)
    add_edge!(G,1,3)
    add_edge!(G,1,4)
    add_edge!(G,2,3)
    add_edge!(G,3,4)
    ordering = topological_sort(G)
    @test ordering == [1,2,3,4]
end
let
    # disjoint graph
    G = DiGraph(4)
    add_edge!(G,1,2)
    add_edge!(G,3,4)
    ordering = topological_sort(G)
    @test length(ordering) == nv(G)
    @test findfirst(ordering .== 1) < findfirst(ordering .== 2)
    @test findfirst(ordering .== 3) < findfirst(ordering .== 4)
end
let
    # cyclic graph (should throw error)
    G = DiGraph(3)
    add_edge!(G,1,2)
    add_edge!(G,2,3)
    add_edge!(G,3,1)
    @test_throws AssertionError topological_sort(G)
end
# `find_index_in_sorted_array` and `insert_to_sorted_array`
let
    array = [1.0,2.0,3.0,4.0,5.0]
    array = insert_to_sorted_array!(array, 3.5)
    @test(array[4] == 3.5)
    array = insert_to_sorted_array!(array, 0.5)
    @test(array[1] == 0.5)
    array = insert_to_sorted_array!(array, 5.5)
    @test(array[end] == 5.5)
end
let
    v = [0,1,2,4,5,6]
    L = length(v)
    x = 3
    @test find_index_in_sorted_array(v,x) == 4
    insert_to_sorted_array!(v,x)
    @test length(v) == L + 1
    v[4] == x
end
let
    v = [0,1,1,1,1,2]
    L = length(v)
    x = 1
    @test find_index_in_sorted_array(v,x) == 2
end
let
    v = Vector{Int}()
    insert_to_sorted_array!(v,1)
    @test length(v) == 1
end


text:
using RandomizedLinAlg
using Test, Random, LinearAlgebra

@testset "ID" begin
    Random.seed!(1)
    M = randn(4,5)
    k = 3

    F = id(M, k, 3)
    @test norm(F.B * F.P - M) ≤ 2svdvals(M)[k + 1]
end


text:
# Raytracer.jl
# Raytracing for the generation of photorealistic images in Julia
# Copyright (c) 2021 Samuele Colombo, Paolo Galli

# Parser utilities

"""
    evaluate_math_expression(token::Token{MathExpression}, scene::Scene)

Replace all identifiers in the mathematical expression stored in the [`MathExpression`](@ref) token and then evaluate it.
"""
function evaluate_math_expression(token::Token{MathExpression}, scene::Scene)
    vars = scene.variables
    expr = token.value.value
    args = map(expr.args[begin + 1: end]) do arg
        if isa(arg, Symbol)
            arg === :TIME && return scene.time
            (type = findfirst(d -> haskey(d, arg), vars)) |> isnothing &&
                throw(UndefinedIdentifier(token.loc, "Undefined variable '$arg' in 'MathExpression'", token.length))
            return vars[type][arg].value
        elseif isa(arg, Expr)
            return evaluate_math_expression(Token(token.loc, MathExpression(arg), token.length), scene)
        else
            hasmethod(isfinite, Tuple{typeof(arg)}) &&(
                isfinite(arg) ||
                    throw(InvalidExpression(token.loc, "'MathExpression' should not return or contain infinite or NaN values.", token.length)))
            return arg
        end
    end
    res = try
        Expr(expr.head, expr.args[begin], args...) |> eval
    catch e
        isa(e, ArgumentError) &&
            throw(InvalidExpression(token.loc, "Julia `eval` has thrown an error:\n`$(e.msg)`\nPlease follow the exception instructions or consult Julia documentation", token.length))
        if isa(e, DomainError)
            msg = "DomainError with " * repr( e.val)
            if isdefined(e, :msg)
                msg *= ":\n" * e.msg
            end
            throw(InvalidExpression(token.loc, "Julia `eval` has thrown an error:\n`$(e.msg)`\nPlease follow the exception instructions or consult Julia documentation", token.length))
        end
        rethrow(e)
    end
    hasmethod(isfinite, Tuple{typeof(res)}) &&
        (isfinite(res) ||
            throw(InvalidExpression(token.loc, "'MathExpression' should not return or contain infinite or NaN values.", token.length)))
    res
end

"""
    function generate_kwargs(stream::InputStream, scene::Scene, kw::NamedTuple)

Return a `Dict{Symbol, Any}` as instructed by the given `kw` `NamedTuple`.

The argument `kw` pairs all the keywords a constructor need with the parsing functions they need.
This function parses keywords and/or values until it reaches the end of the constructor.
The arguments may also be positional, but positional arguments must not follow keyword arguments.
If a keyword is used more than once an exception is thrown.
"""
function generate_kwargs(stream::InputStream, scene::Scene, kw::NamedTuple)
    expect_symbol(stream, Symbol("("))
    kwargs = Dict{Symbol, Any}()
    sizehint!(kwargs, length(kw))
    is_positional_allowed = true
    for i ∈ SOneTo(length(kw))
        token = read_token(stream)
        isa(token.value, LiteralSymbol) && token.value.value == Symbol(")") && break
        unread_token(stream, token)
        key = if isa(token.value, Keyword)
            local key = expect_keyword(stream, keys(kw)).value.value
            haskey(kwargs, key) &&
                throw(InvalidKeyword(token.loc, "Argument `$key` has already been defined", token.length))
            is_positional_allowed = false
            key
        elseif is_positional_allowed
            keys(kw)[i]
        else
            throw(InvalidKeyword(token.loc, "Positional arguments cannot follow keyword arguments", token.length))
        end
        value = kw[key](stream, scene)
        push!(kwargs, key => value)
        expect_symbol(stream, (Symbol(","), Symbol(")"))).value.value == Symbol(")") && break
    end
    kwargs
end

"""
    parse_by_identifier(expected_type::IdTableKey, stream::InputStream, table::IdTable)

If the next token is an [`Identifier`](@ref), check that its type is coherent with `expected_type` and return its value.
"""
function parse_by_identifier(expected_type::IdTableKey, stream::InputStream, table::IdTable)
    next_token = read_token(stream)
    unread_token(stream, next_token)
    if !isa(next_token.value, Identifier)
        return nothing
    end
    id_name = next_token.value.value
    haskey(table[expected_type], id_name) && return table[expected_type][id_name].value
    (type = findfirst(d -> haskey(d, id_name), table)) |> isnothing ||
        throw(WrongTokenType(next_token.loc, "Variable '$id_name' is of type '$type': expected '$expected_type'\nVariable '$id_name' was declared at $(table[type][id_name].loc)", next_token.length))
    throw(UndefinedIdentifier(next_token.loc, "Undefined variable '$id_name'", next_token.length))
end


text:
"""
    (new_trace, accepted) = metropolis_hastings(trace, selection::Selection)

Perform a Metropolis-Hastings update that proposes new values for the selected addresses from the internal proposal (often using ancestral sampling), returning the new trace (which is equal to the previous trace if the move was not accepted) and a Bool indicating whether the move was accepted or not.
"""
function metropolis_hastings(trace, selection::Selection)
    args = get_args(trace)
    argdiffs = map((_) -> NoChange(), args)
    (new_trace, weight) = regenerate(trace, args, argdiffs, selection)
    if log(rand()) < weight
        # accept
        return (new_trace, true)
    else
        # reject
        return (trace, false)
    end
end

"""
    (new_trace, accepted) = metropolis_hastings(trace, proposal::GenerativeFunction, proposal_args::Tuple)

Perform a Metropolis-Hastings update that proposes new values for some subset of random choices in the given trace using the given proposal generative function, returning the new trace (which is equal to the previous trace if the move was not accepted) and a Bool indicating whether the move was accepted or not.

The proposal generative function should take as its first argument the current trace of the model, and remaining arguments `proposal_args`.
If the proposal modifies addresses that determine the control flow in the model, values must be provided by the proposal for any addresses that are newly sampled by the model.
"""
function metropolis_hastings(trace, proposal::GenerativeFunction, proposal_args::Tuple)
    model_args = get_args(trace)
    argdiffs = map((_) -> NoChange(), model_args)
    proposal_args_forward = (trace, proposal_args...,)
    (fwd_choices, fwd_weight, _) = propose(proposal, proposal_args_forward)
    (new_trace, weight, _, discard) = update(trace,
        model_args, argdiffs, fwd_choices)
    proposal_args_backward = (new_trace, proposal_args...,)
    (bwd_weight, _) = assess(proposal, proposal_args_backward, discard)
    alpha = weight - fwd_weight + bwd_weight
    if log(rand()) < alpha
        # accept
        return (new_trace, true)
    else
        # reject
        return (trace, false)
    end
end

"""
    (new_trace, accepted) = metropolis_hastings(trace, proposal::GenerativeFunction, proposal_args::Tuple, involution::Function)

Perform a generalized Metropolis-Hastings update based on an involution (bijection that is its own inverse) on a space of choice maps, returning the new trace (which is equal to the previous trace if the move was not accepted) and a Bool indicating whether the move was accepted or not.

The `involution` Julia function has the following signature:

    (new_trace, bwd_choices::ChoiceMap, weight) = involution(trace, fwd_choices::ChoiceMap, fwd_ret, proposal_args::Tuple)

The generative function `proposal` is executed on arguments `(trace, proposal_args...)`, producing a choice map `fwd_choices` and return value `fwd_ret`.
For each value of model arguments (contained in `trace`) and `proposal_args`, the `involution` function applies an involution that maps the tuple `(get_choices(trace), fwd_choices)` to the tuple `(get_choices(new_trace), bwd_choices)`.
Note that `fwd_ret` is a deterministic function of `fwd_choices` and `proposal_args`.
When only discrete random choices are used, the `weight` must be equal to `get_score(new_trace) - get_score(trace)`.

**Including Continuous Random Choices**
When continuous random choices are used, the `weight` must include an additive term that is the determinant of the the Jacobian of the bijection on the continuous random choices that is obtained by currying the involution on the discrete random choices.
"""
function metropolis_hastings(trace, proposal::GenerativeFunction,
                    proposal_args::Tuple, involution::Function; check_round_trip=false)
    (fwd_choices, fwd_score, fwd_ret) = propose(proposal, (trace, proposal_args...,))
    (new_trace, bwd_choices, weight) = involution(trace, fwd_choices, fwd_ret, proposal_args)
    (bwd_score, bwd_ret) = assess(proposal, (new_trace, proposal_args...), bwd_choices)
    if check_round_trip
        (trace_rt, fwd_choices_rt, weight_rt) = involution(new_trace, bwd_choices, bwd_ret, proposal_args)
        if !isapprox(fwd_choices_rt, fwd_choices)
            println("fwd_choices:")
            println(fwd_choices)
            println("fwd_choices_rt:")
            println(fwd_choices_rt)
            error("Involution round trip check failed")
        end
        if !isapprox(get_choices(trace), get_choices(trace_rt))
            println("get_choices(trace):")
            println(get_choices(trace))
            println("get_choices(trace_rt):")
            println(get_choices(trace_rt))
            error("Involution round trip check failed")
        end
        if !isapprox(weight, -weight_rt)
            println("weight: $weight, -weight_rt: $(-weight_rt)")
            error("Involution round trip check failed")
        end
    end
    if log(rand()) < weight - fwd_score + bwd_score
        # accept
        (new_trace, true)
    else
        # reject
        (trace, false)
    end
end

"""
    (new_trace, accepted) = mh(trace, selection::Selection)
    (new_trace, accepted) = mh(trace, proposal::GenerativeFunction, proposal_args::Tuple)
    (new_trace, accepted) = mh(trace, proposal::GenerativeFunction, proposal_args::Tuple, involution::Function)

Alias for [`metropolis_hastings`](@ref). Perform a Metropolis-Hastings update on the given trace.
"""
const mh = metropolis_hastings

export metropolis_hastings, mh

max_stars_repo_name: cameronraysmith/Gen

max_issues_repo_name: cameronraysmith/Gen

max_forks_repo_name: cameronraysmith/Gen