import re


def get_feat_shape(visual_archi, compress_type=None):
    """
    Get feat shape depending on the training archi and compress type
    """
    if visual_archi == "fasterrcnn":
        # the RCNN model should be trained with min_size=224
        feat_shape = (-1, 2048, 7, 7)
    elif visual_archi == "maskrcnn":
        # the RCNN model should be trained with min_size=800
        feat_shape = (-1, 2048, 10, 10)
    elif visual_archi == "resnet18":
        feat_shape = (-1, 512, 7, 7)
    else:
        raise NotImplementedError("Unknown archi {}".format(visual_archi))

    if compress_type is not None:
        if not re.match(r"\d+x", compress_type):
            raise NotImplementedError("Unknown compress type {}".format(compress_type))
        compress_times = int(compress_type[:-1])
        feat_shape = (
            feat_shape[0],
            feat_shape[1] // compress_times,
            feat_shape[2],
            feat_shape[3],
        )
    return feat_shape