# MigFlow - Copyright (C) <2010-2026>
# <Universite catholique de Louvain (UCL), Belgium
#  Universite de Montpellier, France>
#
# List of the contributors to the development of MigFlow: see AUTHORS file.
# Description and complete License: see LICENSE file.
#
# This program (MigFlow) is free software:
# you can redistribute it and/or modify it under the terms of the GNU Lesser General
# Public License as published by the Free Software Foundation, either version
# 3 of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program (see COPYING and COPYING.LESSER files).  If not,
# see <http://www.gnu.org/licenses/>.

#!/usr/bin/env python

# %%
# Falling particles in a sandglass
# ================================
# This example illustrates the fall of dense solid particles in a
# viscous fluid under gravity in a two-dimensional hourglass.
#
# Keywords
# --------
# DEM, FEM
#
# Description
# -----------
# The test demonstrates:
# - How to initialize a deposit of circular particles.
# - How to couple the particle solver (`scontact`) and the fluid solver (`fluid`).
from migflow import fluid, scontact, time_integration, gmsh_io, volume_coupling as vc
import numpy as np
import os, shutil, subprocess, sys
import random

# %%
# 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
# ------------------------------------------
y0_part = 40e-3  # vertical center of the grains
ly_part = 20e-3  # semi-height of the grains
lx_part = 0.015  # semi-width of the grains
meshfile = "2d_mesh"
subprocess.run(["gmsh", "-2", meshfile + ".geo", "-o", meshfile + ".msh"])

# %%
# Physical Parameters
# -------------------
g = np.array([0, -9.81])  # gravity
rho = 1  # fluid density
rhop = 1500  # grains density
nu = 1e-5  # kinematic viscosity
r = 1e-4  # grains radius

# %%
# Simulation Parameters
# ---------------------
dt = 2e-4  # time step
tEnd = 0.5  # final time
outf = 10  # number of iterations between output files
alpha = 0.8  # 2d-3d correction factor for the fluid-particle volume coupling
use_lmgc90 = False  # enable the use of lmgc90

# %%
# Particle Problem Initialization
# -------------------------------
p = scontact.ParticleProblem(2)

# Loading of the mesh.msh file specifying physical boundaries name
gmsh_io.load_msh_boundaries(p, f"{meshfile}.msh", ["Top", "Box"], material="Material1")

# Generation of the particles in a rectangular grid
x = np.arange(r, lx_part, 2 * r)
y = np.arange(y0_part + r, y0_part + ly_part, 2 * r)
for i in range(x.shape[0]):
    for j in range(y.shape[0]):
        rr = r * (0.95 + random.random() * 0.05)
        p.add_particle((x[i], y[j]), rr, rr**2 * np.pi * rhop, "Material2")

if use_lmgc90:
    from migflow import lmgc90Interface
    print("Deprecated: lmgc90 interface is no longer supported.")
    pass
else:
    p.set_friction_coefficient(0.4, "Material2", "Material2")  # between particles
    p.set_friction_coefficient(
        0.5, "Material1", "Material2"
    )  # between particles and boundaries

# %%
# Fluid Problem Initialization
# ----------------------------
f = fluid.FluidProblem2(g, nu * rho, rho)
gmsh_io.load_msh(f, f"{meshfile}.msh")
f.set_wall_boundary("Top", velocity=[0, 0])
f.set_wall_boundary("Box", velocity=[0, 0])
f.set_mean_pressure(0)

# %%
# Simulation Loop
# ---------------
# Time integration of coupled fluid–particle motion.
t = 0
i = 0
while t < tEnd:
    print(f"{i:4d}, {t:.6g}/{tEnd:.6g}")
    # The 2d-3d correction scales the body volume, so tiers A and B must see the
    # corrected body: its radius and its density both come from that volume.
    # The grains are free and the implicit_euler below is the Patankar step that
    # advances them, so this generation carries a datum. No omega: the legacy
    # call passed p.omega()*0, i.e. a purely translational datum.
    vol = alpha * p.volume()
    geo = vc.get_particles(f, p.position() + p.velocity() * dt,
                           vc.body_radii(vol, 2),
                           density=vc.body_density(p, volume=vol))
    clo = vc.get_particles_closure(f, geo, nu * rho, rho, p.velocity(),
                                   omega=None,
                                   contact_forces=p.contact_forces())
    dat = vc.get_patankar(geo, clo, dt, g, p.velocity())
    vc.set_coupling_geometry(f, geo)
    vc.set_coupling_closure(f, clo)
    vc.set_coupling_datum(f, dat)
    if i % outf == 0:
        p.write_mig(outputdir, t)
        f.write_mig(outputdir, t)
    f.implicit_euler(dt)
    forces = rhop * p.r() ** 2 * np.pi * g + f.compute_node_force()
    time_integration._advance_particles(p, forces, dt, min_nsub=1, contact_tol=1e-3 * r)
    t += dt
    i += 1

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