# %%
# Rotating Drum
# =============
# This example illustrates a circular shear flow generated by a rotating inner
# cylinder (drum) surrounded by an outer stationary wall. Grains are placed in
# the annular gap and interact with the fluid flow through drag and pressure force.
#
# Keywords
# --------
# FEM, DEM, LMGC
#
# Description
# -----------
# The test demonstrates:
# - How to define moving boundary conditions as functions of position.
# - How to generate an initial particle packing in a circular domain.
# - How to couple the particle and fluid solvers in MigFlow.
# - Optional use of `lmgc90` instead of `scontact` for contact resolution.
import os, sys, shutil, subprocess, time, random
import numpy as np
from migflow import fluid, scontact, time_integration, gmsh_io, volume_coupling as vc

use_lmgc90 = False
if use_lmgc90:
    from migflow import lmgc90Interface

# %%
# Output Directory and Mesh Generation
# ------------------------------------
# The output directory is created, and a 2D annular mesh is generated using Gmsh.

outputdir = "output" if len(sys.argv) < 2 else sys.argv[1]
initdir = f"{outputdir}_init"  # per-run: the ladder runs concurrently
geomesh_filename = "2d_mesh.geo"
mesh_filename = f"{outputdir}/mesh.msh"

shutil.rmtree(outputdir, ignore_errors=True)
shutil.rmtree(initdir, ignore_errors=True)
os.makedirs(outputdir)
os.makedirs(initdir)

subprocess.call(["gmsh", "-2", geomesh_filename, "-o", mesh_filename,
                 "-setnumber", "lcfac", os.environ.get("COUETTE_LCFAC", "20")])


# %%
# Initial Particle Generation
# ---------------------------
# Function to create an initial packing of grains within the annular domain.
def genInitialPosition(initdir, r, rout, rin, rhop, use_lmgc90):
    """Generate a particle packing and write it to an initial output file.

    Parameters
    ----------
    initdir : str
        Directory where the initial state is written.
    r : float
        Maximum particle radius.
    rout : float
        Outer radius of the drum.
    rin : float
        Inner radius of the drum.
    rhop : float
        Particle density.
    use_lmgc90 : bool
        If True, uses LMGC90 instead of scontact for contacts.
    """
    # Create particle problem
    p = scontact.ParticleProblem(2)
    print(mesh_filename)
    gmsh_io.load_msh_boundaries(p, mesh_filename, ["Outer", "Inner"], material="Steel")

    # Define grid of possible particle centers
    x = np.arange(rout, -rout, 2.5 * -r)
    x, y = np.meshgrid(x, x)
    R2 = x**2 + y**2

    # Keep only points inside the annular region
    keep = (R2 < (rout - r) ** 2) & (R2 > (rin + r) ** 2)
    x = x[keep]
    y = y[keep]

    # Add particles with slight density variation
    for xi, yi in zip(x, y):
        if yi < rout:
            rhop1 = random.choice([rhop * 0.9, 1.1 * rhop, rhop])
            p.add_particle((xi, yi), r, r**2 * np.pi * rhop1, "Sand")

    p.write_mig(initdir, 0)


# %%
# Physical and Numerical Parameters
# ---------------------------------
g = np.array([0, 0])  # no gravity
rho = 1.253e3  # fluid density [kg/m³]
rhop = 1000  # particle density [kg/m³]
nu = 1e-3  # kinematic viscosity [m²/s]
mu = rho * nu  # dynamic viscosity [Pa·s]
rout = 0.0254  # outer radius [m]
rin = 0.0064  # inner radius [m]
r = 397e-6 / 2  # grain radius [m]

# %%
# Particle Problem Initialization
# -------------------------------
random.seed(0)  # the ladder compares runs: same bed, same masses
genInitialPosition(initdir, r, rout, rin, rhop, use_lmgc90)

if use_lmgc90:
    friction = 0.1
    lmgc90Interface.scontactTolmgc90(initdir, 2, 0, friction)
    p = lmgc90Interface.ParticleProblem(2)
else:
    p = scontact.ParticleProblem(2)
    p.read_mig(initdir, 0)
    p.set_friction_coefficient(0.1, "Sand", "Sand")
    p.set_friction_coefficient(0.1, "Sand", "Steel")
p.set_fixed_contact_geometry(0)


