"""
3D Voxel Shape Generator — VAE-Matched Resolution (8×16×16)
=============================================================
Adapted from v10.2 for Flux 2 VAE latent geometry:
  - GZ=8 (channel dimension), GY=16, GX=16 (spatial dimensions)
  - 2048 voxels per patch (vs 15,625 at 25³)
  - Aspect ratio 1:2:2 matches VAE latent structure
  - All 38 shape classes preserved
  - Selective rasterization (polyhedra get edges, point-classes stay sparse)
  - Shapes generated in native aspect ratio space
  
Multi-scale usage:
  Patches extracted at any scale (4×8×8, 8×16×16, 16×32×32, 32×64×64)
  are resized to canonical 8×16×16 before classification.
"""

import numpy as np
from itertools import combinations

# === Grid dimensions (non-cubic, VAE-matched) ================================
GZ = 8    # channel dimension (thin)
GY = 16   # spatial height
GX = 16   # spatial width
GRID_SHAPE = (GZ, GY, GX)
GRID_VOLUME = GZ * GY * GX  # 2048

# Precompute coordinate grid for vectorized generation
_COORDS = np.mgrid[0:GZ, 0:GY, 0:GX].reshape(3, -1).T.astype(np.float64)

# === Shape classes ============================================================
CLASS_META = {
    # 0D
    "point":            {"dim": 0, "curved": False, "curvature": "none"},
    # 1D
    "line_x":           {"dim": 1, "curved": False, "curvature": "none"},
    "line_y":           {"dim": 1, "curved": False, "curvature": "none"},
    "line_z":           {"dim": 1, "curved": False, "curvature": "none"},
    "line_diag":        {"dim": 1, "curved": False, "curvature": "none"},
    "cross":            {"dim": 1, "curved": False, "curvature": "none"},
    "l_shape":          {"dim": 1, "curved": False, "curvature": "none"},
    "collinear":        {"dim": 1, "curved": False, "curvature": "none"},
    # 2D flat
    "triangle_xy":      {"dim": 2, "curved": False, "curvature": "none"},
    "triangle_xz":      {"dim": 2, "curved": False, "curvature": "none"},
    "triangle_3d":      {"dim": 2, "curved": False, "curvature": "none"},
    "square_xy":        {"dim": 2, "curved": False, "curvature": "none"},
    "square_xz":        {"dim": 2, "curved": False, "curvature": "none"},
    "rectangle":        {"dim": 2, "curved": False, "curvature": "none"},
    "coplanar":         {"dim": 2, "curved": False, "curvature": "none"},
    "plane":            {"dim": 2, "curved": False, "curvature": "none"},
    # 3D flat
    "tetrahedron":      {"dim": 3, "curved": False, "curvature": "none"},
    "pyramid":          {"dim": 3, "curved": False, "curvature": "none"},
    "pentachoron":      {"dim": 3, "curved": False, "curvature": "none"},
    "cube":             {"dim": 3, "curved": False, "curvature": "none"},
    "cuboid":           {"dim": 3, "curved": False, "curvature": "none"},
    "triangular_prism": {"dim": 3, "curved": False, "curvature": "none"},
    "octahedron":       {"dim": 3, "curved": False, "curvature": "none"},
    # 1D curved
    "arc":              {"dim": 1, "curved": True,  "curvature": "convex"},
    "helix":            {"dim": 1, "curved": True,  "curvature": "helical"},
    # 2D curved
    "circle":           {"dim": 2, "curved": True,  "curvature": "convex"},
    "ellipse":          {"dim": 2, "curved": True,  "curvature": "convex"},
    "disc":             {"dim": 2, "curved": True,  "curvature": "convex"},
    # 3D curved
    "sphere":           {"dim": 3, "curved": True,  "curvature": "convex"},
    "hemisphere":       {"dim": 3, "curved": True,  "curvature": "convex"},
    "cylinder":         {"dim": 3, "curved": True,  "curvature": "cylindrical"},
    "cone":             {"dim": 3, "curved": True,  "curvature": "conical"},
    "capsule":          {"dim": 3, "curved": True,  "curvature": "convex"},
    "torus":            {"dim": 3, "curved": True,  "curvature": "toroidal"},
    "shell":            {"dim": 3, "curved": True,  "curvature": "convex"},
    "tube":             {"dim": 3, "curved": True,  "curvature": "cylindrical"},
    "bowl":             {"dim": 3, "curved": True,  "curvature": "concave"},
    "saddle":           {"dim": 3, "curved": True,  "curvature": "hyperbolic"},
}

