module Topologies
import ..BrickMesh
import MPI
using DocStringExtensions

export AbstractTopology, BrickTopology, StackedBrickTopology,
    CubedShellTopology, StackedCubedSphereTopology, isstacked

export grid1d, SingleExponentialStretching, InteriorStretching

"""
    AbstractTopology{dim}

Represents the connectivity of individual elements, with local dimension `dim`.
"""
abstract type AbstractTopology{dim} end

"""
    BoxElementTopology{dim, T} <: AbstractTopology{dim}

The local topology of a larger MPI-distributed topology, represented by
`dim`-dimensional box elements.

This contains the necessary information for the connectivity elements of the
elements on the local process, along with "ghost" elements from neighbouring
processes.

# Fields

$(DocStringExtensions.FIELDS)
"""
struct BoxElementTopology{dim, T} <: AbstractTopology{dim}

  """
  MPI communicator for communicating with neighbouring processes.
  """
  mpicomm::MPI.Comm

  """
  Range of element indices
  """
  elems::UnitRange{Int64}

  """
  Range of real (aka nonghost) element indices
  """
  realelems::UnitRange{Int64}

  """
  Range of ghost element indices
  """
  ghostelems::UnitRange{Int64}

  """
  Ghost element to face is received; `ghostfaces[f,ge] == true` if face `f` of
  ghost element `ge` is received.
  """
  ghostfaces::BitArray{2}

  """
  Array of send element indices
  """
  sendelems::Array{Int64, 1}

  """
  Send element to face is sent; `sendfaces[f,se] == true` if face `f` of send
  element `se` is sent.
  """
  sendfaces::BitArray{2}

  """
  Element to vertex coordinates; `elemtocoord[d,i,e]` is the `d`th coordinate of
  corner `i` of element `e`

  !!! note
      currently coordinates always are of size 3 for `(x1, x2, x3)`
  """
  elemtocoord::Array{T, 3}

  """
  Element to neighboring element; `elemtoelem[f,e]` is the number of the element
  neighboring element `e` across face `f`.  If there is no neighboring element
  then `elemtoelem[f,e] == e`.
  """
  elemtoelem::Array{Int64, 2}

  """
  Element to neighboring element face; `elemtoface[f,e]` is the face number of
  the element neighboring element `e` across face `f`.  If there is no
  neighboring element then `elemtoface[f,e] == f`."
  """
  elemtoface::Array{Int64, 2}

  """
  element to neighboring element order; `elemtoordr[f,e]` is the ordering number
  of the element neighboring element `e` across face `f`.  If there is no
  neighboring element then `elemtoordr[f,e] == 1`.
  """
  elemtoordr::Array{Int64, 2}

  """
  Element to boundary number; `elemtobndy[f,e]` is the boundary number of face
  `f` of element `e`.  If there is a neighboring element then `elemtobndy[f,e]
  == 0`.
  """
  elemtobndy::Array{Int64, 2}

  """
  List of the MPI ranks for the neighboring processes
  """
  nabrtorank::Array{Int64, 1}

  """
  Range in ghost elements to receive for each neighbor
  """
  nabrtorecv::Array{UnitRange{Int64}, 1}

  """
  Range in `sendelems` to send for each neighbor
  """
  nabrtosend::Array{UnitRange{Int64}, 1}

  """
  original order in partitioning
  """
  origsendorder::Array{Int64,1}

  """
  boolean for whether or not this topology has a boundary
  """
  hasboundary::Bool
end

"""
    hasboundary(topology::AbstractTopology)

query function to check whether a topology has a boundary (i.e., not fully
periodic)
"""
hasboundary(topology::AbstractTopology) = topology.hasboundary

if VERSION >= v"1.2-"
  isstacked(::T) where {T<:AbstractTopology} = hasfield(T, :stacksize)
else
  isstacked(::T) where {T<:AbstractTopology} = Base.fieldindex(T, :stacksize,
                                                               false) > 0
end

"""
    BrickTopology{dim, T} <: AbstractTopology{dim}

A simple grid-based topology. This is a convenience wrapper around
[`BoxElementTopology`](@ref).
"""
struct BrickTopology{dim, T} <: AbstractTopology{dim}
  topology::BoxElementTopology{dim, T}
end
Base.getproperty(a::BrickTopology, p::Symbol) =
getproperty(getfield(a, :topology), p)

"""
    CubedShellTopology{T} <: AbstractTopology{2}

A cube-shell topology. This is a convenience wrapper around
[`BoxElementTopology`](@ref).
"""
struct CubedShellTopology{T} <: AbstractTopology{2}
  topology::BoxElementTopology{2, T}
end
Base.getproperty(a::CubedShellTopology, p::Symbol) =
getproperty(getfield(a, :topology), p)

"""
    StackedBrickTopology{dim, T} <: AbstractTopology{dim}

A simple grid-based topology, where all elements on the trailing dimension are
stacked to be contiguous. This is a convenience wrapper around
[`BoxElementTopology`](@ref).
"""
struct StackedBrickTopology{dim, T} <: AbstractTopology{dim}
  topology::BoxElementTopology{dim, T}
  stacksize::Int64
end
function Base.getproperty(a::StackedBrickTopology, p::Symbol)
  return p == :stacksize ? getfield(a, p) :
                           getproperty(getfield(a, :topology), p)
end

"""
    StackedCubedSphereTopology{3, T} <: AbstractTopology{3}

A cube-sphere topology. All elements on the same "vertical" dimension are
stacked to be contiguous. This is a convenience wrapper around
[`BoxElementTopology`](@ref).
"""
struct StackedCubedSphereTopology{T} <: AbstractTopology{3}
  topology::BoxElementTopology{3, T}
  stacksize::Int64
end
function Base.getproperty(a::StackedCubedSphereTopology, p::Symbol)
  return p == :stacksize ? getfield(a, p) :
                           getproperty(getfield(a, :topology), p)
end

""" A wrapper for the BrickTopology """
BrickTopology(mpicomm, Nelems::NTuple{N, Integer}; kw...) where N =
BrickTopology(mpicomm, map(Ne->0:Ne, Nelems); kw...)

"""
    BrickTopology{dim, T}(mpicomm, elemrange; boundary, periodicity)

Generate a brick mesh topology with coordinates given by the tuple `elemrange`
and the periodic dimensions given by the `periodicity` tuple.

The elements of the brick are partitioned equally across the MPI ranks based
on a space-filling curve.

By default boundary faces will be marked with a one and other faces with a
zero.  Specific boundary numbers can also be passed for each face of the brick
in `boundary`.  This will mark the nonperiodic brick faces with the given
boundary number.

# Examples

We can build a 3 by 2 element two-dimensional mesh that is periodic in the
\$x2\$-direction with
```jldoctest brickmesh

using CLIMA.Topologies
using MPI
MPI.Init()
topology = BrickTopology(MPI.COMM_SELF, (2:5,4:6);
                         periodicity=(false,true),
                         boundary=((1,2),(3,4)))
```
This returns the mesh structure for

             x2

              ^
              |
             6-  +-----+-----+-----+
              |  |     |     |     |
              |  |  3  |  4  |  5  |
              |  |     |     |     |
             5-  +-----+-----+-----+
              |  |     |     |     |
              |  |  1  |  2  |  6  |
              |  |     |     |     |
             4-  +-----+-----+-----+
              |
              +--|-----|-----|-----|--> x1
                 2     3     4     5

For example, the (dimension by number of corners by number of elements) array
`elemtocoord` gives the coordinates of the corners of each element.
```jldoctest brickmesh
julia> topology.elemtocoord
2×4×6 Array{Int64,3}:
[:, :, 1] =
 2  3  2  3
 4  4  5  5

[:, :, 2] =
 3  4  3  4
 4  4  5  5

[:, :, 3] =
 2  3  2  3
 5  5  6  6

[:, :, 4] =
 3  4  3  4
 5  5  6  6

[:, :, 5] =
 4  5  4  5
 5  5  6  6

[:, :, 6] =
 4  5  4  5
 4  4  5  5
```
Note that the corners are listed in Cartesian order.

The (number of faces by number of elements) array `elemtobndy` gives the
boundary number for each face of each element.  A zero will be given for
connected faces.
```jldoctest brickmesh
julia> topology.elemtobndy
4×6 Array{Int64,2}:
 1  0  1  0  0  0
 0  0  0  0  2  2
 0  0  0  0  0  0
 0  0  0  0  0  0
```
Note that the faces are listed in Cartesian order.

"""
function BrickTopology(mpicomm, elemrange;
                       boundary=ntuple(j->(1,1), length(elemrange)),
                       periodicity=ntuple(j->false, length(elemrange)),
                       connectivity=:face, ghostsize=1)

  if boundary isa Matrix
    boundary = tuple(mapslices(x -> tuple(x...), boundary, dims=1)...)
  end

  # We cannot handle anything else right now...
  @assert connectivity == :face
  @assert ghostsize == 1

  mpirank = MPI.Comm_rank(mpicomm)
  mpisize = MPI.Comm_size(mpicomm)
  topology = BrickMesh.brickmesh(elemrange, periodicity, part=mpirank+1,
                              numparts=mpisize, boundary=boundary)
  topology = BrickMesh.partition(mpicomm, topology...)
  origsendorder = topology[5]
  topology = BrickMesh.connectmesh(mpicomm, topology[1:4]...)

  dim = length(elemrange)
  T = eltype(topology.elemtocoord)
  return BrickTopology{dim, T}(BoxElementTopology{dim, T}(
              mpicomm, topology.elems, topology.realelems,
              topology.ghostelems, topology.ghostfaces, topology.sendelems,
              topology.sendfaces, topology.elemtocoord, topology.elemtoelem,
              topology.elemtoface, topology.elemtoordr, topology.elemtobndy,
              topology.nabrtorank, topology.nabrtorecv, topology.nabrtosend,
              origsendorder, !minimum(periodicity)))
