# -*- 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.
from typing import List, Tuple

import numpy as np

import megengine._internal as mgb
import megengine.functional as F
from megengine import Graph, jit
from megengine.module import Linear, Module
from megengine.test import assertTensorClose

from .env import modified_environ


class MLP(Module):
    def __init__(self):
        super().__init__()
        self.dense0 = Linear(28, 50)
        self.dense1 = Linear(50, 20)

    def forward(self, x):
        x = self.dense0(x)
        x = F.relu(x)
        x = self.dense1(x)
        return x


def has_gpu(num=1):
    try:
        mgb.comp_node("gpu{}".format(num - 1))
    except mgb.MegBrainError:
        return False

    return True


def randomNp(*args):
    for arg in args:
        assert isinstance(arg, int)
    return np.random.random(args)


def randomTorch(*args):
    import torch  # pylint: disable=import-outside-toplevel

    for arg in args:
        assert isinstance(arg, int)
    return torch.tensor(randomNp(*args), dtype=torch.float32)


def graph_mode(*modes):
    if not set(modes).issubset({"eager", "static"}):
        raise ValueError("graph mode must be in (eager, static)")

    def decorator(func):
        def wrapper(*args, **kwargs):
            if "eager" in set(modes):
                func(*args, **kwargs)
            if "static" in set(modes):
                with 
Graph()