# %%
# 2D rotating wheel with immersed granular flow using PFEM-DEM
# ============================================================
# This example simulates a rotating wheel partially immersed in a fluid.
# Granular particles are deposited above the wheel and fall under gravity,
# interacting with both the rotating structure and the surrounding fluid.
#
# The fluid domain is represented using the PFEM method while the grains
# are modeled using DEM particles.
#
# The test demonstrates:
# - PFEM free-surface simulations with remeshing
# - Coupling between PFEM and DEM
# - Moving immersed boundaries with prescribed motion
# - Fluid-grain interactions around rotating machinery
# - Robustness of the PFEM remeshing strategy under large deformations
#
# Keywords
# --------
# PFEM, DEM, rotating wheel, immersed boundary, granular flow, free surface, 2D

import numpy as np
import sys
import os
import shutil

# %%
# 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, time_integration, scontact, pfem, deposit
from migflow import volume_coupling as vc
from mesh_wheel_blades import generate_blade, spin

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

gmsh.initialize()

# %%
# Geometrical parameters
# ----------------------
# Tank dimensions and rotating wheel geometry.
L_tank = 15.0
H_tank = 12.5
H_fluid = H_tank / 3

# Blades
L_blade = 2.0
H_blade = 2.5
r_le = 0.2
r_te = 0.05
R_wheel = 2.0
c_wheel = np.array([L_tank / 2, 5])
n_blades = 10
omega = np.pi / 8
freq = omega / (2 * np.pi)
U = 1

# %%
# Physical parameters
# -------------------
# Fluid properties and gravity acceleration.
g = np.array([0.0, -9.81])
rho = 1000
mu = 10
sigma = 0.0071
rho_bubble = 1

# %%
# DEM parameters
# --------------
# Granular material properties.
n_particles = int(4e3)
r_particles = 0.03
rhop = 1500
friction = 0.3

# %%
# Mesh parameters
# ---------------
# PFEM remeshing and refinement parameters.
alpha = 1.2
mesh_size = 0.3
geo_mesh_size = 2.0
smin = 0.3 * mesh_size
smax = mesh_size
dmax = 4 * smax

# %%
# Time parameters
# ---------------
# Time integration and CFL parameters.
U_init = U
cfl = 0.1
t = 0
ii = 0
dt = 2e-4
tEnd = 3
print(f"dt = {dt}, U = {U}")
outf = 30

# %%
# Initial fluid mesh
# ------------------
# Initial fluid domain before PFEM remeshing.
gmsh.model.add("ModelInit")
gmsh.model.occ.addRectangle(0, 0, 0, L_tank, H_fluid)
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
# ------------
# Construction of the rotating wheel geometry and boundary tags.
blade_points = generate_blade(L_blade, H_blade, r_le, r_te, c_wheel, R_wheel, 20)
gmsh.option.setNumber("General.Verbosity", 0)
gmsh.model.add("ModelGeo")
d_alpha = 0.0
spin(d_alpha, L_tank, H_tank, c_wheel, R_wheel, n_blades, blade_points)

eps = 1e-6
tag_bottom = gmsh.model.getEntitiesInBoundingBox(
    0 - eps, 0 - eps, -eps, L_tank + eps, 0 + eps, eps, 1
)[0][1]
tag_right = gmsh.model.getEntitiesInBoundingBox(
    L_tank - eps, 0 - eps, -eps, L_tank + eps, H_tank + eps, eps, 1
)[0][1]
tag_top = gmsh.model.getEntitiesInBoundingBox(
    0 - eps, H_tank - eps, -eps, L_tank + eps, H_tank + eps, eps, 1
)[0][1]
tag_left = gmsh.model.getEntitiesInBoundingBox(
    0 - eps, 0 - eps, -eps, 0 + eps, H_tank + eps, eps, 1
)[0][1]
ent_wheel = gmsh.model.getEntitiesInBoundingBox(
    c_wheel[0] - R_wheel - L_blade * 1.5 - eps,
    c_wheel[1] - R_wheel - L_blade * 1.5 - eps,
    -eps,
    c_wheel[0] + R_wheel + L_blade * 1.5 + eps,
    c_wheel[1] + R_wheel + L_blade * 1.5 + eps,
    eps,
    1,
)
tags_wheel = [t[1] for t in ent_wheel]

