# %%
# Bidimensional body Sedimentation (Hexagonal Packing)
# ========================================================
# This example illustrates the sedimentation of dense solid particles in a
# viscous fluid under gravity in two dimensions using MigFlow.
# Particles are initialized in a dense hexagonal packing.
#
# Keywords
# --------
# DEM, FEM, Generation
#
# Description
# -----------
# The test demonstrates:
# - How to generate a 2D rectangular fluid domain with named boundaries in Gmsh.
# - How to initialize a dense cloud of circular particles.
# - How to couple the particle solver (`scontact`) and the fluid solver (`fluid`).
# - How to export derived fields (pressure, velocity, porosity, dynamic pressure) for post-processing in Paraview or similar tools.
import os, sys, shutil
import numpy as np
import gmsh
from migflow import fluid, scontact, time_integration, gmsh_io
from migflow import volume_coupling as vc

gmsh.initialize()

# %%
# Output Directory
# ----------------
# Create a clean output directory for simulation results.
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("outputdir", nargs="?", default="output_2d_depot_hex_body")
parser.add_argument("--tend", type=float, default=3,
                    help="final simulation time [s]")
parser.add_argument("--body", action=argparse.BooleanOptionalAction, default=True,
                    help="the library's UNFITTED coupling (the default, and "
                         "what 2d_depot_hex runs); --no-body substitutes the "
                         "legacy mixture retrieval through the set_bodies hook, "
                         "which is the comparison this case exists to make")
args = parser.parse_args()
outputdir = args.outputdir
shutil.rmtree(outputdir, ignore_errors=True)
os.makedirs(outputdir)

USE_BODY = args.body

# %%
# Geometrical parameters and mesh generation
# ------------------------------------------
# A rectangular 2D mesh is generated using Gmsh with named boundaries.
height = 0.6  # domain height [m]
width = 0.4  # domain width [m]
mesh_size = 0.01  # element size [m]
h = 0.15  # particle bed height [m]
w = 0.4  # particle bed width [m]
origin = [-width / 2, -height / 2]  # mesh origin (bottom-left)
eps = 1e-8  # small numerical tolerance


