import os
import re
import sys
import numpy as np
import torch
import trimesh
from torch.utils.data import DataLoader


path = os.path.abspath(__file__)
for i in range(4):
    path = os.path.dirname(path)
PROJECT_ROOT = path
sys.path.append(PROJECT_ROOT)

from utils.omni_util import OmniUtil
from utils.view_util import ViewUtil
from runners.preprocessors.object_pose.abstract_object_pose_preprocessor import ObjectPosePreprocessor
from configs.config import ConfigManager

from torch.utils.data import Dataset


class ObjectPoseInferenceDataset(Dataset):
    CAMERA_PARAMS_TEMPLATE = "camera_params_{}.json"
    DISTANCE_TEMPLATE = "distance_to_camera_{}.npy"
    RGB_TEMPLATE = "rgb_{}.png"
    MASK_TEMPLATE = "semantic_segmentation_{}.png"
    MASK_LABELS_TEMPLATE = "semantic_segmentation_labels_{}.json"

    def __init__(
        self,
        source="nbv1",
        data_type="sample",
        data_dir="/mnt/h/AI/Datasets",
    ):

        self.data_dir = data_dir
        self.empty_frame = set()
        self.data_path = str(os.path.join(self.data_dir, source, data_type))
        self.scene_list = os.listdir(self.data_path)
        self.data_list = self.get_datalist()
        
        self.object_data_list = self.get_object_datalist()
        self.object_name_list = list(self.object_data_list.keys())
        self.mesh_dir_path = os.path.join(self.data_dir, source,  "objects")
        
        self.meshes = {}
        self.load_all_meshes()
        
    def __len__(self):
        return len(self.data_list)

    def __getitem__(self, index):
        frame_path, target = self.data_list[index]
        frame_data = self.load_frame_data(frame_path=frame_path, object_name=target)
        return frame_data

    def load_all_meshes(self):
        object_name_list = os.listdir(self.mesh_dir_path)
        for object_name in object_name_list:
            mesh_path = os.path.join(self.mesh_dir_path, object_name, "Scan", "Simp.obj")
            mesh = trimesh.load(mesh_path)
            object_model_scale = [0.001, 0.001, 0.001]
            mesh.apply_scale(object_model_scale)
            self.meshes[object_name] = mesh
    
    def get_datalist(self):
        for scene in self.scene_list:
            scene_path = os.path.join(self.data_path, scene)
            file_list = os.listdir(scene_path)
            scene_frame_list = []
            for file in file_list:
                if file.startswith("camera_params"):
                    frame_index = re.findall(r"\d+", file)[0]
                    frame_path = os.path.join(scene_path, frame_index)
                    target_list = OmniUtil.get_object_list(frame_path)
                    for target in target_list:
                        scene_frame_list.append((frame_path,target))
                    if len(target_list) == 0:
                        self.empty_frame.add(frame_path)

        return scene_frame_list
    
    def get_object_datalist(self):
        object_datalist = {}
        for data_item in self.data_list:
            frame_path, target = data_item
            if target not in object_datalist:
                object_datalist[target] = []
            object_datalist[target].append(frame_path)
        return object_datalist
    
    def get_object_data_batch(self, object_name):
        object_data_list = self.object_data_list[object_name]
        batch_data = {"frame_path_list":[],
                      "rgb_batch":[],
                      "depth_batch":[],
                      "seg_batch":[],
                      "gt_pose_batch":[],
                      "K":None,
                      "mesh":None}
        for frame_path in object_data_list:
            frame_data = self.load_frame_data(frame_path, object_name)
            batch_data["frame_path_list"].append(frame_path)
            batch_data["rgb_batch"].append(frame_data["rgb"])
            batch_data["depth_batch"].append(frame_data["depth"])
            batch_data["seg_batch"].append(frame_data["seg"])
            batch_data["gt_pose_batch"].append(frame_data["gt_pose"])
            batch_data["K"] = frame_data["K"]
            batch_data["mesh"] = frame_data["mesh"]
        
        batch_data["rgb_batch"] = np.asarray(batch_data["rgb_batch"],dtype=np.uint8)
        batch_data["depth_batch"] = np.asarray(batch_data["depth_batch"])
        batch_data["seg_batch"] = np.asarray(batch_data["seg_batch"])
        batch_data["gt_pose_batch"] = np.asarray(batch_data["gt_pose_batch"])
        return batch_data
    
    def load_frame_data(self, frame_path, object_name):
        rgb = OmniUtil.get_rgb(frame_path)
        depth = OmniUtil.get_depth(frame_path)
        seg = OmniUtil.get_single_seg(frame_path, object_name)
        K = OmniUtil.get_intrinsic_matrix(frame_path)
        gt_obj_pose = OmniUtil.get_o2c_pose(frame_path, object_name)
        ret_dict  = {
            "frame_path": frame_path,
            "rgb": rgb.astype(np.float32),
            "depth": depth.astype(np.float32),
            "seg": seg,
            "K": K.astype(np.float32),
            "object_name": object_name,
            "mesh": self.meshes[object_name],
            "gt_pose": gt_obj_pose.astype(np.float32)
        }
        return ret_dict
    
class FoundationPosePreprocessor(ObjectPosePreprocessor):
    
    def __init__(self, config_path):
        super().__init__(config_path)
        
    def run(self):
        for dataset_config in self.dataset_list_config:
            dataset = ObjectPoseInferenceDataset( 
                source=dataset_config["source"],
                data_type=dataset_config["data_type"],
                data_dir=dataset_config["data_dir"],
                )
            result = self.prediction(dataset)
            self.save_processed_data(result, dataset_config)
            
    def prediction(self, dataset):
        final_result = {}
        cnt = 0
        for object_name in dataset.object_name_list:
            cnt	+= 1
            print(f"Processing object: {object_name} ({cnt}/{len(dataset.object_name_list)})")
            object_data_batch = dataset.get_object_data_batch(object_name)
            print(f"batch size of object {object_name}: {len(object_data_batch['frame_path_list'])}")
            pose_batch, result_batch = ViewUtil.get_object_pose_batch(
                object_data_batch["K"],
                object_data_batch["mesh"],
                object_data_batch["rgb_batch"],
                object_data_batch["depth_batch"],
                object_data_batch["seg_batch"],
                object_data_batch["gt_pose_batch"],
                self.web_server_config["port"]
            )
            for frame_path, pred_pose,gt_pose,result in zip(object_data_batch["frame_path_list"], pose_batch,object_data_batch["gt_pose_batch"],result_batch):
                if frame_path not in final_result:
                    final_result[frame_path]={}
                final_result[frame_path][object_name] = {"gt_pose":gt_pose.tolist(),"pred_pose":pred_pose.tolist(),"eval_result":result}
        for frame_path in dataset.empty_frame:
            final_result[frame_path] = {}
        return final_result
            
if __name__ == "__main__":
    config_path = os.path.join(PROJECT_ROOT, "configs/server_object_preprocess_config.yaml")
    preprocessor = FoundationPosePreprocessor(config_path)
    preprocessor.run()
    
    
    
    
    