import os
import glob
import pickle
from posixpath import basename
import numpy as np
import h5py
from .base_dumper import BaseDumper

import sys

ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../"))
sys.path.insert(0, ROOT_DIR)
import utils


class scannet(BaseDumper):
    def get_seqs(self):
        self.pair_list = np.loadtxt("../assets/scannet_eval_list.txt", dtype=str)
        self.seq_list = np.unique(
            np.asarray([path.split("/")[0] for path in self.pair_list[:, 0]], dtype=str)
        )
        self.dump_seq, self.img_seq = [], []
        for seq in self.seq_list:
            dump_dir = os.path.join(self.config["feature_dump_dir"], seq)
            cur_img_seq = glob.glob(
                os.path.join(
                    os.path.join(self.config["rawdata_dir"], seq, "img", "*.jpg")
                )
            )
            cur_dump_seq = [
                os.path.join(dump_dir, path.split("/")[-1])
                + "_"
                + self.config["extractor"]["name"]
                + "_"
                + str(self.config["extractor"]["num_kpt"])
                + ".hdf5"
                for path in cur_img_seq
            ]
            self.img_seq += cur_img_seq
            self.dump_seq += cur_dump_seq

    def format_dump_folder(self):
        if not os.path.exists(self.config["feature_dump_dir"]):
            os.mkdir(self.config["feature_dump_dir"])
        for seq in self.seq_list:
            seq_dir = os.path.join(self.config["feature_dump_dir"], seq)
            if not os.path.exists(seq_dir):
                os.mkdir(seq_dir)

    def format_dump_data(self):
        print("Formatting data...")
        self.data = {
            "K1": [],
            "K2": [],
            "R": [],
            "T": [],
            "e": [],
            "f": [],
            "fea_path1": [],
            "fea_path2": [],
            "img_path1": [],
            "img_path2": [],
        }

        for pair in self.pair_list:
            img_path1, img_path2 = pair[0], pair[1]
            seq = img_path1.split("/")[0]
            index1, index2 = int(img_path1.split("/")[-1][:-4]), int(
                img_path2.split("/")[-1][:-4]
            )
            ex1, ex2 = np.loadtxt(
                os.path.join(
                    self.config["rawdata_dir"], seq, "extrinsic", str(index1) + ".txt"
                ),
                dtype=float,
            ), np.loadtxt(
                os.path.join(
                    self.config["rawdata_dir"], seq, "extrinsic", str(index2) + ".txt"
                ),
                dtype=float,
            )
            K1, K2 = np.loadtxt(
                os.path.join(
                    self.config["rawdata_dir"], seq, "intrinsic", str(index1) + ".txt"
                ),
                dtype=float,
            ), np.loadtxt(
                os.path.join(
                    self.config["rawdata_dir"], seq, "intrinsic", str(index2) + ".txt"
                ),
                dtype=float,
            )

            relative_extrinsic = np.matmul(np.linalg.inv(ex2), ex1)
            dR, dt = relative_extrinsic[:3, :3], relative_extrinsic[:3, 3]
            dt /= np.sqrt(np.sum(dt**2))

            e_gt_unnorm = np.reshape(
                np.matmul(
                    np.reshape(
                        utils.evaluation_utils.np_skew_symmetric(
                            dt.astype("float64").reshape(1, 3)
                        ),
                        (3, 3),
                    ),
                    np.reshape(dR.astype("float64"), (3, 3)),
                ),
                (3, 3),
            )
            e_gt = e_gt_unnorm / np.linalg.norm(e_gt_unnorm)
            f_gt_unnorm = np.linalg.inv(K2.T) @ e_gt @ np.linalg.inv(K1)
            f_gt = f_gt_unnorm / np.linalg.norm(f_gt_unnorm)

            self.data["K1"].append(K1), self.data["K2"].append(K2)
            self.data["R"].append(dR), self.data["T"].append(dt)
            self.data["e"].append(e_gt), self.data["f"].append(f_gt)

            dump_seq_dir = os.path.join(self.config["feature_dump_dir"], seq)
            fea_path1, fea_path2 = os.path.join(
                dump_seq_dir,
                img_path1.split("/")[-1]
                + "_"
                + self.config["extractor"]["name"]
                + "_"
                + str(self.config["extractor"]["num_kpt"])
                + ".hdf5",
            ), os.path.join(
                dump_seq_dir,
                img_path2.split("/")[-1]
                + "_"
                + self.config["extractor"]["name"]
                + "_"
                + str(self.config["extractor"]["num_kpt"])
                + ".hdf5",
            )
            self.data["img_path1"].append(img_path1), self.data["img_path2"].append(
                img_path2
            )
            self.data["fea_path1"].append(fea_path1), self.data["fea_path2"].append(
                fea_path2
            )

        self.form_standard_dataset()