import cv2
import numpy as np
import os
import time

def safe_overlay(background, overlay_img, x_offset, y_offset):
    """安全叠加图像（自动处理边界和alpha通道）"""
    bg_height, bg_width = background.shape[:2]
    overlay_height, overlay_width = overlay_img.shape[:2]
    
    # 计算实际可叠加区域
    x1 = max(0, x_offset)
    y1 = max(0, y_offset)
    x2 = min(bg_width, x_offset + overlay_width)
    y2 = min(bg_height, y_offset + overlay_height)
    
    # 计算overlay图像中的对应区域
    ox1 = max(0, -x_offset)
    oy1 = max(0, -y_offset)
    ox2 = overlay_width - max(0, (x_offset + overlay_width) - bg_width)
    oy2 = overlay_height - max(0, (y_offset + overlay_height) - bg_height)
    
    # 确保有可叠加区域
    if x1 >= x2 or y1 >= y2 or ox1 >= ox2 or oy1 >= oy2:
        return background
    
    # 处理alpha通道
    if overlay_img.shape[2] == 4:
        overlay = overlay_img[oy1:oy2, ox1:ox2]
        alpha = overlay[:, :, 3:] / 255.0
        overlay = overlay[:, :, :3]
    else:
        overlay = overlay_img[oy1:oy2, ox1:ox2]
        alpha = 1.0
    
    # 执行叠加
    background[y1:y2, x1:x2] = (overlay * alpha + 
                               background[y1:y2, x1:x2] * (1.0 - alpha))
    return background

def main():
    # 检查并创建output目录
    if not os.path.exists('output'):
        os.makedirs('output')

    # 记录开始总时间
    total_start_time = time.time()
    print("程序开始执行...")

    try:
        # 1. 创建空白背景 (1080x1920)
        start_time = time.time()
        background = np.zeros((1920, 1080, 3), dtype=np.uint8)
        print(f"1. 创建空白背景耗时: {time.time() - start_time:.4f}秒")

        # 2. 处理a.jpg
        start_time = time.time()
        a_img = cv2.imread('a.jpg')
        if a_img is None:
            raise FileNotFoundError("无法找到a.jpg文件")

        # 放大105%并旋转5度
        a_resized = cv2.resize(a_img, None, fx=1.05, fy=1.05, interpolation=cv2.INTER_LINEAR)
        rows, cols = a_resized.shape[:2]
        M = cv2.getRotationMatrix2D((cols/2, rows/2), -5, 1)
        a_rotated = cv2.warpAffine(a_resized, M, (cols, rows))
        a_rotated_rgba = cv2.cvtColor(a_rotated, cv2.COLOR_BGR2BGRA)

        # 计算居中位置并安全叠加
        x_offset = (background.shape[1] - a_rotated.shape[1]) // 2
        y_offset = (background.shape[0] - a_rotated.shape[0]) // 2
        background = safe_overlay(background, a_rotated_rgba, x_offset, y_offset)
        print(f"2. 处理a.jpg耗时: {time.time() - start_time:.4f}秒")

        # 3. 处理b.png
        start_time = time.time()
        b_img = cv2.imread('b.png', cv2.IMREAD_UNCHANGED)
        if b_img is None:
            raise FileNotFoundError("无法找到b.png文件")

        # 缩小70%并旋转15度
        b_resized = cv2.resize(b_img, None, fx=0.7, fy=0.7, interpolation=cv2.INTER_AREA)
        rows, cols = b_resized.shape[:2]
        M = cv2.getRotationMatrix2D((cols/2, rows/2), -15, 1)
        b_rotated = cv2.warpAffine(b_resized, M, (cols, rows))
        
        # 如果缺少alpha通道则添加
        if b_rotated.shape[2] == 3:
            b_rotated_rgba = cv2.cvtColor(b_rotated, cv2.COLOR_BGR2BGRA)
        else:
            b_rotated_rgba = b_rotated

        # 叠加到右下角
        x_offset = background.shape[1] - b_rotated.shape[1] - 50
        y_offset = background.shape[0] - b_rotated.shape[0] - 50
        background = safe_overlay(background, b_rotated_rgba, x_offset, y_offset)
        print(f"3. 处理b.png耗时: {time.time() - start_time:.4f}秒")

        # 4. 处理c.jpg
        start_time = time.time()
        c_img = cv2.imread('c.jpg')
        if c_img is None:
            raise FileNotFoundError("无法找到c.jpg文件")

        # 缩小30%并旋转90度
        c_resized = cv2.resize(c_img, None, fx=0.3, fy=0.3, interpolation=cv2.INTER_AREA)
        rows, cols = c_resized.shape[:2]
        M = cv2.getRotationMatrix2D((cols/2, rows/2), -90, 1)
        c_rotated = cv2.warpAffine(c_resized, M, (cols, rows))
        c_rotated_rgba = cv2.cvtColor(c_rotated, cv2.COLOR_BGR2BGRA)

        # 叠加到左上角
        x_offset, y_offset = 50, 50
        background = safe_overlay(background, c_rotated_rgba, x_offset, y_offset)
        print(f"4. 处理c.jpg耗时: {time.time() - start_time:.4f}秒")

        # 5. 保存结果
        start_time = time.time()
        output_path = os.path.join('output', 'final_result.jpg')
        cv2.imwrite(output_path, background)
        print(f"5. 保存结果耗时: {time.time() - start_time:.4f}秒")
        print(f"\n总处理时间: {time.time() - total_start_time:.4f}秒")
        print(f"结果已保存到: {os.path.abspath(output_path)}")

        # 显示结果（可选）
        cv2.imshow('Final Result', background)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    except Exception as e:
        print(f"\n程序执行出错: {str(e)}")

if __name__ == "__main__":
    main()