def gen_mesh(width, height, mesh_size, origin=np.array([0, 0])):
    """Generate a rectangular 2D mesh with physical boundaries."""
    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 [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(lambda dim, tag, x, y, z, lc: mesh_size)
    gmsh.model.mesh.generate(2)


gen_mesh(width, height, mesh_size, origin)


# %%
# Physical Parameters
# -------------------
g = np.array([0, -9.81])  # gravity [m/s²]
r = 1.5e-3  # particle radius [m]
rhop = 1500  # particle density [kg/m³]
rho = 1000  # fluid density [kg/m³]
nu = 1e-6  # kinematic viscosity [m²/s]
mu = rho * nu  # dynamic viscosity [Pa·s]

# %%
# Particle Problem Initialization
# -------------------------------
# Particles are placed in a hexagonal patch at the top of the domain.
p = scontact.ParticleProblem(2)
p.set_fixed_contact_geometry(0)
gmsh_io.load_msh_boundaries(p, None, ["Top", "Left", "Right", "Bottom"])

# Generate a hexagonal dense packing of particles
step = 2 * r + eps
lx, ly = w, h
y0 = height / 2 - r - eps
x = np.arange(-lx / 2, lx / 2 - 2 * r, step)
y = np.arange(y0, y0 - h, -np.sqrt(3) * step)
x2 = np.arange(-lx / 2 + r, lx / 2 - 2 * r, step)
y2 = np.arange(y0, y0 - h, -np.sqrt(3) * step) - np.sqrt(3) * step / 2
x, y = np.meshgrid(x, y)
x2, y2 = np.meshgrid(x2, y2)
x = np.concatenate([x.ravel(), x2.ravel()])
y = np.concatenate([y.ravel(), y2.ravel()])
x_off = w - ((x.max() + step / 2) - (x.min() - step / 2))
x += r + x_off / 2
# Sort particles by height for body ordering, for better visualization
order = np.argsort(y)
x, y = x[order], y[order]
for xi, yi in zip(x, y):
    p.add_particle((xi, yi), r, np.pi * r**2 * rhop)

# %%
# Fluid Problem Initialization
# ----------------------------
f = fluid.FluidProblem2(g, mu, rho)
gmsh_io.load_msh(f, None)
for wall in ["Bottom", "Left", "Right", "Top"]:
    f.set_wall_boundary(wall)
f.set_strong_boundary("Left", velocity=[0, 0])
f.set_strong_boundary("Right", velocity=[0, 0])
f.set_mean_pressure(0)


def get_fields(fluid):
    """Return derived output fields for visualization."""
    y = fluid.coordinates_fields()[fluid.field_indices(fluid.dimension())][:, 1]
    p1_element = fluid.get_p1_element()
    return {
        "pressure": (fluid.pressure(), p1_element),
        "velocity": (fluid.velocity(), p1_element),
        "porosity": (fluid.porosity().get(), p1_element),
        "u_solid": (fluid.u_solid(), p1_element),
        "dynamic_pressure": (fluid.pressure() - rho * g[1] * y, p1_element),
    }


# %%
# Simulation Parameters
# ---------------------
outf = 1  # number of iterations between outputs
dt = 2.5e-3  # time step [s] -- the real-contact-force prediction needs
# dt*Fc/m below the flow scale; at 5e-3 the case diverges at t ~ 19, after the
# bed has settled and the remaining grains are in free flight (|Fc| = 0)
tEnd = args.tend  # final time [s] (--tend)
t = 0
i = 0

# %%
# Simulation Loop
# ---------------
# Time integration of coupled fluid–particle motion.
mass = np.pi * p.r() ** 2 * rhop


def legacy_tiers(fl, pp, v_full=None):
    """Tiers A and B in the LEGACY RETRIEVAL -- m = 0, div m = 1,
    beta = gamma_Dallavalle/(mu V_b), trace band and both extra diffusivities
    off. Tier C (the Patankar datum) is the SCHEME's job.

    This lives HERE and not in time_integration because no scheme selects the
    retrieval any more: the library runs the unfitted model and the retrieval is
    an instrument a case reaches for when it wants to measure against it."""
    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, rho, vel,
        omega=0 * pp.omega(), contact_forces=pp.contact_forces())
    vc.set_coupling_geometry(fl, geo)
    vc.set_coupling_closure(fl, clo)
    return dict(geo, **clo)


while t < tEnd:
    print(f"{i:4d}, {t:.6g}/{tEnd:.6g}")
    print("----------------- ")
    if i % outf == 0:
        f.write_mig(outputdir, t, get_fields(f))
        p.write_mig(outputdir, t)
    # Coupled step. --body (the default) is the plain library call: tier A
    # geometry, tier B UNFITTED closure, tier C Patankar datum, all three set by
    # name inside iterate_patankar. --no-body substitutes the legacy retrieval
    # for tier B and nothing else, so the two runs differ by one closure and are
    # directly comparable -- which is the whole point of this case.
    #
    # It used to call fl.set_particles_body() through the set_bodies hook. That
    # is the IN-CORE disc-overlap setter, which lands on the legacy mixture
    # kernel (bodies_set + particle_drag_coefficient) and never touches the
    # unfitted model, so "the body form" tested a second CODE PATH rather than a
    # second model and could not be compared with its sibling at all.
    if USE_BODY:
        time_integration.iterate_patankar(
            f, p, dt, 1, 1e-3 * r, external_particles_forces=mass * g)
    else:
        time_integration.iterate_patankar(
            f, p, dt, 1, 1e-3 * r, external_particles_forces=mass * g,
            gravity=g, patankar_datum=True, set_bodies=legacy_tiers,
            get_forces=lambda fl, pp: np.asarray(
                fl.get_forces_on_bodies()).reshape(-1, 2))
    t += dt
    i += 1

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