import open3d as o3d
import numpy as np
import os
import cv2
from evaluation.constants import SCANNET_LABELS, SCANNET_IDS, SCANNET18_LABELS, SCANNET18_IDS, SCANNETPP84_IDS, SCANNETPP84_LABELS, SCANNET20_LABELS, SCANNET20_IDS, ARKIT_LABELS, ARKIT_IDS

class ScanNetDataset:

    def __init__(self, seq_name, root='data/scannet', use_templates=False) -> None:
        self.seq_name = seq_name
        self.use_templates = use_templates
        self.root = os.path.join(root, 'processed', seq_name)
        self.rgb_dir = f'{self.root}/color'
        self.depth_dir = f'{self.root}/depth'
        self.segmentation_dir = f'{self.root}/output/mask'
        self.object_dict_dir = f'{self.root}/output/object'
        self.point_cloud_path = f'{self.root}/{seq_name}_vh_clean_2.ply'
        self.mesh_path = self.point_cloud_path
        self.extrinsics_dir = f'{self.root}/pose'
        self.intrinsic_dir = f'{self.root}/intrinsic'
        self.label_features_dict = None

        self.depth_scale = 1000.0
        self.image_size = self.get_image_size()
        self.depth_size = self.get_depth_shape()


    def get_frame_list(self, stride):
        image_list = os.listdir(self.rgb_dir)
        image_list = sorted(image_list, key=lambda x: int(x.split('.')[0]))

        end = int(image_list[-1].split('.')[0]) + 1
        frame_id_list = [int(a.split('.')[0]) for a in image_list]
        return list(frame_id_list)

    def get_image_size(self):
        image_list = os.listdir(self.rgb_dir)
        image_list = sorted(image_list, key=lambda x: int(x.split('.')[0]))
        image_path = os.path.join(self.rgb_dir, image_list[0])
        image = cv2.imread(image_path)
        return image.shape[:2][::-1]

    def get_depth_shape(self):
        image_list = os.listdir(self.rgb_dir)
        image_list = sorted(image_list, key=lambda x: int(x.split('.')[0]))
        depth_path = os.path.join(self.depth_dir, f"{image_list[0].split('.')[0]}.png")
        depth = cv2.imread(depth_path, -1)
        return depth.shape[:2][::-1]

    def get_intrinsics(self, frame_id):
        intrinsic_path = f'{self.intrinsic_dir}/intrinsic_depth.txt'
        intrinsics = np.loadtxt(intrinsic_path)

        intrinisc_cam_parameters = o3d.camera.PinholeCameraIntrinsic()
        intrinisc_cam_parameters.set_intrinsics(self.image_size[0], self.image_size[1], intrinsics[0, 0], intrinsics[1, 1], intrinsics[0, 2], intrinsics[1, 2])
        return intrinisc_cam_parameters
    

    def get_extrinsic(self, frame_id):
        pose_path = os.path.join(self.extrinsics_dir, str(frame_id) + '.txt')
        pose = np.loadtxt(pose_path)
        return pose
    

    def get_depth(self, frame_id):
        depth_path = os.path.join(self.depth_dir, str(frame_id) + '.png')
        depth = cv2.imread(depth_path, -1)
        depth = depth / self.depth_scale
        depth = depth.astype(np.float32)
        return depth


    def get_rgb(self, frame_id, change_color=True):
        rgb_path = os.path.join(self.rgb_dir, str(frame_id) + '.jpg')
        rgb = cv2.imread(rgb_path)

        if change_color:
            rgb = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB)
        return rgb    


    def get_segmentation(self, frame_id, align_with_depth=False):
        segmentation_path = os.path.join(self.segmentation_dir, f'{frame_id}.png')
        if not os.path.exists(segmentation_path):
            assert False, f"Segmentation not found: {segmentation_path}"
        segmentation = cv2.imread(segmentation_path, cv2.IMREAD_UNCHANGED)
        if align_with_depth:
            segmentation = cv2.resize(segmentation, self.depth_size, interpolation=cv2.INTER_NEAREST)
        return segmentation


    def get_frame_path(self, frame_id):
        rgb_path = os.path.join(self.rgb_dir, str(frame_id) + '.jpg')
        segmentation_path = os.path.join(self.segmentation_dir, f'{frame_id}.png')
        return rgb_path, segmentation_path
    

    def get_label_features(self):
        if self.label_features_dict is None:
            if self.use_templates:
                label_features_dict = np.load(f'data/text_features/scannet_templates.npy', allow_pickle=True).item()
            else:
                label_features_dict = np.load(f'data/text_features/scannet.npy', allow_pickle=True).item()
            self.label_features_dict = label_features_dict
        return self.label_features_dict


    def get_scene_points(self):
        mesh = o3d.io.read_point_cloud(self.point_cloud_path)
        vertices = np.asarray(mesh.points)
        return vertices
    
    
    def get_label_id(self):
        self.class_id = SCANNET_IDS
        self.class_label = SCANNET_LABELS

        self.label2id = {}
        self.id2label = {}
        for label, id in zip(self.class_label, self.class_id):
            self.label2id[label] = id
            self.id2label[id] = label

        return self.label2id, self.id2label
    

