text:
export issplit, multiplication_table, restrict_scalars, center

################################################################################
#
#  Basic field access
#
################################################################################

base_ring(A::AlgAss{T}) where {T} = A.base_ring::parent_type(T)

has_one(A::AlgAss) = A.has_one

iszero(A::AlgAss) = A.iszero

function Generic.dim(A::AlgAss)
  if iszero(A)
    return 0
  end
  return size(multiplication_table(A, copy = false), 1)
end

degree(A::AlgAss) = dim(A)

elem_type(::Type{AlgAss{T}}) where {T} = AlgAssElem{T, AlgAss{T}}

order_type(::AlgAss{fmpq}) = AlgAssAbsOrd{AlgAss{fmpq}, elem_type(AlgAss{fmpq})}
order_type(::Type{AlgAss{fmpq}}) = AlgAssAbsOrd{AlgAss{fmpq}, elem_type(AlgAss{fmpq})}

order_type(::AlgAss{T}) where { T <: NumFieldElem } = AlgAssRelOrd{T, fractional_ideal_type(order_type(parent_type(T)))}
order_type(::Type{AlgAss{T}}) where { T <: NumFieldElem } = AlgAssRelOrd{T, fractional_ideal_type(order_type(parent_type(T)))}

@doc Markdown.doc"""
    multiplication_table(A::AlgAss; copy::Bool = true) -> Array{RingElem, 3}

Given an algebra $A$ this function returns the multiplication table of $A$:
If the function returns $M$ and the basis of $A$ is $e_1,\dots, e_n$ then
it holds $e_i \cdot e_j = \sum_k M[i, j, k] \cdot e_k$.
"""
function multiplication_table(A::AlgAss; copy::Bool = true)
  @assert !iszero(A)
  if copy
    return deepcopy(A.mult_table)
  else
    return A.mult_table
  end
end

################################################################################
#
#  Commutativity
#
################################################################################

iscommutative_known(A::AlgAss) = (A.iscommutative != 0)

@doc Markdown.doc"""
    iscommutative(A::AlgAss) -> Bool

Returns `true` if $A$ is a commutative ring and `false` otherwise.
"""
function iscommutative(A::AlgAss)
  if iscommutative_known(A)
    return A.iscommutative == 1
  end
  for i = 1:dim(A)
    for j = i + 1:dim(A)
      if multiplication_table(A, copy = false)[i, j, :] != multiplication_table(A, copy = false)[j, i, :]
        A.iscommutative = 2
        return false
      end
    end
  end
  A.iscommutative = 1
  return true
end

################################################################################
#
#  Construction
#
################################################################################

# This only works if base_ring(A) is a field (probably)
# Returns (true, one) if there is a one and (false, something) if not.
function find_one(A::AlgAss)
  if iszero(A)
    return true, elem_type(base_ring(A))[]
  end
  n = dim(A)
  M = zero_matrix(base_ring(A), n^2, n)
  c = zero_matrix(base_ring(A), n^2, 1)
  for k = 1:n
    kn = (k - 1)*n
    c[kn + k, 1] = base_ring(A)(1)
    for i = 1:n
      for j = 1:n
        M[i + kn, j] = deepcopy(multiplication_table(A, copy = false)[j, k, i])
      end
    end
  end
  Mc = hcat(M, c)
  rref!(Mc)
  if iszero(Mc[n, n])
    return false, zeros(base_ring(A), n)
  end
  if n != 1 && !iszero(Mc[n + 1, n + 1])
    return false, zeros(base_ring(A), n)
  end
  cc = solve_ut(sub(Mc, 1:n, 1:n), sub(Mc, 1:n, (n + 1):(n + 1)))
  one = elem_type(base_ring(A))[ cc[i, 1] for i = 1:n ]
  return true, one
end

function _zero_algebra(R::Ring)
  A = AlgAss{elem_type(R)}(R)
  A.iszero = true
  A.iscommutative = 1
  A.has_one = true
  A.one = elem_type(R)[]
  return A
end

function AlgAss(R::Ring, mult_table::Array{T, 3}, one::Array{T, 1}) where {T}
  if size(mult_table, 1) == 0
    return _zero_algebra(R)
  end
  return AlgAss{T}(R, mult_table, one)
end

function AlgAss(R::Ring, mult_table::Array{T, 3}) where {T}
  if size(mult_table, 1) == 0
    return _zero_algebra(R)
  end
  A = AlgAss{T}(R)
  A.mult_table = mult_table
  A.iszero = false
  has_one, one = find_one(A)
  A.has_one = has_one
  if has_one
    A.one = one
  end
  return A
end

function AlgAss(R::Ring, d::Int, arr::Array{T, 1}) where {T}
  if d == 0
    return _zero_algebra(R)
  end
  mult_table = Array{T, 3}(undef, d, d, d)
  n = d^2
  for i in 1:d
    for j in 1:d
      for k in 1:d
        mult_table[i, j, k] = arr[(i - 1) * n + (j - 1) * d + k]
      end
    end
  end
  return AlgAss(R, mult_table)
end

# Constructs the algebra R[X]/f
function AlgAss(f::PolyElem)
  R = base_ring(parent(f))
  n = degree(f)
  Rx = parent(f)
  x = gen(Rx)
  B = Array{elem_type(Rx), 1}(undef, 2*n - 1)
  B[1] = Rx(1)
  for i = 2:2*n - 1
    B[i] = mod(B[i - 1]*x, f)
  end
  mult_table = Array{elem_type(R), 3}(undef, n, n, n)
  for i = 1:n
    for j = i:n
      for k = 1:n
        mult_table[i, j, k] = coeff(B[i + j - 1], k - 1)
        mult_table[j, i, k] = coeff(B[i + j - 1], k - 1)
      end
    end
  end
  one = map(R, zeros(Int, n))
  one[1] = R(1)
  A = AlgAss(R, mult_table, one)
  A.iscommutative = 1
  return A
end

function AlgAss(K::AnticNumberField)
  A = AlgAss(K.pol)
  m = AbsAlgAssToNfAbsMor(A, K, identity_matrix(FlintQQ, dim(A)), identity_matrix(FlintQQ, dim(A)))
  A.maps_to_numberfields = [ (K, m) ]
  return A, m
end

# Reduces the rows of M in `rows` modulo N in place.
# Assumes that N is in lowerleft HNF.
function reduce_rows_mod_hnf!(M::fmpz_mat, N::fmpz_mat, rows::Vector{Int})
  for i in rows
    for j = ncols(M):-1:1
      if iszero(M[i, j])
        continue
      end

      t = fdiv(M[i, j], N[j, j])
      for k = 1:j
        M[i, k] = M[i, k] - t*N[j, k]
      end
    end
  end
  return M
end

@doc Markdown.doc"""
    quo(O::NfAbsOrd, I::NfAbsOrdIdl, p::Union{ Int, fmpz })
    quo(O::AlgAssAbsOrd, I::AlgAssAbsOrdIdl, p::Union{ Int, fmpz })
      -> AlgAss, AbsOrdToAlgAssMor

Given an ideal $I$ such that $p \cdot O \subseteq I \subseteq O$ this function
constructs $O/I$ as an algebra over $\mathbb F_p$ together with the projection
map $O \to O/I$.
It is assumed that $p$ is prime.
"""
quo(O::Union{NfAbsOrd, AlgAssAbsOrd}, I::Union{NfAbsOrdIdl, AlgAssAbsOrdIdl}, p::Union{Integer, fmpz}) = AlgAss(O, I, p)

function AlgAss(O::Union{NfAbsOrd, AlgAssAbsOrd}, I::Union{NfAbsOrdIdl, AlgAssAbsOrdIdl}, p::Union{Integer, fmpz})
  @assert order(I) === O

  n = degree(O)
  bmatI = integral_basis_matrix_wrt(I, O, copy = false)

  basis_elts = Vector{Int}()
  for i = 1:n
    if iscoprime(bmatI[i, i], p)
      continue
    end

    push!(basis_elts, i)
  end

  r = length(basis_elts)
  Fp = GF(p, cached = false)

  if r == 0
    A = _zero_algebra(Fp)

    local _image_zero

    let A = A
      function _image_zero(a::Union{ NfAbsOrdElem, AlgAssAbsOrdElem })
        return A()
      end
    end

    local _preimage_zero

    let O = O
      function _preimage_zero(a::AlgAssElem)
        return O()
      end
    end

    OtoA = AbsOrdToAlgAssMor{typeof(O), elem_type(Fp)}(O, A, _image_zero, _preimage_zero)
    return A, OtoA
  end

  BO = basis(O, copy = false)
  mult_table = Array{elem_type(Fp), 3}(undef, r, r, r)
  for i = 1:r
    M = representation_matrix_mod(BO[basis_elts[i]], fmpz(p))
    if r != degree(O)
      M = reduce_rows_mod_hnf!(M, bmatI, basis_elts)
    end
    for j = 1:r
      for k = 1:r
        mult_table[i, j, k] = Fp(M[basis_elts[j], basis_elts[k]])
      end
    end
  end

  if isone(BO[1])
    one = zeros(Fp, r)
    one[1] = Fp(1)
    A = AlgAss(Fp, mult_table, one)
  else
    A = AlgAss(Fp, mult_table)
  end
  if iscommutative(O)
    A.iscommutative = 1
  end

  local _image

  let I = I, A = A, basis_elts = basis_elts, Fp = Fp
    function _image(a::Union{NfAbsOrdElem, AlgAssAbsOrdElem})
      c = coordinates(mod(a, I), copy = false)
      return A([ Fp(c[i]) for i in basis_elts ])
    end
  end

  local _preimage

  let BO = BO, basis_elts = basis_elts, r = r
    function _preimage(a::AlgAssElem)
      return sum(lift(coeffs(a, copy = false)[i])*BO[basis_elts[i]] for i = 1:r)
    end
  end

  OtoA = AbsOrdToAlgAssMor{typeof(O), elem_type(Fp)}(O, A, _image, _preimage)

  return A, OtoA
end

# Requires M to be in lower left HNF
function reduce_vector_mod_hnf(v::fmpz_mat, M::fmpz_mat)
  @assert ncols(v) == nrows(M) && nrows(M) == ncols(M)

  w = Vector{fmpz}(undef, length(v))
  t = fmpz()
  for i in length(v):-1:1
    t = fdiv(v[1, i], M[i, i])
    for j in 1:i
      w[j] = v[1, j] - t*M[i, j]
    end
  end
  return w
end

@doc Markdown.doc"""
    quo(I::NfAbsOrdIdl, J::NfAbsOrdIdl, p::Union{ Int, fmpz })
    quo(I::AlgAssAbsOrdIdl, J::AlgAssAbsOrdIdl, p::Union{ Int, fmpz })
      -> AlgAss, AbsOrdToAlgAssMor

Given an ideal $J$ such that $p \cdot I \subseteq J \subseteq I$ this function
constructs $I/J$ as an algebra over $\mathbb F_p$ together with the projection
map $I \to I/J$.
It is assumed that $p$ is prime.
"""
quo(I::Union{ NfAbsOrdIdl, AlgAssAbsOrdIdl }, J::Union{ NfAbsOrdIdl, AlgAssAbsOrdIdl }, p::Union{ Integer, fmpz }) = AlgAss(I, J, p)

