# 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
# %%
# 2D Particle Transport in Vortex Flows
# =====================================
# This example simulates the transport of small granular particles by a
# two-dimensional vortex flow generated around fixed cylindrical obstacles.
# The coupled fluid–grain problem is solved using MigFlow's FEM–DEM framework.
# %%
# Keywords
# --------
# FEM, DEM, Vortex, Particle Transport
import gmsh
from migflow import fluid, scontact, time_integration, gmsh_io, deposit
import numpy as np
import os, sys, shutil, 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)

# %%
# Physical parameters
# -------------------
g = np.array([0, -9.81])  # gravity
rho = 1000  # fluid density
nu = 5e-5  # kinematic viscosity
rhop = 2500  # grains density
r = 2e-5  # grains mean radius

# %%
# Geometrical parameters
# ----------------------
r_in = 2.5e-4
# Rayon des cercles
lc_in = r_in / 2
# Taille de maille aux frontieres interieures
lc_out = r_in
# Taille de maille aux frontieres exterieures
H = 2.6e-3 + 0.005
# Hauteur totale du domaine fluide
L = 8e-3
# Largeur totale du domaine fluide
c = 1.5e-3
# Demi-distance entre les cercles
h = 2e-3
# Hauteur des centres des cercles par rapport a l'origine

# %%
# Mesh generation
# ---------------
gmsh.initialize()
gmsh.model.add("vortices")
box_tag = gmsh.model.occ.add_rectangle(-L / 2, -H / 2, 0, L, H)
disc0_tag = gmsh.model.occ.add_disk(-c, -H / 2 + h, 0, r_in, r_in)
disc1_tag = gmsh.model.occ.add_disk(+c, -H / 2 + h, 0, r_in, r_in)
gmsh.model.occ.cut([(2, box_tag)], [(2, disc0_tag), (2, disc1_tag)])
gmsh.model.occ.synchronize()


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


gmsh.model.add_physical_group(2, [1], name="domain")
gmsh.model.add_physical_group(1, get_line(-L / 2, +L / 2, +H / 2, +H / 2), name="Top")
gmsh.model.add_physical_group(
    1, get_line(-L / 2, +L / 2, -H / 2, -H / 2), name="Bottom"
)
gmsh.model.add_physical_group(1, get_line(-L / 2, -L / 2, -H / 2, +H / 2), name="Left")
gmsh.model.add_physical_group(1, get_line(+L / 2, +L / 2, -H / 2, +H / 2), name="Right")
gmsh.model.add_physical_group(
    1,
    get_line(-c - r_in, -c + r_in, -H / 2 + h - r_in, -H / 2 + h + r_in),
    name="Disc0",
)
gmsh.model.add_physical_group(
    1,
    get_line(+c - r_in, +c + r_in, -H / 2 + h - r_in, -H / 2 + h + r_in),
    name="Disc1",
)
gmsh.model.mesh.set_size_callback(lambda *args: lc_in)
gmsh.model.mesh.generate(2)

# %%
# Particle problem initialization
# -------------------------------
p = scontact.ParticleProblem(2)
bnd_body = p.add_body((0, 0), 0, 0)
gmsh_io.load_mesh_to_body(p, bnd_body, None)

h_p = h - r_in
bnd_deposit = np.array(
    [
        (-L / 2, -H / 2 + h_p),
        (-L / 2, -H / 2),
        (+L / 2, -H / 2),
        (+L / 2, -H / 2 + h_p),
    ]
)
n_p_max = 1e4
radii = np.random.uniform(0.7 * r, r, int(n_p_max))
print("Generating deposit particle positions...")
xp, rp = deposit.fastdeposit_2d(bnd_deposit, radii, "deposit.svg")
for xi, ri in zip(xp, rp):
    p.add_particle((xi[0], xi[1]), ri, ri**2 * np.pi * rhop)

# %%
# Numerical parameters
# --------------------
tEnd = 3  # final time
t = 0
i = 0

