r"""
Copyright (c) 2021 Intel Corporation
 \file Graph partitioning
 \brief Calls Libra - Vertex-cut based graph partitioner for distirbuted training
 \author Vasimuddin Md <vasimuddin.md@intel.com>,
         Guixiang Ma <guixiang.ma@intel.com>
         Sanchit Misra <sanchit.misra@intel.com>,
         Ramanarayan Mohanty <ramanarayan.mohanty@intel.com>,
         Sasikanth Avancha <sasikanth.avancha@intel.com>
         Nesreen K. Ahmed <nesreen.k.ahmed@intel.com>
"""


import argparse
import csv
import os
import random
import sys
import time
from statistics import mean

import numpy as np

sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
from load_graph import load_ogb

import dgl
from dgl.base import DGLError
from dgl.data import load_data
from dgl.distgnn.partition import partition_graph
from dgl.distgnn.tools import load_proteins

if __name__ == "__main__":
    argparser = argparse.ArgumentParser()
    argparser.add_argument("--dataset", type=str, default="cora")
    argparser.add_argument("--num-parts", type=int, default=2)
    argparser.add_argument("--out-dir", type=str, default="./")
    args = argparser.parse_args()

    dataset = args.dataset
    num_community = args.num_parts
    out_dir = "Libra_result_" + dataset  ## "Libra_result_" prefix is mandatory
    resultdir = os.path.join(args.out_dir, out_dir)

    print("Input dataset for partitioning: ", dataset)
    if args.dataset == "ogbn-products":
        print("Loading ogbn-products")
        G, _ = load_ogb("ogbn-products")
    elif args.dataset == "ogbn-papers100M":
        print("Loading ogbn-papers100M")
        G, _ = load_ogb("ogbn-papers100M")
    elif args.dataset == "proteins":
        G = load_proteins("proteins")
    elif args.dataset == "ogbn-arxiv":
        print("Loading ogbn-arxiv")
        G, _ = load_ogb("ogbn-arxiv")
    else:
        try:
            G = load_data(args)[0]
        except:
            raise DGLError("Error: Dataset {} not found !!!".format(dataset))

    print("Done loading the graph.", flush=True)

    partition_graph(num_community, G, resultdir)