end

""" A wrapper for the StackedBrickTopology """
StackedBrickTopology(mpicomm, Nelems::NTuple{N, Integer}; kw...) where N =
StackedBrickTopology(mpicomm, map(Ne->0:Ne, Nelems); kw...)

"""
    StackedBrickTopology{dim, T}(mpicomm, elemrange; boundary, periodicity)

Generate a stacked brick mesh topology with coordinates given by the tuple
`elemrange` and the periodic dimensions given by the `periodicity` tuple.

The elements are stacked such that the elements associated with range
`elemrange[dim]` are contiguous in the element ordering.

The elements of the brick are partitioned equally across the MPI ranks based
on a space-filling curve.  Further, stacks are not split at MPI boundaries.

By default boundary faces will be marked with a one and other faces with a
zero.  Specific boundary numbers can also be passed for each face of the brick
in `boundary`.  This will mark the nonperiodic brick faces with the given
boundary number.

# Examples

We can build a 3 by 2 element two-dimensional mesh that is periodic in the
\$x2\$-direction with
```jldoctest brickmesh

using CLIMA.Topologies
using MPI
MPI.Init()
topology = StackedBrickTopology(MPI.COMM_SELF, (2:5,4:6);
                                periodicity=(false,true),
                                boundary=((1,2),(3,4)))
```
This returns the mesh structure stacked in the \$x2\$-direction for

             x2

              ^
              |
             6-  +-----+-----+-----+
              |  |     |     |     |
              |  |  2  |  4  |  6  |
              |  |     |     |     |
             5-  +-----+-----+-----+
              |  |     |     |     |
              |  |  1  |  3  |  5  |
              |  |     |     |     |
             4-  +-----+-----+-----+
              |
              +--|-----|-----|-----|--> x1
                 2     3     4     5

For example, the (dimension by number of corners by number of elements) array
`elemtocoord` gives the coordinates of the corners of each element.
```jldoctest brickmesh
julia> topology.elemtocoord
2×4×6 Array{Int64,3}:
[:, :, 1] =
 2  3  2  3
 4  4  5  5

[:, :, 2] =
 2  3  2  3
 5  5  6  6

[:, :, 3] =
 3  4  3  4
 4  4  5  5

[:, :, 4] =
 3  4  3  4
 5  5  6  6

[:, :, 5] =
 4  5  4  5
 4  4  5  5

[:, :, 6] =
 4  5  4  5
 5  5  6  6
```
Note that the corners are listed in Cartesian order.

The (number of faces by number of elements) array `elemtobndy` gives the
boundary number for each face of each element.  A zero will be given for
connected faces.
```jldoctest brickmesh
julia> topology.elemtobndy
4×6 Array{Int64,2}:
 1  0  1  0  0  0
 0  0  0  0  2  2
 0  0  0  0  0  0
 0  0  0  0  0  0
```
Note that the faces are listed in Cartesian order.
"""
function StackedBrickTopology(mpicomm, elemrange;
                       boundary=ntuple(j->(1,1), length(elemrange)),
                       periodicity=ntuple(j->false, length(elemrange)),
                       connectivity=:face, ghostsize=1)

  if boundary isa Matrix
    boundary = tuple(mapslices(x -> tuple(x...), boundary, dims=1)...)
  end

  dim = length(elemrange)

  dim <= 1 && error("Stacked brick topology works for 2D and 3D")

  # Build the base topology
  basetopo = BrickTopology(mpicomm, elemrange[1:dim-1];
                     boundary=boundary[1:dim-1],
                     periodicity=periodicity[1:dim-1],
                     connectivity=connectivity,
                     ghostsize=ghostsize)

  # Use the base topology to build the stacked topology
  stack = elemrange[dim]
  stacksize = length(stack) - 1

  nvert = 2^dim
  nface = 2dim

  nreal = length(basetopo.realelems)*stacksize
  nghost = length(basetopo.ghostelems)*stacksize

  elems=1:(nreal+nghost)
  realelems=1:nreal
  ghostelems=nreal.+(1:nghost)

  sendelems = similar(basetopo.sendelems,
                      length(basetopo.sendelems)*stacksize)
  for i=1:length(basetopo.sendelems), j=1:stacksize
    sendelems[stacksize*(i-1) + j] = stacksize*(basetopo.sendelems[i]-1) + j
  end

  ghostfaces = similar(basetopo.ghostfaces, nface, length(ghostelems))
  ghostfaces .= false

  for i=1:length(basetopo.ghostelems), j=1:stacksize
    e = stacksize*(i-1) + j
    for f = 1:2(dim-1)
      ghostfaces[f, e] = basetopo.ghostfaces[f, i]
    end
  end

  sendfaces = similar(basetopo.sendfaces, nface, length(sendelems))
  sendfaces .= false

  for i=1:length(basetopo.sendelems), j=1:stacksize
    e = stacksize*(i-1) + j
    for f = 1:2(dim-1)
      sendfaces[f, e] = basetopo.sendfaces[f, i]
    end
  end

  elemtocoord = similar(basetopo.elemtocoord, dim, nvert, length(elems))

  for i=1:length(basetopo.elems), j=1:stacksize
    e = stacksize*(i-1) + j

    for v = 1:2^(dim-1)
      for d = 1:(dim-1)
        elemtocoord[d, v, e] = basetopo.elemtocoord[d, v, i]
        elemtocoord[d, 2^(dim-1) + v, e] = basetopo.elemtocoord[d, v, i]
      end

      elemtocoord[dim, v, e] = stack[j]
      elemtocoord[dim, 2^(dim-1) + v, e] = stack[j+1]
    end
  end

  elemtoelem = similar(basetopo.elemtoelem, nface, length(elems))
  elemtoface = similar(basetopo.elemtoface, nface, length(elems))
  elemtoordr = similar(basetopo.elemtoordr, nface, length(elems))
  elemtobndy = similar(basetopo.elemtobndy, nface, length(elems))

  for e=1:length(basetopo.elems)*stacksize, f=1:nface
    elemtoelem[f, e] = e
    elemtoface[f, e] = f
    elemtoordr[f, e] = 1
    elemtobndy[f, e] = 0
  end

  for i=1:length(basetopo.realelems), j=1:stacksize
    e1 = stacksize*(i-1) + j

    for f = 1:2(dim-1)
      e2 = stacksize*(basetopo.elemtoelem[f, i]-1) + j

      elemtoelem[f, e1] = e2
      elemtoface[f, e1] = basetopo.elemtoface[f, i]

      # We assume a simple orientation right now
      @assert basetopo.elemtoordr[f, i] == 1
      elemtoordr[f, e1] = basetopo.elemtoordr[f, i]
    end

    et = stacksize*(i-1) + j + 1
    eb = stacksize*(i-1) + j - 1
    ft = 2(dim-1) + 1
    fb = 2(dim-1) + 2
    ot = 1
    ob = 1

    if j == stacksize
      et = periodicity[dim] ? stacksize*(i-1) + 1 : e1
      ft = periodicity[dim] ? ft : 2(dim-1) + 2
    end
    if j == 1
      eb = periodicity[dim] ? stacksize*(i-1) + stacksize : e1
      fb = periodicity[dim] ? fb : 2(dim-1) + 1
    end

    elemtoelem[2(dim-1)+1, e1] = eb
    elemtoelem[2(dim-1)+2, e1] = et
    elemtoface[2(dim-1)+1, e1] = fb
    elemtoface[2(dim-1)+2, e1] = ft
    elemtoordr[2(dim-1)+1, e1] = ob
    elemtoordr[2(dim-1)+2, e1] = ot
  end

  for i=1:length(basetopo.elems), j=1:stacksize
    e1 = stacksize*(i-1) + j

    for f = 1:2(dim-1)
      elemtobndy[f, e1] = basetopo.elemtobndy[f, i]
    end

    bt = bb = 0

    if j == stacksize
      bt = periodicity[dim] ? bt : boundary[dim][2]
    end
    if j == 1
      bb = periodicity[dim] ? bb : boundary[dim][1]
    end

    elemtobndy[2(dim-1)+1, e1] = bb
    elemtobndy[2(dim-1)+2, e1] = bt
  end

  nabrtorank = basetopo.nabrtorank
  nabrtorecv =
    UnitRange{Int}[UnitRange(stacksize*(first(basetopo.nabrtorecv[n])-1)+1,
                             stacksize*last(basetopo.nabrtorecv[n]))
                   for n = 1:length(nabrtorank)]
  nabrtosend =
    UnitRange{Int}[UnitRange(stacksize*(first(basetopo.nabrtosend[n])-1)+1,
                             stacksize*last(basetopo.nabrtosend[n]))
                   for n = 1:length(nabrtorank)]

  T = eltype(basetopo.elemtocoord)

  StackedBrickTopology{dim, T}(
    BoxElementTopology{dim, T}(
      mpicomm, elems, realelems, ghostelems, ghostfaces, sendelems, sendfaces,
      elemtocoord, elemtoelem, elemtoface, elemtoordr, elemtobndy,
      nabrtorank, nabrtorecv, nabrtosend, basetopo.origsendorder, !minimum(periodicity)),
    stacksize)
