import torch
import numpy as np
from torch.nn import functional as F
from einops import rearrange
import math
from . import save
def get_condition_tensors(cubeface, cube_tensors, new_h=None, new_w=None):

    b,n,c,f,h,w, = cube_tensors.shape
    device = cube_tensors.device
    cube_tensors = rearrange(cube_tensors,'b n c f h w -> (b n f) c h w')
    matrix_ = get_transform_matrix_1(cubeface,cube_tensors.shape, new_h, new_w).to(device=device)
    affine_ = F.grid_sample(cube_tensors, matrix_, mode='bilinear', align_corners=False)
    affine_ = rearrange(affine_,'(b f) c h w -> b c f h w',f=f)
    zero_tensor = torch.zeros_like(affine_)
    
    if cubeface=='L':
        row1 = torch.cat([affine_,zero_tensor],dim=-1)
        row2 = torch.zeros_like(row1)
   
    if cubeface=='F':
        row1 = torch.cat([zero_tensor,affine_],dim=-1)
        row2 = torch.zeros_like(row1)
         
    if cubeface=='B':
        row2 = torch.cat([affine_,zero_tensor],dim=-1)
        row1 = torch.zeros_like(row2)
         
    if cubeface=='R':
        row2 = torch.cat([zero_tensor,affine_],dim=-1)
        row1 = torch.zeros_like(row2)
         
    res = torch.cat([row1,row2],dim=-2)
    return res


def get_mask(cubeface, tensors_shape):
    b,c,h,w = tensors_shape
    if cubeface == 'null':
        mask = torch.zeros((b,1,h,w))
    else:
        mask = torch.ones((b,1,h,w))

    matrix_ = get_transform_matrix_1('F', tensors_shape)
    mask = F.grid_sample(mask, matrix_, mode='bilinear',align_corners=False)

    return mask

def get_transform_matrix_1(cubeface, tensors_shape, scale=None, new_h=None, new_w=None, reverse=False):
    if cubeface == 'F':
        angle = -45
    elif cubeface == 'L':
        angle = 45
    elif cubeface == 'B':
        angle = 135
    elif cubeface == 'R':
        angle = -135
    elif cubeface == 'D':
        angle = -45
    elif cubeface == 'U':
        angle = -135  

    if reverse:
        angle = -angle
    b,c,h,w = tensors_shape
    theta = np.radians(angle)
    if scale is None:
        scale = np.sqrt(2)
    if new_h is None:
        new_h = math.ceil(scale*h/8)*8
    if new_w is None:
        new_w = math.ceil(scale*w/8)*8

    cos_theta = np.cos(theta)*(scale)
    sin_theta = np.sin(theta)*(scale) 
    transform_matrix = torch.tensor(
        [[cos_theta, -sin_theta, 0],
        [sin_theta, cos_theta, 0]]
        ).float().unsqueeze(0).repeat(b,1,1)  # (1, 2, 3)
    grid = F.affine_grid(transform_matrix, (b, c, new_h, new_w), align_corners=False)
    return grid

def get_transform_matrix(cubeface, tensors_shape):
    if cubeface == 'F':
        angle = -45
    elif cubeface == 'L':
        angle = 45
    elif cubeface == 'B':
        angle = 135
    elif cubeface == 'R':
        angle = -135
    elif cubeface == 'D':
        angle = -45

    b,c,h,w = tensors_shape
    theta = np.radians(angle)
    scale = np.sqrt(2)
    new_h = math.ceil(scale*h/8)*8
    new_w = math.ceil(scale*w/8)*8

    cos_theta = np.cos(theta)*(scale/3)
    sin_theta = np.sin(theta)*(scale/3) 
    transform_matrix = torch.tensor(
        [[cos_theta, -sin_theta, 0],
        [sin_theta, cos_theta, 0]]
        ).float().unsqueeze(0).repeat(b,1,1)  # (1, 2, 3)
    grid = F.affine_grid(transform_matrix, (b, c, new_h, new_w), align_corners=False)
    return grid


