import numpy as np
import matplotlib.pyplot as plt
import sys
import os
import trimesh

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils.data_load import DataLoadUtil
from utils.pts import PtsUtil

class visualizeUtil:
    
    @staticmethod
    def save_all_cam_pos_and_cam_axis(root, scene, output_dir):
        length = DataLoadUtil.get_scene_seq_length(root, scene)
        all_cam_pos = []
        all_cam_axis = []
        for i in range(length):
            path = DataLoadUtil.get_path(root, scene, i)
            cam_info = DataLoadUtil.load_cam_info(path, binocular=True)
            cam_pose = cam_info["cam_to_world"]
            cam_pos = cam_pose[:3, 3]
            cam_axis = cam_pose[:3, 2] 
            
            num_samples = 10
            sample_points = [cam_pos + 0.02*t * cam_axis for t in range(num_samples)]
            sample_points = np.array(sample_points)
        
            all_cam_pos.append(cam_pos)
            all_cam_axis.append(sample_points)
        
        all_cam_pos = np.array(all_cam_pos)
        all_cam_axis = np.array(all_cam_axis).reshape(-1, 3)
        np.savetxt(os.path.join(output_dir, "all_cam_pos.txt"), all_cam_pos)
        np.savetxt(os.path.join(output_dir, "all_cam_axis.txt"), all_cam_axis)
        
    @staticmethod
    def save_all_combined_pts(root, scene, output_dir):
        length = DataLoadUtil.get_scene_seq_length(root, scene)
        all_combined_pts = []   
        for i in range(length):
            path = DataLoadUtil.get_path(root, scene, i)
            pts = DataLoadUtil.load_from_preprocessed_pts(path,"npy")
            if pts.shape[0] == 0:
                continue
            all_combined_pts.append(pts)
        all_combined_pts = np.vstack(all_combined_pts)
        downsampled_all_pts = PtsUtil.voxel_downsample_point_cloud(all_combined_pts, 0.001)
        np.savetxt(os.path.join(output_dir, "all_combined_pts.txt"), downsampled_all_pts)

    @staticmethod
    def save_seq_cam_pos_and_cam_axis(root, scene, frame_idx_list, output_dir):
        all_cam_pos = []
        all_cam_axis = []
        for i in frame_idx_list:
            path = DataLoadUtil.get_path(root, scene, i)
            cam_info = DataLoadUtil.load_cam_info(path, binocular=True)
            cam_pose = cam_info["cam_to_world"]
            cam_pos = cam_pose[:3, 3]
            cam_axis = cam_pose[:3, 2] 
            
            num_samples = 10
            sample_points = [cam_pos + 0.02*t * cam_axis for t in range(num_samples)]
            sample_points = np.array(sample_points)
        
            all_cam_pos.append(cam_pos)
            all_cam_axis.append(sample_points)
        
        all_cam_pos = np.array(all_cam_pos)
        all_cam_axis = np.array(all_cam_axis).reshape(-1, 3)
        np.savetxt(os.path.join(output_dir, "seq_cam_pos.txt"), all_cam_pos)
        np.savetxt(os.path.join(output_dir, "seq_cam_axis.txt"), all_cam_axis)
        
    @staticmethod
    def save_seq_combined_pts(root, scene, frame_idx_list, output_dir):
        all_combined_pts = []   
        for i in frame_idx_list:
            path = DataLoadUtil.get_path(root, scene, i)
            pts = DataLoadUtil.load_from_preprocessed_pts(path,"npy")
            if pts.shape[0] == 0:
                continue
            all_combined_pts.append(pts)
        all_combined_pts = np.vstack(all_combined_pts)
        downsampled_all_pts = PtsUtil.voxel_downsample_point_cloud(all_combined_pts, 0.001)
        np.savetxt(os.path.join(output_dir, "seq_combined_pts.txt"), downsampled_all_pts)
        
    @staticmethod
    def save_target_mesh_at_world_space(
        root, model_dir, scene_name, display_table_as_world_space_origin=True
    ):
        scene_info = DataLoadUtil.load_scene_info(root, scene_name)
        target_name = scene_info["target_name"]
        transformation = scene_info[target_name]
        if display_table_as_world_space_origin:
            location = transformation["location"] - DataLoadUtil.get_display_table_top(
                root, scene_name
            )
        else:
            location = transformation["location"]
        rotation_euler = transformation["rotation_euler"]
        pose_mat = trimesh.transformations.euler_matrix(*rotation_euler)
        pose_mat[:3, 3] = location

        mesh = DataLoadUtil.load_mesh_at(model_dir, target_name, pose_mat)
        mesh_dir = os.path.join(root, scene_name, "mesh")
        if not os.path.exists(mesh_dir):
            os.makedirs(mesh_dir)
        model_path = os.path.join(mesh_dir, "world_target_mesh.obj")
        mesh.export(model_path)

    @staticmethod
    def save_points_and_normals(root, scene, frame_idx, output_dir, binocular=False):
        target_mask_label = (0, 255, 0, 255)
        path = DataLoadUtil.get_path(root, scene, frame_idx)
        cam_info = DataLoadUtil.load_cam_info(path, binocular=binocular, display_table_as_world_space_origin=False)
        depth = DataLoadUtil.load_depth(
                path, cam_info["near_plane"], 
                cam_info["far_plane"], 
                binocular=binocular,
            )
        if isinstance(depth, tuple):
            depth = depth[0]
            
        mask = DataLoadUtil.load_seg(path, binocular=binocular, left_only=True)
        normal = DataLoadUtil.load_normal(path, binocular=binocular, left_only=True)
        ''' target points '''
        if mask is None:
            target_mask_img = np.ones_like(depth, dtype=bool)
        else:
            target_mask_img = (mask == target_mask_label).all(axis=-1)
        cam_intrinsic = cam_info["cam_intrinsic"]
        z = depth[target_mask_img]
        i, j = np.nonzero(target_mask_img)
        x = (j - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0]
        y = (i - cam_intrinsic[1, 2]) * z / cam_intrinsic[1, 1]
        
        random_downsample_N = 1000

        points_camera = np.stack((x, y, z), axis=-1).reshape(-1, 3)
        normal_camera = normal[target_mask_img].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:
            print("No target points")
            

        sampled_normal_camera  = normal_camera[idx]
        sampled_visualized_normal = []
        sampled_normal_camera[:, 2] = -sampled_normal_camera[:, 2]
        sampled_normal_camera[:, 1] = -sampled_normal_camera[:, 1]
        num_samples = 10
        for i in range(len(sampled_target_points)):
            sampled_visualized_normal.append([sampled_target_points[i] + 0.02*t * sampled_normal_camera[i] for t in range(num_samples)])
            
        sampled_visualized_normal = np.array(sampled_visualized_normal).reshape(-1, 3)
        np.savetxt(os.path.join(output_dir, "target_pts.txt"), sampled_target_points)
        np.savetxt(os.path.join(output_dir, "target_normal.txt"), sampled_visualized_normal)

    @staticmethod
    def save_pts_nrm(root, scene, frame_idx, output_dir, binocular=False):
        path = DataLoadUtil.get_path(root, scene, frame_idx)
        pts_world = DataLoadUtil.load_from_preprocessed_pts(path, "npy")
        nrm_camera = DataLoadUtil.load_from_preprocessed_nrm(path, "npy")
        cam_info = DataLoadUtil.load_cam_info(path, binocular=binocular)
        cam_to_world = cam_info["cam_to_world"]
        nrm_world = nrm_camera @ cam_to_world[:3, :3].T
        visualized_nrm = []
        num_samples = 10
        for i in range(len(pts_world)):
            for t in range(num_samples):
                visualized_nrm.append(pts_world[i] - 0.02 * t * nrm_world[i])
        
        visualized_nrm = np.array(visualized_nrm)
        np.savetxt(os.path.join(output_dir, "nrm.txt"), visualized_nrm)
        np.savetxt(os.path.join(output_dir, "pts.txt"), pts_world)
        
# ------ Debug ------

if __name__ == "__main__":
    root = r"C:\Document\Local Project\nbv_rec\nbv_reconstruction\temp"
    model_dir = r"H:\\AI\\Datasets\\scaled_object_box_meshes"
    scene = "box"
    output_dir = r"C:\Document\Local Project\nbv_rec\nbv_reconstruction\test"
    
    #visualizeUtil.save_all_cam_pos_and_cam_axis(root, scene, output_dir)
    # visualizeUtil.save_all_combined_pts(root, scene, output_dir)
    # visualizeUtil.save_seq_combined_pts(root, scene, [0, 121, 286, 175, 111,366,45,230,232,225,255,17,199,78,60], output_dir)
    # visualizeUtil.save_seq_cam_pos_and_cam_axis(root, scene, [0, 121, 286, 175, 111,366,45,230,232,225,255,17,199,78,60], output_dir)
    # visualizeUtil.save_target_mesh_at_world_space(root, model_dir, scene)
    #visualizeUtil.save_points_and_normals(root, scene,"10", output_dir, binocular=True)
    visualizeUtil.save_pts_nrm(root, scene, "116", output_dir, binocular=True)   