end

"""
    CubedShellTopology(mpicomm, Nelem, T) <: AbstractTopology{dim}

Generate a cubed shell mesh with the number of elements along each dimension of
the cubes being `Nelem`. This topology actual creates a cube mesh, and the
warping should be done after the grid is created using the `cubedshellwarp`
function. The coordinates of the points will be of type `T`.

The elements of the shell are partitioned equally across the MPI ranks based
on a space-filling curve.

Note that this topology is logically 2-D but embedded in a 3-D space

# Examples

We can build a cubed shell mesh with 10 elements on each cube, total elements is
`10 * 10 * 6 = 600`, with
```jldoctest brickmesh
using CLIMA.Topologies
using MPI
MPI.Init()
topology = CubedShellTopology(MPI.COMM_SELF, 10, Float64)

# Typically the warping would be done after the grid is created, but the cell
# corners could be warped with...

# Shell radius = 1
x1, x2, x3 = ntuple(j->topology.elemtocoord[j, :, :], 3)
for n = 1:length(x1)
   x1[n], x2[n], x3[n] = Topologies.cubedshellwarp(x1[n], x2[n], x3[n])
end

# Shell radius = 10
x1, x2, x3 = ntuple(j->topology.elemtocoord[j, :, :], 3)
for n = 1:length(x1)
  x1[n], x2[n], x3[n] = Topologies.cubedshellwarp(x1[n], x2[n], x3[n], 10)
end
```
"""
function CubedShellTopology(mpicomm, Neside, T; connectivity=:face,
                            ghostsize=1)

  # We cannot handle anything else right now...
  @assert connectivity == :face
  @assert ghostsize == 1

  mpirank = MPI.Comm_rank(mpicomm)
  mpisize = MPI.Comm_size(mpicomm)

  topology = cubedshellmesh(Neside, part=mpirank+1, numparts=mpisize)

  topology = BrickMesh.partition(mpicomm, topology...)#topology(1), topology(2), topology(3), topology(4) ) #topology...)
  origsendorder = topology[5]
  dim, nvert = 3, 4
  elemtovert = topology[1]
  nelem = size(elemtovert, 2)
  elemtocoord = Array{T}(undef, dim, nvert, nelem)
  ind2vert = CartesianIndices((Neside+1, Neside+1, Neside+1))
  for e in 1:nelem
    for n = 1:nvert
      v = elemtovert[n, e]
      i, j, k = Tuple(ind2vert[v])
      elemtocoord[:, n, e] = (2*[i-1, j-1, k-1] .- Neside) / Neside
    end
  end

  topology = BrickMesh.connectmesh(mpicomm, topology[1], elemtocoord, topology[3],
                                topology[4]; dim = 2)

  CubedShellTopology{T}(
    BoxElementTopology{2, T}(
      mpicomm, topology.elems, topology.realelems,
      topology.ghostelems, topology.ghostfaces, topology.sendelems,
      topology.sendfaces, topology.elemtocoord, topology.elemtoelem,
      topology.elemtoface, topology.elemtoordr, topology.elemtobndy,
      topology.nabrtorank, topology.nabrtorecv, topology.nabrtosend, origsendorder, false))
end

"""
    cubedshellmesh(T, Ne; part=1, numparts=1)

Generate a cubed mesh with each of the "cubes" has an `Ne X Ne` grid of
elements.

The mesh can optionally be partitioned into `numparts` and this returns
partition `part`.  This is a simple Cartesian partition and further partitioning
(e.g, based on a space-filling curve) should be done before the mesh is used for
computation.

This mesh returns the cubed spehere in a flatten fashion for the vertex values,
and a remapping is needed to embed the mesh in a 3-D space.

The mesh structures for the cubes is as follows:

```
x2
   ^
   |
4Ne-           +-------+
   |           |       |
   |           |   6   |
   |           |       |
3Ne-           +-------+
   |           |       |
   |           |   5   |
   |           |       |
2Ne-           +-------+
   |           |       |
   |           |   4   |
   |           |       |
 Ne-   +-------+-------+-------+
   |   |       |       |       |
   |   |   1   |   2   |   3   |
   |   |       |       |       |
  0-   +-------+-------+-------+
   |
   +---|-------|-------|------|-> x1
       0      Ne      2Ne    3Ne
```

"""
function cubedshellmesh(Ne; part=1, numparts=1)
  dim = 2
  @assert 1 <= part <= numparts

  globalnelems = 6 * Ne^2

  # How many vertices and faces per element
  nvert = 2^dim # 4
  nface = 2dim  # 4

  # linearly partition to figure out which elements we own
  elemlocal = BrickMesh.linearpartition(prod(globalnelems), part, numparts)

  # elemen to vertex maps which we own
  elemtovert = Array{Int}(undef, nvert, length(elemlocal))
  elemtocoord = Array{Int}(undef, dim, nvert, length(elemlocal))

  nelemcube = Ne^dim # Ne^2

  etoijb = CartesianIndices((Ne, Ne, 6))
  bx = [0 Ne 2Ne  Ne  Ne  Ne]
  by = [0  0   0  Ne 2Ne 3Ne]

  vertmap = LinearIndices((Ne+1, Ne+1, Ne+1))
  for (le, e) = enumerate(elemlocal)
    i, j, blck = Tuple(etoijb[e])
    elemtocoord[1, :, le] = bx[blck] .+ [i-1 i i-1 i]
    elemtocoord[2, :, le] = by[blck] .+ [j-1 j-1 j j]

    for n = 1:4
      ix = i + mod(n-1, 2)
      jx = j + div(n-1, 2)
      # set the vertices like they are the face vertices of a cube
      if blck == 1
        elemtovert[n, le] = vertmap[   1   , Ne+2-ix,      jx]
      elseif blck == 2
        elemtovert[n, le] = vertmap[     ix,   1    ,      jx]
      elseif blck == 3
        elemtovert[n, le] = vertmap[Ne+1   ,      ix,      jx]
      elseif blck == 4
        elemtovert[n, le] = vertmap[     ix,      jx,  Ne+1  ]
      elseif blck == 5
        elemtovert[n, le] = vertmap[     ix, Ne+1   , Ne+2-jx]
      elseif blck == 6
        elemtovert[n, le] = vertmap[     ix, Ne+2-jx,    1   ]
      end
    end
  end

  # no boundaries for a shell
  elemtobndy = zeros(Int, nface, length(elemlocal))

  # no faceconnections for a shell
  faceconnections = Array{Array{Int, 1}}(undef, 0)

  (elemtovert, elemtocoord, elemtobndy, faceconnections, collect(elemlocal))
