import numpy as np
import pandas as pd
import sys, os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from util.read_write_model import read_model
import numpy as np
from pathlib import Path
from datasets._base import (Image_Class, Base_Collection, Line3D, Pose,
                            Camera)

strlist2floatlist = lambda strlist: [float(s) for s in strlist]
strlist2intlist = lambda strlist: [int(s) for s in strlist]

class DataCollection(Base_Collection):
    def __init__(self, args:dict, cfg:dict, mode="train")->None:
        super(DataCollection, self).__init__(args, cfg, mode)
        self.gt_3Dmodels_path = self.args.sfm_dir / f"{self.args.dataset}/{self.args.scene}"
        self.SfM_with_depth = self.cfg.train.use_depth # use SfM labels which has been corrected by depth or not
        self.train_imgs = [] # list of train image names
        self.test_imgs = [] # list of test image names
        self.imgname2limapID = {} # map from image name to limap image id
        self.limapID2imgname = {} # map from limap image id to image name
        # load all images 2D & 3D points and create Image_Class objects
        self.imgname2imgclass = {} # map from image name to Image_Class object
        self.load_all_2Dpoints_by_dataset(self.args.dataset)
    
    def load_all_2Dpoints_by_dataset(self, dataset):
        if dataset == "7scenes":
            self.load_all_2Dpoints_7scenes()
        elif dataset == "Cambridge" or dataset == "indoor6":
            self.load_all_2Dpoints_Cambridge()
        elif dataset == "BKC":
            self.load_all_2Dpoints_BKC()
        elif dataset == "custom":
            self.load_all_2Dpoints_custom()
        else:
            raise NotImplemented
        

    def load_all_2Dpoints_7scenes(self):
        # currently used for 7scenes. 
        # load all 2d & 3d points from colmap output
        path_gt_3Dmodels_full = self.gt_3Dmodels_path/"sfm_sift_full"
        if self.SfM_with_depth:
            print("[INFOR] Using SfM labels corrected by depth.")
            path_gt_3Dmodels_train = self.gt_3Dmodels_path/"sfm_superpoint+superglue+depth"
        else:
            path_gt_3Dmodels_train = self.gt_3Dmodels_path/"sfm_superpoint+superglue"
        testlist_path = path_gt_3Dmodels_full/"list_test.txt"
        cameras_all, images_all, _ = read_model(path=path_gt_3Dmodels_full, ext=".bin")
        _, images_train, points3D_train = read_model(path=path_gt_3Dmodels_train, ext=".bin")
        name2id_train = {image.name: i for i, image in images_train.items()}

        if os.path.exists(testlist_path):
            with open(testlist_path, 'r') as f:
                testlist = f.read().rstrip().split('\n')
        else:
            raise ValueError("Error! Input file/directory {0} not found.".format(testlist_path))
        for id_, image in images_all.items():
            img_name = image.name
            self.imgname2imgclass[img_name] = Image_Class(img_name)
            if image.name in testlist:
                # fill data to TEST img classes 
                self.test_imgs.append(img_name)
                self.imgname2imgclass[img_name].pose = Pose(image.qvec, image.tvec)
                self.imgname2imgclass[img_name].camera = Camera(cameras_all[image.camera_id],
                                                                          iscolmap=True)
            else:
                # fill data to TRAIN img classes
                self.train_imgs.append(img_name)
                self.imgname2imgclass[img_name].pose = Pose(image.qvec, image.tvec)
                image_train = images_train[name2id_train[img_name]]
                self.imgname2imgclass[img_name].points2Ds = image_train.xys
                self.imgname2imgclass[img_name].points3Ds = np.stack([points3D_train[ii].xyz if ii != -1 else 
                                np.array([0,0,0]) for ii in image_train.point3D_ids], 0)
                self.imgname2imgclass[img_name].validPoints = np.stack([1 if ii != -1 else 
                                0 for ii in image_train.point3D_ids], 0)
                self.imgname2imgclass[img_name].camera = Camera(cameras_all[image.camera_id],
                                                                          iscolmap=True)
    
    def load_all_2Dpoints_Cambridge(self):
        # load all 2d & 3d points from colmap output
        path_gt_3Dmodels_full = self.gt_3Dmodels_path/"sfm_sift_full"
        
        # load query_list_with_intrinsics.txt 
        query_list_with_intrinsics = self.gt_3Dmodels_path/"query_list_with_intrinsics.txt"
        if not os.path.exists(query_list_with_intrinsics):
            raise ValueError("Error! Input file/directory {0} not found.".format(query_list_with_intrinsics))
        query_list_with_intrinsics = pd.read_csv(query_list_with_intrinsics, sep=" ", header=None)
        # get test dictionary with its intrinsic 
        testimgname2intrinsic = {query_list_with_intrinsics.iloc[i,0]:list(query_list_with_intrinsics.iloc[i,1:]) 
                                   for i in range(len(query_list_with_intrinsics))}
        
        # load id_to_origin_name.txt 
        import json 
        id_to_origin_name = self.gt_3Dmodels_path / "id_to_origin_name.txt"
        with open(id_to_origin_name, 'r') as f:
            id_to_origin_name = json.load(f)

        originalname2newimgname = {}
        for id, originalname in id_to_origin_name.items():
            id = int(id)
            originalname2newimgname[originalname] = "image{0:08d}.png".format(id)
        
        
        # load the camera model from colmap output
        _, images_all, _ = read_model(path=path_gt_3Dmodels_full, ext=".bin")
        path_gt_3Dmodels_train = self.gt_3Dmodels_path/"sfm_superpoint+superglue"
        cameras_train, images_train, points3D_train = read_model(path=path_gt_3Dmodels_train, ext=".bin")
        name2id_train = {image.name: i for i, image in images_train.items()}
        
        for _, image in images_all.items():
            img_name = image.name
            new_img_name = originalname2newimgname[img_name]
            self.imgname2imgclass[new_img_name] = Image_Class(new_img_name)
            if new_img_name in testimgname2intrinsic:
                # fill data to TEST img classes 
                self.test_imgs.append(new_img_name)
                self.imgname2imgclass[new_img_name].pose = Pose(image.qvec, image.tvec)
                self.imgname2imgclass[new_img_name].camera = Camera(testimgname2intrinsic[new_img_name],
                                                                          iscolmap=False)
            else:
                # fill data to TRAIN img classes
                if new_img_name not in name2id_train:
                    continue
                image_train = images_train[name2id_train[new_img_name]]
                if len(image_train.point3D_ids) == 0:
                    continue
                self.train_imgs.append(new_img_name)
                self.imgname2imgclass[new_img_name].pose = Pose(image.qvec, image.tvec)
                self.imgname2imgclass[new_img_name].points2Ds = image_train.xys
                self.imgname2imgclass[new_img_name].points3Ds = np.stack([points3D_train[ii].xyz if ii != -1 else 
                                np.array([0,0,0]) for ii in image_train.point3D_ids], 0)
                self.imgname2imgclass[new_img_name].validPoints = np.stack([1 if ii != -1 else 
                                0 for ii in image_train.point3D_ids], 0)
                self.imgname2imgclass[new_img_name].camera = Camera(cameras_train[image_train.camera_id],
                                                                          iscolmap=True)
    
    def load_all_2Dpoints_BKC(self):
        # load all 2d & 3d points from colmap output
        path_gt_3Dmodels_full = self.gt_3Dmodels_path/"sfm_sift_full"
        path_gt_3Dmodels_train = self.gt_3Dmodels_path/"sfm_superpoint+superglue"
        alltestlist_path = path_gt_3Dmodels_full/"test_all.txt"
        test_path = path_gt_3Dmodels_full/"test_seq5.txt"
        cameras_all, images_all, _ = read_model(path=path_gt_3Dmodels_full, ext=".bin")
        _, images_train, points3D_train = read_model(path=path_gt_3Dmodels_train, ext=".bin")
        name2id_train = {image.name: i for i, image in images_train.items()}

        if os.path.exists(alltestlist_path):
            with open(alltestlist_path, 'r') as f:
                testlistall = f.read().rstrip().split('\n')
        else:
            raise ValueError("Error! Input file/directory {0} not found.".format(alltestlist_path))
        
        if os.path.exists(test_path):
            with open(test_path, 'r') as f:
                testlist = f.read().rstrip().split('\n')
        else:
            raise ValueError("Error! Input file/directory {0} not found.".format(test_path))
        
        for id_, image in images_all.items():
            img_name = image.name
            self.imgname2imgclass[img_name] = Image_Class(img_name)
            if image.name in testlistall:
                if image.name in testlist:
                    # fill data to TEST img classes 
                    self.test_imgs.append(img_name)
                    self.imgname2imgclass[img_name].pose = Pose(image.qvec, image.tvec)
                    self.imgname2imgclass[img_name].camera = Camera(cameras_all[image.camera_id],
                                                                            iscolmap=True)
            else:
                # fill data to TRAIN img classes
                self.train_imgs.append(img_name)
                self.imgname2imgclass[img_name].pose = Pose(image.qvec, image.tvec)
                image_train = images_train[name2id_train[img_name]]
                self.imgname2imgclass[img_name].points2Ds = image_train.xys
                self.imgname2imgclass[img_name].points3Ds = np.stack([points3D_train[ii].xyz if ii != -1 else 
                                np.array([0,0,0]) for ii in image_train.point3D_ids], 0)
                self.imgname2imgclass[img_name].validPoints = np.stack([1 if ii != -1 else 
                                0 for ii in image_train.point3D_ids], 0)
                self.imgname2imgclass[img_name].camera = Camera(cameras_all[image.camera_id],
                                                                          iscolmap=True)
    
    def load_all_2Dpoints_custom(self):
        # load the camera model from colmap output
        path_gt_3Dmodels_train = self.args.sfm_dir
        cameras_train, images_train, points3D_train = read_model(path=path_gt_3Dmodels_train, ext=".bin")
        name2id_train = {image.name: i for i, image in images_train.items()}
        
        test_path = os.path.join(self.args.dataset_dir, 'test_list.txt')
        if os.path.exists(test_path):
            with open(test_path, 'r') as f:
                testlist = f.read().rstrip().split('\n')
        else:
            raise ValueError("Error! Input file/directory {0} not found.".format(test_path))
        
        for _, image in images_train.items():
            img_name = image.name
            self.imgname2imgclass[img_name] = Image_Class(img_name)
            # fill data to TRAIN img classes
            if img_name not in name2id_train:
                continue
            image_train = images_train[name2id_train[img_name]]
            if len(image_train.point3D_ids) == 0:
                continue
            if img_name in testlist:
                self.test_imgs.append(img_name)
            else:
                self.train_imgs.append(img_name)
                
            self.imgname2imgclass[img_name].pose = Pose(image.qvec, image.tvec)
            self.imgname2imgclass[img_name].points2Ds = image_train.xys
            self.imgname2imgclass[img_name].points3Ds = np.stack([points3D_train[ii].xyz if ii != -1 else
                            np.array([0,0,0]) for ii in image_train.point3D_ids], 0)
            self.imgname2imgclass[img_name].validPoints = np.stack([1 if ii != -1 else
                            0 for ii in image_train.point3D_ids], 0)
            self.imgname2imgclass[img_name].camera = Camera(cameras_train[image_train.camera_id],
                                                                    iscolmap=True)