"""
HK Remote Data Dataset

Author: Yujia Zhang (yujia.zhang.cs@gmail.com)
Please cite our work if the code is helpful to you.
"""

import os
import numpy as np
import torch

from pointcept.utils.cache import shared_dict
from .builder import DATASETS
from .defaults import DefaultDataset
import numpy as np
from scipy.spatial import cKDTree


def find_closest_pair_kdtree(points):
    n = points.shape[0]
    if n < 2:
        return float("inf"), None, None

    tree = cKDTree(points)
    distances, indices = tree.query(points, k=2)
    min_dist = np.min(distances[:, 1])
    min_idx_in_distances = np.argmin(distances[:, 1])
    point1_idx = min_idx_in_distances
    point2_idx = indices[min_idx_in_distances, 1]

    return min_dist, (points[point1_idx], points[point2_idx])


@DATASETS.register_module()
class HKDataset(DefaultDataset):
    def __init__(self, crop_h=630, crop_w=1120, patch_size=14, if_img=True, **kwargs):
        self.crop_h = crop_h
        self.crop_w = crop_w
        self.patch_size = patch_size
        self.patch_h = crop_h // patch_size
        self.patch_w = crop_w // patch_size
        self.if_img = if_img
        super().__init__(**kwargs)

    def get_data(self, idx):
        data_path = self.data_list[idx % len(self.data_list)]
        name = self.get_data_name(idx)
        split = self.get_split_name(idx)
        if self.cache:
            cache_name = f"pointcept-{name}"
            return shared_dict(cache_name)

        data_dict = {}
        assets = os.listdir(data_path)
        for asset in assets:
            if not asset.endswith(".npy"):
                continue
            if asset[:-4] not in self.VALID_ASSETS:
                continue
            data_dict[asset[:-4]] = np.load(os.path.join(data_path, asset))
        data_dict["name"] = name
        data_dict["split"] = split
        if self.if_img:
            data_dict["images"] = torch.empty(
                (0, 3, self.patch_h * self.patch_size, self.patch_w * self.patch_size)
            )
            data_dict["img_num"] = np.array([0], dtype=np.int32)
            data_dict["correspondence"] = np.ones(
                (data_dict["coord"].shape[0], 0, 2), dtype=np.float32
            ) * (-1)
        if "coord" in data_dict.keys():
            data_dict["coord"] = data_dict["coord"].astype(np.float32)

        if "color" in data_dict.keys():
            data_dict["color"] = data_dict["color"].astype(np.float32)

        if "normal" in data_dict.keys():
            data_dict["normal"] = data_dict["normal"].astype(np.float32)

        if "segment" in data_dict.keys():
            data_dict["segment"] = data_dict["segment"].reshape([-1]).astype(np.int32)
        else:
            data_dict["segment"] = (
                np.ones(data_dict["coord"].shape[0], dtype=np.int32) * -1
            )

        if "instance" in data_dict.keys():
            data_dict["instance"] = data_dict["instance"].reshape([-1]).astype(np.int32)
        else:
            data_dict["instance"] = (
                np.ones(data_dict["coord"].shape[0], dtype=np.int32) * -1
            )

        return data_dict