"""Dataset for Graph Isomorphism Network(GIN)
(chen jun): Used for compacted graph kernel dataset in GIN
Data sets include:
MUTAG, COLLAB, IMDBBINARY, IMDBMULTI, NCI1, PROTEINS, PTC, REDDITBINARY, REDDITMULTI5K
https://github.com/weihua916/powerful-gnns/blob/master/dataset.zip
"""

import os
import numpy as np

from .. import backend as F

from .dgl_dataset import DGLBuiltinDataset
from .utils import loadtxt, save_graphs, load_graphs, save_info, load_info, download, extract_archive
from ..utils import retry_method_with_fix
from ..convert import graph as dgl_graph

class GINDataset(DGLBuiltinDataset):
    """Datasets for Graph Isomorphism Network (GIN)
    Adapted from `<https://github.com/weihua916/powerful-gnns/blob/master/dataset.zip>`_.

    The dataset contains the compact format of popular graph kernel datasets.
    For more graph kernel datasets, see :class:`TUDataset`.

    Parameters
    ---------
    name: str
        dataset name, one of
        (``'MUTAG'``, ``'COLLAB'``, \
        ``'IMDBBINARY'``, ``'IMDBMULTI'``, \
        ``'NCI1'``, ``'PROTEINS'``, ``'PTC'``, \
        ``'REDDITBINARY'``, ``'REDDITMULTI5K'``)
    self_loop: bool
        add self to self edge if true
    degree_as_nlabel: bool
        take node degree as label and feature if true

    Examples
    --------
    >>> data = GINDataset(name='MUTAG', self_loop=False)

    The dataset instance is an iterable

    >>> len(data)
    188
    >>> g, label = data[128]
    >>> g
    Graph(num_nodes=13, num_edges=26,
          ndata_schemes={'label': Scheme(shape=(), dtype=torch.int64), 'attr': Scheme(shape=(7,), dtype=torch.float64)}
          edata_schemes={})
    >>> label
    tensor(1)

    Batch the graphs and labels for mini-batch training

    >>> graphs, labels = zip(*[data[i] for i in range(16)])
    >>> batched_graphs = dgl.batch(graphs)
    >>> batched_labels = torch.tensor(labels)
    >>> batched_graphs
    Graph(num_nodes=330, num_edges=748,
          ndata_schemes={'label': Scheme(shape=(), dtype=torch.int64), 'attr': Scheme(shape=(7,), dtype=torch.float64)}
          edata_schemes={})
    """

    def __init__(self, name, self_loop, degree_as_nlabel=False,
                 raw_dir=None, force_reload=False, verbose=False):

        self._name = name  # MUTAG
        gin_url = 'https://raw.githubusercontent.com/weihua916/powerful-gnns/master/dataset.zip'
        self.ds_name = 'nig'

        self.self_loop = self_loop
        self.graphs = []
        self.labels = []

        # relabel
        self.glabel_dict = {}
        self.nlabel_dict = {}
        self.elabel_dict = {}
        self.ndegree_dict = {}

        # global num
        self.N = 0  # total graphs number
        self.n = 0  # total nodes number
        self.m = 0  # total edges number

        # global num of classes
        self.gclasses = 0
        self.nclasses = 0
        self.eclasses = 0
        self.dim_nfeats = 0

        # flags
        self.degree_as_nlabel = degree_as_nlabel
        self.nattrs_flag = False
        self.nlabels_flag = False

        super(GINDataset, self).__init__(name=name, url=gin_url, hash_key=(name, self_loop, degree_as_nlabel),
                                         raw_dir=raw_dir, force_reload=force_reload, verbose=verbose)

    @property
    def raw_path(self):
        return os.path.join(self.raw_dir, 'GINDataset')

    def download(self):
        r""" Automatically download data and extract it.
        """
        zip_file_path = os.path.join(self.raw_dir, 'GINDataset.zip')
        download(self.url, path=zip_file_path)
        extract_archive(zip_file_path, self.raw_path)

    def __len__(self):
        """Return the number of graphs in the dataset."""
        return len(self.graphs)

    def __getitem__(self, idx):
        """Get the idx-th sample.

        Parameters
        ---------
        idx : int
            The sample index.

        Returns
        -------
        (:class:`dgl.Graph`, Tensor)
            The graph and its label.
        """
        return self.graphs[idx], self.labels[idx]

    def _file_path(self):
        return os.path.join(self.raw_dir, "GINDataset", 'dataset', self.name, "{}.txt".format(self.name))

    def process(self):
        """ Loads input dataset from dataset/NAME/NAME.txt file
        """
        if self.verbose:
            print('loading data...')
        self.file = self._file_path()
        with open(self.file, 'r') as f:
            # line_1 == N, total number of graphs
            self.N = int(f.readline().strip())

            for i in range(self.N):
                if (i + 1) % 10 == 0 and self.verbose is True:
                    print('processing graph {}...'.format(i + 1))

                grow = f.readline().strip().split()
                # line_2 == [n_nodes, l] is equal to
                # [node number of a graph, class label of a graph]
                n_nodes, glabel = [int(w) for w in grow]

                # relabel graphs
                if glabel not in self.glabel_dict:
                    mapped = len(self.glabel_dict)
                    self.glabel_dict[glabel] = mapped

                self.labels.append(self.glabel_dict[glabel])

                g = dgl_graph(([], []))
                g.add_nodes(n_nodes)

                nlabels = []  # node labels
                nattrs = []  # node attributes if it has
                m_edges = 0

                for j in range(n_nodes):
                    nrow = f.readline().strip().split()

                    # handle edges and attributes(if has)
                    tmp = int(nrow[1]) + 2  # tmp == 2 + #edges
                    if tmp == len(nrow):
                        # no node attributes
                        nrow = [int(w) for w in nrow]
                        nattr = None
                    elif tmp > len(nrow):
                        nrow = [int(w) for w in nrow[:tmp]]
                        nattr = [float(w) for w in nrow[tmp:]]
                        nattrs.append(nattr)
                    else:
                        raise Exception('edge number is incorrect!')

                    # relabel nodes if it has labels
                    # if it doesn't have node labels, then every nrow[0]==0
                    if not nrow[0] in self.nlabel_dict:
                        mapped = len(self.nlabel_dict)
                        self.nlabel_dict[nrow[0]] = mapped

                    nlabels.append(self.nlabel_dict[nrow[0]])

                    m_edges += nrow[1]
                    g.add_edges(j, nrow[2:])

                    # add self loop
                    if self.self_loop:
                        m_edges += 1
                        g.add_edges(j, j)

                    if (j + 1) % 10 == 0 and self.verbose is True:
                        print(
                            'processing node {} of graph {}...'.format(
                                j + 1, i + 1))
                        print('this node has {} edgs.'.format(
                            nrow[1]))

                if nattrs != []:
                    nattrs = np.stack(nattrs)
                    g.ndata['attr'] = F.tensor(nattrs)
                    self.nattrs_flag = True
                else:
                    nattrs = None

                g.ndata['label'] = F.tensor(nlabels)
                if len(self.nlabel_dict) > 1:
                    self.nlabels_flag = True

                assert g.number_of_nodes() == n_nodes

                # update statistics of graphs
                self.n += n_nodes
                self.m += m_edges

                self.graphs.append(g)

        self.labels = F.tensor(self.labels)
        # if no attr
        if not self.nattrs_flag:
            if self.verbose:
                print('there are no node features in this dataset!')
            label2idx = {}
            # generate node attr by node degree
            if self.degree_as_nlabel:
                if self.verbose:
                    print('generate node features by node degree...')
                nlabel_set = set([])
                for g in self.graphs:
                    # actually this label shouldn't be updated
                    # in case users want to keep it
                    # but usually no features means no labels, fine.
                    g.ndata['label'] = g.in_degrees()
                    # extracting unique node labels
                    nlabel_set = nlabel_set.union(set([F.as_scalar(nl) for nl in g.ndata['label']]))

                nlabel_set = list(nlabel_set)
                # in case the labels/degrees are not continuous number
                self.ndegree_dict = {
                    nlabel_set[i]: i
                    for i in range(len(nlabel_set))
                }
                label2idx = self.ndegree_dict
            # generate node attr by node label
            else:
                if self.verbose:
                    print('generate node features by node label...')
                label2idx = self.nlabel_dict

            for g in self.graphs:
                g.ndata['attr'] = F.tensor(np.zeros((
                    g.number_of_nodes(), len(label2idx))))
                g.ndata['attr'][range(g.number_of_nodes()), [label2idx[F.as_scalar(F.reshape(nl, (1,)))] for nl in g.ndata['label']]] = 1

        # after load, get the #classes and #dim
        self.gclasses = len(self.glabel_dict)
        self.nclasses = len(self.nlabel_dict)
        self.eclasses = len(self.elabel_dict)
        self.dim_nfeats = len(self.graphs[0].ndata['attr'][0])

        if self.verbose:
            print('Done.')
            print(
                """
                -------- Data Statistics --------'
                #Graphs: %d
                #Graph Classes: %d
                #Nodes: %d
                #Node Classes: %d
                #Node Features Dim: %d
                #Edges: %d
                #Edge Classes: %d
                Avg. of #Nodes: %.2f
                Avg. of #Edges: %.2f
                Graph Relabeled: %s
                Node Relabeled: %s
                Degree Relabeled(If degree_as_nlabel=True): %s \n """ % (
                    self.N, self.gclasses, self.n, self.nclasses,
                    self.dim_nfeats, self.m, self.eclasses,
                    self.n / self.N, self.m / self.N, self.glabel_dict,
                    self.nlabel_dict, self.ndegree_dict))

    def save(self):
        graph_path = os.path.join(self.save_path, 'gin_{}_{}.bin'.format(self.name, self.hash))
        info_path = os.path.join(self.save_path, 'gin_{}_{}.pkl'.format(self.name, self.hash))
        label_dict = {'labels': self.labels}
        info_dict = {'N': self.N,
                     'n': self.n,
                     'm': self.m,
                     'self_loop': self.self_loop,
                     'gclasses': self.gclasses,
                     'nclasses': self.nclasses,
                     'eclasses': self.eclasses,
                     'dim_nfeats': self.dim_nfeats,
                     'degree_as_nlabel': self.degree_as_nlabel,
                     'glabel_dict': self.glabel_dict,
                     'nlabel_dict': self.nlabel_dict,
                     'elabel_dict': self.elabel_dict,
                     'ndegree_dict': self.ndegree_dict}
        save_graphs(str(graph_path), self.graphs, label_dict)
        save_info(str(info_path), info_dict)

    def load(self):
        graph_path = os.path.join(self.save_path, 'gin_{}_{}.bin'.format(self.name, self.hash))
        info_path = os.path.join(self.save_path, 'gin_{}_{}.pkl'.format(self.name, self.hash))
        graphs, label_dict = load_graphs(str(graph_path))
        info_dict = load_info(str(info_path))

        self.graphs = graphs
        self.labels = label_dict['labels']

        self.N = info_dict['N']
        self.n = info_dict['n']
        self.m = info_dict['m']
        self.self_loop = info_dict['self_loop']
        self.gclasses = info_dict['gclasses']
        self.nclasses = info_dict['nclasses']
        self.eclasses = info_dict['eclasses']
        self.dim_nfeats = info_dict['dim_nfeats']
        self.glabel_dict = info_dict['glabel_dict']
        self.nlabel_dict = info_dict['nlabel_dict']
        self.elabel_dict = info_dict['elabel_dict']
        self.ndegree_dict = info_dict['ndegree_dict']
        self.degree_as_nlabel = info_dict['degree_as_nlabel']

    def has_cache(self):
        graph_path = os.path.join(self.save_path, 'gin_{}_{}.bin'.format(self.name, self.hash))
        info_path = os.path.join(self.save_path, 'gin_{}_{}.pkl'.format(self.name, self.hash))
        if os.path.exists(graph_path) and os.path.exists(info_path):
            return True
        return False