geoEntities = [tag_bottom, tag_right, tag_top, tag_left] + tags_wheel
gmsh.model.addPhysicalGroup(1, [tag_bottom], -1, "bottom")
gmsh.model.addPhysicalGroup(1, [tag_right], -1, "right")
gmsh.model.addPhysicalGroup(1, [tag_top], -1, "top")
gmsh.model.addPhysicalGroup(1, [tag_left], -1, "left")
gmsh.model.addPhysicalGroup(1, tags_wheel, -1, "wheel")

# %%
# Fluid problem
# -------------
# Definition of the PFEM fluid problem and boundary conditions.
f = fluid.FluidProblem2(g, mu, rho, advection=False)
f.set_strong_boundary("top", velocity_y=0)
f.set_open_boundary("freeSurface", pressure=0, viscous_flag=False)


def wheel_velocity_x(x):
    r_y = x[:, 1] - c_wheel[1]
    return r_y * omega


def wheel_velocity_y(x):
    r_x = x[:, 0] - c_wheel[0]
    return -r_x * omega


def pressure_outflow(x):
    return -rho * g[1] * (H_fluid - x[:, 1])


f.set_wall_boundary("bottom")
f.set_strong_boundary("bottom", velocity=[0, 0])

f.set_wall_boundary("left")
f.set_strong_boundary("left", velocity_x=0)

f.set_wall_boundary("right")
f.set_strong_boundary("right", velocity_x=0)

# Rotating wheel boundary condition
f.set_open_boundary("wheel", velocity=[wheel_velocity_x, wheel_velocity_y])
f.set_strong_boundary("wheel", velocity=[wheel_velocity_x, wheel_velocity_y])

# %%
# DEM problem
# -----------
# Construction of the granular domain and moving wheel contact geometry.
dem = scontact.ParticleProblem(2)
dem.set_fixed_contact_geometry(0)

outer_body = dem.add_body((0, 0), 0, 0)
x0 = np.array([0, 0])
x1 = np.array([L_tank, 0])
x2 = np.array([L_tank, H_tank])
x3 = np.array([0, H_tank])
dem.add_segment_to_body(x0, x1, outer_body, "bnd")
dem.add_segment_to_body(x1, x2, outer_body, "bnd")
dem.add_segment_to_body(x2, x3, outer_body, "bnd")
dem.add_segment_to_body(x3, x0, outer_body, "bnd")
dem.add_particle_to_body(x0, 0, outer_body)
dem.add_particle_to_body(x1, 0, outer_body)
dem.add_particle_to_body(x2, 0, outer_body)
dem.add_particle_to_body(x3, 0, outer_body)

wheel_body = dem.add_body(c_wheel, 0, 0)
for ent in ent_wheel:
    pts = gmsh.model.getBoundary([ent])
    _, x0, _ = gmsh.model.mesh.getNodes(0, pts[0][1])
    _, x1, _ = gmsh.model.mesh.getNodes(0, pts[1][1])
    dem.add_particle_to_body(x0[:2] - c_wheel, 0, wheel_body)
    dem.add_particle_to_body(x1[:2] - c_wheel, 0, wheel_body)
    dem.add_segment_to_body(x0[:2] - c_wheel, x1[:2] - c_wheel, wheel_body, "wheel")
dem.body_omega()[wheel_body] = -omega
dem.set_friction_coefficient(friction)