# Two coupling paths, chosen by what is installed:
#  * overlap importable -> THE UNFITTED MODEL in its legacy-retrieval
#    setting (overlap CSR from migflow.volume_coupling): stable at dt = 5e-3
#    (that setting keeps the contact forces out of the prediction by
#    definition -- it reproduces the previous MigFlow model exactly).
#  * overlap absent -> the mixture path of iterate_patankar with the DEM
#    solver's contact forces in the prediction (library default). Stability is
#    then the time step's job: with r = 2e-5 the lagged bed force injects
#    dt*Fc/m ~ 20x the flow velocity at 5e-3 (measured blow-up i ~ 15) and
#    still 4x at 1e-3 (measured divergence t ~ 0.16); dt = 2.5e-4 rides it
#    (measured plateau vmax ~ 4e-3 through t = 0.2).
try:
    from migflow import overlap  # noqa: F401
    HAVE_SPHERE_PATCH = True
    dt = 5e-3
except ImportError:
    HAVE_SPHERE_PATCH = False
    dt = 2.5e-4
outf = int(round(0.05 / dt))  # a frame every 0.05 s on either path


# Function defining the rotation of the inner boundaries
# Parameters:
#              -x:     coordinates of the point on the circle (x for vertical velocity and y for horizontal velocity)
#              -c:     centre of the circle (go along with x)
#              -A:     amplitude of the periodic signal
#              -T:     periode
#              -s:     sign for the rotation
def Bnd(x, c, A, T, s):
    return s * A * (x[:] - c) * np.sin(t * np.pi * 2.0 / T)


yc = -H / 2 + h  # absolute y-centre
v_tang = 2e-2  # max absolute tangential velocity
omega = v_tang / r_in  # max absolute angular velocity
A = omega  # amplitude of periodic signal
T = 1  # period

# %%
# Fluid problem
# -------------
fluid = fluid.FluidProblem2(g, nu * rho, rho)
gmsh_io.load_msh(fluid, None)
walls = ["Top", "Bottom", "Left", "Right", "Disc0", "Disc1"]
for wall in walls:
    fluid.set_wall_boundary(wall)
fluid.set_mean_pressure(0)
fluid.set_strong_boundary("Disc1", velocity_x=lambda x: Bnd(x[:, 1], +yc, A, T, +1))
fluid.set_strong_boundary("Disc1", velocity_y=lambda x: Bnd(x[:, 0], +c, A, T, -1))
fluid.set_strong_boundary("Disc0", velocity_x=lambda x: Bnd(x[:, 1], +yc, A, T, -1))
fluid.set_strong_boundary("Disc0", velocity_y=lambda x: Bnd(x[:, 0], -c, A, T, +1))

# %%
# Computation loop
# ----------------
mass = p.volume() * rhop
hooks = {}
if HAVE_SPHERE_PATCH:
    # THE UNFITTED MODEL, legacy-retrieval setting: m = 0, div m = 1,
    # beta = gamma_Dallavalle/(mu V_p), trace stabiliser and the extra
    # velocity & pressure diffusivities off (see volume_coupling).
    from migflow import volume_coupling as vc
    _csr = {}

    def set_bodies(fl, pp, v_full=None):
        """Tiers A and B; the datum is the SCHEME's, built by
        iterate_patankar(patankar_datum=True) out of the dict returned here."""
        vel = pp.velocity() if v_full is None else v_full
        geo = vc.get_particles(fl, pp.position(), pp.r(),
                               density=vc.body_density(pp))
        clo = vc.get_particles_closure(fl, geo, mu=nu * rho, rho_f=rho,
                                       velocity=vel,
                                       contact_forces=None)
        vc.set_coupling_geometry(fl, geo)
        vc.set_coupling_closure(fl, clo)
        _csr["csr"] = dict(geo, **clo)
        return _csr["csr"]

    def get_forces(fl, pp):
        return vc.bodies_force(fl, _csr["csr"])

    hooks = dict(set_bodies=set_bodies, get_forces=get_forces,
                 patankar_datum=True)
    set_bodies(fluid, p)  # so the first frame carries the porosity field
else:
    print("overlap not installed: mixture path, dt =", dt, flush=True)

while t < tEnd:
    if i % outf == 0:
        print(f"t/tEnd={t:1.4f}/{tEnd:1.4f}, i={i:1d}", flush=True)
        fluid.write_mig(outputdir, t)
        p.write_mig(outputdir, t)

    time_integration.iterate_patankar(fluid, p, dt, 1, 1e-3 * r,
                                      external_particles_forces=mass * g,
                                      gravity=g, **hooks)

    t += dt
    i += 1

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