# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2021 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 io

import numpy as np
import pytest

import megengine.functional as F
import megengine.module as M
import megengine.utils.comp_graph_tools as cgtools
from megengine import Parameter, Tensor
from megengine.core.tensor import megbrain_graph as G
from megengine.jit.tracing import trace
from megengine.quantization.quantize import quantize, quantize_qat
from megengine.utils.naming import AutoNaming


def _dump_and_load(func, symbolic, keep_opr_name=True):
    AutoNaming.clear()
    func = trace(func, symbolic=symbolic, capture_as_const=True)
    x = Tensor(np.ones(shape=(2, 3)))
    func(x).numpy()
    file = io.BytesIO()
    func.dump(
        file,
        optimize_for_inference=False,
        arg_names=("x",),
        keep_opr_name=keep_opr_name,
        keep_var_name=2,
    )
    file.seek(0)
    outputs = G.load_graph(file).output_vars_list
    ops = cgtools.get_oprs_seq(outputs)
    return ops


@pytest.mark.parametrize("symbolic", [False, True])
def test_auto_naming(symbolic):
    class Simple(M.Module):
        def __init__(self, name):
            super().__init__()
            self.name = name

        def forward(self, x):
            return x + x

    m = Simple("simple")
    op = _dump_and_load(m, symbolic)[-1]
    assert op.name == "simple.ADD"
    assert op.outputs[0].name == "simple.ADD"


@pytest.mark.parametrize("symbolic", [False, True])
def test_user_named_tensor(symbolic):
    class Simple(M.Module):
        def __init__(self, name):
            super().__init__()
            self.name = name
            self.k = Parameter(1.0, name="k")

        def forward(self, x):
            x = x + x
            x.name = "o_x"
            return x

    m = Simple("simple")

    op = _dump_and_load(m, symbolic)[-1]
    assert op.name == "simple.ADD"
    assert op.outputs[0].name == "o_x"


@pytest.mark.parametrize("symbolic", [False, True])
def test_user_named_param(symbolic):
    class Simple(M.Module):
        def __init__(self, name):
            super().__init__()
            self.name = name
            self.k = 
Parameter(2.0, name="k")