function AlgAss(I::Union{ NfAbsOrdIdl, AlgAssAbsOrdIdl }, J::Union{NfAbsOrdIdl, AlgAssAbsOrdIdl}, p::Union{Integer, fmpz})
  @assert order(I) === order(J)

  O = order(I)

  n = degree(O)
  BmatJinI = hnf(basis_matrix(J, copy = false)*basis_mat_inv(I, copy = false), :lowerleft)
  @assert isone(BmatJinI.den) "J is not a subset of I"
  BmatJinI = BmatJinI.num
  basis_elts = Vector{Int}()
  for i = 1:n
    if valuation(BmatJinI[i, i], p) == 0
      continue
    end

    push!(basis_elts, i)
  end

  r = length(basis_elts)
  Fp = GF(p, cached = false)

  if r == 0
    A = _zero_algebra(Fp)

    local _image_zero

    let A = A
      function _image_zero(a::Union{ NfAbsOrdElem, AlgAssAbsOrdElem })
        return A()
      end
    end

    local _preimage_zero

    let O = O
      function _preimage_zero(a::AlgAssElem)
        return O()
      end
    end

    OtoA = AbsOrdToAlgAssMor{typeof(O), elem_type(Fp)}(O, A, _image_zero, _preimage_zero)
    return A, OtoA
  end

  BI = basis(I, copy = false)
  BmatI = basis_matrix(I, copy = false)
  BmatIinv = inv(BmatI)

  mult_table = Array{elem_type(Fp), 3}(undef, r, r, r)
  for i = 1:r
    M = FakeFmpqMat(representation_matrix(_elem_in_algebra(BI[basis_elts[i]], copy = false)))
    M = mul!(M, BmatI, M)
    M = mul!(M, M, BmatIinv)
    @assert M.den == 1
    M = M.num # M is now the representation matrix in the basis of I
    if r != degree(O)
      M = reduce_rows_mod_hnf!(M, BmatJinI, basis_elts)
    end
    for j = 1:r
      for k = 1:r
        mult_table[i, j, k] = Fp(M[basis_elts[j], basis_elts[k]])
      end
    end
  end

  A = AlgAss(Fp, mult_table)
  if iscommutative(O)
    A.iscommutative = 1
  end

  t = FakeFmpqMat(zero_matrix(FlintZZ, 1, n))

  local _image

  let BmatJinI = BmatJinI, I = I, r = r, A = A, t = t, Fp = Fp
    function _image(a::Union{NfAbsOrdElem, AlgAssAbsOrdElem})
      elem_to_mat_row!(t.num, 1, t.den, _elem_in_algebra(a, copy = false))
      t = mul!(t, t, basis_mat_inv(I, copy = false))
      @assert isone(t.den) "Not an element of the domain"
      c = reduce_vector_mod_hnf(t.num, BmatJinI)
      return A([ Fp(c[i]) for i in basis_elts ])
    end
  end

  local _preimage

  let BI = BI, basis_elts = basis_elts, r = r
    function _preimage(a::AlgAssElem)
      return O(sum(lift(coeffs(a, copy = false)[i])*BI[basis_elts[i]] for i = 1:r))
    end
  end

  OtoA = AbsOrdToAlgAssMor{typeof(O), elem_type(Fp)}(O, A, _image, _preimage)

  return A, OtoA
end

#=
Qx, x = QQ["x"];
f = x^2 + 12*x - 92;
K, a = number_field(f, "a");
OK = maximal_order(K);
Ky, y = K["y"];
g = y^2 - 54*y - 73;
L, b = number_field(g, "b");
OL = maximal_order(L);
p = prime_decomposition(OK, 2)[1][1]
=#

# Assume that O is relative order over OK, I is an ideal of O and p is a prime
# ideal of OK with pO \subseteq I. O/I is an OK/p-algebra.
#
# The idea is to compute pseudo-basis of O and I respectively, for which the
# coefficient ideals have zero p-adic valuation. Then we can think in the
# localization at p and do as in the case of principal ideal domains.
@doc Markdown.doc"""
    quo(O::NfRelOrd, I::NfRelOrdIdl, p::Union{ NfAbsOrdIdl, NfRelOrdIdl })
    quo(O::AlgAssRelOrd, I::AlgAssRelOrdIdl, p::Union{ NfAbsOrdIdl, NfRelOrdIdl })
      -> AlgAss, RelOrdToAlgAssMor

Given an ideal $I$ such that $p \cdot O \subseteq I \subseteq O$ this function
constructs $O/I$ as an algebra over the finite field $R/p$, where $R$ is the
order of $p$, together with the projection map $O \to O/I$.
It is assumed that `R == base_ring(O)` and that $p$ is prime.
"""
quo(O::Union{ NfRelOrd{T, S}, AlgAssRelOrd{T, S} }, I::Union{ NfRelOrdIdl{T, S}, AlgAssRelOrdIdl{T, S} }, p::Union{NfOrdIdl, NfRelOrdIdl}) where {T, S} = AlgAss(O, I, p)

function AlgAss(O::Union{ NfRelOrd{T, S}, AlgAssRelOrd{T, S} }, I::Union{ NfRelOrdIdl{T, S}, AlgAssRelOrdIdl{T, S} }, p::Union{NfOrdIdl, NfRelOrdIdl}) where {T, S}

  K = _algebra(O)

  new_basisO, new_basisI, new_bmatO, new_bmatI = coprime_bases(O, I, p)
  new_bmatinvO = inv(new_bmatO)

  Fp, mF = ResidueField(order(p), p)
  mmF = extend(mF, _base_ring(K))
  invmmF = pseudo_inv(mmF)

  basis_elts = Int[]
  reducers = Int[]

  for i in 1:degree(O)
    v = valuation(new_bmatI[i, i], p)
    @assert v >= 0
    if v == 0
      push!(reducers, i)
    else
      push!(basis_elts, i)
    end
  end

  r = length(basis_elts)

  if r == 0
    A = _zero_algebra(Fp)

    local _image_zero

    let A = A
      function _image_zero(a::Union{ NfRelOrdElem, AlgAssRelOrdElem })
        return A()
      end
    end

    local _preimage_zero

    let O = O
      function _preimage_zero(a::AlgAssElem)
        return O()
      end
    end

    OtoA = RelOrdToAlgAssMor(O, A, _image_zero, _preimage_zero)
    return A, OtoA
  end

  reverse!(reducers)

  tmp_matrix = zero_matrix(_base_ring(K), 1, degree(O))

  function _coeff(c)
    for i = 1:degree(O)
      tmp_matrix[1, i] = coeffs(c, copy = false)[i]
    end
    return tmp_matrix*new_bmatinvO
  end

  mult_table = Array{elem_type(Fp), 3}(undef, r, r, r)
  for i in 1:r
    for j in 1:r
      c = new_basisO[basis_elts[i]][1]*new_basisO[basis_elts[j]][1]
      coeffs_c = _coeff(c)

      for k in reducers
        d = -coeffs_c[k]//new_bmatI[k, k]
        c = c + d*new_basisI[k][1]
      end
      coeffs_c = _coeff(c)
      for k in 1:degree(O)
        if !(k in basis_elts)
          @assert iszero(coeffs_c[k])
        end
      end
      for k in 1:r
        mult_table[i, j, k] = mmF(coeffs_c[basis_elts[k]])
      end
    end
  end

  if isone(new_basisO[basis_elts[1]][1])
    one = zeros(Fp, length(basis_elts))
    one[1] = Fp(1)
    A = AlgAss(Fp, mult_table, one)
  else
    A = AlgAss(Fp, mult_table)
  end
  if iscommutative(O)
    A.iscommutative = 1
  end

  local _image

  let A = A, O = O
    function _image(a::Union{ NfRelOrdElem, AlgAssRelOrdElem })
      c = _elem_in_algebra(a, copy = false)
      coeffs_c = _coeff(c)
      for k in reducers
        d = -coeffs_c[k]//new_bmatI[k, k]
        c = c + d*new_basisI[k][1]
      end
      coeffs_c = _coeff(c)
      for k in 1:degree(O)
        if !(k in basis_elts)
          @assert iszero(coeffs_c[k])
        end
      end
      b = A()
      for k in 1:r
        b.coeffs[k] = mmF(coeffs_c[basis_elts[k]])
      end
      return b
    end
  end

  lifted_basis_of_A = []

  for i in basis_elts
    c = coprime_to(new_basisO[i][2], p)
    b = invmmF(inv(mmF(c)))*c*new_basisO[i][1]
    @assert b in O
    push!(lifted_basis_of_A, b)
  end

  local _preimage
  let lifted_basis_of_A = lifted_basis_of_A, O = O, invmmF = invmmF
    function _preimage(v::AlgAssElem)
      return O(sum((invmmF(v.coeffs[i]))*lifted_basis_of_A[i] for i in 1:r))
    end
  end

  OtoA = RelOrdToAlgAssMor(O, A, _image, _preimage)

  return A, OtoA
end

@doc Markdown.doc"""
    quo(I::NfRelOrdIdl, J::NfRelOrdIdl, p::Union{ NfAbsOrdIdl, NfRelOrdIdl })
    quo(I::AlgAssRelOrdIdl, J::AlgAssRelOrdIdl, p::Union{ NfAbsOrdIdl, NfRelOrdIdl })
      -> AlgAss, RelOrdToAlgAssMor

Given an ideal $J$ such that $p \cdot I \subseteq J \subseteq I$ this function
constructs $I/J$ as an algebra over the finite field $R/p$, where $R$ is the
order of $p$, together with the projection map $I \to I/J$.
It is assumed that `order(I) === order(J)` and in particular both should be
defined. Further, it should hold `R == base_ring(order(I))` and $p$ should be
prime.
"""
quo(I::Union{ NfRelOrdIdl{T, S}, AlgAssRelOrdIdl{T, S} }, J::Union{ NfRelOrdIdl{T, S}, AlgAssRelOrdIdl{T, S} }, p::Union{NfOrdIdl, NfRelOrdIdl}) where {T, S} = AlgAss(I, J, p)

function AlgAss(I::Union{ NfRelOrdIdl{T, S}, AlgAssRelOrdIdl{T, S} }, J::Union{ NfRelOrdIdl{T, S}, AlgAssRelOrdIdl{T, S} }, p::Union{NfOrdIdl, NfRelOrdIdl}) where {T, S}
  @assert _algebra(I) === _algebra(J)
  @assert order(I) === order(J)

  O = order(I)
  K = _algebra(I)
  new_basisI, new_basisJ, new_bmatI, new_bmatJinI = coprime_bases(I, J, p)
  bmatinvI = inv(new_bmatI)

  Fp, mF = ResidueField(order(p), p)
  mmF = extend(mF, _base_ring(K))
  invmmF = pseudo_inv(mmF)

  basis_elts = Int[]
  reducers = Int[]

  for i in 1:degree(O)
    v = valuation(new_bmatJinI[i, i], p)
    @assert v >= 0
    if v == 0
      push!(reducers, i)
    else
      push!(basis_elts, i)
    end
  end

  r = length(basis_elts)

  if r == 0
    A = _zero_algebra(Fp)

    local _image_zero

    let A = A
      function _image_zero(a::Union{ NfRelOrdElem, AlgAssRelOrdElem })
        return A()
      end
    end

    local _preimage_zero

    let O = O
      function _preimage_zero(a::AlgAssElem)
        return O()
      end
    end

    OtoA = RelOrdToAlgAssMor(O, A, _image_zero, _preimage_zero)
    return A, OtoA
  end

  reverse!(reducers)

  tmp_matrix = zero_matrix(_base_ring(K), 1, degree(O))

  function _coeff(c)
    for i = 1:degree(O)
      tmp_matrix[1, i] = coeffs(c, copy = false)[i]
    end
    return tmp_matrix*bmatinvI
  end

  mult_table = Array{elem_type(Fp), 3}(undef, r, r, r)

  for i in 1:r
    for j in 1:r
      c = new_basisI[basis_elts[i]][1]*new_basisI[basis_elts[j]][1]
      coeffs = _coeff(c)

      for k in reducers
        d = -coeffs[k]//new_bmatJinI[k, k]
        c = c + d * new_basisJ[k][1]
      end
      coeffs = _coeff(c)
      for k in 1:degree(O)
        if !(k in basis_elts)
          @assert iszero(coeffs[k])
        end
      end
      for k in 1:r
        mult_table[i, j, k] = mmF(coeffs[basis_elts[k]])
      end
    end
  end

  if isone(new_basisI[basis_elts[1]][1])
    one = zeros(Fp, length(basis_elts))
    one[1] = Fp(1)
    A = AlgAss(Fp, mult_table, one)
  else
    A = AlgAss(Fp, mult_table)
  end
  if iscommutative(O)
    A.iscommutative = 1
  end

  local _image

  let O = O, new_bmatJinI = new_bmatJinI, A = A
    function _image(a::Union{ NfRelOrdElem, AlgAssRelOrdElem })
      c = _elem_in_algebra(a, copy = false)
      coeffs = _coeff(c)
      for k in reducers
        d = -coeffs[k]//new_bmatJinI[k, k]
        c = c + d*new_basisJ[k][1]
      end
      coeffs = _coeff(c)
      for k in 1:degree(O)
        if !(k in basis_elts)
          @assert iszero(coeffs[k])
        end
      end
      b = A()
      for k in 1:r
        b.coeffs[k] = mmF(coeffs[basis_elts[k]])
      end
      return b
    end
  end

  lifted_basis_of_A = []

  for i in basis_elts
    c = coprime_to(new_basisI[i][2], p)
    b = invmmF(inv(mmF(c)))*c*new_basisI[i][1]
    @assert O(b) in I
    push!(lifted_basis_of_A, b)
  end

  local _preimage

  let O = O, invmmF = invmmF, lifted_basis_of_A = lifted_basis_of_A
    function _preimage(v::AlgAssElem)
      return O(sum((invmmF(v.coeffs[i])) * lifted_basis_of_A[i] for i in 1:r))
    end
  end

  OtoA = RelOrdToAlgAssMor(O, A, _image, _preimage)

  return A, OtoA
end

