# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
import random

import megengine as mge
import megengine.functional as F
import megengine.module as M
import numpy as np
from megengine.jit import trace


def dump_mge_model(net, data, fpath="test_model", optimize_for_inference=False):
    if mge.__version__ <= "0.6.0":

        @trace(symbolic=True)
        def inference(data, *, net):
            net.eval()
            output = net(data)
            return output

        inference.trace(data, net=net)
        mge_result = inference(data, net=net).numpy()
        inference.dump(
            fpath + ".mge",
            arg_names=["data"],
            optimize_for_inference=optimize_for_inference,
        )
        return mge_result
    else:
        mge_result = net(mge.tensor(data))
        net.eval()
        mge_result = net(mge.tensor(data))

        @
trace(symbolic=True, capture_as_const=True)