import numpy as np
import torch
import random


# Reworked so this matches gluPerspective / glm::perspective, using fovy
def perspective(fovx=0.7854, aspect=1.0, n=0.1, f=1000.0, device=None):
    # y = np.tan(fovy / 2)
    x = np.tan(fovx / 2)
    return torch.tensor([[1/x,         0,            0,              0],
                         [  0, -aspect/x,            0,              0],
                         [  0,         0, -(f+n)/(f-n), -(2*f*n)/(f-n)],
                         [  0,         0,           -1,              0]], dtype=torch.float32, device=device)


def translate(x, y, z, device=None):
    return torch.tensor([[1, 0, 0, x],
                         [0, 1, 0, y],
                         [0, 0, 1, z],
                         [0, 0, 0, 1]], dtype=torch.float32, device=device)


def rotate_x(a, device=None):
    s, c = np.sin(a), np.cos(a)
    return torch.tensor([[1, 0,  0, 0],
                         [0, c, -s, 0],
                         [0, s,  c, 0],
                         [0, 0,  0, 1]], dtype=torch.float32, device=device)


def rotate_y(a, device=None):
    s, c = np.sin(a), np.cos(a)
    return torch.tensor([[ c, 0, s, 0],
                         [ 0, 1, 0, 0],
                         [-s, 0, c, 0],
                         [ 0, 0, 0, 1]], dtype=torch.float32, device=device)


def rotate_z(a, device=None):
    s, c = np.sin(a), np.cos(a)
    return torch.tensor([[c, -s, 0, 0],
                         [s,  c, 0, 0],
                         [0,  0, 1, 0],
                         [0,  0, 0, 1]], dtype=torch.float32, device=device)

@torch.no_grad()
def batch_random_rotation_translation(b, t, device=None):
    m = np.random.normal(size=[b, 3, 3])
    m[:, 1] = np.cross(m[:, 0], m[:, 2])
    m[:, 2] = np.cross(m[:, 0], m[:, 1])
    m = m / np.linalg.norm(m, axis=2, keepdims=True)
    m = np.pad(m, [[0, 0], [0, 1], [0, 1]], mode='constant')
    m[:, 3, 3] = 1.0
    m[:, :3, 3] = np.random.uniform(-t, t, size=[b, 3])
    return torch.tensor(m, dtype=torch.float32, device=device)

@torch.no_grad()
def random_rotation_translation(t, device=None):
    m = np.random.normal(size=[3, 3])
    m[1] = np.cross(m[0], m[2])
    m[2] = np.cross(m[0], m[1])
    m = m / np.linalg.norm(m, axis=1, keepdims=True)
    m = np.pad(m, [[0, 1], [0, 1]], mode='constant')
    m[3, 3] = 1.0
    m[:3, 3] = np.random.uniform(-t, t, size=[3])
    return torch.tensor(m, dtype=torch.float32, device=device)


@torch.no_grad()
def random_rotation(device=None):
    m = np.random.normal(size=[3, 3])
    m[1] = np.cross(m[0], m[2])
    m[2] = np.cross(m[0], m[1])
    m = m / np.linalg.norm(m, axis=1, keepdims=True)
    m = np.pad(m, [[0, 1], [0, 1]], mode='constant')
    m[3, 3] = 1.0
    m[:3, 3] = np.array([0,0,0]).astype(np.float32)
    return torch.tensor(m, dtype=torch.float32, device=device)