function AlgAss(A::Generic.MatAlgebra{T}) where { T <: FieldElem }
  n = A.n
  K = base_ring(A)
  n2 = n^2
  # We use the matrices M_{ij} with a 1 at row i and column j and zeros everywhere else as the basis for A.
  # We sort "column major", so A[i + (j - 1)*n] corresponds to the matrix M_{ij}.
  # M_{ik}*M_{lj} = 0, if k != l, and M_{ik}*M_{kj} = M_{ij}
  mult_table = zeros(K, n2, n2, n2)
  oneK = one(K)
  for j = 0:n:(n2 - n)
    for k = 1:n
      kn = (k - 1)*n
      for i = 1:n
        mult_table[i + kn, k + j, i + j] = oneK
      end
    end
  end
  oneA = zeros(K, n2)
  for i = 1:n
    oneA[i + (i - 1)*n] = oneK
  end
  A = AlgAss(K, mult_table, oneA)
  A.iscommutative = ( n == 1 ? 1 : 2 )
  return A
end

function AlgAss(A::AlgAss)
  R = base_ring(A)
  d = dim(A)
  return A, hom(A, A, identity_matrix(R, d), identity_matrix(R, d))
end

###############################################################################
#
#  String I/O
#
################################################################################

function show(io::IO, A::AlgAss)
  print(io, "Associative algebra of dimension ", dim(A), " over ", base_ring(A))
end

################################################################################
#
#  Deepcopy
#
################################################################################

function Base.deepcopy_internal(A::AlgAss{T}, dict::IdDict) where {T}
  B = AlgAss{T}(base_ring(A))
  for x in fieldnames(typeof(A))
    if x != :base_ring && isdefined(A, x)
      setfield!(B, x, Base.deepcopy_internal(getfield(A, x), dict))
    end
  end
  B.base_ring = A.base_ring
  return B
end

################################################################################
#
#  Subalgebra
#
################################################################################

# Builds a multiplication table for the subalgebra of A with basis matrix B.
# We assume ncols(B) == dim(A).
# A rref of B will be computed IN PLACE! If return_LU is Val{true}, a LU-factorization
# of transpose(rref(B)) is returned.
function _build_subalgebra_mult_table!(A::AlgAss{T}, B::MatElem{T}, return_LU::Type{Val{S}} = Val{false}) where { T, S }
  K = base_ring(A)
  n = dim(A)
  r = rref!(B)
  if r == 0
    if return_LU == Val{true}
      return Array{elem_type(K), 3}(undef, 0, 0, 0), SymmetricGroup(ncols(B))(), zero_matrix(K, 0, 0), zero_matrix(K, 0, 0)
    else
      return Array{elem_type(K), 3}(undef, 0, 0, 0)
    end
  end

  basis = Vector{elem_type(A)}(undef, r)
  for i = 1:r
    basis[i] = elem_from_mat_row(A, B, i)
  end

  _, p, L, U = lu(transpose(B))

  mult_table = Array{elem_type(K), 3}(undef, r, r, r)
  c = A()
  d = zero_matrix(K, n, 1)
  for i = 1:r
    for j = 1:r
      if iscommutative(A) && j < i
        continue
      end
      c = mul!(c, basis[i], basis[j])
      for k = 1:n
        d[p[k], 1] = c.coeffs[k]
      end
      d = solve_lt(L, d)
      d = solve_ut(U, d)
      for k = 1:r
        mult_table[i, j, k] = deepcopy(d[k, 1])
        if iscommutative(A) && i != j
          mult_table[j, i, k] = deepcopy(d[k, 1])
        end
      end
    end
  end

  if return_LU == Val{true}
    return mult_table, p, L, U
  else
    return mult_table
  end
end

@doc Markdown.doc"""
     subalgebra(A::AlgAss, e::AlgAssElem, idempotent::Bool = false,
                action::Symbol = :left)
       -> AlgAss, AbsAlgAssMor

Given an algebra $A$ and an element $e$, this function constructs the algebra
$e \cdot A$ (if `action == :left`) respectively $A \cdot e$ (if `action == :right`)
and a map from this algebra to $A$.
If `idempotent` is `true`, it is assumed that $e$ is idempotent in $A$.
"""
function subalgebra(A::AlgAss{T}, e::AlgAssElem{T, AlgAss{T}}, idempotent::Bool = false, action::Symbol = :left) where {T}
  @assert parent(e) == A
  R = base_ring(A)
  n = dim(A)
  B = representation_matrix(e, action)

  mult_table, p, L, U = _build_subalgebra_mult_table!(A, B, Val{true})
  r = size(mult_table, 1)

  if r == 0
    eA = _zero_algebra(R)
    return eA, hom(eA, A, zero_matrix(R, 0, n))
  end

  # The basis matrix of e*A resp. A*e with respect to A is
  basis_mat_of_eA = sub(B, 1:r, 1:n)

  if idempotent
    c = A()
    d = zero_matrix(R, n, 1)
    for k = 1:n
      d[p[k], 1] = e.coeffs[k]
    end
    d = solve_lt(L, d)
    d = solve_ut(U, d)
    v = Vector{elem_type(R)}(undef, r)
    for i in 1:r
      v[i] = d[i, 1]
    end
    eA = AlgAss(R, mult_table, v)
  else
    eA = AlgAss(R, mult_table)
  end

  if A.iscommutative == 1
    eA.iscommutative = 1
  end

  if idempotent
    # We have the map eA -> A, given by the multiplying with basis_mat_of_eA.
    # But there is also the canonical projection A -> eA, a -> ea.
    # We compute the corresponding matrix.
    B = representation_matrix(e, action)
    C = zero_matrix(R, n, r)
    for i in 1:n
      for k = 1:n
        d[p[k], 1] = B[i, k]
      end
      d = solve_lt(L, d)
      d = solve_ut(U, d)
      for k in 1:r
        C[i, k] = d[k, 1]
      end
    end
    eAtoA = hom(eA, A, basis_mat_of_eA, C)
  else
    eAtoA = hom(eA, A, basis_mat_of_eA)
  end
  return eA, eAtoA
end

@doc Markdown.doc"""
    subalgebra(A::AlgAss, basis::Vector{AlgAssElem}) -> AlgAss, AbsAlgAssMor

Returns the subalgebra of $A$ generated by the elements in `basis` and a map
from this algebra to $A$.
"""
function subalgebra(A::AlgAss{T}, basis::Vector{AlgAssElem{T, AlgAss{T}}}) where T
  M = zero_matrix(base_ring(A), dim(A), dim(A))
  for i = 1:length(basis)
    elem_to_mat_row!(M, i, basis[i])
  end
  mt = _build_subalgebra_mult_table!(A, M)
  B = AlgAss(base_ring(A), mt)
  return B, hom(B, A, sub(M, 1:length(basis), 1:dim(A)))
end

###############################################################################
#
#  Trace Matrix
#
###############################################################################

function _assure_trace_basis(A::AlgAss{T}) where T
  if !isdefined(A, :trace_basis_elem)
    A.trace_basis_elem = Array{T, 1}(undef, dim(A))
    for i=1:length(A.trace_basis_elem)
      A.trace_basis_elem[i]=sum(multiplication_table(A, copy = false)[i,j,j] for j= 1:dim(A))
    end
  end
  return nothing
end

@doc Markdown.doc"""
    trace_matrix(A::AlgAss) -> MatElem

Returns a matrix $M$ over the base ring of $A$ such that
$M_{i, j} = \mathrm{tr}(b_i \cdot b_j)$, where $b_1, \dots, b_n$ is the
basis of $A$.
"""
function trace_matrix(A::AlgAss)
  _assure_trace_basis(A)
  F = base_ring(A)
  n = dim(A)
  M = zero_matrix(F, n, n)
  for i = 1:n
    M[i,i] = tr(A[i]^2)
  end
  for i = 1:n
    for j = i+1:n
      x = tr(A[i]*A[j])
      M[i,j] = x
      M[j,i] = x
    end
  end
  return M
end

###############################################################################
#
#  Center
#
###############################################################################

function _rep_for_center(M::T, A::AlgAss) where T<: MatElem
  n=dim(A)
  for i=1:n
    for j = 1:n
      for k = 1:n
        M[k+(i-1)*n, j] = multiplication_table(A, copy = false)[i, j, k]-multiplication_table(A, copy = false)[j, i, k]
      end
    end
  end
  return nothing
end

@doc Markdown.doc"""
    center(A::AlgAss) -> AlgAss, AbsAlgAssMor

Returns the center $C$ of $A$ and the inclusion $C \to A$.
"""
function center(A::AlgAss{T}) where {T}
  if iscommutative(A)
    B, mB = AlgAss(A)
    return B, mB
  end
  if isdefined(A, :center)
    return A.center::Tuple{AlgAss{T}, morphism_type(AlgAss{T}, AlgAss{T})}
  end
  n=dim(A)
  M=zero_matrix(base_ring(A), n^2, n)
  # I concatenate the difference between the right and left representation matrices.
  _rep_for_center(M,A)
  k,B=nullspace(M)
  res=Array{elem_type(A),1}(undef, k)
  for i=1:k
    res[i]= A(T[B[j,i] for j=1:n])
  end
  C, mC = subalgebra(A, res)
  A.center = C, mC
  return C, mC
end

################################################################################
#
#  Idempotents
#
################################################################################

# See W. Eberly "Computations for Algebras and Group Representations" p. 126.
# TODO: fix the type
function _find_non_trivial_idempotent(A::AlgAss{T}) where { T } #<: Union{gfp_elem, Generic.ResF{fmpz}, fq, fq_nmod} }
  if dim(A) == 1
    error("Dimension of algebra is 1")
  end
  while true
    a = rand(A)
    if isone(a) || iszero(a)
      continue
    end
    mina = minpoly(a)
    if isirreducible(mina)
      if degree(mina) == dim(A)
        error("Algebra is a field")
      end
      continue
    end
    if issquarefree(mina)
      e = _find_idempotent_via_squarefree_poly(A, a, mina)
    else
      e = _find_idempotent_via_non_squarefree_poly(A, a, mina)
    end
    if isone(e) || iszero(e)
      continue
    end
    return e
  end
end

#function _find_idempotent_via_non_squarefree_poly(A::AlgAss{T}, a::AlgAssElem{T}, mina::Union{gfp_poly, gfp_fmpz_poly, fq_poly, fq_nmod_poly}) where { T <: Union{gfp_elem, Generic.ResF{fmpz}, fq, fq_nmod} }
function _find_idempotent_via_non_squarefree_poly(A::AlgAss{T}, a::AlgAssElem{T}, mina) where {T}
  fac = factor(mina)
  if length(fac) == 1
    return zero(A)
  end
  sf_part = one(parent(mina))
  for (k, v) in fac
    mul!(sf_part, sf_part, k)
  end
  b = sf_part(a)
  # This is not really an algebra, only a right sided ideal
  bA, bAtoA = subalgebra(A, b, false, :left)

  # Find an element e of bA such that e*x == x for all x in bA
  M = zero_matrix(base_ring(A), dim(bA), 0)
  for i = 1:dim(bA)
    M = hcat(M, representation_matrix(bA[i], :right))
  end

  N = zero_matrix(base_ring(A), 0, 1)
  for i = 1:dim(bA)
    N = vcat(N, matrix(base_ring(A), dim(bA), 1, coeffs(bA[i])))
  end
  MN = hcat(transpose(M), N)
  r = rref!(MN)
  be = solve_ut(sub(MN, 1:r, 1:dim(bA)), sub(MN, 1:r, (dim(bA) + 1):(dim(bA) + 1)))
  e = bAtoA(bA([ be[i, 1] for i = 1:dim(bA) ]))
  return e
end

# A should be semi-simple
# See W. Eberly "Computations for Algebras and Group Representations" p. 89.
function _extraction_of_idempotents(A::AlgAss, only_one::Bool = false)
  Z, ZtoA = center(A)
  if dim(Z) == 1
    error("Dimension of centre is 1")
  end

  a = ZtoA(rand(Z))
  f = minpoly(a)
  while isirreducible(f)
    if degree(f) == dim(A)
      error("Cannot find idempotents (algebra is a field)")
    end
    a = ZtoA(rand(Z))
    f = minpoly(a)
  end

  fac = factor(f)
  fi = [ k for k in keys(fac.fac) ]
  l = length(fi)
  R = parent(f)
  if only_one
    r = zeros(R, l)
    r[1] = one(R)
    g = crt(r, fi)
    return g(a)
  else
    oneR = one(R)
    zeroR = zero(R)
    gi = Vector{elem_type(R)}(undef, l)
    r = zeros(R, l)
    for i = 1:l
      r[i] = oneR
      gi[i] = crt(r, fi)
      r[i] = zeroR
    end
    return [ g(a) for g in gi ]
  end
end