end

"""
    cubedshellwarp(a, b, c, R = max(abs(a), abs(b), abs(c)))

Given points `(a, b, c)` on the surface of a cube, warp the points out to a
spherical shell of radius `R` based on the equiangular gnomonic grid proposed by
Ronchi, Iacono, Paolucci (1996) <https://doi.org/10.1006/jcph.1996.0047>

```
@article{RonchiIaconoPaolucci1996,
  title={The ``cubed sphere'': a new method for the solution of partial
         differential equations in spherical geometry},
  author={Ronchi, C. and Iacono, R. and Paolucci, P. S.},
  journal={Journal of Computational Physics},
  volume={124},
  number={1},
  pages={93--114},
  year={1996},
  doi={10.1006/jcph.1996.0047}
}
```

"""
function cubedshellwarp(a, b, c, R = max(abs(a), abs(b), abs(c)))

  function f(sR, ξ, η)
    X, Y = tan(π * ξ / 4), tan(π * η / 4)
    x1 = sR / sqrt(X^2 + Y^2 + 1)
    x2, x3 = X * x1, Y * x1
    x1,x2,x3
  end

  fdim = argmax(abs.((a, b, c)))
  if fdim == 1 && a < 0
    # (-R, *, *) : Face I from Ronchi, Iacono, Paolucci (1996)
    x1,x2,x3 = f(-R, b/a, c/a)
  elseif fdim == 2 && b < 0
    # ( *,-R, *) : Face II from Ronchi, Iacono, Paolucci (1996)
    x2,x1,x3 = f(-R, a/b, c/b)
  elseif fdim == 1 && a > 0
    # ( R, *, *) : Face III from Ronchi, Iacono, Paolucci (1996)
    x1,x2,x3 = f(R, b/a, c/a)
  elseif fdim == 2 && b > 0
    # ( *, R, *) : Face IV from Ronchi, Iacono, Paolucci (1996)
    x2,x1,x3 = f(R, a/b, c/b)
  elseif fdim == 3 && c > 0
    # ( *, *, R) : Face V from Ronchi, Iacono, Paolucci (1996)
    x3,x2,x1 = f(R, b/c, a/c)
  elseif fdim == 3 && c < 0
    # ( *, *,-R) : Face VI from Ronchi, Iacono, Paolucci (1996)
    x3,x2,x1 = f(-R, b/c, a/c)
  else
    error("invalid case for cubedshellwarp: $a, $b, $c")
  end

  return x1, x2, x3
end

"""
    cubedshellunwarp(x1, x2, x3)

The inverse of [`cubedshellwarp`](@ref).
"""
function cubedshellunwarp(x1,x2,x3)

  function g(R, X, Y)
    ξ = atan(X) * 4 / pi
    η = atan(Y) * 4 / pi
    R, R*ξ, R*η
  end

  R = hypot(x1,x2,x3)
  fdim = argmax(abs.((x1, x2, x3)))

  if fdim == 1 && x1 < 0
    # (-R, *, *) : Face I from Ronchi, Iacono, Paolucci (1996)
    a,b,c = g(-R, x2/x1, x3/x1)
  elseif fdim == 2 && x2 < 0
    # ( *,-R, *) : Face II from Ronchi, Iacono, Paolucci (1996)
    b,a,c = g(-R, x1/x2, x3/x2)
  elseif fdim == 1 && x1 > 0
    # ( R, *, *) : Face III from Ronchi, Iacono, Paolucci (1996)
    a,b,c = g(R, x2/x1, x3/x1)
  elseif fdim == 2 && x2 > 0
    # ( *, R, *) : Face IV from Ronchi, Iacono, Paolucci (1996)
    b,a,c = g(R, x1/x2, x3/x2)
  elseif fdim == 3 && x3 > 0
    # ( *, *, R) : Face V from Ronchi, Iacono, Paolucci (1996)
    c,b,a = g(R, x2/x3, x1/x3)
  elseif fdim == 3 && x3 < 0
    # ( *, *,-R) : Face VI from Ronchi, Iacono, Paolucci (1996)
    c,b,a = g(-R, x2/x3, x1/x3)
  else
    error("invalid case for cubedshellunwarp: $a, $b, $c")
  end

  return a,b,c
end

"""
   StackedCubedSphereTopology(mpicomm, Nhorz, Rrange;
                              boundary=(1,1)) <: AbstractTopology{3}

Generate a stacked cubed sphere topology with `Nhorz` by `Nhorz` cells for each
horizontal face and `Rrange` is the radius edges of the stacked elements.  This
topology actual creates a cube mesh, and the warping should be done after the
grid is created using the `cubedshellwarp` function. The coordinates of the
points will be of type `eltype(Rrange)`. The inner boundary condition type is
`boundary[1]` and the outer boundary condition type is `boundary[2]`.

The elements are stacked such that the vertical elements are contiguous in the
element ordering.

The elements of the brick are partitioned equally across the MPI ranks based
on a space-filling curve. Further, stacks are not split at MPI boundaries.

# Examples

We can build a cubed sphere mesh with 10 x 10 x 5 elements on each cube, total
elements is `10 * 10 * 5 * 6 = 3000`, with
```jldoctest brickmesh
using CLIMA.Topologies
using MPI
MPI.Init()
Nhorz = 10
Nstack = 5
Rrange = Float64.(accumulate(+,1:Nstack+1))
topology = StackedCubedSphereTopology(MPI.COMM_SELF, Nhorz, Rrange)

x1, x2, x3 = ntuple(j->reshape(topology.elemtocoord[j, :, :],
                            2, 2, 2, length(topology.elems)), 3)
for n = 1:length(x1)
   x1[n], x2[n], x3[n] = Topologies.cubedshellwarp(x1[n], x2[n], x3[n])
end
```
Note that the faces are listed in Cartesian order.
"""
function StackedCubedSphereTopology(mpicomm, Nhorz, Rrange; boundary = (1, 1),
                                    connectivity=:face, ghostsize=1)
  T = eltype(Rrange)

  basetopo = CubedShellTopology(mpicomm, Nhorz, T; connectivity=connectivity,
                                   ghostsize=ghostsize)

  dim = 3
  nvert = 2^dim
  nface = 2dim
  stacksize = length(Rrange) - 1

  nreal = length(basetopo.realelems)*stacksize
  nghost = length(basetopo.ghostelems)*stacksize

  elems=1:(nreal+nghost)
  realelems=1:nreal
  ghostelems=nreal.+(1:nghost)

  sendelems = similar(basetopo.sendelems,
                      length(basetopo.sendelems)*stacksize)

  for i=1:length(basetopo.sendelems), j=1:stacksize
    sendelems[stacksize*(i-1) + j] = stacksize*(basetopo.sendelems[i]-1) + j
  end

  ghostfaces = similar(basetopo.ghostfaces, nface, length(ghostelems))
  ghostfaces .= false

  for i=1:length(basetopo.ghostelems), j=1:stacksize
    e = stacksize*(i-1) + j
    for f = 1:2(dim-1)
      ghostfaces[f, e] = basetopo.ghostfaces[f, i]
    end
  end

  sendfaces = similar(basetopo.sendfaces, nface, length(sendelems))
  sendfaces .= false

  for i=1:length(basetopo.sendelems), j=1:stacksize
    e = stacksize*(i-1) + j
    for f=1:2(dim-1)
      sendfaces[f, e] = basetopo.sendfaces[f, i]
    end
  end

  elemtocoord = similar(basetopo.elemtocoord, dim, nvert, length(elems))

  for i=1:length(basetopo.elems), j=1:stacksize
    # i is base element
    # e is stacked element
    e = stacksize*(i-1) + j

    # v is base vertex
    for v = 1:2^(dim-1)
      for d = 1:dim # dim here since shell is embedded in 3-D
        # v is lower stacked vertex
        elemtocoord[d, v, e] = basetopo.elemtocoord[d, v, i] * Rrange[j]
        # 2^(dim-1) + v is higher stacked vertex
        elemtocoord[d, 2^(dim-1) + v, e] = basetopo.elemtocoord[d, v, i] *
                                           Rrange[j+1]
      end
    end
  end

  elemtoelem = similar(basetopo.elemtoelem, nface, length(elems))
  elemtoface = similar(basetopo.elemtoface, nface, length(elems))
  elemtoordr = similar(basetopo.elemtoordr, nface, length(elems))
  elemtobndy = similar(basetopo.elemtobndy, nface, length(elems))

  for e=1:length(basetopo.elems)*stacksize, f=1:nface
    elemtoelem[f, e] = e
    elemtoface[f, e] = f
    elemtoordr[f, e] = 1
    elemtobndy[f, e] = 0
  end

  for i=1:length(basetopo.realelems), j=1:stacksize
    e1 = stacksize*(i-1) + j

    for f = 1:2(dim-1)
      e2 = stacksize*(basetopo.elemtoelem[f, i]-1) + j

      elemtoelem[f, e1] = e2
      elemtoface[f, e1] = basetopo.elemtoface[f, i]

      # since the basetopo is 2-D we only need to worry about two orientations
      @assert basetopo.elemtoordr[f, i] ∈ (1,2)
      #=
      orientation 1:
      2---3     2---3
      |   | --> |   |
      0---1     0---1
      same:
      (a,b) --> (a,b)

      orientation 3:
      2---3     3---2
      |   | --> |   |
      0---1     1---0
      reverse first index:
      (a,b) --> (N+1-a,b)
      =#
      elemtoordr[f, e1] = basetopo.elemtoordr[f, i] == 1 ? 1 : 3
    end

    # If top or bottom of stack set neighbor to self on respective face
    elemtoelem[2(dim-1)+1, e1] = j == 1         ? e1 : stacksize*(i-1) + j - 1
    elemtoelem[2(dim-1)+2, e1] = j == stacksize ? e1 : stacksize*(i-1) + j + 1

    elemtoface[2(dim-1)+1, e1] = j == 1         ? 2(dim-1) + 1 : 2(dim-1) + 2
    elemtoface[2(dim-1)+2, e1] = j == stacksize ? 2(dim-1) + 2 : 2(dim-1) + 1

    elemtoordr[2(dim-1)+1, e1] = 1
    elemtoordr[2(dim-1)+2, e1] = 1
  end

  # Set the top and bottom boundary condition
  for i=1:length(basetopo.elems)
    eb = stacksize*(i-1) + 1
    et = stacksize*(i-1) + stacksize

    elemtobndy[2(dim-1)+1, eb] = boundary[1]
    elemtobndy[2(dim-1)+2, et] = boundary[2]
  end

  nabrtorank = basetopo.nabrtorank
  nabrtorecv =
    UnitRange{Int}[UnitRange(stacksize*(first(basetopo.nabrtorecv[n])-1)+1,
                             stacksize*last(basetopo.nabrtorecv[n]))
                   for n = 1:length(nabrtorank)]
  nabrtosend =
    UnitRange{Int}[UnitRange(stacksize*(first(basetopo.nabrtosend[n])-1)+1,
                             stacksize*last(basetopo.nabrtosend[n]))
                   for n = 1:length(nabrtorank)]

  StackedCubedSphereTopology{T}(
    BoxElementTopology{3, T}(
      mpicomm, elems, realelems, ghostelems, ghostfaces, sendelems,
      sendfaces, elemtocoord, elemtoelem, elemtoface, elemtoordr, elemtobndy,
      nabrtorank, nabrtorecv, nabrtosend, basetopo.origsendorder, true),
    stacksize)
