from numpy import mintypecode
import utils
import torch
import train_ as train
import pyro
import layers
import torch.nn.functional as F
import numpy as np 


class GraphEnv:
    def __init__(self, dataset, k_hops = 1, cuda = True):
        A, features, labels, _, _, _ = utils.load_data(dataset)
        saved_model, idx_train, idx_test, acc = train.run_gcn(dataset = dataset)
        self.victim_model, self.idx_train, self.idx_test = saved_model, idx_train, idx_test
        self.adj, self.features, self.labels = A.to_dense(), features, labels
        self.cuda = cuda
        self.n_class = len(np.unique(self.labels.cpu().numpy()))
        self.adj[self.adj>0] = 1
        if cuda:
            self.adj = self.adj.cuda()
            self.features = self.features.cuda()
            self.labels = self.labels.cuda()

        for param in self.victim_model.parameters():
            param.requires_grad = False
        
        predictions = self.victim_model(self.features, self.adj)
        self.idx_test = [i for i in self.idx_test if torch.argmax(predictions[i]) == labels[i]]
        self.sub_graphs, self.neighbor_index = self.k_hop_subgraphs(A, k_hops, self.idx_test)
        self.acc = acc

    def get_sub_graph(self, index):
        if self.cuda:
            return self.sub_graphs[index].cuda(), self.neighbor_index[index]
        else:
            return self.sub_graphs[index], self.neighbor_index[index]

    def get_loss(self, features, adj, node_index):
        logits = self.victim_model(features, normalize(adj))[node_index]
        return F.nll_loss(logits.reshape(1, self.n_class), self.labels[node_index].reshape(1))

    def get_reward(self, node_index, previous_adj, current_adj, new_feature):
        previous_logit = self.victim_model(new_feature, normalize(previous_adj))[node_index]
        current_logit = self.victim_model(new_feature, normalize(current_adj))[node_index]
        if self.cuda:
            label = torch.tensor([self.labels[node_index]]).cuda()
        else:
            label = torch.tensor([self.labels[node_index]])
        return  F.nll_loss(current_logit.unsqueeze(0), label) - \
            F.nll_loss(previous_logit.unsqueeze(0), label), utils.accuracy(current_logit.unsqueeze(0), label)
    
    def k_hop_subgraphs(self, adj, k, node_index):
        resulted_adj = adj
        dense_adj = adj.to_dense()
        sub_graphs = []
        neighbor_index = []
        for _ in range(k - 1):
            resulted_adj = torch.sparse.mm(resulted_adj, adj)
        for node in node_index:
            neighbors = resulted_adj[node].coalesce().indices()[0]
            raw_adj = dense_adj[neighbors][:,neighbors]
            raw_adj = raw_adj + raw_adj.t()
            raw_adj[raw_adj > 0] = 1
            sub_graphs.append(raw_adj)
            neighbor_index.append(neighbors)
        return sub_graphs, neighbor_index

    def get_mask(self, node_index, current_adj):
        current_edges = current_adj[-1, :-1]
        node_neighbor = self.adj[node_index].clone()
        node_neighbor[current_edges == 1] = 0.
        return node_neighbor

def normalize(matrix, node_budget = 1):
    res = matrix.clone()
    res[:,-1] = res[:,-1] * node_budget
    rowsum = res.sum(1)
    return res.div(rowsum.unsqueeze(1))


def merge_node(adj, sampled_edge):
    res = F.pad(adj, (0, 1, 0, 1), 'constant', 0)
    res[-1, -1] = 1
    res[-1, :-1] = sampled_edge
    res[:-1, -1] = sampled_edge.t().squeeze()
    return res

class Node_Generator(torch.nn.Module):
    def __init__(self, in_features, hidden):
        super(Node_Generator, self).__init__()
        self.gc1 = layers.GraphConvolution(in_features, hidden)
        self.generator = torch.nn.Linear(3 * hidden, in_features)

    def forward(self, x, adj, node_index):
        h = F.relu(self.gc1(x, adj))
        h = torch.cat((h.sum(0), h.max(0).values, h[node_index]))
        feature_dist = self.generator(h)
        dist = pyro.distributions.RelaxedBernoulliStraightThrough(temperature=torch.tensor(1.).cuda(), logits = feature_dist)
        generated_feature = dist.rsample()
        return generated_feature/(generated_feature.sum() + 0.0001), generated_feature.sum()

class Edge_Sampler(torch.nn.Module):
    def __init__(self, in_features, hidden):
        super(Edge_Sampler, self).__init__()
        self.gc1 = layers.GraphConvolution(in_features, hidden)
        self.gc2 = layers.GraphConvolution(hidden, hidden)
        self.regressor = torch.nn.Linear(in_features + hidden * 2, 1)

    def forward(self, x, adj, masks, node_index):
        adj_ = adj.detach().clone()
        nor_adj = normalize(adj_)
        node_neighbor_mask = masks
        sampled_feature = x[-1].expand((x.shape[0], -1))
        h = F.relu(self.gc1(x, nor_adj))
        h = F.relu(self.gc2(h, nor_adj))
        node_embedding = h[node_index].expand((x.shape[0], -1))
        h = torch.cat((h, sampled_feature, node_embedding), dim = 1)[:-1]
        feature_dist = self.regressor(h).t().squeeze()
        feature_dist = feature_dist + node_neighbor_mask * 10000
        feature_dist = F.softmax(feature_dist, dim = -1)
        dist = pyro.distributions.RelaxedOneHotCategoricalStraightThrough(temperature=torch.tensor(1.).cuda(), probs = feature_dist, validate_args = False)
        sample = dist.rsample()
        log_prob = dist.log_prob(sample)
        adj_[-1, :-1] += sample
        adj_[:-1,-1] += sample
        return adj_, log_prob

class Value_Predictor(torch.nn.Module):
    def __init__(self, feature_dim, out_dim=64, n_class = 7):
        super(Value_Predictor, self).__init__()
        self.gc1 = layers.GraphConvolution(feature_dim, out_dim)
        self.gc2 = layers.GraphConvolution(out_dim, n_class)
        self.regressor = torch.nn.Linear(out_dim + n_class, n_class)
        self.n_class = n_class

    def forward(self, x, adj, node_index, label):
        nor_adj = normalize(adj)
        h1 = F.relu(self.gc1(x, nor_adj))
        logits = self.gc2(h1, nor_adj)
        h = torch.cat((h1[node_index], logits[node_index]), dim = -1)
        h = self.regressor(h)
        return F.nll_loss(h.reshape(1, self.n_class), label.reshape(1))