import pygame
import sys
import math

# Initialize pygame
pygame.init()

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

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

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

# Square properties
square_size = 400
square_center = [WIDTH // 2, HEIGHT // 2]
square_angle = 0
rotation_speed = 0.5  # degrees per frame

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

def rotate_point(point, center, angle_degrees):
    """Rotate a point around a center by a given angle in degrees."""
    angle_radians = math.radians(angle_degrees)

    # Translate point to origin
    x = point[0] - center[0]
    y = point[1] - center[1]

    # Rotate point
    x_rotated = x * math.cos(angle_radians) - y * math.sin(angle_radians)
    y_rotated = x * math.sin(angle_radians) + y * math.cos(angle_radians)

    # Translate point back
    return [x_rotated + center[0], y_rotated + center[1]]

def get_square_vertices(center, size, angle):
    """Get the four vertices of the square after rotation."""
    half_size = size / 2
    vertices = [
        [center[0] - half_size, center[1] - half_size],  # Top-left
        [center[0] + half_size, center[1] - half_size],  # Top-right
        [center[0] + half_size, center[1] + half_size],  # Bottom-right
        [center[0] - half_size, center[1] + half_size]   # Bottom-left
    ]

    # Rotate each vertex
    rotated_vertices = [rotate_point(vertex, center, angle) for vertex in vertices]
    return rotated_vertices

def is_ball_inside_square(ball_pos, vertices):
    """Check if the ball is inside the square using the point-in-polygon algorithm."""
    # Transform to local coordinates (un-rotate)
    local_pos = rotate_point(ball_pos, square_center, -square_angle)

    # In local coordinates, we can just check against axis-aligned square
    half_size = square_size / 2
    return (abs(local_pos[0] - square_center[0]) < half_size - ball_radius and
            abs(local_pos[1] - square_center[1]) < half_size - ball_radius)

def handle_collision(ball_pos, ball_velocity, vertices):
    """Handle collision between ball and square edges."""
    # Transform to local coordinates (un-rotate)
    local_pos = rotate_point(ball_pos, square_center, -square_angle)
    local_vel = rotate_point([ball_velocity[0] + square_center[0],
                             ball_velocity[1] + square_center[1]],
                             square_center, -square_angle)
    local_vel = [local_vel[0] - square_center[0], local_vel[1] - square_center[1]]

    half_size = square_size / 2
    collision = False

    # Check collision with each edge in local coordinates
    if local_pos[0] - ball_radius < square_center[0] - half_size:  # Left edge
        local_vel[0] = abs(local_vel[0])
        local_pos[0] = square_center[0] - half_size + ball_radius
        collision = True
    elif local_pos[0] + ball_radius > square_center[0] + half_size:  # Right edge
        local_vel[0] = -abs(local_vel[0])
        local_pos[0] = square_center[0] + half_size - ball_radius
        collision = True

    if local_pos[1] - ball_radius < square_center[1] - half_size:  # Top edge
        local_vel[1] = abs(local_vel[1])
        local_pos[1] = square_center[1] - half_size + ball_radius
        collision = True
    elif local_pos[1] + ball_radius > square_center[1] + half_size:  # Bottom edge
        local_vel[1] = -abs(local_vel[1])
        local_pos[1] = square_center[1] + half_size - ball_radius
        collision = True

    if collision:
        # Transform back to global coordinates
        new_pos = rotate_point(local_pos, square_center, square_angle)
        new_vel = rotate_point([local_vel[0] + square_center[0],
                               local_vel[1] + square_center[1]],
                               square_center, square_angle)
        new_vel = [new_vel[0] - square_center[0], new_vel[1] - square_center[1]]

        return new_pos, new_vel, True

    return ball_pos, ball_velocity, False

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

    # Clear the screen
    screen.fill(BLACK)

    # Update square rotation
    square_angle += rotation_speed
    if square_angle >= 360:
        square_angle -= 360

    # Get square vertices
    vertices = get_square_vertices(square_center, square_size, square_angle)

    # Update ball position
    ball_pos[0] += ball_velocity[0]
    ball_pos[1] += ball_velocity[1]

    # Handle collision
    ball_pos, ball_velocity, collision = handle_collision(ball_pos, ball_velocity, vertices)

    # Draw the square
    pygame.draw.polygon(screen, BLUE, vertices, 2)

    # 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)

pygame.quit()
sys.exit()