using Test
using AtomicLevels
# ^^^ -- to make it possible to run the test file separately

module GRASPParser
using Test
using AtomicLevels
import AtomicLevels: CSF, csfs
using WignerSymbols
include("grasp/rcsfparser.jl")

compare_with_grasp(f, rcsfout) = @testset "GRASP CSFs: $(rcsfout)" begin
    grasp_csfs = parse_rcsf(joinpath(@__DIR__, "grasp", rcsfout))
    atlev_csfs = f()
    @test length(atlev_csfs) == length(grasp_csfs)
    for atlev_csf in atlev_csfs
        @test atlev_csf in grasp_csfs
    end
end

export compare_with_grasp
end

using .GRASPParser

@testset "Excited configurations" begin
    @testset "Simple" begin
        @test_throws ArgumentError excited_configurations(rc"[Kr] 5s2 5p6", max_excitations=:triples)
        @test_throws ArgumentError excited_configurations(rc"[Kr] 5s2 5p6", min_occupancy=[2,0])
        @test_throws ArgumentError excited_configurations(rc"[Kr] 5s2 5p6", max_occupancy=[2,0])

        @test_throws ArgumentError excited_configurations(rc"[Kr] 5s2 5p6", min_occupancy=[-1,0,0])
        @test_throws ArgumentError excited_configurations(rc"[Kr] 5s2 5p6", max_occupancy=[3,2,4])
    end

    @testset "Ion–continuum" begin
        @test ion_continuum(c"1s2", os"k[s-d]") == [c"1s2", c"1s ks", c"1s kp", c"1s kd"]
    end

    @testset "GRASP comparisons" begin
        @test excited_configurations(c"1s2", os"2[s-p]"...) ==
            [c"1s2", c"1s 2s", c"2s2", c"2p2"]
        @test excited_configurations(c"1s2", os"k[s-p]"...) ==
            [c"1s2", c"1s ks", c"ks2", c"kp2"]

        @test excited_configurations(rc"1s2", ros"2[s-p]"...) ==
            [rc"1s2", rc"1s 2s", rc"2s2", rc"2p-2", rc"2p- 2p", rc"2p2"]
        @test excited_configurations(rc"1s2", ro"2s") == [rc"1s2", rc"1s 2s", rc"2s2"]
        @test excited_configurations(rc"1s2", ro"2p-") == [rc"1s2", rc"2p-2"]

        #= rcsfgenerate input:
        *
        1
        2s(2,i)2p(2,i)

        2s,2p
        0 10
        0
        n
        =#
        compare_with_grasp("rcsf.out.1.txt") do
            csfs(rc"1s2 2s2" ⊗ rcs"2p2")
        end

        #= rcsfgenerate input:
        *
        0
        3s(1,i)3p(3,i)3d(4,i)

        3s,3p,3d
        0 50
        0
        n
        =#
        compare_with_grasp("rcsf.out.2.txt") do
            csfs(rc"3s1" ⊗ rcs"3p3" ⊗ rcs"3d4")
        end

        #= rcsfgenerate input:
        *
        2
        3s(1,i)3p(2,i)3d(2,i)

        3s,3p,3d
        1 51
        0
        n
        =#
        compare_with_grasp("rcsf.out.3.txt") do
            csfs(rc"[Ne]* 3s1" ⊗ rcs"3p2" ⊗ rcs"3d2")
        end

        #= rcsfgenerate input:
        *
        0
        1s(2,*)

        2s,2p
        0 20
        2
        n
        =#
        compare_with_grasp("rcsf.out.4.txt") do
            csfs(excited_configurations(rc"1s2", ro"2s", ro"2p-", ro"2p"))
        end

        #= rcsfgenerate input:
        *
        0
        1s(2,*)

        3s,3p,3d
        0 20
        2
        n
        =#
        compare_with_grasp("rcsf.out.5.txt") do
            excited_configurations(
                rc"1s2",
                ro"2s", ro"2p-", ro"2p", ro"3s", ro"3p-", ro"3p", ro"3d-", ro"3d"
            ) |> csfs
        end

        # TODO:

        #= rcsfgenerate input: (problematic alphas with semicolons etc)
        *
        0
        5f(5,i)

        5s,5p,5d,5f
        1 51
        0
        n
        =#
    end
end


function Base.sort!(df::DataFrame; cols=Any[], alg=nothing,
                    lt=isless, by=identity, rev=false, order=Forward)
    if !(isa(by, Function) || eltype(by) <: Function)
        msg = "'by' must be a Function or a vector of Functions. Perhaps you wanted 'cols'."
        throw(ArgumentError(msg))
    end
    ord = ordering(df, cols, lt, by, rev, order)
    _alg = Sort.defalg(df, ord; alg=alg, cols=cols)
    sort!(df, _alg, ord)
end

function Base.sort!(df::DataFrame, a::Base.Sort.Algorithm, o::Base.Sort.Ordering)
    p = sortperm(df, a, o)
    pp = similar(p)
    c = columns(df)

    for (i,col) in enumerate(c)
        # Check if this column has been sorted already
        if any(j -> c[j]===col, 1:i-1)
            continue
        end

        copy!(pp,p)
        Base.permute!!(col, pp)
    end
    df
end


################################################################################
#
#   NfOrd/Ideal/Prime.jl : Prime ideals in orders of absolute number fields
#
# This file is part of Hecke.
#
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
#
#  Copyright (C) 2015, 2016, 2017 Tommy Hofmann
#  Copyright (C) 2015, 2016, 2017 Claus Fieker
#
################################################################################

export PrimeIdealsSet

doc"""
***
    isramified(O::NfOrd, p::Int) -> Bool

> Returns whether the integer $p$ is ramified in $\mathcal O$.
> It is assumed that $p$ is prime.
"""
function isramified(O::NfOrd, p::Union{Int, fmpz})
  @assert ismaximal_known(O) && ismaximal(O)

  return mod(discriminant(O), p) == 0
end

doc"""
***
    degree(P::NfOrdIdl) -> Int
> The inertia degree of the prime-ideal $P$.
"""
function degree(A::NfOrdIdl)
  @assert isprime(A)
  return A.splitting_type[2]
end

doc"""
***
    ramification_index(P::NfOrdIdl) -> Int
> The ramification index of the prime-ideal $P$.
"""
function ramification_index(A::NfOrdIdl)
  @assert isprime(A)
  return A.splitting_type[1]
end

doc"""
***
    splitting_type(P::NfOrdIdl) -> Int, Int
> The ramification index and inertia degree of the prime ideal $P$.
> First value is the ramificatio index, the second the degree of $P$.
"""
function splitting_type(A::NfOrdIdl)
  @assert isprime(A)
  return A.splitting_type
end

################################################################################
#
#  Prime decomposition
#
################################################################################
doc"""
    intersect_nonindex(f::Map, P::NfOrdIdl) -> NfOrdIdl
> Given a prime ideal $P$ in $K$ and the inclusion map $f:k \to K$
> of number fields, find the unique prime $p$ in $k$ below.
"""
function intersect_nonindex(f::Map, P::NfOrdIdl)
  @assert P.is_prime == 1
  #let g be minpoly of k, G minpoly of K and h in Qt the primitive
  #element of k in K (image of gen(k))
  #then
  #  g(h) = 0 mod G
  k = domain(f)
  K = codomain(f)
  G = K.pol
  Qx = parent(G)
  g = k.pol(gen(Qx))
  h = Qx(f(gen(k)))

  Fp, xp = PolynomialRing(ResidueRing(FlintZZ, Int(minimum(P)), cached=false), cached=false)
  gp = factor(Fp(g))
  hp = Fp(h)
  Gp = gcd(Fp(K(P.gen_two)), Fp(G))
  for (f, e) = gp.fac
    if f(hp) % Gp == 0
      Zk = maximal_order(k)
      p = ideal_from_poly(Zk, Int(minimum(P)), f, 1)
      return p
    end
  end
end

doc"""
    prime_decomposition_nonindex(f::Map, p::NfOrdIdl) -> Array{Tuple{NfOrdIdl, Int}, 1}
> Given a map $f: k\to K$ of number fields defined over $\mathbb Q$ and
> a prime ideal in the maximal order of $k$, find all prime ideals in
> the maximal order of $K$ above.
"""
function prime_decomposition_nonindex(f::Map, p::NfOrdIdl)
  @assert p.is_prime == 1
  k = domain(f)
  K = codomain(f)
  ZK = maximal_order(K)
  G = K.pol
  Qx = parent(G)

  Fp, xp = PolynomialRing(ResidueRing(FlintZZ, Int(minimum(p)), cached=false), cached=false)
  Gp = factor(gcd(Fp(f(K(p.gen_two))), Fp(G)))
  res = []
  Zk = maximal_order(k)
  for (f, e) = Gp.fac
    P = ideal_from_poly(ZK, Int(minimum(p)), f, 1)
    push!(res, (P, e))
  end
  return res
end

doc"""
    lift(K::AnticNumberField, f::nmod_poly) -> nf_elem
> Given a polynomial $f$ over a finite field, lift it to an element of the
> number field $K$. The lift if given by the eleemnt represented by the
> canonical lift of $f$ to a polynomial over the integers.
"""
function lift(K::AnticNumberField, f::nmod_poly)
  if degree(f)>=degree(K)
    f = mod(f, parent(f)(K.pol))
  end
  r = K()
  for i=0:f.length-1
    u = ccall((:nmod_poly_get_coeff_ui, :libflint), UInt, (Ptr{nmod_poly}, Int), &f, i)
    _num_setcoeff!(r, i, u)
  end
  return r
end

##TODO: make fmpz-safe!!!!
#return <p, lift(O, fi> in 2-element normal presentation given the data
function ideal_from_poly(O::NfOrd, p::Int, fi::nmod_poly, ei::Int)
  b = lift(nf(O), fi)
  idl = ideal(O, fmpz(p), O(b, false))
  idl.is_prime = 1
  idl.splitting_type = ei, degree(fi)
  idl.norm = FlintZZ(p)^degree(fi)
  idl.minimum = FlintZZ(p)

  # We have to do something to get 2-normal presentation:
  # if ramified or valuation val(b,P) == 1, (p,b)
  # is a P(p)-normal presentation
  # otherwise we need to take p+b
  # I SHOULD CHECK THAT THIS WORKS

  if !((mod(norm(b),(idl.norm)^2) != 0) || (ei > 1))
    idl.gen_two = idl.gen_two + O(p)
  end

  idl.gens_normal = p
  idl.gens_weakly_normal = true

  # Find an "anti-uniformizer" in case P is unramified
  # We don't call it anti-unfiformizer anymore

  #if ideal.splitting_type[1] == 1
  #  t = parent(f)(lift(Zx, divexact(fmodp, fi)))
  #  ideal.anti_uniformizer = O(K(t), false)
  #end

  if idl.splitting_type[2] == degree(O)
    # Prime number is inert, in particular principal
    idl.is_principal = 1
    idl.princ_gen = O(p)
  end
  return idl
end

doc"""
***
    prime_decomposition(O::NfOrd,
                        p::Integer,
                        degree_limit::Int = 0,
                        lower_limit::Int = 0) -> Array{Tuple{NfOrdIdl, Int}, 1}

> Returns an array of tuples $(\mathfrak p_i,e_i)$ such that $p \mathcal O$ is the product of
> the $\mathfrak p_i^{e_i}$ and $\mathfrak p_i \neq \mathfrak p_j$ for $i \neq j$.
>
> If `degree_limit` is a nonzero integer $k > 0$, then only those prime ideals
> $\mathfrak p$ with $\deg(\mathfrak p) \leq k$ will be returned.
> Similarly if `\lower_limit` is a nonzero integer $l > 0$, then only those prime ideals
> $\mathfrak p$ with $l \leq \deg(\mathfrak p)$ will be returned.
> Note that in this case it may happen that $p\mathcal O$ is not the product of the
> $\mathfrak p_i^{e_i}$.
"""
function prime_decomposition(O::NfAbsOrd{S, T}, p::Union{Integer, fmpz}, degree_limit::Int = 0, lower_limit::Int = 0, cached::Bool = true) where {S, T}
  if typeof(p) == fmpz && nbits(p) < 64
    return prime_decomposition(O, Int(p), degree_limit, lower_limit)
  end

  if mod(index(O),fmpz(p)) == 0 || !issimple(nf(O))
    if cached
      if haskey(O.index_div, fmpz(p))
        lp = O.index_div[fmpz(p)]
        z = Tuple{NfAbsOrdIdl{S, T}, Int}[]
        for (Q, e) in lp
          if degree(Q) <= degree_limit
            push!(z, (Q, e))
          end
        end
        return z
      end
    end
    return prime_dec_index(O, p, degree_limit, lower_limit)
  end
  return prime_dec_nonindex(O, p, degree_limit, lower_limit)
end

function _fac_and_lift(f::fmpz_poly, p, degree_limit, lower_limit)
  Zx = parent(f)
  Zmodpx = PolynomialRing(ResidueRing(FlintZZ, p, cached = false), "y", cached = false)[1]
  fmodp = Zmodpx(f)
  fac = factor(fmodp)
  lifted_fac = Array{Tuple{fmpz_poly, Int}, 1}()
  for (k, v) in fac
    if degree(k) <= degree_limit && degree(k) >= lower_limit
      push!(lifted_fac, (lift(Zx, k), v))
    end
  end
  return lifted_fac
end