end

"""
    grid1d(a, b[, stretch::AbstractGridStretching]; elemsize, nelem)

Discretize the 1D interval [`a`,`b`] into elements.
Exactly one of the following keyword arguments must be provided:
- `elemsize`: the average element size, or
- `nelem`: the number of elements.

The optional `stretch` argument allows stretching, otherwise the element sizes will be uniform.

Returns either a range object or a vector containing the element boundaries.
"""
function grid1d(a, b, stretch=nothing; elemsize=nothing, nelem=nothing)
  xor(nelem === nothing, elemsize === nothing) || error("Either `elemsize` or `nelem` arguments must be provided")
  if elemsize !== nothing
    nelem = round(Int,abs(b-a)/elemsize)
  end
  grid1d(a, b, stretch, nelem)
end
function grid1d(a, b, ::Nothing, nelem)
  range(a, stop=b, length=nelem+1)
end

# TODO: document these
abstract type AbstractGridStretching end

"""
    SingleExponentialStretching(A)

Apply single-exponential stretching: `A > 0` will increase the density of points at the lower boundary, `A < 0` will increase the density at the upper boundary.

# Reference
* "Handbook of Grid Generation" J. F. Thompson, B. K. Soni, N. P. Weatherill (Editors) RCR Press 1999, §3.6.1 Single-Exponential Function
"""
struct SingleExponentialStretching{T} <: AbstractGridStretching
  A::T
end
function grid1d(a::A, b::B, stretch::SingleExponentialStretching, nelem) where {A,B}
  F = float(promote_type(A,B))
  s = range(zero(F), stop=one(F), length=nelem+1)
  a .+ (b-a) .* expm1.(stretch.A .* s) ./ expm1(stretch.A)
end

struct InteriorStretching{T} <: AbstractGridStretching
  attractor::T
end
function grid1d(a::A, b::B, stretch::InteriorStretching, nelem) where {A,B}
  F = float(promote_type(A,B))
  coe = F(2.5)
  s = range(zero(F), stop=one(F), length=nelem+1)
  range(a, stop=b, length=nelem+1) .+ coe .* (stretch.attractor .- (b-a).*s) .* (1 .-s) .* s
end

end


# round to nearest ties to even
round(x::Posit8) = ccall((:p8_roundToInt, SoftPositPath), Posit8, (Posit8,),x)
round(x::Posit16) = ccall((:p16_roundToInt, SoftPositPath), Posit16, (Posit16,),x)
round(x::Posit32) = ccall((:p32_roundToInt, SoftPositPath), Posit32, (Posit32,),x)

# for Posits with 1 exponent bit
for T in (:Posit8_1, :Posit16_1, :Posit24_1)
    @eval begin
        round(x::$T) =  ccall((:pX1_roundToInt, SoftPositPath), $T, ($T,),x)
    end
end

# for Posits with 1 exponent bit
for T in (:Posit8_2, :Posit16_2, :Posit24_2)
    @eval begin
        round(x::$T) =  ccall((:pX2_roundToInt, SoftPositPath), $T, ($T,),x)
    end
end

function round(x::T, r::RoundingMode{:Down}) where {T<:AbstractPosit}
    xr = round(x)
    if xr > x
        return xr-one(T)
    else
        return xr
    end
end

function round(x::T, r::RoundingMode{:Up}) where {T<:AbstractPosit}
    xr = round(x)
    if xr < x
        return xr+one(T)
    else
        return xr
    end
end


# utility functions needed for discerning topology from geometry

using Combinatorics

using dataset_pb2

export bohr_to_angstroms, angstroms_to_bohr, distance_between_atoms, bonded, distances
export add_atom!, add_bond!
export smu_atype_to_atomic_number, number_to_smu_btype, smu_btype_to_number, smu_atom_type_to_max_con
export atomic_number_to_smu_atom_type
export add_atom_by_atomic_number!, add_atom_by_type!
export distances, bonded
export canonical_bond_topology!, same_bond_topology
export is_single_fragment

function bohr_to_angstroms(distance)
  return distance * 0.529177249
end

function angstroms_to_bohr(distance)
  return distance / 0.529177249
end

function distance_between_atoms(geom::dataset_pb2.Geometry, a1, a2)::Float32
  return bohr_to_angstroms(sqrt(
    (geom.atom_positions[a1].x - geom.atom_positions[a2].x) *
    (geom.atom_positions[a1].x - geom.atom_positions[a2].x) +
    (geom.atom_positions[a1].y - geom.atom_positions[a2].y) *
    (geom.atom_positions[a1].y - geom.atom_positions[a2].y) +
    (geom.atom_positions[a1].z - geom.atom_positions[a2].z) *
    (geom.atom_positions[a1].z - geom.atom_positions[a2].z)
  ))
end

_smu_atype_to_atomic_number = Dict{Int32, Int32}(
  BondTopology_AtomType.ATOM_H => 1,
  BondTopology_AtomType.ATOM_C => 6,
  BondTopology_AtomType.ATOM_N => 7,
  BondTopology_AtomType.ATOM_NPOS => 7,
  BondTopology_AtomType.ATOM_O => 8,
  BondTopology_AtomType.ATOM_ONEG => 8,
  BondTopology_AtomType.ATOM_F => 9
)

