import numpy as np
from PytorchBoot.dataset import BaseDataset
import PytorchBoot.namespace as namespace
import PytorchBoot.stereotype as stereotype
from PytorchBoot.config import ConfigManager
from PytorchBoot.utils.log_util import Log
import pickle
import torch
import os
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("seq_reconstruction_dataset_preprocessed")
class SeqReconstructionDatasetPreprocessed(BaseDataset):
    def __init__(self, config):
        super(SeqReconstructionDatasetPreprocessed, self).__init__(config)
        self.config = config
        self.root_dir = config["root_dir"]
        self.real_root_dir = r"/media/hofee/data/data/new_testset"
        self.item_list = os.listdir(self.root_dir)

    def __getitem__(self, index):
        data = pickle.load(open(os.path.join(self.root_dir, self.item_list[index]), "rb"))
        data_item = {
            "first_scanned_pts": np.asarray(data["first_scanned_pts"], dtype=np.float32), # Ndarray(S x Nv x 3)
            "first_scanned_n_to_world_pose_9d": np.asarray(data["first_scanned_n_to_world_pose_9d"], dtype=np.float32), # Ndarray(S x 9)
            "seq_max_coverage_rate": data["seq_max_coverage_rate"], # Float, range(0, 1)
            "best_seq_len": data["best_seq_len"], # Int
            "scene_name": data["scene_name"], # String
            "gt_pts": np.asarray(data["gt_pts"], dtype=np.float32), # Ndarray(N x 3)
            "scene_path": os.path.join(self.real_root_dir, data["scene_name"]), # String
            "O_to_L_pose": np.asarray(data["O_to_L_pose"], dtype=np.float32),
        }
        return data_item

    def __len__(self):
        return len(self.item_list)

# -------------- Debug ---------------- #
if __name__ == "__main__":
    import torch

    seed = 0
    torch.manual_seed(seed)
    np.random.seed(seed)
    '''
    OmniObject3d_test:
    root_dir: "H:\\AI\\Datasets\\packed_test_data"
    model_dir: "H:\\AI\\Datasets\\scaled_object_meshes"
    source: seq_reconstruction_dataset
    split_file: "H:\\AI\\Datasets\\data_list\\OmniObject3d_test.txt"
    type: test
    filter_degree: 75
    eval_list:
      - pose_diff
      - coverage_rate_increase
    ratio: 0.1
    batch_size: 1
    num_workers: 12
    pts_num: 8192
    load_from_preprocess: True
    '''
    config = {
        "root_dir": "/media/hofee/data/data/test_bottle/preprocessed_dataset",
        "source": "seq_reconstruction_dataset",
        "split_file": "H:\\AI\\Datasets\\data_list\\OmniObject3d_test.txt",
        "load_from_preprocess": True,
        "ratio": 1,
        "filter_degree": 75,
        "num_workers": 0,
        "pts_num": 8192,
        "type": "test",
    }
    ds = SeqReconstructionDataset(config)
    print(len(ds))
    print(ds.__getitem__(10))