# Source code for deeprobust.graph.rl.nipa_env

    """
        This part of code is adopted from https://github.com/Hanjun-Dai/graph_adversarial_attack (Copyright (c) 2018 Dai, Hanjun and Li, Hui and Tian, Tian and Huang, Xin and Wang, Lin and Zhu, Jun and Song, Le)
        but modified to be integrated into the repository.
    """

    import os
    import sys
    import numpy as np
    import torch
    import networkx as nx
    import random
    from torch.nn.parameter import Parameter
    import torch.nn as nn
    import torch.nn.functional as F
    import torch.optim as optim
    from tqdm import tqdm
    from copy import deepcopy
    import pickle as cp
    from deeprobust.graph.utils import *
    import scipy.sparse as sp
    from scipy.sparse.linalg.eigen.arpack import eigsh
    from deeprobust.graph import utils
    from deeprobust.graph.rl.env import *

    class NodeInjectionEnv(NodeAttackEnv):
        """Node attack environment. It executes an action and then change the
        environment status (modify the graph).
        """

        def __init__(self, features, labels, idx_train, idx_val, dict_of_lists, classifier, ratio=0.01, parallel_size=1, reward_type='binary'):
            """number of injected nodes: ratio*|V|
               number of modifications: ratio*|V|*|D_avg|
            """
            # super(NodeInjectionEnv, self).__init__(features, labels, all_targets, list_action_space, classifier, num_mod, reward_type)
            super(NodeInjectionEnv, self).__init__(features, labels, idx_val, dict_of_lists, classifier)
            self.parallel_size = parallel_size

            degrees = np.array([len(d) for n, d in dict_of_lists.items()])
            N = len(degrees[degrees > 0])
            avg_degree = degrees.sum() / N
            self.n_injected = len(degrees) - N
            assert self.n_injected == int(ratio * N)

            self.ori_adj_size = N
            self.n_perturbations = int(self.n_injected * avg_degree)
            print("number of perturbations: {}".format(self.n_perturbations))
            self.all_nodes = np.arange(N)
            self.injected_nodes = self.all_nodes[-self.n_injected: ]
            self.previous_acc = [1] * parallel_size

            self.idx_train = np.hstack((idx_train, self.injected_nodes))
            self.idx_val = idx_val

            self.modified_label_list = []
            for i in range(self.parallel_size):
                self.modified_label_list.append(labels[-self.n_injected: ].clone())

        def init_overall_steps(self):
            self.overall_steps = 0
            self.modified_list = []
            for i in range(self.parallel_size):
                self.modified_list.append(ModifiedGraph())

        def setup(self):
            self.n_steps = 0
            self.first_nodes = None
            self.second_nodes = None
            self.rewards = None
            self.binary_rewards = None
            self.list_acc_of_all = []

        def step(self, actions, inference=False):
            '''
                run actions and get reward
            '''
            if self.first_nodes is None: # pick the first node of edge
                assert (self.n_steps + 1) % 3 == 1
                self.first_nodes = actions[:]

            if (self.n_steps + 1) % 3 == 2:
                self.second_nodes = actions[:]
                for i in range(self.parallel_size):
                    # add an edge from the graph
                    self.modified_list[i].add_edge(self.first_nodes[i], actions[i], 1.0)

            if (self.n_steps + 1) % 3 == 0:
                for i in range(self.parallel_size):
                    # change label
                    self.modified_label_list[i][self.first_nodes[i] - self.ori_adj_size] = actions[i]

                self.first_nodes = None
                self.second_nodes = None

            self.n_steps += 1
            self.overall_steps += 1

            if not inference:
                if self.isActionFinished() :
                    rewards = []
                    for i in (range(self.parallel_size)):
                        device = self.labels.device
                        extra_adj = self.modified_list[i].get_extra_adj(device=device)
                        adj = self.classifier.norm_tool.norm_extra(extra_adj)
                        labels = torch.cat((self.labels, self.modified_label_list[i]))
                        # self.classifier.fit(self.features, adj, labels, self.idx_train, self.idx_val, normalize=False)
                        self.classifier.fit(self.features, adj, labels, self.idx_train, self.idx_val, normalize=False, patience=30)
                        output = self.classifier(self.features, adj)
                        loss, correct = loss_acc(output, self.labels, self.idx_val, avg_loss=False)
                        acc = correct.sum()
                        # r = 1 if self.previous_acc[i] - acc > 0.01  else -1
                        r = 1 if self.previous_acc[i] - acc > 0  else -1
                        self.previous_acc[i] = acc
                        rewards.append(r)
                        self.rewards = np.array(rewards).astype(np.float32)

        def sample_pos_rewards(self, num_samples):
            assert self.list_acc_of_all is not None
            cands = []

            for i in range(len(self.list_acc_of_all)):
                succ = np.where( self.list_acc_of_all[i] < 0.9 )[0]

                for j in range(len(succ)):

                    cands.append((i, self.all_targets[succ[j]]))

            if num_samples > len(cands):
                return cands
            random.shuffle(cands)
            return cands[0:num_samples]

        def uniformRandActions(self):
            act_list = []
            for i in range(self.parallel_size):
                if self.first_nodes is None:
                    # a1: choose a node from injected nodes
                    cur_action = np.random.choice(self.injected_nodes)

                if self.first_nodes is not None and self.second_nodes is None:
                    # a2: choose a node from all nodes
                    cur_action = np.random.randint(len(self.list_action_space))
                    while (self.first_nodes[i], cur_action) in self.modified_list[i].edge_set:
                        cur_action = np.random.randint(len(self.list_action_space))

                if self.first_nodes is not None and self.second_nodes is not None:
                    # a3: choose label
                    cur_action = np.random.randint(self.labels.cpu().max() + 1)

                act_list.append(cur_action)
            return act_list

        def isActionFinished(self):
            if (self.n_steps) % 3 == 0 and self.n_steps != 0:
                return True
            return False

        def isTerminal(self):
            if self.overall_steps == 3 * self.n_perturbations:
                return True
            return False

        def getStateRef(self):
            return list(zip(self.modified_list, self.modified_label_list))

        def cloneState(self):
            return list(zip(deepcopy(self.modified_list), deepcopy(self.modified_label_list)))