# %%
# 2D Lituya Bay landslide-generated tsunami simulation with PFEM-DEM
# =======================================
# This example simulates the Lituya bay landslide-generated tsunami using a PFEM-DEM approach.
# The fluid domain is represented using the PFEM method, while the granular landslide is modeled
# using the DEM approach.
# This test case follows works from
#  H.M. Fritz, F. Mohammed, J. Yoo, Lituya Bay landslide impact generated mega-tsunami 50 th anniversary, Tsunami Science Four Years After the 2004 Indian Ocean Tsunami: Part II: Observation and Data Analysis (2009) 153–175.
# The initial particle distribution is generated using a separate script: gen_particles.py

# Keywords
# --------
# PFEM, fluid-grains, dam break, 2D, Lituya bay
#
# Description
# -----------
# The test demonstrates:
# - How to perform a PFEM-DEM simulation with moving boundaries.
import numpy as np
import gmsh
from migflow import fluid, time_integration, scontact, pfem
import shutil
import os
import sys

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

# %%
# Run the particle generation script
# ----------------------------------
import subprocess

subprocess.run(["python3", "gen_particles.py"])


# %%
# Geometrical parameters and mesh generation
# ------------------------------------------
bassin_length = 1342
bassin_depth = 122
theta = np.pi / 4
length_after_bassin = bassin_length
length_before_bassin = bassin_length
r_particles = 4

# %%
# Mesh parameters
# ---------------
alpha = 1.2
mesh_size = 0.01 * bassin_length
geo_mesh_size = 0.05 * bassin_length
smin = 1 * r_particles
smax = 6 * smin
dmax = 6 * smax
tol_bnd = 1e-5 * geo_mesh_size

# %%
# Initial fluid mesh
# ------------------
gmsh.model.add("Model_Init")
disp = 0.0
p0 = gmsh.model.geo.addPoint(
    disp + bassin_depth / np.tan(theta), -bassin_depth, 0.0, mesh_size
)
p1 = gmsh.model.geo.addPoint(
    bassin_length - bassin_depth / np.tan(theta), -bassin_depth, 0.0, mesh_size
)
p2 = gmsh.model.geo.addPoint(bassin_length, 0.0, 0.0, mesh_size)
p3 = gmsh.model.geo.addPoint(disp, 0.0, 0.0, mesh_size)
l0 = gmsh.model.geo.addLine(p0, p1)
l1 = gmsh.model.geo.addLine(p1, p2)
l2 = gmsh.model.geo.addLine(p2, p3)
l3 = gmsh.model.geo.addLine(p3, p0)
cloop = gmsh.model.geo.addCurveLoop([l0, l1, l2, l3])
s = gmsh.model.geo.addPlaneSurface([cloop])
gmsh.model.geo.synchronize()
gmsh.model.mesh.set_size_callback(lambda *args: mesh_size)
gmsh.model.mesh.generate(2)
nodeTags = coords = None
nodeTags, coords, _ = gmsh.model.mesh.get_nodes(2, -1, True)

# %%
# Solid domain
# ------------
gmsh.model.add("ModelGeo")
p0 = gmsh.model.geo.addPoint(
    0.0 + bassin_depth / np.tan(theta), -bassin_depth, 0.0, geo_mesh_size
)
p1 = gmsh.model.geo.addPoint(
    bassin_length - bassin_depth / np.tan(theta), -bassin_depth, 0.0, geo_mesh_size
)
p2 = gmsh.model.geo.addPoint(
    bassin_length - bassin_depth / np.tan(theta) + length_after_bassin,
    -bassin_depth + length_after_bassin / np.tan(theta),
    0.0,
    geo_mesh_size,
)
p3 = gmsh.model.geo.addPoint(
    0.0 + bassin_depth / np.tan(theta) - length_before_bassin,
    -bassin_depth + length_before_bassin / np.tan(theta),
    0.0,
    geo_mesh_size,
)
l0 = gmsh.model.geo.addLine(p0, p1)
l1 = gmsh.model.geo.addLine(p1, p2)
l2 = gmsh.model.geo.addLine(p2, p3)
l3 = gmsh.model.geo.addLine(p3, p0)
cloop = gmsh.model.geo.addCurveLoop([l0, l1, l2, l3])
s = gmsh.model.geo.addPlaneSurface([cloop])
gmsh.model.geo.synchronize()
gmsh.model.add_physical_group(1, [l0], -1, "bottom")
gmsh.model.add_physical_group(1, [l1], -1, "right")
gmsh.model.add_physical_group(1, [l2], -1, "top")
gmsh.model.add_physical_group(1, [l3], -1, "left")
geoEntities = gmsh.model.getEntities(1)
gmsh.model.mesh.generate(2)

