# %%
# 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
import numpy as np
import os, sys, shutil, subprocess

# %%
# Output Directory
# ----------------
# The output directory is prepared for storing the simulation results.
outputdir = "output" 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. `p.load_msh_boundaries(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)
p.load_msh_boundaries(mshfile + ".msh", ["Outer"], material="PVC")

# %%
# Material Properties and Parameters
# ----------------------------------
# In this section, we define the physical properties of the particles
# and the fluid, as well as 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.

itag = p.get_tag_id("Outer")
bodies_segments = p.segments()[p.segments()["tag"] == itag]["body"]
vit_x = v * p.body_position()[bodies_segments, 1] / rout
vit_y = -v * p.body_position()[bodies_segments, 0] / rout
p.body_velocity()[bodies_segments, 0] = vit_x
p.body_velocity()[bodies_segments, 1] = vit_y

# %%
# 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 boundary positions are reset to old_pos_seg to maintain the circular motion.
#    - 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.

old_pos_seg = p.body_position()[bodies_segments, :]

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

while t < tEnd:
    p.body_position()[bodies_segments, :] = old_pos_seg
    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)