#!/usr/bin/env python

# 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/>.

# %%
# Drag on a Fixed Cylinder in a 3D Immersed Granular Flow
# =======================================================
# This example studies the drag force acting on a fixed cylindrical obstacle
# immersed in a granular flow saturated by a viscous fluid.
#
# Keywords
# --------
# Drag, DEM, FEM, Boundary Force, 3D Flow
#
# Description
# -----------
# The test demonstrates:
# - How to set up a mixed granular–fluid simulation in a 3D geometry.
# - How to impose an inflow boundary condition to simulate a stream.
# - How to insert, remove, and constrain particles dynamically.
# - How to compute and record the total drag force acting on a solid obstacle.

import sys
from migflow import scontact, gmsh_io, volume_coupling as vc
from migflow import fluid
from migflow import time_integration

import numpy as np
import os, subprocess
import shutil


# %%
# Utility Functions
# -----------------
def gen_rect(origin, w, h, d, step):
    """
    Generate coordinates on a rectangular grid.

    Parameters
    ----------
    origin : array-like
        Origin of the top-left corner
    w : float
        Width
    h : float
        Height
    d : float
        Depth
    step : float
        Maximal radius
    """
    eps = 1e-8
    x = np.arange(-w / 2 + step - eps, w / 2 - step + eps, 2 * step) + origin[0]
    y = np.arange(step, h - step, 2 * step) + origin[1]
    z = np.arange(step, d - step, 2 * step) + origin[2]
    x, y, z = np.meshgrid(x, y, z)
    return x.reshape(-1), y.reshape(-1), z.reshape(-1)


# %%
# Physical Parameters
# -------------------
vit = -0.05  # stream velocity [m/s]
rhop = 2500.0  # particle density [kg/m³]
rho = 1000.0  # fluid density [kg/m³]
r = 6e-3  # particle radius [m]
g = np.array([0.0, -9.81, 0.0])  # gravity [m/s²]
nu = 1e-6  # fluid kinematic viscosity [m²/s]

# %%
# Geometrical Parameters
# ----------------------
h = 0.3  # particles area height [m]
w = 0.28  # particles area width [m]
d = 0.05  # particles area depth [m]
H = 1.4  # domain height [m]

# %%
# Numerical Parameters
# --------------------
dt = 1e-3  # time step [s]
t = 0.0  # initial time [s]
t_end = 1.0  # final time [s]
ii = 0  # iteration number
outf = 10  # output frequency (iterations)

# %%
# Output Directory and Mesh Generation
# ------------------------------------
# Create a clean output directory and generate 3D mesh using Gmsh.
outputdir = "output_3d_drag" if len(sys.argv) < 2 else sys.argv[1]
shutil.rmtree(outputdir, True)
os.makedirs(outputdir)

mshname = "3d_mesh"
subprocess.run(["gmsh", "-3", f"{mshname}.geo", "-o", f"{mshname}.msh"])

# %%
# Particle Problem Initialization
# -------------------------------
# Load mesh boundaries and define contact properties.
p = scontact.ParticleProblem(3)
gmsh_io.load_msh_boundaries(p, f"{mshname}.msh", ["Inner", "Left", "Right", "Front", "Rear"], material="Steel")

# %%
# Particle Generation
# -------------------
# Initial particle cloud around the cylinder
x, y, z = gen_rect([0.0, -0.7, 0.0], w, 4 * h, d, r)
for xi, yi, zi in zip(x, y, z):
    if xi**2 + yi**2 > (0.025 + r) ** 2:
        mass = 4.0 * np.pi * r**3 * rhop / 3.0
        p.add_particle((xi, yi, zi), r, mass)

# Inflow particles at the top
x, y, z = gen_rect([0.0, 0.5, 0.0], w, 0.5 * h, d, 1.1 * r)
for xi, yi, zi in zip(x, y, z):
    mass = 4.0 * np.pi * r**3 * rhop / 3.0
    p.add_particle((xi, yi, zi), r, mass)

# %%
# Material Properties
# -------------------
p.set_friction_coefficient(0.2, "Sand", "Sand")
p.set_friction_coefficient(0.2, "Sand", "Steel")
p.write_mig(outputdir, 0.0)

