from __future__ import absolute_import, division, print_function, unicode_literals

import copy
from collections import defaultdict

import numpy as np
from caffe2.python import core, utils
from caffe2.python.fb import hardcode_scale_zp


def pairwise(iterable):
    "s -> (s0,s1), (s1,s2), (s2, s3), ..."
    from itertools import tee

    a, b = tee(iterable)
    next(b, None)
    return zip(a, b)


def blob_uses(net, blob):
    u = []
    for i, op in enumerate(net.op):
        if blob in op.input or blob in op.control_input:
            u.append(i)
    return u


def fuse_first_bn(net, params, removed_tensors):
    net = copy.deepcopy(net)
    params = copy.deepcopy(params)

    for ((i, current), (j, next_)) in pairwise(enumerate(net.op)):
        if next_.input[0] != current.output[0]:
            continue

        if current.type not in ("Conv", "ConvTranspose") or next_.type != "SpatialBN":
            continue
        if (
            len(blob_uses(net, current.output[0])) != 1
            and current.output[0] != next_.output[0]
        ):
            # Can't fuse if more than one user unless SpatialBN is inplace
            continue

        # else, can fuse
        conv = current
        bn = next_
        fused_conv = copy.deepcopy(conv)
        fused_conv.output[0] = bn.output[0]
        conv_weight = params[conv.input[1]]
        if len(conv.input) > 2:
            conv_bias = params[conv.input[2]]
        else:
            conv_bias = np.zeros(len(params[bn.input[2]])).astype(np.float32)

        bn_scale = params[bn.input[1]]
        bn_bias = params[bn.input[2]]
        bn_running_mean = params[bn.input[3]]
        bn_running_var = params[bn.input[4]]

        # First, BN computation can be phrased as follows:
        # (X - running_mean) * (1.0 / sqrt(running_var + eps)) *
        # bn_scale + bias
        # Thus, we can rewrite bn_scale as:
        # X * bn_scale * 1.0 / (sqrt(running_var + eps)) + (bias -
        # running_mean * (1.0 / sqrt(running_var + eps)) * bn_scale)
        # Thus, can just have the affine transform
        # X * A + B
        # where
        # A = bn_scale * 1.0 / (sqrt(running_var + eps))
        # B =  (bias - running_mean * (1.0 / sqrt(running_var + eps))
        # * bn_scale)
        eps = 1.0e-5
        for arg in bn.arg:
            if arg.name == "epsilon":
                eps = arg.f
        A = bn_scale * 1.0 / (np.sqrt(bn_running_var + eps))
        B = bn_bias - bn_running_mean * A

        # This identity should hold if we have correctly fused
        # np.testing.assert_array_equal(
        #     params[conv.output[0]] * A + B,
        #     params[bn.output[0]])

        # Now, we have that the computation made is the following:
        # ((X `conv` W) + b) * A + B
        # Then, we can simply fuse this as follows:
        # (X `conv` (W * A)) + b * A + B
        # which is simply
        # (X `conv` Q) + C
        # where

        # Q = W * A
        # C = b * A + B

        # For ConvTranspose, from the view of convolutions as a
        # Toepeliz multiplication, we have W_ = W^T, so the weights
        # are laid out as (R, S, K, K) (vs (S, R, K, K) for a Conv),
        # so the weights broadcast slightly differently. Remember, our
        # BN scale 'B' is of size (S,)

        A_ = (
            A.reshape((-1,) + tuple([1] * (conv_weight.ndim - 1)))
            if conv.type == "Conv"
            else A.reshape((1, -1) + tuple([1] * (conv_weight.ndim - 2)))
        )

        C = conv_bias * A + B
        Q = conv_weight * A_

        assert params[conv.input[1]].shape == Q.shape
        if len(conv.input) > 2:
            assert params[conv.input[2]].shape == C.shape
        else:
            assert bn_bias.shape == C.shape

        params[conv.input[1]] = Q
        if len(conv.input) > 2:
            params[conv.input[2]] = C
        else:
            params[bn.input[2]] = C
            fused_conv.input.append(bn.input[2])

        new_ops = net.op[:i] + [fused_conv] + net.op[j + 1 :]
        del net.op[:]
        removed_tensors.append(bn.input[1])
        if len(conv.input) > 2:
            removed_tensors.append(bn.input[2])
        removed_tensors.append(bn.input[3])
        removed_tensors.append(bn.input[4])
        del params[bn.input[1]]
        if len(conv.input) > 2:
            del params[bn.input[2]]
        del params[bn.input[3]]
        del params[bn.input[4]]
        net.op.extend(new_ops)
        break
    return net, params, removed_tensors


