import math


def squeezenet1_0_fpn_feature_shape_fn(img_shape):
    """ Takes an image_shape as an input to calculate the FPN output sizes
    Ensure that img_shape is of the format (..., H, W)

    Args
        img_shape : image shape as torch.Tensor not torch.Size should have
            H, W as last 2 axis
    Returns
        P3_shape, P4_shape, P5_shape, P6_shape, P7_shape : as 5 (2,) Tensors
    """
    C0_shape = img_shape[-2:]

    C1_shape = (math.floor((C0_shape[0] - 5) / 2), math.floor((C0_shape[1] - 5) / 2))
    C2_shape = (math.ceil((C1_shape[0] - 1) / 2), math.ceil((C1_shape[1] - 1) / 2))

    P3_shape = (math.ceil((C2_shape[0] - 1) / 2), math.ceil((C2_shape[1] - 1) / 2))
    P4_shape = (math.ceil((P3_shape[0] - 1) / 2), math.ceil((P3_shape[1] - 1) / 2))
    P5_shape = (math.ceil((P4_shape[0] - 1) / 2), math.ceil((P4_shape[1] - 1) / 2))

    P6_shape = (math.ceil(P5_shape[0] / 2), math.ceil(P5_shape[1] / 2))
    P7_shape = (math.ceil(P6_shape[0] / 2), math.ceil(P6_shape[1] / 2))

    return C2_shape, P3_shape, P4_shape, P5_shape, P6_shape, P7_shape