function smu_atype_to_atomic_number(smu_type::T)::Int32 where {T<:Integer}
  return _smu_atype_to_atomic_number[smu_type]
end

"""Given an atomic number and charge, return the BondTopology_AtomType
Args:
  atomic_number:
  charge:
Returns:
  BondTopology_AtomType
"""
function atomic_number_to_smu_atom_type(atomic_number::Integer, charge::Integer=0)::Int32
  if atomic_number == 6
    return BondTopology_AtomType.ATOM_C
  elseif atomic_number == 7
    if charge == 0
      return BondTopology_AtomType.ATOM_N
    else
      return BondTopology_AtomType.ATOM_NPOS
    end
  elseif atomic_number == 8
    if charge == 0
      return BondTopology_AtomType.ATOM_O
    else
      return BondTopology_AtomType.ATOM_ONEG
    end
  elseif atomic_number == 9
    return BondTopology_AtomType.ATOM_F
  elseif atomic_number == 1
    return BondTopology_AtomType.ATOM_H
  else
    @error("Unrecognized atomic number $(atomic_number)")
  end
end

_number_to_smu_btype = Dict{Int32, Int32}(
  0 => BondTopology_BondType.BOND_UNDEFINED,
  1 => BondTopology_BondType.BOND_SINGLE,
  2 => BondTopology_BondType.BOND_DOUBLE,
  3 => BondTopology_BondType.BOND_TRIPLE
)

function number_to_smu_btype(n::T)::Int32 where {T<:Integer}
  return _number_to_smu_btype[n]
end

_smu_btype_to_number = Dict{Int32, Int32}(
  BondTopology_BondType.BOND_UNDEFINED => 0,
  BondTopology_BondType.BOND_SINGLE => 1,
  BondTopology_BondType.BOND_DOUBLE => 2,
  BondTopology_BondType.BOND_TRIPLE => 3
)

_smu_atom_type_to_max_con = Dict{Int32, Int32}(
  BondTopology_AtomType.ATOM_H => 1,
  BondTopology_AtomType.ATOM_C => 4,
  BondTopology_AtomType.ATOM_N => 3,
  BondTopology_AtomType.ATOM_NPOS => 3,
  BondTopology_AtomType.ATOM_O => 2,
  BondTopology_AtomType.ATOM_ONEG => 1,
  BondTopology_AtomType.ATOM_F => 1
)

# Return the maximum number of bonds for `smu_atype`.
# Unfortunately the type declaration on the next line does not work.
# function smu_atom_type_to_max_con(smu_atype::BondTopology_AtomType)::Int32

function smu_atom_type_to_max_con(smu_atype::T)::Int32 where {T <:Integer}
  return _smu_atom_type_to_max_con[smu_atype]
end

function smu_btype_to_number(smu_btype::T)::Int32 where {T<:Integer}
  return _smu_btype_to_number[smu_btype]
end

function add_bond!(a1::T, a2::T, btype, bond_topology::dataset_pb2.BondTopology) where {T<:Integer}
  smu_btype = number_to_smu_btype(btype)

  new_bond = BondTopology_Bond(atom_a=a1, atom_b=a2, bond_type=smu_btype)

  if !hasproperty(bond_topology, :bonds)
    setproperty!(bond_topology, :bonds, [new_bond])
  else
    push!(bond_topology.bonds, new_bond)
  end
end

function add_atom_by_atomic_number!(bond_topology::dataset_pb2.BondTopology, atomic_number::T, charge::T=0) where {T<:Integer}
  smu_atype = atomic_number_to_smu_atom_type(atomic_number, charge)
  add_atom_by_type!(bond_topology, smu_atype)
end

function add_atom_by_type!(bond_topology::dataset_pb2.BondTopology, smu_atype::T) where {T <:Integer}
  if !hasproperty(bond_topology, :atoms)
    setproperty!(bond_topology, :atoms, [smu_atype])
  else
    push!(bond_topology.atoms, smu_atype)
  end
end

"""Return an int array of the bonded atoms in `bond_topology`.
The returned array is square[natoms,natoms]. If there is no bond
between the atoms, the value is zero. If the atoms are bonded,
the entry will be the bond type.
Args:
  bond_topology:
Returns:
  Vector of BondType's
"""
function bonded(bond_topology::dataset_pb2.BondTopology)::Array{Int32, 2}
  natoms = length(bond_topology.atoms)
  connected = zeros(Int32, (natoms, natoms))  # to be returned
  # Note need to convert to 1 indexing for Julia
  for bond in bond_topology.bonds
    a1 = bond.atom_a + 1
    a2 = bond.atom_b + 1
    connected[a1, a2] = connected[a2, a1] = bond.bond_type
  end
  return connected
end

"""Given a connection matrix `bonded`, return a vector of
vectors, one for each atom, containing the indices of the
bonded atoms.
Args:
  bonded:
Returns:
"""
function connections(bonded::Array{Int32, 2})::Vector{Vector{Int32}}
  natoms = size(bonded, 1)
  result = Vector{Vector{Int32}}(undef, (natoms))
  for i in 1:natoms
    nbrs = Vector{Int32}()
    for j in 1:natoms
      if bonded[i,j] > 0
        push!(nbrs, j)
      end
    end
    result[i] = nbrs
  end
  return result
end

"""Return a float array of the interatomic distances in `geometry`.
Args:
  geometry:
Returns:
  a numpy array of distances
"""
function distances(geometry)::Array{Float32, 2}
  natoms = length(geometry.atom_positions)
  result = zeros(Float32, (natoms, natoms))
  for atoms in combinations(1:natoms, 2)
    i = atoms[1]
    j = atoms[2]
    result[i, j] = result[j, i] = distance_between_atoms(geometry, i, j)
  end
  return result
end

"""Transform the bonds attribute of `bond_topology` to a canonical form.

Args:
  bond_topology:
Returns:
  BondTopology
"""
function canonical_bond_topology!(bond_topology)
  if length(bond_topology.bonds) < 2
    return
  end

  for bond in bond_topology.bonds
    if bond.atom_a > bond.atom_b
      bond.atom_a, bond.atom_b = bond.atom_b, bond.atom_a
    end
  end

  sort!(bond_topology.bonds, by = b -> (b.atom_a, b.atom_b))
end

Base.:(==)(b1::BondTopology_Bond, b2::BondTopology_Bond)::Bool = b1.atom_a == b2.atom_a &&
                b1.atom_b == b2.atom_b && b1.bond_type == b2.bond_type

"""Return True if bt1 == bt2.
Note that there is no attempt to canonialise the protos.
Args:
Returns:
"""
function same_bond_topology(bt1, bt2)::Bool
  length(bt2.atoms) == length(bt2.atoms) || return false
  length(bt2.bonds) == length(bt2.bonds) || return false

  # https://stackoverflow.com/questions/47564825/check-if-all-the-elements-of-a-julia-array-are-equal
  # all(y->y==x[1], x);

  atoms1 = bt1.atoms
  if ! all(atoms1->atoms1==bt2.atoms[1], bt2.atoms)
    return false
  end

  # This should work, but raises an exception.
  # bonds1 = bt1.bonds
  # bonds2 = bt2.bonds
  # return all(bonds1->bonds1==bonds2[1], bonds2)

  for i in 1:length(bt1.bonds)
    bt1.bonds[i] != bt2.bonds[i] && return false
  end
  return true
end

"""Recusrively visit nodes in the graph defined by `nbrs`.
Args:
  nbrs:
  atom:
  visited:
Returns:
  The number of nodes visited - including `atom`.
"""
function visit(nbrs::Vector, atom::T, visited::Vector{Bool})::Int32 where {T<:Integer}
  visited[atom] = true
  result = 1    # To be returned.
  for nbr in nbrs[atom]
    if visited[nbr] > 0
      continue
    end
    result += visit(nbrs, nbr, visited)
  end

  return result
end

