
import torch
import math
import numpy as np
from scipy.ndimage import map_coordinates
import cv2
# from .get_cubemap import torch_grid_sample

# def _equirect_facetype(h: int, w: int) -> torch.Tensor:
#     """0F 1R 2B 3L 4U 5D"""

#     int_dtype = torch.int64

#     w_ratio = (w - 1) / w
#     h_ratio = (h - 1) / h

#     tp = torch.roll(
#         torch.arange(4)  # 1
#         .repeat_interleave(w // 4)  # 2 same as np.repeat
#         .unsqueeze(0)
#         .transpose(0, 1)  # 3
#         .repeat(1, h)  # 4
#         .view(-1, h)  # 5
#         .transpose(0, 1),  # 6
#         shifts=3 * w // 8,
#         dims=1,
#     )

#     # Prepare ceil mask
#     mask = torch.zeros((h, w // 4), dtype=torch.bool)
#     idx = torch.linspace(-(math.pi * w_ratio), math.pi * w_ratio, w // 4) / 4
#     idx = h // 2 - torch.round(
#         torch.atan(torch.cos(idx)) * h / (math.pi * h_ratio)
#     )
#     idx = idx.type(int_dtype)
#     for i, j in enumerate(idx):
#         mask[:j, i] = 1
#     mask = torch.roll(torch.cat([mask] * 4, 1), 3 * w // 8, 1)

#     tp[mask] = 4
#     tp[torch.flip(mask, dims=(0,))] = 5

#     return tp.type(int_dtype)


# def create_equi_grid(
#     h_out: int,
#     w_out: int,
#     w_face: int,
#     batch: int,
#     dtype: torch.dtype = torch.float32,
#     device: torch.device = torch.device("cpu"),
# ) -> torch.Tensor:
#     w_ratio = (w_out - 1) / w_out
#     h_ratio = (h_out - 1) / h_out
#     theta = torch.linspace(
#         -(math.pi * w_ratio),
#         math.pi * w_ratio,
#         steps=w_out,
#         dtype=dtype,
#         device=device,
#     )
#     phi = torch.linspace(
#         (math.pi * h_ratio) / 2,
#         -(math.pi * h_ratio) / 2,
#         steps=h_out,
#         dtype=dtype,
#         device=device,
#     )
#     phi, theta = torch.meshgrid([phi, theta], indexing="ij")

#     # Get face id to each pixel: 0F 1R 2B 3L 4U 5D
#     tp = _equirect_facetype(h_out, w_out)

#     # xy coordinate map
#     coor_x = torch.zeros((h_out, w_out), dtype=dtype, device=device)
#     coor_y = torch.zeros((h_out, w_out), dtype=dtype, device=device)

#     # FIXME: there's a bug where left section (3L) has artifacts
#     # on top and bottom
#     # It might have to do with 4U or 5D
#     for i in range(6):
#         mask = tp == i

#         if i < 4:
#             coor_x[mask] = 0.5 * torch.tan(theta[mask] - math.pi * i / 2)
#             coor_y[mask] = (
#                 -0.5
#                 * torch.tan(phi[mask])
#                 / torch.cos(theta[mask] - math.pi * i / 2)
#             )
#         elif i == 4:
#             c = 0.5 * torch.tan(math.pi / 2 - phi[mask])
#             coor_x[mask] = c * torch.sin(theta[mask])
#             coor_y[mask] = c * torch.cos(theta[mask])
#         elif i == 5:
#             c = 0.5 * torch.tan(math.pi / 2 - torch.abs(phi[mask]))
#             coor_x[mask] = c * torch.sin(theta[mask])
#             coor_y[mask] = -c * torch.cos(theta[mask])

#     # Final renormalize
#     coor_x = torch.clamp(
#         torch.clamp(coor_x + 0.5, 0, 1) * w_face, 0, w_face - 1
#     )
#     coor_y = torch.clamp(
#         torch.clamp(coor_y + 0.5, 0, 1) * w_face, 0, w_face - 1
#     )

#     # change x axis of the x coordinate map
#     for i in range(6):
#         mask = tp == i
#         coor_x[mask] = coor_x[mask] + w_face * i

#     # repeat batch
#     coor_x = coor_x.repeat(batch, 1, 1) - 0.5
#     coor_y = coor_y.repeat(batch, 1, 1) - 0.5

#     grid = torch.stack((coor_y, coor_x), dim=-3).to(device)
#     return grid

# def cube2equi(cubemap):
#     b,c,h,w = cubemap.shape # w=6h
#     cubemap = cubemap * 255
    