#function _find_idempotent_via_squarefree_poly(A::AlgAss{T}, a::AlgAssElem{T}, mina::Union{gfp_poly, gfp_fmpz_poly, fq_poly, fq_nmod_poly}) where { T <: Union{gfp_elem, Generic.ResF{fmpz}, fq, fq_nmod} }
# TODO: fix the type
function _find_idempotent_via_squarefree_poly(A::AlgAss{T}, a::AlgAssElem{T}, mina) where {T}
  B = AlgAss(mina)
  idemB = _extraction_of_idempotents(B, true)

  e = dot(coeffs(idemB, copy = false), [ a^k for k = 0:(degree(mina) - 1) ])
  return e
end

# TODO: fix the type
function _primitive_idempotents(A::AlgAss{T}) where { T } #<: Union{gfp_elem, Generic.ResF{fmpz}, fq, fq_nmod} }
  if dim(A) == 1
    return [ one(A) ]
  end

  e = _find_non_trivial_idempotent(A)

  idempotents = Vector{elem_type(A)}()

  eA, m1 = subalgebra(A, e, true, :left)
  eAe, m2 = subalgebra(eA, m1\e, true, :right)
  if dim(eAe) == dim(A)
    push!(idempotents, e)
  else
    idems = _primitive_idempotents(eAe)
    append!(idempotents, [ m1(m2(idem)) for idem in idems ])
  end

  f = (1 - e)
  fA, n1 = subalgebra(A, f, true, :left)
  fAf, n2 = subalgebra(fA, n1\f, true, :right)

  if dim(fAf) == dim(A)
    push!(idempotents, f)
  else
    idems = _primitive_idempotents(fAf)
    append!(idempotents, [ n1(n2(idem)) for idem in idems ])
  end

  return idempotents
end

################################################################################
#
#  Matrix Algebra
#
################################################################################

# This computes a "matrix type" basis for A.
# See W. Eberly "Computations for Algebras and Group Representations" p. 121.
# TODO: fix the type
function _matrix_basis(A::AlgAss{T}, idempotents::Vector{S}) where { T, S }#<: Union{gfp_elem, Generic.ResF{fmpz}, fq, fq_nmod}, S <: AlgAssElem{T, AlgAss{T}} }
  k = length(idempotents)
  # Compute a basis e_ij of A (1 <= i, j <= k) with
  # e_11 + e_22 + ... + e_kk = 1 and e_rs*e_tu = \delta_st*e_ru.
  new_basis = Vector{elem_type(A)}(undef, k^2) # saved column major: new_basis[i + (j - 1)*k] = e_ij
  for i = 1:k
    new_basis[i + (i - 1)*k] = idempotents[i]
  end

  a = idempotents[1]
  for i = 2:k
    b = idempotents[i]
    e = a + b
    eA, m1 = subalgebra(A, e, true, :left)
    eAe, m2 = subalgebra(eA, m1\e, true, :right)

    aa = m2\(m1\(a))
    bb = m2\(m1\(b))

    # We compute an element x of eAe which fulfils
    # aa*x == x, bb*x == 0, x*aa == 0 and x*bb == x.
    M1 = representation_matrix(aa - one(eAe), :left)
    M2 = representation_matrix(bb, :left)
    M3 = representation_matrix(aa, :right)
    M4 = representation_matrix(bb - one(eAe), :right)

    M = hcat(M1, M2, M3, M4)
    xx = eAe(left_kernel_basis(M)[1])
    x = m1(m2(xx))

    new_basis[1 + (i - 1)*k] = x # this is e_1i

    # We compute an element y of eAe which fulfils
    # aa*y == 0, bb*y == y, y*aa == y, y*bb == 0, y*xx == bb, xx*y == aa.
    N1 = representation_matrix(aa, :left)
    N2 = representation_matrix(bb - one(eAe), :left)
    N3 = representation_matrix(aa - one(eAe), :right)
    N4 = representation_matrix(bb, :right)
    N5 = representation_matrix(xx, :right)
    N6 = representation_matrix(xx, :left)
    N = hcat(N1, N2, N3, N4, N5, N6)
    NN = zero_matrix(base_ring(A), 4*dim(eAe), 1)
    NN = vcat(NN, matrix(base_ring(A), dim(eAe), 1, coeffs(bb)))
    NN = vcat(NN, matrix(base_ring(A), dim(eAe), 1, coeffs(aa)))
    b, yy = can_solve_with_solution(transpose(N), NN)
    @assert b
    y = m1(m2(eAe([ yy[i, 1] for i = 1:dim(eAe) ])))

    new_basis[i] = y # this is e_i1
  end

  for j = 2:k
    jk = (j - 1)*k
    e1j = new_basis[1 + jk]
    for i = 2:k
      new_basis[i + jk] = new_basis[i]*e1j # this is e_ij
    end
  end
  return new_basis
end

# Assumes that A is central and isomorphic to a matrix algebra of base_ring(A)
# TODO: fix the type
function _as_matrix_algebra(A::AlgAss{T}) where { T } # <: Union{gfp_elem, Generic.ResF{fmpz}, fq, fq_nmod}, S <: AlgAssElem{T, AlgAss{T}} }

  idempotents = _primitive_idempotents(A)
  @assert length(idempotents)^2 == dim(A)
  Fq = base_ring(A)

  B = matrix_algebra(Fq, length(idempotents))

  matrix_basis = _matrix_basis(A, idempotents)

  # matrix_basis is another basis for A. We build the matrix for the basis change.
  M = zero_matrix(Fq, dim(A), dim(A))
  for i = 1:dim(A)
    elem_to_mat_row!(M, i, matrix_basis[i])
  end
  return B, hom(A, B, inv(M), M)
end

################################################################################
#
#  Schur Index at real place
#
################################################################################

function trace_signature(A::AlgAss, P::InfPlc)
  M = trred_matrix(basis(A))
  Ky, y = PolynomialRing(base_ring(A), "y", cached = false)
  f = charpoly(Ky, M)
  npos = number_positive_roots(f, P)
  return (npos, degree(f) - npos)
end

function trace_signature(A::AlgAss{fmpq})
  M = trred_matrix(basis(A))
  Ky, y = PolynomialRing(base_ring(A), "y", cached = false)
  f = charpoly(Ky, M)
  Zx = Globals.Zx
  d = denominator(f)
  npos = number_positive_roots(change_base_ring(FlintZZ, d * f))
  return (npos, degree(f) - npos)
end

function schur_index_at_real_plc(A::AlgAss{fmpq})
  if dim(A) % 4 != 0
    return 1
  end
  x = trace_signature(A)
  n = root(dim(A),2)
  if x[1] == divexact(n*(n+1),2)
    return 1
  else
    return 2
  end
end

function schur_index_at_real_plc(A::AlgAss, P::InfPlc)
  if dim(A) % 4 != 0
    return 1
  end
  x = trace_signature(A, P)
  n = root(dim(A),2)
  if x[1] == divexact(n*(n+1),2)
    return 1
  else
    return 2
  end
end

################################################################################
#
#  Is split
#
################################################################################

@doc Markdown.doc"""
    issplit(A::AlgAss{fmpq}) -> Bool

Given a $\mathbb Q$-algebra $A$, this function returns `true` if $A$ splits
over $\mathbb Q$ and `false` otherwise.
"""
function issplit(A::AlgAss{fmpq})
  O = Order(A, basis(A))
  i = schur_index_at_real_plc(O)
  if i==2
    @vprint :AlgAssOrd 1 "Not split at the infinite prime\n"
    return false
  end
  fac = factor(root(abs(discriminant(O)),2))
  for (p,j) in fac
    O1 = pmaximal_overorder(O, Int(p))
    if valuation(discriminant(O1), Int(p)) != 0
      @vprint :AlgAssOrd 1 "Not split at $p\n"
      return false
    end
  end
  return true
end

function issplit(A::AlgAss, P::InfPlc)
  if iscomplex(P)
    return true
  end
  return schur_index_at_real_plc(A, P) == 1
end

################################################################################
#
#  Eichler condition
#
################################################################################

function ramified_infinite_places(A::AlgAss{nf_elem})
  K = base_ring(A)
  inf_plc = Vector{InfPlc}()
  places = real_places(K)
  for p in places
    if !issplit(A, p)
      push!(inf_plc, p)
    end
  end

  return inf_plc
end

# Tests whether A fulfils the Eichler condition relative to the maximal Z-order
# of base_ring(A)
function iseichler(A::AlgAss{nf_elem})
  @assert issimple(A)
  @assert iscentral(A)
  if dim(A) != 4
    return true
  end
  K = base_ring(A)
  places = real_places(K)
  for p in places
    if issplit(A, p)
      return true
    end
  end
  return false
end

function iseichler(A::AlgAss{fmpq})
  @assert issimple(A)
  @assert iscentral(A)
  if dim(A) != 4
    return true
  end
  O = Order(A, basis(A))
  return schur_index_at_real_plc(O) == 1
end

################################################################################
#
#  Direct product
#
################################################################################

@doc Markdown.doc"""
    direct_product(algebras::AlgAss...; task::Symbol = :sum)
      -> AlgAss, Vector{AbsAlgAssMor}, Vector{AbsAlgAssMor}
    direct_product(algebras::Vector{AlgAss}; task::Symbol = :sum)
      -> AlgAss, Vector{AbsAlgAssMor}, Vector{AbsAlgAssMor}

Returns the algebra $A = A_1 \times \cdots \times A_k$. `task` can be
":sum", ":prod", ":both" or ":none" and determines which canonical maps
are computed as well: ":sum" for the injections, ":prod" for the projections.
"""
function direct_product(algebras::Vector{ <: AlgAss{T} }; task::Symbol = :sum) where T
  return direct_product(algebras..., task = task)
end

function direct_product(algebras::AlgAss{T}...; task::Symbol = :sum) where T
  @assert !isempty(algebras)
  @assert all( A -> base_ring(A) == base_ring(algebras[1]), algebras)
  @assert task in [ :prod, :sum, :both, :none ]

  d = sum( dim(A) for A in algebras )
  mt = zeros(base_ring(algebras[1]), d, d, d)
  offset = 0
  for B in algebras
    dd = dim(B)
    mtB = multiplication_table(B)
    for i = 1:dd
      for j = 1:dd
        for k = 1:dd
          mt[i + offset, j + offset, k + offset] = multiplication_table(B)[i, j, k]
        end
      end
    end
    offset += dd
  end
  A = AlgAss(base_ring(algebras[1]), mt)
  if task == :none
    return A
  end

  if task == :sum || task == :both
    inj = Vector{morphism_type(typeof(A), typeof(algebras[1]))}(undef, length(algebras))
  end
  if task == :prod || task == :both
    proj = Vector{morphism_type(typeof(A), typeof(algebras[1]))}(undef, length(algebras))
  end
  offset = 0
  for i = 1:length(algebras)
    B = algebras[i]
    M = zero_matrix(base_ring(A), dim(A), dim(B))
    for i = 1:dim(B)
      M[i + offset, i] = one(base_ring(A))
    end
    Mt = transpose(M)
    if task == :sum || task == :both
      inj[i] = hom(B, A, Mt, M)
    end
    if task == :prod || task == :both
      proj[i] = hom(A, B, M, Mt)
    end
    offset += dim(B)
  end
  if task == :prod
    return A, proj
  elseif task == :sum
    return A, inj
  else
    return A, proj, inj
  end
end

@doc Markdown.doc"""
    direct_product(fields::AnticNumberFields...)
      -> AlgAss{fmpq}, Vector{AbsAlgAssToNfAbsMor}
    direct_product(fields::Vector{AnticNumberFields})
      -> AlgAss{fmpq}, Vector{AbsAlgAssToNfAbsMor}

Returns the algebra $A = K_1 \times \cdots \times K_k$ and the projection
maps $A ->> K_i$.
"""
function direct_product(fields::Vector{AnticNumberField})
  return direct_product(fields...)
end

function direct_product(fields::AnticNumberField...)
  algebras = Tuple{AlgAss{fmpq}, AbsAlgAssToNfAbsMor{AlgAss{fmpq}, elem_type(AlgAss{fmpq}), AnticNumberField, fmpq_mat}}[ AlgAss(K) for K in fields ]
  A, proj, inj = direct_product([ B for (B, m) in algebras ], task = :both)
  A.decomposition = [ (algebras[i][1], inj[i]) for i = 1:length(algebras) ]
  maps_to_fields = Vector{AbsAlgAssToNfAbsMor{AlgAss{fmpq}, elem_type(AlgAss{fmpq}), AnticNumberField, fmpq_mat}}(undef, length(fields))
  for i = 1:length(fields)
    # Assumes, that the map algebras[i] -> K is given by the identity matrix
    maps_to_fields[i] = AbsAlgAssToNfAbsMor(A, fields[i], proj[i].mat, proj[i].imat)
  end
  A.maps_to_numberfields = Tuple{AnticNumberField, AbsAlgAssToNfAbsMor{AlgAss{fmpq}, elem_type(AlgAss{fmpq}), AnticNumberField, fmpq_mat}}[ (fields[i], maps_to_fields[i]) for i = 1:length(fields) ]
  return A, maps_to_fields
