def cast_ray(origin, direction, depth):
    if depth > MAX_DEPTH:
        return background_color

    hit_point, normal, material = scene_intersect(origin, direction)

    if no_intersection:
        return background_color

    color = Vec3(0, 0, 0)

    # 计算局部光照（漫反射 + 高光）
    for light in lights:
        if not in_shadow(hit_point, light):
            color += compute_lighting(hit_point, normal, material, light)

    # 反射
    if material.is_reflective:
        reflect_dir = reflect(direction, normal)
        reflect_color = cast_ray(hit_point, reflect_dir, depth + 1)
        color += reflect_color * material.reflectivity

    # 折射
    if material.is_transparent:
        refract_dir = refract(direction, normal, material.ior)
        refract_color = cast_ray(hit_point, refract_dir, depth + 1)
        color += refract_color * material.transparency

    return color

def ray_sphere_intersect(origin, direction, center, radius):
    L = center - origin
    tca = L.dot(direction)
    d2 = L.dot(L) - tca * tca

    if d2 > radius * radius:
        return False, 0

    thc = math.sqrt(radius * radius - d2)
    t0 = tca - thc
    t1 = tca + thc

    if t0 < 0:
        t0 = t1
    if t0 < 0:
        return False, 0

    return True, t0

def in_shadow(hit_point, light_position):
    shadow_dir = (light_position - hit_point).normalize()
    shadow_origin = hit_point + normal * epsilon  # 偏移避免自相交
    light_distance = (light_position - hit_point).norm()

    intersect, shadow_hit, _, _ = scene_intersect(shadow_origin, shadow_dir)

    if intersect and (shadow_hit - shadow_origin).norm() < light_distance:
        return True  # 在阴影中
    return False

class Vec3:
    """三维向量：用于点、方向、颜色"""
    def __init__(self, x=0, y=0, z=0): ...
    def dot(self, other): ...      # 点积
    def normalize(self): ...        # 归一化
    def __add__, __sub__, __mul__   # 运算符重载

class Material:
    """材质：颜色、反射率、折射率等"""
    diffuse_color: Vec3
    albedo: (diffuse, specular, reflect, refract)
    specular_exponent: float
    refractive_index: float

class Sphere:
    """球体：圆心、半径、材质"""
    center: Vec3
    radius: float
    material: Material

class Light:
    """点光源：位置、强度"""
    position: Vec3
    intensity: float

def render(width, height, fov):
    framebuffer = []
    for j in range(height):
        for i in range(width):
            # 计算光线方向
            x = (2 * (i + 0.5) / width - 1) * tan(fov/2) * width/height
            y = -(2 * (j + 0.5) / height - 1) * tan(fov/2)
            direction = Vec3(x, y, -1).normalize()

            # 从原点发射光线
            color = cast_ray(Vec3(0,0,0), direction, depth=0)
            framebuffer.append(color)

    return framebuffer

import math
from typing import List, Tuple
from PIL import Image

class Vec3:
    """三维向量类，用于表示点、方向、颜色等"""

    def __init__(self, x: float = 0.0, y: float = 0.0, z: float = 0.0):
        self.x = x
        self.y = y
        self.z = z

    def __add__(self, other: 'Vec3') -> 'Vec3':
        return Vec3(self.x + other.x, self.y + other.y, self.z + other.z)

    def __sub__(self, other: 'Vec3') -> 'Vec3':
        return Vec3(self.x - other.x, self.y - other.y, self.z - other.z)

    def __mul__(self, scalar: float) -> 'Vec3':
        return Vec3(self.x * scalar, self.y * scalar, self.z * scalar)

    def __rmul__(self, scalar: float) -> 'Vec3':
        return self.__mul__(scalar)

    def __neg__(self) -> 'Vec3':
        return Vec3(-self.x, -self.y, -self.z)

    def dot(self, other: 'Vec3') -> float:
        """点积"""
        return self.x * other.x + self.y * other.y + self.z * other.z

    def norm(self) -> float:
        """向量的模（长度）"""
        return math.sqrt(self.dot(self))

    def normalize(self) -> 'Vec3':
        """返回单位向量"""
        n = self.norm()
        if n < 1e-6:
            return Vec3(0, 0, 0)
        return self * (1.0 / n)

    def __repr__(self) -> str:
        return f"Vec3({self.x:.2f}, {self.y:.2f}, {self.z:.2f})"


def reflect(I: Vec3, N: Vec3) -> Vec3:
    """反射向量: I - 2 * (I · N) * N"""
    return I - N * 2.0 * I.dot(N)


