import os
import time
import pickle
import torch


def timelog(func):
    def printtime(*args, **argv):
        print("--------The start of {}.--------".format(func.__name__))
        t1 = time.time()
        print("> Start time: {}".format(time.strftime("%Y-%m-%d %H:%M:%S",time.localtime(t1))))
        returns = func(*args, **argv)
        t2 = time.time()
        print("> End time: {}".format(time.strftime("%Y-%m-%d %H:%M:%S",time.localtime(t2))))
        print("> Time consumption: {}s".format(t2-t1))
        print("--------The end of {}.----------".format(func.__name__))
        return returns
    return printtime


def eye_normalize(adj):
    adj_ = adj + torch.eye(adj.shape[0]).float().cuda()
    rowsum = adj_.sum(1)
    degree_mat_sqrt = torch.diag(torch.pow(rowsum, -0.5))
    adj_normalized = torch.mm(degree_mat_sqrt, torch.mm(adj_, degree_mat_sqrt))
    return adj_normalized


def load_data(data_filepath='facebook.graphs', split=False):
    data_dict = pickle.load(open(data_filepath, 'rb'))
    adj_mats = data_dict['adj_mats']
    attr_vecs = data_dict['attr_vecs']
    graphs = data_dict['graphs']
    if split:
        train_adj_mats = adj_mats[:int(len(adj_mats) * .8)]
        train_attr_vecs = attr_vecs[:int(len(attr_vecs) * .8)]
    else:
        train_adj_mats = adj_mats
        train_attr_vecs = attr_vecs
    test_adj_mats = adj_mats[int(len(adj_mats) * .8):]
    test_attr_vecs = attr_vecs[int(len(attr_vecs) * .8):]
    return graphs, train_adj_mats, test_adj_mats, train_attr_vecs, test_attr_vecs


# @timelog
def dropped_select_edges(adj_prob, adj):
    assert (adj_prob.t() == adj_prob).all(), 'Probability matrix is not SYMMETRY.'
    node_num = adj_prob.shape[0]
    edge_num = int(adj.sum()/2)
    eps = torch.tensor(1e-4).cuda()
    adj_ = (adj_prob - torch.min(adj_prob))/(torch.max(adj_prob) - torch.min(adj_prob) + eps)
    adj_ -= torch.diag(torch.diag(adj_))*2
    indexs = torch.tensor([[[j, i] for i in range(node_num)] for j in range(node_num)]).float().cuda()
    sorte_result = sorted(torch.cat([torch.unsqueeze(adj_, dim=-1), indexs], dim=-1).reshape([-1, 3]), key=lambda x: x[0], reverse=True)
    reconstructed_adj = torch.zeros(adj_.size()).float().cuda()
    for edge in sorte_result[:edge_num]:
        i = int(edge[1])
        j = int(edge[2])
        reconstructed_adj[i, j] = 1.
        reconstructed_adj[j, i] = 1.

    return reconstructed_adj


# @timelog
# def dropped_select_edges(adj_prob, adj):
#    assert (adj_prob.t() == adj_prob).all(), 'Probability matrix is not SYMMETRY.'
#    node_num = adj_prob.shape[0]
#    edge_num = int(adj.sum()/2)
#    adj_ = adj_prob - torch.diag(torch.ones(node_num)*torch.max(adj_prob)*2).cuda()
#    indexs = torch.tensor([[[j, i] for i in range(node_num)] for j in range(node_num)]).float().cuda()
#    sorte_result = sorted(torch.cat([
#                                     torch.unsqueeze(adj_, dim=-1),
#                                     indexs
#                                     ], dim=-1).reshape([-1, 3]), key=lambda x: x[0], reverse=True)
#    reconstructed_adj = torch.zeros(adj_.size()).float().cuda()
#    for edge in sorte_result[:edge_num]:
#        i = int(edge[1])
#        j = int(edge[2])
#        reconstructed_adj[i, j] = 1.
#        reconstructed_adj[j, i] = 1.
#    return reconstructed_adj


def select_edges(adj_prob, adj):
    assert (adj_prob.t() == adj_prob).all(), 'Probability matrix is not SYMMETRY.'
    node_num = adj_prob.shape[0]
    edge_num = int(adj.sum() / 2)
    adj_ = adj_prob - torch.diag(torch.ones(node_num) * torch.max(adj_prob) * 2).cuda()
    sort_index = torch.argsort(adj_.flatten())[-1-edge_num:-1:1]
    reconstructed_adj = torch.zeros(adj_.size()).float().cuda()
    for edge in sort_index:
        i = edge // node_num
        j = edge % node_num
        reconstructed_adj[i, j] = 1.
        reconstructed_adj[j, i] = 1.
    return reconstructed_adj


def cut_edges(adj_prob):
    assert (adj_prob.t() == adj_prob).all(), 'Probability matrix is not SYMMETRY.'
    adj_ = torch.round(adj_prob)
    return adj_ - torch.diag(torch.diag(adj_))

import os
import time
import pickle
import torch


def timelog(func):
    def printtime(*args, **argv):
        print("--------The start of {}.--------".format(func.__name__))
        t1 = time.time()
        print("> Start time: {}".format(time.strftime("%Y-%m-%d %H:%M:%S",time.localtime(t1))))
        returns = func(*args, **argv)
        t2 = time.time()
        print("> End time: {}".format(time.strftime("%Y-%m-%d %H:%M:%S",time.localtime(t2))))
        print("> Time consumption: {}s".format(t2-t1))
        print("--------The end of {}.----------".format(func.__name__))
        return returns
    return printtime


