# %%
# Darcy Flow Through a Square Arrangement of Grains
# =================================================
# This example simulates Darcy flow through a square lattice arrangement
# of circular grains in two dimensions. A fluid velocity is imposed across
# the periodic 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

np.random.seed(42)

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

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

# %%
# Geometrical parameters
# ----------------------
r = 40e-3  # grains radius
L = 1  # domain width
h = r / 2  # mesh size
origin = np.array([-L / 2, -L / 2])  # mesh origin

gmsh.initialize()
gmsh.model.add("mesh")
gmsh.model.occ.add_rectangle(origin[0], origin[1], 0, 2 * L, 2 * L)
gmsh.model.occ.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] + 2 * L, origin[1]]), name="Bottom"
)
gmsh.model.add_physical_group(
    1,
    get_line([origin[0], origin[1] + 2 * L], [origin[0] + 2 * L, origin[1] + 2 * L]),
    name="Top",
)
gmsh.model.add_physical_group(
    1, get_line([origin[0], origin[1]], [origin[0], origin[1] + 2 * L]), name="Left"
)
gmsh.model.add_physical_group(
    1,
    get_line([origin[0] + 2 * L, origin[1]], [origin[0] + 2 * L, origin[1] + 2 * L]),
    name="Right",
)
gmsh.model.add_physical_group(2, [1], name="domain")
gmsh.model.mesh.set_size_callback(lambda dim, tag, x, y, z, lc: h)
gmsh.model.mesh.generate(2)

# %%
# Fluid Problem
# -------------
f = fluid.FluidProblem2(g, mu, rho)
gmsh_io.load_msh(f, None)
f.set_open_boundary("Left", velocity=[V, 0])
f.set_open_boundary("Right", pressure=0.0)
f.set_wall_boundary("Bottom", velocity=[0, 0])
f.set_wall_boundary("Top", velocity=[0, 0])
f.interpolate(velocity_x=V)

# %%
# Particle Problem
# ----------------
p = scontact.ParticleProblem(2)

x_g = np.arange(0 + 2 * r, 2 * L - 2 * r, 3 * r) + origin[0]
y_g = np.arange(0 + 2 * r, 2 * L - 2 * r, 3 * r) + origin[1]
x_g, y_g = np.meshgrid(x_g, y_g)

x_g = x_g.flat
y_g = y_g.flat
polygons = np.array([x_g, y_g]).T.reshape(-1, 1, 2)
grains_size = np.random.uniform(0.6 * r, 1 * r, polygons.shape[0])
polygons = np.repeat(polygons, 4, axis=1)
polygons[:, 0, :] += np.array([-grains_size, -grains_size]).T
polygons[:, 1, :] += np.array([+grains_size, -grains_size]).T
polygons[:, 2, :] += np.array([+grains_size, +grains_size]).T
polygons[:, 3, :] += np.array([-grains_size, +grains_size]).T

volume = (2 * grains_size) ** 2
density = np.full_like(volume, rho)

# Each grain is a SQUARE, and the DEM carries it as one: a body whose four
# corners are radius-0 particles and whose four edges are segments. A single
# disc at the centre is all the renderer would ever see, and the movie would
# draw circles for a bed of squares. The bodies are held
# -- iterate_iqn runs its fixed-grain branch -- so this changes what is
# WRITTEN, not what is solved: the fluid still sees `polygons` through
# set_unfitted_polygons.
# A hair of a radius, not zero: the renderer reconstructs a body's outline
# from its corner discs and its segments, and a radius-0 particle is filtered
# out before it gets there (r > 0 is what tells a grain from a wall marker).
# At 1/1000 of the grain the corner is subpixel, so the square stays sharp.
CORNER_R = 1e-3
for poly, ri in zip(polygons, grains_size):
    x_center = np.mean(poly, axis=0)
    body = p.add_body(x_center, 0.0, 0.0)          # inverse mass/inertia: held
    rel = poly - x_center
    for k in range(4):
        p.add_particle_to_body(rel[k], CORNER_R * ri, body, "Sand")
        p.add_segment_to_body(rel[k], rel[(k + 1) % 4], body, "grain", "Sand")


# %%
# Numerical parameters
# --------------------
t = 0
i = 0
dt = 1e-3
tEnd = 1.0
outf = 10

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


def get_fields(fluid: fluid.FluidProblem2):
    """Return derived output fields for visualization."""
    p1_element = fluid.get_p1_element()
    grad_v = fluid.fields_gradient()[:, :2, :]
    # porosity().get() is 1-D: reshape it too, or (n, 1) * (n,) broadcasts to
    # the (n, n) outer product -- 1.1 GB per frame on a 12k-node mesh, which
    # went to disk as a 105 GB vorticity stream and OOM-killed ParaView.
    vorticity = ((grad_v[:, 1, 0] - grad_v[:, 0, 1]).reshape(-1, 1)
                 * fluid.porosity().get().reshape(-1, 1))
    return {
        "pressure": (fluid.pressure(), p1_element),
        "velocity": (fluid.velocity(), p1_element),
        "porosity": (fluid.porosity().get(), p1_element),
        "vorticity": (vorticity, p1_element),
    }


gamma = 1e2 * np.ones(polygons.shape[0])


def set_polygons(fluid_, particles_, v_full=None):
    """set_bodies hook: the held squares on THE UNFITTED MODEL (overlap
    polygon overlap CSR, m = (x - c)/a with a the half-side -- exact for a
    square, m . n = 1 on every face, div m = 2/a --, the case's drag
    coefficient as the Babuska coarse-end term, trace stabiliser and extra
    diffusivities). m_kind="superellipse" is the smooth alternative if the
    corner set were a problem."""
    vc.set_unfitted_polygons(fluid_, polygons, mu=mu, rho_f=rho, density=density,
                             gamma_d=gamma)


while t < tEnd:
    if i % outf == 0:
        print(f"t/tEnd={t:1.1f}/{tEnd:1.1f}, i={i:1d}")
        f.write_mig(outputdir, t, get_fields(f))
        p.write_mig(outputdir, t)
    # Held polygons: iterate_iqn's fixed-grain branch with the polygon-body
    # hook (contact = -F closure, one solve per step).
    time_integration.iterate_iqn(f, p, dt, fixed_grains=True, set_bodies=set_polygons)
    t += dt
    i += 1

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