def get_cue_ball_path(self):
"""
Sets the cue stick line endpoint.
Will either be at a ball or a cushion.
"""
# If cue ball is currently pocketed, skip
if self.cue_ball is None:
return
# Reset lines
self.object_deflect_line_start = self.object_deflect_line_end = self.cue_deflect_line_end = None
angle = self.cue_angle
nw = Coordinates(self.left, self.top)
se = Coordinates(self.right, self.bottom)
cue_mid_start = self.cue_ball.pos # Line start is cue ball position
cue_mid_end = self.cue_line_end = get_line_endpoint_within_box(
cue_mid_start, angle, nw, se, self.cue_ball.radius)
cue_top_start, cue_top_end = get_parallel_line(cue_mid_start,
cue_mid_end,
self.cue_ball.radius,
True)
cue_bot_start, cue_bot_end = get_parallel_line(cue_mid_start,
cue_mid_end,
self.cue_ball.radius,
False)
# Ghost ball computation
balls_by_distance = list(self.balls.values())
balls_by_distance.sort(
key=lambda b: get_distance(cue_mid_start, b.pos))
for ball in balls_by_distance:
if ball.ball_type is BallType.CUE: continue # Skip the cue ball
if (check_ray_circle_intersection(cue_top_start, cue_top_end,
ball.pos, ball.radius)
or check_ray_circle_intersection(
cue_bot_start, cue_bot_end, ball.pos, ball.radius)):
print("CUE BALL INTERSECTING {}".format(ball))
self.cue_line_end = get_point_on_line_distance_from_point(
cue_mid_start, cue_mid_end, ball.pos, 2 * ball.radius)
# Set object ball deflection line
self.object_deflect_line_start = ball.pos
object_ball_angle = get_angle(ball.pos, self.cue_line_end)
self.object_deflect_line_end = get_line_endpoint_within_box(
ball.pos, object_ball_angle, nw, se, self.cue_ball.radius)
# Set cue ball deflection line
cue_deflect_angle = get_angle(self.object_deflect_line_end,
self.object_deflect_line_start)
cue_object_angle = get_angle(ball.pos, self.cue_ball.pos)
print('self.cue_angle: {}'.format(self.cue_angle))
print('self.cue_angle: {}'.format(cue_object_angle))
if self.cue_angle % 360 == 0:
# Edge case when perfectly to the right
cue_deflect_angle = (cue_deflect_angle + 90) % 360
elif self.cue_angle < cue_object_angle:
print("CUE BALL GOING RIGHT OF OBJECT BALl")
cue_deflect_angle = (cue_deflect_angle - 90) % 360
else:
print("CUE BALL GOING LEFT OF OBJECT BALl")
cue_deflect_angle = (cue_deflect_angle + 90) % 360
self.cue_deflect_line_end = get_line_endpoint_within_box(
self.cue_line_end, cue_deflect_angle, nw, se,
self.cue_ball.radius)
return
def get_angle(self):
"""
Get the angle for this vector.
"""
return util.get_angle(Coordinates(self.x, self.y))
def test_get_parallel_line(self):
# def get_parallel_line(p1: Coordinates, p2: Coordinates, dist: float, top: bool) -> (Coordinates, Coordinates):
######################
# Parallel to x-axis #
######################
p1 = Coordinates(-1, 0)
p2 = Coordinates(1, 0)
dist = 1.0
expected_angle = get_angle(p2, p1)
result_top = get_parallel_line(p1, p2, dist, True)
result_bot = get_parallel_line(p1, p2, dist, False)
result_top_angle = get_angle(result_top[1], result_top[0])
result_bot_angle = get_angle(result_bot[1], result_bot[0])
# Check angles are the same
self.assertEqual(expected_angle, result_top_angle)
self.assertEqual(expected_angle, result_bot_angle)
# Check distance apart
self.assertAlmostEqual(dist, get_distance(p1, result_top[0]),
FLOAT_PLACES)
self.assertAlmostEqual(dist, get_distance(p2, result_top[1]),
FLOAT_PLACES)
######################
# Parallel to y-axis #
######################
p1 = Coordinates(0, -1)
p2 = Coordinates(0, 1)
dist = 1.0
expected_angle = get_angle(p2, p1)
result_top = get_parallel_line(p1, p2, dist, True)
result_bot = get_parallel_line(p1, p2, dist, False)
result_top_angle = get_angle(result_top[1], result_top[0])
result_bot_angle = get_angle(result_bot[1], result_bot[0])
# Check angles are the same
self.assertEqual(expected_angle, result_top_angle)
self.assertEqual(expected_angle, result_bot_angle)
# Check distance apart
self.assertAlmostEqual(dist, get_distance(p1, result_top[0]),
FLOAT_PLACES)
self.assertAlmostEqual(dist, get_distance(p2, result_top[1]),
FLOAT_PLACES)
#####################
# Parallel to y = x #
#####################
p1 = Coordinates(-1, -1)
p2 = Coordinates(1, 1)
dist = 1.0
expected_angle = get_angle(p2, p1)
result_top = get_parallel_line(p1, p2, dist, True)
result_bot = get_parallel_line(p1, p2, dist, False)
result_top_angle = get_angle(result_top[1], result_top[0])
result_bot_angle = get_angle(result_bot[1], result_bot[0])
# Check angles are the same
self.assertEqual(expected_angle, result_top_angle)
self.assertEqual(expected_angle, result_bot_angle)
# Check distance apart
self.assertAlmostEqual(dist, get_distance(p1, result_top[0]),
FLOAT_PLACES)
self.assertAlmostEqual(dist, get_distance(p2, result_top[1]),
FLOAT_PLACES)
def main2():
init()
# Create pool table
nw = coords_from_pygame((TABLE_OFFSET_X, TABLE_OFFSET_Y), HEIGHT)
se = coords_from_pygame(
(TABLE_OFFSET_X + TABLE_LENGTH, TABLE_OFFSET_Y + TABLE_LENGTH / 2),
HEIGHT)
table = PoolTable(nw, se)
while 1:
# Get just the list of balls to iterate easily
balls = list(table.balls.values())
clear_screen()
# Check Pygame events
for event in pygame.event.get():
if event.type == pygame.QUIT:
sys.exit()
elif event.type == pygame.MOUSEBUTTONUP:
# These positions assume origin lower-left
target_pos = coords_from_pygame(pygame.mouse.get_pos(), HEIGHT)
cue_pos = table.cue_ball.pos
table.cue_angle = get_angle(target_pos, cue_pos)
print('AFTER SETTING cue_angle', table.cue_angle)
elif event.type == pygame.KEYDOWN:
if event.key == pygame.K_q:
sys.exit()
elif event.key == pygame.K_b:
# BREAK cue ball
mag = 500.0
force = Vector(mag * np.cos(np.radians(table.cue_angle)),
mag * np.sin(np.radians(table.cue_angle)))
table.balls[BallType.CUE].apply_force(force)
elif event.key == pygame.K_SPACE:
# Strike cue ball
mag = 50.0
force = Vector(mag * np.cos(np.radians(table.cue_angle)),
mag * np.sin(np.radians(table.cue_angle)))
table.balls[BallType.CUE].apply_force(force)
elif event.key == pygame.K_p:
# DEBUG set all speeds to 0
for ball in balls:
ball.vel.x, ball.vel.y = 0, 0
elif event.key == pygame.K_r:
# DEBUG reset
# Create pool table
nw = coords_from_pygame((TABLE_OFFSET_X, TABLE_OFFSET_Y),
HEIGHT)
se = coords_from_pygame(
(TABLE_OFFSET_X + TABLE_LENGTH,
TABLE_OFFSET_Y + TABLE_LENGTH / 2), HEIGHT)
table = PoolTable(nw, se)
# Table time step
table.time_step()
# Draw pool table
draw_pool_table(table)
draw_cue_stick_line(table)
draw_cue_ghost_ball(table)
draw_cue_ball_deflection_line(table)
draw_object_ball_deflection_line(table)
# Draw all pool balls
for ball in balls:
draw_pool_ball(ball)
pygame.display.flip()