import os
import glob
import pickle
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 yfcc(BaseDumper):
    def get_seqs(self):
        data_dir = os.path.join(self.config["rawdata_dir"], "yfcc100m")
        for seq in self.config["data_seq"]:
            for split in self.config["data_split"]:
                split_dir = os.path.join(data_dir, seq, split)
                dump_dir = os.path.join(self.config["feature_dump_dir"], seq, split)
                cur_img_seq = glob.glob(os.path.join(split_dir, "images", "*.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.config["data_seq"]:
            seq_dir = os.path.join(self.config["feature_dump_dir"], seq)
            if not os.path.exists(seq_dir):
                os.mkdir(seq_dir)
            for split in self.config["data_split"]:
                split_dir = os.path.join(seq_dir, split)
                if not os.path.exists(split_dir):
                    os.mkdir(split_dir)

    def format_dump_data(self):
        print("Formatting data...")
        pair_path = os.path.join(self.config["rawdata_dir"], "pairs")
        self.data = {
            "K1": [],
            "K2": [],
            "R": [],
            "T": [],
            "e": [],
            "f": [],
            "fea_path1": [],
            "fea_path2": [],
            "img_path1": [],
            "img_path2": [],
        }

        for seq in self.config["data_seq"]:
            pair_name = os.path.join(pair_path, seq + "-te-1000-pairs.pkl")
            with open(pair_name, "rb") as f:
                pairs = pickle.load(f)

            # generate id list
            seq_dir = os.path.join(self.config["rawdata_dir"], "yfcc100m", seq, "test")
            name_list = np.loadtxt(os.path.join(seq_dir, "images.txt"), dtype=str)
            cam_name_list = np.loadtxt(
                os.path.join(seq_dir, "calibration.txt"), dtype=str
            )

            for cur_pair in pairs:
                index1, index2 = cur_pair[0], cur_pair[1]
                cam1, cam2 = h5py.File(
                    os.path.join(seq_dir, cam_name_list[index1]), "r"
                ), h5py.File(os.path.join(seq_dir, cam_name_list[index2]), "r")
                K1, K2 = cam1["K"][()], cam2["K"][()]
                [w1, h1], [w2, h2] = cam1["imsize"][()][0], cam2["imsize"][()][0]
                cx1, cy1, cx2, cy2 = (
                    (w1 - 1.0) * 0.5,
                    (h1 - 1.0) * 0.5,
                    (w2 - 1.0) * 0.5,
                    (h2 - 1.0) * 0.5,
                )
                K1[0, 2], K1[1, 2], K2[0, 2], K2[1, 2] = cx1, cy1, cx2, cy2

                R1, R2, t1, t2 = (
                    cam1["R"][()],
                    cam2["R"][()],
                    cam1["T"][()].reshape([3, 1]),
                    cam2["T"][()].reshape([3, 1]),
                )
                dR = np.dot(R2, R1.T)
                dt = t2 - np.dot(dR, t1)
                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)

                img_path1, img_path2 = os.path.join(
                    "yfcc100m", seq, "test", name_list[index1]
                ), os.path.join("yfcc100m", seq, "test", name_list[index2])
                dump_seq_dir = os.path.join(
                    self.config["feature_dump_dir"], seq, "test"
                )
                fea_path1, fea_path2 = os.path.join(
                    dump_seq_dir,
                    name_list[index1].split("/")[-1]
                    + "_"
                    + self.config["extractor"]["name"]
                    + "_"
                    + str(self.config["extractor"]["num_kpt"])
                    + ".hdf5",
                ), os.path.join(
                    dump_seq_dir,
                    name_list[index2].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()