def eye_normalize(adj):
    adj_ = adj + torch.eye(adj.shape[0]).float().cuda()
    rowsum = adj_.sum(1)
    degree_mat_sqrt = torch.diag(torch.pow(rowsum, -0.5))
    adj_normalized = torch.mm(degree_mat_sqrt, torch.mm(adj_, degree_mat_sqrt))
    return adj_normalized


def load_data(data_filepath='facebook.graphs', split=False):
    data_dict = pickle.load(open(data_filepath, 'rb'))
    adj_mats = data_dict['adj_mats']
    attr_vecs = data_dict['attr_vecs']
    graphs = data_dict['graphs']
    if split:
        train_adj_mats = adj_mats[:int(len(adj_mats) * .8)]
        train_attr_vecs = attr_vecs[:int(len(attr_vecs) * .8)]
    else:
        train_adj_mats = adj_mats
        train_attr_vecs = attr_vecs
    test_adj_mats = adj_mats[int(len(adj_mats) * .8):]
    test_attr_vecs = attr_vecs[int(len(attr_vecs) * .8):]
    return graphs, train_adj_mats, test_adj_mats, train_attr_vecs, test_attr_vecs


# @timelog
def dropped_select_edges(adj_prob, adj):
    assert (adj_prob.t() == adj_prob).all(), 'Probability matrix is not SYMMETRY.'
    node_num = adj_prob.shape[0]
    edge_num = int(adj.sum()/2)
    eps = torch.tensor(1e-4).cuda()
    adj_ = (adj_prob - torch.min(adj_prob))/(torch.max(adj_prob) - torch.min(adj_prob) + eps)
    adj_ -= torch.diag(torch.diag(adj_))*2
    indexs = torch.tensor([[[j, i] for i in range(node_num)] for j in range(node_num)]).float().cuda()
    sorte_result = sorted(torch.cat([torch.unsqueeze(adj_, dim=-1), indexs], dim=-1).reshape([-1, 3]), key=lambda x: x[0], reverse=True)
    reconstructed_adj = torch.zeros(adj_.size()).float().cuda()
    for edge in sorte_result[:edge_num]:
        i = int(edge[1])
        j = int(edge[2])
        reconstructed_adj[i, j] = 1.
        reconstructed_adj[j, i] = 1.

    return reconstructed_adj


# @timelog
# def dropped_select_edges(adj_prob, adj):
#    assert (adj_prob.t() == adj_prob).all(), 'Probability matrix is not SYMMETRY.'
#    node_num = adj_prob.shape[0]
#    edge_num = int(adj.sum()/2)
#    adj_ = adj_prob - torch.diag(torch.ones(node_num)*torch.max(adj_prob)*2).cuda()
#    indexs = torch.tensor([[[j, i] for i in range(node_num)] for j in range(node_num)]).float().cuda()
#    sorte_result = sorted(torch.cat([
#                                     torch.unsqueeze(adj_, dim=-1),
#                                     indexs
#                                     ], dim=-1).reshape([-1, 3]), key=lambda x: x[0], reverse=True)
#    reconstructed_adj = torch.zeros(adj_.size()).float().cuda()
#    for edge in sorte_result[:edge_num]:
#        i = int(edge[1])
#        j = int(edge[2])
#        reconstructed_adj[i, j] = 1.
#        reconstructed_adj[j, i] = 1.
#    return reconstructed_adj


def select_edges(adj_prob, adj):
    assert (adj_prob.t() == adj_prob).all(), 'Probability matrix is not SYMMETRY.'
    node_num = adj_prob.shape[0]
    edge_num = int(adj.sum() / 2)
    adj_ = adj_prob - torch.diag(torch.ones(node_num) * torch.max(adj_prob) * 2).cuda()
    sort_index = torch.argsort(adj_.flatten())[-1-edge_num:-1:1]
    reconstructed_adj = torch.zeros(adj_.size()).float().cuda()
    for edge in sort_index:
        i = edge // node_num
        j = edge % node_num
        reconstructed_adj[i, j] = 1.
        reconstructed_adj[j, i] = 1.
    return reconstructed_adj


def cut_edges(adj_prob):
    assert (adj_prob.t() == adj_prob).all(), 'Probability matrix is not SYMMETRY.'
    adj_ = torch.round(adj_prob)
    return adj_ - torch.diag(torch.diag(adj_))

# @timelog # @timelog # def dropped_select_edges(adj_prob, adj): #    assert (adj_prob.t() == adj_prob).all(), 'Probability matrix is not SYMMETRY.' #    node_num = adj_prob.shape[0] #    edge_num = int(adj.sum()/2) #    adj_ = adj_prob - torch.diag(torch.ones(node_num)*torch.max(adj_prob)*2).cuda() #    indexs = torch.tensor([[[j, i] for i in range(node_num)] for j in range(node_num)]).float().cuda() #    sorte_result = sorted(torch.cat([ #                                     torch.unsqueeze(adj_, dim=-1), #                                     indexs #                                     ], dim=-1).reshape([-1, 3]), key=lambda x: x[0], reverse=True) #    reconstructed_adj = torch.zeros(adj_.size()).float().cuda() #    for edge in sorte_result[:edge_num]: #        i = int(edge[1]) #        j = int(edge[2]) #        reconstructed_adj[i, j] = 1. #        reconstructed_adj[j, i] = 1. #    return reconstructed_adj