function prime_dec_nonindex(O::NfOrd, p::Union{Integer, fmpz}, degree_limit::Int = 0, lower_limit::Int = 0)
  K = nf(O)
  f = K.pol
  R = parent(f)
  Zx, x = PolynomialRing(FlintIntegerRing(),"x")
  Zf = Zx(f)

  if degree_limit == 0
    degree_limit = degree(K)
  end

  fac = _fac_and_lift(Zf, p, degree_limit, lower_limit)

  result = Array{Tuple{ideal_type(O),Int}}(length(fac))

  for k in 1:length(fac)
    fi = fac[k][1]
    ei = fac[k][2]
    #ideal = ideal_from_poly(O, p, fi, ei)
    t = parent(f)(fi)
    b = K(t)
    ideal = NfAbsOrdIdl(O)
    ideal.gen_one = p
    ideal.gen_two = O(b, false)
    ideal.is_prime = 1
    ideal.splitting_type = ei, degree(fi)
    ideal.norm = FlintZZ(p)^degree(fi)
    ideal.minimum = FlintZZ(p)

    # We have to do something to get 2-normal presentation:
    # if ramified or valuation val(b,P) == 1, (p,b)
    # is a P(p)-normal presentation
    # otherwise we need to take p+b
    # I SHOULD CHECK THAT THIS WORKS

    if !((mod(norm(b),(ideal.norm)^2) != 0) || (ei > 1))
      ideal.gen_two = ideal.gen_two + O(p)
    end

    ideal.gens_normal = p
    ideal.gens_weakly_normal = true

    # Find an "anti-uniformizer" in case P is unramified
    # We don't call it anti-unfiformizer anymore

    #if ideal.splitting_type[1] == 1
    #  t = parent(f)(lift(Zx, divexact(fmodp, fi)))
    #  ideal.anti_uniformizer = O(K(t), false)
    #end

    if length(fac) == 1 && ideal.splitting_type[2] == degree(f)
      # Prime number is inert, in particular principal
      ideal.is_principal = 1
      ideal.princ_gen = O(p)
    end
    result[k] =  (ideal, ei)
    k += 1
  end
  return result
end

function prime_dec_index(O::NfOrd, p::Union{Integer, fmpz}, degree_limit::Int = 0, lower_limit::Int = 0)
  if degree_limit == 0
    degree_limit = degree(O)
  end

  Ip = pradical(O, p)
  A, OtoA = AlgAss(O, Ip, p)
  AtoO = inv(OtoA)
  AA = split(A)

  basisO = basis(O)

  ideals = Vector{NfOrdIdl}()
  m = zero_matrix(FlintZZ, 1, degree(O))
  for (B, BtoA) in AA
    f = dim(B)
    idem = BtoA(B[1]) # Assumes that B == idem*A
    M = representation_matrix(idem)
    ker = left_kernel(M)
    N = basis_mat(Ip)
    for i = 1:length(ker)
      b = elem_in_basis(AtoO(A(ker[i])))
      for j = 1:degree(O)
        m[1, j] = b[j]
      end
      N = vcat(N, m)
    end
    N = sub(_hnf_modular_eldiv(N, fmpz(p), :lowerleft), rows(N) - degree(O) + 1:rows(N), 1:degree(O))
    P = ideal(O, N)
    P.norm = fmpz(p)^f
    P.splitting_type = (0, f)
    #
    fromOtosimplealgebra = Hecke._compose(inv(BtoA), OtoA)
    compute_residue_field_data!(P, fromOtosimplealgebra)
    #primB, minprimB, getcoordpowerbasis = _as_field(B)
    #@assert degree(minprimB) == f
    #P.prim_elem = AtoO(BtoA(primB))
    #P.min_poly_prim_elem = fmpz_poly(fmpz[FlintZZ(coeff(minprimB, i)) for i in 0:degree(minprimB)])
    #P.min_poly_prim_elem.parent = FmpzPolyRing(:$, false)
    #P.basis_in_prim = Vector{fmpz_mat}(degree(O))
    #for i in 1:degree(O)
    #  P.basis_in_prim[i] = zero_matrix(FlintZZ, 1, f)
    #  t = getcoordpowerbasis(BtoA\(OtoA(basisO[i])))
    #  for j in 1:f
    #    P.basis_in_prim[i][1, j] = FlintZZ(t[1, j])
    #  end
    #end
    push!(ideals, P)
  end

  result = Vector{Tuple{NfOrdIdl, Int}}()

  for j in 1:length(ideals)
    P = ideals[j]
    f = P.splitting_type[2]

    if f > degree_limit || f < lower_limit
      continue
    end

    # The following does not work if there is only one prime ideal
    if length(ideals) > 1 && (1-1/BigInt(p))^degree(O) < 0.1
      # This is roughly Algorithm 6.4 of Belabas' "Topics in comptutational algebraic
      # number theory".

      # Compute Vp = P_1 * ... * P_j-1 * P_j+1 * ... P_g

      B, BtoA = AA[j]
      J = ideal(O, AtoO(BtoA(B[1])))
      N = sub(_hnf_modular_eldiv(vcat(basis_mat(Ip), basis_mat(J)), fmpz(p), :lowerleft), degree(O) + 1:2*degree(O), 1:degree(O))
      Vp = ideal(O, N)

      u, v = idempotents(P, Vp)

      x = zero(parent(u))

      if !iszero(mod(norm(u), norm(P)*p))
        x = u
      elseif !iszero(mod(norm(u + p), norm(P)*p))
        x = u + p
      elseif !iszero(mod(norm(u - p), norm(P)*p))
        x = u - p
      else
        for i in 1:degree(O)
          if !iszero(mod(norm(v*basis(P)[i] + u), norm(P)*p))
            x = v*basis(P)[i] + u
          end
        end
      end

      @hassert :NfOrd 1 !iszero(x)
      @hassert :NfOrd 2 O*O(p) + O*x == P

      P.gen_one = p
      P.gen_two = x
      P.gens_normal = p
      P.gens_weakly_normal = 1
    else
      @vprint :NfOrd 1 "Chances for finding second generator: ~$((1-1/p))\n"
      _assure_weakly_normal_presentation(P)
      assure_2_normal(P)
    end

    e = Int(valuation(nf(O)(p), P))
    P.splitting_type = e, f
    P.is_prime = 1
    push!(result, (P, e))
  end
  #=
  if degree_limit >= degree(O)
    O.index_div[fmpz(p)] = result
  end
  =#
  return result
end

function uniformizer(P::NfOrdIdl)
  p = minimum(P)
  if P.gens_normal == p
    return P.gen_two
  else
    if p > 250
      r = 500  # should still have enough elements...
    else
      r = Int(div(p, 2))
    end
    z = rand(P, r)
    while true
      if !iszero(z) && valuation(z, P) == 1
        break
      end
      z = rand(P, r)
    end
    return z
  end
end

# Belabas p. 40
function anti_uniformizer(P::NfOrdIdl)
  if isdefined(P, :anti_uniformizer)
    return P.anti_uniformizer
  else
    p = minimum(P)
    M = representation_matrix(uniformizer(P))
    Mp = MatrixSpace(ResidueRing(FlintZZ, p, cached=false), rows(M), cols(M), false)(M)
    K = kernel(Mp)
    @assert length(K) > 0
    P.anti_uniformizer = elem_in_nf(order(P)(_lift(K[1])))//p
    return P.anti_uniformizer
  end
end

# Don't use the following function. It does not work for index divisors
# TH: Or does it?
function prime_decomposition_type(O::NfOrd, p::Integer)
  if (mod(discriminant(O), p)) != 0 && (mod(fmpz(index(O)), p) != 0)
    K = nf(O)
    f = K.pol
    R = parent(f)
    Zx, x = PolynomialRing(FlintZZ,"x", cached=false)
    Zf = Zx(f)
    fmodp = PolynomialRing(ResidueRing(FlintZZ,p, cached = false), "y", cached = false)[1](Zf)
    fac = factor_shape(fmodp)
    g = sum([ x for x in values(fac)])
    res = Array{Tuple{Int, Int}}(g)
    k = 1
    for (fi, ei) in fac
      for j in 1:ei
        res[k] = (fi, 1)
        k = k + 1
      end
    end
  else
    lp = prime_decomposition(O, p)
    res = Array{Tuple{Int, Int}}(length(lp))
    for i in 1:length(lp)
      res[i] = (lp[i][1].splitting_type[2], lp[i][1].splitting_type[1])
    end
  end
  return res
end

doc"""
***
    prime_ideals_up_to(O::NfOrd,
                       B::Int;
                       degree_limit::Int = 0) -> Array{NfOrdIdl, 1}

> Computes the prime ideals $\mathcal O$ with norm up to $B$.
>
> If `degree_limit` is a nonzero integer $k$, then prime ideals $\mathfrak p$
> with $\deg(\mathfrak p) > k$ will be discarded.
"""
function prime_ideals_up_to(O::NfOrd, B::Int;
                            complete::Bool = false,
                            degree_limit::Int = 0)
  p = 1
  r = NfOrdIdl[]
  while p < B
    p = next_prime(p)
    if p > B
      return r
    end
    if !complete
      deg_lim = Int(floor(log(B)/log(p)))
      if degree_limit >0
        deg_lim = min(deg_lim, degree_limit)
      end
    else
      deg_lim = 0
    end
    @vprint :ClassGroup 2 "decomposing $p ... (bound is $B, deg_lim $deg_lim)\n"
    li = prime_decomposition(O, p, deg_lim)
    for P in li
      push!(r, P[1])
    end
  end
  return r
end

doc"""
***
    prime_ideals_over(O::NfOrd,
                       lp::AbstractArray{Int, 1};
                       degree_limit::Int = 0) -> Array{NfOrdIdl, 1}

> Computes the prime ideals $\mathcal O$ over prime numbers in $lp$.
>
> If `degree_limit` is a nonzero integer $k$, then prime ideals $\mathfrak p$
> with $\deg(\mathfrak p) > k$ will be discarded.
"""
function prime_ideals_over(O::NfOrd,
                           lp::AbstractArray{T};
                           degree_limit::Int = 0) where T <: Union{fmpz, Integer}
  p = 1
  r = NfOrdIdl[]
  for p in lp
    @vprint :ClassGroup 2 "decomposing $p ... (deg_lim $deg_lim)"
    li = prime_decomposition(O, p, degree_limit)
    for P in li
      push!(r, P[1])
    end
  end
  return r
end

doc"""
***
    prime_ideals_up_to(O::NfOrd,
                       B::Int;
                       complete::Bool = false,
                       degree_limit::Int = 0,
                       F::Function,
                       bad::fmpz)

> Computes the prime ideals $\mathcal O$ with norm up to $B$.
>
> If `degree_limit` is a nonzero integer $k$, then prime ideals $\mathfrak p$
> with $\deg(\mathfrak p) > k$ will be discarded.
>
> The function $F$ must be a function on prime numbers not dividing `bad` such that
> $F(p) = \deg(\mathfrak p)$ for all prime ideals $\mathfrak p$ lying above $p$.
"""
function prime_ideals_up_to(O::NfOrd, B::Int, F::Function, bad::fmpz = discriminant(O);
                            complete::Bool = false,
                            degree_limit::Int = 0)
  p = 1
  r = NfOrdIdl[]
  while p < B
    p = next_prime(p)
    if p > B
      return r
    end
    if !complete
      deg_lim = flog(fmpz(B), p) # Int(floor(log(B)/log(p)))
      if degree_limit > 0
        deg_lim = min(deg_lim, degree_limit)
      end
    else
      deg_lim = 0
    end
    @vprint :ClassGroup 2 "decomposing $p ... (bound is $B)"
    if mod(bad, p) == 0
      li = prime_decomposition(O, p, deg_lim)
      for P in li
        push!(r, P[1])
      end
    else
      if F(p) <= deg_lim
        li = prime_decomposition(O, p, deg_lim)
        for P in li
          push!(r, P[1])
        end
      end
    end
  end
  return r
end

################################################################################
#
#  Coprime
#
################################################################################

#TODO: do sth. useful here!!!
function divides(A::NfOrdIdl, B::NfOrdIdl)
  minimum(A) % minimum(B) == 0 || return false
  return valuation(A, B) > 0
end

function coprime_base(A::Array{NfOrdIdl, 1}, p::fmpz)
  #consider A^2 B and A B: if we do gcd with the minimum, we get twice AB
  #so the coprime base is AB
  #however using the p-part of the norm, the coprime basis becomes A, B...
  if iseven(p)
    lp = prime_decomposition(order(A[1]), 2)
    Ap = [x[1] for x = lp if any(y->divides(y, x[1]) > 0, A)]
    a = remove(p, 2)[2]
    if !isone(a)
      Bp = coprime_base(A, a)
      return vcat(Ap, Bp)
    else
      return Ap
    end
  else
    Ap = [gcd(x, p^valuation(norm(x), p)) for x = A if minimum(x) % p == 0]
  end
  return coprime_base_steel(Ap)
end

doc"""
***
    coprime_base(A::Array{NfOrdIdl, 1}) -> Array{NfOrdIdl, 1}
    coprime_base(A::Array{NfOrdElem, 1}) -> Array{NfOrdIdl, 1}
> A coprime base for the (principal) ideals in $A$, ie. the returned array
> generated multiplicatively the same ideals as the input and are pairwise
> coprime.
"""
function coprime_base(A::Array{NfOrdIdl, 1})
  a = collect(Set(map(minimum, A)))
  a = coprime_base(a)
  C = Array{NfOrdIdl, 1}()

  for p = a
    if p == 1
      continue
    end
    cp = coprime_base(A, p)
    append!(C, cp)
  end
  return C
end

function coprime_base(A::Array{NfOrdElem, 1})
  O = parent(A[1])
  return coprime_base([ideal(O, x) for x = A])
end

function integral_split(A::NfOrdIdl)
  return A, ideal(Order(A), fmpz(1))
end

################################################################################
#
#  Factorization into prime ideals
#
################################################################################

#TODO: factoring type??? (with unit)
doc"""
***
    factor(A::NfOrdIdl) -> Dict{NfOrdIdl, Int}

> Computes the prime ideal factorization $A$ as a dictionary, the keys being
> the prime ideal divisors:
> If `lp = factor_dict(A)`, then `keys(lp)` are the prime ideal divisors of A
> and `lp[P]` is the `P`-adic valuation of `A` for all `P` in `keys(lp)`.
"""
factor(A::NfOrdIdl) = factor_dict(A)

function factor_dict(A::NfOrdIdl)
  ## this should be fixed
  lf = factor(minimum(A))
  lF = Dict{NfOrdIdl, Int}()
  n = norm(A)
  O = order(A)
  for (i, (p, v)) in enumerate(lf)
    lP = prime_decomposition(O, p)
    for P in lP
      v = valuation(A, P[1])
      if v != 0
        lF[P[1]] = v
        n = n//norm(P[1])^v
      end
      if n == 1
        return lF
      end
    end
  end
  return lF
