# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
""" Compile And Export MXNET Resnet18 Model With VTA As Backend """
from __future__ import absolute_import, print_function

import os
from os.path import exists
import numpy as np
from mxnet.gluon.model_zoo import vision

import tvm
from tvm import autotvm, relay
from tvm.relay import op, transform

import vta
from vta.top import graph_pack
from vta.top.graphpack import run_opt_pass

# Load VTA parameters from the vta/config/vta_config.json file
ENV = vta.get_env()
assert ENV.target.device_name == "vta"
# Dictionary lookup for when to start/end bit packing
PACK_DICT = {"resnet18_v1": ["nn.max_pool2d", "nn.global_avg_pool2d", None, None],}

# Name of Gluon model to compile
MODEL = "resnet18_v1"
assert MODEL in PACK_DICT

def merge_transform_to_mxnet_model(mod):
    """ Add Image Transform Logic Into Model """
    svalue = np.array([123., 117., 104.])
    sub_data = relay.Constant(tvm.nd.array(svalue)).astype("float32")
    dvalue = np.array([58.395, 57.12, 57.37])
    divide_data = relay.Constant(tvm.nd.array(dvalue)).astype("float32")

    data_shape = (224, 224, 3)
    data = relay.var("data", relay.TensorType(data_shape, "float32"))

    simple_net = relay.expand_dims(data, axis=0, num_newaxis=1)
    # To do, relay not support dynamic shape now, future need to add resize logic
    # simple_net = relay.image.resize(simple_net, (224, 224), "NHWC", "bilinear", "align_corners")
    simple_net = relay.subtract(simple_net, sub_data)
    simple_net = relay.divide(simple_net, divide_data)
    simple_net = relay.transpose(simple_net, ((0, 3, 1, 2)))

    #merge tranform into pretrained model network
    entry = mod["main"]
    anf = run_opt_pass(entry.body, transform.ToANormalForm())
    call = anf.value
    data, weights = call.args
    first_op = op.nn.conv2d(
        simple_net,
        weights,
        strides=call.attrs.strides,
        padding=call.attrs.padding,
        dilation=call.attrs.dilation,
        groups=call.attrs.groups,
        channels=call.attrs.channels,
        kernel_size=call.attrs.kernel_size,
        out_dtype=call.attrs.out_dtype)
    net = relay.expr.Let(anf.var, first_op, anf.body)
    net = run_opt_pass(net, transform.ToGraphNormalForm())

    mod['main'] = net
    return mod

def compile_mxnet_gulon_resnet(_env, _model):
    """ Compile Model """
    # Generate tvm IR from mxnet gluon model
    # Populate the shape and data type dictionary for ImageNet classifier input
    dtype_dict = {"data": 'float32'}
    shape_dict = {"data": (_env.BATCH, 3, 224, 224)}
    # Get off the shelf gluon model, and convert to relay
    gluon_model = vision.get_model(_model, pretrained=True)
    # Start front end compilation
    mod, params = relay.frontend.from_mxnet(gluon_model, shape_dict)
    mod = merge_transform_to_mxnet_model(mod)
    # Update shape and type dictionary
    shape_dict.update({k: v.shape for k, v in params.items()})
    dtype_dict.update({k: str(v.dtype) for k, v in params.items()})

    # Load pre-configured AutoTVM schedules
    with autotvm.tophub.context(_env.target):
        # Perform quantization in Relay
        # Note: We set opt_level to 3 in order to fold batch norm
        with relay.build_config(opt_level=3):
            with relay.quantize.qconfig(global_scale=8.0, skip_conv_layers=[0]):
                mod = relay.quantize.quantize(mod, params=params)
            # Perform graph packing and constant folding for VTA target
            relay_prog = graph_pack(
                mod["main"],
                _env.BATCH,
                _env.BLOCK_IN,
                _env.WGT_WIDTH,
                start_name=PACK_DICT[_model][0],
                stop_name=PACK_DICT[_model][1])

    # Compile Relay program with AlterOpLayout disabled
    with relay.build_config(opt_level=3, disabled_pass={"AlterOpLayout"}):
        with vta.build_config(debug_flag=0):
            graph, lib, params = relay.build(
                relay_prog, target=_env.target,
                params=params, target_host=_env.target_host)

    return graph, lib, params

def export_tvm_compile(graph, lib, params, path):
    """ Export Model"""
    if not exists(path):
        os.makedirs(path)
    lib.save(path+"/lib.o")
    with open(path+"/graph.json", "w") as graphfile:
        graphfile.write(graph)
    with open(path+"/params.params", "wb") as paramfile:
        paramfile.write(relay.save_param_dict(params))

GRAPH, LIB, PARAMS = compile_mxnet_gulon_resnet(ENV, MODEL)
export_tvm_compile(GRAPH, LIB, PARAMS, "./build/model")