from itertools import product

import dgl
import torch
from dgl import DGLGraph
from dgl.contrib.data import load_data
from dgl.data import citation_graph
from runtime.dgl.gat import GAT, GATSPMV
from runtime.dgl.gcn import GCN, GCNSPMV
from runtime.dgl.hidden import HiddenPrint
from runtime.dgl.rgcn import RGCN, RGCNSPMV
from runtime.dgl.train import train_runtime

if torch.cuda.is_available():
    device = torch.device('cuda')
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
    device = torch.device('mps')
else:
    device = torch.device('cpu')

with HiddenPrint():
    Cora = citation_graph.load_cora()
    CiteSeer = citation_graph.load_citeseer()
    PubMed = citation_graph.load_pubmed()
    MUTAG = load_data('mutag')  # fair comparison

# One training run before we start tracking duration to warm up GPU.
g = DGLGraph(Cora.graph)
g.set_n_initializer(dgl.init.zero_initializer)
g.add_edges(g.nodes(), g.nodes())
norm = torch.pow(g.in_degrees().float(), -0.5)
norm[torch.isinf(norm)] = 0
g.ndata['norm'] = norm.unsqueeze(1).to(device)
model = GCNSPMV(g, Cora.features.shape[1], Cora.num_labels).to(device)
train_runtime(model, Cora, epochs=200, device=device)

for d, Net in product([Cora, CiteSeer, PubMed], [GCN, GCNSPMV, GAT, GATSPMV]):
    g = DGLGraph(d.graph)
    g.set_n_initializer(dgl.init.zero_initializer)
    g.add_edges(g.nodes(), g.nodes())
    norm = torch.pow(g.in_degrees().float(), -0.5)
    norm[torch.isinf(norm)] = 0
    g.ndata['norm'] = norm.unsqueeze(1).to(device)
    model = Net(g, d.features.shape[1], d.num_labels).to(device)
    t = train_runtime(model, d, epochs=200, device=device)
    print(f'{d.name} - {Net.__name__}: {t:.2f}s')

for d, Net in product([MUTAG], [RGCN, RGCNSPMV]):
    g = DGLGraph()
    g.add_nodes(d.num_nodes)
    g.add_edges(d.edge_src, d.edge_dst)
    edge_type = torch.from_numpy(d.edge_type).to(device)
    edge_norm = torch.from_numpy(d.edge_norm).to(device)
    g.edata.update({'type': edge_type, 'norm': edge_norm})
    g.ndata['id'] = torch.arange(d.num_nodes, dtype=torch.long, device=device)
    model = Net(g, d.num_nodes, d.num_classes, d.num_rels)
    t = train_runtime(model, d, epochs=200, device=device)
    print(f'{d.name} - {Net.__name__}: {t:.2f}s')
