import pygame
import math
import sys

# Initialize Pygame
pygame.init()

# Screen dimensions
WIDTH, HEIGHT = 600, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Bouncing Ball in Rotating Square")

# Colors
YELLOW = (255, 255, 0)
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)

# Ball properties
ball_radius = 20
ball_pos = [WIDTH // 2, HEIGHT // 2]
ball_velocity = [5, 5]

# Square properties
square_size = 400
square_angle = 0
square_rotation_speed = 1  # Degrees per frame

# Clock for controlling frame rate
clock = pygame.time.Clock()

def rotate_point(point, angle, center):
    """Rotate a point around a center by a given angle."""
    angle_rad = math.radians(angle)
    x = point[0] - center[0]
    y = point[1] - center[1]
    new_x = x * math.cos(angle_rad) - y * math.sin(angle_rad)
    new_y = x * math.sin(angle_rad) + y * math.cos(angle_rad)
    return [new_x + center[0], new_y + center[1]]

def get_rotated_square(center, size, angle):
    """Get the vertices of a rotated square."""
    half_size = size // 2
    corners = [
        [center[0] - half_size, center[1] - half_size],
        [center[0] + half_size, center[1] - half_size],
        [center[0] + half_size, center[1] + half_size],
        [center[0] - half_size, center[1] + half_size],
    ]
    return [rotate_point(corner, angle, center) for corner in corners]

def is_point_in_rotated_square(point, square_corners):
    """Check if a point is inside a rotated square."""
    def cross_product(a, b, c):
        return (b[0] - a[0]) * (c[1] - a[1]) - (b[1] - a[1]) * (c[0] - a[0])

    d1 = cross_product(square_corners[0], square_corners[1], point)
    d2 = cross_product(square_corners[1], square_corners[2], point)
    d3 = cross_product(square_corners[2], square_corners[3], point)
    d4 = cross_product(square_corners[3], square_corners[0], point)

    return (d1 >= 0 and d2 >= 0 and d3 >= 0 and d4 >= 0) or (d1 <= 0 and d2 <= 0 and d3 <= 0 and d4 <= 0)

def bounce_ball(ball_pos, ball_velocity, square_corners):
    """Handle ball bouncing off the edges of the square."""
    if not is_point_in_rotated_square([ball_pos[0] + ball_velocity[0], ball_pos[1]], square_corners):
        ball_velocity[0] *= -1
    if not is_point_in_rotated_square([ball_pos[0], ball_pos[1] + ball_velocity[1]], square_corners):
        ball_velocity[1] *= -1

# Main loop
running = True
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    # Clear the screen
    screen.fill(BLACK)

    # Rotate the square
    square_angle = (square_angle + square_rotation_speed) % 360
    square_corners = get_rotated_square([WIDTH // 2, HEIGHT // 2], square_size, square_angle)

    # Draw the square
    pygame.draw.polygon(screen, WHITE, square_corners, 2)

    # Update ball position and handle bouncing
    ball_pos[0] += ball_velocity[0]
    ball_pos[1] += ball_velocity[1]
    bounce_ball(ball_pos, ball_velocity, square_corners)

    # Draw the ball
    pygame.draw.circle(screen, YELLOW, (int(ball_pos[0]), int(ball_pos[1])), ball_radius)

    # Update the display
    pygame.display.flip()

    # Cap the frame rate
    clock.tick(60)

# Quit Pygame
pygame.quit()
sys.exit()