#     dtype = cubemap.dtype
#     device = cubemap.device
#     grid = create_equi_grid(
#         h_out=2*h,
#         w_out=4*h,
#         w_face=h,
#         batch=b,
#         dtype=dtype,
#         device=device,
#     )
#     out = torch.empty((b, c, 2*h, 4*h), dtype=dtype, device=device)
#     out = torch_grid_sample(
#         img=cubemap, grid=grid, out=out, backend="pure"
#     )
#     out = out / 255
#     return out

def sample_cubefaces(cube_faces, tp, coor_y, coor_x, order):
    cube_faces = cube_faces.copy()
    cube_faces[1] = np.flip(cube_faces[1], 1)
    cube_faces[2] = np.flip(cube_faces[2], 1)
    cube_faces[4] = np.flip(cube_faces[4], 0)

    # Pad up down
    pad_ud = np.zeros((6, 2, cube_faces.shape[2]))
    pad_ud[0, 0] = cube_faces[5, 0, :]
    pad_ud[0, 1] = cube_faces[4, -1, :]
    pad_ud[1, 0] = cube_faces[5, :, -1]
    pad_ud[1, 1] = cube_faces[4, ::-1, -1]
    pad_ud[2, 0] = cube_faces[5, -1, ::-1]
    pad_ud[2, 1] = cube_faces[4, 0, ::-1]
    pad_ud[3, 0] = cube_faces[5, ::-1, 0]
    pad_ud[3, 1] = cube_faces[4, :, 0]
    pad_ud[4, 0] = cube_faces[0, 0, :]
    pad_ud[4, 1] = cube_faces[2, 0, ::-1]
    pad_ud[5, 0] = cube_faces[2, -1, ::-1]
    pad_ud[5, 1] = cube_faces[0, -1, :]
    cube_faces = np.concatenate([cube_faces, pad_ud], 1)

    # Pad left right
    pad_lr = np.zeros((6, cube_faces.shape[1], 2))
    pad_lr[0, :, 0] = cube_faces[1, :, 0]
    pad_lr[0, :, 1] = cube_faces[3, :, -1]
    pad_lr[1, :, 0] = cube_faces[2, :, 0]
    pad_lr[1, :, 1] = cube_faces[0, :, -1]
    pad_lr[2, :, 0] = cube_faces[3, :, 0]
    pad_lr[2, :, 1] = cube_faces[1, :, -1]
    pad_lr[3, :, 0] = cube_faces[0, :, 0]
    pad_lr[3, :, 1] = cube_faces[2, :, -1]
    pad_lr[4, 1:-1, 0] = cube_faces[1, 0, ::-1]
    pad_lr[4, 1:-1, 1] = cube_faces[3, 0, :]
    pad_lr[5, 1:-1, 0] = cube_faces[1, -2, :]
    pad_lr[5, 1:-1, 1] = cube_faces[3, -2, ::-1]
    cube_faces = np.concatenate([cube_faces, pad_lr], 2)

    return map_coordinates(cube_faces, [tp, coor_y, coor_x], order=order, mode='wrap')

def equirect_uvgrid(h, w):
    u = np.linspace(-np.pi, np.pi, num=w, dtype=np.float32)
    v = np.linspace(np.pi, -np.pi, num=h, dtype=np.float32) / 2

    return np.stack(np.meshgrid(u, v), axis=-1)

