# Copyright 2019 The MACE Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import os
import sys
import numpy as np
from utils import config_parser
from utils import util
from utils.util import mace_check
from py_proto import mace_pb2
from transform import base_converter as cvt
from transform import transformer
from visualize import visualize_model

device_type_map = {'cpu': cvt.DeviceType.CPU.value,
                   'gpu': cvt.DeviceType.GPU.value,
                   'dsp': cvt.DeviceType.HEXAGON.value,
                   'hta': cvt.DeviceType.HTA.value,
                   'apu': cvt.DeviceType.APU.value,
                   'cpu+gpu': cvt.DeviceType.CPU.value}

data_format_map = {
    'NONE': cvt.DataFormat.NONE,
    'NHWC': cvt.DataFormat.NHWC,
    'NCHW': cvt.DataFormat.NCHW,
    'OIHW': cvt.DataFormat.OIHW,
}

data_type_map = {
    'float32': mace_pb2.DT_FLOAT,
    'int32': mace_pb2.DT_INT32,
}


def parse_data_type(data_type, quantize):
    if quantize or data_type == 'fp32_fp32':
        return mace_pb2.DT_FLOAT
    else:
        return mace_pb2.DT_HALF


def split_shape(shape):
    if shape.strip() == "":
        return []
    else:
        return shape.split(',')


def parse_int_array_from_str(ints_str):
    return [int(i) for i in split_shape(ints_str)]


def parse_float_array_from_str(floats_str):
    return [float(i) for i in floats_str.split(',')]


def transpose_shape(shape, dst_order):
    t_shape = [0] * len(shape)
    for i in range(len(shape)):
        t_shape[i] = shape[dst_order[i]]
    return t_shape


def to_list(x):
    if isinstance(x, list):
        return x
    else:
        return [x]


def separate_params(mace_model):
    tensors = mace_model.tensors
    params = mace_pb2.NetDef()
    params.tensors.extend(tensors)

    model = mace_model
    del model.tensors[:]
    return model, params


def convert(conf, output):
    if not os.path.exists(output):
        os.mkdir(output)

    for model_name, model_conf in conf["models"].items():
        model_output = output + "/" + model_name
        if not os.path.exists(model_output):
            os.mkdir(model_output)

        subgraph = model_conf["subgraphs"][0]
        del model_conf["subgraphs"]
        model_conf.update(subgraph)

        model_file = util.download_or_get_file(model_conf["model_file_path"],
                                               model_conf[
                                                   "model_sha256_checksum"],
                                               model_output)
        model_conf["model_file_path"] = model_file
        if "weight_file_path" in model_conf:
            weight_file = util.download_or_get_file(
                model_conf["weight_file_path"],
                model_conf["weight_sha256_checksum"], model_output)
            model_conf["weight_file_path"] = weight_file
        # TODO: remove the following after quantize tool is made
        if "quantize_range_file" in model_conf:
            range_file = util.download_or_get_file(
                model_conf["quantize_range_file"],
                "", model_output)
            model_conf["quantize_range_file"] = range_file

        mace_model = convert_model(model_conf)

        try:
            visualizer = visualize_model.ModelVisualizer(model_name,
                                                         mace_model,
                                                         model_output)
            visualizer.save_html()
        except:  # noqa
            print("Failed to visualize model:", sys.exc_info()[0])

        model, params = merge_params(mace_model)

        output_model_file = model_output + "/" + model_name + ".pb"
        output_params_file = model_output + "/" + model_name + ".data"
        with open(output_model_file, "wb") as f:
            f.write(model.SerializeToString())
        with open(output_params_file, "wb") as f:
            f.write(bytearray(params))
        with open(output_model_file + "_txt", "w") as f:
            f.write(str(model))