# %%
# Fluid Problem Initialization
# ----------------------------
# The fluid problem is loaded from the same mesh and coupled to the particle phase.
fluid_problem = fluid.FluidProblem3(g, nu * rho, rho)
gmsh_io.load_msh(fluid_problem, f"{mshname}.msh")

fluid_problem.set_wall_boundary("Inner")
fluid_problem.set_wall_boundary("Left", velocity=[0.0, vit, 0.0])
fluid_problem.set_wall_boundary("Right", velocity=[0.0, vit, 0.0])
fluid_problem.set_wall_boundary("Front", velocity=[0.0, vit, 0.0])
fluid_problem.set_wall_boundary("Rear", velocity=[0.0, vit, 0.0])
fluid_problem.set_open_boundary("Bottom", velocity=[0.0, vit, 0.0])
fluid_problem.set_open_boundary("Top", pressure=0.0)

# Tier A (the body geometry on the mesh) and tier B (the legacy mixture
# closure) only: this primes the coupling so the frame written below carries a
# porosity field. The loop rebuilds both, with the Patankar datum, before every
# solve.
geo = vc.get_particles(fluid_problem, p.position(), p.r(),
                       density=vc.body_density(p))
clo = vc.get_particles_closure(fluid_problem, geo, nu * rho, rho, p.velocity(),
                               omega=p.omega(),
                               contact_forces=p.contact_forces())
vc.set_coupling_geometry(fluid_problem, geo)
vc.set_coupling_closure(fluid_problem, clo)

fluid_problem.write_mig(outputdir, 0.0)

# %%
# Computation Loop
# ----------------
while t < t_end:

    posit = p.position()[p.r()[:, 0] != 0, :]

    # Constrain particles at the bottom
    mask = p.body_position()[:, 1] < -0.5
    p.body_invert_mass()[mask] = 0.0
    p.body_invert_inertia()[mask] = 0.0

    n_dyn = posit.shape[0]
    a = p.body_velocity()[p.n_bodies() - n_dyn :, :]
    b = p.body_omega()[p.n_bodies() - n_dyn :, :]
    c = p.body_invert_mass()[p.n_bodies() - n_dyn :, :]

    fixed = (c == 0)[:, 0]
    a[fixed, :] = [0.0, vit, 0.0]
    b[fixed, :] = [0.0, 0.0, 0.0]

    # Insert new particles at the top when needed
    if np.amax(posit[:, 1]) + r < 0.5:
        for xi, yi, zi in zip(x, y, z):
            mass = 4.0 * np.pi * r**3 * rhop / 3.0
            p.add_particle((xi, yi, zi), r, mass)

    # Remove particles that have fallen outside the domain
    p.remove_bodies_flag(p.body_position()[:, 1] > -0.55)

    # Compute forces
    mass = 4.0 * np.pi * p.r() ** 3 * rhop / 3.0

    # Set particles to the FEM module. The grains are free and the
    # implicit_euler below is the Patankar step that advances them, so this
    # generation carries a datum.
    geo = vc.get_particles(fluid_problem, p.position(), p.r(),
                           density=vc.body_density(p))
    clo = vc.get_particles_closure(fluid_problem, geo, nu * rho, rho,
                                   p.velocity(),
                                   omega=p.omega(),
                                   contact_forces=p.contact_forces())
    dat = vc.get_patankar(geo, clo, dt, g, p.velocity())
    vc.set_coupling_geometry(fluid_problem, geo)
    vc.set_coupling_closure(fluid_problem, clo)
    vc.set_coupling_datum(fluid_problem, dat)

    # Solve fluid and particle dynamics
    fluid_problem.implicit_euler(dt)
    fp = fluid_problem.compute_node_force(dt)

    time_integration._advance_particles(p, g * mass + fp, dt, 1, 1e-3 * r)

    t += dt

    # Write output files
    if ii % outf == 0:
        p.write_mig(outputdir, t)
        fluid_problem.write_mig(outputdir, t)

    ii += 1
    print(f"{ii:d} : {t:.2g}/{t_end:.2g}")

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