class ARKitDataset(ScanNetDataset):
    def __init__(self, seq_name, root='data/arkit_dust3r_posed'):
        super().__init__(seq_name, root)
        self.image_size = self.get_image_size()

    def get_image_size(self):
        image_list = os.listdir(self.rgb_dir)
        image_list = sorted(image_list, key=lambda x: int(x.split('.')[0]))
        image_path = os.path.join(self.rgb_dir, image_list[0])
        image = cv2.imread(image_path)
        return image.shape[:2][::-1]

    def get_frame_list(self, stride):
        image_list = os.listdir(self.rgb_dir)
        image_list = sorted(image_list, key=lambda x: int(x.split('.')[0]))

        end = int(image_list[-1].split('.')[0]) + 1
        frame_id_list = [a.split('.')[0] for a in image_list]
        return list(frame_id_list)
    
    def get_label_id(self):
        self.class_id = ARKIT_IDS
        self.class_label = ARKIT_LABELS

        self.label2id = {}
        self.id2label = {}
        for label, id in zip(self.class_label, self.class_id):
            self.label2id[label] = id
            self.id2label[id] = label

        return self.label2id, self.id2label

    def get_label_features(self):
        label_features_dict = np.load(f'data/text_features/arkit.npy', allow_pickle=True).item()
        return label_features_dict

class ITWDataset(ARKitDataset):

    def get_image_size(self):
        image_list = os.listdir(self.rgb_dir)
        image_list = sorted(image_list, key=lambda x: int(x.split('_')[0]))
        image_path = os.path.join(self.rgb_dir, image_list[0])
        image = cv2.imread(image_path)
        return image.shape[:2][::-1]

    def get_depth_shape(self):
        image_list = os.listdir(self.rgb_dir)
        image_list = sorted(image_list, key=lambda x: int(x.split('_')[0]))
        depth_path = os.path.join(self.depth_dir, f"{image_list[0].split('.')[0]}.png")
        depth = cv2.imread(depth_path, -1)
        return depth.shape[:2][::-1]

    def get_frame_list(self, stride):
        image_list = os.listdir(self.rgb_dir)
        image_list = sorted(image_list, key=lambda x: int(x.split('_')[0]))

        frame_id_list = [a.split('.')[0] for a in image_list]
        return list(frame_id_list)
    
    def get_label_features(self):
        label_features_dict = np.load(f'{self.root}/text_features.npy', allow_pickle=True).item()
        return label_features_dict

    def get_label_id(self):
        text_features = self.get_label_features()

        self.class_label = list(text_features.keys())
        self.class_id = list(range(len(self.class_label)))

        self.label2id = {}
        self.id2label = {}
        for label, id in zip(self.class_label, self.class_id):
            self.label2id[label] = id
            self.id2label[id] = label

        return self.label2id, self.id2label

class WildDataset(ARKitDataset):
    def __init__(self, seq_name, root):
        self.root = os.path.join(root, seq_name)
        self.rgb_dir = f'{self.root}/images'
        self.depth_dir = f'{self.root}/depth'
        self.segmentation_dir = f'{self.root}/output/mask'
        self.object_dict_dir = f'{self.root}/output/object'
        self.point_cloud_path = f'{self.root}/point_cloud.ply'
        self.mesh_path = self.point_cloud_path
        self.extrinsics_dir = f'{self.root}/pose'
        self.intrinsic_dir = f'{self.root}/intrinsic'
        self.label_features_dict = None

        self.depth_scale = 1000.0
        self.image_size = self.get_depth_shape()
        self.depth_size = self.get_depth_shape()

    def get_label_features(self):
        label_features_dict = np.load(f'{self.root}/text_features.npy', allow_pickle=True).item()
        return label_features_dict

    def get_segmentation(self, frame_id, align_with_depth=False):
        segmentation_path = os.path.join(self.segmentation_dir, f'{frame_id}.png')
        if not os.path.exists(segmentation_path):
            assert False, f"Segmentation not found: {segmentation_path}"
        segmentation = cv2.imread(segmentation_path, cv2.IMREAD_UNCHANGED)
        segmentation = cv2.resize(segmentation, self.depth_size, interpolation=cv2.INTER_NEAREST)
        return segmentation
    def get_label_id(self):
        text_features = self.get_label_features()

        self.class_label = list(text_features.keys())
        self.class_id = list(range(len(self.class_label)))

        self.label2id = {}
        self.id2label = {}
        for label, id in zip(self.class_label, self.class_id):
            self.label2id[label] = id
            self.id2label[id] = label

        return self.label2id, self.id2label


