import torch


class OutlierRemoval(object):
    '''
    Class to perform outlier removal based on depth difference in local neighborhood

    Arg(s):
        kernel_size : int
            local neighborhood to consider
        threshold : float
            depth difference threshold
    '''

    def __init__(self, kernel_size=7, threshold=1.5):

        self.kernel_size = kernel_size
        self.threshold = threshold

    def remove_outliers(self, sparse_depth, validity_map):
        '''
        Removes erroneous measurements from sparse depth and validity map

        Arg(s):
            sparse_depth : torch.Tensor[float32]
                N x 1 x H x W tensor sparse depth
            validity_map : torch.Tensor[float32]
                N x 1 x H x W tensor validity map
        Returns:
            torch.Tensor[float32] : N x 1 x H x W sparse depth
            torch.Tensor[float32] : N x 1 x H x W validity map
        '''

        # Replace all zeros with large values
        max_value = 10 * torch.max(sparse_depth)
        sparse_depth_max_filled = torch.where(
            validity_map <= 0,
            torch.full_like(sparse_depth, fill_value=max_value),
            sparse_depth)

        # For each neighborhood find the smallest value
        padding = self.kernel_size // 2
        sparse_depth_max_filled = torch.nn.functional.pad(
            input=sparse_depth_max_filled,
            pad=(padding, padding, padding, padding),
            mode='constant',
            value=max_value)

        min_values = -torch.nn.functional.max_pool2d(
            input=-sparse_depth_max_filled,
            kernel_size=self.kernel_size,
            stride=1,
            padding=0)

        # If measurement differs a lot from minimum value then remove
        validity_map_clean = torch.where(
            min_values < sparse_depth - self.threshold,
            torch.zeros_like(validity_map),
            torch.ones_like(validity_map))

        # Update sparse depth and validity map
        validity_map_clean = validity_map * validity_map_clean
        sparse_depth_clean = sparse_depth * validity_map_clean

        return sparse_depth_clean, validity_map_clean