def fuse_bn(net, params, ignore_failure):
    # Run until we hit a fixed point
    removed_tensors = []
    while True:
        (next_net, next_params, removed_tensors) = fuse_first_bn(
            net, params, removed_tensors
        )
        if len(next_net.op) == len(net.op):
            if any(op.type == "SpatialBN" for op in next_net.op) and not ignore_failure:
                raise Exception(
                    "Model contains SpatialBN op after fusion: %s", next_net
                )
            return (next_net, next_params, removed_tensors)
        net, params, removed_tensors = (next_net, next_params, removed_tensors)


def fuse_first_scale(net, params, removed_tensors):
    net = copy.deepcopy(net)
    params = copy.deepcopy(params)

    for ((i, current), (j, next_)) in pairwise(enumerate(net.op)):
        if next_.input[0] != current.output[0]:
            continue

        if (
            current.type != "SpatialBN"
            or next_.type != "Mul"
            or len(net.op) <= j + 1
            or net.op[j + 1].type != "Add"
        ):
            continue

        # else, can fuse
        bn = current
        mul = next_
        add = net.op[j + 1]

        fused_bn = copy.deepcopy(bn)
        fused_bn.output[0] = add.output[0]
        bn_scale = params[bn.input[1]]
        mul_scale = params[mul.input[1]]
        bn_bias = params[bn.input[2]]
        add_bias = params[add.input[1]]

        params[bn.input[1]] = bn_scale * mul_scale
        params[bn.input[2]] = mul_scale * bn_bias + add_bias

        new_ops = net.op[:i] + [fused_bn] + net.op[j + 2 :]
        del net.op[:]
        removed_tensors.append(mul.input[1])
        removed_tensors.append(add.input[1])
        del params[mul.input[1]]
        del params[add.input[1]]
        net.op.extend(new_ops)
        break
    return net, params, removed_tensors


def fuse_scale(net, params, ignore_failure):
    # Run until we hit a fixed point
    removed_tensors = []
    while True:
        (next_net, next_params, removed_tensors) = fuse_first_scale(
            net, params, removed_tensors
        )
        if len(next_net.op) == len(net.op):
            return (next_net, next_params, removed_tensors)
        net, params, removed_tensors = (next_net, next_params, removed_tensors)


def fuse_first_relu(net, ignore_op_with_output=None):
    net = copy.deepcopy(net)

    for ((i, current), (j, next_)) in pairwise(enumerate(net.op)):
        if next_.input[0] != current.output[0]:
            continue

        if current.type not in ("Conv", "Sum") or next_.type != "Relu":
            continue

        if ignore_op_with_output and current.output[0] in ignore_op_with_output:
            continue

        # else, can fuse
        conv = current
        relu = next_
        fused_conv = copy.deepcopy(conv)
        fused_conv.type = "ConvRelu" if current.type == "Conv" else "SumRelu"
        fused_conv.output[0] = relu.output[0]

        new_ops = net.op[:i] + [fused_conv] + net.op[j + 1 :]
        del net.op[:]
        net.op.extend(new_ops)
        break
    return net


def fuse_relu(net, ignore_failure, ignore_op_with_output=None):
    # Run until we hit a fixed point
    while True:
        next_net = fuse_first_relu(net, ignore_op_with_output)
        if len(next_net.op) == len(net.op):
            if any(op.type == "Relu" for op in next_net.op) and not ignore_failure:
                raise Exception("Model contains Relu op after fusion: %s", next_net)
            return next_net
        net = next_net


def last_producer(ops, blob):
    for (i, op) in reversed(list(enumerate(ops))):
        if op.output[0] == blob:
            return i
    raise ValueError("Failed to find last producer of blob, %s", blob)


def swap_first_concat_relu(net, ignore_op_with_output=None):
    net = copy.deepcopy(net)

    for ((i, current), (j, next_)) in pairwise(enumerate(net.op)):
        if next_.input[0] != current.output[0]:
            continue

        if current.type != "Concat" or next_.type != "Relu":
            continue

        if ignore_op_with_output and current.output[0] in ignore_op_with_output:
            continue

        # else, can swap
        concat = copy.deepcopy(current)
        relu = copy.deepcopy(next_)
        pre_ops = copy.deepcopy(net.op[:i])
        post_ops = copy.deepcopy(net.op[j + 1 :])

        # Delete the Relu after Concat
        concat.output[0] = relu.output[0]

        # Insert Relu after each op that produces inputs to Concat
        for blob in concat.input:
            k = last_producer(pre_ops, blob)
            producer = pre_ops[k]
            assert producer.output[0] == blob
            producer.output[0] = blob + "_pre_relu"

            new_relu = copy.deepcopy(relu)
            new_relu.input[0] = producer.output[0]
            new_relu.output[0] = blob

            pre_ops = pre_ops[: k + 1] + [new_relu] + pre_ops[k + 1 :]

        new_ops = pre_ops + [concat] + post_ops
        del net.op[:]
        net.op.extend(new_ops)
        break
    return net


