# %%
# 2D rising bubble simulation with PFEM
# =====================================
# This example simulates a 2D rising bubble flow using the Particle Finite Element Method (PFEM).
# It uses the bubble boundary condition.
# The reference setup is Hysing et al., "Quantitative benchmark computations of two-dimensional bubble dynamics", Int. J. Numer. Meth. Fluids, 2009.
# %%
# Keywords
# --------
# PFEM, rising bubble, 2D, fluid
#
# %%
# Description
# -----------
# The test demonstrates:
# - How to perform a PFEM simulation in a fixed domain with the bubble boundary condition
import os, sys, shutil
import numpy as np

# %%
# Gmsh pfem branch
# ----------------
# This test requires the alphashapes branch of GMSH to use the Particle Finite Element (PFEM) module.
# PFEM is a lagrangian method where the mesh is regenerated at each time step, making it suitable for problems with large deformations and free surfaces.
sys.path.insert(0, os.environ["GMSH_PFEM_DIR"])

import gmsh
from migflow import fluid, pfem

# %%
# Output Directory
# ----------------
# Create a clean output directory for simulation results.
outputdir = "output_2d_hysing" if len(sys.argv) < 2 else sys.argv[1]
shutil.rmtree(outputdir, ignore_errors=True)
os.makedirs(outputdir)


# %%
# Geometrical parameters and mesh generation
# ------------------------------------------
l_box = 1.0
h_box = 2.0
bubble_radius = 0.25
bubble_center = np.array([0.5, 0.5])
# %%
# Mesh parameters
# ---------------
geo_mesh_size = l_box / 10
mesh_size = l_box / 10
smin = 0.2 * mesh_size
smax = mesh_size
dmax = 5 * smax
alpha = 1.3

# %%
# Initial fluid mesh
# ------------------
gmsh.initialize()

gmsh.model.add("ModelInit")
rect = gmsh.model.occ.addRectangle(0, 0, 0, l_box, h_box)
disk = gmsh.model.occ.addDisk(
    bubble_center[0], bubble_center[1], 0, bubble_radius, bubble_radius
)
gmsh.model.occ.cut([(2, rect)], [(2, disk)])
gmsh.model.occ.synchronize()
gmsh.model.mesh.setSizeCallback(lambda *args: mesh_size)
gmsh.model.mesh.generate(2)
nodeTags, coords, _ = gmsh.model.mesh.getNodes()

# %%
# Solid domain
# ------------
gmsh.model.add("ModelGeo")
gmsh.model.occ.addRectangle(0, 0, 0, l_box, h_box)
gmsh.model.occ.synchronize()
gmsh.model.mesh.setSizeCallback(lambda *args: mesh_size)
gmsh.model.mesh.generate(2)
gmsh.model.addPhysicalGroup(1, [1], -1, "bottom")
gmsh.model.addPhysicalGroup(1, [2], -1, "right")
gmsh.model.addPhysicalGroup(1, [3], -1, "top")
gmsh.model.addPhysicalGroup(1, [4], -1, "left")
geoEntities = gmsh.model.getEntities(1)
# %%
# PFEM mesh and size fields
# -------------------------
gmsh.model.add("ModelFluid")
alphaDomainTag = gmsh.model.addDiscreteEntity(2, -1, [])
for dim, tag in geoEntities:
    gmsh.model.addDiscreteEntity(dim, tag, [])
alphaBoundaryTag = gmsh.model.addDiscreteEntity(1, -1, [])
gmsh.model.mesh.addNodes(2, alphaDomainTag, nodeTags, coords)
gmsh.model.addPhysicalGroup(1, [1], -1, "bottom")
gmsh.model.addPhysicalGroup(1, [2], -1, "right")
gmsh.model.addPhysicalGroup(1, [3], -1, "top")
gmsh.model.addPhysicalGroup(1, [4], -1, "left")
gmsh.model.addPhysicalGroup(1, [alphaBoundaryTag], -1, "freeSurface")
gmsh.model.addPhysicalGroup(2, [alphaDomainTag], -1, "domain")

# %%
# Size fields
# -----------
sizeFieldConstant = gmsh.model.mesh.field.add("Box")
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "VIn", mesh_size)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "VOut", mesh_size)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "XMin", 0.0)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "XMax", l_box)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "YMin", 0.0)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "YMax", h_box)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "Thickness", 0.001)

print("gmsh path : ", gmsh.__file__)

sizeFieldDistFS = gmsh.model.mesh.field.add("AlphaShapeDistance")
gmsh.model.mesh.field.setNumber(sizeFieldDistFS, "Tag", alphaBoundaryTag)
gmsh.model.mesh.field.setNumber(sizeFieldDistFS, "SamplingLength", 0.1 * smin)
sizeFieldRefine = gmsh.model.mesh.field.add("Threshold")
gmsh.model.mesh.field.setNumber(sizeFieldRefine, "InField", sizeFieldDistFS)
gmsh.model.mesh.field.setNumber(sizeFieldRefine, "SizeMin", smin)
gmsh.model.mesh.field.setNumber(sizeFieldRefine, "SizeMax", smax)
gmsh.model.mesh.field.setNumber(sizeFieldRefine, "DistMin", 0.0)
gmsh.model.mesh.field.setNumber(sizeFieldRefine, "DistMax", dmax)
gmsh.model.mesh.computeAlphaShape(
    2,
    alphaDomainTag,
    alphaBoundaryTag,
    "ModelGeo",
    alpha,
    sizeFieldConstant,
    sizeFieldRefine,
    False,
)

