import io

import numpy as np

import megengine.core.tensor.megbrain_graph as G
import megengine.functional as F
import megengine.module as M
import megengine.utils.network_node as N
from megengine.jit.tracing import trace
from megengine.tensor import Tensor
from megengine.utils.comp_graph_tools import GraphInference
from megengine.utils.network import Network as Net
from megengine.utils.network import as_oprnode, set_symbolic_shape
from megengine.utils.network_node import Host2DeviceCopy, VarNode


def test_metadata():
    x = Tensor(0)

    @trace(symbolic=True, capture_as_const=True)
    def fwd(x):
        return x * 2

    fwd(x)

    orig_model = io.BytesIO()
    fwd.dump(orig_model, user_info="test", optimize_for_inference=False)
    orig_model.seek(0)
    graph = Net.load(orig_model)
    assert graph.metadata == {
        "user_info": "test",
        "graph_modified": False,  # False: tracing.dump
        "optimized_for_inference": False,
    }

    orig_model.seek(0)
    graph.dump(
        orig_model,
        user_info={"str": "x", "tensor": x, "module": M.Module, "none": None},
        optimize_for_inference=True,
        enable_nchw4=True,
        enable_ioc16=True,
    )
    orig_model.seek(0)
    graph = Net.load(orig_model)
    assert graph.metadata == {
        "user_info": {"str": "x", "tensor": x, "module": M.Module, "none": None},
        "graph_modified": True,  # True: Network.dump
        "optimized_for_inference": True,
        "enable_nchw4": True,
        "enable_ioc16": True,
    }

    orig_model.seek(0)
    fwd.dump(orig_model, enable_metadata=False)
    orig_model.seek(0)
    graph = Net.load(orig_model)
    assert graph.metadata is None


def test_replace_var():

    a = Tensor([1, 2])
    b = Tensor([3, 4])

    @trace(symbolic=True, capture_as_const=True)
    def fwd(a, b):
        return (a + b) * 2

    fwd(a, b)
    orig_model = io.BytesIO()
    fwd.dump(
        orig_model, arg_names=["a", "b"], output_names="o", optimize_for_inference=False
    )
    orig_model.seek(0)

    graph = Net.load(orig_model)
    vara = graph.var_filter.name("a").as_unique()
    varb = graph.var_filter.name("b").as_unique()

    out = F.mul(vara, varb)
    out = F.relu(out)

    opnode = list(graph.opr_filter.has_input(vara))
    repl_dict = {opnode[0].outputs[0]: out}
    graph.replace_vars(repl_dict)

    modified_model = io.BytesIO()
    graph.dump(modified_model)
    modified_model.seek(0)
    load_graph = GraphInference(modified_model)

    out = load_graph.run(a, b)
    np.testing.assert_equal(out["o"], [6, 16])


def test_replace_opr():

    a = Tensor([1, 2])
    b = Tensor([3, 4])

    @trace(symbolic=True, capture_as_const=True)
    def fwd(a, b):
        return (a + b) * 2

    fwd(a, b)
    orig_model = io.BytesIO()
    fwd.dump(
        orig_model, arg_names=["a", "b"], output_names="o", optimize_for_inference=False
    )
    orig_model.seek(0)

    graph = Net.load(orig_model)
    vara = graph.var_filter.name("a").as_unique()
    varb = graph.var_filter.name("b").as_unique()

    out1 = 
F.sub(vara, varb)