# %%
# Semi-Resolved Avalanche
# =======================
# This example simulates a two-dimensional dense granular avalanche using a
# semi-resolved fluid–particle coupling. In this formulation, the mesh size
# of the fluid is of the same order as the particle diameter, allowing the
# hydrodynamic forces to be partially resolved around individual particles.
#
# A set of polydisperse particles is first deposited inside a rectangular
# domain. The particles interact through frictional contacts,
# while a viscous incompressible fluid surrounds them and exerts
# drag and buoyancy. The simulation illustrates the onset of an avalanche and
# the interplay between solid contact chains and fluid stresses.
# %%
# Keywords
# --------
# FEM, DEM, Semi-Resolved, Adaptation, Friction, Polydispersity
import sys, os, shutil
import gmsh
import numpy as np
from scipy.spatial import KDTree

gmsh.initialize()

# %%
# This testcase relies on the LMGC90 library
sys.path.insert(0, os.environ["LMGC90_DIR"])

from pylmgc90 import pre
from migflow import fluid, scontact, time_integration, gmsh_io, volume_coupling as vc

np.random.seed(0)


# %%
# Output Directory
# ----------------
# The output directory is created (deleted first if it already exists).
outputdir = "output_2d_avalanche_semi_resolved" if len(sys.argv) < 2 else sys.argv[1]
shutil.rmtree(outputdir, ignore_errors=True)
os.makedirs(outputdir)

# %%
# Physical Parameters
# -------------------
# Fluid and particle properties are defined here. The particles are
# denser than the fluid to initiate motion under gravity.
rho = 1000.0  # fluid density [kg/m³]
mu = 1e-3  # fluid viscosity [Pa·s]
gravity = np.array([0.0, -9.81])  # gravity vector [m/s²]
rhop = 2500.0  # particle density [kg/m³]
dmax = 1e-2  # maximum particle diameter [m]
dmin = 5e-3  # minimum particle diameter [m]
polydispersity = dmax / dmin  # size ratio

# %%
# Geometrical Parameters
# ----------------------
# The computational domain is a rectangular box.
# The mesh is generated so that the element size is of the same order
# as the particle diameter (semi-resolved scale).
height = 0.4  # domain height [m]
width = 0.8  # domain width [m]
mesh_size = 10 * dmin  # mesh size [m]
origin = np.array([-width / 2, -height / 2])  # mesh origin [m]
eps = 1e-8  # numerical tolerance [m]
h0 = 0.1875  # initial particle bed height [m]
w0 = 1e-1  # initial particle bed width [m]


# %%
# Mesh Generation
# ---------------
# The Gmsh model is generated with physical groups for the walls and domain.
# These will be used later for setting boundary conditions in the fluid and
# particle solvers.
# The local mesh size is controlled using the distance to the particle cloud.
# This ensures refinement near the granular region and coarser elements
# elsewhere.
lc = 2 * dmax
lcmin = lc / 2
lcmax = lc
distmin = 10 * dmin


class MeshSize:
    def __call__(self, dim, tag, x, y, z, lc):
        dist, _ = self.tree.query([[x, y]])
        alpha = np.clip((dist[0] - distmin) / (1 * distmin), 0, 1)
        return lcmin * (1 - alpha) + lcmax * alpha

    def set_points(self, points):
        self.tree = KDTree(points)


def gen_mesh(width, height, mesh_size, origin=np.array([0, 0])):
    origin = np.asarray(origin)
    gmsh.model.add("box")
    gmsh.model.occ.add_rectangle(origin[0], origin[1], 0, width, height)
    gmsh.model.occ.synchronize()

    def get_line(x0, x1, eps=1e-6):
        r = gmsh.model.get_entities_in_bounding_box(
            x0[0] - eps, x0[1] - eps, -eps, x1[0] + eps, x1[1] + eps, eps, 1
        )
        return list(tag for dim, tag in r)

    h, w = height, width
    gmsh.model.add_physical_group(
        1, get_line(origin + [0, 0], origin + [w, 0]), name="Bottom"
    )
    gmsh.model.add_physical_group(
        1, get_line(origin + [0, h], origin + [w, h]), name="Top"
    )
    gmsh.model.add_physical_group(
        1, get_line(origin + [0, 0], origin + [0, h]), name="Left"
    )
    gmsh.model.add_physical_group(
        1, get_line(origin + [w, 0], origin + [w, h]), name="Right"
    )
    gmsh.model.add_physical_group(2, [1], name="domain")
    gmsh.model.mesh.set_size_callback(mesh_size)
    gmsh.model.mesh.generate(2)


