import numpy as np
import torch
import torch.utils.data as data
import cv2
import os
import h5py
import random

import sys

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

from utils import train_utils, evaluation_utils

torch.multiprocessing.set_sharing_strategy("file_system")


class Offline_Dataset(data.Dataset):
    def __init__(self, config, mode):
        assert mode == "train" or mode == "valid"

        self.config = config
        self.mode = mode
        metadir = (
            os.path.join(config.dataset_path, "valid")
            if mode == "valid"
            else os.path.join(config.dataset_path, "train")
        )

        pair_num_list = np.loadtxt(os.path.join(metadir, "pair_num.txt"), dtype=str)
        self.total_pairs = int(pair_num_list[0, 1])
        self.pair_seq_list, self.accu_pair_num = train_utils.parse_pair_seq(
            pair_num_list
        )

    def collate_fn(self, batch):
        batch_size, num_pts = len(batch), batch[0]["x1"].shape[0]

        data = {}
        dtype = [
            "x1",
            "x2",
            "kpt1",
            "kpt2",
            "desc1",
            "desc2",
            "num_corr",
            "num_incorr1",
            "num_incorr2",
            "e_gt",
            "pscore1",
            "pscore2",
            "img_path1",
            "img_path2",
        ]
        for key in dtype:
            data[key] = []
        for sample in batch:
            for key in dtype:
                data[key].append(sample[key])

        for key in [
            "x1",
            "x2",
            "kpt1",
            "kpt2",
            "desc1",
            "desc2",
            "e_gt",
            "pscore1",
            "pscore2",
        ]:
            data[key] = torch.from_numpy(np.stack(data[key])).float()
        for key in ["num_corr", "num_incorr1", "num_incorr2"]:
            data[key] = torch.from_numpy(np.stack(data[key])).int()

        # kpt augmentation with random homography
        if self.mode == "train" and self.config.data_aug:
            homo_mat = torch.from_numpy(
                train_utils.get_rnd_homography(batch_size)
            ).unsqueeze(1)
            aug_seed = random.random()
            if aug_seed < 0.5:
                x1_homo = torch.cat(
                    [data["x1"], torch.ones([batch_size, num_pts, 1])], dim=-1
                ).unsqueeze(-1)
                x1_homo = torch.matmul(homo_mat.float(), x1_homo.float()).squeeze(-1)
                data["aug_x1"] = x1_homo[:, :, :2] / x1_homo[:, :, 2].unsqueeze(-1)
                data["aug_x2"] = data["x2"]
            else:
                x2_homo = torch.cat(
                    [data["x2"], torch.ones([batch_size, num_pts, 1])], dim=-1
                ).unsqueeze(-1)
                x2_homo = torch.matmul(homo_mat.float(), x2_homo.float()).squeeze(-1)
                data["aug_x2"] = x2_homo[:, :, :2] / x2_homo[:, :, 2].unsqueeze(-1)
                data["aug_x1"] = data["x1"]
        else:
            data["aug_x1"], data["aug_x2"] = data["x1"], data["x2"]
        return data

    def __getitem__(self, index):
        seq = self.pair_seq_list[index]
        index_within_seq = index - self.accu_pair_num[seq]

        with h5py.File(
            os.path.join(self.config.dataset_path, seq, "info.h5py"), "r"
        ) as data:
            R, t = (
                data["dR"][str(index_within_seq)][()],
                data["dt"][str(index_within_seq)][()],
            )
            egt = np.reshape(
                np.matmul(
                    np.reshape(
                        evaluation_utils.np_skew_symmetric(
                            t.astype("float64").reshape(1, 3)
                        ),
                        (3, 3),
                    ),
                    np.reshape(R.astype("float64"), (3, 3)),
                ),
                (3, 3),
            )
            egt = egt / np.linalg.norm(egt)
            K1, K2 = (
                data["K1"][str(index_within_seq)][()],
                data["K2"][str(index_within_seq)][()],
            )
            size1, size2 = (
                data["size1"][str(index_within_seq)][()],
                data["size2"][str(index_within_seq)][()],
            )

            img_path1, img_path2 = (
                data["img_path1"][str(index_within_seq)][()][0].decode(),
                data["img_path2"][str(index_within_seq)][()][0].decode(),
            )
            img_name1, img_name2 = img_path1.split("/")[-1], img_path2.split("/")[-1]
            img_path1, img_path2 = os.path.join(
                self.config.rawdata_path, img_path1
            ), os.path.join(self.config.rawdata_path, img_path2)
            fea_path1, fea_path2 = os.path.join(
                self.config.desc_path, seq, img_name1 + self.config.desc_suffix
            ), os.path.join(
                self.config.desc_path, seq, img_name2 + self.config.desc_suffix
            )
            with h5py.File(fea_path1, "r") as fea1, h5py.File(fea_path2, "r") as fea2:
                desc1, kpt1, pscore1 = (
                    fea1["descriptors"][()],
                    fea1["keypoints"][()][:, :2],
                    fea1["keypoints"][()][:, 2],
                )
                desc2, kpt2, pscore2 = (
                    fea2["descriptors"][()],
                    fea2["keypoints"][()][:, :2],
                    fea2["keypoints"][()][:, 2],
                )
                kpt1, kpt2, desc1, desc2 = (
                    kpt1[: self.config.num_kpt],
                    kpt2[: self.config.num_kpt],
                    desc1[: self.config.num_kpt],
                    desc2[: self.config.num_kpt],
                )

            # normalize kpt
            if self.config.input_normalize == "intrinsic":
                x1, x2 = np.concatenate(
                    [kpt1, np.ones([kpt1.shape[0], 1])], axis=-1
                ), np.concatenate([kpt2, np.ones([kpt2.shape[0], 1])], axis=-1)
                x1, x2 = (
                    np.matmul(np.linalg.inv(K1), x1.T).T[:, :2],
                    np.matmul(np.linalg.inv(K2), x2.T).T[:, :2],
                )
            elif self.config.input_normalize == "img":
                x1, x2 = (kpt1 - size1 / 2) / size1, (kpt2 - size2 / 2) / size2
                S1_inv, S2_inv = np.asarray(
                    [
                        [size1[0], 0, 0.5 * size1[0]],
                        [0, size1[1], 0.5 * size1[1]],
                        [0, 0, 1],
                    ]
                ), np.asarray(
                    [
                        [size2[0], 0, 0.5 * size2[0]],
                        [0, size2[1], 0.5 * size2[1]],
                        [0, 0, 1],
                    ]
                )
                M1, M2 = np.matmul(np.linalg.inv(K1), S1_inv), np.matmul(
                    np.linalg.inv(K2), S2_inv
                )
                egt = np.matmul(np.matmul(M2.transpose(), egt), M1)
                egt = egt / np.linalg.norm(egt)
            else:
                raise NotImplementedError

            corr = data["corr"][str(index_within_seq)][()]
            incorr1, incorr2 = (
                data["incorr1"][str(index_within_seq)][()],
                data["incorr2"][str(index_within_seq)][()],
            )

        # permute kpt
        valid_corr = corr[corr.max(axis=-1) < self.config.num_kpt]
        valid_incorr1, valid_incorr2 = (
            incorr1[incorr1 < self.config.num_kpt],
            incorr2[incorr2 < self.config.num_kpt],
        )
        num_corr, num_incorr1, num_incorr2 = (
            len(valid_corr),
            len(valid_incorr1),
            len(valid_incorr2),
        )
        mask1_invlaid, mask2_invalid = np.ones(x1.shape[0]).astype(bool), np.ones(
            x2.shape[0]
        ).astype(bool)
        mask1_invlaid[valid_corr[:, 0]] = False
        mask2_invalid[valid_corr[:, 1]] = False
        mask1_invlaid[valid_incorr1] = False
        mask2_invalid[valid_incorr2] = False
        invalid_index1, invalid_index2 = (
            np.nonzero(mask1_invlaid)[0],
            np.nonzero(mask2_invalid)[0],
        )

        # random sample from point w/o valid annotation
        cur_kpt1 = self.config.num_kpt - num_corr - num_incorr1
        cur_kpt2 = self.config.num_kpt - num_corr - num_incorr2

        if invalid_index1.shape[0] < cur_kpt1:
            sub_idx1 = np.concatenate(
                [
                    np.arange(len(invalid_index1)),
                    np.random.randint(
                        len(invalid_index1), size=cur_kpt1 - len(invalid_index1)
                    ),
                ]
            )
        if invalid_index1.shape[0] >= cur_kpt1:
            sub_idx1 = np.random.choice(len(invalid_index1), cur_kpt1, replace=False)
        if invalid_index2.shape[0] < cur_kpt2:
            sub_idx2 = np.concatenate(
                [
                    np.arange(len(invalid_index2)),
                    np.random.randint(
                        len(invalid_index2), size=cur_kpt2 - len(invalid_index2)
                    ),
                ]
            )
        if invalid_index2.shape[0] >= cur_kpt2:
            sub_idx2 = np.random.choice(len(invalid_index2), cur_kpt2, replace=False)

        per_idx1, per_idx2 = np.concatenate(
            [valid_corr[:, 0], valid_incorr1, invalid_index1[sub_idx1]]
        ), np.concatenate([valid_corr[:, 1], valid_incorr2, invalid_index2[sub_idx2]])

        pscore1, pscore2 = (
            pscore1[per_idx1][:, np.newaxis],
            pscore2[per_idx2][:, np.newaxis],
        )
        x1, x2 = x1[per_idx1][:, :2], x2[per_idx2][:, :2]
        desc1, desc2 = desc1[per_idx1], desc2[per_idx2]
        kpt1, kpt2 = kpt1[per_idx1], kpt2[per_idx2]

        return {
            "x1": x1,
            "x2": x2,
            "kpt1": kpt1,
            "kpt2": kpt2,
            "desc1": desc1,
            "desc2": desc2,
            "num_corr": num_corr,
            "num_incorr1": num_incorr1,
            "num_incorr2": num_incorr2,
            "e_gt": egt,
            "pscore1": pscore1,
            "pscore2": pscore2,
            "img_path1": img_path1,
            "img_path2": img_path2,
        }

    def __len__(self):
        return self.total_pairs