end

################################################################################
#
#  Functions for index divisor splitting
#
################################################################################

mutable struct quoringalg{T} <: Ring
  base_order::NfOrd
  ideal::NfOrdIdl
  prime::T
  basis::Array{NfOrdElem, 1}

  function quoringalg(O::NfOrd, I::NfOrdIdl, p::T) where {T}

    z = new{T}()
    z.base_order = O
    z.ideal = I
    z.prime = p

    # compute a basis
    Rp = ResidueRing(FlintZZ, p, cached=false)
    Amodp = MatrixSpace(Rp, degree(O), degree(O), false)(basis_mat(I))
    Amodp = vcat(Amodp, zero_matrix(Rp, 1, degree(O)))
    Amodp[1,1] = 1
    Amodp = sub(Amodp, 1:degree(O), 1:degree(O))

    r, B = _rref(Amodp)
    C = zero_matrix(Rp, degree(O)-r, degree(O))
    BB = Array{NfOrdElem}(degree(O) - r)
    pivots = Array{Int}(0)
#    # get he pivots of B
    for i in 1:r
      for j in 1:degree(O)
        if !iszero(B[i,j])
          push!(pivots, j)
          break
        end
      end
    end
    i = 1
    k = 1
    while i <= degree(O)-r
      for j in k:degree(O)
        if !in(j, pivots)
          BB[i] = basis(O)[j]
          C[i,j] = 1
          k = j + 1
          i = i + 1
          break
        end
      end
    end
    insert!(BB, 1, basis(O)[1])
    z.basis = BB
    return z
  end
end

mutable struct quoelem
  parent::quoringalg
  elem::NfOrdElem
  coord::Array{fmpz, 1}

  function quoelem(R::quoringalg, x::NfOrdElem)
    z = new()
    z.parent = R
    z.elem = x

    return z
  end
end

function _kernel_of_frobenius(R::quoringalg)
  O = R.base_order
  BB = R.basis
  p = R.prime
  Rp = ResidueRing(FlintZZ, R.prime, cached=false)
  C = zero_matrix(Rp, length(BB)+1, degree(O))
  D = zero_matrix(Rp, length(BB), degree(O))

  Q, mQ = quo(O, R.ideal)

  function g(x)
    return mQ\(mQ(x)^p)
  end

  for i in 1:length(BB)
    A = elem_in_basis(mod(g(BB[i]) - BB[i], R.ideal))
    for j in 1:degree(O)
      D[i,j] = A[j]
    end
  end

  DD = NfOrdElem[ dot(BB, _lift(r)) for r in kernel(D) ]

  return [ quoelem(R, r) for r in DD ]
end

function _lift(T::Array{Generic.Res{fmpz}, 1})
  return [ z.data for z in T ]
end

function _lift(T::Array{Nemo.nmod, 1})
  return [ fmpz(z.data) for z in T ]
end

function *(x::quoelem, y::quoelem)
  z = mod(x.elem * y.elem, x.parent.ideal)
  return quoelem(x.parent, z)
end

function ^(x::quoelem, y::Int)
  z = mod(x.elem^y, x.parent.ideal)
  return quoelem(x.parent, z)
end

function ^(x::quoelem, y::Union{Integer, fmpz})
  # Do something stupid
  R, m = quo(x.parent.base_order, x.parent.ideal)
  return quoelem(x.parent, (m\(m(x.elem)^y)))
end

function ==(x::quoelem, y::quoelem)
  z = mod(x.elem - y.elem, x.parent.ideal)
  return zero(parent(z)) == z
end

function minpoly(x::quoelem)
  O = x.parent.base_order
  p = x.parent.prime

  Rp = ResidueRing(FlintZZ, p, cached=false)
  A = zero_matrix(Rp, 0, degree(O))
  B = zero_matrix(Rp, 1, degree(O))

  for i in 0:degree(O)
    ar = elem_in_basis((x^i).elem)
    for j in 1:degree(O)
      B[1, j] = ar[j]
    end
    A = vcat(A, B)
    K = kernel(A)
    if length(K) > 0
      @assert length(K) == 1
      f = PolynomialRing(Rp, "x", cached=false)[1](K[1])
      return f
    end
  end
  error("cannot find minpoly")
end

function split(R::quoringalg)
  if length(R.basis) == 1
    return [ R ]
  end
  K = _kernel_of_frobenius(R)
  O = R.base_order
  p = R.prime

  k = length(K)

  if k == 1
    # the algebra is a field over F_p
    # the ideal Ip is a prime ideal!
    return [ R ]
  end

  maxit = 1

  while true
    maxit = maxit + 1
    r = rand(0:p-1, length(K))

    x = quoelem(R, dot([ x.elem for x in K], r))

    if mod((x^p).elem, R.ideal) != mod(x.elem, R.ideal)
      #println("element^p: $(mod((x^p).elem, R.ideal))")
      #println("element: $(mod(x.elem, R.ideal))")
      #println(R.ideal.basis_mat)
      #println(K)
      error("Strange")
    end

    f = minpoly(x)

    if degree(f) < 2
      continue
    end
    @assert  issquarefree(f)

#    # By theory, all factors should have degree 1 # exploit this if p is small!
    fac = factor(f)
    F = first(keys(fac.fac))
    @assert length(fac) == degree(f)
    H = divexact(f,F)
    E, U, V = gcdx(F, H)
    @assert E == 1
    H = U*F;
    idem = O(coeff(H,0).data)
    for i in 1:degree(H)
      idem = idem + coeff(H,i).data*x.elem^i
    end

    I1 = R.ideal + ideal(O, idem)
    I2 = R.ideal + ideal(O, O(1)-idem)

    return vcat(split(quoringalg(O, I1, p)), split(quoringalg(O, I2, p)))
    break
  end
end

################################################################################
#
#  Primality testing
#
################################################################################

doc"""
***
    isprime_known(A::NfOrdIdl) -> Bool

> Returns whether $A$ knows if it is prime.
"""
function isprime_known(A::NfOrdIdl)
  return A.is_prime != 0
end

doc"""
***
    isprime(A::NfOrdIdl) -> Bool

> Returns whether $A$ is a prime ideal.
"""
function isprime(A::NfOrdIdl)
  if isprime_known(A)
    return A.is_prime == 1
  elseif minimum(A) == 0
    A.is_prime = 2
    return false
  end

  (n, p) = ispower(norm(A, Val{false}))

  if !isprime(p)
    A.is_prime = 2
    return false
  end

  p > 2^62 && error("Not implemented (yet)")

  lp = prime_decomposition(order(A), Int(p))

  for (P, f) in lp
    e = valuation(A, P)
    if e == 1 && n == degree(P)
      A.is_prime = 1
      return true
    elseif e == 0
      continue
    else
      A.is_prime = 2
      return false
    end
  end

  error("Something wrong in isprime")
end

################################################################################
#
#  Valuation
#
################################################################################

# CF:
# Classical algorithm of Cohen, but take a valuation element with smaller (?)
# coefficients. Core idea is that the valuation elementt is, originally, den*gen_two(p)^-1
# where gen_two(p) is "small". Actually, we don't care about gen_two, we
# need gen_two^-1 to be small, hence this version.
function val_func_no_index(p::NfOrdIdl)
  P = p.gen_one
  K = nf(order(p))
  pi = inv(p)
  d = denominator(K(pi.num.gen_two))
  @assert gcd(d, P) == 1
  e = K(pi.num.gen_two)*d
  M = zero_matrix(FlintZZ, 1, degree(K))
  elem_to_mat_row!(M, 1, d, e)
  @assert d == 1
  P2 = P^2
  P22 = div(P2, 2)
  for i=1:degree(K)
    x = M[1,i] % P2
    if x>P22
      x -= P2
    end
    M[1,i] = x
  end
  e = elem_from_mat_row(K, M, 1, P)
  # e is still a valuation element, but with smaller coefficients.
  return function(x::nf_elem)
    v = 0
    d = denominator(x)
    x *= d
    x = x*e
    while denominator(x) % P != 0
      v += 1
      mul!(x, x, e)
    end
    return v-valuation(d, P)*p.splitting_type[1]
  end
end

# CF:
# The idea is that valuations are mostly small, eg. in the class group
# algorithm. So this version computes the completion and the embedding into it
# at small precision and can thus compute (small) valuation at the effective
# cost of an mod(nmod_poly, nmod_poly) operation.
# Isn't it nice?
function valuation(a::UInt, b::UInt)
  return ccall((:n_remove, :libflint), Int, (Ref{UInt}, UInt), a, b)
end

#=
function valuation(a::UInt, b::UInt, bi::Cdouble)
  return ccall((:n_remove2_precomp, :libflint), Int, (Ref{UInt}, UInt, Cdouble), a, b, bi)
end
=#

function val_func_no_index_small(p::NfOrdIdl)
  P = p.gen_one
  @assert P <= typemax(UInt)
  K = nf(order(p))
  Rx = PolynomialRing(ResidueRing(FlintZZ, UInt(P), cached=false), cached=false)[1]
  Zx = PolynomialRing(FlintZZ)[1]
  g = Rx(p.gen_two.elem_in_nf)
  f = Rx(K.pol)
  g = gcd!(g, g, f)
  g = lift(Zx, g)
  k = flog(fmpz(typemax(UInt)), P)
  g = hensel_lift(Zx(K.pol), g, P, k)
  Sx = PolynomialRing(ResidueRing(FlintZZ, UInt(P)^k, cached=false), cached=false)[1]
  g = Sx(g)
  h = Sx()
  uP = UInt(P)
  return function(x::nf_elem)
    d = denominator(x)
    nf_elem_to_nmod_poly!(h, x, false) # ignores the denominator
    h = rem!(h, h, g)
    c = Nemo.coeff_raw(h, 0)
    v = c==0 ? typemax(Int) : valuation(c, uP)
    for i=1:degree(h)
      c = Nemo.coeff_raw(h, i)
      v = min(v, c==0 ? typemax(Int) : valuation(c, uP))
    end
    return v-valuation(d, P)
  end
end

function val_func_index(p::NfOrdIdl)
  # We are in the index divisor case. In larger examples, a lot of
  # time is spent computing denominators of order elements.
  # By using the representation matrix to multiply, we can stay in the order
  # and still be fast (faster even than in field).

  pi = inv(p)
  M = representation_matrix(pi.num.gen_two)
  O = order(p)
  P = p.gen_one
  return function(x::nf_elem)
    v = 0
    d, x_mat = integral_split(x, O)
    Nemo.mul!(x_mat, x_mat, M)
    while gcd(content(x_mat), P) == P  # should divide and test in place
      divexact!(x_mat, x_mat, P)
      Nemo.mul!(x_mat, x_mat, M)
      v += 1
    end
    return v-valuation(d, P)*p.splitting_type[1]
  end
end

doc"""
***
    valuation(a::nf_elem, p::NfOrdIdl) -> fmpz
    valuation(a::NfOrdElem, p::NfOrdIdl) -> fmpz
    valuation(a::fmpz, p::NfOrdIdl) -> fmpz

> Computes the $\mathfrak p$-adic valuation of $a$, that is, the largest $i$
> such that $a$ is contained in $\mathfrak p^i$.
"""
function valuation(a::nf_elem, p::NfOrdIdl)
  @hassert :NfOrd 0 !iszero(a)
  #assert(a !=0) # can't handle infinity yet
  if isdefined(p, :valuation)
    return p.valuation(a)::Int
  end
  O = order(p)
  P = p.gen_one

  # for generic ideals
  if p.splitting_type[2] == 0
    #global bad_ideal = p
    p.valuation = function(a::nf_elem)
      d = denominator(a, O)
      x = O(d*a)
      return valuation_naive(O(x), p)::Int - valuation_naive(O(d), p)::Int
    end
    return p.valuation(a)::Int
  end

  if p.splitting_type[1]*p.splitting_type[2] == degree(O)
    p.valuation = function(a::nf_elem)
      return divexact(valuation(norm(a), P)[1], p.splitting_type[2])::Int
    end
  elseif mod(index(O),P) != 0 && p.splitting_type[1] == 1
    if p.gen_one^2 <= typemax(UInt)
      f1 = val_func_no_index_small(p)
      f2 = val_func_no_index(p)
      p.valuation = function(x::nf_elem)
        v = f1(x)
        if v > 100  # can happen ONLY if the precision in the .._small function
                    # was too small.
          return f2(x)::Int
        else
          return v::Int
        end
      end
    else
      p.valuation = val_func_no_index(p)
    end
  else
    p.valuation = val_func_index(p)
  end

  return p.valuation(a)::Int
end

doc"""
***
    valuation(a::nf_elem, p::NfOrdIdl) -> fmpz
    valuation(a::NfOrdElem, p::NfOrdIdl) -> fmpz
    valuation(a::fmpz, p::NfOrdIdl) -> fmpz

> Computes the $\mathfrak p$-adic valuation of $a$, that is, the largest $i$
> such that $a$ is contained in $\mathfrak p^i$.
"""
valuation(a::NfOrdElem, p::NfOrdIdl) = valuation(a.elem_in_nf, p)

doc"""
***
    valuation(a::nf_elem, p::NfOrdIdl) -> fmpz
    valuation(a::NfOrdElem, p::NfOrdIdl) -> fmpz
    valuation(a::fmpz, p::NfOrdIdl) -> fmpz

> Computes the $\mathfrak p$-adic valuation of $a$, that is, the largest $i$
> such that $a$ is contained in $\mathfrak p^i$.
"""
function valuation(a::fmpz, p::NfOrdIdl)
  if p.splitting_type[1] == 0
    return valuation_naive(order(p)(a), p)
  end
  P = p.gen_one
  return valuation(a, P)* p.splitting_type[1]
end
valuation(a::Integer, p::NfOrdIdl) = valuation(fmpz(a), p)

#TODO: some more intelligence here...
function valuation_naive(A::NfOrdIdl, B::NfOrdIdl)
  @assert !isone(B)
  Bi = inv(B)
  i = 0
  C = simplify(A* Bi)
  while denominator(C) == 1
    C = simplify(Bi*C)
    i += 1
  end
  return i
