from sklearn.model_selection import train_test_split

dataset = 'flickr'
sparse_node_threshold = 5


def read_flickr_graph():
    node_dict = dict()
    graph_dict = dict()
    class_dict = dict()
    node_class = dict()
    node_index = 1
    with open('./dataset/' + dataset + '/nodes.csv', "r") as fr:
        lines = fr.readlines()
        for line in lines:
            temp = list(line.strip('\n').split(' '))
            node_dict[temp[0]] = node_index
            node_index += 1
    with open('./dataset/' + dataset + '/nodes.csv', "r") as fr:
        lines = fr.readlines()
        for line in lines:
            temp = list(line.strip('\n').split(' '))
            if node_dict[temp[0]] not in graph_dict.keys():
                graph_dict[node_dict[temp[0]]] = set()
    edge_num = 0
    with open('./dataset/' + dataset + '/edges.csv', "r") as fr:
        lines = fr.readlines()
        for line in lines:
            edge_num += 1
            temp = list(line.strip('\n').split(','))
            graph_dict[node_dict[temp[0]]].add(node_dict[temp[1]])
            graph_dict[node_dict[temp[1]]].add(node_dict[temp[0]])
    class_index = 0
    with open('./dataset/' + dataset + '/groups.csv', "r") as fr:
        lines = fr.readlines()
        for line in lines:
            temp = list(line.strip('\n').split(' '))
            if temp[0] not in class_dict.keys():
                class_dict[temp[0]] = class_index
                class_index += 1
    with open('./dataset/' + dataset + '/group-edges.csv', "r") as fr:
        lines = fr.readlines()
        for line in lines:
            temp = list(line.strip('\n').split(','))
            if node_dict[temp[0]] not in node_class.keys():
                node_class[node_dict[temp[0]]] = set()
                node_class[node_dict[temp[0]]].add(class_dict[temp[1]])
            else:
                node_class[node_dict[temp[0]]].add(class_dict[temp[1]])
    print('dataset ' + dataset)
    print('total node number: %d' % len(graph_dict.keys()))
    print('total edge number: %d' % edge_num)
    print('total class number: %d' % len(class_dict.keys()))
    return node_dict, graph_dict, node_class, class_dict


def read_email_graph():
    node_dict = dict()
    graph_dict = dict()
    node_class = dict()
    classset = set()
    node_index = 1
    with open('./dataset/' + dataset + '/edges.txt', "r") as fr:
        lines = fr.readlines()
        for line in lines:
            temp = list(line.strip('\n').split(' '))
            if temp[0] not in node_dict.keys():
                node_dict[temp[0]] = node_index
                node_index += 1
            if temp[1] not in node_dict.keys():
                node_dict[temp[1]] = node_index
                node_index += 1
    # reverse_node_dict = dict(zip(node_dict.values(), node_dict.keys()))
    # with open('./dataset/' + dataset + '/node_dict.txt', "w") as fw:
    #     for n in node_dict:
    #         fw.write(str(n) + ' ' + str(node_dict) + '\n')
    for n in node_dict.keys():
        graph_dict[node_dict[n]] = set()
    edge_num = 0
    with open('./dataset/' + dataset + '/edges.txt', "r") as fr:
        lines = fr.readlines()
        for line in lines:
            edge_num += 1
            temp = list(line.strip('\n').split(' '))
            graph_dict[node_dict[temp[0]]].add(node_dict[temp[1]])
            graph_dict[node_dict[temp[1]]].add(node_dict[temp[0]])
    with open('./dataset/' + dataset + '/labels.txt', "r") as fr:
        lines = fr.readlines()
        for line in lines:
            temp = list(line.strip('\n').split(' '))
            if temp[0] in node_dict.keys():
                node_class[node_dict[temp[0]]] = temp[1]
                classset.add(temp[1])
    print(len(classset))
    print('dataset ' + dataset)
    print('total node number: %d' % len(graph_dict.keys()))
    print('total edge number: %d' % edge_num)
    return node_dict, graph_dict, node_class


