# pylint: disable=invalid-name, unused-variable, too-many-locals, too-many-arguments, unused-argument
"""Bitserial Conv2D operators"""
from __future__ import absolute_import as _abs
from collections import namedtuple
import numpy as np
import tvm
from topi.transform import concatenate
from .pad import pad
from .util import get_pad_tuple
from ..util import get_const_tuple, get_const_int

# workload description of conv2d
Workload = namedtuple('Workload',
                      ['in_dtype', 'out_dtype', 'height', 'width', 'in_filter', 'out_filter',
                       'hkernel', 'wkernel', 'hpad', 'wpad', 'hstride', 'wstride'])

SpatialPackNCHW = namedtuple('SpatialPack',
                             ['vh', 'vw', 'vc', 'ba', 'bc'])

SpatialPackNHWC = namedtuple('SpatialPack',
                             ['vh', 'vw', 'vc', 'ba', 'bc'])

_WORKLOADS = [
    # workloads of resnet18 on imagenet
    # input_size, input_size, ic, oc, kh, kw, pad, pad, stride, stride
    Workload('uint32', 'int32', 56, 56, 64, 64, 3, 3, 1, 1, 1, 1),
    Workload('uint32', 'int32', 56, 56, 64, 64, 1, 1, 0, 0, 1, 1),
    Workload('uint32', 'int32', 56, 56, 64, 128, 3, 3, 1, 1, 2, 2),
    Workload('uint32', 'int32', 56, 56, 64, 128, 1, 1, 0, 0, 2, 2),
    Workload('uint32', 'int32', 28, 28, 128, 128, 3, 3, 1, 1, 1, 1),
    Workload('uint32', 'int32', 28, 28, 128, 256, 3, 3, 1, 1, 2, 2),
    Workload('uint32', 'int32', 28, 28, 128, 256, 1, 1, 0, 0, 2, 2),
    Workload('uint32', 'int32', 14, 14, 256, 256, 3, 3, 1, 1, 1, 1),
    Workload('uint32', 'int32', 14, 14, 256, 512, 3, 3, 1, 1, 2, 2),
    Workload('uint32', 'int32', 14, 14, 256, 512, 1, 1, 0, 0, 2, 2),
    Workload('uint32', 'int32', 7, 7, 512, 512, 3, 3, 1, 1, 1, 1),

    # workload of alexnet on cifar10
    Workload('int32', 'int32', 27, 27, 96, 192, 5, 5, 2, 2, 1, 1),
    Workload('int32', 'int32', 13, 13, 192, 384, 3, 3, 1, 1, 1, 1),
    Workload('int32', 'int32', 13, 13, 384, 384, 3, 3, 1, 1, 1, 1),
    Workload('int32', 'int32', 13, 13, 384, 256, 3, 3, 1, 1, 1, 1),
]

@tvm.target.generic_func
def bitserial_conv2d(data, kernel, stride, padding, activation_bits, weight_bits,
                     layout='NCHW', pack_dtype='uint32', out_dtype='int32', dorefa=True):
    """Bitserial Conv2D operator.

    Parameters
    ----------
    input : tvm.Tensor
        4-D with shape [batch, in_channel, in_height, in_width] or
                       [batch, in_height, in_width, in_channel]

    filter : tvm.Tensor
        4-D with shape [num_filter, in_channel, filter_height, filter_width] or
		       [filter_height, filter_width, in_channel, num_filter]

    stride : int or a list/tuple of two ints
        stride size, or [stride_height, stride_width]

    padding : int or a list/tuple of two ints
        padding size, or [pad_height, pad_width]

    layout : str
        layout of data

    activation_bits: int
        number of bits used for activations/input elements

    weight_bits: int
        number of bits used for weight elements

    out_dtype: str
        return type of convolution

    pack_dtype: str
        bit packing type

    dorefa: bool
        preform the bitserial dot-product using 2 popcounts (required for DoReFa-Net)

    Returns
    -------
    output : tvm.Tensor
        4-D with shape [batch, out_channel, out_height, out_width] or
                       [batch, out_height, out_width, out_channel]
    """
    # search platform specific declaration first
    # default declaration
    if layout == 'NCHW':
        return spatial_pack_nchw(data, kernel, stride, padding, activation_bits, weight_bits,
                                 pack_dtype=pack_dtype, out_dtype=out_dtype, dorefa=dorefa)
    elif layout == 'NHWC':
        return spatial_pack_nhwc(data, kernel, stride, padding, activation_bits, weight_bits,
                                 pack_dtype=pack_dtype, out_dtype=out_dtype, dorefa=dorefa)
    raise ValueError("not support this layout {} yet".format(layout))

def _get_workload(data, kernel, stride, padding, out_dtype, layout):
    """ Get the workload structure. """
    assert layout == "NCHW" or layout == "NHWC", \
        "Only support layouts NCHW and NHWC"
    if layout == "NCHW":
        _, CI, IH, IW = [x.value for x in data.shape]
        CO, _, KH, KW = [x.value for x in kernel.shape]
    else: # NHWC
        IH, IW = data.shape[1].value, data.shape[2].value
        KH, KW, CI, CO = [x for x in get_const_tuple(kernel.shape)]

    HPAD, WPAD, _, _ = get_pad_tuple(padding, kernel)
    if isinstance(stride, (tuple, list)):
        HSTR, WSTR = stride
    else:
        HSTR, WSTR = stride, stride

    return Workload(data.dtype, out_dtype, IH, IW, CI, CO, KH, KW, HPAD, WPAD, HSTR, WSTR)

@tvm.target.generic_func
def _get_schedule(wkl, layout):
    # pylint: disable=unreachable
    """ Get the platform specific schedule. """
    target = tvm.target.current_target()
    raise RuntimeError(
        "No schedule for current target:{}".format(target))
    # This return has no use, merely to supress pylint warning
    return wkl

def spatial_pack_nchw(data, kernel, stride, padding, in_bits, weight_bits,
                      pack_dtype, out_dtype, dorefa=False):
    """ Compute convolution with pack on spatial axes. """
    assert data.shape[0].value == 1, "spatial pack convolution only support batch size=1"
    data_q = bitpack(data, in_bits, pack_axis=1, bit_axis=0, pack_type=pack_dtype)
    kernel_q = bitpack(kernel, weight_bits, pack_axis=1, bit_axis=0, pack_type=pack_dtype)
    IB, _, CI, H, W = data_q.shape
    KB, CO, _, KH, KW = kernel_q.shape
    HPAD, WPAD, _, _ = get_pad_tuple(padding, kernel)

    if isinstance(stride, (tuple, list)):
        HSTR, WSTR = stride
    else:
        HSTR, WSTR = stride, stride
    HCAT, WCAT = KH-1, KW-1

    wkl = _get_workload(data, kernel, stride, padding, out_dtype, "NCHW")
    sch = _get_schedule(wkl, "NCHW")
    VH = sch.vh
    VW = sch.vw
    VC = sch.vc

    TH = H + 2*HPAD
    TW = W + 2*WPAD
    OH = (H + 2*HPAD - KH) // HSTR + 1
    OW = (W + 2*WPAD - KW) // WSTR + 1

    dshape = (IB, 1, CI, H, W)