import backend as F
from test_utils import parametrize_idtype

import dgl


@parametrize_idtype
def test_heterograph_merge(idtype):
    g1 = (
        dgl.heterograph({("a", "to", "b"): ([0, 1], [1, 0])})
        .astype(idtype)
        .to(F.ctx())
    )
    g1_n_edges = g1.num_edges(etype="to")
    g1.nodes["a"].data["nh"] = F.randn((2, 3))
    g1.nodes["b"].data["nh"] = F.randn((2, 3))
    g1.edges["to"].data["eh"] = F.randn((2, 3))

    g2 = (
        dgl.heterograph({("a", "to", "b"): ([1, 2, 3], [2, 3, 5])})
        .astype(idtype)
        .to(F.ctx())
    )
    g2.nodes["a"].data["nh"] = F.randn((4, 3))
    g2.nodes["b"].data["nh"] = F.randn((6, 3))
    g2.edges["to"].data["eh"] = F.randn((3, 3))
    g2.add_nodes(3, ntype="a")
    g2.add_nodes(3, ntype="b")

    m = dgl.merge([g1, g2])

    # Check g2's edges and nodes were added to g1's in m.
    m_us = F.asnumpy(m.edges()[0][g1_n_edges:])
    g2_us = F.asnumpy(g2.edges()[0])
    assert all(m_us == g2_us)
    m_vs = F.asnumpy(m.edges()[1][g1_n_edges:])
    g2_vs = F.asnumpy(g2.edges()[1])
    assert all(m_vs == g2_vs)
    for ntype in m.ntypes:
        assert m.num_nodes(ntype=ntype) == max(
            g1.num_nodes(ntype=ntype), g2.num_nodes(ntype=ntype)
        )

        # Check g1's node data was updated with g2's in m.
        for key in m.nodes[ntype].data:
            g2_n_nodes = g2.num_nodes(ntype=ntype)
            updated_g1_ndata = F.asnumpy(m.nodes[ntype].data[key][:g2_n_nodes])
            g2_ndata = F.asnumpy(g2.nodes[ntype].data[key])
            assert all((updated_g1_ndata == g2_ndata).flatten())

    # Check g1's edge data was updated with g2's in m.
    for key in m.edges["to"].data:
        updated_g1_edata = F.asnumpy(m.edges["to"].data[key][g1_n_edges:])
        g2_edata = F.asnumpy(g2.edges["to"].data[key])
        assert all((updated_g1_edata == g2_edata).flatten())