def affine_U(U_tensors):
    # b c h w
    U_tensors = torch.cat(torch.chunk(U_tensors,2,dim=-2)[: :-1],dim=-2)
    b,c,H,W = U_tensors.shape
    
    y, x = torch.meshgrid(torch.arange(H), torch.arange(W), indexing='ij')
    x = x.float()
    y = y.float()
    
    # 坐标系参数
    center = torch.tensor([H/2, W/2], device=x.device)
    
    # 计算相对坐标（考虑图像y轴方向）
    dx = x - center[0]
    dy = center[1] - y  # 反转y轴
    
    # 极坐标转换
    theta = torch.atan2(dy, dx)
    r_prime = torch.sqrt(dx**2 + dy**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_max = (H/2) / (denominator)
    # 有效点判断
    valid_mask = r_prime <= r_max
    
    # 计算原图坐标
    r = torch.where(valid_mask, r_prime * ((H/2) / r_max), 0.0)
    px = center[0] + r * cos_theta
    py = center[1] - r * sin_theta  # 注意y轴反转
    
    # 归一化到[-1, 1]
    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.device)
    affine_U_tensors = F.grid_sample(U_tensors, grid, mode='bilinear', align_corners=False)
    affine_U_tensors = torch.cat(torch.chunk(affine_U_tensors,2,dim=-2)[: :-1],dim=-2)
    return affine_U_tensors

def get_cat_tensors(cube_tensors, cube):  # cube_tensors is a list

    F_tensors = cube_tensors[0]
    R_tensors = cube_tensors[1]
    B_tensors = cube_tensors[2]
    L_tensors = cube_tensors[3]
    U_tensors = cube_tensors[4]
    D_tensors = cube_tensors[5]

    b,c,h,w = F_tensors.shape

    cat_tensors = torch.zeros((b,c,3*h,3*w)).to(device=F_tensors.device, dtype=F_tensors.dtype)
    
    if cube == 'D':
        cat_tensors[:, :, 0*h:1*h, 1*w:2*w] = F_tensors  
        cat_tensors[:, :, 1*h:2*h, 0*w:1*w] = torch.rot90(L_tensors, k=1, dims=(2, 3))
        cat_tensors[:, :, 1*h:2*h, 1*w:2*w] = D_tensors
        cat_tensors[:, :, 1*h:2*h, 2*w:3*w] = torch.rot90(R_tensors, k=-1, dims=(2, 3))
        cat_tensors[:, :, 2*h:3*h, 1*w:2*w] = torch.rot90(B_tensors, k=2, dims=(2, 3))
    
    if cube == 'F':
        cat_tensors[:, :, 0*h:1*h, 1*w:2*w] = affine_U(U_tensors)
        cat_tensors[:, :, 1*h:2*h, 0*w:1*w] = L_tensors
        cat_tensors[:, :, 1*h:2*h, 1*w:2*w] = F_tensors
        cat_tensors[:, :, 1*h:2*h, 2*w:3*w] = R_tensors
        cat_tensors[:, :, 2*h:3*h, 1*w:2*w] = D_tensors
    
    if cube == 'L':
        U_tensors = torch.rot90(U_tensors, k=1, dims=(2, 3))
        cat_tensors[:, :, 0*h:1*h, 1*w:2*w] = affine_U(U_tensors)
        cat_tensors[:, :, 1*h:2*h, 0*w:1*w] = B_tensors
        cat_tensors[:, :, 1*h:2*h, 1*w:2*w] = L_tensors
        cat_tensors[:, :, 1*h:2*h, 2*w:3*w] = F_tensors
        cat_tensors[:, :, 2*h:3*h, 1*w:2*w] = torch.rot90(D_tensors, k=-1, dims=(2, 3))
     
    if cube == 'R':
        U_tensors = torch.rot90(U_tensors, k=-1, dims=(2, 3))
        cat_tensors[:, :, 0*h:1*h, 1*w:2*w] = affine_U(U_tensors)
        cat_tensors[:, :, 1*h:2*h, 0*w:1*w] = F_tensors
        cat_tensors[:, :, 1*h:2*h, 1*w:2*w] = R_tensors
        cat_tensors[:, :, 1*h:2*h, 2*w:3*w] = B_tensors
        cat_tensors[:, :, 2*h:3*h, 1*w:2*w] = torch.rot90(D_tensors, k=1, dims=(2, 3))

    if cube == 'B':
        U_tensors = torch.rot90(U_tensors, k=2, dims=(2, 3))
        cat_tensors[:, :, 0*h:1*h, 1*w:2*w] = affine_U(U_tensors)
        cat_tensors[:, :, 1*h:2*h, 0*w:1*w] = R_tensors
        cat_tensors[:, :, 1*h:2*h, 1*w:2*w] = B_tensors
        cat_tensors[:, :, 1*h:2*h, 2*w:3*w] = L_tensors
        cat_tensors[:, :, 2*h:3*h, 1*w:2*w] = torch.rot90(D_tensors, k=2, dims=(2, 3))
    
    return cat_tensors

