# %%
# Bidimensional shear of a dry granular material
# ==============================================
#
# This example simulates a 2D simple shear test of a dry granular material
# confined between two walls with periodic lateral boundaries.
# The top and bottom walls move at constant velocities to impose a shear rate.
# Keywords
# --------
# DEM, Periodic Boundaries, Extra Fields
import os, sys, shutil
import numpy as np
import gmsh
from migflow import scontact, time_integration

np.random.seed(0)

# %%
# Output Directory
# ----------------
# The output directory is (re)created
outputdir = "output" if len(sys.argv) < 2 else sys.argv[1]
shutil.rmtree(outputdir, ignore_errors=True)
os.makedirs(outputdir)

# %%
# Geometrical parameters and mesh generation
# ------------------------------------------
h = 0.1  # domain height
w = 0.2  # domain width
r = 1e-3  # particle radius

# %%
# Physical Parameters
# -------------------
g = np.array([0.0, 0.0])  # gravity (no gravity)
rhop = 1000  # particle density
shear_rate = 10  # imposed shear rate
Pp = 100  # confining pressure
friction_wall_sand = 1.0  # wall-sand friction coefficient
friction_sand_sand = 0.4  # sand-sand friction coefficient

# %%
# Particle problem setup
# ----------------------
# Initialize the particle problem and set up friction coefficients.
p = scontact.ParticleProblem(2)
p.set_fixed_contact_geometry(0)
p.set_friction_coefficient(friction_wall_sand, "wall", "sand")
p.set_friction_coefficient(friction_sand_sand, "sand", "sand")

# %%
# Periodic Boundaries
# -------------------
# Create the periodic entity that store the transformation, then add the periodic segments linked to it.
trans = np.array([w, 0])
ent0 = p.add_periodic_entity(1, trans)
ent1 = p.add_periodic_entity(1, -trans)
p.add_periodic_segment(ent0, np.array([(-w / 2, h), (-w / 2, -h)]))
p.add_periodic_segment(ent1, np.array([(w / 2, h), (w / 2, -h)]))

# %%
# Add top and bottom walls as rigid bodies
# ----------------------------------------
x0, x1 = np.array([-w / 2 - w / 20, 0]), np.array([w / 2 + w / 20, 0])
mass_seg = rhop * np.pi * r * w / 2
top_center = np.array([0, h / 2])
top_body = p.add_body(top_center, mass_seg, 0)
p.add_segment_to_body(x0 + top_center, x1 + top_center, top_body, "Top", "wall")
bottom_center = np.array([0, -h / 2])
bottom_body = p.add_body(bottom_center, mass_seg, 0)
p.add_segment_to_body(x0 + bottom_center, x1 + bottom_center, bottom_body, "Bottom", "wall")


# %%
# Particle initialization
# -----------------------
# Generate hexagonal packing
eps = 1e-4
step = 2 * r + eps
lx, ly = w, h - 4 * r
y0 = -h / 2 + r + eps

x = np.arange(-lx / 2, lx / 2 - 2 * r, step)
y = np.arange(y0, y0 + ly, np.sqrt(3) * step)
x2 = np.arange(-lx / 2 + r, lx / 2 - 2 * r, step)
y2 = y + 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
order = np.argsort(y)
x, y = x[order], y[order]

pbodies = np.array([], np.int32)
for xi, yi in zip(x, y):
    ri = np.random.uniform(0.5, 1.0) * r
    body = p.add_particle((xi, yi), ri, np.pi * ri**2 * rhop)
    pbodies = np.append(pbodies, body)

# %%
# Simulation Parameters
# ---------------------
outf = 5
dt = 1e-3
tEnd = 0.5
t = 0
i = 0

# %%
# Initial Conditions for the walls
# --------------------------------
v_top = shear_rate * p.body_position()[top_body, 1]
v_bottom = shear_rate * p.body_position()[bottom_body, 1]
p.body_velocity()[top_body, 0] = v_top
p.body_velocity()[bottom_body, 0] = v_bottom
p.body_velocity()[pbodies, 0] = p.body_position()[pbodies, 1] * shear_rate


# %%
# Computational Loop
# ------------------
def periodic_map(width, bodies):
    """Remap particles crossing periodic boundaries."""
    p.body_position()[bodies, 0] = (
        np.remainder(p.body_position()[bodies, 0] + 3 * width / 2, width) - width / 2
    )


def accumulate(f_contacts, n_divide):
    """Accumulate boundary forces on top and bottom walls."""
    f_contacts[0] += p.get_boundary_forces("Top") / n_divide
    f_contacts[1] += p.get_boundary_forces("Bottom") / n_divide


forces_particles = np.zeros_like(p.velocity())
forces_body = np.zeros_like(p.body_velocity())
forces_body[top_body, 1] = -Pp * w / 2
forces_body[bottom_body, 1] = Pp * w / 2
xold = p.body_position()[:, 0].copy()
displacement = np.zeros_like(p.position())
while t < tEnd:
    print(f"{i:4d}, {t:.6g}/{tEnd:.6g}")
    periodic_map(w, pbodies)
    if i % outf == 0:
        p.write_mig(outputdir, t, {"displacement": displacement})

    p.body_velocity()[top_body, 0] = v_top
    p.body_velocity()[bottom_body, 0] = v_bottom
    f_contacts = np.zeros((2, 2))

    time_integration._advance_particles(
        p,
        forces_particles,
        dt,
        min_nsub=1,
        max_nsub=5,
        contact_tol=1e-3 * r,
        fbody=forces_body,
        after_sub_iter=lambda n_divide: accumulate(f_contacts, n_divide),
    )

    p.body_position()[top_body, 0] = xold[top_body]
    p.body_position()[bottom_body, 0] = xold[bottom_body]
    displacement += p.velocity() * dt
    t += dt
    i += 1

# %%
# Plot
# ----
# .. code-block:: shell
#
#  python3 -m migflow.plot.migplot output --actors particles
# %%
# Artifacts
# ---------
# - output/animation.mp4
# - output/deposit.svg