# %%
# Particle Problem
# ----------------
# The particle assembly is initialized by randomly depositing polydisperse
# grains inside a rectangular region. Each grain is represented by a circular
# particle with its own radius and mass. Particles are deposited inside the box using the LMGC preprocessor.
# This generates a dense initial packing.
p = scontact.ParticleProblem(2)
n_particles_max = int(2e5)  # upper bound for number of particles
u = np.random.power(1, n_particles_max)  # probability density function [0,1]
d = (
    dmax * dmin / (dmax + u * (dmin - dmax))
)  # random particle diameter following the given "u" law
rp = d / 2
[nb_deposited, x, rp] = pre.depositInBox2D(rp, w0, h0)
xp = x.reshape(-1, 2)
xp += origin[None, :]
for xi, ri in zip(xp, rp):
    p.add_particle(xi, ri, np.pi * ri**2 * rhop, "particle")

# %%
# Mesh Adaptation
# ---------------
# The adaptive mesh size field is built from the particle positions and used
# to generate the mesh.
mesh_size = MeshSize()
mesh_size.set_points(p.position()[p.r()[:, 0] > 0])
gen_mesh(width, height, mesh_size, origin)

# %%
# Boundary and Friction Conditions
# --------------------------------
# Wall boundaries are imported from the mesh and assigned frictional contact
# properties. Particle–wall and particle–particle friction coefficients are
# also defined.
gmsh_io.load_msh_boundaries(p, None, ["Top", "Left", "Right", "Bottom"], material="wall")
p.set_friction_coefficient(1.3, "wall", "particle")
p.set_friction_coefficient(0.4, "particle", "particle")


# %%
# Fluid Problem
# -------------
# The fluid problem is created and initialized from the Gmsh mesh. No-slip
# boundary conditions are imposed on all walls. The coupling is built in tiers:
# tier A is the body geometry on the current mesh, tier B the legacy mixture
# closure on that geometry.
f = fluid.FluidProblem2(gravity, mu, rho)
gmsh_io.load_msh(f, None)
f.set_wall_boundary("Bottom", velocity=[0, 0])
f.set_wall_boundary("Left", velocity=[0, 0])
f.set_wall_boundary("Right", velocity=[0, 0])
f.set_wall_boundary("Top", velocity=[0, 0])
f.set_mean_pressure(0)
# Tiers A and B alone here: they prime the coupling so the first frame carries a
# porosity field. The loop rebuilds both and adds the Patankar datum, which
# needs the time step chosen below.
geo = vc.get_particles(f, p.position(), p.r(), density=vc.body_density(p))
clo = vc.get_particles_closure(f, geo, mu, rho, p.velocity(),
                               omega=p.omega(),
                               contact_forces=p.contact_forces())
vc.set_coupling_geometry(f, geo)
vc.set_coupling_closure(f, clo)

# %%
# Time Integration
# ----------------
# The simulation runs for a short time interval with a fixed time step.
# At each iteration, the fluid phase is solved and the fluid-grain forces are applied to the particles,
# and then system evolves according to the contact dynamics solver.
# The simulation runs until the avalanche stabilizes.
t = 0
i = 0
dt = 1e-3
tEnd = 4
outf = 50
mass = np.pi * p.r() ** 2 * rhop
while t < tEnd:
    print(f"{i = :4d} ----- {t = :g}")
    if i % outf == 0:
        f.write_mig(outputdir, t)
        p.write_mig(outputdir, t)
    f.implicit_euler(dt, check_residual_norm=1)
    fext = gravity * mass + f.compute_node_force()
    time_integration._advance_particles(p, fext, dt, 1, 1e-3 * dmin)
    geo = vc.get_particles(f, p.position(), p.r(), density=vc.body_density(p))
    clo = vc.get_particles_closure(f, geo, mu, rho, p.velocity(),
                                   omega=p.omega(),
                                   contact_forces=p.contact_forces())
    # The grains are free and the next iteration's implicit_euler is the
    # Patankar step that advances them, so this generation carries a datum.
    dat = vc.get_patankar(geo, clo, dt, gravity, p.velocity())
    vc.set_coupling_geometry(f, geo)
    vc.set_coupling_closure(f, clo)
    vc.set_coupling_datum(f, dat)
    t += dt
    i += 1

# %%
# Plot
# ----
# .. code-block:: shell
#
#  python3 -m migflow.plot.migplot output_2d_avalanche_semi_resolved --actors fluid particles --fluid-field pressure
# %%
# Artifacts
# ---------
# - output_2d_avalanche_semi_resolved/animation.mp4