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 *


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):
        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:
            self.shape2_sides = max(self.shape2_sides - 1, 3)
            self.generate_shape2()
        if key == pygame.K_d:
            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]:
            self.shape2.rotate((1/60) * -90)
        if inputs[pygame.K_e]:
            self.shape2.rotate((1/60) * 90)

    def draw_axis(self, surface, normal, overlap, min1, max1, min2, max2):
        offset = scale(*rnormal(*normal), 300)

        colliding = overlap < 0

        if not colliding:
            # separating line
            dist = scale(*normal,
                         (min1 if min1 > min2 else min2) - overlap * 0.5)
            line(surface, BLUE,
                 add(*dist, *offset),
                 add(*dist, *scale(*lnormal(*normal), 1000)))

        # axis
        line(surface, WHITE,
             add(*scale(*reverse(*normal), 1000), *offset),
             add(*scale(*normal, 1000), *offset))

        # shape 1
        line(surface, GREEN if not colliding else RED,
             add(*scale(*normal, min1), *offset),
             add(*scale(*normal, max1), *offset), 8)
        # shape 2
        line(surface, GREEN if not colliding else RED,
             add(*scale(*normal, min2), *offset),
             add(*scale(*normal, max2), *offset), 8)

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

        normals = []

        def add_if_not_exists(normal):
            for existing in normals:
                if is_equal(*existing, *normal):
                    return
            normals.append(normal)

        for normal in self.shape1.get_normals() + self.shape2.get_normals():
            add_if_not_exists(normal)

        def get_min_max(normal, vertices):
            _min = float('inf')
            _max = float('-inf')
            for vertex in vertices:
                proj = dot(*normal, *vertex)
                _min = min(_min, proj)
                _max = max(_max, proj)
            return _min, _max

        overlaps = []
        for normal in normals:
            min1, max1 = get_min_max(normal,
                                     self.shape1.get_translated_vertices())
            min2, max2 = get_min_max(normal,
                                     self.shape2.get_translated_vertices())

            overlap = max(min2 - max1, min1 - max2)
            overlaps.append(overlap)

            self.draw_axis(surface, normal, overlap,
                           min1, max1, min2, max2)

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

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