# %%
# Initial granular deposit
# ------------------------
# Generate a cloud of grains above the rotating wheel.
ratio = 1.2
diameter_max = 2.0 * r_particles
diameter_min = diameter_max / ratio
np.random.seed(0)
u = np.random.power(1, n_particles)
d = diameter_max * diameter_min / (diameter_max + u * (diameter_min - diameter_max))
radii = d / 2
boundaries = np.array(
    [
        [0.2 * L_tank, H_tank],
        [0.2 * L_tank, 0.8 * H_tank],
        [0.8 * L_tank, 0.8 * H_tank],
        [0.8 * L_tank, H_tank],
    ]
)
xp, radii = deposit.fastdeposit_2d(boundaries, radii, "deposit.svg")
eps = 1e-4 * diameter_max
for x, r in zip(xp, radii):
    dem.add_particle(x + eps, r, np.pi * r**2 * rhop)

# %%
# PFEM mesh and size fields
# -------------------------
# Create the PFEM discrete domain and adaptive remeshing fields.
gmsh.model.add("ModelFluid")
alphaDomainTag = gmsh.model.addDiscreteEntity(2, -1, [])
for tag in geoEntities:
    gmsh.model.addDiscreteEntity(1, tag, [])
alphaBoundaryTag = gmsh.model.addDiscreteEntity(1, -1, [])
gmsh.model.mesh.addNodes(2, alphaDomainTag, nodeTags, coords)
gmsh.model.addPhysicalGroup(1, [tag_bottom], -1, "bottom")
gmsh.model.addPhysicalGroup(1, [tag_right], -1, "right")
gmsh.model.addPhysicalGroup(1, [tag_top], -1, "top")
gmsh.model.addPhysicalGroup(1, [tag_left], -1, "left")
gmsh.model.addPhysicalGroup(1, tags_wheel, -1, "wheel")
gmsh.model.addPhysicalGroup(1, [alphaBoundaryTag], -1, "freeSurface")
gmsh.model.addPhysicalGroup(2, [alphaDomainTag], -1, "domain")

# %%
# Size fields
# -----------
# Adaptive refinement near the free surface and rotating wheel.
sizeFieldConstant = gmsh.model.mesh.field.add("Box")
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "VIn", mesh_size)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "VOut", 10 * mesh_size)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "XMin", 0.0)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "XMax", L_tank)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "YMin", 0.0)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "YMax", H_tank)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "Thickness", 0.001)

sizeFieldDistFS = gmsh.model.mesh.field.add("AlphaShapeDistance")
gmsh.model.mesh.field.setNumber(sizeFieldDistFS, "Tag", alphaBoundaryTag)
gmsh.model.mesh.field.setNumber(sizeFieldDistFS, "SamplingLength", 0.25 * smin)
sizeFieldRefine1 = gmsh.model.mesh.field.add("Threshold")
gmsh.model.mesh.field.setNumber(sizeFieldRefine1, "InField", sizeFieldDistFS)
gmsh.model.mesh.field.setNumber(sizeFieldRefine1, "SizeMin", smin)
gmsh.model.mesh.field.setNumber(sizeFieldRefine1, "SizeMax", smax)
gmsh.model.mesh.field.setNumber(sizeFieldRefine1, "DistMin", 0.0)
gmsh.model.mesh.field.setNumber(sizeFieldRefine1, "DistMax", dmax)

sizeFieldBnd = []
for b in tags_wheel:
    sf = gmsh.model.mesh.field.add("AlphaShapeDistance")
    gmsh.model.mesh.field.setNumber(sf, "Tag", b)
    gmsh.model.mesh.field.setNumber(sf, "SamplingLength", 0.25 * smin)
    sizeFieldBnd.append(gmsh.model.mesh.field.add("Threshold"))
    gmsh.model.mesh.field.setNumber(sizeFieldBnd[-1], "InField", sf)
    gmsh.model.mesh.field.setNumber(sizeFieldBnd[-1], "SizeMin", smin)
    gmsh.model.mesh.field.setNumber(sizeFieldBnd[-1], "SizeMax", smax)
    gmsh.model.mesh.field.setNumber(sizeFieldBnd[-1], "DistMin", 0.0)
    gmsh.model.mesh.field.setNumber(sizeFieldBnd[-1], "DistMax", dmax)