def convert_model(conf):
    print(conf)
    platform = conf["platform"]
    mace_check(platform in ['tensorflow', 'caffe', 'onnx'],
               "platform not supported")
    runtime = conf["runtime"]
    mace_check(
        runtime in ['cpu', 'gpu', 'dsp', 'hta', 'apu', 'cpu+gpu'],
        "runtime not supported")

    option = cvt.ConverterOption()
    if "graph_optimize_options" in conf:
        option.transformer_option = conf["graph_optimize_options"]
    option.winograd = conf.get("winograd", 0)
    option.quantize = bool(conf.get("quantize", 0))
    option.quantize_large_weights = bool(conf.get("quantize_large_weights", 0))
    option.quantize_range_file = conf.get("quantize_range_file", "")
    option.change_concat_ranges = bool(conf.get("change_concat_ranges", 0))
    option.cl_mem_type = conf.get("cl_mem_type", "image")
    option.device = device_type_map[conf.get("runtime", "cpu")]
    option.data_type = parse_data_type(conf.get("data_type", "fp16_fp32"),
                                       option.quantize)
    input_tensors = to_list(conf["input_tensors"])
    input_shapes = [parse_int_array_from_str(shape) for shape in
                    to_list(conf["input_shapes"])]
    mace_check(len(input_tensors) == len(input_shapes),
               "input node count and shape count do not match")
    input_count = len(input_tensors)
    input_data_types = [data_type_map[dt] for dt in
                        to_list(conf.get("input_data_types",
                                         ["float32"]))]
    if len(input_data_types) == 1 and input_count > 1:
        input_data_types = [input_data_types[0]] * input_count
    mace_check(len(input_data_types) == input_count,
               "the number of input_data_types should be "
               "the same as input tensors")
    input_data_formats = [data_format_map[df] for df in
                          to_list(conf.get("input_data_formats",
                                           ["NHWC"]))]
    if len(input_data_formats) == 1 and input_count > 1:
        input_data_formats = [input_data_formats[0]] * input_count
    mace_check(len(input_data_formats) == input_count,
               "the number of input_data_formats should be "
               "the same as input tensors")
    input_ranges = [parse_float_array_from_str(r) for r in
                    to_list(conf.get("input_ranges",
                                     ["-1.0,1.0"]))]
    if len(input_ranges) == 1 and input_count > 1:
        input_ranges = [input_ranges[0]] * input_count
    mace_check(len(input_ranges) == input_count,
               "the number of input_ranges should be "
               "the same as input tensors")
    for i in range(len(input_tensors)):
        input_node = cvt.NodeInfo()
        input_node.name = input_tensors[i]
        input_node.shape = input_shapes[i]
        input_node.data_type = input_data_types[i]
        input_node.data_format = input_data_formats[i]
        if (input_node.data_format == cvt.DataFormat.NCHW and len(
              input_node.shape) == 4):
            input_node.shape = transpose_shape(input_node.shape, [0, 2, 3, 1])
            input_node.data_format = cvt.DataFormat.NHWC
        input_node.range = input_ranges[i]
        option.add_input_node(input_node)

    output_tensors = to_list(conf["output_tensors"])
    output_shapes = [parse_int_array_from_str(shape) for shape in
                     to_list(conf["output_shapes"])]
    mace_check(len(output_tensors) == len(output_shapes),
               "output node count and shape count do not match")
    output_count = len(output_tensors)
    output_data_types = [data_type_map[dt] for dt in
                         to_list(conf.get("output_data_types",
                                          ["float32"]))]
    if len(output_data_types) == 1 and output_count > 1:
        output_data_types = [output_data_types[0]] * output_count
    mace_check(len(output_data_types) == output_count,
               "the number of output_data_types should be "
               "the same as output tensors")
    output_data_formats = [data_format_map[df] for df in
                           to_list(conf.get("output_data_formats",
                                            ["NHWC"]))]
    if len(output_data_formats) == 1 and output_count > 1:
        output_data_formats = [output_data_formats[0]] * output_count
    mace_check(len(output_data_formats) == output_count,
               "the number of output_data_formats should be "
               "the same as output tensors")
    for i in range(len(output_tensors)):
        output_node = cvt.NodeInfo()
        output_node.name = output_tensors[i]
        output_node.shape = output_shapes[i]
        output_node.data_type = output_data_types[i]
        output_node.data_format = output_data_formats[i]
        if output_node.data_format == cvt.DataFormat.NCHW and len(
                output_node.shape) == 4:
            output_node.shape = transpose_shape(output_node.shape,
                                                [0, 2, 3, 1])
            output_node.data_format = cvt.DataFormat.NHWC
        option.add_output_node(output_node)

    if "check_tensors" in conf:
        check_tensors = to_list(conf["check_tensors"])
        check_tensors_shapes = [parse_int_array_from_str(shape) for shape in
                                to_list(conf["check_shapes"])]
        mace_check(len(check_tensors) == len(check_tensors_shapes),
                   "check tensors count and shape count do not match.")
        for i in range(len(check_tensors)):
            check_node = cvt.NodeInfo()
            check_node.name = check_tensors[i]
            check_node.shape = check_tensors_shapes[i]
            option.add_check_node(check_node)
    else:
        option.check_nodes = option.output_nodes

    option.build()

    print("Transform model to one that can better run on device")

    if platform == 'tensorflow':
        from transform import tensorflow_converter
        converter = tensorflow_converter.TensorflowConverter(
            option, conf["model_file_path"])
    elif platform == 'caffe':
        from transform import caffe_converter
        converter = caffe_converter.CaffeConverter(option,
                                                   conf["model_file_path"],
                                                   conf["weight_file_path"])
    elif platform == 'onnx':
        from transform import onnx_converter
        converter = onnx_converter.OnnxConverter(option,
                                                 conf["model_file_path"])
    else:
        mace_check(False, "Mace do not support platorm %s yet." % platform)

    output_graph_def = converter.run()
    mace_transformer = transformer.Transformer(
        option, output_graph_def)
    output_graph_def, quantize_activation_info = mace_transformer.run()

    if option.device in [cvt.DeviceType.HEXAGON.value,
                         cvt.DeviceType.HTA.value]:
        from transform import hexagon_converter
        converter = hexagon_converter.HexagonConverter(
            option, output_graph_def, quantize_activation_info)
        output_graph_def = converter.run()
    elif runtime == 'apu':
        mace_check(platform == "tensorflow",
                   "apu only support model from tensorflow")
        from transform import apu_converter
        converter = apu_converter.ApuConverter(
            option, output_graph_def, quantize_activation_info)
        output_graph_def = converter.run()

    return output_graph_def


