import pygame
import math
from pygame.locals import *

# Initialize Pygame
pygame.init()

# Set up display
width, height = 800, 800
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)

# Square parameters
square_center = (width//2, height//2)
square_size = int(min(width, height) * 0.7071)  # Max size to fit rotated square (≈800/√2)
half_square = square_size // 2
rotation_angle = 0  # in radians
rotation_speed = 0.01  # radians per frame

# Ball parameters
ball_radius = 20
ball_pos = [square_center[0], square_center[1]]
ball_vel = [3, 3]  # Initial velocity

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

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

    # Update rotation angle
    rotation_angle += rotation_speed

    # Move the ball
    ball_pos[0] += ball_vel[0]
    ball_pos[1] += ball_vel[1]

    # Collision detection and response
    dx = ball_pos[0] - square_center[0]
    dy = ball_pos[1] - square_center[1]

    cos_theta = math.cos(rotation_angle)
    sin_theta = math.sin(rotation_angle)

    # Transform to local coordinates
    local_x = dx * cos_theta + dy * sin_theta
    local_y = -dx * sin_theta + dy * cos_theta

    # Adjusted half size considering ball radius
    adjusted_half = half_square - ball_radius
    collision = False

    # X-axis collision check
    if local_x > adjusted_half:
        local_x = adjusted_half
        collision = True
    elif local_x < -adjusted_half:
        local_x = -adjusted_half
        collision = True

    # Y-axis collision check
    if local_y > adjusted_half:
        local_y = adjusted_half
        collision = True
    elif local_y < -adjusted_half:
        local_y = -adjusted_half
        collision = True

    if collision:
        # Transform velocity to local coordinates
        vx_local = ball_vel[0] * cos_theta + ball_vel[1] * sin_theta
        vy_local = -ball_vel[0] * sin_theta + ball_vel[1] * cos_theta

        # Reverse velocity components if collision occurred
        if abs(local_x) == adjusted_half:
            vx_local *= -1
        if abs(local_y) == adjusted_half:
            vy_local *= -1

        # Transform velocity back to global coordinates
        ball_vel[0] = vx_local * cos_theta - vy_local * sin_theta
        ball_vel[1] = vx_local * sin_theta + vy_local * cos_theta

        # Update position to stay within bounds
        dx_new = local_x * cos_theta - local_y * sin_theta
        dy_new = local_x * sin_theta + local_y * cos_theta
        ball_pos[0] = square_center[0] + dx_new
        ball_pos[1] = square_center[1] + dy_new

    # Drawing
    screen.fill(BLACK)

    # Draw rotating square
    half = square_size // 2
    corners = [
        (-half, -half),
        (-half, half),
        (half, half),
        (half, -half)
    ]
    rotated_corners = []
    for x, y in corners:
        rot_x = x * cos_theta - y * sin_theta
        rot_y = x * sin_theta + y * cos_theta
        rotated_corners.append((square_center[0] + rot_x, square_center[1] + rot_y))
    pygame.draw.polygon(screen, WHITE, rotated_corners, 2)

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

    pygame.display.flip()
    clock.tick(60)

pygame.quit()