# Source code for deeprobust.graph.data.attacked_data

    import numpy as np
    import scipy.sparse as sp
    import os.path as osp
    import os
    import urllib.request
    import warnings
    import json

    class PtbDataset:
        """Dataset class manages pre-attacked adjacency matrix on different datasets. Currently only support metattack under 5% perturbation. Note metattack is generated by deeprobust/graph/global_attack/metattack.py. While PrePtbDataset provides pre-attacked graph generate by Zugner, https://github.com/danielzuegner/gnn-meta-attack. The attacked graphs are downloaded from https://github.com/ChandlerBang/pytorch-gnn-meta-attack/tree/master/pre-attacked.

        Parameters
        ----------
        root :
            root directory where the dataset should be saved.
        name :
            dataset name. It can be choosen from ['cora', 'citeseer', 'cora_ml', 'polblogs', 'pubmed']
        attack_method :
            currently this class only support metattack. User can pass 'meta', 'metattack' or 'mettack' since all of them will be interpreted as the same attack.
        seed :
            random seed for splitting training/validation/test.

        Examples
        --------

        >>> from deeprobust.graph.data import Dataset, PtbDataset
        >>> data = Dataset(root='/tmp/', name='cora')
        >>> adj, features, labels = data.adj, data.features, data.labels
        >>> perturbed_data = PtbDataset(root='/tmp/',
                                name='cora',
                                attack_method='meta')
        >>> perturbed_adj = perturbed_data.adj

        """

        def __init__(self, root, name, attack_method='mettack'):
            assert attack_method in ['mettack', 'metattack', 'meta'], \
                'Currently the database only stores graphs perturbed by 5% mettack'

            self.name = name.lower()
            assert self.name in ['cora', 'citeseer', 'polblogs'], \
                'Currently only support cora, citeseer, polblogs'

            self.attack_method = 'mettack' # attack_method
            self.url = 'https://raw.githubusercontent.com/ChandlerBang/pytorch-gnn-meta-attack/master/pre-attacked/{}_{}_0.05.npz'.format(self.name, self.attack_method)
            self.root = osp.expanduser(osp.normpath(root))
            self.data_filename = osp.join(root,
                    '{}_{}_0.05.npz'.format(self.name, self.attack_method))
            self.adj = self.load_data()

        def load_data(self):
            if not osp.exists(self.data_filename):
                self.download_npz()
            print('Loading {} dataset perturbed by 0.05 mettack...'.format(self.name))
            adj = sp.load_npz(self.data_filename)
            warnings.warn('''the adjacency matrix is perturbed, using the data splits under seed 15(default seed for deeprobust.graph.data.Dataset), so if you are going to verify the attacking performance, you should use the same data splits''')
            return adj

        def download_npz(self):
            print('Dowloading from {} to {}'.format(self.url, self.data_filename))
            try:
                urllib.request.urlretrieve(self.url, self.data_filename)
            except:
                raise Exception('''Download failed! Make sure you have
                        stable Internet connection and enter the right name''')

    class PrePtbDataset:
        """Dataset class manages pre-attacked adjacency matrix on different datasets. Note metattack is generated by deeprobust/graph/global_attack/metattack.py. While PrePtbDataset provides pre-attacked graph generate by Zugner, https://github.com/danielzuegner/gnn-meta-attack. The attacked graphs are downloaded from https://github.com/ChandlerBang/Pro-GNN/tree/master/meta.

        Parameters
        ----------
        root :
            root directory where the dataset should be saved.
        name :
            dataset name. It can be choosen from ['cora', 'citeseer', 'polblogs', 'pubmed']
        attack_method :
            currently this class only support metattack and  nettack. Note 'meta', 'metattack' or 'mettack' will be interpreted as the same attack.
        seed :
            random seed for splitting training/validation/test.

        Examples
        --------

        >>> from deeprobust.graph.data import Dataset, PrePtbDataset
        >>> data = Dataset(root='/tmp/', name='cora')
        >>> adj, features, labels = data.adj, data.features, data.labels
        >>> # Load meta attacked data
        >>> perturbed_data = PrePtbDataset(root='/tmp/',
                                name='cora',
                                attack_method='meta',
                                ptb_rate=0.05)
        >>> perturbed_adj = perturbed_data.adj
        >>> # Load nettacked data
        >>> perturbed_data = PrePtbDataset(root='/tmp/',
                                name='cora',
                                attack_method='nettack',
                                ptb_rate=1.0)
        >>> perturbed_adj = perturbed_data.adj
        >>> target_nodes = perturbed_data.target_nodes
        """

        def __init__(self, root, name, attack_method='meta', ptb_rate=0.05):

            if attack_method == 'mettack' or attack_method == 'metattack':
                attack_method = 'meta'

            assert attack_method in ['meta', 'nettack'], \
                ' Currently the database only stores graphs perturbed by metattack, nettack'
            # assert attack_method in ['meta'], \
            #     ' Currently the database only stores graphs perturbed by metattack. Will update nettack soon.'

            self.name = name.lower()
            assert self.name in ['cora', 'citeseer', 'polblogs', 'pubmed', 'cora_ml'], \
                'Currently only support cora, citeseer, pubmed, polblogs, cora_ml'

            self.attack_method = attack_method
            self.ptb_rate = ptb_rate
            self.url = 'https://raw.githubusercontent.com/ChandlerBang/Pro-GNN/master/{}/{}_{}_adj_{}.npz'.\
                    format(self.attack_method, self.name, self.attack_method, self.ptb_rate)
            self.root = osp.expanduser(osp.normpath(root))
            self.data_filename = osp.join(root,
                    '{}_{}_adj_{}.npz'.format(self.name, self.attack_method, self.ptb_rate))
            self.target_nodes = None
            self.adj = self.load_data()

        def load_data(self):
            if not osp.exists(self.data_filename):
                self.download_npz()
            print('Loading {} dataset perturbed by {} {}...'.format(self.name, self.ptb_rate, self.attack_method))

            if self.attack_method == 'meta':
                warnings.warn("The pre-attacked graph is perturbed under the data splits provided by ProGNN. So if you are going to verify the attacking performance, you should use the same data splits  (setting='prognn').")
                adj = sp.load_npz(self.data_filename)

            if self.attack_method == 'nettack':
                # assert True, "Will update pre-attacked data by nettack soon"
                warnings.warn("The pre-attacked graph is perturbed under the data splits provided by ProGNN. So if you are going to verify the attacking performance, you should use the same data splits  (setting='prognn').")
                adj = sp.load_npz(self.data_filename)
                self.target_nodes = self.get_target_nodes()
            return adj

        def get_target_nodes(self):
            """Get target nodes incides, which is the nodes with degree > 10 in the test set."""
            url = 'https://raw.githubusercontent.com/ChandlerBang/Pro-GNN/master/nettack/{}_nettacked_nodes.json'.format(self.name)
            json_file = osp.join(self.root,
                    '{}_nettacked_nodes.json'.format(self.name))

            if not osp.exists(json_file):
                self.download_file(url, json_file)
            # with open(f'/mnt/home/jinwei2/Projects/nettack/{dataset}_nettacked_nodes.json', 'r') as f:
            with open(json_file, 'r') as f:
                idx = json.loads(f.read())
            return idx["attacked_test_nodes"]

        def download_file(self, url, file):
            print('Dowloading from {} to {}'.format(url, file))
            try:
                urllib.request.urlretrieve(url, file)
            except:
                raise Exception("Download failed! Make sure you have \
                        stable Internet connection and enter the right name")

        def download_npz(self):
            print('Dowloading from {} to {}'.format(self.url, self.data_filename))
            try:
                urllib.request.urlretrieve(self.url, self.data_filename)
            except:
                raise Exception("Download failed! Make sure you have \
                        stable Internet connection and enter the right name")

    class RandomAttack():

        def __init__(self):
            self.name = 'RandomAttack'

        def attack(self, adj, ratio=0.4):
            print('random attack: ratio=%s' % ratio)
            modified_adj = self._random_add_edges(adj, ratio)
            return modified_adj

        def _random_add_edges(self, adj, add_ratio):

            def sample_zero_forever(mat):
                nonzero_or_sampled = set(zip(*mat.nonzero()))
                while True:
                    t = tuple(np.random.randint(0, mat.shape[0], 2))
                    if t not in nonzero_or_sampled:
                        yield t
                        nonzero_or_sampled.add(t)
                        nonzero_or_sampled.add((t[1], t[0]))

            def sample_zero_n(mat, n=100):
                itr = sample_zero_forever(mat)
                return [next(itr) for _ in range(n)]

            assert np.abs(adj - adj.T).sum() == 0, "Input graph is not symmetric"
            non_zeros = [(x, y) for x,y in np.argwhere(adj != 0) if x < y] # (x, y)

            added = sample_zero_n(adj, n=int(add_ratio * len(non_zeros)))
            for x, y in added:
                adj[x, y] = 1
                adj[y, x] = 1
            return adj

    if __name__ == '__main__':
        perturbed_data = PrePtbDataset(root='/tmp/',
                name='cora',
                attack_method='meta',
                ptb_rate=0.05)
        perturbed_adj = perturbed_data.adj