from functools import reduce


def flatten_shape(input_shape, axis=1, end_axis=-1):
    """ calculate the output shape of this layer using input shape

    Args:
        @input_shape (list of num): a list of number which represents the input shape
        @axis (int): parameter from caffe's Flatten layer
        @end_axis (int): parameter from caffe's Flatten layer

    Returns:
        @output_shape (list of num): a list of numbers represent the output shape
    """

    start_axis = axis
    end_axis = end_axis
    input_shape = list(input_shape)
    if start_axis < 0:
        start_axis += len(input_shape)

    if end_axis < 0:
        end_axis += len(input_shape)

    assert start_axis <= end_axis, 'invalid axis[%d] or end_axis[%d] params'\
            % (start_axis, end_axis)
    output_shape = input_shape[0:start_axis]
    flat_sz = reduce(lambda a, b: a * b, input_shape[start_axis:end_axis])
    output_shape += [flat_sz]
    output_shape += input_shape[end_axis:-1]

    return output_shape