import torch
import cv2
from torch.nn import functional as F
import numpy as np
import math 
from . import cube2equi, get_transform_matrix_1
from PIL import Image
import torchvision.transforms as T


def convert_to_cubemap(L_F_D_B_R):
    L_F_D_B_R = torch.tensor(L_F_D_B_R).permute(0,3,1,2)
    L_F, B_R = torch.chunk(L_F_D_B_R, 2, dim=-2)
    L_, F_ = torch.chunk(L_F, 2, dim=-1)
    B_, R_ = torch.chunk(B_R, 2, dim=-1)
 
    scale = np.sqrt(2)/2
    L_affine = F.grid_sample(L_, get_transform_matrix_1('L',L_.shape,scale,reverse=True), mode='bilinear', align_corners=False)
    F_affine = F.grid_sample(F_, get_transform_matrix_1('F',F_.shape,scale,reverse=True), mode='bilinear', align_corners=False)
    B_affine = F.grid_sample(B_, get_transform_matrix_1('B',B_.shape,scale,reverse=True), mode='bilinear', align_corners=False)
    R_affine = F.grid_sample(R_, get_transform_matrix_1('R',R_.shape,scale,reverse=True), mode='bilinear', align_corners=False)

    scale = np.sqrt(2)/4
    D_affine = F.grid_sample(L_F_D_B_R, get_transform_matrix_1('D',L_F_D_B_R.shape,scale,reverse=True), mode='bilinear', align_corners=False)

    L_ = torch.rot90(torch.flip(L_,dims=[-1]), k=-1, dims=(-2, -1))
    F_ = torch.rot90(torch.flip(F_,dims=[-1]), k=1, dims=(-2, -1))
    B_ = torch.rot90(torch.flip(B_,dims=[-1]), k=1, dims=(-2, -1))
    R_ = torch.rot90(torch.flip(R_,dims=[-1]), k=-1, dims=(-2, -1))
    L_F_U_B_R = torch.cat([torch.cat([L_,F_],dim=-1),torch.cat([B_,R_],dim=-1)],dim=-2)
    L_F_U_B_R = torch.flip(L_F_U_B_R,dims=[-1])
    U_affine = F.grid_sample(L_F_U_B_R, get_transform_matrix_1('U',L_F_U_B_R.shape,scale,reverse=True), mode='bilinear', align_corners=False)
    return torch.cat([F_affine,R_affine,B_affine,L_affine,U_affine,D_affine],dim=-1)