"""Return True if `bond_topology` is a single fragment.
Args:
  bond_topology:
Returns:
  True if `bond_topology` is a single fragment.
"""
function is_single_fragment(bond_topology::BondTopology)::Bool
  natoms = length(bond_topology.atoms)
  nbonds = length(bond_topology.bonds)
  # Some special cases are easy.
  natoms == 1 && return true
  if natoms == 2 && nbonds == 1
    return true
  end
  if natoms == 3 && nbonds == 2
    return true
  end
  if natoms == nbonds && natoms <= 4
    return true
  end

  attached = connections(bonded(bond_topology))
  # Any atom with zero neighbors means a detached atom.
  any(n->length(n) == 0, attached) && return false

  visited = zeros(Bool, natoms)
  # Mark anything with a single connection as visited.
  # Record the index of an atom that has multiple connections.
  a_multiply_connected_atom = -1
  for i in 1:natoms
    if bond_topology.atoms[i] == BondTopology_AtomType.ATOM_H
      visited[i] = 1
      continue
    end

    if length(attached[i]) > 1
      a_multiply_connected_atom = i
      continue
    end

    # A singly connected heavy atom. Mark visited if not part of a two atom fragment.
    if length(attached[attached[i][1]]) > 1
      visited[i] = 1
    end
  end

  # Not sure this can happen.
  a_multiply_connected_atom < 0 && return false

  number_visited = sum(visited) + visit(attached, a_multiply_connected_atom, visited)
  return number_visited == natoms
end


"""
Construct model for HydroGen with FixedOutput Formulation
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, FixedOutput},
    ::Type{S},
) where {H <: PSY.HydroGen, S <: PM.AbstractPowerModel}
    devices = get_available_components(H, sys)

    add_parameters!(container, ActivePowerTimeSeriesParameter, devices, model)
    add_parameters!(container, ReactivePowerTimeSeriesParameter, devices, model)

    # Expression
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerTimeSeriesParameter,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ReactivePowerBalance,
        ReactivePowerTimeSeriesParameter,
        devices,
        model,
        S,
    )
    return
end

function construct_device!(
    ::OptimizationContainer,
    ::PSY.System,
    ::ModelConstructStage,
    ::DeviceModel{H, FixedOutput},
    ::Type{S},
) where {H <: PSY.HydroGen, S <: PM.AbstractPowerModel}
    # FixedOutput doesn't add any constraints to the model. This function covers
    # AbstractPowerModel and AbtractActivePowerModel
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, FixedOutput},
    ::Type{S},
) where {H <: PSY.HydroGen, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(H, sys)

    add_parameters!(container, ActivePowerTimeSeriesParameter, devices, model)

    # Expression
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerTimeSeriesParameter,
        devices,
        model,
        S,
    )
    return
end

"""
Construct model for HydroGen with RunOfRiver Dispatch Formulation
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {
    H <: PSY.HydroGen,
    D <: AbstractHydroDispatchFormulation,
    S <: PM.AbstractPowerModel,
}
    devices = get_available_components(H, sys)

    add_variables!(container, ActivePowerVariable, devices, D())
    add_variables!(container, ReactivePowerVariable, devices, D())
    add_variables!(container, EnergyOutput, devices, D())
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ReactivePowerBalance,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )

    add_parameters!(container, ActivePowerTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )

    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {
    H <: PSY.HydroGen,
    D <: AbstractHydroDispatchFormulation,
    S <: PM.AbstractPowerModel,
}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    add_constraints!(
        container,
        ReactivePowerVariableLimitsConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )
    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)
    add_constraint_dual!(container, sys, model)

    return
end

"""
Construct model for HydroGen with RunOfRiver Dispatch Formulation
with only Active Power.
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {
    H <: PSY.HydroGen,
    D <: AbstractHydroDispatchFormulation,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(H, sys)

    add_variables!(container, ActivePowerVariable, devices, D())
    add_variables!(container, EnergyOutput, devices, D())
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )

    add_parameters!(container, ActivePowerTimeSeriesParameter, devices, model)
    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {
    H <: PSY.HydroGen,
    D <: AbstractHydroDispatchFormulation,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)

    add_constraint_dual!(container, sys, model)
    return
end

"""
Construct model for HydroGen with ReservoirBudget Dispatch Formulation
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, HydroDispatchReservoirBudget},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractPowerModel}
    devices = get_available_components(H, sys)

    add_variables!(container, ActivePowerVariable, devices, HydroDispatchReservoirBudget())
    add_variables!(
        container,
        ReactivePowerVariable,
        devices,
        HydroDispatchReservoirBudget(),
    )
    add_variables!(container, EnergyOutput, devices, HydroDispatchReservoirBudget())
    add_parameters!(container, EnergyBudgetTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ReactivePowerBalance,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, HydroDispatchReservoirBudget},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractPowerModel}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    add_constraints!(
        container,
        ReactivePowerVariableLimitsConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )

    # Energy Budget Constraint
    add_constraints!(container, EnergyBudgetConstraint, devices, model, S)

    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)

    add_constraint_dual!(container, sys, model)
    return
end

"""
Construct model for HydroGen with ReservoirBudget Dispatch Formulation
with only Active Power.
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, HydroDispatchReservoirBudget},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(H, sys)

    add_variables!(container, ActivePowerVariable, devices, HydroDispatchReservoirBudget())
    add_variables!(container, EnergyOutput, devices, HydroDispatchReservoirBudget())
    add_parameters!(container, EnergyBudgetTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )

    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, HydroDispatchReservoirBudget},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    # Energy Budget Constraint
    add_constraints!(container, EnergyBudgetConstraint, devices, model, S)

    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)
    add_constraint_dual!(container, sys, model)

    return
end

"""
Construct model for HydroGen with ReservoirStorage Dispatch Formulation
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, HydroDispatchReservoirStorage},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractPowerModel}
    devices = get_available_components(H, sys)

    add_variables!(container, ActivePowerVariable, devices, HydroDispatchReservoirStorage())
    add_variables!(
        container,
        ReactivePowerVariable,
        devices,
        HydroDispatchReservoirStorage(),
    )
    add_variables!(container, EnergyVariable, devices, HydroDispatchReservoirStorage())
    add_variables!(
        container,
        WaterSpillageVariable,
        devices,
        HydroDispatchReservoirStorage(),
    )
    add_variables!(
        container,
        EnergyShortageVariable,
        devices,
        HydroDispatchReservoirStorage(),
    )
    add_variables!(
        container,
        EnergySurplusVariable,
        devices,
        HydroDispatchReservoirStorage(),
    )
    add_variables!(container, EnergyOutput, devices, HydroDispatchReservoirStorage())
    add_parameters!(container, EnergyTargetTimeSeriesParameter, devices, model)
    add_parameters!(container, InflowTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ReactivePowerBalance,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, HydroDispatchReservoirStorage},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractPowerModel}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    add_constraints!(
        container,
        ReactivePowerVariableLimitsConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )

    add_initial_condition!(
        container,
        devices,
        HydroDispatchReservoirStorage(),
        InitialEnergyLevel(),
    )
    # Energy Balance Constraint
    add_constraints!(container, EnergyBalanceConstraint, devices, model, S)
    add_constraints!(container, EnergyTargetConstraint, devices, model, S)
    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)

    add_constraint_dual!(container, sys, model)
    return
end

"""
Construct model for HydroGen with ReservoirStorage Dispatch Formulation
with only Active Power
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, HydroDispatchReservoirStorage},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(H, sys)

    add_variables!(container, ActivePowerVariable, devices, HydroDispatchReservoirStorage())
    add_variables!(container, EnergyVariable, devices, HydroDispatchReservoirStorage())
    add_variables!(
        container,
        WaterSpillageVariable,
        devices,
        HydroDispatchReservoirStorage(),
    )
    add_variables!(
        container,
        EnergyShortageVariable,
        devices,
        HydroDispatchReservoirStorage(),
    )
    add_variables!(
        container,
        EnergySurplusVariable,
        devices,
        HydroDispatchReservoirStorage(),
    )
    add_variables!(container, EnergyOutput, devices, HydroDispatchReservoirStorage())
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )

    add_parameters!(container, EnergyTargetTimeSeriesParameter, devices, model)
    add_parameters!(container, InflowTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )

    add_feedforward_arguments!(container, model, devices)

    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, HydroDispatchReservoirStorage},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    add_initial_condition!(
        container,
        devices,
        HydroDispatchReservoirStorage(),
        InitialEnergyLevel(),
    )
    # Energy Balance Constraint
    add_constraints!(container, EnergyBalanceConstraint, devices, model, S)
    add_constraints!(container, EnergyTargetConstraint, devices, model, S)
    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)
    add_constraint_dual!(container, sys, model)

    return
end

"""
Construct model for HydroGen with RunOfRiver Commitment Formulation
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {H <: PSY.HydroGen, D <: HydroCommitmentRunOfRiver, S <: PM.AbstractPowerModel}
    devices = get_available_components(H, sys)

    add_variables!(container, ActivePowerVariable, devices, D())
    add_variables!(container, ReactivePowerVariable, devices, D())
    add_variables!(container, OnVariable, devices, D())
    add_variables!(container, EnergyOutput, devices, D())
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )

    add_to_expression!(
        container,
        ReactivePowerBalance,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )

    add_parameters!(container, ActivePowerTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {H <: PSY.HydroGen, D <: HydroCommitmentRunOfRiver, S <: PM.AbstractPowerModel}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    add_constraints!(
        container,
        ReactivePowerVariableLimitsConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )

    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)

    add_constraint_dual!(container, sys, model)
    return