CLASS_NAMES = list(CLASS_META.keys())
NUM_CLASSES = len(CLASS_NAMES)
CLASS_TO_IDX = {n: i for i, n in enumerate(CLASS_NAMES)}

CURVATURE_NAMES = ["none", "convex", "concave", "cylindrical",
                   "conical", "toroidal", "hyperbolic", "helical"]
CURV_TO_IDX = {n: i for i, n in enumerate(CURVATURE_NAMES)}

# Edge topology for rasterized shapes
TRIANGLE_EDGES = [(0,1), (1,2), (2,0)]
QUAD_EDGES = [(0,1), (1,3), (3,2), (2,0)]
TETRA_EDGES = list(combinations(range(4), 2))  # 6 edges
CUBE_EDGES = [(0,1),(0,2),(0,4),(1,3),(1,5),(2,3),(2,6),(3,7),(4,5),(4,6),(5,7),(6,7)]
PYRAMID_EDGES = [(0,1),(1,3),(3,2),(2,0),(0,4),(1,4),(2,4),(3,4)]
PENTA_EDGES = list(combinations(range(5), 2))  # 10 edges
OCTA_EDGES = [(0,1),(0,2),(0,3),(0,4),(5,1),(5,2),(5,3),(5,4),(1,2),(2,3),(3,4),(4,1)]


def rasterize_line(p1, p2):
    """Bresenham-style 3D line rasterization between two points."""
    p1 = np.array(p1, dtype=float)
    p2 = np.array(p2, dtype=float)
    diff = p2 - p1
    n_steps = max(int(np.max(np.abs(diff))) + 1, 2)
    t = np.linspace(0, 1, n_steps)
    pts = p1[None, :] + t[:, None] * diff[None, :]
    pts = np.round(pts).astype(int)
    # Clip to grid bounds
    pts[:, 0] = np.clip(pts[:, 0], 0, GZ - 1)
    pts[:, 1] = np.clip(pts[:, 1], 0, GY - 1)
    pts[:, 2] = np.clip(pts[:, 2], 0, GX - 1)
    return np.unique(pts, axis=0)


def rasterize_edges(vertices, edges):
    """Rasterize a complete wireframe from vertex list and edge topology."""
    all_pts = [vertices]
    for i, j in edges:
        all_pts.append(rasterize_line(vertices[i], vertices[j]))
    return np.unique(np.vstack(all_pts), axis=0)


