# %%
# 3D Hourglass Simulation
# =======================
# This example simulates the discharge of granular particles through a
# three-dimensional hourglass geometry. An initial particle deposit is loaded
# from a pre-generated configuration. Particles fall under gravity through a
# narrow orifice and accumulate at the bottom.
# %%
# Keywords
# --------
# DEM, 3D, Granular Flow, Hourglass
import sys
import os
import shutil
import subprocess

import numpy as np
from migflow import scontact, time_integration

# %%
# Parameters
# ----------

np.random.seed(0)

h = 4.0
l = 1.0
e = 0.5
d = 0.01

outputdir = "output" if len(sys.argv) < 2 else sys.argv[1]
shutil.rmtree(outputdir, ignore_errors=True)
os.makedirs(outputdir, exist_ok=True)

t = 0.0
dt = 5e-4
i = 0
outf = 20
g = np.array([0, 9.81, 0])

# %%
# Load particle problem
# ---------------------
p = scontact.ParticleProblem(3)
# Ready-made initial deposit or generate it if not existing
inputdir = "hourglass-deposit"
if inputdir not in os.listdir():
    # run the initial deposit generation script
    print(f"Generating initial deposit in directory '{inputdir}'...")
    subprocess.call(["python3", "3d_gen_initial_deposit.py"])
else:
    print(f"Using existing directory '{inputdir}' for initial deposit.")
p.read_mig(inputdir, iteration=-1)

# Slightly shrink particles
p.r()[:] *= 0.9

# %%
# Time integration loop
# ---------------------

tEnd = 1.0
while t < tEnd:
    mass = 4 / 3 * np.pi * p.r() ** 3
    forces = mass * g

    time_integration._advance_particles(p, forces, dt, 1, 1e-3 * d)

    if i % outf == 0:
        print("t =", t)
        p.write_mig(outputdir, t)

    t += dt
    i += 1

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