# %%
# 2D Granular column collapse onto a fluid bed
# ============================================
# This example simulates the collapse of a granular column onto a fluid bed using the PFEM-DEM approach.
# It showcases the interaction between granular materials and fluids, capturing the dynamics of the collapse and subsequent fluid-grain interactions.
# The aim is to visualise the type of free surface wave generated by the impact of the grains on the fluid bed.
# This test case follows the setup described in:
# - W. Sarlin, C. Morize, A. Sauret, P. Gondret, Nonlinear regimes of tsunami waves generated by a granular collapse, J. Fluid. Mech. 919 (2021) R6.

# Keywords
# --------
# PFEM, fluid-grains, granular column collapse, 2D
#
# Description
# -----------
# The test demonstrates:
# - How to perform a PFEM-DEM simulation with moving boundaries.
import numpy as np
import sys
import shutil
import os

# %%
# 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

# %%
# 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)
os.makedirs(outputdir)

# %%
# Geometrical parameters and mesh generation
# ------------------------------------------
L = 2.0  # length of the domain
H = 1.0  # height of the domain
H_0_array = np.array([0.39, 0.29, 0.29])  # initial column heights
L_0_array = np.array([0.145, 0.1, 0.05])  # initial column lengths
h_0_array = np.array([0.04, 0.08, 0.25])  # initial fluid heights

H_0 = H_0_array[0]
L_0 = L_0_array[0]
l_0 = L - L_0
h_0 = h_0_array[0]
t_gate = 0.01 * L_0  # gate thickness
y_gate = h_0 + 1e-4  # initial gate position
h_gate = y_gate + 1.2 * H_0  # gate height
u_gate = 1  # m/s; gate upward velocity

# %%
# Mesh parameters
# ---------------
geo_mesh_size = L / 10
mesh_size = 0.01 * l_0
smax = mesh_size
smin = 0.2 * mesh_size
dmax = 6 * smax
alpha = 1.3

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

gmsh.model.add("ModelInit")
gmsh.model.occ.addRectangle(L_0, 0, 0, l_0, h_0)
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")


def set_geo_mesh(y_gate=1e-4):
    name = gmsh.model.getCurrent()
    gmsh.model.setCurrent("ModelGeo")
    dimtag = gmsh.model.getEntities(2)
    gmsh.model.occ.remove(dimtag, True)
    box = gmsh.model.occ.addRectangle(0, 0, 0, L, H)
    step = gmsh.model.occ.addRectangle(0, 0, 0, L_0, h_0)
    gate = gmsh.model.occ.addRectangle(L_0, y_gate, 0, t_gate, h_gate)  # gate
    gmsh.model.occ.cut([(2, box)], [(2, gate), (2, step)])
    gmsh.model.occ.synchronize()
    gmsh.model.mesh.setSizeCallback(lambda *args: geo_mesh_size)
    gmsh.model.mesh.generate(2)
    geoEntities = gmsh.model.getEntities(1)
    gmsh.model.setCurrent(name)
    return geoEntities


geoEntities = set_geo_mesh(y_gate)
gmsh.model.addPhysicalGroup(1, [16], -1, "bottom")
gmsh.model.addPhysicalGroup(1, [15], -1, "right")
gmsh.model.addPhysicalGroup(1, [13], -1, "left")
gmsh.model.addPhysicalGroup(1, [7], -1, "step_top")
gmsh.model.addPhysicalGroup(1, [6], -1, "step_right")
gmsh.model.addPhysicalGroup(1, [9, 10, 11, 12], -1, "gate")

# %%
# 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, [16], -1, "bottom")
gmsh.model.addPhysicalGroup(1, [15], -1, "right")
gmsh.model.addPhysicalGroup(1, [13], -1, "left")
gmsh.model.addPhysicalGroup(1, [7], -1, "step_top")
gmsh.model.addPhysicalGroup(1, [6], -1, "step_right")
gmsh.model.addPhysicalGroup(1, [9, 10, 11, 12], -1, "gate")
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", 10 * mesh_size)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "XMin", 0.0)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "XMax", L)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "YMin", 0.0)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "YMax", H)
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)

bndEntities = []
sizeFieldBnd = []
for b in bndEntities:
    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,
)

# %%
# Physical Parameters
# -------------------
g = np.array([0.0, -9.81])
rho = 1000
mu = 1e-3
rho_bubble = 1
sigma = 0.00728
use_bubble = False

# %%
# DEM parameters
# --------------
volume_corr = 0.8  # volume correction on the grains to account for 2D-3D discrepancies
rhop = 2500
friction = 0.9  # 0.2
friction_wall = 0.9
friction_gate = 0.0
r_particles = 0.0025
n_particles = int(L_0 / (2 * r_particles) * H_0 / (2 * r_particles) * 1.5)

# %%
# Time parameters
# ---------------
cfl = 0.1
U = 10
U_init = 1.0
t = 0
dt = smin / U * cfl
settle_time = 0.03
tEnd = 1.0 + settle_time
outf = 10
print(f"dt = {dt}, U = {U}")

