# %%
# Tetris
# ======
# This example illustrates how to define a particle problem in which
# composite rigid bodies—shaped like *tetriminoes*—fall under gravity
# into an open box. The pieces are generated randomly and interact
# through frictional contact, stacking up in a way reminiscent of
# the classic Tetris game.
#
# Keywords
# --------
# DEM, Rigid bodies, Random generation
#
#
# Description
# -----------
# Each piece (tetrimino) is composed of four circular particles grouped
# into a single rigid body. Pieces are inserted at random positions and
# orientations, and fall under gravity into an open container bounded
# by three fixed walls. The contact solver handles collisions between
# the falling bodies and the boundaries, leading to complex packing and
# rotation dynamics.
import sys, os, shutil, random
import numpy as np
from migflow import scontact

random.seed(0)

# %%
# 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)


# %%
# Piece Generation
# ----------------
# Function to generate and insert a tetrimino piece composed of four
# rigidly connected circular particles.
def add_piece(p: scontact.ParticleProblem, r: float, initial_coord: np.ndarray):
    """Create and insert a tetrimino piece into the MigFlow particle problem.

    Parameters
    ----------
    p : scontact.ParticleProblem
        The particle problem instance to which the piece is added.
    r : float
        Particle radius.
    initial_coord : np.ndarray
        Initial coordinates for the lower-left reference position of the piece.
    """
    # Each piece is defined by a binary mask over a 2x4 rectangular grid.
    models = [
        [[1, 0, 0, 0], [1, 1, 1, 0]],
        [[0, 0, 0, 1], [0, 1, 1, 1]],
        [[1, 1, 1, 1], [0, 0, 0, 0]],
        [[0, 1, 1, 0], [0, 1, 1, 0]],
        [[0, 0, 1, 1], [0, 1, 1, 0]],
        [[0, 1, 0, 0], [1, 1, 1, 0]],
        [[0, 1, 1, 0], [0, 0, 1, 1]],
    ]
    piece = random.choice(models)
    omega = 10 * np.pi * (-1 + 2 * random.random())
    y, x = np.where(piece)
    x = r * (initial_coord[0] + 2 * x)
    y = r * (initial_coord[1] + 2 * y)
    R = np.repeat(r, 4)  # particles radii
    coord = np.column_stack([x, y])  # particle coordinates
    body = p.add_particle_body(coord, R, np.pi * R**2 * rho)
    p.body_omega()[body, 0] = omega  # initial body rotation


# %%
# Physical Parameters
# -------------------
g = np.array([0, -9.81])  # gravity [m/s²]
rho = 1000  # particle density [kg/m³]
r = 0.05  # particle radius [m]
h = 2  # box height [m]
w = 2  # box width [m]
x0 = np.array([-3, 40])  # initial insertion position

# %%
# Particle Problem
# ----------------
# The 2D particle problem is created and the open box boundaries are defined.
# The container is defined using three rigid walls (left, bottom, right).
p = scontact.ParticleProblem(2)
p.set_fixed_contact_geometry(0)
xs0 = np.array([-w / 2, +h / 2])
xs1 = np.array([-w / 2, -h / 2])
xs2 = np.array([+w / 2, -h / 2])
xs3 = np.array([+w / 2, +h / 2])
edges = [[xs0, xs1], [xs1, xs2], [xs2, xs3]]
nodes = [xs0, xs1, xs2, xs3]
bnd_body = p.add_body((0, 0), 0, 0)
for edge in edges:
    p.add_segment_to_body(edge[0], edge[1], bnd_body, "bnd")
for n in nodes:
    p.add_particle_to_body(n, 0, bnd_body)

# %%
# Initial Piece
# -------------
# An initial tetrimino piece is inserted at the top of the domain.
add_piece(p, r, x0)


# %%
# Simulation Parameters
# ---------------------
outf = 20  # number of iterations between output files
dt = 1e-3  # time step [s]
tEnd = 20  # total simulation time [s]
t = 0  # initial time [s]
i = 0  # iteration counter

# %%
# Computational Loop
# ------------------
# Pieces fall into the container under gravity. Every 250 iterations,
# a new piece is generated and inserted.
while t < tEnd:
    if i % outf == 0:
        print(f"{i =:4d}, {t:.6g}/{tEnd:.6g}")
        p.write_mig(outputdir, t)  # write results
    forces = g * (np.pi * p.r() ** 2 * rho)  # external forces (gravity)
    p.iterate(dt, forces, tol=1e-4 * r)  # contact solver
    t += dt
    i += 1
    # Add a new piece periodically
    if i % 250 == 0:
        print("add_piece")
        add_piece(p, r, x0)

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