collision_manager.py

"""
collision_entities.py

Requirements:
    pip install numpy
Optional (for visual debug): pip install PyOpenGL PyOpenGL_accelerate
"""

import numpy as np
from typing import Callable, Optional, List, Tuple

# ---------------------------
# Utilities
# ---------------------------

_EPS = 1e-9

def normalize(v: np.ndarray) -> np.ndarray:
    n = np.linalg.norm(v)
    if n <= _EPS:
        return None
    return v / n

def transform_points(points: np.ndarray, mat: np.ndarray) -> np.ndarray:
    """
    points: (N,3)
    mat: (4,4)
    returns transformed points (N,3)
    """
    n = points.shape[0]
    hom = np.ones((n, 4), dtype=np.float64)
    hom[:, :3] = points
    th = hom @ mat.T
    return th[:, :3] / th[:, 3:4]

# ---------------------------
# Model matrix helpers
# ---------------------------

def mat_identity() -> np.ndarray:
    return np.eye(4, dtype=np.float64)

def mat_translate(tx=0, ty=0, tz=0) -> np.ndarray:
    m = np.eye(4, dtype=np.float64)
    m[:3, 3] = (tx, ty, tz)
    return m

def mat_scale(sx=1, sy=1, sz=1) -> np.ndarray:
    m = np.eye(4, dtype=np.float64)
    m[0,0], m[1,1], m[2,2] = sx, sy, sz
    return m

def mat_rotate_euler(rx=0.0, ry=0.0, rz=0.0) -> np.ndarray:
    """Euler rotations in radians: rx around X, ry around Y, rz around Z (ZYX order)."""
    cx, sx = np.cos(rx), np.sin(rx)
    cy, sy = np.cos(ry), np.sin(ry)
    cz, sz = np.cos(rz), np.sin(rz)
    Rx = np.array([[1,0,0,0],[0,cx,-sx,0],[0,sx,cx,0],[0,0,0,1]], dtype=np.float64)
    Ry = np.array([[cy,0,sy,0],[0,1,0,0],[-sy,0,cy,0],[0,0,0,1]], dtype=np.float64)
    Rz = np.array([[cz,-sz,0,0],[sz,cz,0,0],[0,0,1,0],[0,0,0,1]], dtype=np.float64)
    return Rz @ Ry @ Rx

def mat_shear(xy=0, xz=0, yx=0, yz=0, zx=0, zy=0) -> np.ndarray:
    """
    Creates an affine shear (skew) matrix. 
    Each parameter is the shear factor of the first axis with respect to the second, etc.
    """
    m = np.eye(4, dtype=np.float64)
    m[0,1] = xy
    m[0,2] = xz
    m[1,0] = yx
    m[1,2] = yz
    m[2,0] = zx
    m[2,1] = zy
    return m

# ---------------------------
# Collision entity
# ---------------------------