end

################################################################################
#
#  Quaternion algebras
#
################################################################################

function quaternion_algebra2(K::Field, a::T, b::T) where { T <: FieldElem }

  M = zeros(K, 4, 4, 4)

  M[1, 1, 1] = one(K) # 1*1=1
  M[1, 2, 2] = one(K) # 1*i=i
  M[1, 3, 3] = one(K) # 1*j=j
  M[1, 4, 4] = one(K) # 1*ij=1

  M[2, 1, 2] = one(K)
  M[2, 2, 1] = a
  M[2, 3, 4] = one(K)
  M[2, 4, 3] = a

  M[3, 1, 3] = one(K)
  M[3, 2, 4] = -one(K)
  M[3, 3, 1] = b
  M[3, 4, 2] = -b

  M[4, 1, 4] = one(K)
  M[4, 2, 3] = -a
  M[4, 3, 2] = b
  M[4, 4, 1] = -a*b

  return AlgAss(K, M, [ one(K), zero(K), zero(K), zero(K) ])
end

quaternion_algebra2(K::Field, a::Int, b::Int) = quaternion_algebra2(K, K(a), K(b))

quaternion_algebra2(a::Int, b::Int) = quaternion_algebra2(FlintQQ, fmpq(a), fmpq(b))

################################################################################
#
#  Opposite algebra
#
################################################################################

function opposite_algebra(A::AlgAss)
  K = base_ring(A)
  B = basis(A)
  d = dim(A)
  z = Array{elem_type(K), 3}(undef, d, d, d)
  for i in 1:d
    for j in 1:d
      z[i, j, :] = A.mult_table[j, i, :]
    end
  end
  o = one(A).coeffs

  B = AlgAss(K, z, o)
  return B, hom(A, B, identity_matrix(K, d), identity_matrix(K, d))
end


text:
module ReplicaExchange

include("recore.jl")
include("interface.jl")

end # module

ext: jl

max_stars_repo_name: zenna/expect

max_issues_repo_name: zenna/expect

max_forks_repo_name: zenna/expect


text:
# Test that @mock works within the context of a patch.
#
# WARNING: Do not use the following code as a template as this test is for illustration
# purposes only. Typically the way this problem is handled is by using `@mock foo(...)`
# within the original function declaration of `foo(::AbstractArray)`.
@testset "mock in patch" begin
    foo(arr::AbstractArray{Float64}) = map(foo, arr)  # Typically foo should use @mock here
    foo(x::Float64) = floor(x)

    # Patching only the function that takes a scalar
    patches = Patch[
        @patch foo(x::Float64) = ceil(x)
    ]

    @test (@mock foo(1.6)) == 1.0

    apply(patches) do
        @test (@mock foo(1.6)) == 2.0
        @test (@mock foo([1.6])) == [1.0]  # Ends up calling original function
    end

    # Create a set of patches where @mock is used within a @patch
    patches = Patch[
        @patch foo(arr::AbstractArray{Float64}) = map(x -> (@mock foo(x)), arr)
        @patch foo(x::Float64) = ceil(x)
    ]

    apply(patches) do
        @test (@mock foo(1.6)) == 2.0
        @test (@mock foo([1.6])) == [2.0]
    end
end

# https://github.com/invenia/Mocking.jl/issues/59
@testset "patch calls patch" begin
    f(args...) = 0

    patches = [
       @patch f(a::Function, b) = b
       @patch f(b) = @mock f(() -> nothing, b)
    ]

    apply(patches) do
        @test (@mock f(identity, 1)) == 1
        @test (@mock f(1)) == 1
    end
end


text:
# In covariance adaptation we want to use cholesky factorization instead
# of a full eigendecomposition. Can we speed this up with sparse matrices in
# Julia?

# Objective function to be maximized.
function my_fun1(x)
  exp(-(x-2).^2) + 0.8 * exp(-(x+2).^2)
end

Ns = [10, 100, 1000, 2000]
num_reps = 10
Ds = [0.01, 0.10, 0.50]

