# %%
# Particle Deposition Between Two Inclined Walls
# ==============================================
# This example simulates granular particle deposition between two
# inclined rigid walls. The walls form a V-shape and particles are
# randomly deposited inside the domain using a fast deposition algorithm.
# %%
# Keywords
# --------
# DEM, Deposition, Granular Flow


import shutil
import os, sys
import numpy as np
from migflow import scontact, deposit

np.random.seed(0)

# %%
# Output Directory
# ----------------
# Create a clean output directory for simulation results.
outputdir = "output" if len(sys.argv) < 2 else sys.argv[1]
shutil.rmtree(outputdir, ignore_errors=True)
os.makedirs(outputdir)

# %%
# Geometry Definition
# -------------------
# In this section, we construct the two rigid inclined walls that form
# the boundaries of the deposition domain.
theta1 = np.radians(80)  # left wall angle
theta2 = np.radians(100)  # right wall angle
l = 10  # wall length
w = 1  # lateral offset

x0 = np.array([0, 0])
x1 = np.array([l * np.cos(theta1), l * np.sin(theta1)])
x2 = np.array([l * np.cos(theta2), l * np.sin(theta2)])

s0 = np.array([x0, x1])
s1 = np.array([x0, x2])

s0[:, 0] += w / 2
s1[:, 0] -= w / 2

# %%
# Particle Problem
# ----------------
# This section initializes the 2D discrete element method (DEM) problem that solves the contact dynamics of the particles.
p = scontact.ParticleProblem(2)

# %%
# Material Properties and Simulation Parameters
# ---------------------------------------------
# In this section, we define the physical properties of the particles
# and key simulation parameters.
rhop = 1000  # particle density
r = 1 / 15  # particle radius
g = np.array([0.0, -9.81])  # gravity vector
mu = 0.5  # friction coefficient
p.set_friction_coefficient(mu)

# %%
# DEM Initial Conditions
# ----------------------
# In this section, we generate the boundary and the initial positions of all particles
# %%
# Boundaries
# ----------
bnd = p.add_body(x0, 0, 0)
p.add_segment_to_body(s0[0], s0[1], bnd, "bnd")
p.add_segment_to_body(s1[0], s1[1], bnd, "bnd")
p.add_particle_to_body(s0[0], 0, bnd)
p.add_particle_to_body(s1[0], 0, bnd)
p.add_particle_to_body(s0[1], 0, bnd)
p.add_particle_to_body(s1[1], 0, bnd)

# %%
# Particles
# ---------
# 1. Define a maximum number of particles `n_p_max` and initialize an
#    array of particle radii, all equal to r.
#
# 2. Define the polygon `bnd_deposit`, by vertices, representing the boundaries where
#    particles will be deposited. Be careful that the order of points matters.
#    It must start with the the top left point and go anticlockwise.
#    The polygon must not be closed, nor self-intersecting.
#
# 3. Use `deposit.fastdeposit_2d` to generate particle positions (`xp`)
#    and adjusted radii (`radius`) inside the polygon. This function packs
#    particles efficiently without overlaps and outputs a visualization
#    file 'deposit.svg'.
n_p_max = 10000  # maximum number of particles to insert
radius = np.full(n_p_max, r)  # array of particle radii
bnd_deposit = np.array([s1[1], s1[0], s0[0], s0[1]])
xp, radius = deposit.fastdeposit_2d(bnd_deposit, radius, outputdir + "/deposit.svg")
for xi, ri in zip(xp, radius):
    # A body of its own, with zero inverse inertia so the disc cannot spin.
    body = p.add_body(xi, 1 / (np.pi * ri**2 * rhop), 0)
    p.add_particle_to_body((0, 0), ri, body)

# %%
# Numerical Parameters
# --------------------
t = 0.0  # initial time
i = 0  # iteration counter
dt = 0.001  # time step
tEnd = 2  # end time
outf = 10  # output frequency

# %%
# Simulation Loop
# ---------------
forces = p.r() ** 2 * np.pi * rhop * g
while t < tEnd:
    print(f"{i:4d}, {t:.6g}/{tEnd:.6g}, {dt:.6g}")
    if i % outf == 0:
        p.write_mig(outputdir, t)

    # Advance the simulation by one time step
    p.iterate(dt, forces, tol=1e-4 * r)
    # Remove particles that have fallen below y = 0
    xp = p.body_position()
    outside = xp[:, 1] < 0
    p.remove_bodies_flag(~outside)
    # Advance time and iteration counter
    t += dt
    i += 1

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