# %%
import numpy as np
from pathlib import Path
from pypcd4.pypcd4 import PointCloud
from scipy.spatial.transform import Rotation as R

GT_DTYPE = [
    ("stamp", "S20"),
    ("x", "f8"),
    ("y", "f8"),
    ("z", "f8"),
    ("qx", "f8"),
    ("qy", "f8"),
    ("qz", "f8"),
    ("qw", "f8"),
]


def get_lidar_files(path: Path) -> list[Path]:
    """Get all .bin files in the given directory."""
    return sorted(path.glob("*.pcd"))


# %%
def load_pcd(file_path: Path) -> np.ndarray:
    """Load a .pcd file and return the point cloud as a numpy array."""
    return PointCloud.from_path(file_path).numpy()


# %%
def load_global_poses(file_path: Path) -> tuple[np.ndarray, np.ndarray]:
    def mtx(row):
        _, x, y, z, qx, qy, qz, qw = row
        ret = np.eye(4)
        ret[:3, 3] = [x, y, z]
        ret[:3, :3] = R.from_quat([qx, qy, qz, qw]).as_matrix()
        return ret

    data = np.loadtxt(file_path, skiprows=1, delimiter=",", dtype=GT_DTYPE)
    timestamps = data["stamp"].astype(str)
    poses = np.array([mtx(row) for row in data])
    return timestamps, poses


# %%
def transform_point_cloud(points: np.ndarray, T: np.ndarray) -> np.ndarray:
    """
    Transforms a point cloud using a transformation matrix.
    Args:
        points (np.ndarray): A numpy array of shape (N, 3) containing the points.
        T (np.ndarray): A 4x4 transformation matrix.
    Returns:
        np.ndarray: A numpy array of shape (N, 3) containing the transformed points.
    """
    assert points.shape[1] == 3, "Points should be of shape (N, 3)."
    assert T.shape == (4, 4), "Transformation matrix should be of shape (4, 4)."
    num = points.shape[0]
    points_homogeneous = np.hstack((points, np.ones((num, 1))))
    return (T @ points_homogeneous.T).T[:, :3]


# %% 测试用代码


def _debug_read_all():
    files = get_lidar_files(Path("/ws/data/wild/Karawatha/K-01/Clouds"))
    gstamps, gposes = load_global_poses(
        Path("/ws/data/wild/Karawatha/K-01/poses_aligned.csv")
    )
    clds = []
    for f, pose in zip(files[::10], gposes[::10]):
        cld = load_pcd(f)
        cld_w = transform_point_cloud(cld, pose)
        clds.append(cld_w)
    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")


def _debug_plot_traj():
    import matplotlib.pyplot as plt

    _, poses = load_global_poses(Path("/ws/data/wild/Karawatha/K-01/poses_aligned.csv"))
    plt.scatter(poses[:, 0, 3], poses[:, 1, 3], c=poses[:, 2, 3], s=0.3)


if __name__ == "__main__":
    _debug_read_all()
    _debug_plot_traj()
