import torch


def zero_out_col_span(xfft, col, start_row, end_row=None):
    """
    Zero out values for all data points, channels, in the given column and row
    span.

    :param xfft: input complex tensor
    :param col: the col number to zero out value
    :param start_row: the start row (inclusive, it is zeroed out)
    :param end_row: the end row (exclusive, it is not zeroed out) - by default,
    zero out all the values to the end of the column
    :return: the zero out col for the given row range

    >>> xfft = torch.tensor(
    ... [ # 1 image
    ... [ # 1 channel
    ... [[[8, 0], [2, 5], [3, 0]], # 1st row
    ... [[4, 0], [5, 0], [2, 0]],
    ... [[-1, 1], [1, 0], [1, 0]]]
    ... ]
    ... ])
    >>> result = zero_out_col_span(xfft, 1, 0, 3)
    >>> expect = torch.tensor(
    ... [ # 1 image
    ... [ # 1 channel
    ... [[[8, 0], [0, 0], [3, 0]], # 1st row
    ... [[4, 0], [0.0, 0.0 ], [2, 0]],
    ... [[-1, 1], [0, 0], [1, 0]]]
    ... ]
    ... ])
    >>> # print("result: ", result)
    >>> np.testing.assert_array_almost_equal(result, expect)
    >>> result2 = zero_out_col_span(xfft, 2, 0, 1)
    >>> expect2 = torch.tensor(
    ... [ # 1 image
    ... [ # 1 channel
    ... [[[8, 0], [0, 0], [0, 0]], # 1st row
    ... [[4, 0], [0.0, 0.0], [2, 0]],
    ... [[-1, 1], [0, 0], [1, 0]]]
    ... ]
    ... ])
    >>> np.testing.assert_array_almost_equal(result2, expect2)
    >>> result3 = zero_out_col_span(xfft, 0, 1, 3)
    >>> expect3 = torch.tensor(
    ... [ # 1 image
    ... [ # 1 channel
    ... [[[8, 0], [0, 0], [0, 0]], # 1st row
    ... [[0, 0], [0.0, 0.0], [2, 0]],
    ... [[0, 0], [0, 0], [1, 0]]]
    ... ]
    ... ])
    >>> np.testing.assert_array_almost_equal(result3, expect3)
    """
    if end_row is None:
        # zero out to the end of the column
        end_row = xfft.shape[2]
    if end_row > start_row:
        xfft[:, :, start_row:end_row, col, :] = torch.zeros(xfft.shape[0],
                                                            xfft.shape[1],
                                                            end_row - start_row,
                                                            xfft.shape[4])
    return xfft