# %%
# Darcy Flow Through a Bed of Circular Discs
# ==========================================
# This example simulates Darcy flow through a packed bed of circular
# disc-shaped particles in two dimensions. A pressure gradient is applied
# across the domain and the resulting flow field through the porous medium
# is computed.
# %%
# Keywords
# --------
# FEM, DEM, Darcy, Porous Media
import sys
import os
import gmsh
import shutil
import numpy as np
from migflow import fluid, scontact, gmsh_io, time_integration
from migflow import volume_coupling as vc

outputdir = "output_2d_discs" if len(sys.argv) < 2 else sys.argv[1]
shutil.rmtree(outputdir, ignore_errors=True)
os.makedirs(outputdir)

# %%
# Geometrical parameters
# ----------------------
r = 5e-3  # grains radius
L = 25 * 2 * r  # domain width
H = 25 * 2 * r  # domain height
mesh_size = r / 4  # 0.1                             # mesh size
origin = np.array([-L / 2, -L / 2])  # mesh origin

# %%
# Physical parameters
# -------------------
rho = 1000  # fluid density
rhop = 1500  # grains density
mu = 1  # dynamic viscosity
compacity = 0.5  # Compactness
porosity = 1 - compacity  # porosity
p1 = 0
p0 = 10.0
g = np.array([-(p1 - p0) / (rho * L), 0])  # gravity


gmsh.initialize()
gmsh.model.add("mesh")

x0 = origin
x1 = origin + np.array([L, 0])
x2 = origin + np.array([L, H])
x3 = origin + np.array([0, H])
x = np.array([x0, x1, x2, x3])
edges = np.array([[0, 1], [1, 2], [2, 3], [3, 0]])
for xi in x:
    gmsh.model.geo.add_point(xi[0], xi[1], 0, mesh_size)
for edge in edges:
    gmsh.model.geo.add_line(edge[0] + 1, edge[1] + 1)
gmsh.model.geo.add_curve_loop([1, 2, 3, 4], 1)
gmsh.model.geo.add_plane_surface([1], 1)
gmsh.model.geo.synchronize()


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


gmsh.model.add_physical_group(
    1, get_line([origin[0], origin[1]], [origin[0], origin[1] + L]), name="Left"
)
gmsh.model.add_physical_group(
    1,
    get_line([origin[0] + L, origin[1]], [origin[0] + L, origin[1] + L]),
    name="Right",
)
gmsh.model.add_physical_group(2, [1], name="domain")
transform_y = np.array([[1, 0, 0, 0], [0, 1, 0, -H], [0, 0, 1, 0], [0, 0, 0, 1]])
transform_x = np.array([[1, 0, 0, +L], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])

gmsh.model.mesh.set_periodic(1, [1], [3], transform_y.flatten().tolist())
gmsh.model.mesh.set_periodic(1, [2], [4], transform_x.flatten().tolist())
gmsh.model.mesh.set_size_callback(lambda dim, tag, x, y, z, lc: mesh_size)
gmsh.model.mesh.generate(2)
p = scontact.ParticleProblem(2)
gmsh_io.load_msh_boundaries(p, None)

# Definition of the points where the grains are located
# The particles are first placed on a regular grid
# Then they are centered in the domain
a_g = np.pi * r**2
a_b = L**2
N = max(compacity * a_b / a_g, 1)

e = L / N**0.5
x = np.linspace(-L / 2, L / 2, int(N**0.5))
y = np.linspace(-H / 2, H / 2, int(N**0.5))
x, y = np.meshgrid(x, y)
for xi, yi in zip(x.flat, y.flat):
    p.add_particle((xi, yi), r, r**2 * np.pi * rhop)

f = fluid.FluidProblem2(g, mu, rho)
gmsh_io.load_msh(f, None)
f.set_mean_pressure(0)