class ScannetPP2Dataset(ScanNetDataset):
    def __init__(self, seq_name, root='data/scannetpp_dust3r_posed'):
        super().__init__(seq_name, root)
        self.image_size = self.get_image_size()
        self.depth_size = self.get_depth_shape()

        self.point_cloud_path = f'{self.root}/{seq_name}.ply'

    def get_image_size(self):
        image_list = os.listdir(self.rgb_dir)
        image_list = sorted(image_list, key=lambda x: int(x.split('.')[0].split('_')[1]))
        image_path = os.path.join(self.rgb_dir, image_list[0])
        image = cv2.imread(image_path)
        return image.shape[:2][::-1]

    def get_depth_shape(self):

        image_list = os.listdir(self.rgb_dir)
        image_list = sorted(image_list, key=lambda x: int(x.split('.')[0].split('_')[1]))
        depth_path = os.path.join(self.depth_dir, f"{image_list[0].split('.')[0]}.png")
        depth = cv2.imread(depth_path, -1)
        return depth.shape[:2][::-1]

    def get_frame_list(self, stride):
        image_list = os.listdir(self.rgb_dir)
        image_list = sorted(image_list, key=lambda x: int(x.split('.')[0].split('_')[1]))

        frame_id_list = [a.split('.')[0] for a in image_list]
        return list(frame_id_list)

    def get_segmentation(self, frame_id, align_with_depth=False):
        segmentation_path = os.path.join(self.segmentation_dir, f'{frame_id}.png')
        if not os.path.exists(segmentation_path):
            assert False, f"Segmentation not found: {segmentation_path}"
        segmentation = cv2.imread(segmentation_path, cv2.IMREAD_UNCHANGED)
        segmentation = cv2.resize(segmentation, self.depth_size, interpolation=cv2.INTER_NEAREST)
        return segmentation

    
    def get_label_id(self):
        self.class_id = SCANNETPP84_IDS
        self.class_label = SCANNETPP84_LABELS

        self.label2id = {}
        self.id2label = {}
        for label, id in zip(self.class_label, self.class_id):
            self.label2id[label] = id
            self.id2label[id] = label

        return self.label2id, self.id2label
    
    def get_label_features(self):
        label_features_dict = np.load(f'data/text_features/scannetpp84.npy', allow_pickle=True).item()
        return label_features_dict
    
    def get_depth(self, frame_id):
        depth_path = os.path.join(self.depth_dir, str(frame_id) + '.png')
        depth = cv2.imread(depth_path, -1)
        depth = depth / self.depth_scale
        depth = depth.astype(np.float32)
        return depth

    
    def get_intrinsics(self, frame_id):
        intrinsic_path = f'{self.intrinsic_dir}/intrinsic_depth.txt'
        intrinsics = np.loadtxt(intrinsic_path)

        intrinisc_cam_parameters = o3d.camera.PinholeCameraIntrinsic()
        intrinisc_cam_parameters.set_intrinsics(self.image_size[0], self.image_size[1], intrinsics[0, 0], intrinsics[1, 1], intrinsics[0, 2], intrinsics[1, 2])
        return intrinisc_cam_parameters