def swap_concat_relu(net, ignore_op_with_output=None):
    # Run until we hit a fixed point
    while True:
        next_net = swap_first_concat_relu(net, ignore_op_with_output)
        if len(next_net.op) == len(net.op):
            return next_net
        net = next_net


def add_version_to_conv_bias(net, init_net):
    """
    In architectures such as FPN (https://arxiv.org/abs/1612.03144), few Conv
    ops share the same weight and bias and are run at different scales of
    the input. Since 'bias_scale = input_scale * weight_scale', sharing the
    same bias blob among multiple Conv ops means that we need different bias
    scale for each of the ops. To achieve this, we just duplicate those bias
    blobs that are used by multiple Conv ops before performing int8 rewrite.
    """
    bias_count = defaultdict(int)
    for op in net._net.op:
        if "Conv" in op.type and len(op.input) >= 3:
            bias_count[op.input[2]] += 1

    bias_fill_op = {}
    for op in init_net._net.op:
        if bias_count[op.output[0]] > 1:
            bias_fill_op[op.output[0]] = op

    bias_version = defaultdict(int)
    for op in net._net.op:
        if "Conv" in op.type and len(op.input) >= 3:
            bias = op.input[2]
            if bias_count[bias] <= 1:
                continue

            version = bias_version[bias]
            bias_version[bias] += 1
            if version == 0:
                continue

            new_bias = bias + "_v" + str(version)
            fill_op = copy.deepcopy(bias_fill_op[bias])
            fill_op.output[0] = new_bias
            init_net._net.op.extend([fill_op])
            op.input[2] = new_bias
            net._net.external_input.append(new_bias)


def add_quantization_param_args_(op, q_param):
    op.arg.extend(
        [
            utils.MakeArgument("Y_scale", q_param.scale),
            utils.MakeArgument("Y_zero_point", q_param.zero_point),
        ]
    )


def choose_quantization_params(tensor_min, tensor_max, preserve_sparsity=False):
    if tensor_min < 0 and tensor_max > 0 and preserve_sparsity:
        symmetric_qmin = -(255 // 2 + 1)
        symmetric_qmax = 255 // 2
        max_scale = max(
            abs(tensor_min / symmetric_qmin), abs(tensor_max / symmetric_qmax)
        )
        tensor_min = max_scale * symmetric_qmin
        tensor_max = max_scale * symmetric_qmax

    q_param = hardcode_scale_zp.choose_quantization_params(tensor_min, tensor_max)

    if tensor_min < 0 and tensor_max > 0 and preserve_sparsity:
        q_param = hardcode_scale_zp.QuantizationParam(q_param.scale, 128)

    return q_param


def add_quantization_param_args(op, tensor, preserve_sparsity=False):
    tensor_min = 0 if tensor.size == 0 else tensor.min()
    tensor_max = 0 if tensor.size == 0 else tensor.max()

    q_param = choose_quantization_params(tensor_min, tensor_max, preserve_sparsity)

    add_quantization_param_args_(op, q_param)
    return q_param


def create_int8_given_tensor_fill(tensor, out_blob_name, preserve_sparsity=False):
    """
    Create Int8GivenTensorFill op that quantizes the given tensor and outputs
    an Int8Tensor with out_blob_name.
    """
    op = core.CreateOperator("Int8GivenTensorFill", [], out_blob_name)
    q_param = add_quantization_param_args(op, tensor, preserve_sparsity)
    quantized_tensor = (
        np.around(tensor / q_param.scale).astype(np.int32) + q_param.zero_point
    )
    quantized_tensor = np.maximum(0, np.minimum(quantized_tensor, 255))
    op.arg.extend(
        [
            utils.MakeArgument("values", quantized_tensor.astype(np.uint8).tobytes()),
            utils.MakeArgument("shape", quantized_tensor.shape),
        ]
    )
    return op, q_param


def create_int8_bias_tensor_fill(tensor, out_blob_name, x_q_param, w_q_param):
    """
    Similar to create_int8_given_tensor_fill, but for bias blobs to be stored
    as int32.
    """
    scale = x_q_param.scale * w_q_param.scale
    quantized_tensor = np.around(tensor / scale).astype(np.int32)
    quantized_tensor.reshape(-1)
    op = core.CreateOperator("Int8GivenIntTensorFill", [], out_blob_name)
    op.arg.extend(
        [
            utils.MakeArgument("values", quantized_tensor),
            utils.MakeArgument("shape", quantized_tensor.shape),
        ]
    )
    q_param = hardcode_scale_zp.QuantizationParam(scale, 0)
    add_quantization_param_args_(op, q_param)
    return op