def equirect_facetype(h, w):
    '''
    0F 1R 2B 3L 4U 5D
    '''
    tp = np.roll(np.arange(4).repeat(w // 4)[None, :].repeat(h, 0), 3 * w // 8, 1)

    # Prepare ceil mask
    mask = np.zeros((h, w // 4), bool)
    idx = np.linspace(-np.pi, np.pi, w // 4) / 4
    idx = h // 2 - np.round(np.arctan(np.cos(idx)) * h / np.pi).astype(int)
    for i, j in enumerate(idx):
        mask[:j, i] = 1
    mask = np.roll(np.concatenate([mask] * 4, 1), 3 * w // 8, 1)

    tp[mask] = 4
    tp[np.flip(mask, 0)] = 5

    return tp.astype(np.int32)

def c2e(cubemap, h, w, mode='bilinear', cube_format='horizon'):
    if mode == 'bilinear':
        order = 1
    elif mode == 'nearest':
        order = 0
    else:
        raise NotImplementedError('unknown mode')

    if cube_format == 'horizon':
        pass
    elif cube_format == 'list':
        cubemap = utils.cube_list2h(cubemap)
    elif cube_format == 'dict':
        cubemap = utils.cube_dict2h(cubemap)
    elif cube_format == 'dice':
        cubemap = utils.cube_dice2h(cubemap)
    else:
        raise NotImplementedError('unknown cube_format')
    assert len(cubemap.shape) == 3
    assert cubemap.shape[0] * 6 == cubemap.shape[1]
    assert w % 8 == 0
    face_w = cubemap.shape[0] # h 6*h 3
     
    uv = equirect_uvgrid(h, w)
    u, v = np.split(uv, 2, axis=-1)
    u = u[..., 0]
    v = v[..., 0]
    cube_faces = np.stack(np.split(cubemap, 6, 1), 0)

    # Get face id to each pixel: 0F 1R 2B 3L 4U 5D
    tp = equirect_facetype(h, w)
    coor_x = np.zeros((h, w))
    coor_y = np.zeros((h, w))

    for i in range(4):
        mask = (tp == i)
        coor_x[mask] = 0.5 * np.tan(u[mask] - np.pi * i / 2)
        coor_y[mask] = -0.5 * np.tan(v[mask]) / np.cos(u[mask] - np.pi * i / 2)

    mask = (tp == 4)
    c = 0.5 * np.tan(np.pi / 2 - v[mask])
    coor_x[mask] = c * np.sin(u[mask])
    coor_y[mask] = c * np.cos(u[mask])

    mask = (tp == 5)
    c = 0.5 * np.tan(np.pi / 2 - np.abs(v[mask]))
    coor_x[mask] = c * np.sin(u[mask])
    coor_y[mask] = -c * np.cos(u[mask])

    # Final renormalize
    coor_x = (np.clip(coor_x, -0.5, 0.5) + 0.5) * face_w
    coor_y = (np.clip(coor_y, -0.5, 0.5) + 0.5) * face_w

    equirec = np.stack([
        sample_cubefaces(cube_faces[..., i], tp, coor_y, coor_x, order=order)
        for i in range(cube_faces.shape[3])
    ], axis=-1)

    return equirec

def inpainting_mask(equi_h,equi_w):
    cube_h = equi_h//2
    cube = np.zeros((cube_h, 6*cube_h, 3), dtype=np.uint8)
    U = np.zeros((cube_h, cube_h, 3), dtype=np.uint8)
 
    radius = 32
    center = (radius,radius)
 
    circle = np.zeros((2*radius, 2*radius, 3), dtype=np.uint8)
#    # U[:,:4,:]= 1
    cv2.circle(circle, center, radius, (1, 1, 1), -1)
    circle = 1-circle
    U[:radius,:radius,:] = circle[:radius,:radius,:]
    U[:radius,cube_h-radius:,:] = circle[:radius,radius:,:]
    U[cube_h-radius:,:radius,:] = circle[radius:,:radius,:]
    U[cube_h-radius:,cube_h-radius:,:] = circle[radius:,radius:,:]

    mask_len = cube_h//3
    U[:mask_len,:4,:] = 1
    U[-mask_len:,:4,:] = 1
    U[:mask_len,-4:,:] = 1
    U[-mask_len:,-4:,:] = 1
    U[:4,:mask_len,:] = 1
    U[:4,-mask_len:,:] = 1
    U[-4:,:mask_len,:] = 1
    U[-4:,-mask_len:,:] = 1
    

    cube[:,cube_h*4:cube_h*5,:] = U
    equi = c2e(cube,equi_h,equi_w)
    equi=equi[:,:,0]
    equi[2:,:] = equi[:-2,:].copy()

    return equi
    
def cube2equi(cubemap):
    
    b,c,h,w = cubemap.shape # w=6h
    # alignment 
    cubemap[:,:,:,h:2*h] = torch.flip(cubemap[:,:,:,h:2*h],dims=[-1])
    cubemap[:,:,:,2*h:3*h] = torch.flip(cubemap[:,:,:,2*h:3*h],dims=[-1])
    cubemap[:,:,:,4*h:5*h] = torch.rot90(torch.flip(cubemap[:,:,:,4*h:5*h],dims=[-1]), k=2, dims=(-2, -1))

    equi_h = h*2
    equi_w = h*4
    equi_list = []
    inp_mask = inpainting_mask(equi_h,equi_w).astype(np.uint8)
    
    #equi_list.append(torch.from_numpy(inp_mask).unsqueeze(2).repeat(1,1,3).permute(2,0,1))
    
    for frame in cubemap:
        cube = frame.permute(1,2,0).numpy()
        equi = (c2e(cube,equi_h,equi_w)*255).astype(np.uint8)
      
        equi = cv2.inpaint(equi, inp_mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)
        equi = equi / 255.0
        equi = torch.from_numpy(equi)
        equi_list.append(equi.permute(2,0,1))
    equi = torch.stack(equi_list)
    return equi