# %%
# Numerical parameters
# --------------------
import os as _o
dt = float(_o.environ.get('DARCY_DT', '1e-4'))  # time step
tEnd = float(_o.environ.get('DARCY_TEND', '5e-2'))  # final time -- must
# hold several CTRL_EVERY windows (~49 steps each) so the driver can settle
outf = 10  # number of iterations between output files
t = 0.0
i = 0

# %%
# Time Integration
# ----------------


def get_fields(fluid):
    """Return derived output fields for visualization."""
    x = fluid.coordinates_fields()[fluid.field_indices(fluid.dimension())][:, 0]
    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[0] * x, p1_element),
    }


_csr = {}


def set_closure(fl, pp, geo, v_full=None):
    """Tier B: the model on a geometry that is already set. Re-evaluated every
    step because the FLUID state moves (the slip feeds the Dallavalle blend);
    the geometry underneath it does not.

    gamma_B is a LENGTH and belongs to the BODY: coeff * r with coeff a pure
    number. 1000 * r = 5.0e-3 m here -- written against the radius so it
    survives a change of units.
    """
    vel = pp.velocity() if v_full is None else v_full
    # gamma_b LEFT AT THE LIBRARY DEFAULT (GAMMA_B_COEFF * r, now 2 r -- the
    # drag-equivalent length for this bed; see GAMMA_B_COEFF in volume_coupling).
    #
    # It used to be pinned at 1000 r, which is ~550x the drag-equivalent length
    # this bed actually implies: the correlation branch defines
    # gamma_BD = gamma_D r^2/(mu S_b) = cd_re r/(4 pi), and at cd_re ~ 23 that
    # is 1.8 r. The consequence is not subtle -- betap = mu gamma_B/ell^2 puts
    # term B's cell Peclet at gamma_B/h ~ 5000, the velocity field diverges
    # (max|u_x| ~ 9e5 m/s) and the recovered pressure gradient drifts off the
    # imposed -40. Measured 2026-08-24 over gamma_b = 1000 r ... 0.01 r: dp goes
    # -46.05, -40.117, -40.000024, -40.000017 as the coefficient comes down.
    # DARCY_GAMMA_T=0 forces the trace operator OFF (unset = the library
    # band, which is the shipped default). Nothing else changes.
    _gt = _o.environ.get("DARCY_GAMMA_T")
    clo = vc.get_particles_closure(fl, geo, mu=mu, rho_f=rho, velocity=vel,
                                   **({} if _gt is None else {"gamma_t": float(_gt)}))
    vc.set_coupling_closure(fl, clo)
    # NO tier C: this bed is held, so there is no body momentum to predict --
    # and "no datum" is the absence of a set_coupling_datum call, not a flag.
    _csr["csr"] = dict(geo, **clo)


# TIER A, ONCE FOR THE WHOLE RUN. This bed is HELD -- the grains never move --
# so the overlap of every disc with every element it touches is invariant, and
# rebuilding it every step (which is what a single set_bodies hook has to do)
# repeats the one expensive part of the coupling for nothing. iterate_iqn's
# ``geometry=`` takes the tier-A dict built here and calls only ``set_closure``
# inside the loop.
bed = vc.get_particles(f, p.position(), p.r(), density=vc.body_density(p))
vc.set_coupling_geometry(f, bed)

# %%
# Imposed flow rate
# -----------------
# The cell is periodic in BOTH directions and driven by a body force, so there
# is no inlet on which to impose a velocity. Instead the driver is STEERED: the
# body force is corrected each step until the pore-averaged velocity sits on
# V_target, which turns the measurement around. Dp_drag becomes exact (it is the
# correlation evaluated at a number we chose, carrying no measurement noise) and
# the body force the solver ends up needing IS the answer.
#
# The old form imposed dp = -40 and compared it against a correlation fed by a
# MEASURED velocity, so every bit of noise in that velocity landed in the error.
# Worse, the measurement was polluted (see the U comment below) badly enough to
# flip the sign of Dp_drag.
#
# V_target is chosen so the correlation reproduces exactly the historical
# operating point dp = -40, which keeps the case in the regime it was built for
# and makes the steady-state driver 40/rho = 0.04, the shipped g[0]. So a model
# that agrees with the correlation leaves g untouched, and any drift of g away
# from 0.04 IS the disagreement, read directly.
_voidage = porosity ** (-1.8)
_N_bodies = compacity / (np.pi * r**2)


