# %%
# Rotating Mill Particle Simulation
# =================================
# This example simulates a rotating mill with granular particles inside.
# Particles interact with the mill boundaries and among themselves with
# friction. The simulation demonstrates particle motion under gravity and rotation.
# %%
# Keywords
# --------
# DEM, Friction
from migflow import scontact, time_integration, gmsh_io
import numpy as np
import os, sys, shutil, subprocess

# %%
# Output Directory
# ----------------
# The output directory is prepared for storing the simulation results.
outputdir = "output_2d_mill" if len(sys.argv) < 2 else sys.argv[1]
shutil.rmtree(outputdir, ignore_errors=True)
os.makedirs(outputdir)


# %%
# Mesh Generation
# ---------------
# - "mesh.geo" is an external geometry file defining the problem domain
#   (here, the mill walls). It contains all the geometric entities and constraints.
mshfile = "2d_mesh_mill"
subprocess.call(
    ["gmsh", "-2", f"{mshfile}.geo", "-o", mshfile + ".msh", "-clscale", "1"]
)

# %%
# Particle Problem
# ----------------
# In this section, we define the particle problem and configure the
# boundaries of the rotating mill for the simulation.
#
# 1. `p = scontact.ParticleProblem(2)` creates a 2D discrete element
#    method (DEM) problem. This object will manage all particles,
#    their positions, velocities, contacts, and interactions.
#
# 2. `gmsh_io.load_msh_boundaries(p, None, ["Outer"], material="PVC")`
#    loads the boundary definitions from the current mesh stored in gmsh.
#
#      - `"Outer"` specifies which part of the mesh to use as the mill wall.
#      - `material="PVC"` assigns the material properties to these boundaries, which affects friction and particle-wall interactions.
#
# In short, this block links the mesh boundaries to the DEM simulation,
# allowing the code to know where the rotating walls are and how the
# particles should interact with them.
p = scontact.ParticleProblem(2)
# The drum is ONE massless body: load_msh_to_body puts every wall segment of
# the "Outer" boundary into it, so the mill turns by setting a single angular
# velocity instead of moving each arc by hand (load_msh_boundaries would give
# every arc its own body, and a rotation would have to be re-imposed on each of
# them at every step).
bnd_body = p.add_body((0, 0), 0, 0)  # inverse mass and inertia zero: kinematic
gmsh_io.load_msh_to_body(p, mshfile + ".msh", bnd_body, ["Outer"], material="PVC")

# Moving boundary: the walls travel between detections, so the contact basis
# must be predicted rather than frozen at detection time.
p.set_predict_basis(1)

# %%
# Material Properties and Parameters
# ----------------------------------
# In this section, we define the physical properties of the particles
# and the fluid
#
# Simulation parameters:
#   - `rhop` : density of the particles (used to compute their mass)
#   - `rho` : fluid density
#   - `r` : particle radius
#   - `rout` : radius of the rotating mill
#   - `g` : gravity vector
#   - `Fr` : Froude number, used to calculate the angular velocity
#   - `v` : angular velocity of the mill
#
# Friction coefficients are set to define interactions:
#   - `Glass-Glass` friction governs particle-particle contacts
#   - `Glass-PVC` friction governs particle-boundary contacts
rhop = 2500
rho = 1000
r = 1e-2
rout = 0.5
g = np.array([0.0, -9.81])
Fr = 2.5
v = (Fr * -g[1] * rout) ** 0.5
friction = 0.8
p.set_friction_coefficient(friction, "Glass", "Glass")
p.set_friction_coefficient(friction, "Glass", "PVC")

# %%
# Particles Initial Conditions
# ----------------------------
# In this section, we generate the initial positions of all particles
# inside the rotating mill and assign a mass to each particle.
#
# Steps performed in this block:
#
# 1. A rectangular grid of points is defined over a region of width w and height h.
#    The function gen_rect computes coordinates with spacing equal to the particle radius.
#    This ensures that particles are uniformly distributed and do not overlap initially.
# 2. Each grid point is tested to see if it lies within the mill radius. Only particles
#    inside the mill are added to the simulation.
# 3. When a particle is added using p.add_particle, it is assigned a mass computed as
#    pi * r^2 * rhop, where r is the particle radius and rhop is the particle density.
#    This mass will be used by the DEM solver to compute gravitational and contact forces.

h = 0.6
w = 1.0


def gen_rect(origin, w, h, step):
    eps = 1e-8
    x = np.arange(-w / 2 + step - eps, w / 2 - step + eps, 2 * step) + origin[0]
    y = np.arange(step, h - step, 2 * step) + origin[1]
    x, y = np.meshgrid(x, y)
    return x.reshape(-1), y.reshape(-1)


x, y = gen_rect([0, -rout], w, h, r)
for xi, yi in zip(x, y):
    if xi**2 + yi**2 < (rout - 2 * r) ** 2:
        p.add_particle((xi, yi), r, np.pi * r**2 * rhop, "Glass")

mass = np.pi * p.r() ** 2 * rhop


# %%
# Mill Boundary Circular Rotation
# -------------------------------
# In this section, we initialize the velocities of the particles that belong
# to the rotating mill walls. This setup ensures that the boundary segments
# move consistently with the mill’s rotation from the start of the simulation.
#
# Steps performed in this block:
#
# 1. Retrieve the internal tag ID of the "Outer" boundary using get_tag_id.
# 2. Identify all bodies (segments) in the mesh that belong to this outer boundary.
#    These bodies represent the mill walls.
# 3. Compute the velocity of each boundary particle based on its position:
#
#      - vit_x = v * y / rout
#      - vit_y = -v * x / rout
#
#    This corresponds to a circular rotation around the mill center,
#    where 'v' is the angular velocity and 'rout' is the mill radius.
#    The formulas ensure that each particle moves tangentially to the circle.
# 4. Assign the computed velocities to the boundary particles in the DEM simulation.

# The mill turns clockwise at v / rout: one angular velocity on the one wall
# body, and the DEM carries the whole drum -- geometry included.
p.body_omega()[bnd_body] = -v / rout

# %%
# Time Integration of the Simulation
# ----------------------------------
# In this section, the simulation is advanced over time using a fixed time step.
#
# Steps performed in this block:
#
# 1. The original positions of the mill boundary segments are stored in old_pos_seg
#    to ensure the boundaries can be reset at each time step.
# 2. A while loop iterates until the final simulation time (tEnd) is reached.
# 3. At each iteration:
#
#      - The drum body turns under its own angular velocity, so the walls and
#        the grains they drive stay consistent without any per-step fixup.
#      - The time integration solver (DEM solver) advances the particle positions
#        and velocities, taking into account gravitational forces and contact dynamics.
#      - The simulation time t and iteration counter ii are updated.
#      - Every outf iterations, the simulation state is written to the output directory
#        using p.write_mig, enabling later visualization or analysis.

dt = 1e-3
t = 0.0
tEnd = 2.0
ii = 0
outf = 10

while t < tEnd:
    time_integration.iterate(
        None,
        p,
        dt,
        min_nsub=8,
        external_particles_forces=g * mass,
        contact_tol=1e-4 * r,
    )
    t += dt
    ii += 1
    if ii % outf == 0:
        print(f"Output written at t = {t:.3f}")
        p.write_mig(outputdir, t)

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