def refract(I: Vec3, N: Vec3, refractive_index: float) -> Vec3:
    """折射向量（使用斯涅尔定律）"""
    cosi = -max(-1.0, min(1.0, I.dot(N)))
    etai = 1.0
    etat = refractive_index
    n = N

    if cosi < 0:
        cosi = -cosi
        etai, etat = etat, etai
        n = -N

    eta = etai / etat
    k = 1 - eta * eta * (1 - cosi * cosi)

    if k < 0:
        return Vec3(0, 0, 0)
    else:
        return I * eta + n * (eta * cosi - math.sqrt(k))


class Material:
    """材质类，定义物体表面属性"""

    def __init__(self,
                 diffuse_color: Vec3 = Vec3(1, 1, 1),
                 albedo: Tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.0),
                 specular_exponent: float = 0.0,
                 refractive_index: float = 1.0):
        """
        参数:
            diffuse_color: 漫反射颜色
            albedo: (漫反射, 镜面反射, 反射, 折射) 权重
            specular_exponent: 镜面反射指数（高光强度）
            refractive_index: 折射率
        """
        self.diffuse_color = diffuse_color
        self.albedo = albedo
        self.specular_exponent = specular_exponent
        self.refractive_index = refractive_index


class Sphere:
    """球体类"""

    def __init__(self, center: Vec3, radius: float, material: Material):
        self.center = center
        self.radius = radius
        self.material = material

    def ray_intersect(self, orig: Vec3, dir: Vec3) -> Tuple[bool, float]:
        """
        射线与球体相交检测
        返回: (是否相交, 交点距离)
        """
        L = self.center - orig
        tca = L.dot(dir)
        d2 = L.dot(L) - tca * tca

        if d2 > self.radius * self.radius:
            return False, 0.0

        thc = math.sqrt(self.radius * self.radius - d2)
        t0 = tca - thc
        t1 = tca + thc

        if t0 < 0:
            t0 = t1
        if t0 < 0:
            return False, 0.0

        return True, t0


class Light:
    """点光源类"""

    def __init__(self, position: Vec3, intensity: float = 1.0):
        self.position = position
        self.intensity = intensity


def scene_intersect(orig: Vec3, dir: Vec3, spheres: List[Sphere]) -> Tuple[bool, Vec3, Vec3, Material]:
    """
    检测射线与场景的交点
    返回: (是否相交, 交点位置, 法线, 材质)
    """
    spheres_dist = float('inf')
    hit = Vec3()
    N = Vec3()
    material = Material()

    # 检测与所有球体的相交
    for sphere in spheres:
        intersect, dist = sphere.ray_intersect(orig, dir)
        if intersect and dist < spheres_dist:
            spheres_dist = dist
            hit = orig + dir * dist
            N = (hit - sphere.center).normalize()
            material = sphere.material

    # 添加棋盘地板
    checkerboard_dist = float('inf')
    if abs(dir.y) > 1e-3:
        d = -(orig.y + 4) / dir.y  # 地板在 y = -4
        pt = orig + dir * d
        if d > 0 and abs(pt.x) < 10 and pt.z < -10 and pt.z > -30 and d < spheres_dist:
            checkerboard_dist = d
            hit = pt
            N = Vec3(0, 1, 0)
            # 棋盘格纹理
            color = Vec3(1, 1, 1) if (int(0.5 * hit.x + 1000) + int(0.5 * hit.z)) & 1 else Vec3(1, 0.7, 0.3)
            material.diffuse_color = color
            material.albedo = (0.6, 0.3, 0.1, 0.0)
            material.specular_exponent = 50.0

    return min(spheres_dist, checkerboard_dist) < 1000, hit, N, material