def read_wiki_graph():
    node_dict = dict()
    graph_dict = dict()
    node_class = dict()
    node_index = 1
    with open('./dataset/' + dataset + '/Wiki_edgelist.txt', "r") as fr:
        lines = fr.readlines()
        for line in lines:
            temp = list(line.strip('\n').split('\t'))
            if temp[0] not in node_dict.keys():
                node_dict[temp[0]] = node_index
                node_index += 1
            if temp[1] not in node_dict.keys():
                node_dict[temp[1]] = node_index
                node_index += 1
    # reverse_node_dict = dict(zip(node_dict.values(), node_dict.keys()))
    # with open('./dataset/' + dataset + '/node_dict.txt', "w") as fw:
    #     for n in node_dict:
    #         fw.write(str(n) + ' ' + str(node_dict) + '\n')
    for n in node_dict.keys():
        graph_dict[node_dict[n]] = set()
    edge_num = 0
    with open('./dataset/' + dataset + '/Wiki_edgelist.txt', "r") as fr:
        lines = fr.readlines()
        for line in lines:
            edge_num += 1
            temp = list(line.strip('\n').split('\t'))
            graph_dict[node_dict[temp[0]]].add(node_dict[temp[1]])
            graph_dict[node_dict[temp[1]]].add(node_dict[temp[0]])
    with open('./dataset/' + dataset + '/Wiki_category.txt', "r") as fr:
        lines = fr.readlines()
        for line in lines:
            temp = list(line.strip('\n').split('\t'))
            if temp[0] in node_dict.keys():
                node_class[node_dict[temp[0]]] = temp[1]
    print('dataset ' + dataset)
    print('total node number: %d' % len(graph_dict.keys()))
    print('total edge number: %d' % edge_num)
    return node_dict, graph_dict, node_class


def select_sparse_node(graph_dict):
    sparse_node_set = set()
    for i in graph_dict.keys():
        if len(graph_dict[i]) <= sparse_node_threshold:
            sparse_node_set.add(i)
    return sparse_node_set


def node_type_info(graph_dict, sparse_node_set):
    node_type_dict = dict()
    for n in graph_dict.keys():
        if n in sparse_node_set:
            node_type_dict[n] = 'sparse'
        else:
            neighbor = len(graph_dict[n])
            for adj in graph_dict[n]:
                if adj in sparse_node_set:
                    neighbor -= 1
            if neighbor > sparse_node_threshold:
                node_type_dict[n] = 'dense'
            else:
                node_type_dict[n] = 'middle'
    s_num = 0
    d_num = 0
    m_num = 0
    for n in node_type_dict.keys():
        if node_type_dict[n] == 'sparse':
            s_num += 1
        if node_type_dict[n] == 'dense':
            d_num += 1
        if node_type_dict[n] == 'middle':
            m_num += 1
    return node_type_dict


