geom-demo/geom/demos/sat.py

import pygame
from ..geom import Polygon, generate_polygon
from .. import WIDTH, HEIGHT, get_cursor_pos, get_inputs
from ..colors import *
from ..math import *
from ..draw import *


def compute_separating_axes(p1, p2):
    """Compute the projections upon the separating axes of polygons p1 and p2.

    For efficient collision detection, this could early-exit as soon as one
    axis is found where the projections of p1 and p2 do not overlap. For
    visualization purposes, we compute the projections for all axes and
    return the results."""

    # The list of axes to project upon
    axes = []

    # Iterate the normals of both polygons
    for normal in p1.get_normals() + p2.get_normals():
        # we can trim out normals that are in the same or opposite direction
        # as axes we already have, since they will give us the same overlap
        # information
        if not any(is_colinear(*normal, *axis) for axis in axes):
            axes.append(normal)

    # helper function to find the min/max projections of vertices onto
    # the given axis.

    def get_min_max(axis, vertices):
        projections = [dot(*axis, *vertex) for vertex in vertices]
        return min(projections), max(projections)

    return [
        (axis,
         get_min_max(axis, p1.get_translated_vertices()),
         get_min_max(axis, p2.get_translated_vertices()))
        for axis in axes
    ]


class SeparatingAxisTheorem:

    title = "Separating Axis Theorem"

    def __init__(self):
        self.shape1 = generate_polygon(0, 0, sides=4, radius=60)

        self.shape2 = None
        self.shape2_sides = 4
        self.shape2_radius = 60
        self.generate_shape2()

    def generate_shape2(self):
        """Generates the shape we can move around - a regular polygon."""
        rotation = 45 if self.shape2 is None else self.shape2.rotation
        self.shape2 = generate_polygon(*get_cursor_pos(),
                                       rotation,
                                       sides=self.shape2_sides,
                                       radius=self.shape2_radius)

    def handle_key_down(self, key):
        if key == pygame.K_a:
            # remove a side from the polygon
            self.shape2_sides = max(self.shape2_sides - 1, 3)
            self.generate_shape2()
        if key == pygame.K_d:
            # add a side to the polygon
            self.shape2_sides = min(self.shape2_sides + 1, 31)
            self.generate_shape2()

    def handle_input(self):
        self.shape2.set_position(*get_cursor_pos())

        inputs = get_inputs()
        if inputs[pygame.K_q]:
            # rotate counter-clockwise by 90 degrees every second
            self.shape2.rotate((1/60) * -90)
        if inputs[pygame.K_e]:
            # same as above, but clockwise
            self.shape2.rotate((1/60) * 90)

    def render(self, surface):
        self.handle_input()

        gaps = []

        for result in compute_separating_axes(self.shape1, self.shape2):

            axis, (min1, max1), (min2, max2) = result

            # length of the gap between the two projections
            gap = max(min2 - max1, min1 - max2)
            gaps.append(gap)

            self.draw_axis(surface, axis, gap, min1, max1, min2, max2)

        colliding = all(gap < 0 for gap in gaps)
        if colliding:
            text_screen(surface, WHITE, (10, HEIGHT - 20),
                        f"Min Distance to Resolve: {max(gaps)}")

        self.shape1.draw(surface, YELLOW, 1)
        self.shape2.draw(surface, RED if colliding else WHITE, 1)

    def draw_axis(self, surface, normal, gap,  min1, max1, min2, max2):
        """Draws a separating axis with shape projections, and separating line.
        Uses colours to clearly indicate whether the projections overlap."""

        # since the normal's origin is 0,0 (center of the screen), calcuate
        # an offset to move all drawing off to one side - in this case 300px.
        offset = scale(*rnormal(*normal), 300)

        overlapping = gap < 0

        if not overlapping:
            # how far from the center of the gap is from the origin
            dist = scale(*normal, (min1 if min1 > min2 else min2) - 0.5 * gap)

            # the separating line
            line(surface, CYAN,
                 add(*dist, *offset),
                 add(*dist, *scale(*lnormal(*normal), 1000)))

        # axis, extending off the screen in both directions
        line(surface, WHITE,
             add(*offset, *scale(*reverse(*normal), 1000)),
             add(*offset, *scale(*normal, 1000)))

        # projection 1
        line(surface, RED if overlapping else GREEN,
             add(*offset, *scale(*normal, min1)),
             add(*offset, *scale(*normal, max1)), 8)

        # projection 2
        line(surface, RED if overlapping else GREEN,
             add(*offset, *scale(*normal, min2)),
             add(*offset, *scale(*normal, max2)), 8)