import pygame
import math

# Initialize Pygame
pygame.init()

# Screen dimensions
screen_width = 800
screen_height = 600
screen = pygame.display.set_mode((screen_width, screen_height))
pygame.display.set_caption("Bouncing Yellow Ball in Rotating Square")

# Colors
black = (0, 0, 0)
white = (255, 255, 255)
yellow = (255, 255, 0)
light_gray = (200, 200, 200)

# Ball properties
ball_radius = 20
ball_x = screen_width // 2
ball_y = screen_height // 2
ball_speed_x = 5
ball_speed_y = 3

# Square properties
square_size = 300
square_center_x = screen_width // 2
square_center_y = screen_height // 2
square_rotation_angle = 0
square_rotation_speed = 0.5  # Degrees per frame

# Function to rotate a point around a center
def rotate_point(point, center, angle_degrees):
    angle_radians = math.radians(angle_degrees)
    cos_theta = math.cos(angle_radians)
    sin_theta = math.sin(angle_radians)
    x = point[0] - center[0]
    y = point[1] - center[1]
    rotated_x = x * cos_theta - y * sin_theta
    rotated_y = x * sin_theta + y * cos_theta
    return [rotated_x + center[0], rotated_y + center[1]]

# Function to get square vertices based on rotation
def get_square_vertices(center_x, center_y, size, angle):
    half_size = size / 2
    vertices = [
        [-half_size, -half_size],
        [half_size, -half_size],
        [half_size, half_size],
        [-half_size, half_size]
    ]
    rotated_vertices = []
    for vertex in vertices:
        rotated_vertices.append(rotate_point([vertex[0] + center_x, vertex[1] + center_y], [center_x, center_y], angle))
    return rotated_vertices

# Function to check collision with a square edge and bounce (more robust and no sticking)
def check_square_collision_and_bounce(ball_x, ball_y, ball_radius, ball_speed_x, ball_speed_y, vertices):
    next_ball_x = ball_x + ball_speed_x
    next_ball_y = ball_y + ball_speed_y
    new_ball_speed_x = ball_speed_x
    new_ball_speed_y = ball_speed_y

    for i in range(4):
        p1 = vertices[i]
        p2 = vertices[(i + 1) % 4]

        dx = p2[0] - p1[0]
        dy = p2[1] - p1[1]

        # Vector from p1 to ball center
        ball_vec_x = next_ball_x - p1[0]
        ball_vec_y = next_ball_y - p1[1]

        # Project ball vector onto edge normal (normalized perpendicular vector)
        edge_length_sq = dx*dx + dy*dy
        if edge_length_sq == 0: # Handle degenerate edge
            continue

        normal_x = -dy
        normal_y = dx
        normal_length = math.sqrt(normal_x*normal_x + normal_y*normal_y)
        normal_x /= normal_length
        normal_y /= normal_length

        projection = (ball_vec_x * normal_x + ball_vec_y * normal_y)

        # Closest point on the edge to the ball center
        closest_x = next_ball_x - projection * normal_x
        closest_y = next_ball_y - projection * normal_y

        # Check if closest point is on the segment
        on_segment = True
        dot_product1 = (closest_x - p1[0]) * dx + (closest_y - p1[1]) * dy
        dot_product2 = (closest_x - p2[0]) * dx + (closest_y - p2[1]) * dy
        if dot_product1 < 0 or dot_product2 > 0:
            on_segment = False

        if on_segment:
            distance_sq = (next_ball_x - closest_x)**2 + (next_ball_y - closest_y)**2
            if distance_sq <= ball_radius**2: # Collision!
                if projection < 0: # Ball moving towards the edge
                    dot_product_vel = ball_speed_x * normal_x + ball_speed_y * normal_y
                    new_ball_speed_x -= 2 * dot_product_vel * normal_x
                    new_ball_speed_y -= 2 * dot_product_vel * normal_y
                    # Slightly push ball away to prevent sticking
                    next_ball_x += normal_x * 0.1
                    next_ball_y += normal_y * 0.1
                    return next_ball_x - ball_speed_x, next_ball_y - ball_speed_y, new_ball_speed_x, new_ball_speed_y

    return next_ball_x - ball_speed_x, next_ball_y - ball_speed_y, new_ball_speed_x, new_ball_speed_y


# Game loop
running = True
clock = pygame.time.Clock()

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

    # Update square rotation
    square_rotation_angle += square_rotation_speed

    # Get rotated square vertices
    square_vertices = get_square_vertices(square_center_x, square_center_y, square_size, square_rotation_angle)

    # Ball movement and collision detection
    ball_x, ball_y, ball_speed_x, ball_speed_y = check_square_collision_and_bounce(ball_x, ball_y, ball_radius, ball_speed_x, ball_speed_y, square_vertices)

    ball_x += ball_speed_x
    ball_y += ball_speed_y


    # Keep ball within screen bounds (fallback, but square collision is primary)
    if ball_x - ball_radius < 0:
        ball_x = ball_radius
        ball_speed_x *= -1
    if ball_x + ball_radius > screen_width:
        ball_x = screen_width - ball_radius
        ball_speed_x *= -1
    if ball_y - ball_radius < 0:
        ball_y = ball_radius
        ball_speed_y *= -1
    if ball_y + ball_radius > screen_height:
        ball_y = screen_height - ball_radius
        ball_speed_y *= -1


    # Drawing
    screen.fill(black)

    # Draw rotating square
    pygame.draw.polygon(screen, light_gray, square_vertices, 2) # 2 is line thickness

    # Draw ball
    pygame.draw.circle(screen, yellow, (int(ball_x), int(ball_y)), ball_radius)

    pygame.display.flip()

    clock.tick(60) # Limit frame rate to 60 FPS

pygame.quit()