import copy
import torch

def process_fixpoint_weight_bias(module, weight_quantizer, input_quantizer, output_quantizer=None, bias_bit=32):
    assert(not (module.training or input_quantizer.training)), "only for eval!"
    assert not input_quantizer.is_per_channel or input_quantizer.scale.shape == weight_quantizer.scale.shape, \
        "only for per-tensor activation quantizer or per-channel input activation quantizer before per_channel weight quantizer for depthwise conv"
    assert weight_quantizer.is_symmetric_quant, "only for symmetric weight quantizer"
    assert not hasattr(module, 'bias_fake_quantizer'), "not supported for bias quantizer"
    fixpoint_module = copy.deepcopy(module)

    fixpoint_module.weight, fixpoint_module.bias = \
        process_fixpoint_module_params(fixpoint_module.weight, fixpoint_module.bias,
                                     weight_quantizer,
                                     input_quantizer,
                                     output_quantizer,
                                     bias_bit=bias_bit
                                     )

    return fixpoint_module


def process_fixpoint_module_params(weight, bias, weight_quantizer, input_quantizer, output_quantizer, bias_bit):
    weight_shape = weight.shape
    if weight_quantizer.is_per_channel:
        # perchannel quantizer process
        Sw_shape = [1 for _ in range(len(weight_shape))]
        Sw_shape[weight_quantizer.ch_axis] = weight_shape[weight_quantizer.ch_axis]
        scale = weight_quantizer.scale.reshape(Sw_shape)
    else:
        scale = weight_quantizer.scale

    fixpoint_module_w = torch.clamp(torch.round(weight / scale), weight_quantizer.quant_min, weight_quantizer.quant_max)

    if bias is not None or not input_quantizer.is_symmetric_quant or (
            output_quantizer is not None and not output_quantizer.is_symmetric_quant):
        if bias is None:
            bias_shape = weight_shape[weight_quantizer.ch_axis]
            bias = torch.zeros([bias_shape], dtype=torch.float, device=weight.device)

        SainSw = weight_quantizer.scale * input_quantizer.scale
        fixpoint_module_bias = bias / SainSw

        if weight_quantizer.ch_axis == 0:
            if not input_quantizer.is_per_channel:
                w_sum = fixpoint_module_w.transpose(weight_quantizer.ch_axis, 0).flatten(start_dim=1).sum(dim=1)
                # w_sum = fixpoint_module_w.sum(dim=[1, 2, 3])
                fixpoint_module_bias -= input_quantizer.zero_point.to(torch.float32).round() * w_sum
            else:
                w_sum = fixpoint_module_w.sum(dim=[2, 3])
                fixpoint_module_bias -= (input_quantizer.zero_point.to(torch.float32).round() * w_sum).sum(dim=1)

            if output_quantizer is not None:
                fixpoint_module_bias += output_quantizer.zero_point.to(
                    torch.float32).round() * output_quantizer.scale / SainSw

        elif weight_quantizer.ch_axis == 1:
            if not input_quantizer.is_per_channel:
                w_sum = fixpoint_module_w.sum(dim=[0, 2, 3])
                fixpoint_module_bias -= input_quantizer.zero_point.to(torch.float32).round() * w_sum
            else:
                w_sum = fixpoint_module_w.sum(dim=[2, 3])
                fixpoint_module_bias -= (input_quantizer.zero_point.to(torch.float32).round() * w_sum).sum(dim=0)

            if output_quantizer is not None:
                fixpoint_module_bias += output_quantizer.zero_point.to(
                    torch.float32).round() * output_quantizer.scale / SainSw

        fixpoint_module_bias = torch.round(fixpoint_module_bias)

        b_min = -2 ** (bias_bit - 1)
        b_max = 2 ** (bias_bit - 1) - 1
        fixpoint_module_bias = torch.clamp(fixpoint_module_bias, b_min, b_max)
        return torch.nn.Parameter(fixpoint_module_w), torch.nn.Parameter(fixpoint_module_bias)
    else:
        num_bias = fixpoint_module_w.shape[weight_quantizer.ch_axis] 
        return torch.nn.Parameter(fixpoint_module_w), torch.nn.Parameter(torch.zeros(num_bias), requires_grad=False)


def flatten_args(node):
    flattned_args = []
    if isinstance(node, dict):
        for v in node.values():
            flattned_args.extend(flatten_args(v))
    elif isinstance(node, tuple) or isinstance(node, list):
        for n in node:
            flattned_args.extend(flatten_args(n))
    else:
        flattned_args.extend([node])
    return flattned_args


def merge_quant_params(input_params, weight_params, output_params):
    assert weight_params['is_symmetric_quant'], "symmetric_quant for weight supported only"
    scale = input_params['scale'] * weight_params['scale'] / output_params['scale']
    return {
        'scale': scale,
        'is_per_channel': weight_params['is_per_channel'] or output_params['is_per_channel'],
        'is_symmetric_quant': output_params['is_symmetric_quant'],
        'ch_axis': output_params['ch_axis'] if not weight_params['is_per_channel'] else 1,
        'quant_min': output_params['quant_min'],
        'quant_max': output_params['quant_max'],
        'bias': torch.zeros_like(scale)
    }


def merge_elementwise_quant_params(input_params, output_params, factor=None):
    scale_in = input_params['scale']
    zero_point_in = input_params['zero_point'].round()
    scale_out = output_params['scale']
    zero_point_out = output_params['zero_point'].round()
    scale = scale_in / scale_out
    bias = zero_point_out - zero_point_in * scale
    return {
        'scale': scale if factor is None else scale * factor,
        'bias': bias,
        'is_per_channel': output_params['is_per_channel'],
        'is_symmetric_quant': output_params['is_symmetric_quant'],
        'ch_axis': output_params['ch_axis'],
        'quant_min': output_params['quant_min'],
        'quant_max': output_params['quant_max']
    }


def merge_add_quant_params(input_params_1, input_params_2, output_params):
    scale_in_1 = input_params_1['scale']
    zero_point_in_1 = input_params_1['zero_point'].round()
    scale_in_2 = input_params_2['scale']
    zero_point_in_2 = input_params_2['zero_point'].round()
    scale_out = output_params['scale']
    zero_point_out = output_params['zero_point'].round()

    scale1 = scale_in_1 / scale_out
    scale2 = scale_in_2 / scale_out
    bias = zero_point_out - zero_point_in_1 * scale1 - zero_point_in_2 * scale2
    return {
        'scale1': scale1,
        'scale2': scale2,
        'bias': bias,
        'is_per_channel': output_params['is_per_channel'],
        'is_symmetric_quant': output_params['is_symmetric_quant'],
        'ch_axis': output_params['ch_axis'],
        'quant_min': output_params['quant_min'],
        'quant_max': output_params['quant_max']
    }


def to_quantizer_params(module):
    return {
        'quant_min': module.quant_min,
        'quant_max': module.quant_max,
        'scale': module.scale if hasattr(module, 'scale') else None,
        'zero_point': module.zero_point if hasattr(module, 'zero_point') else None,
        'ch_axis': module.ch_axis,
        'is_per_channel': module.is_per_channel,
        'is_symmetric_quant': module.is_symmetric_quant,
    }