from __future__ import annotations

import open3d as o3d
from open3d.geometry import PointCloud


class RemoveFloor:
    def __init__(self, distance_threshold: float) -> None:
        self.distance_threshold = distance_threshold

    def __call__(self, pcd: PointCloud) -> PointCloud:
        max_bound = pcd.get_max_bound()
        min_bound = pcd.get_min_bound()
        min_bound[2] = min_bound[2] + self.distance_threshold

        bounding_box = o3d.geometry.AxisAlignedBoundingBox(
            min_bound=min_bound,
            max_bound=max_bound,
        )

        pcd = pcd.crop(bounding_box)

        return pcd