end

#TODO: some more intelligence here...
#      in non-maximal orders, interesting ideals cannot be inverted
#      maybe this needs to be checked...
function valuation_naive(x::NfOrdElem, B::NfOrdIdl)
  @assert !isone(B)
  i = 0
  C = B
  while x in C
    i += 1
    C *= B
  end
  return i
end

doc"""
***
    valuation(A::NfOrdIdl, p::NfOrdIdl) -> fmpz

> Computes the $\mathfrak p$-adic valuation of $A$, that is, the largest $i$
> such that $A$ is contained in $\mathfrak p^i$.
"""
function valuation(A::NfOrdIdl, p::NfOrdIdl)
  _assure_weakly_normal_presentation(A)
  if !isdefined(p, :splitting_type) || p.splitting_type[1] == 0 #ie. if p is non-prime...
    return valuation_naive(A, p)
  end
  return min(valuation(A.gen_one, p), valuation(elem_in_nf(A.gen_two), p))
end

################################################################################
#
#   Prime ideals iterator
#
################################################################################

mutable struct PrimeIdealsSet
  order::NfOrd
  from::fmpz
  to::fmpz
  primes::PrimesSet{fmpz}
  currentprime::fmpz
  currentindex::Int
  decomposition::Array{Tuple{NfOrdIdl, Int}, 1}
  proof::Bool
  indexdivisors::Bool
  ramified::Bool
  iscoprimeto::Bool
  coprimeto::NfOrdIdl
  degreebound::Int
  unbound::Bool

  function PrimeIdealsSet(O::NfOrd)
    z = new()
    z.order = O
    z.proof = false
    z.indexdivisors = true
    z.ramified = true
    z.unbound = false
    z.degreebound =  degree(O)
    z.iscoprimeto = false
    return z
  end
end

doc"""
***
    PrimeIdealsSet(O::NfOrd, f, t; proof = false,
                                   indexdivisors = true,
                                   ramified = true,
                                   degreebound = degree(O),
                                   coprimeto = false)

Returns an iterable object $S$ representing the prime ideals $\mathfrak p$
of $\mathcal O$ with $f \leq \min(\mathfrak p) \leq t$.

The optional arguments can be used to exclude index divisors, ramified prime
ideals and to include only prime ideals with degree less or equal than
`degreebound` and which are coprime to `coprimeto`.

If $t=-1$, then the upper bound is infinite.

If `coprimeto` is supplied, it must be either an integer, an element of $\mathcal O$,
or a non-zero ideal of $\mathcal O$.
"""
function PrimeIdealsSet(O::NfOrd, from::T, to::S;
                       proof::Bool = false,
                       indexdivisors::Bool = true,
                       ramified::Bool = true,
                       degreebound::Int = degree(O),
                       coprimeto = false) where {T <: Union{fmpz, Integer},
                                                 S <: Union{fmpz, Integer}}
  from < 0 && error("Lower bound must be non-negative")
  to < -1 && error("Upper bound must be non-negative or -1")

  z = PrimeIdealsSet(O)
  z.from = fmpz(from)
  z.to = fmpz(to)
  z.primes = PrimesSet(z.from, z.to)
  if to == -1
    z.unbound = true
  end
  z.proof = proof
  z.indexdivisors = indexdivisors
  z.ramified = ramified
  z.degreebound = degreebound
  if !(coprimeto isa Bool)
    if coprimeto isa NfOrdIdl
      z.coprimeto = coprimeto
    elseif coprimeto isa Union{Integer, fmpz, NfOrdElem}
      z.coprimeto = ideal(O, coprimeto)
    else
      error("Coprime argument of wrong type ($(typeof(coprimeto)))")
    end
    z.iscoprimeto = true
  end
  return z
end

function Base.start(S::PrimeIdealsSet)
  O = S.order
  pstate = start(S.primes)
  found_prime = false
  while !found_prime
    (p, pstate) = next(S.primes, pstate)
    if !S.indexdivisors && isindex_divisor(O, p)
      continue
    end
    lP = prime_decomposition(O, p)
    j = -1
    for i in 1:length(lP)
      e = lP[i][2]
      if !S.ramified && e > 1
        continue
      end
      P = lP[i][1]
      if P.splitting_type[2] > S.degreebound
        continue
      end
      if S.iscoprimeto && !iscoprime(P, S.coprimeto)
        continue
      end
      j = i
      break
    end
    if j != -1
      S.decomposition = lP
      S.currentprime = p
      S.currentindex = j
      return (p, j)
    end
  end
end

function Base.next(S::PrimeIdealsSet, x)
  pstate = x[1]
  j = x[2]
  Q = S.decomposition[j][1] # This we want to return
  newindex = -1
  lP = S.decomposition
  O = S.order

  # Find the next prime ideal in the current decomposition
  for i in (j + 1):length(S.decomposition)
    e = lP[i][2]
    if !S.ramified && e > 1
      continue
    end
    P = lP[i][1]
    if P.splitting_type[2] > S.degreebound
      continue
    end
    newindex = i
    break
  end

  if newindex != -1
    return Q, (pstate, newindex)
  else
    # We have to change the prime
    found_prime = false
    while !found_prime
      (p, pstate) = next(S.primes, pstate)
      if !S.indexdivisors && isindex_divisor(O, pstate)
        continue
      end
      lP = prime_decomposition(O, pstate)
      j = -1
      for i in 1:length(lP)
        e = lP[i][2]
        if !S.ramified && e > 1
          continue
        end
        P = lP[i][1]
        if P.splitting_type[2] > S.degreebound
          continue
        end
        if S.iscoprimeto && !iscoprime(P, S.coprimeto)
          continue
        end
        j = i
        break
      end
      if j != -1
        S.decomposition = lP
        S.currentprime = p
        S.currentindex = j
        return Q, (pstate, j)
      end
    end
  end
end

function Base.done(S::PrimeIdealsSet, x)
  pstate = x[1]
  index = x[2]
  return !S.unbound && pstate > S.to
end

Base.eltype(::PrimeIdealsSet) = NfOrdIdl

Base.iteratorsize(::Type{PrimeIdealsSet}) = Base.SizeUnknown()


# implements Tries

# `shiftby` is equal to the number of bits required to represent index information
# for one level of the BitmappedTrie.
#
# Here, `shiftby` is 5, which means that the BitmappedTrie Arrays will be length 32.
const shiftby = 5
const trielen = 2^shiftby

abstract type BitmappedTrie{T} end

# Copy elements from one Array to another, up to `n` elements.
#
function copy_to(from::Array{T}, to::Array{T}, n::Int) where {T}
    for i=1:n
        to[i] = from[i]
    end
    to
end

# Copies elements from one Array to another of size `len`.
#
copy_to_len(from::Array{T}, len::Int) where {T} =
    copy_to(from, Array{T,1}(undef, len), min(len, length(from)))

mask(t::BitmappedTrie, i::Int) = (((i - 1) >>> shift(t)) & (trielen - 1)) + 1

Base.lastindex(t::BitmappedTrie) = length(t)

Base.length(t::BitmappedTrie) =
    error("$(typeof(t)) does not implement Base.length")
shift(t::BitmappedTrie) =
    error("$(typeof(t)) does not implement FunctionalCollections.shift")
maxlength(t::BitmappedTrie) =
    error("$(typeof(t)) does not implement FunctionalCollections.maxlength")
arrayof(t::BitmappedTrie) =
    error("$(typeof(t)) does not implement FunctionalCollections.arrayof")

function Base.isequal(t1::BitmappedTrie, t2::BitmappedTrie)
    length(t1)    == length(t2)    &&
    shift(t1)     == shift(t2)     &&
    maxlength(t1) == maxlength(t2) &&
    arrayof(t1)   == arrayof(t2)
end

==(t1::BitmappedTrie, t2::BitmappedTrie) = isequal(t1, t2)

# Dense Bitmapped Tries
# =====================

abstract type DenseBitmappedTrie{T} <: BitmappedTrie{T} end

# Why is the shift value of a DenseLeaf 5 instead of 0, and why does
# the shift value of a DenseNode start at 10?
#
# The PersistentVector implements a "tail" optimization, where it
# inserts appended elements into a tail array until that array is 32
# elements long, only then inserting it into the actual bitmapped
# vector trie. This significantly increases the performance of
# operations that touch the very end of the vector (last, append, pop,
# etc.) because you don't have to traverse the trie.
#
# However, it adds a small amount of complexity to the implementation;
# when you query the trie, you now get back a length-32 array instead
# of the actual element. This is why the DenseLeaf has a shift value
# of 5: it leaves an "extra" 5 bits for the PersistentVector to use to
# index into the array returned from the trie. (This also means that a
# DenseNode has to start at shiftby*2.)

struct DenseNode{T} <: DenseBitmappedTrie{T}
    arr::Vector{DenseBitmappedTrie{T}}
    shift::Int
    length::Int
    maxlength::Int
end

struct DenseLeaf{T} <: DenseBitmappedTrie{T}
    arr::Vector{T}
end
DenseLeaf{T}() where {T} = DenseLeaf{T}(T[])

arrayof(    node::DenseNode) = node.arr
shift(      node::DenseNode) = node.shift
maxlength(  node::DenseNode) = node.maxlength
Base.length(node::DenseNode) = node.length

arrayof(    leaf::DenseLeaf) = leaf.arr
shift(          ::DenseLeaf) = shiftby
maxlength(  leaf::DenseLeaf) = trielen
Base.length(leaf::DenseLeaf) = length(arrayof(leaf))

function promoted(node::DenseBitmappedTrie{T}) where T
    DenseNode{T}(DenseBitmappedTrie{T}[node],
                 shift(node) + shiftby,
                 length(node),
                 maxlength(node) * trielen)
end

function demoted(node::DenseNode{T}) where T
    if shift(node) == shiftby * 2
        DenseLeaf{T}(T[])
    else
        DenseNode{T}(DenseBitmappedTrie{T}[],
                     shift(node) - shiftby,
                     0,
                     round(Int, maxlength(node) / trielen))
    end
end

function witharr(node::DenseNode{T}, arr::Array, lenshift::Int=0) where T
    DenseNode{T}(arr, shift(node), length(node) + lenshift, maxlength(node))
end
witharr(leaf::DenseLeaf{T}, arr::Array) where {T} = DenseLeaf{T}(arr)

function append(leaf::DenseLeaf, el)
    if length(leaf) < maxlength(leaf)
        newarr = copy_to_len(arrayof(leaf), 1 + length(leaf))
        newarr[end] = el
        witharr(leaf, newarr)
    else
        append(promoted(leaf), el)
    end
end
function append(node::DenseNode{T}, el) where T
    if length(node) == 0
        child = append(demoted(node), el)
        witharr(node, DenseBitmappedTrie{T}[child], 1)
    elseif length(node) < maxlength(node)
        if length(arrayof(node)[end]) == maxlength(arrayof(node)[end])
            newarr = copy_to_len(arrayof(node), 1 + length(arrayof(node)))
            newarr[end] = append(demoted(node), el)
            witharr(node, newarr, 1)
        else
            newarr = arrayof(node)[:]
            newarr[end] = append(newarr[end], el)
            witharr(node, newarr, 1)
        end
    else
        append(promoted(node), el)
    end
end
push(leaf::DenseLeaf, el) = append(leaf, el)
push(node::DenseNode, el) = append(node, el)

Base.getindex(leaf::DenseLeaf, i::Int) = arrayof(leaf)[mask(leaf, i)]
Base.getindex(node::DenseNode, i::Int) = arrayof(node)[mask(node, i)][i]

function assoc(leaf::DenseLeaf{T}, i::Int, el) where T
    newarr = arrayof(leaf)[:]
    newarr[mask(leaf, i)] = el
    DenseLeaf{T}(newarr)
end
function assoc(node::DenseNode, i::Int, el)
    newarr = arrayof(node)[:]
    idx = mask(node, i)
    newarr[idx] = assoc(newarr[idx], i, el)
    witharr(node, newarr)
end

peek(bt::DenseBitmappedTrie) = bt[end]

# Pop is usually destructive, but that doesn't make sense for an immutable
# structure, so `pop` is defined to return a Trie without its last
# element. Use `peek` to access the last element.
#
pop(leaf::DenseLeaf) = witharr(leaf, arrayof(leaf)[1:end-1])
function pop(node::DenseNode)
    newarr = arrayof(node)[:]
    newarr[end] = pop(newarr[end])
    witharr(node, newarr, -1)
end

# Sparse Bitmapped Tries
# ======================

abstract type SparseBitmappedTrie{T} <: BitmappedTrie{T} end

struct SparseNode{T} <: SparseBitmappedTrie{T}
    arr::Vector{SparseBitmappedTrie{T}}
    shift::Int
    length::Int
    maxlength::Int
    bitmap::Int
end
SparseNode(T::Type) = SparseNode{T}(SparseBitmappedTrie{T}[], shiftby*7, 0, trielen^7, 0)

struct SparseLeaf{T} <: SparseBitmappedTrie{T}
    arr::Vector{T}
    bitmap::Int
end
SparseLeaf{T}() where {T} = SparseLeaf{T}(T[], 0)

arrayof(    n::SparseNode) = n.arr
shift(      n::SparseNode) = n.shift
maxlength(  n::SparseNode) = n.maxlength
Base.length(n::SparseNode) = n.length

arrayof(    l::SparseLeaf) = l.arr
shift(       ::SparseLeaf) = 0
maxlength(  l::SparseLeaf) = trielen
Base.length(l::SparseLeaf) = length(arrayof(l))

function demoted(n::SparseNode{T}) where T
    shift(n) == shiftby ?
    SparseLeaf{T}(T[], 0) :
    SparseNode{T}(SparseBitmappedTrie{T}[],
                  shift(n) - shiftby,
                  0,
                  round(Int, maxlength(n) / trielen), 0)
end

