# %%
# 2D Fluid-Grains dam break
# =========================
# This example simulates a 2-phase dam break, consisting of a column of fluid and
# partially immersed grains
# We follow the setup described in
# X. Sun, M. Sakai, Y. Yamada, Three-dimensional simulation of a solid–liquid flow by the DEM–SPH method, J. Comput. Phys. (2013)

# Keywords
# --------
# PFEM, fluid-grains, dam break, 2D
#
# Description
# -----------
# The test demonstrates:
# - How to perform a PFEM-DEM simulation with moving boundaries.
# - A validation case for fluid-grains interaction.
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, scontact, time_integration, pfem, deposit
from migflow import volume_coupling as vc

np.random.seed(0)

# %%
# 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_f = 0.05  # length of the fluid column
H_f = 0.1  # height of the fluid column
l = 0.2  # length of the box
h = 1.0  # height of the box
t_gate = 0.05 * L_f  # thickness of the gate
h_gate = h  # height of the gate
y_gate = 1e-8  # initial position of the gate
u_gate = 0.68  # velocity of the gate
L_g = L_f  # length of the grains column
H_g = H_f / 3  # height of the grains column
r_particles = 0.00135  # radius of the grains

# %%
# Mesh parameters
# ---------------
geo_mesh_size = l / 10
mesh_size = 0.1 * L_f
smin = 1.4 * r_particles
smax = 4 * smin
dmax = 4 * smax
alpha = 1.3

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

gmsh.model.add("ModelInit")
gmsh.model.occ.addRectangle(l - L_f, 0, 0, L_f, H_f)
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)
    gate = gmsh.model.occ.addRectangle(
        l - L_f - t_gate, y_gate, 0, t_gate, h_gate
    )  # gate
    gmsh.model.occ.cut([(2, box)], [(2, gate)])
    gmsh.model.occ.synchronize()
    gmsh.model.mesh.setSizeCallback(lambda *args: geo_mesh_size)
    gmsh.model.mesh.generate(2)
    geoEntities = gmsh.model.getEntities(1)
    gmsh.write(outputdir + "/lastGeo.msh")
    gmsh.model.setCurrent(name)
    return geoEntities


geoEntities = set_geo_mesh()
gmsh.model.addPhysicalGroup(1, [12], -1, "bottom")
gmsh.model.addPhysicalGroup(1, [11], -1, "right")
gmsh.model.addPhysicalGroup(1, [6], -1, "left")
gmsh.model.addPhysicalGroup(1, [5, 8, 9], -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)
p0 = gmsh.model.addPhysicalGroup(1, [12], -1, "bottom")
p1 = gmsh.model.addPhysicalGroup(1, [11], -1, "right")
p2 = gmsh.model.addPhysicalGroup(1, [6], -1, "left")
p3 = gmsh.model.addPhysicalGroup(1, [5, 8, 9], -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.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,
)

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

# %%
# DEM parameters
# --------------
rhop = 2500
friction = 0.2  # 0.2
volume_corr = 0.8  # volume corection on the grains to account for 2D-3D discrepancies
mass_grains = 0.2
n_particles = int(mass_grains / (np.pi * r_particles**2 * rhop * L_g))

# %%
# Time parameters
# ---------------
cfl = 0.2
U = 2.5
U_init = 0.3
dt = smin / U * cfl
t = 0
settle_time = 0.06
tEnd = 1.0 + settle_time

# %%
# 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_strong_boundary("bottom", velocity_y=0)
f.set_strong_boundary("right", velocity_x=0)
f.set_strong_boundary("left", velocity_x=0)
f.set_open_boundary("freeSurface", pressure=0, viscous_flag=False)

# %%
# DEM Problem
# -----------
dem = scontact.ParticleProblem(2)
dem.set_friction_coefficient(friction)

# add outer box
outer_body = dem.add_body((0, 0), 0, 0)
x0 = np.array([0, 0])
x1 = np.array([l, 0])
x2 = np.array([l, h])
x3 = np.array([0, h])
dem.add_segment_to_body(x0, x1, outer_body, "bnd", material="wall")
dem.add_segment_to_body(x1, x2, outer_body, "bnd", material="wall")
dem.add_segment_to_body(x2, x3, outer_body, "bnd", material="wall")
dem.add_segment_to_body(x3, x0, outer_body, "bnd", 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)

# add gate
gate_body = dem.add_body((0, 0), 0, 0)
x0 = np.array([l - L_f - t_gate, y_gate])
x1 = np.array([l - L_f, y_gate])
x2 = np.array([l - L_f, h_gate])
x3 = np.array([l - L_f - t_gate, h_gate])
dem.add_segment_to_body(x0, x1, gate_body, "bnd", material="wall")
dem.add_segment_to_body(x1, x2, gate_body, "bnd", material="wall")
dem.add_segment_to_body(x2, x3, gate_body, "bnd", material="wall")
dem.add_segment_to_body(x3, x0, gate_body, "bnd", material="wall")
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


# %%
# initial grains deposition
# -------------------------
ratio = 1.2
diameter_max = 2.0 * r_particles
diameter_min = diameter_max / ratio
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([[l - L_g, 0], [l, 0], [l, 2 * H_g], [l - L_g, 2 * H_g]])
xp, radii = deposit.fastdeposit_2d(boundaries, radii, outputdir + "/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)

# %%
# Simulation Loop
# ---------------
# Time integration of coupled fluid–particle motion.
i = 0
outf = 10
gmsh.option.setNumber("General.Verbosity", 0)
mass = np.pi * dem.r() ** 2 * rhop

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 * smin,
        usePreviousMesh=True,
    )
    oparamcoord = oparamcoord.reshape((-1, 3))
    gmsh.write(outputdir + "/lastMesh.msh")
    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:  # or t > 0.3:
        f.write_mig(outputdir, t)
        dem.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)

    # move gate
    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
    time_integration._advance_particles(dem, external_forces, dt, 1, 1e-4 * r_particles)

    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
    dx = dx.reshape((-1,))

    # update gate position and geo mesh
    if y_gate > h / 2:
        dem.body_velocity()[gate_body, 1] = 0
    y_gate += dem.body_velocity()[gate_body, 1] * dt
    geoEntities = set_geo_mesh(y_gate)

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

    # 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 --actors fluid particles --fluid-field velocity --show-edges 1 --bounds 0 0.2 0 0.2
# %%
# Artifacts
# ---------
# - output/animation.mp4
# - output/deposit.svg