# %%
# Three-Dimensional Cylindrical Poiseuille Flow
# =============================================
# This test reproduces a 3D Poiseuille flow inside a cylindrical pipe using
# MigFlow. The cylindrical geometry is generated with Gmsh through OpenCASCADE,
# and a parabolic velocity profile is prescribed at the inlet. The objective is
# to illustrate how to set up a 3D viscous flow simulation with physical
# boundaries and TFEM (P2/P1) discretization.
#
# Keywords
# --------
# FEM, FEM, Poiseuille flow
#
# Description
# -----------
# This example demonstrates:
# - How to generate a 3D cylindrical domain using Gmsh.
# - How to assign physical boundary groups for inlet, outlet, and walls.
# - How to impose an analytical Poiseuille inflow profile.
# - How to solve the 3D Stokes/Navier–Stokes equations using MigFlow.
# - How to write MIG output for post-processing.

from migflow import fluid, gmsh_io
import numpy as np
import os, sys, shutil
import gmsh

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


# %%
# Geometry
# --------
# The domain is a cylinder of:
# - radius r = h/2
# - length w
# - axis oriented along the x-direction
h = 1  # domain height (diameter)
r = h / 2  # radius
w = 2  # cylinder length
axis = [w, 0, 0]  # cylinder axis direction
origin = [0, 0, 0]
# %%
# Mesh Generation
# ---------------
# - Physical groups auto-detected from Gmsh entities:
#       "Left"   : inlet plane
#       "Right"  : outlet plane
#       "Lateral": cylindrical wall
#       "Domain" : volume
# - Uniform mesh size set by callback `mesh_size = r/2`.
mesh_size = r / 2  # characteristic mesh size
gmsh.initialize()


gmsh.model.occ.add_cylinder(origin[0], origin[1], origin[2], *axis, r)
gmsh.model.occ.synchronize()

# Assign physical groups automatically retrieved from CAD IDs:
gmsh.model.add_physical_group(2, [1], name="Lateral")
gmsh.model.add_physical_group(2, [2], name="Right")
gmsh.model.add_physical_group(2, [3], name="Left")
gmsh.model.add_physical_group(3, [1], name="Domain")

# Set uniform mesh size
gmsh.model.mesh.set_size_callback(lambda dim, tag, x, y, z, lc: mesh_size)

# Generate tetrahedral 3D mesh
gmsh.model.mesh.generate(3)


# %%
# Physical Parameters
# -------------------
g = np.array([0.0, 0.0, 0.0])  # no gravity
mu = 1  # viscosity
rho = 1  # density
dp = 1  # driving pressure parameter (for U_max estimation)


# %%
# Fluid Problem Initialization
# ----------------------------
f = fluid.FluidProblem3(g, mu, rho, p2p1=True)
gmsh_io.load_msh(f, None)

# %%
# Physical Setup
# --------------
# The analytical Poiseuille velocity profile inside a cylindrical pipe is:
#
#     u_x(r) = U_max * (1 - (r/R)^2)
#
# where:
# - r = radial distance to cylinder axis
# - R = cylinder radius
# - U_max derived from dp, μ, geometry
u_max = dp / w * r**2 / (20 * mu)


def u_in(x):
    """Analytical Poiseuille inflow profile projected on velocity x-component."""
    radial_dist = np.linalg.norm(x[:, 1:], axis=1)
    return u_max * (1 - (radial_dist / r) ** 2)


# %%
# Boundary conditions
# -------------------
# - Left   : inflow with Poiseuille profile
# - Right  : open boundary (p = 0)
# - Lateral: no-slip (zero velocity)
f.set_wall_boundary("Lateral", velocity=[0.0, 0.0, 0.0])
f.set_open_boundary("Right", pressure=0)
f.set_open_boundary("Left", velocity=[u_in, 0.0, 0.0])


# %%
# Simulation Parameters
# ---------------------
t = 0
i = 0
dt = 1e-2
tEnd = 10 * dt
outf = 1


# %%
# Time Integration Loop
# ---------------------
while t < tEnd:
    print(f"{i:4d}, {t:.6g}/{tEnd:.6g}, dt={dt:.6g}")

    # Write output every `outf` iterations
    if i % outf == 0:
        f.write_mig(outputdir, t)

    # Implicit Euler solver
    f.implicit_euler(dt, check_residual_norm=1)

    t += dt
    i += 1

# %%
# Plot
# ----
# .. code-block:: shell
#
#  python3 -m migflow.plot.migplot output_3d_poiseuille_p2p1 --actors fluid --fluid-field velocity --clip -z 0
# %%
# Artifacts
# ---------
# - output_3d_poiseuille_p2p1/animation.mp4