bitpos(  t::SparseBitmappedTrie, i::Int) = 1 << (mask(t, i) - 1)
hasindex(t::SparseBitmappedTrie, i::Int) = t.bitmap & bitpos(t, i) != 0
index(   t::SparseBitmappedTrie, i::Int) =
    1 + count_ones(t.bitmap & (bitpos(t, i) - 1))

function update(l::SparseLeaf{T}, i::Int, el::T) where T
    hasi = hasindex(l, i)
    bitmap = bitpos(l, i) | l.bitmap
    idx = index(l, i)
    if hasi
        newarr = arrayof(l)[:]
        newarr[idx] = el
    else
        newarr = vcat(arrayof(l)[1:idx-1], [el], arrayof(l)[idx:end])
    end
    (SparseLeaf{T}(newarr, bitmap), !hasi)
end
function update(n::SparseNode{T}, i::Int, el::T) where T
    bitmap = bitpos(n, i) | n.bitmap
    idx = index(n, i)
    if hasindex(n, i)
        newarr = arrayof(n)[:]
        updated, inc = update(newarr[idx], i, el)
        newarr[idx] = updated
    else
        child, inc = update(demoted(n), i, el)
        newarr = vcat(arrayof(n)[1:idx-1], [child], arrayof(n)[idx:end])
    end
    (SparseNode{T}(newarr,
                   n.shift,
                   inc ? n.length + 1 : n.length,
                   n.maxlength, bitmap),
     inc)
end

Base.get(n::SparseLeaf, i::Int, default) =
    hasindex(n, i) ? arrayof(n)[index(n, i)] : default
Base.get(n::SparseNode, i::Int, default) =
    hasindex(n, i) ? get(arrayof(n)[index(n, i)], i, default) : default

function initial_state(t::SparseBitmappedTrie)
    t.length == 0 && return Int[]
    ones(Int, 1 + round(Int, t.shift / shiftby))
end

function Base.iterate(t::SparseBitmappedTrie, state = initial_state(t))
    if isempty(state)
        return nothing
    else
        item = directindex(t, state)
        while true
            index = pop!(state)
            node = directindex(t, state)
            if length(node) > index
                push!(state, index + 1)
                return item, vcat(state, ones(Int, 1 + round(Int, t.shift / shiftby) -
                                                   length(state)))
            elseif node === arrayof(t)
                return item, Int[]
            end
        end
    end
end

function directindex(t::SparseBitmappedTrie, v::Vector{Int})
    isempty(v) && return arrayof(t)
    local node = arrayof(t)
    for i=v
        node = node[i]
        node = isa(node, SparseBitmappedTrie) ? arrayof(node) : node
    end
    node
end


# mcmc_inference.jl
# 2019-11-11
# David Merrell
#
# Contains functions relevant for inference
# in our network reconstruction setting.

"""
Generic MCMC inference wrapper.
Arguments:
    * gen_model: a Gen probabilistic program
    * model_args: a tuple of arguments for the "gen_model"
    * update_loop_fn: a function which updates the variables in "gen_model"
    * update_loop_args: a tuple of arguments for "update_loop_fn"
    * update_results_fn: a function which updates the data structure holding MCMC results
    * update_results_args: arguments for "update_results_fn"
    * timeout: maximum execution time (in seconds) (default 1 hour)
    * n_steps: maximum number of iterations (default no limit)
    * thinning: only store results once every "thinning"th iteration
"""
function mcmc_inference(gen_model,
                        model_args,
                        observations,
                        update_loop_fn::Function,
                        update_loop_args::Vector,
                        update_results_fn::Function,
                        update_results_args::Vector;
                        timeout::Float64=3600.0,
                        n_steps::Int64=-1,
                        thinning::Int64=1,
                        chain_id::String="")

    # start the timer
    t_start = time()
    t_elapsed = 0.0

    # Generate an initial trace
    tr, _ = Gen.generate(gen_model, model_args,
			 observations)

    # The results we care about
    results = update_results_fn(nothing, tr, update_results_args)

    # Sampling loop
    print_str = string("Sampling for ",string(timeout), " seconds.")
    if n_steps > 0
        print_str = string(print_str, " (Or ", n_steps, " iterations.)")
    else
        n_steps = typemax(Int64)
        print_str = string(print_str, " (No iteration limit.)")
    end
    println(print_str)

    prop_count = 0
    t_print = 0.0
    while t_elapsed < timeout && prop_count <= n_steps

        # thinning loop (if applicable)
        for i=1:thinning

            # Print progress
            if (prop_count > 0) && (t_elapsed - t_print >= 20.0)
                println("[", chain_id, "]\t", prop_count," iterations;\t", round(t_elapsed), " seconds")
                t_print = t_elapsed
            end

            # update the variables
            tr, acc = update_loop_fn(tr, update_loop_args)
            prop_count += 1

            t_elapsed = time() - t_start
        end

        # update the results
        results = update_results_fn(results, tr, update_results_args)

    end

    return results
end


using CSFML.LibCSFML
using ModernGL
include("ModernGL_utils.jl")

# https://www.tomdalling.com/blog/modern-opengl/01-getting-started-in-xcode-and-visual-cpp/

### Init

mode = sfVideoMode(700, 720, 32)
# settings = SF::ContextSettings.new(
#   depth: 24, stencil: 8, antialiasing: 4,
#   major: 3, minor: 0
# )
window = sfRenderWindow_create(mode, "SFML window", sfResize | sfClose, C_NULL)
@assert window != C_NULL
sfWindow_setVerticalSyncEnabled(window,true)
sfWindow_setActive(window, sfTrue)
println(createcontextinfo()) # used by ModernGL_utils;
glEnable(GL_DEPTH_TEST)
glDepthFunc(GL_LESS)
glClearColor(0.4 ,0.2, 0.3, 1.0)

### shaders

vsh = """
    $(get_glsl_version_string())
    //layout(location = 0) in vec3 v_position;
    //layout(location = 1) in vec3 v_color;

    in vec3 v_position;
    //in vec3 v_color;
    //out vec3 initial_Color;
    void main() {
        //initial_Color = v_color;
        gl_Position = vec4(v_position, 1.0);
    }
    """

fsh = """
    $(get_glsl_version_string())
    // in vec3 initial_Color;
    out vec4 outColor;
    void main() {
        outColor = vec4(5.0, 2.0, 3.0, 1.0);
        //outColor = vec4(frag_colour, 1.0 );
    }
    """

vertexShader = createShader(vsh, GL_VERTEX_SHADER)
fragmentShader = createShader(fsh, GL_FRAGMENT_SHADER)

function bindAttrib(shader_program)
    #glBindAttribLocation(shader_program, 0, "v_position");
    #glBindAttribLocation(shader_program, 1, "v_color");
end

program = createShaderProgram( bindAttrib, vertexShader, fragmentShader)
#glUseProgram(program)
#println("program : $program")
print("shader ---> ")
glCheckError()

### vertex, colors and buffers

vao = glGenVertexArray()
glBindVertexArray(vao)

points_vbo = glGenBuffer()
glBindBuffer(GL_ARRAY_BUFFER, points_vbo)

vertexData = GLfloat[
        #  X     Y     Z       U     V
         0.0, 0.8, 0.0,
        -0.8,-0.8, 0.0,
         0.8,-0.8, 0.0
]

glBufferData(GL_ARRAY_BUFFER, sizeof(vertexData), vertexData, GL_STATIC_DRAW)
println("points_vbo : $points_vbo")
glCheckError()

glEnableVertexAttribArray(0);
glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 0 , C_NULL)

glBindBuffer(GL_ARRAY_BUFFER, 0)
glBindVertexArray(0)

###

event_ref = Ref{sfEvent}()
clock = sfClock_create()
# Let's express time in seconds.
function get_time(clock::Ptr{sfClock})
    sfTime_asSeconds(sfClock_getElapsedTime(clock))
end
function restart(clock::Ptr{sfClock})
     sfTime_asSeconds(sfClock_restart(clock))
end

target_frequency = 24.0
running = true
t0 = 9999999
while (running)

    while Bool(sfRenderWindow_isOpen(window)) && running
        frame_timestamp = get_time(clock)
        frame_time = frame_timestamp - t0
        global t0 = frame_timestamp
        actual_frequency = 1 / frame_time

        # process events
        while Bool(sfRenderWindow_pollEvent(window, event_ref))
            # close window : exit
            event_ref[].type == sfEvtClosed && sfRenderWindow_close(window)
            event_ref[].type == sfEvtResized && println("Trigger sfEvtResized.")
            event_ref[].type == sfEvtLostFocus && println("Trigger sfEvtLostFocus.")
            event_ref[].type == sfEvtGainedFocus && println("Trigger sfEvtGainedFocus.")
            event_ref[].type == sfEvtTextEntered && println("Trigger sfEvtTextEntered: $(event_ref[].text.unicode)")
            event_ref[].type == sfEvtKeyPressed && println("Trigger sfEvtKeyPressed: $(event_ref[].key.code)")
            event_ref[].type == sfEvtKeyReleased && println("Trigger sfEvtKeyReleased: $(event_ref[].key.code)")
            event_ref[].type == sfEvtMouseWheelMoved && println("Trigger sfEvtMouseWheelMoved: $(event_ref[].mouseWheel.x), $(event_ref[].mouseWheel.y)")
            event_ref[].type == sfEvtMouseWheelScrolled && println("Trigger sfEvtMouseWheelScrolled: $(event_ref[].mouseWheelScroll.wheel)")
            event_ref[].type == sfEvtMouseButtonPressed && println("Trigger sfEvtMouseButtonPressed: $(event_ref[].mouseButton.button)")
            event_ref[].type == sfEvtMouseButtonReleased && println("Trigger sfEvtMouseButtonReleased: $(event_ref[].mouseButton.x), $(event_ref[].mouseButton.y)")
            event_ref[].type == sfEvtMouseMoved && println("Trigger sfEvtMouseMoved: $(event_ref[].mouseMove.x), $(event_ref[].mouseMove.y)")

            if event_ref[].type == sfEvtClosed
                global running = false
                println("Render window closed.")
            elseif event_ref[].type == sfEvtResized
                glViewport(9, 9, event_ref[].size.width, event_ref[].size.height)
                println("Viewport updated.")
            elseif event_ref[].type == sfEvtKeyPressed
                if event_ref[].key.code == sfKeyQ
                    global running = false
                    println("Quiting application.")
                end
            end
        end
        time_event_processing = round(get_time(clock) - t0, digits=9)

        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)

        glUseProgram(program)
        glBindVertexArray(vao)
        glDrawArrays(GL_TRIANGLES, 0 , 3)
        glBindVertexArray(0)
        glUseProgram(0)

        time_render = round(get_time(clock) - time_event_processing - t0, digits=9)
        sfRenderWindow_display(window)
        time_to_sleep = 1 / target_frequency - time_render - time_event_processing

        if time_to_sleep > 0.005
            sleep(time_to_sleep-0.005)
        end
    end
end

sfRenderWindow_destroy(window)


"""
    CrystalNets.jl

Module for automatic recognition of crystal net topologies.
To use as an executable, run the source file in a shell:
```bash
julia --project=/path/to/CrystalNets/ /path/to/CrystalNets/src/CrystalNets.jl
```
Otherwise, as a module, to try to recognize the net underlying a crystal given in a
chemical file format called FILE, the entry point is the following execution:
```julia
julia> using CrystalNets

julia> determine_topology(FILE)
```

See the documentation at http://molsim.info/CrystalNets.jl/ for further instructions on
on the use of CrystalNets.jl, and how to do a full installation to reduce latency at
startup.
"""
module CrystalNets

export CrystalNet,
       CrystalNet1D, CrystalNet2D, CrystalNet3D,
       UnderlyingNets,
       TopologicalGenome,
       TopologyResult,
       determine_topology,
       determine_topology_dataset,
       guess_topology,
       guess_topology_dataset,
       parse_chemfile,
       topological_genome,
       recognize_topology,
       StructureType,
       Bonding,
       Clustering,
       ClusterKinds

import LinearAlgebra: det, norm, rank, cross
using Base.Threads
import Serialization

using PeriodicGraphEmbeddings
using PeriodicGraphs
using StaticArrays
using Graphs

import Logging

const DOWARN = Base.RefValue{Bool}(false)
const DOEXPORT = Base.RefValue{Bool}(false)

"""
    toggle_warning(to=nothing)

Toggle warnings on (if `to == true`) or off (if `to == false`).
Without an argument, toggle on and off repeatedly at each call.
"""
function toggle_warning(to=nothing)
    global DOWARN
    DOWARN[] = to isa Nothing ? !DOWARN[] : to
end

"""
    toggle_export(to=nothing)

Toggle exports on (if `to == true`) or off (if `to == false`).
Without an argument, toggle on and off repeatedly at each call.
"""
function toggle_export(to=nothing)
    global DOEXPORT
    DOEXPORT[] = to isa Nothing ? !DOEXPORT[] : to
end

function __init__()
    toggle_warning("--no-warn" ∉ ARGS)
    toggle_export("--no-export" ∉ ARGS)
    @static if abspath(PROGRAM_FILE) == @__FILE__
        global DOWARN
        if DOWARN[]
            @info "Proceed to a full installation (see README.md) for better performance."
        end
        @show "EXECUTING JULIA MAIN"
        ret_code = CrystalNets.julia_main()
        exit(ret_code)
    end
    nothing
end

include("utils.jl")
__precompile__(true)
include("specialsolver.jl") # Exact sparse symmetric integer matrix factorization
include("options.jl") # Computation options
include("types.jl") # Main internal type definitions used to represent topologies
include("input.jl") # Crystal file parsing and conversion to an internal type
include("archive.jl") # Manipulation of the topological archive
include("output.jl") # Crystal file exports
include("arithmetics.jl") # Handling of (possibly sparse) integer/rational matrices
include("stability.jl") # Functions related to unstable nets
include("topology.jl") # Main functions of the algorithm
include("query.jl") # Entry point for the user-facing functions
include("executable.jl") # Entry point for the argument parsing of the executable
include("precompile.jl")

end # module CrystalNets

nothing