def _gamma_of_V(V):
    """Dallavalle-Di Felice gamma at interstitial velocity V. Re is built on the
    SUPERFICIAL velocity U = porosity*V, matching the postpro below."""
    Re = porosity * V * rho * 2 * r / mu
    cd_re = _voidage / porosity * (0.63 * np.sqrt(porosity * abs(Re)) + 4.8) ** 2
    return 0.5 * mu * cd_re


def _dp_of_V(V):
    return -_gamma_of_V(V) * _N_bodies * V


# solve _dp_of_V(V) = DP_SET for V by bisection -- monotone in V, and the
# sqrt(|Re|) makes it mildly nonlinear, so no closed form.
DP_SET = -rho * g[0]          # the shipped operating point, -40
_lo, _hi = 0.0, 1.0
while _dp_of_V(_hi) > DP_SET:      # push the bracket out until it straddles
    _hi *= 2.0
for _ in range(200):
    _mid = 0.5 * (_lo + _hi)
    if _dp_of_V(_mid) > DP_SET:
        _lo = _mid
    else:
        _hi = _mid
V_TARGET = float(_o.environ.get("DARCY_VTARGET", 0.5 * (_lo + _hi)))
# (gamma*N/rho) converts a velocity error into the acceleration that would close
# it at steady state, so K = 1 is one Newton step on the force balance and is
# EXACT while the drag is linear in V (it is, at this Re).
K_CTRL = float(_o.environ.get("DARCY_KCTRL", "1.0"))
_ctrl_scale = _gamma_of_V(V_TARGET) * _N_bodies / rho
# THE BED HAS A RELAXATION TIME and the correction must wait for it. Momentum
# per unit volume is rho*porosity*dV/dt = rho*g - gamma*N*V, so
#     tau = rho*porosity/(gamma*N) = porosity/_ctrl_scale
# which here is 9.8e-4 s, i.e. ~10 dt. Correcting EVERY step instead drove the
# integral action ~10x faster than the plant could answer: measured 2026-08-25,
# V overshot to 272x target by step 10 and then oscillated between +/-2e-2 with
# g swinging over +/-40, never settling. Waiting 5*tau between corrections makes
# each one a Newton step on a converged state, and the linear drag makes it land
# in one or two.
_tau = porosity / _ctrl_scale
CTRL_EVERY = int(_o.environ.get("DARCY_CTRL_EVERY", max(1, round(5 * _tau / dt))))
CTRL_TOL = float(_o.environ.get("DARCY_CTRL_TOL", "1e-4"))
print(f"[ctrl] V_target={V_TARGET:.6e}  dp_target={_dp_of_V(V_TARGET):.6f}  "
      f"g0_ff={_ctrl_scale * V_TARGET:.6e}  K={K_CTRL}  "
      f"tau={_tau:.4e}  every={CTRL_EVERY} steps", flush=True)

# node_volume and the porosity weight are fixed (held bed, fixed mesh), so the
# pore average below is assembled once.
_nv = f.node_volume().get()
_nv_sum = np.sum(_nv)


def _pore_velocity():
    """Interstitial (pore-averaged) x-velocity. porosity -> 0 on covered nodes,
    which is what keeps the grain interiors out of the average."""
    _phi = np.asarray(f.porosity().get()).ravel()
    return np.sum(f.velocity()[:, 0] * _phi * _nv) / _nv_sum / porosity