sizeFieldRefine = gmsh.model.mesh.field.add("Min")
gmsh.model.mesh.field.setNumbers(
    sizeFieldRefine, "FieldsList", [sizeFieldRefine1] + sizeFieldBnd
)
gmsh.model.mesh.computeAlphaShape(
    2,
    alphaDomainTag,
    alphaBoundaryTag,
    "ModelGeo",
    alpha,
    sizeFieldConstant,
    sizeFieldRefine,
    False,
)

# %%
# Simulation loop
# ---------------
# Time integration of the coupled PFEM-DEM problem.
gmsh.option.setNumber("General.Verbosity", 0)
while t < tEnd:
    print("ii = ", ii, "; t = ", "%.4f" % t)

    # Update PFEM mesh
    nodetag, oelemtag, oparamcoord, _ = gmsh.model.mesh.computeAlphaShape(
        2,
        alphaDomainTag,
        alphaBoundaryTag,
        "ModelGeo",
        alpha,
        sizeFieldRefine,
        sizeFieldRefine,
        boundaryTolerance=0.001 * smin,
        usePreviousMesh=True,
    )
    oparamcoord = oparamcoord.reshape((-1, 3))

    ordered_node_tags = pfem.prepareMeshForMigflow(
        ii, alphaDomainTag, f, nodetag, oelemtag, oparamcoord
    )
    geo = vc.get_particles(
        f, dem.position(), dem.r(), density=vc.body_density(dem)
    )
    # WHEEL_GAMMA_T=0 forces the trace operator OFF; eta is left at the library
    # schedule, so this is the full unfitted model with the interior fill and
    # without the trace. Unset = the shipped default band.
    _gt = os.environ.get("WHEEL_GAMMA_T")
    clo = vc.get_particles_closure(
        f, geo, mu, rho, dem.velocity(), omega=None,
        contact_forces=dem.contact_forces(),
        **({} if _gt is None else {"gamma_t": float(_gt)}),
    )
    dat = vc.get_patankar(geo, clo, dt, g, dem.velocity())
    vc.set_coupling_geometry(f, geo)
    vc.set_coupling_closure(f, clo)
    vc.set_coupling_datum(f, dat)

    pfem.applySurfaceTension(f, sigma, False, g, rho_bubble)

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

    f.implicit_euler(dt)

    drag = f.compute_node_force()
    external_forces = drag + g * dem.r() ** 2 * np.pi * rhop
    time_integration._advance_particles(dem, external_forces, dt, 1, 1e-3 * r_particles)

    t += dt
    ii += 1

    # Rotate wheel geometry
    d_alpha += omega * dt
    spin(d_alpha, L_tank, H_tank, c_wheel, R_wheel, n_blades, blade_points)

    # Advect PFEM mesh nodes
    dx = np.zeros((f.coordinates().shape[0], 3))
    dx[:, :2] = f.velocity() * dt
    dx = dx.reshape((-1,))
    gmsh.model.mesh.advectMeshNodes(
        2,
        alphaDomainTag,
        alphaBoundaryTag,
        "ModelGeo",
        ordered_node_tags,
        dx,
        0.001 * smin,
    )

    _, newCoords, _ = gmsh.model.mesh.getNodes(2, alphaDomainTag)

    f.set_coordinates(newCoords)

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

    print("new dt = ", dt)

gmsh.finalize()

# %%
# Plot
# ----
# .. code-block:: shell
#
#  python3 -m migflow.plot.migplot output --actors fluid particles --fluid-field velocity --fluid-vmin -2 --fluid-vmax 3

# %%
# Artifacts
# ---------
# - output/animation.mp4
# - output/deposit.svg