import os
import numpy as np
import time
import sys
np.random.seed(0)
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from utils.reconstruction import ReconstructionUtil
from utils.data_load import DataLoadUtil
from utils.pts import PtsUtil

def save_np_pts(path, pts: np.ndarray, file_type="txt"):
    if file_type == "txt":
        np.savetxt(path, pts)
    else:
        np.save(path, pts)


def save_target_points(root, scene, frame_idx, target_points: np.ndarray, file_type="txt"):
    pts_path = os.path.join(root,scene, "pts", f"{frame_idx}.{file_type}")
    if not os.path.exists(os.path.join(root,scene, "pts")):
        os.makedirs(os.path.join(root,scene, "pts"))
    save_np_pts(pts_path, target_points, file_type)

def save_target_normals(root, scene, frame_idx, target_normals: np.ndarray, file_type="txt"):
    pts_path = os.path.join(root,scene, "nrm", f"{frame_idx}.{file_type}")
    if not os.path.exists(os.path.join(root,scene, "nrm")):
        os.makedirs(os.path.join(root,scene, "nrm"))
    save_np_pts(pts_path, target_normals, file_type)
    
def save_scan_points_indices(root, scene, frame_idx, scan_points_indices: np.ndarray, file_type="txt"):
    indices_path = os.path.join(root,scene, "scan_points_indices", f"{frame_idx}.{file_type}")
    if not os.path.exists(os.path.join(root,scene, "scan_points_indices")):
        os.makedirs(os.path.join(root,scene, "scan_points_indices"))
    save_np_pts(indices_path, scan_points_indices, file_type)

def save_scan_points(root, scene, scan_points: np.ndarray):
    scan_points_path = os.path.join(root,scene, "scan_points.txt")
    save_np_pts(scan_points_path, scan_points)

def get_world_points(depth, mask, cam_intrinsic, cam_extrinsic, random_downsample_N):
    z = depth[mask]
    i, j = np.nonzero(mask)
    x = (j - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0]
    y = (i - cam_intrinsic[1, 2]) * z / cam_intrinsic[1, 1]

    points_camera = np.stack((x, y, z), axis=-1).reshape(-1, 3)
    sampled_target_points = PtsUtil.random_downsample_point_cloud(
                points_camera, random_downsample_N
            )
    points_camera_aug = np.concatenate((sampled_target_points, np.ones((sampled_target_points.shape[0], 1))), axis=-1)
    points_camera_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3]
    
    return points_camera_world

def get_world_points_and_normal(depth, mask, normal, cam_intrinsic, cam_extrinsic, random_downsample_N):
    z = depth[mask]
    i, j = np.nonzero(mask)
    x = (j - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0]
    y = (i - cam_intrinsic[1, 2]) * z / cam_intrinsic[1, 1]

    points_camera = np.stack((x, y, z), axis=-1).reshape(-1, 3)
    normal_camera = normal[mask].reshape(-1, 3)
    sampled_target_points, idx = PtsUtil.random_downsample_point_cloud(
                points_camera, random_downsample_N, require_idx=True
            )
    if len(sampled_target_points) == 0:
        return np.zeros((0, 3)), np.zeros((0, 3))
    sampled_normal_camera  = normal_camera[idx]

    points_camera_aug = np.concatenate((sampled_target_points, np.ones((sampled_target_points.shape[0], 1))), axis=-1)
    points_camera_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3]
    
    return points_camera_world, sampled_normal_camera

def get_scan_points_indices(scan_points, mask, display_table_mask_label, cam_intrinsic, cam_extrinsic):
    scan_points_homogeneous = np.hstack((scan_points, np.ones((scan_points.shape[0], 1))))
    points_camera = np.dot(np.linalg.inv(cam_extrinsic), scan_points_homogeneous.T).T[:, :3]
    points_image_homogeneous = np.dot(cam_intrinsic, points_camera.T).T
    points_image_homogeneous /= points_image_homogeneous[:, 2:]
    pixel_x = points_image_homogeneous[:, 0].astype(int)
    pixel_y = points_image_homogeneous[:, 1].astype(int)
    h, w = mask.shape[:2]
    valid_indices = (pixel_x >= 0) & (pixel_x < w) & (pixel_y >= 0) & (pixel_y < h)
    mask_colors = mask[pixel_y[valid_indices], pixel_x[valid_indices]]
    selected_points_indices = np.where((mask_colors == display_table_mask_label).all(axis=-1))[0]
    selected_points_indices = np.where(valid_indices)[0][selected_points_indices]
    return selected_points_indices
    