while t < tEnd:
    print(f"t/tEnd={t:1.4f}/{tEnd:1.4f}, i={i:1d}")
    if i % outf == 0:
        p.write_mig(outputdir, t)
        f.write_mig(outputdir, t, get_fields(f))
    # Held bed on THE UNFITTED MODEL (full setting: overlap overlap CSR,
    # m terms, Babuska/Dallavalle blend, trace stabiliser, extra
    # diffusivities) through iterate_iqn's fixed-grain branch: one solve per
    # step, the same pipeline as a moving case, no bespoke loop.
    time_integration.iterate_iqn(f, p, dt, fixed_grains=True,
                                 check_residual_norm=-1,
                                 geometry=bed, set_closure=set_closure)
    # STEER the driver onto V_TARGET, but only once the bed has answered the
    # previous correction -- see CTRL_EVERY above.
    if i > 0 and i % CTRL_EVERY == 0:
        _V = _pore_velocity()
        _rel = _V / V_TARGET - 1.0
        g[0] += K_CTRL * (V_TARGET - _V) * _ctrl_scale
        f.set_g(g)
        print(f"[ctrl] i={i:4d}  V={_V:.6e}  V/V_target={_V / V_TARGET:.6f}  "
              f"g0={g[0]:.6e}  dp={-rho * g[0]:.6f}", flush=True)
        if abs(_rel) < CTRL_TOL:
            print(f"[ctrl] CONVERGED |V/V_target - 1| = {abs(_rel):.3e} "
                  f"< {CTRL_TOL:g} at i={i}", flush=True)
            break
    i += 1
    t += dt

node_volume = f.node_volume().get()

# U is the SUPERFICIAL (Darcy) velocity, so the per-node porosity has to be in
# the average: without it the mean runs over every node including the ones
# inside the grains, and the covered-node velocity -- whatever the closure
# happens to leave there -- pollutes it. Measured 2026-08-24: dropping the
# trace moved that interior field enough to FLIP THE SIGN of Dp_drag
# (-789 -> +50.6) while the recovered gradient stayed at -40 to 8 digits,
# i.e. the whole error metric was reporting the interior, not the flow.
# phi -> 0 on covered nodes kills that term. It is the nodal porosity, not
# the exact CSR coverage, so U stays a little noisy -- fine for this postpro.
phi = np.asarray(f.porosity().get()).ravel()
U = np.sum(f.velocity()[:, 0] * phi * node_volume) / np.sum(node_volume)
V = U / porosity
Re = V * rho * 2 * r / mu
print("Re : %2.g ---- t : %2.g ---- tEnd : %2.g" % (Re, dt, tEnd))

voidage = porosity ** (-1.8)
cross_section = 2 * r
Re = U * rho * cross_section / mu
# Cd ALONE IS SINGULAR AT REST. Cd = voidage (0.63 + 4.8 (eps Re)^-1/2)^2 blows
# up as Re -> 0 and is NaN for Re < 0, which is why this line printed nan
# whenever the mean velocity came out negative. Cd*Re is the well-defined group
# -- it is what dallavalle_gamma uses inside the library -- and the velocity
# cancels out of the drag slope exactly:
#
#     gamma = Cd (rho d U / 2) = (Cd Re / Re)(rho d U / 2) = mu (Cd Re) / 2
#
# since Re = rho d U / mu. So gamma is U-INDEPENDENT and finite at rest, and
# Dp_drag keeps its linear-in-U Darcy form through V.
cd_re = voidage / porosity * (0.63 * np.sqrt(porosity * abs(Re)) + 4.8) ** 2
gamma = 0.5 * mu * cd_re
vol = np.pi * r**2
N = compacity / vol
# Dp_drag is the correlation at the IMPOSED V, so it carries no measurement
# noise at all. The measured quantity is dp, the driver the solver needed.
Dp_drag = _dp_of_V(V_TARGET)

# The pressure gradient along the flow: fields_gradient is (node, field, dim).
dp = f.fields_gradient()[:, -1, 0]
dp = np.sum(dp * node_volume) / np.sum(node_volume) - rho * g[0]

error = np.sum(abs(dp - Dp_drag) / Dp_drag)
print(dp, Dp_drag, error)

print("Pressure gradient from simulation : ", dp)
print("Pressure gradient from drag model : ", Dp_drag)
print("Relative error : ", error)

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