def post_process_U(U_tensors):
 
    U_tensors_F_B = torch.cat(torch.chunk(U_tensors,2,dim=-2)[: :-1],dim=-2)
    U_tensors_L_R = torch.cat(torch.chunk(U_tensors,2,dim=-1)[: :-1],dim=-1)

    b,c,H,W = U_tensors_F_B.shape
    y, x = torch.meshgrid(torch.arange(H), torch.arange(W), indexing='ij')
    x = x.float()
    y = y.float()
    
    center = torch.tensor([H/2, H/2], device=x.device)
    
    # 计算相对坐标
    dx = x - center[0]
    dy = center[1] - y  # 反转y轴
    
    # 极坐标参数
    theta = torch.atan2(dy, dx)
    r = torch.sqrt(dx**2 + dy**2)
    
    # 圆形有效区域判断
    valid_mask = r <= (H/2)
    
    # 计算菱形映射参数
    cos_theta = torch.cos(theta)
    sin_theta = torch.sin(theta)
    denominator = torch.abs(cos_theta) + torch.abs(sin_theta)
    r_max = (H/2) / (denominator + 1e-6)
    
    # 计算原图坐标
    r_prime = torch.where(valid_mask, r * (r_max / (H/2)), 0.0)
    px = center[0] + r_prime * cos_theta
    py = center[1] - r_prime * sin_theta
    
    # 归一化处理
    grid_x = (px / (W-1)) * 2 - 1
    grid_y = (py / (H-1)) * 2 - 1
    
    # 无效点处理
    grid_x = torch.where(valid_mask, grid_x, -2.0)
    grid_y = torch.where(valid_mask, grid_y, -2.0)
    
    grid = torch.stack((grid_x, grid_y), dim=-1).unsqueeze(0).repeat(b,1,1,1).to(device=U_tensors_F_B.device)
    affine_U_tensors_F_B = F.grid_sample(U_tensors_F_B, grid, mode='bilinear', align_corners=False)
    affine_U_tensors_L_R = F.grid_sample(U_tensors_L_R, grid, mode='bilinear', align_corners=False)

 
    U_FB = torch.cat(torch.chunk(affine_U_tensors_F_B,2,dim=-2)[: :-1],dim=-2)
    U_LR = torch.cat(torch.chunk(affine_U_tensors_L_R,2,dim=-1)[: :-1],dim=-1)
    b_,c_,h_,w_ = U_FB.shape
    FB_mask = generate_diamond_mask(h_)
    FB_mask = np.concatenate((FB_mask[h_//2:,:], FB_mask[0:h_//2,:]), axis=0)
    
    U_affine =  (torch.from_numpy(FB_mask))*U_FB + (torch.from_numpy(1-FB_mask))*U_LR
    # kernel_size = (5, 5)
    # sigma = (2.0, 2.0)
    # transform = T.GaussianBlur(kernel_size=kernel_size, sigma=sigma)
    # U_affine = torch.stack([transform(image) for image in U_affine])
    mask_np = np.zeros((h_, w_)).astype(np.uint8)
    inpaint_size = 8
    h_off = (h_-inpaint_size)//2
    w_off = (w_-inpaint_size)//2
    mask_np[0:inpaint_size,0:inpaint_size] = 1
    mask_np[(h_-inpaint_size):,0:inpaint_size] = 1
    mask_np[0:inpaint_size,(w_-inpaint_size):] = 1
    mask_np[(h_-inpaint_size):,(w_-inpaint_size):] = 1
    mask_np[h_off:h_off+inpaint_size,w_off:w_off+inpaint_size] = 1
    repaired_images = []
    for i in range(U_affine.shape[0]):
        # 将张量转换为 NumPy 数组并将通道移至最后
        image_np = (U_affine[i].permute(1, 2, 0).cpu().numpy()*255).astype(np.uint8)
        
        # 使用 OpenCV 的 inpaint 进行修复
        repaired_np = cv2.inpaint(image_np, mask_np, inpaintRadius=3, flags=cv2.INPAINT_TELEA)

        # 将修复后的结果转换回 PyTorch 张量并将通道移至最前
        repaired_tensor = torch.from_numpy(repaired_np).permute(2, 0, 1)

        # 添加到结果列表中
        repaired_images.append(repaired_tensor)

    U_affine = torch.stack(repaired_images)/255.0
    return U_affine

def post_process_U2(U_tensors):
 
    U_tensors_F_B = torch.cat(torch.chunk(U_tensors,2,dim=-2)[: :-1],dim=-2)
    U_tensors_L_R = torch.cat(torch.chunk(U_tensors,2,dim=-1)[: :-1],dim=-1)

    b,c,H,W = U_tensors_F_B.shape
    y, x = torch.meshgrid(torch.arange(H), torch.arange(W), indexing='ij')
    x = x.float()
    y = y.float()
    
    center = torch.tensor([H/2, H/2], device=x.device)
    
    # 计算相对坐标
    dx = x - center[0]
    dy = center[1] - y  # 反转y轴
    
    # 极坐标参数
    theta = torch.atan2(dy, dx)
    r = torch.sqrt(dx**2 + dy**2)
    
    # 圆形有效区域判断
    valid_mask = r <= (H/2)
    
    # 计算菱形映射参数
    cos_theta = torch.cos(theta)
    sin_theta = torch.sin(theta)
    denominator = torch.abs(cos_theta) + torch.abs(sin_theta)
    r_max = (H/2) / (denominator + 1e-6)
    
    # 计算原图坐标
    r_prime = torch.where(valid_mask, r * (r_max / (H/2)), 0.0)
    px = center[0] + r_prime * cos_theta
    py = center[1] - r_prime * sin_theta
    
    # 归一化处理
    grid_x = (px / (W-1)) * 2 - 1
    grid_y = (py / (H-1)) * 2 - 1
    
    # 无效点处理
    grid_x = torch.where(valid_mask, grid_x, -2.0)
    grid_y = torch.where(valid_mask, grid_y, -2.0)
    
    grid = torch.stack((grid_x, grid_y), dim=-1).unsqueeze(0).repeat(b,1,1,1).to(device=U_tensors_F_B.device)
    affine_U_tensors_F_B = F.grid_sample(U_tensors_F_B, grid, mode='bilinear', align_corners=False)
    affine_U_tensors_L_R = F.grid_sample(U_tensors_L_R, grid, mode='bilinear', align_corners=False)

 
    U_FB = torch.cat(torch.chunk(affine_U_tensors_F_B,2,dim=-2)[: :-1],dim=-2)
    U_LR = torch.cat(torch.chunk(affine_U_tensors_L_R,2,dim=-1)[: :-1],dim=-1)
    
    b_,c_,h_,w_ = U_FB.shape
    c_mask = create_circular_mask(h_,w_)
   
    FB_mask = np.concatenate((c_mask[h_//2:,:], c_mask[0:h_//2,:]), axis=0)
    LR_mask = np.concatenate((c_mask[:,w_//2:], c_mask[:,0:w_//2]), axis=1)
    fusion_mask = (FB_mask & LR_mask)
 
    fusion_part = torch.from_numpy(fusion_mask)*(U_FB+U_LR)/2
    U_affine =  fusion_part+(torch.from_numpy(1-LR_mask))*U_FB + (torch.from_numpy(1-FB_mask))*U_LR
    # kernel_size = (5, 5)
    # sigma = (2.0, 2.0)
    # transform = T.GaussianBlur(kernel_size=kernel_size, sigma=sigma)
    # U_affine = torch.stack([transform(image) for image in U_affine])
    mask_np = create_circle_line_mask(h_,8) 
    mask_np1 = np.concatenate((mask_np[h_//2:,:], mask_np[0:h_//2,:]), axis=0)
    mask_np2 = np.concatenate((mask_np[:,w_//2:], mask_np[:,0:w_//2]), axis=1)
    mask_np = mask_np1|mask_np2

     # mask_np = np.zeros((h_, w_)).astype(np.uint8)
    # inpaint_size = 8
    # h_off = (h_-inpaint_size)//2
    # w_off = (w_-inpaint_size)//2
    # mask_np[0:inpaint_size,0:inpaint_size] = 1
    # mask_np[(h_-inpaint_size):,0:inpaint_size] = 1
    # mask_np[0:inpaint_size,(w_-inpaint_size):] = 1
    # mask_np[(h_-inpaint_size):,(w_-inpaint_size):] = 1
    # mask_np[h_off:h_off+inpaint_size,w_off:w_off+inpaint_size] = 1

    repaired_images = []
    for i in range(U_affine.shape[0]):
        # 将张量转换为 NumPy 数组并将通道移至最后
        image_np = (U_affine[i].permute(1, 2, 0).cpu().numpy()*255).astype(np.uint8)
        

        # 使用 OpenCV 的 inpaint 进行修复
        repaired_np = cv2.inpaint(image_np, mask_np, inpaintRadius=3, flags=cv2.INPAINT_TELEA)

        # 将修复后的结果转换回 PyTorch 张量并将通道移至最前
        repaired_tensor = torch.from_numpy(repaired_np).permute(2, 0, 1)

        # 添加到结果列表中
        repaired_images.append(repaired_tensor)

    U_affine = torch.stack(repaired_images)/255.0
    return U_affine

def post_process_U3(U_tensors):
 
    U_tensors_F_B = torch.cat(torch.chunk(U_tensors,2,dim=-2)[: :-1],dim=-2)
    U_tensors_L_R = torch.cat(torch.chunk(U_tensors,2,dim=-1)[: :-1],dim=-1)

    b,c,H,W = U_tensors_F_B.shape
    y, x = torch.meshgrid(torch.arange(H), torch.arange(W), indexing='ij')
    x = x.float()
    y = y.float()
    
    center = torch.tensor([H/2, H/2], device=x.device)
    
    # 计算相对坐标
    dx = x - center[0]
    dy = center[1] - y  # 反转y轴
    
    # 极坐标参数
    theta = torch.atan2(dy, dx)
    r = torch.sqrt(dx**2 + dy**2)
    
    # 圆形有效区域判断
    valid_mask = r <= (H/2)
    
    # 计算菱形映射参数
    cos_theta = torch.cos(theta)
    sin_theta = torch.sin(theta)
    denominator = torch.abs(cos_theta) + torch.abs(sin_theta)
    r_max = (H/2) / (denominator + 1e-6)
    
    # 计算原图坐标
    r_prime = torch.where(valid_mask, r * (r_max / (H/2)), 0.0)
    px = center[0] + r_prime * cos_theta
    py = center[1] - r_prime * sin_theta
    
    # 归一化处理
    grid_x = (px / (W-1)) * 2 - 1
    grid_y = (py / (H-1)) * 2 - 1
    
    # 无效点处理
    grid_x = torch.where(valid_mask, grid_x, -2.0)
    grid_y = torch.where(valid_mask, grid_y, -2.0)
    
    grid = torch.stack((grid_x, grid_y), dim=-1).unsqueeze(0).repeat(b,1,1,1).to(device=U_tensors_F_B.device)
    affine_U_tensors_F_B = F.grid_sample(U_tensors_F_B, grid, mode='bilinear', align_corners=False)
    affine_U_tensors_L_R = F.grid_sample(U_tensors_L_R, grid, mode='bilinear', align_corners=False)

 
    U_FB = torch.cat(torch.chunk(affine_U_tensors_F_B,2,dim=-2)[: :-1],dim=-2)
    U_LR = torch.cat(torch.chunk(affine_U_tensors_L_R,2,dim=-1)[: :-1],dim=-1)
    
    b_,c_,h_,w_ = U_FB.shape
    c_mask = create_circular_mask(h_,w_)
   
    FB_mask = np.concatenate((c_mask[h_//2:,:], c_mask[0:h_//2,:]), axis=0)
    LR_mask = np.concatenate((c_mask[:,w_//2:], c_mask[:,0:w_//2]), axis=1)
    fusion_mask = (FB_mask & LR_mask)

    inner_radius = (np.sqrt(2)-1)*(h_//2)
    fusion_mask1 = create_gradient_ring(h_,inner_radius=inner_radius,r1=0,r2=1)
    fusion_mask1_1 = np.concatenate((fusion_mask1[h_//2:,:], fusion_mask1[0:h_//2,:]), axis=0)
    fusion_mask1_2 = np.concatenate((fusion_mask1[:,w_//2:], fusion_mask1[:,0:w_//2]), axis=1)
    fusion_part_1 =  (torch.from_numpy(fusion_mask1_1)*U_LR +torch.from_numpy(1-fusion_mask1_1)*U_FB)*fusion_mask 
    fusion_part_2 =  (torch.from_numpy(1-fusion_mask1_2)*U_LR +torch.from_numpy(fusion_mask1_2)*U_FB)*fusion_mask 
    fusion_part = (fusion_part_1 + fusion_part_2)/2
 
    U_affine =  fusion_part+(torch.from_numpy(1-LR_mask))*U_FB + (torch.from_numpy(1-FB_mask))*U_LR
    # kernel_size = (5, 5)
    # sigma = (2.0, 2.0)
    # transform = T.GaussianBlur(kernel_size=kernel_size, sigma=sigma)
    # U_affine = torch.stack([transform(image) for image in U_affine])
    mask_np = create_circle_line_mask(h_,6) 
    mask_np1 = np.concatenate((mask_np[h_//2:,:], mask_np[0:h_//2,:]), axis=0)
    mask_np2 = np.concatenate((mask_np[:,w_//2:], mask_np[:,0:w_//2]), axis=1)
    mask_np = mask_np1|mask_np2
    mask_np[:h_//4,:] = 0
    mask_np[3*(h_//4):,:] = 0
    mask_np[:,:w_//4] = 0
    mask_np[:,3*(w_//4):] = 0
     # mask_np = np.zeros((h_, w_)).astype(np.uint8)
    # inpaint_size = 8
    # h_off = (h_-inpaint_size)//2
    # w_off = (w_-inpaint_size)//2
    # mask_np[0:inpaint_size,0:inpaint_size] = 1
    # mask_np[(h_-inpaint_size):,0:inpaint_size] = 1
    # mask_np[0:inpaint_size,(w_-inpaint_size):] = 1
    # mask_np[(h_-inpaint_size):,(w_-inpaint_size):] = 1
    # mask_np[h_off:h_off+inpaint_size,w_off:w_off+inpaint_size] = 1

    repaired_images = []
    for i in range(U_affine.shape[0]):
        # 将张量转换为 NumPy 数组并将通道移至最后
        image_np = (U_affine[i].permute(1, 2, 0).cpu().numpy()*255).astype(np.uint8)
        
         
        # 使用 OpenCV 的 inpaint 进行修复
        repaired_np = cv2.inpaint(image_np, mask_np, inpaintRadius=3, flags=cv2.INPAINT_TELEA)

        # 将修复后的结果转换回 PyTorch 张量并将通道移至最前
        repaired_tensor = torch.from_numpy(repaired_np).permute(2, 0, 1)

        # 添加到结果列表中
        repaired_images.append(repaired_tensor)
    U_affine = torch.stack(repaired_images)/255.0

    return U_affine

def create_gradient_ring(image_size, inner_radius, r1, r2):
    # Create coordinate grids
    y, x = np.ogrid[:image_size, :image_size]
    center = (image_size//2,image_size//2)
    outer_radius = image_size//2
    # Calculate Euclidean distance from each point to the center
    distance_from_center = np.sqrt((x - center[0])**2 + (y - center[1])**2)
    
    # Create mask with gradient between two circles
    mask = np.zeros((image_size, image_size), dtype=np.float32)
    within_outer_circle = distance_from_center <= outer_radius
    outside_inner_circle = distance_from_center >= inner_radius
    
    # Apply gradient only in the ring section
    ring_section = within_outer_circle & outside_inner_circle
    normalized_distance = (distance_from_center[ring_section] - inner_radius) / (outer_radius - inner_radius)
    mask[ring_section] = r2 + (r1 - r2) * (1 - normalized_distance)
    
    return mask

def create_gradient_circle(image_size, r1, r2):
    # Create coordinate grids
    y, x = np.ogrid[:image_size, :image_size]
    center = (image_size//2,image_size//2)
    radius = image_size//2
    # Calculate Euclidean distance from each point to the center
    distance_from_center = np.sqrt((x - center[0])**2 + (y - center[1])**2)
    
    # Create a mask that defines a circular gradient
    mask = np.zeros((image_size, image_size), dtype=np.float32)
    inside_circle = distance_from_center <= radius
  
    # Interpolate values between r1 (at the perimeter) and r2 (at the center)
    mask[inside_circle] = r1 + (r2 - r1) * (distance_from_center[inside_circle] / radius)
    
    return mask

def create_circle_line_mask(image_size, thickness):
    # Create a blank black image
    mask = np.zeros((image_size, image_size), dtype=np.uint8)
    center = (image_size//2,image_size//2)
    radius = image_size//2
    # Draw a circle on the mask
    cv2.circle(mask, center, radius, (1), thickness=thickness)
    
    return mask

def create_circular_mask(height, width, center=None, radius=None):
    # Default center and radius
    if center is None:  # Compute center if not provided
        center = (int(width / 2), int(height / 2))
    if radius is None:  # Compute radius if not provided
        radius = min(center[0], center[1], width - center[0], height - center[1])

    # Create a grid of coordinates
    Y, X = np.ogrid[:height, :width]

    # Compute the distance from the circle's center
    distance_from_center = np.sqrt((X - center[0])**2 + (Y - center[1])**2)

    # Apply mask
    mask = distance_from_center <= radius

    return mask.astype(np.uint8)

def generate_diamond_mask(size):
    # 创建一个空的掩码，将其初始化为全零（False）
    mask = np.zeros((size, size))
    
    # 计算中间点
    center = size // 2

    for i in range(size):
        for j in range(size):
            # 计算单元格到中心的曼哈顿距离（菱形距离）
            distance = abs(center - i) + abs(center - j)
            
            # 如果距离小于或等于中心的距离，则将该位置设为True
            if distance <= center:
                mask[i, j] = 1

    return mask.astype(np.uint8)

def convert_to_cubemap1(L_F_D_B_R, norm_overlap=True):
    
    L_F_D_B_R = torch.tensor(L_F_D_B_R).permute(0,3,1,2)
    L_F, B_R = torch.chunk(L_F_D_B_R, 2, dim=-2)
    L_, F_ = torch.chunk(L_F, 2, dim=-1)
    B_, R_ = torch.chunk(B_R, 2, dim=-1)
 
    if norm_overlap:
        b,c,h,w = L_.shape
        # L
        L_norm = L_.clone()
        L_norm[:,:,0:h//2,w//2:w] = (L_[:,:,0:h//2,w//2:w] + torch.rot90(F_, k=1, dims=(-2, -1))[:,:,h//2:h,0:w//2])/2
        L_norm[:,:,h//2:h,0:w//2] = (L_[:,:,h//2:h,0:w//2] + torch.rot90(B_, k=-1, dims=(-2, -1))[:,:,0:h//2,w//2:w])/2
        # F
        F_norm = F_.clone()
        F_norm[:,:,0:h//2,0:w//2] = (F_[:,:,0:h//2,0:w//2] + torch.rot90(L_, k=-1, dims=(-2, -1))[:,:,h//2:h,w//2:w])/2
        F_norm[:,:,h//2:h,w//2:w] = (F_[:,:,h//2:h,w//2:w] + torch.rot90(R_, k=1, dims=(-2, -1))[:,:,0:h//2,0:w//2])/2
        # B
        B_norm = B_.clone()
        B_norm[:,:,0:h//2,0:w//2] = (B_[:,:,0:h//2,0:w//2] + torch.rot90(L_, k=1, dims=(-2, -1))[:,:,h//2:h,w//2:w])/2
        B_norm[:,:,h//2:h,w//2:w] = (B_[:,:,h//2:h,w//2:w] + torch.rot90(R_, k=-1, dims=(-2, -1))[:,:,0:h//2,0:w//2])/2
        # R
        R_norm = R_.clone()
        R_norm[:,:,0:h//2,w//2:w] = (R_[:,:,0:h//2,w//2:w] + torch.rot90(F_, k=-1, dims=(-2, -1))[:,:,h//2:h,0:w//2])/2
        R_norm[:,:,h//2:h,0:w//2] = (R_[:,:,h//2:h,0:w//2] + torch.rot90(B_, k=1, dims=(-2, -1))[:,:,0:h//2,w//2:w])/2

        L_ = L_norm
        F_ = F_norm
        B_ = B_norm
        R_ = R_norm
        
    scale = np.sqrt(2)/2
    L_affine = F.grid_sample(L_, get_transform_matrix_1('L',L_.shape,scale,reverse=True), mode='bilinear', align_corners=False)
    F_affine = F.grid_sample(F_, get_transform_matrix_1('F',F_.shape,scale,reverse=True), mode='bilinear', align_corners=False)
    B_affine = F.grid_sample(B_, get_transform_matrix_1('B',B_.shape,scale,reverse=True), mode='bilinear', align_corners=False)
    R_affine = F.grid_sample(R_, get_transform_matrix_1('R',R_.shape,scale,reverse=True), mode='bilinear', align_corners=False)

    scale = np.sqrt(2)/4
    D_affine = F.grid_sample(L_F_D_B_R, get_transform_matrix_1('D',L_F_D_B_R.shape,scale,reverse=True), mode='bilinear', align_corners=False)

    L_ = torch.rot90(torch.flip(L_,dims=[-1]), k=-1, dims=(-2, -1))
    F_ = torch.rot90(torch.flip(F_,dims=[-1]), k=1, dims=(-2, -1))
    B_ = torch.rot90(torch.flip(B_,dims=[-1]), k=1, dims=(-2, -1))
    R_ = torch.rot90(torch.flip(R_,dims=[-1]), k=-1, dims=(-2, -1))
    L_F_U_B_R = torch.cat([torch.cat([L_,F_],dim=-1),torch.cat([B_,R_],dim=-1)],dim=-2)
    L_F_U_B_R = torch.flip(L_F_U_B_R,dims=[-1])
    U_affine = F.grid_sample(L_F_U_B_R, get_transform_matrix_1('U',L_F_U_B_R.shape,scale,reverse=True), mode='bilinear', align_corners=False)
    
    U_affine = post_process_U(U_affine)
   

    
    return torch.cat([F_affine,R_affine,B_affine,L_affine,U_affine,D_affine],dim=-1)

def save_array_as_image(array, filename):
    # Ensure array is in uint8 format if not already
    array = (array * 255).astype(np.uint8)  # Normalize if array values are [0, 1]

    # Convert NumPy array to PIL Image
    image = Image.fromarray(array)

    # Save the image
    image.save(filename)

def generate_pyramid(n):
    # 创建一个空的正方形矩阵
    pyramid = torch.zeros((n, n), dtype=torch.float32)
    
    # 最大的对角线距离
    max_distance = np.sqrt((n - 1)**2 + (n - 1)**2)
    
    for i in range(n):
        for j in range(n):
            # 当前点到左上角 (0,0) 的欧几里德距离
            distance = np.sqrt(i**2 + j**2)
            
            # 计算高度值
            pyramid[i, j] = 1 - (distance / max_distance)
    
    return pyramid

def generate_linear_gradient(n):
    # 创建正方形矩阵
    gradient = torch.zeros((n, n), dtype=torch.float32)
    
    for i in range(n):
        # 线性插值计算从1 (上边) 到 0 (下边) 的值
        gradient[i, :] = 1 - (i / (n - 1))
        
    return (gradient + torch.rot90(gradient))/2

def convert_to_cubemap2(L_F_D_B_R, norm_overlap=True):
    L_F_D_B_R = torch.tensor(L_F_D_B_R).permute(0,3,1,2)
    L_F, B_R = torch.chunk(L_F_D_B_R, 2, dim=-2)
    L_, F_ = torch.chunk(L_F, 2, dim=-1)
    B_, R_ = torch.chunk(B_R, 2, dim=-1)
 
    if norm_overlap:
        b,c,h,w = L_.shape
        weighted_pyramid1 = generate_pyramid(h//2)
        weighted_pyramid2 = generate_linear_gradient(h//2)
 
        # L
        L_norm = L_.clone()
        rot_weighted_pyramid1 = torch.rot90(weighted_pyramid2,k=1)
        rot_weighted_pyramid2 = torch.rot90(weighted_pyramid2,k=-1)
        L_norm[:,:,0:h//2,w//2:w] = rot_weighted_pyramid1*L_[:,:,0:h//2,w//2:w] + (1-rot_weighted_pyramid1)*torch.rot90(F_, k=1, dims=(-2, -1))[:,:,h//2:h,0:w//2]
        L_norm[:,:,h//2:h,0:w//2] = rot_weighted_pyramid2*L_[:,:,h//2:h,0:w//2] + (1-rot_weighted_pyramid2)*torch.rot90(B_, k=-1, dims=(-2, -1))[:,:,0:h//2,w//2:w]
        # F
        F_norm = F_.clone()
        rot_weighted_pyramid1 = torch.rot90(weighted_pyramid2,k=-2)
        rot_weighted_pyramid2 = weighted_pyramid2
        F_norm[:,:,0:h//2,0:w//2] = rot_weighted_pyramid1*F_[:,:,0:h//2,0:w//2] + (1-rot_weighted_pyramid1)*torch.rot90(L_, k=-1, dims=(-2, -1))[:,:,h//2:h,w//2:w]
        F_norm[:,:,h//2:h,w//2:w] = rot_weighted_pyramid2*F_[:,:,h//2:h,w//2:w] + (1-rot_weighted_pyramid2)*torch.rot90(R_, k=1, dims=(-2, -1))[:,:,0:h//2,0:w//2]
        # B
        B_norm = B_.clone()
        rot_weighted_pyramid1 = torch.rot90(weighted_pyramid2,k=-2)
        rot_weighted_pyramid2 = weighted_pyramid2
        B_norm[:,:,0:h//2,0:w//2] = rot_weighted_pyramid1*B_[:,:,0:h//2,0:w//2] + (1-rot_weighted_pyramid1)*torch.rot90(L_, k=1, dims=(-2, -1))[:,:,h//2:h,w//2:w]
        B_norm[:,:,h//2:h,w//2:w] = rot_weighted_pyramid2*B_[:,:,h//2:h,w//2:w] + (1-rot_weighted_pyramid2)*torch.rot90(R_, k=-1, dims=(-2, -1))[:,:,0:h//2,0:w//2]
        # R
        R_norm = R_.clone()
        rot_weighted_pyramid1 = torch.rot90(weighted_pyramid2,k=1)
        rot_weighted_pyramid2 = torch.rot90(weighted_pyramid2,k=-1)
        R_norm[:,:,0:h//2,w//2:w] = rot_weighted_pyramid1*R_[:,:,0:h//2,w//2:w] + (1-rot_weighted_pyramid1)*torch.rot90(F_, k=-1, dims=(-2, -1))[:,:,h//2:h,0:w//2]
        R_norm[:,:,h//2:h,0:w//2] = rot_weighted_pyramid2*R_[:,:,h//2:h,0:w//2] + (1-rot_weighted_pyramid2)*torch.rot90(B_, k=1, dims=(-2, -1))[:,:,0:h//2,w//2:w]

        L_ = L_norm
        F_ = F_norm
        B_ = B_norm
        R_ = R_norm
        
    scale = np.sqrt(2)/2
    
    L_affine = F.grid_sample(L_, get_transform_matrix_1('L',L_.shape,scale,reverse=True), mode='bilinear', align_corners=False)
    F_affine = F.grid_sample(F_, get_transform_matrix_1('F',F_.shape,scale,reverse=True), mode='bilinear', align_corners=False)
    B_affine = F.grid_sample(B_, get_transform_matrix_1('B',B_.shape,scale,reverse=True), mode='bilinear', align_corners=False)
    R_affine = F.grid_sample(R_, get_transform_matrix_1('R',R_.shape,scale,reverse=True), mode='bilinear', align_corners=False)

    scale = np.sqrt(2)/4
    D_affine = F.grid_sample(L_F_D_B_R, get_transform_matrix_1('D',L_F_D_B_R.shape,scale,reverse=True), mode='bilinear', align_corners=False)

    L_ = torch.rot90(torch.flip(L_,dims=[-1]), k=-1, dims=(-2, -1))
    F_ = torch.rot90(torch.flip(F_,dims=[-1]), k=1, dims=(-2, -1))
    B_ = torch.rot90(torch.flip(B_,dims=[-1]), k=1, dims=(-2, -1))
    R_ = torch.rot90(torch.flip(R_,dims=[-1]), k=-1, dims=(-2, -1))
    L_F_U_B_R = torch.cat([torch.cat([L_,F_],dim=-1),torch.cat([B_,R_],dim=-1)],dim=-2)
    L_F_U_B_R = torch.flip(L_F_U_B_R,dims=[-1])
    U_affine = F.grid_sample(L_F_U_B_R, get_transform_matrix_1('U',L_F_U_B_R.shape,scale,reverse=True), mode='bilinear', align_corners=False)
    U_affine = post_process_U3(U_affine)
    return torch.cat([F_affine,R_affine,B_affine,L_affine,U_affine,D_affine],dim=-1)