class CollisionEntity:
    """
    Represents a box-like collision object with:
      - local half extents (hx, hy, hz) centered at local origin
      - an arbitrary 4x4 model matrix (affine) that places it in world space (rotation, skew, scale, translate)
      - optional callback and trigger behavior
      - optional debug grid render that follows transform
    """

    __slots__ = (
        "name", "half_extents", "model_matrix",
        "callback", "trigger_once", "triggered",
        "enabled", "show_grid", "grid_subdiv",
    )

    def __init__(
        self,
        name: Optional[str],
        half_extents=(0.5, 0.5, 0.5),
        model_matrix: Optional[np.ndarray] = None,
        callback: Optional[Callable[['CollisionEntity', 'CollisionEntity'], None]] = None,
        trigger_once: bool = True,
        show_grid: bool = False,
        grid_subdiv: int = 2,
    ):
        self.name = name
        self.half_extents = np.asarray(half_extents, dtype=np.float64)
        self.model_matrix = mat_identity() if model_matrix is None else np.asarray(model_matrix, dtype=np.float64)
        self.callback = callback
        self.trigger_once = bool(trigger_once)
        self.triggered = False
        self.enabled = True
        self.show_grid = bool(show_grid)
        self.grid_subdiv = max(1, int(grid_subdiv))

    # --- local corner generation (centered box) ---
    def local_corners(self) -> np.ndarray:
        hx, hy, hz = self.half_extents
        # Order: (-x,-y,-z), ( x,-y,-z), (-x, y,-z), ( x, y,-z), (-x,-y, z), ( x,-y, z), (-x, y, z), ( x, y, z)
        return np.array([
            [-hx, -hy, -hz],
            [ hx, -hy, -hz],
            [-hx,  hy, -hz],
            [ hx,  hy, -hz],
            [-hx, -hy,  hz],
            [ hx, -hy,  hz],
            [-hx,  hy,  hz],
            [ hx,  y ,  hz] if False else [ hx,  hy,  hz]  # ensure literal
        ], dtype=np.float64)

    def world_corners(self) -> np.ndarray:
        """Return (8,3) array of corners transformed to world space."""
        lc = self.local_corners()
        return transform_points(lc, self.model_matrix)

    def edges_from_corner0(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Return three edge vectors (from corner 0) in world space: e1, e2, e3"""
        w = self.world_corners()
        e1 = w[1] - w[0]  # x-direction
        e2 = w[2] - w[0]  # y-direction
        e3 = w[4] - w[0]  # z-direction
        return e1, e2, e3

    # --- grid generation (local) ---
    def _face_quads_local(self) -> List[Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]]:
        """Return 6 faces as 4-corner quads in local coords (each corner order is consistent)."""
        hx, hy, hz = self.half_extents
        # same corner ordering as local_corners
        c = lambda x,y,z: np.array([x*hx, y*hy, z*hz], dtype=np.float64)
        # faces: -Z, +Z, -Y, +Y, -X, +X (each quad in (u,v) consistent)
        return [
            (c(-1,-1,-1), c(1,-1,-1), c(-1,1,-1), c(1,1,-1)),  # back (-Z)
            (c(-1,-1, 1), c(1,-1, 1), c(-1,1, 1), c(1,1, 1)),  # front (+Z)
            (c(-1,-1,-1), c(1,-1,-1), c(-1,-1,1), c(1,-1,1)),  # bottom (-Y)
            (c(-1, 1,-1), c(1, 1,-1), c(-1, 1,1), c(1, 1,1)),  # top (+Y)
            (c(-1,-1,-1), c(-1,1,-1), c(-1,-1,1), c(-1,1,1)),  # left (-X)
            (c( 1,-1,-1), c( 1,1,-1), c( 1,-1,1), c( 1,1,1)),  # right (+X)
        ]

    def generate_grid_lines_world(self) -> np.ndarray:
        """
        Return an (M,2,3) array of line endpoints in world space representing a grid
        on the box surfaces. Useful for debug rendering.
        """
        faces = self._face_quads_local()
        subdiv = max(1, int(self.grid_subdiv))
        lines = []
        for quad in faces:
            # quad corners: a (u=0,v=0), b (u=1,v=0), c (u=0,v=1), d (u=1,v=1)
            a,b,c,d = quad
            # generate (subdiv+1) lines in u direction and v direction
            for i in range(subdiv+1):
                t = i / subdiv
                # u-line across v
                p0 = a * (1-t) + b * t
                p1 = c * (1-t) + d * t
                lines.append((p0, p1))
            for j in range(subdiv+1):
                s = j / subdiv
                # v-line across u
                p0 = a * (1-s) + c * s
                p1 = b * (1-s) + d * s
                lines.append((p0, p1))
        # transform all line endpoints to world
        pts = np.array([p for pair in lines for p in pair], dtype=np.float64)  # (2*M, 3)
        wpts = transform_points(pts, self.model_matrix)
        # reshape into (M,2,3)
        M = len(lines)
        return wpts.reshape((M, 2, 3))

    # --- simple API helpers ---
    def set_model_matrix(self, mat: np.ndarray):
        self.model_matrix = np.asarray(mat, dtype=np.float64)

    def set_enabled(self, val: bool):
        self.enabled = bool(val)

    def reset_triggered(self):
        self.triggered = False

# ---------------------------
# Collision test (SAT for parallelepipeds)
# ---------------------------

def project_points_on_axis(points: np.ndarray, axis: np.ndarray) -> Tuple[float, float]:
    """
    points: (N,3), axis: (3,)
    returns (min_proj, max_proj)
    """
    projs = points @ axis
    return float(projs.min()), float(projs.max())

def box_face_normals_from_edges(e1: np.ndarray, e2: np.ndarray, e3: np.ndarray) -> List[np.ndarray]:
    # face normals are cross pairs: cross(e1,e2), cross(e1,e3), cross(e2,e3)
    n1 = np.cross(e1, e2)
    n2 = np.cross(e1, e3)
    n3 = np.cross(e2, e3)
    ns = []
    for n in (n1, n2, n3):
        u = normalize(n)
        if u is not None:
            ns.append(u)
    return ns

def parallelepiped_sat_collision(entA: CollisionEntity, entB: CollisionEntity) -> bool:
    """
    SAT between two parallelepipeds (affine-transformed boxes) by using:
        - face normals of each (3 per box)
        - cross products of each edge of A with each edge of B (up to 9)
    Project all 8 corners of both boxes onto each axis. If any axis separates, return False.
    """
    if (not entA.enabled) or (not entB.enabled):
        return False

    A_pts = entA.world_corners()
    B_pts = entB.world_corners()

    # Quick coarse test: bounding spheres (cheap reject)
    centerA = A_pts.mean(axis=0)
    centerB = B_pts.mean(axis=0)
    rA = np.max(np.linalg.norm(A_pts - centerA, axis=1))
    rB = np.max(np.linalg.norm(B_pts - centerB, axis=1))
    if np.sum((centerA - centerB)**2) > (rA + rB)**2:
        return False

    # get edges (not normalized)
    eA1, eA2, eA3 = entA.edges_from_corner0()
    eB1, eB2, eB3 = entB.edges_from_corner0()

    axes = []
    # face normals
    axes += box_face_normals_from_edges(eA1, eA2, eA3)
    axes += box_face_normals_from_edges(eB1, eB2, eB3)

    # cross products of edges
    for ea in (eA1, eA2, eA3):
        for eb in (eB1, eB2, eB3):
            c = np.cross(ea, eb)
            u = normalize(c)
            if u is not None:
                axes.append(u)

    # final axis list deduplicated by direction (tolerance)
    kept = []
    for ax in axes:
        if ax is None:
            continue
        # unify sign
        if ax[0] < -0.0 or (abs(ax[0])<1e-12 and ax[1] < -0.0) or (abs(ax[0])<1e-12 and abs(ax[1])<1e-12 and ax[2] < -0.0):
            ax = -ax
        skip = False
        for k in kept:
            if np.allclose(k, ax, atol=1e-6):
                skip = True
                break
        if not skip:
            kept.append(ax)
    axes = kept

    # project on each axis
    for ax in axes:
        minA, maxA = project_points_on_axis(A_pts, ax)
        minB, maxB = project_points_on_axis(B_pts, ax)
        if maxA < minB - _EPS or maxB < minA - _EPS:
            # separated
            return False
    # no separating axis => intersection
    return True

# ---------------------------
# Simple collision check + trigger helper
# ---------------------------

def check_collision(a: CollisionEntity, b: CollisionEntity, method: str = "sat") -> bool:
    """
    method:
      - 'sat' (default): the accurate SAT/parallelepiped test above (works for rotation/skew/scale).
      - 'aabb' : coarse test using world AABB of transformed corners.
      - 'sphere': extremely cheap bounding sphere test.
    """
    if method == "sphere":
        A = a.world_corners(); B = b.world_corners()
        ca = A.mean(axis=0); cb = B.mean(axis=0)
        ra = np.max(np.linalg.norm(A - ca, axis=1)); rb = np.max(np.linalg.norm(B - cb, axis=1))
        return np.sum((ca - cb)**2) <= (ra + rb)**2
    if method == "aabb":
        A = a.world_corners(); B = b.world_corners()
        amin = A.min(axis=0); amax = A.max(axis=0)
        bmin = B.min(axis=0); bmax = B.max(axis=0)
        return np.all(amax >= bmin) and np.all(bmax >= amin)
    # default:
    return parallelepiped_sat_collision(a, b)

def check_and_trigger(a: CollisionEntity, b: CollisionEntity, method: str = "sat", call_both: bool = True) -> bool:
    """
    Checks collision between a and b using method and runs callbacks.
    - If collided:
        - call a.callback(a, b) if allowed by a.trigger_once
        - if call_both: call b.callback(b, a)
    - Returns True if collision detected.
    """
    collided = check_collision(a, b, method=method)
    if collided:
        if a.callback is not None and (not a.trigger_once or not a.triggered):
            a.callback(a, b)
            a.triggered = True
        if call_both:
            if b.callback is not None and (not b.trigger_once or not b.triggered):
                b.callback(b, a)
                b.triggered = True
    else:
        # do not auto-reset one-shot triggers — leave to user via reset_triggered()
        pass
    return collided

# ---------------------------
# Manager for named entities
# ---------------------------

class CollisionManager:
    __slots__ = ("_by_name",)
    def __init__(self):
        self._by_name = {}

    def register(self, ent: CollisionEntity):
        if ent.name:
            self._by_name[ent.name] = ent

    def unregister(self, ent: CollisionEntity):
        if ent.name and ent.name in self._by_name:
            del self._by_name[ent.name]

    def get(self, name: str) -> Optional[CollisionEntity]:
        return self._by_name.get(name)

    def set_enabled(self, name: str, enabled: bool):
        e = self.get(name)
        if e is not None:
            e.set_enabled(enabled)

    def toggle_grid(self, name: str, on: Optional[bool] = None):
        e = self.get(name)
        if e is not None:
            e.show_grid = not e.show_grid if on is None else bool(on)

    def all_entities(self) -> List[CollisionEntity]:
        return list(self._by_name.values())

# ---------------------------
# Simple immediate-mode debug render (PyOpenGL)
# ---------------------------
# NOTE: this is debug-only and uses immediate mode (glBegin) for simplicity.
# In production you will want to upload VBOs once and reuse them.
try:
    from OpenGL.GL import glBegin, glEnd, glVertex3fv, GL_LINES, glColor3f, glLineWidth
    _HAS_GL = True
except Exception:
    _HAS_GL = False

def render_debug_grid(entity: CollisionEntity, color=(0.2, 1.0, 0.2), linewidth=1.0):
    """
    Render the grid for one entity if OpenGL is available and entity.show_grid is True.
    """
    if not _HAS_GL:
        return
    if not entity.show_grid:
        return
    lines = entity.generate_grid_lines_world()  # (M,2,3)
    glLineWidth(float(max(1.0, linewidth)))
    glColor3f(float(color[0]), float(color[1]), float(color[2]))
    glBegin(GL_LINES)
    for a,b in lines:
        glVertex3fv(a)
        glVertex3fv(b)
    glEnd()

# ---------------------------
# Example usage
# ---------------------------

if __name__ == "__main__":
    # Example callbacks
    def on_door(entity_self, other):
        print(f"[{entity_self.name}] triggered by [{other.name}] -> door open!")

    def on_wall(entity_self, other):
        print(f"[{entity_self.name}] touched {other.name} -> block movement")

    # Build manager
    mgr = CollisionManager()

    # Player entity (small box)
    player = CollisionEntity(name="player",
                             half_extents=(0.4, 1.0, 0.4),
                             model_matrix=mat_translate(0,0,0),
                             callback=None,
                             trigger_once=False,
                             show_grid=False)
    mgr.register(player)

    # Door entity (will trigger once)
    door_mat = mat_translate(5.0, 0.0, 5.0) @ mat_rotate_euler(0, np.pi*0.2, 0)
    door = CollisionEntity(name="door",
                           half_extents=(1.0, 2.0, 0.3),
                           model_matrix=door_mat,
                           callback=on_door,
                           trigger_once=True,
                           show_grid=True,
                           grid_subdiv=3)
    mgr.register(door)

    # Wall entity with a shear (skew) transform
    wall_mat = mat_translate(10.0, 0.0, 0.0) @ mat_shear(xy=0.4) @ mat_rotate_euler(0, 0.2, 0)
    wall = CollisionEntity(name="wall_skew",
                           half_extents=(0.5, 3.0, 10.0),
                           model_matrix=wall_mat,
                           callback=on_wall,
                           trigger_once=False,
                           show_grid=True,
                           grid_subdiv=4)
    mgr.register(wall)

    # Simulate movement and checks (no GL)
    for step in range(12):
        # place player somewhere (translate with some Z)
        pm = mat_translate(step * 1.0, 0.0, step * 0.5)
        player.set_model_matrix(pm)
        collided_door = check_and_trigger(player, door)
        collided_wall = check_and_trigger(player, wall)
        print(f"step {step}: player at {player.world_corners().mean(axis=0)}, door={collided_door}, wall={collided_wall}")

    # Toggle grid off for door by name
    mgr.toggle_grid("door", on=False)
    # Disable wall collisions by name
    mgr.set_enabled("wall_skew", False)