import os
import sys
import kornia
import torch
import numpy as np
import scipy.io as sio
from math import cos, sin
from pytorch3d.utils import cameras_from_opencv_projection
import torchvision as tv

sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
# from superSMPL.IO import readPC2Frame, readOBJ


class Frame:
    def __init__(self, id, frame_nr, garment):
        self.id = id
        self.frame_nr = frame_nr
        self.garment = garment

    def __hash__(self):
        return hash((self.id, self.frame_nr))

    def __eq__(self, other):
        return (self.id, self.frame_nr) == (other.id, other.frame_nr)


def mask_faces(F, mask):
    F_mask = np.take(mask, F.reshape(-1), axis=0)
    F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3

    V_ind = np.zeros(mask.shape, dtype=np.int32)
    V_ind[mask] = np.arange(mask.sum())
    F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))

    return F_new


def read_frames_file(file):
    frames = []
    with open(file, 'r') as f:
        lines = f.readlines()
        for line in lines:
            id, frame_nr, garment = line.split()
            frames.append(Frame(id, int(frame_nr), garment))
    return frames


def get_normals_from_depth(depth_image, depth_mask, camera_intrinsic):
    depth_dilated = kornia.morphology.dilation(depth_image, torch.ones((3,3), device=depth_image.get_device()), engine='convolution')
    depth_dilated = kornia.morphology.dilation(depth_dilated, torch.ones((3,3), device=depth_image.get_device()), engine='convolution')
    depth_dilated = torch.permute(depth_dilated, [0,2,3,1])
    depth_image = torch.permute(depth_image, [0,2,3,1])
    depth_dilated[depth_mask,:] = depth_image[depth_mask,:]
    depth_dilated = torch.permute(depth_dilated, [0,3,1,2])
    normals = kornia.geometry.depth.depth_to_normals(depth_dilated, camera_intrinsic, normalize_points=True)
    normals = torch.permute(normals, [0,2,3,1]) * depth_mask[:,:,:,None]
    return normals


def resize_image(img, mask_croped, mask_resized, size):
    resizer = tv.transforms.Resize(size, interpolation=tv.transforms.functional.InterpolationMode.BILINEAR)
    dilated = kornia.morphology.dilation(img, torch.ones((3,3), device=img.get_device()), engine='convolution')
    dilated = kornia.morphology.dilation(dilated, torch.ones((3,3), device=img.get_device()), engine='convolution')
    dilated = torch.permute(dilated, [0,2,3,1])
    img = torch.permute(img, [0,2,3,1])
    dilated[mask_croped, :] = img[mask_croped, :]
    dilated = torch.permute(dilated, [0,3,1,2])
    out = resizer(dilated)
    out = out * mask_resized[:,None,:,:]
    return out


def center_crop_resize(img, new_center, size):
    pass


""" CAMERA """
def intrinsic(RES_X=640, RES_Y=480):
	f_mm             = 50 # blender default
	sensor_w_mm      = 36 # blender default
	sensor_h_mm = sensor_w_mm * RES_Y / RES_X

	fx_px = f_mm * RES_X / sensor_w_mm
	fy_px = f_mm * RES_Y / sensor_h_mm

	u = RES_X / 2
	v = RES_Y / 2

	return np.array([[fx_px, 0,     u],
					 [0,     fy_px, v],
					 [0,     0,     1]], np.float32)
				

def extrinsic(camLoc):
	R_w2bc = np.array([[0, 1, 0],
					   [0, 0, 1],
					   [1, 0, 0]], np.float32)
	
	T_w2bc = -1 * R_w2bc.dot(camLoc)
	
	R_bc2cv = np.array([[1,  0,  0],
						[0, -1,  0],
						[0,  0, -1]], np.float32)
	
	R_w2cv = R_bc2cv.dot(R_w2bc)
	T_w2cv = R_bc2cv.dot(T_w2bc)

	return np.concatenate((R_w2cv, T_w2cv[:,None]), axis=1)
	

def proj(camLoc):
	return intrinsic().dot(extrinsic(camLoc))


def get_camera(cam, img_size, device, res_x=640, res_y=480):
    size = np.array(img_size, dtype=np.float32)
    ext = torch.tensor(extrinsic(cam).astype(np.float32), device=device)
    K = torch.tensor(intrinsic(res_x, res_y)[None, :, :], device=device)
    cameras = cameras_from_opencv_projection(R=ext[None, :, :3], tvec=ext[None, :, 3], camera_matrix=K, image_size=torch.tensor(size[None, :], device=device))
    return cameras, K


def filter_info(info, frame_nr):
    frame_specific = set(['poses', 'trans'])
    general = set(['shape', 'camLoc', 'zrot', 'gender'])
    out = dict()
    for key in info.keys():
        if key in frame_specific:
            out[key] = info[key][:, frame_nr]
        elif key in general:
            out[key] = info[key]
    
    return out