def merge_params(net_def):
    def tensor_to_bytes(tensor):
        if tensor.data_type == mace_pb2.DT_HALF:
            data = bytearray(
                np.array(tensor.float_data).astype(np.float16).tobytes())
            tensor.data_size = len(tensor.float_data)
        elif tensor.data_type == mace_pb2.DT_FLOAT:
            data = bytearray(
                np.array(tensor.float_data).astype(np.float32).tobytes())
            tensor.data_size = len(tensor.float_data)
        elif tensor.data_type == mace_pb2.DT_INT32:
            data = bytearray(
                np.array(tensor.int32_data).astype(np.int32).tobytes())
            tensor.data_size = len(tensor.int32_data)
        elif tensor.data_type == mace_pb2.DT_UINT8:
            data = bytearray(
                np.array(tensor.int32_data).astype(np.uint8).tolist())
            tensor.data_size = len(tensor.int32_data)
        elif tensor.data_type == mace_pb2.DT_FLOAT16:
            data = bytearray(
                np.array(tensor.float_data).astype(np.float16).tobytes())
            tensor.data_size = len(tensor.float_data)
        else:
            raise Exception('Tensor data type %s not supported' %
                            tensor.data_type)
        return data

    model_data = []
    offset = 0
    for tensor in net_def.tensors:
        raw_data = tensor_to_bytes(tensor)
        if tensor.data_type != mace_pb2.DT_UINT8 and offset % 4 != 0:
            padding = 4 - offset % 4
            model_data.extend(bytearray([0] * padding))
            offset += padding

        tensor.offset = offset
        model_data.extend(raw_data)
        offset += len(raw_data)

    for tensor in net_def.tensors:
        if tensor.data_type == mace_pb2.DT_FLOAT \
                or tensor.data_type == mace_pb2.DT_HALF \
                or tensor.data_type == mace_pb2.DT_FLOAT16:
            del tensor.float_data[:]
        elif tensor.data_type == mace_pb2.DT_INT32:
            del tensor.int32_data[:]
        elif tensor.data_type == mace_pb2.DT_UINT8:
            del tensor.int32_data[:]

    return net_def, model_data


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default="",
        required=True,
        help="the path of model yaml configuration file.")
    parser.add_argument(
        '--output',
        type=str,
        default=".",
        help="output dir")
    flgs, _ = parser.parse_known_args()
    return flgs


if __name__ == '__main__':
    flags = parse_args()
    conf = config_parser.parse(flags.config)
    convert(conf, flags.output)