def generate_train_test_candidate_flickr(graph_dict, sparse_node_set, node_type_dict):
    train_node_candidate_set = set()
    test_node_candidate_set_all = set()
    test_node_candidate_set_1 = set()
    test_node_candidate_set_2 = set()
    test_node_candidate_set_3 = set()
    test_node_candidate_set_4 = set()
    test_node_candidate_set_5 = set()
    test_nodes_list = list()
    test_node_candidate_num = 0
    for n in sparse_node_set:
        dense_num = 0
        middle_num = 0
        sparse_num = 0
        for adj in graph_dict[n]:
            if node_type_dict[adj] == 'dense':
                dense_num += 1
            elif node_type_dict[adj] == 'middle':
                middle_num += 1
            else:
                sparse_num += 1
        if middle_num == 0 and sparse_num == 0:
            test_node_candidate_set_all.add(n)
            if len(graph_dict[n]) == 1:
                test_node_candidate_set_1.add(n)
            if len(graph_dict[n]) == 2:
                test_node_candidate_set_2.add(n)
            if len(graph_dict[n]) == 3:
                test_node_candidate_set_3.add(n)
            if len(graph_dict[n]) == 4:
                test_node_candidate_set_4.add(n)
            if len(graph_dict[n]) == 5:
                test_node_candidate_set_5.add(n)
            test_node_candidate_num += 1
    print(len(test_node_candidate_set_1))
    print(len(test_node_candidate_set_2))
    print(len(test_node_candidate_set_3))
    print(len(test_node_candidate_set_4))
    print(len(test_node_candidate_set_5))
    print(test_node_candidate_num)
    for n in graph_dict.keys():
        if n in sparse_node_set:
            continue
        else:
            if node_type_dict[n] == 'middle':
                continue
            else:
                neighbor = len(graph_dict[n])
                for adj in graph_dict[n]:
                    if adj in sparse_node_set:
                        neighbor -= 1
                if 36 < neighbor <= 56:
                    train_node_candidate_set.add(n)
    # flickr do not have 1 shot node
    # test_nodes_list_1, _, _, _ = train_test_split(list(test_node_candidate_set_1),
    #                                               range(len(test_node_candidate_set_1)), train_size=157,
    #                                               random_state=19)
    # test_nodes_list.extend(test_nodes_list_1)
    test_nodes_list_2, _, _, _ = train_test_split(list(test_node_candidate_set_2),
                                                  range(len(test_node_candidate_set_2)), train_size=320,
                                                  random_state=19)
    test_nodes_list.extend(test_nodes_list_2)
    test_nodes_list_3, _, _, _ = train_test_split(list(test_node_candidate_set_3),
                                                  range(len(test_node_candidate_set_3)), train_size=260,
                                                  random_state=19)
    test_nodes_list.extend(test_nodes_list_3)
    test_nodes_list_4, _, _, _ = train_test_split(list(test_node_candidate_set_4),
                                                  range(len(test_node_candidate_set_4)), train_size=227,
                                                  random_state=19)
    test_nodes_list.extend(test_nodes_list_4)
    test_nodes_list_5, _, _, _ = train_test_split(list(test_node_candidate_set_5),
                                                  range(len(test_node_candidate_set_5)), train_size=193,
                                                  random_state=19)
    test_nodes_list.extend(test_nodes_list_5)
    test_node_candidate_set = set(test_node_candidate_set_all)
    return train_node_candidate_set, test_node_candidate_set


def generate_train_test_candidate_email(graph_dict, sparse_node_set, node_type_dict):
    test_node_candidate_num = 0
    test_node_candidate_set_all = set()
    train_node_candidate_set = set()
    test_node_candidate_set_1 = set()
    test_node_candidate_set_2 = set()
    test_node_candidate_set_3 = set()
    test_node_candidate_set_4 = set()
    test_node_candidate_set_5 = set()
    test_nodes_list = list()
    for n in sparse_node_set:
        dense_num = 0
        middle_num = 0
        sparse_num = 0
        for adj in graph_dict[n]:
            if node_type_dict[adj] == 'dense':
                dense_num += 1
            elif node_type_dict[adj] == 'middle':
                middle_num += 1
            else:
                sparse_num += 1
        if middle_num == 0 and sparse_num == 0:
            test_node_candidate_set_all.add(n)
            if len(graph_dict[n]) == 1:
                test_node_candidate_set_1.add(n)
            if len(graph_dict[n]) == 2:
                test_node_candidate_set_2.add(n)
            if len(graph_dict[n]) == 3:
                test_node_candidate_set_3.add(n)
            if len(graph_dict[n]) == 4:
                test_node_candidate_set_4.add(n)
            if len(graph_dict[n]) == 5:
                test_node_candidate_set_5.add(n)
            test_node_candidate_num += 1
    print(test_node_candidate_num)
    print(len(test_node_candidate_set_1))
    print(len(test_node_candidate_set_2))
    print(len(test_node_candidate_set_3))
    print(len(test_node_candidate_set_4))
    print(len(test_node_candidate_set_5))
    for n in graph_dict.keys():
        if n in sparse_node_set:
            continue
        else:
            if node_type_dict[n] == 'middle':
                continue
            else:
                neighbor = len(graph_dict[n])
                for adj in graph_dict[n]:
                    if adj in sparse_node_set:
                        neighbor -= 1
                if 12 < neighbor <= 32:
                    train_node_candidate_set.add(n)
    test_nodes_list_1, _, _, _ = train_test_split(list(test_node_candidate_set_1),
                                                  range(len(test_node_candidate_set_1)), train_size=43,
                                                  random_state=19)
    test_nodes_list.extend(test_nodes_list_1)
    test_nodes_list_2, _, _, _ = train_test_split(list(test_node_candidate_set_2),
                                                  range(len(test_node_candidate_set_2)), train_size=17,
                                                  random_state=19)
    test_nodes_list.extend(test_nodes_list_2)
    test_nodes_list_3, _, _, _ = train_test_split(list(test_node_candidate_set_3),
                                                  range(len(test_node_candidate_set_3)), train_size=15,
                                                  random_state=19)
    test_nodes_list.extend(test_nodes_list_3)
    test_nodes_list_4, _, _, _ = train_test_split(list(test_node_candidate_set_4),
                                                  range(len(test_node_candidate_set_4)), train_size=11,
                                                  random_state=19)
    test_nodes_list.extend(test_nodes_list_4)
    test_nodes_list_5, _, _, _ = train_test_split(list(test_node_candidate_set_5),
                                                  range(len(test_node_candidate_set_5)), train_size=14,
                                                  random_state=19)
    test_nodes_list.extend(test_nodes_list_5)
    test_node_candidate_set = set(test_nodes_list)
    print('train node candidate num is %d' % len(train_node_candidate_set))
    print('test node candidate num is %d' % len(test_node_candidate_set))
    print('=================================')
    return train_node_candidate_set, test_node_candidate_set


