import os
import json
import time
import subprocess
import tempfile
import shutil
import numpy as np
from utils.data_load import DataLoadUtil
from utils.reconstruction import ReconstructionUtil
from utils.pts import PtsUtil
class RenderUtil:
    target_mask_label = (0, 255, 0)
    display_table_mask_label = (0, 0, 255)
    random_downsample_N = 32768
    min_z = 0.2
    max_z = 0.5

    @staticmethod
    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
    
    @staticmethod
    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
    
    @staticmethod
    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

    @staticmethod
    def render_pts(cam_pose, scene_path, script_path, scan_points, voxel_threshold=0.005, filter_degree=75, nO_to_nL_pose=None, require_full_scene=False):
        #import ipdb; ipdb.set_trace()
        nO_to_world_pose = DataLoadUtil.get_real_cam_O_from_cam_L(cam_pose, nO_to_nL_pose, scene_path=scene_path)

    
        with tempfile.TemporaryDirectory() as temp_dir:
            params = {
                "cam_pose": nO_to_world_pose.tolist(),
                "scene_path": scene_path
            }
            scene_info_path = os.path.join(scene_path, "scene_info.json")
            shutil.copy(scene_info_path, os.path.join(temp_dir, "scene_info.json"))
            params_data_path = os.path.join(temp_dir, "params.json")
            with open(params_data_path, 'w') as f:
                json.dump(params, f) 
            result = subprocess.run([
                '/home/hofee/blender-4.0.2-linux-x64/blender', '-b', '-P', script_path, '--', temp_dir
            ], capture_output=True, text=True)
            #print(result)
            #import ipdb; ipdb.set_trace()
            path = os.path.join(temp_dir, "tmp")
            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 == RenderUtil.target_mask_label).all(axis=-1)
            target_mask_img_R = (mask_R == RenderUtil.target_mask_label).all(axis=-1)

            
            sampled_target_points_L, sampled_target_normal_L = RenderUtil.get_world_points_and_normal(depth_L,target_mask_img_L,normal_L, cam_info["cam_intrinsic"], cam_info["cam_to_world"],  RenderUtil.random_downsample_N)
            sampled_target_points_R = RenderUtil.get_world_points(depth_R, target_mask_img_R, cam_info["cam_intrinsic"], cam_info["cam_to_world_R"],  RenderUtil.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_threshold, 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=(RenderUtil.min_z, RenderUtil.max_z)
                    )
                
            
            scan_points_indices_L = RenderUtil.get_scan_points_indices(scan_points, mask_img_L, RenderUtil.display_table_mask_label, cam_info["cam_intrinsic"], cam_info["cam_to_world"]) 
            scan_points_indices_R = RenderUtil.get_scan_points_indices(scan_points, mask_img_R, RenderUtil.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))
            #import ipdb; ipdb.set_trace()
            return target_points, target_normals, scan_points_indices