o1 pro ball test

test.py

import pygame
import sys
import math

# ---------------------------------------------------
# Helper functions
# ---------------------------------------------------
def rotate_point(px, py, angle):
    """
    Rotate point (px, py) around the origin (0,0) by 'angle' radians.
    Returns the new coordinates (rx, ry).
    """
    cos_a = math.cos(angle)
    sin_a = math.sin(angle)
    rx = px * cos_a - py * sin_a
    ry = px * sin_a + py * cos_a
    return rx, ry

def closest_point_on_segment(px, py, ax, ay, bx, by):
    """
    Returns the closest point (cx, cy) on the line segment A->B to point P.
    (px, py) is P, (ax, ay) is A, (bx, by) is B.
    """
    # Segment AB vector
    abx = bx - ax
    aby = by - ay
    # Vector AP
    apx = px - ax
    apy = py - ay
    
    # Length squared of AB
    ab_len_sq = abx * abx + aby * aby
    if ab_len_sq == 0:
        # A and B are the same point
        return ax, ay
    
    # Projection of AP onto AB, in normalized "t" form
    t = (apx * abx + apy * aby) / ab_len_sq
    # Clamp t to [0,1] so it stays within the segment
    t = max(0, min(1, t))
    
    # Find the projection point
    cx = ax + t * abx
    cy = ay + t * aby
    return cx, cy

# ---------------------------------------------------
# Main script
# ---------------------------------------------------
def main():
    pygame.init()
    WIDTH, HEIGHT = 800, 600
    screen = pygame.display.set_mode((WIDTH, HEIGHT))
    clock = pygame.time.Clock()
    pygame.display.set_caption("Rotating Square with Bouncing Ball")
    
    # Define the square in local coordinates (centered at origin):
    # Clockwise or counter-clockwise, just be consistent
    # Let's define corners in clockwise order.
    square_side = 300
    half_side = square_side / 2
    square_local_points = [
        (-half_side, -half_side),
        ( half_side, -half_side),
        ( half_side,  half_side),
        (-half_side,  half_side),
    ]
    
    # Ball properties
    ball_radius = 20
    ball_x = WIDTH // 2
    ball_y = HEIGHT // 2
    ball_vx = 3.0   # velocity X
    ball_vy = 2.0   # velocity Y
    
    # Rotation angle for the square
    angle = 0.0
    rotation_speed = 0.01  # radians per frame, ~0.57 deg per frame
    
    # Main loop
    running = True
    while running:
        clock.tick(60)  # ~60 FPS
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
        
        # Update ball position
        ball_x += ball_vx
        ball_y += ball_vy
        
        # Rotate square
        angle += rotation_speed
        
        # Compute square corners in world coordinates (screen coordinates)
        # The center of the square is at (WIDTH/2, HEIGHT/2)
        square_world_points = []
        for (lx, ly) in square_local_points:
            rx, ry = rotate_point(lx, ly, angle)
            # Translate to screen center
            rx += WIDTH / 2
            ry += HEIGHT / 2
            square_world_points.append((rx, ry))
        
        # Perform collision detection for each edge in the square
        # The edges are pairs of consecutive points plus the last->first
        for i in range(len(square_world_points)):
            j = (i + 1) % len(square_world_points)
            ax, ay = square_world_points[i]
            bx, by = square_world_points[j]
            
            # Find the closest point on segment AB to the ball center
            cx, cy = closest_point_on_segment(ball_x, ball_y, ax, ay, bx, by)
            
            # Distance from that closest point to ball center
            dx = ball_x - cx
            dy = ball_y - cy
            dist_sq = dx*dx + dy*dy
            if dist_sq < ball_radius * ball_radius:
                dist = math.sqrt(dist_sq)
                if dist == 0:
                    # If the ball center is exactly on the segment,
                    # push it out in some direction; 
                    # normal direction can be from A->center for instance
                    # or just skip to avoid division by zero.
                    continue
                
                # Normal from the closest point to the ball center
                nx = dx / dist
                ny = dy / dist
                # Reflect the velocity
                # v' = v - 2(v . n) n
                dot = ball_vx * nx + ball_vy * ny
                ball_vx = ball_vx - 2 * dot * nx
                ball_vy = ball_vy - 2 * dot * ny
                
                # Push the ball so it's exactly ball_radius away
                overlap = ball_radius - dist
                ball_x += nx * overlap
                ball_y += ny * overlap
        
        # Clear screen
        screen.fill((30, 30, 30))
        
        # Draw rotating square (white)
        pygame.draw.polygon(screen, (255, 255, 255), square_world_points, width=2)
        
        # Draw ball (yellow)
        pygame.draw.circle(screen, (255, 255, 0), (int(ball_x), int(ball_y)), ball_radius)
        
        # Flip display
        pygame.display.flip()
    
    pygame.quit()
    sys.exit()

if __name__ == "__main__":
    main()