import numpy as np
from PIL import Image
import time

def save_image(data, out_path):
    data = data.astype(np.uint8)
    image = Image.fromarray(data, mode='L')
    image.save(out_path)

def sdf_sphere(x, y, cx, cy, r):
    """计算点 (x, y) 到中心为 (cx, cy) 且半径为 r 的圆的符号距离"""
    dx = x - cx
    dy = y - cy
    return np.sqrt(dx * dx + dy * dy) - r

def func1(input_image):
  	"""For循环遍历每一个元素"""
    width = input_image.shape[0]
    height = input_image.shape[1]
    print(f'width: {width}, height: {height}')
    for i in range(width):
        for j in range(height):
            u = i / width
            v = j / height
            if sdf_sphere(u, v, 0.5, 0.5, 0.2) < 0.1:
                input_image[i, j] = 255
    return input_image

def func2(input_image):
  	"""使用np向量操作"""
    width = input_image.shape[0]
    height = input_image.shape[1]
    u, v = np.meshgrid(np.linspace(0, 1, width), np.linspace(0, 1, height), indexing='ij')
    distances = sdf_sphere(u, v, 0.5, 0.5, 0.2)
    input_image[distances < 0.1] = 255
    return input_image

if __name__ == "__main__":
    image1 = np.zeros((512, 512), dtype=np.float32)
    image2 = np.zeros((512, 512), dtype=np.float32)
    start_time = time.time()
    out1 = func1(image1)
    end_time = time.time()
    print(end_time - start_time)	# 0.4647049903869629
    start_time = time.time()
    out2 = func2(image2)
    end_time = time.time()
    print(end_time - start_time)	# 0.00978851318359375
    save_image(out1, 'assets/temp/image.png')
    save_image(out2, 'assets/temp/image2.png')

def sdf_sphere(x, y, cx, cy, r):
    """计算点 (x, y) 到中心为 (cx, cy) 且半径为 r 的圆的符号距离"""
    dx = x - cx
    dy = y - cy
    return np.sqrt(dx * dx + dy * dy) - r

def sdf_square(x, y, cx, cy, size):
    """计算点 (x, y) 到中心为 (cx, cy) 且边长为 size 的正方形的符号距离"""
    dx = np.maximum(np.abs(x - cx) - size / 2, 0)
    dy = np.maximum(np.abs(y - cy) - size / 2, 0)
    return np.sqrt(dx * dx + dy * dy)

def sdf_square_frame(x, y, cx, cy, size):
    dx = np.abs(x - cx) - size / 2
    dy = np.abs(y - cy) - size / 2
    outside_distance = np.sqrt(np.maximum(dx, 0)**2 + np.maximum(dy, 0)**2)
    inside_distance = -np.minimum(np.maximum(dx, dy), 0)
    return outside_distance + inside_distance

arr = np.array([[1, 2], [3, 4]])
result = np.prod(arr)		# 24 = 1 * 2 * 3 *4
# 根据轴
result2 = np.prod(arr, axis=0)	# [3, 8], 3 = 1 * 3, 8 = 2 * 4
# 使用初始值
result3 = np.prod(arr, initial=2)	# 48 = 2 * 1 * 2 * 3 * 4