# %%
# 2D Poiseuille flow with PFEM
# ============================
# This test case simulates a 2D Poiseuille flow using the Particle Finite Element Method (PFEM).
#
# %%
# Keywords
# --------
# PFEM, fluid, Poiseuille, 2D
#
# %%
# Description
# -----------
# The test demonstrates:
# - How to perform a PFEM simulation.
import numpy as np
import shutil
import os
import sys

# %%
# 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_poiseuille_pfem" if len(sys.argv) < 2 else sys.argv[1]
shutil.rmtree(outputdir, ignore_errors=True)
os.makedirs(outputdir)

# %%
# Geometrical parameters and mesh generation
# ------------------------------------------
L_x = 3.0
L_y = 1.0
# %%
# Mesh parameters
# ---------------
lc = 0.05
lc_mesh = 0.1
alpha = 10

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

gmsh.model.add("ModelInit")
gmsh.model.occ.addRectangle(0, -L_y / 2, 0, L_x, L_y)
gmsh.model.occ.synchronize()
gmsh.model.mesh.setSizeCallback(lambda *args: lc)
gmsh.model.mesh.generate(2)
_, x, _ = gmsh.model.mesh.getNodes()

# %%
# Solid domain
# ------------
gmsh.model.add("ModelGeo")
gmsh.model.occ.addRectangle(0, -L_y / 2, 0, L_x, L_y)
gmsh.model.occ.synchronize()
geoEntities = gmsh.model.getEntities(1)
gmsh.model.mesh.setSizeCallback(lambda *args: lc)
gmsh.model.mesh.generate(2)
tag_left = 4
tag_right = 2

# %%
# 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, [], x)
gmsh.model.addPhysicalGroup(1, [1, 3], -1, "lateral_walls")
gmsh.model.addPhysicalGroup(1, [2], -1, "outlet")
gmsh.model.addPhysicalGroup(1, [4], -1, "inlet")
gmsh.model.addPhysicalGroup(1, [alphaBoundaryTag], -1, "freeSurface")

# %%
# Size fields
# -----------
sizeFieldConstant = gmsh.model.mesh.field.add("Box")
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "VIn", lc)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "VOut", 10 * lc)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "XMin", 0.0)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "XMax", L_x)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "YMin", -L_y / 2)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "YMax", L_y / 2)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "Thickness", 0.001)
gmsh.model.addPhysicalGroup(2, [alphaDomainTag], -1, "domain")


# %%
# Physical Parameters
# -------------------
rho = 1000
mu = 1e-3
v_in = 1.0
Re = rho * v_in * L_y / mu
g = np.array([0.0, 0.0])

# %%
# Time parameters
# ---------------
cfl = 0.5
dt = lc / v_in * cfl
tEnd = 10
outf = 10

# %%
# Fluid problem
# -------------
f = fluid.FluidProblem2(g, mu, rho, advection=False)
f.set_wall_boundary("lateral_walls")
f.set_strong_boundary("lateral_walls", velocity=[0, 0])


def U_inlet(x):
    return v_in * (1 - (2 * x[:, 1] / L_y) ** 2)


f.set_open_boundary("inlet", velocity=[U_inlet, 0])
f.set_open_boundary("outlet", pressure=0)

# %%
# Simulation Loop
# ---------------
# Time integration of coupled fluid–particle motion.
t = 0.0
i = 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,
        sizeFieldConstant,
        sizeFieldConstant,
        boundaryTolerance=0.01 * lc,
        usePreviousMesh=False,
    )
    oparamcoord = oparamcoord.reshape((-1, 3))
    ordered_node_tags = pfem.prepareMeshForMigflow(
        i, alphaDomainTag, f, nodetag, oelemtag, oparamcoord
    )

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

    # Solve step
    f.implicit_euler(dt)
    dx = np.zeros((f.coordinates().shape[0], 3))
    dx[:, :2] = f.velocity() * dt

    # Fix the boundary nodes on left and right walls
    _, right_nodes = gmsh.model.mesh.getElementsByType(1, tag_right)
    _, left_nodes = gmsh.model.mesh.getElementsByType(1, tag_left)
    fixed_nodes = np.concatenate((right_nodes, left_nodes))
    fixed_nodes = np.unique(fixed_nodes)
    fixed_nodes_idx = np.searchsorted(ordered_node_tags, fixed_nodes)
    dx[fixed_nodes_idx, :] = [0, 0, 0]

    # Advect mesh nodes
    dx = dx.reshape((-1,))
    gmsh.model.mesh.advectMeshNodes(
        2,
        alphaDomainTag,
        alphaBoundaryTag,
        "ModelGeo",
        ordered_node_tags,
        dx,
        0.01 * lc,
    )
    _, newCoords, _ = gmsh.model.mesh.get_nodes(2)

    newCoords = newCoords.reshape((-1, 3))
    f.set_coordinates(newCoords)

    t += dt
    i += 1

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