# %%
# Rotating Drum with Blades
# =========================
# This test case simulates a 2D rotating drum partially filled with granular
# material. The outer wall contains internal blades that stir the particles
# as the drum rotates slowly around its center. This configuration is used
# to study granular mixing, internal avalanches, and force chains formation
# under steady rotation.
#
# Keywords
# --------
# DEM, Friction, Mixing, Rotation, Rigid Body
#
# Description
# -----------
# The geometry is fully defined in the script: an outer square drum containing
# two internal blades and a fixed central axle. A population of circular grains
# is initialized at the bottom of the drum. As the drum rotates slowly,
# the grains are lifted and then cascade down, forming intermittent avalanches.
# The simulation captures the frictional interactions between the grains and
# the confining boundaries.
from migflow import scontact, time_integration
import os, sys, shutil
import numpy as np

np.random.seed(0)

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

# %%
# Physical Parameters
# -------------------
# Define the characteristic size, mass, and geometry dimensions.
r = 0.025  # particle radius [m]
mass = 1.0  # particle mass [kg]
w = 2.0  # drum width [m]
h = 2.0  # drum height [m]

# %%
# Particle Problem Definition
# ---------------------------
# The particle problem is initialized with frictional properties and
# material definitions for both the drum ("Steel") and the particles ("Sand").
p = scontact.ParticleProblem(2)
p.set_fixed_contact_geometry(0)
p.set_friction_coefficient(1.5, "Sand", "Sand")
p.set_friction_coefficient(1.5, "Sand", "Steel")

# %%
# Geometry Construction
# ---------------------
# The drum is made of rigid segments connected together as a single body.
# Two internal blades are added on opposite sides of the drum.
outer_body = p.add_body((0, 0), 0, 0)
x0 = np.array([+w / 2, +h / 2])
x1 = np.array([+w / 2, -h / 2])
x2 = np.array([-w / 2, -h / 2])
x3 = np.array([-w / 2, +h / 2])
x4 = np.array([-w / 2, 0])
x5 = np.array([-w / 6 - r, 0])
x6 = np.array([+w / 6 + r, 0])
x7 = np.array([+w / 2, 0])
edges = [[x0, x1], [x1, x2], [x2, x3], [x3, x0], [x4, x5], [x6, x7]]
for edge in edges:
    p.add_segment_to_body(edge[0], edge[1], outer_body, "drum", "Steel")
points = [x0, x1, x2, x3, x4, x5, x6, x7]
for point in points:
    p.add_particle_to_body(point, 0.0, outer_body, "Steel")

# %%
# Central Fixed Axis
# ------------------
# A rigid segment representing the central axle is added as a separate body.
x8 = np.array([-w / 6 + 3 * r, 0])
x9 = np.array([+w / 6 - 3 * r, 0])
axis_body = p.add_body((0, 0), 0, 0)
p.add_segment_to_body(x8, x9, axis_body, "axis", "Sand")
p.add_particle_to_body(x8, 0.0, axis_body, "Sand")
p.add_particle_to_body(x9, 0.0, axis_body, "Sand")

# %%
# Particle Initialization
# -----------------------
# A population of particles is generated in a rectangular region at the bottom
# of the drum. Small random perturbations are added to avoid perfect alignment.
p_bodies = np.asarray([], int)
for x in np.arange(-w / 2 + r, w / 2 - r, 3 * r):
    for y in np.arange(-w / 2 + r, -r, 3 * r):
        m = np.random.random()
        body = p.add_particle((x + m * r / 4, y), r - m * r / 4, mass, "Sand")
        p_bodies = np.append(p_bodies, body)

# %%
# Simulation Parameters
# ---------------------
omega = 1e-2  # angular velocity of the drum [rad/s]
dt = 1  # time step [s]
tEnd = 2500  # total simulation time [s]
outf = 10  # output frequency
i = 0  # iteration counter
t = 0  # initial time [s]

# %%
# Time Integration
# ----------------
# The outer drum rotates at a constant angular velocity.
# The granular material is integrated using the MigFlow contact dynamics solver.
f_ext = np.zeros_like(p.velocity())
f_ext[:, 1] -= 2e-4  # constant vertical body force (gravity-like)
p.body_omega()[outer_body] = omega
while t < tEnd:
    print(f"{i = :4d} ----- {t = :g}")
    if i % outf == 0:
        p.write_mig(outputdir, t)
    time_integration._advance_particles(p, f_ext, dt, 1, 1e-4 * r)
    i += 1
    t += dt

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