end

"""
Construct model for HydroGen with RunOfRiver Commitment Formulation
with only Active Power.
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {
    H <: PSY.HydroGen,
    D <: HydroCommitmentRunOfRiver,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(H, sys)

    add_variables!(container, ActivePowerVariable, devices, D())
    add_variables!(container, OnVariable, devices, D())
    add_variables!(container, EnergyOutput, devices, D())
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_parameters!(container, ActivePowerTimeSeriesParameter, devices, model)
    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )

    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {
    H <: PSY.HydroGen,
    D <: HydroCommitmentRunOfRiver,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)

    add_constraint_dual!(container, sys, model)
    return
end

"""
Construct model for HydroGen with ReservoirBudget Commitment Formulation
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {H <: PSY.HydroGen, D <: HydroCommitmentReservoirBudget, S <: PM.AbstractPowerModel}
    devices = get_available_components(H, sys)

    add_variables!(container, ActivePowerVariable, devices, D())
    add_variables!(container, ReactivePowerVariable, devices, D())
    add_variables!(container, OnVariable, devices, D())
    add_variables!(container, EnergyOutput, devices, D())
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerBalance,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )

    add_parameters!(container, EnergyBudgetTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {H <: PSY.HydroGen, D <: HydroCommitmentReservoirBudget, S <: PM.AbstractPowerModel}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    add_constraints!(
        container,
        ReactivePowerVariableLimitsConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )
    # Energy Budget Constraint
    add_constraints!(container, EnergyBudgetConstraint, devices, model, S)

    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)
    add_constraint_dual!(container, sys, model)

    return
end

"""
Construct model for HydroGen with ReservoirBudget Commitment Formulation
with only Active Power.
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {
    H <: PSY.HydroGen,
    D <: HydroCommitmentReservoirBudget,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(H, sys)

    add_variables!(container, ActivePowerVariable, devices, D())
    add_variables!(container, OnVariable, devices, D())
    add_variables!(container, EnergyOutput, devices, D())
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )

    add_parameters!(container, EnergyBudgetTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)
    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_feedforward_arguments!(container, model, devices)

    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {
    H <: PSY.HydroGen,
    D <: HydroCommitmentReservoirBudget,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    # Energy Budget Constraint
    add_constraints!(container, EnergyBudgetConstraint, devices, model, S)

    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)

    add_constraint_dual!(container, sys, model)
    return
end

"""
Construct model for HydroGen with ReservoirStorage Commitment Formulation
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, HydroCommitmentReservoirStorage},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractPowerModel}
    devices = get_available_components(H, sys)

    add_variables!(
        container,
        ActivePowerVariable,
        devices,
        HydroCommitmentReservoirStorage(),
    )
    add_variables!(
        container,
        ReactivePowerVariable,
        devices,
        HydroCommitmentReservoirStorage(),
    )
    add_variables!(container, OnVariable, devices, HydroCommitmentReservoirStorage())
    add_variables!(container, EnergyVariable, devices, HydroCommitmentReservoirStorage())
    add_variables!(
        container,
        WaterSpillageVariable,
        devices,
        HydroCommitmentReservoirStorage(),
    )
    add_variables!(
        container,
        EnergyShortageVariable,
        devices,
        HydroCommitmentReservoirStorage(),
    )
    add_variables!(
        container,
        EnergySurplusVariable,
        devices,
        HydroCommitmentReservoirStorage(),
    )
    add_variables!(container, EnergyOutput, devices, HydroCommitmentReservoirStorage())
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerBalance,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )

    add_parameters!(container, EnergyTargetTimeSeriesParameter, devices, model)
    add_parameters!(container, InflowTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, HydroCommitmentReservoirStorage},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractPowerModel}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    add_constraints!(
        container,
        ReactivePowerVariableLimitsConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )

    add_initial_condition!(
        container,
        devices,
        HydroCommitmentReservoirStorage(),
        InitialEnergyLevel(),
    )
    # Energy Balance Constraint
    add_constraints!(container, EnergyBalanceConstraint, devices, model, S)
    add_constraints!(container, EnergyTargetConstraint, devices, model, S)
    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)
    add_constraint_dual!(container, sys, model)

    return
end

"""
Construct model for HydroGen with ReservoirStorage Dispatch Formulation
with only Active Power
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, HydroCommitmentReservoirStorage},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(H, sys)

    add_variables!(
        container,
        ActivePowerVariable,
        devices,
        HydroCommitmentReservoirStorage(),
    )
    add_variables!(container, OnVariable, devices, HydroCommitmentReservoirStorage())
    add_variables!(container, EnergyVariable, devices, HydroCommitmentReservoirStorage())
    add_variables!(
        container,
        WaterSpillageVariable,
        devices,
        HydroCommitmentReservoirStorage(),
    )
    add_variables!(
        container,
        EnergyShortageVariable,
        devices,
        HydroCommitmentReservoirStorage(),
    )
    add_variables!(
        container,
        EnergySurplusVariable,
        devices,
        HydroCommitmentReservoirStorage(),
    )

    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_variables!(container, EnergyOutput, devices, HydroCommitmentReservoirStorage())
    add_parameters!(container, EnergyTargetTimeSeriesParameter, devices, model)
    add_parameters!(container, InflowTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, HydroCommitmentReservoirStorage},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    add_initial_condition!(
        container,
        devices,
        HydroCommitmentReservoirStorage(),
        InitialEnergyLevel(),
    )

    # Energy Balance Constraint
    add_constraints!(container, EnergyBalanceConstraint, devices, model, S)
    add_constraints!(container, EnergyTargetConstraint, devices, model, S)
    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)
    add_constraint_dual!(container, sys, model)

    return
end

"""
Construct model for HydroPumpedStorage with PumpedStorage Dispatch Formulation
with only Active Power
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, HydroDispatchPumpedStorage},
    ::Type{S},
) where {H <: PSY.HydroPumpedStorage, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(H, sys)

    add_variables!(container, ActivePowerInVariable, devices, HydroDispatchPumpedStorage())
    add_variables!(container, ActivePowerOutVariable, devices, HydroDispatchPumpedStorage())
    add_variables!(container, EnergyVariableUp, devices, HydroDispatchPumpedStorage())
    add_variables!(container, EnergyVariableDown, devices, HydroDispatchPumpedStorage())
    add_variables!(container, WaterSpillageVariable, devices, HydroDispatchPumpedStorage())
    add_variables!(container, EnergyOutput, devices, HydroDispatchPumpedStorage())
    if get_attribute(model, "reservation")
        add_variables!(
            container,
            ReservationVariable,
            devices,
            HydroDispatchPumpedStorage(),
        )
    end

    add_parameters!(container, InflowTimeSeriesParameter, devices, model)
    add_parameters!(container, OutflowTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerInVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerOutVariable,
        devices,
        model,
        S,
    )

    add_expressions!(container, ReserveRangeExpressionLB, devices, model)
    add_expressions!(container, ReserveRangeExpressionUB, devices, model)

    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, HydroDispatchPumpedStorage},
    ::Type{S},
) where {H <: PSY.HydroPumpedStorage, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        OutputActivePowerVariableLimitsConstraint,
        ActivePowerOutVariable,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        InputActivePowerVariableLimitsConstraint,
        ActivePowerInVariable,
        devices,
        model,
        S,
    )

    add_initial_condition!(
        container,
        devices,
        HydroDispatchPumpedStorage(),
        InitialEnergyLevelUp(),
    )
    add_initial_condition!(
        container,
        devices,
        HydroDispatchPumpedStorage(),
        InitialEnergyLevelDown(),
    )

    # Energy Balanace limits
    add_constraints!(container, EnergyCapacityUpConstraint, devices, model, S)
    add_constraints!(container, EnergyCapacityDownConstraint, devices, model, S)
    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)

    add_constraint_dual!(container, sys, model)
    return
end