function emldescriptors(x, q, t, a, b, c, pbc, eta, kappa, emdescriptors, mldescriptors=nothing)

    if length(emdescriptors) > 0
        rcutmax = 0.0
        for i = 1:length(emdescriptors)
            rcutmax = max(rcutmax, emdescriptors[i].rcut)
        end
        ed, edd, edv = empiricaldescriptors(x, q, t, a, b, c, pbc, rcutmax, eta, kappa, emdescriptors)
        sz = size(edd)
        if length(sz)==3
            edd = reshape(edd, (sz[1]*sz[2], sz[3]))
        else
            edd = reshape(edd, (sz[1]*sz[2], 1))
        end
    end

    # bispectrum, bispectrum derivatives, bispectrum virial
    if mldescriptors !== nothing
        if mldescriptors.name == "SNAP"
            rcutmax = 0.0
            for ielem = 1:mldescriptors.ntypes
                rcutmax = max(2.0*mldescriptors.radelem[1+ielem]*mldescriptors.rcutfac,rcutmax)
            end
            mldescriptors.rcutmax = rcutmax
            mld, mldd, mldv = snapdescriptors(x, t, a, b, c, pbc, mldescriptors)
            sz = size(mldd)
            mldd = reshape(mldd, (sz[1]*sz[2], sz[3]))
        end
    end

    if (length(emdescriptors) > 0) & (mldescriptors !== nothing)
        d = [ed[:]; mld[:]]
        dd = cat(edd, mldd, dims=2);
        dv = cat(edv, mldv, dims=2);
        return d, dd, dv
    elseif length(emdescriptors) > 0
        return ed, edd, edv
    elseif (mldescriptors !== nothing)
        return mld, mldd, mldv
    else
        error("empirical descriptors and ML descriptors are empty")
    end
end

function emldescriptors(config, indices, eta, kappa, normalizeenergy, emdescriptors,
                mldescriptors=nothing, potential=nothing)

normalizestress = 0
# cumalative sum of numbers of atoms
nconfigs = length(config.natom)
natom = [0; cumsum(config.natom[:])];

nf = zeros(Int64, length(indices))
be = []
bf = []
bs  = []
Ae = []
Af = []
As = []
for i = 1:length(indices)
    ci = indices[i]
    if ci > nconfigs
        error("Indices are greater than nconfigs")
    end

    x = config.x[:,(natom[ci]+1):natom[ci+1]] # atom positions of configuration ci
    t = config.t[:,(natom[ci]+1):natom[ci+1]] # atom types of configuration ci
    if length(config.q)>0
        q = config.q[:,(natom[ci]+1):natom[ci+1]]
    else
        q = []
    end
    if size(config.a,2) > 1
        a = config.a[:,ci]
        b = config.b[:,ci]
        c = config.c[:,ci]
    else
        a = config.a[:,1]
        b = config.b[:,1]
        c = config.c[:,1]
    end
    if size(config.pbc,2) > 1
        pbc = config.pbc[:,ci]
    else
        pbc = config.pbc[:,1]
    end

    d, dd, dv = emldescriptors(x, q, t, a[:], b[:], c[:], pbc[:], eta, kappa, emdescriptors, mldescriptors)

    if normalizestress==1
        vol = latticevolume(a[:], b[:], c[:])
        dv = (1.6021765e6/vol) * dv
    end

    if potential !== nothing
        rcutmax = 0.0
        for j = 1:length(potential)
            rcutmax = max(rcutmax, potential[j].rcut)
        end
        eref, e, fref, v, sref = empiricalpotential(x, q, t, a[:], b[:], c[:], pbc[:], rcutmax, eta, kappa, potential)
        if normalizestress==1
            vol = latticevolume(a[:], b[:], c[:])
            sref = (1.6021765e6/vol) * sref
        end
    else
        eref = 0.0;
        fref = [0.0];
        sref = [0.0];
    end

    if normalizeenergy==1
        d = reshape(d, (1,length(d)))/config.natom[ci]
    else
        d = reshape(d, (1,length(d)))
    end

    nf[i] = size(dd,1)
    if i == 1
        Ae = d
        Af = dd
        As = dv
    else
        Ae = cat(Ae, d, dims=1)
        Af = cat(Af, dd, dims=1)
        As = cat(As, dv, dims=1)
    end

    if length(config.e)>0
        e = config.e[indices[i]] - eref
        if normalizeenergy==1
            e = e/config.natom[indices[i]]
        end
        be = [be; e]
    end
    if length(config.f)>0
        f = config.f[:,(natom[ci]+1):natom[ci+1]]
        bf = [bf; f[:] .- fref[:]]
    end
    if length(config.stress)>0
        if size(config.stress,1)==9
            ind = [1, 5, 9, 6, 3, 2]
            s = config.stress[ind,ci]
        else
            s = config.stress[:,ci]
        end
        bs = [bs; s[:] .- sref[:]]
    end
end

return Ae, Af, As, be, bf, bs, nf

end

function emldescriptors2(config, indices, eta, kappa, normalizeenergy, normalizestress,
        descriptors, potential=nothing)

# getdescriptors
emdescriptors, mldescriptors = getdescriptors(descriptors)

# cumalative sum of numbers of atoms
nconfigs = length(config.natom)
natom = [0; cumsum(config.natom[:])];

nf = zeros(Int64, length(indices))
be = []
bf = []
bs  = []
Ae = []
Af = []
As = []
for i = 1:length(indices)
    ci = indices[i]
    if ci > nconfigs
        error("Indices are greater than nconfigs")
    end

    x = config.x[:,(natom[ci]+1):natom[ci+1]] # atom positions of configuration ci
    t = config.t[:,(natom[ci]+1):natom[ci+1]] # atom types of configuration ci
    if length(config.q)>0
        q = config.q[:,(natom[ci]+1):natom[ci+1]]
    else
        q = []
    end
    if size(config.a,2) > 1
        a = config.a[:,ci]
        b = config.b[:,ci]
        c = config.c[:,ci]
    else
        a = config.a[:,1]
        b = config.b[:,1]
        c = config.c[:,1]
    end
    if size(config.pbc,2) > 1
        pbc = config.pbc[:,ci]
    else
        pbc = config.pbc[:,1]
    end

    d, dd, dv = emldescriptors(x, q, t, a[:], b[:], c[:], pbc[:], eta, kappa, emdescriptors, mldescriptors)

    if normalizestress==1
        vol = latticevolume(a[:], b[:], c[:])
        dv = (1.6021765e6/vol) * dv
    end

    if (potential !== nothing)
        if (length(potential) > 0)
            rcutmax = 0.0
            for j = 1:length(potential)
                rcutmax = max(rcutmax, potential[j].rcut)
            end
            eref, e, fref, v, sref = empiricalpotential(x, q, t, a[:], b[:], c[:], pbc[:], rcutmax, eta, kappa, potential)
            if normalizestress==1
                vol = latticevolume(a[:], b[:], c[:])
                sref = (1.6021765e6/vol) * sref
            end
        else
            eref = 0.0; fref = [0.0]; sref = [0.0];
        end
    else
        eref = 0.0; fref = [0.0]; sref = [0.0];
    end

    if normalizeenergy==1
        d = reshape(d, (1,length(d)))/config.natom[ci]
    else
        d = reshape(d, (1,length(d)))
    end

    nf[i] = size(dd,1)
    if i == 1
        Ae = d
        Af = dd
        As = dv
    else
        Ae = cat(Ae, d, dims=1)
        Af = cat(Af, dd, dims=1)
        As = cat(As, dv, dims=1)
    end

    if length(config.e)>0
        e = config.e[indices[i]] - eref
        if normalizeenergy==1
            e = e/config.natom[indices[i]]
        end
        be = [be; e]
    end
    if length(config.f)>0
        f = config.f[:,(natom[ci]+1):natom[ci+1]]
        bf = [bf; f[:] .- fref[:]]
    end
    if length(config.stress)>0
        if size(config.stress,1)==9
            ind = [1, 5, 9, 6, 3, 2]
            s = config.stress[ind,ci]
        else
            s = config.stress[:,ci]
        end
        bs = [bs; s[:] .- sref[:]]
    end

end

return Ae, Af, As, be, bf, bs, nf

end


module Plot2D

using GLMakie
export plot2d, buttons

LOG = nothing
const P1 = [Node(Vector{Point2f0}(undef, 6000))]
const P2 = [Node(Vector{Point2f0}(undef, 6000))]

function autoscale(ax, x, y)
    xlims!(ax, x[1], x[end])
    range = maximum(y) - minimum(y)
    if range > 0.0f0
        ylims!(ax, minimum(y)-0.05*range, maximum(y)+0.05*range)
    end
end

function plot2d(se, ax, label1, log, points, field, lower=false)
    global LOG, P1, P2
    LOG=log
    if lower
        P2[1]=points
    else
        P1[1]=points
    end
    unit = ""
    y = [1f0]
    if field == :height
        unit = "[m]"
        y    = log.extlog.z ./ se().zoom
    elseif field == :elevation
        unit = "[°]"
        y    = log.syslog.elevation/pi*180.0
    elseif field == :azimuth
        unit = "[°]"
        y    = log.syslog.azimuth/pi*180.0
    elseif field == :v_reelout
        unit = "[m/s]"
        y    = log.syslog.v_reelout
    elseif field == :force
        unit = "[N]"
        y    = log.syslog.force
    elseif field == :depower
        unit = "[%]"
        y    = log.syslog.depower * 100.0f0
    elseif field == :v_app
        unit = "[m/s]"
        y    = log.syslog.v_app
    elseif field == :l_tether
        unit = "[m]"
        y    = log.syslog.l_tether
    elseif field == :power
        unit = "[kW]"
        y    = log.syslog.v_reelout .* log.syslog.force * 0.001f0
    end
    x       = log.syslog.time
    if field == :power
        label1[] = "mechanical " * string(field) * " " * unit
    else
        label1[] = string(field) * " " * unit
    end
    points[]    =  Point2f0.(x, y)
    autoscale(ax, x, y)
end

function buttons(fig, bg, se, ax1, ax2, label1, label2, reset)
    textsize=14
    btn_height         = Button(fig, label = "height", textsize=textsize)
    btn_elevation      = Button(fig, label = "elevation", textsize=textsize)
    btn_azimuth        = Button(fig, label = "azimuth", textsize=textsize)
    btn_v_reelout      = Button(fig, label = "v_reelout", textsize=textsize)
    btn_force          = Button(fig, label = "force", textsize=textsize)
    btn_depower        = Button(fig, label = "depower", textsize=textsize)
    btn_v_app          = Button(fig, label = "v_app", textsize=textsize)
    btn_l_tether       = Button(fig, label = "l_tether", textsize=textsize)
    btn_power          = Button(fig, label = "power", textsize=textsize)
    sw = Toggle(fig, active = false) # active means plot in the upper area
    bg[1, 1:10] = [btn_height, btn_elevation, btn_azimuth, btn_v_reelout, btn_force, btn_depower, btn_v_app, btn_l_tether, btn_power, sw]

    on(btn_height.clicks) do c
        if sw.active[]
            plot2d(se, ax1, label1, LOG, P1[1], :height)
        else
            plot2d(se, ax2, label2, LOG, P2[1], :height, true)
        end
        reset()
    end
    on(btn_elevation.clicks) do c
        if sw.active[]
            plot2d(se, ax1, label1, LOG, P1[1], :elevation)
        else
            plot2d(se, ax2, label2, LOG, P2[1], :elevation, true)
        end
        reset()
    end
    on(btn_azimuth.clicks) do c
        if sw.active[]
            plot2d(se, ax1, label1, LOG, P1[1], :azimuth)
        else
            plot2d(se, ax2, label2, LOG, P2[1], :azimuth, true)
        end
        sleep(0.05)
        reset()
    end
    on(btn_v_reelout.clicks) do c
        if sw.active[]
            plot2d(se, ax1, label1, LOG, P1[1], :v_reelout)
        else
            plot2d(se, ax2, label2, LOG, P2[1], :v_reelout, true)
        end
        reset()
    end
    on(btn_force.clicks) do c
        if sw.active[]
            plot2d(se, ax1, label1, LOG, P1[1], :force)
        else
            plot2d(se, ax2, label2, LOG, P2[1], :force, true)
        end
        reset()
    end
    on(btn_depower.clicks) do c
        if sw.active[]
            plot2d(se, ax1, label1, LOG, P1[1], :depower)
        else
            plot2d(se, ax2, label2, LOG, P2[1], :depower, true)
        end
        reset()
    end
    on(btn_v_app.clicks) do c
        if sw.active[]
            plot2d(se, ax1, label1, LOG, P1[1], :v_app)
        else
            plot2d(se, ax2, label2, LOG, P2[1], :v_app, true)
        end
        reset()
    end
    on(btn_l_tether.clicks) do c
        if sw.active[]
            plot2d(se, ax1, label1, LOG, P1[1], :l_tether)
        else
            plot2d(se, ax2, label2, LOG, P2[1], :l_tether, true)
        end
        reset()
    end
    on(btn_power.clicks) do c
        if sw.active[]
            plot2d(se, ax1, label1, LOG, P1[1], :power)
        else
            plot2d(se, ax2, label2, LOG, P2[1], :power)
        end
        reset()
    end
end

end


using Random; import Future
using Base
using Plots

include("rendering.jl")

function run_agent!(env::TE, agent::TA, num_episodes::Int, rng::AbstractRNG)::Array{Float64, 1} where {TE<:AbstractEnvironment,TA<:AbstractAgent}
    returns = zeros(num_episodes)
    step::Int = 0
    for episode in 1:num_episodes
        env_reset!(env, rng)
        new_episode!(agent, rng)
        step = 0

        while !is_terminal(env)
            act!(agent, env, rng)
            returns[episode] += env.reward
        end
    end
    return returns
end

function eval_agent(env::TE, agent::TA, num_episodes::Int, rng::AbstractRNG)::Array{Float64, 1} where {TE<:AbstractEnvironment,TA<:AbstractAgent}
    returns = zeros(num_episodes)
    step::Int = 0
    for episode in 1:num_episodes
        env_reset!(env, rng)
        new_episode!(agent, rng)
        step = 0

        while !is_terminal(env)
            act!(agent, env, rng, false)  # set training flag to false
            returns[episode] += env.reward
        end
    end
    return returns