class ScanNet18Dataset:

    def __init__(self, seq_name, root='data/scannet') -> None:
        self.seq_name = seq_name
        self.root = os.path.join(root, 'processed', seq_name)
        self.rgb_dir = f'{self.root}/color'
        self.depth_dir = f'{self.root}/depth'
        self.segmentation_dir = f'{self.root}/output/mask'
        self.object_dict_dir = f'{self.root}/output/object'
        self.point_cloud_path = f'{self.root}/{seq_name}.ply'
        self.mesh_path = self.point_cloud_path
        self.extrinsics_dir = f'{self.root}/pose'
        self.intrinsic_dir = f'{self.root}/intrinsic'

        self.depth_scale = 1000.0
        self.image_size = self.get_image_size()
        self.depth_size = self.get_depth_shape()


    def get_frame_list(self, stride):
        image_list = os.listdir(self.rgb_dir)
        image_list = sorted(image_list, key=lambda x: int(x.split('.')[0]))

        end = int(image_list[-1].split('.')[0]) + 1
        frame_id_list = [a.split('.')[0] for a in image_list]
        return list(frame_id_list)

    def get_image_size(self):
        image_list = os.listdir(self.rgb_dir)
        image_list = sorted(image_list, key=lambda x: int(x.split('.')[0]))
        image_path = os.path.join(self.rgb_dir, image_list[0])
        image = cv2.imread(image_path)
        return image.shape[:2][::-1]

    def get_depth_shape(self):
        image_list = os.listdir(self.rgb_dir)
        image_list = sorted(image_list, key=lambda x: int(x.split('.')[0]))
        depth_path = os.path.join(self.depth_dir, f"{image_list[0].split('.')[0]}.png")
        depth = cv2.imread(depth_path, -1)
        return depth.shape[:2][::-1]

    def get_intrinsics(self, frame_id):
        intrinsic_path = f'{self.intrinsic_dir}/intrinsic_depth.txt'
        intrinsics = np.loadtxt(intrinsic_path)

        intrinisc_cam_parameters = o3d.camera.PinholeCameraIntrinsic()
        intrinisc_cam_parameters.set_intrinsics(self.image_size[0], self.image_size[1], intrinsics[0, 0], intrinsics[1, 1], intrinsics[0, 2], intrinsics[1, 2])
        return intrinisc_cam_parameters
    

    def get_extrinsic(self, frame_id):
        pose_path = os.path.join(self.extrinsics_dir, str(frame_id) + '.txt')
        pose = np.loadtxt(pose_path)
        return pose
    

    def get_depth(self, frame_id):
        depth_path = os.path.join(self.depth_dir, str(frame_id) + '.png')
        depth = cv2.imread(depth_path, -1)
        depth = depth / self.depth_scale
        depth = depth.astype(np.float32)
        return depth


    def get_rgb(self, frame_id, change_color=True):
        rgb_path = os.path.join(self.rgb_dir, str(frame_id) + '.jpg')
        rgb = cv2.imread(rgb_path)

        if change_color:
            rgb = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB)
        return rgb    


    def get_segmentation(self, frame_id, align_with_depth=False):
        segmentation_path = os.path.join(self.segmentation_dir, f'{frame_id}.png')
        if not os.path.exists(segmentation_path):
            assert False, f"Segmentation not found: {segmentation_path}"
        segmentation = cv2.imread(segmentation_path, cv2.IMREAD_UNCHANGED)
        segmentation = cv2.resize(segmentation, self.depth_size, interpolation=cv2.INTER_NEAREST)
        return segmentation


    def get_frame_path(self, frame_id):
        rgb_path = os.path.join(self.rgb_dir, str(frame_id) + '.jpg')
        segmentation_path = os.path.join(self.segmentation_dir, f'{frame_id}.png')
        return rgb_path, segmentation_path
    

    def get_label_features(self):
        label_features_dict = np.load(f'data/text_features/scannet18.npy', allow_pickle=True).item()
        return label_features_dict


    def get_scene_points(self):
        mesh = o3d.io.read_point_cloud(self.point_cloud_path)
        vertices = np.asarray(mesh.points)
        return vertices
    
    
    def get_label_id(self):
        self.class_id = SCANNET18_IDS
        self.class_label = SCANNET18_LABELS

        self.label2id = {}
        self.id2label = {}
        for label, id in zip(self.class_label, self.class_id):
            self.label2id[label] = id
            self.id2label[id] = label

        return self.label2id, self.id2label


class ScanNet20Dataset(ScanNet18Dataset):

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.point_cloud_path = f'{self.root}/{self.seq_name}_vh_clean_2.ply'

    def get_label_features(self):
        label_features_dict = np.load(f'/home/jovyan/users/lemeshko/Indoor/MaskClustering/data/text_features/scannet20.npy', allow_pickle=True).item()
        return label_features_dict
    
    
    def get_label_id(self):
        self.class_id = SCANNET20_IDS
        self.class_label = SCANNET20_LABELS

        self.label2id = {}
        self.id2label = {}
        for label, id in zip(self.class_label, self.class_id):
            self.label2id[label] = id
            self.id2label[id] = label

        return self.label2id, self.id2label