# %%
# Fluid Problem Initialization
# ----------------------------
# The fluid solver is initialized, with a rotating inner boundary and fixed outer wall.
f = fluid.FluidProblem2(g, nu * rho, rho)
gmsh_io.load_msh(f, mesh_filename)
f.set_mean_pressure(0)
f.set_wall_boundary("Outer", velocity=[0, 0])
U_DRIVE = 0.1  # tangential speed the inner wall would have at r = rout [m/s]
f.set_strong_boundary(
    "Inner",
    velocity=[lambda x: (x[:, 1] / rout) * U_DRIVE,
              lambda x: (-x[:, 0] / rout) * U_DRIVE],
)

# Coupling with particles: tier A is the body geometry on the mesh, tier B the
# legacy mixture closure on it. No datum here -- this only primes the coupling
# so the first frame carries a porosity field, and time_integration.iterate
# re-sets all three tiers, datum included, at every step.
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)

# How much of the coupled volume the TRACE BAND actually reaches. gamma_t is
# ramped by sres = 1 - alpha(h/r), which is 0 at h/r >= 2, so on a coarse mesh
# the band can be inert and a sweep over its amplitude would compare identical
# runs. Print it once, before anyone reads the ladder.
_b, _w, _s = geo["body"], geo["w"], clo["s"]
_h = np.zeros(len(_s))
np.maximum.at(_h, _b, geo["max_edge"][np.asarray(geo["eid"], int)])
_rad = np.asarray(p.r()).ravel()        # (n,1) as returned; the masks need 1-D
_m = (_rad > 0) & (_h > 0)
_hr = _h[_m] / _rad[_m]
_ratio = clo["gamma_t"] * mu / clo["h"] ** 2 / np.maximum(clo["gamma_d_entry"], 1e-300)
print(f"[trace] MIGFLOW_TRACE_FRAC={os.environ.get('MIGFLOW_TRACE_FRAC', 'unset')}"
      f"  covered_volume_frac={(_s[_b] * _w).sum() / _w.sum():.4f}"
      f"  grains_with_sres>0={(_s[_m] > 0).mean():.4f}"
      f"  h/r={_hr.min():.2f}/{np.median(_hr):.2f}/{_hr.max():.2f}"
      f"  betat/gamma_v_max={_ratio.max():.3e}", flush=True)

# %%
# Simulation Parameters
# ---------------------
dt = 2.5e-3  # time step [s] -- with the DEM contact forces in the prediction
# (library default) dt*Fc/m must stay below the drive scale: at 5e-3 the
# sheared bed overdrives the grains, the sub-iteration count explodes and a
# step costs >18 min
tEnd = float(os.environ.get("COUETTE_TEND", 10))  # total sim time [s]
outf = 10  # number of iterations between outputs
t = 0  # time [s]
i = 0  # iteration counter
# %%
# Computational Loop
# ------------------
mass = np.pi * p.r() ** 2 * rhop  # particle masses
# Residual gate. iterate_patankar takes ONE Newton step per fluid solve; the
# Patankar prediction is affine, so if the momentum convection is negligible the
# reassembled residual after that step should sit at machine epsilon. A residual
# that GROWS with the step index is the linear system being amplified, which is
# what we are looking for here. -1 disables (the shipped default).
_RESID = float(os.environ.get("COUETTE_RESID", -1))
print(f"[resid] COUETTE_RESID={_RESID:g}", flush=True)

while t < tEnd:
    # The drive is evaluated at the END of the step, which is the velocity the
    # implicit solve is imposing.
    vmax = np.linalg.norm(p.velocity()[_rad > 0], axis=1).max()
    print(f"{i:4d}, {t:.6g}/{tEnd:.6g}, vmax {vmax:.6g}", flush=True)
    if i % outf == 0:
        p.write_mig(outputdir, t)
        f.write_mig(outputdir, t)
    time_integration.iterate_patankar(f, p, dt, 1, 1e-3 * r,
                                    external_particles_forces=mass * g,
                                    check_residual_norm=_RESID)
    t += dt
    i += 1
# %%
# Plot
# ----
# .. code-block:: shell
#
#  python3 -m migflow.plot.migplot output --actors fluid particles --fluid-field velocity
# %%
# Artifacts
# ---------
# - output/animation.mp4