# %%
# Fluid problem
# -------------
f = fluid.FluidProblem2(g, mu, rho, advection=False)
f.set_wall_boundary("bottom")
f.set_wall_boundary("right")


def gate_velocity(t):
    if t > settle_time:
        return u_gate
    return 0


f.set_wall_boundary("gate", velocity=[0, gate_velocity(t)])
f.set_wall_boundary("left")
f.set_wall_boundary("step_top")
f.set_wall_boundary("step_bottom")
f.set_strong_boundary("bottom", velocity=[0, 0])
f.set_strong_boundary("right", velocity_x=0)
f.set_strong_boundary("left", velocity_x=0)
f.set_strong_boundary("step_top", velocity_y=0)
f.set_strong_boundary("step_right", velocity_x=0)
f.set_open_boundary("freeSurface", pressure=0, viscous_flag=False)

# %%
# DEM Problem
# -----------
dem = scontact.ParticleProblem(2)
outer_body = dem.add_body((0, 0), 0, 0)
x0 = np.array([L_0, 0])
x1 = np.array([L, 0])
x2 = np.array([L, H])
x3 = np.array([0, H])
x4 = np.array([0, h_0])
x5 = np.array([L_0, h_0])
dem.add_segment_to_body(x0, x1, outer_body, "Wall", material="Wall")
dem.add_segment_to_body(x1, x2, outer_body, "Wall", material="Wall")
dem.add_segment_to_body(x2, x3, outer_body, "Wall", material="Wall")
dem.add_segment_to_body(x3, x4, outer_body, "Wall", material="Wall")
dem.add_segment_to_body(x4, x5, outer_body, "Wall", material="Wall")
dem.add_segment_to_body(x5, x0, outer_body, "Wall", material="Wall")
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)
dem.add_particle_to_body(x4, 0, outer_body)
dem.add_particle_to_body(x5, 0, outer_body)

gate_body = dem.add_body((0, 0), 0, 0)
x0 = np.array([L_0, y_gate])
x1 = np.array([L_0 + t_gate, y_gate])
x2 = np.array([L_0 + t_gate, h_gate])
x3 = np.array([L_0, h_gate])
dem.add_segment_to_body(x0, x1, gate_body, "Gate", material="Gate")
dem.add_segment_to_body(x1, x2, gate_body, "Gate", material="Gate")
dem.add_segment_to_body(x2, x3, gate_body, "Gate", material="Gate")
dem.add_segment_to_body(x3, x0, gate_body, "Gate", material="Gate")
dem.add_particle_to_body(x0, 0, gate_body)
dem.add_particle_to_body(x1, 0, gate_body)
dem.add_particle_to_body(x2, 0, gate_body)
dem.add_particle_to_body(x3, 0, gate_body)
dem.body_velocity()[gate_body, 1] = u_gate

dem.set_fixed_contact_geometry(0)
dem.set_friction_coefficient(10)

dem.set_friction_coefficient(friction, "Grain", "Grain")
dem.set_friction_coefficient(friction_gate, "Grain", "Gate")
dem.set_friction_coefficient(friction_wall, "Grain", "Wall")

# %%
# initial grains deposition
# -------------------------
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)  # probability density function [0,1]
d = (
    diameter_max * diameter_min / (diameter_max + u * (diameter_min - diameter_max))
)  # random particle diameter following the given "u" law
radii = d / 2
boundaries = np.array([[0, H], [0, h_0], [L_0, h_0], [L_0, H]])
xp, radii = deposit.fastdeposit_2d(boundaries, radii, "deposit.svg")
eps = 1e-4 * diameter_max
for x, r in zip(xp, radii):
    if x[1] < h_0 + H_0:
        dem.add_particle(x + eps, r, np.pi * r**2 * rhop, "Grain")

# %%
# Simulation Loop
# ---------------
# Time integration of coupled fluid–particle motion.
i = 0
gmsh.option.setNumber("General.Verbosity", 0)
while t < tEnd:
    print(f"{i:4d}, {t:.6g}/{tEnd:.6g}, {dt:.6g}")
    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(
        i, alphaDomainTag, f, nodetag, oelemtag, oparamcoord
    )
    vol = dem.volume() * volume_corr
    geo = vc.get_particles(
        f,
        dem.position(),
        vc.body_radii(vol, 2),
        density=vc.body_density(dem, volume=vol,
                                delassus=dem.delassus() * volume_corr),
    )
    clo = vc.get_particles_closure(
        f, geo, mu, rho, dem.velocity(), omega=None,
        contact_forces=dem.contact_forces(),
    )
    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 i % outf == 0:
        f.write_mig(outputdir, t)
        dem.write_mig(outputdir, t)

    f.implicit_euler(dt)

    if t < settle_time:
        dem.body_velocity()[gate_body, 1] = 0
    else:
        dem.body_velocity()[gate_body, 1] = u_gate

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

    if y_gate > H / 3:
        dem.body_velocity()[gate_body, 1] = 0
    y_gate += dem.body_velocity()[gate_body, 1] * dt
    geoEntities = set_geo_mesh(y_gate)

    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.01 * smin,
    )

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

    f.set_coordinates(newCoords)

    t += dt
    i += 1

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

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