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