def get_affine_tensors(cube_tensors, concat=False):
    b,n,c,f,h,w = cube_tensors[0].shape
    device = cube_tensors[0].device
    F_tensors = rearrange(cube_tensors[0],'b n c f h w -> (b n f) c h w') 
    R_tensors = rearrange(cube_tensors[1],'b n c f h w -> (b n f) c h w') 
    B_tensors = rearrange(cube_tensors[2],'b n c f h w -> (b n f) c h w') 
    L_tensors = rearrange(cube_tensors[3],'b n c f h w -> (b n f) c h w') 
    U_tensors = rearrange(cube_tensors[4],'b n c f h w -> (b n f) c h w') 
    D_tensors = rearrange(cube_tensors[5],'b n c f h w -> (b n f) c h w') 

    rearranged_list = [F_tensors,R_tensors,B_tensors,L_tensors,U_tensors,D_tensors]
    tensor_shape = F_tensors.shape

    cat_F = get_cat_tensors(rearranged_list, 'F')
    matrix_F = get_transform_matrix('F',tensor_shape).to(device=device)
    affine_F = F.grid_sample(cat_F, matrix_F, mode='bilinear', align_corners=False)
    affine_F = rearrange(affine_F,'(b f) c h w -> b c f h w',f=f)

    cat_L = get_cat_tensors(rearranged_list, 'L')
    matrix_L = get_transform_matrix('L',tensor_shape).to(device=device)
    affine_L = F.grid_sample(cat_L, matrix_L, mode='bilinear', align_corners=False)
    affine_L = rearrange(affine_L,'(b f) c h w -> b c f h w',f=f)

    cat_B = get_cat_tensors(rearranged_list, 'B')
    matrix_B = get_transform_matrix('B',tensor_shape).to(device=device)
    affine_B = F.grid_sample(cat_B, matrix_B, mode='bilinear', align_corners=False)
    affine_B = rearrange(affine_B,'(b f) c h w -> b c f h w',f=f)

    cat_R = get_cat_tensors(rearranged_list, 'R')
    matrix_R = get_transform_matrix('R',tensor_shape).to(device=device)
    affine_R = F.grid_sample(cat_R, matrix_R, mode='bilinear', align_corners=False)
    affine_R = rearrange(affine_R,'(b f) c h w -> b c f h w',f=f)

    if concat:
        L_F = torch.cat([affine_L,affine_F],dim=-1)
        B_R = torch.cat([affine_B,affine_R],dim=-1)
        L_F_B_R = torch.cat([L_F,B_R],dim=-2)
    else:
        L_F_B_R = torch.stack([affine_L,affine_F,affine_B,affine_R], dim=1)
    # L_F = torch.cat([affine_L,affine_F],dim=-1)
    # B_R = torch.cat([affine_B,affine_R],dim=-1)

    # tensors = torch.cat([L_F,B_R],dim=-2)

    return L_F_B_R