import pygame
import math
import sys

pygame.init()
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Ray Marching - Single Ray with Steps")
clock = pygame.time.Clock()

# Colors
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
GRAY = (180, 180, 180)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)

# SDF primitives
def sdf_circle(p, center, radius):
    return math.hypot(p[0] - center[0], p[1] - center[1]) - radius

def sdf_box(p, center, size):
    dx = abs(p[0] - center[0]) - size[0] / 2
    dy = abs(p[1] - center[1]) - size[1] / 2
    dx = max(dx, 0)
    dy = max(dy, 0)
    return math.hypot(dx, dy)

objects = [
    {"sdf": lambda p: sdf_circle(p, (400, 300), 60), "color": GREEN},
    {"sdf": lambda p: sdf_box(p, (200, 150), (100, 100)), "color": BLUE},
    {"sdf": lambda p: sdf_box(p, (600, 400), (80, 150)), "color": RED},
]

def scene_sdf(p):
    return min(obj["sdf"](p) for obj in objects)

def ray_march(origin, direction, max_steps=100, max_dist=1000, epsilon=1.0):
    p = list(origin)
    total_dist = 0
    steps = []
    for _ in range(max_steps):
        dist = scene_sdf(p)
        steps.append((tuple(p), dist))
        if dist < epsilon or total_dist > max_dist:
            break
        p[0] += direction[0] * dist
        p[1] += direction[1] * dist
        total_dist += dist
    return steps

# Main loop
running = True
while running:
    screen.fill(BLACK)
    mouse_pos = pygame.mouse.get_pos()
    origin = (600, 200)
    dx = mouse_pos[0] - origin[0]
    dy = mouse_pos[1] - origin[1]
    length = math.hypot(dx, dy)
    if length == 0:
        direction = (0, 0)
    else:
        direction = (dx / length, dy / length)

    # Draw scene objects
    for obj in objects:
        if obj["color"] == GREEN:
            pygame.draw.circle(screen, obj["color"], (400, 300), 60, 2)
        elif obj["color"] == BLUE:
            pygame.draw.rect(screen, obj["color"], pygame.Rect(150, 100, 100, 100), 2)
        elif obj["color"] == RED:
            pygame.draw.rect(screen, obj["color"], pygame.Rect(560, 325, 80, 150), 2)

    # Perform ray march
    steps = ray_march(origin, direction)

    # Draw ray steps
    for point, dist in steps:
        pygame.draw.circle(screen, GRAY, (int(point[0]), int(point[1])), int(dist), 1)

    # Draw final point
    if steps:
        pygame.draw.circle(screen, WHITE, (int(steps[-1][0][0]), int(steps[-1][0][1])), 3)

    pygame.draw.line(screen, WHITE, origin, mouse_pos, 1)
    pygame.draw.circle(screen, WHITE, origin, 4)

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

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

pygame.quit()
sys.exit()