experiment with understanding collisions of capsules

capsule.py

from dataclasses import dataclass
from pygame import Vector2
import pygame
import math


@dataclass
class Circle:
    center: Vector2
    radius: float


@dataclass
class Capsule:
    center: Vector2
    radius: float
    half_length: float
    angle: float = 0

    @property
    def length(self) -> float:
        return self.half_length * 2

    @length.setter
    def length(self, value: float) -> None:
        self.half_length = value * 0.5


def get_collision(cap: Capsule, circ: Circle) -> tuple[Vector2, float] | None:
    x_axis = Vector2(
        math.cos(cap.angle), math.sin(cap.angle)
    )
    y_axis = Vector2(
        -math.sin(cap.angle), math.cos(cap.angle)
    )

    # get projected x/y coords, but relative to the capsules center.
    rel = circ.center - cap.center
    circ_x = x_axis.dot(rel)
    circ_y = y_axis.dot(rel)
    proj_circ = Vector2(circ_x, circ_y)

    # with a unit vector, the dot product gives us the circles projected x-coord, on the x_axis.
    # capsule half-length goes along that same axis, so we clamp it to it's half-length range,
    # to figure out where on the "center line" of the capsule the circle is closest to.
    x_coord = pygame.math.clamp(
        circ_x,
        -cap.half_length,
        cap.half_length
    )

    line_point = Vector2(x_coord, 0)
    delta_vec = proj_circ - line_point
    distance_squared = delta_vec.length_squared()
    sum_radius = cap.radius + circ.radius
    if distance_squared > sum_radius ** 2:
        return None

    distance = math.sqrt(distance_squared)  # distance between centers.
    overlap = sum_radius - distance

    if distance == 0.0:
        # arbitrary but stable
        world_normal = y_axis
    else:
        # translate normal back to world space, and return with the overlap
        local_normal = delta_vec / distance
        world_normal = (x_axis * local_normal.x) + (y_axis * local_normal.y)
    return world_normal, overlap


def test(cap: Capsule, circ: Circle):
    print('--')
    print(f'For:\n  {cap}\n  {circ}')
    if val := get_collision(cap, circ):
        normal, overlap = val
        print(f'before: {circ.center}')
        circ.center += normal * overlap
        print(f'after: {circ.center}')
        print(val)
        print(get_collision(cap, circ))
    else:
        print(f'No overlap')


def main():
    print()

    test(
        Capsule(Vector2(0, 0), 5, 10, math.radians(-45)),
        Circle(Vector2(0, 10), 4)
    )

    test(
        Capsule(Vector2(0, 0), 5, 10, math.radians(0)),
        Circle(Vector2(18, 0), 4)
    )

    test(
        Capsule(Vector2(0, 0), 5, 10, math.radians(0)),
        Circle(Vector2(18, -2), 4)
    )

    test(
        Capsule(Vector2(0, 0), 5, 10, math.radians(0)),
        Circle(Vector2(0, 0), 4)
    )


if __name__ == '__main__':
    main()