# %%
# 2D Sand Pile Formation
# ======================
# This example simulates the formation of a sand pile through deposition of
# a large number of polydisperse circular particles between two rigid walls.
# The particles settle under gravity and form a stable pile through frictional
# contact.
# %%
# Keywords
# --------
# DEM, Friction, Granular Pile, Generation
from migflow import scontact, time_integration, deposit
import numpy as np
import shutil, os, sys

for n_particles in [20_000]:
    # Seed for the rng
    np.random.seed(0)

    # %%
    # output directory
    # ----------------
    outputdir = (
        f"output_2_no_inertia_{n_particles}" if len(sys.argv) < 2 else sys.argv[1]
    )
    shutil.rmtree(outputdir, ignore_errors=True)
    os.makedirs(outputdir)

    # %%
    # Physical parameters of the particles
    # ------------------------------------
    rho = 2650  # [kg.m**-3]
    friction_coefficient_particle_particle = 0.4
    friction_coefficient_particle_wall = 1.4

    L = 1e-3  # [m] reference length

    dmin = 0.6 * L  # [m]

    dmax = 2 * dmin

    # Domain size in 2d of the cut of the cylinder
    domain_length = 5e1 * L  # [m]
    domain_height = 7e2 * L  # [m]

    # %%
    # Problem definition
    # ------------------
    problem = scontact.ParticleProblem(2)
    problem.set_fixed_contact_geometry(0)
    problem.set_friction_coefficient(
        friction_coefficient_particle_particle, "Sand", "Sand"
    )
    problem.set_friction_coefficient(friction_coefficient_particle_wall, "Sand", "Wall")

    # Geometry = an immovable body
    floor = problem.add_body(
        (0, 0), 0, 0
    )  # [0, 0] = origin of the reference coordinates of the body
    wall_left = problem.add_body((0, 0), 0, 0)
    wall_right = problem.add_body((0, 0), 0, 0)

    x_shift = domain_length

    # problem.add_segment_to_body([ -20 * domain_length + x_shift, 0 ], [ 20 * domain_length + x_shift, 0 ], floor, "Wall", "Wall") #position = relative to the origin of the body
    x_bounds = np.array([-domain_length * 20 + x_shift, domain_length * 20 + x_shift])
    r_bnd = dmin / 2
    x_p = np.arange(x_bounds[0], x_bounds[1], 1.2 * r_bnd * 2)
    for xp in x_p:
        problem.add_particle_to_body([xp, -r_bnd], r_bnd, floor, "Wall")

    problem.add_segment_to_body(
        [-domain_length / 2 + x_shift, 0],
        [-domain_length / 2 + x_shift, domain_height],
        wall_left,
        "Wall",
        "Wall",
    )
    problem.add_segment_to_body(
        [domain_length / 2 + x_shift, 0],
        [domain_length / 2 + x_shift, domain_height],
        wall_right,
        "Wall",
        "Wall",
    )

    # %%
    # Particle initialisation
    # -----------------------
    problem_particles = np.asarray([], int)
    x0 = np.array([-domain_length / 2 + x_shift, 0.8 * domain_height])
    x1 = np.array([-domain_length / 2 + x_shift, 0])
    x2 = np.array([domain_length / 2 + x_shift, 0])
    x3 = np.array([domain_length / 2 + x_shift, 0.8 * domain_height])
    bnd = np.array([x0, x1, x2, x3])
    ri = np.random.uniform(dmin / 2, dmax / 2, n_particles)
    x, r = deposit.fastdeposit_2d(bnd, 1.1 * ri)
    r /= 1.1
    for xi, ri in zip(x, r):
        mi = np.pi * ri**2 * rho  # [kg] mass of the particle
        inertia = (mi * ri**2) / 2 * 100_000
        new_particle = problem.add_body(xi, 1 / mi, 1 / inertia)
        problem.add_particle_to_body((0, 0), ri, new_particle, "Sand")
        problem_particles = np.append(problem_particles, new_particle)

    # %%
    # Simulation parameters
    # ---------------------
    dt = 1e-4
    tEnd = 0.15#5.0
    outf = int(1e-2 / dt)
    iit = 0
    t = 0
    mg = np.pi * rho * problem.r() ** 2 * [0.0, -9.81]  # [m.s^-2]
    velocity_walls = 0.05

    # problem.set_compute_contact_forces(1)
    # problem.body_invert_inertia()[ : ] = 0

    while t < tEnd:
        print(
            f"{iit=: 04d} -- t/tEnd: {t:.3f}/{tEnd: .3f} -- {t/tEnd*100:02.2f}% completed",
            end="\r",
        )
        if iit % outf == 0:
            problem.write_mig(outputdir, t)
        if t > 0.0:
            problem.body_velocity()[wall_right, 1] = velocity_walls
            problem.body_velocity()[wall_left, 1] = velocity_walls
        time_integration._advance_particles(problem, mg, dt, 5, 1e-4 * dmin / 2)
        t += dt
        iit += 1

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