function non_sparse_sqrtm(C, randsamples)
  sqrtC = sqrtm(C)
  s = sqrtC * randsamples
  C + (s * s')
end

function sparse_cholfact(C, randsamples)
  sqrtC = sparse(cholfact(C))
  s = sqrtC' * randsamples
  C = C + (s * s')
end

for(n in Ns)
  # Non-sparse, sqrtm
  C = eye(n)
  if n <= 2000
    tic()
    for(rep in 1:num_reps)
      C = non_sparse_sqrtm(C, randn(n, n))
    end
    t = toq()
    tpdn = t/(num_reps*n)
    println("Non-sparse, sqrtm, n = $(n), time = $(t/num_reps), time = $(tpdn)")
  end

  for(d in Ds)
    # Sparse, cholfact, randn
    C = speye(n)
    tic()
    for(rep in 1:num_reps)
      C = sparse_cholfact(C, sprandn(n, n, d))
    end
    t = toq()
    tpds = t/(num_reps*n)
    println("Sparse, cholfact, d = $(d), n = $(n), time = $(t/num_reps), time = $(tpds), $(tpdn/tpds) times faster")
  end
end


text:
################################################################################
# setting up
################################################################################
const VT = Pair{UInt32,UInt32}
const VTV = Vector{VT}

################################################################################
# Straight counting sort
################################################################################
function sortperm_int_range2(a, rangelen, minval)
    cnts = fill(0, rangelen)

    # create a count first
    @inbounds for ai in a
        cnts[ai - minval + 1] +=  1
    end

    # create cumulative sum
    cumsum!(cnts, cnts)

    la = length(a)
    res = Vector{Int}(la)
    @inbounds for i in la:-1:1
        ai = a[i] - minval + 1
        c = cnts[ai]
        cnts[ai] -= 1
        res[c] = i
    end

    res
end

################################################################################
# Hybrid MSD and counting sort
# I have shown below that coutning sorts beats radix sort up to 2^11= 2048 unique values
# so it makes sense to create a better one
################################################################################
function fsortperm_msd_hybrid(a, rangelen, minval, RADIX_SIZE;  sortcutoff = -1)
    # println(RADIX_SIZE," ", now())
    @assert 2^32 > rangelen
    @assert 2^32 >= length(a)
    vs = Vector{VT}(length(a))
    @inbounds for i in eachindex(a)
        vs[i] = Pair(UInt32(i),UInt32(a[i]))
    end
    # @time @inbounds vs = [Pair(UInt32(i),UInt32(ai)) for (i, ai) in enumerate(a)]

    iters = Int(ceil(log(2, rangelen)/RADIX_SIZE))
    # @assert iters <= ceil(Integer, sizeof(typeof(vs[1].second))*8/RADIX_SIZE)

    _fsortperm_msd_hybrid(vs, 1, length(vs), rangelen, minval, RADIX_SIZE, iters, sortcutoff = sortcutoff)
end

function _fsortperm_msd_hybrid(vs, lo, hi, rangelen, minval, RADIX_SIZE, iters, ts = similar(vs); sortcutoff = -1)
    len = hi-lo+1
    if len  == 1
        return [Int(vs[lo].first)]
    elseif rangelen <= sortcutoff || len <= sortcutoff
        # println("iters $iters lo:hi $(lo:hi) rangelen $rangelen minval $minval")
        av = [a1.second for a1 in vs[lo:hi]]
        # println("extrema $(extrema(Int.(av)))")
        if lo == 3026 && hi == 4037
            # @show av
            throw(error())
        end
        av_sortperm = sortperm_int_range2(av, rangelen, minval)
        # @show av_sortperm
        # issorted([Int(vs[j+lo-1].second) for j in av_sortperm])
        res = [Int(vs[j+lo-1].first) for j in av_sortperm]
        # println("ok2")
        return res
    end

    RADIX_MASK = UInt32(1<<RADIX_SIZE-1)

    # Init
    bin = zeros(UInt32, 2^RADIX_SIZE)

    # Histogram for each element, radix
    @inbounds for i = lo:hi
        v = vs[i].second - minval
        idx = Int((v >> (iters-1)*RADIX_SIZE) & RADIX_MASK) + 1
        @inbounds bin[idx] += 1
    end

    # Sort!

    # check for degenerate case
    v = vs[hi].second - minval
    idx = Int((v >> (iters-1)*RADIX_SIZE) & RADIX_MASK) + 1

    # are all values the same at this radix?
    # if bin[idx,iters] == len;  continue;  end
    if bin[idx] == len
        if iters == 1
            return [Int(vsi.first) for vsi in vs[lo:hi]]
        else
            return _fsortperm_msd_hybrid(vs, lo, hi, rangelen, minval, RADIX_SIZE, iters-1, ts, sortcutoff = sortcutoff)
        end
    end

    # cbin = cumsum(bin[:])
    cbin  = copy(bin)
    cumsum!(bin, bin)
    cumsumbin = copy(bin)

    # copy the elements to the temporary array first
    @inbounds for i in lo:hi
        ts[i] = vs[i]
    end

    # now assign to original in the right order
    @inbounds for i in lo:hi
        v = ts[i].second - minval
        idx = Int((v >> (iters-1)*RADIX_SIZE) & RADIX_MASK) + 1
        ci = bin[idx]
        vs[lo - 1 + ci] = ts[i]
        bin[idx] -= 1
    end

    if iters == 1
        return [Int(vsi.first) for vsi in vs[lo:hi]]
    end

    # now that it's finished sorting sort each chunk recursive now
    res = Vector{Int}(length(vs))
    start_now = 1
    new_rangelen = 2^((iters-1)*RADIX_SIZE)
    # if false
    #     i = 1
    #     lo = 1
    #     hi = cumsumbin[1]
    #     rangelen = new_rangelen
    #     minval = 1
    #     iters = 1

    #     i = 2
    #     lo = cumsumbin[1] + 1
    #     hi = cumsumbin[2]
    #     hi |> Int
    #     minval = 1+(i-1)new_rangelen

    #     i = 4
    #     lo = cumsumbin[i-1] + 1
    #     hi = cumsumbin[3]
    #     hi |> Int
    #     minval = 1+(i-1)new_rangelen
    #     iters = 1
    # end

    for i = 1:length(cumsumbin)
        if cbin[i] != 0
            # println("iters $iters ith count $i")
            new_minval = minval+(i-1)new_rangelen
            res[start_now:cumsumbin[i]] = _fsortperm_msd_hybrid(vs, start_now + lo - 1, cumsumbin[i] + lo - 1, new_rangelen, new_minval, RADIX_SIZE, iters-1, ts, sortcutoff = sortcutoff)
            start_now = cumsumbin[i] + 1
            # println("Ok")
        end
    end
    res
end

################################################################################
# sorting algorithms from discourse
# https://discourse.julialang.org/t/ironic-observation-about-sort-and-sortperm-speed-for-small-intergers-vs-r/8715/20?u=xiaodai
################################################################################
# a = rand(1:1_000_000, 100_000_000)

function sortperm_int_range_p1(a, rangelen, minval, RADIX_SIZE)
    # println(RADIX_SIZE," ", now())
    @assert 2^32 > rangelen
    @assert 2^32 >= length(a)
    @time vs::VTV = Vector{VT}(length(a))
    @time @inbounds for i in eachindex(a)
        vs[i] = Pair(UInt32(i),UInt32(a[i]))
    end
    # @time @inbounds vs = [Pair(UInt32(i),UInt32(ai)) for (i, ai) in enumerate(a)]
    _sortperm_int_range_p1(vs, rangelen, minval, RADIX_SIZE)
end

function _sortperm_int_range_p1(vs, rangelen, minval, RADIX_SIZE)
    RADIX_MASK = UInt32(1<<RADIX_SIZE-1)
    ts = similar(vs)

    # Init
    lo = 1
    hi = length(vs)
    iters = Int(ceil(log(2, rangelen)/RADIX_SIZE))

    # println("iters: $iters")
    @assert iters <= ceil(Integer, sizeof(typeof(vs[1].second))*8/RADIX_SIZE)
    bin = zeros(UInt32, 2^RADIX_SIZE, iters)

    # Histogram for each element, radix
    @inbounds for i = lo:hi
        v = vs[i].second
        for j = 1:iters
            idx = Int((v >> (j-1)*RADIX_SIZE) & RADIX_MASK) + 1
            @inbounds bin[idx,j] += 1
        end
    end

    # Sort!
    swaps = 0
    len = hi-lo+1
    @inbounds for j = 1:iters
        # Unroll first data iteration, check for degenerate case
        v = vs[hi].second
        idx = Int((v >> (j-1)*RADIX_SIZE) & RADIX_MASK) + 1

        # are all values the same at this radix?
        if bin[idx,j] == len;  continue;  end

        cbin = cumsum(bin[:,j])
        ci = cbin[idx]
        ts[ci] = vs[hi]
        cbin[idx] -= 1

        # Finish the loop...
        @inbounds for i in hi-1:-1:lo
            v = vs[i].second
            idx = Int((v >> (j-1)*RADIX_SIZE) & RADIX_MASK) + 1
            ci = cbin[idx]
            ts[ci] = vs[i]
            cbin[idx] -= 1
        end
        vs,ts = ts,vs
        swaps += 1
    end

    if isodd(swaps)
        vs,ts = ts,vs
        @inbounds for i = lo:hi
            vs[i] = ts[i]
        end
    end
    res = Vector{Int}(length(vs))
    @inbounds for i in eachindex(vs)
        res[i] = Int(vs[i].first)
    end
    res
    # [Int(vs1.first) for vs1 in vs]
end


text:
module covid19abm
using Parameters, Distributions, StatsBase, StaticArrays, Random, Match, DataFrames

# health status
@enum HEALTH SUS LAT PRE ASYMP MILD MISO INF IISO HOS ICU REC DED UNDEF

# define an agent and all agent properties
Base.@kwdef mutable struct Human
    idx::Int64 = 0
    health::HEALTH = SUS
    swap::HEALTH = UNDEF
    sickfrom::HEALTH = UNDEF
    nextday_meetcnt::Int16 = 0 ## how many contacts for a single day
    age::Int16   = 0   # in years. don't really need this but left it incase needed later
    ag::Int16    = 0   # age group
    tis::Int16   = 0   # time in state
    exp::Int16   = 0   # max statetime
    dur::NTuple{4, Int8} = (0, 0, 0, 0)   # Order: (latents, asymps, pres, infs) TURN TO NAMED TUPS LATER
    doi::Int16   = 999   # day of infection.
    iso::Bool = false  ## isolated (limited contacts)
    isovia::Symbol = :null ## isolated via quarantine (:qu), preiso (:pi), intervention measure (:im), or contact tracing (:ct)
    tracing::Bool = false ## are we tracing contacts for this individual?
    tracestart::Int16 = -1 ## when to start tracing, based on values sampled for x.dur
    traceend::Int16 = -1 ## when to end tracing
    tracecontacts::Array{Bool, 1} = [false for _ = 1:10000]
    tracedxp::Int16 = 0 ## the trace is killed after tracedxp amount of days
end
Base.show(io::IO, ::MIME"text/plain", z::Human) = dump(z)

## default system parameters
@with_kw mutable struct ModelParameters @deftype Float64    ## use @with_kw from Parameters
    β = 0.0
    frelasymp::Float64 = 0.11 ## relative transmission of asymptomatic
    seasonal::Bool = false ## seasonal betas or not
    popsize::Int64 = 10000
    prov::Symbol = :newyork
    calibration::Bool = false
    modeltime::Int64 = 300
    initialinf::Int64 = 1
    initialhi::Int64 = 0 ## initial herd immunity, inserts number of REC individuals
    asymp_props::NTuple{5, Float64} = (0.25, 0.25, 0.14, 0.07, 0.07) # percentage going to asymp (age-specific)
    mild_props::NTuple{5, Float64} = (0.95, 0.9, 0.85, 0.6, 0.2)  # percentage of sick individuals going to mild infection stage
    τmild::Int64 = 0 ## days before they self-isolate for mild cases
    fmild::Float64 = 0.0  ## percent of people practice self-isolation
    fsevere::Float64 = 0.0 # fixed at 0.80
    tpreiso::Int64 = 0 ## isolate presymptomatic -- begin time parameter
    fpreiso::Float64 = 0.0 ## isolate presymptomatic -- fraction.
    quar_grp::Int8 = 5 ## isolate entire agegroup -- time onset of simulations fixed
    quar_grp_frac::Float64 = 0.0 ## isolate entire agegroup -- fraction
    ctstrat::Int8 = 0 ## contact tracing strategy
    fctcapture::Float16 = 0.0 ## how many symptomatic people identified
    fcontactst::Float16 = 0.0 ## fraction of contacts being isolated/quarantined
    cdaysback::Int8 = 0 ## number of days to go back and collect contacts
    strat3qdays::Int8 = 0
end

# struct for contact tracing statistics
Base.@kwdef mutable struct ct_data_collect
    total_symp_id::Int64 = 0  # total symptomatic identified
    totaltrace::Int64 = 0     # total contacts traced
    totalisolated::Int64 = 0  # total number of people isolated
    iso_sus::Int64 = 0        # total susceptible isolated IN DAYS
    iso_lat::Int64 = 0        # total latent isolated IN DAYS
    iso_asymp::Int64 = 0      # total asymp isolated IN DAYS
    iso_symp::Int64 = 0       # total symp (mild, inf) isolated IN DAYS
end

## load empty constants  (will be loaded for each process if using pmap, watch for memory issues)
const humans = Array{Human}(undef, 0)
const p = ModelParameters()  ## setup default parameters
const agebraks = @SVector [0:4, 5:19, 20:49, 50:64, 65:99]
const BETAS = Array{Float64, 1}(undef, 0) ## to hold betas (whether fixed or seasonal), array will get resized
const ct_data = ct_data_collect()
export ModelParameters, HEALTH, Human, humans, BETAS

# runs a single simulation
function runsim(simnum, ip::ModelParameters)
    Random.seed!(simnum*726)

    # function runs the `main` function, and collects the data as dataframes.
    hmatrix = main(ip)

    # get infectors counters
    infectors = _count_infectors()

    ct_numbers = (ct_data.total_symp_id, ct_data.totaltrace, ct_data.totalisolated,
                    ct_data.iso_sus, ct_data.iso_lat, ct_data.iso_asymp, ct_data.iso_symp)

    # get simulation age groups
    ags = [x.ag for x in humans] # store a vector of the age group distribution
    all = _collectdf(hmatrix)
    spl = _splitstate(hmatrix, ags)
    ag1 = _collectdf(spl[1])
    ag2 = _collectdf(spl[2])
    ag3 = _collectdf(spl[3])
    ag4 = _collectdf(spl[4])
    ag5 = _collectdf(spl[5])
    insertcols!(all, 1, :sim => simnum); insertcols!(ag1, 1, :sim => simnum); insertcols!(ag2, 1, :sim => simnum);
    insertcols!(ag3, 1, :sim => simnum); insertcols!(ag4, 1, :sim => simnum); insertcols!(ag5, 1, :sim => simnum);
    return (a=all, g1=ag1, g2=ag2, g3=ag3, g4=ag4, g5=ag5, infectors=infectors, ct_numbers=ct_numbers)
end
export runsim

function main(ip::ModelParameters)
    ## datacollection
    # matrix to collect model state for every time step

    # reset the parameters for the simulation scenario
    reset_params(ip)

    hmatrix = zeros(Int16, p.popsize, p.modeltime)
    initialize() # initialize population
    p.popsize != length(humans[1].tracecontacts) && error("forgot to change size of tracecontacts property in humans")

    # insert initial infected agents into the model
    if p.calibration
        insert_infected(PRE, p.initialinf, 4)
    else
        insert_infected(PRE, p.initialinf, 4)
        insert_infected(REC, p.initialhi, 4)
    end

    ## save the preisolation isolation parameter (it gets turned on at time = p.tpreiso)
    _fpreiso = p.fpreiso
    p.fpreiso = 0

    # split population in agegroups
    grps = get_ag_dist()

    # start the time loop
    for st = 1:p.modeltime
        # start of day
        if st == p.tpreiso ## time to introduce testing
            p.fpreiso = _fpreiso
        end
        _get_model_state(st, hmatrix) ## this datacollection needs to be at the start of the for loop
        dyntrans(st, grps)
        sw = time_update()
        # end of day
    end
    return hmatrix ## return the model state as well as the age groups.
end
export main

## Initialization Functions
reset_params() = reset_params(ModelParameters())
function reset_params(ip::ModelParameters)
    # the p is a global const
    # the ip is an incoming different instance of parameters
    # copy the values from ip to p.
    for x in propertynames(p)
        setfield!(p, x, getfield(ip, x))
    end

    # reset the contact tracing data collection structure
    for x in propertynames(ct_data)
        setfield!(ct_data, x, 0)
    end

    # resize and update the BETAS constant array
    init_betas()

    # error checks
    p.popsize == 0 && error("no population size given")
    p.ctstrat ∉ (0, 1, 3) && error("invalid ct strategy")
    # resize the human array to change population size
    resize!(humans, p.popsize)
end
export reset_params

function initialize()
    agedist = get_province_ag(p.prov)
    for i = 1:p.popsize
        humans[i] = Human()              ## create an empty human
        x = humans[i]
        x.idx = i
        x.ag = rand(agedist)
        x.age = rand(agebraks[x.ag])
        x.exp = 999  ## susceptible people don't expire.
        x.dur = sample_epi_durations() # sample epi periods
        if rand() < p.quar_grp_frac && x.ag == p.quar_grp   ## check if elderly need to be quarantined.
            x.iso = true
            x.isovia = :qu
        end
        # initialize the next day counts (this is important in initialization since dyntrans runs first)
        get_nextday_counts(x)
    end
end
export initialize

function init_betas()
    # type stable
    if p.seasonal
        tmp = p.β .* td_seasonality()
    else
        tmp = p.β .* ones(Float64, p.modeltime)
    end
    resize!(BETAS, length(tmp))
    for i = 1:length(tmp)
        BETAS[i] = tmp[i]
    end
end

function td_seasonality()
    # returns a vector of seasonal oscillations
    # type stable
    t = 1:p.modeltime
    a0 = 6.261
    a1 = -11.81
    b1 = 1.817
    w = 0.022 #0.01815
    temp = @. a0 + a1*cos((80-t)*w) + b1*sin((80-t)*w)  #100
    temp = (temp .- 2.5*minimum(temp))./(maximum(temp) .- minimum(temp)); # normalize
    return temp
end

## Data Collection/ Model State functions
function _get_model_state(st, hmatrix)
    # collects the model state (i.e. agent status at time st)
    for i=1:length(humans)
        hmatrix[i, st] = Int(humans[i].health)
    end
end
export _get_model_state

function _collectdf(hmatrix)
    ## takes the output of the humans x time matrix and processes it into a dataframe
    #_names_inci = Symbol.(["lat_inc", "mild_inc", "miso_inc", "inf_inc", "iiso_inc", "hos_inc", "icu_inc", "rec_inc", "ded_inc"])
    #_names_prev = Symbol.(["sus", "lat", "mild", "miso", "inf", "iiso", "hos", "icu", "rec", "ded"])
    mdf_inc, mdf_prev = _get_incidence_and_prev(hmatrix)
    mdf = hcat(mdf_inc, mdf_prev)
    _names_inc = Symbol.(string.((Symbol.(instances(HEALTH)[1:end - 1])), "_INC"))
    _names_prev = Symbol.(string.((Symbol.(instances(HEALTH)[1:end - 1])), "_PREV"))
    _names = vcat(_names_inc..., _names_prev...)
    datf = DataFrame(mdf, _names)

    lat_inc = mdf_inc[:, 2]
    inf_prv = vec(sum(mdf_prev[:, 3:8], dims=2))
    secondarys = lat_inc ./ inf_prv
    replace!(secondarys, NaN => 0)
    replace!(secondarys, Inf => 0)
    insertcols!(datf, 1, :time => 1:p.modeltime) ## add a time column to the resulting dataframe
    insertcols!(datf, 2, :secondarys => secondarys) ## add a time column to the resulting dataframe
    return datf
end

function _splitstate(hmatrix, ags)
    #split the full hmatrix into 4 age groups based on ags (the array of age group of each agent)
    #sizes = [length(findall(x -> x == i, ags)) for i = 1:4]
    matx = []#Array{Array{Int64, 2}, 1}(undef, 4)
    for i = 1:length(agebraks)
        idx = findall(x -> x == i, ags)
        push!(matx, view(hmatrix, idx, :))
    end
    return matx
end
export _splitstate

function _get_incidence_and_prev(hmatrix)
    cols = instances(HEALTH)[1:end - 1] ## don't care about the UNDEF health status
    inc = zeros(Int64, p.modeltime, length(cols))
    pre = zeros(Int64, p.modeltime, length(cols))
    for i = 1:length(cols)
        inc[:, i] = _get_column_incidence(hmatrix, cols[i])
        pre[:, i] = _get_column_prevalence(hmatrix, cols[i])
    end
    return inc, pre
end

function _get_column_incidence(hmatrix, hcol)
    inth = Int(hcol)
    timevec = zeros(Int64, p.modeltime)
    for r in eachrow(hmatrix)
        idx = findfirst(x -> x == inth, r)
        if idx !== nothing
            timevec[idx] += 1
        end
    end
    return timevec
end

function _get_column_prevalence(hmatrix, hcol)
    inth = Int(hcol)
    timevec = zeros(Int64, p.modeltime)
    for (i, c) in enumerate(eachcol(hmatrix))
        idx = findall(x -> x == inth, c)
        if idx !== nothing
            ps = length(c[idx])
            timevec[i] = ps
        end
    end
    return timevec
end

function _count_infectors()
    pre_ctr = asymp_ctr = mild_ctr = inf_ctr = 0
    for x in humans
        if x.health != SUS ## meaning they got sick at some point
            if x.sickfrom == PRE
                pre_ctr += 1
            elseif x.sickfrom == ASYMP
                asymp_ctr += 1
            elseif x.sickfrom == MILD || x.sickfrom == MISO
                mild_ctr += 1
            elseif x.sickfrom == INF || x.sickfrom == IISO
                inf_ctr += 1
            else
                error("sickfrom not set right: $(x.sickfrom)")
            end
        end
    end
    return (pre_ctr, asymp_ctr, mild_ctr, inf_ctr)
end
export _collectdf, _get_incidence_and_prev, _get_column_incidence, _get_column_prevalence, _count_infectors

# Generic Model Functions
function get_province_ag(prov)
    ret = @match prov begin
        :alberta => Distributions.Categorical(@SVector [0.0655, 0.1851, 0.4331, 0.1933, 0.1230])
        :bc => Distributions.Categorical(@SVector [0.0475, 0.1570, 0.3905, 0.2223, 0.1827])
        :canada => Distributions.Categorical(@SVector [0.0540, 0.1697, 0.3915, 0.2159, 0.1689])
        :manitoba => Distributions.Categorical(@SVector [0.0634, 0.1918, 0.3899, 0.1993, 0.1556])
        :newbruns => Distributions.Categorical(@SVector [0.0460, 0.1563, 0.3565, 0.2421, 0.1991])
        :newfdland => Distributions.Categorical(@SVector [0.0430, 0.1526, 0.3642, 0.2458, 0.1944])
        :nwterrito => Distributions.Categorical(@SVector [0.0747, 0.2026, 0.4511, 0.1946, 0.0770])
        :novasco => Distributions.Categorical(@SVector [0.0455, 0.1549, 0.3601, 0.2405, 0.1990])
        :nunavut => Distributions.Categorical(@SVector [0.1157, 0.2968, 0.4321, 0.1174, 0.0380])
        :ontario => Distributions.Categorical(@SVector [0.0519, 0.1727, 0.3930, 0.2150, 0.1674])
        :pei => Distributions.Categorical(@SVector [0.0490, 0.1702, 0.3540, 0.2329, 0.1939])
        :quebec => Distributions.Categorical(@SVector [0.0545, 0.1615, 0.3782, 0.2227, 0.1831])
        :saskat => Distributions.Categorical(@SVector [0.0666, 0.1914, 0.3871, 0.1997, 0.1552])
        :yukon => Distributions.Categorical(@SVector [0.0597, 0.1694, 0.4179, 0.2343, 0.1187])
        :newyork   => Distributions.Categorical(@SVector [0.064000, 0.163000, 0.448000, 0.181000, 0.144000])
        _ => error("shame for not knowing your canadian provinces and territories")
    end
    return ret
end

function get_ag_dist()
    # splits the initialized human pop into its age groups
    grps =  map(x -> findall(y -> y.ag == x, humans), 1:length(agebraks))
    return grps
end

function sample_epi_durations()
    # when a person is sick, samples the
    lat_dist = truncated(LogNormal(log(5.2), 0.1), 4, 7) # truncated between 4 and 7
    pre_dist = truncated(Gamma(1.058, 5/2.3), 0.8, 3)#truncated between 0.8 and 3
    asy_dist = Gamma(5, 1)
    inf_dist = Gamma((3.2)^2/3.7, 3.7/3.2)

    latents = Int.(round.(rand(lat_dist)))
    pres = Int.(round.(rand(pre_dist)))
    latents = latents - pres # ofcourse substract from latents, the presymp periods
    asymps = Int.(ceil.(rand(asy_dist)))
    infs = Int.(ceil.(rand(inf_dist)))
    return (latents, asymps, pres, infs)
end

function insert_infected(health, num, ag)
    ## inserts a number of infected people in the population randomly
    ## this function should resemble move_to_inf()
    ## when health == INF, the ag matters.
    ## for  health == REC, just use all humans
    if health == INF
        l = findall(x -> x.health == SUS && x.ag == ag, humans)
    else
        l = findall(x -> x.health == SUS, humans)
    end
    # we can throw an error here because this function is run right at the start
    length(l) == 0 && error("No susceptibles in the population to change status")
    num > length(l) && error("No susceptibles in the population to change status")

    h = sample(l, num; replace = false)
    @inbounds for i in h
        x = humans[i]
        if health == PRE
            move_to_pre(x) ## the swap may be asymp, mild, or severe, but we can force severe in the time_update function
        elseif health == LAT
            move_to_latent(x)
            #if the person is moving to latent, let them expire immediately.
            #this removes the model artifact of "dipping" at the start of simulations.
            x.tis = x.exp
        elseif health == INF
            move_to_infsimple(x)
        elseif health == REC
            move_to_recovered(x)
        else
            error("can not insert human of health $(health)")
        end
        x.sickfrom = INF # this will add +1 to the INF count in _count_infectors()... keeps the logic simple in that function.
    end

    return h
end
export insert_infected

# Main Time Step of the Model
function time_update()
    # counters to calculate incidence at each time step
    lat=0; pre=0; asymp=0; mild=0; miso=0; inf=0; infiso=0; hos=0; icu=0; rec=0; ded=0;
    for x in humans
        x.tis += 1
        x.doi += 1 # increase day of infection. variable is garbage until person is latent
        if x.tis >= x.exp
            @match Symbol(x.swap) begin
                :LAT  => begin move_to_latent(x); lat += 1; end
                :PRE  => begin move_to_pre(x); pre += 1; end
                :ASYMP => begin move_to_asymp(x); asymp += 1; end
                :MILD => begin move_to_mild(x); mild += 1; end
                :MISO => begin move_to_miso(x); miso += 1; end
                :INF  => begin move_to_inf(x); inf +=1; end
                :IISO => begin move_to_iiso(x); infiso += 1; end
                :HOS  => begin move_to_hospicu(x); hos += 1; end
                :ICU  => begin move_to_hospicu(x); icu += 1; end
                :REC  => begin move_to_recovered(x); rec += 1; end
                :DED  => begin move_to_dead(x); ded += 1; end
                _    => error("swap expired, but no swap set.")
            end
        end
        # run covid-19 functions for other integrated dynamics.
        ct_dynamics(x)

        # get the meet counts for the next day
        get_nextday_counts(x)
    end
    return (lat, mild, miso, inf, infiso, hos, icu, rec, ded)
end
export time_update

@inline _set_isolation(x::Human, iso) = _set_isolation(x, iso, x.isovia)
@inline function _set_isolation(x::Human, iso, via)
    # a helper setter function to not overwrite the isovia property.
    # a person could be isolated in susceptible/latent phase through contact tracing
    # --> in which case it will follow through the natural history of disease
    # --> if the person remains susceptible, then iso = off
    # a person could be isolated in presymptomatic phase through fpreiso
    # --> if x.iso == true from CT and x.isovia == :ct, do not overwrite
    # a person could be isolated in mild/severe phase through fmild, fsevere
    # --> if x.iso == true from CT and x.isovia == :ct, do not overwrite
    # --> if x.iso == true from PRE and x.isovia == :pi, do not overwrite
    x.iso = iso
    x.isovia == :null && (x.isovia = via)
end

function move_to_latent(x::Human)
    ## transfers human h to the incubation period and samples the duration
    x.health = LAT
    x.doi = 0 ## day of infection is reset when person becomes latent
    x.tis = 0   # reset time in state
    x.exp = x.dur[1] # get the latent period
    # the swap to asymptomatic is based on age group.
    # ask seyed for the references
    asymp_pcts = p.asymp_props
    x.swap = rand() < asymp_pcts[x.ag] ? ASYMP : PRE
    ## in calibration mode, latent people never become infectious.
    if p.calibration
        x.swap = LAT
        x.exp = 999
    end
end
export move_to_latent

function move_to_asymp(x::Human)
    ## transfers human h to the asymptomatic stage
    x.health = ASYMP
    x.tis = 0
    x.exp = x.dur[2] # get the presymptomatic period
    x.swap = REC
    # x.iso property remains from either the latent or presymptomatic class
    # if x.iso is true, the asymptomatic individual has limited contacts
end
export move_to_asymp

function move_to_pre(x::Human)
    θ = p.mild_props # percentage of sick individuals going to mild infection stage
    x.health = PRE
    x.tis = 0   # reset time in state
    x.exp = x.dur[3] # get the presymptomatic period

    if rand() < θ[x.ag]
        x.swap = MILD
    else
        x.swap = INF
    end
    # calculate whether person is isolated
    rand() < p.fpreiso && _set_isolation(x, true, :pi)
end
export move_to_pre

function move_to_mild(x::Human)
    ## transfers human h to the mild infection stage for γ days
    x.health = MILD
    x.tis = 0
    x.exp = x.dur[4]
    x.swap = REC
    # x.iso property remains from either the latent or presymptomatic class
    # if x.iso is true, staying in MILD is same as MISO since contacts will be limited.
    # we still need the separation of MILD, MISO because if x.iso is false, then here we have to determine
    # how many days as full contacts before self-isolation
    # NOTE: if need to count non-isolated mild people, this is overestimate as isolated people should really be in MISO all the time
    #   and not go through the mild compartment
    if x.iso || rand() < p.fmild
        x.swap = MISO
        x.exp = p.τmild
    end
end
export move_to_mild

function move_to_miso(x::Human)
    ## transfers human h to the mild isolated infection stage for γ days
    x.health = MISO
    x.swap = REC
    x.tis = 0
    x.exp = x.dur[4] - p.τmild  ## since tau amount of days was already spent as infectious
    _set_isolation(x, true, :mi)
end
export move_to_miso

function move_to_infsimple(x::Human)
    ## transfers human h to the severe infection stage for γ days
    ## simplified function for calibration/general purposes
    x.health = INF
    x.tis = 0
    x.exp = x.dur[4]
    x.swap = REC
    _set_isolation(x, false, :null)
end

function move_to_inf(x::Human)
    ## transfers human h to the severe infection stage for γ days
    ## for swap, check if person will be hospitalized, selfiso, die, or recover

    # h = prob of hospital, c = prob of icu AFTER hospital
    h = (rand(Uniform(0.02, 0.03)), rand(Uniform(0.02, 0.03)), rand(Uniform(0.28, 0.34)), rand(Uniform(0.28, 0.34)), rand(Uniform(0.60, 0.68)))
    c = (rand(Uniform(0.01, 0.015)), rand(Uniform(0.01, 0.015)), rand(Uniform(0.03, 0.05)), rand(Uniform(0.05, 0.20)), rand(Uniform(0.05, 0.15)))
    mh = [0.01/5, 0.01/5, 0.0135/3, 0.01225/1.5, 0.04/2]     # death rate for severe cases.

    # turn off hospitalization in calibration mode
    if p.calibration
        h =  (0, 0, 0, 0, 0)
        c =  (0, 0, 0, 0, 0)
        mh = (0, 0, 0, 0, 0)
    end

    time_to_hospital = Int(round(rand(Uniform(2, 5)))) # duration symptom onset to hospitalization

    x.health = INF
    x.swap = UNDEF
    x.tis = 0
    if rand() < h[x.ag]     # going to hospital or ICU but will spend delta time transmissing the disease with full contacts
        x.exp = time_to_hospital
        x.swap = rand() < c[x.ag] ? ICU : HOS
    else ## no hospital for this lucky (but severe) individual
        if rand() < mh[x.ag]
            x.exp = x.dur[4]
            x.swap = DED
        else
            x.exp = x.dur[4]
            x.swap = REC
            if x.iso || rand() < p.fsevere
                x.exp = 1  ## 1 day isolation for severe cases
                x.swap = IISO
            end
        end
    end
    ## before returning, check if swap is set
    x.swap == UNDEF && error("agent I -> ?")
end

function move_to_iiso(x::Human)
    ## transfers human h to the sever isolated infection stage for γ days
    x.health = IISO
    x.swap = REC
    x.tis = 0     ## reset time in state
    x.exp = x.dur[4] - 1  ## since 1 day was spent as infectious
    _set_isolation(x, true, :mi)
end

function move_to_hospicu(x::Human)
    ## to do: don't sample from distribution unless needed
    ## to do: check whether the parameters seyed gave are for mean/std or shape/scale
    #mh = [0.001, 0.001, 0.00135, 0.01225, 0.04]
    #mc = [0.002, 0.002, 0.0027, 0.0245, 0.08]
    mh = [0.01*2, 0.01*2, 0.0135, 0.01225, 0.03*2]
    mc = 2*mh
    #mh = [0, 0, 0, 0.01225*2, 0.04*2]
    #mc = [0, 0, 0, 0.01225*3.5, 0.04*3.5]

    psiH = Int(round(rand(truncated(Gamma(4.5, 2.75), 8, 17))))
    psiC = Int(round(rand(truncated(Gamma(4.5, 2.75), 8, 17)))) + 2
    muH = Int(round(rand(truncated(Gamma(5.3, 2.1), 9, 15))))
    muC = Int(round(rand(truncated(Gamma(5.3, 2.1), 9, 15)))) + 2

    swaphealth = x.swap
    x.health = swaphealth ## swap either to HOS or ICU
    x.swap = UNDEF
    x.tis = 0
    _set_isolation(x, true) # do not set the isovia property here.

    if swaphealth == HOS
        if rand() < mh[x.ag] ## person will die in the hospital
            x.exp = muH
            x.swap = DED
        else
            x.exp = psiH
            x.swap = REC
        end
    end

    if swaphealth == ICU
        if rand() < mc[x.ag] ## person will die in the ICU
            x.exp = muC
            x.swap = DED
        else
            x.exp = psiC
            x.swap = REC
        end
    end
    ## before returning, check if swap is set
    x.swap == UNDEF && error("agent H -> ?")
end

function move_to_dead(h::Human)
    # no level of alchemy will bring someone back to life.
    h.health = DED
    h.swap = UNDEF
    h.tis = 0
    h.exp = 999 ## stay recovered indefinitely
    h.iso = true # a dead person is isolated
    _set_isolation(h, true)  # do not set the isovia property here.
    # isolation property has no effect in contact dynamics anyways (unless x == SUS)
end

function move_to_recovered(h::Human)
    h.health = REC
    h.swap = UNDEF
    h.tis = 0
    h.exp = 999 ## stay recovered indefinitely
    # do not set isolation property for recovered.
    # if they were isolated through contact tracing, they will unisolate automatically.
    # to do: check other pathways.
    # _set_isolation(h, false)  # do not set the isovia property here.
end

## Contact Tracing Functions
function apply_ct_strategy(y::Human)
    _set_isolation(y, true, :ct)
    ct_data.totalisolated += 1
    if p.ctstrat == 1
        y.tracedxp = 14  # trace isolation will last for 14 days before expiry
    end
    if p.ctstrat == 3
        y.tracedxp = p.strat3qdays  # trace isolation will last for preset amount of days before expiry
    end
end

function ct_dynamics(x::Human)
    # main function for ct dynamics, runs on a daily basis
    # if the person is newly latent (i.e. xh == LAT & doi == 0), calculate tracing times in the future
    # when its time to trace the contacts (i.e. doi == tracingstart), turn tracing property on
    # however only do it for severe cases.
    # now x% of contacts will be traced and stored in memory.
    # now eventually the person will be identified (i.e. when doi == tracingend)
    # then apply isolation to all traced contacts
    # order of if statements matter here

    (p.ctstrat == 0) && return
    xh = x.health
    xs = x.swap
    dur = x.dur
    doi = x.doi

    ## person is newly infectious, calculate tracestart and traceend times
    ## ofcourse asymp is not included, so traceend/start will remain default.
    ## not possible to know whether person will be mild/severe at this case
    ## fctcapture = proportion of infected people that will be contact traced
    if xh == LAT && xs != ASYMP && doi == 0 && rand() < p.fctcapture
        # sample time to id from their symptomatic period
        time_to_id = rand(1:(dur[4]))
        delta = dur[1] + dur[3] + time_to_id #latent + presymp + time to identification time
        # so always delta <= dur[1] + dur[3] + dur[4]
        q = delta - p.cdaysback  ## how many days to trace prior to identification
        x.tracestart = max(0, q) # minimum of zero since possible that delta < p.cdaysback
        x.traceend = delta
        x.tracing = false
        x.tracecontacts .= false  # reset their tracecontacts for new infection
        (x.tracestart > x.traceend) && error("tracestart < traceend")
    end

    ## when the person turns PRE, we know whether they will be in MILD/SEV
    ## only want to capture the severe cases
    # if xh == PRE && doi == dur[1] && xs != INF
    #    x.tracestart = -1
    #    x.traceend = -1
    # end

    ## turn on tracing in the right time window and only for latent, pre, and severe symptomatic
    if doi >= x.tracestart && doi < x.traceend
        x.tracing = true
    end

    ## on day of identification, i.e. doi == x.traceend
    ## isolate all the contacts
    if doi == x.traceend
        x.tracing = false
        apply_ct_strategy(x)   # set isolation true of the traced individual (ofcourse)
        # find x's contacts (this gives us the index of the contacts)
        alltraced = findall(x -> x == true, x.tracecontacts)
        for yidx in alltraced
            if rand() < p.fcontactst  # not all contacts are traced due to imperfect tracing
                apply_ct_strategy(humans[yidx])
            end
        end
        x.tracecontacts .= false
    end

    if x.tracedxp > 0  # person was traced, isolated, and will be for x days
        x.iso == false && error("a traced person should always be isolated until the trace is killed")
        x.tracedxp -= 1

        # collect stats on the traced individuals
        # if xh == INF || xh == MILD || xh == PRE
        #     ct_data.iso_symp += 1
        # elseif xh == LAT
        #     ct_data.iso_lat += 1
        # elseif xh == ASYMP
        #     ct_data.iso_asymp += 1
        # else
        #     ct_data.iso_sus += 1
        # end

        # once the isolation time is over, kill the trace.
        if x.tracedxp == 0
            # else their trace is killed
            _set_isolation(x, false)
            x.isovia = :null # not isolated via contact tracing anymore

            if p.ctstrat == 3
                rn = rand()
                if (x.health == LAT && rn < 0.05) || (x.health == PRE && rn < 0.95) || (x.health == ASYMP && rn < 0.70) || x.health in (MILD, INF, MISO, IISO)
                    # if strategy 3, only those that are tested positive are furthered isolated for 14 days.
                    # this is technically a bug since this if statement can run repeatedly,
                    # but after 14 days none of the conditions are going to be true
                    _set_isolation(x, true, :ct)
                    x.tracedxp = 14
                end
            end
        end
    end
    return
end
export ct_dynamics

## Transmission Dynamics
@inline function _get_betavalue(sys_time, xhealth)
    #bf = p.β ## baseline PRE
    length(BETAS) == 0 && return 0.0
    bf = BETAS[sys_time]
    # values coming from FRASER Figure 2... relative tranmissibilities of different stages.
    if xhealth == ASYMP
        bf = bf * p.frelasymp
    elseif xhealth == MILD || xhealth == MISO
        bf = bf * 0.44
    elseif xhealth == INF || xhealth == IISO
        bf = bf * 0.89
    end
    return bf
end
export _get_betavalue

@inline function get_nextday_counts(x::Human)
    # gets the number of agents x will meet (or recieve contacts from)
    # we can sample this at the start of the simulation to avoid everyday but its type stable and fast
    cnt = 0
    ag = x.ag
    #if person is isolated, they can recieve only 3 maximum contacts 50% of the time
    if x.iso
        cnt = rand() < 0.5 ? 0 : rand(1:3)
    else
        cnt = rand(nbs[ag])  # expensive operation, try to optimize
    end
    if x.health == DED
        cnt = 0
    end
    x.nextday_meetcnt = cnt
    return cnt
end

function dyntrans(sys_time, grps)
    totalmet = 0 # count the total number of contacts (total for day, for all INF contacts)
    totalinf = 0 # count number of new infected
    ## find all the people infectious
    infs = findall(x -> x.health in (PRE, ASYMP, MILD, MISO, INF, IISO), humans)

    # go through every infectious person
    for xid in infs
        x = humans[xid]
        xhealth = x.health
        beta = _get_betavalue(sys_time, xhealth) # get the infected x beta value

        cnts = x.nextday_meetcnt ## how many contacts to give
        cnts == 0 && continue
        gpw = Int.(round.(cm[x.ag]*cnts)) # split the counts over age groups
        for (i, g) in enumerate(gpw)
            # sample the people from each group
            meet = rand(grps[i], g)
            # go through each person
            for j in meet
                y = humans[j]
                ycnt = y.nextday_meetcnt
                ycnt == 0 && continue
                y.nextday_meetcnt = y.nextday_meetcnt - 1 # remove a contact
                totalmet += 1

                # trace contacts
                if x.tracing
                    x.tracecontacts[j] = true
                end

                # tranmission dynamics
                if  y.health == SUS && y.swap == UNDEF
                    if rand() < beta
                        totalinf += 1
                        y.swap = LAT
                        y.exp = y.tis   ## force the move to latent in the next time step.
                        y.sickfrom = xhealth ## stores the infector's status to the infectee's sickfrom
                    end
                end
            end
        end
    end
    return totalmet, totalinf
end
export dyntrans

## Contact Matrix Functions
function contact_matrix()
    # regular contacts, just with 5 age groups.
    #  0-4, 5-19, 20-49, 50-64, 65+
    CM = Array{Array{Float64, 1}, 1}(undef, 5)
    CM[1] = [0.2287, 0.1839, 0.4219, 0.1116, 0.0539]
    CM[2] = [0.0276, 0.5964, 0.2878, 0.0591, 0.0291]
    CM[3] = [0.0376, 0.1454, 0.6253, 0.1423, 0.0494]
    CM[4] = [0.0242, 0.1094, 0.4867, 0.2723, 0.1074]
    CM[5] = [0.0207, 0.1083, 0.4071, 0.2193, 0.2446]
    return CM
end
#
# calibrate for 2.7 r0
# 20% selfisolation, tau 1 and 2.

function negative_binomials()
    ## the means/sd here are calculated using _calc_avgag
    means = [10.21, 16.793, 13.7950, 11.2669, 8.0027]
    sd = [7.65, 11.7201, 10.5045, 9.5935, 6.9638]
    totalbraks = length(means)
    nbinoms = Vector{NegativeBinomial{Float64}}(undef, totalbraks)
    for i = 1:totalbraks
        p = 1 - (sd[i]^2-means[i])/(sd[i]^2)
        r = means[i]^2/(sd[i]^2-means[i])
        nbinoms[i] =  NegativeBinomial(r, p)
    end
    return nbinoms
end
const nbs = negative_binomials()
const cm = contact_matrix()
export negative_binomials, contact_matrix, nbs, cm

## internal functions to do intermediate calculations
function _calc_avgag(lb, hb)
    ## internal function to calculate the mean/sd of the negative binomials
    ## returns a vector of sampled number of contacts between age group lb to age group hb
    dists = _negative_binomials_15ag()[lb:hb]
    totalcon = Vector{Int64}(undef, 0)
    for d in dists
        append!(totalcon, rand(d, 10000))
    end
    return totalcon
end
export _calc_avgag

function _negative_binomials_15ag()
    ## negative binomials 15 agegroups
    AgeMean = Vector{Float64}(undef, 15)
    AgeSD = Vector{Float64}(undef, 15)
    #0-4, 5-9, 10-14, 15-19, 20-24, 25-29, 30-34, 35-39, 40-44, 45-49, 50-54, 55-59, 60-64, 65-69, 70+
    AgeMean = [10.21, 14.81, 18.22, 17.58, 13.57, 13.57, 14.14, 14.14, 13.83, 13.83, 12.3, 12.3, 9.21, 9.21, 6.89]
    AgeSD = [7.65, 10.09, 12.27, 12.03, 10.6, 10.6, 10.15, 10.15, 10.86, 10.86, 10.23, 10.23, 7.96, 7.96, 5.83]

    nbinoms = Vector{NegativeBinomial{Float64}}(undef, 15)
    for i = 1:15
        p = 1 - (AgeSD[i]^2-AgeMean[i])/(AgeSD[i]^2)
        r = AgeMean[i]^2/(AgeSD[i]^2-AgeMean[i])
        nbinoms[i] =  NegativeBinomial(r, p)
    end
    return nbinoms
end
export negative_binomials

## references:
# critical care capacity in Canada https://www.ncbi.nlm.nih.gov/pubmed/25888116
end # module end