import argparse
import glob
import os
import sys

import numpy as np

import onnx
from onnx import numpy_helper


def read_tensorproto_pb_file(filename):
    """Return tuple of tensor name and numpy.ndarray of the data from a pb file containing a TensorProto."""

    tensor = onnx.load_tensor(filename)
    np_array = numpy_helper.to_array(tensor)
    return tensor.name, np_array


def dump_tensorproto_pb_file(filename):
    """Dump the data from a pb file containing a TensorProto."""

    name, data = read_tensorproto_pb_file(filename)
    print("Name: {}".format(name))
    print("Shape: {}".format(data.shape))
    print(data)


def dump_pb(dir_or_filename):
    """Dump the data from either a single .pb file, or all .pb files in a directory.
    All files must contain a serialized TensorProto."""

    if os.path.isdir(dir_or_filename):
        for f in glob.glob(os.path.join(dir_or_filename, '*.pb')):
            print(f)
            dump_tensorproto_pb_file(f)
    else:
        dump_tensorproto_pb_file(dir_or_filename)


def numpy_to_pb(name, np_data, out_filename):
    """Convert numpy data to a protobuf file."""

    tensor = numpy_helper.from_array(np_data, name)
    onnx.save_tensor(tensor, out_filename)


def image_to_numpy(filename, shape, channels_last, add_batch_dim):
    """Convert an image file into a numpy array."""

    import PIL.Image  # from 'Pillow' package

    img = PIL.Image.open(filename)
    if shape:
        img = img.resize(shape, PIL.Image.ANTIALIAS)
    img_as_np = np.array(img).astype(np.float32)
    if not channels_last:
        # HWC to CHW
        img_as_np = np.transpose(img_as_np, (2, 0, 1))

    if add_batch_dim:
        # to NCHW or NHWC
        img_as_np = np.expand_dims(img_as_np, axis=0)

    return img_as_np


def create_random_data(shape, type, minvalue, maxvalue, seed):
    nptype = np.dtype(type)
    np.random.seed(seed)
    return ((maxvalue - minvalue) * np.random.sample(shape) + minvalue).astype(nptype)


def update_name_in_pb(filename, name, output_filename):
    """Update the name of the tensor in the pb file."""

    tensor = onnx.load_tensor(filename)
    tensor.name = name

    if not output_filename:
        output_filename = filename

    onnx.save_tensor(tensor, output_filename)


def get_arg_parser():
    parser = argparse.ArgumentParser(
        description="""
        Utilities for working with the input/output protobuf files used by the ONNX test cases and onnx_test_runner.
        These are expected to only contain a serialized TensorProto.

        dump_pb: Dumps the TensorProto data from an individual pb file, or all pb files in a directory.
        numpy_to_pb: Convert numpy array saved to a file with numpy.save() to a TensorProto, and serialize to a pb file.
        image_to_pb: Convert data from an image file into a TensorProto, and serialize to a pb file.
        random_to_pb: Create a TensorProto with random data, and serialize to a pb file.
        update_name_in_pb: Update the TensorProto.name value in a pb file.
                           Updates the input file unless --output <filename> is specified.
        """,
        formatter_class=argparse.RawDescriptionHelpFormatter
    )

    parser.add_argument('--action', help='Action to perform',
                        choices=['dump_pb', 'numpy_to_pb', 'image_to_pb', 'random_to_pb', 'update_name_in_pb'],
                        required=True)

    parser.add_argument('--input', help='The input filename or directory name')
    parser.add_argument('--name', help='The value to set TensorProto.name to if creating/updating one.')
    parser.add_argument('--output', help='Filename to serialize the TensorProto to.')

    image_to_pb_group = parser.add_argument_group('image_to_pb',
                                                  'image_to_pb specific options')
    image_to_pb_group.add_argument('--resize', default=None, type=lambda s: [int(item) for item in s.split(',')],
                                   help='Provide the shape as comma separated values to resize the image to.'
                                        ' e.g. --shape 200,200')
    image_to_pb_group.add_argument('--channels_last', action='store_true',
                                   help='Transpose image from channels first to channels last.')
    image_to_pb_group.add_argument('--add_batch_dim', action='store_true',
                                   help='Prepend a batch dimension with value of 1 to the shape. '
                                        'i.e. convert from CHW to NCHW')

    random_to_pb_group = parser.add_argument_group('random_to_pb',
                                                   'random_to_pb specific options')
    random_to_pb_group.add_argument('--shape', type=lambda s: [int(item) for item in s.split(',')],
                                    help='Provide the shape as comma separated values e.g. --shape 200,200')
    random_to_pb_group.add_argument('--datatype',
                                    help="numpy dtype value for the data type. e.g. f4=float32, i8=int64. "
                                         "See: https://docs.scipy.org/doc/numpy/reference/arrays.dtypes.html")
    random_to_pb_group.add_argument('--min_value', default=0, type=int,
                                    help="Limit the generated values to this minimum.")
    random_to_pb_group.add_argument('--max_value', default=1, type=int,
                                    help="Limit the generated values to this maximum.")
    random_to_pb_group.add_argument('--seed', default=None, type=int,
                                    help="seed to use for the random values so they're deterministic.")

    return parser


if __name__ == '__main__':
    arg_parser = get_arg_parser()
    args = arg_parser.parse_args()

    if args.action == 'dump_pb':
        if not args.input:
            print("Missing argument. Need input to be specified.", file=sys.stderr)
            sys.exit(-1)
        np.set_printoptions(precision=10)
        dump_pb(args.input)
    elif args.action == 'numpy_to_pb':
        if not args.input or not args.output or not args.name:
            print("Missing argument. Need input, output and name to be specified.", file=sys.stderr)
            sys.exit(-1)
        # read data saved with numpy
        data = np.load(args.input)
        numpy_to_pb(args.name, data, args.output)
    elif args.action == 'image_to_pb':
        if not args.input or not args.output or not args.name:
            print("Missing argument. Need input, output, name to be specified.", file=sys.stderr)
            sys.exit(-1)

        img_np = image_to_numpy(args.input, args.resize, args.channels_last, args.add_batch_dim)
        numpy_to_pb(args.name, img_np, args.output)
    elif args.action == 'random_to_pb':
        if not args.output or not args.shape or not args.datatype or not args.name:
            print("Missing argument. Need output, shape, datatype and name to be specified.", file=sys.stderr)
            sys.exit(-1)

        data = create_random_data(args.shape, args.datatype, args.min_value, args.max_value, args.seed)
        numpy_to_pb(args.name, data, args.output)
    elif args.action == 'update_name_in_pb':
        if not args.input or not args.name:
            print("Missing argument. Need input and name to be specified.", file=sys.stderr)
            sys.exit(-1)

        update_name_in_pb(args.input, args.name, args.output)
    else:
        print("Unknown action.", file=sys.stderr)
        arg_parser.print_help(sys.stderr)