def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"):
    
    ''' configuration '''
    target_mask_label = (0, 255, 0)
    display_table_mask_label=(0, 0, 255)
    random_downsample_N = 32768
    voxel_size=0.003
    filter_degree = 75
    min_z = 0.2
    max_z = 0.5
    
    ''' scan points '''
    display_table_info = DataLoadUtil.get_display_table_info(root, scene)
    radius = display_table_info["radius"]
    
    scan_points = np.asarray(ReconstructionUtil.generate_scan_points(display_table_top=0,display_table_radius=radius))
    
    ''' read frame data(depth|mask|normal) '''
    frame_num = DataLoadUtil.get_scene_seq_length(root, scene)
    for frame_id in range(frame_num):
        print(f"[scene({scene_idx}/{scene_total})|frame({frame_id}/{frame_num})]Processing {scene} frame {frame_id}")
        path = DataLoadUtil.get_path(root, scene, frame_id)
        cam_info = DataLoadUtil.load_cam_info(path, binocular=True)
        depth_L, depth_R = DataLoadUtil.load_depth(
                path, cam_info["near_plane"], 
                cam_info["far_plane"], 
                binocular=True
            )
        mask_L, mask_R = DataLoadUtil.load_seg(path, binocular=True)
        normal_L = DataLoadUtil.load_normal(path, binocular=True, left_only=True)
        ''' target points '''
        mask_img_L = mask_L
        mask_img_R = mask_R

        target_mask_img_L = (mask_L == target_mask_label).all(axis=-1)
        target_mask_img_R = (mask_R == target_mask_label).all(axis=-1)

        
        sampled_target_points_L, sampled_target_normal_L = get_world_points_and_normal(depth_L,target_mask_img_L,normal_L, cam_info["cam_intrinsic"], cam_info["cam_to_world"], random_downsample_N)
        sampled_target_points_R = get_world_points(depth_R, target_mask_img_R, cam_info["cam_intrinsic"], cam_info["cam_to_world_R"], random_downsample_N)


        has_points = sampled_target_points_L.shape[0] > 0 and sampled_target_points_R.shape[0] > 0
        if has_points:
            target_points, overlap_idx = PtsUtil.get_overlapping_points(
                    sampled_target_points_L, sampled_target_points_R, voxel_size, require_idx=True
                )
            sampled_target_normal_L = sampled_target_normal_L[overlap_idx]

        if has_points:
            has_points = target_points.shape[0] > 0

        if has_points:
            target_points, target_normals = PtsUtil.filter_points(
                target_points, sampled_target_normal_L, cam_info["cam_to_world"], theta_limit = filter_degree, z_range=(min_z, max_z)
                )


        ''' scan points indices '''
        scan_points_indices_L = get_scan_points_indices(scan_points, mask_img_L, display_table_mask_label, cam_info["cam_intrinsic"], cam_info["cam_to_world"]) 
        scan_points_indices_R = get_scan_points_indices(scan_points, mask_img_R, display_table_mask_label, cam_info["cam_intrinsic"], cam_info["cam_to_world_R"])
        scan_points_indices = np.intersect1d(scan_points_indices_L, scan_points_indices_R)
        
        if not has_points:
            target_points = np.zeros((0, 3))
            target_normals = np.zeros((0, 3))
            
        save_target_points(root, scene, frame_id, target_points, file_type=file_type)
        save_target_normals(root, scene, frame_id, target_normals, file_type=file_type)
        save_scan_points_indices(root, scene, frame_id, scan_points_indices, file_type=file_type)
    
    save_scan_points(root, scene, scan_points) # The "done" flag of scene preprocess


if __name__ == "__main__":
    #root = "/media/hofee/repository/new_data_with_normal"
    root = r"/media/hofee/data/results/ycb_view_data"
    scene_list = os.listdir(root)
    from_idx = 0 # 1000
    to_idx = len(scene_list) # 1500


    cnt = 0
    import time
    total = to_idx - from_idx
    for scene in scene_list[from_idx:to_idx]:
        start = time.time()
        if os.path.exists(os.path.join(root, scene, "scan_points.txt")):
            print(f"Scene {scene} has been processed")
            cnt+=1
            continue
        save_scene_data(root, scene, cnt, total, file_type="npy")
        cnt+=1
        end = time.time()
        print(f"Time cost: {end-start}")