def cast_ray(orig: Vec3, dir: Vec3, spheres: List[Sphere], lights: List[Light], depth: int = 0) -> Vec3:
    """
    发射光线并计算颜色
    depth: 递归深度（用于反射和折射）
    """
    intersect, point, N, material = scene_intersect(orig, dir, spheres)

    if depth > 4 or not intersect:
        # 背景色（渐变）
        return Vec3(0.2, 0.7, 0.8)

    # 计算反射方向
    reflect_dir = reflect(dir, N).normalize()
    # 计算折射方向
    refract_dir = refract(dir, N, material.refractive_index).normalize()

    # 稍微偏移以避免自相交
    reflect_orig = point + N * 1e-3 if reflect_dir.dot(N) > 0 else point - N * 1e-3
    refract_orig = point - N * 1e-3 if refract_dir.dot(N) < 0 else point + N * 1e-3

    # 递归计算反射和折射颜色
    reflect_color = cast_ray(reflect_orig, reflect_dir, spheres, lights, depth + 1)
    refract_color = cast_ray(refract_orig, refract_dir, spheres, lights, depth + 1)

    # 计算光照
    diffuse_light_intensity = 0.0
    specular_light_intensity = 0.0

    for light in lights:
        light_dir = (light.position - point).normalize()
        light_distance = (light.position - point).norm()

        # 阴影检测
        shadow_orig = point + N * 1e-3 if light_dir.dot(N) > 0 else point - N * 1e-3
        shadow_intersect, shadow_pt, _, _ = scene_intersect(shadow_orig, light_dir, spheres)

        if shadow_intersect and (shadow_pt - shadow_orig).norm() < light_distance:
            continue  # 在阴影中

        # 漫反射（Lambert 余弦定律）
        diffuse_light_intensity += light.intensity * max(0.0, light_dir.dot(N))

        # 镜面反射（Blinn-Phong 模型）
        specular_light_intensity += math.pow(
            max(0.0, -reflect(-light_dir, N).dot(dir)),
            material.specular_exponent
        ) * light.intensity

    # 组合所有光照分量
    color = (material.diffuse_color * diffuse_light_intensity * material.albedo[0] +
             Vec3(1, 1, 1) * specular_light_intensity * material.albedo[1] +
             reflect_color * material.albedo[2] +
             refract_color * material.albedo[3])

    return color


def render(spheres: List[Sphere], lights: List[Light], width: int = 1024, height: int = 768) -> List[Vec3]:
    """
    渲染场景
    返回: 颜色数组
    """
    fov = math.pi / 3.0  # 视场角 60 度
    framebuffer = []

    for j in range(height):
        for i in range(width):
            # 计算射线方向
            x = (2 * (i + 0.5) / width - 1) * math.tan(fov / 2.0) * width / height
            y = -(2 * (j + 0.5) / height - 1) * math.tan(fov / 2.0)
            dir = Vec3(x, y, -1).normalize()

            # 从原点发射射线
            framebuffer.append(cast_ray(Vec3(0, 0, 0), dir, spheres, lights))

        # 显示进度
        if (j + 1) % 50 == 0:
            print(f"渲染进度: {(j + 1) / height * 100:.1f}%")

    return framebuffer


def save_png(filename: str, framebuffer: List[Vec3], width: int, height: int):
    """保存为 PNG 格式图像（使用 Pillow）"""
    if Image is None:
        print(f"跳过保存 {filename}: 需要安装 Pillow 库")
        return

    # 创建 RGB 图像
    img = Image.new('RGB', (width, height))
    pixels = []

    for color in framebuffer:
        # 限制颜色范围在 [0, 1] 并转换为 [0, 255]
        r = max(0, min(255, int(255 * color.x)))
        g = max(0, min(255, int(255 * color.y)))
        b = max(0, min(255, int(255 * color.z)))
        pixels.append((r, g, b))

    img.putdata(pixels)
    img.save(filename)
    print(f"  ✓ {filename}")


def main():
    """主函数 - 创建场景并渲染"""

    # 定义材质
    ivory = Material(
        diffuse_color=Vec3(0.4, 0.4, 0.3),
        albedo=(0.6, 0.3, 0.1, 0.0),
        specular_exponent=50.0
    )

    glass = Material(
        diffuse_color=Vec3(0.6, 0.7, 0.8),
        albedo=(0.0, 0.5, 0.1, 0.8),
        specular_exponent=125.0,
        refractive_index=1.5
    )

    red_rubber = Material(
        diffuse_color=Vec3(0.3, 0.1, 0.1),
        albedo=(0.9, 0.1, 0.0, 0.0),
        specular_exponent=10.0
    )

    mirror = Material(
        diffuse_color=Vec3(1.0, 1.0, 1.0),
        albedo=(0.0, 10.0, 0.8, 0.0),
        specular_exponent=1425.0
    )

    # 创建球体
    spheres = [
        Sphere(Vec3(-3, 0, -16), 2, ivory),
        Sphere(Vec3(-1.0, -1.5, -12), 2, glass),
        Sphere(Vec3(1.5, -0.5, -18), 3, red_rubber),
        Sphere(Vec3(9, 3, -18), 4, mirror)
    ]

    # 创建光源
    lights = [
        Light(Vec3(-20, 20, 20), 1.5),
        Light(Vec3(30, 50, -25), 1.8),
    ]

    # 渲染设置
    width = 2048
    height = 1536

    print("开始渲染...")
    print(f"分辨率: {width}x{height}")
    print(f"球体数量: {len(spheres)}")
    print(f"光源数量: {len(lights)}")
    print()

    # 渲染场景
    framebuffer = render(spheres, lights, width, height)

    # 保存图像
    print("\n保存图像...")
    save_png("output.png", framebuffer, width, height)
    print("\n完成!")


if __name__ == "__main__":
    main()