# -*- 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 contextlib
import os
import tempfile

import numpy as np
import pytest

import megengine as mge
import megengine._internal as mgb
import megengine.module as M
from megengine import jit, tensor
from megengine.core.tensor import Tensor
from megengine.jit import SublinearMemoryConfig
from megengine.test import assertTensorClose


@contextlib.contextmanager
def mkstemp():
    fd, path = tempfile.mkstemp()
    try:
        os.close(fd)
        yield path
    finally:
        os.remove(path)


def load_and_compile(fpath):
    cg, _, outputs = mgb.load_comp_graph_from_file(fpath)
    inputs = mgb.cgtools.get_dep_vars(outputs, "Host2DeviceCopy")
    inputs = sorted(inputs, key=lambda i: i.name)
    outputs = list(map(mgb.copy_output, outputs))
    if len(outputs) == 1:
        (outputs,) = outputs
    return cg.compile(inputs, outputs)


def test_symbolic():
    @jit.trace(symbolic=False)
    def f(x):
        return Tensor(mgb.opr.assert_equal(x._symvar, x._symvar + 1))

    with pytest.raises(mgb.exc.MegBrainError):
        f.trace(0)

    @jit.trace(symbolic=True)
    def f(x):
        return Tensor(mgb.opr.assert_equal(x._symvar, x._symvar + 1))

    f.trace(0)


def test_dump():
    @jit.trace(symbolic=True)
    def f(x, y):
        return x * y

    f.trace(0, 0)

    with mkstemp() as out:
        f.dump(out)
        g = load_and_compile(out)

    np.testing.assert_allclose(g([1, 2, 3], [1, 2, 3]), [1, 4, 9])


def test_goptions():
    @jit.trace(symbolic=True, opt_level=0)
    def f(x):
        return x / x

    @
jit.trace(symbolic=True, opt_level=1)