def generate_train_test_candidate_wiki(graph_dict, sparse_node_set, node_type_dict):
    test_node_candidate_num = 0
    test_node_candidate_set_all = set()
    train_node_candidate_set = set()
    test_node_candidate_set_1 = set()
    test_node_candidate_set_2 = set()
    test_node_candidate_set_3 = set()
    test_node_candidate_set_4 = set()
    test_node_candidate_set_5 = set()
    test_nodes_list = list()
    for n in sparse_node_set:
        dense_num = 0
        middle_num = 0
        sparse_num = 0
        for adj in graph_dict[n]:
            if node_type_dict[adj] == 'dense':
                dense_num += 1
            elif node_type_dict[adj] == 'middle':
                middle_num += 1
            else:
                sparse_num += 1
        if middle_num == 0 and sparse_num == 0:
            test_node_candidate_set_all.add(n)
            if len(graph_dict[n]) == 1:
                test_node_candidate_set_1.add(n)
            if len(graph_dict[n]) == 2:
                test_node_candidate_set_2.add(n)
            if len(graph_dict[n]) == 3:
                test_node_candidate_set_3.add(n)
            if len(graph_dict[n]) == 4:
                test_node_candidate_set_4.add(n)
            if len(graph_dict[n]) == 5:
                test_node_candidate_set_5.add(n)
            test_node_candidate_num += 1
    print(test_node_candidate_num)
    print(len(test_node_candidate_set_1))
    print(len(test_node_candidate_set_2))
    print(len(test_node_candidate_set_3))
    print(len(test_node_candidate_set_4))
    print(len(test_node_candidate_set_5))
    for n in graph_dict.keys():
        if n in sparse_node_set:
            continue
        else:
            if node_type_dict[n] == 'middle':
                continue
            else:
                neighbor = len(graph_dict[n])
                for adj in graph_dict[n]:
                    if adj in sparse_node_set:
                        neighbor -= 1
                if 7 < neighbor <= 20:
                    train_node_candidate_set.add(n)
    test_nodes_list_1, _, _, _ = train_test_split(list(test_node_candidate_set_1),
                                                  range(len(test_node_candidate_set_1)), train_size=30,
                                                  random_state=19)
    test_nodes_list.extend(test_nodes_list_1)
    test_nodes_list_2, _, _, _ = train_test_split(list(test_node_candidate_set_2),
                                                  range(len(test_node_candidate_set_2)), train_size=24,
                                                  random_state=19)
    test_nodes_list.extend(test_nodes_list_2)
    test_nodes_list_3, _, _, _ = train_test_split(list(test_node_candidate_set_3),
                                                  range(len(test_node_candidate_set_3)), train_size=20,
                                                  random_state=19)
    test_nodes_list.extend(test_nodes_list_3)
    test_nodes_list_4, _, _, _ = train_test_split(list(test_node_candidate_set_4),
                                                  range(len(test_node_candidate_set_4)), train_size=13,
                                                  random_state=19)
    test_nodes_list.extend(test_nodes_list_4)
    test_nodes_list_5, _, _, _ = train_test_split(list(test_node_candidate_set_5),
                                                  range(len(test_node_candidate_set_5)), train_size=13,
                                                  random_state=19)
    test_nodes_list.extend(test_nodes_list_5)
    test_node_candidate_set = set(test_nodes_list)
    print('train node candidate num is %d' % len(train_node_candidate_set))
    print('test node candidate num is %d' % len(test_node_candidate_set))
    print('=================================')
    return train_node_candidate_set, test_node_candidate_set