def dot(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    return torch.sum(x*y, -1, keepdim=True)


def length(x: torch.Tensor, eps: float =1e-20) -> torch.Tensor:
    return torch.sqrt(torch.clamp(dot(x,x), min=eps)) # Clamp to avoid nan gradients because grad(sqrt(0)) = NaN


def safe_normalize(x: torch.Tensor, eps: float =1e-20) -> torch.Tensor:
    return x / length(x, eps)


def lr_schedule(iter, warmup_iter, scheduler_decay):
    if iter < warmup_iter:
        return iter / warmup_iter
    return max(0.0, 10 ** (
            -(iter - warmup_iter) * scheduler_decay)) 


def trans_depth(depth):
    depth = depth[0].detach().cpu().numpy()
    valid = depth > 0
    depth[valid] -= depth[valid].min()
    depth[valid] = ((depth[valid] / depth[valid].max()) * 255)
    return depth.astype('uint8')


def nan_to_num(input, nan=0.0, posinf=None, neginf=None, *, out=None):
    assert isinstance(input, torch.Tensor)
    if posinf is None:
        posinf = torch.finfo(input.dtype).max
    if neginf is None:
        neginf = torch.finfo(input.dtype).min
    assert nan == 0
    return torch.clamp(input.unsqueeze(0).nansum(0), min=neginf, max=posinf, out=out)


def load_item(filepath):
    with open(filepath, 'r') as f:
        items = [name.strip() for name in f.readlines()]
    return set(items)

def load_prompt(filepath):
    uuid2prompt = {}
    with open(filepath, 'r') as f:
        for line in f.readlines():
            list_line = line.split(',')
            uuid2prompt[list_line[0]] = ','.join(list_line[1:]).strip()
    return uuid2prompt

def resize_and_center_image(image_tensor, scale=0.95, c = 0, shift = 0, rgb=False, aug_shift = 0):
    if scale == 1:
        return image_tensor
    B, C, H, W = image_tensor.shape
    new_H, new_W = int(H * scale), int(W * scale)
    resized_image = torch.nn.functional.interpolate(image_tensor, size=(new_H, new_W), mode='bilinear', align_corners=False).squeeze(0)
    background = torch.zeros_like(image_tensor) + c
    start_y, start_x = (H - new_H) // 2, (W - new_W) // 2
    if shift == 0:
        background[:, :, start_y:start_y + new_H, start_x:start_x + new_W] = resized_image
    else:
        for i in range(B):
            randx = random.randint(-shift, shift)
            randy = random.randint(-shift, shift)   
            if rgb == True:
                if i == 0 or i==2 or i==4:
                    randx = 0
                    randy = 0 
            background[i, :, start_y+randy:start_y + new_H+randy, start_x+randx:start_x + new_W+randx] = resized_image[i]
    if aug_shift == 0:
        return background  
    for i in range(B):
        for j in range(C):
            background[i, j, :, :] += (random.random() - 0.5)*2 * aug_shift / 255
    return background 
                               
def get_tri(triview_color, dim = 1, blender=True, c = 0, scale=0.95, shift = 0, fix = False, rgb=False, aug_shift = 0):
    # triview_color: [6,C,H,W]
    # rgb is useful when shift is not 0
    triview_color = resize_and_center_image(triview_color, scale=scale, c = c, shift=shift,rgb=rgb, aug_shift = aug_shift)
    if blender is False:
        triview_color0 = torch.rot90(triview_color[0],k=2,dims=[1,2])
        triview_color1 = torch.rot90(triview_color[4],k=1,dims=[1,2]).flip(2).flip(1)
        triview_color2 = torch.rot90(triview_color[5],k=1,dims=[1,2]).flip(2)
        triview_color3 = torch.rot90(triview_color[3],k=2,dims=[1,2]).flip(2)
        triview_color4 = torch.rot90(triview_color[1],k=3,dims=[1,2]).flip(1)
        triview_color5 = torch.rot90(triview_color[2],k=3,dims=[1,2]).flip(1).flip(2)
    else:
        triview_color0 = torch.rot90(triview_color[2],k=2,dims=[1,2])
        triview_color1 = torch.rot90(triview_color[4],k=0,dims=[1,2]).flip(2).flip(1)
        triview_color2 = torch.rot90(torch.rot90(triview_color[0],k=3,dims=[1,2]).flip(2), k=2,dims=[1,2])
        triview_color3 = torch.rot90(torch.rot90(triview_color[5],k=2,dims=[1,2]).flip(2), k=2,dims=[1,2])
        triview_color4 = torch.rot90(triview_color[1],k=2,dims=[1,2]).flip(1).flip(1).flip(2)
        triview_color5 = torch.rot90(triview_color[3],k=1,dims=[1,2]).flip(1).flip(2)
        if fix == True:
            triview_color0[1] = triview_color0[1] * 0
            triview_color0[2] = triview_color0[2] * 0
            triview_color3[1] = triview_color3[1] * 0
            triview_color3[2] = triview_color3[2] * 0

            triview_color1[0] = triview_color1[0] * 0
            triview_color1[1] = triview_color1[1] * 0
            triview_color4[0] = triview_color4[0] * 0
            triview_color4[1] = triview_color4[1] * 0

            triview_color2[0] = triview_color2[0] * 0
            triview_color2[2] = triview_color2[2] * 0
            triview_color5[0] = triview_color5[0] * 0
            triview_color5[2] = triview_color5[2] * 0
    color_tensor1_gt = torch.cat((triview_color0, triview_color1, triview_color2), dim=2)
    color_tensor2_gt = torch.cat((triview_color3, triview_color4, triview_color5), dim=2)
    color_tensor_gt = torch.cat((color_tensor1_gt, color_tensor2_gt), dim = dim)
    return color_tensor_gt