end

function eval_policy(env::TE, policy, γ, num_episodes::Int, rng::AbstractRNG)::Array{Float64, 1} where {TE<:AbstractEnvironment}
    returns = zeros(num_episodes)
    step::Int = 0
    for episode in 1:num_episodes
        env_reset!(env, rng)
        # new_episode!(agent, rng)
        step = 0

        while !is_terminal(env)
            get_action!(policy, env.state, rng)
            r = step!(env, policy.action, rng)
            returns[episode] += γ^step * r
            step += 1
        end
    end
    return returns
end

function run_agent!(env::AbstractEnvironment, agent::AbstractAgent, num_episodes::Int, rng::AbstractRNG, draw::Bool)::Array{Float64, 1}
    returns = zeros(num_episodes)
    step::Int = 0
    for episode in 1:num_episodes
        env_reset!(env, rng)
        new_episode!(agent, rng)
        step = 0
        if draw
            plt = plot(env)
            display(plt);
            gui()
        end

        while !is_terminal(env)
            act!(agent, env, rng)
            returns[episode] += env.reward
            if draw
                plt = plot(env)
                display(plt);
                gui()
            end
        end
    end
    return returns
end

function run_agent(env::AbstractEnvironment, agent::AbstractAgent, num_trials::Int, num_episodes::Int, rng::MersenneTwister)::Array{Float64, 2}
    returns = zeros(num_episodes, num_trials)
    agents = [clone(agent) for i in 1:num_trials]
    envs = [clone(env) for i in 1:num_trials]

    for trial in 1:num_trials
        for episode in 1:num_episodes
            env_reset!(envs[trial], rng)
            new_episode!(agents[trial], rng)

            while !is_terminal(envs[trial])
                act!(agents[trial], envs[trial], rng)
                returns[episode, trial] += envs[trial].reward
            end
        end
    end

    return returns
end

function run_agent_parallel(env::AbstractEnvironment, agent::AbstractAgent, num_trials::Int, num_episodes::Int, rng::MersenneTwister)::Array{Float64, 2}
    agents = [clone(agent) for i in 1:num_trials]
    envs = [clone(env) for i in 1:num_trials]

    rngs = [rng; accumulate(Future.randjump, fill(big(10)^20, num_trials-1), init=rng)]

    returns = zeros(num_episodes, num_trials)

    Threads.@threads for trial in 1:num_trials
        for episode in 1:num_episodes
            env_reset!(envs[trial], rngs[trial])
            new_episode!(agents[trial], rngs[trial])

            while !is_terminal(envs[trial])
                act!(agents[trial], envs[trial], rngs[trial])
                returns[episode, trial] += envs[trial].reward
            end
        end
    end

    return returns
end


# included from Potentials.jl
# part of the module JuLIP.Potentials

using JuLIP: JVec, JMat, neighbourlist
using LinearAlgebra: I
using JuLIP.Chemistry: atomic_number
using NeighbourLists

export ZeroPairPotential, ZBLPotential,
         LennardJones, lennardjones,
         Morse, morse

# For PairPotentials we do something kind of weird - we create an intentional
# stack overflow by redirecting
#    evaluate -> evaluate! -> evaluate -> evaluate! -> ...
# A concrete instance of PairPotential must overload either evaluate or
# evaluate! to break this infinite loop. The advantage though is that either
# can be implemented.

evaluate(V::PairPotential, r::Number, z1, z0) =
      evaluate!(alloc_temp(V, 1), V, r, z1, z0)
evaluate_d(V::PairPotential, r::Number, z1, z0) =
      evaluate_d!(alloc_temp_d(V, 1), V, r, z1, z0)
evaluate_dd(V::PairPotential, r::Number, z1, z0) =
      evaluate_dd!(alloc_temp_dd(V, 1), V, r, z1, z0)
evaluate!(tmp, V::PairPotential, r::Number, z1, z0) =
      evaluate(V, r, z1, z0)
evaluate_d!(tmp, V::PairPotential, r::Number, z1, z0) =
      evaluate_d(V, r, z1, z0)
evaluate_dd!(tmp, V::PairPotential, r::Number, z1, z0) =
      evaluate_dd(V, r, z1, z0)

# evaluate!(tmp, V::PairPotential, r::Union{Number, JVec}) = V(r)
# evaluate_d!(tmp, V::PairPotential, r::Number) = @D V(r)
# evaluate_dd!(tmp, V::PairPotential, r::Number) = @DD V(r)
# evaluate_d!(tmp, V::PairPotential, R::JVec) =
#       evaluate_d!(tmp, V, norm(R), R)
# evaluate_d!(tmp, V::PairPotential, r::Number, R::JVec) = ((@D V(r))/r) * R
# evaluate_dd!(tmp, V::PairPotential, r::Number) = @DD V(r)
# evaluate_dd!(tmp, V::PairPotential, R::JVec) =
#       evaluate_dd!(tmp, V, norm(R), R)
# evaluate_dd!(tmp, V::PairPotential, r::Number, R::JVec) =
#       _hess!(tmp, V, r, R)

function evaluate!(tmp, V::PairPotential,
                   R::AbstractVector{JVec{T}}, Z, z0) where {T}
   Es = zero(T)
   for i = 1:length(R)
      Es += T(0.5) * evaluate!(tmp, V, norm(R[i]), Z[i], z0)
   end
   return Es
end

function evaluate_d!(dEs, tmp, V::PairPotential,
                     R::AbstractVector{JVec{T}}, Z, z0) where {T}
   for i = 1:length(R)
      r = norm(R[i])
      dEs[i] = (T(0.5) * evaluate_d!(tmp, V, r, Z[i], z0) / r) * R[i]
   end
   return dEs
end

function evaluate_dd!(hEs, tmp, V::PairPotential,
                      R::AbstractVector{<: JVec}, Z, z0)
   n = length(R)
   for i = 1:n
      hEs[i,i] = 0.5 * _hess!(tmp, V, norm(R[i]), R[i], Z[i], z0)
   end
   return hEs
end

function _hess!(tmp, V::PairPotential, r::Number, R::JVec, z1, z0)
   R̂ = R/r
   P = R̂ * R̂'
   dV = evaluate_d!(tmp, V, r, z1, z0) / r
   ddV = evaluate_dd!(tmp, V, r, z1, z0)
   return (ddV - dV) * P + dV * I
end

function precon!(hEs, tmp, V::PairPotential,
                 R::AbstractVector{<: JVec{T}}, Z, z0,
                 innerstab=T(0.0)) where {T}
   n = length(R)
   for i = 1:n
      hEs[i,i] = precon!(tmp, V, norm(R[i]), R[i], Z[i], z0, innerstab)
   end
   return hEs
end

# an FF preconditioner for pair potentials
function precon!(tmp, V::PairPotential, r::T, R::JVec{T}, innerstab=T(0.1)
                 ) where {T <: Number}
   r = norm(R)
   dV = evaluate_d!(tmp, V, r)
   ddV = evaluate_dd!(tmp, V, r)
   R̂ = R/r
   return (1-innerstab) * (abs(ddV) * R̂ * R̂' + abs(dV / r) * (I - R̂ * R̂')) +
             innerstab  * (abs(ddV) + abs(dV / r)) * I
end

# ------- Implementation of SimplePairPotential

evaluate!(tmp, V::SimplePairPotential, r::Number, z1, z0) = evaluate!(tmp, V, r)
evaluate_d!(tmp, V::SimplePairPotential, r::Number, z1, z0) = evaluate_d!(tmp, V, r)
evaluate_dd!(tmp, V::SimplePairPotential, r::Number, z1, z0) = evaluate_dd!(tmp, V, r)
evaluate!(tmp, V::SimplePairPotential, r) = evaluate(V, r)
evaluate_d!(tmp, V::SimplePairPotential, r) = evaluate_d(V, r)
evaluate_dd!(tmp, V::SimplePairPotential, r) = evaluate_dd(V, r)
evaluate(V::SimplePairPotential, r::Number) = evaluate!(alloc_temp(V, 1), V, r)
evaluate_d(V::SimplePairPotential, r::Number) = evaluate_d!(alloc_temp_d(V, 1), V, r)
evaluate_dd(V::SimplePairPotential, r::Number) = evaluate_dd!(alloc_temp_dd(V, 1), V, r)

evaluate_d(V::SimplePairPotential, r::Number, R::JVec) =
      (evaluate_d(V, r)/r) * R

function evaluate_dd(V::SimplePairPotential, r::Number, R::JVec)
   dV = evaluate_d(V, r) / r
   ddV = evaluate_dd(V, r)
   return ((ddV - dV)/r^2) * R * R'  + dV * I
end

# ------- Implementation of ExplicitPairPotential

evaluate(p::ExplicitPairPotential, r::Number) = p.f(r)
evaluate_d(p::ExplicitPairPotential, r::Number) = p.f_d(r)
evaluate_dd(p::ExplicitPairPotential, r::Number) = p.f_dd(r)

# ------- Some concrete potentials

"""
`LennardJones(σ, e0):` constructs the 6-12 Lennard-Jones potential [wiki](https://en.wikipedia.org/wiki/Lennard-Jones_potential)

   e0 * 4 * ( (σ/r)¹² - (σ/r)⁶ )

Constructor with kw arguments: `LennardJones(; kwargs...)`

* `e0, σ` : standard LJ parameters, default `e0 = 1.0, σ = 1.0`
* `r0` : equilibrium distance - if `r0` is specified then `σ` is ignored
* `a0` : FCC lattice parameter, if `a0` is specified then `σ` is ignored

(`r0, a0` cannot both be specified at the same time)
"""
LennardJones(σ, e0) = (@analytic r -> e0 * 4.0 * ((σ/r)^(12) - (σ/r)^(6)))

function ljparams(; σ=1.0, e0=1.0, r0 = nothing, a0 = nothing)
   if r0 != nothing && a0 != nothing
      error("`LenndardJones`: cannot specify both `r0` and `a0`")
   end
   if a0 != nothing    # r0 = nn-dist = a0/sqrt(2) in FCC
      r0 = a0 / sqrt(2)
   end
   if r0 != nothing    # standard LJ is minimised at r = 2^(1/6)
      σ =  r0 / 2^(1/6)
   end
   return σ, e0
end

LennardJones(; kwargs...) = LennardJones(ljparams(;kwargs...)...)

"""
`lennardjones(; kwargs...)`

simplified constructor for `LennardJones` (note this is type unstable!)

In addition to the `kwargs` of `LennardJones`, this accepts also

* `rcut` : default `:auto` which gives `rcut = (1.9*σ, 2.7*σ)`. Use
`nothing` or `Inf` to specify no cutoff, or specify a tuple or
array with two elements specifying the lower and upper cut-off radii to be
used with `SplineCutoff`.
"""
function lennardjones(; rcut = :auto, kwargs...)
   σ, e0 = ljparams(; kwargs...)
   if (rcut == nothing || rcut == Inf)
      return LennardJones(σ, e0)
   elseif rcut == :auto
      rcut = (1.9*σ, 2.7*σ)
   end
   return SplineCutoff(rcut[1], rcut[2]) * LennardJones(σ, e0)
end

"""
`Morse(A, e0, r0)` or `Morse(;A=4.0, e0=1.0, r0=1.0)`: constructs a
`PairPotential` for
```
   e0 ( exp( -2 A (r/r0 - 1) ) - 2 exp( - A (r/r0 - 1) ) )
```
"""
Morse(A, e0, r0) = @analytic(
   r -> e0 * ( exp(-(2.0*A) * (r/r0 - 1.0)) - 2.0 * exp(-A * (r/r0 - 1.0)) ) )
Morse(;A=4.0, e0=1.0, r0=1.0) = Morse(A, e0, r0)

"""
`morse(A=4.0, e0=1.0, r0=1.0, rcut=(1.9*r0, 2.7*r0))`

simplified constructor for `Morse` (type unstable)
"""
morse(;A=4.0, e0=1.0, r0=1.0, rcut=(1.9*r0, 2.7*r0)) = (
   (rcut == nothing || rcut == Inf)
         ?  Morse(A, e0, r0)
         :  SplineCutoff(rcut[1], rcut[2]) * Morse(A, e0, r0) )

"""
`ZeroPairPotential()`: creates a potential that just returns zero
"""
struct ZeroPairPotential <: SimplePairPotential end

@pot ZeroPairPotential

evaluate(p::ZeroPairPotential, r::T) where {T <: Number} = T(0.0)
evaluate_d(p::ZeroPairPotential, r::T) where {T <: Number} = T(0.0)
evaluate_dd(p::ZeroPairPotential, r::T) where {T <: Number} = T(0.0)
cutoff(p::ZeroPairPotential) = Bool(0) # the weakest number type

# ====================================================================
#   A product of two pair potentials: primarily used for cutoff mechanisms

"product of two `SimplePairPotential`"
mutable struct SimpleProdPot{P1, P2} <: SimplePairPotential
   p1::P1
   p2::P2
end

@pot SimpleProdPot

import Base.*
*(p1::SimplePairPotential, p2::SimplePairPotential) = SimpleProdPot(p1, p2)
evaluate(p::SimpleProdPot, r::Number) = p.p1(r) * p.p2(r)
evaluate_d(p::SimpleProdPot, r::Number) = (p.p1(r) * (@D p.p2(r)) + (@D p.p1(r)) * p.p2(r))
evaluate_dd(p::SimpleProdPot, r::Number) = (p.p1(r) * (@DD p.p2(r)) +
              2 * (@D p.p1(r)) * (@D p.p2(r)) + (@DD p.p1(r)) * p.p2(r))
cutoff(p::SimpleProdPot) = min(cutoff(p.p1), cutoff(p.p2))

"product of two `PairPotential`"
mutable struct ProdPot{P1, P2} <: PairPotential
   p1::P1
   p2::P2
end

@pot ProdPot

import Base.*
*(p1::PairPotential, p2::PairPotential) = ProdPot(p1, p2)
evaluate(p::ProdPot, r::Number, z1, z0) =
      p.p1(r, z1, z0) * p.p2(r, z1, z0)
