"""
A set of utility functions to get the actual computation complexity of various layers
"""

import math
import warnings

oneGiga = 1e9


def convolutional_layer_direct(imageSize, filterSize, padding='SAME', stride=(1, 1), inGFLOP=True,
                               result_format='mac'):
    """ The input conforms to tensorflow tf.nn.conv2d format

    :param imageSize: (image_height, image_width, in_channels), represented as a tuple
    :param filterSize: (filter_height, filter_width, in_channels, out_channels)
    :param padding:
    :param stride: (height_stride, width_stride)
    :return: The number of operations: accumulation + multiplication
    """
    image_height, image_width, in_channels = imageSize
    assert in_channels == filterSize[2], "image size and filter size must have the same depth."
    filter_height, filter_width, _, out_channels = filterSize
    height_stride, width_stride = stride
    if padding == 'SAME':
        height_padding = (image_height - 1) * height_stride - image_height + filter_height
        width_padding = (image_width - 1) * width_stride - image_width + filter_width
    elif padding == 'VALID':
        height_padding = 0
        width_padding = 0
    else:
        raise ValueError('Unknown padding')

    out_height = (image_height - filter_height + height_padding) // height_stride + 1
    out_width = (image_width - filter_width + width_padding) // width_stride + 1
    # number of operations to get one result
    numOpMac = filter_height * filter_width * in_channels
    numOpPerResult = 2 * numOpMac + in_channels + 1  # 1 is the bias
    total_num_result = out_height * out_width * out_channels
    if result_format == 'mac':
        total_num_ops = total_num_result * numOpMac
    elif result_format == 'op':
        total_num_ops = total_num_result * numOpPerResult
    else:
        raise ValueError('Unknown result format')
    if inGFLOP:
        total_num_ops /= oneGiga
    return total_num_ops


def isPowerOfTwo(num):
    return ((num & (num - 1)) == 0) and num > 0


def findUnitSize(size):
    """
    >>> findUnitSize(3)
    2
    >>> findUnitSize(7)
    2
    """
    assert size >= 3, "It is not valid to use fft when size < 3"
    F = size
    start = int(math.ceil(math.log(F, 2)))
    while True:
        total = 2 ** start
        unitSize = total + 1 - F
        if unitSize >= 2:
            break
        start += 1
    return unitSize


def convolutional_layer_fft(imageSize, filterSize, padding='SAME', stride=(1, 1), inGFLOP=True,
                            result_format='mac', fft_size=None):
    """ The input conforms to tensorflow tf.nn.conv2d format

    :param imageSize: (image_height, image_width, in_channels), represented as a tuple
    :param filterSize: (filter_height, filter_width, in_channels, out_channels)
    :param padding:
    :param stride: (height_stride, width_stride)
    :param fft_size: find the minimum feasible fft size if it is None, (height_fft_size, width_fft_size)
    :return: The number of operations: accumulation + multiplication
    """
    image_height, image_width, in_channels = imageSize
    filter_height, filter_width, _, out_channels = filterSize
    assert in_channels == filterSize[2], "image size and filter size must have the same depth."
    if fft_size is None:
        height_unit_size = findUnitSize(filter_height)
        width_unit_size = findUnitSize(filter_width)
        height_fft_size = height_unit_size + filter_height - 1
        width_fft_size = width_unit_size + filter_width - 1
    else:
        height_fft_size, width_fft_size = fft_size
        height_unit_size = height_fft_size + 1 - filter_height
        width_unit_size = width_fft_size + 1 - filter_width

    # number of operations
    height_logFFTUnitSize = int(math.log(height_fft_size, 2))
    width_logFFTUnitSize = int(math.log(width_fft_size, 2))

    height_stride, width_stride = stride

    if padding == 'SAME':