"""Tetrahedral mesh generation for the 3D PD engine."""
import numpy as np

# 6-tet decomposition of a unit cube (indices into the cube's 8 corners,
# corner order: (i,j,k) bits -> idx = i + 2*j + 4*k)
_CUBE_TETS = np.array([
    [0, 1, 3, 7],
    [0, 3, 2, 7],
    [0, 2, 6, 7],
    [0, 6, 4, 7],
    [0, 4, 5, 7],
    [0, 5, 1, 7],
], dtype=np.int64)


def beam_tet_mesh(nx=8, ny=2, nz=2, dx=0.1):
    """Regular beam of (nx,ny,nz) cells, 6 tets per cell.

    Returns (verts (V,3) float64, tets (T,4) int64). All tets have
    positive orientation (det Dm > 0).
    """
    xs = np.arange(nx + 1) * dx
    ys = np.arange(ny + 1) * dx
    zs = np.arange(nz + 1) * dx
    gx, gy, gz = np.meshgrid(xs, ys, zs, indexing="ij")
    verts = np.stack([gx, gy, gz], axis=-1).reshape(-1, 3)

    def vid(i, j, k):
        return (i * (ny + 1) + j) * (nz + 1) + k

    tets = []
    for i in range(nx):
        for j in range(ny):
            for k in range(nz):
                corners = np.array([
                    vid(i + di, j + dj, k + dk)
                    for dk in (0, 1) for dj in (0, 1) for di in (0, 1)
                ])
                for t in _CUBE_TETS:
                    tets.append(corners[t])
    tets = np.array(tets, dtype=np.int64)

    # enforce positive orientation
    d = verts[tets]
    dm = np.stack([d[:, 1] - d[:, 0], d[:, 2] - d[:, 0], d[:, 3] - d[:, 0]], axis=-1)
    neg = np.linalg.det(dm) < 0
    tets[neg] = tets[neg][:, [0, 2, 1, 3]]
    return verts, tets