r"""
Copyright (c) 2021 Intel Corporation
 \file distgnn/tools/tools.py
 \brief Tools for use in Libra graph partitioner.
 \author Vasimuddin Md <vasimuddin.md@intel.com>
"""

import os
import random

import requests
import torch as th
from scipy.io import mmread

import dgl
from dgl.base import DGLError
from dgl.data.utils import load_graphs, save_graphs, save_tensors


def rep_per_node(prefix, num_community):
    """
    Used on Libra partitioned data.
    This function reports number of split-copes per node (replication) of
    a partitioned graph
    Parameters
    ----------
    prefix: Partition folder location (contains replicationlist.csv)
    num_community: number of partitions or communities
    """
    ifile = os.path.join(prefix, "replicationlist.csv")
    fhandle = open(ifile, "r")
    r_dt = {}

    fline = fhandle.readline()  ## reading first line, contains the comment.
    print(fline)
    for line in fhandle:
        if line[0] == "#":
            raise DGLError("[Bug] Read Hash char in rep_per_node func.")

        node = line.strip("\n")
        if r_dt.get(node, -100) == -100:
            r_dt[node] = 1
        else:
            r_dt[node] += 1

    fhandle.close()
    ## sanity checks
    for v in r_dt.values():
        if v >= num_community:
            raise DGLError(
                "[Bug] Unexpected event in rep_per_node() in tools.py."
            )

    return r_dt


def download_proteins():
    """
    Downloads the proteins dataset
    """
    print("Downloading dataset...")
    print("This might a take while..")
    url = "https://portal.nersc.gov/project/m1982/GNN/"
    file_name = "subgraph3_iso_vs_iso_30_70length_ALL.m100.propermm.mtx"
    url = url + file_name
    try:
        req = requests.get(url)
    except:
        raise DGLError(
            "Error: Failed to download Proteins dataset!! Aborting.."
        )

    with open("proteins.mtx", "wb") as handle:
        handle.write(req.content)


def proteins_mtx2dgl():
    """
    This function converts Proteins dataset from mtx to dgl format.
    """
    print("Converting mtx2dgl..")
    print("This might a take while..")
    a_mtx = mmread("proteins.mtx")
    coo = a_mtx.tocoo()
    u = th.tensor(coo.row, dtype=th.int64)
    v = th.tensor(coo.col, dtype=th.int64)
    g = dgl.DGLGraph()

    g.add_edges(u, v)

    n = g.num_nodes()
    feat_size = 128  ## arbitrary number
    feats = th.empty([n, feat_size], dtype=th.float32)

    ## arbitrary numbers
    train_size = 1000000
    test_size = 500000
    val_size = 5000
    nlabels = 256

    train_mask = th.zeros(n, dtype=th.bool)
    test_mask = th.zeros(n, dtype=th.bool)
    val_mask = th.zeros(n, dtype=th.bool)
    label = th.zeros(n, dtype=th.int64)

    for i in range(train_size):
        train_mask[i] = True

    for i in range(test_size):
        test_mask[train_size + i] = True

    for i in range(val_size):
        val_mask[train_size + test_size + i] = True

    for i in range(n):
        label[i] = random.choice(range(nlabels))

    g.ndata["feat"] = feats
    g.ndata["train_mask"] = train_mask
    g.ndata["test_mask"] = test_mask
    g.ndata["val_mask"] = val_mask
    g.ndata["label"] = label

    return g


def save(g, dataset):
    """
    This function saves input dataset to dgl format
    Parameters
    ----------
    g : graph to be saved
    dataset : output folder name
    """
    print("Saving dataset..")
    part_dir = os.path.join("./" + dataset)
    node_feat_file = os.path.join(part_dir, "node_feat.dgl")
    part_graph_file = os.path.join(part_dir, "graph.dgl")
    os.makedirs(part_dir, mode=0o775, exist_ok=True)
    save_tensors(node_feat_file, g.ndata)
    save_graphs(part_graph_file, [g])
    print("Graph saved successfully !!")


def load_proteins(dataset):
    """
    This function downloads, converts, and load Proteins graph dataset
    Parameter
    ---------
    dataset: output folder name
    """
    part_dir = dataset
    graph_file = os.path.join(part_dir + "/graph.dgl")

    if not os.path.exists("proteins.mtx"):
        download_proteins()
    if not os.path.exists(graph_file):
        g = proteins_mtx2dgl()
        save(g, dataset)
    ## load
    graph = load_graphs(graph_file)[0][0]
    return graph