class ShapeGenerator:
    def __init__(self, seed=42):
        self.rng = np.random.RandomState(seed)

    def _pts_to_result(self, pts):
        """Convert point array to grid + metadata."""
        pts = np.atleast_2d(pts).astype(int)
        # Clip to grid
        pts[:, 0] = np.clip(pts[:, 0], 0, GZ - 1)
        pts[:, 1] = np.clip(pts[:, 1], 0, GY - 1)
        pts[:, 2] = np.clip(pts[:, 2], 0, GX - 1)
        pts = np.unique(pts, axis=0)
        grid = np.zeros(GRID_SHAPE, dtype=np.float32)
        grid[pts[:, 0], pts[:, 1], pts[:, 2]] = 1.0
        return {"grid": grid, "n_occupied": int(pts.shape[0]), "points": pts}

    def _rand_center(self, margin_z=2, margin_yx=3):
        """Random center respecting aspect ratio margins."""
        cz = self.rng.uniform(margin_z, GZ - margin_z)
        cy = self.rng.uniform(margin_yx, GY - margin_yx)
        cx = self.rng.uniform(margin_yx, GX - margin_yx)
        return np.array([cz, cy, cx])

    def _rand_pts_2d(self, n, min_dist=2):
        """Random 2D points in YX plane."""
        for _ in range(100):
            pts = np.column_stack([
                self.rng.randint(1, GY - 1, n),
                self.rng.randint(1, GX - 1, n)])
            dists = [np.linalg.norm(pts[i] - pts[j])
                     for i in range(n) for j in range(i+1, n)]
            if all(d >= min_dist for d in dists):
                return pts
        return None

    def _rand_pts_3d(self, n, min_dist=2):
        """Random 3D points respecting grid bounds."""
        for _ in range(100):
            pts = np.column_stack([
                self.rng.randint(0, GZ, n),
                self.rng.randint(1, GY - 1, n),
                self.rng.randint(1, GX - 1, n)])
            dists = [np.linalg.norm(pts[i] - pts[j])
                     for i in range(n) for j in range(i+1, n)]
            if all(d >= min_dist for d in dists):
                return pts
        return None

    def _rigid(self, name):
        """Generate rotation axis for rigid shapes."""
        axes = [(1,0,0), (0,1,0), (0,0,1)]
        return axes[self.rng.randint(len(axes))]

    def _make(self, name):
        rng = self.rng

        # === 0D ===
        if name == "point":
            z = rng.randint(0, GZ)
            y = rng.randint(0, GY)
            x = rng.randint(0, GX)
            return self._pts_to_result(np.array([[z, y, x]]))

        # === 1D lines ===
        elif name == "line_x":
            z = rng.randint(0, GZ)
            y = rng.randint(0, GY)
            x1, x2 = sorted(rng.choice(GX, 2, replace=False))
            pts = np.array([[z, y, x1], [z, y, x2]])
            return self._pts_to_result(rasterize_line(pts[0], pts[1]))

        elif name == "line_y":
            z = rng.randint(0, GZ)
            x = rng.randint(0, GX)
            y1, y2 = sorted(rng.choice(GY, 2, replace=False))
            pts = np.array([[z, y1, x], [z, y2, x]])
            return self._pts_to_result(rasterize_line(pts[0], pts[1]))

        elif name == "line_z":
            y = rng.randint(0, GY)
            x = rng.randint(0, GX)
            z1, z2 = sorted(rng.choice(GZ, 2, replace=False))
            pts = np.array([[z1, y, x], [z2, y, x]])
            return self._pts_to_result(rasterize_line(pts[0], pts[1]))

        elif name == "line_diag":
            p1 = np.array([rng.randint(0, GZ), rng.randint(0, GY), rng.randint(0, GX)])
            p2 = np.array([rng.randint(0, GZ), rng.randint(0, GY), rng.randint(0, GX)])
            if np.linalg.norm(p1 - p2) < 3:
                return None
            return self._pts_to_result(rasterize_line(p1, p2))

        elif name == "cross":
            c = self._rand_center(margin_z=1, margin_yx=2)
            arm_yx = rng.randint(2, min(6, GY // 2))
            arm_z = rng.randint(1, min(3, GZ // 2))
            pts = []
            ci = np.round(c).astype(int)
            # YX cross
            for dy in range(-arm_yx, arm_yx + 1):
                pts.append([ci[0], np.clip(ci[1] + dy, 0, GY-1), ci[2]])
            for dx in range(-arm_yx, arm_yx + 1):
                pts.append([ci[0], ci[1], np.clip(ci[2] + dx, 0, GX-1)])
            # Z arm
            for dz in range(-arm_z, arm_z + 1):
                pts.append([np.clip(ci[0] + dz, 0, GZ-1), ci[1], ci[2]])
            return self._pts_to_result(np.array(pts))

        elif name == "l_shape":
            c = self._rand_center(margin_z=1, margin_yx=2)
            arm = rng.randint(2, min(5, GY // 2))
            ci = np.round(c).astype(int)
            pts = []
            for dy in range(arm + 1):
                pts.append([ci[0], np.clip(ci[1] + dy, 0, GY-1), ci[2]])
            for dx in range(1, arm + 1):
                pts.append([ci[0], ci[1], np.clip(ci[2] + dx, 0, GX-1)])
            return self._pts_to_result(np.array(pts))

        elif name == "collinear":
            # Identity = 3 discrete points on a line, NOT a line segment
            axis = rng.randint(3)
            gs = [GZ, GY, GX]
            vals = sorted(rng.choice(gs[axis], 3, replace=False))
            fixed = [rng.randint(0, gs[(axis+1)%3]), rng.randint(0, gs[(axis+2)%3])]
            pts = np.zeros((3, 3), dtype=int)
            for i, v in enumerate(vals):
                pts[i, axis] = v
                pts[i, (axis + 1) % 3] = fixed[0]
                pts[i, (axis + 2) % 3] = fixed[1]
            return self._pts_to_result(pts)

        # === 2D flat ===
        elif name == "triangle_xy":
            z = rng.randint(0, GZ)
            pts2d = self._rand_pts_2d(3, min_dist=3)
            if pts2d is None: return None
            return self._pts_to_result(np.column_stack([np.full(3, z), pts2d]))

        elif name == "triangle_xz":
            y = rng.randint(0, GY)
            for _ in range(50):
                pts = np.column_stack([
                    rng.randint(0, GZ, 3),
                    np.full(3, y),
                    rng.randint(1, GX - 1, 3)])
                dists = [np.linalg.norm(pts[i] - pts[j]) for i in range(3) for j in range(i+1, 3)]
                if all(d >= 2 for d in dists):
                    return self._pts_to_result(pts)
            return None

        elif name == "triangle_3d":
            verts = self._rand_pts_3d(3, min_dist=3)
            if verts is None: return None
            return self._pts_to_result(verts)

        elif name == "square_xy":
            z = rng.randint(0, GZ)
            s = rng.randint(3, min(7, GY - 2))
            cy, cx = rng.randint(s, GY - s), rng.randint(s, GX - s)
            verts = np.array([
                [z, cy - s, cx - s], [z, cy - s, cx + s],
                [z, cy + s, cx - s], [z, cy + s, cx + s]])
            return self._pts_to_result(rasterize_edges(verts, QUAD_EDGES))

        elif name == "square_xz":
            y = rng.randint(0, GY)
            s_z = rng.randint(1, min(3, GZ // 2))
            s_x = rng.randint(2, min(6, GX // 2))
            cz, cx = rng.randint(s_z, GZ - s_z), rng.randint(s_x, GX - s_x)
            verts = np.array([
                [cz - s_z, y, cx - s_x], [cz - s_z, y, cx + s_x],
                [cz + s_z, y, cx - s_x], [cz + s_z, y, cx + s_x]])
            return self._pts_to_result(rasterize_edges(verts, QUAD_EDGES))

        elif name == "rectangle":
            z = rng.randint(0, GZ)
            sy = rng.randint(2, min(6, GY // 2))
            sx = rng.randint(2, min(6, GX // 2))
            while abs(sy - sx) < 2:
                sy = rng.randint(2, min(6, GY // 2))
                sx = rng.randint(2, min(6, GX // 2))
            cy, cx = rng.randint(sy, GY - sy), rng.randint(sx, GX - sx)
            verts = np.array([
                [z, cy - sy, cx - sx], [z, cy - sy, cx + sx],
                [z, cy + sy, cx - sx], [z, cy + sy, cx + sx]])
            return self._pts_to_result(rasterize_edges(verts, QUAD_EDGES))

        elif name == "coplanar":
            # Identity = 4 discrete coplanar points, NOT a quadrilateral
            pts = self._rand_pts_3d(4, min_dist=2)
            if pts is None: return None
            axis = rng.randint(3)
            pts[:, axis] = pts[0, axis]
            return self._pts_to_result(pts)

        elif name == "plane":
            axis = rng.randint(3)
            gs = [GZ, GY, GX]
            pos = rng.randint(0, gs[axis])
            thick = rng.randint(1, max(2, gs[axis] // 4) + 1)
            mask = np.zeros(GRID_SHAPE, dtype=np.float32)
            for t in range(thick):
                p = min(pos + t, gs[axis] - 1)
                if axis == 0: mask[p, :, :] = 1
                elif axis == 1: mask[:, p, :] = 1
                else: mask[:, :, p] = 1
            pts = np.argwhere(mask > 0)
            return self._pts_to_result(pts)

        # === 3D polyhedra (rasterized edges) ===
        elif name == "tetrahedron":
            verts = self._rand_pts_3d(4, min_dist=3)
            if verts is None: return None
            return self._pts_to_result(rasterize_edges(verts, TETRA_EDGES))

        elif name == "pyramid":
            base_y = rng.randint(2, GY - 2)
            s = rng.randint(2, min(5, GX // 2))
            cy, cx = rng.randint(s + 1, GY - s - 1), rng.randint(s + 1, GX - s - 1)
            base_z = rng.randint(1, GZ - 2)
            apex_z = rng.randint(0, GZ) if rng.random() < 0.5 else base_z + rng.randint(2, min(4, GZ - base_z))
            apex_z = min(apex_z, GZ - 1)
            verts = np.array([
                [base_z, cy - s, cx - s], [base_z, cy - s, cx + s],
                [base_z, cy + s, cx - s], [base_z, cy + s, cx + s],
                [apex_z, cy, cx]])
            return self._pts_to_result(rasterize_edges(verts, PYRAMID_EDGES))

        elif name == "pentachoron":
            verts = self._rand_pts_3d(5, min_dist=3)
            if verts is None: return None
            return self._pts_to_result(rasterize_edges(verts, PENTA_EDGES))

        elif name == "cube":
            s_z = rng.randint(1, min(3, GZ // 2))
            s_yx = rng.randint(2, min(5, GY // 2))
            c = self._rand_center(margin_z=s_z + 1, margin_yx=s_yx + 1)
            ci = np.round(c).astype(int)
            verts = np.array([
                [ci[0]-s_z, ci[1]-s_yx, ci[2]-s_yx],
                [ci[0]-s_z, ci[1]-s_yx, ci[2]+s_yx],
                [ci[0]-s_z, ci[1]+s_yx, ci[2]-s_yx],
                [ci[0]-s_z, ci[1]+s_yx, ci[2]+s_yx],
                [ci[0]+s_z, ci[1]-s_yx, ci[2]-s_yx],
                [ci[0]+s_z, ci[1]-s_yx, ci[2]+s_yx],
                [ci[0]+s_z, ci[1]+s_yx, ci[2]-s_yx],
                [ci[0]+s_z, ci[1]+s_yx, ci[2]+s_yx]])
            return self._pts_to_result(rasterize_edges(verts, CUBE_EDGES))

        elif name == "cuboid":
            sz = rng.randint(1, min(3, GZ // 2))
            sy = rng.randint(2, min(6, GY // 2))
            sx = rng.randint(2, min(6, GX // 2))
            # Ensure at least one dimension differs significantly
            while abs(sy - sx) < 2 and abs(sz * 2 - sy) < 2:
                sy = rng.randint(2, min(6, GY // 2))
                sx = rng.randint(2, min(6, GX // 2))
            c = self._rand_center(margin_z=sz + 1, margin_yx=max(sy, sx) + 1)
            ci = np.round(c).astype(int)
            verts = np.array([
                [ci[0]-sz, ci[1]-sy, ci[2]-sx], [ci[0]-sz, ci[1]-sy, ci[2]+sx],
                [ci[0]-sz, ci[1]+sy, ci[2]-sx], [ci[0]-sz, ci[1]+sy, ci[2]+sx],
                [ci[0]+sz, ci[1]-sy, ci[2]-sx], [ci[0]+sz, ci[1]-sy, ci[2]+sx],
                [ci[0]+sz, ci[1]+sy, ci[2]-sx], [ci[0]+sz, ci[1]+sy, ci[2]+sx]])
            return self._pts_to_result(rasterize_edges(verts, CUBE_EDGES))

        elif name == "triangular_prism":
            z1, z2 = sorted(rng.choice(GZ, 2, replace=False))
            pts2d = self._rand_pts_2d(3, min_dist=3)
            if pts2d is None: return None
            # Two triangular faces + connecting edges
            top = np.column_stack([np.full(3, z1), pts2d])
            bot = np.column_stack([np.full(3, z2), pts2d])
            verts = np.vstack([top, bot])
            edges = [(0,1),(1,2),(2,0),(3,4),(4,5),(5,3),(0,3),(1,4),(2,5)]
            return self._pts_to_result(rasterize_edges(verts, edges))

        elif name == "octahedron":
            c = self._rand_center(margin_z=2, margin_yx=3)
            ci = np.round(c).astype(int)
            rz = rng.randint(1, min(3, GZ // 2))
            ryx = rng.randint(2, min(5, GY // 2))
            verts = np.array([
                [ci[0], ci[1] + ryx, ci[2]], [ci[0], ci[1], ci[2] + ryx],
                [ci[0], ci[1] - ryx, ci[2]], [ci[0], ci[1], ci[2] - ryx],
                [ci[0] + rz, ci[1], ci[2]], [ci[0] - rz, ci[1], ci[2]]])
            return self._pts_to_result(rasterize_edges(verts, OCTA_EDGES))

        # === 1D curved ===
        elif name == "arc":
            plane = rng.randint(3)
            c = self._rand_center(margin_z=1, margin_yx=2)
            r_main = rng.uniform(2.0, min(5.0, GY / 2 - 1))
            r_z = rng.uniform(1.0, min(3.0, GZ / 2 - 1)) if plane != 0 else r_main
            angle_start = rng.uniform(0, np.pi)
            angle_span = rng.uniform(np.pi / 3, np.pi)
            n_pts = max(8, int(angle_span * r_main))
            t = np.linspace(angle_start, angle_start + angle_span, n_pts)
            if plane == 0:  # YX plane
                pts = np.column_stack([
                    np.full(n_pts, c[0]),
                    c[1] + r_main * np.cos(t),
                    c[2] + r_main * np.sin(t)])
            elif plane == 1:  # ZX plane
                pts = np.column_stack([
                    c[0] + r_z * np.cos(t),
                    np.full(n_pts, c[1]),
                    c[2] + r_main * np.sin(t)])
            else:  # ZY plane
                pts = np.column_stack([
                    c[0] + r_z * np.cos(t),
                    c[1] + r_main * np.sin(t),
                    np.full(n_pts, c[2])])
            pts = np.round(pts).astype(int)
            return self._pts_to_result(pts)

        elif name == "helix":
            c = self._rand_center(margin_z=0, margin_yx=3)
            r = rng.uniform(1.5, min(4.0, GY / 2 - 2))
            turns = rng.uniform(1.0, 2.5)
            n_pts = int(turns * 20)
            t = np.linspace(0, turns * 2 * np.pi, n_pts)
            z_span = GZ - 1
            pts = np.column_stack([
                t / (turns * 2 * np.pi) * z_span,
                c[1] + r * np.cos(t),
                c[2] + r * np.sin(t)])
            pts = np.round(pts).astype(int)
            return self._pts_to_result(pts)

        # === 2D curved ===
        elif name == "circle":
            plane = rng.randint(3)
            c = self._rand_center(margin_z=1, margin_yx=3)
            r = rng.uniform(2.0, min(5.0, GY / 2 - 1))
            n_pts = max(12, int(2 * np.pi * r))
            t = np.linspace(0, 2 * np.pi, n_pts, endpoint=False)
            if plane == 0:
                pts = np.column_stack([
                    np.full(n_pts, c[0]),
                    c[1] + r * np.cos(t),
                    c[2] + r * np.sin(t)])
            elif plane == 1:
                r_z = min(r, GZ / 2 - 1)
                pts = np.column_stack([
                    c[0] + r_z * np.cos(t),
                    np.full(n_pts, c[1]),
                    c[2] + r * np.sin(t)])
            else:
                r_z = min(r, GZ / 2 - 1)
                pts = np.column_stack([
                    c[0] + r_z * np.cos(t),
                    c[1] + r * np.sin(t),
                    np.full(n_pts, c[2])])
            pts = np.round(pts).astype(int)
            return self._pts_to_result(pts)

        elif name == "ellipse":
            c = self._rand_center(margin_z=1, margin_yx=3)
            ry = rng.uniform(2.0, min(5.0, GY / 2 - 1))
            ratio = rng.uniform(1.6, 2.5)
            if rng.random() < 0.5:
                rx = ry / ratio
            else:
                rx = ry * ratio
                rx = min(rx, GX / 2 - 1)
                if rx / ry < 1.6: ry = rx / 1.6
            n_pts = max(16, int(2 * np.pi * max(rx, ry)))
            t = np.linspace(0, 2 * np.pi, n_pts, endpoint=False)
            pts = np.column_stack([
                np.full(n_pts, c[0]),
                c[1] + ry * np.cos(t),
                c[2] + rx * np.sin(t)])
            pts = np.round(pts).astype(int)
            return self._pts_to_result(pts)

        elif name == "disc":
            plane = rng.randint(3)
            c = self._rand_center(margin_z=1, margin_yx=3)
            r = rng.uniform(2.0, min(5.0, GY / 2 - 1))
            if plane == 0:
                mask = ((_COORDS[:, 1] - c[1])**2 + (_COORDS[:, 2] - c[2])**2 <= r**2) & \
                       (np.abs(_COORDS[:, 0] - c[0]) < 0.6)
            elif plane == 1:
                r_z = min(r, GZ / 2 - 1)
                mask = ((_COORDS[:, 0] - c[0])**2 / max(r_z, 0.5)**2 +
                        (_COORDS[:, 2] - c[2])**2 / r**2 <= 1) & \
                       (np.abs(_COORDS[:, 1] - c[1]) < 0.6)
            else:
                r_z = min(r, GZ / 2 - 1)
                mask = ((_COORDS[:, 0] - c[0])**2 / max(r_z, 0.5)**2 +
                        (_COORDS[:, 1] - c[1])**2 / r**2 <= 1) & \
                       (np.abs(_COORDS[:, 2] - c[2]) < 0.6)
            pts = _COORDS[mask].astype(int)
            if len(pts) < 3: return None
            return self._pts_to_result(pts)

        # === 3D curved ===
        elif name == "sphere":
            c = self._rand_center(margin_z=2, margin_yx=3)
            r = rng.uniform(2.0, min(3.5, GZ / 2 - 0.5, GY / 2 - 1))
            # Use ellipsoidal check respecting aspect ratio
            d2 = ((_COORDS[:, 0] - c[0]) / r)**2 + \
                 ((_COORDS[:, 1] - c[1]) / r)**2 + \
                 ((_COORDS[:, 2] - c[2]) / r)**2
            mask = d2 <= 1.0
            pts = _COORDS[mask].astype(int)
            if len(pts) < 4: return None
            return self._pts_to_result(pts)

        elif name == "hemisphere":
            c = self._rand_center(margin_z=2, margin_yx=3)
            r = rng.uniform(2.0, min(3.5, GZ / 2 - 0.5, GY / 2 - 1))
            cut_axis = rng.randint(3)
            d2 = ((_COORDS[:, 0] - c[0]) / r)**2 + \
                 ((_COORDS[:, 1] - c[1]) / r)**2 + \
                 ((_COORDS[:, 2] - c[2]) / r)**2
            mask = d2 <= 1.0
            if cut_axis == 0: mask &= _COORDS[:, 0] >= c[0]
            elif cut_axis == 1: mask &= _COORDS[:, 1] >= c[1]
            else: mask &= _COORDS[:, 2] >= c[2]
            pts = _COORDS[mask].astype(int)
            if len(pts) < 3: return None
            return self._pts_to_result(pts)

        elif name == "cylinder":
            axis = rng.randint(3)
            c = self._rand_center(margin_z=0, margin_yx=3)
            r = rng.uniform(1.5, min(3.0, GY / 2 - 1))
            if axis == 0:
                d2 = (_COORDS[:, 1] - c[1])**2 + (_COORDS[:, 2] - c[2])**2
                mask = d2 <= r**2
            elif axis == 1:
                r_z = min(r, GZ / 2 - 0.5)
                d2 = (_COORDS[:, 0] - c[0])**2 / max(r_z, 0.5)**2 + \
                     (_COORDS[:, 2] - c[2])**2 / r**2
                mask = d2 <= 1.0
            else:
                r_z = min(r, GZ / 2 - 0.5)
                d2 = (_COORDS[:, 0] - c[0])**2 / max(r_z, 0.5)**2 + \
                     (_COORDS[:, 1] - c[1])**2 / r**2
                mask = d2 <= 1.0
            pts = _COORDS[mask].astype(int)
            if len(pts) < 4: return None
            return self._pts_to_result(pts)

        elif name == "cone":
            axis = rng.randint(3)
            c = self._rand_center(margin_z=1, margin_yx=3)
            r = rng.uniform(2.0, min(4.0, GY / 2 - 1))
            gs = [GZ, GY, GX]
            h = gs[axis] - 1
            apex_frac = _COORDS[:, axis] / max(h, 1)
            local_r = r * (1.0 - apex_frac)
            if axis == 0:
                d2 = (_COORDS[:, 1] - c[1])**2 + (_COORDS[:, 2] - c[2])**2
            elif axis == 1:
                d2 = (_COORDS[:, 0] - c[0])**2 + (_COORDS[:, 2] - c[2])**2
            else:
                d2 = (_COORDS[:, 0] - c[0])**2 + (_COORDS[:, 1] - c[1])**2
            mask = d2 <= local_r**2
            pts = _COORDS[mask].astype(int)
            if len(pts) < 4: return None
            return self._pts_to_result(pts)

        elif name == "capsule":
            axis = rng.randint(3)
            c = self._rand_center(margin_z=1, margin_yx=3)
            r = rng.uniform(1.5, min(2.5, GZ / 2 - 0.5, GY / 2 - 1))
            gs = [GZ, GY, GX]
            half_h = rng.uniform(1.0, gs[axis] / 2 - r - 0.5)
            # Cylinder body + spherical caps
            dist_axis = np.abs(_COORDS[:, axis] - c[axis])
            clamped = np.clip(dist_axis - half_h, 0, None)
            perp_axes = [i for i in range(3) if i != axis]
            d2 = clamped**2
            for a in perp_axes:
                d2 += (_COORDS[:, a] - c[a])**2
            mask = d2 <= r**2
            pts = _COORDS[mask].astype(int)
            if len(pts) < 4: return None
            return self._pts_to_result(pts)

        elif name == "torus":
            c = self._rand_center(margin_z=2, margin_yx=4)
            R = rng.uniform(2.5, min(4.0, GY / 2 - 2))
            r = rng.uniform(0.8, min(1.5, GZ / 2 - 0.5, R * 0.5))
            # Torus in YX plane
            d_yx = np.sqrt((_COORDS[:, 1] - c[1])**2 + (_COORDS[:, 2] - c[2])**2)
            d2 = (d_yx - R)**2 + (_COORDS[:, 0] - c[0])**2
            mask = d2 <= r**2
            pts = _COORDS[mask].astype(int)
            if len(pts) < 4: return None
            return self._pts_to_result(pts)

        elif name == "shell":
            c = self._rand_center(margin_z=1, margin_yx=3)
            r = rng.uniform(2.0, min(3.5, GZ / 2 - 0.5, GY / 2 - 1))
            thick = rng.uniform(0.4, 0.8)
            d2 = ((_COORDS[:, 0] - c[0]) / r)**2 + \
                 ((_COORDS[:, 1] - c[1]) / r)**2 + \
                 ((_COORDS[:, 2] - c[2]) / r)**2
            mask = (d2 <= 1.0) & (d2 >= (1.0 - thick)**2)
            pts = _COORDS[mask].astype(int)
            if len(pts) < 4: return None
            return self._pts_to_result(pts)

        elif name == "tube":
            axis = rng.randint(3)
            c = self._rand_center(margin_z=0, margin_yx=3)
            r_out = rng.uniform(2.0, min(3.5, GY / 2 - 1))
            r_in = r_out * rng.uniform(0.4, 0.7)
            if axis == 0:
                d2 = (_COORDS[:, 1] - c[1])**2 + (_COORDS[:, 2] - c[2])**2
            elif axis == 1:
                d2 = (_COORDS[:, 0] - c[0])**2 + (_COORDS[:, 2] - c[2])**2
            else:
                d2 = (_COORDS[:, 0] - c[0])**2 + (_COORDS[:, 1] - c[1])**2
            mask = (d2 <= r_out**2) & (d2 >= r_in**2)
            pts = _COORDS[mask].astype(int)
            if len(pts) < 4: return None
            return self._pts_to_result(pts)

        elif name == "bowl":
            c = self._rand_center(margin_z=1, margin_yx=3)
            r = rng.uniform(2.0, min(3.5, GZ / 2 - 0.5, GY / 2 - 1))
            thick = rng.uniform(0.3, 0.7)
            d2 = ((_COORDS[:, 0] - c[0]) / r)**2 + \
                 ((_COORDS[:, 1] - c[1]) / r)**2 + \
                 ((_COORDS[:, 2] - c[2]) / r)**2
            mask = (d2 <= 1.0) & (d2 >= (1.0 - thick)**2) & (_COORDS[:, 0] <= c[0])
            pts = _COORDS[mask].astype(int)
            if len(pts) < 3: return None
            return self._pts_to_result(pts)

        elif name == "saddle":
            c = self._rand_center(margin_z=2, margin_yx=4)
            scale = rng.uniform(1.5, 3.0)
            dy = (_COORDS[:, 1] - c[1]) / scale
            dx = (_COORDS[:, 2] - c[2]) / scale
            z_saddle = c[0] + (dy**2 - dx**2)
            mask = np.abs(_COORDS[:, 0] - z_saddle) < 0.8
            pts = _COORDS[mask].astype(int)
            if len(pts) < 4: return None
            return self._pts_to_result(pts)

        return None

    def generate(self, name, max_retries=10):
        """Generate one sample with retries."""
        for _ in range(max_retries):
            result = self._make(name)
            if result is not None and result["n_occupied"] > 0:
                return result
        return None

    def generate_dataset(self, n_per_class, seed=None):
        """Generate balanced dataset."""
        if seed is not None:
            self.rng = np.random.RandomState(seed)
        grids, labels, dims, curveds = [], [], [], []
        for cls_idx, name in enumerate(CLASS_NAMES):
            meta = CLASS_META[name]
            count = 0
            while count < n_per_class:
                self.rng = np.random.RandomState(seed * 1000 + cls_idx * n_per_class + count if seed else None)
                result = self.generate(name)
                if result is not None:
                    grids.append(result["grid"])
                    labels.append(cls_idx)
                    dims.append(meta["dim"])
                    curveds.append(1 if meta["curved"] else 0)
                    count += 1
        return {
            "grids": np.array(grids),
            "labels": np.array(labels),
            "dims": np.array(dims),
            "curveds": np.array(curveds),
        }


# === Verification =============================================================
if __name__ == "__main__":
    gen = ShapeGenerator(seed=42)
    print(f"Grid: {GZ}×{GY}×{GX} = {GRID_VOLUME} voxels")
    print(f"Classes: {NUM_CLASSES}")
    print(f"\n{'Shape':20s} {'OK':>4s} {'Avg vox':>8s}")
    print("-" * 36)
    for name in CLASS_NAMES:
        ok = 0; voxels = []
        for trial in range(20):
            gen.rng = np.random.RandomState(trial * 100 + hash(name) % 10000)
            s = gen.generate(name)
            if s:
                ok += 1
                voxels.append(s["n_occupied"])
        avg = np.mean(voxels) if voxels else 0
        status = "✓" if ok >= 15 else "✗"
        print(f"  {status} {name:20s} {ok:2d}/20 {avg:7.1f}")

print(f'\nLoaded {NUM_CLASSES} shape classes, grid={GZ}×{GY}×{GX}')