# %%
from collections.abc import Iterator
from enum import Enum, auto
from pathlib import Path
from typing import overload
from scipy.spatial.transform import Rotation as R

import numpy as np

from .basic import (
    get_lidar_files,
    load_pcd,
    load_global_poses,
    transform_point_cloud,
)

CLD_DIR = "Clouds"
GT_PATH = "poses_aligned.csv"


class FrameID(Enum):
    BASE_LINK = auto()
    WORLD = auto()

    @classmethod
    def help(cls) -> str:
        names = [f"FrameID.{e.name}: {e.__doc__}" for e in FrameID]
        return f"Valid frame_id choices:\n\t{'\n\t'.join(names)}"


FrameID.BASE_LINK.__doc__ = "Vehicle frame"
FrameID.WORLD.__doc__ = "World frame (ground truth poses)"


class PointCloudLoader:
    def __init__(self, root_dir: str, frame_id: FrameID = FrameID.WORLD):
        """
        Data loader for point cloud sequences.

        Args:
            root_dir (str): Root directory containing the dataset.
            frame_id (FrameID, optional): Coordinate frame for point clouds. Defaults to FrameID.BASE_LINK.
        """
        self.root_dir = Path(root_dir)
        self.frame_id = frame_id
        cld_dir = self.root_dir / CLD_DIR
        gt_path = self.root_dir / GT_PATH

        assert self.root_dir.exists(), f"Root directory {self.root_dir} does not exist."
        assert cld_dir.exists(), f"Point cloud directory {cld_dir} does not exist."
        assert gt_path.exists(), f"Ground truth file {gt_path} does not exist."

        self.lidar_files = get_lidar_files(cld_dir)
        self.gstamps, self.gposes = load_global_poses(gt_path)

        if len(self.lidar_files) != len(self.gposes):
            raise ValueError("Number of LiDAR files and poses do not match.")

    def __len__(self) -> int:
        return len(self.lidar_files)

    @property
    def frame_id(self) -> FrameID:
        return self._frame_id

    @frame_id.setter
    def frame_id(self, value: FrameID) -> None:
        if not isinstance(value, FrameID):
            raise TypeError(
                f"frame_id must be an instance of FrameID, got {type(value)}\n"
                f"{FrameID.help()}"
            )
        self._frame_id = value

    def __get_one(self, idx: int) -> tuple[np.ndarray, np.ndarray]:
        if idx < 0 or idx >= len(self):
            raise IndexError(f"Index {idx} out of range [0, {len(self)})")

        cld = load_pcd(self.lidar_files[idx])
        pose = self.gposes[idx]
        if self.frame_id == FrameID.WORLD:
            cld = transform_point_cloud(cld, pose)
        return cld, pose

    @overload
    def __getitem__(self, idx: int) -> tuple[np.ndarray, np.ndarray]:
        """
        Returns the LiDAR point cloud and transformation matrix for a single index.
        Args:
            idx (int): The index of the LiDAR data to retrieve.
        Returns:
            tuple[np.ndarray, np.ndarray]: A tuple containing the point cloud and transformation matrix.
        """
        ...

    @overload
    def __getitem__(self, idx: slice) -> list[tuple[np.ndarray, np.ndarray]]:
        """
        Returns the LiDAR point clouds and transformation matrices for a slice of indices.
        Args:
            idx (slice): A slice object specifying the range of indices to retrieve.
        Returns:
            Iterator[tuple[np.ndarray, np.ndarray]]: An iterator of tuples, each containing a point cloud and transformation matrix.
        """
        ...

    @overload
    def __getitem__(self, idx: np.ndarray) -> list[tuple[np.ndarray, np.ndarray]]:
        """
        Returns the LiDAR point clouds and transformation matrices for an array of indices.
        Args:
            idx (np.ndarray): A numpy array of indices to retrieve.
        Returns:
            Iterator[tuple[np.ndarray, np.ndarray]]: An iterator of tuples, each containing a point cloud and transformation matrix.
        """
        ...

    def __getitem__(
        self, idx: int | slice | np.ndarray
    ) -> tuple[np.ndarray, np.ndarray] | list[tuple[np.ndarray, np.ndarray]]:
        if isinstance(idx, int):
            ret = self.__get_one(idx)
        elif isinstance(idx, slice):
            start, stop, step = idx.indices(len(self))
            ret = list(map(self.__get_one, range(start, stop, step)))
        elif isinstance(idx, np.ndarray):
            if idx.dtype == np.int_:
                ret = list(map(self.__get_one, idx))
            else:
                raise TypeError("Index array must be of integer type.")
        else:
            raise TypeError(
                "Index must be an integer, slice, or numpy array of integers."
            )
        return ret

    def __iter__(self) -> Iterator[tuple[np.ndarray, np.ndarray]]:
        return map(self.__get_one, range(len(self)))


# %% 测试用代码
def _debug_read_all():
    from pypcd4.pypcd4 import PointCloud

    loader = PointCloudDataLoader("/ws//data/wild/Karawatha/K-01")
    clds = []
    for cld, _ in loader[::10]:
        clds.append(cld)
    clds = np.concatenate(clds, axis=0)
    clds = clds[np.random.choice(clds.shape[0], int(1e6), replace=False)]
    PointCloud.from_xyz_points(clds).save("/ws/wild-dataloader/clds.pcd")


if __name__ == "__main__":
    _debug_read_all()