# %%
# PFEM mesh and size fields
# -------------------------
gmsh.model.add("Model_Fluid")
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, [1], -1, "bottom")
gmsh.model.addPhysicalGroup(1, [2], -1, "right")
gmsh.model.addPhysicalGroup(1, [3], -1, "top")
gmsh.model.addPhysicalGroup(1, [4], -1, "left")
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", -np.inf)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "XMax", np.inf)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "YMin", -np.inf)
gmsh.model.mesh.field.setNumber(sizeFieldConstant, "YMax", np.inf)
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)
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
rho_bubble = 1
mu = 1e-3
sigma = 0.0071
use_bubble = False
surfaceTension = False
rhop = 2640

# %%
# Time parameters
# ---------------
U = 100
U_init = U
cfl = 0.2
dt = smin / U * cfl
t = 0
tEnd = 75.0
print(f"dt = {dt}, U = {U}")

# %%
# Fluid problem
# -------------
f = fluid.FluidProblem(2, g, mu, rho, advection=False)
f.set_wall_boundary("bottom")
f.set_wall_boundary("right")
f.set_wall_boundary("top")
f.set_wall_boundary("left")
f.set_strong_boundary("bottom", velocity_y=0)
# f.set_strong_boundary("right", velocity = [0, 0])
# f.set_strong_boundary("top", velocity = [0, 0])
# f.set_strong_boundary("left", velocity = [0, 0])
f.set_open_boundary("freeSurface", pressure=0, viscous_flag=False)


# %%
# DEM problem
# -------------
dem = scontact.ParticleProblem(2)
inputdir = "output_initial_particles"
dem.read_mig(inputdir, 0)
v = np.pi * np.sum(dem.r() ** 2)
mass_paper = 98.5e6  # kg
volume_corr = mass_paper / (rhop * v)

# %%
# Simulation Loop
# ---------------
# Time integration of coupled fluid–particle motion.

i = 0
outf = 20
gmsh.option.setNumber("General.Verbosity", 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,
        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
    )
    f.set_particles(
        dem.delassus() * volume_corr,
        dem.volume() * volume_corr,
        dem.position(),
        dem.velocity(),
        dem.omega() * 0,
        dem.contact_forces(),
    )

    if surfaceTension:
        pfem.applySurfaceTension(f, sigma, False, g, rho_bubble)

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

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

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

    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,))

    # advect nodes and project if needed
    gmsh.model.mesh.advectMeshNodes(
        2,
        alphaDomainTag,
        alphaBoundaryTag,
        "ModelGeo",
        ordered_node_tags,
        dx,
        0.01 * smin,
    )
    newTags, newCoords, _ = gmsh.model.mesh.get_nodes()
    new_order = np.argsort(newTags)
    newCoords = newCoords.reshape((-1, 3))
    newCoords = newCoords[new_order]

    # set new coordinates
    f.set_coordinates(newCoords)

    # update time step
    t += dt
    i += 1
    max_U = np.max([U_init, np.max(np.linalg.norm(f.velocity(), axis=1))])
    dt = smin / max_U * cfl