from PIL import Image, ImageOps
import os

input_pre_path = 'inputs'
output_pre_path = 'outputs'

def resize_to_width(image, target_width):
    w, h = image.size
    ratio = target_width / w
    new_height = int(h * ratio)
    return image.resize((target_width, new_height))

def open_and_correct_image(path):
    img = Image.open(path)
    return ImageOps.exif_transpose(img)

def create_vertical_mosaic(up_list_path, down_list_path, center_path, target_width = 800, interval = 10, bg_color = (255, 255, 255)):
    # 加载图片
    print('加载图片...')
    up_imgs = [open_and_correct_image(os.path.join(input_pre_path, p)) for p in up_list_path]
    down_imgs = [open_and_correct_image(os.path.join(input_pre_path, p)) for p in down_list_path]
    center_img = open_and_correct_image(os.path.join(input_pre_path, center_path))
    print('加载完成')

    # 缩放图片
    print('缩放图片...')
    up_imgs = [resize_to_width(img, target_width) for img in up_imgs]
    down_imgs = [resize_to_width(img, target_width) for img in down_imgs]
    center_img = resize_to_width(center_img, target_width)
    print('缩放完成')

    # 计算高度
    total_height = center_img.height + max(sum(img.height + interval for img in up_imgs), sum(img.height + interval for img in down_imgs)) * 2
    # 创建长图
    mosaic = Image.new('RGB', (target_width, total_height), color=bg_color)

    # 放置图片到对应位置 -----------------
    # 中心图片
    print('绘制中心...')
    mosaic.paste(center_img, (0, (total_height - center_img.height) // 2))
    # 上方图片
    print('绘制上方...')
    y_offset = (total_height - center_img.height) // 2 - sum(img.height + interval for img in up_imgs)
    for img in up_imgs:
        mosaic.paste(img, (0, y_offset))
        y_offset += img.height + interval
    # 下方图片
    print('绘制下方...')
    y_offset = (total_height + center_img.height) // 2 + interval
    for img in down_imgs:
        mosaic.paste(img, (0, y_offset))
        y_offset += img.height + interval

    # 保存 --------------------------------
    return mosaic
    # output_path = os.path.join(output_pre_path, output_name)
    # mosaic.save(output_path)
    # print(f"保存至 {output_path}")

# =============================================================================================================================================

# 贵阳 -------------------------------------------------------
input_pre_path = os.path.join(input_pre_path, '贵阳')
up_img_list = ['D.jpg', 'C.jpg', 'E.jpg', 'H.jpg', 'F.jpg']
down_img_list = ['G.jpg', 'A.jpg', 'B.jpg']
center_img = '中心.jpg'
output_name = '贵阳.jpg'

# 重庆 -------------------------------------------------------
# input_pre_path = os.path.join(input_pre_path, '重庆')
# up_img_list = ['A2.jpg', 'J.jpg', 'D.jpg', 'C.jpg', 'E.jpg']
# down_img_list = ['F.jpg', 'G.jpg', 'H.jpg', 'I2.jpg', 'B.jpg']
# center_img = '中心.jpg'
# output_name = '重庆2.jpg'

# 洛阳 -------------------------------------------------------
# input_pre_path = os.path.join(input_pre_path, '洛阳')
# up_img_list = ['A.jpg', 'D.jpg', 'C.jpg', 'G.jpg']
# down_img_list = ['F.jpg', 'E.jpg', 'H2.jpg', 'B.jpg', 'I.jpg']
# center_img = '中心.jpg'
# output_name = '洛阳2.jpg'

# 上海 -------------------------------------------------------
# input_pre_path = os.path.join(input_pre_path, '上海')
# up_img_list = ['G.jpg', 'J.jpg', 'D.jpg', 'E.jpg', 'C.jpg', 'A.jpg']
# down_img_list = ['B.jpg', 'F.jpg', 'H.jpg', 'I.jpg']
# center_img = '中心.jpg'
# output_name = '上海.jpg'

print('开始处理...')
result = create_vertical_mosaic(up_img_list, down_img_list, center_img, 1580, 30)
output_path = os.path.join(output_pre_path, output_name)
result.save(output_path)
print(f"保存至 {output_path}")

print('处理结束')