import pygame
import math

# Initialize Pygame
pygame.init()

# Screen settings
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
clock = pygame.time.Clock()

# Camera settings
camera_position = (WIDTH // 2, HEIGHT // 2, -500)
view_direction = (0, 0, 1)

# Light source
light_position = (WIDTH // 2, HEIGHT // 2, 200)

# Define a simple sphere
class Sphere:
    def __init__(self, center, radius, color):
        self.center = center
        self.radius = radius
        self.color = color
    
    def intersect(self, origin, direction):
        # Ray-sphere intersection calculation
        oc = (origin[0] - self.center[0], origin[1] - self.center[1], origin[2] - self.center[2])
        a = direction[0] ** 2 + direction[1] ** 2 + direction[2] ** 2
        b = 2 * (oc[0] * direction[0] + oc[1] * direction[1] + oc[2] * direction[2])
        c = oc[0] ** 2 + oc[1] ** 2 + oc[2] ** 2 - self.radius ** 2
        discriminant = b ** 2 - 4 * a * c
        if discriminant > 0:
            t1 = (-b - math.sqrt(discriminant)) / (2 * a)
            t2 = (-b + math.sqrt(discriminant)) / (2 * a)
            if t1 > 0:
                return t1
            if t2 > 0:
                return t2
        return None

# Simple lighting model
def lighting(point, normal, light_pos, color):
    light_dir = (light_pos[0] - point[0], light_pos[1] - point[1], light_pos[2] - point[2])
    light_distance = math.sqrt(light_dir[0] ** 2 + light_dir[1] ** 2 + light_dir[2] ** 2)
    light_dir = (light_dir[0] / light_distance, light_dir[1] / light_distance, light_dir[2] / light_distance)

    # Dot product for diffuse lighting
    dot = max(0, normal[0] * light_dir[0] + normal[1] * light_dir[1] + normal[2] * light_dir[2])
    
    # Simulating simple ambient light
    ambient = 0.1
    diffuse = 0.9 * dot
    intensity = ambient + diffuse

    # Color modification based on intensity
    color = (min(255, color[0] * intensity), min(255, color[1] * intensity), min(255, color[2] * intensity))
    
    return color

# Ray tracing logic
def ray_trace():
    for x in range(0, WIDTH, 4):  # Rasterizing the screen
        for y in range(0, HEIGHT, 4):
            direction = (x - camera_position[0], y - camera_position[1], 0 - camera_position[2])
            length = math.sqrt(direction[0] ** 2 + direction[1] ** 2 + direction[2] ** 2)
            direction = (direction[0] / length, direction[1] / length, direction[2] / length)

            # Check if ray intersects any object
            sphere = Sphere((WIDTH // 2, HEIGHT // 2, 1000), 200, (255, 0, 0))
            t = sphere.intersect(camera_position, direction)

            if t:
                # Calculate point of intersection
                point = (camera_position[0] + direction[0] * t, 
                         camera_position[1] + direction[1] * t,
                         camera_position[2] + direction[2] * t)
                normal = (point[0] - sphere.center[0], point[1] - sphere.center[1], point[2] - sphere.center[2])
                norm_len = math.sqrt(normal[0] ** 2 + normal[1] ** 2 + normal[2] ** 2)
                normal = (normal[0] / norm_len, normal[1] / norm_len, normal[2] / norm_len)
                
                # Apply lighting model
                color = lighting(point, normal, light_position, sphere.color)
                screen.set_at((x, y), color)

# Main loop
running = True
while running:
    screen.fill((0, 0, 0))  # Clear screen
    
    ray_trace()  # Perform ray tracing
    
    pygame.display.flip()  # Update display
    clock.tick(60)  # Limit the frame rate to 60 FPS
    
    # Handle events
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

pygame.quit()