evaluate_d(p::ProdPot, r::Number, z1, z0) =
      (p.p1(r, z1, z0) * (@D p.p2(r, z1, z0)) + (@D p.p1(r, z1, z0)) * p.p2(r, z1, z0))
evaluate_dd(p::ProdPot, r::Number, z1, z0) =
      (   p.p1(r, z1, z0) * (@DD p.p2(r, z1, z0))
        + 2 * (@D p.p1(r, z1, z0)) * (@D p.p2(r, z1, z0))
        + (@DD p.p1(r, z1, z0)) * p.p2(r, z1, z0)
      )
cutoff(p::ProdPot) = min(cutoff(p.p1), cutoff(p.p2))

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

"""
`struct WrappedPairPotential`

wraps a pairpotential using `FunctionWrappers` in order to allow
type-stable storage of multiple potentials. This is the main technique
required at the moment to work with multi-component systems.
Otherwise, this is not advisable since it disables a range of
possible compiler optimisations.
"""
struct WrappedPairPotential <: ExplicitPairPotential
   f::F64fun
   f_d::F64fun
   f_dd::F64fun
   rcut::Float64
end

@pot WrappedPairPotential

cutoff(V::WrappedPairPotential) = V.rcut
# evaluate, etc are all derived from SimplePairPotential

function WrappedPairPotential(V::AnalyticFunction, rcut)
   @assert (0 < rcut < Inf)
   f, f_d, f_dd = let V=V, rc = rcut
      (F64fun(r -> evaluate(V, r) * (r<rc)),
              F64fun(r -> evaluate_d(V, r) * (r<rc)),
              F64fun(r -> evaluate_dd(V, r) * (r<rc)))
   end
   return WrappedPairPotential(f, f_d, f_dd, cutoff(V))
end

function WrappedPairPotential(V::PairPotential)
   @assert (0 < cutoff(V) < Inf)
   f, f_d, f_dd = let V=V, rc = cutoff(V)
      (F64fun(r -> evaluate(V, r) * (r<rc)),
              F64fun(r -> evaluate_d(V, r) * (r<rc)),
              F64fun(r -> evaluate_dd(V, r) * (r<rc)))
   end
   return WrappedPairPotential(f, f_d, f_dd, cutoff(V))
end

# ------------------------------------------------------------------------

# TODO: write more docs + tests for ZBL

"""
Implementation of the ZBL potential to model close approach.
"""
struct ZBLPotential{TV} <: PairPotential
   V::TV   # analytic "inner" potential
end

@pot ZBLPotential

ZBLPotential() = (
   let
      # au = 0.8854 * 0.529 / (Z1^0.23 + Z2^0.23)
      ϵ0 = 0.00552634940621
      C = 1/(4*π*ϵ0)
      E1, E2, E3, E4 = 0.1818, 0.5099, 0.2802, 0.02817
      A1, A2, A3, A4 = 3.2, 0.9423, 0.4028, 0.2016
      V = @analytic(r -> C * (E1*exp(-A1*r) + E2*exp(-A2*r) +
                              E3*exp(-A4*r) + E4*exp(-A4*r) ) / r)
      ZBLPotential(V)
   end)

_zbl_au(Z1, Z2) = (0.8854 * 0.529) / (Z1^0.23 + Z2^0.23)

function evaluate(V::ZBLPotential, r::Number, z1, z0)
   au = _zbl_au(z1, z0)
   return evaluate(V.V, r / au)
end

function evaluate_d(V::ZBLPotential, r::Number, z1, z0)
   au = _zbl_au(z1, z0)
   return evaluate_d(V.V, r / au) / au
end

cutoff(::ZBLPotential) = Inf

write_dict(V::ZBLPotential) = Dict("__id__" => "JuLIP_ZBLPotential")
read_dict(::Val{:JuLIP_ZBLPotential}, D::Dict) = ZBLPotential()


@inline incrementp(A::AbstractStridedPointer{T,3} where T, a::Ptr) = VectorizationBase.increment_ptr(A, a, (Zero(), Zero(), One()))
@inline increment2(B::AbstractStridedPointer{T,2} where T, b::Ptr, ::StaticInt{nᵣ}) where {nᵣ} = VectorizationBase.increment_ptr(B, b, (Zero(), StaticInt{nᵣ}()))
@inline increment1(C::AbstractStridedPointer{T,2} where T, c::Ptr, ::StaticInt{mᵣW}) where {mᵣW} = VectorizationBase.increment_ptr(C, c, (StaticInt{mᵣW}(), Zero()))
macro kernel(pack::Bool, ex::Expr)
  ex.head === :for || throw(ArgumentError("Must be a matmul for loop."))
  mincrements = Expr[:(c = increment1(C, c, mᵣW)), :(ãₚ = incrementp(Ãₚ, ãₚ)), :(m = vsub_nsw(m, mᵣW))]
  # massumes = Expr[:(assume(m < mᵣW)),
  #                 :(assume(VectorizationBase.vgt(ãₚ, VectorizationBase.increment_ptr($(esc(:Ãₚ)), ãₚ, (vsub_nsw($(esc(:M)), mᵣW), LoopVectorization.Zero())), $(esc(:Ãₚ))))),
  #                 :(assume(VectorizationBase.vgt(c, VectorizationBase.increment_ptr($(esc(:C)), c, (vsub_nsw($(esc(:M)), mᵣW), LoopVectorization.Zero())), $(esc(:C)))))]
  offsetprecalc = GlobalRef(VectorizationBase,:offsetprecalc)
  preheader = quote
    mᵣ, nᵣ = matmul_params(Val($(esc(:T))))
    mᵣW = pick_vector_width($(esc(:T))) * mᵣ
    m = $(esc(:M)) % Int32
    n = $(esc(:N)) % Int32
    Ãₚ = $(esc(:Ãₚ))
    B = $offsetprecalc($(esc(:B)), Val{(9,9)}())
    C = $offsetprecalc($(esc(:C)), Val{(9,9)}())
    b = pointer(B); c = pointer(C); ãₚ = pointer(Ãₚ)
  end
  if pack
    push!(mincrements, :(a = increment1(A, a, mᵣW)))
    push!(preheader.args, :(A = $(esc(:A))), :(a = pointer(A)))
    areconstruct = Expr[:($(esc(:A)) = VectorizationBase.reconstruct_ptr(A, a))]
    # push!(massumes, :(assume(VectorizationBase.vgt(a, VectorizationBase.increment_ptr($(esc(:A)), a, (vsub_nsw($(esc(:M)), mᵣW), LoopVectorization.Zero())), $(esc(:A))))))
  else
    Ainit = areconstruct = Expr[]
  end
  lvkern = esc(:(@turbo inline=true $ex))

  loopnest = quote
    let ãₚ = ãₚ, c = c, $(esc(:B)) = VectorizationBase.reconstruct_ptr(B, b), m = m
      while m ≥ mᵣW#VectorizationBase.vle(a, amax, A)
        let $(esc(:M)) = mᵣW, $(esc(:N)) = nᵣ, $(esc(:Ãₚ)) = VectorizationBase.reconstruct_ptr(droplastdim(Ãₚ), ãₚ), $(esc(:C)) = VectorizationBase.reconstruct_ptr(C, c), $(areconstruct...)
          $lvkern
          $(mincrements...)
        end
      end
      if m > zero(Int32)#vne(a, amax, A)
        let $(esc(:M)) = UpperBoundedInteger((m%UInt)%Int, mᵣW - One()), $(esc(:N)) = nᵣ, $(esc(:Ãₚ)) = VectorizationBase.reconstruct_ptr(droplastdim(Ãₚ), ãₚ), $(esc(:C)) = VectorizationBase.reconstruct_ptr(C, c), $(areconstruct...)
          # $(massumes...)
          $lvkern
        end
      end
    end
  end

  if !pack
    loopnest = quote
      while n ≥ nᵣ#VectorizationBase.vle(c, cmax, C)
        $loopnest
        c = increment2(C, c, nᵣ)
        b = increment2(B, b, nᵣ)
        n = vsub_nsw(n, nᵣ)
      end
      if n > zero(Int32)#vne(c, cmax, C)
        let $(esc(:B)) = VectorizationBase.reconstruct_ptr(B, b), m = m
          while m ≥ mᵣW#VectorizationBase.vle(a, amax, A)
            let $(esc(:M)) = mᵣW, $(esc(:N)) = UpperBoundedInteger((n%UInt)%Int, nᵣ - One()), $(esc(:Ãₚ)) = VectorizationBase.reconstruct_ptr(droplastdim(Ãₚ), ãₚ), $(esc(:C)) = VectorizationBase.reconstruct_ptr(C, c), $(areconstruct...)
              # assume(n < nᵣ)
              # assume((VectorizationBase.vgt(c, VectorizationBase.increment_ptr($(esc(:C)), c, (Zero(), vsub_nsw($(esc(:N)), nᵣ))), $(esc(:C)))))
              # assume((VectorizationBase.vgt(b, VectorizationBase.increment_ptr($(esc(:B)), b, (Zero(), vsub_nsw($(esc(:N)), nᵣ))), $(esc(:B)))))
              $lvkern
              $(mincrements...)
            end
          end
          if m > zero(Int32)#vne(a, amax, A)
            let $(esc(:M)) = UpperBoundedInteger((m%UInt)%Int, mᵣW - One()), $(esc(:N)) = UpperBoundedInteger((n%UInt)%Int, nᵣ - One()), $(esc(:Ãₚ)) = VectorizationBase.reconstruct_ptr(droplastdim(Ãₚ), ãₚ), $(esc(:C)) = VectorizationBase.reconstruct_ptr(C, c), $(areconstruct...)
              # $(massumes...)
              # assume(n < nᵣ)
              # assume((VectorizationBase.vgt(c, VectorizationBase.increment_ptr($(esc(:C)), c, (Zero(), vsub_nsw($(esc(:N)), nᵣ))), $(esc(:C)))))
              # assume((VectorizationBase.vgt(b, VectorizationBase.increment_ptr($(esc(:B)), b, (Zero(), vsub_nsw($(esc(:N)), nᵣ))), $(esc(:B)))))
              $lvkern
            end
          end
        end
      end
    end
  end
  Expr(:block, preheader, loopnest)
end
@inline function loopmul!(C, A, B, α, β, M, K, N)
    @turbo for n ∈ CloseOpen(N), m ∈ CloseOpen(M)
        Cₘₙ = zero(eltype(C))
        for k ∈ CloseOpen(K)
            Cₘₙ += A[m,k] * B[k,n]
        end
        C[m,n] = α * Cₘₙ + β * C[m,n]
    end
    nothing
end
@inline function ploopmul!(C::AbstractStridedPointer{T}, Ãₚ, B, α, β, M, K, N) where {T}
  @kernel false for n ∈ CloseOpen(N), m ∈ CloseOpen(M)
    Cₘₙ = zero(eltype(C))
    for k ∈ CloseOpen(K)
      Cₘₙ += Ãₚ[m,k] * B[k,n]
    end
    C[m,n] = α * Cₘₙ + β * C[m,n]
  end
  nothing
end
@inline function packamul!(
    C::AbstractStridedPointer{T}, Ãₚ, A, B,
    α, β, M, K, N
  ) where {T}
  @kernel true for n ∈ CloseOpen(N), m ∈ CloseOpen(M)
    Cₘₙ = zero(eltype(C))
    for k ∈ CloseOpen(K)
      Aₘₖ = A[m,k]
      Cₘₙ += Aₘₖ * B[k,n]
      Ãₚ[m,k] = Aₘₖ
    end
    C[m,n] = α * Cₘₙ + β * C[m,n]
  end
end
@inline function alloc_a_pack(K, ::Val{T}) where {T}
  buffer = first_cache_buffer(Val(T))
  alloc_a_pack(K, buffer), buffer
end
@inline alloc_a_pack(K, buffer::MemoryBuffer) = alloc_a_pack(K, align(pointer(buffer)))
@inline function alloc_a_pack(K, bufferptr::Ptr{T}) where {T}
  mᵣ, nᵣ = matmul_params(Val(T))
  mᵣW = mᵣ * pick_vector_width(T)
  Apack = default_zerobased_stridedpointer(bufferptr, (One(), mᵣW, mᵣW * K)) # mᵣW x K x cld(M, mᵣW)
  Apack
end
function packaloopmul!(
    C::AbstractStridedPointer{T},
    A::AbstractStridedPointer,
    B::AbstractStridedPointer,
    α, β, M, K, N
) where {T}
  Ãₚ, buffer = alloc_a_pack(K, Val(T))
  GC.@preserve buffer begin
    Mᵣ, Nᵣ = matmul_params(Val(T))
    packamul!(C, Ãₚ, A, B, α, β, M, K, Nᵣ)
    ploopmul!(gesp(C, (Zero(), Nᵣ)), Ãₚ, gesp(B, (Zero(), Nᵣ)), α, β, M, K, N - Nᵣ)
  end
  nothing
end

@inline function inlineloopmul!(C, A, B, α, β, M, K, N)
    @turbo inline=true for m ∈ CloseOpen(M), n ∈ CloseOpen(N)
        Cₘₙ = zero(eltype(C))
        for k ∈ CloseOpen(K)
            Cₘₙ += A[m,k] * B[k,n]
        end
        C[m,n]  = α * Cₘₙ + β * C[m,n]
    end
    C
end


#3D primitives
typealias Primitives3D  Union{AbstractGeometry{3}, AbstractMesh}
#2D primitives AKA sprites, since they are shapes mapped onto a 2D rectangle
typealias Sprites Union{AbstractGeometry{2}, Shape, Char, Type}
typealias AllPrimitives Union{AbstractGeometry, Shape, Char}

function assemble(data, shaders, primitive, instanced=nothing)
    if instanced == nothing
    end

end

function convert2gl!(buffer_preference, dict)
    view = Dict{String, String}()
    map!(dict) do key_value
        k,v = key_value
        new_val = if isa(v, AbstractArray)
            gl_convert(buffer_preference, v)
        elseif isa(v, AbstractMesh)
            NativeMesh(v)
        else
            gl_convert(v)
        end
        return k => new_val
    end

end