# Put the nodes on the bubble boundary back onto the exact circle (this is necessary due to the refinement)
_, elnbubble = gmsh.model.mesh.getElementsByType(1, alphaBoundaryTag)
nt_bubble = np.unique(elnbubble)
for n in nt_bubble:
    c, _, _, _ = gmsh.model.mesh.getNode(n)
    rad = np.linalg.norm(np.array([c[0], c[1]]) - bubble_center)
    c[0] = bubble_center[0] + bubble_radius * (c[0] - bubble_center[0]) / rad
    c[1] = bubble_center[1] + bubble_radius * (c[1] - bubble_center[1]) / rad
    gmsh.model.mesh.setNode(n, c, [])

# %%
# Physical Parameters
# -------------------
g = np.array([0.0, -0.98])
rho = 1000
rho_bubble = 100
mu = 10
sigma = 24.5
U_g = np.sqrt(-g[1] * 2 * bubble_radius)
L = 2 * bubble_radius
Re = rho * U_g * L / mu
Eo = rho * U_g**2 * L / sigma
print(f"Re={Re}, Eo={Eo}")

# %%
# Time parameters
# ---------------
cfl = 0.3
U = U_g
U_init = U
dt = mesh_size / U * cfl
t = 0
tEnd = 10.0

# %%
# Fluid problem
# -------------
f = fluid.FluidProblem2(g, mu, rho, advection=False)
f.set_wall_boundary("bottom")
f.set_wall_boundary("right")
f.set_wall_boundary("top")
f.set_wall_boundary("left")
f.set_strong_boundary("bottom", velocity=[0, 0])
f.set_strong_boundary("right", velocity_x=0)
f.set_strong_boundary("top", velocity=[0, 0])
f.set_strong_boundary("left", velocity_x=0)
f.set_open_boundary("freeSurface", pressure=0, viscous_flag=False)

# %%
# Simulation Loop
# ---------------
# Time integration of coupled fluid–particle motion.
i = 0
outf = 10
gmsh.option.setNumber("General.Verbosity", 0)

while t < tEnd:
    print(f"{i:4d}, {t:.6g}/{tEnd:.6g}, {dt:.6g}")
    # Update PFEM mesh
    nodetag, oelemtag, oparamcoord, _ = gmsh.model.mesh.computeAlphaShape(
        2,
        alphaDomainTag,
        alphaBoundaryTag,
        "ModelGeo",
        alpha,
        sizeFieldRefine,
        sizeFieldRefine,
        boundaryTolerance=0.01 * mesh_size,
        usePreviousMesh=True,
    )
    oparamcoord = oparamcoord.reshape((-1, 3))
    gmsh.write(outputdir + "/lastMesh.msh")
    ordered_node_tags = pfem.prepareMeshForMigflow(
        i, alphaDomainTag, f, nodetag, oelemtag, oparamcoord
    )

    pfem.applySurfaceTension(f, sigma, bubbleCondition=True, g=g, rho_bubble=rho_bubble)

    if i % outf == 0:
        f.write_mig(outputdir, t)

    # nodes velocity, be aware that its dimension is (n_nodes, 3) not (n_nodes, 2)
    u_old = np.zeros_like(f.coordinates())
    u_old[:, :2] = f.velocity()

    f.implicit_euler(dt)

    dx = np.zeros((f.coordinates().shape[0], 3))
    u = np.zeros_like(dx)
    u[:, :2] = f.velocity()
    if i == 0:
        u_old = u
    dx = u * dt + 0.5 * (u - u_old) / dt * dt**2

    # advect nodes and project if needed
    gmsh.model.mesh.advectMeshNodes(
        2,
        alphaDomainTag,
        alphaBoundaryTag,
        "ModelGeo",
        ordered_node_tags,
        dx.flatten(),
        0.01 * mesh_size,
    )

    # set new coordinates
    _, newCoords, _ = gmsh.model.mesh.getNodes(2, alphaDomainTag)
    f.set_coordinates(newCoords)

    # update time step
    max_U = np.max([U_init, np.max(np.linalg.norm(f.velocity(), axis=1))])
    dt = smin / max_U * cfl

    t += dt
    i += 1

# %%
# Plot
# ----
# .. code-block:: shell
#
#  python3 -m migflow.plot.migplot output_2d_hysing --actors fluid --fluid-field velocity --fluid-vmin 0.0 --fluid-vmax 0.5 --show-edges 1
# %%
# Artifacts
# ---------
# - output_2d_hysing/animation.mp4