""" load dataset from ogb """

import argparse
from ogb.linkproppred import DglLinkPropPredDataset
import time

def load_from_ogbl_with_name(name):    
    choices = ['ogbl-collab', 'ogbl-ddi', 'ogbl-ppa', 'ogbl-citation']
    assert name in choices, "name must be selected from " + str(choices)
    dataset = DglLinkPropPredDataset(name)
    return dataset[0]

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--name', type=str,
        choices=['ogbl-collab', 'ogbl-ddi', 'ogbl-ppa', 'ogbl-citation'],
        default='ogbl-collab',
        help="name of datasets by ogb")
    args = parser.parse_args()

    print("loading graph... it might take some time")
    name = args.name
    g = load_from_ogbl_with_name(name=name)

    try:
        w = g.edata['edge_weight']
        weighted = True
    except:
        weighted = False

    
    edge_num = g.edges()[0].shape[0]
    src = list(g.edges()[0])
    tgt = list(g.edges()[1])
    if weighted:
        weight = list(g.edata['edge_weight'])

    print("writing...")
    start_time = time.time()
    with open(name + "-net.txt", "w") as f:
        for i in range(edge_num):
            if weighted:
                f.write(str(src[i].item()) + " "\
                    +str(tgt[i].item()) + " "\
                    +str(weight[i].item()) + "\n")
            else:
                f.write(str(src[i].item()) + " "\
                    +str(tgt[i].item()) + " "\
                    +"1\n")
    print("writing used time: %d s" % int(time.time() - start_time))