# ------------------------------------------------------------------------------
#
#  Gmsh Julia tutorial 21
#
#  Mesh partitioning
#
# ------------------------------------------------------------------------------

import gmsh

# Gmsh can partition meshes using different algorithms, e.g. the graph
# partitioner Metis or the `SimplePartition' plugin. For all the partitioning
# algorithms, the relationship between mesh elements and mesh partitions is
# encoded through the creation of new (discrete) elementary entities, called
# "partition entities".
#
# Partition entities behave exactly like other discrete elementary entities; the
# only difference is that they keep track of both a mesh partition index and
# their parent elementary entity.
#
# The major advantage of this approach is that it allows to maintain a full
# boundary representation of the partition entities, which Gmsh creates
# automatically if `Mesh.PartitionCreateTopology' is set.

gmsh.initialize()

# Let us start by creating a simple geometry with two adjacent squares sharing
# an edge:
gmsh.model.add("t21")
gmsh.model.occ.addRectangle(0, 0, 0, 1, 1, 1)
gmsh.model.occ.addRectangle(1, 0, 0, 1, 1, 2)
gmsh.model.occ.fragment([(2, 1)], [(2, 2)])
gmsh.model.occ.synchronize()
gmsh.model.mesh.setSize(gmsh.model.getEntities(0), 0.05)

# We create one physical group for each square, and we mesh the resulting
# geometry:
gmsh.model.addPhysicalGroup(2, [1], 100, "Left")
gmsh.model.addPhysicalGroup(2, [2], 200, "Right")
gmsh.model.mesh.generate(2)

# We now define several ONELAB parameters to fine-tune how the mesh will be
# partitioned:
gmsh.onelab.set("""[
  {
    "type":"number",
    "name":"Parameters/0Mesh partitioner",
    "values":[0],
    "choices":[0, 1],
    "valueLabels":{"Metis":0, "SimplePartition":1}
  },
  {
    "type":"number",
    "name":"Parameters/1Number of partitions",
    "values":[3],
    "min":1,
    "max":256,
    "step":1
  },
  {
    "type":"number",
    "name":"Parameters/2Create partition topology (BRep)?",
    "values":[1],
    "choices":[0, 1]
  },
  {
    "type":"number",
    "name":"Parameters/3Create ghost cells?",
    "values":[0],
    "choices":[0, 1]
  },
  {
    "type":"number",
    "name":"Parameters/3Create new physical groups?",
    "values":[0],
    "choices":[0, 1]
  },
  {
    "type":"number",
    "name":"Parameters/3Write file to disk?",
    "values":[1],
    "choices":[0, 1]
  },
  {
    "type":"number",
    "name":"Parameters/4Write one file per partition?",
    "values":[0],
    "choices":[0, 1]
  }
]""")

function partitionMesh()
    # Number of partitions
    N = Int(gmsh.onelab.getNumber("Parameters/1Number of partitions")[1])

    # Should we create the boundary representation of the partition entities?
    gmsh.option.setNumber("Mesh.PartitionCreateTopology",
                          gmsh.onelab.getNumber(
                              "Parameters/2Create partition topology (BRep)?")[1])

    # Should we create ghost cells?
    gmsh.option.setNumber("Mesh.PartitionCreateGhostCells",
                          gmsh.onelab.getNumber(
                              "Parameters/3Create ghost cells?")[1])

    # Should we automatically create new physical groups on the partition
    # entities?
    gmsh.option.setNumber("Mesh.PartitionCreatePhysicals",
                          gmsh.onelab.getNumber(
                              "Parameters/3Create new physical groups?")[1])

    # Should we keep backward compatibility with pre-Gmsh 4, e.g. to save the
    # mesh in MSH2 format?
    gmsh.option.setNumber("Mesh.PartitionOldStyleMsh2", 0)

    # Should we save one mesh file per partition?
    gmsh.option.setNumber("Mesh.PartitionSplitMeshFiles",
                          gmsh.onelab.getNumber(
                              "Parameters/4Write one file per partition?")[1])

    if gmsh.onelab.getNumber("Parameters/0Mesh partitioner")[1] == 0
        # Use Metis to create N partitions
        gmsh.model.mesh.partition(N)
        # Several options can be set to control Metis: `Mesh.MetisAlgorithm' (1:
        # Recursive, 2: K-way), `Mesh.MetisObjective' (1: min. edge-cut, 2:
        # min. communication volume), `Mesh.PartitionTriWeight' (weight of
        # triangles), `Mesh.PartitionQuadWeight' (weight of quads), ...
    else
        # Use the `SimplePartition' plugin to create chessboard-like partitions
        gmsh.plugin.setNumber("SimplePartition", "NumSlicesX", N)
        gmsh.plugin.setNumber("SimplePartition", "NumSlicesY", 1)
        gmsh.plugin.setNumber("SimplePartition", "NumSlicesZ", 1)
        gmsh.plugin.run("SimplePartition")
    end

    # Save mesh file (or files, if `Mesh.PartitionSplitMeshFiles' is set):
    if gmsh.onelab.getNumber("Parameters/3Write file to disk?")[1] == 1
        gmsh.write("t21.msh")
    end

    # Iterate over partitioned entities and print some info (see the first
    # extended tutorial `x1.jl' for additional information):
    entities = gmsh.model.getEntities()
    for e in entities
        partitions = gmsh.model.getPartitions(e[1], e[2])
        if length(partitions) > 0
            println("Entity ", e, " of type ", gmsh.model.getType(e[1], e[2]))
            println(" - Partition(s): ", partitions)
            println(" - Parent: ", gmsh.model.getParent(e[1], e[2]))
            println(" - Boundary: ", gmsh.model.getBoundary([e]))
        end
    end
end

partitionMesh()

# Launch the GUI and handle the "check" event to re-partition the mesh according
# to the choices made in the GUI
function checkForEvent()
    action = gmsh.onelab.getString("ONELAB/Action")
    if length(action) > 0 && action[1] == "check"
        gmsh.onelab.setString("ONELAB/Action", [""])
        partitionMesh()
        gmsh.graphics.draw()
    end
    return true
end

if !("-nopopup" in ARGS)
    gmsh.fltk.initialize()
    while gmsh.fltk.isAvailable() == 1 && checkForEvent()
        gmsh.fltk.wait()
    end
end

gmsh.finalize()
