import cv2
from cv2 import imread
import numpy as np
from numpy import unique, load
import argparse
import os, os.path as osp
import tqdm
import pdb

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--root', type=str, required=True, help='root dir of one identity')
    parser.add_argument('--cam_file', type=str,  default='camera.npz')

    parser.add_argument('--depth_dir',  help='depth map path', default='depth', type=str)
    parser.add_argument('--out_dir', '-n',  help='outputnormal path', default='normal_depth', type=str)
    parser.add_argument('--mask_dir', help='mask path', default='masks', type=str)
    parser.add_argument('--mask_total_dir',  help='total mask path', default='masks_total', type=str)
    args = parser.parse_args()
    return args

def get_surface_normal_by_depth(depth, fx, fy):
    """
    depth: (h, w) of float, the unit of depth is meter
    K: (3, 3) of float, the depth camere's intrinsic
    """

    dz_dv, dz_du = np.gradient(depth)  # u, v mean the pixel coordinate in the image
    # u*depth = fx*x + cx --> du/dx = fx / dept
    # smooth
    # dz_du = cv2.Sobel(depth, cv2.CV_64F, 1, 0, ksize=5)     
    # dz_dv = cv2.Sobel(depth, cv2.CV_64F, 0, 1, ksize=5)
    du_dx = fx / depth  # x is xyz of camera coordinate
    dv_dy = fy / depth

    dz_dx = dz_du * du_dx
    dz_dy = dz_dv * dv_dy
    # cross-product (1,0,dz_dx)X(0,1,dz_dy) = (-dz_dx, -dz_dy, 1)
    normal_cross = np.dstack((-dz_dx, -dz_dy, np.ones_like(depth)))
    # normalize to unit vector
    normal_unit = normal_cross / np.linalg.norm(normal_cross, axis=2, keepdims=True)
    # set default normal to [0, 0, 1]
    # pdb.set_trace()
    nan_mask = np.isfinite(normal_unit).all(2)
    normal_unit[~np.isfinite(normal_unit).all(2)] = [-1, -1, -1]
    return normal_unit, nan_mask

# def delete_holes(img):
#     im = img.copy()
#     h, w = im.shape[:2]
#     mask = np.zeros((h+2, w+2), dtype=np.uint8)
#     holes = cv2.floodFill(im.copy(), mask, (0, 0), 255)[1]
#     holes = ~holes
#     im[holes == 255] = 255
#     return im
def denoise(img):
    
    kernel = np.ones((11, 11), np.uint8)
    img= cv2.erode(img, kernel, iterations=2)
    img= cv2.dilate(img, kernel, iterations=2)

    return img

def main():
    args = parse_args()
    root = args.root
    cam_file = osp.join(root, args.cam_file)
    depth_dir = osp.join(root, args.depth_dir)
    out_dir = osp.join(root, args.out_dir)
    os.makedirs(out_dir, exist_ok=True)
    mask_dir = osp.join(root, args.mask_dir)
    mask_total_dir = osp.join(root, args.mask_total_dir)
    os.makedirs(mask_total_dir, exist_ok=True)

    cam = np.load(cam_file)
    fx, fy = cam['fx'], cam['fy']

    assert osp.exists(depth_dir), 'depth dir not exists'

    depth_paths = sorted([
        osp.join(depth_dir, f) for f in os.listdir(depth_dir)])
    mask_paths = sorted([osp.join(mask_dir, f) for f in os.listdir(mask_dir)])


    assert depth_paths[0].endswith('.npy'), 'depth dir must be .npy'

    # [-1,1] -> [0, 2]/2 -> [0,1]*255 -> [0,255]
    vis_normal = lambda normal: np.uint8((normal + 1) / 2 * 255)[..., ::-1]
    for idx, depth_p in enumerate(depth_paths):

        # mask is [0, 1]
        mask = cv2.imread(mask_paths[idx], -1)[..., 0] / 255
        assert len(mask.shape) == 2, 'dims of mask are not 2'

        depth = np.load(depth_p) # the unit of depth is meter
        # the val of normal is between [-1, 1]
        normal, nan_mask = get_surface_normal_by_depth(depth, fx, fy)
        mask_total = nan_mask * mask
        #mask_total = denoise(mask_total)
        # normal should be filtered by mask and nan_mask
        normal = vis_normal(normal) * mask_total[..., None]
        
        cv2.imwrite(osp.join(out_dir, f'{idx:05d}.png'), normal)
        cv2.imwrite(osp.join(mask_total_dir, f'{idx:05d}.png'), mask_total.astype(np.uint8) * 255)
        #cv2.imwrite(osp.join(root, 'nan_mask.png'), nan_mask.astype(np.uint8))
        #xxx
        


if __name__ == '__main__':
    #os.makedirs(args.normalPath, exist_ok = True)
    main()
    '''    if args.totalmaskPath:
        os.makedirs(args.totalmaskPath, exist_ok = True)
    for img_name in tqdm.tqdm(sorted(os.listdir(args.depthPath))):
        img_path = os.path.join(args.depthPath, img_name)
        depth = cv2.imread(img_path, -1) / 1000.0
        normal, nan_mask = get_surface_normal_by_depth(depth)
        
        vis_normal = lambda normal: np.uint8((normal + 1) / 2 * 255)[..., ::-1]
        normal = vis_normal(normal)
        if args.maskPath:
            mask_path = os.path.join(args.maskPath, img_name[:-4] + '_mask' + img_name[-4:])
            mask = cv2.imread(mask_path)
            mask = cv2.rotate(mask, cv2.ROTATE_90_CLOCKWISE)
            total_mask = mask[..., 0] * nan_mask
            # cv2.imwrite(os.path.join(args.totalmaskPath, img_name), np.uint8(nan_mask*255))
            # cv2.imwrite(os.path.join(args.totalmaskPath, img_name), total_mask)
            kernel = np.ones((11, 11), np.uint8)
            total_mask = cv2.erode(total_mask, kernel, iterations=1)
            total_mask = cv2.dilate(total_mask, kernel, iterations=1)
            # pdb.set_trace()
            normal = np.uint8(total_mask[..., None]/255) * normal
        
        if args.totalmaskPath:
            # cv2.imwrite(os.path.join(args.totalmaskPath, img_name), delete_holes(total_mask))
            total_mask = cv2.rotate(total_mask, cv2.ROTATE_90_COUNTERCLOCKWISE)
            cv2.imwrite(os.path.join(args.totalmaskPath, img_name), total_mask)
            # pdb.set_trace()

        normal_path = os.path.join(args.normalPath, img_name)
        normal = cv2.rotate(normal, cv2.ROTATE_90_COUNTERCLOCKWISE)
        cv2.imwrite(normal_path, normal)
  
    '''