def generate_train_test_set(train_node_candidate_set, test_node_candidate_set):
    train_node_set = train_node_candidate_set
    return train_node_set, test_node_candidate_set


def save_graph(graph_dict, graph_dir):
    with open(graph_dir, 'w') as fw:
        for key in graph_dict.keys():
            fw.write(str(key))
            for item in graph_dict[key]:
                fw.write(' ' + str(item))
            fw.write('\n')


def save_node(node_set, node_save_dir):
    with open(node_save_dir, 'w') as fw:
        for n in node_set:
            fw.write(str(n))
            fw.write('\n')


def save_dense_graph(graph_dict, graph_dir, node_type_dict):
    with open(graph_dir, 'w') as fw:
        for key in graph_dict.keys():
            fw.write(str(key))
            for item in graph_dict[key]:
                if node_type_dict[item] == 'dense':
                    fw.write(' ' + str(item))
            fw.write('\n')


def save_node_class_info(node_info, save_dir):
    with open(save_dir, 'w') as fw:
        for key in node_info.keys():
            fw.write(str(key))
            for item in node_info[key]:
                fw.write(' ' + str(item))
            fw.write('\n')


def train_flickr():
    node2id, graph, node_class_info, class2id = read_flickr_graph()
    sparse_node = select_sparse_node(graph)
    node_type = node_type_info(graph, sparse_node)
    train_node_candidate, test_node_candidate = generate_train_test_candidate_flickr(graph, sparse_node, node_type)
    train_node, test_node = generate_train_test_set(train_node_candidate, test_node_candidate)
    save_node(train_node, './data/' + dataset + '/train.csv')
    save_node(test_node, './data/' + dataset + '/test.csv')
    save_node_class_info(node_class_info, './data/' + dataset + '/node_class.txt')
    save_graph(graph, './data/' + dataset + '/graph.adjlist')
    save_dense_graph(graph, './data/' + dataset + '/graph_dense.adjlist', node_type)


def train_email():
    node2id, graph, node_class_info = read_email_graph()
    sparse_node = select_sparse_node(graph)
    node_type = node_type_info(graph, sparse_node)
    train_node_candidate, test_node_candidate = generate_train_test_candidate_email(graph, sparse_node, node_type)
    train_node, test_node = generate_train_test_set(train_node_candidate, test_node_candidate)
    save_node(train_node, './data/' + dataset + '/train.csv')
    save_node(test_node, './data/' + dataset + '/test.csv')
    save_node_class_info(node_class_info, './data/' + dataset + '/node_class.txt')
    save_graph(graph, './data/' + dataset + '/graph.adjlist')
    save_dense_graph(graph, './data/' + dataset + '/graph_dense.adjlist', node_type)


def train_wiki():
    node2id, graph, node_class_info = read_wiki_graph()
    sparse_node = select_sparse_node(graph)
    node_type = node_type_info(graph, sparse_node)
    train_node_candidate, test_node_candidate = generate_train_test_candidate_wiki(graph, sparse_node, node_type)
    train_node, test_node = generate_train_test_set(train_node_candidate, test_node_candidate)
    save_node(train_node, './data/' + dataset + '/train.csv')
    save_node(test_node, './data/' + dataset + '/test.csv')
    save_node_class_info(node_class_info, './data/' + dataset + '/node_class.txt')
    save_graph(graph, './data/' + dataset + '/graph.adjlist')
    save_dense_graph(graph, './data/' + dataset + '/graph_dense.adjlist', node_type)


if __name__ == '__main__':
    if dataset == 'flickr':
        train_flickr()
    elif dataset == 'email':
        train_email()
    else:
        train_wiki()
