import numpy as np
from PytorchBoot.dataset import BaseDataset
import PytorchBoot.stereotype as stereotype
from torch.nn.utils.rnn import pad_sequence

import sys
sys.path.append(r"C:\Document\Local Project\nbv_rec\nbv_reconstruction")

from utils.data_load import DataLoadUtil
from utils.pose import PoseUtil
from utils.pts import PtsUtil

@stereotype.dataset("nbv_reconstruction_dataset")
class NBVReconstructionDataset(BaseDataset):
    def __init__(self, config):
        super(NBVReconstructionDataset, self).__init__(config)
        self.config = config
        self.root_dir = config["root_dir"]
        self.split_file_path = config["split_file"]
        self.scene_name_list = self.load_scene_name_list()
        self.datalist = self.get_datalist()
        self.pts_num = config["pts_num"]

    def load_scene_name_list(self):
        scene_name_list = []
        with open(self.split_file_path, "r") as f:
            for line in f:
                scene_name = line.strip()
                scene_name_list.append(scene_name)
        return scene_name_list
    
    def get_datalist(self):
        datalist = []
        for scene_name in self.scene_name_list:
            label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name)
            label_data = DataLoadUtil.load_label(label_path)
            for data_pair in label_data["data_pairs"]:
                scanned_views = data_pair[0]
                next_best_view = data_pair[1]
                max_coverage_rate = label_data["max_coverage_rate"]
                datalist.append(
                    {
                        "scanned_views": scanned_views,
                        "next_best_view": next_best_view,
                        "max_coverage_rate": max_coverage_rate,
                        "scene_name": scene_name,
                    }
                )
        return datalist

    def __getitem__(self, index):
        data_item_info = self.datalist[index]
        scanned_views = data_item_info["scanned_views"]
        nbv = data_item_info["next_best_view"]
        max_coverage_rate = data_item_info["max_coverage_rate"]
        scene_name = data_item_info["scene_name"]
        scanned_views_pts, scanned_coverages_rate, scanned_n_to_1_pose = [], [], []
        first_frame_idx = scanned_views[0][0]
        first_frame_to_world = DataLoadUtil.load_cam_info(DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx))["cam_to_world"]
        for view in scanned_views:
            frame_idx = view[0]
            coverage_rate = view[1]
            view_path = DataLoadUtil.get_path(self.root_dir, scene_name, frame_idx)
            depth = DataLoadUtil.load_depth(view_path)
            cam_info = DataLoadUtil.load_cam_info(view_path)
            mask = DataLoadUtil.load_seg(view_path)
            frame_curr_to_world = cam_info["cam_to_world"]
            n_to_1_pose = np.dot(np.linalg.inv(first_frame_to_world), frame_curr_to_world)
            target_point_cloud = DataLoadUtil.get_target_point_cloud(depth, cam_info["cam_intrinsic"], n_to_1_pose, mask)["points_world"]
            downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(target_point_cloud, self.pts_num)
            scanned_views_pts.append(downsampled_target_point_cloud)
            scanned_coverages_rate.append(coverage_rate)          
            n_to_1_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(n_to_1_pose[:3,:3]))
            n_to_1_trans = n_to_1_pose[:3,3]
            n_to_1_9d = np.concatenate([n_to_1_6d, n_to_1_trans], axis=0)
            scanned_n_to_1_pose.append(n_to_1_9d)

        nbv_idx, nbv_coverage_rate = nbv[0], nbv[1]
        nbv_path = DataLoadUtil.get_path(self.root_dir, scene_name, nbv_idx)
        cam_info = DataLoadUtil.load_cam_info(nbv_path)
        best_frame_to_world = cam_info["cam_to_world"]
        best_to_1_pose = np.dot(np.linalg.inv(first_frame_to_world), best_frame_to_world)
        best_to_1_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(best_to_1_pose[:3,:3]))
        best_to_1_trans = best_to_1_pose[:3,3]
        best_to_1_9d = np.concatenate([best_to_1_6d, best_to_1_trans], axis=0)
    
        data_item = {
            "scanned_pts": np.asarray(scanned_views_pts,dtype=np.float32),
            "scanned_coverage_rate": np.asarray(scanned_coverages_rate,dtype=np.float32),
            "scanned_n_to_1_pose_9d": np.asarray(scanned_n_to_1_pose,dtype=np.float32),
            "best_coverage_rate": nbv_coverage_rate,
            "best_to_1_pose_9d": best_to_1_9d,
            "max_coverage_rate": max_coverage_rate,
            "scene_name": scene_name
        }
        
        return data_item

    def __len__(self):
        return len(self.datalist)
    
    def get_collate_fn(self):
        def collate_fn(batch):
            scanned_pts = [item['scanned_pts'] for item in batch]
            scanned_n_to_1_pose_9d = [item['scanned_n_to_1_pose_9d'] for item in batch]
            rest = {}
            for key in batch[0].keys():
                if key in ['scanned_pts', 'scanned_n_to_1_pose_9d']:
                    continue
                if isinstance(batch[0][key], torch.Tensor):
                    rest[key] = torch.stack([item[key] for item in batch])
                elif isinstance(batch[0][key], str):
                    rest[key] = [item[key] for item in batch]
                else:
                    rest[key] = [item[key] for item in batch] 
            return {
                'scanned_pts': scanned_pts,
                'scanned_n_to_1_pose_9d': scanned_n_to_1_pose_9d,
                **rest
            }
        return collate_fn

if __name__ == "__main__":
    import torch
    seed = 0
    torch.manual_seed(seed)
    np.random.seed(seed)
    config = {
        "root_dir": "C:\\Document\\Local Project\\nbv_rec\\data\\sample",
        "split_file": "C:\\Document\\Local Project\\nbv_rec\\data\\OmniObject3d_train.txt",
        "ratio": 0.5,
        "batch_size": 2,
        "num_workers": 0,
        "pts_num": 2048
    }
    ds = NBVReconstructionDataset(config)
    print(len(ds))
    dl = ds.get_loader(shuffle=True)
    for idx, data in enumerate(dl):
        cnt=0
        print(data["scene_name"])
        print(data["scanned_coverage_rate"])
        print(data["best_coverage_rate"])
        for pts in data["scanned_pts"][0]:
            #np.savetxt(f"pts_{cnt}.txt", pts)
            cnt+=1
        #np.savetxt("best_pts.txt", best_pts)
        for key, value in data.items():
            if isinstance(value, torch.Tensor):
                print(key, ":" ,value.shape)
            else:
                print(key, ":" ,len(value))
            if key == "scanned_n_to_1_pose_9d":
                for val in value:
                    print(val.shape)
            if key == "scanned_pts":
                for val in value:
                    print(val.shape)
            
                
        print()