"""Portable INT8 weight storage with FP32 arithmetic (no dynamic activation quantization)."""
import numpy as np
import onnx
from onnx import numpy_helper, helper, TensorProto


def quantize_weights(source, destination):
    model = onnx.load(str(source))
    initializers = []
    dequantizers = []
    gather_weights = {node.input[0] for node in model.graph.node if node.op_type == 'Gather'}
    for tensor in model.graph.initializer:
        if tensor.data_type != TensorProto.FLOAT or len(tensor.dims) != 2:
            initializers.append(tensor)
            continue
        array = numpy_helper.to_array(tensor)
        axis = 0 if tensor.name in gather_weights else 1
        scale = np.maximum(np.max(np.abs(array), axis=1-axis) / 127, np.finfo(np.float32).tiny)
        expanded = scale[:, None] if axis == 0 else scale[None, :]
        quantized = np.clip(np.rint(array / expanded), -127, 127).astype(np.int8)
        weight, factor, zero = [tensor.name + suffix for suffix in ('_int8', '_scale', '_zero')]
        initializers.extend([numpy_helper.from_array(quantized, weight),
                             numpy_helper.from_array(scale, factor),
                             numpy_helper.from_array(np.zeros_like(scale, dtype=np.int8), zero)])
        dequantizers.append(helper.make_node('DequantizeLinear', [weight, factor, zero],
                                             [tensor.name], axis=axis, name=tensor.name+'_dequantize'))
    del model.graph.initializer[:]
    model.graph.initializer.extend(initializers)
    nodes = list(model.graph.node)
    del model.graph.node[:]
    model.graph.node.extend(dequantizers + nodes)
    onnx.save_model(model, str(destination))

if __name__ == '__main__':
    import sys
    quantize_weights(sys.argv[1], sys.argv[2])