import pygame
import math

# Initialize pygame
pygame.init()

# Screen dimensions
width, height = 800, 600
screen = pygame.display.set_mode((width, height))
pygame.display.set_caption("Spinning Triangle with Bouncing Ball")

# Colors
black = (0, 0, 0)
white = (255, 255, 255)
red = (255, 0, 0)

# Triangle vertices (initial, equilateral triangle centered)
triangle_size = 150
triangle_points = [
    [0, -triangle_size],
    [triangle_size * math.sqrt(3) / 2, triangle_size / 2],
    [-triangle_size * math.sqrt(3) / 2, triangle_size / 2]
]

# Calculate centroid of the triangle
centroid_x = sum(p[0] for p in triangle_points) / 3
centroid_y = sum(p[1] for p in triangle_points) / 3
centroid = [centroid_x, centroid_y]

# Function to rotate a point around the centroid
def rotate_point(point, angle, center):
    """Rotates a point around a center by a given angle in radians."""
    px, py = point
    cx, cy = center
    angle_rad = math.radians(angle)
    rotated_x = math.cos(angle_rad) * (px - cx) - math.sin(angle_rad) * (py - cy) + cx
    rotated_y = math.sin(angle_rad) * (px - cx) + math.cos(angle_rad) * (py - cy) + cy
    return [rotated_x, rotated_y]

# Ball properties
ball_radius = 15
ball_color = red
ball_pos = [0, 0] # Start at centroid initially
ball_speed = 3
ball_velocity = [ball_speed, ball_speed] # Initial velocity

# Rotation speed of the triangle
rotation_speed = 0.5 # degrees per frame
rotation_angle = 0

def is_point_in_triangle(point, triangle_vertices):
    """Checks if a point is inside a triangle using barycentric coordinates."""
    x, y = point
    x1, y1 = triangle_vertices[0]
    x2, y2 = triangle_vertices[1]
    x3, y3 = triangle_vertices[2]

    denominator = (y2 - y3) * (x1 - x3) + (x3 - x2) * (y1 - y3)
    if denominator == 0: return False  # Triangle is degenerate

    a = ((y2 - y3) * (x - x3) + (x3 - x2) * (y - y3)) / denominator
    b = ((y3 - y1) * (x - x3) + (x1 - x3) * (y - y3)) / denominator
    c = 1 - a - b

    return 0 <= a <= 1 and 0 <= b <= 1 and 0 <= c <= 1

def distance_point_to_segment(point, segment_p1, segment_p2):
    """Calculates the shortest distance from a point to a line segment."""
    px, py = point
    x1, y1 = segment_p1
    x2, y2 = segment_p2

    l2 = (x1 - x2)**2 + (y1 - y2)**2
    if l2 == 0: return math.sqrt((px - x1)**2 + (py - y1)**2)   # p1 == p2 case

    t = ((px - x1) * (x2 - x1) + (py - y1) * (y2 - y1)) / l2
    t = max(0, min(1, t)) # Clamp t to be within the segment [0, 1]

    closest_x = x1 + t * (x2 - x1)
    closest_y = y1 + t * (y2 - y1)

    return math.sqrt((px - closest_x)**2 + (py - closest_y)**2), [closest_x, closest_y]


def reflect_velocity(ball_pos, ball_velocity, segment_p1, segment_p2):
    """Reflects the ball's velocity off a line segment."""
    x1, y1 = segment_p1
    x2, y2 = segment_p2

    # Vector representing the segment
    segment_vector_x = x2 - x1
    segment_vector_y = y2 - y1

    # Normal vector to the segment (normalized)
    normal_x = -segment_vector_y
    normal_y = segment_vector_x
    norm = math.sqrt(normal_x**2 + normal_y**2)
    normal_x /= norm
    normal_y /= norm

    # Dot product of velocity and normal
    dot_product = ball_velocity[0] * normal_x + ball_velocity[1] * normal_y

    # Reflection formula: v_reflected = v - 2 * (v . n) * n
    reflected_velocity_x = ball_velocity[0] - 2 * dot_product * normal_x
    reflected_velocity_y = ball_velocity[1] - 2 * dot_product * normal_y

    return [reflected_velocity_x, reflected_velocity_y]


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

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

    # Rotate the triangle
    rotated_triangle_points = [rotate_point(p, rotation_angle, centroid) for p in triangle_points]
    rotation_angle += rotation_speed

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

    # Collision detection with triangle edges
    collided = False
    for i in range(3):
        p1 = rotated_triangle_points[i]
        p2 = rotated_triangle_points[(i + 1) % 3]
        distance, closest_point = distance_point_to_segment(ball_pos, p1, p2)

        if distance <= ball_radius:
            ball_velocity = reflect_velocity(ball_pos, ball_velocity, p1, p2)
            collided = True
            # Slightly move ball out of collision to prevent sticking
            push_out_vector_x = ball_pos[0] - closest_point[0]
            push_out_vector_y = ball_pos[1] - closest_point[1]
            push_out_norm = max(distance, 1e-9) # Avoid division by zero if distance is zero
            ball_pos[0] += push_out_vector_x * (ball_radius - distance + 0.1) / push_out_norm
            ball_pos[1] += push_out_vector_y * (ball_radius - distance + 0.1) / push_out_norm
            break # Only reflect off one edge at a time per frame for simplicity


    # Keep ball inside triangle (as a fallback, although collision should handle this)
    if not is_point_in_triangle(ball_pos, rotated_triangle_points):
        # If somehow ball gets outside, reposition it to centroid and reverse velocity (basic fix)
        ball_pos = [0, 0] # Reset to centroid - you could do something smarter
        ball_velocity[0] *= -1
        ball_velocity[1] *= -1


    # Clear the screen
    screen.fill(black)

    # Draw the spinning triangle
    triangle_screen_points = [
        (int(p[0] + width / 2), int(p[1] + height / 2)) for p in rotated_triangle_points
    ]
    pygame.draw.polygon(screen, white, triangle_screen_points, 2) # Draw triangle outline

    # Draw the bouncing ball
    pygame.draw.circle(screen, ball_color, (int(ball_pos[0] + width / 2), int(ball_pos[1] + height / 2)), ball_radius)

    # Update the display
    pygame.display.flip()

    # Control frame rate
    clock.tick(60)

pygame.quit()