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 with Speed Controls")

# 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) / math.sqrt(2))  # Critical calculation , patch #1: adjust square size to random window size
half_square = square_size // 2
rotation_angle = 0
rotation_speed = 0.01

# Ball parameters
ball_radius = 20
ball_pos = [square_center[0], square_center[1]]
ball_vel = [3.0, 3.0]  # Use floats for smoother acceleration, ball velocity/speed

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

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

    # Get pressed keys state
    keys = pygame.key.get_pressed()

    # Handle square rotation controls,    #patch #2
    if keys[K_LEFT]:
        rotation_speed = -abs(rotation_speed)       #Counter clock wise rotation
    if keys[K_RIGHT]:
        rotation_speed = abs(rotation_speed)        #Clock wise rotation
    if keys[K_UP]:
        rotation_speed *= 1.02  # Smoother acceleration
    if keys[K_DOWN]:
        rotation_speed *= 0.98  # Smoother deceleration

    # Handle ball speed controls
    if keys[K_a] or keys[K_s]:
        current_speed = math.hypot(ball_vel[0], ball_vel[1])
        if current_speed > 0:
            direction = (ball_vel[0] / current_speed, ball_vel[1] / current_speed)
        else:
            direction = (1, 0)

        if keys[K_a]:  # Accelerate
            current_speed += 0.5  # More noticeable per-frame change
        if keys[K_s]:  # Decelerate
            current_speed = max(current_speed - 0.5, 0.5)  # Minimum speed

        ball_vel[0] = direction[0] * current_speed
        ball_vel[1] = direction[1] * current_speed

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