# %%
# 3D Mill
# =======
# This example demonstrates the workings of a 3-dimensional rolling drum.
# Spherical particles are added to a rotating drum, demonstratin.
# %%
# Keywords
# --------
# DEM, 3D, Rotating Drum, Granular Mixing
import sys, os, shutil, subprocess
import numpy as np
from migflow import scontact, time_integration, gmsh_io

# %%
# Output Directory
# ----------------
# The output directory is prepared for storing the simulation results.
outputdir = "output_3d_mill" if len(sys.argv) < 2 else sys.argv[1]
shutil.rmtree(outputdir, ignore_errors=True)
os.makedirs(outputdir)
meshfile = "3d_mesh_mill"
subprocess.call(
    ["gmsh", "-3", f"{meshfile}.geo", "-o", meshfile + ".msh", "-clscale", "1"]
)

# %%
# Geometrical Parameters
# ----------------------
# We define the area in which the particles will be created inside the mill.
rout = 0.5  # Outer radius of the mill
h = 0.6  # Height
w = 1  # Width
d = 0.05  # Depth

# %%
# Material Properties and Parameters
# ----------------------------------
rho = 2500  # Density
r = 9e-3  # Particle radius
g = np.array([0.0, -9.81, 0.0])  # Gravity vector
friction_coefficient_particle_particle = 0.3  # Friction coefficient between particles
friction_coefficient_particle_wall = (
    0.6  # Friction coefficient between particles and wall
)

# %%
# Particle Problem
# ----------------
# We begin by defining the 3D particle problem.
# A friction coefficient is set for both particle-particle and particle-wall interactions based on material types.
p = scontact.ParticleProblem(3)
p.set_friction_coefficient(friction_coefficient_particle_particle, "Glass", "Glass")
p.set_friction_coefficient(friction_coefficient_particle_wall, "Glass", "PVC")


# %%
# Particles generation
# --------------------
# We generate the particles on a rectangular grid inside the mill
def gen_rect(origin, w, h, d, step):
    """Generate coodinates on a rectangular grid

    Args:
        origin (list[float, float, float]): origin of top left corner
        w (float): width
        h (float): height
        d (float): depth
        step (float): maximal radius

    Returns:
        Positions where to place the
    """

    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]
    z = np.arange(step, d - step, 2 * step) + origin[2]
    x, y, z = np.meshgrid(x, y, z)
    return x.reshape(-1), y.reshape(-1), z.reshape(-1)


origin = [0, -rout, 0]
x, y, z = gen_rect(origin, w, h, d, 1.3 * r)
for xi, yi, zi in zip(x, y, z):
    if xi**2 + yi**2 < (rout - r) ** 2:
        p.add_particle((xi, yi, zi), r, 4 * r**3 * np.pi * rho / 3, "Glass")

# %%
# Boundary Conditions
# -------------------
# Load the cylindrical mill mesh as an additional body
bnd_body = p.add_body((0, 0, 0), 0, 0)
gmsh_io.load_mesh_to_body(p, bnd_body, f"{meshfile}.msh", ["Outer", "Front", "Rear"], material="PVC")

# %%
# Mill Rotation Parameters
rpm = 30  # revolutions per minute
omega = rpm * 2 * np.pi / 60  # rad/s

# %%
# Time Integration
# ----------------
# The simulation runs for a short time interval with a fixed time step.
# At each iteration, the gravitational forces are applied to the particles,
# and the system evolves according to the contact dynamics solver.
p.body_omega()[bnd_body, 2] = omega

t = 0.0
i = 0
dt = 2e-4
tEnd = 1
outf = 50
mass = 4 * np.pi * p.r() ** 3 * rho / 3
forces = mass * g
while t < tEnd:
    print(f"{i = :5d}, {t:.6g}/{tEnd:.6g}")
    if (i % outf) == 0:
        p.write_mig(outputdir, t)
    time_integration._advance_particles(p, forces, dt, 1, 1e-4 * r)
    t += dt
    i += 1

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